Add support for Iris-3B

toolkit/models/registry.py

@@ -44,6 +44,7 @@ ARCH_REGISTRY: Dict[str, dict] = {
     "flux2_klein_4b": {"modality": "image", "model": {"name_or_path": "black-forest-labs/FLUX.2-klein-base-4B", "quantize_te": True}, "sample": dict(IMG)},
     "flux2_klein_9b": {"modality": "image", "model": {"name_or_path": "black-forest-labs/FLUX.2-klein-base-9B", "quantize": True, "quantize_te": True}, "sample": dict(IMG)},
     "prx_pixel": {"modality": "image", "model": {"name_or_path": "Photoroom/prxpixel-t2i", "quantize_te": True}, "sample": dict(IMG)},
+    "iris3b": {"modality": "image", "model": {"name_or_path": "speridlabs/iris-3b", "quantize_te": True}, "sample": {**IMG, "num_inference_steps": 30, "guidance_scale": 3.0}},
     "f-lite": {"modality": "image", "model": {"name_or_path": "Freepik/F-Lite", "quantize": True, "quantize_te": True}, "sample": dict(IMG)},
     "sd1": {"modality": "image", "model": {"name_or_path": "stable-diffusion-v1-5/stable-diffusion-v1-5"}, "sample": {"width": 512, "height": 512, "num_inference_steps": 20, "guidance_scale": 7.5}},
     "sdxl": {"modality": "image", "model": {"name_or_path": "stabilityai/stable-diffusion-xl-base-1.0"}, "sample": {**IMG, "guidance_scale": 6.0}},

toolkit/models/v2/text_encoders/qwen3_vl.py

@@ -44,6 +44,23 @@ class Qwen3VLTextEncoder(Qwen3VLForConditionalGeneration, OstrisTransformersMixi
     def get_transformer_block_names(cls):
         return ["model.language_model.layers"]
 
+    @classmethod
+    def convert_state_dict_on_load(cls, state_dict):
+        """Single-file (ComfyUI repack) checkpoints: accept the older comfy key
+        layout (``model.layers.*`` / ``visual.*``) next to the transformers one,
+        and materialize the tied ``lm_head`` the files omit."""
+        out = {}
+        for key, value in state_dict.items():
+            if key.startswith("model.layers.") or key in ("model.embed_tokens.weight", "model.norm.weight"):
+                key = "model.language_model." + key[len("model.") :]
+            elif key.startswith("visual."):
+                key = "model." + key
+            out[key] = value
+        embed = out.get("model.language_model.embed_tokens.weight")
+        if "lm_head.weight" not in out and embed is not None:
+            out["lm_head.weight"] = embed
+        return out
+
     def drop_vision_tower(self):
         """Text-only conditioning: the vision tower is dead weight — drop it to
         free VRAM and skip loading its (bf16-slow) Conv3d patch_embed."""