inference.py

"""InfoVis inference: one slide image from an instruction and up to four images.

Single case:
    python inference.py --transformer checkpoints/InfoVis --prompt "A title slide ..." \
        [--images ref.png asset1.png] --width 1376 --height 768 --out slide.png

VisSlideBench (writes <out_dir>/<task>/<id>.png; shard over GPUs with --shard/--num_shards):
    python inference.py --transformer checkpoints/InfoVis --bench_root VisSlideBench \
        --task multi_image --out_dir outputs/infovis

Sampling follows the paper: 50 Euler steps, CFG 4.0 with an empty negative
prompt, 4,096 text tokens padded to full length and masked in attention, BF16,
seed = 2233 + case index. Output size follows the evaluation runs of the
released model: 1280x720 for t2i, 1376x768 for template and multi_image, and
the source slide size for edit (the record's width/height, floored to
multiples of 16).
"""

import argparse
import json
import os

import numpy as np
import torch
from diffusers import Flux2KleinPipeline, Flux2Transformer2DModel
from diffusers.pipelines.flux2.pipeline_flux2_klein import compute_empirical_mu, retrieve_timesteps
from PIL import Image, ImageOps

from infovis.modeling import encode_text, tokenize, vae_decode


def load_pipeline(base_model, transformer=None, device="cuda"):
    pipe = Flux2KleinPipeline.from_pretrained(base_model, torch_dtype=torch.bfloat16)
    if transformer:
        pipe.register_modules(transformer=Flux2Transformer2DModel.from_pretrained(transformer, torch_dtype=torch.bfloat16))
    # Casting the whole pipeline also moves fp32 buffers (Qwen3 rotary inv_freq) to BF16, as in training.
    return pipe.to(device, dtype=torch.bfloat16)


def prepare_condition_images(pipe, images):
    """Cap each image at 1024x1024 pixels and center-crop to multiples of 32."""
    out = []
    for img in images:
        pipe.image_processor.check_image_input(img)
        if img.width * img.height > 1024 * 1024:
            img = pipe.image_processor._resize_to_target_area(img, 1024 * 1024)
        m = pipe.vae_scale_factor * 2
        out.append(pipe.image_processor.preprocess(img, height=img.height // m * m, width=img.width // m * m,
                                                   resize_mode="crop"))
    return out


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


@torch.no_grad()
def generate(pipe, prompt, images=None, width=1376, height=768, steps=50, guidance=4.0, seed=2233,
             max_tokens=4096, text_layer_offset=0):
    device, tr = pipe.device, pipe.transformer
    embeds, txt_ids, txt_mask = _embed(pipe, prompt, max_tokens, text_layer_offset)
    if guidance > 1:
        neg_embeds, neg_ids, neg_mask = _embed(pipe, "", max_tokens, text_layer_offset)

    generator = torch.Generator(device=device).manual_seed(seed)
    latents, latent_ids = pipe.prepare_latents(1, tr.config.in_channels // 4, height, width, embeds.dtype,
                                               device, generator)
    cond, cond_ids = None, None
    if images:
        cond, cond_ids = pipe.prepare_image_latents(prepare_condition_images(pipe, images), 1, generator,
                                                    device, pipe.vae.dtype)

    sigmas = np.linspace(1.0, 1 / steps, steps)
    mu = compute_empirical_mu(image_seq_len=latents.shape[1], num_steps=steps)
    timesteps, _ = retrieve_timesteps(pipe.scheduler, steps, device, sigmas=sigmas, mu=mu)
    pipe.scheduler.set_begin_index(0)

    n_out = latents.shape[1]
    img_ids = latent_ids if cond is None else torch.cat([latent_ids, cond_ids], 1)

    def predict(x, e, ids_txt, mask_txt, t):
        image_mask = torch.ones(x.shape[:2], dtype=torch.bool, device=device)
        attn = torch.cat([mask_txt, image_mask], 1)[:, None, None, :]
        return tr(hidden_states=x, timestep=t / 1000, guidance=None, encoder_hidden_states=e, txt_ids=ids_txt,
                  img_ids=img_ids, joint_attention_kwargs={"attention_mask": attn}, return_dict=False)[0][:, :n_out]

    for t in pipe.progress_bar(timesteps):
        x = latents if cond is None else torch.cat([latents, cond], 1)
        x, tt = x.to(tr.dtype), t.expand(1).to(latents.dtype)
        v = predict(x, embeds, txt_ids, txt_mask, tt)
        if guidance > 1:
            v_neg = predict(x, neg_embeds, neg_ids, neg_mask, tt)
            v = v_neg + guidance * (v - v_neg)
        latents = pipe.scheduler.step(v, t, latents, return_dict=False)[0]

    latents = Flux2KleinPipeline._unpack_latents_with_ids(latents, latent_ids)
    return pipe.image_processor.postprocess(vae_decode(pipe.vae, latents), output_type="pil")[0]


