import { CommandContext } from "grammy"; import { StatelessQuestion } from "grammy_stateless_question"; import { debug } from "std/log/mod.ts"; import { getConfig } from "../app/config.ts"; import { generationQueue } from "../app/generationQueue.ts"; import { formatUserChat } from "../utils/formatUserChat.ts"; import { ErisContext } from "./mod.ts"; import { getPngInfo, parsePngInfo, PngInfo } from "./parsePngInfo.ts"; import { adminStore } from "../app/adminStore.ts"; export const txt2imgQuestion = new StatelessQuestion( "txt2img", async (ctx) => { if (!ctx.message.text) return; await txt2img(ctx, ctx.message.text, false); }, ); export async function txt2imgCommand(ctx: CommandContext) { await txt2img(ctx, ctx.match, true); } async function txt2img(ctx: ErisContext, match: string, includeRepliedTo: boolean): Promise { if (!ctx.from || !ctx.message) { return; } if (ctx.from.is_bot) { return; } const config = await getConfig(); let priority = 0; if (config.pausedReason != null) { await ctx.reply(`I'm paused: ${config.pausedReason || "No reason given"}`, { reply_to_message_id: ctx.message?.message_id, }); return; } const [admin] = await adminStore.getBy("tgUserId", { value: ctx.from.id }); if (admin) { priority = 1; } else { const jobs = await generationQueue.getAllJobs(); if (jobs.length >= config.maxJobs) { await ctx.reply(`The queue is full. Try again later. (Max queue size: ${config.maxJobs})`, { reply_to_message_id: ctx.message?.message_id, }); return; } const userJobs = jobs.filter((job) => job.state.from.id === ctx.message?.from?.id); if (userJobs.length >= config.maxUserJobs) { await ctx.reply(`You already have ${userJobs.length} jobs in the queue. Try again later.`, { reply_to_message_id: ctx.message?.message_id, }); return; } } let params: Partial = {}; const repliedToMsg = ctx.message.reply_to_message; 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()); params = parsePngInfo(getPngInfo(buffer) ?? "", params); } const repliedToText = repliedToMsg?.text || repliedToMsg?.caption; const isReply = includeRepliedTo && repliedToText; if (isReply) { // TODO: remove bot command from replied to text params = parsePngInfo(repliedToText, params); } params = parsePngInfo(match, params, true); if (isReply) { const parsedInfo = parsePngInfo(repliedToText, undefined, true); if (parsedInfo.prompt !== params.prompt) { params.seed = parsedInfo.seed ?? -1; } } if (!params.prompt) { await ctx.reply( "Please tell me what you want to see." + txt2imgQuestion.messageSuffixMarkdown(), { reply_markup: { force_reply: true, selective: true }, parse_mode: "Markdown", reply_to_message_id: ctx.message?.message_id, }, ); return; } const replyMessage = await ctx.reply("Accepted. You are now in queue.", { reply_to_message_id: ctx.message?.message_id, }); await generationQueue.pushJob({ task: { type: "txt2img", params }, from: ctx.message.from, chat: ctx.message.chat, requestMessage: ctx.message, replyMessage: replyMessage, }, { retryCount: 3, retryDelayMs: 10_000, priority: priority }); debug(`Generation (txt2img) enqueued for ${formatUserChat(ctx.message)}`); }