Fix ComfyExtension types (#778)

src/extensions/core/clipspace.ts

@@ -170,6 +170,7 @@ export class ClipspaceDialog extends ComfyDialog {
 app.registerExtension({
   name: 'Comfy.Clipspace',
   init(app) {
+    // @ts-expect-error Move to ComfyApp
     app.openClipspace = function () {
       if (!ClipspaceDialog.instance) {
         ClipspaceDialog.instance = new ClipspaceDialog()

src/extensions/core/rerouteNode.ts

@@ -74,6 +74,7 @@ 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) {
@@ -112,6 +113,7 @@ 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') {
@@ -164,6 +166,7 @@ app.registerExtension({
             for (const l of node.outputs[0].links || []) {
               const link = app.graph.links[l]
               if (link) {
+                // @ts-expect-error Fix litegraph types
                 link.color = color
 
                 if (app.configuringGraph) continue
@@ -177,6 +180,7 @@ app.registerExtension({
                   }
                   if (!targetWidget) {
                     targetWidget = targetNode.widgets?.find(
+                      // @ts-expect-error fix widget types
                       (w) => w.name === targetInput.widget.name
                     )
                   }
@@ -209,6 +213,7 @@ app.registerExtension({
           if (inputNode) {
             const link = app.graph.links[inputNode.inputs[0].link]
             if (link) {
+              // @ts-expect-error Fix litegraph types
               link.color = color
             }
           }

src/extensions/core/slotDefaults.ts

@@ -30,8 +30,7 @@ app.registerExtension({
   slot_types_default_in: {},
   async beforeRegisterNodeDef(nodeType, nodeData, app) {
     var nodeId = nodeData.name
-    var inputs = []
-    inputs = nodeData['input']['required'] //only show required inputs to reduce the mess also not logical to create node with optional inputs
+    const inputs = nodeData['input']['required'] //only show required inputs to reduce the mess also not logical to create node with optional inputs
     for (const inputKey in inputs) {
       var input = inputs[inputKey]
       if (typeof input[0] !== 'string') continue

src/extensions/core/widgetInputs.ts

@@ -826,6 +826,7 @@ app.registerExtension({
     }
 
     function isNodeAtPos(pos) {
+      // @ts-expect-error Fix litegraph types
       for (const n of app.graph._nodes) {
         if (n.pos[0] === pos[0] && n.pos[1] === pos[1]) {
           return true

src/scripts/app.ts

@@ -2871,7 +2871,7 @@ export class ComfyApp {
    * Registers a Comfy web extension with the app
    * @param {ComfyExtension} extension
    */
-  registerExtension(extension) {
+  registerExtension(extension: ComfyExtension) {
     if (!extension.name) {
       throw new Error("Extensions must have a 'name' property.")
     }

src/types/comfy.d.ts

@@ -2,6 +2,16 @@ import { LGraphNode, IWidget } from './litegraph'
 import { ComfyApp } from '../scripts/app'
 import type { ComfyNodeDef } from '@/types/apiTypes'
 
+export type Widgets = Record<
+  string,
+  (
+    node,
+    inputName,
+    inputData,
+    app?: ComfyApp
+  ) => { widget?: IWidget; minWidth?: number; minHeight?: number }
+>
+
 export interface ComfyExtension {
   /**
    * The name of the extension
@@ -11,12 +21,12 @@ export interface ComfyExtension {
    * Allows any initialisation, e.g. loading resources. Called after the canvas is created but before nodes are added
    * @param app The ComfyUI app instance
    */
-  init?(app: ComfyApp): Promise<void>
+  init?(app: ComfyApp): Promise<void> | void
   /**
    * Allows any additional setup, called after the application is fully set up and running
    * @param app The ComfyUI app instance
    */
-  setup?(app: ComfyApp): Promise<void>
+  setup?(app: ComfyApp): Promise<void> | void
   /**
    * Called before nodes are registered with the graph
    * @param defs The collection of node definitions, add custom ones or edit existing ones
@@ -25,25 +35,13 @@ export interface ComfyExtension {
   addCustomNodeDefs?(
     defs: Record<string, ComfyObjectInfo>,
     app: ComfyApp
-  ): Promise<void>
+  ): Promise<void> | void
   /**
    * Allows the extension to add custom widgets
    * @param app The ComfyUI app instance
    * @returns An array of {[widget name]: widget data}
    */
-  getCustomWidgets?(
-    app: ComfyApp
-  ): Promise<
-    Record<
-      string,
-      (
-        node,
-        inputName,
-        inputData,
-        app
-      ) => { widget?: IWidget; minWidth?: number; minHeight?: number }
-    >
-  >
+  getCustomWidgets?(app: ComfyApp): Promise<Widgets> | Widgets
   /**
    * Allows the extension to add additional handling to the node before it is registered with **LGraph**
    * @param nodeType The node class (not an instance)
@@ -54,7 +52,7 @@ export interface ComfyExtension {
     nodeType: typeof LGraphNode,
     nodeData: ComfyObjectInfo,
     app: ComfyApp
-  ): Promise<void>
+  ): Promise<void> | void
 
   /**
    * Allows the extension to modify the node definitions before they are used in the Vue app
@@ -71,21 +69,23 @@ export interface ComfyExtension {
    *
    * @param app The ComfyUI app instance
    */
-  registerCustomNodes?(app: ComfyApp): Promise<void>
+  registerCustomNodes?(app: ComfyApp): Promise<void> | void
   /**
    * Allows the extension to modify a node that has been reloaded onto the graph.
    * If you break something in the backend and want to patch workflows in the frontend
    * This is the place to do this
    * @param node The node that has been loaded
    * @param app The ComfyUI app instance
    */
-  loadedGraphNode?(node: LGraphNode, app: ComfyApp)
+  loadedGraphNode?(node: LGraphNode, app: ComfyApp): void
   /**
    * Allows the extension to run code after the constructor of the node
    * @param node The node that has been created
    * @param app The ComfyUI app instance
    */
-  nodeCreated?(node: LGraphNode, app: ComfyApp)
+  nodeCreated?(node: LGraphNode, app: ComfyApp): void
+
+  [key: string]: any
 }
 
 export type ComfyObjectInfo = {