diff --git a/ai_diffusion/backend/comfy_client.py b/ai_diffusion/backend/comfy_client.py index 95d292934..e11029ee8 100644 --- a/ai_diffusion/backend/comfy_client.py +++ b/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 } diff --git a/ai_diffusion/backend/resources.py b/ai_diffusion/backend/resources.py index 2830ee27d..c3406a1fb 100644 --- a/ai_diffusion/backend/resources.py +++ b/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 = [ diff --git a/ai_diffusion/backend/workflow.py b/ai_diffusion/backend/workflow.py index 825cab01e..f1b7e00d4 100644 --- a/ai_diffusion/backend/workflow.py +++ b/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}") diff --git a/ai_diffusion/icons/sd-version-krea2-dark.svg b/ai_diffusion/icons/sd-version-krea2-dark.svg new file mode 100644 index 000000000..811fa2af3 --- /dev/null +++ b/ai_diffusion/icons/sd-version-krea2-dark.svg @@ -0,0 +1,11 @@ + + + + + + + + + + + diff --git a/ai_diffusion/icons/sd-version-krea2-light.svg b/ai_diffusion/icons/sd-version-krea2-light.svg new file mode 100644 index 000000000..657764064 --- /dev/null +++ b/ai_diffusion/icons/sd-version-krea2-light.svg @@ -0,0 +1,11 @@ + + + + + + + + + + + diff --git a/ai_diffusion/ui/theme.py b/ai_diffusion/ui/theme.py index 175fb0740..9b514863a 100644 --- a/ai_diffusion/ui/theme.py +++ b/ai_diffusion/ui/theme.py @@ -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") diff --git a/tests/test_resources.py b/tests/test_resources.py index 71d66a042..b07c47aec 100644 --- a/tests/test_resources.py +++ b/tests/test_resources.py @@ -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)