Harden setup script with better error handling

- Downloads to .part files, moves on success (no corrupt files)
- Skips existing non-empty files, removes empty/corrupt ones
- Ctrl+C trap cleans up partial downloads
- Preflight checks for required commands and directory
- CUDA verification after PyTorch install
- Symlink helper that warns if target is missing
- Summary at end lists all failed downloads and node deps
- Removed set -e in favor of explicit error handling

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

setup-vastai.sh

@@ -1,5 +1,5 @@
 #!/usr/bin/env bash
-set -euo pipefail
+set -uo pipefail
 
 # ── ComfyUI vast.ai Setup Script ──────────────────────────────
 # Run from inside the cloned ComfyUI directory.
@@ -9,49 +9,117 @@ set -euo pipefail
 
 SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
 MODELS_DIR="$SCRIPT_DIR/models"
+FAILED_DOWNLOADS=()
+FAILED_NODES=()
+CURRENT_DOWNLOAD=""
+
+cleanup() {
+    echo ""
+    if [ -n "$CURRENT_DOWNLOAD" ] && [ -f "$CURRENT_DOWNLOAD" ]; then
+        rm -f "$CURRENT_DOWNLOAD"
+        err "Cleaned up partial download: $(basename "$CURRENT_DOWNLOAD")"
+    fi
+    err "Setup interrupted. Re-run the script to resume — completed downloads will be skipped."
+    exit 130
+}
+trap cleanup INT TERM
 
 # ── Helpers ───────────────────────────────────────────────────
 log()  { echo -e "\033[1;32m[setup]\033[0m $*"; }
+warn() { echo -e "\033[1;33m[warn]\033[0m  $*" >&2; }
 err()  { echo -e "\033[1;31m[error]\033[0m $*" >&2; }
+die()  { err "$*"; exit 1; }
+
 dl() {
-    local url="$1" dest="$2"
+    local url="$1" dest="$2" name
+    name="$(basename "$dest")"
     if [ -f "$dest" ]; then
-        log "  Already exists: $(basename "$dest")"
+        local size
+        size=$(stat -c%s "$dest" 2>/dev/null || stat -f%z "$dest" 2>/dev/null || echo 0)
+        if [ "$size" -gt 0 ]; then
+            log "  Skip (exists): $name"
             return 0
+        else
+            warn "  Removing empty file: $name"
+            rm -f "$dest"
+        fi
     fi
     mkdir -p "$(dirname "$dest")"
-    log "  Downloading: $(basename "$dest")"
+    log "  Downloading: $name"
-    wget -q --show-progress -O "$dest" "$url"
+    local tmp="${dest}.part"
+    CURRENT_DOWNLOAD="$tmp"
+    if wget -q --show-progress -O "$tmp" "$url"; then
+        mv "$tmp" "$dest"
+        CURRENT_DOWNLOAD=""
+    else
+        rm -f "$tmp"
+        CURRENT_DOWNLOAD=""
+        err "  Failed to download: $name"
+        FAILED_DOWNLOADS+=("$name ($url)")
+        return 1
+    fi
 }
 
+symlink() {
+    local target="$1" link="$2"
+    mkdir -p "$(dirname "$link")"
+    if [ ! -f "$target" ]; then
+        warn "  Symlink target missing: $(basename "$target") — skipping"
+        return 1
+    fi
+    ln -sf "$target" "$link"
+}
+
+# ── Preflight Checks ────────────────────────────────────────
 cd "$SCRIPT_DIR"
 
+[ -f "main.py" ] || die "Not in a ComfyUI directory (main.py not found). Run from inside the repo."
+
+for cmd in git python3 wget; do
+    command -v "$cmd" &>/dev/null || die "'$cmd' is required but not found."
+done
+
+log "Starting ComfyUI setup in: $SCRIPT_DIR"
+
 # ── Git LFS ──────────────────────────────────────────────────
-if command -v git-lfs &>/dev/null || git lfs version &>/dev/null; then
+if git lfs version &>/dev/null; then
     log "Pulling LFS files (Civitai loras)..."
     git lfs install > /dev/null 2>&1
-    git lfs pull
+    git lfs pull || warn "LFS pull failed — Civitai loras may be missing"
 else
     log "git-lfs not found, installing..."
