From 645897f8b8c6f26178dcdf31a7b0c39a4a130491 Mon Sep 17 00:00:00 2001 From: Chenlei Hu Date: Sat, 26 Oct 2024 15:28:14 -0400 Subject: [PATCH] [Refactor] Make node badge a vue component (#1317) * [Refactor] Make node badge a vue component * Simplify badge logic * nit --- src/components/graph/GraphCanvas.vue | 2 + src/components/graph/NodeBadge.vue | 93 +++++++++++++++++ src/components/graph/TitleEditor.vue | 1 - src/extensions/core/index.ts | 1 - src/extensions/core/nodeBadge.ts | 150 --------------------------- src/stores/nodeDefStore.ts | 21 +++- 6 files changed, 115 insertions(+), 153 deletions(-) create mode 100644 src/components/graph/NodeBadge.vue delete mode 100644 src/extensions/core/nodeBadge.ts diff --git a/src/components/graph/GraphCanvas.vue b/src/components/graph/GraphCanvas.vue index 8d5abe5067..8438ac89e8 100644 --- a/src/components/graph/GraphCanvas.vue +++ b/src/components/graph/GraphCanvas.vue @@ -17,6 +17,7 @@ + diff --git a/src/components/graph/TitleEditor.vue b/src/components/graph/TitleEditor.vue index 5fb4aa3ad2..0188276d7c 100644 --- a/src/components/graph/TitleEditor.vue +++ b/src/components/graph/TitleEditor.vue @@ -106,7 +106,6 @@ const canvasEventHandler = (event: LiteGraphCanvasEvent) => { const [x, y] = group.pos const e = event.detail.originalEvent - // @ts-expect-error LiteGraphCanvasEvent is not typed const relativeY = e.canvasY - y // Only allow editing if the click is on the title bar if (relativeY > group.titleHeight) { diff --git a/src/extensions/core/index.ts b/src/extensions/core/index.ts index 772dd5f3db..fe0f331782 100644 --- a/src/extensions/core/index.ts +++ b/src/extensions/core/index.ts @@ -21,4 +21,3 @@ import './uploadImage' import './webcamCapture' import './widgetInputs' import './uploadAudio' -import './nodeBadge' diff --git a/src/extensions/core/nodeBadge.ts b/src/extensions/core/nodeBadge.ts deleted file mode 100644 index 64a76235ae..0000000000 --- a/src/extensions/core/nodeBadge.ts +++ /dev/null @@ -1,150 +0,0 @@ -// @ts-strict-ignore -import { app, type ComfyApp } from '@/scripts/app' -import type { ComfyExtension } from '@/types/comfy' -import type { LGraphNode } from '@comfyorg/litegraph' -import { LGraphBadge } from '@comfyorg/litegraph' -import { useSettingStore } from '@/stores/settingStore' -import { computed, ComputedRef, watch } from 'vue' -import { NodeBadgeMode, NodeSource, NodeSourceType } from '@/types/nodeSource' -import _ from 'lodash' -import { getColorPalette, defaultColorPalette } from './colorPalette' -import { BadgePosition } from '@comfyorg/litegraph' -import type { Palette } from '@/types/colorPalette' -import { useNodeDefStore } from '@/stores/nodeDefStore' - -function getNodeSource(node: LGraphNode): NodeSource | null { - const nodeDef = node.constructor.nodeData - // Frontend-only nodes don't have nodeDef - if (!nodeDef) { - return null - } - const nodeDefStore = useNodeDefStore() - return nodeDefStore.nodeDefsByName[nodeDef.name]?.nodeSource ?? null -} - -function isCoreNode(node: LGraphNode) { - return getNodeSource(node)?.type === NodeSourceType.Core -} - -function badgeTextVisible(node: LGraphNode, badgeMode: NodeBadgeMode): boolean { - return ( - badgeMode === NodeBadgeMode.None || - (isCoreNode(node) && badgeMode === NodeBadgeMode.HideBuiltIn) - ) -} - -function getNodeIdBadgeText(node: LGraphNode, nodeIdBadgeMode: NodeBadgeMode) { - return badgeTextVisible(node, nodeIdBadgeMode) ? '' : `#${node.id}` -} - -function getNodeSourceBadgeText( - node: LGraphNode, - nodeSourceBadgeMode: NodeBadgeMode -) { - const nodeSource = getNodeSource(node) - return badgeTextVisible(node, nodeSourceBadgeMode) - ? '' - : nodeSource?.badgeText ?? '' -} - -function getNodeLifeCycleBadgeText( - node: LGraphNode, - nodeLifeCycleBadgeMode: NodeBadgeMode -) { - let text = '' - const nodeDef = node.constructor.nodeData - - // Frontend-only nodes don't have nodeDef - if (!nodeDef) { - return '' - } - - if (nodeDef.deprecated) { - text = '[DEPR]' - } - - if (nodeDef.experimental) { - text = '[BETA]' - } - - return badgeTextVisible(node, nodeLifeCycleBadgeMode) ? '' : text -} - -class NodeBadgeExtension implements ComfyExtension { - name = 'Comfy.NodeBadge' - - constructor( - public nodeIdBadgeMode: ComputedRef | null = null, - public nodeSourceBadgeMode: ComputedRef | null = null, - public nodeLifeCycleBadgeMode: ComputedRef | null = null, - public colorPalette: ComputedRef | null = null - ) {} - - init(app: ComfyApp) { - const settingStore = useSettingStore() - this.nodeSourceBadgeMode = computed( - () => - settingStore.get('Comfy.NodeBadge.NodeSourceBadgeMode') as NodeBadgeMode - ) - this.nodeIdBadgeMode = computed( - () => settingStore.get('Comfy.NodeBadge.NodeIdBadgeMode') as NodeBadgeMode - ) - this.nodeLifeCycleBadgeMode = computed( - () => - settingStore.get( - 'Comfy.NodeBadge.NodeLifeCycleBadgeMode' - ) as NodeBadgeMode - ) - this.colorPalette = computed(() => - getColorPalette(settingStore.get('Comfy.ColorPalette')) - ) - - watch(this.nodeSourceBadgeMode, () => { - app.graph.setDirtyCanvas(true, true) - }) - - watch(this.nodeIdBadgeMode, () => { - app.graph.setDirtyCanvas(true, true) - }) - watch(this.nodeLifeCycleBadgeMode, () => { - app.graph.setDirtyCanvas(true, true) - }) - } - - nodeCreated(node: LGraphNode, app: ComfyApp) { - 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( - [ - getNodeIdBadgeText(node, this.nodeIdBadgeMode.value), - getNodeLifeCycleBadgeText( - node, - this.nodeLifeCycleBadgeMode.value - ), - getNodeSourceBadgeText(node, this.nodeSourceBadgeMode.value) - ] - .filter((s) => s.length > 0) - .join(' '), - { - length: 31 - } - ), - fgColor: - this.colorPalette.value.colors.litegraph_base?.BADGE_FG_COLOR || - defaultColorPalette.colors.litegraph_base.BADGE_FG_COLOR, - bgColor: - this.colorPalette.value.colors.litegraph_base?.BADGE_BG_COLOR || - defaultColorPalette.colors.litegraph_base.BADGE_BG_COLOR - }) - ) - - node.badges.push(() => badge.value) - } -} - -app.registerExtension(new NodeBadgeExtension()) diff --git a/src/stores/nodeDefStore.ts b/src/stores/nodeDefStore.ts index 4256f2a68f..406a0b7904 100644 --- a/src/stores/nodeDefStore.ts +++ b/src/stores/nodeDefStore.ts @@ -13,7 +13,12 @@ import { TreeNode } from 'primevue/treenode' import { buildTree } from '@/utils/treeUtil' import { computed, ref } from 'vue' import axios from 'axios' -import { type NodeSource, getNodeSource } from '@/types/nodeSource' +import { + type NodeSource, + NodeSourceType, + getNodeSource +} from '@/types/nodeSource' +import type { LGraphNode } from '@comfyorg/litegraph' export interface BaseInputSpec { name: string @@ -205,6 +210,16 @@ export class ComfyNodeDefImpl { const nodeFrequency = nodeFrequencyStore.getNodeFrequencyByName(this.name) return [scores[0], -nodeFrequency, ...scores.slice(1)] } + + get isCoreNode(): boolean { + return this.nodeSource.type === NodeSourceType.Core + } + + get nodeLifeCycleBadgeText(): string { + if (this.deprecated) return '[DEPR]' + if (this.experimental) return '[BETA]' + return '' + } } export const SYSTEM_NODE_DEFS: Record = { @@ -331,6 +346,10 @@ export const useNodeDefStore = defineStore('nodeDef', { }, inputIsWidget(spec: BaseInputSpec) { return this.getWidgetType(spec.type, spec.name) !== null + }, + fromLGraphNode(node: LGraphNode): ComfyNodeDefImpl | null { + // Frontend-only nodes don't have nodeDef + return this.nodeDefsByName[node.constructor.nodeData?.name] ?? null } } })