From 3ec60ded7dc014d30005f36dd5cd373ce15d7d74 Mon Sep 17 00:00:00 2001 From: Peyton DeNiro <25550995+pydn@users.noreply.github.com> Date: Sun, 10 May 2026 00:31:02 -0500 Subject: [PATCH] fix: use _load_module_temp for bootstrap modules to avoid sys.modules conflicts bootstrap_comfyui_runtime() now uses _load_module_temp() with temporary module names (_bootstrap_options, _bootstrap_cli_args, _bootstrap_cuda_malloc) instead of real module names. This prevents cached stale modules from sys.modules conflicting with ComfyUI's internal import chain (e.g. nodes.py -> comfy.sd -> comfy.model_management -> comfy.cli_args). Fixes the text-to-image runtime e2e failure: ImportError: cannot import name 'args' from 'comfy.cli_args' --- .../generator/generated_helpers.py | 2 + comfyui_to_python/generator/render.py | 2 + comfyui_to_python/node_runtime.py | 42 +++-- tests/runtime/generated/text-to-image.py | 178 ++++++++++++++---- tests/test_upscale_model_loader_export.py | 8 +- 5 files changed, 177 insertions(+), 55 deletions(-) diff --git a/comfyui_to_python/generator/generated_helpers.py b/comfyui_to_python/generator/generated_helpers.py index 1a070ad..c0b6e9b 100644 --- a/comfyui_to_python/generator/generated_helpers.py +++ b/comfyui_to_python/generator/generated_helpers.py @@ -2,6 +2,7 @@ from ..node_runtime import ( _find_from_extension_location, _is_comfyui_directory, _load_module, + _load_module_temp, add_comfyui_directory_to_sys_path, add_extra_model_paths, bootstrap_comfyui_runtime, @@ -15,6 +16,7 @@ __all__ = [ "_find_from_extension_location", "_is_comfyui_directory", "_load_module", + "_load_module_temp", "add_comfyui_directory_to_sys_path", "add_extra_model_paths", "bootstrap_comfyui_runtime", diff --git a/comfyui_to_python/generator/render.py b/comfyui_to_python/generator/render.py index 963617b..707bed7 100644 --- a/comfyui_to_python/generator/render.py +++ b/comfyui_to_python/generator/render.py @@ -9,6 +9,7 @@ from .generated_helpers import ( _find_from_extension_location, _is_comfyui_directory, _load_module, + _load_module_temp, add_comfyui_directory_to_sys_path, add_extra_model_paths, bootstrap_comfyui_runtime, @@ -35,6 +36,7 @@ class WorkflowRenderer: func_strings = [] for func in [ _load_module, + _load_module_temp, _is_comfyui_directory, _find_from_extension_location, get_value_at_index, diff --git a/comfyui_to_python/node_runtime.py b/comfyui_to_python/node_runtime.py index 798a7b2..61a7f72 100644 --- a/comfyui_to_python/node_runtime.py +++ b/comfyui_to_python/node_runtime.py @@ -126,6 +126,17 @@ def add_extra_model_paths() -> None: log.debug("Could not find the extra_model_paths config file.") +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 bootstrap_comfyui_runtime() -> None: """Mirror the allocator-related ComfyUI startup steps before torch import.""" add_comfyui_directory_to_sys_path() @@ -134,21 +145,28 @@ def bootstrap_comfyui_runtime() -> None: log.debug("bootstrap_comfyui_runtime: ComfyUI path not found") return - options_mod = _load_module( - "comfy.options", os.path.join(comfyui_path, "comfy", "options.py") + # 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() - cli_args_mod = _load_module( - "comfy.cli_args", os.path.join(comfyui_path, "comfy", "cli_args.py") + 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 + if args is None: + return + if os.name == "nt": os.environ["MIMALLOC_PURGE_DELAY"] = "0" - if args is not None and args.default_device is not None: + if args.default_device is not None: default_dev = args.default_device devices = list(range(32)) devices.remove(default_dev) @@ -157,23 +175,19 @@ def bootstrap_comfyui_runtime() -> None: os.environ["CUDA_VISIBLE_DEVICES"] = str(devices) os.environ["HIP_VISIBLE_DEVICES"] = str(devices) - if args is not None and args.cuda_device is not None: + if args.cuda_device is not None: os.environ["CUDA_VISIBLE_DEVICES"] = str(args.cuda_device) os.environ["HIP_VISIBLE_DEVICES"] = str(args.cuda_device) os.environ["ASCEND_RT_VISIBLE_DEVICES"] = str(args.cuda_device) - if args is not None and args.oneapi_device_selector is not None: + if args.oneapi_device_selector is not None: os.environ["ONEAPI_DEVICE_SELECTOR"] = args.oneapi_device_selector - if ( - args is not None - and args.deterministic - and "CUBLAS_WORKSPACE_CONFIG" not in os.environ - ): + if args.deterministic and "CUBLAS_WORKSPACE_CONFIG" not in os.environ: os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8" - cuda_malloc_mod = _load_module( - "cuda_malloc", os.path.join(comfyui_path, "cuda_malloc.py") + cuda_malloc_mod = _load_module_temp( + "_bootstrap_cuda_malloc", os.path.join(comfyui_path, "cuda_malloc.py") ) if ( cuda_malloc_mod is not None diff --git a/tests/runtime/generated/text-to-image.py b/tests/runtime/generated/text-to-image.py index 99c66dc..cb109f6 100644 --- a/tests/runtime/generated/text-to-image.py +++ b/tests/runtime/generated/text-to-image.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" diff --git a/tests/test_upscale_model_loader_export.py b/tests/test_upscale_model_loader_export.py index 9c6e4f8..64be1e6 100644 --- a/tests/test_upscale_model_loader_export.py +++ b/tests/test_upscale_model_loader_export.py @@ -146,10 +146,12 @@ class UpscaleModelLoaderExportTest(unittest.TestCase): self.assertIn("def bootstrap_comfyui_runtime()", generated) self.assertIn("def cleanup_comfyui_runtime(", generated) - # Hardened imports use _load_module(), not bare 'import comfy.options' - self.assertIn('"comfy.options"', generated) + # Hardened imports use _load_module_temp() with temporary names, + # not bare 'import comfy.options' — avoids sys.modules pollution + self.assertIn("_load_module_temp", generated) + self.assertIn("options.py", generated) self.assertIn("enable_args_parsing()", generated) - self.assertIn('"cuda_malloc"', generated) + self.assertIn("cuda_malloc.py", generated) self.assertNotIn("\nbootstrap_comfyui_runtime()\n", generated) self.assertIn( "def main(unload_models: bool | None = None):\n"