-    apt-get update -qq && apt-get install -y -qq git-lfs > /dev/null 2>&1
+    if apt-get update -qq && apt-get install -y -qq git-lfs > /dev/null 2>&1; then
         git lfs install > /dev/null 2>&1
-    git lfs pull
+        git lfs pull || warn "LFS pull failed — Civitai loras may be missing"
+    else
+        warn "Could not install git-lfs — Civitai loras will not be available"
+    fi
 fi
 
 # ── Python Environment ───────────────────────────────────────
 log "Setting up Python environment..."
 if [ ! -d "$SCRIPT_DIR/venv" ]; then
-    python3 -m venv venv
+    python3 -m venv venv || die "Failed to create virtual environment"
 fi
+# shellcheck disable=SC1091
 source venv/bin/activate
 
 # Install PyTorch with CUDA
-log "Installing PyTorch..."
+log "Installing PyTorch with CUDA..."
-pip install --quiet --upgrade torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu124
+pip install --quiet --upgrade torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu124 \
+    || die "Failed to install PyTorch — cannot continue without it"
+
+# Verify CUDA is available
+if python3 -c "import torch; assert torch.cuda.is_available(), 'no cuda'" 2>/dev/null; then
+    log "PyTorch CUDA: OK ($(python3 -c 'import torch; print(torch.cuda.get_device_name(0))'))"
+else
+    warn "PyTorch installed but CUDA not available — generation will be very slow"
+fi
 
 # Install ComfyUI requirements
 log "Installing ComfyUI requirements..."
-pip install --quiet --upgrade -r requirements.txt
+pip install --quiet --upgrade -r requirements.txt \
+    || die "Failed to install ComfyUI requirements"
 
 # ── Custom Node Requirements ─────────────────────────────────
 log "Installing custom node requirements..."
