feat: harden pipeline UX — preprocess, rembg alpha, remix/lightbox
Downscale to 2000px JPEG before rembg, random blend opacity, timeout retries, per-variant remix prefill, lightbox, and longer SFTP idle. Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
+37
-6
@@ -5,26 +5,57 @@ the `rembg[cli]` extra (aiohttp, gradio, watchdog, ...) just to run a
|
||||
single background removal. Calling the `remove()` API directly from a
|
||||
tiny script keeps the image smaller while still giving pipeline.py a
|
||||
subprocess boundary to apply `nice` and a hard timeout to.
|
||||
|
||||
Defaults target the event box (≈8GB RAM): u2net + alpha matting.
|
||||
`birefnet-general` (~928MB weights) OOMs here; override via REMBG_MODEL
|
||||
/ REMBG_ALPHA if the host has more headroom.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import sys
|
||||
import traceback
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
def _env_bool(name: str, default: bool) -> bool:
|
||||
raw = os.environ.get(name)
|
||||
if raw is None:
|
||||
return default
|
||||
return raw.strip().lower() in {"1", "true", "yes", "on"}
|
||||
|
||||
|
||||
def main() -> int:
|
||||
if len(sys.argv) != 3:
|
||||
print("usage: python -m app.rembg_cli <input> <output>", file=sys.stderr)
|
||||
return 2
|
||||
|
||||
from rembg import remove
|
||||
# Keep model-download progress bars from polluting error parsing upstream.
|
||||
os.environ.setdefault("TQDM_DISABLE", "1")
|
||||
|
||||
input_path = Path(sys.argv[1])
|
||||
output_path = Path(sys.argv[2])
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
output_path.write_bytes(remove(input_path.read_bytes()))
|
||||
return 0
|
||||
try:
|
||||
from rembg import new_session, remove
|
||||
|
||||
input_path = Path(sys.argv[1])
|
||||
output_path = Path(sys.argv[2])
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
model = os.environ.get("REMBG_MODEL", "u2net")
|
||||
alpha = _env_bool("REMBG_ALPHA", True)
|
||||
session = new_session(model)
|
||||
output_path.write_bytes(
|
||||
remove(
|
||||
input_path.read_bytes(),
|
||||
session=session,
|
||||
alpha_matting=alpha,
|
||||
)
|
||||
)
|
||||
return 0
|
||||
except Exception as exc: # noqa: BLE001 - surface real cause to parent process
|
||||
print(f"{type(exc).__name__}: {exc}", file=sys.stderr)
|
||||
traceback.print_exc(file=sys.stderr)
|
||||
return 1
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
Reference in New Issue
Block a user