Manage widget definitions with Pinia store (#1510)

* Fix compile

* nit

* Remove extensions.test

* nit

src/stores/extensionStore.ts

@@ -7,6 +7,7 @@ import { useSettingStore } from './settingStore'
 import { app } from '@/scripts/app'
 import { useMenuItemStore } from './menuItemStore'
 import { useBottomPanelStore } from './workspace/bottomPanelStore'
+import { useWidgetStore } from './widgetStore'
 
 export const useExtensionStore = defineStore('extension', () => {
   // For legacy reasons, the name uniquely identifies an extension
@@ -50,6 +51,16 @@ export const useExtensionStore = defineStore('extension', () => {
     useMenuItemStore().loadExtensionMenuCommands(extension)
     useSettingStore().loadExtensionSettings(extension)
     useBottomPanelStore().registerExtensionBottomPanelTabs(extension)
+    if (extension.getCustomWidgets) {
+      // TODO(huchenlei): We should deprecate the async return value of
+      // getCustomWidgets.
+      ;(async () => {
+        if (extension.getCustomWidgets) {
+          const widgets = await extension.getCustomWidgets(app)
+          useWidgetStore().registerCustomWidgets(widgets)
+        }
+      })()
+    }
     /*
      * Extensions are currently stored in both extensionStore and app.extensions.
      * Legacy jest tests still depend on app.extensions being populated.

src/stores/nodeDefStore.ts

@@ -323,20 +323,6 @@ export const useNodeDefStore = defineStore('nodeDef', () => {
     nodeDefsByName.value[nodeDef.name] = nodeDefImpl
     nodeDefsByDisplayName.value[nodeDef.display_name] = nodeDefImpl
   }
-  function getWidgetType(type: string, inputName: string) {
-    if (type === 'COMBO') {
-      return 'COMBO'
-    } else if (`${type}:${inputName}` in widgets.value) {
-      return `${type}:${inputName}`
-    } else if (type in widgets.value) {
-      return type
-    } else {
-      return null
-    }
-  }
-  function inputIsWidget(spec: BaseInputSpec) {
-    return getWidgetType(spec.type, spec.name) !== null
-  }
   function fromLGraphNode(node: LGraphNode): ComfyNodeDefImpl | null {
     // Frontend-only nodes don't have nodeDef
     return nodeDefsByName.value[node.constructor?.nodeData?.name] ?? null
@@ -345,7 +331,6 @@ export const useNodeDefStore = defineStore('nodeDef', () => {
   return {
     nodeDefsByName,
     nodeDefsByDisplayName,
-    widgets,
     showDeprecated,
     showExperimental,
 
@@ -356,8 +341,6 @@ export const useNodeDefStore = defineStore('nodeDef', () => {
 
     updateNodeDefs,
     addNodeDef,
-    getWidgetType,
-    inputIsWidget,
     fromLGraphNode
   }
 })

src/stores/widgetStore.ts

@@ -0,0 +1,45 @@
+import { ComfyWidgets, ComfyWidgetConstructor } from '@/scripts/widgets'
+import { defineStore } from 'pinia'
+import { ref, computed } from 'vue'
+import type { BaseInputSpec } from './nodeDefStore'
+
+export const useWidgetStore = defineStore('widget', () => {
+  const coreWidgets = ComfyWidgets
+  const customWidgets = ref<Record<string, ComfyWidgetConstructor>>({})
+  const widgets = computed(() => ({
+    ...customWidgets.value,
+    ...coreWidgets
+  }))
+
+  function getWidgetType(type: string, inputName: string) {
+    if (type === 'COMBO') {
+      return 'COMBO'
+    } else if (`${type}:${inputName}` in widgets.value) {
+      return `${type}:${inputName}`
+    } else if (type in widgets.value) {
+      return type
+    } else {
+      return null
+    }
+  }
+
+  function inputIsWidget(spec: BaseInputSpec) {
+    return getWidgetType(spec.type, spec.name) !== null
+  }
+
+  function registerCustomWidgets(
+    newWidgets: Record<string, ComfyWidgetConstructor>
+  ) {
+    customWidgets.value = {
+      ...customWidgets.value,
+      ...newWidgets
+    }
+  }
+
+  return {
+    widgets,
+    getWidgetType,
+    inputIsWidget,
+    registerCustomWidgets
+  }
+})