From 02b7527f46f87ee4c7322ec48bea3de90d421fcb Mon Sep 17 00:00:00 2001 From: Nathan Date: Sun, 21 Jul 2024 19:43:12 -0700 Subject: [PATCH] cleanup + bugfixes --- js/chatbot/{chatbot.js => chatbot.js.dep} | 0 js/chatbot/{edgeTypes.js => edgeTypes.js.dep} | 2 - js/chatbot/ext.js | 34 +-- js/chatbot/workflow.js | 171 ++------------- js/comfy/ext.js | 5 + js/resource/local.js | 5 + python/chat/classes.py | 194 ++++++++++++++++++ python/chat/comfy_types.py | 35 ++++ python/chat/format.py | 19 ++ python/chat/graph.py | 59 ++++++ python/custom_routes.py | 56 +++-- 11 files changed, 385 insertions(+), 195 deletions(-) rename js/chatbot/{chatbot.js => chatbot.js.dep} (100%) rename js/chatbot/{edgeTypes.js => edgeTypes.js.dep} (99%) create mode 100644 python/chat/classes.py create mode 100644 python/chat/comfy_types.py create mode 100644 python/chat/format.py create mode 100644 python/chat/graph.py diff --git a/js/chatbot/chatbot.js b/js/chatbot/chatbot.js.dep similarity index 100% rename from js/chatbot/chatbot.js rename to js/chatbot/chatbot.js.dep diff --git a/js/chatbot/edgeTypes.js b/js/chatbot/edgeTypes.js.dep similarity index 99% rename from js/chatbot/edgeTypes.js rename to js/chatbot/edgeTypes.js.dep index 1a94310..51b9256 100644 --- a/js/chatbot/edgeTypes.js +++ b/js/chatbot/edgeTypes.js.dep @@ -383,8 +383,6 @@ function scanNodeInputOutput() { } } - - console.log(LiteGraph.registered_node_types); if (LiteGraph.registered_node_types) { // Get the first node in the list for (const [key, nodeType] of Object.entries( diff --git a/js/chatbot/ext.js b/js/chatbot/ext.js index ce9010f..940607c 100644 --- a/js/chatbot/ext.js +++ b/js/chatbot/ext.js @@ -1,15 +1,12 @@ // import { api, app } from './comfy/comfy.js' import { app } from "../comfy/comfy.js"; -import { - getBotResponse, - isValidWorkflow, - parseWorkflowFromBot, -} from "./chatbot.js"; -import { loadGraphFromPrompt } from "./workflow.js"; +import { loadGraphFromPrompt, isValidWorkflow } from "./workflow.js"; import { setButtonDefault } from "../button/ui.js"; import { authDialog } from "../auth/index.js"; import { getApiToken } from "../utils.js"; +import { local } from '../resource/index.js'; + const CHAT_BUTTON_ID = "chat-button"; const CHAT_BOX_ID = "chat-box"; const CHAT_INPUT_ID = "chat-input"; @@ -214,16 +211,13 @@ const sendMessage = async () => { chatInput.value = ""; try { - const response = await getBotResponse(message); let botMessage = ""; - const parsedBotResponse = parseWorkflowFromBot(response); + // fetch response to localhost + const { nodes: parsedBotResponse } = await local.sendMessage({ message }); - if (isValidWorkflow(parsedBotResponse)) { - botMessage = JSON.stringify(parsedBotResponse); - } else { - botMessage = response.responses.bot; - } + console.log(parsedBotResponse); + botMessage = JSON.stringify(parsedBotResponse); const botMessageElement = document.createElement("div"); botMessageElement.innerText = botMessage; @@ -251,7 +245,7 @@ const sendMessage = async () => { } }; -const registerChat = () => { +export const registerChat = () => { // With js, append an bottom that opens up a chatbox in the bottom right hand corner const chatButton = createChatButton(); const chatBox = createChatBox(); @@ -261,15 +255,3 @@ const registerChat = () => { chatButton.addEventListener("click", toggleChatBox); }; -/** @typedef {import('../../../web/types/comfy.js').ComfyExtension} ComfyExtension*/ -/** @type {ComfyExtension} */ -const ext = { - name: "nathannlu.ComfyUI-Chat", - - // ComfyUI extension init - init() { - registerChat(); - }, -}; - -app.registerExtension(ext); diff --git a/js/chatbot/workflow.js b/js/chatbot/workflow.js index 0e923ef..1f7b995 100644 --- a/js/chatbot/workflow.js +++ b/js/chatbot/workflow.js @@ -1,75 +1,6 @@ import { app } from "../comfy/comfy.js"; -import { getInputType } from "./edgeTypes.js"; -import { - getNodesInputsOutputs -} from '../utils.js' - - -var nodesInputsOutputs = {}; -getNodesInputsOutputs().then((nodes) => { - nodesInputsOutputs = nodes; -}) - - -function getInputIndex(node, inputName) { - return node.inputs.findIndex( - (input) => input.name === getInputType(inputName) - ); -} - -function validateObject(objectA, objectB) { - const validatedObject = {}; - - // If it is undefined, just return the default values - if (!objectB) { - for (const key in objectA) { - const defaultValue = objectA[key][1]?.default; - if (defaultValue !== undefined) { - validatedObject[key] = defaultValue; - } - } - - return validatedObject; - } - - for (const key in objectA) { - if (objectB.hasOwnProperty(key)) { - // Check the type - const [_t, meta] = objectA[key]; - const defaultValue = meta?.default; - const type = typeof defaultValue; - - if (typeof objectB[key] === type || type == 'undefined') { - validatedObject[key] = objectB[key]; - } else { - validatedObject[key] = defaultValue; - } - } else { - // Use default value if key is missing in ObjectB - const defaultValue = objectA[key][1]?.default; - if (defaultValue !== undefined) { - validatedObject[key] = defaultValue; - } - } - } - - return validatedObject; -} - - -// translates chatgpt input to immediately loadable graph format -function transformInputToGraph(input) { - const baseGraph = { - last_node_id: input.nodes.length, - last_link_id: input.edges.length, - nodes: [], - links: [], - groups: [], - config: {}, - extra: {}, - version: 0.4, - }; +function positionNodes(nodes) { const xStep = 500; // step for node positions const yStep = 500; // step for node positions const yBase = 100; // base y position @@ -79,42 +10,23 @@ function transformInputToGraph(input) { let y = yBase; // starting y post let columnCounter = 0; // counter for columns - // Create nodes with positions in order - input.nodes.forEach((node, index) => { + + for (let i = 0; i < nodes.length; i++) { + const node = nodes[i]; // Calculate node size - const registeredNodeType = LiteGraph.registered_node_types[node.class_type]; - - // TODO: [width, height] - const nodeSize = registeredNodeType.prototype.computeSize() - - const newNode = { - id: node.id, - type: node.class_type, - pos: [x, y], - flags: {}, - order: index, - mode: 0, - inputs: [], - outputs: [], - properties: {}, - widgets_values: node.widget_values ? Object.values(node.widget_values) : [] - }; - - // Validate inputs - const inputTypes = nodesInputsOutputs[node.class_type].input_types; - // check if it has node.class_type - if (!LiteGraph.registered_node_types[node.class_type]) { - console.error(`Node type ${node.class_type} not found`); + /* + if (!LiteGraph.registered_node_types[node.type]) { + console.error(`Node type ${node.type} not found`); return; } + const registeredNodeType = LiteGraph.registered_node_types[node.class_type]; + const nodeSize = registeredNodeType.prototype.computeSize() // [width, height] + */ - // TODO: validated input/output - const widget_values = validateObject(inputTypes, node.widget_values); - - baseGraph.nodes.push(newNode); - nodesById[node.id] = newNode; + node.pos.push(x) + node.pos.push(y) if (columnCounter % 3 === 0) { x += xStep; @@ -124,61 +36,16 @@ function transformInputToGraph(input) { y += yStep; } columnCounter++; - }); - - // Insert edges - input.edges.forEach((edgeData, index) => { - const fromNode = nodesById[edgeData.from]; - const toNode = nodesById[edgeData.to]; - - // Adding output to fromNode - let outputIndex = edgeData.from_output; - // add inputs - let inputIndex = getInputIndex(toNode, edgeData.to_input); - - // if input doesnt already exist - if (inputIndex < 0) { - toNode.inputs.push({ - name: edgeData.to_input, - type: getInputType(edgeData.to_input), - link: index + 1, - }); - } - - // add outputs - fromNode.outputs.push({ - name: getInputType(edgeData.to_input), - // name: edgeData.input, - type: getInputType(edgeData.to_input), - links: [], - slot_index: fromNode.outputs.length, - }); - - - outputIndex = fromNode.outputs.length - 1; - inputIndex = fromNode.inputs.length - 1; - - // add links - const link = [ - index + 1, // Link id - fromNode.id, // Origin node id - edgeData.from_output, - // getOutputIndex(fromNode, edgeData.to_input), // origin index - toNode.id, // Destination node id - getInputIndex(toNode, edgeData.to_input), // destination index - getInputType(edgeData.to_input), - ]; - - fromNode.outputs[outputIndex].links.push(link[0]); - baseGraph.links.push(link); - }); - - - return baseGraph; + } } export function loadGraphFromPrompt(generatedWorkflow) { - const translatedData = transformInputToGraph(generatedWorkflow); + const translatedData = generatedWorkflow; + positionNodes(translatedData.nodes); app.loadGraphData(translatedData, true); app.graph.change(); } + +export function isValidWorkflow(workflow) { + return workflow && workflow.nodes && workflow.edges; +} diff --git a/js/comfy/ext.js b/js/comfy/ext.js index 2b8b1c3..08c0038 100644 --- a/js/comfy/ext.js +++ b/js/comfy/ext.js @@ -2,6 +2,10 @@ import { app } from "./comfy.js"; import { addInterface } from "../button/index.js"; import { endpoint } from "../constants.js"; import { ComfyCloud } from "../node/index.js"; +import { registerChat } from "../chatbot/ext.js"; +import { + getNodesInputsOutputs +} from '../utils.js' export const ext = { name: "nathannlu.ComfyCloud", @@ -9,6 +13,7 @@ export const ext = { init() { addInterface(); addPing(); + registerChat(); }, // Add in node that keeps track of workflow_name diff --git a/js/resource/local.js b/js/resource/local.js index 3e83596..538329c 100644 --- a/js/resource/local.js +++ b/js/resource/local.js @@ -40,6 +40,11 @@ const local = { validatePathDirectory: { method: "POST", path: "/comfy-cloud/validate-path" + }, + + sendMessage: { + method: "POST", + path: "/comfy-cloud/send-message" } } diff --git a/python/chat/classes.py b/python/chat/classes.py new file mode 100644 index 0000000..e8c7970 --- /dev/null +++ b/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 + + diff --git a/python/chat/comfy_types.py b/python/chat/comfy_types.py new file mode 100644 index 0000000..1d898bd --- /dev/null +++ b/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() diff --git a/python/chat/format.py b/python/chat/format.py new file mode 100644 index 0000000..fce6aa4 --- /dev/null +++ b/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 + + diff --git a/python/chat/graph.py b/python/chat/graph.py new file mode 100644 index 0000000..691666c --- /dev/null +++ b/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 diff --git a/python/custom_routes.py b/python/custom_routes.py index 75a5f77..0e728ab 100644 --- a/python/custom_routes.py +++ b/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)