Add runtime E2E validation harness
This commit is contained in:
Peyton DeNiro
+47
-14
@@ -196,6 +196,7 @@ class CodeGenerator:
|
||||
def generate_workflow(
|
||||
self,
|
||||
load_order: List,
|
||||
workflow_data: Dict,
|
||||
queue_size: int = 10,
|
||||
) -> str:
|
||||
"""Generate the execution code based on the load order.
|
||||
@@ -209,7 +210,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 +246,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)
|
||||
|
||||
@@ -268,6 +270,11 @@ class CodeGenerator:
|
||||
and "unique_id" in input_types["hidden"].keys()
|
||||
):
|
||||
inputs["unique_id"] = random.randint(1, 2**64)
|
||||
if "hidden" in input_types.keys():
|
||||
if "prompt" in input_types["hidden"].keys():
|
||||
inputs["prompt"] = {"variable_name": "prompt"}
|
||||
if "extra_pnginfo" in input_types["hidden"].keys():
|
||||
inputs["extra_pnginfo"] = {"variable_name": "extra_pnginfo"}
|
||||
elif class_def_params is not None:
|
||||
if "unique_id" in class_def_params:
|
||||
inputs["unique_id"] = random.randint(1, 2**64)
|
||||
@@ -302,6 +309,7 @@ class CodeGenerator:
|
||||
import_statements,
|
||||
special_functions_code,
|
||||
code,
|
||||
workflow_data,
|
||||
queue_size,
|
||||
custom_nodes,
|
||||
)
|
||||
@@ -361,18 +369,20 @@ class CodeGenerator:
|
||||
|
||||
def assemble_python_code(
|
||||
self,
|
||||
import_statements: set,
|
||||
import_statements: Dict[str, set],
|
||||
speical_functions_code: List[str],
|
||||
code: List[str],
|
||||
workflow_data: Dict,
|
||||
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 original workflow data for metadata preservation.
|
||||
queue_size (int): Number of photos that will be generated by the script.
|
||||
custom_nodes (bool): Whether to include custom nodes in the code.
|
||||
|
||||
@@ -392,6 +402,7 @@ class CodeGenerator:
|
||||
# Define static import statements required for the script
|
||||
static_imports = (
|
||||
[
|
||||
"import json",
|
||||
"import os",
|
||||
"import random",
|
||||
"import sys",
|
||||
@@ -399,7 +410,17 @@ 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 = "
|
||||
'{"workflow": {"source": "workflow_api", "payload": json.loads('
|
||||
+ json.dumps(json.dumps(workflow_data))
|
||||
+ ")}}"
|
||||
),
|
||||
]
|
||||
)
|
||||
# Check if custom nodes should be included
|
||||
if custom_nodes:
|
||||
@@ -408,9 +429,10 @@ 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}")
|
||||
# Assemble the main function code, including custom nodes if applicable
|
||||
main_function_code = (
|
||||
"def main():\n\t"
|
||||
@@ -430,20 +452,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
|
||||
@@ -596,7 +629,7 @@ 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, queue_size=self.queue_size
|
||||
)
|
||||
|
||||
# Step 5: Write the generated code to a 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
|
||||
]
|
||||
}
|
||||
}
|
||||
}
|
||||
+86
@@ -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
|
||||
]
|
||||
}
|
||||
}
|
||||
}
|
||||
+11
@@ -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"
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
]
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,674 @@
|
||||
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"
|
||||
BASELINE_COMFYUI_PATH = Path("/home/peyton/code/ComfyUI")
|
||||
UPSTREAM_SMOKE_COMFYUI_PATH = Path("/home/peyton/code/ComfyUI-upstream-smoke")
|
||||
COMFYUI_OUTPUT_DIRNAME = "output"
|
||||
COMFYUI_INPUT_DIRNAME = "input"
|
||||
|
||||
if str(ROOT) not in sys.path:
|
||||
sys.path.insert(0, str(ROOT))
|
||||
|
||||
|
||||
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 = {
|
||||
"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, runtime, or upstream-smoke validation for committed fixtures."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--tier",
|
||||
choices=("fast", "runtime", "upstream-smoke"),
|
||||
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 == "smoke":
|
||||
return ["upscale-model-loader", "text-to-image"]
|
||||
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 = BASELINE_COMFYUI_PATH
|
||||
elif tier == "upstream-smoke":
|
||||
runtime_path = UPSTREAM_SMOKE_COMFYUI_PATH
|
||||
else:
|
||||
return os.environ.get("COMFYUI_PATH", "")
|
||||
|
||||
if not runtime_path.is_dir():
|
||||
raise ValidationFailure(
|
||||
"environment/setup failure",
|
||||
f"Expected ComfyUI checkout at {runtime_path}, but it was not found.",
|
||||
)
|
||||
|
||||
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
|
||||
)
|
||||
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:
|
||||
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
|
||||
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)
|
||||
|
||||
should_execute = execute or tier in {"runtime", "upstream-smoke"}
|
||||
if should_execute and tier in {"runtime", "upstream-smoke"}:
|
||||
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 in {"runtime", "upstream-smoke"}:
|
||||
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 in {"runtime", "upstream-smoke"}:
|
||||
print_download_plan(fixture, runtime_path)
|
||||
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())
|
||||
Reference in New Issue
Block a user