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