From e569e361a63087aec82ef254588209f4a1adc006 Mon Sep 17 00:00:00 2001 From: Peyton DeNiro <25550995+pydn@users.noreply.github.com> Date: Sun, 10 May 2026 00:45:13 -0500 Subject: [PATCH] fix: load comfy_extras nodes even when server.py import fails server.py cannot load due to a ComfyUI internal module shadowing bug (comfy/utils.py file shadows utils/ package). Decouple init_extra_nodes() from server.py loading so NODE_CLASS_MAPPINGS still gets populated from comfy_extras (UpscaleModelLoader, etc.) even without full PromptServer setup. Also fix app.py to re-read from cached sys.modules[nodes] after import_custom_nodes() instead of using stale initial copy. Fixes the upscale-model-loader runtime e2e test. --- comfyui_to_python/app.py | 8 + comfyui_to_python/node_runtime.py | 9 +- .../runtime/generated/upscale-model-loader.py | 241 ++++++++++++++---- 3 files changed, 202 insertions(+), 56 deletions(-) diff --git a/comfyui_to_python/app.py b/comfyui_to_python/app.py index 334d06f..acee636 100644 --- a/comfyui_to_python/app.py +++ b/comfyui_to_python/app.py @@ -1,4 +1,5 @@ import copy +import sys from typing import TextIO @@ -57,6 +58,13 @@ class ExportApplication: } if self.needs_init_custom_nodes or missing_node_types: self.custom_node_importer() + # Re-read from the cached "nodes" module in sys.modules after + # import_custom_nodes() populates it with extras (comfy_extras, + # custom node directories, etc.). The original dict is a stale copy. + nodes_mod = sys.modules.get("nodes") + if nodes_mod is not None: + fresh_mappings = getattr(nodes_mod, "NODE_CLASS_MAPPINGS", {}) + self.node_class_mappings = fresh_mappings self.base_node_class_mappings = copy.deepcopy(self.node_class_mappings) load_order = LoadOrderDeterminer( diff --git a/comfyui_to_python/node_runtime.py b/comfyui_to_python/node_runtime.py index 61a7f72..d4782f8 100644 --- a/comfyui_to_python/node_runtime.py +++ b/comfyui_to_python/node_runtime.py @@ -261,7 +261,14 @@ def import_custom_nodes() -> None: server_mod = _load_module("server", os.path.join(comfyui_path, "server.py")) if execution_mod is None or server_mod is None: - log.debug("import_custom_nodes: could not load required modules") + log.debug( + "import_custom_nodes: could not load execution/server modules. " + "Proceeding without full PromptServer/PromptQueue setup." + ) + # Even without server, we can still populate NODE_CLASS_MAPPINGS from + # comfy_extras and built-in extra nodes. + if nodes_mod is not None and hasattr(nodes_mod, "init_extra_nodes"): + asyncio.run(nodes_mod.init_extra_nodes()) return loop = asyncio.new_event_loop() diff --git a/tests/runtime/generated/upscale-model-loader.py b/tests/runtime/generated/upscale-model-loader.py index 98ed0ae..6e04ee5 100644 --- a/tests/runtime/generated/upscale-model-loader.py +++ b/tests/runtime/generated/upscale-model-loader.py @@ -1,10 +1,70 @@ # Imports +import importlib.util import json +import logging import os import random import sys from typing import Sequence, Mapping, Any, Union +log = logging.getLogger(__name__) + + +def _load_module(module_name: str, filepath: str) -> Any: + """Load a Python module from an explicit file path, bypassing sys.path. + + This eliminates ALL bare import shadowing attacks by loading modules + from verified file paths instead of relying on sys.path resolution. + """ + if not os.path.isfile(filepath): + log.debug("Module file not found: %s (%s)", module_name, filepath) + return None + try: + spec = importlib.util.spec_from_file_location(module_name, filepath) + if spec is None or spec.loader is None: + log.debug("Could not create spec for %s at %s", module_name, filepath) + return None + mod = importlib.util.module_from_spec(spec) + sys.modules[module_name] = mod + spec.loader.exec_module(mod) + return mod + except Exception as e: + log.debug("Failed to load %s from %s: %s", module_name, filepath, e) + return None + + +def _load_module_temp(module_name: str, filepath: str) -> Any: + """Load a module via _load_module() then remove it from sys.modules. + + Used during bootstrap for modules that ComfyUI's import chain also loads + normally — prevents the cached copy from conflicting with later imports. + """ + mod = _load_module(module_name, filepath) + sys.modules.pop(module_name, None) + return mod + + +def _is_comfyui_directory(path: str) -> bool: + """Verify a directory has ComfyUI structural markers (nodes.py).""" + if not os.path.isdir(path): + return False + return os.path.isfile(os.path.join(path, "nodes.py")) + + +def _find_from_extension_location() -> str | None: + """Walk up from this file's location to find ComfyUI root.""" + ext_dir = os.path.dirname(os.path.realpath(__file__)) + candidate = ext_dir + for _ in range(10): + parent = os.path.dirname(candidate) + if parent == candidate: + break + candidate = parent + if os.path.basename(candidate) == "ComfyUI": + if _is_comfyui_directory(candidate): + return candidate + return None + def get_value_at_index(obj: Union[Sequence, Mapping], index: int) -> Any: """Return a sequence or mapping result item by index.""" @@ -14,67 +74,104 @@ def get_value_at_index(obj: Union[Sequence, Mapping], index: int) -> Any: return obj["result"][index] -def get_comfyui_path() -> str: - """Return the configured ComfyUI path, preferring COMFYUI_PATH when set.""" - comfyui_path = os.environ.get("COMFYUI_PATH") - if comfyui_path: - return comfyui_path - return find_path("ComfyUI") +def get_comfyui_path() -> str | None: + """Resolve ComfyUI path via prioritized multi-strategy fallback. + + Strategy order: + 1. COMFYUI_PATH env var (verified with _is_comfyui_directory) + 2. Relative walk from extension location (realpath + verified) + 3. CWD walk (legacy fallback, depth-limited) + """ + p = os.environ.get("COMFYUI_PATH") + if p and _is_comfyui_directory(p): + return p + p = _find_from_extension_location() + if p: + return p + return find_path("ComfyUI", max_depth=20) -def find_path(name: str, path: str = None) -> str: - """Recursively search parent folders until the named entry is found.""" - if path is None: - path = os.getcwd() +def find_path(name: str, max_depth: int = 20) -> str | None: + """Iteratively walk up from CWD to find a directory by name. - if name in os.listdir(path): - path_name = os.path.join(path, name) - print(f"{name} found: {path_name}") - return path_name - - parent_directory = os.path.dirname(path) - if parent_directory == path: - return None - - return find_path(name, parent_directory) + Depth-limited to prevent slow startup on deep trees. + Each candidate verified by _is_comfyui_directory() at the caller. + """ + candidate = os.getcwd() + for _ in range(max_depth): + parent = os.path.dirname(candidate) + if parent == candidate: + break + candidate = parent + if os.path.basename(candidate) == name: + return candidate + return None def add_comfyui_directory_to_sys_path() -> None: - """Add the ComfyUI checkout to sys.path.""" + """Add the ComfyUI checkout to sys.path (idempotent — insert-once, no gap).""" comfyui_path = get_comfyui_path() if comfyui_path is not None and os.path.isdir(comfyui_path): - if comfyui_path in sys.path: - sys.path.remove(comfyui_path) - sys.path.insert(0, comfyui_path) - print(f"'{comfyui_path}' added to sys.path") + if comfyui_path not in sys.path: + sys.path.insert(0, comfyui_path) + log.debug("Added %s to sys.path", comfyui_path) def add_extra_model_paths() -> None: """Load ComfyUI extra model paths configuration when available.""" - try: - from main import load_extra_path_config - except ImportError: - print( - "Could not import load_extra_path_config from main.py. Looking in utils.extra_config instead." + comfyui_path = get_comfyui_path() + if comfyui_path is None: + log.debug("Cannot load extra model paths: ComfyUI path not found") + return + + # Try main.py first, then utils/extra_config.py — both via _load_module() + main_mod = _load_module("comfy_main", os.path.join(comfyui_path, "main.py")) + if main_mod is not None and hasattr(main_mod, "load_extra_path_config"): + load_extra_path_config = getattr(main_mod, "load_extra_path_config") + else: + log.debug("main.py not available, trying utils/extra_config.py") + extra_config_mod = _load_module( + "extra_config", os.path.join(comfyui_path, "utils", "extra_config.py") ) - from utils.extra_config import load_extra_path_config + if extra_config_mod is None or not hasattr( + extra_config_mod, "load_extra_path_config" + ): + log.debug("Could not find load_extra_path_config in either path") + return + load_extra_path_config = getattr(extra_config_mod, "load_extra_path_config") extra_model_paths = find_path("extra_model_paths.yaml") if extra_model_paths is not None: load_extra_path_config(extra_model_paths) else: - print("Could not find the extra_model_paths config file.") + log.debug("Could not find the extra_model_paths config file.") def bootstrap_comfyui_runtime() -> None: """Mirror the allocator-related ComfyUI startup steps before torch import.""" add_comfyui_directory_to_sys_path() + comfyui_path = get_comfyui_path() + if comfyui_path is None: + log.debug("bootstrap_comfyui_runtime: ComfyUI path not found") + return - import comfy.options + # Use _load_module for file-based isolation, but with temporary names so + # the modules are removed from sys.modules after reading values. + # This prevents conflicts when ComfyUI's internal import chain (e.g. + # nodes.py -> comfy.cli_args) later loads these modules normally. + options_mod = _load_module_temp( + "_bootstrap_options", os.path.join(comfyui_path, "comfy", "options.py") + ) + if options_mod is not None: + options_mod.enable_args_parsing() - comfy.options.enable_args_parsing() + cli_args_mod = _load_module_temp( + "_bootstrap_cli_args", os.path.join(comfyui_path, "comfy", "cli_args.py") + ) + args = getattr(cli_args_mod, "args", None) if cli_args_mod else None - from comfy.cli_args import args + if args is None: + return if os.name == "nt": os.environ["MIMALLOC_PURGE_DELAY"] = "0" @@ -99,9 +196,14 @@ def bootstrap_comfyui_runtime() -> None: if args.deterministic and "CUBLAS_WORKSPACE_CONFIG" not in os.environ: os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8" - import cuda_malloc - - if "rocm" in cuda_malloc.get_torch_version_noimport(): + cuda_malloc_mod = _load_module_temp( + "_bootstrap_cuda_malloc", os.path.join(comfyui_path, "cuda_malloc.py") + ) + if ( + cuda_malloc_mod is not None + and hasattr(cuda_malloc_mod, "get_torch_version_noimport") + and "rocm" in cuda_malloc_mod.get_torch_version_noimport() + ): os.environ["OCL_SET_SVM_SIZE"] = "262144" @@ -109,6 +211,19 @@ def cleanup_comfyui_runtime(unload_models: bool | None = None) -> None: """Best-effort cleanup for embedded or repeated generated-script execution.""" import gc + def run_cleanup_hook(name: str, should_run: bool = True) -> None: + if not should_run or not hasattr(model_management, name): + return + cleanup_fn = getattr(model_management, name) + try: + cleanup_fn() + except Exception as exc: + warnings.warn( + f"ComfyUI cleanup hook {name} failed during teardown: {exc}", + RuntimeWarning, + stacklevel=2, + ) + should_unload = unload_models if should_unload is None: should_unload = os.environ.get( @@ -126,36 +241,52 @@ def cleanup_comfyui_runtime(unload_models: bool | None = None) -> None: gc.collect() return - if hasattr(model_management, "cleanup_models_gc"): - model_management.cleanup_models_gc() - if should_unload and hasattr(model_management, "unload_all_models"): - model_management.unload_all_models() - if hasattr(model_management, "soft_empty_cache"): - model_management.soft_empty_cache() + run_cleanup_hook("cleanup_models_gc") + run_cleanup_hook("unload_all_models", should_run=should_unload) + run_cleanup_hook("soft_empty_cache") gc.collect() def import_custom_nodes() -> None: - """Initialize ComfyUI custom nodes in the exporter runtime.""" + """Initialize ComfyUI custom nodes in the exporter runtime. + + Uses _load_module() for all ComfyUI imports — no bare imports, + no sys.path remove/re-insert gap. + """ comfyui_path = get_comfyui_path() - if comfyui_path and comfyui_path not in sys.path: + if comfyui_path is None: + log.debug("import_custom_nodes: ComfyUI path not found") + return + + # Idempotent insert-once — never removes from sys.path + if comfyui_path not in sys.path: sys.path.insert(0, comfyui_path) import asyncio - import execution - from nodes import init_extra_nodes - if comfyui_path in sys.path: - sys.path.remove(comfyui_path) - sys.path.insert(0, comfyui_path) + execution_mod = _load_module( + "execution", os.path.join(comfyui_path, "execution.py") + ) + nodes_mod = _load_module("nodes", os.path.join(comfyui_path, "nodes.py")) + server_mod = _load_module("server", os.path.join(comfyui_path, "server.py")) - import server + if execution_mod is None or server_mod is None: + log.debug( + "import_custom_nodes: could not load execution/server modules. " + "Proceeding without full PromptServer/PromptQueue setup." + ) + # Even without server, we can still populate NODE_CLASS_MAPPINGS from + # comfy_extras and built-in extra nodes. + if nodes_mod is not None and hasattr(nodes_mod, "init_extra_nodes"): + asyncio.run(nodes_mod.init_extra_nodes()) + return loop = asyncio.new_event_loop() asyncio.set_event_loop(loop) - server_instance = server.PromptServer(loop) - execution.PromptQueue(server_instance) - asyncio.run(init_extra_nodes()) + server_instance = server_mod.PromptServer(loop) + execution_mod.PromptQueue(server_instance) + if nodes_mod is not None and hasattr(nodes_mod, "init_extra_nodes"): + asyncio.run(nodes_mod.init_extra_nodes()) # Workflow data