init

msdiffusion_native/__init__.py

@@ -0,0 +1,4 @@
+"""Native MS-Diffusion nodes: pure torch + ComfyUI core."""
+from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
+
+__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]

msdiffusion_native/core.py

@@ -0,0 +1,546 @@
+"""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

msdiffusion_native/nodes.py

@@ -0,0 +1,369 @@
+"""ComfyUI nodes for the native MS-Diffusion chain.
+
+Intended graph (everything except the three MS nodes is a stock ComfyUI node):
+
+  CheckpointLoaderSimple -> MODEL / CLIP / VAE
+  LoadImage xN -> native ImageBatch (chained N-1 times) -> CLIPVisionEncode
+                                                      -> CLIP_VISION_OUTPUT
+  CLIPTextEncode x2 -> positive / negative CONDITIONING
+  MSDiffusionLoader ------------------------------> MS_ADAPTER
+  MSDiffusionConditioning ------------------------> MS_EMBEDS (+ clean prompt)
+  MSDiffusionApply(MODEL, MS_ADAPTER, MS_EMBEDS) --> patched MODEL
+  EmptyLatentImage -> KSampler -> VAEDecode -> SaveImage
+
+No diffusers / transformers / omegaconf / einops / PIL imports anywhere.
+"""
+
+import os
+import re
+
+import torch
+import folder_paths
+
+from .core import (MSAdapter, MSEmbeds, adapter_hidden_sizes, build_block_weights,
+                   enumerate_attn2, make_attn2_patch, new_state, NUM_TOKENS,
+                   NUM_DUMMY_TOKENS)
+
+ADAPTER_FOLDER = "ms_adapter"
+_ADAPTER_EXTS = (".bin", ".safetensors", ".pt", ".pth", ".ckpt")
+
+
+def _register_folders():
+    """Register models/ipadapter and models/photomaker as one virtual folder.
+
+    The upstream plugin hard-coded models/photomaker; the user's actual file is
+    in models/ipadapter. Support both without touching ComfyUI source.
+    """
+    try:
+        if ADAPTER_FOLDER not in folder_paths.folder_names_and_paths:
+            folder_paths.folder_names_and_paths[ADAPTER_FOLDER] = ([], set())
+        for sub in ("ipadapter", "photomaker"):
+            p = os.path.join(folder_paths.models_dir, sub)
+            if os.path.isdir(p):
+                folder_paths.add_model_folder_path(ADAPTER_FOLDER, p)
+    except Exception:
+        pass
+
+
+_register_folders()
+
+
+def _adapter_files():
+    try:
+        files = [f for f in folder_paths.get_filename_list(ADAPTER_FOLDER)
+                 if f.lower().endswith(_ADAPTER_EXTS)]
+    except Exception:
+        files = []
+    if not files:
+        # fallback scan, keeps the node usable even if registration was blocked
+        for sub in ("ipadapter", "photomaker"):
+            d = os.path.join(folder_paths.models_dir, sub)
+            if os.path.isdir(d):
+                files += [sub + "/" + f for f in sorted(os.listdir(d))
+                          if f.lower().endswith(_ADAPTER_EXTS)]
+    files = sorted(set(files)) or ["none"]
+    # models/ipadapter often holds many unrelated IP-Adapter weights; put the
+    # real MS-Diffusion checkpoint first so the default is the correct file.
+    preferred = [f for f in files if os.path.basename(f).lower().startswith("ms_adapter")]
+    return preferred + [f for f in files if f not in preferred]
+
+
+def _adapter_path(name):
+    try:
+        return folder_paths.get_full_path_or_raise(ADAPTER_FOLDER, name)
+    except Exception:
+        return os.path.join(folder_paths.models_dir, name.replace("/", os.sep))
+
+
+# ---------------------------------------------------------------------------
+# phrase token position (port of msdiffusion/utils.py, no transformers import)
+# ---------------------------------------------------------------------------
+
+def _token_encoder(clip):
+    """Return f(word) -> len(HF tokenizer.encode(word)).
+
+    ComfyUI already owns a CLIP tokenizer instance; we just borrow it, so this
+    plugin adds no transformers import of its own.
+    """
+    try:
+        tk = clip.tokenizer.clip_l.tokenizer
+        return lambda w: len(tk.encode(w))
+    except Exception:
+        pass
+    # fallback: derive from ComfyUI's own token/weight pairs
+    def from_comfy(w):
+        pairs = clip.tokenize(w)["l"][0]
+        n = 0
+        for tid, _ in pairs:
+            if int(tid) == 49407:      # EOS == pad token
+                break
+            n += 1
+        return max(n, 1) + 1          # [BOS] + word tokens => n-1 + 2
+    return from_comfy
+
+
+def _is_equal_words(pr_words, ph_words):
+    if len(pr_words) != len(ph_words):
+        return False
+    for pr, ph in zip(pr_words, ph_words):
+        if "-" + ph not in pr and ph != re.sub(r'[.!?,:]$', '', pr):
+            return False
+    return True
+
+
+def _get_phrase_idx(enc_len, phrase, prompt, num=0):
+    phrase_words = phrase.split()
+    if len(phrase_words) == 0:
+        return [0, 0]
+    prompt_words = prompt.split()
+    start, end = 1, 0
+    for i in range(len(prompt_words)):
+        if _is_equal_words(prompt_words[i:i + len(phrase_words)], phrase_words):
+            if num != 0:
+                num -= 1
+                continue
+            end = start
+            prompt_words[i + len(phrase_words) - 1] = re.sub(
+                r'[.!?,:]$', '', prompt_words[i + len(phrase_words) - 1])
+            for j in range(i, i + len(phrase_words)):
+                end += enc_len(prompt_words[j]) - 2
+            break
+        else:
+            start += enc_len(prompt_words[i]) - 2
+    if end == 0:
+        return [0, 0]
+    return [start, end]
+
+
+def _parse_boxes(text, num_ref, enabled):
+    """Upstream: mask_threshold == 0 disables layout guidance -> all-zero boxes."""
+    zeros = [[0.0, 0.0, 0.0, 0.0] for _ in range(num_ref)]
+    if not enabled or not text.strip():
+        return zeros
+    groups = [g.strip() for g in text.strip().split(";") if g.strip()]
+    parsed = []
+    for g in groups:
+        vals = [float(x.strip()) for x in g.split(",")]
+        if len(vals) != 4:
+            raise ValueError("MS-Diffusion: each box must have 4 floats (x1,y1,x2,y2); got %r" % g)
+        parsed.append(vals)
+    if len(parsed) < num_ref:
+        # A zero box masks that subject's 16 image tokens completely (-10000 bias),
+        # silently dropping it. Upstream raises here instead of guessing; do the same.
+        raise ValueError(
+            "MS-Diffusion: %d subjects but only %d boxes. Provide one "
+            "'x1,y1,x2,y2' per subject separated by ';' (or set mask_threshold=0 "
+            "to disable layout guidance)." % (num_ref, len(parsed)))
+    return parsed[:num_ref]
+
+
+def _clip_l_pooled(clip, phrase):
+    """CLIP-L 768-d pooled output (the grounding phrase embedding).
+
+    SDXL's CONDITIONING pooled_output is clip_g (1280) and cannot be used here.
+    """
+    clip.load_model()
+    dev = clip.patcher.load_device
+    try:
+        clip.cond_stage_model.clip_l.set_clip_options({"execution_device": dev})
+    except Exception:
+        pass
+    tokens = clip.tokenize(phrase)["l"]
+    with torch.no_grad():
+        out = clip.cond_stage_model.clip_l.encode_token_weights(tokens)
+    pooled = out[1]
+    return pooled.to(torch.float32).cpu()
+
+
+# ---------------------------------------------------------------------------
+# Nodes
+# ---------------------------------------------------------------------------
+
+class MSDiffusionLoader:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {"required": {"ms_adapter": (_adapter_files(),)}}
+
+    RETURN_TYPES = ("MS_ADAPTER",)
+    RETURN_NAMES = ("ms_adapter",)
+    FUNCTION = "load"
+    CATEGORY = "MSdiffusion/Native"
+    DESCRIPTION = ("Load ms_adapter.bin from models/ipadapter or models/photomaker. "
+                   "Pure torch; no diffusers/transformers needed.")
+
+    def load(self, ms_adapter):
+        if ms_adapter == "none":
+            raise RuntimeError("MS-Diffusion: no ms_adapter.bin found in models/ipadapter or models/photomaker")
+        path = _adapter_path(ms_adapter)
+        sizes = adapter_hidden_sizes(path)
+        adapter = MSAdapter(sizes).load(path)
+        print("[MS-Diffusion] loaded %s (%d cross-attn groups)" % (path, len(sizes)))
+        return (adapter,)
+
+
+class MSDiffusionConditioning:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {"required": {
+            "ms_adapter": ("MS_ADAPTER",),
+            "clip": ("CLIP",),
+            "clip_vision_output": ("CLIP_VISION_OUTPUT",),
+            "prompt": ("STRING", {"multiline": True,
+                                  "default": "best quality, high quality, a [dog] on the beach"}),
+            "boxes": ("STRING", {"multiline": True,
+                                 "default": "0.25, 0.25, 0.75, 0.75",
+                                 "tooltip": "one 'x1,y1,x2,y2' per subject, separated by ';'. "
+                                            "Ignored when mask_threshold is 0."}),
+            "mask_threshold": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.05,
+                                         "tooltip": "0 disables layout guidance entirely (upstream behaviour). "
+                                                    ">0 enables box guidance + pseudo attention mask."}),
+            "drop_grounding_tokens": ("BOOLEAN", {"default": False}),
+        }}
+
+    RETURN_TYPES = ("MS_EMBEDS", "STRING")
+    RETURN_NAMES = ("ms_embeds", "prompt_clean")
+    FUNCTION = "encode"
+    CATEGORY = "MSdiffusion/Native"
+    DESCRIPTION = ("Build MS-Diffusion image prompt tokens from CLIP-Vision penultimate "
+                   "hidden states + CLIP-L phrase pooled embeddings + layout boxes. "
+                   "Connect prompt_clean to your positive CLIPTextEncode so the two match.")
+
+    def encode(self, ms_adapter, clip, clip_vision_output, prompt, boxes,
+               mask_threshold, drop_grounding_tokens):
+        import comfy.model_management as mm
+
+        phrases = [p.strip() for p in re.findall(r'\[(.*?)\]', prompt)]
+        if not phrases:
+            raise ValueError("MS-Diffusion: subject names must be enclosed in [ ] inside the prompt")
+        for p in phrases:
+            if " " in p:
+                raise ValueError("MS-Diffusion: subject %r must be a single word" % p)
+        phrases = list(dict.fromkeys(phrases))
+        clean_prompt = prompt.replace("]", " ").replace("[", " ")
+
+        hidden = clip_vision_output["penultimate_hidden_states"]
+        if hidden.dim() != 3:
+            raise ValueError("MS-Diffusion: CLIP-Vision output must be penultimate_hidden_states [B,257,1664]")
+        num_ref = hidden.shape[0]
+        if num_ref != len(phrases):
+            raise ValueError(
+                "MS-Diffusion: got %d subject images but %d [subject] names in prompt (%s)"
+                % (num_ref, len(phrases), phrases))
+        if hidden.shape[-1] != 1664:
+            raise ValueError("MS-Diffusion: expected CLIP-ViT-bigG hidden size 1664, got %d. "
+                             "Use clip_vision_g.safetensors." % hidden.shape[-1])
+
+        layout_on = mask_threshold > 0.0
+        box_vals = _parse_boxes(boxes, num_ref, layout_on)
+        boxes_t = torch.tensor([box_vals], dtype=torch.float32)          # [1, num_ref, 4]
+        batch_boxes = boxes_t.view(num_ref, 4)                           # Resampler grounding input
+
+        phrase_embeds = torch.cat([_clip_l_pooled(clip, p) for p in phrases], dim=0)
+        if phrase_embeds.shape != (num_ref, 768):
+            raise RuntimeError("MS-Diffusion: CLIP-L pooled shape %s, expected (%d,768)"
+                               % (tuple(phrase_embeds.shape), num_ref))
+
+        enc_len = _token_encoder(clip)
+        seen = {}
+        pidx = []
+        for p in phrases:
+            n = seen.get(p, 0)
+            seen[p] = n + 1
+            pidx.append(_get_phrase_idx(enc_len, p, clean_prompt, num=n))
+        phrase_idxes = torch.tensor([pidx], dtype=torch.int32)           # [1, num_ref, 2]
+
+        dev = mm.get_torch_device()
+        dtype = torch.float16
+        with torch.no_grad():
+            ms_adapter.image_proj.to(device=dev, dtype=dtype)
+            x = hidden.to(device=dev, dtype=dtype)
+            grounding = {
+                "boxes": batch_boxes.to(device=dev, dtype=dtype),
+                "phrase_embeds": phrase_embeds.to(device=dev, dtype=dtype),
+                "drop_grounding_tokens": [1 if drop_grounding_tokens else 0],
+            }
+            ip_tokens = ms_adapter.image_proj(x, grounding_kwargs=grounding)
+            ip_tokens = ip_tokens.view(1, -1, ip_tokens.shape[-2], ip_tokens.shape[-1])
+            ip_tokens = ip_tokens.view(1, ip_tokens.shape[-3] * ip_tokens.shape[-2], ip_tokens.shape[-1])
+            dummy = ms_adapter.dummy_image_tokens.to(device=dev, dtype=dtype)
+            image_prompt_embeds = torch.cat([dummy, ip_tokens], dim=1).float().cpu()
+            ms_adapter.image_proj.to(device=torch.device("cpu"))
+            ms_adapter._placed = None
+
+        assert image_prompt_embeds.shape == (1, NUM_DUMMY_TOKENS + NUM_TOKENS * num_ref, 2048)
+        embeds = MSEmbeds(
+            image_prompt_embeds=image_prompt_embeds,
+            boxes=boxes_t,
+            phrase_idxes=phrase_idxes,
+            num_ref=num_ref,
+            scale=1.0,                                   # set by MSDiffusionApply
+            mask_threshold=None if not layout_on else float(mask_threshold),
+            start_step=5,                               # set by MSDiffusionApply
+        )
+        return (embeds, clean_prompt)
+
+
+class MSDiffusionApply:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {"required": {
+            "model": ("MODEL",),
+            "ms_adapter": ("MS_ADAPTER",),
+            "ms_embeds": ("MS_EMBEDS",),
+            "scale": ("FLOAT", {"default": 0.8, "min": 0.0, "max": 2.0, "step": 0.05}),
+            "start_step": ("INT", {"default": 5, "min": 0, "max": 1024,
+                                   "tooltip": "pseudo attention mask kicks in after this many "
+                                              "denoising steps (only when mask_threshold > 0)"}),
+        }}
+
+    RETURN_TYPES = ("MODEL",)
+    FUNCTION = "apply"
+    CATEGORY = "MSdiffusion/Native"
+    DESCRIPTION = ("Patch every native SDXL attn2 with MS-Diffusion dual-branch "
+                   "cross-attention. Output is a normal MODEL: use stock KSampler / "
+                   "ControlNet / FreeU / LoRA nodes downstream.")
+
+    def apply(self, model, ms_adapter, ms_embeds, scale, start_step):
+        import copy
+        m = model.clone()
+        desc = enumerate_attn2(m.model.diffusion_model)
+        if not desc:
+            raise RuntimeError("MS-Diffusion: no attn2 modules found. This adapter requires a stock SDXL UNet.")
+        perm = build_block_weights(desc, ms_adapter.hidden_sizes)
+
+        emb = copy.copy(ms_embeds)
+        emb.scale = float(scale)
+        emb.start_step = int(start_step)
+        sid = new_state(ms_adapter, emb)
+
+        # IMPORTANT: set_model_patch_replace() REPLACES the whole
+        # transformer_options dict and silently drops previously attached keys.
+        # Build patches_replace in place on the cloned options instead.
+        to = m.model_options["transformer_options"]
+        pr = dict(to.get("patches_replace", {}))
+        a2 = dict(pr.get("attn2", {}))
+        for pos, (key3, _inner, _section) in enumerate(desc):
+            a2[key3] = make_attn2_patch(sid, perm[pos], is_sentinel=(pos == 0))
+        pr["attn2"] = a2
+        to["patches_replace"] = pr
+        print("[MS-Diffusion] patched %d attn2 layers (scale=%.2f, layout=%s)"
+              % (len(desc), emb.scale, "on" if emb.mask_threshold is not None else "off"))
+        return (m,)
+
+
+# No custom image-batch node: ComfyUI's native `ImageBatch` already stacks
+# images into [N,H,W,3] and aligns size/channels. Chain it for N > 2.
+# (The old `MS_Object_img_Batch` / `MSDiffusionImageBatch` ids are intentionally
+#  gone; replace them with the native ImageBatch node in old graphs.)
+
+NODE_CLASS_MAPPINGS = {
+    "MSDiffusionLoader": MSDiffusionLoader,
+    "MSDiffusionConditioning": MSDiffusionConditioning,
+    "MSDiffusionApply": MSDiffusionApply,
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "MSDiffusionLoader": "MS-Diffusion Loader (Native)",
+    "MSDiffusionConditioning": "MS-Diffusion Conditioning (Native)",
+    "MSDiffusionApply": "MS-Diffusion Apply to MODEL (Native)",
+}