Update version.

Add generated runtime cleanup after workflow execution

README.md

@@ -13,6 +13,7 @@ This project supports:
 ## Install
 
 Choose the setup that matches how you want to use the project.
+This project supports Python 3.12 and newer.
 
 ### Web UI extension (`File -> Save As Script`)
 
@@ -65,12 +66,13 @@ In current ComfyUI builds, `Save As Script` is typically available under:
 `File -> Save As Script`
 
 The command downloads a generated `.py` file.
+The current UI export uses the default filename `workflow_api.py` so it works in ComfyUI Desktop without relying on `prompt()`.
 
 ![Save As Script](images/save_as_script.png)
 
 Notes:
 - menu placement can differ between frontend versions
-- ComfyUI Desktop may fail on the current filename prompt flow; use the CLI flow below if that happens
+- the Web UI export uses a fixed default filename rather than asking for one interactively
 
 ## CLI Export
 
@@ -79,18 +81,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`
@@ -112,14 +120,24 @@ The generated script is a workflow export. It does not automatically turn workfl
 
 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.
 
+Generated scripts reuse ComfyUI's runtime argument parser during bootstrap, so common ComfyUI memory flags such as `--highvram`, `--normalvram`, `--lowvram`, `--novram`, `--cpu`, and `--disable-smart-memory` can be passed directly to the exported `.py` file.
+
+Lifecycle notes:
+- exported scripts are single-shot workflow runners, not long-lived ComfyUI prompt servers
+- they do not implement Web UI prompt/result caching across repeated service calls
+- exported `main()` now performs best-effort ComfyUI model/cache cleanup in a `finally` block
+- set `COMFYUI_TOPYTHON_UNLOAD_MODELS=1` or call `main(unload_models=True)` if an embedded or repeated-call host should aggressively unload models after each run instead of preserving them for reuse
+
 ## Troubleshooting
 
+- unsupported Python version:
+  use Python 3.12 or newer, then rerun `uv sync`
 - `Save As Script` not visible:
   check your current ComfyUI menu/frontend version and look under `File`
 - `Save As Script` not visible after restart:
   make sure this repo is discoverable by ComfyUI through `custom_nodes` by cloning it into `ComfyUI/custom_nodes`, symlinking it there, or adding an external `custom_nodes` path in `extra_model_paths.yaml`
-- Desktop says `prompt()` is unsupported:
-  use the CLI export flow instead
+- save uses the default filename:
+  rename `workflow_api.py` after download if you want a different local filename
 - ComfyUI cannot be found:
   set `COMFYUI_PATH`
 - models or paths are missing at runtime:

__init__.py

@@ -13,18 +13,10 @@ sys.path.append(ext_dir)
 try:
     import black
 except ImportError:
-    print("Unable to import requirements for ComfyUI-SaveAsScript.")
-    print("Installing...")
-
-    import importlib
-
-    spec = importlib.util.spec_from_file_location(
-        "impact_install", os.path.join(os.path.dirname(__file__), "install.py")
-    )
-    impact_install = importlib.util.module_from_spec(spec)
-    spec.loader.exec_module(impact_install)
-
-    print("Successfully installed. Hopefully, at least.")
+    raise ImportError(
+        "ComfyUI-to-Python-Extension requires the project dependencies to be installed. "
+        f"Run 'uv sync' in {ext_dir} with Python 3.12+ before loading this extension."
+    ) from None
 
 # Prevent reimporting of custom nodes
 os.environ["RUNNING_IN_COMFYUI"] = "TRUE"

comfyui_to_python.py

@@ -1,773 +1,5 @@
-import copy
-import glob
-import inspect
-import json
-import keyword
-import os
-import random
-import sys
-import re
-from typing import Dict, List, Any, Callable, Tuple, TextIO
-from argparse import ArgumentParser
-
-import black
-
-# Add current directory to Python path
-sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
-from comfyui_to_python_utils import (
-    import_custom_nodes,
-    find_path,
-    get_comfyui_path,
-    add_comfyui_directory_to_sys_path,
-    add_extra_model_paths,
-    get_value_at_index,
-)
-
-
-DEFAULT_INPUT_FILE = "workflow_api.json"
-DEFAULT_OUTPUT_FILE = "workflow_api.py"
-DEFAULT_QUEUE_SIZE = 10
-
-
-def get_node_class_mappings() -> Dict:
-    """Load ComfyUI node mappings on demand.
-
-    Tests that inject explicit node mappings should not need a full ComfyUI runtime
-    just to import this module.
-    """
-    add_comfyui_directory_to_sys_path()
-    from nodes import NODE_CLASS_MAPPINGS
-
-    return NODE_CLASS_MAPPINGS
-
-
-class FileHandler:
-    """Handles reading and writing files.
-
-    This class provides methods to read JSON data from an input file and write code to an output file.
-    """
-
-    @staticmethod
-    def read_json_file(file_path: str | TextIO, encoding: str = "utf-8") -> dict:
-        """
-        Reads a JSON file and returns its contents as a dictionary.
-
-        Args:
-            file_path (str): The path to the JSON file.
-
-        Returns:
-            dict: The contents of the JSON file as a dictionary.
-
-        Raises:
-            FileNotFoundError: If the file is not found, it lists all JSON files in the directory of the file path.
-            ValueError: If the file is not a valid JSON.
-        """
-
-        if hasattr(file_path, "read"):
-            return json.load(file_path)
-        with open(file_path, "r", encoding="utf-8") as file:
-            data = json.load(file)
-        return data
-
-    @staticmethod
-    def write_code_to_file(file_path: str | TextIO, code: str) -> None:
-        """Write the specified code to a Python file.
-
-        Args:
-            file_path (str): The path to the Python file.
-            code (str): The code to write to the file.
-
-        Returns:
-            None
-        """
-        if isinstance(file_path, str):
-            # Extract directory from the filename
-            directory = os.path.dirname(file_path)
-
-            # If the directory does not exist, create it
-            if directory and not os.path.exists(directory):
-                os.makedirs(directory)
-
-            # Save the code to a .py file
-            with open(file_path, "w", encoding="utf-8") as file:
-                file.write(code)
-        else:
-            file_path.write(code)
-
-
-class LoadOrderDeterminer:
-    """Determine the load order of each key in the provided dictionary.
-
-    This class places the nodes without node dependencies first, then ensures that any node whose
-    result is used in another node will be added to the list in the order it should be executed.
-
-    Attributes:
-        data (Dict): The dictionary for which to determine the load order.
-        node_class_mappings (Dict): Mappings of node classes.
-    """
-
-    def __init__(self, data: Dict, node_class_mappings: Dict):
-        """Initialize the LoadOrderDeterminer with the given data and node class mappings.
-
-        Args:
-            data (Dict): The dictionary for which to determine the load order.
-            node_class_mappings (Dict): Mappings of node classes.
-        """
-        self.data = data
-        self.node_class_mappings = node_class_mappings
-        self.visited = {}
-        self.load_order = []
-        self.is_special_function = False
-
-    def determine_load_order(self) -> List[Tuple[str, Dict, bool]]:
-        """Determine the load order for the given data.
-
-        Returns:
-            List[Tuple[str, Dict, bool]]: A list of tuples representing the load order.
-        """
-        self._load_special_functions_first()
-        self.is_special_function = False
-        for key in self.data:
-            if key not in self.visited:
-                self._dfs(key)
-        return self.load_order
-
-    def _dfs(self, key: str) -> None:
-        """Depth-First Search function to determine the load order.
-
-        Args:
-            key (str): The key from which to start the DFS.
-
-        Returns:
-            None
-        """
-        # Mark the node as visited.
-        self.visited[key] = True
-        inputs = self.data[key]["inputs"]
-        # Loop over each input key.
-        for input_key, val in inputs.items():
-            # If the value is a list and the first item in the list has not been visited yet,
-            # then recursively apply DFS on the dependency.
-            if isinstance(val, list) and val[0] not in self.visited:
-                self._dfs(val[0])
-        # Add the key and its corresponding data to the load order list.
-        self.load_order.append((key, self.data[key], self.is_special_function))
-
-    def _load_special_functions_first(self) -> None:
-        """Load functions without dependencies, loaderes, and encoders first.
-
-        Returns:
-            None
-        """
-        # Iterate over each key in the data to check for loader keys.
-        for key in self.data:
-            class_def = self.node_class_mappings[self.data[key]["class_type"]]()
-            # Check if the class is a loader class or meets specific conditions.
-            if (
-                class_def.CATEGORY == "loaders"
-                or class_def.FUNCTION in ["encode"]
-                or not any(
-                    isinstance(val, list) for val in self.data[key]["inputs"].values()
-                )
-            ):
-                self.is_special_function = True
-                # If the key has not been visited, perform a DFS from that key.
-                if key not in self.visited:
-                    self._dfs(key)
-
-
-class CodeGenerator:
-    """Generates Python code for a workflow based on the load order.
-
-    Attributes:
-        node_class_mappings (Dict): Mappings of node classes.
-        base_node_class_mappings (Dict): Base mappings of node classes.
-    """
-
-    def __init__(self, node_class_mappings: Dict, base_node_class_mappings: Dict):
-        """Initialize the CodeGenerator with given node class mappings.
-
-        Args:
-            node_class_mappings (Dict): Mappings of node classes.
-            base_node_class_mappings (Dict): Base mappings of node classes.
-        """
-        self.node_class_mappings = node_class_mappings
-        self.base_node_class_mappings = base_node_class_mappings
-
-    @staticmethod
-    def sanitize_node_id(node_id: str) -> str:
-        """Convert node IDs into variable-safe tokens without collapsing separators."""
-        sanitized = re.sub(r"[^a-z0-9_]", "_", str(node_id).lower().strip())
-        sanitized = re.sub(r"_+", "_", sanitized).strip("_")
-        if not sanitized:
-            sanitized = "node"
-        return sanitized
-
-    def generate_workflow(
-        self,
-        load_order: List,
-        workflow_data: Dict,
-        metadata_workflow_data: Dict | None = None,
-        queue_size: int = 10,
-    ) -> str:
-        """Generate the execution code based on the load order.
-
-        Args:
-            load_order (List): A list of tuples representing the load order.
-            queue_size (int): The number of photos that will be created by the script.
-
-        Returns:
-            str: Generated execution code as a string.
-        """
-        # Create the necessary data structures to hold imports and generated code
-        import_statements, executed_variables, special_functions_code, code = (
-            {"nodes": {"NODE_CLASS_MAPPINGS"}},
-            {},
-            [],
-            [],
-        )
-        # This dictionary will store the names of the objects that we have already initialized
-        initialized_objects = {}
-
-        custom_nodes = False
-        # Loop over each dictionary in the load order list
-        for idx, data, is_special_function in load_order:
-            # Generate class definition and inputs from the data
-            inputs, class_type = data["inputs"], data["class_type"]
-            input_types = self.node_class_mappings[class_type].INPUT_TYPES()
-            class_def = self.node_class_mappings[class_type]()
-
-            # If required inputs are not present, skip the node as it will break the code if passed through to the script
-            missing_required_variable = False
-            if "required" in input_types.keys():
-                for required in input_types["required"]:
-                    if required not in inputs.keys():
-                        missing_required_variable = True
-            if missing_required_variable:
-                continue
-
-            # If the class hasn't been initialized yet, initialize it and generate the import statements
-            if class_type not in initialized_objects:
-                # No need to use preview image nodes since we are executing the script in a terminal
-                if class_type == "PreviewImage":
-                    continue
-
-                class_type, import_statement, class_code = self.get_class_info(
-                    class_type
-                )
-                initialized_objects[class_type] = self.clean_variable_name(class_type)
-                if class_type in self.base_node_class_mappings.keys():
-                    module_name, import_name = import_statement
-                    import_statements.setdefault(module_name, set()).add(import_name)
-                if 'NODE_CLASS_MAPPINGS["' in class_code:
-                    custom_nodes = True
-                special_functions_code.append(class_code)
-
-            # Get all possible parameters for class_def
-            class_def_params = self.get_function_parameters(
-                getattr(class_def, class_def.FUNCTION)
-            )
-            no_params = class_def_params is None
-
-            # Remove any keyword arguments from **inputs if they are not in class_def_params
-            inputs = {
-                key: value
-                for key, value in inputs.items()
-                if no_params or key in class_def_params
-            }
-            # Deal with hidden variables
-            hidden_inputs = input_types.get("hidden", {})
-            if (
-                "unique_id" in hidden_inputs
-                and (no_params or "unique_id" in class_def_params)
-            ):
-                inputs["unique_id"] = random.randint(1, 2**64)
-            if "prompt" in hidden_inputs and (no_params or "prompt" in class_def_params):
-                inputs["prompt"] = {"variable_name": "prompt"}
-            if "extra_pnginfo" in hidden_inputs and (
-                no_params or "extra_pnginfo" in class_def_params
-            ):
-                inputs["extra_pnginfo"] = {"variable_name": "extra_pnginfo"}
-            if "hidden" not in input_types and class_def_params is not None:
-                if "unique_id" in class_def_params:
-                    inputs["unique_id"] = random.randint(1, 2**64)
-
-            # Create executed variable and generate code
-            executed_variables[idx] = (
-                f"{self.clean_variable_name(class_type)}_"
-                f"{self.sanitize_node_id(str(idx))}"
-            )
-            inputs = self.update_inputs(inputs, executed_variables)
-            seed_sync_code = self.create_prompt_seed_sync_code(
-                idx, inputs, is_special_function
-            )
-
-            if is_special_function:
-                if seed_sync_code:
-                    special_functions_code.extend(seed_sync_code)
-                special_functions_code.append(
-                    self.create_function_call_code(
-                        initialized_objects[class_type],
-                        class_def.FUNCTION,
-                        executed_variables[idx],
-                        is_special_function,
-                        **inputs,
-                    )
-                )
-            else:
-                if seed_sync_code:
-                    code.extend(seed_sync_code)
-                code.append(
-                    self.create_function_call_code(
-                        initialized_objects[class_type],
-                        class_def.FUNCTION,
-                        executed_variables[idx],
-                        is_special_function,
-                        **inputs,
-                    )
-                )
-
-        # Generate final code by combining imports and code, and wrap them in a main function
-        final_code = self.assemble_python_code(
-            import_statements,
-            special_functions_code,
-            code,
-            workflow_data,
-            metadata_workflow_data,
-            queue_size,
-            custom_nodes,
-        )
-
-        return final_code
-
-    def create_function_call_code(
-        self,
-        obj_name: str,
-        func: str,
-        variable_name: str,
-        is_special_function: bool,
-        **kwargs,
-    ) -> str:
-        """Generate Python code for a function call.
-
-        Args:
-            obj_name (str): The name of the initialized object.
-            func (str): The function to be called.
-            variable_name (str): The name of the variable that the function result should be assigned to.
-            is_special_function (bool): Determines the code indentation.
-            **kwargs: The keyword arguments for the function.
-
-        Returns:
-            str: The generated Python code.
-        """
-        args = ", ".join(self.format_arg(key, value) for key, value in kwargs.items())
-
-        # Generate the Python code
-        code = f"{variable_name} = {obj_name}.{func}({args})\n"
-
-        # If the code contains dependencies and is not a loader or encoder, indent the code because it will be placed inside
-        # of a for loop
-        if not is_special_function:
-            code = f"\t{code}"
-
-        return code
-
-    def create_prompt_seed_sync_code(
-        self, node_id: str, inputs: Dict, is_special_function: bool
-    ) -> List[str]:
-        """Generate code that keeps prompt metadata aligned with randomized seeds."""
-        seed_sync_lines = []
-        for key in ("seed", "noise_seed"):
-            if key not in inputs:
-                continue
-            randomized_seed_variable = (
-                f"node_{self.sanitize_node_id(str(node_id))}_{self.clean_variable_name(key)}"
-            )
-            seed_sync_lines.append(
-                f'{randomized_seed_variable} = prompt["{node_id}"]["inputs"]["{key}"] = random.randint(1, 2**64)'
-            )
-            inputs[key] = {"variable_name": randomized_seed_variable}
-
-        if not seed_sync_lines:
-            return []
-
-        indentation = "" if is_special_function else "\t"
-        return [f"{indentation}{line}\n" for line in seed_sync_lines]
-
-    def format_arg(self, key: str, value: any) -> str:
-        """Formats arguments based on key and value.
-
-        Args:
-            key (str): Argument key.
-            value (any): Argument value.
-
-        Returns:
-            str: Formatted argument as a string.
-        """
-        value_code = self.format_arg_value(key, value)
-        if key.isidentifier() and not keyword.iskeyword(key):
-            return f"{key}={value_code}"
-        return f"**{{{json.dumps(key)}: {value_code}}}"
-
-    @staticmethod
-    def format_arg_value(key: str, value: any) -> str:
-        """Formats an argument value as Python source."""
-        if isinstance(value, dict) and "variable_name" in value:
-            return value["variable_name"]
-        if key == "noise_seed" or key == "seed":
-            return "random.randint(1, 2**64)"
-        if isinstance(value, str):
-            return json.dumps(value)
-        return repr(value)
-
-    def assemble_python_code(
-        self,
-        import_statements: Dict[str, set],
-        speical_functions_code: List[str],
-        code: List[str],
-        workflow_data: Dict,
-        metadata_workflow_data: Dict | None,
-        queue_size: int,
-        custom_nodes=False,
-    ) -> str:
-        """Generates the final code string.
-
-        Args:
-            import_statements (Dict[str, set]): Import statements grouped by module.
-            speical_functions_code (List[str]): A list of special functions code strings.
-            code (List[str]): A list of code strings.
-            workflow_data (Dict): The API workflow data used for runtime prompt execution.
-            metadata_workflow_data (Dict | None): The workflow metadata to embed into saved outputs.
-            queue_size (int): Number of photos that will be generated by the script.
-            custom_nodes (bool): Whether to include custom nodes in the code.
-
-        Returns:
-            str: Generated final code as a string.
-        """
-        if metadata_workflow_data is None:
-            extra_pnginfo_code = "extra_pnginfo = None"
-        else:
-            extra_pnginfo_code = (
-                "extra_pnginfo = "
-                '{"workflow": json.loads('
-                + json.dumps(json.dumps(metadata_workflow_data))
-                + ")}"
-            )
-
-        # Get the source code of the utils functions as a string
-        func_strings = []
-        for func in [
-            get_value_at_index,
-            get_comfyui_path,
-            find_path,
-            add_comfyui_directory_to_sys_path,
-            add_extra_model_paths,
-        ]:
-            func_strings.append(f"\n{inspect.getsource(func)}")
-        # Define static import statements required for the script
-        static_imports = (
-            [
-                "import json",
-                "import os",
-                "import random",
-                "import sys",
-                "from typing import Sequence, Mapping, Any, Union",
-                "import torch",
-            ]
-            + func_strings
-            + [
-                "\n\nadd_comfyui_directory_to_sys_path()\nadd_extra_model_paths()\n",
-                f"workflow = json.loads({json.dumps(json.dumps(workflow_data))})",
-                "prompt = json.loads(json.dumps(workflow))",
-                extra_pnginfo_code,
-            ]
-        )
-        # Check if custom nodes should be included
-        if custom_nodes:
-            static_imports.append(f"\n{inspect.getsource(import_custom_nodes)}\n")
-            custom_nodes = "import_custom_nodes()\n\t"
-        else:
-            custom_nodes = ""
-        # Create import statements for node classes
-        imports_code = []
-        for module_name in sorted(import_statements.keys()):
-            class_names = ", ".join(sorted(import_statements[module_name]))
-            imports_code.append(f"from {module_name} import {class_names}")
-        special_functions_body = "\n\t\t".join(speical_functions_code) or "pass"
-        loop_body = "\n\t\t".join(code) or "\tpass"
-        # Assemble the main function code, including custom nodes if applicable
-        main_function_code = (
-            "def main():\n\t"
-            + f"{custom_nodes}with torch.inference_mode():\n\t\t"
-            + special_functions_body
-            + f"\n\n\t\tfor q in range({queue_size}):\n\t\t"
-            + loop_body
-        )
-        # Concatenate all parts to form the final code
-        final_code = "\n".join(
-            static_imports
-            + imports_code
-            + ["", main_function_code, "", 'if __name__ == "__main__":', "\tmain()"]
-        )
-        # Format the final code according to PEP 8 using the Black library
-        final_code = black.format_str(final_code, mode=black.Mode())
-
-        return final_code
-
-    def get_class_info(self, class_type: str) -> Tuple[str, Tuple[str, str], str]:
-        """Generates and returns necessary information about class type.
-
-        Args:
-            class_type (str): Class type.
-
-        Returns:
-            Tuple[str, Tuple[str, str], str]: Updated class type, import statement and initialization code.
-        """
-        class_obj = self.base_node_class_mappings.get(class_type)
-        module_name = "nodes"
-        if class_obj is not None:
-            module_name = class_obj.__module__
-        variable_name = self.clean_variable_name(class_type)
-        is_importable_module = bool(
-            module_name
-            and "/" not in module_name
-            and "\\" not in module_name
-            and all(part.isidentifier() for part in module_name.split("."))
-        )
-        if class_type in self.base_node_class_mappings.keys() and is_importable_module:
-            import_statement = (module_name, class_type)
-            class_code = f"{variable_name} = {class_type.strip()}()"
-        else:
-            import_statement = ("nodes", "NODE_CLASS_MAPPINGS")
-            class_code = f'{variable_name} = NODE_CLASS_MAPPINGS["{class_type}"]()'
-
-        return class_type, import_statement, class_code
-
-    @staticmethod
-    def clean_variable_name(class_type: str) -> str:
-        """
-        Remove any characters from variable name that could cause errors running the Python script.
-
-        Args:
-            class_type (str): Class type.
-
-        Returns:
-            str: Cleaned variable name with no special characters or spaces
-        """
-        # Convert to lowercase and replace spaces with underscores
-        clean_name = class_type.lower().strip().replace("-", "_").replace(" ", "_")
-
-        # Remove characters that are not letters, numbers, or underscores
-        clean_name = re.sub(r"[^a-z0-9_]", "", clean_name)
-
-        # Ensure that it doesn't start with a number
-        if clean_name[0].isdigit():
-            clean_name = "_" + clean_name
-
-        return clean_name
-
-    def get_function_parameters(self, func: Callable) -> List:
-        """Get the names of a function's parameters.
-
-        Args:
-            func (Callable): The function whose parameters we want to inspect.
-
-        Returns:
-            List: A list containing the names of the function's parameters.
-        """
-        signature = inspect.signature(func)
-        parameters = {
-            name: param.default if param.default != param.empty else None
-            for name, param in signature.parameters.items()
-        }
-        catch_all = any(
-            param.kind == inspect.Parameter.VAR_KEYWORD
-            for param in signature.parameters.values()
-        )
-        return list(parameters.keys()) if not catch_all else None
-
-    def update_inputs(self, inputs: Dict, executed_variables: Dict) -> Dict:
-        """Update inputs based on the executed variables.
-
-        Args:
-            inputs (Dict): Inputs dictionary to update.
-            executed_variables (Dict): Dictionary storing executed variable names.
-
-        Returns:
-            Dict: Updated inputs dictionary.
-        """
-        for key in inputs.keys():
-            if (
-                isinstance(inputs[key], list)
-                and inputs[key][0] in executed_variables.keys()
-            ):
-                inputs[key] = {
-                    "variable_name": f"get_value_at_index({executed_variables[inputs[key][0]]}, {inputs[key][1]})"
-                }
-        return inputs
-
-
-class ComfyUItoPython:
-    """Main workflow to generate Python code from a workflow_api.json file.
-
-    Attributes:
-        input_file (str): Path to the input JSON file.
-        output_file (str): Path to the output Python file.
-        queue_size (int): The number of photos that will be created by the script.
-        node_class_mappings (Dict): Mappings of node classes.
-        base_node_class_mappings (Dict): Base mappings of node classes.
-    """
-
-    def __init__(
-        self,
-        workflow: str = "",
-        frontend_workflow: str | Dict | None = None,
-        input_file: str = "",
-        output_file: str | TextIO = "",
-        queue_size: int = 1,
-        node_class_mappings: Dict | None = None,
-        needs_init_custom_nodes: bool = False,
-    ):
-        """Initialize the ComfyUItoPython class with the given parameters. Exactly one of workflow or input_file must be specified.
-        Args:
-            workflow (str): The workflow's JSON.
-            input_file (str): Path to the input JSON file.
-            output_file (str | TextIO): Path to the output file or a file-like object.
-            queue_size (int): The number of times a workflow will be executed by the script. Defaults to 1.
-            node_class_mappings (Dict | None): Mappings of node classes. Defaults to the current
-                ComfyUI NODE_CLASS_MAPPINGS when not provided.
-            needs_init_custom_nodes (bool): Whether to initialize custom nodes. Defaults to False.
-        """
-        if input_file and workflow:
-            raise ValueError("Can't provide both input_file and workflow")
-        elif not input_file and not workflow:
-            raise ValueError("Needs input_file or workflow")
-
-        if not output_file:
-            raise ValueError("Needs output_file")
-
-        self.workflow = workflow
-        self.frontend_workflow = frontend_workflow
-        self.input_file = input_file
-        self.output_file = output_file
-        self.queue_size = queue_size
-        self.node_class_mappings = (
-            node_class_mappings
-            if node_class_mappings is not None
-            else get_node_class_mappings()
-        )
-        self.needs_init_custom_nodes = needs_init_custom_nodes
-
-        self.base_node_class_mappings = copy.deepcopy(self.node_class_mappings)
-        self.execute()
-
-    def execute(self):
-        """Execute the main workflow to generate Python code.
-
-        Returns:
-            None
-        """
-        # Step 1: Read JSON data from the input file
-        if self.input_file:
-            data = FileHandler.read_json_file(self.input_file)
-        else:
-            data = json.loads(self.workflow)
-
-        metadata_workflow_data = None
-        if self.frontend_workflow:
-            if isinstance(self.frontend_workflow, str):
-                metadata_workflow_data = json.loads(self.frontend_workflow)
-            else:
-                metadata_workflow_data = self.frontend_workflow
-
-        # Step 2: Initialize extra/custom nodes when requested or when the workflow references
-        # a node class that is not currently loaded in the runtime.
-        missing_node_types = {
-            node_data["class_type"]
-            for node_data in data.values()
-            if node_data["class_type"] not in self.node_class_mappings
-        }
-        if self.needs_init_custom_nodes or missing_node_types:
-            import_custom_nodes()
-            self.base_node_class_mappings = copy.deepcopy(self.node_class_mappings)
-
-        # Step 3: Determine the load order
-        load_order_determiner = LoadOrderDeterminer(data, self.node_class_mappings)
-        load_order = load_order_determiner.determine_load_order()
-
-        # Step 4: Generate the workflow code
-        code_generator = CodeGenerator(
-            self.node_class_mappings, self.base_node_class_mappings
-        )
-        generated_code = code_generator.generate_workflow(
-            load_order,
-            data,
-            metadata_workflow_data,
-            queue_size=self.queue_size,
-        )
-
-        # Step 5: Write the generated code to a file
-        FileHandler.write_code_to_file(self.output_file, generated_code)
-
-        print(f"Code successfully generated and written to {self.output_file}")
-
-
-def run(
-    input_file: str = DEFAULT_INPUT_FILE,
-    output_file: str = DEFAULT_OUTPUT_FILE,
-    queue_size: int = DEFAULT_QUEUE_SIZE,
-) -> None:
-    """Generate Python code from a ComfyUI workflow_api.json file.
-
-    Args:
-        input_file (str): Path to the input JSON file. Defaults to "workflow_api.json".
-        output_file (str): Path to the output Python file.
-            Defaults to "workflow_api.py".
-        queue_size (int): The number of times a workflow will be executed by the script.
-            Defaults to 1.
-
-    Returns:
-        None
-    """
-    ComfyUItoPython(
-        input_file=input_file,
-        output_file=output_file,
-        queue_size=queue_size,
-        needs_init_custom_nodes=True,
-    )
-
-
-def main() -> None:
-    """Main function to generate Python code from a ComfyUI workflow_api.json file."""
-    parser = ArgumentParser(
-        description="Generate Python    code from a ComfyUI workflow_api.json file."
-    )
-    parser.add_argument(
-        "-f",
-        "--input_file",
-        type=str,
-        help="path to the input JSON file",
-        default=DEFAULT_INPUT_FILE,
-    )
-    parser.add_argument(
-        "-o",
-        "--output_file",
-        type=str,
-        help="path to the output Python file",
-        default=DEFAULT_OUTPUT_FILE,
-    )
-    parser.add_argument(
-        "-q",
-        "--queue_size",
-        type=int,
-        help="number of times the workflow will be executed by default",
-        default=DEFAULT_QUEUE_SIZE,
-    )
-    pargs = parser.parse_args()
-    run(**vars(pargs))
-    print("Done.")
+from comfyui_to_python.cli import main
 
 
 if __name__ == "__main__":
-    """Run the main function."""
     main()

comfyui_to_python/__init__.py

@@ -0,0 +1,55 @@
+from typing import TextIO
+
+from .app import ExportApplication
+from .cli import DEFAULT_INPUT_FILE, DEFAULT_OUTPUT_FILE, DEFAULT_QUEUE_SIZE, main
+from .node_runtime import get_node_class_mappings, import_custom_nodes
+
+
+class ComfyUItoPython:
+    """Public compatibility facade for the exporter package."""
+
+    def __init__(
+        self,
+        workflow: str = "",
+        frontend_workflow: str | dict | None = None,
+        input_file: str = "",
+        output_file: str | TextIO = "",
+        queue_size: int = 1,
+        node_class_mappings: dict | None = None,
+        needs_init_custom_nodes: bool = False,
+    ):
+        self._app = ExportApplication(
+            workflow=workflow,
+            frontend_workflow=frontend_workflow,
+            input_file=input_file,
+            output_file=output_file,
+            queue_size=queue_size,
+            node_class_mappings=node_class_mappings,
+            needs_init_custom_nodes=needs_init_custom_nodes,
+            node_mapping_loader=get_node_class_mappings,
+            custom_node_importer=import_custom_nodes,
+        )
+        self._app.execute()
+
+
+def run(
+    input_file: str = DEFAULT_INPUT_FILE,
+    output_file: str = DEFAULT_OUTPUT_FILE,
+    queue_size: int = DEFAULT_QUEUE_SIZE,
+) -> None:
+    """Generate Python code from a ComfyUI workflow_api.json file."""
+    ComfyUItoPython(
+        input_file=input_file,
+        output_file=output_file,
+        queue_size=queue_size,
+        needs_init_custom_nodes=True,
+    )
+
+
+__all__ = [
+    "ComfyUItoPython",
+    "run",
+    "main",
+    "get_node_class_mappings",
+    "import_custom_nodes",
+]

comfyui_to_python/__main__.py

@@ -0,0 +1,5 @@
+from .cli import main
+
+
+if __name__ == "__main__":
+    main()

comfyui_to_python/app.py

@@ -0,0 +1,75 @@
+import copy
+
+from typing import TextIO
+
+from .generator.planner import WorkflowPlanner
+from .generator.render import WorkflowRenderer
+from .io import write_python_output
+from .load_order import LoadOrderDeterminer
+from .workflow_loader import load_frontend_workflow_data, load_workflow_data
+
+
+class ExportApplication:
+    """High-level exporter orchestration."""
+
+    def __init__(
+        self,
+        workflow: str = "",
+        frontend_workflow: str | dict | None = None,
+        input_file: str = "",
+        output_file: str | TextIO = "",
+        queue_size: int = 1,
+        node_class_mappings: dict | None = None,
+        needs_init_custom_nodes: bool = False,
+        node_mapping_loader=None,
+        custom_node_importer=None,
+    ):
+        if input_file and workflow:
+            raise ValueError("Can't provide both input_file and workflow")
+        if not input_file and not workflow:
+            raise ValueError("Needs input_file or workflow")
+        if not output_file:
+            raise ValueError("Needs output_file")
+
+        self.workflow = workflow
+        self.frontend_workflow = frontend_workflow
+        self.input_file = input_file
+        self.output_file = output_file
+        self.queue_size = queue_size
+        self.node_mapping_loader = node_mapping_loader
+        self.custom_node_importer = custom_node_importer
+        self.node_class_mappings = (
+            node_class_mappings
+            if node_class_mappings is not None
+            else self.node_mapping_loader()
+        )
+        self.needs_init_custom_nodes = needs_init_custom_nodes
+        self.base_node_class_mappings = copy.deepcopy(self.node_class_mappings)
+
+    def execute(self) -> None:
+        data = load_workflow_data(self.workflow, self.input_file)
+        metadata_workflow_data = load_frontend_workflow_data(self.frontend_workflow)
+
+        missing_node_types = {
+            node_data["class_type"]
+            for node_data in data.values()
+            if node_data["class_type"] not in self.node_class_mappings
+        }
+        if self.needs_init_custom_nodes or missing_node_types:
+            self.custom_node_importer()
+            self.base_node_class_mappings = copy.deepcopy(self.node_class_mappings)
+
+        load_order = LoadOrderDeterminer(
+            data, self.node_class_mappings
+        ).determine_load_order()
+        plan = WorkflowPlanner(
+            self.node_class_mappings, self.base_node_class_mappings
+        ).build_plan(
+            load_order,
+            data,
+            metadata_workflow_data,
+            queue_size=self.queue_size,
+        )
+        generated_code = WorkflowRenderer().render(plan)
+        write_python_output(self.output_file, generated_code)
+        print(f"Code successfully generated and written to {self.output_file}")

comfyui_to_python/cli.py

@@ -0,0 +1,42 @@
+from argparse import ArgumentParser
+
+DEFAULT_INPUT_FILE = "workflow_api.json"
+DEFAULT_OUTPUT_FILE = "workflow_api.py"
+DEFAULT_QUEUE_SIZE = 10
+
+
+def build_argument_parser() -> ArgumentParser:
+    parser = ArgumentParser(
+        description="Generate Python    code from a ComfyUI workflow_api.json file."
+    )
+    parser.add_argument(
+        "-f",
+        "--input_file",
+        type=str,
+        help="path to the input JSON file",
+        default=DEFAULT_INPUT_FILE,
+    )
+    parser.add_argument(
+        "-o",
+        "--output_file",
+        type=str,
+        help="path to the output Python file",
+        default=DEFAULT_OUTPUT_FILE,
+    )
+    parser.add_argument(
+        "-q",
+        "--queue_size",
+        type=int,
+        help="number of times the workflow will be executed by default",
+        default=DEFAULT_QUEUE_SIZE,
+    )
+    return parser
+
+
+def main() -> None:
+    from . import run
+
+    parser = build_argument_parser()
+    pargs = parser.parse_args()
+    run(**vars(pargs))
+    print("Done.")

comfyui_to_python/generator/__init__.py

@@ -0,0 +1,4 @@
+from .planner import WorkflowPlanner
+from .render import WorkflowRenderer
+
+__all__ = ["WorkflowPlanner", "WorkflowRenderer"]

comfyui_to_python/generator/generated_helpers.py

@@ -0,0 +1,19 @@
+from ..node_runtime import (
+    add_comfyui_directory_to_sys_path,
+    add_extra_model_paths,
+    bootstrap_comfyui_runtime,
+    cleanup_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",
+    "cleanup_comfyui_runtime",
+    "find_path",
+    "get_comfyui_path",
+    "get_value_at_index",
+]

comfyui_to_python/generator/model.py

@@ -0,0 +1,12 @@
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class GenerationPlan:
+    import_statements: dict[str, set[str]]
+    special_functions_code: list[str]
+    loop_code: list[str]
+    workflow_data: dict
+    metadata_workflow_data: dict | None
+    queue_size: int
+    custom_nodes: bool

comfyui_to_python/generator/planner.py

@@ -0,0 +1,261 @@
+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()
+            input_value_types = self.get_input_value_types(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, input_value_types, 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,
+                    input_value_types=input_value_types,
+                    **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,
+        input_value_types: dict[str, str] | None = None,
+        **kwargs,
+    ) -> str:
+        args = ", ".join(
+            self.format_arg(key, value, (input_value_types or {}).get(key))
+            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,
+        input_value_types: dict[str, str],
+        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)}"
+            )
+            randomized_seed_code = self.get_randomized_seed_code(
+                input_value_types.get(key)
+            )
+            seed_sync_lines.append(
+                f'{randomized_seed_variable} = prompt["{node_id}"]["inputs"]["{key}"] = {randomized_seed_code}'
+            )
+            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, input_value_type: str | None = None) -> str:
+        value_code = self.format_arg_value(key, value, input_value_type)
+        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, input_value_type: str | None = None
+    ) -> str:
+        if isinstance(value, dict) and "variable_name" in value:
+            return value["variable_name"]
+        if key == "noise_seed" or key == "seed":
+            return WorkflowPlanner.get_randomized_seed_code(input_value_type)
+        if isinstance(value, str):
+            return json.dumps(value)
+        return repr(value)
+
+    @staticmethod
+    def get_input_value_types(input_types: dict) -> dict[str, str]:
+        value_types = {}
+        for section in ("required", "optional", "hidden"):
+            for key, value in input_types.get(section, {}).items():
+                if isinstance(value, tuple) and value:
+                    value_types[key] = value[0]
+                elif isinstance(value, str):
+                    value_types[key] = value
+        return value_types
+
+    @staticmethod
+    def get_randomized_seed_code(input_value_type: str | None) -> str:
+        randomized_seed_code = "random.randint(1, 2**64)"
+        if input_value_type == "STRING":
+            return f"str({randomized_seed_code})"
+        return randomized_seed_code
+
+    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

comfyui_to_python/generator/render.py

@@ -0,0 +1,148 @@
+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,
+    cleanup_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,
+            cleanup_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(unload_models: bool | None = None):",
+            "    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",
+                "",
+                "    try:",
+                "        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()
+        )
+        execution_section.extend(
+            [
+                "    finally:",
+                "        cleanup_comfyui_runtime(unload_models=unload_models)",
+            ]
+        )
+
+        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()

comfyui_to_python/io.py

@@ -0,0 +1,23 @@
+import json
+import os
+from typing import TextIO
+
+
+def load_json_input(file_path: str | TextIO, encoding: str = "utf-8") -> dict:
+    """Read workflow JSON from a file path or file-like object."""
+    if hasattr(file_path, "read"):
+        return json.load(file_path)
+    with open(file_path, "r", encoding=encoding) as file:
+        return json.load(file)
+
+
+def write_python_output(file_path: str | TextIO, code: str) -> None:
+    """Write generated Python to a file path or file-like object."""
+    if isinstance(file_path, str):
+        directory = os.path.dirname(file_path)
+        if directory and not os.path.exists(directory):
+            os.makedirs(directory)
+        with open(file_path, "w", encoding="utf-8") as file:
+            file.write(code)
+        return
+    file_path.write(code)

comfyui_to_python/load_order.py

@@ -0,0 +1,43 @@
+from typing import Dict
+
+
+class LoadOrderDeterminer:
+    """Determine workflow execution order with loader-like nodes prioritized."""
+
+    def __init__(self, data: Dict, node_class_mappings: Dict):
+        self.data = data
+        self.node_class_mappings = node_class_mappings
+        self.visited = {}
+        self.load_order = []
+        self.is_special_function = False
+
+    def determine_load_order(self) -> list[tuple[str, Dict, bool]]:
+        self._load_special_functions_first()
+        self.is_special_function = False
+        for key in self.data:
+            if key not in self.visited:
+                self._dfs(key)
+        return self.load_order
+
+    def _dfs(self, key: str) -> None:
+        self.visited[key] = True
+        inputs = self.data[key]["inputs"]
+        for value in inputs.values():
+            if isinstance(value, list) and value[0] not in self.visited:
+                self._dfs(value[0])
+        self.load_order.append((key, self.data[key], self.is_special_function))
+
+    def _load_special_functions_first(self) -> None:
+        for key in self.data:
+            class_def = self.node_class_mappings[self.data[key]["class_type"]]()
+            if (
+                class_def.CATEGORY == "loaders"
+                or class_def.FUNCTION in ["encode"]
+                or not any(
+                    isinstance(value, list)
+                    for value in self.data[key]["inputs"].values()
+                )
+            ):
+                self.is_special_function = True
+                if key not in self.visited:
+                    self._dfs(key)

comfyui_to_python/node_runtime.py

@@ -0,0 +1,172 @@
+import os
+import sys
+import warnings
+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 cleanup_comfyui_runtime(unload_models: bool | None = None) -> None:
+    """Best-effort cleanup for embedded or repeated generated-script execution."""
+    import gc
+
+    def run_cleanup_hook(name: str, should_run: bool = True) -> None:
+        if not should_run or not hasattr(model_management, name):
+            return
+        cleanup_fn = getattr(model_management, name)
+        try:
+            cleanup_fn()
+        except Exception as exc:
+            warnings.warn(
+                f"ComfyUI cleanup hook {name} failed during teardown: {exc}",
+                RuntimeWarning,
+                stacklevel=2,
+            )
+
+    should_unload = unload_models
+    if should_unload is None:
+        should_unload = os.environ.get("COMFYUI_TOPYTHON_UNLOAD_MODELS", "").lower() in {
+            "1",
+            "true",
+            "yes",
+            "on",
+        }
+
+    try:
+        import comfy.model_management as model_management
+    except ModuleNotFoundError:
+        gc.collect()
+        return
+
+    run_cleanup_hook("cleanup_models_gc")
+    run_cleanup_hook("unload_all_models", should_run=should_unload)
+    run_cleanup_hook("soft_empty_cache")
+    gc.collect()
+
+
+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]

comfyui_to_python/workflow_loader.py

@@ -0,0 +1,19 @@
+import json
+
+from .io import load_json_input
+
+
+def load_workflow_data(workflow: str, input_file: str):
+    """Load workflow data from inline JSON or an input file."""
+    if input_file:
+        return load_json_input(input_file)
+    return json.loads(workflow)
+
+
+def load_frontend_workflow_data(frontend_workflow: str | dict | None):
+    """Load optional frontend workflow metadata."""
+    if not frontend_workflow:
+        return None
+    if isinstance(frontend_workflow, str):
+        return json.loads(frontend_workflow)
+    return frontend_workflow

comfyui_to_python_utils.py

@@ -1,123 +1,19 @@
-import os
-from typing import Sequence, Mapping, Any, Union
-import sys
+from comfyui_to_python.node_runtime import (
+    add_comfyui_directory_to_sys_path,
+    add_extra_model_paths,
+    bootstrap_comfyui_runtime,
+    find_path,
+    get_comfyui_path,
+    get_value_at_index,
+    import_custom_nodes,
+)
 
-
-def get_comfyui_path() -> str:
-    """Return the configured ComfyUI path, preferring COMFYUI_PATH when set."""
-    comfyui_path = os.environ.get("COMFYUI_PATH")
-    if comfyui_path:
-        return comfyui_path
-    return find_path("ComfyUI")
-
-
-def import_custom_nodes() -> None:
-    """Find all custom nodes in the custom_nodes folder and add those node objects to NODE_CLASS_MAPPINGS
-
-    This function sets up a new asyncio event loop, initializes the PromptServer,
-    creates a PromptQueue, and initializes the custom nodes.
-    """
-    comfyui_path = get_comfyui_path()
-    if comfyui_path and comfyui_path not in sys.path:
-        sys.path.insert(0, comfyui_path)
-
-    import asyncio
-    import execution
-    from nodes import init_extra_nodes
-    if comfyui_path in sys.path:
-        sys.path.remove(comfyui_path)
-    sys.path.insert(0, comfyui_path)
-    import server
-
-    # Creating a new event loop and setting it as the default loop
-    loop = asyncio.new_event_loop()
-    asyncio.set_event_loop(loop)
-
-    # Creating an instance of PromptServer with the loop
-    server_instance = server.PromptServer(loop)
-    execution.PromptQueue(server_instance)
-
-    # Initializing custom nodes
-    asyncio.run(init_extra_nodes())
-
-
-def find_path(name: str, path: str = None) -> str:
-    """
-    Recursively looks at parent folders starting from the given path until it finds the given name.
-    Returns the path as a Path object if found, or None otherwise.
-    """
-    # If no path is given, use the current working directory
-    if path is None:
-        path = os.getcwd()
-
-    # Check if the current directory contains the name
-    if name in os.listdir(path):
-        path_name = os.path.join(path, name)
-        print(f"{name} found: {path_name}")
-        return path_name
-
-    # Get the parent directory
-    parent_directory = os.path.dirname(path)
-
-    # If the parent directory is the same as the current directory, we've reached the root and stop the search
-    if parent_directory == path:
-        return None
-
-    # Recursively call the function with the parent directory
-    return find_path(name, parent_directory)
-
-
-def add_comfyui_directory_to_sys_path() -> None:
-    """
-    Add 'ComfyUI' to the sys.path
-    """
-    comfyui_path = get_comfyui_path()
-    if comfyui_path is not None and os.path.isdir(comfyui_path):
-        if comfyui_path in sys.path:
-            sys.path.remove(comfyui_path)
-        sys.path.insert(0, comfyui_path)
-        print(f"'{comfyui_path}' added to sys.path")
-
-
-def add_extra_model_paths() -> None:
-    """
-    Parse the optional extra_model_paths.yaml file and add the parsed paths to the sys.path.
-    """
-    try:
-        from main import load_extra_path_config
-    except ImportError:
-        print(
-            "Could not import load_extra_path_config from main.py. Looking in utils.extra_config instead."
-        )
-        from utils.extra_config import load_extra_path_config
-
-    extra_model_paths = find_path("extra_model_paths.yaml")
-
-    if extra_model_paths is not None:
-        load_extra_path_config(extra_model_paths)
-    else:
-        print("Could not find the extra_model_paths config file.")
-
-
-def get_value_at_index(obj: Union[Sequence, Mapping], index: int) -> Any:
-    """Returns the value at the given index of a sequence or mapping.
-
-    If the object is a sequence (like list or string), returns the value at the given index.
-    If the object is a mapping (like a dictionary), returns the value at the index-th key.
-
-    Some return a dictionary, in these cases, we look for the "results" key
-
-    Args:
-        obj (Union[Sequence, Mapping]): The object to retrieve the value from.
-        index (int): The index of the value to retrieve.
-
-    Returns:
-        Any: The value at the given index.
-
-    Raises:
-        IndexError: If the index is out of bounds for the object and the object is not a mapping.
-    """
-    try:
-        return obj[index]
-    except KeyError:
-        return obj["result"][index]
+__all__ = [
+    "add_comfyui_directory_to_sys_path",
+    "add_extra_model_paths",
+    "bootstrap_comfyui_runtime",
+    "find_path",
+    "get_comfyui_path",
+    "get_value_at_index",
+    "import_custom_nodes",
+]

js/save-as-script.js

@@ -1,6 +1,25 @@
 import { api } from "../../scripts/api.js";
 import { app } from "../../scripts/app.js";
-import { $el } from "../../scripts/ui.js";
+
+const DEFAULT_SCRIPT_FILENAME = "workflow_api.py";
+const DEFAULT_WORKFLOW_NAME = "workflow_api.json";
+
+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",
@@ -23,15 +42,12 @@ const extension = {
 		});
 	},
 	savePythonScript() {
-		var filename = prompt("Save script as:");
-		if(filename === undefined || filename === null || filename === "") {
-			return
-		}
-		
+		const filename = DEFAULT_SCRIPT_FILENAME;
+
 		app.graphToPrompt().then(async (p) => {
 			const frontendWorkflow = p.workflow ?? app.graph.serialize();
 			const json = JSON.stringify({
-				name: filename + ".json",
+				name: DEFAULT_WORKFLOW_NAME,
 				workflow: JSON.stringify(p.output, null, 2),
 				frontend_workflow: JSON.stringify(frontendWorkflow, null, 2),
 			}, null, 2); // convert the data to a JSON string
@@ -39,10 +55,6 @@ const extension = {
 			if(response.status == 200) {
 				const blob = new Blob([await response.text()], {type: "text/python;charset=utf-8"});
 				const url = URL.createObjectURL(blob);
-				if(!filename.endsWith(".py")) {
-					filename += ".py";
-				}
-
 				const a = $el("a", {
 					href: url,
 					download: filename,

pyproject.toml

@@ -1,8 +1,9 @@
 [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.1.0"
 license = { text = "MIT License" }
+requires-python = ">=3.12"
 dependencies = ["black"]
 
 [project.urls]

tests/README.md

@@ -5,6 +5,8 @@ This directory contains two test paths:
 - unit tests for exporter behavior
 - runtime validation for committed workflow fixtures
 
+Use the lightweight path for routine contributor validation, and use the runtime path when you need end-to-end confidence against a real ComfyUI checkout.
+
 ## Prerequisites
 
 - Run commands from the repo root.
@@ -30,12 +32,18 @@ The current runtime-capable fixtures require these model files in the target Com
 
 ## Unit Tests
 
-Run the current unit test module:
+Run the exporter-focused unit test module:
 
 ```bash
 uv run python -m unittest tests.test_upscale_model_loader_export
 ```
 
+Run the runtime-harness unit test module:
+
+```bash
+uv run python -m unittest tests.test_runtime_validation_harness
+```
+
 Run all `unittest`-discoverable tests under `tests`:
 
 ```bash
