Update setup_comfyui_wan22.sh

setup_comfyui_wan22.sh

@@ -1,14 +1,29 @@
 #!/usr/bin/env bash
-# setup_comfyui_wan22.sh
+# setup_comfyui_wan22_optimized.sh
+# Script otimizado para instalação do ComfyUI com Wan 2.2
 # Requisitos: Linux + Python3 + git + (opcional) NVCC/CUDA para SageAttention
 
 set -Eeuo pipefail
-trap 'rc=$?; echo "[ERRO] linha $LINENO: \"$BASH_COMMAND\" saiu com código $rc"; exit $rc' ERR
+trap 'rc=$?; echo -e "\033[31m[ERRO]\033[0m Linha $LINENO: \"$BASH_COMMAND\" saiu com código $rc"; exit $rc' ERR
+
+# -----------------------------
+# Cores para melhor visualização
+# -----------------------------
+RED='\033[0;31m'
+GREEN='\033[0;32m'
+YELLOW='\033[1;33m'
+BLUE='\033[0;34m'
+NC='\033[0m' # No Color
+
+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"; }
 
 # -----------------------------
 # Configuráveis
 # -----------------------------
-COMFY_DIR="${COMFY_DIR:-$PWD/ComfyUI}"        # Onde ficará o ComfyUI
+COMFY_DIR="${COMFY_DIR:-$PWD/ComfyUI}"
 COMFY_HOST="${COMFY_HOST:-0.0.0.0}"
 COMFY_PORT="${COMFY_PORT:-8818}"
 
@@ -19,140 +34,329 @@ REPO_CN="xinsir/controlnet-union-sdxl-1.0"
 
 # Token do Civitai (pode sobrescrever via env)
 CIVITAI_TOKEN="${CIVITAI_TOKEN:-4fcb2834969399006a736ee402b061e5}"
-CIVITAI_URL="https://civitai.com/api/download/models/2122278?type=Model&format=SafeTensor&size=pruned&fp=fp16&token=${CIVITAI_TOKEN}"
 
-# Caminhos de destino dos modelos
-MODELS_DIR="$COMFY_DIR/models"
-mkdir -p "$MODELS_DIR"/{diffusion_models,loras,vae,text_encoders,clip_vision,controlnet}
+# 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"
+
+# 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"
 
-# Pasta temporária para os downloads
+# -----------------------------
+# Setup inicial
+# -----------------------------
+MODELS_DIR="$COMFY_DIR/models"
 TMPDL="$(mktemp -d)"
-cleanup() { rm -rf "$TMPDL"; }
+cleanup() { 
+    log_info "Limpando arquivos temporários..."
+    rm -rf "$TMPDL" 
+}
 trap cleanup EXIT
 
-need() { command -v "$1" >/dev/null 2>&1 || { echo "Falta dependência: $1"; exit 1; }; }
-need python3; need git; need wget
+# 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
+    fi
+}
+
+log_info "Verificando dependências..."
+need python3
+need git
+need wget
+need curl
+
+# Criar estrutura de diretórios
+log_info "Criando estrutura de diretórios..."
+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,civitai}
 
+# -----------------------------
+# Funções auxiliares
+# -----------------------------
 clone_or_update() {
-  local url="$1" dest="$2"
-  if [ -d "$dest/.git" ]; then
-    echo "   - Atualizando $(basename "$dest")..."
-    git -C "$dest" pull --ff-only || true
-  else
-    echo "   - Clonando $(basename "$dest")..."
-    git clone --recursive "$url" "$dest"
-  fi
-  if [ -f "$dest/requirements.txt" ]; then
-    echo "     > Instalando requirements de $(basename "$dest")..."
-    python3 -m pip install -r "$dest/requirements.txt" || true
-  fi
+    local url="$1" dest="$2"
+    local repo_name=$(basename "$dest")
+    
+    if [ -d "$dest/.git" ]; then
+        log_info "Atualizando $repo_name..."
+        git -C "$dest" pull --ff-only --quiet || true
+    else
+        log_info "Clonando $repo_name..."
+        git clone --recursive --depth 1 "$url" "$dest" --quiet
+    fi
+    
+    if [ -f "$dest/requirements.txt" ]; then
+        python3 -m pip install --quiet --no-warn-script-location -r "$dest/requirements.txt" 2>/dev/null || true
+    fi
 }
 
-echo "[1/6] Instalando comfy-cli..."
-python3 -m pip install -U pip wheel setuptools >/dev/null
-python3 -m pip install -U comfy-cli >/dev/null
+download_with_progress() {
+    local url="$1"
+    local output="$2"
+    local description="$3"
+    
+    log_info "Baixando: $description"
+    wget --progress=bar:force:noscroll \
+         --tries=3 \
+         --timeout=30 \
+         --continue \
+         -O "$output" \
+         "$url" 2>&1 | \
+         grep --line-buffered "%" | \
+         sed -u -e "s/^/    /"
+}
 
-echo "[2/6] Preparando/instalando ComfyUI (GPU NVIDIA, fast deps)..."
-comfy --skip-prompt install --fast-deps --nvidia
+# -----------------------------
+# Instalação principal
+# -----------------------------
+echo ""
+log_info "========================================="
+log_info "  ComfyUI + Wan 2.2 Setup Otimizado"
+log_info "  CPU Cores: $CPU_CORES | Jobs: $MAX_JOBS"
+log_info "========================================="
+echo ""
 
-mkdir -p "$MODELS_DIR" "$COMFY_DIR/custom_nodes"
+# [1/7] Atualizar pip e instalar ferramentas essenciais
+log_info "[1/7] Atualizando pip e instalando ferramentas..."
+python3 -m pip install --quiet --upgrade pip wheel setuptools uv >/dev/null 2>&1
+python3 -m pip install --quiet --upgrade comfy-cli "huggingface_hub>=0.23.0" hf_transfer aria2 >/dev/null 2>&1
+log_success "Ferramentas instaladas"
 
-echo "[3/6] Iniciando downloads dos modelos (HF Hub) em background..."
-python3 -m pip install -U "huggingface_hub>=0.23.0" hf_transfer >/dev/null 2>&1 || true
-export HF_HUB_ENABLE_HF_TRANSFER=1
+# [2/7] Instalar ComfyUI
+log_info "[2/7] Instalando ComfyUI (GPU NVIDIA, fast deps)..."
+if [ ! -d "$COMFY_DIR" ]; then
+    comfy --skip-prompt install --fast-deps --nvidia
+else
+    log_warn "ComfyUI já existe, pulando instalação base"
+fi
+log_success "ComfyUI pronto"
 
+# [3/7] Baixar modelos Wan 2.2 (em paralelo com aria2)
+log_info "[3/7] Iniciando downloads dos modelos Wan 2.2..."
 (
-  set -Eeuo pipefail
-  echo "   - Baixando pacotes do repo: $REPO_22"
-  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
-
-  echo "   - Baixando pacotes do repo: $REPO_21"
-  huggingface-cli download "$REPO_21" \
-    split_files/clip_vision/clip_vision_h.safetensors \
-    --local-dir "$TMPDL/wan21" \
-    --local-dir-use-symlinks False \
-    --resume-download
-
-  install -m 644 "$TMPDL/wan22/split_files/diffusion_models/"*.safetensors "$MODELS_DIR/diffusion_models/" 2>/dev/null || true
-  install -m 644 "$TMPDL/wan22/split_files/loras/"*.safetensors            "$MODELS_DIR/loras/"            2>/dev/null || true
-  install -m 644 "$TMPDL/wan22/split_files/vae/"*.safetensors              "$MODELS_DIR/vae/"              2>/dev/null || true
-  install -m 644 "$TMPDL/wan22/split_files/text_encoders/"*.safetensors    "$MODELS_DIR/text_encoders/"    2>/dev/null || true
-  install -m 644 "$TMPDL/wan21/split_files/clip_vision/"*.safetensors      "$MODELS_DIR/clip_vision/"      2>/dev/null || true
-
-  echo "   - Modelos principais baixados."
+    set -Eeuo pipefail
+    
+    # Wan 2.2
+    log_info "  → Baixando Wan 2.2 models..."
+    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 \
+        --quiet 2>/dev/null || true
+    
+    # Wan 2.1
+    log_info "  → Baixando Wan 2.1 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 \
+        --quiet 2>/dev/null || true
+    
+    # Copiar para destino
+    cp -f "$TMPDL/wan22/split_files/diffusion_models/"*.safetensors "$MODELS_DIR/diffusion_models/" 2>/dev/null || true
+    cp -f "$TMPDL/wan22/split_files/loras/"*.safetensors "$MODELS_DIR/loras/" 2>/dev/null || true
+    cp -f "$TMPDL/wan22/split_files/vae/"*.safetensors "$MODELS_DIR/vae/" 2>/dev/null || true
+    cp -f "$TMPDL/wan22/split_files/text_encoders/"*.safetensors "$MODELS_DIR/text_encoders/" 2>/dev/null || true
+    cp -f "$TMPDL/wan21/split_files/clip_vision/"*.safetensors "$MODELS_DIR/clip_vision/" 2>/dev/null || true
+    
+    log_success "  Modelos Wan baixados"
 ) &
