From 320924da9132ded4b8e719567be134bbda771ae8 Mon Sep 17 00:00:00 2001 From: Vargoshi 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, shouldParseSeed?: boolean): Partial { +export function parsePngInfo( + pngInfo: string, + baseParams?: Partial, + shouldParseSeed?: boolean, +): Partial { const tags = pngInfo.split(/[,;]+|\.+\s|\n/u); let part: "prompt" | "negative_prompt" | "params" = "prompt"; const params: Partial = {}; @@ -104,8 +108,8 @@ export function parsePngInfo(pngInfo: string, baseParams?: Partial, 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 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` : "", -- 2.43.5