diff --git a/main.py b/main.py index 4780a9c69..2126763cf 100644 --- a/main.py +++ b/main.py @@ -69,6 +69,13 @@ def execute_prestartup_script(): if args.disable_all_custom_nodes: return + # Get the environment variable and convert it to a list (if not empty) + included_custom_nodes = os.environ.get("INCLUDED_CUSTOM_NODES", "").strip() + included_custom_nodes = ( + [name.strip() for name in included_custom_nodes.split(",")] + if included_custom_nodes else None + ) + node_paths = folder_paths.get_folder_paths("custom_nodes") for custom_node_path in node_paths: possible_modules = os.listdir(custom_node_path) @@ -79,6 +86,10 @@ def execute_prestartup_script(): if os.path.isfile(module_path) or module_path.endswith(".disabled") or module_path == "__pycache__": continue + # Skip modules that are not in INCLUDED_CUSTOM_NODES + if included_custom_nodes and possible_module not in included_custom_nodes: + continue + script_path = os.path.join(module_path, "prestartup_script.py") if os.path.exists(script_path): time_before = time.perf_counter() diff --git a/nodes.py b/nodes.py index 8c1720c1a..92e05b9ca 100644 --- a/nodes.py +++ b/nodes.py @@ -2184,6 +2184,13 @@ def init_external_custom_nodes(): Returns: None """ + # Get the environment variable and convert it to a list (if not empty) + included_custom_nodes = os.environ.get("INCLUDED_CUSTOM_NODES", "").strip() + included_custom_nodes = ( + [name.strip() for name in included_custom_nodes.split(",")] + if included_custom_nodes else None + ) + base_node_names = set(NODE_CLASS_MAPPINGS.keys()) node_paths = folder_paths.get_folder_paths("custom_nodes") node_import_times = [] @@ -2193,6 +2200,10 @@ def init_external_custom_nodes(): possible_modules.remove("__pycache__") for possible_module in possible_modules: + # Skip modules that are not in INCLUDED_CUSTOM_NODES + if included_custom_nodes and possible_module not in included_custom_nodes: + continue + 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 if module_path.endswith(".disabled"): continue