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