rti/model.py

"""Region Token Interface (RTI): a frozen MiniT2I backbone whose middle blocks run
on R region tokens instead of all N patches."""
from __future__ import annotations

import math
from typing import Optional

import torch
import torch.nn as nn
import torch.nn.functional as F

from .minit2i import (
    rotate_half, RMSNorm, SwiGLUMlp, TextRotaryEmbedding1D,
    get_2d_sincos_pos_embed, MiniT2I_models, MiniT2I,
)
from .lora import LoRALinear


def _vision_rope_angles(side: int, head_dim: int, device, theta: float = 10000.0):
    """Per-token 2D rope PHASES [side*side, head_dim] (row-major)."""
    half = head_dim // 2
    q = half // 2
    freqs = 1.0 / (theta ** (torch.arange(0, q, device=device, dtype=torch.float32) / q))
    ys, xs = torch.meshgrid(torch.arange(side, device=device, dtype=torch.float32),
                            torch.arange(side, device=device, dtype=torch.float32), indexing="ij")
    ay = ys.reshape(-1)[:, None] * freqs[None]
    ax = xs.reshape(-1)[:, None] * freqs[None]
    ang = torch.cat([ay, ax], dim=-1)
    return torch.cat([ang, ang], dim=-1)


def _vision_rope_tables(side: int, head_dim: int, device, dtype, theta: float = 10000.0):
    ang = _vision_rope_angles(side, head_dim, device, theta)
    return ang.cos().to(dtype), ang.sin().to(dtype)


def _text_rope_tables(length: int, head_dim: int, device, dtype, theta: float = 10000.0):
    half = head_dim // 2
    freqs = 1.0 / (theta ** (torch.arange(0, half, device=device, dtype=torch.float32) / half))
    ang = torch.arange(length, device=device, dtype=torch.float32)[:, None] * freqs[None]
    ang = torch.cat([ang, ang], dim=-1)
    return ang.cos().to(dtype), ang.sin().to(dtype)


def _apply_img_rope(x, cos, sin):
    """x [B, N, H, hd]; cos/sin [B, N, hd] -> rotary-applied x."""
    cos = cos.unsqueeze(2).to(x.dtype)
    sin = sin.unsqueeze(2).to(x.dtype)
    return x * cos + rotate_half(x) * sin


_HILBERT_CACHE: dict = {}


