fix: prevent OOM kills during background removal on CPU

Skip alpha matting on CPU (pymatting's sparse matrices are the main
memory hog), auto-downscale images above 2048px before sending to
rembg, and retry with the lighter u2net model when OOM is detected.
This commit is contained in:
SnapOtter
2026-04-28 02:20:48 +08:00
parent 4f6fd707a1
commit c6c78dc76f
3 changed files with 74 additions and 14 deletions
+9 -4
View File
@@ -94,18 +94,23 @@ def main():
input_data = f.read()
emit_progress(30, "Analyzing image")
use_alpha_matting = device != "cpu"
try:
output_data = remove(
input_data,
session=session,
alpha_matting=True,
alpha_matting=use_alpha_matting,
alpha_matting_foreground_threshold=240,
alpha_matting_background_threshold=10,
)
except Exception as e:
raise RuntimeError(
f"Alpha matting failed: {e}. Try again without alpha matting or with a different model."
) from 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
emit_progress(80, "Background removed")