Update setup.py

setup.py

@@ -1,135 +1,181 @@
 # setup.py
+#
+# Copyright (C) August 4, 2025  Carlos Rodrigues dos Santos
+#
+# Versão 2.3.0 (Setup Robusto e Idempotente)
+# - Verifica a existência de repositórios e arquivos de modelo antes de baixar.
+# - Pula downloads se os artefatos já existirem, sem gerar erros.
+# - Unifica o download de todas as dependências (Git, LTX Models, SeedVR Models).
 
 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 ---
+# --- Configuração Geral ---
 DEPS_DIR = Path("/data")
-MODELS_ROOT = DEPS_DIR / "ltx_models" # Diretório centralizado para modelos
 
-# URLs explícitas para os repositórios
+# --- Configuração Específica LTX-Video ---
+LTX_VIDEO_REPO_DIR = DEPS_DIR / "LTX-Video"
+
+# --- Configuração Específica SeedVR ---
+SEEDVR_MODELS_DIR = DEPS_DIR / "SeedVR"
+
+# --- Repositórios para Clonar ---
 REPOS_TO_CLONE = {
     "LTX-Video": "https://huggingface.co/spaces/Lightricks/ltx-video-distilled",
-    "SeedVR_Space": "https://huggingface.co/spaces/ByteDance-Seed/SeedVR2-3B",
+    "SeedVR": "https://huggingface.co/spaces/ByteDance-Seed/SeedVR2-3B",
     "MMAudio": "https://github.com/hkchengrex/MMAudio.git"
 }
 
-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
-        
+def run_command(command, cwd=None):
+    """Executa um comando no terminal e lida com erros."""
     print(f"Executando: {' '.join(command)}")
     try:
-        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)
+        subprocess.run(
+            command,
+            check=True,
+            cwd=cwd,
+            stdin=subprocess.DEVNULL,
+        )
     except subprocess.CalledProcessError as e:
-        print(f"ERRO: O comando falhou com o código de saída {e.returncode}")
+        print(f"ERRO: O comando falhou com o código de saída {e.returncode}\nStderr: {e.stderr}")
         sys.exit(1)
     except FileNotFoundError:
-        print(f"ERRO: O comando '{command[0]}' não foi encontrado.")
-        sys.exit(1)
-
-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() 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", "-q", repo_url, str(repo_path)])
-            print(f"'{repo_name}' clonado com sucesso.")
-
-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.")
+        print(f"ERRO: O comando '{command[0]}' não foi encontrado. Certifique-se de que o git está instalado e no seu PATH.")
         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"
-    ltx_configs_dir = "/app/configs"
-
+# --- Funções de Download (LTX-Video) ---
+
+def _load_ltx_config():
+    """Carrega o arquivo de configuração YAML do LTX-Video."""
+    print("--- Carregando Configuração do LTX-Video ---")
+    base = LTX_VIDEO_REPO_DIR / "configs"
+    candidates = [
+        base / "ltxv-13b-0.9.8-dev-fp8.yaml",
+        base / "ltxv-13b-0.9.8-distilled-fp8.yaml",
+        base / "ltxv-13b-0.9.8-distilled.yaml",
+    ]
+    for cfg_path in candidates:
+        if cfg_path.exists():
+            print(f"Configuração encontrada: {cfg_path}")
+            with open(cfg_path, "r") as file:
+                return yaml.safe_load(file)
+    
+    fallback_path = base / "ltxv-13b-0.9.8-distilled-fp8.yaml"
+    print(f"AVISO: Nenhuma configuração preferencial encontrada. Usando fallback: {fallback_path}")
+    if not fallback_path.exists():
+        print(f"ERRO: Arquivo de configuração fallback '{fallback_path}' não encontrado.")
+        return None
+        
+    with open(fallback_path, "r") as file:
+        return yaml.safe_load(file)
 
-    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.")
+def _download_ltx_models(config):
+    """Baixa os modelos principais do LTX-Video, pulando os que já existem."""
+    print("\n--- Verificando Modelos do LTX-Video ---")
+    LTX_REPO = "Lightricks/LTX-Video"
+    
+    if "checkpoint_path" not in config or "spatial_upscaler_model_path" not in config:
+        print("ERRO: Chaves de modelo não encontradas no arquivo de configuração do LTX.")
         sys.exit(1)
 
-    required_model_files = set()
-    for cfg_path in ltx_configs_dir.glob("*.yaml"):
+    models_to_download = {
+        config["checkpoint_path"]: "checkpoint principal",
+        config["spatial_upscaler_model_path"]: "upscaler espacial"
+    }
+    
+    # O hf_hub_download já verifica o cache, mas vamos verificar o diretório final para clareza
+    # e para garantir que a lógica seja explícita.
+    for filename, description in models_to_download.items():
+        # A biblioteca huggingface_hub gerencia o local exato, então confiamos nela.
+        # A verificação aqui é para garantir que o download seja tentado.
+        print(f"Garantindo a existência do {description}: {filename}...")
         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"])
