Update setup_comfyui_wan22.sh

setup_comfyui_wan22.sh

@@ -1,6 +1,6 @@
 #!/usr/bin/env bash
-# setup_comfyui_wan22_fixed.sh
-# Script corrigido para instalação do ComfyUI com Wan 2.2
+# setup_comfyui_wan22_final.sh
+# Script final corrigido para instalação do ComfyUI com Wan 2.2
 # Requisitos: Linux + Python3 + git + (opcional) NVCC/CUDA para SageAttention
 
 set -Eeuo pipefail
@@ -23,7 +23,8 @@ log_error() { echo -e "${RED}[✗]${NC} $1"; }
 # -----------------------------
 # Configuráveis
 # -----------------------------
-COMFY_DIR="${COMFY_DIR:-$PWD/ComfyUI}"
+# IMPORTANTE: comfy-cli SEMPRE instala em /root/comfy/ComfyUI
+COMFY_DIR="/root/comfy/ComfyUI"  # Caminho fixo onde comfy-cli instala
 COMFY_HOST="${COMFY_HOST:-0.0.0.0}"
 COMFY_PORT="${COMFY_PORT:-8818}"
 
@@ -37,16 +38,30 @@ CIVITAI_TOKEN="${CIVITAI_TOKEN:-4fcb2834969399006a736ee402b061e5}"
 
 # Configurações de performance
 export HF_HUB_ENABLE_HF_TRANSFER=1
-export HF_TRANSFER_CONCURRENCY=8
-export UV_SYSTEM_PYTHON=1
-export PYTORCH_CUDA_ALLOC_CONF="expandable_segments:True"
+export HF_TRANSFER_CONCURRENCY=16  # Aumentado para melhor performance
 
 # Detecção de CPU cores para compilação
 CPU_CORES=$(nproc 2>/dev/null || sysctl -n hw.ncpu 2>/dev/null || echo 4)
-HALF_CORES=$((CPU_CORES / 2))
-export MAX_JOBS=$HALF_CORES
-export EXT_PARALLEL=$HALF_CORES
-export NVCC_APPEND_FLAGS="--threads 8"
+
+# Otimização para máquinas com muitos cores (como 112 cores)
+if [ $CPU_CORES -gt 32 ]; then
+    # Para máquinas high-end, usar configurações agressivas
+    export MAX_JOBS=32  # Limite do CUDA
+    export EXT_PARALLEL=16  # Aumentado
+    export NVCC_APPEND_FLAGS="--threads 16"  # Aumentado
+    COMPILE_JOBS=32
+    log_info "Detectado sistema high-end: $CPU_CORES cores, usando MAX_JOBS=$MAX_JOBS"
+else
+    # Para máquinas normais
+    HALF_CORES=$((CPU_CORES / 2))
+    export MAX_JOBS=$HALF_CORES
+    export EXT_PARALLEL=4
+    export NVCC_APPEND_FLAGS="--threads 8"
+    COMPILE_JOBS=$HALF_CORES
+fi
+
+export UV_SYSTEM_PYTHON=1
+export PYTORCH_CUDA_ALLOC_CONF="expandable_segments:True"
 
 # -----------------------------
 # Setup inicial
