Support for Qwen-Image, Qwen-Image-Edit, and Qwen-Image-Edit-Plus

ai_diffusion/client.py

@@ -104,9 +104,11 @@ class Quantization(Enum):
     svdq = 1
 
     @staticmethod
-    def from_string(s: str):
+    def from_string(s: str, filename: str | None = None):
         if s == "svdq":
             return Quantization.svdq
+        elif filename and "qwen" in filename and "svdq" in filename:
+            return Quantization.svdq
         else:
             return Quantization.none
 
@@ -351,6 +353,7 @@ def resolve_arch(style: Style, client: Client | ClientModels | None = None):
         checkpoint = style.preferred_checkpoint(models.checkpoints.keys())
         if checkpoint != "not-found":
             arch = models.arch_of(checkpoint)
+
     elif style.checkpoints:
         arch = style.architecture.resolve(style.checkpoints[0])
 

ai_diffusion/comfy_client.py

@@ -412,7 +412,7 @@ 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", "none"), filename),
                     info.get("is_inpaint", False),
                     info.get("is_refiner", False),
                 )
@@ -625,7 +625,7 @@ def _find_text_encoder_models(model_list: Sequence[str]):
     kind = ResourceKind.text_encoder
     return {
         resource_id(kind, Arch.all, te): _find_model(model_list, kind, Arch.all, te)
-        for te in ["clip_l", "clip_g", "t5"]
+        for te in ["clip_l", "clip_g", "t5", "qwen"]
     }
 
 

ai_diffusion/comfy_workflow.py

@@ -499,12 +499,24 @@ def load_inpaint_model(self, model_name: str):
     def load_fooocus_inpaint(self, head: str, patch: str):
         return self.add_cached("INPAINT_LoadFooocusInpaint", 1, head=head, patch=patch)
 
-    def nunchaku_load_diffusion_model(self, model_path: str, cache_threshold: float):
+    def nunchaku_load_flux_diffusion_model(self, model_path: str, cache_threshold: float):
         return self.add_cached(
             "NunchakuFluxDiTLoader", 1, model_path=model_path, cache_threshold=cache_threshold
         )
 
-    def nunchaku_load_lora(self, model: Output, name: str, strength: float):
+    def nunchaku_load_qwen_diffusion_model(
+        self, model_name: str, cpu_offload: str, num_blocks_on_gpu: int, use_pin_memory: str
+    ):
+        return self.add_cached(
+            "NunchakuQwenImageDiTLoader",
+            1,
+            model_name=model_name,
+            cpu_offload=cpu_offload,
+            num_blocks_on_gpu=num_blocks_on_gpu,
+            use_pin_memory=use_pin_memory,
+        )
+
+    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
         )
@@ -587,6 +599,31 @@ def instruct_pix_to_pix_conditioning(
     def reference_latent(self, conditioning: Output, latent: Output):
         return self.add("ReferenceLatent", 1, conditioning=conditioning, latent=latent)
 
+    def text_encode_qwen_image_edit(
+        self, clip: Output, vae: Output | None, image: Output, prompt: str | Output
+    ):
+        return self.add(
+            "TextEncodeQwenImageEdit", 1, clip=clip, vae=vae, image=image, prompt=prompt
+        )
+
+    def text_encode_qwen_image_edit_plus(
+        self, clip: Output, vae: Output | None, images: list[Output], prompt: str | Output
+    ):
+        image1 = images[0] if len(images) > 0 else None
+        image2 = images[1] if len(images) > 1 else None
+        image3 = images[2] if len(images) > 2 else None
+
+        return self.add(
+            "TextEncodeQwenImageEditPlus",
+            1,
+            clip=clip,
+            vae=vae,
+            image1=image1,
+            image2=image2,
+            image3=image3,
+            prompt=prompt,
+        )
+
     def background_region(self, conditioning: Output):
         return self.add("ETN_BackgroundRegion", 1, conditioning=conditioning)
 

ai_diffusion/control.py

@@ -148,7 +148,7 @@ def _update_is_supported(self):
                     self.error_text = _("The server is missing the ClipVision model") + f" {search}"
                     is_supported = False
 
-            if self.mode.is_ip_adapter and models.arch is Arch.flux_k:
+            if self.mode.is_ip_adapter and models.arch.is_edit:
                 is_supported = True  # Reference images are merged into the conditioning context
             elif self.mode.is_ip_adapter and models.ip_adapter.find(self.mode) is None:
                 search_path = resources.search_path(ResourceKind.ip_adapter, models.arch, self.mode)