Add BEN2 and BiRefNet-HR as selectable Remove Background models

BEN2 becomes the new default local backend (clean cutouts, strong on
hair/fur edges), with BiRefNet-HR available as a high-res/print
alternate and U2Net kept as the lightweight fallback. Both are
MIT-licensed and download weights from HuggingFace on first use
(cached via the existing hf_cache bind mount), unlike U2Net/SAM which
need an explicit download script.

- config: new BG_REMOVAL_MODEL setting (default "ben2")
- tools.py: remove-background-base64 now tries local backends in
  order (request.model override > BG_REMOVAL_MODEL > ben2/u2net),
  falling back to rembg's birefnet-general session as a last resort
- requirements.gpu.txt / Dockerfile.gpu: add ben2 + transformers deps
  needed for the new backends, with a build-time smoke test for ben2
- frontend: model dropdown in the Remove Background dialog, threaded
  through api.js to the new request field

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Ro4PwQKvSc3CH19LSN21Ht

.env.example

@@ -171,6 +171,12 @@ AUTO_DOWNLOAD_SAM=true
 # When false: Skips download, Remove Background falls back to rembg (if installed)
 AUTO_DOWNLOAD_U2NET=true
 
+# Background removal model (Remove Background tool) — used when request.model="auto"
+# Options: ben2 (default — best for clean cutouts, hair/edges), birefnet-hr
+# (best for high-res/print work, slower), u2net (lightweight, always-on fallback)
+# ben2 and birefnet-hr download weights from HuggingFace on first use (GPU image only).
+BG_REMOVAL_MODEL=ben2
+
 # Allow users to select model per-edit
 ALLOW_MODEL_OVERRIDE=true
 

Dockerfile.gpu

@@ -54,6 +54,10 @@ RUN echo "BUILDID=$BUILDID" && pip install --no-cache-dir -r requirements.gpu.tx
 RUN python -c "from rembg import remove; print('rembg OK')" \
     || echo "WARNING: rembg unavailable — Remove Background disabled"
 
+# Smoke-test ben2 (weights download from HuggingFace on first use)
+RUN python -c "import ben2; print('ben2 OK')" \
+    || echo "WARNING: ben2 unavailable — Remove Background falls back to U2Net/rembg"
+
 # Copy backend application
 COPY backend/ .
 

README.md

@@ -111,7 +111,7 @@ Override the auto-selected model with `HF_MODEL_TXT2IMG`, `HF_MODEL_INPAINT` in
 - **Prepare for Print** — one-click: AI upscale to target DPI + fit to frame
 - **Fit to Frame** — resize/crop/AI-extend to standard print sizes
 - **Expand Canvas (Outpaint)** — AI extends the image in any direction
-- **Remove Background** — one-click background removal
+- **Remove Background** — one-click background removal (BEN2 by default, BiRefNet-HR or U2Net selectable)
 
 ### Print presets
 Frame sizes: 4×6, 5×7, 8×10, 11×14, 16×20, 18×24, 20×24, 24×36 (portrait + landscape)  
@@ -201,21 +201,27 @@ docker compose -f docker-compose.gpu.yml logs | grep -i sam
 
 If Docker created `./data/` as root and you can't write there without `sudo`, you can also use root's curl as above — the container reads the file regardless of owner.
 
-**Remove Background fails ("Install u2net or rembg")**
+**Remove Background fails ("Install ben2, u2net, or rembg")**
 
-The U2Net model auto-downloads (~176MB) from GitHub on first use, same as SAM. If that download fails (DNS/firewall, see above) and `rembg` isn't installed either, you'll see this error. Fix it the same way — download directly on the host:
+Remove Background tries, in order: the model set by `BG_REMOVAL_MODEL` (default `ben2`), then the other local models, then `rembg` as a last resort. You'll see this error only if all of them fail.
 
