Update setup.py

setup.py

@@ -1,63 +1,133 @@
 # setup.py
-#
-# Copyright (C) August 4, 2025  Carlos Rodrigues dos Santos
-#
-# Versão 2.0.0 (Clonagem Anônima e Robusta)
-# - Usa URLs HTTPS explícitas e anônimas para evitar que o Git tente
-#   usar credenciais em cache desnecessariamente para repositórios públicos.
 
 import os
 import subprocess
 import sys
 from pathlib import Path
+import yaml
+
+# Adiciona a importação do hf_hub_download
+from huggingface_hub import hf_hub_download
 
 # --- Configuração ---
 DEPS_DIR = Path("/data")
+MODELS_ROOT = DEPS_DIR / "ltx_models" # Diretório centralizado para modelos
 
-# URLs explícitas e anônimas para os repositórios públicos
+# URLs explícitas para os repositórios
 REPOS_TO_CLONE = {
     "LTX-Video": "https://huggingface.co/spaces/Lightricks/ltx-video-distilled",
     "SeedVR_Space": "https://huggingface.co/spaces/ByteDance-Seed/SeedVR2-3B",
     "MMAudio": "https://github.com/hkchengrex/MMAudio.git"
 }
 
-def run_command(command, cwd=None):
-    """Executa um comando no terminal e lida com erros."""
+def run_command(command, cwd=None, is_pip=False):
+    """Executa um comando no terminal e exibe a saída em tempo real."""
+    if is_pip:
+        command = [sys.executable, "-m", "pip"] + command
+        
     print(f"Executando: {' '.join(command)}")
     try:
-        # Redireciona o stdin para DEVNULL para garantir que o git não tente pedir senha
-        subprocess.run(
-            command, 
-            check=True, 
-            cwd=cwd, 
-            stdin=subprocess.DEVNULL,
-        )
+        with subprocess.Popen(
+            command, cwd=cwd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
+            text=True, encoding='utf-8', bufsize=1
+        ) as process:
+            for line in iter(process.stdout.readline, ''):
+                print(line, end='')
+            
+            return_code = process.wait()
+            if return_code != 0:
+                raise subprocess.CalledProcessError(return_code, command)
     except subprocess.CalledProcessError as e:
         print(f"ERRO: O comando falhou com o código de saída {e.returncode}")
-        # stderr é capturado automaticamente se check=True falhar
-        print(f"Stderr: {e.stderr}") 
         sys.exit(1)
     except FileNotFoundError:
-        print(f"ERRO: O comando '{command[0]}' não foi encontrado. Certifique-se de que o git está instalado e no seu PATH.")
+        print(f"ERRO: O comando '{command[0]}' não foi encontrado.")
         sys.exit(1)
 
-def main():
-    print("--- Iniciando Setup do Ambiente ADUC-SDR ---")
-    
-    DEPS_DIR.mkdir(exist_ok=True)
-    
+def clone_repos():
+    """ETAPA 1: Garante que todos os repositórios de código necessários existam."""
+    print("--- ETAPA 1: Verificando e clonando repositórios de código ---")
     for repo_name, repo_url in REPOS_TO_CLONE.items():
         repo_path = DEPS_DIR / repo_name
-        if repo_path.exists():
-            print(f"Repositório '{repo_name}' já existe. Pulando a clonagem.")
+        if repo_path.exists() and (repo_path / ".git").exists():
+            print(f"Repositório '{repo_name}' já existe. Pulando.")
         else:
             print(f"Clonando '{repo_name}' de {repo_url}...")
-            run_command(["git", "clone", "--depth", "1", repo_url, str(repo_path)])
+            run_command(["git", "clone", "--depth", "1", "-q", repo_url, str(repo_path)])
             print(f"'{repo_name}' clonado com sucesso.")
 
-    print("\n--- Setup do Ambiente Concluído com Sucesso! ---")
-    print("Você agora pode iniciar a aplicação principal (ex: python app.py).")
+def install_ltx_dependencies():
+    """ETAPA 2: Instala as dependências Python do LTX-Video."""
+    print("\n--- ETAPA 2: Instalando dependências Python do LTX-Video ---")
+    ltx_repo_path = DEPS_DIR / "LTX-Video"
+    if ltx_repo_path.exists():
+        run_command(["install", "-e", ".[inference]"], cwd=str(ltx_repo_path), is_pip=True)
+        print("Dependências do LTX-Video instaladas com sucesso.")
+    else:
+        print("ERRO CRÍTICO: Diretório LTX-Video não encontrado. Não foi possível instalar dependências.")
+        sys.exit(1)
+
+def download_ltx_models():
+    """
+    ETAPA 3: Verifica e baixa os modelos LTX necessários, lendo os configs
+    do repositório já clonado.
+    """
+    print("\n--- ETAPA 3: Verificando e baixando modelos LTX ---")
+    LTX_REPO_ID = "Lightricks/LTX-Video" # Repositório dos modelos no Hub
+    ltx_configs_dir = DEPS_DIR / "LTX-Video" / "configs"
+    
+    if not ltx_configs_dir.is_dir():
+        print(f"ERRO CRÍTICO: Diretório de configs '{ltx_configs_dir}' não encontrado. O clone do LTX-Video pode ter falhado.")
+        sys.exit(1)
+
+    required_model_files = set()
+    for cfg_path in ltx_configs_dir.glob("*.yaml"):
+        try:
+            with open(cfg_path, 'r') as f:
+                config = yaml.safe_load(f)
+                if config and config.get("checkpoint_path"):
+                    required_model_files.add(config["checkpoint_path"])
+                if config and config.get("spatial_upscaler_model_path"):
+                    required_model_files.add(config["spatial_upscaler_model_path"])
+        except Exception as e:
+            print(f"AVISO: Falha ao ler o arquivo de configuração {cfg_path}: {e}")
+
+    if not required_model_files:
+        print("AVISO: Nenhum arquivo de modelo foi encontrado nos arquivos de configuração. Verifique os configs.")
+        return
+
+    print(f"Modelos necessários encontrados: {required_model_files}")
+    MODELS_ROOT.mkdir(exist_ok=True)
+    
+    for model_filename in required_model_files:
+        local_path = MODELS_ROOT / model_filename
+        if local_path.exists():
+            print(f"Modelo '{model_filename}' já existe localmente. Pulando.")
+        else:
+            print(f"Baixando modelo '{model_filename}'...")
+            try:
+                hf_hub_download(
+                    repo_id=LTX_REPO_ID,
+                    filename=model_filename,
+                    local_dir=str(MODELS_ROOT),
+                    local_dir_use_symlinks=False,
+                    token=os.getenv("HF_TOKEN")
+                )
+                print(f"'{model_filename}' baixado com sucesso.")
+            except Exception as e:
+                print(f"ERRO: Falha ao baixar '{model_filename}'. Verifique o nome do arquivo e a conexão. Erro: {e}")
+
+def main():
+    print("--- Iniciando Setup Completo do Ambiente ADUC-SDR ---")
+    DEPS_DIR.mkdir(exist_ok=True)
+    
+    # Executa as etapas na ordem correta
+    clone_repos()
+    install_ltx_dependencies()
+    download_ltx_models()
+
+    print("\n--- ✅ Setup do Ambiente Concluído com Sucesso! ---")
+    print("Todos os repositórios, dependências e modelos estão prontos.")
 
 if __name__ == "__main__":
-    main()
-    
+    main()