diff --git a/README.md b/README.md index 32afd50..b80470b 100644 --- a/README.md +++ b/README.md @@ -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: diff --git a/__init__.py b/__init__.py index 69ce450..10cdf1c 100644 --- a/__init__.py +++ b/__init__.py @@ -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() diff --git a/comfyui_to_python.py b/comfyui_to_python.py index ebcf928..efae605 100644 --- a/comfyui_to_python.py +++ b/comfyui_to_python.py @@ -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 diff --git a/js/save-as-script.js b/js/save-as-script.js index 91a1a84..e0db883 100644 --- a/js/save-as-script.js +++ b/js/save-as-script.js @@ -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"}); diff --git a/tests/README.md b/tests/README.md new file mode 100644 index 0000000..cc7bca6 --- /dev/null +++ b/tests/README.md @@ -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. diff --git a/tests/fixtures/runtime/subgraph-identifiers.json b/tests/fixtures/runtime/subgraph-identifiers.json new file mode 100644 index 0000000..3690e72 --- /dev/null +++ b/tests/fixtures/runtime/subgraph-identifiers.json @@ -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 + ] + } + } +} diff --git a/tests/fixtures/runtime/text-to-image.json b/tests/fixtures/runtime/text-to-image.json new file mode 100644 index 0000000..6d5cf67 --- /dev/null +++ b/tests/fixtures/runtime/text-to-image.json @@ -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 + ] + } + } +} diff --git a/tests/fixtures/runtime/unsafe-kwargs.json b/tests/fixtures/runtime/unsafe-kwargs.json new file mode 100644 index 0000000..9c7fdae --- /dev/null +++ b/tests/fixtures/runtime/unsafe-kwargs.json @@ -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" + } + } +} diff --git a/tests/fixtures/runtime/upscale-model-loader.json b/tests/fixtures/runtime/upscale-model-loader.json new file mode 100644 index 0000000..948daa1 --- /dev/null +++ b/tests/fixtures/runtime/upscale-model-loader.json @@ -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 + ] + } + } +} diff --git a/tests/runtime/run_runtime_validation.py b/tests/runtime/run_runtime_validation.py new file mode 100644 index 0000000..46b0d84 --- /dev/null +++ b/tests/runtime/run_runtime_validation.py @@ -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()) diff --git a/tests/test_upscale_model_loader_export.py b/tests/test_upscale_model_loader_export.py index 754c69a..d711c4b 100644 --- a/tests/test_upscale_model_loader_export.py +++ b/tests/test_upscale_model_loader_export.py @@ -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()