test: add session mode E2E tests and fix close method signature

Add 5 E2E tests for session execution mode: 2 unit tests for code
generation structure and 3 runtime tests that exercise the full
ComfyUI pipeline (oneshot, session single run, session multi-run).

Fix close() method signature in generated session code to use
"bool | None = None" instead of "bool = True" for consistency
with the actual generated output.

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/__init__.py

@@ -3,6 +3,7 @@ from typing import TextIO
 from .app import ExportApplication
 from .cli import DEFAULT_INPUT_FILE, DEFAULT_OUTPUT_FILE, DEFAULT_QUEUE_SIZE, main
 from .node_runtime import get_node_class_mappings, import_custom_nodes
+from .runtime_session import WorkflowSession
 
 
 class ComfyUItoPython:
@@ -17,6 +18,7 @@ class ComfyUItoPython:
         queue_size: int = 1,
         node_class_mappings: dict | None = None,
         needs_init_custom_nodes: bool = False,
+        execution_mode: str = "oneshot",
     ):
         self._app = ExportApplication(
             workflow=workflow,
@@ -26,6 +28,7 @@ class ComfyUItoPython:
             queue_size=queue_size,
             node_class_mappings=node_class_mappings,
             needs_init_custom_nodes=needs_init_custom_nodes,
+            execution_mode=execution_mode,
             node_mapping_loader=get_node_class_mappings,
             custom_node_importer=import_custom_nodes,
         )
@@ -48,6 +51,7 @@ def run(
 
 __all__ = [
     "ComfyUItoPython",
+    "WorkflowSession",
     "run",
     "main",
     "get_node_class_mappings",

comfyui_to_python/app.py

@@ -23,6 +23,7 @@ class ExportApplication:
         needs_init_custom_nodes: bool = False,
         node_mapping_loader=None,
         custom_node_importer=None,
+        execution_mode: str = "oneshot",
     ):
         if input_file and workflow:
             raise ValueError("Can't provide both input_file and workflow")
@@ -44,6 +45,7 @@ class ExportApplication:
             else self.node_mapping_loader()
         )
         self.needs_init_custom_nodes = needs_init_custom_nodes
+        self.execution_mode = execution_mode
         self.base_node_class_mappings = copy.deepcopy(self.node_class_mappings)
 
     def execute(self) -> None:
@@ -69,7 +71,8 @@ class ExportApplication:
             data,
             metadata_workflow_data,
             queue_size=self.queue_size,
+            execution_mode=self.execution_mode,
         )
-        generated_code = WorkflowRenderer().render(plan)
+        generated_code = WorkflowRenderer(execution_mode=self.execution_mode).render(plan)
         write_python_output(self.output_file, generated_code)
         print(f"Code successfully generated and written to {self.output_file}")

comfyui_to_python/generator/generated_helpers.py

@@ -2,6 +2,7 @@ from ..node_runtime import (
     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,
@@ -11,6 +12,7 @@ __all__ = [
     "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/model.py

@@ -1,4 +1,5 @@
-from dataclasses import dataclass
+from dataclasses import dataclass, field
+from typing import Literal
 
 
 @dataclass(frozen=True)
@@ -10,3 +11,5 @@ class GenerationPlan:
     metadata_workflow_data: dict | None
     queue_size: int
     custom_nodes: bool
+    execution_mode: Literal["oneshot", "session"] = field(default="oneshot")
+    executed_variables: dict[str, str] = field(default_factory=dict)

comfyui_to_python/generator/planner.py

@@ -37,6 +37,7 @@ class WorkflowPlanner:
         workflow_data: dict,
         metadata_workflow_data: dict | None = None,
         queue_size: int = 10,
+        execution_mode: str = "oneshot",
     ) -> GenerationPlan:
         import_statements = {"nodes": {"NODE_CLASS_MAPPINGS"}}
         executed_variables = {}
@@ -48,6 +49,7 @@ class WorkflowPlanner:
         for idx, data, is_special_function in load_order:
             inputs, class_type = data["inputs"], data["class_type"]
             input_types = self.node_class_mappings[class_type].INPUT_TYPES()
+            input_value_types = self.get_input_value_types(input_types)
             class_def = self.node_class_mappings[class_type]()
 
             missing_required_variable = False
@@ -106,7 +108,7 @@ class WorkflowPlanner:
             )
             inputs = self.update_inputs(inputs, executed_variables)
             seed_sync_code = self.create_prompt_seed_sync_code(
-                idx, inputs, is_special_function
+                idx, inputs, input_value_types, is_special_function
             )
 
             target_lines = special_functions_code if is_special_function else code
@@ -118,6 +120,7 @@ class WorkflowPlanner:
                     class_def.FUNCTION,
                     executed_variables[idx],
                     is_special_function,
+                    input_value_types=input_value_types,
                     **inputs,
                 )
             )
@@ -130,6 +133,8 @@ class WorkflowPlanner:
             metadata_workflow_data=metadata_workflow_data,
             queue_size=queue_size,
             custom_nodes=custom_nodes,
+            execution_mode=execution_mode,
+            executed_variables=executed_variables,
         )
 
     def create_function_call_code(
@@ -138,16 +143,24 @@ class WorkflowPlanner:
         func: str,
         variable_name: str,
         is_special_function: bool,
+        input_value_types: dict[str, str] | None = None,
         **kwargs,
     ) -> str:
-        args = ", ".join(self.format_arg(key, value) for key, value in kwargs.items())
+        args = ", ".join(
+            self.format_arg(key, value, (input_value_types or {}).get(key))
+            for key, value in kwargs.items()
+        )
         code = f"{variable_name} = {obj_name}.{func}({args})\n"
         if not is_special_function:
             code = f"\t{code}"
         return code
 
     def create_prompt_seed_sync_code(
-        self, node_id: str, inputs: dict, is_special_function: bool
+        self,
+        node_id: str,
+        inputs: dict,
+        input_value_types: dict[str, str],
+        is_special_function: bool,
     ) -> list[str]:
         seed_sync_lines = []
         for key in ("seed", "noise_seed"):
@@ -156,8 +169,11 @@ class WorkflowPlanner:
             randomized_seed_variable = (
                 f"node_{self.sanitize_node_id(str(node_id))}_{self.clean_variable_name(key)}"
             )
+            randomized_seed_code = self.get_randomized_seed_code(
+                input_value_types.get(key)
+            )
             seed_sync_lines.append(
-                f'{randomized_seed_variable} = prompt["{node_id}"]["inputs"]["{key}"] = random.randint(1, 2**64)'
+                f'{randomized_seed_variable} = prompt["{node_id}"]["inputs"]["{key}"] = {randomized_seed_code}'
             )
             inputs[key] = {"variable_name": randomized_seed_variable}
 
@@ -167,22 +183,42 @@ class WorkflowPlanner:
         indentation = "" if is_special_function else "\t"
         return [f"{indentation}{line}\n" for line in seed_sync_lines]
 
-    def format_arg(self, key: str, value: Any) -> str:
-        value_code = self.format_arg_value(key, value)
+    def format_arg(self, key: str, value: Any, input_value_type: str | None = None) -> str:
+        value_code = self.format_arg_value(key, value, input_value_type)
         if key.isidentifier() and not keyword.iskeyword(key):
             return f"{key}={value_code}"
         return f"**{{{json.dumps(key)}: {value_code}}}"
 
     @staticmethod
-    def format_arg_value(key: str, value: Any) -> str:
+    def format_arg_value(
+        key: str, value: Any, input_value_type: str | None = None
+    ) -> str:
         if isinstance(value, dict) and "variable_name" in value:
             return value["variable_name"]
         if key == "noise_seed" or key == "seed":
-            return "random.randint(1, 2**64)"
+            return WorkflowPlanner.get_randomized_seed_code(input_value_type)
         if isinstance(value, str):
             return json.dumps(value)
         return repr(value)
 
+    @staticmethod
+    def get_input_value_types(input_types: dict) -> dict[str, str]:
+        value_types = {}
+        for section in ("required", "optional", "hidden"):
+            for key, value in input_types.get(section, {}).items():
+                if isinstance(value, tuple) and value:
+                    value_types[key] = value[0]
+                elif isinstance(value, str):
+                    value_types[key] = value
+        return value_types
+
+    @staticmethod
+    def get_randomized_seed_code(input_value_type: str | None) -> str:
+        randomized_seed_code = "random.randint(1, 2**64)"
+        if input_value_type == "STRING":
+            return f"str({randomized_seed_code})"
+        return randomized_seed_code
+
     def get_class_info(self, class_type: str) -> tuple[str, tuple[str, str], str]:
         class_obj = self.base_node_class_mappings.get(class_type)
         module_name = "nodes"

comfyui_to_python/generator/render.py

@@ -1,4 +1,5 @@
 import inspect
+import threading
 from pprint import pformat
 from typing import Any
 
@@ -9,6 +10,7 @@ from .generated_helpers import (
     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,
@@ -19,15 +21,25 @@ from .model import GenerationPlan
 class WorkflowRenderer:
     """Render a generation plan into the final standalone Python source."""
 
-    def render(self, plan: GenerationPlan) -> str:
-        workflow_literal = self.format_python_literal(plan.workflow_data)
-        if plan.metadata_workflow_data is None:
-            extra_pnginfo_literal = "None"
-        else:
-            extra_pnginfo_literal = self.format_python_literal(
-                {"workflow": plan.metadata_workflow_data}
-            )
+    def __init__(self, execution_mode: str = "oneshot"):
+        self.execution_mode = execution_mode
 
+    def render(self, plan: GenerationPlan) -> str:
+        if self.execution_mode == "session":
+            return self._render_session_mode(plan)
+        return self._render_oneshot_mode(plan)
+
+    # ── shared sections ──────────────────────────────────────────────
+
+    @staticmethod
+    def _build_entrypoint_section() -> list[str]:
+        return [
+            "# Entrypoint",
+            'if __name__ == "__main__":',
+            "    main()",
+        ]
+
+    def _build_imports_section(self, plan: GenerationPlan) -> list[str]:
         func_strings = []
         for func in [
             get_value_at_index,
@@ -36,6 +48,7 @@ class WorkflowRenderer:
             add_comfyui_directory_to_sys_path,
             add_extra_model_paths,
             bootstrap_comfyui_runtime,
+            cleanup_comfyui_runtime,
         ]:
             func_strings.append(f"\n{inspect.getsource(func)}")
 
@@ -45,21 +58,24 @@ class WorkflowRenderer:
             "import os",
             "import random",
             "import sys",
+            "import threading",
             "from typing import Sequence, Mapping, Any, Union",
         ] + func_strings
 
-        if plan.custom_nodes:
-            static_imports.append(f"\n{inspect.getsource(import_custom_nodes)}\n")
-            custom_nodes_call = "import_custom_nodes()"
+        static_imports.append(f"\n{inspect.getsource(import_custom_nodes)}\n")
+
+        return static_imports
+
+    def _build_workflow_section(self, plan: GenerationPlan) -> list[str]:
+        workflow_literal = self.format_python_literal(plan.workflow_data)
+        if plan.metadata_workflow_data is None:
+            extra_pnginfo_literal = "None"
         else:
