mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2026-01-10 06:10:50 +08:00
98ae55b059
backends and other changes - remove uv.lock since it will not be used in most cases for installation - add cli args to prevent some custom nodes from installing packages at runtime - temp directories can now be shared between workers without being deleted - propcache yanked is now in the dependencies - fix configuration arguments loading in some tests
294 lines
13 KiB
Python
294 lines
13 KiB
Python
from __future__ import annotations
|
|
|
|
import fnmatch
|
|
import importlib
|
|
import importlib.util
|
|
import logging
|
|
import os
|
|
import sys
|
|
import time
|
|
import types
|
|
from contextlib import contextmanager, nullcontext
|
|
from os.path import join, basename, dirname, isdir, isfile, exists, abspath, split, splitext, realpath
|
|
from typing import Iterable, Any, Generator
|
|
from unittest.mock import patch, MagicMock
|
|
|
|
from comfy_compatibility.vanilla import prepare_vanilla_environment, patch_pip_install_subprocess_run, patch_pip_install_popen
|
|
from . import base_nodes
|
|
from .comfyui_v3_package_imports import _comfy_entrypoint_upstream_v3_imports
|
|
from .package_typing import ExportedNodes
|
|
from ..cmd import folder_paths
|
|
from ..component_model.plugins import prompt_server_instance_routes
|
|
from ..distributed.server_stub import ServerStub
|
|
from ..execution_context import current_execution_context
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class StreamToLogger:
|
|
"""
|
|
File-like stream object that redirects writes to a logger instance.
|
|
This is used to capture print() statements from modules during import.
|
|
"""
|
|
|
|
def __init__(self, logger: logging.Logger, log_level=logging.INFO):
|
|
self.logger = logger
|
|
self.log_level = log_level
|
|
|
|
def write(self, buf):
|
|
# Process each line from the buffer. Print statements usually end with a newline.
|
|
for line in buf.rstrip().splitlines():
|
|
# Log the line, removing any trailing whitespace
|
|
self.logger.log(self.log_level, line.rstrip())
|
|
|
|
def flush(self):
|
|
# The logger handles its own flushing, so this can be a no-op.
|
|
pass
|
|
|
|
@property
|
|
def encoding(self):
|
|
return "utf-8"
|
|
|
|
|
|
class _PromptServerStub(ServerStub):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.routes = prompt_server_instance_routes
|
|
self.on_prompt_handlers = []
|
|
|
|
def add_on_prompt_handler(self, handler):
|
|
# todo: these need to be added to a real prompt server if the loading order is behaving in a complex way
|
|
self.on_prompt_handlers.append(handler)
|
|
|
|
def send_sync(self, *args, **kwargs):
|
|
logger.warning(f"Node tried to send a message over the websocket while importing, args={args} kwargs={kwargs}")
|
|
|
|
|
|
def _vanilla_load_importing_execute_prestartup_script(node_paths: Iterable[str]) -> None:
|
|
def execute_script(script_path):
|
|
module_name = splitext(script_path)[0]
|
|
try:
|
|
with _stdout_intercept(module_name):
|
|
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:
|
|
logger.error(f"Failed to execute startup-script: {script_path}", exc_info=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):
|
|
if "comfyui-manager" in module_path.lower():
|
|
os.environ['COMFYUI_PATH'] = str(folder_paths.base_path)
|
|
os.environ['COMFYUI_FOLDERS_BASE_PATH'] = str(folder_paths.models_dir)
|
|
# Monkey-patch ComfyUI-Manager's security check to prevent it from crashing on startup
|
|
# and its logging handler to prevent it from taking over logging.
|
|
glob_path = join(module_path, "glob")
|
|
glob_path_added = False
|
|
original_add_handler = logging.Logger.addHandler
|
|
|
|
def no_op_add_handler(self, handler):
|
|
logger.info(f"Skipping addHandler for {type(handler).__name__} during ComfyUI-Manager prestartup.")
|
|
|
|
try:
|
|
sys.path.insert(0, glob_path)
|
|
glob_path_added = True
|
|
# Patch security_check
|
|
import security_check # pylint: disable=import-error
|
|
original_check = security_check.security_check
|
|
|
|
def patched_security_check():
|
|
try:
|
|
return original_check()
|
|
except Exception as e:
|
|
logger.error(f"ComfyUI-Manager security_check failed but was caught gracefully: {e}", exc_info=e)
|
|
|
|
security_check.security_check = patched_security_check
|
|
logger.debug("Patched ComfyUI-Manager's security_check to fail gracefully.")
|
|
|
|
# Patch logging
|
|
logging.Logger.addHandler = no_op_add_handler
|
|
logger.debug("Patched logging.Logger.addHandler to prevent ComfyUI-Manager from adding a logging handler.")
|
|
|
|
time_before = time.perf_counter()
|
|
success = execute_script(script_path)
|
|
node_prestartup_times.append((time.perf_counter() - time_before, module_path, success))
|
|
except Exception as e:
|
|
logger.error(f"Failed to patch and execute ComfyUI-Manager's prestartup script: {e}", exc_info=e)
|
|
finally:
|
|
if glob_path_added and glob_path in sys.path:
|
|
sys.path.remove(glob_path)
|
|
logging.Logger.addHandler = original_add_handler
|
|
else:
|
|
time_before = time.perf_counter()
|
|
success = execute_script(script_path)
|
|
node_prestartup_times.append((time.perf_counter() - time_before, module_path, success))
|
|
|
|
|
|
@contextmanager
|
|
def _exec_mitigations(module: types.ModuleType, module_path: str) -> Generator[ExportedNodes, Any, None]:
|
|
if module.__name__.lower() in (
|
|
"comfyui-manager",
|
|
"comfyui_ryanonyheinside",
|
|
"comfyui-easy-use",
|
|
"comfyui_custom_nodes_alekpet",
|
|
):
|
|
from ..cmd import folder_paths
|
|
old_file = folder_paths.__file__
|
|
|
|
try:
|
|
# mitigate path
|
|
new_path = join(abspath(join(dirname(old_file), "..", "..")), basename(old_file))
|
|
config = current_execution_context()
|
|
|
|
block_installation = config and config.configuration and config.configuration.block_runtime_package_installation
|
|
with (
|
|
patch.object(folder_paths, "__file__", new_path),
|
|
# mitigate packages installing things dynamically
|
|
patch_pip_install_subprocess_run() if block_installation else nullcontext(),
|
|
patch_pip_install_popen() if block_installation else nullcontext(),
|
|
):
|
|
yield ExportedNodes()
|
|
finally:
|
|
# todo: mitigate "/manager/reboot"
|
|
# todo: mitigate process_wrap
|
|
# todo: unfortunately, we shouldn't restore the patches here, they will have to be applied forever.
|
|
# concurrent.futures.ThreadPoolExecutor = _ThreadPoolExecutor
|
|
# threading.Thread.start = original_thread_start
|
|
logger.info(f"Exec mitigations were applied for {module.__name__}, due to using the folder_paths.__file__ symbol and manipulating EXTENSION_WEB_DIRS")
|
|
else:
|
|
yield ExportedNodes()
|
|
|
|
|
|
@contextmanager
|
|
def _stdout_intercept(name: str):
|
|
original_stdout = sys.stdout
|
|
|
|
try:
|
|
module_logger = logging.getLogger(name)
|
|
sys.stdout = StreamToLogger(module_logger, logging.INFO)
|
|
yield
|
|
finally:
|
|
sys.stdout = original_stdout
|
|
|
|
|
|
def _vanilla_load_custom_nodes_1(module_path, ignore: set = None) -> ExportedNodes:
|
|
if ignore is None:
|
|
ignore = set()
|
|
exported_nodes = ExportedNodes()
|
|
module_name = basename(module_path)
|
|
if isfile(module_path):
|
|
sp = splitext(module_path)
|
|
module_name = sp[0]
|
|
try:
|
|
if isfile(module_path):
|
|
module_spec = importlib.util.spec_from_file_location(module_name, module_path)
|
|
module_dir = split(module_path)[0]
|
|
else:
|
|
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
|
|
|
|
with _exec_mitigations(module, module_path) as mitigated_exported_nodes, _stdout_intercept(module_name):
|
|
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 = 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:
|
|
for name in module.NODE_CLASS_MAPPINGS:
|
|
if name not in ignore:
|
|
exported_nodes.NODE_CLASS_MAPPINGS[name] = module.NODE_CLASS_MAPPINGS[name]
|
|
if hasattr(module, "NODE_DISPLAY_NAME_MAPPINGS") and getattr(module,
|
|
"NODE_DISPLAY_NAME_MAPPINGS") is not None:
|
|
exported_nodes.NODE_DISPLAY_NAME_MAPPINGS.update(module.NODE_DISPLAY_NAME_MAPPINGS)
|
|
else:
|
|
logger.error(f"Skip {module_path} module for custom nodes due to the lack of NODE_CLASS_MAPPINGS.")
|
|
|
|
exported_nodes.update(_comfy_entrypoint_upstream_v3_imports(module))
|
|
except Exception as e:
|
|
logger.error(f"Cannot import {module_path} module for custom nodes:", exc_info=e)
|
|
return exported_nodes
|
|
|
|
|
|
def _vanilla_load_custom_nodes_2(node_paths: Iterable[str]) -> ExportedNodes:
|
|
from ..cli_args import args
|
|
base_node_names = set(base_nodes.NODE_CLASS_MAPPINGS.keys())
|
|
node_import_times = []
|
|
exported_nodes = ExportedNodes()
|
|
for custom_node_path in node_paths:
|
|
if not exists(custom_node_path) or not isdir(custom_node_path):
|
|
continue
|
|
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 = join(custom_node_path, possible_module)
|
|
if isfile(module_path) and splitext(module_path)[1] != ".py": continue
|
|
if module_path.endswith(".disabled"): continue
|
|
if args.disable_all_custom_nodes and possible_module not in args.whitelist_custom_nodes:
|
|
logger.info(f"Skipping {possible_module} due to disable_all_custom_nodes and whitelist_custom_nodes")
|
|
continue
|
|
if any(fnmatch.fnmatch(possible_module, pattern) for pattern in args.blacklist_custom_nodes):
|
|
logger.info(f"Skipping {possible_module} due to blacklist_custom_nodes")
|
|
continue
|
|
time_before = time.perf_counter()
|
|
possible_exported_nodes = _vanilla_load_custom_nodes_1(module_path, ignore=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, import_succeeded))
|
|
exported_nodes.update(possible_exported_nodes)
|
|
|
|
if len(node_import_times) > 0:
|
|
for n in sorted(node_import_times):
|
|
if n[2]:
|
|
import_message = ""
|
|
else:
|
|
import_message = " (IMPORT FAILED)"
|
|
logger.debug(f"{n[0]:6.1f} seconds{import_message}: {n[1]}")
|
|
return exported_nodes
|
|
|
|
|
|
def mitigated_import_of_vanilla_custom_nodes() -> ExportedNodes:
|
|
# only vanilla custom nodes will ever go into the custom_nodes directory
|
|
# this mitigation puts files that custom nodes expects are at the root of the repository back where they should be
|
|
# found. we're in the middle of executing the import of execution and server, in all likelihood, so like all things,
|
|
# the way community custom nodes is pretty radioactive
|
|
# there's a lot of subtle details here, and unfortunately, once this is called, there are some things that have
|
|
# to be activated later, in different places, to make all the hacks necessary for custom nodes to work
|
|
prepare_vanilla_environment()
|
|
|
|
from ..cmd import folder_paths
|
|
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"))]
|
|
|
|
node_paths = frozenset(abspath(custom_node_path) for custom_node_path in node_paths)
|
|
_vanilla_load_importing_execute_prestartup_script(node_paths)
|
|
vanilla_custom_nodes = _vanilla_load_custom_nodes_2(node_paths)
|
|
return vanilla_custom_nodes
|