forked from pinks/eris
1
0
Fork 0
eris/queue.ts

112 lines
3.5 KiB
TypeScript

import { InputFile, InputMediaBuilder } from "./deps.ts";
import { bot, getGlobalSession } from "./bot.ts";
import { formatOrdinal } from "./intl.ts";
import { SdProgressResponse, SdRequest, txt2img } from "./sd.ts";
import { extFromMimeType, mimeTypeFromBase64 } from "./mimeType.ts";
export const queue: Job[] = [];
interface Job {
params: Partial<SdRequest>;
userId: number;
userName: string;
chatId: number;
chatName: string;
requestMessageId: number;
statusMessageId: number;
}
export async function processQueue() {
while (true) {
const job = queue.shift();
if (!job) {
await new Promise((resolve) => setTimeout(resolve, 1000));
continue;
}
for (const [index, job] of queue.entries()) {
const place = index + 1;
await bot.api
.editMessageText(
job.chatId,
job.statusMessageId,
`You are ${formatOrdinal(place)} in queue.`,
)
.catch(() => {});
}
try {
await bot.api
.deleteMessage(job.chatId, job.statusMessageId)
.catch(() => {});
const progressMessage = await bot.api.sendMessage(
job.chatId,
"Generating your prompt now...",
{ reply_to_message_id: job.requestMessageId },
);
const onProgress = (progress: SdProgressResponse) => {
bot.api
.editMessageText(
job.chatId,
progressMessage.message_id,
`Generating your prompt now... ${
Math.round(
progress.progress * 100,
)
}%`,
)
.catch(() => {});
};
const config = await getGlobalSession();
const response = await txt2img(
config.sdApiUrl,
{ ...config.defaultParams, ...job.params },
onProgress,
);
console.log(
`Generated ${response.images.length} images (${
response.images
.map((image) => (image.length / 1024).toFixed(0) + "kB")
.join(", ")
}) for ${job.userName} in ${job.chatName}: ${job.params.prompt?.replace(/\s+/g, " ")}`,
);
await bot.api.editMessageText(
job.chatId,
progressMessage.message_id,
`Uploading your images...`,
).catch(() => {});
const inputFiles = await Promise.all(
response.images.map(async (imageBase64, idx) => {
const mimeType = mimeTypeFromBase64(imageBase64);
const imageBlob = await fetch(`data:${mimeType};base64,${imageBase64}`).then((resp) =>
resp.blob()
);
console.log(
`Uploading image ${idx + 1} of ${response.images.length} (${
(imageBlob.size / 1024).toFixed(0)
}kB)`,
);
return InputMediaBuilder.photo(
new InputFile(imageBlob, `${idx}.${extFromMimeType(mimeType)}`),
);
}),
);
await bot.api.sendMediaGroup(job.chatId, inputFiles, {
reply_to_message_id: job.requestMessageId,
});
await bot.api
.deleteMessage(job.chatId, progressMessage.message_id)
.catch(() => {});
console.log(`${queue.length} jobs remaining`);
} catch (err) {
console.error(
`Failed to generate image for ${job.userName} in ${job.chatName}: ${job.params.prompt} - ${err}`,
);
await bot.api
.sendMessage(job.chatId, err.toString(), {
reply_to_message_id: job.requestMessageId,
})
.catch(() => bot.api.sendMessage(job.chatId, err.toString()))
.catch(() => {});
}
}
}