From 890c839b4db02cec95d2583dca4898081ff3c741 Mon Sep 17 00:00:00 2001
From: space-nuko <24979496+space-nuko@users.noreply.github.com>
Date: Tue, 2 May 2023 14:58:02 -0700
Subject: [PATCH] temp

---
 src/lib/IComfyInputSlot.ts                |   1 +
 src/lib/components/ComfyApp.svelte        |   8 +-
 src/lib/components/WidgetContainer.svelte |   5 +-
 src/lib/nodes/ComfyWidgetNodes.ts         | 136 ++++++++++++++++++----
 src/lib/stores/layoutState.ts             |  81 +++++--------
 src/lib/widgets/ComboWidget.svelte        |  29 +++--
 src/lib/widgets/RangeWidget.svelte        |  33 +++---
 src/lib/widgets/TextWidget.svelte         |  23 ++--
 8 files changed, 202 insertions(+), 114 deletions(-)

diff --git a/src/lib/IComfyInputSlot.ts b/src/lib/IComfyInputSlot.ts
index 4f6b090..0765207 100644
--- a/src/lib/IComfyInputSlot.ts
+++ b/src/lib/IComfyInputSlot.ts
@@ -1,5 +1,6 @@
 import type { INodeInputSlot } from "@litegraph-ts/core";
 
+// TODO generalize
 export type ComfyInputConfig = {
     min?: number,
     max?: number,
diff --git a/src/lib/components/ComfyApp.svelte b/src/lib/components/ComfyApp.svelte
index a3ffc7d..3ed254a 100644
--- a/src/lib/components/ComfyApp.svelte
+++ b/src/lib/components/ComfyApp.svelte
@@ -110,10 +110,10 @@
      app.eventBus.on("configured", nodeState.configureFinished);
      app.eventBus.on("cleared", nodeState.clear);
 
-     // app.eventBus.on("nodeAdded", layoutState.nodeAdded);
-     // app.eventBus.on("nodeRemoved", layoutState.nodeRemoved);
-     // app.eventBus.on("configured", layoutState.configureFinished);
-     // app.eventBus.on("cleared", layoutState.clear);
+     app.eventBus.on("nodeAdded", layoutState.nodeAdded);
+     app.eventBus.on("nodeRemoved", layoutState.nodeRemoved);
+     app.eventBus.on("configured", layoutState.configureFinished);
+     app.eventBus.on("cleared", layoutState.clear);
 
      app.eventBus.on("autosave", doAutosave);
      app.eventBus.on("restored", doRestore);
diff --git a/src/lib/components/WidgetContainer.svelte b/src/lib/components/WidgetContainer.svelte
index 44302a1..67f542c 100644
--- a/src/lib/components/WidgetContainer.svelte
+++ b/src/lib/components/WidgetContainer.svelte
@@ -27,7 +27,6 @@
  }
  else if (dragItem.type === "widget") {
      widget = dragItem as WidgetLayout;
-     widgetState = nodeState.findWidgetByName(widget.nodeId, widget.widgetName)
      container = null;
  }
 
@@ -47,9 +46,9 @@
 {:else if widget}
     <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.attrs.associatedNode}
+        class:is-executing={$queueState.runningNodeId && $queueState.runningNodeId == widget.node.id}
         >
