547 lines
24 KiB
Python
547 lines
24 KiB
Python
"""MS-Diffusion native core.
|
|
|
|
Pure torch + ComfyUI core APIs. No diffusers, no transformers, no einops,
|
|
no omegaconf, no PIL. Everything here is a line-by-line port of the upstream
|
|
MS-Diffusion inference path (D:/Downloads/MS-Diffusion-main):
|
|
|
|
msdiffusion/models/projection.py -> FourierEmbedder / Resampler
|
|
msdiffusion/models/attention_processor.py-> MaskedIPAttnProcessor2_0
|
|
msdiffusion/models/model.py -> MSAdapter.generate (conditioning assembly)
|
|
|
|
The SDXL UNet itself, text encoding, CLIP-Vision, KSampler and VAE are all
|
|
left to ComfyUI native nodes. This file only supplies the *patch* that turns
|
|
native attn2 into the MS-Diffusion dual-branch cross-attention.
|
|
|
|
Key semantic facts (verified against real weights on 2026-11-06):
|
|
*_ms_adapter.bin* has exactly 70 (to_k_ip, to_v_ip) pairs, indexed
|
|
1,3,...,139 in diffusers attn_processors order. That order is
|
|
down_blocks -> up_blocks -> mid_block (module *registration* order),
|
|
whereas ComfyUI *executes* input_blocks -> middle_block -> output_blocks.
|
|
The two must be permuted, see build_block_weights().
|
|
* diffusers and ComfyUI use the SAME head partition for SDXL:
|
|
inner_dim 640 -> 10 heads, 1280 -> 20 heads, head_dim 64.
|
|
"""
|
|
|
|
import math
|
|
import os
|
|
import uuid
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
NUM_TOKENS = 16 # image tokens per subject
|
|
NUM_DUMMY_TOKENS = 4 # learnable "background" image tokens
|
|
TEXT_TOKENS = 77 # upstream hard-codes this split point
|
|
VISION_DIM = 1664 # CLIP-ViT-bigG/14 penultimate hidden size
|
|
CROSS_DIM = 2048 # SDXL cross_attention_dim
|
|
PHRASE_DIM = 768 # CLIP-L pooled output
|
|
RESAMPLER_DIM = 1280
|
|
FOURIER_FREQS = 8
|
|
HEAD_DIM = 64 # SDXL attention head dim (both frameworks)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# projection.py port (einops removed; runtime never touches masked_mean/Rearrange)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
class FourierEmbedder(nn.Module):
|
|
def __init__(self, num_freqs=64, temperature=100):
|
|
super().__init__()
|
|
self.num_freqs = num_freqs
|
|
self.temperature = temperature
|
|
freq_bands = temperature ** (torch.arange(num_freqs) / num_freqs)
|
|
self.register_buffer("freq_bands", freq_bands[None, None], persistent=False)
|
|
|
|
def __call__(self, x):
|
|
x = self.freq_bands * x.unsqueeze(-1)
|
|
return torch.stack((x.sin(), x.cos()), dim=-1).permute(0, 2, 3, 1).reshape(x.shape[0], -1)
|
|
|
|
|
|
def reshape_tensor(x, heads):
|
|
bs, length, width = x.shape
|
|
x = x.view(bs, length, heads, -1).transpose(1, 2)
|
|
return x.reshape(bs, heads, length, -1)
|
|
|
|
|
|
class PerceiverAttention(nn.Module):
|
|
def __init__(self, dim, dim_head=64, heads=8):
|
|
super().__init__()
|
|
self.scale = dim_head ** -0.5
|
|
self.dim_head = dim_head
|
|
self.heads = heads
|
|
inner_dim = dim_head * heads
|
|
self.norm1 = nn.LayerNorm(dim)
|
|
self.norm2 = nn.LayerNorm(dim)
|
|
self.to_q = nn.Linear(dim, inner_dim, bias=False)
|
|
self.to_kv = nn.Linear(dim, inner_dim * 2, bias=False)
|
|
self.to_out = nn.Linear(inner_dim, dim, bias=False)
|
|
|
|
def forward(self, x, latents):
|
|
x = self.norm1(x)
|
|
latents = self.norm2(latents)
|
|
b, l, _ = latents.shape
|
|
q = self.to_q(latents)
|
|
kv_input = torch.cat((x, latents), dim=-2)
|
|
k, v = self.to_kv(kv_input).chunk(2, dim=-1)
|
|
q = reshape_tensor(q, self.heads)
|
|
k = reshape_tensor(k, self.heads)
|
|
v = reshape_tensor(v, self.heads)
|
|
# upstream stability trick: scale BOTH operands, float32 softmax
|
|
scale = 1 / math.sqrt(math.sqrt(self.dim_head))
|
|
weight = (q * scale) @ (k * scale).transpose(-2, -1)
|
|
weight = torch.softmax(weight.float(), dim=-1).type(weight.dtype)
|
|
out = weight @ v
|
|
out = out.permute(0, 2, 1, 3).reshape(b, l, -1)
|
|
return self.to_out(out)
|
|
|
|
|
|
def feed_forward(dim, mult=4):
|
|
inner_dim = int(dim * mult)
|
|
return nn.Sequential(
|
|
nn.LayerNorm(dim),
|
|
nn.Linear(dim, inner_dim, bias=False),
|
|
nn.GELU(),
|
|
nn.Linear(inner_dim, dim, bias=False),
|
|
)
|
|
|
|
|
|
class Resampler(nn.Module):
|
|
"""Grounding-mode Resampler, key-for-key identical to upstream projection.py."""
|
|
|
|
def __init__(self, dim=RESAMPLER_DIM, depth=4, dim_head=64, heads=20,
|
|
num_queries=NUM_TOKENS, embedding_dim=VISION_DIM,
|
|
output_dim=CROSS_DIM, ff_mult=4, max_seq_len=257,
|
|
phrase_embeddings_dim=PHRASE_DIM, fourier_freqs=FOURIER_FREQS):
|
|
super().__init__()
|
|
self.num_queries = num_queries
|
|
self.grounding_token_num = num_queries
|
|
self.dim = dim
|
|
self.latent_init_mode = "grounding"
|
|
self.grounding_latents = nn.Parameter(torch.randn(1, num_queries, dim) / dim ** 0.5)
|
|
self.fourier_embedder = FourierEmbedder(num_freqs=fourier_freqs)
|
|
grounding_embedding_dim = phrase_embeddings_dim + fourier_freqs * 2 * 4
|
|
self.latent_proj = nn.Sequential(
|
|
nn.Linear(grounding_embedding_dim, grounding_embedding_dim * 2),
|
|
nn.GELU(),
|
|
nn.Linear(grounding_embedding_dim * 2, dim * num_queries),
|
|
)
|
|
self.latent_norm = nn.LayerNorm(dim)
|
|
self.proj_in = nn.Linear(embedding_dim, dim)
|
|
self.attention_norm = nn.LayerNorm(dim)
|
|
self.proj_out = nn.Linear(dim, output_dim)
|
|
self.norm_out = nn.LayerNorm(output_dim)
|
|
self.layers = nn.ModuleList([])
|
|
for _ in range(depth):
|
|
self.layers.append(nn.ModuleList([
|
|
PerceiverAttention(dim=dim, dim_head=dim_head, heads=heads),
|
|
feed_forward(dim=dim, mult=ff_mult),
|
|
]))
|
|
|
|
def forward(self, x, grounding_kwargs=None, shortcut=False, scale=1.0):
|
|
boxes = grounding_kwargs["boxes"]
|
|
phrase_embeds = grounding_kwargs["phrase_embeds"]
|
|
fourier_embeds = self.fourier_embedder(boxes)
|
|
grounding_embeds = torch.cat((phrase_embeds, fourier_embeds), dim=-1)
|
|
|
|
drop_grounding_tokens = grounding_kwargs["drop_grounding_tokens"]
|
|
num_ref = x.shape[0] // len(drop_grounding_tokens)
|
|
drop_grounding_tokens = [item for item in drop_grounding_tokens for _ in range(num_ref)]
|
|
|
|
latents = self.latent_proj(grounding_embeds)
|
|
latents = latents.view(-1, self.grounding_token_num, self.dim)
|
|
latents = self.latent_norm(latents)
|
|
|
|
drop_num = len([item for item in drop_grounding_tokens if item == 1])
|
|
if drop_num > 0:
|
|
latents_ = []
|
|
learnable_latents = self.grounding_latents.repeat(drop_num, 1, 1)
|
|
cur_idx = 0
|
|
for latent, flag in zip(latents, drop_grounding_tokens):
|
|
if flag == 1:
|
|
latent = learnable_latents[cur_idx]
|
|
cur_idx += 1
|
|
latents_.append(latent)
|
|
latents = torch.stack(latents_)
|
|
|
|
x = self.proj_in(x)
|
|
init_latents = latents
|
|
for attn, ff in self.layers:
|
|
latents = attn(x, latents) + latents
|
|
latents = ff(latents) + latents
|
|
latents = self.attention_norm(latents)
|
|
latents = self.proj_out(latents)
|
|
if shortcut:
|
|
latents = init_latents + latents * scale
|
|
return self.norm_out(latents)
|
|
|
|
|
|
class IPProjection(nn.Module):
|
|
def __init__(self, hidden_size):
|
|
super().__init__()
|
|
self.to_k_ip = nn.Linear(CROSS_DIM, hidden_size, bias=False)
|
|
self.to_v_ip = nn.Linear(CROSS_DIM, hidden_size, bias=False)
|
|
|
|
|
|
class MSAdapter(nn.Module):
|
|
"""Self-contained adapter. Keys are 1:1 with ms_adapter.bin after index remap."""
|
|
|
|
def __init__(self, hidden_sizes):
|
|
super().__init__()
|
|
self.hidden_sizes = list(hidden_sizes)
|
|
self.image_proj = Resampler()
|
|
self.adapters = nn.ModuleList([IPProjection(h) for h in self.hidden_sizes])
|
|
self.dummy_image_tokens = nn.Parameter(torch.randn(1, NUM_DUMMY_TOKENS, CROSS_DIM))
|
|
self._placed = None
|
|
|
|
def load(self, path):
|
|
if os.path.isdir(path):
|
|
path = os.path.join(path, "ms_adapter.bin")
|
|
sd = torch.load(path, map_location="cpu", weights_only=True)
|
|
missing_top = {"image_proj", "ms_adapter", "dummy_image_tokens"} - set(sd)
|
|
if missing_top:
|
|
raise RuntimeError("MS-Diffusion checkpoint missing groups: %s" % sorted(missing_top))
|
|
|
|
self.image_proj.load_state_dict(sd["image_proj"], strict=True)
|
|
proj_sd = {}
|
|
for i in range(len(self.hidden_sizes)):
|
|
src = 2 * i + 1 # 1,3,...,139
|
|
for name in ("to_k_ip", "to_v_ip"):
|
|
proj_sd["%d.%s.weight" % (i, name)] = sd["ms_adapter"]["%d.%s.weight" % (src, name)]
|
|
self.adapters.load_state_dict(proj_sd, strict=True)
|
|
assert sd["dummy_image_tokens"].shape == self.dummy_image_tokens.shape
|
|
with torch.no_grad():
|
|
self.dummy_image_tokens.copy_(sd["dummy_image_tokens"])
|
|
for i, h in enumerate(self.hidden_sizes):
|
|
w = self.adapters[i].to_k_ip.weight
|
|
if w.shape != (h, CROSS_DIM):
|
|
raise RuntimeError("adapter %d shape %s != (%d,%d)" % (i, tuple(w.shape), h, CROSS_DIM))
|
|
return self
|
|
|
|
@torch.no_grad()
|
|
def ensure_on(self, device, dtype):
|
|
tag = (str(device), str(dtype))
|
|
if self._placed == tag:
|
|
return
|
|
self.to(device=device, dtype=dtype)
|
|
self._placed = tag
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Conditioning assembly (port of MSAdapter.generate, without diffusers pipeline)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
class MSEmbeds:
|
|
"""Everything the attn2 patch needs, independent of CONDITIONING."""
|
|
|
|
def __init__(self, image_prompt_embeds, boxes, phrase_idxes, num_ref,
|
|
scale, mask_threshold, start_step, subject_scales=None):
|
|
self.image_prompt_embeds = image_prompt_embeds # [1, 4+16*num_ref, 2048] fp32
|
|
self.boxes = boxes # [1, num_ref, 4] fp32
|
|
self.phrase_idxes = phrase_idxes # [1, num_ref, 2] int
|
|
self.num_ref = num_ref
|
|
self.scale = float(scale)
|
|
self.mask_threshold = mask_threshold # None disables pseudo mask
|
|
self.start_step = int(start_step)
|
|
self.subject_scales = subject_scales
|
|
|
|
|
|
def minmax_normalize(batch_maps):
|
|
min_val = batch_maps.min(dim=-1, keepdim=True)[0].min(dim=-2, keepdim=True)[0]
|
|
max_val = batch_maps.max(dim=-1, keepdim=True)[0].max(dim=-2, keepdim=True)[0]
|
|
return (batch_maps - min_val) / (max_val - min_val + 1e-5)
|
|
|
|
|
|
def build_box_masks(boxes, seq_len_q, dtype, device):
|
|
"""Port of MaskedIPAttnProcessor2_0.prepare_attention_mask_qk (boxes path)."""
|
|
if boxes is None:
|
|
return None, None, None
|
|
batch_size, num_ref, _ = boxes.shape
|
|
nrow = int(seq_len_q ** 0.5)
|
|
box_idxes_start = torch.floor(boxes[:, :, 0:2] * nrow)
|
|
box_idxes_end = torch.ceil(boxes[:, :, 2:4] * nrow)
|
|
box_idxes = torch.cat([box_idxes_start, box_idxes_end], dim=-1)
|
|
box_masks = []
|
|
dummy_attention_mask = torch.ones((batch_size, seq_len_q), dtype=dtype, device=device)
|
|
for box_idx in box_idxes.unbind(dim=1):
|
|
xs, ys, xe, ye = box_idx.unbind(dim=1)
|
|
x_indices = torch.arange(nrow).unsqueeze(0).expand(batch_size, -1).to(device)
|
|
y_indices = torch.arange(nrow).unsqueeze(0).expand(batch_size, -1).to(device)
|
|
x_mask = ((x_indices >= xs.unsqueeze(1)) & (x_indices < xe.unsqueeze(1))).to(dtype)
|
|
y_mask = ((y_indices >= ys.unsqueeze(1)) & (y_indices < ye.unsqueeze(1))).to(dtype)
|
|
box_mask = torch.bmm(y_mask.unsqueeze(2), x_mask.unsqueeze(1)).reshape(batch_size, -1)
|
|
box_masks.append(box_mask)
|
|
dummy_attention_mask = torch.clamp(dummy_attention_mask - box_mask, min=0)
|
|
|
|
post_dummy = dummy_attention_mask.to(torch.bool)
|
|
image_mask = torch.stack(box_masks, dim=-1)
|
|
image_mask = image_mask.repeat_interleave(NUM_TOKENS, dim=-1)
|
|
image_mask = (1 - image_mask.to(dtype)) * -10000.0
|
|
dummy_mask = dummy_attention_mask.unsqueeze(-1).repeat_interleave(NUM_DUMMY_TOKENS, dim=-1)
|
|
dummy_mask = (1 - dummy_mask) * -10000.0
|
|
image_mask = torch.cat([dummy_mask, image_mask], dim=-1)
|
|
return image_mask, post_dummy, box_masks
|
|
|
|
|
|
def build_pseudo_masks(maps, mask_threshold, seq_len_q, dtype, device, num_ref):
|
|
"""Port of get_psuedo_attention_mask(). `maps` is a list of [B,num_ref,h,w]."""
|
|
text_maps = minmax_normalize(torch.stack(maps).mean(dim=0))
|
|
bsz, nr, h, w = text_maps.shape
|
|
seq_len_q = h * w
|
|
text_maps = text_maps.view(bsz, nr, -1).transpose(1, 2)
|
|
pseudo = (text_maps > mask_threshold).to(dtype)
|
|
pseudo_dummy = torch.ones((bsz, seq_len_q), dtype=dtype, device=device)
|
|
for i in range(nr):
|
|
pseudo_dummy = torch.clamp(pseudo_dummy - pseudo[..., i], min=0)
|
|
post_dummy = pseudo_dummy.to(torch.bool)
|
|
pseudo = pseudo.repeat_interleave(NUM_TOKENS, dim=-1)
|
|
pseudo = (1 - pseudo) * -10000.0
|
|
pseudo_dummy = pseudo_dummy.unsqueeze(-1).repeat_interleave(NUM_DUMMY_TOKENS, dim=-1)
|
|
pseudo_dummy = (1 - pseudo_dummy) * -10000.0
|
|
return torch.cat([pseudo_dummy, pseudo], dim=-1), post_dummy
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# attn2 patch
|
|
# ---------------------------------------------------------------------------
|
|
|
|
# Module-level registry: ComfyUI's ModelPatcher.clone() deep-copies dict/list in
|
|
# model_options but passes arbitrary objects (functions, dicts of state) through
|
|
# unchanged. Closures capture only a string id, so cloned patchers share state.
|
|
_STATES = {}
|
|
|
|
|
|
def new_state(adapter, embeds):
|
|
sid = uuid.uuid4().hex
|
|
_STATES[sid] = {
|
|
"adapter": adapter,
|
|
"embeds": embeds,
|
|
"maps": {}, # per layer index -> list of text attention maps
|
|
"last_sigma": -1.0,
|
|
"calls": 0,
|
|
}
|
|
# keep the registry bounded across many node executions
|
|
if len(_STATES) > 64:
|
|
for k in list(_STATES)[:-32]:
|
|
_STATES.pop(k, None)
|
|
return sid
|
|
|
|
|
|
def _select_per_batch(t, batch, cond_or_uncond, zero_for_uncond=True):
|
|
"""Expand a [1,...] tensor to `batch` rows, zeroing uncond rows.
|
|
|
|
Upstream builds `cat([zeros_like(x), x])` for negative/positive. ComfyUI
|
|
instead reports the role of every batch chunk in transformer_options
|
|
["cond_or_uncond"] (0=cond, 1=uncond), so we build the same layout without
|
|
relying on a fixed chunk order.
|
|
"""
|
|
t = t.expand(batch, *t.shape[1:])
|
|
if not zero_for_uncond or cond_or_uncond is None:
|
|
return t
|
|
labels = list(cond_or_uncond)
|
|
if len(labels) == 0:
|
|
return t
|
|
if len(labels) != batch:
|
|
repeat = batch // len(labels)
|
|
if repeat * len(labels) != batch:
|
|
return t
|
|
labels = [x for x in labels for _ in range(repeat)]
|
|
keep = torch.tensor([1.0 if int(c) == 0 else 0.0 for c in labels],
|
|
device=t.device, dtype=t.dtype).view(batch, *([1] * (t.dim() - 1)))
|
|
return t * keep
|
|
|
|
|
|
def make_attn2_patch(state_id, layer_index, is_sentinel=False):
|
|
"""Return a ComfyUI `patches_replace["attn2"][(block, n, depth)]` callable.
|
|
|
|
Signature is (q, k, v, extra_options); q/k/v are already projected flat
|
|
tensors [B, L, inner_dim]. The enclosing BasicTransformerBlock applies
|
|
attn2.to_out and the residual, exactly like diffusers' processor.
|
|
"""
|
|
|
|
def attn2_patch(q, k, v, extra_options):
|
|
st = _STATES[state_id]
|
|
adapter = st["adapter"]
|
|
emb = st["embeds"]
|
|
|
|
# Detect a fresh sampling run by the (monotonically decreasing) sigma.
|
|
# All 70 patches see this; the first one resets, the rest no-op.
|
|
sigmas = extra_options.get("sigmas")
|
|
if sigmas is not None:
|
|
try:
|
|
sig = float(torch.as_tensor(sigmas).flatten()[0])
|
|
if sig > st["last_sigma"] + 1e-6:
|
|
st["maps"] = {}
|
|
st["calls"] = 0
|
|
st["last_sigma"] = sig
|
|
except Exception:
|
|
pass
|
|
if is_sentinel:
|
|
st["calls"] += 1
|
|
|
|
device, dtype = q.device, q.dtype
|
|
adapter.ensure_on(device, dtype)
|
|
bsz, seq_q, inner_dim = q.shape
|
|
heads = inner_dim // HEAD_DIM
|
|
head_dim = inner_dim // heads
|
|
scale_q = head_dim ** -0.5
|
|
use_pseudo = emb.mask_threshold is not None
|
|
|
|
def split_heads(t):
|
|
return t.view(t.shape[0], -1, heads, head_dim).transpose(1, 2)
|
|
|
|
qt = split_heads(q)
|
|
|
|
# ---- text branch -------------------------------------------------
|
|
kt, vt = split_heads(k), split_heads(v)
|
|
if not use_pseudo:
|
|
text_out = F.scaled_dot_product_attention(qt, kt, vt)
|
|
else:
|
|
scores = torch.matmul(qt * scale_q, kt.transpose(-1, -2))
|
|
text_probs = torch.softmax(scores.float(), dim=-1).to(dtype)
|
|
# upstream get_text_attention_maps() works on [B*heads, Lq, Lk]
|
|
flat = text_probs.reshape(bsz * heads, seq_q, -1).view(bsz, heads, seq_q, -1)
|
|
# Upstream builds cat([zeros, real]) for uncond/cond: uncond boxes
|
|
# and phrase indices are all zero. Reproduce that exactly.
|
|
boxes = _select_per_batch(emb.boxes.to(device=device, dtype=dtype), bsz,
|
|
extra_options.get("cond_or_uncond"))
|
|
pidx = _select_per_batch(emb.phrase_idxes.to(device=device, dtype=dtype), bsz,
|
|
extra_options.get("cond_or_uncond")).to(torch.int)
|
|
num_ref = boxes.shape[1]
|
|
h = w = int(seq_q ** 0.5)
|
|
maps = []
|
|
for i in range(bsz):
|
|
per_ref = []
|
|
for j in range(num_ref):
|
|
start_idx = int(pidx[i, j, 0].item())
|
|
end_idx = int(pidx[i, j, 1].item())
|
|
if start_idx == 0 and end_idx == 0:
|
|
per_ref.append(torch.zeros(seq_q, dtype=dtype, device=device))
|
|
else:
|
|
amap = flat[i, :, :, start_idx:end_idx]
|
|
per_ref.append(torch.mean(torch.mean(amap, dim=-1), dim=0))
|
|
maps.append(torch.stack(per_ref))
|
|
st["maps"].setdefault(layer_index, []).append(
|
|
torch.stack(maps).reshape(bsz, num_ref, h, w))
|
|
text_out = torch.matmul(text_probs, vt)
|
|
text_out = text_out.transpose(1, 2).reshape(bsz, seq_q, inner_dim).to(dtype)
|
|
|
|
# ---- image branch ------------------------------------------------
|
|
ip = emb.image_prompt_embeds.to(device=device, dtype=dtype)
|
|
ip = _select_per_batch(ip, bsz, extra_options.get("cond_or_uncond"))
|
|
ipk = split_heads(F.linear(ip, adapter.adapters[layer_index].to_k_ip.weight))
|
|
ipv = split_heads(F.linear(ip, adapter.adapters[layer_index].to_v_ip.weight))
|
|
|
|
boxes = _select_per_batch(emb.boxes.to(device=device, dtype=dtype), bsz,
|
|
extra_options.get("cond_or_uncond"))
|
|
rf_mask, dummy_mask, _ = build_box_masks(boxes, seq_q, dtype, device)
|
|
if use_pseudo:
|
|
maps_here = st["maps"].get(layer_index, [])
|
|
if len(maps_here) >= emb.start_step:
|
|
pm, pd = build_pseudo_masks(maps_here, emb.mask_threshold,
|
|
seq_q, dtype, device, boxes.shape[1])
|
|
rf_mask, dummy_mask = pm, pd
|
|
# [B, Lq, N] -> [B*heads, Lq, N]
|
|
if rf_mask is not None:
|
|
rf_mask = rf_mask.repeat_interleave(heads, dim=0).reshape(bsz * heads, seq_q, -1)
|
|
if dummy_mask is not None:
|
|
dummy_mask = dummy_mask.repeat_interleave(heads, dim=0)
|
|
|
|
qf = qt.reshape(bsz * heads, seq_q, head_dim)
|
|
kf = ipk.reshape(bsz * heads, -1, head_dim)
|
|
vf = ipv.reshape(bsz * heads, -1, head_dim)
|
|
scores = torch.bmm(qf * scale_q, kf.transpose(-1, -2))
|
|
if rf_mask is not None:
|
|
scores = scores + rf_mask
|
|
ip_probs = torch.softmax(scores.float(), dim=-1).to(dtype)
|
|
if dummy_mask is not None:
|
|
ip_probs = torch.where(dummy_mask.unsqueeze(-1), torch.zeros_like(ip_probs), ip_probs)
|
|
if emb.subject_scales is not None:
|
|
ss = torch.tensor(emb.subject_scales, dtype=dtype, device=device)
|
|
ss = ss.unsqueeze(0).unsqueeze(0).repeat_interleave(NUM_TOKENS, dim=-1)
|
|
dss = torch.ones((1, 1, 1), dtype=dtype, device=device).repeat_interleave(NUM_DUMMY_TOKENS, dim=-1)
|
|
ip_probs = ip_probs * torch.cat([dss, ss], dim=-1)
|
|
ip_out = torch.bmm(ip_probs, vf)
|
|
ip_out = ip_out.reshape(bsz, heads, seq_q, head_dim).transpose(1, 2).reshape(bsz, seq_q, inner_dim)
|
|
|
|
if emb.subject_scales is None:
|
|
return text_out + emb.scale * ip_out
|
|
return text_out + ip_out
|
|
|
|
return attn2_patch
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Weight -> ComfyUI block mapping
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def enumerate_attn2(diffusion_model):
|
|
"""Return attn2 descriptors in ComfyUI *execution* order.
|
|
|
|
Each item: (key3, inner_dim, section) where key3 is the exact
|
|
(block_name, block_number, transformer_index) key ComfyUI looks up in
|
|
patches_replace["attn2"].
|
|
"""
|
|
import re
|
|
out = []
|
|
for name, module in diffusion_model.named_modules():
|
|
if not name.endswith(".attn2"):
|
|
continue
|
|
m = re.match(r"^(input_blocks|middle_block|output_blocks)\.(\d+)\..*?transformer_blocks\.(\d+)\.attn2$", name)
|
|
if m is None:
|
|
continue
|
|
section = {"input_blocks": "input", "middle_block": "middle", "output_blocks": "output"}[m.group(1)]
|
|
block_name = section
|
|
block_number = 0 if section == "middle" else int(m.group(2))
|
|
depth = int(m.group(3))
|
|
inner_dim = module.to_q.out_features
|
|
out.append(((block_name, block_number, depth), inner_dim, section))
|
|
# named_modules() already yields registration order == execution order
|
|
return out
|
|
|
|
|
|
def build_block_weights(attn2_desc, hidden_sizes):
|
|
"""Permute the 70 checkpoint groups into ComfyUI execution order.
|
|
|
|
Checkpoint order (diffusers registration): input | output | middle
|
|
ComfyUI execution order: input | middle | output
|
|
"""
|
|
n_input = sum(1 for x in attn2_desc if x[2] == "input")
|
|
n_middle = sum(1 for x in attn2_desc if x[2] == "middle")
|
|
n_output = sum(1 for x in attn2_desc if x[2] == "output")
|
|
total = n_input + n_middle + n_output
|
|
if total != len(hidden_sizes):
|
|
raise RuntimeError(
|
|
"MS-Diffusion: UNet has %d attn2 modules (%d input / %d middle / %d output) "
|
|
"but the checkpoint has %d adapter groups. This adapter only supports "
|
|
"stock SDXL (70 cross-attention layers)." % (total, n_input, n_middle, n_output, len(hidden_sizes)))
|
|
reordered = (list(range(0, n_input))
|
|
+ list(range(n_input + n_output, n_input + n_output + n_middle))
|
|
+ list(range(n_input, n_input + n_output)))
|
|
for pos, ckpt_idx in enumerate(reordered):
|
|
want = attn2_desc[pos][1]
|
|
got = hidden_sizes[ckpt_idx]
|
|
if want != got:
|
|
raise RuntimeError(
|
|
"MS-Diffusion: adapter mapping mismatch at attn2 #%d: UNet inner_dim=%d "
|
|
"but checkpoint group %d has hidden_size=%d. Refusing to run with wrong weights."
|
|
% (pos, want, ckpt_idx, got))
|
|
return reordered
|
|
|
|
|
|
def adapter_hidden_sizes(checkpoint_path):
|
|
"""Read just the 70 to_k_ip shapes so the loader node can build the module."""
|
|
sd = torch.load(checkpoint_path, map_location="cpu", weights_only=True)
|
|
if not isinstance(sd, dict) or "ms_adapter" not in sd or "image_proj" not in sd:
|
|
raise RuntimeError(
|
|
"MS-Diffusion: %s is not an ms_adapter.bin checkpoint (needs "
|
|
"'image_proj', 'ms_adapter' and 'dummy_image_tokens' groups). "
|
|
"Pick ms_adapter.bin, not a generic IP-Adapter file." % checkpoint_path)
|
|
group = sd["ms_adapter"]
|
|
idx = sorted({int(k.split(".")[0]) for k in group})
|
|
try:
|
|
return [group["%d.to_k_ip.weight" % i].shape[0] for i in idx]
|
|
except KeyError as e:
|
|
raise RuntimeError("MS-Diffusion: unexpected ms_adapter key layout (%s). "
|
|
"Expected 70 to_k_ip/to_v_ip pairs." % e) from e
|