141 lines
3.9 KiB
JavaScript
141 lines
3.9 KiB
JavaScript
import { app } from "../../scripts/app.js";
|
|
|
|
const EXTENSION = "ethanfel.prompt_builder.seed_control";
|
|
const NODE_NAME = "SxCPSeedControl";
|
|
const SEED_AXES = [
|
|
"category",
|
|
"subcategory",
|
|
"content",
|
|
"person",
|
|
"scene",
|
|
"pose",
|
|
"role",
|
|
"expression",
|
|
"composition",
|
|
];
|
|
|
|
function widget(node, name) {
|
|
return node.widgets?.find((w) => w.name === name);
|
|
}
|
|
|
|
function isSeedControlNode(node) {
|
|
return node?._sxcpSeedControlNode || node?.comfyClass === NODE_NAME || node?.type === NODE_NAME;
|
|
}
|
|
|
|
function resizeNode(node) {
|
|
const size = node.computeSize?.();
|
|
if (size) node.setSize?.(size);
|
|
app.graph?.setDirtyCanvas(true, true);
|
|
}
|
|
|
|
function randomSeed() {
|
|
if (globalThis.crypto?.getRandomValues) {
|
|
const values = new Uint32Array(1);
|
|
globalThis.crypto.getRandomValues(values);
|
|
return values[0];
|
|
}
|
|
return Math.floor(Math.random() * 0x100000000);
|
|
}
|
|
|
|
function setWidgetValue(node, w, value) {
|
|
if (!w) return;
|
|
const oldValue = w.value;
|
|
w.value = value;
|
|
w.callback?.(value, app.canvas, node);
|
|
node.onWidgetChanged?.(w.name, value, oldValue, w);
|
|
}
|
|
|
|
function materializeNode(node) {
|
|
const changes = [];
|
|
for (const axis of SEED_AXES) {
|
|
const modeWidget = widget(node, `${axis}_seed_mode`);
|
|
const seedWidget = widget(node, `${axis}_seed`);
|
|
if (modeWidget?.value !== "random" || !seedWidget) continue;
|
|
|
|
const seed = randomSeed();
|
|
changes.push({
|
|
node,
|
|
modeWidget,
|
|
previousMode: modeWidget.value,
|
|
seedWidget,
|
|
previousSeed: seedWidget.value,
|
|
seed,
|
|
});
|
|
setWidgetValue(node, seedWidget, seed);
|
|
setWidgetValue(node, modeWidget, "fixed");
|
|
}
|
|
if (changes.length) resizeNode(node);
|
|
return changes;
|
|
}
|
|
|
|
function materializeAllForQueue() {
|
|
const changes = [];
|
|
for (const node of app.graph?._nodes || []) {
|
|
if (isSeedControlNode(node)) changes.push(...materializeNode(node));
|
|
}
|
|
if (changes.length) app.graph?.setDirtyCanvas(true, true);
|
|
return changes;
|
|
}
|
|
|
|
function restoreRandomModes(changes) {
|
|
for (const change of changes) {
|
|
setWidgetValue(change.node, change.modeWidget, change.previousMode);
|
|
}
|
|
if (changes.length) app.graph?.setDirtyCanvas(true, true);
|
|
}
|
|
|
|
function lockRandomSeeds(node) {
|
|
const changes = materializeNode(node);
|
|
if (!changes.length) {
|
|
alert("No random seed modes to lock.");
|
|
return;
|
|
}
|
|
resizeNode(node);
|
|
}
|
|
|
|
function setupNode(node) {
|
|
node._sxcpSeedControlNode = true;
|
|
if (!node._sxcpLockRandomSeedsButton) {
|
|
node._sxcpLockRandomSeedsButton = node.addWidget("button", "Lock Random Seeds Now", null, () => lockRandomSeeds(node));
|
|
}
|
|
resizeNode(node);
|
|
}
|
|
|
|
app.registerExtension({
|
|
name: EXTENSION,
|
|
|
|
async setup() {
|
|
if (app._sxcpSeedControlQueuePatched) return;
|
|
const originalQueuePrompt = app.queuePrompt;
|
|
if (!originalQueuePrompt) return;
|
|
|
|
app._sxcpSeedControlQueuePatched = true;
|
|
app.queuePrompt = async function () {
|
|
const randomSeedChanges = materializeAllForQueue();
|
|
try {
|
|
return await originalQueuePrompt.apply(this, arguments);
|
|
} finally {
|
|
restoreRandomModes(randomSeedChanges);
|
|
}
|
|
};
|
|
},
|
|
|
|
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);
|
|
setupNode(this);
|
|
return result;
|
|
};
|
|
|
|
const onConfigure = nodeType.prototype.onConfigure;
|
|
nodeType.prototype.onConfigure = function () {
|
|
const result = onConfigure?.apply(this, arguments);
|
|
queueMicrotask(() => setupNode(this));
|
|
return result;
|
|
};
|
|
},
|
|
});
|