diff --git a/comfyui_to_python/generator/render.py b/comfyui_to_python/generator/render.py index 0b95a89..4614649 100644 --- a/comfyui_to_python/generator/render.py +++ b/comfyui_to_python/generator/render.py @@ -58,11 +58,11 @@ class WorkflowRenderer: "import os", "import random", "import sys", + "import threading", "from typing import Sequence, Mapping, Any, Union", ] + func_strings - if plan.custom_nodes: - static_imports.append(f"\n{inspect.getsource(import_custom_nodes)}\n") + static_imports.append(f"\n{inspect.getsource(import_custom_nodes)}\n") return static_imports @@ -182,6 +182,13 @@ class WorkflowRenderer: 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:", @@ -214,17 +221,31 @@ class WorkflowRenderer: " 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():", ] + 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( @@ -240,6 +261,26 @@ class WorkflowRenderer: ) 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:", @@ -249,39 +290,20 @@ class WorkflowRenderer: ] ) - # 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()", - ] - ) + # 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 diff --git a/tests/runtime/run_runtime_validation.py b/tests/runtime/run_runtime_validation.py index 235f99a..2e5447b 100644 --- a/tests/runtime/run_runtime_validation.py +++ b/tests/runtime/run_runtime_validation.py @@ -320,6 +320,12 @@ def parse_args() -> argparse.Namespace: "--generated-path", help=argparse.SUPPRESS, ) + parser.add_argument( + "--execution-mode", + default="oneshot", + choices=("oneshot", "session"), + help=argparse.SUPPRESS, + ) args = parser.parse_args() if not args.internal_export and not args.tier: parser.error("--tier is required unless --internal-export is used.") @@ -401,6 +407,7 @@ def export_workflow( fixture: FixtureConfig, tier: str, runtime_path: str, + execution_mode: str = "oneshot", ) -> tuple[str, str]: from comfyui_to_python import ComfyUItoPython @@ -409,6 +416,7 @@ def export_workflow( kwargs = { "workflow": workflow, "output_file": output, + "execution_mode": execution_mode, } if tier == "fast" and fixture.fast_mapping_factory is not None: kwargs["node_class_mappings"] = fixture.fast_mapping_factory() @@ -479,6 +487,41 @@ def export_workflow_in_runtime_env(fixture: FixtureConfig, runtime_path: str) -> return generated_path.read_text(encoding="utf-8") +def export_session_workflow_in_runtime_env( + workflow_json: str, + execution_mode: str = "session", +) -> str: + """Export a session workflow via subprocess in the ComfyUI runtime env. + + Unlike ``export_workflow_in_runtime_env`` this does not require a fixture + config – it receives workflow JSON directly and re-enters the runtime + interpreter so ``ComfyUItoPython`` can import ComfyUI's nodes. + """ + runtime_path = os.environ.get("COMFYUI_PATH", "") + runtime_python = get_runtime_python(runtime_path) + + with tempfile.NamedTemporaryFile( + suffix=".json", mode="w", delete=False, encoding="utf-8" + ) as wf: + wf.write(workflow_json) + wf_path = wf.name + + try: + temp_fixture = FixtureConfig( + name="session-mode-export", + path=Path(wf_path), + ) + _, generated = export_workflow( + fixture=temp_fixture, + tier="runtime", + runtime_path=runtime_path, + execution_mode=execution_mode, + ) + return generated + finally: + os.unlink(wf_path) + + def validate_generated_python(generated_code: str, fixture_name: str) -> None: try: ast.parse(generated_code) @@ -671,8 +714,10 @@ def main() -> int: fixture=fixture, tier="runtime", runtime_path=os.environ.get("COMFYUI_PATH", ""), + execution_mode=args.execution_mode, ) output_path.write_text(generated_code, encoding="utf-8") + print(generated_code, end="") return 0 try: diff --git a/tests/test_session_e2e.py b/tests/test_session_e2e.py new file mode 100644 index 0000000..85a0972 --- /dev/null +++ b/tests/test_session_e2e.py @@ -0,0 +1,370 @@ +"""End-to-end tests for session execution mode feature. + +Covers: +- Code generation correctness (session vs oneshot) +- Runtime execution of generated scripts +- Multi-run session behavior +""" + +import ast +import json +import os +import subprocess +import sys +import tempfile +import unittest +from io import StringIO +from pathlib import Path + +ROOT = Path(__file__).resolve().parents[1] +COMFYUI_PATH = os.environ.get("COMFYUI_PATH", str(ROOT.parent / "ComfyUI")) + +# Minimal mock nodes for export +class KSamplerMock: + CATEGORY = "sampling" + FUNCTION = "sample" + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "seed": ("INT", {"default": 0}), + "steps": ("INT", {"default": 20}), + "cfg": ("FLOAT", {"default": 8.0}), + "sampler_name": (["euler", "heun"],), + "scheduler": (["normal"],), + "denoise": ("FLOAT", {"default": 1.0}), + "model": ("MODEL",), + "positive": ("CONDITIONING",), + "negative": ("CONDITIONING",), + "latent_image": ("LATENT",), + } + } + def sample(self, seed, steps, cfg, sampler_name, scheduler, denoise, model, positive, negative, latent_image): + return ({"samples": latent_image},) + +class CheckpointLoaderMock: + CATEGORY = "loaders" + FUNCTION = "load_checkpoint" + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "ckpt_name": ("STRING",), + } + } + def load_checkpoint(self, ckpt_name): + return (None, None, None) + +class VAEDecodeMock: + CATEGORY = "latent" + FUNCTION = "decode" + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "samples": ("LATENT",), + "vae": ("VAE",), + } + } + def decode(self, samples, vae): + return (samples,) + +class CLIPTextEncodeMock: + CATEGORY = "conditioning" + FUNCTION = "encode" + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "text": ("STRING",), + "clip": ("CLIP",), + } + } + def encode(self, text, clip): + return ([],) + +class EmptyLatentImageMock: + CATEGORY = "latent" + FUNCTION = "generate" + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "width": ("INT", {"default": 512}), + "height": ("INT", {"default": 512}), + "batch_size": ("INT", {"default": 1}), + } + } + def generate(self, width, height, batch_size): + return ({"samples": {}},) + +class SaveImageMock: + CATEGORY = "image" + FUNCTION = "save" + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "images": ("IMAGE",), + "filename_prefix": ("STRING",), + } + } + def save(self, images, filename_prefix): + return () + +NODENAMES = { + "CheckpointLoaderSimple": CheckpointLoaderMock, + "CLIPTextEncode": CLIPTextEncodeMock, + "KSampler": KSamplerMock, + "VAEDecode": VAEDecodeMock, + "EmptyLatentImage": EmptyLatentImageMock, + "SaveImage": SaveImageMock, +} + +TEXT_TO_IMAGE_WORKFLOW = { + "1": { + "class_type": "CheckpointLoaderSimple", + "inputs": {"ckpt_name": "v1-5-pruned-emaonly-fp16.safetensors"}, + }, + "2": { + "class_type": "CLIPTextEncode", + "inputs": { + "text": "a small cottage in a meadow, soft daylight", + "clip": ["1", 1], + }, + }, + "3": { + "class_type": "CLIPTextEncode", + "inputs": { + "text": "blurry, low quality", + "clip": ["1", 1], + }, + }, + "4": { + "class_type": "EmptyLatentImage", + "inputs": {"width": 512, "height": 512, "batch_size": 1}, + }, + "5": { + "class_type": "KSampler", + "inputs": { + "seed": 1, "steps": 4, "cfg": 7, "sampler_name": "euler", + "scheduler": "normal", "denoise": 1, + "model": ["1", 0], "positive": ["2", 0], + "negative": ["3", 0], "latent_image": ["4", 0], + }, + }, + "6": { + "class_type": "VAEDecode", + "inputs": {"samples": ["5", 0], "vae": ["1", 2]}, + }, + "7": { + "class_type": "SaveImage", + "inputs": {"filename_prefix": "E2E_session_mode", "images": ["6", 0]}, + }, +} + + +def _get_runtime_python(): + """Get the ComfyUI Python interpreter for running generated scripts.""" + rt_python = Path(COMFYUI_PATH) / ".venv" / "bin" / "python" + if rt_python.is_file(): + return str(rt_python) + return sys.executable + + +def _export_session_workflow_in_runtime_env( + workflow_json, + execution_mode="session", +): + """Export workflow via subprocess in ComfyUI runtime env. + + Re-enters the runtime interpreter so ``ComfyUItoPython`` can import + ComfyUI's nodes.py (which requires torch, not available in the test venv). + """ + runtime_python = _get_runtime_python() + env = os.environ.copy() + env["COMFYUI_PATH"] = COMFYUI_PATH + env["PYTHONPATH"] = os.pathsep.join([str(ROOT), env.get("PYTHONPATH", "")]).rstrip( + os.pathsep + ) + + tmp_path = tempfile.mktemp(suffix=".py") + + try: + result = subprocess.run( + [ + runtime_python, + str(Path(__file__).resolve().parents[0] / "runtime" / "run_runtime_validation.py"), + "--internal-export", + "text-to-image", + "--execution-mode", + execution_mode, + "--generated-path", + tmp_path, + ], + cwd=ROOT, + env=env, + capture_output=True, + text=True, + ) + if result.returncode != 0: + output = (result.stderr or result.stdout or "").strip() + if "Missing runtime dependency" not in output and "ModuleNotFoundError" not in output: + classification = "repo regression" + else: + classification = "environment/setup failure" + raise RuntimeError( + f"Runtime export failed: [{classification}] {output}" + ) + # Read generated code from the temp file (stdout is polluted by + # ComfyUI runtime prints such as the sys.path line). + return Path(tmp_path).read_text() + finally: + os.unlink(tmp_path) + + +def _export_workflow(execution_mode="oneshot"): + """Export workflow to a string using ComfyUItoPython (unit tests, mock nodes).""" + from comfyui_to_python import ComfyUItoPython + output = StringIO() + ComfyUItoPython( + workflow=json.dumps(TEXT_TO_IMAGE_WORKFLOW), + output_file=output, + node_class_mappings=NODENAMES, + execution_mode=execution_mode, + ) + return output.getvalue() + + +class SessionCodeGenerationTest(unittest.TestCase): + """Unit tests for session mode code generation.""" + + def test_oneshot_code_has_bootstrap_helpers_no_session(self): + """Oneshot mode: bootstrap/cleanup helpers present, no WorkflowSession class.""" + generated = _export_workflow("oneshot") + ast.parse(generated) + self.assertIn("bootstrap_comfyui_runtime()", generated) + self.assertIn("cleanup_comfyui_runtime(", generated) + self.assertNotIn("class WorkflowSession", generated) + self.assertNotIn("session.run()", generated) + + def test_session_code_has_workflow_session_class_and_session_methods(self): + """Session mode: WorkflowSession class with run() and close() present.""" + generated = _export_workflow("session") + ast.parse(generated) + self.assertIn("class WorkflowSession", generated) + self.assertIn("def run(self)", generated) + self.assertIn("def close(self, unload_models", generated) + self.assertIn("session.run()", generated) + self.assertIn("session.close(", generated) + + +class SessionModeExecutionTest(unittest.TestCase): + """E2E tests for session mode script execution.""" + + @unittest.skipIf(not Path(COMFYUI_PATH).is_dir(), "ComfyUI checkout not available") + def test_oneshot_e2e_text_to_image(self): + """Oneshot mode: generate and run text-to-image workflow, verify PNG output.""" + with tempfile.TemporaryDirectory() as tmpdir: + script_path = Path(tmpdir) / "oneshot.py" + env = os.environ.copy() + env["COMFYUI_PATH"] = COMFYUI_PATH + env["PYTHONPATH"] = os.pathsep.join([str(ROOT), env.get("PYTHONPATH", "")]) + runtime_py = _get_runtime_python() + + # Export inside ComfyUI env so node mappings resolve + generated = _export_session_workflow_in_runtime_env( + json.dumps(TEXT_TO_IMAGE_WORKFLOW), + execution_mode="oneshot", + ) + script_path.write_text(generated) + + # Run + result = subprocess.run( + [runtime_py, str(script_path), "--cpu"], + cwd=ROOT, + env=env, + capture_output=True, + text=True, + timeout=300, + ) + self.assertEqual(result.returncode, 0, f"stderr: {result.stderr}") + + # Check output + output_dir = Path(COMFYUI_PATH) / "output" + new_outputs = list(output_dir.glob("E2E_text_to_image*.png")) + self.assertTrue(len(new_outputs) > 0, "No PNG output produced") + + @unittest.skipIf(not Path(COMFYUI_PATH).is_dir(), "ComfyUI checkout not available") + def test_session_e2e_text_to_image(self): + """Session mode: generate session-mode script, verify WorkflowSession present and run.""" + with tempfile.TemporaryDirectory() as tmpdir: + script_path = Path(tmpdir) / "session.py" + env = os.environ.copy() + env["COMFYUI_PATH"] = COMFYUI_PATH + env["PYTHONPATH"] = os.pathsep.join([str(ROOT), env.get("PYTHONPATH", "")]) + runtime_py = _get_runtime_python() + + # Export inside ComfyUI env + generated = _export_session_workflow_in_runtime_env( + json.dumps(TEXT_TO_IMAGE_WORKFLOW), + execution_mode="session", + ) + script_path.write_text(generated) + + # Verify code structure + self.assertIn("class WorkflowSession", generated) + + # Run + result = subprocess.run( + [runtime_py, str(script_path), "--cpu"], + cwd=ROOT, + env=env, + capture_output=True, + text=True, + timeout=300, + ) + self.assertEqual(result.returncode, 0, f"stderr: {result.stderr}") + + # Check output + output_dir = Path(COMFYUI_PATH) / "output" + new_outputs = list(output_dir.glob("E2E_text_to_image*.png")) + self.assertTrue(len(new_outputs) > 0, "No PNG output produced") + + @unittest.skipIf(not Path(COMFYUI_PATH).is_dir(), "ComfyUI checkout not available") + def test_session_e2e_multiple_runs(self): + """Session mode: generate script that calls session.run() 3x in a row, verify no crash.""" + with tempfile.TemporaryDirectory() as tmpdir: + script_path = Path(tmpdir) / "session_multi.py" + env = os.environ.copy() + env["COMFYUI_PATH"] = COMFYUI_PATH + env["PYTHONPATH"] = os.pathsep.join([str(ROOT), env.get("PYTHONPATH", "")]) + runtime_py = _get_runtime_python() + + # Export session mode inside ComfyUI env + generated = _export_session_workflow_in_runtime_env( + json.dumps(TEXT_TO_IMAGE_WORKFLOW), + execution_mode="session", + ) + + # Modify to run 3x + modified = generated.replace( + " session.run()", + " session.run()\n session.run()\n session.run()", + ) + script_path.write_text(modified) + + # Run + result = subprocess.run( + [runtime_py, str(script_path), "--cpu"], + cwd=ROOT, + env=env, + capture_output=True, + text=True, + timeout=600, + ) + self.assertEqual(result.returncode, 0, f"stderr: {result.stderr}") + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_session_renderer.py b/tests/test_session_renderer.py index 14acd06..0e8aae4 100644 --- a/tests/test_session_renderer.py +++ b/tests/test_session_renderer.py @@ -137,7 +137,7 @@ class SessionRendererTest(unittest.TestCase): ) generated = output.getvalue() - self.assertIn("def close(self, unload_models: bool = True)", generated) + self.assertIn("def close(self, unload_models: bool | None = None)", generated) def test_session_mode_oneshot_generates_same_code(self): workflow = {