import yaml
import os
from pathlib import Path
from huggingface_hub import hf_hub_download

def download_ltx_models():
    """
    独立下载LTX-Video模型的脚本
    保持与主程序相同的路径和配置
    """
    
    # 读取配置文件
    config_file_path = "configs/ltxv-13b-0.9.7-distilled.yaml"
    
    if not os.path.exists(config_file_path):
        print(f"错误: 配置文件 {config_file_path} 不存在")
        print("请确保配置文件在正确的位置")
        return False
    
    with open(config_file_path, "r") as file:
        PIPELINE_CONFIG_YAML = yaml.safe_load(file)
    
    # 设置常量
    LTX_REPO = "Lightricks/LTX-Video"
    models_dir = "downloaded_models_gradio_cpu_init"
    
    # 创建模型目录
    Path(models_dir).mkdir(parents=True, exist_ok=True)
    print(f"模型下载目录: {Path(models_dir).resolve()}")
    
    try:
        # 下载主模型
        print("\n开始下载主模型...")
        print(f"模型文件: {PIPELINE_CONFIG_YAML['checkpoint_path']}")
        
        distilled_model_actual_path = hf_hub_download(
            repo_id=LTX_REPO,
            filename=PIPELINE_CONFIG_YAML["checkpoint_path"],
            local_dir=models_dir,
            local_dir_use_symlinks=False
        )
        print(f"✅ 主模型下载完成: {distilled_model_actual_path}")
        
        # 下载空间上采样器模型
        print("\n开始下载空间上采样器模型...")
        SPATIAL_UPSCALER_FILENAME = PIPELINE_CONFIG_YAML["spatial_upscaler_model_path"]
        print(f"模型文件: {SPATIAL_UPSCALER_FILENAME}")
        
        spatial_upscaler_actual_path = hf_hub_download(
            repo_id=LTX_REPO,
            filename=SPATIAL_UPSCALER_FILENAME,
            local_dir=models_dir,
            local_dir_use_symlinks=False
        )
        print(f"✅ 空间上采样器模型下载完成: {spatial_upscaler_actual_path}")
        
        # 显示下载摘要
        print("\n" + "="*60)
        print("模型下载完成摘要:")
        print("="*60)
        print(f"下载目录: {models_dir}")
        print(f"主模型: {os.path.basename(distilled_model_actual_path)}")
        print(f"上采样器: {os.path.basename(spatial_upscaler_actual_path)}")
        
        # 检查文件大小
        main_size = os.path.getsize(distilled_model_actual_path) / (1024**3)  # GB
        upscaler_size = os.path.getsize(spatial_upscaler_actual_path) / (1024**3)  # GB
        total_size = main_size + upscaler_size
        
        print(f"\n文件大小:")
        print(f"主模型: {main_size:.2f} GB")
        print(f"上采样器: {upscaler_size:.2f} GB")
        print(f"总计: {total_size:.2f} GB")
        
        return True
        
    except Exception as e:
        print(f"\n❌ 下载过程中出现错误: {e}")
        print("可能的解决方案:")
        print("1. 检查网络连接")
        print("2. 确认Hugging Face访问权限")
        print("3. 检查磁盘空间是否足够")
        return False

def check_models_exist():
    """
    检查模型是否已经存在
    """
    config_file_path = "configs/ltxv-13b-0.9.7-distilled.yaml"
    
    if not os.path.exists(config_file_path):
        return False
    
    with open(config_file_path, "r") as file:
        config = yaml.safe_load(file)
    
    models_dir = "downloaded_models_gradio_cpu_init"
    main_model = os.path.join(models_dir, config["checkpoint_path"])
    upscaler_model = os.path.join(models_dir, config["spatial_upscaler_model_path"])
    
    main_exists = os.path.exists(main_model)
    upscaler_exists = os.path.exists(upscaler_model)
    
    print("模型存在性检查:")
    print(f"主模型: {'✅ 存在' if main_exists else '❌ 不存在'}")
    print(f"上采样器: {'✅ 存在' if upscaler_exists else '❌ 不存在'}")
    
    return main_exists and upscaler_exists

def main():
    print("LTX-Video 模型下载器")
    print("="*40)
    
    # 检查模型是否已存在
    if check_models_exist():
        print("\n所有模型已存在，无需重新下载。")
        choice = input("是否要重新下载？(y/N): ").lower().strip()
        if choice != 'y':
            print("取消下载。")
            return
    
    print("\n开始下载模型...")
    success = download_ltx_models()
    
    if success:
        print("\n🎉 所有模型下载成功！")
        print("现在可以运行主程序了。")
    else:
        print("\n💥 模型下载失败，请检查错误信息并重试。")

if __name__ == "__main__":
    main()