fix(ai): advance the progress bar during upscale and background removal (#608)

RealESRGAN's enhance() and rembg's remove() run in one opaque call, so the
progress bar froze at 30% for the whole inference. Add a time-based
heartbeat that advances the bar in a background thread while the model runs
and stops when it returns, so the bar moves instead of freezing. Verified
end to end on a GPU box (forced CPU): a 55s upscale emitted 26 steady ticks
then completed; background removal too.

Fixes #591

packages/ai/python/progress_heartbeat.py

@@ -0,0 +1,35 @@
+"""Slowly-rising progress heartbeat for opaque model calls.
+
+RealESRGAN and rembg run the model in a single call with no per-step
+callback, so the progress bar would otherwise sit frozen for the whole
+inference. run_with_heartbeat advances the bar in a background thread while
+the call runs, purely to show the job is still alive, and stops the moment
+the call returns (or raises). It never reaches ``end``; the caller emits the
+real completion value once the work is done.
+"""
+import threading
+
+
+def run_with_heartbeat(fn, emit, start, end, stage, interval=2.0):
+    """Run ``fn()`` while emitting rising progress via ``emit(pct, stage)``.
+
+    Advances one percent every ``interval`` seconds from ``start`` toward
+    ``end`` (never past ``end - 1``). Returns ``fn()``'s value and propagates
+    any exception unchanged.
+    """
+    stop = threading.Event()
+
+    def beat():
+        pct = start
+        while not stop.wait(interval):
+            if pct < end - 1:
+                pct += 1
+                emit(pct, stage)
+
+    thread = threading.Thread(target=beat, daemon=True)
+    thread.start()
+    try:
+        return fn()
+    finally:
+        stop.set()
+        thread.join(timeout=1.0)

packages/ai/python/remove_bg.py

@@ -234,22 +234,27 @@ def main():
 
         emit_progress(30, "Analyzing image")
         use_alpha_matting = device != "cpu"
-        try:
+
-            output_data = remove(
+        # remove() runs the whole model in one opaque call with no per-step
-                input_data,
+        # callback, so advance the bar in the background to show the job is
-                session=session,
+        # alive instead of freezing at 30% (#591).
-                alpha_matting=use_alpha_matting,
+        from progress_heartbeat import run_with_heartbeat
-                alpha_matting_foreground_threshold=240,
+
-                alpha_matting_background_threshold=10,
+        def _remove():
-            )
+            try:
-        except Exception as e:
+                return remove(
-            if use_alpha_matting:
+                    input_data,
-                emit_progress(35, "Retrying without alpha matting")
+                    session=session,
-                output_data = remove(input_data, session=session, alpha_matting=False)
+                    alpha_matting=use_alpha_matting,
-            else:
+                    alpha_matting_foreground_threshold=240,
-                raise RuntimeError(
+                    alpha_matting_background_threshold=10,
-                    f"Background removal failed: {e}"
+                )
-                ) from e
+            except Exception as e:
+                if use_alpha_matting:
+                    return remove(input_data, session=session, alpha_matting=False)
+                raise RuntimeError(f"Background removal failed: {e}") from e
+
+        output_data = run_with_heartbeat(_remove, emit_progress, 30, 80, "Analyzing image")
 
         emit_progress(80, "Background removed")
 

packages/ai/python/upscale.py

@@ -144,20 +144,29 @@ def main():
                     img_array = np.array(img.convert("RGB"))
                     emit_progress(30, "Enhancing image with AI")
 
-                    try:
+                    # enhance() runs the whole model in one opaque call with no
-                        output_array, _ = upsampler.enhance(img_array, outscale=scale)
+                    # per-tile callback, so advance the bar in the background to
-                    except RuntimeError as oom_err:
+                    # show the job is alive instead of freezing at 30% (#591).
-                        if "out of memory" not in str(oom_err).lower():
+                    from progress_heartbeat import run_with_heartbeat
-                            raise
+
-                        torch.cuda.empty_cache()
+                    def _enhance():
-                        print(
+                        try:
-                            f"[upscale] OOM with tile={tile_size}, retrying with tile=256",
+                            return upsampler.enhance(img_array, outscale=scale)
-                            file=sys.stderr,
+                        except RuntimeError as oom_err:
-                            flush=True,
+                            if "out of memory" not in str(oom_err).lower():
-                        )
+                                raise
-                        emit_progress(35, "Retrying with smaller tiles")
+                            torch.cuda.empty_cache()
-                        upsampler.tile = 256
+                            print(
-                        output_array, _ = upsampler.enhance(img_array, outscale=scale)
+                                f"[upscale] OOM with tile={tile_size}, retrying with tile=256",
+                                file=sys.stderr,
+                                flush=True,
+                            )
+                            upsampler.tile = 256
+                            return upsampler.enhance(img_array, outscale=scale)
+
+                    output_array, _ = run_with_heartbeat(
+                        _enhance, emit_progress, 30, 80, "Enhancing image with AI"
+                    )
 
                     emit_progress(80, "AI enhancement complete")
                     result = Image.fromarray(output_array)

tests/unit/ai/progress-heartbeat.test.ts

@@ -0,0 +1,54 @@
+import { spawnSync } from "node:child_process";
+import { dirname, resolve } from "node:path";
+import { fileURLToPath } from "node:url";
+import { describe, expect, it } from "vitest";
+import { hasPython, pythonBin } from "../../helpers/python-gate.js";
+
+const here = dirname(fileURLToPath(import.meta.url));
+const pyDir = resolve(here, "../../../packages/ai/python");
+
+// The heartbeat is pure stdlib (threading), so this runs anywhere python3 is
+// present, no AI bundle required. It guards the contract the upscale / remove-bg
+// scripts rely on: advance the bar during an opaque model call, stop the moment
+// it returns, pass the return value through, and propagate exceptions (#591).
+const SCRIPT = `
+import sys, time
+sys.path.insert(0, ${JSON.stringify(pyDir)})
+from progress_heartbeat import run_with_heartbeat
+
+emitted = []
+def emit(pct, stage): emitted.append(pct)
+
+def slow(): time.sleep(1.6); return "RESULT"
+value = run_with_heartbeat(slow, emit, 30, 80, "Working", interval=0.5)
+assert value == "RESULT", "must return the wrapped call's value"
+assert emitted, "must emit at least one tick during a slow call"
+assert emitted[0] == 31, ("first tick is start+1", emitted)
+assert emitted == sorted(emitted), ("monotonic", emitted)
+assert max(emitted) <= 79, ("never reaches end", emitted)
+
+before = len(emitted)
+time.sleep(1.0)
+assert len(emitted) == before, "heartbeat must stop once the call returns"
+
+def boom(): raise ValueError("boom")
+try:
+    run_with_heartbeat(boom, emit, 30, 80, "Working", interval=0.2)
+    raise SystemExit("exception was swallowed")
+except ValueError:
+    pass
+
+print("OK")
+`;
+
+describe.skipIf(!hasPython)("progress heartbeat", () => {
+  it("advances during a slow call, stops after, returns value, propagates errors", () => {
+    const res = spawnSync(pythonBin as string, ["-c", SCRIPT], {
+      encoding: "utf8",
+      timeout: 30_000,
+    });
+    expect(res.stderr).toBe("");
+    expect(res.status).toBe(0);
+    expect(res.stdout).toContain("OK");
+  }, 40_000);
+});