nodes/pydn-ComfyUI-to-Python-Extension
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix: use _load_module_temp for bootstrap modules to avoid sys.modules conflicts

Browse Source 
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'
This commit is contained in:
Peyton DeNiro
2026-05-10 00:31:02 -05:00
parent 49649a6745
commit 3ec60ded7d
5 changed files with 177 additions and 55 deletions
Download Patch File Download Diff File Expand all files Collapse all files
comfyui_to_python/generator/generated_helpers.py
+2
View File
@@ -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",
comfyui_to_python/generator/render.py
+2
View File
@@ -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,
comfyui_to_python/node_runtime.py
+28 -14
View File
@@ -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
tests/runtime/generated/text-to-image.py
+140 -38
View File
@@ -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"
tests/test_upscale_model_loader_export.py
+5 -3
View File
@@ -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"