From b7cd8a13aff93e68b0f8de791dc04136e97d1291 Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Fri, 16 Oct 2026 12:19:06 -0600 Subject: [PATCH] Add support to train MOSS Image Base as an edit model --- .../diffusion_models/moss_image/moss_image.py | 251 ++++++++++++++++-- .../moss_image/src/pipeline.py | 20 +- .../moss_image/src/transformer.py | 225 +++++++++++++--- extensions_built_in/diffusion_models/ui.tsx | 53 ++++ testing/test_model_loading.py | 14 + toolkit/models/registry.py | 1 + toolkit/models/v2/text_encoders/qwen35.py | 42 ++- 7 files changed, 548 insertions(+), 58 deletions(-) diff --git a/extensions_built_in/diffusion_models/moss_image/moss_image.py b/extensions_built_in/diffusion_models/moss_image/moss_image.py index dd1d27c..8a5d230 100644 --- a/extensions_built_in/diffusion_models/moss_image/moss_image.py +++ b/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: ``, 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): diff --git a/extensions_built_in/diffusion_models/moss_image/src/pipeline.py b/extensions_built_in/diffusion_models/moss_image/src/pipeline.py index 8044344..292aa92 100644 --- a/extensions_built_in/diffusion_models/moss_image/src/pipeline.py +++ b/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: diff --git a/extensions_built_in/diffusion_models/moss_image/src/transformer.py b/extensions_built_in/diffusion_models/moss_image/src/transformer.py index 26c033b..a84050e 100644 --- a/extensions_built_in/diffusion_models/moss_image/src/transformer.py +++ b/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] diff --git a/extensions_built_in/diffusion_models/ui.tsx b/extensions_built_in/diffusion_models/ui.tsx index bf6613a..209d66b 100644 --- a/extensions_built_in/diffusion_models/ui.tsx +++ b/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", diff --git a/testing/test_model_loading.py b/testing/test_model_loading.py index 8663e52..deb815c 100644 --- a/testing/test_model_loading.py +++ b/testing/test_model_loading.py @@ -237,6 +237,20 @@ MODEL_TESTS = { }, "sample": {**IMG, "num_inference_steps": 20, "guidance_scale": 3.0}, }, + "moss_image:base_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-Base", + "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: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 diff --git a/toolkit/models/registry.py b/toolkit/models/registry.py index 67c9623..b858878 100644 --- a/toolkit/models/registry.py +++ b/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: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}}, "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}}, diff --git a/toolkit/models/v2/text_encoders/qwen35.py b/toolkit/models/v2/text_encoders/qwen35.py index cf0b5d0..925c745 100644 --- a/toolkit/models/v2/text_encoders/qwen35.py +++ b/toolkit/models/v2/text_encoders/qwen35.py @@ -1,6 +1,7 @@ -from transformers import Qwen3_5TextModel +from transformers import Qwen3_5Model, Qwen3_5TextModel from .._mixin import OstrisTransformersMixin +from .qwen3_vl import patch_qwen_vl_patch_embed QWEN35_COMFY_REPO = "Comfy-Org/Qwen3.5" @@ -54,3 +55,42 @@ class Qwen35TextOnlyEncoder(Qwen3_5TextModel, OstrisTransformersMixin): @classmethod def convert_state_dict_on_save(cls, state_dict): return {f"model.language_model.{k}": v for k, v in state_dict.items()} + + +class Qwen35VLEncoder(Qwen3_5Model, OstrisTransformersMixin): + """Qwen3.5 language tower + vision tower (moss_image edit: reference images + ride in the prompt). Same comfy files as Qwen35TextOnlyEncoder; only the + lm_head and mtp head are dropped. Text-only inputs give the same states + as the text-only class.""" + + aitk_subfolder = "text_encoder" + aitk_tokenizer_subfolder = "tokenizer" + aitk_cast_quantized_load = True + + aitk_comfy_repo = QWEN35_COMFY_REPO + aitk_comfy_weight_names = Qwen35TextOnlyEncoder.aitk_comfy_weight_names + + _DROPPED = ("lm_head.", "mtp.") + + @classmethod + def get_transformer_block_names(cls): + return ["language_model.layers", "visual.blocks"] + + @classmethod + def convert_state_dict_on_load(cls, state_dict): + out = {} + for key, value in state_dict.items(): + if key.startswith(cls._DROPPED): + continue + if key.startswith("model."): + key = key[len("model.") :] + out[key] = value + return out + + @classmethod + def convert_state_dict_on_save(cls, state_dict): + return {f"model.{k}": v for k, v in state_dict.items()} + + def patch_vision_patch_embed(self) -> int: + """kernel==stride Conv3d patch_embed has no fast bf16 path; swap for a GEMM.""" + return patch_qwen_vl_patch_embed(self)