app.py

"""InfoVis demo on Hugging Face Spaces (ZeroGPU).

One model for text-to-slide, template-conditioned and multi-image grounded generation, and slide
editing. Sampling is the evaluation code of the paper (inference.py): Euler steps, CFG 4.0 with an
empty negative prompt, BF16. The only change is that the prompt is not padded to 4,096 tokens: padded
tokens are masked out of attention, so the image is the same up to floating-point rounding, and
generation is faster.
"""

import json
import random
import time

import gradio as gr
import spaces
import torch
from diffusers import AutoencoderKLFlux2, FlowMatchEulerDiscreteScheduler, Flux2KleinPipeline, Flux2Transformer2DModel
from PIL import Image, ImageOps
from transformers import Qwen2TokenizerFast, Qwen3ForCausalLM

import inference
from infovis.modeling import encode_text, tokenize

BASE = "black-forest-labs/FLUX.2-klein-base-9B"
MODEL = "InfoVis/InfoVis"
DTYPE = torch.bfloat16

# Text encoder, VAE, tokenizer and scheduler come from the base model; the transformer is InfoVis.
pipe = Flux2KleinPipeline(
    scheduler=FlowMatchEulerDiscreteScheduler.from_pretrained(BASE, subfolder="scheduler"),
    vae=AutoencoderKLFlux2.from_pretrained(BASE, subfolder="vae", torch_dtype=DTYPE),
    text_encoder=Qwen3ForCausalLM.from_pretrained(BASE, subfolder="text_encoder", torch_dtype=DTYPE),
    tokenizer=Qwen2TokenizerFast.from_pretrained(BASE, subfolder="tokenizer"),
    transformer=Flux2Transformer2DModel.from_pretrained(MODEL, torch_dtype=DTYPE),
)
# Casting the whole pipeline also moves the Qwen3 rotary buffer to BF16, as in training and evaluation.
pipe.to("cuda", dtype=DTYPE)


def _embed_unpadded(pipe, prompt, max_tokens, offset):
    t = tokenize(pipe.tokenizer, [prompt], max_tokens, padding="longest")
    return encode_text(pipe.text_encoder, t["input_ids"].to(pipe.device), t["attention_mask"].to(pipe.device),
                       offset=offset)


inference._embed = _embed_unpadded

T2I, TEMPLATE, MULTI, EDIT = "Text-to-slide", "Template-conditioned", "Multi-image grounded", "Slide editing"
MAX_OUT_PIXELS = 1376 * 768


