From e663f4db8865131e280a62ecd728a58e048bc6ed Mon Sep 17 00:00:00 2001
From: space-nuko <24979496+space-nuko@users.noreply.github.com>
Date: Thu, 4 May 2023 16:25:23 -0500
Subject: [PATCH] AfterQueued action

---
 litegraph                         |  2 +-
 src/lib/components/ComfyApp.ts    | 31 +++++++++++++++++-------
 src/lib/nodes/ComfyActionNodes.ts | 40 +++++++++++++++++++++++++++++++
 src/lib/nodes/ComfyGraphNode.ts   |  3 ++-
 4 files changed, 65 insertions(+), 11 deletions(-)

diff --git a/litegraph b/litegraph
index 93b0ed1..0e4bbd2 160000
--- a/litegraph
+++ b/litegraph
@@ -1 +1 @@
-Subproject commit 93b0ed15b7f5ca70ab074ebe7d8b39d3e5963818
+Subproject commit 0e4bbd2169661f1ea47d7abccdba071bc0ba1c9a
diff --git a/src/lib/components/ComfyApp.ts b/src/lib/components/ComfyApp.ts
index 69ed594..22193eb 100644
--- a/src/lib/components/ComfyApp.ts
+++ b/src/lib/components/ComfyApp.ts
@@ -43,6 +43,19 @@ export type SerializedAppState = {
     canvas: SerializedGraphCanvasState
 }
 
+/** [link origin, link index] | value */
+export type SerializedPromptInput = [string, number] | any
+
+export type SerializedPromptInputs = {
+    inputs: Record<string, SerializedPromptInput>,
+    class_type: string
+}
+
+export type SerializedPrompt = {
+    workflow: SerializedLGraph,
+    output: Record<string, SerializedPromptInputs>
+}
+
 export type Progress = {
     value: number,
     max: number
@@ -402,7 +415,7 @@ export default class ComfyApp {
      * Converts the current graph workflow for sending to the API
      * @returns The workflow and node links
      */
-    async graphToPrompt() {
+    async graphToPrompt(): Promise<SerializedPrompt> {
         // Run frontend-only logic
         this.lGraph.runStep(1)
 
@@ -410,14 +423,16 @@ export default class ComfyApp {
 
         const output = {};
         // Process nodes in order of execution
-        for (const node of this.lGraph.computeExecutionOrder<ComfyGraphNode>(false, null)) {
-            const n = workflow.nodes.find((n) => n.id === node.id);
+        for (const node_ of this.lGraph.computeExecutionOrder<ComfyGraphNode>(false, null)) {
+            const n = workflow.nodes.find((n) => n.id === node_.id);
 
-            if (!node.isBackendNode) {
-                console.debug("Not serializing node: ", node.type)
+            if (!node_.isBackendNode) {
+                console.debug("Not serializing node: ", node_.type)
                 continue;
             }
 
+            const node = node_ as ComfyBackendNode;
+
             if (node.mode === NodeMode.NEVER) {
                 // Don't serialize muted nodes
                 continue;
@@ -444,10 +459,8 @@ export default class ComfyApp {
                     if ("config" in inp)
                         serialize = (inp as IComfyInputSlot).serialize
 
-                    let isBackendNode = false;
+                    let isBackendNode = node.isBackendNode;
                     let isInputBackendNode = false;
-                    if ("isBackendNode" in node)
-                        isBackendNode = (node as ComfyGraphNode).isBackendNode;
                     if ("isBackendNode" in inputNode)
                         isInputBackendNode = (inputNode as ComfyGraphNode).isBackendNode;
 
@@ -544,7 +557,7 @@ export default class ComfyApp {
                     for (const n of p.workflow.nodes) {
                         const node = this.lGraph.getNodeById(n.id);
                         if ("afterQueued" in node) {
-                            (node as ComfyGraphNode).afterQueued();
+                            (node as ComfyGraphNode).afterQueued(p);
                         }
                     }
 
diff --git a/src/lib/nodes/ComfyActionNodes.ts b/src/lib/nodes/ComfyActionNodes.ts
index 091019e..e58a2ba 100644
--- a/src/lib/nodes/ComfyActionNodes.ts
+++ b/src/lib/nodes/ComfyActionNodes.ts
@@ -1,6 +1,46 @@
 import { LiteGraph, type ContextMenuItem, type LGraphNode, type Vector2, LConnectionKind, LLink, LGraphCanvas, type SlotType, TitleMode, type SlotLayout, BuiltInSlotType, type ITextWidget } from "@litegraph-ts/core";
 import ComfyGraphNode from "./ComfyGraphNode";
 import { Watch } from "@litegraph-ts/nodes-basic";
+import type { SerializedPrompt } from "$lib/components/ComfyApp";
+
+export interface ComfyAfterQueuedAction extends Record<any, any> {
+    prompt: SerializedPrompt
+}
+
+export class ComfyAfterQueuedAction extends ComfyGraphNode {
+    override properties: ComfyCopyActionProperties = {
+        prompt: null
+    }
+
+    static slotLayout: SlotLayout = {
+        outputs: [
+            { name: "afterQueued", type: BuiltInSlotType.EVENT },
+            { name: "prompt", type: "*" }
+        ],
+    }
+
+    override onPropertyChanged(property: string, value: any, prevValue?: any) {
+        if (property === "value") {
+            this.setOutputData(0, this.properties.prompt)
+        }
+    }
+
+    override onExecute() {
+        this.setOutputData(0, this.properties.prompt)
+    }
+
+    override afterQueued(p: SerializedPrompt) {
+        this.setProperty("value", p)
+        this.triggerSlot(0, "bang")
+    }
+}
+
+LiteGraph.registerNodeType({
+    class: ComfyAfterQueuedAction,
+    title: "Comfy.AfterQueuedAction",
+    desc: "Triggers a 'bang' event when a prompt is queued.",
+    type: "actions/after_queued"
+})
 
 export interface ComfyCopyActionProperties extends Record<any, any> {
     value: any
diff --git a/src/lib/nodes/ComfyGraphNode.ts b/src/lib/nodes/ComfyGraphNode.ts
index 9bfb588..bbe87d1 100644
--- a/src/lib/nodes/ComfyGraphNode.ts
+++ b/src/lib/nodes/ComfyGraphNode.ts
@@ -1,9 +1,10 @@
+import type { SerializedPrompt } from "$lib/components/ComfyApp";
 import type ComfyWidget from "$lib/components/widgets/ComfyWidget";
 import { LGraph, LGraphNode } from "@litegraph-ts/core";
 
 export default class ComfyGraphNode extends LGraphNode {
     isBackendNode?: boolean;
 
-    afterQueued?(): void;
+    afterQueued?(prompt: SerializedPrompt): void;
     onExecuted?(output: any): void;
 }