Update setup.py

setup.py

@@ -2,179 +2,175 @@
 #
 # 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).
+# Versão 3.0.0 (Setup Unificado com LTX, SeedVR e VINCIE)
+# - 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.
 
 import os
 import subprocess
 import sys
 from pathlib import Path
 import yaml
-from huggingface_hub import hf_hub_download
+from huggingface_hub import hf_hub_download, snapshot_download
 
-# --- Configuração Geral ---
+# --- Configuração de Paths e Cache ---
+# Assume-se que /data é um volume persistente montado no contêiner.
 DEPS_DIR = Path("/data")
+CACHE_DIR = DEPS_DIR / ".cache" / "huggingface"
 
-# --- Configuração Específica LTX-Video ---
+# --- Configuração 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"
+VINCIE_CKPT_DIR = DEPS_DIR / "ckpt" / "VINCIE-3B"
 
-# --- Configuração Específica SeedVR ---
-SEEDVR_MODELS_DIR = DEPS_DIR / "SeedVR"
-
-# --- Repositórios para Clonar ---
+# --- Repositórios Git para Clonar ---
 REPOS_TO_CLONE = {
     "LTX-Video": "https://huggingface.co/spaces/Lightricks/ltx-video-distilled",
     "SeedVR": "https://github.com/numz/ComfyUI-SeedVR2_VideoUpscaler",
-    "MMAudio": "https://github.com/hkchengrex/MMAudio.git"
+    "VINCIE": "https://github.com/ByteDance-Seed/VINCIE",
 }
 
 def run_command(command, cwd=None):
-    """Executa um comando no terminal e lida com erros."""
+    """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,
-            stdin=subprocess.DEVNULL,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            text=True,
         )
     except subprocess.CalledProcessError as e:
-        print(f"ERRO: O comando falhou com o código de saída {e.returncode}\nStderr: {e.stderr}")
+        print(f"ERRO: O comando falhou com o código {e.returncode}")
+        print(f"Stderr:\n{e.stderr.strip()}")
         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: Comando '{command[0]}' não encontrado. Verifique se o git está instalado no contêiner.")
         sys.exit(1)
 
-# --- 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)
+    config_file = LTX_VIDEO_REPO_DIR / "configs" / "ltxv-13b-0.9.8-distilled-fp8.yaml"
     
-    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.")
+    if not config_file.exists():
+        print(f"ERRO: Arquivo de configuração do LTX não encontrado em '{config_file}'")
         return None
         
-    with open(fallback_path, "r") as file:
+    print(f"Configuração LTX encontrada: {config_file}")
+    with open(config_file, "r") as file:
         return yaml.safe_load(file)
 
-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)
-
-    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}...")
+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=LTX_REPO, filename=filename,
-                local_dir=os.getenv("HF_HOME"), cache_dir=os.getenv("HF_HOME_CACHE"), token=os.getenv("HF_TOKEN")
+                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, # Salva o arquivo real, não um link simbólico
+                token=os.getenv("HF_TOKEN"),
             )
-            print(f"{description.capitalize()} está disponível.")
+            print(f"-> '{filename}' está disponível.")
         except Exception as e:
-            print(f"ERRO ao baixar o {description}: {e}")
-            sys.exit(1)
-
-
-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)
-    
-    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=repo_id,
-                    filename=filename,
-                    local_dir=str(SEEDVR_MODELS_DIR),
-                    cache_dir=os.getenv("HF_HOME_CACHE"),
-                    token=os.getenv("HF_TOKEN"),
-                )
-                print(f"'{filename}' baixado com sucesso.")
-            except Exception as e:
-                print(f"ERRO ao baixar o modelo SeedVR '{filename}': {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)
-        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 do Ambiente ADUC-SDR (Versão Robusta) ---")
+    """Função principal que orquestra todo o processo de setup."""
+    print("--- Iniciando Setup do Ambiente ADUC-SDR (LTX + SeedVR + VINCIE) ---")
     DEPS_DIR.mkdir(exist_ok=True)
+    CACHE_DIR.mkdir(parents=True, exist_ok=True)
     
     # --- ETAPA 1: Clonar Repositórios ---
-    print("\n--- ETAPA 1: Clonando Repositórios Git ---")
+    print("\n--- ETAPA 1: Verificando 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.")
+        if repo_path.is_dir():
+            print(f"Repositório '{repo_name}' já existe em '{repo_path}'. 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.")
+            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)
-    
+    # --- ETAPA 2: Baixar Modelos do LTX-Video e suas Dependências ---
+    print("\n--- ETAPA 2: Verificando Modelos LTX-Video e Dependências ---")
     ltx_config = _load_ltx_config()
-    if ltx_config:
-        _download_ltx_models(ltx_config)
-    else:
+    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(
+        repo_id="Lightricks/LTX-Video",
+        filenames=[ltx_config.get("checkpoint_path"), ltx_config.get("spatial_upscaler_model_path")],
+        cache_dir=CACHE_DIR
+    )
+    
+    dependency_repos = [
+        ltx_config.get("text_encoder_model_name_or_path"),
+        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)
 
     # --- ETAPA 3: Baixar Modelos do SeedVR ---
-    print("\n--- ETAPA 3: Preparando Modelos SeedVR ---")
-    _download_seedvr_models()
+    print("\n--- ETAPA 3: Verificando Modelos SeedVR ---")
+    SEEDVR_MODELS_DIR.mkdir(parents=True, exist_ok=True)
+    seedvr_files = {
+        "seedvr2_ema_7b_fp16.safetensors": "MonsterMMORPG/SeedVR2_SECourses",
+        "seedvr2_ema_7b_sharp_fp16.safetensors": "MonsterMMORPG/SeedVR2_SECourses",
+        "ema_vae_fp16.safetensors": "MonsterMMORPG/SeedVR2_SECourses",
+    }
+    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)
+        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)
 
-    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.")
+    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.")
 
 if __name__ == "__main__":
     main()