从 Hugging Face 加载模型列表并缓存,如果 HF 超时或失败,则尝试使用 ModelScope
()
| 237 | |
| 238 | |
| 239 | def load_hf_models_cache(): |
| 240 | """从 Hugging Face 加载模型列表并缓存,如果 HF 超时或失败,则尝试使用 ModelScope""" |
| 241 | # 超时时间(秒) |
| 242 | HF_TIMEOUT = 30 |
| 243 | |
| 244 | for repo_id in HF_MODELS_CACHE.keys(): |
| 245 | files = None |
| 246 | source = None |
| 247 | |
| 248 | # 首先尝试从 ModelScope 获取 |
| 249 | try: |
| 250 | if MS_AVAILABLE: |
| 251 | logger.info(f"Loading models from ModelScope {repo_id}...") |
| 252 | api = HubApi() |
| 253 | # ModelScope API 获取文件列表 |
| 254 | model_files = api.get_model_files(model_id=repo_id, recursive=True) |
| 255 | # 提取文件路径 |
| 256 | files = [file["Path"] for file in model_files if file.get("Type") == "blob"] |
| 257 | source = "ModelScope" |
| 258 | logger.info(f"Successfully loaded models from ModelScope {repo_id}") |
| 259 | except: # noqa E722 |
| 260 | # 如果 ModelScope 失败,尝试从 Hugging Face 获取(带超时) |
| 261 | if files is None and HF_AVAILABLE: |
| 262 | logger.info(f"Loading models from Hugging Face {repo_id}...") |
| 263 | api = HfApi() |
| 264 | |
| 265 | # 使用线程池执行器设置超时 |
| 266 | with concurrent.futures.ThreadPoolExecutor() as executor: |
| 267 | future = executor.submit(list_repo_files, repo_id=repo_id, repo_type="model") |
| 268 | files = future.result(timeout=HF_TIMEOUT) |
| 269 | source = "Hugging Face" |
| 270 | |
| 271 | # 处理文件列表 |
| 272 | if files: |
| 273 | model_names = process_files(files, repo_id) |
| 274 | HF_MODELS_CACHE[repo_id] = model_names |
| 275 | logger.info(f"Loaded {len(HF_MODELS_CACHE[repo_id])} models from {source} {repo_id}") |
| 276 | else: |
| 277 | logger.warning(f"No files retrieved from {repo_id}, setting empty cache") |
| 278 | HF_MODELS_CACHE[repo_id] = [] |
| 279 | |
| 280 | |
| 281 | def get_hf_models(repo_id, prefix_filter=None, keyword_filter=None): |