@@ -49,6 +57,7 @@ The runtime harness lives at `tests/runtime/run_runtime_validation.py`.
 ### Fast Tier
 
 Fast tier validates export behavior against committed fixtures without requiring a full ComfyUI runtime for every fixture.
+This is the recommended default validation lane for routine changes.
 
 Run all fast-tier compatible fixtures:
 
@@ -65,6 +74,7 @@ uv run python tests/runtime/run_runtime_validation.py --tier fast --fixture unsa
 ### Runtime Tier
 
 Runtime tier exports inside a real ComfyUI checkout and executes generated Python for runtime-capable fixtures.
+This is the heavier validation lane for changes that need end-to-end runtime confidence.
 
 Run all runtime-capable fixtures:
 
@@ -92,11 +102,15 @@ Current committed fixtures:
 - `text-to-image`
 - `unsafe-kwargs`
 - `subgraph-identifiers`
+- `reused-node-class-branches`
+- `secondary-output-selection`
 
 Notes:
 
 - `--tier runtime` only runs fixtures marked runtime-capable.
 - `--tier fast` only runs fixtures with local test mappings.
+- `reused-node-class-branches` protects repeated node-class usage and branch wiring in the fast tier.
+- `secondary-output-selection` protects non-zero output index wiring in the fast tier.
 
 ## Troubleshooting
 

