cleanup + bugfixes

python/chat/classes.py

@@ -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
+
+

python/chat/comfy_types.py

@@ -0,0 +1,35 @@
+from nodes import NODE_CLASS_MAPPINGS
+
+def get_inputs_outputs():
+    node_inputs_outputs = {}
+
+    for key in NODE_CLASS_MAPPINGS:
+        data = NODE_CLASS_MAPPINGS[key]()
+        input_types = data.INPUT_TYPES()
+        output_types = data.RETURN_TYPES
+
+        node_inputs_outputs[key] = {
+            "input_types": input_types["required"],
+            "output_types": output_types
+        }
+
+    return node_inputs_outputs
+
+def prepare_input_output_types():
+    all_input_output_types = {}
+    for key, value in input_output_types.items():
+        class_inputs = value["input_types"] 
+        class_outputs = value["output_types"]
+
+        for key, value in class_inputs.items():
+            arg = key
+            typ = value[0]
+            all_input_output_types[arg] = typ
+    return all_input_output_types
+
+def get_type_from_input_name(input_name):
+    return all_input_output_types[input_name]
+
+
+input_output_types = get_inputs_outputs()
+all_input_output_types = prepare_input_output_types()

python/chat/format.py

@@ -0,0 +1,19 @@
+import dataclasses
+import json
+from .classes import ComfyUIWorkflowGraph
+
+def preprocess_bot_response(bot_response):
+    # Remove triple ticks
+    bot_response = bot_response.replace("```json", "")
+    bot_response = bot_response.replace("```", "")
+    return json.loads(bot_response)
+
+def postprocess(graph: ComfyUIWorkflowGraph) -> dict:
+    # Clean up intermediate values
+    for node in graph.nodes:
+        node.__pre_serialize__()
+
+    serialized = dataclasses.asdict(graph)
+    return serialized
+
+

python/chat/graph.py

@@ -0,0 +1,59 @@
+from .classes import ComfyUIWorkflowGraph, ComfyUIWorkflowGraphNode, ComfyUIWorkflowGraphLink, LLMGeneratedGraph, _get_widget_values
+from .comfy_types import input_output_types, get_type_from_input_name
+from nodes import NODE_CLASS_MAPPINGS
+
+
+def transform_input_to_graph(llm_graph) -> ComfyUIWorkflowGraph:
+    llm_graph = LLMGeneratedGraph(
+        nodes=llm_graph["nodes"],
+        edges=llm_graph["edges"]
+    )
+    last_node_id = llm_graph.get_last_node_id()
+    last_link_id = llm_graph.get_last_link_id()
+
+    comfy_graph = ComfyUIWorkflowGraph(
+        last_node_id=llm_graph.get_last_node_id(),
+        last_link_id=llm_graph.get_last_link_id(),
+    )
+
+    # Add nodes
+    for index, node in enumerate(llm_graph.nodes):
+        widget_values = _get_widget_values(node)
+        new_node = ComfyUIWorkflowGraphNode(
+            id=node["id"],
+            type=node["class_type"],
+            order=index,
+            widgets_values=widget_values
+        )
+
+        comfy_graph.add_node(new_node)
+
+    for index, edge in enumerate(llm_graph.edges):
+        link_id = index + 1
+        to_node = comfy_graph.get_node_by_id(edge['to'])
+        from_node = comfy_graph.get_node_by_id(edge['from'])
+
+        to_node.add_input(
+            edge['to_input'],
+            get_type_from_input_name(edge['to_input']),
+            link_id,
+        )
+        from_node.add_output(
+            link_id,
+            edge['from_output']
+        )
+
+        # Create links
+        link = ComfyUIWorkflowGraphLink(
+            id=link_id,
+            from_node_id=edge['from'],
+            from_node_output_slot=edge['from_output'],
+            to_node_id=edge['to'],
+            to_node_input_slot=to_node._get_input_slot_index_from_name(edge['to_input']),
+            from_node_output_type=from_node._get_output_type_from_slot_index(edge['from_output'])
+        )
+
+        comfy_graph.add_link(link)
+
+
+    return comfy_graph

python/custom_routes.py

@@ -26,7 +26,8 @@ from .utils.task import task_create, task_set_status, task_set_progress, task_ge
 from .utils.requirements import update_requirements
 from .utils.custom_nodes import get_custom_node_list_silent
 
-from nodes import NODE_CLASS_MAPPINGS
+from .chat.graph import transform_input_to_graph
+from .chat.format import preprocess_bot_response, postprocess
 
 
 @server.PromptServer.instance.routes.post("/comfy-cloud/validate-input-path")
@@ -105,23 +106,48 @@ async def get_custom_nodes_list(request):
     custom_nodes = get_custom_node_list_silent()
     return web.json_response({'custom_nodes': custom_nodes}, content_type='application/json')
 
+@server.PromptServer.instance.routes.post("/comfy-cloud/send-message")
+async def chat(request):
+    try:
+        json_data = await request.json()
+        auth_header = request.headers.get('Authorization')
 
-@server.PromptServer.instance.routes.get("/comfy-cloud/nodes-inputs-outputs")
-async def get_nodes_inputs_outputs(request):
-    node_inputs_outputs = {}
-
-    for key in NODE_CLASS_MAPPINGS:
-        data = NODE_CLASS_MAPPINGS[key]()
-        input_types = data.INPUT_TYPES()
-        output_types = data.RETURN_TYPES
-
-        node_inputs_outputs[key] = {
-            "input_types": input_types["required"],
-            "output_types": output_types
+        message = json_data["message"]
+        origin = "Comfy Cloud"
+        url = "https://comfyui-cloud-a0cf78bd8c3d.herokuapp.com/chat/send-message"
+        payload = {
+            "message": message,
+            "origin": origin
+        }
+        headers = {
+            "Content-Type": "application/json",
+            "Authorization": auth_header
         }
-    return web.json_response({'nodes': node_inputs_outputs}, content_type='application/json')
 
-            
+        response_data = None
+        response = requests.post(url, json=payload, headers=headers)
+        if response.status_code == 200:
+            response_data = response.json()
+        else:
+            raise Exception("Failed to send message to API")
+
+        bot_response = response_data["responses"]["bot"]
+
+        # Parse response 
+        bot_response = preprocess_bot_response(bot_response)
+
+        # Construct comfyui node
+        graph = transform_input_to_graph(bot_response)
+        data = postprocess(graph)
+        response = requests.post(url, json=payload, headers=headers)
+
+        return web.json_response({'nodes': data}, content_type='application/json')
+
+    except Exception as e:
+        print("Error:", e)
+        return web.json_response({ "error": e }, status=400)
+
+
 async def upload_task_execution(task_id, file_specs, workflow_id):
     try:
         task_set_status(task_id, TaskStatus.HASHING)