-        <svelte:component this={getComponentForWidgetState(widgetState)} item={widgetState} />
+        <svelte:component this={widget.node.svelteComponentType} node={widget.node} />
     </div>
     {#if showHandles}
         <div class="handle handle-widget" style="z-index: {zIndex+100}" data-drag-item-id={widget.id} on:mousedown={startDrag} on:touchstart={startDrag} on:mouseup={stopDrag} on:touchend={stopDrag}/>
diff --git a/src/lib/nodes/ComfyWidgetNodes.ts b/src/lib/nodes/ComfyWidgetNodes.ts
index 008f394..7a8d20c 100644
--- a/src/lib/nodes/ComfyWidgetNodes.ts
+++ b/src/lib/nodes/ComfyWidgetNodes.ts
@@ -6,9 +6,10 @@ import TextWidget from "$lib/widgets/TextWidget.svelte";
 import type { SvelteComponentDev } from "svelte/internal";
 import { ComfyWidgets } from "$lib/widgets";
 import { Watch } from "@litegraph-ts/nodes-basic";
+import type IComfyInputSlot from "$lib/IComfyInputSlot";
+import { writable, type Unsubscriber, type Writable } from "svelte/store";
 
 export interface ComfyWidgetProperties extends Record<string, any> {
-    value: any,
 }
 
 /*
@@ -16,14 +17,20 @@ export interface ComfyWidgetProperties extends Record<string, any> {
  * widget is changed, the value of the first output in the node is updated
  * in the litegraph instance.
  */
-export abstract class ComfyWidgetNode extends ComfyGraphNode {
+export abstract class ComfyWidgetNode<T> extends ComfyGraphNode {
     abstract properties: ComfyWidgetProperties;
 
+    value: Writable<T>
+    unsubscribe: Unsubscriber;
+
     /** 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;
+
     override isBackendNode = false;
     override serialize_widgets = true;
 
@@ -31,9 +38,10 @@ export abstract class ComfyWidgetNode extends ComfyGraphNode {
 
     override size: Vector2 = [60, 40];
 
-    constructor(name?: string) {
-        super(name)
+    constructor(name?: string, value: T) {
         const color = LGraphCanvas.node_colors["blue"]
+        super(name)
+        this.value = writable(value)
         this.color ||= color.color
         this.bgColor ||= color.bgColor
         this.displayWidget = this.addWidget<ITextWidget>(
@@ -41,16 +49,15 @@ export abstract class ComfyWidgetNode extends ComfyGraphNode {
             "Value",
             ""
         );
+        this.unsubscribe = this.value.subscribe(this.onValueUpdated.bind(this))
     }
 
-    override onPropertyChanged(name: string, value: any) {
-        if (name == "value") {
-            this.displayWidget.value = Watch.toString(value)
-        }
+    private onValueUpdated(value: any) {
+        this.displayWidget.value = Watch.toString(value)
     }
 
     setValue(value: any) {
-        this.setProperty("value", value)
+        this.value.set(value)
     }
 
     override onExecute() {
@@ -60,15 +67,54 @@ export abstract class ComfyWidgetNode extends ComfyGraphNode {
 
         // TODO send event to linked nodes
     }
+
+    onConnectOutput(
+        outputIndex: number,
+        inputType: INodeInputSlot["type"],
+        input: INodeInputSlot,
+        inputNode: LGraphNode,
+        inputIndex: number
+    ): boolean {
+        if (this.autoConfig) {
+            // Copy properties from default config in input slot
+            if ("config" in input) {
+                const comfyInput = input as IComfyInputSlot;
+                for (const key in comfyInput.config)
+                    this.setProperty(key, comfyInput.config[key])
+            }
+            console.debug("Property copy", input, this.properties)
+        }
+
+        return true;
+    }
+
+    clampConfig() {
+        for (const link of this.getOutputLinks(0)) {
+            const node = this.graph._nodes_by_id[link.target_id]
+            if (node) {
+                const input = node.inputs[link.target_slot]
+                if (input && "config" in input)
+                    this.clampOneConfig(input as IComfyInputSlot)
+            }
+        }
+    }
+
+    clampOneConfig(input: IComfyInputSlot) {}
 }
 
 export interface ComfySliderProperties extends ComfyWidgetProperties {
-    value: number
+    min: number,
+    max: number,
+    step: number,
+    precision: number
 }
 
-export class ComfySliderNode extends ComfyWidgetNode {
+export class ComfySliderNode extends ComfyWidgetNode<number> {
     override properties: ComfySliderProperties = {
-        value: 0
+        min: 0,
+        max: 10,
+        step: 1,
+        precision: 1
     }
 
     override svelteComponentType = RangeWidget
@@ -79,6 +125,17 @@ export class ComfySliderNode extends ComfyWidgetNode {
             { name: "value", type: "number" }
         ]
     }
+
+    constructor(name?: string) {
+        super(name, 0)
+    }
+
+    override clampOneConfig(input: IComfyInputSlot) {
+        this.setProperty("min", Math.max(this.properties.min, input.config.min))
+        this.setProperty("max", Math.max(this.properties.max, input.config.max))
+        this.setProperty("step", Math.max(this.properties.step, input.config.step))
+        this.setProperty("value", Math.max(Math.min(this.properties.value, this.properties.max), this.properties.min))
+    }
 }
 
 LiteGraph.registerNodeType({
@@ -89,14 +146,12 @@ LiteGraph.registerNodeType({
 })
 
 export interface ComfyComboProperties extends ComfyWidgetProperties {
-    options: string[],
-    value: string
+    options: string[]
 }
 
-export class ComfyComboNode extends ComfyWidgetNode {
+export class ComfyComboNode extends ComfyWidgetNode<string> {
     override properties: ComfyComboProperties = {
-        options: ["*"],
-        value: "*"
+        options: ["*"]
     }
 
     static slotLayout: SlotLayout = {
@@ -107,6 +162,43 @@ export class ComfyComboNode extends ComfyWidgetNode {
 
     override svelteComponentType = ComboWidget
     override inputWidgetTypes = ["combo", "enum"]
+
+    constructor(name?: string) {
+        super(name, "*")
+    }
+
+    onConnectOutput(
+        outputIndex: number,
+        inputType: INodeInputSlot["type"],
+        input: INodeInputSlot,
+        inputNode: LGraphNode,
+        inputIndex: number
+    ): boolean {
+        if (!super.onConnectOutput(outputIndex, inputType, input, inputNode, inputIndex))
+            return false;
+
+        const thisProps = this.properties;
+        const otherProps = inputNode.properties;
+
+        // Ensure combo options match
+        if (!(otherProps.options instanceof Array))
+            return false;
+        if (otherProps.options.length !== thisProps.options.length)
+            return false;
+        if (otherProps.find((v, i) => thisProps[i] !== v))
+            return false;
+
+        return true;
+    }
+
+    override clampOneConfig(input: IComfyInputSlot) {
+        if (!(this.properties.value in input.config.values)) {
+            if (input.config.values.length === 0)
+                this.setProperty("value", "")
+            else
+                this.setProperty("value", input.config.values[0])
+        }
+    }
 }
 
 LiteGraph.registerNodeType({
@@ -117,12 +209,12 @@ LiteGraph.registerNodeType({
 })
 
 export interface ComfyTextProperties extends ComfyWidgetProperties {
-    value: string
+    multiline: boolean;
 }
 
-export class ComfyTextNode extends ComfyWidgetNode {
+export class ComfyTextNode extends ComfyWidgetNode<string> {
     override properties: ComfyTextProperties = {
-        value: ""
+        multiline: false
     }
 
     static slotLayout: SlotLayout = {
@@ -133,6 +225,10 @@ export class ComfyTextNode extends ComfyWidgetNode {
 
     override svelteComponentType = TextWidget
     override inputWidgetTypes = ["text"]
+
+    constructor(name?: string) {
+        super(name, "")
+    }
 }
 
 LiteGraph.registerNodeType({
diff --git a/src/lib/stores/layoutState.ts b/src/lib/stores/layoutState.ts
index 8a5c33a..1a4dc77 100644
--- a/src/lib/stores/layoutState.ts
+++ b/src/lib/stores/layoutState.ts
@@ -5,6 +5,7 @@ import type { LGraphNode, IWidget, LGraph } from "@litegraph-ts/core"
 import nodeState from "$lib/state/nodeState";
 import type { NodeStateStore } from './nodeState';
  import { dndzone, SHADOW_PLACEHOLDER_ITEM_ID } from 'svelte-dnd-action';
+import type { ComfyWidgetNode } from '$lib/nodes';
 
 type DragItemEntry = {
     dragItem: IDragItem,
@@ -59,16 +60,6 @@ const ALL_ATTRIBUTES: AttributesSpecList = [
                 editable: true,
             },
         ]
-    },
-    {
-        categoryName: "behavior",
-        specs: [
-            {
-                name: "associatedNode",
-                type: "number",
-                editable: false,
-            },
-        ]
     }
 ];
 export { ALL_ATTRIBUTES };
@@ -77,8 +68,7 @@ export type Attributes = {
     direction: string,
     title: string,
     showTitle: boolean,
-    classes: string,
-    associatedNode: number | null
+    classes: string
 }
 
 export interface IDragItem {
@@ -94,8 +84,7 @@ export interface ContainerLayout extends IDragItem {
 
 export interface WidgetLayout extends IDragItem {
     type: "widget",
-    nodeId: number,
-    widgetName: string
+    node: ComfyWidgetNode
 }
 
 type DragItemID = string;
@@ -121,7 +110,8 @@ const store: Writable<LayoutState> = writable({
     allItems: {},
     currentId: 0,
     currentSelection: [],
-    isMenuOpen: false
+    isMenuOpen: false,
+    isConfiguring: true
 })
 addContainer(null, { direction: "horizontal", showTitle: false });
 
@@ -155,7 +145,6 @@ function addContainer(parentId: DragItemID | null, attrs: Partial<Attributes> =
             showTitle: true,
             direction: "vertical",
             classes: "",
-            associatedNode: null,
             ...attrs
         }
     }
@@ -173,31 +162,25 @@ function addContainer(parentId: DragItemID | null, attrs: Partial<Attributes> =
     return dragItem;
 }
 
-function addWidget(parentId: DragItemID, node: LGraphNode, widget: IWidget<any, any>, attrs: Partial<Attributes> = {}, index: number = -1): WidgetLayout {
+function addWidget(parent: ContainerLayout, node: ComfyWidgetNode, attrs: Partial<Attributes> = {}, index: number = -1): WidgetLayout {
     const state = get(store);
+    const widgetName = "Widget"
     const dragItem: WidgetLayout = {
         type: "widget",
         id: `${state.currentId++}`,
-        nodeId: node.id,
-        widgetName: widget.name, // TODO name and displayName
+        node: node,
         attrs: {
-            title: widget.name,
+            title: widgetName,
             showTitle: true,
             direction: "horizontal",
             classes: "",
-            associatedNode: null,
             ...attrs
         }
     }
-    const parent = state.allItems[parentId]
-    const entry: DragItemEntry = { dragItem, children: [], parent: parent.dragItem };
+    const parentEntry = state.allItems[parent.id]
+    const entry: DragItemEntry = { dragItem, children: [], parent: parentEntry.dragItem };
     state.allItems[dragItem.id] = entry;
-    parent.children ||= []
-    if (index !== -1)
-        parent.children.splice(index, 0, dragItem)
-    else
-        parent.children.push(dragItem)
-    store.set(state)
+    moveItem(dragItem, parent)
     return dragItem;
 }
 
@@ -220,7 +203,17 @@ function nodeAdded(node: LGraphNode) {
         return;
 
     const parent = findDefaultContainerForInsertion();
-    // Add default node panel containing all widgets
+
+    // Two cases where we want to add nodes:
+    // 1. User adds a new UI node, so we should instantiate its widget in the frontend.
+    // 2. User adds a node with inputs that can be filled by frontend widgets.
+    // Depending on config, this means we should instantiate default UI nodes connected to those inputs.
+
+    if ("svelteComponentType" in node) {
+        addWidget(parent, node as ComfyWidgetNode);
+    }
+
+    // Add default node panel with all widgets autoinstantiated
     // if (node.widgets && node.widgets.length > 0) {
     //     const container = addContainer(parent.id, { title: node.title, direction: "vertical", associatedNode: node.id });
     //     for (const widget of node.widgets) {
@@ -248,32 +241,12 @@ function nodeRemoved(node: LGraphNode) {
 
     console.debug("[layoutState] nodeRemoved", node)
 
-    // Remove widgets bound to the node
     let del = Object.entries(state.allItems).filter(pair =>
         pair[1].dragItem.type === "widget"
-        && pair[1].dragItem.attrs.associatedNode === node.id)
-    for (const item of del) {
-        const [id, dragItem] = item;
-        console.debug("[layoutState] Remove widget", id, state.allItems[id])
-        removeEntry(state, id)
-    }
+        && (pair[1].dragItem as WidgetLayout).node.id === node.id)
 
-    const isAssociatedContainer = (dragItem: IDragItem) =>
-        dragItem.type === "container"
-        && dragItem.attrs.associatedNode === node.id;
-
-    let delContainers = []
-
-    // Remove widget from all children lists
-  // TODO just use parent.children
-    for (const entry of Object.values(state.allItems)) {
-        if (entry.children?.length === 0 && isAssociatedContainer(entry.dragItem))
-            delContainers.push(entry.dragItem.id)
-    }
-
-    // Remove empty containers bound to the node
-    for (const id of delContainers) {
-        console.debug("[layoutState] Remove container", id, state.allItems[id])
+    for (const pair of del) {
+        const [id, dragItem] = pair;
         removeEntry(state, id)
     }
 
@@ -290,7 +263,7 @@ function configureFinished(graph: LGraph) {
     const left = addContainer(state.root.id, { direction: "vertical", showTitle: false });
     const right = addContainer(state.root.id, { direction: "vertical", showTitle: false });
 
-    for (const node of graph.computeExecutionOrder(false, null)) {
+    for (const node of graph._nodes_in_order) {
         nodeAdded(node)
     }
 
diff --git a/src/lib/widgets/ComboWidget.svelte b/src/lib/widgets/ComboWidget.svelte
index b8a36bb..c4e2f5e 100644
--- a/src/lib/widgets/ComboWidget.svelte
+++ b/src/lib/widgets/ComboWidget.svelte
@@ -2,32 +2,39 @@
  import type { WidgetDrawState, WidgetUIState, WidgetUIStateStore } from "$lib/stores/nodeState";
  import { BlockTitle } from "@gradio/atoms";
  import { Dropdown } from "@gradio/form";
- import { get } from "svelte/store";
  import Select from 'svelte-select';
- export let item: WidgetUIState | null = null;
+ import type { ComfyComboNode } from "$lib/nodes/index";
+ import { type WidgetLayout } from "$lib/stores/layoutState";
+ import { get, type Writable } from "svelte/store";
+ export let widget: WidgetLayout | null = null;
+ let node: ComfyComboNode | null = null;
+ let nodeValue: Writable<string> | null = null;
  let itemValue: WidgetUIStateStore | null = null;
  let option: any;
 
- $: if(item) {
-     if (!itemValue)
-         itemValue = item.value;
-     if (!option)
-         option = get(item.value);
+ $: if(widget) {
+     node = widget.node as ComfyComboNode
+     nodeValue = node.value;
+     updateOption(); // don't react on option
  };
 
+ function updateOption() {
+     option = get(nodeValue);
+ }
+
  $: if (option && itemValue) {
      $itemValue = option.value
  }
 </script>
 
 <div class="wrapper gr-combo">
-    {#if item !== null && option !== undefined}
+    {#if node !== null && option !== undefined}
         <label>
-            <BlockTitle show_label={true}>{item.widget.name}</BlockTitle>
+            <BlockTitle show_label={true}>{widget.attrs.title}</BlockTitle>
             <Select
                 bind:value={option}
-                bind:items={item.widget.options.values}
-                disabled={item.widget.options.values.length === 0}
+                bind:items={node.properties.values}
+                disabled={node.properties.values.length === 0}
                 clearable={false}
                 on:change
                 on:select
diff --git a/src/lib/widgets/RangeWidget.svelte b/src/lib/widgets/RangeWidget.svelte
index f0fd27d..2c86236 100644
--- a/src/lib/widgets/RangeWidget.svelte
+++ b/src/lib/widgets/RangeWidget.svelte
@@ -1,35 +1,38 @@
 <script lang="ts">
- import type { WidgetUIState, WidgetUIStateStore } from "$lib/stores/nodeState";
+ import type { ComfySliderNode } from "$lib/nodes/index";
+ import { type WidgetLayout } from "$lib/stores/layoutState";
  import { Range } from "@gradio/form";
- import { get } from "svelte/store";
- export let item: WidgetUIState | null = null;
- let itemValue: WidgetUIStateStore | null = null;
+ import { get, type Writable } from "svelte/store";
+ export let widget: WidgetLayout | null = null;
+ let node: ComfySliderNode | null = null;
+ let nodeValue: Writable<number> | null = null;
  let option: number | null = null;
 
- $: if (item) {
-     itemValue = item.value;
+ $: if(widget) {
+     node = widget.node
+     nodeValue = node.value;
      updateOption(); // don't react on option
- }
+ };
 
  function updateOption() {
-     option = get(itemValue);
+     option = get(nodeValue);
  }
 
  function onRelease(e: Event) {
-     if (itemValue && option) {
-         $itemValue = option
+     if (nodeValue && option) {
+         $nodeValue = option
      }
  }
 </script>
 
 <div class="wrapper gr-range">
-    {#if item !== null && option !== null}
+    {#if node !== null && option !== null}
         <Range
             bind:value={option}
-            minimum={item.widget.options.min}
-            maximum={item.widget.options.max}
-            step={item.widget.options.step}
-            label={item.widget.name}
+            minimum={node.properties.min}
+            maximum={node.properties.max}
+            step={node.properties.step}
+            label={widget.attrs.title}
             show_label={true}
             on:release={onRelease}
             on:change
diff --git a/src/lib/widgets/TextWidget.svelte b/src/lib/widgets/TextWidget.svelte
index 960ac4c..58a54bc 100644
--- a/src/lib/widgets/TextWidget.svelte
+++ b/src/lib/widgets/TextWidget.svelte
@@ -1,18 +1,27 @@
 <script lang="ts">
- import type { WidgetUIState, WidgetUIStateStore } from "$lib/stores/nodeState";
+ import type { WidgetUIStateStore } from "$lib/stores/nodeState";
  import { TextBox } from "@gradio/form";
- export let item: WidgetUIState | null = null;
+ import type { ComfyComboNode } from "$lib/nodes/index";
+ import { type WidgetLayout } from "$lib/stores/layoutState";
+ import { get, type Writable } from "svelte/store";
+ export let widget: WidgetLayout | null = null;
+ let node: ComfyComboNode | null = null;
+ let nodeValue: Writable<string> | null = null;
  let itemValue: WidgetUIStateStore | null = null;
- $: if (item) { itemValue = item.value; }
+
+ $: if(widget) {
+     node = widget.node as ComfyComboNode
+     nodeValue = node.value;
+ };
 </script>
 
 <div class="wrapper gr-textbox">
-    {#if item !== null && itemValue !== null}
+    {#if node !== null && nodeValue !== null}
         <TextBox
             bind:value={$itemValue}
-            label={item.widget.name}
-            lines={item.widget.options.multiline ? 5 : 1}
-            max_lines={item.widget.options.multiline ? 5 : 1}
+            label={widget.attrs.title}
+            lines={node.properties.multiline ? 5 : 1}
+            max_lines={node.properties.multiline ? 5 : 1}
             show_label={true}
             on:change
             on:submit