tests/fixtures/runtime/reused-node-class-branches.json

@@ -0,0 +1,27 @@
+{
+  "1:10": {
+    "class_type": "PassthroughText",
+    "inputs": {
+      "text": "left branch"
+    }
+  },
+  "1:20": {
+    "class_type": "PassthroughText",
+    "inputs": {
+      "text": "right branch"
+    }
+  },
+  "2": {
+    "class_type": "JoinText",
+    "inputs": {
+      "left": [
+        "1:10",
+        0
+      ],
+      "right": [
+        "1:20",
+        0
+      ]
+    }
+  }
+}

tests/fixtures/runtime/secondary-output-selection.json

@@ -0,0 +1,17 @@
+{
+  "1": {
+    "class_type": "SplitText",
+    "inputs": {
+      "text": "alpha|omega"
+    }
+  },
+  "2": {
+    "class_type": "PassthroughText",
+    "inputs": {
+      "text": [
+        "1",
+        1
+      ]
+    }
+  }
+}

tests/fixtures/unit/generator_codegen/string-seed-node.json

@@ -0,0 +1,8 @@
+{
+  "1": {
+    "class_type": "StringSeedNode",
+    "inputs": {
+      "seed": "seed-placeholder"
+    }
+  }
+}

tests/fixtures/unit/generator_codegen/subgraph-upscaler-identifiers.json

