Update image_Studio.sh

image_Studio.sh

@@ -1,49 +1,42 @@
 #!/usr/bin/env bash
-# setup_comfyui_custom_v2.sh
-# Script otimizado com suporte para RTX 5090, estado persistente e downloads do Mega
+# setup_comfyui_custom_v2.sh (FIX DL HANG + BLACKWELL)
+# Otimizado: timeouts/retries, Civitai headers, estado robusto, Blackwell nightly
 
 set -euo pipefail
 
-# Detectar se está rodando via pipe (não-interativo)
-if [ ! -t 0 ] || [ ! -t 1 ]; then
-    INTERACTIVE=false
-else
-    INTERACTIVE=true
-fi
-
-# -----------------------------
-# Cores para output
-# -----------------------------
-RED='\033[0;31m'
-GREEN='\033[0;32m'
-YELLOW='\033[1;33m'
-BLUE='\033[0;34m'
-CYAN='\033[0;36m'
-NC='\033[0m'
+# --------------------------------
+# Modo interativo?
+# --------------------------------
+if [ ! -t 0 ] || [ ! -t 1 ]; then INTERACTIVE=false; else INTERACTIVE=true; fi
 
+# --------------------------------
+# Cores
+# --------------------------------
+RED='\033[0;31m'; GREEN='\033[0;32m'; YELLOW='\033[1;33m'; BLUE='\033[0;34m'; CYAN='\033[0;36m'; NC='\033[0m'
 log_info() { echo -e "${BLUE}[INFO]${NC} $1"; }
 log_success() { echo -e "${GREEN}[✓]${NC} $1"; }
 log_warn() { echo -e "${YELLOW}[!]${NC} $1"; }
 log_error() { echo -e "${RED}[✗]${NC} $1"; }
 log_progress() { echo -e "${CYAN}[→]${NC} $1"; }
 
-# -----------------------------
-# Configuração
-# -----------------------------
+# --------------------------------
+# Config
+# --------------------------------
 COMFY_DIR="/root/comfy/ComfyUI"
 MODELS_DIR="$COMFY_DIR/models"
 COMFY_HOST="${COMFY_HOST:-0.0.0.0}"
 COMFY_PORT="${COMFY_PORT:-8818}"
+
+# Mantemos as keys hardcoded como pedido
 CIVITAI_TOKEN="${CIVITAI_TOKEN:-4fcb2834969399006a736ee402b061e5}"
 HF_TOKEN="${HF_TOKEN:-}"
 
-# Diretório para salvar estado
 STATE_DIR="/root/.comfyui_setup_state"
 STATE_FILE="$STATE_DIR/installation_state.txt"
 COMPLETED_DOWNLOADS="$STATE_DIR/completed_downloads.txt"
 COMPLETED_NODES="$STATE_DIR/completed_nodes.txt"
 
-# Configurações de performance
+# Performance
 export MAX_JOBS=32
 export NVCC_APPEND_FLAGS="--threads 8"
 export UV_SYSTEM_PYTHON=1
@@ -51,678 +44,446 @@ export PYTORCH_CUDA_ALLOC_CONF="expandable_segments:True"
 export HF_HUB_ENABLE_HF_TRANSFER=1
 export HF_TRANSFER_CONCURRENCY=16
 export EXT_PARALLEL=4
+export PIP_DEFAULT_TIMEOUT=60
 
-# -----------------------------
-# Sistema de Estado Persistente
-# -----------------------------
-init_state() {
-    mkdir -p "$STATE_DIR"
-    touch "$STATE_FILE" "$COMPLETED_DOWNLOADS" "$COMPLETED_NODES"
-}
+# Rede (timeouts e headers)
+DOWNLOAD_TIMEOUT="${DOWNLOAD_TIMEOUT:-900}"          # 15min hard cap por artefato
+CONNECT_TIMEOUT="${CONNECT_TIMEOUT:-20}"             # conexão
+RETRY_COUNT="${RETRY_COUNT:-3}"
+RETRY_WAIT="${RETRY_WAIT:-3}"
+UA="${UA:-Mozilla/5.0 (X11; Linux x86_64) curl/8}"
+CIVITAI_REF="https://civitai.com/"
 
-save_state() {
-    local step="$1"
-    echo "$step" > "$STATE_FILE"
-    log_info "Estado salvo: $step"
-}
+# --------------------------------
+# Estado
+# --------------------------------
+init_state(){ mkdir -p "$STATE_DIR"; :> "$STATE_FILE"; touch "$COMPLETED_DOWNLOADS" "$COMPLETED_NODES"; }
+save_state(){ echo "$1" > "$STATE_FILE"; log_info "Estado salvo: $1"; }
+get_state(){ [[ -f "$STATE_FILE" && "$(cat "$STATE_FILE" || echo 0)" =~ ^[0-9]+$ ]] && cat "$STATE_FILE" || echo 0; }
+reset_state(){ log_warn "Resetando estado da instalação..."; rm -rf "$STATE_DIR"; init_state; }
 
