diff --git a/comfy/nodes/vanilla_node_importing.py b/comfy/nodes/vanilla_node_importing.py index 911546bbb..478b6ee58 100644 --- a/comfy/nodes/vanilla_node_importing.py +++ b/comfy/nodes/vanilla_node_importing.py @@ -1,37 +1,106 @@ from __future__ import annotations import importlib +import logging import os +import shutil import sys import time import types -from typing import Dict +from contextlib import contextmanager +from typing import Dict, List +from os.path import join, basename, dirname, isdir, isfile, exists, abspath, split, splitext, realpath from . import base_nodes from .package_typing import ExportedNodes +def _vanilla_load_importing_execute_prestartup_script(node_paths: List[str]) -> None: + def execute_script(script_path): + module_name = splitext(script_path)[0] + try: + spec = importlib.util.spec_from_file_location(module_name, script_path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return True + except Exception as e: + print(f"Failed to execute startup-script: {script_path} / {e}") + return False + + node_prestartup_times = [] + for custom_node_path in node_paths: + # patched + if not isdir(custom_node_path): + continue + # end patch + possible_modules = os.listdir(custom_node_path) + + for possible_module in possible_modules: + module_path = join(custom_node_path, possible_module) + if isfile(module_path) or module_path.endswith(".disabled") or module_path == "__pycache__": + continue + + script_path = join(module_path, "prestartup_script.py") + if exists(script_path): + time_before = time.perf_counter() + success = execute_script(script_path) + node_prestartup_times.append((time.perf_counter() - time_before, module_path, success)) + if len(node_prestartup_times) > 0: + print("\nPrestartup times for custom nodes:") + for n in sorted(node_prestartup_times): + if n[2]: + import_message = "" + else: + import_message = " (PRESTARTUP FAILED)" + print("{:6.1f} seconds{}:".format(n[0], import_message), n[1]) + print() + + +@contextmanager +def _exec_mitigations(module: types.ModuleType, module_path: str) -> ExportedNodes: + if module.__name__ == "ComfyUI-Manager": + from ..cmd import folder_paths + old_file = folder_paths.__file__ + + try: + # mitigate path + new_path = join(abspath(join(dirname(old_file), "..", "..")), basename(old_file)) + folder_paths.__file__ = new_path + # mitigate JS copy + sys.modules['nodes'].EXTENSION_WEB_DIRS = {} + yield ExportedNodes() + finally: + folder_paths.__file__ = old_file + # todo: mitigate "/manager/reboot" + # todo: mitigate process_wrap + else: + yield ExportedNodes() + + def _vanilla_load_custom_nodes_1(module_path, ignore=set()) -> ExportedNodes: exported_nodes = ExportedNodes() - module_name = os.path.basename(module_path) - if os.path.isfile(module_path): - sp = os.path.splitext(module_path) + module_name = basename(module_path) + if isfile(module_path): + sp = splitext(module_path) module_name = sp[0] try: - if os.path.isfile(module_path): + if isfile(module_path): module_spec = importlib.util.spec_from_file_location(module_name, module_path) - module_dir = os.path.split(module_path)[0] + module_dir = split(module_path)[0] else: - module_spec = importlib.util.spec_from_file_location(module_name, os.path.join(module_path, "__init__.py")) + module_spec = importlib.util.spec_from_file_location(module_name, join(module_path, "__init__.py")) module_dir = module_path module = importlib.util.module_from_spec(module_spec) sys.modules[module_name] = module - module_spec.loader.exec_module(module) + + with _exec_mitigations(module, module_path) as mitigated_exported_nodes: + module_spec.loader.exec_module(module) + exported_nodes.update(mitigated_exported_nodes) if hasattr(module, "WEB_DIRECTORY") and getattr(module, "WEB_DIRECTORY") is not None: - web_dir = os.path.abspath(os.path.join(module_dir, getattr(module, "WEB_DIRECTORY"))) - if os.path.isdir(web_dir): + web_dir = abspath(join(module_dir, getattr(module, "WEB_DIRECTORY"))) + if isdir(web_dir): exported_nodes.EXTENSION_WEB_DIRS[module_name] = web_dir if hasattr(module, "NODE_CLASS_MAPPINGS") and getattr(module, "NODE_CLASS_MAPPINGS") is not None: @@ -52,27 +121,27 @@ def _vanilla_load_custom_nodes_1(module_path, ignore=set()) -> ExportedNodes: return exported_nodes -def _vanilla_load_custom_nodes_2() -> ExportedNodes: - from ..cmd import folder_paths +def _vanilla_load_custom_nodes_2(node_paths: List[str]) -> ExportedNodes: base_node_names = set(base_nodes.NODE_CLASS_MAPPINGS.keys()) - node_paths = folder_paths.get_folder_paths("custom_nodes") node_import_times = [] exported_nodes = ExportedNodes() for custom_node_path in node_paths: - if not os.path.exists(custom_node_path) or not os.path.isdir(custom_node_path): + if not exists(custom_node_path) or not isdir(custom_node_path): continue - possible_modules = os.listdir(os.path.realpath(custom_node_path)) + possible_modules = os.listdir(realpath(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 + module_path = join(custom_node_path, possible_module) + if isfile(module_path) and splitext(module_path)[1] != ".py": continue if module_path.endswith(".disabled"): continue time_before = time.perf_counter() possible_exported_nodes = _vanilla_load_custom_nodes_1(module_path, base_node_names) + # comfyui-manager mitigation + import_succeeded = len(possible_exported_nodes.NODE_CLASS_MAPPINGS) > 0 or "ComfyUI-Manager" in module_path node_import_times.append( - (time.perf_counter() - time_before, module_path, len(possible_exported_nodes.NODE_CLASS_MAPPINGS) > 0)) + (time.perf_counter() - time_before, module_path, import_succeeded)) exported_nodes.update(possible_exported_nodes) if len(node_import_times) > 0: @@ -116,5 +185,13 @@ def mitigated_import_of_vanilla_custom_nodes() -> ExportedNodes: setattr(comfy_extras, module_short_name, module) comfy_extras_mitigation[f'comfy_extras.{module_short_name}'] = module sys.modules.update(comfy_extras_mitigation) - vanilla_custom_nodes = _vanilla_load_custom_nodes_2() + node_paths = folder_paths.get_folder_paths("custom_nodes") + + potential_git_dir_parent = join(dirname(__file__), "..", "..") + is_git_repository = exists(join(potential_git_dir_parent, ".git")) + if is_git_repository: + node_paths += [abspath(join(potential_git_dir_parent, "custom_nodes"))] + + _vanilla_load_importing_execute_prestartup_script(node_paths) + vanilla_custom_nodes = _vanilla_load_custom_nodes_2(node_paths) return vanilla_custom_nodes diff --git a/comfy/web/extensions/.gitignore b/comfy/web/extensions/.gitignore new file mode 100644 index 000000000..1c75447a9 --- /dev/null +++ b/comfy/web/extensions/.gitignore @@ -0,0 +1,3 @@ +* +!core/ +!logging.js.example \ No newline at end of file