From 71a3abf5180743923748bad16aed86129749bcf2 Mon Sep 17 00:00:00 2001
From: space-nuko <24979496+space-nuko@users.noreply.github.com>
Date: Thu, 4 May 2023 11:42:21 -0500
Subject: [PATCH] Auto add new widgets for backend nodes

---
 litegraph                                 |  2 +-
 src/lib/ComfyGraph.ts                     | 92 +++++++++++++++++++++++
 src/lib/GraphSync.ts                      | 11 ++-
 src/lib/IComfyInputSlot.ts                |  6 +-
 src/lib/api.ts                            |  4 +-
 src/lib/components/ComfyApp.svelte        |  4 +-
 src/lib/components/ComfyApp.ts            | 80 +++-----------------
 src/lib/components/WidgetContainer.svelte |  4 +-
 src/lib/nodes/ComfyWidgetNodes.ts         |  5 --
 src/lib/stores/layoutState.ts             | 31 ++++++--
 src/lib/stores/uiState.ts                 |  2 +
 src/lib/widgets.ts                        |  9 ++-
 src/lib/widgets/ComboWidget.svelte        |  1 +
 13 files changed, 150 insertions(+), 101 deletions(-)
 create mode 100644 src/lib/ComfyGraph.ts

diff --git a/litegraph b/litegraph
index ac3621c..f09cba0 160000
--- a/litegraph
+++ b/litegraph
@@ -1 +1 @@
-Subproject commit ac3621c4bdb9b5dd7a9f0a4b0b07599613dee42a
+Subproject commit f09cba0804836d506c98235fb680ec7024e3e69a
diff --git a/src/lib/ComfyGraph.ts b/src/lib/ComfyGraph.ts
new file mode 100644
index 0000000..1c7d6eb
--- /dev/null
+++ b/src/lib/ComfyGraph.ts
@@ -0,0 +1,92 @@
+import { LConnectionKind, LGraph, LGraphNode, type INodeSlot, type SlotIndex, LiteGraph, getStaticProperty, type LGraphAddNodeOptions } from "@litegraph-ts/core";
+import GraphSync from "./GraphSync";
+import EventEmitter from "events";
+import type TypedEmitter from "typed-emitter";
+import layoutState from "./stores/layoutState";
+import uiState from "./stores/uiState";
+import { get } from "svelte/store";
+import type ComfyGraphNode from "./nodes/ComfyGraphNode";
+import type IComfyInputSlot from "./IComfyInputSlot";
+
+type ComfyGraphEvents = {
+    configured: (graph: LGraph) => void
+    nodeAdded: (node: LGraphNode) => void
+    nodeRemoved: (node: LGraphNode) => void
+    nodeConnectionChanged: (kind: LConnectionKind, node: LGraphNode, slot: SlotIndex, targetNode: LGraphNode, targetSlot: SlotIndex) => void
+    cleared: () => void
+    beforeChange: (graph: LGraph, param: any) => void
+    afterChange: (graph: LGraph, param: any) => void
+}
+
+export default class ComfyGraph extends LGraph {
+    graphSync: GraphSync;
+    eventBus: TypedEmitter<ComfyGraphEvents> = new EventEmitter() as TypedEmitter<ComfyGraphEvents>;
+
+    constructor() {
+        super();
+        this.graphSync = new GraphSync(this)
+    }
+
+    override onConfigure() {
+        console.debug("Configured");
+        this.eventBus.emit("configured", this);
+    }
+
+    override onBeforeChange(graph: LGraph, info: any) {
+        console.debug("BeforeChange", info);
+        this.eventBus.emit("beforeChange", graph, info);
+    }
+
+    override onAfterChange(graph: LGraph, info: any) {
+        console.debug("AfterChange", info);
+        this.eventBus.emit("afterChange", graph, info);
+    }
+
+    override onNodeAdded(node: LGraphNode, options: LGraphAddNodeOptions) {
+        layoutState.nodeAdded(node)
+        this.graphSync.onNodeAdded(node);
+
+        if ("comfyClass" in node                // Is this a comfy node
+            && !("svelteComponentType" in node) // ...and not also a ComfyWidgetNode
+            && !options.addedByDeserialize      // ...and we're not trying to deserialize an existing workflow
+            && get(uiState).autoAddUI) {
+            console.debug("[ComfyGraph] AutoAdd UI")
+            const comfyNode = node as ComfyGraphNode;
+            const widgetNodesAdded = []
+            for (let index = 0; index < comfyNode.inputs.length; index++) {
+                const input = comfyNode.inputs[index];
+                if ("config" in input) {
+                    const comfyInput = input as IComfyInputSlot;
+                    if (comfyInput.defaultWidgetNode) {
+                        const widgetNode = LiteGraph.createNode(comfyInput.defaultWidgetNode)
+                        const inputPos = comfyNode.getConnectionPos(true, index);
+                        this.add(widgetNode)
+                        widgetNode.connect(0, comfyNode, index);
+                        widgetNode.collapse();
+                        widgetNode.pos = [inputPos[0] - 140, inputPos[1] + LiteGraph.NODE_SLOT_HEIGHT / 2];
+                        widgetNodesAdded.push(widgetNode)
+                    }
+                }
+            }
+            const dragItems = widgetNodesAdded.map(wn => get(layoutState).allItemsByNode[wn.id]?.dragItem).filter(di => di)
+            console.debug("[ComfyGraph] Group new widgets", dragItems)
+            layoutState.groupItems(dragItems, comfyNode.comfyClass)
+        }
+
+        console.debug("Added", node);
+        this.eventBus.emit("nodeAdded", node);
+    }
+
+    override onNodeRemoved(node: LGraphNode) {
+        layoutState.nodeRemoved(node);
+        this.graphSync.onNodeRemoved(node);
+
+        console.debug("Removed", node);
+        this.eventBus.emit("nodeRemoved", node);
+    }
+
+    override onNodeConnectionChange(kind: LConnectionKind, node: LGraphNode, slot: SlotIndex, targetNode: LGraphNode, targetSlot: SlotIndex) {
+        console.debug("ConnectionChange", node);
+        this.eventBus.emit("nodeConnectionChanged", kind, node, slot, targetNode, targetSlot);
+    }
+}
diff --git a/src/lib/GraphSync.ts b/src/lib/GraphSync.ts
index 4b18344..9614594 100644
--- a/src/lib/GraphSync.ts
+++ b/src/lib/GraphSync.ts
@@ -2,6 +2,7 @@ import type { LGraph, LGraphNode } from "@litegraph-ts/core";
 import type ComfyApp from "./components/ComfyApp";
 import type { Unsubscriber, Writable } from "svelte/store";
 import type { ComfyWidgetNode } from "./nodes";
