Upload comfyui_setup_script.sh

comfyui_setup_script.sh

@@ -23,11 +23,26 @@ readonly NC='\033[0m' # No Color
 # Configurações
 readonly SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
 readonly LOG_FILE="${SCRIPT_DIR}/setup.log"
-readonly COMFY_DIR="/root/comfy"
+
+# Detectar diretório base automaticamente
+detect_base_dir() {
+    if [[ -d "/teamspace/studios" ]]; then
+        echo "/teamspace/studios/this_studio"
+    elif [[ -d "/workspace" ]]; then
+        echo "/workspace"
+    elif [[ $EUID -eq 0 ]]; then
+        echo "/root"
+    else
+        echo "$HOME"
+    fi
+}
+
+readonly BASE_DIR="$(detect_base_dir)"
+readonly COMFY_DIR="${BASE_DIR}/comfy"
 readonly TOKEN="0b698de03658fb8df2eeef90bad3ab9a"
 readonly LISTEN_PORT="8818"
 
-# URLs dos arquivos
+# URLs dos arquivos com nomes específicos
 declare -A DOWNLOADS=(
     ["checkpoint"]="https://civitai.com/api/download/models/1761560?type=Model&format=SafeTensor&size=pruned&fp=fp16&token=${TOKEN}"
     ["lora1"]="https://civitai.com/api/download/models/1426079?type=Model&format=SafeTensor&token=${TOKEN}"
@@ -46,6 +61,13 @@ declare -A DOWNLOAD_PATHS=(
     ["ultralytics_bbox_99coins"]="${COMFY_DIR}/ComfyUI/models/ultralytics/bbox"
 )
 
+# Nomes específicos para os arquivos
+declare -A DOWNLOAD_NAMES=(
+    ["snapshot"]="snapshotBRABEX.json"
+    ["upscaler_animesharpv3"]="2x-AnimeSharpV3.pth"
+    ["ultralytics_bbox_99coins"]="99coins_anime_girl_face_m_seg.pt"
+)
+
 # =============================================================================
 # FUNÇÕES UTILITÁRIAS
 # =============================================================================
@@ -102,12 +124,34 @@ install_system_packages() {
     log "STEP" "Instalando pacotes do sistema..."
     
     if command -v apt-get &> /dev/null; then
-        apt-get update -qq
-        apt-get install -y -qq aria2 wget curl
-        log "SUCCESS" "Pacotes do sistema instalados com sucesso"
-    else
-        log "WARN" "apt-get não encontrado, pulando instalação de pacotes do sistema"
+        # Tentar instalar com sudo se não for root
+        if [[ $EUID -ne 0 ]]; then
+            if command -v sudo &> /dev/null; then
+                log "INFO" "Tentando instalar com sudo..."
+                if sudo apt-get update -qq && sudo apt-get install -y -qq aria2 wget curl; then
+                    log "SUCCESS" "Pacotes do sistema instalados com sucesso (com sudo)"
+                    return 0
+                else
+                    log "WARN" "Falha ao instalar com sudo, tentando sem privilégios elevados..."
+                fi
+            else
+                log "WARN" "sudo não disponível, pulando instalação de pacotes do sistema"
+            fi
+        else
+            apt-get update -qq
+            apt-get install -y -qq aria2 wget curl
+            log "SUCCESS" "Pacotes do sistema instalados com sucesso"
+            return 0
+        fi
     fi
+    
+    # Verificar se aria2c já está disponível
+    if command -v aria2c &> /dev/null; then
+        log "SUCCESS" "aria2c já está disponível"
+        return 0
+    fi
+    
+    log "WARN" "aria2c não está disponível, usando apenas wget/curl para downloads"
 }
 
 install_comfy_cli() {
@@ -161,13 +205,16 @@ create_directories() {
     local dirs=(
         "${COMFY_DIR}/ComfyUI/models/checkpoints"
         "${COMFY_DIR}/ComfyUI/models/loras"
-        "${COMFY_DIR}/ComfyUI/models/upscale_models"  # Diretório adicionado
-        "${COMFY_DIR}/ComfyUI/models/ultralytics/bbox" # Diretório adicionado
+        "${COMFY_DIR}/ComfyUI/models/upscale_models"
+        "${COMFY_DIR}/ComfyUI/models/ultralytics/bbox"
     )
     
     for dir in "${dirs[@]}"; do
-        mkdir -p "$dir"
-        log "INFO" "Diretório criado: $dir"
+        if mkdir -p "$dir" 2>/dev/null; then
+            log "INFO" "Diretório criado: $dir"
+        else
+            log "WARN" "Não foi possível criar: $dir (pode já existir)"
+        fi
     done
 }
 
@@ -185,32 +232,60 @@ download_files() {
         
         log "INFO" "Baixando $name..."
         
-        if [[ "$name" == "snapshot" ]]; then
-            # Usar wget para o snapshot
-            if wget --quiet --show-progress --progress=bar:force:noscroll \
-                   --timeout=30 --tries=3 \
-                   "$url" -P "$path"; then
+        # Determinar nome do arquivo
+        local filename=""
+        if [[ -n "${DOWNLOAD_NAMES[$name]:-}" ]]; then
+            filename="${DOWNLOAD_NAMES[$name]}"
+        fi
+        
+        # Escolher método de download baseado na disponibilidade
+        if command -v aria2c &> /dev/null && [[ "$name" != "snapshot" ]]; then
+            # Usar aria2c para modelos grandes (exceto snapshot)
+            local aria_opts=("--console-log-level=warn" "--summary-interval=10" "--download-result=hide"
+                           "-c" "-s" "16" "-x" "16" "-k" "1M" "--timeout=60" "--retry-wait=5" "--max-tries=3"
+                           "-d" "$path")
+            
+            if [[ -n "$filename" ]]; then
+                aria_opts+=("-o" "$filename")
+            fi
+            
+            if aria2c "${aria_opts[@]}" "$url"; then
                 log "SUCCESS" "Download concluído: $name"
                 return 0
             else
-                log "ERROR" "Falha no download: $name"
-                return 1
+                log "WARN" "aria2c falhou para $name, tentando com wget..."
+            fi
+        fi
+        
+        # Fallback para wget
+        local wget_opts=("--quiet" "--show-progress" "--progress=bar:force:noscroll"
+                        "--timeout=30" "--tries=3" "--continue")
+        
+        if [[ -n "$filename" ]]; then
+            if wget "${wget_opts[@]}" "$url" -O "${path}/${filename}"; then
+                log "SUCCESS" "Download concluído: $name (wget)"
+                return 0
             fi
         else
-            # Usar aria2c para modelos grandes
-            if aria2c --console-log-level=warn \
-                     --summary-interval=10 \
-                     --download-result=hide \
-                     -c -s 16 -x 16 -k 1M \
-                     --timeout=60 --retry-wait=5 --max-tries=3 \
-                     -d "$path" "$url"; then
-                log "SUCCESS" "Download concluído: $name"
+            if wget "${wget_opts[@]}" "$url" -P "$path"; then
+                log "SUCCESS" "Download concluído: $name (wget)"
                 return 0
-            else
-                log "ERROR" "Falha no download: $name"
-                return 1
             fi
         fi
+        
+        # Último recurso: curl
+        log "WARN" "wget falhou para $name, tentando com curl..."
+        local output_file="${path}/${filename:-$(basename "$url" | cut -d'?' -f1)}"
+        
+        if curl -L --fail --silent --show-error --retry 3 --retry-delay 5 \
+               --connect-timeout 30 --max-time 1800 \
+               -o "$output_file" "$url"; then
+            log "SUCCESS" "Download concluído: $name (curl)"
+            return 0
+        else
+            log "ERROR" "Falha no download: $name"
+            return 1
+        fi
     }
     
     # Iniciar downloads em paralelo
@@ -347,6 +422,9 @@ main() {
         fi
     fi
     
+    log "INFO" "Diretório base detectado: $BASE_DIR"
+    log "INFO" "ComfyUI será instalado em: $COMFY_DIR"
+    
     # Executar etapas da instalação
     local steps=(
         "check_dependencies"
@@ -398,10 +476,12 @@ trap 'log "ERROR" "Script interrompido pelo usuário"; exit 1' INT TERM
 # EXECUÇÃO
 # =============================================================================
 
-# Verificar se está sendo executado como root (recomendado para RunPod/VastAI)
+# Verificar se está sendo executado como root (informativo apenas)
 if [[ $EUID -ne 0 ]]; then
-    log "WARN" "Script não está sendo executado como root"
-    log "WARN" "Algumas operações podem falhar"
+    log "INFO" "Script sendo executado como usuário não-root (normal para Lightning AI)"
+    log "INFO" "Usando diretório base: $BASE_DIR"
+else
+    log "INFO" "Script sendo executado como root"
 fi
 
 # Executar função principal