Update image_Studio.sh

image_Studio.sh

@@ -1,107 +1,60 @@
 #!/usr/bin/env bash
-# setup_comfyui_custom_v2.sh
-# FIX: instalar SEM pré-criar $COMFY_DIR (evita falso positivo do comfy-cli)
-# + timeouts/retries em downloads, headers p/ Civitai, estado robusto
-# + Blackwell (5090/5080) com torch/vision nightly via --pre --upgrade --no-cache-dir
-# + comfy-cli mantido para instalar/rodar/instalar nodes
+# setup_comfyui_simple.sh
+# Script simples para instalação do ComfyUI
 
 set -euo pipefail
 
-# --------------------------------
-# Interatividade
-# --------------------------------
-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'
+NC='\033[0m'
 
-# --------------------------------
-# Cores/logs
-# --------------------------------
-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_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"; }
+log_warn() { echo -e "${YELLOW}[!]${NC} $1"; }
+log_error() { echo -e "${RED}[✗]${NC} $1"; }
 
-# --------------------------------
-# Config
-# --------------------------------
+# Configuração
 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 (pedido do usuário)
-CIVITAI_TOKEN="${CIVITAI_TOKEN:-4fcb2834969399006a736ee402b061e5}"
-HF_TOKEN="${HF_TOKEN:-}"
-
-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"
+COMFY_HOST="0.0.0.0"
+COMFY_PORT="8818"
+CIVITAI_TOKEN="4fcb2834969399006a736ee402b061e5"
 
 # Performance
 export MAX_JOBS=32
