Update wanForProduction.sh

wanForProduction.sh

@@ -1,12 +1,13 @@
 #!/usr/bin/env bash
-# setup_comfyui_combined.sh
-# Script baseado no CODE1 com melhorias e modelos/nodes do CODE2
+# setup_comfyui_simple.sh (base do CODE1, ajustado)
+# - troca de modelos/nodes pelos do CODE2
+# - adiciona SageAttention
+# - remove tratamento 5090/Blackwell
+# - melhora download_hf para paths com subpastas e /resolve/
 
 set -euo pipefail
 
-# -----------------------------
-# Cores para output
-# -----------------------------
+# Cores
 RED='\033[0;31m'
 GREEN='\033[0;32m'
 YELLOW='\033[1;33m'
@@ -18,69 +19,66 @@ 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="${COMFY_DIR:-/root/comfy/ComfyUI}"
+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:-}"
-
-# 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"
+COMFY_HOST="0.0.0.0"
+COMFY_PORT="8818"
+CIVITAI_TOKEN="4fcb2834969399006a736ee402b061e5"
+
+# Performance
+export MAX_JOBS=32
 export HF_HUB_ENABLE_HF_TRANSFER=1
-export HF_TRANSFER_CONCURRENCY=${HF_TRANSFER_CONCURRENCY:-16}
-export EXT_PARALLEL=${EXT_PARALLEL:-4}
-
-# -----------------------------
-# 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"
-    "hf://Kijai/WanVideo_comfy_fp8_scaled/I2V/Wan2_2-I2V-A14B-LOW_fp8_e5m2_scaled_KJ.safetensors|diffusion_models|Wan2_2-I2V-A14B-LOW_fp8_e5m2_scaled_KJ.safetensors"
-    "hf://Kijai/WanVideo_comfy/Wan22-Lightning/Wan2.2-Lightning_I2V-A14B-4steps-lora_HIGH_fp16.safetensors|loras|Wan2.2-Lightning_I2V-A14B-4steps-lora_HIGH_fp16.safetensors"
-    "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"
-    "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"
-    "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"
-    "https://huggingface.co/ABDALLALSWAITI/Upscalers/resolve/main/anime/2x-AnimeSharpV2_MoSR_Soft.pth|upscale_models|2x-AnimeSharpV2_MoSR_Soft.pth"
-)
 