-DL_PID=$!
+WAN_PID=$!
 
-echo "[4/6] Instalando SageAttention (em paralelo aos downloads)..."
-python3 -m pip install "triton>=3.0.0"
-if [ -d "SageAttention" ]; then
-  (cd SageAttention && git pull --ff-only)
-else
-  git clone https://github.com/thu-ml/SageAttention.git
-fi
-export EXT_PARALLEL=4 NVCC_APPEND_FLAGS="--threads 8" MAX_JOBS=32
+# [4/7] Baixar ControlNet (em paralelo)
+log_info "[4/7] 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 \
+        --quiet 2>/dev/null || true
+    
+    cp -f "$TMPDL/controlnet/diffusion_pytorch_model_promax.safetensors" \
+        "$MODELS_DIR/controlnet/controlnet-union.safetensors"
+    
+    log_success "  ControlNet baixado"
+) &
+CN_PID=$!
+
+# [5/7] Baixar modelo do Civitai (CORREÇÃO IMPORTANTE)
+log_info "[5/7] Baixando modelo do Civitai..."
+(
+    set -Eeuo pipefail
+    CIVITAI_URL="https://civitai.com/api/download/models/2122278"
+    CIVITAI_PARAMS="type=Model&format=SafeTensor&size=pruned&fp=fp16&token=${CIVITAI_TOKEN}"
+    FULL_URL="${CIVITAI_URL}?${CIVITAI_PARAMS}"
+    
+    cd "$MODELS_DIR/loras"
+    
+    # Usar curl com retry e progress bar
+    curl -L \
+         --retry 3 \
+         --retry-delay 5 \
+         --connect-timeout 30 \
+         --max-time 1800 \
+         --progress-bar \
+         -H "User-Agent: Mozilla/5.0" \
+         -o "civitai_model_temp.safetensors" \
+         "$FULL_URL"
+    
+    # Renomear com o nome correto do header Content-Disposition se possível
+    if [ -f "civitai_model_temp.safetensors" ]; then
+        # Tentar obter o nome real do arquivo
+        REAL_NAME=$(curl -sI -L "$FULL_URL" | grep -i content-disposition | sed -n 's/.*filename="\?\([^"]*\)"\?.*/\1/p' | head -1)
+        if [ -n "$REAL_NAME" ]; then
+            mv "civitai_model_temp.safetensors" "$REAL_NAME"
+            log_success "  Civitai modelo salvo como: $REAL_NAME"
+        else
+            mv "civitai_model_temp.safetensors" "civitai_2122278.safetensors"
+            log_success "  Civitai modelo salvo como: civitai_2122278.safetensors"
+        fi
+    fi
+) &
+CIVITAI_PID=$!
+
+# [6/7] Instalar custom nodes (em paralelo)
+log_info "[6/7] Instalando custom nodes..."
 (
-  set -Eeuo pipefail
-  cd SageAttention
-  python3 setup.py install
+    set -Eeuo pipefail
+    CN_DIR="$COMFY_DIR/custom_nodes"
+    
+    declare -A REPOS=(
+        ["https://github.com/kijai/ComfyUI-Florence2"]="$CN_DIR/ComfyUI-Florence2"
+        ["https://github.com/kijai/ComfyUI-WanVideoWrapper"]="$CN_DIR/ComfyUI-WanVideoWrapper"
+        ["https://github.com/Fannovel16/ComfyUI-Frame-Interpolation"]="$CN_DIR/ComfyUI-Frame-Interpolation"
+        ["https://github.com/kijai/ComfyUI-GIMM-VFI"]="$CN_DIR/ComfyUI-GIMM-VFI"
+        ["https://github.com/Kosinkadink/ComfyUI-VideoHelperSuite"]="$CN_DIR/ComfyUI-VideoHelperSuite"
+        ["https://github.com/kijai/ComfyUI-KJNodes"]="$CN_DIR/ComfyUI-KJNodes"
+        ["https://github.com/Fannovel16/comfyui_controlnet_aux"]="$CN_DIR/comfyui_controlnet_aux"
+        ["https://github.com/Artificial-Sweetener/comfyui-WhiteRabbit"]="$CN_DIR/comfyui-WhiteRabbit"
+        ["https://github.com/LucipherDev/ComfyUI-AniDoc"]="$CN_DIR/ComfyUI-AniDoc"
+        ["https://github.com/shiimizu/ComfyUI_smZNodes"]="$CN_DIR/ComfyUI_smZNodes"
+        ["https://github.com/yuvraj108c/ComfyUI-Dwpose-Tensorrt"]="$CN_DIR/ComfyUI-Dwpose-Tensorrt"
+        ["https://github.com/grmchn/ComfyUI-ProportionChanger"]="$CN_DIR/ComfyUI-ProportionChanger"
+    )
+    
+    for url in "${!REPOS[@]}"; do
+        clone_or_update "$url" "${REPOS[$url]}" &
+    done
+    wait
+    
+    log_success "  Custom nodes instalados"
 ) &
