From a0f665c5f6f362e642402c8dd0344f69e081f8f7 Mon Sep 17 00:00:00 2001 From: atmaranto <35581227+atmaranto@users.noreply.github.com> Date: Mon, 18 Mar 2024 16:12:13 -0500 Subject: [PATCH] Allow argument to determine save location --- comfyui_to_python.py | 27 ++++++++++++-- comfyui_to_python_utils.py | 76 ++++++++++++++++++++++++++++++++++++++ 2 files changed, 99 insertions(+), 4 deletions(-) diff --git a/comfyui_to_python.py b/comfyui_to_python.py index a0f9597..a28498c 100644 --- a/comfyui_to_python.py +++ b/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 diff --git a/comfyui_to_python_utils.py b/comfyui_to_python_utils.py index dc7bc62..d51d879 100644 --- a/comfyui_to_python_utils.py +++ b/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