Use ComfyNodeDefImpl on nodeDefStore (#235)

Author: huchenlei

src/scripts/app.ts

@@ -36,7 +36,7 @@ import { StorageLocation } from '@/types/settingTypes'
 import '@comfyorg/litegraph/css/litegraph.css'
 import '../assets/css/style.css'
 import { ExtensionManager } from '@/types/extensionTypes'
-import { useNodeDefStore } from '@/stores/nodeDefStore'
+import { SYSTEM_NODE_DEFS, useNodeDefStore } from '@/stores/nodeDefStore'
 
 export const ANIM_PREVIEW_WIDGET = '$$comfy_animation_preview'
 
@@ -1968,7 +1968,7 @@ export class ComfyApp {
     await this.registerNodesFromDefs(defs)
     if (this.vueAppReady) {
       const nodeDefStore = useNodeDefStore()
-      nodeDefStore.addNodeDefs(Object.values(defs))
+      nodeDefStore.updateNodeDefs([...Object.values(defs), ...SYSTEM_NODE_DEFS])
       nodeDefStore.updateWidgets(this.widgets)
     }
     await this.#invokeExtensionsAsync('registerCustomNodes')

src/stores/nodeDefStore.ts

@@ -216,41 +216,34 @@ export const SYSTEM_NODE_DEFS: ComfyNodeDef[] = [
   }
 ]
 
-const SYSTEM_NODE_DEFS_BY_NAME = SYSTEM_NODE_DEFS.reduce((acc, nodeDef) => {
-  acc[nodeDef.name] = nodeDef
-  return acc
-}, {}) as Record<string, ComfyNodeDef>
-
 interface State {
-  nodeDefsByName: Record<string, ComfyNodeDef>
+  nodeDefsByName: Record<string, ComfyNodeDefImpl>
   widgets: Record<string, ComfyWidgetConstructor>
 }
 
 export const useNodeDefStore = defineStore('nodeDef', {
   state: (): State => ({
-    nodeDefsByName: SYSTEM_NODE_DEFS_BY_NAME,
+    nodeDefsByName: {},
     widgets: {}
   }),
   getters: {
     nodeDefs(state) {
       return Object.values(state.nodeDefsByName)
     },
     nodeSearchService(state) {
-      return new NodeSearchService(
-        Object.values(state.nodeDefsByName).map((nodeDef) =>
-          plainToClass(ComfyNodeDefImpl, nodeDef)
-        )
-      )
+      return new NodeSearchService(Object.values(state.nodeDefsByName))
     }
   },
   actions: {
-    addNodeDef(nodeDef: ComfyNodeDef) {
-      this.nodeDefsByName[nodeDef.name] = nodeDef
-    },
-    addNodeDefs(nodeDefs: ComfyNodeDef[]) {
+    updateNodeDefs(nodeDefs: ComfyNodeDef[]) {
+      const newNodeDefsByName: { [key: string]: ComfyNodeDefImpl } = {}
       for (const nodeDef of nodeDefs) {
-        this.nodeDefsByName[nodeDef.name] = nodeDef
+        newNodeDefsByName[nodeDef.name] = plainToClass(
+          ComfyNodeDefImpl,
+          nodeDef
+        )
       }
+      this.nodeDefsByName = newNodeDefsByName
     },
     updateWidgets(widgets: Record<string, ComfyWidgetConstructor>) {
       this.widgets = widgets