Claude
3a977c7438
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.
93 lines
3.6 KiB
Bash
Executable File
93 lines
3.6 KiB
Bash
Executable File
#!/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 "=================================================="
|