ideogram-ai/ideogram4

1import os2 3os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")4# outlines_core ships an @torch.compile bitmask kernel dynamo can't trace (torch.device const) -> noisy5# WON'T CONVERT spam on every local upsample. We never use torch.compile at runtime, so disable dynamo.6os.environ.setdefault("TORCHDYNAMO_DISABLE", "1")7 8# diffusers (with Ideogram4 support) is pip-installed from the PR — see requirements.txt. No bundled source.9 10import json11import random12import time13from threading import Thread14 15import gradio as gr16import requests17import spaces18import torch19from huggingface_hub import hf_hub_download20 21from diffusers import Ideogram4Pipeline22 23# Runtime shim (keeps the bundled diffusers pristine): cu130-era bitsandbytes returns Params4bit.shape as a24# plain tuple, but diffusers' check_quantized_param_shape calls .numel() on it. math.prod handles both, so25# this is a no-op once diffusers/bnb fix it upstream.26import math  # noqa: E40227 28from diffusers.quantizers.bitsandbytes.bnb_quantizer import BnB4BitDiffusersQuantizer  # noqa: E40229 30 31def _check_quantized_param_shape(self, param_name, current_param, loaded_param):32    n = math.prod(tuple(current_param.shape))33    inferred_shape = (n,) if "bias" in param_name else ((n + 1) // 2, 1)34    if tuple(loaded_param.shape) != tuple(inferred_shape):35        raise ValueError(f"Expected flattened shape of {param_name} to be {inferred_shape}, got {tuple(loaded_param.shape)}.")36    return True37 38 39BnB4BitDiffusersQuantizer.check_quantized_param_shape = _check_quantized_param_shape40 41MODEL_ID = "ideogram-ai/ideogram-4-nf4"42LM_HEAD_REPO = "multimodalart/qwen3-vl-8b-instruct-lm-head"43AOTI_REPO = "multimodalart/i4-block-aoti"44AOTI_BLOCK_FILE = "Ideogram4TransformerBlock/package.pt2"45MAX_SEED = 2**31 - 146 47# Prompt upsampling: Ideogram's hosted magic-prompt (default) with the local Qwen graft as fallback.48IDEOGRAM_MAGIC_PROMPT_URL = "https://api.ideogram.ai/v1/ideogram-v4/magic-prompt"49IDEOGRAM_API_KEY = os.environ.get("IDEOGRAM_API_KEY")50UPSAMPLERS = ["Ideogram (remote)", "Qwen (local)"]51 52# V4 presets (forward step-order: main CFG 7.0 -> polish 3.0).53MODES = {54    "Turbo · 12 steps": dict(num_inference_steps=12, guidance_schedule=(7.0,) * 11 + (3.0,) * 1, mu=0.5, std=1.75),55    "Default · 20 steps": dict(num_inference_steps=20, guidance_schedule=(7.0,) * 18 + (3.0,) * 2, mu=0.0, std=1.75),56    "Quality · 48 steps": dict(num_inference_steps=48, guidance_schedule=(7.0,) * 45 + (3.0,) * 3, mu=0.0, std=1.5),57}58 59# --- Pipeline: dequantize both transformers nf4 -> bf16 in the parent (CPU) so AOTI can bind its weight-less60# graph to real bf16 weights (this is repo cold start, which is fine; function cold start stays fast). ---61t = time.perf_counter()62pipe = Ideogram4Pipeline.from_pretrained(MODEL_ID, torch_dtype=torch.bfloat16)63pipe.transformer.dequantize()64pipe.unconditional_transformer.dequantize()65pipe.to("cuda")66print(f"[timing] pipeline load + dequant: {time.perf_counter() - t:.1f}s", flush=True)67 68# The local prompt-enhancer LM head is grafted lazily by `pipe.upsample_prompt` on first use (onto the worker's69# GPU), so no explicit load is needed here. Local is only the fallback; Ideogram's remote API is the default.70 71# Pre-fetch the AOTI package AND pre-warm torch-inductor's CPU-ISA probe in the PARENT (repo cold start). The72# probe (valid_vec_isa_list) compiles test programs (~20s) the first time aoti_blocks_load builds a LazyAOTIModel;73# doing it once here means every ZeroGPU fork inherits the functools.cache, so per-worker (function cold start)74# aoti_blocks_load is just the ~instant block patch instead of a ~20s compile.75try:76    hf_hub_download(AOTI_REPO, "package.pt2", subfolder="Ideogram4TransformerBlock")77    from torch._inductor.cpu_vec_isa import valid_vec_isa_list78 79    t = time.perf_counter()80    valid_vec_isa_list()81    print(f"[timing] vec-isa prewarm (parent): {time.perf_counter() - t:.1f}s", flush=True)82    AOTI_OK = True83except Exception as e:84    AOTI_OK = False85    print(f"[aoti] prefetch/prewarm failed, running eager: {e!r}", flush=True)86 87_AOTI_APPLIED = False88 89 90def _apply_aoti():91    """Patch the compiled block onto every Ideogram4TransformerBlock of both transformers (once per worker).92 93    `aoti_blocks_load` is lazy (binds forward, defers the .so to first diffusion step) and CPU-only, so this is94    safe to run in a background thread overlapping the (transformer-idle) upsampling step."""95    global _AOTI_APPLIED96    if _AOTI_APPLIED or not AOTI_OK:97        return98    try:99        t = time.perf_counter()100        spaces.aoti_blocks_load(pipe.transformer, AOTI_REPO)101        spaces.aoti_blocks_load(pipe.unconditional_transformer, AOTI_REPO)102        _AOTI_APPLIED = True103        print(f"[timing] aoti_blocks_load (both transformers): {time.perf_counter() - t:.2f}s", flush=True)104    except Exception as e:  # never let a bind hiccup block generation105        print(f"[aoti] apply failed, running eager: {e!r}", flush=True)106 107 108def remote_upsample(prompt, width, height):109    """Rewrite the prompt into Ideogram's native JSON caption via the hosted magic-prompt API."""110    d = math.gcd(width, height) or 1111    aspect_ratio = f"{width // d}x{height // d}"  # Ideogram's WxH form112    resp = requests.post(113        IDEOGRAM_MAGIC_PROMPT_URL,114        headers={"Api-Key": IDEOGRAM_API_KEY, "Content-Type": "application/json"},115        json={"text_prompt": prompt, "aspect_ratio": aspect_ratio},116        timeout=120,117    )118    resp.raise_for_status()119    jp = resp.json().get("json_prompt")120    if not jp:121        raise RuntimeError("Ideogram API returned no json_prompt")122    jp.pop("aspect_ratio", None)123    for el in jp.get("compositional_deconstruction", {}).get("elements", []):124        if isinstance(el, dict):125            el.pop("bbox", None)126    return json.dumps(jp, ensure_ascii=False, separators=(",", ":"))127 128 129# --- Dynamic GPU duration ---------------------------------------------------------------------------------130# Per-step diffusion time, linear in image tokens between the two measured anchors (1024 @ 1.10 it/s,131# 2048 @ 6 s/it). The chord overestimates in between, so it's a safe budget; clamped low for small images.132# Remote upsample is a network call done OFF the GPU (in `generate`), so it isn't budgeted here.133_TOK_1024, _TOK_2048 = (1024 // 16) ** 2, (2048 // 16) ** 2  # 4096, 16384 image tokens134_PS_1024, _PS_2048 = 1.0 / 1.10, 6.0  # measured seconds/iteration135_PS_B = (_PS_2048 - _PS_1024) / (_TOK_2048 - _TOK_1024)136_PS_A = _PS_1024 - _PS_B * _TOK_1024137LOCAL_UPSAMPLE_S = 15  # local Qwen graft+generate (~12s) with headroom138DIFFUSION_OVERHEAD_S = 8  # .so dlopen + block patch + cudnn setup on a cold worker's first forward139DURATION_MARGIN = 1.3140 141 142def _per_step(width, height):143    return max(0.2, _PS_A + _PS_B * ((int(width) // 16) * (int(height) // 16)))144 145 146def _gpu_duration(final_prompt, mode, width, height, seed, do_local, progress=None):147    steps = MODES.get(mode, MODES["Default · 20 steps"])["num_inference_steps"]148    budget = steps * _per_step(width, height) + DIFFUSION_OVERHEAD_S149    if do_local:150        budget += LOCAL_UPSAMPLE_S151    return max(60, int(math.ceil(budget * DURATION_MARGIN)))152 153 154@spaces.GPU(duration=_gpu_duration, size="xlarge")155def _gpu_generate(final_prompt, mode, width, height, seed, do_local, progress=gr.Progress(track_tqdm=True)):156    # Overlap the AOTI block-patch with the (transformer-idle) local upsample, if any.157    aoti_thread = Thread(target=_apply_aoti, daemon=True)158    aoti_thread.start()159    if do_local:160        progress(0.0, desc="✍️ Upsampling (local Qwen)…")161        t = time.perf_counter()162        try:163            final_prompt = pipe.upsample_prompt(164                final_prompt, height=int(height), width=int(width), lm_head_repo_id=LM_HEAD_REPO165            )[0]166            print(f"[timing] upsample local: {time.perf_counter() - t:.2f}s", flush=True)167        except Exception as e:168            print(f"[upsample] local failed: {e!r}", flush=True)169            gr.Warning("Local upsampler unavailable — generating from the raw prompt.")170    aoti_thread.join()  # ensure blocks are patched before the diffusion loop171 172    progress(0.0, desc="🎨 Generating image…")173    generator = torch.Generator(device="cuda").manual_seed(int(seed))174    preset = MODES.get(mode, MODES["Default · 20 steps"])175    t = time.perf_counter()176    image = pipe(prompt=final_prompt, width=int(width), height=int(height), generator=generator, **preset).images[0]177    print(f"[timing] diffusion ({mode}): {time.perf_counter() - t:.2f}s", flush=True)178 179    try:180        caption = json.loads(final_prompt)181    except Exception:182        caption = {"prompt": final_prompt}183    return image, int(seed), caption184 185 186def generate(187    prompt,188    mode="Default · 20 steps",189    upsampler=UPSAMPLERS[0],190    width=1024,191    height=1024,192    seed=0,193    randomize_seed=False,194    progress=gr.Progress(track_tqdm=True),195):196    if randomize_seed or seed < 0:197        seed = random.randint(0, MAX_SEED)198 199    # Remote upsample is a network call -> run it here, OFF the GPU. Fall back to local (on-GPU) on failure.200    final_prompt, do_local = prompt, True201    if upsampler == UPSAMPLERS[0] and IDEOGRAM_API_KEY:202        progress(0.0, desc="✍️ Upsampling (Ideogram)…")203        t = time.perf_counter()204        try:205            final_prompt = remote_upsample(prompt, int(width), int(height))206            do_local = False207            print(f"[timing] upsample remote (off-GPU): {time.perf_counter() - t:.2f}s", flush=True)208        except Exception as e:209            print(f"[upsample] remote failed, falling back to local: {e!r}", flush=True)210            gr.Warning("Ideogram API unavailable — using the local Qwen upsampler.")211 212    return _gpu_generate(final_prompt, mode, width, height, seed, do_local)213 214 215@spaces.GPU(size="xlarge")216def _warmup():217    """Warm the local upsampler (lazy LM-head graft) on the startup worker (no diffusion)."""218    _apply_aoti()  # no-op while AOTI is disabled219    t = time.perf_counter()220    pipe.upsample_prompt("a red apple on a wooden table", height=1024, width=1024, lm_head_repo_id=LM_HEAD_REPO)221    print(f"[timing] warmup upsample: {time.perf_counter() - t:.2f}s", flush=True)222 223 224try:225    _warmup()226except Exception as e:  # a flaky ZeroGPU worker must not take down the Space227    print(f"[warmup] failed (will warm lazily on first request): {e!r}", flush=True)228 229CSS='''230.dark .gradio-container { color: var(--body-text-color); }231'''232with gr.Blocks(theme=gr.themes.Citrus(), title="Ideogram 4", css=CSS) as demo:233    gr.Markdown(234        "# Ideogram 4\n"235        "Ideogram's first open-weights model — a 9.3B-parameter text-to-image foundation model at the "236        "forefront of design, with best-in-class text rendering.\n\n"237        "[Model](https://huggingface.co/ideogram-ai/ideogram-4-nf4) · "238        "[Model (fp8)](https://huggingface.co/ideogram-ai/ideogram-4-fp8) · "239        "[Blog](https://ideogram.ai/blog/ideogram-4.0/)"240    )241 242    with gr.Row():243        with gr.Column():244            prompt = gr.Textbox(label="Prompt", value="a ginger cat wearing a tiny wizard hat reading a spellbook", lines=3)245            mode = gr.Radio(choices=list(MODES.keys()), value="Default · 20 steps", label="Mode (speed ↔ quality)")246            run = gr.Button("Generate", variant="primary")247            with gr.Accordion("Advanced", open=False):248                upsampler = gr.Radio(249                    choices=UPSAMPLERS,250                    value=UPSAMPLERS[0],251                    label="Prompt upsampler",252                    info="Rewrite into Ideogram's native JSON caption. Remote (Ideogram) preferred; falls back to local.",253                )254                with gr.Row():255                    width = gr.Slider(512, 2048, value=1024, step=64, label="Width")256                    height = gr.Slider(512, 2048, value=1024, step=64, label="Height")257                with gr.Row():258                    seed = gr.Number(label="Seed", value=0, precision=0)259                    randomize = gr.Checkbox(label="Randomize seed", value=True)260        with gr.Column():261            out_image = gr.Image(label="Output", type="pil")262            out_caption = gr.JSON(label="Caption fed to the model (upsampled when enabled)")263 264    gr.Examples(265        examples=[266            ["a ginger cat wearing a tiny wizard hat reading a spellbook"],267            ["an isometric illustration of a tiny city floating in the clouds"],268            ["a golden retriever on a skateboard"],269        ],270        inputs=[prompt],271        outputs=[out_image, seed, out_caption],272        fn=generate,273        cache_examples=True,274        cache_mode="lazy",275    )276 277    run.click(278        generate,279        inputs=[prompt, mode, upsampler, width, height, seed, randomize],280        outputs=[out_image, seed, out_caption],281    )282 283demo.launch()284