feat: add session execution mode with WorkflowSession and WorkflowSessionRuntime

Add session-mode code generation that produces a reusable WorkflowSession class
instead of a simple oneshot script. The session mode generates code with an
in-class run/close lifecycle for warm ComfyUI reuse.

Also refactors render.py to extract shared code sections into common methods,
reducing duplication between oneshot and session renderers.

Adds WorkflowSessionRuntime with cleanup policies (per_run, session, manual)
and reset_every_n_runs support, plus 72 new tests covering session rendering,
runtime session, and session export pipeline.

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/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 = {}
@@ -132,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(

comfyui_to_python/generator/render.py

@@ -1,4 +1,5 @@
 import inspect
+import threading
 from pprint import pformat
 from typing import Any
 
@@ -20,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,
@@ -52,16 +63,19 @@ class WorkflowRenderer:
 
         if plan.custom_nodes:
             static_imports.append(f"\n{inspect.getsource(import_custom_nodes)}\n")
-            custom_nodes_call = "import_custom_nodes()"
+
+        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}",
@@ -74,52 +88,16 @@ class WorkflowRenderer:
             "extra_pnginfo = build_extra_pnginfo()",
         ]
 
-        execution_section = [
-            "# Workflow execution",
-            "def main(unload_models: bool | None = None):",
-            "    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",
-                "",
-                "    try:",
-                "        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()
-        )
-        execution_section.extend(
-            [
-                "    finally:",
-                "        cleanup_comfyui_runtime(unload_models=unload_models)",
-            ]
-        )
+    # ── 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
             + [""]
@@ -129,6 +107,198 @@ 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]:
+        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()",
+            "",
+            "            if self._cleanup_policy == 'session' and not self._custom_nodes_initialized:",
+            "                self._custom_nodes_initialized = True",
+            "                import_custom_nodes()",
+            "",
+            "            prompt = json.loads(json.dumps(build_workflow()))",
+            "            extra_pnginfo = build_extra_pnginfo()",
+            "",
+            "            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())
+
+        lines.extend(
+            [
+                "            finally:",
+                "                if self._cleanup_policy == 'per_run':",
+                "                    cleanup_comfyui_runtime(unload_models=True)",
+                "",
+            ]
+        )
+
+        # Build outputs collection: outputs = {node_id: var_name, ...}
+        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(
+            [
+                "",
+                "    def close(self, unload_models: bool = True) -> None:",
+                '        """Close the session and optionally unload models."""',
+                "        with self._lock:",
+                "            if self._closed:",
+                "                return",
+                "            self._closed = True",
+                "            cleanup_comfyui_runtime(unload_models=unload_models)",
+                "            import gc",
+                "            gc.collect()",
+            ]
+        )
+
+        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/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)