Try to make core extensions subgraph-aware

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 link = this.graph.links[linkId];
-								const node = app.graph.getNodeById(link.target_id);
+								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 link = this.graph.links[linkId];
-							const node = app.graph.getNodeById(link.origin_id);
+							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() {

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

web/extensions/core/widgetInputs.js

@@ -143,8 +143,8 @@ app.registerExtension({
 			return r;
 		};
 
-		function isNodeAtPos(pos) {
+		function isNodeAtPos(graph, pos) {
-			for (const n of app.graph._nodes) {
+			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;

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

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) {
+		for(let node of this.graph.iterateNodesRecursive()) {
-			const node = this.graph._nodes[nodeNum];
-
 			const def = defs[node.type];
 
 			// HOTFIX: The current patch is designed to prevent the rest of the code from breaking due to primitive nodes,

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

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) {
+		for (const n of graphcanvas.graph.iterateNodes()) {
-			n = graph._nodes[n];
 			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)