Update image_Studio.sh

image_Studio.sh

@@ -1,393 +1,724 @@
 #!/usr/bin/env bash
-# setup_comfyui_simplificado.sh
-# Objetivo: script simples, previsível e resiliente para instalar/atualizar ComfyUI
-# - Ambiente isolado em venv (sem poluir Python do sistema)
-# - Torch/Torchvision/Torchaudio sempre compatíveis (um único passo)
-# - Estado persistente com retomada de steps, downloads e nodes
-# - Downloads com retry/timeout e suporte Mega/HF/HTTP(s)
-# - Layout e variáveis fáceis de editar (manifesto inline)
+# setup_comfyui_custom_v2.sh
+# Script otimizado com suporte para RTX 5090, estado persistente e downloads do Mega
 
-set -Eeuo pipefail
+set -euo pipefail
 
-# ====== Configuração básica ======
-COMFY_DIR=${COMFY_DIR:-"/root/comfy/ComfyUI"}
-VENV_DIR=${VENV_DIR:-"$COMFY_DIR/.venv"}
-PYTHON="$VENV_DIR/bin/python"
-PIP="$VENV_DIR/bin/pip"
-COMFY_HOST=${COMFY_HOST:-"0.0.0.0"}
-COMFY_PORT=${COMFY_PORT:-"8818"}
+# Detectar se está rodando via pipe (não-interativo)
+if [ ! -t 0 ] || [ ! -t 1 ]; then
+    INTERACTIVE=false
+else
+    INTERACTIVE=true
+fi
 
-# Tokens opcionais
-CIVITAI_TOKEN=${CIVITAI_TOKEN:-""}   # NUNCA hardcode aqui. Passe por env.
-HF_TOKEN=${HF_TOKEN:-""}
+# -----------------------------
+# 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'
 
-# Diretórios de modelo
+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
+# -----------------------------
+COMFY_DIR="/root/comfy/ComfyUI"
 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:-}"
 
-# Estado
-STATE_DIR=${STATE_DIR:-"/root/.comfyui_setup_state"}
-STATE_FILE="$STATE_DIR/step.txt"
-DONE_DL="$STATE_DIR/completed_downloads.txt"
-DONE_NODE="$STATE_DIR/completed_nodes.txt"
-mkdir -p "$STATE_DIR"
-: > "$DONE_DL" 2>/dev/null || true
-: > "$DONE_NODE" 2>/dev/null || true
+# 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"
 
-# Tempo & rede
-RETRIES=${RETRIES:-3}
-CONNECT_TIMEOUT=${CONNECT_TIMEOUT:-30}    # seg
-DL_TIMEOUT=${DL_TIMEOUT:-3600}            # seg
+# Configurações de 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
 
-# UI
-BLUE='\033[0;34m'; GREEN='\033[0;32m'; YELLOW='\033[1;33m'; RED='\033[0;31m'; CYAN='\033[0;36m'; NC='\033[0m'
-log_info(){ echo -e "${BLUE}[INFO]${NC} $*"; }
-log_ok(){   echo -e "${GREEN}[✓]${NC} $*"; }
-log_warn(){ echo -e "${YELLOW}[!]${NC} $*"; }
-log_err(){  echo -e "${RED}[✗]${NC} $*"; }
-log_step(){ echo -e "${CYAN}[→]${NC} $*"; }
+# -----------------------------
+# Sistema de Estado Persistente
+# -----------------------------
+init_state() {
+    mkdir -p "$STATE_DIR"
+    touch "$STATE_FILE" "$COMPLETED_DOWNLOADS" "$COMPLETED_NODES"
+}
 
-# ====== Manifestos (edite à vontade) ======
-# Formato: "URL|subpasta_em_models|nome_arquivo(opcional)"
-# Use mega:// para links do Mega. Para HF você pode usar hf://<repo>/<path>
-read -r -d '' DOWNLOADS << 'EOF'
-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|
-mega://https://mega.nz/file/gIRTFQSQ#no6Ay3JLE9LVRi7ib9O-Jc0CW7XmG046kCgpCzDg1tY|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
-EOF
+save_state() {
+    local step="$1"
+    echo "$step" > "$STATE_FILE"
+    log_info "Estado salvo: $step"
+}
 
-# Lista de nodes (clones git). Edite à vontade.
-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
-)
+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
+}
 
