Fix recursive model loading

Ensure we find checkpoints within subdirectories.
This commit is contained in:
d8ahazard 2022-09-30 09:28:28 -05:00
parent 19eb1467f1
commit ca87c09c0e

View File

@ -1,3 +1,4 @@
import glob
import os
import shutil
import importlib
@ -41,7 +42,7 @@ def load_models(model_path: str, model_url: str = None, command_path: str = None
for place in places:
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)
if os.path.isdir(full_path):
continue
@ -50,6 +51,7 @@ def load_models(model_path: str, model_url: str = None, command_path: str = None
if extension not in ext_filter:
continue
if file not in output:
print(f"FILE: {full_path}")
output.append(full_path)
if model_url is not None and len(output) == 0: