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

Auto-download U2Net model on startup; remove dead PyTorch U2Net module

Browse Source 
U2Net only ever downloaded lazily on the first Remove Background click,
unlike SAM which retries on every container start. If that one attempt
failed (DNS/firewall) the model was never fetched again, surfacing as
"No background removal method available. Install u2net or rembg."

Mirrors the existing SAM auto-download/AUTO_DOWNLOAD_SAM pattern for
U2Net, and documents manual host-side recovery in the README.

Also deletes backend/app/services/u2net_model.py (hand-written
U2NET/U2NETP PyTorch classes) — unused since tools.py switched to
cv2.dnn.readNetFromONNX for background removal.
This commit is contained in:
Claude
2026-06-18 13:10:07 +00:00
parent 34642161f4
commit a4a9dbc33a
5 changed files with 55 additions and 502 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/entrypoint.sh
+33 -2
View File
@@ -5,8 +5,9 @@
# This script runs when the container starts. It:
# 1. Initializes the database
# 2. Downloads SAM model automatically (can be disabled with AUTO_DOWNLOAD_SAM=false)
# 3. Downloads sample eye images if the catalog is empty
# 4. Starts the FastAPI server
# 3. Downloads U2Net model automatically (can be disabled with AUTO_DOWNLOAD_U2NET=false)
# 4. Downloads sample eye images if the catalog is empty
# 5. Starts the FastAPI server
# =============================================================================
set -e
@@ -61,6 +62,36 @@ else
fi
fi
echo ""
echo "Checking U2Net model (Remove Background)..."
echo "------------------------------------------"
if [ -f "/app/data/models/u2net.onnx" ] || [ -f "/app/data/models/u2netp.onnx" ]; then
echo "✓ U2Net model found - Remove Background will use local AI (free, offline)"
else
# Auto-download U2Net unless explicitly disabled
AUTO_DOWNLOAD_U2NET="${AUTO_DOWNLOAD_U2NET:-true}"
if [ "$AUTO_DOWNLOAD_U2NET" = "true" ]; then
echo "U2Net model not found. Downloading automatically..."
echo "(This is a one-time ~176MB download that persists across rebuilds)"
echo ""
python /scripts/download_u2net_model.py u2net || {
echo ""
echo "⚠ U2Net download failed (non-fatal)"
echo " Remove Background will fall back to rembg (if installed)"
echo " To retry later: docker exec -it ai-photo-edit-backend python /scripts/download_u2net_model.py u2net"
}
else
echo ""
echo "⚠ U2Net model not found (AUTO_DOWNLOAD_U2NET=false)"
echo ""
echo " Remove Background will fall back to rembg (if installed)"
echo ""
echo " To enable FREE offline Remove Background, run:"
echo " docker exec -it ai-photo-edit-backend python /scripts/download_u2net_model.py u2net"
echo ""
fi
fi
echo ""
echo "Checking GPU capabilities..."
echo "------------------------------------------"