From 30ee669f5c9d4b9173fd4a247a97620f1d5e2f3a Mon Sep 17 00:00:00 2001 From: Christian Byrne Date: Fri, 30 May 2025 02:05:41 -0700 Subject: [PATCH] [refactor] Refactor file handling (#3955) --- src/scripts/app.ts | 192 +++---------- src/utils/fileHandlers.ts | 336 ++++++++++++++++++++++ tests-ui/tests/utils/fileHandlers.test.ts | 127 ++++++++ 3 files changed, 495 insertions(+), 160 deletions(-) create mode 100644 src/utils/fileHandlers.ts create mode 100644 tests-ui/tests/utils/fileHandlers.test.ts diff --git a/src/scripts/app.ts b/src/scripts/app.ts index ad4ed91fc..8f34e420d 100644 --- a/src/scripts/app.ts +++ b/src/scripts/app.ts @@ -30,12 +30,6 @@ import { isComboInputSpecV1, isComboInputSpecV2 } from '@/schemas/nodeDefSchema' -import { getFromWebmFile } from '@/scripts/metadata/ebml' -import { getGltfBinaryMetadata } from '@/scripts/metadata/gltf' -import { getFromIsobmffFile } from '@/scripts/metadata/isobmff' -import { getMp3Metadata } from '@/scripts/metadata/mp3' -import { getOggMetadata } from '@/scripts/metadata/ogg' -import { getSvgMetadata } from '@/scripts/metadata/svg' import { useDialogService } from '@/services/dialogService' import { useExtensionService } from '@/services/extensionService' import { useLitegraphService } from '@/services/litegraphService' @@ -72,13 +66,7 @@ import { deserialiseAndCreate } from '@/utils/vintageClipboard' import { type ComfyApi, PromptExecutionError, api } from './api' import { defaultGraph } from './defaultGraph' import { pruneWidgets } from './domWidget' -import { - getFlacMetadata, - getLatentMetadata, - getPngMetadata, - getWebpMetadata, - importA1111 -} from './pnginfo' +import { importA1111 } from './pnginfo' import { $el, ComfyUI } from './ui' import { ComfyAppMenu } from './ui/menu/index' import { clone } from './utils' @@ -1274,6 +1262,7 @@ export class ComfyApp { * @param {File} file */ async handleFile(file: File) { + const { getFileHandler } = await import('@/utils/fileHandlers') const removeExt = (f: string) => { if (!f) return f const p = f.lastIndexOf('.') @@ -1281,161 +1270,44 @@ export class ComfyApp { return f.substring(0, p) } const fileName = removeExt(file.name) - if (file.type === 'image/png') { - const pngInfo = await getPngMetadata(file) - if (pngInfo?.workflow) { - await this.loadGraphData( - JSON.parse(pngInfo.workflow), - true, - true, - fileName - ) - } else if (pngInfo?.prompt) { - this.loadApiJson(JSON.parse(pngInfo.prompt), fileName) - } else if (pngInfo?.parameters) { - // Note: Not putting this in `importA1111` as it is mostly not used - // by external callers, and `importA1111` has no access to `app`. + + // Get the appropriate file handler for this file type + const fileHandler = getFileHandler(file) + + if (!fileHandler) { + // No handler found for this file type + this.showErrorOnFileLoad(file) + return + } + + try { + // Process the file using the handler + const { workflow, prompt, parameters, jsonTemplateData } = + await fileHandler(file) + + if (workflow) { + // We have a workflow, load it + await this.loadGraphData(workflow, true, true, fileName) + } else if (prompt) { + // We have a prompt in API format, load it + this.loadApiJson(prompt, fileName) + } else if (parameters) { + // We have A1111 parameters, import them useWorkflowService().beforeLoadNewGraph() - importA1111(this.graph, pngInfo.parameters) + importA1111(this.graph, parameters) useWorkflowService().afterLoadNewGraph( fileName, this.graph.serialize() as unknown as ComfyWorkflowJSON ) + } else if (jsonTemplateData) { + // We have template data from JSON + this.loadTemplateData(jsonTemplateData) } else { + // No usable data found in the file this.showErrorOnFileLoad(file) } - } else if (file.type === 'image/webp') { - const pngInfo = await getWebpMetadata(file) - // Support loading workflows from that webp custom node. - const workflow = pngInfo?.workflow || pngInfo?.Workflow - const prompt = pngInfo?.prompt || pngInfo?.Prompt - - if (workflow) { - this.loadGraphData(JSON.parse(workflow), true, true, fileName) - } else if (prompt) { - this.loadApiJson(JSON.parse(prompt), fileName) - } else { - this.showErrorOnFileLoad(file) - } - } else if (file.type === 'audio/mpeg') { - const { workflow, prompt } = await getMp3Metadata(file) - if (workflow) { - this.loadGraphData(workflow, true, true, fileName) - } else if (prompt) { - this.loadApiJson(prompt, fileName) - } else { - this.showErrorOnFileLoad(file) - } - } else if (file.type === 'audio/ogg') { - const { workflow, prompt } = await getOggMetadata(file) - if (workflow) { - this.loadGraphData(workflow, true, true, fileName) - } else if (prompt) { - this.loadApiJson(prompt, fileName) - } else { - this.showErrorOnFileLoad(file) - } - } else if (file.type === 'audio/flac' || file.type === 'audio/x-flac') { - const pngInfo = await getFlacMetadata(file) - const workflow = pngInfo?.workflow || pngInfo?.Workflow - const prompt = pngInfo?.prompt || pngInfo?.Prompt - - if (workflow) { - this.loadGraphData(JSON.parse(workflow), true, true, fileName) - } else if (prompt) { - this.loadApiJson(JSON.parse(prompt), fileName) - } else { - this.showErrorOnFileLoad(file) - } - } else if (file.type === 'video/webm') { - const webmInfo = await getFromWebmFile(file) - if (webmInfo.workflow) { - this.loadGraphData(webmInfo.workflow, true, true, fileName) - } else if (webmInfo.prompt) { - this.loadApiJson(webmInfo.prompt, fileName) - } else { - this.showErrorOnFileLoad(file) - } - } else if ( - file.type === 'video/mp4' || - file.name?.endsWith('.mp4') || - file.name?.endsWith('.mov') || - file.name?.endsWith('.m4v') || - file.type === 'video/quicktime' || - file.type === 'video/x-m4v' - ) { - const mp4Info = await getFromIsobmffFile(file) - if (mp4Info.workflow) { - this.loadGraphData(mp4Info.workflow, true, true, fileName) - } else if (mp4Info.prompt) { - this.loadApiJson(mp4Info.prompt, fileName) - } - } else if (file.type === 'image/svg+xml' || file.name?.endsWith('.svg')) { - const svgInfo = await getSvgMetadata(file) - if (svgInfo.workflow) { - this.loadGraphData(svgInfo.workflow, true, true, fileName) - } else if (svgInfo.prompt) { - this.loadApiJson(svgInfo.prompt, fileName) - } else { - this.showErrorOnFileLoad(file) - } - } else if ( - file.type === 'model/gltf-binary' || - file.name?.endsWith('.glb') - ) { - const gltfInfo = await getGltfBinaryMetadata(file) - if (gltfInfo.workflow) { - this.loadGraphData(gltfInfo.workflow, true, true, fileName) - } else if (gltfInfo.prompt) { - this.loadApiJson(gltfInfo.prompt, fileName) - } else { - this.showErrorOnFileLoad(file) - } - } else if ( - file.type === 'application/json' || - file.name?.endsWith('.json') - ) { - const reader = new FileReader() - reader.onload = async () => { - const readerResult = reader.result as string - const jsonContent = JSON.parse(readerResult) - if (jsonContent?.templates) { - this.loadTemplateData(jsonContent) - } else if (this.isApiJson(jsonContent)) { - this.loadApiJson(jsonContent, fileName) - } else { - await this.loadGraphData( - JSON.parse(readerResult), - true, - true, - fileName - ) - } - } - reader.readAsText(file) - } else if ( - file.name?.endsWith('.latent') || - file.name?.endsWith('.safetensors') - ) { - const info = await getLatentMetadata(file) - // TODO define schema to LatentMetadata - // @ts-expect-error - if (info.workflow) { - await this.loadGraphData( - // @ts-expect-error - JSON.parse(info.workflow), - true, - true, - fileName - ) - // @ts-expect-error - } else if (info.prompt) { - // @ts-expect-error - this.loadApiJson(JSON.parse(info.prompt)) - } else { - this.showErrorOnFileLoad(file) - } - } else { + } catch (error) { + console.error('Error processing file:', error) this.showErrorOnFileLoad(file) } } diff --git a/src/utils/fileHandlers.ts b/src/utils/fileHandlers.ts new file mode 100644 index 000000000..01896c1ae --- /dev/null +++ b/src/utils/fileHandlers.ts @@ -0,0 +1,336 @@ +/** + * Maps MIME types and file extensions to handler functions for extracting + * workflow data from various file formats. Uses supportedWorkflowFormats.ts + * as the source of truth for supported formats. + */ +import { + AUDIO_WORKFLOW_FORMATS, + DATA_WORKFLOW_FORMATS, + IMAGE_WORKFLOW_FORMATS, + MODEL_WORKFLOW_FORMATS, + VIDEO_WORKFLOW_FORMATS +} from '@/constants/supportedWorkflowFormats' +import type { + ComfyApiWorkflow, + ComfyWorkflowJSON +} from '@/schemas/comfyWorkflowSchema' +import { getFromWebmFile } from '@/scripts/metadata/ebml' +import { getGltfBinaryMetadata } from '@/scripts/metadata/gltf' +import { getFromIsobmffFile } from '@/scripts/metadata/isobmff' +import { getMp3Metadata } from '@/scripts/metadata/mp3' +import { getOggMetadata } from '@/scripts/metadata/ogg' +import { getSvgMetadata } from '@/scripts/metadata/svg' +import { + getFlacMetadata, + getLatentMetadata, + getPngMetadata, + getWebpMetadata +} from '@/scripts/pnginfo' + +/** + * Type for the file handler function + */ +export type WorkflowFileHandler = (file: File) => Promise<{ + workflow?: ComfyWorkflowJSON + prompt?: ComfyApiWorkflow + parameters?: string + jsonTemplateData?: any // For template JSON data +}> + +/** + * Maps MIME types to file handlers for loading workflows from different file formats + */ +export const mimeTypeHandlers = new Map() + +/** + * Maps file extensions to file handlers for loading workflows + * Used as a fallback when MIME type detection fails + */ +export const extensionHandlers = new Map() + +/** + * Handler for PNG files + */ +const handlePngFile: WorkflowFileHandler = async (file) => { + const pngInfo = await getPngMetadata(file) + return { + workflow: pngInfo?.workflow ? JSON.parse(pngInfo.workflow) : undefined, + prompt: pngInfo?.prompt ? JSON.parse(pngInfo.prompt) : undefined, + parameters: pngInfo?.parameters + } +} + +/** + * Handler for WebP files + */ +const handleWebpFile: WorkflowFileHandler = async (file) => { + const pngInfo = await getWebpMetadata(file) + // Support loading workflows from that webp custom node. + const workflow = pngInfo?.workflow || pngInfo?.Workflow + const prompt = pngInfo?.prompt || pngInfo?.Prompt + + return { + workflow: workflow ? JSON.parse(workflow) : undefined, + prompt: prompt ? JSON.parse(prompt) : undefined + } +} + +/** + * Handler for SVG files + */ +const handleSvgFile: WorkflowFileHandler = async (file) => { + const svgInfo = await getSvgMetadata(file) + return { + workflow: svgInfo.workflow, + prompt: svgInfo.prompt + } +} + +/** + * Handler for MP3 files + */ +const handleMp3File: WorkflowFileHandler = async (file) => { + const { workflow, prompt } = await getMp3Metadata(file) + return { workflow, prompt } +} + +/** + * Handler for OGG files + */ +const handleOggFile: WorkflowFileHandler = async (file) => { + const { workflow, prompt } = await getOggMetadata(file) + return { workflow, prompt } +} + +/** + * Handler for FLAC files + */ +const handleFlacFile: WorkflowFileHandler = async (file) => { + const pngInfo = await getFlacMetadata(file) + const workflow = pngInfo?.workflow || pngInfo?.Workflow + const prompt = pngInfo?.prompt || pngInfo?.Prompt + + return { + workflow: workflow ? JSON.parse(workflow) : undefined, + prompt: prompt ? JSON.parse(prompt) : undefined + } +} + +/** + * Handler for WebM files + */ +const handleWebmFile: WorkflowFileHandler = async (file) => { + const webmInfo = await getFromWebmFile(file) + return { + workflow: webmInfo.workflow, + prompt: webmInfo.prompt + } +} + +/** + * Handler for MP4/MOV/M4V files + */ +const handleMp4File: WorkflowFileHandler = async (file) => { + const mp4Info = await getFromIsobmffFile(file) + return { + workflow: mp4Info.workflow, + prompt: mp4Info.prompt + } +} + +/** + * Handler for GLB files + */ +const handleGlbFile: WorkflowFileHandler = async (file) => { + const gltfInfo = await getGltfBinaryMetadata(file) + return { + workflow: gltfInfo.workflow, + prompt: gltfInfo.prompt + } +} + +/** + * Handler for JSON files + */ +const handleJsonFile: WorkflowFileHandler = async (file) => { + // For JSON files, we need to preserve the exact behavior from app.ts + // This code intentionally mirrors the original implementation + return new Promise((resolve, reject) => { + const reader = new FileReader() + reader.onload = () => { + try { + const readerResult = reader.result as string + const jsonContent = JSON.parse(readerResult) + + if (jsonContent?.templates) { + // This case will be handled separately in handleFile + resolve({ + workflow: undefined, + prompt: undefined, + jsonTemplateData: jsonContent + }) + } else if (isApiJson(jsonContent)) { + // API JSON format + resolve({ workflow: undefined, prompt: jsonContent }) + } else { + // Regular workflow JSON + resolve({ workflow: JSON.parse(readerResult), prompt: undefined }) + } + } catch (error) { + reject(error) + } + } + reader.onerror = () => reject(reader.error) + reader.readAsText(file) + }) +} + +/** + * Handler for .latent and .safetensors files + */ +const handleLatentFile: WorkflowFileHandler = async (file) => { + // Preserve the exact behavior from app.ts for latent files + const info = await getLatentMetadata(file) + + // Direct port of the original code, preserving behavior for TS compatibility + if (info && typeof info === 'object' && 'workflow' in info && info.workflow) { + return { + workflow: JSON.parse(info.workflow as string), + prompt: undefined + } + } else if ( + info && + typeof info === 'object' && + 'prompt' in info && + info.prompt + ) { + return { + workflow: undefined, + prompt: JSON.parse(info.prompt as string) + } + } else { + return { workflow: undefined, prompt: undefined } + } +} + +/** + * Helper function to determine if a JSON object is in the API JSON format + */ +function isApiJson(data: unknown) { + return ( + typeof data === 'object' && + data !== null && + Object.values(data as Record).every((v) => v.class_type) + ) +} + +// Register image format handlers +IMAGE_WORKFLOW_FORMATS.mimeTypes.forEach((mimeType) => { + if (mimeType === 'image/png') { + mimeTypeHandlers.set(mimeType, handlePngFile) + } else if (mimeType === 'image/webp') { + mimeTypeHandlers.set(mimeType, handleWebpFile) + } else if (mimeType === 'image/svg+xml') { + mimeTypeHandlers.set(mimeType, handleSvgFile) + } +}) + +IMAGE_WORKFLOW_FORMATS.extensions.forEach((ext) => { + if (ext === '.png') { + extensionHandlers.set(ext, handlePngFile) + } else if (ext === '.webp') { + extensionHandlers.set(ext, handleWebpFile) + } else if (ext === '.svg') { + extensionHandlers.set(ext, handleSvgFile) + } +}) + +// Register audio format handlers +AUDIO_WORKFLOW_FORMATS.mimeTypes.forEach((mimeType) => { + if (mimeType === 'audio/mpeg') { + mimeTypeHandlers.set(mimeType, handleMp3File) + } else if (mimeType === 'audio/ogg') { + mimeTypeHandlers.set(mimeType, handleOggFile) + } else if (mimeType === 'audio/flac' || mimeType === 'audio/x-flac') { + mimeTypeHandlers.set(mimeType, handleFlacFile) + } +}) + +AUDIO_WORKFLOW_FORMATS.extensions.forEach((ext) => { + if (ext === '.mp3') { + extensionHandlers.set(ext, handleMp3File) + } else if (ext === '.ogg') { + extensionHandlers.set(ext, handleOggFile) + } else if (ext === '.flac') { + extensionHandlers.set(ext, handleFlacFile) + } +}) + +// Register video format handlers +VIDEO_WORKFLOW_FORMATS.mimeTypes.forEach((mimeType) => { + if (mimeType === 'video/webm') { + mimeTypeHandlers.set(mimeType, handleWebmFile) + } else if ( + mimeType === 'video/mp4' || + mimeType === 'video/quicktime' || + mimeType === 'video/x-m4v' + ) { + mimeTypeHandlers.set(mimeType, handleMp4File) + } +}) + +VIDEO_WORKFLOW_FORMATS.extensions.forEach((ext) => { + if (ext === '.webm') { + extensionHandlers.set(ext, handleWebmFile) + } else if (ext === '.mp4' || ext === '.mov' || ext === '.m4v') { + extensionHandlers.set(ext, handleMp4File) + } +}) + +// Register 3D model format handlers +MODEL_WORKFLOW_FORMATS.mimeTypes.forEach((mimeType) => { + if (mimeType === 'model/gltf-binary') { + mimeTypeHandlers.set(mimeType, handleGlbFile) + } +}) + +MODEL_WORKFLOW_FORMATS.extensions.forEach((ext) => { + if (ext === '.glb') { + extensionHandlers.set(ext, handleGlbFile) + } +}) + +// Register data format handlers +DATA_WORKFLOW_FORMATS.mimeTypes.forEach((mimeType) => { + if (mimeType === 'application/json') { + mimeTypeHandlers.set(mimeType, handleJsonFile) + } +}) + +DATA_WORKFLOW_FORMATS.extensions.forEach((ext) => { + if (ext === '.json') { + extensionHandlers.set(ext, handleJsonFile) + } else if (ext === '.latent' || ext === '.safetensors') { + extensionHandlers.set(ext, handleLatentFile) + } +}) + +/** + * Gets the appropriate file handler for a given file based on mime type or extension + */ +export function getFileHandler(file: File): WorkflowFileHandler | null { + // First try to match by MIME type + if (file.type && mimeTypeHandlers.has(file.type)) { + return mimeTypeHandlers.get(file.type) || null + } + + // If no MIME type match, try to match by file extension + if (file.name) { + const extension = '.' + file.name.split('.').pop()?.toLowerCase() + if (extension && extensionHandlers.has(extension)) { + return extensionHandlers.get(extension) || null + } + } + + return null +} diff --git a/tests-ui/tests/utils/fileHandlers.test.ts b/tests-ui/tests/utils/fileHandlers.test.ts new file mode 100644 index 000000000..1d84fdc2a --- /dev/null +++ b/tests-ui/tests/utils/fileHandlers.test.ts @@ -0,0 +1,127 @@ +import { describe, expect, it } from 'vitest' + +import { + AUDIO_WORKFLOW_FORMATS, + DATA_WORKFLOW_FORMATS, + IMAGE_WORKFLOW_FORMATS, + MODEL_WORKFLOW_FORMATS, + VIDEO_WORKFLOW_FORMATS +} from '../../../src/constants/supportedWorkflowFormats' +import { + extensionHandlers, + getFileHandler, + mimeTypeHandlers +} from '../../../src/utils/fileHandlers' + +describe('fileHandlers', () => { + describe('handler registrations', () => { + it('should register handlers for all image MIME types', () => { + IMAGE_WORKFLOW_FORMATS.mimeTypes.forEach((mimeType) => { + expect(mimeTypeHandlers.has(mimeType)).toBe(true) + expect(mimeTypeHandlers.get(mimeType)).toBeTypeOf('function') + }) + }) + + it('should register handlers for all image extensions', () => { + IMAGE_WORKFLOW_FORMATS.extensions.forEach((ext) => { + expect(extensionHandlers.has(ext)).toBe(true) + expect(extensionHandlers.get(ext)).toBeTypeOf('function') + }) + }) + + it('should register handlers for all audio MIME types', () => { + AUDIO_WORKFLOW_FORMATS.mimeTypes.forEach((mimeType) => { + expect(mimeTypeHandlers.has(mimeType)).toBe(true) + expect(mimeTypeHandlers.get(mimeType)).toBeTypeOf('function') + }) + }) + + it('should register handlers for all audio extensions', () => { + AUDIO_WORKFLOW_FORMATS.extensions.forEach((ext) => { + expect(extensionHandlers.has(ext)).toBe(true) + expect(extensionHandlers.get(ext)).toBeTypeOf('function') + }) + }) + + it('should register handlers for all video MIME types', () => { + VIDEO_WORKFLOW_FORMATS.mimeTypes.forEach((mimeType) => { + expect(mimeTypeHandlers.has(mimeType)).toBe(true) + expect(mimeTypeHandlers.get(mimeType)).toBeTypeOf('function') + }) + }) + + it('should register handlers for all video extensions', () => { + VIDEO_WORKFLOW_FORMATS.extensions.forEach((ext) => { + expect(extensionHandlers.has(ext)).toBe(true) + expect(extensionHandlers.get(ext)).toBeTypeOf('function') + }) + }) + + it('should register handlers for all 3D model MIME types', () => { + MODEL_WORKFLOW_FORMATS.mimeTypes.forEach((mimeType) => { + expect(mimeTypeHandlers.has(mimeType)).toBe(true) + expect(mimeTypeHandlers.get(mimeType)).toBeTypeOf('function') + }) + }) + + it('should register handlers for all 3D model extensions', () => { + MODEL_WORKFLOW_FORMATS.extensions.forEach((ext) => { + expect(extensionHandlers.has(ext)).toBe(true) + expect(extensionHandlers.get(ext)).toBeTypeOf('function') + }) + }) + + it('should register handlers for all data MIME types', () => { + DATA_WORKFLOW_FORMATS.mimeTypes.forEach((mimeType) => { + expect(mimeTypeHandlers.has(mimeType)).toBe(true) + expect(mimeTypeHandlers.get(mimeType)).toBeTypeOf('function') + }) + }) + + it('should register handlers for all data extensions', () => { + DATA_WORKFLOW_FORMATS.extensions.forEach((ext) => { + expect(extensionHandlers.has(ext)).toBe(true) + expect(extensionHandlers.get(ext)).toBeTypeOf('function') + }) + }) + }) + + describe('getFileHandler', () => { + it('should return a handler when a file with a recognized MIME type is provided', () => { + const file = new File(['test content'], 'test.png', { type: 'image/png' }) + const handler = getFileHandler(file) + expect(handler).not.toBeNull() + expect(handler).toBeTypeOf('function') + }) + + it('should return a handler when a file with a recognized extension but no MIME type is provided', () => { + // File with empty MIME type but recognizable extension + const file = new File(['test content'], 'test.webp', { type: '' }) + const handler = getFileHandler(file) + expect(handler).not.toBeNull() + expect(handler).toBeTypeOf('function') + }) + + it('should return null when no handler is found for the file type', () => { + const file = new File(['test content'], 'test.txt', { + type: 'text/plain' + }) + const handler = getFileHandler(file) + expect(handler).toBeNull() + }) + + it('should prioritize MIME type over extension when both are present and different', () => { + // A file with a JSON MIME type but SVG extension + const file = new File(['{}'], 'test.svg', { type: 'application/json' }) + const handler = getFileHandler(file) + + // Make a shadow copy of the handlers for comparison + const jsonHandler = mimeTypeHandlers.get('application/json') + const svgHandler = extensionHandlers.get('.svg') + + // The handler should match the JSON handler, not the SVG handler + expect(handler).toBe(jsonHandler) + expect(handler).not.toBe(svgHandler) + }) + }) +})