From 3a977c7438e6b838131b032c052490ace0615287 Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 18 Jun 2026 16:13:11 +0000 Subject: [PATCH] Add prefetch-models.sh to download AI models on host when container DNS is blocked MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Lets SAM, U2Net, BEN2, and BiRefNet-HR (optionally SDXL via --sdxl) be downloaded outside Docker into ./data/, which is already bind-mounted into the GPU container — so a blocked container network no longer blocks first-run setup. Reuses the existing dual-mode download_sam_model.py and download_u2net_model.py as-is. For the HuggingFace Hub models, sets HF_HOME (rather than --cache-dir) so the host-side cache layout matches the container's default ~/.cache/huggingface resolution exactly, avoiding a path-nesting mismatch between the two. Wired into install-local-gpu.sh's completion banner and bring-up-local-gpu.sh's header, and referenced from the relevant README troubleshooting sections and the hf_cache bind-mount comment in docker-compose.gpu.yml. --- README.md | 23 +++++++++-- bring-up-local-gpu.sh | 4 ++ docker-compose.gpu.yml | 7 ++-- install-local-gpu.sh | 4 ++ prefetch-models.sh | 92 ++++++++++++++++++++++++++++++++++++++++++ 5 files changed, 123 insertions(+), 7 deletions(-) create mode 100755 prefetch-models.sh diff --git a/README.md b/README.md index e8bab06..30c200a 100644 --- a/README.md +++ b/README.md @@ -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. diff --git a/bring-up-local-gpu.sh b/bring-up-local-gpu.sh index 7a13cec..cbbdf2e 100755 --- a/bring-up-local-gpu.sh +++ b/bring-up-local-gpu.sh @@ -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 diff --git a/docker-compose.gpu.yml b/docker-compose.gpu.yml index fbd4225..eea2f0a 100644 --- a/docker-compose.gpu.yml +++ b/docker-compose.gpu.yml @@ -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) diff --git a/install-local-gpu.sh b/install-local-gpu.sh index bb6be1d..78b6c9b 100755 --- a/install-local-gpu.sh +++ b/install-local-gpu.sh @@ -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 "==================================================" diff --git a/prefetch-models.sh b/prefetch-models.sh new file mode 100755 index 0000000..0a39bb8 --- /dev/null +++ b/prefetch-models.sh @@ -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 "=================================================="