Update provision.sh

provision.sh

@@ -1,146 +1,392 @@
 #!/bin/bash
+# =============================================================================
+# ComfyUI Provisioning Script for Vast.ai
+#
+# Unified setup: installs custom nodes, downloads all models (LoRAs, VAE,
+# checkpoints) from HuggingFace (private + public) and CivitAI.
+#
+# Uses aria2c for fast parallel downloads (16 connections per file).
+#
+# Environment variables (set in Vast.ai template Docker Options):
+#   HF_TOKEN       - Required. HuggingFace token for private repos.
+#   CIVITAI_TOKEN  - Required. CivitAI API token for gated models.
+#
+# Usage:
+#   # Set in Vast.ai template environment:
+#   HF_TOKEN=hf_xxx CIVITAI_TOKEN=xxx ./setup-comfyui-vastai.sh
+#
+#   # Or source the ai-dock venv first (if using ai-dock template):
+#   source /venv/main/bin/activate && ./setup-comfyui-vastai.sh
+# =============================================================================
 
-echo "Starting provisioning script..."
-
-############################################
-# 0. HUGGING FACE AUTH
-############################################
-if [ -z "$HF_TOKEN" ]; then
-    echo "⚠️ HF_TOKEN not set — private Hugging Face repos will fail."
-else
-    echo "Logging into Hugging Face..."
-    pip install -q huggingface_hub
-    huggingface-cli login --token $HF_TOKEN
+set -euo pipefail
+
+# =============================================================================
+# Configuration
+# =============================================================================
+
+WORKSPACE="${WORKSPACE:-/workspace}"
+COMFYUI_DIR="${COMFYUI_DIR:-${WORKSPACE}/ComfyUI}"
+LORA_DIR="${COMFYUI_DIR}/models/loras"
+VAE_DIR="${COMFYUI_DIR}/models/vae"
+CKPT_DIR="${COMFYUI_DIR}/models/checkpoints"
+NODES_DIR="${COMFYUI_DIR}/custom_nodes"
+TEXT_ENCODER_DIR="${COMFYUI_DIR}/models/text_encoders"
+CLIP_DIR="${COMFYUI_DIR}/models/clip"
+DIFFUSION_DIR="${COMFYUI_DIR}/models/diffusion_models"
+ULTRALYTICS_DIR="${COMFYUI_DIR}/models/ultralytics/bbox"
+UPSCALE_DIR="${COMFYUI_DIR}/models/upscale_models"
+
+# Activate venv if available (ai-dock template)
+if [ -f /venv/main/bin/activate ]; then
+    source /venv/main/bin/activate
 fi
 
