Subgraph ergonomic improvements & node def types

src/lib/ComfyNodeDef.ts

@@ -0,0 +1,52 @@
+import { range } from "./utils"
+import ComfyWidgets from "./widgets"
+
+export type ComfyNodeDef = {
+    name: string
+    display_name?: string
+    category: string
+    input: ComfyNodeDefInputs
+    /** Output type like "LATENT" or "IMAGE" */
+    output: string[]
+    output_name: string[]
+    output_is_list: boolean[]
+}
+
+export type ComfyNodeDefInputs = {
+    required: Record<string, ComfyNodeDefInput>,
+    optional?: Record<string, ComfyNodeDefInput>
+}
+export type ComfyNodeDefInput = [ComfyNodeDefInputType, ComfyNodeDefInputOptions | null]
+export type ComfyNodeDefInputType = string[] | keyof typeof ComfyWidgets | string
+export type ComfyNodeDefInputOptions = {
+    forceInput?: boolean
+}
+
+// TODO when comfy refactors
+export type ComfyNodeDefOutput = {
+    type: string,
+    name: string,
+    is_list?: boolean
+}
+
+export function isBackendNodeDefInputType(inputName: string, type: ComfyNodeDefInputType): type is string {
+    return !Array.isArray(type) && !(type in ComfyWidgets) && !(`${type}:${inputName}` in ComfyWidgets);
+}
+
+export function iterateNodeDefInputs(def: ComfyNodeDef): Iterable<[string, ComfyNodeDefInput]> {
+    var inputs = def.input.required
+    if (def.input.optional != null) {
+        inputs = Object.assign({}, def.input.required, def.input.optional)
+    }
+    return Object.entries(inputs);
+}
+
+export function iterateNodeDefOutputs(def: ComfyNodeDef): Iterable<ComfyNodeDefOutput> {
+    return range(def.output.length).map(i => {
+        return {
+            type: def.output[i],
+            name: def.output_name[i] || def.output[i],
+            is_list: def.output_is_list[i],
+        }
+    })
+}

src/lib/api.ts

@@ -4,6 +4,7 @@ import EventEmitter from "events";
 import type { ComfyExecutionResult, ComfyImageLocation } from "./nodes/ComfyWidgetNodes";
 import type { SerializedLGraph, UUID } from "@litegraph-ts/core";
 import type { SerializedLayoutState } from "./stores/layoutState";
