Support nunchaku-z-image-turbo

Author: fukc

ai_diffusion/comfy_client.py

@@ -478,7 +478,9 @@ def parse_model_info(models: dict, model_format: FileFormat):
                 (
                     filename,
                     Arch.from_string(info["base_model"], info.get("type", "eps"), filename),
-                    Quantization.from_string(info.get("quant", "none")),
+                    Quantization.from_string(info.get("quant",
+                        "svdq" if filename.startswith("svdq-") else "none")
+                    ),
                     info.get("is_inpaint", False),
                     info.get("is_refiner", False),
                 )

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: