From 612f99bda60e91d29c6ece027af80b30622120a0 Mon Sep 17 00:00:00 2001
From: flyingshutter <alois1.schneider@googlemail.com>
Date: Fri, 7 Apr 2023 19:32:27 +0200
Subject: [PATCH] move Node to frontend

---
 nodes.py                          | 17 -----------
 web/extensions/core/batchIndex.js | 48 +++++++++++++++++++++++++++++++
 web/scripts/widgets.js            |  1 -
 3 files changed, 48 insertions(+), 18 deletions(-)
 create mode 100644 web/extensions/core/batchIndex.js

diff --git a/nodes.py b/nodes.py
index 70be9bab6..5c3b3a4ee 100644
--- a/nodes.py
+++ b/nodes.py
@@ -1059,22 +1059,6 @@ class ImagePadForOutpaint:
 
         return (new_image, mask)
 
-class BatchIndex:
-    @classmethod
-    def INPUT_TYPES(cls):
-        return {
-            "required": {
-                "batch_index": ("INT", {}),
-            },
-        }
-
-    RETURN_TYPES = ("INT", "FLOAT")
-    FUNCTION = "getBatchIndex"
-    CATEGORY = "utils"
-
-    def getBatchIndex(self, batch_index):
-        return (batch_index, batch_index)
-
 
 NODE_CLASS_MAPPINGS = {
     "KSampler": KSampler,
@@ -1118,7 +1102,6 @@ NODE_CLASS_MAPPINGS = {
     "unCLIPCheckpointLoader": unCLIPCheckpointLoader,
     "CheckpointLoader": CheckpointLoader,
     "DiffusersLoader": DiffusersLoader,
-    "BatchIndex": BatchIndex,
 }
 
 def load_custom_node(module_path):
diff --git a/web/extensions/core/batchIndex.js b/web/extensions/core/batchIndex.js
new file mode 100644
index 000000000..8ba608880
--- /dev/null
+++ b/web/extensions/core/batchIndex.js
@@ -0,0 +1,48 @@
+import { ComfyWidgets } from "/scripts/widgets.js";
+import { app } from "/scripts/app.js";
+
+class BatchInfo {
+    constructor() {
+        this.addOutput("iteration", "INT");
+        
+        let widget = (ComfyWidgets["INT:batch_index"](this, "iteration", ["INT",{}], app) || {}).widget;
+        
+        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, {});
+                    }
+                }
+            }
+        }
+    }
+}
+
+app.registerExtension({
+	name: "Comfy.BatchInfo",
+    
+    registerCustomNodes() {
+        LiteGraph.registerNodeType(
+			"BatchInfo",
+			Object.assign(BatchInfo, {
+				title: "BatchInfo",
+			})
+		);
+		BatchInfo.category = "utils";
+    },
+});
\ No newline at end of file
diff --git a/web/scripts/widgets.js b/web/scripts/widgets.js
index aa5219865..ad8f01786 100644
--- a/web/scripts/widgets.js
+++ b/web/scripts/widgets.js
@@ -48,7 +48,6 @@ 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.forbidConvertToInput = true;
 	batchIndex.disabled = true;
 
 	batchIndex.onInitBatch = () => {