-```bash
-mkdir -p ./data/models
-sudo curl -L -o ./data/models/u2net.onnx \
-  https://github.com/danielgatis/rembg/releases/download/v0.0.0/u2net.onnx
-```
+- **ben2 / birefnet-hr** (GPU image only) download their weights from HuggingFace on first use, cached under `./data/hf_cache`. If that download fails (DNS/firewall, see above), check the logs for the specific error:
+  ```bash
+  docker compose -f docker-compose.gpu.yml logs -f | grep -iE "ben2|birefnet"
+  ```
+- **u2net** auto-downloads (~176MB) from GitHub on first use, same as SAM. If that fails too, download it directly on the host:
+  ```bash
+  mkdir -p ./data/models
+  sudo curl -L -o ./data/models/u2net.onnx \
+    https://github.com/danielgatis/rembg/releases/download/v0.0.0/u2net.onnx
+  ```
+  The file is ~176 MB. Once it exists at `./data/models/u2net.onnx`, the next "Remove Background" click picks it up — no rebuild or restart needed. Verify with:
+  ```bash
+  docker compose logs -f | grep -i u2net
+  # Should show: "U2Net model loaded successfully with OpenCV DNN"
+  ```
 
-The file is ~176 MB. Once it exists at `./data/models/u2net.onnx`, the next "Remove Background" click picks it up — no rebuild or restart needed. Verify with:
-```bash
-docker compose logs -f | grep -i u2net
-# Should show: "U2Net model loaded successfully with OpenCV DNN"
-```
+You can also pick a specific model per-edit from the Remove Background dialog's model dropdown, overriding `BG_REMOVAL_MODEL` for that one call.
 
 **AI models not downloading (container DNS blocked)**
 

backend/app/config.py

@@ -46,6 +46,11 @@ class Settings(BaseSettings):
     # Allow per-edit model override
     allow_model_override: bool = True
 
+    # Remove Background — preferred local model when request.model="auto"
+    # Options: ben2 (default, best for clean cutouts/hair), birefnet-hr (best
+    # for high-res/print work), u2net (lightweight, smallest download)
+    bg_removal_model: str = "ben2"
+
     # Local GPU diffusion (AI_PROVIDER=local_gpu)
     auto_download_models: bool = True      # download HF models on first use
     local_gpu_max_pipelines: int = 2       # max diffusion pipelines kept in GPU memory

backend/app/routers/tools.py

@@ -35,6 +35,7 @@ class InpaintRequest(BaseModel):
 
 class RemoveBackgroundRequest(BaseModel):
     image: str  # Base64 encoded image
+    model: Optional[str] = "auto"  # "auto", "ben2", "birefnet-hr", "u2net", "rembg"
 
 
 @router.post("/smart-select-base64")
@@ -128,36 +129,54 @@ async def inpaint_base64(request: InpaintRequest):
         raise HTTPException(status_code=500, detail=str(e))
 
 
-class RemoveBackgroundRequestV2(BaseModel):
-    image: str  # Base64 encoded image
-    model: Optional[str] = "auto"  # "auto", "u2net", "rembg", "birefnet"
-
-
 @router.post("/remove-background-base64")
 async def remove_background_base64(request: RemoveBackgroundRequest):
     """
     Remove background from a base64 encoded image.
-    Tries multiple methods: U2Net (direct), rembg with BiRefNet, rembg default.
+
+    request.model selects the backend:
+      - "auto" (default): BG_REMOVAL_MODEL setting first, then falls back
+        through the other local models, then rembg as a last resort.
+      - "ben2" / "birefnet-hr" / "u2net": use only that local model.
+      - "rembg": skip local models, use rembg directly.
+
     Returns base64 encoded PNG with transparent background.
     Used by miniPaint frontend.
     """
     try:
+        from app.config import settings
+
         # Decode base64 image
         image_bytes = base64.b64decode(request.image)
         img = Image.open(BytesIO(image_bytes)).convert('RGB')
 
+        local_backends = {
+            "ben2": _remove_background_ben2,
+            "birefnet-hr": _remove_background_birefnet_hr,
+            "u2net": _remove_background_u2net,
+        }
+
+        if request.model in local_backends:
+            order = [request.model]
+        elif request.model == "rembg":
+            order = []
+        else:
+            preferred = settings.bg_removal_model if settings.bg_removal_model in local_backends else "ben2"
+            order = [preferred] + [name for name in ("ben2", "u2net") if name != preferred]
+
         result_bytes = None
         method_used = None
 
