Skip to content

Support lazy-init Linear layers in Z-Image patching#817

Open
kj863257 wants to merge 2 commits into
nunchaku-ai:mainfrom
kj863257:main
Open

Support lazy-init Linear layers in Z-Image patching#817
kj863257 wants to merge 2 commits into
nunchaku-ai:mainfrom
kj863257:main

Conversation

@kj863257

@kj863257 kj863257 commented Apr 7, 2026

Copy link
Copy Markdown

Motivation

On recent ComfyUI Windows builds, some Linear modules can be lazily initialized and temporarily have weight=None
before model weights are fully loaded.

In the Z-Image loading path, ComfyUI-nunchaku patches attention and feed-forward modules before weight loading
completes. The current implementation assumes linear.weight is always available and directly accesses
linear.weight.dtype and linear.weight.device, which can raise:

AttributeError: 'NoneType' object has no attribute 'dtype'

This PR makes the Z-Image patching path compatible with lazy-initialized Linear modules while preserving the
existing behavior when real weights are already present.

This PR is paired with the corresponding nunchaku PR:

  • nunchaku: Support lazy-initialized Linear modules in SVDQW4A4Linear.from_linear
  • Link: <NUNCHAKU_PR_URL>

Modifications

  • Add a helper in models/zimage.py to resolve dtype and device from a possibly lazy-initialized Linear module.
  • Keep the existing behavior unchanged when linear.weight is available.
  • Fall back to the provided torch_dtype, or weight_comfy_model_dtype, and finally CPU when linear.weight is None.
  • Apply this logic to fuse_to_svdquant_linear() so the Z-Image feed-forward patching path no longer depends on
    linear.weight.dtype or linear.weight.device being immediately available.
  • Run pre-commit on models/zimage.py.

No new workflow test case is added because this PR does not introduce a new workflow, model asset, custom node, or
output behavior. It only makes the existing Z-Image patching path tolerant to lazy-initialized Linear modules.

Checklist

  • Code is formatted using Pre-Commit hooks (run pre-commit run --all-files).
  • Relevant unit tests are added in the tests/workflows directory following the guidance in
    the Contribution Guide.
  • Reference images are uploaded to PR comments and URLs are added to test_cases.json.
  • Additional test data (if needed) is registered in test_data/inputs.yaml.
  • Additional models (if needed) are registered in scripts/download_models.py
    and test_data/models.yaml.
  • Additional custom nodes (if needed) are added to [.github/workflows/pr-test.yaml](../.github/workflows/pr-
    test.yaml).
  • For reviewers: If you're only helping merge the main branch and haven't contributed code to this PR, please
    remove yourself as a co-author when merging.
  • Please feel free to join our Discord or [WeChat](https://huggingface.co/
    datasets/nunchaku-tech/cdn/blob/main/nunchaku/assets/wechat.jpg) to discuss your PR.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant