# download_models.py
import os
import logging
from huggingface_hub import hf_hub_download

# Configuração do log para ser claro e informativo
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

def download_model(repo_id, filename, local_dir, desc):
    """
    Função auxiliar para baixar um modelo, verificando se ele já existe
    para evitar downloads repetidos.
    """
    # Garante que o diretório de destino existe
    os.makedirs(local_dir, exist_ok=True)
    
    model_path = os.path.join(local_dir, filename)
    
    if os.path.exists(model_path):
        logging.info(f"Modelo '{desc}' já existe em: {model_path}")
        return

    logging.info(f"Baixando modelo '{desc}' de '{repo_id}' para '{local_dir}'...")
    try:
        hf_hub_download(
            repo_id=repo_id,
            filename=filename,
            local_dir=local_dir,
            #local_dir_use_symlinks=False, # Importante para ambientes de contêiner
            #resume_download=True
        )
        logging.info(f"Download de '{desc}' concluído com sucesso.")
    except Exception as e:
        logging.error(f"Falha CRÍTICA ao baixar '{desc}'. Erro: {e}")
        # Lançar o erro para parar a inicialização se um modelo crítico faltar
        raise

def main():
    """
    Ponto de entrada para baixar todos os modelos de IA necessários para a Aduc-SDR.
    """
    logging.info("--- Iniciando verificação e download de modelos ---")
    
    # 1. Modelos para LTX-Video
    # Usando o modelo FP8 otimizado, como solicitado
    ltx_models_dir = "/app/ckpts"
    download_model(
        repo_id="Lightricks/LTX-Video",
        filename="ltxv-13b-0.9.8-dev-fp8.safetensors",
        local_dir=ltx_models_dir,
        desc="LTX Main Model (FP8)"
    )
    download_model(
        repo_id="Lightricks/LTX-Video",
        filename="ltxv-spatial-upscaler-0.9.8.safetensors",
        local_dir=ltx_models_dir,
        desc="LTX Spatial Upscaler"
    )

    # 2. Modelos para SeedVR
    # Usando a versão FP16 otimizada para economizar VRAM
    seedvr_models_dir = "/app/ckpts"
    download_model(
        repo_id="batuhanince/seedvr_3b_fp16",
        filename="seedvr2_ema_3b_fp16.safetensors",
        local_dir=seedvr_models_dir,
        desc="SeedVR 3B (FP16)"
    )
    download_model(
        repo_id="batuhanince/seedvr_3b_fp16",
        filename="ema_vae_fp16.safetensors",
        local_dir=seedvr_models_dir,
        desc="SeedVR VAE (FP16)"
    )
    download_model(
        repo_id="ByteDance-Seed/SeedVR2-3B",
        filename="pos_emb.pt",
        local_dir=seedvr_models_dir,
        desc="SeedVR Positive Embedding"
    )
    download_model(
        repo_id="ByteDance-Seed/SeedVR2-3B",
        filename="neg_emb.pt",
        local_dir=seedvr_models_dir,
        desc="SeedVR Negative Embedding"
    )

    # 3. Modelos para Wan2.2 (Adicione os modelos específicos conforme necessário)
    # Exemplo:
    # wan_models_dir = "/app/Wan2.2/ckpts"
    # download_model("Wan-Video/Wan2.2", "model_file.safetensors", wan_models_dir, "Wan2.2 Main Model")
    
    # 4. Modelos para MMAudio (Geralmente são baixados pelo próprio código, mas podemos forçar aqui se necessário)
    # ...

    logging.info("--- Verificação de modelos concluída ---")

if __name__ == "__main__":
    main()