"""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