Spaces:
Running
on
Zero
Running
on
Zero
File size: 3,554 Bytes
3d5837a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 |
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);
};
}
},
});
|