upload_ckpt.py

import os
from huggingface_hub import HfApi, login
from pathlib import Path

# 登录到 Hugging Face（需要先设置 HF_TOKEN 环境变量或手动输入token）
login()

# 初始化 HfApi
api = HfApi()

# 配置参数
repo_id = "lihong-cs/rhosckpt"  # 替换为您的用户名和模型名
local_ckpt_path = "/home/lixingyu/workspace_lxy2/workspace/robotpolicy/src/ckpt"  # 替换为您的检查点文件路径
repo_type = "model"  # 可以是 "model", "dataset", 或 "space"

# 方法1: 上传单个文件
def upload_single_file():
    ckpt_file = "model.pth"  # 或 "checkpoint.pt", "model.safetensors" 等
    
    api.upload_file(
        path_or_fileobj=os.path.join(local_ckpt_path, ckpt_file),
        path_in_repo=ckpt_file,
        repo_id=repo_id,
        repo_type=repo_type,
        commit_message=f"Upload {ckpt_file}"
    )
    print(f"文件 {ckpt_file} 上传完成！")

# 方法2: 上传整个文件夹
def upload_folder():
    """上传整个检查点文件夹"""
    api.upload_folder(
        folder_path=local_ckpt_path,
        repo_id=repo_id,
        repo_type=repo_type,
        commit_message="Upload checkpoint folder"
    )
    print(f"文件夹 {local_ckpt_path} 上传完成！")

# 方法3: 创建仓库并上传（如果仓库不存在）
def create_and_upload():
    """创建仓库并上传检查点"""
    try:
        # 创建仓库（如果不存在）
        api.create_repo(
            repo_id=repo_id,
            repo_type=repo_type,
            exist_ok=True,  # 如果仓库已存在则不报错
            private=False   # 设置为 True 创建私有仓库
        )
        print(f"仓库 {repo_id} 创建成功或已存在")
        
        # 上传文件夹
        upload_folder()
        
    except Exception as e:
        print(f"上传过程中出错: {e}")

# 方法4: 带进度条的上传
def upload_with_progress():
    """带进度显示的上传"""
    from tqdm import tqdm
    
    # 获取所有要上传的文件
    files_to_upload = []
    for root, dirs, files in os.walk(local_ckpt_path):
        for file in files:
            file_path = os.path.join(root, file)
            relative_path = os.path.relpath(file_path, local_ckpt_path)
            files_to_upload.append((file_path, relative_path))
    
    # 逐个上传文件
    for file_path, relative_path in tqdm(files_to_upload, desc="上传文件"):
        try:
            api.upload_file(
                path_or_fileobj=file_path,
                path_in_repo=relative_path,
                repo_id=repo_id,
                repo_type=repo_type,
                commit_message=f"Upload {relative_path}"
            )
        except Exception as e:
            print(f"上传 {relative_path} 时出错: {e}")

if __name__ == "__main__":
    # 检查本地路径是否存在
    if not os.path.exists(local_ckpt_path):
        print(f"错误: 本地路径 {local_ckpt_path} 不存在")
        exit(1)
    
    # 选择上传方式
    print("请选择上传方式:")
    print("1. 上传单个文件")
    print("2. 上传整个文件夹")
    print("3. 创建仓库并上传")
    print("4. 带进度条上传")
    
    choice = input("请输入选择 (1-4): ")
    
    if choice == "1":
        upload_single_file()
    elif choice == "2":
        upload_folder()
    elif choice == "3":
        create_and_upload()
    elif choice == "4":
        upload_with_progress()
    else:
        print("无效选择")