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/components/graph/GraphCanvas.vue

@@ -32,7 +32,8 @@ import {
   LGraphGroup,
   DragAndScale,
   LGraphCanvas,
-  ContextMenu
+  ContextMenu,
+  LGraphBadge
 } from '@comfyorg/litegraph'
 import type { RenderedTreeExplorerNode } from '@/types/treeExplorerTypes'
 import { useNodeBookmarkStore } from '@/stores/nodeBookmarkStore'
@@ -98,6 +99,7 @@ onMounted(async () => {
   window['DragAndScale'] = DragAndScale
   window['LGraphCanvas'] = LGraphCanvas
   window['ContextMenu'] = ContextMenu
+  window['LGraphBadge'] = LGraphBadge
 
   comfyApp.vueAppReady = true
 

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

src/scripts/app.ts

@@ -51,6 +51,7 @@ import { useToastStore } from '@/stores/toastStore'
 import { ModelStore, useModelStore } from '@/stores/modelStore'
 import type { ToastMessageOptions } from 'primevue/toast'
 import { useWorkspaceStore } from '@/stores/workspaceStateStore'
+import { ComfyLGraphNode } from '@/types/comfyLGraphNode'
 
 export const ANIM_PREVIEW_WIDGET = '$$comfy_animation_preview'
 
@@ -2007,12 +2008,12 @@ export class ComfyApp {
 
   async registerNodeDef(nodeId: string, nodeData: ComfyNodeDef) {
     const self = this
-    const node = class ComfyNode extends LGraphNode {
+    const node: new () => ComfyLGraphNode = class ComfyNode extends LGraphNode {
       static comfyClass? = nodeData.name
       // TODO: change to "title?" once litegraph.d.ts has been updated
       static title = nodeData.display_name || nodeData.name
       static nodeData? = nodeData
-      static category?: string
+      static category: string = nodeData.category
 
       constructor(title?: string) {
         super(title)
@@ -2083,7 +2084,6 @@ export class ComfyApp {
         app.#invokeExtensionsAsync('nodeCreated', this)
       }
     }
-    // @ts-expect-error
     node.prototype.comfyClass = nodeData.name
 
     this.#addNodeContextMenuHandler(node)
@@ -2092,7 +2092,6 @@ export class ComfyApp {
 
     await this.#invokeExtensionsAsync('beforeRegisterNodeDef', node, nodeData)
     LiteGraph.registerNodeType(nodeId, node)
-    node.category = nodeData.category
   }
 
   async registerNodesFromDefs(defs: Record<string, ComfyNodeDef>) {

src/stores/settingStore.ts

@@ -10,6 +10,7 @@
 import { app } from '@/scripts/app'
 import { ComfySettingsDialog } from '@/scripts/ui/settings'
 import { Settings } from '@/types/apiTypes'
+import { NodeBadgeMode } from '@/types/nodeSource'
 import {
   LinkReleaseTriggerAction,
   LinkReleaseTriggerMode
@@ -330,6 +331,22 @@ export const useSettingStore = defineStore('setting', {
         options: ['en', 'zh'],
         defaultValue: navigator.language.split('-')[0] || 'en'
       })
+
+      app.ui.settings.addSetting({
+        id: 'Comfy.NodeBadge.NodeSourceBadgeMode',
+        name: 'Node source badge mode',
+        type: 'combo',
+        options: Object.values(NodeBadgeMode),
+        defaultValue: NodeBadgeMode.HideBuiltIn
+      })
+
+      app.ui.settings.addSetting({
+        id: 'Comfy.NodeBadge.NodeIdBadgeMode',
+        name: 'Node ID badge mode',
+        type: 'combo',
+        options: [NodeBadgeMode.None, NodeBadgeMode.ShowAll],
+        defaultValue: NodeBadgeMode.ShowAll
+      })
     },
 
     set<K extends keyof Settings>(key: K, value: Settings[K]) {

src/types/apiTypes.ts

@@ -3,6 +3,7 @@ import { zComfyWorkflow, zNodeId } from './comfyWorkflow'
 import { fromZodError } from 'zod-validation-error'
 import { colorPalettesSchema } from './colorPalette'
 import { LinkReleaseTriggerAction } from './searchBoxTypes'
+import { NodeBadgeMode } from './nodeSource'
 
 const zNodeType = z.string()
 const zQueueIndex = z.number()
@@ -424,6 +425,10 @@ const zLinkReleaseTriggerAction = z.enum(
   Object.values(LinkReleaseTriggerAction) as [string, ...string[]]
 )
 
+const zNodeBadgeMode = z.enum(
+  Object.values(NodeBadgeMode) as [string, ...string[]]
+)
+
 const zSettings = z.record(z.any()).and(
   z
     .object({
@@ -484,7 +489,9 @@ const zSettings = z.record(z.any()).and(
       'Comfy.Workflow.ModelDownload.AllowedSources': z.array(z.string()),
       'Comfy.Workflow.ModelDownload.AllowedSuffixes': z.array(z.string()),
       'Comfy.Node.DoubleClickTitleToEdit': z.boolean(),
-      'Comfy.Window.UnloadConfirmation': z.boolean()
+      'Comfy.Window.UnloadConfirmation': z.boolean(),
+      'Comfy.NodeBadge.NodeSourceBadgeMode': zNodeBadgeMode,
+      'Comfy.NodeBadge.NodeIdBadgeMode': zNodeBadgeMode
     })
     .optional()
 )

src/types/colorPalette.ts

@@ -52,7 +52,9 @@ const litegraphBaseSchema = z
     WIDGET_SECONDARY_TEXT_COLOR: z.string(),
     LINK_COLOR: z.string(),
     EVENT_LINK_COLOR: z.string(),
-    CONNECTING_LINK_COLOR: z.string()
+    CONNECTING_LINK_COLOR: z.string(),
+    BADGE_FG_COLOR: z.string().optional(),
+    BADGE_BG_COLOR: z.string().optional()
   })
   .passthrough()
 

src/types/comfyLGraphNode.ts

@@ -0,0 +1,11 @@
+import type { LGraphNode } from '@comfyorg/litegraph'
+import type { ComfyNodeDef } from './apiTypes'
+
+export declare class ComfyLGraphNode extends LGraphNode {
+  static comfyClass: string
+  static title: string
+  static nodeData?: ComfyNodeDef
+  static category: string
+
+  constructor(title?: string)
+}

src/types/nodeSource.ts

@@ -3,6 +3,7 @@ export type NodeSource = {
   type: NodeSourceType
   className: string
   displayText: string
+  badgeText: string
 }
 
 export const getNodeSource = (python_module: string): NodeSource => {
@@ -11,15 +12,23 @@ export const getNodeSource = (python_module: string): NodeSource => {
     return {
       type: 'core',
       className: 'comfy-core',
-      displayText: 'Comfy Core'
+      displayText: 'Comfy Core',
+      badgeText: '🦊'
     }
   } else if (modules[0] === 'custom_nodes') {
     return {
       type: 'custom_nodes',
       className: 'comfy-custom-nodes',
-      displayText: modules[1]
+      displayText: modules[1],
+      badgeText: modules[1]
     }
   } else {
     throw new Error(`Unknown node source: ${python_module}`)
   }
 }
+
+export enum NodeBadgeMode {
+  None = 'None',
+  ShowAll = 'Show all',
+  HideBuiltIn = 'Hide built-in'
+}