[Refactor] Convert NodeBadge.vue to composable (#3883)

Author: huchenlei

src/components/graph/GraphCanvas.vue

@@ -27,7 +27,6 @@
     class="w-full h-full touch-none"
   />
 
-  <NodeBadge />
   <NodeTooltip v-if="tooltipEnabled" />
   <NodeSearchboxPopover />
 
@@ -53,7 +52,6 @@ import BottomPanel from '@/components/bottomPanel/BottomPanel.vue'
 import SubgraphBreadcrumb from '@/components/breadcrumb/SubgraphBreadcrumb.vue'
 import DomWidgets from '@/components/graph/DomWidgets.vue'
 import GraphCanvasMenu from '@/components/graph/GraphCanvasMenu.vue'
-import NodeBadge from '@/components/graph/NodeBadge.vue'
 import NodeTooltip from '@/components/graph/NodeTooltip.vue'
 import SelectionOverlay from '@/components/graph/SelectionOverlay.vue'
 import SelectionToolbox from '@/components/graph/SelectionToolbox.vue'
@@ -62,6 +60,7 @@ import NodeSearchboxPopover from '@/components/searchbox/NodeSearchBoxPopover.vu
 import SideToolbar from '@/components/sidebar/SideToolbar.vue'
 import SecondRowWorkflowTabs from '@/components/topbar/SecondRowWorkflowTabs.vue'
 import { useChainCallback } from '@/composables/functional/useChainCallback'
+import { useNodeBadge } from '@/composables/node/useNodeBadge'
 import { useCanvasDrop } from '@/composables/useCanvasDrop'
 import { useContextMenuTranslation } from '@/composables/useContextMenuTranslation'
 import { useCopy } from '@/composables/useCopy'
@@ -254,6 +253,7 @@ const workflowPersistence = useWorkflowPersistence()
 // @ts-expect-error fixme ts strict error
 useCanvasDrop(canvasRef)
 useLitegraphSettings()
+useNodeBadge()
 
 onMounted(async () => {
   useGlobalLitegraph()

src/components/graph/NodeBadge.vue

@@ -0,0 +1,122 @@
+import {
+  BadgePosition,
+  LGraphBadge,
+  type LGraphNode
+} from '@comfyorg/litegraph'
+import _ from 'lodash'
+import { computed, onMounted, watch } from 'vue'
+
+import { app } from '@/scripts/app'
+import { useExtensionStore } from '@/stores/extensionStore'
+import { ComfyNodeDefImpl, useNodeDefStore } from '@/stores/nodeDefStore'
+import { useSettingStore } from '@/stores/settingStore'
+import { useColorPaletteStore } from '@/stores/workspace/colorPaletteStore'
+import { NodeBadgeMode } from '@/types/nodeSource'
+
+/**
+ * Add LGraphBadge to LGraphNode based on settings.
+ *
+ * Following badges are added:
+ * - Node ID badge
+ * - Node source badge
+ * - Node life cycle badge
+ * - API node credits badge
+ */
+export const useNodeBadge = () => {
+  const settingStore = useSettingStore()
+  const extensionStore = useExtensionStore()
+  const colorPaletteStore = useColorPaletteStore()
+
+  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 nodeDefStore = useNodeDefStore()
+  function badgeTextVisible(
+    nodeDef: ComfyNodeDefImpl | null,
+    badgeMode: NodeBadgeMode
+  ): boolean {
+    return !(
+      badgeMode === NodeBadgeMode.None ||
+      (nodeDef?.isCoreNode && badgeMode === NodeBadgeMode.HideBuiltIn)
+    )
+  }
+
+  onMounted(() => {
+    extensionStore.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:
+              colorPaletteStore.completedActivePalette.colors.litegraph_base
+                .BADGE_FG_COLOR,
+            bgColor:
+              colorPaletteStore.completedActivePalette.colors.litegraph_base
+                .BADGE_BG_COLOR
+          })
+        })
+
+        node.badges.push(() => badge.value)
+
+        if (node.constructor.nodeData?.api_node) {
+          const creditsBadge = computed(() => {
+            return new LGraphBadge({
+              text: '',
+              iconOptions: {
+                unicode: '\ue96b',
+                fontFamily: 'PrimeIcons',
+                color: '#FABC25',
+                bgColor: '#353535',
+                fontSize: 8
+              },
+              fgColor:
+                colorPaletteStore.completedActivePalette.colors.litegraph_base
+                  .BADGE_FG_COLOR,
+              bgColor:
+                colorPaletteStore.completedActivePalette.colors.litegraph_base
+                  .BADGE_BG_COLOR
+            })
+          })
+
+          node.badges.push(() => creditsBadge.value)
+        }
+      }
+    })
+  })
+}