Add MOSS-Image SFT edit training

extensions_built_in/diffusion_models/ui.tsx

@@ -1841,6 +1841,59 @@ export const AI_TOOLKIT_UI_MODELS: ModelArch[] = [
     disableSections: ["network.conv"],
     additionalSections: ["model.layer_offloading", "model.low_vram"],
   },
+  {
+    name: "moss_image:edit",
+    label: "MOSS-Image SFT [Edit Training]",
+    generateNameOverride: "MOSS-Image SFT Edit",
+    group: "experimental",
+    defaults: {
+      // default updates when [selected, unselected] in the UI
+      "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 + vision tower (references ride in the prompt)
+      "config.process[0].model.quantize_te": [true, false],
+      "config.process[0].model.qtype_te": ["convrot8", "qfloat8"],
+      "config.process[0].model.low_vram": [true, false],
+      // prompt embeds depend on the control images, so the TE stays loaded
+      "config.process[0].train.unload_text_encoder": [false, 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],
+      "config.process[0].sample.guidance_scale": [3, 4],
+      "config.process[0].sample.sample_steps": [30, 25],
+      "config.process[0].train.timestep_type": ["shift", "sigmoid"],
+      // edit: references as clean t=0 latents + in the prompt; kv_cache:
+      // references attend only to themselves, so sampling computes their K/V
+      // once (must be trained that way)
+      "config.process[0].model.model_kwargs": [
+        {
+          edit: true,
+          kv_cache: true,
+          match_target_res: true,
+          uncond_skip_layers: [22],
+        },
+        {},
+      ],
+      "config.process[0].network.network_kwargs.ignore_if_contains": [
+        ["layers.22."],
+        [],
+      ],
+    },
+    disableSections: ["network.conv", "train.unload_text_encoder"],
+    additionalSections: [
+      "datasets.multi_control_paths",
+      "sample.multi_ctrl_imgs",
+      "model.low_vram",
+      "model.layer_offloading",
+      "model.qie.match_target_res",
+      "model.model_kwargs.kv_cache",
+    ],
+  },
   {
     name: "moss_image:base_edit",
     label: "MOSS-Image Base [Edit Training]",

testing/test_model_loading.py

@@ -237,6 +237,20 @@ MODEL_TESTS = {
         },
         "sample": {**IMG, "num_inference_steps": 20, "guidance_scale": 3.0},
     },
+    "moss_image:edit": {
+        # edit LoRA path: vision-tower prompt + t=0 reference latents, isolated
+        # references with their K/V cached across sampling steps
+        "model": {
+            "name_or_path": "OpenMOSS-Team/MOSS-Image-SFT",
+            "quantize": True,
+            "quantize_te": True,
+            "qtype": "convrot8",
+            "qtype_te": "convrot8",
+            "model_kwargs": {"edit": True, "kv_cache": True, "match_target_res": True},
+        },
+        "sample": {**IMG, "num_inference_steps": 20, "guidance_scale": 3.0},
+        "needs_control_image": True,
+    },
     "moss_image:base_edit": {
         # edit LoRA path: vision-tower prompt + t=0 reference latents, isolated
         # references with their K/V cached across sampling steps

toolkit/models/registry.py

@@ -58,6 +58,7 @@ ARCH_REGISTRY: Dict[str, dict] = {
     "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}},
+    "moss_image:edit": {"modality": "image", "model": {"name_or_path": "OpenMOSS-Team/MOSS-Image-SFT", "quantize": True, "qtype": "convrot8", "quantize_te": True, "qtype_te": "convrot8", "model_kwargs": {"edit": True, "kv_cache": True, "match_target_res": True, "uncond_skip_layers": [22]}}, "sample": {"width": 1024, "height": 1024, "num_inference_steps": 30, "guidance_scale": 3.0}, "needs_control_image": True},
     "moss_image:base_edit": {"modality": "image", "model": {"name_or_path": "OpenMOSS-Team/MOSS-Image-Base", "quantize": True, "qtype": "convrot8", "quantize_te": True, "qtype_te": "convrot8", "model_kwargs": {"edit": True, "kv_cache": True, "match_target_res": True, "uncond_skip_layers": [22]}}, "sample": {"width": 1024, "height": 1024, "num_inference_steps": 30, "guidance_scale": 3.0}, "needs_control_image": True},
     "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}},