Allow argument to determine save location

comfyui_to_python.py

@@ -11,7 +11,18 @@ from typing import Dict, List, Any, Callable, Tuple, TextIO
 import black
 
 
-from comfyui_to_python_utils import import_custom_nodes, find_path, add_comfyui_directory_to_sys_path, add_extra_model_paths, get_value_at_index, parse_arg
+from comfyui_to_python_utils import import_custom_nodes, find_path, add_comfyui_directory_to_sys_path, add_extra_model_paths,\
+                                    get_value_at_index, parse_arg, save_image_wrapper
+
+PACKAGED_FUNCTIONS = [
+    get_value_at_index,
+    find_path,
+    add_comfyui_directory_to_sys_path,
+    add_extra_model_paths,
+    import_custom_nodes,
+    save_image_wrapper,
+    parse_arg
+]
 
 add_comfyui_directory_to_sys_path()
 from nodes import NODE_CLASS_MAPPINGS
@@ -304,7 +315,7 @@ class CodeGenerator:
         """
         # Get the source code of the utils functions as a string
         func_strings = []
-        for func in [get_value_at_index, find_path, add_comfyui_directory_to_sys_path, add_extra_model_paths, import_custom_nodes, parse_arg]:
+        for func in PACKAGED_FUNCTIONS:
             func_strings.append(f'\n{inspect.getsource(func)}')
         
         argparse_code = f'\nparser = argparse.ArgumentParser(description="A converted ComfyUI workflow. Required inputs listed below. Values passed should be valid JSON (assumes string if not valid JSON).")\n'
@@ -312,6 +323,8 @@ class CodeGenerator:
             argparse_code += f'parser.add_argument("{input_name}", help="{arg_desc} (autogenerated)")\n'
         argparse_code += f'parser.add_argument("--queue-size", "-q", type=int, default={queue_size}, help="How many times the workflow will be executed (default: {queue_size})")\n'
         argparse_code += f'parser.add_argument("--comfyui-directory", "-c", default=None, help="Where to look for ComfyUI (default: current directory)")\n'
+        argparse_code += f'parser.add_argument("--output", "-o", default=None, help="The location to save the output image -- a file path, a directory, or - for stdout (default: the ComfyUI output directory)")\n'
+        argparse_code += f'parser.add_argument("--disable-metadata", action="store_true", help="Disables writing workflow metadata to the outputs")\n'
         argparse_code += 'args = parser.parse_args()\nsys.argv = [sys.argv[0]]\n'
         
         # Define static import statements required for the script
@@ -347,10 +360,16 @@ class CodeGenerator:
         """
         import_statement = class_type
         variable_name = self.clean_variable_name(class_type)
+        before = ""
+        after = ""
+        if class_type.strip() == 'SaveImage':
+            before = 'save_image_wrapper('
+            after = ')'
+        
         if self.can_be_imported(class_type):
-            class_code = f'{variable_name} = {class_type.strip()}()'
+            class_code = f'{variable_name} = {before}{class_type.strip()}{after}()'
         else:
-            class_code = f'{variable_name} = NODE_CLASS_MAPPINGS["{class_type}"]()'
+            class_code = f'{variable_name} = {before}NODE_CLASS_MAPPINGS["{class_type}"]{after}()'
 
         return class_type, import_statement, class_code