Update wanForProduction.sh

wanForProduction.sh

@@ -1,10 +1,12 @@
 #!/usr/bin/env bash
 # setup_comfyui_combined.sh
-# Script combinado para instalação do ComfyUI com Wan 2.2
+# Script baseado no CODE1 com melhorias e modelos/nodes do CODE2
 
 set -euo pipefail
 
-# Cores
+# -----------------------------
+# Cores para output
+# -----------------------------
 RED='\033[0;31m'
 GREEN='\033[0;32m'
 YELLOW='\033[1;33m'
@@ -16,25 +18,31 @@ log_success() { echo -e "${GREEN}[✓]${NC} $1"; }
 log_warn() { echo -e "${YELLOW}[!]${NC} $1"; }
 log_error() { echo -e "${RED}[✗]${NC} $1"; }
 
+command_exists() { command -v "$1" >/dev/null 2>&1; }
+
+# -----------------------------
 # Configuração
-COMFY_DIR="/root/comfy/ComfyUI"
+# -----------------------------
+COMFY_DIR="${COMFY_DIR:-/root/comfy/ComfyUI}"
 MODELS_DIR="$COMFY_DIR/models"
-COMFY_HOST="0.0.0.0"
-COMFY_PORT="8818"
-CIVITAI_TOKEN="4fcb2834969399006a736ee402b061e5"
-
-# Performance
-export MAX_JOBS=32
-export NVCC_APPEND_FLAGS="--threads 8"
+COMFY_HOST="${COMFY_HOST:-0.0.0.0}"
+COMFY_PORT="${COMFY_PORT:-8818}"
+CIVITAI_TOKEN="${CIVITAI_TOKEN:-4fcb2834969399006a736ee402b061e5}"
+HF_TOKEN="${HF_TOKEN:-}"
+
+# Configurações de performance
+export MAX_JOBS=${MAX_JOBS:-32}
+export NVCC_APPEND_FLAGS="${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 HF_TRANSFER_CONCURRENCY=${HF_TRANSFER_CONCURRENCY:-16}
+export EXT_PARALLEL=${EXT_PARALLEL:-4}
 
