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>

.github/workflows/test-browser-exp.yaml

@@ -24,6 +24,11 @@ jobs:
         repository: "Comfy-Org/ComfyUI_frontend"
         path: "ComfyUI_frontend"
         ref: ${{ github.head_ref }}
+    - name: Checkout ComfyUI_devtools
+      uses: actions/checkout@v4
+      with:
+        repository: "Comfy-Org/ComfyUI_devtools"
+        path: "ComfyUI/custom_nodes/ComfyUI_devtools"
     - uses: actions/setup-node@v3
       with:
         node-version: lts/*

browser_tests/ComfyPage.ts

@@ -3,7 +3,7 @@ import { test as base } from '@playwright/test'
 import dotenv from 'dotenv'
 dotenv.config()
 import * as fs from 'fs'
-import * as path from 'path'
+import { NodeBadgeMode } from '../src/types/nodeSource'
 
 interface Position {
   x: number
@@ -202,6 +202,13 @@ export class ComfyPage {
     // Reset view to force re-rendering of canvas. So that info fields like fps
     // become hidden.
     await this.resetView()
+
+    // Hide all badges by default.
+    await this.setSetting('Comfy.NodeBadge.NodeIdBadgeMode', NodeBadgeMode.None)
+    await this.setSetting(
+      'Comfy.NodeBadge.NodeSourceBadgeMode',
+      NodeBadgeMode.None
+    )
   }
 
   public assetPath(fileName: string) {

browser_tests/assets/execution_error.json

@@ -5,10 +5,10 @@
     {
       "id": 14,
       "type": "PreviewImage",
-      "pos": [
-        858,
-        -41
-      ],
+      "pos": {
+        "0": 300,
+        "1": 60
+      },
       "size": {
         "0": 213.8594970703125,
         "1": 50.65289306640625
@@ -23,6 +23,7 @@
           "link": 15
         }
       ],
+      "outputs": [],
       "properties": {
         "Node name for S&R": "PreviewImage"
       }
@@ -30,10 +31,10 @@
     {
       "id": 17,
       "type": "DevToolsErrorRaiseNode",
-      "pos": [
-        477,
-        -40
-      ],
+      "pos": {
+        "0": 20,
+        "1": 60
+      },
       "size": {
         "0": 210,
         "1": 26
@@ -41,6 +42,7 @@
       "flags": {},
       "order": 0,
       "mode": 0,
+      "inputs": [],
       "outputs": [
         {
           "name": "IMAGE",
@@ -71,10 +73,10 @@
   "config": {},
   "extra": {
     "ds": {
-      "scale": 1.2100000000000006,
+      "scale": 1,
       "offset": [
-        -266.1038310281165,
-        337.94335447664554
+        117.20766722169206,
+        472.69035116826046
       ]
     }
   },

browser_tests/nodeBadge.spec.ts

@@ -0,0 +1,79 @@
+import { expect } from '@playwright/test'
+import { comfyPageFixture as test } from './ComfyPage'
+import type { ComfyApp } from '../src/scripts/app'
+import { NodeBadgeMode } from '../src/types/nodeSource'
+
+test.describe('Node Badge', () => {
+  test('Can add badge', async ({ comfyPage }) => {
+    await comfyPage.page.evaluate(() => {
+      const LGraphBadge = window['LGraphBadge']
+      const app = window['app'] as ComfyApp
+      const graph = app.graph
+      // @ts-expect-error - accessing private property
+      const nodes = graph._nodes
+
+      for (const node of nodes) {
+        node.badges = [new LGraphBadge({ text: 'Test Badge' })]
+      }
+
+      graph.setDirtyCanvas(true, true)
+    })
+
+    await expect(comfyPage.canvas).toHaveScreenshot('node-badge.png')
+  })
+
+  test('Can add multiple badges', async ({ comfyPage }) => {
+    await comfyPage.page.evaluate(() => {
+      const LGraphBadge = window['LGraphBadge']
+      const app = window['app'] as ComfyApp
+      const graph = app.graph
+      // @ts-expect-error - accessing private property
+      const nodes = graph._nodes
+
+      for (const node of nodes) {
+        node.badges = [
+          new LGraphBadge({ text: 'Test Badge 1' }),
+          new LGraphBadge({ text: 'Test Badge 2' })
+        ]
+      }
+
+      graph.setDirtyCanvas(true, true)
+    })
+
+    await expect(comfyPage.canvas).toHaveScreenshot('node-badge-multiple.png')
+  })
+
+  test('Can add badge left-side', async ({ comfyPage }) => {
+    await comfyPage.page.evaluate(() => {
+      const LGraphBadge = window['LGraphBadge']
+      const app = window['app'] as ComfyApp
+      const graph = app.graph
+      // @ts-expect-error - accessing private property
+      const nodes = graph._nodes
+
+      for (const node of nodes) {
+        node.badges = [new LGraphBadge({ text: 'Test Badge' })]
+        // @ts-expect-error - Enum value
+        node.badgePosition = 'top-left'
+      }
+
+      graph.setDirtyCanvas(true, true)
+    })
+
+    await expect(comfyPage.canvas).toHaveScreenshot('node-badge-left.png')
+  })
+})
+
+test.describe('Node source badge', () => {
+  Object.values(NodeBadgeMode).forEach(async (mode) => {
+    test(`Shows node badges (${mode})`, async ({ comfyPage }) => {
+      // Execution error workflow has both custom node and core node.
+      await comfyPage.loadWorkflow('execution_error')
+      await comfyPage.setSetting('Comfy.NodeBadge.NodeSourceBadgeMode', mode)
+      await comfyPage.setSetting('Comfy.NodeBadge.NodeIdBadgeMode', mode)
+      await comfyPage.nextFrame()
+      await comfyPage.resetView()
+      await expect(comfyPage.canvas).toHaveScreenshot(`node-badge-${mode}.png`)
+    })
+  })
+})

package-lock.json

@@ -9,7 +9,7 @@
       "version": "1.2.48",
       "dependencies": {
         "@atlaskit/pragmatic-drag-and-drop": "^1.2.1",
-        "@comfyorg/litegraph": "^0.7.65",
+        "@comfyorg/litegraph": "^0.7.67",
         "@primevue/themes": "^4.0.5",
         "@vitejs/plugin-vue": "^5.0.5",
         "@vueuse/core": "^11.0.0",
@@ -1909,9 +1909,9 @@
       "dev": true
     },
     "node_modules/@comfyorg/litegraph": {
-      "version": "0.7.65",
-      "resolved": "https://registry.npmjs.org/@comfyorg/litegraph/-/litegraph-0.7.65.tgz",
-      "integrity": "sha512-Yau14XPptHRmk/2My46mfkQBrZhZgtib75mghgwPelt2oRrk+O4OrhtbPQRzYoll7LmJAA7cHA9YI+zZbJ6IaA==",
+      "version": "0.7.67",
+      "resolved": "https://registry.npmjs.org/@comfyorg/litegraph/-/litegraph-0.7.67.tgz",
+      "integrity": "sha512-X8eRpBmSGTahJteNFDG9P0IsHXOk4QDU3p3iWPhk0rGfTnl4RZ8YcJ8MVo7zRgF3qxxX/Tcw4RpelhnjBJe4Gg==",
       "license": "MIT"
     },
     "node_modules/@cspotcode/source-map-support": {

package.json

@@ -62,7 +62,7 @@
   },
   "dependencies": {
     "@atlaskit/pragmatic-drag-and-drop": "^1.2.1",
-    "@comfyorg/litegraph": "^0.7.65",
+    "@comfyorg/litegraph": "^0.7.67",
     "@primevue/themes": "^4.0.5",
     "@vitejs/plugin-vue": "^5.0.5",
     "@vueuse/core": "^11.0.0",

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'
+}