Add support to train MOSS Image Base as an edit model

extensions_built_in/diffusion_models/moss_image/moss_image.py

@@ -18,13 +18,17 @@ stock Qwen3.5-9B from a comfy file.
 Flow convention matches ai-toolkit (sigma 1 = noise, target = noise - clean).
 """
 
+import math
 import os
 from typing import TYPE_CHECKING, List, Optional
 
 import torch
+import torch.nn.functional as F
+from PIL import Image
 from safetensors import safe_open
 from safetensors.torch import load_file
-from transformers import Qwen2Tokenizer
+from torchvision.transforms.functional import to_tensor
+from transformers import AutoImageProcessor, Qwen2Tokenizer
 
 from toolkit.accelerator import unwrap_model
 from toolkit.advanced_prompt_embeds import AdvancedPromptEmbeds
@@ -32,7 +36,10 @@ 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.text_encoders.qwen35 import (
+    Qwen35TextOnlyEncoder,
+    Qwen35VLEncoder,
+)
 from toolkit.models.v2.vae.flux2_kl import AutoEncoder, convert_diffusers_state_dict
 from toolkit.samplers.custom_flowmatch_sampler import (
     CustomFlowMatchEulerDiscreteScheduler,
@@ -98,6 +105,33 @@ def _hub_tokenizer_folder(repo_id: str) -> str:
     return os.path.join(root, "tokenizer")
 
 
+def _hub_file_folder(repo_id: str, filename: str) -> str:
+    """Folder holding one hub file, cache first (loaded from the folder so no
+    other file of the repo is resolved)."""
+    from huggingface_hub import hf_hub_download
+
+    try:
+        path = hf_hub_download(repo_id, filename, token=HF_TOKEN, local_files_only=True)
+    except Exception:
+        path = hf_hub_download(repo_id, filename, token=HF_TOKEN)
+    return os.path.dirname(path)
+
+
+def _drop_repeats(images):
+    """ctrl_img_1 defaults to ctrl_img, so a single reference can arrive twice."""
+    kept = []
+    for image in images:
+        previous = kept[-1] if kept else None
+        if (
+            previous is not None
+            and previous.shape == image.shape
+            and torch.equal(previous, image)
+        ):
+            continue
+        kept.append(image)
+    return kept
+
+
 class MossImageModel(BaseModel):
     arch = "moss_image"
 
@@ -128,6 +162,18 @@ class MossImageModel(BaseModel):
             kw.get("uncond_skip_layers", DEFAULT_UNCOND_SKIP_LAYERS) or ()
         )
 
+        # model_kwargs.edit: krea2-style reference conditioning for edit LoRAs
+        # (the vendor model is t2i only). References ride in the prompt through
+        # the Qwen3.5 vision tower and as clean t=0 latents after the text.
+        # kv_cache: train with references attending only to themselves so
+        # their K/V are computed once per sample at inference.
+        self.is_edit = bool(kw.get("edit", False))
+        self.kv_cache = bool(kw.get("kv_cache", False))
+        self.encode_control_in_text_embeddings = self.is_edit
+        self.has_multiple_control_images = self.is_edit
+        self.use_raw_control_images = self.is_edit
+        self.image_processor = None
+
     @staticmethod
     def get_train_scheduler():
         return CustomFlowMatchEulerDiscreteScheduler(
@@ -194,7 +240,9 @@ class MossImageModel(BaseModel):
         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(
+        # edit keeps the vision tower for the reference images
+        te_cls = Qwen35VLEncoder if self.is_edit else Qwen35TextOnlyEncoder
+        text_encoder = te_cls.load_model(
             te_path,
             dtype=te_kwargs["dtype"],
             config_path=config_path,
@@ -203,6 +251,12 @@ class MossImageModel(BaseModel):
             quantize_on_load=False,
             use_comfy_weights=te_kwargs["use_comfy_weights"],
         )
+        if self.is_edit:
+            text_encoder.patch_vision_patch_embed()
+            # MOSS ships no image processor; Qwen's (same template + vocab)
+            self.image_processor = AutoImageProcessor.from_pretrained(
+                _hub_file_folder(QWEN35_9B_REPO, "preprocessor_config.json")
+            )
         text_encoder.eval()
         text_encoder.requires_grad_(False)
         return text_encoder
@@ -294,6 +348,27 @@ class MossImageModel(BaseModel):
         gen_config.width = int(gen_config.width // sc * sc)
         gen_config.height = int(gen_config.height // sc * sc)
 
+        # reference latents; the prompt embeds already carry the references
+        # through the vision tower (ctrl_img_1 mirrors ctrl_img when unset)
+        ref_latents = None
+        if self.is_edit:
+            ctrl_paths = []
+            for path in (
+                gen_config.ctrl_img,
+                gen_config.ctrl_img_1,
+                gen_config.ctrl_img_2,
+                gen_config.ctrl_img_3,
+            ):
+                if path is not None and path not in ctrl_paths:
+                    ctrl_paths.append(path)
+            if ctrl_paths:
+                ref_latents = [
+                    self._encode_ref_latents(
+                        [to_tensor(Image.open(p).convert("RGB")) for p in ctrl_paths],
+                        target_pixels=gen_config.width * gen_config.height,
+                    )
+                ]
+
         result = pipeline(
             conditional_embeds=conditional_embeds,
             unconditional_embeds=unconditional_embeds,
@@ -305,9 +380,86 @@ class MossImageModel(BaseModel):
             generator=generator,
             shift=self.sample_shift,
             uncond_skip_layers=self.uncond_skip_layers,
+            ref_latents=ref_latents,
+            isolate_refs=self.kv_cache,
         )
         return result[0]
 
+    # ------------------------------------------------------------------
+    # Reference images (edit)
+    # ------------------------------------------------------------------
+    def _ref_target_pixels(self, target_pixels: Optional[int]) -> int:
+        """Area each reference is fitted to: control_image_max_pixels (1 MP
+        cap, only scales down), or the target's area with match_target_res."""
+        kw = self.model_config.model_kwargs
+        if kw.get("match_target_res", False) and target_pixels:
+            return int(target_pixels)
+        return int(kw.get("control_image_max_pixels", 1024 * 1024))
+
+    def _encode_ref_latents(
+        self, control_tensors, target_pixels: Optional[int] = None
+    ) -> List[torch.Tensor]:
+        """[0, 1] (C, H, W) / (1, C, H, W) references -> (128, h, w) latents,
+        aspect kept, snapped to the 16 px grid."""
+        sc = self.get_bucket_divisibility()
+        budget = self._ref_target_pixels(target_pixels)
+        match = self.model_config.model_kwargs.get("match_target_res", False)
+        latents = []
+        for img in control_tensors:
+            if img.dim() == 3:
+                img = img.unsqueeze(0)
+            img = img.to(self.device_torch, dtype=torch.float32)
+            h, w = img.shape[2], img.shape[3]
+            if match or h * w > budget:
+                ratio = h / w
+                new_h = math.sqrt(budget * ratio)
+                new_w = new_h / ratio
+            else:
+                new_h, new_w = float(h), float(w)
+            new_h = max(sc, int(round(new_h / sc)) * sc)
+            new_w = max(sc, int(round(new_w / sc)) * sc)
+            if (new_h, new_w) != (h, w):
+                img = F.interpolate(img, size=(new_h, new_w), mode="bilinear")
+            latent = self.encode_images(
+                img * 2 - 1, device=self.device_torch, dtype=self.torch_dtype
+            )
+            latents.append(latent[0])
+        return latents
+
+    def _batch_ref_latents(
+        self, batch: "DataLoaderBatchDTO", batch_size: int, target_pixels: int
+    ) -> Optional[List[List[torch.Tensor]]]:
+        control_list = batch.control_tensor_list
+        if control_list is None and batch.control_tensor is not None:
+            control_list = [batch.control_tensor[b : b + 1] for b in range(batch_size)]
+        if control_list is None:
+            return None
+        if len(control_list) != batch_size:
+            raise ValueError("control tensor list length does not match batch size")
+        return [
+            self._encode_ref_latents(controls, target_pixels=target_pixels)
+            for controls in control_list
+        ]
+
+    def _prep_vlm_images(self, images: List[torch.Tensor]) -> List[torch.Tensor]:
+        """References for the vision tower: downscaled (never up) to
+        vlm_max_pixels (384^2 = 144 tokens); detail comes from the latents."""
+        target = int(self.model_config.model_kwargs.get("vlm_max_pixels", 384 * 384))
+        out = []
+        for img in images:
+            if img.dim() == 4:
+                img = img[0]
+            img = img.to(self.device_torch, dtype=torch.float32)
+            h, w = img.shape[1], img.shape[2]
+            scale = min(1.0, math.sqrt(target / (h * w)))
+            nh, nw = max(round(h * scale), 32), max(round(w * scale), 32)
+            if (nh, nw) != (h, w):
+                img = F.interpolate(
+                    img[None], size=(nh, nw), mode="bicubic", antialias=True
+                )[0].clamp(0, 1)
+            out.append(img)
+        return out
+
     # ------------------------------------------------------------------
     # Training hooks
     # ------------------------------------------------------------------
