xlarge GPU + bf16 transformers + upsampling default-on + gr.JSON caption + seed-in-field + non-fatal warmup

app.py

@@ -7,6 +7,7 @@ os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
 _HERE = os.path.dirname(os.path.abspath(__file__))
 sys.path.insert(0, os.path.join(_HERE, "diffusers_src", "src"))
 
+import json
 import random
 from typing import List, Literal, Union
 
@@ -159,7 +160,7 @@ def upsample_prompt(prompt: str, width: int, height: int) -> str:
     )[0].strip()
 
 
-@spaces.GPU(duration=240)
+@spaces.GPU(duration=240, size="xlarge")
 def generate(
     prompt: str,
     mode: str,
@@ -179,6 +180,7 @@ def generate(
             gr.Warning("`outlines` is not installed — upsampling without structural constraints.")
         final_prompt = upsample_prompt(prompt, int(width), int(height))
 
+    _ensure_bf16_transformers()
     generator = torch.Generator(device="cuda").manual_seed(int(seed))
     preset = MODES.get(mode, MODES["Default · 20 steps"])
     image = pipe(
@@ -188,14 +190,19 @@ def generate(
         generator=generator,
         **preset,
     ).images[0]
-    return image, seed, final_prompt
+    try:
+        caption = json.loads(final_prompt)
+    except Exception:
+        caption = {"prompt": final_prompt}
+    return image, seed, caption
 
 
-@spaces.GPU
+@spaces.GPU(size="xlarge")
 def _warmup():
     """Force the upsampler + pipeline onto GPU and warm their kernels at STARTUP, so request #1
     isn't slow. On ZeroGPU, module-level loading is CPU-only; GPU placement + JIT warmup otherwise
     happen on the first request."""
+    _ensure_bf16_transformers()
     try:
         if ENHANCER is not None:
             upsample_prompt("a red apple on a wooden table", 1024, 1024)
@@ -210,7 +217,10 @@ def _warmup():
         print(f"[warmup] pipeline warmup skipped: {e!r}", flush=True)
 
 
-_warmup()
+try:
+    _warmup()
+except Exception as e:  # a flaky ZeroGPU worker (e.g. ECC) must not take down the Space
+    print(f"[warmup] failed (will warm lazily on first request): {e!r}", flush=True)
 
 
 with gr.Blocks(theme=gr.themes.Citrus(), title="Ideogram 4 (NF4) — diffusers preview") as demo:
@@ -238,7 +248,7 @@ with gr.Blocks(theme=gr.themes.Citrus(), title="Ideogram 4 (NF4) — diffusers p
             with gr.Accordion("Advanced", open=False):
                 enhance = gr.Checkbox(
                     label="Prompt upsampling (Outlines)",
-                    value=False,
+                    value=True,
                     info="Rewrite the prompt into Ideogram's native JSON caption before generating."
                     + ("" if OUTLINES_AVAILABLE else "  ⚠ outlines not installed — runs unconstrained."),
                 )
@@ -250,16 +260,12 @@ with gr.Blocks(theme=gr.themes.Citrus(), title="Ideogram 4 (NF4) — diffusers p
                     randomize = gr.Checkbox(label="Randomize seed", value=True)
         with gr.Column():
             out_image = gr.Image(label="Output", type="pil")
-            out_seed = gr.Number(label="Seed used", interactive=False, precision=0)
-            out_caption = gr.Textbox(
-                label="Caption fed to the model (upsampled when enabled)",
-                lines=4,
-            )
+            out_caption = gr.JSON(label="Caption fed to the model (upsampled when enabled)")
 
     run.click(
         generate,
         inputs=[prompt, mode, enhance, width, height, seed, randomize],
-        outputs=[out_image, out_seed, out_caption],
+        outputs=[out_image, seed, out_caption],
     )
 
 demo.queue().launch()