From f9ef1b61553477c855aefc1b947dce5ef289baf1 Mon Sep 17 00:00:00 2001
From: space-nuko <24979496+space-nuko@users.noreply.github.com>
Date: Mon, 15 May 2023 18:34:10 -0500
Subject: [PATCH] Fix output display when executed

---
 execution.py       |  8 ++++----
 web/scripts/app.js | 50 ++++++++++++++++++++++------------------------
 2 files changed, 28 insertions(+), 30 deletions(-)

diff --git a/execution.py b/execution.py
index 06a1f40c7..e555333fa 100644
--- a/execution.py
+++ b/execution.py
@@ -192,20 +192,20 @@ def get_output_data(obj, input_data_all_batches, server, unique_id, prompt_id):
         if len(uis) > 0:
             output_ui = {k: [y for x in uis for y in x[k]] for k in uis[0].keys()}
 
+        all_outputs.append(output)
+        all_outputs_ui.append(output_ui)
+
         # update the UI after each batch finishes
         if server.client_id is not None:
             message = {
                 "node": unique_id,
-                "output": output_ui,
+                "output": all_outputs_ui, # list of outputs so far
                 "prompt_id": prompt_id,
                 "batch_num": batch_num,
                 "total_batches": total_batches
             }
             server.send_sync("executed", message, server.client_id)
 
-        all_outputs.append(output)
-        all_outputs_ui.append(output_ui)
-
     return all_outputs, all_outputs_ui
 
 def recursive_execute(server, prompt, outputs, current_item, extra_data, executed, prompt_id, outputs_ui):
diff --git a/web/scripts/app.js b/web/scripts/app.js
index 66504d326..95979198d 100644
--- a/web/scripts/app.js
+++ b/web/scripts/app.js
@@ -355,32 +355,30 @@ export class ComfyApp {
 			if (!this.flags.collapsed) {
 				// Outputs returned by the frontend for each node are a list, one for each combinatorial batch run.
 				// With no combinatorial outputs, it's a list of length 1.
-				// For now select only the first batch output.
-				const batchOutputs = app.nodeOutputs[this.id + ""];
-				if (batchOutputs && batchOutputs.length > 0) {
-					const output = batchOutputs[0];
-					if (output && output.images) {
-						if (this.images !== output.images) {
-							this.images = output.images;
-							this.imgs = null;
-							this.imageIndex = null;
-							Promise.all(
-								output.images.map((src) => {
-									return new Promise((r) => {
-										const img = new Image();
-										img.onload = () => r(img);
-										img.onerror = () => r(null);
-										img.src = "/view?" + new URLSearchParams(src).toString();
-									});
-								})
-							).then((imgs) => {
-								if (this.images === output.images) {
-									this.imgs = imgs.filter(Boolean);
-									this.setSizeForImage?.();
-									app.graph.setDirtyCanvas(true);
-								}
-							});
-						}
+				const outputs = app.nodeOutputs[this.id + ""]
+				if (outputs && this.batchOutputs !== outputs) {
+                  this.batchOutputs = outputs;
+					const batchImages = outputs.filter(Boolean).flatMap(o => o.images || []);
+					if (this.images !== batchImages) {
+						this.images = batchImages;
+						this.imgs = null;
+						this.imageIndex = null;
+						Promise.all(
+							batchImages.map((src) => {
+								return new Promise((r) => {
+									const img = new Image();
+									img.onload = () => r(img);
+									img.onerror = () => r(null);
+									img.src = "/view?" + new URLSearchParams(src).toString();
+								});
+							})
+						).then((imgs) => {
+							if (this.images === batchImages) {
+								this.imgs = imgs.filter(Boolean);
+								this.setSizeForImage?.();
+								app.graph.setDirtyCanvas(true);
+							}
+						});
 					}
 				}