app.py

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


# Measured on the live Space (100 steps, cfg 5.0): B/16 R=512 ≈ 5.7s, R=256 ≈
# 5.2s, L/16 dense ≈ 7.2s. The UI caps steps at 150 (1.5×), so 60s covers the
# worst case plus ZeroGPU cold-start weight streaming with margin.
@spaces.GPU(duration=60)
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)