@@ -333,36 +485,95 @@ class MossImageModel(BaseModel):
             t = timestep.to(device, dtype=torch.float32).reshape(-1)
             if t.shape[0] != batch_size:
                 t = t.expand(batch_size)
+            ref_latents = None
+            if self.is_edit and batch is not None:
+                _, _, lh, lw = latent_model_input.shape
+                ref_latents = self._batch_ref_latents(
+                    batch, batch_size, target_pixels=lh * lw * 16 * 16
+                )
 
         return self.transformer(
-            latent_model_input.to(device, dtype), t, embeds, mask
+            latent_model_input.to(device, dtype),
+            t,
+            embeds,
+            mask,
+            ref_latents=ref_latents,
+            isolate_refs=self.kv_cache,
         )
 
+    def _chat(self, content: str) -> str:
+        return self.tokenizer.apply_chat_template(
+            [{"role": "user", "content": content}],
+            tokenize=False,
+            add_generation_prompt=True,
+            enable_thinking=True,
+        )
+
+    def _encode_with_images(self, text: str, images: List[torch.Tensor]) -> torch.Tensor:
+        """Prompt with references ahead of it (``Picture N: <vision>``, the
+        Qwen-Image-Edit-Plus layout) -> (L, 4096) final states."""
+        pixels = self.image_processor(
+            images=self._prep_vlm_images(images), return_tensors="pt", do_rescale=False
+        )
+        grid = pixels["image_grid_thw"]
+        counts = (grid.prod(-1) // self.image_processor.merge_size**2).tolist()
+        content = "".join(
+            f"Picture {i + 1}: <|vision_start|>{'<|image_pad|>' * n}<|vision_end|>"
+            for i, n in enumerate(counts)
+        )
+        # no truncation: the image-pad runs must stay whole
+        input_ids = self.tokenizer(self._chat(content + text), return_tensors="pt")[
+            "input_ids"
+        ].to(self.device_torch)
+        mm_token_type_ids = (input_ids == self.text_encoder.config.image_token_id).int()
+        if int(mm_token_type_ids.sum()) != sum(counts):
+            raise ValueError("image placeholder count does not match the image grids")
+        out = self.text_encoder(
+            input_ids=input_ids,
+            pixel_values=pixels["pixel_values"].to(self.device_torch, self.torch_dtype),
+            image_grid_thw=grid.to(self.device_torch),
+            mm_token_type_ids=mm_token_type_ids,
+            use_cache=False,
+        )
+        return out.last_hidden_state[0]
+
     @torch.no_grad()
-    def get_prompt_embeds(self, prompt) -> AdvancedPromptEmbeds:
+    def get_prompt_embeds(self, prompt, control_images=None) -> AdvancedPromptEmbeds:
         if isinstance(prompt, str):
             prompt = [prompt]
         if self.text_encoder.device == torch.device("cpu"):
             self.text_encoder.to(self.device_torch)
 
+        # control images -> one list per prompt: a (B, C, H, W) batch, a list of
+        # per-sample lists, or a flat list for a single prompt
+        if control_images is not None and self.is_edit:
+            if isinstance(control_images, torch.Tensor):
+                control_images = [[control_images[i]] for i in range(control_images.shape[0])]
+            elif len(control_images) > 0 and not isinstance(control_images[0], list):
+                control_images = [control_images]
+            if len(control_images) == 1 and len(prompt) > 1:
+                control_images = control_images * len(prompt)
+            if len(control_images) != len(prompt):
+                raise ValueError("number of prompts must match number of control image sets")
+            control_images = [_drop_repeats(c) if c else None for c in control_images]
+        else:
+            control_images = [None] * len(prompt)
+
         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))
