提交 ca87c09c 编写于 作者: D d8ahazard

Fix recursive model loading

Ensure we find checkpoints within subdirectories.
上级 19eb1467
import glob
import os import os
import shutil import shutil
import importlib import importlib
...@@ -41,7 +42,7 @@ def load_models(model_path: str, model_url: str = None, command_path: str = None ...@@ -41,7 +42,7 @@ def load_models(model_path: str, model_url: str = None, command_path: str = None
for place in places: for place in places:
if os.path.exists(place): if os.path.exists(place):
for file in os.listdir(place): for file in glob.iglob(place + '**/**', recursive=True):
full_path = os.path.join(place, file) full_path = os.path.join(place, file)
if os.path.isdir(full_path): if os.path.isdir(full_path):
continue continue
...@@ -50,6 +51,7 @@ def load_models(model_path: str, model_url: str = None, command_path: str = None ...@@ -50,6 +51,7 @@ def load_models(model_path: str, model_url: str = None, command_path: str = None
if extension not in ext_filter: if extension not in ext_filter:
continue continue
if file not in output: if file not in output:
print(f"FILE: {full_path}")
output.append(full_path) output.append(full_path)
if model_url is not None and len(output) == 0: if model_url is not None and len(output) == 0:
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册