From c13fe47de617b15ae583b0fdd8f8c48f13615512 Mon Sep 17 00:00:00 2001
From: space-nuko <24979496+space-nuko@users.noreply.github.com>
Date: Mon, 15 May 2023 18:41:15 -0500
Subject: [PATCH] Fix outputs shifting around when restoring from history

---
 execution.py       | 11 +++++++++--
 web/scripts/app.js |  2 --
 2 files changed, 9 insertions(+), 4 deletions(-)

diff --git a/execution.py b/execution.py
index e555333fa..6435f4341 100644
--- a/execution.py
+++ b/execution.py
@@ -195,11 +195,15 @@ def get_output_data(obj, input_data_all_batches, server, unique_id, prompt_id):
         all_outputs.append(output)
         all_outputs_ui.append(output_ui)
 
+        outputs_ui_to_send = None
+        if any(all_outputs_ui):
+            outputs_ui_to_send = all_outputs_ui
+
         # update the UI after each batch finishes
         if server.client_id is not None:
             message = {
                 "node": unique_id,
-                "output": all_outputs_ui, # list of outputs so far
+                "output": outputs_ui_to_send,
                 "prompt_id": prompt_id,
                 "batch_num": batch_num,
                 "total_batches": total_batches
@@ -233,7 +237,10 @@ def recursive_execute(server, prompt, outputs, current_item, extra_data, execute
 
     output_data_from_batches, output_ui_from_batches = get_output_data(obj, input_data_all_batches, server, unique_id, prompt_id)
     outputs[unique_id] = output_data_from_batches
-    outputs_ui[unique_id] = output_ui_from_batches
+    if any(output_ui_from_batches):
+        outputs_ui[unique_id] = output_ui_from_batches
+    elif unique_id in outputs_ui:
+        outputs_ui.pop(unique_id)
     executed.add(unique_id)
 
 def recursive_will_execute(prompt, outputs, current_item):
diff --git a/web/scripts/app.js b/web/scripts/app.js
index 95979198d..39fa93540 100644
--- a/web/scripts/app.js
+++ b/web/scripts/app.js
@@ -353,8 +353,6 @@ export class ComfyApp {
 
 		node.prototype.onDrawBackground = function (ctx) {
 			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.
 				const outputs = app.nodeOutputs[this.id + ""]
 				if (outputs && this.batchOutputs !== outputs) {
                   this.batchOutputs = outputs;