app.py

import os

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
os.environ.setdefault("TORCHDYNAMO_DISABLE", "1")

import json
import math
import random
import time
from threading import Thread

import gradio as gr
import spaces
import torch
from huggingface_hub import hf_hub_download
from transformers import AutoModel

from diffusers import Ideogram4Pipeline
from diffusers.quantizers.bitsandbytes.bnb_quantizer import BnB4BitDiffusersQuantizer


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

MODEL_ID = "ideogram-ai/ideogram-4-nf4"
AOTI_REPO = "multimodalart/i4-block-aoti"
TEXT_ENCODER_ID = os.environ.get("TEXT_ENCODER_ID", "huihui-ai/Huihui-Qwen3-VL-8B-Instruct-abliterated")
MAX_SEED = 2**31 - 1
HF_TOKEN = os.environ.get("HF_TOKEN")

MODES = {
    "Turbo · 12 steps": dict(num_inference_steps=12, final_guidance_steps=1, mu=0.5, std=1.75),
    "Default · 20 steps": dict(num_inference_steps=20, final_guidance_steps=2, mu=0.0, std=1.75),
    "Quality · 48 steps": dict(num_inference_steps=48, final_guidance_steps=3, mu=0.0, std=1.5),
}
DEFAULT_MAIN_GUIDANCE = 7.0
DEFAULT_FINAL_GUIDANCE = 3.0

DEFAULT_CAPTION = {
    "high_level_description": "A clean poster announcing a small experimental image generation lab.",
    "style_description": {
        "aesthetics": "minimal, precise, modern graphic design, generous whitespace",
        "lighting": "even soft studio lighting",
        "medium": "graphic_design",
        "art_style": "flat vector poster, crisp sans-serif typography",
        "color_palette": ["#F9FAFB", "#111827", "#2563EB", "#F97316"],
    },
    "compositional_deconstruction": {
        "background": "A warm off-white poster background with a subtle paper texture.",
        "elements": [
            {
                "type": "text",
                "bbox": [250, 130, 430, 870],
                "text": "IDEOGRAM 4",
                "desc": "Large bold black uppercase title text centered near the upper half.",
                "color_palette": ["#111827"],
            },
            {
                "type": "text",
                "bbox": [470, 240, 580, 760],
                "text": "JSON LAB",
                "desc": "Medium blue uppercase subtitle text centered under the main title.",
                "color_palette": ["#2563EB"],
            },
            {
                "type": "obj",
                "bbox": [660, 330, 760, 670],
                "desc": "A thin orange rounded rectangle outline used as a design accent.",
                "color_palette": ["#F97316"],
            },
        ],
    },
}


def dumps_caption(caption):
    return json.dumps(caption, ensure_ascii=False, separators=(",", ":"), indent=2)


def normalize_caption(raw_caption):
    try:
        caption = json.loads(raw_caption, strict=False)
    except Exception as e:
        raise gr.Error(f"JSON parse error: {e}") from e

    if not isinstance(caption, dict):
        raise gr.Error("Top-level JSON must be an object.")

    if "compositional_deconstruction" not in caption:
        gr.Warning("compositional_deconstruction is missing. The model accepts any string, but this is outside the usual Ideogram 4 caption format.")

    return json.dumps(caption, ensure_ascii=False, separators=(",", ":")), caption


def build_preset(mode, main_guidance=DEFAULT_MAIN_GUIDANCE, final_guidance=DEFAULT_FINAL_GUIDANCE):
    preset = dict(MODES.get(mode, MODES["Default · 20 steps"]))
    steps = int(preset.pop("num_inference_steps"))
    final_steps = min(int(preset.pop("final_guidance_steps")), steps)
    main_steps = steps - final_steps
    guidance_schedule = (float(main_guidance),) * main_steps + (float(final_guidance),) * final_steps
    preset.update(num_inference_steps=steps, guidance_schedule=guidance_schedule)
    return preset


