From 806520572b592c9768d3b018d8dc792e12b084b2 Mon Sep 17 00:00:00 2001 From: xiaoshuisheng Date: Mon, 11 Dec 2023 11:02:12 +0800 Subject: [PATCH] [TEST|ADD] add flow-control inputs and outputs --- web/lib/litegraph.core.js | 5 + web/scripts/api.js | 34 +++++- web/scripts/app.js | 249 +++++++++++++++++++++++++++++++++++++- 3 files changed, 281 insertions(+), 7 deletions(-) diff --git a/web/lib/litegraph.core.js b/web/lib/litegraph.core.js index f571edb30..fc5340f89 100644 --- a/web/lib/litegraph.core.js +++ b/web/lib/litegraph.core.js @@ -227,6 +227,8 @@ } } + console.log(type); + console.log(base_class); this.registered_node_types[type] = base_class; if (base_class.constructor.name) { this.Nodes[classname] = base_class; @@ -4327,6 +4329,9 @@ target_node.id, target_slot ); + console.log("new link"); + console.trace(); + console.log(link_info); //add to graph links list this.graph.links[link_info.id] = link_info; diff --git a/web/scripts/api.js b/web/scripts/api.js index 9aa7528af..da08f5473 100644 --- a/web/scripts/api.js +++ b/web/scripts/api.js @@ -12,6 +12,9 @@ class ComfyApi extends EventTarget { } fetchApi(route, options) { + console.log("FetchAPI"); + console.log(route); + console.log(options); return fetch(this.apiURL(route), options); } @@ -26,6 +29,7 @@ class ComfyApi extends EventTarget { #pollQueue() { setInterval(async () => { try { + console.log("#pollQueue"); const resp = await this.fetchApi("/prompt"); const status = await resp.json(); this.dispatchEvent(new CustomEvent("status", { detail: status })); @@ -178,7 +182,31 @@ class ComfyApi extends EventTarget { */ async getNodeDefs() { const resp = await this.fetchApi("/object_info", { cache: "no-store" }); - return await resp.json(); + let node_defs = await resp.json(); + console.log(node_defs); + try{ + for (let node_name in node_defs) + { + let node_def = node_defs[node_name]; + // add input + if(!("optional" in node_def.input)) + { + node_def.input["optional"] = {}; + } + node_def.input["optional"]["FROM"] = ['FLOW']; + + // add output + node_def.output.unshift("FLOW"); + node_def.output_is_list.unshift(false); + node_def.output_name.unshift("TO"); + } + } catch(error) + { + console.log("err happends"); + console.log(error); + } + console.log(node_defs); + return node_defs; } /** @@ -186,7 +214,7 @@ class ComfyApi extends EventTarget { * @param {number} number The index at which to queue the prompt, passing -1 will insert the prompt at the front of the queue * @param {object} prompt The prompt data to queue */ - async queuePrompt(number, { output, workflow }) { + async queuePrompt(number, { output, workflow, flows }) { const body = { client_id: this.clientId, prompt: output, @@ -199,6 +227,8 @@ class ComfyApi extends EventTarget { body.number = number; } + console.log("queuePrompt"); + console.log(body); const res = await this.fetchApi("/prompt", { method: "POST", headers: { diff --git a/web/scripts/app.js b/web/scripts/app.js index 5faf41fb3..eaf48f62b 100644 --- a/web/scripts/app.js +++ b/web/scripts/app.js @@ -1396,6 +1396,8 @@ export class ComfyApp { } async registerNodeDef(nodeId, nodeData) { + console.log("registerNodeDef"); + console.log(nodeData); const self = this; const node = Object.assign( function ComfyNode() { @@ -1418,7 +1420,8 @@ export class ComfyApp { } } else { // Node connection inputs - this.addInput(inputName, type); + const inputShape = type == "FLOW"? LiteGraph.ARROW_SHAPE : LiteGraph.CIRCLE_SHAPE; + this.addInput(inputName, type, { shape: inputShape }); widgetCreated = false; } @@ -1436,7 +1439,8 @@ export class ComfyApp { let output = nodeData["output"][o]; if(output instanceof Array) output = "COMBO"; const outputName = nodeData["output_name"][o] || output; - const outputShape = nodeData["output_is_list"][o] ? LiteGraph.GRID_SHAPE : LiteGraph.CIRCLE_SHAPE ; + // const outputShape = nodeData["output_is_list"][o] ? LiteGraph.GRID_SHAPE : LiteGraph.CIRCLE_SHAPE; + const outputShape = output == "FLOW"? LiteGraph.ARROW_SHAPE : (nodeData["output_is_list"][o] ? LiteGraph.GRID_SHAPE : LiteGraph.CIRCLE_SHAPE) ; this.addOutput(outputName, output, { shape: outputShape }); } @@ -1556,6 +1560,205 @@ export class ComfyApp { }); } + + calculateFlowConnection(graphData) + { + if("flows" in graphData) + { + return graphData.flows; + } + + // id-nodes, id-links + var nodes = {}; + var links = {}; + for(const node of graphData.nodes) + { + nodes[node.id] = node; + } + for(const link of graphData.links) + { + links[link[0]] = link; + } + + // in-degree info + var in_degree = {}; + var flow_order = []; + for (const cur_node of graphData.nodes) + { + // in-degree info + var degree = 0; + // current node has inputs + if('inputs' in cur_node) + { + for (const inp of cur_node.inputs) + { + // only connected input + if(inp.link != null) + { + ++degree; + } + } + } + + // update in-degree info + in_degree[cur_node.id] = degree; + + if(degree == 0) + { + flow_order.push(cur_node); + } + } + + + // calculate flow connection + var idx = 0; + while(idx < flow_order.length) + { + let cur_node = flow_order[idx++]; + if('outputs' in cur_node) + { + for(const output of cur_node.outputs) + { + if(output.links === null || output.links.length == 0) + { + continue; + } + + // + for(const link_id of output.links) + { + let link = links[link_id]; + if(cur_node.id != link[3]) + { + --in_degree[link[3]]; + if(in_degree[link[3]] == 0) + { + flow_order.push(nodes[link[3]]); + } + } + } + } + } + } + + // update flows data + var flows = {}; + idx = 0; + while(idx < flow_order.length) + { + if(idx == flow_order.length -1) + { + flows[flow_order[idx].id] = null; + } + else{ + flows[flow_order[idx].id] = flow_order[idx+1].id; + } + ++idx; + } + graphData["flows"] = flows; + + return flows; + } + + + initFlowControlConnection(graphData) + { + this.calculateFlowConnection(graphData); + + // No nodes exist, just return + if(graphData.nodes.length==0) + { + return graphData; + } + + // If flow-control inputs/outputs exist, just return + for (const node of graphData.nodes) + { + if("inputs" in node) + { + if (node.inputs[0].name == "FROM") + { + return graphData; + } + else{ + break; + } + } + } + + // add flow inputs & outputs + for(const cur_node of graphData.nodes) + { + if(!('inputs' in cur_node)) + { + cur_node.inputs = []; + } + cur_node.inputs.unshift({name:'FROM', type:"FLOW", link:null, shape: LiteGraph.ARROW_SHAPE}); + if(!('outputs' in cur_node)) + { + cur_node.outputs = []; + } + for(const output of cur_node.outputs) + { + output.slot_index +=1; + } + cur_node.outputs.unshift({name:"TO", type: "FLOW", slot_index: 0, links:null, shape: LiteGraph.ARROW_SHAPE}); + + } + + // update link info + for (const link of graphData.links) + { + link[2] +=1; + link[4] +=1; + } + + // max link id + var max_link_id = 0; + for(const link of graphData.links) + { + if(max_link_id < link[0]) + {max_link_id = link[0];} + } + + // id-nodes, id-links + var nodes = {}; + for(const node of graphData.nodes) + { + nodes[node.id] = node; + } + + // add links & flows + for(let from_id in graphData.flows) + { + let to_id = graphData.flows[from_id]; + if(to_id == null) + continue; + + let link_id = ++max_link_id; + let from_node = nodes[from_id]; + let to_node = nodes[to_id]; + + var link = [link_id, from_id, 0, to_id, 0, "FLOW"]; + // {id: link_id, + // origin_id: from_id, + // origin_slot: from_node.outputs.length -1, + // target_id: to_id, + // target_slot: to_node.inputs.length -1, + // type: "FLOW" + // // _data: null + // }; + console.log("Add Link:"); + console.log(link); + graphData.links.push(link); + + from_node.outputs[0].links = [link_id]; + to_node.inputs[0].link = link_id; + } + + return graphData; + } + /** * Populates the graph with the specified workflow data * @param {*} graphData A serialized graph object @@ -1593,6 +1796,10 @@ export class ComfyApp { } try { + console.trace(); + graphData = this.initFlowControlConnection(graphData); + console.log("Graph Data"); + console.log(graphData); this.graph.configure(graphData); } catch (error) { let errorHint = []; @@ -1700,7 +1907,9 @@ export class ComfyApp { } } - const workflow = this.graph.serialize(); + const _workflow = this.graph.serialize(); + console.log("GraphToPrompt, originalPrompt"); + console.log(_workflow); const output = {}; // Process nodes in order of execution for (const outerNode of this.graph.computeExecutionOrder(false)) { @@ -1771,7 +1980,7 @@ export class ComfyApp { if (parent?.updateLink) { link = parent.updateLink(link); } - inputs[node.inputs[i].name] = [String(link.origin_id), parseInt(link.origin_slot)]; + inputs[node.inputs[i].name] = [String(link.origin_id), parseInt(link.origin_slot) - 1]; } } } @@ -1795,7 +2004,37 @@ export class ComfyApp { } } - return { workflow, output }; + // Get flows from workflow + let workflow = JSON.parse(JSON.stringify(_workflow)); + var flows = {}; + for (const link of workflow.links) + { + flows[link[1]] = link[3]; + } + workflow["flows"] = flows; + + // Remove flow control info from workflow + // link + for (const link of workflow.links) + { + --link[2]; + --link[4]; + } + workflow.links = workflow.links.filter(element=>element[5]!="FLOW"); + // input & output + for (const node of workflow.nodes) + { + node.inputs.shift(); + node.outputs.shift(); + } + + // Remove flow control info from output + for (let node_id in output) + { + delete output[node_id].inputs.FROM; + } + + return { workflow, output, flows }; } #formatPromptError(error) {