Tha gallery widget

package.json

@@ -39,6 +39,8 @@
     "@gradio/form": "workspace:*",
     "@gradio/icons": "workspace:*",
     "@gradio/theme": "workspace:*",
+    "@gradio/gallery": "workspace:*",
+    "@gradio/upload": "workspace:*",
     "@gradio/utils": "workspace:*",
     "@litegraph-ts/core": "workspace:*",
     "@litegraph-ts/nodes-basic": "workspace:*",

pnpm-lock.yaml

@@ -8,8 +8,10 @@ importers:
       '@gradio/button': workspace:*
       '@gradio/client': ^0.0.1
       '@gradio/form': workspace:*
+      '@gradio/gallery': workspace:*
       '@gradio/icons': workspace:*
       '@gradio/theme': workspace:*
+      '@gradio/upload': workspace:*
       '@gradio/utils': workspace:*
       '@litegraph-ts/core': workspace:*
       '@litegraph-ts/nodes-basic': workspace:*
@@ -41,8 +43,10 @@ importers:
       '@gradio/button': link:gradio/js/button
       '@gradio/client': 0.0.1
       '@gradio/form': link:gradio/js/form
+      '@gradio/gallery': link:gradio/js/gallery
       '@gradio/icons': link:gradio/js/icons
       '@gradio/theme': link:gradio/js/theme
+      '@gradio/upload': link:gradio/js/upload
       '@gradio/utils': link:gradio/js/utils
       '@litegraph-ts/core': link:litegraph/packages/core
       '@litegraph-ts/nodes-basic': link:litegraph/packages/nodes-basic

src/lib/ComfyGraphNode.ts

@@ -1,5 +0,0 @@
-import { LGraphNode } from "@litegraph-ts/core";
-
-export default class ComfyGraphNode extends LGraphNode {
-    onExecuted?(output: any): void;
-}

src/lib/api.ts

