AI/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2026-10-24 02:49:04 +03:00
Code Activity

Add attention bias conditioning and image strength for MOSS-Image edit (#17320)

Browse Source
This commit is contained in:
Jukka Seppänen
2026-10-23 20:52:38 +03:00
committed by GitHub
parent 8baf68d53d
commit 9fe7ef764d
3 changed files with 42 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files

30
comfy/ldm/moss_image/model.py
View File

@@ -59,7 +59,7 @@ def rope_split_half(x, pe):
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, ref_kv=None, kv_out=None):
def attention(q, k, v, q_norm, k_norm, pe, heads, comfy_attention, transformer_options, ref_kv=None, kv_out=None, mask=None):
B, N, _ = q.shape
q = q.view(B, N, heads, -1)
k = k.view(B, N, heads, -1)
@@ -79,7 +79,7 @@ def attention(q, k, v, q_norm, k_norm, pe, heads, comfy_attention, transformer_o
# 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)
return optimized_attention(q, k, v, heads, mask=mask, transformer_options=transformer_options, preferred_attention=comfy_attention)
class TextBlock(nn.Module):
@@ -97,9 +97,9 @@ class TextBlock(nn.Module):
self.ffn_up = operations.Linear(dim, ffn_dim, dtype=dtype, device=device)
self.ffn_down = operations.Linear(ffn_dim, dim, dtype=dtype, device=device)
def forward(self, x, pe, transformer_options={}):
def forward(self, x, pe, transformer_options={}, mask=None):
q, k, v = self.qkv(self.norm1(x)).chunk(3, dim=-1)
x = x + self.proj(attention(q, k, v, self.q_norm, self.k_norm, pe, self.num_heads, self.comfy_attention, transformer_options))
x = x + self.proj(attention(q, k, v, self.q_norm, self.k_norm, pe, self.num_heads, self.comfy_attention, transformer_options, mask=mask))
y = self.norm2(x)
return x + self.ffn_down(situ(self.ffn_up(y), self.ffn_gate(y)))
@@ -110,12 +110,12 @@ class TextEmbedder(nn.Module):
self.proj = operations.Linear(text_dim, dim, dtype=dtype, device=device)
self.blocks = nn.ModuleList([TextBlock(dim, ffn_dim, num_heads, dtype=dtype, device=device, operations=operations) for _ in range(2)])
def forward(self, x, pe, transformer_options={}):
def forward(self, x, pe, transformer_options={}, mask=None):
x = self.proj(x)
prefetch_queue = comfy.model_prefetch.make_prefetch_queue(list(self.blocks), x.device, transformer_options)
for block in self.blocks:
comfy.model_prefetch.prefetch_queue_pop(prefetch_queue, x.device, block, x.dtype)
x = block(x, pe, transformer_options)
x = block(x, pe, transformer_options, mask)
comfy.model_prefetch.prefetch_queue_pop(prefetch_queue, x.device, None)
return x
@@ -136,14 +136,14 @@ 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={}, ref=None, ref_kv=None, kv_out=None):
def forward(self, x, mod, pe, transformer_options={}, ref=None, ref_kv=None, kv_out=None, mask=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, ref_kv, kv_out)
out = attention(q, k, v, self.q_norm, self.k_norm, pe, self.num_heads, self.comfy_attention, transformer_options, ref_kv, kv_out, mask)
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
@@ -212,7 +212,7 @@ 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={}, ref_latents=[], ref_latents_method=None, ref_kv=None, **kwargs):
def _forward(self, x, timestep, context, transformer_options={}, ref_latents=[], ref_latents_method=None, ref_kv=None, context_attention_bias=None, **kwargs):
B, _, H, W = x.shape
dtype, device = x.dtype, x.device
transformer_options = transformer_options.copy()
@@ -237,7 +237,13 @@ class MOSSImage(nn.Module):
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:H * W + text_len], transformer_options)
mask = text_mask = None
if context_attention_bias is not None:
# additive attention logit bias on the text tokens (prompt or image weights), image and reference keys stay unbiased
text_mask = context_attention_bias[:, None, None]
mask = torch.nn.functional.pad(context_attention_bias, (H * W, sum(r.shape[-2] * r.shape[-1] for r in refs) + (0 if ref_kv is None else ref_kv.shape[3])))[:, None, None]
txt = self.text_embedder(context, pe[:, H * W:H * W + text_len], transformer_options, text_mask)
# 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] + [self.video_embedder(r.flatten(2).transpose(1, 2)) for r in refs], dim=1)
del txt, refs
@@ -252,10 +258,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"], ref, None if ref_kv is None else ref_kv[:, i])}
return {"img": block(args["img"], args["mod"], args["pe"], args["transformer_options"], ref, None if ref_kv is None else ref_kv[:, i], mask=mask)}
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, ref, None if ref_kv is None else ref_kv[:, i])
x = block(x, mod, pe, transformer_options, ref, None if ref_kv is None else ref_kv[:, i], mask=mask)
comfy.model_prefetch.prefetch_queue_pop(prefetch_queue, device, None, malloc_scope="block")
comfy.model_prefetch.malloc_graph_end()

