app.py

import os

# High-res pixel-space ops + large attention activations on a 9.28B DiT -> transient alloc spikes.
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
# outlines_core ships an @torch.compile bitmask kernel dynamo can't trace -> noisy WON'T CONVERT
# spam on every local upsample. We never torch.compile at runtime, so disable dynamo.
os.environ.setdefault("TORCHDYNAMO_DISABLE", "1")

import json
import math
import random
import time

import gradio as gr
import spaces
import torch

from diffusers import Ideogram4Pipeline, Ideogram4Transformer2DModel

try:
    from diffusers import Ideogram4PromptEnhancerHead
    _HAS_HEAD = True
except Exception:  # pragma: no cover - older diffusers without the enhancer head
    _HAS_HEAD = False

# Runtime shim (matches the official Ideogram 4 Space): cu130-era bitsandbytes returns Params4bit.shape
# as a plain tuple, but diffusers' check_quantized_param_shape calls .numel() on it. math.prod handles
# both, so this is a no-op once diffusers/bnb fix it upstream. Needed to load the nf4-quantized text
# encoder from the shared components repo.
from diffusers.quantizers.bitsandbytes.bnb_quantizer import BnB4BitDiffusersQuantizer  # noqa: E402


def _check_quantized_param_shape(self, param_name, current_param, loaded_param):
    n = math.prod(tuple(current_param.shape))
    inferred_shape = (n,) if "bias" in param_name else ((n + 1) // 2, 1)
    if tuple(loaded_param.shape) != tuple(inferred_shape):
        raise ValueError(
            f"Expected flattened shape of {param_name} to be {inferred_shape}, got {tuple(loaded_param.shape)}."
        )
    return True


BnB4BitDiffusersQuantizer.check_quantized_param_shape = _check_quantized_param_shape

# --- Config --------------------------------------------------------------------------------------------
# fal's Fast checkpoint is a TRANSFORMER-ONLY release: a QAD/CFG-distilled Ideogram 4 transformer that
# folds guidance into a single conditional branch (20-step, no runtime CFG). The shared inference
# components (VAE, Qwen3-VL text encoder, tokenizer, scheduler) come from Ideogram AI's public gated
# repo — neither of ITS transformers is used.
FAST_REPO = "fal/ideogram-v4-fast"
COMPONENTS_REPO = "ideogram-ai/ideogram-4-nf4-diffusers"
COMPONENTS_REVISION = "1874bc70267ba2c823a7239e1d70dd308c8d64dc"
LM_HEAD_REPO = "diffusers/qwen3-vl-8b-instruct-lm-head"
HF_TOKEN = os.environ.get("HF_TOKEN")  # both repos are gated -> read them with the Space secret
MAX_SEED = 2**31 - 1

# Fast schedule at 1024 (from the fal model card): 20 steps, mu=0.0, std=1.75, no runtime CFG.
FAST_STEPS = 20
FAST_MU = 0.0
FAST_STD = 1.75

# Aspect-ratio presets -> (width, height). All multiples of 64.
ASPECT_RATIOS = {
    "1:1 · 1024×1024": (1024, 1024),
    "3:2 · 1216×832": (1216, 832),
    "2:3 · 832×1216": (832, 1216),
    "16:9 · 1344×768": (1344, 768),
    "9:16 · 768×1344": (768, 1344),
}
DEFAULT_RATIO = "1:1 · 1024×1024"

# --- Load the pipeline ---------------------------------------------------------------------------------
# The Fast transformer stores its (pre-pack) weights in a loadable bf16-serialized form. Intended quality
# is on an FP4 execution path; here we run bf16 on ZeroGPU (the model card notes bf16 may be lower quality
# than native FP4). Good enough for an interactive demo.
t0 = time.perf_counter()
fast_transformer = Ideogram4Transformer2DModel.from_pretrained(
    FAST_REPO,
    subfolder="transformer",
    torch_dtype=torch.bfloat16,
    token=HF_TOKEN,
)
print(f"[timing] fal Fast transformer load: {time.perf_counter() - t0:.1f}s", flush=True)

# Optional local prompt enhancer (Qwen3-VL LM head grafted onto the text encoder). Free, on-device.
enhancer_head = None
if _HAS_HEAD:
    try:
        enhancer_head = Ideogram4PromptEnhancerHead.from_pretrained(
            LM_HEAD_REPO, torch_dtype=torch.bfloat16, token=HF_TOKEN
        )
    except Exception as e:
        print(f"[enhancer] LM-head load failed (raw prompt only): {e!r}", flush=True)

t0 = time.perf_counter()
# Reuse the SAME Fast transformer for the unconditional slot: with a constant guidance_scale of 1.0 the
# velocity blend is v = 1.0*pos_v + 0.0*neg_v == pos_v, so the unconditional branch output is discarded
# entirely — mathematically identical to fal's single conditional branch, no extra VRAM (one module,
# registered twice). This keeps us fully compatible with stock diffusers `main`.
pipe = Ideogram4Pipeline.from_pretrained(
    COMPONENTS_REPO,
    revision=COMPONENTS_REVISION,
    transformer=fast_transformer,
    unconditional_transformer=fast_transformer,
    prompt_enhancer_head=enhancer_head,
    torch_dtype=torch.bfloat16,
    token=HF_TOKEN,
)
pipe.to("cuda")
print(f"[timing] pipeline assemble + to(cuda): {time.perf_counter() - t0:.1f}s", flush=True)

# --- AoTI: load the precompiled repeated block (compiled offline by the companion Space) --------------
# torch.compile (JIT) isn't supported on ZeroGPU, so we use PyTorch AOTInductor via the official
# `spaces` helper. The Ideogram4TransformerBlock was AoT-compiled once by
# hugging-apps/ideogram-v4-fast-demo-aoti-compile and published to the AOTI_REPO model repo as
# {BlockClassName}/package.pt2. spaces.aoti_blocks_load downloads it and patches every repeated block
# so serving cold starts pay no compilation cost. Weights stay runtime inputs, so this is a no-op-safe
# graph swap. If the artifact is missing/stale/incompatible we fall back to the eager block cleanly.
AOTI_REPO = os.environ.get("AOTI_REPO", "hugging-apps/ideogram-v4-fast-demo-aoti")
try:
    t0 = time.perf_counter()
    spaces.aoti_blocks_load(pipe.transformer, AOTI_REPO)
    print(f"[aoti] loaded precompiled block from {AOTI_REPO} in {time.perf_counter() - t0:.1f}s", flush=True)
except Exception as e:
    print(f"[aoti] load failed ({e!r}); running eager (no AoTI speedup)", flush=True)


def _looks_like_json(text):
    s = (text or "").strip()
    return s.startswith("{") and s.endswith("}")


# --- Warm the local prompt enhancer on the startup worker (forks inherit the graft) --------------------
@spaces.GPU(duration=120, size="xlarge")
def _warmup():
    if enhancer_head is not None:
        pipe.upsample_prompt("a red apple on a wooden table", height=1024, width=1024)


if enhancer_head is not None:
    try:
        _warmup()
        print("[enhancer] prompt enhancer grafted", flush=True)
    except Exception as e:
        print(f"[enhancer] warmup failed (will graft lazily on first request): {e!r}", flush=True)


# --- Dynamic GPU budget (per-step time scales with image tokens) ---------------------------------------
_TOK_1024, _TOK_2048 = (1024 // 16) ** 2, (2048 // 16) ** 2
_PS_1024, _PS_2048 = 1.0, 6.0  # measured/estimated seconds per denoising step
_PS_B = (_PS_2048 - _PS_1024) / (_TOK_2048 - _TOK_1024)
_PS_A = _PS_1024 - _PS_B * _TOK_1024


def _per_step(width, height):
    return max(0.3, _PS_A + _PS_B * ((int(width) // 16) * (int(height) // 16)))


def _gpu_duration(prompt, aspect_ratio=DEFAULT_RATIO, enhance=True, seed=0, randomize_seed=True, progress=None):
    width, height = ASPECT_RATIOS.get(aspect_ratio, ASPECT_RATIOS[DEFAULT_RATIO])
    budget = FAST_STEPS * _per_step(width, height) * 1.4 + 20
    if enhance:
        budget += 20  # local prompt upsampling (grafted Qwen head)
    return max(60, int(math.ceil(budget)))


@spaces.GPU(duration=_gpu_duration, size="xlarge")
def generate(
    prompt,
    aspect_ratio=DEFAULT_RATIO,
    enhance=True,
    seed=0,
    randomize_seed=True,
    progress=gr.Progress(track_tqdm=True),
):
    """Generate an image with Ideogram 4 Fast (by fal) from a text prompt or a structured JSON caption.

    Args:
        prompt: A plain-text prompt (expanded into Ideogram's JSON caption when `enhance` is on) or a
            complete structured JSON caption fed to the model verbatim.
        aspect_ratio: One of the preset aspect-ratio / resolution labels.
        enhance: Expand a plain-text prompt into Ideogram's structured JSON caption before generation.
        seed: RNG seed (ignored when `randomize_seed` is on).
        randomize_seed: Draw a fresh random seed each run.
    """
    if not prompt or not prompt.strip():
        raise gr.Error("Please enter a prompt.")
    if randomize_seed or seed is None or int(seed) < 0:
        seed = random.randint(0, MAX_SEED)
    seed = int(seed)
    width, height = ASPECT_RATIOS.get(aspect_ratio, ASPECT_RATIOS[DEFAULT_RATIO])
    generator = torch.Generator(device="cuda").manual_seed(seed)

    # Ideogram 4 is trained on structured JSON captions. If the user typed JSON, honour it verbatim;
    # otherwise upsample the plain prompt into a native caption (best quality). Toggle off to feed raw.
    final_prompt = prompt.strip()
    if _looks_like_json(final_prompt):
        pass  # already a JSON caption
    elif enhance and enhancer_head is not None:
        progress(0.0, desc="✍️ Writing the JSON caption…")
        try:
            final_prompt = pipe.upsample_prompt(
                final_prompt, height=height, width=width, generator=generator
            )[0]
        except Exception as e:
            print(f"[enhancer] failed, using raw prompt: {e!r}", flush=True)
            gr.Warning("Prompt enhancer unavailable — generating from the raw prompt.")

    progress(0.0, desc="🎨 Generating…")
    t = time.perf_counter()
    image = pipe(
        prompt=final_prompt,
        width=width,
        height=height,
        num_inference_steps=FAST_STEPS,
        guidance_scale=1.0,  # single conditional branch (no runtime CFG) — matches fal Fast
        guidance_schedule=None,  # must clear the pipeline's default schedule when guidance_scale is set
        mu=FAST_MU,
        std=FAST_STD,
        generator=generator,
    ).images[0]
    dt = time.perf_counter() - t
    print(f"[timing] {FAST_STEPS}-step generation: {dt:.2f}s", flush=True)

    try:
        caption = json.loads(final_prompt)
    except (TypeError, ValueError):
        caption = {"prompt": final_prompt}
    return image, seed, caption, f"{FAST_STEPS} steps · {dt:.1f}s"


CSS = """
#col-container { max-width: 1200px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""

DEFAULT_PROMPT = "a bold typographic poster that reads 'FAST BY FAL' in black and electric orange on warm white paper"

with gr.Blocks(title="Ideogram 4 Fast · by fal") as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            "# Ideogram 4 Fast ⚡ — by fal\n"
            "[**fal/ideogram-v4-fast**](https://huggingface.co/fal/ideogram-v4-fast) is a speed-distilled "
            "Ideogram 4 checkpoint: **20 steps, one transformer, no runtime CFG**. Ideogram 4 is trained on "
            "**structured JSON captions**, so a plain prompt is expanded into one on-device (Qwen3-VL) before "
            "generation — or paste your own JSON caption.\n\n"
            "[Fast model](https://huggingface.co/fal/ideogram-v4-fast) · "
            "[Base Ideogram 4](https://huggingface.co/ideogram-ai/ideogram-4-nf4-diffusers) · "
            "[fal blog](https://blog.fal.ai/serving-sub-second-ideogram-v4-without-quality-loss/)\n\n"
            "> Note: runs in **bf16** on ZeroGPU. fal's intended quality is on an FP4 execution path, so "
            "results here may differ slightly from the production endpoint."
        )

        with gr.Row():
            with gr.Column():
                prompt = gr.Textbox(
                    label="Prompt",
                    value=DEFAULT_PROMPT,
                    lines=3,
                    info="Plain text (auto-expanded to a JSON caption) or a full structured JSON caption.",
                )
                run = gr.Button("Generate", variant="primary")
                aspect_ratio = gr.Radio(
                    choices=list(ASPECT_RATIOS.keys()), value=DEFAULT_RATIO, label="Aspect ratio"
                )
                with gr.Accordion("Advanced settings", open=False):
                    enhance = gr.Checkbox(
                        label="Enhance prompt → JSON caption",
                        value=True,
                        info="Ideogram 4 is trained on structured captions. On = best quality (recommended). "
                        "Ignored when the prompt is already JSON.",
                    )
                    with gr.Row():
                        seed = gr.Number(label="Seed", value=0, precision=0)
                        randomize = gr.Checkbox(label="Randomize seed", value=True)
            with gr.Column():
                out_image = gr.Image(label="Output", type="pil")
                with gr.Row():
                    out_seed = gr.Number(label="Seed used", precision=0, interactive=False)
                    out_time = gr.Textbox(label="Generation", interactive=False)
                out_caption = gr.JSON(label="Caption fed to the model")

        gr.Examples(
            examples=[
                ["a bold typographic poster that reads 'FAST BY FAL' in black and electric orange on warm white paper"],
                ["a ginger cat wearing a tiny wizard hat reading a spellbook, storybook illustration"],
                ["an isometric illustration of a tiny city floating in the clouds"],
                ["a vintage travel poster for Kyoto in autumn, deco typography"],
                ["a golden retriever on a skateboard, studio photo"],
            ],
            inputs=[prompt],
            outputs=[out_image, out_seed, out_caption, out_time],
            fn=generate,
            cache_examples=True,
            cache_mode="lazy",
        )

    run.click(
        generate,
        inputs=[prompt, aspect_ratio, enhance, seed, randomize],
        outputs=[out_image, out_seed, out_caption, out_time],
        api_name="generate",
    )

demo.launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True)