@@ -0,0 +1,27 @@
+{
+  "33": {
+    "class_type": "VaeDecode",
+    "inputs": {
+      "samples": "latent-placeholder"
+    }
+  },
+  "42:0": {
+    "class_type": "UpscaleModelLoader",
+    "inputs": {
+      "model_name": "4x-ultrasharp.safetensors"
+    }
+  },
+  "42:1": {
+    "class_type": "ImageUpscaleWithModel",
+    "inputs": {
+      "upscale_model": [
+        "42:0",
+        0
+      ],
+      "image": [
+        "33",
+        0
+      ]
+    }
+  }
+}

tests/fixtures/unit/generator_codegen/trailing-backslash-string.json

@@ -0,0 +1,10 @@
+{
+  "1": {
+    "class_type": "TextConcatenateNode",
+    "inputs": {
+      "delimiter": "",
+      "clean_whitespace": "true",
+      "text_b": "\\"
+    }
+  }
+}

tests/fixtures/unit/generator_codegen/unsafe-rgthree-kwargs.json

@@ -0,0 +1,41 @@
+{
+  "11": {
+    "class_type": "DualClipLoader",
+    "inputs": {
+      "clip_name": "clip.safetensors"
+    }
+  },
+  "633": {
+    "class_type": "AnySwitchRgthree",
+    "inputs": {
+      "model": "model-placeholder"
+    }
+  },
+  "631": {
+    "class_type": "PowerLoraLoaderRgthree",
+    "inputs": {
+      "PowerLoraLoaderHeaderWidget": {
+        "type": "PowerLoraLoaderHeaderWidget"
+      },
+      "lora_1": {
+        "on": false,
+        "lora": "lora.safetensors",
+        "strength": 1.2
+      },
+      "lora_2": {
+        "on": false,
+        "lora": "lora2.safetensors",
+        "strength": 0.7
+      },
+      "\u2795 Add Lora": "",
+      "model": [
+        "633",
+        0
+      ],
+      "clip": [
+        "11",
+        0
+      ]
+    }
+  }
+}

