Peyton DeNiro
79312a25bb
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.
341 lines
13 KiB
Python
341 lines
13 KiB
Python
import inspect
|
|
import threading
|
|
from pprint import pformat
|
|
from typing import Any
|
|
|
|
import black
|
|
|
|
from ..node_runtime import import_custom_nodes
|
|
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,
|
|
)
|
|
from .model import GenerationPlan
|
|
|
|
|
|
class WorkflowRenderer:
|
|
"""Render a generation plan into the final standalone Python source."""
|
|
|
|
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,
|
|
get_comfyui_path,
|
|
find_path,
|
|
add_comfyui_directory_to_sys_path,
|
|
add_extra_model_paths,
|
|
bootstrap_comfyui_runtime,
|
|
cleanup_comfyui_runtime,
|
|
]:
|
|
func_strings.append(f"\n{inspect.getsource(func)}")
|
|
|
|
static_imports = [
|
|
"# Imports",
|
|
"import json",
|
|
"import os",
|
|
"import random",
|
|
"import sys",
|
|
"import threading",
|
|
"from typing import Sequence, Mapping, Any, Union",
|
|
] + func_strings
|
|
|
|
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:
|
|
extra_pnginfo_literal = self.format_python_literal(
|
|
{"workflow": plan.metadata_workflow_data}
|
|
)
|
|
|
|
return [
|
|
"# Workflow data",
|
|
"def build_workflow() -> dict[str, Any]:",
|
|
f" return {workflow_literal}",
|
|
"",
|
|
"def build_extra_pnginfo() -> dict[str, Any] | None:",
|
|
f" return {extra_pnginfo_literal}",
|
|
"",
|
|
"workflow = build_workflow()",
|
|
"prompt = json.loads(json.dumps(workflow))",
|
|
"extra_pnginfo = build_extra_pnginfo()",
|
|
]
|
|
|
|
# ── oneshot renderer ─────────────────────────────────────────────
|
|
|
|
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(
|
|
imports_section
|
|
+ [""]
|
|
+ workflow_section
|
|
+ [""]
|
|
+ execution_section
|
|
+ [""]
|
|
+ entrypoint_section
|
|
)
|
|
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)
|
|
|
|
@staticmethod
|
|
def build_function_body(
|
|
code_lines: list[str], empty_fallback: str, indentation: str = " "
|
|
) -> str:
|
|
if not code_lines:
|
|
return f"{indentation}{empty_fallback}"
|
|
formatted_lines = []
|
|
for line in code_lines:
|
|
stripped_line = line.lstrip()
|
|
if not stripped_line.endswith("\n"):
|
|
stripped_line += "\n"
|
|
formatted_lines.append(f"{indentation}{stripped_line}")
|
|
return "".join(formatted_lines).rstrip()
|