Shorten node source package name by remove ComfyUI prefix/suffix (#883)

* Shorten node source package name by remove ComfyUI prefix/suffix

* Update test expectations [skip ci]

---------

Co-authored-by: github-actions <github-actions@github.com>

src/components/node/NodeSourceChip.vue

@@ -1,29 +0,0 @@
-<template>
-  <Chip :class="nodeSource.className">
-    {{ nodeSource.displayText }}
-  </Chip>
-</template>
-
-<script setup lang="ts">
-import { getNodeSource } from '@/types/nodeSource'
-import Chip from 'primevue/chip'
-import { computed } from 'vue'
-
-const props = defineProps({
-  python_module: {
-    type: String,
-    required: true
-  }
-})
-
-const nodeSource = computed(() => getNodeSource(props.python_module))
-</script>
-
-<style scoped>
-.comfy-core,
-.comfy-custom-nodes,
-.comfy-unknown {
-  font-size: small;
-  font-weight: lighter;
-}
-</style>

src/components/searchbox/NodeSearchItem.vue

@@ -38,17 +38,20 @@
         :value="formatNumberWithSuffix(nodeFrequency, { roundToInt: true })"
         severity="secondary"
       />
-      <NodeSourceChip
-        v-if="nodeDef.python_module !== undefined"
-        :python_module="nodeDef.python_module"
-      />
+      <Chip
+        v-if="nodeDef.nodeSource.type !== NodeSourceType.Unknown"
+        class="text-sm font-light"
+      >
+        {{ nodeDef.nodeSource.displayText }}
+      </Chip>
     </div>
   </div>
 </template>
 
 <script setup lang="ts">
 import Tag from 'primevue/tag'
-import NodeSourceChip from '@/components/node/NodeSourceChip.vue'
+import Chip from 'primevue/chip'
+import { NodeSourceType } from '@/types/nodeSource'
 import { ComfyNodeDefImpl, useNodeFrequencyStore } from '@/stores/nodeDefStore'
 import { highlightQuery } from '@/utils/formatUtil'
 import { computed } from 'vue'

src/extensions/core/nodeBadge.ts

@@ -4,23 +4,23 @@ import type { ComfyLGraphNode } from '@/types/comfyLGraphNode'
 import { LGraphBadge } from '@comfyorg/litegraph'
 import { useSettingStore } from '@/stores/settingStore'
 import { computed, ComputedRef, watch } from 'vue'
-import {
-  getNodeSource as getNodeSourceFromPythonModule,
-  NodeBadgeMode
-} from '@/types/nodeSource'
+import { NodeBadgeMode, NodeSource, NodeSourceType } from '@/types/nodeSource'
 import _ from 'lodash'
 import { getColorPalette, defaultColorPalette } from './colorPalette'
 import { BadgePosition } from '@comfyorg/litegraph'
 import type { Palette } from '@/types/colorPalette'
+import type { ComfyNodeDef } from '@/types/apiTypes'
+import { useNodeDefStore } from '@/stores/nodeDefStore'
 
-function getNodeSource(node: ComfyLGraphNode) {
-  const pythonModule = (node.constructor as typeof ComfyLGraphNode).nodeData
-    ?.python_module
-  return pythonModule ? getNodeSourceFromPythonModule(pythonModule) : null
+function getNodeSource(node: ComfyLGraphNode): NodeSource | null {
+  const nodeDef = (node.constructor as typeof ComfyLGraphNode)
+    .nodeData as ComfyNodeDef
+  const nodeDefStore = useNodeDefStore()
+  return nodeDefStore.nodeDefsByName[nodeDef.name]?.nodeSource ?? null
 }
 
 function isCoreNode(node: ComfyLGraphNode) {
-  return getNodeSource(node)?.type === 'core'
+  return getNodeSource(node)?.type === NodeSourceType.Core
 }
 
 function badgeTextVisible(

src/services/nodeSearchService.ts

@@ -1,5 +1,4 @@
 import { ComfyNodeDefImpl } from '@/stores/nodeDefStore'
-import { getNodeSource } from '@/types/nodeSource'
 import Fuse, { IFuseOptions, FuseSearchOptions } from 'fuse.js'
 import _ from 'lodash'
 
@@ -202,7 +201,7 @@ export class NodeSourceFilter extends NodeFilter<string> {
   public readonly longInvokeSequence = 'source'
 
   public override getNodeOptions(node: ComfyNodeDefImpl): string[] {
-    return [getNodeSource(node.python_module).displayText]
+    return [node.nodeSource.displayText]
   }
 }
 

src/stores/nodeDefStore.ts

@@ -10,6 +10,7 @@ import { TreeNode } from 'primevue/treenode'
 import { buildTree } from '@/utils/treeUtil'
 import { computed, ref } from 'vue'
 import axios from 'axios'
+import { type NodeSource, getNodeSource } from '@/types/nodeSource'
 
 export class BaseInputSpec<T = any> {
   name: string
@@ -194,6 +195,12 @@ export class ComfyNodeDefImpl {
   @Transform(({ obj }) => ComfyNodeDefImpl.transformOutputSpec(obj))
   output: ComfyOutputsSpec
 
+  @Transform(({ obj }) => getNodeSource(obj.python_module), {
+    toClassOnly: true
+  })
+  @Expose()
+  nodeSource: NodeSource
+
   private static transformOutputSpec(obj: any): ComfyOutputsSpec {
     const { output, output_is_list, output_name, output_tooltips } = obj
     const result = output.map((type: string | any[], index: number) => {

src/types/nodeSource.ts

@@ -1,4 +1,9 @@
-export type NodeSourceType = 'core' | 'custom_nodes'
+export enum NodeSourceType {
+  Core = 'core',
+  CustomNodes = 'custom_nodes',
+  Unknown = 'unknown'
+}
+
 export type NodeSource = {
   type: NodeSourceType
   className: string
@@ -6,24 +11,41 @@ export type NodeSource = {
   badgeText: string
 }
 
-export const getNodeSource = (python_module: string): NodeSource => {
+const UNKNOWN_NODE_SOURCE: NodeSource = {
+  type: NodeSourceType.Unknown,
+  className: 'comfy-unknown',
+  displayText: 'Unknown',
+  badgeText: '?'
+}
+
+const shortenNodeName = (name: string) => {
+  return name
+    .replace(/^(ComfyUI-|ComfyUI_|Comfy-|Comfy_)/, '')
+    .replace(/(-ComfyUI|_ComfyUI|-Comfy|_Comfy)$/, '')
+}
+
+export const getNodeSource = (python_module?: string): NodeSource => {
+  if (!python_module) {
+    return UNKNOWN_NODE_SOURCE
+  }
   const modules = python_module.split('.')
   if (['nodes', 'comfy_extras'].includes(modules[0])) {
     return {
-      type: 'core',
+      type: NodeSourceType.Core,
       className: 'comfy-core',
       displayText: 'Comfy Core',
       badgeText: '🦊'
     }
   } else if (modules[0] === 'custom_nodes') {
+    const displayName = shortenNodeName(modules[1])
     return {
-      type: 'custom_nodes',
+      type: NodeSourceType.CustomNodes,
       className: 'comfy-custom-nodes',
-      displayText: modules[1],
-      badgeText: modules[1]
+      displayText: displayName,
+      badgeText: displayName
     }
   } else {
-    throw new Error(`Unknown node source: ${python_module}`)
+    return UNKNOWN_NODE_SOURCE
   }
 }