From aa29284a76d0782a9dbbab0d4e204b353ce5873a Mon Sep 17 00:00:00 2001
From: flyingshutter <alois1.schneider@googlemail.com>
Date: Fri, 7 Apr 2023 21:31:09 +0200
Subject: [PATCH] add batchCount ouptut have all widget related stuff in the
 extension instead of widgets.js

---
 web/extensions/core/batchIndex.js | 48 ++++++++++++++++++++-----------
 web/scripts/app.js                |  2 +-
 web/scripts/widgets.js            | 19 ------------
 3 files changed, 33 insertions(+), 36 deletions(-)

diff --git a/web/extensions/core/batchIndex.js b/web/extensions/core/batchIndex.js
index 8ba608880..386a1d7a7 100644
--- a/web/extensions/core/batchIndex.js
+++ b/web/extensions/core/batchIndex.js
@@ -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, {});
+                            }
+                        }
                     }
                 }
             }
diff --git a/web/scripts/app.js b/web/scripts/app.js
index bd1762dfa..e1a459cf0 100644
--- a/web/scripts/app.js
+++ b/web/scripts/app.js
@@ -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);
 									}
 								}
 							}
diff --git a/web/scripts/widgets.js b/web/scripts/widgets.js
index ad8f01786..d1a9c6c6e 100644
--- a/web/scripts/widgets.js
+++ b/web/scripts/widgets.js
@@ -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) {