nodes/pydn-ComfyUI-to-Python-Extension
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Defer generated ComfyUI bootstrap until execution

Browse Source
This commit is contained in:
Peyton DeNiro
2026-03-29 14:40:17 -05:00
parent d751147ed9
commit 441dcfebdf
3 changed files with 68 additions and 44 deletions
Download Patch File Download Diff File Expand all files Collapse all files
comfyui_to_python/generator/render.py
+13 -9
View File
@@ -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=" "
tests/runtime/generated/text-to-image.py
+17 -20
View File
@@ -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(
tests/test_upscale_model_loader_export.py
+38 -15
View File
@@ -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)