+            hf_hub_download(
+                repo_id=LTX_REPO, filename=filename,
+                local_dir=os.getenv("HF_HOME"), cache_dir=os.getenv("HF_HOME_CACHE"), token=os.getenv("HF_TOKEN")
+            )
+            print(f"{description.capitalize()} está disponível.")
         except Exception as e:
-            print(f"AVISO: Falha ao ler o arquivo de configuração {cfg_path}: {e}")
+            print(f"ERRO ao baixar o {description}: {e}")
+            sys.exit(1)
 
-    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)
+def _download_seedvr_models():
+    """Baixa os modelos do SeedVR, pulando os que já existem."""
+    print(f"\n--- Verificando Checkpoints do SeedVR em {SEEDVR_MODELS_DIR} ---")
+    SEEDVR_MODELS_DIR.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}'...")
+    model_files = {
+        "seedvr2_ema_7b_fp16.safetensors": "MonsterMMORPG/SeedVR2_SECourses",
+        "seedvr2_ema_7b_sharp_fp16.safetensors": "MonsterMMORPG/SeedVR2_SECourses",
+        "seedvr2_ema_3b_fp16.safetensors": "MonsterMMORPG/SeedVR2_SECourses",
+        "ema_vae_fp16.safetensors": "MonsterMMORPG/SeedVR2_SECourses",
+        "pos_emb.pt": "ByteDance-Seed/SeedVR2-3B",
+        "neg_emb.pt": "ByteDance-Seed/SeedVR2-3B"
+    }
+    
+    for filename, repo_id in model_files.items():
+        local_path = SEEDVR_MODELS_DIR / filename
+        if not local_path.is_file(): # Verifica se é um arquivo
+            print(f"Baixando {filename} de {repo_id}...")
             try:
                 hf_hub_download(
-                    repo_id=LTX_REPO_ID,
-                    filename=model_filename,
-                    local_dir=str(MODELS_ROOT),
+                    repo_id=repo_id,
+                    filename=filename,
+                    local_dir=str(SEEDVR_MODELS_DIR),
                     cache_dir=os.getenv("HF_HOME_CACHE"),
                     token=os.getenv("HF_TOKEN"),
+                    local_dir_use_symlinks=False
                 )
-                print(f"'{model_filename}' baixado com sucesso.")
+                print(f"'{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}")
+                print(f"ERRO ao baixar o modelo SeedVR '{filename}': {e}")
+                sys.exit(1)
+        else:
+            print(f"Arquivo '{filename}' já existe. Pulando.")
+    print("Checkpoints do SeedVR estão no local correto.")
+
+# --- Função Principal ---
 
 def main():
-    print("--- Iniciando Setup Completo do Ambiente ADUC-SDR ---")
+    print("--- Iniciando Setup do Ambiente ADUC-SDR (Versão Robusta) ---")
     DEPS_DIR.mkdir(exist_ok=True)
     
-    # Executa as etapas na ordem correta
-    clone_repos()
-    #install_ltx_dependencies()
-    download_ltx_models()
+    # --- ETAPA 1: Clonar Repositórios ---
+    print("\n--- ETAPA 1: Clonando Repositórios Git ---")
+    for repo_name, repo_url in REPOS_TO_CLONE.items():
+        repo_path = DEPS_DIR / repo_name
+        if repo_path.is_dir(): # Verifica se é um diretório
+            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)])
+            print(f"'{repo_name}' clonado com sucesso.")
+
+    # --- ETAPA 2: Baixar Modelos do LTX-Video ---
+    print("\n--- ETAPA 2: Preparando Modelos LTX-Video ---")
+    if not LTX_VIDEO_REPO_DIR.is_dir():
+        print(f"ERRO: Diretório '{LTX_VIDEO_REPO_DIR}' não encontrado. Execute a clonagem primeiro.")
+        sys.exit(1)
+    
+    ltx_config = _load_ltx_config()
+    if ltx_config:
+        _download_ltx_models(ltx_config)
+    else:
+        print("ERRO: Não foi possível carregar a configuração do LTX-Video. Abortando.")
+        sys.exit(1)
+
+    # --- ETAPA 3: Baixar Modelos do SeedVR ---
+    print("\n--- ETAPA 3: Preparando Modelos SeedVR ---")
+    _download_seedvr_models()
 
-    print("\n--- ✅ Setup do Ambiente Concluído com Sucesso! ---")
-    print("Todos os repositórios, dependências e modelos estão prontos.")
+    print("\n\n--- Setup do Ambiente Concluído com Sucesso! ---")
+    print("Todos os repositórios e modelos necessários foram verificados e estão prontos.")
+    print("Você agora pode iniciar a aplicação principal.")
 
 if __name__ == "__main__":
     main()