Spaces:

hugging-apps
/

rti-pixel-diffusion

Running on Zero

Upload folder using huggingface_hub

2c92732 verified 4 days ago

8.93 kB

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

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