-HF_HEADER="Authorization: Bearer $HF_TOKEN"
-AUTH_HEADER="Authorization: Bearer 50d5f37a60d5bb4b1148bbc7addfe2a0"
-
-############################################
-# 1. DIRECTORIES
-############################################
-mkdir -p /workspace/ComfyUI/models/{text_encoders,checkpoints,vae,clip,loras,diffusion_models,upscale_models}
-mkdir -p /workspace/ComfyUI/models/ultralytics/bbox
-
-############################################
-# 2. HELPER FOR PRIVATE HF DOWNLOADS
-############################################
-hf_wget () {
-    DEST=$1
-    URL=$2
-    if [ ! -f "$DEST" ]; then
-        echo "Downloading $(basename $DEST)"
-        wget --header="$HF_HEADER" -O "$DEST" "$URL"
+# =============================================================================
+# Model Definitions
+# =============================================================================
+
+# HuggingFace LoRAs — private repo (requires HF_TOKEN)
+# Format: "filename|url"
+HF_PRIVATE_LORAS=(
+    "Z-Isla_3K_0002_000003500.safetensors|https://huggingface.co/cozyga/test/resolve/main/Z-Isla_3K_0002_000003500.safetensors"
+    "z-kyla-v1.safetensors|https://huggingface.co/cozyga/test/resolve/main/z-kyla-v1.safetensors"
+    "m1a_zit_3.safetensors|https://huggingface.co/cozyga/test/resolve/main/m1a_zit_3.safetensors"
+    "c3lia_v1_000003600.safetensors|https://huggingface.co/cozyga/test/resolve/main/c3lia_v1_000003600.safetensors"
+    "c3lia_v1_000003900.safetensors|https://huggingface.co/cozyga/test/resolve/main/c3lia_v1_000003900.safetensors"
+    "qwen_c3lia_000003600.safetensors|https://huggingface.co/cozyga/test/resolve/main/qwen_c3lia_000003600.safetensors"
+    "s0fia_v2_000004800.safetensors|https://huggingface.co/cozyga/test/resolve/main/s0fia_v2_000004800.safetensors"
+    "s0fia_v2_000003700.safetensors|https://huggingface.co/cozyga/test/resolve/main/s0fia_v2_000003700.safetensors"
+    "s0fia_v2_000004300.safetensors|https://huggingface.co/cozyga/test/resolve/main/s0fia_v2_000004300.safetensors"
+)
+
+# HuggingFace LoRAs — public repos
+HF_PUBLIC_LORAS=(
+    "l0ki_zit_v1_000003000.safetensors|https://huggingface.co/XLB666/qwenzitloras/resolve/main/l0ki_zit_v1_000003000.safetensors"
+    "l0ki_qwen_v1_000002800.safetensors|https://huggingface.co/XLB666/qwenzitloras/resolve/main/l0ki_qwen_v1_000002800.safetensors"
+    "dmd2_sdxl_4step_lora_fp16.safetensors|https://huggingface.co/tianweiy/DMD2/resolve/main/dmd2_sdxl_4step_lora_fp16.safetensors"
+)
+
+# Additional HuggingFace model groups (added from user's list)
+HF_TEXT_ENCODERS=(
+    "t5xxl_fp8_e4m3fn_scaled.safetensors|https://huggingface.co/comfyanonymous/flux_text_encoders/resolve/main/t5xxl_fp8_e4m3fn_scaled.safetensors"
+    "qwen_2.5_vl_7b_fp8_scaled.safetensors|https://huggingface.co/Comfy-Org/Qwen-Image_ComfyUI/resolve/main/split_files/text_encoders/qwen_2.5_vl_7b_fp8_scaled.safetensors"
+)
+
+HF_CLIPS=(
+    "qwen_3_4b.safetensors|https://huggingface.co/Comfy-Org/z_image_turbo/resolve/main/split_files/text_encoders/qwen_3_4b.safetensors"
+)
+
+HF_CHECKPOINTS_EXTRA=(
+    "Qwen-Rapid-AIO-SFW-v23.safetensors|https://huggingface.co/Phr00t/Qwen-Image-Edit-Rapid-AIO/resolve/main/v23/Qwen-Rapid-AIO-SFW-v23.safetensors"
+    "Qwen-Rapid-AIO-NSFW-v23.safetensors|https://huggingface.co/Phr00t/Qwen-Image-Edit-Rapid-AIO/resolve/main/v23/Qwen-Rapid-AIO-NSFW-v23.safetensors"
+    "lustifySDXLNSFW_endgame.safetensors|https://huggingface.co/cozyga/test/resolve/main/lustifySDXLNSFW_endgame.safetensors"
+)
+
+HF_VAES_EXTRA=(
+    "qwen_image_vae.safetensors|https://huggingface.co/Comfy-Org/Qwen-Image_ComfyUI/resolve/main/split_files/vae/qwen_image_vae.safetensors"
+    "flux_vae.safetensors|https://huggingface.co/StableDiffusionVN/Flux/resolve/main/Vae/flux_vae.safetensors"
+)
+
+HF_DIFFUSION=(
+    "z_image_turbo_bf16.safetensors|https://huggingface.co/Comfy-Org/z_image_turbo/resolve/main/split_files/diffusion_models/z_image_turbo_bf16.safetensors"
+)
+
+HF_ULTRALYTICS=(
+    "FacesV1.pt|https://huggingface.co/ashllay/YOLO_Models/resolve/e07b01219ff1807e1885015f439d788b038f49bd/bbox/FacesV1.pt"
+    "nipple.pt|https://huggingface.co/cozyga/test/resolve/main/nipple.pt"
+    "pussyV2.pt|https://huggingface.co/cozyga/test/resolve/main/pussyV2.pt"
+)
+
+HF_UPSCALE=(
+    "1x-ITF-SkinDiffDetail-Lite-v1.pth|https://huggingface.co/cozyga/test/resolve/main/1x-ITF-SkinDiffDetail-Lite-v1.pth"
+)
+
+# CivitAI LoRAs (requires CIVITAI_TOKEN)
+# Format: "filename|url"
+# Explicit filenames avoid the CivitAI redirect long-filename bug
+CIVITAI_LORAS=(
+    "zittitslider.safetensors|https://civitai.com/api/download/models/2478366"
+    "QwenEditVersatilePoses.safetensors|https://civitai.com/api/download/models/2457989?type=Model&format=SafeTensor"
+    "sdxlcloseuppussy.safetensors|https://civitai.com/api/download/models/2376235?type=Model&format=SafeTensor"
+    "amateurphotography.safetensors|https://civitai.com/api/download/models/2363467"
+    "nextscene.safetensors|https://civitai.com/api/download/models/2337214?type=Model&format=SafeTensor"
+)
+
+# CivitAI VAE models
+CIVITAI_VAE=(
+    "zImage_vae.safetensors|https://civitai.com/api/download/models/2442479"
+)
+
+# CivitAI Checkpoints
+CIVITAI_CHECKPOINTS=(
+    # Pony Diffusion V6 XL (or whichever model ID 798204 maps to — verify this)
+    "checkpoint.safetensors|https://civitai.com/api/download/models/798204?type=Model&format=SafeTensor&size=full&fp=fp16"
+    "LustifyENDGAME.safetensors|https://civitai.com/api/download/models/1094291?type=Model&format=SafeTensor&size=pruned&fp=fp16"
+    "analogMadnessXL.safetensors|https://civitai.com/api/download/models/2207703?type=Model&format=SafeTensor&size=full&fp=fp16"
+)
+
+# Custom nodes to install/update
+NODES=(
+    "https://github.com/ltdrdata/ComfyUI-Impact-Pack"
+    "https://github.com/chflame163/ComfyUI_LayerStyle"
+    "https://github.com/rgthree/rgthree-comfy"
+    "https://github.com/yolain/ComfyUI-Easy-Use"
+    "https://github.com/numz/ComfyUI-SeedVR2_VideoUpscaler"
+    "https://github.com/cubiq/ComfyUI_essentials"
+    "https://github.com/ltdrdata/ComfyUI-Impact-Subpack"
+    "https://github.com/EllangoK/ComfyUI-post-processing-nodes"
+    "https://github.com/wallish77/wlsh_nodes"
+    "https://github.com/chrisgoringe/cg-image-filter"
+    "https://github.com/moonwhaler/comfyui-seedvr2-tilingups"
+)
+
+# Extra pip packages (installed after nodes)
+PIP_PACKAGES=(
+    # "package-name"
+)
+
+# =============================================================================
+# Functions
+# =============================================================================
+
+print_header() {
+    echo ""
+    echo "=============================================="
+    echo "  ComfyUI Provisioning — Vast.ai"
+    echo "=============================================="
+    echo "  Workspace:  ${WORKSPACE}"
+    echo "  ComfyUI:    ${COMFYUI_DIR}"
+    echo "  HF_TOKEN:   ${HF_TOKEN:+set}${HF_TOKEN:-MISSING}"
+    echo "  CIVITAI:    ${CIVITAI_TOKEN:+set}${CIVITAI_TOKEN:-MISSING}"
+    echo "=============================================="
+    echo ""
+}
+
+# Install aria2c if not present (vast.ai containers are typically Debian/Ubuntu)
+install_aria2() {
+    if command -v aria2c &>/dev/null; then
+        return 0
+    fi
+    echo "Installing aria2c for fast parallel downloads..."
+    if [ "$(id -u)" -eq 0 ]; then
+        apt-get update -qq && apt-get install -y -qq aria2
+    elif command -v sudo &>/dev/null; then
+        sudo apt-get update -qq && sudo apt-get install -y -qq aria2
+    else
+        echo "ERROR: aria2 not found and cannot install (not root, no sudo)."
+        echo "Falling back to wget."
+        return 1
+    fi
+}
+
+# Validate that a HuggingFace token is valid
+validate_hf_token() {
+    if [ -z "${HF_TOKEN:-}" ]; then
+        echo "WARNING: HF_TOKEN is not set. Private HuggingFace downloads will fail."
+        return 1
+    fi
+    local status
+    status=$(curl -s -o /dev/null -w "%{http_code}" \
+        -H "Authorization: Bearer ${HF_TOKEN}" \
+        "https://huggingface.co/api/whoami-v2")
+    if [ "$status" -eq 200 ]; then
+        echo "HF_TOKEN: validated"
+        return 0
     else
-        echo "Exists: $(basename $DEST)"
+        echo "WARNING: HF_TOKEN validation failed (HTTP ${status}). Private downloads may fail."
+        return 1
     fi
 }
 
-############################################
-# 3. MODEL DOWNLOADS (Hugging Face)
-############################################
-echo "Checking Hugging Face models..."
-
-[ ! -f /workspace/ComfyUI/models/text_encoders/t5xxl_fp8_e4m3fn_scaled.safetensors ] && \
-hf_wget /workspace/ComfyUI/models/text_encoders/t5xxl_fp8_e4m3fn_scaled.safetensors \
-https://huggingface.co/comfyanonymous/flux_text_encoders/resolve/main/t5xxl_fp8_e4m3fn_scaled.safetensors
-
-[ ! -f /workspace/ComfyUI/models/checkpoints/Qwen-Rapid-AIO-SFW-v23.safetensors ] && \
-hf_wget /workspace/ComfyUI/models/checkpoints/Qwen-Rapid-AIO-SFW-v23.safetensors \
-https://huggingface.co/Phr00t/Qwen-Image-Edit-Rapid-AIO/resolve/main/v23/Qwen-Rapid-AIO-SFW-v23.safetensors
-
-[ ! -f /workspace/ComfyUI/models/checkpoints/Qwen-Rapid-AIO-NSFW-v23.safetensors ] && \
-hf_wget /workspace/ComfyUI/models/checkpoints/Qwen-Rapid-AIO-NSFW-v23.safetensors \
-https://huggingface.co/Phr00t/Qwen-Image-Edit-Rapid-AIO/resolve/main/v23/Qwen-Rapid-AIO-NSFW-v23.safetensors
-
-[ ! -f /workspace/ComfyUI/models/vae/qwen_image_vae.safetensors ] && \
-hf_wget /workspace/ComfyUI/models/vae/qwen_image_vae.safetensors \
-https://huggingface.co/Comfy-Org/Qwen-Image_ComfyUI/resolve/main/split_files/vae/qwen_image_vae.safetensors
-
-[ ! -f /workspace/ComfyUI/models/clip/qwen_3_4b.safetensors ] && \
-hf_wget /workspace/ComfyUI/models/clip/qwen_3_4b.safetensors \
-https://huggingface.co/Comfy-Org/z_image_turbo/resolve/main/split_files/text_encoders/qwen_3_4b.safetensors
-
-[ ! -f /workspace/ComfyUI/models/text_encoders/qwen_2.5_vl_7b_fp8_scaled.safetensors ] && \
-hf_wget /workspace/ComfyUI/models/text_encoders/qwen_2.5_vl_7b_fp8_scaled.safetensors \
-https://huggingface.co/Comfy-Org/Qwen-Image_ComfyUI/resolve/main/split_files/text_encoders/qwen_2.5_vl_7b_fp8_scaled.safetensors
-
-[ ! -f /workspace/ComfyUI/models/diffusion_models/z_image_turbo_bf16.safetensors ] && \
-hf_wget /workspace/ComfyUI/models/diffusion_models/z_image_turbo_bf16.safetensors \
-https://huggingface.co/Comfy-Org/z_image_turbo/resolve/main/split_files/diffusion_models/z_image_turbo_bf16.safetensors
-
-[ ! -f /workspace/ComfyUI/models/vae/flux_vae.safetensors ] && \
-hf_wget /workspace/ComfyUI/models/vae/flux_vae.safetensors \
-https://huggingface.co/StableDiffusionVN/Flux/resolve/main/Vae/flux_vae.safetensors
-
-[ ! -f /workspace/ComfyUI/models/ultralytics/bbox/FacesV1.pt ] && \
-hf_wget /workspace/ComfyUI/models/ultralytics/bbox/FacesV1.pt \
-https://huggingface.co/ashllay/YOLO_Models/resolve/e07b01219ff1807e1885015f439d788b038f49bd/bbox/FacesV1.pt
-
-[ ! -f /workspace/ComfyUI/models/upscale_models/1x-ITF-SkinDiffDetail-Lite-v1.pth ] && \
-hf_wget /workspace/ComfyUI/models/upscale_models/1x-ITF-SkinDiffDetail-Lite-v1.pth \
-https://huggingface.co/cozyga/test/resolve/main/1x-ITF-SkinDiffDetail-Lite-v1.pth
-
-############################################
-# CIVITAI MODELS
-############################################
-[ ! -f /workspace/ComfyUI/models/checkpoints/LustifyENDGAME.safetensors ] && \
-curl -L -H "$AUTH_HEADER" -o /workspace/ComfyUI/models/checkpoints/LustifyENDGAME.safetensors \
-"https://civitai.com/api/download/models/1094291?type=Model&format=SafeTensor&size=pruned&fp=fp16"
-
-[ ! -f /workspace/ComfyUI/models/checkpoints/analogMadnessXL.safetensors ] && \
-curl -L -H "$AUTH_HEADER" -o /workspace/ComfyUI/models/checkpoints/analogMadnessXL.safetensors \
-"https://civitai.com/api/download/models/2207703?type=Model&format=SafeTensor&size=full&fp=fp16"
-
-############################################
-# REMAINING FILES (UNCHANGED)
-############################################
-wget --content-disposition -P /workspace/ComfyUI/models/ultralytics/bbox/ \
-https://huggingface.co/cozyga/test/resolve/main/nipple.pt \
-https://huggingface.co/cozyga/test/resolve/main/pussyV2.pt
-
-wget --content-disposition -P /workspace/ComfyUI/models/checkpoints/ \
-https://huggingface.co/cozyga/test/resolve/main/lustifySDXLNSFW_endgame.safetensors
-
-wget --content-disposition -P /workspace/ComfyUI/models/loras/ \
-https://huggingface.co/cozyga/test/resolve/main/dmd2_sdxl_4step_lora_fp16.safetensors
-
-############################################
-# (Your LoRAs + Civitai section stays identical)
-############################################
-LORA_DIR="/workspace/ComfyUI/models/loras"
-[ ! -f "$LORA_DIR/Z-Isla_3K_0002_000003500.safetensors" ] && wget -P "$LORA_DIR" https://huggingface.co/cozyga/test/resolve/main/Z-Isla_3K_0002_000003500.safetensors
-[ ! -f "$LORA_DIR/z-kyla-v1.safetensors" ] && wget -P "$LORA_DIR" https://huggingface.co/cozyga/test/resolve/main/z-kyla-v1.safetensors
-[ ! -f "$LORA_DIR/m1a_zit_3.safetensors" ] && wget -P "$LORA_DIR" https://huggingface.co/cozyga/test/resolve/main/m1a_zit_3.safetensors
-[ ! -f "$LORA_DIR/c3lia_v1_000003600.safetensors" ] && wget -P "$LORA_DIR" https://huggingface.co/cozyga/test/resolve/main/c3lia_v1_000003600.safetensors
-[ ! -f "$LORA_DIR/c3lia_v1_000003900.safetensors" ] && wget -P "$LORA_DIR" https://huggingface.co/cozyga/test/resolve/main/c3lia_v1_000003900.safetensors
-[ ! -f "$LORA_DIR/qwen_c3lia_000003600.safetensors" ] && wget -P "$LORA_DIR" https://huggingface.co/cozyga/test/resolve/main/qwen_c3lia_000003600.safetensors
-[ ! -f "$LORA_DIR/l0ki_zit_v1_000003000.safetensors" ] && wget -P "$LORA_DIR" https://huggingface.co/XLB666/qwenzitloras/resolve/main/l0ki_zit_v1_000003000.safetensors
-[ ! -f "$LORA_DIR/l0ki_qwen_v1_000002800.safetensors" ] && wget -P "$LORA_DIR" https://huggingface.co/XLB666/qwenzitloras/resolve/main/l0ki_qwen_v1_000002800.safetensors
-[ ! -f "$LORA_DIR/s0fia_v2_000004800.safetensors" ] && wget -P "$LORA_DIR" https://huggingface.co/cozyga/test/resolve/main/s0fia_v2_000004800.safetensors
-[ ! -f "$LORA_DIR/s0fia_v2_000003700.safetensors" ] && wget -P "$LORA_DIR" https://huggingface.co/cozyga/test/resolve/main/s0fia_v2_000003700.safetensors
-[ ! -f "$LORA_DIR/s0fia_v2_000004300.safetensors" ] && wget -P "$LORA_DIR" https://huggingface.co/cozyga/test/resolve/main/s0fia_v2_000004300.safetensors
-[ ! -f "$LORA_DIR/dmd2_sdxl_4step_lora_fp16.safetensors" ] && wget -P "$LORA_DIR" https://huggingface.co/tianweiy/DMD2/resolve/main/dmd2_sdxl_4step_lora_fp16.safetensors
-[ ! -f /workspace/ComfyUI/models/loras/zittitslider.safetensors ] && \
-    curl -L -H "$AUTH_HEADER" -o /workspace/ComfyUI/models/loras/zittitslider.safetensors "https://civitai.com/api/download/models/2478366"
-[ ! -f /workspace/ComfyUI/models/loras/QwenEditVersatilePoses.safetensors ] && \
-    curl -L -H "$AUTH_HEADER" -o /workspace/ComfyUI/models/loras/QwenEditVersatilePoses.safetensors "https://civitai.com/api/download/models/2457989?type=Model&format=SafeTensor"
-[ ! -f /workspace/ComfyUI/models/loras/sdxlcloseuppussy.safetensors ] && \
-    curl -L -H "$AUTH_HEADER" -o /workspace/ComfyUI/models/loras/sdxlcloseuppussy.safetensors "https://civitai.com/api/download/models/2376235?type=Model&format=SafeTensor"
-[ ! -f /workspace/ComfyUI/models/vae/zImage_vae.safetensors ] && \
-    curl -L -H "$AUTH_HEADER" -o /workspace/ComfyUI/models/vae/zImage_vae.safetensors "https://civitai.com/api/download/models/2442479"
-[ ! -f /workspace/ComfyUI/models/loras/amateurphotography.safetensors ] && \
-    curl -L -H "$AUTH_HEADER" -o /workspace/ComfyUI/models/loras/amateurphotography.safetensors "https://civitai.com/api/download/models/2363467"
-[ ! -f /workspace/ComfyUI/models/loras/nextscene.safetensors ] && \
-    curl -L -H "$AUTH_HEADER" -o /workspace/ComfyUI/models/loras/nextscene.safetensors "https://civitai.com/api/download/models/2337214?type=Model&format=SafeTensor"
-
-############################################
-# UPDATE & START
-############################################
-echo "Provisioning complete. Updating ComfyUI..."
-
-cd /workspace/ComfyUI || exit
-git pull
-
-echo "ComfyUI is updated. Starting now..."
-python main.py --listen
+# Validate that a CivitAI token is valid
+validate_civitai_token() {
+    if [ -z "${CIVITAI_TOKEN:-}" ]; then
+        echo "WARNING: CIVITAI_TOKEN is not set. CivitAI downloads will fail."
+        return 1
+    fi
+    local status
+    status=$(curl -s -o /dev/null -w "%{http_code}" \
+        -H "Authorization: Bearer ${CIVITAI_TOKEN}" \
+        "https://civitai.com/api/v1/models?hidden=1&limit=1")
+    if [ "$status" -eq 200 ]; then
+        echo "CIVITAI_TOKEN: validated"
+        return 0
+    else
+        echo "WARNING: CIVITAI_TOKEN validation failed (HTTP ${status}). CivitAI downloads may fail."
+        return 1
+    fi
+}
+
+# Download a file if it doesn't already exist
+# Uses aria2c with 16 parallel connections for speed
+# Falls back to wget if aria2c is unavailable
+#
+# Usage: download_if_missing <dir> <filename> <url> [auth_header]
+download_if_missing() {
+    local dir="$1"
+    local filename="$2"
+    local url="$3"
+    local auth_header="${4:-}"
+
+    mkdir -p "$dir"
+
+    if [ -f "${dir}/${filename}" ]; then
+        echo "  [skip] ${filename}"
+        return 0
+    fi
+
+    echo "  [downloading] ${filename}..."
+
+    if command -v aria2c &>/dev/null; then
+        if [ -n "$auth_header" ]; then
+            aria2c -x 16 -s 16 -k 1M -c \
+                --header="Authorization: Bearer ${auth_header}" \
+                -d "$dir" -o "$filename" "$url" 2>&1 | tail -1
+        else
+            aria2c -x 16 -s 16 -k 1M -c \
+                -d "$dir" -o "$filename" "$url" 2>&1 | tail -1
+        fi
+    else
+        # Fallback: wget
+        if [ -n "$auth_header" ]; then
+            wget -q --show-progress \
+                --header="Authorization: Bearer ${auth_header}" \
+                -O "${dir}/${filename}" "$url"
+        else
+            wget -q --show-progress -O "${dir}/${filename}" "$url"
+        fi
+    fi
+
+    # Verify file was actually downloaded (not a 0-byte error)
+    if [ ! -s "${dir}/${filename}" ]; then
+        echo "  [FAILED] ${filename} — downloaded file is empty, removing"
+        rm -f "${dir}/${filename}"
+        return 1
+    fi
+}
+
+# Download a batch of files from an array of "filename|url" entries
+# Usage: download_batch <dir> <auth_token> <array_entries...>
+download_batch() {
+    local dir="$1"
+    local auth_token="$2"
+    shift 2
+
+    local failed=0
+    for entry in "$@"; do
+        local filename="${entry%%|*}"
+        local url="${entry#*|}"
+        download_if_missing "$dir" "$filename" "$url" "$auth_token" || ((failed++)) || true
+    done
+
+    if [ "$failed" -gt 0 ]; then
+        echo "  WARNING: ${failed} download(s) failed in ${dir}"
+    fi
+}
+
+# Install/update custom nodes
+install_nodes() {
+    if [ ${#NODES[@]} -eq 0 ]; then
+        return 0
+    fi
+
+    echo "=== Installing Custom Nodes ==="
+    mkdir -p "$NODES_DIR"
+
+    for repo in "${NODES[@]}"; do
+        local dirname="${repo##*/}"
+        local path="${NODES_DIR}/${dirname}"
+        local requirements="${path}/requirements.txt"
+
+        if [ -d "$path" ]; then
+            echo "  [update] ${dirname}"
+            (cd "$path" && git pull -q)
+        else
+            echo "  [clone] ${dirname}"
+            git clone --recursive -q "$repo" "$path"
+        fi
+
+        if [ -f "$requirements" ]; then
+            pip install --no-cache-dir -q -r "$requirements"
+        fi
+    done
+    echo ""
+}
+
+# Install extra pip packages
+install_pip_packages() {
+    if [ ${#PIP_PACKAGES[@]} -eq 0 ]; then
+        return 0
+    fi
+    echo "=== Installing pip packages ==="
+    pip install --no-cache-dir -q "${PIP_PACKAGES[@]}"
+    echo ""
+}
+
+# =============================================================================
+# Main
+# =============================================================================
+
+print_header
+
+# Install aria2c for fast downloads
+install_aria2 || true
+
+# Validate tokens (non-fatal — script continues with warnings)
+validate_hf_token || true
+validate_civitai_token || true
+echo ""
+
+# Create all model directories
+mkdir -p "$LORA_DIR" "$VAE_DIR" "$CKPT_DIR"
+
+# --- Custom Nodes ---
+install_nodes
+
+# --- Pip Packages ---
+install_pip_packages
+
+# --- HuggingFace Private LoRAs ---
+echo "=== HuggingFace Private LoRAs (cozyga/test) ==="
+download_batch "$LORA_DIR" "${HF_TOKEN:-}" "${HF_PRIVATE_LORAS[@]}"
+echo ""
+
+# --- HuggingFace Public LoRAs ---
+echo "=== HuggingFace Public LoRAs ==="
+download_batch "$LORA_DIR" "${HF_TOKEN:-}" "${HF_PUBLIC_LORAS[@]}"
+echo ""
+
+# --- CivitAI LoRAs ---
+echo "=== CivitAI LoRAs ==="
+download_batch "$LORA_DIR" "${CIVITAI_TOKEN:-}" "${CIVITAI_LORAS[@]}"
+echo ""
+
+# --- CivitAI VAE ---
+echo "=== CivitAI VAE ==="
+download_batch "$VAE_DIR" "${CIVITAI_TOKEN:-}" "${CIVITAI_VAE[@]}"
+echo ""
+
+# --- CivitAI Checkpoints ---
+echo "=== CivitAI Checkpoints ==="
+download_batch "$CKPT_DIR" "${CIVITAI_TOKEN:-}" "${CIVITAI_CHECKPOINTS[@]}"
+echo ""
+
+# =============================================================================
+# Summary
+# =============================================================================
+
+echo "=============================================="
+echo "  Provisioning complete"
+echo "=============================================="
+echo ""
+echo "LoRAs:"
+ls -1 "$LORA_DIR"/*.safetensors 2>/dev/null | wc -l | xargs -I{} echo "  {} files in ${LORA_DIR}"
+echo ""
+echo "VAE:"
+ls -1 "$VAE_DIR"/*.safetensors 2>/dev/null | wc -l | xargs -I{} echo "  {} files in ${VAE_DIR}"
+echo ""
+echo "Checkpoints:"
+ls -1 "$CKPT_DIR"/*.safetensors 2>/dev/null | wc -l | xargs -I{} echo "  {} files in ${CKPT_DIR}"
+echo ""
+echo "Application will start now."