-export NVCC_APPEND_FLAGS="--threads 8"
-export UV_SYSTEM_PYTHON=1
-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
-
-# Rede
-DOWNLOAD_TIMEOUT="${DOWNLOAD_TIMEOUT:-900}"   # 15 min por artefato
-CONNECT_TIMEOUT="${CONNECT_TIMEOUT:-20}"
-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/"
-
-# --------------------------------
-# Estado
-# --------------------------------
-init_state(){ mkdir -p "$STATE_DIR"; touch "$STATE_FILE" "$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; }
-
-normalize_entry(){ # "url|type|filename"
-  local url="$1" type="$2" filename="${3:-}"
-  url="$(echo -n "$url" | xargs)"; type="$(echo -n "$type" | xargs)"; filename="$(echo -n "$filename" | xargs)"
-  echo "${url}|${type}|${filename}"
-}
-mark_download_complete(){ echo "$1" >> "$COMPLETED_DOWNLOADS"; }
-is_download_complete(){ [[ -f "$COMPLETED_DOWNLOADS" ]] && grep -Fxq "$1" "$COMPLETED_DOWNLOADS"; }
-mark_node_complete(){ echo "$1" >> "$COMPLETED_NODES"; }
-is_node_complete(){ [[ -f "$COMPLETED_NODES" ]] && grep -Fxq "$1" "$COMPLETED_NODES"; }
-
-# --------------------------------
-# Downloads (lista)
-# --------------------------------
-readonly DOWNLOAD_FILES=(
-  "mega://https://mega.nz/file/gIRTFQSQ#no6Ay3JLE9LVRi7ib9O-Jc0CW7XmG046kCgpCzDg1tY|loras|"
-
-  "https://civitai.com/api/download/models/2122278?type=Model&format=SafeTensor&size=pruned&fp=fp16|checkpoints|raehoshiIllustXL_v60.safetensors"
-
-  "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|"
-
-  "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"
-
-  "https://huggingface.co/adbrasi/testedownload/resolve/main/99coins_anime_girl_face_m_seg.pt|ultralytics/bbox|99coins_anime_girl_face_m_seg.pt"
+
+# Lista de downloads - Formato: "URL|TIPO|NOME_OPCIONAL"
+DOWNLOAD_FILES=(
+    # Mega primeiro (sem nome forçado - usar nome original)
+    "mega://https://mega.nz/file/gIRTFQSQ#no6Ay3JLE9LVRi7ib9O-Jc0CW7XmG046kCgpCzDg1tY|loras|"
+    
+    # Checkpoints
+    "https://civitai.com/api/download/models/2122278?type=Model&format=SafeTensor&size=pruned&fp=fp16|checkpoints|raehoshiIllustXL_v60.safetensors"
+    
+    # LoRAs 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|"
+    
+    # Upscalers
+    "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/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"
+    
+    # Ultralytics
+    "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
-readonly CUSTOM_NODES=(
+CUSTOM_NODES=(
   "https://github.com/adbrasi/huggpackreator"
   "https://github.com/adbrasi/packreator_processor"
   "https://github.com/adbrasi/Packreator_manager"
@@ -143,361 +96,183 @@ readonly CUSTOM_NODES=(
   "https://github.com/CoreyCorza/ComfyUI-CRZnodes"
 )
 
-readonly SPECIAL_NODES=("Civicomfy" "ComfyUI-RMBG")
-
-# --------------------------------
-# 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 (MEGA pode não funcionar)"
-  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 HuggingFace
+download_hf() {
+    local url="$1"
+    local target_dir="$2"
+    local filename="$3"
+    
+    if [ -n "$filename" ] && [ -f "$target_dir/$filename" ]; then
+        log_success "Já existe: $filename"
+        return 0
+    fi
+    
+    log_info "Baixando HF: $filename"
+    
+    if command -v hf >/dev/null 2>&1; then
+        cd "$target_dir"
+        hf download "$(dirname "$url")" "$(basename "$url")" --local-dir . 2>/dev/null
+        cd - >/dev/null
+    else
+        wget -q --show-progress -c -O "$target_dir/$filename" "https://huggingface.co/$url"
+    fi
 }
 
-download_hf(){ # via hf/huggingface-cli com timeout
-  local repo="$1" file_path="$2" target_dir="$3" filename="$4"
-  [[ -z "$filename" ]] && filename="$(basename "$file_path")"
-  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 "HF: $repo/$file_path -> $filename"
-  local tmp; tmp="$(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 "$tmp" --local-dir-use-symlinks False >/dev/null 2>&1; then ok=true; fi
-  fi
-  if ! $ok && command_exists huggingface-cli; 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 "$tmp" --local-dir-use-symlinks False >/dev/null 2>&1; then ok=true; fi
-  fi
-
-  if $ok; then
-    local f; f="$(find "$tmp" -type f \( -name "*.safetensors" -o -name "*.pth" -o -name "*.pt" -o -name "*.bin" -o -name "*.onnx" \) | head -1)"
-    if [[ -n "$f" ]]; then mv "$f" "$out"; rm -rf "$tmp"; log_success "HF OK: $filename"; return 0; fi
-  fi
-  rm -rf "$tmp"; log_error "HF falhou: $filename"; return 1
+# Download Mega
+download_mega() {
+    local url="$1"
+    local target_dir="$2"
+    
+    # Remover prefixo mega://
+    url="${url#mega://}"
+    
+    log_info "Baixando do Mega..."
+    
+    cd "$target_dir"
+    # Usar megadl simples sem forçar nome
+    if megadl "$url" 2>/dev/null; then
+        log_success "Mega download OK"
+    else
+        log_warn "Mega download falhou"
+    fi
+    cd - >/dev/null
 }
 
-download_mega(){ # com timeout
-  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 normal
+download_file() {
+    local url="$1"
+    local target_dir="$2"
+    local filename="$3"
+    
+    # Adicionar token Civitai se necessário
+    if [[ "$url" == *"civitai.com"* ]]; then
+        url="${url}&token=${CIVITAI_TOKEN}"
+    fi
+    
+    if [ -n "$filename" ] && [ -f "$target_dir/$filename" ]; then
+        log_success "Já existe: $filename"
+        return 0
+    fi
+    
+    log_info "Baixando: ${filename:-arquivo}"
+    
+    if command -v aria2c >/dev/null 2>&1; then
+        aria2c -c -x 4 -s 4 --console-log-level=warn --dir="$target_dir" --out="$filename" "$url"
+    elif [ -n "$filename" ]; then
+        wget -q --show-progress -c -O "$target_dir/$filename" "$url"
+    else
+        cd "$target_dir"
+        wget -q --show-progress -c "$url"
+        cd - >/dev/null
+    fi
 }
 
-detect_remote_name(){ # pega filename do cabeçalho
-  local url="$1" 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"
+# Processar downloads
+process_downloads() {
+    for entry in "${DOWNLOAD_FILES[@]}"; do
+        IFS='|' read -r url type filename <<< "$entry"
+        
+        # Criar diretório
+        local target_dir="$MODELS_DIR/$type"
+        mkdir -p "$target_dir"
+        
+        # Baixar arquivo
+        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"
+        else
+            download_file "$url" "$target_dir" "$filename"
+        fi
+    done
 }
 
-download_file(){ # CURL -> WGET -> ARIA2C (todos com timeout & retries)
-  local url="$1" target_dir="$2" filename="${3:-}"
-  url="$(add_civitai_token "$url")"
-  [[ -z "$filename" ]] && filename="$(detect_remote_name "$url")"
-  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"
-
-  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
+# Clone git repo
+clone_repo() {
+    local url="$1"
+    local dest="$2"
+    
+    if [ -d "$dest/.git" ]; then
+        git -C "$dest" pull --ff-only 2>/dev/null || true
     else
-      log_warn "CURL falhou para $filename"
+        git clone --depth 1 "$url" "$dest" 2>/dev/null || true
     fi
-  fi
-
-  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"
+    
+    # Instalar requirements se existir
+    if [ -f "$dest/requirements.txt" ]; then
+        pip install -q -r "$dest/requirements.txt" 2>/dev/null || true
     fi
-  fi
+}
 
-  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_warn "ARIA2C falhou para $filename"
-    fi
-  fi
+# ========== INSTALAÇÃO ==========
 
-  log_error "Falha final em: $url"
-  return 1
-}
+log_info "========================================="
+log_info " ComfyUI Setup"
+log_info "========================================="
 
-process_downloads(){
-  local total="${#DOWNLOAD_FILES[@]}" current=0 failed=0
-  log_info "Total de downloads: $total"
-  log_info "Iniciando loop de downloads..."
+# 1. Instalar dependências
+log_info "[1/5] Instalando dependências..."
 
-  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")"
+apt-get update -qq
+apt-get install -y -qq aria2 megatools git wget curl 2>/dev/null
 
-    if is_download_complete "$key"; then
-      log_success "[$current/$total] Já baixado: ${filename:-$(basename "$url" | cut -d'?' -f1)}"; continue
-    fi
+pip install --upgrade pip wheel setuptools -q
+pip install --upgrade "huggingface_hub[cli,hf_transfer]" comfy-cli -q
 
-    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="$(echo "$url" | cut -d'/' -f1-2)"
-      local file_path="$(echo "$url" | cut -d'/' -f3-)"
-      download_hf "$repo" "$file_path" "$target_dir" "$filename" && ok=0
-    elif [[ "$url" == https://huggingface.co/*/resolve/* ]]; then
-      download_file "$url" "$target_dir" "$filename" && ok=0
-    else
-      download_file "$url" "$target_dir" "$filename" && ok=0
-    fi
+log_success "Dependências instaladas"
 
-    if [[ $ok -eq 0 ]]; then
-      mark_download_complete "$key"
-    else
-      ((failed++)); log_warn "[$current/$total] Falhou, segue..."
-    fi
-  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
-}
+# 2. Instalar ComfyUI
+log_info "[2/5] Instalando ComfyUI..."
 
-clone_or_update(){
-  local url="$1" dest="$2" node_name="$(basename "$dest")"
-  if is_node_complete "$url"; then log_success "Node já instalado: $node_name"; return 0; fi
+if [ -f "$COMFY_DIR/main.py" ]; then
+    log_warn "ComfyUI já existe"
+else
+    comfy --skip-prompt install --fast-deps --nvidia
+fi
 
-  if [[ -d "$dest/.git" ]]; then
-    # sanity: garante que é o repo esperado
-    if ! git -C "$dest" remote -v | grep -qiE 'comfy|comfyui|github\.com'; then
-      log_warn "$node_name tem git estranho; removendo..."
-      rm -rf "$dest"
-    fi
-  fi
-
-  if [[ -d "$dest/.git" ]]; then
-    log_info "Atualizando $node_name..."
-    timeout 180 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 300 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
-}
+log_success "ComfyUI instalado"
 
-# --------------------------------
-# 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
-}
+# 3. Corrigir PyTorch para RTX 5090 se detectado
+log_info "[3/5] Verificando GPU..."
 
-# --------------------------------
-# Main
-# --------------------------------
-main(){
-  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 via comfy-cli (sem pré-criar $COMFY_DIR)
-  if [[ "$s" -lt "3" ]]; then
-    log_info "[3/6] Instalando ComfyUI (comfy-cli)..."
-
-    if [[ -d "$COMFY_DIR" ]]; then
-      if [[ -d "$COMFY_DIR/.git" || -f "$COMFY_DIR/main.py" ]]; then
-        log_warn "ComfyUI já parece instalado em $COMFY_DIR"
-      else
-        log_warn "Diretório $COMFY_DIR existe porém não é ComfyUI válido. Removendo..."
-        rm -rf "$COMFY_DIR"
-      fi
-    fi
+GPU_INFO=$(nvidia-smi --query-gpu=name --format=csv,noheader 2>/dev/null || echo "")
 
-    # git sanity: se existir git mas não for o repo esperado, remove
-    if [[ -d "$COMFY_DIR/.git" ]]; then
-      if ! git -C "$COMFY_DIR" remote -v 2>/dev/null | grep -qi 'comfyanonymous/ComfyUI'; then
-        log_warn "Git em $COMFY_DIR não aponta para comfyanonymous/ComfyUI. Removendo..."
-        rm -rf "$COMFY_DIR"
-      fi
-    fi
+if [[ "$GPU_INFO" == *"5090"* ]] || [[ "$GPU_INFO" == *"5080"* ]]; then
+    log_warn "RTX 5090/5080 detectada - instalando PyTorch nightly"
+    pip install --force-reinstall --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu128
+else
+    log_info "GPU: ${GPU_INFO:-Não detectada}"
+fi
 
-    if [[ ! -f "$COMFY_DIR/main.py" ]]; then
-      comfy --workspace "$COMFY_DIR" --skip-prompt install --nvidia --fast-deps || {
-        log_error "ComfyUI não instalou (comfy-cli erro)"; exit 1; }
-    fi
+log_success "PyTorch configurado"
 
-    [[ -f "$COMFY_DIR/main.py" ]] || { log_error "ComfyUI não encontrado após instalar"; exit 1; }
-    log_success "ComfyUI instalado em: $COMFY_DIR"
-    save_state "3"
-  fi
-
-  # STEP 4: Blackwell nightly (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 (pedido do usuário)"
-      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_info "GPU: ${gpu_info:-Não detectada} - mantendo PyTorch padrão"
-    fi
-    save_state "4"
-  fi
-
-  # STEP 5: modelos (só agora criamos $MODELS_DIR)
-  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
-    log_success "Custom nodes instalados"; save_state "6"
-  fi
-
-  # Verificação final
-  log_info "Verificando instalação..."
-  [[ -f "$COMFY_DIR/main.py" ]] && log_success "✓ ComfyUI instalado" || log_error "✗ ComfyUI não encontrado"
-
-  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
-  done
+# 4. Baixar modelos
+log_info "[4/5] Baixando modelos..."
 
-  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"
+process_downloads
 
-  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"
+log_success "Downloads processados"
 
-  save_state "6"
+# 5. Instalar custom nodes
+log_info "[5/5] Instalando custom nodes..."
 
-  echo ""; log_success "========================================="; log_success " Instalação concluída!"; log_success "========================================="
+CN_DIR="$COMFY_DIR/custom_nodes"
+mkdir -p "$CN_DIR"
 
-  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
+for repo in "${CUSTOM_NODES[@]}"; do
+    node_name=$(basename "$repo")
+    clone_repo "$repo" "$CN_DIR/$node_name"
+done
 
-  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 o script novamente para reinstalar." ;;
-    3) log_info "Saindo..." ;;
-    *) cd "$COMFY_DIR"; exec comfy launch -- --listen "$COMFY_HOST" --port "$COMFY_PORT" ;;
-  esac
-}
+log_success "Custom nodes instalados"
+
+# Verificar e iniciar
+log_info "========================================="
+log_success "Instalação concluída!"
+log_info "Iniciando ComfyUI..."
+log_info "URL: http://localhost:$COMFY_PORT"
+log_info "========================================="
 
-main "$@"
+cd "$COMFY_DIR"
+exec comfy launch -- --listen "$COMFY_HOST" --port "$COMFY_PORT"