import os
from huggingface_hub import HfApi

# 🔥 1. 设置你的仓库 ID
REPO_ID = "LancetRobotics/DeCo-MAE"

# 🔥 2. 设置文件路径 (根据你之前的训练记录)
# SOTA 最佳模型
MODEL_SOTA_PATH = "/root/autodl-tmp/checkpoints_final/final_sota_best.pth"
# Zero-Shot 最佳模型
MODEL_ZS_PATH = "/root/autodl-tmp/checkpoints_zeroshot/zeroshot_model.pth"
# 冷却模型 (如果你有的话)
MODEL_COOL_PATH = "/root/autodl-tmp/checkpoints_cooldown/cooldown_best.pth"

# 项目代码目录
CODE_DIR = "/root/hri30"

print(f"🚀 开始上传到 Hugging Face: {REPO_ID}")
api = HfApi()

# ==========================================
# 任务 A: 上传模型权重 (大文件)
# ==========================================
print("\n📦 正在上传模型权重 (可能需要几分钟)...")

# 1. 上传 Final SOTA 模型
if os.path.exists(MODEL_SOTA_PATH):
    print(f"  -> Uploading: {os.path.basename(MODEL_SOTA_PATH)}")
    api.upload_file(
        path_or_fileobj=MODEL_SOTA_PATH,
        path_in_repo="final_sota_best.pth",
        repo_id=REPO_ID,
        repo_type="model"
    )

# 2. 上传 Zero-Shot 模型
if os.path.exists(MODEL_ZS_PATH):
    print(f"  -> Uploading: {os.path.basename(MODEL_ZS_PATH)}")
    api.upload_file(
        path_or_fileobj=MODEL_ZS_PATH,
        path_in_repo="zeroshot_model.pth",
        repo_id=REPO_ID,
        repo_type="model"
    )

# 3. 上传 Cooldown 模型 (如果有)
if os.path.exists(MODEL_COOL_PATH):
    print(f"  -> Uploading: {os.path.basename(MODEL_COOL_PATH)}")
    api.upload_file(
        path_or_fileobj=MODEL_COOL_PATH,
        path_in_repo="cooldown_best.pth",
        repo_id=REPO_ID,
        repo_type="model"
    )

# ==========================================
# 任务 B: 上传代码和图片
# ==========================================
print("\n📂 正在上传代码和图片...")

# 我们只上传 .py, .md, .txt 和 fig 文件夹，忽略 .git, __pycache__ 等垃圾文件
api.upload_folder(
    folder_path=CODE_DIR,
    repo_id=REPO_ID,
    repo_type="model",
    allow_patterns=["*.py", "*.md", "*.txt", "fig/*", "LICENSE"],
    ignore_patterns=[".git/*", "__pycache__/*", "*.pth", "wandb/*"] # 忽略本地的pth防止重复
)

print("\n✅ 所有文件上传成功！")
print(f"👉 查看地址: https://huggingface.co/{REPO_ID}")