-# ====== Estado ======
-get_state(){ [[ -s "$STATE_FILE" ]] && cat "$STATE_FILE" || echo 0; }
-save_state(){ printf '%s\n' "$1" >"$STATE_FILE.tmp" && mv -f "$STATE_FILE.tmp" "$STATE_FILE"; log_info "Estado salvo: $1"; }
-mark_done(){ local f="$1"; shift; local key="$*"; grep -Fxq "$key" "$f" 2>/dev/null || echo "$key" >> "$f"; }
-was_done(){ local f="$1"; shift; local key="$*"; grep -Fxq "$key" "$f" 2>/dev/null; }
+mark_download_complete() {
+    local url="$1"
+    echo "$url" >> "$COMPLETED_DOWNLOADS"
+}
 
-# ====== Dependências base ======
-need(){ command -v "$1" >/dev/null 2>&1; }
-ensure_base(){
-  log_step "[1/7] Instalando dependências do SO (git, curl, wget, aria2, megatools, python3-venv)…"
-  export DEBIAN_FRONTEND=noninteractive
-  apt-get update -qq
-  apt-get install -y -qq git curl wget aria2 megatools python3 python3-venv ca-certificates
-  log_ok "Dependências do SO OK"
-  save_state 1
+is_download_complete() {
+    local url="$1"
+    grep -Fxq "$url" "$COMPLETED_DOWNLOADS" 2>/dev/null
 }
 
-# ====== Ambiente Python isolado ======
-ensure_venv(){
-  log_step "[2/7] Preparando venv…"
-  if [[ ! -x "$PYTHON" ]]; then
-    mkdir -p "$COMFY_DIR"
-    python3 -m venv "$VENV_DIR"
-  fi
-  "$PIP" install -q --upgrade pip wheel setuptools
-  # Login opcional no HF
-  if [[ -n "$HF_TOKEN" ]]; then
-    "$PYTHON" -m pip install -q huggingface_hub[cli,hf_transfer]
-    "$VENV_DIR/bin/hf" auth login --token "$HF_TOKEN" || log_warn "HF auth falhou (seguindo sem login)"
-  fi
-  log_ok "Venv OK"
-  save_state 2
+mark_node_complete() {
+    local node="$1"
+    echo "$node" >> "$COMPLETED_NODES"
 }
 
-# ====== Torch consistente (1 passo) ======
-install_torch(){
-  log_step "[3/7] Instalando PyTorch/Torchvision/Torchaudio…"
-  local gpu_name="$(nvidia-smi --query-gpu=name --format=csv,noheader 2>/dev/null || true)"
-  local use_nightly=0
-  if echo "$gpu_name" | grep -Eqi "5090|5080|blackwell"; then use_nightly=1; fi
+is_node_complete() {
+    local node="$1"
+    grep -Fxq "$node" "$COMPLETED_NODES" 2>/dev/null
+}
 
-  if (( use_nightly )); then
-    log_warn "GPU Blackwell detectada ($gpu_name). Usando *nightly cu128* compatível."
-    "$PIP" install --pre -q \
-      --extra-index-url https://download.pytorch.org/whl/nightly/cu128 \
-      torch torchvision torchaudio
-  else
-    log_info "GPU: ${gpu_name:-Não detectada}. Usando *estável cu128*."
-    "$PIP" install -q \
-      --index-url https://download.pytorch.org/whl/cu128 \
-      torch torchvision torchaudio
-  fi
-  log_ok "Torch stack OK"
-  save_state 3
+reset_state() {
+    log_warn "Resetando estado da instalação..."
+    rm -rf "$STATE_DIR"
+    init_state
 }
 
-# ====== ComfyUI ======
-install_comfy(){
-  log_step "[4/7] Instalando/Atualizando ComfyUI…"
-  if [[ -d "$COMFY_DIR/.git" ]]; then
-    git -C "$COMFY_DIR" pull --ff-only || log_warn "git pull falhou (mantendo versão atual)"
-  else
-    git clone --recursive https://github.com/comfyanonymous/ComfyUI "$COMFY_DIR"
-  fi
-  "$PIP" install -q -r "$COMFY_DIR/requirements.txt"
-  # Manager (útil p/ gerir nodes)
-  local mdir="$COMFY_DIR/custom_nodes/ComfyUI-Manager"
-  if [[ -d "$mdir/.git" ]]; then
-    git -C "$mdir" pull --ff-only || true
-  else
-    git clone --depth 1 https://github.com/ltdrdata/ComfyUI-Manager "$mdir"
-  fi
-  log_ok "ComfyUI OK"
-  save_state 4
+# -----------------------------
+# Lista de downloads (ATUALIZADA - sem ControlNet)
+# -----------------------------
+readonly DOWNLOAD_FILES=(
+    # 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|"
+    
+    # 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"
+)
+
+# Custom nodes para instalar (ATUALIZADA com novos 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"
+)
+
+# Custom nodes que precisam de instalação especial via comfy-cli
+readonly SPECIAL_NODES=(
+    "Civicomfy"
+    "ComfyUI-RMBG"
+)
+
+# -----------------------------
+# Funções auxiliares
+# -----------------------------
+command_exists() {
+    command -v "$1" >/dev/null 2>&1
 }
 
-# ====== Downloads ======
-add_civitai_token(){
-  local url="$1"
-  if [[ "$url" == *"civitai.com/api/download"* && "$url" != *"token="* ]]; then
-    if [[ -n "$CIVITAI_TOKEN" ]]; then
-      printf '%s&token=%s' "$url" "$CIVITAI_TOKEN"
-      return 0
-    else
-      log_warn "Civitai sem token: pulando $url"
-      return 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
-  fi
-  printf '%s' "$url"
 }
 
-http_download(){ # url dir [filename]
-  local url="$1"; local dir="$2"; local name="${3:-}"
-  mkdir -p "$dir"
-  local final_url
-  if ! final_url=$(add_civitai_token "$url"); then return 1; fi
-  if need aria2c; then
-    local args=("-c" "-x16" "-s16" "-k1M" "--max-tries=$RETRIES" "--retry-wait=5" "--timeout=$CONNECT_TIMEOUT" "--console-log-level=warn" "--summary-interval=10")
-    if [[ -n "$name" ]]; then
-      aria2c "${args[@]}" --dir="$dir" --out="$name" "$final_url" && return 0
-    else
-      aria2c "${args[@]}" --dir="$dir" "$final_url" && return 0
+# 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
-  fi
-  if need wget; then
-    local wargs=("-c" "--tries=$RETRIES" "--timeout=$CONNECT_TIMEOUT" "--read-timeout=$DL_TIMEOUT")
-    if [[ -n "$name" ]]; then
-      wget -q "${wargs[@]}" -O "$dir/$name" "$final_url" && return 0
+    
+    # Remover prefixo mega:// se houver
+    url="${url#mega://}"
+    
+    log_info "Baixando do Mega para $target_dir"
+    
+    # Usar megadl simples como você sempre usa
+    cd "$target_dir"
+    if megadl "$url"; then
+        # Se especificou nome, renomear
+        if [ -n "$filename" ]; then
+            # Encontrar o arquivo baixado (o mais recente)
+            local downloaded=$(ls -t | head -1)
+            if [ -n "$downloaded" ] && [ "$downloaded" != "$filename" ]; then
+                mv "$downloaded" "$filename"
+            fi
+        fi
+        log_success "Download do Mega concluído"
+        cd - >/dev/null
+        return 0
     else
-      wget -q "${wargs[@]}" --content-disposition -P "$dir" "$final_url" && return 0
+        log_error "Falha no download do Mega"
+        cd - >/dev/null
+        return 1
     fi
-  fi
-  # curl fallback
-  local cargs=("-L" "--retry" "$RETRIES" "--connect-timeout" "$CONNECT_TIMEOUT" "--max-time" "$DL_TIMEOUT" "-C" "-")
-  if [[ -n "$name" ]]; then
-    curl -sS "${cargs[@]}" -o "$dir/$name" "$final_url" && return 0
-  else
-    # tentar extrair nome do header
-    local detected=$(curl -sI -L "$final_url" | sed -n 's/.*filename="\([^"]\+\)".*/\1/p' | tr -d '\r' | head -1)
-    [[ -z "$detected" ]] && detected="download_$(date +%s)"
-    curl -sS "${cargs[@]}" -o "$dir/$detected" "$final_url" && return 0
-  fi
-  return 1
 }
 
-mega_download(){ # mega_url dir [filename]
-  local url="$1"; local dir="$2"; local name="${3:-}"
-  mkdir -p "$dir"
-  url="${url#mega://}"
-  if ! need megadl; then log_err "megadl não encontrado"; return 1; fi
-  log_info "megadl -> $url"
-  if timeout "$DL_TIMEOUT" megadl --path "$dir" "$url"; then
-    if [[ -n "$name" ]]; then
-      # renomeia para nome fixo, mantendo o arquivo único baixado
-      local got=$(find "$dir" -maxdepth 1 -type f -printf '%f\n' | tail -n 1)
-      [[ -n "$got" ]] && mv -f "$dir/$got" "$dir/$name"
+# 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
-    return 0
-  fi
-  return 1
 }
 
-hf_download(){ # hf://repo/path/to/file dir [filename]
-  local spec="$1"; local dir="$2"; local name="${3:-}"
-  mkdir -p "$dir"
-  spec="${spec#hf://}"
-  local repo="${spec%%/*/*}"
-  local path="${spec#*/ */}"
-  # Uso simples: se tiver a CLI, usa; senão cai no http_download se passar uma URL de 'resolve/main'
-  if "$VENV_DIR/bin/hf" --version >/dev/null 2>&1; then
-    "$VENV_DIR/bin/hf" download "$repo" "$path" --local-dir "$dir" --local-dir-use-symlinks False
-    if [[ -n "$name" ]]; then
-      local got=$(find "$dir" -type f \( -name "*.safetensors" -o -name "*.pth" -o -name "*.pt" -o -name "*.onnx" -o -name "*.bin" \) | head -1)
-      [[ -n "$got" ]] && mv -f "$got" "$dir/$name"
+# 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
-    return 0
-  fi
-  log_warn "hf CLI não disponível; tente fornecer URL de resolve/main"
-  return 1
+    
+    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
 }
 
-subdir_for(){
-  case "$1" in
-    checkpoints|diffusion_models|loras|vae|vae_approx|text_encoders|clip_vision|upscale_models) echo "$MODELS_DIR/$1" ;;
-    ultralytics/bbox) echo "$MODELS_DIR/ultralytics/bbox" ;;
-    *) echo "$MODELS_DIR/$1" ;;
-  esac
+# 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
+    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
+    
+    log_error "Falha ao baixar: $url"
+    return 1
 }
 
