Update wanForProduction.sh

wanForProduction.sh

@@ -1,13 +1,10 @@
 #!/usr/bin/env bash
-# setup_comfyui_simple.sh (base do CODE1, ajustado)
-# - troca de modelos/nodes pelos do CODE2
-# - adiciona SageAttention
-# - remove tratamento 5090/Blackwell
-# - melhora download_hf para paths com subpastas e /resolve/
+# setup_comfyui_wan22.sh
+# Script baseado no CODE1 com modelos/nodes do CODE2 e melhorias de robustez
 
 set -euo pipefail
 
-# Cores
+# Cores para output
 RED='\033[0;31m'
 GREEN='\033[0;32m'
 YELLOW='\033[1;33m'
@@ -20,44 +17,40 @@ log_warn() { echo -e "${YELLOW}[!]${NC} $1"; }
 log_error() { echo -e "${RED}[✗]${NC} $1"; }
 
 # Configuração
-COMFY_DIR="/root/comfy/ComfyUI"
+COMFY_DIR="${COMFY_DIR:-/root/comfy/ComfyUI}"
 MODELS_DIR="$COMFY_DIR/models"
-COMFY_HOST="0.0.0.0"
-COMFY_PORT="8818"
-CIVITAI_TOKEN="4fcb2834969399006a736ee402b061e5"
+COMFY_HOST="${COMFY_HOST:-0.0.0.0}"
+COMFY_PORT="${COMFY_PORT:-8818}"
+CIVITAI_TOKEN="${CIVITAI_TOKEN:-}"
+HF_TOKEN="${HF_TOKEN:-}"
 
 # Performance
-export MAX_JOBS=32
+export MAX_JOBS="${MAX_JOBS:-32}"
 export HF_HUB_ENABLE_HF_TRANSFER=1
+export HF_TRANSFER_CONCURRENCY="${HF_TRANSFER_CONCURRENCY:-16}"
+export NVCC_APPEND_FLAGS="${NVCC_APPEND_FLAGS:---threads 8}"
+export UV_SYSTEM_PYTHON=1
+export PYTORCH_CUDA_ALLOC_CONF="expandable_segments:True"
 
