fix: pnginfo iTXt

This commit is contained in:
Vargoshi 2023-10-15 22:18:39 +02:00
parent 2f059ceaff
commit 320924da91
2 changed files with 14 additions and 3 deletions

View File

@ -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":

View File

@ -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` : "",