Add support for MOSS Image

toolkit/models/registry.py

@@ -56,6 +56,7 @@ ARCH_REGISTRY: Dict[str, dict] = {
     "ltx2.3": {"modality": "video", "model": {"name_or_path": "Lightricks/LTX-2.3/ltx-2.3-22b-dev.safetensors", "quantize": True, "quantize_te": True}, "sample": {**VID, "num_inference_steps": 25, "guidance_scale": 3.0, "num_frames": 49, "fps": 24}},
     "ltx2.5": {"modality": "video", "model": {"name_or_path": "Lightricks/LTX-2.5", "quantize": True, "quantize_te": True}, "sample": {**VID, "num_inference_steps": 25, "guidance_scale": 3.0, "num_frames": 49, "fps": 24}},
     "flux3": {"modality": "image", "model": {"name_or_path": "black-forest-labs/flux-3-image-klein-base", "quantize": True, "qtype": "convrot8", "quantize_te": True, "qtype_te": "convrot8"}, "sample": {"width": 1024, "height": 1024, "num_inference_steps": 28, "guidance_scale": 3.5}},
+    "moss_image": {"modality": "image", "model": {"name_or_path": "OpenMOSS-Team/MOSS-Image-SFT", "quantize": True, "qtype": "convrot8", "quantize_te": True, "qtype_te": "convrot8"}, "sample": {"width": 1024, "height": 1024, "num_inference_steps": 30, "guidance_scale": 3.0}},
     "ace_step_15": {"modality": "audio", "model": {"name_or_path": "ostris/ace_step_1.5_ComfyUI_files/ace_step_1.5_base_aio.safetensors", "quantize": True, "quantize_te": True}, "sample": {"width": 512, "height": 512, "num_inference_steps": 20, "guidance_scale": 4.0}},
     "qwen25_omni": {"modality": "text", "model": {"name_or_path": "ai-toolkit/Qwen2.5-Omni-7B/qwen2_5_omni_7b_convrot8.safetensors", "quantize": True, "qtype": "convrot8"}, "sample": {"width": 512, "height": 512, "num_inference_steps": 1, "guidance_scale": 1.0}},
     "yue2": {"modality": "audio", "model": {"name_or_path": "Comfy-Org/YuE2/checkpoints/yue2_3b_int8_convrot.safetensors", "quantize": True, "qtype": "convrot8"}, "sample": {"width": 512, "height": 512, "num_inference_steps": 32, "guidance_scale": 1.0, "duration": 120}},

toolkit/models/v2/text_encoders/qwen35.py

@@ -0,0 +1,56 @@
+from transformers import Qwen3_5TextModel
+
+from .._mixin import OstrisTransformersMixin
+
+QWEN35_COMFY_REPO = "Comfy-Org/Qwen3.5"
+
+
+class Qwen35TextOnlyEncoder(Qwen3_5TextModel, OstrisTransformersMixin):
+    """The Qwen3.5 language tower alone (moss_image). Text-only inputs, so the
+    default arange positions match the multimodal model's text path. A stock
+    Qwen repo id resolves to the Comfy-Org single-file repack under
+    MODELS_PATH/text_encoders/ (config from the Qwen repo); its lm_head, mtp
+    head and vision tower are dropped on load."""
+
+    aitk_subfolder = "text_encoder"
+    aitk_tokenizer_subfolder = "tokenizer"
+    # the int8 repacks keep A_log / gated-norm weights in fp32; one dtype
+    # throughout, like the bf16 path
+    aitk_cast_quantized_load = True
+
+    aitk_comfy_repo = QWEN35_COMFY_REPO
+    aitk_comfy_weight_names = {
+        "Qwen/Qwen3.5-2B": ["text_encoders/qwen3.5_2b_bf16.safetensors"],
+        "Qwen/Qwen3.5-4B": ["text_encoders/qwen3.5_4b_bf16.safetensors"],
+        "Qwen/Qwen3.5-9B": ["text_encoders/qwen3.5_9b_bf16.safetensors"],
+    }
+
+    _PREFIXES = ("model.language_model.", "language_model.", "model.")
+    _DROPPED = ("lm_head.", "mtp.", "model.visual.", "visual.")
+
+    @classmethod
+    def get_transformer_block_names(cls):
+        return ["layers"]
+
+    @classmethod
+    def aitk_load_config(cls, path, subfolder=None):
+        config = super().aitk_load_config(path, subfolder=subfolder)
+        # the Qwen repos carry the multimodal Qwen3_5Config
+        return getattr(config, "text_config", config)
+
+    @classmethod
+    def convert_state_dict_on_load(cls, state_dict):
+        out = {}
+        for key, value in state_dict.items():
+            if key.startswith(cls._DROPPED):
+                continue
+            for prefix in cls._PREFIXES:
+                if key.startswith(prefix):
+                    key = key[len(prefix) :]
+                    break
+            out[key] = value
+        return out
+
+    @classmethod
+    def convert_state_dict_on_save(cls, state_dict):
+        return {f"model.language_model.{k}": v for k, v in state_dict.items()}