Basic vanilla workflow converter

2023-05-21 14:11:12 -05:00
parent c3ab3aa69a
commit cad8b420f6
9 changed files with 329 additions and 86 deletions
--- a/src/lib/ComfyGraph.ts
+++ b/src/lib/ComfyGraph.ts
@@ -10,7 +10,6 @@ import type { ComfyBackendNode } from "./nodes/ComfyBackendNode";
 import type { ComfyComboNode, ComfyWidgetNode } from "./nodes/widgets";
 import selectionState from "./stores/selectionState";
 import type { WritableLayoutStateStore } from "./stores/layoutStates";
-import type { WorkflowInstID } from "./components/ComfyApp";
 import layoutStates from "./stores/layoutStates";
 import type { ComfyWorkflow } from "./stores/workflowState";
 import workflowState from "./stores/workflowState";
--- a/src/lib/components/ComfyApp.ts
+++ b/src/lib/components/ComfyApp.ts
@@ -21,7 +21,7 @@ import type ComfyGraphNode from "$lib/nodes/ComfyGraphNode";
 import queueState from "$lib/stores/queueState";
 import { type SvelteComponentDev } from "svelte/internal";
 import type IComfyInputSlot from "$lib/IComfyInputSlot";
-import type { LayoutState, SerializedLayoutState, WritableLayoutStateStore } from "$lib/stores/layoutStates";
+import { defaultWorkflowAttributes, type LayoutState, type SerializedLayoutState, type WritableLayoutStateStore } from "$lib/stores/layoutStates";
 import { toast } from '@zerodevx/svelte-toast'
 import ComfyGraph from "$lib/ComfyGraph";
 import { ComfyBackendNode } from "$lib/nodes/ComfyBackendNode";
@@ -44,6 +44,7 @@ import selectionState from "$lib/stores/selectionState";
 import layoutStates from "$lib/stores/layoutStates";
 import { ComfyWorkflow, type WorkflowAttributes, type WorkflowInstID } from "$lib/stores/workflowState";
 import workflowState from "$lib/stores/workflowState";
+import convertVanillaWorkflow, { type ComfyVanillaWorkflow } from "$lib/convertVanillaWorkflow";
 import convertVanillaWorkflow from "$lib/convertVanillaWorkflow";

 export const COMFYBOX_SERIAL_VERSION = 1;
@@ -149,6 +150,11 @@ function isVanillaWorkflow(data: any): data is SerializedLGraph {
    return data != null && (typeof data === "object") && data.last_node_id != null;
 }

+type BackendNodeDef = {
+    ctor: new (title?: string) => ComfyBackendNode,
+    nodeDef: ComfyNodeDef
+}
+
 export default class ComfyApp {
    api: ComfyAPI;

@@ -239,13 +245,13 @@ export default class ComfyApp {
        this.lCanvas.draw(true, true);
    }

-    serialize(workflow: ComfyWorkflow): SerializedAppState {
+    serialize(workflow: ComfyWorkflow, canvas?: SerializedGraphCanvasState): SerializedAppState {
        const layoutState = layoutStates.getLayout(workflow.id);
        if (layoutState == null)
            throw new Error("Workflow has no layout!")

        const { graph, layout, attrs } = workflow.serialize(layoutState);
-        const canvas = this.lCanvas.serialize();
+        canvas ||= this.lCanvas.serialize();

        return {
            comfyBoxWorkflow: true,
@@ -259,6 +265,22 @@ export default class ComfyApp {
        }
    }

+
+    convertVanillaWorkflow(workflow: ComfyVanillaWorkflow): SerializedAppState {
+        const attrs: WorkflowAttributes = {
+            ...defaultWorkflowAttributes,
+            title: "ComfyUI Workflow"
+        }
+
+        const canvas: SerializedGraphCanvasState = {
+            offset: [0, 0],
+            scale: 1
+        }
+
+        const comfyBoxWorkflow = convertVanillaWorkflow(workflow, attrs);
+        return this.serialize(comfyBoxWorkflow, canvas);
+    }
+
    saveStateToLocalStorage() {
        try {
            uiState.update(s => { s.isSavingToLocalStorage = true; return s; })
@@ -307,7 +329,11 @@ export default class ComfyApp {
    static node_type_overrides: Record<string, typeof ComfyBackendNode> = {}
    static widget_type_overrides: Record<string, typeof SvelteComponentDev> = {}

+    static knownBackendNodes: Record<string, BackendNodeDef> = {}
+
    private async registerNodes(defs: Record<ComfyNodeID, ComfyNodeDef>) {
+        ComfyApp.knownBackendNodes = {}
+
        // Register a node for each definition
        for (const [nodeId, nodeDef] of Object.entries(defs)) {
            const typeOverride = ComfyApp.node_type_overrides[nodeId]
@@ -330,6 +356,10 @@ export default class ComfyApp {

            LiteGraph.registerNodeType(node);
            node.category = nodeDef.category;
+            ComfyApp.knownBackendNodes[nodeId] = {
+                ctor,
+                nodeDef
+            }

            ComfyApp.registerDefaultSlotHandlers(nodeId, nodeDef)
        }
@@ -550,10 +580,10 @@ export default class ComfyApp {
    }

    async openVanillaWorkflow(data: SerializedLGraph) {
-        const converted = convertVanillaWorkflow(data)
+        const converted = this.convertVanillaWorkflow(data)
        console.info("WORKFLWO", converted)
        notify("Converted ComfyUI workflow to ComfyBox format.", { type: "info" })
-        // await this.openWorkflow(JSON.parse(pngInfo.workflow));
+        await this.openWorkflow(converted);
    }

    setActiveWorkflow(id: WorkflowInstID) {
--- a/src/lib/components/WidgetContainer.svelte
+++ b/src/lib/components/WidgetContainer.svelte
@@ -87,7 +87,7 @@
                 class:is-executing={$queueState.runningNodeID && $queueState.runningNodeID == widget.node.id}
                 class:hidden={hidden}
            >
-                <svelte:component this={widget.node.svelteComponentType} {layoutState} {widget} {isMobile} />
+                <svelte:component this={widget.node.svelteComponentType} {widget} {isMobile} />
            </div>
            {#if hidden && edit}
                <div class="handle handle-hidden" class:hidden={!edit} />
--- a/src/lib/convertVanillaWorkflow.ts
+++ b/src/lib/convertVanillaWorkflow.ts
@@ -1,51 +1,209 @@
-import { LGraph, type SerializedLGraph } from "@litegraph-ts/core";
+import { LGraph, type INodeInputSlot, type SerializedLGraph, type LinkID, type UUID, type NodeID, LiteGraph, BuiltInSlotType, type SerializedLGraphNode, type Vector2 } from "@litegraph-ts/core";
 import type { SerializedAppState } from "./components/ComfyApp";
 import layoutStates, { defaultWorkflowAttributes, type DragItemID, type SerializedDragEntry, type SerializedLayoutState } from "./stores/layoutStates";
-import type { WorkflowAttributes } from "./stores/workflowState";
+import { ComfyWorkflow, type WorkflowAttributes } from "./stores/workflowState";
 import type { SerializedGraphCanvasState } from "./ComfyGraphCanvas";
+import ComfyApp from "./components/ComfyApp";
+import { iterateNodeDefInputs } from "./ComfyNodeDef";
+import type { ComfyNodeDefInput } from "./ComfyNodeDef";
+import type IComfyInputSlot from "./IComfyInputSlot";
+import ComfyWidgets from "./widgets"
+import type { SerializedComfyWidgetNode } from "./nodes/widgets/ComfyWidgetNode";
+import { v4 as uuidv4 } from "uuid"
+import type ComfyWidgetNode from "./nodes/widgets/ComfyWidgetNode";

 /*
 * The workflow type used by base ComfyUI
 */
 export type ComfyVanillaWorkflow = SerializedLGraph;

-function addLayoutToVanillaWorkflow(workflow: ComfyVanillaWorkflow): SerializedLayoutState {
-    // easier to create a real layout first and then serialize it, then have to
-    // deal with manually constructing the serialized state from the ground up
-    const layoutState = layoutStates.createRaw();
-    const graph = new LGraph();
-    graph.configure(workflow)
-
-    for (const node of graph.iterateNodesInOrder()) {
-        console.warn("NODE", node)
-    }
-
-    return layoutState.serialize()
+/*
+ * The settings for a widget converted to an input slot via the widgetInputs.js
+ * frontend extension.
+ */
+type ComfyUIConvertedWidget = {
+    name: string,
+    config: ComfyNodeDefInput
 }

-export default function convertVanillaWorkflow(workflow: ComfyVanillaWorkflow): SerializedAppState {
-    const attrs: WorkflowAttributes = {
-        ...defaultWorkflowAttributes,
-        title: "ComfyUI Workflow"
-    }
-
-    const canvas: SerializedGraphCanvasState = {
-        offset: [0, 0],
-        scale: 1
-    }
-
-    const layout = addLayoutToVanillaWorkflow(workflow);
-
-    const appState: SerializedAppState = {
-        comfyBoxWorkflow: true,
-        createdBy: "ComfyUI", // ???
-        version: 1,
-        workflow,
-        attrs,
-        canvas,
-        layout,
-        commitHash: null
-    }
-
-    return appState
+interface IComfyUINodeInputSlot extends INodeInputSlot {
+    widget?: ComfyUIConvertedWidget
+}
+
+/*
+ * ComfyUI frontend nodes that should be converted directly to another type.
+ */
+const vanillaToComfyBoxNodeMapping: Record<string, string> = {
+    "Reroute": "utils/reroute"
+}
+
+/*
+ * Version of LGraphNode.getConnectionPos but for serialized nodes.
+ *
+ * TODO handle other node types! (horizontal, hardcoded slot pos, collapsed...)
+ */
+function getConnectionPos(node: SerializedLGraphNode, is_input: boolean, slotNumber: number, out: Vector2 = [0, 0]): Vector2 {
+    var offset = LiteGraph.NODE_SLOT_HEIGHT * 0.5;
+
+    if (is_input) {
+        out[0] = node.pos[0] + offset;
+    } else {
+        out[0] = node.pos[0] + this.size[0] + 1 - offset;
+    }
+    out[1] =
+        node.pos[1] +
+        (slotNumber + 0.7) * LiteGraph.NODE_SLOT_HEIGHT +
+        ((node.constructor as any).slot_start_y || 0);
+    return out;
+}
+
+export default function convertVanillaWorkflow(vanillaWorkflow: ComfyVanillaWorkflow, attrs: WorkflowAttributes): ComfyWorkflow {
+    const [comfyBoxWorkflow, layoutState] = ComfyWorkflow.create();
+    const { root, left, right } = layoutState.initDefaultLayout();
+
+    // TODO will need to convert IDs to UUIDs
+    const idToUUID: Record<NodeID | LinkID, UUID> = {}
+
+    for (const [id, node] of Object.entries(vanillaWorkflow.nodes)) {
+        const newType = vanillaToComfyBoxNodeMapping[node.type];
+        if (newType != null) {
+            node.type = newType;
+        }
+
+        // ComfyUI uses widgets on the node itself to change values. These are
+        // all made into input/output slots in ComfyBox. So we must convert
+        // serialized widgets into ComfyWidgetNodes, add new inputs/outputs,
+        // then attach the new nodes to the slots
+
+        const def = ComfyApp.knownBackendNodes[node.type];
+        if (def == null) {
+            console.error("Unknown backend node", node.type)
+            continue;
+        }
+
+        const group = layoutState.addContainer(left, { title: node.title })
+
+        for (const [inputName, [inputType, inputOpts]] of iterateNodeDefInputs(def.nodeDef)) {
+            // Detect if this input was a widget converted to an input
+            const convertedWidget = node.inputs?.find((i: IComfyUINodeInputSlot) => {
+                return i.widget?.name === inputName;
+            })
+
+            if (convertedWidget != null) {
+                // This input is an extra input slot on the node that should be
+                // accounted for.
+                const [value] = node.widgets_values.splice(0, 1);
+            }
+            else {
+                // This input is a widget, it should be converted to an input
+                // connected to a ComfyWidgetNode.
+
+                let widgetNodeType = null;
+                let widgetInputType = null;
+
+                if (Array.isArray(inputType)) {
+                    // Combo options of string[]
+                    widgetInputType = "string"
+                    widgetNodeType = "ui/combo";
+                }
+                else if (inputType in ComfyWidgets) {
+                    // Widget type
+                    const widgetFactory = ComfyWidgets[inputType]
+                    widgetInputType = widgetFactory.inputType
+                    widgetNodeType = widgetFactory.nodeType;
+                }
+                else if ("${inputType}:{inputName}" in ComfyWidgets) {
+                    // Widget type override for input of type with given name ("seed", "noise_seed")
+                    const widgetFactory = ComfyWidgets["${inputType}:{inputName}"]
+                    widgetInputType = widgetFactory.inputType
+                    widgetNodeType = widgetFactory.nodeType;
+                }
+                else {
+                    // Backend type, we can safely ignore this
+                    continue
+                }
+
+                const newInput: IComfyInputSlot = {
+                    name: inputName,
+                    link: null,
+                    type: widgetInputType,
+                    config: inputOpts,
+                    defaultWidgetNode: null,
+                    widgetNodeType,
+                    serialize: true,
+                    properties: {}
+                }
+
+                node.inputs ||= []
+                node.inputs.push(newInput);
+                const connInputIndex = node.inputs.length - 1;
+
+                // Now get the widget value.
+                const [value] = node.widgets_values.splice(0, 1);
+
+                const comfyWidgetNode = LiteGraph.createNode<ComfyWidgetNode>(widgetNodeType);
+                comfyWidgetNode.flags.collapsed = true;
+                const size: Vector2 = [0, 0];
+
+                // Compute collapsed size, sinze computeSize() ignores the collapsed flag
+                // LiteGraph only computes it if the node is rendered
+                const fontSize = LiteGraph.NODE_TEXT_SIZE;
+                size[0] = Math.min(
+                    comfyWidgetNode.size[0],
+                    comfyWidgetNode.title.length * fontSize +
+                    LiteGraph.NODE_TITLE_HEIGHT * 2
+                );
+
+                const serWidgetNode = comfyWidgetNode.serialize() as SerializedComfyWidgetNode;
+                serWidgetNode.comfyValue = value;
+                serWidgetNode.shownOutputProperties = {};
+                getConnectionPos(node, true, connInputIndex, serWidgetNode.pos);
+                serWidgetNode.pos[0] -= size[0] - 20;
+                serWidgetNode.pos[1] += LiteGraph.NODE_TITLE_HEIGHT / 2;
+                vanillaWorkflow.nodes.push(serWidgetNode)
+
+                layoutState.addWidget(group, comfyWidgetNode)
+
+                const connOutputIndex = serWidgetNode.outputs?.findIndex(o => o.name === comfyWidgetNode.outputSlotName)
+                if (connOutputIndex != null) {
+                    // Time to link
+                    const connOutput = serWidgetNode.outputs[connOutputIndex]
+                    const newLinkID = uuidv4();
+                    newInput.link = newLinkID
+                    connOutput.links ||= []
+                    connOutput.links.push(newLinkID);
+                    vanillaWorkflow.links ||= []
+                    vanillaWorkflow.links.push([newLinkID, serWidgetNode.id, connOutputIndex, node.id, connInputIndex, widgetInputType])
+                }
+                else {
+                    console.error("[convertVanillaWorkflow] No output to connect converted widget into!", comfyWidgetNode.outputSlotName, node)
+                }
+            }
+        }
+
+        // Add OUTPUT event slot to output nodes
+        // TODO needs to be generalized!
+        if (["PreviewImage", "SaveImage"].indexOf(node.type) !== -1) {
+            node.outputs ||= []
+            node.outputs.push({
+                name: "OUTPUT",
+                type: BuiltInSlotType.EVENT,
+                color_off: "rebeccapurple",
+                color_on: "rebeccapurple",
+                links: [],
+                properties: {},
+            })
+        }
+    }
+
+    const layout = layoutState.serialize();
+    comfyBoxWorkflow.deserialize(layoutState, { graph: vanillaWorkflow, attrs, layout })
+
+    for (const node of comfyBoxWorkflow.graph.iterateNodesInOrder()) {
+        if ((node as any).isBackendNode) {
+            console.warn("BENDNODE", node)
+        }
+    }
+
+    return comfyBoxWorkflow
 }
--- a/src/lib/nodes/ComfyBackendNode.ts
+++ b/src/lib/nodes/ComfyBackendNode.ts
@@ -2,7 +2,7 @@ import LGraphCanvas from "@litegraph-ts/core/src/LGraphCanvas";
 import ComfyGraphNode from "./ComfyGraphNode";
 import ComfyWidgets from "$lib/widgets"
 import type { ComfyWidgetNode } from "$lib/nodes/widgets";
-import { BuiltInSlotShape, BuiltInSlotType, type SerializedLGraphNode } from "@litegraph-ts/core";
+import { BuiltInSlotShape, BuiltInSlotType, LiteGraph, type SerializedLGraphNode } from "@litegraph-ts/core";
 import type IComfyInputSlot from "$lib/IComfyInputSlot";
 import type { ComfyInputConfig } from "$lib/IComfyInputSlot";
 import { iterateNodeDefOutputs, type ComfyNodeDef, iterateNodeDefInputs } from "$lib/ComfyNodeDef";
@@ -13,20 +13,24 @@ import type { SerializedPromptOutput } from "$lib/utils";
 */
 export class ComfyBackendNode extends ComfyGraphNode {
    comfyClass: string;
+    comfyNodeDef: ComfyNodeDef;
    displayName: string | null;

-    constructor(title: string, comfyClass: string, nodeData: any) {
+    constructor(title: string, comfyClass: string, nodeDef: ComfyNodeDef) {
        super(title)
        this.type = comfyClass; // XXX: workaround dependency in LGraphNode.addInput()
-        this.displayName = nodeData.display_name;
+        this.displayName = nodeDef.display_name;
+        this.comfyNodeDef = nodeDef;
        this.comfyClass = comfyClass;
        this.isBackendNode = true;

        const color = LGraphCanvas.node_colors["yellow"];
-        this.color = color.color
-        this.bgColor = color.bgColor
+        if (this.color == null)
+            this.color = color.color
+        if (this.bgColor == null)
+            this.bgColor = color.bgColor

-        this.setup(nodeData)
+        this.setup(nodeDef)

        // ComfyUI has no obvious way to identify if a node will return outputs back to the frontend based on its properties.
        // It just returns a hash like { "ui": { "images": results } } internally.
@@ -55,13 +59,13 @@ export class ComfyBackendNode extends ComfyGraphNode {
            } else {
                if (Array.isArray(type)) {
                    // Enums
-                    Object.assign(config, ComfyWidgets.COMBO(this, inputName, inputData) || {});
+                    Object.assign(config, ComfyWidgets.COMBO.callback(this, inputName, inputData) || {});
                } else if (`${type}:${inputName}` in ComfyWidgets) {
                    // Support custom ComfyWidgets by Type:Name
-                    Object.assign(config, ComfyWidgets[`${type}:${inputName}`](this, inputName, inputData) || {});
+                    Object.assign(config, ComfyWidgets[`${type}:${inputName}`].callback(this, inputName, inputData) || {});
                } else if (type in ComfyWidgets) {
                    // Standard type ComfyWidgets
-                    Object.assign(config, ComfyWidgets[type](this, inputName, inputData) || {});
+                    Object.assign(config, ComfyWidgets[type].callback(this, inputName, inputData) || {});
                } else {
                    // Node connection inputs (backend)
                    this.addInput(inputName, type);
--- a/src/lib/nodes/ComfyGraphNode.ts
+++ b/src/lib/nodes/ComfyGraphNode.ts
@@ -317,6 +317,11 @@ export default class ComfyGraphNode extends LGraphNode {
    }

    override onConfigure(o: SerializedLGraphNode) {
+        if (this.inputs.length != (o.inputs || []).length || this.outputs.length != (o.outputs || []).length) {
+            console.error("Expected node slot size mismatch when deserializing!", o.type, "ours", this.inputs, this.outputs, "theirs", o.inputs, o.outputs)
+            return;
+        }
+
        // Save the litegraph type of the default ComfyWidgetNode for each input.
        for (let index = 0; index < this.inputs.length; index++) {
            const input = this.inputs[index]
--- a/src/lib/nodes/widgets/ComfyWidgetNode.ts
+++ b/src/lib/nodes/widgets/ComfyWidgetNode.ts
@@ -15,6 +15,11 @@ export type AutoConfigOptions = {
    setWidgetTitle?: boolean
 }

+export type SerializedComfyWidgetNode = {
+    comfyValue?: any
+    shownOutputProperties?: ComfyWidgetNode["shownOutputProperties"]
+} & SerializedLGraphNode
+
 /*
 * NOTE: If you want to add a new widget but it has the same input/output type
 * as another one of the existing widgets, best to create a new "variant" of
@@ -35,6 +40,11 @@ export interface ComfyWidgetProperties extends ComfyGraphNodeProperties {
    defaultValue: any
 }

+export type ShownOutputProperty = {
+    type: string,
+    outputName: string
+}
+
 /*
 * A node that is tied to a UI widget in the frontend. When the frontend's
 * widget is changed, the value of the first output in the node is updated
@@ -72,7 +82,7 @@ export default abstract class ComfyWidgetNode<T = any> extends ComfyGraphNode {
    // shownInputProperties: string[] = []

    /** Names of properties to add as outputs */
-    private shownOutputProperties: Record<string, { type: string, outputName: string }> = {}
+    private shownOutputProperties: Record<string, ShownOutputProperty> = {}
    outputProperties: { name: string, type: string }[] = []

    override isBackendNode = false;
@@ -328,7 +338,7 @@ export default abstract class ComfyWidgetNode<T = any> extends ComfyGraphNode {

    clampOneConfig(input: IComfyInputSlot) { }

-    override onSerialize(o: SerializedLGraphNode) {
+    override onSerialize(o: SerializedComfyWidgetNode) {
        (o as any).comfyValue = get(this.value);
        (o as any).shownOutputProperties = this.shownOutputProperties
        super.onSerialize(o);
--- a/src/lib/stores/layoutStates.ts
+++ b/src/lib/stores/layoutStates.ts
@@ -657,13 +657,19 @@ export interface WidgetLayout extends IDragItem {
    node: ComfyWidgetNode
 }

+export type DefaultLayout = {
+    root: ContainerLayout,
+    left: ContainerLayout,
+    right: ContainerLayout,
+}
+
 export type DragItemID = UUID;

 type LayoutStateOps = {
    workflow: ComfyWorkflow | null,

-    addContainer: (parent: ContainerLayout | null, attrs: Partial<Attributes>, index?: number) => ContainerLayout,
-    addWidget: (parent: ContainerLayout, node: ComfyWidgetNode, attrs: Partial<Attributes>, index?: number) => WidgetLayout,
+    addContainer: (parent: ContainerLayout | null, attrs?: Partial<Attributes>, index?: number) => ContainerLayout,
+    addWidget: (parent: ContainerLayout, node: ComfyWidgetNode, attrs?: Partial<Attributes>, index?: number) => WidgetLayout,
    findDefaultContainerForInsertion: () => ContainerLayout | null,
    updateChildren: (parent: IDragItem, children: IDragItem[]) => IDragItem[],
    nodeAdded: (node: LGraphNode, options: LGraphAddNodeOptions) => void,
@@ -675,7 +681,7 @@ type LayoutStateOps = {
    findLayoutForNode: (nodeId: ComfyNodeID) => IDragItem | null,
    serialize: () => SerializedLayoutState,
    deserialize: (data: SerializedLayoutState, graph: LGraph) => void,
-    initDefaultLayout: () => void,
+    initDefaultLayout: () => DefaultLayout,
    onStartConfigure: () => void
    notifyWorkflowModified: () => void
 }
@@ -1036,7 +1042,7 @@ function createRaw(workflow: ComfyWorkflow | null = null): WritableLayoutStateSt
        return found.dragItem as WidgetLayout
    }

-    function initDefaultLayout() {
+    function initDefaultLayout(): DefaultLayout {
        store.set({
            root: null,
            allItems: {},
@@ -1056,6 +1062,8 @@ function createRaw(workflow: ComfyWorkflow | null = null): WritableLayoutStateSt
        })

        console.debug("[layoutState] initDefault")
+
+        return { root, left, right }
    }

    function serialize(): SerializedLayoutState {
--- a/src/lib/widgets.ts
+++ b/src/lib/widgets.ts
@@ -4,7 +4,16 @@ import type { ComfyInputConfig } from "./IComfyInputSlot";
 import { ComfyComboNode, ComfyNumberNode, ComfyTextNode } from "./nodes/widgets";
 import type { ComfyNodeDefInput } from "./ComfyNodeDef";

-type WidgetFactory = (node: LGraphNode, inputName: string, inputData: ComfyNodeDefInput) => IComfyInputSlot;
+type WidgetFactoryCallback = (node: LGraphNode, inputName: string, inputData: ComfyNodeDefInput) => IComfyInputSlot;
+
+type WidgetFactory = {
+    /* Creates the input */
+    callback: WidgetFactoryCallback,
+    /* Input type as used by litegraph */
+    inputType: string,
+    /* Node type to instantiate */
+    nodeType: string
+}

 function getNumberDefaults(inputData: ComfyNodeDefInput, defaultStep: number): ComfyInputConfig {
    let defaultValue = inputData[1].default;
@@ -34,34 +43,54 @@ function addComfyInput(node: LGraphNode, inputName: string, extraInfo: Partial<I
    return input;
 }

-const FLOAT: WidgetFactory = (node: LGraphNode, inputName: string, inputData: ComfyNodeDefInput): IComfyInputSlot => {
-    const config = getNumberDefaults(inputData, 0.5);
-    return addComfyInput(node, inputName, { type: "number", config, defaultWidgetNode: ComfyNumberNode })
+const FLOAT: WidgetFactory = {
+    callback: (node: LGraphNode, inputName: string, inputData: ComfyNodeDefInput): IComfyInputSlot => {
+        const config = getNumberDefaults(inputData, 0.5);
+        return addComfyInput(node, inputName, { type: "number", config, defaultWidgetNode: ComfyNumberNode })
+    },
+    inputType: "number",
+    nodeType: "ui/number"
 }

-const INT: WidgetFactory = (node: LGraphNode, inputName: string, inputData: ComfyNodeDefInput): IComfyInputSlot => {
-    const config = getNumberDefaults(inputData, 1);
-    return addComfyInput(node, inputName, { type: "number", config, defaultWidgetNode: ComfyNumberNode })
-};
-
-const STRING: WidgetFactory = (node: LGraphNode, inputName: string, inputData: ComfyNodeDefInput): IComfyInputSlot => {
-    const defaultValue = inputData[1].default || "";
-    const multiline = !!inputData[1].multiline;
-
-    return addComfyInput(node, inputName, { type: "string", config: { defaultValue, multiline }, defaultWidgetNode: ComfyTextNode })
-};
-
-const COMBO: WidgetFactory = (node: LGraphNode, inputName: string, inputData: ComfyNodeDefInput): IComfyInputSlot => {
-    const type = inputData[0] as string[];
-    let defaultValue = type[0];
-    if (inputData[1] && inputData[1].default) {
-        defaultValue = inputData[1].default;
-    }
-    return addComfyInput(node, inputName, { type: "string", config: { values: type, defaultValue }, defaultWidgetNode: ComfyComboNode })
+const INT: WidgetFactory = {
+    callback: (node: LGraphNode, inputName: string, inputData: ComfyNodeDefInput): IComfyInputSlot => {
+        const config = getNumberDefaults(inputData, 1);
+        return addComfyInput(node, inputName, { type: "number", config, defaultWidgetNode: ComfyNumberNode })
+    },
+    nodeType: "ui/number",
+    inputType: "number"
 }

-const IMAGEUPLOAD: WidgetFactory = (node: LGraphNode, inputName: string, inputData: ComfyNodeDefInput): IComfyInputSlot => {
-    return addComfyInput(node, inputName, { type: "number", config: {} })
+const STRING: WidgetFactory = {
+    callback: (node: LGraphNode, inputName: string, inputData: ComfyNodeDefInput): IComfyInputSlot => {
+        const defaultValue = inputData[1].default || "";
+        const multiline = !!inputData[1].multiline;
+
+        return addComfyInput(node, inputName, { type: "string", config: { defaultValue, multiline }, defaultWidgetNode: ComfyTextNode })
+    },
+    inputType: "number",
+    nodeType: "ui/text"
+}
+
+const COMBO: WidgetFactory = {
+    callback: (node: LGraphNode, inputName: string, inputData: ComfyNodeDefInput): IComfyInputSlot => {
+        const type = inputData[0] as string[];
+        let defaultValue = type[0];
+        if (inputData[1] && inputData[1].default) {
+            defaultValue = inputData[1].default;
+        }
+        return addComfyInput(node, inputName, { type: "string", config: { values: type, defaultValue }, defaultWidgetNode: ComfyComboNode })
+    },
+    inputType: "number",
+    nodeType: "ui/combo"
+}
+
+const IMAGEUPLOAD: WidgetFactory = {
+    callback: (node: LGraphNode, inputName: string, inputData: ComfyNodeDefInput): IComfyInputSlot => {
+        return addComfyInput(node, inputName, { type: "number", config: {} })
+    },
+    inputType: "COMFY_IMAGES",
+    nodeType: "ui/image_upload"
 }

 export type WidgetRepository = Record<string, WidgetFactory>