diff --git a/README.md b/README.md index b80470b..d4ce020 100644 --- a/README.md +++ b/README.md @@ -79,18 +79,24 @@ Notes: 3. Run the exporter: ```bash -uv run python comfyui_to_python.py +uv run python -m comfyui_to_python ``` Options: ```bash -uv run python comfyui_to_python.py \ +uv run python -m comfyui_to_python \ --input_file workflow_api.json \ --output_file workflow_api.py \ --queue_size 10 ``` +The legacy wrapper still works if you prefer it: + +```bash +uv run python comfyui_to_python.py +``` + Flags: - `--input_file`: input workflow JSON, default `workflow_api.json` - `--output_file`: output Python file, default `workflow_api.py` diff --git a/comfyui_to_python.py b/comfyui_to_python.py index efae605..4ff8101 100644 --- a/comfyui_to_python.py +++ b/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() diff --git a/comfyui_to_python/__init__.py b/comfyui_to_python/__init__.py new file mode 100644 index 0000000..18aafa1 --- /dev/null +++ b/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", +] diff --git a/comfyui_to_python/__main__.py b/comfyui_to_python/__main__.py new file mode 100644 index 0000000..2f05ddc --- /dev/null +++ b/comfyui_to_python/__main__.py @@ -0,0 +1,5 @@ +from .cli import main + + +if __name__ == "__main__": + main() diff --git a/comfyui_to_python/app.py b/comfyui_to_python/app.py new file mode 100644 index 0000000..334d06f --- /dev/null +++ b/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}") diff --git a/comfyui_to_python/cli.py b/comfyui_to_python/cli.py new file mode 100644 index 0000000..458d4c8 --- /dev/null +++ b/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.") diff --git a/comfyui_to_python/generator/__init__.py b/comfyui_to_python/generator/__init__.py new file mode 100644 index 0000000..c7be5af --- /dev/null +++ b/comfyui_to_python/generator/__init__.py @@ -0,0 +1,4 @@ +from .planner import WorkflowPlanner +from .render import WorkflowRenderer + +__all__ = ["WorkflowPlanner", "WorkflowRenderer"] diff --git a/comfyui_to_python/generator/generated_helpers.py b/comfyui_to_python/generator/generated_helpers.py new file mode 100644 index 0000000..88f5f2f --- /dev/null +++ b/comfyui_to_python/generator/generated_helpers.py @@ -0,0 +1,17 @@ +from ..node_runtime import ( + add_comfyui_directory_to_sys_path, + add_extra_model_paths, + bootstrap_comfyui_runtime, + find_path, + get_comfyui_path, + get_value_at_index, +) + +__all__ = [ + "add_comfyui_directory_to_sys_path", + "add_extra_model_paths", + "bootstrap_comfyui_runtime", + "find_path", + "get_comfyui_path", + "get_value_at_index", +] diff --git a/comfyui_to_python/generator/model.py b/comfyui_to_python/generator/model.py new file mode 100644 index 0000000..ffe5e89 --- /dev/null +++ b/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 diff --git a/comfyui_to_python/generator/planner.py b/comfyui_to_python/generator/planner.py new file mode 100644 index 0000000..23922d9 --- /dev/null +++ b/comfyui_to_python/generator/planner.py @@ -0,0 +1,228 @@ +import inspect +import json +import keyword +import random +import re +from typing import Any, Callable + +from .model import GenerationPlan + + +class WorkflowPlanner: + """Convert ordered workflow nodes into a renderer-ready generation plan.""" + + def __init__(self, node_class_mappings: dict, base_node_class_mappings: dict): + self.node_class_mappings = node_class_mappings + self.base_node_class_mappings = base_node_class_mappings + + @staticmethod + def sanitize_node_id(node_id: str) -> str: + sanitized = re.sub(r"[^a-z0-9_]", "_", str(node_id).lower().strip()) + sanitized = re.sub(r"_+", "_", sanitized).strip("_") + if not sanitized: + sanitized = "node" + return sanitized + + @staticmethod + def clean_variable_name(class_type: str) -> str: + clean_name = class_type.lower().strip().replace("-", "_").replace(" ", "_") + clean_name = re.sub(r"[^a-z0-9_]", "", clean_name) + if clean_name[0].isdigit(): + clean_name = "_" + clean_name + return clean_name + + def build_plan( + self, + load_order: list[tuple[str, dict, bool]], + workflow_data: dict, + metadata_workflow_data: dict | None = None, + queue_size: int = 10, + ) -> GenerationPlan: + import_statements = {"nodes": {"NODE_CLASS_MAPPINGS"}} + executed_variables = {} + special_functions_code = [] + code = [] + initialized_objects = {} + custom_nodes = False + + for idx, data, is_special_function in load_order: + inputs, class_type = data["inputs"], data["class_type"] + input_types = self.node_class_mappings[class_type].INPUT_TYPES() + class_def = self.node_class_mappings[class_type]() + + missing_required_variable = False + if "required" in input_types.keys(): + for required in input_types["required"]: + if required not in inputs.keys(): + missing_required_variable = True + if missing_required_variable: + continue + + if class_type not in initialized_objects: + if class_type == "PreviewImage": + continue + + class_type, import_statement, class_code = self.get_class_info( + class_type + ) + initialized_objects[class_type] = self.clean_variable_name(class_type) + if class_type in self.base_node_class_mappings.keys(): + module_name, import_name = import_statement + import_statements.setdefault(module_name, set()).add(import_name) + if 'NODE_CLASS_MAPPINGS["' in class_code: + custom_nodes = True + special_functions_code.append(class_code) + + class_def_params = self.get_function_parameters( + getattr(class_def, class_def.FUNCTION) + ) + no_params = class_def_params is None + + inputs = { + key: value + for key, value in inputs.items() + if no_params or key in class_def_params + } + + hidden_inputs = input_types.get("hidden", {}) + if ( + "unique_id" in hidden_inputs + and (no_params or "unique_id" in class_def_params) + ): + inputs["unique_id"] = random.randint(1, 2**64) + if "prompt" in hidden_inputs and (no_params or "prompt" in class_def_params): + inputs["prompt"] = {"variable_name": "prompt"} + if "extra_pnginfo" in hidden_inputs and ( + no_params or "extra_pnginfo" in class_def_params + ): + inputs["extra_pnginfo"] = {"variable_name": "extra_pnginfo"} + if "hidden" not in input_types and class_def_params is not None: + if "unique_id" in class_def_params: + inputs["unique_id"] = random.randint(1, 2**64) + + executed_variables[idx] = ( + f"{self.clean_variable_name(class_type)}_" + f"{self.sanitize_node_id(str(idx))}" + ) + inputs = self.update_inputs(inputs, executed_variables) + seed_sync_code = self.create_prompt_seed_sync_code( + idx, inputs, is_special_function + ) + + target_lines = special_functions_code if is_special_function else code + if seed_sync_code: + target_lines.extend(seed_sync_code) + target_lines.append( + self.create_function_call_code( + initialized_objects[class_type], + class_def.FUNCTION, + executed_variables[idx], + is_special_function, + **inputs, + ) + ) + + return GenerationPlan( + import_statements=import_statements, + special_functions_code=special_functions_code, + loop_code=code, + workflow_data=workflow_data, + metadata_workflow_data=metadata_workflow_data, + queue_size=queue_size, + custom_nodes=custom_nodes, + ) + + def create_function_call_code( + self, + obj_name: str, + func: str, + variable_name: str, + is_special_function: bool, + **kwargs, + ) -> str: + args = ", ".join(self.format_arg(key, value) for key, value in kwargs.items()) + code = f"{variable_name} = {obj_name}.{func}({args})\n" + if not is_special_function: + code = f"\t{code}" + return code + + def create_prompt_seed_sync_code( + self, node_id: str, inputs: dict, is_special_function: bool + ) -> list[str]: + seed_sync_lines = [] + for key in ("seed", "noise_seed"): + if key not in inputs: + continue + randomized_seed_variable = ( + f"node_{self.sanitize_node_id(str(node_id))}_{self.clean_variable_name(key)}" + ) + seed_sync_lines.append( + f'{randomized_seed_variable} = prompt["{node_id}"]["inputs"]["{key}"] = random.randint(1, 2**64)' + ) + inputs[key] = {"variable_name": randomized_seed_variable} + + if not seed_sync_lines: + return [] + + indentation = "" if is_special_function else "\t" + return [f"{indentation}{line}\n" for line in seed_sync_lines] + + def format_arg(self, key: str, value: Any) -> str: + value_code = self.format_arg_value(key, value) + if key.isidentifier() and not keyword.iskeyword(key): + return f"{key}={value_code}" + return f"**{{{json.dumps(key)}: {value_code}}}" + + @staticmethod + def format_arg_value(key: str, value: Any) -> str: + if isinstance(value, dict) and "variable_name" in value: + return value["variable_name"] + if key == "noise_seed" or key == "seed": + return "random.randint(1, 2**64)" + if isinstance(value, str): + return json.dumps(value) + return repr(value) + + def get_class_info(self, class_type: str) -> tuple[str, tuple[str, str], str]: + class_obj = self.base_node_class_mappings.get(class_type) + module_name = "nodes" + if class_obj is not None: + module_name = class_obj.__module__ + variable_name = self.clean_variable_name(class_type) + is_importable_module = bool( + module_name + and "/" not in module_name + and "\\" not in module_name + and all(part.isidentifier() for part in module_name.split(".")) + ) + if class_type in self.base_node_class_mappings.keys() and is_importable_module: + import_statement = (module_name, class_type) + class_code = f"{variable_name} = {class_type.strip()}()" + else: + import_statement = ("nodes", "NODE_CLASS_MAPPINGS") + class_code = f'{variable_name} = NODE_CLASS_MAPPINGS["{class_type}"]()' + return class_type, import_statement, class_code + + @staticmethod + def get_function_parameters(func: Callable) -> list | None: + signature = inspect.signature(func) + parameters = { + name: param.default if param.default != param.empty else None + for name, param in signature.parameters.items() + } + catch_all = any( + param.kind == inspect.Parameter.VAR_KEYWORD + for param in signature.parameters.values() + ) + return list(parameters.keys()) if not catch_all else None + + def update_inputs(self, inputs: dict, executed_variables: dict) -> dict: + for key in inputs.keys(): + if ( + isinstance(inputs[key], list) + and inputs[key][0] in executed_variables.keys() + ): + inputs[key] = { + "variable_name": f"get_value_at_index({executed_variables[inputs[key][0]]}, {inputs[key][1]})" + } + return inputs diff --git a/comfyui_to_python/generator/render.py b/comfyui_to_python/generator/render.py new file mode 100644 index 0000000..a9fe5ae --- /dev/null +++ b/comfyui_to_python/generator/render.py @@ -0,0 +1,135 @@ +import inspect +from pprint import pformat +from typing import Any + +import black + +from ..node_runtime import import_custom_nodes +from .generated_helpers import ( + add_comfyui_directory_to_sys_path, + add_extra_model_paths, + bootstrap_comfyui_runtime, + find_path, + get_comfyui_path, + get_value_at_index, +) +from .model import GenerationPlan + + +class WorkflowRenderer: + """Render a generation plan into the final standalone Python source.""" + + def render(self, plan: GenerationPlan) -> str: + workflow_literal = self.format_python_literal(plan.workflow_data) + if plan.metadata_workflow_data is None: + extra_pnginfo_literal = "None" + else: + extra_pnginfo_literal = self.format_python_literal( + {"workflow": plan.metadata_workflow_data} + ) + + func_strings = [] + for func in [ + get_value_at_index, + get_comfyui_path, + find_path, + add_comfyui_directory_to_sys_path, + add_extra_model_paths, + bootstrap_comfyui_runtime, + ]: + func_strings.append(f"\n{inspect.getsource(func)}") + + static_imports = [ + "# Imports", + "import json", + "import os", + "import random", + "import sys", + "from typing import Sequence, Mapping, Any, Union", + ] + func_strings + static_imports.append("\n# Runtime support") + static_imports.append( + "\nbootstrap_comfyui_runtime()\nimport torch\nadd_extra_model_paths()\n" + ) + + 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"] + if imports_code: + execution_section.extend(["# Node imports", *imports_code, ""]) + execution_section.extend( + [ + "def main():", + " with torch.inference_mode():", + ] + ) + if custom_nodes_call: + execution_section.append(f" {custom_nodes_call}") + execution_section.extend( + self.build_function_body( + plan.special_functions_code, "pass", indentation=" " + ).splitlines() + ) + execution_section.append(f" for q in range({plan.queue_size}):") + execution_section.extend( + self.build_function_body( + plan.loop_code, "pass", indentation=" " + ).splitlines() + ) + + entrypoint_section = [ + "# Entrypoint", + 'if __name__ == "__main__":', + " main()", + ] + + final_code = "\n".join( + static_imports + + [""] + + workflow_section + + [""] + + execution_section + + [""] + + entrypoint_section + ) + return black.format_str(final_code, mode=black.Mode()) + + @staticmethod + def format_python_literal(value: Any) -> str: + return pformat(value, sort_dicts=False) + + @staticmethod + def build_function_body( + code_lines: list[str], empty_fallback: str, indentation: str = " " + ) -> str: + if not code_lines: + return f"{indentation}{empty_fallback}" + formatted_lines = [] + for line in code_lines: + stripped_line = line.lstrip() + if not stripped_line.endswith("\n"): + stripped_line += "\n" + formatted_lines.append(f"{indentation}{stripped_line}") + return "".join(formatted_lines).rstrip() diff --git a/comfyui_to_python/io.py b/comfyui_to_python/io.py new file mode 100644 index 0000000..cbe27c2 --- /dev/null +++ b/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) diff --git a/comfyui_to_python/load_order.py b/comfyui_to_python/load_order.py new file mode 100644 index 0000000..59a75f4 --- /dev/null +++ b/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) diff --git a/comfyui_to_python/node_runtime.py b/comfyui_to_python/node_runtime.py new file mode 100644 index 0000000..894e317 --- /dev/null +++ b/comfyui_to_python/node_runtime.py @@ -0,0 +1,133 @@ +import os +import sys +from typing import Any, Mapping, Sequence, Union + + +def find_path(name: str, path: str = None) -> str: + """Recursively search parent folders until the named entry is found.""" + if path is None: + path = os.getcwd() + + if name in os.listdir(path): + path_name = os.path.join(path, name) + print(f"{name} found: {path_name}") + return path_name + + parent_directory = os.path.dirname(path) + if parent_directory == path: + return None + + return find_path(name, parent_directory) + + +def get_comfyui_path() -> str: + """Return the configured ComfyUI path, preferring COMFYUI_PATH when set.""" + comfyui_path = os.environ.get("COMFYUI_PATH") + if comfyui_path: + return comfyui_path + return find_path("ComfyUI") + + +def add_comfyui_directory_to_sys_path() -> None: + """Add the ComfyUI checkout to sys.path.""" + comfyui_path = get_comfyui_path() + if comfyui_path is not None and os.path.isdir(comfyui_path): + if comfyui_path in sys.path: + sys.path.remove(comfyui_path) + sys.path.insert(0, comfyui_path) + print(f"'{comfyui_path}' added to sys.path") + + +def add_extra_model_paths() -> None: + """Load ComfyUI extra model paths configuration when available.""" + try: + from main import load_extra_path_config + except ImportError: + print( + "Could not import load_extra_path_config from main.py. Looking in utils.extra_config instead." + ) + from utils.extra_config import load_extra_path_config + + extra_model_paths = find_path("extra_model_paths.yaml") + if extra_model_paths is not None: + load_extra_path_config(extra_model_paths) + else: + print("Could not find the extra_model_paths config file.") + + +def bootstrap_comfyui_runtime() -> None: + """Mirror the allocator-related ComfyUI startup steps before torch import.""" + add_comfyui_directory_to_sys_path() + + import comfy.options + + comfy.options.enable_args_parsing() + + from comfy.cli_args import args + + if os.name == "nt": + os.environ["MIMALLOC_PURGE_DELAY"] = "0" + + if args.default_device is not None: + default_dev = args.default_device + devices = list(range(32)) + devices.remove(default_dev) + devices.insert(0, default_dev) + devices = ",".join(map(str, devices)) + os.environ["CUDA_VISIBLE_DEVICES"] = str(devices) + os.environ["HIP_VISIBLE_DEVICES"] = str(devices) + + if args.cuda_device is not None: + os.environ["CUDA_VISIBLE_DEVICES"] = str(args.cuda_device) + os.environ["HIP_VISIBLE_DEVICES"] = str(args.cuda_device) + os.environ["ASCEND_RT_VISIBLE_DEVICES"] = str(args.cuda_device) + + if args.oneapi_device_selector is not None: + os.environ["ONEAPI_DEVICE_SELECTOR"] = args.oneapi_device_selector + + if args.deterministic and "CUBLAS_WORKSPACE_CONFIG" not in os.environ: + os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8" + + import cuda_malloc + + if "rocm" in cuda_malloc.get_torch_version_noimport(): + os.environ["OCL_SET_SVM_SIZE"] = "262144" + + +def import_custom_nodes() -> None: + """Initialize ComfyUI custom nodes in the exporter runtime.""" + comfyui_path = get_comfyui_path() + if comfyui_path and comfyui_path not in sys.path: + sys.path.insert(0, comfyui_path) + + import asyncio + import execution + from nodes import init_extra_nodes + + if comfyui_path in sys.path: + sys.path.remove(comfyui_path) + sys.path.insert(0, comfyui_path) + + import server + + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + server_instance = server.PromptServer(loop) + execution.PromptQueue(server_instance) + asyncio.run(init_extra_nodes()) + + +def get_node_class_mappings() -> dict: + """Load ComfyUI node mappings on demand.""" + add_comfyui_directory_to_sys_path() + from nodes import NODE_CLASS_MAPPINGS + + return NODE_CLASS_MAPPINGS + + +def get_value_at_index(obj: Union[Sequence, Mapping], index: int) -> Any: + """Return a sequence or mapping result item by index.""" + try: + return obj[index] + except KeyError: + return obj["result"][index] diff --git a/comfyui_to_python/workflow_loader.py b/comfyui_to_python/workflow_loader.py new file mode 100644 index 0000000..a937b9b --- /dev/null +++ b/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 diff --git a/comfyui_to_python_utils.py b/comfyui_to_python_utils.py index e0bfa46..529f173 100644 --- a/comfyui_to_python_utils.py +++ b/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", +]