auto_uoload.py 2.59 KB
Newer Older
wangkx1's avatar
init  
wangkx1 committed
1
2
3
4
5
6
7
8
9
10
from pycsghub.upload_large_folder.main import upload_large_folder_internal
from pycsghub.upload_large_folder.main import create_repo
from pycsghub.csghub_api import CsgHubApi
import os
import argparse


BASE_URL = "http://10.20.100.5:4997"
TOKEN = "a2683e780fee4aa4b16cbedee0f75a45"

wangkx1's avatar
init  
wangkx1 committed
11
def upload(local_path, repo_id, repo_type="model", force=False):
wangkx1's avatar
init  
wangkx1 committed
12
13
14
15
16
17
18
19
    csg_api = CsgHubApi()
    use_full_repo_id = f"root/{repo_id}"

    repoExist, branchExist = csg_api.repo_branch_exists(
        repo_id=use_full_repo_id, repo_type=repo_type, revision="v1",
        endpoint=BASE_URL, token=TOKEN)
    
    if repoExist:
wangkx1's avatar
init  
wangkx1 committed
20
21
22
23
24
        if not force:
            print(f"{repo_id} has exists. Use --force to update it.")
            return
        else:
            print(f"{repo_id} has exists. Updating it because --force is set.")
wangkx1's avatar
init  
wangkx1 committed
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
    else:
        create_repo(
            api=csg_api,
            repo_id=use_full_repo_id,
            repo_type=repo_type,
            revision="v1",
            endpoint=BASE_URL,
            token=TOKEN
        )

    upload_large_folder_internal(
        repo_id=use_full_repo_id,
        local_path=local_path,
        repo_type=repo_type,
        revision="v1",
        endpoint=BASE_URL,
        token=TOKEN,
        allow_patterns=None,
wangkx1's avatar
init  
wangkx1 committed
43
        ignore_patterns=[".*"],  # 忽略所有隐藏目录和文件(以点开头的)
wangkx1's avatar
init  
wangkx1 committed
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
        num_workers=1,
        print_report=False,
        print_report_every=1,
    )


if __name__ == "__main__":
    # 创建命令行参数解析器
    parser = argparse.ArgumentParser(description="自动上传模型到CSGHUB")
    
    # 添加本地模型路径参数
    parser.add_argument(
        "local_path", 
        type=str, 
        help="本地模型文件夹的路径"
    )
    
wangkx1's avatar
init  
wangkx1 committed
61
62
63
64
65
66
67
    # 添加强制更新参数
    parser.add_argument(
        "--force",
        action="store_true",
        help="强制更新远程仓库,即使仓库已存在"
    )
    
wangkx1's avatar
init  
wangkx1 committed
68
69
    # 解析命令行参数
    args = parser.parse_args()
wangkx1's avatar
init  
wangkx1 committed
70

wangkx1's avatar
init  
wangkx1 committed
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
    # 获取本地模型路径
    local_path = args.local_path
    
    # 判断本地模型路径是否存在
    if not os.path.exists(local_path):
        print(f"错误: 本地模型路径 '{local_path}' 不存在")
        exit(1)
    
    if not os.path.isdir(local_path):
        print(f"错误: 本地模型路径 '{local_path}' 不是一个文件夹")
        exit(1)
    
    # 从本地路径中提取仓库名称(使用文件夹名)
    repo_id = os.path.basename(local_path)
    
    print(f"开始上传本地模型: {local_path}")
    print(f"目标仓库: {repo_id}")
    
    # 执行上传
wangkx1's avatar
init  
wangkx1 committed
90
    upload(local_path, repo_id, force=args.force)