mirror of
https://github.com/snapotter-hq/SnapOtter.git
synced 2026-08-03 07:46:42 +02:00
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
This commit is contained in:
@@ -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)
|
||||
Reference in New Issue
Block a user