Compare commits
2 Commits
5722238c06
...
1c55ae70af
Author | SHA1 | Date |
---|---|---|
pinks | 1c55ae70af | |
pinks | 20b8cb52f2 |
192
bot.ts
192
bot.ts
|
@ -1,83 +1,41 @@
|
||||||
import {
|
import { autoQuote, bold, Bot, Context, hydrateReply, ParseModeFlavor } from "./deps.ts";
|
||||||
autoQuote,
|
import { fmt } from "./intl.ts";
|
||||||
autoRetry,
|
import { getAllJobs, pushJob } from "./queue.ts";
|
||||||
bold,
|
import { mySession, MySessionFlavor } from "./session.ts";
|
||||||
Bot,
|
|
||||||
Context,
|
|
||||||
DenoKVAdapter,
|
|
||||||
fmt,
|
|
||||||
hydrateReply,
|
|
||||||
ParseModeFlavor,
|
|
||||||
session,
|
|
||||||
SessionFlavor,
|
|
||||||
} from "./deps.ts";
|
|
||||||
import { fmtArray, formatOrdinal } from "./intl.ts";
|
|
||||||
import { queue } from "./queue.ts";
|
|
||||||
import { SdRequest } from "./sd.ts";
|
|
||||||
|
|
||||||
type AppContext = ParseModeFlavor<Context> & SessionFlavor<SessionData>;
|
export type MyContext = ParseModeFlavor<Context> & MySessionFlavor;
|
||||||
|
export const bot = new Bot<MyContext>(Deno.env.get("TG_BOT_TOKEN") ?? "");
|
||||||
interface SessionData {
|
|
||||||
global: {
|
|
||||||
adminUsernames: string[];
|
|
||||||
pausedReason: string | null;
|
|
||||||
sdApiUrl: string;
|
|
||||||
maxUserJobs: number;
|
|
||||||
maxJobs: number;
|
|
||||||
defaultParams?: Partial<SdRequest>;
|
|
||||||
};
|
|
||||||
user: {
|
|
||||||
steps: number;
|
|
||||||
detail: number;
|
|
||||||
batchSize: number;
|
|
||||||
};
|
|
||||||
}
|
|
||||||
|
|
||||||
export const bot = new Bot<AppContext>(Deno.env.get("TG_BOT_TOKEN") ?? "");
|
|
||||||
bot.use(autoQuote);
|
bot.use(autoQuote);
|
||||||
bot.use(hydrateReply);
|
bot.use(hydrateReply);
|
||||||
bot.api.config.use(autoRetry({ maxRetryAttempts: 5, maxDelaySeconds: 60 }));
|
bot.use(mySession);
|
||||||
|
|
||||||
const db = await Deno.openKv("./app.db");
|
// Automatically retry bot requests if we get a 429 error
|
||||||
|
bot.api.config.use(async (prev, method, payload, signal) => {
|
||||||
const getDefaultGlobalSession = (): SessionData["global"] => ({
|
let remainingAttempts = 5;
|
||||||
adminUsernames: (Deno.env.get("ADMIN_USERNAMES") ?? "").split(",").filter(Boolean),
|
while (true) {
|
||||||
pausedReason: null,
|
const result = await prev(method, payload, signal);
|
||||||
sdApiUrl: Deno.env.get("SD_API_URL") ?? "http://127.0.0.1:7860/",
|
if (result.ok) return result;
|
||||||
maxUserJobs: 3,
|
if (result.error_code !== 429 || remainingAttempts <= 0) return result;
|
||||||
maxJobs: 20,
|
remainingAttempts -= 1;
|
||||||
defaultParams: {
|
const retryAfterMs = (result.parameters?.retry_after ?? 30) * 1000;
|
||||||
batch_size: 1,
|
await new Promise((resolve) => setTimeout(resolve, retryAfterMs));
|
||||||
n_iter: 1,
|
}
|
||||||
width: 128 * 2,
|
|
||||||
height: 128 * 3,
|
|
||||||
steps: 20,
|
|
||||||
cfg_scale: 9,
|
|
||||||
send_images: true,
|
|
||||||
negative_prompt: "boring_e621_fluffyrock_v4 boring_e621_v4",
|
|
||||||
},
|
|
||||||
});
|
});
|
||||||
|
|
||||||
bot.use(session<SessionData, AppContext>({
|
// if error happened, try to reply to the user with the error
|
||||||
type: "multi",
|
bot.use(async (ctx, next) => {
|
||||||
global: {
|
try {
|
||||||
getSessionKey: () => "global",
|
await next();
|
||||||
initial: getDefaultGlobalSession,
|
} catch (err) {
|
||||||
storage: new DenoKVAdapter(db),
|
try {
|
||||||
},
|
await ctx.reply(`Handling update failed: ${err}`, {
|
||||||
user: {
|
reply_to_message_id: ctx.message?.message_id,
|
||||||
initial: () => ({
|
});
|
||||||
steps: 20,
|
} catch {
|
||||||
detail: 8,
|
throw err;
|
||||||
batchSize: 2,
|
}
|
||||||
}),
|
}
|
||||||
},
|
});
|
||||||
}));
|
|
||||||
|
|
||||||
export async function getGlobalSession(): Promise<SessionData["global"]> {
|
|
||||||
const entry = await db.get<SessionData["global"]>(["sessions", "global"]);
|
|
||||||
return entry.value ?? getDefaultGlobalSession();
|
|
||||||
}
|
|
||||||
|
|
||||||
bot.api.setMyShortDescription("I can generate furry images from text");
|
bot.api.setMyShortDescription("I can generate furry images from text");
|
||||||
bot.api.setMyDescription(
|
bot.api.setMyDescription(
|
||||||
|
@ -99,12 +57,13 @@ bot.command("txt2img", async (ctx) => {
|
||||||
if (config.pausedReason != null) {
|
if (config.pausedReason != null) {
|
||||||
return ctx.reply(`I'm paused: ${config.pausedReason || "No reason given"}`);
|
return ctx.reply(`I'm paused: ${config.pausedReason || "No reason given"}`);
|
||||||
}
|
}
|
||||||
if (queue.length >= config.maxJobs) {
|
const jobs = await getAllJobs();
|
||||||
|
if (jobs.length >= config.maxJobs) {
|
||||||
return ctx.reply(
|
return ctx.reply(
|
||||||
`The queue is full. Try again later. (Max queue size: ${config.maxJobs})`,
|
`The queue is full. Try again later. (Max queue size: ${config.maxJobs})`,
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
const jobCount = queue.filter((job) => job.userId === ctx.from.id).length;
|
const jobCount = jobs.filter((job) => job.user.id === ctx.from.id).length;
|
||||||
if (jobCount >= config.maxUserJobs) {
|
if (jobCount >= config.maxUserJobs) {
|
||||||
return ctx.reply(
|
return ctx.reply(
|
||||||
`You already have ${config.maxUserJobs} jobs in queue. Try again later.`,
|
`You already have ${config.maxUserJobs} jobs in queue. Try again later.`,
|
||||||
|
@ -113,36 +72,50 @@ bot.command("txt2img", async (ctx) => {
|
||||||
if (!ctx.match) {
|
if (!ctx.match) {
|
||||||
return ctx.reply("Please describe what you want to see after the command");
|
return ctx.reply("Please describe what you want to see after the command");
|
||||||
}
|
}
|
||||||
const place = queue.length + 1;
|
pushJob({
|
||||||
const queueMessage = await ctx.reply(`You are ${formatOrdinal(place)} in queue.`);
|
|
||||||
const userName = [ctx.from.first_name, ctx.from.last_name].filter(Boolean).join(" ");
|
|
||||||
const chatName = ctx.chat.type === "supergroup" || ctx.chat.type === "group"
|
|
||||||
? ctx.chat.title
|
|
||||||
: "private chat";
|
|
||||||
queue.push({
|
|
||||||
params: { prompt: ctx.match },
|
params: { prompt: ctx.match },
|
||||||
userId: ctx.from.id,
|
user: ctx.from,
|
||||||
userName,
|
chat: ctx.chat,
|
||||||
chatId: ctx.chat.id,
|
requestMessage: ctx.message,
|
||||||
chatName,
|
status: { type: "idle" },
|
||||||
requestMessageId: ctx.message.message_id,
|
|
||||||
statusMessageId: queueMessage.message_id,
|
|
||||||
});
|
});
|
||||||
console.log(`Enqueued job for ${userName} in chat ${chatName}`);
|
console.log(
|
||||||
|
`Enqueued job ${jobs.length + 1} for ${ctx.from.first_name} in ${ctx.chat.type} chat:`,
|
||||||
|
ctx.match.replace(/\s+/g, " "),
|
||||||
|
"\n",
|
||||||
|
);
|
||||||
});
|
});
|
||||||
|
|
||||||
bot.command("queue", (ctx) => {
|
bot.command("queue", async (ctx) => {
|
||||||
if (queue.length === 0) return ctx.reply("Queue is empty");
|
let jobs = await getAllJobs();
|
||||||
return ctx.replyFmt(
|
const getMessageText = () => {
|
||||||
fmt`Current queue:\n\n${
|
if (jobs.length === 0) return fmt`Queue is empty.`;
|
||||||
fmtArray(
|
const sortedJobs = [];
|
||||||
queue.map((job, index) =>
|
let place = 0;
|
||||||
fmt`${bold(index + 1)}. ${bold(job.userName)} in ${bold(job.chatName)}`
|
for (const job of jobs) {
|
||||||
),
|
if (job.status.type === "idle") place += 1;
|
||||||
"\n",
|
sortedJobs.push({ ...job, place });
|
||||||
|
}
|
||||||
|
return fmt`Current queue:\n\n${
|
||||||
|
sortedJobs.map((job) =>
|
||||||
|
fmt`${job.place}. ${bold(job.user.first_name)} in ${job.chat.type} chat ${
|
||||||
|
job.status.type === "processing" ? `(${(job.status.progress * 100).toFixed(0)}%)` : ""
|
||||||
|
}\n`
|
||||||
)
|
)
|
||||||
}`,
|
}`;
|
||||||
);
|
};
|
||||||
|
const message = await ctx.replyFmt(getMessageText());
|
||||||
|
handleFutureUpdates();
|
||||||
|
async function handleFutureUpdates() {
|
||||||
|
for (let idx = 0; idx < 12; idx++) {
|
||||||
|
await new Promise((resolve) => setTimeout(resolve, 5000));
|
||||||
|
jobs = await getAllJobs();
|
||||||
|
const formattedMessage = getMessageText();
|
||||||
|
await ctx.api.editMessageText(ctx.chat.id, message.message_id, formattedMessage.text, {
|
||||||
|
entities: formattedMessage.entities,
|
||||||
|
}).catch(() => undefined);
|
||||||
|
}
|
||||||
|
}
|
||||||
});
|
});
|
||||||
|
|
||||||
bot.command("pause", (ctx) => {
|
bot.command("pause", (ctx) => {
|
||||||
|
@ -272,14 +245,17 @@ bot.command("setsdparam", (ctx) => {
|
||||||
bot.command("sdparams", (ctx) => {
|
bot.command("sdparams", (ctx) => {
|
||||||
if (!ctx.from?.username) return;
|
if (!ctx.from?.username) return;
|
||||||
const config = ctx.session.global;
|
const config = ctx.session.global;
|
||||||
return ctx.replyFmt(fmt`Current config:\n\n${
|
return ctx.replyFmt(
|
||||||
fmtArray(
|
fmt`Current config:\n\n${
|
||||||
Object.entries(config.defaultParams ?? {}).map(([key, value]) =>
|
Object.entries(config.defaultParams ?? {}).map(([key, value]) =>
|
||||||
fmt`${bold(key)} = ${String(value)}`
|
fmt`${bold(key)} = ${String(value)}\n`
|
||||||
),
|
)
|
||||||
"\n",
|
}`,
|
||||||
)
|
);
|
||||||
}`);
|
});
|
||||||
|
|
||||||
|
bot.command("crash", () => {
|
||||||
|
throw new Error("Crash command used");
|
||||||
});
|
});
|
||||||
|
|
||||||
bot.catch((err) => {
|
bot.catch((err) => {
|
||||||
|
|
4
deps.ts
4
deps.ts
|
@ -2,5 +2,5 @@ export * from "https://deno.land/x/grammy@v1.18.1/mod.ts";
|
||||||
export * from "https://deno.land/x/grammy_autoquote@v1.1.2/mod.ts";
|
export * from "https://deno.land/x/grammy_autoquote@v1.1.2/mod.ts";
|
||||||
export * from "https://deno.land/x/grammy_parse_mode@1.7.1/mod.ts";
|
export * from "https://deno.land/x/grammy_parse_mode@1.7.1/mod.ts";
|
||||||
export * from "https://deno.land/x/grammy_storages@v2.3.1/denokv/src/mod.ts";
|
export * from "https://deno.land/x/grammy_storages@v2.3.1/denokv/src/mod.ts";
|
||||||
export { autoRetry } from "https://esm.sh/@grammyjs/auto-retry@1.1.1";
|
export * as types from "https://deno.land/x/grammy_types@v3.2.0/mod.ts";
|
||||||
export * from "https://deno.land/x/zod/mod.ts";
|
export * from "https://deno.land/x/ulid@v0.3.0/mod.ts";
|
||||||
|
|
44
intl.ts
44
intl.ts
|
@ -8,26 +8,36 @@ export function formatOrdinal(n: number) {
|
||||||
return `${n}th`;
|
return `${n}th`;
|
||||||
}
|
}
|
||||||
|
|
||||||
|
type DeepArray<T> = Array<T | DeepArray<T>>;
|
||||||
|
type StringLikes = DeepArray<FormattedString | string | number | null | undefined>;
|
||||||
|
|
||||||
/**
|
/**
|
||||||
* Like `fmt` from `grammy_parse_mode` but accepts an array instead of template string.
|
* Like `fmt` from `grammy_parse_mode` but additionally accepts arrays.
|
||||||
* @see https://deno.land/x/grammy_parse_mode@1.7.1/format.ts?source=#L182
|
* @see https://deno.land/x/grammy_parse_mode@1.7.1/format.ts?source=#L182
|
||||||
*/
|
*/
|
||||||
export function fmtArray(
|
export const fmt = (
|
||||||
stringLikes: FormattedString[],
|
rawStringParts: TemplateStringsArray | StringLikes,
|
||||||
separator = "",
|
...stringLikes: StringLikes
|
||||||
): FormattedString {
|
): FormattedString => {
|
||||||
let text = "";
|
let text = "";
|
||||||
const entities: ConstructorParameters<typeof FormattedString>[1] = [];
|
const entities: ConstructorParameters<typeof FormattedString>[1][] = [];
|
||||||
for (let i = 0; i < stringLikes.length; i++) {
|
|
||||||
const stringLike = stringLikes[i];
|
const length = Math.max(rawStringParts.length, stringLikes.length);
|
||||||
entities.push(
|
for (let i = 0; i < length; i++) {
|
||||||
...stringLike.entities.map((e) => ({
|
for (let stringLike of [rawStringParts[i], stringLikes[i]]) {
|
||||||
...e,
|
if (Array.isArray(stringLike)) {
|
||||||
offset: e.offset + text.length,
|
stringLike = fmt(stringLike);
|
||||||
})),
|
}
|
||||||
);
|
if (stringLike instanceof FormattedString) {
|
||||||
text += stringLike.toString();
|
entities.push(
|
||||||
if (i < stringLikes.length - 1) text += separator;
|
...stringLike.entities.map((e) => ({
|
||||||
|
...e,
|
||||||
|
offset: e.offset + text.length,
|
||||||
|
})),
|
||||||
|
);
|
||||||
|
}
|
||||||
|
if (stringLike != null) text += stringLike.toString();
|
||||||
|
}
|
||||||
}
|
}
|
||||||
return new FormattedString(text, entities);
|
return new FormattedString(text, entities);
|
||||||
}
|
};
|
||||||
|
|
3
main.ts
3
main.ts
|
@ -1,8 +1,9 @@
|
||||||
import "https://deno.land/std@0.201.0/dotenv/load.ts";
|
import "https://deno.land/std@0.201.0/dotenv/load.ts";
|
||||||
import { bot } from "./bot.ts";
|
import { bot } from "./bot.ts";
|
||||||
import { processQueue } from "./queue.ts";
|
import { processQueue, returnHangedJobs } from "./queue.ts";
|
||||||
|
|
||||||
await Promise.all([
|
await Promise.all([
|
||||||
bot.start(),
|
bot.start(),
|
||||||
processQueue(),
|
processQueue(),
|
||||||
|
returnHangedJobs(),
|
||||||
]);
|
]);
|
||||||
|
|
202
queue.ts
202
queue.ts
|
@ -1,111 +1,167 @@
|
||||||
import { InputFile, InputMediaBuilder } from "./deps.ts";
|
import { InputFile, InputMediaBuilder, types } from "./deps.ts";
|
||||||
import { bot, getGlobalSession } from "./bot.ts";
|
import { bot } from "./bot.ts";
|
||||||
|
import { getGlobalSession } from "./session.ts";
|
||||||
import { formatOrdinal } from "./intl.ts";
|
import { formatOrdinal } from "./intl.ts";
|
||||||
import { SdProgressResponse, SdRequest, txt2img } from "./sd.ts";
|
import { SdRequest, txt2img } from "./sd.ts";
|
||||||
import { extFromMimeType, mimeTypeFromBase64 } from "./mimeType.ts";
|
import { extFromMimeType, mimeTypeFromBase64 } from "./mimeType.ts";
|
||||||
|
import { Model, Store } from "./store.ts";
|
||||||
export const queue: Job[] = [];
|
|
||||||
|
|
||||||
interface Job {
|
interface Job {
|
||||||
params: Partial<SdRequest>;
|
params: Partial<SdRequest>;
|
||||||
userId: number;
|
user: types.User;
|
||||||
userName: string;
|
chat: types.Chat.PrivateChat | types.Chat.GroupChat | types.Chat.SupergroupChat;
|
||||||
chatId: number;
|
requestMessage: types.Message & types.Message.TextMessage;
|
||||||
chatName: string;
|
statusMessage?: types.Message & types.Message.TextMessage;
|
||||||
requestMessageId: number;
|
status: { type: "idle" } | { type: "processing"; progress: number; updatedDate: Date };
|
||||||
statusMessageId: number;
|
}
|
||||||
|
|
||||||
|
const db = await Deno.openKv("./app.db");
|
||||||
|
|
||||||
|
const jobStore = new Store<Job>(db, "job");
|
||||||
|
|
||||||
|
export async function pushJob(job: Job) {
|
||||||
|
await jobStore.create(job);
|
||||||
|
}
|
||||||
|
|
||||||
|
async function takeJob(): Promise<Model<Job> | null> {
|
||||||
|
const jobs = await jobStore.list();
|
||||||
|
const job = jobs.find((job) => job.value.status.type === "idle");
|
||||||
|
if (!job) return null;
|
||||||
|
await job.update({ status: { type: "processing", progress: 0, updatedDate: new Date() } });
|
||||||
|
return job;
|
||||||
|
}
|
||||||
|
|
||||||
|
export async function getAllJobs(): Promise<Array<Job>> {
|
||||||
|
return await jobStore.list().then((jobs) => jobs.map((job) => job.value));
|
||||||
}
|
}
|
||||||
|
|
||||||
export async function processQueue() {
|
export async function processQueue() {
|
||||||
while (true) {
|
while (true) {
|
||||||
const job = queue.shift();
|
const job = await takeJob();
|
||||||
if (!job) {
|
if (!job) {
|
||||||
await new Promise((resolve) => setTimeout(resolve, 1000));
|
await new Promise((resolve) => setTimeout(resolve, 1000));
|
||||||
continue;
|
continue;
|
||||||
}
|
}
|
||||||
for (const [index, job] of queue.entries()) {
|
let place = 0;
|
||||||
const place = index + 1;
|
for (const job of await jobStore.list()) {
|
||||||
await bot.api
|
if (job.value.status.type === "idle") place += 1;
|
||||||
.editMessageText(
|
if (place === 0) continue;
|
||||||
job.chatId,
|
const statusMessageText = `You are ${formatOrdinal(place)} in queue.`;
|
||||||
job.statusMessageId,
|
if (!job.value.statusMessage) {
|
||||||
`You are ${formatOrdinal(place)} in queue.`,
|
await bot.api.sendMessage(job.value.chat.id, statusMessageText, {
|
||||||
|
reply_to_message_id: job.value.requestMessage.message_id,
|
||||||
|
}).catch(() => undefined)
|
||||||
|
.then((message) => job.update({ statusMessage: message }));
|
||||||
|
} else {
|
||||||
|
await bot.api.editMessageText(
|
||||||
|
job.value.chat.id,
|
||||||
|
job.value.statusMessage.message_id,
|
||||||
|
statusMessageText,
|
||||||
)
|
)
|
||||||
.catch(() => {});
|
.catch(() => undefined);
|
||||||
|
}
|
||||||
}
|
}
|
||||||
try {
|
try {
|
||||||
await bot.api
|
if (job.value.statusMessage) {
|
||||||
.deleteMessage(job.chatId, job.statusMessageId)
|
await bot.api
|
||||||
.catch(() => {});
|
.deleteMessage(job.value.chat.id, job.value.statusMessage?.message_id)
|
||||||
const progressMessage = await bot.api.sendMessage(
|
.catch(() => undefined)
|
||||||
job.chatId,
|
.then(() => job.update({ statusMessage: undefined }));
|
||||||
|
}
|
||||||
|
await bot.api.sendMessage(
|
||||||
|
job.value.chat.id,
|
||||||
"Generating your prompt now...",
|
"Generating your prompt now...",
|
||||||
{ reply_to_message_id: job.requestMessageId },
|
{ reply_to_message_id: job.value.requestMessage.message_id },
|
||||||
);
|
).then((message) => job.update({ statusMessage: message }));
|
||||||
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 config = await getGlobalSession();
|
||||||
const response = await txt2img(
|
const response = await txt2img(
|
||||||
config.sdApiUrl,
|
config.sdApiUrl,
|
||||||
{ ...config.defaultParams, ...job.params },
|
{ ...config.defaultParams, ...job.value.params },
|
||||||
onProgress,
|
(progress) => {
|
||||||
|
job.update({
|
||||||
|
status: { type: "processing", progress: progress.progress, updatedDate: new Date() },
|
||||||
|
});
|
||||||
|
if (job.value.statusMessage) {
|
||||||
|
bot.api
|
||||||
|
.editMessageText(
|
||||||
|
job.value.chat.id,
|
||||||
|
job.value.statusMessage.message_id,
|
||||||
|
`Generating your prompt now... ${
|
||||||
|
Math.round(
|
||||||
|
progress.progress * 100,
|
||||||
|
)
|
||||||
|
}%`,
|
||||||
|
)
|
||||||
|
.catch(() => undefined);
|
||||||
|
}
|
||||||
|
},
|
||||||
);
|
);
|
||||||
console.log(
|
console.log(
|
||||||
`Generated ${response.images.length} images (${
|
`Finished job for ${job.value.user.first_name} in ${job.value.chat.type} chat`,
|
||||||
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(
|
if (job.value.statusMessage) {
|
||||||
job.chatId,
|
await bot.api.editMessageText(
|
||||||
progressMessage.message_id,
|
job.value.chat.id,
|
||||||
`Uploading your images...`,
|
job.value.statusMessage.message_id,
|
||||||
).catch(() => {});
|
`Uploading your images...`,
|
||||||
|
).catch(() => undefined);
|
||||||
|
}
|
||||||
const inputFiles = await Promise.all(
|
const inputFiles = await Promise.all(
|
||||||
response.images.map(async (imageBase64, idx) => {
|
response.images.map(async (imageBase64, idx) => {
|
||||||
const mimeType = mimeTypeFromBase64(imageBase64);
|
const mimeType = mimeTypeFromBase64(imageBase64);
|
||||||
const imageBlob = await fetch(`data:${mimeType};base64,${imageBase64}`).then((resp) =>
|
const imageBlob = await fetch(`data:${mimeType};base64,${imageBase64}`)
|
||||||
resp.blob()
|
.then((resp) => resp.blob());
|
||||||
);
|
|
||||||
console.log(
|
|
||||||
`Uploading image ${idx + 1} of ${response.images.length} (${
|
|
||||||
(imageBlob.size / 1024).toFixed(0)
|
|
||||||
}kB)`,
|
|
||||||
);
|
|
||||||
return InputMediaBuilder.photo(
|
return InputMediaBuilder.photo(
|
||||||
new InputFile(imageBlob, `${idx}.${extFromMimeType(mimeType)}`),
|
new InputFile(imageBlob, `image_${idx}.${extFromMimeType(mimeType)}`),
|
||||||
);
|
);
|
||||||
}),
|
}),
|
||||||
);
|
);
|
||||||
await bot.api.sendMediaGroup(job.chatId, inputFiles, {
|
if (job.value.statusMessage) {
|
||||||
reply_to_message_id: job.requestMessageId,
|
await bot.api
|
||||||
|
.deleteMessage(job.value.chat.id, job.value.statusMessage.message_id)
|
||||||
|
.catch(() => undefined).then(() => job.update({ statusMessage: undefined }));
|
||||||
|
}
|
||||||
|
await bot.api.sendMediaGroup(job.value.chat.id, inputFiles, {
|
||||||
|
reply_to_message_id: job.value.requestMessage.message_id,
|
||||||
});
|
});
|
||||||
await bot.api
|
await job.delete();
|
||||||
.deleteMessage(job.chatId, progressMessage.message_id)
|
|
||||||
.catch(() => {});
|
|
||||||
console.log(`${queue.length} jobs remaining`);
|
|
||||||
} catch (err) {
|
} catch (err) {
|
||||||
console.error(
|
console.error(
|
||||||
`Failed to generate image for ${job.userName} in ${job.chatName}: ${job.params.prompt} - ${err}`,
|
`Failed to generate an image for ${job.value.user.first_name} in ${job.value.chat.type} chat: ${err}`,
|
||||||
);
|
);
|
||||||
await bot.api
|
const errorMessage = await bot.api
|
||||||
.sendMessage(job.chatId, err.toString(), {
|
.sendMessage(job.value.chat.id, err.toString(), {
|
||||||
reply_to_message_id: job.requestMessageId,
|
reply_to_message_id: job.value.requestMessage.message_id,
|
||||||
})
|
})
|
||||||
.catch(() => bot.api.sendMessage(job.chatId, err.toString()))
|
.catch(() => undefined);
|
||||||
.catch(() => {});
|
if (errorMessage) {
|
||||||
|
if (job.value.statusMessage) {
|
||||||
|
await bot.api
|
||||||
|
.deleteMessage(job.value.chat.id, job.value.statusMessage.message_id)
|
||||||
|
.catch(() => undefined)
|
||||||
|
.then(() => job.update({ statusMessage: undefined }));
|
||||||
|
}
|
||||||
|
job.update({ status: { type: "idle" } });
|
||||||
|
} else {
|
||||||
|
await job.delete();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
export async function returnHangedJobs() {
|
||||||
|
while (true) {
|
||||||
|
await new Promise((resolve) => setTimeout(resolve, 5000));
|
||||||
|
const jobs = await jobStore.list();
|
||||||
|
for (const job of jobs) {
|
||||||
|
if (job.value.status.type === "idle") continue;
|
||||||
|
// if job wasn't updated for 1 minute, return it to the queue
|
||||||
|
if (job.value.status.updatedDate.getTime() < Date.now() - 60 * 1000) {
|
||||||
|
console.log(
|
||||||
|
`Returning hanged job for ${job.value.user.first_name} in ${job.value.chat.type} chat`,
|
||||||
|
);
|
||||||
|
await job.update({ status: { type: "idle" } });
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
|
@ -0,0 +1,78 @@
|
||||||
|
import { Context, DenoKVAdapter, session, SessionFlavor } from "./deps.ts";
|
||||||
|
import { SdRequest } from "./sd.ts";
|
||||||
|
|
||||||
|
export type MySessionFlavor = SessionFlavor<SessionData>;
|
||||||
|
|
||||||
|
export interface SessionData {
|
||||||
|
global: GlobalData;
|
||||||
|
chat: ChatData;
|
||||||
|
user: UserData;
|
||||||
|
}
|
||||||
|
|
||||||
|
export interface GlobalData {
|
||||||
|
adminUsernames: string[];
|
||||||
|
pausedReason: string | null;
|
||||||
|
sdApiUrl: string;
|
||||||
|
maxUserJobs: number;
|
||||||
|
maxJobs: number;
|
||||||
|
defaultParams?: Partial<SdRequest>;
|
||||||
|
}
|
||||||
|
|
||||||
|
export interface ChatData {
|
||||||
|
language: string;
|
||||||
|
}
|
||||||
|
|
||||||
|
export interface UserData {
|
||||||
|
steps: number;
|
||||||
|
detail: number;
|
||||||
|
batchSize: number;
|
||||||
|
}
|
||||||
|
|
||||||
|
const globalDb = await Deno.openKv("./app.db");
|
||||||
|
|
||||||
|
const globalDbAdapter = new DenoKVAdapter<GlobalData>(globalDb);
|
||||||
|
|
||||||
|
const getDefaultGlobalData = (): GlobalData => ({
|
||||||
|
adminUsernames: (Deno.env.get("ADMIN_USERNAMES") ?? "").split(",").filter(Boolean),
|
||||||
|
pausedReason: null,
|
||||||
|
sdApiUrl: Deno.env.get("SD_API_URL") ?? "http://127.0.0.1:7860/",
|
||||||
|
maxUserJobs: 3,
|
||||||
|
maxJobs: 20,
|
||||||
|
defaultParams: {
|
||||||
|
batch_size: 1,
|
||||||
|
n_iter: 1,
|
||||||
|
width: 128 * 2,
|
||||||
|
height: 128 * 3,
|
||||||
|
steps: 20,
|
||||||
|
cfg_scale: 9,
|
||||||
|
send_images: true,
|
||||||
|
negative_prompt: "boring_e621_fluffyrock_v4 boring_e621_v4",
|
||||||
|
},
|
||||||
|
});
|
||||||
|
|
||||||
|
export const mySession = session<SessionData, Context & MySessionFlavor>({
|
||||||
|
type: "multi",
|
||||||
|
global: {
|
||||||
|
getSessionKey: () => "global",
|
||||||
|
initial: getDefaultGlobalData,
|
||||||
|
storage: globalDbAdapter,
|
||||||
|
},
|
||||||
|
chat: {
|
||||||
|
initial: () => ({
|
||||||
|
language: "en",
|
||||||
|
}),
|
||||||
|
},
|
||||||
|
user: {
|
||||||
|
getSessionKey: (ctx) => ctx.from?.id.toFixed(),
|
||||||
|
initial: () => ({
|
||||||
|
steps: 20,
|
||||||
|
detail: 8,
|
||||||
|
batchSize: 2,
|
||||||
|
}),
|
||||||
|
},
|
||||||
|
});
|
||||||
|
|
||||||
|
export async function getGlobalSession(): Promise<GlobalData> {
|
||||||
|
const data = await globalDbAdapter.read("global");
|
||||||
|
return data ?? getDefaultGlobalData();
|
||||||
|
}
|
|
@ -0,0 +1,76 @@
|
||||||
|
import { ulid } from "./deps.ts";
|
||||||
|
|
||||||
|
export class Store<T extends object> {
|
||||||
|
constructor(
|
||||||
|
private readonly db: Deno.Kv,
|
||||||
|
private readonly storeKey: Deno.KvKeyPart,
|
||||||
|
) {
|
||||||
|
}
|
||||||
|
|
||||||
|
async create(value: T): Promise<Model<T>> {
|
||||||
|
const id = ulid();
|
||||||
|
await this.db.set([this.storeKey, id], value);
|
||||||
|
return new Model(this.db, this.storeKey, id, value);
|
||||||
|
}
|
||||||
|
|
||||||
|
async get(id: Deno.KvKeyPart): Promise<Model<T> | null> {
|
||||||
|
const entry = await this.db.get<T>([this.storeKey, id]);
|
||||||
|
if (entry.versionstamp == null) return null;
|
||||||
|
return new Model(this.db, this.storeKey, id, entry.value);
|
||||||
|
}
|
||||||
|
|
||||||
|
async list(): Promise<Array<Model<T>>> {
|
||||||
|
const models: Array<Model<T>> = [];
|
||||||
|
for await (const entry of this.db.list<T>({ prefix: [this.storeKey] })) {
|
||||||
|
models.push(new Model(this.db, this.storeKey, entry.key[1], entry.value));
|
||||||
|
}
|
||||||
|
return models;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
export class Model<T extends object> {
|
||||||
|
#value: T;
|
||||||
|
|
||||||
|
constructor(
|
||||||
|
private readonly db: Deno.Kv,
|
||||||
|
private readonly storeKey: Deno.KvKeyPart,
|
||||||
|
private readonly entryKey: Deno.KvKeyPart,
|
||||||
|
value: T,
|
||||||
|
) {
|
||||||
|
this.#value = value;
|
||||||
|
}
|
||||||
|
|
||||||
|
get value(): T {
|
||||||
|
return this.#value;
|
||||||
|
}
|
||||||
|
|
||||||
|
async get(): Promise<T | null> {
|
||||||
|
const entry = await this.db.get<T>([this.storeKey, this.entryKey]);
|
||||||
|
if (entry.versionstamp == null) return null;
|
||||||
|
this.#value = entry.value;
|
||||||
|
return entry.value;
|
||||||
|
}
|
||||||
|
|
||||||
|
async set(value: T): Promise<T> {
|
||||||
|
await this.db.set([this.storeKey, this.entryKey], value);
|
||||||
|
this.#value = value;
|
||||||
|
return value;
|
||||||
|
}
|
||||||
|
|
||||||
|
async update(value: Partial<T> | ((value: T) => T)): Promise<T | null> {
|
||||||
|
const entry = await this.db.get<T>([this.storeKey, this.entryKey]);
|
||||||
|
if (entry.versionstamp == null) return null;
|
||||||
|
if (typeof value === "function") {
|
||||||
|
entry.value = value(entry.value);
|
||||||
|
} else {
|
||||||
|
entry.value = { ...entry.value, ...value };
|
||||||
|
}
|
||||||
|
await this.db.set([this.storeKey, this.entryKey], entry.value);
|
||||||
|
this.#value = entry.value;
|
||||||
|
return entry.value;
|
||||||
|
}
|
||||||
|
|
||||||
|
async delete(): Promise<void> {
|
||||||
|
await this.db.delete([this.storeKey, this.entryKey]);
|
||||||
|
}
|
||||||
|
}
|
Loading…
Reference in New Issue