Upload hf_space_cap_checker.py

hf_space_cap_checker.py

@@ -0,0 +1,109 @@
+import requests
+import re
+import ast
+import json  # 🌟 加上这一行
+
+def calculate_dynamic_size():
+    # ==========================================
+    # 1. 真实链接配置 (使用 raw 原始文件下载链接)
+    # ==========================================
+    py_url = "https://huggingface.co/spaces/waveshaker/universal-lora-studio/resolve/main/model_config.py"
+    json_url = "https://huggingface.co/waveshaker/comfy_models/resolve/main/studio_config.json"
+
+    print("🔍 1. 正在从 Hugging Face 下载并解析 model_config.py ...")
+    try:
+        py_res = requests.get(py_url, timeout=15)
+        py_res.raise_for_status()
+        py_text = py_res.text
+        tree = ast.parse(py_text)
+        local_config = {}
+        for node in ast.walk(tree):
+            if isinstance(node, ast.Assign) and getattr(node.targets[0], 'id', '') == 'MODEL_CONFIG':
+                local_config = ast.literal_eval(node.value)
+                break
+    except Exception as e:
+        print(f"❌ 解析 model_config.py 失败: {e}")
+        return
+
+    print("🔍 2. 正在从 Hugging Face 下载并解析 studio_config.json ...")
+    try:
+        json_res = requests.get(json_url, timeout=15)
+        json_res.raise_for_status()
+        json_str = json_res.text
+        # 🌟 容错修复：你截图里的 JSON 在 krea2 的 unets 里少了一个逗号，这里自动补上
+        json_str = re.sub(r'\}\s*\{', '}, {', json_str)
+        remote_config = json.loads(json_str)
+    except Exception as e:
+        print(f"⚠️ 解析 studio_config.json 失败: {e}，将仅使用本地兜底配置。")
+        remote_config = {}
+
+    # ==========================================
+    # 2. 提取并去重 (利用 Set)
+    # ==========================================
+    unique_files = set()
+    def add_file(repo, subfolder, filename):
+        if repo and filename:
+            clean_sub = subfolder.strip("/") if subfolder else ""
+            unique_files.add((repo, clean_sub, filename))
+
+    # 提取 py 配置
+    for model_name, model_data in local_config.get("models", {}).items():
+        # 1. TE & VAE
+        for repo, path in model_data.get("companion_downloads", []):
+            parts = path.split("/")
+            add_file(repo, "/".join(parts[:-1]), parts[-1])
+        # 2. 兜底 UNet
+        for name, info in model_data.get("fallback_unets", {}).items():
+            add_file(info.get("repo"), info.get("subfolder"), info.get("filename"))
+        # 3. 兜底 LoRA
+        for name, info in model_data.get("fallback_base_loras", {}).items():
+            add_file(info.get("repo"), info.get("subfolder"), info.get("filename"))
+
+    # 提取远程 JSON
+    for root_key, data in remote_config.items():
+        for category in ["unets", "base_loras"]:
+            for item in data.get(category, []):
+                add_file(item.get("repo"), item.get("subfolder"), item.get("filename"))
+
+    print(f"\n🧹 去重后共发现 {len(unique_files)} 个独立文件。开始极速计算尺寸...\n")
+
+    # ==========================================
+    # 3. 极速查询 API 并计算
+    # ==========================================
+    total_bytes = 0
+    folder_cache = {}
+
+    def get_file_size(repo, subfolder, filename):
+        cache_key = f"{repo}/{subfolder}"
+        if cache_key not in folder_cache:
+            api_url = f"https://huggingface.co/api/models/{repo}/tree/main/{subfolder}"
+            try:
+                res = requests.get(api_url, timeout=10)
+                if res.status_code == 200:
+                    folder_cache[cache_key] = {
+                        item['path'].split('/')[-1]: item.get('size', 0) 
+                        for item in res.json()
+                    }
+                else:
+                    folder_cache[cache_key] = {}
+            except:
+                folder_cache[cache_key] = {}
+        return folder_cache[cache_key].get(filename, 0)
+
+    for repo, subfolder, filename in sorted(unique_files):
+        size = get_file_size(repo, subfolder, filename)
+        if size > 0:
+            total_bytes += size
+            print(f"✅ {repo:<25} | {filename[:45]:<45} | {size / (1024**3):>5.2f} GB")
+        else:
+            print(f"❌ {repo:<25} | {filename[:45]:<45} | 获取失败")
+
+    total_gb = total_bytes / (1024 ** 3)
+    print(f"\n" + "="*70)
+    print(f"📊 动态计算完成！")
+    print(f"💾 去重后文件数: {len(unique_files)}")
+    print(f"💾 总计需要下载: {total_gb:.2f} GB")
+    print(f"="*70)
+
+if __name__ == "__main__":
+    calculate_dynamic_size()