@@ -60,14 +60,14 @@ export default class ComfyAPI extends EventTarget {
         let opened = false;
         let existingSession = sessionStorage["Comfy.SessionId"] || "";
         if (existingSession) {
-            existingSession = "/" + existingSession;
+            existingSession = "?clientId=" + existingSession;
         }
 
         const hostname = this.hostname || location.hostname;
         const port = this.port || location.port;
 
         this.socket = new WebSocket(
-            `ws${window.location.protocol === "https:" ? "s" : ""}://${hostname}:${port}/ws?clientId=${existingSession}`
+            `ws${window.location.protocol === "https:" ? "s" : ""}://${hostname}:${port}/ws${existingSession}`
         );
 
         this.socket.addEventListener("open", () => {

src/lib/components/ComfyApp.svelte

@@ -52,6 +52,10 @@
  let graphResizeTimer: typeof Timer = -1;
 
  function doAutosave(graph: LGraph): void {
+     // We will merge in the state of the frontend instead of what's inside the
+     // LGraph structure.
+     const frontendState = get(widgetState);
+
      const serializedGraph = graph.serialize()
      const serializedPaneOrder = uiPane.serialize()
 

src/lib/components/ComfyApp.ts

@@ -12,6 +12,9 @@ import * as basic from "@litegraph-ts/nodes-basic"
 import * as nodes from "$lib/nodes/index"
 import ComfyGraphCanvas from "$lib/ComfyGraphCanvas";
 import type ComfyGraphNode from "$lib/ComfyGraphNode";
+import type { WidgetStateStore, WidgetUIState } from "$lib/stores/widgetState";
+import * as widgets from "$lib/widgets/index"
+import type ComfyWidget from "$lib/widgets/ComfyWidget";
 
 LiteGraph.catch_exceptions = false;
 
@@ -87,10 +90,13 @@ export default class ComfyApp {
 
         LiteGraph.release_link_on_empty_shows_menu = true;
         LiteGraph.alt_drag_do_clone_nodes = true;
+        LiteGraph.ignore_all_widget_events = true;
 
         this.lGraph.start();
 
         // await this.#invokeExtensionsAsync("init");
+        this.registerNodeTypeOverrides();
+        this.registerWidgetTypeOverrides();
         await this.registerNodes();
 
         // Load previous workflow
@@ -120,6 +126,7 @@ export default class ComfyApp {
         this.addPasteHandler();
         this.addKeyboardHandler();
 
+
         // await this.#invokeExtensionsAsync("setup");
 
         // Ensure the canvas fills the window
@@ -187,6 +194,17 @@ export default class ComfyApp {
         this.lCanvas.onClear = this.canvasOnClear.bind(this);
     }
 
+    static node_type_overrides: Record<string, typeof ComfyGraphNode> = {}
+    static widget_type_overrides: Record<string, Function> = {}
+
+    private registerNodeTypeOverrides() {
+        ComfyApp.node_type_overrides["SaveImage"] = nodes.ComfySaveImageNode;
+    }
+
+    private registerWidgetTypeOverrides() {
+        ComfyApp.widget_type_overrides["comfy/gallery"] = widgets.ComfyGalleryWidget_Svelte;
+    }
+
     private async registerNodes() {
         const app = this;
 
@@ -206,7 +224,12 @@ export default class ComfyApp {
         for (const nodeId in defs) {
             const nodeData = defs[nodeId];
 
-            const ctor = class extends LGraphNode {
+            const typeOverride = ComfyApp.node_type_overrides[nodeId]
+            if (typeOverride)
+                console.debug("Attaching custom type to received node:", nodeId, typeOverride)
+            const baseClass: typeof LGraphNode = typeOverride || LGraphNode;
+
+            const ctor = class extends baseClass {
                 constructor(title?: string) {
                     super(title);
                     this.type = nodeId; // XXX: workaround dependency in LGraphNode.addInput()
@@ -432,13 +455,13 @@ export default class ComfyApp {
      * Converts the current graph workflow for sending to the API
      * @returns The workflow and node links
      */
-    async graphToPrompt(frontendState: Record<number, any[]> = {}) {
+    async graphToPrompt(frontendState: WidgetStateStore = {}) {
         const workflow = this.lGraph.serialize();
 
         const output = {};
         // Process nodes in order of execution
         for (const node of this.lGraph.computeExecutionOrder<ComfyGraphNodeExecutable>(false, null)) {
-            const fromFrontend = frontendState[node.id];
+            const fromFrontend: WidgetUIState[] | null = frontendState[node.id];
 
             const n = workflow.nodes.find((n) => n.id === node.id);
 
@@ -465,7 +488,7 @@ export default class ComfyApp {
                     const widget = widgets[i];
                     if (!widget.options || widget.options.serialize !== false) {
                         let value = widget.serializeValue ? await widget.serializeValue(n, i) : widget.value;
-                        if (fromFrontend) {
+                        if (fromFrontend && !fromFrontend[i].isVirtual) {
                             value = fromFrontend[i].value;
                         }
                         inputs[widget.name] = value
@@ -514,7 +537,7 @@ export default class ComfyApp {
         return { workflow, output };
     }
 
-    async queuePrompt(num: number, batchCount: number = 1, frontendState: Record<number, any[]> = {}) {
+    async queuePrompt(num: number, batchCount: number = 1, frontendState: WidgetStateStore = {}) {
         this.queueItems.push({ num, batchCount });
 
         // Only have one action process the items so each one gets a unique seed correctly

src/lib/components/ComfyPane.svelte

@@ -2,10 +2,10 @@
  import { onDestroy } from "svelte";
  import { Block, BlockTitle } from "@gradio/atoms";
  import { Move } from 'radix-icons-svelte';
- import ComboWidget from "./widgets/ComboWidget.svelte";
+ import ComboWidget from "$lib/widgets/ComboWidget.svelte";
- import RangeWidget from "./widgets/RangeWidget.svelte";
+ import RangeWidget from "$lib/widgets/RangeWidget.svelte";
- import TextWidget from "./widgets/TextWidget.svelte";
+ import TextWidget from "$lib/widgets/TextWidget.svelte";
- import widgetState from "$lib/stores/widgetState";
+ import widgetState, { type WidgetUIState } from "$lib/stores/widgetState";
 
  import { dndzone, SHADOW_ITEM_MARKER_PROPERTY_NAME } from 'svelte-dnd-action';
 
@@ -13,6 +13,7 @@
  // notice - fade in works fine but don't add svelte's fade-out (known issue)
  import {cubicIn} from 'svelte/easing';
  import { flip } from 'svelte/animate';
+	import ComfyApp from "./ComfyApp";
 
  export let dragItems = [];
  let dragDisabled = true;
@@ -40,6 +41,30 @@
  });
 
  onDestroy(unsubscribe);
+
+ function getComponentForWidgetState(item: WidgetUIState): any {
+     let ctor: any = null;
+
+     // custom widgets with TypeScript sources
+     if (item.isVirtual) {
+         let override = ComfyApp.widget_type_overrides[item.widget.type]
+         if (override) {
+             return override;
+         }
+     }
+
+     // litegraph.ts built-in widgets
+     switch (item.widget.type) {
+             case "combo":
+             return ComboWidget;
+             case "number":
+             return RangeWidget;
+             case "text":
+             return TextWidget;
+     }
+
+     return null;
+ }
 </script>
 
 
@@ -59,14 +84,8 @@
                 <label for={id}>
                     <BlockTitle>{node.title}</BlockTitle>
                 </label>
-                {#each $widgetState[id] as item, i}
+                {#each $widgetState[id] as item}
-                    {#if item.widget.type == "combo"}
+                    <svelte:component this={getComponentForWidgetState(item)} {item} />
-                        <ComboWidget {item} />
-                    {:else if item.widget.type == "number"}
-                        <RangeWidget {item} />
-                    {:else if item.widget.type == "text"}
-                        <TextWidget {item} />
-                    {/if}
                     {#if dragItem[SHADOW_ITEM_MARKER_PROPERTY_NAME]}
                         <div in:fade={{duration:200, easing: cubicIn}} class='drag-item-shadow'/>
                     {/if}

src/lib/nodes/ComfyGraphNode.ts

@@ -1,5 +1,14 @@
+import type ComfyWidget from "$lib/components/widgets/ComfyWidget";
 import { LGraphNode } from "@litegraph-ts/core";
 
 export default class ComfyGraphNode extends LGraphNode {
     isVirtualNode: boolean = false;
+
+    /*
+     * Widgets that aren't a part of the graph, but are used for rendering
+     * purposes only.
+     */
+    virtualWidgets: ComfyWidget[] = [];
+
+    onExecuted?(output: any): void;
 }

src/lib/nodes/ComfySaveImageNode.ts

@@ -0,0 +1,32 @@
+import ComfyGalleryWidget, { type ComfyGalleryEntry } from "$lib/widgets/ComfyGalleryWidget";
+import ComfyGraphNode from "./ComfyGraphNode";
+
+export type ComfyImageResult = {
+    filename: string,
+    subfolder: string,
+    type: "output" | "temp"
+}
+export type ComfyImageExecOutput = {
+    images: ComfyImageResult[]
+}
+
+export default class ComfySaveImageNode extends ComfyGraphNode {
+    private _imageResults: Array<ComfyImageResult> = [];
+    private _galleryWidget: ComfyGalleryWidget;
+
+    constructor(title?: any) {
+        super(title)
+        this._galleryWidget = new ComfyGalleryWidget("Images", this._imageResults, this);
+        this.virtualWidgets.push(this._galleryWidget)
+    }
+
+    override onExecuted(output: ComfyImageExecOutput) {
+        this._imageResults = Array.from(output.images); // TODO append?
+        const galleryItems = this._imageResults.map(r => {
+            // TODO
+            let entry: ComfyGalleryEntry = ["http://localhost:8188/" + r.type + "/" + r.filename, null]
+            return entry
+        });
+        this._galleryWidget.setValue(galleryItems)
+    }
+}

src/lib/nodes/index.ts

@@ -1 +1,2 @@
 export { default as ComfyReroute } from "./ComfyReroute"
+export { default as ComfySaveImageNode } from "./ComfySaveImageNode"

src/lib/stores/widgetState.ts

@@ -1,11 +1,14 @@
 import { writable, get } from 'svelte/store';
 import type { LGraph, LGraphNode, IWidget } from "@litegraph-ts/core";
 import type { Readable, Writable } from 'svelte/store';
+import type ComfyGraphNode from '$lib/nodes/ComfyGraphNode';
+import type ComfyWidget from '$lib/widgets/ComfyWidget';
 
 export type WidgetUIState = {
     node: LGraphNode,
     widget: IWidget,
-    value: any
+    value: any,
+    isVirtual: boolean
 }
 
 type NodeID = number;
@@ -14,9 +17,12 @@ type WidgetStateOps = {
     nodeAdded: (node: LGraphNode) => void,
     nodeRemoved: (node: LGraphNode) => void,
     configureFinished: (graph: LGraph) => void,
+    widgetStateChanged: (widget: ComfyWidget<any, any>) => void,
     clear: () => void,
 }
-type WidgetStateStore = Writable<Record<NodeID, WidgetUIState[]>> & WidgetStateOps;
+
+export type WidgetStateStore = Record<NodeID, WidgetUIState[]>;
+type WritableWidgetStateStore = Writable<WidgetStateStore> & WidgetStateOps;
 
 const store: Writable<Record<NodeID, WidgetUIState[]>> = writable({})
 
@@ -25,31 +31,59 @@ function clear() {
 }
 
 function nodeAdded(node: LGraphNode) {
-    if (node.widgets) {
     let state = get(store)
+
+    if (node.widgets) {
         for (const widget of node.widgets) {
             if (!state[node.id])
                 state[node.id] = []
-            state[node.id].push({ node, widget, value: widget.value })
+            state[node.id].push({ node, widget, value: widget.value, isVirtual: false })
+        }
+    }
+
+    if ("virtualWidgets" in node) {
+        const comfyNode = node as ComfyGraphNode;
+        for (const widget of comfyNode.virtualWidgets) {
+            if (!state[comfyNode.id])
+                state[comfyNode.id] = []
+            state[comfyNode.id].push({ node, widget, value: widget.value, isVirtual: true })
+        }
     }
 
     store.set(state);
 }
-}
 
 function nodeRemoved(node: LGraphNode) {
-    let state = get(store)
+    const state = get(store)
     delete state[node.id]
     store.set(state)
 }
 
+function widgetStateChanged(widget: ComfyWidget<any, any>) {
+    const state = get(store)
+    const entries = state[widget.node.id]
+    if (entries) {
+        let widgetState = entries.find(e => e.widget === widget);
+        if (widgetState) {
+            widgetState.value = widget.value;
+            store.set(state);
+        }
+        else {
+            console.error("Widget state changed and node was found, but widget was not found in state!", widget, widget.node, entries)
+        }
+    }
+    else {
+        console.error("Widget state changed but node was not found in state!", widget, widget.node)
+    }
+}
+
 function configureFinished(graph: LGraph) {
     let state = get(store);
 
     for (const node of graph.computeExecutionOrder(false, null)) {
         if (node.widgets_values) {
             for (const [i, value] of node.widgets_values.entries()) {
-                if (i < state[node.id].length) {
+                if (i < state[node.id].length && !state[node.id][i].isVirtual) { // Virtual widgets always come after real widgets
                     state[node.id][i].value = value;
                 }
                 else {
@@ -63,11 +97,13 @@ function configureFinished(graph: LGraph) {
     store.set(state)
 }
 
-export default
+const widgetStateStore: WritableWidgetStateStore =
 {
     ...store,
     nodeAdded,
     nodeRemoved,
+    widgetStateChanged,
     configureFinished,
     clear
-    } as WidgetStateStore;
+}
+export default widgetStateStore;

src/lib/widgets/ComfyGalleryWidget.svelte

@@ -0,0 +1,28 @@
+<script lang="ts">
+ import type { WidgetUIState } from "$lib/stores/widgetState";
+	import { Block } from "@gradio/atoms";
+ import { Gallery } from "@gradio/gallery";
+ export let item: WidgetUIState | null = null;
+</script>
+
+<div class="wrapper comfy-gallery-widget">
+    {#if item}
+        <Block variant="solid" padding={false}>
+            <Gallery
+                bind:value={item.value}
+                label={item.widget.name}
+                show_label={true}
+                root={""}
+                root_url={""}
+                on:select
+            />
+        </Block>
+    {/if}
+</div>
+
+<style>
+ .wrapper {
+     padding: 2px;
+     width: 100%;
+ }
+</style>

src/lib/widgets/ComfyGalleryWidget.ts

@@ -0,0 +1,12 @@
+import type { WidgetPanelOptions } from "@litegraph-ts/core";
+import ComfyWidget from "./ComfyWidget";
+import type { ComfyImageResult } from "$lib/nodes/ComfySaveImageNode";
+
+export type ComfyGalleryEntry = [string, string | null]; // <img> src and alt/title, gradio format
+
+export interface ComfyGalleryWidgetOptions extends WidgetPanelOptions {
+}
+
+export default class ComfyGalleryWidget extends ComfyWidget<ComfyGalleryWidgetOptions, ComfyGalleryEntry[]> {
+    override type = "comfy/gallery";
+}

src/lib/widgets/ComfyWidget.ts

@@ -0,0 +1,39 @@
+import type ComfyGraphNode from "$lib/nodes/ComfyGraphNode";
+import type { IWidget, LGraphNode, SerializedLGraphNode, Vector2, WidgetCallback, WidgetTypes } from "@litegraph-ts/core";
+import widgetState from "$lib/stores/widgetState";
+
+export default abstract class ComfyWidget<T = any, V = any> implements IWidget<T, V> {
+    name: string;
+    value: V;
+    node: ComfyGraphNode;
+
+    constructor(name: string, value: V, node: ComfyGraphNode) {
+        this.name = name;
+        this.value = value
+        this.node = node;
+    }
+
+    options?: T;
+    type?: WidgetTypes | string | any;
+    y?: number;
+    property?: string;
+    last_y?: number;
+    width?: number;
+    clicked?: boolean;
+    marker?: boolean;
+    disabled?: boolean;
+    callback?: WidgetCallback<this>;
+
+    setValue(value: V) {
+        this.value = value;
+        widgetState.widgetStateChanged(this);
+    }
+
+    draw?(ctx: CanvasRenderingContext2D, node: LGraphNode, width: number, posY: number, height: number): void;
+
+    mouse?(event: MouseEvent, pos: Vector2, node: LGraphNode): boolean;
+
+    computeSize?(width: number): [number, number];
+
+    serializeValue?(serialized: SerializedLGraphNode<LGraphNode>, slot: number): Promise<any>;
+}

src/lib/widgets/index.ts

@@ -0,0 +1,2 @@
+export { default as ComfyGalleryWidget } from "./ComfyGalleryWidget"
+export { default as ComfyGalleryWidget_Svelte } from "./ComfyGalleryWidget.svelte"