From eb02561906c2c673f054975e7b199aa8b9bfb2cd Mon Sep 17 00:00:00 2001
From: space-nuko <24979496+space-nuko@users.noreply.github.com>
Date: Mon, 5 Jun 2023 15:45:08 -0500
Subject: [PATCH] Basic preview support

(as of latest PR commit)
---
 src/lib/api.ts                                | 97 ++++++++++++-------
 src/lib/components/ComfyApp.ts                |  4 +
 .../components/ComfyBoxWorkflowsView.svelte   |  3 +
 src/lib/nodes/widgets/ComfyGalleryNode.ts     |  6 +-
 src/lib/stores/layoutStates.ts                |  8 ++
 src/lib/stores/queueState.ts                  | 37 +++++++
 src/lib/stores/uiState.ts                     |  2 +
 src/lib/widgets/GalleryWidget.svelte          | 77 ++++++++++++++-
 8 files changed, 198 insertions(+), 36 deletions(-)

diff --git a/src/lib/api.ts b/src/lib/api.ts
index 492aa2b..ec54872 100644
--- a/src/lib/api.ts
+++ b/src/lib/api.ts
@@ -101,6 +101,7 @@ export type ComfyUIPromptExtraData = {
 }
 
 type ComfyAPIEvents = {
+    // JSON
     status: (status: ComfyAPIStatusResponse | null, error?: Error | null) => void,
     progress: (progress: Progress) => void,
     reconnecting: () => void,
@@ -111,6 +112,9 @@ type ComfyAPIEvents = {
     execution_cached: (promptID: PromptID, nodes: ComfyNodeID[]) => void,
     execution_interrupted: (error: ComfyInterruptedError) => void,
     execution_error: (error: ComfyExecutionError) => void,
+
+    // Binary
+    b_preview: (imageBlob: Blob) => void
 }
 
 export default class ComfyAPI {
@@ -126,7 +130,7 @@ export default class ComfyAPI {
     }
 
     /**
-     * Poll status  for colab and other things that don't support websockets.
+     * Poll status for colab and other things that don't support websockets.
      */
     private pollQueue() {
         setInterval(async () => {
@@ -176,6 +180,7 @@ export default class ComfyAPI {
         this.socket = new WebSocket(
             `ws${window.location.protocol === "https:" ? "s" : ""}://${hostname}:${port}/ws${existingSession}`
         );
+        this.socket.binaryType = "arraybuffer";
 
         this.socket.addEventListener("open", () => {
             opened = true;
@@ -204,38 +209,64 @@ export default class ComfyAPI {
 
         this.socket.addEventListener("message", (event) => {
             try {
-                const msg = JSON.parse(event.data);
-                switch (msg.type) {
-                    case "status":
-                        if (msg.data.sid) {
-                            this.clientId = msg.data.sid;
-                            sessionStorage["Comfy.SessionId"] = this.clientId;
-                        }
-                        this.eventBus.emit("status", { execInfo: { queueRemaining: msg.data.status.exec_info.queue_remaining } });
-                        break;
-                    case "progress":
-                        this.eventBus.emit("progress", msg.data as Progress);
-                        break;
-                    case "executing":
-                        this.eventBus.emit("executing", msg.data.prompt_id, msg.data.node);
-                        break;
-                    case "executed":
-                        this.eventBus.emit("executed", msg.data.prompt_id, msg.data.node, msg.data.output);
-                        break;
-                    case "execution_start":
-                        this.eventBus.emit("execution_start", msg.data.prompt_id);
-                        break;
-                    case "execution_cached":
-                        this.eventBus.emit("execution_cached", msg.data.prompt_id, msg.data.nodes);
-                        break;
-                    case "execution_interrupted":
-                        this.eventBus.emit("execution_interrupted", msg.data);
-                        break;
-                    case "execution_error":
-                        this.eventBus.emit("execution_error", msg.data);
-                        break;
-                    default:
-                        console.warn("Unhandled message:", event.data);
+                if (event.data instanceof ArrayBuffer) {
+                    const view = new DataView(event.data);
+                    const eventType = view.getUint32(0);
+                    const buffer = event.data.slice(4);
+                    switch (eventType) {
+                        case 1:
+                            const view2 = new DataView(event.data);
+                            const imageType = view2.getUint32(0)
+                            let imageMime: string
+                            switch (imageType) {
+                                case 1:
+                                default:
+                                    imageMime = "image/jpeg";
+                                    break;
+                                case 2:
+                                    imageMime = "image/png"
+                            }
+                            const imageBlob = new Blob([buffer.slice(4)], { type: imageMime });
+                            this.eventBus.emit("b_preview", imageBlob);
+                            break;
+                        default:
+                            throw new Error(`Unknown binary websocket message of type ${eventType}`);
+                    }
+                }
+                else {
+                    const msg = JSON.parse(event.data);
+                    switch (msg.type) {
+                        case "status":
+                            if (msg.data.sid) {
+                                this.clientId = msg.data.sid;
+                                sessionStorage["Comfy.SessionId"] = this.clientId;
+                            }
+                            this.eventBus.emit("status", { execInfo: { queueRemaining: msg.data.status.exec_info.queue_remaining } });
+                            break;
+                        case "progress":
+                            this.eventBus.emit("progress", msg.data as Progress);
+                            break;
+                        case "executing":
+                            this.eventBus.emit("executing", msg.data.prompt_id, msg.data.node);
+                            break;
+                        case "executed":
+                            this.eventBus.emit("executed", msg.data.prompt_id, msg.data.node, msg.data.output);
+                            break;
+                        case "execution_start":
+                            this.eventBus.emit("execution_start", msg.data.prompt_id);
+                            break;
+                        case "execution_cached":
+                            this.eventBus.emit("execution_cached", msg.data.prompt_id, msg.data.nodes);
+                            break;
+                        case "execution_interrupted":
+                            this.eventBus.emit("execution_interrupted", msg.data);
+                            break;
+                        case "execution_error":
+                            this.eventBus.emit("execution_error", msg.data);
+                            break;
+                        default:
+                            console.warn("Unhandled message:", event.data);
+                    }
                 }
             } catch (error) {
                 console.error("Error handling message", event.data, error);
diff --git a/src/lib/components/ComfyApp.ts b/src/lib/components/ComfyApp.ts
index 97abc08..0f0369c 100644
--- a/src/lib/components/ComfyApp.ts
+++ b/src/lib/components/ComfyApp.ts
@@ -651,6 +651,10 @@ export default class ComfyApp {
             }
         });
 
+        this.api.addEventListener("b_preview", (imageBlob: Blob) => {
+            queueState.previewUpdated(imageBlob);
+        });
+
         const config = get(configState);
 
         if (config.pollSystemStatsInterval > 0) {
diff --git a/src/lib/components/ComfyBoxWorkflowsView.svelte b/src/lib/components/ComfyBoxWorkflowsView.svelte
index 4fd5430..722a1dd 100644
--- a/src/lib/components/ComfyBoxWorkflowsView.svelte
+++ b/src/lib/components/ComfyBoxWorkflowsView.svelte
@@ -386,6 +386,9 @@
                 <span style="display: inline-flex !important; padding: 0 0.75rem;">
                     <Checkbox label="Auto-Add UI" bind:value={$uiState.autoAddUI}/>
                 </span>
+                <span style="display: inline-flex !important; padding: 0 0.75rem;">
+                    <Checkbox label="Hide Previews" bind:value={$uiState.hidePreviews}/>
+                </span>
                 <!-- <span class="label" for="ui-edit-mode">
                      <BlockTitle>UI Edit mode</BlockTitle>
                      <select id="ui-edit-mode" name="ui-edit-mode" bind:value={$uiState.uiEditMode}>
diff --git a/src/lib/nodes/widgets/ComfyGalleryNode.ts b/src/lib/nodes/widgets/ComfyGalleryNode.ts
index 1267d18..0896a67 100644
--- a/src/lib/nodes/widgets/ComfyGalleryNode.ts
+++ b/src/lib/nodes/widgets/ComfyGalleryNode.ts
@@ -9,7 +9,8 @@ import ComfyWidgetNode from "./ComfyWidgetNode";
 export interface ComfyGalleryProperties extends ComfyWidgetProperties {
     index: number | null,
     updateMode: "replace" | "append",
-    autoSelectOnUpdate: boolean
+    autoSelectOnUpdate: boolean,
+    showPreviews: boolean
 }
 
 export default class ComfyGalleryNode extends ComfyWidgetNode<ComfyBoxImageMetadata[]> {
@@ -18,7 +19,8 @@ export default class ComfyGalleryNode extends ComfyWidgetNode<ComfyBoxImageMetad
         defaultValue: [],
         index: 0,
         updateMode: "replace",
-        autoSelectOnUpdate: true
+        autoSelectOnUpdate: true,
+        showPreviews: true
     }
 
     static slotLayout: SlotLayout = {
diff --git a/src/lib/stores/layoutStates.ts b/src/lib/stores/layoutStates.ts
index bae1f34..31c4457 100644
--- a/src/lib/stores/layoutStates.ts
+++ b/src/lib/stores/layoutStates.ts
@@ -615,6 +615,14 @@ const ALL_ATTRIBUTES: AttributesSpecList = [
                 validNodeTypes: ["ui/gallery"],
                 defaultValue: true
             },
+            {
+                name: "showPreviews",
+                type: "boolean",
+                location: "nodeProps",
+                editable: true,
+                validNodeTypes: ["ui/gallery"],
+                defaultValue: true
+            },
 
             // ImageUpload
             {
diff --git a/src/lib/stores/queueState.ts b/src/lib/stores/queueState.ts
index 9e45774..f716630 100644
--- a/src/lib/stores/queueState.ts
+++ b/src/lib/stores/queueState.ts
@@ -22,6 +22,7 @@ type QueueStateOps = {
     executionCached: (promptID: PromptID, nodes: ComfyNodeID[]) => void,
     executionError: (error: ComfyExecutionError) => CompletedQueueEntry | null,
     progressUpdated: (progress: Progress) => void
+    previewUpdated: (imageBlob: Blob) => void
     getQueueEntry: (promptID: PromptID) => QueueEntry | null;
     afterQueued: (workflowID: WorkflowInstID, promptID: PromptID, number: number, prompt: SerializedPromptInputsAll, extraData: any) => void
     queueItemDeleted: (type: QueueItemType, id: PromptID) => void;
@@ -88,6 +89,11 @@ export type QueueState = {
      */
     runningNodeID: ComfyNodeID | null;
 
+    /*
+     * Currently executing prompt if any
+     */
+    runningPromptID: PromptID | null;
+
     /*
      * Nodes which should be rendered as "executing" in the frontend (green border).
      * This includes the running node and all its parent subgraphs
@@ -98,6 +104,12 @@ export type QueueState = {
      * Progress for the current node reported by the frontend
      */
     progress: Progress | null,
+
+    /*
+     * Image preview URL
+     */
+    previewURL: string | null,
+
     /**
      * If true, user pressed the "Interrupt" button in the frontend. Disable the
      * button and wait until the next prompt starts running to re-enable it
@@ -115,6 +127,7 @@ const store: Writable<QueueState> = writable({
     runningNodeID: null,
     executingNodes: new Set(),
     progress: null,
+    preview: null,
     isInterrupting: false
 })
 
@@ -171,6 +184,19 @@ function progressUpdated(progress: Progress) {
     })
 }
 
+function previewUpdated(imageBlob: Blob) {
+    console.debug("[queueState] previewUpdated", imageBlob?.type)
+    store.update(s => {
+        if (s.runningNodeID == null) {
+            s.previewURL = null;
+            return s;
+        }
+
+        s.previewURL = URL.createObjectURL(imageBlob);
+        return s;
+    })
+}
+
 function statusUpdated(status: ComfyAPIStatusResponse | null) {
     console.debug("[queueState] statusUpdated", status)
     store.update((s) => {
@@ -296,6 +322,7 @@ function executingUpdated(promptID: PromptID, runningNodeID: ComfyNodeID | null)
                 entry.nodesRan.add(runningNodeID)
             }
             s.runningNodeID = runningNodeID;
+            s.runningPromptID = promptID;
 
             if (entry?.extraData?.workflowID) {
                 const workflow = workflowState.getWorkflow(entry.extraData.workflowID);
@@ -337,7 +364,9 @@ function executingUpdated(promptID: PromptID, runningNodeID: ComfyNodeID | null)
                 console.debug("[queueState] Could not find in pending! (executingUpdated)", promptID)
             }
             s.progress = null;
+            s.previewURL = null;
             s.runningNodeID = null;
+            s.runningPromptID = null;
             s.executingNodes.clear();
         }
         entry_ = entry;
@@ -362,7 +391,9 @@ function executionCached(promptID: PromptID, nodes: ComfyNodeID[]) {
         }
         s.isInterrupting = false; // TODO move to start
         s.progress = null;
+        s.previewURL = null;
         s.runningNodeID = null;
+        s.runningPromptID = null;
         s.executingNodes.clear();
         return s
     })
@@ -380,7 +411,9 @@ function executionError(error: ComfyExecutionError): CompletedQueueEntry | null
             console.error("[queueState] Could not find in pending! (executionError)", error.prompt_id)
         }
         s.progress = null;
+        s.previewURL = null;
         s.runningNodeID = null;
+        s.runningPromptID = null;
         s.executingNodes.clear();
         return s
     })
@@ -416,6 +449,7 @@ function executionStart(promptID: PromptID) {
         }
         s.isInterrupting = false;
         s.runningNodeID = null;
+        s.runningPromptID = promptID;
         s.executingNodes.clear();
         return s
     })
@@ -480,7 +514,9 @@ function queueCleared(type: QueueItemType) {
             s.queuePending.set([]);
             s.queueRemaining = 0;
             s.runningNodeID = null;
+            s.runningPromptID = null;
             s.progress = null;
+            s.previewURL = null;
             s.executingNodes.clear();
         }
         else {
@@ -535,6 +571,7 @@ const queueStateStore: WritableQueueStateStore =
     historyUpdated,
     statusUpdated,
     progressUpdated,
+    previewUpdated,
     executionStart,
     executingUpdated,
     executionCached,
diff --git a/src/lib/stores/uiState.ts b/src/lib/stores/uiState.ts
index e28bc90..2fb7e96 100644
--- a/src/lib/stores/uiState.ts
+++ b/src/lib/stores/uiState.ts
@@ -10,6 +10,7 @@ export type UIState = {
     autoAddUI: boolean,
     uiUnlocked: boolean,
     uiEditMode: UIEditMode,
+    hidePreviews: boolean,
 
     reconnecting: boolean,
     forceSaveUserState: boolean | null,
@@ -30,6 +31,7 @@ const store: Writable<UIState> = writable(
         autoAddUI: true,
         uiUnlocked: false,
         uiEditMode: "widgets",
+        hidePreviews: false,
 
         reconnecting: false,
         forceSaveUserState: null,
diff --git a/src/lib/widgets/GalleryWidget.svelte b/src/lib/widgets/GalleryWidget.svelte
index c820d34..0ce01ee 100644
--- a/src/lib/widgets/GalleryWidget.svelte
+++ b/src/lib/widgets/GalleryWidget.svelte
@@ -11,7 +11,11 @@
  import { clamp, comfyBoxImageToComfyURL, type ComfyBoxImageMetadata } from "$lib/utils";
  import { f7 } from "framework7-svelte";
  import type { ComfyGalleryNode } from "$lib/nodes/widgets";
-	import { showMobileLightbox } from "$lib/components/utils";
+ import { showMobileLightbox } from "$lib/components/utils";
+ import queueState from "$lib/stores/queueState";
+ import uiState from "$lib/stores/uiState";
+ import { loadImage } from "./utils";
+ import Spinner from "$lib/components/Spinner.svelte";
 
  export let widget: WidgetLayout | null = null;
  export let isMobile: boolean = false;
@@ -25,6 +29,47 @@
 
  $: widget && setNodeValue(widget);
 
+ function tagsMatch(tags: string[] | null): boolean {
+     if(tags != null && tags.length > 0)
+         return node.properties.tags.length > 0 && node.properties.tags.every(t => tags.includes(t));
+     else
+         return node.properties.tags.length === 0;
+ }
+
+ let previewURL: string | null;
+ let previewImage: HTMLImageElement | null = null;
+ let previewElem: HTMLImageElement | null = null
+ $: {
+     previewURL = $queueState.previewURL;
+
+     if (previewURL && $queueState.runningPromptID != null && !$uiState.hidePreviews && node.properties.showPreviews) {
+         const queueEntry = queueState.getQueueEntry($queueState.runningPromptID)
+         if (queueEntry != null) {
+             const tags = queueEntry.extraData?.extra_pnginfo?.comfyBoxPrompt?.subgraphs;
+             if (tagsMatch(tags)) {
+                 loadImage(previewURL).then((img) => {
+                     previewImage = img;
+                 })
+             }
+             else {
+                 previewImage = null;
+             }
+         }
+         else {
+             previewImage = null;
+         }
+     }
+     else {
+         previewImage = null;
+     }
+ }
+
+ function showPreview() {
+ }
+
+ function hidePreview() {
+ }
+
  function setNodeValue(widget: WidgetLayout) {
      if (widget) {
          node = widget.node as ComfyGalleryNode
@@ -34,6 +79,8 @@
          imageHeight = node.imageHeight
          selected_image = node.selectedImage;
          forceSelectImage = node.forceSelectImage;
+         previewURL = null;
+         previewImage = null;
 
          if ($nodeValue != null) {
              if (node.properties.index < 0 || node.properties.index >= $nodeValue.length) {
@@ -108,6 +155,11 @@
         <div class="wrapper comfy-gallery-widget gradio-gallery" style={widget.attrs.style || ""}>
             <Block variant="solid" padding={false}>
                 <div class="padding">
+                    {#if previewImage}
+                        <div class="comfy-gallery-preview" on:mouseover={hidePreview} on:mouseout={showPreview} >
+                            <img src={previewImage.src} bind:this={previewElem} on:mouseout={showPreview} />
+                        </div>
+                    {/if}
                     <Gallery
                         value={images}
                         label={widget.attrs.title}
@@ -153,6 +205,29 @@
              }
          }
      }
+
+     &:hover .comfy-gallery-preview {
+         opacity: 0%;
+     }
+ }
+
+ .comfy-gallery-preview {
+     position: absolute;
+     top: 0;
+     left: 0;
+     width: 100%;
+     height: 100%;
+     z-index: var(--layer-top);
+     pointer-events: none;
+     transition: opacity 0.1s linear;
+     opacity: 100%;
+
+     > img {
+         width: var(--size-full);
+         height: var(--size-full);
+         object-fit: contain;
+         border: 5px dashed var(--secondary-400);
+     }
  }
 
  .padding {