Add support for MOSS Image

README.md

@@ -11,6 +11,7 @@ AI Toolkit is an easy to use all in one training suite for diffusion models. I t
 - [black-forest-labs/FLUX.2-dev](https://huggingface.co/black-forest-labs/FLUX.2-dev) (FLUX.2)
 - [black-forest-labs/FLUX.2-klein-base-4B](https://huggingface.co/black-forest-labs/FLUX.2-klein-base-4B) (FLUX.2-klein-base-4B)
 - [black-forest-labs/FLUX.2-klein-base-9B](https://huggingface.co/black-forest-labs/FLUX.2-klein-base-9B) (FLUX.2-klein-base-9B)
+- [black-forest-labs/flux-3-image-klein-base](https://huggingface.co/black-forest-labs/flux-3-image-klein-base) (FLUX.3 Image Klein)
 - [ostris/Flex.1-alpha](https://huggingface.co/ostris/Flex.1-alpha) (Flex.1)
 - [ostris/Flex.2-preview](https://huggingface.co/ostris/Flex.2-preview) (Flex.2)
 - [lodestones/Chroma1-Base](https://huggingface.co/lodestones/Chroma1-Base) (Chroma)
@@ -39,6 +40,7 @@ AI Toolkit is an easy to use all in one training suite for diffusion models. I t
 - [krea/Krea-2-Turbo](https://huggingface.co/krea/Krea-2-Turbo) (Krea 2 Turbo)
 - [inclusionAI/Ming-Image-0.1-Design](https://huggingface.co/inclusionAI/Ming-Image-0.1-Design) (Ming-Image 0.1 Design)
 - [microsoft/Mage-Flow-Base](https://huggingface.co/microsoft/Mage-Flow-Base) (Mage-Flow)
+- [OpenMOSS-Team/MOSS-Image-SFT](https://huggingface.co/OpenMOSS-Team/MOSS-Image-SFT) (MOSS-Image SFT)
 
 ### Instruction / Edit
 - [black-forest-labs/FLUX.1-Kontext-dev](https://huggingface.co/black-forest-labs/FLUX.1-Kontext-dev) (FLUX.1-Kontext-dev)

extensions_built_in/diffusion_models/__init__.py

@@ -24,6 +24,7 @@ from .boogu_image import BooguImageModel, BooguImageEditModel
 from .mageflow import MageFlowModel, MageFlowEditModel
 from .minimax_h3 import MinimaxH3Model, MinimaxH3Ref2VAModel, MinimaxH3FastModel
 from .flux3 import Flux3Model
+from .moss_image import MossImageModel
 
 AI_TOOLKIT_MODELS = [
     # put a list of models here
@@ -67,4 +68,5 @@ AI_TOOLKIT_MODELS = [
     MinimaxH3Ref2VAModel,
     MinimaxH3FastModel,
     Flux3Model,
+    MossImageModel,
 ]

extensions_built_in/diffusion_models/moss_image/__init__.py

@@ -0,0 +1,3 @@
+from .moss_image import MossImageModel
+
+__all__ = ["MossImageModel"]

extensions_built_in/diffusion_models/moss_image/moss_image.py

@@ -0,0 +1,400 @@
+"""MOSS-Image (OpenMOSS-Team/MOSS-Image-SFT) for ai-toolkit.
+
+  - DiT: 42-block single-stream transformer, 3072 wide, 128 latent channels
+    (``src/transformer.py``, Comfy-Org/MOSS-Image key layout)
+  - VAE: FLUX.2, 2x2 packed and BatchNorm normalized (the toolkit's BFL
+    ``AutoEncoder``; same packing as MOSS). The shared comfy flux2-vae file
+    (full decoder) replaces the vendor's small-decoder copy by default
+  - text: Qwen3.5-9B final-normed last hidden state of the chat-templated
+    prompt (thinking on, generation prompt appended), one unpadded forward per
+    prompt; stored per item as (L, 4096) in an ``AdvancedPromptEmbeds``
+
+``model.name_or_path`` defaults to the vendor repo; ``Comfy-Org/MOSS-Image``
+(the repack), a local diffusers dir or a ``.safetensors`` path (local or
+``org/repo/file``) also load. Comfy files on disk take precedence for the
+vendor ids; see ``src/checkpoints.py`` for sourcing. The text encoder is
+stock Qwen3.5-9B from a comfy file.
+
+Flow convention matches ai-toolkit (sigma 1 = noise, target = noise - clean).
+"""
+
+import os
+from typing import TYPE_CHECKING, List, Optional
+
+import torch
+from safetensors import safe_open
+from safetensors.torch import load_file
+from transformers import Qwen2Tokenizer
+
+from toolkit.accelerator import unwrap_model
+from toolkit.advanced_prompt_embeds import AdvancedPromptEmbeds
+from toolkit.basic import flush
+from toolkit.config_modules import GenerateImageConfig, ModelConfig
+from toolkit.metadata import get_meta_for_safetensors
+from toolkit.models.base_model import BaseModel
+from toolkit.models.v2.text_encoders.qwen35 import Qwen35TextOnlyEncoder
+from toolkit.models.v2.vae.flux2_kl import AutoEncoder, convert_diffusers_state_dict
+from toolkit.samplers.custom_flowmatch_sampler import (
+    CustomFlowMatchEulerDiscreteScheduler,
+)
+
+from .src.checkpoints import (
+    COMFY_REPO,
+    QWEN35_9B_REPO,
+    SFT_REPO,
+    VENDOR_VAE_FILE,
+    comfy_text_encoder_file,
+    comfy_transformer_file,
+    comfy_vae_file,
+    resolve_hub_file,
+)
+from .src.pipeline import (
+    DEFAULT_SHIFT,
+    DEFAULT_UNCOND_SKIP_LAYERS,
+    MossImagePipeline,
+    stack_text_embeds,
+)
+from .src.transformer import MossImageTransformer
+
+if TYPE_CHECKING:
+    from toolkit.data_transfer_object.data_loader import DataLoaderBatchDTO
+
+HF_TOKEN = os.getenv("HF_TOKEN", None)
+MAX_TEXT_TOKENS = 4096
+
+scheduler_config = {
+    "num_train_timesteps": 1000,
+    "use_dynamic_shifting": False,
+}
+
+
+def _is_hub_repo(path: str) -> bool:
+    return (
+        not os.path.exists(path)
+        and not path.endswith(".safetensors")
+        and path.count("/") == 1
+    )
+
+
+class MossImageModel(BaseModel):
+    arch = "moss_image"
+
+    def __init__(
+        self,
+        device,
+        model_config: ModelConfig,
+        dtype="bf16",
+        custom_pipeline=None,
+        noise_scheduler=None,
+        **kwargs,
+    ):
+        super().__init__(
+            device, model_config, dtype, custom_pipeline, noise_scheduler, **kwargs
+        )
+        self.is_flow_matching = True
+        self.is_transformer = True
+        self.target_lora_modules = ["MossImageTransformer"]
+        self.supports_model_paths = True
+        self.use_old_lokr_format = False
+        self.lora_keys_use_comfy_prefix = True
+        self.latent_space_version = "moss_image"
+
+        kw = self.model_config.model_kwargs
+        self.max_text_length = int(kw.get("max_text_length", MAX_TEXT_TOKENS))
+        self.sample_shift = float(kw.get("sample_shift", DEFAULT_SHIFT))
+        self.uncond_skip_layers = tuple(
+            kw.get("uncond_skip_layers", DEFAULT_UNCOND_SKIP_LAYERS) or ()
+        )
+
+    @staticmethod
+    def get_train_scheduler():
+        return CustomFlowMatchEulerDiscreteScheduler(
+            shift=DEFAULT_SHIFT, **scheduler_config
+        )
+
+    def get_bucket_divisibility(self):
+        return 16  # f8 conv + 2x2 packing
+
+    # ------------------------------------------------------------------
+    # Loading
+    # ------------------------------------------------------------------
+    def _load_transformer(self) -> MossImageTransformer:
+        self.print_and_status_update("Loading transformer (MossImageTransformer)")
+        base = self.model_config.name_or_path
+        qtype = self.model_config.qtype if self.model_config.quantize else None
+        if os.path.isdir(base):
+            folder = os.path.join(base, "transformer")
+            if not os.path.isdir(folder):
+                folder = base
+            state_dict = {}
+            for name in sorted(os.listdir(folder)):
+                if name.endswith(".safetensors"):
+                    with safe_open(os.path.join(folder, name), "pt") as f:
+                        for key in f.keys():
+                            state_dict[key] = f.get_tensor(key)
+            if not state_dict:
+                raise FileNotFoundError(f"no MOSS-Image transformer weights in {folder}")
+            return MossImageTransformer.load_from_state_dict(state_dict, self.torch_dtype)
+        if base.endswith(".safetensors"):
+            path = resolve_hub_file(
+                base, hf_token=HF_TOKEN, status_fn=self.print_and_status_update
+            )
+        else:
+            path = comfy_transformer_file(
+                base,
+                qtype=qtype,
+                hf_token=HF_TOKEN,
+                status_fn=self.print_and_status_update,
+            )
+        self.print_and_status_update(f" - transformer weights: {path}")
+        return MossImageTransformer.load_model(
+            path, dtype=self.torch_dtype, qtype=qtype, quantize_on_load=False
+        )
+
+    def _load_text_encoder(self):
+        te_path = self.model_config.te_name_or_path
+        if te_path is None:
+            te_path = self.model_config.model_kwargs.get(
+                "text_encoder_path", QWEN35_9B_REPO
+            )
+        te_kwargs = self.component_load_kwargs("te")
+        config_path = None
+        if te_path == QWEN35_9B_REPO and te_kwargs["use_comfy_weights"]:
+            te_path = comfy_text_encoder_file(
+                qtype=te_kwargs["qtype"],
+                hf_token=HF_TOKEN,
+                status_fn=self.print_and_status_update,
+            )
+        elif te_path.endswith(".safetensors"):
+            te_path = resolve_hub_file(
+                te_path, hf_token=HF_TOKEN, status_fn=self.print_and_status_update
+            )
+        if te_path.endswith(".safetensors"):
+            config_path = QWEN35_9B_REPO
+        self.print_and_status_update(f"Loading Qwen3.5 text encoder from {te_path}")
+        text_encoder = Qwen35TextOnlyEncoder.load_model(
+            te_path,
+            dtype=te_kwargs["dtype"],
+            config_path=config_path,
+            subfolder="",
+            qtype=te_kwargs["qtype"],  # comfy candidate ranking only
+            quantize_on_load=False,
+            use_comfy_weights=te_kwargs["use_comfy_weights"],
+        )
+        text_encoder.eval()
+        text_encoder.requires_grad_(False)
+        return text_encoder
+
+    def _load_tokenizer(self):
+        base = self.model_config.name_or_path
+        source = SFT_REPO
+        if os.path.isdir(os.path.join(base, "tokenizer")) or (
+            _is_hub_repo(base) and base != COMFY_REPO
+        ):
+            source = base
+        # the tokenizer's declared class: AutoTokenizer also resolves a
+        # config.json the folder lacks, which fails offline
+        return Qwen35TextOnlyEncoder._local_first(
+            Qwen2Tokenizer.from_pretrained, source, subfolder="tokenizer"
+        )
+
+    def _load_vae(self) -> AutoEncoder:
+        self.print_and_status_update("Loading VAE")
+        base = self.model_config.name_or_path
+        path = self.model_config.vae_path
+        if path is None and os.path.isdir(base):
+            local = os.path.join(base, VENDOR_VAE_FILE)
+            if os.path.exists(local):
+                path = local
+        if path is None:
+            path = comfy_vae_file(
+                hf_token=HF_TOKEN, status_fn=self.print_and_status_update
+            )
+        elif not os.path.exists(path):
+            path = resolve_hub_file(
+                path, hf_token=HF_TOKEN, status_fn=self.print_and_status_update
+            )
+        state_dict = load_file(path)
+        if any(k.startswith("encoder.down_blocks.") for k in state_dict):
+            state_dict = convert_diffusers_state_dict(state_dict)
+        # config sniffed from the checkpoint (small decoder)
+        vae = AutoEncoder.load_from_state_dict(state_dict, self.vae_torch_dtype)
+        return vae.eval().requires_grad_(False)
+
+    def load_model(self):
+        self.print_and_status_update("Loading MOSS-Image model")
+
+        transformer = self._load_transformer()
+        # quantize + offload + placement, all driven by model_config
+        transformer.aitk_post_load(**self.component_load_kwargs("transformer"))
+        flush()
+
+        text_encoder = self._load_text_encoder()
+        text_encoder.aitk_post_load(**self.component_load_kwargs("te"))
+        flush()
+        tokenizer = self._load_tokenizer()
+
+        vae = self._load_vae()
+        vae.to(self.vae_device_torch)
+
+        self.noise_scheduler = MossImageModel.get_train_scheduler()
+
+        self.vae = vae
+        self.text_encoder = text_encoder
+        self.tokenizer = tokenizer
+        self.model = transformer
+        self.pipeline = MossImagePipeline(self)
+        self.print_and_status_update("Model Loaded")
+
+    # ------------------------------------------------------------------
+    # Generation (training previews)
+    # ------------------------------------------------------------------
+    def get_generation_pipeline(self):
+        return MossImagePipeline(self)
+
+    def generate_single_image(
+        self,
+        pipeline: MossImagePipeline,
+        gen_config: GenerateImageConfig,
+        conditional_embeds: AdvancedPromptEmbeds,
+        unconditional_embeds: AdvancedPromptEmbeds,
+        generator: torch.Generator,
+        extra: dict,
+    ):
+        if self.model.device == torch.device("cpu"):
+            self.model.to(self.device_torch)
+
+        sc = self.get_bucket_divisibility()
+        gen_config.width = int(gen_config.width // sc * sc)
+        gen_config.height = int(gen_config.height // sc * sc)
+
+        result = pipeline(
+            conditional_embeds=conditional_embeds,
+            unconditional_embeds=unconditional_embeds,
+            height=gen_config.height,
+            width=gen_config.width,
+            num_inference_steps=gen_config.num_inference_steps,
+            guidance_scale=gen_config.guidance_scale,
+            latents=gen_config.latents,
+            generator=generator,
+            shift=self.sample_shift,
+            uncond_skip_layers=self.uncond_skip_layers,
+        )
+        return result[0]
+
+    # ------------------------------------------------------------------
+    # Training hooks
+    # ------------------------------------------------------------------
+    def get_noise_prediction(
+        self,
+        latent_model_input: torch.Tensor,  # (B, 128, h, w)
+        timestep: torch.Tensor,  # 0..1000 scale
+        text_embeddings: AdvancedPromptEmbeds,
+        batch: "DataLoaderBatchDTO" = None,
+        **kwargs,
+    ):
+        device = self.device_torch
+        dtype = self.torch_dtype
+        if self.model.device == torch.device("cpu"):
+            self.model.to(device)
+
+        batch_size = latent_model_input.shape[0]
+        embeds_list = text_embeddings.text_embeds
+        if len(embeds_list) == 1 and batch_size > 1:
+            embeds_list = embeds_list * batch_size
+        with torch.no_grad():
+            embeds, mask = stack_text_embeds(embeds_list, device, dtype)
+            t = timestep.to(device, dtype=torch.float32).reshape(-1)
+            if t.shape[0] != batch_size:
+                t = t.expand(batch_size)
+
+        return self.transformer(
+            latent_model_input.to(device, dtype), t, embeds, mask
+        )
+
+    @torch.no_grad()
+    def get_prompt_embeds(self, prompt) -> AdvancedPromptEmbeds:
+        if isinstance(prompt, str):
+            prompt = [prompt]
+        if self.text_encoder.device == torch.device("cpu"):
+            self.text_encoder.to(self.device_torch)
+
+        embeds = []
+        # one unpadded forward per prompt, like the vendor (padding would leak
+        # into the gated-delta recurrence)
+        for text in prompt:
+            formatted = self.tokenizer.apply_chat_template(
+                [{"role": "user", "content": text}],
+                tokenize=False,
+                add_generation_prompt=True,
+                enable_thinking=True,
+            )
+            input_ids = self.tokenizer(
+                formatted,
+                return_tensors="pt",
+                truncation=True,
+                max_length=self.max_text_length,
+            )["input_ids"].to(self.device_torch)
+            out = self.text_encoder(input_ids=input_ids, use_cache=False)
+            embeds.append(out.last_hidden_state[0].to(self.torch_dtype))
+        return AdvancedPromptEmbeds(text_embeds=embeds)
+
+    def get_loss_target(self, *args, **kwargs):
+        noise = kwargs.get("noise")
+        batch = kwargs.get("batch")
+        return (noise - batch.latents).detach()
+
+    def get_model_has_grad(self):
+        return False
+
+    def get_te_has_grad(self):
+        return False
+
+    # ------------------------------------------------------------------
+    # VAE encode / decode
+    # ------------------------------------------------------------------
+    @torch.no_grad()
+    def encode_images(self, image_list, device=None, dtype=None):
+        """Images (C, H, W) in [-1, 1] -> normalized latents (B, 128, H/16, W/16)."""
+        if device is None:
+            device = self.vae_device_torch
+        if dtype is None:
+            dtype = self.vae_torch_dtype
+        if self.vae.device == torch.device("cpu"):
+            self.vae.to(device)
+        batch = torch.stack([img.to(device, dtype=self.vae.dtype) for img in image_list])
+        return self.vae.encode(batch).to(device, dtype=dtype)
+
+    @torch.no_grad()
+    def decode_latents(self, latents: torch.Tensor, device=None, dtype=None):
+        if device is None:
+            device = self.vae_device_torch
+        if dtype is None:
+            dtype = self.vae_torch_dtype
+        if self.vae.device == torch.device("cpu"):
+            self.vae.to(device)
+        images = self.vae.decode(latents.to(device, self.vae.dtype))
+        return images.to(device, dtype)
+
+    # ------------------------------------------------------------------
+    # Saving / bookkeeping
+    # ------------------------------------------------------------------
+    def save_model(self, output_path, meta, save_dtype):
+        if not output_path.endswith(".safetensors"):
+            output_path = output_path + ".safetensors"
+        transformer: MossImageTransformer = unwrap_model(self.model)
+        transformer.save_model(
+            output_path,
+            dtype=save_dtype,
+            metadata=get_meta_for_safetensors(meta, name=self.arch),
+            quantized=False,
+        )
+
+    def get_base_model_version(self):
+        return "moss_image"
+
+    def get_transformer_block_names(self) -> Optional[List[str]]:
+        return MossImageTransformer.get_transformer_block_names()
+
+    def get_quantization_exclude_modules(self):
+        return MossImageTransformer.get_quantization_exclude_modules()

extensions_built_in/diffusion_models/moss_image/src/checkpoints.py

@@ -0,0 +1,274 @@
+"""Where MOSS-Image weights come from.
+
+Weights live in the ComfyUI layout under MODELS_PATH, named like the
+Comfy-Org/MOSS-Image repack (DiT: the vendor's keys with q/k/v fused).
+
+- ``OpenMOSS-Team/MOSS-Image-SFT`` / ``-Base`` (vendor, the default): a
+  repack DiT file already on disk wins (best precision for the requested
+  qtype); otherwise the vendor's diffusers shards are downloaded once and
+  merged into ``diffusion_models/moss_image_{sft,base}_bf16.safetensors`` in
+  the Comfy layout. Other hub repos (fine-tunes) merge into a file named
+  after the repo, never substituted.
+- ``Comfy-Org/MOSS-Image``: the repack's SFT DiT (int8 convrot only; a local
+  merged bf16 SFT also counts).
+- ``org/repo/path/file.safetensors``: that file, downloaded to its comfy spot.
+
+The vendor text_encoder/ is stock Qwen3.5-9B (A_log / gated norm stored in
+bf16), so the text encoder is a comfy file: the repack's int8 convrot or
+Comfy-Org/Qwen3.5's bf16, whichever is local (ranked for the requested qtype),
+else the one matching the qtype is downloaded. The VAE is the shared FLUX.2
+file: same encoder + BatchNorm stats as the vendor vae/ (which is the BFL
+weights rounded to bf16), full decoder in place of the small one. The
+tokenizer comes from the vendor repo.
+"""
+
+import json
+import math
+import os
+import shutil
+import struct
+from typing import Callable, List, Optional
+
+import torch
+from safetensors import safe_open
+
+from toolkit.models.v2.resolver import (
+    comfy_precision_rank,
+    resolve_comfy_candidates,
+    resolve_comfy_file,
+)
+from toolkit.paths import MODELS_PATH
+
+SFT_REPO = "OpenMOSS-Team/MOSS-Image-SFT"
+BASE_REPO = "OpenMOSS-Team/MOSS-Image-Base"
+COMFY_REPO = "Comfy-Org/MOSS-Image"
+QWEN35_9B_REPO = "Qwen/Qwen3.5-9B"
+
+VENDOR_VAE_FILE = "vae/diffusion_pytorch_model.safetensors"
+VENDOR_TRANSFORMER_INDEX = "transformer/diffusion_pytorch_model.safetensors.index.json"
+COMFY_VAE_FILE = "vae/flux2-vae.safetensors"
+
+# variant -> (files in the repack, extra local-only candidates). The merged
+# vendor SFT is the bf16 SFT the repack does not ship.
+COMFY_TRANSFORMER_FILES = {
+    "sft": (
+        ["diffusion_models/moss_image_sft_int8_convrot_attn.safetensors"],
+        ["diffusion_models/moss_image_sft_bf16.safetensors"],
+    ),
+    "base": (
+        [
+            "diffusion_models/moss_image_base_int8_convrot_attn.safetensors",
+            "diffusion_models/moss_image_base_bf16.safetensors",
+        ],
+        [],
+    ),
+}
+REPO_VARIANT = {SFT_REPO: "sft", BASE_REPO: "base", COMFY_REPO: "sft"}
+
+# (repo, file) candidates for the Qwen3.5-9B text encoder
+COMFY_TEXT_ENCODER_FILES = [
+    (COMFY_REPO, "text_encoders/qwen3.5_9b_int8_convrot.safetensors"),
+    ("Comfy-Org/Qwen3.5", "text_encoders/qwen3.5_9b_bf16.safetensors"),
+]
+
+
+def comfy_transformer_rel(repo_id: str) -> str:
+    name = repo_id.split("/")[-1].lower().replace("-", "_").replace(".", "_")
+    return f"diffusion_models/{name}_bf16.safetensors"
+
+
+def _staging_dir(category: str) -> str:
+    # same filesystem as the target so the final rename is a move; removed
+    # afterwards together with the hub's .cache bookkeeping
+    return os.path.join(MODELS_PATH, category, ".moss_image_download")
+
+
+def resolve_hub_file(
+    ref: str,
+    hf_token: Optional[str] = None,
+    status_fn: Optional[Callable[[str], None]] = None,
+) -> str:
+    """A local path, or 'org/repo/rel/file.safetensors' found at / downloaded
+    to its comfy spot under MODELS_PATH."""
+    if os.path.isfile(ref):
+        return ref
+    parts = ref.split("/")
+    if len(parts) < 3:
+        raise FileNotFoundError(
+            f"{ref} is not a local file or an 'org/repo/path/file.safetensors' reference"
+        )
+    return resolve_comfy_candidates(
+        ["/".join(parts[2:])],
+        repo_id="/".join(parts[:2]),
+        hf_token=hf_token,
+        status_fn=status_fn,
+    )
+
+
+def comfy_transformer_file(
+    repo_id: str,
+    qtype: Optional[str] = None,
+    hf_token: Optional[str] = None,
+    status_fn: Optional[Callable[[str], None]] = None,
+) -> str:
+    variant = REPO_VARIANT.get(repo_id)
+    if variant is not None:
+        repo_files, local_extra = COMFY_TRANSFORMER_FILES[variant]
+        candidates = repo_files + local_extra
+        if repo_id != COMFY_REPO and not (qtype or "").lower().startswith("convrot"):
+            # int8 only stands in for the vendor bf16 when it attaches as-is
+            candidates = [c for c in candidates if "int8" not in c]
+        found = resolve_comfy_candidates(
+            candidates, repo_id=COMFY_REPO, local_only=True, qtype=qtype
+        )
+        if found is not None:
+            return found
+        if repo_id == COMFY_REPO:
+            return resolve_comfy_candidates(
+                repo_files,
+                repo_id=COMFY_REPO,
+                hf_token=hf_token,
+                status_fn=status_fn,
+                qtype=qtype,
+            )
+
+    rel_path = comfy_transformer_rel(repo_id)
+    found = resolve_comfy_file(rel_path, repo_id, local_only=True)
+    if found is not None:
+        return found
+    from huggingface_hub import hf_hub_download
+
+    target = os.path.join(MODELS_PATH, rel_path)
+    staging = _staging_dir("diffusion_models")
+    if status_fn is not None:
+        status_fn(f"Downloading the {repo_id} transformer to {target}")
+    index_path = hf_hub_download(
+        repo_id, VENDOR_TRANSFORMER_INDEX, token=hf_token, local_dir=staging
+    )
+    with open(index_path) as f:
+        shard_names = sorted(set(json.load(f)["weight_map"].values()))
+    shards = [
+        hf_hub_download(
+            repo_id, f"transformer/{name}", token=hf_token, local_dir=staging
+        )
+        for name in shard_names
+    ]
+    merge_safetensors(shards, target, fuse_qkv=True)
+    shutil.rmtree(staging, ignore_errors=True)
+    return target
+
+
+def comfy_text_encoder_file(
+    qtype: Optional[str] = None,
+    hf_token: Optional[str] = None,
+    status_fn: Optional[Callable[[str], None]] = None,
+) -> str:
+    rels = [rel for _, rel in COMFY_TEXT_ENCODER_FILES]
+    found = resolve_comfy_candidates(rels, repo_id=COMFY_REPO, local_only=True, qtype=qtype)
+    if found is not None:
+        return found
+    repo_id, rel = min(
+        COMFY_TEXT_ENCODER_FILES, key=lambda c: comfy_precision_rank(c[1], qtype=qtype)
+    )
+    return resolve_comfy_candidates(
+        [rel], repo_id=repo_id, hf_token=hf_token, status_fn=status_fn
+    )
+
+
+def comfy_vae_file(
+    hf_token: Optional[str] = None,
+    status_fn: Optional[Callable[[str], None]] = None,
+) -> str:
+    return resolve_comfy_candidates(
+        [COMFY_VAE_FILE], repo_id=COMFY_REPO, hf_token=hf_token, status_fn=status_fn
+    )
+
+
+_CAST_TO_BF16 = ("F32", "F16", "F64")
+
+
+def merge_safetensors(
+    shards: List[str],
+    out_path: str,
+    fuse_qkv: bool = False,
+    chunk_bytes: int = 64 << 20,
+):
+    """Concatenate safetensors shards into one bf16 file. bf16 (and non-float)
+    tensors are copied as raw bytes; other float tensors are cast to bf16.
+    ``fuse_qkv`` writes each ``*.proj.{q,k,v}.weight`` triple as one
+    ``*.proj.qkv.weight`` (row concat of contiguous tensors = byte concat).
+    Written to ``.part`` and renamed, so ``out_path`` may be one of the shards."""
+    tensors = {}
+    for path in shards:
+        with open(path, "rb") as f:
+            header_len = struct.unpack("<Q", f.read(8))[0]
+            header = json.loads(f.read(header_len))
+        header.pop("__metadata__", None)
+        for name, meta in header.items():
+            start, end = meta["data_offsets"]
+            tensors[name] = (meta["dtype"], meta["shape"], path, 8 + header_len + start, end - start)
+
+    # out name -> (dtype, shape, [(src name, path, start, size)])
+    entries = {}
+    for name, (dtype, shape, path, start, size) in tensors.items():
+        if fuse_qkv and name.endswith((".proj.k.weight", ".proj.v.weight")):
+            continue
+        if fuse_qkv and name.endswith(".proj.q.weight"):
+            base = name[: -len("q.weight")]
+            parts = [base + f"{p}.weight" for p in ("q", "k", "v")]
+            srcs = [tensors[p] for p in parts]
+            if any(s[0] != dtype or s[1][1:] != shape[1:] for s in srcs):
+                raise ValueError(f"{base}q/k/v parts do not fuse")
+            fused_shape = [sum(s[1][0] for s in srcs)] + list(shape[1:])
+            entries[base + "qkv.weight"] = (
+                dtype,
+                fused_shape,
+                [(p, s[2], s[3], s[4]) for p, s in zip(parts, srcs)],
+            )
+            continue
+        entries[name] = (dtype, shape, [(name, path, start, size)])
+
+    out_header, offset = {}, 0
+    for name in sorted(entries):
+        dtype, shape, parts = entries[name]
+        size = sum(p[3] for p in parts)
+        if dtype in _CAST_TO_BF16:
+            dtype, size = "BF16", 2 * math.prod(shape)
+        out_header[name] = {
+            "dtype": dtype,
+            "shape": shape,
+            "data_offsets": [offset, offset + size],
+        }
+        offset += size
+    out_header["__metadata__"] = {"format": "pt"}
+    blob = json.dumps(out_header, separators=(",", ":")).encode("utf-8")
+    blob += b" " * (-len(blob) % 8)
+
+    os.makedirs(os.path.dirname(out_path), exist_ok=True)
+    part_path = out_path + ".part"
+    handles = {}
+    try:
+        with open(part_path, "wb") as out:
+            out.write(struct.pack("<Q", len(blob)))
+            out.write(blob)
+            for name in sorted(entries):
+                dtype, _, parts = entries[name]
+                for src_name, path, start, size in parts:
+                    if dtype in _CAST_TO_BF16:
+                        with safe_open(path, "pt") as f:
+                            tensor = f.get_tensor(src_name).to(torch.bfloat16)
+                        out.write(tensor.contiguous().view(torch.uint8).numpy().tobytes())
+                        continue
+                    src = handles.get(path)
+                    if src is None:
+                        src = handles[path] = open(path, "rb")
+                    src.seek(start)
+                    remaining = size
+                    while remaining:
+                        chunk = src.read(min(chunk_bytes, remaining))
+                        out.write(chunk)
+                        remaining -= len(chunk)
+    finally:
+        for src in handles.values():
+            src.close()
+    os.replace(part_path, out_path)

extensions_built_in/diffusion_models/moss_image/src/pipeline.py

@@ -0,0 +1,143 @@
+"""MOSS-Image sampling (training previews / inference engine).
+
+Matches the vendor pipeline: shift-5 sigma grid snapped to the nearest of the
+1000 training sigmas, fp32 Euler on fp32 noise, CFG 3 whose negative branch
+skips DiT block 22 (two forwards per step), empty negative prompt. Flow
+convention is the toolkit's: sigma 1 = noise, the DiT predicts noise - clean.
+"""
+
+from typing import List, Optional, Sequence
+
+import torch
+import torch.nn.functional as F
+from PIL import Image
+from diffusers.utils.torch_utils import randn_tensor
+
+DEFAULT_SHIFT = 5.0
+DEFAULT_GUIDANCE = 3.0
+DEFAULT_STEPS = 30
+DEFAULT_UNCOND_SKIP_LAYERS = (22,)
+NUM_TRAIN_TIMESTEPS = 1000
+
+
+def moss_sigmas(
+    num_steps: int,
+    shift: float = DEFAULT_SHIFT,
+    training_shift: float = 1.0,
+    total: int = NUM_TRAIN_TIMESTEPS,
+) -> torch.Tensor:
+    """The vendor scheduler's sigmas (MOSSImageScheduler.set_timesteps) with
+    the terminal 0 appended: a shifted grid snapped to the training grid,
+    consecutive duplicates removed (so fewer than num_steps are possible)."""
+
+    def grid(count, s):
+        lower = s / (total + s - 1)
+        values = torch.linspace(1, lower, count)
+        return s * values / (1 + (s - 1) * values)
+
+    sigmas = grid(num_steps, shift)
+    training = grid(total, training_shift).flip(0)
+    upper = torch.searchsorted(training, sigmas).clamp(max=total - 1)
+    lower = (upper - 1).clamp(min=0)
+    nearest = torch.where(
+        (sigmas - training[lower]).abs() <= (training[upper] - sigmas).abs(),
+        lower,
+        upper,
+    )
+    sigmas = torch.unique_consecutive(training[nearest])
+    return torch.cat([sigmas, sigmas.new_zeros(1)])
+
+
+def stack_text_embeds(
+    embeds: List[torch.Tensor], device: torch.device, dtype: torch.dtype
+):
+    """Per-item (L_i, D) prompt states -> right-padded (B, S, D) + (B, S) mask."""
+    length = max(e.shape[0] for e in embeds)
+    stacked = torch.stack(
+        [F.pad(e.to(device, dtype), (0, 0, 0, length - e.shape[0])) for e in embeds]
+    )
+    mask = torch.stack(
+        [
+            torch.arange(length, device=device) < e.shape[0]
+            for e in embeds
+        ]
+    )
+    return stacked, mask
+
+
+class MossImagePipeline:
+    """Lightweight sampler used by ai-toolkit's preview generation."""
+
+    def __init__(self, model):
+        # the MossImageModel, for its transformer + encode/decode
+        self.model = model
+
+    @property
+    def device(self):
+        return self.model.device_torch
+
+    def to(self, *args, **kwargs):
+        return self
+
+    def set_progress_bar_config(self, **kwargs):
+        pass
+
+    @torch.no_grad()
+    def __call__(
+        self,
+        conditional_embeds,
+        unconditional_embeds,
+        height: int = 1024,
+        width: int = 1024,
+        num_inference_steps: int = DEFAULT_STEPS,
+        guidance_scale: float = DEFAULT_GUIDANCE,
+        latents: Optional[torch.Tensor] = None,
+        generator: Optional[torch.Generator] = None,
+        shift: float = DEFAULT_SHIFT,
+        uncond_skip_layers: Sequence[int] = DEFAULT_UNCOND_SKIP_LAYERS,
+        **kwargs,
+    ):
+        model = self.model
+        device = model.device_torch
+        dtype = model.torch_dtype
+        transformer = model.transformer
+
+        do_cfg = unconditional_embeds is not None and guidance_scale != 1.0
+        cond, cond_mask = stack_text_embeds(conditional_embeds.text_embeds, device, dtype)
+        if do_cfg:
+            uncond, uncond_mask = stack_text_embeds(
+                unconditional_embeds.text_embeds, device, dtype
+            )
+
+        shape = (1, transformer.params.latent_channels, height // 16, width // 16)
+        if latents is None:
+            # fp32 noise: bf16 noise promoted afterwards cannot regain precision
+            latents = randn_tensor(shape, generator=generator, device=device, dtype=torch.float32)
+        latents = latents.to(device, torch.float32)
+
+        sigmas = moss_sigmas(num_inference_steps, shift=shift).to(device)
+        num_steps = len(sigmas) - 1
+        for step_index in range(num_steps):
+            sigma, sigma_next = sigmas[step_index], sigmas[step_index + 1]
+            timestep = (sigma * NUM_TRAIN_TIMESTEPS).reshape(1)
+            # CFG combines fp32 predictions even though the DiT runs in bf16
+            pred = transformer(latents, timestep, cond, cond_mask).float()
+            if do_cfg:
+                pred_uncond = transformer(
+                    latents,
+                    timestep,
+                    uncond,
+                    uncond_mask,
+                    skip_layers=uncond_skip_layers,
+                ).float()
+                pred = pred_uncond + guidance_scale * (pred - pred_uncond)
+            if getattr(model, "sample_step_hook", None) is not None:
+                # inference engine preview: x0 estimate
+                model._emit_sample_step(latents - sigma * pred, step_index, num_steps)
+            latents = latents + (sigma_next - sigma) * pred
+
+        images = model.decode_latents(latents, device=device, dtype=torch.float32)
+        images = images.float().clamp(-1.0, 1.0)
+        images = ((images + 1.0) * 127.5).round().to(torch.uint8)
+        images = images.permute(0, 2, 3, 1).cpu().numpy()
+        return [Image.fromarray(arr) for arr in images]

extensions_built_in/diffusion_models/moss_image/src/transformer.py

@@ -0,0 +1,340 @@
+"""MOSS-Image DiT (OpenMOSS-Team/MOSS-Image-SFT / -Base).
+
+Port of ``MOSSImageTransformer`` from the vendor's ``pipeline_moss_image.py``
+in the Comfy-Org/MOSS-Image key layout: the vendor's names with each block's
+q/k/v fused into ``proj.qkv`` (vendor shards and split community repacks are
+fused on load), so Comfy files load as-is and LoRA keys match ComfyUI. One stream: image tokens followed by text tokens
+through 42 blocks that share one AdaLN modulation; the text first goes through
+a 2-block refiner. The vendor evaluates every item alone at its true text
+length; here a batch runs together with padded text masked out of every
+attention (identical results) and the unmasked path when lengths match.
+Activations follow the weight dtype exactly like the vendor's bf16 inference:
+fp32 timestep sinusoid and RoPE phases, everything else in the compute dtype.
+"""
+
+from dataclasses import dataclass
+from typing import Iterable, Optional
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+from torch.utils.checkpoint import checkpoint
+
+from toolkit.models.v2._mixin import OstrisModelMixin
+
+
+@dataclass
+class MossImageParams:
+    dim: int = 3072
+    ffn_dim: int = 8192
+    num_heads: int = 24
+    num_layers: int = 42
+    text_dim: int = 4096
+    freq_dim: int = 256
+    latent_channels: int = 128
+    rope_base: float = 10000.0
+
+
+class _RMSNorm(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x):
+        y = F.rms_norm(x.float(), (x.shape[-1],), eps=1e-6)
+        return y.to(x.dtype) * self.weight
+
+
+def _situ(hidden, gate):
+    return (4.0 * torch.tanh(gate / 4.0) * torch.sigmoid(gate)) * (
+        25.0 * torch.tanh(hidden / 25.0)
+    )
+
+
+def _rotate(x, rope):
+    sin, cos = rope.chunk(2, -1)
+    sin = torch.cat((sin, sin), -1).unsqueeze(-2)
+    cos = torch.cat((cos, cos), -1).unsqueeze(-2)
+    a, b = x.chunk(2, -1)
+    return x * cos + torch.cat((-b, a), -1) * sin
+
+
+def _attention(q, k, v, rope, heads, mask):
+    shape = (*q.shape[:2], heads, q.shape[-1] // heads)
+    q = _rotate(q.reshape(shape), rope).transpose(1, 2)
+    k = _rotate(k.reshape(shape), rope).transpose(1, 2)
+    v = v.reshape(shape).transpose(1, 2)
+    out = F.scaled_dot_product_attention(q.to(v.dtype), k.to(v.dtype), v, attn_mask=mask)
+    return out.transpose(1, 2).flatten(2)
+
+
+class _TextBlock(nn.Module):
+    def __init__(self, dim, ffn_dim, heads):
+        super().__init__()
+        self.heads = heads
+        self.norm1, self.norm2 = _RMSNorm(dim), _RMSNorm(dim)
+        self.q_norm, self.k_norm = _RMSNorm(dim // heads), _RMSNorm(dim // heads)
+        self.qkv = nn.Linear(dim, 3 * dim)
+        self.proj = nn.Linear(dim, dim)
+        self.ffn_gate, self.ffn_up = nn.Linear(dim, ffn_dim), nn.Linear(dim, ffn_dim)
+        self.ffn_down = nn.Linear(ffn_dim, dim)
+
+    def forward(self, x, rope, mask):
+        q, k, v = self.qkv(self.norm1(x)).chunk(3, -1)
+        q = self.q_norm(q.unflatten(-1, (self.heads, -1))).flatten(2)
+        k = self.k_norm(k.unflatten(-1, (self.heads, -1))).flatten(2)
+        x = x + self.proj(_attention(q, k, v, rope, self.heads, mask))
+        y = self.norm2(x)
+        return x + self.ffn_down(_situ(self.ffn_up(y), self.ffn_gate(y)))
+
+
+class _TextProjection(nn.Module):
+    def __init__(self, text_dim, dim, ffn_dim, heads):
+        super().__init__()
+        self.proj = nn.Linear(text_dim, dim)
+        self.blocks = nn.ModuleList([_TextBlock(dim, ffn_dim, heads) for _ in range(2)])
+
+
+class _Block(nn.Module):
+    def __init__(self, dim, ffn_dim, heads):
+        super().__init__()
+        self.heads = heads
+        self.norm1, self.norm2 = _RMSNorm(dim), _RMSNorm(dim)
+        self.q_norm, self.k_norm = _RMSNorm(dim // heads), _RMSNorm(dim // heads)
+        # q/k/v fused (rows q, k, v) like the ComfyUI files; the vendor keeps
+        # them split (fused on load)
+        self.proj = nn.ModuleDict(
+            {
+                "qkv": nn.Linear(dim, 3 * dim, bias=False),
+                "o": nn.Linear(dim, dim, bias=False),
+                "ffn_up_gate": nn.Linear(dim, 2 * ffn_dim, bias=False),
+                "ffn_down": nn.Linear(ffn_dim, dim, bias=False),
+            }
+        )
+
+    def forward(self, x, rope, modulation, mask):
+        shift_a, scale_a, gate_a, shift_f, scale_f, gate_f = modulation.unbind(1)
+        y = self.norm1(x) * (1 + scale_a[:, None]) + shift_a[:, None]
+        q, k, v = self.proj["qkv"](y).chunk(3, -1)
+        q = self.q_norm(q.unflatten(-1, (self.heads, -1))).flatten(2)
+        k = self.k_norm(k.unflatten(-1, (self.heads, -1))).flatten(2)
+        attn = _attention(q, k, v, rope, self.heads, mask)
+        x = x + self.proj["o"](attn) * gate_a[:, None]
+        y = self.norm2(x) * (1 + scale_f[:, None]) + shift_f[:, None]
+        hidden, gate = self.proj["ffn_up_gate"](y).chunk(2, -1)
+        return x + self.proj["ffn_down"](_situ(hidden, gate)) * gate_f[:, None]
+
+
+class MossImageTransformer(nn.Module, OstrisModelMixin):
+    @classmethod
+    def aitk_config_from_state_dict(cls, state_dict):
+        # sniffed from unquantized tensors only (block linears may be int8)
+        video = state_dict["video_embedder.weight"]
+        dim = video.shape[0]
+        return MossImageParams(
+            dim=dim,
+            ffn_dim=state_dict["text_embedder.blocks.0.ffn_up.weight"].shape[0],
+            num_heads=dim // state_dict["layers.0.q_norm.weight"].shape[0],
+            num_layers=sum(
+                1
+                for k in state_dict
+                if k.startswith("layers.") and k.endswith(".norm1.weight")
+            ),
+            text_dim=state_dict["text_embedder.proj.weight"].shape[1],
+            freq_dim=state_dict["time_embedding.0.weight"].shape[1],
+            latent_channels=video.shape[1],
+        )
+
+    @classmethod
+    def convert_state_dict_on_load(cls, state_dict):
+        # derived constants; recomputed in fp32 every forward
+        state_dict = {k: v for k, v in state_dict.items() if k != "rope.bands"}
+        fuse_split_qkv(state_dict)
+        return state_dict
+
+    @classmethod
+    def convert_state_dict_on_save(cls, state_dict):
+        # the vendor's strict loader expects the (unused) frequency table
+        state_dict = dict(state_dict)
+        head_dim = state_dict["layers.0.q_norm.weight"].shape[0]
+        state_dict["rope.bands"] = _rope_bands("cpu", head_dim, MossImageParams.rope_base)
+        return state_dict
+
+    @classmethod
+    def get_transformer_block_names(cls):
+        return ["layers"]
+
+    @classmethod
+    def get_quantization_exclude_modules(cls):
+        # embedders, modulation, text refiner and output head stay full precision
+        return [
+            "video_embedder",
+            "text_embedder*",
+            "time_embedding*",
+            "shared_block_modulation",
+            "video_head_modulation",
+            "final_linear_video",
+        ]
+
+    def __init__(self, params: Optional[MossImageParams] = None):
+        super().__init__()
+        if params is None:
+            params = MossImageParams()
+        if params.dim % params.num_heads or (params.dim // params.num_heads) % 8:
+            raise ValueError("dim must divide into heads with head_dim divisible by 8")
+        self.params = params
+        dim, ffn_dim, heads = params.dim, params.ffn_dim, params.num_heads
+        self.video_embedder = nn.Linear(params.latent_channels, dim)
+        self.text_embedder = _TextProjection(params.text_dim, dim, ffn_dim, heads)
+        self.time_embedding = nn.Sequential(
+            nn.Linear(params.freq_dim, dim), nn.SiLU(), nn.Linear(dim, dim)
+        )
+        self.shared_block_modulation = nn.Linear(dim, 6 * dim)
+        self.layers = nn.ModuleList(
+            [_Block(dim, ffn_dim, heads) for _ in range(params.num_layers)]
+        )
+        self.final_norm_video = _RMSNorm(dim)
+        self.final_linear_video = nn.Linear(dim, params.latent_channels, bias=False)
+        self.video_head_modulation = nn.Linear(dim, 2 * dim)
+        self.gradient_checkpointing = False
+        self._bands = {}
+
+    @property
+    def device(self):
+        return self.final_norm_video.weight.device
+
+    @property
+    def dtype(self):
+        # norms are never quantized; the block linears may be
+        return self.final_norm_video.weight.dtype
+
+    def enable_gradient_checkpointing(self, enable: bool = True):
+        self.gradient_checkpointing = enable
+
+    def disable_gradient_checkpointing(self):
+        self.gradient_checkpointing = False
+
+    def _ckpt(self, module, *args):
+        if self.gradient_checkpointing and torch.is_grad_enabled():
+            return checkpoint(module, *args, use_reentrant=False)
+        return module(*args)
+
+    def _rope(self, lengths, height, width, text_len, device):
+        """(B, H*W + S, head_dim) fp32 rope; coordinates (r, f, h, w): image
+        tokens at (0, item text length, centered h, centered w), text token i
+        at (0, i, 0, 0)."""
+        bands = self._bands.get(device)
+        if bands is None:
+            bands = _rope_bands(
+                device, self.params.dim // self.params.num_heads, self.params.rope_base
+            )
+            self._bands[device] = bands
+        batch = lengths.shape[0]
+        h, w = torch.meshgrid(
+            torch.arange(height, device=device, dtype=torch.float32),
+            torch.arange(width, device=device, dtype=torch.float32),
+            indexing="ij",
+        )
+        image = torch.zeros(batch, height * width, 4, device=device)
+        image[..., 1] = lengths.float()[:, None]
+        image[..., 2] = (h - (height - 1) / 2).flatten()
+        image[..., 3] = (w - (width - 1) / 2).flatten()
+        text = torch.zeros(batch, text_len, 4, device=device)
+        text[..., 1] = torch.arange(text_len, device=device, dtype=torch.float32)
+        phase = torch.cat((image, text), 1).unsqueeze(-1) * bands
+        return torch.cat((phase.sin(), phase.cos()), -2).flatten(-2)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,  # (B, latent_channels, H, W)
+        timestep: torch.Tensor,  # 0..1000, 1000 = pure noise
+        encoder_hidden_states: torch.Tensor,  # (B, S, text_dim)
+        encoder_attention_mask: Optional[torch.Tensor] = None,  # (B, S) bool
+        skip_layers: Optional[Iterable[int]] = None,
+    ):
+        skip = frozenset(skip_layers or ())
+        batch, _, height, width = hidden_states.shape
+        num_image = height * width
+        device = hidden_states.device
+        dtype = self.dtype
+
+        context = encoder_hidden_states
+        mask = encoder_attention_mask
+        if mask is None:
+            mask = torch.ones(context.shape[:2], dtype=torch.bool, device=device)
+        mask = mask.to(device=device, dtype=torch.bool)
+        lengths = mask.sum(1)
+        if bool((lengths == 0).any()):
+            raise ValueError("each sample needs at least one valid text token")
+        text_len = int(lengths.max())
+        padded = bool((lengths < text_len).any())
+        if padded or not bool(mask[:, :text_len].all()):
+            # valid tokens first, in order (the vendor selects them by mask)
+            order = torch.argsort((~mask).to(torch.int8), dim=1, stable=True)
+            context = torch.gather(
+                context, 1, order[..., None].expand(-1, -1, context.shape[-1])
+            )
+        context = context[:, :text_len]
+        valid = torch.arange(text_len, device=device)[None] < lengths[:, None]
+
+        times = torch.as_tensor(timestep, device=device, dtype=torch.float32)
+        times = times.reshape(-1).expand(batch)
+        half = self.params.freq_dim // 2
+        freq = 10000 ** (
+            -torch.arange(half, device=device, dtype=torch.float64) / half
+        )
+        phase = times.double()[:, None] * freq
+        embedding = torch.cat((phase.cos(), phase.sin()), -1).float()
+        time = self.time_embedding(embedding.to(dtype)).float()
+        mod = self.shared_block_modulation(time.to(dtype)).view(batch, 6, -1)
+        head = self.video_head_modulation(time.to(dtype)).view(batch, 2, -1)
+
+        rope = self._rope(lengths, height, width, text_len, device).to(dtype)
+
+        text_mask = joint_mask = None
+        if padded:
+            text_mask = valid[:, None, None, :]
+            joint_mask = torch.cat(
+                (torch.ones(batch, num_image, dtype=torch.bool, device=device), valid), 1
+            )[:, None, None, :]
+
+        video = hidden_states.flatten(2).transpose(1, 2).to(dtype)
+        video = self.video_embedder(video.contiguous())
+        text = self.text_embedder.proj(context.to(dtype))
+        text_rope = rope[:, num_image:]
+        for block in self.text_embedder.blocks:
+            text = self._ckpt(block, text, text_rope, text_mask)
+
+        x = torch.cat((video, text), 1)
+        for index, layer in enumerate(self.layers):
+            if index not in skip:
+                x = self._ckpt(layer, x, rope, mod, joint_mask)
+
+        x = self.final_norm_video(x[:, :num_image])
+        x = x * (1 + head[:, 1, None]) + head[:, 0, None]
+        x = self.final_linear_video(x)
+        return x.transpose(1, 2).reshape(batch, -1, height, width)
+
+
+def fuse_split_qkv(state_dict):
+    """Vendor / community layout (``layers.N.proj.{q,k,v}``, plain or comfy
+    quantized) -> the fused ComfyUI ``layers.N.proj.qkv``. In place."""
+    from toolkit.util.comfy_quant_import import fuse_split_quantized_keys
+
+    for key in [k for k in state_dict if k.endswith(".proj.q.weight")]:
+        base = key[: -len(".q.weight")]
+        parts = [f"{base}.{p}" for p in ("q", "k", "v")]
+        if f"{base}.q.comfy_quant" in state_dict:
+            fuse_split_quantized_keys(state_dict, parts, f"{base}.qkv")
+        else:
+            state_dict[f"{base}.qkv.weight"] = torch.cat(
+                [state_dict.pop(f"{p}.weight") for p in parts], dim=0
+            )
+    return state_dict
+
+
+def _rope_bands(device, head_dim: int, base: float):
+    # built on cpu in fp32 like the vendor's training, then moved
+    num_bands = head_dim // 8
+    exponent = torch.arange(num_bands, dtype=torch.float32) / num_bands
+    return (1 / (base**exponent)).to(device)

extensions_built_in/diffusion_models/ui.tsx

@@ -1754,6 +1754,49 @@ export const AI_TOOLKIT_UI_MODELS: ModelArch[] = [
       "model.qie.match_target_res",
     ],
   },
+  {
+    name: "moss_image",
+    label: "MOSS-Image SFT",
+    group: "image",
+    defaults: {
+      // default updates when [selected, unselected] in the UI
+      // the vendor repo: its transformer is merged once into the comfy models
+      // folder; a .safetensors path (fine-tune / int8 convrot repack) loads as named
+      "config.process[0].model.name_or_path": [
+        "OpenMOSS-Team/MOSS-Image-SFT",
+        defaultNameOrPath,
+      ],
+      "config.process[0].model.quantize": [true, false],
+      "config.process[0].model.qtype": ["convrot8", "qfloat8"],
+      // Qwen3.5-9B text encoder
+      "config.process[0].model.quantize_te": [true, false],
+      "config.process[0].model.qtype_te": ["convrot8", "qfloat8"],
+      "config.process[0].model.low_vram": [true, false],
+      "config.process[0].sample.sampler": ["flowmatch", "flowmatch"],
+      "config.process[0].train.noise_scheduler": ["flowmatch", "flowmatch"],
+      "config.process[0].sample.width": [1024, 1024],
+      "config.process[0].sample.height": [1024, 1024],
+      // vendor defaults: CFG 3 with block 22 skipped on the negative branch
+      "config.process[0].sample.guidance_scale": [3, 4],
+      "config.process[0].sample.sample_steps": [30, 25],
+      "config.process[0].train.timestep_type": ["shift", "sigmoid"],
+      // sampling: the CFG negative branch skips these DiT blocks (zero-based)
+      "config.process[0].model.model_kwargs": [
+        {
+          uncond_skip_layers: [22],
+        },
+        {},
+      ],
+      // keep the LoRA off the block the negative branch skips, so CFG's
+      // skipped-block contrast stays the base model's
+      "config.process[0].network.network_kwargs.ignore_if_contains": [
+        ["layers.22."],
+        [],
+      ],
+    },
+    disableSections: ["network.conv"],
+    additionalSections: ["model.layer_offloading", "model.low_vram"],
+  },
   {
     name: "flux2_klein_4b",
     label: "FLUX.2-klein-base-4B",

extensions_built_in/inference_engine/latent_preview.py

@@ -51,6 +51,7 @@ ARCH_LATENT_FORMAT = {
     "cogview4": "Flux",
     "prx_pixel": None,  # pixel space
     "flux3": "Flux3Image",
+    "moss_image": "Flux2",  # FLUX.2 encoder, same packing
     "iris3b": None,  # pixel space
 }
 

testing/test_model_loading.py

@@ -215,6 +215,30 @@ MODEL_TESTS = {
         "sample": {**IMG, "num_inference_steps": 12, "guidance_scale": 1.0},
         "needs_control_image": True,
     },
+    "moss_image": {
+        "model": {
+            "name_or_path": "OpenMOSS-Team/MOSS-Image-SFT",
+            "quantize": True,
+            "quantize_te": True,
+            "qtype": "convrot8",
+            "qtype_te": "convrot8",
+        },
+        # vendor defaults are 50 steps, CFG 3 (negative branch skips block 22)
+        "sample": {**IMG, "num_inference_steps": 20, "guidance_scale": 3.0},
+    },
+    "moss_image:comfy": {
+        # the Comfy-Org repack files named outright (a bare repo id prefers any
+        # local bf16): int8 convrot SFT DiT + Qwen3.5-9B TE attach as-is
+        "model": {
+            "name_or_path": "Comfy-Org/MOSS-Image/diffusion_models/moss_image_sft_int8_convrot_attn.safetensors",
+            "te_name_or_path": "Comfy-Org/MOSS-Image/text_encoders/qwen3.5_9b_int8_convrot.safetensors",
+            "quantize": True,
+            "quantize_te": True,
+            "qtype": "convrot8",
+            "qtype_te": "convrot8",
+        },
+        "sample": {**IMG, "num_inference_steps": 20, "guidance_scale": 3.0},
+    },
     # ---- legacy monolith archs (components adopted into v2 on load) ----
     "sd1": {
         "model": {"name_or_path": "stable-diffusion-v1-5/stable-diffusion-v1-5"},

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()}

version.py

@@ -1 +1 @@
-VERSION = "0.13.26"
+VERSION = "0.13.27"