Support nunchaku-z-image-turbo (#2237)

* Support nunchaku-z-image-turbo

* Don't check filename

---------

Co-authored-by: fukc <fukc@github.com>

Author: fukc-gihtub

ai_diffusion/comfy_workflow.py

@@ -568,6 +568,16 @@ def nunchaku_load_qwen_diffusion_model(self, model_name: str, num_blocks_on_gpu=
             use_pin_memory="disable",
         )
 
+    def nunchaku_load_zimage_diffusion_model(self, model_name: str, num_blocks_on_gpu=1):
+        return self.add_cached(
+            "NunchakuZImageDiTLoader",
+            1,
+            model_name=model_name,
+            cpu_offload="auto",
+            num_blocks_on_gpu=num_blocks_on_gpu,
+            use_pin_memory="disable",
+        )
+
     def nunchaku_load_flux_lora(self, model: Output, name: str, strength: float):
         return self.add(
             "NunchakuFluxLoraLoader", 1, model=model, lora_name=name, lora_strength=strength

ai_diffusion/workflow.py

@@ -105,6 +105,8 @@ def load_checkpoint_with_lora(w: ComfyWorkflow, checkpoint: CheckpointInput, mod
                 model = w.nunchaku_load_flux_diffusion_model(
                     model_info.filename, cache_threshold=cache
                 )
+            elif model_info.arch == Arch.zimage:
+                model = w.nunchaku_load_zimage_diffusion_model(model_info.filename)
             elif model_info.arch.is_qwen_like:
                 model = w.nunchaku_load_qwen_diffusion_model(model_info.filename)
             else: