fix: embed _find_file into generated scripts

add_extra_model_paths() uses _find_file() which must be available in
generated standalone scripts.  Add it to generated_helpers __all__ so the
render step embeds it alongside the other runtime helpers.

.gitignore

@@ -145,3 +145,8 @@ cython_debug/
 *.sql
 *.sqlite
 *.xml
+
+docs/
+.agents/
+.codex/
+AGENTS.md

README.md

@@ -13,22 +13,20 @@ This project supports:
 ## Install
 
 Choose the setup that matches how you want to use the project.
+This project supports Python 3.12 and newer.
 
 ### Web UI extension (`File -> Save As Script`)
 
 For ComfyUI to recognize this project as an extension, the repo must be discoverable through ComfyUI's `custom_nodes` search paths.
 
-Use one of these setups:
+Clone directly into `ComfyUI/custom_nodes`:
 
-1. Clone directly into `ComfyUI/custom_nodes`
 ```bash
 cd /path/to/ComfyUI/custom_nodes
 git clone https://github.com/pydn/ComfyUI-to-Python-Extension.git
-cd ComfyUI-to-Python-Extension
-uv sync
 ```
 
-2. Keep the repo elsewhere, then either:
+Or keep the repo elsewhere, then either:
 - symlink it into `ComfyUI/custom_nodes`
 - add its parent directory to ComfyUI's `custom_nodes` search paths via `extra_model_paths.yaml`
 
@@ -37,10 +35,29 @@ Example symlink setup:
 git clone https://github.com/pydn/ComfyUI-to-Python-Extension.git
 cd /path/to/ComfyUI/custom_nodes
 ln -s /path/to/ComfyUI-to-Python-Extension ComfyUI-to-Python-Extension
-cd /path/to/ComfyUI-to-Python-Extension
-uv sync
 ```
 
+Then install this extension into the same Python environment that launches ComfyUI.
+The `pyproject.toml` file declares the package dependencies, but those dependencies still need to be installed into ComfyUI's runtime Python.
+
+If you run ComfyUI from a source checkout with `uv`:
+
+```bash
+cd /path/to/ComfyUI
+uv pip install -e ./custom_nodes/ComfyUI-to-Python-Extension
+uv run python main.py
+```
+
+If you use the Windows portable build:
+
+```
+cd C:\path\to\ComfyUI_windows_portable\ComfyUI\custom_nodes\ComfyUI-to-Python-Extension
+..\..\..\python_embeded\python.exe -m pip install -e .
+```
+
+Running `uv sync` inside `ComfyUI-to-Python-Extension` creates this extension's own `.venv`.
+ComfyUI does not automatically import dependencies from that `.venv`; it imports custom nodes with the Python interpreter used to launch ComfyUI.
+
 After installation, restart ComfyUI.
 
 ### CLI exporter / generated scripts
@@ -55,9 +72,18 @@ export COMFYUI_PATH=/path/to/ComfyUI
 ```
 
 `COMFYUI_PATH` helps the exporter and generated scripts find the ComfyUI codebase. It does not, by itself, register this repo as a ComfyUI extension for the Web UI.
+It also does not install ComfyUI runtime dependencies such as `torch` into the current Python environment.
 
 `COMFYUI_PATH` is checked first. If it is not set, the exporter falls back to searching parent directories for a folder named `ComfyUI`.
 
+If the CLI fails with `ModuleNotFoundError: No module named 'torch'`, run the command with the same Python environment that launches ComfyUI, or install ComfyUI's runtime dependencies into the environment you are using for the CLI.
+
+For Windows portable builds, run the CLI with ComfyUI's embedded Python from the extension directory:
+
+```
+..\..\..\python_embeded\python.exe -m comfyui_to_python --input_file ".\workflow_api.json" --output_file ".\workflow_api.py"
+```
+
 ## Web UI Export
 
 In current ComfyUI builds, `Save As Script` is typically available under:
@@ -65,12 +91,13 @@ In current ComfyUI builds, `Save As Script` is typically available under:
 `File -> Save As Script`
 
 The command downloads a generated `.py` file.
+The current UI export uses the default filename `workflow_api.py` so it works in ComfyUI Desktop without relying on `prompt()`.
 
 ![Save As Script](images/save_as_script.png)
 
 Notes:
 - menu placement can differ between frontend versions
-- ComfyUI Desktop may fail on the current filename prompt flow; use the CLI flow below if that happens
+- the Web UI export uses a fixed default filename rather than asking for one interactively
 
 ## CLI Export
 
@@ -118,14 +145,30 @@ The generated script is a workflow export. It does not automatically turn workfl
 
 Scripts exported directly from `File -> Save As Script` in the ComfyUI UI already include the frontend workflow metadata needed for drag-and-drop reimport. Images saved by those scripts can be dropped back into ComfyUI and reopen with the original workflow metadata.
 
+Generated scripts reuse ComfyUI's runtime argument parser during bootstrap, so common ComfyUI memory flags such as `--highvram`, `--normalvram`, `--lowvram`, `--novram`, `--cpu`, and `--disable-smart-memory` can be passed directly to the exported `.py` file.
+
+Lifecycle notes:
+- exported scripts are single-shot workflow runners, not long-lived ComfyUI prompt servers
+- they do not implement Web UI prompt/result caching across repeated service calls
+- exported `main()` now performs best-effort ComfyUI model/cache cleanup in a `finally` block
+- set `COMFYUI_TOPYTHON_UNLOAD_MODELS=1` or call `main(unload_models=True)` if an embedded or repeated-call host should aggressively unload models after each run instead of preserving them for reuse
+
 ## Troubleshooting
 
+- unsupported Python version:
+  use Python 3.12 or newer in the environment that runs the extension, then reinstall the extension dependencies there
+- Web UI import fails after `uv sync`:
+  `uv sync` in this repo installs dependencies into this repo's `.venv`, but ComfyUI loads custom nodes with its own Python environment. Install the extension into the Python interpreter that launches ComfyUI.
+- Windows portable import fails after `uv sync`:
+  ComfyUI portable uses its bundled `python_embeded` interpreter. From the extension directory, run `..\..\..\python_embeded\python.exe -m pip install -e .`, then restart ComfyUI.
+- CLI fails with `No module named 'torch'`:
+  the extension `.venv` may not have ComfyUI's runtime dependencies. Either run the CLI from the Python environment that launches ComfyUI, or make sure the target ComfyUI environment is installed and `COMFYUI_PATH` points to it.
 - `Save As Script` not visible:
   check your current ComfyUI menu/frontend version and look under `File`
 - `Save As Script` not visible after restart:
   make sure this repo is discoverable by ComfyUI through `custom_nodes` by cloning it into `ComfyUI/custom_nodes`, symlinking it there, or adding an external `custom_nodes` path in `extra_model_paths.yaml`
-- Desktop says `prompt()` is unsupported:
-  use the CLI export flow instead
+- save uses the default filename:
+  rename `workflow_api.py` after download if you want a different local filename
 - ComfyUI cannot be found:
   set `COMFYUI_PATH`
 - models or paths are missing at runtime:

__init__.py

@@ -13,18 +13,10 @@ sys.path.append(ext_dir)
 try:
     import black
 except ImportError:
-    print("Unable to import requirements for ComfyUI-SaveAsScript.")
-    print("Installing...")
-
-    import importlib
-
-    spec = importlib.util.spec_from_file_location(
-        "impact_install", os.path.join(os.path.dirname(__file__), "install.py")
-    )
-    impact_install = importlib.util.module_from_spec(spec)
-    spec.loader.exec_module(impact_install)
-
-    print("Successfully installed. Hopefully, at least.")
+    raise ImportError(
+        "ComfyUI-to-Python-Extension requires the project dependencies to be installed. "
+        f"Run 'uv sync' in {ext_dir} with Python 3.12+ before loading this extension."
+    ) from None
 
 # Prevent reimporting of custom nodes
 os.environ["RUNNING_IN_COMFYUI"] = "TRUE"

comfyui_to_python/app.py

@@ -1,4 +1,5 @@
 import copy
+import sys
 
 from typing import TextIO
 
@@ -57,7 +58,16 @@ class ExportApplication:
         }
         if self.needs_init_custom_nodes or missing_node_types:
             self.custom_node_importer()
-            self.base_node_class_mappings = copy.deepcopy(self.node_class_mappings)
+            # 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
+            # Leave self.base_node_class_mappings unchanged — it represents the
+            # pre-custom-node baseline used by WorkflowPlanner to decide whether
+            # a node gets a direct import or a NODE_CLASS_MAPPINGS dict lookup.
 
         load_order = LoadOrderDeterminer(
             data, self.node_class_mappings

comfyui_to_python/generator/generated_helpers.py

@@ -1,16 +1,28 @@
 from ..node_runtime import (
+    _find_file,
+    _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,
+    cleanup_comfyui_runtime,
     find_path,
     get_comfyui_path,
     get_value_at_index,
 )
 
 __all__ = [
+    "_find_file",
+    "_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",
+    "cleanup_comfyui_runtime",
     "find_path",
     "get_comfyui_path",
     "get_value_at_index",

comfyui_to_python/generator/render.py

@@ -1,18 +1,14 @@
 import inspect
+import logging
 from pprint import pformat
 from typing import Any
 
 import black
 
+log = logging.getLogger(__name__)
+
 from ..node_runtime import import_custom_nodes
-from .generated_helpers import (
-    add_comfyui_directory_to_sys_path,
-    add_extra_model_paths,
-    bootstrap_comfyui_runtime,
-    find_path,
-    get_comfyui_path,
-    get_value_at_index,
-)
+from . import generated_helpers
 from .model import GenerationPlan
 
 
@@ -28,24 +24,29 @@ class WorkflowRenderer:
                 {"workflow": plan.metadata_workflow_data}
             )
 
+        # Auto-discover helpers from generated_helpers.__all__ so new helpers
+        # are picked up without manually updating this list.
         func_strings = []
-        for func in [
-            get_value_at_index,
-            get_comfyui_path,
-            find_path,
-            add_comfyui_directory_to_sys_path,
-            add_extra_model_paths,
-            bootstrap_comfyui_runtime,
-        ]:
+        for name in generated_helpers.__all__:
+            func = getattr(generated_helpers, name, None)
+            if func is None:
+                log.warning(
+                    "Helper '%s' missing from generated_helpers — skipping", name
+                )
+                continue
             func_strings.append(f"\n{inspect.getsource(func)}")
 
         static_imports = [
             "# 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__)",
         ] + func_strings
 
         if plan.custom_nodes:
@@ -74,7 +75,7 @@ class WorkflowRenderer:
 
         execution_section = [
             "# Workflow execution",
-            "def main():",
+            "def main(unload_models: bool | None = None):",
             "    bootstrap_comfyui_runtime()",
             "    add_extra_model_paths()",
         ]
@@ -88,20 +89,27 @@ class WorkflowRenderer:
                 "",
                 "    import torch",
                 "",
-                "    with torch.inference_mode():",
+                "    try:",
+                "        with torch.inference_mode():",
             ]
         )
         execution_section.extend(
             self.build_function_body(
-                plan.special_functions_code, "pass", indentation="        "
+                plan.special_functions_code, "pass", indentation="            "
             ).splitlines()
         )
-        execution_section.append(f"        for q in range({plan.queue_size}):")
+        execution_section.append(f"            for q in range({plan.queue_size}):")
         execution_section.extend(
             self.build_function_body(
-                plan.loop_code, "pass", indentation="            "
+                plan.loop_code, "pass", indentation="                "
             ).splitlines()
         )
+        execution_section.extend(
+            [
+                "    finally:",
+                "        cleanup_comfyui_runtime(unload_models=unload_models)",
+            ]
+        )
 
         entrypoint_section = [
             "# Entrypoint",

comfyui_to_python/node_runtime.py

@@ -1,69 +1,226 @@
+import importlib.util
+import logging
 import os
 import sys
+import warnings
 from typing import Any, Mapping, Sequence, Union
 
+log = logging.getLogger(__name__)
 
-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()
 
-    if name in os.listdir(path):
-        path_name = os.path.join(path, name)
-        print(f"{name} found: {path_name}")
-        return path_name
+def _is_comfyui_directory(path: str) -> bool:
+    """Verify a directory has ComfyUI structural markers.
 
-    parent_directory = os.path.dirname(path)
-    if parent_directory == path:
+    Checks for nodes.py, main.py, and the comfy/ subdirectory to raise
+    the bar against spoofing via a directory with only a single marker file.
+    """
+    if not os.path.isdir(path):
+        return False
+    return (
+        os.path.isfile(os.path.join(path, "nodes.py"))
+        and os.path.isfile(os.path.join(path, "main.py"))
+        and os.path.isdir(os.path.join(path, "comfy"))
+    )
+
+
+def _load_module(module_name: str, filepath: str) -> Any:
+    """Load a Python module from an explicit file path, bypassing sys.path.
+
+    Significantly reduces bare import shadowing risk by loading modules
+    from verified file paths instead of relying on sys.path resolution.
+    If exec_module() raises, the partially-loaded module is removed from
+    sys.modules so subsequent calls start fresh.
+    """
+    # Return cached module if already loaded — prevents re-execution
+    # that would reset state (e.g. NODE_CLASS_MAPPINGS after init_extra_nodes).
+    if module_name in sys.modules:
+        return sys.modules[module_name]
+    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
+        try:
+            spec.loader.exec_module(mod)
+        except BaseException:
+            sys.modules.pop(module_name, None)
+            raise
+        return mod
+    except BaseException as e:
+        log.debug("Failed to load %s from %s: %s", module_name, filepath, e)
+        sys.modules.pop(module_name, None)  # Also clean up on edge-case exceptions
         return None
 
-    return find_path(name, parent_directory)
+
+def _find_from_extension_location() -> str | None:
+    """Walk up from this file's location to find ComfyUI root.
+
+    Checks the starting directory first before walking upward.
+    """
+    ext_dir = os.path.dirname(os.path.realpath(__file__))
+    candidate = ext_dir
+    for _ in range(10):
+        if os.path.basename(candidate) == "ComfyUI":
+            if _is_comfyui_directory(candidate):
+                return candidate
+        parent = os.path.dirname(candidate)
+        if parent == candidate:
+            break
+        candidate = parent
+    return None
 
 
-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 _find_file(name: str, max_depth: int = 20) -> str | None:
+    """Walk up from CWD to find a file by name.
+
+    Unlike find_path() which searches for directories, this checks
+    os.path.isfile() at each level. Returns full path to the file or None.
+
+    Checks CWD first before walking upward.
+    """
+    candidate = os.getcwd()
+    for _ in range(max_depth):
+        filepath = os.path.join(candidate, name)
+        if os.path.isfile(filepath):
+            return filepath
+        parent = os.path.dirname(candidate)
+        if parent == candidate:
+            break
+        candidate = parent
+    return None
+
+
+def find_path(name: str, max_depth: int = 20) -> str | None:
+    """Iteratively walk up from CWD to find a directory by name.
+
+    Checks CWD first before walking upward. 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):
+        if os.path.basename(candidate) == name:
+            return candidate
+        parent = os.path.dirname(candidate)
+        if parent == candidate:
+            break
+        candidate = parent
+    return None
+
+
+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
+    p = find_path("ComfyUI", max_depth=20)
+    if p and _is_comfyui_directory(p):
+        return p
+    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 — always at index 0).
+
+    If already present but lower in sys.path, removes and re-inserts at front
+    so bare imports always resolve to this copy first.
+    """
     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")
+        log.debug("Added %s to sys.path[0]", 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."
-        )
-        from utils.extra_config import load_extra_path_config
+    comfyui_path = get_comfyui_path()
+    if comfyui_path is None:
+        log.debug("Cannot load extra model paths: ComfyUI path not found")
+        return
 
-    extra_model_paths = find_path("extra_model_paths.yaml")
+    # 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")
+        )
+        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_file("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 _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()
+    comfyui_path = get_comfyui_path()
+    if comfyui_path is None:
+        log.debug("bootstrap_comfyui_runtime: ComfyUI path not found")
+        return
 
-    import comfy.options
+    # Load comfy.options under its canonical name so that enable_args_parsing()
+    # mutates the sys.modules["comfy.options"] cache entry. When cli_args.py
+    # later imports comfy.options by canonical name, it sees the mutation.
+    options_mod = _load_module(
+        "comfy.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 imports comfy.options — so it must be cached under canonical name
+    # at this point. We remove both from sys.modules only after cli_args has
+    # loaded and consumed the args value, to prevent conflicts with later normal
+    # import chains (e.g. nodes.py -> comfy.cli_args).
+    cli_args_mod = _load_module(
+        "comfy.cli_args", os.path.join(comfyui_path, "comfy", "cli_args.py")
+    )
 
-    from comfy.cli_args import args
+    args = getattr(cli_args_mod, "args", None) if cli_args_mod else None
+
+    # Now safe to remove bootstrap copies — args value already captured.
+    sys.modules.pop("comfy.options", None)
+    sys.modules.pop("comfy.cli_args", None)
+
+    if args is None:
+        return
 
     if os.name == "nt":
         os.environ["MIMALLOC_PURGE_DELAY"] = "0"
@@ -88,41 +245,132 @@ 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"
 
 
+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(
+            "COMFYUI_TOPYTHON_UNLOAD_MODELS", ""
+        ).lower() in {
+            "1",
+            "true",
+            "yes",
+            "on",
+        }
+
+    try:
+        import comfy.model_management as model_management
+    except ModuleNotFoundError:
+        gc.collect()
+        return
+
+    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"))
 
-    import server
+    # nodes.py inserts comfy/ subdirectory into sys.path, which shadows the
+    # top-level utils/ package (comfy/utils.py vs utils/). This breaks
+    # server.py → app.frontend_management → from utils.install_util import ...
+    # Filter it out temporarily so server.py loads cleanly.
+    comfy_subdir = os.path.join(comfyui_path, "comfy")
+    original_sys_path = list(sys.path)
+    sys.path[:] = [p for p in sys.path if p != comfy_subdir]
+
+    try:
+        server_mod = _load_module("server", os.path.join(comfyui_path, "server.py"))
+    finally:
+        # Restore sys.path so nodes and other modules that need comfy/ still work.
+        # Guaranteed even if _load_module raises mid-execution.
+        sys.path[:] = original_sys_path
+
+    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())
 
 
 def get_node_class_mappings() -> dict:
-    """Load ComfyUI node mappings on demand."""
-    add_comfyui_directory_to_sys_path()
-    from nodes import NODE_CLASS_MAPPINGS
+    """Load ComfyUI node mappings on demand via _load_module().
 
-    return NODE_CLASS_MAPPINGS
+    Reuses the cached "nodes" module from sys.modules if already loaded
+    (e.g. by import_custom_nodes) to avoid resetting NODE_CLASS_MAPPINGS.
+    """
+    add_comfyui_directory_to_sys_path()
+    comfyui_path = get_comfyui_path()
+    if comfyui_path is None:
+        log.debug("get_node_class_mappings: ComfyUI path not found")
+        return {}
+
+    # Reuse cached module if already loaded to avoid resetting mappings.
+    nodes_mod = sys.modules.get("nodes")
+    if nodes_mod is None:
+        nodes_mod = _load_module("nodes", os.path.join(comfyui_path, "nodes.py"))
+    if nodes_mod is None:
+        return {}
+    return getattr(nodes_mod, "NODE_CLASS_MAPPINGS", {})
 
 
 def get_value_at_index(obj: Union[Sequence, Mapping], index: int) -> Any:

js/save-as-script.js

@@ -1,6 +1,9 @@
 import { api } from "../../scripts/api.js";
 import { app } from "../../scripts/app.js";
 
+const DEFAULT_SCRIPT_FILENAME = "workflow_api.py";
+const DEFAULT_WORKFLOW_NAME = "workflow_api.json";
+
 function $el(tag, options = {}) {
 	const element = document.createElement(tag);
 	const { parent, style, ...props } = options;
@@ -39,15 +42,12 @@ const extension = {
 		});
 	},
 	savePythonScript() {
-		var filename = prompt("Save script as:");
-		if(filename === undefined || filename === null || filename === "") {
-			return
-		}
-		
+		const filename = DEFAULT_SCRIPT_FILENAME;
+
 		app.graphToPrompt().then(async (p) => {
 			const frontendWorkflow = p.workflow ?? app.graph.serialize();
 			const json = JSON.stringify({
-				name: filename + ".json",
+				name: DEFAULT_WORKFLOW_NAME,
 				workflow: JSON.stringify(p.output, null, 2),
 				frontend_workflow: JSON.stringify(frontendWorkflow, null, 2),
 			}, null, 2); // convert the data to a JSON string
@@ -55,10 +55,6 @@ const extension = {
 			if(response.status == 200) {
 				const blob = new Blob([await response.text()], {type: "text/python;charset=utf-8"});
 				const url = URL.createObjectURL(blob);
-				if(!filename.endsWith(".py")) {
-					filename += ".py";
-				}
-
 				const a = $el("a", {
 					href: url,
 					download: filename,

pyproject.toml

@@ -1,10 +1,17 @@
 [project]
 name = "comfyui-to-python-extension"
 description = "This custom node allows you to generate pure python code from your ComfyUI workflow with the click of a button. Great for rapid experimentation or production deployment."
-version = "2.0.0"
-license = { text = "MIT License" }
+version = "2.1.0"
+license = "MIT"
+requires-python = ">=3.12"
 dependencies = ["black"]
 
+[tool.setuptools.packages.find]
+include = ["comfyui_to_python*"]
+
+[project.optional-dependencies]
+dev = ["ruff"]
+
 [project.urls]
 Repository = "https://github.com/pydn/ComfyUI-to-Python-Extension"
 #  Used by Comfy Registry https://comfyregistry.org

tests/runtime/generated/text-to-image.py

@@ -1,80 +1,154 @@
 # 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 get_value_at_index(obj: Union[Sequence, Mapping], index: int) -> Any:
-    """Return a sequence or mapping result item by index."""
+
+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 _is_comfyui_directory(path: str) -> bool:
+    """Verify a directory has ComfyUI structural markers.
+
+    Checks for nodes.py, main.py, and the comfy/ subdirectory to raise
+    the bar against spoofing via a directory with only a single marker file.
+    """
+    if not os.path.isdir(path):
+        return False
+    return (
+        os.path.isfile(os.path.join(path, "nodes.py"))
+        and os.path.isfile(os.path.join(path, "main.py"))
+        and os.path.isdir(os.path.join(path, "comfy"))
+    )
+
+
+def _load_module(module_name: str, filepath: str) -> Any:
+    """Load a Python module from an explicit file path, bypassing sys.path.
+
+    Significantly reduces bare import shadowing risk by loading modules
+    from verified file paths instead of relying on sys.path resolution.
+    If exec_module() raises, the partially-loaded module is removed from
+    sys.modules so subsequent calls start fresh.
+    """
+    # Return cached module if already loaded — prevents re-execution
+    # that would reset state (e.g. NODE_CLASS_MAPPINGS after init_extra_nodes).
+    if module_name in sys.modules:
+        return sys.modules[module_name]
+    if not os.path.isfile(filepath):
+        log.debug("Module file not found: %s (%s)", module_name, filepath)
+        return None
     try:
-        return obj[index]
-    except KeyError:
-        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 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()
-
-    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:
+        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
+        try:
+            spec.loader.exec_module(mod)
+        except BaseException:
+            sys.modules.pop(module_name, None)
+            raise
+        return mod
+    except BaseException as e:
+        log.debug("Failed to load %s from %s: %s", module_name, filepath, e)
+        sys.modules.pop(module_name, None)  # Also clean up on edge-case exceptions
         return None
 
-    return find_path(name, parent_directory)
+
+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 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,12 +173,102 @@ 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"
 
 
+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(
+            "COMFYUI_TOPYTHON_UNLOAD_MODELS", ""
+        ).lower() in {
+            "1",
+            "true",
+            "yes",
+            "on",
+        }
+
+    try:
+        import comfy.model_management as model_management
+    except ModuleNotFoundError:
+        gc.collect()
+        return
+
+    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 find_path(name: str, max_depth: int = 20) -> str | None:
+    """Iteratively walk up from CWD to find a directory by name.
+
+    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 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
+    p = find_path("ComfyUI", max_depth=20)
+    if p and _is_comfyui_directory(p):
+        return p
+    return None
+
+
+def get_value_at_index(obj: Union[Sequence, Mapping], index: int) -> Any:
+    """Return a sequence or mapping result item by index."""
+    try:
+        return obj[index]
+    except KeyError:
+        return obj["result"][index]
+
+
 # Workflow data
 def build_workflow() -> dict[str, Any]:
     return {
@@ -163,7 +327,7 @@ extra_pnginfo = build_extra_pnginfo()
 
 
 # Workflow execution
-def main():
+def main(unload_models: bool | None = None):
     bootstrap_comfyui_runtime()
     add_extra_model_paths()
 
@@ -180,51 +344,54 @@ def main():
 
     import torch
 
-    with torch.inference_mode():
-        checkpointloadersimple = CheckpointLoaderSimple()
-        checkpointloadersimple_1 = checkpointloadersimple.load_checkpoint(
-            ckpt_name="v1-5-pruned-emaonly-fp16.safetensors"
-        )
-        cliptextencode = CLIPTextEncode()
-        cliptextencode_2 = cliptextencode.encode(
-            text="a small cottage in a meadow, soft daylight",
-            clip=get_value_at_index(checkpointloadersimple_1, 1),
-        )
-        cliptextencode_3 = cliptextencode.encode(
-            text="blurry, low quality",
-            clip=get_value_at_index(checkpointloadersimple_1, 1),
-        )
-        emptylatentimage = EmptyLatentImage()
-        emptylatentimage_4 = emptylatentimage.generate(
-            width=512, height=512, batch_size=1
-        )
-        ksampler = KSampler()
-        vaedecode = VAEDecode()
-        saveimage = SaveImage()
-        for q in range(1):
-            node_5_seed = prompt["5"]["inputs"]["seed"] = random.randint(1, 2**64)
-            ksampler_5 = ksampler.sample(
-                seed=node_5_seed,
-                steps=4,
-                cfg=7,
-                sampler_name="euler",
-                scheduler="normal",
-                denoise=1,
-                model=get_value_at_index(checkpointloadersimple_1, 0),
-                positive=get_value_at_index(cliptextencode_2, 0),
-                negative=get_value_at_index(cliptextencode_3, 0),
-                latent_image=get_value_at_index(emptylatentimage_4, 0),
+    try:
+        with torch.inference_mode():
+            checkpointloadersimple = CheckpointLoaderSimple()
+            checkpointloadersimple_1 = checkpointloadersimple.load_checkpoint(
+                ckpt_name="v1-5-pruned-emaonly-fp16.safetensors"
             )
-            vaedecode_6 = vaedecode.decode(
-                samples=get_value_at_index(ksampler_5, 0),
-                vae=get_value_at_index(checkpointloadersimple_1, 2),
+            cliptextencode = CLIPTextEncode()
+            cliptextencode_2 = cliptextencode.encode(
+                text="a small cottage in a meadow, soft daylight",
+                clip=get_value_at_index(checkpointloadersimple_1, 1),
             )
-            saveimage_7 = saveimage.save_images(
-                filename_prefix="E2E_text_to_image",
-                images=get_value_at_index(vaedecode_6, 0),
-                prompt=prompt,
-                extra_pnginfo=extra_pnginfo,
+            cliptextencode_3 = cliptextencode.encode(
+                text="blurry, low quality",
+                clip=get_value_at_index(checkpointloadersimple_1, 1),
             )
+            emptylatentimage = EmptyLatentImage()
+            emptylatentimage_4 = emptylatentimage.generate(
+                width=512, height=512, batch_size=1
+            )
+            ksampler = KSampler()
+            vaedecode = VAEDecode()
+            saveimage = SaveImage()
+            for q in range(1):
+                node_5_seed = prompt["5"]["inputs"]["seed"] = random.randint(1, 2**64)
+                ksampler_5 = ksampler.sample(
+                    seed=node_5_seed,
+                    steps=4,
+                    cfg=7,
+                    sampler_name="euler",
+                    scheduler="normal",
+                    denoise=1,
+                    model=get_value_at_index(checkpointloadersimple_1, 0),
+                    positive=get_value_at_index(cliptextencode_2, 0),
+                    negative=get_value_at_index(cliptextencode_3, 0),
+                    latent_image=get_value_at_index(emptylatentimage_4, 0),
+                )
+                vaedecode_6 = vaedecode.decode(
+                    samples=get_value_at_index(ksampler_5, 0),
+                    vae=get_value_at_index(checkpointloadersimple_1, 2),
+                )
+                saveimage_7 = saveimage.save_images(
+                    filename_prefix="E2E_text_to_image",
+                    images=get_value_at_index(vaedecode_6, 0),
+                    prompt=prompt,
+                    extra_pnginfo=extra_pnginfo,
+                )
+    finally:
+        cleanup_comfyui_runtime(unload_models=unload_models)
 
 
 # Entrypoint

tests/runtime/generated/upscale-model-loader.py

@@ -1,80 +1,154 @@
 # 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 get_value_at_index(obj: Union[Sequence, Mapping], index: int) -> Any:
-    """Return a sequence or mapping result item by index."""
+
+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 _is_comfyui_directory(path: str) -> bool:
+    """Verify a directory has ComfyUI structural markers.
+
+    Checks for nodes.py, main.py, and the comfy/ subdirectory to raise
+    the bar against spoofing via a directory with only a single marker file.
+    """
+    if not os.path.isdir(path):
+        return False
+    return (
+        os.path.isfile(os.path.join(path, "nodes.py"))
+        and os.path.isfile(os.path.join(path, "main.py"))
+        and os.path.isdir(os.path.join(path, "comfy"))
+    )
+
+
+def _load_module(module_name: str, filepath: str) -> Any:
+    """Load a Python module from an explicit file path, bypassing sys.path.
+
+    Significantly reduces bare import shadowing risk by loading modules
+    from verified file paths instead of relying on sys.path resolution.
+    If exec_module() raises, the partially-loaded module is removed from
+    sys.modules so subsequent calls start fresh.
+    """
+    # Return cached module if already loaded — prevents re-execution
+    # that would reset state (e.g. NODE_CLASS_MAPPINGS after init_extra_nodes).
+    if module_name in sys.modules:
+        return sys.modules[module_name]
+    if not os.path.isfile(filepath):
+        log.debug("Module file not found: %s (%s)", module_name, filepath)
+        return None
     try:
-        return obj[index]
-    except KeyError:
-        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 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()
-
-    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:
+        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
+        try:
+            spec.loader.exec_module(mod)
+        except BaseException:
+            sys.modules.pop(module_name, None)
+            raise
+        return mod
+    except BaseException as e:
+        log.debug("Failed to load %s from %s: %s", module_name, filepath, e)
+        sys.modules.pop(module_name, None)  # Also clean up on edge-case exceptions
         return None
 
-    return find_path(name, parent_directory)
+
+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 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,33 +173,156 @@ 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"
 
 
+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(
+            "COMFYUI_TOPYTHON_UNLOAD_MODELS", ""
+        ).lower() in {
+            "1",
+            "true",
+            "yes",
+            "on",
+        }
+
+    try:
+        import comfy.model_management as model_management
+    except ModuleNotFoundError:
+        gc.collect()
+        return
+
+    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 find_path(name: str, max_depth: int = 20) -> str | None:
+    """Iteratively walk up from CWD to find a directory by name.
+
+    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 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
+    p = find_path("ComfyUI", max_depth=20)
+    if p and _is_comfyui_directory(p):
+        return p
+    return None
+
+
+def get_value_at_index(obj: Union[Sequence, Mapping], index: int) -> Any:
+    """Return a sequence or mapping result item by index."""
+    try:
+        return obj[index]
+    except KeyError:
+        return obj["result"][index]
+
+
 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"))
 
-    import server
+    # nodes.py inserts comfy/ subdirectory into sys.path, which shadows the
+    # top-level utils/ package (comfy/utils.py vs utils/). This breaks
+    # server.py → app.frontend_management → from utils.install_util import ...
+    # Filter it out temporarily so server.py loads cleanly.
+    comfy_subdir = os.path.join(comfyui_path, "comfy")
+    original_sys_path = list(sys.path)
+    sys.path[:] = [p for p in sys.path if p != comfy_subdir]
+
+    try:
+        server_mod = _load_module("server", os.path.join(comfyui_path, "server.py"))
+    finally:
+        # Restore sys.path so nodes and other modules that need comfy/ still work.
+        # Guaranteed even if _load_module raises mid-execution.
+        sys.path[:] = original_sys_path
+
+    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
@@ -160,7 +357,7 @@ extra_pnginfo = build_extra_pnginfo()
 
 
 # Workflow execution
-def main():
+def main(unload_models: bool | None = None):
     bootstrap_comfyui_runtime()
     add_extra_model_paths()
     import_custom_nodes()
@@ -170,26 +367,29 @@ def main():
 
     import torch
 
-    with torch.inference_mode():
-        loadimage = LoadImage()
-        loadimage_1 = loadimage.load_image(image="e2e_upscale_input.png")
-        upscalemodelloader = NODE_CLASS_MAPPINGS["UpscaleModelLoader"]()
-        upscalemodelloader_2 = upscalemodelloader.EXECUTE_NORMALIZED(
-            model_name="RealESRGAN_x4plus.safetensors"
-        )
-        imageupscalewithmodel = NODE_CLASS_MAPPINGS["ImageUpscaleWithModel"]()
-        saveimage = SaveImage()
-        for q in range(1):
-            imageupscalewithmodel_3 = imageupscalewithmodel.EXECUTE_NORMALIZED(
-                upscale_model=get_value_at_index(upscalemodelloader_2, 0),
-                image=get_value_at_index(loadimage_1, 0),
-            )
-            saveimage_4 = saveimage.save_images(
-                filename_prefix="E2E_upscale_model_loader",
-                images=get_value_at_index(imageupscalewithmodel_3, 0),
-                prompt=prompt,
-                extra_pnginfo=extra_pnginfo,
+    try:
+        with torch.inference_mode():
+            loadimage = LoadImage()
+            loadimage_1 = loadimage.load_image(image="e2e_upscale_input.png")
+            upscalemodelloader = NODE_CLASS_MAPPINGS["UpscaleModelLoader"]()
+            upscalemodelloader_2 = upscalemodelloader.EXECUTE_NORMALIZED(
+                model_name="RealESRGAN_x4plus.safetensors"
             )
+            imageupscalewithmodel = NODE_CLASS_MAPPINGS["ImageUpscaleWithModel"]()
+            saveimage = SaveImage()
+            for q in range(1):
+                imageupscalewithmodel_3 = imageupscalewithmodel.EXECUTE_NORMALIZED(
+                    upscale_model=get_value_at_index(upscalemodelloader_2, 0),
+                    image=get_value_at_index(loadimage_1, 0),
+                )
+                saveimage_4 = saveimage.save_images(
+                    filename_prefix="E2E_upscale_model_loader",
+                    images=get_value_at_index(imageupscalewithmodel_3, 0),
+                    prompt=prompt,
+                    extra_pnginfo=extra_pnginfo,
+                )
+    finally:
+        cleanup_comfyui_runtime(unload_models=unload_models)
 
 
 # Entrypoint

tests/test_app_base_mappings.py

@@ -0,0 +1,165 @@
+"""Tests for ExportApplication base_node_class_mappings behavior.
+
+Verifies that after custom node import, base_node_class_mappings retains the
+pre-custom-node baseline so WorkflowPlanner correctly distinguishes between
+built-in nodes (direct import) and custom/extras nodes (NODE_CLASS_MAPPINGS lookup).
+"""
+
+import sys
+import unittest
+from io import StringIO
+
+
+def _make_workflow():
+    """Minimal workflow JSON string for testing."""
+    return '{"1": {"class_type": "CheckpointLoaderSimple", "inputs": {"ckpt_name": "test.safetensors"}}}'
+
+
+class TestBaseNodeClassMappings(unittest.TestCase):
+    """Tests for base_node_class_mappings stability across custom node reloads."""
+
+    def tearDown(self):
+        # Clean up any sys.modules pollution to avoid cross-test contamination
+        if "nodes" in sys.modules:
+            del sys.modules["nodes"]
+
+    def test_base_mappings_unchanged_after_custom_node_import(self):
+        """Given custom nodes are loaded, base_node_class_mappings stays at original baseline."""
+        from comfyui_to_python.app import ExportApplication
+
+        class FakeNode:
+            CATEGORY = "loaders"
+            FUNCTION = "execute"
+
+            @staticmethod
+            def INPUT_TYPES():
+                return {"required": {"ckpt_name": ("STRING",)}}
+
+            def execute(self, ckpt_name: str):
+                return ({},)
+
+        class FakeCustomNode:
+            CATEGORY = "custom"
+            FUNCTION = "execute"
+
+            @staticmethod
+            def INPUT_TYPES():
+                return {"required": {"ckpt_name": ("STRING",)}}
+
+            def execute(self, ckpt_name: str):
+                return ({},)
+
+        initial_mappings = {"CheckpointLoaderSimple": FakeNode}
+        after_import_mappings = dict(initial_mappings)
+        # Simulate custom nodes adding a new class
+        after_import_mappings["CustomNode"] = FakeCustomNode
+
+        call_count = [0]
+
+        def mock_node_loader():
+            return initial_mappings
+
+        def mock_custom_node_importer():
+            call_count[0] += 1
+            # Simulate populating sys.modules["nodes"].NODE_CLASS_MAPPINGS
+            nodes_mod_type = type("NodesModule", (), {})()
+            nodes_mod_type.NODE_CLASS_MAPPINGS = after_import_mappings
+            sys.modules["nodes"] = nodes_mod_type
+
+        output = StringIO()
+        app = ExportApplication(
+            workflow=_make_workflow(),
+            output_file=output,
+            node_mapping_loader=mock_node_loader,
+            custom_node_importer=mock_custom_node_importer,
+            needs_init_custom_nodes=True,
+        )
+
+        # base_node_class_mappings should be a deep copy of initial (pre-custom-node) state
+        self.assertEqual(app.base_node_class_mappings, initial_mappings)
+        app.execute()
+
+        # After execute(), base_node_class_mappings must still be the original baseline
+        # (not overwritten with after_import_mappings which includes custom nodes)
+        self.assertEqual(
+            app.base_node_class_mappings,
+            initial_mappings,
+            "base_node_class_mappings should remain at pre-custom-node baseline",
+        )
+
+    def test_base_mappings_is_deep_copy_not_shared_reference(self):
+        """Given node mappings change in-place, base_node_class_mappings is independent."""
+        from comfyui_to_python.app import ExportApplication
+
+        class FakeNode2:
+            CATEGORY = "loaders"
+            FUNCTION = "execute"
+
+            @staticmethod
+            def INPUT_TYPES():
+                return {"required": {"ckpt_name": ("STRING",)}}
+
+            def execute(self, ckpt_name: str):
+                return ({},)
+
+        initial_mappings = {"CheckpointLoaderSimple": FakeNode2}
+
+        def mock_node_loader():
+            return dict(initial_mappings)
+
+        output = StringIO()
+        app = ExportApplication(
+            workflow=_make_workflow(),
+            output_file=output,
+            node_mapping_loader=mock_node_loader,
+            custom_node_importer=lambda: None,
+        )
+
+        # Mutate the live mappings
+        app.node_class_mappings["NewKey"] = "value"
+
+        # base_node_class_mappings must not reflect the mutation (deep copy)
+        self.assertNotIn(
+            "NewKey",
+            app.base_node_class_mappings,
+            "base_node_class_mappings must be a deep copy, not shared reference",
+        )
+
+
+class TestRenderHardenedImports(unittest.TestCase):
+    """Integration test: render output contains hardened _load_module pattern."""
+
+    def test_generated_output_contains_exec_failure_cleanup(self):
+        """Given a GenerationPlan, rendered output includes try/except BaseException around exec_module()."""
+        from comfyui_to_python.generator.render import WorkflowRenderer
+        from comfyui_to_python.generator.model import GenerationPlan
+
+        plan = GenerationPlan(
+            workflow_data={"1": {"class_type": "CheckpointLoaderSimple"}},
+            metadata_workflow_data=None,
+            custom_nodes=False,
+            import_statements={},
+            special_functions_code=[],
+            loop_code=["result = some_node()"],
+            queue_size=1,
+        )
+        renderer = WorkflowRenderer()
+        generated = renderer.render(plan)
+
+        # The _load_module in rendered output must have the exec failure cleanup pattern
+        self.assertIn("try:", generated)
+        self.assertIn("spec.loader.exec_module(mod)", generated)
+        self.assertIn(
+            "except BaseException:",
+            generated,
+            "Must catch BaseException for edge-case safety",
+        )
+        self.assertIn(
+            "sys.modules.pop(module_name, None)",
+            generated,
+            "Must clean up sys.modules on exec failure",
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_import_path_resolution.py

@@ -0,0 +1,712 @@
+"""Tests for hardened import path resolution and importlib isolation layer.
+
+Covers Approach B from docs/specs/harden-import-path-resolution.html:
+- _is_comfyui_directory() structural verification
+- _load_module() centralized importlib isolation
+- get_comfyui_path() multi-strategy resolution
+- sys.path idempotence (no remove/re-insert gap)
+- Security: shadowing attack resistance
+- Logging hygiene (no print() path leaks)
+"""
+
+import os
+import sys
+import tempfile
+import unittest
+from unittest.mock import patch
+
+# Real ComfyUI checkout for integration-style unit tests (parent of this repo)
+_REAL_COMFYUI = os.path.realpath(
+    os.path.join(os.path.dirname(__file__), "..", "..", "ComfyUI")
+)
+
+
+def _make_fake_comfyui(tmpdir: str) -> str:
+    """Create a minimal fake ComfyUI directory in tmpdir with structural markers."""
+    comfyui_dir = os.path.join(tmpdir, "ComfyUI")
+    os.makedirs(comfyui_dir)
+    # Structural markers checked by _is_comfyui_directory()
+    with open(os.path.join(comfyui_dir, "nodes.py"), "w") as f:
+        f.write("# fake nodes.py\n")
+    with open(os.path.join(comfyui_dir, "main.py"), "w") as f:
+        f.write("# fake main.py\n")
+    os.makedirs(os.path.join(comfyui_dir, "comfy"))
+    return comfyui_dir
+
+
+class TestLoadModule(unittest.TestCase):
+    """Tests for _load_module() centralized importlib isolation."""
+
+    def tearDown(self):
+        # Clean up sys.modules cache between tests to avoid cross-test pollution
+        mod_names_to_clean = [
+            name for name in sys.modules if name.startswith("_test_mod_")
+        ]
+        for name in mod_names_to_clean:
+            del sys.modules[name]
+
+    def test_loads_module_from_explicit_file_path(self):
+        """Given a valid .py file, loads the module and returns it with expected attributes."""
+        from comfyui_to_python.node_runtime import _load_module
+
+        # Use a simple self-contained module (nodes.py requires torch)
+        with tempfile.TemporaryDirectory() as tmpdir:
+            mod_file = os.path.join(tmpdir, "mymodule.py")
+            with open(mod_file, "w") as f:
+                f.write("VALUE = 42\nNAME = 'test_module'\n")
+
+            mod = _load_module("_test_mod_simple", mod_file)
+        self.assertIsNotNone(mod)
+        self.assertEqual(mod.VALUE, 42)
+        self.assertEqual(mod.NAME, "test_module")
+
+    def test_returns_none_for_missing_file(self):
+        """Given a nonexistent file path, returns None without raising."""
+        from comfyui_to_python.node_runtime import _load_module
+
+        result = _load_module("_test_mod_missing", "/nonexistent/path/module.py")
+        self.assertIsNone(result)
+
+
+class TestShadowingResistance(unittest.TestCase):
+    """Security: _load_module() must ignore shadowed modules at sys.path[0]."""
+
+    def tearDown(self):
+        # Clean up sys.modules cache between tests
+        mod_names_to_clean = [
+            name for name in sys.modules if name.startswith("_sec_test_")
+        ]
+        for name in mod_names_to_clean:
+            del sys.modules[name]
+
+    def test_ignores_shadowed_main_at_syspath(self):
+        """Given a malicious main.py at sys.path[0], _load_module() loads from verified path."""
+        from comfyui_to_python.node_runtime import _load_module
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            # Legitimate module
+            legit_file = os.path.join(tmpdir, "legit_main.py")
+            with open(legit_file, "w") as f:
+                f.write("ORIGIN = 'legitimate'\n")
+
+            # Malicious shadow at sys.path[0]
+            attack_dir = os.path.join(tmpdir, "attack")
+            os.makedirs(attack_dir)
+            attack_file = os.path.join(attack_dir, "main.py")
+            with open(attack_file, "w") as f:
+                f.write("ORIGIN = 'malicious'\n")
+
+            # Insert malicious path at sys.path[0]
+            old_path = sys.path[:]
+            try:
+                sys.path.insert(0, attack_dir)
+                # _load_module with explicit legit path must return legitimate module
+                mod = _load_module("_sec_test_main", legit_file)
+            finally:
+                sys.path[:] = old_path
+
+        self.assertIsNotNone(mod)
+        self.assertEqual(mod.ORIGIN, "legitimate")
+
+    def test_ignores_shadowed_nodes_at_syspath(self):
+        """Given a malicious nodes.py at sys.path[0], _load_module() loads from verified path."""
+        from comfyui_to_python.node_runtime import _load_module
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            # Legitimate module
+            legit_file = os.path.join(tmpdir, "legit_nodes.py")
+            with open(legit_file, "w") as f:
+                f.write("NODE_CLASS_MAPPINGS = {'RealNode': 'real'}\n")
+
+            # Malicious shadow at sys.path[0]
+            attack_dir = os.path.join(tmpdir, "attack")
+            os.makedirs(attack_dir)
+            attack_file = os.path.join(attack_dir, "nodes.py")
+            with open(attack_file, "w") as f:
+                f.write("NODE_CLASS_MAPPINGS = {'FakeNode': 'malicious'}\n")
+
+            old_path = sys.path[:]
+            try:
+                sys.path.insert(0, attack_dir)
+                mod = _load_module("_sec_test_nodes", legit_file)
+            finally:
+                sys.path[:] = old_path
+
+        self.assertIsNotNone(mod)
+        self.assertEqual(mod.NODE_CLASS_MAPPINGS, {"RealNode": "real"})
+
+
+@unittest.skipUnless(os.path.isdir(_REAL_COMFYUI), "requires ../ComfyUI checkout")
+class TestSysPathIdempotence(unittest.TestCase):
+    """Tests for sys.path idempotence — no remove/re-insert gap."""
+
+    def test_add_comfyui_idempotent_no_remove_insert(self):
+        """Repeated calls to add_comfyui_directory_to_sys_path() are insert-once only."""
+        from comfyui_to_python.node_runtime import (
+            add_comfyui_directory_to_sys_path,
+        )
+
+        with patch.dict("os.environ", {"COMFYUI_PATH": _REAL_COMFYUI}, clear=False):
+            before_length = len(sys.path)
+            add_comfyui_directory_to_sys_path()
+            after_first = len(sys.path)
+
+            # Second call — must not change sys.path at all (no remove+re-insert)
+            add_comfyui_directory_to_sys_path()
+            after_second = len(sys.path)
+
+        self.assertEqual(after_first, after_second, "sys.path changed on second call")
+        # ComfyUI should be on sys.path
+        self.assertIn(_REAL_COMFYUI, sys.path)
+
+    def test_already_at_index_zero_is_noop(self):
+        """Given ComfyUI already at sys.path[0], add_comfyui_directory_to_sys_path() is a no-op."""
+        from comfyui_to_python.node_runtime import (
+            add_comfyui_directory_to_sys_path,
+        )
+
+        original_path = sys.path[:]
+        try:
+            with patch.dict("os.environ", {"COMFYUI_PATH": _REAL_COMFYUI}, clear=False):
+                # Put ComfyUI at index 0 manually
+                if _REAL_COMFYUI in sys.path:
+                    sys.path.remove(_REAL_COMFYUI)
+                sys.path.insert(0, _REAL_COMFYUI)
+                original_index = sys.path.index(_REAL_COMFYUI)
+
+                add_comfyui_directory_to_sys_path()
+                new_index = sys.path.index(_REAL_COMFYUI)
+        finally:
+            sys.path[:] = original_path
+
+        self.assertEqual(
+            original_index, new_index, "Index should not change when already at [0]"
+        )
+        self.assertEqual(new_index, 0)
+
+    def test_promotes_comfyui_to_sys_path_zero(self):
+        """Given ComfyUI at sys.path[5], add_comfyui_directory_to_sys_path() moves it to index 0."""
+        from comfyui_to_python.node_runtime import (
+            add_comfyui_directory_to_sys_path,
+        )
+
+        original_path = sys.path[:]
+        try:
+            with patch.dict("os.environ", {"COMFYUI_PATH": _REAL_COMFYUI}, clear=False):
+                # Insert ComfyUI deep in sys.path (simulating PYTHONPATH placement)
+                if _REAL_COMFYUI in sys.path:
+                    sys.path.remove(_REAL_COMFYUI)
+                sys.path.insert(5, _REAL_COMFYUI)
+
+                add_comfyui_directory_to_sys_path()
+                self.assertEqual(
+                    sys.path[0],
+                    _REAL_COMFYUI,
+                    "ComfyUI must be promoted to sys.path[0] when already present lower down",
+                )
+        finally:
+            sys.path[:] = original_path
+
+    def test_import_custom_nodes_no_gap_window(self):
+        """import_custom_nodes() must never remove ComfyUI from sys.path during execution."""
+        import inspect
+        from comfyui_to_python.node_runtime import import_custom_nodes
+
+        source = inspect.getsource(import_custom_nodes)
+        self.assertNotIn(
+            "sys.path.remove",
+            source,
+            "import_custom_nodes() must not remove ComfyUI from sys.path (gap window)",
+        )
+
+    def test_no_print_statements_in_node_runtime(self):
+        """node_runtime.py must use logging, not print(), to avoid path leaks to stdout."""
+        import inspect
+        from comfyui_to_python import node_runtime
+
+        source = inspect.getsource(node_runtime)
+        self.assertNotIn(
+            "print(",
+            source,
+            "node_runtime.py must not use print() — use logging.DEBUG instead",
+        )
+
+
+class TestGeneratedScriptIsolation(unittest.TestCase):
+    """Tests for generated script bootstrap (Phase 3)."""
+
+    def test_generated_script_embeds_load_module(self):
+        """render.py must embed _load_module() in generated scripts."""
+        from comfyui_to_python.generator.render import WorkflowRenderer
+        from comfyui_to_python.generator.model import GenerationPlan
+
+        plan = GenerationPlan(
+            workflow_data={"1": {"class_type": "CheckpointLoaderSimple"}},
+            metadata_workflow_data=None,
+            custom_nodes=False,
+            import_statements={},
+            special_functions_code=[],
+            loop_code=["result = some_node()"],
+            queue_size=1,
+        )
+        renderer = WorkflowRenderer()
+        generated = renderer.render(plan)
+
+        self.assertIn(
+            "_load_module", generated, "Generated script must embed _load_module()"
+        )
+        self.assertIn(
+            "spec_from_file_location",
+            generated,
+            "Generated script must use importlib isolation",
+        )
+
+    def test_generated_script_no_bare_comfyui_imports(self):
+        """Generated scripts must have zero bare imports of ComfyUI internals."""
+        from comfyui_to_python.generator.render import WorkflowRenderer
+        from comfyui_to_python.generator.model import GenerationPlan
+
+        plan = GenerationPlan(
+            workflow_data={"1": {"class_type": "CheckpointLoaderSimple"}},
+            metadata_workflow_data=None,
+            custom_nodes=True,
+            import_statements={},
+            special_functions_code=[],
+            loop_code=["result = some_node()"],
+            queue_size=1,
+        )
+        renderer = WorkflowRenderer()
+        generated = renderer.render(plan)
+
+        # These bare imports must NOT appear in generated output
+        bare_imports = [
+            "from main import",
+            "from utils.extra_config import",
+            "import execution",
+            "import server",
+            "from nodes import",
+            "import comfy.options",
+            "from comfy.cli_args import",
+        ]
+        for bare in bare_imports:
+            self.assertNotIn(
+                bare,
+                generated,
+                f"Generated script must not contain bare import: {bare}",
+            )
+
+
+class TestGetComfyuiPath(unittest.TestCase):
+    """Tests for get_comfyui_path() multi-strategy resolution."""
+
+    def test_resolves_via_valid_comfyui_path_env(self):
+        """Given COMFYUI_PATH env var set to valid ComfyUI dir, returns that path."""
+        from comfyui_to_python.node_runtime import get_comfyui_path
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            fake_comfyui = _make_fake_comfyui(tmpdir)
+            original_cwd = os.getcwd()
+            try:
+                os.chdir(fake_comfyui)  # strategy 2 would find it, but env wins first
+            finally:
+                pass
+            with patch.dict("os.environ", {"COMFYUI_PATH": fake_comfyui}, clear=False):
+                result = get_comfyui_path()
+            os.chdir(original_cwd)
+        self.assertEqual(result, fake_comfyui)
+
+    def test_falls_through_to_next_strategy_when_env_invalid(self):
+        """Given invalid COMFYUI_PATH (no nodes.py), rejects and falls through."""
+        from comfyui_to_python.node_runtime import get_comfyui_path
+
+        with patch.dict(
+            "os.environ", {"COMFYUI_PATH": "/tmp/not-comfyui"}, clear=False
+        ):
+            result = get_comfyui_path()
+        # The invalid path must be rejected (not returned as-is)
+        self.assertNotEqual(result, "/tmp/not-comfyui")
+        # In standalone/dev setups, no strategy may succeed -> None is valid
+        # In ComfyUI/custom_nodes/ install, relative strategy would find it
+
+
+class TestIsComfyuiDirectory(unittest.TestCase):
+    """Tests for _is_comfyui_directory() structural verification."""
+
+    def test_returns_true_for_real_comfyui_checkout(self):
+        """Given a valid ComfyUI directory with structural markers, returns True."""
+        from comfyui_to_python.node_runtime import _is_comfyui_directory
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            fake_comfyui = _make_fake_comfyui(tmpdir)
+            self.assertTrue(_is_comfyui_directory(fake_comfyui))
+
+    def test_rejects_directory_lacking_markers(self):
+        """Given a directory without nodes.py, returns False."""
+        from comfyui_to_python.node_runtime import _is_comfyui_directory
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            # Empty directory has no ComfyUI markers
+            self.assertFalse(_is_comfyui_directory(tmpdir))
+
+    def test_rejects_nonexistent_path(self):
+        """Given a nonexistent path, returns False."""
+        from comfyui_to_python.node_runtime import _is_comfyui_directory
+
+        self.assertFalse(_is_comfyui_directory("/nonexistent/path/to/comfyui"))
+
+
+class TestLoadModuleFailureCleanup(unittest.TestCase):
+    """Tests for _load_module() behavior when exec_module() raises."""
+
+    def tearDown(self):
+        mod_names_to_clean = [
+            name for name in sys.modules if name.startswith("_test_fail_")
+        ]
+        for name in mod_names_to_clean:
+            del sys.modules[name]
+
+    def test_broken_module_removed_from_sys_modules_on_failure(self):
+        """Given a module that raises during exec, it is removed from sys.modules."""
+        from comfyui_to_python.node_runtime import _load_module
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            mod_file = os.path.join(tmpdir, "badmodule.py")
+            with open(mod_file, "w") as f:
+                f.write("raise RuntimeError('mid-execution failure')\n")
+
+            result = _load_module("_test_fail_bad", mod_file)
+        self.assertIsNone(result)
+        self.assertNotIn(
+            "_test_fail_bad",
+            sys.modules,
+            "Failed module must not remain cached in sys.modules",
+        )
+
+    def test_broken_module_retried_cleanly(self):
+        """Given a failed load, a subsequent load attempt starts fresh (not from cache)."""
+        from comfyui_to_python.node_runtime import _load_module
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            bad_file = os.path.join(tmpdir, "bad.py")
+            good_file = os.path.join(tmpdir, "good.py")
+
+            # First load: bad module that raises
+            with open(bad_file, "w") as f:
+                f.write("raise RuntimeError('fail')\n")
+            result_bad = _load_module("_test_fail_retry", bad_file)
+            self.assertIsNone(result_bad)
+
+            # Overwrite with good module
+            os.remove(bad_file)
+            with open(good_file, "w") as f:
+                f.write("VALUE = 99\n")
+
+            # Reload same name from a different (working) file
+            result_good = _load_module("_test_fail_retry", good_file)
+        self.assertIsNotNone(result_good)
+        self.assertEqual(result_good.VALUE, 99)
+
+
+class TestCanonicalOptionsLoading(unittest.TestCase):
+    """Tests for canonical module loading of comfy.options in bootstrap."""
+
+    def tearDown(self):
+        # Clean up sys.modules cache
+        mod_names_to_clean = [
+            name for name in sys.modules if name.startswith("_test_opts_")
+        ]
+        for name in mod_names_to_clean:
+            del sys.modules[name]
+        if "comfy.options" in sys.modules:
+            del sys.modules["comfy.options"]
+
+    def test_load_under_canonical_name_caches_in_sys_modules(self):
+        """When _load_module() uses canonical name, the module is cached in sys.modules."""
+        from comfyui_to_python.node_runtime import _load_module
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            mod_file = os.path.join(tmpdir, "options.py")
+            with open(mod_file, "w") as f:
+                f.write("args_parsing = False\n")
+
+            # Load under canonical name
+            mod = _load_module("comfy.options", mod_file)
+        self.assertIsNotNone(mod)
+        # The canonical key must be in sys.modules so later imports see mutations
+        self.assertIn(
+            "comfy.options",
+            sys.modules,
+            "Canonical module load must cache under canonical name in sys.modules",
+        )
+
+
+class TestLoadModuleTemp(unittest.TestCase):
+    """Tests for _load_module_temp() removal from sys.modules."""
+
+    def tearDown(self):
+        mod_names_to_clean = [
+            name for name in sys.modules if name.startswith("_test_temp_")
+        ]
+        for name in mod_names_to_clean:
+            del sys.modules[name]
+
+    def test_module_removed_from_sys_modules_after_temp_load(self):
+        """After _load_module_temp(), the module key is absent from sys.modules."""
+        from comfyui_to_python.node_runtime import _load_module_temp
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            mod_file = os.path.join(tmpdir, "tempmod.py")
+            with open(mod_file, "w") as f:
+                f.write("VALUE = 123\n")
+
+            mod = _load_module_temp("_test_temp_mod", mod_file)
+        self.assertIsNotNone(mod)
+        self.assertEqual(mod.VALUE, 123)
+        self.assertNotIn(
+            "_test_temp_mod",
+            sys.modules,
+            "Temp module must be removed from sys.modules after load",
+        )
+
+    def test_module_removed_on_temp_load_failure(self):
+        """After _load_module_temp() fails, the module key is absent from sys.modules."""
+        from comfyui_to_python.node_runtime import _load_module_temp
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            mod_file = os.path.join(tmpdir, "badtemp.py")
+            with open(mod_file, "w") as f:
+                f.write("raise RuntimeError('fail')\n")
+
+            result = _load_module_temp("_test_temp_bad", mod_file)
+        self.assertIsNone(result)
+        self.assertNotIn(
+            "_test_temp_bad",
+            sys.modules,
+            "Failed temp module must be removed from sys.modules",
+        )
+
+
+class TestRepeatedNodeLoad(unittest.TestCase):
+    """Tests that repeated _load_module('nodes', ...) calls respect the cache."""
+
+    def tearDown(self):
+        if "_test_nodes" in sys.modules:
+            del sys.modules["_test_nodes"]
+
+    def test_repeated_load_preserves_cached_attributes(self):
+        """Second _load_module() for same name returns cached module, not a fresh copy."""
+        from comfyui_to_python.node_runtime import _load_module
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            mod_file = os.path.join(tmpdir, "mynodes.py")
+            with open(mod_file, "w") as f:
+                f.write("NODE_CLASS_MAPPINGS = {'A': 1}\n")
+
+            first = _load_module("_test_nodes", mod_file)
+            self.assertIsNotNone(first)
+
+            # Mutate the cached module to simulate init_extra_nodes()
+            first.NODE_CLASS_MAPPINGS["B"] = 2
+
+            # Second load with same name must return cached copy, preserving mutation
+            second = _load_module("_test_nodes", mod_file)
+        self.assertIs(first, second)
+        self.assertEqual(second.NODE_CLASS_MAPPINGS.get("B"), 2)
+
+
+class TestFindFile(unittest.TestCase):
+    """Tests for _find_file() — file discovery via directory walk."""
+
+    def test_finds_file_at_cwd_level(self):
+        """Given a file at CWD, _find_file() returns the full path to it."""
+        from comfyui_to_python.node_runtime import _find_file
+
+        original_cwd = os.getcwd()
+        with tempfile.TemporaryDirectory() as tmpdir:
+            target_file = os.path.join(tmpdir, "extra_model_paths.yaml")
+            with open(target_file, "w") as f:
+                f.write("config: value\n")
+
+            try:
+                os.chdir(tmpdir)
+                result = _find_file("extra_model_paths.yaml")
+            finally:
+                os.chdir(original_cwd)
+
+        self.assertEqual(result, target_file)
+
+    def test_finds_file_in_parent_directory(self):
+        """Given a file in a parent directory, _find_file() walks up and finds it."""
+        from comfyui_to_python.node_runtime import _find_file
+
+        original_cwd = os.getcwd()
+        with tempfile.TemporaryDirectory() as tmpdir:
+            target_file = os.path.join(tmpdir, "extra_model_paths.yaml")
+            with open(target_file, "w") as f:
+                f.write("config: value\n")
+
+            # Create a subdirectory and chdir into it
+            subdir = os.path.join(tmpdir, "subdir", "deeper")
+            os.makedirs(subdir)
+
+            try:
+                os.chdir(subdir)
+                result = _find_file("extra_model_paths.yaml")
+            finally:
+                os.chdir(original_cwd)
+
+        self.assertEqual(result, target_file)
+
+    def test_returns_none_when_file_not_found(self):
+        """Given a filename that doesn't exist within max_depth, returns None."""
+        from comfyui_to_python.node_runtime import _find_file
+
+        original_cwd = os.getcwd()
+        with tempfile.TemporaryDirectory() as tmpdir:
+            try:
+                os.chdir(tmpdir)
+                result = _find_file("nonexistent_config.yaml")
+            finally:
+                os.chdir(original_cwd)
+        self.assertIsNone(result)
+
+    def test_returns_none_for_directory_not_file(self):
+        """Given a directory with matching name, returns None (wants a file)."""
+        from comfyui_to_python.node_runtime import _find_file
+
+        original_cwd = os.getcwd()
+        with tempfile.TemporaryDirectory() as tmpdir:
+            # Create a directory (not file) with the target name
+            fake_dir = os.path.join(tmpdir, "extra_model_paths.yaml")
+            os.makedirs(fake_dir)
+
+            try:
+                os.chdir(tmpdir)
+                result = _find_file("extra_model_paths.yaml")
+            finally:
+                os.chdir(original_cwd)
+
+        self.assertIsNone(result)
+
+
+class TestFindPathCwdFirst(unittest.TestCase):
+    """Tests for find_path() checking CWD before walking to parent."""
+
+    def test_find_path_returns_cwd_when_cwd_is_match(self):
+        """Given CWD is named 'ComfyUI', find_path('ComfyUI') returns CWD itself."""
+        from comfyui_to_python.node_runtime import find_path
+
+        original_cwd = os.getcwd()
+        with tempfile.TemporaryDirectory() as tmpdir:
+            comfyui_dir = os.path.join(tmpdir, "ComfyUI")
+            os.makedirs(comfyui_dir)
+            try:
+                os.chdir(comfyui_dir)
+                result = find_path("ComfyUI")
+            finally:
+                os.chdir(original_cwd)
+
+        self.assertEqual(result, comfyui_dir)
+
+    def test_find_path_checks_cwd_before_parent(self):
+        """When CWD matches but parent also matches, CWD is returned (check-first)."""
+        from comfyui_to_python.node_runtime import find_path
+
+        original_cwd = os.getcwd()
+        with tempfile.TemporaryDirectory() as tmpdir:
+            # Create nested dirs: outer/ComfyUI/inner where inner is named ComfyUI too
+            outer_comfy = os.path.join(tmpdir, "ComfyUI")
+            os.makedirs(outer_comfy)
+            # Rename so we can create an inner dir also named ComfyUI
+            os.rename(outer_comfy, os.path.join(tmpdir, "ParentComfy"))
+            parent_comfy = os.path.join(tmpdir, "ParentComfy")
+            child_dir = os.path.join(parent_comfy, "subdir")
+            os.makedirs(child_dir)
+
+            try:
+                os.chdir(child_dir)
+                # When searching for ParentComfy, should find it at parent level
+                result = find_path("ParentComfy")
+            finally:
+                os.chdir(original_cwd)
+
+        self.assertEqual(result, parent_comfy)
+
+
+class TestGetComfyuiPathThirdStrategy(unittest.TestCase):
+    """Tests for get_comfyui_path() third-strategy verification."""
+
+    def test_fallback_rejects_non_comfyui_directory(self):
+        """Given a directory named ComfyUI without nodes.py, third strategy rejects it."""
+        from comfyui_to_python.node_runtime import get_comfyui_path
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            # Create a fake "ComfyUI" directory (no nodes.py)
+            fake_dir = os.path.join(tmpdir, "ComfyUI")
+            os.makedirs(fake_dir)
+
+            original_cwd = os.getcwd()
+            try:
+                os.chdir(tmpdir)
+                with patch.dict("os.environ", {}, clear=False):
+                    # Strip COMFYUI_PATH so fallback strategies activate
+                    env_copy = dict(os.environ)
+                    env_copy.pop("COMFYUI_PATH", None)
+                    with patch.dict("os.environ", env_copy, clear=True):
+                        result = get_comfyui_path()
+            finally:
+                os.chdir(original_cwd)
+
+        # Must NOT return the fake directory
+        self.assertNotEqual(result, fake_dir)
+
+
+@unittest.skipUnless(os.path.isdir(_REAL_COMFYUI), "requires ../ComfyUI checkout")
+class TestImportCustomNodesSysPathRestoration(unittest.TestCase):
+    """Tests for sys.path restoration in import_custom_nodes()."""
+
+    def test_sys_path_restored_on_crash(self):
+        """Given a crash during import_custom_nodes(), sys.path is still restored.
+
+        The crash fires on the "server" module load (after comfy/ filtering)
+        so that the try/finally sys.path restoration path is actually exercised.
+        """
+        from comfyui_to_python.node_runtime import import_custom_nodes
+
+        original_path = sys.path[:]
+        call_order = []
+
+        def fake_load(name, path):
+            call_order.append(name)
+            if name == "server":
+                raise RuntimeError("simulated crash during server load")
+            return None  # Let execution/nodes return None (falls through gracefully)
+
+        with (
+            patch(
+                "comfyui_to_python.node_runtime._load_module",
+                side_effect=fake_load,
+            ),
+            patch(
+                "comfyui_to_python.node_runtime.get_comfyui_path",
+                return_value=_REAL_COMFYUI if os.path.isdir(_REAL_COMFYUI) else None,
+            ),
+        ):
+            try:
+                import_custom_nodes()
+            except RuntimeError:
+                pass  # Expected crash
+
+        # When ComfyUI path is found, "execution" and "nodes" should be called
+        # before reaching the server module load inside try/finally.
+        self.assertIn("execution", call_order, "Expected 'execution' to be loaded")
+        self.assertIn("nodes", call_order, "Expected 'nodes' to be loaded")
+        self.assertIn("server", call_order, "Expected 'server' to trigger the crash")
+
+        # sys.path should still contain all original entries (no dangling state)
+        for entry in original_path:
+            self.assertIn(entry, sys.path)
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_node_runtime_cleanup.py

@@ -0,0 +1,116 @@
+import sys
+import types
+import unittest
+import warnings
+from unittest.mock import Mock, patch
+
+from comfyui_to_python.node_runtime import cleanup_comfyui_runtime
+
+
+class NodeRuntimeCleanupTest(unittest.TestCase):
+    def test_cleanup_releases_caches_without_forcing_model_unload_by_default(self):
+        comfy_module = types.ModuleType("comfy")
+        comfy_module.__path__ = []
+        model_management = types.ModuleType("comfy.model_management")
+        model_management.cleanup_models_gc = Mock()
+        model_management.unload_all_models = Mock()
+        model_management.soft_empty_cache = Mock()
+        comfy_module.model_management = model_management
+
+        with patch.dict(
+            sys.modules,
+            {
+                "comfy": comfy_module,
+                "comfy.model_management": model_management,
+            },
+        ), patch.dict("os.environ", {}, clear=False):
+            cleanup_comfyui_runtime()
+
+        model_management.cleanup_models_gc.assert_called_once_with()
+        model_management.soft_empty_cache.assert_called_once_with()
+        model_management.unload_all_models.assert_not_called()
+
+    def test_cleanup_can_force_model_unload_from_environment(self):
+        comfy_module = types.ModuleType("comfy")
+        comfy_module.__path__ = []
+        model_management = types.ModuleType("comfy.model_management")
+        model_management.cleanup_models_gc = Mock()
+        model_management.unload_all_models = Mock()
+        model_management.soft_empty_cache = Mock()
+        comfy_module.model_management = model_management
+
+        with patch.dict(
+            sys.modules,
+            {
+                "comfy": comfy_module,
+                "comfy.model_management": model_management,
+            },
+        ), patch.dict(
+            "os.environ",
+            {"COMFYUI_TOPYTHON_UNLOAD_MODELS": "true"},
+            clear=False,
+        ):
+            cleanup_comfyui_runtime()
+
+        model_management.unload_all_models.assert_called_once_with()
+
+    def test_cleanup_suppresses_hook_failures_and_warns(self):
+        comfy_module = types.ModuleType("comfy")
+        comfy_module.__path__ = []
+        model_management = types.ModuleType("comfy.model_management")
+        model_management.cleanup_models_gc = Mock(side_effect=RuntimeError("gc failed"))
+        model_management.unload_all_models = Mock(side_effect=RuntimeError("unload failed"))
+        model_management.soft_empty_cache = Mock(side_effect=RuntimeError("cache failed"))
+        comfy_module.model_management = model_management
+
+        with patch.dict(
+            sys.modules,
+            {
+                "comfy": comfy_module,
+                "comfy.model_management": model_management,
+            },
+        ), warnings.catch_warnings(record=True) as caught:
+            warnings.simplefilter("always")
+            cleanup_comfyui_runtime(unload_models=True)
+
+        model_management.cleanup_models_gc.assert_called_once_with()
+        model_management.unload_all_models.assert_called_once_with()
+        model_management.soft_empty_cache.assert_called_once_with()
+        self.assertEqual(len(caught), 3)
+        self.assertEqual(
+            [str(warning.message) for warning in caught],
+            [
+                "ComfyUI cleanup hook cleanup_models_gc failed during teardown: gc failed",
+                "ComfyUI cleanup hook unload_all_models failed during teardown: unload failed",
+                "ComfyUI cleanup hook soft_empty_cache failed during teardown: cache failed",
+            ],
+        )
+
+    def test_cleanup_does_not_mask_active_workflow_exception(self):
+        comfy_module = types.ModuleType("comfy")
+        comfy_module.__path__ = []
+        model_management = types.ModuleType("comfy.model_management")
+        model_management.cleanup_models_gc = Mock(side_effect=RuntimeError("cleanup failed"))
+        model_management.soft_empty_cache = Mock()
+        comfy_module.model_management = model_management
+
+        with patch.dict(
+            sys.modules,
+            {
+                "comfy": comfy_module,
+                "comfy.model_management": model_management,
+            },
+        ), warnings.catch_warnings(record=True):
+            warnings.simplefilter("always")
+            with self.assertRaisesRegex(ValueError, "workflow failed"):
+                try:
+                    raise ValueError("workflow failed")
+                finally:
+                    cleanup_comfyui_runtime()
+
+        model_management.cleanup_models_gc.assert_called_once_with()
+        model_management.soft_empty_cache.assert_called_once_with()
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_project_contracts.py

@@ -0,0 +1,33 @@
+import tomllib
+import unittest
+from pathlib import Path
+
+
+REPO_ROOT = Path(__file__).resolve().parent.parent
+
+
+class ProjectContractsTest(unittest.TestCase):
+    def test_project_declares_supported_python_floor(self):
+        pyproject = tomllib.loads((REPO_ROOT / "pyproject.toml").read_text(encoding="utf-8"))
+
+        self.assertEqual(pyproject["project"]["requires-python"], ">=3.12")
+
+    def test_readme_documents_python_support_and_default_save_filename(self):
+        readme = (REPO_ROOT / "README.md").read_text(encoding="utf-8")
+
+        self.assertIn("This project supports Python 3.12 and newer.", readme)
+        self.assertIn("default filename `workflow_api.py`", readme)
+
+    def test_extension_import_path_requires_uv_sync_instead_of_running_install_py(self):
+        init_text = (REPO_ROOT / "__init__.py").read_text(encoding="utf-8")
+
+        self.assertIn("Run 'uv sync'", init_text)
+        self.assertNotIn("spec_from_file_location", init_text)
+        self.assertNotIn("Successfully installed. Hopefully, at least.", init_text)
+
+    def test_frontend_save_flow_uses_deterministic_filename_without_prompt(self):
+        save_as_script = (REPO_ROOT / "js" / "save-as-script.js").read_text(encoding="utf-8")
+
+        self.assertIn('const DEFAULT_SCRIPT_FILENAME = "workflow_api.py";', save_as_script)
+        self.assertNotIn("prompt(", save_as_script)
+

tests/test_upscale_model_loader_export.py

@@ -145,20 +145,29 @@ class UpscaleModelLoaderExportTest(unittest.TestCase):
         generated = output.getvalue()
 
         self.assertIn("def bootstrap_comfyui_runtime()", generated)
-        self.assertIn("import comfy.options", generated)
-        self.assertIn("comfy.options.enable_args_parsing()", generated)
-        self.assertIn("import cuda_malloc", generated)
+        self.assertIn("def cleanup_comfyui_runtime(", 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.py", generated)
         self.assertNotIn("\nbootstrap_comfyui_runtime()\n", generated)
         self.assertIn(
-            "def main():\n    bootstrap_comfyui_runtime()\n    add_extra_model_paths()",
+            "def main(unload_models: bool | None = None):\n"
+            "    bootstrap_comfyui_runtime()\n"
+            "    add_extra_model_paths()",
             generated,
         )
         self.assertLess(
-            generated.index("def bootstrap_comfyui_runtime()"), generated.index("def main():")
+            generated.index("def bootstrap_comfyui_runtime()"),
+            generated.index("def main(unload_models: bool | None = None):"),
         )
-        main_section = generated[generated.index("def main():") :]
+        main_section = generated[
+            generated.index("def main(unload_models: bool | None = None):") :
+        ]
         self.assertIn(
-            "def main():\n"
+            "def main(unload_models: bool | None = None):\n"
             "    bootstrap_comfyui_runtime()\n"
             "    add_extra_model_paths()",
             main_section,
@@ -171,7 +180,13 @@ class UpscaleModelLoaderExportTest(unittest.TestCase):
             main_section.index("add_extra_model_paths()"),
             main_section.index("import torch"),
         )
-        self.assertLess(generated.index("import cuda_malloc"), generated.index("import torch"))
+        # cuda_malloc is loaded via _load_module() inside bootstrap_comfyui_runtime()
+        # which runs before 'import torch' in main() — ordering preserved
+        self.assertIn("cuda_malloc", generated)
+        self.assertIn(
+            "    finally:\n        cleanup_comfyui_runtime(unload_models=unload_models)",
+            main_section,
+        )
 
     def test_generated_module_import_does_not_parse_cli_args(self):
         workflow = {
@@ -373,10 +388,13 @@ class UpscaleModelLoaderExportTest(unittest.TestCase):
 
             input_file.write_text(json.dumps(workflow), encoding="utf-8")
 
-            with patch(
-                "comfyui_to_python.get_node_class_mappings",
-                return_value={"LoadImage": LoadImage},
-            ), patch("comfyui_to_python.import_custom_nodes"):
+            with (
+                patch(
+                    "comfyui_to_python.get_node_class_mappings",
+                    return_value={"LoadImage": LoadImage},
+                ),
+                patch("comfyui_to_python.import_custom_nodes"),
+            ):
                 run(
                     input_file=str(input_file),
                     output_file=str(output_file),
@@ -431,18 +449,20 @@ class UpscaleModelLoaderExportTest(unittest.TestCase):
         self.assertIn("# Entrypoint", generated)
         self.assertIn("def build_workflow()", generated)
         self.assertIn("def build_extra_pnginfo()", generated)
-        self.assertIn("def main()", generated)
+        self.assertIn("def main(unload_models: bool | None = None)", generated)
         self.assertIn("bootstrap_comfyui_runtime()", generated)
+        self.assertIn("cleanup_comfyui_runtime(unload_models=unload_models)", generated)
         self.assertNotIn("def initialize_workflow()", generated)
         self.assertNotIn("def run_once(", generated)
         self.assertIn("with torch.inference_mode():", generated)
+        self.assertIn("finally:", generated)
         self.assertIn("for q in range(1):", generated)
         self.assertIn("workflow = build_workflow()", generated)
         self.assertIn("extra_pnginfo = build_extra_pnginfo()", generated)
         self.assertNotIn('workflow = json.loads("', generated)
         self.assertLess(
             generated.index("def build_workflow()"),
-            generated.index("def main()"),
+            generated.index("def main(unload_models: bool | None = None)"),
         )
 
     def test_hidden_metadata_kwargs_follow_function_signature(self):

uv.lock

@@ -57,14 +57,23 @@ wheels = [
 
 [[package]]
 name = "comfyui-to-python-extension"
-version = "2.0.0"
+version = "2.1.0"
 source = { virtual = "." }
 dependencies = [
     { name = "black" },
 ]
 
+[package.optional-dependencies]
+dev = [
+    { name = "ruff" },
+]
+
 [package.metadata]
-requires-dist = [{ name = "black" }]
+requires-dist = [
+    { name = "black" },
+    { name = "ruff", marker = "extra == 'dev'" },
+]
+provides-extras = ["dev"]
 
 [[package]]
 name = "mypy-extensions"
@@ -130,3 +139,28 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/39/2a/cbbf9250020a4a8dd53ba83a46c097b69e5eb49dd14e708f496f548c6612/pytokens-0.4.1-cp314-cp314t-win_amd64.whl", hash = "sha256:584c80c24b078eec1e227079d56dc22ff755e0ba8654d8383b2c549107528918", size = 116287, upload-time = "2026-01-30T01:03:30.912Z" },
     { url = "https://files.pythonhosted.org/packages/c6/78/397db326746f0a342855b81216ae1f0a32965deccfd7c830a2dbc66d2483/pytokens-0.4.1-py3-none-any.whl", hash = "sha256:26cef14744a8385f35d0e095dc8b3a7583f6c953c2e3d269c7f82484bf5ad2de", size = 13729, upload-time = "2026-01-30T01:03:45.029Z" },
 ]
+
+[[package]]
+name = "ruff"
+version = "0.15.12"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/99/43/3291f1cc9106f4c63bdce7a8d0df5047fe8422a75b091c16b5e9355e0b11/ruff-0.15.12.tar.gz", hash = "sha256:ecea26adb26b4232c0c2ca19ccbc0083a68344180bba2a600605538ce51a40a6", size = 4643852, upload-time = "2026-04-24T18:17:14.305Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/c3/6e/e78ffb61d4686f3d96ba3df2c801161843746dcbcbb17a1e927d4829312b/ruff-0.15.12-py3-none-linux_armv6l.whl", hash = "sha256:f86f176e188e94d6bdbc09f09bfd9dc729059ad93d0e7390b5a73efe19f8861c", size = 10640713, upload-time = "2026-04-24T18:17:22.841Z" },
+    { url = "https://files.pythonhosted.org/packages/ae/08/a317bc231fb9e7b93e4ef3089501e51922ff88d6936ce5cf870c4fe55419/ruff-0.15.12-py3-none-macosx_10_12_x86_64.whl", hash = "sha256:e3bcd123364c3770b8e1b7baaf343cc99a35f197c5c6e8af79015c666c423a6c", size = 11069267, upload-time = "2026-04-24T18:17:30.105Z" },
+    { url = "https://files.pythonhosted.org/packages/aa/a4/f828e9718d3dce1f5f11c39c4f65afd32783c8b2aebb2e3d259e492c47bd/ruff-0.15.12-py3-none-macosx_11_0_arm64.whl", hash = "sha256:fe87510d000220aa1ed530d4448a7c696a0cae1213e5ec30e5874287b66557b5", size = 10397182, upload-time = "2026-04-24T18:17:07.177Z" },
+    { url = "https://files.pythonhosted.org/packages/71/e0/3310fc6d1b5e1fdea22bf3b1b807c7e187b581021b0d7d4514cccdb5fb71/ruff-0.15.12-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:84a1630093121375a3e2a95b4a6dc7b59e2b4ee76216e32d81aae550a832d002", size = 10758012, upload-time = "2026-04-24T18:16:55.759Z" },
+    { url = "https://files.pythonhosted.org/packages/11/c1/a606911aee04c324ddaa883ae418f3569792fd3c4a10c50e0dd0a2311e1e/ruff-0.15.12-py3-none-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:fb129f40f114f089ebe0ca56c0d251cf2061b17651d464bb6478dc01e69f11f5", size = 10447479, upload-time = "2026-04-24T18:16:51.677Z" },
+    { url = "https://files.pythonhosted.org/packages/9d/68/4201e8444f0894f21ab4aeeaee68aa4f10b51613514a20d80bd628d57e88/ruff-0.15.12-py3-none-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:b0c862b172d695db7598426b8af465e7e9ac00a3ea2a3630ee67eb82e366aaa6", size = 11234040, upload-time = "2026-04-24T18:17:16.529Z" },
+    { url = "https://files.pythonhosted.org/packages/34/ff/8a6d6cf4ccc23fd67060874e832c18919d1557a0611ebef03fdb01fff11e/ruff-0.15.12-py3-none-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:2849ea9f3484c3aca43a82f484210370319e7170df4dfe4843395ddf6c57bc33", size = 12087377, upload-time = "2026-04-24T18:17:04.944Z" },
+    { url = "https://files.pythonhosted.org/packages/85/f6/c669cf73f5152f623d34e69866a46d5e6185816b19fcd5b6dd8a2d299922/ruff-0.15.12-py3-none-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:9e77c7e51c07fe396826d5969a5b846d9cd4c402535835fb6e21ce8b28fef847", size = 11367784, upload-time = "2026-04-24T18:17:25.409Z" },
+    { url = "https://files.pythonhosted.org/packages/e8/39/c61d193b8a1daaa8977f7dea9e8d8ba866e02ea7b65d32f6861693aa4c12/ruff-0.15.12-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:83b2f4f2f3b1026b5fb449b467d9264bf22067b600f7b6f41fc5958909f449d0", size = 11344088, upload-time = "2026-04-24T18:17:12.258Z" },
+    { url = "https://files.pythonhosted.org/packages/c2/8d/49afab3645e31e12c590acb6d3b5b69d7aab5b81926dbaf7461f9441f37a/ruff-0.15.12-py3-none-manylinux_2_31_riscv64.whl", hash = "sha256:9ba3b8f1afd7e2e43d8943e55f249e13f9682fde09711644a6e7290eb4f3e339", size = 11271770, upload-time = "2026-04-24T18:17:02.457Z" },
+    { url = "https://files.pythonhosted.org/packages/46/06/33f41fe94403e2b755481cdfb9b7ef3e4e0ed031c4581124658d935d52b4/ruff-0.15.12-py3-none-musllinux_1_2_aarch64.whl", hash = "sha256:e852ba9fdc890655e1d78f2df1499efbe0e54126bd405362154a75e2bde159c5", size = 10719355, upload-time = "2026-04-24T18:17:27.648Z" },
+    { url = "https://files.pythonhosted.org/packages/0d/59/18aa4e014debbf559670e4048e39260a85c7fcee84acfd761ac01e7b8d35/ruff-0.15.12-py3-none-musllinux_1_2_armv7l.whl", hash = "sha256:dd8aed930da53780d22fc70bdf84452c843cf64f8cb4eb38984319c24c5cd5fd", size = 10462758, upload-time = "2026-04-24T18:17:32.347Z" },
+    { url = "https://files.pythonhosted.org/packages/25/e7/cc9f16fd0f3b5fddcbd7ec3d6ae30c8f3fde1047f32a4093a98d633c6570/ruff-0.15.12-py3-none-musllinux_1_2_i686.whl", hash = "sha256:01da3988d225628b709493d7dc67c3b9b12c0210016b08690ef9bd27970b262b", size = 10953498, upload-time = "2026-04-24T18:17:20.674Z" },
+    { url = "https://files.pythonhosted.org/packages/72/7a/a9ba7f98c7a575978698f4230c5e8cc54bbc761af34f560818f933dafa0c/ruff-0.15.12-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:9cae0f92bd5700d1213188b31cd3bdd2b315361296d10b96b8e2337d3d11f53e", size = 11447765, upload-time = "2026-04-24T18:17:09.755Z" },
+    { url = "https://files.pythonhosted.org/packages/ea/f9/0ae446942c846b8266059ad8a30702a35afae55f5cdc54c5adf8d7afdc27/ruff-0.15.12-py3-none-win32.whl", hash = "sha256:d0185894e038d7043ba8fd6aee7499ece6462dc0ea9f1e260c7451807c714c20", size = 10657277, upload-time = "2026-04-24T18:17:18.591Z" },
+    { url = "https://files.pythonhosted.org/packages/33/f1/9614e03e1cdcbf9437570b5400ced8a720b5db22b28d8e0f1bda429f660d/ruff-0.15.12-py3-none-win_amd64.whl", hash = "sha256:c87a162d61ab3adca47c03f7f717c68672edec7d1b5499e652331780fe74950d", size = 11837758, upload-time = "2026-04-24T18:17:00.113Z" },
+    { url = "https://files.pythonhosted.org/packages/c0/98/6beb4b351e472e5f4c4613f7c35a5290b8be2497e183825310c4c3a3984b/ruff-0.15.12-py3-none-win_arm64.whl", hash = "sha256:a538f7a82d061cee7be55542aca1d86d1393d55d81d4fcc314370f4340930d4f", size = 11120821, upload-time = "2026-04-24T18:16:57.979Z" },
+]