t = time.perf_counter()
if TEXT_ENCODER_ID:
    text_encoder = AutoModel.from_pretrained(
        TEXT_ENCODER_ID,
        torch_dtype=torch.bfloat16,
        token=HF_TOKEN,
        low_cpu_mem_usage=True,
    )
    print(f"[model] using alternate text encoder: {TEXT_ENCODER_ID}", flush=True)
else:
    text_encoder = None

pipe = Ideogram4Pipeline.from_pretrained(
    MODEL_ID,
    text_encoder=text_encoder,
    torch_dtype=torch.bfloat16,
    token=HF_TOKEN,
)
pipe.transformer.dequantize()
pipe.unconditional_transformer.dequantize()
pipe.to("cuda")
print(f"[timing] pipeline load + dequant: {time.perf_counter() - t:.1f}s", flush=True)

try:
    hf_hub_download(AOTI_REPO, "package.pt2", subfolder="Ideogram4TransformerBlock")
    from torch._inductor.cpu_vec_isa import valid_vec_isa_list

    t = time.perf_counter()
    valid_vec_isa_list()
    print(f"[timing] vec-isa prewarm (parent): {time.perf_counter() - t:.1f}s", flush=True)
    AOTI_OK = True
except Exception as e:
    AOTI_OK = False
    print(f"[aoti] prefetch/prewarm failed, running eager: {e!r}", flush=True)

_AOTI_APPLIED = False


def _apply_aoti():
    global _AOTI_APPLIED
    if _AOTI_APPLIED or not AOTI_OK:
        return
    try:
        t = time.perf_counter()
        spaces.aoti_blocks_load(pipe.transformer, AOTI_REPO)
        spaces.aoti_blocks_load(pipe.unconditional_transformer, AOTI_REPO)
        _AOTI_APPLIED = True
        print(f"[timing] aoti_blocks_load (both transformers): {time.perf_counter() - t:.2f}s", flush=True)
    except Exception as e:
        print(f"[aoti] apply failed, running eager: {e!r}", flush=True)


