diff --git a/src/composables/usePaste.test.ts b/src/composables/usePaste.test.ts index 4e1ac3503c5..83da44ee707 100644 --- a/src/composables/usePaste.test.ts +++ b/src/composables/usePaste.test.ts @@ -7,8 +7,8 @@ import type { } from '@/lib/litegraph/src/litegraph' import { LiteGraph } from '@/lib/litegraph/src/litegraph' import { app } from '@/scripts/app' -import { isImageNode } from '@/utils/litegraphUtil' -import { pasteImageNode, usePaste } from './usePaste' +import { createNode, isImageNode } from '@/utils/litegraphUtil' +import { pasteImageNode, pasteImageNodes, usePaste } from './usePaste' function createMockNode() { return { @@ -86,6 +86,7 @@ vi.mock('@/lib/litegraph/src/litegraph', () => ({ })) vi.mock('@/utils/litegraphUtil', () => ({ + createNode: vi.fn(), isAudioNode: vi.fn(), isImageNode: vi.fn(), isVideoNode: vi.fn() @@ -103,30 +104,28 @@ describe('pasteImageNode', () => { ) }) - it('should create new LoadImage node when no image node provided', () => { + it('should create new LoadImage node when no image node provided', async () => { const mockNode = createMockNode() - vi.mocked(LiteGraph.createNode).mockReturnValue( - mockNode as unknown as LGraphNode - ) + vi.mocked(createNode).mockResolvedValue(mockNode as unknown as LGraphNode) const file = createImageFile() const dataTransfer = createDataTransfer([file]) - pasteImageNode(mockCanvas as unknown as LGraphCanvas, dataTransfer.items) + await pasteImageNode( + mockCanvas as unknown as LGraphCanvas, + dataTransfer.items + ) - expect(LiteGraph.createNode).toHaveBeenCalledWith('LoadImage') - expect(mockNode.pos).toEqual([100, 200]) - expect(mockCanvas.graph!.add).toHaveBeenCalledWith(mockNode) - expect(mockCanvas.graph!.change).toHaveBeenCalled() + expect(createNode).toHaveBeenCalledWith(mockCanvas, 'LoadImage') expect(mockNode.pasteFile).toHaveBeenCalledWith(file) }) - it('should use existing image node when provided', () => { + it('should use existing image node when provided', async () => { const mockNode = createMockNode() const file = createImageFile() const dataTransfer = createDataTransfer([file]) - pasteImageNode( + await pasteImageNode( mockCanvas as unknown as LGraphCanvas, dataTransfer.items, mockNode as unknown as LGraphNode @@ -136,13 +135,13 @@ describe('pasteImageNode', () => { expect(mockNode.pasteFiles).toHaveBeenCalledWith([file]) }) - it('should handle multiple image files', () => { + it('should handle multiple image files', async () => { const mockNode = createMockNode() const file1 = createImageFile('test1.png') const file2 = createImageFile('test2.jpg', 'image/jpeg') const dataTransfer = createDataTransfer([file1, file2]) - pasteImageNode( + await pasteImageNode( mockCanvas as unknown as LGraphCanvas, dataTransfer.items, mockNode as unknown as LGraphNode @@ -152,11 +151,11 @@ describe('pasteImageNode', () => { expect(mockNode.pasteFiles).toHaveBeenCalledWith([file1, file2]) }) - it('should do nothing when no image files present', () => { + it('should do nothing when no image files present', async () => { const mockNode = createMockNode() const dataTransfer = createDataTransfer() - pasteImageNode( + await pasteImageNode( mockCanvas as unknown as LGraphCanvas, dataTransfer.items, mockNode as unknown as LGraphNode @@ -166,13 +165,13 @@ describe('pasteImageNode', () => { expect(mockNode.pasteFiles).not.toHaveBeenCalled() }) - it('should filter non-image items', () => { + it('should filter non-image items', async () => { const mockNode = createMockNode() const imageFile = createImageFile() const textFile = new File([''], 'test.txt', { type: 'text/plain' }) const dataTransfer = createDataTransfer([textFile, imageFile]) - pasteImageNode( + await pasteImageNode( mockCanvas as unknown as LGraphCanvas, dataTransfer.items, mockNode as unknown as LGraphNode @@ -183,6 +182,48 @@ describe('pasteImageNode', () => { }) }) +describe('pasteImageNodes', () => { + beforeEach(() => { + vi.clearAllMocks() + }) + + it('should create multiple nodes for multiple files', async () => { + const mockNode1 = createMockNode() + const mockNode2 = createMockNode() + vi.mocked(createNode) + .mockResolvedValueOnce(mockNode1 as unknown as LGraphNode) + .mockResolvedValueOnce(mockNode2 as unknown as LGraphNode) + + const file1 = createImageFile('test1.png') + const file2 = createImageFile('test2.jpg', 'image/jpeg') + const fileList = createDataTransfer([file1, file2]).files + + const result = await pasteImageNodes( + mockCanvas as unknown as LGraphCanvas, + fileList + ) + + expect(createNode).toHaveBeenCalledTimes(2) + expect(createNode).toHaveBeenNthCalledWith(1, mockCanvas, 'LoadImage') + expect(createNode).toHaveBeenNthCalledWith(2, mockCanvas, 'LoadImage') + expect(mockNode1.pasteFile).toHaveBeenCalledWith(file1) + expect(mockNode2.pasteFile).toHaveBeenCalledWith(file2) + expect(result).toEqual([mockNode1, mockNode2]) + }) + + it('should handle empty file list', async () => { + const fileList = createDataTransfer([]).files + + const result = await pasteImageNodes( + mockCanvas as unknown as LGraphCanvas, + fileList + ) + + expect(createNode).not.toHaveBeenCalled() + expect(result).toEqual([]) + }) +}) + describe('usePaste', () => { beforeEach(() => { vi.clearAllMocks() @@ -195,9 +236,7 @@ describe('usePaste', () => { it('should handle image paste', async () => { const mockNode = createMockNode() - vi.mocked(LiteGraph.createNode).mockReturnValue( - mockNode as unknown as LGraphNode - ) + vi.mocked(createNode).mockResolvedValue(mockNode as unknown as LGraphNode) usePaste() @@ -207,7 +246,7 @@ describe('usePaste', () => { document.dispatchEvent(event) await vi.waitFor(() => { - expect(LiteGraph.createNode).toHaveBeenCalledWith('LoadImage') + expect(createNode).toHaveBeenCalledWith(mockCanvas, 'LoadImage') expect(mockNode.pasteFile).toHaveBeenCalledWith(file) }) }) diff --git a/src/composables/usePaste.ts b/src/composables/usePaste.ts index 1809eb838d7..d37e4d812fb 100644 --- a/src/composables/usePaste.ts +++ b/src/composables/usePaste.ts @@ -6,9 +6,42 @@ import type { ComfyWorkflowJSON } from '@/platform/workflow/validation/schemas/w import { useCanvasStore } from '@/renderer/core/canvas/canvasStore' import { app } from '@/scripts/app' import { useWorkspaceStore } from '@/stores/workspaceStore' -import { isAudioNode, isImageNode, isVideoNode } from '@/utils/litegraphUtil' +import { createNode, isAudioNode, isImageNode, isVideoNode } from '@/utils/litegraphUtil' import { shouldIgnoreCopyPaste } from '@/workbench/eventHelpers' +export function cloneDataTransfer(original: DataTransfer): DataTransfer { + const persistent = new DataTransfer(); + + // Copy string data + for (const type of original.types) { + const data = original.getData(type); + if (data) { + persistent.setData(type, data); + } + } + + // Copy files + for (const file of original.files) { + persistent.items.add(file); + } + + // Also handle any file-kind items that might not be in .files + for (const item of original.items) { + if (item.kind === 'file') { + const file = item.getAsFile(); + if (file) { + persistent.items.add(file); + } + } + } + + // Preserve dropEffect and effectAllowed + persistent.dropEffect = original.dropEffect; + persistent.effectAllowed = original.effectAllowed; + + return persistent; +} + function pasteClipboardItems(data: DataTransfer): boolean { const rawData = data.getData('text/html') const match = rawData.match(/data-metadata="([A-Za-z0-9+/=]+)"/)?.[1] @@ -48,27 +81,37 @@ function pasteItemsOnNode( ) } -export function pasteImageNode( +export async function pasteImageNode( canvas: LGraphCanvas, items: DataTransferItemList, imageNode: LGraphNode | null = null -): void { - const { - graph, - graph_mouse: [posX, posY] - } = canvas - +): Promise { if (!imageNode) { // No image node selected: add a new one - const newNode = LiteGraph.createNode('LoadImage') - if (newNode) { - newNode.pos = [posX, posY] - imageNode = graph?.add(newNode) ?? null - } - graph?.change() + imageNode = await createNode(canvas, 'LoadImage') } pasteItemsOnNode(items, imageNode, 'image') + return imageNode +} + +export async function pasteImageNodes( + canvas: LGraphCanvas, + fileList: FileList +): Promise { + const nodes: LGraphNode[] = []; + + for (const file of fileList) { + const transfer = new DataTransfer() + transfer.items.add(file) + const imageNode = await pasteImageNode(canvas, transfer.items) + + if (imageNode) { + nodes.push(imageNode) + } + } + + return nodes } /** @@ -93,6 +136,7 @@ export const usePaste = () => { const { graph } = canvas let data: DataTransfer | string | null = e.clipboardData if (!data) throw new Error('No clipboard data on clipboard event') + data = cloneDataTransfer(data) const { items } = data @@ -114,7 +158,7 @@ export const usePaste = () => { // Look for image paste data for (const item of items) { if (item.type.startsWith('image/')) { - pasteImageNode(canvas as LGraphCanvas, items, imageNode) + await pasteImageNode(canvas as LGraphCanvas, items, imageNode) return } else if (item.type.startsWith('video/')) { if (!videoNode) { diff --git a/src/scripts/app.ts b/src/scripts/app.ts index a9427f85bf5..0c2e33d852b 100644 --- a/src/scripts/app.ts +++ b/src/scripts/app.ts @@ -77,6 +77,7 @@ import { } from '@/utils/graphTraversalUtil' import { executeWidgetsCallback, + createNode, fixLinkInputSlots, isImageNode } from '@/utils/litegraphUtil' @@ -97,7 +98,7 @@ import { type ComfyWidgetConstructor } from './widgets' import { ensureCorrectLayoutScale } from '@/renderer/extensions/vueNodes/layout/ensureCorrectLayoutScale' import { extractFileFromDragEvent } from '@/utils/eventUtils' import { getWorkflowDataFromFile } from '@/scripts/metadata/parser' -import { pasteImageNode } from '@/composables/usePaste' +import { pasteImageNode, pasteImageNodes } from '@/composables/usePaste' export const ANIM_PREVIEW_WIDGET = '$$comfy_animation_preview' @@ -542,7 +543,13 @@ export class ComfyApp { const workspace = useWorkspaceStore() try { workspace.spinner = true - await this.handleFile(fileMaybe, 'file_drop') + if (fileMaybe instanceof File) { + await this.handleFile(fileMaybe, 'file_drop') + } + + if (fileMaybe instanceof FileList) { + await this.handleFileList(fileMaybe) + } } finally { workspace.spinner = false } @@ -1437,7 +1444,9 @@ export class ComfyApp { if (file.type.startsWith('image')) { const transfer = new DataTransfer() transfer.items.add(file) - pasteImageNode(this.canvas, transfer.items) + console.log('transfer items', transfer.items) + const imageNode = await createNode(this.canvas, 'LoadImage') + await pasteImageNode(this.canvas, transfer.items, imageNode) return } @@ -1516,6 +1525,42 @@ export class ComfyApp { this.showErrorOnFileLoad(file) } + + /** + * Loads multiple files and connects to a batch node + * @param {FileList} fileList + */ + async handleFileList(fileList: FileList) { + if (fileList[0].type.startsWith('image')) { + const imageNodes = await pasteImageNodes(this.canvas, fileList) + const batchImagesNode = await createNode(this.canvas, 'BatchImagesNode') + if (!batchImagesNode) return + + this.positionNodes(imageNodes, batchImagesNode) + + Array.from(imageNodes).forEach((imageNode, index) => { + imageNode.connect(0, batchImagesNode, index) + }) + } + } + + /** + * Positions batched nodes in drag and drop + * @param nodes + * @param batchNode + */ + positionNodes(nodes: LGraphNode[], batchNode: LGraphNode) { + const [x, y, width, height] = nodes[0].getBounding() + batchNode.pos = [ x + width + 100, y + 30 ] + + nodes.forEach((node, index) => { + if (index > 0) { + node.pos = [ x, y + (height * index) + (25 * (index + 1))]; + } + this.canvas.graph?.change() + }); + } + // @deprecated isApiJson(data: unknown): data is ComfyApiWorkflow { if (!_.isObject(data) || Array.isArray(data)) { diff --git a/src/utils/__tests__/eventUtils.test.ts b/src/utils/__tests__/eventUtils.test.ts index 2fc51ac677a..ca20da96c91 100644 --- a/src/utils/__tests__/eventUtils.test.ts +++ b/src/utils/__tests__/eventUtils.test.ts @@ -33,6 +33,45 @@ describe('eventUtils', () => { expect(actual).toBe(fileWithWorkflowMaybeWhoKnows) }) + it('should handle drops with multiple image files', async () => { + const imageFile1 = new File([new Uint8Array()], 'image1.png', { + type: 'image/png' + }) + const imageFile2 = new File([new Uint8Array()], 'image2.jpg', { + type: 'image/jpeg' + }) + + const dataTransfer = new DataTransfer() + dataTransfer.items.add(imageFile1) + dataTransfer.items.add(imageFile2) + + const event = new FakeDragEvent('drop', { dataTransfer }) + + const actual = await extractFileFromDragEvent(event) + expect(actual).toBeDefined() + expect((actual as FileList).length).toBe(2) + expect((actual as FileList)[0]).toBe(imageFile1) + expect((actual as FileList)[1]).toBe(imageFile2) + }) + + it('should return undefined when dropping multiple non-image files', async () => { + const file1 = new File([new Uint8Array()], 'file1.txt', { + type: 'text/plain' + }) + const file2 = new File([new Uint8Array()], 'file2.txt', { + type: 'text/plain' + }) + + const dataTransfer = new DataTransfer() + dataTransfer.items.add(file1) + dataTransfer.items.add(file2) + + const event = new FakeDragEvent('drop', { dataTransfer }) + + const actual = await extractFileFromDragEvent(event) + expect(actual).toBe(undefined) + }) + // Skip until we can setup MSW it.skip('should handle drops with URLs', async () => { const urlWithWorkflow = 'https://fakewebsite.notreal/fake_workflow.json' diff --git a/src/utils/eventUtils.ts b/src/utils/eventUtils.ts index 133ccd709bd..63a0af8963a 100644 --- a/src/utils/eventUtils.ts +++ b/src/utils/eventUtils.ts @@ -1,14 +1,14 @@ export async function extractFileFromDragEvent( event: DragEvent -): Promise { +): Promise { if (!event.dataTransfer) return - // Dragging from Chrome->Firefox there is a file but its a bmp, so ignore that - if ( - event.dataTransfer.files.length && - event.dataTransfer.files[0].type !== 'image/bmp' - ) { - return event.dataTransfer.files[0] + const { files } = event.dataTransfer + // Dragging from Chrome->Firefox there is a file, but it's a bmp, so ignore it + if (files.length === 1 && files[0].type !== 'image/bmp') { + return files[0] + } else if (files.length > 1 && Array.from(files).every(hasImageType)) { + return files } // Try loading the first URI in the transfer list @@ -25,3 +25,5 @@ export async function extractFileFromDragEvent( const blob = await response.blob() return new File([blob], uri, { type: blob.type }) } + +const hasImageType = ({ type }: File): Boolean => type.startsWith('image'); diff --git a/src/utils/litegraphUtil.test.ts b/src/utils/litegraphUtil.test.ts index 36a6692ee92..f491394bed9 100644 --- a/src/utils/litegraphUtil.test.ts +++ b/src/utils/litegraphUtil.test.ts @@ -1,13 +1,151 @@ -import { describe, expect, it } from 'vitest' +import { beforeEach, describe, expect, it, vi } from 'vitest' +import type { LGraphCanvas, LGraphNode } from '@/lib/litegraph/src/litegraph' +import { LiteGraph } from '@/lib/litegraph/src/litegraph' import type { ISerialisedGraph } from '@/lib/litegraph/src/types/serialisation' import type { IWidget } from '@/lib/litegraph/src/types/widgets' import type { InputSpec } from '@/schemas/nodeDef/nodeDefSchemaV2' import { compressWidgetInputSlots, + createNode, migrateWidgetsValues } from '@/utils/litegraphUtil' +vi.mock('@/lib/litegraph/src/litegraph', () => ({ + LiteGraph: { + createNode: vi.fn() + } +})) + +vi.mock('@/platform/updates/common/toastStore', () => ({ + useToastStore: vi.fn(() => ({ + addAlert: vi.fn() + })) +})) + +vi.mock('@/i18n', () => ({ + t: vi.fn((key) => key) +})) + +describe('createNode', () => { + let mockCanvas: Partial + let mockGraph: any + + beforeEach(() => { + vi.clearAllMocks() + mockGraph = { + add: vi.fn((node) => node), + change: vi.fn() + } + mockCanvas = { + graph: mockGraph, + graph_mouse: [100, 200] + } + }) + + it('should create a node successfully', async () => { + const mockNode = { + pos: [0, 0] + } as unknown as LGraphNode + + vi.mocked(LiteGraph.createNode).mockImplementation( + (_name, _title, options: any) => { + setTimeout(() => options?.onNodeCreated?.(), 0) + return mockNode + } + ) + + const result = await createNode(mockCanvas as LGraphCanvas, 'LoadImage') + + expect(LiteGraph.createNode).toHaveBeenCalledWith( + 'LoadImage', + 'LoadImage', + { + onNodeCreated: expect.any(Function) + } + ) + expect(mockNode.pos).toEqual([100, 200]) + expect(mockGraph.add).toHaveBeenCalledWith(mockNode) + expect(mockGraph.change).toHaveBeenCalled() + expect(result).toBe(mockNode) + }) + + it('should return null when name is empty', async () => { + const result = await createNode(mockCanvas as LGraphCanvas, '') + + expect(LiteGraph.createNode).not.toHaveBeenCalled() + expect(result).toBeNull() + }) + + it('should return null when name is falsy', async () => { + const result = await createNode(mockCanvas as LGraphCanvas, null as any) + + expect(LiteGraph.createNode).not.toHaveBeenCalled() + expect(result).toBeNull() + }) + + it('should handle node creation failure and show toast', async () => { + const { useToastStore } = + await import('@/platform/updates/common/toastStore') + const mockAddAlert = vi.fn() + vi.mocked(useToastStore).mockReturnValue({ addAlert: mockAddAlert } as any) + + vi.mocked(LiteGraph.createNode).mockImplementation( + (_name, _title, options: any) => { + setTimeout(() => options?.onNodeCreated?.(), 0) + return null + } + ) + + const result = await createNode(mockCanvas as LGraphCanvas, 'InvalidNode') + + expect(mockAddAlert).toHaveBeenCalledWith('assetBrowser.failedToCreateNode') + expect(result).toBeNull() + }) + + it('should handle graph being null', async () => { + const mockNode = { + pos: [0, 0] + } as unknown as LGraphNode + + mockCanvas.graph = null + + vi.mocked(LiteGraph.createNode).mockImplementation( + (_name, _title, options: any) => { + setTimeout(() => options?.onNodeCreated?.(), 0) + return mockNode + } + ) + + const result = await createNode(mockCanvas as LGraphCanvas, 'LoadImage') + + expect(mockNode.pos).toEqual([100, 200]) + expect(result).toBeNull() + }) + + it('should set position based on canvas graph_mouse', async () => { + const mockCanvasWithDifferentPos = { + ...mockCanvas, + graph_mouse: [250, 350] + } + + const mockNode = { + pos: [0, 0] + } as unknown as LGraphNode + + vi.mocked(LiteGraph.createNode).mockImplementation( + (_name, _title, options: any) => { + setTimeout(() => options?.onNodeCreated?.(), 0) + return mockNode + } + ) + + await createNode(mockCanvasWithDifferentPos as LGraphCanvas, 'LoadAudio') + + expect(mockNode.pos).toEqual([250, 350]) + }) +}) + describe('migrateWidgetsValues', () => { it('should remove widget values for forceInput inputs', () => { const inputDefs: Record = { diff --git a/src/utils/litegraphUtil.ts b/src/utils/litegraphUtil.ts index 898b3ef237f..9797439717f 100644 --- a/src/utils/litegraphUtil.ts +++ b/src/utils/litegraphUtil.ts @@ -1,6 +1,11 @@ import _ from 'es-toolkit/compat' -import type { ColorOption, LGraph } from '@/lib/litegraph/src/litegraph' +import type { + ColorOption, + LGraph, + LGraphCanvas, +} from '@/lib/litegraph/src/litegraph' +import { LiteGraph } from '@/lib/litegraph/src/litegraph' import { LGraphGroup, LGraphNode, @@ -17,6 +22,8 @@ import type { IComboWidget } from '@/lib/litegraph/src/types/widgets' import type { InputSpec } from '@/schemas/nodeDef/nodeDefSchemaV2' +import { useToastStore } from '@/platform/updates/common/toastStore' +import { t } from '@/i18n' type ImageNode = LGraphNode & { imgs: HTMLImageElement[] | undefined } type VideoNode = LGraphNode & { @@ -24,6 +31,36 @@ type VideoNode = LGraphNode & { imgs: HTMLVideoElement[] | undefined } +/** + * Promisify Litegraph.createNode with ComfyUI + * @param canvas + * @param name + */ +export async function createNode( + canvas: LGraphCanvas, + name: string +): Promise { + if (!name) { + return null + } + + const { graph, graph_mouse: [ posX, posY ] } = canvas + const newNode = await new Promise((resolve) => { + const createdNode = LiteGraph.createNode(name) + setTimeout(resolve, 0, createdNode) + }) + + if (newNode) { + newNode!.pos = [ posX, posY ] + const addedNode = graph?.add(newNode!) ?? null + graph?.change() + return addedNode + } else { + useToastStore().addAlert(t('assetBrowser.failedToCreateNode')) + return null + } +} + export function isImageNode(node: LGraphNode | undefined): node is ImageNode { if (!node) return false return (