feat(3-some-png-images-that-are-not-generated-by-comfyui-seem-not-to-…

app/utils/exif.ts

@@ -36,6 +36,9 @@ export function getPngMetadata(buffer: ArrayBuffer): Record<string, string> {
         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<string, string>
-): Uint8Array<ArrayBufferLike> {
+): 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;
 }

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);
+  }
+});