tests/fixtures/unit/generator_codegen/windows-path-string.json

@@ -0,0 +1,8 @@
+{
+  "1": {
+    "class_type": "WindowsPathNode",
+    "inputs": {
+      "path": "C:\\ComfyUI\\models\\upscale_models\\RealESRGAN_x4plus.safetensors"
+    }
+  }
+}

tests/runtime/generated/text-to-image.py

@@ -0,0 +1,275 @@
+# 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 cleanup_comfyui_runtime(unload_models: bool | None = None) -> None:
+    """Best-effort cleanup for embedded or repeated generated-script execution."""
+    import gc
+
+    def run_cleanup_hook(name: str, should_run: bool = True) -> None:
+        if not should_run or not hasattr(model_management, name):
+            return
+        cleanup_fn = getattr(model_management, name)
+        try:
+            cleanup_fn()
+        except Exception as exc:
+            warnings.warn(
+                f"ComfyUI cleanup hook {name} failed during teardown: {exc}",
+                RuntimeWarning,
+                stacklevel=2,
+            )
+
+    should_unload = unload_models
+    if should_unload is None:
+        should_unload = os.environ.get(
+            "COMFYUI_TOPYTHON_UNLOAD_MODELS", ""
+        ).lower() in {
+            "1",
+            "true",
+            "yes",
+            "on",
+        }
+
+    try:
+        import comfy.model_management as model_management
+    except ModuleNotFoundError:
+        gc.collect()
+        return
+
+    run_cleanup_hook("cleanup_models_gc")
+    run_cleanup_hook("unload_all_models", should_run=should_unload)
+    run_cleanup_hook("soft_empty_cache")
+    gc.collect()
+
+
+# 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(unload_models: bool | None = None):
+    bootstrap_comfyui_runtime()
+    add_extra_model_paths()
+
+    # Node imports
+    from nodes import (
+        CLIPTextEncode,
+        CheckpointLoaderSimple,
+        EmptyLatentImage,
+        KSampler,
+        NODE_CLASS_MAPPINGS,
+        SaveImage,
+        VAEDecode,
+    )
+
+    import torch
+
+    try:
+        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,
+                )
+    finally:
+        cleanup_comfyui_runtime(unload_models=unload_models)
+
+
+# Entrypoint
+if __name__ == "__main__":
+    main()

