From 74fc7b772656a59b344508480632d9d45f9127de Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Sun, 16 Apr 2023 01:36:15 -0400 Subject: [PATCH] custom_nodes paths can now be set in the extra_model_paths.yaml --- extra_model_paths.yaml.example | 2 +- folder_paths.py | 7 +++++-- main.py | 15 ++++++++------- nodes.py | 19 ++++++++++--------- 4 files changed, 24 insertions(+), 19 deletions(-) diff --git a/extra_model_paths.yaml.example b/extra_model_paths.yaml.example index af784fd6..f421f54d 100644 --- a/extra_model_paths.yaml.example +++ b/extra_model_paths.yaml.example @@ -18,6 +18,6 @@ a111: #other_ui: # base_path: path/to/ui # checkpoints: models/checkpoints - +# custom_nodes: path/custom_nodes diff --git a/folder_paths.py b/folder_paths.py index ab335934..61f446c9 100644 --- a/folder_paths.py +++ b/folder_paths.py @@ -12,8 +12,8 @@ except: folder_names_and_paths = {} - -models_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)), "models") +base_path = os.path.dirname(os.path.realpath(__file__)) +models_dir = os.path.join(base_path, "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"]) @@ -28,6 +28,9 @@ folder_names_and_paths["diffusers"] = ([os.path.join(models_dir, "diffusers")], 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) +folder_names_and_paths["custom_nodes"] = ([os.path.join(base_path, "custom_nodes")], []) + + output_directory = os.path.join(os.path.dirname(os.path.realpath(__file__)), "output") temp_directory = os.path.join(os.path.dirname(os.path.realpath(__file__)), "temp") input_directory = os.path.join(os.path.dirname(os.path.realpath(__file__)), "input") diff --git a/main.py b/main.py index 9c0a3d8a..02c700eb 100644 --- a/main.py +++ b/main.py @@ -81,6 +81,14 @@ if __name__ == "__main__": server = server.PromptServer(loop) q = execution.PromptQueue(server) + extra_model_paths_config_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), "extra_model_paths.yaml") + if os.path.isfile(extra_model_paths_config_path): + load_extra_path_config(extra_model_paths_config_path) + + if args.extra_model_paths_config: + for config_path in itertools.chain(*args.extra_model_paths_config): + load_extra_path_config(config_path) + init_custom_nodes() server.add_routes() hijack_progress(server) @@ -91,13 +99,6 @@ if __name__ == "__main__": dont_print = args.dont_print_server - extra_model_paths_config_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), "extra_model_paths.yaml") - if os.path.isfile(extra_model_paths_config_path): - load_extra_path_config(extra_model_paths_config_path) - - if args.extra_model_paths_config: - for config_path in itertools.chain(*args.extra_model_paths_config): - load_extra_path_config(config_path) if args.output_directory: output_dir = os.path.abspath(args.output_directory) diff --git a/nodes.py b/nodes.py index e6ad9434..c775da00 100644 --- a/nodes.py +++ b/nodes.py @@ -1178,15 +1178,16 @@ def load_custom_node(module_path): print(f"Cannot import {module_path} module for custom nodes:", e) def load_custom_nodes(): - CUSTOM_NODE_PATH = os.path.join(os.path.dirname(os.path.realpath(__file__)), "custom_nodes") - possible_modules = os.listdir(CUSTOM_NODE_PATH) - if "__pycache__" in possible_modules: - possible_modules.remove("__pycache__") - - for possible_module in possible_modules: - module_path = os.path.join(CUSTOM_NODE_PATH, possible_module) - if os.path.isfile(module_path) and os.path.splitext(module_path)[1] != ".py": continue - load_custom_node(module_path) + node_paths = folder_paths.get_folder_paths("custom_nodes") + for custom_node_path in node_paths: + possible_modules = os.listdir(custom_node_path) + if "__pycache__" in possible_modules: + possible_modules.remove("__pycache__") + + for possible_module in possible_modules: + module_path = os.path.join(custom_node_path, possible_module) + if os.path.isfile(module_path) and os.path.splitext(module_path)[1] != ".py": continue + load_custom_node(module_path) def init_custom_nodes(): load_custom_nodes()