From 60aade0bb0cbc298547fceb136cd195cea67800b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jukka=20Sepp=C3=A4nen?= <40791699+kijai@users.noreply.github.com> Date: Thu, 22 Oct 2026 22:17:59 +0300 Subject: [PATCH] Support MOSS-Image edit LoRAs (#17296) * Support MOSS-image edit loras * Update model.py * Add reference methods to the reference node too --- comfy/ldm/moss_image/model.py | 110 +++++++++++++++++++++++++------ comfy/lora.py | 3 + comfy/model_base.py | 22 +++++++ comfy_extras/nodes_flux.py | 2 +- comfy_extras/nodes_moss_image.py | 44 +++++++++++++ 5 files changed, 159 insertions(+), 22 deletions(-) diff --git a/comfy/ldm/moss_image/model.py b/comfy/ldm/moss_image/model.py index 3bf5bbbae..b6b58bea6 100644 --- a/comfy/ldm/moss_image/model.py +++ b/comfy/ldm/moss_image/model.py @@ -7,6 +7,7 @@ import comfy.model_prefetch import comfy.ops import comfy.patcher_extension import comfy.quant_ops +import comfy.utils from comfy.ldm.flux.layers import timestep_embedding from comfy.ldm.modules.attention import AttentionTensorContainer, ComfyAttention, optimized_attention @@ -33,13 +34,32 @@ def gated_residual(x, y, gate): return x.addcmul_(y, gate) +def modulate(x, norm, mod, i, ref): + # mod[i] is the shift and mod[i + 1] the scale; reference tokens from ref[0] on take the t=0 modulation ref[1] + if ref is None: + return modulated_rms_norm(x, norm, mod[i + 1], mod[i]) + s, ref_mod = ref[0], ref[1] + return torch.cat((modulated_rms_norm(x[:, :s], norm, mod[i + 1], mod[i]), modulated_rms_norm(x[:, s:], norm, ref_mod[i + 1], ref_mod[i])), dim=1) + + +def residual(x, y, mod, i, ref): + if ref is None: + return gated_residual(x, y, mod[i]) + s, ref_mod = ref[0], ref[1] + if comfy.model_management.in_training: + return torch.cat((gated_residual(x[:, :s], y[:, :s], mod[i]), gated_residual(x[:, s:], y[:, s:], ref_mod[i])), dim=1) + gated_residual(x[:, :s], y[:, :s], mod[i]) + gated_residual(x[:, s:], y[:, s:], ref_mod[i]) + return x + + def rope_split_half(x, pe): # differentiable split-half RoPE for training, the kitchen ops have no autograd x_ = x.to(pe.dtype).unflatten(-1, (2, -1)).movedim(-2, -1).unsqueeze(-2) return (pe[..., 0] * x_[..., 0] + pe[..., 1] * x_[..., 1]).movedim(-1, -2).flatten(-2).type_as(x) -def attention(q, k, v, q_norm, k_norm, pe, heads, comfy_attention, transformer_options): +def attention(q, k, v, q_norm, k_norm, pe, heads, comfy_attention, transformer_options, ref_kv=None, kv_out=None): B, N, _ = q.shape q = q.view(B, N, heads, -1) k = k.view(B, N, heads, -1) @@ -51,7 +71,14 @@ def attention(q, k, v, q_norm, k_norm, pe, heads, comfy_attention, transformer_o q, k = comfy.quant_ops.ck.rms_rope_split_half(q, k, pe, q_scale, k_scale, q_norm.eps) comfy.ops.uncast_bias_weight(q_norm, q_scale, None, q_stream) comfy.ops.uncast_bias_weight(k_norm, k_scale, None, k_stream) - q, k, v = AttentionTensorContainer(q.flatten(2)), AttentionTensorContainer(k.flatten(2)), AttentionTensorContainer(v) + q, k = q.flatten(2), k.flatten(2) + if kv_out is not None: + kv_out[:, 0].copy_(k) + kv_out[:, 1].copy_(v) + if ref_kv is not None: + # the references' keys and values, computed once + k, v = torch.cat((k, ref_kv[:, 0]), dim=1), torch.cat((v, ref_kv[:, 1]), dim=1) + q, k, v = AttentionTensorContainer(q), AttentionTensorContainer(k), AttentionTensorContainer(v) return optimized_attention(q, k, v, heads, transformer_options=transformer_options, preferred_attention=comfy_attention) @@ -109,18 +136,18 @@ class Block(nn.Module): self.proj["ffn_up_gate"] = operations.Linear(dim, 2 * ffn_dim, bias=False, dtype=dtype, device=device) self.proj["ffn_down"] = operations.Linear(ffn_dim, dim, bias=False, dtype=dtype, device=device) - def forward(self, x, mod, pe, transformer_options={}): - shift_a, scale_a, gate_a, shift_f, scale_f, gate_f = mod - y = modulated_rms_norm(x, self.norm1, scale_a, shift_a) + def forward(self, x, mod, pe, transformer_options={}, ref=None, ref_kv=None, kv_out=None): + # ref: (first reference token, their t=0 modulation) when the references are in the sequence + y = modulate(x, self.norm1, mod, 0, ref) if self.fused_qkv: q, k, v = self.proj["qkv"](y).chunk(3, dim=-1) else: q, k, v = self.proj["q"](y), self.proj["k"](y), self.proj["v"](y) - out = attention(q, k, v, self.q_norm, self.k_norm, pe, self.num_heads, self.comfy_attention, transformer_options) - x = gated_residual(x, self.proj["o"](out), gate_a) - y = modulated_rms_norm(x, self.norm2, scale_f, shift_f) + out = attention(q, k, v, self.q_norm, self.k_norm, pe, self.num_heads, self.comfy_attention, transformer_options, ref_kv, kv_out) + x = residual(x, self.proj["o"](out), mod, 2, ref) + y = modulate(x, self.norm2, mod, 3, ref) # situ on the [hidden | gate] halves; int8 weights fold it into the activation quantization - return gated_residual(x, comfy.ops.linear_input_act(self.proj["ffn_down"], self.proj["ffn_up_gate"](y), "situ"), gate_f) + return residual(x, comfy.ops.linear_input_act(self.proj["ffn_down"], self.proj["ffn_up_gate"](y), "situ"), mod, 5, ref) class MOSSImage(nn.Module): @@ -139,19 +166,45 @@ class MOSSImage(nn.Module): self.final_linear_video = operations.Linear(dim, latent_channels, bias=False, dtype=dtype, device=device) self.video_head_modulation = operations.Linear(dim, 2 * dim, dtype=dtype, device=device) - def rope(self, h, w, text_len, device, transformer_options={}): - # (r, f, h, w): the image sits after the text on f, its rows and columns centered + def rope(self, h, w, text_len, device, transformer_options={}, ref_shapes=(), f_offset=0): + # (r, f, h, w): the image sits after the text on f, its rows and columns centered; reference i on r = i + 1, f = 0 rope_options = transformer_options.get("rope_options", {}) # NTK-aware axis scales, zero collapses the axis scales = (1.0, 1.0, rope_options.get("scale_y", 1.0), rope_options.get("scale_x", 1.0)) - ids = torch.zeros(1, h * w + text_len, 4, device=device, dtype=torch.float32) - ids[0, :h * w, 1] = text_len + rope_options.get("shift_t", 0.0) + n = h * w + text_len + ids = torch.zeros(1, n + sum(rh * rw for rh, rw in ref_shapes), 4, device=device, dtype=torch.float32) + ids[0, :h * w, 1] = text_len + f_offset + rope_options.get("shift_t", 0.0) ids[0, :h * w, 2] = ((torch.arange(h, device=device, dtype=torch.float32) - (h - 1) / 2) * float(scales[2] != 0) + rope_options.get("shift_y", 0.0)).repeat_interleave(w) ids[0, :h * w, 3] = ((torch.arange(w, device=device, dtype=torch.float32) - (w - 1) / 2) * float(scales[3] != 0) + rope_options.get("shift_x", 0.0)).repeat(h) - ids[0, h * w:, 1] = torch.arange(text_len, device=device, dtype=torch.float32) + ids[0, h * w:n, 1] = torch.arange(text_len, device=device, dtype=torch.float32) + f_offset + start = n + for i, (rh, rw) in enumerate(ref_shapes): + ref = ids[0, n:n + rh * rw] + ref[:, 0] = i + 1 + ref[:, 2] = (torch.arange(rh, device=device, dtype=torch.float32) - (rh - 1) / 2).repeat_interleave(rw) + ref[:, 3] = (torch.arange(rw, device=device, dtype=torch.float32) - (rw - 1) / 2).repeat(rh) + n += rh * rw pe = torch.cat([comfy.ldm.flux.math.rope(ids[..., i], d, 10000 * (s or 1.0) ** (-d / (d - 2))) for i, (d, s) in enumerate(zip(self.rope_axes_dim, scales))], dim=-3) + if len(ref_shapes) > 0: + # the references keep the trained frequencies, Scale ROPE stretches only the generated image + pe[:, start:] = torch.cat([comfy.ldm.flux.math.rope(ids[:, start:, i], d, 10000) for i, d in enumerate(self.rope_axes_dim)], dim=-3) return pe.unsqueeze(1).transpose(1, 2).contiguous() + def preprocess_refs(self, refs, transformer_options={}): + # references that attend only to each other at t = 0 depend on themselves alone: every block's keys and values for them + batch = max(ref.shape[0] for ref in refs) + x = torch.cat([self.video_embedder(comfy.utils.repeat_to_batch_size(ref, batch).flatten(2).transpose(1, 2)) for ref in refs], dim=1) + pe = self.rope(0, 0, 0, x.device, transformer_options, [ref.shape[-2:] for ref in refs]) + vec = self.time_embedding(timestep_embedding(torch.zeros(1, device=x.device), 256).to(x.dtype)) + mod = self.shared_block_modulation(vec).unflatten(-1, (6, -1)).unsqueeze(2).unbind(1) + kv = x.new_empty((x.shape[0], len(self.layers), 2) + x.shape[1:]) + prefetch_queue = comfy.model_prefetch.make_prefetch_queue(list(self.layers), x.device, transformer_options) + for i, block in enumerate(self.layers): + comfy.model_prefetch.prefetch_queue_pop(prefetch_queue, x.device, block, x.dtype) + x = block(x, mod, pe, transformer_options, kv_out=kv[:, i]) + comfy.model_prefetch.prefetch_queue_pop(prefetch_queue, x.device, None) + return kv + def forward(self, x, timestep, context, transformer_options={}, **kwargs): return comfy.patcher_extension.WrapperExecutor.new_class_executor( self._forward, @@ -159,20 +212,35 @@ class MOSSImage(nn.Module): comfy.patcher_extension.get_all_wrappers(comfy.patcher_extension.WrappersMP.DIFFUSION_MODEL, transformer_options) ).execute(x, timestep, context, transformer_options, **kwargs) - def _forward(self, x, timestep, context, transformer_options={}, **kwargs): + def _forward(self, x, timestep, context, transformer_options={}, ref_latents=[], ref_latents_method=None, ref_kv=None, **kwargs): B, _, H, W = x.shape dtype, device = x.dtype, x.device transformer_options = transformer_options.copy() - pe = self.rope(H, W, context.shape[1], device, transformer_options) + text_len = context.shape[1] + if len(ref_latents) > 0 and ref_latents_method not in ("index", "index_timestep_zero"): + # every step when extra_conds could not precompute them (training, hooks) + ref_kv, ref_latents = self.preprocess_refs(ref_latents, transformer_options), [] + if ref_kv is not None: + # one set of references serves the whole batch + ref_kv = ref_kv.expand(B, -1, -1, -1, -1) if ref_kv.shape[0] == 1 else comfy.utils.repeat_to_batch_size(ref_kv, B) + refs = [comfy.utils.repeat_to_batch_size(ref, B) for ref in ref_latents] + # ai-toolkit puts the references on the image's f, Moss-Image-Edit at the text start; shifting the rest keeps them at f = 0, where their keys were computed + f_offset = -text_len if (len(refs) > 0 or ref_kv is not None) and ref_latents_method != "index_timestep_zero_text_start" else 0 + pe = self.rope(H, W, text_len, device, transformer_options, [ref.shape[-2:] for ref in refs], f_offset) vec = self.time_embedding(timestep_embedding(timestep, 256).to(dtype)) mod = self.shared_block_modulation(vec).unflatten(-1, (6, -1)).unsqueeze(2).unbind(1) head_shift, head_scale = self.video_head_modulation(vec).unflatten(-1, (2, -1)).unsqueeze(2).unbind(1) + ref = None + if len(refs) > 0 and ref_latents_method != "index": + # ai-toolkit without kv_cache: the references ride along at t = 0 and attend to everything ("index": at the image's timestep) + ref_vec = self.time_embedding(timestep_embedding(torch.zeros_like(timestep), 256).to(dtype)) + ref = (H * W + text_len, self.shared_block_modulation(ref_vec).unflatten(-1, (6, -1)).unsqueeze(2).unbind(1)) - txt = self.text_embedder(context, pe[:, H * W:], transformer_options) + txt = self.text_embedder(context, pe[:, H * W:H * W + text_len], transformer_options) # allocated outside the memory compiler region, the blocks update it in place - x = torch.cat((self.video_embedder(x.flatten(2).transpose(1, 2)), txt), dim=1) - del txt + x = torch.cat([self.video_embedder(x.flatten(2).transpose(1, 2)), txt] + [self.video_embedder(r.flatten(2).transpose(1, 2)) for r in refs], dim=1) + del txt, refs blocks_replace = transformer_options.get("patches_replace", {}).get("dit", {}) transformer_options["total_blocks"] = len(self.layers) @@ -184,10 +252,10 @@ class MOSSImage(nn.Module): transformer_options["block_index"] = i if ("single_block", i) in blocks_replace: def block_wrap(args): - return {"img": block(args["img"], args["mod"], args["pe"], args["transformer_options"])} + return {"img": block(args["img"], args["mod"], args["pe"], args["transformer_options"], ref, None if ref_kv is None else ref_kv[:, i])} x = blocks_replace[("single_block", i)]({"img": x, "vec": vec, "mod": mod, "pe": pe, "transformer_options": transformer_options}, {"original_block": block_wrap})["img"] else: - x = block(x, mod, pe, transformer_options) + x = block(x, mod, pe, transformer_options, ref, None if ref_kv is None else ref_kv[:, i]) comfy.model_prefetch.prefetch_queue_pop(prefetch_queue, device, None, malloc_scope="block") comfy.model_prefetch.malloc_graph_end() diff --git a/comfy/lora.py b/comfy/lora.py index 384ca91b6..b25843814 100644 --- a/comfy/lora.py +++ b/comfy/lora.py @@ -358,11 +358,14 @@ def model_lora_keys_unet(model, key_map={}): if isinstance(model, comfy.model_base.MOSSImage): for k in sdk: + if k.startswith("diffusion_model.") and k.endswith(".weight"): + key_map[k[len("diffusion_model."):-len(".weight")]] = k # bare module paths if k.endswith(".proj.qkv.weight"): # fused q/k/v files; LoRAs address the separate projections third = sd[k].shape[0] // 3 for i, proj in enumerate(("q", "k", "v")): key_lora = "{}.{}".format(k[:-len(".qkv.weight")], proj) key_map[key_lora] = (k, (0, i * third, third)) + key_map[key_lora[len("diffusion_model."):]] = (k, (0, i * third, third)) key_map["lora_unet_{}".format(key_lora[len("diffusion_model."):].replace(".", "_"))] = (k, (0, i * third, third)) if isinstance(model, comfy.model_base.Krea2): diff --git a/comfy/model_base.py b/comfy/model_base.py index 8ee80152e..7a1e1247a 100644 --- a/comfy/model_base.py +++ b/comfy/model_base.py @@ -2800,6 +2800,7 @@ class Krea2(BaseModel): class MOSSImage(BaseModel): def __init__(self, model_config, model_type=ModelType.FLOW, device=None): super().__init__(model_config, model_type, device=device, unet_model=comfy.ldm.moss_image.model.MOSSImage) + self.memory_usage_factor_conds = ("ref_latents",) def extra_conds(self, **kwargs): out = super().extra_conds(**kwargs) @@ -2807,6 +2808,27 @@ class MOSSImage(BaseModel): # rope positions follow the true text length, so no repeat padding: cond and uncond of different lengths run separately if cross_attn is not None: out['c_crossattn'] = comfy.conds.CONDRegular(cross_attn) + ref_latents = kwargs.get("reference_latents", None) + if ref_latents: + # control adapters widen video_embedder, the references carry no control + channels = self.diffusion_model.video_embedder.weight.shape[1] + ref_latents = [torch.nn.functional.pad(self.process_latent_in(lat), (0, 0, 0, 0, 0, channels - lat.shape[1])) for lat in ref_latents] + ref_latents_method = kwargs.get("reference_latents_method", None) + if ref_latents_method is not None: + out['ref_latents_method'] = comfy.conds.CONDConstant(ref_latents_method) + if ref_latents_method not in ("index", "index_timestep_zero") and torch.is_inference_mode_enabled() and kwargs.get("hooks") is None: + # references that attend only to each other at t = 0: their keys and values are computed once, not every step; equal ones batch without copies + device, dtype = kwargs["device"], self.get_dtype_inference() + out['ref_kv'] = comfy.conds.CONDConstant(self.diffusion_model.preprocess_refs([lat.to(device=device, dtype=dtype) for lat in ref_latents])) + else: + out['ref_latents'] = comfy.conds.CONDList(ref_latents) + return out + + def extra_conds_shapes(self, **kwargs): + out = {} + ref_latents = kwargs.get("reference_latents", None) + if ref_latents is not None: + out['ref_latents'] = list([1, 128, sum(map(lambda a: math.prod(a.size()), ref_latents)) // 128]) return out def concat_cond(self, **kwargs): diff --git a/comfy_extras/nodes_flux.py b/comfy_extras/nodes_flux.py index 1871f385c..708dcd66d 100644 --- a/comfy_extras/nodes_flux.py +++ b/comfy_extras/nodes_flux.py @@ -162,7 +162,7 @@ class FluxKontextMultiReferenceLatentMethod(io.ComfyNode): io.Conditioning.Input("conditioning"), io.Combo.Input( "reference_latents_method", - options=["offset", "index", "uxo/uno", "index_timestep_zero"], + options=["offset", "index", "uxo/uno", "index_timestep_zero", "index_timestep_zero_isolated", "index_timestep_zero_text_start"], advanced=True, ), ], diff --git a/comfy_extras/nodes_moss_image.py b/comfy_extras/nodes_moss_image.py index 23732072f..84d8ba085 100644 --- a/comfy_extras/nodes_moss_image.py +++ b/comfy_extras/nodes_moss_image.py @@ -90,12 +90,56 @@ class MOSSImageControl(io.ComfyNode): node_helpers.conditioning_set_values_with_timestep_range(negative, values, start_percent, end_percent)) +# edit LoRA layouts: the reference method and whether the text encoder also sees the references +EDIT_FORMATS = { + "ai-toolkit": ("index_timestep_zero_isolated", True), + "ai-toolkit (kv_cache off)": ("index_timestep_zero", True), + "Moss-Image-Edit": ("index_timestep_zero_text_start", False), +} + + +class TextEncodeMOSSImageEdit(io.ComfyNode): + @classmethod + def define_schema(cls): + return io.Schema( + node_id="TextEncodeMOSSImageEdit", + display_name="Text Encode MOSS-Image Edit", + category="model/conditioning/moss_image", + description="Encodes an edit instruction with reference images for a MOSS-Image edit LoRA, use it for the negative prompt too. The references are used at the size they come in, scale them first.", + inputs=[ + io.Clip.Input("clip"), + io.String.Input("prompt", multiline=True, dynamic_prompts=True), + io.Vae.Input("vae"), + io.Combo.Input("edit_format", options=list(EDIT_FORMATS), tooltip="The layout the edit LoRA was trained with: ai-toolkit's MOSS-Image Edit Training (its kv_cache option on or off) or Moss-Image-Edit. With kv_cache off the references also attend to the image being generated, so they are recomputed every step."), + io.Autogrow.Input("images", template=io.Autogrow.TemplatePrefix(input=io.Image.Input("image"), prefix="image_", min=1)), + ], + outputs=[io.Conditioning.Output()], + is_experimental=True, + ) + + @classmethod + def execute(cls, clip, prompt, vae, edit_format, images) -> io.NodeOutput: + method, vision = EDIT_FORMATS[edit_format] + images = [images[name] for name in sorted(images) if images[name] is not None] + vl = [] + if vision: + # the text encoder sees each reference downscaled to at most 384x384 pixels, after a "Picture N:" marker + for image in images: + scale = min(1.0, math.sqrt(384 * 384 / (image.shape[1] * image.shape[2]))) + vl.append(comfy.utils.common_upscale(image.movedim(-1, 1), round(image.shape[2] * scale), round(image.shape[1] * scale), "area", "disabled").movedim(1, -1)) + prompt = "".join("Picture {}: <|vision_start|><|image_pad|><|vision_end|>".format(i + 1) for i in range(len(vl))) + prompt + conditioning = clip.encode_from_tokens_scheduled(clip.tokenize(prompt, images=vl)) + refs = [vae.encode(image[:, :, :, :3]) for image in images] + return io.NodeOutput(node_helpers.conditioning_set_values(conditioning, {"reference_latents": refs, "reference_latents_method": method})) + + class MOSSImageExtension(ComfyExtension): @override async def get_node_list(self) -> list[type[io.ComfyNode]]: return [ MOSSImageLayoutPrompt, MOSSImageControl, + TextEncodeMOSSImageEdit, ]