diff --git a/browser_tests/assets/workflowInMedia/workflow_itxt.png b/browser_tests/assets/workflowInMedia/workflow_itxt.png new file mode 100644 index 000000000..95ce757bc Binary files /dev/null and b/browser_tests/assets/workflowInMedia/workflow_itxt.png differ diff --git a/browser_tests/tests/loadWorkflowInMedia.spec.ts b/browser_tests/tests/loadWorkflowInMedia.spec.ts index 754f373a7..1b426c413 100644 --- a/browser_tests/tests/loadWorkflowInMedia.spec.ts +++ b/browser_tests/tests/loadWorkflowInMedia.spec.ts @@ -16,6 +16,7 @@ test.describe( 'no_workflow.webp', 'large_workflow.webp', 'workflow_prompt_parameters.png', + 'workflow_itxt.png', 'workflow.webm', // Skipped due to 3d widget unstable visual result. // 3d widget shows grid after fully loaded. diff --git a/browser_tests/tests/loadWorkflowInMedia.spec.ts-snapshots/workflow-itxt-png-chromium-linux.png b/browser_tests/tests/loadWorkflowInMedia.spec.ts-snapshots/workflow-itxt-png-chromium-linux.png new file mode 100644 index 000000000..2d007b4e2 Binary files /dev/null and b/browser_tests/tests/loadWorkflowInMedia.spec.ts-snapshots/workflow-itxt-png-chromium-linux.png differ diff --git a/src/scripts/metadata/png.test.ts b/src/scripts/metadata/png.test.ts new file mode 100644 index 000000000..dcc47bf04 --- /dev/null +++ b/src/scripts/metadata/png.test.ts @@ -0,0 +1,245 @@ +import { describe, expect, it } from 'vitest' + +import { getFromPngBuffer } from './png' + +function createPngWithChunk( + chunkType: string, + keyword: string, + content: string, + options: { + compressionFlag?: number + compressionMethod?: number + languageTag?: string + translatedKeyword?: string + } = {} +): ArrayBuffer { + const { + compressionFlag = 0, + compressionMethod = 0, + languageTag = '', + translatedKeyword = '' + } = options + + const signature = new Uint8Array([ + 0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a + ]) + const typeBytes = new TextEncoder().encode(chunkType) + const keywordBytes = new TextEncoder().encode(keyword) + const contentBytes = new TextEncoder().encode(content) + + let chunkData: Uint8Array + if (chunkType === 'iTXt') { + const langBytes = new TextEncoder().encode(languageTag) + const transBytes = new TextEncoder().encode(translatedKeyword) + const totalLength = + keywordBytes.length + + 1 + + 2 + + langBytes.length + + 1 + + transBytes.length + + 1 + + contentBytes.length + + chunkData = new Uint8Array(totalLength) + let pos = 0 + chunkData.set(keywordBytes, pos) + pos += keywordBytes.length + chunkData[pos++] = 0 + chunkData[pos++] = compressionFlag + chunkData[pos++] = compressionMethod + chunkData.set(langBytes, pos) + pos += langBytes.length + chunkData[pos++] = 0 + chunkData.set(transBytes, pos) + pos += transBytes.length + chunkData[pos++] = 0 + chunkData.set(contentBytes, pos) + } else { + chunkData = new Uint8Array(keywordBytes.length + 1 + contentBytes.length) + chunkData.set(keywordBytes, 0) + chunkData[keywordBytes.length] = 0 + chunkData.set(contentBytes, keywordBytes.length + 1) + } + + const lengthBytes = new Uint8Array(4) + new DataView(lengthBytes.buffer).setUint32(0, chunkData.length, false) + + const crc = new Uint8Array(4) + + const iendType = new TextEncoder().encode('IEND') + const iendLength = new Uint8Array(4) + const iendCrc = new Uint8Array(4) + + const total = signature.length + 4 + 4 + chunkData.length + 4 + 4 + 4 + 0 + 4 + const result = new Uint8Array(total) + + let offset = 0 + result.set(signature, offset) + offset += signature.length + + result.set(lengthBytes, offset) + offset += 4 + result.set(typeBytes, offset) + offset += 4 + result.set(chunkData, offset) + offset += chunkData.length + result.set(crc, offset) + offset += 4 + + result.set(iendLength, offset) + offset += 4 + result.set(iendType, offset) + offset += 4 + result.set(iendCrc, offset) + + return result.buffer +} + +describe('getFromPngBuffer', () => { + it('returns empty object for invalid PNG', async () => { + const invalidData = new ArrayBuffer(8) + const result = await getFromPngBuffer(invalidData) + expect(result).toEqual({}) + }) + + it('parses tEXt chunk', async () => { + const workflow = '{"nodes":[]}' + const buffer = createPngWithChunk('tEXt', 'workflow', workflow) + const result = await getFromPngBuffer(buffer) + expect(result['workflow']).toBe(workflow) + }) + + it('parses comf chunk', async () => { + const prompt = '{"1":{"class_type":"Test"}}' + const buffer = createPngWithChunk('comf', 'prompt', prompt) + const result = await getFromPngBuffer(buffer) + expect(result['prompt']).toBe(prompt) + }) + + it('parses uncompressed iTXt chunk', async () => { + const workflow = '{"nodes":[{"id":1}]}' + const buffer = createPngWithChunk('iTXt', 'workflow', workflow, { + compressionFlag: 0, + compressionMethod: 0 + }) + const result = await getFromPngBuffer(buffer) + expect(result['workflow']).toBe(workflow) + }) + + it('parses iTXt chunk with language tag and translated keyword', async () => { + const workflow = '{"test":"value"}' + const buffer = createPngWithChunk('iTXt', 'workflow', workflow, { + compressionFlag: 0, + languageTag: 'en', + translatedKeyword: 'Workflow' + }) + const result = await getFromPngBuffer(buffer) + expect(result['workflow']).toBe(workflow) + }) + + it('parses compressed iTXt chunk', async () => { + const workflow = '{"nodes":[{"id":1,"type":"KSampler"}]}' + const contentBytes = new TextEncoder().encode(workflow) + + const stream = new CompressionStream('deflate') + const writer = stream.writable.getWriter() + await writer.write(contentBytes) + await writer.close() + + const reader = stream.readable.getReader() + const chunks: Uint8Array[] = [] + while (true) { + const { done, value } = await reader.read() + if (done) break + chunks.push(value) + } + const compressedBytes = new Uint8Array( + chunks.reduce((acc, c) => acc + c.length, 0) + ) + let pos = 0 + for (const chunk of chunks) { + compressedBytes.set(chunk, pos) + pos += chunk.length + } + + const buffer = createPngWithCompressedITXt( + 'workflow', + compressedBytes, + '', + '' + ) + const result = await getFromPngBuffer(buffer) + expect(result['workflow']).toBe(workflow) + }) +}) + +function createPngWithCompressedITXt( + keyword: string, + compressedContent: Uint8Array, + languageTag: string, + translatedKeyword: string +): ArrayBuffer { + const signature = new Uint8Array([ + 0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a + ]) + const typeBytes = new TextEncoder().encode('iTXt') + const keywordBytes = new TextEncoder().encode(keyword) + const langBytes = new TextEncoder().encode(languageTag) + const transBytes = new TextEncoder().encode(translatedKeyword) + + const totalLength = + keywordBytes.length + + 1 + + 2 + + langBytes.length + + 1 + + transBytes.length + + 1 + + compressedContent.length + + const chunkData = new Uint8Array(totalLength) + let pos = 0 + chunkData.set(keywordBytes, pos) + pos += keywordBytes.length + chunkData[pos++] = 0 + chunkData[pos++] = 1 + chunkData[pos++] = 0 + chunkData.set(langBytes, pos) + pos += langBytes.length + chunkData[pos++] = 0 + chunkData.set(transBytes, pos) + pos += transBytes.length + chunkData[pos++] = 0 + chunkData.set(compressedContent, pos) + + const lengthBytes = new Uint8Array(4) + new DataView(lengthBytes.buffer).setUint32(0, chunkData.length, false) + + const crc = new Uint8Array(4) + const iendType = new TextEncoder().encode('IEND') + const iendLength = new Uint8Array(4) + const iendCrc = new Uint8Array(4) + + const total = signature.length + 4 + 4 + chunkData.length + 4 + 4 + 4 + 0 + 4 + const result = new Uint8Array(total) + + let offset = 0 + result.set(signature, offset) + offset += signature.length + result.set(lengthBytes, offset) + offset += 4 + result.set(typeBytes, offset) + offset += 4 + result.set(chunkData, offset) + offset += chunkData.length + result.set(crc, offset) + offset += 4 + result.set(iendLength, offset) + offset += 4 + result.set(iendType, offset) + offset += 4 + result.set(iendCrc, offset) + + return result.buffer +} diff --git a/src/scripts/metadata/png.ts b/src/scripts/metadata/png.ts index 2e911796a..526d9b2d7 100644 --- a/src/scripts/metadata/png.ts +++ b/src/scripts/metadata/png.ts @@ -1,26 +1,59 @@ +async function decompressZlib( + data: Uint8Array +): Promise> { + const stream = new DecompressionStream('deflate') + const writer = stream.writable.getWriter() + try { + await writer.write(data) + await writer.close() + } finally { + writer.releaseLock() + } + + const reader = stream.readable.getReader() + const chunks: Uint8Array[] = [] + let totalLength = 0 + + try { + while (true) { + const { done, value } = await reader.read() + if (done) break + chunks.push(value) + totalLength += value.length + } + } finally { + reader.releaseLock() + } + + const result = new Uint8Array(totalLength) + let offset = 0 + for (const chunk of chunks) { + result.set(chunk, offset) + offset += chunk.length + } + return result +} + /** @knipIgnoreUnusedButUsedByCustomNodes */ -export function getFromPngBuffer(buffer: ArrayBuffer): Record { - // Get the PNG data as a Uint8Array +export async function getFromPngBuffer( + buffer: ArrayBuffer +): Promise> { const pngData = new Uint8Array(buffer) const dataView = new DataView(pngData.buffer) - // Check that the PNG signature is present if (dataView.getUint32(0) !== 0x89504e47) { console.error('Not a valid PNG file') return {} } - // Start searching for chunks after the PNG signature let offset = 8 - let txt_chunks: Record = {} - // Loop through the chunks in the PNG file + const txt_chunks: Record = {} + while (offset < pngData.length) { - // Get the length of the chunk const length = dataView.getUint32(offset) - // Get the chunk type const type = String.fromCharCode(...pngData.slice(offset + 4, offset + 8)) - if (type === 'tEXt' || type == 'comf' || type === 'iTXt') { - // Get the keyword + + if (type === 'tEXt' || type === 'comf' || type === 'iTXt') { let keyword_end = offset + 8 while (pngData[keyword_end] !== 0) { keyword_end++ @@ -28,11 +61,48 @@ export function getFromPngBuffer(buffer: ArrayBuffer): Record { const keyword = String.fromCharCode( ...pngData.slice(offset + 8, keyword_end) ) - // Get the text - const contentArraySegment = pngData.slice( - keyword_end + 1, - offset + 8 + length - ) + + let textStart = keyword_end + 1 + let isCompressed = false + let compressionMethod = 0 + + if (type === 'iTXt') { + const chunkEnd = offset + 8 + length + isCompressed = pngData[textStart] === 1 + compressionMethod = pngData[textStart + 1] + textStart += 2 + + while (pngData[textStart] !== 0 && textStart < chunkEnd) { + textStart++ + } + if (textStart < chunkEnd) textStart++ + + while (pngData[textStart] !== 0 && textStart < chunkEnd) { + textStart++ + } + if (textStart < chunkEnd) textStart++ + } + + let contentArraySegment = pngData.slice(textStart, offset + 8 + length) + + if (isCompressed) { + if (compressionMethod === 0) { + try { + contentArraySegment = await decompressZlib(contentArraySegment) + } catch (e) { + console.error(`Failed to decompress iTXt chunk "${keyword}":`, e) + offset += 12 + length + continue + } + } else { + console.warn( + `Unsupported compression method ${compressionMethod} for iTXt chunk "${keyword}"` + ) + offset += 12 + length + continue + } + } + const contentJson = new TextDecoder('utf-8').decode(contentArraySegment) txt_chunks[keyword] = contentJson } @@ -42,14 +112,21 @@ export function getFromPngBuffer(buffer: ArrayBuffer): Record { return txt_chunks } -export function getFromPngFile(file: File) { - return new Promise>((r) => { +export async function getFromPngFile( + file: File +): Promise> { + return new Promise>((resolve, reject) => { const reader = new FileReader() - reader.onload = (event) => { - // @ts-expect-error fixme ts strict error - r(getFromPngBuffer(event.target.result as ArrayBuffer)) + reader.onload = async (event) => { + const buffer = event.target?.result + if (!(buffer instanceof ArrayBuffer)) { + reject(new Error('Failed to read file as ArrayBuffer')) + return + } + const result = await getFromPngBuffer(buffer) + resolve(result) } - + reader.onerror = () => reject(reader.error) reader.readAsArrayBuffer(file) }) }