From 981a5eb27524b90cdeabeb7a3974f53ebef2aea6 Mon Sep 17 00:00:00 2001
From: rgthree <regis.gaughan@gmail.com>
Date: Fri, 1 Dec 2023 22:20:25 -0500
Subject: [PATCH] Only swap seed widget values if ComfyUI is currently
 processing a queue. Fixes #62

---
 ts/rgthree.ts  | 12 +++++++++---
 ts/seed.ts     |  7 +++++++
 web/rgthree.js | 12 +++++++++---
 web/seed.js    |  4 ++++
 4 files changed, 29 insertions(+), 6 deletions(-)

diff --git a/ts/rgthree.ts b/ts/rgthree.ts
index 6d9c95a..5e97056 100644
--- a/ts/rgthree.ts
+++ b/ts/rgthree.ts
@@ -102,6 +102,8 @@ class Rgthree {
   monitorBadLinksAlerted = false;
   monitorLinkTimeout: number|null = null;
 
+  processingQueue = false;
+
   async clearAllMessages() {
     let container = document.querySelector('.rgthree-top-messages-container');
     container && (container.innerHTML = '');
@@ -194,9 +196,13 @@ class Rgthree {
     const queuePrompt = app.queuePrompt as Function;
     app.queuePrompt = async function() {
       that.fireEvent('queue', {});
-      let promise = queuePrompt.apply(app, [...arguments]);
-      that.fireEvent('queue-end', {});
-      return promise;
+      that.processingQueue = true;
+      try {
+        await queuePrompt.apply(app, [...arguments])
+      } finally {
+        that.processingQueue = false;
+        that.fireEvent('queue-end', {});
+      }
     }
 
     const graphToPrompt = app.graphToPrompt as Function;
diff --git a/ts/seed.ts b/ts/seed.ts
index 728ad67..2bf0dea 100644
--- a/ts/seed.ts
+++ b/ts/seed.ts
@@ -14,6 +14,7 @@ import type {
 } from "./typings/litegraph.js";
 import type { ComfyApp, ComfyObjectInfo, ComfyWidget, ComfyGraphNode } from "./typings/comfy.js";
 import { RgthreeBaseNode } from "./base_node.js";
+import { rgthree } from "./rgthree.js";
 
 declare const LiteGraph: typeof TLiteGraph;
 declare const LGraphNode: typeof TLGraphNode;
@@ -121,6 +122,12 @@ class SeedControl {
      */
     this.seedWidget.serializeValue = async (node: SerializedLGraphNode, index: number) => {
       const inputSeed = this.seedWidget.value;
+      // Only actually swap and set the value when we're currently queuing. Some other nodes, like
+      // cg-use-everywhere, serializes the graph even when not queing, and we don't want to swap
+      // widget values in these cases.
+      if (!rgthree.processingQueue) {
+        return inputSeed;
+      }
       this.serializedCtx = {
         inputSeed: this.seedWidget.value,
       };
diff --git a/web/rgthree.js b/web/rgthree.js
index c7ba03b..ece78a4 100644
--- a/web/rgthree.js
+++ b/web/rgthree.js
@@ -125,6 +125,7 @@ class Rgthree {
         this.logger = new LogSession("[rgthree]");
         this.monitorBadLinksAlerted = false;
         this.monitorLinkTimeout = null;
+        this.processingQueue = false;
         this.eventsToFns = new Map();
         window.addEventListener("keydown", (e) => {
             this.ctrlKey = !!e.ctrlKey;
@@ -142,9 +143,14 @@ class Rgthree {
         const queuePrompt = app.queuePrompt;
         app.queuePrompt = async function () {
             that.fireEvent('queue', {});
-            let promise = queuePrompt.apply(app, [...arguments]);
-            that.fireEvent('queue-end', {});
-            return promise;
+            that.processingQueue = true;
+            try {
+                await queuePrompt.apply(app, [...arguments]);
+            }
+            finally {
+                that.processingQueue = false;
+                that.fireEvent('queue-end', {});
+            }
         };
         const graphToPrompt = app.graphToPrompt;
         app.graphToPrompt = async function () {
diff --git a/web/seed.js b/web/seed.js
index baf18aa..b744dfa 100644
--- a/web/seed.js
+++ b/web/seed.js
@@ -1,5 +1,6 @@
 import { app } from "../../scripts/app.js";
 import { ComfyWidgets } from "../../scripts/widgets.js";
+import { rgthree } from "./rgthree.js";
 const LAST_SEED_BUTTON_LABEL = "♻️ (Use Last Queued Seed)";
 const SPECIAL_SEED_RANDOM = -1;
 const SPECIAL_SEED_INCREMENT = -2;
@@ -54,6 +55,9 @@ class SeedControl {
         this.lastSeedButton.disabled = true;
         this.seedWidget.serializeValue = async (node, index) => {
             const inputSeed = this.seedWidget.value;
+            if (!rgthree.processingQueue) {
+                return inputSeed;
+            }
             this.serializedCtx = {
                 inputSeed: this.seedWidget.value,
             };