[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/extensions/core/nodeBadge.ts

@@ -1,150 +0,0 @@
-// @ts-strict-ignore
-import { app, type ComfyApp } from '@/scripts/app'
-import type { ComfyExtension } from '@/types/comfy'
-import type { LGraphNode } from '@comfyorg/litegraph'
-import { LGraphBadge } from '@comfyorg/litegraph'
-import { useSettingStore } from '@/stores/settingStore'
-import { computed, ComputedRef, watch } from 'vue'
-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 { useNodeDefStore } from '@/stores/nodeDefStore'
-
-function getNodeSource(node: LGraphNode): NodeSource | null {
-  const nodeDef = node.constructor.nodeData
-  // Frontend-only nodes don't have nodeDef
-  if (!nodeDef) {
-    return null
-  }
-  const nodeDefStore = useNodeDefStore()
-  return nodeDefStore.nodeDefsByName[nodeDef.name]?.nodeSource ?? null
-}
-
-function isCoreNode(node: LGraphNode) {
-  return getNodeSource(node)?.type === NodeSourceType.Core
-}
-
-function badgeTextVisible(node: LGraphNode, badgeMode: NodeBadgeMode): boolean {
-  return (
-    badgeMode === NodeBadgeMode.None ||
-    (isCoreNode(node) && badgeMode === NodeBadgeMode.HideBuiltIn)
-  )
-}
-
-function getNodeIdBadgeText(node: LGraphNode, nodeIdBadgeMode: NodeBadgeMode) {
-  return badgeTextVisible(node, nodeIdBadgeMode) ? '' : `#${node.id}`
-}
-
-function getNodeSourceBadgeText(
-  node: LGraphNode,
-  nodeSourceBadgeMode: NodeBadgeMode
-) {
-  const nodeSource = getNodeSource(node)
-  return badgeTextVisible(node, nodeSourceBadgeMode)
-    ? ''
-    : nodeSource?.badgeText ?? ''
-}
-
-function getNodeLifeCycleBadgeText(
-  node: LGraphNode,
-  nodeLifeCycleBadgeMode: NodeBadgeMode
-) {
-  let text = ''
-  const nodeDef = node.constructor.nodeData
-
-  // Frontend-only nodes don't have nodeDef
-  if (!nodeDef) {
-    return ''
-  }
-
-  if (nodeDef.deprecated) {
-    text = '[DEPR]'
-  }
-
-  if (nodeDef.experimental) {
-    text = '[BETA]'
-  }
-
-  return badgeTextVisible(node, nodeLifeCycleBadgeMode) ? '' : text
-}
-
-class NodeBadgeExtension implements ComfyExtension {
-  name = 'Comfy.NodeBadge'
-
-  constructor(
-    public nodeIdBadgeMode: ComputedRef<NodeBadgeMode> | null = null,
-    public nodeSourceBadgeMode: ComputedRef<NodeBadgeMode> | null = null,
-    public nodeLifeCycleBadgeMode: ComputedRef<NodeBadgeMode> | null = null,
-    public colorPalette: ComputedRef<Palette> | null = null
-  ) {}
-
-  init(app: ComfyApp) {
-    const settingStore = useSettingStore()
-    this.nodeSourceBadgeMode = computed(
-      () =>
-        settingStore.get('Comfy.NodeBadge.NodeSourceBadgeMode') as NodeBadgeMode
-    )
-    this.nodeIdBadgeMode = computed(
-      () => settingStore.get('Comfy.NodeBadge.NodeIdBadgeMode') as NodeBadgeMode
-    )
-    this.nodeLifeCycleBadgeMode = computed(
-      () =>
-        settingStore.get(
-          'Comfy.NodeBadge.NodeLifeCycleBadgeMode'
-        ) as NodeBadgeMode
-    )
-    this.colorPalette = computed(() =>
-      getColorPalette(settingStore.get('Comfy.ColorPalette'))
-    )
-
-    watch(this.nodeSourceBadgeMode, () => {
-      app.graph.setDirtyCanvas(true, true)
-    })
-
-    watch(this.nodeIdBadgeMode, () => {
-      app.graph.setDirtyCanvas(true, true)
-    })
-    watch(this.nodeLifeCycleBadgeMode, () => {
-      app.graph.setDirtyCanvas(true, true)
-    })
-  }
-
-  nodeCreated(node: LGraphNode, app: ComfyApp) {
-    node.badgePosition = BadgePosition.TopRight
-    // @ts-expect-error Disable ComfyUI-Manager's badge drawing by setting badge_enabled to true. Remove this when ComfyUI-Manager's badge drawing is removed.
-    node.badge_enabled = true
-
-    const badge = computed(
-      () =>
-        new LGraphBadge({
-          text: _.truncate(
-            [
-              getNodeIdBadgeText(node, this.nodeIdBadgeMode.value),
-              getNodeLifeCycleBadgeText(
-                node,
-                this.nodeLifeCycleBadgeMode.value
-              ),
-              getNodeSourceBadgeText(node, this.nodeSourceBadgeMode.value)
-            ]
-              .filter((s) => s.length > 0)
-              .join(' '),
-            {
-              length: 31
-            }
-          ),
-          fgColor:
-            this.colorPalette.value.colors.litegraph_base?.BADGE_FG_COLOR ||
-            defaultColorPalette.colors.litegraph_base.BADGE_FG_COLOR,
-          bgColor:
-            this.colorPalette.value.colors.litegraph_base?.BADGE_BG_COLOR ||
-            defaultColorPalette.colors.litegraph_base.BADGE_BG_COLOR
-        })
-    )
-
-    node.badges.push(() => badge.value)
-  }
-}
-
-app.registerExtension(new NodeBadgeExtension())

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
     }
   }
 })