"""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), }