Update provision.sh

provision.sh

@@ -19,7 +19,7 @@
 #   source /venv/main/bin/activate && ./setup-comfyui-vastai.sh
 # =============================================================================
 
-set -euo pipefail
+set -uo pipefail
 
 # =============================================================================
 # Configuration
@@ -30,12 +30,12 @@ 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"
+TEXT_ENC_DIR="${COMFYUI_DIR}/models/text_encoders"
 DIFFUSION_DIR="${COMFYUI_DIR}/models/diffusion_models"
-ULTRALYTICS_DIR="${COMFYUI_DIR}/models/ultralytics/bbox"
 UPSCALE_DIR="${COMFYUI_DIR}/models/upscale_models"
+ULTRALYTICS_DIR="${COMFYUI_DIR}/models/ultralytics/bbox"
+NODES_DIR="${COMFYUI_DIR}/custom_nodes"
 
 # Activate venv if available (ai-dock template)
 if [ -f /venv/main/bin/activate ]; then
@@ -67,74 +67,83 @@ HF_PUBLIC_LORAS=(
     "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=(
+# 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=(
+    "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"
+)
+
+# --- HuggingFace public models (various directories) ---
+
+HF_PUBLIC_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=(
+HF_PUBLIC_CLIP=(
     "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_PUBLIC_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"
 )
 
-HF_VAES_EXTRA=(
+HF_PUBLIC_VAE=(
     "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_PUBLIC_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"
+    "Qwen-Rapid-AIO-NSFW-v23.safetensors|https://huggingface.co/Phr00t/Qwen-Image-Edit-Rapid-AIO/resolve/main/v23/Qwen-Rapid-AIO-NSFW-v23.safetensors"
 )
 
-HF_ULTRALYTICS=(
+HF_PUBLIC_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"
-)
+# --- HuggingFace private models — cozyga/test (non-lora) ---
 
-# 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"
+HF_PRIVATE_CHECKPOINTS=(
+    "lustifySDXLNSFW_endgame.safetensors|https://huggingface.co/cozyga/test/resolve/main/lustifySDXLNSFW_endgame.safetensors"
 )
 
-# CivitAI VAE models
-CIVITAI_VAE=(
-    "zImage_vae.safetensors|https://civitai.com/api/download/models/2442479"
+HF_PRIVATE_UPSCALE=(
+    "1x-ITF-SkinDiffDetail-Lite-v1.pth|https://huggingface.co/cozyga/test/resolve/main/1x-ITF-SkinDiffDetail-Lite-v1.pth"
 )
 
-# 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"
+HF_PRIVATE_ULTRALYTICS=(
+    "nipple.pt|https://huggingface.co/cozyga/test/resolve/main/nipple.pt"
+    "pussyV2.pt|https://huggingface.co/cozyga/test/resolve/main/pussyV2.pt"
 )
 
 # Custom nodes to install/update
 NODES=(
+    "https://github.com/ltdrdata/ComfyUI-Manager"
     "https://github.com/ltdrdata/ComfyUI-Impact-Pack"
-    "https://github.com/chflame163/ComfyUI_LayerStyle"
-    "https://github.com/rgthree/rgthree-comfy"
+    "https://github.com/ltdrdata/ComfyUI-Impact-Subpack"
+    "https://github.com/cubiq/ComfyUI_essentials"
     "https://github.com/yolain/ComfyUI-Easy-Use"
+    "https://github.com/rgthree/rgthree-comfy"
+    "https://github.com/chflame163/ComfyUI_LayerStyle"
     "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"
@@ -301,14 +310,17 @@ install_nodes() {
 
         if [ -d "$path" ]; then
             echo "  [update] ${dirname}"
-            (cd "$path" && git pull -q)
+            (cd "$path" && git pull -q) || echo "  [WARN] git pull failed for ${dirname}, skipping"
         else
             echo "  [clone] ${dirname}"
-            git clone --recursive -q "$repo" "$path"
+            git clone --recursive -q "$repo" "$path" || {
+                echo "  [WARN] git clone failed for ${dirname}, skipping"
+                continue
+            }
         fi
 
         if [ -f "$requirements" ]; then
-            pip install --no-cache-dir -q -r "$requirements"
+            pip install --no-cache-dir -q -r "$requirements" || echo "  [WARN] pip install failed for ${dirname}"
         fi
     done
     echo ""
@@ -339,7 +351,8 @@ validate_civitai_token || true
 echo ""
 
 # Create all model directories
-mkdir -p "$LORA_DIR" "$VAE_DIR" "$CKPT_DIR"
+mkdir -p "$LORA_DIR" "$VAE_DIR" "$CKPT_DIR" "$CLIP_DIR" "$TEXT_ENC_DIR" \
+         "$DIFFUSION_DIR" "$UPSCALE_DIR" "$ULTRALYTICS_DIR"
 
 # --- Custom Nodes ---
 install_nodes
@@ -372,6 +385,51 @@ echo "=== CivitAI Checkpoints ==="
 download_batch "$CKPT_DIR" "${CIVITAI_TOKEN:-}" "${CIVITAI_CHECKPOINTS[@]}"
 echo ""
 
+# --- HuggingFace Public Text Encoders ---
+echo "=== HuggingFace Text Encoders ==="
+download_batch "$TEXT_ENC_DIR" "${HF_TOKEN:-}" "${HF_PUBLIC_TEXT_ENCODERS[@]}"
+echo ""
+
+# --- HuggingFace Public CLIP ---
+echo "=== HuggingFace CLIP ==="
+download_batch "$CLIP_DIR" "${HF_TOKEN:-}" "${HF_PUBLIC_CLIP[@]}"
+echo ""
+
+# --- HuggingFace Public Diffusion Models ---
+echo "=== HuggingFace Diffusion Models ==="
+download_batch "$DIFFUSION_DIR" "${HF_TOKEN:-}" "${HF_PUBLIC_DIFFUSION_MODELS[@]}"
+echo ""
+
+# --- HuggingFace Public VAE ---
+echo "=== HuggingFace VAE ==="
+download_batch "$VAE_DIR" "${HF_TOKEN:-}" "${HF_PUBLIC_VAE[@]}"
+echo ""
+
+# --- HuggingFace Public Checkpoints ---
+echo "=== HuggingFace Checkpoints (Qwen) ==="
+download_batch "$CKPT_DIR" "${HF_TOKEN:-}" "${HF_PUBLIC_CHECKPOINTS[@]}"
+echo ""
+
+# --- HuggingFace Public Ultralytics ---
+echo "=== HuggingFace Ultralytics (YOLO) ==="
+download_batch "$ULTRALYTICS_DIR" "${HF_TOKEN:-}" "${HF_PUBLIC_ULTRALYTICS[@]}"
+echo ""
+
+# --- HuggingFace Private Checkpoints (cozyga/test) ---
+echo "=== HuggingFace Private Checkpoints (cozyga/test) ==="
+download_batch "$CKPT_DIR" "${HF_TOKEN:-}" "${HF_PRIVATE_CHECKPOINTS[@]}"
+echo ""
+
+# --- HuggingFace Private Upscale Models (cozyga/test) ---
+echo "=== HuggingFace Private Upscale Models (cozyga/test) ==="
+download_batch "$UPSCALE_DIR" "${HF_TOKEN:-}" "${HF_PRIVATE_UPSCALE[@]}"
+echo ""
+
+# --- HuggingFace Private Ultralytics (cozyga/test) ---
+echo "=== HuggingFace Private Ultralytics (cozyga/test) ==="
+download_batch "$ULTRALYTICS_DIR" "${HF_TOKEN:-}" "${HF_PRIVATE_ULTRALYTICS[@]}"
+echo ""
+
 # =============================================================================
 # Summary
 # =============================================================================
@@ -380,13 +438,19 @@ 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}"
+for label_dir in \
+    "LoRAs|${LORA_DIR}" \
+    "Checkpoints|${CKPT_DIR}" \
+    "VAE|${VAE_DIR}" \
+    "CLIP|${CLIP_DIR}" \
+    "Text Encoders|${TEXT_ENC_DIR}" \
+    "Diffusion Models|${DIFFUSION_DIR}" \
+    "Upscale Models|${UPSCALE_DIR}" \
+    "Ultralytics|${ULTRALYTICS_DIR}"; do
+    label="${label_dir%%|*}"
+    dir="${label_dir#*|}"
+    count=$(find "$dir" -maxdepth 1 -type f 2>/dev/null | wc -l | tr -d ' ')
+    echo "  ${label}: ${count} files"
+done
 echo ""
-echo "Application will start now."
+echo "Application will start now."