-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"
-    "https://mega.nz/file/cdAg1Q7D#tQVrer2sB-M7XdYiKAqr7kRRfJw4pXCdliOg2aPJUGE"
-    "https://mega.nz/file/5RQFiAAJ#QCWrHkuJjXlFufR18g5R9YXJGqJ1pQp7xbqNRgSYUjc"
-    "https://mega.nz/file/tABgxLTS#NNlPV8NPc7R52MIjABMFUb41EOFDxFoszBxOrPMN13Q"
-    "https://mega.nz/file/4MxBgJ5Y#R8OSby0u-t-QFBnpnOtdjN3xuiVRvHa4_OEzTwWoGKQ"
-    "https://mega.nz/file/dJJxiSRR#o3XsYfSvb04xg7gNVdfHb3XTVHgvx0ZyXe9N8VQ9Lfw"
-    "https://mega.nz/file/wcAFhRJZ#ifZcckts0wPq1FG693tYLkpFrvjXvab1ZQeL6DuL1T8"
-    "https://mega.nz/file/wNIgmRwb#vn6rKM3QuTDpWVTLEBDsKjIkNkaX2wLKeauKNeVkcFs"
-    "https://mega.nz/file/hZhhhCQS#4je0Im2g1kEvTCuPotjVWkA-g6XICkqMb-mKxMt6R3Y"
-)
+# ===============================
+# Lista de downloads (do CODE2)
+# ===============================
+# Formato: "URL|TIPO|NOME_OPCIONAL"
+DOWNLOAD_FILES=(
+    # Wan 2.2 (Kijai) - use URLs SEM /resolve/ para usar o CLI da HF
+    "https://huggingface.co/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"
+    "https://huggingface.co/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"
+    "https://huggingface.co/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"
+    "https://huggingface.co/Kijai/WanVideo_comfy_fp8_scaled/I2V/Wan2_2-I2V-A14B-LOW_fp8_e5m2_scaled_KJ.safetensors|diffusion_models|Wan2_2-I2V-A14B-LOW_fp8_e5m2_scaled_KJ.safetensors"
+    "https://huggingface.co/Kijai/WanVideo_comfy/Wan22-Lightning/Wan2.2-Lightning_I2V-A14B-4steps-lora_HIGH_fp16.safetensors|loras|Wan2.2-Lightning_I2V-A14B-4steps-lora_HIGH_fp16.safetensors"
+    "https://huggingface.co/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"
+    "https://huggingface.co/Kijai/WanVideo_comfy/Wan2_1_VAE_fp32.safetensors|vae|Wan2_1_VAE_fp32.safetensors"
+    "https://huggingface.co/Kijai/WanVideo_comfy/umt5-xxl-enc-bf16.safetensors|text_encoders|umt5-xxl-enc-bf16.safetensors"
+
+    # Wan 2.1
+    "https://huggingface.co/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"
+    "https://huggingface.co/Comfy-Org/Wan_2.2_ComfyUI_Repackaged/split_files/vae/wan_2.1_vae.safetensors|vae|wan_2.1_vae.safetensors"
+    "https://huggingface.co/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
+    "https://huggingface.co/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"
+    "https://huggingface.co/Comfy-Org/Wan_2.1_ComfyUI_repackaged/split_files/clip_vision/clip_vision_h.safetensors|clip_vision|clip_vision_h.safetensors"
+
+    # Upscaler extra (com /resolve/ para baixar via HTTP direto)
+    "https://huggingface.co/ABDALLALSWAITI/Upscalers/resolve/main/anime/2x-AnimeSharpV2_MoSR_Soft.pth|upscale_models|2x-AnimeSharpV2_MoSR_Soft.pth"
 
-readonly MEGA_UPSCALE=(
-    "https://mega.nz/file/qdQwAJSC#rPY2Hg4x-_RbHiHUBFW-YU8xZUqPpq5Gpl1uPFiJvW4"
-    "https://mega.nz/file/KIh2marA#CehVNXxQhGfp1bqUL0K1YPWNqrCUoLPCZGq2W6XqmWo"
+    # LoRAs via Mega.nz (do CODE2) — prefixo mega:// para usar função existente
+    "mega://https://mega.nz/file/QFxREJpZ#0dge8TyRgb52lMFKdsk-Ac-rwcoG49hvBHvaqNsZkuk|loras|"
+    "mega://https://mega.nz/file/lIhQETjT#nKlQ01Euw6fUxo8P5bfwlQrxAkWEedXQQOqHv5u4yrM|loras|"
+    "mega://https://mega.nz/file/UIAESACY#Qlei1Pj5Nwno3Sz-wsTekvD-YgCkN1A5QbFKoRHXG9E|loras|"
+    "mega://https://mega.nz/file/cdAg1Q7D#tQVrer2sB-M7XdYiKAqr7kRRfJw4pXCdliOg2aPJUGE|loras|"
+    "mega://https://mega.nz/file/5RQFiAAJ#QCWrHkuJjXlFufR18g5R9YXJGqJ1pQp7xbqNRgSYUjc|loras|"
+    "mega://https://mega.nz/file/tABgxLTS#NNlPV8NPc7R52MIjABMFUb41EOFDxFoszBxOrPMN13Q|loras|"
+    "mega://https://mega.nz/file/4MxBgJ5Y#R8OSby0u-t-QFBnpnOtdjN3xuiVRvHa4_OEzTwWoGKQ|loras|"
+    "mega://https://mega.nz/file/dJJxiSRR#o3XsYfSvb04xg7gNVdfHb3XTVHgvx0ZyXe9N8VQ9Lfw|loras|"
+    "mega://https://mega.nz/file/wcAFhRJZ#ifZcckts0wPq1FG693tYLkpFrvjXvab1ZQeL6DuL1T8|loras|"
+    "mega://https://mega.nz/file/wNIgmRwb#vn6rKM3QuTDpWVTLEBDsKjIkNkaX2wLKeauKNeVkcFs|loras|"
+    "mega://https://mega.nz/file/hZhhhCQS#4je0Im2g1kEvTCuPotjVWkA-g6XICkqMb-mKxMt6R3Y|loras|"
+
+    # Upscale via Mega.nz (do CODE2)
+    "mega://https://mega.nz/file/qdQwAJSC#rPY2Hg4x-_RbHiHUBFW-YU8xZUqPpq5Gpl1uPFiJvW4|upscale_models|"
+    "mega://https://mega.nz/file/KIh2marA#CehVNXxQhGfp1bqUL0K1YPWNqrCUoLPCZGq2W6XqmWo|upscale_models|"
 )
 
-readonly RIFE_MODELS=()
-
-readonly CUSTOM_NODES=(
+# ===============================
+# Custom nodes (do CODE2)
+# ===============================
+CUSTOM_NODES=(
     "https://github.com/kijai/ComfyUI-Florence2"
     "https://github.com/adbrasi/groqrouter"
     "https://github.com/kijai/ComfyUI-WanVideoWrapper"
@@ -107,452 +105,220 @@ readonly CUSTOM_NODES=(
     "https://github.com/grmchn/ComfyUI-ProportionChanger"
 )
 
-# -----------------------------
-# 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
-        log_warn "apt-get não disponível; pulando update."
-    fi
-}
-
-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
-        log_warn "apt-get não disponível; não foi possível instalar: $*"
-    fi
-}
-
-stat_bytes() {
-    local path="$1"
-    stat -c%s "$path" 2>/dev/null || echo 0
-}
-
-add_civitai_token() {
-    local url="$1"
-    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
-        printf '%s' "$url"
-    fi
-}
+# ===============================
+# Funções de download (mesmas do CODE1, porém download_hf melhorado)
+# ===============================
 
 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 url="$1"
+    local target_dir="$2"
+    local filename="$3"
 
-    local target_file="$target_dir/$filename"
-    if [[ -f "$target_file" ]] && (( $(stat_bytes "$target_file") > 1000000 )); then
-        log_success "Arquivo já existe: $filename"
+    if [ -n "$filename" ] && [ -f "$target_dir/$filename" ]; then
+        log_success "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
+    log_info "Baixando HF: ${filename:-$(basename "$url")}"
 
-    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
+    # Se a URL vier com /resolve/, usa HTTP direto (evita problemas com 'hf download')
+    if [[ "$url" == *"/resolve/"* ]]; then
+        mkdir -p "$target_dir"
+        wget -q --show-progress -c -O "$target_dir/$filename" "https://huggingface.co/$url"
+        return $?
     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
+    # Parse correto: repo_id = 2 primeiros segmentos; file_path = resto
+    local repo_id file_path
+    repo_id="$(echo "$url" | cut -d'/' -f1-2)"
+    file_path="$(echo "$url" | cut -d'/' -f3-)"
 
-    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
+    if command -v hf >/dev/null 2>&1; then
+        mkdir -p "$target_dir"
+        (cd "$target_dir" && HF_HUB_ENABLE_HF_TRANSFER=1 hf download "$repo_id" "$file_path" --local-dir . >/dev/null 2>&1) || {
+            log_warn "hf download falhou; tentando HTTP"
+            wget -q --show-progress -c -O "$target_dir/$filename" "https://huggingface.co/$repo_id/resolve/main/$file_path" || \
+            log_error "Download falhou: $filename"
+        }
+    else
+        mkdir -p "$target_dir"
+        wget -q --show-progress -c -O "$target_dir/$filename" "https://huggingface.co/$repo_id/resolve/main/$file_path" || \
+        log_error "Download falhou: $filename"
     fi
-
-    rm -rf "$temp_dir"
-    log_error "Falha ao baixar: $filename"
-    return 1
 }
 
