[FEAT] support for z-image-turbo

[devgdovg]

Motivation

Supports for https://huggingface.co/Tongyi-MAI/Z-Image-Turbo and maybe future Z-Image-Base & Z-Image-Edit

Modifications

Checklist

• Code is formatted using Pre-Commit hooks.
• Relevant unit tests are added in the tests directory following the guidance in Contribution Guide.
• Documentation and example scripts in examples are updated if necessary.
• Throughput/latency benchmarks and quality evaluations are included where applicable.
• 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 Slack, Discord or WeChat to discuss your PR.

[atgctg]

Hi, I would like to test the FP4 version. Could you please publish the SVDQ safetensors on Huggingface?

[juntaosun]

Could you please publish the SVDQ safetensors on Huggingface?

[lmxyy]

Could you please publish the SVDQ safetensors on Huggingface?

let me test first.

[devgdovg]

Test cases added. @lmxyy
But before model checkpoints are published to huggingface hub, we should modify the path value to a temporary local directory in the test cases in order to run.

[lmxyy]

Could you also update the docs correspondingly?

[devgdovg]

Could you also update the docs correspondingly?

ok.

[juntaosun]

@devgdovg Does it support Lora loading for z-image-turbo?

[praveenmaniyan]

@devgdovg really appreciate your efforts in getting this functionality work. Can we get an update on this implementation. Is the model ready, or by when can we have this feature available?

Thanks in advance.

[vgabbo]

I want to thank you for the contribution, really appreciate efforts. I will keep following PR, Z-turbo in nunchaku is very good news.

[juntaosun]

Is the model ready, or by when can we have this feature available?

[devgdovg]

@devgdovg Does it support Lora loading for z-image-turbo?

not yet in this PR.

[vgabbo]

Sorry to disturb, I was curious, is this faster than original one @devgdovg ? How much time does it take, compared to other models? Nunchaku is really strong, sometimes it offers specialized kernels and so I was curious to understand if they are used also in z-turbo. Thank you.

[devgdovg]

Sorry to disturb, I was curious, is this faster than original one @devgdovg ? How much time does it take, compared to other models? Nunchaku is really strong, sometimes it offers specialized kernels and so I was curious to understand if they are used also in z-turbo. Thank you.

According to the testing I've done so far, Nunchaku optimized z-image-turbo model is definitely faster compared to fp16/fp8 models provided by Comfy-Org. Am waiting for @lmxyy to review all the stuff, including code/docs/models.

[zwukong]

merge一下吧，模型已经上传了

[zwukong]

256版本报错

File "M:\ComfyUI\312_cu128\ComfyUI\custom_nodes\ComfyUI-nunchaku\nodes\models\zimage.py", line 153, in _load
    patch_scale_key(model.diffusion_model, patched_sd)
  File "M:\ComfyUI\312_cu128\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

[atgctg]

This should probably be set to num_inference_steps=9 to match the official example: https://github.com/Tongyi-MAI/Z-Image?tab=readme-ov-file#-quick-start

[lmxyy]

passing tests, merged

[Ph0rk0z]

256版本报错

File "M:\ComfyUI\312_cu128\ComfyUI\custom_nodes\ComfyUI-nunchaku\nodes\models\zimage.py", line 153, in _load
    patch_scale_key(model.diffusion_model, patched_sd)
  File "M:\ComfyUI\312_cu128\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

Same problem with 256. Will try 128. Skip assert did not work.

edit: 128 has the problem too. I commented out the code that calls this function and it inferences at least. also the FP16_downscale patch isn't compatible.

[badraymen]

Plz sir don't forget the controlnet of zimage in the pack

[Heliumrich]

Apparently the quality is poor.
Hopefully it will work better with non-distilled base model. (and with lora support ootb!)