Files
SnapOtter/packages/ai/python/progress_heartbeat.py
T
SnapOtterandGitHub e56edc659f 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
2026-07-21 19:21:45 +08:00

36 lines
1.2 KiB
Python

"""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)