+import type ComfyGraph from "./ComfyGraph";
 
 type WidgetSubStore = {
     store: WidgetUIStateStore,
@@ -28,13 +29,11 @@ export default class GraphSync {
     // nodeId -> widgetSubStore
     private stores: Record<string, WidgetSubStore> = {}
 
-    constructor(app: ComfyApp) {
-        this.graph = app.lGraph;
-        app.eventBus.on("nodeAdded", this.onNodeAdded.bind(this));
-        app.eventBus.on("nodeRemoved", this.onNodeRemoved.bind(this));
+    constructor(graph: ComfyGraph) {
+        this.graph = graph;
     }
 
-    private onNodeAdded(node: LGraphNode) {
+    onNodeAdded(node: LGraphNode) {
         // TODO assumes only a single graph's widget state.
 
         if ("svelteComponentType" in node) {
@@ -44,7 +43,7 @@ export default class GraphSync {
         this.graph.setDirtyCanvas(true, true);
     }
 
-    private onNodeRemoved(node: LGraphNode) {
+    onNodeRemoved(node: LGraphNode) {
         if ("svelteComponentType" in node) {
             this.removeStore(node as ComfyWidgetNode);
         }
diff --git a/src/lib/IComfyInputSlot.ts b/src/lib/IComfyInputSlot.ts
index 0765207..23dc382 100644
--- a/src/lib/IComfyInputSlot.ts
+++ b/src/lib/IComfyInputSlot.ts
@@ -1,4 +1,5 @@
 import type { INodeInputSlot } from "@litegraph-ts/core";
+import type { ComfyWidgetNode } from "./nodes";
 
 // TODO generalize
 export type ComfyInputConfig = {
@@ -12,6 +13,7 @@ export type ComfyInputConfig = {
 }
 
 export default interface IComfyInputSlot extends INodeInputSlot {
-    serialize: boolean
-    config: ComfyInputConfig // stores range min/max/step, etc.
+    serialize: boolean,
+    defaultWidgetNode?: new (name?: string) => ComfyWidgetNode,
+    config: ComfyInputConfig, // stores range min/max/step, etc.
 }
diff --git a/src/lib/api.ts b/src/lib/api.ts
index b7fd412..159d455 100644
--- a/src/lib/api.ts
+++ b/src/lib/api.ts
@@ -231,7 +231,7 @@ export default class ComfyAPI extends EventTarget {
             };
         } catch (error) {
             console.error(error);
-            return { Running: [], Pending: [] };
+            return { Running: [], Pending: [], error };
         }
     }
 
@@ -245,7 +245,7 @@ export default class ComfyAPI extends EventTarget {
             return { History: Object.values(await res.json()) };
         } catch (error) {
             console.error(error);
-            return { History: [] };
+            return { History: [], error };
         }
     }
 
diff --git a/src/lib/components/ComfyApp.svelte b/src/lib/components/ComfyApp.svelte
index 8391532..3eb8191 100644
--- a/src/lib/components/ComfyApp.svelte
+++ b/src/lib/components/ComfyApp.svelte
@@ -107,9 +107,6 @@
          })
      }
 
-     app.eventBus.on("nodeAdded", layoutState.nodeAdded);
-     app.eventBus.on("nodeRemoved", layoutState.nodeRemoved);
-
      app.api.addEventListener("status", (ev: CustomEvent) => {
          queueState.statusUpdated(ev.detail as ComfyAPIStatus);
      });
@@ -176,6 +173,7 @@
         </Button>
         <Checkbox label="Lock Nodes" bind:value={$uiState.nodesLocked}/>
         <Checkbox label="Disable Interaction" bind:value={$uiState.graphLocked}/>
+        <Checkbox label="Auto-Add UI" bind:value={$uiState.autoAddUI}/>
         <label for="enable-ui-editing">Enable UI Editing</label>
         <select id="enable-ui-editing" name="enable-ui-editing" bind:value={$uiState.uiEditMode}>
             <option value="disabled">Disabled</option>
diff --git a/src/lib/components/ComfyApp.ts b/src/lib/components/ComfyApp.ts
index ebdd582..1a57908 100644
--- a/src/lib/components/ComfyApp.ts
+++ b/src/lib/components/ComfyApp.ts
@@ -13,13 +13,13 @@ import * as nodes from "$lib/nodes/index"
 import ComfyGraphCanvas, { type SerializedGraphCanvasState } from "$lib/ComfyGraphCanvas";
 import type ComfyGraphNode from "$lib/nodes/ComfyGraphNode";
 import * as widgets from "$lib/widgets/index"
-import type ComfyWidget from "$lib/widgets/ComfyWidget";
 import queueState from "$lib/stores/queueState";
-import GraphSync from "$lib/GraphSync";
 import type { SvelteComponentDev } from "svelte/internal";
 import type IComfyInputSlot from "$lib/IComfyInputSlot";
 import type { SerializedLayoutState } from "$lib/stores/layoutState";
 import layoutState from "$lib/stores/layoutState";
+import { toast } from '@zerodevx/svelte-toast'
+import ComfyGraph from "$lib/ComfyGraph";
 
 export const COMFYBOX_SERIAL_VERSION = 1;
 
@@ -40,17 +40,6 @@ export type SerializedAppState = {
     canvas: SerializedGraphCanvasState
 }
 
-type ComfyAppEvents = {
-    configured: (graph: LGraph) => void
-    nodeAdded: (node: LGraphNode) => void
-    nodeRemoved: (node: LGraphNode) => void
-    nodeConnectionChanged: (kind: LConnectionKind, node: LGraphNode, slot: INodeSlot, targetNode: LGraphNode, targetSlot: INodeSlot) => void
-    cleared: () => void
-    beforeChange: (graph: LGraph, param: any) => void
-    afterChange: (graph: LGraph, param: any) => void
-    restored: (workflow: SerializedAppState) => void
-}
-
 export type Progress = {
     value: number,
     max: number
@@ -61,12 +50,10 @@ export default class ComfyApp {
     rootEl: HTMLDivElement | null = null;
     canvasEl: HTMLCanvasElement | null = null;
     canvasCtx: CanvasRenderingContext2D | null = null;
-    lGraph: LGraph | null = null;
+    lGraph: ComfyGraph | null = null;
     lCanvas: ComfyGraphCanvas | null = null;
     dropZone: HTMLElement | null = null;
     nodeOutputs: Record<string, any> = {};
-    eventBus: TypedEmitter<ComfyAppEvents> = new EventEmitter() as TypedEmitter<ComfyAppEvents>;
-    graphSync: GraphSync;
 
     dragOverNode: LGraphNode | null = null;
     shiftDown: boolean = false;
@@ -82,12 +69,9 @@ export default class ComfyApp {
     async setup(): Promise<void> {
         this.rootEl = document.getElementById("main") as HTMLDivElement;
         this.canvasEl = document.getElementById("graph-canvas") as HTMLCanvasElement;
-        this.lGraph = new LGraph();
+        this.lGraph = new ComfyGraph();
         this.lCanvas = new ComfyGraphCanvas(this, this.canvasEl);
         this.canvasCtx = this.canvasEl.getContext("2d");
-        this.graphSync = new GraphSync(this);
-
-        this.addGraphLifecycleHooks();
 
         LiteGraph.release_link_on_empty_shows_menu = true;
         LiteGraph.alt_drag_do_clone_nodes = true;
@@ -118,7 +102,7 @@ export default class ComfyApp {
         }
 
         // Save current workflow automatically
-        setInterval(this.saveStateToLocalStorage.bind(this), 1000);
+        // setInterval(this.saveStateToLocalStorage.bind(this), 1000);
 
         this.addApiUpdateHandlers();
         this.addDropHandler();
@@ -147,58 +131,12 @@ export default class ComfyApp {
         this.lCanvas.draw(true, true);
     }
 
-    private graphOnConfigure() {
-        console.debug("Configured");
-        this.eventBus.emit("configured", this.lGraph);
-    }
-
-    private graphOnBeforeChange(graph: LGraph, info: any) {
-        console.debug("BeforeChange", info);
-        this.eventBus.emit("beforeChange", graph, info);
-    }
-
-    private graphOnAfterChange(graph: LGraph, info: any) {
-        console.debug("AfterChange", info);
-        this.eventBus.emit("afterChange", graph, info);
-    }
-
-    private graphOnNodeAdded(node: LGraphNode) {
-        console.debug("Added", node);
-        this.eventBus.emit("nodeAdded", node);
-    }
-
-    private graphOnNodeRemoved(node: LGraphNode) {
-        console.debug("Removed", node);
-        this.eventBus.emit("nodeRemoved", node);
-    }
-
-    private graphOnNodeConnectionChange(kind: LConnectionKind, node: LGraphNode, slot: INodeSlot, targetNode: LGraphNode, targetSlot: INodeSlot) {
-        console.debug("ConnectionChange", node);
-        this.eventBus.emit("nodeConnectionChanged", kind, node, slot, targetNode, targetSlot);
-    }
-
-    private canvasOnClear() {
-        console.debug("CanvasClear");
-        this.eventBus.emit("cleared");
-    }
-
     saveStateToLocalStorage() {
         const savedWorkflow = this.serialize();
         const json = JSON.stringify(savedWorkflow);
         localStorage.setItem("workflow", json)
     }
 
-    private addGraphLifecycleHooks() {
-        this.lGraph.onConfigure = this.graphOnConfigure.bind(this);
-        this.lGraph.onBeforeChange = this.graphOnBeforeChange.bind(this);
-        this.lGraph.onAfterChange = this.graphOnAfterChange.bind(this);
-        this.lGraph.onNodeAdded = this.graphOnNodeAdded.bind(this);
-        this.lGraph.onNodeRemoved = this.graphOnNodeRemoved.bind(this);
-        this.lGraph.onNodeConnectionChange = this.graphOnNodeConnectionChange.bind(this);
-
-        this.lCanvas.onClear = this.canvasOnClear.bind(this);
-    }
-
     static node_type_overrides: Record<string, typeof ComfyGraphNode> = {}
     static widget_type_overrides: Record<string, typeof SvelteComponentDev> = {}
 
@@ -642,7 +580,13 @@ export default class ComfyApp {
                         await this.api.queuePrompt(num, p);
                     } catch (error) {
                         // this.ui.dialog.show(error.response || error.toString());
-                        console.error(error.response || error.toString())
+                        const mes = error.response || error.toString()
+                        toast.push(`Error queuing prompt:\n${mes}`, {
+                            theme: {
+                                '--toastBackground': 'var(--color-red-500)',
+                            }
+                        })
+                        console.error("Error queuing prompt", mes, num, p)
                         break;
                     }
 
diff --git a/src/lib/components/WidgetContainer.svelte b/src/lib/components/WidgetContainer.svelte
index 2e425af..0b8759b 100644
--- a/src/lib/components/WidgetContainer.svelte
+++ b/src/lib/components/WidgetContainer.svelte
@@ -32,7 +32,7 @@
                && !$layoutState.isMenuOpen
 
 
- $: if ($queueState && widget) {
+ $: if ($queueState && widget && widget.node) {
      dragItem.isNodeExecuting = $queueState.runningNodeId === widget.node.id;
  }
 </script>
@@ -40,7 +40,7 @@
 
 {#if container}
     <BlockContainer {container} {classes} {zIndex} {showHandles} />
-{:else if widget}
+{:else if widget && widget.node}
     <div class="widget" class:widget-edit-outline={$uiState.uiEditMode === "widgets" && zIndex > 1}
         class:selected={$uiState.uiEditMode !== "disabled" && $layoutState.currentSelection.includes(widget.id)}
         class:is-executing={$queueState.runningNodeId && $queueState.runningNodeId == widget.node.id}
diff --git a/src/lib/nodes/ComfyWidgetNodes.ts b/src/lib/nodes/ComfyWidgetNodes.ts
index 29e13de..b4f090f 100644
--- a/src/lib/nodes/ComfyWidgetNodes.ts
+++ b/src/lib/nodes/ComfyWidgetNodes.ts
@@ -29,8 +29,6 @@ export abstract class ComfyWidgetNode<T = any> extends ComfyGraphNode {
 
     /** Svelte class for the frontend logic */
     abstract svelteComponentType: typeof SvelteComponentDev
-    /** Compatible litegraph widget types that can be connected to this node */
-    abstract inputWidgetTypes: string[]
 
     /** If false, user manually set min/max/step, and should not be autoinherited from connected input */
     autoConfig: boolean = true;
@@ -159,7 +157,6 @@ export class ComfySliderNode extends ComfyWidgetNode<number> {
     }
 
     override svelteComponentType = RangeWidget
-    override inputWidgetTypes = ["number", "slider"]
 
     static slotLayout: SlotLayout = {
         outputs: [
@@ -203,7 +200,6 @@ export class ComfyComboNode extends ComfyWidgetNode<string> {
     }
 
     override svelteComponentType = ComboWidget
-    override inputWidgetTypes = ["combo", "enum"]
 
     constructor(name?: string) {
         super(name, "A")
@@ -269,7 +265,6 @@ export class ComfyTextNode extends ComfyWidgetNode<string> {
     }
 
     override svelteComponentType = TextWidget
-    override inputWidgetTypes = ["text"]
 
     constructor(name?: string) {
         super(name, "")
diff --git a/src/lib/stores/layoutState.ts b/src/lib/stores/layoutState.ts
index 616e17b..4713252 100644
--- a/src/lib/stores/layoutState.ts
+++ b/src/lib/stores/layoutState.ts
@@ -14,6 +14,7 @@ type DragItemEntry = {
 export type LayoutState = {
     root: IDragItem | null,
     allItems: Record<DragItemID, DragItemEntry>,
+    allItemsByNode: Record<number, DragItemEntry>,
     currentId: number,
     currentSelection: DragItemID[],
     isConfiguring: boolean,
@@ -94,7 +95,7 @@ type LayoutStateOps = {
     updateChildren: (parent: IDragItem, children: IDragItem[]) => IDragItem[],
     nodeAdded: (node: LGraphNode) => void,
     nodeRemoved: (node: LGraphNode) => void,
-    groupItems: (dragItems: IDragItem[]) => ContainerLayout,
+    groupItems: (dragItems: IDragItem[], title: string) => ContainerLayout,
     ungroup: (container: ContainerLayout) => void,
     getCurrentSelection: () => IDragItem[],
     findLayoutForNode: (nodeId: number) => IDragItem | null;
@@ -108,6 +109,7 @@ export type WritableLayoutStateStore = Writable<LayoutState> & LayoutStateOps;
 const store: Writable<LayoutState> = writable({
     root: null,
     allItems: {},
+    allItemsByNode: {},
     currentId: 0,
     currentSelection: [],
     isMenuOpen: false,
@@ -175,6 +177,7 @@ function addWidget(parent: ContainerLayout, node: ComfyWidgetNode, attrs: Partia
     const parentEntry = state.allItems[parent.id]
     const entry: DragItemEntry = { dragItem, children: [], parent: null };
     state.allItems[dragItem.id] = entry;
+    state.allItemsByNode[node.id] = entry;
     console.debug("[layoutState] addWidget", state)
     moveItem(dragItem, parent)
     return dragItem;
@@ -230,6 +233,10 @@ function removeEntry(state: LayoutState, id: DragItemID) {
         const parentEntry = state.allItems[parent.id];
         parentEntry.children = parentEntry.children.filter(item => item.id !== id)
     }
+    if (entry.dragItem.type === "widget") {
+        const widget = entry.dragItem as WidgetLayout;
+        delete state.allItemsByNode[widget.node.id]
+    }
     delete state.allItems[id]
 }
 
@@ -286,7 +293,7 @@ function getCurrentSelection(): IDragItem[] {
     return state.currentSelection.map(id => state.allItems[id].dragItem)
 }
 
-function groupItems(dragItems: IDragItem[]): ContainerLayout {
+function groupItems(dragItems: IDragItem[], title: string = "Group"): ContainerLayout {
     if (dragItems.length === 0)
         return;
 
@@ -303,7 +310,7 @@ function groupItems(dragItems: IDragItem[]): ContainerLayout {
             index = indexFound
     }
 
-    const container = addContainer(parent as ContainerLayout, { title: "Group" }, index)
+    const container = addContainer(parent as ContainerLayout, { title }, index)
 
     for (const item of dragItems) {
         moveItem(item, container)
@@ -420,22 +427,29 @@ function serialize(): SerializedLayoutState {
 
 function deserialize(data: SerializedLayoutState, graph: LGraph) {
     const allItems: Record<DragItemID, DragItemEntry> = {}
+    const allItemsByNode: Record<number, DragItemEntry> = {}
     for (const pair of Object.entries(data.allItems)) {
         const [id, entry] = pair;
+
         const dragItem: IDragItem = {
             type: entry.dragItem.type,
             id: entry.dragItem.id,
             attrs: entry.dragItem.attrs
         };
-        if (dragItem.type === "widget") {
-            const widget = dragItem as WidgetLayout;
-            widget.node = graph.getNodeById(entry.dragItem.nodeId) as ComfyWidgetNode
-        }
-        allItems[id] = {
+
+        const dragEntry: DragItemEntry = {
             dragItem,
             children: [],
             parent: null
         }
+
+        allItems[id] = dragEntry
+
+        if (dragItem.type === "widget") {
+            const widget = dragItem as WidgetLayout;
+            widget.node = graph.getNodeById(entry.dragItem.nodeId) as ComfyWidgetNode
+            allItemsByNode[entry.dragItem.nodeId] = dragEntry
+        }
     }
 
     // reconnect parent/child tree
@@ -457,6 +471,7 @@ function deserialize(data: SerializedLayoutState, graph: LGraph) {
     const state: LayoutState = {
         root,
         allItems,
+        allItemsByNode,
         currentId: data.currentId,
         currentSelection: [],
         isMenuOpen: false,
diff --git a/src/lib/stores/uiState.ts b/src/lib/stores/uiState.ts
index 09c9892..e9422be 100644
--- a/src/lib/stores/uiState.ts
+++ b/src/lib/stores/uiState.ts
@@ -8,6 +8,7 @@ export type UIState = {
     app: ComfyApp,
     nodesLocked: boolean,
     graphLocked: boolean,
+    autoAddUI: boolean,
     uiEditMode: UIEditMode
 }
 
@@ -17,6 +18,7 @@ const store: WritableUIStateStore = writable(
         app: null,
         graphLocked: false,
         nodesLocked: false,
+        autoAddUI: true,
         uiEditMode: "disabled",
     })
 
diff --git a/src/lib/widgets.ts b/src/lib/widgets.ts
index 0bc4464..c020e9e 100644
--- a/src/lib/widgets.ts
+++ b/src/lib/widgets.ts
@@ -4,6 +4,7 @@ import ComfyValueControlWidget from "./widgets/ComfyValueControlWidget";
 import type { ComfyInputConfig } from "./IComfyInputSlot";
 import type IComfyInputSlot from "./IComfyInputSlot";
 import { BuiltInSlotShape } from "@litegraph-ts/core";
+import { ComfyComboNode, ComfySliderNode, ComfyTextNode } from "./nodes";
 
 type WidgetFactory = (node: LGraphNode, inputName: string, inputData: any, app: ComfyApp) => IComfyInputSlot;
 
@@ -29,19 +30,19 @@ function addComfyInput(node: LGraphNode, inputName: string, extraInfo: Partial<I
 
 const FLOAT: WidgetFactory = (node: LGraphNode, inputName: string, inputData: any): IComfyInputSlot => {
     const config = getNumberDefaults(inputData, 0.5);
-    return addComfyInput(node, inputName, { type: "number", config })
+    return addComfyInput(node, inputName, { type: "number", config, defaultWidgetNode: ComfySliderNode })
 }
 
 const INT: WidgetFactory = (node: LGraphNode, inputName: string, inputData: any): IComfyInputSlot => {
     const config = getNumberDefaults(inputData, 1);
-    return addComfyInput(node, inputName, { type: "number", config })
+    return addComfyInput(node, inputName, { type: "number", config, defaultWidgetNode: ComfySliderNode })
 };
 
 const STRING: WidgetFactory = (node: LGraphNode, inputName: string, inputData: any, app: ComfyApp): IComfyInputSlot => {
     const defaultValue = inputData[1].default || "";
     const multiline = !!inputData[1].multiline;
 
-    return addComfyInput(node, inputName, { type: "string", config: { defaultValue, multiline } })
+    return addComfyInput(node, inputName, { type: "string", config: { defaultValue, multiline }, defaultWidgetNode: ComfyTextNode })
 };
 
 const COMBO: WidgetFactory = (node: LGraphNode, inputName: string, inputData: any): IComfyInputSlot => {
@@ -50,7 +51,7 @@ const COMBO: WidgetFactory = (node: LGraphNode, inputName: string, inputData: an
     if (inputData[1] && inputData[1].default) {
         defaultValue = inputData[1].default;
     }
-    return addComfyInput(node, inputName, { type: "string", config: { values: type, defaultValue } })
+    return addComfyInput(node, inputName, { type: "string", config: { values: type, defaultValue }, defaultWidgetNode: ComfyComboNode })
 }
 
 const IMAGEUPLOAD: WidgetFactory = (node: LGraphNode, inputName: string, inputData: any, app): IComfyInputSlot => {
diff --git a/src/lib/widgets/ComboWidget.svelte b/src/lib/widgets/ComboWidget.svelte
index 423785d..9bc7d94 100644
--- a/src/lib/widgets/ComboWidget.svelte
+++ b/src/lib/widgets/ComboWidget.svelte
@@ -13,6 +13,7 @@
  $: widget && setNodeValue(widget);
 
  $: if (nodeValue !== null && (!$propsChanged || $propsChanged)) {
+     $nodeValue = option
      setOption($nodeValue)
      setNodeValue(widget)
      node.properties = node.properties