Add support for including custom nodes via environment variable

This commit is contained in:
younes127001 2025-02-06 02:28:51 +01:00
parent f1059b0b82
commit c2e4b96f85
2 changed files with 24 additions and 0 deletions

12
main.py
View File

@ -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,11 @@ 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:
logging.debug(f"Skipping {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()

View File

@ -2144,6 +2144,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 = []
@ -2153,6 +2160,11 @@ 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:
logging.debug(f"Skipping {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