root/PaintPlus
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

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

Browse Source 
Claude/focused maxwell 3kgh3x
This commit is contained in:
Outis
2026-06-18 10:39:24 -04:00
committed by GitHub
parent 668728348b 38e80f8fb8
commit f58c96c20f
10 changed files with 169 additions and 114 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.env.example
+6
View File
@@ -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
+4
View File
@@ -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
+19 -13
View File
@@ -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
+5
View File
@@ -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
+112 -15
View File
@@ -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
+11
View File
@@ -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
+1
View File
@@ -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
+6 -1
View File
@@ -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
+5 -3
View File
@@ -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
-82
View File
@@ -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..."