add batchCount ouptut

have all widget related stuff in the extension instead of widgets.js
This commit is contained in:
flyingshutter 2023-04-07 21:31:09 +02:00
parent 612f99bda6
commit aa29284a76
3 changed files with 33 additions and 36 deletions

View File

@ -4,28 +4,44 @@ import { app } from "/scripts/app.js";
class BatchInfo {
constructor() {
this.addOutput("iteration", "INT");
this.addOutput("batchCount", "INT");
let widget = (ComfyWidgets["INT:batch_index"](this, "iteration", ["INT",{}], app) || {}).widget;
const batchIteration = (ComfyWidgets["INT"](this, "iteration", ["INT",{}], app) || {}).widget;
batchIteration.disabled = true;
const batchCount = (ComfyWidgets["INT"](this, "batchCount", ["INT",{}], app) || {}).widget;
batchCount.disabled = true;
batchIteration.onInitBatch = (batchSize) => {
batchIteration.value = 0;
batchCount.value = batchSize;
};
batchIteration.afterQueued = () => {
batchIteration.value += 1;
};
this.serialize_widgets = true;
this.isVirtualNode = true;
}
applyToGraph() {
if (!this.outputs[0].links?.length) return;
// For each output link copy our value over the original widget value
for (const l of this.outputs[0].links) {
const linkInfo = app.graph.links[l];
const node = this.graph.getNodeById(linkInfo.target_id);
const input = node.inputs[linkInfo.target_slot];
const widgetName = input.widget.name;
if (widgetName) {
const widget = node.widgets.find((w) => w.name === widgetName);
if (widget) {
widget.value = this.widgets[0].value;
if (widget.callback) {
widget.callback(widget.value, app.canvas, node, app.canvas.graph_mouse, {});
for (const idx in this.outputs) {
// For each output link copy our value over the original widget value
if (this.outputs[idx].links?.length) {
for (const l of this.outputs[idx].links) {
const linkInfo = app.graph.links[l];
const node = this.graph.getNodeById(linkInfo.target_id);
const input = node.inputs[linkInfo.target_slot];
const widgetName = input.widget.name;
if (widgetName) {
const widget = node.widgets.find((w) => w.name === widgetName);
if (widget) {
widget.value = this.widgets[idx].value;
if (widget.callback) {
widget.callback(widget.value, app.canvas, node, app.canvas.graph_mouse, {});
}
}
}
}
}

View File

@ -997,7 +997,7 @@ class ComfyApp {
// Allow widgets to run callbacks on firts iteration of a batch
// e.g. random seed after every gen
if (widget.onInitBatch) {
widget.onInitBatch();
widget.onInitBatch(batchCount);
}
}
}

View File

@ -43,24 +43,6 @@ function seedWidget(node, inputName, inputData) {
return { widget: seed, randomize };
}
export function batchIndexWidget(node, inputName, inputData) {
const { val, config } = getNumberDefaults(inputData, 1);
Object.assign(config, { precision: 0 });
const batchIndex = node.addWidget("number", inputName, val, () => {}, config)
batchIndex.disabled = true;
batchIndex.onInitBatch = () => {
batchIndex.value = 0;
};
batchIndex.afterQueued = () => {
batchIndex.value += 1;
};
return batchIndex;
}
const MultilineSymbol = Symbol();
const MultilineResizeSymbol = Symbol();
@ -215,7 +197,6 @@ function addMultilineWidget(node, name, opts, app) {
}
export const ComfyWidgets = {
"INT:batch_index": batchIndexWidget,
"INT:seed": seedWidget,
"INT:noise_seed": seedWidget,
FLOAT(node, inputName, inputData) {