cleanup + bugfixes

2024-07-21 19:43:12 -07:00
parent 6a9e89dd7f
commit 02b7527f46
11 changed files with 385 additions and 195 deletions
@@ -0,0 +1,194 @@
+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 = 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
+
+