Skip to content

Commit a0f665c

Browse files
committed
Allow argument to determine save location
1 parent 8c7cbf6 commit a0f665c

File tree

2 files changed

+99
-4
lines changed

2 files changed

+99
-4
lines changed

comfyui_to_python.py

Lines changed: 23 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,18 @@
1111
import black
1212

1313

14-
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
14+
from comfyui_to_python_utils import import_custom_nodes, find_path, add_comfyui_directory_to_sys_path, add_extra_model_paths,\
15+
get_value_at_index, parse_arg, save_image_wrapper
16+
17+
PACKAGED_FUNCTIONS = [
18+
get_value_at_index,
19+
find_path,
20+
add_comfyui_directory_to_sys_path,
21+
add_extra_model_paths,
22+
import_custom_nodes,
23+
save_image_wrapper,
24+
parse_arg
25+
]
1526

1627
add_comfyui_directory_to_sys_path()
1728
from nodes import NODE_CLASS_MAPPINGS
@@ -304,14 +315,16 @@ def assemble_python_code(self, import_statements: set, speical_functions_code: L
304315
"""
305316
# Get the source code of the utils functions as a string
306317
func_strings = []
307-
for func in [get_value_at_index, find_path, add_comfyui_directory_to_sys_path, add_extra_model_paths, import_custom_nodes, parse_arg]:
318+
for func in PACKAGED_FUNCTIONS:
308319
func_strings.append(f'\n{inspect.getsource(func)}')
309320

310321
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'
311322
for i, (input_name, arg_desc) in enumerate(arg_inputs):
312323
argparse_code += f'parser.add_argument("{input_name}", help="{arg_desc} (autogenerated)")\n'
313324
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'
314325
argparse_code += f'parser.add_argument("--comfyui-directory", "-c", default=None, help="Where to look for ComfyUI (default: current directory)")\n'
326+
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'
327+
argparse_code += f'parser.add_argument("--disable-metadata", action="store_true", help="Disables writing workflow metadata to the outputs")\n'
315328
argparse_code += 'args = parser.parse_args()\nsys.argv = [sys.argv[0]]\n'
316329

317330
# Define static import statements required for the script
@@ -347,10 +360,16 @@ def get_class_info(self, class_type: str) -> Tuple[str, str, str]:
347360
"""
348361
import_statement = class_type
349362
variable_name = self.clean_variable_name(class_type)
363+
before = ""
364+
after = ""
365+
if class_type.strip() == 'SaveImage':
366+
before = 'save_image_wrapper('
367+
after = ')'
368+
350369
if self.can_be_imported(class_type):
351-
class_code = f'{variable_name} = {class_type.strip()}()'
370+
class_code = f'{variable_name} = {before}{class_type.strip()}{after}()'
352371
else:
353-
class_code = f'{variable_name} = NODE_CLASS_MAPPINGS["{class_type}"]()'
372+
class_code = f'{variable_name} = {before}NODE_CLASS_MAPPINGS["{class_type}"]{after}()'
354373

355374
return class_type, import_statement, class_code
356375

comfyui_to_python_utils.py

Lines changed: 76 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -112,3 +112,79 @@ def parse_arg(s: str):
112112
return json.loads(s)
113113
except json.JSONDecodeError:
114114
return s
115+
116+
def save_image_wrapper(cls):
117+
if args.output is None: return cls
118+
119+
from PIL import Image, ImageOps, ImageSequence
120+
from PIL.PngImagePlugin import PngInfo
121+
122+
import numpy as np
123+
124+
class WrappedSaveImage(cls):
125+
counter = 0
126+
127+
def save_images(self, images, filename_prefix="ComfyUI", prompt=None, extra_pnginfo=None):
128+
if args.output is None:
129+
return super().save_images(images, filename_prefix, prompt, extra_pnginfo)
130+
else:
131+
if len(images) > 1 and args.output == "-":
132+
raise ValueError("Cannot save multiple images to stdout")
133+
filename_prefix += self.prefix_append
134+
135+
results = list()
136+
for (batch_number, image) in enumerate(images):
137+
i = 255. * image.cpu().numpy()
138+
img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
139+
metadata = None
140+
if not args.disable_metadata:
141+
metadata = PngInfo()
142+
if prompt is not None:
143+
metadata.add_text("prompt", json.dumps(prompt))
144+
if extra_pnginfo is not None:
145+
for x in extra_pnginfo:
146+
metadata.add_text(x, json.dumps(extra_pnginfo[x]))
147+
148+
if args.output == "-":
149+
img.save(sys.stdout.buffer, format="png", pnginfo=metadata, compress_level=self.compress_level)
150+
else:
151+
subfolder = ""
152+
if len(images) == 1:
153+
if os.path.isdir(args.output):
154+
subfolder = args.output
155+
file = "output.png"
156+
else:
157+
subfolder, file = os.path.split(args.output)
158+
if subfolder == "":
159+
subfolder = os.getcwd()
160+
else:
161+
if os.path.isdir(args.output):
162+
subfolder = args.output
163+
file = "output"
164+
elif os.pathsep in args.output:
165+
subfolder, file = os.path.split(args.output)
166+
else:
167+
file = args.output
168+
if subfolder == "":
169+
subfolder = os.getcwd()
170+
171+
files = os.listdir(subfolder)
172+
file_pattern = file
173+
while True:
174+
filename_with_batch_num = file_pattern.replace("%batch_num%", str(batch_number))
175+
file = f"{filename_with_batch_num}_{self.counter:05}_.png"
176+
self.counter += 1
177+
178+
if file not in files:
179+
break
180+
181+
img.save(os.path.join(subfolder, file), pnginfo=metadata, compress_level=self.compress_level)
182+
results.append({
183+
"filename": file,
184+
"subfolder": subfolder,
185+
"type": self.type
186+
})
187+
188+
return {"ui": {"images": results}}
189+
190+
return WrappedSaveImage

0 commit comments

Comments
 (0)