198 lines
6.3 KiB
Python
198 lines
6.3 KiB
Python
from dataclasses import dataclass, field
|
|
from typing import List, Dict, Any
|
|
from nodes import NODE_CLASS_MAPPINGS
|
|
|
|
@dataclass
|
|
class LLMGeneratedGraph:
|
|
nodes: dict
|
|
edges: dict
|
|
|
|
def get_last_link_id(self):
|
|
return len(self.edges)
|
|
|
|
def get_last_node_id(self):
|
|
return len(self.edges)
|
|
|
|
|
|
# TODO: move to inside ComfyUIWorkflowGraphNodeClass
|
|
def _get_widget_values(node):
|
|
class_type = node['class_type']
|
|
registered_node = NODE_CLASS_MAPPINGS[class_type]()
|
|
|
|
node_widget_values = []
|
|
if 'widget_values' in node:
|
|
node_widget_values = node['widget_values']
|
|
|
|
registered_node_inputs_dict = registered_node.INPUT_TYPES()["required"] # dictionary
|
|
final_widget_values = {}
|
|
|
|
# ComfyUI represents a dropdown widget value with
|
|
# index 0 of the tuple as a list
|
|
def is_list_selection_widget_value(value):
|
|
return type(value[0]) == list
|
|
|
|
# Prepare the widget values dictionary for the registered custom node's input types
|
|
registered_node_widget_values = {}
|
|
for key, value in registered_node_inputs_dict.items():
|
|
if len(value) > 1 or is_list_selection_widget_value(value):
|
|
registered_node_widget_values[key] = value
|
|
|
|
# Build widget values dictionary
|
|
for key, value in registered_node_widget_values.items():
|
|
if key in node_widget_values:
|
|
# If the key exists in the LLM generated output, use that value
|
|
final_widget_values[key] = node_widget_values[key]
|
|
|
|
else:
|
|
# Load final_widget_values with the default
|
|
if is_list_selection_widget_value(value):
|
|
# If it is a dropdown, just choose the first value of the list
|
|
final_widget_values[key] = value[0][0]
|
|
else:
|
|
final_widget_values[key] = value[1]["default"]
|
|
|
|
final_widget_values_list = list(final_widget_values.values())
|
|
|
|
# edge case:
|
|
# For some reason, KSampler widget values do not account for 'control_after_generate'.
|
|
if node['class_type'] == "KSampler":
|
|
final_widget_values_list.insert(1, "randomize")
|
|
|
|
|
|
return final_widget_values_list
|
|
|
|
|
|
|
|
@dataclass
|
|
class ComfyUIWorkflowGraphNode:
|
|
id: int
|
|
type: str
|
|
order: int
|
|
mode: int = 0
|
|
pos: List[int] = field(default_factory=list)
|
|
flags: Dict[str, Any] = field(default_factory=dict)
|
|
inputs: List[object] = field(default_factory=list)
|
|
outputs: List[object] = field(default_factory=list)
|
|
properties: Dict[str, Any] = field(default_factory=dict)
|
|
widgets_values: List[Any] = field(default_factory=list)
|
|
|
|
# This is the object from NODE_CLASS_MAPPINGS and the
|
|
# internal object used by ComfyUI
|
|
_registered_node: Any = field(init=False)
|
|
|
|
def __post_init__(self):
|
|
node = NODE_CLASS_MAPPINGS[self.type]()
|
|
self._registered_node = node
|
|
|
|
# Prepare input types
|
|
registered_node_inputs_dict = self._registered_node.INPUT_TYPES()["required"] # dictionary
|
|
registered_node_inputs_list = list(registered_node_inputs_dict.keys())
|
|
|
|
for index, (key, val) in enumerate(registered_node_inputs_dict.items()):
|
|
if len(val) > 1 or type(val[0]) == list:
|
|
continue
|
|
|
|
self.inputs.append({
|
|
'name': key,
|
|
'type': val[0],
|
|
'link': None,
|
|
})
|
|
|
|
# Prepare return types
|
|
registered_node_outputs_tuple = self._registered_node.RETURN_TYPES
|
|
for index, ret_typ in enumerate(registered_node_outputs_tuple):
|
|
self.outputs.append({
|
|
'name': ret_typ,
|
|
'type': ret_typ,
|
|
'links': [],
|
|
'slot_index': index
|
|
})
|
|
|
|
|
|
def __pre_serialize__(self):
|
|
# _registered_node is a class from nodes.py
|
|
# and does not need to be sent to the front end
|
|
self._registered_node = None
|
|
|
|
def _get_output_type_from_slot_index(self, slot_index):
|
|
registered_node = self._registered_node
|
|
registered_node_outputs_tuple = registered_node.RETURN_TYPES
|
|
print(registered_node_outputs_tuple, slot_index)
|
|
return registered_node_outputs_tuple[slot_index]
|
|
|
|
def _get_input_slot_index_from_name(self, key):
|
|
registered_node = self._registered_node
|
|
registered_node_inputs_dict = registered_node.INPUT_TYPES()["required"] # dictionary
|
|
registered_node_inputs_list = list(registered_node_inputs_dict.keys())
|
|
|
|
print(registered_node_inputs_list)
|
|
|
|
# find the index of the key in the list
|
|
return registered_node_inputs_list.index(key)
|
|
|
|
def add_input(self, name: str, typ: str, link: int):
|
|
# find by name
|
|
for input in self.inputs:
|
|
if input['name'] == name:
|
|
input['link'] = link
|
|
|
|
|
|
def add_output(self, link_id, slot_index):
|
|
# Append link_id to existing slot index if it exists
|
|
for output in self.outputs:
|
|
if output['slot_index'] == slot_index:
|
|
output['links'].append(link_id)
|
|
|
|
|
|
@dataclass
|
|
class ComfyUIWorkflowGraphLink:
|
|
id: int
|
|
from_node_id: int
|
|
from_node_output_slot: int
|
|
to_node_id: int
|
|
to_node_input_slot: int
|
|
from_node_output_type: str
|
|
|
|
def to_list(self):
|
|
return [
|
|
self.id,
|
|
self.from_node_id,
|
|
self.from_node_output_slot,
|
|
self.to_node_id,
|
|
self.to_node_input_slot,
|
|
self.from_node_output_type
|
|
]
|
|
|
|
@dataclass
|
|
class ComfyUIWorkflowGraph:
|
|
"""
|
|
This is the structure accepted by ComfyUI's LoadGraph method.
|
|
"""
|
|
|
|
last_node_id: int
|
|
last_link_id: int
|
|
nodes: List[ComfyUIWorkflowGraphNode] = field(default_factory=list)
|
|
links: List[object] = field(default_factory=list)
|
|
groups: List[object] = field(default_factory=list)
|
|
config: Dict[str, object] = field(default_factory=dict)
|
|
extra: Dict[str, object] = field(default_factory=dict)
|
|
version: float = 0.4
|
|
|
|
# For easy access to nodes by id
|
|
_nodes_dict: Dict[str, int] = field(init=False, default_factory=dict)
|
|
|
|
def add_node(self, node: ComfyUIWorkflowGraphNode):
|
|
self.nodes.append(node)
|
|
node_index = len(self.nodes) - 1
|
|
self._nodes_dict[node.id] = node_index
|
|
|
|
def add_link(self, link: ComfyUIWorkflowGraphLink):
|
|
self.links.append(link.to_list())
|
|
|
|
def get_node_by_id(self, node_id):
|
|
index = self._nodes_dict.get(node_id, None)
|
|
node = self.nodes[index]
|
|
return node
|
|
|
|
|