From 77c8ad6e02f59bd593f1981bf23414d2112731ff Mon Sep 17 00:00:00 2001
From: huchenlei <huchenlei@proton.me>
Date: Fri, 21 Feb 2025 14:05:10 -0500
Subject: [PATCH] [Refactor] Move Widget.beforeQueued invocation from
 graphToPrompt to queuePrompt

---
 src/scripts/app.ts         | 28 ++++++++++++----------------
 src/utils/executionUtil.ts |  8 --------
 src/utils/litegraphUtil.ts | 11 +++++++++++
 3 files changed, 23 insertions(+), 24 deletions(-)

diff --git a/src/scripts/app.ts b/src/scripts/app.ts
index 4b45c507c..d761097df 100644
--- a/src/scripts/app.ts
+++ b/src/scripts/app.ts
@@ -39,7 +39,7 @@ import {
 import { ExtensionManager } from '@/types/extensionTypes'
 import { ColorAdjustOptions, adjustColor } from '@/utils/colorUtil'
 import { graphToPrompt } from '@/utils/executionUtil'
-import { isImageNode } from '@/utils/litegraphUtil'
+import { executeWidgetsCallback, isImageNode } from '@/utils/litegraphUtil'
 import { deserialiseAndCreate } from '@/utils/vintageClipboard'
 
 import { type ComfyApi, api } from './api'
@@ -1260,7 +1260,7 @@ export class ComfyApp {
     return '(unknown error)'
   }
 
-  async queuePrompt(number, batchCount = 1) {
+  async queuePrompt(number: number, batchCount: number = 1): Promise<boolean> {
     this.#queueItems.push({ number, batchCount })
 
     // Only have one action process the items so each one gets a unique seed correctly
@@ -1276,8 +1276,11 @@ export class ComfyApp {
         ;({ number, batchCount } = this.#queueItems.pop())
 
         for (let i = 0; i < batchCount; i++) {
-          const p = await this.graphToPrompt()
+          // Allow widgets to run callbacks before a prompt has been queued
+          // e.g. random seed before every gen
+          executeWidgetsCallback(this.graph.nodes, 'beforeQueued')
 
+          const p = await this.graphToPrompt()
           try {
             const res = await api.queuePrompt(number, p)
             this.lastNodeErrors = res.node_errors
@@ -1303,19 +1306,12 @@ export class ComfyApp {
             break
           }
 
-          for (const n of p.workflow.nodes) {
-            const node = this.graph.getNodeById(n.id)
-            if (node.widgets) {
-              for (const widget of node.widgets) {
-                // Allow widgets to run callbacks after a prompt has been queued
-                // e.g. random seed after every gen
-                if (widget.afterQueued) {
-                  widget.afterQueued()
-                }
-              }
-            }
-          }
-
+          // Allow widgets to run callbacks after a prompt has been queued
+          // e.g. random seed after every gen
+          executeWidgetsCallback(
+            p.workflow.nodes.map((n) => this.graph.getNodeById(n.id)),
+            'afterQueued'
+          )
           this.canvas.draw(true, true)
           await this.ui.queue.update()
         }
diff --git a/src/utils/executionUtil.ts b/src/utils/executionUtil.ts
index fa7bc00f9..a759623fe 100644
--- a/src/utils/executionUtil.ts
+++ b/src/utils/executionUtil.ts
@@ -15,14 +15,6 @@ export const graphToPrompt = async (
   const { sortNodes = false } = options
 
   for (const outerNode of graph.computeExecutionOrder(false)) {
-    if (outerNode.widgets) {
-      for (const widget of outerNode.widgets) {
-        // Allow widgets to run callbacks before a prompt has been queued
-        // e.g. random seed before every gen
-        widget.beforeQueued?.()
-      }
-    }
-
     const innerNodes = outerNode.getInnerNodes
       ? outerNode.getInnerNodes()
       : [outerNode]
diff --git a/src/utils/litegraphUtil.ts b/src/utils/litegraphUtil.ts
index 96b0f18c6..870e88114 100644
--- a/src/utils/litegraphUtil.ts
+++ b/src/utils/litegraphUtil.ts
@@ -42,3 +42,14 @@ export const getItemsColorOption = (items: unknown[]): ColorOption | null => {
     ? _.head(colorOptions)!
     : null
 }
+
+export function executeWidgetsCallback(
+  nodes: LGraphNode[],
+  callbackName: 'onRemove' | 'beforeQueued' | 'afterQueued'
+) {
+  for (const node of nodes) {
+    for (const widget of node.widgets ?? []) {
+      widget[callbackName]?.()
+    }
+  }
+}