|
11 | 11 | import black
|
12 | 12 |
|
13 | 13 |
|
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 | +] |
15 | 26 |
|
16 | 27 | add_comfyui_directory_to_sys_path()
|
17 | 28 | from nodes import NODE_CLASS_MAPPINGS
|
@@ -304,14 +315,16 @@ def assemble_python_code(self, import_statements: set, speical_functions_code: L
|
304 | 315 | """
|
305 | 316 | # Get the source code of the utils functions as a string
|
306 | 317 | 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: |
308 | 319 | func_strings.append(f'\n{inspect.getsource(func)}')
|
309 | 320 |
|
310 | 321 | 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'
|
311 | 322 | for i, (input_name, arg_desc) in enumerate(arg_inputs):
|
312 | 323 | argparse_code += f'parser.add_argument("{input_name}", help="{arg_desc} (autogenerated)")\n'
|
313 | 324 | 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'
|
314 | 325 | 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' |
315 | 328 | argparse_code += 'args = parser.parse_args()\nsys.argv = [sys.argv[0]]\n'
|
316 | 329 |
|
317 | 330 | # Define static import statements required for the script
|
@@ -347,10 +360,16 @@ def get_class_info(self, class_type: str) -> Tuple[str, str, str]:
|
347 | 360 | """
|
348 | 361 | import_statement = class_type
|
349 | 362 | 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 | + |
350 | 369 | 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}()' |
352 | 371 | 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}()' |
354 | 373 |
|
355 | 374 | return class_type, import_statement, class_code
|
356 | 375 |
|
|
0 commit comments