First draft of major code refactoring and adding support for custom nodes.
This commit is contained in:
Peyton
+252
-276
@@ -1,6 +1,5 @@
|
||||
import glob
|
||||
import copy
|
||||
import inspect
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import sys
|
||||
@@ -9,6 +8,9 @@ import black
|
||||
|
||||
from typing import Dict, List, Any, Callable
|
||||
|
||||
from utils import import_custom_nodes, read_json_file, write_code_to_file, add_comfyui_directories_to_sys_path
|
||||
|
||||
|
||||
sys.path.append('../')
|
||||
|
||||
from nodes import NODE_CLASS_MAPPINGS
|
||||
@@ -16,318 +18,292 @@ from nodes import NODE_CLASS_MAPPINGS
|
||||
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
|
||||
|
||||
def read_json_file(file_path: str) -> dict:
|
||||
"""
|
||||
Reads a JSON file and returns its contents as a dictionary.
|
||||
|
||||
Args:
|
||||
file_path (str): The path to the JSON file.
|
||||
|
||||
Returns:
|
||||
dict: The contents of the JSON file as a dictionary.
|
||||
|
||||
Raises:
|
||||
FileNotFoundError: If the file is not found, it lists all JSON files in the directory of the file path.
|
||||
ValueError: If the file is not a valid JSON.
|
||||
"""
|
||||
|
||||
try:
|
||||
with open(file_path, 'r') as file:
|
||||
data = json.load(file)
|
||||
return data
|
||||
|
||||
except FileNotFoundError:
|
||||
# Get the directory from the file_path
|
||||
directory = os.path.dirname(file_path)
|
||||
|
||||
# If the directory is an empty string (which means file is in the current directory),
|
||||
# get the current working directory
|
||||
if not directory:
|
||||
directory = os.getcwd()
|
||||
|
||||
# Find all JSON files in the directory
|
||||
json_files = glob.glob(f"{directory}/*.json")
|
||||
|
||||
# Format the list of JSON files as a string
|
||||
json_files_str = "\n".join(json_files)
|
||||
|
||||
raise FileNotFoundError(f"\n\nFile not found: {file_path}. JSON files in the directory:\n{json_files_str}")
|
||||
|
||||
except json.JSONDecodeError:
|
||||
raise ValueError(f"Invalid JSON format in file: {file_path}")
|
||||
|
||||
|
||||
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.
|
||||
class ComfyUItoPython:
|
||||
def __init__(self, input='workflow_api.json', queue_size=10):
|
||||
self.input = input
|
||||
self.queue_size = queue_size
|
||||
self.BASE_NODE_CLASS_MAPPINGS = copy.deepcopy(NODE_CLASS_MAPPINGS)
|
||||
self.generate_code()
|
||||
|
||||
Args:
|
||||
data (Dict):
|
||||
The dictionary for which to determine the load order.
|
||||
|
||||
Returns:
|
||||
List:
|
||||
A list of tuples where each tuple contains a key, its corresponding dictionary,
|
||||
and a boolean indicating whether or not the function is dependent on the output of
|
||||
a previous function, ordered by load order.
|
||||
"""
|
||||
|
||||
# Create a dictionary to keep track of visited nodes.
|
||||
visited = {}
|
||||
# Create a list to store the load order for functions
|
||||
load_order = []
|
||||
# Boolean to indicate whether or not the class is a loader class that should not be
|
||||
# reloaded during every loop
|
||||
is_loader = False
|
||||
|
||||
def dfs(key: str) -> None:
|
||||
def generate_code(self):
|
||||
"""
|
||||
Depth-First Search function.
|
||||
Main function to be executed.
|
||||
"""
|
||||
import_custom_nodes()
|
||||
# Load JSON data from the input file
|
||||
prompt = read_json_file(self.input)
|
||||
load_order = self.determine_load_order(prompt)
|
||||
output_file = self.input.replace('.json', '.py')
|
||||
code = self.generate_workflow(load_order, filename=output_file, queue_size=self.queue_size)
|
||||
logging.info(code)
|
||||
|
||||
def determine_load_order(self, 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:
|
||||
key (str): The key from which to start the DFS.
|
||||
data (Dict):
|
||||
The dictionary for which to determine the load order.
|
||||
|
||||
Returns:
|
||||
None
|
||||
List:
|
||||
A list of tuples where each tuple contains a key, its corresponding dictionary,
|
||||
and a boolean indicating whether or not the function is dependent on the output of
|
||||
a previous function, ordered by load order.
|
||||
"""
|
||||
# Mark the node as visited.
|
||||
visited[key] = True
|
||||
inputs = data[key]['inputs']
|
||||
|
||||
# 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])
|
||||
# Create a dictionary to keep track of visited nodes.
|
||||
visited = {}
|
||||
# Create a list to store the load order for functions
|
||||
load_order = []
|
||||
# Boolean to indicate whether or not the class is a loader class that should not be
|
||||
# reloaded during every loop
|
||||
is_loader = False
|
||||
|
||||
# Add the key and its corresponding data to the load order list.
|
||||
load_order.append((key, data[key], is_loader))
|
||||
def dfs(key: str) -> None:
|
||||
"""
|
||||
Depth-First Search function.
|
||||
|
||||
# Load Loader keys first
|
||||
for key in data:
|
||||
class_def = NODE_CLASS_MAPPINGS[data[key]['class_type']]()
|
||||
if class_def.CATEGORY == 'loaders' or class_def.FUNCTION in ['encode'] or not any(isinstance(val, list) for val in data[key]['inputs'].values()):
|
||||
is_loader = True
|
||||
Args:
|
||||
key (str): The key from which to start the DFS.
|
||||
|
||||
Returns:
|
||||
None
|
||||
"""
|
||||
# Mark the node as visited.
|
||||
visited[key] = True
|
||||
inputs = data[key]['inputs']
|
||||
|
||||
# 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], is_loader))
|
||||
|
||||
# Load Loader keys first
|
||||
for key in data:
|
||||
class_def = NODE_CLASS_MAPPINGS[data[key]['class_type']]()
|
||||
if class_def.CATEGORY == 'loaders' or class_def.FUNCTION in ['encode'] or not any(isinstance(val, list) for val in data[key]['inputs'].values()):
|
||||
is_loader = True
|
||||
# If the key has not been visited, perform a DFS from that key.
|
||||
if key not in visited:
|
||||
dfs(key)
|
||||
|
||||
# Reset is_loader bool
|
||||
is_loader = False
|
||||
# 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)
|
||||
|
||||
# Reset is_loader bool
|
||||
is_loader = False
|
||||
# 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
|
||||
return load_order
|
||||
|
||||
|
||||
def create_function_call_code(obj_name: str, func: str, variable_name: str, is_loader: bool, **kwargs) -> str:
|
||||
"""
|
||||
This function generates Python code for a function call.
|
||||
def create_function_call_code(self, obj_name: str, func: str, variable_name: str, is_loader: bool, **kwargs) -> str:
|
||||
"""
|
||||
This function generates Python code for a function call.
|
||||
|
||||
Args:
|
||||
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.
|
||||
is_loader (bool): Determines the code indentation.
|
||||
**kwargs: The keyword arguments for the function.
|
||||
Args:
|
||||
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.
|
||||
is_loader (bool): Determines the code indentation.
|
||||
**kwargs: The keyword arguments for the function.
|
||||
|
||||
Returns:
|
||||
str: The generated Python code.
|
||||
"""
|
||||
Returns:
|
||||
str: The generated Python code.
|
||||
"""
|
||||
|
||||
def format_arg(key: str, value: any) -> str:
|
||||
"""Formats arguments based on key and value."""
|
||||
if key == 'noise_seed':
|
||||
return f'{key}=random.randint(1, 2**64)'
|
||||
elif isinstance(value, str):
|
||||
value = value.replace("\n", "\\n")
|
||||
return f'{key}="{value}"'
|
||||
elif isinstance(value, dict) and 'variable_name' in value:
|
||||
return f'{key}={value["variable_name"]}'
|
||||
return f'{key}={value}'
|
||||
def format_arg(key: str, value: any) -> str:
|
||||
"""Formats arguments based on key and value."""
|
||||
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}"'
|
||||
elif isinstance(value, dict) and 'variable_name' in value:
|
||||
return f'{key}={value["variable_name"]}'
|
||||
return f'{key}={value}'
|
||||
|
||||
args = ', '.join(format_arg(key, value) for key, value in kwargs.items())
|
||||
args = ', '.join(format_arg(key, value) for key, value in kwargs.items())
|
||||
|
||||
# Generate the Python code
|
||||
code = f'{variable_name} = {obj_name}.{func}({args})\n'
|
||||
# Generate the Python code
|
||||
code = f'{variable_name} = {obj_name}.{func}({args})\n'
|
||||
|
||||
# If the code contains dependencies, indent the code because it will be placed inside
|
||||
# of a for loop
|
||||
if not is_loader:
|
||||
code = f'\t{code}'
|
||||
# If the code contains dependencies, indent the code because it will be placed inside
|
||||
# of a for loop
|
||||
if not is_loader:
|
||||
code = f'\t{code}'
|
||||
|
||||
return code
|
||||
return code
|
||||
|
||||
|
||||
def update_inputs(inputs: Dict, executed_variables: Dict) -> Dict:
|
||||
"""
|
||||
Update inputs based on the executed variables.
|
||||
def update_inputs(self, inputs: Dict, executed_variables: Dict) -> Dict:
|
||||
"""
|
||||
Update inputs based on the executed variables.
|
||||
|
||||
Args:
|
||||
inputs (Dict): Inputs dictionary to update.
|
||||
executed_variables (Dict): Dictionary storing executed variable names.
|
||||
Args:
|
||||
inputs (Dict): Inputs dictionary to update.
|
||||
executed_variables (Dict): Dictionary storing executed variable names.
|
||||
|
||||
Returns:
|
||||
Dict: Updated inputs dictionary.
|
||||
"""
|
||||
for key in inputs.keys():
|
||||
if isinstance(inputs[key], list) and inputs[key][0] in executed_variables.keys():
|
||||
inputs[key] = {'variable_name': f"{executed_variables[inputs[key][0]]}[{inputs[key][1]}]"}
|
||||
return inputs
|
||||
Returns:
|
||||
Dict: Updated inputs dictionary.
|
||||
"""
|
||||
for key in inputs.keys():
|
||||
if isinstance(inputs[key], list) and inputs[key][0] in executed_variables.keys():
|
||||
inputs[key] = {'variable_name': f"get_value_at_index({executed_variables[inputs[key][0]]}, {inputs[key][1]})"}
|
||||
return inputs
|
||||
|
||||
|
||||
def get_class_info(class_type: str) -> (str, str, str):
|
||||
"""
|
||||
Generates and returns necessary information about class type.
|
||||
def get_class_info(self, class_type: str) -> (str, str, str):
|
||||
"""
|
||||
Generates and returns necessary information about class type.
|
||||
|
||||
Args:
|
||||
class_type (str): 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
|
||||
"""
|
||||
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, loader_code: List[str], code: List[str], queue_size: int) -> str:
|
||||
"""
|
||||
Generates final code string.
|
||||
|
||||
Args:
|
||||
import_statements (set): A set of unique import statements
|
||||
code (List[str]): A list of code strings
|
||||
queue_size (int): Number of photos that will be generated by the script.
|
||||
|
||||
Returns:
|
||||
final_code (str): Generated final code as a string
|
||||
"""
|
||||
static_imports = ['import random', 'import torch']
|
||||
imports_code = [f"from nodes import {', '.join([class_name for class_name in import_statements])}" ]
|
||||
main_function_code = f"def main():\n\t" + 'with torch.inference_mode():\n\t\t' + '\n\t\t'.join(loader_code) + f'\n\n\t\tfor q in range({queue_size}):\n\t\t' + '\n\t\t'.join(code)
|
||||
final_code = '\n'.join(static_imports + ['import sys\nsys.path.append("../")'] + imports_code + ['', main_function_code, '', 'if __name__ == "__main__":', '\tmain()'])
|
||||
final_code = black.format_str(final_code, mode=black.Mode())
|
||||
|
||||
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 get_function_parameters(func: Callable) -> List:
|
||||
"""Get the names of a function's parameters.
|
||||
|
||||
Args:
|
||||
func (Callable): The function whose parameters we want to inspect.
|
||||
|
||||
Returns:
|
||||
List: A list containing the names of the function's parameters.
|
||||
"""
|
||||
signature = inspect.signature(func)
|
||||
parameters = {name: param.default if param.default != param.empty else None
|
||||
for name, param in signature.parameters.items()}
|
||||
return list(parameters.keys())
|
||||
|
||||
|
||||
def generate_workflow(load_order: List, filename: str = 'generated_code_workflow.py', queue_size: int = 10) -> str:
|
||||
"""
|
||||
Generate the execution code based on the load order.
|
||||
|
||||
Args:
|
||||
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'.
|
||||
queue_size (int): The number of photos that will be created by the script.
|
||||
|
||||
Returns:
|
||||
str: Generated execution code as a string.
|
||||
"""
|
||||
|
||||
# Create the necessary data structures to hold imports and generated code
|
||||
import_statements, executed_variables, loader_code, 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, is_loader in load_order:
|
||||
|
||||
# Generate class definition and inputs from the data
|
||||
inputs, class_type = data['inputs'], data['class_type']
|
||||
class_def = NODE_CLASS_MAPPINGS[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)
|
||||
loader_code.append(class_code)
|
||||
|
||||
# Get all possible parameters for class_def
|
||||
class_def_params = get_function_parameters(getattr(class_def, class_def.FUNCTION))
|
||||
|
||||
# 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}
|
||||
|
||||
# Create executed variable and generate code
|
||||
executed_variables[idx] = f'{class_type.lower()}_{idx}'
|
||||
inputs = update_inputs(inputs, executed_variables)
|
||||
|
||||
if is_loader:
|
||||
loader_code.append(create_function_call_code(initialized_objects[class_type], class_def.FUNCTION, executed_variables[idx], is_loader, **inputs))
|
||||
Returns:
|
||||
class_type (str): Updated class type
|
||||
import_statement (str): Import statement string
|
||||
class_code (str): Class initialization code
|
||||
"""
|
||||
import_statement = class_type
|
||||
if class_type in self.BASE_NODE_CLASS_MAPPINGS.keys():
|
||||
class_code = f'{class_type.lower()} = {class_type}()'
|
||||
else:
|
||||
code.append(create_function_call_code(initialized_objects[class_type], class_def.FUNCTION, executed_variables[idx], is_loader, **inputs))
|
||||
class_code = f'{class_type.lower()} = NODE_CLASS_MAPPINGS["{class_type}"]()'
|
||||
|
||||
# Generate final code by combining imports and code, and wrap them in a main function
|
||||
final_code = assemble_python_code(import_statements, loader_code, code, queue_size)
|
||||
|
||||
# Save the code to a .py file
|
||||
write_code_to_file(filename, final_code)
|
||||
|
||||
return final_code
|
||||
return class_type, import_statement, class_code
|
||||
|
||||
|
||||
def main(input, queue_size=10):
|
||||
"""
|
||||
Main function to be executed.
|
||||
"""
|
||||
# Load JSON data from the input file
|
||||
prompt = read_json_file(input)
|
||||
load_order = determine_load_order(prompt)
|
||||
output_file = input.replace('.json', '.py')
|
||||
code = generate_workflow(load_order, filename=output_file, queue_size=queue_size)
|
||||
logging.info(code)
|
||||
def assemble_python_code(self, import_statements: set, loader_code: List[str], code: List[str], queue_size: int, custom_nodes=False) -> str:
|
||||
"""
|
||||
Generates final code string.
|
||||
|
||||
Args:
|
||||
import_statements (set): A set of unique import statements
|
||||
code (List[str]): A list of code strings
|
||||
queue_size (int): Number of photos that will be generated by the script.
|
||||
custom_nodes (bool): Whether to include custom nodes in the code.
|
||||
|
||||
Returns:
|
||||
final_code (str): Generated final code as a string
|
||||
"""
|
||||
# Get the source code of the function as a string
|
||||
add_comfyui_directories_to_sys_path_code = inspect.getsource(add_comfyui_directories_to_sys_path)
|
||||
# Define static import statements required for the script
|
||||
static_imports = ['import os', 'import random', 'import sys', 'import torch', f'\n{add_comfyui_directories_to_sys_path_code}',
|
||||
'\n\nadd_comfyui_directories_to_sys_path()']
|
||||
# Check if custom nodes should be included
|
||||
if custom_nodes:
|
||||
static_imports.append('\nfrom utils import import_custom_nodes, get_value_at_index\n')
|
||||
custom_nodes = 'import_custom_nodes()\n\t'
|
||||
else:
|
||||
custom_nodes = ''
|
||||
# Create import statements for node classes
|
||||
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(loader_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()'])
|
||||
# Format the final code according to PEP 8 using the Black library
|
||||
final_code = black.format_str(final_code, mode=black.Mode())
|
||||
|
||||
return final_code
|
||||
|
||||
|
||||
|
||||
def get_function_parameters(self, func: Callable) -> List:
|
||||
"""Get the names of a function's parameters.
|
||||
|
||||
Args:
|
||||
func (Callable): The function whose parameters we want to inspect.
|
||||
|
||||
Returns:
|
||||
List: A list containing the names of the function's parameters.
|
||||
"""
|
||||
signature = inspect.signature(func)
|
||||
parameters = {name: param.default if param.default != param.empty else None
|
||||
for name, param in signature.parameters.items()}
|
||||
return list(parameters.keys())
|
||||
|
||||
|
||||
def generate_workflow(self, load_order: List, filename: str = 'generated_code_workflow.py', queue_size: int = 10) -> str:
|
||||
"""
|
||||
Generate the execution code based on the load order.
|
||||
|
||||
Args:
|
||||
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'.
|
||||
queue_size (int): The number of photos that will be created by the script.
|
||||
|
||||
Returns:
|
||||
str: Generated execution code as a string.
|
||||
"""
|
||||
|
||||
# Create the necessary data structures to hold imports and generated code
|
||||
import_statements, executed_variables, loader_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_loader in load_order:
|
||||
|
||||
# Generate class definition and inputs from the data
|
||||
inputs, class_type = data['inputs'], data['class_type']
|
||||
class_def = NODE_CLASS_MAPPINGS[class_type]()
|
||||
|
||||
# If the class hasn't been initialized yet, initialize it and generate the import statements
|
||||
if class_type not in initialized_objects:
|
||||
# No need to use preview image nodes since we are executing the script in a terminal
|
||||
if class_type == 'PreviewImage':
|
||||
continue
|
||||
|
||||
class_type, import_statement, class_code = self.get_class_info(class_type)
|
||||
initialized_objects[class_type] = class_type.lower()
|
||||
if class_type in self.BASE_NODE_CLASS_MAPPINGS.keys():
|
||||
import_statements.add(import_statement)
|
||||
if class_type not in self.BASE_NODE_CLASS_MAPPINGS.keys():
|
||||
custom_nodes = True
|
||||
loader_code.append(class_code)
|
||||
|
||||
# Get all possible parameters for class_def
|
||||
class_def_params = self.get_function_parameters(getattr(class_def, class_def.FUNCTION))
|
||||
|
||||
# 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}
|
||||
|
||||
# Create executed variable and generate code
|
||||
executed_variables[idx] = f'{class_type.lower()}_{idx}'
|
||||
inputs = self.update_inputs(inputs, executed_variables)
|
||||
|
||||
if is_loader:
|
||||
loader_code.append(self.create_function_call_code(initialized_objects[class_type], class_def.FUNCTION, executed_variables[idx], is_loader, **inputs))
|
||||
else:
|
||||
code.append(self.create_function_call_code(initialized_objects[class_type], class_def.FUNCTION, executed_variables[idx], is_loader, **inputs))
|
||||
|
||||
# Generate final code by combining imports and code, and wrap them in a main function
|
||||
final_code = self.assemble_python_code(import_statements, loader_code, code, queue_size, custom_nodes)
|
||||
|
||||
# Save the code to a .py file
|
||||
write_code_to_file(filename, final_code)
|
||||
|
||||
return final_code
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
input = 'workflow_api.json'
|
||||
input = 'workflow_api_serge.json'
|
||||
queue_size = 10
|
||||
main(input, queue_size)
|
||||
ComfyUItoPython(input=input, queue_size=queue_size)
|
||||
|
||||
@@ -0,0 +1,150 @@
|
||||
import asyncio
|
||||
import json
|
||||
import glob
|
||||
import os
|
||||
from typing import Sequence, Mapping, Any, Union
|
||||
import sys
|
||||
|
||||
sys.path.append('../')
|
||||
|
||||
import execution
|
||||
from nodes import init_custom_nodes
|
||||
|
||||
|
||||
def read_json_file(file_path: str) -> dict:
|
||||
"""
|
||||
Reads a JSON file and returns its contents as a dictionary.
|
||||
|
||||
Args:
|
||||
file_path (str): The path to the JSON file.
|
||||
|
||||
Returns:
|
||||
dict: The contents of the JSON file as a dictionary.
|
||||
|
||||
Raises:
|
||||
FileNotFoundError: If the file is not found, it lists all JSON files in the directory of the file path.
|
||||
ValueError: If the file is not a valid JSON.
|
||||
"""
|
||||
|
||||
try:
|
||||
with open(file_path, 'r') as file:
|
||||
data = json.load(file)
|
||||
return data
|
||||
|
||||
except FileNotFoundError:
|
||||
# Get the directory from the file_path
|
||||
directory = os.path.dirname(file_path)
|
||||
|
||||
# If the directory is an empty string (which means file is in the current directory),
|
||||
# get the current working directory
|
||||
if not directory:
|
||||
directory = os.getcwd()
|
||||
|
||||
# Find all JSON files in the directory
|
||||
json_files = glob.glob(f"{directory}/*.json")
|
||||
|
||||
# Format the list of JSON files as a string
|
||||
json_files_str = "\n".join(json_files)
|
||||
|
||||
raise FileNotFoundError(f"\n\nFile not found: {file_path}. JSON files in the directory:\n{json_files_str}")
|
||||
|
||||
except json.JSONDecodeError:
|
||||
raise ValueError(f"Invalid JSON format in file: {file_path}")
|
||||
|
||||
|
||||
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 import_custom_nodes() -> None:
|
||||
"""Find all custom nodes in the custom_nodes folder and add those node objects to NODE_CLASS_MAPPINGS
|
||||
|
||||
This function sets up a new asyncio event loop, initializes the PromptServer,
|
||||
creates a PromptQueue, and initializes the custom nodes.
|
||||
"""
|
||||
import server
|
||||
|
||||
# Creating a new event loop and setting it as the default loop
|
||||
loop = asyncio.new_event_loop()
|
||||
asyncio.set_event_loop(loop)
|
||||
|
||||
# Creating an instance of PromptServer with the loop
|
||||
server_instance = server.PromptServer(loop)
|
||||
execution.PromptQueue(server_instance)
|
||||
|
||||
# Initializing custom nodes
|
||||
init_custom_nodes()
|
||||
|
||||
|
||||
def add_comfyui_directories_to_sys_path() -> None:
|
||||
"""
|
||||
Recursively looks at parent folders starting from the current working directory until it finds 'ComfyUI' and 'ComfyUI-to-Python-Extension'.
|
||||
Once found, the directories are added to sys.path.
|
||||
"""
|
||||
start_path = os.getcwd() # Get the current working directory
|
||||
|
||||
def search_directory(path: str) -> None:
|
||||
# Check if the current directory contains 'ComfyUI' or 'ComfyUI-to-Python-Extension'
|
||||
for directory_name in ['ComfyUI', 'ComfyUI-to-Python-Extension']:
|
||||
if directory_name in os.listdir(path):
|
||||
directory_path = os.path.join(path, directory_name)
|
||||
sys.path.append(directory_path)
|
||||
print(f"'{directory_name}' found and added to sys.path: {directory_path}")
|
||||
|
||||
# Get the parent directory
|
||||
parent_directory = os.path.dirname(path)
|
||||
|
||||
# If the parent directory is the same as the current directory, we've reached the root and stop the search
|
||||
if parent_directory == path:
|
||||
return
|
||||
|
||||
# Recursively call the function with the parent directory
|
||||
search_directory(parent_directory)
|
||||
|
||||
# Start the search from the current working directory
|
||||
search_directory(start_path)
|
||||
|
||||
# Example usage
|
||||
add_comfyui_directories_to_sys_path()
|
||||
|
||||
|
||||
|
||||
def get_value_at_index(obj: Union[Sequence, Mapping], index: int) -> Any:
|
||||
"""Returns the value at the given index of a sequence or mapping.
|
||||
|
||||
If the object is a sequence (like list or string), returns the value at the given index.
|
||||
If the object is a mapping (like a dictionary), returns the value at the index-th key.
|
||||
|
||||
Some return a dictionary, in these cases, we look for the "results" key
|
||||
|
||||
Args:
|
||||
obj (Union[Sequence, Mapping]): The object to retrieve the value from.
|
||||
index (int): The index of the value to retrieve.
|
||||
|
||||
Returns:
|
||||
Any: The value at the given index.
|
||||
|
||||
Raises:
|
||||
IndexError: If the index is out of bounds for the object and the object is not a mapping.
|
||||
"""
|
||||
try:
|
||||
return obj[index]
|
||||
except KeyError:
|
||||
return obj['result'][index]
|
||||
Reference in New Issue
Block a user