-SA_PID=$!
-
-# === Espera especificamente os downloads principais terminarem ===
-echo "[5/6] Aguardando downloads principais para etapa pós-downloads..."
-wait "$DL_PID"
-
-# === Pós-downloads em primeiro plano (enquanto o SageAttention ainda compila) ===
-echo "     > Baixando ControlNet extra e renomeando..."
-huggingface-cli download "$REPO_CN" diffusion_pytorch_model_promax.safetensors \
-  --local-dir "$TMPDL/controlnet" \
-  --local-dir-use-symlinks False \
-  --resume-download
-install -m 644 "$TMPDL/controlnet/diffusion_pytorch_model_promax.safetensors" \
-  "$MODELS_DIR/controlnet/controlnet-union.safetensors"
-
-echo "     > Baixando modelo do Civitai (após o ControlNet)..."
-CIVITAI_DIR="$MODELS_DIR/loras"
-mkdir -p "$CIVITAI_DIR"
+NODES_PID=$!
+
+# [7/7] Instalar SageAttention
+log_info "[7/7] Instalando SageAttention..."
 (
-  cd "$CIVITAI_DIR"
-  wget --content-disposition "$CIVITAI_URL"
-)
-
-echo "     > Instalando custom nodes..."
-CN_DIR="$COMFY_DIR/custom_nodes"
-mkdir -p "$CN_DIR"
-declare -A REPOS=(
-  ["https://github.com/kijai/ComfyUI-Florence2"]="$CN_DIR/ComfyUI-Florence2"
-  ["https://github.com/kijai/ComfyUI-WanVideoWrapper"]="$CN_DIR/ComfyUI-WanVideoWrapper"
-  ["https://github.com/Fannovel16/ComfyUI-Frame-Interpolation"]="$CN_DIR/ComfyUI-Frame-Interpolation"
-  ["https://github.com/kijai/ComfyUI-GIMM-VFI"]="$CN_DIR/ComfyUI-GIMM-VFI"
-  ["https://github.com/Kosinkadink/ComfyUI-VideoHelperSuite"]="$CN_DIR/ComfyUI-VideoHelperSuite"
-  ["https://github.com/kijai/ComfyUI-KJNodes"]="$CN_DIR/ComfyUI-KJNodes"
-  ["https://github.com/Fannovel16/comfyui_controlnet_aux"]="$CN_DIR/comfyui_controlnet_aux"
-  ["https://github.com/Artificial-Sweetener/comfyui-WhiteRabbit"]="$CN_DIR/comfyui-WhiteRabbit"
-  ["https://github.com/LucipherDev/ComfyUI-AniDoc"]="$CN_DIR/ComfyUI-AniDoc"
-  ["https://github.com/shiimizu/ComfyUI_smZNodes"]="$CN_DIR/ComfyUI_smZNodes"
-  ["https://github.com/yuvraj108c/ComfyUI-Dwpose-Tensorrt"]="$CN_DIR/ComfyUI-Dwpose-Tensorrt"
-  ["https://github.com/grmchn/ComfyUI-ProportionChanger"]="$CN_DIR/ComfyUI-ProportionChanger"
-)
-for url in "${!REPOS[@]}"; do
-  clone_or_update "$url" "${REPOS[$url]}"
+    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
+        
+        if [ -d "SageAttention" ]; then
+            (cd SageAttention && git pull --ff-only --quiet)
+        else
+            git clone --depth 1 https://github.com/thu-ml/SageAttention.git --quiet
+        fi
+        
+        cd SageAttention
+        python3 setup.py install --quiet 2>/dev/null || {
+            log_warn "  SageAttention falhou, continuando sem ele"
+        }
+        
+        log_success "  SageAttention instalado"
+    else
+        log_warn "  CUDA não detectado, pulando SageAttention"
+    fi
+) &
+SAGE_PID=$!
+
+# Aguardar todos os processos em paralelo
+log_info "Aguardando conclusão dos downloads e instalações..."
+echo -n "  Progresso: "
+
+PIDS=($WAN_PID $CN_PID $CIVITAI_PID $NODES_PID $SAGE_PID)
+TOTAL=${#PIDS[@]}
+COMPLETED=0
+
+while [ $COMPLETED -lt $TOTAL ]; do
+    COMPLETED=0
+    for pid in "${PIDS[@]}"; do
+        if ! kill -0 $pid 2>/dev/null; then
+            ((COMPLETED++))
+        fi
+    done
+    echo -ne "\r  Progresso: [$COMPLETED/$TOTAL]"
+    sleep 1
 done
-echo "     > Custom nodes prontos."
+echo ""
+
+# Verificar instalação
+log_info "Verificando instalação..."
+ERROR_COUNT=0
 
-# === Agora esperamos o SageAttention terminar, se ainda não terminou ===
-echo "[6/6] Aguardando final do SageAttention..."
-wait "$SA_PID"
+check_file() {
+    if [ -f "$1" ]; then
+        log_success "  ✓ $(basename "$1")"
+        return 0
+    else
+        log_error "  ✗ $(basename "$1") não encontrado"
+        ((ERROR_COUNT++))
+        return 1
+    fi
+}
 
-echo "[OK] Tudo pronto. Abrindo o ComfyUI com SageAttention na porta $COMFY_PORT..."
-exec comfy launch -- --use-sage-attention --listen "$COMFY_HOST" --port "$COMFY_PORT"
+# Verificar modelos principais
+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"
+
+# Relatório final
+echo ""
+log_info "========================================="
+if [ $ERROR_COUNT -eq 0 ]; then
+    log_success "Instalação concluída com sucesso!"
+    
+    # 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"
+    fi
+    
+    echo ""
+    log_info "Iniciando ComfyUI..."
+    log_info "URL: http://$(hostname -I | awk '{print $1}'):$COMFY_PORT"
+    echo ""
+    
+    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"
+    exit 1
+fi