comfyanonymous · ljleb · Mar 14, 2023 · Mar 14, 2023 · Mar 16, 2023 · Mar 19, 2023
@@ -8,3 +8,5 @@ temp/
 custom_nodes/
 !custom_nodes/example_node.py.example
 extra_model_paths.yaml
+venv/
+.idea/
@@ -1,4 +1,3 @@
-import os
 from comfy_extras.chainner_models import model_loading
 from comfy.sd import load_torch_file
 import model_management
@@ -17,7 +16,7 @@ def INPUT_TYPES(s):
     CATEGORY = "loaders"
 
     def load_model(self, model_name):
-        model_path = folder_paths.get_full_path("upscale_models", model_name)
+        model_path = folder_paths.get_full_path("upscale_model", model_name)
         sd = load_torch_file(model_path)
         out = model_loading.load_state_dict(sd).eval()
         return (out, )

@@ -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 @@
     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"])
+def register_model_type(model_tag, supported_extensions):
+    if model_tag not in model_search_locations:
+        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_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 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 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('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"))
@@ -14,6 +14,19 @@
         print("Valid Command line Arguments:")
         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.")
+
+        dummy_path = os.path.join('path', 'to', 'dir')
+        print(f"\t--checkpoint-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.")
         print("\t--use-pytorch-cross-attention\tUse the new pytorch 2.0 cross attention function.")
@@ -61,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)
@@ -105,10 +130,7 @@ def load_extra_path_config(yaml_path):
     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:

@@ -198,8 +198,8 @@ def INPUT_TYPES(s):
     CATEGORY = "loaders"
 
     def load_checkpoint(self, config_name, ckpt_name, output_vae=True, output_clip=True):
-        config_path = folder_paths.get_full_path("configs", config_name)
-        ckpt_path = folder_paths.get_full_path("checkpoints", ckpt_name)
+        config_path = folder_paths.get_full_path("config", config_name)
+        ckpt_path = folder_paths.get_full_path("checkpoint", ckpt_name)
         return comfy.sd.load_checkpoint(config_path, ckpt_path, output_vae=True, output_clip=True, embedding_directory=folder_paths.get_folder_paths("embeddings"))
 
 class CheckpointLoaderSimple:
@@ -213,7 +213,7 @@ def INPUT_TYPES(s):
     CATEGORY = "loaders"
 
     def load_checkpoint(self, ckpt_name, output_vae=True, output_clip=True):
-        ckpt_path = folder_paths.get_full_path("checkpoints", ckpt_name)
+        ckpt_path = folder_paths.get_full_path("checkpoint", ckpt_name)
         out = comfy.sd.load_checkpoint_guess_config(ckpt_path, output_vae=True, output_clip=True, embedding_directory=folder_paths.get_folder_paths("embeddings"))
         return out
 
@@ -248,7 +248,7 @@ def INPUT_TYPES(s):
     CATEGORY = "loaders"
 
     def load_lora(self, model, clip, lora_name, strength_model, strength_clip):
-        lora_path = folder_paths.get_full_path("loras", lora_name)
+        lora_path = folder_paths.get_full_path("lora", lora_name)
         model_lora, clip_lora = comfy.sd.load_lora_for_models(model, clip, lora_path, strength_model, strength_clip)
         return (model_lora, clip_lora)
 
@@ -381,7 +381,7 @@ def INPUT_TYPES(s):
     CATEGORY = "loaders"
 
     def load_style_model(self, style_model_name):
-        style_model_path = folder_paths.get_full_path("style_models", style_model_name)
+        style_model_path = folder_paths.get_full_path("style_model", style_model_name)
         style_model = comfy.sd.load_style_model(style_model_path)
         return (style_model,)