-download_file() {
+download_mega() {
     local url="$1"
     local target_dir="$2"
-    local filename="$3"
-
-    url="$(add_civitai_token "$url")"
-    mkdir -p "$target_dir"
-
-    local final_name="$filename"
-    if [[ -z "$final_name" ]]; then
-        final_name="$(basename "${url%%\?*}")"
-    fi
 
-    if [[ -f "$target_dir/$final_name" ]] && (( $(stat_bytes "$target_dir/$final_name") > 500000 )); then
-        log_success "Arquivo já existe: $final_name"
-        return 0
-    fi
+    url="${url#mega://}" # remove prefixo
 
-    log_info "Baixando: $final_name"
+    log_info "Baixando do Mega..."
 
-    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
-        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
+    cd "$target_dir"
+    if megadl "$url" 2>/dev/null; then
+        log_success "Mega download OK"
+    else
+        log_warn "Mega download falhou"
     fi
-
-    log_error "Falha ao baixar: $url"
-    return 1
+    cd - >/dev/null
 }
 
-download_mega() {
+download_file() {
     local url="$1"
     local target_dir="$2"
+    local filename="$3"
 
-    url="${url#mega://}"
-    mkdir -p "$target_dir"
+    # Adicionar token Civitai se necessário
+    if [[ "$url" == *"civitai.com"* ]] && [[ "$url" != *"token="* ]]; then
+        url="${url}&token=${CIVITAI_TOKEN}"
+    fi
 
-    if ! command_exists megadl; then
-        log_warn "megadl não encontrado; tentando instalar megatools."
-        apt_update_safe
-        apt_install_safe megatools
+    if [ -n "$filename" ] && [ -f "$target_dir/$filename" ]; then
+        log_success "Já existe: $filename"
+        return 0
     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
+    log_info "Baixando: ${filename:-arquivo}"
+
+    if command -v aria2c >/dev/null 2>&1; then
+        if [ -n "$filename" ]; then
+            aria2c -c -x 4 -s 4 --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 4 -s 4 --console-log-level=warn --dir="$target_dir" "$url" || \
+            (cd "$target_dir" && wget -q --show-progress -c "$url") || \
+            log_error "Download falhou: $url"
         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"
     fi
-
-    log_error "Falha ao baixar de Mega.nz: $url"
-    return 1
 }
 
 process_downloads() {
-    local failed=0
     for entry in "${DOWNLOAD_FILES[@]}"; do
-        IFS='|' read -r raw_url type filename <<< "$entry"
-        raw_url="$(echo "$raw_url" | xargs)"
-        type="$(echo "$type" | xargs)"
-        filename="$(echo "$filename" | xargs)"
+        IFS='|' read -r url type filename <<< "$entry"
 
         local target_dir="$MODELS_DIR/$type"
         mkdir -p "$target_dir"
 
-        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++))
+        if [[ "$url" == mega://* ]]; then
+            download_mega "$url" "$target_dir"
+        elif [[ "$url" == *.huggingface.co/* ]]; then
+            local hf_path="${url#*huggingface.co/}"
+            download_hf "$hf_path" "$target_dir" "$filename"
+        else
+            download_file "$url" "$target_dir" "$filename"
         fi
     done
-
-    return $failed
 }
 
-clone_or_update() {
+clone_repo() {
     local url="$1"
     local dest="$2"
-    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
+    if [ -d "$dest/.git" ]; then
+        git -C "$dest" pull --ff-only 2>/dev/null || true
     else
-        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
+        git clone --depth 1 "$url" "$dest" 2>/dev/null || true
     fi
 
-    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
+    if [ -f "$dest/requirements.txt" ]; then
+        pip install -q -r "$dest/requirements.txt" 2>/dev/null || true
     fi
 }
 
-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 ==========
 
-# -----------------------------
-# Instalação principal
-# -----------------------------
 log_info "========================================="
-log_info " ComfyUI + Wan 2.2 Setup (Merged)"
+log_info " ComfyUI Setup"
 log_info "========================================="
 
-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"
+# 1. Dependências
+log_info "[1/6] Instalando dependências..."
 
-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
-    log_warn "Sem CLI oficial do HuggingFace; tentaremos fallback HTTP"
-fi
+apt-get update -qq
+apt-get install -y -qq aria2 megatools git wget curl 2>/dev/null
 
-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"
+pip install --upgrade pip wheel setuptools -q
+pip install --upgrade "huggingface_hub[cli,hf_transfer]" comfy-cli -q
 
-log_info "[4/9] Instalando ComfyUI..."
-if [[ -f "$COMFY_DIR/main.py" ]]; then
-    log_warn "ComfyUI já existe; pulando instalação base"
+log_success "Dependências instaladas"
+
+# 2. ComfyUI
+log_info "[2/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 corretamente"
-        exit 1
-    fi
 fi
-log_success "ComfyUI disponível em $COMFY_DIR"
 
-install_sageattention
+log_success "ComfyUI instalado"
 
-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
+# 3. PyTorch CUDA 12.8 (substitui o bloco de Blackwell/5090 do CODE1)
+log_info "[3/6] Instalando PyTorch (CUDA 12.8)..."
+pip install --upgrade torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu128 -q
+log_success "PyTorch com CUDA 12.8 instalado"
 
-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"
+# 4. SageAttention (do CODE2)
+log_info "[4/6] Instalando SageAttention..."
+if python3 - <<'PY' 2>/dev/null; then
+import torch, sys
+sys.exit(0 if torch.cuda.is_available() else 1)
+PY
+then
+    pip install -q "https://huggingface.co/adbrasi/comfywheel/resolve/main/sageattention-2.2.0-cp311-cp311-linux_x86_64.whl" || log_warn "Falha ao instalar SageAttention"
+    if python3 -c 'import sageattention' >/dev/null 2>&1; then
+        log_success "SageAttention instalado"
+    else
+        log_warn "SageAttention não pôde ser importado (seguindo sem)"
+    fi
 else
-    log_warn "Lista RIFE vazia — pulando"
+    log_warn "CUDA não detectada, pulando SageAttention"
 fi
 
-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)"
+# 5. Modelos
+log_info "[5/6] Baixando modelos..."
+process_downloads
+log_success "Downloads processados"
+
+# 6. Custom nodes
+log_info "[6/6] Instalando custom nodes..."
 
-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++))
-    log_info "[$current_node/$total_nodes] Processando $(basename "$repo")"
-    clone_or_update "$repo" "$CN_DIR/$(basename "$repo")"
+    node_name=$(basename "$repo")
+    log_info "[$current_node/$total_nodes] Instalando: $node_name"
+    clone_repo "$repo" "$CN_DIR/$node_name"
 done
-log_success "Custom nodes instalados"
 
-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
-    done
-    return $missing
-}
+log_success "Custom nodes instalados"
 
+# Iniciar
+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..."
 
+cd "$COMFY_DIR"
+
+# Habilita SageAttention se presente
 SAGE_FLAG=""
-if python3 - <<'PY' >/dev/null 2>&1; then
-import sageattention
-PY
+if python3 -c 'import sageattention' >/dev/null 2>&1; then
     SAGE_FLAG="--use-sage-attention"
 fi
 
-cd "$COMFY_DIR"
-log_info "Iniciando ComfyUI em http://$COMFY_HOST:$COMFY_PORT"
 exec comfy launch -- ${SAGE_FLAG:+$SAGE_FLAG} --listen "$COMFY_HOST" --port "$COMFY_PORT"