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:
parent
21938a4e17
commit
1cc0b80d89
14
scripts/install_rmbg.sh
Executable file
14
scripts/install_rmbg.sh
Executable 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"
|
||||
18
server/operators/bg_remove_local/manifest.json
Normal file
18
server/operators/bg_remove_local/manifest.json
Normal 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)"}
|
||||
}
|
||||
}
|
||||
}
|
||||
68
server/operators/bg_remove_local/run.py
Normal file
68
server/operators/bg_remove_local/run.py
Normal 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}}]}))
|
||||
Loading…
Reference in New Issue
Block a user