From 73810d3204eb95a28abf10e45dcf04c4480ca592 Mon Sep 17 00:00:00 2001 From: nameless Date: Fri, 10 Nov 2023 01:39:59 +0000 Subject: [PATCH] fix: use exifreader for png info - replace pnginfo module - allow jpeg for tag extraction - rename pnginfo to getprompt --- bot/mod.ts | 4 ++-- bot/parsePngInfo.ts | 26 ++++++++++++++++++-------- bot/pnginfoCommand.ts | 4 ++-- bot/txt2imgCommand.ts | 2 +- deno.json | 3 +-- deno.lock | 4 ++-- 6 files changed, 26 insertions(+), 17 deletions(-) diff --git a/bot/mod.ts b/bot/mod.ts index ad3b35b..7052f5d 100644 --- a/bot/mod.ts +++ b/bot/mod.ts @@ -115,7 +115,7 @@ bot.api.setMyDescription( bot.api.setMyCommands([ { command: "txt2img", description: "Generate image from text" }, { command: "img2img", description: "Generate image from image" }, - { command: "pnginfo", description: "Show generation parameters of an image" }, + { command: "getprompt", description: "Try to extract prompt from raw file" }, { command: "queue", description: "Show the current queue" }, { command: "cancel", description: "Cancel all your requests" }, ]); @@ -160,7 +160,7 @@ bot.use(txt2imgQuestion.middleware()); bot.command("img2img", img2imgCommand); bot.use(img2imgQuestion.middleware()); -bot.command("pnginfo", pnginfoCommand); +bot.command("getprompt", pnginfoCommand); bot.use(pnginfoQuestion.middleware()); bot.command("queue", queueCommand); diff --git a/bot/parsePngInfo.ts b/bot/parsePngInfo.ts index 9bf2d40..b751fb3 100644 --- a/bot/parsePngInfo.ts +++ b/bot/parsePngInfo.ts @@ -1,12 +1,22 @@ -import { decode } from "png_chunk_text"; -import extractChunks from "png_chunks_extract"; +import * as ExifReader from "exifreader"; -export function getPngInfo(pngData: Uint8Array): string | undefined { - return extractChunks(pngData) - .filter((chunk) => chunk.name === "tEXt") - .map((chunk) => decode(chunk.data)) - .find((textChunk) => textChunk.keyword === "parameters") - ?.text; +export function getPngInfo(pngData: ArrayBuffer): string | undefined { + const image = ExifReader.load(pngData); + + if (image.UserComment && image.UserComment.value) { + // JPEG image + return String.fromCharCode.apply( + 0, + (image.UserComment.value as number[]).filter((char: number) => char != 0), + ) + .replace("UNICODE", ""); + } else if (image.parameters && image.parameters.description) { + // PNG image + return image.parameters.description; + } else { + // Unknown image type + return undefined; + } } export interface PngInfo { diff --git a/bot/pnginfoCommand.ts b/bot/pnginfoCommand.ts index bf54230..20c2be9 100644 --- a/bot/pnginfoCommand.ts +++ b/bot/pnginfoCommand.ts @@ -20,7 +20,7 @@ async function pnginfo(ctx: ErisContext, includeRepliedTo: boolean): Promise resp.arrayBuffer()); - const params = parsePngInfo(getPngInfo(new Uint8Array(buffer)) ?? ""); + const params = parsePngInfo(getPngInfo(buffer) ?? "Nothing found.", undefined, true); const paramsText = fmt([ `${params.prompt}\n`, diff --git a/bot/txt2imgCommand.ts b/bot/txt2imgCommand.ts index 39f7277..1863728 100644 --- a/bot/txt2imgCommand.ts +++ b/bot/txt2imgCommand.ts @@ -60,7 +60,7 @@ async function txt2img(ctx: ErisContext, match: string, includeRepliedTo: boolea if (includeRepliedTo && repliedToMsg?.document?.mime_type === "image/png") { const file = await ctx.api.getFile(repliedToMsg.document.file_id); const buffer = await fetch(file.getUrl()).then((resp) => resp.arrayBuffer()); - params = parsePngInfo(getPngInfo(new Uint8Array(buffer)) ?? "", params); + params = parsePngInfo(getPngInfo(buffer) ?? "", params); } const repliedToText = repliedToMsg?.text || repliedToMsg?.caption; diff --git a/deno.json b/deno.json index b23b395..1b838a6 100644 --- a/deno.json +++ b/deno.json @@ -11,6 +11,7 @@ "async": "https://deno.land/x/async@v2.0.2/mod.ts", "date-fns": "https://cdn.skypack.dev/date-fns@2.30.0?dts", "date-fns/utc": "https://cdn.skypack.dev/@date-fns/utc@1.1.0?dts", + "exifreader": "https://esm.sh/exifreader@4.14.1", "file_type": "https://esm.sh/file-type@18.5.0", "grammy": "https://lib.deno.dev/x/grammy@1/mod.ts", "grammy_autoquote": "https://lib.deno.dev/x/grammy_autoquote@1/mod.ts", @@ -23,8 +24,6 @@ "kvfs": "https://deno.land/x/kvfs@v0.1.0/mod.ts", "kvmq": "https://deno.land/x/kvmq@v0.3.0/mod.ts", "openapi_fetch": "https://esm.sh/openapi-fetch@0.7.6", - "png_chunk_text": "https://esm.sh/png-chunk-text@1.0.0", - "png_chunks_extract": "https://esm.sh/png-chunks-extract@1.0.0", "react": "https://esm.sh/react@18.2.0?dev", "react-dom/client": "https://esm.sh/react-dom@18.2.0/client?external=react&dev", "react-flip-move": "https://esm.sh/react-flip-move@3.0.5?external=react&dev", diff --git a/deno.lock b/deno.lock index ab58e99..81245d0 100644 --- a/deno.lock +++ b/deno.lock @@ -303,10 +303,9 @@ "https://esm.sh/@grammyjs/types@2.12.1": "eebe448d3bf3d4fdaeacee50e31b9e3947ce965d2591f1036e5c0273cb7aec36", "https://esm.sh/@twind/core@1.1.3": "bf3e64c13de95b061a721831bf2aed322f6e7335848b06d805168c3212686c1d", "https://esm.sh/@twind/preset-tailwind@1.1.4": "29f5e4774c344ff08d2636e3d86382060a8b40d6bce1c158b803df1823c2804c", + "https://esm.sh/exifreader@4.14.1": "d0f21973393b0d1a6ed329dac8fcfb2f87ce47fe40b8172e205e7d6d85790bb6", "https://esm.sh/file-type@18.5.0": "f01f6eddb05d365925545e26bd449f934dc042bead882b2795c35c4f48d690a3", "https://esm.sh/openapi-fetch@0.7.6": "43eff6df93e773801cb6ade02b0f47c05d0bfe2fb044adbf2b94fa74ff30d35b", - "https://esm.sh/png-chunk-text@1.0.0": "08beb86f31b5ff70240650fe095b6c4e4037e5c1d5917f9338e7633cae680356", - "https://esm.sh/png-chunks-extract@1.0.0": "da06bbd3c08199d72ab16354abe5ffd2361cb891ccbb44d757a0a0a4fbfa12b5", "https://esm.sh/react-dom@18.2.0/client?external=react&dev": "2eb39c339720d727591fd55fb44ffcb6f14b06812af0a71e7a2268185b5b6e73", "https://esm.sh/react-flip-move@3.0.5?external=react&dev": "4390c0777a0bec583d3e6cb5e4b33831ac937d670d894a20e4f192ce8cd21bae", "https://esm.sh/react-intl@6.4.7?external=react&dev": "60e68890e2c5ef3c02d37a89c53e056b1bbd1c8467a7aae62f0b634abc7a8a5f", @@ -333,6 +332,7 @@ "https://esm.sh/v133/@remix-run/router@1.9.0/X-ZS9yZWFjdA/denonext/router.development.mjs": "97f96f031a7298b0afbbadb23177559ee9b915314d7adf0b9c6a8d7f452c70e7", "https://esm.sh/v133/client-only@0.0.1/X-ZS9yZWFjdA/denonext/client-only.development.mjs": "b7efaef2653f7f628d084e24a4d5a857c9cd1a5e5060d1d9957e185ee98c8a28", "https://esm.sh/v133/crc-32@0.3.0/denonext/crc-32.mjs": "92bbd96cd5a92e45267cf4b3d3928b9355d16963da1ba740637fb10f1daca590", + "https://esm.sh/v133/exifreader@4.14.1/denonext/exifreader.mjs": "691e1c1d1337ccaf092bf39115fdac56bf69e93259d04bb03985e918202317ab", "https://esm.sh/v133/file-type@18.5.0/denonext/file-type.mjs": "785cac1bc363448647871e1b310422e01c9cfb7b817a68690710b786d3598311", "https://esm.sh/v133/hoist-non-react-statics@3.3.2/X-ZS9yZWFjdA/denonext/hoist-non-react-statics.development.mjs": "fcafac9e3c33810f18ecb43dfc32ce80efc88adc63d3f49fd7ada0665d2146c6", "https://esm.sh/v133/ieee754@1.2.1/denonext/ieee754.mjs": "9ec2806065f50afcd4cf3f3f2f38d93e777a92a5954dda00d219f57d4b24832f",