def _hilbert_perm(side: int, device) -> torch.Tensor:
    """Row-major token indices in Hilbert-curve order for a side x side grid
    (side = 2^n). perm[d] is the flat index of the d-th point on the curve."""
    key = (side, str(device))
    if key not in _HILBERT_CACHE:
        idx = []
        for d in range(side * side):
            rx = ry = 0
            x = y = 0
            t = d
            s = 1
            while s < side:
                rx = 1 & (t // 2)
                ry = 1 & (t ^ rx)
                if ry == 0:
                    if rx == 1:
                        x = s - 1 - x
                        y = s - 1 - y
                    x, y = y, x
                x += s * rx
                y += s * ry
                t //= 4
                s *= 2
            idx.append(y * side + x)
        _HILBERT_CACHE[key] = torch.tensor(idx, dtype=torch.long, device=device)
    return _HILBERT_CACHE[key]


def _attn_fast(q, k, v, scale, key_bias=None):
    """SDPA/flash joint attention; q/k/v [B, L, H, d] -> [B, L, H, d]."""
    qq, kk, vv = (t.transpose(1, 2) for t in (q, k, v))
    mask = None
    if key_bias is not None:
        mask = key_bias[:, None, None, :].to(qq.dtype)
    out = F.scaled_dot_product_attention(qq, kk, vv, attn_mask=mask, scale=scale)
    return out.transpose(1, 2)


class RegionInterface(nn.Module):
    """Read: soft-pool each region's members into one token. Write: give each
    member its own delta from its region's delta. Initialised to mean-pool and
    broadcast."""

    def __init__(self, dim, n_size=11):
        super().__init__()
        self.read_score = nn.Linear(dim, 1)
        self.size_emb = nn.Embedding(n_size, dim)
        self.write = nn.Linear(2 * dim, dim)          # [member | region_delta] -> member delta
        nn.init.zeros_(self.read_score.weight)
        nn.init.zeros_(self.read_score.bias)          # -> mean
        nn.init.zeros_(self.size_emb.weight)          # -> no-op
        with torch.no_grad():                         # write -> copy the region-delta half
            self.write.weight.zero_()
            self.write.weight[:, dim:].copy_(torch.eye(dim))
            self.write.bias.zero_()

    def pool(self, xc, lab, R, cos_c, sin_c):
        """xc [B,li,D] in curve order, lab [B,li] in [0,R) -> tokens [B,R,D] and
        their rope, taken at the members' mean position. Runs in fp32."""
        b, li, D = xc.shape
        odt = xc.dtype
        xf = xc.float()
        w = self.read_score(xc).squeeze(-1).float()               # [B,li]
        cmax = xf.new_full((b, R), -1e9).scatter_reduce_(1, lab, w, reduce="amax", include_self=False)
        e = (w - torch.gather(cmax, 1, lab)).exp()
        denom = torch.gather(xf.new_zeros(b, R).scatter_add_(1, lab, e), 1, lab).clamp_min(1e-9)
        wsoft = (e / denom).unsqueeze(-1)                         # [B,li,1] segment-softmax
        idxD = lab.unsqueeze(-1).expand(-1, -1, D)
        cell = xf.new_zeros(b, R, D).scatter_add_(1, idxD, xf * wsoft)
        cnt = xf.new_zeros(b, R).scatter_add_(1, lab, torch.ones_like(w))
        sz = self.size_emb(cnt.clamp_min(1).log2().long()
                           .clamp_max(self.size_emb.num_embeddings - 1)).float()
        cell = cell + sz
        hd = cos_c.shape[-1]
        cch = lab.unsqueeze(-1).expand(-1, -1, hd)
        cden = cnt.clamp_min(1).unsqueeze(-1)
        cos_cell = xf.new_zeros(b, R, hd).scatter_add_(1, cch, cos_c.float()) / cden
        sin_cell = xf.new_zeros(b, R, hd).scatter_add_(1, cch, sin_c.float()) / cden
        return cell.to(odt), cos_cell.to(odt), sin_cell.to(odt)

    def unpool(self, xc, cell_delta, lab):
        """cell_delta [B,R,D] -> per-member delta [B,li,D]."""
        b, li, D = xc.shape
        cell_per_tok = torch.gather(cell_delta, 1, lab.unsqueeze(-1).expand(-1, -1, D))
        return self.write(torch.cat([xc, cell_per_tok], dim=-1))


class DoubleBlock(nn.Module):
    """MMDiT double-stream block; same parameters as the MiniT2I backbone."""

    def __init__(self, hidden, txt_hidden, num_heads, head_dim, mlp_ratio):
        super().__init__()
        self.num_heads, self.head_dim = num_heads, head_dim
        self.scale = head_dim ** -0.5
        inner = num_heads * head_dim
        self.img_norm1 = RMSNorm(hidden)
        self.img_norm2 = RMSNorm(hidden)
        self.txt_norm1 = RMSNorm(txt_hidden)
        self.txt_norm2 = RMSNorm(txt_hidden)
        self.img_qkv = nn.Linear(hidden, inner * 3)
        self.txt_qkv = nn.Linear(txt_hidden, inner * 3)
        self.q_norm = RMSNorm(head_dim)
        self.k_norm = RMSNorm(head_dim)
        self.text_rope = TextRotaryEmbedding1D(head_dim)
        self.img_attn_proj = nn.Linear(inner, hidden)
        self.txt_attn_proj = nn.Linear(inner, txt_hidden)
        self.img_mlp = SwiGLUMlp(hidden, int(hidden * mlp_ratio))
        self.txt_mlp = SwiGLUMlp(txt_hidden, int(txt_hidden * mlp_ratio))

    def _qkv(self, x_norm, qkv, N):
        b = x_norm.shape[0]
        t = qkv(x_norm).reshape(b, N, 3, self.num_heads, self.head_dim)
        return t[:, :, 0], t[:, :, 1], t[:, :, 2]

    def forward(self, x, txt, cos_full, sin_full):
        b, li, _ = x.shape
        lt = txt.shape[1]
        q_i, k_i, v_i = self._qkv(self.img_norm1(x), self.img_qkv, li)
        q_t, k_t, v_t = self._qkv(self.txt_norm1(txt), self.txt_qkv, lt)
        q_i, k_i = self.q_norm(q_i), self.k_norm(k_i)
        q_t, k_t = self.q_norm(q_t), self.k_norm(k_t)
        q_t, k_t = self.text_rope(q_t), self.text_rope(k_t)
        cos, sin = cos_full[None].expand(b, -1, -1), sin_full[None].expand(b, -1, -1)
        q_i, k_i = _apply_img_rope(q_i, cos, sin), _apply_img_rope(k_i, cos, sin)
        q = torch.cat([q_t, q_i], dim=1)
        k = torch.cat([k_t, k_i], dim=1)
        v = torch.cat([v_t, v_i], dim=1)
        out = _attn_fast(q, k, v, self.scale)
        x = x + self.img_attn_proj(out[:, lt:].reshape(b, li, -1))
        txt = txt + self.txt_attn_proj(out[:, :lt].reshape(b, lt, -1))
        x = x + self.img_mlp(self.img_norm2(x))
        txt = txt + self.txt_mlp(self.txt_norm2(txt))
        return x, txt


class RTIMiniT2I(nn.Module):
    """MiniT2I with the Region Token Interface over blocks `span[0]..span[1]`."""

    def __init__(
        self,
        base: MiniT2I,
        *,
        span: tuple[int, int] = (3, 13),
        capacity: float = 0.125,
        lora_rank: int = 32,
        lora_alpha: Optional[float] = None,
        lora_dropout: float = 0.0,
    ):
        super().__init__()
        net = base.net
        cfg = net.cfg
        self.cfg = cfg
        self.span = tuple(span)
        self.num_heads, self.head_dim = cfg.num_heads, cfg.head_dim
        self.side = net.latent_img_size
        self.patch_size = cfg.patch_size
        self.in_channels = cfg.in_channels
        self.hidden = cfg.hidden_size
        self.capacity = float(capacity)
        self.span_R: Optional[int] = None                   # explicit budget; else capacity*N

        self.img_embedder = net.img_embedder
        self.txt_embedder = net.txt_embedder
        self.mask_token = net.mask_token
        self.t_embedder = net.t_embedder
        self.pooled_embedder = net.pooled_embedder
        self.txt_preamble_blocks = net.txt_preamble_blocks
        self.final_layer = net.final_layer

        self.blocks = nn.ModuleList()
        for i in range(cfg.depth_double):
            blk = DoubleBlock(cfg.hidden_size, cfg.txt_hidden_size, cfg.num_heads,
                              cfg.head_dim, cfg.mlp_ratio)
            miss, unexp = blk.load_state_dict(net.double_blocks[i].state_dict(), strict=False)
            assert not unexp, f"block[{i}] init unexpected keys {unexp}"
            self.blocks.append(blk)

        self.iface = RegionInterface(cfg.hidden_size)
        self._setup_lora_and_freeze(lora_rank, lora_alpha, lora_dropout)
        for p in self.iface.parameters():
            p.requires_grad = True

    def _setup_lora_and_freeze(self, rank, alpha, dropout):
        for p in self.parameters():
            p.requires_grad = False

        def _wrap(blk):
            for name in ("img_qkv", "txt_qkv", "img_attn_proj", "txt_attn_proj"):
                setattr(blk, name, LoRALinear(getattr(blk, name), rank=rank, alpha=alpha, dropout=dropout))
            for mlp in (blk.img_mlp, blk.txt_mlp):
                mlp.w1 = LoRALinear(mlp.w1, rank=rank, alpha=alpha, dropout=dropout)
                mlp.w2 = LoRALinear(mlp.w2, rank=rank, alpha=alpha, dropout=dropout)
                mlp.w3 = LoRALinear(mlp.w3, rank=rank, alpha=alpha, dropout=dropout)

        for blk in self.blocks:
            _wrap(blk)
        for blk in self.txt_preamble_blocks:                 # text preamble also LoRA
            blk.qkv = LoRALinear(blk.qkv, rank=rank, alpha=alpha, dropout=dropout)
            blk.attn_proj = LoRALinear(blk.attn_proj, rank=rank, alpha=alpha, dropout=dropout)
            for w in ("w1", "w2", "w3"):
                setattr(blk.mlp, w, LoRALinear(getattr(blk.mlp, w), rank=rank, alpha=alpha, dropout=dropout))
        self.final_layer.linear = LoRALinear(self.final_layer.linear, rank=rank, alpha=alpha, dropout=dropout)

    def trainable_state(self):
        """LoRA and interface tensors, i.e. what a checkpoint stores."""
        return {k: v for k, v in self.state_dict().items()
                if ("lora_" in k) or k.startswith("iface.")}

    @torch.no_grad()
    def region_labels(self, x, R):
        """Region ids [B, li] in curve order: cut at the R-1 largest adjacent
        feature distances along the curve."""
        b, li, _ = x.shape
        perm = _hilbert_perm(int(math.sqrt(li)), x.device)
        xp = x[:, perm].float()
        bs = (xp[:, 1:] - xp[:, :-1]).norm(dim=-1)
        cuts = bs.topk(R - 1, dim=1).indices
        st = torch.zeros(b, li, device=x.device)
        st.scatter_(1, cuts + 1, 1.0)
        return st.cumsum(dim=1).long()

    @staticmethod
    def _core_on_regions(blk, cell, txt, cos_cell, sin_cell):
        """Run one trunk block on the R region tokens, using the pooled rope."""
        b, R, _ = cell.shape
        lt = txt.shape[1]
        q_i, k_i, v_i = blk._qkv(blk.img_norm1(cell), blk.img_qkv, R)
        q_t, k_t, v_t = blk._qkv(blk.txt_norm1(txt), blk.txt_qkv, lt)
        q_i, k_i = blk.q_norm(q_i), blk.k_norm(k_i)
        q_t, k_t = blk.q_norm(q_t), blk.k_norm(k_t)
        q_t, k_t = blk.text_rope(q_t), blk.text_rope(k_t)
        q_i, k_i = _apply_img_rope(q_i, cos_cell, sin_cell), _apply_img_rope(k_i, cos_cell, sin_cell)
        q = torch.cat([q_t, q_i], 1)
        k = torch.cat([k_t, k_i], 1)
        v = torch.cat([v_t, v_i], 1)
        out = _attn_fast(q, k, v, blk.scale)
        cell = cell + blk.img_attn_proj(out[:, lt:].reshape(b, R, -1))
        txt = txt + blk.txt_attn_proj(out[:, :lt].reshape(b, lt, -1))
        cell = cell + blk.img_mlp(blk.img_norm2(cell))
        txt = txt + blk.txt_mlp(blk.txt_norm2(txt))
        return cell, txt

    def _decode(self, x, txt):
        combined = torch.cat([txt, x], dim=1)
        out = self.final_layer(combined)
        img_out = out[:, txt.shape[1]:, :]
        b, p, c = img_out.shape[0], self.patch_size, self.in_channels
        h = w = int(math.sqrt(img_out.shape[1]))
        img_out = img_out.reshape(b, h, w, p, p, c)
        img_out = torch.einsum("nhwpqc->nchpwq", img_out)
        return img_out.reshape(b, c, h * p, w * p)

    def forward_xpred(self, img, t, context, mask, *, capacity_override: Optional[float] = None,
                      budget: Optional[int] = None):
        """Predict the clean image. `budget` (or `capacity_override`) sets R."""
        if img.ndim == 4 and img.shape[1] != self.in_channels:
            img = img.permute(0, 3, 1, 2)
        mask = mask.to(device=context.device)
        context = torch.where(mask[:, :, None] > 0.5, context, self.mask_token.to(dtype=context.dtype))
        x = self.img_embedder(img)
        pos = get_2d_sincos_pos_embed(self.hidden, self.side, x.device, x.dtype)
        x = x + pos[None]
        txt = self.txt_embedder(context.to(dtype=self.txt_embedder.weight.dtype))
        for blk in self.txt_preamble_blocks:
            txt = blk(txt)
        cos_full, sin_full = _vision_rope_tables(self.side, self.head_dim, x.device, x.dtype)

        b, li, _ = x.shape
        s0, s1 = self.span
        if budget is not None:
            R = int(budget)
        elif self.span_R is not None:
            R = int(self.span_R)
        else:
            cap = capacity_override if capacity_override is not None else self.capacity
            R = int(round(cap * li))
        R = max(1, min(R, li))

        if R >= li:                                          # identity budget -> dense
            for blk in self.blocks:
                x, txt = blk(x, txt, cos_full, sin_full)
            self.last_labels = None
            return self._decode(x, txt)

        for i in range(s0):                                  # prelude, dense
            x, txt = self.blocks[i](x, txt, cos_full, sin_full)

        # Read
        perm = _hilbert_perm(int(math.sqrt(li)), x.device)
        lab = self.region_labels(x, R)
        xc = x[:, perm]
        cos_c = cos_full[perm][None].expand(b, -1, -1)
        sin_c = sin_full[perm][None].expand(b, -1, -1)
        cell, cos_cell, sin_cell = self.iface.pool(xc, lab, R, cos_c, sin_c)
        cell_in = cell
        self.last_labels = lab.detach()

        # core on the R region tokens
        for i in range(s0, s1 + 1):
            cell, txt = self._core_on_regions(self.blocks[i], cell, txt, cos_cell, sin_cell)

        # Write
        x = x.index_copy(1, perm, xc + self.iface.unpool(xc, cell - cell_in, lab))

        for i in range(s1 + 1, len(self.blocks)):            # coda, dense
            x, txt = self.blocks[i](x, txt, cos_full, sin_full)
        return self._decode(x, txt)

    def forward(self, img, t, context, mask, **kw):
        return self.forward_xpred(img, t, context, mask, **kw)


def build_rti(variant: str = "MiniT2I-B/16", *, span: Optional[tuple[int, int]] = None,
              capacity: float = 0.125, lora_rank: int = 32, lora_alpha=None,
              lora_dropout: float = 0.0, pretrained: bool = True, cache_dir=None,
              ckpt: Optional[str] = None, device="cpu", dtype=None):
    """Build RTI on a pretrained MiniT2I backbone and optionally load an adapter
    checkpoint. `span` and `lora_rank` are read from the checkpoint when present."""
    meta = {}
    if ckpt is not None:
        _raw = torch.load(ckpt, map_location="cpu", weights_only=False)
        meta = _raw.get("meta", {}) if isinstance(_raw, dict) else {}
        del _raw
    if span is None:
        span = tuple(meta["span"]) if "span" in meta else ((3, 13) if "B/16" in variant else (3, 19))
    if "lora_rank" in meta:
        lora_rank = int(meta["lora_rank"])
    if "variant" in meta and meta["variant"] != variant:
        raise ValueError(f"checkpoint is for {meta['variant']}, but --variant is {variant}")
    base = MiniT2I_models[variant]()
    if pretrained:
        from .minit2i import load_pretrained
        load_pretrained(variant, base, cache_dir=cache_dir, device="cpu")
    model = RTIMiniT2I(base, span=span, capacity=capacity, lora_rank=lora_rank,
                       lora_alpha=lora_alpha, lora_dropout=lora_dropout)
    if ckpt is not None:
        sd = torch.load(ckpt, map_location="cpu", weights_only=False)
        sd = sd.get("state_dict") or sd.get("ema") or sd.get("model") or sd
        sd = {k: v for k, v in sd.items() if not k.startswith("teacher_")}
        # legacy key name
        sd = {(k.replace("span_iface.", "iface.", 1) if k.startswith("span_iface.") else k): v
              for k, v in sd.items()}
        miss, unexp = model.load_state_dict(sd, strict=False)
        assert not unexp, f"checkpoint has keys this model does not define: {unexp[:8]}"
        n_loaded = len(sd)
        print(f"[rti] loaded {n_loaded} adapter tensors from {ckpt}")
    model = model.to(device)
    if dtype is not None:
        model = model.to(dtype)
    return model


RTI_models = {
    "RTI-B/16": lambda **kw: build_rti("MiniT2I-B/16", **kw),
    "RTI-L/16": lambda **kw: build_rti("MiniT2I-L/16", **kw),
}