bg_remove_local: local RMBG-2.0 background removal on MPS

New operator: briaai/RMBG-2.0 image->transparent cutout, fully local on Apple
Silicon (weights unlocked by owner HF license). Free local alternative to the
cloud fal_bg_remove; the recommended pre-pass before SF3D / image->3D. Verified:
clean 1024px cutout of the FLUX astrolabe (alpha 0-255, 86.6% removed, thin rings
preserved). Same OpenMP guardrails as sf3d. scripts/install_rmbg.sh; venvs/rmbg.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
MODELBEAST 2026-07-13 01:32:38 +10:00
parent 21938a4e17
commit 1cc0b80d89
3 changed files with 100 additions and 0 deletions

14
scripts/install_rmbg.sh Executable file
View File

@ -0,0 +1,14 @@
#!/usr/bin/env bash
# briaai/RMBG-2.0 local background removal on Apple Silicon (MPS). Weights are
# HF-gated (owner has access); downloaded on first run via the HF token in Settings.
set -euo pipefail
cd "$(dirname "$0")/.."
UV=/opt/homebrew/bin/uv
echo "[rmbg] creating venv (python 3.11) ..."
$UV venv --python 3.11 venvs/rmbg
PY="$(pwd)/venvs/rmbg/bin/python"
echo "[rmbg] installing deps (transformers + RMBG-2.0 remote-code deps) ..."
$UV pip install --python "$PY" torch torchvision transformers pillow numpy kornia timm scikit-image einops
echo "[rmbg] verifying ..."
"$PY" -c "import torch, transformers, kornia, timm; print('[rmbg] torch', torch.__version__, 'mps', torch.backends.mps.is_available(), '| transformers', transformers.__version__)"
echo "[rmbg] done"

View File

@ -0,0 +1,18 @@
{
"id": "bg_remove_local",
"name": "Remove Background (local)",
"category": "image-prep",
"description": "Image → clean subject cutout (transparent PNG) entirely on this Mac via briaai/RMBG-2.0 (MPS). Free local alternative to fal_bg_remove — run before SF3D / image→3D for a big quality jump. Weights unlocked by the owner's HF license.",
"accepts": ["image"],
"produces": ["image"],
"resources": "gpu",
"entry": "run.py",
"python": "/Users/m3ultra/Documents/MODELBEAST/venvs/rmbg/bin/python",
"params_schema": {
"type": "object",
"properties": {
"resolution": {"type": "integer", "enum": [1024, 2048], "default": 1024, "description": "Processing resolution (2048 = finer edges, slower)"},
"background": {"type": "string", "enum": ["transparent", "white", "black"], "default": "transparent", "description": "Fill for removed background (transparent RGBA, or a solid matte)"}
}
}
}

View File

@ -0,0 +1,68 @@
import argparse
import json
import os
import sys
from pathlib import Path
# Same libomp collision guardrails as sf3d (kornia/timm/torch each may link one).
os.environ.setdefault("PYTORCH_ENABLE_MPS_FALLBACK", "1")
os.environ.setdefault("KMP_DUPLICATE_LIB_OK", "TRUE")
os.environ.setdefault("OMP_NUM_THREADS", "1")
os.environ.setdefault("MKL_NUM_THREADS", "1")
import torch # noqa: E402
from PIL import Image # noqa: E402
from torchvision import transforms # noqa: E402
from transformers import AutoModelForImageSegmentation # noqa: E402
ap = argparse.ArgumentParser()
ap.add_argument("--input", action="append", default=[])
ap.add_argument("--outdir", required=True)
ap.add_argument("--params", default="{}")
a = ap.parse_args()
p = json.loads(a.params)
if not a.input:
print("ERROR: no input image")
sys.exit(1)
res = int(p.get("resolution", 1024))
bg = p.get("background", "transparent")
device = "mps" if torch.backends.mps.is_available() else "cpu"
print(f"loading briaai/RMBG-2.0 on {device} (first run downloads weights) ...", flush=True)
model = AutoModelForImageSegmentation.from_pretrained("briaai/RMBG-2.0", trust_remote_code=True)
try:
torch.set_float32_matmul_precision("high")
except Exception:
pass
model.eval().to(device)
transform = transforms.Compose([
transforms.Resize((res, res)),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
])
src = Path(a.input[0])
image = Image.open(src).convert("RGB")
inp = transform(image).unsqueeze(0).to(device)
print("running segmentation ...", flush=True)
with torch.no_grad():
preds = model(inp)[-1].sigmoid().cpu()
mask = transforms.ToPILImage()(preds[0].squeeze()).resize(image.size)
outdir = Path(a.outdir)
name = f"{src.stem}_cutout.png"
if bg == "transparent":
out = image.copy()
out.putalpha(mask)
else:
fill = (255, 255, 255) if bg == "white" else (0, 0, 0)
canvas = Image.new("RGB", image.size, fill)
canvas.paste(image, (0, 0), mask)
out = canvas
out.save(outdir / name)
print(f"done: {name}", flush=True)
(outdir / "result.json").write_text(json.dumps(
{"outputs": [{"path": name, "meta": {"tool": "rmbg-2.0", "resolution": res, "background": bg}}]}))