Module API and --output

Author: atmaranto

README.md

@@ -6,6 +6,31 @@ I plan to do some fresh install testing when I have time.
 
 This is a fork maintained by Anthony Maranto of the original [ComfyUI-To-Python-Extension](https://github.com/pydn/ComfyUI-to-Python-Extension) by Peyton DeNiro. It provides a more robust command-line interface and the ability to export your current workflow as a script directly from the ComfyUI web interface.
 
+Once exported, this script can be run to run the workflow without a frontend, or it can be imported and the `main()` function can be used to call the workflow programatically.
+
+### New Feauture: Module Support
+
+Now, scripts exported with SaS can be imported as modules! Once you have a script exported, you can use it like:
+```python
+>>> import exported_script
+>>> results = exported_script.main("A prompt that would be sent to the command-line arguments", queue_size=1)
+```
+
+The first `SaveImage` node reached will instead *return* the output to the calling function. 
+
+### New Feature: Custom Output Path
+
+When running the exported script normally, you can now specify an `--output` option that will override the default path of `SaveImage` nodes.
+If only a single image is exported by the node, then the path will be used verbatim. Otherwise, the path will be used as a prefix, and `_#####.png` will be appended
+to ensure uniqueness. Note that files *will be clobbered* if only one image is exported.
+If the path is a directory, the `SaveImage` node's `filename_prefix` will be used.
+
+If `-` is selected as the output path, normal ComfyUI output will be piped to stderr and the resultant image will be piped to stdout, allowing one to use the script
+like:
+```bash
+python3 script.py "A painting of outer space" --output - --queue-size 1 > image.png
+```
+
 ### Usage (Web)
 
 Upon installation, there will be a button labeled "Save as Script" on the interface, pictured below:
@@ -58,19 +83,34 @@ positional arguments:
 options:
   -h, --help            show this help message and exit
   --queue-size QUEUE_SIZE, -q QUEUE_SIZE
-                        How many times the workflow will be executed (default: 10)
+                        How many times the workflow will be executed (default: 1)
   --comfyui-directory COMFYUI_DIRECTORY, -c COMFYUI_DIRECTORY
                         Where to look for ComfyUI (default: current directory)
+  --output OUTPUT, -o OUTPUT
+                        The location to save the output image. Either a file path, a directory, or - for stdout (default: the ComfyUI output directory)
+  --disable-metadata    Disables writing workflow metadata to the outputs
 ```
 
 Arguments are new. **If you have any suggestions on how to improve them or on how to effectively specify defaults in the workflow and override in the command-line**, feel free to suggest that in an Issue.
 
+#### Passing Arguments to ComfyUI
+
+In case you want to pass anything to the ComfyUI server as an argument, you can use `--` to indicate you're done with SaS arguments and are now passing ComfyUI arguments.
+For instance:
+
+```bash
+python3 script.py "A painting of outer space" --queue-size 1 -- --cpu
+```
+
 ### Other Changes
 
 #### Bugfixes
 - Windows paths are now properly escaped.
 - I also fixed what seemed to be a minor bug with exporting certain Crystools nodes, possibly due to their unusual name.
 
+#### TODO
+- Improve compatibility with module API
+
 ## Old Description of ComfyUI-to-Python-Extension (usage altered)
 
 The `ComfyUI-to-Python-Extension` is a powerful tool that translates [ComfyUI](https://github.com/comfyanonymous/ComfyUI) workflows into executable Python code. Designed to bridge the gap between ComfyUI's visual interface and Python's programming environment, this script facilitates the seamless transition from design to code execution. Whether you're a data scientist, a software developer, or an AI enthusiast, this tool streamlines the process of implementing ComfyUI workflows in Python.

comfyui_to_python.py

@@ -183,7 +183,7 @@ def can_be_imported(self, import_name: str):
 
         return False
 
-    def generate_workflow(self, load_order: List, queue_size: int = 10) -> str:
+    def generate_workflow(self, load_order: List, queue_size: int = 1) -> str:
         """Generate the execution code based on the load order.
 
         Args:
@@ -246,10 +246,19 @@ def generate_workflow(self, load_order: List, queue_size: int = 10) -> str:
             executed_variables[idx] = f'{self.clean_variable_name(class_type)}_{idx}'
             inputs = self.update_inputs(inputs, executed_variables)
 
-            if is_special_function:
-                special_functions_code.append(self.create_function_call_code(initialized_objects[class_type], class_def.FUNCTION, executed_variables[idx], is_special_function, **inputs))
+            if class_type == 'SaveImage':
+                save_code = self.create_function_call_code(initialized_objects[class_type], class_def.FUNCTION, executed_variables[idx], is_special_function, **inputs).strip()
+                return_code = ['if __name__ != "__main__":', '\treturn dict(' + ', '.join(self.format_arg(key, value) for key, value in inputs.items()) + ')', 'else:', '\t' + save_code]
+
+                if is_special_function:
+                    special_functions_code.extend(return_code)
+                else:
+                    code.extend(return_code) ### This should presumably NEVER occur for a valid workflow
             else:
-                code.append(self.create_function_call_code(initialized_objects[class_type], class_def.FUNCTION, executed_variables[idx], is_special_function, **inputs))
+                if is_special_function:
+                    special_functions_code.append(self.create_function_call_code(initialized_objects[class_type], class_def.FUNCTION, executed_variables[idx], is_special_function, **inputs))
+                else:
+                    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, arg_inputs, code, queue_size, custom_nodes)
@@ -274,11 +283,6 @@ def create_function_call_code(self, obj_name: str, func: str, variable_name: str
         # Generate the Python code
         code = f'{variable_name} = {obj_name}.{func}({args})\n'
 
-        # If the code contains dependencies and is not a loader or encoder, indent the code because it will be placed inside
-        # of a for loop
-        if not is_special_function:
-            code = f'\t{code}'
-
         return code
 
     def format_arg(self, key: str, value: any) -> str:
@@ -300,7 +304,7 @@ def format_arg(self, key: str, value: any) -> str:
             return f'{key}={value["variable_name"]}'
         return f'{key}={value}'
 
-    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:
+    def assemble_python_code(self, import_statements: set, special_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:
@@ -318,32 +322,79 @@ def assemble_python_code(self, import_statements: set, speical_functions_code: L
         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'
+        argparse_code = [f'parser = argparse.ArgumentParser(description="A converted ComfyUI workflow. Required inputs listed below. Values passed should be valid JSON (assumes string if not valid JSON).")']
         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 += 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'
+            argparse_code.append(f'parser.add_argument("{input_name}", help="{arg_desc} (autogenerated)")\n')
+        argparse_code.append(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.append('parser.add_argument("--comfyui-directory", "-c", default=None, help="Where to look for ComfyUI (default: current directory)")\n')
+        argparse_code.append(f'parser.add_argument("--output", "-o", default=None, help="The location to save the output image. Either a file path, a directory, or - for stdout (default: the ComfyUI output directory)")\n')
+        argparse_code.append(f'parser.add_argument("--disable-metadata", action="store_true", help="Disables writing workflow metadata to the outputs")\n')
+        argparse_code.append('''
+comfy_args = [sys.argv[0]]
+if "--" in sys.argv:
+    idx = sys.argv.index("--")
+    comfy_args += sys.argv[idx+1:]
+    sys.argv = sys.argv[:idx]
+
+args = None
+if __name__ == "__main__":
+    args = parser.parse_args()
+    sys.argv = comfy_args
+if args is not None and args.output is not None and args.output == "-":
+    ctx = contextlib.redirect_stdout(sys.stderr)
+else:
+    ctx = contextlib.nullcontext()
+''')
 
         # Define static import statements required for the script
-        static_imports = ['import os', 'import random', 'import sys', 'import json', 'import argparse', 'from typing import Sequence, Mapping, Any, Union', 
-                          'import torch'] + func_strings + [argparse_code]
+        static_imports = ['import os', 'import random', 'import sys', 'import json', 'import argparse', 'import contextlib', '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])}", '']
+        newline_doubletab = '\n\t\t' # You can't use backslashes in f-strings
+        newline_tripletab = '\n\t\t\t' # Same
         # 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(args.queue_size):\n\t\t' + '\n\t\t'.join(code)
+        main_function_code = f"""
+_custom_nodes_imported = {str(not custom_nodes)}
+_custom_path_added = False
+
+def main(*func_args, **func_kwargs):
+    global args, _custom_nodes_imported, _custom_path_added
+    if __name__ == "__main__":
+        if args is None:
+            args = parser.parse_args()
+    else:
+        defaults = dict((arg, parser.get_default(arg)) for arg in ['queue_size', 'comfyui_directory', 'output', 'disable_metadata'])
+        ordered_args = dict(zip({[input_name for input_name, _ in arg_inputs]}, func_args))
+
+        all_args = dict()
+        all_args.update(defaults)
+        all_args.update(ordered_args)
+        all_args.update(func_kwargs)
+
+        args = argparse.Namespace(**all_args)
+    
+    with ctx:
+        if not _custom_path_added:
+            add_comfyui_directory_to_sys_path()
+            add_extra_model_paths()
+
+            _custom_path_added = True
+        
+        if not _custom_nodes_imported:
+            import_custom_nodes()
+
+            _custom_nodes_imported = True
+        
+        from nodes import {', '.join([class_name for class_name in import_statements])}
+    
+    with torch.inference_mode(), ctx:
+        {newline_doubletab.join(special_functions_code)}
+        for q in range(args.queue_size):
+            {newline_tripletab.join(code)}""".replace("    ", "\t")
         # 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 + [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())
 
@@ -363,7 +414,7 @@ def get_class_info(self, class_type: str) -> Tuple[str, str, str]:
         before = ""
         after = ""
         if class_type.strip() == 'SaveImage':
-            before = 'save_image_wrapper('
+            before = 'save_image_wrapper(' + 'ctx, '
             after = ')'
 
         if self.can_be_imported(class_type):
@@ -437,15 +488,15 @@ class ComfyUItoPython:
         base_node_class_mappings (Dict): Base mappings of node classes.
     """
 
-    def __init__(self, workflow: str = "", input_file: str = "", output_file: (str | TextIO) = "", queue_size: int = 10, node_class_mappings: Dict = NODE_CLASS_MAPPINGS,
+    def __init__(self, workflow: str = "", input_file: str = "", output_file: (str | TextIO) = "", queue_size: int = 1, node_class_mappings: Dict = NODE_CLASS_MAPPINGS,
                  needs_init_custom_nodes: bool = False):
         """Initialize the ComfyUItoPython class with the given parameters. Exactly one of workflow or input_file must be specified.
 
         Args:
             workflow (str): The workflow's JSON.
             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 10.
+            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.
             needs_init_custom_nodes (bool): Whether to initialize custom nodes. Defaults to False.
         """

comfyui_to_python_utils.py

@@ -106,14 +106,17 @@ def get_value_at_index(obj: Union[Sequence, Mapping], index: int) -> Any:
     except KeyError:
         return obj['result'][index]
 
-def parse_arg(s: str):
+def parse_arg(s: Any):
     """ Parses a JSON string, returning it unchanged if the parsing fails. """
+    if __name__ == "__main__" or not isinstance(s, str):
+        return s
+    
     try:
         return json.loads(s)
     except json.JSONDecodeError:
         return s
 
-def save_image_wrapper(cls):
+def save_image_wrapper(context, cls):
     if args.output is None: return cls
 
     from PIL import Image, ImageOps, ImageSequence
@@ -146,7 +149,14 @@ def save_images(self, images, filename_prefix="ComfyUI", prompt=None, extra_pngi
                                 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)
+                        # Hack to briefly restore stdout
+                        if context is not None:
+                            context.__exit__(None, None, None) 
+                        try:
+                            img.save(sys.stdout.buffer, format="png", pnginfo=metadata, compress_level=self.compress_level)
+                        finally:
+                            if context is not None:
+                                context.__enter__()
                     else:
                         subfolder = ""
                         if len(images) == 1:
@@ -160,7 +170,7 @@ def save_images(self, images, filename_prefix="ComfyUI", prompt=None, extra_pngi
                         else:
                             if os.path.isdir(args.output):
                                 subfolder = args.output
-                                file = "output"
+                                file = filename_prefix
                             else:
                                 subfolder, file = os.path.split(args.output)
 
@@ -171,7 +181,7 @@ def save_images(self, images, filename_prefix="ComfyUI", prompt=None, extra_pngi
                             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"
+                                file = f"{filename_with_batch_num}_{self.counter:05}.png"
                                 self.counter += 1
 
                                 if file not in files: