diff --git a/app/utils/exif.ts b/app/utils/exif.ts index 17df180..d0348d4 100644 --- a/app/utils/exif.ts +++ b/app/utils/exif.ts @@ -36,6 +36,9 @@ export function getPngMetadata(buffer: ArrayBuffer): Record { offset + 8 + length ); const contentJson = new TextDecoder("utf-8").decode(contentArraySegment); + + if (txt_chunks[keyword]) + console.warn(`Duplicated keyword ${keyword} has been overwritten`); txt_chunks[keyword] = contentJson; } @@ -71,7 +74,7 @@ ref: png chunk struct: export function setPngMetadata( buffer: ArrayBuffer, new_txt_chunks: Record -): Uint8Array { +): Uint8Array { // Get the PNG data as a Uint8Array const pngData = new Uint8Array(buffer); const newPngChunks: Uint8Array[] = []; @@ -96,9 +99,7 @@ export function setPngMetadata( if (type === "tEXt" || type == "comf" || type === "iTXt") { // Get the keyword let keyword_end = offset + 8; - while (pngData[keyword_end] !== 0) { - keyword_end++; - } + while (pngData[keyword_end] !== 0) keyword_end++; const keyword = String.fromCharCode( ...pngData.slice(offset + 8, keyword_end) ); @@ -163,6 +164,7 @@ export function setPngMetadata( newPngChunk.set(encoded, 8); dataView.setUint32(8 + chunkLength, chunkCRC32); newPngChunks.push(newPngChunk); + delete new_txt_chunks[keyword]; //mark used } } else { // if this keyword is not in new_txt_chunks, @@ -177,7 +179,28 @@ export function setPngMetadata( offset += 12 + length; } - // Concatenate the new PNG chunks + // If no EXIF section was found, add new metadata chunks + Object.entries(new_txt_chunks).map(([keyword, content]) => { + // console.log(`Adding exif section for ${keyword}`); + const encoded = new TextEncoder().encode(keyword + "\x00" + content); + const chunkLength = encoded.length; + const chunkType = new TextEncoder().encode("tEXt"); + + // Calculate crc32 + const crcTarget = new Uint8Array(chunkType.length + encoded.length); + crcTarget.set(chunkType, 0); + crcTarget.set(encoded, chunkType.length); + const chunkCRC32 = crc32FromArrayBuffer(crcTarget); + + const newPngChunk = new Uint8Array(8 + chunkLength + 4); + const dataView = new DataView(newPngChunk.buffer); + dataView.setUint32(0, chunkLength); + newPngChunk.set(chunkType, 4); + newPngChunk.set(encoded, 8); + dataView.setUint32(8 + chunkLength, chunkCRC32); + newPngChunks.push(newPngChunk); + }); + const newPngData = concatUint8Arrays(newPngChunks); return newPngData; } diff --git a/tests/Blank_2025-03-06-input.png b/tests/Blank_2025-03-06-input.png new file mode 100644 index 0000000..f75a6e7 Binary files /dev/null and b/tests/Blank_2025-03-06-input.png differ diff --git a/tests/ComfyUI_temp_nrhrt_00001_.png b/tests/ComfyUI_00001.png similarity index 100% rename from tests/ComfyUI_temp_nrhrt_00001_.png rename to tests/ComfyUI_00001.png diff --git a/tests/ComfyUI_temp_nrhrt_00001_.png.workflow.json b/tests/ComfyUI_00001.png.workflow.json similarity index 100% rename from tests/ComfyUI_temp_nrhrt_00001_.png.workflow.json rename to tests/ComfyUI_00001.png.workflow.json diff --git a/tests/exif-png.test.ts b/tests/exif-png.test.ts index 378faeb..30de58b 100644 --- a/tests/exif-png.test.ts +++ b/tests/exif-png.test.ts @@ -2,7 +2,7 @@ import { getPngMetadata, setPngMetadata } from "@/app/utils/exif"; import { glob } from "glob"; it("extract png workflow", async () => { - const pngs = await glob("./tests/*.png"); + const pngs = await glob("./tests/ComfyUI_*.png"); expect(pngs.length).toBeGreaterThanOrEqual(1); for await (const filename of pngs) { @@ -18,7 +18,7 @@ it("extract png workflow", async () => { }); it("set png workflow", async () => { - const pngs = await glob("./tests/*.png"); + const pngs = await glob("./tests/ComfyUI_*.png"); expect(pngs.length).toBeGreaterThanOrEqual(1); for await (const filename of pngs) { @@ -37,3 +37,34 @@ it("set png workflow", async () => { expect(workflow_expect).toEqual(workflow_actual); } }); + +it("extract blank png workflow", async () => { + const pngs = await glob("./tests/Blank_*.png"); + expect(pngs.length).toBeGreaterThanOrEqual(1); + + for await (const filename of pngs) { + const png = Bun.file(filename); + const exif = getPngMetadata(await png.arrayBuffer()); + expect(exif.workflow).toBe(undefined); + } +}); + +it("set blank png workflow", async () => { + const pngs = await glob("./tests/Blank_*.png"); + expect(pngs.length).toBeGreaterThanOrEqual(1); + + for await (const filename of pngs) { + const png = Bun.file(filename); + + const newWorkflow = '{"test":"hello, snomiao"}'; + const buffer2 = setPngMetadata(await png.arrayBuffer(), { + workflow: newWorkflow, + }); + const file2 = new File([buffer2], png.name!); + + const exif2 = getPngMetadata(await file2.arrayBuffer()); + const workflow_actual = JSON.stringify(JSON.parse(exif2.workflow)); + const workflow_expect = JSON.stringify(JSON.parse(newWorkflow)); + expect(workflow_expect).toEqual(workflow_actual); + } +});