Add prefetch-models.sh to download AI models on host when container DNS is blocked

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.
This commit is contained in:
Claude
2026-06-18 16:13:11 +00:00
parent 38e80f8fb8
commit 3a977c7438
5 changed files with 123 additions and 7 deletions
+20 -3
View File
@@ -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.