diff --git a/src/composables/useNodeFileInput.ts b/src/composables/useNodeFileInput.ts new file mode 100644 index 000000000..6902bd5b1 --- /dev/null +++ b/src/composables/useNodeFileInput.ts @@ -0,0 +1,45 @@ +interface FileInputOptions { + accept?: string + allow_batch?: boolean + fileFilter?: (file: File) => boolean + onSelect: (files: File[]) => void +} + +/** + * Creates a file input for a node. + */ +export function useNodeFileInput(options: FileInputOptions) { + const { + accept, + allow_batch = false, + fileFilter = () => true, + onSelect + } = options + + const fileInput = document.createElement('input') + fileInput.type = 'file' + fileInput.accept = accept ?? '*' + fileInput.multiple = allow_batch + fileInput.style.visibility = 'hidden' + + fileInput.onchange = () => { + if (fileInput.files?.length) { + const files = Array.from(fileInput.files).filter(fileFilter) + if (files.length) onSelect(files) + } + } + + document.body.append(fileInput) + + /** + * Shows the system file picker dialog for selecting files. + */ + function openFileSelection() { + fileInput.click() + } + + return { + fileInput, + openFileSelection + } +} diff --git a/src/composables/useNodeImage.ts b/src/composables/useNodeImage.ts index 659a57749..a1a1b84bf 100644 --- a/src/composables/useNodeImage.ts +++ b/src/composables/useNodeImage.ts @@ -2,25 +2,16 @@ 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 +export const useNodeImage = (node: LGraphNode) => { 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]) - } + nodeOutputStore.setNodeOutputs(node, output) node.setSizeForImage?.() node.graph?.setDirtyCanvas(true) } diff --git a/src/composables/useNodeImageUpload.ts b/src/composables/useNodeImageUpload.ts index 87ebf28d3..a353f536a 100644 --- a/src/composables/useNodeImageUpload.ts +++ b/src/composables/useNodeImageUpload.ts @@ -1,21 +1,14 @@ import type { LGraphNode } from '@comfyorg/litegraph' +import { useNodeDragAndDrop } from '@/composables/useNodeDragAndDrop' +import { useNodeFileInput } from '@/composables/useNodeFileInput' +import { useNodePaste } from '@/composables/useNodePaste' 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) @@ -38,13 +31,17 @@ const uploadFile = async (file: File, isPasted: boolean) => { interface ImageUploadOptions { fileFilter?: (file: File) => boolean onUploadComplete: (paths: string[]) => void + allow_batch?: boolean } +/** + * Adds image upload to a node via drag & drop, paste, and file input. + */ export const useNodeImageUpload = ( node: LGraphNode, options: ImageUploadOptions ) => { - const { fileFilter = () => true, onUploadComplete } = options + const { fileFilter, onUploadComplete, allow_batch } = options const isPastedFile = (file: File): boolean => file.name === 'image.png' && @@ -60,48 +57,33 @@ 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 + } + // 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 - } + onDrop: handleUploadBatch }) // Handle paste useNodePaste(node, { fileFilter, - onPaste: async (file) => { - const path = await handleUpload(file) - if (path) { - onUploadComplete([path]) - } - return path - } + allow_batch, + onPaste: handleUploadBatch }) // 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) + const { openFileSelection } = useNodeFileInput({ + fileFilter, + allow_batch, + accept: ACCEPTED_IMAGE_TYPES, + onSelect: handleUploadBatch + }) - return { - fileInput, - handleUpload - } + return { openFileSelection, handleUpload } } diff --git a/src/composables/useNodePaste.ts b/src/composables/useNodePaste.ts index 47b242ebb..2b5a9ab04 100644 --- a/src/composables/useNodePaste.ts +++ b/src/composables/useNodePaste.ts @@ -1,10 +1,11 @@ import type { LGraphNode } from '@comfyorg/litegraph' -type PasteHandler = (file: File) => Promise +type PasteHandler = (files: File[]) => Promise interface NodePasteOptions { onPaste: PasteHandler fileFilter?: (file: File) => boolean + allow_batch?: boolean } /** @@ -14,12 +15,15 @@ export const useNodePaste = ( node: LGraphNode, options: NodePasteOptions ) => { - const { onPaste, fileFilter = () => true } = options + const { onPaste, fileFilter = () => true, allow_batch = false } = options - node.pasteFile = function (file: File) { - if (!fileFilter(file)) return false + node.pasteFiles = function (files: File[]) { + const filteredFiles = Array.from(files).filter(fileFilter) + if (!filteredFiles.length) return false - onPaste(file).then((result) => { + const paste = allow_batch ? filteredFiles : filteredFiles.slice(0, 1) + + onPaste(paste).then((result) => { if (!result) return }) return true diff --git a/src/composables/usePaste.ts b/src/composables/usePaste.ts index 31b8d71f3..9871daa8e 100644 --- a/src/composables/usePaste.ts +++ b/src/composables/usePaste.ts @@ -28,7 +28,7 @@ export const usePaste = () => { // Did you mean 'Clipboard'?ts(2551) // TODO: Not sure what the code wants to do. let data = e.clipboardData || window.clipboardData - const items = data.items + const items: DataTransferItemList = data.items // Look for image paste data for (const item of items) { @@ -54,7 +54,13 @@ export const usePaste = () => { graph.change() } const blob = item.getAsFile() - imageNode?.pasteFile?.(blob) + if (blob) imageNode?.pasteFile?.(blob) + + imageNode?.pasteFiles?.( + Array.from(items) + .map((i) => i.getAsFile()) + .filter((f) => f !== null) + ) return } } diff --git a/src/composables/useValueTransform.ts b/src/composables/useValueTransform.ts new file mode 100644 index 000000000..be4b7a03d --- /dev/null +++ b/src/composables/useValueTransform.ts @@ -0,0 +1,31 @@ +/** + * Creates a getter/setter pair that transforms values on access if they have changed. + * Does not observe deep changes. + * + * @example + * const { get, set } = useValueTransform( + * items => items.map(formatPath) + * ) + * + * Object.defineProperty(obj, 'value', { get, set }) + */ +export function useValueTransform( + transform: (value: Internal) => External, + initialValue: Internal +) { + let internalValue: Internal = initialValue + let cachedValue: External = transform(initialValue) + let isChanged = false + + return { + get: () => { + if (!isChanged) return cachedValue + cachedValue = transform(internalValue) + return cachedValue + }, + set: (value: Internal) => { + isChanged = true + internalValue = value + } + } +} diff --git a/src/composables/widgets/useImageUploadWidget.ts b/src/composables/widgets/useImageUploadWidget.ts index 05b774e4e..60b54027d 100644 --- a/src/composables/widgets/useImageUploadWidget.ts +++ b/src/composables/widgets/useImageUploadWidget.ts @@ -3,25 +3,23 @@ import { IComboWidget } from '@comfyorg/litegraph/dist/types/widgets' import { useNodeImage } from '@/composables/useNodeImage' import { useNodeImageUpload } from '@/composables/useNodeImageUpload' +import { useValueTransform } from '@/composables/useValueTransform' import type { ComfyWidgetConstructor } from '@/scripts/widgets' import type { ComfyApp } from '@/types' import type { InputSpec, ResultItem } from '@/types/apiTypes' import { createAnnotatedPath } from '@/utils/formatUtil' +import { addToComboValues } from '@/utils/litegraphUtil' + +type InternalFile = string | ResultItem +type InternalValue = InternalFile | InternalFile[] +type ExposedValue = string | string[] 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) +const findFileComboWidget = (node: LGraphNode, inputName: string) => + node.widgets!.find((w) => w.name === inputName) as IComboWidget & { + value: ExposedValue } -} export const useImageUploadWidget = () => { const widgetConstructor: ComfyWidgetConstructor = ( @@ -30,43 +28,46 @@ export const useImageUploadWidget = () => { inputData: InputSpec, app: ComfyApp ) => { - // 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 inputOptions = inputData[1] ?? {} + const { imageInputName, allow_batch, image_folder = 'input' } = inputOptions + + const { showImage } = useNodeImage(node) + + const fileComboWidget = findFileComboWidget(node, imageInputName) const initialFile = `${fileComboWidget.value}` - const { showImage } = useNodeImage(node, { allowBatch: allow_batch }) + const formatPath = (value: InternalFile) => + createAnnotatedPath(value, { rootFolder: image_folder }) - let internalValue: string | ResultItem = initialFile + const transform = (internalValue: InternalValue): ExposedValue => { + if (!internalValue) return initialFile + if (Array.isArray(internalValue)) + return allow_batch + ? internalValue.map(formatPath) + : formatPath(internalValue[0]) + return formatPath(internalValue) + } - // 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 (!internalValue) return initialFile - if (typeof internalValue === 'string') - return createAnnotatedPath(internalValue, { - rootFolder: image_folder - }) - if (!internalValue.filename) return initialFile - return createAnnotatedPath(internalValue) - } - }) + Object.defineProperty( + fileComboWidget, + 'value', + useValueTransform(transform, initialFile) + ) // Setup file upload handling - const { fileInput } = useNodeImageUpload(node, { + const { openFileSelection } = useNodeImageUpload(node, { + allow_batch, fileFilter: isImageFile, onUploadComplete: (output) => { output.forEach((path) => addToComboValues(fileComboWidget, path)) - fileComboWidget.value = output[0] + // @ts-expect-error litegraph combo value type does not support arrays yet + fileComboWidget.value = output fileComboWidget.callback?.(output) } }) // Create the button widget for selecting the files const uploadWidget = node.addWidget('button', inputName, 'image', () => - fileInput.click() + openFileSelection() ) uploadWidget.label = 'choose file to upload' // @ts-expect-error serialize is not typed diff --git a/src/extensions/core/uploadImage.ts b/src/extensions/core/uploadImage.ts index c50b663d3..e3bb7c12d 100644 --- a/src/extensions/core/uploadImage.ts +++ b/src/extensions/core/uploadImage.ts @@ -1,22 +1,41 @@ -import { ComfyNodeDef, InputSpec } from '@/types/apiTypes' +import { ComfyNodeDef, InputSpec, isComboInputSpecV1 } from '@/types/apiTypes' import { app } from '../../scripts/app' // Adds an upload button to the nodes +const isImageComboInput = (inputSpec: InputSpec) => { + const [inputName, inputOptions] = inputSpec + if (!inputOptions || inputOptions['image_upload'] !== true) return false + return isComboInputSpecV1(inputSpec) || inputName === 'COMBO' +} + +const createUploadInput = ( + imageInputName: string, + imageInputOptions: InputSpec +): InputSpec => [ + 'IMAGEUPLOAD', + { + ...imageInputOptions[1], + imageInputName + } +] + app.registerExtension({ name: 'Comfy.UploadImage', beforeRegisterNodeDef(nodeType, nodeData: ComfyNodeDef) { - // Check if there is a required input named 'image' in the nodeData - const imageInputSpec: InputSpec | undefined = - nodeData?.input?.required?.image + const { input } = nodeData ?? {} + const { required } = input ?? {} + if (!required) return - // Get the config from the image input spec if it exists - const config = imageInputSpec?.[1] ?? {} - const { image_upload = false, image_folder = 'input' } = config + const found = Object.entries(required).find(([_, input]) => + isImageComboInput(input) + ) - if (image_upload && nodeData?.input?.required) { - nodeData.input.required.upload = ['IMAGEUPLOAD', { image_folder }] + // If image combo input found, attach upload input + if (found) { + const [inputName, inputSpec] = found + required.upload = createUploadInput(inputName, inputSpec) } } }) diff --git a/src/types/apiTypes.ts b/src/types/apiTypes.ts index 7222edc56..86e9c6426 100644 --- a/src/types/apiTypes.ts +++ b/src/types/apiTypes.ts @@ -349,6 +349,7 @@ const zComboInputProps = zBaseInputSpecValue.extend({ control_after_generate: z.boolean().optional(), image_upload: z.boolean().optional(), image_folder: z.enum(['input', 'output', 'temp']).optional(), + allow_batch: z.boolean().optional(), remote: zRemoteWidgetConfig.optional() }) diff --git a/src/types/litegraph-augmentation.d.ts b/src/types/litegraph-augmentation.d.ts index fe79a54f3..bbf9c2a04 100644 --- a/src/types/litegraph-augmentation.d.ts +++ b/src/types/litegraph-augmentation.d.ts @@ -122,6 +122,8 @@ declare module '@comfyorg/litegraph' { imageOffset?: number /** Callback for pasting an image file into the node */ pasteFile?(file: File): void + /** Callback for pasting multiple files into the node */ + pasteFiles?(files: File[]): void } } diff --git a/src/utils/litegraphUtil.ts b/src/utils/litegraphUtil.ts index ccbce6d75..fdbbd4132 100644 --- a/src/utils/litegraphUtil.ts +++ b/src/utils/litegraphUtil.ts @@ -1,4 +1,5 @@ import type { IWidget, LGraphNode } from '@comfyorg/litegraph' +import type { IComboWidget } from '@comfyorg/litegraph/dist/types/widgets' export function isImageNode(node: LGraphNode) { return ( @@ -8,3 +9,11 @@ export function isImageNode(node: LGraphNode) { node.widgets.findIndex((obj: IWidget) => obj.name === 'image') >= 0) ) } + +export function addToComboValues(widget: IComboWidget, value: string) { + if (!widget.options) widget.options = { values: [] } + if (!widget.options.values) widget.options.values = [] + if (!widget.options.values.includes(value)) { + widget.options.values.push(value) + } +}