Add initial Krea 2 model support (#2541)

* Add initial Krea 2 model support

* ruff format

---------

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

Author: fukc-gihtub

ai_diffusion/backend/comfy_client.py

@@ -748,7 +748,17 @@ def find_model(model_list: Sequence[str], id: ResourceId):
 
 def _find_text_encoder_models(model_list: Sequence[str]):
     kind = ResourceKind.text_encoder
-    tes = ["clip_l", "clip_g", "t5", "qwen", "qwen_3_06b", "qwen_3_4b", "qwen_3_8b", "ministral"]
+    tes = [
+        "clip_l",
+        "clip_g",
+        "t5",
+        "qwen",
+        "qwen_3_06b",
+        "qwen_3_4b",
+        "qwen_3_8b",
+        "qwen3vl_4b",
+        "ministral",
+    ]
     return {
         resource_id(kind, Arch.all, te): _find_model(model_list, kind, Arch.all, te) for te in tes
     }

ai_diffusion/backend/resources.py

@@ -98,6 +98,7 @@ class Arch(Enum):
     anima = "Anima"
     zimage = "Z-Image"
     ernie = "ERNIE Image"
+    krea2 = "Krea 2"
 
     auto = "Automatic"
     all = "All"
@@ -142,6 +143,8 @@ def from_string(string: str, model_type: str = "eps", filename: str | None = Non
             return Arch.zimage
         if string in {"ernie-image", "ernie_image"}:
             return Arch.ernie
+        if string == "krea2":
+            return Arch.krea2
         return None
 
     @staticmethod
@@ -249,6 +252,8 @@ def text_encoders(self):
                 return ["qwen_3_4b"]
             case Arch.ernie:
                 return ["ministral"]
+            case Arch.krea2:
+                return ["qwen3vl_4b"]
         raise ValueError(f"Unsupported architecture: {self}")
 
     @staticmethod
@@ -271,6 +276,7 @@ def list():
             Arch.anima,
             Arch.zimage,
             Arch.ernie,
+            Arch.krea2,
         ]
 
 
@@ -818,6 +824,7 @@ def is_required(kind: ResourceKind, arch: Arch, identifier: ControlMode | Upscal
     resource_id(ResourceKind.text_encoder, Arch.all, "qwen_3_8b"): ["qwen_3_8b", "qwen3-8b", "qwen3_8b"],
     resource_id(ResourceKind.text_encoder, Arch.all, "qwen_3_06b"): ["qwen_3_06b", "qwen3-06b", "qwen3_06b"],
     resource_id(ResourceKind.text_encoder, Arch.all, "ministral"): ["ministral-3-3b", "ministral"],
+    resource_id(ResourceKind.text_encoder, Arch.all, "qwen3vl_4b"): ["qwen3vl_4b"],
     resource_id(ResourceKind.vae, Arch.sd15, "default"): ["vae-ft-mse-840000-ema"],
     resource_id(ResourceKind.vae, Arch.sdxl, "default"): ["sdxl_vae"],
     resource_id(ResourceKind.vae, Arch.illu, "default"): ["sdxl_vae"],
@@ -835,6 +842,7 @@ def is_required(kind: ResourceKind, arch: Arch, identifier: ControlMode | Upscal
     resource_id(ResourceKind.vae, Arch.anima, "default"): ["qwen_image"],
     resource_id(ResourceKind.vae, Arch.zimage, "default"): ["z-image", "flux-", "flux_", "flux/", "flux1", "ae.s"],
     resource_id(ResourceKind.vae, Arch.ernie, "default"): ["flux2"],
+    resource_id(ResourceKind.vae, Arch.krea2, "default"): ["qwen_image"],
 }
 # fmt: on
 
@@ -870,6 +878,8 @@ def is_required(kind: ResourceKind, arch: Arch, identifier: ControlMode | Upscal
     ResourceId(ResourceKind.vae, Arch.flux2_9b, "default"),
     ResourceId(ResourceKind.text_encoder, Arch.ernie, "ministral"),
     ResourceId(ResourceKind.vae, Arch.ernie, "default"),
+    ResourceId(ResourceKind.text_encoder, Arch.krea2, "qwen3vl_4b"),
+    ResourceId(ResourceKind.vae, Arch.krea2, "default"),
 }
 
 recommended_resource_ids = [

ai_diffusion/backend/workflow.py

@@ -175,6 +175,8 @@ def load_checkpoint_with_lora(w: ComfyWorkflow, checkpoint: CheckpointInput, mod
                 clip = w.load_clip(te["qwen_3_4b"], type="lumina2")
             case Arch.ernie:
                 clip = w.load_clip(te["ministral"], type="flux2")
+            case Arch.krea2:
+                clip = w.load_clip(te["qwen3vl_4b"], type="krea2")
             case _:
                 raise RuntimeError(f"No text encoder for model architecture {arch.name}")
 

ai_diffusion/icons/sd-version-krea2-dark.svg

@@ -86,6 +86,8 @@ def checkpoint_icon(arch: Arch, format: FileFormat | None = None, client: Client
         return icon("sd-version-anima")
     elif arch is Arch.ernie:
         return icon("sd-version-ernie")
+    elif arch is Arch.krea2:
+        return icon("sd-version-krea2")
     else:
         log.warning(f"Unresolved SD version {arch}, cannot fetch icon")
         return icon("warning")

ai_diffusion/icons/sd-version-krea2-light.svg

@@ -55,6 +55,7 @@ def test_resource_ids_exist():
             Arch.flux2_9b,
             Arch.anima,
             Arch.ernie,
+            Arch.krea2,
         ):
             continue  # no model downloads yet
         model = res.find_resource(resource_id)