diff --git a/setup-vastai.sh b/setup-vastai.sh index 721cda73..9d36d264 100755 --- a/setup-vastai.sh +++ b/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; } -dl() { - local url="$1" dest="$2" +die() { err "$*"; exit 1; } + +dl() { + local url="$1" dest="$2" name + name="$(basename "$dest")" if [ -f "$dest" ]; then - log " Already exists: $(basename "$dest")" - return 0 + 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")" - wget -q --show-progress -O "$dest" "$url" + log " Downloading: $name" + 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 - git lfs install > /dev/null 2>&1 - git lfs pull + 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 || 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..." -pip install --quiet --upgrade torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu124 +log "Installing PyTorch with CUDA..." +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" -ln -sf "$MODELS_DIR/clip/t5xxl_fp8_e4m3fn.safetensors" "$MODELS_DIR/text_encoders/t5xxl_fp8_e4m3fn.safetensors" +symlink "$MODELS_DIR/clip/clip_l.safetensors" "$MODELS_DIR/text_encoders/clip_l.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 "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"