Spaces:
Running on Zero
Running on Zero
| """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.")} | |
| 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() | |
| 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), | |
| } | |