fix: img2img default image size from original (#17)

fixes #2

Reviewed-on: #17
Co-authored-by: pinks <lisq@cock.li>
Co-committed-by: pinks <lisq@cock.li>
This commit is contained in:
pinks 2023-09-24 15:19:44 +00:00 committed by pinks
parent ff5320bcf4
commit db5b8a6474
1 changed files with 22 additions and 18 deletions

View File

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