@@ -59,19 +127,23 @@ for req in custom_nodes/*/requirements.txt; do
     [ -f "$req" ] || continue
     node_name="$(basename "$(dirname "$req")")"
     log "  $node_name"
-    pip install --quiet --upgrade -r "$req" 2>/dev/null || err "  Failed some deps in $node_name (may be OK)"
+    if ! pip install --quiet --upgrade -r "$req" 2>/dev/null; then
+        warn "  Some deps failed for $node_name"
+        FAILED_NODES+=("$node_name")
+    fi
 done
 
 # ── Download Models ───────────────────────────────────────────
 log "Downloading models from HuggingFace..."
+log ""
 
 # Checkpoints
-log "Checkpoints..."
+log "[ Checkpoints ]"
 dl "https://huggingface.co/Comfy-Org/flux1-dev/resolve/main/flux1-dev-fp8.safetensors" \
    "$MODELS_DIR/checkpoints/flux1-dev-fp8.safetensors"
 
 # CLIP / Text Encoders
-log "CLIP & Text Encoders..."
+log "[ CLIP & Text Encoders ]"
 dl "https://huggingface.co/comfyanonymous/flux_text_encoders/resolve/main/clip_l.safetensors" \
    "$MODELS_DIR/clip/clip_l.safetensors"
 dl "https://huggingface.co/comfyanonymous/flux_text_encoders/resolve/main/t5xxl_fp8_e4m3fn.safetensors" \
@@ -79,52 +151,52 @@ dl "https://huggingface.co/comfyanonymous/flux_text_encoders/resolve/main/t5xxl_
 dl "https://huggingface.co/QuanSun/EVA-CLIP/resolve/main/EVA02_CLIP_L_336_psz14_s6B.pt" \
    "$MODELS_DIR/clip/EVA02_CLIP_L_336_psz14_s6B.pt"
 
-# Duplicate clip/text_encoders (same files, different paths)
+# Symlink text_encoders -> clip (same files)
 mkdir -p "$MODELS_DIR/text_encoders"
-ln -sf "$MODELS_DIR/clip/clip_l.safetensors" "$MODELS_DIR/text_encoders/clip_l.safetensors"
+symlink "$MODELS_DIR/clip/clip_l.safetensors" "$MODELS_DIR/text_encoders/clip_l.safetensors"
-ln -sf "$MODELS_DIR/clip/t5xxl_fp8_e4m3fn.safetensors" "$MODELS_DIR/text_encoders/t5xxl_fp8_e4m3fn.safetensors"
+symlink "$MODELS_DIR/clip/t5xxl_fp8_e4m3fn.safetensors" "$MODELS_DIR/text_encoders/t5xxl_fp8_e4m3fn.safetensors"
 
 # CLIP Vision
-log "CLIP Vision..."
+log "[ CLIP Vision ]"
 dl "https://huggingface.co/Comfy-Org/sigclip_vision_384/resolve/main/sigclip_vision_patch14_384.safetensors" \
    "$MODELS_DIR/clip_vision/sigclip_vision_patch14_384.safetensors"
 
 # ControlNet
-log "ControlNet..."
+log "[ ControlNet ]"
 dl "https://huggingface.co/Shakker-Labs/FLUX.1-dev-ControlNet-Union-Pro/resolve/main/diffusion_pytorch_model.safetensors" \
    "$MODELS_DIR/controlnet/flux_controlnet_pro.safetensors"
 
 # VAE
-log "VAE..."
+log "[ VAE ]"
 dl "https://huggingface.co/black-forest-labs/FLUX.1-schnell/resolve/main/ae.safetensors" \
    "$MODELS_DIR/vae/ae.safetensors"
 
 # Upscale
-log "Upscale models..."
+log "[ Upscale ]"
 dl "https://huggingface.co/uwg/upscaler/resolve/main/ESRGAN/4x_NMKD-Siax_200k.pth" \
    "$MODELS_DIR/upscale_models/4x_NMKD-Siax_200k.pth"
 
 # PuLID
-log "PuLID..."
+log "[ PuLID ]"
 dl "https://huggingface.co/guozinan/PuLID/resolve/main/pulid_flux_v0.9.0.safetensors" \
    "$MODELS_DIR/pulid/pulid_flux_v0.9.0.safetensors"
 
 # SAM / SAM2
-log "SAM models..."
+log "[ SAM ]"
 dl "https://huggingface.co/Kijai/sam2-safetensors/resolve/main/sam2_hiera_base_plus.safetensors" \
    "$MODELS_DIR/sam2/sam2_hiera_base_plus.safetensors"
 dl "https://huggingface.co/datasets/Gourieff/ReActor/resolve/main/models/sams/sam_vit_b_01ec64.pth" \
    "$MODELS_DIR/sams/sam_vit_b_01ec64.pth"
 
 # Facexlib
-log "Facexlib..."
+log "[ Facexlib ]"
 dl "https://huggingface.co/leonelhs/facexlib/resolve/main/detection_Resnet50_Final.pth" \
    "$MODELS_DIR/facexlib/detection_Resnet50_Final.pth"
 dl "https://huggingface.co/leonelhs/facexlib/resolve/main/parsing_bisenet.pth" \
    "$MODELS_DIR/facexlib/parsing_bisenet.pth"
 
 # InsightFace (antelopev2)
-log "InsightFace (antelopev2)..."
+log "[ InsightFace (antelopev2) ]"
 mkdir -p "$MODELS_DIR/insightface/models/antelopev2"
 dl "https://huggingface.co/DIAMONIK7777/antelopev2/resolve/main/1k3d68.onnx" \
    "$MODELS_DIR/insightface/models/antelopev2/1k3d68.onnx"
@@ -138,7 +210,7 @@ dl "https://huggingface.co/DIAMONIK7777/antelopev2/resolve/main/scrfd_10g_bnkps.
    "$MODELS_DIR/insightface/models/antelopev2/scrfd_10g_bnkps.onnx"
 
 # Ultralytics (YOLO)
-log "Ultralytics..."
+log "[ Ultralytics ]"
 dl "https://huggingface.co/Bingsu/adetailer/resolve/main/face_yolov8m.pt" \
    "$MODELS_DIR/ultralytics/face_yolov8m.pt"
 dl "https://huggingface.co/Bingsu/adetailer/resolve/main/hand_yolov8s.pt" \
@@ -147,14 +219,14 @@ dl "https://huggingface.co/Bingsu/adetailer/resolve/main/person_yolov8m-seg.pt"
    "$MODELS_DIR/ultralytics/person_yolov8m-seg.pt"
 
 # USO (LoRA + projector)
-log "USO models..."
+log "[ USO ]"
 dl "https://huggingface.co/Comfy-Org/USO_1.0_Repackaged/resolve/main/split_files/loras/uso-flux1-dit-lora-v1.safetensors" \
    "$MODELS_DIR/loras/uso-flux1-dit-lora-v1.safetensors"
 dl "https://huggingface.co/Comfy-Org/USO_1.0_Repackaged/resolve/main/split_files/model_patches/uso-flux1-projector-v1.safetensors" \
    "$MODELS_DIR/model_patches/uso-flux1-projector-v1.safetensors"
 
 # Florence-2-base (LLM)
-log "Florence-2-base..."
+log "[ Florence-2-base ]"
 mkdir -p "$MODELS_DIR/LLM/Florence-2-base"
 FLORENCE_BASE="https://huggingface.co/microsoft/Florence-2-base/resolve/main"
 for f in model.safetensors pytorch_model.bin config.json configuration_florence2.py \
@@ -164,18 +236,43 @@ for f in model.safetensors pytorch_model.bin config.json configuration_florence2
 done
 
 # ── Create Symlinks (matching local layout) ──────────────────
-log "Creating model symlinks..."
-mkdir -p "$MODELS_DIR/diffusion_models/FLUX" "$MODELS_DIR/unet" "$MODELS_DIR/xlabs/controlnets"
-ln -sf "$MODELS_DIR/checkpoints/flux1-dev-fp8.safetensors" "$MODELS_DIR/diffusion_models/flux1-dev-fp8.safetensors"
-ln -sf "$MODELS_DIR/checkpoints/flux1-dev-fp8.safetensors" "$MODELS_DIR/diffusion_models/FLUX/flux1-dev-fp8.safetensors"
-ln -sf "$MODELS_DIR/checkpoints/flux1-dev-fp8.safetensors" "$MODELS_DIR/unet/flux1-dev-fp8.safetensors"
-ln -sf "$MODELS_DIR/controlnet/flux_controlnet_pro.safetensors" "$MODELS_DIR/diffusion_models/flux_controlnet_pro.safetensors"
-ln -sf "$MODELS_DIR/controlnet/flux_controlnet_pro.safetensors" "$MODELS_DIR/xlabs/controlnets/flux_controlnet_pro.safetensors"
-
-# ── Done ──────────────────────────────────────────────────────
-log "Setup complete!"
 log ""
+log "Creating model symlinks..."
+symlink "$MODELS_DIR/checkpoints/flux1-dev-fp8.safetensors" "$MODELS_DIR/diffusion_models/flux1-dev-fp8.safetensors"
+symlink "$MODELS_DIR/checkpoints/flux1-dev-fp8.safetensors" "$MODELS_DIR/diffusion_models/FLUX/flux1-dev-fp8.safetensors"
+symlink "$MODELS_DIR/checkpoints/flux1-dev-fp8.safetensors" "$MODELS_DIR/unet/flux1-dev-fp8.safetensors"
+symlink "$MODELS_DIR/controlnet/flux_controlnet_pro.safetensors" "$MODELS_DIR/diffusion_models/flux_controlnet_pro.safetensors"
+symlink "$MODELS_DIR/controlnet/flux_controlnet_pro.safetensors" "$MODELS_DIR/xlabs/controlnets/flux_controlnet_pro.safetensors"
+
+# ── Summary ──────────────────────────────────────────────────
+echo ""
+echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
+
+if [ ${#FAILED_DOWNLOADS[@]} -gt 0 ]; then
+    err "Failed downloads (${#FAILED_DOWNLOADS[@]}):"
+    for f in "${FAILED_DOWNLOADS[@]}"; do
+        err "  - $f"
+    done
+    echo ""
+fi
+
+if [ ${#FAILED_NODES[@]} -gt 0 ]; then
+    warn "Custom nodes with dependency issues (${#FAILED_NODES[@]}):"
+    for n in "${FAILED_NODES[@]}"; do
+        warn "  - $n"
+    done
+    echo ""
+fi
+
+if [ ${#FAILED_DOWNLOADS[@]} -eq 0 ] && [ ${#FAILED_NODES[@]} -eq 0 ]; then
+    log "Setup complete — no errors!"
+else
+    log "Setup complete with issues (see above)."
+fi
+
+echo ""
 log "To start ComfyUI:"
 log "  cd $SCRIPT_DIR"
 log "  source venv/bin/activate"
 log "  python main.py --listen 0.0.0.0 --lowvram --fp8_e4m3fn-text-enc"
+echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"