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'
This commit is contained in:
Peyton DeNiro
@@ -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"
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user