+        for text, images in zip(prompt, control_images):
+            if images:
+                states = self._encode_with_images(text, images)
+            else:
+                input_ids = self.tokenizer(
+                    self._chat(text),
+                    return_tensors="pt",
+                    truncation=True,
+                    max_length=self.max_text_length,
+                )["input_ids"].to(self.device_torch)
+                states = self.text_encoder(input_ids=input_ids, use_cache=False).last_hidden_state[0]
+            embeds.append(states.to(self.torch_dtype))
         return AdvancedPromptEmbeds(text_embeds=embeds)
 
     def get_loss_target(self, *args, **kwargs):

extensions_built_in/diffusion_models/moss_image/src/pipeline.py

@@ -95,6 +95,8 @@ class MossImagePipeline:
         generator: Optional[torch.Generator] = None,
         shift: float = DEFAULT_SHIFT,
         uncond_skip_layers: Sequence[int] = DEFAULT_UNCOND_SKIP_LAYERS,
+        ref_latents: Optional[List[List[torch.Tensor]]] = None,
+        isolate_refs: bool = False,
         **kwargs,
     ):
         model = self.model
@@ -115,13 +117,27 @@ class MossImagePipeline:
             latents = randn_tensor(shape, generator=generator, device=device, dtype=torch.float32)
         latents = latents.to(device, torch.float32)
 
