From 3e3a7fd9ce60355ade17467cdf96417a4df3ed2a Mon Sep 17 00:00:00 2001 From: space-nuko <24979496+space-nuko@users.noreply.github.com> Date: Tue, 6 Jun 2023 13:20:05 -0500 Subject: [PATCH] Try to make core extensions subgraph-aware --- web/extensions/core/rerouteNode.js | 22 +++++++++---------- web/extensions/core/saveImageExtraOutput.js | 4 ++-- web/extensions/core/widgetInputs.js | 13 +++++------ web/lib/litegraph.core.js | 24 +++++++++++++++++++++ web/scripts/app.js | 14 ++++++------ web/scripts/ui.js | 2 ++ web/scripts/widgets.js | 5 +---- 7 files changed, 53 insertions(+), 31 deletions(-) diff --git a/web/extensions/core/rerouteNode.js b/web/extensions/core/rerouteNode.js index c31f63cd0..51afb2aab 100644 --- a/web/extensions/core/rerouteNode.js +++ b/web/extensions/core/rerouteNode.js @@ -22,12 +22,12 @@ app.registerExtension({ // Prevent multiple connections to different types when we have no input if (connected && type === LiteGraph.OUTPUT) { // Ignore wildcard nodes as these will be updated to real types - const types = new Set(this.outputs[0].links.map((l) => app.graph.links[l].type).filter((t) => t !== "*")); + const types = new Set(this.outputs[0].links.map((l) => this.graph.links[l].type).filter((t) => t !== "*")); if (types.size > 1) { for (let i = 0; i < this.outputs[0].links.length - 1; i++) { const linkId = this.outputs[0].links[i]; - const link = app.graph.links[linkId]; - const node = app.graph.getNodeById(link.target_id); + const link = this.graph.links[linkId]; + const node = this.graph.getNodeById(link.target_id); node.disconnectInput(link.target_slot); } } @@ -42,8 +42,8 @@ app.registerExtension({ updateNodes.unshift(currentNode); const linkId = currentNode.inputs[0].link; if (linkId !== null) { - const link = app.graph.links[linkId]; - const node = app.graph.getNodeById(link.origin_id); + const link = this.graph.links[linkId]; + const node = this.graph.getNodeById(link.origin_id); const type = node.constructor.type; if (type === "Reroute") { if (node === this) { @@ -76,12 +76,12 @@ app.registerExtension({ const outputs = (currentNode.outputs ? currentNode.outputs[0].links : []) || []; if (outputs.length) { for (const linkId of outputs) { - const link = app.graph.links[linkId]; + const link = this.graph.links[linkId]; // When disconnecting sometimes the link is still registered if (!link) continue; - const node = app.graph.getNodeById(link.target_id); + const node = this.graph.getNodeById(link.target_id); const type = node.constructor.type; if (type === "Reroute") { @@ -118,7 +118,7 @@ app.registerExtension({ node.applyOrientation(); for (const l of node.outputs[0].links || []) { - const link = app.graph.links[l]; + const link = this.graph.links[l]; if (link) { link.color = color; } @@ -126,7 +126,7 @@ app.registerExtension({ } if (inputNode) { - const link = app.graph.links[inputNode.inputs[0].link]; + const link = this.graph.links[inputNode.inputs[0].link]; if (link) { link.color = color; } @@ -158,7 +158,7 @@ app.registerExtension({ } this.size = this.computeSize(); this.applyOrientation(); - app.graph.setDirtyCanvas(true, true); + this.graph.setDirtyCanvas(true, true); }, }, { @@ -190,7 +190,7 @@ app.registerExtension({ } else { delete this.inputs[0].pos; } - app.graph.setDirtyCanvas(true, true); + this.graph.setDirtyCanvas(true, true); } computeSize() { diff --git a/web/extensions/core/saveImageExtraOutput.js b/web/extensions/core/saveImageExtraOutput.js index 6032d4cc7..c80c879bd 100644 --- a/web/extensions/core/saveImageExtraOutput.js +++ b/web/extensions/core/saveImageExtraOutput.js @@ -55,10 +55,10 @@ app.registerExtension({ } // Find node with matching S&R property name - let nodes = app.graph._nodes.filter((n) => n.properties?.["Node name for S&R"] === split[0]); + let nodes = Array.from(app.graph.iterateNodesRecursive()).filter((n) => n.properties?.["Node name for S&R"] === split[0]); // If we cant, see if there is a node with that title if (!nodes.length) { - nodes = app.graph._nodes.filter((n) => n.title === split[0]); + nodes = Array.from(app.graph.iterateNodesRecursive()).filter((n) => n.title === split[0]); } if (!nodes.length) { console.warn("Unable to find node", split[0]); diff --git a/web/extensions/core/widgetInputs.js b/web/extensions/core/widgetInputs.js index c356655b0..ce1267cef 100644 --- a/web/extensions/core/widgetInputs.js +++ b/web/extensions/core/widgetInputs.js @@ -143,8 +143,8 @@ app.registerExtension({ return r; }; - function isNodeAtPos(pos) { - for (const n of app.graph._nodes) { + function isNodeAtPos(graph, pos) { + for (const n of graph.iterateNodes()) { if (n.pos[0] === pos[0] && n.pos[1] === pos[1]) { return true; } @@ -168,11 +168,12 @@ app.registerExtension({ // Create a primitive node const node = LiteGraph.createNode("PrimitiveNode"); - app.graph.add(node); + const graph = LGraphCanvas.active_canvas.graph; + graph.add(node); // Calculate a position that wont directly overlap another node const pos = [this.pos[0] - node.size[0] - 30, this.pos[1]]; - while (isNodeAtPos(pos)) { + while (isNodeAtPos(graph, pos)) { pos[1] += LiteGraph.NODE_TITLE_HEIGHT; } @@ -203,7 +204,7 @@ app.registerExtension({ function get_links(node) { let links = []; for (const l of node.outputs[0].links) { - const linkInfo = app.graph.links[l]; + const linkInfo = node.graph.links[l]; const n = node.graph.getNodeById(linkInfo.target_id); if (n.type == "Reroute") { links = links.concat(get_links(n)); @@ -217,7 +218,7 @@ app.registerExtension({ let links = get_links(this); // For each output link copy our value over the original widget value for (const l of links) { - const linkInfo = app.graph.links[l]; + const linkInfo = this.graph.links[l]; const node = this.graph.getNodeById(linkInfo.target_id); const input = node.inputs[linkInfo.target_slot]; const widgetName = input.widget.name; diff --git a/web/lib/litegraph.core.js b/web/lib/litegraph.core.js index 10da9d5aa..f95a95114 100644 --- a/web/lib/litegraph.core.js +++ b/web/lib/litegraph.core.js @@ -1401,6 +1401,30 @@ return this.elapsed_time; }; + + /** + * Iterates all nodes in this graph *excluding* subgraphs. + */ + LGraph.prototype.iterateNodes = function() { + const nodes = this._nodes_in_order ? this._nodes_in_order : this._nodes || []; + return nodes; + } + + /** + * Iterates all nodes in this graph and subgraphs. + */ + LGraph.prototype.iterateNodesRecursive = function() { + const nodes_ = this._nodes_in_order ? this._nodes_in_order : this._nodes || []; + let nodes = [] + for (const node of nodes_) { + nodes.push(node) + if (node.subgraph != null) { + nodes = nodes.concat(node.subgraph.iterateNodesRecursive()); + } + } + return nodes; + } + /** * Sends an event to all the nodes, useful to trigger stuff * @method sendEventToAllNodes diff --git a/web/scripts/app.js b/web/scripts/app.js index 460748b0e..6a121bc53 100644 --- a/web/scripts/app.js +++ b/web/scripts/app.js @@ -190,7 +190,7 @@ export class ComfyApp { } } - app.graph.setDirtyCanvas(true); + node.graph.setDirtyCanvas(true); } } @@ -437,7 +437,7 @@ export class ComfyApp { if ((!output || this.images === output.images) && (!preview || this.preview === preview)) { this.imgs = imgs.filter(Boolean); this.setSizeForImage?.(); - app.graph.setDirtyCanvas(true); + node.graph?.setDirtyCanvas(true); } }); } @@ -940,7 +940,7 @@ export class ComfyApp { this.canvas.onGraphAttached = function(graph) { console.warn("canvas ongraphattached") - for (const node of graph._nodes) { + for (const node of graph.iterateNodes()) { if (node.onGraphAttached) node.onGraphAttached() } @@ -948,7 +948,7 @@ export class ComfyApp { this.canvas.onGraphDetached = function(graph) { console.warn("canvas ongraphdetached") - for (const node of graph._nodes) { + for (const node of graph.iterateNodes()) { if (node.onGraphDetached) node.onGraphDetached() } @@ -1289,7 +1289,7 @@ export class ComfyApp { return; } - for (const node of this.graph._nodes) { + for (const node of this.graph.iterateNodesRecursive()) { const size = node.computeSize(); size[0] = Math.max(node.size[0], size[0]); size[1] = Math.max(node.size[1], size[1]); @@ -1573,9 +1573,7 @@ export class ComfyApp { async refreshComboInNodes() { const defs = await api.getNodeDefs(); - for(let nodeNum in this.graph._nodes) { - const node = this.graph._nodes[nodeNum]; - + for(let node of this.graph.iterateNodesRecursive()) { const def = defs[node.type]; // HOTFIX: The current patch is designed to prevent the rest of the code from breaking due to primitive nodes, diff --git a/web/scripts/ui.js b/web/scripts/ui.js index a26eedec3..5919a25a1 100644 --- a/web/scripts/ui.js +++ b/web/scripts/ui.js @@ -602,12 +602,14 @@ export class ComfyUI { $el("button", { id: "comfy-clipspace-button", textContent: "Clipspace", onclick: () => app.openClipspace() }), $el("button", { id: "comfy-clear-button", textContent: "Clear", onclick: () => { if (!confirmClear.value || confirm("Clear workflow?")) { + LGraphCanvas.active_canvas.closeAllSubgraphs(); app.clean(); app.graph.clear(); } }}), $el("button", { id: "comfy-load-default-button", textContent: "Load Default", onclick: () => { if (!confirmClear.value || confirm("Load default workflow?")) { + LGraphCanvas.active_canvas.closeAllSubgraphs(); app.loadGraphData() } }}), diff --git a/web/scripts/widgets.js b/web/scripts/widgets.js index 24017f2c0..8d9597d15 100644 --- a/web/scripts/widgets.js +++ b/web/scripts/widgets.js @@ -200,8 +200,7 @@ function addMultilineWidget(node, name, opts, app) { if (graphcanvas == null || graphcanvas.graph != node.graph) return - for (let n in graphcanvas.graph._nodes) { - n = graph._nodes[n]; + for (const n of graphcanvas.graph.iterateNodes()) { for (let w in n.widgets) { let wid = n.widgets[w]; if (Object.hasOwn(wid, "inputEl")) { @@ -223,7 +222,6 @@ function addMultilineWidget(node, name, opts, app) { const onGraphAttached = node.onGraphAttached; node.onGraphAttached = function() { - console.error("ONGRAPHATTACHKED", widget) widget.inputEl.style.display = "block"; if (onGraphAttached) onGraphAttached.apply(this, arguments) @@ -231,7 +229,6 @@ function addMultilineWidget(node, name, opts, app) { const onGraphDetached = node.onGraphDetached; node.onGraphDetached = function() { - console.error("ONGRAPHDETACHED", widget) widget.inputEl.style.display = "none"; if (onGraphDetached) onGraphDetached.apply(this, arguments)