Dynamic GPU duration (per-step model); move remote upsample off-GPU

app.py

@@ -126,52 +126,48 @@ def remote_upsample(prompt, width, height):
     return json.dumps(jp, ensure_ascii=False, separators=(",", ":"))
 
 
-@spaces.GPU(duration=120, size="xlarge")
-def generate(
-    prompt,
-    mode="Default · 20 steps",
-    upsampler=UPSAMPLERS[0],
-    width=1024,
-    height=1024,
-    seed=0,
-    randomize_seed=False,
-    progress=gr.Progress(track_tqdm=True),
-):
-    t_enter = time.perf_counter()
-    if randomize_seed or seed < 0:
-        seed = random.randint(0, MAX_SEED)
-
-    # Overlap the AOTI block-patch with upsampling: the transformer is idle while we upsample.
+# --- 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.
+# Remote upsample is a network call done OFF the GPU (in `generate`), so it isn't budgeted here.
+_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(30, int(math.ceil(budget * DURATION_MARGIN)))
+
+
+@spaces.GPU(duration=_gpu_duration, size="xlarge")
+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()
-
-    # Always upsample. Prefer Ideogram's hosted magic-prompt; fall back to the local Qwen graft on any failure.
-    use_remote = upsampler == UPSAMPLERS[0] and bool(IDEOGRAM_API_KEY)
-    final_prompt = prompt
-    if use_remote:
-        progress(0.0, desc="✍️ Upsampling (Ideogram)…")
-        t = time.perf_counter()
-        try:
-            final_prompt = remote_upsample(prompt, int(width), int(height))
-            print(f"[timing] upsample remote: {time.perf_counter() - t:.2f}s", flush=True)
-        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.")
-            use_remote = False
-    if not use_remote:
+    if do_local:
         progress(0.0, desc="✍️ Upsampling (local Qwen)…")
         t = time.perf_counter()
         try:
             final_prompt = pipe.upsample_prompt(
-                prompt, height=int(height), width=int(width), lm_head_repo_id=LM_HEAD_REPO
+                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
-    print(f"[timing] pre-diffusion (enter -> ready): {time.perf_counter() - t_enter:.2f}s", flush=True)
 
     progress(0.0, desc="🎨 Generating image…")
     generator = torch.Generator(device="cuda").manual_seed(int(seed))
@@ -184,7 +180,36 @@ def generate(
         caption = json.loads(final_prompt)
     except Exception:
         caption = {"prompt": final_prompt}
-    return image, seed, caption
+    return image, int(seed), caption
+
+
+def generate(
+    prompt,
+    mode="Default · 20 steps",
+    upsampler=UPSAMPLERS[0],
+    width=1024,
+    height=1024,
+    seed=0,
+    randomize_seed=False,
+    progress=gr.Progress(track_tqdm=True),
+):
+    if randomize_seed or seed < 0:
+        seed = random.randint(0, MAX_SEED)
+
+    # Remote upsample is a network call -> run it here, OFF the GPU. Fall back to local (on-GPU) on failure.
+    final_prompt, do_local = prompt, True
+    if upsampler == UPSAMPLERS[0] and IDEOGRAM_API_KEY:
+        progress(0.0, desc="✍️ Upsampling (Ideogram)…")
+        t = time.perf_counter()
+        try:
+            final_prompt = remote_upsample(prompt, int(width), int(height))
+            do_local = False
+            print(f"[timing] upsample remote (off-GPU): {time.perf_counter() - t:.2f}s", flush=True)
+        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.")
+
+    return _gpu_generate(final_prompt, mode, width, height, seed, do_local)
 
 
 @spaces.GPU(size="xlarge")
@@ -201,10 +226,8 @@ try:
 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)
 
-CSS='''
-.dark .gradio-container { color: var(--body-text-color); }
-'''
-with gr.Blocks(theme=gr.themes.Citrus(), title="Ideogram 4", css=CSS) as demo:
+
+with gr.Blocks(theme=gr.themes.Citrus(), title="Ideogram 4") as demo:
     gr.Markdown(
         "# Ideogram 4\n"
         "Ideogram's first open-weights model — a 9.3B-parameter text-to-image foundation model at the "
@@ -255,4 +278,4 @@ with gr.Blocks(theme=gr.themes.Citrus(), title="Ideogram 4", css=CSS) as demo:
         outputs=[out_image, seed, out_caption],
     )
 
-demo.launch()
+demo.queue().launch()