[Refactor] Make node badge a vue component (#1317)

* [Refactor] Make node badge a vue component

* Simplify badge logic

* nit

src/components/graph/GraphCanvas.vue

@@ -17,6 +17,7 @@
   </teleport>
   <NodeSearchboxPopover />
   <NodeTooltip v-if="tooltipEnabled" />
+  <NodeBadge />
 </template>
 
 <script setup lang="ts">
@@ -26,6 +27,7 @@ import BottomPanel from '@/components/bottomPanel/BottomPanel.vue'
 import LiteGraphCanvasSplitterOverlay from '@/components/LiteGraphCanvasSplitterOverlay.vue'
 import NodeSearchboxPopover from '@/components/searchbox/NodeSearchBoxPopover.vue'
 import NodeTooltip from '@/components/graph/NodeTooltip.vue'
+import NodeBadge from '@/components/graph/NodeBadge.vue'
 import { ref, computed, onMounted, watchEffect } from 'vue'
 import { app as comfyApp } from '@/scripts/app'
 import { useSettingStore } from '@/stores/settingStore'

src/components/graph/NodeBadge.vue

@@ -0,0 +1,93 @@
+<template>
+  <div>
+    <!-- This component does not render anything visible. -->
+  </div>
+</template>
+
+<script setup lang="ts">
+import { computed, onMounted, watch } from 'vue'
+import { useSettingStore } from '@/stores/settingStore'
+import {
+  defaultColorPalette,
+  getColorPalette
+} from '@/extensions/core/colorPalette'
+import { app } from '@/scripts/app'
+import type { LGraphNode } from '@comfyorg/litegraph'
+import { BadgePosition } from '@comfyorg/litegraph'
+import { LGraphBadge } from '@comfyorg/litegraph'
+import _ from 'lodash'
+import { NodeBadgeMode } from '@/types/nodeSource'
+import { ComfyNodeDefImpl, useNodeDefStore } from '@/stores/nodeDefStore'
+
+const settingStore = useSettingStore()
+const nodeSourceBadgeMode = computed(
+  () => settingStore.get('Comfy.NodeBadge.NodeSourceBadgeMode') as NodeBadgeMode
+)
+const nodeIdBadgeMode = computed(
+  () => settingStore.get('Comfy.NodeBadge.NodeIdBadgeMode') as NodeBadgeMode
+)
+const nodeLifeCycleBadgeMode = computed(
+  () =>
+    settingStore.get('Comfy.NodeBadge.NodeLifeCycleBadgeMode') as NodeBadgeMode
+)
+
+watch([nodeSourceBadgeMode, nodeIdBadgeMode, nodeLifeCycleBadgeMode], () => {
+  app.graph?.setDirtyCanvas(true, true)
+})
+
+const colorPalette = computed(() =>
+  getColorPalette(settingStore.get('Comfy.ColorPalette'))
+)
+
+const nodeDefStore = useNodeDefStore()
+function badgeTextVisible(
+  nodeDef: ComfyNodeDefImpl | null,
+  badgeMode: NodeBadgeMode
+): boolean {
+  return !(
+    badgeMode === NodeBadgeMode.None ||
+    (nodeDef?.isCoreNode && badgeMode === NodeBadgeMode.HideBuiltIn)
+  )
+}
+
+onMounted(() => {
+  app.registerExtension({
+    name: 'Comfy.NodeBadge',
+    nodeCreated(node: LGraphNode) {
+      node.badgePosition = BadgePosition.TopRight
+
+      const badge = computed(() => {
+        const nodeDef = nodeDefStore.fromLGraphNode(node)
+        return new LGraphBadge({
+          text: _.truncate(
+            [
+              badgeTextVisible(nodeDef, nodeIdBadgeMode.value)
+                ? `#${node.id}`
+                : '',
+              badgeTextVisible(nodeDef, nodeLifeCycleBadgeMode.value)
+                ? nodeDef?.nodeLifeCycleBadgeText ?? ''
+                : '',
+              badgeTextVisible(nodeDef, nodeSourceBadgeMode.value)
+                ? nodeDef?.nodeSource?.badgeText ?? ''
+                : ''
+            ]
+              .filter((s) => s.length > 0)
+              .join(' '),
+            {
+              length: 31
+            }
+          ),
+          fgColor:
+            colorPalette.value.colors.litegraph_base?.BADGE_FG_COLOR ||
+            defaultColorPalette.colors.litegraph_base.BADGE_FG_COLOR,
+          bgColor:
+            colorPalette.value.colors.litegraph_base?.BADGE_BG_COLOR ||
+            defaultColorPalette.colors.litegraph_base.BADGE_BG_COLOR
+        })
+      })
+
+      node.badges.push(() => badge.value)
+    }
+  })
+})
+</script>

src/components/graph/TitleEditor.vue

@@ -106,7 +106,6 @@ const canvasEventHandler = (event: LiteGraphCanvasEvent) => {
     const [x, y] = group.pos
 
     const e = event.detail.originalEvent
-    // @ts-expect-error LiteGraphCanvasEvent is not typed
     const relativeY = e.canvasY - y
     // Only allow editing if the click is on the title bar
     if (relativeY > group.titleHeight) {

src/extensions/core/index.ts

@@ -21,4 +21,3 @@ import './uploadImage'
 import './webcamCapture'
 import './widgetInputs'
 import './uploadAudio'
-import './nodeBadge'

src/stores/nodeDefStore.ts

@@ -13,7 +13,12 @@ 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'
+import {
+  type NodeSource,
+  NodeSourceType,
+  getNodeSource
+} from '@/types/nodeSource'
+import type { LGraphNode } from '@comfyorg/litegraph'
 
 export interface BaseInputSpec<T = any> {
   name: string
@@ -205,6 +210,16 @@ export class ComfyNodeDefImpl {
     const nodeFrequency = nodeFrequencyStore.getNodeFrequencyByName(this.name)
     return [scores[0], -nodeFrequency, ...scores.slice(1)]
   }
+
+  get isCoreNode(): boolean {
+    return this.nodeSource.type === NodeSourceType.Core
+  }
+
+  get nodeLifeCycleBadgeText(): string {
+    if (this.deprecated) return '[DEPR]'
+    if (this.experimental) return '[BETA]'
+    return ''
+  }
 }
 
 export const SYSTEM_NODE_DEFS: Record<string, ComfyNodeDef> = {
@@ -331,6 +346,10 @@ export const useNodeDefStore = defineStore('nodeDef', {
     },
     inputIsWidget(spec: BaseInputSpec) {
       return this.getWidgetType(spec.type, spec.name) !== null
+    },
+    fromLGraphNode(node: LGraphNode): ComfyNodeDefImpl | null {
+      // Frontend-only nodes don't have nodeDef
+      return this.nodeDefsByName[node.constructor.nodeData?.name] ?? null
     }
   }
 })