Add `--no-custom-node` cmd flag (#3903)

* Add --no-custom-node cmd flag

* nit
This commit is contained in:
Chenlei Hu 2024-07-01 17:54:03 -04:00 committed by GitHub
parent b82d67d5bf
commit 9dd549e253
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
4 changed files with 37 additions and 9 deletions

View File

@ -118,6 +118,7 @@ parser.add_argument("--quick-test-for-ci", action="store_true", help="Quick test
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.")

View File

@ -2,9 +2,11 @@ import os
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["gligen"] = ([os.path.join(models_dir, "gligen")], suppor
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)

13
main.py
View File

@ -5,6 +5,8 @@ import os
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_prestartup_script():
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 @@ import yaml
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 @@ if __name__ == "__main__":
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()

View File

@ -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: