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 { export function getPngInfo(pngData: Uint8Array): string | undefined {
return extractChunks(pngData) return extractChunks(pngData)
.filter((chunk) => chunk.name === "tEXt") .filter((chunk) => chunk.name === "tEXt" || chunk.name === "iTXt")
.map((chunk) => decode(chunk.data)) .map((chunk) => decode(chunk.data))
.find((textChunk) => textChunk.keyword === "parameters") .find((textChunk) => textChunk.keyword === "parameters")
?.text; ?.text;
@ -25,7 +25,11 @@ interface PngInfoExtra extends PngInfo {
upscale?: number; 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); const tags = pngInfo.split(/[,;]+|\.+\s|\n/u);
let part: "prompt" | "negative_prompt" | "params" = "prompt"; let part: "prompt" | "negative_prompt" | "params" = "prompt";
const params: Partial<PngInfoExtra> = {}; const params: Partial<PngInfoExtra> = {};
@ -104,8 +108,8 @@ export function parsePngInfo(pngInfo: string, baseParams?: Partial<PngInfo>, sho
if (shouldParseSeed) { if (shouldParseSeed) {
const seed = Number(value.trim()); const seed = Number(value.trim());
params.seed = seed; params.seed = seed;
break;
} }
break;
} }
case "model": case "model":
case "modelhash": 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 buffer = await fetch(file.getUrl()).then((resp) => resp.arrayBuffer());
const params = parsePngInfo(getPngInfo(new Uint8Array(buffer)) ?? ""); 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([ const paramsText = fmt([
`${params.prompt}\n`, `${params.prompt}\n`,
params.negative_prompt ? fmt`${bold("Negative prompt:")} ${params.negative_prompt}\n` : "", params.negative_prompt ? fmt`${bold("Negative prompt:")} ${params.negative_prompt}\n` : "",