forked from pinks/eris
1
0
Fork 0
eris/bot/img2imgCommand.ts

116 lines
3.5 KiB
TypeScript
Raw Normal View History

2023-09-24 13:08:35 +00:00
import { CommandContext } from "grammy";
import { StatelessQuestion } from "grammy_stateless_question";
import { maxBy } from "std/collections";
2023-09-23 18:49:05 +00:00
import { getConfig } from "../app/config.ts";
2023-09-24 13:08:35 +00:00
import { generationQueue } from "../app/generationQueue.ts";
import { parsePngInfo, PngInfo } from "../sd/parsePngInfo.ts";
import { formatUserChat } from "../utils/formatUserChat.ts";
import { ErisContext, logger } from "./mod.ts";
2023-09-12 01:57:44 +00:00
2023-09-24 13:08:35 +00:00
export const img2imgQuestion = new StatelessQuestion<ErisContext>(
2023-09-12 01:57:44 +00:00
"img2img",
async (ctx, state) => {
// todo: also save original image size in state
await img2img(ctx, ctx.message.text, false, state);
},
);
2023-09-24 13:08:35 +00:00
export async function img2imgCommand(ctx: CommandContext<ErisContext>) {
2023-09-12 01:57:44 +00:00
await img2img(ctx, ctx.match, true);
}
async function img2img(
2023-09-24 13:08:35 +00:00
ctx: ErisContext,
2023-09-12 01:57:44 +00:00
match: string | undefined,
includeRepliedTo: boolean,
fileId?: string,
): Promise<void> {
if (!ctx.message?.from?.id) {
await ctx.reply("I don't know who you are");
return;
}
2023-09-22 02:59:22 +00:00
const config = await getConfig();
if (config.pausedReason != null) {
await ctx.reply(`I'm paused: ${config.pausedReason || "No reason given"}`);
2023-09-12 01:57:44 +00:00
return;
}
2023-09-22 02:59:22 +00:00
const jobs = await generationQueue.getAllJobs();
if (jobs.length >= config.maxJobs) {
2023-09-12 01:57:44 +00:00
await ctx.reply(
2023-09-22 02:59:22 +00:00
`The queue is full. Try again later. (Max queue size: ${config.maxJobs})`,
2023-09-12 01:57:44 +00:00
);
return;
}
2023-09-22 02:59:22 +00:00
const userJobs = jobs.filter((job) => job.state.from.id === ctx.message?.from?.id);
if (userJobs.length >= config.maxUserJobs) {
2023-09-12 01:57:44 +00:00
await ctx.reply(
2023-09-22 02:59:22 +00:00
`You already have ${config.maxUserJobs} jobs in queue. Try again later.`,
2023-09-12 01:57:44 +00:00
);
return;
}
let params: Partial<PngInfo> = {};
const repliedToMsg = ctx.message.reply_to_message;
if (includeRepliedTo && repliedToMsg?.photo) {
const photos = repliedToMsg.photo;
2023-09-24 13:08:35 +00:00
const biggestPhoto = maxBy(photos, (p) => p.width * p.height);
2023-09-12 01:57:44 +00:00
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;
}
if (ctx.message.photo) {
const photos = ctx.message.photo;
2023-09-24 13:08:35 +00:00
const biggestPhoto = maxBy(photos, (p) => p.width * p.height);
2023-09-12 01:57:44 +00:00
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;
}
const repliedToText = repliedToMsg?.text || repliedToMsg?.caption;
if (includeRepliedTo && repliedToText) {
// TODO: remove bot command from replied to text
2023-09-16 11:49:12 +00:00
params = parsePngInfo(repliedToText, params);
2023-09-12 01:57:44 +00:00
}
2023-09-16 11:49:12 +00:00
params = parsePngInfo(match ?? "", params);
2023-09-12 01:57:44 +00:00
if (!fileId) {
await ctx.reply(
"Please show me a picture to repaint." +
img2imgQuestion.messageSuffixMarkdown(),
{ reply_markup: { force_reply: true, selective: true }, parse_mode: "Markdown" },
);
return;
}
if (!params.prompt) {
await ctx.reply(
"Please describe the picture you want to repaint." +
img2imgQuestion.messageSuffixMarkdown(fileId),
{ reply_markup: { force_reply: true, selective: true }, parse_mode: "Markdown" },
);
return;
}
const replyMessage = await ctx.reply("Accepted. You are now in queue.");
2023-09-22 02:59:22 +00:00
await generationQueue.pushJob({
2023-09-12 01:57:44 +00:00
task: { type: "img2img", params, fileId },
from: ctx.message.from,
chat: ctx.message.chat,
2023-09-22 02:59:22 +00:00
requestMessage: ctx.message,
replyMessage: replyMessage,
}, { retryCount: 3, repeatDelayMs: 10_000 });
2023-09-12 01:57:44 +00:00
logger().debug(`Generation (img2img) enqueued for ${formatUserChat(ctx.message)}`);
2023-09-12 01:57:44 +00:00
}