From df11c99393ac70ccd5849b0b34e14ffa6572d36d Mon Sep 17 00:00:00 2001 From: bymyself Date: Sun, 16 Feb 2025 08:09:02 -0700 Subject: [PATCH] Refactor node image upload and preview (#2580) Co-authored-by: huchenlei --- src/composables/useNodeDragAndDrop.ts | 49 ++++ src/composables/useNodeImage.ts | 31 +++ src/composables/useNodeImageUpload.ts | 107 +++++++++ src/composables/useNodePaste.ts | 27 +++ .../widgets/useImageUploadWidget.ts | 214 +++++------------- src/services/litegraphService.ts | 45 ++-- src/stores/imagePreviewStore.ts | 85 +++++++ src/utils/formatUtil.ts | 19 ++ src/utils/imageUtil.ts | 12 + 9 files changed, 403 insertions(+), 186 deletions(-) create mode 100644 src/composables/useNodeDragAndDrop.ts create mode 100644 src/composables/useNodeImage.ts create mode 100644 src/composables/useNodeImageUpload.ts create mode 100644 src/composables/useNodePaste.ts create mode 100644 src/stores/imagePreviewStore.ts create mode 100644 src/utils/imageUtil.ts diff --git a/src/composables/useNodeDragAndDrop.ts b/src/composables/useNodeDragAndDrop.ts new file mode 100644 index 000000000..90562d816 --- /dev/null +++ b/src/composables/useNodeDragAndDrop.ts @@ -0,0 +1,49 @@ +import type { LGraphNode } from '@comfyorg/litegraph' + +type DragHandler = (e: DragEvent) => boolean +type DropHandler = (files: File[]) => Promise + +interface DragAndDropOptions { + onDragOver?: DragHandler + onDrop: DropHandler + fileFilter?: (file: File) => boolean +} + +/** + * Adds drag and drop file handling to a node + */ +export const useNodeDragAndDrop = ( + node: LGraphNode, + options: DragAndDropOptions +) => { + const { onDragOver, onDrop, fileFilter = () => true } = options + + const hasFiles = (items: DataTransferItemList) => + !!Array.from(items).find((f) => f.kind === 'file') + + const filterFiles = (files: FileList) => Array.from(files).filter(fileFilter) + + const hasValidFiles = (files: FileList) => filterFiles(files).length > 0 + + const isDraggingFiles = (e: DragEvent | undefined) => { + if (!e?.dataTransfer?.items) return false + return onDragOver?.(e) ?? hasFiles(e.dataTransfer.items) + } + + const isDraggingValidFiles = (e: DragEvent | undefined) => { + if (!e?.dataTransfer?.files) return false + return hasValidFiles(e.dataTransfer.files) + } + + node.onDragOver = isDraggingFiles + + node.onDragDrop = function (e: DragEvent) { + if (!isDraggingValidFiles(e)) return false + + const files = filterFiles(e.dataTransfer!.files) + onDrop(files).then((results) => { + if (!results?.length) return + }) + return true + } +} diff --git a/src/composables/useNodeImage.ts b/src/composables/useNodeImage.ts new file mode 100644 index 000000000..659a57749 --- /dev/null +++ b/src/composables/useNodeImage.ts @@ -0,0 +1,31 @@ +import type { LGraphNode } from '@comfyorg/litegraph' + +import { useNodeOutputStore } from '@/stores/imagePreviewStore' + +interface NodeImageOptions { + allowBatch?: boolean +} + +/** + * Attaches a preview image to a node. + */ +export const useNodeImage = (node: LGraphNode, options: NodeImageOptions) => { + const { allowBatch = false } = options + const nodeOutputStore = useNodeOutputStore() + + /** Displays output image(s) on the node. */ + function showImage(output: string | string[]) { + if (!output) return + if (allowBatch || typeof output === 'string') { + nodeOutputStore.setNodeOutputs(node, output) + } else { + nodeOutputStore.setNodeOutputs(node, output[0]) + } + node.setSizeForImage?.() + node.graph?.setDirtyCanvas(true) + } + + return { + showImage + } +} diff --git a/src/composables/useNodeImageUpload.ts b/src/composables/useNodeImageUpload.ts new file mode 100644 index 000000000..87ebf28d3 --- /dev/null +++ b/src/composables/useNodeImageUpload.ts @@ -0,0 +1,107 @@ +import type { LGraphNode } from '@comfyorg/litegraph' + +import { api } from '@/scripts/api' +import { useToastStore } from '@/stores/toastStore' + +import { useNodeDragAndDrop } from './useNodeDragAndDrop' +import { useNodePaste } from './useNodePaste' + +const ACCEPTED_IMAGE_TYPES = 'image/jpeg,image/png,image/webp' +const PASTED_IMAGE_EXPIRY_MS = 2000 + +const createFileInput = () => { + const fileInput = document.createElement('input') + fileInput.type = 'file' + fileInput.accept = ACCEPTED_IMAGE_TYPES + return fileInput +} + +const uploadFile = async (file: File, isPasted: boolean) => { + const body = new FormData() + body.append('image', file) + if (isPasted) body.append('subfolder', 'pasted') + + const resp = await api.fetchApi('/upload/image', { + method: 'POST', + body + }) + + if (resp.status !== 200) { + useToastStore().addAlert(resp.status + ' - ' + resp.statusText) + return + } + + const data = await resp.json() + return data.subfolder ? `${data.subfolder}/${data.name}` : data.name +} + +interface ImageUploadOptions { + fileFilter?: (file: File) => boolean + onUploadComplete: (paths: string[]) => void +} + +export const useNodeImageUpload = ( + node: LGraphNode, + options: ImageUploadOptions +) => { + const { fileFilter = () => true, onUploadComplete } = options + + const isPastedFile = (file: File): boolean => + file.name === 'image.png' && + file.lastModified - Date.now() < PASTED_IMAGE_EXPIRY_MS + + const handleUpload = async (file: File) => { + try { + const path = await uploadFile(file, isPastedFile(file)) + if (!path) return + return path + } catch (error) { + useToastStore().addAlert(String(error)) + } + } + + // Handle drag & drop + useNodeDragAndDrop(node, { + fileFilter, + onDrop: async (files) => { + const paths = await Promise.all(files.map(handleUpload)) + const validPaths = paths.filter((p): p is string => !!p) + if (validPaths.length) { + onUploadComplete(validPaths) + } + return validPaths + } + }) + + // Handle paste + useNodePaste(node, { + fileFilter, + onPaste: async (file) => { + const path = await handleUpload(file) + if (path) { + onUploadComplete([path]) + } + return path + } + }) + + // Handle file input + const fileInput = createFileInput() + fileInput.onchange = async () => { + if (fileInput.files?.length) { + const paths = await Promise.all( + Array.from(fileInput.files).filter(fileFilter).map(handleUpload) + ) + const validPaths = paths.filter((p): p is string => !!p) + if (validPaths.length) { + onUploadComplete(validPaths) + } + } + } + document.body.append(fileInput) + + return { + fileInput, + handleUpload + } +} diff --git a/src/composables/useNodePaste.ts b/src/composables/useNodePaste.ts new file mode 100644 index 000000000..47b242ebb --- /dev/null +++ b/src/composables/useNodePaste.ts @@ -0,0 +1,27 @@ +import type { LGraphNode } from '@comfyorg/litegraph' + +type PasteHandler = (file: File) => Promise + +interface NodePasteOptions { + onPaste: PasteHandler + fileFilter?: (file: File) => boolean +} + +/** + * Adds paste handling to a node + */ +export const useNodePaste = ( + node: LGraphNode, + options: NodePasteOptions +) => { + const { onPaste, fileFilter = () => true } = options + + node.pasteFile = function (file: File) { + if (!fileFilter(file)) return false + + onPaste(file).then((result) => { + if (!result) return + }) + return true + } +} diff --git a/src/composables/widgets/useImageUploadWidget.ts b/src/composables/widgets/useImageUploadWidget.ts index f05d9be92..05b774e4e 100644 --- a/src/composables/widgets/useImageUploadWidget.ts +++ b/src/composables/widgets/useImageUploadWidget.ts @@ -1,11 +1,27 @@ import type { LGraphNode } from '@comfyorg/litegraph' -import type { IStringWidget } from '@comfyorg/litegraph/dist/types/widgets' +import { IComboWidget } from '@comfyorg/litegraph/dist/types/widgets' -import { api } from '@/scripts/api' +import { useNodeImage } from '@/composables/useNodeImage' +import { useNodeImageUpload } from '@/composables/useNodeImageUpload' import type { ComfyWidgetConstructor } from '@/scripts/widgets' -import { useToastStore } from '@/stores/toastStore' import type { ComfyApp } from '@/types' -import type { InputSpec } from '@/types/apiTypes' +import type { InputSpec, ResultItem } from '@/types/apiTypes' +import { createAnnotatedPath } from '@/utils/formatUtil' + +const isImageFile = (file: File) => file.type.startsWith('image/') + +const findFileComboWidget = (node: LGraphNode, inputData: InputSpec) => + node.widgets?.find( + (w) => w.name === (inputData[1]?.widget ?? 'image') && w.type === 'combo' + ) as IComboWidget & { value: string } + +const addToComboValues = (widget: IComboWidget, path: string) => { + if (!widget.options) widget.options = { values: [] } + if (!widget.options.values) widget.options.values = [] + if (!widget.options.values.includes(path)) { + widget.options.values.push(path) + } +} export const useImageUploadWidget = () => { const widgetConstructor: ComfyWidgetConstructor = ( @@ -14,174 +30,64 @@ export const useImageUploadWidget = () => { inputData: InputSpec, app: ComfyApp ) => { - const imageWidget = node.widgets?.find( - (w) => w.name === (inputData[1]?.widget ?? 'image') - ) as IStringWidget - const { image_folder = 'input' } = inputData[1] ?? {} + // TODO: specify upload widget via input spec rather than input name + const fileComboWidget = findFileComboWidget(node, inputData) + const { allow_batch, image_folder = 'input' } = inputData[1] ?? {} + const initialFile = `${fileComboWidget.value}` + const { showImage } = useNodeImage(node, { allowBatch: allow_batch }) - function showImage(name: string) { - const img = new Image() - img.onload = () => { - node.imgs = [img] - app.graph.setDirtyCanvas(true) - } - const folder_separator = name.lastIndexOf('/') - let subfolder = '' - if (folder_separator > -1) { - subfolder = name.substring(0, folder_separator) - name = name.substring(folder_separator + 1) - } - img.src = api.apiURL( - `/view?filename=${encodeURIComponent(name)}&type=${image_folder}&subfolder=${subfolder}${app.getPreviewFormatParam()}${app.getRandParam()}` - ) - node.setSizeForImage?.() - } + let internalValue: string | ResultItem = initialFile - const default_value = imageWidget.value - Object.defineProperty(imageWidget, 'value', { - set: function (value) { - this._real_value = value + // Setup getter/setter that transforms from `ResultItem` to string and formats paths + Object.defineProperty(fileComboWidget, 'value', { + set: function (value: string | ResultItem) { + internalValue = value }, - get: function () { - if (!this._real_value) { - return default_value - } - - let value = this._real_value - if (value.filename) { - const real_value = value - value = '' - if (real_value.subfolder) { - value = real_value.subfolder + '/' - } - - value += real_value.filename - - if (real_value.type && real_value.type !== 'input') - value += ` [${real_value.type}]` - } - return value + if (!internalValue) return initialFile + if (typeof internalValue === 'string') + return createAnnotatedPath(internalValue, { + rootFolder: image_folder + }) + if (!internalValue.filename) return initialFile + return createAnnotatedPath(internalValue) } }) - // Add our own callback to the combo widget to render an image when it changes + // Setup file upload handling + const { fileInput } = useNodeImageUpload(node, { + fileFilter: isImageFile, + onUploadComplete: (output) => { + output.forEach((path) => addToComboValues(fileComboWidget, path)) + fileComboWidget.value = output[0] + fileComboWidget.callback?.(output) + } + }) + + // Create the button widget for selecting the files + const uploadWidget = node.addWidget('button', inputName, 'image', () => + fileInput.click() + ) + uploadWidget.label = 'choose file to upload' + // @ts-expect-error serialize is not typed + uploadWidget.serialize = false + // TODO: Explain this? // @ts-expect-error LGraphNode.callback is not typed + // Add our own callback to the combo widget to render an image when it changes const cb = node.callback - imageWidget.callback = function (...args) { - showImage(imageWidget.value) - if (cb) { - return cb.apply(this, args) - } + fileComboWidget.callback = function (...args) { + showImage(fileComboWidget.value) + if (cb) return cb.apply(this, args) } // On load if we have a value then render the image // The value isnt set immediately so we need to wait a moment // No change callbacks seem to be fired on initial setting of the value requestAnimationFrame(() => { - if (imageWidget.value) { - showImage(imageWidget.value) - } + showImage(fileComboWidget.value) }) - // Add types for upload parameters - async function uploadFile(file: File, updateNode: boolean, pasted = false) { - try { - // Wrap file in formdata so it includes filename - const body = new FormData() - body.append('image', file) - if (pasted) body.append('subfolder', 'pasted') - const resp = await api.fetchApi('/upload/image', { - method: 'POST', - body - }) - - if (resp.status === 200) { - const data = await resp.json() - // Add the file to the dropdown list and update the widget value - let path = data.name - if (data.subfolder) path = data.subfolder + '/' + path - - if (!imageWidget.options) { - imageWidget.options = { values: [] } - } - if (!imageWidget.options.values) { - imageWidget.options.values = [] - } - if (!imageWidget.options.values.includes(path)) { - imageWidget.options.values.push(path) - } - - if (updateNode) { - showImage(path) - imageWidget.value = path - } - } else { - useToastStore().addAlert(resp.status + ' - ' + resp.statusText) - } - } catch (error) { - useToastStore().addAlert(String(error)) - } - } - - const fileInput = document.createElement('input') - Object.assign(fileInput, { - type: 'file', - accept: 'image/jpeg,image/png,image/webp', - style: 'display: none', - onchange: async () => { - // Add null check for files - if (fileInput.files && fileInput.files.length) { - await uploadFile(fileInput.files[0], true) - } - } - }) - document.body.append(fileInput) - - // Create the button widget for selecting the files - const uploadWidget = node.addWidget('button', inputName, 'image', () => { - fileInput.click() - }) - uploadWidget.label = 'choose file to upload' - // @ts-expect-error IWidget.serialize is not typed - uploadWidget.serialize = false - - // Add handler to check if an image is being dragged over our node - node.onDragOver = function (e: DragEvent) { - if (e.dataTransfer && e.dataTransfer.items) { - const image = [...e.dataTransfer.items].find((f) => f.kind === 'file') - return !!image - } - - return false - } - - // On drop upload files - node.onDragDrop = function (e: DragEvent) { - console.log('onDragDrop called') - let handled = false - if (e.dataTransfer?.files) { - for (const file of e.dataTransfer.files) { - if (file.type.startsWith('image/')) { - uploadFile(file, !handled) - handled = true - } - } - } - return handled - } - - node.pasteFile = function (file: File) { - if (file.type.startsWith('image/')) { - const is_pasted = - file.name === 'image.png' && file.lastModified - Date.now() < 2000 - uploadFile(file, true, is_pasted) - return true - } - return false - } - return { widget: uploadWidget } } diff --git a/src/services/litegraphService.ts b/src/services/litegraphService.ts index a897d605b..da132ec1e 100644 --- a/src/services/litegraphService.ts +++ b/src/services/litegraphService.ts @@ -10,15 +10,16 @@ import { Vector2 } from '@comfyorg/litegraph' import { IBaseWidget, IWidget } from '@comfyorg/litegraph/dist/types/widgets' import { st } from '@/i18n' -import { api } from '@/scripts/api' import { ANIM_PREVIEW_WIDGET, ComfyApp, app } from '@/scripts/app' import { $el } from '@/scripts/ui' import { calculateImageGrid, createImageHost } from '@/scripts/ui/imagePreview' import { useCanvasStore } from '@/stores/graphStore' +import { useNodeOutputStore } from '@/stores/imagePreviewStore' import { useToastStore } from '@/stores/toastStore' -import { ComfyNodeDef, ExecutedWsMessage } from '@/types/apiTypes' +import { ComfyNodeDef } from '@/types/apiTypes' import type { NodeId } from '@/types/comfyWorkflow' import { normalizeI18nKey } from '@/utils/formatUtil' +import { is_all_same_aspect_ratio } from '@/utils/imageUtil' import { isImageNode } from '@/utils/litegraphUtil' import { useExtensionService } from './extensionService' @@ -404,30 +405,22 @@ export const useLitegraphService = () => { ) { if (this.flags.collapsed) return - const imgURLs: (string[] | string)[] = [] - let imagesChanged = false - - const output: ExecutedWsMessage['output'] = app.nodeOutputs[this.id + ''] - if (output?.images && this.images !== output.images) { - this.animatedImages = output?.animated?.find(Boolean) + let imgURLs: string[] = [] + const nodeOutputStore = useNodeOutputStore() + const output = nodeOutputStore.getNodeOutputs(this) + let imagesChanged = nodeOutputStore.isImagesChanged(this) + if (output && imagesChanged) { + this.animatedImages = output.animated?.find(Boolean) this.images = output.images - imagesChanged = true - const preview = this.animatedImages ? '' : app.getPreviewFormatParam() - - for (const params of output.images) { - const imgUrlPart = new URLSearchParams(params).toString() - const rand = app.getRandParam() - const imgUrl = api.apiURL(`/view?${imgUrlPart}${preview}${rand}`) - imgURLs.push(imgUrl) - } + imgURLs = nodeOutputStore.getNodeImageUrls(this) } - const preview = app.nodePreviewImages[this.id + ''] + const preview = nodeOutputStore.getNodePreviews(this) if (this.preview !== preview) { this.preview = preview imagesChanged = true if (preview != null) { - imgURLs.push(preview) + imgURLs.push(...preview) } } @@ -454,22 +447,10 @@ export const useLitegraphService = () => { } }) } else { - this.imgs = null + this.imgs = undefined } } - const is_all_same_aspect_ratio = (imgs: HTMLImageElement[]) => { - // assume: imgs.length >= 2 - const ratio = imgs[0].naturalWidth / imgs[0].naturalHeight - - for (let i = 1; i < imgs.length; i++) { - const this_ratio = imgs[i].naturalWidth / imgs[i].naturalHeight - if (ratio != this_ratio) return false - } - - return true - } - // Nothing to do if (!this.imgs?.length) return diff --git a/src/stores/imagePreviewStore.ts b/src/stores/imagePreviewStore.ts new file mode 100644 index 000000000..5a1a5185f --- /dev/null +++ b/src/stores/imagePreviewStore.ts @@ -0,0 +1,85 @@ +import { LGraphNode } from '@comfyorg/litegraph' +import { defineStore } from 'pinia' + +import { api } from '@/scripts/api' +import { ExecutedWsMessage, ResultItem } from '@/types/apiTypes' + +const toOutputs = ( + filenames: string[], + type: string +): ExecutedWsMessage['output'] => { + return { + images: filenames.map((image) => { + return { filename: image, subfolder: '', type } + }) + } +} + +const getPreviewParam = (node: LGraphNode) => { + if (node.animatedImages) return '' + return app.getPreviewFormatParam() +} + +export const useNodeOutputStore = defineStore('nodeOutput', () => { + function getNodeOutputs(node: LGraphNode): ExecutedWsMessage['output'] { + return app.nodeOutputs[node.id + ''] + } + + function getNodePreviews(node: LGraphNode): string[] { + return app.nodePreviewImages[node.id + ''] + } + + function getNodeImageUrls(node: LGraphNode): string[] { + const outputs = getNodeOutputs(node) + if (!outputs?.images?.length) return [] + + return outputs.images.map((image) => { + const imgUrlPart = new URLSearchParams(image) + const rand = app.getRandParam() + const previewParam = getPreviewParam(node) + return api.apiURL(`/view?${imgUrlPart}${previewParam}${rand}`) + }) + } + + /** + * Checks if the node's images have changed from what's stored + * @returns true if images have changed, false otherwise + */ + function isImagesChanged(node: LGraphNode): boolean { + const currentImages = node.images || [] + const { images: newImages } = getNodeOutputs(node) ?? {} + if (!newImages?.length) return false + + return currentImages !== newImages + } + + function setNodeOutputs( + node: LGraphNode, + filenames: string | string[] | ResultItem, + options: { folder?: string } = {} + ) { + if (!filenames) return + + const { folder = 'input' } = options + const nodeId = node.id + '' + + if (typeof filenames === 'string') { + app.nodeOutputs[nodeId] = toOutputs([filenames], folder) + } else if (!Array.isArray(filenames)) { + app.nodeOutputs[nodeId] = filenames + } else { + const resultItems = toOutputs(filenames, folder) + if (!resultItems?.images?.length) return + + app.nodeOutputs[nodeId] = resultItems + } + } + + return { + getNodeOutputs, + getNodeImageUrls, + getNodePreviews, + setNodeOutputs, + isImagesChanged + } +}) diff --git a/src/utils/formatUtil.ts b/src/utils/formatUtil.ts index e408cfaaa..bf4e2093c 100644 --- a/src/utils/formatUtil.ts +++ b/src/utils/formatUtil.ts @@ -1,3 +1,5 @@ +import { ResultItem } from '@/types/apiTypes' + export function formatCamelCase(str: string): string { // Check if the string is camel case const isCamelCase = /^([A-Z][a-z]*)+$/.test(str) @@ -213,3 +215,20 @@ export function isValidUrl(url: string): boolean { return false } } + +const createAnnotation = (rootFolder = 'input'): string => + rootFolder !== 'input' ? ` [${rootFolder}]` : '' + +const createPath = (filename: string, subfolder = ''): string => + subfolder ? `${subfolder}/${filename}` : filename + +/** Creates annotated filepath in format used by folder_paths.py */ +export function createAnnotatedPath( + item: string | ResultItem, + options: { rootFolder?: string; subfolder?: string } = {} +): string { + const { rootFolder = 'input', subfolder } = options + if (typeof item === 'string') + return `${createPath(item, subfolder)}${createAnnotation(rootFolder)}` + return `${createPath(item.filename ?? '', item.subfolder)}${createAnnotation(item.type)}` +} diff --git a/src/utils/imageUtil.ts b/src/utils/imageUtil.ts new file mode 100644 index 000000000..5e6a43eb0 --- /dev/null +++ b/src/utils/imageUtil.ts @@ -0,0 +1,12 @@ +export const is_all_same_aspect_ratio = (imgs: HTMLImageElement[]): boolean => { + if (!imgs.length || imgs.length === 1) return true + + const ratio = imgs[0].naturalWidth / imgs[0].naturalHeight + + for (let i = 1; i < imgs.length; i++) { + const this_ratio = imgs[i].naturalWidth / imgs[i].naturalHeight + if (ratio != this_ratio) return false + } + + return true +}