ComfyUI Manager now starts successfully, but needs more mitigations:

- /manager/reboot needs to use a different approach to restart the
   currently running Python process.
 - runpy should be used for install.py invocations
This commit is contained in:
doctorpangloss 2024-03-29 16:25:29 -07:00
parent 8f548d4d19
commit bd87697fdf
2 changed files with 99 additions and 19 deletions

View File

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

3
comfy/web/extensions/.gitignore vendored Normal file
View File

@ -0,0 +1,3 @@
*
!core/
!logging.js.example