Update image_Studio.sh

image_Studio.sh

@@ -1,6 +1,6 @@
 #!/usr/bin/env bash
-# setup_comfyui_custom.sh
-# Script modificado para instalação do ComfyUI com custom nodes e modelos personalizados
+# setup_comfyui_custom_v2.sh
+# Script otimizado com suporte para RTX 5090, estado persistente e downloads do Mega
 
 set -euo pipefail
 
@@ -11,12 +11,14 @@ 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
@@ -26,7 +28,13 @@ MODELS_DIR="$COMFY_DIR/models"
 COMFY_HOST="${COMFY_HOST:-0.0.0.0}"
 COMFY_PORT="${COMFY_PORT:-8818}"
 CIVITAI_TOKEN="${CIVITAI_TOKEN:-4fcb2834969399006a736ee402b061e5}"
-HF_TOKEN="${HF_TOKEN:-}"  # Token opcional do HuggingFace para login
+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
 export MAX_JOBS=32
@@ -38,14 +46,58 @@ export HF_TRANSFER_CONCURRENCY=16
 export EXT_PARALLEL=4
 
 # -----------------------------
-# Lista de downloads
+# Sistema de Estado Persistente
+# -----------------------------
+init_state() {
+    mkdir -p "$STATE_DIR"
+    touch "$STATE_FILE" "$COMPLETED_DOWNLOADS" "$COMPLETED_NODES"
+}
+
+save_state() {
+    local step="$1"
+    echo "$step" > "$STATE_FILE"
+    log_info "Estado salvo: $step"
+}
+
+get_state() {
+    if [ -f "$STATE_FILE" ]; then
+        cat "$STATE_FILE"
+    else
+        echo "0"
+    fi
+}
+
+mark_download_complete() {
+    local url="$1"
+    echo "$url" >> "$COMPLETED_DOWNLOADS"
+}
+
+is_download_complete() {
+    local url="$1"
+    grep -Fxq "$url" "$COMPLETED_DOWNLOADS" 2>/dev/null
+}
+
+mark_node_complete() {
+    local node="$1"
+    echo "$node" >> "$COMPLETED_NODES"
+}
+
+is_node_complete() {
+    local node="$1"
+    grep -Fxq "$node" "$COMPLETED_NODES" 2>/dev/null
+}
+
+reset_state() {
+    log_warn "Resetando estado da instalação..."
+    rm -rf "$STATE_DIR"
+    init_state
+}
+
+# -----------------------------
+# Lista de downloads (ATUALIZADA - sem ControlNet)
 # -----------------------------
