Separate nodes

browser_tests/nodeBadge.spec.ts

@@ -1,7 +1,7 @@
 import { expect } from '@playwright/test'
 import { comfyPageFixture as test } from './ComfyPage'
 import type { ComfyApp } from '../src/scripts/app'
-import { NodeSourceBadgeMode } from '../src/types/nodeSource'
+import { NodeBadgeMode } from '../src/types/nodeSource'
 
 test.describe('Node Badge', () => {
   test('Can add badge', async ({ comfyPage }) => {
@@ -65,16 +65,15 @@ test.describe('Node Badge', () => {
 })
 
 test.describe('Node source badge', () => {
-  Object.values(NodeSourceBadgeMode).forEach(async (mode) => {
-    test(`Shows node source name (${mode})`, async ({ comfyPage }) => {
+  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.Node.NodeSourceBadgeMode', mode)
+      await comfyPage.setSetting('Comfy.Node.NodeIdBadgeMode', mode)
       await comfyPage.nextFrame()
       await comfyPage.resetView()
-      await expect(comfyPage.canvas).toHaveScreenshot(
-        `node-source-badge-${mode}.png`
-      )
+      await expect(comfyPage.canvas).toHaveScreenshot(`node-badge-${mode}.png`)
     })
   })
 })

src/extensions/core/nodeSourceBadge.ts

@@ -3,70 +3,61 @@ import type { ComfyExtension } from '@/types/comfy'
 import type { ComfyLGraphNode } from '@/types/comfyLGraphNode'
 import { LGraphBadge } from '@comfyorg/litegraph'
 import { useSettingStore } from '@/stores/settingStore'
-import { computed } from 'vue'
-import { getNodeSource, NodeSourceBadgeMode } from '@/types/nodeSource'
+import { computed, watch } from 'vue'
+import {
+  getNodeSource as getNodeSourceFromPythonModule,
+  NodeBadgeMode
+} from '@/types/nodeSource'
 import _ from 'lodash'
 import { colorPalettes } from './colorPalette'
 import { BadgePosition } from '@comfyorg/litegraph'
 
 const settingStore = useSettingStore()
 const nodeSourceBadgeMode = computed(
-  () =>
-    settingStore.get('Comfy.Node.NodeSourceBadgeMode') as NodeSourceBadgeMode
+  () => settingStore.get('Comfy.Node.NodeSourceBadgeMode') as NodeBadgeMode
+)
+const nodeIdBadgeMode = computed(
+  () => settingStore.get('Comfy.Node.NodeIdBadgeMode') as NodeBadgeMode
 )
 const colorPalette = computed(
   () => colorPalettes[settingStore.get('Comfy.ColorPalette')]
 )
 const defaultColorPalette = colorPalettes['dark']
 
-function getNodeIdBadgeText(
-  node: ComfyLGraphNode,
-  nodeSourceBadgeMode: NodeSourceBadgeMode
-) {
-  if (nodeSourceBadgeMode === NodeSourceBadgeMode.None) {
-    return ''
-  }
+watch(nodeSourceBadgeMode, () => {
+  app.graph.setDirtyCanvas(true, true)
+})
 
-  const nodeId = node.id
-  if (
-    nodeSourceBadgeMode === NodeSourceBadgeMode.IdNickname ||
-    nodeSourceBadgeMode === NodeSourceBadgeMode.IdNicknameHideBuiltIn
-  ) {
-    return `#${nodeId}`
-  }
+watch(nodeIdBadgeMode, () => {
+  app.graph.setDirtyCanvas(true, true)
+})
 
-  return ''
+function getNodeSource(node: ComfyLGraphNode) {
+  const pythonModule = (node.constructor as typeof ComfyLGraphNode).nodeData
+    ?.python_module
+  return getNodeSourceFromPythonModule(pythonModule)
 }
 
-function getNodeSourceBadgeText(
-  node: ComfyLGraphNode,
-  nodeSourceBadgeMode: NodeSourceBadgeMode
-) {
-  if (nodeSourceBadgeMode === NodeSourceBadgeMode.None) {
-    return ''
-  }
+function isCoreNode(node: ComfyLGraphNode) {
+  return getNodeSource(node).type === 'core'
+}
 
-  const pythonModule = (node.constructor as typeof ComfyLGraphNode).nodeData
-    ?.python_module
-  const nodeSource = getNodeSource(pythonModule)
-  if (
-    nodeSource.type === 'core' &&
-    nodeSourceBadgeMode === NodeSourceBadgeMode.NicknameHideBuiltIn
-  ) {
-    return ''
-  }
-  return nodeSource.badgeText
+function getNodeIdBadge(node: ComfyLGraphNode, nodeIdBadgeMode: NodeBadgeMode) {
+  return nodeIdBadgeMode === NodeBadgeMode.None ||
+    (isCoreNode(node) && nodeIdBadgeMode === NodeBadgeMode.HideBuiltIn)
+    ? ''
+    : `#${node.id}`
 }
 
-function getBadgeText(
+function getNodeSourceBadge(
   node: ComfyLGraphNode,
-  nodeSourceBadgeMode: NodeSourceBadgeMode
+  nodeSourceBadgeMode: NodeBadgeMode
 ) {
-  const nodeIdBadgeText = getNodeIdBadgeText(node, nodeSourceBadgeMode)
-  const nodeSourceBadgeText = getNodeSourceBadgeText(node, nodeSourceBadgeMode)
-  return [nodeIdBadgeText, nodeSourceBadgeText]
-    .filter((text) => text.length > 0)
-    .join(' ')
+  const nodeSource = getNodeSource(node)
+  return nodeSourceBadgeMode === NodeBadgeMode.None ||
+    (isCoreNode(node) && nodeSourceBadgeMode === NodeBadgeMode.HideBuiltIn)
+    ? ''
+    : nodeSource.badgeText
 }
 
 class NodeSourceBadgeExtension implements ComfyExtension {
@@ -75,10 +66,11 @@ class NodeSourceBadgeExtension implements ComfyExtension {
     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(
+
+    const idBadge = computed(
       () =>
         new LGraphBadge({
-          text: _.truncate(getBadgeText(node, nodeSourceBadgeMode.value), {
+          text: _.truncate(getNodeIdBadge(node, nodeIdBadgeMode.value), {
             length: 25
           }),
           fgColor:
@@ -90,7 +82,26 @@ class NodeSourceBadgeExtension implements ComfyExtension {
         })
     )
 
-    node.badges.push(() => badge.value)
+    const sourceBadge = computed(
+      () =>
+        new LGraphBadge({
+          text: _.truncate(
+            getNodeSourceBadge(node, nodeSourceBadgeMode.value),
+            {
+              length: 25
+            }
+          ),
+          fgColor:
+            colorPalette.value.colors.litegraph_base?.BADGE_FG_COLOR ||
+            defaultColorPalette.colors.litegraph_base.BADGE_FG_COLOR,
+          bgColor:
+            colorPalette.value.colors.litegraph_base?.BADGE_BG_COLOR ||
+            defaultColorPalette.colors.litegraph_base.BADGE_BG_COLOR
+        })
+    )
+
+    node.badges.push(() => idBadge.value)
+    node.badges.push(() => sourceBadge.value)
   }
 }
 

src/stores/settingStore.ts

@@ -10,7 +10,7 @@
 import { app } from '@/scripts/app'
 import { ComfySettingsDialog } from '@/scripts/ui/settings'
 import { Settings } from '@/types/apiTypes'
-import { NodeSourceBadgeMode } from '@/types/nodeSource'
+import { NodeBadgeMode } from '@/types/nodeSource'
 import {
   LinkReleaseTriggerAction,
   LinkReleaseTriggerMode
@@ -336,8 +336,16 @@ export const useSettingStore = defineStore('setting', {
         id: 'Comfy.Node.NodeSourceBadgeMode',
         name: 'Node source badge mode',
         type: 'combo',
-        options: Object.values(NodeSourceBadgeMode),
-        defaultValue: NodeSourceBadgeMode.None
+        options: Object.values(NodeBadgeMode),
+        defaultValue: NodeBadgeMode.HideBuiltIn
+      })
+
+      app.ui.settings.addSetting({
+        id: 'Comfy.Node.NodeIdBadgeMode',
+        name: 'Node ID badge mode',
+        type: 'combo',
+        options: Object.values(NodeBadgeMode),
+        defaultValue: NodeBadgeMode.HideBuiltIn
       })
     },
 

src/types/apiTypes.ts

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

src/types/nodeSource.ts

@@ -27,10 +27,8 @@ export const getNodeSource = (python_module: string): NodeSource => {
   }
 }
 
-export enum NodeSourceBadgeMode {
+export enum NodeBadgeMode {
   None = 'None',
-  Nickname = 'Nickname',
-  NicknameHideBuiltIn = 'Nickname (hide built-in)',
-  IdNickname = '#ID Nickname',
-  IdNicknameHideBuiltIn = '#ID Nickname (hide built-in)'
+  ShowAll = 'Show all',
+  HideBuiltIn = 'Hide built-in'
 }