Compare commits

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