-# Lista de downloads - Formato: "URL|TIPO|NOME_OPCIONAL"
-DOWNLOAD_FILES=(
-    # Wan 2.2 models
+# -----------------------------
+# Listas de downloads (modelos do CODE2)
+# -----------------------------
+readonly DOWNLOAD_FILES=(
     "hf://Kijai/WanVideo_comfy/Wan22-Lightning/Wan2.2-Lightning_T2V-v1.1-A14B-4steps-lora_HIGH_fp16.safetensors|loras|Wan2.2-Lightning_T2V-v1.1-A14B-4steps-lora_HIGH_fp16.safetensors"
     "hf://Kijai/WanVideo_comfy/Wan22-Lightning/Wan2.2-Lightning_T2V-v1.1-A14B-4steps-lora_LOW_fp16.safetensors|loras|Wan2.2-Lightning_T2V-v1.1-A14B-4steps-lora_LOW_fp16.safetensors"
     "hf://Kijai/WanVideo_comfy_fp8_scaled/I2V/Wan2_2-I2V-A14B-HIGH_fp8_e5m2_scaled_KJ.safetensors|diffusion_models|Wan2_2-I2V-A14B-HIGH_fp8_e5m2_scaled_KJ.safetensors"
@@ -43,22 +51,15 @@ DOWNLOAD_FILES=(
     "hf://Kijai/WanVideo_comfy/Wan22-Lightning/Wan2.2-Lightning_I2V-A14B-4steps-lora_LOW_fp16.safetensors|loras|Wan2.2-Lightning_I2V-A14B-4steps-lora_LOW_fp16.safetensors"
     "hf://Kijai/WanVideo_comfy/Wan2_1_VAE_fp32.safetensors|vae|Wan2_1_VAE_fp32.safetensors"
     "hf://Kijai/WanVideo_comfy/umt5-xxl-enc-bf16.safetensors|text_encoders|umt5-xxl-enc-bf16.safetensors"
-    
-    # Wan 2.1 models
     "hf://Kijai/WanVideo_comfy/Lightx2v/lightx2v_I2V_14B_480p_cfg_step_distill_rank256_bf16.safetensors|loras|lightx2v_I2V_14B_480p_cfg_step_distill_rank256_bf16.safetensors"
     "hf://Comfy-Org/Wan_2.2_ComfyUI_Repackaged/split_files/vae/wan_2.1_vae.safetensors|vae|wan_2.1_vae.safetensors"
     "hf://Comfy-Org/Wan_2.2_ComfyUI_Repackaged/split_files/text_encoders/umt5_xxl_fp16.safetensors|text_encoders|umt5_xxl_fp16.safetensors"
-    
-    # Wan 2.1 clip vision
     "hf://Kijai/WanVideo_comfy/open-clip-xlm-roberta-large-vit-huge-14_visual_fp32.safetensors|clip_vision|open-clip-xlm-roberta-large-vit-huge-14_visual_fp32.safetensors"
     "hf://Comfy-Org/Wan_2.1_ComfyUI_repackaged/split_files/clip_vision/clip_vision_h.safetensors|clip_vision|clip_vision_h.safetensors"
-    
-    # Upscaler extra
     "https://huggingface.co/ABDALLALSWAITI/Upscalers/resolve/main/anime/2x-AnimeSharpV2_MoSR_Soft.pth|upscale_models|2x-AnimeSharpV2_MoSR_Soft.pth"
 )
 
-# Mega downloads
-MEGA_LORAS=(
+readonly MEGA_LORAS=(
     "https://mega.nz/file/QFxREJpZ#0dge8TyRgb52lMFKdsk-Ac-rwcoG49hvBHvaqNsZkuk"
     "https://mega.nz/file/lIhQETjT#nKlQ01Euw6fUxo8P5bfwlQrxAkWEedXQQOqHv5u4yrM"
     "https://mega.nz/file/UIAESACY#Qlei1Pj5Nwno3Sz-wsTekvD-YgCkN1A5QbFKoRHXG9E"
@@ -72,13 +73,14 @@ MEGA_LORAS=(
     "https://mega.nz/file/hZhhhCQS#4je0Im2g1kEvTCuPotjVWkA-g6XICkqMb-mKxMt6R3Y"
 )
 
-MEGA_UPSCALE=(
+readonly MEGA_UPSCALE=(
     "https://mega.nz/file/qdQwAJSC#rPY2Hg4x-_RbHiHUBFW-YU8xZUqPpq5Gpl1uPFiJvW4"
     "https://mega.nz/file/KIh2marA#CehVNXxQhGfp1bqUL0K1YPWNqrCUoLPCZGq2W6XqmWo"
 )
 
-# Custom nodes do CODE2
-CUSTOM_NODES=(
+readonly RIFE_MODELS=()
+
+readonly CUSTOM_NODES=(
     "https://github.com/kijai/ComfyUI-Florence2"
     "https://github.com/adbrasi/groqrouter"
     "https://github.com/kijai/ComfyUI-WanVideoWrapper"
@@ -105,259 +107,452 @@ CUSTOM_NODES=(
     "https://github.com/grmchn/ComfyUI-ProportionChanger"
 )
 
-# Download HuggingFace
-download_hf() {
-    local url="$1"
-    local target_dir="$2"
-    local filename="$3"
-    
-    # Extrair repo e file_path da URL hf://
-    url="${url#hf://}"
-    local repo="$(echo "$url" | cut -d'/' -f1-2)"
-    local file_path="$(echo "$url" | cut -d'/' -f3-)"
-    
-    if [ -n "$filename" ] && [ -f "$target_dir/$filename" ]; then
-        log_success "Já existe: $filename"
-        return 0
-    fi
-    
-    log_info "Baixando HF: $filename"
-    
-    local temp_dir
-    temp_dir="$(mktemp -d)"
-    
-    if command -v hf >/dev/null 2>&1; then
-        HF_HUB_ENABLE_HF_TRANSFER=1 hf download "$repo" "$file_path" \
-            --local-dir "$temp_dir" --local-dir-use-symlinks False 2>/dev/null || {
-            rm -rf "$temp_dir"
-            log_error "Falha ao baixar: $filename"
-            return 1
-        }
-    elif command -v huggingface-cli >/dev/null 2>&1; then
-        HF_HUB_ENABLE_HF_TRANSFER=1 huggingface-cli download "$repo" "$file_path" \
-            --local-dir "$temp_dir" --local-dir-use-symlinks False 2>/dev/null || {
-            rm -rf "$temp_dir"
-            log_error "Falha ao baixar: $filename"
-            return 1
-        }
+# -----------------------------
+# Funções auxiliares
+# -----------------------------
+apt_update_safe() {
+    if command_exists apt-get; then
+        export DEBIAN_FRONTEND=noninteractive
+        apt-get update -y || log_warn "apt-get update falhou; seguindo mesmo assim."
     else
-        wget -q --show-progress -c -O "$target_dir/$filename" "https://huggingface.co/$url"
-        rm -rf "$temp_dir"
-        return $?
+        log_warn "apt-get não disponível; pulando update."
     fi
-    
-    # Mover arquivo baixado para destino
-    local downloaded_file
-    downloaded_file="$(find "$temp_dir" -type f \( -name "*.safetensors" -o -name "*.pth" -o -name "*.pt" -o -name "*.bin" \) | head -1)"
-    if [ -n "$downloaded_file" ]; then
-        mv "$downloaded_file" "$target_dir/$filename"
-        rm -rf "$temp_dir"
-        log_success "Download concluído: $filename"
-        return 0
+}
+
+apt_install_safe() {
+    if command_exists apt-get; then
+        export DEBIAN_FRONTEND=noninteractive
+        apt-get install -y "$@" || log_warn "Falha ao instalar pacote(s): $*"
     else
-        rm -rf "$temp_dir"
-        log_error "Arquivo não encontrado após download: $filename"
-        return 1
+        log_warn "apt-get não disponível; não foi possível instalar: $*"
     fi
 }
 
-# Download Mega
-download_mega() {
+stat_bytes() {
+    local path="$1"
+    stat -c%s "$path" 2>/dev/null || echo 0
+}
+
+add_civitai_token() {
     local url="$1"
-    local target_dir="$2"
-    
-    log_info "Baixando do Mega..."
-    
-    cd "$target_dir"
-    if megadl "$url" 2>/dev/null; then
-        log_success "Mega download OK"
+    if [[ "$url" == *"civitai.com/api/download"* ]] && [[ "$url" != *"token="* ]] && [[ -n "$CIVITAI_TOKEN" ]]; then
+        if [[ "$url" == *"?"* ]]; then
+            printf '%s&token=%s' "$url" "$CIVITAI_TOKEN"
+        else
+            printf '%s?token=%s' "$url" "$CIVITAI_TOKEN"
+        fi
     else
-        log_warn "Mega download falhou"
+        printf '%s' "$url"
     fi
-    cd - >/dev/null
 }
 
-# Download normal
+download_hf() {
+    local repo="$1"
+    local file_path="$2"
+    local target_dir="$3"
+    local filename="$4"
+
+    if [[ -z "$filename" ]]; then
+        filename="$(basename "$file_path")"
+    fi
+
+    local target_file="$target_dir/$filename"
+    if [[ -f "$target_file" ]] && (( $(stat_bytes "$target_file") > 1000000 )); then
+        log_success "Arquivo já existe: $filename"
+        return 0
+    fi
+
+    log_info "Baixando de HF: $filename"
+    local temp_dir
+    temp_dir="$(mktemp -d)"
+    local 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; then
+            success=true
+        else
+            log_warn "hf download falhou; tentando fallback."
+        fi
+    fi
+
+    if [[ "$success" == false ]] && command_exists huggingface-cli; then
+        if HF_HUB_ENABLE_HF_TRANSFER=1 huggingface-cli download "$repo" "$file_path" --local-dir "$temp_dir" --local-dir-use-symlinks False; then
+            success=true
+        else
+            log_warn "huggingface-cli falhou; tentando download direto."
+        fi
+    fi
+
+    if [[ "$success" == false ]]; then
+        local direct_url="https://huggingface.co/${repo}/resolve/main/${file_path}"
+        local header=""
+        if [[ -n "$HF_TOKEN" ]]; then
+            header="Authorization: Bearer ${HF_TOKEN}"
+        fi
+        if command_exists aria2c; then
+            local aria_cmd=(aria2c -c -s 16 -x 16 -k 1M --console-log-level=warn --summary-interval=10)
+            if [[ -n "$header" ]]; then
+                aria_cmd+=(--header="$header")
+            fi
+            aria_cmd+=(--dir="$temp_dir" --out="$filename" "$direct_url")
+            if "${aria_cmd[@]}"; then
+                success=true
+            else
+                log_warn "aria2c direto falhou; tentando curl/wget."
+            fi
+        fi
+        if [[ "$success" == false ]] && command_exists curl; then
+            local curl_cmd=(curl -L -# -o "$temp_dir/$filename")
+            if [[ -n "$header" ]]; then
+                curl_cmd+=(-H "$header")
+            fi
+            curl_cmd+=("$direct_url")
+            if "${curl_cmd[@]}"; then
+                success=true
+            fi
+        fi
+        if [[ "$success" == false ]] && command_exists wget; then
+            local wget_cmd=(wget -q --show-progress -c -O "$temp_dir/$filename")
+            if [[ -n "$header" ]]; then
+                wget_cmd+=(--header="$header")
+            fi
+            wget_cmd+=("$direct_url")
+            if "${wget_cmd[@]}"; then
+                success=true
+            fi
+        fi
+    fi
+
+    if [[ "$success" == true ]]; then
+        local downloaded_file
+        downloaded_file="$(find "$temp_dir" -maxdepth 1 -type f | head -n 1)"
+        if [[ -n "$downloaded_file" ]]; then
+            mkdir -p "$target_dir"
+            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
+}
+
 download_file() {
     local url="$1"
     local target_dir="$2"
     local filename="$3"
-    
-    # Adicionar token Civitai se necessário
-    if [[ "$url" == *"civitai.com"* ]] && [[ "$url" != *"token="* ]]; then
-        url="${url}&token=${CIVITAI_TOKEN}"
+
+    url="$(add_civitai_token "$url")"
+    mkdir -p "$target_dir"
+
+    local final_name="$filename"
+    if [[ -z "$final_name" ]]; then
+        final_name="$(basename "${url%%\?*}")"
     fi
-    
-    if [ -n "$filename" ] && [ -f "$target_dir/$filename" ]; then
-        log_success "Já existe: $filename"
+
+    if [[ -f "$target_dir/$final_name" ]] && (( $(stat_bytes "$target_dir/$final_name") > 500000 )); then
+        log_success "Arquivo já existe: $final_name"
         return 0
     fi
-    
-    log_info "Baixando: ${filename:-arquivo}"
-    
-    # Tentar com aria2c primeiro
-    if command -v aria2c >/dev/null 2>&1; then
-        if [ -n "$filename" ]; then
-            aria2c -c -x 16 -s 16 -k 1M --console-log-level=warn --dir="$target_dir" --out="$filename" "$url" || \
-            wget -q --show-progress -c -O "$target_dir/$filename" "$url" || \
-            log_error "Download falhou: $filename"
-        else
-            aria2c -c -x 16 -s 16 -k 1M --console-log-level=warn --dir="$target_dir" "$url" || \
-            (cd "$target_dir" && wget -q --show-progress -c "$url") || \
-            log_error "Download falhou: $url"
+
+    log_info "Baixando: $final_name"
+
+    if command_exists aria2c; then
+        local aria_opts=(-c -s 16 -x 16 -k 1M --console-log-level=warn --summary-interval=10)
+        if aria2c "${aria_opts[@]}" --dir="$target_dir" --out="$final_name" "$url"; then
+            return 0
         fi
-    elif [ -n "$filename" ]; then
-        wget -q --show-progress -c -O "$target_dir/$filename" "$url" || \
-        log_error "Download falhou: $filename"
-    else
-        (cd "$target_dir" && wget -q --show-progress -c "$url") || \
-        log_error "Download falhou: $url"
+        log_warn "aria2c falhou; tentando wget/curl."
     fi
+
+    if command_exists wget; then
+        if wget -q --show-progress -c -O "$target_dir/$final_name" "$url"; then
+            return 0
+        fi
+        log_warn "wget falhou; tentando curl."
+    fi
+
+    if command_exists curl; then
+        if curl -L -# -C - -o "$target_dir/$final_name" "$url"; then
+            return 0
+        fi
+    fi
+
+    log_error "Falha ao baixar: $url"
+    return 1
+}
+
+download_mega() {
+    local url="$1"
+    local target_dir="$2"
+
+    url="${url#mega://}"
+    mkdir -p "$target_dir"
+
+    if ! command_exists megadl; then
+        log_warn "megadl não encontrado; tentando instalar megatools."
+        apt_update_safe
+        apt_install_safe megatools
+    fi
+
+    if command_exists megadl; then
+        log_info "Baixando de Mega.nz: $url"
+        if (cd "$target_dir" && megadl "$url"); then
+            log_success "Download Mega.nz concluído"
+            return 0
+        fi
+    fi
+
+    log_error "Falha ao baixar de Mega.nz: $url"
+    return 1
 }
 
-# Processar downloads
 process_downloads() {
+    local failed=0
     for entry in "${DOWNLOAD_FILES[@]}"; do
-        IFS='|' read -r url type filename <<< "$entry"
-        
-        # Criar diretório
+        IFS='|' read -r raw_url type filename <<< "$entry"
+        raw_url="$(echo "$raw_url" | xargs)"
+        type="$(echo "$type" | xargs)"
+        filename="$(echo "$filename" | xargs)"
+
         local target_dir="$MODELS_DIR/$type"
         mkdir -p "$target_dir"
-        
-        # Baixar arquivo
-        if [[ "$url" == hf://* ]]; then
-            download_hf "$url" "$target_dir" "$filename"
-        else
-            download_file "$url" "$target_dir" "$filename"
+
+        local attempt success=false
+        local label="$filename"
+        if [[ -z "$label" ]]; then
+            label="$(basename "${raw_url%%\?*}")"
+        fi
+        if [[ -z "$label" || "$label" == "." ]]; then
+            label="$raw_url"
+        fi
+        for attempt in 1 2 3; do
+            if [[ "$raw_url" == mega://* ]]; then
+                if download_mega "$raw_url" "$target_dir"; then
+                    success=true
+                fi
+            elif [[ "$raw_url" == hf://* ]]; then
+                local trimmed="${raw_url#hf://}"
+                local repo="$(echo "$trimmed" | cut -d'/' -f1-2)"
+                local file_path="$(echo "$trimmed" | cut -d'/' -f3-)"
+                if download_hf "$repo" "$file_path" "$target_dir" "$filename"; then
+                    success=true
+                fi
+            else
+                if download_file "$raw_url" "$target_dir" "$filename"; then
+                    success=true
+                fi
+            fi
+
+            if [[ "$success" == true ]]; then
+                break
+            fi
+
+            log_warn "Tentativa $attempt falhou para $label; aguardando antes de tentar novamente."
+            sleep $((2 * attempt))
+        done
+
+        if [[ "$success" == false ]]; then
+            log_error "Falhou após múltiplas tentativas: $label"
+            ((failed++))
         fi
     done
+
+    return $failed
 }
 
-# Clone git repo
-clone_repo() {
+clone_or_update() {
     local url="$1"
     local dest="$2"
-    
-    if [ -d "$dest/.git" ]; then
-        git -C "$dest" pull --ff-only 2>/dev/null || true
+    local node_name
+    node_name="$(basename "$dest")"
+
+    mkdir -p "$(dirname "$dest")"
+
+    if [[ -d "$dest/.git" ]]; then
+        log_info "Atualizando $node_name..."
+        if ! timeout 120 git -C "$dest" pull --ff-only --rebase=false; then
+            log_warn "Timeout/erro ao atualizar $node_name"
+        fi
     else
-        git clone --depth 1 "$url" "$dest" 2>/dev/null || true
+        log_info "Clonando $node_name..."
+        if ! timeout 120 git clone --recursive --depth 1 "$url" "$dest"; then
+            log_warn "Falha ao clonar $node_name"
+            return
+        fi
     fi
-    
-    # Instalar requirements se existir
-    if [ -f "$dest/requirements.txt" ]; then
-        pip install -q -r "$dest/requirements.txt" 2>/dev/null || true
+
+    if [[ -f "$dest/requirements.txt" ]]; then
+        if ! timeout 180 python3 -m pip install --no-warn-script-location -q -r "$dest/requirements.txt"; then
+            log_warn "Falha ao instalar requirements para $node_name"
+        fi
     fi
 }
 
-# ========== INSTALAÇÃO ==========
+install_sageattention() {
+    log_info "[5/9] Instalando SageAttention..."
+    if python3 - <<'PY' 2>/dev/null | grep -q True; then
+import torch
+print(torch.cuda.is_available())
+PY
+        if python3 -m pip install "https://huggingface.co/adbrasi/comfywheel/resolve/main/sageattention-2.2.0-cp311-cp311-linux_x86_64.whl" -q; then
+            if python3 - <<'PY' >/dev/null 2>&1; then
+import sageattention
+PY
+                log_success "SageAttention instalado"
+            else
+                log_warn "SageAttention instalado mas não pôde ser importado"
+            fi
+        else
+            log_warn "Falha ao instalar SageAttention"
+        fi
+    else
+        log_warn "CUDA não detectado; pulando SageAttention"
+    fi
+}
 
+# -----------------------------
+# Instalação principal
+# -----------------------------
 log_info "========================================="
-log_info " ComfyUI + Wan 2.2 Setup"
+log_info " ComfyUI + Wan 2.2 Setup (Merged)"
 log_info "========================================="
 
-# 1. Instalar dependências
-log_info "[1/7] Instalando dependências..."
-
-apt-get update -qq
-apt-get install -y -qq aria2 megatools git wget curl 2>/dev/null
-
-pip install --upgrade pip wheel setuptools -q
-pip install --upgrade "huggingface_hub[cli,hf_transfer]>=0.26.0" comfy-cli -q
-
-log_success "Dependências instaladas"
-
-# 2. Instalar ComfyUI
-log_info "[2/7] Instalando ComfyUI..."
+log_info "[1/9] Preparando dependências do sistema..."
+if ! command_exists aria2c || ! command_exists megadl; then
+    apt_update_safe
+    command_exists aria2c || apt_install_safe aria2
+    command_exists megadl || apt_install_safe megatools
+fi
+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 de sistema OK"
 
-if [ -f "$COMFY_DIR/main.py" ]; then
-    log_warn "ComfyUI já existe"
+log_info "[2/9] Atualizando pip e 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" ]] && command_exists huggingface-cli; then
+    huggingface-cli login --token "$HF_TOKEN" --add-to-git-credential || log_warn "Falha ao autenticar no HuggingFace"
+fi
+if command_exists hf; then
+    log_success "Ferramenta 'hf' disponível"
+elif command_exists huggingface-cli; then
+    log_warn "Usando 'huggingface-cli' para downloads"
 else
-    comfy --skip-prompt install --fast-deps --nvidia
+    log_warn "Sem CLI oficial do HuggingFace; tentaremos fallback HTTP"
 fi
 
-log_success "ComfyUI instalado"
-
-# 3. Instalar PyTorch com CUDA 12.8
-log_info "[3/7] Instalando PyTorch com CUDA 12.8..."
-pip install --upgrade torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu128
+log_info "[3/9] Instalando PyTorch CUDA 12.8..."
+python3 -m pip install --upgrade torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu128
 log_success "PyTorch com CUDA 12.8 instalado"
 
-# 4. Instalar SageAttention
-log_info "[4/7] Instalando SageAttention..."
-if python3 -c "import torch; print(torch.cuda.is_available())" 2>/dev/null | grep -q True; then
-    pip install "https://huggingface.co/adbrasi/comfywheel/resolve/main/sageattention-2.2.0-cp311-cp311-linux_x86_64.whl" -q || log_warn "Falha ao instalar SageAttention"
-    python3 -c "import sageattention" >/dev/null 2>&1 || log_warn "SageAttention não pôde ser importado"
-    log_success "SageAttention instalado"
+log_info "[4/9] Instalando ComfyUI..."
+if [[ -f "$COMFY_DIR/main.py" ]]; then
+    log_warn "ComfyUI já existe; pulando instalação base"
 else
-    log_warn "CUDA não detectado, pulando SageAttention"
+    comfy --skip-prompt install --fast-deps --nvidia
+    if [[ ! -f "$COMFY_DIR/main.py" ]]; then
+        log_error "ComfyUI não foi instalado corretamente"
+        exit 1
+    fi
 fi
+log_success "ComfyUI disponível em $COMFY_DIR"
 
-# 5. Baixar modelos
-log_info "[5/7] Baixando modelos..."
+install_sageattention
 
-# Criar diretórios
+log_info "[6/9] Baixando modelos principais..."
 mkdir -p "$MODELS_DIR"/{diffusion_models,loras,vae,vae_approx,text_encoders,clip_vision,controlnet,upscale_models,checkpoints}
+if process_downloads; then
+    log_success "Modelos baixados (ou já presentes)"
+else
+    log_warn "Alguns modelos falharam. Consulte logs acima e tente novamente."
+fi
 
-# Processar downloads principais
-process_downloads
-
-# Baixar Mega LoRAs
-LORA_DIR="$MODELS_DIR/loras"
-for mega_url in "${MEGA_LORAS[@]}"; do
-    download_mega "$mega_url" "$LORA_DIR"
-done
-
-# Baixar Mega Upscale
-UPSCALE_DIR="$MODELS_DIR/upscale_models"
-for mega_url in "${MEGA_UPSCALE[@]}"; do
-    download_mega "$mega_url" "$UPSCALE_DIR"
-done
-
-log_success "Downloads processados"
+log_info "[7/9] Baixando modelos RIFE (opcional)..."
+if (( ${#RIFE_MODELS[@]} )); then
+    RIFE_DIR="$MODELS_DIR/tensorrt/rife"
+    mkdir -p "$RIFE_DIR"
+    for url in "${RIFE_MODELS[@]}"; do
+        fname="$(basename "${url%%\?*}")"
+        download_file "$url" "$RIFE_DIR" "$fname" || log_warn "Falha ao baixar RIFE: $fname"
+    done
+    log_success "Modelos RIFE processados"
+else
+    log_warn "Lista RIFE vazia — pulando"
+fi
 
-# 6. Instalar custom nodes
-log_info "[6/7] Instalando custom nodes..."
+log_info "[8/9] Baixando arquivos do Mega.nz..."
+if (( ${#MEGA_LORAS[@]} )); then
+    LORA_DIR="$MODELS_DIR/loras"
+    mkdir -p "$LORA_DIR"
+    for url in "${MEGA_LORAS[@]}"; do
+        download_mega "$url" "$LORA_DIR" || log_warn "Mega.nz falhou: $url"
+    done
+fi
+if (( ${#MEGA_UPSCALE[@]} )); then
+    UPSCALE_DIR="$MODELS_DIR/upscale_models"
+    mkdir -p "$UPSCALE_DIR"
+    for url in "${MEGA_UPSCALE[@]}"; do
+        download_mega "$url" "$UPSCALE_DIR" || log_warn "Mega.nz falhou: $url"
+    done
+fi
+log_success "Downloads Mega.nz concluídos (quando disponíveis)"
 
+log_info "[9/9] Instalando custom nodes..."
 CN_DIR="$COMFY_DIR/custom_nodes"
 mkdir -p "$CN_DIR"
-
 total_nodes=${#CUSTOM_NODES[@]}
 current_node=0
-
 for repo in "${CUSTOM_NODES[@]}"; do
     ((current_node++))
-    node_name=$(basename "$repo")
-    log_info "[$current_node/$total_nodes] Instalando: $node_name"
-    clone_repo "$repo" "$CN_DIR/$node_name"
+    log_info "[$current_node/$total_nodes] Processando $(basename "$repo")"
+    clone_or_update "$repo" "$CN_DIR/$(basename "$repo")"
 done
-
 log_success "Custom nodes instalados"
 
-# 7. Verificação final
-log_info "[7/7] Verificando instalação..."
-
-echo ""
-log_info "Modelos disponíveis:"
-for dir in diffusion_models loras vae vae_approx text_encoders clip_vision controlnet upscale_models checkpoints; do
-    if [ -d "$MODELS_DIR/$dir" ]; then
-        local count
-        count=$(find "$MODELS_DIR/$dir" -type f \( -name "*.safetensors" -o -name "*.pth" -o -name "*.pt" -o -name "*.bin" \) 2>/dev/null | wc -l)
-        if [ "$count" -gt 0 ]; then 
-            log_success " $dir: $count arquivo(s)"
+check_critical_files() {
+    local files=(
+        "$COMFY_DIR/main.py"
+        "$MODELS_DIR/vae/wan_2.1_vae.safetensors"
+        "$MODELS_DIR/vae/Wan2_1_VAE_fp32.safetensors"
+        "$MODELS_DIR/text_encoders/umt5_xxl_fp16.safetensors"
+        "$MODELS_DIR/text_encoders/umt5-xxl-enc-bf16.safetensors"
+        "$MODELS_DIR/clip_vision/clip_vision_h.safetensors"
+    )
+    local missing=0
+    for file in "${files[@]}"; do
+        if [[ -f "$file" ]]; then
+            log_success "✓ $(basename "$file")"
+        else
+            log_warn "✗ $(basename "$file") não encontrado"
+            ((missing++))
         fi
-    fi
-done
+    done
+    return $missing
+}
 
-# Verificar e iniciar com SageAttention se disponível
-log_info "========================================="
-log_success "Instalação concluída!"
-log_info "Iniciando ComfyUI..."
-log_info "URL: http://localhost:$COMFY_PORT"
 log_info "========================================="
+if check_critical_files; then
+    log_success "Instalação concluída com sucesso!"
+else
+    log_warn "Alguns arquivos essenciais não foram encontrados. Veja avisos acima."
+fi
+log_info "Preparando para iniciar ComfyUI..."
 
 SAGE_FLAG=""
-python3 -c "import sageattention" >/dev/null 2>&1 && SAGE_FLAG="--use-sage-attention" || true
+if python3 - <<'PY' >/dev/null 2>&1; then
+import sageattention
+PY
+    SAGE_FLAG="--use-sage-attention"
+fi
 
 cd "$COMFY_DIR"
-exec comfy launch -- ${SAGE_FLAG:+$SAGE_FLAG} --listen "$COMFY_HOST" --port "$COMFY_PORT"
+log_info "Iniciando ComfyUI em http://$COMFY_HOST:$COMFY_PORT"
+exec comfy launch -- ${SAGE_FLAG:+$SAGE_FLAG} --listen "$COMFY_HOST" --port "$COMFY_PORT"