3
comfy/model_base.py
View File

@@ -2840,6 +2840,9 @@ 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)
context_attention_bias = kwargs.get("context_attention_bias", None)
if context_attention_bias is not None:
out['context_attention_bias'] = comfy.conds.CONDRegular(context_attention_bias)
ref_latents = kwargs.get("reference_latents", None)
if ref_latents:
# control adapters widen video_embedder, the references carry no control

26
comfy_extras/nodes_moss_image.py
View File

@@ -5,6 +5,7 @@ import re
import torch
from typing_extensions import override
import comfy.text_encoders.qwen_vl
import comfy.utils
import node_helpers
from comfy_api.latest import ComfyExtension, io
@@ -105,30 +106,45 @@ class TextEncodeMOSSImageEdit(io.ComfyNode):
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.",
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. Without a VAE, the images only go to the text encoder.",
inputs=[
io.Clip.Input("clip"),
io.String.Input("prompt", multiline=True, dynamic_prompts=True),
io.Vae.Input("vae"),
io.Vae.Input("vae", optional=True, tooltip="Encodes the images as reference latents for the edit LoRA."),
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)),
io.Float.Input("image_strength", default=1.0, min=0.0, max=100.0, step=0.05, advanced=True, tooltip="Scales how much the model attends to the images the text encoder sees (ai-toolkit formats, or no VAE): 0 ignores them, 1 is unchanged, above 1 emphasizes them."),
],
outputs=[io.Conditioning.Output()],
is_experimental=True,
)
@classmethod
def execute(cls, clip, prompt, vae, edit_format, images) -> io.NodeOutput:
def execute(cls, clip, prompt, edit_format, images, vae=None, image_strength=1.0) -> 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:
if (vision or vae is None) and image_strength > 0:
# 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))
tokens = clip.tokenize(prompt, images=vl)
conditioning = clip.encode_from_tokens_scheduled(tokens)
if len(vl) > 0 and image_strength != 1.0:
# log(strength) on each image's vision tokens scales their share of attention, Qwen3.5 merges 2x2 patches of 16 px into one token
weight = math.log(image_strength)
bias = []
for t in next(iter(tokens.values()))[0]:
if isinstance(t[0], dict):
_, grid = comfy.text_encoders.qwen_vl.process_qwen2vl_images(t[0]["data"], patch_size=16)
bias += [weight] * (int(grid[0, 1] * grid[0, 2]) // 4)
else:
bias.append(0.0)
conditioning = node_helpers.conditioning_set_values(conditioning, {"context_attention_bias": torch.tensor([bias])})
if vae is None:
return io.NodeOutput(conditioning)
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}))