From 1630a1ab791dab5444ebfba372ea06e69d5b38ec Mon Sep 17 00:00:00 2001 From: Peyton DeNiro <25550995+pydn@users.noreply.github.com> Date: Sun, 29 Mar 2026 14:08:18 -0500 Subject: [PATCH 1/6] Strengthen test suite robustness --- tests/README.md | 16 +- .../runtime/reused-node-class-branches.json | 27 ++ .../runtime/secondary-output-selection.json | 17 + tests/runtime/generated/text-to-image.py | 235 ++++++++++++++ .../runtime/generated/upscale-model-loader.py | 200 ++++++++++++ tests/runtime/run_runtime_validation.py | 50 +++ tests/test_runtime_validation_harness.py | 305 ++++++++++++++++++ tests/test_upscale_model_loader_export.py | 112 ++++++- 8 files changed, 957 insertions(+), 5 deletions(-) create mode 100644 tests/fixtures/runtime/reused-node-class-branches.json create mode 100644 tests/fixtures/runtime/secondary-output-selection.json create mode 100644 tests/runtime/generated/text-to-image.py create mode 100644 tests/runtime/generated/upscale-model-loader.py create mode 100644 tests/test_runtime_validation_harness.py diff --git a/tests/README.md b/tests/README.md index cc7bca6..af5dfd2 100644 --- a/tests/README.md +++ b/tests/README.md @@ -5,6 +5,8 @@ This directory contains two test paths: - unit tests for exporter behavior - runtime validation for committed workflow fixtures +Use the lightweight path for routine contributor validation, and use the runtime path when you need end-to-end confidence against a real ComfyUI checkout. + ## Prerequisites - Run commands from the repo root. @@ -30,12 +32,18 @@ The current runtime-capable fixtures require these model files in the target Com ## Unit Tests -Run the current unit test module: +Run the exporter-focused unit test module: ```bash uv run python -m unittest tests.test_upscale_model_loader_export ``` +Run the runtime-harness unit test module: + +```bash +uv run python -m unittest tests.test_runtime_validation_harness +``` + Run all `unittest`-discoverable tests under `tests`: ```bash @@ -49,6 +57,7 @@ The runtime harness lives at `tests/runtime/run_runtime_validation.py`. ### Fast Tier Fast tier validates export behavior against committed fixtures without requiring a full ComfyUI runtime for every fixture. +This is the recommended default validation lane for routine changes. Run all fast-tier compatible fixtures: @@ -65,6 +74,7 @@ uv run python tests/runtime/run_runtime_validation.py --tier fast --fixture unsa ### Runtime Tier Runtime tier exports inside a real ComfyUI checkout and executes generated Python for runtime-capable fixtures. +This is the heavier validation lane for changes that need end-to-end runtime confidence. Run all runtime-capable fixtures: @@ -92,11 +102,15 @@ Current committed fixtures: - `text-to-image` - `unsafe-kwargs` - `subgraph-identifiers` +- `reused-node-class-branches` +- `secondary-output-selection` Notes: - `--tier runtime` only runs fixtures marked runtime-capable. - `--tier fast` only runs fixtures with local test mappings. +- `reused-node-class-branches` protects repeated node-class usage and branch wiring in the fast tier. +- `secondary-output-selection` protects non-zero output index wiring in the fast tier. ## Troubleshooting diff --git a/tests/fixtures/runtime/reused-node-class-branches.json b/tests/fixtures/runtime/reused-node-class-branches.json new file mode 100644 index 0000000..27c61ec --- /dev/null +++ b/tests/fixtures/runtime/reused-node-class-branches.json @@ -0,0 +1,27 @@ +{ + "1:10": { + "class_type": "PassthroughText", + "inputs": { + "text": "left branch" + } + }, + "1:20": { + "class_type": "PassthroughText", + "inputs": { + "text": "right branch" + } + }, + "2": { + "class_type": "JoinText", + "inputs": { + "left": [ + "1:10", + 0 + ], + "right": [ + "1:20", + 0 + ] + } + } +} diff --git a/tests/fixtures/runtime/secondary-output-selection.json b/tests/fixtures/runtime/secondary-output-selection.json new file mode 100644 index 0000000..db9a055 --- /dev/null +++ b/tests/fixtures/runtime/secondary-output-selection.json @@ -0,0 +1,17 @@ +{ + "1": { + "class_type": "SplitText", + "inputs": { + "text": "alpha|omega" + } + }, + "2": { + "class_type": "PassthroughText", + "inputs": { + "text": [ + "1", + 1 + ] + } + } +} diff --git a/tests/runtime/generated/text-to-image.py b/tests/runtime/generated/text-to-image.py new file mode 100644 index 0000000..68568a2 --- /dev/null +++ b/tests/runtime/generated/text-to-image.py @@ -0,0 +1,235 @@ +# Imports +import json +import os +import random +import sys +from typing import Sequence, Mapping, Any, Union + + +def get_value_at_index(obj: Union[Sequence, Mapping], index: int) -> Any: + """Return a sequence or mapping result item by index.""" + try: + return obj[index] + except KeyError: + return obj["result"][index] + + +def get_comfyui_path() -> str: + """Return the configured ComfyUI path, preferring COMFYUI_PATH when set.""" + comfyui_path = os.environ.get("COMFYUI_PATH") + if comfyui_path: + return comfyui_path + return find_path("ComfyUI") + + +def find_path(name: str, path: str = None) -> str: + """Recursively search parent folders until the named entry is found.""" + if path is None: + path = os.getcwd() + + if name in os.listdir(path): + path_name = os.path.join(path, name) + print(f"{name} found: {path_name}") + return path_name + + parent_directory = os.path.dirname(path) + if parent_directory == path: + return None + + return find_path(name, parent_directory) + + +def add_comfyui_directory_to_sys_path() -> None: + """Add the ComfyUI checkout to sys.path.""" + comfyui_path = get_comfyui_path() + if comfyui_path is not None and os.path.isdir(comfyui_path): + if comfyui_path in sys.path: + sys.path.remove(comfyui_path) + sys.path.insert(0, comfyui_path) + print(f"'{comfyui_path}' added to sys.path") + + +def add_extra_model_paths() -> None: + """Load ComfyUI extra model paths configuration when available.""" + try: + from main import load_extra_path_config + except ImportError: + print( + "Could not import load_extra_path_config from main.py. Looking in utils.extra_config instead." + ) + from utils.extra_config import load_extra_path_config + + extra_model_paths = find_path("extra_model_paths.yaml") + if extra_model_paths is not None: + load_extra_path_config(extra_model_paths) + else: + print("Could not find the extra_model_paths config file.") + + +def bootstrap_comfyui_runtime() -> None: + """Mirror the allocator-related ComfyUI startup steps before torch import.""" + add_comfyui_directory_to_sys_path() + + import comfy.options + + comfy.options.enable_args_parsing() + + from comfy.cli_args import args + + if os.name == "nt": + os.environ["MIMALLOC_PURGE_DELAY"] = "0" + + if args.default_device is not None: + default_dev = args.default_device + devices = list(range(32)) + devices.remove(default_dev) + devices.insert(0, default_dev) + devices = ",".join(map(str, devices)) + os.environ["CUDA_VISIBLE_DEVICES"] = str(devices) + os.environ["HIP_VISIBLE_DEVICES"] = str(devices) + + if args.cuda_device is not None: + os.environ["CUDA_VISIBLE_DEVICES"] = str(args.cuda_device) + os.environ["HIP_VISIBLE_DEVICES"] = str(args.cuda_device) + os.environ["ASCEND_RT_VISIBLE_DEVICES"] = str(args.cuda_device) + + if args.oneapi_device_selector is not None: + os.environ["ONEAPI_DEVICE_SELECTOR"] = args.oneapi_device_selector + + if args.deterministic and "CUBLAS_WORKSPACE_CONFIG" not in os.environ: + os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8" + + import cuda_malloc + + if "rocm" in cuda_malloc.get_torch_version_noimport(): + 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 { + "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_text_to_image", "images": ["6", 0]}, + }, + } + + +def build_extra_pnginfo() -> dict[str, Any] | None: + return None + + +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(): + with torch.inference_mode(): + checkpointloadersimple = CheckpointLoaderSimple() + checkpointloadersimple_1 = checkpointloadersimple.load_checkpoint( + ckpt_name="v1-5-pruned-emaonly-fp16.safetensors" + ) + cliptextencode = CLIPTextEncode() + cliptextencode_2 = cliptextencode.encode( + text="a small cottage in a meadow, soft daylight", + clip=get_value_at_index(checkpointloadersimple_1, 1), + ) + cliptextencode_3 = cliptextencode.encode( + text="blurry, low quality", + clip=get_value_at_index(checkpointloadersimple_1, 1), + ) + emptylatentimage = EmptyLatentImage() + emptylatentimage_4 = emptylatentimage.generate( + width=512, height=512, batch_size=1 + ) + ksampler = KSampler() + vaedecode = VAEDecode() + saveimage = SaveImage() + for q in range(1): + node_5_seed = prompt["5"]["inputs"]["seed"] = random.randint(1, 2**64) + ksampler_5 = ksampler.sample( + seed=node_5_seed, + steps=4, + cfg=7, + sampler_name="euler", + scheduler="normal", + denoise=1, + model=get_value_at_index(checkpointloadersimple_1, 0), + positive=get_value_at_index(cliptextencode_2, 0), + negative=get_value_at_index(cliptextencode_3, 0), + latent_image=get_value_at_index(emptylatentimage_4, 0), + ) + vaedecode_6 = vaedecode.decode( + samples=get_value_at_index(ksampler_5, 0), + vae=get_value_at_index(checkpointloadersimple_1, 2), + ) + saveimage_7 = saveimage.save_images( + filename_prefix="E2E_text_to_image", + images=get_value_at_index(vaedecode_6, 0), + prompt=prompt, + extra_pnginfo=extra_pnginfo, + ) + + +# Entrypoint +if __name__ == "__main__": + main() diff --git a/tests/runtime/generated/upscale-model-loader.py b/tests/runtime/generated/upscale-model-loader.py new file mode 100644 index 0000000..b7b0a09 --- /dev/null +++ b/tests/runtime/generated/upscale-model-loader.py @@ -0,0 +1,200 @@ +# Imports +import json +import os +import random +import sys +from typing import Sequence, Mapping, Any, Union + + +def get_value_at_index(obj: Union[Sequence, Mapping], index: int) -> Any: + """Return a sequence or mapping result item by index.""" + try: + return obj[index] + except KeyError: + return obj["result"][index] + + +def get_comfyui_path() -> str: + """Return the configured ComfyUI path, preferring COMFYUI_PATH when set.""" + comfyui_path = os.environ.get("COMFYUI_PATH") + if comfyui_path: + return comfyui_path + return find_path("ComfyUI") + + +def find_path(name: str, path: str = None) -> str: + """Recursively search parent folders until the named entry is found.""" + if path is None: + path = os.getcwd() + + if name in os.listdir(path): + path_name = os.path.join(path, name) + print(f"{name} found: {path_name}") + return path_name + + parent_directory = os.path.dirname(path) + if parent_directory == path: + return None + + return find_path(name, parent_directory) + + +def add_comfyui_directory_to_sys_path() -> None: + """Add the ComfyUI checkout to sys.path.""" + comfyui_path = get_comfyui_path() + if comfyui_path is not None and os.path.isdir(comfyui_path): + if comfyui_path in sys.path: + sys.path.remove(comfyui_path) + sys.path.insert(0, comfyui_path) + print(f"'{comfyui_path}' added to sys.path") + + +def add_extra_model_paths() -> None: + """Load ComfyUI extra model paths configuration when available.""" + try: + from main import load_extra_path_config + except ImportError: + print( + "Could not import load_extra_path_config from main.py. Looking in utils.extra_config instead." + ) + from utils.extra_config import load_extra_path_config + + extra_model_paths = find_path("extra_model_paths.yaml") + if extra_model_paths is not None: + load_extra_path_config(extra_model_paths) + else: + print("Could not find the extra_model_paths config file.") + + +def bootstrap_comfyui_runtime() -> None: + """Mirror the allocator-related ComfyUI startup steps before torch import.""" + add_comfyui_directory_to_sys_path() + + import comfy.options + + comfy.options.enable_args_parsing() + + from comfy.cli_args import args + + if os.name == "nt": + os.environ["MIMALLOC_PURGE_DELAY"] = "0" + + if args.default_device is not None: + default_dev = args.default_device + devices = list(range(32)) + devices.remove(default_dev) + devices.insert(0, default_dev) + devices = ",".join(map(str, devices)) + os.environ["CUDA_VISIBLE_DEVICES"] = str(devices) + os.environ["HIP_VISIBLE_DEVICES"] = str(devices) + + if args.cuda_device is not None: + os.environ["CUDA_VISIBLE_DEVICES"] = str(args.cuda_device) + os.environ["HIP_VISIBLE_DEVICES"] = str(args.cuda_device) + os.environ["ASCEND_RT_VISIBLE_DEVICES"] = str(args.cuda_device) + + if args.oneapi_device_selector is not None: + os.environ["ONEAPI_DEVICE_SELECTOR"] = args.oneapi_device_selector + + if args.deterministic and "CUBLAS_WORKSPACE_CONFIG" not in os.environ: + os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8" + + import cuda_malloc + + if "rocm" in cuda_malloc.get_torch_version_noimport(): + os.environ["OCL_SET_SVM_SIZE"] = "262144" + + +# Runtime support + +bootstrap_comfyui_runtime() +import torch + +add_extra_model_paths() + + +def import_custom_nodes() -> None: + """Initialize ComfyUI custom nodes in the exporter runtime.""" + comfyui_path = get_comfyui_path() + if comfyui_path and comfyui_path not in sys.path: + sys.path.insert(0, comfyui_path) + + import asyncio + import execution + from nodes import init_extra_nodes + + if comfyui_path in sys.path: + sys.path.remove(comfyui_path) + sys.path.insert(0, comfyui_path) + + import server + + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + server_instance = server.PromptServer(loop) + execution.PromptQueue(server_instance) + asyncio.run(init_extra_nodes()) + + +# Workflow data +def build_workflow() -> dict[str, Any]: + return { + "1": {"class_type": "LoadImage", "inputs": {"image": "e2e_upscale_input.png"}}, + "2": { + "class_type": "UpscaleModelLoader", + "inputs": {"model_name": "RealESRGAN_x4plus.safetensors"}, + }, + "3": { + "class_type": "ImageUpscaleWithModel", + "inputs": {"upscale_model": ["2", 0], "image": ["1", 0]}, + }, + "4": { + "class_type": "SaveImage", + "inputs": { + "filename_prefix": "E2E_upscale_model_loader", + "images": ["3", 0], + }, + }, + } + + +def build_extra_pnginfo() -> dict[str, Any] | None: + return None + + +workflow = build_workflow() +prompt = json.loads(json.dumps(workflow)) +extra_pnginfo = build_extra_pnginfo() + +# Workflow execution +# Node imports +from nodes import LoadImage, NODE_CLASS_MAPPINGS, SaveImage + + +def main(): + with torch.inference_mode(): + import_custom_nodes() + loadimage = LoadImage() + loadimage_1 = loadimage.load_image(image="e2e_upscale_input.png") + upscalemodelloader = NODE_CLASS_MAPPINGS["UpscaleModelLoader"]() + upscalemodelloader_2 = upscalemodelloader.EXECUTE_NORMALIZED( + model_name="RealESRGAN_x4plus.safetensors" + ) + imageupscalewithmodel = NODE_CLASS_MAPPINGS["ImageUpscaleWithModel"]() + saveimage = SaveImage() + for q in range(1): + imageupscalewithmodel_3 = imageupscalewithmodel.EXECUTE_NORMALIZED( + upscale_model=get_value_at_index(upscalemodelloader_2, 0), + image=get_value_at_index(loadimage_1, 0), + ) + saveimage_4 = saveimage.save_images( + filename_prefix="E2E_upscale_model_loader", + images=get_value_at_index(imageupscalewithmodel_3, 0), + prompt=prompt, + extra_pnginfo=extra_pnginfo, + ) + + +# Entrypoint +if __name__ == "__main__": + main() diff --git a/tests/runtime/run_runtime_validation.py b/tests/runtime/run_runtime_validation.py index 46b0d84..235f99a 100644 --- a/tests/runtime/run_runtime_validation.py +++ b/tests/runtime/run_runtime_validation.py @@ -109,6 +109,40 @@ class PassthroughText: return (text,) +class SplitText: + CATEGORY = "utils" + FUNCTION = "split" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "text": ("STRING",), + } + } + + def split(self, text): + left, right = text.split("|", 1) + return (left, right) + + +class JoinText: + CATEGORY = "utils" + FUNCTION = "join" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "left": ("STRING",), + "right": ("STRING",), + } + } + + def join(self, left, right): + return (f"{left}::{right}",) + + class UpscaleModelLoader: CATEGORY = "loaders" FUNCTION = "load_model" @@ -234,6 +268,22 @@ FIXTURES = { "PassthroughText": PassthroughText, }, ), + "reused-node-class-branches": FixtureConfig( + name="reused-node-class-branches", + path=FIXTURE_DIR / "reused-node-class-branches.json", + mapping_factory=lambda: { + "PassthroughText": PassthroughText, + "JoinText": JoinText, + }, + ), + "secondary-output-selection": FixtureConfig( + name="secondary-output-selection", + path=FIXTURE_DIR / "secondary-output-selection.json", + mapping_factory=lambda: { + "SplitText": SplitText, + "PassthroughText": PassthroughText, + }, + ), } diff --git a/tests/test_runtime_validation_harness.py b/tests/test_runtime_validation_harness.py new file mode 100644 index 0000000..0e30253 --- /dev/null +++ b/tests/test_runtime_validation_harness.py @@ -0,0 +1,305 @@ +import struct +import tempfile +import unittest +import zlib +from pathlib import Path +from unittest.mock import patch + +from tests.runtime.run_runtime_validation import ( + FixtureConfig, + ModelRequirement, + ValidationFailure, + check_models, + ensure_runtime_path, + execute_generated_python, + load_fixture_names, + parse_png_info, + validate_generated_python, +) + + +def make_png_bytes( + width: int, + height: int, + text_chunks: list[tuple[bytes, bytes]] | None = None, +) -> bytes: + def chunk(chunk_type: bytes, data: bytes) -> bytes: + crc = zlib.crc32(chunk_type + data) & 0xFFFFFFFF + return ( + struct.pack(">I", len(data)) + + chunk_type + + data + + struct.pack(">I", crc) + ) + + ihdr = chunk(b"IHDR", struct.pack(">IIBBBBB", width, height, 8, 2, 0, 0, 0)) + text_chunks = text_chunks or [] + idat = chunk( + b"IDAT", + zlib.compress(b"\x00" + (b"\x00\x00\x00" * width)), + ) + return b"".join( + [ + b"\x89PNG\r\n\x1a\n", + ihdr, + *[chunk(chunk_type, payload) for chunk_type, payload in text_chunks], + idat, + chunk(b"IEND", b""), + ] + ) + + +class RuntimeValidationHarnessTest(unittest.TestCase): + def test_load_fixture_names_all_returns_registered_names(self): + fixture_names = load_fixture_names("all") + + self.assertGreaterEqual(len(fixture_names), 4) + self.assertEqual(fixture_names[0], "upscale-model-loader") + self.assertIn("unsafe-kwargs", fixture_names) + + def test_load_fixture_names_rejects_unknown_fixture(self): + with self.assertRaises(ValidationFailure) as context: + load_fixture_names("missing-fixture") + + self.assertEqual(context.exception.classification, "fixture bug") + self.assertIn("Unknown fixture", context.exception.message) + + def test_ensure_runtime_path_fast_tier_returns_env_without_validation(self): + with patch.dict("os.environ", {"COMFYUI_PATH": "/does/not/exist"}, clear=False): + runtime_path = ensure_runtime_path("fast") + + self.assertEqual(runtime_path, "/does/not/exist") + + @patch("tests.runtime.run_runtime_validation.get_comfyui_path", return_value="") + def test_ensure_runtime_path_runtime_tier_requires_valid_checkout(self, _mock_path): + with self.assertRaises(ValidationFailure) as context: + ensure_runtime_path("runtime") + + self.assertEqual(context.exception.classification, "environment/setup failure") + self.assertIn("Could not find a valid ComfyUI checkout", context.exception.message) + + def test_check_models_returns_only_missing_requirements(self): + with tempfile.TemporaryDirectory() as tmpdir: + runtime_path = Path(tmpdir) + present_dir = runtime_path / "models" / "checkpoints" + present_dir.mkdir(parents=True) + (present_dir / "present.safetensors").write_text("ok", encoding="utf-8") + fixture = FixtureConfig( + name="model-checks", + path=Path("unused.json"), + model_requirements=( + ModelRequirement( + filename="present.safetensors", + relative_dir="models/checkpoints", + source_url="https://example.invalid/present", + ), + ModelRequirement( + filename="missing.safetensors", + relative_dir="models/checkpoints", + source_url="https://example.invalid/missing", + ), + ), + ) + + missing = check_models(fixture, str(runtime_path)) + + self.assertEqual([item.filename for item in missing], ["missing.safetensors"]) + + def test_validate_generated_python_accepts_valid_python(self): + validate_generated_python("value = 1\n", "valid-fixture") + + def test_validate_generated_python_reports_syntax_regression(self): + with self.assertRaises(ValidationFailure) as context: + validate_generated_python("def broken(:\n", "broken-fixture") + + self.assertEqual(context.exception.classification, "repo regression") + self.assertIn("broken-fixture", context.exception.message) + + def test_parse_png_info_reads_dimensions_and_text_chunks(self): + compressed_text = zlib.compress(b"workflow data") + png_bytes = make_png_bytes( + width=3, + height=2, + text_chunks=[ + (b"tEXt", b"prompt\x00hello"), + (b"zTXt", b"workflow\x00\x00" + compressed_text), + (b"iTXt", b"comment\x00\x00\x00\x00\x00unicode text"), + ], + ) + + with tempfile.TemporaryDirectory() as tmpdir: + image_path = Path(tmpdir) / "artifact.png" + image_path.write_bytes(png_bytes) + + width, height, metadata = parse_png_info(image_path) + + self.assertEqual((width, height), (3, 2)) + self.assertEqual(metadata["prompt"], "hello") + self.assertEqual(metadata["workflow"], "workflow data") + self.assertEqual(metadata["comment"], "unicode text") + + def test_parse_png_info_rejects_non_png_files(self): + with tempfile.TemporaryDirectory() as tmpdir: + image_path = Path(tmpdir) / "artifact.bin" + image_path.write_bytes(b"not-a-png") + + with self.assertRaises(ValidationFailure) as context: + parse_png_info(image_path) + + self.assertEqual(context.exception.classification, "environment/setup failure") + self.assertIn("Expected PNG output", context.exception.message) + + def test_parse_png_info_requires_dimensions(self): + with tempfile.TemporaryDirectory() as tmpdir: + image_path = Path(tmpdir) / "artifact.png" + image_path.write_bytes( + b"\x89PNG\r\n\x1a\n" + + struct.pack(">I", 0) + + b"IEND" + + struct.pack(">I", 0) + ) + + with self.assertRaises(ValidationFailure) as context: + parse_png_info(image_path) + + self.assertEqual(context.exception.classification, "environment/setup failure") + self.assertIn("Could not read PNG dimensions", context.exception.message) + + @patch("tests.runtime.run_runtime_validation.get_runtime_python", return_value="/usr/bin/python") + @patch("tests.runtime.run_runtime_validation.subprocess.run") + def test_execute_generated_python_classifies_missing_torch_as_environment_failure( + self, + mock_run, + _mock_runtime_python, + ): + mock_run.return_value.returncode = 1 + mock_run.return_value.stderr = "ModuleNotFoundError: No module named 'torch'" + mock_run.return_value.stdout = "" + fixture = FixtureConfig( + name="runtime-fixture", + path=Path("unused.json"), + runtime_capable=True, + ) + + with tempfile.TemporaryDirectory() as tmpdir: + (Path(tmpdir) / "output").mkdir() + with self.assertRaises(ValidationFailure) as context: + execute_generated_python("print('hello')\n", fixture, tmpdir) + + self.assertEqual(context.exception.classification, "environment/setup failure") + + @patch("tests.runtime.run_runtime_validation.get_runtime_python", return_value="/usr/bin/python") + @patch("tests.runtime.run_runtime_validation.subprocess.run") + def test_execute_generated_python_classifies_missing_files_as_environment_failure( + self, + mock_run, + _mock_runtime_python, + ): + mock_run.return_value.returncode = 1 + mock_run.return_value.stderr = "No such file or directory: missing.png" + mock_run.return_value.stdout = "" + fixture = FixtureConfig( + name="runtime-fixture", + path=Path("unused.json"), + runtime_capable=True, + ) + + with tempfile.TemporaryDirectory() as tmpdir: + (Path(tmpdir) / "output").mkdir() + with self.assertRaises(ValidationFailure) as context: + execute_generated_python("print('hello')\n", fixture, tmpdir) + + self.assertEqual(context.exception.classification, "environment/setup failure") + + @patch("tests.runtime.run_runtime_validation.get_runtime_python", return_value="/usr/bin/python") + @patch("tests.runtime.run_runtime_validation.subprocess.run") + def test_execute_generated_python_classifies_other_failures_as_repo_regression( + self, + mock_run, + _mock_runtime_python, + ): + mock_run.return_value.returncode = 1 + mock_run.return_value.stderr = "ValueError: broken workflow" + mock_run.return_value.stdout = "" + fixture = FixtureConfig( + name="runtime-fixture", + path=Path("unused.json"), + runtime_capable=True, + ) + + with tempfile.TemporaryDirectory() as tmpdir: + (Path(tmpdir) / "output").mkdir() + with self.assertRaises(ValidationFailure) as context: + execute_generated_python("print('hello')\n", fixture, tmpdir) + + self.assertEqual(context.exception.classification, "repo regression") + + @patch("tests.runtime.run_runtime_validation.get_runtime_python", return_value="/usr/bin/python") + @patch("tests.runtime.run_runtime_validation.subprocess.run") + def test_execute_generated_python_requires_fresh_matching_artifact( + self, + mock_run, + _mock_runtime_python, + ): + mock_run.return_value.returncode = 0 + mock_run.return_value.stderr = "" + mock_run.return_value.stdout = "" + fixture = FixtureConfig( + name="runtime-fixture", + path=Path("unused.json"), + runtime_capable=True, + filename_prefix="expected_prefix", + ) + + with tempfile.TemporaryDirectory() as tmpdir: + (Path(tmpdir) / "output").mkdir() + with self.assertRaises(ValidationFailure) as context: + execute_generated_python("print('hello')\n", fixture, tmpdir) + + self.assertEqual(context.exception.classification, "repo regression") + self.assertIn("did not produce a new output file", context.exception.message) + + @patch("tests.runtime.run_runtime_validation.get_runtime_python", return_value="/usr/bin/python") + @patch("tests.runtime.run_runtime_validation.validate_output_artifact") + @patch("tests.runtime.run_runtime_validation.subprocess.run") + def test_execute_generated_python_validates_newest_matching_artifact( + self, + mock_run, + mock_validate_output, + _mock_runtime_python, + ): + mock_run.return_value.returncode = 0 + mock_run.return_value.stderr = "" + mock_run.return_value.stdout = "" + fixture = FixtureConfig( + name="runtime-fixture", + path=Path("unused.json"), + runtime_capable=True, + filename_prefix="expected_prefix", + ) + + with tempfile.TemporaryDirectory() as tmpdir: + output_dir = Path(tmpdir) / "output" + output_dir.mkdir() + older = output_dir / "expected_prefix_00001_.png" + older.write_bytes(make_png_bytes(1, 1)) + with patch( + "tests.runtime.run_runtime_validation.subprocess.run", + side_effect=self._write_runtime_artifact(mock_run.return_value, output_dir), + ): + execute_generated_python("print('hello')\n", fixture, tmpdir) + + validated_path = mock_validate_output.call_args[0][1] + self.assertEqual(validated_path.name, "expected_prefix_00002_.png") + + @staticmethod + def _write_runtime_artifact(result, output_dir: Path): + def side_effect(*_args, **_kwargs): + (output_dir / "expected_prefix_00002_.png").write_bytes(make_png_bytes(2, 2)) + return result + + return side_effect + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_upscale_model_loader_export.py b/tests/test_upscale_model_loader_export.py index d711c4b..5101c20 100644 --- a/tests/test_upscale_model_loader_export.py +++ b/tests/test_upscale_model_loader_export.py @@ -115,6 +115,51 @@ class HiddenPromptSeedNode: class UpscaleModelLoaderExportTest(unittest.TestCase): + def test_top_level_module_preserves_exporter_entrypoints(self): + from comfyui_to_python import main + + self.assertTrue(callable(ComfyUItoPython)) + self.assertTrue(callable(run)) + self.assertTrue(callable(main)) + + def test_export_places_comfyui_bootstrap_before_torch_import(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.assertIn("def bootstrap_comfyui_runtime()", generated) + 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"), + ) + def test_upscale_workflow_uses_direct_upscale_model_loader_init(self): workflow = { "1": { @@ -196,10 +241,11 @@ class UpscaleModelLoaderExportTest(unittest.TestCase): generated = output.getvalue() - self.assertIn('extra_pnginfo = {', generated) - self.assertIn('"workflow": json.loads(', generated) + self.assertIn("def build_extra_pnginfo()", generated) + self.assertIn('"workflow": {', generated) self.assertIn('"version": 0.4', generated) self.assertIn('"nodes": []', generated) + self.assertNotIn('"workflow": json.loads(', generated) self.assertNotIn('"source": "workflow_api"', generated) def test_export_without_frontend_workflow_leaves_png_workflow_metadata_absent(self): @@ -223,7 +269,8 @@ class UpscaleModelLoaderExportTest(unittest.TestCase): generated = output.getvalue() - self.assertIn("extra_pnginfo = None", generated) + self.assertIn("def build_extra_pnginfo()", generated) + self.assertIn("return None", generated) self.assertNotIn('"workflow": json.loads(', generated) def test_export_preserves_unique_variable_names_for_subgraph_node_ids(self): @@ -296,9 +343,66 @@ class UpscaleModelLoaderExportTest(unittest.TestCase): generated = output_file.read_text(encoding="utf-8") - self.assertIn("extra_pnginfo = None", generated) + self.assertIn("def build_extra_pnginfo()", generated) + self.assertIn("return None", generated) self.assertNotIn('"workflow": json.loads(', generated) + def test_export_structures_generated_script_into_readable_sections(self): + workflow = { + "1": { + "class_type": "LoadImage", + "inputs": { + "image": "example.png", + }, + }, + "2": { + "class_type": "UpscaleModelLoader", + "inputs": { + "model_name": "RealESRGAN_x4plus.safetensors", + }, + }, + "3": { + "class_type": "ImageUpscaleWithModel", + "inputs": { + "upscale_model": ["2", 0], + "image": ["1", 0], + }, + }, + } + + output = StringIO() + ComfyUItoPython( + workflow=json.dumps(workflow), + output_file=output, + node_class_mappings={ + "LoadImage": LoadImage, + "UpscaleModelLoader": UpscaleModelLoader, + "ImageUpscaleWithModel": ImageUpscaleWithModel, + }, + ) + + 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.assertNotIn("def initialize_workflow()", generated) + self.assertNotIn("def run_once(", generated) + self.assertIn("with torch.inference_mode():", generated) + self.assertIn("for q in range(1):", generated) + self.assertIn("workflow = build_workflow()", generated) + self.assertIn("extra_pnginfo = build_extra_pnginfo()", generated) + self.assertNotIn('workflow = json.loads("', generated) + self.assertLess( + generated.index("def build_workflow()"), + generated.index("def main()"), + ) + def test_hidden_metadata_kwargs_follow_function_signature(self): workflow = { "1": { From 9e7e474b05977c5d60ccdb000be50d706a6bd1f7 Mon Sep 17 00:00:00 2001 From: Peyton DeNiro <25550995+pydn@users.noreply.github.com> Date: Sun, 29 Mar 2026 14:09:24 -0500 Subject: [PATCH 2/6] Refactor the exporter into a package-backed module layout. Split the previous monolithic implementation into focused modules for CLI handling, workflow loading, runtime bootstrap, load ordering, planning, rendering, and file I/O. Keep the public compatibility surface intact by preserving the legacy top-level wrapper while adding as the documented entrypoint. This reduces cross-cutting coupling inside the exporter and makes the runtime/bootstrap behavior easier to evolve without carrying unrelated concerns in one file. --- README.md | 10 +- comfyui_to_python.py | 770 +----------------- comfyui_to_python/__init__.py | 55 ++ comfyui_to_python/__main__.py | 5 + comfyui_to_python/app.py | 75 ++ comfyui_to_python/cli.py | 42 + comfyui_to_python/generator/__init__.py | 4 + .../generator/generated_helpers.py | 17 + comfyui_to_python/generator/model.py | 12 + comfyui_to_python/generator/planner.py | 228 ++++++ comfyui_to_python/generator/render.py | 135 +++ comfyui_to_python/io.py | 23 + comfyui_to_python/load_order.py | 43 + comfyui_to_python/node_runtime.py | 133 +++ comfyui_to_python/workflow_loader.py | 19 + comfyui_to_python_utils.py | 140 +--- 16 files changed, 818 insertions(+), 893 deletions(-) create mode 100644 comfyui_to_python/__init__.py create mode 100644 comfyui_to_python/__main__.py create mode 100644 comfyui_to_python/app.py create mode 100644 comfyui_to_python/cli.py create mode 100644 comfyui_to_python/generator/__init__.py create mode 100644 comfyui_to_python/generator/generated_helpers.py create mode 100644 comfyui_to_python/generator/model.py create mode 100644 comfyui_to_python/generator/planner.py create mode 100644 comfyui_to_python/generator/render.py create mode 100644 comfyui_to_python/io.py create mode 100644 comfyui_to_python/load_order.py create mode 100644 comfyui_to_python/node_runtime.py create mode 100644 comfyui_to_python/workflow_loader.py diff --git a/README.md b/README.md index b80470b..d4ce020 100644 --- a/README.md +++ b/README.md @@ -79,18 +79,24 @@ Notes: 3. Run the exporter: ```bash -uv run python comfyui_to_python.py +uv run python -m comfyui_to_python ``` Options: ```bash -uv run python comfyui_to_python.py \ +uv run python -m comfyui_to_python \ --input_file workflow_api.json \ --output_file workflow_api.py \ --queue_size 10 ``` +The legacy wrapper still works if you prefer it: + +```bash +uv run python comfyui_to_python.py +``` + Flags: - `--input_file`: input workflow JSON, default `workflow_api.json` - `--output_file`: output Python file, default `workflow_api.py` diff --git a/comfyui_to_python.py b/comfyui_to_python.py index efae605..4ff8101 100644 --- a/comfyui_to_python.py +++ b/comfyui_to_python.py @@ -1,773 +1,5 @@ -import copy -import glob -import inspect -import json -import keyword -import os -import random -import sys -import re -from typing import Dict, List, Any, Callable, Tuple, TextIO -from argparse import ArgumentParser - -import black - -# Add current directory to Python path -sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) -from comfyui_to_python_utils import ( - import_custom_nodes, - find_path, - get_comfyui_path, - add_comfyui_directory_to_sys_path, - add_extra_model_paths, - get_value_at_index, -) - - -DEFAULT_INPUT_FILE = "workflow_api.json" -DEFAULT_OUTPUT_FILE = "workflow_api.py" -DEFAULT_QUEUE_SIZE = 10 - - -def get_node_class_mappings() -> Dict: - """Load ComfyUI node mappings on demand. - - Tests that inject explicit node mappings should not need a full ComfyUI runtime - just to import this module. - """ - add_comfyui_directory_to_sys_path() - from nodes import NODE_CLASS_MAPPINGS - - return NODE_CLASS_MAPPINGS - - -class FileHandler: - """Handles reading and writing files. - - This class provides methods to read JSON data from an input file and write code to an output file. - """ - - @staticmethod - def read_json_file(file_path: str | TextIO, encoding: str = "utf-8") -> dict: - """ - Reads a JSON file and returns its contents as a dictionary. - - Args: - file_path (str): The path to the JSON file. - - Returns: - dict: The contents of the JSON file as a dictionary. - - Raises: - FileNotFoundError: If the file is not found, it lists all JSON files in the directory of the file path. - ValueError: If the file is not a valid JSON. - """ - - if hasattr(file_path, "read"): - return json.load(file_path) - with open(file_path, "r", encoding="utf-8") as file: - data = json.load(file) - return data - - @staticmethod - def write_code_to_file(file_path: str | TextIO, code: str) -> None: - """Write the specified code to a Python file. - - Args: - file_path (str): The path to the Python file. - code (str): The code to write to the file. - - Returns: - None - """ - if isinstance(file_path, str): - # Extract directory from the filename - directory = os.path.dirname(file_path) - - # If the directory does not exist, create it - if directory and not os.path.exists(directory): - os.makedirs(directory) - - # Save the code to a .py file - with open(file_path, "w", encoding="utf-8") as file: - file.write(code) - else: - file_path.write(code) - - -class LoadOrderDeterminer: - """Determine the load order of each key in the provided dictionary. - - This class places the nodes without node dependencies first, then ensures that any node whose - result is used in another node will be added to the list in the order it should be executed. - - Attributes: - data (Dict): The dictionary for which to determine the load order. - node_class_mappings (Dict): Mappings of node classes. - """ - - def __init__(self, data: Dict, node_class_mappings: Dict): - """Initialize the LoadOrderDeterminer with the given data and node class mappings. - - Args: - data (Dict): The dictionary for which to determine the load order. - node_class_mappings (Dict): Mappings of node classes. - """ - self.data = data - self.node_class_mappings = node_class_mappings - self.visited = {} - self.load_order = [] - self.is_special_function = False - - def determine_load_order(self) -> List[Tuple[str, Dict, bool]]: - """Determine the load order for the given data. - - Returns: - List[Tuple[str, Dict, bool]]: A list of tuples representing the load order. - """ - self._load_special_functions_first() - self.is_special_function = False - for key in self.data: - if key not in self.visited: - self._dfs(key) - return self.load_order - - def _dfs(self, key: str) -> None: - """Depth-First Search function to determine the load order. - - Args: - key (str): The key from which to start the DFS. - - Returns: - None - """ - # Mark the node as visited. - self.visited[key] = True - inputs = self.data[key]["inputs"] - # Loop over each input key. - for input_key, val in inputs.items(): - # If the value is a list and the first item in the list has not been visited yet, - # then recursively apply DFS on the dependency. - if isinstance(val, list) and val[0] not in self.visited: - self._dfs(val[0]) - # Add the key and its corresponding data to the load order list. - self.load_order.append((key, self.data[key], self.is_special_function)) - - def _load_special_functions_first(self) -> None: - """Load functions without dependencies, loaderes, and encoders first. - - Returns: - None - """ - # Iterate over each key in the data to check for loader keys. - for key in self.data: - class_def = self.node_class_mappings[self.data[key]["class_type"]]() - # Check if the class is a loader class or meets specific conditions. - if ( - class_def.CATEGORY == "loaders" - or class_def.FUNCTION in ["encode"] - or not any( - isinstance(val, list) for val in self.data[key]["inputs"].values() - ) - ): - self.is_special_function = True - # If the key has not been visited, perform a DFS from that key. - if key not in self.visited: - self._dfs(key) - - -class CodeGenerator: - """Generates Python code for a workflow based on the load order. - - Attributes: - node_class_mappings (Dict): Mappings of node classes. - base_node_class_mappings (Dict): Base mappings of node classes. - """ - - def __init__(self, node_class_mappings: Dict, base_node_class_mappings: Dict): - """Initialize the CodeGenerator with given node class mappings. - - Args: - node_class_mappings (Dict): Mappings of node classes. - base_node_class_mappings (Dict): Base mappings of node classes. - """ - self.node_class_mappings = node_class_mappings - self.base_node_class_mappings = base_node_class_mappings - - @staticmethod - def sanitize_node_id(node_id: str) -> str: - """Convert node IDs into variable-safe tokens without collapsing separators.""" - sanitized = re.sub(r"[^a-z0-9_]", "_", str(node_id).lower().strip()) - sanitized = re.sub(r"_+", "_", sanitized).strip("_") - if not sanitized: - sanitized = "node" - return sanitized - - def generate_workflow( - self, - load_order: List, - workflow_data: Dict, - metadata_workflow_data: Dict | None = None, - queue_size: int = 10, - ) -> str: - """Generate the execution code based on the load order. - - Args: - load_order (List): A list of tuples representing the load order. - queue_size (int): The number of photos that will be created by the script. - - Returns: - str: Generated execution code as a string. - """ - # Create the necessary data structures to hold imports and generated code - import_statements, executed_variables, special_functions_code, code = ( - {"nodes": {"NODE_CLASS_MAPPINGS"}}, - {}, - [], - [], - ) - # This dictionary will store the names of the objects that we have already initialized - initialized_objects = {} - - custom_nodes = False - # Loop over each dictionary in the load order list - for idx, data, is_special_function in load_order: - # Generate class definition and inputs from the data - inputs, class_type = data["inputs"], data["class_type"] - input_types = self.node_class_mappings[class_type].INPUT_TYPES() - class_def = self.node_class_mappings[class_type]() - - # If required inputs are not present, skip the node as it will break the code if passed through to the script - missing_required_variable = False - if "required" in input_types.keys(): - for required in input_types["required"]: - if required not in inputs.keys(): - missing_required_variable = True - if missing_required_variable: - continue - - # If the class hasn't been initialized yet, initialize it and generate the import statements - if class_type not in initialized_objects: - # No need to use preview image nodes since we are executing the script in a terminal - if class_type == "PreviewImage": - continue - - class_type, import_statement, class_code = self.get_class_info( - class_type - ) - initialized_objects[class_type] = self.clean_variable_name(class_type) - if class_type in self.base_node_class_mappings.keys(): - module_name, import_name = import_statement - import_statements.setdefault(module_name, set()).add(import_name) - if 'NODE_CLASS_MAPPINGS["' in class_code: - custom_nodes = True - special_functions_code.append(class_code) - - # Get all possible parameters for class_def - class_def_params = self.get_function_parameters( - getattr(class_def, class_def.FUNCTION) - ) - no_params = class_def_params is None - - # Remove any keyword arguments from **inputs if they are not in class_def_params - inputs = { - key: value - for key, value in inputs.items() - if no_params or key in class_def_params - } - # Deal with hidden variables - hidden_inputs = input_types.get("hidden", {}) - if ( - "unique_id" in hidden_inputs - and (no_params or "unique_id" in class_def_params) - ): - inputs["unique_id"] = random.randint(1, 2**64) - if "prompt" in hidden_inputs and (no_params or "prompt" in class_def_params): - inputs["prompt"] = {"variable_name": "prompt"} - if "extra_pnginfo" in hidden_inputs and ( - no_params or "extra_pnginfo" in class_def_params - ): - inputs["extra_pnginfo"] = {"variable_name": "extra_pnginfo"} - if "hidden" not in input_types and class_def_params is not None: - if "unique_id" in class_def_params: - inputs["unique_id"] = random.randint(1, 2**64) - - # Create executed variable and generate code - executed_variables[idx] = ( - f"{self.clean_variable_name(class_type)}_" - f"{self.sanitize_node_id(str(idx))}" - ) - inputs = self.update_inputs(inputs, executed_variables) - seed_sync_code = self.create_prompt_seed_sync_code( - idx, inputs, is_special_function - ) - - if is_special_function: - if seed_sync_code: - special_functions_code.extend(seed_sync_code) - special_functions_code.append( - self.create_function_call_code( - initialized_objects[class_type], - class_def.FUNCTION, - executed_variables[idx], - is_special_function, - **inputs, - ) - ) - else: - if seed_sync_code: - code.extend(seed_sync_code) - code.append( - self.create_function_call_code( - initialized_objects[class_type], - class_def.FUNCTION, - executed_variables[idx], - is_special_function, - **inputs, - ) - ) - - # Generate final code by combining imports and code, and wrap them in a main function - final_code = self.assemble_python_code( - import_statements, - special_functions_code, - code, - workflow_data, - metadata_workflow_data, - queue_size, - custom_nodes, - ) - - return final_code - - def create_function_call_code( - self, - obj_name: str, - func: str, - variable_name: str, - is_special_function: bool, - **kwargs, - ) -> str: - """Generate Python code for a function call. - - Args: - obj_name (str): The name of the initialized object. - func (str): The function to be called. - variable_name (str): The name of the variable that the function result should be assigned to. - is_special_function (bool): Determines the code indentation. - **kwargs: The keyword arguments for the function. - - Returns: - str: The generated Python code. - """ - args = ", ".join(self.format_arg(key, value) for key, value in kwargs.items()) - - # Generate the Python code - code = f"{variable_name} = {obj_name}.{func}({args})\n" - - # If the code contains dependencies and is not a loader or encoder, indent the code because it will be placed inside - # of a for loop - if not is_special_function: - code = f"\t{code}" - - return code - - def create_prompt_seed_sync_code( - self, node_id: str, inputs: Dict, is_special_function: bool - ) -> List[str]: - """Generate code that keeps prompt metadata aligned with randomized seeds.""" - seed_sync_lines = [] - for key in ("seed", "noise_seed"): - if key not in inputs: - continue - randomized_seed_variable = ( - f"node_{self.sanitize_node_id(str(node_id))}_{self.clean_variable_name(key)}" - ) - seed_sync_lines.append( - f'{randomized_seed_variable} = prompt["{node_id}"]["inputs"]["{key}"] = random.randint(1, 2**64)' - ) - inputs[key] = {"variable_name": randomized_seed_variable} - - if not seed_sync_lines: - return [] - - indentation = "" if is_special_function else "\t" - return [f"{indentation}{line}\n" for line in seed_sync_lines] - - def format_arg(self, key: str, value: any) -> str: - """Formats arguments based on key and value. - - Args: - key (str): Argument key. - value (any): Argument value. - - Returns: - str: Formatted argument as a string. - """ - value_code = self.format_arg_value(key, value) - if key.isidentifier() and not keyword.iskeyword(key): - return f"{key}={value_code}" - return f"**{{{json.dumps(key)}: {value_code}}}" - - @staticmethod - def format_arg_value(key: str, value: any) -> str: - """Formats an argument value as Python source.""" - if isinstance(value, dict) and "variable_name" in value: - return value["variable_name"] - if key == "noise_seed" or key == "seed": - return "random.randint(1, 2**64)" - if isinstance(value, str): - return json.dumps(value) - return repr(value) - - def assemble_python_code( - self, - import_statements: Dict[str, set], - speical_functions_code: List[str], - code: List[str], - workflow_data: Dict, - metadata_workflow_data: Dict | None, - queue_size: int, - custom_nodes=False, - ) -> str: - """Generates the final code string. - - Args: - import_statements (Dict[str, set]): Import statements grouped by module. - speical_functions_code (List[str]): A list of special functions code strings. - code (List[str]): A list of code strings. - workflow_data (Dict): The API workflow data used for runtime prompt execution. - metadata_workflow_data (Dict | None): The workflow metadata to embed into saved outputs. - queue_size (int): Number of photos that will be generated by the script. - custom_nodes (bool): Whether to include custom nodes in the code. - - Returns: - str: Generated final code as a string. - """ - if metadata_workflow_data is None: - extra_pnginfo_code = "extra_pnginfo = None" - else: - extra_pnginfo_code = ( - "extra_pnginfo = " - '{"workflow": json.loads(' - + json.dumps(json.dumps(metadata_workflow_data)) - + ")}" - ) - - # Get the source code of the utils functions as a string - func_strings = [] - for func in [ - get_value_at_index, - get_comfyui_path, - find_path, - add_comfyui_directory_to_sys_path, - add_extra_model_paths, - ]: - func_strings.append(f"\n{inspect.getsource(func)}") - # Define static import statements required for the script - static_imports = ( - [ - "import json", - "import os", - "import random", - "import sys", - "from typing import Sequence, Mapping, Any, Union", - "import torch", - ] - + func_strings - + [ - "\n\nadd_comfyui_directory_to_sys_path()\nadd_extra_model_paths()\n", - f"workflow = json.loads({json.dumps(json.dumps(workflow_data))})", - "prompt = json.loads(json.dumps(workflow))", - extra_pnginfo_code, - ] - ) - # Check if custom nodes should be included - if custom_nodes: - static_imports.append(f"\n{inspect.getsource(import_custom_nodes)}\n") - custom_nodes = "import_custom_nodes()\n\t" - else: - custom_nodes = "" - # Create import statements for node classes - 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}") - special_functions_body = "\n\t\t".join(speical_functions_code) or "pass" - loop_body = "\n\t\t".join(code) or "\tpass" - # Assemble the main function code, including custom nodes if applicable - main_function_code = ( - "def main():\n\t" - + f"{custom_nodes}with torch.inference_mode():\n\t\t" - + special_functions_body - + f"\n\n\t\tfor q in range({queue_size}):\n\t\t" - + loop_body - ) - # Concatenate all parts to form the final code - final_code = "\n".join( - static_imports - + imports_code - + ["", main_function_code, "", 'if __name__ == "__main__":', "\tmain()"] - ) - # Format the final code according to PEP 8 using the Black library - final_code = black.format_str(final_code, mode=black.Mode()) - - return final_code - - def get_class_info(self, class_type: str) -> Tuple[str, Tuple[str, str], str]: - """Generates and returns necessary information about class type. - - Args: - class_type (str): Class type. - - Returns: - Tuple[str, Tuple[str, str], str]: Updated class type, import statement and initialization code. - """ - class_obj = self.base_node_class_mappings.get(class_type) - module_name = "nodes" - if class_obj is not None: - module_name = class_obj.__module__ - variable_name = self.clean_variable_name(class_type) - is_importable_module = bool( - module_name - and "/" not in module_name - and "\\" not in module_name - and all(part.isidentifier() for part in module_name.split(".")) - ) - if class_type in self.base_node_class_mappings.keys() and is_importable_module: - import_statement = (module_name, class_type) - class_code = f"{variable_name} = {class_type.strip()}()" - else: - import_statement = ("nodes", "NODE_CLASS_MAPPINGS") - class_code = f'{variable_name} = NODE_CLASS_MAPPINGS["{class_type}"]()' - - return class_type, import_statement, class_code - - @staticmethod - def clean_variable_name(class_type: str) -> str: - """ - Remove any characters from variable name that could cause errors running the Python script. - - Args: - class_type (str): Class type. - - Returns: - str: Cleaned variable name with no special characters or spaces - """ - # Convert to lowercase and replace spaces with underscores - clean_name = class_type.lower().strip().replace("-", "_").replace(" ", "_") - - # Remove characters that are not letters, numbers, or underscores - clean_name = re.sub(r"[^a-z0-9_]", "", clean_name) - - # Ensure that it doesn't start with a number - if clean_name[0].isdigit(): - clean_name = "_" + clean_name - - return clean_name - - def get_function_parameters(self, func: Callable) -> List: - """Get the names of a function's parameters. - - Args: - func (Callable): The function whose parameters we want to inspect. - - Returns: - List: A list containing the names of the function's parameters. - """ - signature = inspect.signature(func) - parameters = { - name: param.default if param.default != param.empty else None - for name, param in signature.parameters.items() - } - catch_all = any( - param.kind == inspect.Parameter.VAR_KEYWORD - for param in signature.parameters.values() - ) - return list(parameters.keys()) if not catch_all else None - - def update_inputs(self, inputs: Dict, executed_variables: Dict) -> Dict: - """Update inputs based on the executed variables. - - Args: - inputs (Dict): Inputs dictionary to update. - executed_variables (Dict): Dictionary storing executed variable names. - - Returns: - Dict: Updated inputs dictionary. - """ - for key in inputs.keys(): - if ( - isinstance(inputs[key], list) - and inputs[key][0] in executed_variables.keys() - ): - inputs[key] = { - "variable_name": f"get_value_at_index({executed_variables[inputs[key][0]]}, {inputs[key][1]})" - } - return inputs - - -class ComfyUItoPython: - """Main workflow to generate Python code from a workflow_api.json file. - - Attributes: - input_file (str): Path to the input JSON file. - output_file (str): Path to the output Python file. - queue_size (int): The number of photos that will be created by the script. - node_class_mappings (Dict): Mappings of node classes. - base_node_class_mappings (Dict): Base mappings of node classes. - """ - - def __init__( - self, - workflow: str = "", - frontend_workflow: str | Dict | None = None, - input_file: str = "", - output_file: str | TextIO = "", - queue_size: int = 1, - node_class_mappings: Dict | None = None, - needs_init_custom_nodes: bool = False, - ): - """Initialize the ComfyUItoPython class with the given parameters. Exactly one of workflow or input_file must be specified. - Args: - workflow (str): The workflow's JSON. - input_file (str): Path to the input JSON file. - output_file (str | TextIO): Path to the output file or a file-like object. - queue_size (int): The number of times a workflow will be executed by the script. Defaults to 1. - node_class_mappings (Dict | None): Mappings of node classes. Defaults to the current - ComfyUI NODE_CLASS_MAPPINGS when not provided. - needs_init_custom_nodes (bool): Whether to initialize custom nodes. Defaults to False. - """ - if input_file and workflow: - raise ValueError("Can't provide both input_file and workflow") - elif not input_file and not workflow: - raise ValueError("Needs input_file or workflow") - - if not output_file: - raise ValueError("Needs output_file") - - self.workflow = workflow - self.frontend_workflow = frontend_workflow - self.input_file = input_file - self.output_file = output_file - self.queue_size = queue_size - self.node_class_mappings = ( - node_class_mappings - if node_class_mappings is not None - else get_node_class_mappings() - ) - self.needs_init_custom_nodes = needs_init_custom_nodes - - self.base_node_class_mappings = copy.deepcopy(self.node_class_mappings) - self.execute() - - def execute(self): - """Execute the main workflow to generate Python code. - - Returns: - None - """ - # Step 1: Read JSON data from the input file - if self.input_file: - data = FileHandler.read_json_file(self.input_file) - else: - data = json.loads(self.workflow) - - metadata_workflow_data = None - if self.frontend_workflow: - if isinstance(self.frontend_workflow, str): - metadata_workflow_data = json.loads(self.frontend_workflow) - else: - metadata_workflow_data = self.frontend_workflow - - # Step 2: Initialize extra/custom nodes when requested or when the workflow references - # a node class that is not currently loaded in the runtime. - missing_node_types = { - node_data["class_type"] - for node_data in data.values() - if node_data["class_type"] not in self.node_class_mappings - } - if self.needs_init_custom_nodes or missing_node_types: - import_custom_nodes() - self.base_node_class_mappings = copy.deepcopy(self.node_class_mappings) - - # Step 3: Determine the load order - load_order_determiner = LoadOrderDeterminer(data, self.node_class_mappings) - load_order = load_order_determiner.determine_load_order() - - # Step 4: Generate the workflow code - code_generator = CodeGenerator( - self.node_class_mappings, self.base_node_class_mappings - ) - generated_code = code_generator.generate_workflow( - load_order, - data, - metadata_workflow_data, - queue_size=self.queue_size, - ) - - # Step 5: Write the generated code to a file - FileHandler.write_code_to_file(self.output_file, generated_code) - - print(f"Code successfully generated and written to {self.output_file}") - - -def run( - input_file: str = DEFAULT_INPUT_FILE, - output_file: str = DEFAULT_OUTPUT_FILE, - queue_size: int = DEFAULT_QUEUE_SIZE, -) -> None: - """Generate Python code from a ComfyUI workflow_api.json file. - - Args: - input_file (str): Path to the input JSON file. Defaults to "workflow_api.json". - output_file (str): Path to the output Python file. - Defaults to "workflow_api.py". - queue_size (int): The number of times a workflow will be executed by the script. - Defaults to 1. - - Returns: - None - """ - ComfyUItoPython( - input_file=input_file, - output_file=output_file, - queue_size=queue_size, - needs_init_custom_nodes=True, - ) - - -def main() -> None: - """Main function to generate Python code from a ComfyUI workflow_api.json file.""" - parser = ArgumentParser( - description="Generate Python code from a ComfyUI workflow_api.json file." - ) - parser.add_argument( - "-f", - "--input_file", - type=str, - help="path to the input JSON file", - default=DEFAULT_INPUT_FILE, - ) - parser.add_argument( - "-o", - "--output_file", - type=str, - help="path to the output Python file", - default=DEFAULT_OUTPUT_FILE, - ) - parser.add_argument( - "-q", - "--queue_size", - type=int, - help="number of times the workflow will be executed by default", - default=DEFAULT_QUEUE_SIZE, - ) - pargs = parser.parse_args() - run(**vars(pargs)) - print("Done.") +from comfyui_to_python.cli import main if __name__ == "__main__": - """Run the main function.""" main() diff --git a/comfyui_to_python/__init__.py b/comfyui_to_python/__init__.py new file mode 100644 index 0000000..18aafa1 --- /dev/null +++ b/comfyui_to_python/__init__.py @@ -0,0 +1,55 @@ +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 + + +class ComfyUItoPython: + """Public compatibility facade for the exporter package.""" + + def __init__( + self, + workflow: str = "", + frontend_workflow: str | dict | None = None, + input_file: str = "", + output_file: str | TextIO = "", + queue_size: int = 1, + node_class_mappings: dict | None = None, + needs_init_custom_nodes: bool = False, + ): + self._app = ExportApplication( + workflow=workflow, + frontend_workflow=frontend_workflow, + input_file=input_file, + output_file=output_file, + queue_size=queue_size, + node_class_mappings=node_class_mappings, + needs_init_custom_nodes=needs_init_custom_nodes, + node_mapping_loader=get_node_class_mappings, + custom_node_importer=import_custom_nodes, + ) + self._app.execute() + + +def run( + input_file: str = DEFAULT_INPUT_FILE, + output_file: str = DEFAULT_OUTPUT_FILE, + queue_size: int = DEFAULT_QUEUE_SIZE, +) -> None: + """Generate Python code from a ComfyUI workflow_api.json file.""" + ComfyUItoPython( + input_file=input_file, + output_file=output_file, + queue_size=queue_size, + needs_init_custom_nodes=True, + ) + + +__all__ = [ + "ComfyUItoPython", + "run", + "main", + "get_node_class_mappings", + "import_custom_nodes", +] diff --git a/comfyui_to_python/__main__.py b/comfyui_to_python/__main__.py new file mode 100644 index 0000000..2f05ddc --- /dev/null +++ b/comfyui_to_python/__main__.py @@ -0,0 +1,5 @@ +from .cli import main + + +if __name__ == "__main__": + main() diff --git a/comfyui_to_python/app.py b/comfyui_to_python/app.py new file mode 100644 index 0000000..334d06f --- /dev/null +++ b/comfyui_to_python/app.py @@ -0,0 +1,75 @@ +import copy + +from typing import TextIO + +from .generator.planner import WorkflowPlanner +from .generator.render import WorkflowRenderer +from .io import write_python_output +from .load_order import LoadOrderDeterminer +from .workflow_loader import load_frontend_workflow_data, load_workflow_data + + +class ExportApplication: + """High-level exporter orchestration.""" + + def __init__( + self, + workflow: str = "", + frontend_workflow: str | dict | None = None, + input_file: str = "", + output_file: str | TextIO = "", + queue_size: int = 1, + node_class_mappings: dict | None = None, + needs_init_custom_nodes: bool = False, + node_mapping_loader=None, + custom_node_importer=None, + ): + if input_file and workflow: + raise ValueError("Can't provide both input_file and workflow") + if not input_file and not workflow: + raise ValueError("Needs input_file or workflow") + if not output_file: + raise ValueError("Needs output_file") + + self.workflow = workflow + self.frontend_workflow = frontend_workflow + self.input_file = input_file + self.output_file = output_file + self.queue_size = queue_size + self.node_mapping_loader = node_mapping_loader + self.custom_node_importer = custom_node_importer + self.node_class_mappings = ( + node_class_mappings + if node_class_mappings is not None + else self.node_mapping_loader() + ) + self.needs_init_custom_nodes = needs_init_custom_nodes + self.base_node_class_mappings = copy.deepcopy(self.node_class_mappings) + + def execute(self) -> None: + data = load_workflow_data(self.workflow, self.input_file) + metadata_workflow_data = load_frontend_workflow_data(self.frontend_workflow) + + missing_node_types = { + node_data["class_type"] + for node_data in data.values() + if node_data["class_type"] not in self.node_class_mappings + } + if self.needs_init_custom_nodes or missing_node_types: + self.custom_node_importer() + self.base_node_class_mappings = copy.deepcopy(self.node_class_mappings) + + load_order = LoadOrderDeterminer( + data, self.node_class_mappings + ).determine_load_order() + plan = WorkflowPlanner( + self.node_class_mappings, self.base_node_class_mappings + ).build_plan( + load_order, + data, + metadata_workflow_data, + queue_size=self.queue_size, + ) + generated_code = WorkflowRenderer().render(plan) + write_python_output(self.output_file, generated_code) + print(f"Code successfully generated and written to {self.output_file}") diff --git a/comfyui_to_python/cli.py b/comfyui_to_python/cli.py new file mode 100644 index 0000000..458d4c8 --- /dev/null +++ b/comfyui_to_python/cli.py @@ -0,0 +1,42 @@ +from argparse import ArgumentParser + +DEFAULT_INPUT_FILE = "workflow_api.json" +DEFAULT_OUTPUT_FILE = "workflow_api.py" +DEFAULT_QUEUE_SIZE = 10 + + +def build_argument_parser() -> ArgumentParser: + parser = ArgumentParser( + description="Generate Python code from a ComfyUI workflow_api.json file." + ) + parser.add_argument( + "-f", + "--input_file", + type=str, + help="path to the input JSON file", + default=DEFAULT_INPUT_FILE, + ) + parser.add_argument( + "-o", + "--output_file", + type=str, + help="path to the output Python file", + default=DEFAULT_OUTPUT_FILE, + ) + parser.add_argument( + "-q", + "--queue_size", + type=int, + help="number of times the workflow will be executed by default", + default=DEFAULT_QUEUE_SIZE, + ) + return parser + + +def main() -> None: + from . import run + + parser = build_argument_parser() + pargs = parser.parse_args() + run(**vars(pargs)) + print("Done.") diff --git a/comfyui_to_python/generator/__init__.py b/comfyui_to_python/generator/__init__.py new file mode 100644 index 0000000..c7be5af --- /dev/null +++ b/comfyui_to_python/generator/__init__.py @@ -0,0 +1,4 @@ +from .planner import WorkflowPlanner +from .render import WorkflowRenderer + +__all__ = ["WorkflowPlanner", "WorkflowRenderer"] diff --git a/comfyui_to_python/generator/generated_helpers.py b/comfyui_to_python/generator/generated_helpers.py new file mode 100644 index 0000000..88f5f2f --- /dev/null +++ b/comfyui_to_python/generator/generated_helpers.py @@ -0,0 +1,17 @@ +from ..node_runtime import ( + add_comfyui_directory_to_sys_path, + add_extra_model_paths, + bootstrap_comfyui_runtime, + find_path, + get_comfyui_path, + get_value_at_index, +) + +__all__ = [ + "add_comfyui_directory_to_sys_path", + "add_extra_model_paths", + "bootstrap_comfyui_runtime", + "find_path", + "get_comfyui_path", + "get_value_at_index", +] diff --git a/comfyui_to_python/generator/model.py b/comfyui_to_python/generator/model.py new file mode 100644 index 0000000..ffe5e89 --- /dev/null +++ b/comfyui_to_python/generator/model.py @@ -0,0 +1,12 @@ +from dataclasses import dataclass + + +@dataclass(frozen=True) +class GenerationPlan: + import_statements: dict[str, set[str]] + special_functions_code: list[str] + loop_code: list[str] + workflow_data: dict + metadata_workflow_data: dict | None + queue_size: int + custom_nodes: bool diff --git a/comfyui_to_python/generator/planner.py b/comfyui_to_python/generator/planner.py new file mode 100644 index 0000000..23922d9 --- /dev/null +++ b/comfyui_to_python/generator/planner.py @@ -0,0 +1,228 @@ +import inspect +import json +import keyword +import random +import re +from typing import Any, Callable + +from .model import GenerationPlan + + +class WorkflowPlanner: + """Convert ordered workflow nodes into a renderer-ready generation plan.""" + + def __init__(self, node_class_mappings: dict, base_node_class_mappings: dict): + self.node_class_mappings = node_class_mappings + self.base_node_class_mappings = base_node_class_mappings + + @staticmethod + def sanitize_node_id(node_id: str) -> str: + sanitized = re.sub(r"[^a-z0-9_]", "_", str(node_id).lower().strip()) + sanitized = re.sub(r"_+", "_", sanitized).strip("_") + if not sanitized: + sanitized = "node" + return sanitized + + @staticmethod + def clean_variable_name(class_type: str) -> str: + clean_name = class_type.lower().strip().replace("-", "_").replace(" ", "_") + clean_name = re.sub(r"[^a-z0-9_]", "", clean_name) + if clean_name[0].isdigit(): + clean_name = "_" + clean_name + return clean_name + + def build_plan( + self, + load_order: list[tuple[str, dict, bool]], + workflow_data: dict, + metadata_workflow_data: dict | None = None, + queue_size: int = 10, + ) -> GenerationPlan: + import_statements = {"nodes": {"NODE_CLASS_MAPPINGS"}} + executed_variables = {} + special_functions_code = [] + code = [] + initialized_objects = {} + custom_nodes = False + + for idx, data, is_special_function in load_order: + inputs, class_type = data["inputs"], data["class_type"] + input_types = self.node_class_mappings[class_type].INPUT_TYPES() + class_def = self.node_class_mappings[class_type]() + + missing_required_variable = False + if "required" in input_types.keys(): + for required in input_types["required"]: + if required not in inputs.keys(): + missing_required_variable = True + if missing_required_variable: + continue + + if class_type not in initialized_objects: + if class_type == "PreviewImage": + continue + + class_type, import_statement, class_code = self.get_class_info( + class_type + ) + initialized_objects[class_type] = self.clean_variable_name(class_type) + if class_type in self.base_node_class_mappings.keys(): + module_name, import_name = import_statement + import_statements.setdefault(module_name, set()).add(import_name) + if 'NODE_CLASS_MAPPINGS["' in class_code: + custom_nodes = True + special_functions_code.append(class_code) + + class_def_params = self.get_function_parameters( + getattr(class_def, class_def.FUNCTION) + ) + no_params = class_def_params is None + + inputs = { + key: value + for key, value in inputs.items() + if no_params or key in class_def_params + } + + hidden_inputs = input_types.get("hidden", {}) + if ( + "unique_id" in hidden_inputs + and (no_params or "unique_id" in class_def_params) + ): + inputs["unique_id"] = random.randint(1, 2**64) + if "prompt" in hidden_inputs and (no_params or "prompt" in class_def_params): + inputs["prompt"] = {"variable_name": "prompt"} + if "extra_pnginfo" in hidden_inputs and ( + no_params or "extra_pnginfo" in class_def_params + ): + inputs["extra_pnginfo"] = {"variable_name": "extra_pnginfo"} + if "hidden" not in input_types and class_def_params is not None: + if "unique_id" in class_def_params: + inputs["unique_id"] = random.randint(1, 2**64) + + executed_variables[idx] = ( + f"{self.clean_variable_name(class_type)}_" + f"{self.sanitize_node_id(str(idx))}" + ) + inputs = self.update_inputs(inputs, executed_variables) + seed_sync_code = self.create_prompt_seed_sync_code( + idx, inputs, is_special_function + ) + + target_lines = special_functions_code if is_special_function else code + if seed_sync_code: + target_lines.extend(seed_sync_code) + target_lines.append( + self.create_function_call_code( + initialized_objects[class_type], + class_def.FUNCTION, + executed_variables[idx], + is_special_function, + **inputs, + ) + ) + + return GenerationPlan( + import_statements=import_statements, + special_functions_code=special_functions_code, + loop_code=code, + workflow_data=workflow_data, + metadata_workflow_data=metadata_workflow_data, + queue_size=queue_size, + custom_nodes=custom_nodes, + ) + + def create_function_call_code( + self, + obj_name: str, + func: str, + variable_name: str, + is_special_function: bool, + **kwargs, + ) -> str: + args = ", ".join(self.format_arg(key, value) for key, value in kwargs.items()) + code = f"{variable_name} = {obj_name}.{func}({args})\n" + if not is_special_function: + code = f"\t{code}" + return code + + def create_prompt_seed_sync_code( + self, node_id: str, inputs: dict, is_special_function: bool + ) -> list[str]: + seed_sync_lines = [] + for key in ("seed", "noise_seed"): + if key not in inputs: + continue + randomized_seed_variable = ( + f"node_{self.sanitize_node_id(str(node_id))}_{self.clean_variable_name(key)}" + ) + seed_sync_lines.append( + f'{randomized_seed_variable} = prompt["{node_id}"]["inputs"]["{key}"] = random.randint(1, 2**64)' + ) + inputs[key] = {"variable_name": randomized_seed_variable} + + if not seed_sync_lines: + return [] + + indentation = "" if is_special_function else "\t" + return [f"{indentation}{line}\n" for line in seed_sync_lines] + + def format_arg(self, key: str, value: Any) -> str: + value_code = self.format_arg_value(key, value) + if key.isidentifier() and not keyword.iskeyword(key): + return f"{key}={value_code}" + return f"**{{{json.dumps(key)}: {value_code}}}" + + @staticmethod + def format_arg_value(key: str, value: Any) -> str: + if isinstance(value, dict) and "variable_name" in value: + return value["variable_name"] + if key == "noise_seed" or key == "seed": + return "random.randint(1, 2**64)" + if isinstance(value, str): + return json.dumps(value) + return repr(value) + + def get_class_info(self, class_type: str) -> tuple[str, tuple[str, str], str]: + class_obj = self.base_node_class_mappings.get(class_type) + module_name = "nodes" + if class_obj is not None: + module_name = class_obj.__module__ + variable_name = self.clean_variable_name(class_type) + is_importable_module = bool( + module_name + and "/" not in module_name + and "\\" not in module_name + and all(part.isidentifier() for part in module_name.split(".")) + ) + if class_type in self.base_node_class_mappings.keys() and is_importable_module: + import_statement = (module_name, class_type) + class_code = f"{variable_name} = {class_type.strip()}()" + else: + import_statement = ("nodes", "NODE_CLASS_MAPPINGS") + class_code = f'{variable_name} = NODE_CLASS_MAPPINGS["{class_type}"]()' + return class_type, import_statement, class_code + + @staticmethod + def get_function_parameters(func: Callable) -> list | None: + signature = inspect.signature(func) + parameters = { + name: param.default if param.default != param.empty else None + for name, param in signature.parameters.items() + } + catch_all = any( + param.kind == inspect.Parameter.VAR_KEYWORD + for param in signature.parameters.values() + ) + return list(parameters.keys()) if not catch_all else None + + def update_inputs(self, inputs: dict, executed_variables: dict) -> dict: + for key in inputs.keys(): + if ( + isinstance(inputs[key], list) + and inputs[key][0] in executed_variables.keys() + ): + inputs[key] = { + "variable_name": f"get_value_at_index({executed_variables[inputs[key][0]]}, {inputs[key][1]})" + } + return inputs diff --git a/comfyui_to_python/generator/render.py b/comfyui_to_python/generator/render.py new file mode 100644 index 0000000..a9fe5ae --- /dev/null +++ b/comfyui_to_python/generator/render.py @@ -0,0 +1,135 @@ +import inspect +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, + 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 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} + ) + + 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, + ]: + func_strings.append(f"\n{inspect.getsource(func)}") + + static_imports = [ + "# Imports", + "import json", + "import os", + "import random", + "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") + custom_nodes_call = "import_custom_nodes()" + else: + custom_nodes_call = None + + 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 = [ + "# 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()", + ] + + execution_section = ["# Workflow execution"] + if imports_code: + execution_section.extend(["# Node imports", *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=" " + ).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() + ) + + entrypoint_section = [ + "# Entrypoint", + 'if __name__ == "__main__":', + " main()", + ] + + final_code = "\n".join( + static_imports + + [""] + + workflow_section + + [""] + + execution_section + + [""] + + entrypoint_section + ) + return black.format_str(final_code, mode=black.Mode()) + + @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() diff --git a/comfyui_to_python/io.py b/comfyui_to_python/io.py new file mode 100644 index 0000000..cbe27c2 --- /dev/null +++ b/comfyui_to_python/io.py @@ -0,0 +1,23 @@ +import json +import os +from typing import TextIO + + +def load_json_input(file_path: str | TextIO, encoding: str = "utf-8") -> dict: + """Read workflow JSON from a file path or file-like object.""" + if hasattr(file_path, "read"): + return json.load(file_path) + with open(file_path, "r", encoding=encoding) as file: + return json.load(file) + + +def write_python_output(file_path: str | TextIO, code: str) -> None: + """Write generated Python to a file path or file-like object.""" + if isinstance(file_path, str): + directory = os.path.dirname(file_path) + if directory and not os.path.exists(directory): + os.makedirs(directory) + with open(file_path, "w", encoding="utf-8") as file: + file.write(code) + return + file_path.write(code) diff --git a/comfyui_to_python/load_order.py b/comfyui_to_python/load_order.py new file mode 100644 index 0000000..59a75f4 --- /dev/null +++ b/comfyui_to_python/load_order.py @@ -0,0 +1,43 @@ +from typing import Dict + + +class LoadOrderDeterminer: + """Determine workflow execution order with loader-like nodes prioritized.""" + + def __init__(self, data: Dict, node_class_mappings: Dict): + self.data = data + self.node_class_mappings = node_class_mappings + self.visited = {} + self.load_order = [] + self.is_special_function = False + + def determine_load_order(self) -> list[tuple[str, Dict, bool]]: + self._load_special_functions_first() + self.is_special_function = False + for key in self.data: + if key not in self.visited: + self._dfs(key) + return self.load_order + + def _dfs(self, key: str) -> None: + self.visited[key] = True + inputs = self.data[key]["inputs"] + for value in inputs.values(): + if isinstance(value, list) and value[0] not in self.visited: + self._dfs(value[0]) + self.load_order.append((key, self.data[key], self.is_special_function)) + + def _load_special_functions_first(self) -> None: + for key in self.data: + class_def = self.node_class_mappings[self.data[key]["class_type"]]() + if ( + class_def.CATEGORY == "loaders" + or class_def.FUNCTION in ["encode"] + or not any( + isinstance(value, list) + for value in self.data[key]["inputs"].values() + ) + ): + self.is_special_function = True + if key not in self.visited: + self._dfs(key) diff --git a/comfyui_to_python/node_runtime.py b/comfyui_to_python/node_runtime.py new file mode 100644 index 0000000..894e317 --- /dev/null +++ b/comfyui_to_python/node_runtime.py @@ -0,0 +1,133 @@ +import os +import sys +from typing import Any, Mapping, Sequence, Union + + +def find_path(name: str, path: str = None) -> str: + """Recursively search parent folders until the named entry is found.""" + if path is None: + path = os.getcwd() + + if name in os.listdir(path): + path_name = os.path.join(path, name) + print(f"{name} found: {path_name}") + return path_name + + parent_directory = os.path.dirname(path) + if parent_directory == path: + return None + + return find_path(name, parent_directory) + + +def get_comfyui_path() -> str: + """Return the configured ComfyUI path, preferring COMFYUI_PATH when set.""" + comfyui_path = os.environ.get("COMFYUI_PATH") + if comfyui_path: + return comfyui_path + return find_path("ComfyUI") + + +def add_comfyui_directory_to_sys_path() -> None: + """Add the ComfyUI checkout to sys.path.""" + comfyui_path = get_comfyui_path() + if comfyui_path is not None and os.path.isdir(comfyui_path): + if comfyui_path in sys.path: + sys.path.remove(comfyui_path) + sys.path.insert(0, comfyui_path) + print(f"'{comfyui_path}' added to sys.path") + + +def add_extra_model_paths() -> None: + """Load ComfyUI extra model paths configuration when available.""" + try: + from main import load_extra_path_config + except ImportError: + print( + "Could not import load_extra_path_config from main.py. Looking in utils.extra_config instead." + ) + from utils.extra_config import load_extra_path_config + + extra_model_paths = find_path("extra_model_paths.yaml") + if extra_model_paths is not None: + load_extra_path_config(extra_model_paths) + else: + print("Could not find the extra_model_paths config file.") + + +def bootstrap_comfyui_runtime() -> None: + """Mirror the allocator-related ComfyUI startup steps before torch import.""" + add_comfyui_directory_to_sys_path() + + import comfy.options + + comfy.options.enable_args_parsing() + + from comfy.cli_args import args + + if os.name == "nt": + os.environ["MIMALLOC_PURGE_DELAY"] = "0" + + if args.default_device is not None: + default_dev = args.default_device + devices = list(range(32)) + devices.remove(default_dev) + devices.insert(0, default_dev) + devices = ",".join(map(str, devices)) + os.environ["CUDA_VISIBLE_DEVICES"] = str(devices) + os.environ["HIP_VISIBLE_DEVICES"] = str(devices) + + if args.cuda_device is not None: + os.environ["CUDA_VISIBLE_DEVICES"] = str(args.cuda_device) + os.environ["HIP_VISIBLE_DEVICES"] = str(args.cuda_device) + os.environ["ASCEND_RT_VISIBLE_DEVICES"] = str(args.cuda_device) + + if args.oneapi_device_selector is not None: + os.environ["ONEAPI_DEVICE_SELECTOR"] = args.oneapi_device_selector + + if args.deterministic and "CUBLAS_WORKSPACE_CONFIG" not in os.environ: + os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8" + + import cuda_malloc + + if "rocm" in cuda_malloc.get_torch_version_noimport(): + os.environ["OCL_SET_SVM_SIZE"] = "262144" + + +def import_custom_nodes() -> None: + """Initialize ComfyUI custom nodes in the exporter runtime.""" + comfyui_path = get_comfyui_path() + if comfyui_path and comfyui_path not in sys.path: + sys.path.insert(0, comfyui_path) + + import asyncio + import execution + from nodes import init_extra_nodes + + if comfyui_path in sys.path: + sys.path.remove(comfyui_path) + sys.path.insert(0, comfyui_path) + + import server + + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + server_instance = server.PromptServer(loop) + execution.PromptQueue(server_instance) + asyncio.run(init_extra_nodes()) + + +def get_node_class_mappings() -> dict: + """Load ComfyUI node mappings on demand.""" + add_comfyui_directory_to_sys_path() + from nodes import NODE_CLASS_MAPPINGS + + return NODE_CLASS_MAPPINGS + + +def get_value_at_index(obj: Union[Sequence, Mapping], index: int) -> Any: + """Return a sequence or mapping result item by index.""" + try: + return obj[index] + except KeyError: + return obj["result"][index] diff --git a/comfyui_to_python/workflow_loader.py b/comfyui_to_python/workflow_loader.py new file mode 100644 index 0000000..a937b9b --- /dev/null +++ b/comfyui_to_python/workflow_loader.py @@ -0,0 +1,19 @@ +import json + +from .io import load_json_input + + +def load_workflow_data(workflow: str, input_file: str): + """Load workflow data from inline JSON or an input file.""" + if input_file: + return load_json_input(input_file) + return json.loads(workflow) + + +def load_frontend_workflow_data(frontend_workflow: str | dict | None): + """Load optional frontend workflow metadata.""" + if not frontend_workflow: + return None + if isinstance(frontend_workflow, str): + return json.loads(frontend_workflow) + return frontend_workflow diff --git a/comfyui_to_python_utils.py b/comfyui_to_python_utils.py index e0bfa46..529f173 100644 --- a/comfyui_to_python_utils.py +++ b/comfyui_to_python_utils.py @@ -1,123 +1,19 @@ -import os -from typing import Sequence, Mapping, Any, Union -import sys +from comfyui_to_python.node_runtime import ( + add_comfyui_directory_to_sys_path, + add_extra_model_paths, + bootstrap_comfyui_runtime, + find_path, + get_comfyui_path, + get_value_at_index, + import_custom_nodes, +) - -def get_comfyui_path() -> str: - """Return the configured ComfyUI path, preferring COMFYUI_PATH when set.""" - comfyui_path = os.environ.get("COMFYUI_PATH") - if comfyui_path: - return comfyui_path - return find_path("ComfyUI") - - -def import_custom_nodes() -> None: - """Find all custom nodes in the custom_nodes folder and add those node objects to NODE_CLASS_MAPPINGS - - This function sets up a new asyncio event loop, initializes the PromptServer, - creates a PromptQueue, and initializes the custom nodes. - """ - comfyui_path = get_comfyui_path() - if comfyui_path and comfyui_path not in sys.path: - sys.path.insert(0, comfyui_path) - - import asyncio - import execution - from nodes import init_extra_nodes - if comfyui_path in sys.path: - sys.path.remove(comfyui_path) - sys.path.insert(0, comfyui_path) - import server - - # Creating a new event loop and setting it as the default loop - loop = asyncio.new_event_loop() - asyncio.set_event_loop(loop) - - # Creating an instance of PromptServer with the loop - server_instance = server.PromptServer(loop) - execution.PromptQueue(server_instance) - - # Initializing custom nodes - asyncio.run(init_extra_nodes()) - - -def find_path(name: str, path: str = None) -> str: - """ - Recursively looks at parent folders starting from the given path until it finds the given name. - Returns the path as a Path object if found, or None otherwise. - """ - # If no path is given, use the current working directory - if path is None: - path = os.getcwd() - - # Check if the current directory contains the name - if name in os.listdir(path): - path_name = os.path.join(path, name) - print(f"{name} found: {path_name}") - return path_name - - # Get the parent directory - parent_directory = os.path.dirname(path) - - # If the parent directory is the same as the current directory, we've reached the root and stop the search - if parent_directory == path: - return None - - # Recursively call the function with the parent directory - return find_path(name, parent_directory) - - -def add_comfyui_directory_to_sys_path() -> None: - """ - Add 'ComfyUI' to the sys.path - """ - comfyui_path = get_comfyui_path() - if comfyui_path is not None and os.path.isdir(comfyui_path): - if comfyui_path in sys.path: - sys.path.remove(comfyui_path) - sys.path.insert(0, comfyui_path) - print(f"'{comfyui_path}' added to sys.path") - - -def add_extra_model_paths() -> None: - """ - Parse the optional extra_model_paths.yaml file and add the parsed paths to the sys.path. - """ - try: - from main import load_extra_path_config - except ImportError: - print( - "Could not import load_extra_path_config from main.py. Looking in utils.extra_config instead." - ) - from utils.extra_config import load_extra_path_config - - extra_model_paths = find_path("extra_model_paths.yaml") - - if extra_model_paths is not None: - load_extra_path_config(extra_model_paths) - else: - print("Could not find the extra_model_paths config file.") - - -def get_value_at_index(obj: Union[Sequence, Mapping], index: int) -> Any: - """Returns the value at the given index of a sequence or mapping. - - If the object is a sequence (like list or string), returns the value at the given index. - If the object is a mapping (like a dictionary), returns the value at the index-th key. - - Some return a dictionary, in these cases, we look for the "results" key - - Args: - obj (Union[Sequence, Mapping]): The object to retrieve the value from. - index (int): The index of the value to retrieve. - - Returns: - Any: The value at the given index. - - Raises: - IndexError: If the index is out of bounds for the object and the object is not a mapping. - """ - try: - return obj[index] - except KeyError: - return obj["result"][index] +__all__ = [ + "add_comfyui_directory_to_sys_path", + "add_extra_model_paths", + "bootstrap_comfyui_runtime", + "find_path", + "get_comfyui_path", + "get_value_at_index", + "import_custom_nodes", +] From d751147ed95ed5b7b59d782fd65dd6bf870b4a6b Mon Sep 17 00:00:00 2001 From: Peyton DeNiro <25550995+pydn@users.noreply.github.com> Date: Sun, 29 Mar 2026 14:29:27 -0500 Subject: [PATCH 3/6] Update frontend extension to avoid legacy ui.js import --- js/save-as-script.js | 18 +++++++++++++++++- 1 file changed, 17 insertions(+), 1 deletion(-) diff --git a/js/save-as-script.js b/js/save-as-script.js index e0db883..9de431f 100644 --- a/js/save-as-script.js +++ b/js/save-as-script.js @@ -1,6 +1,22 @@ import { api } from "../../scripts/api.js"; import { app } from "../../scripts/app.js"; -import { $el } from "../../scripts/ui.js"; + +function $el(tag, options = {}) { + const element = document.createElement(tag); + const { parent, style, ...props } = options; + + if (style) { + Object.assign(element.style, style); + } + + Object.assign(element, props); + + if (parent) { + parent.appendChild(element); + } + + return element; +} const extension = { name: "Comfy.SaveAsScript", 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 4/6] 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) From 42d0526661409c77ae0120bb2eea6944fe119119 Mon Sep 17 00:00:00 2001 From: Peyton DeNiro <25550995+pydn@users.noreply.github.com> Date: Sun, 29 Mar 2026 15:15:52 -0500 Subject: [PATCH 5/6] Fix generated workflow torch import order Move generated torch imports until after add_extra_model_paths runs so exported scripts do not trigger ComfyUI's early-torch warning during startup. Update export assertions and checked-in runtime generated fixtures to match the new main() ordering. --- comfyui_to_python/generator/render.py | 4 +-- tests/runtime/generated/text-to-image.py | 4 +-- .../runtime/generated/upscale-model-loader.py | 23 ++++++++--------- tests/test_upscale_model_loader_export.py | 25 ++++++++++++++++--- 4 files changed, 36 insertions(+), 20 deletions(-) diff --git a/comfyui_to_python/generator/render.py b/comfyui_to_python/generator/render.py index e3c6238..1c531bc 100644 --- a/comfyui_to_python/generator/render.py +++ b/comfyui_to_python/generator/render.py @@ -76,8 +76,6 @@ class WorkflowRenderer: "# Workflow execution", "def main():", " bootstrap_comfyui_runtime()", - " import torch", - "", " add_extra_model_paths()", ] if custom_nodes_call: @@ -87,6 +85,8 @@ class WorkflowRenderer: execution_section.extend(f" {line}" for line in imports_code) execution_section.extend( [ + "", + " import torch", "", " with torch.inference_mode():", ] diff --git a/tests/runtime/generated/text-to-image.py b/tests/runtime/generated/text-to-image.py index 5f9b0e0..e57de44 100644 --- a/tests/runtime/generated/text-to-image.py +++ b/tests/runtime/generated/text-to-image.py @@ -165,8 +165,6 @@ extra_pnginfo = build_extra_pnginfo() # Workflow execution def main(): bootstrap_comfyui_runtime() - import torch - add_extra_model_paths() # Node imports @@ -180,6 +178,8 @@ def main(): VAEDecode, ) + import torch + with torch.inference_mode(): checkpointloadersimple = CheckpointLoaderSimple() checkpointloadersimple_1 = checkpointloadersimple.load_checkpoint( diff --git a/tests/runtime/generated/upscale-model-loader.py b/tests/runtime/generated/upscale-model-loader.py index b7b0a09..ed56d7a 100644 --- a/tests/runtime/generated/upscale-model-loader.py +++ b/tests/runtime/generated/upscale-model-loader.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() - - def import_custom_nodes() -> None: """Initialize ComfyUI custom nodes in the exporter runtime.""" comfyui_path = get_comfyui_path() @@ -166,14 +158,19 @@ workflow = build_workflow() prompt = json.loads(json.dumps(workflow)) extra_pnginfo = build_extra_pnginfo() + # Workflow execution -# Node imports -from nodes import LoadImage, NODE_CLASS_MAPPINGS, SaveImage - - def main(): + bootstrap_comfyui_runtime() + add_extra_model_paths() + import_custom_nodes() + + # Node imports + from nodes import LoadImage, NODE_CLASS_MAPPINGS, SaveImage + + import torch + with torch.inference_mode(): - import_custom_nodes() loadimage = LoadImage() loadimage_1 = loadimage.load_image(image="e2e_upscale_input.png") upscalemodelloader = NODE_CLASS_MAPPINGS["UpscaleModelLoader"]() diff --git a/tests/test_upscale_model_loader_export.py b/tests/test_upscale_model_loader_export.py index d5d2643..804de12 100644 --- a/tests/test_upscale_model_loader_export.py +++ b/tests/test_upscale_model_loader_export.py @@ -149,9 +149,28 @@ class UpscaleModelLoaderExportTest(unittest.TestCase): self.assertIn("comfy.options.enable_args_parsing()", generated) self.assertIn("import cuda_malloc", generated) 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.assertIn( + "def main():\n bootstrap_comfyui_runtime()\n add_extra_model_paths()", + generated, + ) + self.assertLess( + generated.index("def bootstrap_comfyui_runtime()"), generated.index("def main():") + ) + main_section = generated[generated.index("def main():") :] + self.assertIn( + "def main():\n" + " bootstrap_comfyui_runtime()\n" + " add_extra_model_paths()", + main_section, + ) + self.assertLess( + main_section.index("bootstrap_comfyui_runtime()"), + main_section.index("add_extra_model_paths()"), + ) + self.assertLess( + main_section.index("add_extra_model_paths()"), + main_section.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): From c1cf4e7585de55eda8a80dd4a52a180a240931ff Mon Sep 17 00:00:00 2001 From: Peyton DeNiro <25550995+pydn@users.noreply.github.com> Date: Sun, 29 Mar 2026 15:31:45 -0500 Subject: [PATCH 6/6] version update. --- pyproject.toml | 2 +- uv.lock | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 095912b..b100249 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "comfyui-to-python-extension" description = "This custom node allows you to generate pure python code from your ComfyUI workflow with the click of a button. Great for rapid experimentation or production deployment." -version = "1.3.2" +version = "2.0.0" license = { text = "MIT License" } dependencies = ["black"] diff --git a/uv.lock b/uv.lock index 545dd9f..55f3b37 100644 --- a/uv.lock +++ b/uv.lock @@ -57,7 +57,7 @@ wheels = [ [[package]] name = "comfyui-to-python-extension" -version = "1.3.2" +version = "2.0.0" source = { virtual = "." } dependencies = [ { name = "black" },