AnyDiffuse / utils /model_utils.py
zerhero's picture
refactor
0509eb5
raw
history blame
682 Bytes
def get_model_list(directory_path):
"""
:param directory_path:
:return:
"""
import os
model_list: list = []
valid_extensions = {
'.ckpt',
'.pt',
'.pth',
'.safetensors',
'.bin'
}
for filename in os.listdir(directory_path):
if os.path.splitext(filename)[1] in valid_extensions:
name_without_extension = os.path.splitext(filename)[0]
file_path = os.path.join(directory_path, filename)
# model_list.append((name_without_extension, file_path))
model_list.append(file_path)
print('\033[34mFILE: ' + file_path + '\033[0m')
return model_list