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