-process_downloads(){
-  log_step "[5/7] Baixando modelos (com retry/timeout)…"
-  mkdir -p "$MODELS_DIR"/{diffusion_models,loras,vae,vae_approx,text_encoders,clip_vision,upscale_models,checkpoints,ultralytics/bbox}
-  local IFS=$'\n'
-  local i=0; local ok=0; local fail=0; local skip=0
-  for line in $DOWNLOADS; do
-    [[ -z "$line" || "$line" =~ ^# ]] && continue
-    ((i++))
-    IFS='|' read -r url kind name <<<"$line"
-    url="$(echo "$url" | xargs)"; kind="$(echo "$kind" | xargs)"; name="$(echo "${name:-}" | xargs)"
-    local dst=$(subdir_for "$kind")
-    if was_done "$DONE_DL" "$line"; then log_ok "[$i] já baixado: ${name:-$(basename "$url")}"; ((skip++)); continue; fi
-    log_step "[$i] ${name:-$(basename "$url")} → $kind"
-    if [[ "$url" == mega://* ]]; then
-      if mega_download "$url" "$dst" "$name"; then mark_done "$DONE_DL" "$line"; ((ok++)); else ((fail++)); log_warn "Falha Mega, seguindo"; fi
-    elif [[ "$url" == hf://* ]]; then
-      if hf_download "$url" "$dst" "$name"; then mark_done "$DONE_DL" "$line"; ((ok++)); else ((fail++)); fi
+# Processar lista de downloads
+process_downloads() {
+    local total=${#DOWNLOAD_FILES[@]}
+    local current=0
+    local failed_count=0
+    
+    log_info "Total de downloads: $total"
+    
+    for entry in "${DOWNLOAD_FILES[@]}"; do
+        ((current++))
+        
+        # Verificar se já foi baixado
+        if is_download_complete "$entry"; then
+            log_success "[$current/$total] Já baixado, pulando..."
+            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" ;;
+            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)"
     else
-      if http_download "$url" "$dst" "$name"; then mark_done "$DONE_DL" "$line"; ((ok++)); else ((fail++)); fi
+        log_success "Todos os $total downloads completados!"
     fi
-  done
-  log_info "Downloads: novos=$ok, pulados=$skip, falhas=$fail"
-  save_state 5
-}
-
-# ====== Custom Nodes ======
-clone_or_update(){ # url dest
-  local url="$1"; local dest="$2"; local key="node:$url"
-  if was_done "$DONE_NODE" "$key" && [[ -d "$dest" ]]; then log_ok "Node OK: $(basename "$dest")"; return 0; fi
-  if [[ -d "$dest/.git" ]]; then
-    git -C "$dest" pull --ff-only || log_warn "pull falhou: $(basename "$dest")"
-  else
-    git clone --depth 1 "$url" "$dest" || { log_warn "clone falhou: $url"; return 0; }
-  fi
-  # requirements do node (best-effort)
-  if [[ -f "$dest/requirements.txt" ]]; then
-    "$PIP" install -q -r "$dest/requirements.txt" || log_warn "pip req falhou em $(basename "$dest")"
-  fi
-  mark_done "$DONE_NODE" "$key"
+    
+    return 0
 }
 
-install_nodes(){
-  log_step "[6/7] Instalando/Atualizando custom nodes…"
-  local CN_DIR="$COMFY_DIR/custom_nodes"
-  mkdir -p "$CN_DIR"
-  local total=${#CUSTOM_NODES[@]}
-  local idx=0
-  for repo in "${CUSTOM_NODES[@]}"; do
-    ((idx++))
-    log_step "[$idx/$total] $(basename "$repo")"
-    clone_or_update "$repo" "$CN_DIR/$(basename "$repo")"
-  done
-  log_ok "Nodes OK"
-  save_state 6
+# 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
+    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
+        }
+    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
+        }
+    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"
+        }
+    fi
+    
+    mark_node_complete "$url"
+    return 0
 }
 
-# ====== Verificação final & execução ======
-final_check(){
-  log_step "[7/7] Verificação final…"
-  [[ -f "$COMFY_DIR/main.py" ]] && log_ok "ComfyUI presente" || { log_err "ComfyUI ausente"; exit 1; }
-  # contagens
-  for d in checkpoints loras vae vae_approx text_encoders clip_vision upscale_models ultralytics; do
-    local p="$MODELS_DIR/$d"; [[ -d "$p" ]] || continue
-    local c=$(find "$p" -type f \( -name "*.safetensors" -o -name "*.pth" -o -name "*.pt" -o -name "*.ckpt" -o -name "*.onnx" \) 2>/dev/null | wc -l || echo 0)
-    (( c > 0 )) && log_ok "$d: $c arquivo(s)"
-  done
-  local gpu="$(nvidia-smi --query-gpu=name --format=csv,noheader 2>/dev/null || echo "Não detectada")"
-  log_info "GPU: $gpu"
-  save_state 7
-}
 
-launch(){
-  log_info "Iniciando ComfyUI em http://$COMFY_HOST:$COMFY_PORT"
-  cd "$COMFY_DIR"
-  exec "$PYTHON" main.py --listen "$COMFY_HOST" --port "$COMFY_PORT"
+# -----------------------------
+# 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
+    
+    return 0
 }
 
-# ====== Fluxo principal ======
-main(){
-  log_info "==============================="
-  log_info " ComfyUI Setup (Simplificado)"
-  log_info "==============================="
-
-  local st=$(get_state)
-  log_info "Step atual: $st"
-
-  (( st < 1 )) && ensure_base
-  (( $(get_state) < 2 )) && ensure_venv
-  (( $(get_state) < 3 )) && install_torch
-  (( $(get_state) < 4 )) && install_comfy
-  (( $(get_state) < 5 )) && process_downloads
-  (( $(get_state) < 6 )) && install_nodes
-  final_check
-
-  # Auto-start se não-interativo (pipe) ou se passar --start
-  if [[ ! -t 0 || ! -t 1 || "${1:-}" == "--start" ]]; then
-    launch
-  else
-    echo
+# -----------------------------
+# 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 - corrigindo PyTorch"
+            # Instalar torch, torchvision e torchaudio JUNTOS em um comando só
+            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}
+        
+        process_downloads
+        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"
+    else
+        log_error "✗ ComfyUI não encontrado"
+    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
+    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"
+    fi
+    
+    echo ""
     echo "Opções:"
     echo "  1) Iniciar ComfyUI agora"
-    echo "  2) Sair"
-    read -r -t 10 -p "Escolha [1-2] (padrão 1 em 10s): " ch || ch=1
-    case "$ch" in
-      1|"") launch;;
-      2) log_info "Saindo…";;
-      *) launch;;
+    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
-  fi
 }
 
-main "$@"
+# Executar instalação
+main "$@"