From 8e1d3f3baa8496a578df8557bdcf889a4fba65ea Mon Sep 17 00:00:00 2001
From: Chenlei Hu <chenlei.hu@mail.utoronto.ca>
Date: Thu, 25 Jul 2024 23:15:03 -0400
Subject: [PATCH] Use ComfyNodeDefImpl on NodePreview component (#228)

* store widgets

* Use node def impl
---
 src/components/NodePreview.vue   | 64 ++++++--------------------------
 src/components/NodeSearchBox.vue |  5 ++-
 src/scripts/app.ts               |  8 ++--
 src/stores/nodeDefStore.ts       | 26 ++++++++++++-
 4 files changed, 45 insertions(+), 58 deletions(-)

diff --git a/src/components/NodePreview.vue b/src/components/NodePreview.vue
index a32ccfb2e..6ad7149f4 100644
--- a/src/components/NodePreview.vue
+++ b/src/components/NodePreview.vue
@@ -22,7 +22,7 @@ https://github.com/Nuked88/ComfyUI-N-Sidebar/blob/7ae7da4a9761009fb6629bc04c6830
         <div class="_sb_col">{{ slotInput ? slotInput.name : '' }}</div>
         <div class="_sb_col middle-column"></div>
         <div class="_sb_col _sb_inherit">
-          {{ slotOutput ? slotOutput.name : '' }}
+          {{ slotOutput ? slotOutput.display_name : '' }}
         </div>
         <div class="_sb_col">
           <div v-if="slotOutput" :class="['_sb_dot', slotOutput.type]"></div>
@@ -34,7 +34,7 @@ https://github.com/Nuked88/ComfyUI-N-Sidebar/blob/7ae7da4a9761009fb6629bc04c6830
         <div class="_sb_col _sb_arrow">&#x25C0;</div>
         <div class="_sb_col">{{ widgetInput.name }}</div>
         <div class="_sb_col middle-column"></div>
-        <div class="_sb_col _sb_inherit">{{ widgetInput.defaultValue }}</div>
+        <div class="_sb_col _sb_inherit">{{ widgetInput.default }}</div>
         <div class="_sb_col _sb_arrow">&#x25B6;</div>
       </div>
     </div>
@@ -45,67 +45,27 @@ https://github.com/Nuked88/ComfyUI-N-Sidebar/blob/7ae7da4a9761009fb6629bc04c6830
 </template>
 
 <script setup lang="ts">
-import { app } from '@/scripts/app'
-import { type ComfyNodeDef } from '@/types/apiTypes'
+import { ComfyNodeDefImpl, useNodeDefStore } from '@/stores/nodeDefStore'
 import _ from 'lodash'
-import { PropType } from 'vue'
 
 const props = defineProps({
   nodeDef: {
-    type: Object as PropType<ComfyNodeDef>,
+    type: ComfyNodeDefImpl,
     required: true
   }
 })
 
-const nodeDef = props.nodeDef as ComfyNodeDef
+const nodeDefStore = useNodeDefStore()
 
-// --------------------------------------------------
-// TODO: Move out to separate file
-interface IComfyNodeInputDef {
-  name: string
-  type: string
-  widgetType: string | null
-  defaultValue: any
-}
-
-interface IComfyNodeOutputDef {
-  name: string | null
-  type: string
-  isList: boolean
-}
-
-const allInputs = Object.assign(
-  {},
-  nodeDef.input.required || {},
-  nodeDef.input.optional || {}
+const nodeDef = props.nodeDef
+const allInputDefs = nodeDef.input.all
+const allOutputDefs = Object.values(nodeDef.output)
+const slotInputDefs = allInputDefs.filter(
+  (input) => !nodeDefStore.inputIsWidget(input)
 )
-const allInputDefs: IComfyNodeInputDef[] = Object.entries(allInputs).map(
-  ([inputName, inputData]) => {
-    return {
-      name: inputName,
-      type: inputData[0],
-      widgetType: app.getWidgetType(inputData, inputName),
-      defaultValue:
-        inputData[1]?.default ||
-        (inputData[0] instanceof Array ? inputData[0][0] : '')
-    }
-  }
+const widgetInputDefs = allInputDefs.filter((input) =>
+  nodeDefStore.inputIsWidget(input)
 )
-
-const allOutputDefs: IComfyNodeOutputDef[] = _.zip(
-  nodeDef.output,
-  nodeDef.output_name || [],
-  nodeDef.output_is_list || []
-).map(([outputType, outputName, isList]) => {
-  return {
-    name: outputName,
-    type: outputType instanceof Array ? 'COMBO' : outputType,
-    isList: isList
-  }
-})
-
-const slotInputDefs = allInputDefs.filter((input) => !input.widgetType)
-const widgetInputDefs = allInputDefs.filter((input) => !!input.widgetType)
 </script>
 
 <style scoped>
diff --git a/src/components/NodeSearchBox.vue b/src/components/NodeSearchBox.vue
index 4bac5c1db..1979a56a4 100644
--- a/src/components/NodeSearchBox.vue
+++ b/src/components/NodeSearchBox.vue
@@ -2,7 +2,7 @@
   <div class="comfy-vue-node-search-container">
     <div class="comfy-vue-node-preview-container">
       <NodePreview
-        :nodeDef="hoveredSuggestion"
+        :nodeDef="plainToClass(ComfyNodeDefImpl, hoveredSuggestion)"
         :key="hoveredSuggestion?.name || ''"
         v-if="hoveredSuggestion"
       />
@@ -62,7 +62,8 @@ import NodeSourceChip from '@/components/NodeSourceChip.vue'
 import { ComfyNodeDef } from '@/types/apiTypes'
 import { type FilterAndValue } from '@/services/nodeSearchService'
 import NodePreview from './NodePreview.vue'
-import { useNodeDefStore } from '@/stores/nodeDefStore'
+import { ComfyNodeDefImpl, useNodeDefStore } from '@/stores/nodeDefStore'
+import { plainToClass } from 'class-transformer'
 
 const props = defineProps({
   filters: {
diff --git a/src/scripts/app.ts b/src/scripts/app.ts
index 6f25ab8b4..302d5ee58 100644
--- a/src/scripts/app.ts
+++ b/src/scripts/app.ts
@@ -1965,10 +1965,12 @@ export class ComfyApp {
   async registerNodes() {
     // Load node definitions from the backend
     const defs = await api.getNodeDefs()
-    if (this.vueAppReady) {
-      useNodeDefStore().addNodeDefs(Object.values(defs))
-    }
     await this.registerNodesFromDefs(defs)
+    if (this.vueAppReady) {
+      const nodeDefStore = useNodeDefStore()
+      nodeDefStore.addNodeDefs(Object.values(defs))
+      nodeDefStore.updateWidgets(this.widgets)
+    }
     await this.#invokeExtensionsAsync('registerCustomNodes')
   }
 
diff --git a/src/stores/nodeDefStore.ts b/src/stores/nodeDefStore.ts
index 20067e30b..13af2fed2 100644
--- a/src/stores/nodeDefStore.ts
+++ b/src/stores/nodeDefStore.ts
@@ -2,6 +2,7 @@ import { NodeSearchService } from '@/services/nodeSearchService'
 import { ComfyNodeDef } from '@/types/apiTypes'
 import { defineStore } from 'pinia'
 import { Type, Transform, plainToClass } from 'class-transformer'
+import { ComfyWidgetConstructor } from '@/scripts/widgets'
 
 export class BaseInputSpec<T = any> {
   name: string
@@ -124,6 +125,10 @@ export class ComfyInputsSpec {
         return plainToClass(CustomInputSpec, { name, type, ...spec })
     }
   }
+
+  get all() {
+    return [...Object.values(this.required), ...Object.values(this.optional)]
+  }
 }
 
 export class ComfyOutputSpec {
@@ -224,11 +229,13 @@ const SYSTEM_NODE_DEFS_BY_NAME = SYSTEM_NODE_DEFS.reduce((acc, nodeDef) => {
 
 interface State {
   nodeDefsByName: Record<string, ComfyNodeDef>
+  widgets: Record<string, ComfyWidgetConstructor>
 }
 
 export const useNodeDefStore = defineStore('nodeDef', {
   state: (): State => ({
-    nodeDefsByName: SYSTEM_NODE_DEFS_BY_NAME
+    nodeDefsByName: SYSTEM_NODE_DEFS_BY_NAME,
+    widgets: {}
   }),
   getters: {
     nodeDefs(state) {
@@ -246,6 +253,23 @@ export const useNodeDefStore = defineStore('nodeDef', {
       for (const nodeDef of nodeDefs) {
         this.nodeDefsByName[nodeDef.name] = nodeDef
       }
+    },
+    updateWidgets(widgets: Record<string, ComfyWidgetConstructor>) {
+      this.widgets = widgets
+    },
+    getWidgetType(type: string, inputName: string) {
+      if (type === 'COMBO') {
+        return 'COMBO'
+      } else if (`${type}:${inputName}` in this.widgets) {
+        return `${type}:${inputName}`
+      } else if (type in this.widgets) {
+        return type
+      } else {
+        return null
+      }
+    },
+    inputIsWidget(spec: BaseInputSpec) {
+      return this.getWidgetType(spec.type, spec.name) !== null
     }
   }
 })