External/SnapOtter
1
0
Fork 0
mirror of https://github.com/snapotter-hq/SnapOtter.git synced 2026-08-03 07:46:42 +02:00
Code Issues Packages Projects Releases Wiki Activity

fix(ai): advance the progress bar during upscale and background removal (#608)

Browse Source 
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
Download Patch File Download Diff File Expand all files Collapse all files
packages/ai/python/upscale.py
+23 -14
View File
@@ -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)