From 9dd549e253f0e2c5e6360f9349b0872af18d8962 Mon Sep 17 00:00:00 2001 From: Chenlei Hu Date: Mon, 1 Jul 2024 17:54:03 -0400 Subject: [PATCH] Add `--no-custom-node` cmd flag (#3903) * Add --no-custom-node cmd flag * nit --- comfy/cli_args.py | 1 + folder_paths.py | 8 +++++--- main.py | 13 +++++++++++-- nodes.py | 24 ++++++++++++++++++++---- 4 files changed, 37 insertions(+), 9 deletions(-) diff --git a/comfy/cli_args.py b/comfy/cli_args.py index fb0d37ce750..b72bf3998ae 100644 --- a/comfy/cli_args.py +++ b/comfy/cli_args.py @@ -118,6 +118,7 @@ class LatentPreviewMethod(enum.Enum): parser.add_argument("--windows-standalone-build", action="store_true", help="Windows standalone build: Enable convenient things that most people using the standalone windows build will probably enjoy (like auto opening the page on startup).") parser.add_argument("--disable-metadata", action="store_true", help="Disable saving prompt metadata in files.") +parser.add_argument("--disable-all-custom-nodes", action="store_true", help="Disable loading all custom nodes.") parser.add_argument("--multi-user", action="store_true", help="Enables per-user storage.") diff --git a/folder_paths.py b/folder_paths.py index 234b734095e..2cf45f12a34 100644 --- a/folder_paths.py +++ b/folder_paths.py @@ -2,9 +2,11 @@ import time import logging -supported_pt_extensions = set(['.ckpt', '.pt', '.bin', '.pth', '.safetensors', '.pkl']) +supported_pt_extensions: set[str] = set(['.ckpt', '.pt', '.bin', '.pth', '.safetensors', '.pkl']) -folder_names_and_paths = {} +SupportedFileExtensionsType = set[str] +ScanPathType = list[str] +folder_names_and_paths: dict[str, tuple[ScanPathType, SupportedFileExtensionsType]] = {} base_path = os.path.dirname(os.path.realpath(__file__)) models_dir = os.path.join(base_path, "models") @@ -26,7 +28,7 @@ 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")], []) +folder_names_and_paths["custom_nodes"] = ([os.path.join(base_path, "custom_nodes")], set()) folder_names_and_paths["hypernetworks"] = ([os.path.join(models_dir, "hypernetworks")], supported_pt_extensions) diff --git a/main.py b/main.py index a374f2b124a..2957dd2ffd6 100644 --- a/main.py +++ b/main.py @@ -5,6 +5,8 @@ import importlib.util import folder_paths import time +from comfy.cli_args import args + def execute_prestartup_script(): def execute_script(script_path): @@ -18,6 +20,9 @@ def execute_script(script_path): print(f"Failed to execute startup-script: {script_path} / {e}") return False + if args.disable_all_custom_nodes: + return + node_paths = folder_paths.get_folder_paths("custom_nodes") for custom_node_path in node_paths: possible_modules = os.listdir(custom_node_path) @@ -76,7 +81,7 @@ def execute_script(script_path): import execution import server from server import BinaryEventTypes -from nodes import init_custom_nodes +from nodes import init_builtin_extra_nodes, init_external_custom_nodes import comfy.model_management def cuda_malloc_warning(): @@ -214,7 +219,11 @@ def load_extra_path_config(yaml_path): for config_path in itertools.chain(*args.extra_model_paths_config): load_extra_path_config(config_path) - init_custom_nodes() + init_builtin_extra_nodes() + if not args.disable_all_custom_nodes: + init_external_custom_nodes() + else: + logging.info("Skipping loading of custom nodes") cuda_malloc_warning() diff --git a/nodes.py b/nodes.py index 04775d2906f..a0d2178f463 100644 --- a/nodes.py +++ b/nodes.py @@ -1925,7 +1925,16 @@ def load_custom_node(module_path, ignore=set()): logging.warning(f"Cannot import {module_path} module for custom nodes: {e}") return False -def load_custom_nodes(): +def init_external_custom_nodes(): + """ + Initializes the external custom nodes. + + This function loads custom nodes from the specified folder paths and imports them into the application. + It measures the import times for each custom node and logs the results. + + Returns: + None + """ base_node_names = set(NODE_CLASS_MAPPINGS.keys()) node_paths = folder_paths.get_folder_paths("custom_nodes") node_import_times = [] @@ -1952,7 +1961,16 @@ def load_custom_nodes(): logging.info("{:6.1f} seconds{}: {}".format(n[0], import_message, n[1])) logging.info("") -def init_custom_nodes(): +def init_builtin_extra_nodes(): + """ + Initializes the built-in extra nodes in ComfyUI. + + This function loads the extra node files located in the "comfy_extras" directory and imports them into ComfyUI. + If any of the extra node files fail to import, a warning message is logged. + + Returns: + None + """ extras_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)), "comfy_extras") extras_files = [ "nodes_latent.py", @@ -1999,8 +2017,6 @@ def init_custom_nodes(): if not load_custom_node(os.path.join(extras_dir, node_file)): import_failed.append(node_file) - load_custom_nodes() - if len(import_failed) > 0: logging.warning("WARNING: some comfy_extras/ nodes did not import correctly. This may be because they are missing some dependencies.\n") for node in import_failed: