mirror of
https://github.com/snapotter-hq/SnapOtter.git
synced 2026-08-03 07:46:42 +02:00
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
36 lines
1.2 KiB
Python
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)
|