Frank Schwenk
5a3118f590
env_utils.env_bool replaces duplicated truthy parsing in config and rembg_cli. logging_config.configure_logging unifies entry-point setup. Co-authored-by: Cursor <cursoragent@cursor.com>
58 lines
1.8 KiB
Python
58 lines
1.8 KiB
Python
"""Minimal `rembg` runner, invoked as its own subprocess.
|
|
|
|
We deliberately don't shell out to the official `rembg` CLI: that needs
|
|
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
|
|
|
|
from .env_utils import env_bool
|
|
|
|
|
|
def main() -> int:
|
|
if len(sys.argv) != 3:
|
|
print("usage: python -m app.rembg_cli <input> <output>", file=sys.stderr)
|
|
return 2
|
|
|
|
# Keep model-download progress bars from polluting error parsing upstream.
|
|
os.environ.setdefault("TQDM_DISABLE", "1")
|
|
|
|
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__":
|
|
raise SystemExit(main())
|