diff --git a/src/platform/settings/constants/coreSettings.ts b/src/platform/settings/constants/coreSettings.ts index 95e892c49..0d9cc4e41 100644 --- a/src/platform/settings/constants/coreSettings.ts +++ b/src/platform/settings/constants/coreSettings.ts @@ -1190,5 +1190,16 @@ export const CORE_SETTINGS: SettingParams[] = [ type: 'boolean', defaultValue: false, versionAdded: '1.39.0' + }, + { + id: 'Comfy.NodeReplacement.Enabled', + category: ['Comfy', 'Workflow', 'NodeReplacement'], + name: 'Enable automatic node replacement', + tooltip: + 'When enabled, missing nodes can be automatically replaced with their newer equivalents if a replacement mapping exists.', + type: 'boolean', + defaultValue: true, + experimental: true, + versionAdded: '1.40.0' } ] diff --git a/src/schemas/apiSchema.ts b/src/schemas/apiSchema.ts index e9ac9e454..4524fed77 100644 --- a/src/schemas/apiSchema.ts +++ b/src/schemas/apiSchema.ts @@ -315,6 +315,7 @@ const zSettings = z.object({ 'Comfy.Node.MiddleClickRerouteNode': z.boolean(), 'Comfy.Node.ShowDeprecated': z.boolean(), 'Comfy.Node.ShowExperimental': z.boolean(), + 'Comfy.NodeReplacement.Enabled': z.boolean(), 'Comfy.Pointer.ClickBufferTime': z.number(), 'Comfy.Pointer.ClickDrift': z.number(), 'Comfy.Pointer.DoubleClickTime': z.number(), diff --git a/src/scripts/app.ts b/src/scripts/app.ts index 796965b27..78d1fd57f 100644 --- a/src/scripts/app.ts +++ b/src/scripts/app.ts @@ -62,6 +62,7 @@ import { useNodeOutputStore } from '@/stores/imagePreviewStore' import { KeyComboImpl, useKeybindingStore } from '@/stores/keybindingStore' import { useModelStore } from '@/stores/modelStore' import { SYSTEM_NODE_DEFS, useNodeDefStore } from '@/stores/nodeDefStore' +import { useNodeReplacementStore } from '@/stores/nodeReplacementStore' import { useSubgraphStore } from '@/stores/subgraphStore' import { useWidgetStore } from '@/stores/widgetStore' import { useWorkspaceStore } from '@/stores/workspaceStore' @@ -759,6 +760,7 @@ export class ComfyApp { await useWorkspaceStore().workflow.syncWorkflows() //Doesn't need to block. Blueprints will load async void useSubgraphStore().fetchSubgraphs() + void useNodeReplacementStore().load() await useExtensionService().loadExtensions() this.addProcessKeyHandler() diff --git a/src/services/nodeReplacementService.ts b/src/services/nodeReplacementService.ts new file mode 100644 index 000000000..9bde47099 --- /dev/null +++ b/src/services/nodeReplacementService.ts @@ -0,0 +1,13 @@ +import type { NodeReplacementResponse } from '@/types/nodeReplacementTypes' + +import { api } from '@/scripts/api' + +export async function fetchNodeReplacements(): Promise { + const response = await api.fetchApi('/node_replacements') + if (!response.ok) { + throw new Error( + `Failed to fetch node replacements: ${response.status} ${response.statusText}` + ) + } + return response.json() +} diff --git a/src/stores/nodeReplacementStore.test.ts b/src/stores/nodeReplacementStore.test.ts new file mode 100644 index 000000000..c20c43fe7 --- /dev/null +++ b/src/stores/nodeReplacementStore.test.ts @@ -0,0 +1,171 @@ +import { createPinia, setActivePinia } from 'pinia' +import { beforeEach, describe, expect, it, vi } from 'vitest' + +import { useSettingStore } from '@/platform/settings/settingStore' +import { useNodeReplacementStore } from '@/stores/nodeReplacementStore' + +vi.mock('@/platform/settings/settingStore', () => ({ + useSettingStore: vi.fn() +})) + +function mockSettingStore(enabled: boolean) { + vi.mocked(useSettingStore, { partial: true }).mockReturnValue({ + get: vi.fn().mockImplementation((key: string) => { + if (key === 'Comfy.NodeReplacement.Enabled') { + return enabled + } + return false + }) + }) +} + +describe('useNodeReplacementStore', () => { + let store: ReturnType + + beforeEach(() => { + setActivePinia(createPinia()) + mockSettingStore(true) + store = useNodeReplacementStore() + }) + + it('should initialize with empty replacements', () => { + expect(store.replacements).toEqual({}) + expect(store.isLoaded).toBe(false) + }) + + describe('getReplacementFor', () => { + it('should return first replacement for existing node type', () => { + store.replacements = { + OldNode: [ + { + new_node_id: 'NewNodeA', + old_node_id: 'OldNode', + old_widget_ids: null, + input_mapping: null, + output_mapping: null + }, + { + new_node_id: 'NewNodeB', + old_node_id: 'OldNode', + old_widget_ids: null, + input_mapping: null, + output_mapping: null + } + ] + } + + const result = store.getReplacementFor('OldNode') + + expect(result).not.toBeNull() + expect(result?.new_node_id).toBe('NewNodeA') + }) + + it('should return null for non-existing node type', () => { + store.replacements = { + OldNode: [ + { + new_node_id: 'NewNode', + old_node_id: 'OldNode', + old_widget_ids: null, + input_mapping: null, + output_mapping: null + } + ] + } + + const result = store.getReplacementFor('NonExistentNode') + + expect(result).toBeNull() + }) + + it('should return null for empty replacement array', () => { + store.replacements = { + OldNode: [] + } + + const result = store.getReplacementFor('OldNode') + + expect(result).toBeNull() + }) + + it('should return null when feature is disabled', () => { + mockSettingStore(false) + store.replacements = { + OldNode: [ + { + new_node_id: 'NewNode', + old_node_id: 'OldNode', + old_widget_ids: null, + input_mapping: null, + output_mapping: null + } + ] + } + + const result = store.getReplacementFor('OldNode') + + expect(result).toBeNull() + }) + }) + + describe('hasReplacement', () => { + it('should return true when replacement exists', () => { + store.replacements = { + OldNode: [ + { + new_node_id: 'NewNode', + old_node_id: 'OldNode', + old_widget_ids: null, + input_mapping: null, + output_mapping: null + } + ] + } + + expect(store.hasReplacement('OldNode')).toBe(true) + }) + + it('should return false when node type does not exist', () => { + store.replacements = {} + + expect(store.hasReplacement('NonExistentNode')).toBe(false) + }) + + it('should return false when replacement array is empty', () => { + store.replacements = { + OldNode: [] + } + + expect(store.hasReplacement('OldNode')).toBe(false) + }) + + it('should return false when feature is disabled', () => { + mockSettingStore(false) + store.replacements = { + OldNode: [ + { + new_node_id: 'NewNode', + old_node_id: 'OldNode', + old_widget_ids: null, + input_mapping: null, + output_mapping: null + } + ] + } + + expect(store.hasReplacement('OldNode')).toBe(false) + }) + }) + + describe('isEnabled', () => { + it('should return true when setting is enabled', () => { + mockSettingStore(true) + expect(store.isEnabled()).toBe(true) + }) + + it('should return false when setting is disabled', () => { + mockSettingStore(false) + expect(store.isEnabled()).toBe(false) + }) + }) +}) diff --git a/src/stores/nodeReplacementStore.ts b/src/stores/nodeReplacementStore.ts new file mode 100644 index 000000000..60bae2035 --- /dev/null +++ b/src/stores/nodeReplacementStore.ts @@ -0,0 +1,49 @@ +import type { + NodeReplacement, + NodeReplacementResponse +} from '@/types/nodeReplacementTypes' + +import { defineStore } from 'pinia' +import { ref } from 'vue' + +import { useSettingStore } from '@/platform/settings/settingStore' +import { fetchNodeReplacements } from '@/services/nodeReplacementService' + +export const useNodeReplacementStore = defineStore('nodeReplacement', () => { + const replacements = ref({}) + const isLoaded = ref(false) + + async function load() { + if (isLoaded.value) return + + try { + replacements.value = await fetchNodeReplacements() + isLoaded.value = true + } catch (error) { + console.error('Failed to load node replacements:', error) + } + } + + function isEnabled(): boolean { + return useSettingStore().get('Comfy.NodeReplacement.Enabled') + } + + function getReplacementFor(nodeType: string): NodeReplacement | null { + if (!isEnabled()) return null + return replacements.value[nodeType]?.[0] ?? null + } + + function hasReplacement(nodeType: string): boolean { + if (!isEnabled()) return false + return !!replacements.value[nodeType]?.length + } + + return { + replacements, + isLoaded, + load, + isEnabled, + getReplacementFor, + hasReplacement + } +}) diff --git a/src/types/nodeReplacementTypes.ts b/src/types/nodeReplacementTypes.ts new file mode 100644 index 000000000..efbfbba70 --- /dev/null +++ b/src/types/nodeReplacementTypes.ts @@ -0,0 +1,29 @@ +interface InputAssignOldId { + assign_type: 'old_id' + old_id: string +} + +interface InputAssignSetValue { + assign_type: 'set_value' + value: unknown +} + +interface InputMap { + new_id: string + assign: InputAssignOldId | InputAssignSetValue +} + +interface OutputMap { + new_idx: number + old_idx: number +} + +export interface NodeReplacement { + new_node_id: string + old_node_id: string + old_widget_ids: string[] | null + input_mapping: InputMap[] | null + output_mapping: OutputMap[] | null +} + +export type NodeReplacementResponse = Record