+import type { ComfyNodeDef } from "./ComfyNodeDef";
 
 export type ComfyPromptRequest = {
     client_id?: string,
@@ -226,7 +227,7 @@ export default class ComfyAPI {
      * Loads node object definitions for the graph
      * @returns The node definitions
      */
-    async getNodeDefs(): Promise<any> {
+    async getNodeDefs(): Promise<Record<ComfyNodeID, ComfyNodeDef>> {
         return fetch(this.getBackendUrl() + "/object_info", { cache: "no-store" })
             .then(resp => resp.json())
     }

src/lib/components/ComfyApp.ts

@@ -1,4 +1,4 @@
-import { LiteGraph, LGraph, LGraphCanvas, LGraphNode, type LGraphNodeConstructor, type LGraphNodeExecutable, type SerializedLGraph, type SerializedLGraphGroup, type SerializedLGraphNode, type SerializedLLink, NodeMode, type Vector2, BuiltInSlotType, type INodeInputSlot, type NodeID } from "@litegraph-ts/core";
+import { LiteGraph, LGraph, LGraphCanvas, LGraphNode, type LGraphNodeConstructor, type LGraphNodeExecutable, type SerializedLGraph, type SerializedLGraphGroup, type SerializedLGraphNode, type SerializedLLink, NodeMode, type Vector2, BuiltInSlotType, type INodeInputSlot, type NodeID, type NodeTypeSpec, type NodeTypeOpts } from "@litegraph-ts/core";
 import type { LConnectionKind, INodeSlot } from "@litegraph-ts/core";
 import ComfyAPI, { type ComfyAPIStatusResponse, type ComfyBoxPromptExtraData, type ComfyPromptRequest, type ComfyNodeID, type PromptID } from "$lib/api"
 import { getPngMetadata, importA1111 } from "$lib/pnginfo";
@@ -34,13 +34,10 @@ import configState from "$lib/stores/configState";
 import { blankGraph } from "$lib/defaultGraph";
 import type { ComfyExecutionResult } from "$lib/nodes/ComfyWidgetNodes";
 import ComfyPromptSerializer from "./ComfyPromptSerializer";
+import { iterateNodeDefInputs, type ComfyNodeDef, isBackendNodeDefInputType, iterateNodeDefOutputs } from "$lib/ComfyNodeDef";
 
 export const COMFYBOX_SERIAL_VERSION = 1;
 
-LiteGraph.catch_exceptions = false;
-LiteGraph.CANVAS_GRID_SIZE = 32;
-LiteGraph.default_subgraph_lgraph_factory = () => new ComfyGraph();
-
 if (typeof window !== "undefined") {
     // Load default visibility
     nodes.ComfyReroute.setDefaultTextVisibility(!!localStorage["Comfy.ComfyReroute.DefaultVisibility"]);
@@ -130,8 +127,6 @@ export default class ComfyApp {
         this.lCanvas.allow_dragnodes = uiUnlocked;
         this.lCanvas.allow_interaction = uiUnlocked;
 
-        (window as any).LiteGraph = LiteGraph;
-
         // await this.#invokeExtensionsAsync("init");
         await this.registerNodes();
 
@@ -205,15 +200,11 @@ export default class ComfyApp {
     static widget_type_overrides: Record<string, typeof SvelteComponentDev> = {}
 
     private async registerNodes() {
-        const app = this;
-
         // Load node definitions from the backend
         const defs = await this.api.getNodeDefs();
 
         // Register a node for each definition
-        for (const nodeId in defs) {
+        for (const [nodeId, nodeDef] of Object.entries(defs)) {
-            const nodeData = defs[nodeId];
-
             const typeOverride = ComfyApp.node_type_overrides[nodeId]
             if (typeOverride)
                 console.debug("Attaching custom type to received node:", nodeId, typeOverride)
@@ -221,19 +212,42 @@ export default class ComfyApp {
 
             const ctor = class extends baseClass {
                 constructor(title?: string) {
-                    super(title, nodeId, nodeData);
+                    super(title, nodeId, nodeDef);
                 }
             }
 
             const node: LGraphNodeConstructor = {
                 class: ctor,
-                title: nodeData.display_name || nodeData.name,
+                title: nodeDef.display_name || nodeDef.name,
                 type: nodeId,
                 desc: `ComfyNode: ${nodeId}`
             }
 
             LiteGraph.registerNodeType(node);
-            node.category = nodeData.category;
+            node.category = nodeDef.category;
+
+            ComfyApp.registerDefaultSlotHandlers(nodeId, nodeDef)
+        }
+    }
+
+    static registerDefaultSlotHandlers(nodeId: string, nodeDef: ComfyNodeDef) {
+        const nodeTypeSpec: NodeTypeOpts = {
+            node: nodeId,
+            title: nodeDef.display_name || nodeDef.name,
+            properties: null,
+            inputs: null,
+            outputs: null
+        }
+
+        for (const [inputName, [inputType, _inputOpts]] of iterateNodeDefInputs(nodeDef)) {
+            if (isBackendNodeDefInputType(inputName, inputType)) {
+                LiteGraph.slot_types_default_out[inputType] ||= []
+                LiteGraph.slot_types_default_out[inputType].push(nodeTypeSpec)
+            }
+        }
+        for (const output of iterateNodeDefOutputs(nodeDef)) {
+            LiteGraph.slot_types_default_in[output.type] ||= []
+            LiteGraph.slot_types_default_in[output.type].push(nodeTypeSpec)
         }
     }
 

src/lib/init.ts

@@ -0,0 +1,19 @@
+import ComfyGraph from '$lib/ComfyGraph';
+import { LGraphCanvas, LiteGraph, Subgraph } from '@litegraph-ts/core';
+
+export function configureLitegraph(isMobile: boolean = false) {
+    LiteGraph.catch_exceptions = false;
+    LiteGraph.use_uuids = true;
+    LiteGraph.CANVAS_GRID_SIZE = 32;
+
+    if (isMobile) {
+        LiteGraph.dialog_close_on_mouse_leave = false;
+        LiteGraph.search_hide_on_mouse_leave = false;
+        LiteGraph.pointerevents_method = "pointer";
+    }
+
+    Subgraph.default_lgraph_factory = () => new ComfyGraph;
+
+    (window as any).LiteGraph = LiteGraph;
+    (window as any).LGraphCanvas = LGraphCanvas;
+}

src/lib/nodes/ComfyBackendNode.ts

@@ -5,6 +5,7 @@ import type { ComfyWidgetNode, ComfyExecutionResult } from "./ComfyWidgetNodes";
 import { BuiltInSlotShape, BuiltInSlotType, 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";
 
 /*
  * Base class for any node with configuration sent by the backend.
@@ -37,22 +38,19 @@ export class ComfyBackendNode extends ComfyGraphNode {
     // comfy class -> input name -> input config
     private static defaultInputConfigs: Record<string, Record<string, ComfyInputConfig>> = {}
 
-    private setup(nodeData: any) {
+    private setup(nodeDef: ComfyNodeDef) {
-        var inputs = nodeData["input"]["required"];
-        if (nodeData["input"]["optional"] != undefined) {
-            inputs = Object.assign({}, nodeData["input"]["required"], nodeData["input"]["optional"])
-        }
-
         ComfyBackendNode.defaultInputConfigs[this.type] = {}
 
-        for (const inputName in inputs) {
+        for (const [inputName, inputData] of iterateNodeDefInputs(nodeDef)) {
             const config: Partial<IComfyInputSlot> = {};
 
-            const inputData = inputs[inputName];
+            const [type, opts] = inputData;
-            const type = inputData[0];
 
-            if (inputData[1]?.forceInput) {
+            if (opts?.forceInput) {
-                this.addInput(inputName, type);
+                if (Array.isArray(type)) {
+                    throw new Error(`Can't have forceInput set to true for an enum type! ${type}`)
+                }
+                this.addInput(inputName, type as string);
             } else {
                 if (Array.isArray(type)) {
                     // Enums
@@ -73,11 +71,9 @@ export class ComfyBackendNode extends ComfyGraphNode {
                 ComfyBackendNode.defaultInputConfigs[this.type][inputName] = (config as IComfyInputSlot).config
         }
 
-        for (const o in nodeData["output"]) {
+        for (const output of iterateNodeDefOutputs(nodeDef)) {
-            const output = nodeData["output"][o];
+            const outputShape = output.is_list ? BuiltInSlotShape.GRID_SHAPE : BuiltInSlotShape.CIRCLE_SHAPE;
-            const outputName = nodeData["output_name"][o] || output;
+            this.addOutput(output.name, output.type, { shape: outputShape });
-            const outputShape = nodeData["output_is_list"][o] ? BuiltInSlotShape.GRID_SHAPE : BuiltInSlotShape.CIRCLE_SHAPE;
-            this.addOutput(outputName, output, { shape: outputShape });
         }
 
         this.serialize_widgets = false;