89 lines
3.3 KiB
JavaScript
89 lines
3.3 KiB
JavaScript
import { app } from "/scripts/app.js";
|
|
|
|
// Configuration for each batch divider node type
|
|
const BATCH_DIVIDER_NODES = {
|
|
"ImageBatchDivider": { outputPrefix: "batch_", outputType: "IMAGE" },
|
|
"AudioBatchDivider": { outputPrefix: "audio_", outputType: "AUDIO" }
|
|
};
|
|
|
|
app.registerExtension({
|
|
name: "Distributed.BatchDividers",
|
|
async nodeCreated(node) {
|
|
const config = BATCH_DIVIDER_NODES[node.comfyClass];
|
|
if (!config) return;
|
|
|
|
try {
|
|
const updateOutputs = () => {
|
|
if (!node.widgets) return;
|
|
|
|
const divideByWidget = node.widgets.find(w => w.name === "divide_by");
|
|
if (!divideByWidget) return;
|
|
|
|
const divideBy = parseInt(divideByWidget.value, 10) || 1;
|
|
const totalOutputs = divideBy;
|
|
|
|
// Ensure outputs array exists
|
|
if (!node.outputs) node.outputs = [];
|
|
|
|
// Remove excess outputs
|
|
while (node.outputs.length > totalOutputs) {
|
|
node.removeOutput(node.outputs.length - 1);
|
|
}
|
|
|
|
// Add missing outputs
|
|
while (node.outputs.length < totalOutputs) {
|
|
const outputIndex = node.outputs.length + 1;
|
|
node.addOutput(`${config.outputPrefix}${outputIndex}`, config.outputType);
|
|
}
|
|
|
|
if (node.setDirty) node.setDirty(true);
|
|
};
|
|
|
|
// Initial update with delay to allow workflow loading
|
|
setTimeout(updateOutputs, 200);
|
|
|
|
// Find the widget and set up responsive handlers
|
|
const divideByWidget = node.widgets.find(w => w.name === "divide_by");
|
|
if (divideByWidget) {
|
|
const originalCallback = divideByWidget.callback;
|
|
divideByWidget.callback = (value) => {
|
|
updateOutputs();
|
|
if (originalCallback) originalCallback.call(divideByWidget, value);
|
|
};
|
|
|
|
if (divideByWidget.inputEl) {
|
|
divideByWidget.inputEl.addEventListener('input', updateOutputs);
|
|
}
|
|
|
|
const observer = new MutationObserver(updateOutputs);
|
|
if (divideByWidget.element) {
|
|
observer.observe(divideByWidget.element, { attributes: true, childList: true, subtree: true });
|
|
}
|
|
|
|
node._batchDividerCleanup = () => {
|
|
observer.disconnect();
|
|
if (divideByWidget.inputEl) {
|
|
divideByWidget.inputEl.removeEventListener('input', updateOutputs);
|
|
}
|
|
divideByWidget.callback = originalCallback;
|
|
};
|
|
}
|
|
|
|
const originalConfigure = node.configure;
|
|
node.configure = function(data) {
|
|
const result = originalConfigure ? originalConfigure.call(this, data) : undefined;
|
|
updateOutputs();
|
|
return result;
|
|
};
|
|
} catch (error) {
|
|
console.error(`Error in ${node.comfyClass} extension:`, error);
|
|
}
|
|
},
|
|
|
|
nodeBeforeRemove(node) {
|
|
if (BATCH_DIVIDER_NODES[node.comfyClass] && node._batchDividerCleanup) {
|
|
node._batchDividerCleanup();
|
|
}
|
|
}
|
|
});
|