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.
This commit is contained in:
Peyton DeNiro
@@ -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",
|
||||
|
||||
@@ -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}")
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
@@ -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()
|
||||
@@ -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
|
||||
@@ -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()
|
||||
@@ -0,0 +1,242 @@
|
||||
import threading
|
||||
import unittest
|
||||
from unittest.mock import MagicMock, patch, call
|
||||
|
||||
from comfyui_to_python.runtime_session import WorkflowSessionRuntime
|
||||
|
||||
|
||||
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 TestWorkflowSessionRuntimeRunCount(unittest.TestCase):
|
||||
"""Tests for run count tracking."""
|
||||
|
||||
def _make_runtime(self):
|
||||
return WorkflowSessionRuntime()
|
||||
|
||||
@patch("comfyui_to_python.runtime_session.WorkflowSessionRuntime._do_run")
|
||||
def test_run_increments_run_count(self, mock_do_run):
|
||||
runtime = self._make_runtime()
|
||||
mock_do_run.return_value = {"result": [1]}
|
||||
runtime.run()
|
||||
self.assertEqual(runtime.run_count, 1)
|
||||
runtime.run()
|
||||
self.assertEqual(runtime.run_count, 2)
|
||||
|
||||
|
||||
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_does_not_clear_bootstrapped(self, mock_do_run):
|
||||
runtime = self._make_runtime()
|
||||
runtime.bootstrapped = True
|
||||
mock_do_run.side_effect = RuntimeError("simulated failure")
|
||||
with self.assertRaises(RuntimeError):
|
||||
runtime.run()
|
||||
self.assertTrue(runtime.bootstrapped)
|
||||
|
||||
@patch("comfyui_to_python.runtime_session.WorkflowSessionRuntime._do_run")
|
||||
def test_exception_preserves_custom_nodes_initialized(self, mock_do_run):
|
||||
runtime = self._make_runtime()
|
||||
runtime.custom_nodes_initialized = True
|
||||
mock_do_run.side_effect = RuntimeError("simulated failure")
|
||||
with self.assertRaises(RuntimeError):
|
||||
runtime.run()
|
||||
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)
|
||||
|
||||
|
||||
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_caches_instances(self):
|
||||
runtime = self._make_runtime()
|
||||
node_class = MagicMock()
|
||||
runtime._ensure_node_instances({"TestNode": node_class})
|
||||
self.assertIn("TestNode", runtime.node_instances)
|
||||
|
||||
def test_ensure_node_instances_creates_instance(self):
|
||||
runtime = self._make_runtime()
|
||||
node_class = MagicMock()
|
||||
runtime._ensure_node_instances({"TestNode": node_class})
|
||||
node_class.assert_called_once()
|
||||
|
||||
def test_ensure_node_instances_reuses_cached(self):
|
||||
runtime = self._make_runtime()
|
||||
node_class = MagicMock()
|
||||
runtime._ensure_node_instances({"TestNode": node_class})
|
||||
runtime._ensure_node_instances({"TestNode": node_class})
|
||||
node_class.assert_called_once()
|
||||
|
||||
|
||||
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_with_unload(self, mock_gc, mock_cleanup):
|
||||
runtime = self._make_runtime()
|
||||
runtime._do_close(unload_models=True)
|
||||
mock_cleanup.assert_called_once_with(unload_models=True)
|
||||
|
||||
@patch("comfyui_to_python.runtime_session.cleanup_comfyui_runtime")
|
||||
@patch("comfyui_to_python.runtime_session.gc")
|
||||
def test_do_close_calls_cleanup_without_unload(self, mock_gc, mock_cleanup):
|
||||
runtime = self._make_runtime()
|
||||
runtime._do_close(unload_models=False)
|
||||
mock_cleanup.assert_called_once_with(unload_models=False)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -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 = True)", 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()
|
||||
@@ -0,0 +1,107 @@
|
||||
import threading
|
||||
import unittest
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
from comfyui_to_python import WorkflowSession
|
||||
|
||||
|
||||
class TestWorkflowSessionInit(unittest.TestCase):
|
||||
"""Tests for WorkflowSession initialization."""
|
||||
|
||||
def test_init_default_cleanup_policy(self):
|
||||
session = WorkflowSession()
|
||||
self.assertEqual(session._runtime._cleanup_policy, "session")
|
||||
|
||||
def test_init_accepts_per_run_policy(self):
|
||||
session = WorkflowSession(cleanup_policy="per_run")
|
||||
self.assertEqual(session._runtime._cleanup_policy, "per_run")
|
||||
|
||||
def test_init_accepts_manual_policy(self):
|
||||
session = WorkflowSession(cleanup_policy="manual")
|
||||
self.assertEqual(session._runtime._cleanup_policy, "manual")
|
||||
|
||||
def test_init_accepts_reset_every_n_runs(self):
|
||||
session = WorkflowSession(reset_every_n_runs=5)
|
||||
self.assertEqual(session._runtime._reset_every_n_runs, 5)
|
||||
|
||||
def test_init_wraps_runtime(self):
|
||||
session = WorkflowSession()
|
||||
self.assertIsInstance(session._runtime, object)
|
||||
self.assertTrue(hasattr(session._runtime, "run"))
|
||||
self.assertTrue(hasattr(session._runtime, "close"))
|
||||
self.assertTrue(hasattr(session._runtime, "clear_runtime_cache"))
|
||||
|
||||
|
||||
class TestWorkflowSessionRun(unittest.TestCase):
|
||||
"""Tests for WorkflowSession.run() delegation."""
|
||||
|
||||
@patch("comfyui_to_python.runtime_session.WorkflowSessionRuntime._do_run")
|
||||
def test_run_delegates_to_runtime(self, mock_do_run):
|
||||
mock_do_run.return_value = {"result": [1, 2, 3]}
|
||||
session = WorkflowSession()
|
||||
result = session.run()
|
||||
self.assertEqual(result, {"result": [1, 2, 3]})
|
||||
|
||||
@patch("comfyui_to_python.runtime_session.WorkflowSessionRuntime._do_run")
|
||||
def test_run_with_workflow_data(self, mock_do_run):
|
||||
workflow_data = {"1": {"class_type": "TestNode", "inputs": {"x": 1}}}
|
||||
node_mappings = {"TestNode": MagicMock()}
|
||||
session = WorkflowSession()
|
||||
session.run(workflow_data=workflow_data, node_class_mappings=node_mappings)
|
||||
mock_do_run.assert_called_once()
|
||||
|
||||
@patch("comfyui_to_python.runtime_session.WorkflowSessionRuntime._do_run")
|
||||
def test_run_with_extra_pnginfo(self, mock_do_run):
|
||||
extra_pnginfo = {"workflow": {"metadata": "test"}}
|
||||
session = WorkflowSession()
|
||||
session.run(extra_pnginfo=extra_pnginfo)
|
||||
mock_do_run.assert_called_once()
|
||||
|
||||
@patch("comfyui_to_python.runtime_session.WorkflowSessionRuntime._do_run")
|
||||
def test_run_passes_all_parameters(self, mock_do_run):
|
||||
workflow_data = {"1": {"class_type": "TestNode", "inputs": {}}}
|
||||
node_mappings = {"TestNode": MagicMock()}
|
||||
extra_pnginfo = {"workflow": {}}
|
||||
session = WorkflowSession()
|
||||
session.run(
|
||||
workflow_data=workflow_data,
|
||||
node_class_mappings=node_mappings,
|
||||
extra_pnginfo=extra_pnginfo,
|
||||
)
|
||||
mock_do_run.assert_called_once()
|
||||
|
||||
|
||||
class TestWorkflowSessionClose(unittest.TestCase):
|
||||
"""Tests for WorkflowSession.close() delegation."""
|
||||
|
||||
def test_close_delegates_to_runtime(self):
|
||||
session = WorkflowSession()
|
||||
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_close_delegates_with_unload_false(self):
|
||||
session = WorkflowSession()
|
||||
with patch.object(session._runtime, "close") as mock_close:
|
||||
session.close(unload_models=False)
|
||||
mock_close.assert_called_once_with(unload_models=False)
|
||||
|
||||
def test_run_after_close_raises(self):
|
||||
session = WorkflowSession()
|
||||
session.close()
|
||||
with self.assertRaises(RuntimeError):
|
||||
session.run()
|
||||
|
||||
|
||||
class TestWorkflowSessionClearCache(unittest.TestCase):
|
||||
"""Tests for WorkflowSession.clear_runtime_cache() delegation."""
|
||||
|
||||
@patch("comfyui_to_python.runtime_session.WorkflowSessionRuntime.clear_runtime_cache")
|
||||
def test_clear_cache_delegates_to_runtime(self, mock_clear):
|
||||
session = WorkflowSession()
|
||||
session.clear_runtime_cache()
|
||||
mock_clear.assert_called_once()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,270 @@
|
||||
import json
|
||||
import unittest
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
from comfyui_to_python.runtime_session import WorkflowSessionRuntime
|
||||
|
||||
|
||||
class StubNode:
|
||||
"""Stub node class for testing."""
|
||||
|
||||
FUNCTION = "execute"
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {"required": {"value": ("STRING",)}}
|
||||
|
||||
def execute(self, value):
|
||||
return (f"result:{value}",)
|
||||
|
||||
|
||||
class TestWorkflowSessionRuntimeDoRun(unittest.TestCase):
|
||||
"""Tests for WorkflowSessionRuntime._do_run() workflow execution."""
|
||||
|
||||
@patch("comfyui_to_python.runtime_session.WorkflowSessionRuntime._do_close")
|
||||
def test_run_executes_workflow_nodes(self, mock_do_close):
|
||||
workflow_data = {
|
||||
"1": {
|
||||
"class_type": "StubNode",
|
||||
"inputs": {"value": "test"},
|
||||
}
|
||||
}
|
||||
node_mappings = {"StubNode": StubNode}
|
||||
runtime = WorkflowSessionRuntime()
|
||||
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 = WorkflowSessionRuntime()
|
||||
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 = WorkflowSessionRuntime()
|
||||
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_unknown):
|
||||
workflow_data = {
|
||||
"1": {
|
||||
"class_type": "UnknownNode",
|
||||
"inputs": {"value": "test"},
|
||||
}
|
||||
}
|
||||
node_mappings = {"StubNode": StubNode}
|
||||
runtime = WorkflowSessionRuntime()
|
||||
runtime._workflow_data = workflow_data
|
||||
runtime._node_class_mappings = node_mappings
|
||||
|
||||
result = runtime.run()
|
||||
|
||||
# Unknown node is skipped, but the call still returns empty dict
|
||||
# since no node instances were created
|
||||
self.assertEqual(result, {})
|
||||
|
||||
def test_run_node_returns_tuple_is_converted_to_list(self):
|
||||
workflow_data = {
|
||||
"1": {
|
||||
"class_type": "StubNode",
|
||||
"inputs": {"value": "multi"},
|
||||
}
|
||||
}
|
||||
node_mappings = {"StubNode": StubNode}
|
||||
runtime = WorkflowSessionRuntime()
|
||||
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_preserves_extra_pnginfo(self):
|
||||
workflow_data = {
|
||||
"1": {
|
||||
"class_type": "StubNode",
|
||||
"inputs": {"value": "test"},
|
||||
}
|
||||
}
|
||||
node_mappings = {"StubNode": StubNode}
|
||||
extra_pnginfo = {"workflow": {"test": True}}
|
||||
runtime = WorkflowSessionRuntime()
|
||||
runtime._workflow_data = workflow_data
|
||||
runtime._node_class_mappings = node_mappings
|
||||
runtime._extra_pnginfo = extra_pnginfo
|
||||
runtime.node_instances = {"StubNode": StubNode()}
|
||||
runtime._node_classes = {"StubNode": StubNode}
|
||||
|
||||
result = runtime.run()
|
||||
|
||||
self.assertIn("1", result)
|
||||
|
||||
def test_run_persists_parameters_across_calls(self):
|
||||
workflow_data = {
|
||||
"1": {
|
||||
"class_type": "StubNode",
|
||||
"inputs": {"value": "persist"},
|
||||
}
|
||||
}
|
||||
node_mappings = {"StubNode": StubNode}
|
||||
runtime = WorkflowSessionRuntime()
|
||||
|
||||
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):
|
||||
workflow_data = {
|
||||
"1": {
|
||||
"class_type": "StubNode",
|
||||
"inputs": {"value": "test"},
|
||||
}
|
||||
}
|
||||
node_mappings = {"StubNode": StubNode}
|
||||
runtime = WorkflowSessionRuntime()
|
||||
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()
|
||||
|
||||
# Session should still be usable
|
||||
node_mappings["StubNode"] = StubNode
|
||||
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_run_increments_run_count(self, mock_ensure):
|
||||
workflow_data = {
|
||||
"1": {
|
||||
"class_type": "StubNode",
|
||||
"inputs": {"value": "test"},
|
||||
}
|
||||
}
|
||||
node_mappings = {"StubNode": StubNode}
|
||||
runtime = WorkflowSessionRuntime()
|
||||
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, 2)
|
||||
|
||||
@patch(
|
||||
"comfyui_to_python.runtime_session.WorkflowSessionRuntime._ensure_node_instances"
|
||||
)
|
||||
def test_session_policy_does_not_clear_cache(self, mock_ensure):
|
||||
workflow_data = {
|
||||
"1": {
|
||||
"class_type": "StubNode",
|
||||
"inputs": {"value": "test"},
|
||||
}
|
||||
}
|
||||
node_mappings = {"StubNode": StubNode}
|
||||
runtime = WorkflowSessionRuntime(cleanup_policy="session")
|
||||
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):
|
||||
workflow_data = {
|
||||
"1": {
|
||||
"class_type": "StubNode",
|
||||
"inputs": {"value": "test"},
|
||||
}
|
||||
}
|
||||
node_mappings = {"StubNode": StubNode}
|
||||
runtime = WorkflowSessionRuntime(cleanup_policy="per_run")
|
||||
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):
|
||||
workflow_data = {
|
||||
"1": {
|
||||
"class_type": "StubNode",
|
||||
"inputs": {"value": "test"},
|
||||
}
|
||||
}
|
||||
node_mappings = {"StubNode": StubNode}
|
||||
runtime = WorkflowSessionRuntime(cleanup_policy="manual")
|
||||
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()
|
||||
|
||||
def test_node_instances_are_cached(self):
|
||||
workflow_data = {
|
||||
"1": {
|
||||
"class_type": "StubNode",
|
||||
"inputs": {"value": "test"},
|
||||
}
|
||||
}
|
||||
node_mappings = {"StubNode": StubNode}
|
||||
runtime = WorkflowSessionRuntime()
|
||||
runtime._workflow_data = workflow_data
|
||||
runtime._node_class_mappings = node_mappings
|
||||
|
||||
with patch(
|
||||
"comfyui_to_python.runtime_session.WorkflowSessionRuntime._do_run",
|
||||
return_value={"1": ["result:test"]},
|
||||
):
|
||||
runtime.run()
|
||||
runtime.run()
|
||||
|
||||
# Node class should only be instantiated once
|
||||
self.assertIn("StubNode", runtime.node_instances)
|
||||
self.assertEqual(len(runtime.node_instances), 1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user