From 320924da9132ded4b8e719567be134bbda771ae8 Mon Sep 17 00:00:00 2001
From: Vargoshi <marcinflisak1998@gmail.com>
Date: Sun, 15 Oct 2023 22:18:39 +0200
Subject: [PATCH] fix: pnginfo iTXt

---
 bot/parsePngInfo.ts   | 10 +++++++---
 bot/pnginfoCommand.ts |  7 +++++++
 2 files changed, 14 insertions(+), 3 deletions(-)

diff --git a/bot/parsePngInfo.ts b/bot/parsePngInfo.ts
index 45ada41..2fbfda7 100644
--- a/bot/parsePngInfo.ts
+++ b/bot/parsePngInfo.ts
@@ -3,7 +3,7 @@ import extractChunks from "png_chunks_extract";
 
 export function getPngInfo(pngData: Uint8Array): string | undefined {
   return extractChunks(pngData)
-    .filter((chunk) => chunk.name === "tEXt")
+    .filter((chunk) => chunk.name === "tEXt" || chunk.name === "iTXt")
     .map((chunk) => decode(chunk.data))
     .find((textChunk) => textChunk.keyword === "parameters")
     ?.text;
@@ -25,7 +25,11 @@ interface PngInfoExtra extends PngInfo {
   upscale?: number;
 }
 
-export function parsePngInfo(pngInfo: string, baseParams?: Partial<PngInfo>, shouldParseSeed?: boolean): Partial<PngInfo> {
+export function parsePngInfo(
+  pngInfo: string,
+  baseParams?: Partial<PngInfo>,
+  shouldParseSeed?: boolean,
+): Partial<PngInfo> {
   const tags = pngInfo.split(/[,;]+|\.+\s|\n/u);
   let part: "prompt" | "negative_prompt" | "params" = "prompt";
   const params: Partial<PngInfoExtra> = {};
@@ -104,8 +108,8 @@ export function parsePngInfo(pngInfo: string, baseParams?: Partial<PngInfo>, sho
           if (shouldParseSeed) {
             const seed = Number(value.trim());
             params.seed = seed;
-            break;
           }
+          break;
         }
         case "model":
         case "modelhash":
diff --git a/bot/pnginfoCommand.ts b/bot/pnginfoCommand.ts
index d44ed81..41f4724 100644
--- a/bot/pnginfoCommand.ts
+++ b/bot/pnginfoCommand.ts
@@ -36,6 +36,13 @@ async function pnginfo(ctx: ErisContext, includeRepliedTo: boolean): Promise<voi
   const buffer = await fetch(file.getUrl()).then((resp) => resp.arrayBuffer());
   const params = parsePngInfo(getPngInfo(new Uint8Array(buffer)) ?? "");
 
+  if (!params.prompt) {
+    ctx.reply("File doesn't contain any parameters.", {
+      reply_to_message_id: ctx.message?.message_id,
+    });
+    return;
+  }
+
   const paramsText = fmt([
     `${params.prompt}\n`,
     params.negative_prompt ? fmt`${bold("Negative prompt:")} ${params.negative_prompt}\n` : "",