Merge pull request #64 from outis1one/claude/focused-maxwell-3kgh3x
Add prefetch-models.sh to download AI models on host when container D…
This commit is contained in:
Outis
@@ -161,6 +161,9 @@ EditmaskwithAI/
|
||||
├── docker-compose.dev.yml # Dev with hot reload
|
||||
├── Dockerfile
|
||||
├── Dockerfile.gpu
|
||||
├── install-local-gpu.sh # One-time GPU host setup
|
||||
├── bring-up-local-gpu.sh # Start/stop the GPU container
|
||||
├── prefetch-models.sh # Download AI models on the host (DNS-blocked workaround)
|
||||
└── .env.example
|
||||
```
|
||||
|
||||
@@ -187,6 +190,8 @@ docker compose -f docker-compose.gpu.yml logs -f | grep -E "local_gpu|Error|Fail
|
||||
If the container can't reach `dl.fbaipublicfiles.com` (you'll see `Errno -3 Name or service not known` in the logs), download SAM directly on the host and let the bind mount make it visible to the container — no rebuild needed:
|
||||
|
||||
```bash
|
||||
./prefetch-models.sh
|
||||
# or manually:
|
||||
mkdir -p ./data/models
|
||||
# sudo needed if ./data/ was created by Docker (root-owned):
|
||||
sudo curl -L -o ./data/models/sam_vit_b_01ec64.pth \
|
||||
@@ -205,12 +210,14 @@ If Docker created `./data/` as root and you can't write there without `sudo`, yo
|
||||
|
||||
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.
|
||||
|
||||
- **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:
|
||||
- **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), run `./prefetch-models.sh` to fetch both directly on the host, or 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
|
||||
./prefetch-models.sh
|
||||
# or manually:
|
||||
mkdir -p ./data/models
|
||||
sudo curl -L -o ./data/models/u2net.onnx \
|
||||
https://github.com/danielgatis/rembg/releases/download/v0.0.0/u2net.onnx
|
||||
@@ -225,7 +232,17 @@ You can also pick a specific model per-edit from the Remove Background dialog's
|
||||
|
||||
**AI models not downloading (container DNS blocked)**
|
||||
|
||||
If you ran `./install-local-gpu.sh`, this is already permanently fixed. Otherwise, the container's host firewall is blocking outbound DNS from the Docker bridge — apply the fix manually (does **not** affect container isolation):
|
||||
Easiest fix: download the models on the host instead of inside the container — they land in `./data/`, which is already bind-mounted into the container, so it picks them up with no rebuild:
|
||||
|
||||
```bash
|
||||
./prefetch-models.sh # SAM + U2Net + BEN2 + BiRefNet-HR (~1.5GB)
|
||||
./prefetch-models.sh --sdxl # also Text→Image / AI Edit models (~13GB)
|
||||
./bring-up-local-gpu.sh
|
||||
```
|
||||
|
||||
If that also fails to reach the network, the problem is host-level (firewall/DNS), not Docker-specific — see your network/firewall configuration.
|
||||
|
||||
Alternatively, if you ran `./install-local-gpu.sh`, container DNS is already permanently fixed via a systemd-managed iptables rule. If you skipped that script, apply the same fix manually (does **not** affect container isolation):
|
||||
|
||||
```bash
|
||||
sudo iptables -I DOCKER-USER -p udp --dport 53 -j ACCEPT
|
||||
@@ -237,7 +254,7 @@ The container will now resolve hostnames and download models automatically (~13
|
||||
docker compose -f docker-compose.gpu.yml logs -f | grep -E "local_gpu|Cached|failed"
|
||||
```
|
||||
|
||||
**Alternative: download with a Docker helper container** (no iptables, no host Python needed):
|
||||
**Alternative: download with a Docker helper container** (no host Python needed):
|
||||
|
||||
```bash
|
||||
# Inpainting model (~6.5 GB) — needed for AI Edit, Make less symmetrical, etc.
|
||||
|
||||
@@ -4,6 +4,10 @@
|
||||
# Run this each time you want to start the app.
|
||||
# Run ./install-local-gpu.sh once first on a new machine.
|
||||
#
|
||||
# If models fail to download inside the container (DNS/firewall blocked),
|
||||
# run ./prefetch-models.sh first — it downloads them on the host into
|
||||
# ./data/, which this container already bind-mounts, so no rebuild needed.
|
||||
#
|
||||
# Usage:
|
||||
# ./bring-up-local-gpu.sh # start (detached, rebuild if needed)
|
||||
# ./bring-up-local-gpu.sh --no-build # start without rebuilding
|
||||
|
||||
@@ -73,10 +73,9 @@ services:
|
||||
# Persistent project data
|
||||
- ./data:/app/data
|
||||
# HuggingFace model cache — bind mount so models can be pre-downloaded on the host.
|
||||
# If container DNS is blocked, download on the host and the container picks them up:
|
||||
# pip install huggingface-hub
|
||||
# huggingface-cli download diffusers/stable-diffusion-xl-1.0-inpainting-0.1 \
|
||||
# --cache-dir ./data/hf_cache
|
||||
# If container DNS is blocked, run ./prefetch-models.sh on the host first —
|
||||
# it downloads BEN2/BiRefNet-HR (and optionally SDXL with --sdxl) straight
|
||||
# into this directory, and the container picks them up on next start.
|
||||
# To free disk space: rm -rf ./data/hf_cache
|
||||
- ./data/hf_cache:/root/.cache/huggingface
|
||||
# Scripts (for exec access)
|
||||
|
||||
@@ -162,4 +162,8 @@ echo ""
|
||||
echo "=================================================="
|
||||
echo " Setup complete."
|
||||
echo " Start the app with: ./bring-up-local-gpu.sh"
|
||||
echo ""
|
||||
echo " If models fail to download inside the container (DNS/firewall"
|
||||
echo " blocked), fetch them on the host first instead:"
|
||||
echo " ./prefetch-models.sh"
|
||||
echo "=================================================="
|
||||
|
||||
Executable
+92
@@ -0,0 +1,92 @@
|
||||
#!/usr/bin/env bash
|
||||
# prefetch-models.sh — download AI models on the host, outside Docker.
|
||||
#
|
||||
# Use this when the container's outbound DNS/network is blocked (see
|
||||
# README troubleshooting) and models can't be downloaded at container
|
||||
# startup. Downloads land under ./data/, which both compose files already
|
||||
# bind-mount into the container — so the container picks them up on next
|
||||
# start with no rebuild and no in-container network access required.
|
||||
#
|
||||
# Usage:
|
||||
# ./prefetch-models.sh # SAM + U2Net + BEN2 + BiRefNet-HR (~1.5GB)
|
||||
# ./prefetch-models.sh --sdxl # also prefetch SDXL base + inpaint (~13GB)
|
||||
#
|
||||
# Safe to re-run: every download here skips files that already exist
|
||||
# (HuggingFace Hub) or are already present (SAM/U2Net).
|
||||
|
||||
set -euo pipefail
|
||||
|
||||
cd "$(dirname "$0")"
|
||||
|
||||
if ! command -v python3 &>/dev/null; then
|
||||
echo "✗ python3 is required on the host for this script (Docker is not used here)." >&2
|
||||
echo " Ubuntu/Debian: sudo apt install python3 python3-pip" >&2
|
||||
exit 1
|
||||
fi
|
||||
|
||||
if ! mkdir -p data/models data/hf_cache 2>/dev/null; then
|
||||
echo "✗ Could not create ./data/models or ./data/hf_cache." >&2
|
||||
echo " If ./data/ was already created by Docker (root-owned), re-run with sudo:" >&2
|
||||
echo " sudo ./prefetch-models.sh $*" >&2
|
||||
exit 1
|
||||
fi
|
||||
|
||||
PREFETCH_SDXL=0
|
||||
if [ "${1:-}" = "--sdxl" ]; then
|
||||
PREFETCH_SDXL=1
|
||||
fi
|
||||
|
||||
echo "=================================================="
|
||||
echo " Prefetching AI models (host-side, no Docker)"
|
||||
echo "=================================================="
|
||||
|
||||
echo ""
|
||||
echo "── SAM (Smart Select) ───────────────────────────────"
|
||||
python3 scripts/download_sam_model.py vit_b
|
||||
|
||||
echo ""
|
||||
echo "── U2Net (Remove Background fallback) ──────────────"
|
||||
python3 scripts/download_u2net_model.py u2net
|
||||
|
||||
echo ""
|
||||
echo "── HuggingFace Hub models (BEN2, BiRefNet-HR) ───────"
|
||||
|
||||
if ! python3 -c "import huggingface_hub" &>/dev/null; then
|
||||
echo "Installing huggingface_hub (lightweight — no torch/GPU needed for this step)..."
|
||||
python3 -m pip install --quiet --user "huggingface_hub>=0.23.0"
|
||||
fi
|
||||
|
||||
# HF_HOME must match what the container resolves by default: the bind mount
|
||||
# maps ./data/hf_cache -> /root/.cache/huggingface, and the container never
|
||||
# sets HF_HOME explicitly, so it defaults to ~/.cache/huggingface there.
|
||||
# huggingface_hub itself appends "/hub" to HF_HOME to get the actual cache
|
||||
# root (HF_HUB_CACHE) — setting HF_HOME here (instead of passing --cache-dir
|
||||
# or cache_dir=... directly) lets both sides derive that "/hub" nesting the
|
||||
# same way, rather than us hardcoding it and risking a mismatch.
|
||||
export HF_HOME="$(pwd)/data/hf_cache"
|
||||
|
||||
PREFETCH_SDXL="$PREFETCH_SDXL" python3 - << 'PYEOF'
|
||||
import os
|
||||
from huggingface_hub import snapshot_download
|
||||
|
||||
repos = ["PramaLLC/BEN2", "zhengpeng7/BiRefNet_HR"]
|
||||
if os.environ.get("PREFETCH_SDXL") == "1":
|
||||
repos += [
|
||||
"stabilityai/stable-diffusion-xl-base-1.0",
|
||||
"diffusers/stable-diffusion-xl-1.0-inpainting-0.1",
|
||||
]
|
||||
|
||||
for repo_id in repos:
|
||||
print(f"\nDownloading {repo_id} ...")
|
||||
snapshot_download(repo_id=repo_id, ignore_patterns=["*.msgpack", "flax_*", "tf_*"])
|
||||
print(f" done: {repo_id}")
|
||||
PYEOF
|
||||
|
||||
echo ""
|
||||
echo "=================================================="
|
||||
echo " Done. Models cached under ./data/models and ./data/hf_cache"
|
||||
if [ "$PREFETCH_SDXL" != "1" ]; then
|
||||
echo " (SDXL not included — re-run with --sdxl to also prefetch txt2img/inpaint, ~13GB)"
|
||||
fi
|
||||
echo " Start the app: ./bring-up-local-gpu.sh"
|
||||
echo "=================================================="
|
||||
Reference in New Issue
Block a user