tests/runtime/generated/upscale-model-loader.py

@@ -0,0 +1,230 @@
+# 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 cleanup_comfyui_runtime(unload_models: bool | None = None) -> None:
+    """Best-effort cleanup for embedded or repeated generated-script execution."""
+    import gc
+
+    should_unload = unload_models
+    if should_unload is None:
+        should_unload = os.environ.get(
+            "COMFYUI_TOPYTHON_UNLOAD_MODELS", ""
+        ).lower() in {
+            "1",
+            "true",
+            "yes",
+            "on",
+        }
+
+    try:
+        import comfy.model_management as model_management
+    except ModuleNotFoundError:
+        gc.collect()
+        return
+
+    if hasattr(model_management, "cleanup_models_gc"):
+        model_management.cleanup_models_gc()
+    if should_unload and hasattr(model_management, "unload_all_models"):
+        model_management.unload_all_models()
+    if hasattr(model_management, "soft_empty_cache"):
+        model_management.soft_empty_cache()
+    gc.collect()
+
+
+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(unload_models: bool | None = None):
+    bootstrap_comfyui_runtime()
+    add_extra_model_paths()
+    import_custom_nodes()
+
+    # Node imports
+    from nodes import LoadImage, NODE_CLASS_MAPPINGS, SaveImage
+
+    import torch
+
+    try:
+        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,
+                )
+    finally:
+        cleanup_comfyui_runtime(unload_models=unload_models)
+
+
+# Entrypoint
+if __name__ == "__main__":
+    main()

tests/runtime/run_runtime_validation.py

@@ -109,6 +109,40 @@ class PassthroughText:
         return (text,)
 
 
+class SplitText:
+    CATEGORY = "utils"
+    FUNCTION = "split"
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "text": ("STRING",),
+            }
+        }
+
+    def split(self, text):
+        left, right = text.split("|", 1)
+        return (left, right)
+
+
+class JoinText:
+    CATEGORY = "utils"
+    FUNCTION = "join"
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "left": ("STRING",),
+                "right": ("STRING",),
+            }
+        }
+
+    def join(self, left, right):
+        return (f"{left}::{right}",)
+
+
 class UpscaleModelLoader:
     CATEGORY = "loaders"
     FUNCTION = "load_model"
@@ -234,6 +268,22 @@ FIXTURES = {
             "PassthroughText": PassthroughText,
         },
     ),
+    "reused-node-class-branches": FixtureConfig(
+        name="reused-node-class-branches",
+        path=FIXTURE_DIR / "reused-node-class-branches.json",
+        mapping_factory=lambda: {
+            "PassthroughText": PassthroughText,
+            "JoinText": JoinText,
+        },
+    ),
+    "secondary-output-selection": FixtureConfig(
+        name="secondary-output-selection",
+        path=FIXTURE_DIR / "secondary-output-selection.json",
+        mapping_factory=lambda: {
+            "SplitText": SplitText,
+            "PassthroughText": PassthroughText,
+        },
+    ),
 }
 
 

tests/test_generator_codegen_issue_regressions.py

@@ -0,0 +1,312 @@
+import json
+import unittest
+from io import StringIO
+from pathlib import Path
+
+from comfyui_to_python import ComfyUItoPython
+
+
+FIXTURE_DIR = Path(__file__).parent / "fixtures" / "unit" / "generator_codegen"
+
+
+class AnySwitchRgthree:
+    CATEGORY = "utils"
+    FUNCTION = "switch"
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "model": ("MODEL",),
+            }
+        }
+
+    def switch(self, model):
+        return (model,)
+
+
+class DualClipLoader:
+    CATEGORY = "loaders"
+    FUNCTION = "load_clip"
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "clip_name": ("STRING",),
+            }
+        }
+
+    def load_clip(self, clip_name):
+        return (clip_name,)
+
+
+class PowerLoraLoaderRgthree:
+    CATEGORY = "loaders"
+    FUNCTION = "load_loras"
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "PowerLoraLoaderHeaderWidget": ("DICT",),
+                "model": ("MODEL",),
+                "clip": ("CLIP",),
+            }
+        }
+
+    def load_loras(self, **kwargs):
+        return (kwargs,)
+
+
+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 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 VaeDecode:
+    CATEGORY = "latent"
+    FUNCTION = "decode"
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "samples": ("LATENT",),
+            }
+        }
+
+    def decode(self, samples):
+        return (samples,)
+
+
+class WindowsPathNode:
+    CATEGORY = "paths"
+    FUNCTION = "open_path"
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "path": ("STRING",),
+            }
+        }
+
+    def open_path(self, path):
+        return (path,)
+
+
+class TextConcatenateNode:
+    CATEGORY = "text"
+    FUNCTION = "text_concatenate"
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "delimiter": ("STRING",),
+                "clean_whitespace": ("STRING",),
+                "text_b": ("STRING",),
+            }
+        }
+
+    def text_concatenate(self, delimiter, clean_whitespace, text_b):
+        return (delimiter, clean_whitespace, text_b)
+
+
+class StringSeedNode:
+    CATEGORY = "sampling"
+    FUNCTION = "sample"
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "seed": ("STRING",),
+            },
+            "hidden": {
+                "prompt": "PROMPT",
+            },
+        }
+
+    def sample(self, seed, prompt):
+        return (seed, prompt)
+
+
+def load_fixture(name: str) -> dict:
+    return json.loads((FIXTURE_DIR / name).read_text(encoding="utf-8"))
+
+
+def export_workflow(workflow: dict, node_class_mappings: dict) -> str:
+    output = StringIO()
+    ComfyUItoPython(
+        workflow=json.dumps(workflow),
+        output_file=output,
+        node_class_mappings=node_class_mappings,
+    )
+    return output.getvalue()
+
+
+class GeneratorCodegenIssueRegressionTest(unittest.TestCase):
+    def test_export_uses_dictionary_expansion_for_rgthree_symbol_heavy_input_names(self):
+        generated = export_workflow(
+            load_fixture("unsafe-rgthree-kwargs.json"),
+            {
+                "AnySwitchRgthree": AnySwitchRgthree,
+                "DualClipLoader": DualClipLoader,
+                "PowerLoraLoaderRgthree": PowerLoraLoaderRgthree,
+            },
+        )
+
+        self.assertIn(
+            'powerloraloaderrgthree_631 = powerloraloaderrgthree.load_loras(',
+            generated,
+        )
+        self.assertIn(
+            'PowerLoraLoaderHeaderWidget={"type": "PowerLoraLoaderHeaderWidget"}',
+            generated,
+        )
+        self.assertIn('**{"\\u2795 Add Lora": ""}', generated)
+        self.assertNotIn('➕ Add Lora=""', generated)
+
+    def test_export_sanitizes_subgraph_identifiers_for_upscaler_workflows(self):
+        generated = export_workflow(
+            load_fixture("subgraph-upscaler-identifiers.json"),
+            {
+                "VaeDecode": VaeDecode,
+                "UpscaleModelLoader": UpscaleModelLoader,
+                "ImageUpscaleWithModel": ImageUpscaleWithModel,
+            },
+        )
+
+        self.assertIn("upscalemodelloader_42_0 = upscalemodelloader.load_model(", generated)
+        self.assertIn(
+            "imageupscalewithmodel_42_1 = imageupscalewithmodel.upscale(",
+            generated,
+        )
+        self.assertIn(
+            "upscale_model=get_value_at_index(upscalemodelloader_42_0, 0)",
+            generated,
+        )
+        self.assertNotIn("imageupscalewithmodel_42:1", generated)
+        self.assertNotIn("upscalemodelloader_42:0", generated)
+
+    def test_export_preserves_windows_style_model_paths(self):
+        generated = export_workflow(
+            load_fixture("windows-path-string.json"),
+            {
+                "WindowsPathNode": WindowsPathNode,
+            },
+        )
+
+        globals_dict = {"__name__": "generated_workflow_module"}
+        exec(generated, globals_dict)
+
+        self.assertEqual(
+            globals_dict["build_workflow"]()["1"]["inputs"]["path"],
+            r"C:\ComfyUI\models\upscale_models\RealESRGAN_x4plus.safetensors",
+        )
+
+    def test_export_preserves_trailing_backslash_string_literals(self):
+        generated = export_workflow(
+            load_fixture("trailing-backslash-string.json"),
+            {
+                "TextConcatenateNode": TextConcatenateNode,
+            },
+        )
+
+        globals_dict = {"__name__": "generated_workflow_module"}
+        exec(generated, globals_dict)
+
+        self.assertEqual(
+            globals_dict["build_workflow"]()["1"]["inputs"]["text_b"],
+            "\\",
+        )
+
+    def test_export_randomizes_string_seed_inputs_as_strings(self):
+        generated = export_workflow(
+            load_fixture("string-seed-node.json"),
+            {
+                "StringSeedNode": StringSeedNode,
+            },
+        )
+
+        self.assertIn(
+            'node_1_seed = prompt["1"]["inputs"]["seed"] = str(random.randint(1, 2**64))',
+            generated,
+        )
+        self.assertIn("seed=node_1_seed", generated)
+
+    def test_issue_cluster_regressions_render_parseable_python(self):
+        workflows = [
+            (
+                    load_fixture("unsafe-rgthree-kwargs.json"),
+                    {
+                        "AnySwitchRgthree": AnySwitchRgthree,
+                        "DualClipLoader": DualClipLoader,
+                        "PowerLoraLoaderRgthree": PowerLoraLoaderRgthree,
+                    },
+                ),
+            (
+                load_fixture("subgraph-upscaler-identifiers.json"),
+                {
+                    "VaeDecode": VaeDecode,
+                    "UpscaleModelLoader": UpscaleModelLoader,
+                    "ImageUpscaleWithModel": ImageUpscaleWithModel,
+                },
+            ),
+            (
+                load_fixture("trailing-backslash-string.json"),
+                {
+                    "TextConcatenateNode": TextConcatenateNode,
+                },
+            ),
+            (
+                load_fixture("windows-path-string.json"),
+                {
+                    "WindowsPathNode": WindowsPathNode,
+                },
+            ),
+            (
+                load_fixture("string-seed-node.json"),
+                {
+                    "StringSeedNode": StringSeedNode,
+                },
+            ),
+        ]
+
+        for workflow, mapping in workflows:
+            generated = export_workflow(workflow, mapping)
+            compile(generated, "<generated_workflow>", "exec")
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_node_runtime_cleanup.py