@@ -93,68 +108,117 @@ clone_or_update() {
     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
+    if python3 -c "import hf_transfer" 2>/dev/null; then
+        log_info "  → Usando HF Transfer (rápido) para $repo..."
+        HF_HUB_ENABLE_HF_TRANSFER=1 huggingface-cli download "$repo" \
+            "${files[@]}" \
+            --local-dir "$output_dir" \
+            --local-dir-use-symlinks False \
+            --resume-download 2>/dev/null || {
+            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
+        }
+    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 Corrigido"
-log_info "  CPU Cores: $CPU_CORES | Jobs: $MAX_JOBS"
+log_info "  ComfyUI + Wan 2.2 Setup Final"
+log_info "  CPU Cores: $CPU_CORES | Compile Jobs: $COMPILE_JOBS"
+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..."
 python3 -m pip install --quiet --upgrade pip wheel setuptools >/dev/null 2>&1
-python3 -m pip install --quiet --upgrade comfy-cli "huggingface_hub>=0.23.0" hf_transfer aria2 >/dev/null 2>&1
+
+# Instalar HF Transfer primeiro para downloads mais rápidos
+python3 -m pip install --quiet --upgrade hf_transfer >/dev/null 2>&1 || true
+
+# Depois instalar o resto
+python3 -m pip install --quiet --upgrade comfy-cli "huggingface_hub>=0.23.0" aria2 >/dev/null 2>&1
 log_success "Ferramentas instaladas"
 
-# [2/8] Instalar ComfyUI - SEMPRE PRIMEIRO!
+# [2/8] Instalar ComfyUI
 log_info "[2/8] Instalando ComfyUI (GPU NVIDIA, fast deps)..."
 
-# Verificação correta se ComfyUI já existe
+# 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 em $COMFY_DIR..."
-    # Remove diretório parcialmente criado se existir
+    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
+    # Instala ComfyUI (sempre vai para /root/comfy/ComfyUI)
     comfy --skip-prompt install --fast-deps --nvidia
     
-    # Verifica se instalou corretamente
+    # Verifica se instalou no local correto
     if [ ! -f "$COMFY_DIR/main.py" ]; then
-        log_error "Falha na instalação do ComfyUI!"
-        exit 1
+        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
+        fi
     fi
 fi
 
-# Criar estrutura de diretórios APÓS garantir que ComfyUI existe
+# Criar estrutura de diretórios
 MODELS_DIR="$COMFY_DIR/models"
-log_info "Criando estrutura de diretórios..."
+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 e estrutura criada"
+log_success "ComfyUI pronto em: $COMFY_DIR"
 
 # [3/8] Baixar modelos Wan 2.2
-log_info "[3/8] Iniciando downloads dos modelos Wan 2.2..."
+log_info "[3/8] Iniciando download dos modelos Wan 2.2..."
 (
     set -Eeuo pipefail
     
-    log_info "  → Baixando Wan 2.2 models (isso pode demorar)..."
-    huggingface-cli download "$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 \
-        --local-dir "$TMPDL/wan22" \
-        --local-dir-use-symlinks False \
-        --resume-download
+    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
     cp -f "$TMPDL/wan22/split_files/diffusion_models/"*.safetensors "$MODELS_DIR/diffusion_models/" 2>/dev/null || true
@@ -171,12 +235,8 @@ log_info "[4/8] Baixando Wan 2.1 clip vision..."
 (
     set -Eeuo pipefail
     
-    log_info "  → Baixando clip vision..."
-    huggingface-cli download "$REPO_21" \
-        split_files/clip_vision/clip_vision_h.safetensors \
-        --local-dir "$TMPDL/wan21" \
-        --local-dir-use-symlinks False \
-        --resume-download
+    hf_download_with_fallback "$REPO_21" \
+        "split_files/clip_vision/clip_vision_h.safetensors"
     
     cp -f "$TMPDL/wan21/split_files/clip_vision/"*.safetensors "$MODELS_DIR/clip_vision/" 2>/dev/null || true
     
@@ -189,11 +249,8 @@ log_info "[5/8] Baixando ControlNet Union..."
 (
     set -Eeuo pipefail
     
-    huggingface-cli download "$REPO_CN" \
-        diffusion_pytorch_model_promax.safetensors \
-        --local-dir "$TMPDL/controlnet" \
-        --local-dir-use-symlinks False \
-        --resume-download
+    hf_download_with_fallback "$REPO_CN" \
+        "diffusion_pytorch_model_promax.safetensors"
     
     cp -f "$TMPDL/controlnet/diffusion_pytorch_model_promax.safetensors" \
         "$MODELS_DIR/controlnet/controlnet-union.safetensors"
@@ -209,37 +266,73 @@ log_info "[6/8] Baixando modelo do Civitai..."
     
     cd "$MODELS_DIR/loras"
     
-    # URL correta sem aspas que quebram os parâmetros
+    # URL correta completa
     CIVITAI_URL="https://civitai.com/api/download/models/2122278?type=Model&format=SafeTensor&size=pruned&fp=fp16&token=${CIVITAI_TOKEN}"
     
-    # Usar wget com content-disposition para pegar o nome correto
-    wget --content-disposition \
-         --tries=3 \
-         --timeout=60 \
-         --continue \
-         --progress=bar:force:noscroll \
-         "$CIVITAI_URL" 2>&1 | grep --line-buffered "%" | sed -u -e "s/^/    /" || {
-        log_warn "  Tentando com curl como fallback..."
+    # Tentar com wget primeiro (melhor para nomes de arquivo)
+    if ! wget --content-disposition \
+              --tries=3 \
+              --timeout=60 \
+              --continue \
+              --progress=bar:force:noscroll \
+              "$CIVITAI_URL" 2>&1 | grep --line-buffered "%" | sed -u -e "s/^/    /"; then
+        
+        log_warn "  Wget falhou, tentando com curl..."
         curl -L \
              --retry 3 \
              --connect-timeout 30 \
              --progress-bar \
              -o "flux_unleashed_v2.safetensors" \
              "$CIVITAI_URL"
-    }
+    fi
     
     log_success "  Modelo Civitai baixado"
 ) &
 CIVITAI_PID=$!
 
-# [7/8] Aguardar downloads principais antes de instalar custom nodes
-log_info "[7/8] Aguardando downloads principais..."
-wait $WAN22_PID $WAN21_PID $CN_PID $CIVITAI_PID
+# [7/8] Instalar SageAttention (em paralelo com downloads)
+log_info "[7/8] Preparando SageAttention (paralelo com downloads)..."
+(
+    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
+        python3 -m pip install --quiet "triton>=3.0.0" 2>/dev/null || true
+        
+        # Clonar SageAttention
+        SAGE_DIR="$TMPDL/SageAttention"
+        if [ ! -d "$SAGE_DIR" ]; then
+            git clone --depth 1 https://github.com/thu-ml/SageAttention.git "$SAGE_DIR" --quiet
+        fi
+        
+        cd "$SAGE_DIR"
+        
+        # Compilar com configurações otimizadas para muitos cores
+        log_info "  Compilando SageAttention com $MAX_JOBS jobs paralelos..."
+        python3 setup.py install 2>&1 | grep -E "(Building|Compiling|Installing)" || true
+        
+        # Testar se instalou
+        if python3 -c "import sageattention" 2>/dev/null; then
+            log_success "  SageAttention instalado com sucesso"
+        else
+            log_warn "  SageAttention não foi instalado corretamente"
+        fi
+    else
+        log_warn "  CUDA não detectado, pulando SageAttention"
+    fi
+) &
+SAGE_PID=$!
 
-# [8/8] Agora sim instalar custom nodes (ComfyUI já existe com certeza)
-log_info "[8/8] Instalando custom nodes e SageAttention..."
+# Aguardar downloads principais
+log_info "Aguardando downloads principais (em paralelo com SageAttention)..."
+wait $WAN22_PID $WAN21_PID $CN_PID $CIVITAI_PID
 
-# Custom nodes
+# [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"
@@ -265,51 +358,19 @@ log_info "[8/8] Instalando custom nodes e SageAttention..."
         ["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"
+        clone_or_update "$url" "$node_path" &
     done
+    wait
     
     log_success "  Custom nodes instalados"
-) &
-NODES_PID=$!
-
-# SageAttention
-(
-    set -Eeuo pipefail
-    
-    # Verificar se CUDA está disponível
-    if command -v nvcc >/dev/null 2>&1; then
-        log_info "  CUDA detectado, compilando SageAttention..."
-        
-        python3 -m pip install --quiet "triton>=3.0.0" 2>/dev/null || true
-        
-        # Clonar no diretório temporário primeiro
-        SAGE_DIR="$TMPDL/SageAttention"
-        if [ ! -d "$SAGE_DIR" ]; then
-            git clone --depth 1 https://github.com/thu-ml/SageAttention.git "$SAGE_DIR" --quiet
-        fi
-        
-        cd "$SAGE_DIR"
-        python3 setup.py install 2>/dev/null || {
-            log_warn "  SageAttention falhou na compilação, continuando sem ele"
-        }
-        
-        # Testar se instalou
-        if python3 -c "import sageattention" 2>/dev/null; then
-            log_success "  SageAttention instalado com sucesso"
-        else
-            log_warn "  SageAttention não foi instalado corretamente"
-        fi
-    else
-        log_warn "  CUDA não detectado, pulando SageAttention"
-    fi
-) &
-SAGE_PID=$!
+)
 
-# Aguardar custom nodes e SageAttention
-wait $NODES_PID $SAGE_PID
+# Aguardar SageAttention terminar
+wait $SAGE_PID
 
 # Verificar instalação
 log_info "Verificando instalação..."
@@ -351,14 +412,15 @@ echo ""
 log_info "========================================="
 if [ $ERROR_COUNT -eq 0 ]; then
     log_success "Instalação concluída com sucesso!"
+    log_info "ComfyUI instalado em: $COMFY_DIR"
     
     # Verificar se SageAttention está disponível
     SAGE_FLAG=""
     if python3 -c "import sageattention" 2>/dev/null; then
         SAGE_FLAG="--use-sage-attention"
-        log_success "SageAttention disponível"
+        log_success "SageAttention disponível (compilado com $MAX_JOBS jobs)"
     else
-        log_info "SageAttention não disponível (normal se não tiver CUDA)"
+        log_info "SageAttention não disponível"
     fi
     
     echo ""
@@ -372,7 +434,7 @@ if [ $ERROR_COUNT -eq 0 ]; then
     log_info "URL: http://localhost:$COMFY_PORT"
     echo ""
     
-    # Mudar para o diretório do ComfyUI antes de executar
+    # Mudar para o diretório do ComfyUI e executar
     cd "$COMFY_DIR"
     exec comfy launch -- $SAGE_FLAG --listen "$COMFY_HOST" --port "$COMFY_PORT"
 else