Update app.py

app.py

@@ -34,8 +34,64 @@ def _save_image(image) -> dict:
     }
 
 
-@spaces.GPU
-def generate_image(
+def _estimate_duration(prompt, height, width, num_inference_steps, seed, randomize_seed) -> int:
+    """Rough wall-clock estimate (seconds) for one Z-Image-Turbo call.
+
+    ZeroGPU's default per-call duration is 60s. Requesting less than you need
+    *raises* queue priority (shorter tasks get scheduled sooner) and — crucially
+    for a busy shared Space — frees the GPU slot for the next visitor far
+    sooner than holding it for a full minute, so far fewer users hit the
+    Space's "reached its GPU limit" rejection. Scaled by pixel count and steps;
+    clamped to a small floor/ceiling so a runaway slider can't starve the queue
+    or under-budget a big call.
+
+    Signature mirrors the GPU function exactly because @spaces.GPU passes the
+    decorated function's inputs straight through to the duration callable.
+
+    See: https://huggingface.co/docs/hub/en/spaces-zerogpu#duration-management
+    """
+    pixels = max(int(height), 1) * max(int(width), 1)
+    # ~0.4s/step at 1024^2, linear-ish in pixels. Big calls still need headroom.
+    per_step = 0.4 * (pixels / (1024 * 1024))
+    seconds = int(num_inference_steps) * per_step
+    return max(20, min(int(seconds) + 15, 120))
+
+
+def _friendly_gpu_error(err: Exception) -> str:
+    """Turn ZeroGPU's terse allocator rejections into a clear, honest message.
+
+    'Space app has reached its GPU limit' is a *Space-level* capacity rejection
+    (the shared ZeroGPU pool is saturated), not a per-user quota wall — it
+    reproduces regardless of inputs, account tier, or sign-in state. Don't make
+    an upgrade claim whose truth we can't pin down, so the message is neutral:
+    shared GPU at capacity, retry shortly.
+    """
+    msg = (str(err) or "").lower()
+    capacity_hints = (
+        "gpu limit", "reached its gpu limit", "gpu quota", "out of quota",
+        "quota", "no gpu", "could not allocate", "gpu is busy", "too many",
+        "concurrent",
+    )
+    if any(h in msg for h in capacity_hints):
+        return (
+            "⛔ This demo's shared GPU is at capacity right now — it's not a "
+            "problem with your prompt or your account. The GPU pool is fully "
+            "booked by other users at the moment. Please wait a minute and "
+            "retry; demand clears between bursts."
+        )
+    if "out of memory" in msg or "oom" in msg or "cuda" in msg:
+        return (
+            "💥 Image generation ran out of GPU memory. Try a smaller "
+            "Height/Width or fewer Inference Steps, then retry."
+        )
+    return (
+        "⚠️ Image generation failed. Please try again in a moment — if it "
+        "keeps happening, simplify your prompt or lower the resolution."
+    )
+
+
+@spaces.GPU(duration=_estimate_duration)
+def _generate_image_gpu(
     prompt: str,
     height: int,
     width: int,
@@ -43,15 +99,9 @@ def generate_image(
     seed: int,
     randomize_seed: bool,
 ):
-    """Generate an image from a prompt using Z-Image-Turbo on ZeroGPU.
-
-    Bound to the workflow canvas as a `fn` operator node — the workflow
-    calls this Python function directly server-side, so the entire pipeline
-    (frontend + ZeroGPU) lives in a single Space.
-
-    Returns (image_dict, seed_used). The image is serialized to a /gradio_api
-    file URL so JSON serialization across the fn bridge succeeds; the executor's
-    `fromGradioOutput` turns the dict back into an image port value.
+    """The GPU-decorated worker. Runs only under a ZeroGPU allocation; the
+    allocator raises *before* this body if no GPU can be granted, which is why
+    the rewording lives in the plain `generate_image` wrapper below, not here.
     """
     if not prompt or not prompt.strip():
         raise gr.Error("Please enter a prompt.")
@@ -72,7 +122,40 @@ def generate_image(
     return _save_image(image), int(seed)
 
 
-# The workflow (workflow.json) wires this function as a `fn` operator:
+def generate_image(
+    prompt: str,
+    height: int,
+    width: int,
+    num_inference_steps: int,
+    seed: int,
+    randomize_seed: bool,
+):
+    """Workflow-facing wrapper around the GPU worker.
+
+    Bound to the canvas as a `fn` operator node — the workflow calls this
+    Python function directly server-side, so the entire pipeline (frontend +
+    ZeroGPU) lives in a single Space. This non-GPU wrapper catches rejections
+    from the `@spaces.GPU` allocator (which fire before the worker body runs)
+    and rewords them into a clear, honest user-facing message.
+
+    Returns (image_dict, seed_used). The image is serialized to a /gradio_api
+    file URL so JSON serialization across the fn bridge succeeds; the
+    executor's `fromGradioOutput` turns the dict back into an image port value.
+    """
+    try:
+        return _generate_image_gpu(
+            prompt, height, width, num_inference_steps, seed, randomize_seed
+        )
+    except gr.Error:
+        # Already a user-facing validation message (e.g. empty prompt) — pass
+        # it through unchanged.
+        raise
+    except Exception as e:
+        # Allocator rejection (GPU limit / quota / OOM / etc.) — reword.
+        raise gr.Error(_friendly_gpu_error(e)) from e
+
+
+# The workflow (workflow.json) wires `generate_image` as a `fn` operator:
 #   Prompt, Height, Width, Inference Steps, Seed, Randomize Seed ─▶
 #   generate_image (fn operator, kind="fn") ─▶ Output Image, Seed Used
 #