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

Refactored generate workflow function for better modularity and cleaned up doc strings and type hints.

Browse Source
This commit is contained in:
Peyton
2023-08-03 20:27:47 -05:00
parent 342742cd88
commit 4ea765d8f7
1 changed files with 136 additions and 67 deletions
Download Patch File Download Diff File Expand all files Collapse all files
generate_python_code.py
+136 -67
View File
@@ -1,6 +1,6 @@
import sys
import json
from typing import Dict, List, Union, Tuple, Any
from typing import Dict, List, Any
sys.path.append('../')
# from execution import PromptExecutor
@@ -27,57 +27,73 @@ def read_json_file(file_path: str) -> dict:
raise ValueError(f"Invalid JSON format in file: {file_path}")
def determine_load_order(data: Dict[str, Dict[str, Union[Dict[str, Union[int, List[Union[str, int]]]], str]]]) -> List[Tuple[str, Dict[str, Union[Dict[str, Union[int, List[Union[str, int]]]], str]]]]:
"""Determine the load order of each key in the provided dictionary.
def determine_load_order(data: Dict) -> List:
"""
Determine the load order of each key in the provided dictionary. This code will place the
nodes without node dependencies first, then ensure that any node whose result is used
in another node will be added to the list in the order it should be executed.
Args:
data (Dict[str, Dict[str, Union[Dict[str, Union[int, List[Union[str, int]]]], str]]]):
The dictionary for which to determine the load order.
data (Dict):
The dictionary for which to determine the load order.
Returns:
List[Tuple[str, Dict[str, Union[Dict[str, Union[int, List[Union[str, int]]]], str]]]]:
A list of tuples where each tuple contains a key and its corresponding dictionary, ordered by load order.
List:
A list of tuples where each tuple contains a key and its corresponding dictionary, ordered by load order.
"""
# Visited dictionary to keep track of visited nodes
# Create a dictionary to keep track of visited nodes.
visited = {}
# Create a list to store the load order.
load_order = []
def dfs(key: str):
"""Depth-First Search function.
def dfs(key: str) -> None:
"""
Depth-First Search function.
Args:
key (str): The key from which to start the DFS.
key (str): The key from which to start the DFS.
Returns:
None.
None
"""
# Mark the node as visited
# Mark the node as visited.
visited[key] = True
inputs = data[key]['inputs']
for input_key, val in inputs.items():
# Check if the input value is a list and not already visited
if isinstance(val, list) and val[0] not in visited:
dfs(val[0])
# Add the key to the load order list
load_order.append((key, data[key]))
# Loop over each key in the data
# Loop over each input key.
for input_key, val in inputs.items():
# If the value is a list and the first item in the list (which should be a key)
# has not been visited yet, then recursively apply dfs on the dependency.
if isinstance(val, list) and val[0] not in visited:
dfs(val[0])
# Add the key and its corresponding data to the load order list.
load_order.append((key, data[key]))
# Loop over each key in the data.
for key in data:
# If the key is not visited, perform a DFS from that key
# Check if the key has dependencies
if not any(isinstance(val, list) for val in data[key]['inputs'].values()):
# Add keys with no dependencies to load order and mark as visited
load_order.append((key, data[key]))
visited[key] = True
# Loop over each key in the data.
for key in data:
# If the key has not been visited, perform a DFS from that key.
if key not in visited:
dfs(key)
return load_order
def generate_code(obj: object, func: str, variable_name: str, **kwargs) -> str:
def create_function_call_code(obj_name: str, func: str, variable_name: str, **kwargs) -> str:
"""
This function generates Python code for a function call.
Args:
obj (object): The object whose method is to be called.
obj_name (str): The name of the initialized object.
func (str): The function to be called.
variable_name (str): The name of the variable that the function result should be assigned to.
**kwargs: The keyword arguments for the function.
@@ -85,28 +101,23 @@ def generate_code(obj: object, func: str, variable_name: str, **kwargs) -> str:
Returns:
str: The generated Python code.
"""
# Get the class name of the object
obj_name = obj.__class__.__name__
# Convert the function arguments into a string
# If the value is a string, it is surrounded by quotes
# If the value is a dictionary and has key 'variable_name', its value is used as the arg value
# For images argument and obj_name has 'SaveImage', '.detach' is appended at the end
args = ', '.join(
f'{key}="{value}"' if isinstance(value, str)
else f'{key}={value["variable_name"]}.detach()' if key == 'images' and "SaveImage" in obj_name and isinstance(value, dict) and 'variable_name' in value
else f'{key}={value["variable_name"]}.detach()' if key == 'images' and "saveimage" in obj_name and isinstance(value, dict) and 'variable_name' in value
else f'{key}={value["variable_name"]}' if isinstance(value, dict) and 'variable_name' in value
else f'{key}={value}' for key, value in kwargs.items()
)
# Generate the Python code
code = f'{variable_name} = {obj_name}().{func}({args})'
code = f'{variable_name} = {obj_name}.{func}({args})'
return code
def update_inputs(inputs: Dict[str, Any], executed_variables: Dict[str, str]) -> Dict[str, Any]:
"""
Update inputs based on the executed variables.
@@ -123,58 +134,116 @@ def update_inputs(inputs: Dict[str, Any], executed_variables: Dict[str, str]) ->
inputs[key] = {'variable_name': f"{executed_variables[inputs[key][0]]}[{inputs[key][1]}]"}
return inputs
def generate_execution_code(load_order: List[Dict[str, Any]], filename: str = 'generated_code_workflow.py') -> str:
def get_class_info(class_type: str) -> (str, str, str):
"""
Generates and returns necessary information about class type.
Args:
class_type (str): Class type
Returns:
class_type (str): Updated class type
import_statement (str): Import statement string
class_code (str): Class initialization code
"""
# If the class is 'VAEDecode', adjust the class name
if class_type == 'VAEDecode':
class_type = 'VAEDecodeTiled'
import_statement = class_type
class_code = f'{class_type.lower()} = {class_type}()'
return class_type, import_statement, class_code
def assemble_python_code(import_statements: set, code: List[str]) -> str:
"""
Generates final code string.
Args:
import_statements (set): A set of unique import statements
code (List[str]): A list of code strings
Returns:
final_code (str): Generated final code as a string
"""
imports_code = [f"from nodes import {class_name}" for class_name in import_statements]
main_function_code = f"def main():\n\t" + '\n\t'.join(code)
final_code = '\n'.join(['import sys\nsys.path.append("../")'] + imports_code + ['', main_function_code, '', 'if __name__ == "__main__":', '\tmain()'])
return final_code
def write_code_to_file(filename: str, code: str) -> None:
"""
Writes given code to a .py file. If the directory does not exist, it creates it.
Args:
filename (str): The name of the Python file to save the code to.
code (str): The code to save.
"""
# Extract directory from the filename
directory = os.path.dirname(filename)
# If the directory does not exist, create it
if directory and not os.path.exists(directory):
os.makedirs(directory)
# Save the code to a .py file
with open(filename, 'w') as file:
file.write(code)
def generate_workflow(load_order: List, filename: str = 'generated_code_workflow.py') -> str:
"""
Generate the execution code based on the load order.
This function will organize the necessary imports at the top, ensure the
code is formatted in a readable way, wrap it into a main function, and
save it as a .py file in the same directory.
Args:
load_order (List[Dict[str, Any]]): A list of dictionaries representing the
load order. Each dictionary should contain keys 'inputs', 'class_type', and
optionally other keys with corresponding values.
filename (str, optional): The name of the Python file to which the code should be saved.
Defaults to 'generated_code.py'.
load_order (List): A list of tuples representing the load order.
filename (str): The name of the Python file to which the code should be saved.
Defaults to 'generated_code_workflow.py'.
Returns:
str: Generated execution code as a string.
"""
# Collect the imports required and the code separately
import_statements = set()
executed_variables = {}
code = []
# Create the necessary data structures to hold imports and generated code
import_statements, executed_variables, code = set(), {}, []
# This dictionary will store the names of the objects that we have already initialized
initialized_objects = {}
# Loop over each dictionary in the load order list
for idx, data in load_order:
inputs = data['inputs']
class_type = data['class_type']
# Generate class definition and inputs from the data
inputs, class_type = data['inputs'], data['class_type']
class_def = NODE_CLASS_MAPPINGS[class_type]()
# Keep track of required imports
import_statements.add(class_type)
# If the class hasn't been initialized yet, initialize it and generate the import statements
if class_type not in initialized_objects:
class_type, import_statement, class_code = get_class_info(class_type)
initialized_objects[class_type] = class_type.lower()
import_statements.add(import_statement)
code.append(class_code)
# Create executed variable and generate code
executed_variables[idx] = f'{class_type.lower()}_{idx}'
inputs = update_inputs(inputs, executed_variables)
code.append(generate_code(class_def, class_def.FUNCTION, executed_variables[idx], **inputs))
# Convert the import statements into code
imports_code = [f"from nodes import {class_name}" for class_name in import_statements]
# Combine import statements and code, and wrap them in a main function
main_function_code = f"def main():\n\t" + '\n\t'.join(code)
full_code = '\n'.join(['import sys\nsys.path.append("../")'] + imports_code + ['', main_function_code, '', 'if __name__ == "__main__":', '\tmain()'])
code.append(create_function_call_code(initialized_objects[class_type], class_def.FUNCTION, executed_variables[idx], **inputs))
# Generate final code by combining imports and code, and wrap them in a main function
final_code = assemble_python_code(import_statements, code)
# Save the code to a .py file
with open(filename, 'w') as file:
file.write(full_code)
return full_code
write_code_to_file(filename, final_code)
return final_code
if __name__ == '__main__':
prompt = read_json_file('workflow_api.json')
load_order = determine_load_order(prompt)
code = generate_execution_code(load_order)
code = generate_workflow(load_order)
print(code)