From 557adae33e4ed1de0ecded92d4264c3bee143699 Mon Sep 17 00:00:00 2001 From: Peyton DeNiro <25550995+pydn@users.noreply.github.com> Date: Sat, 28 Mar 2026 13:08:34 -0500 Subject: [PATCH] Fix upscale export runtime compatibility --- comfyui_to_python.py | 71 ++++++++++++++++++++++++++++++++------------ 1 file changed, 52 insertions(+), 19 deletions(-) diff --git a/comfyui_to_python.py b/comfyui_to_python.py index 10e6b79..febb7be 100644 --- a/comfyui_to_python.py +++ b/comfyui_to_python.py @@ -16,20 +16,30 @@ 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, ) -add_comfyui_directory_to_sys_path() -from nodes import NODE_CLASS_MAPPINGS - 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. @@ -199,7 +209,7 @@ class CodeGenerator: """ # Create the necessary data structures to hold imports and generated code import_statements, executed_variables, special_functions_code, code = ( - set(["NODE_CLASS_MAPPINGS"]), + {"nodes": {"NODE_CLASS_MAPPINGS"}}, {}, [], [], @@ -235,8 +245,9 @@ class CodeGenerator: ) initialized_objects[class_type] = self.clean_variable_name(class_type) if class_type in self.base_node_class_mappings.keys(): - import_statements.add(import_statement) - if class_type not in self.base_node_class_mappings.keys(): + module_name, import_name = import_statement + import_statements.setdefault(module_name, set()).add(import_name) + if 'NODE_CLASS_MAPPINGS["' in class_code: custom_nodes = True special_functions_code.append(class_code) @@ -289,7 +300,11 @@ class CodeGenerator: # 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, queue_size, custom_nodes + import_statements, + special_functions_code, + code, + queue_size, + custom_nodes, ) return final_code @@ -347,7 +362,7 @@ class CodeGenerator: def assemble_python_code( self, - import_statements: set, + import_statements: Dict[str, set], speical_functions_code: List[str], code: List[str], queue_size: int, @@ -356,7 +371,7 @@ class CodeGenerator: """Generates the final code string. Args: - import_statements (set): A set of unique import statements. + import_statements (Dict[str, set]): Import statements grouped by module. speical_functions_code (List[str]): A list of special functions code strings. code (List[str]): A list of code strings. queue_size (int): Number of photos that will be generated by the script. @@ -369,6 +384,7 @@ class CodeGenerator: func_strings = [] for func in [ get_value_at_index, + get_comfyui_path, find_path, add_comfyui_directory_to_sys_path, add_extra_model_paths, @@ -393,9 +409,10 @@ class CodeGenerator: else: custom_nodes = "" # Create import statements for node classes - imports_code = [ - f"from nodes import {', '.join([class_name for class_name in import_statements])}" - ] + imports_code = [] + for module_name in sorted(import_statements.keys()): + class_names = ", ".join(sorted(import_statements[module_name])) + imports_code.append(f"from {module_name} import {class_names}") # Assemble the main function code, including custom nodes if applicable main_function_code = ( "def main():\n\t" @@ -415,20 +432,31 @@ class CodeGenerator: return final_code - def get_class_info(self, class_type: str) -> Tuple[str, str, str]: + def get_class_info(self, class_type: str) -> Tuple[str, Tuple[str, str], str]: """Generates and returns necessary information about class type. Args: class_type (str): Class type. Returns: - Tuple[str, str, str]: Updated class type, import statement string, class initialization code. + Tuple[str, Tuple[str, str], str]: Updated class type, import statement and initialization code. """ - import_statement = class_type + 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) - if class_type in self.base_node_class_mappings.keys(): + is_importable_module = bool( + module_name + and "/" not in module_name + and "\\" not in module_name + and all(part.isidentifier() for part in module_name.split(".")) + ) + if class_type in self.base_node_class_mappings.keys() and is_importable_module: + import_statement = (module_name, class_type) class_code = f"{variable_name} = {class_type.strip()}()" else: + import_statement = ("nodes", "NODE_CLASS_MAPPINGS") class_code = f'{variable_name} = NODE_CLASS_MAPPINGS["{class_type}"]()' return class_type, import_statement, class_code @@ -514,7 +542,7 @@ class ComfyUItoPython: input_file: str = "", output_file: str | TextIO = "", queue_size: int = 1, - node_class_mappings: Dict = NODE_CLASS_MAPPINGS, + 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. @@ -523,7 +551,8 @@ class ComfyUItoPython: 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): Mappings of node classes. Defaults to NODE_CLASS_MAPPINGS. + 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: @@ -538,7 +567,11 @@ class ComfyUItoPython: self.input_file = input_file self.output_file = output_file self.queue_size = queue_size - self.node_class_mappings = node_class_mappings + 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)