diff --git a/src/composables/node/useNodeImageUpload.test.ts b/src/composables/node/useNodeImageUpload.test.ts new file mode 100644 index 0000000000..ca755faa4a --- /dev/null +++ b/src/composables/node/useNodeImageUpload.test.ts @@ -0,0 +1,175 @@ +import { beforeEach, describe, expect, it, vi } from 'vitest' + +import type { LGraphNode } from '@/lib/litegraph/src/litegraph' + +const { mockFetchApi, mockAddAlert, mockUpdateInputs } = vi.hoisted(() => ({ + mockFetchApi: vi.fn(), + mockAddAlert: vi.fn(), + mockUpdateInputs: vi.fn() +})) + +let capturedDragOnDrop: (files: File[]) => Promise + +vi.mock('@/composables/node/useNodeDragAndDrop', () => ({ + useNodeDragAndDrop: ( + _node: LGraphNode, + opts: { onDrop: typeof capturedDragOnDrop } + ) => { + capturedDragOnDrop = opts.onDrop + } +})) + +vi.mock('@/composables/node/useNodeFileInput', () => ({ + useNodeFileInput: () => ({ openFileSelection: vi.fn() }) +})) + +vi.mock('@/composables/node/useNodePaste', () => ({ + useNodePaste: vi.fn() +})) + +vi.mock('@/i18n', () => ({ + t: (key: string) => key +})) + +vi.mock('@/platform/updates/common/toastStore', () => ({ + useToastStore: () => ({ addAlert: mockAddAlert }) +})) + +vi.mock('@/scripts/api', () => ({ + api: { fetchApi: mockFetchApi } +})) + +vi.mock('@/stores/assetsStore', () => ({ + useAssetsStore: () => ({ updateInputs: mockUpdateInputs }) +})) + +function createMockNode(): LGraphNode { + return { + isUploading: false, + imgs: [new Image()], + graph: { setDirtyCanvas: vi.fn() }, + size: [300, 400] + } as unknown as LGraphNode +} + +function createFile(name = 'test.png'): File { + return new File(['data'], name, { type: 'image/png' }) +} + +function successResponse(name: string, subfolder?: string) { + return { + status: 200, + json: () => Promise.resolve({ name, subfolder }) + } +} + +function failResponse(status = 500) { + return { + status, + statusText: 'Server Error' + } +} + +describe('useNodeImageUpload', () => { + let node: LGraphNode + let onUploadComplete: (paths: string[]) => void + let onUploadStart: (files: File[]) => void + let onUploadError: () => void + + beforeEach(async () => { + vi.resetModules() + vi.clearAllMocks() + node = createMockNode() + onUploadComplete = vi.fn() + onUploadStart = vi.fn() + onUploadError = vi.fn() + + const { useNodeImageUpload } = await import('./useNodeImageUpload') + useNodeImageUpload(node, { + onUploadComplete, + onUploadStart, + onUploadError, + folder: 'input' + }) + }) + + it('sets isUploading true during upload and false after', async () => { + mockFetchApi.mockResolvedValueOnce(successResponse('test.png')) + + const promise = capturedDragOnDrop([createFile()]) + expect(node.isUploading).toBe(true) + + await promise + expect(node.isUploading).toBe(false) + }) + + it('clears node.imgs on upload start', async () => { + mockFetchApi.mockResolvedValueOnce(successResponse('test.png')) + + const promise = capturedDragOnDrop([createFile()]) + expect(node.imgs).toBeUndefined() + + await promise + }) + + it('calls onUploadStart with files', async () => { + mockFetchApi.mockResolvedValueOnce(successResponse('test.png')) + const files = [createFile()] + + await capturedDragOnDrop(files) + expect(onUploadStart).toHaveBeenCalledWith(files) + }) + + it('calls onUploadComplete with valid paths on success', async () => { + mockFetchApi.mockResolvedValueOnce(successResponse('test.png')) + + await capturedDragOnDrop([createFile()]) + expect(onUploadComplete).toHaveBeenCalledWith(['test.png']) + }) + + it('includes subfolder in returned path', async () => { + mockFetchApi.mockResolvedValueOnce(successResponse('test.png', 'pasted')) + + await capturedDragOnDrop([createFile()]) + expect(onUploadComplete).toHaveBeenCalledWith(['pasted/test.png']) + }) + + it('calls onUploadError when all uploads fail', async () => { + mockFetchApi.mockResolvedValueOnce(failResponse()) + + await capturedDragOnDrop([createFile()]) + expect(onUploadError).toHaveBeenCalled() + expect(onUploadComplete).not.toHaveBeenCalled() + }) + + it('resets isUploading even when upload fails', async () => { + mockFetchApi.mockRejectedValueOnce(new Error('Network error')) + + await capturedDragOnDrop([createFile()]) + expect(node.isUploading).toBe(false) + }) + + it('rejects concurrent uploads with a toast', async () => { + mockFetchApi.mockImplementation( + () => + new Promise((resolve) => + setTimeout(() => resolve(successResponse('a.png')), 50) + ) + ) + + const first = capturedDragOnDrop([createFile('a.png')]) + const second = await capturedDragOnDrop([createFile('b.png')]) + + expect(second).toEqual([]) + expect(mockAddAlert).toHaveBeenCalledWith('g.uploadAlreadyInProgress') + + await first + }) + + it('calls setDirtyCanvas on start and finish', async () => { + mockFetchApi.mockResolvedValueOnce(successResponse('test.png')) + + await capturedDragOnDrop([createFile()]) + expect(node.graph?.setDirtyCanvas).toHaveBeenCalledTimes(2) + }) +}) diff --git a/src/composables/node/useNodeImageUpload.ts b/src/composables/node/useNodeImageUpload.ts index 96d450c1d9..587d67bcdc 100644 --- a/src/composables/node/useNodeImageUpload.ts +++ b/src/composables/node/useNodeImageUpload.ts @@ -1,6 +1,7 @@ import { useNodeDragAndDrop } from '@/composables/node/useNodeDragAndDrop' import { useNodeFileInput } from '@/composables/node/useNodeFileInput' import { useNodePaste } from '@/composables/node/useNodePaste' +import { t } from '@/i18n' import type { LGraphNode } from '@/lib/litegraph/src/litegraph' import { useToastStore } from '@/platform/updates/common/toastStore' import type { ResultItemType } from '@/schemas/apiSchema' @@ -62,6 +63,8 @@ interface ImageUploadOptions { * @example 'input', 'output', 'temp' */ folder?: ResultItemType + onUploadStart?: (files: File[]) => void + onUploadError?: () => void } /** @@ -90,10 +93,29 @@ export const useNodeImageUpload = ( } const handleUploadBatch = async (files: File[]) => { - const paths = await Promise.all(files.map(handleUpload)) - const validPaths = paths.filter((p): p is string => !!p) - if (validPaths.length) onUploadComplete(validPaths) - return validPaths + if (node.isUploading) { + useToastStore().addAlert(t('g.uploadAlreadyInProgress')) + return [] + } + node.isUploading = true + + try { + node.imgs = undefined + node.graph?.setDirtyCanvas(true) + options.onUploadStart?.(files) + + const paths = await Promise.all(files.map(handleUpload)) + const validPaths = paths.filter((p): p is string => !!p) + if (validPaths.length) { + onUploadComplete(validPaths) + } else { + options.onUploadError?.() + } + return validPaths + } finally { + node.isUploading = false + node.graph?.setDirtyCanvas(true) + } } // Handle drag & drop diff --git a/src/extensions/core/uploadAudio.ts b/src/extensions/core/uploadAudio.ts index 2dc453b691..8b6f65ab26 100644 --- a/src/extensions/core/uploadAudio.ts +++ b/src/extensions/core/uploadAudio.ts @@ -43,7 +43,7 @@ async function uploadFile( file: File, updateNode: boolean, pasted: boolean = false -) { +): Promise { try { // Wrap file in formdata so it includes filename const body = new FormData() @@ -76,12 +76,15 @@ async function uploadFile( // Manually trigger the callback to update VueNodes audioWidget.callback?.(path) } + return true } else { useToastStore().addAlert(resp.status + ' - ' + resp.statusText) + return false } } catch (error) { // @ts-expect-error fixme ts strict error useToastStore().addAlert(error) + return false } } @@ -232,7 +235,17 @@ app.registerExtension({ const handleUpload = async (files: File[]) => { if (files?.length) { - uploadFile(audioWidget, audioUIWidget, files[0], true) + const previousValue = audioWidget.value + audioWidget.value = files[0].name + const success = await uploadFile( + audioWidget, + audioUIWidget, + files[0], + true + ) + if (!success) { + audioWidget.value = previousValue + } } return files } diff --git a/src/locales/en/main.json b/src/locales/en/main.json index 8b15f13cf1..c72ce2089c 100644 --- a/src/locales/en/main.json +++ b/src/locales/en/main.json @@ -174,6 +174,7 @@ "control_after_generate": "control after generate", "control_before_generate": "control before generate", "choose_file_to_upload": "choose file to upload", + "uploadAlreadyInProgress": "Upload already in progress", "capture": "capture", "nodes": "Nodes", "nodesCount": "{count} nodes | {count} node | {count} nodes", diff --git a/src/renderer/extensions/vueNodes/widgets/composables/useImagePreviewWidget.ts b/src/renderer/extensions/vueNodes/widgets/composables/useImagePreviewWidget.ts index 90beab2be2..a76a3cb0f6 100644 --- a/src/renderer/extensions/vueNodes/widgets/composables/useImagePreviewWidget.ts +++ b/src/renderer/extensions/vueNodes/widgets/composables/useImagePreviewWidget.ts @@ -43,6 +43,36 @@ function scheduleDeferredImageRender() { }) } +const TWO_PI = Math.PI * 2 +const SPINNER_ARC_LENGTH = Math.PI * 1.5 + +function renderUploadSpinner( + ctx: CanvasRenderingContext2D, + node: LGraphNode, + shiftY: number, + computedHeight: number | undefined +) { + const dw = node.size[0] + const dh = computedHeight ?? 220 + const centerX = dw / 2 + const centerY = shiftY + dh / 2 + const radius = 16 + const angle = ((Date.now() % 1000) / 1000) * TWO_PI + + ctx.save() + ctx.strokeStyle = LiteGraph.NODE_TEXT_COLOR + ctx.lineWidth = 3 + ctx.lineCap = 'round' + ctx.beginPath() + ctx.arc(centerX, centerY, radius, angle, angle + SPINNER_ARC_LENGTH) + ctx.stroke() + ctx.restore() + + // Schedule next frame to keep spinner animating continuously. + // Only runs while node.isUploading is true (checked by caller). + node.graph?.setDirtyCanvas(true) +} + const renderPreview = ( ctx: CanvasRenderingContext2D, node: LGraphNode, @@ -51,6 +81,11 @@ const renderPreview = ( ) => { if (!node.size) return + if (node.isUploading) { + renderUploadSpinner(ctx, node, shiftY, computedHeight) + return + } + const canvas = useCanvasStore().getCanvas() const mouse = canvas.graph_mouse @@ -65,6 +100,8 @@ const renderPreview = ( } const imgs = node.imgs ?? [] + if (imgs.length === 0) return + let { imageIndex } = node const numImages = imgs.length if (numImages === 1 && !imageIndex) { diff --git a/src/renderer/extensions/vueNodes/widgets/composables/useImageUploadWidget.ts b/src/renderer/extensions/vueNodes/widgets/composables/useImageUploadWidget.ts index e38c57f400..4ced75dad4 100644 --- a/src/renderer/extensions/vueNodes/widgets/composables/useImageUploadWidget.ts +++ b/src/renderer/extensions/vueNodes/widgets/composables/useImageUploadWidget.ts @@ -54,12 +54,27 @@ export const useImageUploadWidget = () => { createAnnotatedPath(value, { rootFolder: image_folder }) // Setup file upload handling + let rollback: (() => void) | undefined const { openFileSelection } = useNodeImageUpload(node, { allow_batch, fileFilter, accept, folder, + onUploadStart: (files) => { + if (files.length > 0) { + const prev = fileComboWidget.value + fileComboWidget.value = files[0].name + rollback = () => { + fileComboWidget.value = prev + } + } + }, + onUploadError: () => { + rollback?.() + rollback = undefined + }, onUploadComplete: (output) => { + rollback = undefined const annotated = output.map(formatPath) annotated.forEach((path) => { addToComboValues(fileComboWidget, path) @@ -88,6 +103,7 @@ export const useImageUploadWidget = () => { // Add our own callback to the combo widget to render an image when it changes fileComboWidget.callback = function () { + node.imgs = undefined nodeOutputStore.setNodeOutputs(node, String(fileComboWidget.value), { isAnimated }) diff --git a/src/types/litegraph-augmentation.d.ts b/src/types/litegraph-augmentation.d.ts index 5900fa5c64..8ee82a1d42 100644 --- a/src/types/litegraph-augmentation.d.ts +++ b/src/types/litegraph-augmentation.d.ts @@ -175,6 +175,8 @@ declare module '@/lib/litegraph/src/litegraph' { videoContainer?: HTMLElement /** Whether the node's preview media is loading */ isLoading?: boolean + /** Whether a file is being uploaded to this node */ + isUploading?: boolean /** The content type of the node's preview media */ previewMediaType?: 'image' | 'video' | 'audio' | 'model' /** If true, output images are stored but not rendered below the node */