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

119 lines
3.8 KiB
TypeScript
Raw Normal View History

2023-09-24 13:08:35 +00:00
import { CommandContext } from "grammy";
import { StatelessQuestion } from "grammy_stateless_question";
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 { getPngInfo, parsePngInfo, PngInfo } from "../sd/parsePngInfo.ts";
import { formatUserChat } from "../utils/formatUserChat.ts";
import { ErisContext, logger } from "./mod.ts";
2023-09-10 18:56:17 +00:00
2023-09-24 13:08:35 +00:00
export const txt2imgQuestion = new StatelessQuestion<ErisContext>(
2023-09-10 23:59:33 +00:00
"txt2img",
async (ctx) => {
if (!ctx.message.text) return;
2023-09-11 16:03:07 +00:00
await txt2img(ctx, ctx.message.text, false);
2023-09-10 23:59:33 +00:00
},
);
2023-09-24 13:08:35 +00:00
export async function txt2imgCommand(ctx: CommandContext<ErisContext>) {
2023-09-10 23:59:33 +00:00
await txt2img(ctx, ctx.match, true);
}
2023-09-24 13:08:35 +00:00
async function txt2img(ctx: ErisContext, match: string, includeRepliedTo: boolean): Promise<void> {
2023-09-10 23:59:33 +00:00
if (!ctx.message?.from?.id) {
2023-09-24 19:58:09 +00:00
await ctx.reply("I don't know who you are", { reply_to_message_id: ctx.message?.message_id });
2023-09-11 09:22:15 +00:00
return;
2023-09-10 18:56:17 +00:00
}
2023-09-10 23:59:33 +00:00
2023-09-22 02:59:22 +00:00
const config = await getConfig();
if (config.pausedReason != null) {
2023-09-24 19:58:09 +00:00
await ctx.reply(`I'm paused: ${config.pausedReason || "No reason given"}`, {
reply_to_message_id: ctx.message?.message_id,
});
2023-09-11 09:22:15 +00:00
return;
2023-09-10 18:56:17 +00:00
}
2023-09-10 23:59:33 +00:00
2023-09-22 02:59:22 +00:00
const jobs = await generationQueue.getAllJobs();
if (jobs.length >= config.maxJobs) {
2023-09-24 19:58:09 +00:00
await ctx.reply(`The queue is full. Try again later. (Max queue size: ${config.maxJobs})`, {
reply_to_message_id: ctx.message?.message_id,
});
2023-09-11 09:22:15 +00:00
return;
2023-09-10 18:56:17 +00:00
}
2023-09-10 23:59:33 +00:00
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-24 19:58:09 +00:00
await ctx.reply(`You already have ${config.maxUserJobs} jobs in queue. Try again later.`, {
reply_to_message_id: ctx.message?.message_id,
});
2023-09-11 09:22:15 +00:00
return;
2023-09-10 18:56:17 +00:00
}
2023-09-10 23:59:33 +00:00
2023-09-12 01:57:44 +00:00
let params: Partial<PngInfo> = {};
2023-09-11 20:43:12 +00:00
2023-09-10 23:59:33 +00:00
const repliedToMsg = ctx.message.reply_to_message;
2023-09-11 20:43:12 +00:00
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());
2023-09-16 11:49:12 +00:00
params = parsePngInfo(getPngInfo(new Uint8Array(buffer)) ?? "", params);
2023-09-11 20:43:12 +00:00
}
2023-09-10 23:59:33 +00:00
const repliedToText = repliedToMsg?.text || repliedToMsg?.caption;
if (includeRepliedTo && repliedToText) {
2023-09-11 16:03:07 +00:00
// TODO: remove bot command from replied to text
2023-09-16 11:49:12 +00:00
params = parsePngInfo(repliedToText, params);
2023-09-10 23:59:33 +00:00
}
2023-09-11 20:43:12 +00:00
2023-09-16 11:49:12 +00:00
params = parsePngInfo(match, params);
2023-09-11 20:43:12 +00:00
if (ctx.message.text.trim() == "help") {
await ctx.reply(`
Usage: /txt2img <tags> [options]
<>: required
[]: optional
available options:
negative:<tags>
steps:<number>
detail:<number>
size:<number>x<number>
upscale:<multiplier number>
denoising:<number> between 0 and 1 with two decimal places
Example:
/txt2img dog, anthro, woodland steps:40 detail:7 size:500x600 upscale:2 denoising:0.57 negative:badyiffymix, boring_e621_fluffyrock_v4
`, {
reply_to_message_id: ctx.message?.message_id,
});
}
2023-09-10 18:56:17 +00:00
if (!params.prompt) {
2023-09-11 09:22:15 +00:00
await ctx.reply(
2023-09-10 23:59:33 +00:00
"Please tell me what you want to see." +
txt2imgQuestion.messageSuffixMarkdown(),
2023-09-24 19:58:09 +00:00
{
reply_markup: { force_reply: true, selective: true },
parse_mode: "Markdown",
reply_to_message_id: ctx.message?.message_id,
},
2023-09-10 23:59:33 +00:00
);
2023-09-11 09:22:15 +00:00
return;
2023-09-10 18:56:17 +00:00
}
2023-09-10 23:59:33 +00:00
2023-09-24 19:58:09 +00:00
const replyMessage = await ctx.reply("Accepted. You are now in queue.", {
reply_to_message_id: ctx.message?.message_id,
});
2023-09-10 23:59:33 +00:00
2023-09-22 02:59:22 +00:00
await generationQueue.pushJob({
2023-09-12 01:57:44 +00:00
task: { type: "txt2img", params },
from: ctx.message.from,
chat: ctx.message.chat,
2023-09-22 02:59:22 +00:00
requestMessage: ctx.message,
replyMessage: replyMessage,
}, { retryCount: 3, retryDelayMs: 10_000 });
2023-09-10 23:59:33 +00:00
logger().debug(`Generation (txt2img) enqueued for ${formatUserChat(ctx.message)}`);
2023-09-10 18:56:17 +00:00
}