Report a time-based brew estimate so the progress bar moves on ZeroGPU (Stage 3)

The @spaces.GPU worker is a separate process, so it can't push per-chunk progress back to /api/progress - the brew bar froze at 0% then jumped to 100% on the Space. /api/progress now returns a smooth time-based estimate (fractional done, capped at 92%) computed from the brew's start + a rough per-length budget, snapping to 100% when the tape lands. Local/stub paths keep their real per-chunk progress (the estimate only kicks in while a ZeroGPU brew is active).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

app.py

@@ -23,6 +23,7 @@ import io
 import json
 import os
 import threading
+import time
 import wave
 from pathlib import Path
 
@@ -81,6 +82,11 @@ ALLOWED_SECONDS = (30, 60, 90)
 DEFAULT_SECONDS = int(os.getenv("LOFINITY_DURATION", "30"))
 OLLAMA_URL = os.getenv("OLLAMA_URL", "http://localhost:11434")
 OLLAMA_MODEL = os.getenv("OLLAMA_MODEL", "llama3.2:3b")
+# A ZeroGPU brew renders in a separate GPU worker process, so /api/progress can't
+# read real per-chunk progress from it; it reports a time-based estimate from this
+# rough per-length budget instead (tunable; only affects the cosmetic brew bar).
+GPU_WARMUP_S = 5.0  # enrichment + cold-start/queue allowance before audio flows
+GPU_SECS_PER_CHUNK = 10.0  # rough GPU render time per 30s chunk
 
 app = Server(title="LoFinity")
 
@@ -89,6 +95,10 @@ app = Server(title="LoFinity")
 # frontend polls /api/progress to fill its brewing bar.
 _PROGRESS = {"done": 0, "total": 1}
 
+# A ZeroGPU brew's wall-clock start + estimated total, so /api/progress can report
+# a smooth time-based estimate while the GPU worker is busy (see progress()).
+_BREW = {"active": False, "start": 0.0, "est": 1.0, "total": 1}
+
 # --- prompt enrichment --------------------------------------------------------
 
 ENRICH_SYSTEM = """\
@@ -419,9 +429,10 @@ def _gpu_budget(prompt: str, seconds: int = CHUNK_S) -> int:
 def gpu_brew(prompt: str, seconds: int = CHUNK_S) -> tuple:
     """ZeroGPU entry point — enrichment (MiniCPM) AND MusicGen on the real GPU in
     a single acquisition. Takes the raw vibe and returns
-    (music_prompt, title, bed, samples, rate). No progress_cb: this body runs in
-    a separate GPU worker, so _PROGRESS can't reach /api/progress yet (Stage 3) —
-    the brewing garden jumps to done on the Space. This path is Space-only."""
+    (music_prompt, title, bed, samples, rate). It runs in a separate GPU worker
+    process and can't push per-chunk progress back to the web process, so
+    /api/progress reports a time-based estimate for the bar. This path is
+    Space-only."""
     music_prompt, title, bed = enrich_prompt(prompt)
     samples, rate = musicgen_engine(music_prompt, seconds)
     return music_prompt, title, bed, samples, rate
@@ -463,11 +474,20 @@ def generate_song(prompt: str, seconds: int = DEFAULT_SECONDS) -> dict:
     _PROGRESS.update(done=0, total=chunks)
 
     if IS_ZEROGPU and ENGINE != "stub":
-        # On ZeroGPU both enrichment (MiniCPM) and MusicGen need the real GPU, so
-        # they share ONE @spaces.GPU acquisition. That worker runs in a separate
-        # process, so progress can't stream back yet (Stage 3): the bar jumps.
-        print(f"[lofinity] brewing on GPU :: {prompt!r} ({seconds}s)")
-        music_prompt, title, bed, samples, rate = gpu_brew(prompt, seconds)
+        # On ZeroGPU enrichment (MiniCPM) and MusicGen share ONE @spaces.GPU
+        # acquisition in a separate worker process, which can't push real progress
+        # back — so /api/progress reports a smooth time-based ESTIMATE driven by
+        # this brew's start + budget (capped <100% until the tape actually lands).
+        est = GPU_WARMUP_S + GPU_SECS_PER_CHUNK * chunks
+        _BREW.update(active=True, start=time.monotonic(), est=est, total=chunks)
+        print(f"[lofinity] brewing on GPU :: {prompt!r} ({seconds}s, ~{est:.0f}s est)")
+        try:
+            music_prompt, title, bed, samples, rate = gpu_brew(prompt, seconds)
+        finally:
+            # top the bar off BEFORE clearing active, so a poll landing in between
+            # reads 100% (from _PROGRESS), never the 0% this brew started at
+            _PROGRESS.update(done=chunks, total=chunks)
+            _BREW.update(active=False)
         print(f"[lofinity] brewed {title!r} :: {music_prompt} [+ {bed}]")
     else:
         # Local / stub: enrich in-process (Ollama or fallback), then render with
@@ -493,7 +513,15 @@ def generate_song(prompt: str, seconds: int = DEFAULT_SECONDS) -> dict:
 
 @app.get("/api/progress")
 def progress() -> dict:
-    """Chunks finished / total for the brew currently running, for the UI bar."""
+    """Progress for the brewing bar. Local/stub report real per-chunk progress via
+    _PROGRESS. A ZeroGPU brew runs in a separate GPU worker that can't push
+    progress back, so report a smooth time-based ESTIMATE instead: a fractional
+    `done` (the frontend fills the bar to done/total) capped below 100% until the
+    real tape lands and _PROGRESS tops it off."""
+    if _BREW["active"] and _BREW["est"] > 0:
+        elapsed = time.monotonic() - _BREW["start"]
+        frac = min(0.92, elapsed / _BREW["est"])
+        return {"done": round(frac * _BREW["total"], 3), "total": _BREW["total"]}
     return dict(_PROGRESS)