eris/bot/txt2imgCommand.ts

105 lines
3.4 KiB
TypeScript
Raw Normal View History

2023-09-10 23:59:33 +00:00
import { Grammy, GrammyStatelessQ } from "../deps.ts";
2023-09-10 18:56:17 +00:00
import { formatUserChat } from "../utils.ts";
import { jobStore } from "../db/jobStore.ts";
2023-09-12 01:57:44 +00:00
import { getPngInfo, parsePngInfo, PngInfo } from "../sd.ts";
2023-09-10 18:56:17 +00:00
import { Context, logger } from "./mod.ts";
2023-09-11 16:03:07 +00:00
export const txt2imgQuestion = new GrammyStatelessQ.StatelessQuestion<Context>(
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-10 18:56:17 +00:00
export async function txt2imgCommand(ctx: Grammy.CommandContext<Context>) {
2023-09-10 23:59:33 +00:00
await txt2img(ctx, ctx.match, true);
}
async function txt2img(ctx: Context, match: string, includeRepliedTo: boolean): Promise<void> {
if (!ctx.message?.from?.id) {
2023-09-11 09:22:15 +00:00
await ctx.reply("I don't know who you are");
return;
2023-09-10 18:56:17 +00:00
}
2023-09-10 23:59:33 +00:00
2023-09-11 17:08:12 +00:00
if (ctx.session.global.pausedReason != null) {
await ctx.reply(`I'm paused: ${ctx.session.global.pausedReason || "No reason given"}`);
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-10 18:56:17 +00:00
const jobs = await jobStore.getBy("status.type", "waiting");
2023-09-11 17:08:12 +00:00
if (jobs.length >= ctx.session.global.maxJobs) {
2023-09-11 09:22:15 +00:00
await ctx.reply(
2023-09-11 17:08:12 +00:00
`The queue is full. Try again later. (Max queue size: ${ctx.session.global.maxJobs})`,
2023-09-10 18:56:17 +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-12 01:57:44 +00:00
const userJobs = jobs.filter((job) => job.value.from.id === ctx.message?.from?.id);
2023-09-11 17:08:12 +00:00
if (userJobs.length >= ctx.session.global.maxUserJobs) {
2023-09-11 09:22:15 +00:00
await ctx.reply(
2023-09-11 17:08:12 +00:00
`You already have ${ctx.session.global.maxUserJobs} jobs in queue. Try again later.`,
2023-09-10 18:56:17 +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-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());
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"),
};
}
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-10 23:59:33 +00:00
const originalParams = parsePngInfo(repliedToText);
params = {
...originalParams,
...params,
prompt: [originalParams.prompt, params.prompt].filter(Boolean).join("\n"),
2023-09-11 20:43:12 +00:00
negative_prompt: [originalParams.negative_prompt, params.negative_prompt]
.filter(Boolean).join("\n"),
2023-09-10 23:59:33 +00:00
};
}
2023-09-11 20:43:12 +00:00
const messageParams = parsePngInfo(match);
params = {
...params,
...messageParams,
prompt: [params.prompt, messageParams.prompt].filter(Boolean).join("\n"),
};
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(),
{ reply_markup: { force_reply: true, selective: true }, parse_mode: "Markdown" },
);
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
const replyMessage = await ctx.reply("Accepted. You are now in queue.");
2023-09-10 18:56:17 +00:00
await jobStore.create({
2023-09-12 01:57:44 +00:00
task: { type: "txt2img", params },
from: ctx.message.from,
chat: ctx.message.chat,
requestMessageId: ctx.message.message_id,
replyMessageId: replyMessage.message_id,
2023-09-10 18:56:17 +00:00
status: { type: "waiting" },
});
2023-09-10 23:59:33 +00:00
logger().debug(`Job enqueued for ${formatUserChat(ctx.message)}`);
2023-09-10 18:56:17 +00:00
}