def output_size(task, images, aspect):
    """Evaluation sizes: 1280x720 (text-to-slide), 1376x768 (template, multi-image), source size (edit)."""
    if task == EDIT:
        w, h = images[0].size
        r = min(1.0, (MAX_OUT_PIXELS / (w * h)) ** 0.5)
        return max(256, int(w * r) // 16 * 16), max(256, int(h * r) // 16 * 16)
    if aspect == "4:3":
        return 1184, 880
    return (1280, 720) if task == T2I else (1376, 768)


def _tokens(w, h):
    return (w // 16) * (h // 16)


def _input_tokens(img):
    w, h = img.size
    if w * h > 1024 * 1024:  # inputs are capped at 1024x1024 pixels and cropped to multiples of 32
        r = (1024 * 1024 / (w * h)) ** 0.5
        w, h = int(w * r), int(h * r)
    return _tokens(w // 32 * 32, h // 32 * 32)


def _duration(instruction, images, width, height, steps, guidance, seed):
    # Measured on ZeroGPU: about 1.1 s per step for a 1280x720 slide without input images (3,600 image tokens).
    # Long prompts (e.g. style_json) add text tokens to the sequence; about 3.5 characters per token.
    tokens = _tokens(width, height) + sum(_input_tokens(i) for i in images) + len(instruction) / 3.5
    estimate = 10 + steps * 1.1 * (tokens / 3600) ** 1.15
    return int(min(1.4 * estimate + 10, 240))


@spaces.GPU(duration=_duration)
def _generate(instruction, images, width, height, steps, guidance, seed):
    t0 = time.time()
    image = inference.generate(pipe, instruction, images or None, width=width, height=height, steps=steps,
                               guidance=guidance, seed=seed)
    print(f"[generate] {len(images)} input(s), {width}x{height}, {steps} steps: {time.time() - t0:.1f}s", flush=True)
    return image


def collect_images(task, img1, img2, img3, img4):
    if task == T2I:
        return []
    candidates = [img1] if task in (TEMPLATE, EDIT) else [img1, img2, img3, img4]
    return [ImageOps.exif_transpose(img.convert("RGB")) for img in candidates if img is not None]


def run(task, instruction, img1, img2, img3, img4, steps, guidance, seed, randomize, aspect,
        progress=gr.Progress(track_tqdm=True)):
    instruction = (instruction or "").strip()
    if not instruction:
        raise gr.Error("Please enter an instruction.")
    images = collect_images(task, img1, img2, img3, img4)
    if task == TEMPLATE and not images:
        raise gr.Error("Upload a reference slide.")
    if task == MULTI and not images:
        raise gr.Error("Upload at least one asset image.")
    if task == EDIT and not images:
        raise gr.Error("Upload the slide to edit.")
    if randomize:
        seed = random.randint(0, 2**31 - 1)
    width, height = output_size(task, images, aspect)
    image = _generate(instruction, images, width, height, int(steps), float(guidance), int(seed))
    return gr.update(value=image, label="Result"), int(seed)


def update_inputs(task):
    label = {TEMPLATE: "Reference slide (style)", MULTI: "Asset 1", EDIT: "Slide to edit"}.get(task, "Image")
    more = gr.update(visible=task == MULTI)
    return gr.update(visible=task != T2I, label=label), more, more, more, gr.update(visible=task != EDIT)


# Good cases from the appendix of the paper, two per task.
CASES = json.load(open("examples/cases.json"))
EXAMPLES = [[c["task"], c["instruction"]] + (c["images"] + [None] * 4)[:4] for c in CASES]
EXAMPLE_LABELS = [c["label"] for c in CASES]
PAPER_OUTPUTS = {c["instruction"]: c["output"] for c in CASES}
GALLERY = [(c["output"], c["label"]) for c in CASES]


def show_paper_output(task, instruction, img1, img2, img3, img4):
    """Clicking an example shows the output from the paper; no GPU is used."""
    path = PAPER_OUTPUTS.get((instruction or "").strip())
    return gr.update(value=path, label="Output from the paper (click Generate to run InfoVis)"), None


HEADER = """
# InfoVis: Unified Slide Image Generation and Editing

One 9B model for text-to-slide generation, template-conditioned generation, multi-image grounded generation, and
instruction-based slide editing.

[📄 Paper](https://arxiv.org/abs/2610.16784) | [🌐 Project page](https://vectorspacelab.github.io/InfoVis/) | [💻 Code](https://github.com/VectorSpaceLab/InfoVis) |
[🤗 Model](https://huggingface.co/InfoVis/InfoVis) | [📊 VisSlideBench](https://huggingface.co/datasets/InfoVis/VisSlideBench) |
[📦 Training data](https://huggingface.co/datasets/InfoVis/Data)
"""

TERMS = """
**Terms of use.** By using this demo, you agree to use it only for non-commercial research and to follow the
[FLUX Non-Commercial License](https://huggingface.co/InfoVis/InfoVis/blob/main/LICENSE.md). Do not use it to create
sexual content, content involving minors, non-consensual or deceptive images of real people, or any violent, hateful,
harassing, or otherwise unlawful content. All outputs are AI-generated, and you are responsible for the content you
create and how you use it.
"""

FOOTER = """
**Tips.** InfoVis follows detailed instructions: describe the title, the text, the layout, colors, and visual
elements. Text in quotes is rendered as written. For template-conditioned generation, the reference slide sets the
style; for multi-image grounded generation, the uploaded images are placed on the slide; for editing, describe the
change to make.

**Settings.** Output size follows the evaluation setup of the paper: 1280x720 for text-to-slide, 1376x768 for
template-conditioned and multi-image grounded generation, and the source size (up to about 1 MP) for editing. The
demo defaults to 30 sampling steps to save GPU time; the paper uses 50.

**License.** InfoVis is a derivative of FLUX.2-klein-base-9B and is released under the
[FLUX Non-Commercial License](https://huggingface.co/InfoVis/InfoVis/blob/main/LICENSE.md). This demo is for
non-commercial research use.
"""

with gr.Blocks(title="InfoVis") as demo:
    gr.Markdown(HEADER)
    gr.Markdown(TERMS)
    with gr.Row():
        with gr.Column(scale=5):
            task = gr.Radio([T2I, TEMPLATE, MULTI, EDIT], value=T2I, label="Task")
            instruction = gr.Textbox(label="Instruction", lines=6,
                                     placeholder="Describe the slide to generate, or the edit to make.")
            img1 = gr.Image(type="pil", label="Image", visible=False, height=260)
            with gr.Row():
                img2 = gr.Image(type="pil", label="Asset 2", visible=False, height=180)
                img3 = gr.Image(type="pil", label="Asset 3", visible=False, height=180)
                img4 = gr.Image(type="pil", label="Asset 4", visible=False, height=180)
            aspect = gr.Radio(["16:9", "4:3"], value="16:9", label="Aspect ratio")
            with gr.Accordion("Advanced settings", open=False):
                steps = gr.Slider(10, 50, value=30, step=1, label="Sampling steps (the paper uses 50)")
                guidance = gr.Slider(1.0, 8.0, value=4.0, step=0.5, label="Guidance scale")
                with gr.Row():
                    seed = gr.Number(value=2233, precision=0, label="Seed")
                    randomize = gr.Checkbox(value=True, label="Random seed")
            run_button = gr.Button("Generate", variant="primary")
        with gr.Column(scale=6):
            output = gr.Image(type="pil", format="png", label="Result", height=480)
            used_seed = gr.Number(label="Seed used", precision=0, interactive=False)
    gr.Examples(examples=EXAMPLES, inputs=[task, instruction, img1, img2, img3, img4], cache_examples=False,
                example_labels=EXAMPLE_LABELS, label="Examples (good cases from the paper)",
                fn=show_paper_output, outputs=[output, used_seed], run_on_click=True)
    gr.Gallery(value=GALLERY, label="Good cases from the paper (InfoVis outputs)", columns=4, height="auto",
               object_fit="contain", allow_preview=True, interactive=False)
    gr.Markdown(FOOTER)

    task.change(update_inputs, task, [img1, img2, img3, img4, aspect])
    run_button.click(run, [task, instruction, img1, img2, img3, img4, steps, guidance, seed, randomize, aspect],
                     [output, used_seed])

demo.queue(max_size=20).launch()