From e90773fdc8be582db6e0d05bd211d9158dcf7bd1 Mon Sep 17 00:00:00 2001
From: space-nuko <24979496+space-nuko@users.noreply.github.com>
Date: Mon, 22 May 2023 20:42:41 -0500
Subject: [PATCH] More control over how gallery selects images from events

---
 litegraph                                     |  2 +-
 src/lib/components/ComfyQueue.svelte          |  9 +++-
 src/lib/components/DropZone.svelte            |  1 +
 src/lib/components/ImageUpload.svelte         | 14 ++---
 .../components/gradio/gallery/Gallery.svelte  | 10 +++-
 .../modal/ReceiveOutputTargets.svelte         |  2 +-
 src/lib/nodes/ComfyNoChangeEvent.ts           |  2 +-
 .../actions/ComfyExecuteSubgraphAction.ts     |  7 +++
 .../nodes/actions/ComfyReceiveOutputNode.ts   |  2 +-
 src/lib/nodes/widgets/ComfyGalleryNode.ts     | 54 +++++++++++++++----
 src/lib/nodes/widgets/ComfyWidgetNode.ts      |  4 +-
 src/lib/stores/layoutStates.ts                |  8 +++
 src/lib/utils.ts                              |  2 +
 src/lib/widgets/GalleryWidget.svelte          |  4 ++
 14 files changed, 96 insertions(+), 25 deletions(-)

diff --git a/litegraph b/litegraph
index 3ce3a47..1a77c46 160000
--- a/litegraph
+++ b/litegraph
@@ -1 +1 @@
-Subproject commit 3ce3a478717580ad01778cf78d89285dba51dae4
+Subproject commit 1a77c461ad33546a41bda3d8d983dabadf9d588a
diff --git a/src/lib/components/ComfyQueue.svelte b/src/lib/components/ComfyQueue.svelte
index 4c18de4..31317b5 100644
--- a/src/lib/components/ComfyQueue.svelte
+++ b/src/lib/components/ComfyQueue.svelte
@@ -80,11 +80,18 @@
      updateFromQueue();
      changed = false;
  }
- else if (mode === "history" && (changed  || ($queueCompleted && $queueCompleted.length != _entries.length))) {
+ else if (mode === "history" && (changed || ($queueCompleted && $queueCompleted.length != _entries.length))) {
      updateFromHistory();
      changed = false;
  }
 
+ $: if (mode === "queue" && !$queuePending && !$queueRunning) {
+     _queuedEntries = []
+     _runningEntries = []
+     _entries = [];
+     changed = true
+ }
+
  async function deleteEntry(entry: QueueUIEntry, event: MouseEvent) {
      event.preventDefault();
      event.stopImmediatePropagation()
diff --git a/src/lib/components/DropZone.svelte b/src/lib/components/DropZone.svelte
index 3c9fad4..2acf441 100644
--- a/src/lib/components/DropZone.svelte
+++ b/src/lib/components/DropZone.svelte
@@ -53,6 +53,7 @@
     <div id="dropzone"
          class="dropzone"
          bind:this={dropZone}
+         on:pointerdown={hideDropZone}
          on:dragover={allowDrag}
          on:dragleave={hideDropZone}
          on:drop={handleDrop}
diff --git a/src/lib/components/ImageUpload.svelte b/src/lib/components/ImageUpload.svelte
index af0a680..b423d9a 100644
--- a/src/lib/components/ImageUpload.svelte
+++ b/src/lib/components/ImageUpload.svelte
@@ -200,12 +200,14 @@
             elem_id="comfy-image-upload-block"
             {elem_classes}
         >
-            <BlockLabel
-                label={label}
-                show_label={label != ""}
-                Icon={FileIcon}
-                float={label != ""}
-            />
+            {#if label != ""}
+                <BlockLabel
+                    label={label}
+                    show_label={label != ""}
+                    Icon={FileIcon}
+                    float={label != ""}
+                />
+            {/if}
             {#if value && value.length > 0 && !pending_upload}
                 {@const firstImage = value[0]}
                 <ModifyUpload on:clear={handle_clear} absolute />
diff --git a/src/lib/components/gradio/gallery/Gallery.svelte b/src/lib/components/gradio/gallery/Gallery.svelte
index 8108ca5..e5e6184 100644
--- a/src/lib/components/gradio/gallery/Gallery.svelte
+++ b/src/lib/components/gradio/gallery/Gallery.svelte
@@ -51,12 +51,20 @@
 	let prevValue: string[] | FileData[] | null = value;
 	export let selected_image: number | null = null;
 	let old_selected_image: number | null = null;
+	export let forceSelectImage: boolean | null = null;
 
 	$: if (prevValue !== value) {
 		// When value is falsy (clear button or first load),
 		// style.preview determines the selected image
 		if (was_reset) {
-			selected_image = style.preview && value?.length ? 0 : null;
+            let selectImage: boolean = Boolean(style.preview && value?.length)
+            let selectedIndex = 0;
+            if (forceSelectImage != null) {
+                selectImage = forceSelectImage
+                selectedIndex = selected_image;
+                forceSelectImage = null;
+            }
+			selected_image = selectImage ? selectedIndex : null;
 			was_reset = false;
 			// Otherwise we keep the selected_image the same if the
 			// gallery has at least as many elements as it did before
diff --git a/src/lib/components/modal/ReceiveOutputTargets.svelte b/src/lib/components/modal/ReceiveOutputTargets.svelte
index 9e1b28c..818892f 100644
--- a/src/lib/components/modal/ReceiveOutputTargets.svelte
+++ b/src/lib/components/modal/ReceiveOutputTargets.svelte
@@ -23,7 +23,7 @@
     {#if receiveTargets.length > 0}
         {#each receiveTargets as { workflow, targetNodes }}
             <Block>
-                <BlockTitle>{workflow.attrs.title}</BlockTitle>
+                <BlockTitle>Workflow: <b>{workflow.attrs.title}</b></BlockTitle>
                 {#each targetNodes as targetNode}
                     <Block>
                         <div class="target">
diff --git a/src/lib/nodes/ComfyNoChangeEvent.ts b/src/lib/nodes/ComfyNoChangeEvent.ts
index f385ddb..c03fcd7 100644
--- a/src/lib/nodes/ComfyNoChangeEvent.ts
+++ b/src/lib/nodes/ComfyNoChangeEvent.ts
@@ -17,7 +17,7 @@ export default class ComfyNoChangeEvent extends ComfyGraphNode {
         }
         else {
             param = {
-                value: param,
+                __widgetValue__: param,
                 noChangedEvent: true
             }
         }
diff --git a/src/lib/nodes/actions/ComfyExecuteSubgraphAction.ts b/src/lib/nodes/actions/ComfyExecuteSubgraphAction.ts
index 69b8b41..44fb92f 100644
--- a/src/lib/nodes/actions/ComfyExecuteSubgraphAction.ts
+++ b/src/lib/nodes/actions/ComfyExecuteSubgraphAction.ts
@@ -31,6 +31,13 @@ export default class ComfyExecuteSubgraphAction extends ComfyGraphNode {
             this.setProperty("tag", tag)
     }
 
+    override getTitle(): string {
+        if (this.flags.collapsed) {
+            return "Execute: " + String(this.properties.targetTag);
+        }
+        return this.title;
+    }
+
     override onAction(action: any, param: any) {
         const tag = this.getInputData(1) || this.properties.targetTag;
 
diff --git a/src/lib/nodes/actions/ComfyReceiveOutputNode.ts b/src/lib/nodes/actions/ComfyReceiveOutputNode.ts
index e5bc654..8539899 100644
--- a/src/lib/nodes/actions/ComfyReceiveOutputNode.ts
+++ b/src/lib/nodes/actions/ComfyReceiveOutputNode.ts
@@ -27,7 +27,7 @@ export default class ComfyReceiveOutputNode extends ComfyGraphNode {
 
     static slotLayout: SlotLayout = {
         outputs: [
-            { name: "received", type: BuiltInSlotType.EVENT }
+            { name: "received", type: BuiltInSlotType.EVENT, options: { color_off: "fuchsia", color_on: "fuchsia" } }
         ]
     }
 
diff --git a/src/lib/nodes/widgets/ComfyGalleryNode.ts b/src/lib/nodes/widgets/ComfyGalleryNode.ts
index 59c4ab1..2d2e579 100644
--- a/src/lib/nodes/widgets/ComfyGalleryNode.ts
+++ b/src/lib/nodes/widgets/ComfyGalleryNode.ts
@@ -1,5 +1,5 @@
 import { parseWhateverIntoImageMetadata, type ComfyBoxImageMetadata, type ComfyUploadImageType } from "$lib/utils";
-import { BuiltInSlotType, LiteGraph, type IComboWidget, type ITextWidget, type PropertyLayout, type SlotLayout, type INumberWidget } from "@litegraph-ts/core";
+import { BuiltInSlotType, LiteGraph, type IComboWidget, type ITextWidget, type PropertyLayout, type SlotLayout, type INumberWidget, clamp } from "@litegraph-ts/core";
 import { get, writable, type Writable } from "svelte/store";
 
 import GalleryWidget from "$lib/widgets/GalleryWidget.svelte";
@@ -9,6 +9,7 @@ import ComfyWidgetNode from "./ComfyWidgetNode";
 export interface ComfyGalleryProperties extends ComfyWidgetProperties {
     index: number | null,
     updateMode: "replace" | "append",
+    autoSelectOnUpdate: boolean
 }
 
 export default class ComfyGalleryNode extends ComfyWidgetNode<ComfyBoxImageMetadata[]> {
@@ -17,6 +18,7 @@ export default class ComfyGalleryNode extends ComfyWidgetNode<ComfyBoxImageMetad
         defaultValue: [],
         index: 0,
         updateMode: "replace",
+        autoSelectOnUpdate: true
     }
 
     static slotLayout: SlotLayout = {
@@ -42,17 +44,18 @@ export default class ComfyGalleryNode extends ComfyWidgetNode<ComfyBoxImageMetad
 
     selectedFilename: string | null = null;
 
-    selectedIndexWidget: INumberWidget;
+    selectedIndexWidget: ITextWidget;
     modeWidget: IComboWidget;
 
     imageWidth: Writable<number> = writable(0);
     imageHeight: Writable<number> = writable(0);
 
     selectedImage: Writable<number | null> = writable(null);
+    forceSelectImage: Writable<boolean | null> = writable(null);
 
     constructor(name?: string) {
         super(name, [])
-        this.selectedIndexWidget = this.addWidget("number", "Selected", get(this.selectedImage))
+        this.selectedIndexWidget = this.addWidget("text", "Selected", String(get(this.selectedImage)))
         this.selectedIndexWidget.disabled = true;
         this.modeWidget = this.addWidget("combo", "Mode", this.properties.updateMode, null, { property: "updateMode", values: ["replace", "append"] })
     }
@@ -69,7 +72,7 @@ export default class ComfyGalleryNode extends ComfyWidgetNode<ComfyBoxImageMetad
         this.setOutputData(0, value)
         this.setOutputData(1, index)
 
-        this.selectedIndexWidget.value = index;
+        this.selectedIndexWidget.value = String(index);
 
         if (index != null && value && value[index] != null) {
             const image = value[index];
@@ -86,29 +89,58 @@ export default class ComfyGalleryNode extends ComfyWidgetNode<ComfyBoxImageMetad
         return `Images: ${value?.length || 0}`
     }
 
+    override setValue(value: ComfyBoxImageMetadata[], noChangedEvent: boolean = false) {
+        super.setValue(value, noChangedEvent)
+
+        const newIndex = this._newSelectedIndex;
+        this._newSelectedIndex = null;
+
+        if (newIndex != null) {
+            this.selectedImage.set(newIndex)
+            this.forceSelectImage.set(true)
+        }
+    }
+
+    private _newSelectedIndex: number | null = null;
+
     override parseValue(param: any): ComfyBoxImageMetadata[] {
         if (param == null)
             return []
 
+        let updateMode = this.properties.updateMode;
+        let selectedIndex: number | null = null;
+        if (typeof param === "object" && param.galleryImages != null) {
+            selectedIndex = param.selectedIndex;
+            updateMode = param.updateMode || updateMode;
+            param = param.galleryImages;
+        }
+
         const meta = parseWhateverIntoImageMetadata(param) || [];
 
         console.debug("[ComfyGalleryNode] Received output!", param)
 
-        if (this.properties.updateMode === "append") {
+        if (updateMode === "append") {
             const currentValue = get(this.value)
-            if (meta.length > 0)
-                this.selectedImage.set(currentValue.length);
+            if (meta.length > 0 && (selectedIndex != null || this.properties.autoSelectOnUpdate)) {
+                let index = selectedIndex
+                if (index == null)
+                    index = 0;
+                this._newSelectedIndex = clamp(currentValue.length + index, 0, currentValue.length + meta.length - 1)
+            }
             return currentValue.concat(meta)
         }
         else {
             this.notifyPropsChanged();
+            if (meta.length > 0 && (selectedIndex != null || this.properties.autoSelectOnUpdate)) {
+                let index = selectedIndex;
+                if (index == null)
+                    index = get(this.selectedImage)
+                if (index != null)
+                    this._newSelectedIndex = clamp(index, 0, meta.length - 1)
+            }
             return meta;
         }
     }
-
-    override setValue(value: any, noChangedEvent: boolean = false) {
-        super.setValue(value, noChangedEvent)
-    }
 }
 
 LiteGraph.registerNodeType({
diff --git a/src/lib/nodes/widgets/ComfyWidgetNode.ts b/src/lib/nodes/widgets/ComfyWidgetNode.ts
index 277cfd8..33ef5d6 100644
--- a/src/lib/nodes/widgets/ComfyWidgetNode.ts
+++ b/src/lib/nodes/widgets/ComfyWidgetNode.ts
@@ -239,8 +239,8 @@ export default abstract class ComfyWidgetNode<T = any> extends ComfyGraphNode {
         if (action === this.storeActionName) {
             let noChangedEvent = false;
             let value = param;
-            if (param != null && typeof param === "object" && "value" in param) {
-                value = param.value
+            if (param != null && typeof param === "object" && "__widgetValue__" in param) {
+                value = param.__widgetValue__
                 if ("noChangedEvent" in param)
                     noChangedEvent = Boolean(param.noChangedEvent)
             }
diff --git a/src/lib/stores/layoutStates.ts b/src/lib/stores/layoutStates.ts
index 652edcb..5831094 100644
--- a/src/lib/stores/layoutStates.ts
+++ b/src/lib/stores/layoutStates.ts
@@ -566,6 +566,14 @@ const ALL_ATTRIBUTES: AttributesSpecList = [
                 values: ["replace", "append"],
                 defaultValue: "replace"
             },
+            {
+                name: "autoSelectOnUpdate",
+                type: "boolean",
+                location: "nodeProps",
+                editable: true,
+                validNodeTypes: ["ui/gallery"],
+                defaultValue: true
+            },
 
             // Radio
             {
diff --git a/src/lib/utils.ts b/src/lib/utils.ts
index d9eb7cb..18467ea 100644
--- a/src/lib/utils.ts
+++ b/src/lib/utils.ts
@@ -9,6 +9,8 @@ import workflowState from "./stores/workflowState";
 import { ImageViewer } from "./ImageViewer";
 
 export function clamp(n: number, min: number, max: number): number {
+    if (max <= min)
+        return min;
     return Math.min(Math.max(n, min), max)
 }
 
diff --git a/src/lib/widgets/GalleryWidget.svelte b/src/lib/widgets/GalleryWidget.svelte
index 600a7b0..7873448 100644
--- a/src/lib/widgets/GalleryWidget.svelte
+++ b/src/lib/widgets/GalleryWidget.svelte
@@ -20,6 +20,7 @@
  let imageWidth: Writable<number> = writable(0);
  let imageHeight: Writable<number> = writable(0);
  let selected_image: Writable<number | null> = writable(null);
+ let forceSelectImage: Writable<boolean | null> = writable(null);
 
  $: widget && setNodeValue(widget);
 
@@ -31,6 +32,7 @@
          imageWidth = node.imageWidth
          imageHeight = node.imageHeight
          selected_image = node.selectedImage;
+         forceSelectImage = node.forceSelectImage;
 
          if ($nodeValue != null) {
              if (node.properties.index < 0 || node.properties.index >= $nodeValue.length) {
@@ -43,6 +45,7 @@
  let style: Styles = {
      grid_cols: [isMobile ? 2 : 3],
      object_fit: "cover",
+     // preview: true
  }
  let element: HTMLDivElement;
 
@@ -134,6 +137,7 @@
                         bind:imageWidth={$imageWidth}
                         bind:imageHeight={$imageHeight}
                         bind:selected_image={$selected_image}
+                        bind:forceSelectImage={$forceSelectImage}
                     />
                 </div>
             </Block>