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(
-                input_data,
-                session=session,
-                alpha_matting=use_alpha_matting,
-                alpha_matting_foreground_threshold=240,
-                alpha_matting_background_threshold=10,
-            )
-        except Exception as e:
-            if use_alpha_matting:
-                emit_progress(35, "Retrying without alpha matting")
-                output_data = remove(input_data, session=session, alpha_matting=False)
-            else:
-                raise RuntimeError(
-                    f"Background removal failed: {e}"
-                ) from e
+
+        # remove() runs the whole model in one opaque call with no per-step
+        # callback, so advance the bar in the background to show the job is
+        # alive instead of freezing at 30% (#591).
+        from progress_heartbeat import run_with_heartbeat
+
+        def _remove():
+            try:
+                return remove(
+                    input_data,
+                    session=session,
+                    alpha_matting=use_alpha_matting,
+                    alpha_matting_foreground_threshold=240,
+                    alpha_matting_background_threshold=10,
+                )
+            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:
-                        output_array, _ = upsampler.enhance(img_array, outscale=scale)
-                    except RuntimeError as oom_err:
-                        if "out of memory" not in str(oom_err).lower():
-                            raise
-                        torch.cuda.empty_cache()
-                        print(
-                            f"[upscale] OOM with tile={tile_size}, retrying with tile=256",
-                            file=sys.stderr,
-                            flush=True,
-                        )
-                        emit_progress(35, "Retrying with smaller tiles")
-                        upsampler.tile = 256
-                        output_array, _ = upsampler.enhance(img_array, outscale=scale)
+                    # enhance() runs the whole model in one opaque call with no
+                    # per-tile callback, so advance the bar in the background to
+                    # show the job is alive instead of freezing at 30% (#591).
+                    from progress_heartbeat import run_with_heartbeat
+
+                    def _enhance():
+                        try:
+                            return upsampler.enhance(img_array, outscale=scale)
+                        except RuntimeError as oom_err:
+                            if "out of memory" not in str(oom_err).lower():
+                                raise
+                            torch.cuda.empty_cache()
+                            print(
+                                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)