mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2026-02-08 12:32:32 +08:00
add back any model path cli arg
This commit is contained in:
ljleb
parent
cae4a7fb06
commit
250420fb7b
@ -1,4 +1,5 @@
|
||||
import os
|
||||
from collections import namedtuple
|
||||
|
||||
supported_ckpt_extensions = set(['.ckpt', '.pth'])
|
||||
supported_pt_extensions = set(['.ckpt', '.pt', '.bin', '.pth'])
|
||||
@ -10,31 +11,24 @@ except:
|
||||
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")
|
||||
folder_names_and_paths["checkpoints"] = ([os.path.join(models_dir, "checkpoints")], supported_ckpt_extensions)
|
||||
folder_names_and_paths["configs"] = ([os.path.join(models_dir, "configs")], [".yaml"])
|
||||
|
||||
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 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_folder_path(folder_name, full_folder_path):
|
||||
global folder_names_and_paths
|
||||
if folder_name in folder_names_and_paths:
|
||||
folder_names_and_paths[folder_name][0].append(full_folder_path)
|
||||
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 get_folder_paths(folder_name):
|
||||
return folder_names_and_paths[folder_name][0][:]
|
||||
|
||||
def recursive_search(directory):
|
||||
result = []
|
||||
@ -44,26 +38,50 @@ def recursive_search(directory):
|
||||
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(folder_name, filename):
|
||||
global folder_names_and_paths
|
||||
folders = folder_names_and_paths[folder_name]
|
||||
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(folder_name):
|
||||
global folder_names_and_paths
|
||||
def get_filename_list(model_tag):
|
||||
global model_search_locations
|
||||
output_list = set()
|
||||
folders = folder_names_and_paths[folder_name]
|
||||
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('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
41
main.py
@ -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--port 8188\t\t\tSet the listen port.")
|
||||
|
||||
s = os.path.sep
|
||||
print(f"\t--ckpt-dir path{s}to{s}dir\t\t\tAdd a path to a checkpoints dir.")
|
||||
print(f"\t--clip-dir path{s}to{s}dir\t\t\tAdd a path to a clip dir.")
|
||||
print(f"\t--clip-vision-dir path{s}to{s}dir\tAdd a path to a clip vision dir.")
|
||||
print(f"\t--controlnet-dir path{s}to{s}dir\tAdd a path to a controlnet checkpoints dir.")
|
||||
print(f"\t--embed-dir path{s}to{s}dir\t\t\tAdd a path to an embeddings dir.")
|
||||
print(f"\t--lora-dir path{s}to{s}dir\t\t\tAdd a path to a lora models dir.")
|
||||
print(f"\t--style-model-dir path{s}to{s}dir\tAdd a path to a style models dir.")
|
||||
print(f"\t--t2i-dir path{s}to{s}dir\t\t\tAdd a path to a T2I style adapter dir.")
|
||||
print(f"\t--upscaler-dir path{s}to{s}dir\t\tAdd a path to an upscale models dir.")
|
||||
print(f"\t--vae-dir path{s}to{s}dir\t\t\tAdd a path to a vae dir.")
|
||||
dummy_path = os.path.join('path', 'to', 'dir')
|
||||
print(f"\t--ckpt-dir {dummy_path}\t\t\tAdd a path to a checkpoint directory.")
|
||||
print(f"\t--config-dir {dummy_path}\t\t\tAdd a path to a model config directory.")
|
||||
print(f"\t--clip-dir {dummy_path}\t\t\tAdd a path to a clip directory.")
|
||||
print(f"\t--clip-vision-dir {dummy_path}\tAdd a path to a clip vision directory.")
|
||||
print(f"\t--controlnet-dir {dummy_path}\tAdd a path to a controlnet or T2I checkpoint directory.")
|
||||
print(f"\t--embedding-dir {dummy_path}\t\t\tAdd a path to an embedding directory.")
|
||||
print(f"\t--lora-dir {dummy_path}\t\t\tAdd a path to a lora model directory.")
|
||||
print(f"\t--style-model-dir {dummy_path}\tAdd a path to a style model directory.")
|
||||
print(f"\t--upscale-model-dir {dummy_path}\t\tAdd a path to an upscale model directory.")
|
||||
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--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):
|
||||
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):
|
||||
with open(yaml_path, 'r') as stream:
|
||||
config = yaml.safe_load(stream)
|
||||
@ -92,7 +104,7 @@ def load_extra_path_config(yaml_path):
|
||||
if base_path is not None:
|
||||
full_path = os.path.join(base_path, 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__":
|
||||
cleanup_temp()
|
||||
@ -118,10 +130,7 @@ if __name__ == "__main__":
|
||||
if os.path.isfile(extra_model_paths_config_path):
|
||||
load_extra_path_config(extra_model_paths_config_path)
|
||||
|
||||
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])
|
||||
add_model_search_paths()
|
||||
|
||||
port = 8188
|
||||
try:
|
||||
|
||||
Loading…
Reference in New Issue
Block a user