Add support for Moss Image Base

README.md

@@ -41,6 +41,7 @@ AI Toolkit is an easy to use all in one training suite for diffusion models. I t
 - [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)
+- [OpenMOSS-Team/MOSS-Image-Base](https://huggingface.co/OpenMOSS-Team/MOSS-Image-Base) (MOSS-Image Base)
 
 ### 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/moss_image/src/checkpoints.py

@@ -3,12 +3,14 @@
 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
+- ``OpenMOSS-Team/MOSS-Image-SFT`` / ``-Base`` (vendor, the UI defaults): 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.
+  qtype; int8 only for convrot). Otherwise Base downloads the repack's bf16
+  (byte-identical to the vendor shards); SFT, which the repack has no bf16
+  of, downloads the vendor's diffusers shards once and merges them into
+  ``diffusion_models/moss_image_sft_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.
@@ -64,6 +66,9 @@ COMFY_TRANSFORMER_FILES = {
     ),
 }
 REPO_VARIANT = {SFT_REPO: "sft", BASE_REPO: "base", COMFY_REPO: "sft"}
+# vendor repo -> the repack's file holding its weights byte-identical (same
+# name and layout the shard merge writes): downloaded instead of merging
+VENDOR_MIRROR = {BASE_REPO: "diffusion_models/moss_image_base_bf16.safetensors"}
 
 # (repo, file) candidates for the Qwen3.5-9B text encoder
 COMFY_TEXT_ENCODER_FILES = [
@@ -131,6 +136,13 @@ def comfy_transformer_file(
                 status_fn=status_fn,
                 qtype=qtype,
             )
+        if repo_id in VENDOR_MIRROR:
+            return resolve_comfy_candidates(
+                [VENDOR_MIRROR[repo_id]],
+                repo_id=COMFY_REPO,
+                hf_token=hf_token,
+                status_fn=status_fn,
+            )
 
     rel_path = comfy_transformer_rel(repo_id)
     found = resolve_comfy_file(rel_path, repo_id, local_only=True)

extensions_built_in/diffusion_models/ui.tsx

@@ -1797,6 +1797,50 @@ export const AI_TOOLKIT_UI_MODELS: ModelArch[] = [
     disableSections: ["network.conv"],
     additionalSections: ["model.layer_offloading", "model.low_vram"],
   },
+  {
+    name: "moss_image:base",
+    label: "MOSS-Image Base",
+    generateNameOverride: "MOSS-Image Base",
+    group: "image",
+    defaults: {
+      // default updates when [selected, unselected] in the UI
+      // the vendor repo: Comfy-Org/MOSS-Image base files on disk win, else the
+      // repack's bf16 (identical weights) is downloaded
+      "config.process[0].model.name_or_path": [
+        "OpenMOSS-Team/MOSS-Image-Base",
+        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",

testing/test_model_loading.py

@@ -226,6 +226,17 @@ MODEL_TESTS = {
         # vendor defaults are 50 steps, CFG 3 (negative branch skips block 22)
         "sample": {**IMG, "num_inference_steps": 20, "guidance_scale": 3.0},
     },
+    "moss_image:base": {
+        # vendor id; local Comfy-Org base files win, else the repack's bf16
+        "model": {
+            "name_or_path": "OpenMOSS-Team/MOSS-Image-Base",
+            "quantize": True,
+            "quantize_te": True,
+            "qtype": "convrot8",
+            "qtype_te": "convrot8",
+        },
+        "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

toolkit/models/registry.py

@@ -57,6 +57,7 @@ ARCH_REGISTRY: Dict[str, dict] = {
     "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}},
+    "moss_image:base": {"modality": "image", "model": {"name_or_path": "OpenMOSS-Team/MOSS-Image-Base", "quantize": True, "qtype": "convrot8", "quantize_te": True, "qtype_te": "convrot8", "model_kwargs": {"uncond_skip_layers": [22]}}, "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}},