Node source/id badge (#781)

* Add basic node badge

* Node source badge

* Prevent manager badge rendering

* Update litegraph (Badge support)

* Add playwright tests

* Separate nodes

* nit

* Checkout devtools repo for browser test expectation CI

* Fix failing unittests

* Rename setting

* Hide all badges in playwright tests

* Handle group node

* Update test expectations [skip ci]

* Fix unittest

---------

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

src/extensions/core/colorPalette.ts

@@ -5,7 +5,7 @@ import { LGraphCanvas, LiteGraph } from '@comfyorg/litegraph'
 
 // Manage color palettes
 
-const colorPalettes: ColorPalettes = {
+export const colorPalettes: ColorPalettes = {
   dark: {
     id: 'dark',
     name: 'Dark (Default)',
@@ -52,7 +52,10 @@ const colorPalettes: ColorPalettes = {
 
         LINK_COLOR: '#9A9',
         EVENT_LINK_COLOR: '#A86',
-        CONNECTING_LINK_COLOR: '#AFA'
+        CONNECTING_LINK_COLOR: '#AFA',
+
+        BADGE_FG_COLOR: '#FFF',
+        BADGE_BG_COLOR: '#0F1F0F'
       },
       comfy_base: {
         'fg-color': '#fff',
@@ -114,7 +117,10 @@ const colorPalettes: ColorPalettes = {
 
         LINK_COLOR: '#4CAF50',
         EVENT_LINK_COLOR: '#FF9800',
-        CONNECTING_LINK_COLOR: '#2196F3'
+        CONNECTING_LINK_COLOR: '#2196F3',
+
+        BADGE_FG_COLOR: '#000',
+        BADGE_BG_COLOR: '#FFF'
       },
       comfy_base: {
         'fg-color': '#222',

src/extensions/core/contextMenuFilter.ts

@@ -17,9 +17,11 @@ const ext = {
         const filter = document.createElement('input')
         filter.classList.add('comfy-context-menu-filter')
         filter.placeholder = 'Filter list'
+        // @ts-expect-error
         ctx.root.prepend(filter)
 
         const items = Array.from(
+          // @ts-expect-error
           ctx.root.querySelectorAll('.litemenu-entry')
         ) as HTMLElement[]
         let displayedItems = [...items]
@@ -61,14 +63,18 @@ const ext = {
           }
 
           const positionList = () => {
+            // @ts-expect-error
             const rect = ctx.root.getBoundingClientRect()
 
             // If the top is off-screen then shift the element with scaling applied
             if (rect.top < 0) {
               const scale =
                 1 -
+                // @ts-expect-error
                 ctx.root.getBoundingClientRect().height / ctx.root.clientHeight
+              // @ts-expect-error
               const shift = (ctx.root.clientHeight * scale) / 2
+              // @ts-expect-error
               ctx.root.style.top = -shift + 'px'
             }
           }
@@ -139,6 +145,7 @@ const ext = {
               let top = options.event.clientY - 10
 
               const bodyRect = document.body.getBoundingClientRect()
+              // @ts-expect-error
               const rootRect = ctx.root.getBoundingClientRect()
               if (
                 bodyRect.height &&
@@ -147,6 +154,7 @@ const ext = {
                 top = Math.max(0, bodyRect.height - rootRect.height - 10)
               }
 
+              // @ts-expect-error
               ctx.root.style.top = top + 'px'
               positionList()
             }

src/extensions/core/index.ts

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

src/extensions/core/nodeBadge.ts

@@ -0,0 +1,117 @@
+import { app, type ComfyApp } from '@/scripts/app'
+import type { ComfyExtension } from '@/types/comfy'
+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 _ from 'lodash'
+import { colorPalettes } from './colorPalette'
+import { BadgePosition } from '@comfyorg/litegraph'
+import type { Palette } from '@/types/colorPalette'
+
+function getNodeSource(node: ComfyLGraphNode) {
+  const pythonModule = (node.constructor as typeof ComfyLGraphNode).nodeData
+    ?.python_module
+  return pythonModule ? getNodeSourceFromPythonModule(pythonModule) : null
+}
+
+function isCoreNode(node: ComfyLGraphNode) {
+  return getNodeSource(node)?.type === 'core'
+}
+
+function getNodeIdBadge(node: ComfyLGraphNode, nodeIdBadgeMode: NodeBadgeMode) {
+  return nodeIdBadgeMode === NodeBadgeMode.None ||
+    (isCoreNode(node) && nodeIdBadgeMode === NodeBadgeMode.HideBuiltIn)
+    ? ''
+    : `#${node.id}`
+}
+
+function getNodeSourceBadge(
+  node: ComfyLGraphNode,
+  nodeSourceBadgeMode: NodeBadgeMode
+) {
+  const nodeSource = getNodeSource(node)
+  return nodeSourceBadgeMode === NodeBadgeMode.None ||
+    (isCoreNode(node) && nodeSourceBadgeMode === NodeBadgeMode.HideBuiltIn)
+    ? ''
+    : nodeSource?.badgeText ?? ''
+}
+
+class NodeBadgeExtension implements ComfyExtension {
+  name = 'Comfy.NodeBadge'
+
+  constructor(
+    public nodeIdBadgeMode: ComputedRef<NodeBadgeMode> | null = null,
+    public nodeSourceBadgeMode: ComputedRef<NodeBadgeMode> | null = null,
+    public colorPalette: ComputedRef<Palette> | null = null,
+    public defaultColorPalette: Palette | null = null
+  ) {}
+
+  init(app: ComfyApp) {
+    if (!app.vueAppReady) {
+      return
+    }
+
+    const settingStore = useSettingStore()
+    this.nodeSourceBadgeMode = computed(
+      () =>
+        settingStore.get('Comfy.NodeBadge.NodeSourceBadgeMode') as NodeBadgeMode
+    )
+    this.nodeIdBadgeMode = computed(
+      () => settingStore.get('Comfy.NodeBadge.NodeIdBadgeMode') as NodeBadgeMode
+    )
+    this.colorPalette = computed(
+      () => colorPalettes[settingStore.get('Comfy.ColorPalette')]
+    )
+    this.defaultColorPalette = colorPalettes['dark']
+
+    watch(this.nodeSourceBadgeMode, () => {
+      app.graph.setDirtyCanvas(true, true)
+    })
+
+    watch(this.nodeIdBadgeMode, () => {
+      app.graph.setDirtyCanvas(true, true)
+    })
+  }
+
+  nodeCreated(node: ComfyLGraphNode, app: ComfyApp) {
+    if (!app.vueAppReady) {
+      return
+    }
+
+    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(
+            [
+              getNodeIdBadge(node, this.nodeIdBadgeMode.value),
+              getNodeSourceBadge(node, this.nodeSourceBadgeMode.value)
+            ]
+              .filter((s) => s.length > 0)
+              .join(' '),
+            {
+              length: 25
+            }
+          ),
+          fgColor:
+            this.colorPalette.value.colors.litegraph_base?.BADGE_FG_COLOR ||
+            this.defaultColorPalette.colors.litegraph_base.BADGE_FG_COLOR,
+          bgColor:
+            this.colorPalette.value.colors.litegraph_base?.BADGE_BG_COLOR ||
+            this.defaultColorPalette.colors.litegraph_base.BADGE_BG_COLOR
+        })
+    )
+
+    node.badges.push(() => badge.value)
+  }
+}
+
+app.registerExtension(new NodeBadgeExtension())