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:
SnapOtter
2026-07-21 19:21:45 +08:00
committed by GitHub
parent e3c93333ea
commit e56edc659f
4 changed files with 133 additions and 30 deletions
+21 -16
View File
@@ -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")