@@ -0,0 +1,116 @@
+import sys
+import types
+import unittest
+import warnings
+from unittest.mock import Mock, patch
+
+from comfyui_to_python.node_runtime import cleanup_comfyui_runtime
+
+
+class NodeRuntimeCleanupTest(unittest.TestCase):
+    def test_cleanup_releases_caches_without_forcing_model_unload_by_default(self):
+        comfy_module = types.ModuleType("comfy")
+        comfy_module.__path__ = []
+        model_management = types.ModuleType("comfy.model_management")
+        model_management.cleanup_models_gc = Mock()
+        model_management.unload_all_models = Mock()
+        model_management.soft_empty_cache = Mock()
+        comfy_module.model_management = model_management
+
+        with patch.dict(
+            sys.modules,
+            {
+                "comfy": comfy_module,
+                "comfy.model_management": model_management,
+            },
+        ), patch.dict("os.environ", {}, clear=False):
+            cleanup_comfyui_runtime()
+
+        model_management.cleanup_models_gc.assert_called_once_with()
+        model_management.soft_empty_cache.assert_called_once_with()
+        model_management.unload_all_models.assert_not_called()
+
+    def test_cleanup_can_force_model_unload_from_environment(self):
+        comfy_module = types.ModuleType("comfy")
+        comfy_module.__path__ = []
+        model_management = types.ModuleType("comfy.model_management")
+        model_management.cleanup_models_gc = Mock()
+        model_management.unload_all_models = Mock()
+        model_management.soft_empty_cache = Mock()
+        comfy_module.model_management = model_management
+
+        with patch.dict(
+            sys.modules,
+            {
+                "comfy": comfy_module,
+                "comfy.model_management": model_management,
+            },
+        ), patch.dict(
+            "os.environ",
+            {"COMFYUI_TOPYTHON_UNLOAD_MODELS": "true"},
+            clear=False,
+        ):
+            cleanup_comfyui_runtime()
+
+        model_management.unload_all_models.assert_called_once_with()
+
+    def test_cleanup_suppresses_hook_failures_and_warns(self):
+        comfy_module = types.ModuleType("comfy")
+        comfy_module.__path__ = []
+        model_management = types.ModuleType("comfy.model_management")
+        model_management.cleanup_models_gc = Mock(side_effect=RuntimeError("gc failed"))
+        model_management.unload_all_models = Mock(side_effect=RuntimeError("unload failed"))
+        model_management.soft_empty_cache = Mock(side_effect=RuntimeError("cache failed"))
+        comfy_module.model_management = model_management
+
+        with patch.dict(
+            sys.modules,
+            {
+                "comfy": comfy_module,
+                "comfy.model_management": model_management,
+            },
+        ), warnings.catch_warnings(record=True) as caught:
+            warnings.simplefilter("always")
+            cleanup_comfyui_runtime(unload_models=True)
+
+        model_management.cleanup_models_gc.assert_called_once_with()
+        model_management.unload_all_models.assert_called_once_with()
+        model_management.soft_empty_cache.assert_called_once_with()
+        self.assertEqual(len(caught), 3)
+        self.assertEqual(
+            [str(warning.message) for warning in caught],
+            [
+                "ComfyUI cleanup hook cleanup_models_gc failed during teardown: gc failed",
+                "ComfyUI cleanup hook unload_all_models failed during teardown: unload failed",
+                "ComfyUI cleanup hook soft_empty_cache failed during teardown: cache failed",
+            ],
+        )
+
+    def test_cleanup_does_not_mask_active_workflow_exception(self):
+        comfy_module = types.ModuleType("comfy")
+        comfy_module.__path__ = []
+        model_management = types.ModuleType("comfy.model_management")
+        model_management.cleanup_models_gc = Mock(side_effect=RuntimeError("cleanup failed"))
+        model_management.soft_empty_cache = Mock()
+        comfy_module.model_management = model_management
+
+        with patch.dict(
+            sys.modules,
+            {
+                "comfy": comfy_module,
+                "comfy.model_management": model_management,
+            },
+        ), warnings.catch_warnings(record=True):
+            warnings.simplefilter("always")
+            with self.assertRaisesRegex(ValueError, "workflow failed"):
+                try:
+                    raise ValueError("workflow failed")
+                finally:
+                    cleanup_comfyui_runtime()
+
+        model_management.cleanup_models_gc.assert_called_once_with()
+        model_management.soft_empty_cache.assert_called_once_with()
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_project_contracts.py

@@ -0,0 +1,33 @@
+import tomllib
+import unittest
+from pathlib import Path
+
+
+REPO_ROOT = Path(__file__).resolve().parent.parent
+
+
+class ProjectContractsTest(unittest.TestCase):
+    def test_project_declares_supported_python_floor(self):
+        pyproject = tomllib.loads((REPO_ROOT / "pyproject.toml").read_text(encoding="utf-8"))
+
+        self.assertEqual(pyproject["project"]["requires-python"], ">=3.12")
+
+    def test_readme_documents_python_support_and_default_save_filename(self):
+        readme = (REPO_ROOT / "README.md").read_text(encoding="utf-8")
+
+        self.assertIn("This project supports Python 3.12 and newer.", readme)
+        self.assertIn("default filename `workflow_api.py`", readme)
+
+    def test_extension_import_path_requires_uv_sync_instead_of_running_install_py(self):
+        init_text = (REPO_ROOT / "__init__.py").read_text(encoding="utf-8")
+
+        self.assertIn("Run 'uv sync'", init_text)
+        self.assertNotIn("spec_from_file_location", init_text)
+        self.assertNotIn("Successfully installed. Hopefully, at least.", init_text)
+
+    def test_frontend_save_flow_uses_deterministic_filename_without_prompt(self):
+        save_as_script = (REPO_ROOT / "js" / "save-as-script.js").read_text(encoding="utf-8")
+
+        self.assertIn('const DEFAULT_SCRIPT_FILENAME = "workflow_api.py";', save_as_script)
+        self.assertNotIn("prompt(", save_as_script)
+

tests/test_runtime_validation_harness.py

