Add bidirectional autoscaling index switch

2026-06-25 08:33:35 +02:00
parent a60258dc4c
commit 55477bd826
3 changed files with 294 additions and 0 deletions
@@ -0,0 +1,147 @@
+import { app } from "../../scripts/app.js";
+
+const EXTENSION = "ethanfel.prompt_builder.index_switch_slots";
+const NODE_NAME = "SxCPIndexSwitch";
+const MAX_INPUTS = 64;
+
+function isSwitchInput(input) {
+    return /^input_\d+$/.test(input?.name || "");
+}
+
+function isSwitchOutput(output) {
+    return /^output_\d+$/.test(output?.name || "");
+}
+
+function slotNumber(slot) {
+    const match = String(slot?.name || "").match(/\d+$/);
+    return match ? Number(match[0]) : -1;
+}
+
+function resizeNode(node) {
+    const size = node.computeSize?.();
+    if (size) node.setSize?.(size);
+    app.graph?.setDirtyCanvas(true, true);
+}
+
+function addSwitchInput(node, number) {
+    if (number < 1 || number > MAX_INPUTS) return;
+    const name = `input_${number}`;
+    if (!node.inputs?.some((input) => input.name === name)) {
+        node.addInput(name, "*");
+    }
+}
+
+function removeSwitchInput(node, number) {
+    const inputIndex = node.inputs?.findIndex((input) => input.name === `input_${number}`) ?? -1;
+    if (inputIndex >= 0 && !node.inputs[inputIndex]?.link) {
+        node.removeInput(inputIndex);
+    }
+}
+
+function addSwitchOutput(node, number) {
+    if (number < 1 || number > MAX_INPUTS) return;
+    const name = `output_${number}`;
+    if (!node.outputs?.some((output) => output.name === name)) {
+        node.addOutput(name, "*");
+    }
+}
+
+function removeSwitchOutput(node, number) {
+    const outputIndex = node.outputs?.findIndex((output) => output.name === `output_${number}`) ?? -1;
+    if (outputIndex >= 0 && !(node.outputs[outputIndex]?.links?.length)) {
+        node.removeOutput(outputIndex);
+    }
+}
+
+function trimInputTail(node) {
+    for (let number = MAX_INPUTS; number > 1; number--) {
+        const input = node.inputs?.find((slot) => slot.name === `input_${number}`);
+        const previous = node.inputs?.find((slot) => slot.name === `input_${number - 1}`);
+        if (!input?.link && !previous?.link) removeSwitchInput(node, number);
+    }
+}
+
+function trimOutputTail(node) {
+    for (let number = MAX_INPUTS; number > 1; number--) {
+        const output = node.outputs?.find((slot) => slot.name === `output_${number}`);
+        const previous = node.outputs?.find((slot) => slot.name === `output_${number - 1}`);
+        if (!(output?.links?.length) && !(previous?.links?.length)) removeSwitchOutput(node, number);
+    }
+}
+
+function setupNodeSlots(node) {
+    addSwitchInput(node, 1);
+    addSwitchOutput(node, 1);
+    for (let number = 2; number <= MAX_INPUTS; number++) {
+        const input = node.inputs?.find((slot) => slot.name === `input_${number}`);
+        if (!input?.link) removeSwitchInput(node, number);
+        const output = node.outputs?.find((slot) => slot.name === `output_${number}`);
+        if (!(output?.links?.length)) removeSwitchOutput(node, number);
+    }
+    trimInputTail(node);
+    trimOutputTail(node);
+    resizeNode(node);
+}
+
+function maybeGrowInput(node) {
+    const switchInputs = (node.inputs || []).filter(isSwitchInput);
+    const last = switchInputs.reduce((max, input) => Math.max(max, slotNumber(input)), 1);
+    const lastInput = node.inputs?.find((slot) => slot.name === `input_${last}`);
+    if (lastInput?.link && last < MAX_INPUTS) {
+        addSwitchInput(node, last + 1);
+        resizeNode(node);
+    }
+}
+
+function maybeGrowOutput(node) {
+    const switchOutputs = (node.outputs || []).filter(isSwitchOutput);
+    const last = switchOutputs.reduce((max, output) => Math.max(max, slotNumber(output)), 1);
+    const lastOutput = node.outputs?.find((slot) => slot.name === `output_${last}`);
+    if (lastOutput?.links?.length && last < MAX_INPUTS) {
+        addSwitchOutput(node, last + 1);
+        resizeNode(node);
+    }
+}
+
+app.registerExtension({
+    name: EXTENSION,
+
+    async beforeRegisterNodeDef(nodeType, nodeData) {
+        if (nodeData.name !== NODE_NAME) return;
+
+        const onNodeCreated = nodeType.prototype.onNodeCreated;
+        nodeType.prototype.onNodeCreated = function () {
+            const result = onNodeCreated?.apply(this, arguments);
+            queueMicrotask(() => setupNodeSlots(this));
+            return result;
+        };
+
+        const onConfigure = nodeType.prototype.onConfigure;
+        nodeType.prototype.onConfigure = function () {
+            const result = onConfigure?.apply(this, arguments);
+            queueMicrotask(() => setupNodeSlots(this));
+            return result;
+        };
+
+        const onConnectionsChange = nodeType.prototype.onConnectionsChange;
+        nodeType.prototype.onConnectionsChange = function (type, index, connected, linkInfo) {
+            const result = onConnectionsChange?.apply(this, arguments);
+            if (!linkInfo) return result;
+            const slot = type === LiteGraph.INPUT ? this.inputs?.[index] : this.outputs?.[index];
+            if (type === LiteGraph.INPUT && isSwitchInput(slot)) {
+                if (connected) maybeGrowInput(this);
+                else {
+                    trimInputTail(this);
+                    resizeNode(this);
+                }
+            } else if (isSwitchOutput(slot)) {
+                if (connected) maybeGrowOutput(this);
+                else {
+                    trimOutputTail(this);
+                    resizeNode(this);
+                }
+            }
+            return result;
+        };
+    },
+});