Create direct JSON caption Ideogram 4 lab

README.md

@@ -1,13 +1,14 @@
 ---
-title: Ideogram4 Json Lab
-emoji: 🦀
-colorFrom: purple
-colorTo: green
+title: Ideogram 4 JSON Lab
+emoji: 🧪
+colorFrom: indigo
+colorTo: pink
 sdk: gradio
-sdk_version: 6.16.0
-python_version: '3.12'
+sdk_version: 6.15.2
 app_file: app.py
-pinned: false
+short_description: Direct JSON caption lab for Ideogram 4
+python_version: '3.12'
+startup_duration_timeout: 1h
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+Direct structured JSON caption experiments for Ideogram 4.

app.py

@@ -0,0 +1,253 @@
+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 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"
+MAX_SEED = 2**31 - 1
+
+MODES = {
+    "Turbo · 12 steps": dict(num_inference_steps=12, guidance_schedule=(7.0,) * 11 + (3.0,) * 1, mu=0.5, std=1.75),
+    "Default · 20 steps": dict(num_inference_steps=20, guidance_schedule=(7.0,) * 18 + (3.0,) * 2, mu=0.0, std=1.75),
+    "Quality · 48 steps": dict(num_inference_steps=48, guidance_schedule=(7.0,) * 45 + (3.0,) * 3, mu=0.0, std=1.5),
+}
+
+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)
+    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
+
+
+t = time.perf_counter()
+pipe = Ideogram4Pipeline.from_pretrained(MODEL_ID, torch_dtype=torch.bfloat16)
+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, 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, 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 = MODES.get(mode, MODES["Default · 20 steps"])
+    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}): {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, 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)
+    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."
+    )
+
+    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():
+                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],
+        outputs=[out_image, seed, out_caption, out_text],
+    )
+
+demo.launch()

requirements.txt

@@ -0,0 +1,9 @@
+# diffusers with Ideogram4 support — PR #13860 head, pinned to the verified SHA (AOTI artifact matches this block).
+diffusers @ git+https://github.com/huggingface/diffusers.git@04b197eece42bfc88d1814b20e07987d94cccaa7
+transformers==5.8.0
+peft==0.19.1
+accelerate==1.10.1
+bitsandbytes
+sentencepiece
+outlines==1.3.0
+pydantic>=2