_TOK_1024, _TOK_2048 = (1024 // 16) ** 2, (2048 // 16) ** 2
_PS_1024, _PS_2048 = 1.0 / 1.10, 6.0
_PS_B = (_PS_2048 - _PS_1024) / (_TOK_2048 - _TOK_1024)
_PS_A = _PS_1024 - _PS_B * _TOK_1024
DIFFUSION_OVERHEAD_S = 8
DURATION_MARGIN = 1.3


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


def _gpu_duration(caption_text, mode, width, height, seed, main_guidance, final_guidance, progress=None):
    steps = MODES.get(mode, MODES["Default · 20 steps"])["num_inference_steps"]
    budget = steps * _per_step(width, height) + DIFFUSION_OVERHEAD_S
    return max(60, int(math.ceil(budget * DURATION_MARGIN)))


@spaces.GPU(duration=_gpu_duration, size="xlarge")
def _gpu_generate(caption_text, mode, width, height, seed, main_guidance, final_guidance, progress=gr.Progress(track_tqdm=True)):
    aoti_thread = Thread(target=_apply_aoti, daemon=True)
    aoti_thread.start()
    aoti_thread.join()

    progress(0.0, desc="Generating image")
    generator = torch.Generator(device="cuda").manual_seed(int(seed))
    preset = build_preset(mode, main_guidance, final_guidance)
    t = time.perf_counter()
    image = pipe(prompt=caption_text, width=int(width), height=int(height), generator=generator, **preset).images[0]
    print(
        f"[timing] diffusion ({mode}, guidance={float(main_guidance):.2f}->{float(final_guidance):.2f}): "
        f"{time.perf_counter() - t:.2f}s",
        flush=True,
    )
    return image


def generate(
    caption_json,
    mode="Default · 20 steps",
    width=1024,
    height=1024,
    seed=0,
    randomize_seed=False,
    main_guidance=DEFAULT_MAIN_GUIDANCE,
    final_guidance=DEFAULT_FINAL_GUIDANCE,
    progress=gr.Progress(track_tqdm=True),
):
    caption_text, parsed_caption = normalize_caption(caption_json)
    if randomize_seed or seed < 0:
        seed = random.randint(0, MAX_SEED)
    image = _gpu_generate(caption_text, mode, width, height, seed, main_guidance, final_guidance)
    return image, int(seed), parsed_caption, caption_text


CSS = """
.gradio-container { max-width: 1280px !important; }
textarea { font-family: ui-monospace, SFMono-Regular, Menlo, Consolas, monospace !important; }
"""

with gr.Blocks(theme=gr.themes.Citrus(), title="Ideogram 4 JSON Lab", css=CSS) as demo:
    gr.Markdown(
        "# Ideogram 4 JSON Lab\n"
        "Direct structured JSON caption input for Ideogram 4. No remote magic prompt, no local Qwen prompt upsampling.\n\n"
        f"Text encoder: `{TEXT_ENCODER_ID or 'ideogram-ai/ideogram-4-nf4 bundled Qwen3-VL'}`"
    )

    with gr.Row():
        with gr.Column(scale=6):
            caption = gr.Textbox(label="JSON caption", value=dumps_caption(DEFAULT_CAPTION), lines=28)
            with gr.Row():
                mode = gr.Radio(choices=list(MODES.keys()), value="Default · 20 steps", label="Mode")
                width = gr.Slider(512, 2048, value=1024, step=64, label="Width")
                height = gr.Slider(512, 2048, value=1024, step=64, label="Height")
            with gr.Row():
                main_guidance = gr.Slider(0.0, 9.0, value=DEFAULT_MAIN_GUIDANCE, step=0.25, label="Main guidance")
                final_guidance = gr.Slider(0.0, 9.0, value=DEFAULT_FINAL_GUIDANCE, step=0.25, label="Final guidance")
            with gr.Row():
                seed = gr.Number(label="Seed", value=0, precision=0)
                randomize = gr.Checkbox(label="Randomize seed", value=False)
                run = gr.Button("Generate", variant="primary")
        with gr.Column(scale=5):
            out_image = gr.Image(label="Output", type="pil")
            out_caption = gr.JSON(label="Parsed JSON caption")
            out_text = gr.Textbox(label="Compact caption string sent to model", lines=8)

    gr.Examples(
        examples=[
            [dumps_caption(DEFAULT_CAPTION)],
            [
                dumps_caption(
                    {
                        "high_level_description": "A square package label for a fictional tea brand called BLUE HARBOR.",
                        "style_description": {
                            "aesthetics": "premium, calm, balanced, Japanese-inspired packaging design",
                            "lighting": "even studio light",
                            "medium": "graphic_design",
                            "art_style": "flat vector label design with refined serif typography",
                            "color_palette": ["#F8FAFC", "#0F172A", "#2563EB", "#94A3B8", "#EAB308"],
                        },
                        "compositional_deconstruction": {
                            "background": "A clean ivory square label with a thin navy border.",
                            "elements": [
                                {
                                    "type": "text",
                                    "bbox": [170, 180, 300, 820],
                                    "text": "BLUE HARBOR",
                                    "desc": "Elegant navy serif uppercase brand name centered at the top.",
                                    "color_palette": ["#0F172A"],
                                },
                                {
                                    "type": "obj",
                                    "bbox": [360, 320, 650, 680],
                                    "desc": "A simple blue line illustration of ocean waves inside a gold circular seal.",
                                    "color_palette": ["#2563EB", "#EAB308"],
                                },
                                {
                                    "type": "text",
                                    "bbox": [720, 250, 810, 750],
                                    "text": "EARL GREY",
                                    "desc": "Small spaced navy sans-serif product text centered near the bottom.",
                                    "color_palette": ["#0F172A"],
                                },
                            ],
                        },
                    }
                )
            ],
        ],
        inputs=[caption],
    )

    run.click(
        generate,
        inputs=[caption, mode, width, height, seed, randomize, main_guidance, final_guidance],
        outputs=[out_image, seed, out_caption, out_text],
    )

demo.queue(api_open=False, max_size=8, default_concurrency_limit=1).launch(footer_links=["settings"])