From 25ca5fa37f0ee74fc6b21a911787687c0ad5874b Mon Sep 17 00:00:00 2001
From: pythongosssss <125205205+pythongosssss@users.noreply.github.com>
Date: Mon, 27 Mar 2023 19:05:23 +0100
Subject: [PATCH] Add widget search & replace

---
 saveimage-extra-output/README.md              |  7 ++
 .../saveImageExtraOutput.js                   | 67 +++++++++++++++++++
 2 files changed, 74 insertions(+)
 create mode 100644 saveimage-extra-output/README.md
 create mode 100644 saveimage-extra-output/saveImageExtraOutput.js

diff --git a/saveimage-extra-output/README.md b/saveimage-extra-output/README.md
new file mode 100644
index 0000000..ec7ff7d
--- /dev/null
+++ b/saveimage-extra-output/README.md
@@ -0,0 +1,7 @@
+Allows adding widget values to filenames  
+Search and replace is done in the format `%node.widgetname%`  
+where `node` is a node S&R property name or a node title  
+e.g. to output the model name: `%CheckpointLoaderSimple.ckpt_name%`  
+
+To allow for shorter names or duplicates, you can either use the change node title and use that or right click any node -> Properties -> Node name for S&R  
+And set that to a simpler thing, e.g. ckpt
\ No newline at end of file
diff --git a/saveimage-extra-output/saveImageExtraOutput.js b/saveimage-extra-output/saveImageExtraOutput.js
new file mode 100644
index 0000000..b352ff5
--- /dev/null
+++ b/saveimage-extra-output/saveImageExtraOutput.js
@@ -0,0 +1,67 @@
+import { app } from "/scripts/app.js";
+
+// Adds functionality to add widget values to SaveImage node outputs
+
+app.registerExtension({
+	name: "pysssss.SaveImageExtraOutput",
+	async beforeRegisterNodeDef(nodeType, nodeData, app) {
+		if (nodeData.name === "SaveImage") {
+			const onNodeCreated = nodeType.prototype.onNodeCreated;
+			nodeType.prototype.onNodeCreated = function () {
+				const r = onNodeCreated ? onNodeCreated.apply(this, arguments) : undefined;
+
+				const widget = this.widgets.find((w) => w.name === "filename_prefix");
+				widget.serializeValue = () => {
+					return widget.value.replace(/%([^%]+)%/g, function (match, text) {
+						const split = text.split(".");
+						if (split.length !== 2) {
+							if (text !== "width" && text !== "height") {
+								// Dont warn on standard replacements
+								console.warn("Invalid replacement pattern", text);
+							}
+							return match;
+						}
+
+						// Find node with matching S&R property name
+						let nodes = app.graph._nodes.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]);
+						}
+						if (!nodes.length) {
+							console.warn("Unable to find node", split[0]);
+							return match;
+						}
+
+						if (nodes.length > 1) {
+							console.warn("Multiple nodes matched", split[0], "using first match");
+						}
+
+						const node = nodes[0];
+
+						const widget = node.widgets?.find((w) => w.name === split[1]);
+						if (!widget) {
+							console.warn("Unable to find widget", split[1], "on node", split[0], node);
+							return match;
+						}
+
+						return ((widget.value ?? "") + "").replaceAll(/\/|\\/g, "_");
+					});
+				};
+
+				return r;
+			};
+		} else {
+			const onNodeCreated = nodeType.prototype.onNodeCreated;
+			nodeType.prototype.onNodeCreated = function () {
+				const r = onNodeCreated ? onNodeCreated.apply(this, arguments) : undefined;
+
+				if (!Object.keys(this.properties || {}).length) {
+					this.addProperty("Node name for S&R", this.title, "string");
+				}
+
+				return r;
+			};
+		}
+	},
+});