|
| 1 | +import { |
| 2 | + ComfyApiWorkflow, |
| 3 | + ComfyWorkflowJSON |
| 4 | +} from '@/schemas/comfyWorkflowSchema' |
| 5 | +import { |
| 6 | + ASCII, |
| 7 | + ComfyMetadata, |
| 8 | + ComfyMetadataTags, |
| 9 | + IsobmffBoxContentRange |
| 10 | +} from '@/types/metadataTypes' |
| 11 | + |
| 12 | +const MAX_READ_BYTES = 2 * 1024 * 1024 |
| 13 | +const BOX_TYPES = { |
| 14 | + USER_DATA: [0x75, 0x64, 0x74, 0x61], |
| 15 | + META_DATA: [0x6d, 0x65, 0x74, 0x61], |
| 16 | + ITEM_LIST: [0x69, 0x6c, 0x73, 0x74], |
| 17 | + KEYS: [0x6b, 0x65, 0x79, 0x73], |
| 18 | + DATA: [0x64, 0x61, 0x74, 0x61], |
| 19 | + MOVIE: [0x6d, 0x6f, 0x6f, 0x76] |
| 20 | +} |
| 21 | +const SIZES = { |
| 22 | + HEADER: 8, |
| 23 | + VERSION: 4, |
| 24 | + LOCALE: 4, |
| 25 | + ITEM_MIN: 8 |
| 26 | +} |
| 27 | + |
| 28 | +const bufferMatchesBoxType = ( |
| 29 | + data: Uint8Array, |
| 30 | + pos: number, |
| 31 | + boxType: number[] |
| 32 | +): boolean => { |
| 33 | + if (pos + 4 > data.length) return false |
| 34 | + |
| 35 | + for (let i = 0; i < 4; i++) { |
| 36 | + if (data[pos + i] !== boxType[i]) return false |
| 37 | + } |
| 38 | + return true |
| 39 | +} |
| 40 | + |
| 41 | +const readUint32 = (data: Uint8Array, pos: number): number => { |
| 42 | + if (pos + 4 > data.length) return 0 |
| 43 | + return ( |
| 44 | + (data[pos] << 24) | |
| 45 | + (data[pos + 1] << 16) | |
| 46 | + (data[pos + 2] << 8) | |
| 47 | + data[pos + 3] |
| 48 | + ) |
| 49 | +} |
| 50 | + |
| 51 | +const findIsobmffBoxByType = ( |
| 52 | + data: Uint8Array, |
| 53 | + startPos: number, |
| 54 | + endPos: number, |
| 55 | + boxType: number[] |
| 56 | +): IsobmffBoxContentRange => { |
| 57 | + for (let pos = startPos; pos < endPos - 8; pos++) { |
| 58 | + const size = readUint32(data, pos) |
| 59 | + if (size < SIZES.ITEM_MIN) continue // Minimum size is 8 bytes |
| 60 | + |
| 61 | + if (bufferMatchesBoxType(data, pos + 4, boxType)) |
| 62 | + return { start: pos + SIZES.HEADER, end: pos + size } // Skip header |
| 63 | + |
| 64 | + // If type doesn't match, ensure size is valid before skipping |
| 65 | + if (pos + size > endPos) return null |
| 66 | + |
| 67 | + pos += size - 1 // Skip to the next potential box start |
| 68 | + } |
| 69 | + return null |
| 70 | +} |
| 71 | + |
| 72 | +const extractJson = (data: Uint8Array, start: number, end: number): any => { |
| 73 | + let jsonStart = start |
| 74 | + while (jsonStart < end && data[jsonStart] !== ASCII.OPEN_BRACE) { |
| 75 | + jsonStart++ |
| 76 | + } |
| 77 | + if (jsonStart >= end) return null |
| 78 | + |
| 79 | + try { |
| 80 | + const jsonText = new TextDecoder().decode(data.slice(jsonStart, end)) |
| 81 | + return JSON.parse(jsonText) |
| 82 | + } catch { |
| 83 | + return null |
| 84 | + } |
| 85 | +} |
| 86 | + |
| 87 | +const readUtf8String = (data: Uint8Array, start: number, end: number): string => |
| 88 | + new TextDecoder().decode(data.slice(start, end)) |
| 89 | + |
| 90 | +const parseKeysBox = ( |
| 91 | + data: Uint8Array, |
| 92 | + keysBoxStart: number, |
| 93 | + keysBoxEnd: number |
| 94 | +): Map<number, string> => { |
| 95 | + const keysMap = new Map<number, string>() |
| 96 | + let pos = keysBoxStart + 4 // Skip version/flags |
| 97 | + if (pos + 4 > keysBoxEnd) return keysMap |
| 98 | + |
| 99 | + const entryCount = readUint32(data, pos) |
| 100 | + pos += 4 |
| 101 | + |
| 102 | + for (let i = 1; i <= entryCount; i++) { |
| 103 | + // Keys are 1-indexed |
| 104 | + if (pos + SIZES.HEADER > keysBoxEnd) break |
| 105 | + |
| 106 | + const keySize = readUint32(data, pos) |
| 107 | + pos += SIZES.HEADER |
| 108 | + |
| 109 | + const keyNameEnd = pos + keySize - SIZES.HEADER |
| 110 | + if (keySize < SIZES.ITEM_MIN || keyNameEnd > keysBoxEnd) break |
| 111 | + |
| 112 | + const keyName = readUtf8String(data, pos, keyNameEnd) |
| 113 | + keysMap.set(i, keyName) |
| 114 | + pos = keyNameEnd |
| 115 | + } |
| 116 | + return keysMap |
| 117 | +} |
| 118 | + |
| 119 | +const extractMetadataValueFromDataBox = ( |
| 120 | + data: Uint8Array, |
| 121 | + dataBoxStart: number, |
| 122 | + dataBoxEnd: number, |
| 123 | + keyName: string |
| 124 | +): ComfyWorkflowJSON | ComfyApiWorkflow | null => { |
| 125 | + const valueStart = dataBoxStart + SIZES.VERSION + SIZES.LOCALE |
| 126 | + if (valueStart >= dataBoxEnd) return null |
| 127 | + |
| 128 | + const lowerKeyName = keyName.toLowerCase() |
| 129 | + if ( |
| 130 | + lowerKeyName === ComfyMetadataTags.PROMPT.toLowerCase() || |
| 131 | + lowerKeyName === ComfyMetadataTags.WORKFLOW.toLowerCase() |
| 132 | + ) { |
| 133 | + return extractJson(data, valueStart, dataBoxEnd) || null |
| 134 | + } |
| 135 | + return null |
| 136 | +} |
| 137 | + |
| 138 | +const parseIlstItem = ( |
| 139 | + data: Uint8Array, |
| 140 | + itemStart: number, |
| 141 | + itemEnd: number, |
| 142 | + keysMap: Map<number, string>, |
| 143 | + metadata: ComfyMetadata |
| 144 | +) => { |
| 145 | + if (itemStart + SIZES.HEADER > itemEnd) return |
| 146 | + |
| 147 | + const itemIndex = readUint32(data, itemStart + 4) |
| 148 | + const keyName = keysMap.get(itemIndex) |
| 149 | + if (!keyName) return |
| 150 | + |
| 151 | + const dataBox = findIsobmffBoxByType( |
| 152 | + data, |
| 153 | + itemStart + SIZES.HEADER, |
| 154 | + itemEnd, |
| 155 | + BOX_TYPES.DATA |
| 156 | + ) |
| 157 | + if (dataBox) { |
| 158 | + const value = extractMetadataValueFromDataBox( |
| 159 | + data, |
| 160 | + dataBox.start, |
| 161 | + dataBox.end, |
| 162 | + keyName |
| 163 | + ) |
| 164 | + if (value !== null) { |
| 165 | + metadata[keyName.toLowerCase() as keyof ComfyMetadata] = value |
| 166 | + } |
| 167 | + } |
| 168 | +} |
| 169 | + |
| 170 | +const parseIlstBox = ( |
| 171 | + data: Uint8Array, |
| 172 | + ilstStart: number, |
| 173 | + ilstEnd: number, |
| 174 | + keysMap: Map<number, string>, |
| 175 | + metadata: ComfyMetadata |
| 176 | +) => { |
| 177 | + let pos = ilstStart |
| 178 | + while (pos < ilstEnd - SIZES.HEADER) { |
| 179 | + const itemSize = readUint32(data, pos) |
| 180 | + if (itemSize <= SIZES.HEADER || pos + itemSize > ilstEnd) break // Invalid item size |
| 181 | + parseIlstItem(data, pos, pos + itemSize, keysMap, metadata) |
| 182 | + pos += itemSize |
| 183 | + } |
| 184 | +} |
| 185 | + |
| 186 | +const findUserDataBox = (data: Uint8Array): IsobmffBoxContentRange => { |
| 187 | + let userDataBox: IsobmffBoxContentRange = null |
| 188 | + |
| 189 | + // Metadata can be in 'udta' at top level or inside 'moov' |
| 190 | + userDataBox = findIsobmffBoxByType(data, 0, data.length, BOX_TYPES.USER_DATA) |
| 191 | + |
| 192 | + if (!userDataBox) { |
| 193 | + const moovBox = findIsobmffBoxByType(data, 0, data.length, BOX_TYPES.MOVIE) |
| 194 | + if (moovBox) { |
| 195 | + userDataBox = findIsobmffBoxByType( |
| 196 | + data, |
| 197 | + moovBox.start, |
| 198 | + moovBox.end, |
| 199 | + BOX_TYPES.USER_DATA |
| 200 | + ) |
| 201 | + } |
| 202 | + } |
| 203 | + return userDataBox |
| 204 | +} |
| 205 | + |
| 206 | +const parseIsobmffMetadata = (data: Uint8Array): ComfyMetadata => { |
| 207 | + const metadata: ComfyMetadata = {} |
| 208 | + const userDataBox = findUserDataBox(data) |
| 209 | + if (!userDataBox) return metadata |
| 210 | + |
| 211 | + const metaBox = findIsobmffBoxByType( |
| 212 | + data, |
| 213 | + userDataBox.start, |
| 214 | + userDataBox.end, |
| 215 | + BOX_TYPES.META_DATA |
| 216 | + ) |
| 217 | + if (!metaBox) return metadata |
| 218 | + |
| 219 | + const metaContentStart = metaBox.start + SIZES.VERSION |
| 220 | + const keysBox = findIsobmffBoxByType( |
| 221 | + data, |
| 222 | + metaContentStart, |
| 223 | + metaBox.end, |
| 224 | + BOX_TYPES.KEYS |
| 225 | + ) |
| 226 | + if (!keysBox) return metadata |
| 227 | + |
| 228 | + const keysMap = parseKeysBox(data, keysBox.start, keysBox.end) |
| 229 | + if (keysMap.size === 0) return metadata // keys box is empty or failed to parse |
| 230 | + |
| 231 | + const ilstBox = findIsobmffBoxByType( |
| 232 | + data, |
| 233 | + metaContentStart, |
| 234 | + metaBox.end, |
| 235 | + BOX_TYPES.ITEM_LIST |
| 236 | + ) |
| 237 | + if (!ilstBox) return metadata |
| 238 | + |
| 239 | + parseIlstBox(data, ilstBox.start, ilstBox.end, keysMap, metadata) |
| 240 | + |
| 241 | + return metadata |
| 242 | +} |
| 243 | + |
| 244 | +/** |
| 245 | + * Extracts ComfyUI Workflow metadata from an ISO Base Media File Format (ISOBMFF) file |
| 246 | + * (e.g., MP4, MOV) by parsing the `udta.meta.keys` and `udta.meta.ilst` boxes. |
| 247 | + * @param file - The file to extract metadata from. |
| 248 | + */ |
| 249 | +export function getFromIsobmffFile(file: File): Promise<ComfyMetadata> { |
| 250 | + return new Promise<ComfyMetadata>((resolve) => { |
| 251 | + const reader = new FileReader() |
| 252 | + reader.onload = (event: ProgressEvent<FileReader>) => { |
| 253 | + if (!event.target?.result) { |
| 254 | + resolve({}) |
| 255 | + return |
| 256 | + } |
| 257 | + |
| 258 | + try { |
| 259 | + const data = new Uint8Array(event.target.result as ArrayBuffer) |
| 260 | + resolve(parseIsobmffMetadata(data)) |
| 261 | + } catch (e) { |
| 262 | + console.error('Parser: Error parsing ISOBMFF metadata:', e) |
| 263 | + resolve({}) |
| 264 | + } |
| 265 | + } |
| 266 | + reader.onerror = (err) => { |
| 267 | + console.error('FileReader: Error reading ISOBMFF file:', err) |
| 268 | + resolve({}) |
| 269 | + } |
| 270 | + reader.readAsArrayBuffer(file.slice(0, MAX_READ_BYTES)) |
| 271 | + }) |
| 272 | +} |
0 commit comments