diff --git a/src/composables/widgets/useTextUploadWidget.ts b/src/composables/widgets/useTextUploadWidget.ts new file mode 100644 index 000000000..6b4a9e093 --- /dev/null +++ b/src/composables/widgets/useTextUploadWidget.ts @@ -0,0 +1,136 @@ +import type { IComboWidget } from '@comfyorg/litegraph/dist/types/widgets' + +import { useNodeDragAndDrop } from '@/composables/node/useNodeDragAndDrop' +import { useNodeFileInput } from '@/composables/node/useNodeFileInput' +import { useNodePaste } from '@/composables/node/useNodePaste' +import { t } from '@/i18n' +import { api } from '@/scripts/api' +import type { ComfyWidgetConstructor } from '@/scripts/widgets' +import { useToastStore } from '@/stores/toastStore' +import { addToComboValues } from '@/utils/litegraphUtil' + +// Support only txt and pdf files +const isTextFile = (file: File): boolean => + file.type === 'text/plain' || file.name.toLowerCase().endsWith('.txt') + +const isPdfFile = (file: File): boolean => + file.type === 'application/pdf' || file.name.toLowerCase().endsWith('.pdf') + +const isSupportedFile = (file: File): boolean => + isTextFile(file) || isPdfFile(file) + +// Upload a text file and return the file path +async function uploadTextFile(file: File): Promise { + try { + const body = new FormData() + body.append('image', file) // Using standard field name for compatibility + + const resp = await api.fetchApi('/upload/image', { + method: 'POST', + body + }) + + if (resp.status === 200) { + const data = await resp.json() + let path = data.name + if (data.subfolder) { + path = data.subfolder + '/' + path + } + return path + } else { + useToastStore().addAlert(resp.status + ' - ' + resp.statusText) + } + } catch (error) { + useToastStore().addAlert(String(error)) + } + + return null +} + +export const useTextUploadWidget = (): ComfyWidgetConstructor => { + return (node, inputName, inputData) => { + // Early return with empty button if no valid inputs + if ( + !node.widgets || + !inputData[1] || + typeof inputData[1].textInputName !== 'string' + ) { + const emptyButton = node.addWidget('button', inputName, '', () => {}) + return { widget: emptyButton } + } + + // Get the name of the text input widget from the spec + const textInputName = inputData[1].textInputName as string + + // Find the combo widget that will store the file path + const textWidget = node.widgets.find((w) => w.name === textInputName) as + | IComboWidget + | undefined + + if (!textWidget) { + console.error(`Text widget with name "${textInputName}" not found`) + const fallbackButton = node.addWidget('button', inputName, '', () => {}) + return { widget: fallbackButton } + } + + // Ensure options and values are initialized + if (!textWidget.options) { + textWidget.options = { values: [] } + } else if (!textWidget.options.values) { + textWidget.options.values = [] + } + + // Handle the file upload + const handleFileUpload = async (files: File[]): Promise => { + if (!files.length) return files + + const file = files[0] + if (!isSupportedFile(file)) { + useToastStore().addAlert(t('toastMessages.invalidFileType')) + return files + } + + const path = await uploadTextFile(file) + if (path && textWidget) { + addToComboValues(textWidget, path) + textWidget.value = path + if (textWidget.callback) { + textWidget.callback(path) + } + } + + return files + } + + // Set up file input for upload button + const { openFileSelection } = useNodeFileInput(node, { + accept: '.txt,.pdf,text/plain,application/pdf', + onSelect: handleFileUpload + }) + + // Set up drag and drop + useNodeDragAndDrop(node, { + fileFilter: isSupportedFile, + onDrop: handleFileUpload + }) + + // Set up paste + useNodePaste(node, { + fileFilter: isSupportedFile, + onPaste: handleFileUpload + }) + + // Create upload button widget + const uploadWidget = node.addWidget( + 'button', + inputName, + '', + openFileSelection, + { serialize: false } + ) + + uploadWidget.label = t('g.choose_file_to_upload') + + return { widget: uploadWidget } + } +} diff --git a/src/extensions/core/index.ts b/src/extensions/core/index.ts index 0d39c207a..e53b12c29 100644 --- a/src/extensions/core/index.ts +++ b/src/extensions/core/index.ts @@ -18,5 +18,6 @@ import './simpleTouchSupport' import './slotDefaults' import './uploadAudio' import './uploadImage' +import './uploadText' import './webcamCapture' import './widgetInputs' diff --git a/src/extensions/core/uploadText.ts b/src/extensions/core/uploadText.ts new file mode 100644 index 000000000..edf699a2e --- /dev/null +++ b/src/extensions/core/uploadText.ts @@ -0,0 +1,50 @@ +import { + ComfyNodeDef, + InputSpec, + isComboInputSpecV1 +} from '@/schemas/nodeDefSchema' + +import { app } from '../../scripts/app' + +// Adds an upload button to text/pdf nodes + +const isTextUploadComboInput = (inputSpec: InputSpec) => { + const [inputName, inputOptions] = inputSpec + if (!inputOptions) return false + + const isUploadInput = inputOptions['text_upload'] === true + + return ( + isUploadInput && (isComboInputSpecV1(inputSpec) || inputName === 'COMBO') + ) +} + +const createUploadInput = ( + textInputName: string, + textInputOptions: InputSpec +): InputSpec => [ + 'TEXTUPLOAD', + { + ...textInputOptions[1], + textInputName + } +] + +app.registerExtension({ + name: 'Comfy.UploadText', + beforeRegisterNodeDef(_nodeType, nodeData: ComfyNodeDef) { + const { input } = nodeData ?? {} + const { required } = input ?? {} + if (!required) return + + const found = Object.entries(required).find(([_, input]) => + isTextUploadComboInput(input) + ) + + // If text/pdf combo input found, attach upload input + if (found) { + const [inputName, inputSpec] = found + required.upload = createUploadInput(inputName, inputSpec) + } + } +}) diff --git a/src/schemas/nodeDefSchema.ts b/src/schemas/nodeDefSchema.ts index a3624695e..83129d450 100644 --- a/src/schemas/nodeDefSchema.ts +++ b/src/schemas/nodeDefSchema.ts @@ -76,6 +76,7 @@ export const zComboInputOptions = zBaseInputOptions.extend({ allow_batch: z.boolean().optional(), video_upload: z.boolean().optional(), animated_image_upload: z.boolean().optional(), + text_upload: z.boolean().optional(), options: z.array(zComboOption).optional(), remote: zRemoteWidgetConfig.optional(), /** Whether the widget is a multi-select widget. */ diff --git a/src/scripts/widgets.ts b/src/scripts/widgets.ts index 90f5e54f7..ddd013da8 100644 --- a/src/scripts/widgets.ts +++ b/src/scripts/widgets.ts @@ -12,6 +12,7 @@ import { useImageUploadWidget } from '@/composables/widgets/useImageUploadWidget import { useIntWidget } from '@/composables/widgets/useIntWidget' import { useMarkdownWidget } from '@/composables/widgets/useMarkdownWidget' import { useStringWidget } from '@/composables/widgets/useStringWidget' +import { useTextUploadWidget } from '@/composables/widgets/useTextUploadWidget' import { t } from '@/i18n' import { transformInputSpecV1ToV2 } from '@/schemas/nodeDef/migration' import type { InputSpec as InputSpecV2 } from '@/schemas/nodeDef/nodeDefSchemaV2' @@ -289,5 +290,6 @@ export const ComfyWidgets: Record = { STRING: transformWidgetConstructorV2ToV1(useStringWidget()), MARKDOWN: transformWidgetConstructorV2ToV1(useMarkdownWidget()), COMBO: transformWidgetConstructorV2ToV1(useComboWidget()), - IMAGEUPLOAD: useImageUploadWidget() + IMAGEUPLOAD: useImageUploadWidget(), + TEXTUPLOAD: useTextUploadWidget() }