setup.py

# setup.py
#
# Copyright (C) August 4, 2025  Carlos Rodrigues dos Santos
#
# 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.
# - Usa snapshot_download para baixar dependências de forma eficiente e correta.

import os
import subprocess
import sys
from pathlib import Path
import yaml
from huggingface_hub import hf_hub_download, snapshot_download

# ==============================================================================
# --- 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"

# --- 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"
VINCIE_CKPT_DIR = DEPS_DIR / "ckpt" / "VINCIE-3B"

# --- 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",
    "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:
        subprocess.run(
            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}\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.")
        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 _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,
                allow_patterns=allow_patterns, 
                token=os.getenv("HF_TOKEN"),
            )
        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():
    """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)
    
    # --- ETAPA 1: Clonar Repositórios ---
    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():
            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.")

    # --- 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)
        
    _ensure_hf_model(
        repo_id="Lightricks/LTX-Video",
        filenames=[
            ltx_config.get("checkpoint_path"), 
            ltx_config.get("spatial_upscaler_model_path") # <-- Adicione esta linha
        ]
    )
    
    _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, enhancer_repos):
         _ensure_hf_model(repo_id=repo_id, allow_patterns=["*.json", "*.safetensors", "*.bin"])

    # --- ETAPA 3: Baixar Modelos SeedVR ---
    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():
             _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)
    _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.")

if __name__ == "__main__":
    main()