setup.py

# 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
from huggingface_hub import hf_hub_download

# --- Configuração Geral ---
DEPS_DIR = Path("/data")

# --- 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": "https://github.com/numz/ComfyUI-SeedVR2_VideoUpscaler",
    "MMAudio": "https://github.com/hkchengrex/MMAudio.git"
}

def run_command(command, cwd=None):
    """Executa um comando no terminal e lida com erros."""
    print(f"Executando: {' '.join(command)}")
    try:
        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}\nStderr: {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.")
        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)
    
    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)

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}...")
        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")
            )
            print(f"{description.capitalize()} 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}")
                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) ---")
    DEPS_DIR.mkdir(exist_ok=True)
    
    # --- 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\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()