From 091434dc5474070532646e1f645aedcb660cfc43 Mon Sep 17 00:00:00 2001
From: pythongosssss <125205205+pythongosssss@users.noreply.github.com>
Date: Wed, 16 Aug 2023 20:57:16 +0100
Subject: [PATCH] add repeater

---
 js/repeater.js | 123 +++++++++++++++++++++++++++++++++++++++++++++++++
 py/repeater.py |  46 ++++++++++++++++++
 2 files changed, 169 insertions(+)
 create mode 100644 js/repeater.js
 create mode 100644 py/repeater.py

diff --git a/js/repeater.js b/js/repeater.js
new file mode 100644
index 0000000..23432c4
--- /dev/null
+++ b/js/repeater.js
@@ -0,0 +1,123 @@
+import { app } from "../../../scripts/app.js";
+
+const REPEATER = "Repeater|pysssss";
+
+app.registerExtension({
+	name: "pysssss.Repeater",
+	init() {
+		const graphToPrompt = app.graphToPrompt;
+		app.graphToPrompt = async function () {
+			const res = await graphToPrompt.apply(this, arguments);
+
+			const id = Date.now() + "_";
+			let u = 0;
+
+			let newNodes = {};
+			const newRepeaters = {};
+			for (const nodeId in res.output) {
+				let output = res.output[nodeId];
+				if (output.class_type === REPEATER) {
+					const isMulti = output.inputs.output === "multi";
+					if (output.inputs.node_mode === "create") {
+						// We need to clone the input for every repeat
+						const orig = res.output[output.inputs.source[0]];
+						if (isMulti) {
+							if (!newRepeaters[nodeId]) {
+								newRepeaters[nodeId] = [];
+								newRepeaters[nodeId][output.inputs.repeats - 1] = nodeId;
+							}
+						}
+						for (let i = 0; i < output.inputs.repeats - 1; i++) {
+							const clonedInputId = id + ++u;
+
+							if (isMulti) {
+								// If multi create we need to clone the repeater too
+								newNodes[clonedInputId] = structuredClone(orig);
+
+								output = structuredClone(output);
+
+								const clonedRepeaterId = id + ++u;
+								newNodes[clonedRepeaterId] = output;
+								output.inputs["source"][0] = clonedInputId;
+
+								newRepeaters[nodeId][i] = clonedRepeaterId;
+							} else {
+								newNodes[clonedInputId] = orig;
+							}
+							output.inputs[clonedInputId] = [clonedInputId, output.inputs.source[1]];
+						}
+					} else if (isMulti) {
+						newRepeaters[nodeId] = Array(output.inputs.repeats).fill(nodeId);
+					}
+				}
+			}
+
+			Object.assign(res.output, newNodes);
+			newNodes = {};
+
+			for (const nodeId in res.output) {
+				const output = res.output[nodeId];
+				for (const k in output.inputs) {
+					const v = output.inputs[k];
+					if (v instanceof Array) {
+						const repeaterId = v[0];
+						const source = newRepeaters[repeaterId];
+						if (source) {
+							v[0] = source.pop();
+							v[1] = 0;
+						}
+					}
+				}
+			}
+
+			// Object.assign(res.output, newNodes);
+
+			return res;
+		};
+	},
+	beforeRegisterNodeDef(nodeType, nodeData, app) {
+		if (nodeData.name === REPEATER) {
+			const SETUP_OUTPUTS = Symbol();
+			nodeType.prototype[SETUP_OUTPUTS] = function (repeats) {
+				if (repeats == null) {
+					repeats = this.widgets[0].value;
+				}
+				while (this.outputs.length > repeats) {
+					this.removeOutput(repeats);
+				}
+				const id = Date.now() + "_";
+				let u = 0;
+				while (this.outputs.length < repeats) {
+					this.addOutput(id + ++u, "*", { label: "*" });
+				}
+			};
+
+			const onAdded = nodeType.prototype.onAdded;
+			nodeType.prototype.onAdded = function () {
+				const self = this;
+				const repeatsCb = this.widgets[0].callback;
+				this.widgets[0].callback = async function () {
+					const v = (await repeatsCb?.apply(this, arguments)) ?? this.value;
+					if (self.widgets[1].value === "multi") {
+						self[SETUP_OUTPUTS](v);
+					}
+					return v;
+				};
+
+				const outputCb = this.widgets[1].callback;
+				this.widgets[1].callback = async function () {
+					const v = (await outputCb?.apply(this, arguments)) ?? this.value;
+					if (v === "single") {
+						self.outputs[0].shape = 6;
+						self[SETUP_OUTPUTS](1);
+					} else {
+						delete self.outputs[0].shape;
+						self[SETUP_OUTPUTS]();
+					}
+					return v;
+				};
+				return onAdded?.apply(this, arguments);
+			};
+		}
+	},
+});
diff --git a/py/repeater.py b/py/repeater.py
new file mode 100644
index 0000000..1d83c67
--- /dev/null
+++ b/py/repeater.py
@@ -0,0 +1,46 @@
+# Hack: string type that is always equal in not equal comparisons
+class AnyType(str):
+    def __ne__(self, __value: object) -> bool:
+        return False
+
+
+# Our any instance wants to be a wildcard string
+any = AnyType("*")
+
+
+class Repeater:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required": {
+            "source": (any, {}),
+            "repeats": ("INT", {"min": 0, "max": 5000, "default": 2}),
+            "output": (["single", "multi"], {}),
+            "node_mode": (["reuse", "create"], {}),
+        }}
+
+    RETURN_TYPES = (any,)
+    FUNCTION = "repeat"
+    OUTPUT_NODE = False
+    OUTPUT_IS_LIST = (True,)
+
+    CATEGORY = "utils"
+
+    def repeat(self, repeats, output, node_mode, **kwargs):
+        if output == "multi":
+            # Multi outputs are split to indiviual nodes on the frontend when serializing
+            return ([kwargs["source"]],)
+        elif node_mode == "reuse":
+            # When reusing we have a single input node, repeat that N times
+            return ([kwargs["source"]] * repeats,)
+        else:
+            # When creating new nodes, they'll be added dynamically when the graph is serialized
+            return ((list(kwargs.values())),)
+
+
+NODE_CLASS_MAPPINGS = {
+    "Repeater|pysssss": Repeater,
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "Repeater|pysssss": "Repeater 🐍",
+}