nodes/pydn-ComfyUI-to-Python-Extension
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Module API and --output

Browse Source
This commit is contained in:
atmaranto
2024-03-19 01:27:41 -05:00
parent fc76d349eb
commit e42b98048f
3 changed files with 138 additions and 37 deletions
Download Patch File Download Diff File Expand all files Collapse all files
comfyui_to_python.py
+82 -31
View File
@@ -183,7 +183,7 @@ class CodeGenerator:
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 @@ class CodeGenerator:
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 @@ class CodeGenerator:
# 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 @@ class CodeGenerator:
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 @@ class CodeGenerator:
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 @@ class CodeGenerator:
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,7 +488,7 @@ 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.
@@ -445,7 +496,7 @@ class ComfyUItoPython:
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.
"""