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