File size: 5,307 Bytes
01b25a3
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
import json
import urllib.request
import urllib.error
import sys
from huggingface_hub import HfApi

# ================= 配置区 =================
SPACE_ID = "waveshaker/character_loras_gallery"
MODEL_REPO_ID = "waveshaker/character_loras"
# ==========================================

OUTPUT_FILE = "models.json"

def fetch_url_json(url):
    print(f"正在请求: {url}")
    try:
        req = urllib.request.Request(url, headers={'User-Agent': 'Mozilla/5.0'})
        with urllib.request.urlopen(req, timeout=60) as response:
            return json.loads(response.read().decode())
    except urllib.error.HTTPError as e:
        if e.code == 401:
            print(f"\n❌ 401 错误: 访问被拒绝。请确保 Space 和 Model 仓库是 Public (公开) 状态。")
        elif e.code == 404:
            print(f"\n❌ 404 错误: 找不到文件,可能是首次生成。")
        else:
            print(f"\n❌ HTTP 错误 {e.code}")
        return None
    except Exception as e:
        print(f"\n❌ 网络请求失败: {e}")
        return None

def get_all_repo_files(repo_id):
    api = HfApi()
    try:
        all_files = list(api.list_repo_files(repo_id=repo_id, repo_type="model"))
    except Exception as e:
        print(f"❌ 获取文件列表失败: {e}")
        return None
    
    return [f for f in all_files if f.endswith('.safetensors')]

def custom_sort_key(file_path):
    parts = file_path.split('/')
    top_dir = parts[0] if len(parts) > 0 else ""
    sub_dir = parts[1] if len(parts) > 1 else ""

    top_order = {"krea2": 1, "anima": 2, "flux1_dev": 3, "pony": 4, "sd15": 5}
    sub_order = {"CN": 1, "KR": 2, "JP": 3, "Fictional": 4, "WE": 5}

    top_weight = top_order.get(top_dir, 999)
    sub_weight = sub_order.get(sub_dir, 999)

    return (top_weight, sub_weight, file_path)

def main():
    # 1. 获取 Model Repo 里的所有文件
    print("\n--- 步骤 1: 拉取 Model Repo 文件列表 ---")
    remote_files = get_all_repo_files(MODEL_REPO_ID)
    
    if remote_files is None:
        print("拉取 Model Repo 失败,脚本终止。")
        sys.exit(1)
        
    print(f"✅ 找到 {len(remote_files)} 个 safetensors 文件。")

    # 2. 从 Space 获取现有的 models.json
    print("\n--- 步骤 2: 拉取 Space 现有的 models.json ---")
    space_subdomain = SPACE_ID.replace("/", "-").replace("_", "-").lower()
    config_url = f"https://{space_subdomain}.static.hf.space/models.json"
    
    existing_data = fetch_url_json(config_url)
    existing_items = {}
    
    if existing_data:
        print("✅ 成功拉取现有的 models.json!")
        existing_items = existing_data.get("items", {})
        
        if not existing_items and "titles" in existing_data:
            print("⚠️ 检测到旧版 titles 结构,正在自动升级为新版 items 结构...")
            old_titles = existing_data.get("titles", {})
            for k, v in old_titles.items():
                existing_items[k] = {"title": v, "trigger": "", "dataset": ""}
    else:
        print("⚠️ 无法获取现有的 models.json,将从头生成。")

    # 3. 比对与合并
    print("\n--- 步骤 3: 比对并合并 ---")
    merged_items = {}
    
    old_keys = set(existing_items.keys())
    newly_added_paths = [] # 👈 新增:用来记录新加入了哪些路径
    
    for file_path in remote_files:
        if file_path in existing_items:
            merged_items[file_path] = existing_items[file_path]
            old_keys.remove(file_path)
        else:
            # 生成空字典并记录路径
            merged_items[file_path] = {
                "title": "",
                "trigger": "",
                "dataset": ""
            }
            newly_added_paths.append(file_path) 

    deleted_count = len(old_keys)
    if deleted_count > 0:
        print(f"⚠️ 发现 {deleted_count} 个在云端已删除的文件,已剔除。")

    # 4. 按自定义规则排序
    print("正在进行自定义排序...")
    sorted_items = {k: merged_items[k] for k in sorted(merged_items.keys(), key=custom_sort_key)}

    # 5. 保存本地文件
    output_data = {"items": sorted_items}
    with open(OUTPUT_FILE, 'w', encoding='utf-8') as f:
        json.dump(output_data, f, indent=2, ensure_ascii=False)

    # 6. 终端输出总结与待填清单
    print(f"\n🎉 处理完成!已保存到本地: {OUTPUT_FILE}")
    print(f"  - 远端真实文件总数: {len(remote_files)}")
    print(f"  - 本地保留的旧数据: {len(remote_files) - len(newly_added_paths)}")
    print(f"  - 本次新增待填写: {len(newly_added_paths)}")
    print(f"  - 清理已失效文件: {deleted_count}")

    # 👇 核心修改点:如果存在新增,直接在终端打印出来
    if newly_added_paths:
        print("\n" + "="*60)
        print("⚠️ 以下新增文件需要你手动填写 title / trigger / dataset:")
        print("="*60)
        for path in sorted(newly_added_paths, key=custom_sort_key):
            print(f"👉 {path}")
        print("="*60)
        print("👉 请打开 models.json,搜索上述路径,将对应的空字符串补齐。然后上传覆盖 Space。")
    else:
        print("\n✅ 没有新增文件,无需手动修改。")

if __name__ == "__main__":
    main()