def open_image(path):
    return ImageOps.exif_transpose(Image.open(path).convert("RGB"))


def run_bench(pipe, args):
    with open(os.path.join(args.bench_root, "data", f"{args.task}.jsonl")) as f:
        records = [json.loads(line) for line in f]
    records = records[args.shard::args.num_shards]
    out_dir = os.path.join(args.out_dir, args.task)
    os.makedirs(out_dir, exist_ok=True)
    for r in records:
        path = os.path.join(out_dir, f"{r['id']}.png")
        if os.path.exists(path):
            continue
        inputs = [r["source_image"]] if args.task == "edit" else r["input_images"]
        size = {"t2i": (1280, 720), "edit": (r["width"], r["height"])}.get(args.task, (1376, 768))
        width, height = (args.width, args.height) if args.width else size
        img = generate(pipe, r["instruction"], [open_image(os.path.join(args.bench_root, p)) for p in inputs],
                       width=width, height=height, steps=args.steps, guidance=args.guidance,
                       seed=args.seed + r["index"], max_tokens=args.max_tokens,
                       text_layer_offset=args.text_layer_offset)
        img.save(path)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--base_model", default="black-forest-labs/FLUX.2-klein-base-9B")
    ap.add_argument("--transformer", default=None, help="InfoVis transformer folder (diffusers format)")
    ap.add_argument("--prompt")
    ap.add_argument("--images", nargs="*", default=[])
    ap.add_argument("--out", default="slide.png")
    ap.add_argument("--bench_root")
    ap.add_argument("--task", choices=["t2i", "template", "multi_image", "edit"])
    ap.add_argument("--out_dir", default="outputs")
    ap.add_argument("--shard", type=int, default=0)
    ap.add_argument("--num_shards", type=int, default=1)
    ap.add_argument("--width", type=int, default=None, help="override output width")
    ap.add_argument("--height", type=int, default=None, help="override output height")
    ap.add_argument("--steps", type=int, default=50)
    ap.add_argument("--guidance", type=float, default=4.0)
    ap.add_argument("--seed", type=int, default=2233)
    ap.add_argument("--max_tokens", type=int, default=4096)
    ap.add_argument("--device", default="cuda")
    ap.add_argument("--text_layer_offset", type=int, default=0, choices=[0, 1],
                    help="0: Qwen3 hidden_states[9,18,27] (used for the paper's evaluation); "
                         "1: outputs of decoder layers 9/18/27 (used during training)")
    args = ap.parse_args()

    pipe = load_pipeline(args.base_model, args.transformer, args.device)
    if args.bench_root:
        run_bench(pipe, args)
    else:
        img = generate(pipe, args.prompt, [open_image(p) for p in args.images], width=args.width or 1376,
                       height=args.height or 768, steps=args.steps, guidance=args.guidance, seed=args.seed,
                       max_tokens=args.max_tokens, text_layer_offset=args.text_layer_offset)
        img.save(args.out)


if __name__ == "__main__":
    main()