fix: use BiRefNet-Lite as default model and fix JSON parsing

- Switch default from birefnet-general (973MB, 4min) to
  birefnet-general-lite (faster, still SOTA quality)
- Fix Python script stdout pollution — progress messages now go to
  stderr so the JSON result parser doesn't break
- Pre-download birefnet-general-lite in Docker build
This commit is contained in:
Siddharth Kumar Sah
2026-03-22 20:07:39 +08:00
parent 77ee8469d8
commit 2bdc367767
3 changed files with 13 additions and 10 deletions
+10 -7
View File
@@ -8,25 +8,26 @@ def main():
output_path = sys.argv[2]
settings = json.loads(sys.argv[3]) if len(sys.argv) > 3 else {}
model = settings.get("model", "birefnet-general")
model = settings.get("model", "birefnet-general-lite")
bg_color = settings.get("backgroundColor", "")
try:
from rembg import remove, new_session
from PIL import Image
import io
print(json.dumps({"progress": "loading_model"}), flush=True)
# Progress messages go to stderr (stdout reserved for JSON result)
sys.stderr.write(f"Loading model: {model}\n")
sys.stderr.flush()
# Create a session with the selected model
session = new_session(model)
sys.stderr.write("Processing image...\n")
sys.stderr.flush()
with open(input_path, "rb") as f:
input_data = f.read()
print(json.dumps({"progress": "processing"}), flush=True)
# Try with alpha matting first for better edges
# Try with alpha matting for better edges, fall back without
try:
output_data = remove(
input_data,
@@ -40,6 +41,8 @@ def main():
# If a background color is specified, composite onto it
if bg_color and bg_color.startswith("#"):
from PIL import Image
img = Image.open(io.BytesIO(output_data)).convert("RGBA")
hex_color = bg_color.lstrip("#")
r = int(hex_color[0:2], 16)