Add ERNIE Image model support

Author: VERIGEN

ai_diffusion/comfy_client.py

@@ -751,7 +751,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", "qwen", "qwen_3_06b", "qwen_3_4b", "qwen_3_8b"]
+        for te in ["clip_l", "clip_g", "t5", "qwen", "qwen_3_06b", "qwen_3_4b", "qwen_3_8b", "ministral"]
     }
 
 

ai_diffusion/comfy_workflow.py

@@ -613,7 +613,7 @@ def empty_latent_image(self, extent: Extent, arch: Arch, batch_size=1):
         w, h = extent.width, extent.height
         if arch.is_flux_like or arch.is_qwen_like or arch in (Arch.sd3, Arch.chroma, Arch.zimage):
             return self.add("EmptySD3LatentImage", 1, width=w, height=h, batch_size=batch_size)
-        if arch.is_flux2:
+        if arch.is_flux2 or arch is Arch.ernie:
             return self.add("EmptyFlux2LatentImage", 1, width=w, height=h, batch_size=batch_size)
         else:
             return self.add("EmptyLatentImage", 1, width=w, height=h, batch_size=batch_size)

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

@@ -97,6 +97,7 @@ class Arch(Enum):
     qwen_l = "Qwen Layered"
     anima = "Anima"
     zimage = "Z-Image"
+    ernie = "ERNIE Image"
 
     auto = "Automatic"
     all = "All"
@@ -139,6 +140,8 @@ def from_string(string: str, model_type: str = "eps", filename: str | None = Non
             return Arch.anima
         if string in {"z-image", "zimage"}:
             return Arch.zimage
+        if string in {"ernie-image", "ernie_image"}:
+            return Arch.ernie
         return None
 
     @staticmethod
@@ -244,6 +247,8 @@ def text_encoders(self):
                 return ["qwen_3_06b"]
             case Arch.zimage:
                 return ["qwen_3_4b"]
+            case Arch.ernie:
+                return ["ministral"]
         raise ValueError(f"Unsupported architecture: {self}")
 
     @staticmethod
@@ -265,6 +270,7 @@ def list():
             Arch.qwen_l,
             Arch.anima,
             Arch.zimage,
+            Arch.ernie,
         ]
 
 
@@ -799,6 +805,7 @@ def is_required(kind: ResourceKind, arch: Arch, identifier: ControlMode | Upscal
     resource_id(ResourceKind.text_encoder, Arch.all, "qwen_3_4b"): ["qwen_3_4b", "qwen3-4b", "qwen3_4b", "qwen_3", "qwen-3"],
     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.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"],
@@ -815,6 +822,7 @@ def is_required(kind: ResourceKind, arch: Arch, identifier: ControlMode | Upscal
     resource_id(ResourceKind.vae, Arch.qwen_l, "default"): ["qwen_image_layered_vae"],
     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"],
 }
 # fmt: on
 
@@ -848,6 +856,8 @@ def is_required(kind: ResourceKind, arch: Arch, identifier: ControlMode | Upscal
     ResourceId(ResourceKind.vae, Arch.zimage, "default"),
     ResourceId(ResourceKind.vae, Arch.flux2_4b, "default"),
     ResourceId(ResourceKind.vae, Arch.flux2_9b, "default"),
+    ResourceId(ResourceKind.text_encoder, Arch.ernie, "ministral"),
+    ResourceId(ResourceKind.vae, Arch.ernie, "default"),
 }
 
 recommended_resource_ids = [

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

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

ai_diffusion/resources.py

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

ai_diffusion/ui/theme.py

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