feat: extract pnginfo

bot/mod.ts

@@ -1,8 +1,9 @@
-import { Grammy, GrammyAutoQuote, GrammyParseMode, Log } from "../deps.ts";
+import { Grammy, GrammyAutoQuote, GrammyFiles, GrammyParseMode, Log } from "../deps.ts";
 import { formatUserChat } from "../utils.ts";
 import { session, SessionFlavor } from "./session.ts";
 import { queueCommand } from "./queueCommand.ts";
 import { txt2imgCommand, txt2imgQuestion } from "./txt2imgCommand.ts";
+import { pnginfoCommand, pnginfoQuestion } from "./pnginfoCommand.ts";
 
 export const logger = () => Log.getLogger();
 
@@ -12,7 +13,9 @@ type WithRetryApi<T extends Grammy.RawApi> = {
     : T[M];
 };
 
-export type Context = GrammyParseMode.ParseModeFlavor<Grammy.Context> & SessionFlavor;
+export type Context =
+  & GrammyFiles.FileFlavor<GrammyParseMode.ParseModeFlavor<Grammy.Context>>
+  & SessionFlavor;
 export const bot = new Grammy.Bot<Context, Grammy.Api<WithRetryApi<Grammy.RawApi>>>(
   Deno.env.get("TG_BOT_TOKEN") ?? "",
 );
@@ -20,9 +23,7 @@ bot.use(GrammyAutoQuote.autoQuote);
 bot.use(GrammyParseMode.hydrateReply);
 bot.use(session);
 
