Relands "Add support for LiteGraph to convert to classes" (#667)

* Revert "Revert "Add support for LiteGraph to convert to classes (#334)" (#386)"

This reverts commit d607f6c7f7.

* Format

src/extensions/core/contextMenuFilter.ts

@@ -8,19 +8,19 @@ const ext = {
   init() {
     const ctxMenu = LiteGraph.ContextMenu
     // @ts-expect-error
-    // TODO Very hacky way to modify Litegraph behaviour. Fix this later.
+    // TODO Very hacky way to modify Litegraph behaviour. Fix ctx later.
     LiteGraph.ContextMenu = function (values, options) {
-      const ctx = ctxMenu.call(this, values, options)
+      const ctx = new ctxMenu(values, options)
 
       // If we are a dark menu (only used for combo boxes) then add a filter input
-      if (options?.className === 'dark' && values?.length > 10) {
+      if (options?.className === 'dark' && values?.length > 4) {
         const filter = document.createElement('input')
         filter.classList.add('comfy-context-menu-filter')
         filter.placeholder = 'Filter list'
-        this.root.prepend(filter)
+        ctx.root.prepend(filter)
 
         const items = Array.from(
-          this.root.querySelectorAll('.litemenu-entry')
+          ctx.root.querySelectorAll('.litemenu-entry')
         ) as HTMLElement[]
         let displayedItems = [...items]
         let itemCount = displayedItems.length
@@ -61,16 +61,15 @@ const ext = {
           }
 
           const positionList = () => {
-            const rect = this.root.getBoundingClientRect()
+            const rect = ctx.root.getBoundingClientRect()
 
             // If the top is off-screen then shift the element with scaling applied
             if (rect.top < 0) {
               const scale =
                 1 -
-                this.root.getBoundingClientRect().height /
-                  this.root.clientHeight
-              const shift = (this.root.clientHeight * scale) / 2
-              this.root.style.top = -shift + 'px'
+                ctx.root.getBoundingClientRect().height / ctx.root.clientHeight
+              const shift = (ctx.root.clientHeight * scale) / 2
+              ctx.root.style.top = -shift + 'px'
             }
           }
 
@@ -109,7 +108,7 @@ const ext = {
                 selectedItem?.click()
                 break
               case 'Escape':
-                this.close()
+                ctx.close()
                 break
             }
           })
@@ -140,7 +139,7 @@ const ext = {
               let top = options.event.clientY - 10
 
               const bodyRect = document.body.getBoundingClientRect()
-              const rootRect = this.root.getBoundingClientRect()
+              const rootRect = ctx.root.getBoundingClientRect()
               if (
                 bodyRect.height &&
                 top > bodyRect.height - rootRect.height - 10
@@ -148,7 +147,7 @@ const ext = {
                 top = Math.max(0, bodyRect.height - rootRect.height - 10)
               }
 
-              this.root.style.top = top + 'px'
+              ctx.root.style.top = top + 'px'
               positionList()
             }
           })

src/extensions/core/noteNode.ts

@@ -1,32 +1,32 @@
 import { LiteGraph, LGraphCanvas } from '@comfyorg/litegraph'
 import { app } from '../../scripts/app'
 import { ComfyWidgets } from '../../scripts/widgets'
+import { LGraphNode } from '@comfyorg/litegraph'
 // Node that add notes to your project
 
 app.registerExtension({
   name: 'Comfy.NoteNode',
   registerCustomNodes() {
-    class NoteNode {
+    class NoteNode extends LGraphNode {
       static category: string
 
       color = LGraphCanvas.node_colors.yellow.color
       bgcolor = LGraphCanvas.node_colors.yellow.bgcolor
       groupcolor = LGraphCanvas.node_colors.yellow.groupcolor
-      properties: { text: string }
-      serialize_widgets: boolean
       isVirtualNode: boolean
       collapsable: boolean
       title_mode: number
 
-      constructor() {
+      constructor(title?: string) {
+        super(title)
         if (!this.properties) {
           this.properties = { text: '' }
         }
         ComfyWidgets.STRING(
-          // @ts-expect-error
-          // Should we extends LGraphNode?
+          // Should we extends LGraphNode?  Yesss
           this,
           '',
+          // @ts-expect-error
           ['', { default: this.properties.text, multiline: true }],
           app
         )
@@ -40,7 +40,6 @@ app.registerExtension({
 
     LiteGraph.registerNodeType(
       'Note',
-      // @ts-expect-error
       Object.assign(NoteNode, {
         title_mode: LiteGraph.NORMAL_TITLE,
         title: 'Note',

src/extensions/core/rerouteNode.ts

@@ -11,11 +11,12 @@ app.registerExtension({
       __outputType?: string
     }
 
-    class RerouteNode {
+    class RerouteNode extends LGraphNode {
       static category: string | undefined
       static defaultVisibility = false
 
-      constructor() {
+      constructor(title?: string) {
+        super(title)
         if (!this.properties) {
           this.properties = {}
         }

src/extensions/core/widgetInputs.ts

@@ -1,8 +1,8 @@
 import { ComfyWidgets, addValueControlWidgets } from '../../scripts/widgets'
 import { app } from '../../scripts/app'
 import { applyTextReplacements } from '../../scripts/utils'
-import { LiteGraph } from '@comfyorg/litegraph'
-import type { LGraphNode, INodeInputSlot, IWidget } from '@comfyorg/litegraph'
+import { LiteGraph, LGraphNode } from '@comfyorg/litegraph'
+import type { INodeInputSlot, IWidget } from '@comfyorg/litegraph'
 
 const CONVERTED_TYPE = 'converted-widget'
 const VALID_TYPES = ['STRING', 'combo', 'number', 'toggle', 'BOOLEAN']
@@ -13,11 +13,12 @@ const TARGET = Symbol() // Used for reroutes to specify the real target widget
 interface PrimitiveNode extends LGraphNode {}
 
 const replacePropertyName = 'Run widget replace on values'
-class PrimitiveNode {
+class PrimitiveNode extends LGraphNode {
   controlValues: any[]
   lastType: string
   static category: string
-  constructor() {
+  constructor(title?: string) {
+    super(title)
     this.addOutput('connect to widget input', '*')
     this.serialize_widgets = true
     this.isVirtualNode = true

src/scripts/app.ts

@@ -2001,8 +2001,15 @@ export class ComfyApp {
 
   async registerNodeDef(nodeId: string, nodeData: ComfyNodeDef) {
     const self = this
-    const node = Object.assign(
-      function ComfyNode() {
+    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
+      static nodeData? = nodeData
+      static category?: string
+
+      constructor(title?: string) {
+        super(title)
         var inputs = nodeData['input']['required']
         if (nodeData['input']['optional'] != undefined) {
           inputs = Object.assign(
@@ -2068,13 +2075,9 @@ export class ComfyApp {
         this.serialize_widgets = true
 
         app.#invokeExtensionsAsync('nodeCreated', this)
-      },
-      {
-        title: nodeData.display_name || nodeData.name,
-        comfyClass: nodeData.name,
-        nodeData
       }
-    )
+    }
+    // @ts-expect-error
     node.prototype.comfyClass = nodeData.name
 
     this.#addNodeContextMenuHandler(node)
@@ -2082,9 +2085,7 @@ export class ComfyApp {
     this.#addNodeKeyHandler(node)
 
     await this.#invokeExtensionsAsync('beforeRegisterNodeDef', node, nodeData)
-    // @ts-expect-error
     LiteGraph.registerNodeType(nodeId, node)
-    // @ts-expect-error
     node.category = nodeData.category
   }
 

src/scripts/domWidget.ts

@@ -260,7 +260,6 @@ export function addDomClippingSetting(): void {
   })
 }
 
-//@ts-ignore
 LGraphNode.prototype.addDOMWidget = function (
   name: string,
   type: string,

src/types/litegraph-augmentation.d.ts

@@ -14,6 +14,13 @@ declare module '@comfyorg/litegraph' {
      * If the node is a frontend only node and should not be serialized into the prompt.
      */
     isVirtualNode?: boolean
+
+    addDOMWidget(
+      name: string,
+      type: string,
+      element: HTMLElement,
+      options: Record<string, any>
+    ): DOMWidget
   }
 
   interface IWidget<TValue = any, TOptions = any> {
@@ -67,4 +74,8 @@ declare module '@comfyorg/litegraph' {
       slotPos: Vector2
     ): number
   }
+
+  interface ContextMenu {
+    root?: HTMLDivElement
+  }
 }