diff --git a/scripts/install_rmbg.sh b/scripts/install_rmbg.sh new file mode 100755 index 0000000..e95f6a6 --- /dev/null +++ b/scripts/install_rmbg.sh @@ -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" diff --git a/server/operators/bg_remove_local/manifest.json b/server/operators/bg_remove_local/manifest.json new file mode 100644 index 0000000..be7cabc --- /dev/null +++ b/server/operators/bg_remove_local/manifest.json @@ -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)"} + } + } +} diff --git a/server/operators/bg_remove_local/run.py b/server/operators/bg_remove_local/run.py new file mode 100644 index 0000000..04c54bc --- /dev/null +++ b/server/operators/bg_remove_local/run.py @@ -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}}]}))