z image dit loader node doesnt work

[abisssalina-droid]

i have built nunchaku from source for cu 130 torch 12.9 python 13

error log
Using pytorch attention in VAE
Using pytorch attention in VAE
VAE load device: cuda:0, offload device: cpu, dtype: torch.bfloat16
FETCH ComfyRegistry Data: 50/115
FETCH ComfyRegistry Data: 55/115
gguf qtypes: F32 (145), Q6_K (49), Q5_K (30), Q4_K (154), IQ4_XS (20)
FETCH ComfyRegistry Data: 60/115
CLIP/text encoder model load device: cuda:0, offload device: cpu, current: cpu, dtype: torch.float16
Requested to load ZImageTEModel_
loaded completely; 3615.68 MB usable, 3602.16 MB loaded, full load: True
FETCH ComfyRegistry Data: 65/115
model weight dtype torch.float16, manual cast: None
model_type FLOW
!!! Exception during processing !!!
Traceback (most recent call last):
File "D:\ai\ComfyUI_windows_portable_nvidia\ComfyUI_windows_portable\ComfyUI\execution.py", line 516, in execute
output_data, output_ui, has_subgraph, has_pending_tasks = await get_output_data(prompt_id, unique_id, obj, input_data_all, execution_block_cb=execution_block_cb, pre_execute_cb=pre_execute_cb, v3_data=v3_data)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "D:\ai\ComfyUI_windows_portable_nvidia\ComfyUI_windows_portable\ComfyUI\execution.py", line 330, in get_output_data
return_values = await _async_map_node_over_list(prompt_id, unique_id, obj, input_data_all, obj.FUNCTION, allow_interrupt=True, execution_block_cb=execution_block_cb, pre_execute_cb=pre_execute_cb, v3_data=v3_data)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "D:\ai\ComfyUI_windows_portable_nvidia\ComfyUI_windows_portable\ComfyUI\execution.py", line 304, in _async_map_node_over_list
await process_inputs(input_dict, i)
File "D:\ai\ComfyUI_windows_portable_nvidia\ComfyUI_windows_portable\ComfyUI\execution.py", line 292, in process_inputs
result = f(**inputs)
File "D:\ai\ComfyUI_windows_portable_nvidia\ComfyUI_windows_portable\ComfyUI\custom_nodes\ComfyUI-nunchaku\nodes\models\zimage.py", line 216, in load_model
model = _load(sd, metadata=metadata)
File "D:\ai\ComfyUI_windows_portable_nvidia\ComfyUI_windows_portable\ComfyUI\custom_nodes\ComfyUI-nunchaku\nodes\models\zimage.py", line 153, in _load
patch_scale_key(model.diffusion_model, patched_sd)
~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "D:\ai\ComfyUI_windows_portable_nvidia\ComfyUI_windows_portable\python_embeded\Lib\site-packages\nunchaku\models\transformers\utils.py", line 168, in patch_scale_key
assert state_dict[k].dtype == state_dict_from_checkpoint[k].dtype
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
AssertionError

[maksmaisak]

The node fails to import for me. I'm on the dev branch (this commit) on cuda 12.6, torch 2.9.1, python 3.13.6

Nunchaku version: 1.1.0.dev20251111
ComfyUI-nunchaku version: 1.1.0dev

(I also tried Nunchaku 1.0.0 and 1.0.2, the errors below are unaffected)

Nodes `NunchakuZImageDiTLoader` import failed:
Traceback (most recent call last):
  File "D:\AI\ComfyUI_windows_portable\ComfyUI\custom_nodes\ComfyUI-nunchaku-dev\__init__.py", line 144, in <module>
    from .nodes.models.zimage import NunchakuZImageDiTLoader
  File "D:\AI\ComfyUI_windows_portable\ComfyUI\custom_nodes\ComfyUI-nunchaku-dev\nodes\models\zimage.py", line 12, in <module>
    from nunchaku.models.transformers.utils import patch_scale_key
ImportError: cannot import name 'patch_scale_key' from 'nunchaku.models.transformers.utils' (D:\AI\ComfyUI_windows_portable\python_embeded\Lib\site-packages\nunchaku\models\transformers\utils.py)
'nunchaku_versions.json' not found. Node will start in minimal mode. Use 'update node' to fetch versions.

[trollver9000]

whole dev branch is buggy now and incompatible with newest comfuyi, let the devs solve this after they lookup new comfy code to make it compatible

[Ph0rk0z]

This actually fixed it for me without having to comment that stuff out: nunchaku-ai/nunchaku#833

I am on linux tho. Now there is lora support and more comfy fixes in the PRs. Eager to see if I can use this vs BF16 weights + cublas_ops.

[1506086927]

nunchaku\nunchaku\models\transformers\utils.py:

    for k in state_dict.keys():
        if k not in state_dict_from_checkpoint:
            assert ".wcscales" in k
            state_dict_from_checkpoint[k] = torch.ones_like(state_dict[k])
        else:
            if state_dict[k].dtype != state_dict_from_checkpoint[k].dtype:
                state_dict_from_checkpoint[k] = state_dict_from_checkpoint[k].to(state_dict[k].dtype)

[fukc-gihtub]

z-image-turbo seems to work with the latest nunchaku (v1.1.0)

[LiJT]

Nope it isn't working with nunchaku-1.1.0+torch2.9-cp313-cp313-win_amd64 wheel installed
and I tried with fp4 model of Z-image with latest 1.1.0 still failed

nunchaku\nunchaku\models\transformers\utils.py:

    for k in state_dict.keys():
        if k not in state_dict_from_checkpoint:
            assert ".wcscales" in k
            state_dict_from_checkpoint[k] = torch.ones_like(state_dict[k])
        else:
            if state_dict[k].dtype != state_dict_from_checkpoint[k].dtype:
                state_dict_from_checkpoint[k] = state_dict_from_checkpoint[k].to(state_dict[k].dtype)

I can run with this! but the generated result has four image in an array! WHy tho???

[Ph0rk0z]

It's working perfectly with lora support now when I merge the right PR.

I merge this into dev: #739

I merge this into nunchaku: nunchaku-ai/nunchaku#833

Same speed as cublas_ops + BF16. I don't have the visual degradation anymore either.

[trollver9000]

ok its way too early to install this and frankenstein the code, i wish comfy and nunchaku devs read each others code instead of fighting against each other with each new version.
Maybe 2 weeks from now it will be solved, for now this node does not work for z , z loader nunchaku does not show up.

[mcmonkey4eva]

Merging the PRs that @Ph0rk0z mentioned seems the right method, but for the record I was able to cheat to get it work by just editing the file in site-packages to comment out the bad assert

[Ph0rk0z]

When I edit out the assert, the quality was worse.

[DavideAlidosi]

nunchaku\nunchaku\models\transformers\utils.py:

    for k in state_dict.keys():
        if k not in state_dict_from_checkpoint:
            assert ".wcscales" in k
            state_dict_from_checkpoint[k] = torch.ones_like(state_dict[k])
        else:
            if state_dict[k].dtype != state_dict_from_checkpoint[k].dtype:
                state_dict_from_checkpoint[k] = state_dict_from_checkpoint[k].to(state_dict[k].dtype)

This fixed it on my side, but I can't make LoRAs work.

[m8rr]

nunchaku-ai/nunchaku#841 (comment)

In my case, it worked fine after disabling fast fp16

[praveenmaniyan]

@m8rr I use the command "python main.py --lowvram --use-sage-attention", and do not use any of the fast options. How can we disable fast fp16?

[Ph0rk0z]

simply do not add --fast fp16_accumulation