+        # edit: references ride in both branches. Isolated references get
+        # their K/V computed on the first step and reused after; one cache per
+        # branch (the reference RoPE follows the prompt length, and the
+        # negative branch skips blocks)
+        has_refs = ref_latents is not None and any(len(r) > 0 for r in ref_latents)
+        ref_kwargs = {}
+        cond_cache = uncond_cache = None
+        if has_refs:
+            ref_kwargs = {"ref_latents": ref_latents, "isolate_refs": isolate_refs}
+            if isolate_refs:
+                cond_cache, uncond_cache = {}, {}
+
         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()
+            pred = transformer(
+                latents, timestep, cond, cond_mask, ref_kv_cache=cond_cache, **ref_kwargs
+            ).float()
             if do_cfg:
                 pred_uncond = transformer(
                     latents,
@@ -129,6 +145,8 @@ class MossImagePipeline:
                     uncond,
                     uncond_mask,
                     skip_layers=uncond_skip_layers,
+                    ref_kv_cache=uncond_cache,
+                    **ref_kwargs,
                 ).float()
                 pred = pred_uncond + guidance_scale * (pred - pred_uncond)
             if getattr(model, "sample_step_hook", None) is not None:

extensions_built_in/diffusion_models/moss_image/src/transformer.py

@@ -10,10 +10,15 @@ 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.
+
+Edit LoRAs (not part of the vendor model): clean reference
+latents ride after the text tokens with t=0 modulation, each on its own index
+of the RoPE axis the vendor leaves at 0, optionally isolated so their per-block
+K/V can be cached across sampling steps.
 """
 
 from dataclasses import dataclass
-from typing import Iterable, Optional
+from typing import Iterable, List, Optional
 
 import torch
 import torch.nn.functional as F
@@ -59,11 +64,24 @@ def _rotate(x, rope):
     return x * cos + torch.cat((-b, a), -1) * sin
 
 
-def _attention(q, k, v, rope, heads, mask):
+def _attention(q, k, v, rope, heads, mask, ref_span=None, kv_capture=None, kv_cache=None):
     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)
+    if kv_capture is not None:
+        # post-RoPE reference K/V, reused as extra keys on later steps. Copies:
+        # slices would pin the full-sequence K and the whole fused qkv output
+        # (~4x the cache) for the rest of sampling
+        kv_capture.append(
+            tuple(
+                t[:, :, ref_span[0] : ref_span[1]].clone(memory_format=torch.contiguous_format)
+                for t in (k, v)
+            )
+        )
+    if kv_cache is not None:
+        k = torch.cat((k, kv_cache[0].to(k.dtype)), 2)
+        v = torch.cat((v, kv_cache[1].to(v.dtype)), 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)
 
@@ -112,17 +130,46 @@ class _Block(nn.Module):
             }
         )
 
-    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]
+    def forward(self, x, rope, modulation, mask, ref_span=None, kv_capture=None, kv_cache=None):
+        # modulation (B, 6, D) [shift_a, scale_a, gate_a, shift_f, scale_f,
+        # gate_f], or (mod, ref_mod, split): tokens [split:] are clean
+        # references modulated at t=0
+        if isinstance(modulation, tuple):
+            mod, ref_mod, split = modulation
+            m, r = mod.unbind(1), ref_mod.unbind(1)
+
+            def affine(h, scale, shift):
+                return torch.cat(
+                    (
+                        h[:, :split] * (1 + m[scale][:, None]) + m[shift][:, None],
+                        h[:, split:] * (1 + r[scale][:, None]) + r[shift][:, None],
+                    ),
+                    1,
+                )
+
+            def gate(h, g):
+                return torch.cat(
+                    (h[:, :split] * m[g][:, None], h[:, split:] * r[g][:, None]), 1
+                )
+
+        else:
+            m = modulation.unbind(1)
+
+            def affine(h, scale, shift):
+                return h * (1 + m[scale][:, None]) + m[shift][:, None]
+
+            def gate(h, g):
+                return h * m[g][:, None]
+
+        y = affine(self.norm1(x), 1, 0)
         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]
+        attn = _attention(q, k, v, rope, self.heads, mask, ref_span, kv_capture, kv_cache)
+        x = x + gate(self.proj["o"](attn), 2)
+        y = affine(self.norm2(x), 4, 3)
+        hidden, ffn_gate = self.proj["ffn_up_gate"](y).chunk(2, -1)
+        return x + gate(self.proj["ffn_down"](_situ(hidden, ffn_gate)), 5)
 
 
 class MossImageTransformer(nn.Module, OstrisModelMixin):
@@ -219,30 +266,83 @@ class MossImageTransformer(nn.Module, OstrisModelMixin):
             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)."""
