wangbo/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2026-02-16 16:32:34 +08:00
Code Issues Actions 11 Packages Projects Releases Wiki Activity

add back any model path cli arg

Browse Source
This commit is contained in:
ljleb 2023-03-19 19:44:05 -04:00
parent cae4a7fb06
commit 250420fb7b
2 changed files with 70 additions and 43 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

72
folder_paths.py
View File

@ -1,4 +1,5 @@
import os import os
from collections import namedtuple
supported_ckpt_extensions = set(['.ckpt', '.pth']) supported_ckpt_extensions = set(['.ckpt', '.pth'])
supported_pt_extensions = set(['.ckpt', '.pt', '.bin', '.pth']) supported_pt_extensions = set(['.ckpt', '.pt', '.bin', '.pth'])
@ -10,31 +11,24 @@ except:
print("Could not import safetensors, safetensors support disabled.") print("Could not import safetensors, safetensors support disabled.")
folder_names_and_paths = {} SearchLocation = namedtuple('SearchLocation', ['paths', 'extensions'])
model_search_locations = {}
models_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)), "models") def register_model_type(model_tag, supported_extensions):
folder_names_and_paths["checkpoints"] = ([os.path.join(models_dir, "checkpoints")], supported_ckpt_extensions) if model_tag not in model_search_locations:
folder_names_and_paths["configs"] = ([os.path.join(models_dir, "configs")], [".yaml"]) model_search_locations[model_tag] = SearchLocation([], supported_extensions)
folder_names_and_paths["loras"] = ([os.path.join(models_dir, "loras")], supported_pt_extensions)
folder_names_and_paths["vae"] = ([os.path.join(models_dir, "vae")], supported_pt_extensions)
folder_names_and_paths["clip"] = ([os.path.join(models_dir, "clip")], supported_pt_extensions)
folder_names_and_paths["clip_vision"] = ([os.path.join(models_dir, "clip_vision")], supported_pt_extensions)
folder_names_and_paths["style_models"] = ([os.path.join(models_dir, "style_models")], supported_pt_extensions)
folder_names_and_paths["embeddings"] = ([os.path.join(models_dir, "embeddings")], supported_pt_extensions)
folder_names_and_paths["controlnet"] = ([os.path.join(models_dir, "controlnet"), os.path.join(models_dir, "t2i_adapter")], supported_pt_extensions)
folder_names_and_paths["upscale_models"] = ([os.path.join(models_dir, "upscale_models")], supported_pt_extensions)
def add_model_folder_path(folder_name, full_folder_path): def add_model_search_path(model_tag, full_folder_path):
global folder_names_and_paths global model_search_locations
if folder_name in folder_names_and_paths: if model_tag in model_search_locations:
folder_names_and_paths[folder_name][0].append(full_folder_path) model_search_locations[model_tag].paths.append(full_folder_path)
def get_folder_paths(model_tag):
return model_search_locations[model_tag][0][:]
def get_folder_paths(folder_name):
return folder_names_and_paths[folder_name][0][:]
def recursive_search(directory): def recursive_search(directory):
result = [] result = []
@ -44,26 +38,50 @@ def recursive_search(directory):
result.append(os.path.join(root, filepath).replace(os.path.join(directory,''),'')) result.append(os.path.join(root, filepath).replace(os.path.join(directory,''),''))
return result return result
def filter_files_extensions(files, extensions): def filter_files_extensions(files, extensions):
return sorted(list(filter(lambda a: os.path.splitext(a)[-1].lower() in extensions, files))) return sorted(list(filter(lambda a: os.path.splitext(a)[-1].lower() in extensions, files)))
def get_full_path(model_tag, filename):
def get_full_path(folder_name, filename): global model_search_locations
global folder_names_and_paths folders = model_search_locations[model_tag]
folders = folder_names_and_paths[folder_name]
for x in folders[0]: for x in folders[0]:
full_path = os.path.join(x, filename) full_path = os.path.join(x, filename)
if os.path.isfile(full_path): if os.path.isfile(full_path):
return full_path return full_path
def get_filename_list(folder_name): def get_filename_list(model_tag):
global folder_names_and_paths global model_search_locations
output_list = set() output_list = set()
folders = folder_names_and_paths[folder_name] folders = model_search_locations[model_tag]
for x in folders[0]: for x in folders[0]:
output_list.update(filter_files_extensions(recursive_search(x), folders[1])) output_list.update(filter_files_extensions(recursive_search(x), folders[1]))
return sorted(list(output_list)) return sorted(list(output_list))
# default model types
register_model_type('ckpt', 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('ckpt', 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"))

41
main.py
View File

@ -15,17 +15,17 @@ if __name__ == "__main__":
print("\t--listen\t\t\tListen on 0.0.0.0 so the UI can be accessed from other computers.") print("\t--listen\t\t\tListen on 0.0.0.0 so the UI can be accessed from other computers.")
print("\t--port 8188\t\t\tSet the listen port.") print("\t--port 8188\t\t\tSet the listen port.")
s = os.path.sep dummy_path = os.path.join('path', 'to', 'dir')
print(f"\t--ckpt-dir path{s}to{s}dir\t\t\tAdd a path to a checkpoints dir.") print(f"\t--ckpt-dir {dummy_path}\t\t\tAdd a path to a checkpoint directory.")
print(f"\t--clip-dir path{s}to{s}dir\t\t\tAdd a path to a clip dir.") print(f"\t--config-dir {dummy_path}\t\t\tAdd a path to a model config directory.")
print(f"\t--clip-vision-dir path{s}to{s}dir\tAdd a path to a clip vision dir.") print(f"\t--clip-dir {dummy_path}\t\t\tAdd a path to a clip directory.")
print(f"\t--controlnet-dir path{s}to{s}dir\tAdd a path to a controlnet checkpoints dir.") print(f"\t--clip-vision-dir {dummy_path}\tAdd a path to a clip vision directory.")
print(f"\t--embed-dir path{s}to{s}dir\t\t\tAdd a path to an embeddings dir.") print(f"\t--controlnet-dir {dummy_path}\tAdd a path to a controlnet or T2I checkpoint directory.")
print(f"\t--lora-dir path{s}to{s}dir\t\t\tAdd a path to a lora models dir.") print(f"\t--embedding-dir {dummy_path}\t\t\tAdd a path to an embedding directory.")
print(f"\t--style-model-dir path{s}to{s}dir\tAdd a path to a style models dir.") print(f"\t--lora-dir {dummy_path}\t\t\tAdd a path to a lora model directory.")
print(f"\t--t2i-dir path{s}to{s}dir\t\t\tAdd a path to a T2I style adapter dir.") print(f"\t--style-model-dir {dummy_path}\tAdd a path to a style model directory.")
print(f"\t--upscaler-dir path{s}to{s}dir\t\tAdd a path to an upscale models dir.") print(f"\t--upscale-model-dir {dummy_path}\t\tAdd a path to an upscale model directory.")
print(f"\t--vae-dir path{s}to{s}dir\t\t\tAdd a path to a vae dir.") print(f"\t--vae-dir {dummy_path}\t\t\tAdd a path to a vae directory.")
print("\t--dont-upcast-attention\t\tDisable upcasting of attention \n\t\t\t\t\tcan boost speed but increase the chances of black images.\n") print("\t--dont-upcast-attention\t\tDisable upcasting of attention \n\t\t\t\t\tcan boost speed but increase the chances of black images.\n")
print("\t--use-split-cross-attention\tUse the split cross attention optimization instead of the sub-quadratic one.\n\t\t\t\t\tIgnored when xformers is used.") print("\t--use-split-cross-attention\tUse the split cross attention optimization instead of the sub-quadratic one.\n\t\t\t\t\tIgnored when xformers is used.")
@ -74,6 +74,18 @@ def cleanup_temp():
if os.path.exists(temp_dir): if os.path.exists(temp_dir):
shutil.rmtree(temp_dir, ignore_errors=True) shutil.rmtree(temp_dir, ignore_errors=True)
def add_model_search_paths():
if '--extra-model-paths-config' in sys.argv:
indices = [(i + 1) for i in range(len(sys.argv) - 1) if sys.argv[i] == '--extra-model-paths-config']
for i in indices:
load_extra_path_config(sys.argv[i])
for model_tag, search_location in folder_paths.model_search_locations.items():
model_path_arg = f"--{model_tag.replace('_', '-')}-dir"
extra_paths = [sys.argv[i + 1] for i, e in enumerate(sys.argv[:-1]) if e == model_path_arg]
for extra_path in extra_paths:
folder_paths.add_model_search_path(model_tag, extra_path)
def load_extra_path_config(yaml_path): def load_extra_path_config(yaml_path):
with open(yaml_path, 'r') as stream: with open(yaml_path, 'r') as stream:
config = yaml.safe_load(stream) config = yaml.safe_load(stream)
@ -92,7 +104,7 @@ def load_extra_path_config(yaml_path):
if base_path is not None: if base_path is not None:
full_path = os.path.join(base_path, full_path) full_path = os.path.join(base_path, full_path)
print("Adding extra search path", x, full_path) print("Adding extra search path", x, full_path)
folder_paths.add_model_folder_path(x, full_path) folder_paths.add_model_search_path(x, full_path)
if __name__ == "__main__": if __name__ == "__main__":
cleanup_temp() cleanup_temp()
@ -118,10 +130,7 @@ if __name__ == "__main__":
if os.path.isfile(extra_model_paths_config_path): if os.path.isfile(extra_model_paths_config_path):
load_extra_path_config(extra_model_paths_config_path) load_extra_path_config(extra_model_paths_config_path)
if '--extra-model-paths-config' in sys.argv: add_model_search_paths()
indices = [(i + 1) for i in range(len(sys.argv) - 1) if sys.argv[i] == '--extra-model-paths-config']
for i in indices:
load_extra_path_config(sys.argv[i])
port = 8188 port = 8188
try: try: