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

@@ -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()
 

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):

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):

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,
                 ),
             ],

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,
         ]