From db5b8a647468e1ceadcaba4f4504fa0932388240 Mon Sep 17 00:00:00 2001
From: pinks <lisq@cock.li>
Date: Sun, 24 Sep 2023 15:19:44 +0000
Subject: [PATCH] fix: img2img default image size from original (#17)

fixes #2

Reviewed-on: https://git.foxo.me/pinks/eris/pulls/17
Co-authored-by: pinks <lisq@cock.li>
Co-committed-by: pinks <lisq@cock.li>
---
 bot/img2imgCommand.ts | 40 ++++++++++++++++++++++------------------
 1 file changed, 22 insertions(+), 18 deletions(-)

diff --git a/bot/img2imgCommand.ts b/bot/img2imgCommand.ts
index 50a195f..12c7cb8 100644
--- a/bot/img2imgCommand.ts
+++ b/bot/img2imgCommand.ts
@@ -7,10 +7,12 @@ import { parsePngInfo, PngInfo } from "../sd/parsePngInfo.ts";
 import { formatUserChat } from "../utils/formatUserChat.ts";
 import { ErisContext, logger } from "./mod.ts";
 
+type QuestionState = { fileId?: string; params?: Partial<PngInfo> };
+
 export const img2imgQuestion = new StatelessQuestion<ErisContext>(
   "img2img",
-  async (ctx, state) => {
-    // todo: also save original image size in state
+  async (ctx, stateString) => {
+    const state: QuestionState = JSON.parse(stateString);
     await img2img(ctx, ctx.message.text, false, state);
   },
 );
@@ -23,7 +25,7 @@ async function img2img(
   ctx: ErisContext,
   match: string | undefined,
   includeRepliedTo: boolean,
-  fileId?: string,
+  state: QuestionState = {},
 ): Promise<void> {
   if (!ctx.message?.from?.id) {
     await ctx.reply("I don't know who you are");
@@ -53,49 +55,51 @@ async function img2img(
     return;
   }
 
-  let params: Partial<PngInfo> = {};
-
   const repliedToMsg = ctx.message.reply_to_message;
 
   if (includeRepliedTo && repliedToMsg?.photo) {
     const photos = repliedToMsg.photo;
     const biggestPhoto = maxBy(photos, (p) => p.width * p.height);
     if (!biggestPhoto) throw new Error("Message was a photo but had no photos?");
-    fileId = biggestPhoto.file_id;
-    params.width = biggestPhoto.width;
-    params.height = biggestPhoto.height;
+    state.fileId = biggestPhoto.file_id;
+    if (!state.params) state.params = {};
+    state.params.width = biggestPhoto.width;
+    state.params.height = biggestPhoto.height;
   }
 
   if (ctx.message.photo) {
     const photos = ctx.message.photo;
     const biggestPhoto = maxBy(photos, (p) => p.width * p.height);
     if (!biggestPhoto) throw new Error("Message was a photo but had no photos?");
-    fileId = biggestPhoto.file_id;
-    params.width = biggestPhoto.width;
-    params.height = biggestPhoto.height;
+    state.fileId = biggestPhoto.file_id;
+    if (!state.params) state.params = {};
+    state.params.width = biggestPhoto.width;
+    state.params.height = biggestPhoto.height;
   }
 
   const repliedToText = repliedToMsg?.text || repliedToMsg?.caption;
   if (includeRepliedTo && repliedToText) {
     // TODO: remove bot command from replied to text
-    params = parsePngInfo(repliedToText, params);
+    state.params = parsePngInfo(repliedToText, state.params);
   }
 
-  params = parsePngInfo(match ?? "", params);
+  if (match) {
+    state.params = parsePngInfo(match, state.params);
+  }
 
-  if (!fileId) {
+  if (!state.fileId) {
     await ctx.reply(
       "Please show me a picture to repaint." +
-        img2imgQuestion.messageSuffixMarkdown(),
+        img2imgQuestion.messageSuffixMarkdown(JSON.stringify(state satisfies QuestionState)),
       { reply_markup: { force_reply: true, selective: true }, parse_mode: "Markdown" },
     );
     return;
   }
 
-  if (!params.prompt) {
+  if (!state.params?.prompt) {
     await ctx.reply(
       "Please describe the picture you want to repaint." +
-        img2imgQuestion.messageSuffixMarkdown(fileId),
+        img2imgQuestion.messageSuffixMarkdown(JSON.stringify(state satisfies QuestionState)),
       { reply_markup: { force_reply: true, selective: true }, parse_mode: "Markdown" },
     );
     return;
@@ -104,7 +108,7 @@ async function img2img(
   const replyMessage = await ctx.reply("Accepted. You are now in queue.");
 
   await generationQueue.pushJob({
-    task: { type: "img2img", params, fileId },
+    task: { type: "img2img", fileId: state.fileId, params: state.params },
     from: ctx.message.from,
     chat: ctx.message.chat,
     requestMessage: ctx.message,