Update setup_comfyui_wan22.sh

setup_comfyui_wan22.sh

@@ -1,590 +1,411 @@
 #!/usr/bin/env bash
-# setup_comfyui_wan22_verbose.sh
-# Script com logs completos para instalação do ComfyUI com Wan 2.2
-# Requisitos: Linux + Python3 + git + (opcional) NVCC/CUDA para SageAttention
+# setup_comfyui_wan22_clean.sh
+# Script limpo e funcional para instalação do ComfyUI com Wan 2.2
 
-set -Eeuo pipefail
-trap 'rc=$?; echo -e "\033[31m[ERRO]\033[0m Linha $LINENO: \"$BASH_COMMAND\" saiu com código $rc"; exit $rc' ERR
+set -euo pipefail
 
 # -----------------------------
-# Cores para melhor visualização
+# Cores para output
 # -----------------------------
 RED='\033[0;31m'
 GREEN='\033[0;32m'
 YELLOW='\033[1;33m'
 BLUE='\033[0;34m'
-CYAN='\033[0;36m'
-MAGENTA='\033[0;35m'
-NC='\033[0m' # No Color
+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_debug() { echo -e "${CYAN}[DEBUG]${NC} $1"; }
-log_cmd() { echo -e "${MAGENTA}[CMD]${NC} $1"; }
 
 # -----------------------------
-# Configuráveis
+# Configuração
 # -----------------------------
-# IMPORTANTE: comfy-cli SEMPRE instala em /root/comfy/ComfyUI
-COMFY_DIR="/root/comfy/ComfyUI"  # Caminho fixo onde comfy-cli instala
+COMFY_DIR="/root/comfy/ComfyUI"
+MODELS_DIR="$COMFY_DIR/models"
 COMFY_HOST="${COMFY_HOST:-0.0.0.0}"
 COMFY_PORT="${COMFY_PORT:-8818}"
-
-# Repositórios no HF
-REPO_22="Comfy-Org/Wan_2.2_ComfyUI_Repackaged"
-REPO_21="Comfy-Org/Wan_2.1_ComfyUI_repackaged"
-REPO_CN="xinsir/controlnet-union-sdxl-1.0"
-
-# Token do Civitai
 CIVITAI_TOKEN="${CIVITAI_TOKEN:-4fcb2834969399006a736ee402b061e5}"
 
-# Configurações de performance
+# Configurações de performance (fixas, sem detecção complexa)
+export MAX_JOBS=16
+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  # Aumentado para melhor performance
+export HF_TRANSFER_CONCURRENCY=16
 
 # -----------------------------
-# Detectar capacidades do CUDA
+# Lista de downloads
 # -----------------------------
-detect_cuda_capabilities() {
-    if command -v nvidia-smi >/dev/null 2>&1; then
-        log_info "Detectando capacidades do CUDA..."
-        
-        # Obter informações da GPU
-        GPU_NAME=$(nvidia-smi --query-gpu=name --format=csv,noheader | head -n1)
-        GPU_MEMORY=$(nvidia-smi --query-gpu=memory.total --format=csv,noheader,nounits | head -n1)
-        GPU_COUNT=$(nvidia-smi --query-gpu=count --format=csv,noheader | wc -l)
-        
-        log_debug "GPU: $GPU_NAME ($GPU_MEMORY MB) x$GPU_COUNT"
-        
-        # Detectar compute capability
-        if command -v nvcc >/dev/null 2>&1; then
-            CUDA_VERSION=$(nvcc --version | grep "release" | awk '{print $6}' | cut -d',' -f1)
-            log_debug "CUDA Version: $CUDA_VERSION"
-            
-            # Obter compute capability via nvidia-ml-py se disponível
-            COMPUTE_CAP=$(python3 -c "
-try:
-    import pynvml
-    pynvml.nvmlInit()
-    handle = pynvml.nvmlDeviceGetHandleByIndex(0)
-    major, minor = pynvml.nvmlDeviceGetCudaComputeCapability(handle)
-    print(f'{major}.{minor}')
-except:
-    print('unknown')
-" 2>/dev/null || echo "unknown")
-            
-            if [ "$COMPUTE_CAP" != "unknown" ]; then
-                log_debug "Compute Capability: $COMPUTE_CAP"
-            fi
-            
-            # Determinar MAX_JOBS baseado na GPU e memória
-            # RTX 4090/A100: 128 SMs -> pode usar 32 jobs
-            # RTX 3090: 82 SMs -> pode usar 24 jobs
-            # RTX 3080: 68 SMs -> pode usar 16 jobs
-            # Regra geral: ~4 jobs por GB de VRAM, máximo 32
-            
-            if [ "$GPU_MEMORY" -gt 0 ]; then
-                SUGGESTED_JOBS=$((GPU_MEMORY / 1024 / 4))  # 4 jobs por GB
-                if [ $SUGGESTED_JOBS -gt 32 ]; then
-                    SUGGESTED_JOBS=32  # Limite do CUDA
-                elif [ $SUGGESTED_JOBS -lt 4 ]; then
-                    SUGGESTED_JOBS=4   # Mínimo razoável
-                fi
-                
-                # Para GPUs conhecidas, usar valores otimizados
-                if [[ "$GPU_NAME" == *"4090"* ]] || [[ "$GPU_NAME" == *"A100"* ]]; then
-                    SUGGESTED_JOBS=32
-                    SUGGESTED_THREADS=16
-                elif [[ "$GPU_NAME" == *"3090"* ]] || [[ "$GPU_NAME" == *"A6000"* ]]; then
-                    SUGGESTED_JOBS=24
-                    SUGGESTED_THREADS=12
-                elif [[ "$GPU_NAME" == *"3080"* ]] || [[ "$GPU_NAME" == *"A5000"* ]]; then
-                    SUGGESTED_JOBS=16
-                    SUGGESTED_THREADS=8
-                else
-                    SUGGESTED_THREADS=$((SUGGESTED_JOBS / 2))
-                fi
-                
-                export MAX_JOBS=$SUGGESTED_JOBS
-                export NVCC_APPEND_FLAGS="--threads $SUGGESTED_THREADS"
-                export EXT_PARALLEL=$((SUGGESTED_JOBS / 2))
-                
-                log_success "Configuração CUDA otimizada:"
-                log_debug "  MAX_JOBS=$MAX_JOBS"
-                log_debug "  NVCC threads=$SUGGESTED_THREADS"
-                log_debug "  EXT_PARALLEL=$EXT_PARALLEL"
-            fi
-        fi
-    else
-        log_warn "nvidia-smi não encontrado, usando configuração padrão"
-    fi
-}
-
-# Detecção de CPU cores para compilação
-CPU_CORES=$(nproc 2>/dev/null || sysctl -n hw.ncpu 2>/dev/null || echo 4)
-log_debug "CPU Cores detectados: $CPU_CORES"
-
-# Configuração inicial baseada em CPU (será sobrescrita pela detecção CUDA se disponível)
-if [ $CPU_CORES -gt 32 ]; then
-    export MAX_JOBS=${MAX_JOBS:-32}
-    export EXT_PARALLEL=${EXT_PARALLEL:-16}
-    export NVCC_APPEND_FLAGS=${NVCC_APPEND_FLAGS:-"--threads 16"}
-    COMPILE_JOBS=32
-else
-    HALF_CORES=$((CPU_CORES / 2))
-    export MAX_JOBS=${MAX_JOBS:-$HALF_CORES}
-    export EXT_PARALLEL=${EXT_PARALLEL:-4}
-    export NVCC_APPEND_FLAGS=${NVCC_APPEND_FLAGS:-"--threads 8"}
-    COMPILE_JOBS=$HALF_CORES
-fi
-
-# Detectar e otimizar para CUDA
-detect_cuda_capabilities
+# Formato: "URL|TIPO|NOME_ARQUIVO_OPCIONAL"
+# TIPO pode ser: checkpoints, loras, vae, text_encoders, clip_vision, controlnet, upscale_models
+readonly DOWNLOAD_FILES=(
+    # Wan 2.2 models
+    "hf://Comfy-Org/Wan_2.2_ComfyUI_Repackaged/split_files/diffusion_models/wan2.2_fun_control_low_noise_14B_fp8_scaled.safetensors|checkpoints|"
+    "hf://Comfy-Org/Wan_2.2_ComfyUI_Repackaged/split_files/diffusion_models/wan2.2_fun_control_high_noise_14B_fp8_scaled.safetensors|checkpoints|"
+    "hf://Comfy-Org/Wan_2.2_ComfyUI_Repackaged/split_files/loras/wan2.2_i2v_lightx2v_4steps_lora_v1_high_noise.safetensors|loras|"
+    "hf://Comfy-Org/Wan_2.2_ComfyUI_Repackaged/split_files/loras/wan2.2_i2v_lightx2v_4steps_lora_v1_low_noise.safetensors|loras|"
+    "hf://Comfy-Org/Wan_2.2_ComfyUI_Repackaged/split_files/vae/wan_2.1_vae.safetensors|vae|"
+    "hf://Comfy-Org/Wan_2.2_ComfyUI_Repackaged/split_files/text_encoders/umt5_xxl_fp16.safetensors|text_encoders|"
+    
+    # Wan 2.1 clip vision
+    "hf://Comfy-Org/Wan_2.1_ComfyUI_repackaged/split_files/clip_vision/clip_vision_h.safetensors|clip_vision|"
+    
+    # ControlNet
+    "hf://xinsir/controlnet-union-sdxl-1.0/diffusion_pytorch_model_promax.safetensors|controlnet|controlnet-union.safetensors"
+    
+    # Civitai model
+    "https://civitai.com/api/download/models/2122278?type=Model&format=SafeTensor&size=pruned&fp=fp16|loras|"
+)
 
-export UV_SYSTEM_PYTHON=1
-export PYTORCH_CUDA_ALLOC_CONF="expandable_segments:True"
+# Custom nodes para instalar
+readonly CUSTOM_NODES=(
+    "https://github.com/kijai/ComfyUI-Florence2"
+    "https://github.com/kijai/ComfyUI-WanVideoWrapper"
+    "https://github.com/Fannovel16/ComfyUI-Frame-Interpolation"
+    "https://github.com/kijai/ComfyUI-GIMM-VFI"
+    "https://github.com/Kosinkadink/ComfyUI-VideoHelperSuite"
+    "https://github.com/kijai/ComfyUI-KJNodes"
+    "https://github.com/Fannovel16/comfyui_controlnet_aux"
+    "https://github.com/Artificial-Sweetener/comfyui-WhiteRabbit"
+    "https://github.com/LucipherDev/ComfyUI-AniDoc"
+    "https://github.com/shiimizu/ComfyUI_smZNodes"
+    "https://github.com/yuvraj108c/ComfyUI-Dwpose-Tensorrt"
+    "https://github.com/grmchn/ComfyUI-ProportionChanger"
+)
 
 # -----------------------------
-# Setup inicial
+# Funções auxiliares
 # -----------------------------
-TMPDL="$(mktemp -d)"
-cleanup() { 
-    log_info "Limpando arquivos temporários..."
-    rm -rf "$TMPDL" 
+command_exists() {
+    command -v "$1" >/dev/null 2>&1
 }
-trap cleanup EXIT
 
-# Verificação de dependências
-need() { 
-    if ! command -v "$1" >/dev/null 2>&1; then
-        log_error "Falta dependência: $1"
-        log_info "Instale com: apt-get install $1"
-        exit 1
+# Adicionar token do Civitai se necessário
+add_civitai_token() {
+    local url="$1"
+    if [[ "$url" == *"civitai.com/api/download"* ]] && [[ "$url" != *"token="* ]]; then
+        echo "${url}&token=${CIVITAI_TOKEN}"
+    else
+        echo "$url"
     fi
 }
 
-log_info "Verificando dependências..."
-need python3
-need git
-need wget
-need curl
-
-# -----------------------------
-# Funções auxiliares
-# -----------------------------
-clone_or_update() {
-    local url="$1" dest="$2"
-    local repo_name=$(basename "$dest")
+# Download HuggingFace usando HF Transfer ou huggingface-cli
+download_hf() {
+    local repo="$1"
+    local file_path="$2"
+    local target_dir="$3"
+    local filename="$4"
     
-    if [ -d "$dest/.git" ]; then
-        log_info "Atualizando $repo_name..."
-        log_cmd "git -C \"$dest\" pull --ff-only"
-        git -C "$dest" pull --ff-only || true
-    else
-        log_info "Clonando $repo_name..."
-        log_cmd "git clone --recursive --depth 1 \"$url\" \"$dest\""
-        git clone --recursive --depth 1 "$url" "$dest"
+    # Extrair nome do arquivo se não fornecido
+    if [ -z "$filename" ]; then
+        filename=$(basename "$file_path")
     fi
     
-    if [ -f "$dest/requirements.txt" ]; then
-        log_debug "Instalando requirements para $repo_name..."
-        python3 -m pip install --no-warn-script-location -r "$dest/requirements.txt"
+    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
-}
-
-# Função para download com HF Transfer ou fallback
-hf_download_with_fallback() {
-    local repo="$1"
-    shift  # Remove primeiro argumento
-    local files=("$@")  # Resto são os arquivos
-    local output_dir="$TMPDL/$(basename $repo)"
     
-    # Tentar primeiro com HF Transfer
+    log_info "Baixando de HF: $filename"
+    
+    # Tentar com HF Transfer primeiro
     if python3 -c "import hf_transfer" 2>/dev/null; then
-        log_info "  → Usando HF Transfer (rápido) para $repo..."
-        log_cmd "HF_HUB_ENABLE_HF_TRANSFER=1 huggingface-cli download \"$repo\" ..."
-        
-        HF_HUB_ENABLE_HF_TRANSFER=1 huggingface-cli download "$repo" \
-            "${files[@]}" \
-            --local-dir "$output_dir" \
+        HF_HUB_ENABLE_HF_TRANSFER=1 huggingface-cli download "$repo" "$file_path" \
+            --local-dir "$target_dir" \
             --local-dir-use-symlinks False \
-            --resume-download || {
-            log_warn "  HF Transfer falhou, tentando método padrão..."
-            HF_HUB_ENABLE_HF_TRANSFER=0 huggingface-cli download "$repo" \
-                "${files[@]}" \
-                --local-dir "$output_dir" \
-                --local-dir-use-symlinks False \
-                --resume-download
+            --resume-download 2>/dev/null && {
+            # Renomear se necessário
+            if [ "$filename" != "$(basename "$file_path")" ]; then
+                mv "$target_dir/$file_path" "$target_file" 2>/dev/null || true
+            fi
+            return 0
         }
-    else
-        log_info "  → Usando download padrão para $repo..."
-        huggingface-cli download "$repo" \
-            "${files[@]}" \
-            --local-dir "$output_dir" \
-            --local-dir-use-symlinks False \
-            --resume-download
     fi
-}
-
-# -----------------------------
-# Instalação principal
-# -----------------------------
-echo ""
-log_info "========================================="
-log_info "  ComfyUI + Wan 2.2 Setup (Verbose)"
-log_info "  CPU: $CPU_CORES cores | Compile Jobs: $COMPILE_JOBS"
-log_info "  CUDA MAX_JOBS: ${MAX_JOBS:-não detectado}"
-log_info "  ComfyUI Dir: $COMFY_DIR"
-log_info "========================================="
-echo ""
-
-# [1/8] Atualizar pip e instalar ferramentas essenciais
-log_info "[1/8] Atualizando pip e instalando ferramentas..."
-log_cmd "python3 -m pip install --upgrade pip wheel setuptools"
-python3 -m pip install --upgrade pip wheel setuptools
-
-# Instalar HF Transfer primeiro para downloads mais rápidos
-log_cmd "python3 -m pip install --upgrade hf_transfer"
-python3 -m pip install --upgrade hf_transfer || true
-
-# Depois instalar o resto
-log_cmd "python3 -m pip install --upgrade comfy-cli huggingface_hub>=0.23.0 aria2"
-python3 -m pip install --upgrade comfy-cli "huggingface_hub>=0.23.0" aria2
-log_success "Ferramentas instaladas"
-
-# [2/8] Instalar ComfyUI
-log_info "[2/8] Instalando ComfyUI (GPU NVIDIA, fast deps)..."
-
-# Verificar se ComfyUI já existe no caminho correto
-if [ -f "$COMFY_DIR/main.py" ] && [ -d "$COMFY_DIR/models" ]; then
-    log_warn "ComfyUI já existe em $COMFY_DIR, pulando instalação base"
-else
-    log_info "Instalando ComfyUI (será instalado em /root/comfy/ComfyUI)..."
     
-    # Remove instalação parcial se existir
-    [ -d "$COMFY_DIR" ] && rm -rf "$COMFY_DIR"
-    
-    # Instala ComfyUI (sempre vai para /root/comfy/ComfyUI)
-    log_cmd "comfy --skip-prompt install --fast-deps --nvidia"
-    comfy --skip-prompt install --fast-deps --nvidia
-    
-    # Verifica se instalou no local correto
-    if [ ! -f "$COMFY_DIR/main.py" ]; then
-        log_error "ComfyUI não foi encontrado em $COMFY_DIR após instalação!"
-        log_info "Verificando outros locais possíveis..."
-        
-        # Tentar encontrar onde foi instalado
-        for possible_path in "/root/comfy/ComfyUI" "$HOME/comfy/ComfyUI" "/workspace/ComfyUI"; do
-            if [ -f "$possible_path/main.py" ]; then
-                log_info "ComfyUI encontrado em: $possible_path"
-                COMFY_DIR="$possible_path"
-                break
-            fi
-        done
-        
-        if [ ! -f "$COMFY_DIR/main.py" ]; then
-            log_error "Não foi possível encontrar ComfyUI instalado!"
-            exit 1
+    # Fallback para método padrão
+    huggingface-cli download "$repo" "$file_path" \
+        --local-dir "$target_dir" \
+        --local-dir-use-symlinks False \
+        --resume-download && {
+        # Renomear se necessário
+        if [ "$filename" != "$(basename "$file_path")" ]; then
+            mv "$target_dir/$file_path" "$target_file" 2>/dev/null || true
         fi
-    fi
-fi
-
-# Criar estrutura de diretórios
-MODELS_DIR="$COMFY_DIR/models"
-log_info "Criando estrutura de diretórios em $MODELS_DIR..."
-mkdir -p "$MODELS_DIR"/{diffusion_models,loras,vae,text_encoders,clip_vision,controlnet}
-mkdir -p "$COMFY_DIR/custom_nodes"
-mkdir -p "$TMPDL"/{wan22,wan21,controlnet}
-
-log_success "ComfyUI pronto em: $COMFY_DIR"
-
-# [3/8] Baixar modelos Wan 2.2
-log_info "[3/8] Iniciando download dos modelos Wan 2.2..."
-(
-    set -Eeuo pipefail
-    
-    hf_download_with_fallback "$REPO_22" \
-        "split_files/diffusion_models/wan2.2_fun_control_low_noise_14B_fp8_scaled.safetensors" \
-        "split_files/diffusion_models/wan2.2_fun_control_high_noise_14B_fp8_scaled.safetensors" \
-        "split_files/loras/wan2.2_i2v_lightx2v_4steps_lora_v1_high_noise.safetensors" \
-        "split_files/loras/wan2.2_i2v_lightx2v_4steps_lora_v1_low_noise.safetensors" \
-        "split_files/vae/wan_2.1_vae.safetensors" \
-        "split_files/text_encoders/umt5_xxl_fp16.safetensors"
-    
-    # Copiar para destino
-    log_debug "Copiando modelos Wan 2.2 para $MODELS_DIR..."
-    cp -fv "$TMPDL/wan22/split_files/diffusion_models/"*.safetensors "$MODELS_DIR/diffusion_models/" 2>/dev/null || true
-    cp -fv "$TMPDL/wan22/split_files/loras/"*.safetensors "$MODELS_DIR/loras/" 2>/dev/null || true
-    cp -fv "$TMPDL/wan22/split_files/vae/"*.safetensors "$MODELS_DIR/vae/" 2>/dev/null || true
-    cp -fv "$TMPDL/wan22/split_files/text_encoders/"*.safetensors "$MODELS_DIR/text_encoders/" 2>/dev/null || true
+        return 0
+    }
     
-    log_success "  Modelos Wan 2.2 baixados"
-) &
-WAN22_PID=$!
+    return 1
+}
 
-# [4/8] Baixar Wan 2.1 clip vision
-log_info "[4/8] Baixando Wan 2.1 clip vision..."
-(
-    set -Eeuo pipefail
-    
-    hf_download_with_fallback "$REPO_21" \
-        "split_files/clip_vision/clip_vision_h.safetensors"
+# Download genérico com aria2c, wget ou curl
+download_file() {
+    local url="$1"
+    local target_dir="$2"
+    local filename="$3"
     
-    log_debug "Copiando clip vision para $MODELS_DIR/clip_vision..."
-    cp -fv "$TMPDL/wan21/split_files/clip_vision/"*.safetensors "$MODELS_DIR/clip_vision/" 2>/dev/null || true
+    # Adicionar token Civitai se necessário
+    url=$(add_civitai_token "$url")
     
-    log_success "  Clip vision baixado"
-) &
-WAN21_PID=$!
-
-# [5/8] Baixar ControlNet
-log_info "[5/8] Baixando ControlNet Union..."
-(
-    set -Eeuo pipefail
+    # 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
+    fi
     
-    hf_download_with_fallback "$REPO_CN" \
-        "diffusion_pytorch_model_promax.safetensors"
+    log_info "Baixando: ${filename:-$(basename "$url" | cut -d'?' -f1)}"
     
-    log_debug "Copiando ControlNet para $MODELS_DIR/controlnet..."
-    # O arquivo pode estar direto no diretório ou em subdiretório
-    if [ -f "$TMPDL/controlnet/diffusion_pytorch_model_promax.safetensors" ]; then
-        cp -fv "$TMPDL/controlnet/diffusion_pytorch_model_promax.safetensors" \
-            "$MODELS_DIR/controlnet/controlnet-union.safetensors"
-    elif [ -f "$TMPDL/xinsir/controlnet-union-sdxl-1.0/diffusion_pytorch_model_promax.safetensors" ]; then
-        cp -fv "$TMPDL/xinsir/controlnet-union-sdxl-1.0/diffusion_pytorch_model_promax.safetensors" \
-            "$MODELS_DIR/controlnet/controlnet-union.safetensors"
-    else
-        # Procurar o arquivo em qualquer lugar do diretório temporário
-        FOUND_FILE=$(find "$TMPDL" -name "diffusion_pytorch_model_promax.safetensors" -type f | head -n1)
-        if [ -n "$FOUND_FILE" ]; then
-            log_debug "ControlNet encontrado em: $FOUND_FILE"
-            cp -fv "$FOUND_FILE" "$MODELS_DIR/controlnet/controlnet-union.safetensors"
+    # 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
+            aria2c $aria_opts --dir="$target_dir" --out="$filename" "$url" && return 0
         else
-            log_warn "ControlNet não encontrado no diretório temporário"
+            aria2c $aria_opts --dir="$target_dir" "$url" && return 0
         fi
     fi
     
-    log_success "  ControlNet baixado"
-) &
-CN_PID=$!
-
-# [6/8] Baixar modelo do Civitai
-log_info "[6/8] Baixando modelo do Civitai..."
-(
-    set -Eeuo pipefail
-    
-    cd "$MODELS_DIR/loras"
+    # Fallback para wget
+    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
     
-    # URL correta completa
-    CIVITAI_URL="https://civitai.com/api/download/models/2122278?type=Model&format=SafeTensor&size=pruned&fp=fp16&token=${CIVITAI_TOKEN}"
+    # Último recurso: curl
+    if command_exists curl; then
+        if [ -n "$filename" ]; then
+            curl -L -# -C - -o "$target_dir/$filename" "$url" && return 0
+        else
+            # Tentar obter nome do arquivo do header
+            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_debug "Baixando de: $CIVITAI_URL"
+    log_error "Falha ao baixar: $url"
+    return 1
+}
+
+# Processar lista de downloads
+process_downloads() {
+    local failed_count=0
     
-    # Usar curl direto pois wget está falhando com nomes longos
-    # Primeiro obter o nome real do arquivo
-    log_debug "Obtendo nome do arquivo..."
-    HEADERS=$(curl -sI -L --max-time 30 "$CIVITAI_URL")
-    FILENAME=$(echo "$HEADERS" | grep -i "content-disposition" | sed -n 's/.*filename="\([^"]*\)".*/\1/p' | tr -d '\r')
+    for entry in "${DOWNLOAD_FILES[@]}"; do
+        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/diffusion_models" ;;
+            loras) target_dir="$MODELS_DIR/loras" ;;
+            vae) target_dir="$MODELS_DIR/vae" ;;
+            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" ;;
+            *) target_dir="$MODELS_DIR/$type" ;;
+        esac
+        
+        # Criar diretório se não existir
+        mkdir -p "$target_dir"
+        
+        # Processar URL
+        if [[ "$url" == hf://* ]]; then
+            # HuggingFace download
+            url="${url#hf://}"  # Remove prefixo hf://
+            
+            # Separar repo e caminho
+            repo=$(echo "$url" | cut -d'/' -f1-2)
+            file_path=$(echo "$url" | cut -d'/' -f3-)
+            
+            download_hf "$repo" "$file_path" "$target_dir" "$filename" || ((failed_count++))
+        else
+            # Download normal (HTTP/HTTPS)
+            download_file "$url" "$target_dir" "$filename" || ((failed_count++))
+        fi
+    done
     
-    if [ -z "$FILENAME" ]; then
-        # Tentar outro padrão de header
-        FILENAME=$(echo "$HEADERS" | grep -i "content-disposition" | sed -n "s/.*filename='\([^']*\)'.*/\1/p" | tr -d '\r')
-    fi
+    return $failed_count
+}
+
+# Clonar ou atualizar repositório git
+clone_or_update() {
+    local url="$1"
+    local dest="$2"
     
-    if [ -z "$FILENAME" ]; then
-        # Se ainda não conseguiu, usar nome padrão baseado no ID do modelo
-        FILENAME="civitai_model_2122278.safetensors"
-        log_warn "Não foi possível detectar nome do arquivo, usando: $FILENAME"
+    if [ -d "$dest/.git" ]; then
+        log_info "Atualizando $(basename "$dest")..."
+        git -C "$dest" pull --ff-only 2>/dev/null || true
     else
-        log_success "Nome do arquivo detectado: $FILENAME"
+        log_info "Clonando $(basename "$dest")..."
+        git clone --recursive --depth 1 "$url" "$dest" 2>/dev/null || return 1
     fi
     
-    # Baixar o arquivo
-    log_cmd "curl -L --progress-bar \"$CIVITAI_URL\" -o \"$FILENAME\""
+    # Instalar requirements se existir
+    if [ -f "$dest/requirements.txt" ]; then
+        python3 -m pip install --no-warn-script-location -q -r "$dest/requirements.txt" 2>/dev/null || true
+    fi
+}
+
+# -----------------------------
+# Instalação principal
+# -----------------------------
+echo ""
+log_info "========================================="
+log_info "  ComfyUI + Wan 2.2 Setup (Clean)"
+log_info "========================================="
+echo ""
+
+# [1] Verificar dependências
+log_info "[1/6] Verificando dependências..."
+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"
+
+# [2] Atualizar pip e instalar ferramentas
+log_info "[2/6] Preparando ferramentas Python..."
+python3 -m pip install --upgrade pip wheel setuptools -q
+python3 -m pip install --upgrade comfy-cli "huggingface_hub>=0.23.0" hf_transfer aria2 -q
+log_success "Ferramentas instaladas"
+
+# [3] Instalar ComfyUI
+log_info "[3/6] 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
     
-    if curl -L \
-         --retry 3 \
-         --connect-timeout 60 \
-         --max-time 1800 \
-         --progress-bar \
-         -o "$FILENAME" \
-         "$CIVITAI_URL"; then
-        log_success "  Modelo Civitai baixado: $FILENAME"
-    else
-        log_error "  Falha ao baixar modelo Civitai"
+    # Verificar instalação
+    if [ ! -f "$COMFY_DIR/main.py" ]; then
+        log_error "ComfyUI não foi instalado corretamente!"
+        exit 1
     fi
-) &
-CIVITAI_PID=$!
+fi
+log_success "ComfyUI instalado em: $COMFY_DIR"
 
-# [7/8] Instalar SageAttention (em paralelo com downloads)
-log_info "[7/8] Preparando SageAttention (paralelo com downloads)..."
+# [4] Instalar SageAttention (em background)
+log_info "[4/6] Instalando SageAttention..."
 (
-    set -Eeuo pipefail
-    
-    # Verificar se CUDA está disponível
-    if command -v nvcc >/dev/null 2>&1; then
-        log_info "  CUDA detectado, instalando triton e preparando SageAttention..."
-        log_info "  Usando configuração otimizada: MAX_JOBS=$MAX_JOBS, EXT_PARALLEL=$EXT_PARALLEL"
-        
-        # Instalar triton primeiro
-        log_cmd "python3 -m pip install \"triton>=3.0.0\""
-        python3 -m pip install "triton>=3.0.0"
+    if command_exists nvcc; then
+        python3 -m pip install "triton>=3.0.0" nvidia-ml-py -q
         
-        # Instalar pynvml para melhor detecção de GPU (ignorar warning de deprecation)
-        log_cmd "python3 -m pip install nvidia-ml-py"
-        python3 -m pip install nvidia-ml-py 2>&1 | grep -v "FutureWarning" || true
-        
-        # Clonar SageAttention
-        SAGE_DIR="$TMPDL/SageAttention"
-        if [ ! -d "$SAGE_DIR" ]; then
-            log_cmd "git clone --depth 1 https://github.com/thu-ml/SageAttention.git \"$SAGE_DIR\""
-            git clone --depth 1 https://github.com/thu-ml/SageAttention.git "$SAGE_DIR"
-        fi
+        SAGE_DIR="/tmp/SageAttention"
+        rm -rf "$SAGE_DIR"
+        git clone --depth 1 https://github.com/thu-ml/SageAttention.git "$SAGE_DIR" 2>/dev/null
         
         cd "$SAGE_DIR"
+        python3 setup.py install --quiet 2>/dev/null
         
-        # Compilar com configurações otimizadas
-        log_info "  Compilando SageAttention com MAX_JOBS=$MAX_JOBS..."
-        log_debug "  NVCC flags: $NVCC_APPEND_FLAGS"
-        
-        # Mostrar output completo da compilação
-        log_cmd "python3 setup.py install"
-        python3 setup.py install 2>&1 | grep -v "SetuptoolsDeprecationWarning" | grep -v "FutureWarning" || true
-        
-        log_success "  SageAttention instalado"
+        log_success "SageAttention instalado"
     else
-        log_warn "  CUDA não detectado, pulando SageAttention"
+        log_warn "CUDA não detectado, pulando SageAttention"
     fi
 ) &
 SAGE_PID=$!
 
-# Aguardar downloads principais
-log_info "Aguardando downloads principais (em paralelo com SageAttention)..."
-wait $WAN22_PID $WAN21_PID $CN_PID $CIVITAI_PID
-
-# [8/8] Instalar custom nodes (após downloads, mas ainda em paralelo com SageAttention)
-log_info "[8/8] Instalando custom nodes..."
-(
-    set -Eeuo pipefail
-    CN_DIR="$COMFY_DIR/custom_nodes"
-    
-    # Verificar se o diretório existe
-    if [ ! -d "$CN_DIR" ]; then
-        log_error "Diretório custom_nodes não existe! Criando..."
-        mkdir -p "$CN_DIR"
-    fi
-    
-    declare -A REPOS=(
-        ["https://github.com/kijai/ComfyUI-Florence2"]="ComfyUI-Florence2"
-        ["https://github.com/kijai/ComfyUI-WanVideoWrapper"]="ComfyUI-WanVideoWrapper"
-        ["https://github.com/Fannovel16/ComfyUI-Frame-Interpolation"]="ComfyUI-Frame-Interpolation"
-        ["https://github.com/kijai/ComfyUI-GIMM-VFI"]="ComfyUI-GIMM-VFI"
-        ["https://github.com/Kosinkadink/ComfyUI-VideoHelperSuite"]="ComfyUI-VideoHelperSuite"
-        ["https://github.com/kijai/ComfyUI-KJNodes"]="ComfyUI-KJNodes"
-        ["https://github.com/Fannovel16/comfyui_controlnet_aux"]="comfyui_controlnet_aux"
-        ["https://github.com/Artificial-Sweetener/comfyui-WhiteRabbit"]="comfyui-WhiteRabbit"
-        ["https://github.com/LucipherDev/ComfyUI-AniDoc"]="ComfyUI-AniDoc"
-        ["https://github.com/shiimizu/ComfyUI_smZNodes"]="ComfyUI_smZNodes"
-        ["https://github.com/yuvraj108c/ComfyUI-Dwpose-Tensorrt"]="ComfyUI-Dwpose-Tensorrt"
-        ["https://github.com/grmchn/ComfyUI-ProportionChanger"]="ComfyUI-ProportionChanger"
-    )
-    
-    # Clonar todos em paralelo
-    for url in "${!REPOS[@]}"; do
-        node_name="${REPOS[$url]}"
-        node_path="$CN_DIR/$node_name"
-        clone_or_update "$url" "$node_path" &
-    done
-    wait
-    
-    log_success "  Custom nodes instalados"
-)
-
-# Aguardar SageAttention terminar
-log_info "Aguardando conclusão da instalação do SageAttention..."
-wait $SAGE_PID
+# [5] Baixar modelos
+log_info "[5/6] Baixando modelos..."
+mkdir -p "$MODELS_DIR"/{diffusion_models,loras,vae,text_encoders,clip_vision,controlnet,upscale_models}
 
-# Verificar instalação
-log_info "Verificando instalação..."
-ERROR_COUNT=0
+# Verificar se aria2c está disponível
+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)"
+    log_info "Instale aria2c para downloads mais rápidos: apt-get install aria2"
+fi
 
-check_file() {
-    if [ -f "$1" ]; then
-        SIZE=$(du -h "$1" | cut -f1)
-        log_success "  ✓ $(basename "$1") ($SIZE)"
-        return 0
-    else
-        log_error "  ✗ $(basename "$1") não encontrado"
-        ((ERROR_COUNT++))
-        return 1
-    fi
-}
+# Processar downloads
+if process_downloads; then
+    log_success "Todos os modelos baixados com sucesso"
+else
+    log_warn "Alguns downloads falharam, mas continuando..."
+fi
 
-# Verificar ComfyUI
-log_debug "Verificando ComfyUI..."
-check_file "$COMFY_DIR/main.py"
+# [6] Instalar custom nodes
+log_info "[6/6] Instalando custom nodes..."
+CN_DIR="$COMFY_DIR/custom_nodes"
+mkdir -p "$CN_DIR"
 
-# Verificar modelos principais
-log_info "Verificando modelos baixados..."
-check_file "$MODELS_DIR/diffusion_models/wan2.2_fun_control_low_noise_14B_fp8_scaled.safetensors"
-check_file "$MODELS_DIR/diffusion_models/wan2.2_fun_control_high_noise_14B_fp8_scaled.safetensors"
-check_file "$MODELS_DIR/vae/wan_2.1_vae.safetensors"
-check_file "$MODELS_DIR/text_encoders/umt5_xxl_fp16.safetensors"
-check_file "$MODELS_DIR/clip_vision/clip_vision_h.safetensors"
-check_file "$MODELS_DIR/controlnet/controlnet-union.safetensors"
+for repo_url in "${CUSTOM_NODES[@]}"; do
+    node_name=$(basename "$repo_url")
+    clone_or_update "$repo_url" "$CN_DIR/$node_name" &
+done
+wait
 
-# Verificar se algum modelo Civitai foi baixado
-log_debug "Verificando modelos Civitai em $MODELS_DIR/loras..."
-CIVITAI_FILES=$(ls -1 "$MODELS_DIR/loras/"*.safetensors 2>/dev/null || true)
-CIVITAI_COUNT=$(echo "$CIVITAI_FILES" | grep -c safetensors || echo 0)
+log_success "Custom nodes instalados"
 
-if [ "$CIVITAI_COUNT" -gt 0 ]; then
-    log_success "  ✓ $CIVITAI_COUNT modelo(s) Civitai encontrado(s):"
-    echo "$CIVITAI_FILES" | while read file; do
-        [ -n "$file" ] && log_debug "    - $(basename "$file") ($(du -h "$file" | cut -f1))"
-    done
-else
-    log_warn "  ⚠ Nenhum modelo Civitai encontrado (não crítico)"
-fi
+# Aguardar SageAttention
+wait $SAGE_PID 2>/dev/null || true
 
-# Relatório final
+# Verificação final
 echo ""
 log_info "========================================="
-if [ $ERROR_COUNT -eq 0 ]; then
-    log_success "Instalação concluída com sucesso!"
-    log_info "ComfyUI instalado em: $COMFY_DIR"
+log_info "Verificando instalação..."
+
+# Verificar arquivos principais
+check_critical_files() {
+    local files=(
+        "$COMFY_DIR/main.py"
+        "$MODELS_DIR/diffusion_models/wan2.2_fun_control_low_noise_14B_fp8_scaled.safetensors"
+        "$MODELS_DIR/diffusion_models/wan2.2_fun_control_high_noise_14B_fp8_scaled.safetensors"
+        "$MODELS_DIR/vae/wan_2.1_vae.safetensors"
+        "$MODELS_DIR/text_encoders/umt5_xxl_fp16.safetensors"
+    )
     
-    # Sempre usar flag do SageAttention (ComfyUI ignora se não estiver instalado)
-    SAGE_FLAG="--use-sage-attention"
-    log_info "Flag SageAttention habilitada (será usada se disponível)"
+    local missing=0
+    for file in "${files[@]}"; do
+        if [ -f "$file" ]; then
+            log_success "✓ $(basename "$file")"
+        else
+            log_error "✗ $(basename "$file") não encontrado"
+            ((missing++))
+        fi
+    done
     
-    # Mostrar informações do sistema
-    log_debug "Configuração final:"
-    log_debug "  CPU: $CPU_CORES cores"
-    if command -v nvidia-smi >/dev/null 2>&1; then
-        log_debug "  GPU: $(nvidia-smi --query-gpu=name --format=csv,noheader | head -n1)"
-        log_debug "  VRAM: $(nvidia-smi --query-gpu=memory.total --format=csv,noheader | head -n1)"
-    fi
+    return $missing
+}
+
+if check_critical_files; then
+    log_success "Instalação concluída com sucesso!"
     
     echo ""
     log_info "Iniciando ComfyUI..."
+    log_info "URL: http://localhost:$COMFY_PORT"
     
-    # Tentar obter IP local
-    LOCAL_IP=$(hostname -I 2>/dev/null | awk '{print $1}')
-    if [ -n "$LOCAL_IP" ]; then
-        log_info "URL Local: http://$LOCAL_IP:$COMFY_PORT"
+    # Tentar usar SageAttention se disponível
+    SAGE_FLAG=""
+    if python3 -c "import sageattention" 2>/dev/null; then
+        SAGE_FLAG="--use-sage-attention"
+        log_info "SageAttention habilitado"
     fi
-    log_info "URL: http://localhost:$COMFY_PORT"
-    echo ""
     
-    # Mudar para o diretório do ComfyUI e executar
     cd "$COMFY_DIR"
-    log_cmd "comfy launch -- $SAGE_FLAG --listen \"$COMFY_HOST\" --port \"$COMFY_PORT\""
     exec comfy launch -- $SAGE_FLAG --listen "$COMFY_HOST" --port "$COMFY_PORT"
 else
-    log_error "Instalação completada com $ERROR_COUNT erro(s)"
-    log_info "Verifique os erros acima e execute o script novamente"
+    log_error "Arquivos críticos faltando. Verifique os erros acima."
     exit 1
 fi