-# ===============================
-# Lista de downloads (do CODE2)
-# ===============================
-# Formato: "URL|TIPO|NOME_OPCIONAL"
+# Lista de downloads (formato: URL|TIPO|NOME_OPCIONAL)
 DOWNLOAD_FILES=(
-    # Wan 2.2 (Kijai) - use URLs SEM /resolve/ para usar o CLI da HF
-    "https://huggingface.co/Kijai/WanVideo_comfy/Wan22-Lightning/Wan2.2-Lightning_T2V-v1.1-A14B-4steps-lora_HIGH_fp16.safetensors|loras|Wan2.2-Lightning_T2V-v1.1-A14B-4steps-lora_HIGH_fp16.safetensors"
-    "https://huggingface.co/Kijai/WanVideo_comfy/Wan22-Lightning/Wan2.2-Lightning_T2V-v1.1-A14B-4steps-lora_LOW_fp16.safetensors|loras|Wan2.2-Lightning_T2V-v1.1-A14B-4steps-lora_LOW_fp16.safetensors"
-    "https://huggingface.co/Kijai/WanVideo_comfy_fp8_scaled/I2V/Wan2_2-I2V-A14B-HIGH_fp8_e5m2_scaled_KJ.safetensors|diffusion_models|Wan2_2-I2V-A14B-HIGH_fp8_e5m2_scaled_KJ.safetensors"
-    "https://huggingface.co/Kijai/WanVideo_comfy_fp8_scaled/I2V/Wan2_2-I2V-A14B-LOW_fp8_e5m2_scaled_KJ.safetensors|diffusion_models|Wan2_2-I2V-A14B-LOW_fp8_e5m2_scaled_KJ.safetensors"
-    "https://huggingface.co/Kijai/WanVideo_comfy/Wan22-Lightning/Wan2.2-Lightning_I2V-A14B-4steps-lora_HIGH_fp16.safetensors|loras|Wan2.2-Lightning_I2V-A14B-4steps-lora_HIGH_fp16.safetensors"
-    "https://huggingface.co/Kijai/WanVideo_comfy/Wan22-Lightning/Wan2.2-Lightning_I2V-A14B-4steps-lora_LOW_fp16.safetensors|loras|Wan2.2-Lightning_I2V-A14B-4steps-lora_LOW_fp16.safetensors"
-    "https://huggingface.co/Kijai/WanVideo_comfy/Wan2_1_VAE_fp32.safetensors|vae|Wan2_1_VAE_fp32.safetensors"
-    "https://huggingface.co/Kijai/WanVideo_comfy/umt5-xxl-enc-bf16.safetensors|text_encoders|umt5-xxl-enc-bf16.safetensors"
-
-    # Wan 2.1
-    "https://huggingface.co/Kijai/WanVideo_comfy/Lightx2v/lightx2v_I2V_14B_480p_cfg_step_distill_rank256_bf16.safetensors|loras|lightx2v_I2V_14B_480p_cfg_step_distill_rank256_bf16.safetensors"
-    "https://huggingface.co/Comfy-Org/Wan_2.2_ComfyUI_Repackaged/split_files/vae/wan_2.1_vae.safetensors|vae|wan_2.1_vae.safetensors"
-    "https://huggingface.co/Comfy-Org/Wan_2.2_ComfyUI_Repackaged/split_files/text_encoders/umt5_xxl_fp16.safetensors|text_encoders|umt5_xxl_fp16.safetensors"
-
-    # Wan 2.1 clip vision
-    "https://huggingface.co/Kijai/WanVideo_comfy/open-clip-xlm-roberta-large-vit-huge-14_visual_fp32.safetensors|clip_vision|open-clip-xlm-roberta-large-vit-huge-14_visual_fp32.safetensors"
-    "https://huggingface.co/Comfy-Org/Wan_2.1_ComfyUI_repackaged/split_files/clip_vision/clip_vision_h.safetensors|clip_vision|clip_vision_h.safetensors"
-
-    # Upscaler extra (com /resolve/ para baixar via HTTP direto)
-    "https://huggingface.co/ABDALLALSWAITI/Upscalers/resolve/main/anime/2x-AnimeSharpV2_MoSR_Soft.pth|upscale_models|2x-AnimeSharpV2_MoSR_Soft.pth"
-
-    # LoRAs via Mega.nz (do CODE2) — prefixo mega:// para usar função existente
+    # Wan 2.2 principais (HF)
+    "hf://Kijai/WanVideo_comfy/Wan22-Lightning/Wan2.2-Lightning_T2V-v1.1-A14B-4steps-lora_HIGH_fp16.safetensors|loras|Wan2.2-Lightning_T2V-v1.1-A14B-4steps-lora_HIGH_fp16.safetensors"
+    "hf://Kijai/WanVideo_comfy/Wan22-Lightning/Wan2.2-Lightning_T2V-v1.1-A14B-4steps-lora_LOW_fp16.safetensors|loras|Wan2.2-Lightning_T2V-v1.1-A14B-4steps-lora_LOW_fp16.safetensors"
+    "hf://Kijai/WanVideo_comfy_fp8_scaled/I2V/Wan2_2-I2V-A14B-HIGH_fp8_e5m2_scaled_KJ.safetensors|diffusion_models|Wan2_2-I2V-A14B-HIGH_fp8_e5m2_scaled_KJ.safetensors"
+    "hf://Kijai/WanVideo_comfy_fp8_scaled/I2V/Wan2_2-I2V-A14B-LOW_fp8_e5m2_scaled_KJ.safetensors|diffusion_models|Wan2_2-I2V-A14B-LOW_fp8_e5m2_scaled_KJ.safetensors"
+    "hf://Kijai/WanVideo_comfy/Wan22-Lightning/Wan2.2-Lightning_I2V-A14B-4steps-lora_HIGH_fp16.safetensors|loras|Wan2.2-Lightning_I2V-A14B-4steps-lora_HIGH_fp16.safetensors"
+    "hf://Kijai/WanVideo_comfy/Wan22-Lightning/Wan2.2-Lightning_I2V-A14B-4steps-lora_LOW_fp16.safetensors|loras|Wan2.2-Lightning_I2V-A14B-4steps-lora_LOW_fp16.safetensors"
+    "hf://Kijai/WanVideo_comfy/Wan2_1_VAE_fp32.safetensors|vae|Wan2_1_VAE_fp32.safetensors"
+    "hf://Kijai/WanVideo_comfy/umt5-xxl-enc-bf16.safetensors|text_encoders|umt5-xxl-enc-bf16.safetensors"
+    "hf://Kijai/WanVideo_comfy/Lightx2v/lightx2v_I2V_14B_480p_cfg_step_distill_rank256_bf16.safetensors|loras|lightx2v_I2V_14B_480p_cfg_step_distill_rank256_bf16.safetensors"
+    "hf://Comfy-Org/Wan_2.2_ComfyUI_Repackaged/split_files/vae/wan_2.1_vae.safetensors|vae|wan_2.1_vae.safetensors"
+    "hf://Comfy-Org/Wan_2.2_ComfyUI_Repackaged/split_files/text_encoders/umt5_xxl_fp16.safetensors|text_encoders|umt5_xxl_fp16.safetensors"
+    "hf://Kijai/WanVideo_comfy/open-clip-xlm-roberta-large-vit-huge-14_visual_fp32.safetensors|clip_vision|open-clip-xlm-roberta-large-vit-huge-14_visual_fp32.safetensors"
+    "hf://Comfy-Org/Wan_2.1_ComfyUI_repackaged/split_files/clip_vision/clip_vision_h.safetensors|clip_vision|clip_vision_h.safetensors"
+    "hf://ABDALLALSWAITI/Upscalers/anime/2x-AnimeSharpV2_MoSR_Soft.pth|upscale_models|2x-AnimeSharpV2_MoSR_Soft.pth"
+
+    # Wan 2.2 extras (Mega.nz)
     "mega://https://mega.nz/file/QFxREJpZ#0dge8TyRgb52lMFKdsk-Ac-rwcoG49hvBHvaqNsZkuk|loras|"
     "mega://https://mega.nz/file/lIhQETjT#nKlQ01Euw6fUxo8P5bfwlQrxAkWEedXQQOqHv5u4yrM|loras|"
     "mega://https://mega.nz/file/UIAESACY#Qlei1Pj5Nwno3Sz-wsTekvD-YgCkN1A5QbFKoRHXG9E|loras|"
@@ -69,15 +62,11 @@ DOWNLOAD_FILES=(
     "mega://https://mega.nz/file/wcAFhRJZ#ifZcckts0wPq1FG693tYLkpFrvjXvab1ZQeL6DuL1T8|loras|"
     "mega://https://mega.nz/file/wNIgmRwb#vn6rKM3QuTDpWVTLEBDsKjIkNkaX2wLKeauKNeVkcFs|loras|"
     "mega://https://mega.nz/file/hZhhhCQS#4je0Im2g1kEvTCuPotjVWkA-g6XICkqMb-mKxMt6R3Y|loras|"
-
-    # Upscale via Mega.nz (do CODE2)
     "mega://https://mega.nz/file/qdQwAJSC#rPY2Hg4x-_RbHiHUBFW-YU8xZUqPpq5Gpl1uPFiJvW4|upscale_models|"
     "mega://https://mega.nz/file/KIh2marA#CehVNXxQhGfp1bqUL0K1YPWNqrCUoLPCZGq2W6XqmWo|upscale_models|"
 )
 
-# ===============================
 # Custom nodes (do CODE2)
-# ===============================
 CUSTOM_NODES=(
     "https://github.com/kijai/ComfyUI-Florence2"
     "https://github.com/adbrasi/groqrouter"
@@ -105,63 +94,101 @@ CUSTOM_NODES=(
     "https://github.com/grmchn/ComfyUI-ProportionChanger"
 )
 
-# ===============================
-# Funções de download (mesmas do CODE1, porém download_hf melhorado)
-# ===============================
-
 download_hf() {
-    local url="$1"
+    local hf_path="$1"
     local target_dir="$2"
     local filename="$3"
 
-    if [ -n "$filename" ] && [ -f "$target_dir/$filename" ]; then
-        log_success "Já existe: $filename"
-        return 0
+    if [[ -z "$filename" ]]; then
+        filename="$(basename "$hf_path")"
     fi
 
-    log_info "Baixando HF: ${filename:-$(basename "$url")}"
+    mkdir -p "$target_dir"
+    if [[ -f "$target_dir/$filename" ]]; then
+        local existing_size
+        existing_size=$(stat -c%s "$target_dir/$filename" 2>/dev/null || echo 0)
+        if [[ "$existing_size" -gt 1048576 ]]; then
+            log_success "Já existe: $filename"
+            return 0
+        fi
+    fi
 
-    # Se a URL vier com /resolve/, usa HTTP direto (evita problemas com 'hf download')
-    if [[ "$url" == *"/resolve/"* ]]; then
-        mkdir -p "$target_dir"
-        wget -q --show-progress -c -O "$target_dir/$filename" "https://huggingface.co/$url"
-        return $?
+    IFS='/' read -r owner repo rest <<< "$hf_path"
+    if [[ -z "$owner" || -z "$repo" || -z "$rest" ]]; then
+        log_error "Caminho HuggingFace inválido: $hf_path"
+        return 1
     fi
 
-    # Parse correto: repo_id = 2 primeiros segmentos; file_path = resto
-    local repo_id file_path
-    repo_id="$(echo "$url" | cut -d'/' -f1-2)"
-    file_path="$(echo "$url" | cut -d'/' -f3-)"
+    local repo_id="$owner/$repo"
+    local attempt success temp_dir
+    for attempt in 1 2 3; do
+        temp_dir="$(mktemp -d)"
+        success=false
 
-    if command -v hf >/dev/null 2>&1; then
-        mkdir -p "$target_dir"
-        (cd "$target_dir" && HF_HUB_ENABLE_HF_TRANSFER=1 hf download "$repo_id" "$file_path" --local-dir . >/dev/null 2>&1) || {
-            log_warn "hf download falhou; tentando HTTP"
-            wget -q --show-progress -c -O "$target_dir/$filename" "https://huggingface.co/$repo_id/resolve/main/$file_path" || \
-            log_error "Download falhou: $filename"
-        }
-    else
-        mkdir -p "$target_dir"
-        wget -q --show-progress -c -O "$target_dir/$filename" "https://huggingface.co/$repo_id/resolve/main/$file_path" || \
-        log_error "Download falhou: $filename"
-    fi
-}
+        if command -v hf >/dev/null 2>&1; then
+            if HF_HUB_ENABLE_HF_TRANSFER=1 hf download "$repo_id" "$rest" --local-dir "$temp_dir" --local-dir-use-symlinks False; then
+                success=true
+            fi
+        fi
 
-download_mega() {
-    local url="$1"
-    local target_dir="$2"
+        if [[ "$success" == false ]] && command -v huggingface-cli >/dev/null 2>&1; then
+            if HF_HUB_ENABLE_HF_TRANSFER=1 huggingface-cli download "$repo_id" "$rest" --local-dir "$temp_dir" --local-dir-use-symlinks False; then
+                success=true
+            fi
+        fi
 
-    url="${url#mega://}" # remove prefixo
+        if [[ "$success" == false ]]; then
+            local direct_url="https://huggingface.co/${repo_id}/resolve/main/${rest}"
+            if command -v aria2c >/dev/null 2>&1; then
+                local aria_cmd=(aria2c -c -s 16 -x 16 -k 1M --console-log-level=warn --summary-interval=10 --dir "$temp_dir" --out "$filename")
+                if [[ -n "$HF_TOKEN" ]]; then
+                    aria_cmd+=(--header="Authorization: Bearer $HF_TOKEN")
+                fi
+                aria_cmd+=("$direct_url")
+                if "${aria_cmd[@]}"; then
+                    success=true
+                fi
+            fi
+            if [[ "$success" == false ]] && command -v curl >/dev/null 2>&1; then
+                local curl_cmd=(curl -L -# -o "$temp_dir/$filename")
+                if [[ -n "$HF_TOKEN" ]]; then
+                    curl_cmd+=(-H "Authorization: Bearer $HF_TOKEN")
+                fi
+                curl_cmd+=("$direct_url")
+                if "${curl_cmd[@]}"; then
+                    success=true
+                fi
+            fi
+            if [[ "$success" == false ]] && command -v wget >/dev/null 2>&1; then
+                local wget_cmd=(wget -q --show-progress -c -O "$temp_dir/$filename")
+                if [[ -n "$HF_TOKEN" ]]; then
+                    wget_cmd+=(--header="Authorization: Bearer $HF_TOKEN")
+                fi
+                wget_cmd+=("$direct_url")
+                if "${wget_cmd[@]}"; then
+                    success=true
+                fi
+            fi
+        fi
 
-    log_info "Baixando do Mega..."
+        if [[ "$success" == true ]]; then
+            local downloaded
+            downloaded="$(find "$temp_dir" -maxdepth 1 -type f | head -n 1)"
+            if [[ -n "$downloaded" ]]; then
+                mv "$downloaded" "$target_dir/$filename"
+                rm -rf "$temp_dir"
+                log_success "Download concluído: $filename"
+                return 0
+            fi
+        fi
 
-    cd "$target_dir"
-    if megadl "$url" 2>/dev/null; then
-        log_success "Mega download OK"
-    else
-        log_warn "Mega download falhou"
-    fi
-    cd - >/dev/null
+        rm -rf "$temp_dir"
+        log_warn "Tentativa $attempt falhou para $filename"
+        sleep $((2 * attempt))
+    done
+
+    log_error "Falha ao baixar: $filename"
+    return 1
 }
 
 download_file() {
@@ -169,156 +196,241 @@ download_file() {
     local target_dir="$2"
     local filename="$3"
 
-    # Adicionar token Civitai se necessário
-    if [[ "$url" == *"civitai.com"* ]] && [[ "$url" != *"token="* ]]; then
-        url="${url}&token=${CIVITAI_TOKEN}"
+    if [[ "$url" == *"civitai.com/api/download"* ]] && [[ "$url" != *"token="* ]] && [[ -n "$CIVITAI_TOKEN" ]]; then
+        if [[ "$url" == *"?"* ]]; then
+            url="${url}&token=${CIVITAI_TOKEN}"
+        else
+            url="${url}?token=${CIVITAI_TOKEN}"
+        fi
     fi
+    mkdir -p "$target_dir"
 
-    if [ -n "$filename" ] && [ -f "$target_dir/$filename" ]; then
-        log_success "Já existe: $filename"
-        return 0
+    local final_name="$filename"
+    if [[ -z "$final_name" ]]; then
+        final_name="$(basename "${url%%\?*}")"
     fi
 
-    log_info "Baixando: ${filename:-arquivo}"
+    if [[ -f "$target_dir/$final_name" ]]; then
+        local existing_size
+        existing_size=$(stat -c%s "$target_dir/$final_name" 2>/dev/null || echo 0)
+        if [[ "$existing_size" -gt 524288 ]]; then
+            log_success "Já existe: $final_name"
+            return 0
+        fi
+    fi
 
-    if command -v aria2c >/dev/null 2>&1; then
-        if [ -n "$filename" ]; then
-            aria2c -c -x 4 -s 4 --console-log-level=warn --dir="$target_dir" --out="$filename" "$url" || \
-            wget -q --show-progress -c -O "$target_dir/$filename" "$url" || \
-            log_error "Download falhou: $filename"
-        else
-            aria2c -c -x 4 -s 4 --console-log-level=warn --dir="$target_dir" "$url" || \
-            (cd "$target_dir" && wget -q --show-progress -c "$url") || \
-            log_error "Download falhou: $url"
+    local attempt success
+    for attempt in 1 2 3; do
+        success=false
+        if command -v aria2c >/dev/null 2>&1; then
+            if aria2c -c -s 16 -x 16 -k 1M --console-log-level=warn --summary-interval=10 --dir "$target_dir" --out "$final_name" "$url"; then
+                success=true
+            fi
         fi
-    elif [ -n "$filename" ]; then
-        wget -q --show-progress -c -O "$target_dir/$filename" "$url" || \
-        log_error "Download falhou: $filename"
-    else
-        (cd "$target_dir" && wget -q --show-progress -c "$url") || \
-        log_error "Download falhou: $url"
+        if [[ "$success" == false ]] && command -v wget >/dev/null 2>&1; then
+            if wget -q --show-progress -c -O "$target_dir/$final_name" "$url"; then
+                success=true
+            fi
+        fi
+        if [[ "$success" == false ]] && command -v curl >/dev/null 2>&1; then
+            if curl -L -# -C - -o "$target_dir/$final_name" "$url"; then
+                success=true
+            fi
+        fi
+        if [[ "$success" == true ]]; then
+            log_success "Download concluído: $final_name"
+            return 0
+        fi
+        log_warn "Tentativa $attempt falhou para $final_name"
+        sleep $((2 * attempt))
+    done
+
+    log_error "Falha ao baixar: $final_name"
+    return 1
+}
+
+download_mega() {
+    local url="$1"
+    local target_dir="$2"
+
+    url="${url#mega://}"
+    mkdir -p "$target_dir"
+
+    if ! command -v megadl >/dev/null 2>&1; then
+        log_warn "megadl não encontrado; tente instalar megatools manualmente"
+        return 1
     fi
+
+    local attempt
+    for attempt in 1 2 3; do
+        if (cd "$target_dir" && megadl "$url"); then
+            log_success "Download Mega.nz concluído"
+            return 0
+        fi
+        log_warn "Tentativa $attempt falhou para Mega.nz"
+        sleep $((2 * attempt))
+    done
+
+    log_error "Falha ao baixar de Mega.nz: $url"
+    return 1
 }
 
 process_downloads() {
+    local failed=0
     for entry in "${DOWNLOAD_FILES[@]}"; do
-        IFS='|' read -r url type filename <<< "$entry"
-
+        IFS='|' read -r raw_url type filename <<< "$entry"
         local target_dir="$MODELS_DIR/$type"
         mkdir -p "$target_dir"
 
-        if [[ "$url" == mega://* ]]; then
-            download_mega "$url" "$target_dir"
-        elif [[ "$url" == *.huggingface.co/* ]]; then
-            local hf_path="${url#*huggingface.co/}"
-            download_hf "$hf_path" "$target_dir" "$filename"
+        if [[ "$raw_url" == mega://* ]]; then
+            if ! download_mega "$raw_url" "$target_dir"; then
+                ((failed++))
+            fi
+        elif [[ "$raw_url" == hf://* ]]; then
+            local hf_path="${raw_url#hf://}"
+            if ! download_hf "$hf_path" "$target_dir" "$filename"; then
+                ((failed++))
+            fi
         else
-            download_file "$url" "$target_dir" "$filename"
+            if ! download_file "$raw_url" "$target_dir" "$filename"; then
+                ((failed++))
+            fi
         fi
     done
+
+    if (( failed > 0 )); then
+        return 1
+    fi
+    return 0
 }
 
 clone_repo() {
     local url="$1"
     local dest="$2"
 
-    if [ -d "$dest/.git" ]; then
-        git -C "$dest" pull --ff-only 2>/dev/null || true
+    if [[ -d "$dest/.git" ]]; then
+        git -C "$dest" pull --ff-only 2>/dev/null || log_warn "Falha ao atualizar $(basename "$dest")"
     else
-        git clone --depth 1 "$url" "$dest" 2>/dev/null || true
+        git clone --depth 1 "$url" "$dest" 2>/dev/null || log_warn "Falha ao clonar $(basename "$dest")"
     fi
 
-    if [ -f "$dest/requirements.txt" ]; then
-        pip install -q -r "$dest/requirements.txt" 2>/dev/null || true
+    if [[ -f "$dest/requirements.txt" ]]; then
+        pip install -q -r "$dest/requirements.txt" 2>/dev/null || log_warn "Falha ao instalar requirements de $(basename "$dest")"
     fi
 }
 
 # ========== INSTALAÇÃO ==========
-
 log_info "========================================="
-log_info " ComfyUI Setup"
+log_info " ComfyUI + Wan 2.2 Setup"
 log_info "========================================="
 
 # 1. Dependências
-log_info "[1/6] Instalando dependências..."
-
-apt-get update -qq
-apt-get install -y -qq aria2 megatools git wget curl 2>/dev/null
-
-pip install --upgrade pip wheel setuptools -q
-pip install --upgrade "huggingface_hub[cli,hf_transfer]" comfy-cli -q
-
-log_success "Dependências instaladas"
+log_info "[1/6] Instalando dependências do sistema..."
+if command -v apt-get >/dev/null 2>&1; then
+    apt-get update -qq || log_warn "apt-get update falhou; seguindo assim mesmo"
+    apt-get install -y -qq aria2 megatools git wget curl || log_warn "apt-get install retornou erro"
+else
+    log_warn "apt-get indisponível; certifique-se de que aria2, megatools, git, wget e curl estão instalados"
+fi
+pip install --upgrade pip wheel setuptools -q || log_warn "Falha na atualização do pip"
+pip install --upgrade "huggingface_hub[cli,hf_transfer]" comfy-cli -q || log_warn "Falha instalando huggingface_hub/comfy-cli"
+log_success "Dependências tratadas"
 
-# 2. ComfyUI
+# 2. Instalar ComfyUI
 log_info "[2/6] Instalando ComfyUI..."
-
-if [ -f "$COMFY_DIR/main.py" ]; then
-    log_warn "ComfyUI já existe"
+if [[ -f "$COMFY_DIR/main.py" ]]; then
+    log_warn "ComfyUI já presente"
 else
-    comfy --skip-prompt install --fast-deps --nvidia
+    comfy --skip-prompt install --fast-deps --nvidia || log_error "Instalação do ComfyUI falhou"
 fi
+log_success "ComfyUI pronto"
 
-log_success "ComfyUI instalado"
+# 3. Instalar PyTorch (CUDA 12.8)
+log_info "[3/6] Instalando PyTorch CUDA 12.8..."
+python3 -m pip install --upgrade torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu128 || log_error "Falha ao instalar PyTorch"
+log_success "PyTorch configurado"
 
-# 3. PyTorch CUDA 12.8 (substitui o bloco de Blackwell/5090 do CODE1)
-log_info "[3/6] Instalando PyTorch (CUDA 12.8)..."
-pip install --upgrade torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu128 -q
-log_success "PyTorch com CUDA 12.8 instalado"
-
-# 4. SageAttention (do CODE2)
+# 4. Instalar SageAttention (condicional)
 log_info "[4/6] Instalando SageAttention..."
-if python3 - <<'PY' 2>/dev/null; then
+if python3 - <<'PY' >/dev/null 2>&1
 import torch, sys
 sys.exit(0 if torch.cuda.is_available() else 1)
 PY
 then
-    pip install -q "https://huggingface.co/adbrasi/comfywheel/resolve/main/sageattention-2.2.0-cp311-cp311-linux_x86_64.whl" || log_warn "Falha ao instalar SageAttention"
-    if python3 -c 'import sageattention' >/dev/null 2>&1; then
-        log_success "SageAttention instalado"
+    if python3 -m pip install "https://huggingface.co/adbrasi/comfywheel/resolve/main/sageattention-2.2.0-cp311-cp311-linux_x86_64.whl" -q; then
+        if python3 - <<'PY' >/dev/null 2>&1
+import sageattention
+PY
+        then
+            log_success "SageAttention instalado"
+        else
+            log_warn "SageAttention instalado, mas não pôde ser importado"
+        fi
     else
-        log_warn "SageAttention não pôde ser importado (seguindo sem)"
+        log_warn "Falha ao instalar SageAttention"
     fi
 else
-    log_warn "CUDA não detectada, pulando SageAttention"
+    log_warn "CUDA indisponível; pulando SageAttention"
 fi
 
-# 5. Modelos
+# 5. Downloads de modelos
 log_info "[5/6] Baixando modelos..."
-process_downloads
-log_success "Downloads processados"
+mkdir -p "$MODELS_DIR"
+if process_downloads; then
+    log_success "Downloads concluídos (ou já presentes)"
+else
+    log_warn "Alguns downloads falharam; reexecute o script para tentar novamente"
+fi
 
 # 6. Custom nodes
 log_info "[6/6] Instalando custom nodes..."
-
 CN_DIR="$COMFY_DIR/custom_nodes"
 mkdir -p "$CN_DIR"
-
-total_nodes=${#CUSTOM_NODES[@]}
-current_node=0
-
+node_total=${#CUSTOM_NODES[@]}
+node_index=0
 for repo in "${CUSTOM_NODES[@]}"; do
-    ((current_node++))
-    node_name=$(basename "$repo")
-    log_info "[$current_node/$total_nodes] Instalando: $node_name"
-    clone_repo "$repo" "$CN_DIR/$node_name"
+    ((node_index++))
+    log_info "[$node_index/$node_total] $(basename "$repo")"
+    clone_repo "$repo" "$CN_DIR/$(basename "$repo")"
 done
+log_success "Custom nodes processados"
 
-log_success "Custom nodes instalados"
-
-# Iniciar
-log_info "========================================="
-log_success "Instalação concluída!"
-log_info "Iniciando ComfyUI..."
-log_info "URL: http://localhost:$COMFY_PORT"
+# Resumo final
 log_info "========================================="
+log_info "Verificando arquivos críticos..."
+critical_missing=0
+for required in \
+    "$COMFY_DIR/main.py" \
+    "$MODELS_DIR/vae/wan_2.1_vae.safetensors" \
+    "$MODELS_DIR/vae/Wan2_1_VAE_fp32.safetensors" \
+    "$MODELS_DIR/text_encoders/umt5_xxl_fp16.safetensors" \
+    "$MODELS_DIR/text_encoders/umt5-xxl-enc-bf16.safetensors" \
+    "$MODELS_DIR/clip_vision/clip_vision_h.safetensors"; do
+    if [[ -f "$required" ]]; then
+        log_success "✓ $(basename "$required")"
+    else
+        log_warn "✗ $(basename "$required") não encontrado"
+        ((critical_missing++))
+    fi
+done
 
-cd "$COMFY_DIR"
+if (( critical_missing == 0 )); then
+    log_success "Instalação concluída com sucesso!"
+else
+    log_warn "Instalação finalizada com pendências; verifique avisos acima"
+fi
 
-# Habilita SageAttention se presente
 SAGE_FLAG=""
-if python3 -c 'import sageattention' >/dev/null 2>&1; then
+if python3 - <<'PY' >/dev/null 2>&1
+import sageattention
+PY
+then
     SAGE_FLAG="--use-sage-attention"
 fi
 
+log_info "========================================="
+log_info "Iniciando ComfyUI em http://${COMFY_HOST}:${COMFY_PORT}"
+log_info "========================================="
+
+cd "$COMFY_DIR"
 exec comfy launch -- ${SAGE_FLAG:+$SAGE_FLAG} --listen "$COMFY_HOST" --port "$COMFY_PORT"