Add TS types / merge ComfyLGraphNode (#902)

* Add TS type for LGraphNodeConstructor

* Add TS type & move shared prop to parent

* Add TS types - Comfy augmentations

* nit - TS type

* Merge ComfyLGNode into existing augmentations

* nit - fix missed explicit type on import

src/extensions/core/nodeBadge.ts

@@ -1,6 +1,6 @@
 import { app, type ComfyApp } from '@/scripts/app'
 import type { ComfyExtension } from '@/types/comfy'
-import type { ComfyLGraphNode } from '@/types/comfyLGraphNode'
+import type { LGraphNode } from '@comfyorg/litegraph'
 import { LGraphBadge } from '@comfyorg/litegraph'
 import { useSettingStore } from '@/stores/settingStore'
 import { computed, ComputedRef, watch } from 'vue'
@@ -12,9 +12,8 @@ import type { Palette } from '@/types/colorPalette'
 import type { ComfyNodeDef } from '@/types/apiTypes'
 import { useNodeDefStore } from '@/stores/nodeDefStore'
 
-function getNodeSource(node: ComfyLGraphNode): NodeSource | null {
+function getNodeSource(node: LGraphNode): NodeSource | null {
-  const nodeDef = (node.constructor as typeof ComfyLGraphNode)
+  const nodeDef = node.constructor.nodeData
-    .nodeData as ComfyNodeDef
   // Frontend-only nodes don't have nodeDef
   if (!nodeDef) {
     return null
@@ -23,29 +22,23 @@ function getNodeSource(node: ComfyLGraphNode): NodeSource | null {
   return nodeDefStore.nodeDefsByName[nodeDef.name]?.nodeSource ?? null
 }
 
-function isCoreNode(node: ComfyLGraphNode) {
+function isCoreNode(node: LGraphNode) {
   return getNodeSource(node)?.type === NodeSourceType.Core
 }
 
-function badgeTextVisible(
+function badgeTextVisible(node: LGraphNode, badgeMode: NodeBadgeMode): boolean {
-  node: ComfyLGraphNode,
-  badgeMode: NodeBadgeMode
-): boolean {
   return (
     badgeMode === NodeBadgeMode.None ||
     (isCoreNode(node) && badgeMode === NodeBadgeMode.HideBuiltIn)
   )
 }
 
-function getNodeIdBadgeText(
+function getNodeIdBadgeText(node: LGraphNode, nodeIdBadgeMode: NodeBadgeMode) {
-  node: ComfyLGraphNode,
-  nodeIdBadgeMode: NodeBadgeMode
-) {
   return badgeTextVisible(node, nodeIdBadgeMode) ? '' : `#${node.id}`
 }
 
 function getNodeSourceBadgeText(
-  node: ComfyLGraphNode,
+  node: LGraphNode,
   nodeSourceBadgeMode: NodeBadgeMode
 ) {
   const nodeSource = getNodeSource(node)
@@ -55,11 +48,11 @@ function getNodeSourceBadgeText(
 }
 
 function getNodeLifeCycleBadgeText(
-  node: ComfyLGraphNode,
+  node: LGraphNode,
   nodeLifeCycleBadgeMode: NodeBadgeMode
 ) {
   let text = ''
-  const nodeDef = (node.constructor as typeof ComfyLGraphNode).nodeData
+  const nodeDef = node.constructor.nodeData
 
   // Frontend-only nodes don't have nodeDef
   if (!nodeDef) {
@@ -118,7 +111,7 @@ class NodeBadgeExtension implements ComfyExtension {
     })
   }
 
-  nodeCreated(node: ComfyLGraphNode, app: ComfyApp) {
+  nodeCreated(node: LGraphNode, app: ComfyApp) {
     node.badgePosition = BadgePosition.TopRight
     // @ts-expect-error Disable ComfyUI-Manager's badge drawing by setting badge_enabled to true. Remove this when ComfyUI-Manager's badge drawing is removed.
     node.badge_enabled = true

src/extensions/core/nodeTemplates.ts

@@ -378,7 +378,6 @@ app.registerExtension({
             const nodeIds = Object.keys(app.canvas.selected_nodes)
             for (let i = 0; i < nodeIds.length; i++) {
               const node = app.graph.getNodeById(nodeIds[i])
-              // @ts-expect-error
               const nodeData = node?.constructor.nodeData
 
               let groupData = GroupNodeHandler.getGroupData(node)

src/extensions/core/rerouteNode.ts

@@ -74,7 +74,6 @@ app.registerExtension({
               const link = app.graph.links[linkId]
               if (!link) return
               const node = app.graph.getNodeById(link.origin_id)
-              // @ts-expect-error Nodes that extend LGraphNode will not have a static type property
               const type = node.constructor.type
               if (type === 'Reroute') {
                 if (node === this) {
@@ -113,7 +112,6 @@ app.registerExtension({
                 if (!link) continue
 
                 const node = app.graph.getNodeById(link.target_id)
-                // @ts-expect-error Nodes that extend LGraphNode will not have a static type property
                 const type = node.constructor.type
 
                 if (type === 'Reroute') {
@@ -179,7 +177,6 @@ app.registerExtension({
                   }
                   if (!targetWidget) {
                     targetWidget = targetNode.widgets?.find(
-                      // @ts-expect-error fix widget types
                       (w) => w.name === targetInput.widget.name
                     )
                   }

src/scripts/app.ts

@@ -51,7 +51,6 @@ import { useToastStore } from '@/stores/toastStore'
 import { ModelStore, useModelStore } from '@/stores/modelStore'
 import type { ToastMessageOptions } from 'primevue/toast'
 import { useWorkspaceStore } from '@/stores/workspaceStateStore'
-import { ComfyLGraphNode } from '@/types/comfyLGraphNode'
 import { useExecutionStore } from '@/stores/executionStore'
 
 export const ANIM_PREVIEW_WIDGET = '$$comfy_animation_preview'
@@ -1935,7 +1934,6 @@ export class ComfyApp {
         {
           name,
           display_name: name,
-          // @ts-expect-error
           category: node.category || '__frontend_only__',
           input: { required: {}, optional: {} },
           output: [],
@@ -1989,7 +1987,7 @@ export class ComfyApp {
 
   async registerNodeDef(nodeId: string, nodeData: ComfyNodeDef) {
     const self = this
-    const node: new () => ComfyLGraphNode = class ComfyNode extends LGraphNode {
+    const node = class ComfyNode extends LGraphNode {
       static comfyClass? = nodeData.name
       // TODO: change to "title?" once litegraph.d.ts has been updated
       static title = nodeData.display_name || nodeData.name
@@ -2074,7 +2072,6 @@ export class ComfyApp {
     await this.#invokeExtensionsAsync('beforeRegisterNodeDef', node, nodeData)
     LiteGraph.registerNodeType(nodeId, node)
     // Note: Do not move this to the class definition, it will be overwritten
-    // @ts-expect-error
     node.category = nodeData.category
   }
 

src/types/comfyLGraphNode.ts

@@ -1,11 +0,0 @@
-import type { LGraphNode } from '@comfyorg/litegraph'
-import type { ComfyNodeDef } from './apiTypes'
-
-export declare class ComfyLGraphNode extends LGraphNode {
-  static comfyClass: string
-  static title: string
-  static nodeData?: ComfyNodeDef
-  static category?: string
-
-  constructor(title?: string)
-}

src/types/litegraph-augmentation.d.ts

@@ -1,14 +1,33 @@
 import '@comfyorg/litegraph'
+import type { ComfyNodeDef } from '@/types/apiTypes'
+import type { LLink } from '@comfyorg/litegraph'
 
 /**
  *  ComfyUI extensions of litegraph
  */
 declare module '@comfyorg/litegraph' {
+  interface LGraphNodeConstructor<T extends LGraphNode = LGraphNode> {
+    type?: string
+    comfyClass: string
+    title: string
+    nodeData?: ComfyNodeDef
+    category?: string
+    new (): T
+  }
+
   interface LGraphNode {
+    constructor: LGraphNodeConstructor
+
     /**
      * Callback fired on each node after the graph is configured
      */
     onAfterGraphConfigured?(): void
+    onNodeCreated?(this: LGraphNode): void
+    setInnerNodes?(nodes: LGraphNode[]): void
+    applyToGraph?(extraLinks?: LLink[]): void
+    updateLink?(link: LLink): LLink | null
+
+    comfyClass?: string
 
     /**
      * If the node is a frontend only node and should not be serialized into the prompt.
@@ -24,25 +43,30 @@ declare module '@comfyorg/litegraph' {
   }
 
   interface IWidget<TValue = any, TOptions = any> {
+    type?: string
+
     /**
      * Allows for additional cleanup when removing a widget when converting to input.
      */
     onRemove?(): void
 
+    serializeValue?(node?: LGraphNode, i?: string)
+    beforeQueued?(): void
+
     /**
      * DOM element used for the widget
      */
     element?: HTMLElement
 
     tooltip?: string
+
+    origType?: IWidget['type']
+    origComputeSize?: IWidget['computeSize']
+    origSerializeValue?: IWidget['serializeValue']
   }
 
-  interface INodeOutputSlot {
+  interface INodeSlot {
-    widget?: unknown
+    widget?: unknown & { name?: string }
-  }
-
-  interface INodeInputSlot {
-    widget?: unknown
   }
 }
 

src/types/litegraph-core-augmentation.d.ts

@@ -6,6 +6,7 @@ declare module '@comfyorg/litegraph' {
   interface LiteGraphExtended {
     search_filter_enabled: boolean
     middle_click_slot_add_default_node: boolean
+    registered_node_types: Record<string, LGraphNodeConstructor>
     registered_slot_out_types: Record<string, { nodes: string[] }>
     registered_slot_in_types: Record<string, { nodes: string[] }>
     slot_types_out: string[]