Upload folder using huggingface_hub

README.md

@@ -1,13 +1,45 @@
 ---
-title: Rti Pixel Diffusion
-emoji: 🌖
-colorFrom: red
+title: RTI - Elastic Token Compression
+emoji: 🧩
+colorFrom: pink
 colorTo: purple
 sdk: gradio
 sdk_version: 6.30.0
-python_version: '3.12'
 app_file: app.py
-pinned: false
+short_description: Pixel-space t2i with a region-token budget dial
+python_version: "3.12"
+startup_duration_timeout: 30m
+license: mit
+models:
+  - eduardzamfir/RTI
+  - MiniT2I/MiniT2I
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# RTI — Elastic Token Compression for Pixel-Space Diffusion Transformers
+
+Interactive demo for **RTI (Region Token Interface)** — [arXiv:2608.29281](https://arxiv.org/abs/2608.29281).
+A frozen MiniT2I pixel-space diffusion backbone runs its middle blocks on
+`R ≪ N` content-shaped region tokens instead of all patches. One adapter
+checkpoint serves **every** budget: the region count was drawn at random during
+fine-tuning, so you can set R at inference to move along the quality/compute
+curve.
+
+**Budget dial (R at 512×512, N = 1024 patch tokens):**
+
+| R | Meaning | Speedup (paper) |
+|---|---|---|
+| 1024 | dense backbone | 1.0× |
+| 512 | matches dense quality | ~2.1× |
+| 256 | small quality drop | ~2.6× |
+| 64 | extreme compression | — |
+
+Two released arms are preloaded — **RTI-B/16** and **RTI-L/16** (LoRA + the
+small interface on the frozen MiniT2I trunk) — sharing a single FLAN-T5-Large
+prompt encoder. Prompts are encoded at fp32, sampling runs bf16 Euler
+flow-matching with CFG 5.0 over 100 steps, exactly as in the authors'
+reference implementation (`rti.sampling.sample`).
+
+The `rti/` package in this Space is the inference subset of
+[eduardzamfir/RTI](https://github.com/eduardzamfir/RTI) (MIT license). Example
+prompts are the ones the authors showcase on the
+[project page](https://eduardzamfir.github.io/rti).

app.py

@@ -0,0 +1,241 @@
+"""Gradio demo for RTI — Elastic Token Compression for Pixel-Space Diffusion
+Transformers (arXiv:2608.29281).
+
+One adapter checkpoint (LoRA + interface on a frozen MiniT2I backbone) serves
+EVERY region-token budget: pick R at inference to trade quality for compute.
+"""
+
+import os
+
+os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
+
+import spaces  # MUST come before any torch / CUDA-touching import
+
+import random
+import time
+
+import torch
+import gradio as gr
+from huggingface_hub import hf_hub_download
+from transformers import AutoTokenizer, T5EncoderModel
+
+from rti import build_rti
+from rti.sampling import sample
+
+ADAPTER_REPO = "eduardzamfir/RTI"
+DTYPE = torch.bfloat16
+MAX_SEED = 2**31 - 1
+IMAGE_SIZE = 512
+N_PATCH_TOKENS = 1024  # 32x32 patch grid at 512/16 — R = 1024 means dense
+
+# Authors' released arms: B/16 (62 MB adapter) and L/16 (137 MB adapter).
+VARIANTS = {
+    "RTI-B/16 (base)": ("MiniT2I-B/16", "rti_b16_40k.pt"),
+    "RTI-L/16 (large)": ("MiniT2I-L/16", "rti_l16_40k.pt"),
+}
+
+CSS = """
+#col-container { max-width: 1100px; margin: 0 auto; }
+.dark .gradio-container { color: var(--body-text-color); }
+"""
+
+# Prompts the authors showcase on the RTI project page / paper figures.
+EXAMPLES = [
+    ["A polished white marble statue of a young woman with intricately carved curls and deeply folded drapery, displayed in a museum under warm directional light against a dark background, fine chisel marks visible in the stone."],
+    ["Golden-hour street photograph of a rain-soaked crosswalk in Tokyo, umbrellas and neon reflections, shot on Portra 400 with a 35mm Summicron, shallow depth of field."],
+    ["A lone red fox standing in deep snow at dawn, individual guard hairs and whiskers sharp, breath visible, flat white field and pale sky behind, wildlife photograph with a long lens."],
+    ["Reflection in an astronaut's gold visor: the surface of Mars with a rocket lifting off in the middle distance, exhaust plume and a rolling dust cloud, the astronaut's own gloved hand raised, ultra detailed, black sky."],
+    ["A small friendly robot in a bright research lab holding a white sign that says exactly HELLO in bold black letters. No other readable text."],
+    ["The Mona Lisa rendered as a Minecraft scene, blocky voxel textures."],
+]
+
+
+def _load(variant_label):
+    variant, ckpt_name = VARIANTS[variant_label]
+    ckpt = hf_hub_download(ADAPTER_REPO, ckpt_name)
+    model = build_rti(
+        variant, ckpt=ckpt, device="cuda", dtype=DTYPE
+    ).eval()
+    return model
+
+
+print("[app] loading RTI backbones (adapters + MiniT2I trunk)…", flush=True)
+MODELS = {label: _load(label) for label in VARIANTS}
+print("[app] models loaded", flush=True)
+
+# Shared FLAN-T5-Large prompt encoder (fp32, as in the authors' reference code).
+TEXT_ENCODER = "google/flan-t5-large"
+print("[app] loading FLAN-T5-Large prompt encoder…", flush=True)
+TOKENIZER = AutoTokenizer.from_pretrained(TEXT_ENCODER)
+TEXT_ENC = (
+    T5EncoderModel.from_pretrained(TEXT_ENCODER, torch_dtype=torch.float32)
+    .to("cuda")
+    .eval()
+)
+print("[app] prompt encoder loaded", flush=True)
+
+
+@torch.no_grad()
+def _encode(prompts):
+    tok = TOKENIZER(
+        list(prompts),
+        padding="max_length",
+        truncation=True,
+        max_length=256,
+        return_tensors="pt",
+    )
+    out = TEXT_ENC(
+        input_ids=tok.input_ids.to("cuda"),
+        attention_mask=tok.attention_mask.to("cuda"),
+    )
+    return out.last_hidden_state.float(), tok.attention_mask.float()
+
+
+@spaces.GPU(duration=120)
+def generate(
+    prompt: str,
+    variant_label: str = "RTI-B/16 (base)",
+    budget: int = 512,
+    guidance_scale: float = 5.0,
+    num_inference_steps: int = 100,
+    seed: int = 0,
+    randomize_seed: bool = True,
+    progress=gr.Progress(track_tqdm=True),
+):
+    """Generate a 512x512 image from a text prompt at a chosen region-token budget.
+
+    Args:
+        prompt: what to generate.
+        variant_label: which released RTI adapter to use (B/16 or L/16).
+        budget: region-token count R for the model's core blocks. 1024 = dense
+            backbone (1.0x), 512 ≈ 1.8-2.1x faster, 256 ≈ 2.2-2.6x faster, down
+            to 64 for maximum compression.
+        guidance_scale: classifier-free guidance scale (paper uses 5.0).
+        num_inference_steps: Euler flow-matching steps (paper uses 100).
+        seed: RNG seed for reproducibility.
+        randomize_seed: draw a fresh seed each run instead of using `seed`.
+    """
+    if not prompt or not prompt.strip():
+        raise gr.Error("Please enter a prompt.")
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+    seed = int(seed)
+    budget = int(budget)
+    num_inference_steps = int(num_inference_steps)
+
+    model = MODELS[variant_label]
+    ctx, mask = _encode([prompt.strip()])
+    g = torch.Generator(device="cuda").manual_seed(seed)
+
+    t0 = time.perf_counter()
+    imgs = sample(
+        model,
+        ctx,
+        mask,
+        budget=budget,
+        steps=num_inference_steps,
+        cfg_scale=float(guidance_scale),
+        image_size=IMAGE_SIZE,
+        noise_scale=2.0,
+        t_eps=0.001,
+        cfg_interval=(0.0, 1.0),
+        generator=g,
+        device="cuda",
+        dtype=DTYPE,
+        progress=True,
+    )
+    elapsed = time.perf_counter() - t0
+
+    from torchvision.transforms.functional import to_pil_image
+
+    image = to_pil_image(imgs[0].clamp(0, 1))
+    mode = "dense (R = N = 1024)" if budget >= N_PATCH_TOKENS else f"R = {budget}"
+    return image, seed, f"{elapsed:.1f}s · {mode}"
+
+
+with gr.Blocks(theme=gr.themes.Citrus(), css=CSS) as demo:
+    gr.Markdown(
+        """
+        # 🧩 RTI — Elastic Token Compression for Pixel-Space Diffusion Transformers
+        One adapter checkpoint serves **every** region-token budget. Set the budget
+        dial below to move along the quality/compute curve: **R = 1024** runs the
+        dense backbone, **512** matches dense quality at ~2× speed, **256** trades a
+        little quality for ~2.6× speed, down to **64** for extreme compression.
+        Images are 512×512, generated pixel-space (no VAE).
+        [Paper](https://arxiv.org/abs/2608.29281) ·
+        [Code](https://github.com/eduardzamfir/RTI) ·
+        [Adapters](https://huggingface.co/eduardzamfir/RTI) ·
+        [Backbone: MiniT2I](https://huggingface.co/MiniT2I/MiniT2I)
+        """
+    )
+    with gr.Row():
+        with gr.Column(scale=1):
+            prompt = gr.Textbox(
+                label="Prompt",
+                placeholder="Describe the image you want to generate…",
+                lines=3,
+            )
+            variant_label = gr.Radio(
+                choices=list(VARIANTS.keys()),
+                value="RTI-B/16 (base)",
+                label="Model",
+            )
+            budget = gr.Slider(
+                label="Region-token budget R (compute dial)",
+                minimum=64,
+                maximum=1024,
+                step=64,
+                value=512,
+                info="1024 = dense backbone · 512 ≈ 2.1× speed · 256 ≈ 2.6× speed",
+            )
+            run_btn = gr.Button("Generate", variant="primary")
+            with gr.Accordion("Advanced settings", open=False):
+                guidance_scale = gr.Slider(
+                    label="Guidance scale",
+                    minimum=1.0,
+                    maximum=10.0,
+                    step=0.1,
+                    value=5.0,
+                )
+                num_inference_steps = gr.Slider(
+                    label="Inference steps",
+                    minimum=10,
+                    maximum=150,
+                    step=1,
+                    value=100,
+                )
+                with gr.Row():
+                    seed = gr.Slider(
+                        label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0
+                    )
+                    randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
+        with gr.Column(scale=1):
+            output = gr.Image(label="Result", height=512)
+            status = gr.Markdown(label="Status")
+
+    gr.Examples(
+        examples=EXAMPLES,
+        inputs=[prompt],
+        outputs=[output, seed, status],
+        fn=generate,
+        cache_examples=True,
+        cache_mode="lazy",
+    )
+
+    inputs = [
+        prompt,
+        variant_label,
+        budget,
+        guidance_scale,
+        num_inference_steps,
+        seed,
+        randomize_seed,
+    ]
+    run_btn.click(
+        generate, inputs=inputs, outputs=[output, seed, status], api_name="generate"
+    )
+    prompt.submit(
+        generate, inputs=inputs, outputs=[output, seed, status], api_name="generate"
+    )
+
+demo.launch(mcp_server=True)

requirements.txt

@@ -0,0 +1,10 @@
+torch
+torchvision
+transformers>=4.40
+safetensors
+pillow
+numpy
+tqdm
+einops
+accelerate
+sentencepiece

rti/__init__.py

@@ -0,0 +1,6 @@
+"""Region Token Interface (RTI) for pixel-space diffusion transformers."""
+from .model import RTIMiniT2I, RegionInterface, build_rti, RTI_models
+from .minit2i import MiniT2I, MiniT2I_models, load_pretrained
+
+__all__ = ["RTIMiniT2I", "RegionInterface", "build_rti", "RTI_models",
+           "MiniT2I", "MiniT2I_models", "load_pretrained"]

rti/lora.py

@@ -0,0 +1,58 @@
+"""LoRA adapter for a frozen nn.Linear."""
+from __future__ import annotations
+
+import math
+
+import torch
+import torch.nn as nn
+
+
+class LoRALinear(nn.Module):
+    def __init__(
+        self,
+        base_linear: nn.Linear,
+        rank: int = 8,
+        alpha: float | None = None,
+        dropout: float = 0.0,
+    ):
+        super().__init__()
+        if rank <= 0:
+            raise ValueError(f"LoRA rank must be positive; got {rank}")
+        if alpha is None:
+            alpha = float(rank)
+        self.rank = rank
+        self.alpha = float(alpha)
+        self.scale = self.alpha / rank
+        self.base = base_linear
+
+        in_f = base_linear.in_features
+        out_f = base_linear.out_features
+        # Match base's device/dtype so post-construction LoRA wrapping
+        # (e.g. apply_lora_and_freeze after the model is on GPU) doesn't leave
+        # the adapter params on CPU.
+        device = base_linear.weight.device
+        dtype  = base_linear.weight.dtype
+        self.lora_A = nn.Parameter(torch.empty(rank, in_f, device=device, dtype=dtype))
+        self.lora_B = nn.Parameter(torch.zeros(out_f, rank, device=device, dtype=dtype))
+        nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
+
+        self.lora_dropout = nn.Dropout(dropout) if dropout > 0 else nn.Identity()
+        # When False the adapter is skipped and forward == base(x). Lets callers
+        # evaluate the pristine frozen base (e.g. a self-distillation teacher)
+        # without unwrapping. Defaults True — backward-compatible.
+        self.enabled = True
+
+    @property
+    def in_features(self) -> int:
+        return self.base.in_features
+
+    @property
+    def out_features(self) -> int:
+        return self.base.out_features
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        out = self.base(x)
+        if not self.enabled:
+            return out
+        update = self.lora_dropout(x) @ self.lora_A.t() @ self.lora_B.t()
+        return out + update * self.scale

rti/minit2i.py

@@ -0,0 +1,553 @@
+"""MiniT2I: pixel-space MMDiT text-to-image backbone.
+
+Port of the released model (https://huggingface.co/MiniT2I/MiniT2I) that matches
+the published checkpoints key for key, so RoPE and block internals must not be
+"simplified".
+
+There is no AdaLN: conditioning runs entirely through the joint text+image
+attention stream, with frozen FLAN-T5-Large embeddings as the text stream. The
+time and pooled-text `vec` exist in the checkpoint but the blocks ignore them.
+
+Flow matching: x_t = t*img + (1-t)*noise with noise scaled by 2; the net predicts
+x0, and v = (x0 - x_t) / clamp(1-t). t=1 is clean, t=0 is noise.
+"""
+from __future__ import annotations
+
+import math
+import os
+from dataclasses import dataclass, replace
+from typing import Optional
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+MINIT2I_HF_REPO = "MiniT2I/MiniT2I"
+
+
+# ---------------------------------------------------------------------------
+# Primitives (verbatim from upstream mmdit.py)
+# ---------------------------------------------------------------------------
+def rotate_half(x: torch.Tensor) -> torch.Tensor:
+    # NOTE: interleaved variant (reshape → 2 → -1), NOT the half-split variant.
+    # The checkpoint's RoPE depends on this exact form.
+    x1, x2 = x.reshape(*x.shape[:-1], 2, -1).unbind(dim=-2)
+    return torch.cat((-x2, x1), dim=-1)
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(dim))
+        self.eps = eps
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        y = x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps)
+        return y * self.weight
+
+
+class TimestepEmbedder(nn.Module):
+    """Kept for state-dict compatibility; its output is unused in the forward."""
+
+    def __init__(self, hidden_size: int, frequency_embedding_size: int = 256):
+        super().__init__()
+        self.frequency_embedding_size = frequency_embedding_size
+        self.mlp = nn.Sequential(
+            nn.Linear(frequency_embedding_size, hidden_size),
+            nn.SiLU(),
+            nn.Linear(hidden_size, hidden_size),
+        )
+
+    def forward(self, t: torch.Tensor) -> torch.Tensor:
+        half = self.frequency_embedding_size // 2
+        freqs = torch.exp(
+            -math.log(10000.0)
+            * torch.arange(half, device=t.device, dtype=torch.float32)
+            / half
+        )
+        args = t.float()[:, None] * freqs[None]
+        emb = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+        return self.mlp(emb.to(dtype=self.mlp[0].weight.dtype))
+
+
+class BottleneckPatchEmbed(nn.Module):
+    """proj1: (in_ch -> pca) strided conv; proj2: (pca -> D) 1x1 conv."""
+
+    def __init__(self, img_size=512, patch_size=16, in_channels=3, pca_channels=128, hidden_size=768):
+        super().__init__()
+        self.img_size = img_size
+        self.patch_size = patch_size
+        self.proj1 = nn.Conv2d(in_channels, pca_channels, kernel_size=patch_size, stride=patch_size, bias=False)
+        self.proj2 = nn.Conv2d(pca_channels, hidden_size, kernel_size=1, stride=1, bias=True)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.proj2(self.proj1(x))
+        return x.flatten(2).transpose(1, 2)
+
+
+class SwiGLUMlp(nn.Module):
+    def __init__(self, in_features: int, hidden_features: int):
+        super().__init__()
+        hidden_dim = (hidden_features + 7) // 8 * 8
+        self.w1 = nn.Linear(in_features, hidden_dim, bias=False)
+        self.w3 = nn.Linear(in_features, hidden_dim, bias=False)
+        self.w2 = nn.Linear(hidden_dim, in_features, bias=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.w2(F.silu(self.w1(x)) * self.w3(x))
+
+
+class TextRotaryEmbedding1D(nn.Module):
+    def __init__(self, head_dim: int, theta: float = 10000.0):
+        super().__init__()
+        self.head_dim = head_dim
+        self.theta = theta
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        b, length, h, d = x.shape
+        inv = 1.0 / (self.theta ** (torch.arange(0, d, 2, device=x.device, dtype=torch.float32) / d))
+        pos = torch.arange(length, device=x.device, dtype=torch.float32)
+        angles = torch.einsum("l,f->lf", pos, inv)
+        angles = torch.cat([angles, angles], dim=-1)
+        cos = angles.cos().to(dtype=x.dtype)
+        sin = angles.sin().to(dtype=x.dtype)
+        return x * cos[None, :, None, :] + rotate_half(x) * sin[None, :, None, :]
+
+
+class VisionRotaryEmbeddingFast(nn.Module):
+    def __init__(self, head_dim: int, theta: float = 10000.0):
+        super().__init__()
+        self.dim = head_dim // 2
+        self.theta = theta
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        length = x.shape[1]
+        side = int(math.sqrt(length))
+        if side * side != length:
+            raise ValueError(f"image token length must be square, got {length}")
+        freqs = 1.0 / (
+            self.theta
+            ** (torch.arange(0, self.dim, 2, device=x.device, dtype=torch.float32)[: self.dim // 2] / self.dim)
+        )
+        t = torch.arange(side, device=x.device, dtype=torch.float32)
+        base = torch.einsum("l,f->lf", t, freqs)
+        f_h, f_w = torch.broadcast_tensors(base[:, None, :], base[None, :, :])
+        angles = torch.cat([f_h, f_w], dim=-1)
+        angles = torch.cat([angles, angles], dim=-1).reshape(length, -1)
+        cos = angles.cos().to(dtype=x.dtype)
+        sin = angles.sin().to(dtype=x.dtype)
+        return x * cos[None, :, None, :] + rotate_half(x) * sin[None, :, None, :]
+
+
+class MultiModalRotaryEmbeddingFast(nn.Module):
+    def __init__(self, head_dim: int):
+        super().__init__()
+        self.text_rope = TextRotaryEmbedding1D(head_dim)
+        self.vision_rope = VisionRotaryEmbeddingFast(head_dim)
+
+    def forward(self, x: torch.Tensor, txt_len: int) -> torch.Tensor:
+        txt = self.text_rope(x[:, :txt_len])
+        img = self.vision_rope(x[:, txt_len:])
+        return torch.cat([txt, img], dim=1)
+
+
+# ---------------------------------------------------------------------------
+# Blocks
+# ---------------------------------------------------------------------------
+class PlainTextTransformerBlock(nn.Module):
+    def __init__(self, hidden_size=768, num_heads=12, head_dim=64, mlp_ratio=2.6667):
+        super().__init__()
+        self.num_heads = num_heads
+        self.head_dim = head_dim
+        inner_dim = num_heads * head_dim
+        self.norm1 = RMSNorm(hidden_size)
+        self.norm2 = RMSNorm(hidden_size)
+        self.qkv = nn.Linear(hidden_size, inner_dim * 3)
+        self.attn_proj = nn.Linear(inner_dim, hidden_size)
+        self.mlp = SwiGLUMlp(hidden_size, int(hidden_size * mlp_ratio))
+        self.q_norm = RMSNorm(head_dim)
+        self.k_norm = RMSNorm(head_dim)
+        self.rope = TextRotaryEmbedding1D(head_dim)
+
+    def forward(self, txt: torch.Tensor) -> torch.Tensor:
+        b, length, _ = txt.shape
+        qkv = self.qkv(self.norm1(txt)).reshape(b, length, 3, self.num_heads, self.head_dim)
+        q, k, v = qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2]
+        q = self.rope(self.q_norm(q))
+        k = self.rope(self.k_norm(k))
+        attn = torch.einsum("bqhd,bkhd->bhqk", q, k) * (self.head_dim ** -0.5)
+        out = torch.einsum("bhqk,bkhd->bqhd", attn.softmax(dim=-1), v).reshape(b, length, -1)
+        txt = txt + self.attn_proj(out)
+        txt = txt + self.mlp(self.norm2(txt))
+        return txt
+
+
+class DoubleStreamDiTBlock(nn.Module):
+    def __init__(self, hidden_size=768, txt_hidden_size=768, num_heads=12, head_dim=64, mlp_ratio=2.6667):
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.txt_hidden_size = txt_hidden_size
+        self.num_heads = num_heads
+        self.head_dim = head_dim
+        inner_dim = num_heads * head_dim
+        self.img_norm1 = RMSNorm(hidden_size)
+        self.img_norm2 = RMSNorm(hidden_size)
+        self.txt_norm1 = RMSNorm(txt_hidden_size)
+        self.txt_norm2 = RMSNorm(txt_hidden_size)
+        self.img_qkv = nn.Linear(hidden_size, inner_dim * 3)
+        self.txt_qkv = nn.Linear(txt_hidden_size, inner_dim * 3)
+        self.q_norm = RMSNorm(head_dim)
+        self.k_norm = RMSNorm(head_dim)
+        self.rope = MultiModalRotaryEmbeddingFast(head_dim)
+        self.img_attn_proj = nn.Linear(inner_dim, hidden_size)
+        self.txt_attn_proj = nn.Linear(inner_dim, txt_hidden_size)
+        self.img_mlp = SwiGLUMlp(hidden_size, int(hidden_size * mlp_ratio))
+        self.txt_mlp = SwiGLUMlp(txt_hidden_size, int(txt_hidden_size * mlp_ratio))
+
+    def forward(self, x: torch.Tensor, txt: torch.Tensor, vec=None):
+        b, li, _ = x.shape
+        lt = txt.shape[1]
+        x_norm = self.img_norm1(x)
+        txt_norm = self.txt_norm1(txt)
+        qkv_i = self.img_qkv(x_norm).reshape(b, li, 3, self.num_heads, self.head_dim)
+        qkv_t = self.txt_qkv(txt_norm).reshape(b, lt, 3, self.num_heads, self.head_dim)
+        q_i, k_i, v_i = qkv_i[:, :, 0], qkv_i[:, :, 1], qkv_i[:, :, 2]
+        q_t, k_t, v_t = qkv_t[:, :, 0], qkv_t[:, :, 1], qkv_t[:, :, 2]
+        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 = self.rope(torch.cat([q_t, q_i], dim=1), txt_len=lt)
+        k = self.rope(torch.cat([k_t, k_i], dim=1), txt_len=lt)
+        v = torch.cat([v_t, v_i], dim=1)
+        attn = torch.einsum("bqhd,bkhd->bhqk", q, k) * (self.head_dim ** -0.5)
+        out = torch.einsum("bhqk,bkhd->bqhd", attn.softmax(dim=-1), v)
+        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 FinalLayer(nn.Module):
+    def __init__(self, hidden_size=768, patch_size=16, out_channels=3):
+        super().__init__()
+        self.patch_size = patch_size
+        self.out_channels = out_channels
+        self.norm_final = RMSNorm(hidden_size)
+        self.linear = nn.Linear(hidden_size, patch_size * patch_size * out_channels)
+
+    def forward(self, x: torch.Tensor, vec=None) -> torch.Tensor:
+        return self.linear(self.norm_final(x))
+
+
+def get_1d_sincos_pos_embed(embed_dim: int, pos: torch.Tensor) -> torch.Tensor:
+    omega = torch.arange(embed_dim // 2, device=pos.device, dtype=torch.float32)
+    omega = 1.0 / (10000 ** (omega / (embed_dim / 2.0)))
+    out = torch.einsum("m,d->md", pos.reshape(-1), omega)
+    return torch.cat([out.sin(), out.cos()], dim=1)
+
+
+def get_2d_sincos_pos_embed(embed_dim: int, grid_size: int, device, dtype) -> torch.Tensor:
+    grid_h = torch.arange(grid_size, device=device, dtype=torch.float32)
+    grid_w = torch.arange(grid_size, device=device, dtype=torch.float32)
+    grid = torch.meshgrid(grid_w, grid_h, indexing="xy")
+    grid = torch.stack(grid, dim=0).reshape(2, 1, grid_size, grid_size)
+    emb_h = get_1d_sincos_pos_embed(embed_dim // 2, grid[0])
+    emb_w = get_1d_sincos_pos_embed(embed_dim // 2, grid[1])
+    return torch.cat([emb_h, emb_w], dim=1).to(dtype=dtype)
+
+
+# ---------------------------------------------------------------------------
+# Config + net
+# ---------------------------------------------------------------------------
+@dataclass
+class MMJiTConfig:
+    image_size: int = 512
+    patch_size: int = 16
+    in_channels: int = 3
+    txt_input_size: int = 1024          # FLAN-T5-Large hidden size
+    hidden_size: int = 768
+    txt_hidden_size: int = 768
+    cond_vec_size: int = 768
+    depth_double: int = 17
+    txt_preamble_depth: int = 2
+    num_heads: int = 12
+    head_dim: int = 64
+    mlp_ratio: float = 2.6667
+    pca_channels: int = 128
+    prompt_length: int = 256
+    n_T: int = 100
+    prediction: str = "x"
+    sampler: str = "euler"
+    cfg_channels: int = 3
+    cfg_interval: tuple = (0.0, 1.0)
+    llm: str = "google/flan-t5-large"
+
+
+class MMJiT(nn.Module):
+    """The `net` sub-tree. State-dict matches the `model.net.<...>` checkpoint keys."""
+
+    def __init__(self, cfg: MMJiTConfig):
+        super().__init__()
+        self.cfg = cfg
+        self.latent_img_size = cfg.image_size // cfg.patch_size
+        self.img_embedder = BottleneckPatchEmbed(
+            cfg.image_size, cfg.patch_size, cfg.in_channels, cfg.pca_channels, cfg.hidden_size
+        )
+        self.txt_embedder = nn.Linear(cfg.txt_input_size, cfg.txt_hidden_size, bias=False)
+        self.mask_token = nn.Parameter(torch.zeros(1, 1, cfg.txt_input_size))
+        self.t_embedder = TimestepEmbedder(cfg.cond_vec_size)
+        self.pooled_embedder = nn.Linear(cfg.txt_input_size, cfg.cond_vec_size, bias=False)
+        self.txt_preamble_blocks = nn.ModuleList(
+            [
+                PlainTextTransformerBlock(cfg.txt_hidden_size, cfg.num_heads, cfg.head_dim, cfg.mlp_ratio)
+                for _ in range(cfg.txt_preamble_depth)
+            ]
+        )
+        self.double_blocks = nn.ModuleList(
+            [
+                DoubleStreamDiTBlock(
+                    cfg.hidden_size, cfg.txt_hidden_size, cfg.num_heads, cfg.head_dim, cfg.mlp_ratio
+                )
+                for _ in range(cfg.depth_double)
+            ]
+        )
+        self.final_layer = FinalLayer(cfg.hidden_size, cfg.patch_size, cfg.in_channels)
+
+    def unpatchify(self, x: torch.Tensor) -> torch.Tensor:
+        b = x.shape[0]
+        p = self.cfg.patch_size
+        c = self.cfg.in_channels
+        h = w = int(math.sqrt(x.shape[1]))
+        x = x.reshape(b, h, w, p, p, c)
+        x = torch.einsum("nhwpqc->nchpwq", x)
+        return x.reshape(b, c, h * p, w * p)
+
+    def forward(self, img: torch.Tensor, t: torch.Tensor, context: torch.Tensor, attn_mask: torch.Tensor) -> torch.Tensor:
+        if img.ndim == 4 and img.shape[1] != self.cfg.in_channels:
+            img = img.permute(0, 3, 1, 2)
+        attn_mask = attn_mask.to(device=context.device)
+        context = torch.where(attn_mask[:, :, None] > 0.5, context, self.mask_token.to(dtype=context.dtype))
+        x = self.img_embedder(img)
+        pos = get_2d_sincos_pos_embed(self.cfg.hidden_size, self.latent_img_size, x.device, x.dtype)
+        x = x + pos[None]
+        t_vec = self.t_embedder(t)
+        txt = self.txt_embedder(context.to(dtype=self.txt_embedder.weight.dtype))
+        pooled_text = context.mean(dim=1)
+        vec = t_vec + self.pooled_embedder(pooled_text.to(dtype=self.pooled_embedder.weight.dtype))
+        for block in self.txt_preamble_blocks:
+            txt = block(txt)
+        for block in self.double_blocks:
+            x, txt = block(x, txt, vec)
+        combined = torch.cat([txt, x], dim=1)
+        out = self.final_layer(combined, vec)
+        img_out = out[:, txt.shape[1]:, :]
+        return self.unpatchify(img_out)
+
+
+# ---------------------------------------------------------------------------
+# Wrapper (mirrors upstream DiffusionModel): flow-matching + Euler/CFG sampler
+# ---------------------------------------------------------------------------
+class MiniT2I(nn.Module):
+    """Holds `self.net = MMJiT(cfg)` under the same name the checkpoint uses.
+
+    `forward` returns the x0 prediction (`predict='x'`), so the repo's existing
+    x→velocity machinery applies. `sample()` runs the native Euler + CFG loop.
+    """
+
+    NOISE_SCALE = 2.0
+    T_EPS = 0.001
+
+    def __init__(self, cfg: MMJiTConfig):
+        super().__init__()
+        self.cfg = cfg
+        self.net = MMJiT(cfg)
+
+    @property
+    def image_size(self) -> int:
+        return self.cfg.image_size
+
+    def forward(self, img: torch.Tensor, t: torch.Tensor, context: torch.Tensor, attn_mask: torch.Tensor) -> torch.Tensor:
+        return self.net(img, t, context, attn_mask)
+
+    def pred_velocity(self, x, t, context, mask):
+        x0 = self.net(x, t, context, mask)
+        return (x0 - x) / torch.clamp(1 - t[:, None, None, None], min=self.T_EPS)
+
+    def cfg_velocity(self, x, t, context, mask, cfg_scale: float):
+        b = x.shape[0]
+        xx = torch.cat([x, x], dim=0)
+        tt = torch.cat([t, t], dim=0)
+        cc = torch.cat([context, context], dim=0)
+        mm = torch.cat([mask, torch.zeros_like(mask)], dim=0)   # null branch = all-masked
+        out = self.pred_velocity(xx, tt, cc, mm)
+        cond, uncond = out[:b], out[b:]
+        lo, hi = self.cfg.cfg_interval
+        use_cfg = ((t >= lo) & (t <= hi)).to(out.dtype)[:, None, None, None]
+        scale = torch.where(use_cfg > 0, out.new_tensor(cfg_scale), out.new_tensor(1.0))
+        return uncond + (cond - uncond) * scale
+
+    @torch.no_grad()
+    def sample(self, context, mask, cfg_scale: float = 2.5, steps: Optional[int] = None,
+               generator=None, progress: bool = False) -> torch.Tensor:
+        """Euler flow-matching sampler. Returns pixels in [-1, 1]."""
+        steps = int(steps or self.cfg.n_T)
+        b = context.shape[0]
+        device = context.device
+        dtype = next(self.parameters()).dtype
+        x = torch.randn(
+            b, self.cfg.in_channels, self.cfg.image_size, self.cfg.image_size,
+            generator=generator, device=device, dtype=dtype,
+        ) * self.NOISE_SCALE
+        timesteps = torch.linspace(0.0, 1.0, steps + 1, device=device, dtype=dtype)
+        iterator = range(steps)
+        if progress:
+            from tqdm.auto import tqdm
+            iterator = tqdm(iterator)
+        context = context.to(dtype)
+        mask = mask.to(dtype)
+        for i in iterator:
+            t_cur = timesteps[i].expand(b)
+            t_next = timesteps[i + 1].expand(b)
+            if cfg_scale != 1.0:
+                v = self.cfg_velocity(x, t_cur, context, mask, cfg_scale)
+            else:
+                v = self.pred_velocity(x, t_cur, context, mask)
+            x = x + (t_next - t_cur)[:, None, None, None] * v
+        return x.clamp(-1, 1)
+
+
+# ---------------------------------------------------------------------------
+# Registry + pretrained loader
+# ---------------------------------------------------------------------------
+# cfg_vec_size / txt_hidden_size follow hidden_size for both released variants.
+MiniT2I_configs = {
+    "MiniT2I-B/16": MMJiTConfig(
+        hidden_size=768, txt_hidden_size=768, cond_vec_size=768,
+        depth_double=17, num_heads=12, head_dim=64, mlp_ratio=2.6667,
+    ),
+    "MiniT2I-L/16": MMJiTConfig(
+        hidden_size=1248, txt_hidden_size=1248, cond_vec_size=1248,
+        depth_double=23, num_heads=24, head_dim=52, mlp_ratio=2.7051282051282053,
+    ),
+}
+
+MiniT2I_HF_SUBFOLDER = {
+    "MiniT2I-B/16": "minit2i-b-16",
+    "MiniT2I-L/16": "minit2i-l-16",
+}
+
+
+def _build(variant: str):
+    base_cfg = MiniT2I_configs[variant]
+
+    def ctor(**overrides) -> MiniT2I:
+        # Accept common repo kwargs (image_size, in_channels, patch_size, ...) as
+        # dataclass field overrides; ignore anything not in MMJiTConfig.
+        fields = {f for f in base_cfg.__dataclass_fields__}
+        cfg = replace(base_cfg, **{k: v for k, v in overrides.items() if k in fields})
+        return MiniT2I(cfg)
+
+    return ctor
+
+
+MiniT2I_models = {name: _build(name) for name in MiniT2I_configs}
+
+
+def default_cfg_scale(variant: str) -> float:
+    """Recommended classifier-free guidance scale per released variant."""
+    return 6.0 if "L/" in variant else 2.5
+
+
+def _cache_dir(cache_dir: Optional[str]) -> Optional[str]:
+    # L/16 is 3.65 GB — keep off the boot nvme by default (see CLAUDE.md).
+    return cache_dir or os.environ.get("MINIT2I_CACHE") or None
+
+
+@torch.no_grad()
+def load_pretrained(
+    variant: str,
+    model: Optional[MiniT2I] = None,
+    *,
+    cache_dir: Optional[str] = None,
+    device: str | torch.device = "cpu",
+    dtype: Optional[torch.dtype] = None,
+) -> MiniT2I:
+    """Download the released diffusers safetensors and load them into `model`.
+
+    Checkpoint keys are prefixed `model.net.`; we strip that and load into
+    `model.net` with strict=True (the unused t/pooled embedders exist in the
+    checkpoint, so strict passes cleanly).
+    """
+    from huggingface_hub import hf_hub_download
+    from safetensors.torch import load_file
+
+    if variant not in MiniT2I_HF_SUBFOLDER:
+        raise KeyError(f"unknown MiniT2I variant {variant!r}; choices: {list(MiniT2I_HF_SUBFOLDER)}")
+    if model is None:
+        model = MiniT2I_models[variant]()
+
+    sub = MiniT2I_HF_SUBFOLDER[variant]
+    path = hf_hub_download(
+        MINIT2I_HF_REPO,
+        f"{sub}/transformer/diffusion_pytorch_model.safetensors",
+        cache_dir=_cache_dir(cache_dir),
+    )
+    raw = load_file(path)
+    prefix = "model.net."
+    net_sd = {}
+    for k, v in raw.items():
+        if k.startswith(prefix):
+            net_sd[k[len(prefix):]] = v
+        elif k.startswith("net."):
+            net_sd[k[len("net."):]] = v
+        else:
+            net_sd[k] = v
+    missing, unexpected = model.net.load_state_dict(net_sd, strict=True)
+    assert not missing and not unexpected, f"key mismatch: missing={missing} unexpected={unexpected}"
+
+    model = model.to(device)
+    if dtype is not None:
+        model = model.to(dtype)
+    return model
+
+
+@torch.no_grad()
+def load_pretrained_repatched(
+    variant: str,
+    model: MiniT2I,
+    *,
+    cache_dir: Optional[str] = None,
+    device: str | torch.device = "cpu",
+    dtype: Optional[torch.dtype] = None,
+) -> tuple[MiniT2I, list]:
+    """Load the released $512$/p16 weights into a RE-PATCHED model (e.g. p=32 @1024).
+
+    Both p16@512 and p32@1024 tokenize to a $32{\\times}32$ grid, so the entire
+    transformer trunk (blocks, text path, rope, embedders) transfers by shape.
+    Only the patch conv (`img_embedder.proj1`) and the output head
+    (`final_layer.linear`) depend on patch size and are left at init --- these must
+    be finetuned. Returns (model, list of re-initialised keys)."""
+    from huggingface_hub import hf_hub_download
+    from safetensors.torch import load_file
+    sub = MiniT2I_HF_SUBFOLDER[variant]
+    path = hf_hub_download(
+        MINIT2I_HF_REPO, f"{sub}/transformer/diffusion_pytorch_model.safetensors",
+        cache_dir=_cache_dir(cache_dir))
+    raw = load_file(path)
+    net_sd = {}
+    for k, v in raw.items():
+        for pre in ("model.net.", "net."):
+            if k.startswith(pre):
+                k = k[len(pre):]; break
+        net_sd[k] = v
+    tgt = model.net.state_dict()
+    keep = {k: v for k, v in net_sd.items() if k in tgt and tgt[k].shape == v.shape}
+    reinit = sorted(set(tgt) - set(keep))
+    model.net.load_state_dict(keep, strict=False)
+    if int(os.environ.get("RANK", "0")) == 0:
+        print(f"[repatch] transferred {len(keep)}/{len(tgt)} trunk tensors; "
+              f"re-init {len(reinit)}: {reinit}", flush=True)
+    model = model.to(device)
+    if dtype is not None:
+        model = model.to(dtype)
+    return model, reinit

rti/minit2i_text.py

@@ -0,0 +1,63 @@
+"""FLAN-T5-Large prompt encoder used by MiniT2I."""
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+
+_T5_NAME = "google/flan-t5-large"
+_SINGLETON: dict[tuple, "T5PromptEncoder"] = {}
+
+
+class T5PromptEncoder:
+    def __init__(
+        self,
+        model_name: str = _T5_NAME,
+        device: str | torch.device = "cuda",
+        dtype: torch.dtype = torch.float32,
+        max_len: int = 256,
+        cache_dir: Optional[str] = None,
+    ):
+        from transformers import AutoTokenizer, T5EncoderModel
+
+        self.max_len = max_len
+        self.device = torch.device(device)
+        self.dtype = dtype
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name, cache_dir=cache_dir)
+        self.model = (
+            T5EncoderModel.from_pretrained(model_name, cache_dir=cache_dir)
+            .to(self.device, dtype=dtype)
+            .eval()
+        )
+        for p in self.model.parameters():
+            p.requires_grad_(False)
+
+    @torch.no_grad()
+    def encode(self, prompts: list[str], max_len: Optional[int] = None):
+        max_len = int(max_len or self.max_len)
+        tok = self.tokenizer(
+            list(prompts),
+            padding="max_length",
+            truncation=True,
+            max_length=max_len,
+            return_tensors="pt",
+        )
+        input_ids = tok.input_ids.to(self.device)
+        attn_mask = tok.attention_mask.to(self.device)
+        out = self.model(input_ids=input_ids, attention_mask=attn_mask)
+        context = out.last_hidden_state.to(self.dtype)
+        return context, attn_mask.to(self.dtype)
+
+
+def get_prompt_encoder(
+    model_name: str = _T5_NAME,
+    device: str | torch.device = "cuda",
+    dtype: torch.dtype = torch.float32,
+    max_len: int = 256,
+    cache_dir: Optional[str] = None,
+) -> T5PromptEncoder:
+    """Cached accessor so repeated calls reuse a single loaded T5 encoder."""
+    key = (model_name, str(device), dtype, max_len, cache_dir)
+    if key not in _SINGLETON:
+        _SINGLETON[key] = T5PromptEncoder(model_name, device, dtype, max_len, cache_dir)
+    return _SINGLETON[key]

rti/model.py

@@ -0,0 +1,410 @@
+"""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),
+}

rti/sampling.py

@@ -0,0 +1,42 @@
+"""Euler flow-matching sampler with interval classifier-free guidance."""
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+
+
+@torch.no_grad()
+def sample(model, ctx, mask, *, budget: int, steps: int = 100, cfg_scale: float = 5.0,
+           image_size: int = 512, noise_scale: float = 2.0, t_eps: float = 0.001,
+           cfg_interval: tuple[float, float] = (0.0, 1.0), generator=None,
+           device: Optional[str] = None, dtype=torch.bfloat16, progress: bool = False):
+    """model: an RTIMiniT2I. ctx, mask: encoded prompts [B,L,D] and [B,L].
+    budget: region-token budget R; R >= N runs dense.
+    Returns images in [0,1], [B,3,image_size,image_size], on CPU."""
+    device = device or next(model.parameters()).device
+    B = ctx.shape[0]
+    ctx = ctx.to(device, dtype)
+    mask = mask.to(device, dtype)
+    null = torch.zeros_like(mask)
+    x = torch.randn(B, 3, image_size, image_size, generator=generator,
+                    device=device, dtype=dtype) * noise_scale
+    ts = torch.linspace(0.0, 1.0, steps + 1, device=device, dtype=dtype)
+    lo, hi = cfg_interval
+    it = range(steps)
+    if progress:
+        from tqdm.auto import tqdm
+        it = tqdm(it)
+    with torch.autocast(device.type if hasattr(device, "type") else "cuda", dtype=dtype):
+        for i in it:
+            t0 = float(ts[i])
+            tt = torch.full((2 * B,), t0, device=device, dtype=x.dtype)
+            x2 = torch.cat([x, x])
+            x0 = model.forward_xpred(x2, tt, torch.cat([ctx, ctx]),
+                                     torch.cat([mask, null]), budget=budget)
+            v = (x0 - x2) / max(1.0 - t0, t_eps)
+            v_cond, v_uncond = v.chunk(2, 0)
+            s = cfg_scale if lo <= t0 <= hi else 1.0
+            v = v_uncond + (v_cond - v_uncond) * s
+            x = x + (ts[i + 1] - ts[i]) * v
+    return ((x.float().clamp(-1, 1) + 1) / 2).cpu()