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.
2026-04-19 12:06:46 -05:00
parent 6cdcc235a0
commit 08ad731038
13 changed files with 1472 additions and 61 deletions
@@ -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)