Spaces:
Running
on
Zero
Running
on
Zero
File size: 683 Bytes
169b607 0509eb5 169b607 e90d194 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 |
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
|