-bot.catch((err) => {
+bot.api.config.use(GrammyFiles.hydrateFiles(bot.token));
-  logger().error(`Handling update from ${formatUserChat(err.ctx)} failed: ${err}`);
-});
 
 // Automatically retry bot requests if we get a "too many requests" or telegram internal error
 bot.api.config.use(async (prev, method, payload, signal) => {
@@ -46,6 +47,10 @@ bot.api.config.use(async (prev, method, payload, signal) => {
   }
 });
 
+bot.catch((err) => {
+  logger().error(`Handling update from ${formatUserChat(err.ctx)} failed: ${err}`);
+});
+
 // if error happened, try to reply to the user with the error
 bot.use(async (ctx, next) => {
   try {
@@ -68,6 +73,7 @@ bot.api.setMyDescription(
 );
 bot.api.setMyCommands([
   { command: "txt2img", description: "Generate an image" },
+  { command: "pnginfo", description: "Show generation parameters of an image" },
   { command: "queue", description: "Show the current queue" },
 ]);
 
@@ -76,6 +82,9 @@ bot.command("start", (ctx) => ctx.reply("Hello! Use the /txt2img command to gene
 bot.command("txt2img", txt2imgCommand);
 bot.use(txt2imgQuestion.middleware());
 
+bot.command("pnginfo", pnginfoCommand);
+bot.use(pnginfoQuestion.middleware());
+
 bot.command("queue", queueCommand);
 
 bot.command("pause", (ctx) => {

bot/pnginfoCommand.ts

@@ -0,0 +1,51 @@
+import { Grammy, GrammyParseMode, GrammyStatelessQ } from "../deps.ts";
+import { fmt } from "../utils.ts";
+
+import { getPngInfo, parsePngInfo } from "../sd.ts";
+import { Context } from "./mod.ts";
+
+export const pnginfoQuestion = new GrammyStatelessQ.StatelessQuestion<Context>(
+  "pnginfo",
+  async (ctx) => {
+    await pnginfo(ctx, false);
+  },
+);
+
+export async function pnginfoCommand(ctx: Grammy.CommandContext<Context>) {
+  await pnginfo(ctx, true);
+}
+
+async function pnginfo(ctx: Context, includeRepliedTo: boolean): Promise<void> {
+  const document = ctx.message?.document ||
+    (includeRepliedTo ? ctx.message?.reply_to_message?.document : undefined);
+
+  if (document?.mime_type !== "image/png") {
+    await ctx.reply(
+      "Please send me a PNG file." +
+        pnginfoQuestion.messageSuffixMarkdown(),
+      { reply_markup: { force_reply: true, selective: true }, parse_mode: "Markdown" },
+    );
+    return;
+  }
+
+  const file = await ctx.api.getFile(document.file_id);
+  const buffer = await fetch(file.getUrl()).then((resp) => resp.arrayBuffer());
+  const params = parsePngInfo(getPngInfo(new Uint8Array(buffer)) ?? "");
+
+  const { bold } = GrammyParseMode;
+
+  const paramsText = fmt([
+    `${params.prompt}\n`,
+    params.negative_prompt ? fmt`${bold("Negative prompt:")} ${params.negative_prompt}\n` : "",
+    params.steps ? fmt`${bold("Steps:")} ${params.steps}, ` : "",
+    params.sampler_name ? fmt`${bold("Sampler:")} ${params.sampler_name}, ` : "",
+    params.cfg_scale ? fmt`${bold("CFG scale:")} ${params.cfg_scale}, ` : "",
+    params.seed ? fmt`${bold("Seed:")} ${params.seed}, ` : "",
+    params.width && params.height ? fmt`${bold("Size")}: ${params.width}x${params.height}` : "",
+  ]);
+
+  await ctx.reply(paramsText.text, {
+    reply_to_message_id: ctx.message?.message_id,
+    entities: paramsText.entities,
+  });
+}

bot/txt2imgCommand.ts

@@ -1,7 +1,7 @@
 import { Grammy, GrammyStatelessQ } from "../deps.ts";
 import { formatUserChat } from "../utils.ts";
 import { jobStore } from "../db/jobStore.ts";
-import { parsePngInfo } from "../sd.ts";
+import { getPngInfo, parsePngInfo, SdTxt2ImgRequest } from "../sd.ts";
 import { Context, logger } from "./mod.ts";
 
 export const txt2imgQuestion = new GrammyStatelessQ.StatelessQuestion<Context>(
@@ -43,8 +43,23 @@ async function txt2img(ctx: Context, match: string, includeRepliedTo: boolean):
     return;
   }
 
-  let params = parsePngInfo(match);
+  let params: Partial<SdTxt2ImgRequest> = {};
+
   const repliedToMsg = ctx.message.reply_to_message;
+
+  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());
+    const fileParams = parsePngInfo(getPngInfo(new Uint8Array(buffer)) ?? "");
+    params = {
+      ...params,
+      ...fileParams,
+      prompt: [params.prompt, fileParams.prompt].filter(Boolean).join("\n"),
+      negative_prompt: [params.negative_prompt, fileParams.negative_prompt]
+        .filter(Boolean).join("\n"),
+    };
+  }
+
   const repliedToText = repliedToMsg?.text || repliedToMsg?.caption;
   if (includeRepliedTo && repliedToText) {
     // TODO: remove bot command from replied to text
@@ -53,8 +68,18 @@ async function txt2img(ctx: Context, match: string, includeRepliedTo: boolean):
       ...originalParams,
       ...params,
       prompt: [originalParams.prompt, params.prompt].filter(Boolean).join("\n"),
+      negative_prompt: [originalParams.negative_prompt, params.negative_prompt]
+        .filter(Boolean).join("\n"),
     };
   }
+
+  const messageParams = parsePngInfo(match);
+  params = {
+    ...params,
+    ...messageParams,
+    prompt: [params.prompt, messageParams.prompt].filter(Boolean).join("\n"),
+  };
+
   if (!params.prompt) {
     await ctx.reply(
       "Please tell me what you want to see." +

deps.ts

@@ -12,6 +12,7 @@ export * as GrammyAutoQuote from "https://deno.land/x/grammy_autoquote@v1.1.2/mo
 export * as GrammyParseMode from "https://deno.land/x/grammy_parse_mode@1.7.1/mod.ts";
 export * as GrammyKvStorage from "https://deno.land/x/grammy_storages@v2.3.1/denokv/src/mod.ts";
 export * as GrammyStatelessQ from "https://deno.land/x/grammy_stateless_question_alpha@v3.0.3/mod.ts";
+export * as GrammyFiles from "https://deno.land/x/grammy_files@v1.0.4/mod.ts";
 export * as FileType from "npm:file-type@18.5.0";
 // @deno-types="./types/png-chunks-extract.d.ts"
 export * as PngChunksExtract from "npm:png-chunks-extract@1.0.0";

pnginfo.test.ts

@@ -0,0 +1,83 @@
+import {
+  assert,
+  assertEquals,
+  assertMatch,
+} from "https://deno.land/std@0.135.0/testing/asserts.ts";
+import { parsePngInfo } from "./sd.ts";
+
+Deno.test("parses pnginfo", async (t) => {
+  await t.step("1", () => {
+    const params = parsePngInfo(
+      `female, red fox, pink hair, (long hair), eyeshadow, lipstick, armband, midriff, leather shorts, (fishnet) legwear, fingerless gloves, <3, tongue out, 
+      presenting breasts, ((flashing)) breasts, shirt lift, raised shirt, public exposure, 
+      lgbt pride, pride colors, pride color clothing, flag \\(object\\), pride march, public, 
+      braeburned, keadonger, zaush, jishinu, pixelsketcher, detailed background, insane detail, soft shading, masterpiece
+      Negative prompt:  bad-hands-5, boring_e621
+      Steps: 40, Sampler: Euler a, CFG scale: 8, Seed: 2843818575, Size: 768x768, Model hash: e2d72a81a3, Model: bb95FurryMix_v90, Denoising strength: 0.28,
+      SD upscale overlap: 64, SD upscale upscaler: Lanczos, Version: v1.4.0`,
+    );
+    assertMatch(params.prompt ?? "", /\blong hair\b/);
+    assertMatch(params.prompt ?? "", /\bflag \\\(object\\\)/);
+    assertMatch(params.prompt ?? "", /\bmasterpiece\b/);
+    assert(!params.prompt?.includes("2843818575"));
+    assertMatch(params.negative_prompt ?? "", /\bbad-hands-5\b/);
+    assertMatch(params.negative_prompt ?? "", /\bboring_e621\b/);
+    assert(!params.negative_prompt?.includes("2843818575"));
+    assertEquals(params.steps, 40);
+    assertEquals(params.cfg_scale, 8);
+    assertEquals(params.width, 768);
+    assertEquals(params.height, 768);
+  });
+
+  await t.step("2", () => {
+    const params = parsePngInfo(
+      `anthro, female, wolf:1.2, long hair, fluffy tail, thick thighs, 
+      wraps, loincloth, tribal clothing, tribal body markings, bone necklace, feathers in hair, skimpy, 
+      holding spear, weapon, crouching, digitigrade, action pose, perspective, motion lines, forest, hunting, female pred, grin, angry,
+      by kenket, by ruaidri, by keadonger, by braeburned, by twiren, detailed fur, hires, masterpiece, <lora:add_detail:0.8>
+      Negative prompt: boring_e621_fluffyrock
+      Steps: 40, Sampler: Euler a, CFG scale: 7, Seed: 2876880391, Size: 1536x1536, Model hash: 06ac6055bd,
+      Model: fluffyrock-576-704-832-960-1088-lion-low-lr-e61-terminal-snr-e34, Denoising strength: 0.2, 
+      Ultimate SD upscale upscaler: None, Ultimate SD upscale tile_width: 768, Ultimate SD upscale tile_height: 768, 
+      Ultimate SD upscale mask_blur: 8, Ultimate SD upscale padding: 48, 
+      Lora hashes: "add_detail: 7c6bad76eb54", Version: v1.3.2`,
+    );
+    assertMatch(params.prompt ?? "", /\bwolf\b/);
+    assertMatch(params.prompt ?? "", /\bdigitigrade\b/);
+    assertMatch(params.prompt ?? "", /\bby ruaidri\b/);
+    assert(!params.prompt?.includes("7c6bad76eb54"));
+    assert(!params.prompt?.includes("tile_width"));
+    assert(!params.prompt?.includes("boring_e621_fluffyrock"));
+    assertMatch(params.negative_prompt ?? "", /\bboring_e621_fluffyrock\b/);
+    assert(!params.negative_prompt?.includes("7c6bad76eb54"));
+    assert(!params.negative_prompt?.includes("tile_width"));
+    assert(!params.negative_prompt?.includes("add_detail"));
+    assertEquals(params.steps, 40);
+    assertEquals(params.cfg_scale, 7);
+    assertEquals(params.width, 1536);
+    assertEquals(params.height, 1536);
+  });
+
+  await t.step("3", () => {
+    const params = parsePngInfo(
+      `anthro, female, red fox, long black hair with pink highlights, bra, underwear, digitigrade, pawpads, foot focus, foot fetish, sitting, 4 toes, 
+      sticker, outline, simple background, by braeburned, by alibi-cami, by ultrabondagefairy, by dripponi, <lora:easy_sticker:0.5>
+      Negative prompt: boring_e621_v4, happy, smile, 
+      Steps: 40, Sampler: Euler a, CFG scale: 7, Seed: 3154849350, Size: 512x512, Model hash: fd926f7598, Model: bb95FurryMix_v100,
+      Denoising strength: 0.65, Lora hashes: "easy_sticker: 2c98dc945091", Version: v1.3.2`,
+    );
+    assertMatch(params.prompt ?? "", /\bbra\b/);
+    assertMatch(params.prompt ?? "", /\blora:easy_sticker\b/);
+    assert(!params.prompt?.includes("smile"));
+    assert(!params.prompt?.includes("Euler a"));
+    assert(!params.prompt?.includes("bb95FurryMix_v100"));
+    assertMatch(params.negative_prompt ?? "", /\bboring_e621_v4\b/);
+    assert(!params.negative_prompt?.includes("simple background"));
+    assert(!params.negative_prompt?.includes("easy_sticker"));
+    assert(!params.negative_prompt?.includes("bb95FurryMix_v100"));
+    assertEquals(params.steps, 40);
+    assertEquals(params.cfg_scale, 7);
+    assertEquals(params.width, 512);
+    assertEquals(params.height, 512);
+  });
+});

sd.ts

@@ -249,15 +249,20 @@ export function parsePngInfo(pngInfo: string): Partial<SdTxt2ImgRequest> {
   const prompt: string[] = [];
   const negativePrompt: string[] = [];
   for (const tag of tags) {
-    const paramValuePair = tag.trim().match(/^(\w+\s*\w*):\s+([\d\w. ]+)\s*$/u);
+    const paramValuePair = tag.trim().match(/^(\w+\s*\w*):\s+(.*)$/u);
     if (paramValuePair) {
       const [, param, value] = paramValuePair;
       switch (param.replace(/\s+/u, "").toLowerCase()) {
+        case "positiveprompt":
+        case "positive":
         case "prompt":
+        case "pos":
           part = "prompt";
           prompt.push(value.trim());
           break;
         case "negativeprompt":
+        case "negative":
+        case "neg":
           part = "negative_prompt";
           negativePrompt.push(value.trim());
           break;
@@ -269,6 +274,7 @@ export function parsePngInfo(pngInfo: string): Partial<SdTxt2ImgRequest> {
           break;
         }
         case "cfgscale":
+        case "cfg":
         case "detail": {
           part = "params";
           const cfgScale = Number(value.trim());
@@ -288,6 +294,17 @@ export function parsePngInfo(pngInfo: string): Partial<SdTxt2ImgRequest> {
           }
           break;
         }
+        case "seed":
+        case "model":
+        case "modelhash":
+        case "modelname":
+        case "sampler":
+        case "denoisingstrength":
+        case "denoising":
+        case "denoise":
+          part = "params";
+          // ignore for now
+          break;
         default:
           break;
       }