From bfbbad692f157180fd7966dae08bc26fd556cee3 Mon Sep 17 00:00:00 2001
From: space-nuko <24979496+space-nuko@users.noreply.github.com>
Date: Fri, 12 May 2023 17:11:57 -0500
Subject: [PATCH] Gallery fix

---
 src/lib/ImageViewer.ts               | 14 +++---
 src/lib/components/ComfyApp.svelte   |  1 -
 src/lib/components/ComfyApp.ts       | 10 ++--
 src/lib/components/ComfyQueue.svelte | 53 +++++++++++---------
 src/lib/notify.ts                    |  5 ++
 src/lib/stores/queueState.ts         | 73 +++++++++++++++++++---------
 6 files changed, 100 insertions(+), 56 deletions(-)

diff --git a/src/lib/ImageViewer.ts b/src/lib/ImageViewer.ts
index eeb68db..4d6837c 100644
--- a/src/lib/ImageViewer.ts
+++ b/src/lib/ImageViewer.ts
@@ -43,11 +43,12 @@ export class ImageViewer {
     }
 
     static selected_gallery_button(gallery: HTMLDivElement): [HTMLButtonElement | null, number] {
-        var allCurrentButtons = gallery.querySelectorAll('.preview > .thumbnails > .thumbnail-item.thumbnail-small.selected');
+        var allCurrentButtons = gallery.querySelectorAll('.preview > .thumbnails > .thumbnail-item.thumbnail-small');
+        console.log(allCurrentButtons)
         var visibleCurrentButton = null;
         let index = -1;
         allCurrentButtons.forEach((elem, i) => {
-            if (elem.parentElement.offsetParent) {
+            if (elem.parentElement.offsetParent && elem.classList.contains("selected")) {
                 visibleCurrentButton = elem;
                 index = i;
             }
@@ -104,11 +105,9 @@ export class ImageViewer {
         this.setModalImageSrc(selectedImageUrl)
 
         if (this.currentGallery) {
-            var galleryButtons = ImageViewer.all_gallery_buttons(this.currentGallery);
-            var [_currentButton, index] = ImageViewer.selected_gallery_button(this.currentGallery);
-
-            if (index != -1) {
-                const nextButton = galleryButtons[negmod((index + offset), galleryButtons.length)]
+            const galleryButtons = ImageViewer.all_gallery_buttons(this.currentGallery);
+            const nextButton = galleryButtons[this.selectedIndex];
+            if (nextButton) {
                 nextButton.click()
             }
         }
@@ -173,6 +172,7 @@ export class ImageViewer {
 
             let urls = ImageViewer.get_gallery_urls(galleryElem)
             const [_currentButton, index] = ImageViewer.selected_gallery_button(galleryElem)
+            console.warn("Gallery!", index, urls, galleryElem)
 
             this.showModal(urls, index, galleryElem)
             evt.stopPropagation();
diff --git a/src/lib/components/ComfyApp.svelte b/src/lib/components/ComfyApp.svelte
index 43cced4..68b2a9a 100644
--- a/src/lib/components/ComfyApp.svelte
+++ b/src/lib/components/ComfyApp.svelte
@@ -284,7 +284,6 @@
     <LightboxModal />
     <input bind:this={fileInput} id="comfy-file-input" type="file" accept=".json" on:change={loadWorkflow} />
 </div>
-
 <SvelteToast options={toastOptions} />
 
 <style lang="scss">
diff --git a/src/lib/components/ComfyApp.ts b/src/lib/components/ComfyApp.ts
index 600726b..b01fced 100644
--- a/src/lib/components/ComfyApp.ts
+++ b/src/lib/components/ComfyApp.ts
@@ -373,7 +373,7 @@ export default class ComfyApp {
 
             // Queue prompt using ctrl or command + enter
             if ((e.ctrlKey || e.metaKey) && (e.key === "Enter" || e.keyCode === 13 || e.keyCode === 10)) {
-                this.queuePrompt(e.shiftKey ? -1 : 0);
+                this.runDefaultQueueAction();
             }
         });
         window.addEventListener("keyup", (e) => {
@@ -525,7 +525,6 @@ export default class ComfyApp {
 
         if (get(layoutState).attrs.queuePromptButtonRunWorkflow) {
             this.queuePrompt(0, 1);
-            notify("Prompt queued.");
         }
     }
 
@@ -757,7 +756,13 @@ export default class ComfyApp {
 
                     try {
                         const response = await this.api.queuePrompt(request);
+
+                        // BUG: This can cause race conditions updating frontend state
+                        // since we don't have a prompt ID until the backend
+                        // returns!
                         promptID = response.promptID;
+                        queueState.afterQueued(promptID, num, p.output, extraData)
+
                         error = response.error;
                     } catch (error) {
                         error = error.toString();
@@ -779,7 +784,6 @@ export default class ComfyApp {
                     }
 
                     this.lCanvas.draw(true, true);
-                    queueState.afterQueued(promptID, num, p.output, extraData)
                 }
             }
         } finally {
diff --git a/src/lib/components/ComfyQueue.svelte b/src/lib/components/ComfyQueue.svelte
index 4540ef8..54d7625 100644
--- a/src/lib/components/ComfyQueue.svelte
+++ b/src/lib/components/ComfyQueue.svelte
@@ -6,10 +6,10 @@
  import { convertComfyOutputToComfyURL, convertFilenameToComfyURL, getNodeInfo } from "$lib/utils"
  import type { Writable } from "svelte/store";
  import type { QueueItemType } from "$lib/api";
-	import { ImageViewer } from "$lib/ImageViewer";
-	import { Button } from "@gradio/button";
-	import type ComfyApp from "./ComfyApp";
-	import { tick } from "svelte";
+ import { ImageViewer } from "$lib/ImageViewer";
+ import { Button } from "@gradio/button";
+ import type ComfyApp from "./ComfyApp";
+ import { tick } from "svelte";
 
  let queuePending: Writable<QueueEntry[]> | null = null;
  let queueRunning: Writable<QueueEntry[]> | null = null;
@@ -22,7 +22,7 @@
      date?: string,
      status: QueueEntryStatus | "pending" | "running",
      images?: string[], // URLs
-     error?: string
+     details?: string // shown in a tooltip on hover
  }
 
  $: if ($queueState) {
@@ -80,17 +80,18 @@
          submessage,
          date,
          status: "pending",
-         images
+         images,
      }
  }
 
  function convertCompletedEntry(entry: CompletedQueueEntry): QueueUIEntry {
      const result = convertEntry(entry.entry);
      result.status = entry.status;
-     result.error = entry.error;
 
-     if (result.status === "all_cached")
-         result.submessage = "(Execution was cached)"
+     if (entry.message)
+         result.submessage = entry.message
+     if (entry.error)
+         result.details = entry.error
 
      return result;
  }
@@ -105,10 +106,10 @@
  }
 
  async function updateFromHistory() {
-     _entries = $queueCompleted.map(convertCompletedEntry);
+     _entries = $queueCompleted.map(convertCompletedEntry).reverse();
      if (queueList) {
          await tick(); // Wait for list size to be recalculated
-         queueList.scroll({ top: queueList.scrollHeight })
+         queueList.scrollTo(0, 0);
      }
      console.warn("[ComfyQueue] BUILDHISTORY", _entries, $queueCompleted)
  }
@@ -137,7 +138,7 @@
 </script>
 
 <div class="queue">
-    <div class="queue-entries" bind:this={queueList}>
+    <div class="queue-entries {mode}-mode" bind:this={queueList}>
         {#if _entries.length > 0}
             {#each _entries as entry}
                 <div class="queue-entry {entry.status}">
@@ -156,13 +157,13 @@
                             {entry.submessage}
                         </div>
                     </div>
-                    <div class="queue-entry-rest">
-                        {#if entry.date != null}
-                            <span class="queue-entry-queued-at">
-                                {entry.date}
-                            </span>
-                        {/if}
-                    </div>
+                </div>
+                <div class="queue-entry-rest">
+                    {#if entry.date != null}
+                        <span class="queue-entry-queued-at">
+                            {entry.date}
+                        </span>
+                    {/if}
                 </div>
             {/each}
         {:else}
@@ -230,9 +231,17 @@
  }
 
  .queue-entries {
-     overflow-y: scroll;
      height: $queue-height;
      max-height: $queue-height;
+     display: flex;
+     overflow-y: auto;
+     flex-flow: column nowrap;
+
+     &.queue-mode > :first-child {
+         // elements stick to bottom in queue mode only
+         // next element in queue is on the bottom
+         margin-top: auto !important;
+     }
 
      > .queue-empty {
          display: flex;
@@ -328,11 +337,11 @@
 
  .queue-entry-queued-at {
      width: auto;
-     font-size: 10px;
+     font-size: 12px;
      position:absolute;
      right: 0px;
      bottom: 0px;
-     padding: 0.0rem 0.4rem;
+     padding: 0.4rem 0.6rem;
      color: var(--body-text-color);
  }
 
diff --git a/src/lib/notify.ts b/src/lib/notify.ts
index 531c0ba..7808ef7 100644
--- a/src/lib/notify.ts
+++ b/src/lib/notify.ts
@@ -45,6 +45,11 @@ function notifyToast(text: string, options: NotifyOptions) {
             '--toastBackground': 'var(--color-blue-500)',
         }
     }
+    else if (options.type === "warning") {
+        toastOptions.theme = {
+            '--toastBackground': 'var(--color-yellow-500)',
+        }
+    }
     else if (options.type === "error") {
         toastOptions.theme = {
             '--toastBackground': 'var(--color-red-500)',
diff --git a/src/lib/stores/queueState.ts b/src/lib/stores/queueState.ts
index d2b7f97..4e2c9a6 100644
--- a/src/lib/stores/queueState.ts
+++ b/src/lib/stores/queueState.ts
@@ -1,13 +1,14 @@
 import type { ComfyAPIHistoryEntry, ComfyAPIHistoryItem, ComfyAPIHistoryResponse, ComfyAPIQueueResponse, ComfyAPIStatusResponse, ComfyPromptExtraData, NodeID, PromptID } from "$lib/api";
-import type { Progress, SerializedPrompt, SerializedPromptInputsAll, SerializedPromptOutputs } from "$lib/components/ComfyApp";
+import type { Progress, SerializedPromptInputsAll, SerializedPromptOutputs } from "$lib/components/ComfyApp";
 import type { GalleryOutput } from "$lib/nodes/ComfyWidgetNodes";
+import notify from "$lib/notify";
 import { get, writable, type Writable } from "svelte/store";
 
 export type QueueItem = {
     name: string
 }
 
-export type QueueEntryStatus = "success" | "error" | "all_cached" | "unknown";
+export type QueueEntryStatus = "success" | "error" | "interrupted" | "all_cached" | "unknown";
 
 type QueueStateOps = {
     queueUpdated: (resp: ComfyAPIQueueResponse) => void,
@@ -35,11 +36,15 @@ export type QueueEntry = {
     /* Prompt outputs, collected while the prompt is still executing */
     outputs: SerializedPromptOutputs,
 
+    /* Nodes in of the workflow that have finished running so far. */
+    nodesRan: Set<NodeID>,
+    cachedNodes: Set<NodeID>
 }
 
 export type CompletedQueueEntry = {
     entry: QueueEntry,
     status: QueueEntryStatus,
+    message?: string,
     error?: string,
 }
 
@@ -72,7 +77,9 @@ function toQueueEntry(resp: ComfyAPIHistoryItem): QueueEntry {
         prompt,
         extraData,
         goodOutputs,
-        outputs: {}
+        outputs: {},
+        nodesRan: new Set(), // TODO can ComfyUI send this too?
+        cachedNodes: new Set()
     }
 }
 
@@ -122,27 +129,28 @@ function statusUpdated(status: ComfyAPIStatusResponse | null) {
     })
 }
 
-function findEntryInPending(promptID: PromptID): [number, QueueEntry, Writable<QueueEntry[]>] | null {
+function findEntryInPending(promptID: PromptID): [number, QueueEntry | null, Writable<QueueEntry[]> | null] {
     const state = get(store);
     let index = get(state.queuePending).findIndex(e => e.promptID === promptID)
-    if (index)
+    if (index !== -1)
         return [index, get(state.queuePending)[index], state.queuePending]
 
     index = get(state.queueRunning).findIndex(e => e.promptID === promptID)
-    if (index)
+    if (index !== -1)
         return [index, get(state.queueRunning)[index], state.queueRunning]
 
-    return null
+    return [-1, null, null]
 }
 
-function moveToCompleted(index: number, queue: Writable<QueueEntry[]>, status: QueueEntryStatus, error?: string) {
+function moveToCompleted(index: number, queue: Writable<QueueEntry[]>, status: QueueEntryStatus, message?: string, error?: string) {
     const state = get(store)
 
     const entry = get(queue)[index];
+    console.debug("[queueState] Move to completed", entry.promptID, index, status, message, error)
     entry.finishedAt = new Date() // Now
     queue.update(qp => { qp.splice(index, 1); return qp });
     state.queueCompleted.update(qc => {
-        const completed: CompletedQueueEntry = { entry, status, error }
+        const completed: CompletedQueueEntry = { entry, status, message, error }
         qc.push(completed)
         return qc
     })
@@ -150,21 +158,36 @@ function moveToCompleted(index: number, queue: Writable<QueueEntry[]>, status: Q
     store.set(state)
 }
 
-function executingUpdated(promptID: PromptID | null, runningNodeID: NodeID | null) {
+function executingUpdated(promptID: PromptID, runningNodeID: NodeID | null) {
     console.debug("[queueState] executingUpdated", promptID, runningNodeID)
     store.update((s) => {
         s.progress = null;
+
+        const [index, entry, queue] = findEntryInPending(promptID);
         if (runningNodeID != null) {
+            if (entry != null) {
+                entry.nodesRan.add(runningNodeID)
+            }
             s.runningNodeID = parseInt(runningNodeID);
         }
-        else if (promptID != null) {
+        else {
             // Prompt finished executing.
-            const [index, entry, queue] = findEntryInPending(promptID);
-            if (entry) {
-                moveToCompleted(index, queue, "success")
+            if (entry != null) {
+                const totalNodesInPrompt = Object.keys(entry.prompt).length
+                if (entry.cachedNodes.size >= Object.keys(entry.prompt).length) {
+                    notify("Prompt was cached, nothing to run.", { type: "warning" })
+                    moveToCompleted(index, queue, "all_cached", "(Execution was cached)");
+                }
+                else if (entry.nodesRan.size >= totalNodesInPrompt) {
+                    moveToCompleted(index, queue, "success")
+                }
+                else {
+                    notify("Interrupted prompt.")
+                    moveToCompleted(index, queue, "interrupted", `Interrupted after ${entry.nodesRan.size}/${totalNodesInPrompt} nodes`)
+                }
             }
             else {
-                console.error("[queueState] Could not find in pending! (executingUpdated)", promptID)
+                console.debug("[queueState] Could not find in pending! (executingUpdated)", promptID)
             }
             s.progress = null;
             s.runningNodeID = null;
@@ -177,13 +200,14 @@ function executionCached(promptID: PromptID, nodes: NodeID[]) {
     console.debug("[queueState] executionCached", promptID, nodes)
     store.update(s => {
         const [index, entry, queue] = findEntryInPending(promptID);
-        if (entry) {
-            if (nodes.length >= Object.keys(entry.prompt).length) {
-                moveToCompleted(index, queue, "all_cached");
+        if (entry != null) {
+            for (const nodeID of nodes) {
+                entry.nodesRan.add(nodeID);
+                entry.cachedNodes.add(nodeID);
             }
         }
         else {
-            console.error("[queueState] Could not find in pending! (executionCached)", promptID)
+            console.error("[queueState] Could not find in pending! (executionCached)", promptID, "pending", JSON.stringify(get(get(store).queuePending).map(p => p.promptID)), "running", JSON.stringify(get(get(store).queueRunning).map(p => p.promptID)))
         }
         s.progress = null;
         s.runningNodeID = null;
@@ -195,8 +219,8 @@ function executionError(promptID: PromptID, message: string) {
     console.debug("[queueState] executionError", promptID, message)
     store.update(s => {
         const [index, entry, queue] = findEntryInPending(promptID);
-        if (entry) {
-            moveToCompleted(index, queue, "error", message)
+        if (entry != null) {
+            moveToCompleted(index, queue, "error", "Error executing", message)
         }
         else {
             console.error("[queueState] Could not find in pending! (executionError)", promptID)
@@ -218,9 +242,12 @@ function afterQueued(promptID: PromptID, number: number, prompt: SerializedPromp
             prompt,
             extraData,
             goodOutputs: [],
-            outputs: {}
+            outputs: {},
+            nodesRan: new Set(),
+            cachedNodes: new Set()
         }
         s.queuePending.update(qp => { qp.push(entry); return qp })
+        console.debug("[queueState] ADD PROMPT", promptID)
         return s
     })
 }
@@ -229,7 +256,7 @@ function onExecuted(promptID: PromptID, nodeID: NodeID, output: GalleryOutput) {
     console.debug("[queueState] onExecuted", promptID, nodeID, output)
     store.update(s => {
         const [index, entry, queue] = findEntryInPending(promptID)
-        if (entry) {
+        if (entry != null) {
             entry.outputs[nodeID] = output;
             queue.set(get(queue))
         }