Spaces:
Running on Zero
Running on Zero
| import os | |
| os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") | |
| # outlines_core ships an @torch.compile bitmask kernel dynamo can't trace (torch.device const) -> noisy | |
| # WON'T CONVERT spam on every local upsample. We never use torch.compile at runtime, so disable dynamo. | |
| os.environ.setdefault("TORCHDYNAMO_DISABLE", "1") | |
| import base64 | |
| import io | |
| import json | |
| import random | |
| import time | |
| from threading import Thread | |
| import gradio as gr | |
| import requests | |
| import spaces | |
| import torch | |
| from huggingface_hub import hf_hub_download | |
| from diffusers import Ideogram4Pipeline | |
| # Runtime shim (keeps the bundled diffusers pristine): 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. | |
| import math # noqa: E402 | |
| 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 | |
| MODEL_ID = "ideogram-ai/ideogram-4-nf4" | |
| LM_HEAD_REPO = "multimodalart/qwen3-vl-8b-instruct-lm-head" | |
| AOTI_REPO = "multimodalart/i4-block-aoti" | |
| AOTI_BLOCK_FILE = "Ideogram4TransformerBlock/package.pt2" | |
| MAX_SEED = 2**31 - 1 | |
| # Prompt upsampling: Ideogram's hosted magic-prompt (default) with the local Qwen graft as fallback, | |
| # plus "None" — the Studio JSON (or raw text) goes to the model verbatim. | |
| IDEOGRAM_MAGIC_PROMPT_URL = "https://api.ideogram.ai/v1/ideogram-v4/magic-prompt" | |
| IDEOGRAM_API_KEY = os.environ.get("IDEOGRAM_API_KEY") | |
| UPSAMPLER_REMOTE = "Ideogram (remote)" | |
| UPSAMPLER_LOCAL = "Qwen (local)" | |
| UPSAMPLERS = [UPSAMPLER_REMOTE, UPSAMPLER_LOCAL] | |
| # V4 presets (forward step-order: main CFG 7.0 -> polish 3.0). | |
| 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), | |
| } | |
| # --- Pipeline: dequantize both transformers nf4 -> bf16 in the parent (CPU) so AOTI can bind its weight-less | |
| # graph to real bf16 weights (this is repo cold start, which is fine; function cold start stays fast). --- | |
| 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) | |
| # Pre-fetch the AOTI package AND pre-warm torch-inductor's CPU-ISA probe in the PARENT (repo cold start). The | |
| # probe (valid_vec_isa_list) compiles test programs (~20s) the first time aoti_blocks_load builds a LazyAOTIModel; | |
| # doing it once here means every ZeroGPU fork inherits the functools.cache, so per-worker (function cold start) | |
| # aoti_blocks_load is just the ~instant block patch instead of a ~20s compile. | |
| 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(): | |
| """Patch the compiled block onto every Ideogram4TransformerBlock of both transformers (once per worker).""" | |
| 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: # never let a bind hiccup block generation | |
| print(f"[aoti] apply failed, running eager: {e!r}", flush=True) | |
| def remote_upsample(prompt, width, height): | |
| """Rewrite the prompt into Ideogram's native JSON caption via the hosted magic-prompt API. | |
| Unlike the plain demo, bbox entries are KEPT — the Studio canvas editor uses them for layout control.""" | |
| d = math.gcd(width, height) or 1 | |
| aspect_ratio = f"{width // d}x{height // d}" # Ideogram's WxH form | |
| resp = requests.post( | |
| IDEOGRAM_MAGIC_PROMPT_URL, | |
| headers={"Api-Key": IDEOGRAM_API_KEY, "Content-Type": "application/json"}, | |
| json={"text_prompt": prompt, "aspect_ratio": aspect_ratio}, | |
| timeout=120, | |
| ) | |
| resp.raise_for_status() | |
| jp = resp.json().get("json_prompt") | |
| if not jp: | |
| raise RuntimeError("Ideogram API returned no json_prompt") | |
| jp.pop("aspect_ratio", None) | |
| return jp | |
| # --- Dynamic GPU duration --------------------------------------------------------------------------------- | |
| # Per-step diffusion time, linear in image tokens between the two measured anchors (1024 @ 1.10 it/s, | |
| # 2048 @ 6 s/it). The chord overestimates in between, so it's a safe budget; clamped low for small images. | |
| _TOK_1024, _TOK_2048 = (1024 // 16) ** 2, (2048 // 16) ** 2 # 4096, 16384 image tokens | |
| _PS_1024, _PS_2048 = 1.0 / 1.10, 6.0 # measured seconds/iteration | |
| _PS_B = (_PS_2048 - _PS_1024) / (_TOK_2048 - _TOK_1024) | |
| _PS_A = _PS_1024 - _PS_B * _TOK_1024 | |
| LOCAL_UPSAMPLE_S = 15 # local Qwen graft+generate (~12s) with headroom | |
| DIFFUSION_OVERHEAD_S = 8 # .so dlopen + block patch + cudnn setup on a cold worker's first forward | |
| 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(final_prompt, mode, width, height, seed, do_local, progress=None): | |
| steps = MODES.get(mode, MODES["Default · 20 steps"])["num_inference_steps"] | |
| budget = steps * _per_step(width, height) + DIFFUSION_OVERHEAD_S | |
| if do_local: | |
| budget += LOCAL_UPSAMPLE_S | |
| return max(60, int(math.ceil(budget * DURATION_MARGIN))) | |
| def _gpu_generate(final_prompt, mode, width, height, seed, do_local, progress=gr.Progress(track_tqdm=True)): | |
| # Overlap the AOTI block-patch with the (transformer-idle) local upsample, if any. | |
| aoti_thread = Thread(target=_apply_aoti, daemon=True) | |
| aoti_thread.start() | |
| if do_local: | |
| progress(0.0, desc="✍️ Upsampling (local Qwen)…") | |
| t = time.perf_counter() | |
| try: | |
| final_prompt = pipe.upsample_prompt( | |
| final_prompt, height=int(height), width=int(width), lm_head_repo_id=LM_HEAD_REPO | |
| )[0] | |
| print(f"[timing] upsample local: {time.perf_counter() - t:.2f}s", flush=True) | |
| except Exception as e: | |
| print(f"[upsample] local failed: {e!r}", flush=True) | |
| gr.Warning("Local upsampler unavailable — generating from the raw prompt.") | |
| aoti_thread.join() # ensure blocks are patched before the diffusion loop | |
| 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=final_prompt, width=int(width), height=int(height), generator=generator, **preset).images[0] | |
| print(f"[timing] diffusion ({mode}): {time.perf_counter() - t:.2f}s", flush=True) | |
| try: | |
| caption = json.loads(final_prompt) | |
| except Exception: | |
| caption = {"high_level_description": final_prompt} | |
| return image, int(seed), caption | |
| def _gpu_upsample(prompt, width, height, progress=gr.Progress(track_tqdm=True)): | |
| """Prompt-only drafting with the local Qwen graft (no diffusion).""" | |
| progress(0.0, desc="✍️ Upsampling (local Qwen)…") | |
| t = time.perf_counter() | |
| out = pipe.upsample_prompt(prompt, height=int(height), width=int(width), lm_head_repo_id=LM_HEAD_REPO)[0] | |
| print(f"[timing] upsample-only local: {time.perf_counter() - t:.2f}s", flush=True) | |
| return out | |
| def _warmup(): | |
| """Warm the local upsampler (lazy LM-head graft) on the startup worker (no diffusion).""" | |
| _apply_aoti() | |
| t = time.perf_counter() | |
| pipe.upsample_prompt("a red apple on a wooden table", height=1024, width=1024, lm_head_repo_id=LM_HEAD_REPO) | |
| print(f"[timing] warmup upsample: {time.perf_counter() - t:.2f}s", flush=True) | |
| try: | |
| _warmup() | |
| except Exception as e: # a flaky ZeroGPU worker must not take down the Space | |
| print(f"[warmup] failed (will warm lazily on first request): {e!r}", flush=True) | |
| # --- Studio glue ------------------------------------------------------------------------------------------- | |
| def _img_to_data_url(image): | |
| buf = io.BytesIO() | |
| image.save(buf, format="PNG") | |
| return "data:image/png;base64," + base64.b64encode(buf.getvalue()).decode() | |
| def _studio_has_content(j): | |
| if not isinstance(j, dict): | |
| return False | |
| cd = j.get("compositional_deconstruction") or {} | |
| return bool(j.get("high_level_description") or cd.get("elements") or cd.get("background")) | |
| def generate( | |
| mode, | |
| width, | |
| height, | |
| seed, | |
| randomize_seed, | |
| evt: gr.EventData = None, | |
| progress=gr.Progress(track_tqdm=True), | |
| ): | |
| if randomize_seed or seed is None or seed < 0: | |
| seed = random.randint(0, MAX_SEED) | |
| studio = None | |
| if evt is not None: | |
| try: | |
| studio = json.loads(evt.state_json) | |
| except Exception: | |
| studio = None | |
| if not _studio_has_content(studio): | |
| raise gr.Error("The Studio JSON is empty — draft it with ✨ Generate prompt or fill in the editor first.") | |
| final_prompt = json.dumps(studio, ensure_ascii=False, separators=(",", ":")) | |
| image, seed, caption = _gpu_generate(final_prompt, mode, width, height, seed, False) | |
| editor_update = gr.update( | |
| value=caption, | |
| image_url=_img_to_data_url(image), | |
| img_width=int(width), | |
| img_height=int(height), | |
| ) | |
| return editor_update, image, seed | |
| def generate_prompt( | |
| prompt, | |
| upsampler, | |
| width, | |
| height, | |
| progress=gr.Progress(track_tqdm=True), | |
| ): | |
| prompt = (prompt or "").strip() | |
| if not prompt: | |
| raise gr.Error("Type a text prompt first — the upsampler expands it into Ideogram's JSON caption.") | |
| if upsampler == UPSAMPLER_REMOTE and IDEOGRAM_API_KEY: | |
| progress(0.0, desc="✍️ Upsampling (Ideogram)…") | |
| try: | |
| return gr.update(value=remote_upsample(prompt, int(width), int(height))) | |
| except Exception as e: | |
| print(f"[upsample] remote failed, falling back to local: {e!r}", flush=True) | |
| gr.Warning("Ideogram API unavailable — using the local Qwen upsampler.") | |
| raw = _gpu_upsample(prompt, width, height) | |
| try: | |
| jp = json.loads(raw) | |
| except Exception: | |
| jp = {"high_level_description": raw} | |
| return gr.update(value=jp) | |
| # --- Studio editor: custom gr.HTML component --------------------------------------------------------------- | |
| DEFAULT_PROMPT_JSON = { | |
| "high_level_description": "", | |
| "style_description": {"aesthetics": "", "lighting": "", "medium": "", "art_style": "", "color_palette": []}, | |
| "compositional_deconstruction": {"background": "", "elements": []}, | |
| } | |
| EDITOR_CSS = """ | |
| .i4-root { display:flex; flex-direction:column; gap:10px; } | |
| .i4-toolbar { display:flex; justify-content:space-between; align-items:center; gap:8px; flex-wrap:wrap; } | |
| .i4-tabs { display:flex; gap:4px; } | |
| .i4-tab { padding:6px 14px; border-radius:8px; border:1px solid var(--border-color-primary); background:var(--block-background-fill); color:var(--body-text-color); cursor:pointer; font-size:13px; } | |
| .i4-tab.on { border-color:var(--color-accent, #6366f1); color:var(--color-accent, #6366f1); font-weight:700; } | |
| .i4-actions { display:flex; gap:8px; } | |
| .i4-btn { padding:8px 18px; border-radius:8px; border:1px solid var(--border-color-primary); cursor:pointer; background:var(--block-background-fill); color:var(--body-text-color); font-weight:600; font-size:13px; } | |
| .i4-btn.primary { background:var(--button-primary-background-fill, #ea580c); color:var(--button-primary-text-color, #fff); border-color:transparent; } | |
| .i4-btn:hover, .i4-tab:hover, .i4-mini:hover { filter:brightness(1.08); } | |
| .i4-cols { display:flex; gap:14px; align-items:flex-start; flex-wrap:wrap; } | |
| .i4-left { flex:1.5; min-width:320px; display:flex; flex-direction:column; gap:8px; } | |
| .i4-right { flex:1; min-width:260px; display:flex; flex-direction:column; gap:8px; } | |
| .i4-canvas { position:relative; width:100%; border:1.5px dashed var(--border-color-primary); border-radius:8px; overflow:hidden; touch-action:none; user-select:none; -webkit-user-select:none; cursor:crosshair; background:repeating-conic-gradient(var(--background-fill-secondary) 0% 25%, transparent 0% 50%) 50% / 24px 24px; } | |
| .i4-box { position:absolute; border:2px solid #ff3366; background:rgba(255,51,102,.12); cursor:grab; box-sizing:border-box; border-radius:2px; } | |
| .i4-box.text { border-color:#06b6d4; background:rgba(6,182,212,.12); } | |
| .i4-box.sel { border-color:#6366f1; background:rgba(99,102,241,.2); z-index:5; } | |
| .i4-boxlab { position:absolute; top:0; left:0; font-size:10px; line-height:1.2; background:rgba(0,0,0,.65); color:#fff; padding:1px 5px; border-radius:0 0 4px 0; white-space:nowrap; pointer-events:none; max-width:95%; overflow:hidden; text-overflow:ellipsis; } | |
| .i4-handle { position:absolute; width:12px; height:12px; right:-6px; bottom:-6px; background:#fff; border:1px solid #333; cursor:nwse-resize; border-radius:3px; } | |
| .i4-boxedit { position:absolute; top:2px; left:2px; width:calc(100% - 26px); min-width:70px; max-width:calc(100% - 4px); box-sizing:border-box; font-size:11px; line-height:1.3; padding:2px 5px; border:none; border-radius:3px; background:rgba(0,0,0,.7); color:#fff; outline:none; font-family:inherit; cursor:text; } | |
| .i4-boxedit::placeholder { color:rgba(255,255,255,.55); } | |
| .i4-boxtype { position:absolute; top:2px; right:2px; width:18px; height:18px; font-size:11px; line-height:1; border:none; border-radius:3px; background:rgba(0,0,0,.7); color:#fff; cursor:pointer; padding:0; } | |
| .i4-boxtype:hover { background:rgba(0,0,0,.9); } | |
| .i4-panel { border:1px solid var(--border-color-primary); border-radius:8px; padding:12px; background:var(--block-background-fill); display:flex; flex-direction:column; gap:5px; } | |
| .i4-panel-title { font-weight:700; font-size:13px; margin-bottom:4px; } | |
| .i4-panel label { font-size:11px; font-weight:600; text-transform:uppercase; letter-spacing:.04em; opacity:.7; margin-top:6px; } | |
| .i4-panel input[type=text], .i4-panel textarea, .i4-panel select { width:100%; box-sizing:border-box; padding:6px 8px; border:1px solid var(--border-color-primary); border-radius:6px; background:var(--input-background-fill, var(--background-fill-secondary)); color:var(--body-text-color); font-size:13px; font-family:inherit; } | |
| .i4-pills { display:flex; gap:6px; } | |
| .i4-pill { padding:3px 14px; border-radius:20px; border:1px solid var(--border-color-primary); cursor:pointer; background:transparent; color:var(--body-text-color); font-size:12px; } | |
| .i4-pill.on { background:var(--color-accent, #6366f1); color:#fff; border-color:transparent; } | |
| .i4-palrow { display:flex; gap:6px; align-items:center; } | |
| .i4-color { width:38px; height:28px; padding:1px; border:1px solid var(--border-color-primary); border-radius:6px; background:none; cursor:pointer; } | |
| .i4-mini { padding:4px 12px; font-size:12px; border-radius:6px; border:1px solid var(--border-color-primary); background:var(--block-background-fill); color:var(--body-text-color); cursor:pointer; } | |
| .i4-mini.danger { color:#ff3366; border-color:#ff3366; margin-top:10px; } | |
| .i4-ellist { display:flex; flex-direction:column; gap:2px; max-height:380px; overflow-y:auto; border:1px solid var(--border-color-primary); border-radius:8px; padding:8px 10px; background:var(--block-background-fill); } | |
| .i4-eledit { margin:4px 0 10px 18px; border-left:3px solid var(--color-accent, #6366f1); background:var(--background-fill-secondary); } | |
| .i4-elrow { display:flex; align-items:center; gap:8px; padding:3px 6px; border-radius:6px; cursor:pointer; font-size:12.5px; color:var(--body-text-color); } | |
| .i4-elrow:hover { background:var(--background-fill-secondary); } | |
| .i4-elrow.sel { outline:1.5px solid var(--color-accent, #6366f1); } | |
| .i4-eltype { font-size:11px; color:#ff3366; flex-shrink:0; width:12px; text-align:center; } | |
| .i4-eltype.text { color:#06b6d4; } | |
| .i4-elname { white-space:nowrap; overflow:hidden; text-overflow:ellipsis; flex:1; } | |
| .i4-elrow .i4-mini { padding:1px 8px; font-size:11px; flex-shrink:0; } | |
| .i4-swatches { display:flex; flex-wrap:wrap; gap:4px; min-height:4px; } | |
| .i4-swatch { width:26px; height:26px; border-radius:5px; border:1px solid rgba(0,0,0,.35); cursor:pointer; color:transparent; display:flex; align-items:center; justify-content:center; font-size:13px; } | |
| .i4-swatch:hover { color:#fff; text-shadow:0 0 3px #000; } | |
| .i4-json { width:100%; min-height:440px; font-family:var(--font-mono, ui-monospace, monospace); font-size:12px; box-sizing:border-box; padding:10px; border:1px solid var(--border-color-primary); border-radius:8px; background:var(--input-background-fill, var(--background-fill-secondary)); color:var(--body-text-color); } | |
| .i4-jsonbar { display:flex; gap:8px; align-items:center; margin-top:6px; } | |
| .i4-jsonmsg { font-size:12px; opacity:.85; } | |
| .i4-hint { font-size:12px; opacity:.6; } | |
| """ | |
| EDITOR_JS = """ | |
| const root = element.querySelector('.i4-root'); | |
| const deep = (o) => JSON.parse(JSON.stringify(o)); | |
| function normalize(v){ | |
| let s; | |
| try { s = (v && typeof v === 'object') ? deep(v) : (typeof v === 'string' && v.trim() ? JSON.parse(v) : {}); } | |
| catch (e) { s = { high_level_description: String(v || '') }; } | |
| if (!s || typeof s !== 'object' || Array.isArray(s)) s = { high_level_description: String(v || '') }; | |
| s.high_level_description = s.high_level_description || ''; | |
| s.style_description = (s.style_description && typeof s.style_description === 'object') ? s.style_description : {}; | |
| s.compositional_deconstruction = (s.compositional_deconstruction && typeof s.compositional_deconstruction === 'object') ? s.compositional_deconstruction : {}; | |
| if (s.compositional_deconstruction.background === undefined) s.compositional_deconstruction.background = ''; | |
| if (!Array.isArray(s.compositional_deconstruction.elements)) s.compositional_deconstruction.elements = []; | |
| return s; | |
| } | |
| let state = normalize(props.value); | |
| let lastPushed = JSON.stringify(props.value ?? null); | |
| let selIdx = null; | |
| let tab = 'visual'; | |
| let drag = null; | |
| let jsonMsg = ''; | |
| const sd = () => { state.style_description = state.style_description || {}; return state.style_description; }; | |
| const cd = () => { | |
| state.compositional_deconstruction = state.compositional_deconstruction || {}; | |
| const c = state.compositional_deconstruction; | |
| if (!Array.isArray(c.elements)) c.elements = []; | |
| return c; | |
| }; | |
| const els = () => cd().elements; | |
| const isPhoto = () => Object.prototype.hasOwnProperty.call(sd(), 'photo'); | |
| function esc(s){ return String(s ?? '').replace(/&/g,'&').replace(/</g,'<').replace(/>/g,'>').replace(/"/g,'"'); } | |
| function getField(f){ | |
| if (f === 'hld') return state.high_level_description || ''; | |
| if (f === 'bg') return cd().background || ''; | |
| if (f === 'stylefield') return isPhoto() ? (sd().photo || '') : (sd().art_style || ''); | |
| if (f.indexOf('sd.') === 0) return sd()[f.slice(3)] || ''; | |
| return ''; | |
| } | |
| function setField(f, v){ | |
| if (f === 'hld') state.high_level_description = v; | |
| else if (f === 'bg') cd().background = v; | |
| else if (f === 'stylefield') { if (isPhoto()) sd().photo = v; else sd().art_style = v; } | |
| else if (f.indexOf('sd.') === 0) sd()[f.slice(3)] = v; | |
| } | |
| // Ideogram key-order convention: photo mode -> photo before medium; art mode -> medium before art_style. | |
| function orderSD(x){ | |
| const o = {}; | |
| const put = (k) => { const v = x[k]; if (v !== undefined && v !== null && v !== '' && !(Array.isArray(v) && !v.length)) o[k] = v; }; | |
| put('aesthetics'); put('lighting'); | |
| if (Object.prototype.hasOwnProperty.call(x, 'photo')) { put('photo'); put('medium'); } | |
| else { put('medium'); put('art_style'); } | |
| put('color_palette'); | |
| for (const k of Object.keys(x)) { | |
| const v = x[k]; | |
| if (!(k in o) && v !== undefined && v !== null && v !== '' && !(Array.isArray(v) && !v.length)) o[k] = v; | |
| } | |
| return o; | |
| } | |
| function clean(){ | |
| const s = deep(state); | |
| const out = {}; | |
| if (s.high_level_description) out.high_level_description = s.high_level_description; | |
| const sdo = orderSD(s.style_description || {}); | |
| if (Object.keys(sdo).length) out.style_description = sdo; | |
| const cdo = {}; | |
| const scd = s.compositional_deconstruction || {}; | |
| if (scd.background) cdo.background = scd.background; | |
| if (Array.isArray(scd.elements) && scd.elements.length) cdo.elements = scd.elements.map(el => { const c = { ...el }; if (c.type !== 'text') delete c.text; return c; }); | |
| if (Object.keys(cdo).length) out.compositional_deconstruction = cdo; | |
| for (const k of Object.keys(s)) { | |
| if (!(k in out) && k !== 'high_level_description' && k !== 'style_description' && k !== 'compositional_deconstruction') out[k] = s[k]; | |
| } | |
| return out; | |
| } | |
| const serialize = (indent) => JSON.stringify(clean(), null, indent); | |
| function fieldHTML(label, f, kind){ | |
| const v = esc(getField(f)); | |
| if (kind === 'ta') return `<label>${label}</label><textarea data-f="${f}" rows="2">${v}</textarea>`; | |
| return `<label>${label}</label><input type="text" data-f="${f}" value="${v}">`; | |
| } | |
| function swatchesHTML(list, pal){ | |
| return (list || []).map(h => `<span class="i4-swatch" data-act="rmcolor" data-pal="${pal}" data-hex="${esc(h)}" style="background:${esc(h)}" title="${esc(h)}">×</span>`).join(''); | |
| } | |
| function visualHTML(){ | |
| const photo = isPhoto(); | |
| return `<div class="i4-cols"> | |
| <div class="i4-left"> | |
| <div class="i4-canvas"></div> | |
| <div class="i4-hint">Drag on the canvas to add an element box, then type its description right on it · the ▢/T button toggles object vs text · drag to move · corner handle to resize. Elements without a box render freely — 📍 place them to control their layout.</div> | |
| <div class="i4-ellist" id="i4-ellist"></div> | |
| </div> | |
| <div class="i4-right"> | |
| <div class="i4-panel"> | |
| <div class="i4-panel-title">Prompt structure</div> | |
| ${fieldHTML('High-level description', 'hld', 'ta')} | |
| ${fieldHTML('Background', 'bg', 'ta')} | |
| ${fieldHTML('Aesthetics', 'sd.aesthetics')} | |
| ${fieldHTML('Lighting', 'sd.lighting')} | |
| <label>Style mode</label> | |
| <div class="i4-pills"> | |
| <button class="i4-pill ${photo ? 'on' : ''}" data-act="stylemode" data-m="photo">photo</button> | |
| <button class="i4-pill ${photo ? '' : 'on'}" data-act="stylemode" data-m="art">art style</button> | |
| </div> | |
| ${photo ? fieldHTML('Photo', 'stylefield') + fieldHTML('Medium', 'sd.medium') : fieldHTML('Medium', 'sd.medium') + fieldHTML('Art style', 'stylefield')} | |
| <label>Color palette (max 16)</label> | |
| <div class="i4-palrow"><input type="color" class="i4-color" data-pal="global"><button class="i4-mini" data-act="addcolor" data-pal="global">Add</button></div> | |
| <div class="i4-swatches">${swatchesHTML(sd().color_palette, 'global')}</div> | |
| </div> | |
| </div> | |
| </div>`; | |
| } | |
| function jsonHTML(){ | |
| return `<div> | |
| <textarea class="i4-json" spellcheck="false">${esc(serialize(2))}</textarea> | |
| <div class="i4-jsonbar"> | |
| <button class="i4-mini" data-act="applyjson">Apply JSON</button> | |
| <button class="i4-mini" data-act="copyjson">Copy</button> | |
| <span class="i4-jsonmsg">${esc(jsonMsg)}</span> | |
| </div> | |
| </div>`; | |
| } | |
| function renderAll(){ | |
| root.innerHTML = ` | |
| <div class="i4-toolbar"> | |
| <div class="i4-tabs"> | |
| <button class="i4-tab ${tab === 'visual' ? 'on' : ''}" data-act="tab" data-tab="visual">Visual editor</button> | |
| <button class="i4-tab ${tab === 'json' ? 'on' : ''}" data-act="tab" data-tab="json">JSON</button> | |
| </div> | |
| <div class="i4-actions"> | |
| <button class="i4-btn primary" data-act="gen-image">🎨 Generate image</button> | |
| </div> | |
| </div> | |
| <div class="i4-body">${tab === 'visual' ? visualHTML() : jsonHTML()}</div>`; | |
| if (tab === 'visual') { renderCanvas(); renderElList(); } | |
| } | |
| const hasBbox = (el) => Array.isArray(el.bbox) && el.bbox.length === 4; | |
| const elLabel = (el, i) => el.type === 'text' ? '“' + (el.text || '') + '”' : (el.desc || 'object ' + (i + 1)); | |
| function renderCanvas(){ | |
| const c = root.querySelector('.i4-canvas'); | |
| if (!c) return; | |
| const W = props.img_width || 1024, H = props.img_height || 1024; | |
| c.style.aspectRatio = W + ' / ' + H; | |
| if (props.image_url) { c.style.backgroundImage = "url('" + props.image_url + "')"; c.style.backgroundSize = '100% 100%'; } | |
| else { c.style.backgroundImage = ''; c.style.backgroundSize = ''; } | |
| c.innerHTML = els().map((el, i) => { | |
| if (!hasBbox(el)) return ''; | |
| const b = el.bbox; | |
| const isText = el.type === 'text'; | |
| const inner = i === selIdx | |
| ? `<input class="i4-boxedit" data-bef="${isText ? 'text' : 'desc'}" data-idx="${i}" value="${esc(isText ? (el.text || '') : (el.desc || ''))}" placeholder="${isText ? 'text to render…' : 'describe this element…'}" spellcheck="false"><button class="i4-boxtype" data-act="boxtype" data-idx="${i}" title="Toggle object / text">${isText ? 'T' : '▢'}</button>` | |
| : `<span class="i4-boxlab">${esc(String(elLabel(el, i)).slice(0, 60))}</span>`; | |
| return `<div class="i4-box ${isText ? 'text' : ''} ${i === selIdx ? 'sel' : ''}" data-idx="${i}" | |
| style="top:${b[0] / 10}%;left:${b[1] / 10}%;height:${(b[2] - b[0]) / 10}%;width:${(b[3] - b[1]) / 10}%"> | |
| ${inner}<span class="i4-handle" data-idx="${i}"></span></div>`; | |
| }).join(''); | |
| } | |
| function elEditorHTML(el){ | |
| return `<div class="i4-panel i4-eledit"> | |
| <label>Type</label> | |
| <select data-ef="type"><option value="obj" ${el.type !== 'text' ? 'selected' : ''}>Object</option><option value="text" ${el.type === 'text' ? 'selected' : ''}>Text</option></select> | |
| ${el.type === 'text' ? `<label>Text content</label><input type="text" data-ef="text" value="${esc(el.text || '')}">` : ''} | |
| <label>Description</label><textarea data-ef="desc" rows="2">${esc(el.desc || '')}</textarea> | |
| <label>Color palette (max 5)</label> | |
| <div class="i4-palrow"><input type="color" class="i4-color" data-pal="box"><button class="i4-mini" data-act="addcolor" data-pal="box">Add</button></div> | |
| <div class="i4-swatches">${swatchesHTML(el.color_palette, 'box')}</div> | |
| <button class="i4-mini danger" data-act="delbox">Delete element</button> | |
| </div>`; | |
| } | |
| function renderElList(){ | |
| const l = root.querySelector('#i4-ellist'); | |
| if (!l) return; | |
| if (!els().length) { l.innerHTML = ''; l.style.display = 'none'; return; } | |
| l.style.display = 'flex'; | |
| l.innerHTML = '<div class="i4-panel-title">Elements</div>' + els().map((el, i) => ` | |
| <div class="i4-elrow ${i === selIdx ? 'sel' : ''}" data-act="selel" data-idx="${i}"> | |
| <span class="i4-eltype ${el.type === 'text' ? 'text' : ''}">${el.type === 'text' ? 'T' : '▢'}</span> | |
| <span class="i4-elname">${esc(String(elLabel(el, i)).slice(0, 70))}</span> | |
| ${hasBbox(el) ? '' : `<button class="i4-mini" data-act="place" data-idx="${i}">📍 place</button>`} | |
| </div>${i === selIdx ? elEditorHTML(el) : ''}`).join(''); | |
| } | |
| function positionBox(i){ | |
| const d = root.querySelector(`.i4-box[data-idx="${i}"]`); | |
| const el = els()[i]; | |
| if (!d || !el) return; | |
| const b = el.bbox; | |
| d.style.top = b[0] / 10 + '%'; d.style.left = b[1] / 10 + '%'; | |
| d.style.height = (b[2] - b[0]) / 10 + '%'; d.style.width = (b[3] - b[1]) / 10 + '%'; | |
| } | |
| function select(i){ | |
| selIdx = i; | |
| renderCanvas(); | |
| renderElList(); | |
| } | |
| element.addEventListener('click', (e) => { | |
| const t = e.target.closest('[data-act]'); | |
| if (!t) return; | |
| const act = t.dataset.act; | |
| if (act === 'tab') { tab = t.dataset.tab; jsonMsg = ''; renderAll(); } | |
| else if (act === 'gen-image') { trigger('generate_image', { state_json: JSON.stringify(clean()) }); } | |
| else if (act === 'stylemode') { | |
| const m = t.dataset.m, s = sd(); | |
| if (m === 'photo' && !isPhoto()) { s.photo = s.photo || s.art_style || ''; delete s.art_style; s.medium = 'photograph'; } | |
| else if (m === 'art' && isPhoto()) { s.art_style = s.art_style || s.photo || ''; delete s.photo; if (s.medium === 'photograph') s.medium = ''; } | |
| renderAll(); | |
| } | |
| else if (act === 'addcolor') { | |
| const pal = t.dataset.pal; | |
| const picker = root.querySelector(`.i4-color[data-pal="${pal}"]`); | |
| if (!picker) return; | |
| const hex = picker.value.toUpperCase(); | |
| if (pal === 'global') { | |
| const arr = sd().color_palette = sd().color_palette || []; | |
| if (arr.length < 16 && !arr.includes(hex)) arr.push(hex); | |
| const sw = picker.closest('.i4-panel').querySelector('.i4-swatches'); | |
| if (sw) sw.innerHTML = swatchesHTML(arr, 'global'); | |
| } else if (selIdx !== null && els()[selIdx]) { | |
| const el = els()[selIdx]; | |
| const arr = el.color_palette = el.color_palette || []; | |
| if (arr.length < 5 && !arr.includes(hex)) arr.push(hex); | |
| renderElList(); | |
| } | |
| } | |
| else if (act === 'rmcolor') { | |
| const pal = t.dataset.pal, hex = t.dataset.hex; | |
| if (pal === 'global') { | |
| sd().color_palette = (sd().color_palette || []).filter(c => c !== hex); | |
| t.parentElement.innerHTML = swatchesHTML(sd().color_palette, 'global'); | |
| } else if (selIdx !== null && els()[selIdx]) { | |
| const el = els()[selIdx]; | |
| el.color_palette = (el.color_palette || []).filter(c => c !== hex); | |
| renderElList(); | |
| } | |
| } | |
| else if (act === 'selel') { | |
| const i = Number(t.dataset.idx); | |
| if (i === selIdx) { selIdx = null; renderCanvas(); renderElList(); } | |
| else select(i); | |
| } | |
| else if (act === 'boxtype') { | |
| const i = Number(t.dataset.idx); | |
| const el = els()[i]; | |
| if (el) { | |
| el.type = el.type === 'text' ? 'obj' : 'text'; | |
| renderCanvas(); | |
| renderElList(); | |
| const inp = root.querySelector(`.i4-box[data-idx="${i}"] .i4-boxedit`); | |
| if (inp) inp.focus(); | |
| } | |
| } | |
| else if (act === 'place') { | |
| const i = Number(t.dataset.idx); | |
| const el = els()[i]; | |
| if (el) { | |
| const o = (i % 5) * 40; | |
| el.bbox = [200 + o, 200 + o, 650 + o, 650 + o]; | |
| renderCanvas(); | |
| renderElList(); | |
| select(i); | |
| } | |
| } | |
| else if (act === 'delbox') { | |
| if (selIdx !== null) { els().splice(selIdx, 1); selIdx = null; renderCanvas(); renderElList(); } | |
| } | |
| else if (act === 'applyjson') { | |
| const ta = root.querySelector('.i4-json'); | |
| try { state = normalize(JSON.parse(ta.value)); selIdx = null; jsonMsg = '✓ applied'; } | |
| catch (err) { jsonMsg = '✗ ' + err.message; } | |
| renderAll(); | |
| } | |
| else if (act === 'copyjson') { | |
| navigator.clipboard && navigator.clipboard.writeText(serialize(2)); | |
| jsonMsg = '✓ copied'; | |
| const msg = root.querySelector('.i4-jsonmsg'); | |
| if (msg) msg.textContent = jsonMsg; | |
| } | |
| }); | |
| element.addEventListener('input', (e) => { | |
| const ds = e.target.dataset || {}; | |
| if (ds.f) { setField(ds.f, e.target.value); return; } | |
| const syncRow = (i, el) => { | |
| const row = root.querySelector(`.i4-elrow[data-idx="${i}"] .i4-elname`); | |
| if (row) row.textContent = String(elLabel(el, i)).slice(0, 70); | |
| }; | |
| if (ds.bef !== undefined) { | |
| const i = Number(ds.idx); | |
| const el = els()[i]; | |
| if (!el) return; | |
| el[ds.bef] = e.target.value; | |
| syncRow(i, el); | |
| const lf = root.querySelector(`.i4-eledit [data-ef="${ds.bef}"]`); | |
| if (lf) lf.value = e.target.value; | |
| return; | |
| } | |
| const ef = ds.ef; | |
| if (ef && selIdx !== null && els()[selIdx]) { | |
| const el = els()[selIdx]; | |
| if (ef === 'type') { el.type = e.target.value; renderCanvas(); renderElList(); } | |
| else { | |
| el[ef] = e.target.value; | |
| syncRow(selIdx, el); | |
| const inp = root.querySelector(`.i4-box[data-idx="${selIdx}"] .i4-boxedit[data-bef="${ef}"]`); | |
| if (inp) inp.value = e.target.value; | |
| } | |
| } | |
| }); | |
| element.addEventListener('pointerdown', (e) => { | |
| const canvas = root.querySelector('.i4-canvas'); | |
| if (!canvas || !canvas.contains(e.target)) return; | |
| if (e.target.closest('.i4-boxedit, .i4-boxtype')) return; | |
| e.preventDefault(); | |
| const r = canvas.getBoundingClientRect(); | |
| const px = (e.clientX - r.left) / r.width * 1000; | |
| const py = (e.clientY - r.top) / r.height * 1000; | |
| const handle = e.target.closest('.i4-handle'); | |
| const boxEl = e.target.closest('.i4-box'); | |
| if (handle) { | |
| const i = Number(handle.dataset.idx); | |
| select(i); | |
| drag = { kind: 'resize', i, px, py, b: els()[i].bbox.slice() }; | |
| } else if (boxEl) { | |
| const i = Number(boxEl.dataset.idx); | |
| select(i); | |
| drag = { kind: 'move', i, px, py, b: els()[i].bbox.slice() }; | |
| } else { | |
| const el = { type: 'obj', bbox: [Math.round(py), Math.round(px), Math.round(py), Math.round(px)], desc: '' }; | |
| els().push(el); | |
| select(els().length - 1); | |
| drag = { kind: 'draw', i: els().length - 1, px, py, b: el.bbox.slice() }; | |
| } | |
| }); | |
| window.addEventListener('pointermove', (e) => { | |
| if (!drag) return; | |
| const canvas = root.querySelector('.i4-canvas'); | |
| if (!canvas) return; | |
| const r = canvas.getBoundingClientRect(); | |
| const px = (e.clientX - r.left) / r.width * 1000; | |
| const py = (e.clientY - r.top) / r.height * 1000; | |
| const clamp = (v) => Math.max(0, Math.min(1000, Math.round(v))); | |
| const el = els()[drag.i]; | |
| if (!el) { drag = null; return; } | |
| const b = drag.b; | |
| if (drag.kind === 'draw') { | |
| el.bbox = [clamp(Math.min(drag.py, py)), clamp(Math.min(drag.px, px)), clamp(Math.max(drag.py, py)), clamp(Math.max(drag.px, px))]; | |
| } else if (drag.kind === 'move') { | |
| const h = b[2] - b[0], w = b[3] - b[1]; | |
| const ny = Math.max(0, Math.min(1000 - h, Math.round(b[0] + (py - drag.py)))); | |
| const nx = Math.max(0, Math.min(1000 - w, Math.round(b[1] + (px - drag.px)))); | |
| el.bbox = [ny, nx, ny + h, nx + w]; | |
| } else { | |
| el.bbox = [b[0], b[1], clamp(Math.max(b[0] + 10, b[2] + (py - drag.py))), clamp(Math.max(b[1] + 10, b[3] + (px - drag.px)))]; | |
| } | |
| positionBox(drag.i); | |
| }); | |
| window.addEventListener('pointerup', () => { | |
| if (!drag) return; | |
| const el = els()[drag.i]; | |
| if (drag.kind === 'draw' && el && ((el.bbox[2] - el.bbox[0]) < 15 || (el.bbox[3] - el.bbox[1]) < 15)) { | |
| els().splice(drag.i, 1); | |
| selIdx = null; | |
| renderCanvas(); | |
| renderElList(); | |
| } else if (drag.kind === 'draw') { | |
| const inp = root.querySelector(`.i4-box[data-idx="${drag.i}"] .i4-boxedit`); | |
| if (inp) inp.focus(); | |
| } | |
| drag = null; | |
| }); | |
| watch('value', () => { | |
| const s = JSON.stringify(props.value ?? null); | |
| if (s === lastPushed) return; | |
| lastPushed = s; | |
| state = normalize(props.value); | |
| selIdx = null; | |
| tab = 'visual'; | |
| renderAll(); | |
| }); | |
| watch(['image_url', 'img_width', 'img_height'], () => { if (tab === 'visual') renderCanvas(); }); | |
| renderAll(); | |
| """ | |
| class StudioEditor(gr.HTML): | |
| def __init__(self, value=None, **kwargs): | |
| # setdefault (not hardcode): gradio re-instantiates the class with updated props as kwargs | |
| # when an event returns gr.update(image_url=...), which would otherwise collide. | |
| kwargs.setdefault("image_url", "") | |
| kwargs.setdefault("img_width", 1024) | |
| kwargs.setdefault("img_height", 1024) | |
| kwargs.setdefault("html_template", '<div class="i4-root"></div>') | |
| kwargs.setdefault("css_template", EDITOR_CSS) | |
| kwargs.setdefault("js_on_load", EDITOR_JS) | |
| if value is None: | |
| value = json.loads(json.dumps(DEFAULT_PROMPT_JSON)) | |
| super().__init__(value=value, **kwargs) | |
| def api_info(self): | |
| return {"type": "object", "description": "Ideogram 4 structured JSON prompt"} | |
| CSS = """ | |
| .dark .gradio-container { color: var(--body-text-color); } | |
| """ | |
| with gr.Blocks(title="Ideogram 4 Studio") as demo: | |
| gr.Markdown( | |
| "# Ideogram 4 Studio\n" | |
| "A studio workspace for Ideogram's open-weights model: draft a structured JSON prompt from text " | |
| "(✨ Generate prompt), refine it in the visual or JSON editor, then render it (🎨 Generate image) — " | |
| "the image always renders the Studio JSON exactly as edited, and lands on the canvas so you can " | |
| "move, resize and re-describe each element and regenerate.\n\n" | |
| "[Model](https://huggingface.co/ideogram-ai/ideogram-4-nf4) · " | |
| "[Plain demo](https://huggingface.co/spaces/ideogram-ai/ideogram4) · " | |
| "[Blog](https://ideogram.ai/blog/ideogram-4.0/)" | |
| ) | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| prompt = gr.Textbox( | |
| label="Prompt", | |
| value="a ginger cat wearing a tiny wizard hat reading a spellbook", | |
| lines=3, | |
| info="✨ Generate prompt expands this into Ideogram's JSON caption and fills the editor.", | |
| ) | |
| upsampler = gr.Radio( | |
| choices=UPSAMPLERS, | |
| value=UPSAMPLER_REMOTE, | |
| label="Prompt enhancement", | |
| info="Which upsampler drafts the JSON caption that populates the editor.", | |
| ) | |
| gen_prompt_btn = gr.Button("✨ Generate prompt", variant="primary") | |
| gr.Examples( | |
| examples=[ | |
| ["a ginger cat wearing a tiny wizard hat reading a spellbook"], | |
| ["an isometric illustration of a tiny city floating in the clouds"], | |
| ["a movie poster for 'THE LAST SUMMER' with dramatic golden-hour lighting"], | |
| ], | |
| inputs=[prompt], | |
| ) | |
| with gr.Column(scale=2): | |
| editor = StudioEditor() | |
| with gr.Row(): | |
| mode = gr.Radio( | |
| choices=list(MODES.keys()), value="Default · 20 steps", label="Mode (speed ↔ quality)", scale=2 | |
| ) | |
| with gr.Accordion("Advanced", open=False): | |
| with gr.Row(): | |
| 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=True) | |
| with gr.Accordion("Generated image (full resolution)", open=False): | |
| out_image = gr.Image(label="Generated image", type="pil", interactive=False, show_label=False) | |
| gen_prompt_btn.click( | |
| generate_prompt, | |
| inputs=[prompt, upsampler, width, height], | |
| outputs=[editor], | |
| ) | |
| editor.generate_image( | |
| generate, | |
| inputs=[mode, width, height, seed, randomize], | |
| outputs=[editor, out_image, seed], | |
| ) | |
| width.change(lambda v: gr.update(img_width=int(v)), width, editor, show_progress="hidden") | |
| height.change(lambda v: gr.update(img_height=int(v)), height, editor, show_progress="hidden") | |
| demo.launch(theme=gr.themes.Citrus(), css=CSS) | |