From 441dcfebdf5b098395af0bf83efb9f30724680f6 Mon Sep 17 00:00:00 2001 From: Peyton DeNiro <25550995+pydn@users.noreply.github.com> Date: Sun, 29 Mar 2026 14:40:17 -0500 Subject: [PATCH] Defer generated ComfyUI bootstrap until execution --- comfyui_to_python/generator/render.py | 22 ++++++---- tests/runtime/generated/text-to-image.py | 37 ++++++++-------- tests/test_upscale_model_loader_export.py | 53 ++++++++++++++++------- 3 files changed, 68 insertions(+), 44 deletions(-) diff --git a/comfyui_to_python/generator/render.py b/comfyui_to_python/generator/render.py index a9fe5ae..e3c6238 100644 --- a/comfyui_to_python/generator/render.py +++ b/comfyui_to_python/generator/render.py @@ -47,10 +47,6 @@ class WorkflowRenderer: "import sys", "from typing import Sequence, Mapping, Any, Union", ] + func_strings - static_imports.append("\n# Runtime support") - static_imports.append( - "\nbootstrap_comfyui_runtime()\nimport torch\nadd_extra_model_paths()\n" - ) if plan.custom_nodes: static_imports.append(f"\n{inspect.getsource(import_custom_nodes)}\n") @@ -76,17 +72,25 @@ class WorkflowRenderer: "extra_pnginfo = build_extra_pnginfo()", ] - execution_section = ["# Workflow execution"] + execution_section = [ + "# Workflow execution", + "def main():", + " bootstrap_comfyui_runtime()", + " import torch", + "", + " add_extra_model_paths()", + ] + if custom_nodes_call: + execution_section.append(f" {custom_nodes_call}") if imports_code: - execution_section.extend(["# Node imports", *imports_code, ""]) + execution_section.extend(["", " # Node imports"]) + execution_section.extend(f" {line}" for line in imports_code) execution_section.extend( [ - "def main():", + "", " with torch.inference_mode():", ] ) - if custom_nodes_call: - execution_section.append(f" {custom_nodes_call}") execution_section.extend( self.build_function_body( plan.special_functions_code, "pass", indentation=" " diff --git a/tests/runtime/generated/text-to-image.py b/tests/runtime/generated/text-to-image.py index 68568a2..5f9b0e0 100644 --- a/tests/runtime/generated/text-to-image.py +++ b/tests/runtime/generated/text-to-image.py @@ -105,14 +105,6 @@ def bootstrap_comfyui_runtime() -> None: os.environ["OCL_SET_SVM_SIZE"] = "262144" -# Runtime support - -bootstrap_comfyui_runtime() -import torch - -add_extra_model_paths() - - # Workflow data def build_workflow() -> dict[str, Any]: return { @@ -169,20 +161,25 @@ workflow = build_workflow() prompt = json.loads(json.dumps(workflow)) extra_pnginfo = build_extra_pnginfo() + # Workflow execution -# Node imports -from nodes import ( - CLIPTextEncode, - CheckpointLoaderSimple, - EmptyLatentImage, - KSampler, - NODE_CLASS_MAPPINGS, - SaveImage, - VAEDecode, -) - - def main(): + bootstrap_comfyui_runtime() + import torch + + add_extra_model_paths() + + # Node imports + from nodes import ( + CLIPTextEncode, + CheckpointLoaderSimple, + EmptyLatentImage, + KSampler, + NODE_CLASS_MAPPINGS, + SaveImage, + VAEDecode, + ) + with torch.inference_mode(): checkpointloadersimple = CheckpointLoaderSimple() checkpointloadersimple_1 = checkpointloadersimple.load_checkpoint( diff --git a/tests/test_upscale_model_loader_export.py b/tests/test_upscale_model_loader_export.py index 5101c20..d5d2643 100644 --- a/tests/test_upscale_model_loader_export.py +++ b/tests/test_upscale_model_loader_export.py @@ -1,4 +1,5 @@ import json +import sys import tempfile import unittest from io import StringIO @@ -122,7 +123,7 @@ class UpscaleModelLoaderExportTest(unittest.TestCase): self.assertTrue(callable(run)) self.assertTrue(callable(main)) - def test_export_places_comfyui_bootstrap_before_torch_import(self): + def test_export_defers_comfyui_bootstrap_until_main(self): workflow = { "1": { "class_type": "LoadImage", @@ -147,19 +148,41 @@ class UpscaleModelLoaderExportTest(unittest.TestCase): self.assertIn("import comfy.options", generated) self.assertIn("comfy.options.enable_args_parsing()", generated) self.assertIn("import cuda_malloc", generated) - self.assertLess( - generated.index("def bootstrap_comfyui_runtime()"), - generated.index("import torch"), - ) - self.assertLess( - generated.index("bootstrap_comfyui_runtime()"), - generated.index("import torch"), - ) - self.assertLess( - generated.index("import cuda_malloc"), - generated.index("import torch"), + self.assertNotIn("\nbootstrap_comfyui_runtime()\n", generated) + self.assertIn("def main():\n bootstrap_comfyui_runtime()\n import torch", generated) + self.assertLess(generated.index("def bootstrap_comfyui_runtime()"), generated.index("def main():")) + self.assertLess(generated.index("bootstrap_comfyui_runtime()"), generated.index("import torch")) + self.assertLess(generated.index("import cuda_malloc"), generated.index("import torch")) + + def test_generated_module_import_does_not_parse_cli_args(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("\nimport torch\n", generated) + + globals_dict = {"__name__": "generated_workflow_module"} + with patch.object(sys, "argv", ["generated_workflow.py", "--wrapper-flag"]): + exec(generated, globals_dict) + + self.assertTrue(callable(globals_dict["main"])) + def test_upscale_workflow_uses_direct_upscale_model_loader_init(self): workflow = { "1": { @@ -196,8 +219,8 @@ class UpscaleModelLoaderExportTest(unittest.TestCase): generated = output.getvalue() - self.assertIn("from nodes import NODE_CLASS_MAPPINGS", generated) - self.assertIn(f"from {LoadImage.__module__} import (", generated) + self.assertIn(" from nodes import NODE_CLASS_MAPPINGS", generated) + self.assertIn(f" from {LoadImage.__module__} import (", generated) self.assertIn("LoadImage,", generated) self.assertIn("UpscaleModelLoader,", generated) self.assertIn("ImageUpscaleWithModel,", generated) @@ -384,13 +407,13 @@ class UpscaleModelLoaderExportTest(unittest.TestCase): generated = output.getvalue() self.assertIn("# Imports", generated) - self.assertIn("# Runtime support", generated) self.assertIn("# Workflow data", generated) self.assertIn("# Workflow execution", generated) self.assertIn("# Entrypoint", generated) self.assertIn("def build_workflow()", generated) self.assertIn("def build_extra_pnginfo()", generated) self.assertIn("def main()", generated) + self.assertIn("bootstrap_comfyui_runtime()", generated) self.assertNotIn("def initialize_workflow()", generated) self.assertNotIn("def run_once(", generated) self.assertIn("with torch.inference_mode():", generated)