-get_state() {
-    if [ -f "$STATE_FILE" ]; then
-        local state=$(cat "$STATE_FILE" 2>/dev/null)
-        # Verificar se é um número válido
-        if [[ "$state" =~ ^[0-9]+$ ]]; then
-            echo "$state"
-        else
-            echo "0"
-        fi
-    else
-        echo "0"
-    fi
+# normaliza "url|type|filename" (trim e remove espaços fantasma)
+normalize_entry(){
+  local url="$1"; local type="$2"; local filename="$3"
+  url="$(echo -n "$url" | xargs)"; type="$(echo -n "$type" | xargs)"; filename="$(echo -n "${filename:-}" | xargs)"
+  echo "${url}|${type}|${filename}"
 }
+mark_download_complete(){ local key="$1"; echo "$key" >> "$COMPLETED_DOWNLOADS"; }
+is_download_complete(){ local key="$1"; [[ -f "$COMPLETED_DOWNLOADS" ]] && grep -Fxq "$key" "$COMPLETED_DOWNLOADS"; }
+mark_node_complete(){ echo "$1" >> "$COMPLETED_NODES"; }
+is_node_complete(){ [[ -f "$COMPLETED_NODES" ]] && grep -Fxq "$1" "$COMPLETED_NODES"; }
 
-mark_download_complete() {
-    local url="$1"
-    echo "$url" >> "$COMPLETED_DOWNLOADS"
-}
+# --------------------------------
+# Downloads (lista)
+# --------------------------------
+readonly DOWNLOAD_FILES=(
+  "mega://https://mega.nz/file/gIRTFQSQ#no6Ay3JLE9LVRi7ib9O-Jc0CW7XmG046kCgpCzDg1tY|loras|"
 
-is_download_complete() {
-    local url="$1"
-    grep -Fxq "$url" "$COMPLETED_DOWNLOADS" 2>/dev/null
-}
+  "https://civitai.com/api/download/models/2122278?type=Model&format=SafeTensor&size=pruned&fp=fp16|checkpoints|raehoshiIllustXL_v60.safetensors"
 
-mark_node_complete() {
-    local node="$1"
-    echo "$node" >> "$COMPLETED_NODES"
-}
-
-is_node_complete() {
-    local node="$1"
-    grep -Fxq "$node" "$COMPLETED_NODES" 2>/dev/null
-}
+  "https://civitai.com/api/download/models/1268294?type=Model&format=SafeTensor|loras|"
+  "https://civitai.com/api/download/models/1715330?type=Model&format=SafeTensor|loras|"
+  "https://civitai.com/api/download/models/1499397?type=Model&format=SafeTensor|loras|"
+  "https://civitai.com/api/download/models/1779002?type=Model&format=SafeTensor|loras|"
+  "https://civitai.com/api/download/models/1114313?type=Model&format=SafeTensor|loras|"
+  "https://civitai.com/api/download/models/1780244?type=Model&format=SafeTensor|loras|"
 
-reset_state() {
-    log_warn "Resetando estado da instalação..."
-    rm -rf "$STATE_DIR"
-    init_state
-}
+  "https://huggingface.co/Kim2091/AnimeSharp/resolve/main/4x-AnimeSharp.pth|upscale_models|4x-AnimeSharp.pth"
+  "https://huggingface.co/Kim2091/AnimeSharpV3/resolve/main/2x-AnimeSharpV3.pth|upscale_models|2x-AnimeSharpV3.pth"
+  "https://huggingface.co/ABDALLALSWAITI/Upscalers/resolve/main/anime/2x-AnimeSharpV2_MoSR_Soft.pth|upscale_models|2x-AnimeSharpV2_MoSR_Soft.pth"
+  "https://huggingface.co/Kim2091/UltraSharpV2/resolve/main/4x-UltraSharpV2.pth|upscale_models|4x-UltraSharpV2.pth"
+  "https://huggingface.co/Kim2091/UltraSharpV2/resolve/main/4x-UltraSharpV2_Lite.pth|upscale_models|4x-UltraSharpV2_Lite.pth"
+  "https://huggingface.co/Kim2091/UltraSharpV2/resolve/main/4x-UltraSharpV2_fp32_op17.onnx|upscale_models|4x-UltraSharpV2_fp32_op17.onnx"
 
-# -----------------------------
-# Lista de downloads (ATUALIZADA - sem ControlNet)
-# -----------------------------
-readonly DOWNLOAD_FILES=(
-    # LoRA from Mega (PRIMEIRO para falhar rápido se houver problema)
-    "mega://https://mega.nz/file/gIRTFQSQ#no6Ay3JLE9LVRi7ib9O-Jc0CW7XmG046kCgpCzDg1tY|loras|"
-    
-    # Checkpoint mantido
-    "https://civitai.com/api/download/models/2122278?type=Model&format=SafeTensor&size=pruned&fp=fp16|checkpoints|raehoshiIllustXL_v60.safetensors"
-    
-    # LoRAs from Civitai
-    "https://civitai.com/api/download/models/1268294?type=Model&format=SafeTensor|loras|"
-    "https://civitai.com/api/download/models/1715330?type=Model&format=SafeTensor|loras|"
-    "https://civitai.com/api/download/models/1499397?type=Model&format=SafeTensor|loras|"
-    "https://civitai.com/api/download/models/1779002?type=Model&format=SafeTensor|loras|"
-    "https://civitai.com/api/download/models/1114313?type=Model&format=SafeTensor|loras|"
-    "https://civitai.com/api/download/models/1780244?type=Model&format=SafeTensor|loras|"
-    
-    # Upscale models - Coleção completa
-    "https://huggingface.co/Kim2091/AnimeSharp/resolve/main/4x-AnimeSharp.pth|upscale_models|4x-AnimeSharp.pth"
-    "https://huggingface.co/Kim2091/AnimeSharpV3/resolve/main/2x-AnimeSharpV3.pth|upscale_models|2x-AnimeSharpV3.pth"
-    "https://huggingface.co/ABDALLALSWAITI/Upscalers/resolve/main/anime/2x-AnimeSharpV2_MoSR_Soft.pth|upscale_models|2x-AnimeSharpV2_MoSR_Soft.pth"
-    "https://huggingface.co/Kim2091/UltraSharpV2/resolve/main/4x-UltraSharpV2.pth|upscale_models|4x-UltraSharpV2.pth"
-    "https://huggingface.co/Kim2091/UltraSharpV2/resolve/main/4x-UltraSharpV2_Lite.pth|upscale_models|4x-UltraSharpV2_Lite.pth"
-    "https://huggingface.co/Kim2091/UltraSharpV2/resolve/main/4x-UltraSharpV2_fp32_op17.onnx|upscale_models|4x-UltraSharpV2_fp32_op17.onnx"
-    
-    # Ultralytics model
-    "https://huggingface.co/adbrasi/testedownload/resolve/main/99coins_anime_girl_face_m_seg.pt|ultralytics/bbox|99coins_anime_girl_face_m_seg.pt"
+  "https://huggingface.co/adbrasi/testedownload/resolve/main/99coins_anime_girl_face_m_seg.pt|ultralytics/bbox|99coins_anime_girl_face_m_seg.pt"
 )
 
-# Custom nodes para instalar (ATUALIZADA com novos nodes)
+# Custom nodes
 readonly CUSTOM_NODES=(
-    "https://github.com/adbrasi/huggpackreator"
-    "https://github.com/adbrasi/packreator_processor"
-    "https://github.com/adbrasi/Packreator_manager"
-    "https://github.com/adbrasi/cezarsave34"
-    "https://github.com/adbrasi/pageonetor"
-    "https://github.com/adbrasi/pakreatorio"
-    "https://github.com/adbrasi/WaterMark_bumbumzin"
-    "https://github.com/adbrasi/marcadaguita"
-    "https://github.com/adbrasi/randomico"
-    "https://github.com/adbrasi/groqrouter"
-    "https://github.com/adbrasi/find_charakito"
-    "https://github.com/adbrasi/randomsizito"
-    "https://github.com/adbrasi/importex"
-    "https://github.com/adbrasi/storitadifusita"
-    "https://github.com/adbrasi/attentionPPM"
-    "https://github.com/ClownsharkBatwing/RES4LYF"
-    "https://github.com/Suzie1/ComfyUI_Comfyroll_CustomNodes"
-    "https://github.com/sipherxyz/comfyui-art-venture"
-    "https://github.com/pamparamm/sd-perturbed-attention"
-    "https://github.com/KoreTeknology/ComfyUI-Universal-Styler"
-    "https://github.com/WASasquatch/was-node-suite-comfyui"
-    "https://github.com/chflame163/ComfyUI_LayerStyle"
-    "https://github.com/ltdrdata/ComfyUI-Impact-Pack"
-    "https://github.com/pythongosssss/ComfyUI-Custom-Scripts"
-    "https://github.com/rgthree/rgthree-comfy"
-    "https://github.com/ssitu/ComfyUI_UltimateSDUpscale"
-    "https://github.com/adbrasi/Importador"
-    "https://github.com/adbrasi/GetFirstTag"
-    "https://github.com/adbrasi/comfydodi"
-    "https://github.com/omar92/ComfyUI-QualityOfLifeSuit_Omar92"
-    "https://github.com/Cezarsaint/blacklisto"
-    "https://github.com/TinyTerra/ComfyUI_tinyterraNodes"
-    "https://github.com/ltdrdata/ComfyUI-Impact-Subpack"
-    "https://github.com/Cezarsaint/rand0micoUploaderLoven"
-    "https://github.com/adbrasi/pixivmosaic"
-    "https://github.com/adbrasi/futfilter"
-    "https://github.com/Artificial-Sweetener/comfyui-WhiteRabbit"
-    "https://github.com/shiimizu/ComfyUI_smZNodes"
-    "https://github.com/CoreyCorza/ComfyUI-CRZnodes"
+  "https://github.com/adbrasi/huggpackreator"
+  "https://github.com/adbrasi/packreator_processor"
+  "https://github.com/adbrasi/Packreator_manager"
+  "https://github.com/adbrasi/cezarsave34"
+  "https://github.com/adbrasi/pageonetor"
+  "https://github.com/adbrasi/pakreatorio"
+  "https://github.com/adbrasi/WaterMark_bumbumzin"
+  "https://github.com/adbrasi/marcadaguita"
+  "https://github.com/adbrasi/randomico"
+  "https://github.com/adbrasi/groqrouter"
+  "https://github.com/adbrasi/find_charakito"
+  "https://github.com/adbrasi/randomsizito"
+  "https://github.com/adbrasi/importex"
+  "https://github.com/adbrasi/storitadifusita"
+  "https://github.com/adbrasi/attentionPPM"
+  "https://github.com/ClownsharkBatwing/RES4LYF"
+  "https://github.com/Suzie1/ComfyUI_Comfyroll_CustomNodes"
+  "https://github.com/sipherxyz/comfyui-art-venture"
+  "https://github.com/pamparamm/sd-perturbed-attention"
+  "https://github.com/KoreTeknology/ComfyUI-Universal-Styler"
+  "https://github.com/WASasquatch/was-node-suite-comfyui"
+  "https://github.com/chflame163/ComfyUI_LayerStyle"
+  "https://github.com/ltdrdata/ComfyUI-Impact-Pack"
+  "https://github.com/pythongosssss/ComfyUI-Custom-Scripts"
+  "https://github.com/rgthree/rgthree-comfy"
+  "https://github.com/ssitu/ComfyUI_UltimateSDUpscale"
+  "https://github.com/adbrasi/Importador"
+  "https://github.com/adbrasi/GetFirstTag"
+  "https://github.com/adbrasi/comfydodi"
+  "https://github.com/omar92/ComfyUI-QualityOfLifeSuit_Omar92"
+  "https://github.com/Cezarsaint/blacklisto"
+  "https://github.com/TinyTerra/ComfyUI_tinyterraNodes"
+  "https://github.com/ltdrdata/ComfyUI-Impact-Subpack"
+  "https://github.com/Cezarsaint/rand0micoUploaderLoven"
+  "https://github.com/adbrasi/pixivmosaic"
+  "https://github.com/adbrasi/futfilter"
+  "https://github.com/Artificial-Sweetener/comfyui-WhiteRabbit"
+  "https://github.com/shiimizu/ComfyUI_smZNodes"
+  "https://github.com/CoreyCorza/ComfyUI-CRZnodes"
 )
 
-# Custom nodes que precisam de instalação especial via comfy-cli
-readonly SPECIAL_NODES=(
-    "Civicomfy"
-    "ComfyUI-RMBG"
-)
+readonly SPECIAL_NODES=("Civicomfy" "ComfyUI-RMBG")
 
-# -----------------------------
-# Funções auxiliares
-# -----------------------------
-command_exists() {
-    command -v "$1" >/dev/null 2>&1
+# --------------------------------
+# Helpers
+# --------------------------------
+command_exists(){ command -v "$1" >/dev/null 2>&1; }
+
+install_megatools(){
+  if ! command_exists megadl; then
+    log_info "Instalando megatools..."
+    apt-get update -qq && apt-get install -y -qq megatools >/dev/null 2>&1 || true
+    command_exists megadl && log_success "Megatools OK" || log_warn "Megatools falhou (downloads do Mega podem não funcionar)"
+  fi
 }
 
-# Instalar megatools se necessário
-install_megatools() {
-    if ! command_exists megadl; then
-        log_info "Instalando megatools para downloads do Mega..."
-        apt-get update -qq && apt-get install -y -qq megatools 2>/dev/null
-        if command_exists megadl; then
-            log_success "Megatools instalado com sucesso"
-        else
-            log_warn "Falha ao instalar megatools, downloads do Mega não estarão disponíveis"
-        fi
-    fi
+add_civitai_token(){
+  local url="$1"
+  if [[ "$url" == *"civitai.com/api/download"* ]] && [[ "$url" != *"token="* ]]; then
+    echo "${url}&token=${CIVITAI_TOKEN}"
+  else
+    echo "$url"
+  fi
 }
 
-# Download do Mega
-download_mega() {
-    local url="$1"
-    local target_dir="$2"
-    local filename="$3"
-    
-    if ! command_exists megadl; then
-        log_error "megatools não instalado"
-        return 1
-    fi
-    
-    # Remover prefixo mega:// se houver
-    url="${url#mega://}"
-    
-    log_info "Executando: megadl '$url'"
-    
-    # Usar megadl com timeout
-    cd "$target_dir"
-    if timeout 300 megadl "$url"; then
-        log_success "Download do Mega concluído"
-        cd - >/dev/null
-        return 0
-    else
-        log_error "Mega falhou ou timeout (5 min)"
-        cd - >/dev/null
-        return 1
-    fi
+# HuggingFace (CLI) com timeout
+download_hf(){
+  local repo="$1" file_path="$2" target_dir="$3" filename="$4"
+  [[ -z "$filename" ]] && filename="$(basename "$file_path")"
+  local target_file="$target_dir/$filename"
+  if [[ -s "$target_file" && $(stat -c%s "$target_file" 2>/dev/null || echo 0) -gt 1000000 ]]; then
+    log_success "Arquivo já existe: $filename"; return 0
+  fi
+
+  log_info "Baixando de HF: $repo/$file_path -> $filename"
+  local temp_dir; temp_dir="$(mktemp -d)"
+  local ok=false
+
+  if command_exists hf; then
+    if timeout "$DOWNLOAD_TIMEOUT" env HF_HUB_ENABLE_HF_TRANSFER=1 hf download "$repo" "$file_path" \
+      --local-dir "$temp_dir" --local-dir-use-symlinks False >/dev/null 2>&1; then ok=true; fi
+  fi
+
+  if [[ "$ok" = false && $(command_exists huggingface-cli; echo $?) -eq 0 ]]; then
+    log_warn "Tentando com huggingface-cli..."
+    if timeout "$DOWNLOAD_TIMEOUT" env HF_HUB_ENABLE_HF_TRANSFER=1 huggingface-cli download "$repo" "$file_path" \
+      --local-dir "$temp_dir" --local-dir-use-symlinks False >/dev/null 2>&1; then ok=true; fi
+  fi
+
+  if "$ok"; then
+    local f; f="$(find "$temp_dir" -type f \( -name "*.safetensors" -o -name "*.pth" -o -name "*.pt" -o -name "*.bin" -o -name "*.onnx" \) | head -1)"
+    if [[ -n "$f" ]]; then mv "$f" "$target_file"; rm -rf "$temp_dir"; log_success "HF OK: $filename"; return 0; fi
+  fi
+
+  rm -rf "$temp_dir"; log_error "HF falhou: $filename"; return 1
 }
 
-# Adicionar token do Civitai se necessário
-add_civitai_token() {
-    local url="$1"
-    if [[ "$url" == *"civitai.com/api/download"* ]] && [[ "$url" != *"token="* ]]; then
-        echo "${url}&token=${CIVITAI_TOKEN}"
-    else
-        echo "$url"
-    fi
+# MEGA com timeout
+download_mega(){
+  local url="$1" target_dir="$2"
+  command_exists megadl || { log_error "megatools não instalado"; return 1; }
+  url="${url#mega://}"
+  log_info "MEGA: $url"
+  ( cd "$target_dir" && timeout "$DOWNLOAD_TIMEOUT" megadl -q "$url" ) && { log_success "MEGA OK"; return 0; }
+  log_error "MEGA falhou/timeout"; return 1
 }
 
-# Download HuggingFace usando hf com hf_transfer
-download_hf() {
-    local repo="$1"
-    local file_path="$2"
-    local target_dir="$3"
-    local filename="$4"
-    
-    if [ -z "$filename" ]; then
-        filename=$(basename "$file_path")
-    fi
-    
-    local target_file="$target_dir/$filename"
-    
-    if [ -f "$target_file" ] && [ $(stat -c%s "$target_file" 2>/dev/null || echo 0) -gt 1000000 ]; then
-        log_success "Arquivo já existe: $filename"
-        return 0
-    fi
-    
-    log_info "Baixando de HF: $filename"
-    
-    local temp_dir=$(mktemp -d)
-    local download_success=false
-    
-    if command_exists hf; then
-        if HF_HUB_ENABLE_HF_TRANSFER=1 hf download "$repo" "$file_path" \
-            --local-dir "$temp_dir" \
-            --local-dir-use-symlinks False 2>/dev/null; then
-            download_success=true
-        fi
-    fi
-    
-    if [ "$download_success" = false ] && command_exists huggingface-cli; then
-        log_warn "Tentando com huggingface-cli..."
-        if HF_HUB_ENABLE_HF_TRANSFER=1 huggingface-cli download "$repo" "$file_path" \
-            --local-dir "$temp_dir" \
-            --local-dir-use-symlinks False 2>/dev/null; then
-            download_success=true
-        fi
-    fi
-    
-    if [ "$download_success" = true ]; then
-        local downloaded_file=$(find "$temp_dir" -type f \( -name "*.safetensors" -o -name "*.pth" -o -name "*.pt" -o -name "*.bin" -o -name "*.onnx" \) 2>/dev/null | head -1)
-        
-        if [ -n "$downloaded_file" ]; then
-            mv "$downloaded_file" "$target_file"
-            rm -rf "$temp_dir"
-            log_success "Download concluído: $filename"
-            return 0
-        fi
-    fi
-    
-    rm -rf "$temp_dir"
-    log_error "Falha ao baixar: $filename"
-    return 1
+# Detecta nome via cabeçalho
+detect_remote_name(){
+  local url="$1"
+  local name
+  name="$(timeout 60 curl -I -sSL -4 -A "$UA" -e "$CIVITAI_REF" "$url" \
+          | awk -F'filename=' 'BEGIN{IGNORECASE=1}/content-disposition/{gsub(/\"|;|\r/,"",$2); print $2}' | tail -1)"
+  [[ -z "$name" ]] && name="$(basename "$(echo "$url" | cut -d'?' -f1)")"
+  echo "$name"
 }
 
-# Download genérico com aria2c, wget ou curl
-download_file() {
-    local url="$1"
-    local target_dir="$2"
-    local filename="$3"
-    
-    url=$(add_civitai_token "$url")
-    
-    if [ -n "$filename" ] && [ -f "$target_dir/$filename" ] && [ $(stat -c%s "$target_dir/$filename" 2>/dev/null || echo 0) -gt 1000000 ]; then
-        log_success "Arquivo já existe: $filename"
-        return 0
-    fi
-    
-    log_info "Baixando: ${filename:-$(basename "$url" | cut -d'?' -f1)}"
-    
-    if command_exists aria2c; then
-        local aria_opts="-c -s 16 -x 16 -k 1M --console-log-level=warn --summary-interval=10"
-        if [ -n "$filename" ]; then
-            aria2c $aria_opts --dir="$target_dir" --out="$filename" "$url" && return 0
-        else
-            aria2c $aria_opts --dir="$target_dir" "$url" && return 0
-        fi
-    fi
-    
-    if command_exists wget; then
-        if [ -n "$filename" ]; then
-            wget -q --show-progress -c -O "$target_dir/$filename" "$url" && return 0
-        else
-            wget -q --show-progress -c --content-disposition -P "$target_dir" "$url" && return 0
-        fi
+# Baixador genérico (CURL -> WGET -> ARIA2C), todos com timeout & retries
+download_file(){
+  local url="$1" target_dir="$2" filename="$3"
+  url="$(add_civitai_token "$url")"
+
+  if [[ -z "$filename" ]]; then filename="$(detect_remote_name "$url")"; fi
+  local out="$target_dir/$filename"
+
+  if [[ -s "$out" && $(stat -c%s "$out" 2>/dev/null || echo 0) -gt 1000000 ]]; then
+    log_success "Arquivo já existe: $filename"; return 0
+  fi
+
+  log_info "HTTP: ${filename} (-> $target_dir)"
+
+  # 1) CURL
+  if command_exists curl; then
+    if timeout "$DOWNLOAD_TIMEOUT" curl -fL -4 -A "$UA" -e "$CIVITAI_REF" \
+      --connect-timeout "$CONNECT_TIMEOUT" --retry "$RETRY_COUNT" --retry-delay "$RETRY_WAIT" \
+      -C - -o "$out" "$url"; then
+      log_success "CURL OK: $filename"; return 0
+    else
+      log_warn "CURL falhou para $filename"
     fi
-    
-    if command_exists curl; then
-        if [ -n "$filename" ]; then
-            curl -L -# -C - -o "$target_dir/$filename" "$url" && return 0
-        else
-            local headers=$(curl -sI -L "$url")
-            local detected_name=$(echo "$headers" | grep -i "content-disposition" | sed -n 's/.*filename="\([^"]*\)".*/\1/p' | tr -d '\r')
-            if [ -z "$detected_name" ]; then
-                detected_name="downloaded_$(date +%s).safetensors"
-            fi
-            curl -L -# -C - -o "$target_dir/$detected_name" "$url" && return 0
-        fi
+  fi
+
+  # 2) WGET
+  if command_exists wget; then
+    if timeout "$DOWNLOAD_TIMEOUT" wget -q --show-progress -4 --referer="$CIVITAI_REF" --user-agent="$UA" \
+      --timeout="$CONNECT_TIMEOUT" --tries="$RETRY_COUNT" -c -O "$out" "$url"; then
+      log_success "WGET OK: $filename"; return 0
+    else
+      log_warn "WGET falhou para $filename"
     fi
-    
-    log_error "Falha ao baixar: $url"
-    return 1
-}
+  fi
 
-# Processar lista de downloads
-process_downloads() {
-    local total=${#DOWNLOAD_FILES[@]}
-    local current=0
-    local failed_count=0
-    
-    log_info "Total de downloads: $total"
-    log_info "Iniciando loop de downloads..."
-    
-    for entry in "${DOWNLOAD_FILES[@]}"; do
-        ((current++))
-        
-        log_info "[$current/$total] Processando entrada..."
-        
-        # Verificar se já foi baixado
-        if is_download_complete "$entry"; then
-            log_success "[$current/$total] Já baixado, pulando..."
-            continue
-        fi
-        
-        log_info "[$current/$total] Preparando download..."
-        
-        IFS='|' read -r url type filename <<< "$entry"
-        
-        # Limpar espaços
-        url=$(echo "$url" | xargs)
-        type=$(echo "$type" | xargs)
-        filename=$(echo "$filename" | xargs)
-        
-        log_info "[$current/$total] URL: ${url:0:60}..."
-        log_info "[$current/$total] Tipo: $type"
-        
-        # Determinar diretório de destino
-        local target_dir="$MODELS_DIR"
-        case "$type" in
-            checkpoints) target_dir="$MODELS_DIR/checkpoints" ;;
-            diffusion_models) target_dir="$MODELS_DIR/diffusion_models" ;;
-            loras) target_dir="$MODELS_DIR/loras" ;;
-            vae) target_dir="$MODELS_DIR/vae" ;;
-            vae_approx) target_dir="$MODELS_DIR/vae_approx" ;;
-            text_encoders) target_dir="$MODELS_DIR/text_encoders" ;;
-            clip_vision) target_dir="$MODELS_DIR/clip_vision" ;;
-            upscale_models) target_dir="$MODELS_DIR/upscale_models" ;;
-            ultralytics/bbox) target_dir="$MODELS_DIR/ultralytics/bbox" ;;
-            *) target_dir="$MODELS_DIR/$type" ;;
-        esac
-        
-        mkdir -p "$target_dir"
-        
-        # Executar download apropriado
-        local download_result=1
-        
-        if [[ "$url" == mega://* ]]; then
-            log_info "[$current/$total] Download do Mega..."
-            download_mega "$url" "$target_dir" "$filename" && download_result=0
-        elif [[ "$url" == hf://* ]]; then
-            log_info "[$current/$total] Download HuggingFace..."
-            url="${url#hf://}"
-            local repo=$(echo "$url" | cut -d'/' -f1-2)
-            local file_path=$(echo "$url" | cut -d'/' -f3-)
-            download_hf "$repo" "$file_path" "$target_dir" "$filename" && download_result=0
-        else
-            log_info "[$current/$total] Download HTTP..."
-            download_file "$url" "$target_dir" "$filename" && download_result=0
-        fi
-        
-        if [ $download_result -eq 0 ]; then
-            mark_download_complete "$entry"
-        else
-            ((failed_count++))
-            log_warn "Falhou, continuando..."
-        fi
-    done
-    
-    if [ $failed_count -gt 0 ]; then
-        log_warn "Downloads completos: $((total - failed_count))/$total (${failed_count} falharam)"
+  # 3) ARIA2C
+  if command_exists aria2c; then
+    local aria_opts="-c -s 16 -x 16 -k 1M --file-allocation=none --allow-overwrite=true --auto-file-renaming=false \
+      --summary-interval=10 --console-log-level=warn --timeout=$CONNECT_TIMEOUT --retry-wait=$RETRY_WAIT --max-tries=$RETRY_COUNT --referer=$CIVITAI_REF --user-agent=$UA"
+    if timeout "$DOWNLOAD_TIMEOUT" aria2c $aria_opts --dir="$target_dir" --out="$filename" "$url"; then
+      log_success "ARIA2C OK: $filename"; return 0
     else
-        log_success "Todos os $total downloads completados!"
+      log_warn "ARIA2C falhou para $filename"
     fi
-    
-    return 0
+  fi
+
+  log_error "Falha final em: $url"
+  return 1
 }
 
-# Clonar ou atualizar repositório git
-clone_or_update() {
-    local url="$1"
-    local dest="$2"
-    local node_name=$(basename "$dest")
-    
-    if is_node_complete "$url"; then
-        log_success "Node já instalado: $node_name"
-        return 0
+# Processa a lista de downloads com chave normalizada + logs
+process_downloads(){
+  local total="${#DOWNLOAD_FILES[@]}"; local current=0; local failed=0
+  log_info "Total de downloads: $total"
+  log_info "Iniciando loop de downloads..."
+
+  for raw in "${DOWNLOAD_FILES[@]}"; do
+    ((current++))
+    IFS='|' read -r url type filename <<< "$raw"
+    url="$(echo -n "$url" | xargs)"; type="$(echo -n "$type" | xargs)"; filename="$(echo -n "${filename:-}" | xargs)"
+    local key; key="$(normalize_entry "$url" "$type" "$filename")"
+
+    if is_download_complete "$key"; then
+      log_success "[$current/$total] Já baixado: ${filename:-$(basename "$url" | cut -d'?' -f1)}"; continue
     fi
-    
-    if [ -d "$dest/.git" ]; then
-        log_info "Atualizando $node_name..."
-        timeout 60 git -C "$dest" pull --ff-only 2>/dev/null || {
-            log_warn "Timeout ou erro ao atualizar $node_name"
-            return 0
-        }
+
+    local target_dir="$MODELS_DIR"
+    case "$type" in
+      checkpoints) target_dir="$MODELS_DIR/checkpoints" ;;
+      loras) target_dir="$MODELS_DIR/loras" ;;
+      upscale_models) target_dir="$MODELS_DIR/upscale_models" ;;
+      ultralytics/bbox) target_dir="$MODELS_DIR/ultralytics/bbox" ;;
+      *) target_dir="$MODELS_DIR/$type" ;;
+    esac
+    mkdir -p "$target_dir"
+
+    log_info "[$current/$total] $type :: $(basename "${filename:-$url}")"
+
+    local ok=1
+    if [[ "$url" == mega://* ]]; then
+      download_mega "$url" "$target_dir" && ok=0
+    elif [[ "$url" == hf://* ]]; then
+      url="${url#hf://}"; local repo="${url%%/*/*}"; repo="${url%/*/*}"; # não usado aqui, mantido por compat
+      local repo2="$(echo "$url" | cut -d'/' -f1-2)"
+      local file_path="$(echo "$url" | cut -d'/' -f3-)"
+      download_hf "$repo2" "$file_path" "$target_dir" "$filename" && ok=0
+    elif [[ "$url" == https://huggingface.co/*/resolve/* ]]; then
+      # resolve/ direto por HTTP (mais rápido se público)
+      download_file "$url" "$target_dir" "$filename" && ok=0
     else
-        log_info "Clonando $node_name..."
-        timeout 60 git clone --recursive --depth 1 "$url" "$dest" 2>/dev/null || {
-            log_warn "Falha ao clonar $node_name"
-            return 0
-        }
+      download_file "$url" "$target_dir" "$filename" && ok=0
     fi
-    
-    if [ -f "$dest/requirements.txt" ]; then
-        timeout 120 python3 -m pip install --no-warn-script-location -q -r "$dest/requirements.txt" 2>/dev/null || {
-            log_warn "Falha ao instalar requirements para $node_name"
-        }
+
+    if [[ $ok -eq 0 ]]; then
+      mark_download_complete "$key"
+    else
+      ((failed++)); log_warn "[$current/$total] Falhou, segue..."
     fi
-    
-    mark_node_complete "$url"
-    return 0
+  done
+
+  if [[ $failed -gt 0 ]]; then
+    log_warn "Downloads: $((total - failed))/$total OK ($failed falharam)"
+  else
+    log_success "Todos os $total downloads completados!"
+  fi
+  return 0
 }
 
+# Git clone/update com pip reqs e timeouts
+clone_or_update(){
+  local url="$1" dest="$2" node_name="$(basename "$dest")"
 
-# -----------------------------
-# Menu de opções
-# -----------------------------
-show_menu() {
-    echo ""
-    log_info "========================================="
-    log_info " ComfyUI Custom Setup v2.0"
-    log_info "========================================="
-    echo ""
-    
-    local current_state=$(get_state)
-    
-    # Verificar se o estado está vazio ou inválido
-    if [ -z "$current_state" ] || ! [[ "$current_state" =~ ^[0-9]+$ ]]; then
-        log_info "Iniciando nova instalação..."
-        reset_state
-        return 0
-    fi
-    
-    # Se não for interativo, sempre continuar ou começar do zero
-    if [ "$INTERACTIVE" = false ]; then
-        if [ "$current_state" != "0" ] && [ "$current_state" -lt "8" ]; then
-            log_info "Modo não-interativo: continuando instalação anterior (Step: $current_state)"
-        else
-            log_info "Modo não-interativo: iniciando nova instalação"
-        fi
-        return 0
-    fi
-    
-    if [ "$current_state" != "0" ] && [ "$current_state" -lt "8" ]; then
-        log_warn "Instalação anterior detectada (Step: $current_state)"
-        echo ""
-        echo "Opções:"
-        echo "  1) Continuar instalação de onde parou"
-        echo "  2) Reiniciar instalação do zero"
-        echo "  3) Sair"
-        echo ""
-        
-        # Timeout de 10 segundos, padrão é continuar
-        read -t 10 -p "Escolha uma opção [1-3] (padrão: 1 em 10s): " choice || choice="1"
-        
-        case $choice in
-            1|"") return 0 ;;
-            2) reset_state; return 0 ;;
-            3) exit 0 ;;
-            *) log_warn "Opção inválida, continuando..."; return 0 ;;
-        esac
-    fi
-    
+  if is_node_complete "$url"; then log_success "Node já instalado: $node_name"; return 0; fi
+
+  if [[ -d "$dest/.git" ]]; then
+    log_info "Atualizando $node_name..."
+    timeout 120 git -C "$dest" pull --ff-only >/dev/null 2>&1 || log_warn "Timeout/erro ao atualizar $node_name"
+  else
+    log_info "Clonando $node_name..."
+    timeout 180 git clone --recursive --depth 1 "$url" "$dest" >/dev/null 2>&1 || { log_warn "Falha ao clonar $node_name"; return 0; }
+  fi
+
+  if [[ -f "$dest/requirements.txt" ]]; then
+    timeout 300 python3 -m pip install --no-warn-script-location -q -r "$dest/requirements.txt" >/dev/null 2>&1 || \
+      log_warn "Falha em requirements de $node_name"
+  fi
+
+  mark_node_complete "$url"
+  return 0
+}
+
+# --------------------------------
+# Menu
+# --------------------------------
+show_menu(){
+  echo ""; log_info "========================================="; log_info " ComfyUI Custom Setup v2.0"; log_info "========================================="; echo ""
+  local s; s="$(get_state)"
+
+  if [[ -z "$s" || ! "$s" =~ ^[0-9]+$ ]]; then log_info "Iniciando nova instalação..."; reset_state; return 0; fi
+
+  if [[ "$INTERACTIVE" = false ]]; then
+    if [[ "$s" != "0" && "$s" -lt "8" ]]; then log_info "Modo não-interativo: continuando (Step: $s)"; else log_info "Modo não-interativo: nova instalação"; fi
     return 0
+  fi
+
+  if [[ "$s" != "0" && "$s" -lt "8" ]]; then
+    log_warn "Instalação anterior detectada (Step: $s)"
+    echo -e "\nOpções:\n  1) Continuar\n  2) Reiniciar\n  3) Sair\n"
+    read -t 10 -p "Escolha [1-3] (padrão 1): " choice || choice="1"
+    case "$choice" in
+      1|"") return 0 ;;
+      2) reset_state; return 0 ;;
+      3) exit 0 ;;
+      *) log_warn "Opção inválida, continuando..."; return 0 ;;
+    esac
+  fi
 }
 