-# Formato: "URL|TIPO|NOME_ARQUIVO_OPCIONAL"
 readonly DOWNLOAD_FILES=(
-    # Checkpoints from Civitai
-    "https://civitai.com/api/download/models/1761560?type=Model&format=SafeTensor&size=pruned&fp=fp16|checkpoints|"
-    "https://civitai.com/api/download/models/1410435?type=Model&format=SafeTensor&size=pruned&fp=fp16|checkpoints|"
-    "https://civitai.com/api/download/models/1522463?type=Model&format=SafeTensor&size=full&fp=fp16|checkpoints|"
+    # Checkpoint mantido
     "https://civitai.com/api/download/models/2122278?type=Model&format=SafeTensor&size=pruned&fp=fp16|checkpoints|raehoshiIllustXL_v60.safetensors"
     
     # LoRAs from Civitai
@@ -56,19 +108,22 @@ readonly DOWNLOAD_FILES=(
     "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
+    # LoRA from Mega (nome automático pelo megadl)
+    "mega://https://mega.nz/file/gIRTFQSQ#no6Ay3JLE9LVRi7ib9O-Jc0CW7XmG046kCgpCzDg1tY|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"
-    
-    # ControlNet
-    "hf://xinsir/controlnet-union-sdxl-1.0/diffusion_pytorch_model_promax.safetensors|controlnet|controlnet-union.safetensors"
 )
 
-# Custom nodes para instalar
+# Custom nodes para instalar (ATUALIZADA com novos nodes)
 readonly CUSTOM_NODES=(
     "https://github.com/adbrasi/huggpackreator"
     "https://github.com/adbrasi/packreator_processor"
@@ -82,6 +137,10 @@ readonly CUSTOM_NODES=(
     "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"
@@ -101,8 +160,6 @@ readonly CUSTOM_NODES=(
     "https://github.com/ltdrdata/ComfyUI-Impact-Subpack"
     "https://github.com/Cezarsaint/rand0micoUploaderLoven"
     "https://github.com/adbrasi/pixivmosaic"
-    "https://github.com/adbrasi/storitadifusita"
-    "https://github.com/adbrasi/attentionPPM"
     "https://github.com/adbrasi/futfilter"
     "https://github.com/Artificial-Sweetener/comfyui-WhiteRabbit"
     "https://github.com/shiimizu/ComfyUI_smZNodes"
@@ -122,6 +179,73 @@ command_exists() {
     command -v "$1" >/dev/null 2>&1
 }
 
+# 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
+}
+
+# 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, pulando download do Mega"
+        return 1
+    fi
+    
+    # Limpar URL - remover prefixo mega:// se houver
+    url="${url#mega://}"
+    
+    log_info "Baixando do Mega..."
+    log_info "  URL: $url"
+    log_info "  Destino: $target_dir"
+    
+    # Criar diretório temporário para download
+    local temp_dir=$(mktemp -d)
+    cd "$temp_dir"
+    
+    # Baixar com megadl (ele salva com o nome correto automaticamente)
+    if timeout 1800 megadl "$url" 2>/dev/null; then
+        # Encontrar o arquivo baixado
+        local downloaded_file=$(find . -type f -maxdepth 1 2>/dev/null | grep -v "^\.$" | head -1)
+        
+        if [ -n "$downloaded_file" ]; then
+            local actual_filename=$(basename "$downloaded_file")
+            
+            # Se foi especificado um nome E não está vazio, renomear
+            if [ -n "$filename" ] && [ "$filename" != "" ]; then
+                mv "$downloaded_file" "$target_dir/$filename"
+                log_success "Download concluído: $filename"
+            else
+                # Manter o nome original do Mega
+                mv "$downloaded_file" "$target_dir/$actual_filename"
+                log_success "Download concluído: $actual_filename"
+            fi
+            cd - >/dev/null
+            rm -rf "$temp_dir"
+            return 0
+        else
+            log_error "Arquivo baixado mas não encontrado no diretório temporário"
+        fi
+    else
+        log_error "Falha ao baixar do Mega (timeout ou erro)"
+    fi
+    
+    cd - >/dev/null
+    rm -rf "$temp_dir"
+    return 1
+}
+
 # Adicionar token do Civitai se necessário
 add_civitai_token() {
     local url="$1"
@@ -139,27 +263,20 @@ download_hf() {
     local target_dir="$3"
     local filename="$4"
     
-    # Se filename não foi fornecido, extrair do caminho
     if [ -z "$filename" ]; then
         filename=$(basename "$file_path")
     fi
     
     local target_file="$target_dir/$filename"
     
-    # Se arquivo já existe e tem tamanho razoável, pular
     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"
-    log_info "  Repo: $repo"
-    log_info "  Path: $file_path"
     
-    # Criar diretório temporário para download
     local temp_dir=$(mktemp -d)
-    
-    # Tentar primeiro com 'hf download' (novo comando)
     local download_success=false
     
     if command_exists hf; then
@@ -170,7 +287,6 @@ download_hf() {
         fi
     fi
     
-    # Se 'hf' não funcionou, tentar com huggingface-cli
     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" \
@@ -181,17 +297,13 @@ download_hf() {
     fi
     
     if [ "$download_success" = true ]; then
-        # Procurar o arquivo baixado (pode estar em subdiretórios)
-        local downloaded_file=$(find "$temp_dir" -type f \( -name "*.safetensors" -o -name "*.pth" -o -name "*.pt" -o -name "*.bin" \) 2>/dev/null | head -1)
+        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
-            # Mover para o destino correto com o nome correto
             mv "$downloaded_file" "$target_file"
             rm -rf "$temp_dir"
             log_success "Download concluído: $filename"
             return 0
-        else
-            log_error "Arquivo baixado mas não encontrado em $temp_dir"
         fi
     fi
     
@@ -206,10 +318,8 @@ download_file() {
     local target_dir="$2"
     local filename="$3"
     
-    # Adicionar token Civitai se necessário
     url=$(add_civitai_token "$url")
     
-    # Se arquivo já existe e tem tamanho razoável, pular
     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
@@ -217,7 +327,6 @@ download_file() {
     
     log_info "Baixando: ${filename:-$(basename "$url" | cut -d'?' -f1)}"
     
-    # Tentar com aria2c primeiro (mais rápido)
     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
@@ -227,7 +336,6 @@ download_file() {
         fi
     fi
     
-    # Fallback para wget
     if command_exists wget; then
         if [ -n "$filename" ]; then
             wget -q --show-progress -c -O "$target_dir/$filename" "$url" && return 0
@@ -236,7 +344,6 @@ download_file() {
         fi
     fi
     
-    # Último recurso: curl
     if command_exists curl; then
         if [ -n "$filename" ]; then
             curl -L -# -C - -o "$target_dir/$filename" "$url" && return 0
@@ -254,19 +361,32 @@ download_file() {
     return 1
 }
 
-# Processar lista de downloads
+# Processar lista de downloads (NÃO PARA EM FALHAS)
 process_downloads() {
+    local total=${#DOWNLOAD_FILES[@]}
+    local current=0
     local failed_count=0
+    local skipped_count=0
+    
+    log_info "Total de downloads: $total"
     
     for entry in "${DOWNLOAD_FILES[@]}"; do
+        ((current++))
+        log_progress "[$current/$total] Processando download..."
+        
+        # Verificar se já foi baixado
+        if is_download_complete "$entry"; then
+            log_success "Download já completado anteriormente, pulando..."
+            ((skipped_count++))
+            continue
+        fi
+        
         IFS='|' read -r url type filename <<< "$entry"
         
-        # Limpar espaços
         url=$(echo "$url" | xargs)
         type=$(echo "$type" | xargs)
         filename=$(echo "$filename" | xargs)
         
-        # Determinar diretório de destino
         local target_dir="$MODELS_DIR"
         case "$type" in
             checkpoints) target_dir="$MODELS_DIR/checkpoints" ;;
@@ -276,32 +396,40 @@ process_downloads() {
             vae_approx) target_dir="$MODELS_DIR/vae_approx" ;;
             text_encoders) target_dir="$MODELS_DIR/text_encoders" ;;
             clip_vision) target_dir="$MODELS_DIR/clip_vision" ;;
-            controlnet) target_dir="$MODELS_DIR/controlnet" ;;
             upscale_models) target_dir="$MODELS_DIR/upscale_models" ;;
             ultralytics/bbox) target_dir="$MODELS_DIR/ultralytics/bbox" ;;
             *) target_dir="$MODELS_DIR/$type" ;;
         esac
         
-        # Criar diretório se não existir
         mkdir -p "$target_dir"
         
-        # Processar URL
-        if [[ "$url" == hf://* ]]; then
+        # Determinar tipo de download e executar
+        local download_result=1
+        
+        if [[ "$url" == mega://* ]]; then
+            # Download do Mega
+            download_mega "$url" "$target_dir" "$filename" && download_result=0
+        elif [[ "$url" == hf://* ]]; then
             # HuggingFace download
-            url="${url#hf://}"  # Remove prefixo hf://
-            
-            # Separar repo e caminho do arquivo
+            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" || ((failed_count++))
+            download_hf "$repo" "$file_path" "$target_dir" "$filename" && download_result=0
         else
-            # Download normal (HTTP/HTTPS)
-            download_file "$url" "$target_dir" "$filename" || ((failed_count++))
+            # Download normal
+            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 "Download falhou, mas continuando com o próximo..."
         fi
     done
     
-    return $failed_count
+    log_info "Downloads completados: $((total - failed_count - skipped_count)) novos, $skipped_count já existentes, $failed_count falharam"
+    return 0  # Sempre retorna sucesso para não parar o script
 }
 
 # Clonar ou atualizar repositório git
@@ -310,197 +438,280 @@ clone_or_update() {
     local dest="$2"
     local node_name=$(basename "$dest")
     
-    # Timeout de 60 segundos para operações git
+    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 60 git -C "$dest" pull --ff-only 2>/dev/null || {
             log_warn "Timeout ou erro ao atualizar $node_name"
-            return 0  # Não falhar, apenas avisar
+            return 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  # Não falhar, apenas avisar
+            return 0
         }
     fi
     
-    # Instalar requirements se existir (com timeout)
     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"
         }
     fi
     
+    mark_node_complete "$url"
+    return 0
+}
+
+# Detectar GPU e ajustar PyTorch se necessário
+detect_and_fix_gpu() {
+    log_info "Detectando GPU..."
+    
+    # Verificar se é uma GPU Blackwell (RTX 5090, etc)
+    local gpu_info=$(nvidia-smi --query-gpu=name --format=csv,noheader 2>/dev/null || echo "")
+    
+    if [[ "$gpu_info" == *"5090"* ]] || [[ "$gpu_info" == *"5080"* ]] || [[ "$gpu_info" == *"Blackwell"* ]]; then
+        log_warn "GPU Blackwell detectada (${gpu_info})"
+        log_info "Instalando PyTorch nightly para suporte completo..."
+        
+        # Instalar PyTorch nightly para Blackwell
+        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
+        python3 -m pip install --pre --upgrade --no-cache-dir torchaudio --extra-index-url https://download.pytorch.org/whl/nightly/cu128
+        
+        log_success "PyTorch nightly instalado para Blackwell"
+    else
+        log_info "GPU padrão detectada: ${gpu_info:-Não detectada}"
+        log_info "Usando PyTorch stable"
+    fi
+}
+
+# -----------------------------
+# Menu de opções
+# -----------------------------
+show_menu() {
+    echo ""
+    log_info "========================================="
+    log_info " ComfyUI Custom Setup v2.0"
+    log_info "========================================="
+    echo ""
+    
+    local current_state=$(get_state)
+    
+    if [ "$current_state" != "0" ]; 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 ""
+        read -p "Escolha uma opção [1-3]: " choice
+        
+        case $choice in
+            1) return 0 ;;
+            2) reset_state; return 0 ;;
+            3) exit 0 ;;
+            *) echo "Opção inválida"; exit 1 ;;
+        esac
+    fi
+    
     return 0
 }
 
 # -----------------------------
 # Instalação principal
 # -----------------------------
-echo ""
-log_info "========================================="
-log_info " ComfyUI Custom Setup"
-log_info "========================================="
-echo ""
-
-# [1] Verificar e instalar dependências
-log_info "[1/7] Verificando e instalando dependências..."
-
-# Instalar aria2c do sistema se não existir
-if ! command_exists aria2c; then
-    log_info "Instalando aria2c do sistema..."
-    apt-get update -qq && apt-get install -y -qq aria2 2>/dev/null
-fi
-
-for cmd in python3 git wget curl; do
-    if ! command_exists "$cmd"; then
-        log_error "Dependência faltando: $cmd"
-        exit 1
+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/8] 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
-done
-log_success "Dependências OK"
-
-# [2] Atualizar pip e instalar ferramentas
-log_info "[2/7] Preparando ferramentas Python..."
-python3 -m pip install --upgrade pip wheel setuptools -q
-
-# Instalar huggingface_hub com CLI e hf_transfer para velocidade máxima
-python3 -m pip install --upgrade "huggingface_hub[cli,hf_transfer]>=0.26.0" comfy-cli -q
-
-# Configurar token HF se disponível (opcional)
-HF_TOKEN="${HF_TOKEN:-}"
-if [ -n "$HF_TOKEN" ]; then
-    if huggingface-cli login --token "$HF_TOKEN" 2>/dev/null; then
-        log_success "Login HF configurado"
-    else
-        log_warn "Falha ao configurar login HF, continuando sem autenticação"
+    
+    # [STEP 2] Atualizar pip e instalar ferramentas
+    if [ "$current_state" -lt "2" ]; then
+        log_info "[2/8] 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
-fi
-
-# Verificar comandos disponíveis
-if command_exists hf; then
-    log_success "Comando 'hf' disponível (recomendado)"
-elif command_exists huggingface-cli; then
-    log_warn "Usando 'huggingface-cli' (deprecated, mas funcional)"
-else
-    log_error "Nenhum comando HF disponível!"
-fi
-
-log_success "Ferramentas instaladas"
-
-# [3] Instalar PyTorch com CUDA 12.8
-log_info "[3/7] Instalando PyTorch com CUDA 12.8..."
-python3 -m pip install --upgrade torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu128
-log_success "PyTorch com CUDA 12.8 instalado"
-
-# [4] Instalar ComfyUI
-log_info "[4/7] Instalando ComfyUI..."
-if [ -f "$COMFY_DIR/main.py" ]; then
-    log_warn "ComfyUI já existe, pulando instalação base"
-else
-    comfy --skip-prompt install --fast-deps --nvidia
-    
-    # Verificar instalação
-    if [ ! -f "$COMFY_DIR/main.py" ]; then
-        log_error "ComfyUI não foi instalado corretamente!"
-        exit 1
+    
+    # [STEP 3] Instalar PyTorch inicial
+    if [ "$current_state" -lt "3" ]; then
+        log_info "[3/8] Instalando PyTorch com CUDA 12.8..."
+        python3 -m pip install --upgrade torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu128
+        log_success "PyTorch instalado"
+        save_state "3"
     fi
-fi
-log_success "ComfyUI instalado em: $COMFY_DIR"
-
-# [5] Baixar modelos
-log_info "[5/7] Baixando modelos..."
-mkdir -p "$MODELS_DIR"/{diffusion_models,loras,vae,vae_approx,text_encoders,clip_vision,controlnet,upscale_models,checkpoints,ultralytics/bbox}
-
-# Verificar ferramentas de download
-if command_exists aria2c; then
-    log_success "Usando aria2c para downloads (rápido)"
-else
-    log_warn "aria2c não encontrado, usando wget/curl (mais lento)"
-fi
-
-# Processar downloads
-if process_downloads; then
-    log_success "Todos os modelos baixados com sucesso"
-else
-    log_warn "Alguns downloads falharam, mas continuando..."
-fi
-
-# [6] Instalar custom nodes
-log_info "[6/7] Instalando custom nodes..."
-CN_DIR="$COMFY_DIR/custom_nodes"
-mkdir -p "$CN_DIR"
-
-# Instalar nodes do GitHub
-for repo_url in "${CUSTOM_NODES[@]}"; do
-    node_name=$(basename "$repo_url")
-    clone_or_update "$repo_url" "$CN_DIR/$node_name"
-done
-
-# Instalar nodes especiais via comfy-cli
-for node_name in "${SPECIAL_NODES[@]}"; do
-    log_info "Instalando $node_name via comfy-cli..."
-    comfy node install "$node_name" 2>/dev/null || log_warn "Falha ao instalar $node_name"
-done
-
-log_success "Custom nodes instalados"
-
-# [7] Verificação final
-echo ""
-log_info "[7/7] Verificando instalação..."
-
-# Verificar arquivos principais
-check_critical_files() {
-    local files=(
-        "$COMFY_DIR/main.py"
-    )
-    
-    local missing=0
-    for file in "${files[@]}"; do
-        if [ -f "$file" ]; then
-            log_success "✓ $(basename "$file")"
+    
+    # [STEP 4] Instalar ComfyUI
+    if [ "$current_state" -lt "4" ]; then
+        log_info "[4/8] Instalando ComfyUI..."
+        if [ -f "$COMFY_DIR/main.py" ]; then
+            log_warn "ComfyUI já existe"
         else
-            log_error "✗ $(basename "$file") não encontrado"
-            ((missing++))
+            comfy --skip-prompt install --fast-deps --nvidia
+            
+            if [ ! -f "$COMFY_DIR/main.py" ]; then
+                log_error "ComfyUI não foi instalado corretamente!"
+                exit 1
+            fi
         fi
-    done
+        log_success "ComfyUI instalado"
+        save_state "4"
+    fi
+    
+    # [STEP 5] Detectar e corrigir para GPU Blackwell se necessário
+    if [ "$current_state" -lt "5" ]; then
+        log_info "[5/8] Verificando compatibilidade de GPU..."
+        detect_and_fix_gpu
+        save_state "5"
+    fi
+    
+    # [STEP 6] Baixar modelos
+    if [ "$current_state" -lt "6" ]; then
+        log_info "[6/8] Baixando modelos..."
+        mkdir -p "$MODELS_DIR"/{diffusion_models,loras,vae,vae_approx,text_encoders,clip_vision,upscale_models,checkpoints,ultralytics/bbox}
+        
+        process_downloads  # Não falha mesmo se alguns downloads falharem
+        
+        log_success "Processo de download concluído"
+        save_state "6"
+    fi
+    
+    # [STEP 7] Instalar custom nodes
+    if [ "$current_state" -lt "7" ]; then
+        log_info "[7/8] Instalando custom nodes..."
+        CN_DIR="$COMFY_DIR/custom_nodes"
+        mkdir -p "$CN_DIR"
+        
+        local total_nodes=${#CUSTOM_NODES[@]}
+        local current_node=0
+        
+        for repo_url in "${CUSTOM_NODES[@]}"; do
+            ((current_node++))
+            node_name=$(basename "$repo_url")
+            log_progress "[$current_node/$total_nodes] Instalando: $node_name"
+            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 "7"
+    fi
+    
+    # [STEP 8] Verificação final
+    log_info "[8/8] Verificando instalação..."
+    
+    echo ""
+    log_info "Estatísticas da instalação:"
+    
+    # Verificar ComfyUI
+    if [ -f "$COMFY_DIR/main.py" ]; then
+        log_success "✓ ComfyUI instalado"
+    else
+        log_error "✗ ComfyUI não encontrado"
+    fi
     
-    # Listar modelos baixados
+    # Contar modelos
     echo ""
-    log_info "Modelos disponíveis:"
-    for dir in checkpoints loras vae vae_approx text_encoders clip_vision controlnet upscale_models ultralytics; do
+    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" \) 2>/dev/null | wc -l)
+            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
     done
     
-    # Listar custom nodes instalados
+    # Contar nodes
     echo ""
-    log_info "Custom nodes instalados:"
-    local node_count=$(find "$CN_DIR" -maxdepth 1 -type d | wc -l)
-    log_success "  Total: $((node_count - 1)) nodes"
+    local node_count=$(find "$CN_DIR" -maxdepth 1 -type d 2>/dev/null | wc -l)
+    log_success "Custom nodes: $((node_count - 1)) instalados"
     
-    return $missing
-}
-
-if check_critical_files; then
-    log_success "Instalação concluída com sucesso!"
+    # Verificar GPU
     echo ""
-    log_info "========================================="
-    log_info "Iniciando ComfyUI..."
-    log_info "URL: http://localhost:$COMFY_PORT"
-    log_info "========================================="
+    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 "8"
     
-    cd "$COMFY_DIR"
-    exec comfy launch -- --listen "$COMFY_HOST" --port "$COMFY_PORT"
-else
-    log_error "Alguns arquivos críticos não foram encontrados!"
-    log_info "Tentando iniciar ComfyUI mesmo assim..."
+    echo ""
+    log_success "========================================="
+    log_success " Instalação concluída!"
+    log_success "========================================="
     
-    cd "$COMFY_DIR"
-    exec comfy launch -- --listen "$COMFY_HOST" --port "$COMFY_PORT"
-fi
+    echo ""
+    echo "Opções:"
+    echo "  1) Iniciar ComfyUI agora"
+    echo "  2) Resetar estado da instalação"
+    echo "  3) Sair"
+    echo ""
+    read -p "Escolha uma opção [1-3]: " final_choice
+    
+    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_warn "Opção inválida"
+            ;;
+    esac
+}
+
+# Executar instalação
+main "$@"