Try to make core extensions subgraph-aware

This commit is contained in:
space-nuko 2023-06-06 13:20:05 -05:00
parent 32ae383ad5
commit 3e3a7fd9ce
7 changed files with 53 additions and 31 deletions

View File

@ -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() {

View File

@ -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]);

View File

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

View File

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

View File

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

View File

@ -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()
}
}}),

View File

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