diff --git a/comfyui_to_python.py b/comfyui_to_python.py index b5cdc42..9ea6487 100644 --- a/comfyui_to_python.py +++ b/comfyui_to_python.py @@ -172,7 +172,7 @@ class CodeGenerator: return False - def generate_workflow(self, load_order: List, filename: str = 'generated_code_workflow.py', queue_size: int = 10) -> str: + def generate_workflow(self, load_order: List, queue_size: int = 10) -> str: """Generate the execution code based on the load order. Args: @@ -185,16 +185,25 @@ class CodeGenerator: 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 = set(['NODE_CLASS_MAPPINGS']), {}, [], [] + import_statements, executed_variables, arg_inputs, special_functions_code, code = set(['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'] + + missing = [] + for i, input in enumerate(self.node_class_mappings[class_type].INPUT_TYPES().get("required", {}).keys()): + if input not in inputs: + input_var = f"{input}{len(arg_inputs)+1}" + arg_inputs.append((input_var, f"Argument {i}, input `{input}` for node \\\"{data['_meta'].get('title', class_type)}\\\" id {idx}")) + print("WARNING: Missing required input", input, "for", class_type) + print("That will be CLI arg " + str(len(arg_inputs))) + missing.append((input_var, len(arg_inputs))) + class_def = self.node_class_mappings[class_type]() # If the class hasn't been initialized yet, initialize it and generate the import statements @@ -216,6 +225,8 @@ class CodeGenerator: # Remove any keyword arguments from **inputs if they are not in class_def_params inputs = {key: value for key, value in inputs.items() if key in class_def_params} + for input, arg in missing: + inputs[input] = {"variable_name": f"argv." + input} # Deal with hidden variables if 'unique_id' in class_def_params: inputs['unique_id'] = random.randint(1, 2**64) @@ -230,7 +241,7 @@ class CodeGenerator: 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, queue_size, custom_nodes) + final_code = self.assemble_python_code(import_statements, special_functions_code, arg_inputs, code, queue_size, custom_nodes) return final_code @@ -272,13 +283,12 @@ class CodeGenerator: if key == 'noise_seed' or key == 'seed': return f'{key}=random.randint(1, 2**64)' elif isinstance(value, str): - value = value.replace("\n", "\\n").replace('"', "'") - return f'{key}="{value}"' + return f'{key}={repr(value)}' elif isinstance(value, dict) and 'variable_name' in value: return f'{key}={value["variable_name"]}' return f'{key}={value}' - def assemble_python_code(self, import_statements: set, speical_functions_code: List[str], code: List[str], queue_size: int, custom_nodes=False) -> str: + def assemble_python_code(self, import_statements: set, speical_functions_code: List[str], arg_inputs: List[Tuple[str, str]], code: List[str], queue_size: int, custom_nodes=False) -> str: """Generates the final code string. Args: @@ -295,22 +305,31 @@ class CodeGenerator: func_strings = [] for func in [get_value_at_index, find_path, add_comfyui_directory_to_sys_path, add_extra_model_paths, import_custom_nodes]: func_strings.append(f'\n{inspect.getsource(func)}') + + argparse_code = "" + if arg_inputs: + argparse_code = f'argv = sys.argv[:]\n\nparser = argparse.ArgumentParser(description="A converted ComfyUI workflow. Required inputs listed below. Values passed should be in JSON")\n' + for i, (input_name, arg_desc) in enumerate(arg_inputs): + argparse_code += f'parser.add_argument("{input_name}", help="{arg_desc} (autogenerated)")\n' + argparse_code += 'args = parser.parse_args()\nsys.argv = [sys.argv[0]]\n' + # Define static import statements required for the script - static_imports = ['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'] + static_imports = ['import os', 'import random', 'import sys', 'import argparse', 'from typing import Sequence, Mapping, Any, Union', + 'import torch'] + func_strings + [argparse_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 = '' + static_imports += ['\n\nadd_comfyui_directory_to_sys_path()\nadd_extra_model_paths()\n'] # Create import statements for node classes - imports_code = [f"from nodes import {', '.join([class_name for class_name in import_statements])}" ] + imports_code = [f"from nodes import {', '.join([class_name for class_name in import_statements])}", ''] # Assemble the main function code, including custom nodes if applicable main_function_code = "def main():\n\t" + f'{custom_nodes}with torch.inference_mode():\n\t\t' + '\n\t\t'.join(speical_functions_code) \ + f'\n\n\t\tfor q in range({queue_size}):\n\t\t' + '\n\t\t'.join(code) # Concatenate all parts to form the final code - final_code = '\n'.join(static_imports + imports_code + ['', main_function_code, '', 'if __name__ == "__main__":', '\tmain()']) + 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()) @@ -453,7 +472,7 @@ class ComfyUItoPython: # 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, filename=self.output_file, queue_size=self.queue_size) + generated_code = code_generator.generate_workflow(load_order, queue_size=self.queue_size) # Step 5: Write the generated code to a file FileHandler.write_code_to_file(self.output_file, generated_code)