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
     

comfyui_to_python_utils.py

@@ -112,3 +112,79 @@ def parse_arg(s: str):
         return json.loads(s)
     except json.JSONDecodeError:
         return s
+
+def save_image_wrapper(cls):
+    if args.output is None: return cls
+    
+    from PIL import Image, ImageOps, ImageSequence
+    from PIL.PngImagePlugin import PngInfo
+
+    import numpy as np
+
+    class WrappedSaveImage(cls):
+        counter = 0
+
+        def save_images(self, images, filename_prefix="ComfyUI", prompt=None, extra_pnginfo=None):
+            if args.output is None:
+                return super().save_images(images, filename_prefix, prompt, extra_pnginfo)
+            else:
+                if len(images) > 1 and args.output == "-":
+                    raise ValueError("Cannot save multiple images to stdout")
+                filename_prefix += self.prefix_append
+                
+                results = list()
+                for (batch_number, image) in enumerate(images):
+                    i = 255. * image.cpu().numpy()
+                    img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
+                    metadata = None
+                    if not args.disable_metadata:
+                        metadata = PngInfo()
+                        if prompt is not None:
+                            metadata.add_text("prompt", json.dumps(prompt))
+                        if extra_pnginfo is not None:
+                            for x in extra_pnginfo:
+                                metadata.add_text(x, json.dumps(extra_pnginfo[x]))
+                    
+                    if args.output == "-":
+                        img.save(sys.stdout.buffer, format="png", pnginfo=metadata, compress_level=self.compress_level)
+                    else:
+                        subfolder = ""
+                        if len(images) == 1:
+                            if os.path.isdir(args.output):
+                                subfolder = args.output
+                                file = "output.png"
+                            else:
+                                subfolder, file = os.path.split(args.output)
+                                if subfolder == "":
+                                    subfolder = os.getcwd()
+                        else:
+                            if os.path.isdir(args.output):
+                                subfolder = args.output
+                                file = "output"
+                            elif os.pathsep in args.output:
+                                subfolder, file = os.path.split(args.output)
+                            else:
+                                file = args.output
+                            if subfolder == "":
+                                subfolder = os.getcwd()
+                            
+                            files = os.listdir(subfolder)
+                            file_pattern = file
+                            while True:
+                                filename_with_batch_num = file_pattern.replace("%batch_num%", str(batch_number))
+                                file = f"{filename_with_batch_num}_{self.counter:05}_.png"
+                                self.counter += 1
+
+                                if file not in files:
+                                    break
+                        
+                        img.save(os.path.join(subfolder, file), pnginfo=metadata, compress_level=self.compress_level)
+                        results.append({
+                            "filename": file,
+                            "subfolder": subfolder,
+                            "type": self.type
+                        })
+
+                return {"ui": {"images": results}}
+    
+    return WrappedSaveImage