diff --git a/src/components/graph/GraphCanvas.vue b/src/components/graph/GraphCanvas.vue
index 8327aa058..e4cf189ad 100644
--- a/src/components/graph/GraphCanvas.vue
+++ b/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()
diff --git a/src/components/graph/NodeBadge.vue b/src/components/graph/NodeBadge.vue
deleted file mode 100644
index f52aabb40..000000000
--- a/src/components/graph/NodeBadge.vue
+++ /dev/null
@@ -1,114 +0,0 @@
-<template>
-  <div>
-    <!-- This component does not render anything visible. -->
-  </div>
-</template>
-
-<script setup lang="ts">
-import {
-  BadgePosition,
-  LGraphBadge,
-  type LGraphNode
-} from '@comfyorg/litegraph'
-import _ from 'lodash'
-import { computed, onMounted, watch } from 'vue'
-
-import { app } from '@/scripts/app'
-import { ComfyNodeDefImpl, useNodeDefStore } from '@/stores/nodeDefStore'
-import { useSettingStore } from '@/stores/settingStore'
-import { useColorPaletteStore } from '@/stores/workspace/colorPaletteStore'
-import { NodeBadgeMode } from '@/types/nodeSource'
-
-const settingStore = useSettingStore()
-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(() => {
-  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:
-            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)
-      }
-    }
-  })
-})
-</script>
diff --git a/src/composables/node/useNodeBadge.ts b/src/composables/node/useNodeBadge.ts
new file mode 100644
index 000000000..be629bd03
--- /dev/null
+++ b/src/composables/node/useNodeBadge.ts
@@ -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)
+        }
+      }
+    })
+  })
+}