Update setup.py

setup.py

@@ -2,9 +2,9 @@
 #
 # Copyright (C) August 4, 2025  Carlos Rodrigues dos Santos
 #
-# Versão 3.0.0 (Setup Unificado com LTX, SeedVR e VINCIE)
+# Versão 3.1.0 (Setup Unificado com LTX, SeedVR e VINCIE com Cache Robusto)
 # - Orquestra a instalação de todos os repositórios e modelos para a suíte ADUC-SDR.
-# - Garante que todos os downloads usem um cache persistente para inicializações rápidas.
+# - Usa snapshot_download para baixar dependências de forma eficiente e correta.
 
 import os
 import subprocess
@@ -13,12 +13,14 @@ from pathlib import Path
 import yaml
 from huggingface_hub import hf_hub_download, snapshot_download
 
-# --- Configuração de Paths e Cache ---
-# Assume-se que /data é um volume persistente montado no contêiner.
+# ==============================================================================
+# --- CONFIGURAÇÃO DE PATHS E CACHE ---
+# ==============================================================================
+# Assume que /data é um volume persistente montado no contêiner.
 DEPS_DIR = Path("/data")
 CACHE_DIR = DEPS_DIR / ".cache" / "huggingface"
 
-# --- Configuração dos Módulos da Aplicação ---
+# --- Paths dos Módulos da Aplicação ---
 LTX_VIDEO_REPO_DIR = DEPS_DIR / "LTX-Video"
 SEEDVR_MODELS_DIR = DEPS_DIR / "models" / "SeedVR"
 VINCIE_REPO_DIR = DEPS_DIR / "VINCIE"
@@ -31,65 +33,69 @@ REPOS_TO_CLONE = {
     "VINCIE": "https://github.com/ByteDance-Seed/VINCIE",
 }
 
+# ==============================================================================
+# --- FUNÇÕES AUXILIARES ---
+# ==============================================================================
+
 def run_command(command, cwd=None):
     """Executa um comando no terminal de forma segura e com logs claros."""
     print(f"Executando: {' '.join(command)}")
     try:
-        # Usamos check=True para lançar uma exceção se o comando falhar.
         subprocess.run(
-            command,
-            check=True,
-            cwd=cwd,
-            stdout=subprocess.PIPE,
-            stderr=subprocess.PIPE,
-            text=True,
+            command, check=True, cwd=cwd,
+            stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True,
         )
     except subprocess.CalledProcessError as e:
-        print(f"ERRO: O comando falhou com o código {e.returncode}")
-        print(f"Stderr:\n{e.stderr.strip()}")
+        print(f"ERRO: O comando falhou com o código {e.returncode}\nStderr:\n{e.stderr.strip()}")
         sys.exit(1)
     except FileNotFoundError:
-        print(f"ERRO: Comando '{command[0]}' não encontrado. Verifique se o git está instalado no contêiner.")
+        print(f"ERRO: Comando '{command[0]}' não encontrado. Verifique se o git está instalado.")
         sys.exit(1)
 
 def _load_ltx_config():
     """Carrega o arquivo de configuração YAML do LTX-Video."""
     print("--- Carregando Configuração do LTX-Video ---")
     config_file = LTX_VIDEO_REPO_DIR / "configs" / "ltxv-13b-0.9.8-distilled-fp8.yaml"
-    
     if not config_file.exists():
         print(f"ERRO: Arquivo de configuração do LTX não encontrado em '{config_file}'")
         return None
-        
     print(f"Configuração LTX encontrada: {config_file}")
     with open(config_file, "r") as file:
         return yaml.safe_load(file)
 
-def _download_models_from_hub(repo_id, filenames, cache_dir, local_dir=None):
-    """Função auxiliar genérica para baixar arquivos do Hugging Face Hub."""
-    for filename in filenames:
-        if not filename: continue
-        
-        print(f"Verificando/Baixando '{filename}' de '{repo_id}'...")
-        try:
-            hf_hub_download(
-                repo_id=repo_id,
-                filename=filename,
-                cache_dir=str(cache_dir),
+def _ensure_hf_model(repo_id, filenames=None, allow_patterns=None, local_dir=None):
+    """Função genérica para baixar um ou mais arquivos (hf_hub_download) ou um snapshot (snapshot_download)."""
+    if not repo_id: return
+    
+    print(f"Verificando/Baixando modelo do repositório: '{repo_id}'...")
+    try:
+        if filenames: # Baixa arquivos específicos
+            for filename in filenames:
+                if not filename: continue
+                hf_hub_download(
+                    repo_id=repo_id, filename=filename, cache_dir=str(CACHE_DIR),
+                    local_dir=str(local_dir) if local_dir else None,
+                    #local_dir_use_symlinks=False, 
+                    token=os.getenv("HF_TOKEN"),
+                )
+        else: # Baixa um snapshot (partes de um repositório)
+            snapshot_download(
+                repo_id=repo_id, cache_dir=str(CACHE_DIR), 
                 local_dir=str(local_dir) if local_dir else None,
-                local_dir_use_symlinks=False, # Salva o arquivo real, não um link simbólico
+                allow_patterns=allow_patterns, 
                 token=os.getenv("HF_TOKEN"),
             )
-            print(f"-> '{filename}' está disponível.")
-        except Exception as e:
-            if "404" in str(e) or "Not Found" in str(e):
-                 print(f"AVISO: Arquivo '{filename}' não encontrado em '{repo_id}'. Isso pode ser normal se for um ponteiro para um diretório (como em Text Encoders). Continuando.")
-            else:
-                print(f"ERRO CRÍTICO ao baixar '{filename}': {e}")
-                sys.exit(1)
+        print(f"-> Modelo '{repo_id}' está disponível.")
+    except Exception as e:
+        print(f"ERRO CRÍTICO ao baixar o modelo '{repo_id}': {e}")
+        sys.exit(1)
+
+# ==============================================================================
+# --- FUNÇÃO PRINCIPAL DE SETUP ---
+# ==============================================================================
 
 def main():
-    """Função principal que orquestra todo o processo de setup."""
+    """Orquestra todo o processo de setup do ambiente."""
     print("--- Iniciando Setup do Ambiente ADUC-SDR (LTX + SeedVR + VINCIE) ---")
     DEPS_DIR.mkdir(exist_ok=True)
     CACHE_DIR.mkdir(parents=True, exist_ok=True)
@@ -105,28 +111,31 @@ def main():
             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 e suas Dependências ---
+    # --- ETAPA 2: Baixar Modelos LTX-Video e Dependências ---
     print("\n--- ETAPA 2: Verificando Modelos LTX-Video e Dependências ---")
     ltx_config = _load_ltx_config()
     if not ltx_config:
         print("ERRO: Não foi possível carregar a configuração do LTX-Video. Abortando.")
         sys.exit(1)
         
-    _download_models_from_hub(
+    _ensure_hf_model(
         repo_id="Lightricks/LTX-Video",
-        filenames=[ltx_config.get("checkpoint_path"), ltx_config.get("spatial_upscaler_model_path")],
-        cache_dir=CACHE_DIR
+        filenames=[ltx_config.get("checkpoint_path"), ltx_config.get("spatial_upscaler_model_path")]
     )
     
-    dependency_repos = [
-        ltx_config.get("text_encoder_model_name_or_path"),
+    _ensure_hf_model(
+        repo_id=ltx_config.get("text_encoder_model_name_or_path"),
+        allow_patterns=["*.json", "*.safetensors"]
+    )
+    
+    enhancer_repos = [
         ltx_config.get("prompt_enhancer_image_caption_model_name_or_path"),
         ltx_config.get("prompt_enhancer_llm_model_name_or_path"),
     ]
-    for repo_id in filter(None, dependency_repos):
-         _download_models_from_hub(repo_id=repo_id, filenames=["config.json"], cache_dir=CACHE_DIR)
+    for repo_id in filter(None, enhancer_repos):
+         _ensure_hf_model(repo_id=repo_id, allow_patterns=["*.json", "*.safetensors", "*.bin"])
 
-    # --- ETAPA 3: Baixar Modelos do SeedVR ---
+    # --- ETAPA 3: Baixar Modelos SeedVR ---
     print("\n--- ETAPA 3: Verificando Modelos SeedVR ---")
     SEEDVR_MODELS_DIR.mkdir(parents=True, exist_ok=True)
     seedvr_files = {
@@ -136,38 +145,24 @@ def main():
     }
     for filename, repo_id in seedvr_files.items():
         if not (SEEDVR_MODELS_DIR / filename).is_file():
-             _download_models_from_hub(repo_id=repo_id, filenames=[filename], cache_dir=CACHE_DIR, local_dir=SEEDVR_MODELS_DIR)
+             _ensure_hf_model(repo_id=repo_id, filenames=[filename], local_dir=SEEDVR_MODELS_DIR)
         else:
             print(f"Arquivo SeedVR '{filename}' já existe. Pulando.")
 
     # --- ETAPA 4: Baixar Modelos VINCIE ---
     print("\n--- ETAPA 4: Verificando Modelos VINCIE ---")
     VINCIE_CKPT_DIR.mkdir(parents=True, exist_ok=True)
-    try:
-        print(f"Verificando snapshot do VINCIE em '{VINCIE_CKPT_DIR}'...")
-        # snapshot_download é ideal para baixar um repositório inteiro.
-        snapshot_download(
-            repo_id="ByteDance-Seed/VINCIE-3B",
-            local_dir=str(VINCIE_CKPT_DIR),
-            cache_dir=str(CACHE_DIR),
-            force_download=False, # Não baixa novamente se o diretório já estiver completo
-            token=os.getenv("HF_TOKEN")
-        )
-        print("-> Modelos VINCIE estão disponíveis.")
-        
-        # Cria o symlink de compatibilidade, se necessário, para que o código do VINCIE encontre os modelos
-        repo_ckpt_dir = VINCIE_REPO_DIR / "ckpt"
-        repo_ckpt_dir.mkdir(parents=True, exist_ok=True)
-        link = repo_ckpt_dir / "VINCIE-3B"
-        if not link.exists():
-            link.symlink_to(VINCIE_CKPT_DIR.resolve(), target_is_directory=True)
-            print(f"-> Symlink de compatibilidade VINCIE criado: '{link}' -> '{VINCIE_CKPT_DIR.resolve()}'")
-        else:
-            print(f"-> Symlink de compatibilidade VINCIE já existe.")
-
-    except Exception as e:
-        print(f"ERRO CRÍTICO ao baixar modelos VINCIE: {e}")
-        sys.exit(1)
+    _ensure_hf_model(repo_id="ByteDance-Seed/VINCIE-3B", local_dir=VINCIE_CKPT_DIR)
+    
+    # Cria o symlink de compatibilidade, se necessário
+    repo_ckpt_dir = VINCIE_REPO_DIR / "ckpt"
+    repo_ckpt_dir.mkdir(parents=True, exist_ok=True)
+    link = repo_ckpt_dir / "VINCIE-3B"
+    if not link.exists():
+        link.symlink_to(VINCIE_CKPT_DIR.resolve(), target_is_directory=True)
+        print(f"-> Symlink de compatibilidade VINCIE criado: '{link}' -> '{VINCIE_CKPT_DIR.resolve()}'")
+    else:
+        print(f"-> Symlink de compatibilidade VINCIE já existe.")
 
     print("\n\n--- ✅ Setup Completo do Ambiente ADUC-SDR Concluído com Sucesso! ---")
     print("Todos os repositórios e modelos foram verificados e estão prontos para uso.")