diff --git a/src/scripts/app.ts b/src/scripts/app.ts index da803eb810..b8aba7a89c 100644 --- a/src/scripts/app.ts +++ b/src/scripts/app.ts @@ -273,6 +273,12 @@ export class ComfyApp { useExtensionService().invokeExtensions('onNodeOutputsUpdated', value) } + /** + * If the user has specified a preferred format to receive preview images in, + * this function will return that format as a url query param. + * If the node's outputs are not images, this param should not be used, as it will + * force the server to load the output file as an image. + */ getPreviewFormatParam() { let preview_format = useSettingStore().get('Comfy.PreviewFormat') if (preview_format) return `&preview=${preview_format}` diff --git a/src/stores/imagePreviewStore.ts b/src/stores/imagePreviewStore.ts index ce95be3a77..83f5f8bcdd 100644 --- a/src/stores/imagePreviewStore.ts +++ b/src/stores/imagePreviewStore.ts @@ -3,6 +3,7 @@ import { defineStore } from 'pinia' import { ExecutedWsMessage, ResultItem } from '@/schemas/apiSchema' import { api } from '@/scripts/api' +import { app } from '@/scripts/app' import { parseFilePath } from '@/utils/formatUtil' import { isVideoNode } from '@/utils/litegraphUtil' @@ -17,11 +18,6 @@ const createOutputs = ( } } -const getPreviewParam = (node: LGraphNode): string => { - if (node.animatedImages || isVideoNode(node)) return '' - return app.getPreviewFormatParam() -} - export const useNodeOutputStore = defineStore('nodeOutput', () => { const getNodeId = (node: LGraphNode): string => node.id.toString() @@ -35,6 +31,41 @@ export const useNodeOutputStore = defineStore('nodeOutput', () => { return app.nodePreviewImages[getNodeId(node)] } + /** + * Check if a node's outputs includes images that should/can be loaded normally + * by PIL. + */ + const isImageOutputs = ( + node: LGraphNode, + outputs: ExecutedWsMessage['output'] + ): boolean => { + // If animated webp/png or video outputs, return false + if (node.animatedImages || isVideoNode(node)) return false + + // If no images, return false + if (!outputs?.images?.length) return false + + // If svg images, return false + if (outputs.images.some((image) => image.filename?.endsWith('svg'))) + return false + + return true + } + + /** + * Get the preview param for the node's outputs. + * + * If the output is an image, use the user's preferred format (from settings). + * For non-image outputs, return an empty string, as including the preview param + * will force the server to load the output file as an image. + */ + function getPreviewParam( + node: LGraphNode, + outputs: ExecutedWsMessage['output'] + ): string { + return isImageOutputs(node, outputs) ? app.getPreviewFormatParam() : '' + } + function getNodeImageUrls(node: LGraphNode): string[] | undefined { const previews = getNodePreviews(node) if (previews?.length) return previews @@ -43,7 +74,7 @@ export const useNodeOutputStore = defineStore('nodeOutput', () => { if (!outputs?.images?.length) return const rand = app.getRandParam() - const previewParam = getPreviewParam(node) + const previewParam = getPreviewParam(node, outputs) return outputs.images.map((image) => { const imgUrlPart = new URLSearchParams(image) @@ -78,6 +109,7 @@ export const useNodeOutputStore = defineStore('nodeOutput', () => { getNodeOutputs, getNodeImageUrls, getNodePreviews, - setNodeOutputs + setNodeOutputs, + getPreviewParam } }) diff --git a/tests-ui/tests/store/imagePreviewStore.test.ts b/tests-ui/tests/store/imagePreviewStore.test.ts new file mode 100644 index 0000000000..0a464ba815 --- /dev/null +++ b/tests-ui/tests/store/imagePreviewStore.test.ts @@ -0,0 +1,98 @@ +import { LGraphNode } from '@comfyorg/litegraph' +import { createPinia, setActivePinia } from 'pinia' +import { beforeEach, describe, expect, it, vi } from 'vitest' + +import { ExecutedWsMessage } from '@/schemas/apiSchema' +import { app } from '@/scripts/app' +import { useNodeOutputStore } from '@/stores/imagePreviewStore' +import * as litegraphUtil from '@/utils/litegraphUtil' + +vi.mock('@/utils/litegraphUtil', () => ({ + isVideoNode: vi.fn() +})) + +vi.mock('@/scripts/app', () => ({ + app: { + getPreviewFormatParam: vi.fn(() => '&format=test_webp') + } +})) + +const createMockNode = (overrides: Partial = {}): LGraphNode => + ({ + id: 1, + type: 'TestNode', + ...overrides + }) as LGraphNode + +const createMockOutputs = ( + images?: ExecutedWsMessage['output']['images'] +): ExecutedWsMessage['output'] => ({ images }) + +describe('imagePreviewStore getPreviewParam', () => { + beforeEach(() => { + setActivePinia(createPinia()) + vi.clearAllMocks() + vi.mocked(litegraphUtil.isVideoNode).mockReturnValue(false) + }) + + it('should return empty string if node.animatedImages is true', () => { + const store = useNodeOutputStore() + // @ts-expect-error `animatedImages` property is not typed + const node = createMockNode({ animatedImages: true }) + const outputs = createMockOutputs([{ filename: 'img.png' }]) + expect(store.getPreviewParam(node, outputs)).toBe('') + expect(vi.mocked(app).getPreviewFormatParam).not.toHaveBeenCalled() + }) + + it('should return empty string if isVideoNode returns true', () => { + const store = useNodeOutputStore() + vi.mocked(litegraphUtil.isVideoNode).mockReturnValue(true) + const node = createMockNode() + const outputs = createMockOutputs([{ filename: 'img.png' }]) + expect(store.getPreviewParam(node, outputs)).toBe('') + expect(vi.mocked(app).getPreviewFormatParam).not.toHaveBeenCalled() + }) + + it('should return empty string if outputs.images is undefined', () => { + const store = useNodeOutputStore() + const node = createMockNode() + const outputs: ExecutedWsMessage['output'] = {} + expect(store.getPreviewParam(node, outputs)).toBe('') + expect(vi.mocked(app).getPreviewFormatParam).not.toHaveBeenCalled() + }) + + it('should return empty string if outputs.images is empty', () => { + const store = useNodeOutputStore() + const node = createMockNode() + const outputs = createMockOutputs([]) + expect(store.getPreviewParam(node, outputs)).toBe('') + expect(vi.mocked(app).getPreviewFormatParam).not.toHaveBeenCalled() + }) + + it('should return empty string if outputs.images contains SVG images', () => { + const store = useNodeOutputStore() + const node = createMockNode() + const outputs = createMockOutputs([{ filename: 'img.svg' }]) + expect(store.getPreviewParam(node, outputs)).toBe('') + expect(vi.mocked(app).getPreviewFormatParam).not.toHaveBeenCalled() + }) + + it('should return format param for standard image outputs', () => { + const store = useNodeOutputStore() + const node = createMockNode() + const outputs = createMockOutputs([{ filename: 'img.png' }]) + expect(store.getPreviewParam(node, outputs)).toBe('&format=test_webp') + expect(vi.mocked(app).getPreviewFormatParam).toHaveBeenCalledTimes(1) + }) + + it('should return format param for multiple standard images', () => { + const store = useNodeOutputStore() + const node = createMockNode() + const outputs = createMockOutputs([ + { filename: 'img1.png' }, + { filename: 'img2.jpg' } + ]) + expect(store.getPreviewParam(node, outputs)).toBe('&format=test_webp') + expect(vi.mocked(app).getPreviewFormatParam).toHaveBeenCalledTimes(1) + }) +})