+    def _rope_from_coords(self, coords):
+        """(B, L, 4) fp32 (r, f, h, w) coordinates -> (B, L, head_dim) rope."""
+        device = coords.device
         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]
+        phase = coords.unsqueeze(-1) * bands
+        return torch.cat((phase.sin(), phase.cos()), -2).flatten(-2)
+
+    @staticmethod
+    def _image_coords(index, frame, height, width, device):
+        """(H*W, 4) coordinates of one latent grid: (index, frame, centered h,
+        centered w). The vendor's image has index 0; references use 1, 2, ..."""
         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)
+        coords = torch.zeros(height * width, 4, device=device)
+        coords[:, 0] = index
+        coords[:, 1] = frame
+        coords[:, 2] = (h - (height - 1) / 2).flatten()
+        coords[:, 3] = (w - (width - 1) / 2).flatten()
+        return coords
+
+    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)."""
+        batch = lengths.shape[0]
+        image = self._image_coords(0, 0, height, width, device)
+        image = image.unsqueeze(0).repeat(batch, 1, 1)
         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)
+        return self._rope_from_coords(torch.cat((image, text), 1))
+
+    def _pack_refs(self, ref_latents, lengths, device, dtype):
+        """Per-sample lists of clean (C, h, w) reference latents -> right-padded
+        tokens (B, R, C), rope (B, R, head_dim) and valid mask (B, R).
+        Reference i sits on RoPE axis r = i + 1 (always 0 for the vendor's
+        image and text) with the image's frame coordinate and its own centered
+        grid."""
+        tokens, coords = [], []
+        for refs, length in zip(ref_latents, lengths.tolist()):
+            sample_tokens, sample_coords = [], []
+            for i, latent in enumerate(refs):
+                _, h, w = latent.shape
+                sample_tokens.append(latent.flatten(1).transpose(0, 1))
+                sample_coords.append(self._image_coords(i + 1, length, h, w, device))
+            if sample_tokens:
+                tokens.append(torch.cat(sample_tokens).to(device, dtype))
+                coords.append(torch.cat(sample_coords))
+            else:
+                tokens.append(torch.zeros(0, self.params.latent_channels, device=device, dtype=dtype))
+                coords.append(torch.zeros(0, 4, device=device))
+        num_ref = max(t.shape[0] for t in tokens)
+        valid = torch.stack(
+            [torch.arange(num_ref, device=device) < t.shape[0] for t in tokens]
+        )
+        tokens = torch.stack([F.pad(t, (0, 0, 0, num_ref - t.shape[0])) for t in tokens])
+        coords = torch.stack([F.pad(c, (0, 0, 0, num_ref - c.shape[0])) for c in coords])
+        return tokens, self._rope_from_coords(coords), valid
+
+    def _time_modulation(self, times, dtype):
+        half = self.params.freq_dim // 2
+        freq = 10000 ** (
+            -torch.arange(half, device=times.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(times.shape[0], 6, -1)
+        return mod, time
 
     def forward(
         self,
@@ -251,8 +351,23 @@ class MossImageTransformer(nn.Module, OstrisModelMixin):
         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,
+        ref_latents: Optional[List[List[torch.Tensor]]] = None,  # per-sample (C, h, w)
+        isolate_refs: bool = False,
+        ref_kv_cache: Optional[dict] = None,
     ):
+        """``ref_latents``: clean reference latents appended after the text at
+        t=0 (edit LoRAs; the base model never saw them). ``isolate_refs``:
+        reference queries attend only to reference keys, so their K/V do not
+        depend on the step and can be cached. ``ref_kv_cache`` (inference,
+        needs ``isolate_refs``): an empty dict is filled with every block's
+        reference K/V on this call; once filled, later calls skip the
+        reference tokens and attend to the cached K/V instead. One cache per
+        prompt / skip-layer combination (the reference RoPE follows the text
+        length)."""
         skip = frozenset(skip_layers or ())
+        if ref_kv_cache is not None and not isolate_refs:
+            raise ValueError("ref_kv_cache needs isolate_refs: cached K/V are only step-invariant when references attend to themselves")
+        reuse_refs = ref_kv_cache is not None and ref_kv_cache.get("kv") is not None
         batch, _, height, width = hidden_states.shape
         num_image = height * width
         device = hidden_states.device
@@ -279,36 +394,74 @@ class MossImageTransformer(nn.Module, OstrisModelMixin):
 
         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)
+        mod, time = self._time_modulation(times, dtype)
         head = self.video_head_modulation(time.to(dtype)).view(batch, 2, -1)
 
         rope = self._rope(lengths, height, width, text_len, device).to(dtype)
+        num_live = num_image + text_len
 
-        text_mask = joint_mask = None
+        num_ref, ref_valid, modulation = 0, None, mod
+        if reuse_refs:
+            ref_valid = ref_kv_cache["mask"]
+        elif ref_latents is not None and any(len(r) > 0 for r in ref_latents):
+            ref_tokens, ref_rope, ref_valid = self._pack_refs(
+                ref_latents, lengths, device, dtype
+            )
+            num_ref = ref_tokens.shape[1]
+            ref_mod, _ = self._time_modulation(torch.zeros_like(times), dtype)
+            modulation = (mod, ref_mod, num_live)
+            rope = torch.cat((rope, ref_rope.to(dtype)), 1)
+
+        text_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, :]
+        live_valid = torch.cat(
+            (torch.ones(batch, num_image, dtype=torch.bool, device=device), valid), 1
+        )
+        joint_mask = None
+        if num_ref and isolate_refs:
+            key_valid = torch.cat((live_valid, ref_valid), 1)
+            is_ref = torch.arange(num_live + num_ref, device=device) >= num_live
+            allowed = key_valid[:, None, :] & (~is_ref[:, None] | is_ref[None, :])[None]
+            # every query keeps at least itself (padded rows would be all-masked)
+            allowed = allowed | torch.eye(num_live + num_ref, dtype=torch.bool, device=device)
+            joint_mask = allowed[:, None]
+        elif num_ref or reuse_refs:
+            # live queries see every real reference key (cached ones included)
+            key_valid = torch.cat((live_valid, ref_valid), 1)
+            if not bool(key_valid.all()):
+                joint_mask = key_valid[:, None, None, :]
+        elif padded:
+            joint_mask = live_valid[:, 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:]
+        text_rope = rope[:, num_image:num_live]
         for block in self.text_embedder.blocks:
             text = self._ckpt(block, text, text_rope, text_mask)
 
-        x = torch.cat((video, text), 1)
+        parts = [video, text]
+        if num_ref:
+            parts.append(self.video_embedder(ref_tokens))
+        x = torch.cat(parts, 1)
+
+        capture = ref_kv_cache is not None and not reuse_refs and num_ref > 0
+        captured = {}
         for index, layer in enumerate(self.layers):
-            if index not in skip:
-                x = self._ckpt(layer, x, rope, mod, joint_mask)
+            if index in skip:
+                continue
+            if capture:
+                slot = []
+                x = layer(x, rope, modulation, joint_mask, (num_live, num_live + num_ref), slot)
+                captured[index] = slot[0]
+            elif reuse_refs:
+                x = layer(x, rope, modulation, joint_mask, kv_cache=ref_kv_cache["kv"][index])
+            else:
+                x = self._ckpt(layer, x, rope, modulation, joint_mask)
+        if capture:
+            ref_kv_cache["kv"] = captured
+            ref_kv_cache["mask"] = ref_valid
 
         x = self.final_norm_video(x[:, :num_image])
         x = x * (1 + head[:, 1, None]) + head[:, 0, None]

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:base_edit",
+    label: "MOSS-Image Base [Edit Training]",
+    generateNameOverride: "MOSS-Image Base Edit",
+    group: "experimental",
+    defaults: {
+      // default updates when [selected, unselected] in the UI
+      "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 + 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: "flux2_klein_4b",
     label: "FLUX.2-klein-base-4B",