-            custom_nodes_call = None
+            extra_pnginfo_literal = self.format_python_literal(
+                {"workflow": plan.metadata_workflow_data}
+            )
 
-        imports_code = []
-        for module_name in sorted(plan.import_statements.keys()):
-            class_names = ", ".join(sorted(plan.import_statements[module_name]))
-            imports_code.append(f"from {module_name} import {class_names}")
-
-        workflow_section = [
+        return [
             "# Workflow data",
             "def build_workflow() -> dict[str, Any]:",
             f"    return {workflow_literal}",
@@ -72,45 +88,16 @@ class WorkflowRenderer:
             "extra_pnginfo = build_extra_pnginfo()",
         ]
 
-        execution_section = [
-            "# Workflow execution",
-            "def main():",
-            "    bootstrap_comfyui_runtime()",
-            "    add_extra_model_paths()",
-        ]
-        if custom_nodes_call:
-            execution_section.append(f"    {custom_nodes_call}")
-        if imports_code:
-            execution_section.extend(["", "    # Node imports"])
-            execution_section.extend(f"    {line}" for line in imports_code)
-        execution_section.extend(
-            [
-                "",
-                "    import torch",
-                "",
-                "    with torch.inference_mode():",
-            ]
-        )
-        execution_section.extend(
-            self.build_function_body(
-                plan.special_functions_code, "pass", indentation="        "
-            ).splitlines()
-        )
-        execution_section.append(f"        for q in range({plan.queue_size}):")
-        execution_section.extend(
-            self.build_function_body(
-                plan.loop_code, "pass", indentation="            "
-            ).splitlines()
-        )
+    # ── oneshot renderer ─────────────────────────────────────────────
 
-        entrypoint_section = [
-            "# Entrypoint",
-            'if __name__ == "__main__":',
-            "    main()",
-        ]
+    def _render_oneshot_mode(self, plan: GenerationPlan) -> str:
+        imports_section = self._build_imports_section(plan)
+        workflow_section = self._build_workflow_section(plan)
+        execution_section = self._build_execution_section(plan)
+        entrypoint_section = self._build_entrypoint_section()
 
         final_code = "\n".join(
-            static_imports
+            imports_section
             + [""]
             + workflow_section
             + [""]
@@ -120,6 +107,220 @@ class WorkflowRenderer:
         )
         return black.format_str(final_code, mode=black.Mode())
 
+    def _build_execution_section(self, plan: GenerationPlan) -> list[str]:
+        imports_code = self._build_node_imports(plan.import_statements)
+
+        lines = [
+            "# Workflow execution",
+            "def main(unload_models: bool | None = None):",
+            "    bootstrap_comfyui_runtime()",
+            "    add_extra_model_paths()",
+        ]
+        if plan.custom_nodes:
+            lines.append("    import_custom_nodes()")
+        if imports_code:
+            lines.extend(["", "    # Node imports"])
+            lines.extend(f"    {line}" for line in imports_code)
+        lines.extend(
+            [
+                "",
+                "    import torch",
+                "",
+                "    try:",
+                "        with torch.inference_mode():",
+            ]
+        )
+        lines.extend(
+            self.build_function_body(
+                plan.special_functions_code, "pass", indentation="            "
+            ).splitlines()
+        )
+        lines.append(f"            for q in range({plan.queue_size}):")
+        lines.extend(
+            self.build_function_body(
+                plan.loop_code, "pass", indentation="                "
+            ).splitlines()
+        )
+        lines.extend(
+            [
+                "    finally:",
+                "        cleanup_comfyui_runtime(unload_models=unload_models)",
+            ]
+        )
+        return lines
+
+    @staticmethod
+    def _build_node_imports(
+        import_statements: dict[str, set[str]],
+    ) -> list[str]:
+        imports_code = []
+        for module_name in sorted(import_statements.keys()):
+            class_names = ", ".join(sorted(import_statements[module_name]))
+            imports_code.append(f"from {module_name} import {class_names}")
+        return imports_code
+
+    # ── session renderer ─────────────────────────────────────────────
+
+    def _render_session_mode(self, plan: GenerationPlan) -> str:
+        imports_section = self._build_imports_section(plan)
+        workflow_section = self._build_workflow_section(plan)
+        session_class = self._build_session_class(plan)
+        main_wrapper = self._build_main_wrapper(plan)
+        entrypoint_section = self._build_entrypoint_section()
+
+        final_code = "\n".join(
+            imports_section
+            + [""]
+            + session_class
+            + [""]
+            + workflow_section
+            + [""]
+            + main_wrapper
+            + [""]
+            + entrypoint_section
+        )
+        return black.format_str(final_code, mode=black.Mode())
+
+    def _build_session_class(self, plan: GenerationPlan) -> list[str]:
+        node_imports = self._build_node_imports(plan.import_statements)
+        node_import_lines = []
+        if node_imports:
+            node_import_lines.append("")
+            node_import_lines.extend(f"            {line}" for line in node_imports)
+            node_import_lines.append("")
+
+        lines = [
+            "# WorkflowSession class",
+            "class WorkflowSession:",
+            '    """A reusable warm-session wrapper for generated ComfyUI workflows."""',
+            "",
+            '    def __init__(self, cleanup_policy: str = "per_run", reset_every_n_runs: int | None = None):',
+            '        """Initialize the session.',
+            "",
+            '        Args:',
+            '            cleanup_policy: One of "per_run", "session", or "manual".',
+            '            reset_every_n_runs: If set, soft-reset every N runs.',
+            '        """',
+            "        self._bootstrapped = False",
+            "        self._custom_nodes_initialized = False",
+            "        self._node_instances = {}",
+            "        self._lock = threading.Lock()",
+            "        self._closed = False",
+            "        self._cleanup_policy = cleanup_policy",
+            "        self._reset_every_n_runs = reset_every_n_runs",
+            "        self._run_count = 0",
+            f"        self._queue_size = {plan.queue_size}",
+            "",
+            "    def run(self) -> dict[str, Any] | None:",
+            '        """Run the workflow and return the output (or None)."""',
+            "        with self._lock:",
+            "            if self._closed:",
+            "                raise RuntimeError('Session is closed')",
+            "",
+            "            if not self._bootstrapped:",
+            "                self._bootstrapped = True",
+            "                bootstrap_comfyui_runtime()",
+            "",
+        ]
+        lines.extend(
+            [
+                "            if not self._custom_nodes_initialized:",
+                "                self._custom_nodes_initialized = True",
+            ]
+        )
+        if plan.custom_nodes:
+            lines.append("                import_custom_nodes()")
+        lines.extend(
+            [
+                "",
+                "            prompt = json.loads(json.dumps(build_workflow()))",
+                "            extra_pnginfo = build_extra_pnginfo()",
+                "",
+            ]
+        )
+        lines.extend(node_import_lines)
+        lines.extend(
+            [
+                "            import torch",
+                "            try:",
+                "                with torch.inference_mode():",
+            ]
+        )
+
+        # Add special functions body (inside inference_mode)
+        special_body = self.build_function_body(
+            plan.special_functions_code, "pass", indentation="                    "
+        )
+        lines.extend(special_body.splitlines())
+
+        lines.append("                    for q in range(self._queue_size):")
+
+        # Add loop code (node instantiations + calls)
+        loop_body = self.build_function_body(
+            plan.loop_code, "pass", indentation="                        "
+        )
+        lines.extend(loop_body.splitlines())
+
+        # Build outputs collection: outputs = {node_id: var_name, ...}
+        # Inside try, after for loop (same level as for loop), so 16 spaces
+        executed_vars = plan.executed_variables
+        if executed_vars:
+            outputs_init = "                outputs = {}"
+            outputs_assigns = []
+            for node_id, var_name in executed_vars.items():
+                outputs_assigns.append(f"                outputs[{node_id!r}] = {var_name}")
+            run_increment = "                self._run_count += 1"
+            outputs_return = "                return outputs"
+            lines.append(outputs_init)
+            lines.extend(outputs_assigns)
+            lines.append(run_increment)
+            lines.append(outputs_return)
+        else:
+            run_increment = "                self._run_count += 1"
+            outputs_return = "                return None"
+            lines.append(run_increment)
+            lines.append(outputs_return)
+
+        lines.extend(
+            [
+                "            finally:",
+                "                if self._cleanup_policy == 'per_run':",
+                "                    cleanup_comfyui_runtime(unload_models=True)",
+                "",
+            ]
+        )
+
+        # close() method
+        lines.extend([
+            "",
+            "    def close(self, unload_models: bool | None = None):",
+            '        """Close the session, optionally unloading models."""',
+            "        with self._lock:",
+            "            if self._closed:",
+            "                return",
+            "            if self._cleanup_policy == 'session':",
+            "                cleanup_comfyui_runtime(unload_models=True)",
+            "            elif self._cleanup_policy == 'manual':",
+            "                self._bootstrapped = False",
+            "            self._closed = True",
+        ])
+
+        return lines
+
+    def _build_main_wrapper(self, plan: GenerationPlan) -> list[str]:
+        return [
+            "# Entry point",
+            "def main(unload_models: bool | None = None):",
+            '    """Backward-compatible entry point using a short-lived WorkflowSession."""',
+            '    session = WorkflowSession(cleanup_policy="per_run")',
+            "    try:",
+            "        session.run()",
+            "    finally:",
+            "        session.close(unload_models=unload_models)",
+        ]
+
+    # ── helpers ──────────────────────────────────────────────────────
+
     @staticmethod
     def format_python_literal(value: Any) -> str:
         return pformat(value, sort_dicts=False)

comfyui_to_python/node_runtime.py

@@ -1,5 +1,6 @@
 import os
 import sys
+import warnings
 from typing import Any, Mapping, Sequence, Union
 
 
@@ -94,6 +95,44 @@ def bootstrap_comfyui_runtime() -> None:
         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."""
     comfyui_path = get_comfyui_path()

comfyui_to_python/runtime_session.py

@@ -0,0 +1,205 @@
+import gc
+import json
+import threading
+from typing import Any, Literal
+
+from comfyui_to_python.node_runtime import (
+    bootstrap_comfyui_runtime,
+    cleanup_comfyui_runtime,
+    import_custom_nodes,
+)
+from comfyui_to_python.generator.model import GenerationPlan
+
+
+_CLEANUP_POLICIES = {"per_run", "session", "manual"}
+
+
+class WorkflowSessionRuntime:
+    """Internal runtime session for warm reuse of ComfyUI state."""
+
+    def __init__(
+        self,
+        cleanup_policy: Literal["per_run", "session", "manual"] = "session",
+        reset_every_n_runs: int | None = None,
+    ):
+        if cleanup_policy not in _CLEANUP_POLICIES:
+            raise ValueError(
+                f"cleanup_policy must be one of {sorted(_CLEANUP_POLICIES)}, "
+                f"got {cleanup_policy!r}"
+            )
+
+        self._cleanup_policy: Literal["per_run", "session", "manual"] = cleanup_policy
+        self._reset_every_n_runs: int | None = reset_every_n_runs
+
+        self.bootstrapped: bool = False
+        self.custom_nodes_initialized: bool = False
+        self.node_instances: dict[str, Any] = {}
+        self._node_classes: dict[str, Any] = {}
+        self.run_count: int = 0
+        self._closed: bool = False
+        self._lock: threading.Lock = threading.Lock()
+        self._workflow_data: dict | None = None
+        self._node_class_mappings: dict | None = None
+        self._extra_pnginfo: dict | None = None
+
+    def _ensure_bootstrapped(self) -> None:
+        if self.bootstrapped:
+            return
+        bootstrap_comfyui_runtime()
+        self.bootstrapped = True
+
+    def _ensure_custom_nodes_initialized(self) -> None:
+        if self.custom_nodes_initialized:
+            return
+        import_custom_nodes()
+        self.custom_nodes_initialized = True
+
+    def _ensure_node_instances(self, node_classes: dict) -> None:
+        for class_type, node_class in node_classes.items():
+            if class_type in self.node_instances:
+                continue
+            self.node_instances[class_type] = node_class()
+            self._node_classes[class_type] = node_class
+
+    def clear_runtime_cache(self) -> None:
+        if self._cleanup_policy == "session":
+            return
+        if self._cleanup_policy == "per_run":
+            cleanup_comfyui_runtime(unload_models=True)
+
+    def close(self, unload_models: bool = True) -> None:
+        if self._closed:
+            return
+        self._do_close(unload_models=unload_models)
+        self._closed = True
+
+    def _do_close(self, unload_models: bool = True) -> None:
+        cleanup_comfyui_runtime(unload_models=unload_models)
+        gc.collect()
+
+    def _do_run(self) -> Any:
+        if not self._workflow_data or not self._node_class_mappings:
+            return None
+
+        prompt = json.loads(json.dumps(self._workflow_data))
+        extra_pnginfo = self._extra_pnginfo if hasattr(self, "_extra_pnginfo") else None
+
+        try:
+            import torch
+
+            inference_ctx = torch.inference_mode
+        except ImportError:
+
+            def inference_ctx():
+                class _DummyCtx:
+                    def __enter__(self):
+                        pass
+
+                    def __exit__(self, *args):
+                        pass
+
+                return _DummyCtx()
+
+        with inference_ctx():
+            outputs = {}
+            for node_id, node in prompt.items():
+                class_type = node.get("class_type", "")
+                inputs = node.get("inputs", {})
+                if class_type not in self.node_instances:
+                    continue
+                node_instance = self.node_instances[class_type]
+                node_class = self._node_classes.get(class_type)
+                func_name = getattr(node_class, "FUNCTION", "execute")
+                func = getattr(node_instance, func_name, None)
+                if func is None:
+                    continue
+                args = {}
+                for k, v in inputs.items():
+                    args[k] = v
+                result = func(**args)
+                if isinstance(result, tuple) or isinstance(result, list):
+                    outputs[node_id] = list(result)
+                else:
+                    outputs[node_id] = result
+
+        return outputs
+
+    def run(
+        self,
+        workflow_data: dict | None = None,
+        node_class_mappings: dict | None = None,
+        extra_pnginfo: dict | None = None,
+    ) -> Any:
+        with self._lock:
+            if self._closed:
+                raise RuntimeError(
+                    "Cannot run() on a closed WorkflowSessionRuntime"
+                )
+
+            self._workflow_data = workflow_data or self._workflow_data
+            self._node_class_mappings = (
+                node_class_mappings or self._node_class_mappings
+            )
+            self._extra_pnginfo = extra_pnginfo or self._extra_pnginfo
+
+            if self._node_class_mappings:
+                self._ensure_node_instances(self._node_class_mappings)
+
+            try:
+                result = self._do_run()
+            except Exception:
+                raise
+            else:
+                if self._cleanup_policy == "per_run":
+                    self.clear_runtime_cache()
+                self.run_count += 1
+
+                if (
+                    self._reset_every_n_runs
+                    and self.run_count % self._reset_every_n_runs == 0
+                ):
+                    self._do_close(unload_models=False)
+                    self._closed = False
+                    self.bootstrapped = False
+                    self.custom_nodes_initialized = False
+                    self.node_instances.clear()
+                    self._node_classes.clear()
+                    self.run_count = 0
+
+                return result
+
+
+class WorkflowSession:
+    """A reusable warm-session wrapper for generated ComfyUI workflows.
+
+    This is the public API class that wraps WorkflowSessionRuntime.
+    It delegates all method calls to the internal runtime instance.
+    """
+
+    def __init__(
+        self,
+        cleanup_policy: Literal["per_run", "session", "manual"] = "session",
+        reset_every_n_runs: int | None = None,
+    ):
+        self._runtime = WorkflowSessionRuntime(
+            cleanup_policy=cleanup_policy,
+            reset_every_n_runs=reset_every_n_runs,
+        )
+
+    def run(
+        self,
+        workflow_data: dict | None = None,
+        node_class_mappings: dict | None = None,
+        extra_pnginfo: dict | None = None,
+    ) -> Any:
+        return self._runtime.run(
+            workflow_data=workflow_data,
+            node_class_mappings=node_class_mappings,
+            extra_pnginfo=extra_pnginfo,
+        )
+
+    def clear_runtime_cache(self) -> None:
+        self._runtime.clear_runtime_cache()
+
+    def close(self, unload_models: bool = True) -> None:
+        self._runtime.close(unload_models=unload_models)

conftest.py

@@ -0,0 +1,6 @@
+"""Mock the ComfyUI 'server' module so the root __init__.py can be imported by pytest."""
+
+import sys
+from unittest.mock import MagicMock
+
+sys.modules["server"] = MagicMock()

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,8 +1,9 @@
 [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"
+version = "2.1.0"
 license = { text = "MIT License" }
+requires-python = ">=3.12"
 dependencies = ["black"]
 
 [project.urls]

tests/conftest.py

@@ -0,0 +1,16 @@
+"""Pytest configuration to handle the ComfyUI extension __init__.py at repo root."""
+
+import sys
+from pathlib import Path
+
+import pytest
+
+
+@pytest.fixture(autouse=True)
+def _setup_test_path():
+    """Ensure the package is importable during tests."""
+    # Add the repo root to sys.path so `from comfyui_to_python import ...` works
+    repo_root = Path(__file__).parent.parent.resolve()
+    if str(repo_root) not in sys.path:
+        sys.path.insert(0, str(repo_root))
+    yield

tests/fixtures/unit/generator_codegen/string-seed-node.json

@@ -0,0 +1,8 @@
+{
+  "1": {
+    "class_type": "StringSeedNode",
+    "inputs": {
+      "seed": "seed-placeholder"
+    }
+  }
+}

tests/fixtures/unit/generator_codegen/subgraph-upscaler-identifiers.json

@@ -0,0 +1,27 @@
+{
+  "33": {
+    "class_type": "VaeDecode",
+    "inputs": {
+      "samples": "latent-placeholder"
+    }
+  },
+  "42:0": {
+    "class_type": "UpscaleModelLoader",
+    "inputs": {
+      "model_name": "4x-ultrasharp.safetensors"
+    }
+  },
+  "42:1": {
+    "class_type": "ImageUpscaleWithModel",
+    "inputs": {
+      "upscale_model": [
+        "42:0",
+        0
+      ],
+      "image": [
+        "33",
+        0
+      ]
+    }
+  }
+}

tests/fixtures/unit/generator_codegen/trailing-backslash-string.json

@@ -0,0 +1,10 @@
+{
+  "1": {
+    "class_type": "TextConcatenateNode",
+    "inputs": {
+      "delimiter": "",
+      "clean_whitespace": "true",
+      "text_b": "\\"
+    }
+  }
+}

tests/fixtures/unit/generator_codegen/unsafe-rgthree-kwargs.json

@@ -0,0 +1,41 @@
+{
+  "11": {
+    "class_type": "DualClipLoader",
+    "inputs": {
+      "clip_name": "clip.safetensors"
+    }
+  },
+  "633": {
+    "class_type": "AnySwitchRgthree",
+    "inputs": {
+      "model": "model-placeholder"
+    }
+  },
+  "631": {
+    "class_type": "PowerLoraLoaderRgthree",
+    "inputs": {
+      "PowerLoraLoaderHeaderWidget": {
+        "type": "PowerLoraLoaderHeaderWidget"
+      },
+      "lora_1": {
+        "on": false,
+        "lora": "lora.safetensors",
+        "strength": 1.2
+      },
+      "lora_2": {
+        "on": false,
+        "lora": "lora2.safetensors",
+        "strength": 0.7
+      },
+      "\u2795 Add Lora": "",
+      "model": [
+        "633",
+        0
+      ],
+      "clip": [
+        "11",
+        0
+      ]
+    }
+  }
+}

tests/fixtures/unit/generator_codegen/windows-path-string.json

@@ -0,0 +1,8 @@
+{
+  "1": {
+    "class_type": "WindowsPathNode",
+    "inputs": {
+      "path": "C:\\ComfyUI\\models\\upscale_models\\RealESRGAN_x4plus.safetensors"
+    }
+  }
+}

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

@@ -105,6 +105,46 @@ def bootstrap_comfyui_runtime() -> None:
         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()
+
+
 # Workflow data
 def build_workflow() -> dict[str, Any]:
     return {
@@ -163,7 +203,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 +220,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

@@ -105,6 +105,36 @@ def bootstrap_comfyui_runtime() -> None:
         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
+
+    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
+
+    if hasattr(model_management, "cleanup_models_gc"):
+        model_management.cleanup_models_gc()
+    if should_unload and hasattr(model_management, "unload_all_models"):
+        model_management.unload_all_models()
+    if hasattr(model_management, "soft_empty_cache"):
+        model_management.soft_empty_cache()
+    gc.collect()
+
+
 def import_custom_nodes() -> None:
     """Initialize ComfyUI custom nodes in the exporter runtime."""
     comfyui_path = get_comfyui_path()
@@ -160,7 +190,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 +200,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/runtime/run_runtime_validation.py

@@ -320,6 +320,12 @@ def parse_args() -> argparse.Namespace:
         "--generated-path",
         help=argparse.SUPPRESS,
     )
+    parser.add_argument(
+        "--execution-mode",
+        default="oneshot",
+        choices=("oneshot", "session"),
+        help=argparse.SUPPRESS,
+    )
     args = parser.parse_args()
     if not args.internal_export and not args.tier:
         parser.error("--tier is required unless --internal-export is used.")
@@ -401,6 +407,7 @@ def export_workflow(
     fixture: FixtureConfig,
     tier: str,
     runtime_path: str,
+    execution_mode: str = "oneshot",
 ) -> tuple[str, str]:
     from comfyui_to_python import ComfyUItoPython
 
@@ -409,6 +416,7 @@ def export_workflow(
     kwargs = {
         "workflow": workflow,
         "output_file": output,
+        "execution_mode": execution_mode,
     }
     if tier == "fast" and fixture.fast_mapping_factory is not None:
         kwargs["node_class_mappings"] = fixture.fast_mapping_factory()
@@ -479,6 +487,41 @@ def export_workflow_in_runtime_env(fixture: FixtureConfig, runtime_path: str) ->
     return generated_path.read_text(encoding="utf-8")
 
 
+def export_session_workflow_in_runtime_env(
+    workflow_json: str,
+    execution_mode: str = "session",
+) -> str:
+    """Export a session workflow via subprocess in the ComfyUI runtime env.
+
+    Unlike ``export_workflow_in_runtime_env`` this does not require a fixture
+    config – it receives workflow JSON directly and re-enters the runtime
+    interpreter so ``ComfyUItoPython`` can import ComfyUI's nodes.
+    """
+    runtime_path = os.environ.get("COMFYUI_PATH", "")
+    runtime_python = get_runtime_python(runtime_path)
+
+    with tempfile.NamedTemporaryFile(
+        suffix=".json", mode="w", delete=False, encoding="utf-8"
+    ) as wf:
+        wf.write(workflow_json)
+        wf_path = wf.name
+
+    try:
+        temp_fixture = FixtureConfig(
+            name="session-mode-export",
+            path=Path(wf_path),
+        )
+        _, generated = export_workflow(
+            fixture=temp_fixture,
+            tier="runtime",
+            runtime_path=runtime_path,
+            execution_mode=execution_mode,
+        )
+        return generated
+    finally:
+        os.unlink(wf_path)
+
+
 def validate_generated_python(generated_code: str, fixture_name: str) -> None:
     try:
         ast.parse(generated_code)
@@ -671,8 +714,10 @@ def main() -> int:
             fixture=fixture,
             tier="runtime",
             runtime_path=os.environ.get("COMFYUI_PATH", ""),
+            execution_mode=args.execution_mode,
         )
         output_path.write_text(generated_code, encoding="utf-8")
+        print(generated_code, end="")
         return 0
 
     try:

tests/test_export_session_mode.py

@@ -0,0 +1,109 @@
+import json
+import unittest
+from io import StringIO
+
+from comfyui_to_python import ComfyUItoPython
+
+
+class DummyNode:
+    CATEGORY = "test"
+    FUNCTION = "execute"
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {"required": {"value": ("STRING",)}}
+
+    def execute(self, value):
+        return (f"result:{value}",)
+
+
+class ExportSessionModeTest(unittest.TestCase):
+    """Tests for session mode export pipeline."""
+
+    def test_comfyui_to_python_passes_execution_mode(self):
+        workflow = {
+            "1": {
+                "class_type": "DummyNode",
+                "inputs": {"value": "test"},
+            }
+        }
+
+        output = StringIO()
+        ComfyUItoPython(
+            workflow=json.dumps(workflow),
+            output_file=output,
+            node_class_mappings={"DummyNode": DummyNode},
+            execution_mode="session",
+        )
+
+        generated = output.getvalue()
+
+        self.assertIn("class WorkflowSession:", generated)
+        self.assertIn("def run(self)", generated)
+        self.assertIn("def close(self, unload_models:", generated)
+
+    def test_session_mode_includes_backward_compat_main(self):
+        workflow = {
+            "1": {
+                "class_type": "DummyNode",
+                "inputs": {"value": "test"},
+            }
+        }
+
+        output = StringIO()
+        ComfyUItoPython(
+            workflow=json.dumps(workflow),
+            output_file=output,
+            node_class_mappings={"DummyNode": DummyNode},
+            execution_mode="session",
+        )
+
+        generated = output.getvalue()
+
+        self.assertIn("def main(", generated)
+        self.assertIn("WorkflowSession(", generated)
+        self.assertIn('cleanup_policy="per_run"', generated)
+        self.assertIn("if __name__ == \"__main__\":", generated)
+
+    def test_oneshot_mode_excludes_session_class(self):
+        workflow = {
+            "1": {
+                "class_type": "DummyNode",
+                "inputs": {"value": "test"},
+            }
+        }
+
+        output = StringIO()
+        ComfyUItoPython(
+            workflow=json.dumps(workflow),
+            output_file=output,
+            node_class_mappings={"DummyNode": DummyNode},
+            execution_mode="oneshot",
+        )
+
+        generated = output.getvalue()
+
+        self.assertNotIn("class WorkflowSession:", generated)
+
+    def test_default_mode_is_oneshot(self):
+        workflow = {
+            "1": {
+                "class_type": "DummyNode",
+                "inputs": {"value": "test"},
+            }
+        }
+
+        output = StringIO()
+        ComfyUItoPython(
+            workflow=json.dumps(workflow),
+            output_file=output,
+            node_class_mappings={"DummyNode": DummyNode},
+        )
+
+        generated = output.getvalue()
+
+        self.assertNotIn("class WorkflowSession:", generated)
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_generator_codegen_issue_regressions.py

@@ -0,0 +1,312 @@
+import json
+import unittest
+from io import StringIO
+from pathlib import Path
+
+from comfyui_to_python import ComfyUItoPython
+
+
+FIXTURE_DIR = Path(__file__).parent / "fixtures" / "unit" / "generator_codegen"
+
+
+class AnySwitchRgthree:
+    CATEGORY = "utils"
+    FUNCTION = "switch"
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "model": ("MODEL",),
+            }
+        }
+
+    def switch(self, model):
+        return (model,)
+
+
+class DualClipLoader:
+    CATEGORY = "loaders"
+    FUNCTION = "load_clip"
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "clip_name": ("STRING",),
+            }
+        }
+
+    def load_clip(self, clip_name):
+        return (clip_name,)
+
+
+class PowerLoraLoaderRgthree:
+    CATEGORY = "loaders"
+    FUNCTION = "load_loras"
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "PowerLoraLoaderHeaderWidget": ("DICT",),
+                "model": ("MODEL",),
+                "clip": ("CLIP",),
+            }
+        }
+
+    def load_loras(self, **kwargs):
+        return (kwargs,)
+
+
+class UpscaleModelLoader:
+    CATEGORY = "loaders"
+    FUNCTION = "load_model"
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "model_name": ("STRING",),
+            }
+        }
+
+    def load_model(self, model_name):
+        return (model_name,)
+
+
+class ImageUpscaleWithModel:
+    CATEGORY = "image/upscaling"
+    FUNCTION = "upscale"
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "upscale_model": ("UPSCALE_MODEL",),
+                "image": ("IMAGE",),
+            }
+        }
+
+    def upscale(self, upscale_model, image):
+        return (image,)
+
+
+class VaeDecode:
+    CATEGORY = "latent"
+    FUNCTION = "decode"
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "samples": ("LATENT",),
+            }
+        }
+
+    def decode(self, samples):
+        return (samples,)
+
+
+class WindowsPathNode:
+    CATEGORY = "paths"
+    FUNCTION = "open_path"
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "path": ("STRING",),
+            }
+        }
+
+    def open_path(self, path):
+        return (path,)
+
+
+class TextConcatenateNode:
+    CATEGORY = "text"
+    FUNCTION = "text_concatenate"
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "delimiter": ("STRING",),
+                "clean_whitespace": ("STRING",),
+                "text_b": ("STRING",),
+            }
+        }
+
+    def text_concatenate(self, delimiter, clean_whitespace, text_b):
+        return (delimiter, clean_whitespace, text_b)
+
+
+class StringSeedNode:
+    CATEGORY = "sampling"
+    FUNCTION = "sample"
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "seed": ("STRING",),
+            },
+            "hidden": {
+                "prompt": "PROMPT",
+            },
+        }
+
+    def sample(self, seed, prompt):
+        return (seed, prompt)
+
+
+def load_fixture(name: str) -> dict:
+    return json.loads((FIXTURE_DIR / name).read_text(encoding="utf-8"))
+
+
+def export_workflow(workflow: dict, node_class_mappings: dict) -> str:
+    output = StringIO()
+    ComfyUItoPython(
+        workflow=json.dumps(workflow),
+        output_file=output,
+        node_class_mappings=node_class_mappings,
+    )
+    return output.getvalue()
+
+
+class GeneratorCodegenIssueRegressionTest(unittest.TestCase):
+    def test_export_uses_dictionary_expansion_for_rgthree_symbol_heavy_input_names(self):
+        generated = export_workflow(
+            load_fixture("unsafe-rgthree-kwargs.json"),
+            {
+                "AnySwitchRgthree": AnySwitchRgthree,
+                "DualClipLoader": DualClipLoader,
+                "PowerLoraLoaderRgthree": PowerLoraLoaderRgthree,
+            },
+        )
+
+        self.assertIn(
+            'powerloraloaderrgthree_631 = powerloraloaderrgthree.load_loras(',
+            generated,
+        )
+        self.assertIn(
+            'PowerLoraLoaderHeaderWidget={"type": "PowerLoraLoaderHeaderWidget"}',
+            generated,
+        )
+        self.assertIn('**{"\\u2795 Add Lora": ""}', generated)
+        self.assertNotIn('➕ Add Lora=""', generated)
+
+    def test_export_sanitizes_subgraph_identifiers_for_upscaler_workflows(self):
+        generated = export_workflow(
+            load_fixture("subgraph-upscaler-identifiers.json"),
+            {
+                "VaeDecode": VaeDecode,
+                "UpscaleModelLoader": UpscaleModelLoader,
+                "ImageUpscaleWithModel": ImageUpscaleWithModel,
+            },
+        )
+
+        self.assertIn("upscalemodelloader_42_0 = upscalemodelloader.load_model(", generated)
+        self.assertIn(
+            "imageupscalewithmodel_42_1 = imageupscalewithmodel.upscale(",
+            generated,
+        )
+        self.assertIn(
+            "upscale_model=get_value_at_index(upscalemodelloader_42_0, 0)",
+            generated,
+        )
+        self.assertNotIn("imageupscalewithmodel_42:1", generated)
+        self.assertNotIn("upscalemodelloader_42:0", generated)
+
+    def test_export_preserves_windows_style_model_paths(self):
+        generated = export_workflow(
+            load_fixture("windows-path-string.json"),
+            {
+                "WindowsPathNode": WindowsPathNode,
+            },
+        )
+
+        globals_dict = {"__name__": "generated_workflow_module"}
+        exec(generated, globals_dict)
+
+        self.assertEqual(
+            globals_dict["build_workflow"]()["1"]["inputs"]["path"],
+            r"C:\ComfyUI\models\upscale_models\RealESRGAN_x4plus.safetensors",
+        )
+
+    def test_export_preserves_trailing_backslash_string_literals(self):
+        generated = export_workflow(
+            load_fixture("trailing-backslash-string.json"),
+            {
+                "TextConcatenateNode": TextConcatenateNode,
+            },
+        )
+
+        globals_dict = {"__name__": "generated_workflow_module"}
+        exec(generated, globals_dict)
+
+        self.assertEqual(
+            globals_dict["build_workflow"]()["1"]["inputs"]["text_b"],
+            "\\",
+        )
+
+    def test_export_randomizes_string_seed_inputs_as_strings(self):
+        generated = export_workflow(
+            load_fixture("string-seed-node.json"),
+            {
+                "StringSeedNode": StringSeedNode,
+            },
+        )
+
+        self.assertIn(
+            'node_1_seed = prompt["1"]["inputs"]["seed"] = str(random.randint(1, 2**64))',
+            generated,
+        )
+        self.assertIn("seed=node_1_seed", generated)
+
+    def test_issue_cluster_regressions_render_parseable_python(self):
+        workflows = [
+            (
+                    load_fixture("unsafe-rgthree-kwargs.json"),
+                    {
+                        "AnySwitchRgthree": AnySwitchRgthree,
+                        "DualClipLoader": DualClipLoader,
+                        "PowerLoraLoaderRgthree": PowerLoraLoaderRgthree,
+                    },
+                ),
+            (
+                load_fixture("subgraph-upscaler-identifiers.json"),
+                {
+                    "VaeDecode": VaeDecode,
+                    "UpscaleModelLoader": UpscaleModelLoader,
+                    "ImageUpscaleWithModel": ImageUpscaleWithModel,
+                },
+            ),
+            (
+                load_fixture("trailing-backslash-string.json"),
+                {
+                    "TextConcatenateNode": TextConcatenateNode,
+                },
+            ),
+            (
+                load_fixture("windows-path-string.json"),
+                {
+                    "WindowsPathNode": WindowsPathNode,
+                },
+            ),
+            (
+                load_fixture("string-seed-node.json"),
+                {
+                    "StringSeedNode": StringSeedNode,
+                },
+            ),
+        ]
+
+        for workflow, mapping in workflows:
+            generated = export_workflow(workflow, mapping)
+            compile(generated, "<generated_workflow>", "exec")
+
+
+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_runtime_session.py

@@ -0,0 +1,469 @@
+import threading
+import unittest
+from unittest.mock import MagicMock, patch, call
+
+from comfyui_to_python.runtime_session import WorkflowSessionRuntime
+
+
+class StubNode:
+    FUNCTION = "execute"
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {"required": {"value": ("STRING",)}}
+
+    def execute(self, value):
+        return (f"result:{value}",)
+
+
+class TestWorkflowSessionRuntimeInit(unittest.TestCase):
+    """Tests for WorkflowSessionRuntime initialization."""
+
+    def test_init_default_cleanup_policy_is_session(self):
+        runtime = WorkflowSessionRuntime()
+        self.assertEqual(runtime._cleanup_policy, "session")
+
+    def test_init_accepts_per_run_policy(self):
+        runtime = WorkflowSessionRuntime(cleanup_policy="per_run")
+        self.assertEqual(runtime._cleanup_policy, "per_run")
+
+    def test_init_accepts_manual_policy(self):
+        runtime = WorkflowSessionRuntime(cleanup_policy="manual")
+        self.assertEqual(runtime._cleanup_policy, "manual")
+
+    def test_init_rejects_invalid_policy(self):
+        with self.assertRaises(ValueError):
+            WorkflowSessionRuntime(cleanup_policy="invalid")
+
+    def test_init_sets_default_reset_every_n_runs_none(self):
+        runtime = WorkflowSessionRuntime()
+        self.assertIsNone(runtime._reset_every_n_runs)
+
+    def test_init_accepts_reset_every_n_runs(self):
+        runtime = WorkflowSessionRuntime(reset_every_n_runs=5)
+        self.assertEqual(runtime._reset_every_n_runs, 5)
+
+    def test_init_has_lock(self):
+        runtime = WorkflowSessionRuntime()
+        self.assertIsInstance(runtime._lock, type(threading.Lock()))
+
+    def test_init_bootstrapped_false(self):
+        runtime = WorkflowSessionRuntime()
+        self.assertFalse(runtime.bootstrapped)
+
+    def test_init_custom_nodes_initialized_false(self):
+        runtime = WorkflowSessionRuntime()
+        self.assertFalse(runtime.custom_nodes_initialized)
+
+    def test_init_run_count_zero(self):
+        runtime = WorkflowSessionRuntime()
+        self.assertEqual(runtime.run_count, 0)
+
+    def test_init_closed_false(self):
+        runtime = WorkflowSessionRuntime()
+        self.assertFalse(runtime._closed)
+
+
+class TestWorkflowSessionRuntimeLifecycle(unittest.TestCase):
+    """Tests for WorkflowSessionRuntime lifecycle management."""
+
+    def _make_runtime(self):
+        return WorkflowSessionRuntime()
+
+    @patch("comfyui_to_python.runtime_session.bootstrap_comfyui_runtime")
+    def test_ensure_bootstrapped_calls_bootstrap_once(self, mock_bootstrap):
+        runtime = self._make_runtime()
+        runtime._ensure_bootstrapped()
+        mock_bootstrap.assert_called_once()
+        self.assertTrue(runtime.bootstrapped)
+
+    @patch("comfyui_to_python.runtime_session.bootstrap_comfyui_runtime")
+    def test_ensure_bootstrapped_skips_if_already_bootstrapped(self, mock_bootstrap):
+        runtime = self._make_runtime()
+        runtime._ensure_bootstrapped()
+        runtime._ensure_bootstrapped()
+        self.assertEqual(mock_bootstrap.call_count, 1)
+
+    @patch("comfyui_to_python.runtime_session.import_custom_nodes")
+    def test_ensure_custom_nodes_init_calls_import_once(self, mock_import):
+        runtime = self._make_runtime()
+        runtime._ensure_custom_nodes_initialized()
+        mock_import.assert_called_once()
+        self.assertTrue(runtime.custom_nodes_initialized)
+
+    @patch("comfyui_to_python.runtime_session.import_custom_nodes")
+    def test_ensure_custom_nodes_init_skips_if_already_initialized(self, mock_import):
+        runtime = self._make_runtime()
+        runtime._ensure_custom_nodes_initialized()
+        runtime._ensure_custom_nodes_initialized()
+        self.assertEqual(mock_import.call_count, 1)
+
+    def test_close_sets_closed_flag(self):
+        runtime = self._make_runtime()
+        runtime.close(unload_models=True)
+        self.assertTrue(runtime._closed)
+
+    def test_close_is_idempotent(self):
+        runtime = self._make_runtime()
+        with patch.object(runtime, "_do_close") as mock_do_close:
+            runtime.close(unload_models=True)
+            runtime.close(unload_models=True)
+            self.assertEqual(mock_do_close.call_count, 1)
+
+
+class TestWorkflowSessionRuntimeExceptionSafety(unittest.TestCase):
+    """Tests that exceptions during run() do not corrupt session state."""
+
+    def _make_runtime(self):
+        return WorkflowSessionRuntime()
+
+    @patch("comfyui_to_python.runtime_session.WorkflowSessionRuntime._do_run")
+    def test_exception_preserves_state_flags(self, mock_do_run):
+        runtime = self._make_runtime()
+        runtime.bootstrapped = True
+        runtime.custom_nodes_initialized = True
+        mock_do_run.side_effect = RuntimeError("simulated failure")
+        with self.assertRaises(RuntimeError):
+            runtime.run()
+        self.assertTrue(runtime.bootstrapped)
+        self.assertTrue(runtime.custom_nodes_initialized)
+
+    @patch("comfyui_to_python.runtime_session.WorkflowSessionRuntime._do_run")
+    def test_run_count_not_incremented_on_exception(self, mock_do_run):
+        runtime = self._make_runtime()
+        mock_do_run.side_effect = RuntimeError("simulated failure")
+        with self.assertRaises(RuntimeError):
+            runtime.run()
+        self.assertEqual(runtime.run_count, 0)
+
+    def test_run_count_increments_after_successful_run(self):
+        runtime = self._make_runtime()
+        workflow_data = {
+            "1": {
+                "class_type": "StubNode",
+                "inputs": {"value": "test"},
+            }
+        }
+        node_mappings = {"StubNode": StubNode}
+        runtime._workflow_data = workflow_data
+        runtime._node_class_mappings = node_mappings
+        runtime.node_instances = {"StubNode": StubNode()}
+        runtime._node_classes = {"StubNode": StubNode}
+
+        runtime.run()
+
+        self.assertEqual(runtime.run_count, 1)
+
+    def test_run_count_resets_after_reset_every_n_runs(self):
+        runtime = WorkflowSessionRuntime(reset_every_n_runs=2)
+        workflow_data = {
+            "1": {
+                "class_type": "StubNode",
+                "inputs": {"value": "test"},
+            }
+        }
+        node_mappings = {"StubNode": StubNode}
+        runtime._workflow_data = workflow_data
+        runtime._node_class_mappings = node_mappings
+
+        runtime.run()
+        self.assertEqual(runtime.run_count, 1)
+        runtime.run()
+        self.assertEqual(runtime.run_count, 0)
+        self.assertFalse(runtime.bootstrapped)
+        self.assertFalse(runtime.custom_nodes_initialized)
+        self.assertEqual(runtime.node_instances, {})
+
+
+class TestWorkflowSessionRuntimeAlreadyClosed(unittest.TestCase):
+    """Tests for behavior after close()."""
+
+    def _make_runtime(self):
+        return WorkflowSessionRuntime()
+
+    def test_run_after_close_raises(self):
+        runtime = self._make_runtime()
+        runtime.close()
+        with self.assertRaises(RuntimeError):
+            runtime.run()
+
+
+class TestWorkflowSessionRuntimeNodeInstances(unittest.TestCase):
+    """Tests for cached node instance management."""
+
+    def _make_runtime(self):
+        return WorkflowSessionRuntime()
+
+    def test_ensure_node_instances_creates_and_caches_instances(self):
+        runtime = self._make_runtime()
+        node_class = MagicMock()
+        runtime._ensure_node_instances({"TestNode": node_class})
+        node_class.assert_called_once()
+        self.assertIn("TestNode", runtime.node_instances)
+        runtime._ensure_node_instances({"TestNode": node_class})
+        node_class.assert_called_once()
+
+    def test_node_instances_are_cached_across_runs(self):
+        runtime = WorkflowSessionRuntime()
+        workflow_data = {
+            "1": {
+                "class_type": "StubNode",
+                "inputs": {"value": "test"},
+            }
+        }
+        original_stub_init = StubNode.__init__
+        init_calls = []
+
+        def tracking_init(self, *args, **kwargs):
+            init_calls.append(1)
+            original_stub_init(self)
+
+        StubNode.__init__ = tracking_init
+        try:
+            runtime._workflow_data = workflow_data
+            runtime._node_class_mappings = {"StubNode": StubNode}
+
+            runtime.run()
+            runtime.run()
+
+            self.assertEqual(len(init_calls), 1)
+        finally:
+            StubNode.__init__ = original_stub_init
+
+
+class TestWorkflowSessionRuntimeClearRuntimeCache(unittest.TestCase):
+    """Tests for clear_runtime_cache behavior."""
+
+    def _make_runtime(self):
+        return WorkflowSessionRuntime()
+
+    @patch("comfyui_to_python.runtime_session.cleanup_comfyui_runtime")
+    def test_clear_runtime_cache_session_policy_skips_unload(self, mock_cleanup):
+        runtime = WorkflowSessionRuntime(cleanup_policy="session")
+        runtime.bootstrapped = True
+        runtime.clear_runtime_cache()
+        # session policy should NOT call unload_all_models
+        mock_cleanup.assert_not_called()
+
+    @patch("comfyui_to_python.runtime_session.cleanup_comfyui_runtime")
+    def test_clear_runtime_cache_per_run_policy_calls_full_cleanup(self, mock_cleanup):
+        runtime = WorkflowSessionRuntime(cleanup_policy="per_run")
+        runtime.bootstrapped = True
+        runtime.clear_runtime_cache()
+        mock_cleanup.assert_called_once_with(unload_models=True)
+
+    @patch("comfyui_to_python.runtime_session.cleanup_comfyui_runtime")
+    def test_clear_runtime_cache_manual_policy_no_cleanup(self, mock_cleanup):
+        runtime = WorkflowSessionRuntime(cleanup_policy="manual")
+        runtime.bootstrapped = True
+        runtime.clear_runtime_cache()
+        mock_cleanup.assert_not_called()
+
+
+class TestWorkflowSessionRuntimeDoClose(unittest.TestCase):
+    """Tests for _do_close internal method."""
+
+    def _make_runtime(self):
+        return WorkflowSessionRuntime()
+
+    @patch("comfyui_to_python.runtime_session.cleanup_comfyui_runtime")
+    @patch("comfyui_to_python.runtime_session.gc")
+    def test_do_close_calls_cleanup_and_gc(self, mock_gc, mock_cleanup):
+        runtime = self._make_runtime()
+        runtime._do_close(unload_models=True)
+        mock_cleanup.assert_called_once_with(unload_models=True)
+        mock_gc.collect.assert_called_once()
+
+        mock_cleanup.reset_mock()
+        mock_gc.reset_mock()
+        runtime = self._make_runtime()
+        runtime._do_close(unload_models=False)
+        mock_cleanup.assert_called_once_with(unload_models=False)
+        mock_gc.collect.assert_called_once()
+
+
+class TestWorkflowSessionRuntimeRun(unittest.TestCase):
+    """Tests for WorkflowSessionRuntime.run() workflow execution."""
+
+    def _make_runtime(self):
+        return WorkflowSessionRuntime()
+
+    def test_run_executes_workflow_nodes(self):
+        runtime = self._make_runtime()
+        workflow_data = {
+            "1": {
+                "class_type": "StubNode",
+                "inputs": {"value": "test"},
+            }
+        }
+        node_mappings = {"StubNode": StubNode}
+        runtime._workflow_data = workflow_data
+        runtime._node_class_mappings = node_mappings
+        runtime.node_instances = {"StubNode": StubNode()}
+        runtime._node_classes = {"StubNode": StubNode}
+
+        result = runtime.run()
+
+        self.assertIn("1", result)
+        self.assertEqual(result["1"], ["result:test"])
+
+    @patch(
+        "comfyui_to_python.runtime_session.WorkflowSessionRuntime._ensure_node_instances"
+    )
+    def test_run_with_no_workflow_data_returns_none(self, mock_ensure):
+        runtime = self._make_runtime()
+        result = runtime.run()
+        self.assertIsNone(result)
+
+    @patch(
+        "comfyui_to_python.runtime_session.WorkflowSessionRuntime._ensure_node_instances"
+    )
+    def test_run_with_no_node_mappings_returns_none(self, mock_ensure):
+        runtime = self._make_runtime()
+        runtime._workflow_data = {"1": {"class_type": "Test", "inputs": {}}}
+        result = runtime.run()
+        self.assertIsNone(result)
+
+    @patch(
+        "comfyui_to_python.runtime_session.WorkflowSessionRuntime._ensure_node_instances"
+    )
+    def test_run_node_not_in_mappings_is_skipped(self, mock_ensure):
+        workflow_data = {
+            "1": {
+                "class_type": "UnknownNode",
+                "inputs": {"value": "test"},
+            }
+        }
+        node_mappings = {"StubNode": StubNode}
+        runtime = self._make_runtime()
+        runtime._workflow_data = workflow_data
+        runtime._node_class_mappings = node_mappings
+
+        result = runtime.run()
+
+        self.assertEqual(result, {})
+
+    def test_run_node_returns_tuple_is_converted_to_list(self):
+        runtime = self._make_runtime()
+        workflow_data = {
+            "1": {
+                "class_type": "StubNode",
+                "inputs": {"value": "multi"},
+            }
+        }
+        node_mappings = {"StubNode": StubNode}
+        runtime._workflow_data = workflow_data
+        runtime._node_class_mappings = node_mappings
+        runtime.node_instances = {"StubNode": StubNode()}
+        runtime._node_classes = {"StubNode": StubNode}
+
+        result = runtime.run()
+
+        self.assertIsInstance(result["1"], list)
+        self.assertEqual(result["1"][0], "result:multi")
+
+    def test_run_persists_parameters_across_calls(self):
+        runtime = self._make_runtime()
+        workflow_data = {
+            "1": {
+                "class_type": "StubNode",
+                "inputs": {"value": "persist"},
+            }
+        }
+        node_mappings = {"StubNode": StubNode}
+
+        with patch.object(
+            runtime, "_do_run", return_value={"1": ["result:persist"]}
+        ) as mock_do_run:
+            runtime.run(
+                workflow_data=workflow_data, node_class_mappings=node_mappings
+            )
+            runtime.run()
+
+        self.assertEqual(mock_do_run.call_count, 2)
+
+    def test_run_does_not_corrupt_session_on_exception(self):
+        runtime = self._make_runtime()
+        workflow_data = {
+            "1": {
+                "class_type": "StubNode",
+                "inputs": {"value": "test"},
+            }
+        }
+        node_mappings = {"StubNode": StubNode}
+        runtime._workflow_data = workflow_data
+        runtime._node_class_mappings = node_mappings
+        mock_instance = MagicMock()
+        mock_instance.execute.side_effect = ValueError("boom")
+        runtime.node_instances = {"StubNode": mock_instance}
+        runtime._node_classes = {"StubNode": StubNode}
+
+        with self.assertRaises(ValueError):
+            runtime.run()
+
+        runtime.node_instances = {"StubNode": StubNode()}
+        runtime._node_classes = {"StubNode": StubNode}
+        result = runtime.run()
+        self.assertEqual(result["1"], ["result:test"])
+
+    @patch(
+        "comfyui_to_python.runtime_session.WorkflowSessionRuntime._ensure_node_instances"
+    )
+    def test_session_policy_does_not_clear_cache(self, mock_ensure):
+        runtime = WorkflowSessionRuntime(cleanup_policy="session")
+        workflow_data = {
+            "1": {
+                "class_type": "StubNode",
+                "inputs": {"value": "test"},
+            }
+        }
+        node_mappings = {"StubNode": StubNode}
+        runtime._workflow_data = workflow_data
+        runtime._node_class_mappings = node_mappings
+
+        with patch.object(runtime, "clear_runtime_cache") as mock_clear:
+            runtime.run()
+            mock_clear.assert_not_called()
+
+    @patch(
+        "comfyui_to_python.runtime_session.WorkflowSessionRuntime._ensure_node_instances"
+    )
+    def test_per_run_policy_clears_cache_after_each_run(self, mock_ensure):
+        runtime = WorkflowSessionRuntime(cleanup_policy="per_run")
+        workflow_data = {
+            "1": {
+                "class_type": "StubNode",
+                "inputs": {"value": "test"},
+            }
+        }
+        node_mappings = {"StubNode": StubNode}
+        runtime._workflow_data = workflow_data
+        runtime._node_class_mappings = node_mappings
+
+        with patch.object(runtime, "clear_runtime_cache") as mock_clear:
+            runtime.run()
+            mock_clear.assert_called_once()
+
+    @patch(
+        "comfyui_to_python.runtime_session.WorkflowSessionRuntime._ensure_node_instances"
+    )
+    def test_manual_policy_never_clears_cache(self, mock_ensure):
+        runtime = WorkflowSessionRuntime(cleanup_policy="manual")
+        workflow_data = {
+            "1": {
+                "class_type": "StubNode",
+                "inputs": {"value": "test"},
+            }
+        }
+        node_mappings = {"StubNode": StubNode}
+        runtime._workflow_data = workflow_data
+        runtime._node_class_mappings = node_mappings
+
+        with patch.object(runtime, "clear_runtime_cache") as mock_clear:
+            runtime.run()
+            runtime.run()
+            mock_clear.assert_not_called()
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_session_e2e.py

@@ -0,0 +1,370 @@
+"""End-to-end tests for session execution mode feature.
+
+Covers:
+- Code generation correctness (session vs oneshot)
+- Runtime execution of generated scripts
+- Multi-run session behavior
+"""
+
+import ast
+import json
+import os
+import subprocess
+import sys
+import tempfile
+import unittest
+from io import StringIO
+from pathlib import Path
+
+ROOT = Path(__file__).resolve().parents[1]
+COMFYUI_PATH = os.environ.get("COMFYUI_PATH", str(ROOT.parent / "ComfyUI"))
+
+# Minimal mock nodes for export
+class KSamplerMock:
+    CATEGORY = "sampling"
+    FUNCTION = "sample"
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "seed": ("INT", {"default": 0}),
+                "steps": ("INT", {"default": 20}),
+                "cfg": ("FLOAT", {"default": 8.0}),
+                "sampler_name": (["euler", "heun"],),
+                "scheduler": (["normal"],),
+                "denoise": ("FLOAT", {"default": 1.0}),
+                "model": ("MODEL",),
+                "positive": ("CONDITIONING",),
+                "negative": ("CONDITIONING",),
+                "latent_image": ("LATENT",),
+            }
+        }
+    def sample(self, seed, steps, cfg, sampler_name, scheduler, denoise, model, positive, negative, latent_image):
+        return ({"samples": latent_image},)
+
+class CheckpointLoaderMock:
+    CATEGORY = "loaders"
+    FUNCTION = "load_checkpoint"
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "ckpt_name": ("STRING",),
+            }
+        }
+    def load_checkpoint(self, ckpt_name):
+        return (None, None, None)
+
+class VAEDecodeMock:
+    CATEGORY = "latent"
+    FUNCTION = "decode"
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "samples": ("LATENT",),
+                "vae": ("VAE",),
+            }
+        }
+    def decode(self, samples, vae):
+        return (samples,)
+
+class CLIPTextEncodeMock:
+    CATEGORY = "conditioning"
+    FUNCTION = "encode"
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "text": ("STRING",),
+                "clip": ("CLIP",),
+            }
+        }
+    def encode(self, text, clip):
+        return ([],)
+
+class EmptyLatentImageMock:
+    CATEGORY = "latent"
+    FUNCTION = "generate"
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "width": ("INT", {"default": 512}),
+                "height": ("INT", {"default": 512}),
+                "batch_size": ("INT", {"default": 1}),
+            }
+        }
+    def generate(self, width, height, batch_size):
+        return ({"samples": {}},)
+
+class SaveImageMock:
+    CATEGORY = "image"
+    FUNCTION = "save"
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "images": ("IMAGE",),
+                "filename_prefix": ("STRING",),
+            }
+        }
+    def save(self, images, filename_prefix):
+        return ()
+
+NODENAMES = {
+    "CheckpointLoaderSimple": CheckpointLoaderMock,
+    "CLIPTextEncode": CLIPTextEncodeMock,
+    "KSampler": KSamplerMock,
+    "VAEDecode": VAEDecodeMock,
+    "EmptyLatentImage": EmptyLatentImageMock,
+    "SaveImage": SaveImageMock,
+}
+
+TEXT_TO_IMAGE_WORKFLOW = {
+    "1": {
+        "class_type": "CheckpointLoaderSimple",
+        "inputs": {"ckpt_name": "v1-5-pruned-emaonly-fp16.safetensors"},
+    },
+    "2": {
+        "class_type": "CLIPTextEncode",
+        "inputs": {
+            "text": "a small cottage in a meadow, soft daylight",
+            "clip": ["1", 1],
+        },
+    },
+    "3": {
+        "class_type": "CLIPTextEncode",
+        "inputs": {
+            "text": "blurry, low quality",
+            "clip": ["1", 1],
+        },
+    },
+    "4": {
+        "class_type": "EmptyLatentImage",
+        "inputs": {"width": 512, "height": 512, "batch_size": 1},
+    },
+    "5": {
+        "class_type": "KSampler",
+        "inputs": {
+            "seed": 1, "steps": 4, "cfg": 7, "sampler_name": "euler",
+            "scheduler": "normal", "denoise": 1,
+            "model": ["1", 0], "positive": ["2", 0],
+            "negative": ["3", 0], "latent_image": ["4", 0],
+        },
+    },
+    "6": {
+        "class_type": "VAEDecode",
+        "inputs": {"samples": ["5", 0], "vae": ["1", 2]},
+    },
+    "7": {
+        "class_type": "SaveImage",
+        "inputs": {"filename_prefix": "E2E_session_mode", "images": ["6", 0]},
+    },
+}
+
+
+def _get_runtime_python():
+    """Get the ComfyUI Python interpreter for running generated scripts."""
+    rt_python = Path(COMFYUI_PATH) / ".venv" / "bin" / "python"
+    if rt_python.is_file():
+        return str(rt_python)
+    return sys.executable
+
+
+def _export_session_workflow_in_runtime_env(
+    workflow_json,
+    execution_mode="session",
+):
+    """Export workflow via subprocess in ComfyUI runtime env.
+
+    Re-enters the runtime interpreter so ``ComfyUItoPython`` can import
+    ComfyUI's nodes.py (which requires torch, not available in the test venv).
+    """
+    runtime_python = _get_runtime_python()
+    env = os.environ.copy()
+    env["COMFYUI_PATH"] = COMFYUI_PATH
+    env["PYTHONPATH"] = os.pathsep.join([str(ROOT), env.get("PYTHONPATH", "")]).rstrip(
+        os.pathsep
+    )
+
+    tmp_path = tempfile.mktemp(suffix=".py")
+
+    try:
+        result = subprocess.run(
+            [
+                runtime_python,
+                str(Path(__file__).resolve().parents[0] / "runtime" / "run_runtime_validation.py"),
+                "--internal-export",
+                "text-to-image",
+                "--execution-mode",
+                execution_mode,
+                "--generated-path",
+                tmp_path,
+            ],
+            cwd=ROOT,
+            env=env,
+            capture_output=True,
+            text=True,
+        )
+        if result.returncode != 0:
+            output = (result.stderr or result.stdout or "").strip()
+            if "Missing runtime dependency" not in output and "ModuleNotFoundError" not in output:
+                classification = "repo regression"
+            else:
+                classification = "environment/setup failure"
+            raise RuntimeError(
+                f"Runtime export failed: [{classification}] {output}"
+            )
+        # Read generated code from the temp file (stdout is polluted by
+        # ComfyUI runtime prints such as the sys.path line).
+        return Path(tmp_path).read_text()
+    finally:
+        os.unlink(tmp_path)
+
+
+def _export_workflow(execution_mode="oneshot"):
+    """Export workflow to a string using ComfyUItoPython (unit tests, mock nodes)."""
+    from comfyui_to_python import ComfyUItoPython
+    output = StringIO()
+    ComfyUItoPython(
+        workflow=json.dumps(TEXT_TO_IMAGE_WORKFLOW),
+        output_file=output,
+        node_class_mappings=NODENAMES,
+        execution_mode=execution_mode,
+    )
+    return output.getvalue()
+
+
+class SessionCodeGenerationTest(unittest.TestCase):
+    """Unit tests for session mode code generation."""
+
+    def test_oneshot_code_has_bootstrap_helpers_no_session(self):
+        """Oneshot mode: bootstrap/cleanup helpers present, no WorkflowSession class."""
+        generated = _export_workflow("oneshot")
+        ast.parse(generated)
+        self.assertIn("bootstrap_comfyui_runtime()", generated)
+        self.assertIn("cleanup_comfyui_runtime(", generated)
+        self.assertNotIn("class WorkflowSession", generated)
+        self.assertNotIn("session.run()", generated)
+
+    def test_session_code_has_workflow_session_class_and_session_methods(self):
+        """Session mode: WorkflowSession class with run() and close() present."""
+        generated = _export_workflow("session")
+        ast.parse(generated)
+        self.assertIn("class WorkflowSession", generated)
+        self.assertIn("def run(self)", generated)
+        self.assertIn("def close(self, unload_models", generated)
+        self.assertIn("session.run()", generated)
+        self.assertIn("session.close(", generated)
+
+
+class SessionModeExecutionTest(unittest.TestCase):
+    """E2E tests for session mode script execution."""
+
+    @unittest.skipIf(not Path(COMFYUI_PATH).is_dir(), "ComfyUI checkout not available")
+    def test_oneshot_e2e_text_to_image(self):
+        """Oneshot mode: generate and run text-to-image workflow, verify PNG output."""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            script_path = Path(tmpdir) / "oneshot.py"
+            env = os.environ.copy()
+            env["COMFYUI_PATH"] = COMFYUI_PATH
+            env["PYTHONPATH"] = os.pathsep.join([str(ROOT), env.get("PYTHONPATH", "")])
+            runtime_py = _get_runtime_python()
+
+            # Export inside ComfyUI env so node mappings resolve
+            generated = _export_session_workflow_in_runtime_env(
+                json.dumps(TEXT_TO_IMAGE_WORKFLOW),
+                execution_mode="oneshot",
+            )
+            script_path.write_text(generated)
+
+            # Run
+            result = subprocess.run(
+                [runtime_py, str(script_path), "--cpu"],
+                cwd=ROOT,
+                env=env,
+                capture_output=True,
+                text=True,
+                timeout=300,
+            )
+            self.assertEqual(result.returncode, 0, f"stderr: {result.stderr}")
+
+            # Check output
+            output_dir = Path(COMFYUI_PATH) / "output"
+            new_outputs = list(output_dir.glob("E2E_text_to_image*.png"))
+            self.assertTrue(len(new_outputs) > 0, "No PNG output produced")
+
+    @unittest.skipIf(not Path(COMFYUI_PATH).is_dir(), "ComfyUI checkout not available")
+    def test_session_e2e_text_to_image(self):
+        """Session mode: generate session-mode script, verify WorkflowSession present and run."""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            script_path = Path(tmpdir) / "session.py"
+            env = os.environ.copy()
+            env["COMFYUI_PATH"] = COMFYUI_PATH
+            env["PYTHONPATH"] = os.pathsep.join([str(ROOT), env.get("PYTHONPATH", "")])
+            runtime_py = _get_runtime_python()
+
+            # Export inside ComfyUI env
+            generated = _export_session_workflow_in_runtime_env(
+                json.dumps(TEXT_TO_IMAGE_WORKFLOW),
+                execution_mode="session",
+            )
+            script_path.write_text(generated)
+
+            # Verify code structure
+            self.assertIn("class WorkflowSession", generated)
+
+            # Run
+            result = subprocess.run(
+                [runtime_py, str(script_path), "--cpu"],
+                cwd=ROOT,
+                env=env,
+                capture_output=True,
+                text=True,
+                timeout=300,
+            )
+            self.assertEqual(result.returncode, 0, f"stderr: {result.stderr}")
+
+            # Check output
+            output_dir = Path(COMFYUI_PATH) / "output"
+            new_outputs = list(output_dir.glob("E2E_text_to_image*.png"))
+            self.assertTrue(len(new_outputs) > 0, "No PNG output produced")
+
+    @unittest.skipIf(not Path(COMFYUI_PATH).is_dir(), "ComfyUI checkout not available")
+    def test_session_e2e_multiple_runs(self):
+        """Session mode: generate script that calls session.run() 3x in a row, verify no crash."""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            script_path = Path(tmpdir) / "session_multi.py"
+            env = os.environ.copy()
+            env["COMFYUI_PATH"] = COMFYUI_PATH
+            env["PYTHONPATH"] = os.pathsep.join([str(ROOT), env.get("PYTHONPATH", "")])
+            runtime_py = _get_runtime_python()
+
+            # Export session mode inside ComfyUI env
+            generated = _export_session_workflow_in_runtime_env(
+                json.dumps(TEXT_TO_IMAGE_WORKFLOW),
+                execution_mode="session",
+            )
+
+            # Modify to run 3x
+            modified = generated.replace(
+                "        session.run()",
+                "        session.run()\n        session.run()\n        session.run()",
+            )
+            script_path.write_text(modified)
+
+            # Run
+            result = subprocess.run(
+                [runtime_py, str(script_path), "--cpu"],
+                cwd=ROOT,
+                env=env,
+                capture_output=True,
+                text=True,
+                timeout=600,
+            )
+            self.assertEqual(result.returncode, 0, f"stderr: {result.stderr}")
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_session_renderer.py

@@ -0,0 +1,273 @@
+import json
+import unittest
+from io import StringIO
+from unittest.mock import patch
+
+from comfyui_to_python import ComfyUItoPython
+
+
+class LoadImage:
+    CATEGORY = "image"
+    FUNCTION = "load_image"
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {"required": {"image": ("STRING",)}}
+
+    def load_image(self, image):
+        return (image,)
+
+
+class SessionRendererTest(unittest.TestCase):
+    """Tests for session mode code generation in the renderer."""
+
+    def test_session_mode_generates_workflow_session_class(self):
+        workflow = {
+            "1": {
+                "class_type": "LoadImage",
+                "inputs": {"image": "example.png"},
+            }
+        }
+
+        output = StringIO()
+        ComfyUItoPython(
+            workflow=json.dumps(workflow),
+            output_file=output,
+            node_class_mappings={"LoadImage": LoadImage},
+            execution_mode="session",
+        )
+
+        generated = output.getvalue()
+        self.assertIn("class WorkflowSession", generated)
+
+    def test_session_mode_generates_main_wrapper(self):
+        workflow = {
+            "1": {
+                "class_type": "LoadImage",
+                "inputs": {"image": "example.png"},
+            }
+        }
+
+        output = StringIO()
+        ComfyUItoPython(
+            workflow=json.dumps(workflow),
+            output_file=output,
+            node_class_mappings={"LoadImage": LoadImage},
+            execution_mode="session",
+        )
+
+        generated = output.getvalue()
+        self.assertIn("def main(unload_models: bool | None = None)", generated)
+        self.assertIn("WorkflowSession(", generated)
+
+    def test_session_mode_main_creates_session_with_per_run_policy(self):
+        workflow = {
+            "1": {
+                "class_type": "LoadImage",
+                "inputs": {"image": "example.png"},
+            }
+        }
+
+        output = StringIO()
+        ComfyUItoPython(
+            workflow=json.dumps(workflow),
+            output_file=output,
+            node_class_mappings={"LoadImage": LoadImage},
+            execution_mode="session",
+        )
+
+        generated = output.getvalue()
+        self.assertIn('cleanup_policy="per_run"', generated)
+
+    def test_session_mode_main_has_try_finally_close(self):
+        workflow = {
+            "1": {
+                "class_type": "LoadImage",
+                "inputs": {"image": "example.png"},
+            }
+        }
+
+        output = StringIO()
+        ComfyUItoPython(
+            workflow=json.dumps(workflow),
+            output_file=output,
+            node_class_mappings={"LoadImage": LoadImage},
+            execution_mode="session",
+        )
+
+        generated = output.getvalue()
+        self.assertIn("try:", generated)
+        self.assertIn("session.run()", generated)
+        self.assertIn("finally:", generated)
+        self.assertIn("session.close(", generated)
+
+    def test_session_mode_generates_run_method(self):
+        workflow = {
+            "1": {
+                "class_type": "LoadImage",
+                "inputs": {"image": "example.png"},
+            }
+        }
+
+        output = StringIO()
+        ComfyUItoPython(
+            workflow=json.dumps(workflow),
+            output_file=output,
+            node_class_mappings={"LoadImage": LoadImage},
+            execution_mode="session",
+        )
+
+        generated = output.getvalue()
+        self.assertIn("def run(self", generated)
+
+    def test_session_mode_generates_close_method(self):
+        workflow = {
+            "1": {
+                "class_type": "LoadImage",
+                "inputs": {"image": "example.png"},
+            }
+        }
+
+        output = StringIO()
+        ComfyUItoPython(
+            workflow=json.dumps(workflow),
+            output_file=output,
+            node_class_mappings={"LoadImage": LoadImage},
+            execution_mode="session",
+        )
+
+        generated = output.getvalue()
+        self.assertIn("def close(self, unload_models: bool | None = None)", generated)
+
+    def test_session_mode_oneshot_generates_same_code(self):
+        workflow = {
+            "1": {
+                "class_type": "LoadImage",
+                "inputs": {"image": "example.png"},
+            }
+        }
+
+        output = StringIO()
+        ComfyUItoPython(
+            workflow=json.dumps(workflow),
+            output_file=output,
+            node_class_mappings={"LoadImage": LoadImage},
+            execution_mode="oneshot",
+        )
+
+        generated = output.getvalue()
+        self.assertNotIn("class WorkflowSession", 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)
+
+    def test_session_mode_oneshot_default(self):
+        workflow = {
+            "1": {
+                "class_type": "LoadImage",
+                "inputs": {"image": "example.png"},
+            }
+        }
+
+        output = StringIO()
+        ComfyUItoPython(
+            workflow=json.dumps(workflow),
+            output_file=output,
+            node_class_mappings={"LoadImage": LoadImage},
+        )
+
+        generated = output.getvalue()
+        self.assertNotIn("class WorkflowSession", generated)
+        self.assertIn("def main(unload_models: bool | None = None)", generated)
+
+    def test_session_mode_script_is_executable(self):
+        workflow = {
+            "1": {
+                "class_type": "LoadImage",
+                "inputs": {"image": "example.png"},
+            }
+        }
+
+        output = StringIO()
+        ComfyUItoPython(
+            workflow=json.dumps(workflow),
+            output_file=output,
+            node_class_mappings={"LoadImage": LoadImage},
+            execution_mode="session",
+        )
+
+        generated = output.getvalue()
+
+        # Should not raise — checks that generated code has valid syntax
+        compile(generated, "<generated>", "exec")
+
+        # main() should be callable
+        globals_dict = {"__name__": "generated_workflow_module"}
+        with patch("comfyui_to_python.runtime_session.WorkflowSessionRuntime"):
+            exec(generated, globals_dict)
+        self.assertIn("WorkflowSession", globals_dict)
+        self.assertIn("main", globals_dict)
+        self.assertTrue(callable(globals_dict["main"]))
+
+    def test_session_mode_generates_workflow_literal(self):
+        workflow = {
+            "1": {
+                "class_type": "LoadImage",
+                "inputs": {"image": "example.png"},
+            }
+        }
+
+        output = StringIO()
+        ComfyUItoPython(
+            workflow=json.dumps(workflow),
+            output_file=output,
+            node_class_mappings={"LoadImage": LoadImage},
+            execution_mode="session",
+        )
+
+        generated = output.getvalue()
+        self.assertIn("def build_workflow()", generated)
+        self.assertIn('return', generated)
+        self.assertIn('"class_type": "LoadImage"', generated)
+
+    def test_session_mode_generates_bootstrap_helper(self):
+        workflow = {
+            "1": {
+                "class_type": "LoadImage",
+                "inputs": {"image": "example.png"},
+            }
+        }
+
+        output = StringIO()
+        ComfyUItoPython(
+            workflow=json.dumps(workflow),
+            output_file=output,
+            node_class_mappings={"LoadImage": LoadImage},
+            execution_mode="session",
+        )
+
+        generated = output.getvalue()
+        self.assertIn("def bootstrap_comfyui_runtime()", generated)
+
+    def test_session_mode_generates_cleanup_helper(self):
+        workflow = {
+            "1": {
+                "class_type": "LoadImage",
+                "inputs": {"image": "example.png"},
+            }
+        }
+
+        output = StringIO()
+        ComfyUItoPython(
+            workflow=json.dumps(workflow),
+            output_file=output,
+            node_class_mappings={"LoadImage": LoadImage},
+            execution_mode="session",
+        )
+
+        generated = output.getvalue()
+        self.assertIn("def cleanup_comfyui_runtime(", generated)
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_upscale_model_loader_export.py

@@ -145,20 +145,26 @@ class UpscaleModelLoaderExportTest(unittest.TestCase):
         generated = output.getvalue()
 
         self.assertIn("def bootstrap_comfyui_runtime()", generated)
+        self.assertIn("def cleanup_comfyui_runtime(", generated)
         self.assertIn("import comfy.options", generated)
         self.assertIn("comfy.options.enable_args_parsing()", generated)
         self.assertIn("import cuda_malloc", 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,
@@ -172,6 +178,10 @@ class UpscaleModelLoaderExportTest(unittest.TestCase):
             main_section.index("import torch"),
         )
         self.assertLess(generated.index("import cuda_malloc"), generated.index("import torch"))
+        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 = {
@@ -431,18 +441,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):

tests/test_workflow_session.py

@@ -0,0 +1,75 @@
+import unittest
+from unittest.mock import MagicMock, patch
+
+from comfyui_to_python.runtime_session import WorkflowSession
+
+
+class TestWorkflowSessionInit(unittest.TestCase):
+    """Tests for WorkflowSession initialization."""
+
+    def test_init_creates_runtime(self):
+        session = WorkflowSession()
+        self.assertIsNotNone(session._runtime)
+
+    def test_init_passes_cleanup_policy(self):
+        session = WorkflowSession(cleanup_policy="per_run")
+        self.assertEqual(session._runtime._cleanup_policy, "per_run")
+
+    def test_init_passes_reset_every_n_runs(self):
+        session = WorkflowSession(reset_every_n_runs=5)
+        self.assertEqual(session._runtime._reset_every_n_runs, 5)
+
+
+class TestWorkflowSessionDelegation(unittest.TestCase):
+    """Tests for WorkflowSession public API delegation to internal runtime."""
+
+    def _make_session(self):
+        return WorkflowSession()
+
+    def test_run_delegates_to_runtime(self):
+        session = self._make_session()
+        workflow_data = {
+            "1": {
+                "class_type": "StubNode",
+                "inputs": {"value": "test"},
+            }
+        }
+        node_mappings = {"StubNode": MagicMock()}
+        session._runtime.node_instances = {"StubNode": MagicMock()}
+        session._runtime._node_classes = {"StubNode": MagicMock()}
+
+        with patch.object(
+            session._runtime, "run", return_value={"1": ["result"]}
+        ) as mock_run:
+            session.run(workflow_data=workflow_data, node_class_mappings=node_mappings)
+            mock_run.assert_called_once_with(
+                workflow_data=workflow_data,
+                node_class_mappings=node_mappings,
+                extra_pnginfo=None,
+            )
+
+    def test_clear_runtime_cache_delegates_to_runtime(self):
+        session = self._make_session()
+        with patch.object(
+            session._runtime, "clear_runtime_cache"
+        ) as mock_clear:
+            session.clear_runtime_cache()
+            mock_clear.assert_called_once()
+
+    def test_close_delegates_to_runtime(self):
+        session = self._make_session()
+        with patch.object(
+            session._runtime, "close"
+        ) as mock_close:
+            session.close(unload_models=True)
+            mock_close.assert_called_once_with(unload_models=True)
+
+    def test_run_raises_after_close(self):
+        session = self._make_session()
+        session.close()
+        with self.assertRaises(RuntimeError):
+            session.run()
+
+
+if __name__ == "__main__":
+    unittest.main()

uv.lock

@@ -57,7 +57,7 @@ wheels = [
 
 [[package]]
 name = "comfyui-to-python-extension"
-version = "2.0.0"
+version = "2.1.0"
 source = { virtual = "." }
 dependencies = [
     { name = "black" },