Fixed Windows paths, missing inputs used as args

Author: atmaranto

comfyui_to_python.py

@@ -172,7 +172,7 @@ def can_be_imported(self, import_name: str):
 
         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 @@ def generate_workflow(self, load_order: List, filename: str = 'generated_code_wo
             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 @@ def generate_workflow(self, load_order: List, filename: str = 'generated_code_wo
 
             # 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 @@ def generate_workflow(self, load_order: List, filename: str = 'generated_code_wo
                 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 @@ def format_arg(self, key: str, value: any) -> str:
         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 @@ def assemble_python_code(self, import_statements: set, speical_functions_code: L
         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 @@ def execute(self):
 
         # 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)