Merge pull request #63 from outis1one/claude/focused-maxwell-3kgh3x

Claude/focused maxwell 3kgh3x

.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
+- **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:
-mkdir -p ./data/models
+  ```bash
-sudo curl -L -o ./data/models/u2net.onnx \
+  docker compose -f docker-compose.gpu.yml logs -f | grep -iE "ben2|birefnet"
-  https://github.com/danielgatis/rembg/releases/download/v0.0.0/u2net.onnx
+  ```
-```
+- **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:
+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.
-```bash
-docker compose logs -f | grep -i u2net
-# Should show: "U2Net model loaded successfully with OpenCV DNN"
-```
 
 **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)
+        for name in order:
-        try:
+            try:
-            result_bytes = await _remove_background_u2net(img)
+                result_bytes = await local_backends[name](img)
-            method_used = "u2net"
+                method_used = name
-        except Exception as e:
+                break
-            print(f"U2Net failed: {e}")
+            except Exception as e:
+                print(f"{name} failed: {e}")
 
-        # Fall back to rembg if U2Net failed
+        # rembg is the universal last resort (also reachable directly via model="rembg")
-        if result_bytes is None:
+        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,
             }),
         });
 

install-local-gpu.sh

@@ -147,88 +147,6 @@ echo "✓ Docker restarted"
 systemctl start docker-dns-fix.service
 echo "✓ DNS fix applied"
 
-# ── 5. Verify GPU access ─────────────────────────────────────────────────────
-echo ""
-echo "Verifying GPU access inside Docker..."
-if docker run --rm --gpus all nvidia/cuda:12.1.0-base-ubuntu22.04 nvidia-smi &>/dev/null; then
-    echo "✓ GPU is accessible inside Docker"
-else
-    echo "⚠ GPU check failed. Try rebooting if the driver was just installed."
-    echo "  Manual check: nvidia-smi"
-fi
-
-echo ""
-echo "=================================================="
-echo "  Setup complete."
-echo "  Start the app with:  ./bring-up-local-gpu.sh"
-echo "=================================================="
-
-if command -v nvidia-ctk &>/dev/null; then
-    echo "✓ nvidia-container-toolkit already installed — skipping"
-else
-    echo "Installing nvidia-container-toolkit..."
-    . /etc/os-release
-    case "$ID" in
-        ubuntu|debian)
-            curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey \
-                | gpg --dearmor -o /usr/share/keyrings/nvidia-ctk.gpg
-            curl -fsSL "https://nvidia.github.io/libnvidia-container/stable/deb/nvidia-container-toolkit.list" \
-                | sed 's#deb https://#deb [signed-by=/usr/share/keyrings/nvidia-ctk.gpg] https://#g' \
-                | tee /etc/apt/sources.list.d/nvidia-container-toolkit.list
-            apt-get update -qq
-            apt-get install -y nvidia-container-toolkit
-            ;;
-        rhel|fedora|rocky|centos|almalinux)
-            dnf install -y nvidia-container-toolkit
-            ;;
-        *)
-            echo "⚠ Unrecognised distro ($ID). Install nvidia-container-toolkit manually."
-            echo "  See: https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/install-guide.html"
-            ;;
-    esac
-fi
-
-nvidia-ctk runtime configure --runtime=docker
-
-# ── 2. Permanent Docker DNS fix via systemd ───────────────────────────────────
-# Adds a rule to the DOCKER-USER iptables chain so containers can resolve
-# hostnames. Runs after docker.service on every boot. Does NOT touch ufw.
-echo ""
-echo "Installing docker-dns-fix systemd service..."
-
-cat > /etc/systemd/system/docker-dns-fix.service << 'EOF'
-[Unit]
-Description=Allow Docker containers to resolve DNS (DOCKER-USER iptables rule)
-After=docker.service
-Requires=docker.service
-BindsTo=docker.service
-
-[Service]
-Type=oneshot
-ExecStart=/bin/sh -c \
-  'iptables -C DOCKER-USER -p udp --dport 53 -j ACCEPT 2>/dev/null || \
-   iptables -I DOCKER-USER -p udp --dport 53 -j ACCEPT'
-RemainAfterExit=yes
-
-[Install]
-WantedBy=multi-user.target
-EOF
-
-systemctl daemon-reload
-systemctl enable docker-dns-fix.service
-echo "✓ docker-dns-fix.service installed and enabled"
-
-# ── 3. Restart Docker ─────────────────────────────────────────────────────────
-echo ""
-echo "Restarting Docker..."
-systemctl restart docker
-sleep 2
-echo "✓ Docker restarted"
-
-# ── 4. Apply DNS rule now (don't wait for next boot) ─────────────────────────
-systemctl start docker-dns-fix.service
-echo "✓ DNS fix applied"
-
 # ── 5. Verify GPU access ─────────────────────────────────────────────────────
 echo ""
 echo "Verifying GPU access inside Docker..."