nodes/atmaranto-ComfyUI-SaveAsScript
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Fixed Windows paths, missing inputs used as args

Browse Source
This commit is contained in:
atmaranto
2024-03-09 04:59:32 -06:00
parent f0bc65ed1b
commit 793854d21b
1 changed files with 31 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files
comfyui_to_python.py
+31 -12
View File
@@ -172,7 +172,7 @@ class CodeGenerator:
return False
def generate_workflow(self, load_order: List, filename: str = 'generated_code_workflow.py', queue_size: int = 10) -> str:
def generate_workflow(self, load_order: List, queue_size: int = 10) -> str:
"""Generate the execution code based on the load order.
Args:
@@ -185,16 +185,25 @@ class CodeGenerator:
str: Generated execution code as a string.
"""
# Create the necessary data structures to hold imports and generated code
import_statements, executed_variables, special_functions_code, code = set(['NODE_CLASS_MAPPINGS']), {}, [], []
import_statements, executed_variables, arg_inputs, special_functions_code, code = set(['NODE_CLASS_MAPPINGS']), {}, [], [], []
# This dictionary will store the names of the objects that we have already initialized
initialized_objects = {}
custom_nodes = False
# Loop over each dictionary in the load order list
for idx, data, is_special_function in load_order:
# Generate class definition and inputs from the data
inputs, class_type = data['inputs'], data['class_type']
missing = []
for i, input in enumerate(self.node_class_mappings[class_type].INPUT_TYPES().get("required", {}).keys()):
if input not in inputs:
input_var = f"{input}{len(arg_inputs)+1}"
arg_inputs.append((input_var, f"Argument {i}, input `{input}` for node \\\"{data['_meta'].get('title', class_type)}\\\" id {idx}"))
print("WARNING: Missing required input", input, "for", class_type)
print("That will be CLI arg " + str(len(arg_inputs)))
missing.append((input_var, len(arg_inputs)))
class_def = self.node_class_mappings[class_type]()
# If the class hasn't been initialized yet, initialize it and generate the import statements
@@ -216,6 +225,8 @@ class CodeGenerator:
# Remove any keyword arguments from **inputs if they are not in class_def_params
inputs = {key: value for key, value in inputs.items() if key in class_def_params}
for input, arg in missing:
inputs[input] = {"variable_name": f"argv." + input}
# Deal with hidden variables
if 'unique_id' in class_def_params:
inputs['unique_id'] = random.randint(1, 2**64)
@@ -230,7 +241,7 @@ class CodeGenerator:
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, code, queue_size, custom_nodes)
final_code = self.assemble_python_code(import_statements, special_functions_code, arg_inputs, code, queue_size, custom_nodes)
return final_code
@@ -272,13 +283,12 @@ class CodeGenerator:
if key == 'noise_seed' or key == 'seed':
return f'{key}=random.randint(1, 2**64)'
elif isinstance(value, str):
value = value.replace("\n", "\\n").replace('"', "'")
return f'{key}="{value}"'
return f'{key}={repr(value)}'
elif isinstance(value, dict) and 'variable_name' in value:
return f'{key}={value["variable_name"]}'
return f'{key}={value}'
def assemble_python_code(self, import_statements: set, speical_functions_code: List[str], code: List[str], queue_size: int, custom_nodes=False) -> str:
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:
"""Generates the final code string.
Args:
@@ -295,22 +305,31 @@ class CodeGenerator:
func_strings = []
for func in [get_value_at_index, find_path, add_comfyui_directory_to_sys_path, add_extra_model_paths, import_custom_nodes]:
func_strings.append(f'\n{inspect.getsource(func)}')
argparse_code = ""
if arg_inputs:
argparse_code = f'argv = sys.argv[:]\n\nparser = argparse.ArgumentParser(description="A converted ComfyUI workflow. Required inputs listed below. Values passed should be in JSON")\n'
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 += 'args = parser.parse_args()\nsys.argv = [sys.argv[0]]\n'
# Define static import statements required for the script
static_imports = ['import os', 'import random', 'import sys', 'from typing import Sequence, Mapping, Any, Union',
'import torch'] + func_strings + ['\n\nadd_comfyui_directory_to_sys_path()\nadd_extra_model_paths()\n']
static_imports = ['import os', 'import random', 'import sys', 'import argparse', '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])}" ]
imports_code = [f"from nodes import {', '.join([class_name for class_name in import_statements])}", '']
# 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({queue_size}):\n\t\t' + '\n\t\t'.join(code)
# 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 + imports_code + [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())
@@ -453,7 +472,7 @@ class ComfyUItoPython:
# Step 4: Generate the workflow code
code_generator = CodeGenerator(self.node_class_mappings, self.base_node_class_mappings)
generated_code = code_generator.generate_workflow(load_order, filename=self.output_file, queue_size=self.queue_size)
generated_code = code_generator.generate_workflow(load_order, queue_size=self.queue_size)
# Step 5: Write the generated code to a file
FileHandler.write_code_to_file(self.output_file, generated_code)