mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2026-01-11 14:50:49 +08:00
88 lines
3.4 KiB
Python
88 lines
3.4 KiB
Python
import os
|
|
from collections import namedtuple
|
|
|
|
supported_ckpt_extensions = set(['.ckpt', '.pth'])
|
|
supported_pt_extensions = set(['.ckpt', '.pt', '.bin', '.pth'])
|
|
try:
|
|
import safetensors.torch
|
|
supported_ckpt_extensions.add('.safetensors')
|
|
supported_pt_extensions.add('.safetensors')
|
|
except:
|
|
print("Could not import safetensors, safetensors support disabled.")
|
|
|
|
|
|
SearchLocation = namedtuple('SearchLocation', ['paths', 'extensions'])
|
|
model_search_locations = {}
|
|
|
|
|
|
def register_model_type(model_tag, supported_extensions):
|
|
if model_tag not in model_search_locations:
|
|
model_search_locations[model_tag] = SearchLocation([], supported_extensions)
|
|
|
|
|
|
def add_model_search_path(model_tag, full_folder_path):
|
|
global model_search_locations
|
|
if model_tag in model_search_locations:
|
|
model_search_locations[model_tag].paths.append(full_folder_path)
|
|
|
|
|
|
def get_folder_paths(model_tag):
|
|
return model_search_locations[model_tag][0][:]
|
|
|
|
|
|
def recursive_search(directory):
|
|
result = []
|
|
for root, subdir, file in os.walk(directory, followlinks=True):
|
|
for filepath in file:
|
|
#we os.path,join directory with a blank string to generate a path separator at the end.
|
|
result.append(os.path.join(root, filepath).replace(os.path.join(directory,''),''))
|
|
return result
|
|
|
|
|
|
def filter_files_extensions(files, extensions):
|
|
return sorted(list(filter(lambda a: os.path.splitext(a)[-1].lower() in extensions, files)))
|
|
|
|
|
|
def get_full_path(model_tag, filename):
|
|
global model_search_locations
|
|
folders = model_search_locations[model_tag]
|
|
for x in folders[0]:
|
|
full_path = os.path.join(x, filename)
|
|
if os.path.isfile(full_path):
|
|
return full_path
|
|
|
|
|
|
def get_filename_list(model_tag):
|
|
global model_search_locations
|
|
output_list = set()
|
|
folders = model_search_locations[model_tag]
|
|
for x in folders[0]:
|
|
output_list.update(filter_files_extensions(recursive_search(x), folders[1]))
|
|
return sorted(list(output_list))
|
|
|
|
|
|
# default model types
|
|
register_model_type('checkpoint', supported_ckpt_extensions)
|
|
register_model_type('config', ['.yaml'])
|
|
register_model_type('lora', supported_pt_extensions)
|
|
register_model_type('vae', supported_pt_extensions)
|
|
register_model_type('clip', supported_pt_extensions)
|
|
register_model_type('clip_vision', supported_pt_extensions)
|
|
register_model_type('style_model', supported_pt_extensions)
|
|
register_model_type('embedding', supported_pt_extensions)
|
|
register_model_type('controlnet', supported_pt_extensions)
|
|
register_model_type('upscale_model', supported_pt_extensions)
|
|
|
|
# default search paths
|
|
models_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)), "models")
|
|
add_model_search_path('checkpoint', os.path.join(models_dir, "checkpoints"))
|
|
add_model_search_path('config', os.path.join(models_dir, "configs"))
|
|
add_model_search_path('lora', os.path.join(models_dir, "loras"))
|
|
add_model_search_path('vae', os.path.join(models_dir, "vae"))
|
|
add_model_search_path('clip', os.path.join(models_dir, "clip"))
|
|
add_model_search_path('clip_vision', os.path.join(models_dir, "clip_vision"))
|
|
add_model_search_path('style_model', os.path.join(models_dir, "style_models"))
|
|
add_model_search_path('embedding', os.path.join(models_dir, "embeddings"))
|
|
add_model_search_path('controlnet', os.path.join(models_dir, "controlnet"))
|
|
add_model_search_path('upscale_model', os.path.join(models_dir, "upscale_models"))
|