diff --git a/browser_tests/assets/workflow.avif b/browser_tests/assets/workflow.avif new file mode 100644 index 000000000..8d8d77ced Binary files /dev/null and b/browser_tests/assets/workflow.avif differ diff --git a/browser_tests/tests/loadWorkflowInMedia.spec.ts b/browser_tests/tests/loadWorkflowInMedia.spec.ts index 215f977e7..1250d3385 100644 --- a/browser_tests/tests/loadWorkflowInMedia.spec.ts +++ b/browser_tests/tests/loadWorkflowInMedia.spec.ts @@ -15,7 +15,8 @@ test.describe('Load Workflow in Media', () => { 'workflow.mp4', 'workflow.mov', 'workflow.m4v', - 'workflow.svg' + 'workflow.svg', + 'workflow.avif' ] fileNames.forEach(async (fileName) => { test(`Load workflow in ${fileName} (drop from filesystem)`, async ({ diff --git a/src/constants/supportedWorkflowFormats.ts b/src/constants/supportedWorkflowFormats.ts index f00bd3f05..ac0bee186 100644 --- a/src/constants/supportedWorkflowFormats.ts +++ b/src/constants/supportedWorkflowFormats.ts @@ -6,8 +6,8 @@ * All supported image formats that can contain workflow data */ export const IMAGE_WORKFLOW_FORMATS = { - extensions: ['.png', '.webp', '.svg'], - mimeTypes: ['image/png', 'image/webp', 'image/svg+xml'] + extensions: ['.png', '.webp', '.svg', '.avif'], + mimeTypes: ['image/png', 'image/webp', 'image/svg+xml', 'image/avif'] } /** diff --git a/src/scripts/app.ts b/src/scripts/app.ts index 2061a5f6a..257a30cf6 100644 --- a/src/scripts/app.ts +++ b/src/scripts/app.ts @@ -80,6 +80,7 @@ import { deserialiseAndCreate } from '@/utils/vintageClipboard' import { type ComfyApi, PromptExecutionError, api } from './api' import { defaultGraph } from './defaultGraph' import { + getAvifMetadata, getFlacMetadata, getLatentMetadata, getPngMetadata, @@ -1351,6 +1352,16 @@ export class ComfyApp { } else { this.showErrorOnFileLoad(file) } + } else if (file.type === 'image/avif') { + const { workflow, prompt } = await getAvifMetadata(file) + + 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 === 'image/webp') { const pngInfo = await getWebpMetadata(file) // Support loading workflows from that webp custom node. diff --git a/src/scripts/metadata/avif.ts b/src/scripts/metadata/avif.ts new file mode 100644 index 000000000..d23adf4bb --- /dev/null +++ b/src/scripts/metadata/avif.ts @@ -0,0 +1,412 @@ +import { + type AvifIinfBox, + type AvifIlocBox, + type AvifInfeBox, + ComfyMetadata, + ComfyMetadataTags, + type IsobmffBoxContentRange +} from '@/types/metadataTypes' + +const readNullTerminatedString = ( + dataView: DataView, + start: number, + end: number +): { str: string; length: number } => { + let length = 0 + while (start + length < end && dataView.getUint8(start + length) !== 0) { + length++ + } + const str = new TextDecoder('utf-8').decode( + new Uint8Array(dataView.buffer, dataView.byteOffset + start, length) + ) + return { str, length: length + 1 } // Include null terminator +} + +const parseInfeBox = (dataView: DataView, start: number): AvifInfeBox => { + const version = dataView.getUint8(start) + const flags = dataView.getUint32(start) & 0xffffff + let offset = start + 4 + + let item_ID: number, item_protection_index: number, item_type: string + + if (version >= 2) { + if (version === 2) { + item_ID = dataView.getUint16(offset) + offset += 2 + } else { + item_ID = dataView.getUint32(offset) + offset += 4 + } + + item_protection_index = dataView.getUint16(offset) + offset += 2 + item_type = String.fromCharCode( + ...new Uint8Array(dataView.buffer, dataView.byteOffset + offset, 4) + ) + offset += 4 + + const { str: item_name, length: name_len } = readNullTerminatedString( + dataView, + offset, + dataView.byteLength + ) + offset += name_len + + const content_type = readNullTerminatedString( + dataView, + offset, + dataView.byteLength + ).str + + return { + box_header: { size: 0, type: 'infe' }, // Size is dynamic + version, + flags, + item_ID, + item_protection_index, + item_type, + item_name, + content_type + } + } + throw new Error(`Unsupported infe box version: ${version}`) +} + +const parseIinfBox = ( + dataView: DataView, + range: IsobmffBoxContentRange +): AvifIinfBox => { + if (!range) throw new Error('iinf box not found') + + const version = dataView.getUint8(range.start) + const flags = dataView.getUint32(range.start) & 0xffffff + let offset = range.start + 4 + + const entry_count = + version === 0 ? dataView.getUint16(offset) : dataView.getUint32(offset) + offset += version === 0 ? 2 : 4 + + const entries: AvifInfeBox[] = [] + for (let i = 0; i < entry_count; i++) { + const boxSize = dataView.getUint32(offset) + const boxType = String.fromCharCode( + ...new Uint8Array(dataView.buffer, dataView.byteOffset + offset + 4, 4) + ) + + if (boxType === 'infe') { + const infe = parseInfeBox(dataView, offset + 8) + infe.box_header.size = boxSize + entries.push(infe) + } + offset += boxSize + } + + return { + box_header: { size: range.end - range.start + 8, type: 'iinf' }, + version, + flags, + entry_count, + entries + } +} + +const parseIlocBox = ( + dataView: DataView, + range: IsobmffBoxContentRange +): AvifIlocBox => { + if (!range) throw new Error('iloc box not found') + + const version = dataView.getUint8(range.start) + const flags = dataView.getUint32(range.start) & 0xffffff + let offset = range.start + 4 + + const sizes = dataView.getUint8(offset++) + const offset_size = (sizes >> 4) & 0x0f + const length_size = sizes & 0x0f + + const base_offset_size = (dataView.getUint8(offset) >> 4) & 0x0f + const index_size = + version === 1 || version === 2 ? dataView.getUint8(offset) & 0x0f : 0 + offset++ + + const item_count = + version < 2 ? dataView.getUint16(offset) : dataView.getUint32(offset) + offset += version < 2 ? 2 : 4 + + const items = [] + for (let i = 0; i < item_count; i++) { + const item_ID = + version < 2 ? dataView.getUint16(offset) : dataView.getUint32(offset) + offset += version < 2 ? 2 : 4 + + if (version === 1 || version === 2) { + offset += 2 // construction_method + } + + const data_reference_index = dataView.getUint16(offset) + offset += 2 + + const base_offset = base_offset_size > 0 ? dataView.getUint32(offset) : 0 // Simplified + offset += base_offset_size + + const extent_count = dataView.getUint16(offset) + offset += 2 + + const extents = [] + for (let j = 0; j < extent_count; j++) { + if ((version === 1 || version === 2) && index_size > 0) { + offset += index_size + } + const extent_offset = dataView.getUint32(offset) // Simplified + offset += offset_size + const extent_length = dataView.getUint32(offset) // Simplified + offset += length_size + extents.push({ extent_offset, extent_length }) + } + items.push({ + item_ID, + data_reference_index, + base_offset, + extent_count, + extents + }) + } + + return { + box_header: { size: range.end - range.start + 8, type: 'iloc' }, + version, + flags, + offset_size, + length_size, + base_offset_size, + index_size, + item_count, + items + } +} + +function findBox( + dataView: DataView, + start: number, + end: number, + type: string +): IsobmffBoxContentRange { + let offset = start + while (offset < end) { + if (offset + 8 > end) break + const boxLength = dataView.getUint32(offset) + const boxType = String.fromCharCode( + ...new Uint8Array(dataView.buffer, dataView.byteOffset + offset + 4, 4) + ) + + if (boxLength === 0) break + + if (boxType === type) { + return { start: offset + 8, end: offset + boxLength } + } + if (offset + boxLength > end) break + offset += boxLength + } + return null +} + +function parseAvifMetadata(buffer: ArrayBuffer): ComfyMetadata { + const metadata: ComfyMetadata = {} + const dataView = new DataView(buffer) + + if ( + dataView.getUint32(4) !== 0x66747970 || + dataView.getUint32(8) !== 0x61766966 + ) { + console.error('Not a valid AVIF file') + return {} + } + + const metaBox = findBox(dataView, 0, dataView.byteLength, 'meta') + if (!metaBox) return {} + + const metaBoxContentStart = metaBox.start + 4 // Skip version and flags + + const iinfBoxRange = findBox( + dataView, + metaBoxContentStart, + metaBox.end, + 'iinf' + ) + const iinf = parseIinfBox(dataView, iinfBoxRange) + + const exifInfe = iinf.entries.find((e) => e.item_type === 'Exif') + if (!exifInfe) return {} + + const ilocBoxRange = findBox( + dataView, + metaBoxContentStart, + metaBox.end, + 'iloc' + ) + const iloc = parseIlocBox(dataView, ilocBoxRange) + + const exifIloc = iloc.items.find((i) => i.item_ID === exifInfe.item_ID) + if (!exifIloc || exifIloc.extents.length === 0) return {} + + const exifExtent = exifIloc.extents[0] + const itemData = new Uint8Array( + buffer, + exifExtent.extent_offset, + exifExtent.extent_length + ) + + let tiffHeaderOffset = -1 + for (let i = 0; i < itemData.length - 4; i++) { + if ( + (itemData[i] === 0x4d && + itemData[i + 1] === 0x4d && + itemData[i + 2] === 0x00 && + itemData[i + 3] === 0x2a) || // MM* + (itemData[i] === 0x49 && + itemData[i + 1] === 0x49 && + itemData[i + 2] === 0x2a && + itemData[i + 3] === 0x00) // II* + ) { + tiffHeaderOffset = i + break + } + } + + if (tiffHeaderOffset !== -1) { + const exifData = itemData.subarray(tiffHeaderOffset) + const data: Record = parseExifData(exifData) + for (const key in data) { + const value = data[key] + if (typeof value === 'string') { + if (key === 'usercomment') { + try { + const metadataJson = JSON.parse(value) + if (metadataJson.prompt) { + metadata[ComfyMetadataTags.PROMPT] = metadataJson.prompt + } + if (metadataJson.workflow) { + metadata[ComfyMetadataTags.WORKFLOW] = metadataJson.workflow + } + } catch (e) { + console.error('Failed to parse usercomment JSON', e) + } + } else { + const [metadataKey, ...metadataValueParts] = value.split(':') + const metadataValue = metadataValueParts.join(':').trim() + if ( + metadataKey.toLowerCase() === + ComfyMetadataTags.PROMPT.toLowerCase() || + metadataKey.toLowerCase() === + ComfyMetadataTags.WORKFLOW.toLowerCase() + ) { + try { + const jsonValue = JSON.parse(metadataValue) + metadata[metadataKey.toLowerCase() as keyof ComfyMetadata] = + jsonValue + } catch (e) { + console.error(`Failed to parse JSON for ${metadataKey}`, e) + } + } + } + } + } + } else { + console.log('Warning: TIFF header not found in EXIF data.') + } + + return metadata +} + +// @ts-expect-error fixme ts strict error +export function parseExifData(exifData) { + // Check for the correct TIFF header (0x4949 for little-endian or 0x4D4D for big-endian) + const isLittleEndian = String.fromCharCode(...exifData.slice(0, 2)) === 'II' + + // Function to read 16-bit and 32-bit integers from binary data + // @ts-expect-error fixme ts strict error + function readInt(offset, isLittleEndian, length) { + let arr = exifData.slice(offset, offset + length) + if (length === 2) { + return new DataView(arr.buffer, arr.byteOffset, arr.byteLength).getUint16( + 0, + isLittleEndian + ) + } else if (length === 4) { + return new DataView(arr.buffer, arr.byteOffset, arr.byteLength).getUint32( + 0, + isLittleEndian + ) + } + } + + // Read the offset to the first IFD (Image File Directory) + const ifdOffset = readInt(4, isLittleEndian, 4) + + // @ts-expect-error fixme ts strict error + function parseIFD(offset) { + const numEntries = readInt(offset, isLittleEndian, 2) + const result = {} + + // @ts-expect-error fixme ts strict error + for (let i = 0; i < numEntries; i++) { + const entryOffset = offset + 2 + i * 12 + const tag = readInt(entryOffset, isLittleEndian, 2) + const type = readInt(entryOffset + 2, isLittleEndian, 2) + const numValues = readInt(entryOffset + 4, isLittleEndian, 4) + const valueOffset = readInt(entryOffset + 8, isLittleEndian, 4) + + // Read the value(s) based on the data type + let value + if (type === 2) { + // ASCII string + value = new TextDecoder('utf-8').decode( + // @ts-expect-error fixme ts strict error + exifData.subarray(valueOffset, valueOffset + numValues - 1) + ) + } + + // @ts-expect-error fixme ts strict error + result[tag] = value + } + + return result + } + + // Parse the first IFD + const ifdData = parseIFD(ifdOffset) + return ifdData +} + +export function getFromAvifFile(file: File): Promise> { + return new Promise>((resolve) => { + const reader = new FileReader() + reader.onload = (event) => { + const buffer = event.target?.result as ArrayBuffer + if (!buffer) { + resolve({}) + return + } + + try { + const comfyMetadata = parseAvifMetadata(buffer) + const result: Record = {} + if (comfyMetadata.prompt) { + result.prompt = JSON.stringify(comfyMetadata.prompt) + } + if (comfyMetadata.workflow) { + result.workflow = JSON.stringify(comfyMetadata.workflow) + } + resolve(result) + } catch (e) { + console.error('Parser: Error parsing AVIF metadata:', e) + resolve({}) + } + } + reader.onerror = (err) => { + console.error('FileReader: Error reading AVIF file:', err) + resolve({}) + } + reader.readAsArrayBuffer(file) + }) +} diff --git a/src/scripts/pnginfo.ts b/src/scripts/pnginfo.ts index 7c7f6f597..689537f7e 100644 --- a/src/scripts/pnginfo.ts +++ b/src/scripts/pnginfo.ts @@ -1,6 +1,7 @@ import { LiteGraph } from '@comfyorg/litegraph' import { api } from './api' +import { getFromAvifFile } from './metadata/avif' import { getFromFlacFile } from './metadata/flac' import { getFromPngFile } from './metadata/png' @@ -13,6 +14,10 @@ export function getFlacMetadata(file: File): Promise> { return getFromFlacFile(file) } +export function getAvifMetadata(file: File): Promise> { + return getFromAvifFile(file) +} + // @ts-expect-error fixme ts strict error function parseExifData(exifData) { // Check for the correct TIFF header (0x4949 for little-endian or 0x4D4D for big-endian) diff --git a/src/types/metadataTypes.ts b/src/types/metadataTypes.ts index 9170ef6d4..c63324196 100644 --- a/src/types/metadataTypes.ts +++ b/src/types/metadataTypes.ts @@ -85,3 +85,57 @@ export type GltfJsonData = { * Null if the box was not found. */ export type IsobmffBoxContentRange = { start: number; end: number } | null + +export type AvifInfeBox = { + box_header: { + size: number + type: 'infe' + } + version: number + flags: number + item_ID: number + item_protection_index: number + item_type: string + item_name: string + content_type?: string + content_encoding?: string +} + +export type AvifIinfBox = { + box_header: { + size: number + type: 'iinf' + } + version: number + flags: number + entry_count: number + entries: AvifInfeBox[] +} + +export type AvifIlocItemExtent = { + extent_offset: number + extent_length: number +} + +export type AvifIlocItem = { + item_ID: number + data_reference_index: number + base_offset: number + extent_count: number + extents: AvifIlocItemExtent[] +} + +export type AvifIlocBox = { + box_header: { + size: number + type: 'iloc' + } + version: number + flags: number + offset_size: number + length_size: number + base_offset_size: number + index_size: number + item_count: number + items: AvifIlocItem[] +}