@@ -0,0 +1,305 @@
+import struct
+import tempfile
+import unittest
+import zlib
+from pathlib import Path
+from unittest.mock import patch
+
+from tests.runtime.run_runtime_validation import (
+    FixtureConfig,
+    ModelRequirement,
+    ValidationFailure,
+    check_models,
+    ensure_runtime_path,
+    execute_generated_python,
+    load_fixture_names,
+    parse_png_info,
+    validate_generated_python,
+)
+
+
+def make_png_bytes(
+    width: int,
+    height: int,
+    text_chunks: list[tuple[bytes, bytes]] | None = None,
+) -> bytes:
+    def chunk(chunk_type: bytes, data: bytes) -> bytes:
+        crc = zlib.crc32(chunk_type + data) & 0xFFFFFFFF
+        return (
+            struct.pack(">I", len(data))
+            + chunk_type
+            + data
+            + struct.pack(">I", crc)
+        )
+
+    ihdr = chunk(b"IHDR", struct.pack(">IIBBBBB", width, height, 8, 2, 0, 0, 0))
+    text_chunks = text_chunks or []
+    idat = chunk(
+        b"IDAT",
+        zlib.compress(b"\x00" + (b"\x00\x00\x00" * width)),
+    )
+    return b"".join(
+        [
+            b"\x89PNG\r\n\x1a\n",
+            ihdr,
+            *[chunk(chunk_type, payload) for chunk_type, payload in text_chunks],
+            idat,
+            chunk(b"IEND", b""),
+        ]
+    )
+
+
+class RuntimeValidationHarnessTest(unittest.TestCase):
+    def test_load_fixture_names_all_returns_registered_names(self):
+        fixture_names = load_fixture_names("all")
+
+        self.assertGreaterEqual(len(fixture_names), 4)
+        self.assertEqual(fixture_names[0], "upscale-model-loader")
+        self.assertIn("unsafe-kwargs", fixture_names)
+
+    def test_load_fixture_names_rejects_unknown_fixture(self):
+        with self.assertRaises(ValidationFailure) as context:
+            load_fixture_names("missing-fixture")
+
+        self.assertEqual(context.exception.classification, "fixture bug")
+        self.assertIn("Unknown fixture", context.exception.message)
+
+    def test_ensure_runtime_path_fast_tier_returns_env_without_validation(self):
+        with patch.dict("os.environ", {"COMFYUI_PATH": "/does/not/exist"}, clear=False):
+            runtime_path = ensure_runtime_path("fast")
+
+        self.assertEqual(runtime_path, "/does/not/exist")
+
+    @patch("tests.runtime.run_runtime_validation.get_comfyui_path", return_value="")
+    def test_ensure_runtime_path_runtime_tier_requires_valid_checkout(self, _mock_path):
+        with self.assertRaises(ValidationFailure) as context:
+            ensure_runtime_path("runtime")
+
+        self.assertEqual(context.exception.classification, "environment/setup failure")
+        self.assertIn("Could not find a valid ComfyUI checkout", context.exception.message)
+
+    def test_check_models_returns_only_missing_requirements(self):
+        with tempfile.TemporaryDirectory() as tmpdir:
+            runtime_path = Path(tmpdir)
+            present_dir = runtime_path / "models" / "checkpoints"
+            present_dir.mkdir(parents=True)
+            (present_dir / "present.safetensors").write_text("ok", encoding="utf-8")
+            fixture = FixtureConfig(
+                name="model-checks",
+                path=Path("unused.json"),
+                model_requirements=(
+                    ModelRequirement(
+                        filename="present.safetensors",
+                        relative_dir="models/checkpoints",
+                        source_url="https://example.invalid/present",
+                    ),
+                    ModelRequirement(
+                        filename="missing.safetensors",
+                        relative_dir="models/checkpoints",
+                        source_url="https://example.invalid/missing",
+                    ),
+                ),
+            )
+
+            missing = check_models(fixture, str(runtime_path))
+
+        self.assertEqual([item.filename for item in missing], ["missing.safetensors"])
+
+    def test_validate_generated_python_accepts_valid_python(self):
+        validate_generated_python("value = 1\n", "valid-fixture")
+
+    def test_validate_generated_python_reports_syntax_regression(self):
+        with self.assertRaises(ValidationFailure) as context:
+            validate_generated_python("def broken(:\n", "broken-fixture")
+
+        self.assertEqual(context.exception.classification, "repo regression")
+        self.assertIn("broken-fixture", context.exception.message)
+
+    def test_parse_png_info_reads_dimensions_and_text_chunks(self):
+        compressed_text = zlib.compress(b"workflow data")
+        png_bytes = make_png_bytes(
+            width=3,
+            height=2,
+            text_chunks=[
+                (b"tEXt", b"prompt\x00hello"),
+                (b"zTXt", b"workflow\x00\x00" + compressed_text),
+                (b"iTXt", b"comment\x00\x00\x00\x00\x00unicode text"),
+            ],
+        )
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            image_path = Path(tmpdir) / "artifact.png"
+            image_path.write_bytes(png_bytes)
+
+            width, height, metadata = parse_png_info(image_path)
+
+        self.assertEqual((width, height), (3, 2))
+        self.assertEqual(metadata["prompt"], "hello")
+        self.assertEqual(metadata["workflow"], "workflow data")
+        self.assertEqual(metadata["comment"], "unicode text")
+
+    def test_parse_png_info_rejects_non_png_files(self):
+        with tempfile.TemporaryDirectory() as tmpdir:
+            image_path = Path(tmpdir) / "artifact.bin"
+            image_path.write_bytes(b"not-a-png")
+
+            with self.assertRaises(ValidationFailure) as context:
+                parse_png_info(image_path)
+
+        self.assertEqual(context.exception.classification, "environment/setup failure")
+        self.assertIn("Expected PNG output", context.exception.message)
+
+    def test_parse_png_info_requires_dimensions(self):
+        with tempfile.TemporaryDirectory() as tmpdir:
+            image_path = Path(tmpdir) / "artifact.png"
+            image_path.write_bytes(
+                b"\x89PNG\r\n\x1a\n"
+                + struct.pack(">I", 0)
+                + b"IEND"
+                + struct.pack(">I", 0)
+            )
+
+            with self.assertRaises(ValidationFailure) as context:
+                parse_png_info(image_path)
+
+        self.assertEqual(context.exception.classification, "environment/setup failure")
+        self.assertIn("Could not read PNG dimensions", context.exception.message)
+
+    @patch("tests.runtime.run_runtime_validation.get_runtime_python", return_value="/usr/bin/python")
+    @patch("tests.runtime.run_runtime_validation.subprocess.run")
+    def test_execute_generated_python_classifies_missing_torch_as_environment_failure(
+        self,
+        mock_run,
+        _mock_runtime_python,
+    ):
+        mock_run.return_value.returncode = 1
+        mock_run.return_value.stderr = "ModuleNotFoundError: No module named 'torch'"
+        mock_run.return_value.stdout = ""
+        fixture = FixtureConfig(
+            name="runtime-fixture",
+            path=Path("unused.json"),
+            runtime_capable=True,
+        )
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            (Path(tmpdir) / "output").mkdir()
+            with self.assertRaises(ValidationFailure) as context:
+                execute_generated_python("print('hello')\n", fixture, tmpdir)
+
+        self.assertEqual(context.exception.classification, "environment/setup failure")
+
+    @patch("tests.runtime.run_runtime_validation.get_runtime_python", return_value="/usr/bin/python")
+    @patch("tests.runtime.run_runtime_validation.subprocess.run")
+    def test_execute_generated_python_classifies_missing_files_as_environment_failure(
+        self,
+        mock_run,
+        _mock_runtime_python,
+    ):
+        mock_run.return_value.returncode = 1
+        mock_run.return_value.stderr = "No such file or directory: missing.png"
+        mock_run.return_value.stdout = ""
+        fixture = FixtureConfig(
+            name="runtime-fixture",
+            path=Path("unused.json"),
+            runtime_capable=True,
+        )
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            (Path(tmpdir) / "output").mkdir()
+            with self.assertRaises(ValidationFailure) as context:
+                execute_generated_python("print('hello')\n", fixture, tmpdir)
+
+        self.assertEqual(context.exception.classification, "environment/setup failure")
+
+    @patch("tests.runtime.run_runtime_validation.get_runtime_python", return_value="/usr/bin/python")
+    @patch("tests.runtime.run_runtime_validation.subprocess.run")
+    def test_execute_generated_python_classifies_other_failures_as_repo_regression(
+        self,
+        mock_run,
+        _mock_runtime_python,
+    ):
+        mock_run.return_value.returncode = 1
+        mock_run.return_value.stderr = "ValueError: broken workflow"
+        mock_run.return_value.stdout = ""
+        fixture = FixtureConfig(
+            name="runtime-fixture",
+            path=Path("unused.json"),
+            runtime_capable=True,
+        )
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            (Path(tmpdir) / "output").mkdir()
+            with self.assertRaises(ValidationFailure) as context:
+                execute_generated_python("print('hello')\n", fixture, tmpdir)
+
+        self.assertEqual(context.exception.classification, "repo regression")
+
+    @patch("tests.runtime.run_runtime_validation.get_runtime_python", return_value="/usr/bin/python")
+    @patch("tests.runtime.run_runtime_validation.subprocess.run")
+    def test_execute_generated_python_requires_fresh_matching_artifact(
+        self,
+        mock_run,
+        _mock_runtime_python,
+    ):
+        mock_run.return_value.returncode = 0
+        mock_run.return_value.stderr = ""
+        mock_run.return_value.stdout = ""
+        fixture = FixtureConfig(
+            name="runtime-fixture",
+            path=Path("unused.json"),
+            runtime_capable=True,
+            filename_prefix="expected_prefix",
+        )
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            (Path(tmpdir) / "output").mkdir()
+            with self.assertRaises(ValidationFailure) as context:
+                execute_generated_python("print('hello')\n", fixture, tmpdir)
+
+        self.assertEqual(context.exception.classification, "repo regression")
+        self.assertIn("did not produce a new output file", context.exception.message)
+
+    @patch("tests.runtime.run_runtime_validation.get_runtime_python", return_value="/usr/bin/python")
+    @patch("tests.runtime.run_runtime_validation.validate_output_artifact")
+    @patch("tests.runtime.run_runtime_validation.subprocess.run")
+    def test_execute_generated_python_validates_newest_matching_artifact(
+        self,
+        mock_run,
+        mock_validate_output,
+        _mock_runtime_python,
+    ):
+        mock_run.return_value.returncode = 0
+        mock_run.return_value.stderr = ""
+        mock_run.return_value.stdout = ""
+        fixture = FixtureConfig(
+            name="runtime-fixture",
+            path=Path("unused.json"),
+            runtime_capable=True,
+            filename_prefix="expected_prefix",
+        )
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            output_dir = Path(tmpdir) / "output"
+            output_dir.mkdir()
+            older = output_dir / "expected_prefix_00001_.png"
+            older.write_bytes(make_png_bytes(1, 1))
+            with patch(
+                "tests.runtime.run_runtime_validation.subprocess.run",
+                side_effect=self._write_runtime_artifact(mock_run.return_value, output_dir),
+            ):
+                execute_generated_python("print('hello')\n", fixture, tmpdir)
+
+        validated_path = mock_validate_output.call_args[0][1]
+        self.assertEqual(validated_path.name, "expected_prefix_00002_.png")
+
+    @staticmethod
+    def _write_runtime_artifact(result, output_dir: Path):
+        def side_effect(*_args, **_kwargs):
+            (output_dir / "expected_prefix_00002_.png").write_bytes(make_png_bytes(2, 2))
+            return result
+
+        return side_effect
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_upscale_model_loader_export.py

@@ -1,4 +1,5 @@
 import json
+import sys
 import tempfile
 import unittest
 from io import StringIO
@@ -115,6 +116,102 @@ 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("def cleanup_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(unload_models: bool | None = None):\n"
+            "    bootstrap_comfyui_runtime()\n"
+            "    add_extra_model_paths()",
+            generated,
+        )
+        self.assertLess(
+            generated.index("def bootstrap_comfyui_runtime()"),
+            generated.index("def main(unload_models: bool | None = None):"),
+        )
+        main_section = generated[
+            generated.index("def main(unload_models: bool | None = None):") :
+        ]
+        self.assertIn(
+            "def main(unload_models: bool | None = None):\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"))
+        self.assertIn(
+            "    finally:\n        cleanup_comfyui_runtime(unload_models=unload_models)",
+            main_section,
+        )
+
+    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 +248,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 +293,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 +321,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 +395,68 @@ 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(unload_models: bool | None = None)", generated)
+        self.assertIn("bootstrap_comfyui_runtime()", generated)
+        self.assertIn("cleanup_comfyui_runtime(unload_models=unload_models)", generated)
+        self.assertNotIn("def initialize_workflow()", generated)
+        self.assertNotIn("def run_once(", generated)
+        self.assertIn("with torch.inference_mode():", generated)
+        self.assertIn("finally:", 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(unload_models: bool | None = None)"),
+        )
+
     def test_hidden_metadata_kwargs_follow_function_signature(self):
         workflow = {
             "1": {

uv.lock

@@ -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.1.0"
 source = { virtual = "." }
 dependencies = [
     { name = "black" },