-        # Try U2Net first (direct implementation, no rembg dependency issues)
-        try:
-            result_bytes = await _remove_background_u2net(img)
-            method_used = "u2net"
-        except Exception as e:
-            print(f"U2Net failed: {e}")
+        for name in order:
+            try:
+                result_bytes = await local_backends[name](img)
+                method_used = name
+                break
+            except Exception as e:
+                print(f"{name} failed: {e}")
 
-        # Fall back to rembg if U2Net failed
-        if result_bytes is None:
+        # rembg is the universal last resort (also reachable directly via model="rembg")
+        if result_bytes is None and request.model in ("auto", "rembg"):
             try:
                 from rembg import remove, new_session
                 try:
@@ -175,7 +194,7 @@ async def remove_background_base64(request: RemoveBackgroundRequest):
         if result_bytes is None:
             raise HTTPException(
                 status_code=500,
-                detail="No background removal method available. Install u2net or rembg."
+                detail="No background removal method available. Install ben2, u2net, or rembg."
             )
 
         # Convert result to base64
@@ -328,6 +347,84 @@ async def _remove_background_u2net(img: Image.Image) -> bytes:
     return buffer.getvalue()
 
 
+# Global BEN2 model cache
+_ben2_model = None
+
+
+async def _remove_background_ben2(img: Image.Image) -> bytes:
+    """
+    Remove background using BEN2 (Confidence Guided Matting) — clean cutouts,
+    strong on hair/fur edges. MIT licensed. Downloads weights from HF Hub on
+    first use (cached under the hf_cache bind mount).
+    """
+    global _ben2_model
+
+    if _ben2_model is None:
+        import torch
+        from ben2 import AutoModel as Ben2AutoModel
+
+        device = 'cuda' if torch.cuda.is_available() else 'cpu'
+        print(f"Loading BEN2_Base model on {device} (first run downloads ~170MB from HuggingFace)")
+        _ben2_model = Ben2AutoModel.from_pretrained("PramaLLC/BEN2")
+        _ben2_model.to(device).eval()
+        print("BEN2_Base model loaded")
+
+    result = _ben2_model.inference(img.convert('RGB'), refine_foreground=False)
+
+    buffer = BytesIO()
+    result.save(buffer, format='PNG')
+    return buffer.getvalue()
+
+
+# Global BiRefNet-HR model cache
+_birefnet_hr_model = None
+_birefnet_hr_device = None
+
+
+async def _remove_background_birefnet_hr(img: Image.Image) -> bytes:
+    """
+    Remove background using BiRefNet-HR (2048x2048, MIT licensed) — best for
+    high-resolution / print work. Downloads weights from HF Hub on first use.
+    """
+    global _birefnet_hr_model, _birefnet_hr_device
+
+    import torch
+    from torchvision import transforms
+
+    if _birefnet_hr_model is None:
+        from transformers import AutoModelForImageSegmentation
+
+        _birefnet_hr_device = 'cuda' if torch.cuda.is_available() else 'cpu'
+        print(f"Loading BiRefNet-HR model on {_birefnet_hr_device} (first run downloads ~900MB from HuggingFace)")
+        _birefnet_hr_model = AutoModelForImageSegmentation.from_pretrained(
+            'zhengpeng7/BiRefNet_HR', trust_remote_code=True
+        )
+        _birefnet_hr_model.to(_birefnet_hr_device).eval()
+        print("BiRefNet-HR model loaded")
+
+    original_size = img.size
+    rgb_img = img.convert('RGB')
+
+    transform = transforms.Compose([
+        transforms.Resize((2048, 2048)),
+        transforms.ToTensor(),
+        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
+    ])
+    input_tensor = transform(rgb_img).unsqueeze(0).to(_birefnet_hr_device)
+
+    with torch.no_grad():
+        preds = _birefnet_hr_model(input_tensor)[-1].sigmoid().cpu()
+
+    mask = transforms.ToPILImage()(preds[0].squeeze()).resize(original_size, Image.Resampling.LANCZOS)
+
+    result = rgb_img.convert('RGBA')
+    result.putalpha(mask)
+
+    buffer = BytesIO()
+    result.save(buffer, format='PNG')
+    return buffer.getvalue()
+
+
 @router.post("/remove-background")
 async def remove_background(
     project_id: Optional[int] = Form(None),

backend/requirements.gpu.txt

@@ -31,3 +31,14 @@ sentencepiece>=0.2.0
 # Install post-container-start if needed:
 #   pip install xformers --index-url https://download.pytorch.org/whl/cu121
 # xformers
+
+# Background removal — BEN2 (default, clean cutouts/hair) + BiRefNet-HR
+# (high-res/print alternate). Both MIT-licensed. Verified against upstream
+# source: neither requires torch>=2.5 despite the BiRefNet repo's own
+# requirements.txt floor — that pin is for its training/eval scripts, not
+# the inference path used here. Weights download from HuggingFace on first
+# use (cached via the hf_cache bind mount, same as the diffusion models).
+ben2 @ git+https://github.com/PramaLLC/BEN2.git
+timm>=1.0.10
+einops>=0.6.0
+kornia>=0.7.0

docker-compose.gpu.yml

@@ -121,6 +121,7 @@ services:
       - CORS_ORIGINS=*
       - AUTO_DOWNLOAD_SAM=${AUTO_DOWNLOAD_SAM:-true}
       - AUTO_DOWNLOAD_U2NET=${AUTO_DOWNLOAD_U2NET:-true}
+      - BG_REMOVAL_MODEL=${BG_REMOVAL_MODEL:-ben2}
 
     # ── NVIDIA GPU passthrough ────────────────────────────────────────────────
     # Requires nvidia-container-toolkit; see prerequisites at top of this file.

frontend/src/js/modules/image/remove_background.js

@@ -49,6 +49,11 @@ class Image_remove_background_class {
             title: 'Remove Background',
             params: [
                 { name: "info", title: "AI will detect the main subject and remove the background.", type: "label" },
+                {
+                    name: "model", title: "Model:", value: "auto", type: "select",
+                    values: ["auto", "ben2", "birefnet-hr", "u2net"],
+                    comment: "auto = best available (BEN2 by default). BiRefNet-HR is slower but sharper on high-res/print work.",
+                },
                 { name: "new_layer", title: "Create as new layer:", value: true },
                 { name: "trim_result", title: "Trim transparent edges:", value: false },
             ],
@@ -74,7 +79,7 @@ class Image_remove_background_class {
             var imageData = canvas.toDataURL('image/png').split(',')[1];
 
             // Call backend API
-            var result = await apiService.removeBackground(imageData);
+            var result = await apiService.removeBackground(imageData, params.model);
 
             // Create image from result
             var resultImage = new Image();

frontend/src/js/services/api.js

@@ -72,11 +72,12 @@ class ApiService {
     }
 
     /**
-     * Remove background from image using AI (rembg)
+     * Remove background from image using AI (BEN2 / BiRefNet-HR / U2Net / rembg)
      * @param {string} imageData - Base64 encoded image data
-     * @returns {Promise<{result: string, width: number, height: number}>} - Base64 encoded result with transparency
+     * @param {string} [model='auto'] - "auto", "ben2", "birefnet-hr", "u2net", or "rembg"
+     * @returns {Promise<{result: string, width: number, height: number, method: string}>} - Base64 encoded result with transparency
      */
-    async removeBackground(imageData) {
+    async removeBackground(imageData, model = 'auto') {
         const response = await fetch(`${this.baseUrl}/tools/remove-background-base64`, {
             method: 'POST',
             headers: {
@@ -84,6 +85,7 @@ class ApiService {
             },
             body: JSON.stringify({
                 image: imageData,
+                model: model,
             }),
         });