#!/usr/bin/env python3
"""
快速上传 KeyVID 模型到 Hugging Face Hub
针对大文件和高速上传优化
"""

from pathlib import Path
from huggingface_hub import HfApi, upload_folder
from tqdm import tqdm
import time
import os

# Model repository name on Hugging Face
MODEL_ID = "RyanWW/KeyVID"
KEYVID_PATH = "/dockerx/groups/KeyVID_hf_model"

def main():
    print("🚀 KeyVID 快速上传到 Hugging Face")
    print(f"📦 Repository: {MODEL_ID}")
    print(f"📁 Directory: {KEYVID_PATH}\n")
    
    # 检查认证
    try:
        api = HfApi()
        print("✅ Hugging Face 认证成功\n")
    except Exception as e:
        print("❌ 需要 Hugging Face 认证")
        print("请运行: huggingface-cli login")
        print("或设置环境变量: export HF_TOKEN=your_token")
        return
    
    keyvid_dir = Path(KEYVID_PATH)
    if not keyvid_dir.exists():
        print(f"❌ 目录不存在: {KEYVID_PATH}")
        return
    
    # 优化的排除模式（只上传必要文件）
    ignore_patterns = [
        "**/__pycache__/**",
        "**/.git/**",
        "**/*.pyc",
        "**/.DS_Store",
        "**/save_results/**",  # 排除结果文件
        "**/*.log",
        "**/*.tmp",
        "**/upload*.py",
        "**/.gitignore",
        "**/.gitattributes",
    ]
    
    # 显示上传策略
    print("⚙️  上传配置:")
    print("   - multi_commits=True (分批提交，加快大文件)")
    print("   - 自动处理大文件 (使用LFS)")
    print("   - 跳过不必要的文件\n")
    
    # 确认
    response = input("❓ 开始上传? (y/n): ").strip().lower()
    if response != 'y':
        print("❌ 已取消")
        return
    
    print("\n⬆️  开始上传...")
    print("💡 提示: 大文件上传可能需要较长时间，请耐心等待\n")
    
    start_time = time.time()
    
    try:
        upload_folder(
            folder_path=str(keyvid_dir),
            repo_id=MODEL_ID,
            repo_type="model",
            ignore_patterns=ignore_patterns,
            commit_message="Upload KeyVID model files",
            # 关键优化参数
            multi_commits=True,          # 分批提交，不需要等所有文件
            multi_commits_verbose=True,  # 显示详细进度
            allow_patterns=None,         # 不限制模式，上传所有匹配文件
        )
        
        elapsed = time.time() - start_time
        print(f"\n✅ 上传完成!")
        print(f"⏱️  耗时: {elapsed/60:.2f} 分钟 ({elapsed:.2f} 秒)")
        print(f"🔗 查看: https://huggingface.co/{MODEL_ID}")
        
    except KeyboardInterrupt:
        print("\n⚠️  上传被用户中断")
        print("💡 可以稍后重新运行脚本，Hugging Face 支持断点续传")
    except Exception as e:
        print(f"\n❌ 上传错误: {e}")
        print("\n💡 建议:")
        print("   1. 检查网络连接")
        print("   2. 检查 Hugging Face token 是否有效")
        print("   3. 检查仓库权限 (repo_id: {MODEL_ID})")
        print("   4. 对于大文件，可能需要较长时间")

if __name__ == "__main__":
    main()