-# -----------------------------
-# Instalação principal
-# -----------------------------
-main() {
-    # Inicializar sistema de estado
-    init_state
-    
-    # Mostrar menu
-    show_menu
-    
-    local current_state=$(get_state)
-    
-    # [STEP 1] Verificar e instalar dependências
-    if [ "$current_state" -lt "1" ]; then
-        log_info "[1/6] Verificando e instalando dependências..."
-        
-        # Instalar aria2c e megatools
-        if ! command_exists aria2c; then
-            log_info "Instalando aria2c..."
-            apt-get update -qq && apt-get install -y -qq aria2 2>/dev/null
-        fi
-        
-        install_megatools
-        
-        for cmd in python3 git wget curl; do
-            if ! command_exists "$cmd"; then
-                log_error "Dependência faltando: $cmd"
-                exit 1
-            fi
-        done
-        
-        log_success "Dependências OK"
-        save_state "1"
-    fi
-    
-    # [STEP 2] Atualizar pip e instalar ferramentas
-    if [ "$current_state" -lt "2" ]; then
-        log_info "[2/6] Preparando ferramentas Python..."
-        python3 -m pip install --upgrade pip wheel setuptools -q
-        python3 -m pip install --upgrade "huggingface_hub[cli,hf_transfer]>=0.26.0" comfy-cli -q
-        
-        if [ -n "$HF_TOKEN" ]; then
-            huggingface-cli login --token "$HF_TOKEN" 2>/dev/null || log_warn "Login HF falhou"
-        fi
-        
-        log_success "Ferramentas instaladas"
-        save_state "2"
-    fi
-    
-    # [STEP 3] Instalar ComfyUI
-    if [ "$current_state" -lt "3" ]; then
-        log_info "[3/6] Instalando ComfyUI..."
-        if [ -f "$COMFY_DIR/main.py" ]; then
-            log_warn "ComfyUI já existe"
-        else
-            comfy --skip-prompt install --fast-deps --nvidia
-            
-            if [ ! -f "$COMFY_DIR/main.py" ]; then
-                log_error "ComfyUI não foi instalado!"
-                exit 1
-            fi
-        fi
-        log_success "ComfyUI instalado"
-        save_state "3"
-    fi
-    
-    # [STEP 4] Corrigir PyTorch para Blackwell se detectado
-    if [ "$current_state" -lt "4" ]; then
-        log_info "[4/6] Verificando GPU..."
-        
-        # Detectar GPU
-        local gpu_info=$(nvidia-smi --query-gpu=name --format=csv,noheader 2>/dev/null || echo "")
-        
-        if [[ "$gpu_info" == *"5090"* ]] || [[ "$gpu_info" == *"5080"* ]]; then
-            log_warn "RTX ${gpu_info} detectada - instalando PyTorch nightly"
-            # FORÇAR reinstalação do PyTorch nightly
-            pip uninstall -y torch torchvision torchaudio 2>/dev/null || true
-            pip install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu128
-            log_success "PyTorch nightly instalado"
-        else
-            log_info "GPU: ${gpu_info:-Não detectada} - mantendo PyTorch do ComfyUI"
-        fi
-        save_state "4"
-    fi
-    
-    # [STEP 5] Baixar modelos
-    # [STEP 5] Baixar modelos
-    if [ "$current_state" -lt "5" ]; then
-        log_info "[5/6] Baixando modelos..."
-        mkdir -p "$MODELS_DIR"/{diffusion_models,loras,vae,vae_approx,text_encoders,clip_vision,upscale_models,checkpoints,ultralytics/bbox}
-        
-        log_info "Chamando process_downloads..."
-        process_downloads
-        log_info "process_downloads terminou"
-        
-        log_success "Downloads processados"
-        save_state "5"
-    fi
-    
-    # [STEP 6] Instalar custom nodes
-    if [ "$current_state" -lt "6" ]; then
-        log_info "[6/6] Instalando custom nodes..."
-        CN_DIR="$COMFY_DIR/custom_nodes"
-        mkdir -p "$CN_DIR"
-        
-        for repo_url in "${CUSTOM_NODES[@]}"; do
-            node_name=$(basename "$repo_url")
-            clone_or_update "$repo_url" "$CN_DIR/$node_name"
-        done
-        
-        for node_name in "${SPECIAL_NODES[@]}"; do
-            if ! is_node_complete "special:$node_name"; then
-                log_info "Instalando $node_name via comfy-cli..."
-                comfy node install "$node_name" 2>/dev/null && mark_node_complete "special:$node_name" || log_warn "Falha: $node_name"
-            fi
-        done
-        
-        log_success "Custom nodes instalados"
-        save_state "6"
-    fi
-    
-    # Verificação final
-    log_info "Verificando instalação..."
-    
-    echo ""
-    log_info "Estatísticas da instalação:"
-    
-    # Verificar ComfyUI
-    if [ -f "$COMFY_DIR/main.py" ]; then
-        log_success "✓ ComfyUI instalado"
+# --------------------------------
+# Main
+# --------------------------------
+main(){
+  mkdir -p "$COMFY_DIR" "$MODELS_DIR"; init_state
+  show_menu
+
+  local s; s="$(get_state)"
+
+  # STEP 1: deps
+  if [[ "$s" -lt "1" ]]; then
+    log_info "[1/6] Dependências..."
+    if ! command_exists aria2c; then apt-get update -qq && apt-get install -y -qq aria2 >/dev/null 2>&1 || true; fi
+    install_megatools
+    for cmd in python3 git wget curl; do command_exists "$cmd" || { log_error "Falta: $cmd"; exit 1; }; done
+    log_success "Dependências OK"; save_state "1"
+  fi
+
+  # STEP 2: pip tools
+  if [[ "$s" -lt "2" ]]; then
+    log_info "[2/6] Ferramentas Python..."
+    python3 -m pip install --upgrade --no-cache-dir pip wheel setuptools >/dev/null 2>&1
+    python3 -m pip install --upgrade --no-cache-dir "huggingface_hub[cli,hf_transfer]>=0.26.0" comfy-cli >/dev/null 2>&1
+    if [[ -n "$HF_TOKEN" ]]; then huggingface-cli login --token "$HF_TOKEN" >/dev/null 2>&1 || log_warn "Login HF falhou"; fi
+    log_success "Ferramentas OK"; save_state "2"
+  fi
+
+  # STEP 3: ComfyUI
+  if [[ "$s" -lt "3" ]]; then
+    log_info "[3/6] Instalando ComfyUI..."
+    if [[ -f "$COMFY_DIR/main.py" ]]; then log_warn "ComfyUI já existe"; else comfy --skip-prompt install --fast-deps --nvidia || { log_error "ComfyUI não instalou"; exit 1; }; fi
+    log_success "ComfyUI instalado"; save_state "3"
+  fi
+
+  # STEP 4: Blackwell nightly (se 5090/5080)
+  if [[ "$s" -lt "4" ]]; then
+    log_info "[4/6] Verificando GPU..."
+    local gpu_info; gpu_info="$(nvidia-smi --query-gpu=name --format=csv,noheader 2>/dev/null || echo "")"
+    if [[ "$gpu_info" == *"5090"* || "$gpu_info" == *"5080"* ]]; then
+      log_warn "RTX ${gpu_info:-Blackwell} detectada - aplicando PyTorch nightly (como você pediu)"
+      # alinhado ao pedido: --pre --upgrade --no-cache-dir + extra-index nightly
+      python3 -m pip install --pre --upgrade --no-cache-dir torch --extra-index-url https://download.pytorch.org/whl/nightly/cu128
+      python3 -m pip install --pre --upgrade --no-cache-dir torchvision --extra-index-url https://download.pytorch.org/whl/nightly/cu128
+      log_success "PyTorch nightly (cu128) pronto"
     else
-        log_error "✗ ComfyUI não encontrado"
+      log_info "GPU: ${gpu_info:-Não detectada} - mantendo PyTorch padrão"
     fi
-    
-    # Contar modelos
-    echo ""
-    log_info "Modelos instalados:"
-    for dir in checkpoints loras vae vae_approx text_encoders clip_vision upscale_models ultralytics; do
-        if [ -d "$MODELS_DIR/$dir" ]; then
-            local count=$(find "$MODELS_DIR/$dir" -type f \( -name "*.safetensors" -o -name "*.pth" -o -name "*.pt" -o -name "*.ckpt" -o -name "*.onnx" \) 2>/dev/null | wc -l)
-            if [ $count -gt 0 ]; then
-                log_success "  $dir: $count arquivo(s)"
-            fi
-        fi
+    save_state "4"
+  fi
+
+  # STEP 5: modelos
+  if [[ "$s" -lt "5" ]]; then
+    log_info "[5/6] Baixando modelos..."
+    mkdir -p "$MODELS_DIR"/{loras,upscale_models,checkpoints,ultralytics/bbox}
+    process_downloads
+    log_success "Downloads processados"; save_state "5"
+  fi
+
+  # STEP 6: custom nodes
+  if [[ "$s" -lt "6" ]]; then
+    log_info "[6/6] Instalando custom nodes..."
+    local CN_DIR="$COMFY_DIR/custom_nodes"; mkdir -p "$CN_DIR"
+    for repo_url in "${CUSTOM_NODES[@]}"; do clone_or_update "$repo_url" "$CN_DIR/$(basename "$repo_url")"; done
+    for node_name in "${SPECIAL_NODES[@]}"; do
+      if ! is_node_complete "special:$node_name"; then
+        log_info "Instalando $node_name via comfy-cli..."
+        comfy node install "$node_name" >/dev/null 2>&1 && mark_node_complete "special:$node_name" || log_warn "Falha: $node_name"
+      fi
     done
-    
-    # Contar nodes
-    echo ""
-    local node_count=$(find "$CN_DIR" -maxdepth 1 -type d 2>/dev/null | wc -l)
-    log_success "Custom nodes: $((node_count - 1)) instalados"
-    
-    # Verificar GPU
-    echo ""
-    local gpu_info=$(nvidia-smi --query-gpu=name --format=csv,noheader 2>/dev/null || echo "Não detectada")
-    log_info "GPU: $gpu_info"
-    
-    # Marcar instalação como completa
-    save_state "6"
-    
-    echo ""
-    log_success "========================================="
-    log_success " Instalação concluída!"
-    log_success "========================================="
-    
-    # Se não for interativo ou se for via argumento, iniciar direto
-    if [ "$INTERACTIVE" = false ] || [ "${1:-}" = "--start" ]; then
-        log_info "Iniciando ComfyUI automaticamente..."
-        log_info "URL: http://localhost:$COMFY_PORT"
-        cd "$COMFY_DIR"
-        exec comfy launch -- --listen "$COMFY_HOST" --port "$COMFY_PORT"
+    log_success "Custom nodes instalados"; save_state "6"
+  fi
+
+  # Verificação final
+  log_info "Verificando instalação..."
+  if [[ -f "$COMFY_DIR/main.py" ]]; then log_success "✓ ComfyUI instalado"; else log_error "✗ ComfyUI não encontrado"; fi
+
+  echo ""; log_info "Modelos instalados:"
+  for dir in checkpoints loras vae vae_approx text_encoders clip_vision upscale_models ultralytics; do
+    if [[ -d "$MODELS_DIR/$dir" ]]; then
+      local c; c="$(find "$MODELS_DIR/$dir" -type f \( -name "*.safetensors" -o -name "*.pth" -o -name "*.pt" -o -name "*.ckpt" -o -name "*.onnx" \) 2>/dev/null | wc -l)"
+      [[ "$c" -gt 0 ]] && log_success "  $dir: $c arquivo(s)"
     fi
-    
-    echo ""
-    echo "Opções:"
-    echo "  1) Iniciar ComfyUI agora"
-    echo "  2) Resetar estado da instalação"
-    echo "  3) Sair"
-    echo ""
-    
-    # Timeout de 10 segundos, padrão é iniciar
-    read -t 10 -p "Escolha uma opção [1-3] (padrão: 1 em 10s): " final_choice || final_choice="1"
-    
-    case $final_choice in
-        1|"")
-            log_info "Iniciando ComfyUI..."
-            log_info "URL: http://localhost:$COMFY_PORT"
-            cd "$COMFY_DIR"
-            exec comfy launch -- --listen "$COMFY_HOST" --port "$COMFY_PORT"
-            ;;
-        2)
-            reset_state
-            log_success "Estado resetado. Execute o script novamente para reinstalar."
-            ;;
-        3)
-            log_info "Saindo..."
-            ;;
-        *)
-            log_info "Iniciando ComfyUI (opção padrão)..."
-            cd "$COMFY_DIR"
-            exec comfy launch -- --listen "$COMFY_HOST" --port "$COMFY_PORT"
-            ;;
-    esac
+  done
+
+  echo ""; local CN_DIR="$COMFY_DIR/custom_nodes"; local node_count; node_count="$(find "$CN_DIR" -maxdepth 1 -type d 2>/dev/null | wc -l)"
+  log_success "Custom nodes: $((node_count - 1)) instalados"
+
+  echo ""; local gpu_info; gpu_info="$(nvidia-smi --query-gpu=name --format=csv,noheader 2>/dev/null || echo "Não detectada")"
+  log_info "GPU: $gpu_info"
+
+  save_state "6"
+
+  echo ""; log_success "========================================="; log_success " Instalação concluída!"; log_success "========================================="
+
+  if [[ "$INTERACTIVE" = false || "${1:-}" == "--start" ]]; then
+    log_info "Iniciando ComfyUI automaticamente..."
+    log_info "URL: http://localhost:$COMFY_PORT"
+    cd "$COMFY_DIR"
+    exec comfy launch -- --listen "$COMFY_HOST" --port "$COMFY_PORT"
+  fi
+
+  echo -e "\nOpções:\n  1) Iniciar ComfyUI agora\n  2) Resetar estado\n  3) Sair\n"
+  read -t 10 -p "Escolha [1-3] (padrão 1): " final_choice || final_choice="1"
+  case "$final_choice" in
+    1|"") log_info "Iniciando ComfyUI..."; cd "$COMFY_DIR"; exec comfy launch -- --listen "$COMFY_HOST" --port "$COMFY_PORT" ;;
+    2) reset_state; log_success "Estado resetado. Rode de novo para reinstalar." ;;
+    3) log_info "Saindo..." ;;
+    *) cd "$COMFY_DIR"; exec comfy launch -- --listen "$COMFY_HOST" --port "$COMFY_PORT" ;;
+  esac
 }
 
-# Executar instalação
-main "$@"
+main "$@"