From 781b9e7cdcaa57e13f42673aa0a23b80317bf77a Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 18 Jun 2026 13:57:28 +0000 Subject: [PATCH 1/2] Fix install-local-gpu.sh running its entire setup body twice Commit 5c85987 appended a second copy of the toolkit-install/DNS-fix/ GPU-verify steps instead of inserting the new prerequisite-checks section once. Removed the duplicate, keeping the more informative GPU-check failure message from the second copy. --- install-local-gpu.sh | 82 -------------------------------------------- 1 file changed, 82 deletions(-) diff --git a/install-local-gpu.sh b/install-local-gpu.sh index 41b0793..bb6be1d 100755 --- a/install-local-gpu.sh +++ b/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..." From 38e80f8fb876a7276fe383dc0dfb91239bea3b33 Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 18 Jun 2026 14:20:24 +0000 Subject: [PATCH 2/2] 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 Claude-Session: https://claude.ai/code/session_01Ro4PwQKvSc3CH19LSN21Ht --- .env.example | 6 + Dockerfile.gpu | 4 + README.md | 32 +++-- backend/app/config.py | 5 + backend/app/routers/tools.py | 127 +++++++++++++++--- backend/requirements.gpu.txt | 11 ++ docker-compose.gpu.yml | 1 + .../src/js/modules/image/remove_background.js | 7 +- frontend/src/js/services/api.js | 8 +- 9 files changed, 169 insertions(+), 32 deletions(-) diff --git a/.env.example b/.env.example index 64570c6..836db6f 100644 --- a/.env.example +++ b/.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 diff --git a/Dockerfile.gpu b/Dockerfile.gpu index a27584a..b922cef 100644 --- a/Dockerfile.gpu +++ b/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/ . diff --git a/README.md b/README.md index 41f2692..e8bab06 100644 --- a/README.md +++ b/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)** diff --git a/backend/app/config.py b/backend/app/config.py index cfe29e1..55b6d46 100644 --- a/backend/app/config.py +++ b/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 diff --git a/backend/app/routers/tools.py b/backend/app/routers/tools.py index 03a6fc7..4a0cbce 100644 --- a/backend/app/routers/tools.py +++ b/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), diff --git a/backend/requirements.gpu.txt b/backend/requirements.gpu.txt index 02af631..fd25964 100644 --- a/backend/requirements.gpu.txt +++ b/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 diff --git a/docker-compose.gpu.yml b/docker-compose.gpu.yml index eda3690..fbd4225 100644 --- a/docker-compose.gpu.yml +++ b/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. diff --git a/frontend/src/js/modules/image/remove_background.js b/frontend/src/js/modules/image/remove_background.js index 64e3050..923479a 100644 --- a/frontend/src/js/modules/image/remove_background.js +++ b/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(); diff --git a/frontend/src/js/services/api.js b/frontend/src/js/services/api.js index b78d6e9..f8c88b3 100644 --- a/frontend/src/js/services/api.js +++ b/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, }), });