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

Merge pull request #149 from pydn/codex/runtime-e2e-validation

Browse Source 
Add runtime E2E validation harness and fixtures
This commit is contained in:
Peyton DeNiro
2026-03-28 18:08:24 -05:00
committed by GitHub
parent cb5917eeed 4370a78728
commit ced63d1a00
11 changed files with 1362 additions and 32 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+2
View File
@@ -110,6 +110,8 @@ export COMFYUI_PATH=/path/to/ComfyUI
The generated script is a workflow export. It does not automatically turn workflow inputs into command-line arguments.
Scripts exported directly from `File -> Save As Script` in the ComfyUI UI already include the frontend workflow metadata needed for drag-and-drop reimport. Images saved by those scripts can be dropped back into ComfyUI and reopen with the original workflow metadata.
## Troubleshooting
- `Save As Script` not visible:
__init__.py
+6 -1
View File
@@ -45,9 +45,14 @@ async def save_as_script(request):
data = await request.json()
name = data["name"]
workflow = data["workflow"]
frontend_workflow = data.get("frontend_workflow")
sio = StringIO()
ComfyUItoPython(workflow=workflow, output_file=sio)
ComfyUItoPython(
workflow=workflow,
frontend_workflow=frontend_workflow,
output_file=sio,
)
sio.seek(0)
data = sio.read()
comfyui_to_python.py
+134 -27
View File
@@ -2,6 +2,7 @@ import copy
import glob
import inspect
import json
import keyword
import os
import random
import sys
@@ -193,9 +194,20 @@ class CodeGenerator:
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.
@@ -209,7 +221,7 @@ class CodeGenerator:
"""
# Create the necessary data structures to hold imports and generated code
import_statements, executed_variables, special_functions_code, code = (
set(["NODE_CLASS_MAPPINGS"]),
{"nodes": {"NODE_CLASS_MAPPINGS"}},
{},
[],
[],
@@ -245,8 +257,9 @@ class CodeGenerator:
)
initialized_objects[class_type] = self.clean_variable_name(class_type)
if class_type in self.base_node_class_mappings.keys():
import_statements.add(import_statement)
if class_type not 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)
@@ -263,20 +276,35 @@ class CodeGenerator:
if no_params or key in class_def_params
}
# Deal with hidden variables
hidden_inputs = input_types.get("hidden", {})
if (
"hidden" in input_types.keys()
and "unique_id" in input_types["hidden"].keys()
"unique_id" in hidden_inputs
and (no_params or "unique_id" in class_def_params)
):
inputs["unique_id"] = random.randint(1, 2**64)
elif class_def_params is not None:
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)}_{idx}"
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],
@@ -287,6 +315,8 @@ class CodeGenerator:
)
)
else:
if seed_sync_code:
code.extend(seed_sync_code)
code.append(
self.create_function_call_code(
initialized_objects[class_type],
@@ -302,6 +332,8 @@ class CodeGenerator:
import_statements,
special_functions_code,
code,
workflow_data,
metadata_workflow_data,
queue_size,
custom_nodes,
)
@@ -340,6 +372,28 @@ class CodeGenerator:
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.
@@ -350,35 +404,56 @@ class CodeGenerator:
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 f"{key}=random.randint(1, 2**64)"
elif isinstance(value, str):
value = value.replace("\n", "\\n").replace('"', "'")
return f'{key}="{value}"'
elif isinstance(value, dict) and "variable_name" in value:
return f'{key}={value["variable_name"]}'
return f"{key}={value}"
return "random.randint(1, 2**64)"
if isinstance(value, str):
return json.dumps(value)
return repr(value)
def assemble_python_code(
self,
import_statements: set,
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 (set): A set of unique import statements.
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 [
@@ -392,6 +467,7 @@ class CodeGenerator:
# Define static import statements required for the script
static_imports = (
[
"import json",
"import os",
"import random",
"import sys",
@@ -399,7 +475,12 @@ class CodeGenerator:
"import torch",
]
+ func_strings
+ ["\n\nadd_comfyui_directory_to_sys_path()\nadd_extra_model_paths()\n"]
+ [
"\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:
@@ -408,16 +489,19 @@ class CodeGenerator:
else:
custom_nodes = ""
# Create import statements for node classes
imports_code = [
f"from nodes import {', '.join([class_name for class_name in import_statements])}"
]
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"
+ "\n\t\t".join(speical_functions_code)
+ special_functions_body
+ f"\n\n\t\tfor q in range({queue_size}):\n\t\t"
+ "\n\t\t".join(code)
+ loop_body
)
# Concatenate all parts to form the final code
final_code = "\n".join(
@@ -430,20 +514,31 @@ class CodeGenerator:
return final_code
def get_class_info(self, class_type: str) -> Tuple[str, str, str]:
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, str, str]: Updated class type, import statement string, class initialization code.
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)
import_statement = class_type
if class_type in self.base_node_class_mappings.keys():
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
@@ -526,6 +621,7 @@ class ComfyUItoPython:
def __init__(
self,
workflow: str = "",
frontend_workflow: str | Dict | None = None,
input_file: str = "",
output_file: str | TextIO = "",
queue_size: int = 1,
@@ -551,6 +647,7 @@ class ComfyUItoPython:
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
@@ -576,6 +673,13 @@ class ComfyUItoPython:
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 = {
@@ -596,7 +700,10 @@ class ComfyUItoPython:
self.node_class_mappings, self.base_node_class_mappings
)
generated_code = code_generator.generate_workflow(
load_order, queue_size=self.queue_size
load_order,
data,
metadata_workflow_data,
queue_size=self.queue_size,
)
# Step 5: Write the generated code to a file
js/save-as-script.js
+6 -1
View File
@@ -29,7 +29,12 @@ const extension = {
}
app.graphToPrompt().then(async (p) => {
const json = JSON.stringify({name: filename + ".json", workflow: JSON.stringify(p.output, null, 2)}, null, 2); // convert the data to a JSON string
const frontendWorkflow = p.workflow ?? app.graph.serialize();
const json = JSON.stringify({
name: filename + ".json",
workflow: JSON.stringify(p.output, null, 2),
frontend_workflow: JSON.stringify(frontendWorkflow, null, 2),
}, null, 2); // convert the data to a JSON string
var response = await api.fetchApi(`/saveasscript`, { method: "POST", body: json });
if(response.status == 200) {
const blob = new Blob([await response.text()], {type: "text/python;charset=utf-8"});
tests/README.md
+112
View File
@@ -0,0 +1,112 @@
# Tests
This directory contains two test paths:
- unit tests for exporter behavior
- runtime validation for committed workflow fixtures
## Prerequisites
- Run commands from the repo root.
- Use `uv` for all project commands.
- Install the repo environment first:
```bash
uv sync
```
- Runtime validation also needs a working ComfyUI checkout.
Set `COMFYUI_PATH` to that checkout, or run the tests from a location where a parent directory contains `ComfyUI`.
- Runtime-tier tests do not download missing models automatically.
If a required model is missing, the fixture fails with `model provisioning failure`.
Use `--print-download-plan` to print the manual download commands.
### Runtime Model Requirements
The current runtime-capable fixtures require these model files in the target ComfyUI checkout:
- `text-to-image`: `models/checkpoints/v1-5-pruned-emaonly-fp16.safetensors`
- `upscale-model-loader`: `models/upscale_models/RealESRGAN_x4plus.safetensors`
## Unit Tests
Run the current unit test module:
```bash
uv run python -m unittest tests.test_upscale_model_loader_export
```
Run all `unittest`-discoverable tests under `tests`:
```bash
uv run python -m unittest discover -s tests
```
## Runtime Tests
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.
Run all fast-tier compatible fixtures:
```bash
uv run python tests/runtime/run_runtime_validation.py --tier fast --fixture all
```
Run one fixture:
```bash
uv run python tests/runtime/run_runtime_validation.py --tier fast --fixture unsafe-kwargs
```
### Runtime Tier
Runtime tier exports inside a real ComfyUI checkout and executes generated Python for runtime-capable fixtures.
Run all runtime-capable fixtures:
```bash
uv run python tests/runtime/run_runtime_validation.py --tier runtime --fixture all
```
Run one runtime-capable fixture:
```bash
uv run python tests/runtime/run_runtime_validation.py --tier runtime --fixture text-to-image
```
Print download commands for missing models instead of failing immediately:
```bash
uv run python tests/runtime/run_runtime_validation.py --tier runtime --fixture all --print-download-plan
```
## Runtime Fixture Names
Current committed fixtures:
- `upscale-model-loader`
- `text-to-image`
- `unsafe-kwargs`
- `subgraph-identifiers`
Notes:
- `--tier runtime` only runs fixtures marked runtime-capable.
- `--tier fast` only runs fixtures with local test mappings.
## Troubleshooting
- `No module named ...`:
run `uv sync` in this repo, and for runtime-tier failures also make sure the target ComfyUI checkout has its own dependencies installed.
- `Could not find a valid ComfyUI checkout for runtime validation.`:
set `COMFYUI_PATH` to your ComfyUI checkout, or run the tests from a directory layout where a parent folder contains `ComfyUI`.
- `Missing models for ...`:
runtime-tier tests do not fetch models for you. Rerun with `--print-download-plan` to print the expected `curl` commands and target model directories, then install the files manually.
- `No selected fixtures are runtime-capable for this tier.`:
choose a runtime-capable fixture such as `text-to-image` or `upscale-model-loader`.
- Generated script execution fails because files or models are missing:
confirm the required models exist under the target ComfyUI checkout and that any staged inputs can be copied into its `input/` directory.
tests/fixtures/runtime/subgraph-identifiers.json
Vendored
+19
View File
@@ -0,0 +1,19 @@
{
"173:132": {
"class_type": "RegexReplace",
"inputs": {
"text": "abc123",
"pattern": "\\d+",
"replace": "X"
}
},
"174:133": {
"class_type": "PassthroughText",
"inputs": {
"text": [
"173:132",
0
]
}
}
}
tests/fixtures/runtime/text-to-image.json
Vendored
+86
View File
@@ -0,0 +1,86 @@
{
"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
]
}
}
}
tests/fixtures/runtime/unsafe-kwargs.json
Vendored
+11
View File
@@ -0,0 +1,11 @@
{
"1": {
"class_type": "FlexibleNode",
"inputs": {
"safe_name": "kept-readable",
"class": "reserved-word",
"\u2795 Add Lora": "symbol-heavy-key",
"spaces and-hyphens": "still-unsafe"
}
}
}
tests/fixtures/runtime/upscale-model-loader.json
Vendored
+37
View File
@@ -0,0 +1,37 @@
{
"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
]
}
}
}
tests/runtime/run_runtime_validation.py
+691
View File
@@ -0,0 +1,691 @@
import argparse
import ast
import json
import os
import shutil
import struct
import subprocess
import sys
import tempfile
import zlib
from dataclasses import dataclass
from io import StringIO
from pathlib import Path
from typing import Callable
ROOT = Path(__file__).resolve().parents[2]
FIXTURE_DIR = ROOT / "tests" / "fixtures" / "runtime"
GENERATED_DIR = ROOT / "tests" / "runtime" / "generated"
COMFYUI_OUTPUT_DIRNAME = "output"
COMFYUI_INPUT_DIRNAME = "input"
if str(ROOT) not in sys.path:
sys.path.insert(0, str(ROOT))
from comfyui_to_python_utils import get_comfyui_path
class ValidationFailure(RuntimeError):
def __init__(self, classification: str, message: str):
super().__init__(message)
self.classification = classification
self.message = message
@dataclass(frozen=True)
class FixtureConfig:
name: str
path: Path
mapping_factory: Callable[[], dict] | None = None
fast_mapping_factory: Callable[[], dict] | None = None
runtime_capable: bool = False
filename_prefix: str | None = None
expected_min_dimensions: tuple[int, int] | None = None
metadata_markers: tuple[str, ...] = ()
model_requirements: tuple["ModelRequirement", ...] = ()
staged_inputs: tuple["StagedInput", ...] = ()
@dataclass(frozen=True)
class ModelRequirement:
filename: str
relative_dir: str
source_url: str
@dataclass(frozen=True)
class StagedInput:
source_path: Path
destination_name: str
class FlexibleNode:
CATEGORY = "utils"
FUNCTION = "run"
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"safe_name": ("STRING",),
}
}
def run(self, **kwargs):
return (kwargs,)
class RegexReplace:
CATEGORY = "utils"
FUNCTION = "replace"
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"text": ("STRING",),
"pattern": ("STRING",),
"replace": ("STRING",),
}
}
def replace(self, text, pattern, replace):
return (text,)
class PassthroughText:
CATEGORY = "utils"
FUNCTION = "run"
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"text": ("STRING",),
}
}
def run(self, text):
return (text,)
class UpscaleModelLoader:
CATEGORY = "loaders"
FUNCTION = "load_model"
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"model_name": ("STRING",),
}
}
def load_model(self, model_name):
return (model_name,)
class LoadImage:
CATEGORY = "image"
FUNCTION = "load_image"
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"image": ("STRING",),
}
}
def load_image(self, image):
return (image,)
class ImageUpscaleWithModel:
CATEGORY = "image/upscaling"
FUNCTION = "upscale"
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"upscale_model": ("UPSCALE_MODEL",),
"image": ("IMAGE",),
}
}
def upscale(self, upscale_model, image):
return (image,)
class SaveImageNode:
CATEGORY = "image"
FUNCTION = "save_images"
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"images": ("IMAGE",),
"filename_prefix": ("STRING",),
}
}
def save_images(self, images, filename_prefix):
return ()
FIXTURES = {
# Runtime-capable fixtures are exported inside a real ComfyUI checkout and
# then executed end-to-end. Fast fixtures use local stub mappings so the
# exporter can be validated without importing the full runtime.
"upscale-model-loader": FixtureConfig(
name="upscale-model-loader",
path=FIXTURE_DIR / "upscale-model-loader.json",
fast_mapping_factory=lambda: {
"LoadImage": LoadImage,
"UpscaleModelLoader": UpscaleModelLoader,
"ImageUpscaleWithModel": ImageUpscaleWithModel,
"SaveImage": SaveImageNode,
},
runtime_capable=True,
filename_prefix="E2E_upscale_model_loader",
expected_min_dimensions=(1000, 1000),
metadata_markers=("UpscaleModelLoader", "E2E_upscale_model_loader"),
model_requirements=(
ModelRequirement(
filename="RealESRGAN_x4plus.safetensors",
relative_dir="models/upscale_models",
source_url="https://huggingface.co/Comfy-Org/Real-ESRGAN_repackaged/resolve/main/RealESRGAN_x4plus.safetensors",
),
),
staged_inputs=(
StagedInput(
source_path=ROOT / "images" / "save_as_script.png",
destination_name="e2e_upscale_input.png",
),
),
),
"text-to-image": FixtureConfig(
name="text-to-image",
path=FIXTURE_DIR / "text-to-image.json",
runtime_capable=True,
filename_prefix="E2E_text_to_image",
expected_min_dimensions=(512, 512),
metadata_markers=("a small cottage in a meadow", "CheckpointLoaderSimple"),
model_requirements=(
ModelRequirement(
filename="v1-5-pruned-emaonly-fp16.safetensors",
relative_dir="models/checkpoints",
source_url="https://huggingface.co/Comfy-Org/stable-diffusion-v1-5-archive/resolve/main/v1-5-pruned-emaonly-fp16.safetensors",
),
),
),
"unsafe-kwargs": FixtureConfig(
name="unsafe-kwargs",
path=FIXTURE_DIR / "unsafe-kwargs.json",
mapping_factory=lambda: {"FlexibleNode": FlexibleNode},
),
"subgraph-identifiers": FixtureConfig(
name="subgraph-identifiers",
path=FIXTURE_DIR / "subgraph-identifiers.json",
mapping_factory=lambda: {
"RegexReplace": RegexReplace,
"PassthroughText": PassthroughText,
},
),
}
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(
description="Run fast or runtime validation for committed fixtures."
)
parser.add_argument(
"--tier",
choices=("fast", "runtime"),
help="Validation tier to run.",
)
parser.add_argument(
"--fixture",
default="all",
help="Fixture name or 'all'.",
)
parser.add_argument(
"--execute",
action="store_true",
help="Execute generated Python for runtime-capable fixtures.",
)
parser.add_argument(
"--print-download-plan",
action="store_true",
help="Print model download commands for missing models.",
)
parser.add_argument(
"--internal-export",
choices=tuple(FIXTURES.keys()),
help=argparse.SUPPRESS,
)
parser.add_argument(
"--generated-path",
help=argparse.SUPPRESS,
)
args = parser.parse_args()
if not args.internal_export and not args.tier:
parser.error("--tier is required unless --internal-export is used.")
return args
def load_fixture_names(selection: str) -> list[str]:
if selection == "all":
return list(FIXTURES.keys())
if selection not in FIXTURES:
raise ValidationFailure("fixture bug", f"Unknown fixture '{selection}'.")
return [selection]
def ensure_runtime_path(tier: str) -> str:
if tier == "runtime":
runtime_path = get_comfyui_path()
else:
return os.environ.get("COMFYUI_PATH", "")
if not runtime_path or not Path(runtime_path).is_dir():
raise ValidationFailure(
"environment/setup failure",
"Could not find a valid ComfyUI checkout for runtime validation. "
"Set COMFYUI_PATH or run the tests from a location where a "
"parent directory contains ComfyUI.",
)
return str(runtime_path)
def get_runtime_python(runtime_path: str) -> str:
runtime_python = Path(runtime_path) / ".venv" / "bin" / "python"
if runtime_python.is_file():
return str(runtime_python)
return sys.executable
def get_fixture(fixture_name: str) -> FixtureConfig:
return FIXTURES[fixture_name]
def check_models(fixture: FixtureConfig, runtime_path: str) -> list[ModelRequirement]:
runtime_root = Path(runtime_path)
missing = []
for requirement in fixture.model_requirements:
target = runtime_root / requirement.relative_dir / requirement.filename
if not target.is_file():
missing.append(requirement)
return missing
def print_download_plan(fixture: FixtureConfig, runtime_path: str) -> None:
runtime_root = Path(runtime_path)
for requirement in check_models(fixture, runtime_path):
target_dir = runtime_root / requirement.relative_dir
print(
"download:",
f"mkdir -p {target_dir} && curl -L {requirement.source_url} -o {target_dir / requirement.filename}",
)
def stage_inputs(fixture: FixtureConfig, runtime_path: str) -> None:
runtime_input_dir = Path(runtime_path) / COMFYUI_INPUT_DIRNAME
runtime_input_dir.mkdir(parents=True, exist_ok=True)
for staged_input in fixture.staged_inputs:
if not staged_input.source_path.is_file():
raise ValidationFailure(
"fixture bug",
f"Missing staged input source {staged_input.source_path} for {fixture.name}.",
)
shutil.copyfile(
staged_input.source_path,
runtime_input_dir / staged_input.destination_name,
)
def export_workflow(
fixture: FixtureConfig,
tier: str,
runtime_path: str,
) -> tuple[str, str]:
from comfyui_to_python import ComfyUItoPython
workflow = fixture.path.read_text(encoding="utf-8")
output = StringIO()
kwargs = {
"workflow": workflow,
"output_file": output,
}
if tier == "fast" and fixture.fast_mapping_factory is not None:
kwargs["node_class_mappings"] = fixture.fast_mapping_factory()
elif fixture.mapping_factory is not None:
kwargs["node_class_mappings"] = fixture.mapping_factory()
else:
os.environ["COMFYUI_PATH"] = runtime_path
try:
ComfyUItoPython(**kwargs)
except ModuleNotFoundError as exc:
missing_module = exc.name or "unknown"
raise ValidationFailure(
"environment/setup failure",
f"Missing runtime dependency '{missing_module}' while exporting {fixture.name}.",
) from exc
except KeyError as exc:
raise ValidationFailure(
"repo regression",
f"Exporter failed to resolve workflow data for {fixture.name}: {exc}",
) from exc
except Exception as exc:
raise ValidationFailure(
"repo regression",
f"Exporter failed for {fixture.name}: {exc}",
) from exc
return workflow, output.getvalue()
def export_workflow_in_runtime_env(fixture: FixtureConfig, runtime_path: str) -> str:
runtime_python = get_runtime_python(runtime_path)
GENERATED_DIR.mkdir(parents=True, exist_ok=True)
generated_path = GENERATED_DIR / f"{fixture.name}.py"
env = os.environ.copy()
env["COMFYUI_PATH"] = runtime_path
env["PYTHONPATH"] = os.pathsep.join([str(ROOT), env.get("PYTHONPATH", "")]).rstrip(
os.pathsep
)
# Re-enter this script under the runtime interpreter so export happens with
# the target ComfyUI checkout on sys.path, not just the repo's current venv.
result = subprocess.run(
[
runtime_python,
str(Path(__file__).resolve()),
"--internal-export",
fixture.name,
"--generated-path",
str(generated_path),
],
cwd=ROOT,
env=env,
capture_output=True,
text=True,
)
if result.returncode != 0:
output = (result.stderr or result.stdout or "").strip()
classification = "environment/setup failure"
if (
"Missing runtime dependency" not in output
and "ModuleNotFoundError" not in output
):
classification = "repo regression"
raise ValidationFailure(
classification,
f"Runtime export failed for {fixture.name}: {output}",
)
return generated_path.read_text(encoding="utf-8")
def validate_generated_python(generated_code: str, fixture_name: str) -> None:
try:
ast.parse(generated_code)
except SyntaxError as exc:
raise ValidationFailure(
"repo regression",
f"Generated Python is not valid for {fixture_name}: {exc}",
) from exc
def parse_png_info(image_path: Path) -> tuple[int, int, dict[str, str]]:
with image_path.open("rb") as handle:
signature = handle.read(8)
if signature != b"\x89PNG\r\n\x1a\n":
raise ValidationFailure(
"environment/setup failure",
f"Expected PNG output for {image_path.name}, got a different file format.",
)
width = height = None
text_data: dict[str, str] = {}
while True:
length_bytes = handle.read(4)
if not length_bytes:
break
length = struct.unpack(">I", length_bytes)[0]
chunk_type = handle.read(4)
chunk_data = handle.read(length)
handle.read(4)
if chunk_type == b"IHDR":
width, height = struct.unpack(">II", chunk_data[:8])
elif chunk_type == b"tEXt":
key, value = chunk_data.split(b"\x00", 1)
text_data[key.decode("latin-1")] = value.decode("latin-1")
elif chunk_type == b"zTXt":
key, compressed = chunk_data.split(b"\x00", 1)
text_data[key.decode("latin-1")] = zlib.decompress(
compressed[1:]
).decode("latin-1")
elif chunk_type == b"iTXt":
parts = chunk_data.split(b"\x00", 5)
if len(parts) == 6:
key = parts[0].decode("utf-8")
compressed_flag = parts[1]
value = parts[5]
if compressed_flag == b"\x01":
value = zlib.decompress(value)
text_data[key] = value.decode("utf-8")
elif chunk_type == b"IEND":
break
if width is None or height is None:
raise ValidationFailure(
"environment/setup failure",
f"Could not read PNG dimensions from {image_path.name}.",
)
return width, height, text_data
def validate_output_artifact(
fixture: FixtureConfig,
output_path: Path,
) -> None:
# Read PNG metadata directly so artifact validation does not depend on
# optional imaging libraries inside the runtime environment.
width, height, metadata = parse_png_info(output_path)
if fixture.expected_min_dimensions is not None:
min_width, min_height = fixture.expected_min_dimensions
if width < min_width or height < min_height:
raise ValidationFailure(
"repo regression",
f"Output dimensions for {fixture.name} were {width}x{height}, expected at least {min_width}x{min_height}.",
)
metadata_blob = "\n".join(
[output_path.name] + [f"{key}={value}" for key, value in metadata.items()]
)
for marker in fixture.metadata_markers:
if marker not in metadata_blob:
raise ValidationFailure(
"repo regression",
f"Output metadata for {fixture.name} did not contain expected marker '{marker}'.",
)
def execute_generated_python(
generated_code: str,
fixture: FixtureConfig,
runtime_path: str,
) -> None:
if not fixture.runtime_capable:
raise ValidationFailure(
"fixture bug",
f"Fixture {fixture.name} is not marked runtime-capable.",
)
with tempfile.TemporaryDirectory() as tmpdir:
tmp_path = Path(tmpdir) / f"{fixture.name}.py"
tmp_path.write_text(generated_code, encoding="utf-8")
output_dir = Path(runtime_path) / COMFYUI_OUTPUT_DIRNAME
# Compare against the pre-run snapshot so validation can prove this
# execution created a fresh artifact instead of reusing an old output.
existing_outputs = set(output_dir.glob("*.png"))
env = os.environ.copy()
env["COMFYUI_PATH"] = runtime_path
env["PYTHONPATH"] = os.pathsep.join(
[str(ROOT), env.get("PYTHONPATH", "")]
).rstrip(os.pathsep)
runtime_python = get_runtime_python(runtime_path)
result = subprocess.run(
[runtime_python, str(tmp_path)],
cwd=ROOT,
env=env,
capture_output=True,
text=True,
)
if result.returncode == 0:
if fixture.filename_prefix is None:
return
new_outputs = [
path
for path in output_dir.glob(f"{fixture.filename_prefix}*.png")
if path not in existing_outputs
]
if not new_outputs:
raise ValidationFailure(
"repo regression",
f"Generated script for {fixture.name} did not produce a new output file with prefix {fixture.filename_prefix}.",
)
newest_output = max(new_outputs, key=lambda path: path.stat().st_mtime)
validate_output_artifact(fixture, newest_output)
return
stderr = (result.stderr or "").strip()
stdout = (result.stdout or "").strip()
output = stderr or stdout or "generated script exited with a non-zero status"
lower_output = output.lower()
if "no module named 'torch'" in lower_output:
classification = "environment/setup failure"
elif "no such file or directory" in lower_output or "not found" in lower_output:
classification = "environment/setup failure"
else:
classification = "repo regression"
raise ValidationFailure(
classification,
f"Generated script execution failed for {fixture.name}: {output}",
)
def run_fixture(fixture: FixtureConfig, tier: str, execute: bool, runtime_path: str) -> str:
if tier == "fast":
_, generated_code = export_workflow(fixture, tier, runtime_path)
else:
missing_models = check_models(fixture, runtime_path)
if missing_models:
raise ValidationFailure(
"model provisioning failure",
"Missing models for "
f"{fixture.name}: "
+ ", ".join(
f"{item.relative_dir}/{item.filename}" for item in missing_models
),
)
stage_inputs(fixture, runtime_path)
generated_code = export_workflow_in_runtime_env(fixture, runtime_path)
validate_generated_python(generated_code, fixture.name)
# Fast keeps execution opt-in because its stub node mappings are intended for
# export coverage only. Runtime always executes the generated script.
should_execute = execute or tier == "runtime"
if should_execute and tier == "runtime":
execute_generated_python(generated_code, fixture, runtime_path)
return "pass"
def main() -> int:
args = parse_args()
if args.internal_export:
fixture = get_fixture(args.internal_export)
output_path = Path(args.generated_path)
_, generated_code = export_workflow(
fixture=fixture,
tier="runtime",
runtime_path=os.environ.get("COMFYUI_PATH", ""),
)
output_path.write_text(generated_code, encoding="utf-8")
return 0
try:
runtime_path = ensure_runtime_path(args.tier)
fixture_names = load_fixture_names(args.fixture)
requested = [FIXTURES[name] for name in fixture_names]
if args.tier == "fast":
requested = [
fixture
for fixture in requested
if fixture.fast_mapping_factory is not None
or fixture.mapping_factory is not None
]
if not requested:
raise ValidationFailure(
"fixture bug",
"No selected fixtures are fast-tier compatible.",
)
elif args.tier == "runtime":
requested = [fixture for fixture in requested if fixture.runtime_capable]
if not requested:
raise ValidationFailure(
"fixture bug",
"No selected fixtures are runtime-capable for this tier.",
)
failures: list[tuple[str, str, str]] = []
for fixture in requested:
try:
if args.print_download_plan and args.tier == "runtime":
missing_models = check_models(fixture, runtime_path)
for requirement in missing_models:
target_dir = Path(runtime_path) / requirement.relative_dir
print(
"download:",
f"mkdir -p {target_dir} && curl -L {requirement.source_url} -o {target_dir / requirement.filename}",
)
if missing_models:
print(f"{fixture.name}: download-plan")
continue
status = run_fixture(fixture, args.tier, args.execute, runtime_path)
print(f"{fixture.name}: {status}")
except ValidationFailure as exc:
failures.append((fixture.name, exc.classification, exc.message))
print(
f"{fixture.name}: fail ({exc.classification})",
file=sys.stderr,
)
print(exc.message, file=sys.stderr)
if failures:
classifications = ", ".join(
f"{name}={classification}" for name, classification, _ in failures
)
print(f"classification: {classifications}", file=sys.stderr)
return 1
except ValidationFailure as exc:
print(f"classification: {exc.classification}", file=sys.stderr)
print(exc.message, file=sys.stderr)
return 1
return 0
if __name__ == "__main__":
sys.exit(main())
tests/test_upscale_model_loader_export.py
+258 -3
View File
@@ -1,8 +1,11 @@
import json
import tempfile
import unittest
from io import StringIO
from pathlib import Path
from unittest.mock import patch
from comfyui_to_python import ComfyUItoPython
from comfyui_to_python import ComfyUItoPython, run
class UpscaleModelLoader:
@@ -54,6 +57,63 @@ class ImageUpscaleWithModel:
return (image,)
class MergeImages:
CATEGORY = "image"
FUNCTION = "merge"
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"left": ("IMAGE",),
"right": ("IMAGE",),
}
}
def merge(self, left, right):
return (left, right)
class HiddenMetadataFilteredNode:
CATEGORY = "image"
FUNCTION = "save"
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"image": ("IMAGE",),
},
"hidden": {
"prompt": "PROMPT",
"extra_pnginfo": "EXTRA_PNGINFO",
},
}
def save(self, image):
return (image,)
class HiddenPromptSeedNode:
CATEGORY = "sampling"
FUNCTION = "sample"
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"seed": ("INT",),
"noise_seed": ("INT",),
},
"hidden": {
"prompt": "PROMPT",
},
}
def sample(self, seed, noise_seed, prompt):
return (seed, noise_seed, prompt)
class UpscaleModelLoaderExportTest(unittest.TestCase):
def test_upscale_workflow_uses_direct_upscale_model_loader_init(self):
workflow = {
@@ -91,11 +151,11 @@ class UpscaleModelLoaderExportTest(unittest.TestCase):
generated = output.getvalue()
self.assertIn("from nodes 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)
self.assertIn("NODE_CLASS_MAPPINGS,", generated)
self.assertIn("loadimage = LoadImage()", generated)
self.assertIn("upscalemodelloader = UpscaleModelLoader()", generated)
self.assertIn(
@@ -104,6 +164,201 @@ class UpscaleModelLoaderExportTest(unittest.TestCase):
)
self.assertNotIn('NODE_CLASS_MAPPINGS["UpscaleModelLoader"]()', generated)
def test_frontend_workflow_metadata_is_embedded_for_reimport(self):
workflow = {
"1": {
"class_type": "LoadImage",
"inputs": {
"image": "example.png",
},
}
}
frontend_workflow = {
"version": 0.4,
"last_node_id": 1,
"last_link_id": 0,
"nodes": [],
"links": [],
"groups": [],
"config": {},
"extra": {},
}
output = StringIO()
ComfyUItoPython(
workflow=json.dumps(workflow),
frontend_workflow=json.dumps(frontend_workflow),
output_file=output,
node_class_mappings={
"LoadImage": LoadImage,
},
)
generated = output.getvalue()
self.assertIn('extra_pnginfo = {', generated)
self.assertIn('"workflow": json.loads(', generated)
self.assertIn('"version": 0.4', generated)
self.assertIn('"nodes": []', generated)
self.assertNotIn('"source": "workflow_api"', generated)
def test_export_without_frontend_workflow_leaves_png_workflow_metadata_absent(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("extra_pnginfo = None", generated)
self.assertNotIn('"workflow": json.loads(', generated)
def test_export_preserves_unique_variable_names_for_subgraph_node_ids(self):
workflow = {
"1:23": {
"class_type": "LoadImage",
"inputs": {
"image": "left.png",
},
},
"12:3": {
"class_type": "LoadImage",
"inputs": {
"image": "right.png",
},
},
"20": {
"class_type": "MergeImages",
"inputs": {
"left": ["1:23", 0],
"right": ["12:3", 0],
},
},
}
output = StringIO()
ComfyUItoPython(
workflow=json.dumps(workflow),
output_file=output,
node_class_mappings={
"LoadImage": LoadImage,
"MergeImages": MergeImages,
},
)
generated = output.getvalue()
self.assertIn("loadimage_1_23 = loadimage.load_image(", generated)
self.assertIn("loadimage_12_3 = loadimage.load_image(", generated)
self.assertIn("mergeimages_20 = mergeimages.merge(", generated)
self.assertIn("left=get_value_at_index(loadimage_1_23, 0)", generated)
self.assertIn("right=get_value_at_index(loadimage_12_3, 0)", generated)
def test_run_cli_export_leaves_png_workflow_metadata_absent(self):
workflow = {
"1": {
"class_type": "LoadImage",
"inputs": {
"image": "example.png",
},
}
}
with tempfile.TemporaryDirectory() as tmpdir:
tmpdir_path = Path(tmpdir)
input_file = tmpdir_path / "workflow_api.json"
output_file = tmpdir_path / "workflow_api.py"
input_file.write_text(json.dumps(workflow), encoding="utf-8")
with patch(
"comfyui_to_python.get_node_class_mappings",
return_value={"LoadImage": LoadImage},
), patch("comfyui_to_python.import_custom_nodes"):
run(
input_file=str(input_file),
output_file=str(output_file),
queue_size=1,
)
generated = output_file.read_text(encoding="utf-8")
self.assertIn("extra_pnginfo = None", generated)
self.assertNotIn('"workflow": json.loads(', generated)
def test_hidden_metadata_kwargs_follow_function_signature(self):
workflow = {
"1": {
"class_type": "HiddenMetadataFilteredNode",
"inputs": {
"image": "example.png",
},
}
}
output = StringIO()
ComfyUItoPython(
workflow=json.dumps(workflow),
output_file=output,
node_class_mappings={
"HiddenMetadataFilteredNode": HiddenMetadataFilteredNode,
},
)
generated = output.getvalue()
self.assertIn("hiddenmetadatafilterednode_1 =", generated)
self.assertNotIn("prompt=prompt", generated)
self.assertNotIn("extra_pnginfo=extra_pnginfo", generated)
def test_randomized_seed_inputs_update_prompt_metadata_before_execution(self):
workflow = {
"1": {
"class_type": "HiddenPromptSeedNode",
"inputs": {
"seed": 1,
"noise_seed": 2,
},
}
}
output = StringIO()
ComfyUItoPython(
workflow=json.dumps(workflow),
output_file=output,
node_class_mappings={
"HiddenPromptSeedNode": HiddenPromptSeedNode,
},
)
generated = output.getvalue()
self.assertIn(
'node_1_seed = prompt["1"]["inputs"]["seed"] = random.randint(1, 2**64)',
generated,
)
self.assertIn(
'node_1_noise_seed = prompt["1"]["inputs"]["noise_seed"] = random.randint(',
generated,
)
self.assertIn("1, 2**64", generated)
self.assertIn("seed=node_1_seed", generated)
self.assertIn("noise_seed=node_1_noise_seed", generated)
self.assertIn("prompt=prompt", generated)
if __name__ == "__main__":
unittest.main()