forked from pinks/eris
1
0
Fork 0

Compare commits

..

58 Commits
main ... main

Author SHA1 Message Date
lisq 6ccf8a04d8 chore: update deno.json 2024-03-02 17:51:58 +01:00
lisq 02e3c22b83 chore: remove deno.lock 2024-03-02 17:44:30 +01:00
lisq 03b0c2bc89 perf: faster queue status updates 2024-03-02 17:42:36 +01:00
pinks 5b6a1a3471 refactor: rewrite API to Elysia (#25)
https://elysiajs.com/
Reviewed-on: pinks/eris#25
Co-authored-by: pinks <lisq@cock.li>
Co-committed-by: pinks <lisq@cock.li>
2023-11-20 02:14:14 +00:00
pinks ffd3d0dc8d fix: priority for img2img 2023-11-18 00:50:34 +01:00
pinks ec4a3c893e fix: clearing paused reason 2023-11-18 00:49:39 +01:00
pinks e7d292df60 fix type error 2023-11-17 19:22:35 +01:00
pinks fad14f685e feat: show position in upload queue 2023-11-17 19:00:56 +01:00
nameless 2f62c17e32 review fixes 2023-11-12 02:33:35 +00:00
nameless 602a63f2c9 feat: user daily stats 2023-11-11 04:13:26 +03:00
nameless 8ced57c175 refactor: move sendChatAction
resolves pinks/eris#18
2023-11-11 04:13:26 +03:00
nameless d22848691c feat: remove url from error 2023-11-11 04:13:26 +03:00
nameless af7b370321 feat: tagcount endpoint
removes tagCountMap from user endpoint, returns it only via admin-only endpoint
2023-11-11 04:13:26 +03:00
nameless 5b61576315 fix: withUser -> withAdmin 2023-11-11 04:13:26 +03:00
nameless e270a3ab1f refactor: optimize jobs endpoint output 2023-11-11 04:13:26 +03:00
nameless 84ad11b709 refactor: webui adjustments
- remove unnecessary margin in admins and workers page
- use rounded-md instead of rounded-xl for queue page
- output "No worker/admin/jobs" as item in a list
2023-11-11 04:13:26 +03:00
nameless ac63e39373 feat: admin priority
ignore queue limits and assign higher priority on admin jobs
2023-11-11 04:13:26 +03:00
nameless 73810d3204 fix: use exifreader for png info
- replace pnginfo module
- allow jpeg for tag extraction
- rename pnginfo to getprompt
2023-11-11 04:13:26 +03:00
nameless f4e7b5e1a7 change txt2img seed behaviour
will ignore seed only on /txt2img or new requests
2023-11-11 04:13:15 +03:00
pinks f678324889 add badges to readme 2023-11-06 18:38:41 +01:00
pinks 6d6174ab48 silence getUpdates error 2023-10-27 21:50:20 +02:00
pinks a3bada0da2 sort most used tags 2023-10-27 21:49:46 +02:00
pinks 3a13d5c140 fix: slow queue position updates 2023-10-23 02:40:58 +02:00
pinks dfa94e219d increase upload concurrency 2023-10-23 02:39:19 +02:00
pinks a35f07b036 feat: managing admins in webui 2023-10-23 02:39:01 +02:00
pinks f7b29dd150 update deps 2023-10-21 12:40:42 +02:00
pinks 4f2371fa8b exactOptionalPropertyTypes 2023-10-19 23:37:03 +02:00
pinks f020499f4d dialog animations 2023-10-19 03:18:56 +02:00
pinks 11d8a66c18 update README 2023-10-18 23:11:39 +02:00
pinks 22d9c518be add env var for db path 2023-10-18 23:11:25 +02:00
pinks 4ba6f494d2 sort config 2023-10-18 23:08:18 +02:00
pinks 2665fa1c02 add robots.txt 2023-10-17 15:03:27 +02:00
pinks 083f6bc01c simplify api routes 2023-10-17 15:03:14 +02:00
Vargoshi 2f059ceaff fix errors 2023-10-15 21:47:08 +02:00
Vargoshi cd3b53018a fix: ignore other bots (#22)
fixes #5

Reviewed-on: pinks/eris#22
Co-authored-by: Vargoshi <marcinflisak1998@gmail.com>
Co-committed-by: Vargoshi <marcinflisak1998@gmail.com>
2023-10-15 19:41:41 +00:00
pinks 5a241ff13c sort import map 2023-10-15 21:14:24 +02:00
pinks 2a3674b0c4 allow configuring log level 2023-10-15 21:13:38 +02:00
pinks 843f6d9103 log upload type and bytes 2023-10-15 13:12:31 +02:00
nameless 07fc0c71f2 feat: parse seed on new txt2img request (#20)
Reviewed-on: pinks/eris#20
Co-authored-by: nameless <nameless@noreply@foxo.me>
Co-committed-by: nameless <nameless@noreply@foxo.me>
2023-10-14 10:57:46 +00:00
Vargoshi d9f8966ad8 fix: login bot url 2023-10-13 21:18:19 +02:00
Vargoshi 5bb2a13f35 fix windows path 2023-10-13 20:56:16 +02:00
pinks 2258a484e1 commit deno.lock 2023-10-13 20:05:23 +02:00
pinks f0a570c1b9 perf: parallel tg update processing 2023-10-13 18:10:16 +02:00
pinks 0b85c99c92 feat: manage workers via webui 2023-10-13 13:47:57 +02:00
pinks a251d5e965 fix README 2023-10-12 11:39:19 +02:00
pinks 6313c5d67e feat: workers api 2023-10-11 03:59:52 +02:00
pinks 55c53ac565 fix html 2023-10-11 03:56:08 +02:00
pinks 477cb03ea2 add sampler to settings page 2023-10-10 18:21:46 +02:00
pinks 0e88bc5053 feat: home page with counters 2023-10-10 18:21:25 +02:00
pinks bf75cce20c feat: stats api 2023-10-09 21:03:31 +02:00
pinks a8d36db20b update rest api 2023-10-08 23:23:54 +02:00
pinks c0354ef679 add catalan language 2023-10-08 20:23:31 +02:00
pinks 99540d4035 feat: admin ui 2023-10-05 11:00:51 +02:00
pinks adebd34db8 faster progress refresh 2023-10-01 23:47:29 +02:00
pinks 01261b6f3a fix: language flags 2023-09-26 13:25:13 +02:00
pinks 82fdd0f21c feat: webui initial impl 2023-09-26 12:43:36 +02:00
pinks 70a9b1180d fix: sd gen request crash 2023-09-26 11:52:11 +02:00
pinks 7ebdf6eae4 fix: ignore edit message errors 2023-09-26 00:01:50 +02:00
60 changed files with 3378 additions and 340 deletions

4
.gitignore vendored
View File

@ -1,4 +1,4 @@
.env .env
app.db* *.db
*.db-*
deno.lock deno.lock
updateConfig.ts

View File

@ -1,5 +1,11 @@
# Eris the Bot # Eris the Bot
[![Website](https://img.shields.io/website?url=https%3A%2F%2Feris.lisq.eu%2F)](https://eris.lisq.eu/)
![Unique users](https://img.shields.io/badge/dynamic/json?url=https%3A%2F%2Feris.lisq.eu%2Fapi%2Fstats&query=%24.userCount&label=unique%20users)
![Generated images](https://img.shields.io/badge/dynamic/json?url=https%3A%2F%2Feris.lisq.eu%2Fapi%2Fstats&query=%24.imageCount&label=images%20generated)
![Processed steps](https://img.shields.io/badge/dynamic/json?url=https%3A%2F%2Feris.lisq.eu%2Fapi%2Fstats&query=%24.stepCount&label=steps%20processed)
![Painted pixels](https://img.shields.io/badge/dynamic/json?url=https%3A%2F%2Feris.lisq.eu%2Fapi%2Fstats&query=%24.pixelCount&label=pixels%20painted)
Telegram bot for generating images from text. Telegram bot for generating images from text.
## Requirements ## Requirements
@ -13,21 +19,27 @@ You can put these in `.env` file or pass them as environment variables.
- `TG_BOT_TOKEN` - Telegram bot token. Get yours from [@BotFather](https://t.me/BotFather). - `TG_BOT_TOKEN` - Telegram bot token. Get yours from [@BotFather](https://t.me/BotFather).
Required. Required.
- `SD_API_URL` - URL to Stable Diffusion API. Only used on first run. Default: - `DENO_KV_PATH` - [Deno KV](https://deno.land/api?s=Deno.openKv&unstable) database file path. A
`http://127.0.0.1:7860/` temporary file is used by default.
- `TG_ADMIN_USERS` - Comma separated list of usernames of users that can use admin commands. Only - `LOG_LEVEL` - [Log level](https://deno.land/std@0.201.0/log/mod.ts?s=LogLevels). Default: `INFO`.
used on first run. Optional.
## Running ## Running
- Start stable diffusion webui: `cd sd-webui`, `./webui.sh --api` 1. Start Eris: `deno task start`
- Start bot: `deno task start` 2. Visit [Eris WebUI](http://localhost:5999/) and login via Telegram.
3. Promote yourself to admin in the Eris WebUI.
4. Start Stable Diffusion WebUI: `./webui.sh --api` (in SD WebUI directory)
5. Add a new worker in the Eris WebUI.
## Codegen ## Codegen
The Stable Diffusion API in `sd/sdApi.ts` is auto-generated. To regenerate it, first start your SD The Stable Diffusion API types are auto-generated. To regenerate them, first start your SD WebUI
WebUI with `--nowebui --api`, and then run: with `--nowebui --api`, and then run `deno task generate`
```sh ## Project structure
deno run npm:openapi-typescript http://localhost:7861/openapi.json -o sd/sdApi.ts
``` - `/api` - Eris API served at `http://localhost:5999/api/`.
- `/app` - Queue handling and other core processes.
- `/bot` - Handling bot commands and other updates from Telegram API.
- `/ui` - Eris WebUI frontend files served at `http://localhost:5999/`.
- `/util` - Utility functions shared by other parts.

128
api/adminsRoute.ts Normal file
View File

@ -0,0 +1,128 @@
import { info } from "std/log/mod.ts";
import { Elysia, Static, t } from "elysia";
import { Admin, adminSchema, adminStore } from "../app/adminStore.ts";
import { getUser, withSessionAdmin, withSessionUser } from "./getUser.ts";
import { TkvEntry } from "../utils/Tkv.ts";
const adminDataSchema = t.Intersect([adminSchema, t.Object({ tgUserId: t.Number() })]);
export type AdminData = Static<typeof adminDataSchema>;
function getAdminData(adminEntry: TkvEntry<["admins", number], Admin>): AdminData {
return { tgUserId: adminEntry.key[1], ...adminEntry.value };
}
export const adminsRoute = new Elysia()
.get(
"",
async () => {
const adminEntries = await Array.fromAsync(adminStore.list({ prefix: ["admins"] }));
const admins = adminEntries.map(getAdminData);
return admins;
},
{
response: t.Array(adminDataSchema),
},
)
.post(
"",
async ({ query, body, set }) => {
return withSessionAdmin({ query, set }, async (sessionUser, sessionAdminEntry) => {
const newAdminUser = await getUser(body.tgUserId);
const newAdminKey = ["admins", body.tgUserId] as const;
const newAdminValue = { promotedBy: sessionAdminEntry.key[1] };
const newAdminResult = await adminStore.atomicSet(newAdminKey, null, newAdminValue);
if (!newAdminResult.ok) {
set.status = 409;
return "User is already an admin";
}
info(`User ${sessionUser.first_name} promoted user ${newAdminUser.first_name} to admin`);
return getAdminData({ ...newAdminResult, key: newAdminKey, value: newAdminValue });
});
},
{
query: t.Object({ sessionId: t.String() }),
body: t.Object({
tgUserId: t.Number(),
}),
response: {
200: adminDataSchema,
401: t.Literal("Must be logged in"),
403: t.Literal("Must be an admin"),
409: t.Literal("User is already an admin"),
},
},
)
.post(
"/promote_self",
// if there are no admins, allow any user to promote themselves
async ({ query, set }) => {
return withSessionUser({ query, set }, async (sessionUser) => {
const adminEntries = await Array.fromAsync(adminStore.list({ prefix: ["admins"] }));
if (adminEntries.length !== 0) {
set.status = 409;
return "You are not allowed to promote yourself";
}
const newAdminKey = ["admins", sessionUser.id] as const;
const newAdminValue = { promotedBy: null };
const newAdminResult = await adminStore.set(newAdminKey, newAdminValue);
info(`User ${sessionUser.first_name} promoted themselves to admin`);
return getAdminData({ ...newAdminResult, key: newAdminKey, value: newAdminValue });
});
},
{
query: t.Object({ sessionId: t.String() }),
response: {
200: adminDataSchema,
401: t.Literal("Must be logged in"),
403: t.Literal("Must be an admin"),
409: t.Literal("You are not allowed to promote yourself"),
},
},
)
.get(
"/:adminId",
async ({ params, set }) => {
const adminEntry = await adminStore.get(["admins", Number(params.adminId)]);
if (!adminEntry.versionstamp) {
set.status = 404;
return "Admin not found";
}
return getAdminData(adminEntry);
},
{
params: t.Object({ adminId: t.String() }),
response: {
200: adminDataSchema,
404: t.Literal("Admin not found"),
},
},
)
.delete(
"/:adminId",
async ({ params, query, set }) => {
return withSessionAdmin({ query, set }, async (sessionUser) => {
const deletedAdminEntry = await adminStore.get(["admins", Number(params.adminId)]);
if (!deletedAdminEntry.versionstamp) {
set.status = 404;
return "Admin not found";
}
const deletedAdminUser = await getUser(deletedAdminEntry.key[1]);
await adminStore.delete(["admins", Number(params.adminId)]);
info(
`User ${sessionUser.first_name} demoted user ${deletedAdminUser.first_name} from admin`,
);
return null;
});
},
{
params: t.Object({ adminId: t.String() }),
query: t.Object({ sessionId: t.String() }),
response: {
200: t.Null(),
401: t.Literal("Must be logged in"),
403: t.Literal("Must be an admin"),
404: t.Literal("Admin not found"),
},
},
);

14
api/botRoute.ts Normal file
View File

@ -0,0 +1,14 @@
import { Elysia, t } from "elysia";
import { bot } from "../bot/mod.ts";
export const botRoute = new Elysia()
.get(
"",
async () => {
const username = bot.botInfo.username;
return { username };
},
{
response: t.Object({ username: t.String() }),
},
);

45
api/getUser.ts Normal file
View File

@ -0,0 +1,45 @@
import { Chat } from "grammy_types";
import { Admin, adminStore } from "../app/adminStore.ts";
import { bot } from "../bot/mod.ts";
import { sessions } from "./sessionsRoute.ts";
import { TkvEntry } from "../utils/Tkv.ts";
export async function withSessionUser<O>(
{ query, set }: { query: { sessionId: string }; set: { status?: string | number } },
cb: (sessionUser: Chat.PrivateGetChat) => Promise<O>,
) {
const session = sessions.get(query.sessionId);
if (!session?.userId) {
set.status = 401;
return "Must be logged in";
}
const user = await getUser(session.userId);
return cb(user);
}
export async function withSessionAdmin<O>(
{ query, set }: { query: { sessionId: string }; set: { status?: string | number } },
cb: (
sessionUser: Chat.PrivateGetChat,
sessionAdminEntry: TkvEntry<["admins", number], Admin>,
) => Promise<O>,
) {
const session = sessions.get(query.sessionId);
if (!session?.userId) {
set.status = 401;
return "Must be logged in";
}
const sessionUser = await getUser(session.userId);
const sessionAdminEntry = await adminStore.get(["admins", sessionUser.id]);
if (!sessionAdminEntry.versionstamp) {
set.status = 403;
return "Must be an admin";
}
return cb(sessionUser, sessionAdminEntry);
}
export async function getUser(userId: number): Promise<Chat.PrivateGetChat> {
const chat = await bot.api.getChat(userId);
if (chat.type !== "private") throw new Error("Chat is not private");
return chat;
}

40
api/jobsRoute.ts Normal file
View File

@ -0,0 +1,40 @@
import { Elysia, t } from "elysia";
import { generationQueue } from "../app/generationQueue.ts";
export const jobsRoute = new Elysia()
.get(
"",
async () => {
const allJobs = await generationQueue.getAllJobs();
return allJobs.map((job) => ({
id: job.id.join(":"),
place: job.place,
state: {
from: {
language_code: job.state.from.language_code ?? null,
first_name: job.state.from.first_name,
last_name: job.state.from.last_name ?? null,
username: job.state.from.username ?? null,
},
progress: job.state.progress ?? null,
workerInstanceKey: job.state.workerInstanceKey ?? null,
},
}));
},
{
response: t.Array(t.Object({
id: t.String(),
place: t.Number(),
state: t.Object({
from: t.Object({
language_code: t.Nullable(t.String()),
first_name: t.String(),
last_name: t.Nullable(t.String()),
username: t.Nullable(t.String()),
}),
progress: t.Nullable(t.Number()),
workerInstanceKey: t.Nullable(t.String()),
}),
})),
},
);

23
api/mod.ts Normal file
View File

@ -0,0 +1,23 @@
import { route } from "reroute";
import { serveSpa } from "serve_spa";
import { api } from "./serveApi.ts";
import { fromFileUrl } from "std/path/mod.ts";
export async function serveUi() {
const server = Deno.serve({ port: 5999 }, (request) =>
route(request, {
"/api/*": (request) => api.fetch(request),
"/*": (request) =>
serveSpa(request, {
fsRoot: fromFileUrl(new URL("../ui/", import.meta.url)),
indexFallback: true,
importMapFile: "../deno.json",
aliasMap: {
"/utils/*": "../utils/",
},
log: (_request, response) => response.status >= 400,
}),
}));
await server.finished;
}

39
api/paramsRoute.ts Normal file
View File

@ -0,0 +1,39 @@
import { Elysia, t } from "elysia";
import { info } from "std/log/mod.ts";
import { defaultParamsSchema, getConfig, setConfig } from "../app/config.ts";
import { withSessionAdmin } from "./getUser.ts";
export const paramsRoute = new Elysia()
.get(
"",
async () => {
const config = await getConfig();
return config.defaultParams;
},
{
response: {
200: defaultParamsSchema,
},
},
)
.patch(
"",
async ({ query, body, set }) => {
return withSessionAdmin({ query, set }, async (user) => {
const config = await getConfig();
info(`User ${user.first_name} updated default params: ${JSON.stringify(body)}`);
const defaultParams = { ...config.defaultParams, ...body };
await setConfig({ defaultParams });
return config.defaultParams;
});
},
{
query: t.Object({ sessionId: t.String() }),
body: defaultParamsSchema,
response: {
200: defaultParamsSchema,
401: t.Literal("Must be logged in"),
403: t.Literal("Must be an admin"),
},
},
);

32
api/serveApi.ts Normal file
View File

@ -0,0 +1,32 @@
import { Elysia } from "elysia";
import { swagger } from "elysia/swagger";
import { adminsRoute } from "./adminsRoute.ts";
import { botRoute } from "./botRoute.ts";
import { jobsRoute } from "./jobsRoute.ts";
import { paramsRoute } from "./paramsRoute.ts";
import { sessionsRoute } from "./sessionsRoute.ts";
import { statsRoute } from "./statsRoute.ts";
import { usersRoute } from "./usersRoute.ts";
import { workersRoute } from "./workersRoute.ts";
export const api = new Elysia()
.use(
swagger({
path: "/docs",
swaggerOptions: { url: "docs/json" } as never,
documentation: {
info: { title: "Eris API", version: "0.1" },
servers: [{ url: "/api" }],
},
}),
)
.group("/admins", (api) => api.use(adminsRoute))
.group("/bot", (api) => api.use(botRoute))
.group("/jobs", (api) => api.use(jobsRoute))
.group("/sessions", (api) => api.use(sessionsRoute))
.group("/settings/params", (api) => api.use(paramsRoute))
.group("/stats", (api) => api.use(statsRoute))
.group("/users", (api) => api.use(usersRoute))
.group("/workers", (api) => api.use(workersRoute));
export type Api = typeof api;

43
api/sessionsRoute.ts Normal file
View File

@ -0,0 +1,43 @@
import { Elysia, NotFoundError, t } from "elysia";
import { ulid } from "ulid";
export const sessions = new Map<string, Session>();
export interface Session {
userId?: number | undefined;
}
export const sessionsRoute = new Elysia()
.post(
"",
async () => {
const id = ulid();
const session: Session = {};
sessions.set(id, session);
return { id, userId: session.userId ?? null };
},
{
response: t.Object({
id: t.String(),
userId: t.Nullable(t.Number()),
}),
},
)
.get(
"/:sessionId",
async ({ params }) => {
const id = params.sessionId!;
const session = sessions.get(id);
if (!session) {
throw new NotFoundError("Session not found");
}
return { id, userId: session.userId ?? null };
},
{
params: t.Object({ sessionId: t.String() }),
response: t.Object({
id: t.String(),
userId: t.Nullable(t.Number()),
}),
},
);

134
api/statsRoute.ts Normal file
View File

@ -0,0 +1,134 @@
import { Elysia, t } from "elysia";
import { subMinutes } from "date-fns";
import { dailyStatsSchema, getDailyStats } from "../app/dailyStatsStore.ts";
import { generationStore } from "../app/generationStore.ts";
import { globalStats } from "../app/globalStats.ts";
import { getUserDailyStats, userDailyStatsSchema } from "../app/userDailyStatsStore.ts";
import { getUserStats, userStatsSchema } from "../app/userStatsStore.ts";
import { withSessionAdmin } from "./getUser.ts";
const STATS_INTERVAL_MIN = 3;
export const statsRoute = new Elysia()
.get(
"",
async () => {
const after = subMinutes(new Date(), STATS_INTERVAL_MIN);
const generations = await generationStore.getAll({ after });
const imagesPerMinute = generations.length / STATS_INTERVAL_MIN;
const stepsPerMinute = generations
.map((generation) => generation.value.info?.steps ?? 0)
.reduce((sum, steps) => sum + steps, 0) / STATS_INTERVAL_MIN;
const pixelsPerMinute = generations
.map((generation) =>
(generation.value.info?.width ?? 0) * (generation.value.info?.height ?? 0)
)
.reduce((sum, pixels) => sum + pixels, 0) / STATS_INTERVAL_MIN;
const pixelStepsPerMinute = generations
.map((generation) =>
(generation.value.info?.width ?? 0) * (generation.value.info?.height ?? 0) *
(generation.value.info?.steps ?? 0)
)
.reduce((sum, pixelSteps) => sum + pixelSteps, 0) / STATS_INTERVAL_MIN;
return {
imageCount: globalStats.imageCount,
stepCount: globalStats.stepCount,
pixelCount: globalStats.pixelCount,
pixelStepCount: globalStats.pixelStepCount,
userCount: globalStats.userIds.length,
imagesPerMinute,
stepsPerMinute,
pixelsPerMinute,
pixelStepsPerMinute,
};
},
{
response: {
200: t.Object({
imageCount: t.Number(),
stepCount: t.Number(),
pixelCount: t.Number(),
pixelStepCount: t.Number(),
userCount: t.Number(),
imagesPerMinute: t.Number(),
stepsPerMinute: t.Number(),
pixelsPerMinute: t.Number(),
pixelStepsPerMinute: t.Number(),
}),
},
},
)
.get(
"/daily/:year/:month/:day",
async ({ params }) => {
return getDailyStats(params.year, params.month, params.day);
},
{
params: t.Object({
year: t.Number(),
month: t.Number(),
day: t.Number(),
}),
response: {
200: dailyStatsSchema,
},
},
)
.get(
"/users/:userId",
async ({ params }) => {
const userId = params.userId;
// deno-lint-ignore no-unused-vars
const { tagCountMap, ...stats } = await getUserStats(userId);
return stats;
},
{
params: t.Object({ userId: t.Number() }),
response: {
200: t.Omit(userStatsSchema, ["tagCountMap"]),
},
},
)
.get(
"/users/:userId/tagcount",
async ({ params, query, set }) => {
return withSessionAdmin({ query, set }, async () => {
const stats = await getUserStats(params.userId);
return {
tagCountMap: stats.tagCountMap,
timestamp: stats.timestamp,
};
});
},
{
params: t.Object({ userId: t.Number() }),
query: t.Object({ sessionId: t.String() }),
response: {
200: t.Pick(userStatsSchema, ["tagCountMap", "timestamp"]),
401: t.Literal("Must be logged in"),
403: t.Literal("Must be an admin"),
},
},
)
.get(
"/users/:userId/daily/:year/:month/:day",
async ({ params }) => {
return getUserDailyStats(params.userId, params.year, params.month, params.day);
},
{
params: t.Object({
userId: t.Number(),
year: t.Number(),
month: t.Number(),
day: t.Number(),
}),
response: {
200: userDailyStatsSchema,
},
},
);

55
api/usersRoute.ts Normal file
View File

@ -0,0 +1,55 @@
import { Elysia, t } from "elysia";
import { adminSchema, adminStore } from "../app/adminStore.ts";
import { bot } from "../bot/mod.ts";
import { getUser } from "./getUser.ts";
export const usersRoute = new Elysia()
.get(
"/:userId/photo",
async ({ params }) => {
const user = await getUser(Number(params.userId));
if (!user.photo) {
throw new Error("User has no photo");
}
const photoFile = await bot.api.getFile(user.photo.small_file_id);
const photoData = await fetch(
`https://api.telegram.org/file/bot${bot.token}/${photoFile.file_path}`,
).then((resp) => {
if (!resp.ok) {
throw new Error("Failed to fetch photo");
}
return resp;
}).then((resp) => resp.arrayBuffer());
return new Response(new File([photoData], "avatar.jpg", { type: "image/jpeg" }));
},
{
params: t.Object({ userId: t.String() }),
},
)
.get(
"/:userId",
async ({ params }) => {
const user = await getUser(Number(params.userId));
const adminEntry = await adminStore.get(["admins", user.id]);
return {
id: user.id,
first_name: user.first_name,
last_name: user.last_name ?? null,
username: user.username ?? null,
bio: user.bio ?? null,
admin: adminEntry.value ?? null,
};
},
{
params: t.Object({ userId: t.String() }),
response: t.Object({
id: t.Number(),
first_name: t.String(),
last_name: t.Nullable(t.String()),
username: t.Nullable(t.String()),
bio: t.Nullable(t.String()),
admin: t.Nullable(adminSchema),
}),
},
);

248
api/workersRoute.ts Normal file
View File

@ -0,0 +1,248 @@
import { subMinutes } from "date-fns";
import { Model } from "indexed_kv";
import createOpenApiFetch from "openapi_fetch";
import { info } from "std/log/mod.ts";
import { Elysia, NotFoundError, Static, t } from "elysia";
import { activeGenerationWorkers } from "../app/generationQueue.ts";
import { generationStore } from "../app/generationStore.ts";
import * as SdApi from "../app/sdApi.ts";
import {
WorkerInstance,
workerInstanceSchema,
workerInstanceStore,
} from "../app/workerInstanceStore.ts";
import { getAuthHeader } from "../utils/getAuthHeader.ts";
import { omitUndef } from "../utils/omitUndef.ts";
import { withSessionAdmin } from "./getUser.ts";
const workerResponseSchema = t.Intersect([
t.Object({ id: t.String() }),
t.Omit(workerInstanceSchema, ["sdUrl", "sdAuth"]),
t.Object({
isActive: t.Boolean(),
imagesPerMinute: t.Number(),
stepsPerMinute: t.Number(),
pixelsPerMinute: t.Number(),
pixelStepsPerMinute: t.Number(),
}),
]);
export type WorkerResponse = Static<typeof workerResponseSchema>;
const workerRequestSchema = t.Omit(workerInstanceSchema, ["lastOnlineTime", "lastError"]);
const STATS_INTERVAL_MIN = 10;
async function getWorkerResponse(workerInstance: Model<WorkerInstance>): Promise<WorkerResponse> {
const after = subMinutes(new Date(), STATS_INTERVAL_MIN);
const generations = await generationStore.getBy("workerInstanceKey", {
value: workerInstance.value.key,
after: after,
});
const imagesPerMinute = generations.length / STATS_INTERVAL_MIN;
const stepsPerMinute = generations
.map((generation) => generation.value.info?.steps ?? 0)
.reduce((sum, steps) => sum + steps, 0) / STATS_INTERVAL_MIN;
const pixelsPerMinute = generations
.map((generation) => (generation.value.info?.width ?? 0) * (generation.value.info?.height ?? 0))
.reduce((sum, pixels) => sum + pixels, 0) / STATS_INTERVAL_MIN;
const pixelStepsPerMinute = generations
.map((generation) =>
(generation.value.info?.width ?? 0) * (generation.value.info?.height ?? 0) *
(generation.value.info?.steps ?? 0)
)
.reduce((sum, pixelSteps) => sum + pixelSteps, 0) / STATS_INTERVAL_MIN;
return omitUndef({
id: workerInstance.id,
key: workerInstance.value.key,
name: workerInstance.value.name,
lastError: workerInstance.value.lastError,
lastOnlineTime: workerInstance.value.lastOnlineTime,
isActive: activeGenerationWorkers.get(workerInstance.id)?.isProcessing ?? false,
imagesPerMinute,
stepsPerMinute,
pixelsPerMinute,
pixelStepsPerMinute,
});
}
export const workersRoute = new Elysia()
.get(
"",
async () => {
const workerInstances = await workerInstanceStore.getAll();
const workers = await Promise.all(workerInstances.map(getWorkerResponse));
return workers;
},
{
response: t.Array(workerResponseSchema),
},
)
.post(
"",
async ({ query, body, set }) => {
return withSessionAdmin({ query, set }, async (sessionUser) => {
const workerInstance = await workerInstanceStore.create(body);
info(`User ${sessionUser.first_name} created worker ${workerInstance.value.name}`);
return await getWorkerResponse(workerInstance);
});
},
{
query: t.Object({ sessionId: t.String() }),
body: workerRequestSchema,
response: {
200: workerResponseSchema,
401: t.Literal("Must be logged in"),
403: t.Literal("Must be an admin"),
},
},
)
.get(
"/:workerId",
async ({ params }) => {
const workerInstance = await workerInstanceStore.getById(params.workerId);
if (!workerInstance) {
throw new NotFoundError("Worker not found");
}
return await getWorkerResponse(workerInstance);
},
{
params: t.Object({ workerId: t.String() }),
response: {
200: workerResponseSchema,
},
},
)
.patch(
"/:workerId",
async ({ params, query, body, set }) => {
const workerInstance = await workerInstanceStore.getById(params.workerId);
if (!workerInstance) {
throw new NotFoundError("Worker not found");
}
return withSessionAdmin({ query, set }, async (sessionUser) => {
info(
`User ${sessionUser.first_name} updated worker ${workerInstance.value.name}: ${
JSON.stringify(body)
}`,
);
await workerInstance.update(body);
return await getWorkerResponse(workerInstance);
});
},
{
params: t.Object({ workerId: t.String() }),
query: t.Object({ sessionId: t.String() }),
body: t.Partial(workerRequestSchema),
response: {
200: workerResponseSchema,
401: t.Literal("Must be logged in"),
403: t.Literal("Must be an admin"),
},
},
)
.delete(
"/:workerId",
async ({ params, query, set }) => {
const workerInstance = await workerInstanceStore.getById(params.workerId);
if (!workerInstance) {
throw new Error("Worker not found");
}
return withSessionAdmin({ query, set }, async (sessionUser) => {
info(`User ${sessionUser.first_name} deleted worker ${workerInstance.value.name}`);
await workerInstance.delete();
return null;
});
},
{
params: t.Object({ workerId: t.String() }),
query: t.Object({ sessionId: t.String() }),
response: {
200: t.Null(),
401: t.Literal("Must be logged in"),
403: t.Literal("Must be an admin"),
},
},
)
.get(
"/:workerId/loras",
async ({ params }) => {
const workerInstance = await workerInstanceStore.getById(params.workerId);
if (!workerInstance) {
throw new NotFoundError("Worker not found");
}
const sdClient = createOpenApiFetch<SdApi.paths>({
baseUrl: workerInstance.value.sdUrl,
headers: getAuthHeader(workerInstance.value.sdAuth),
});
const lorasResponse = await sdClient.GET("/sdapi/v1/loras", {});
if (lorasResponse.error) {
throw new Error(
`Loras request failed: ${lorasResponse["error"]}`,
);
}
const loras = (lorasResponse.data as Lora[]).map((lora) => ({
name: lora.name,
alias: lora.alias ?? null,
}));
return loras;
},
{
params: t.Object({ workerId: t.String() }),
response: t.Array(
t.Object({
name: t.String(),
alias: t.Nullable(t.String()),
}),
),
},
)
.get(
"/:workerId/models",
async ({ params }) => {
const workerInstance = await workerInstanceStore.getById(params.workerId);
if (!workerInstance) {
throw new NotFoundError("Worker not found");
}
const sdClient = createOpenApiFetch<SdApi.paths>({
baseUrl: workerInstance.value.sdUrl,
headers: getAuthHeader(workerInstance.value.sdAuth),
});
const modelsResponse = await sdClient.GET("/sdapi/v1/sd-models", {});
if (modelsResponse.error) {
throw new Error(
`Models request failed: ${modelsResponse["error"]}`,
);
}
const models = modelsResponse.data.map((model) => ({
title: model.title,
modelName: model.model_name,
hash: model.hash ?? null,
sha256: model.sha256 ?? null,
}));
return models;
},
{
params: t.Object({ workerId: t.String() }),
response: t.Array(
t.Object({
title: t.String(),
modelName: t.String(),
hash: t.Nullable(t.String()),
sha256: t.Nullable(t.String()),
}),
),
},
);
export interface Lora {
name: string;
alias: string | null;
metadata: object;
}

11
app/adminStore.ts Normal file
View File

@ -0,0 +1,11 @@
import { Static, t } from "elysia";
import { db } from "./db.ts";
import { Tkv } from "../utils/Tkv.ts";
export const adminSchema = t.Object({
promotedBy: t.Nullable(t.Number()),
});
export type Admin = Static<typeof adminSchema>;
export const adminStore = new Tkv<["admins", number], Admin>(db);

View File

@ -1,53 +1,45 @@
import * as SdApi from "../sd/sdApi.ts"; import { Static, t } from "elysia";
import { Tkv } from "../utils/Tkv.ts";
import { db } from "./db.ts"; import { db } from "./db.ts";
export interface ConfigData { export const defaultParamsSchema = t.Partial(t.Object({
adminUsernames: string[]; batch_size: t.Number(),
pausedReason: string | null; n_iter: t.Number(),
maxUserJobs: number; width: t.Number(),
maxJobs: number; height: t.Number(),
defaultParams?: Partial< steps: t.Number(),
| SdApi.components["schemas"]["StableDiffusionProcessingTxt2Img"] cfg_scale: t.Number(),
| SdApi.components["schemas"]["StableDiffusionProcessingImg2Img"] sampler_name: t.String(),
>; negative_prompt: t.String(),
sdInstances: SdInstanceData[]; }));
}
export interface SdInstanceData { export type DefaultParams = Static<typeof defaultParamsSchema>;
id: string;
name?: string;
api: { url: string; auth?: string };
maxResolution: number;
}
const getDefaultConfig = (): ConfigData => ({ export const configSchema = t.Object({
adminUsernames: Deno.env.get("TG_ADMIN_USERS")?.split(",") ?? [], pausedReason: t.Nullable(t.String()),
pausedReason: null, maxUserJobs: t.Number(),
maxUserJobs: 3, maxJobs: t.Number(),
maxJobs: 20, defaultParams: defaultParamsSchema,
defaultParams: {
batch_size: 1,
n_iter: 1,
width: 512,
height: 768,
steps: 30,
cfg_scale: 10,
negative_prompt: "boring_e621_fluffyrock_v4 boring_e621_v4",
},
sdInstances: [
{
id: "local",
api: { url: Deno.env.get("SD_API_URL") ?? "http://127.0.0.1:7860/" },
maxResolution: 1024 * 1024,
},
],
}); });
export async function getConfig(): Promise<ConfigData> { export type Config = Static<typeof configSchema>;
const configEntry = await db.get<ConfigData>(["config"]);
return configEntry.value ?? getDefaultConfig(); export const configStore = new Tkv<["config"], Config>(db);
const defaultConfig: Config = {
pausedReason: null,
maxUserJobs: Infinity,
maxJobs: Infinity,
defaultParams: {},
};
export async function getConfig(): Promise<Config> {
const configEntry = await configStore.get(["config"]);
return { ...defaultConfig, ...configEntry.value };
} }
export async function setConfig(config: ConfigData): Promise<void> { export async function setConfig<K extends keyof Config>(newConfig: Pick<Config, K>): Promise<void> {
await db.set(["config"], config); const configEntry = await configStore.get(["config"]);
const config = { ...defaultConfig, ...configEntry.value, ...newConfig };
await configStore.atomicSet(["config"], configEntry.versionstamp, config);
} }

69
app/dailyStatsStore.ts Normal file
View File

@ -0,0 +1,69 @@
import { hoursToMilliseconds, isSameDay, minutesToMilliseconds } from "date-fns";
import { UTCDateMini } from "date-fns/utc";
import { Static, t } from "elysia";
import { info } from "std/log/mod.ts";
import { db } from "./db.ts";
import { generationStore } from "./generationStore.ts";
import { kvMemoize } from "../utils/kvMemoize.ts";
export const dailyStatsSchema = t.Object({
userIds: t.Array(t.Number()),
imageCount: t.Number(),
stepCount: t.Number(),
pixelCount: t.Number(),
pixelStepCount: t.Number(),
timestamp: t.Number(),
});
export type DailyStats = Static<typeof dailyStatsSchema>;
export const getDailyStats = kvMemoize(
db,
["dailyStats"],
async (year: number, month: number, day: number): Promise<DailyStats> => {
const userIdSet = new Set<number>();
let imageCount = 0;
let stepCount = 0;
let pixelCount = 0;
let pixelStepCount = 0;
const after = new Date(Date.UTC(year, month - 1, day));
const before = new Date(Date.UTC(year, month - 1, day + 1));
info(`Calculating daily stats for ${year}-${month}-${day}`);
for await (
const generation of generationStore.listAll({ after, before })
) {
userIdSet.add(generation.value.from.id);
imageCount++;
stepCount += generation.value.info?.steps ?? 0;
pixelCount += (generation.value.info?.width ?? 0) * (generation.value.info?.height ?? 0);
pixelStepCount += (generation.value.info?.width ?? 0) *
(generation.value.info?.height ?? 0) *
(generation.value.info?.steps ?? 0);
}
return {
userIds: [...userIdSet],
imageCount,
stepCount,
pixelCount,
pixelStepCount,
timestamp: Date.now(),
};
},
{
// expire in 1 minute if was calculated on the same day, otherwise 7-14 days.
expireIn: (result, year, month, day) => {
const requestDate = new UTCDateMini(year, month - 1, day);
const calculatedDate = new UTCDateMini(result.timestamp);
return isSameDay(requestDate, calculatedDate)
? minutesToMilliseconds(1)
: hoursToMilliseconds(24 * 7 + Math.random() * 24 * 7);
},
// should cache if the stats are non-zero
shouldCache: (result) =>
result.userIds.length > 0 || result.imageCount > 0 || result.pixelCount > 0,
},
);

View File

@ -1,4 +1,4 @@
import { KvFs } from "kvfs"; import { KvFs } from "kvfs";
export const db = await Deno.openKv("./app.db"); export const db = await Deno.openKv(Deno.env.get("DENO_KV_PATH"));
export const fs = new KvFs(db); export const fs = new KvFs(db);

View File

@ -2,22 +2,23 @@ import { promiseState } from "async";
import { Chat, Message, User } from "grammy_types"; import { Chat, Message, User } from "grammy_types";
import { JobData, Queue, Worker } from "kvmq"; import { JobData, Queue, Worker } from "kvmq";
import createOpenApiClient from "openapi_fetch"; import createOpenApiClient from "openapi_fetch";
import { delay } from "std/async"; import { delay } from "std/async/delay.ts";
import { decode, encode } from "std/encoding/base64"; import { decode, encode } from "std/encoding/base64.ts";
import { getLogger } from "std/log"; import { debug, error, info } from "std/log/mod.ts";
import { ulid } from "ulid"; import { ulid } from "ulid";
import { bot } from "../bot/mod.ts"; import { bot } from "../bot/mod.ts";
import { SdError } from "../sd/SdError.ts"; import { PngInfo } from "../bot/parsePngInfo.ts";
import { PngInfo } from "../sd/parsePngInfo.ts";
import * as SdApi from "../sd/sdApi.ts";
import { formatOrdinal } from "../utils/formatOrdinal.ts"; import { formatOrdinal } from "../utils/formatOrdinal.ts";
import { formatUserChat } from "../utils/formatUserChat.ts"; import { formatUserChat } from "../utils/formatUserChat.ts";
import { getConfig, SdInstanceData } from "./config.ts"; import { getAuthHeader } from "../utils/getAuthHeader.ts";
import { omitUndef } from "../utils/omitUndef.ts";
import { SdError } from "./SdError.ts";
import { getConfig } from "./config.ts";
import { db, fs } from "./db.ts"; import { db, fs } from "./db.ts";
import { SdGenerationInfo } from "./generationStore.ts"; import { SdGenerationInfo } from "./generationStore.ts";
import * as SdApi from "./sdApi.ts";
import { uploadQueue } from "./uploadQueue.ts"; import { uploadQueue } from "./uploadQueue.ts";
import { workerInstanceStore } from "./workerInstanceStore.ts";
const logger = () => getLogger();
interface GenerationJob { interface GenerationJob {
task: task:
@ -34,7 +35,7 @@ interface GenerationJob {
chat: Chat; chat: Chat;
requestMessage: Message; requestMessage: Message;
replyMessage: Message; replyMessage: Message;
sdInstanceId?: string; workerInstanceKey?: string;
progress?: number; progress?: number;
} }
@ -45,17 +46,17 @@ export const activeGenerationWorkers = new Map<string, Worker<GenerationJob>>();
/** /**
* Initializes queue workers for each SD instance when they become online. * Initializes queue workers for each SD instance when they become online.
*/ */
export async function processGenerationQueue() { export async function processGenerationQueue(): Promise<never> {
while (true) { while (true) {
const config = await getConfig(); for await (const workerInstance of workerInstanceStore.listAll()) {
const activeWorker = activeGenerationWorkers.get(workerInstance.id);
for (const sdInstance of config.sdInstances) { if (activeWorker?.isProcessing) {
const activeWorker = activeGenerationWorkers.get(sdInstance.id); continue;
if (activeWorker?.isProcessing) continue; }
const workerSdClient = createOpenApiClient<SdApi.paths>({ const workerSdClient = createOpenApiClient<SdApi.paths>({
baseUrl: sdInstance.api.url, baseUrl: workerInstance.value.sdUrl,
headers: { "Authorization": sdInstance.api.auth }, headers: getAuthHeader(workerInstance.value.sdAuth),
}); });
// check if worker is up // check if worker is up
@ -69,18 +70,23 @@ export async function processGenerationQueue() {
return response; return response;
}) })
.catch((error) => { .catch((error) => {
logger().warning(`Worker ${sdInstance.id} is down: ${error}`); const cleanedErrorMessage = error.message.replace(/url \([^)]+\)/, "");
workerInstance.update({ lastError: { message: cleanedErrorMessage, time: Date.now() } })
.catch(() => undefined);
debug(`Worker ${workerInstance.value.key} is down: ${error}`);
}); });
if (!activeWorkerStatus?.data) { if (!activeWorkerStatus?.data) {
continue; continue;
} }
// create worker // create worker
const newWorker = generationQueue.createWorker(async ({ state }, updateJob) => { const newWorker = generationQueue.createWorker(async ({ state }, updateJob) => {
await processGenerationJob(state, updateJob, sdInstance); await processGenerationJob(state, updateJob, workerInstance.id);
}); });
newWorker.addEventListener("error", (e) => { newWorker.addEventListener("error", (e) => {
logger().error( error(
`Generation failed for ${formatUserChat(e.detail.job.state)}: ${e.detail.error}`, `Generation failed for ${formatUserChat(e.detail.job.state)}: ${e.detail.error}`,
); );
bot.api.sendMessage( bot.api.sendMessage(
@ -94,13 +100,20 @@ export async function processGenerationQueue() {
allow_sending_without_reply: true, allow_sending_without_reply: true,
}, },
).catch(() => undefined); ).catch(() => undefined);
if (e.detail.error instanceof SdError) { newWorker.stopProcessing();
newWorker.stopProcessing(); workerInstance.update({ lastError: { message: e.detail.error.message, time: Date.now() } })
} .catch(() => undefined);
info(`Stopped worker ${workerInstance.value.key}`);
}); });
newWorker.addEventListener("complete", () => {
workerInstance.update({ lastOnlineTime: Date.now() }).catch(() => undefined);
});
await workerInstance.update({ lastOnlineTime: Date.now() });
newWorker.processJobs(); newWorker.processJobs();
activeGenerationWorkers.set(sdInstance.id, newWorker); activeGenerationWorkers.set(workerInstance.id, newWorker);
logger().info(`Started worker ${sdInstance.id}`); info(`Started worker ${workerInstance.value.key}`);
} }
await delay(60_000); await delay(60_000);
} }
@ -112,39 +125,37 @@ export async function processGenerationQueue() {
async function processGenerationJob( async function processGenerationJob(
state: GenerationJob, state: GenerationJob,
updateJob: (job: Partial<JobData<GenerationJob>>) => Promise<void>, updateJob: (job: Partial<JobData<GenerationJob>>) => Promise<void>,
sdInstance: SdInstanceData, workerInstanceId: string,
) { ) {
const startDate = new Date(); const startDate = new Date();
const workerSdClient = createOpenApiClient<SdApi.paths>({ const config = await getConfig();
baseUrl: sdInstance.api.url, const workerInstance = await workerInstanceStore.getById(workerInstanceId);
headers: { "Authorization": sdInstance.api.auth }, if (!workerInstance) {
}); throw new Error(`Unknown workerInstanceId: ${workerInstanceId}`);
state.sdInstanceId = sdInstance.id;
logger().debug(`Generation started for ${formatUserChat(state)}`);
await updateJob({ state: state });
// check if bot can post messages in this chat
const chat = await bot.api.getChat(state.chat.id);
if (
(chat.type === "group" || chat.type === "supergroup") &&
(!chat.permissions?.can_send_messages || !chat.permissions?.can_send_photos)
) {
throw new Error("Bot doesn't have permissions to send photos in this chat");
} }
const workerSdClient = createOpenApiClient<SdApi.paths>({
baseUrl: workerInstance.value.sdUrl,
headers: getAuthHeader(workerInstance.value.sdAuth),
});
state.workerInstanceKey = workerInstance.value.key;
state.progress = 0;
debug(`Generation started for ${formatUserChat(state)}`);
await updateJob({ state: state });
// edit the existing status message // edit the existing status message
await bot.api.editMessageText( await bot.api.editMessageText(
state.replyMessage.chat.id, state.replyMessage.chat.id,
state.replyMessage.message_id, state.replyMessage.message_id,
`Generating your prompt now... 0% using ${sdInstance.name || sdInstance.id}`, `Generating your prompt now... 0% using ${
workerInstance.value.name || workerInstance.value.key
}`,
{ maxAttempts: 1 }, { maxAttempts: 1 },
); ).catch(() => undefined);
// reduce size if worker can't handle the resolution // reduce size if worker can't handle the resolution
const config = await getConfig();
const size = limitSize( const size = limitSize(
{ ...config.defaultParams, ...state.task.params }, omitUndef({ ...config.defaultParams, ...state.task.params }),
sdInstance.maxResolution, 1024 * 1024,
); );
function limitSize( function limitSize(
{ width, height }: { width?: number; height?: number }, { width, height }: { width?: number; height?: number },
@ -164,18 +175,18 @@ async function processGenerationJob(
// start generating the image // start generating the image
const responsePromise = state.task.type === "txt2img" const responsePromise = state.task.type === "txt2img"
? workerSdClient.POST("/sdapi/v1/txt2img", { ? workerSdClient.POST("/sdapi/v1/txt2img", {
body: { body: omitUndef({
...config.defaultParams, ...config.defaultParams,
...state.task.params, ...state.task.params,
...size, ...size,
negative_prompt: state.task.params.negative_prompt negative_prompt: state.task.params.negative_prompt
? state.task.params.negative_prompt ? state.task.params.negative_prompt
: config.defaultParams?.negative_prompt, : config.defaultParams?.negative_prompt,
}, }),
}) })
: state.task.type === "img2img" : state.task.type === "img2img"
? workerSdClient.POST("/sdapi/v1/img2img", { ? workerSdClient.POST("/sdapi/v1/img2img", {
body: { body: omitUndef({
...config.defaultParams, ...config.defaultParams,
...state.task.params, ...state.task.params,
...size, ...size,
@ -191,7 +202,7 @@ async function processGenerationJob(
).then((resp) => resp.arrayBuffer()), ).then((resp) => resp.arrayBuffer()),
), ),
], ],
}, }),
}) })
: undefined; : undefined;
@ -199,7 +210,12 @@ async function processGenerationJob(
throw new Error(`Unknown task type: ${state.task.type}`); throw new Error(`Unknown task type: ${state.task.type}`);
} }
// we await the promise only after it finishes
// so we need to add catch callback to not crash the process before that
responsePromise.catch(() => undefined);
// poll for progress while the generation request is pending // poll for progress while the generation request is pending
do { do {
const progressResponse = await workerSdClient.GET("/sdapi/v1/progress", { const progressResponse = await workerSdClient.GET("/sdapi/v1/progress", {
params: {}, params: {},
@ -213,21 +229,20 @@ async function processGenerationJob(
); );
} }
state.progress = progressResponse.data.progress; if (progressResponse.data.progress > state.progress) {
await updateJob({ state: state }); state.progress = progressResponse.data.progress;
await updateJob({ state: state });
await bot.api.editMessageText(
state.replyMessage.chat.id,
state.replyMessage.message_id,
`Generating your prompt now... ${
(progressResponse.data.progress * 100).toFixed(0)
}% using ${workerInstance.value.name || workerInstance.value.key}`,
{ maxAttempts: 1 },
).catch(() => undefined);
}
await bot.api.sendChatAction(state.chat.id, "upload_photo", { maxAttempts: 1 }) await Promise.race([delay(2_000), responsePromise]).catch(() => undefined);
.catch(() => undefined);
await bot.api.editMessageText(
state.replyMessage.chat.id,
state.replyMessage.message_id,
`Generating your prompt now... ${(progressResponse.data.progress * 100).toFixed(0)}% using ${
sdInstance.name || sdInstance.id
}`,
{ maxAttempts: 1 },
).catch(() => undefined);
await Promise.race([delay(3000), responsePromise]).catch(() => undefined);
} while (await promiseState(responsePromise) === "pending"); } while (await promiseState(responsePromise) === "pending");
// check response // check response
@ -257,48 +272,56 @@ async function processGenerationJob(
from: state.from, from: state.from,
requestMessage: state.requestMessage, requestMessage: state.requestMessage,
replyMessage: state.replyMessage, replyMessage: state.replyMessage,
sdInstanceId: sdInstance.id, workerInstanceKey: workerInstance.value.key,
startDate, startDate,
endDate: new Date(), endDate: new Date(),
imageKeys, imageKeys,
info, info,
}, { retryCount: 5, retryDelayMs: 10000 }); }, { retryCount: 5, retryDelayMs: 10000 });
const uploadQueueSize = await uploadQueue.getAllJobs().then((jobs) =>
jobs.filter((job) => job.status !== "processing").length
);
// change status message to uploading images // change status message to uploading images
await bot.api.editMessageText( await bot.api.editMessageText(
state.replyMessage.chat.id, state.replyMessage.chat.id,
state.replyMessage.message_id, state.replyMessage.message_id,
`Uploading your images...`, uploadQueueSize > 10
? `You are ${formatOrdinal(uploadQueueSize)} in upload queue.`
: `Uploading your images...`,
{ maxAttempts: 1 }, { maxAttempts: 1 },
).catch(() => undefined); ).catch(() => undefined);
logger().debug(`Generation finished for ${formatUserChat(state)}`); // send upload photo action
await bot.api.sendChatAction(state.chat.id, "upload_photo", { maxAttempts: 1 })
.catch(() => undefined);
debug(`Generation finished for ${formatUserChat(state)}`);
} }
/** /**
* Handles queue updates and updates the status message. * Handles queue updates and updates the status message.
*/ */
export async function updateGenerationQueue() { export async function updateGenerationQueue(): Promise<never> {
while (true) { while (true) {
const jobs = await generationQueue.getAllJobs(); const jobs = await generationQueue.getAllJobs();
let index = 0;
for (const job of jobs) { await Promise.all(jobs.map(async (job) => {
if (job.lockUntil > new Date()) { if (job.status === "processing") {
// job is currently being processed, the worker will update its status message // if the job is processing, the worker will update its status message
continue; return;
} }
if (!job.state.replyMessage) {
// no status message, nothing to update return await bot.api.editMessageText(
continue;
}
index++;
await bot.api.editMessageText(
job.state.replyMessage.chat.id, job.state.replyMessage.chat.id,
job.state.replyMessage.message_id, job.state.replyMessage.message_id,
`You are ${formatOrdinal(index)} in queue.`, `You are ${formatOrdinal(job.place)} in queue.`,
{ maxAttempts: 1 }, { maxAttempts: 1 },
).catch(() => undefined); ).catch(() => {});
} }));
await delay(3000);
// delay between updates based on the number of jobs
await delay(Math.min(15_000, Math.max(3_000, jobs.length * 100)));
} }
} }

View File

@ -5,10 +5,10 @@ import { db } from "./db.ts";
export interface GenerationSchema { export interface GenerationSchema {
from: User; from: User;
chat: Chat; chat: Chat;
sdInstanceId?: string; sdInstanceId?: string | undefined; // TODO: change to workerInstanceKey
info?: SdGenerationInfo; info?: SdGenerationInfo | undefined;
startDate?: Date; startDate?: Date | undefined;
endDate?: Date; endDate?: Date | undefined;
} }
/** /**
@ -48,6 +48,7 @@ export interface SdGenerationInfo {
type GenerationIndices = { type GenerationIndices = {
fromId: number; fromId: number;
chatId: number; chatId: number;
workerInstanceKey: string;
}; };
export const generationStore = new Store<GenerationSchema, GenerationIndices>( export const generationStore = new Store<GenerationSchema, GenerationIndices>(
@ -57,6 +58,7 @@ export const generationStore = new Store<GenerationSchema, GenerationIndices>(
indices: { indices: {
fromId: { getValue: (item) => item.from.id }, fromId: { getValue: (item) => item.from.id },
chatId: { getValue: (item) => item.chat.id }, chatId: { getValue: (item) => item.chat.id },
workerInstanceKey: { getValue: (item) => item.sdInstanceId ?? "" },
}, },
}, },
); );

57
app/globalStats.ts Normal file
View File

@ -0,0 +1,57 @@
import { addDays } from "date-fns";
import { decodeTime } from "ulid";
import { getDailyStats } from "./dailyStatsStore.ts";
import { generationStore } from "./generationStore.ts";
export interface GlobalStats {
userIds: number[];
imageCount: number;
stepCount: number;
pixelCount: number;
pixelStepCount: number;
timestamp: number;
}
export const globalStats: GlobalStats = await getGlobalStats();
async function getGlobalStats(): Promise<GlobalStats> {
// find the year/month/day of the first generation
const startDate = await generationStore.getAll({}, { limit: 1 })
.then((generations) => generations[0]?.id)
.then((generationId) => generationId ? new Date(decodeTime(generationId)) : new Date());
// iterate to today and sum up stats
const userIdSet = new Set<number>();
let imageCount = 0;
let stepCount = 0;
let pixelCount = 0;
let pixelStepCount = 0;
const tomorrow = addDays(new Date(), 1);
for (
let date = startDate;
date < tomorrow;
date = addDays(date, 1)
) {
const dailyStats = await getDailyStats(
date.getUTCFullYear(),
date.getUTCMonth() + 1,
date.getUTCDate(),
);
for (const userId of dailyStats.userIds) userIdSet.add(userId);
imageCount += dailyStats.imageCount;
stepCount += dailyStats.stepCount;
pixelCount += dailyStats.pixelCount;
pixelStepCount += dailyStats.pixelStepCount;
}
return {
userIds: [...userIdSet],
imageCount,
stepCount,
pixelCount,
pixelStepCount,
timestamp: Date.now(),
};
}

View File

@ -3,21 +3,20 @@ import { InputFile, InputMediaBuilder } from "grammy";
import { bold, fmt } from "grammy_parse_mode"; import { bold, fmt } from "grammy_parse_mode";
import { Chat, Message, User } from "grammy_types"; import { Chat, Message, User } from "grammy_types";
import { Queue } from "kvmq"; import { Queue } from "kvmq";
import { format } from "std/fmt/duration"; import { format } from "std/fmt/duration.ts";
import { getLogger } from "std/log"; import { debug, error } from "std/log/mod.ts";
import { bot } from "../bot/mod.ts"; import { bot } from "../bot/mod.ts";
import { formatUserChat } from "../utils/formatUserChat.ts"; import { formatUserChat } from "../utils/formatUserChat.ts";
import { db, fs } from "./db.ts"; import { db, fs } from "./db.ts";
import { generationStore, SdGenerationInfo } from "./generationStore.ts"; import { generationStore, SdGenerationInfo } from "./generationStore.ts";
import { globalStats } from "./globalStats.ts";
const logger = () => getLogger();
interface UploadJob { interface UploadJob {
from: User; from: User;
chat: Chat; chat: Chat;
requestMessage: Message; requestMessage: Message;
replyMessage: Message; replyMessage: Message;
sdInstanceId: string; workerInstanceKey?: string;
startDate: Date; startDate: Date;
endDate: Date; endDate: Date;
imageKeys: Deno.KvKey[]; imageKeys: Deno.KvKey[];
@ -55,19 +54,23 @@ export async function processUploadQueue() {
fmt`${bold("CFG scale:")} ${state.info.cfg_scale}, `, fmt`${bold("CFG scale:")} ${state.info.cfg_scale}, `,
fmt`${bold("Seed:")} ${state.info.seed}, `, fmt`${bold("Seed:")} ${state.info.seed}, `,
fmt`${bold("Size")}: ${state.info.width}x${state.info.height}, `, fmt`${bold("Size")}: ${state.info.width}x${state.info.height}, `,
fmt`${bold("Worker")}: ${state.sdInstanceId}, `, state.workerInstanceKey ? fmt`${bold("Worker")}: ${state.workerInstanceKey}, ` : "",
fmt`${bold("Time taken")}: ${format(jobDurationMs, { ignoreZero: true })}`, fmt`${bold("Time taken")}: ${format(jobDurationMs, { ignoreZero: true })}`,
] ]
: [], : [],
]); ]);
// parse files from reply JSON // parse files from reply JSON
let size = 0;
const types = new Set<string>();
const inputFiles = await Promise.all( const inputFiles = await Promise.all(
state.imageKeys.map(async (fileKey, idx) => { state.imageKeys.map(async (fileKey, idx) => {
const imageBuffer = await fs.get(fileKey).then((entry) => entry.value); const imageBuffer = await fs.get(fileKey).then((entry) => entry.value);
if (!imageBuffer) throw new Error("File not found"); if (!imageBuffer) throw new Error("File not found");
const imageType = await fileTypeFromBuffer(imageBuffer); const imageType = await fileTypeFromBuffer(imageBuffer);
if (!imageType) throw new Error("Image has unknown type"); if (!imageType) throw new Error("Image has unknown type");
size += imageBuffer.byteLength;
types.add(imageType.ext);
return InputMediaBuilder.photo( return InputMediaBuilder.photo(
new InputFile(imageBuffer, `image${idx}.${imageType.ext}`), new InputFile(imageBuffer, `image${idx}.${imageType.ext}`),
// if it can fit, add caption for first photo // if it can fit, add caption for first photo
@ -89,7 +92,7 @@ export async function processUploadQueue() {
// send caption in separate message if it couldn't fit // send caption in separate message if it couldn't fit
if (caption.text.length > 1024 && caption.text.length <= 4096) { if (caption.text.length > 1024 && caption.text.length <= 4096) {
await bot.api.sendMessage(state.chat.id, caption.text, { await bot.api.sendMessage(state.chat.id, caption.text, {
reply_to_message_id: resultMessages[0].message_id, reply_to_message_id: resultMessages[0]!.message_id,
allow_sending_without_reply: true, allow_sending_without_reply: true,
entities: caption.entities, entities: caption.entities,
maxWait: 60, maxWait: 60,
@ -103,19 +106,36 @@ export async function processUploadQueue() {
await generationStore.create({ await generationStore.create({
from: state.from, from: state.from,
chat: state.chat, chat: state.chat,
sdInstanceId: state.sdInstanceId, sdInstanceId: state.workerInstanceKey,
startDate: state.startDate, startDate: state.startDate,
endDate: new Date(), endDate: new Date(),
info: state.info, info: state.info,
}); });
// update live stats
{
globalStats.imageCount++;
globalStats.stepCount += state.info.steps;
globalStats.pixelCount += state.info.width * state.info.height;
globalStats.pixelStepCount += state.info.width * state.info.height * state.info.steps;
const userIdSet = new Set(globalStats.userIds);
userIdSet.add(state.from.id);
globalStats.userIds = [...userIdSet];
}
debug(
`Uploaded ${state.imageKeys.length} ${[...types].join(",")} images (${
Math.trunc(size / 1024)
}kB) for ${formatUserChat(state)}`,
);
// delete the status message // delete the status message
await bot.api.deleteMessage(state.replyMessage.chat.id, state.replyMessage.message_id) await bot.api.deleteMessage(state.replyMessage.chat.id, state.replyMessage.message_id)
.catch(() => undefined); .catch(() => undefined);
}, { concurrency: 3 }); }, { concurrency: 10 });
uploadWorker.addEventListener("error", (e) => { uploadWorker.addEventListener("error", (e) => {
logger().error(`Upload failed for ${formatUserChat(e.detail.job.state)}: ${e.detail.error}`); error(`Upload failed for ${formatUserChat(e.detail.job.state)}: ${e.detail.error}`);
bot.api.sendMessage( bot.api.sendMessage(
e.detail.job.state.requestMessage.chat.id, e.detail.job.state.requestMessage.chat.id,
`Upload failed: ${e.detail.error}\n\n` + `Upload failed: ${e.detail.error}\n\n` +

View File

@ -0,0 +1,59 @@
import { Static, t } from "elysia";
import { kvMemoize } from "../utils/kvMemoize.ts";
import { db } from "./db.ts";
import { generationStore } from "./generationStore.ts";
import { hoursToMilliseconds, isSameDay, minutesToMilliseconds } from "date-fns";
import { UTCDateMini } from "date-fns/utc";
export const userDailyStatsSchema = t.Object({
imageCount: t.Number(),
pixelCount: t.Number(),
pixelStepCount: t.Number(),
timestamp: t.Number(),
});
export type UserDailyStats = Static<typeof userDailyStatsSchema>;
export const getUserDailyStats = kvMemoize(
db,
["userDailyStats"],
async (userId: number, year: number, month: number, day: number): Promise<UserDailyStats> => {
let imageCount = 0;
let pixelCount = 0;
let pixelStepCount = 0;
for await (
const generation of generationStore.listBy("fromId", {
after: new Date(Date.UTC(year, month - 1, day)),
before: new Date(Date.UTC(year, month - 1, day + 1)),
value: userId,
})
) {
imageCount++;
pixelCount += (generation.value.info?.width ?? 0) * (generation.value.info?.height ?? 0);
pixelStepCount += (generation.value.info?.width ?? 0) *
(generation.value.info?.height ?? 0) *
(generation.value.info?.steps ?? 0);
}
return {
imageCount,
pixelCount,
pixelStepCount,
timestamp: Date.now(),
};
},
{
// expire in 1 minute if was calculated on the same day, otherwise 7-14 days.
expireIn: (result, year, month, day) => {
const requestDate = new UTCDateMini(year, month - 1, day);
const calculatedDate = new UTCDateMini(result.timestamp);
return isSameDay(requestDate, calculatedDate)
? minutesToMilliseconds(1)
: hoursToMilliseconds(24 * 7 + Math.random() * 24 * 7);
},
// should cache if the stats are non-zero
shouldCache: (result) =>
result.imageCount > 0 || result.pixelCount > 0 || result.pixelStepCount > 0,
},
);

122
app/userStatsStore.ts Normal file
View File

@ -0,0 +1,122 @@
import { minutesToMilliseconds } from "date-fns";
import { Store } from "indexed_kv";
import { info } from "std/log/mod.ts";
import { Static, t } from "elysia";
import { db } from "./db.ts";
import { generationStore } from "./generationStore.ts";
import { kvMemoize } from "../utils/kvMemoize.ts";
import { sortBy } from "std/collections/sort_by.ts";
export const userStatsSchema = t.Object({
userId: t.Number(),
imageCount: t.Number(),
stepCount: t.Number(),
pixelCount: t.Number(),
pixelStepCount: t.Number(),
tagCountMap: t.Record(t.String(), t.Number()),
timestamp: t.Number(),
});
export type UserStats = Static<typeof userStatsSchema>;
type UserStatsIndices = {
userId: number;
imageCount: number;
pixelCount: number;
};
const userStatsStore = new Store<UserStats, UserStatsIndices>(
db,
"userStats",
{
indices: {
userId: { getValue: (item) => item.userId },
imageCount: { getValue: (item) => item.imageCount },
pixelCount: { getValue: (item) => item.pixelCount },
},
},
);
export const getUserStats = kvMemoize(
db,
["userStats"],
async (userId: number): Promise<UserStats> => {
let imageCount = 0;
let stepCount = 0;
let pixelCount = 0;
let pixelStepCount = 0;
const tagCountMap = new Map<string, number>();
info(`Calculating user stats for ${userId}`);
for await (
const generation of generationStore.listBy("fromId", { value: userId })
) {
imageCount++;
stepCount += generation.value.info?.steps ?? 0;
pixelCount += (generation.value.info?.width ?? 0) * (generation.value.info?.height ?? 0);
pixelStepCount += (generation.value.info?.width ?? 0) *
(generation.value.info?.height ?? 0) *
(generation.value.info?.steps ?? 0);
const tags = generation.value.info?.prompt
// split on punctuation and newlines
.split(/[,;.]\s+|\n/)
// remove `:weight` syntax
.map((tag) => tag.replace(/:[\d\.]+/g, ""))
// remove `(tag)` and `[tag]` syntax
.map((tag) => tag.replace(/[()[\]]/g, " "))
// collapse multiple whitespace to one
.map((tag) => tag.replace(/\s+/g, " "))
// trim whitespace
.map((tag) => tag.trim())
// remove empty tags
.filter((tag) => tag.length > 0)
// lowercase tags
.map((tag) => tag.toLowerCase()) ??
// default to empty array
[];
for (const tag of tags) {
const count = tagCountMap.get(tag) ?? 0;
tagCountMap.set(tag, count + 1);
}
}
const tagCountObj = Object.fromEntries(
sortBy(
Array.from(tagCountMap.entries()),
([_tag, count]) => -count,
).filter(([_tag, count]) => count >= 3),
);
return {
userId,
imageCount,
stepCount,
pixelCount,
pixelStepCount,
tagCountMap: tagCountObj,
timestamp: Date.now(),
};
},
{
// expire in random time between 5-10 minutes
expireIn: () => minutesToMilliseconds(5 + Math.random() * 5),
// override default set/get behavior to use userStatsStore
override: {
get: async (_key, [userId]) => {
const items = await userStatsStore.getBy("userId", { value: userId }, { reverse: true });
return items[0]?.value;
},
set: async (_key, [userId], value, options) => {
// delete old stats
for await (const item of userStatsStore.listBy("userId", { value: userId })) {
await item.delete();
}
// set new stats
await userStatsStore.create(value, options);
},
},
},
);

View File

@ -0,0 +1,24 @@
import { Store } from "indexed_kv";
import { Static, t } from "elysia";
import { db } from "./db.ts";
export const workerInstanceSchema = t.Object({
key: t.String(),
name: t.Nullable(t.String()),
sdUrl: t.String(),
sdAuth: t.Nullable(t.Object({
user: t.String(),
password: t.String(),
})),
lastOnlineTime: t.Optional(t.Number()),
lastError: t.Optional(t.Object({
message: t.String(),
time: t.Number(),
})),
});
export type WorkerInstance = Static<typeof workerInstanceSchema>;
export const workerInstanceStore = new Store<WorkerInstance>(db, "workerInstances", {
indices: {},
});

View File

@ -1,19 +1,20 @@
import { CommandContext } from "grammy"; import { CommandContext } from "grammy";
import { bold, fmt, FormattedString } from "grammy_parse_mode"; import { bold, fmt, FormattedString } from "grammy_parse_mode";
import { distinctBy } from "std/collections"; import { distinctBy } from "std/collections/distinct_by.ts";
import { getConfig } from "../app/config.ts"; import { error, info } from "std/log/mod.ts";
import { adminStore } from "../app/adminStore.ts";
import { generationStore } from "../app/generationStore.ts"; import { generationStore } from "../app/generationStore.ts";
import { ErisContext, logger } from "./mod.ts";
import { formatUserChat } from "../utils/formatUserChat.ts"; import { formatUserChat } from "../utils/formatUserChat.ts";
import { ErisContext } from "./mod.ts";
export async function broadcastCommand(ctx: CommandContext<ErisContext>) { export async function broadcastCommand(ctx: CommandContext<ErisContext>) {
if (!ctx.from?.username) { if (!ctx.from?.username) {
return ctx.reply("I don't know who you are."); return ctx.reply("I don't know who you are.");
} }
const config = await getConfig(); const adminEntry = await adminStore.get(["admins", ctx.from.id]);
if (!config.adminUsernames.includes(ctx.from.username)) { if (!adminEntry.versionstamp) {
return ctx.reply("Only a bot admin can use this command."); return ctx.reply("Only a bot admin can use this command.");
} }
@ -39,16 +40,17 @@ export async function broadcastCommand(ctx: CommandContext<ErisContext>) {
const replyMessage = await ctx.replyFmt(getMessage(), { const replyMessage = await ctx.replyFmt(getMessage(), {
reply_to_message_id: ctx.message?.message_id, reply_to_message_id: ctx.message?.message_id,
allow_sending_without_reply: true,
}); });
// send message to each user // send message to each user
for (const gen of gens) { for (const gen of gens) {
try { try {
await ctx.api.sendMessage(gen.value.from.id, text); await ctx.api.sendMessage(gen.value.from.id, text);
logger().info(`Broadcasted to ${formatUserChat({ from: gen.value.from })}`); info(`Broadcasted to ${formatUserChat({ from: gen.value.from })}`);
sentCount++; sentCount++;
} catch (err) { } catch (err) {
logger().error(`Broadcasting to ${formatUserChat({ from: gen.value.from })} failed: ${err}`); error(`Broadcasting to ${formatUserChat({ from: gen.value.from })} failed: ${err}`);
errors.push(fmt`${bold(formatUserChat({ from: gen.value.from }))} - ${err.message}\n`); errors.push(fmt`${bold(formatUserChat({ from: gen.value.from }))} - ${err.message}\n`);
} }
const fmtMessage = getMessage(); const fmtMessage = getMessage();

View File

@ -1,4 +1,5 @@
import { generationQueue } from "../app/generationQueue.ts"; import { generationQueue } from "../app/generationQueue.ts";
import { omitUndef } from "../utils/omitUndef.ts";
import { ErisContext } from "./mod.ts"; import { ErisContext } from "./mod.ts";
export async function cancelCommand(ctx: ErisContext) { export async function cancelCommand(ctx: ErisContext) {
@ -7,7 +8,11 @@ export async function cancelCommand(ctx: ErisContext) {
.filter((job) => job.lockUntil < new Date()) .filter((job) => job.lockUntil < new Date())
.filter((j) => j.state.from.id === ctx.from?.id); .filter((j) => j.state.from.id === ctx.from?.id);
for (const job of userJobs) await generationQueue.deleteJob(job.id); for (const job of userJobs) await generationQueue.deleteJob(job.id);
await ctx.reply(`Cancelled ${userJobs.length} jobs`, { await ctx.reply(
reply_to_message_id: ctx.message?.message_id, `Cancelled ${userJobs.length} jobs`,
}); omitUndef({
reply_to_message_id: ctx.message?.message_id,
allow_sending_without_reply: true,
}),
);
} }

View File

@ -1,11 +1,12 @@
import { CommandContext } from "grammy"; import { CommandContext } from "grammy";
import { StatelessQuestion } from "grammy_stateless_question"; import { StatelessQuestion } from "grammy_stateless_question";
import { maxBy } from "std/collections"; import { maxBy } from "std/collections/max_by.ts";
import { debug } from "std/log/mod.ts";
import { getConfig } from "../app/config.ts"; import { getConfig } from "../app/config.ts";
import { generationQueue } from "../app/generationQueue.ts"; import { generationQueue } from "../app/generationQueue.ts";
import { parsePngInfo, PngInfo } from "../sd/parsePngInfo.ts";
import { formatUserChat } from "../utils/formatUserChat.ts"; import { formatUserChat } from "../utils/formatUserChat.ts";
import { ErisContext, logger } from "./mod.ts"; import { ErisContext } from "./mod.ts";
import { parsePngInfo, PngInfo } from "./parsePngInfo.ts";
type QuestionState = { fileId?: string; params?: Partial<PngInfo> }; type QuestionState = { fileId?: string; params?: Partial<PngInfo> };
@ -27,10 +28,11 @@ async function img2img(
includeRepliedTo: boolean, includeRepliedTo: boolean,
state: QuestionState = {}, state: QuestionState = {},
): Promise<void> { ): Promise<void> {
if (!ctx.message?.from?.id) { if (!ctx.from || !ctx.message) {
await ctx.reply("I don't know who you are", { return;
reply_to_message_id: ctx.message?.message_id, }
});
if (ctx.from.is_bot) {
return; return;
} }
@ -53,7 +55,7 @@ async function img2img(
const userJobs = jobs.filter((job) => job.state.from.id === ctx.message?.from?.id); const userJobs = jobs.filter((job) => job.state.from.id === ctx.message?.from?.id);
if (userJobs.length >= config.maxUserJobs) { if (userJobs.length >= config.maxUserJobs) {
await ctx.reply(`You already have ${config.maxUserJobs} jobs in queue. Try again later.`, { await ctx.reply(`You already have ${userJobs.length} jobs in queue. Try again later.`, {
reply_to_message_id: ctx.message?.message_id, reply_to_message_id: ctx.message?.message_id,
}); });
return; return;
@ -127,7 +129,7 @@ async function img2img(
chat: ctx.message.chat, chat: ctx.message.chat,
requestMessage: ctx.message, requestMessage: ctx.message,
replyMessage: replyMessage, replyMessage: replyMessage,
}, { retryCount: 3, repeatDelayMs: 10_000 }); }, { priority: 0, retryCount: 3, repeatDelayMs: 10_000 });
logger().debug(`Generation (img2img) enqueued for ${formatUserChat(ctx.message)}`); debug(`Generation (img2img) enqueued for ${formatUserChat(ctx.message)}`);
} }

View File

@ -1,9 +1,11 @@
import { Api, Bot, Context, RawApi, session, SessionFlavor } from "grammy"; import { Api, Bot, Context, RawApi, session, SessionFlavor } from "grammy";
import { FileFlavor, hydrateFiles } from "grammy_files"; import { FileFlavor, hydrateFiles } from "grammy_files";
import { hydrateReply, ParseModeFlavor } from "grammy_parse_mode"; import { hydrateReply, ParseModeFlavor } from "grammy_parse_mode";
import { getLogger } from "std/log"; import { run, sequentialize } from "grammy_runner";
import { getConfig, setConfig } from "../app/config.ts"; import { error, info, warning } from "std/log/mod.ts";
import { sessions } from "../api/sessionsRoute.ts";
import { formatUserChat } from "../utils/formatUserChat.ts"; import { formatUserChat } from "../utils/formatUserChat.ts";
import { omitUndef } from "../utils/omitUndef.ts";
import { broadcastCommand } from "./broadcastCommand.ts"; import { broadcastCommand } from "./broadcastCommand.ts";
import { cancelCommand } from "./cancelCommand.ts"; import { cancelCommand } from "./cancelCommand.ts";
import { img2imgCommand, img2imgQuestion } from "./img2imgCommand.ts"; import { img2imgCommand, img2imgQuestion } from "./img2imgCommand.ts";
@ -11,19 +13,17 @@ import { pnginfoCommand, pnginfoQuestion } from "./pnginfoCommand.ts";
import { queueCommand } from "./queueCommand.ts"; import { queueCommand } from "./queueCommand.ts";
import { txt2imgCommand, txt2imgQuestion } from "./txt2imgCommand.ts"; import { txt2imgCommand, txt2imgQuestion } from "./txt2imgCommand.ts";
export const logger = () => getLogger();
interface SessionData { interface SessionData {
chat: ErisChatData; chat: ErisChatData;
user: ErisUserData; user: ErisUserData;
} }
interface ErisChatData { interface ErisChatData {
language?: string; language?: string | undefined;
} }
interface ErisUserData { interface ErisUserData {
params?: Record<string, string>; params?: Record<string, string> | undefined;
} }
export type ErisContext = export type ErisContext =
@ -46,6 +46,9 @@ export const bot = new Bot<ErisContext, ErisApi>(
); );
bot.use(hydrateReply); bot.use(hydrateReply);
bot.use(sequentialize((ctx) => ctx.chat?.id.toString()));
bot.use(session<SessionData, ErisContext>({ bot.use(session<SessionData, ErisContext>({
type: "multi", type: "multi",
chat: { chat: {
@ -72,7 +75,7 @@ bot.api.config.use(async (prev, method, payload, signal) => {
if (attempt >= maxAttempts) return result; if (attempt >= maxAttempts) return result;
const retryAfter = result.parameters?.retry_after ?? (attempt * 5); const retryAfter = result.parameters?.retry_after ?? (attempt * 5);
if (retryAfter > maxWait) return result; if (retryAfter > maxWait) return result;
logger().warning( warning(
`${method} (attempt ${attempt}) failed: ${result.error_code} ${result.description}`, `${method} (attempt ${attempt}) failed: ${result.error_code} ${result.description}`,
); );
await new Promise((resolve) => setTimeout(resolve, retryAfter * 1000)); await new Promise((resolve) => setTimeout(resolve, retryAfter * 1000));
@ -80,7 +83,7 @@ bot.api.config.use(async (prev, method, payload, signal) => {
}); });
bot.catch((err) => { bot.catch((err) => {
logger().error( error(
`Handling update from ${formatUserChat(err.ctx)} failed: ${err.name} ${err.message}`, `Handling update from ${formatUserChat(err.ctx)} failed: ${err.name} ${err.message}`,
); );
}); });
@ -91,9 +94,13 @@ bot.use(async (ctx, next) => {
await next(); await next();
} catch (err) { } catch (err) {
try { try {
await ctx.reply(`Handling update failed: ${err}`, { await ctx.reply(
reply_to_message_id: ctx.message?.message_id, `Handling update failed: ${err}`,
}); omitUndef({
reply_to_message_id: ctx.message?.message_id,
allow_sending_without_reply: true,
}),
);
} catch { } catch {
throw err; throw err;
} }
@ -108,12 +115,44 @@ bot.api.setMyDescription(
bot.api.setMyCommands([ bot.api.setMyCommands([
{ command: "txt2img", description: "Generate image from text" }, { command: "txt2img", description: "Generate image from text" },
{ command: "img2img", description: "Generate image from image" }, { command: "img2img", description: "Generate image from image" },
{ command: "pnginfo", description: "Show generation parameters of an image" }, { command: "pnginfo", description: "Try to extract prompt from raw file" },
{ command: "queue", description: "Show the current queue" }, { command: "queue", description: "Show the current queue" },
{ command: "cancel", description: "Cancel all your requests" }, { command: "cancel", description: "Cancel all your requests" },
]); ]);
bot.command("start", (ctx) => ctx.reply("Hello! Use the /txt2img command to generate an image")); bot.command("start", async (ctx) => {
if (ctx.match) {
const id = ctx.match.trim();
const session = sessions.get(id);
if (session == null) {
await ctx.reply(
"Login failed: Invalid session ID",
omitUndef({
reply_to_message_id: ctx.message?.message_id,
}),
);
return;
}
session.userId = ctx.from?.id;
sessions.set(id, session);
info(`User ${formatUserChat(ctx)} logged in`);
// TODO: show link to web ui
await ctx.reply(
"Login successful! You can now return to the WebUI.",
omitUndef({
reply_to_message_id: ctx.message?.message_id,
}),
);
return;
}
await ctx.reply(
"Hello! Use the /txt2img command to generate an image",
omitUndef({
reply_to_message_id: ctx.message?.message_id,
}),
);
});
bot.command("txt2img", txt2imgCommand); bot.command("txt2img", txt2imgCommand);
bot.use(txt2imgQuestion.middleware()); bot.use(txt2imgQuestion.middleware());
@ -130,30 +169,11 @@ bot.command("cancel", cancelCommand);
bot.command("broadcast", broadcastCommand); bot.command("broadcast", broadcastCommand);
bot.command("pause", async (ctx) => {
if (!ctx.from?.username) return;
const config = await getConfig();
if (!config.adminUsernames.includes(ctx.from.username)) return;
if (config.pausedReason != null) {
return ctx.reply(`Already paused: ${config.pausedReason}`);
}
config.pausedReason = ctx.match ?? "No reason given";
await setConfig(config);
logger().warning(`Bot paused by ${ctx.from.first_name} because ${config.pausedReason}`);
return ctx.reply("Paused");
});
bot.command("resume", async (ctx) => {
if (!ctx.from?.username) return;
const config = await getConfig();
if (!config.adminUsernames.includes(ctx.from.username)) return;
if (config.pausedReason == null) return ctx.reply("Already running");
config.pausedReason = null;
await setConfig(config);
logger().info(`Bot resumed by ${ctx.from.first_name}`);
return ctx.reply("Resumed");
});
bot.command("crash", () => { bot.command("crash", () => {
throw new Error("Crash command used"); throw new Error("Crash command used");
}); });
export async function runBot() {
const runner = run(bot, { runner: { silent: true } });
await runner.task();
}

View File

@ -1,12 +1,24 @@
import { decode } from "png_chunk_text"; import * as ExifReader from "exifreader";
import extractChunks from "png_chunks_extract";
export function getPngInfo(pngData: Uint8Array): string | undefined { export function getPngInfo(pngData: ArrayBuffer): string | undefined {
return extractChunks(pngData) const info = ExifReader.load(pngData);
.filter((chunk) => chunk.name === "tEXt")
.map((chunk) => decode(chunk.data)) if (info.UserComment?.value && Array.isArray(info.UserComment.value)) {
.find((textChunk) => textChunk.keyword === "parameters") // JPEG image
?.text; return String.fromCharCode(
...info.UserComment.value
.filter((char): char is number => typeof char == "number")
.filter((char) => char !== 0),
).replace("UNICODE", "");
}
if (info.parameters?.description) {
// PNG image
return info.parameters.description;
}
// Unknown image type
return undefined;
} }
export interface PngInfo { export interface PngInfo {
@ -25,7 +37,11 @@ interface PngInfoExtra extends PngInfo {
upscale?: number; upscale?: number;
} }
export function parsePngInfo(pngInfo: string, baseParams?: Partial<PngInfo>): Partial<PngInfo> { export function parsePngInfo(
pngInfo: string,
baseParams?: Partial<PngInfo>,
shouldParseSeed?: boolean,
): Partial<PngInfo> {
const tags = pngInfo.split(/[,;]+|\.+\s|\n/u); const tags = pngInfo.split(/[,;]+|\.+\s|\n/u);
let part: "prompt" | "negative_prompt" | "params" = "prompt"; let part: "prompt" | "negative_prompt" | "params" = "prompt";
const params: Partial<PngInfoExtra> = {}; const params: Partial<PngInfoExtra> = {};
@ -34,7 +50,7 @@ export function parsePngInfo(pngInfo: string, baseParams?: Partial<PngInfo>): Pa
for (const tag of tags) { for (const tag of tags) {
const paramValuePair = tag.trim().match(/^(\w+\s*\w*):\s+(.*)$/u); const paramValuePair = tag.trim().match(/^(\w+\s*\w*):\s+(.*)$/u);
if (paramValuePair) { if (paramValuePair) {
const [, param, value] = paramValuePair; const [_match, param = "", value = ""] = paramValuePair;
switch (param.replace(/\s+/u, "").toLowerCase()) { switch (param.replace(/\s+/u, "").toLowerCase()) {
case "positiveprompt": case "positiveprompt":
case "positive": case "positive":
@ -67,7 +83,7 @@ export function parsePngInfo(pngInfo: string, baseParams?: Partial<PngInfo>): Pa
case "size": case "size":
case "resolution": { case "resolution": {
part = "params"; part = "params";
const [width, height] = value.trim() const [width = 0, height = 0] = value.trim()
.split(/\s*[x,]\s*/u, 2) .split(/\s*[x,]\s*/u, 2)
.map((v) => v.trim()) .map((v) => v.trim())
.map(Number); .map(Number);
@ -99,7 +115,16 @@ export function parsePngInfo(pngInfo: string, baseParams?: Partial<PngInfo>): Pa
params.denoising_strength = denoisingStrength; params.denoising_strength = denoisingStrength;
break; break;
} }
case "seed": case "seed": {
part = "params";
if (shouldParseSeed) {
const seed = Number(value.trim());
if (Number.isFinite(seed)) {
params.seed = seed;
}
}
break;
}
case "model": case "model":
case "modelhash": case "modelhash":
case "modelname": case "modelname":

View File

@ -1,8 +1,9 @@
import { CommandContext } from "grammy"; import { CommandContext } from "grammy";
import { bold, fmt } from "grammy_parse_mode"; import { bold, fmt } from "grammy_parse_mode";
import { StatelessQuestion } from "grammy_stateless_question"; import { StatelessQuestion } from "grammy_stateless_question";
import { getPngInfo, parsePngInfo } from "../sd/parsePngInfo.ts"; import { omitUndef } from "../utils/omitUndef.ts";
import { ErisContext } from "./mod.ts"; import { ErisContext } from "./mod.ts";
import { getPngInfo, parsePngInfo } from "./parsePngInfo.ts";
export const pnginfoQuestion = new StatelessQuestion<ErisContext>( export const pnginfoQuestion = new StatelessQuestion<ErisContext>(
"pnginfo", "pnginfo",
@ -19,22 +20,31 @@ async function pnginfo(ctx: ErisContext, includeRepliedTo: boolean): Promise<voi
const document = ctx.message?.document || const document = ctx.message?.document ||
(includeRepliedTo ? ctx.message?.reply_to_message?.document : undefined); (includeRepliedTo ? ctx.message?.reply_to_message?.document : undefined);
if (document?.mime_type !== "image/png") { if (document?.mime_type !== "image/png" && document?.mime_type !== "image/jpeg") {
await ctx.reply( await ctx.reply(
"Please send me a PNG file." + "Please send me a PNG or JPEG file." +
pnginfoQuestion.messageSuffixMarkdown(), pnginfoQuestion.messageSuffixMarkdown(),
{ omitUndef(
reply_markup: { force_reply: true, selective: true }, {
parse_mode: "Markdown", reply_markup: { force_reply: true, selective: true },
reply_to_message_id: ctx.message?.message_id, parse_mode: "Markdown",
}, reply_to_message_id: ctx.message?.message_id,
} as const,
),
); );
return; return;
} }
const file = await ctx.api.getFile(document.file_id); const file = await ctx.api.getFile(document.file_id);
const buffer = await fetch(file.getUrl()).then((resp) => resp.arrayBuffer()); const buffer = await fetch(file.getUrl()).then((resp) => resp.arrayBuffer());
const params = parsePngInfo(getPngInfo(new Uint8Array(buffer)) ?? ""); const info = getPngInfo(buffer);
if (!info) {
return void await ctx.reply(
"No info found in file.",
omitUndef({ reply_to_message_id: ctx.message?.message_id }),
);
}
const params = parsePngInfo(info, undefined, true);
const paramsText = fmt([ const paramsText = fmt([
`${params.prompt}\n`, `${params.prompt}\n`,
@ -46,8 +56,11 @@ async function pnginfo(ctx: ErisContext, includeRepliedTo: boolean): Promise<voi
params.width && params.height ? fmt`${bold("Size")}: ${params.width}x${params.height}` : "", params.width && params.height ? fmt`${bold("Size")}: ${params.width}x${params.height}` : "",
]); ]);
await ctx.reply(paramsText.text, { await ctx.reply(
entities: paramsText.entities, paramsText.text,
reply_to_message_id: ctx.message?.message_id, omitUndef({
}); entities: paramsText.entities,
reply_to_message_id: ctx.message?.message_id,
}),
);
} }

View File

@ -1,21 +1,25 @@
import { CommandContext } from "grammy"; import { CommandContext } from "grammy";
import { bold, fmt } from "grammy_parse_mode"; import { bold, fmt } from "grammy_parse_mode";
import { getConfig } from "../app/config.ts";
import { activeGenerationWorkers, generationQueue } from "../app/generationQueue.ts"; import { activeGenerationWorkers, generationQueue } from "../app/generationQueue.ts";
import { workerInstanceStore } from "../app/workerInstanceStore.ts";
import { getFlagEmoji } from "../utils/getFlagEmoji.ts"; import { getFlagEmoji } from "../utils/getFlagEmoji.ts";
import { omitUndef } from "../utils/omitUndef.ts";
import { ErisContext } from "./mod.ts"; import { ErisContext } from "./mod.ts";
export async function queueCommand(ctx: CommandContext<ErisContext>) { export async function queueCommand(ctx: CommandContext<ErisContext>) {
let formattedMessage = await getMessageText(); let formattedMessage = await getMessageText();
const queueMessage = await ctx.replyFmt(formattedMessage, { const queueMessage = await ctx.replyFmt(
disable_notification: true, formattedMessage,
reply_to_message_id: ctx.message?.message_id, omitUndef({
}); disable_notification: true,
reply_to_message_id: ctx.message?.message_id,
}),
);
handleFutureUpdates().catch(() => undefined); handleFutureUpdates().catch(() => undefined);
async function getMessageText() { async function getMessageText() {
const config = await getConfig();
const allJobs = await generationQueue.getAllJobs(); const allJobs = await generationQueue.getAllJobs();
const workerInstances = await workerInstanceStore.getAll();
const processingJobs = allJobs const processingJobs = allJobs
.filter((job) => job.lockUntil > new Date()).map((job) => ({ ...job, index: 0 })); .filter((job) => job.lockUntil > new Date()).map((job) => ({ ...job, index: 0 }));
const waitingJobs = allJobs const waitingJobs = allJobs
@ -30,24 +34,17 @@ export async function queueCommand(ctx: CommandContext<ErisContext>) {
`${job.index}. `, `${job.index}. `,
fmt`${bold(job.state.from.first_name)} `, fmt`${bold(job.state.from.first_name)} `,
job.state.from.last_name ? fmt`${bold(job.state.from.last_name)} ` : "", job.state.from.last_name ? fmt`${bold(job.state.from.last_name)} ` : "",
job.state.from.username ? `(@${job.state.from.username}) ` : "",
getFlagEmoji(job.state.from.language_code) ?? "", getFlagEmoji(job.state.from.language_code) ?? "",
job.state.chat.type === "private" ? " in private chat " : ` in ${job.state.chat.title} `, job.index === 0 && job.state.progress && job.state.workerInstanceKey
job.state.chat.type !== "private" && job.state.chat.type !== "group" && ? `(${(job.state.progress * 100).toFixed(0)}% using ${job.state.workerInstanceKey}) `
job.state.chat.username
? `(@${job.state.chat.username}) `
: "",
job.index === 0 && job.state.progress && job.state.sdInstanceId
? `(${(job.state.progress * 100).toFixed(0)}% using ${job.state.sdInstanceId}) `
: "", : "",
"\n", "\n",
]) ])
: ["Queue is empty.\n"], : ["Queue is empty.\n"],
"\nActive workers:\n", "\nActive workers:\n",
...config.sdInstances.flatMap((sdInstance) => [ ...workerInstances.flatMap((workerInstace) => [
activeGenerationWorkers.get(sdInstance.id)?.isProcessing ? "✅ " : "☠️ ", activeGenerationWorkers.get(workerInstace.id)?.isProcessing ? "✅ " : "☠️ ",
fmt`${bold(sdInstance.name || sdInstance.id)} `, fmt`${bold(workerInstace.value.name || workerInstace.value.key)} `,
`(max ${(sdInstance.maxResolution / 1000000).toFixed(1)} Mpx) `,
"\n", "\n",
]), ]),
]); ]);

View File

@ -1,10 +1,12 @@
import { CommandContext } from "grammy"; import { CommandContext } from "grammy";
import { StatelessQuestion } from "grammy_stateless_question"; import { StatelessQuestion } from "grammy_stateless_question";
import { debug } from "std/log/mod.ts";
import { getConfig } from "../app/config.ts"; import { getConfig } from "../app/config.ts";
import { generationQueue } from "../app/generationQueue.ts"; import { generationQueue } from "../app/generationQueue.ts";
import { getPngInfo, parsePngInfo, PngInfo } from "../sd/parsePngInfo.ts";
import { formatUserChat } from "../utils/formatUserChat.ts"; import { formatUserChat } from "../utils/formatUserChat.ts";
import { ErisContext, logger } from "./mod.ts"; import { ErisContext } from "./mod.ts";
import { getPngInfo, parsePngInfo, PngInfo } from "./parsePngInfo.ts";
import { adminStore } from "../app/adminStore.ts";
export const txt2imgQuestion = new StatelessQuestion<ErisContext>( export const txt2imgQuestion = new StatelessQuestion<ErisContext>(
"txt2img", "txt2img",
@ -19,12 +21,16 @@ export async function txt2imgCommand(ctx: CommandContext<ErisContext>) {
} }
async function txt2img(ctx: ErisContext, match: string, includeRepliedTo: boolean): Promise<void> { async function txt2img(ctx: ErisContext, match: string, includeRepliedTo: boolean): Promise<void> {
if (!ctx.message?.from?.id) { if (!ctx.from || !ctx.message) {
await ctx.reply("I don't know who you are", { reply_to_message_id: ctx.message?.message_id }); return;
}
if (ctx.from.is_bot) {
return; return;
} }
const config = await getConfig(); const config = await getConfig();
let priority = 0;
if (config.pausedReason != null) { if (config.pausedReason != null) {
await ctx.reply(`I'm paused: ${config.pausedReason || "No reason given"}`, { await ctx.reply(`I'm paused: ${config.pausedReason || "No reason given"}`, {
@ -33,20 +39,26 @@ async function txt2img(ctx: ErisContext, match: string, includeRepliedTo: boolea
return; return;
} }
const jobs = await generationQueue.getAllJobs(); const adminEntry = await adminStore.get(["admins", ctx.from.id]);
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 (adminEntry.versionstamp) {
if (userJobs.length >= config.maxUserJobs) { priority = 1;
await ctx.reply(`You already have ${config.maxUserJobs} jobs in queue. Try again later.`, { } else {
reply_to_message_id: ctx.message?.message_id, const jobs = await generationQueue.getAllJobs();
}); if (jobs.length >= config.maxJobs) {
return; 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<PngInfo> = {}; let params: Partial<PngInfo> = {};
@ -56,38 +68,25 @@ async function txt2img(ctx: ErisContext, match: string, includeRepliedTo: boolea
if (includeRepliedTo && repliedToMsg?.document?.mime_type === "image/png") { if (includeRepliedTo && repliedToMsg?.document?.mime_type === "image/png") {
const file = await ctx.api.getFile(repliedToMsg.document.file_id); const file = await ctx.api.getFile(repliedToMsg.document.file_id);
const buffer = await fetch(file.getUrl()).then((resp) => resp.arrayBuffer()); const buffer = await fetch(file.getUrl()).then((resp) => resp.arrayBuffer());
params = parsePngInfo(getPngInfo(new Uint8Array(buffer)) ?? "", params); params = parsePngInfo(getPngInfo(buffer) ?? "", params);
} }
const repliedToText = repliedToMsg?.text || repliedToMsg?.caption; const repliedToText = repliedToMsg?.text || repliedToMsg?.caption;
if (includeRepliedTo && repliedToText) { const isReply = includeRepliedTo && repliedToText;
if (isReply) {
// TODO: remove bot command from replied to text // TODO: remove bot command from replied to text
params = parsePngInfo(repliedToText, params); params = parsePngInfo(repliedToText, params);
} }
params = parsePngInfo(match, params); params = parsePngInfo(match, params, true);
if (ctx.message.text.trim() == "help") { if (isReply) {
await ctx.reply(` const parsedInfo = parsePngInfo(repliedToText, undefined, true);
Usage: /txt2img <tags> [options] if (parsedInfo.prompt !== params.prompt) {
<>: required params.seed = parsedInfo.seed ?? -1;
[]: 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,
});
} }
}
if (!params.prompt) { if (!params.prompt) {
await ctx.reply( await ctx.reply(
@ -112,7 +111,7 @@ async function txt2img(ctx: ErisContext, match: string, includeRepliedTo: boolea
chat: ctx.message.chat, chat: ctx.message.chat,
requestMessage: ctx.message, requestMessage: ctx.message,
replyMessage: replyMessage, replyMessage: replyMessage,
}, { retryCount: 3, retryDelayMs: 10_000 }); }, { retryCount: 3, retryDelayMs: 10_000, priority: priority });
logger().debug(`Generation (txt2img) enqueued for ${formatUserChat(ctx.message)}`); debug(`Generation (txt2img) enqueued for ${formatUserChat(ctx.message)}`);
} }

63
deno.json Normal file
View File

@ -0,0 +1,63 @@
{
"compilerOptions": {
"exactOptionalPropertyTypes": true,
"jsx": "react",
"noUncheckedIndexedAccess": true
},
"fmt": {
"lineWidth": 100
},
"imports": {
"async": "https://deno.land/x/async@v2.0.2/mod.ts",
"date-fns": "https://cdn.skypack.dev/date-fns@2.30.0?dts",
"date-fns/utc": "https://cdn.skypack.dev/@date-fns/utc@1.1.0?dts",
"elysia": "https://esm.sh/elysia@0.7.21?dev",
"elysia/eden": "https://esm.sh/@elysiajs/eden@0.7.4?external=elysia&dev",
"elysia/swagger": "https://esm.sh/@elysiajs/swagger@0.7.4?external=elysia&dev",
"exifreader": "https://esm.sh/exifreader@4.14.1",
"file_type": "https://esm.sh/file-type@18.5.0",
"grammy": "https://lib.deno.dev/x/grammy@1/mod.ts",
"grammy_autoquote": "https://lib.deno.dev/x/grammy_autoquote@1/mod.ts",
"grammy_files": "https://lib.deno.dev/x/grammy_files@1/mod.ts",
"grammy_parse_mode": "https://lib.deno.dev/x/grammy_parse_mode@1/mod.ts",
"grammy_runner": "https://lib.deno.dev/x/grammy_runner@2/mod.ts",
"grammy_stateless_question": "https://lib.deno.dev/x/grammy_stateless_question_alpha@3/mod.ts",
"grammy_types": "https://lib.deno.dev/x/grammy_types@3/mod.ts",
"indexed_kv": "https://deno.land/x/indexed_kv@v0.6.1/mod.ts",
"kvfs": "https://deno.land/x/kvfs@v0.1.0/mod.ts",
"kvmq": "https://deno.land/x/kvmq@v0.3.0/mod.ts",
"openapi_fetch": "https://esm.sh/openapi-fetch@0.7.6",
"react": "https://esm.sh/react@18.2.0?dev",
"react-dom/client": "https://esm.sh/react-dom@18.2.0/client?external=react&dev",
"react-flip-move": "https://esm.sh/react-flip-move@3.0.5?external=react&dev",
"react-intl": "https://esm.sh/react-intl@6.4.7?external=react&dev",
"react-router-dom": "https://esm.sh/react-router-dom@6.16.0?external=react&dev",
"reroute": "https://deno.land/x/reroute@v0.1.0/mod.ts",
"serve_spa": "https://deno.land/x/serve_spa@v0.2.0/mod.ts",
"std/async/": "https://deno.land/std@0.201.0/async/",
"std/collections/": "https://deno.land/std@0.202.0/collections/",
"std/dotenv/": "https://deno.land/std@0.201.0/dotenv/",
"std/encoding/": "https://deno.land/std@0.202.0/encoding/",
"std/fmt/": "https://deno.land/std@0.202.0/fmt/",
"std/log/": "https://deno.land/std@0.201.0/log/",
"std/path/": "https://deno.land/std@0.204.0/path/",
"swr": "https://esm.sh/swr@2.2.4?external=react&dev",
"swr/mutation": "https://esm.sh/swr@2.2.4/mutation?external=react&dev",
"twind/core": "https://esm.sh/@twind/core@1.1.3",
"twind/preset-tailwind": "https://esm.sh/@twind/preset-tailwind@1.1.4",
"ulid": "https://deno.land/x/ulid@v0.3.0/mod.ts"
},
"lint": {
"rules": {
"exclude": [
"require-await"
]
}
},
"tasks": {
"check": "deno check --unstable-kv main.ts && deno check --unstable-kv ui/main.tsx",
"generate": "deno run npm:openapi-typescript http://localhost:7861/openapi.json -o app/sdApi.ts",
"start": "deno run --unstable-kv --allow-env --allow-read=. --allow-write=db --allow-net main.ts",
"test": "deno test"
}
}

View File

@ -1,30 +0,0 @@
{
"tasks": {
"start": "deno run --unstable --allow-env --allow-read --allow-write --allow-net main.ts"
},
"fmt": {
"lineWidth": 100
},
"imports": {
"std/log": "https://deno.land/std@0.201.0/log/mod.ts",
"std/async": "https://deno.land/std@0.201.0/async/mod.ts",
"std/fmt/duration": "https://deno.land/std@0.202.0/fmt/duration.ts",
"std/collections": "https://deno.land/std@0.202.0/collections/mod.ts",
"std/encoding/base64": "https://deno.land/std@0.202.0/encoding/base64.ts",
"async": "https://deno.land/x/async@v2.0.2/mod.ts",
"ulid": "https://deno.land/x/ulid@v0.3.0/mod.ts",
"indexed_kv": "https://deno.land/x/indexed_kv@v0.4.0/mod.ts",
"kvmq": "https://deno.land/x/kvmq@v0.2.0/mod.ts",
"kvfs": "https://deno.land/x/kvfs@v0.1.0/mod.ts",
"grammy": "https://lib.deno.dev/x/grammy@1/mod.ts",
"grammy_types": "https://lib.deno.dev/x/grammy_types@3/mod.ts",
"grammy_autoquote": "https://lib.deno.dev/x/grammy_autoquote@1/mod.ts",
"grammy_parse_mode": "https://lib.deno.dev/x/grammy_parse_mode@1/mod.ts",
"grammy_stateless_question": "https://lib.deno.dev/x/grammy_stateless_question_alpha@3/mod.ts",
"grammy_files": "https://lib.deno.dev/x/grammy_files@1/mod.ts",
"file_type": "https://esm.sh/file-type@18.5.0",
"png_chunks_extract": "https://esm.sh/png-chunks-extract@1.0.0",
"png_chunk_text": "https://esm.sh/png-chunk-text@1.0.0",
"openapi_fetch": "https://esm.sh/openapi-fetch@0.7.6"
}
}

20
main.ts
View File

@ -1,18 +1,26 @@
import "https://deno.land/std@0.201.0/dotenv/load.ts"; /// <reference lib="deno.unstable" />
import { handlers, setup } from "std/log"; import "std/dotenv/load.ts";
import { ConsoleHandler } from "std/log/handlers.ts";
import { LevelName, setup } from "std/log/mod.ts";
import { serveUi } from "./api/mod.ts";
import { runAllTasks } from "./app/mod.ts"; import { runAllTasks } from "./app/mod.ts";
import { bot } from "./bot/mod.ts"; import { runBot } from "./bot/mod.ts";
const logLevel = Deno.env.get("LOG_LEVEL")?.toUpperCase() as LevelName ?? "INFO";
// setup logging
setup({ setup({
handlers: { handlers: {
console: new handlers.ConsoleHandler("DEBUG"), console: new ConsoleHandler(logLevel),
}, },
loggers: { loggers: {
default: { level: "DEBUG", handlers: ["console"] }, default: { level: logLevel, handlers: ["console"] },
}, },
}); });
// run parts of the app
await Promise.all([ await Promise.all([
bot.start(), runBot(),
runAllTasks(), runAllTasks(),
serveUi(),
]); ]);

242
ui/AdminsPage.tsx Normal file
View File

@ -0,0 +1,242 @@
import React, { useRef } from "react";
import useSWR, { useSWRConfig } from "swr";
import { AdminData } from "../api/adminsRoute.ts";
import { fetchApi, handleResponse } from "./apiClient.ts";
export function AdminsPage(props: { sessionId: string | null }) {
const { sessionId } = props;
const { mutate } = useSWRConfig();
const addDialogRef = useRef<HTMLDialogElement>(null);
const getSession = useSWR(
sessionId ? ["/sessions/:sessionId", { params: { sessionId } }] as const : null,
(args) => fetchApi(...args).then(handleResponse),
);
const getUser = useSWR(
getSession.data?.userId
? ["/users/:userId", { params: { userId: String(getSession.data.userId) } }] as const
: null,
(args) => fetchApi(...args).then(handleResponse),
);
const getAdmins = useSWR(
["/admins", { method: "GET" }] as const,
(args) => fetchApi(...args).then(handleResponse),
);
return (
<>
{getUser.data && getAdmins.data && getAdmins.data.length === 0 && (
<button
className="button-filled"
onClick={() => {
fetchApi("/admins/promote_self", {
method: "POST",
query: { sessionId: sessionId ?? "" },
}).then(handleResponse).then(() => mutate(() => true));
}}
>
Promote me to admin
</button>
)}
{getAdmins.data?.length
? (
<ul className="flex flex-col gap-2">
{getAdmins.data.map((admin) => (
<AdminListItem key={admin.tgUserId} admin={admin} sessionId={sessionId} />
))}
</ul>
)
: getAdmins.data?.length === 0
? (
<ul className="flex flex-col gap-2">
<li className="flex flex-col gap-2 rounded-md bg-zinc-100 dark:bg-zinc-800 p-2">
<p key="no-admins" className="text-center text-gray-500">No admins.</p>
</li>
</ul>
)
: getAdmins.error
? <p className="alert">Loading admins failed</p>
: <div className="spinner self-center" />}
{getUser.data?.admin && (
<button
className="button-filled"
onClick={() => addDialogRef.current?.showModal()}
>
Add admin
</button>
)}
<AddAdminDialog
dialogRef={addDialogRef}
sessionId={sessionId}
/>
</>
);
}
function AddAdminDialog(props: {
dialogRef: React.RefObject<HTMLDialogElement>;
sessionId: string | null;
}) {
const { dialogRef, sessionId } = props;
const { mutate } = useSWRConfig();
return (
<dialog ref={dialogRef} className="dialog animate-pop-in backdrop-animate-fade-in">
<form
method="dialog"
className="flex flex-col gap-4 p-4"
onSubmit={(e) => {
e.preventDefault();
const data = new FormData(e.target as HTMLFormElement);
fetchApi("/admins", {
method: "POST",
query: { sessionId: sessionId! },
body: {
tgUserId: Number(data.get("tgUserId") as string),
},
}).then(handleResponse).then(() => mutate(() => true));
dialogRef.current?.close();
}}
>
<label className="flex flex-col items-stretch gap-1">
<span className="text-sm">
Telegram user ID
</span>
<input
className="input-text"
type="text"
name="tgUserId"
required
pattern="-?\d+"
/>
</label>
<div className="flex gap-2 items-center justify-end">
<button
type="button"
className="button-outlined"
onClick={() => dialogRef.current?.close()}
>
Close
</button>
<button type="submit" className="button-filled">
Save
</button>
</div>
</form>
</dialog>
);
}
function AdminListItem(props: { admin: AdminData; sessionId: string | null }) {
const { admin, sessionId } = props;
const deleteDialogRef = useRef<HTMLDialogElement>(null);
const getAdminUser = useSWR(
["/users/:userId", { params: { userId: String(admin.tgUserId) } }] as const,
(args) => fetchApi(...args).then(handleResponse),
);
const getSession = useSWR(
sessionId ? ["/sessions/:sessionId", { params: { sessionId } }] as const : null,
(args) => fetchApi(...args).then(handleResponse),
);
const getUser = useSWR(
getSession.data?.userId
? ["/users/:userId", { params: { userId: String(getSession.data.userId) } }] as const
: null,
(args) => fetchApi(...args).then(handleResponse),
);
return (
<li className="flex flex-col gap-2 rounded-md bg-zinc-100 dark:bg-zinc-800 p-2">
<p className="font-bold">
{getAdminUser.data?.first_name ?? admin.tgUserId} {getAdminUser.data?.last_name}{" "}
{getAdminUser.data?.username
? (
<a href={`https://t.me/${getAdminUser.data.username}`} className="link">
@{getAdminUser.data.username}
</a>
)
: null}
</p>
{getAdminUser.data?.bio
? (
<p>
{getAdminUser.data?.bio}
</p>
)
: null}
{getUser.data?.admin && (
<p className="flex gap-2">
<button
className="button-outlined"
onClick={() => deleteDialogRef.current?.showModal()}
>
Delete
</button>
</p>
)}
<DeleteAdminDialog
dialogRef={deleteDialogRef}
adminId={admin.tgUserId}
sessionId={sessionId}
/>
</li>
);
}
function DeleteAdminDialog(props: {
dialogRef: React.RefObject<HTMLDialogElement>;
adminId: number;
sessionId: string | null;
}) {
const { dialogRef, adminId, sessionId } = props;
const { mutate } = useSWRConfig();
return (
<dialog
ref={dialogRef}
className="dialog animate-pop-in backdrop-animate-fade-in"
>
<form
method="dialog"
className="flex flex-col gap-4 p-4"
onSubmit={(e) => {
e.preventDefault();
fetchApi("/admins/:adminId", {
method: "DELETE",
query: { sessionId: sessionId! },
params: { adminId: String(adminId) },
}).then(handleResponse).then(() => mutate(() => true));
dialogRef.current?.close();
}}
>
<p>
Are you sure you want to delete this admin?
</p>
<div className="flex gap-2 items-center justify-end">
<button
type="button"
className="button-outlined"
onClick={() => dialogRef.current?.close()}
>
Close
</button>
<button type="submit" className="button-filled">
Delete
</button>
</div>
</form>
</dialog>
);
}

46
ui/App.tsx Normal file
View File

@ -0,0 +1,46 @@
import React, { useEffect } from "react";
import { Route, Routes } from "react-router-dom";
import { AdminsPage } from "./AdminsPage.tsx";
import { AppHeader } from "./AppHeader.tsx";
import { QueuePage } from "./QueuePage.tsx";
import { SettingsPage } from "./SettingsPage.tsx";
import { StatsPage } from "./StatsPage.tsx";
import { WorkersPage } from "./WorkersPage.tsx";
import { fetchApi, handleResponse } from "./apiClient.ts";
import { useLocalStorage } from "./useLocalStorage.ts";
export function App() {
// store session ID in the local storage
const [sessionId, setSessionId] = useLocalStorage("sessionId");
// initialize a new session when there is no session ID
useEffect(() => {
if (!sessionId) {
fetchApi("/sessions", { method: "POST" }).then((resp) => resp).then(handleResponse).then(
(session) => {
console.log("Initialized session", session.id);
setSessionId(session.id);
},
);
}
}, [sessionId]);
return (
<>
<AppHeader
className="self-stretch"
sessionId={sessionId}
onLogOut={() => setSessionId(null)}
/>
<div className="self-center w-full max-w-screen-md flex flex-col items-stretch gap-4 p-4">
<Routes>
<Route path="/" element={<StatsPage />} />
<Route path="/admins" element={<AdminsPage sessionId={sessionId} />} />
<Route path="/workers" element={<WorkersPage sessionId={sessionId} />} />
<Route path="/queue" element={<QueuePage />} />
<Route path="/settings" element={<SettingsPage sessionId={sessionId} />} />
</Routes>
</div>
</>
);
}

135
ui/AppHeader.tsx Normal file
View File

@ -0,0 +1,135 @@
import React, { ReactNode } from "react";
import { NavLink } from "react-router-dom";
import useSWR from "swr";
import { cx } from "twind/core";
import { API_URL, fetchApi, handleResponse } from "./apiClient.ts";
function NavTab(props: { to: string; children: ReactNode }) {
return (
<NavLink
className={({ isActive }) => cx("tab", isActive && "tab-active")}
to={props.to}
>
{props.children}
</NavLink>
);
}
export function AppHeader(props: {
className?: string;
sessionId: string | null;
onLogOut: () => void;
}) {
const { className, sessionId, onLogOut } = props;
const getSession = useSWR(
sessionId ? ["/sessions/:sessionId", { method: "GET", params: { sessionId } }] as const : null,
(args) => fetchApi(...args).then(handleResponse),
{ onError: () => onLogOut() },
);
const getUser = useSWR(
getSession.data?.userId
? ["/users/:userId", {
method: "GET",
params: { userId: String(getSession.data.userId) },
}] as const
: null,
(args) => fetchApi(...args).then(handleResponse),
);
const getBot = useSWR(
["/bot", { method: "GET" }] as const,
(args) => fetchApi(...args).then(handleResponse),
);
const getUserPhoto = useSWR(
getSession.data?.userId
? ["/users/:userId/photo", {
method: "GET",
params: { userId: String(getSession.data.userId) },
}] as const
: null,
() =>
// elysia fetch can't download file
fetch(`${API_URL}/users/${getSession.data?.userId}/photo`)
.then((response) => {
if (!response.ok) throw new Error(response.statusText);
return response;
})
.then((response) => response.blob())
.then((blob) => blob ? URL.createObjectURL(blob) : null),
);
return (
<header
className={cx(
"bg-zinc-50 dark:bg-zinc-800 shadow-md flex items-center gap-2 px-4 py-1",
className,
)}
>
{/* logo */}
<img className="w-12 h-12" src="/favicon.png" alt="logo" />
{/* tabs */}
<nav className="flex-grow self-stretch flex items-stretch justify-center gap-2">
<NavTab to="/">
Stats
</NavTab>
<NavTab to="/admins">
Admins
</NavTab>
<NavTab to="/workers">
Workers
</NavTab>
<NavTab to="/queue">
Queue
</NavTab>
<NavTab to="/settings">
Settings
</NavTab>
</nav>
{/* loading indicator */}
{getSession.isLoading || getUser.isLoading ? <div className="spinner" /> : null}
{/* user avatar */}
{getUser.data
? getUserPhoto.data
? (
<img
src={getUserPhoto.data}
alt="avatar"
className="w-9 h-9 rounded-full"
/>
)
: (
<div className="w-9 h-9 rounded-full bg-zinc-400 dark:bg-zinc-500 flex items-center justify-center text-white text-2xl font-bold select-none">
{getUser.data.first_name.at(0)?.toUpperCase()}
</div>
)
: null}
{/* login/logout button */}
{!getSession.isLoading && !getUser.isLoading && getBot.data && sessionId
? (
getUser.data
? (
<button className="button-outlined" onClick={() => onLogOut()}>
Logout
</button>
)
: (
<a
className="button-filled"
href={`https://t.me/${getBot.data.username}?start=${sessionId}`}
target="_blank"
>
Login
</a>
)
)
: null}
</header>
);
}

89
ui/Counter.tsx Normal file
View File

@ -0,0 +1,89 @@
import React from "react";
import { cx } from "twind/core";
function CounterDigit(props: { value: number; transitionDurationMs?: number | undefined }) {
const { value, transitionDurationMs = 1500 } = props;
const rads = -(Math.floor(value) % 1_000_000) * 2 * Math.PI * 0.1;
return (
<span className="w-[1em] relative">
{Array.from({ length: 10 }).map((_, i) => (
<span
key={i}
className="absolute inset-0 grid place-items-center transition-transform duration-1000 ease-in-out [backface-visibility:hidden]"
style={{
transform: `rotateX(${rads + i * 2 * Math.PI * 0.1}rad)`,
transformOrigin: `center center 2em`,
transitionDuration: `${transitionDurationMs}ms`,
}}
>
{i}
</span>
))}
</span>
);
}
const Spacer = () =>
<span className="border-l-1 mx-[0.1em] border-zinc-200 dark:border-zinc-700" />;
const CounterText = (props: { children: React.ReactNode }) => (
<span className="self-center px-[0.1em]">
{props.children}
</span>
);
export function Counter(props: {
value: number;
digits: number;
fractionDigits?: number | undefined;
transitionDurationMs?: number | undefined;
className?: string | undefined;
postfix?: string | undefined;
}) {
const { value, digits, fractionDigits = 0, transitionDurationMs, className, postfix } = props;
return (
<span
className={cx(
"inline-flex h-[1.5em] items-stretch border-1 border-zinc-300 dark:border-zinc-700 shadow-inner rounded-md overflow-hidden",
className,
)}
>
{Array.from({ length: digits })
.flatMap((_, i) => [
i > 0 && i % 3 === 0 ? <Spacer key={`spacer${i}`} /> : null,
<CounterDigit
key={`digit${i}`}
value={value / 10 ** i}
transitionDurationMs={transitionDurationMs}
/>,
])
.reverse()}
{fractionDigits > 0 && (
<>
<Spacer />
<CounterText>.</CounterText>
<Spacer />
{Array.from({ length: fractionDigits })
.flatMap((_, i) => [
i > 0 && i % 3 === 0 ? <Spacer key={`fractionSpacer${i}`} /> : null,
<CounterDigit
key={`fractionDigit${i}`}
value={value * 10 ** (i + 1)}
transitionDurationMs={transitionDurationMs}
/>,
])}
</>
)}
{postfix && (
<>
<Spacer />
<CounterText>{postfix}</CounterText>
</>
)}
</span>
);
}

27
ui/Progress.tsx Normal file
View File

@ -0,0 +1,27 @@
import React from "react";
import { cx } from "twind/core";
export function Progress(props: { value: number; className?: string }) {
const { value, className } = props;
return (
<div
className={cx(
"flex items-stretch overflow-hidden rounded-md bg-zinc-200 text-xs text-white dark:bg-zinc-800",
className,
)}
>
<div
className="bg-stripes flex items-center justify-center overflow-hidden rounded-md bg-sky-500 transition-[flex-grow] duration-1000"
style={{ flexGrow: value }}
>
{(value * 100).toFixed(0)}%
</div>
<div
className="flex items-center justify-center overflow-hidden transition-[flex-grow] duration-500"
style={{ flexGrow: 1 - value }}
>
</div>
</div>
);
}

57
ui/QueuePage.tsx Normal file
View File

@ -0,0 +1,57 @@
import React from "react";
import FlipMove from "react-flip-move";
import useSWR from "swr";
import { getFlagEmoji } from "../utils/getFlagEmoji.ts";
import { Progress } from "./Progress.tsx";
import { fetchApi, handleResponse } from "./apiClient.ts";
export function QueuePage() {
const getJobs = useSWR(
["/jobs", {}] as const,
(args) => fetchApi(...args).then(handleResponse),
{ refreshInterval: 2000 },
);
return (
<FlipMove
typeName={"ul"}
className="flex flex-col items-stretch gap-2 p-2 bg-zinc-200 dark:bg-zinc-800 rounded-md"
enterAnimation="fade"
leaveAnimation="fade"
>
{getJobs.data && getJobs.data.length === 0
? <li key="no-jobs" className="text-center text-gray-500">Queue is empty.</li>
: (
getJobs.data?.map((job) => (
<li
className="flex items-baseline gap-2 bg-zinc-100 dark:bg-zinc-700 px-2 py-1 rounded-md"
key={job.id}
>
<span className="">{job.place}.</span>
<span>{getFlagEmoji(job.state.from.language_code ?? undefined)}</span>
<strong>{job.state.from.first_name} {job.state.from.last_name}</strong>
{job.state.from.username
? (
<a
className="link"
href={`https://t.me/${job.state.from.username}`}
target="_blank"
>
@{job.state.from.username}
</a>
)
: null}
<span className="flex-grow self-center h-full">
{job.state.progress != null && (
<Progress className="w-full h-full" value={job.state.progress} />
)}
</span>
<span className="text-sm text-zinc-500 dark:text-zinc-400">
{job.state.workerInstanceKey}
</span>
</li>
))
)}
</FlipMove>
);
}

277
ui/SettingsPage.tsx Normal file
View File

@ -0,0 +1,277 @@
import React, { useState } from "react";
import useSWR from "swr";
import { cx } from "twind/core";
import { fetchApi, handleResponse } from "./apiClient.ts";
import { omitUndef } from "../utils/omitUndef.ts";
export function SettingsPage(props: { sessionId: string | null }) {
const { sessionId } = props;
const getSession = useSWR(
sessionId ? ["/sessions/:sessionId", { params: { sessionId } }] as const : null,
(args) => fetchApi(...args).then(handleResponse),
);
const getUser = useSWR(
getSession.data?.userId
? ["/users/:userId", { params: { userId: String( getSession.data.userId) } }] as const
: null,
(args) => fetchApi(...args).then(handleResponse),
);
const getParams = useSWR(
["/settings/params", {}] as const,
(args) => fetchApi(...args).then(handleResponse),
);
const [newParams, setNewParams] = useState<Partial<typeof getParams.data>>({});
const [patchParamsError, setPatchParamsError] = useState<string>();
return (
<form
className="flex flex-col items-stretch gap-4"
onSubmit={(e) => {
e.preventDefault();
getParams.mutate(() =>
fetchApi("/settings/params", {
method: "PATCH",
query: { sessionId: sessionId ?? "" },
body: omitUndef(newParams ?? {}),
}).then(handleResponse)
)
.then(() => setNewParams({}))
.catch((e) => setPatchParamsError(String(e)));
}}
>
<label className="flex flex-col items-stretch gap-1">
<span className="text-sm">
Negative prompt {newParams?.negative_prompt != null ? "(Changed)" : ""}
</span>
<textarea
className="input-text"
disabled={getParams.isLoading || !getUser.data?.admin}
value={newParams?.negative_prompt ??
getParams.data?.negative_prompt ??
""}
onChange={(e) =>
setNewParams((params) => ({
...params,
negative_prompt: e.target.value,
}))}
/>
</label>
<label className="flex flex-col items-stretch gap-1">
<span className="text-sm">
Sampler {newParams?.sampler_name != null ? "(Changed)" : ""}
</span>
<input
className="input-text"
disabled={getParams.isLoading || !getUser.data?.admin}
value={newParams?.sampler_name ??
getParams.data?.sampler_name ??
""}
onChange={(e) =>
setNewParams((params) => ({
...params,
sampler_name: e.target.value,
}))}
/>
</label>
<label className="flex flex-col items-stretch gap-1">
<span className="text-sm">
Steps {newParams?.steps != null ? "(Changed)" : ""}
</span>
<span className="flex items-center gap-1">
<input
className="input-text w-20"
type="number"
min={5}
max={50}
step={5}
disabled={getParams.isLoading || !getUser.data?.admin}
value={newParams?.steps ??
getParams.data?.steps ??
0}
onChange={(e) =>
setNewParams((params) => ({
...params,
steps: Number(e.target.value),
}))}
/>
<input
className="input-range flex-grow"
type="range"
min={5}
max={50}
step={5}
disabled={getParams.isLoading || !getUser.data?.admin}
value={newParams?.steps ??
getParams.data?.steps ??
0}
onChange={(e) =>
setNewParams((params) => ({
...params,
steps: Number(e.target.value),
}))}
/>
</span>
</label>
<label className="flex flex-col items-stretch gap-1">
<span className="text-sm">
Detail {newParams?.cfg_scale != null ? "(Changed)" : ""}
</span>
<span className="flex items-center gap-1">
<input
className="input-text w-20"
type="number"
min={1}
max={20}
step={1}
disabled={getParams.isLoading || !getUser.data?.admin}
value={newParams?.cfg_scale ??
getParams.data?.cfg_scale ??
0}
onChange={(e) =>
setNewParams((params) => ({
...params,
cfg_scale: Number(e.target.value),
}))}
/>
<input
className="input-range flex-grow"
type="range"
min={1}
max={20}
step={1}
disabled={getParams.isLoading || !getUser.data?.admin}
value={newParams?.cfg_scale ??
getParams.data?.cfg_scale ??
0}
onChange={(e) =>
setNewParams((params) => ({
...params,
cfg_scale: Number(e.target.value),
}))}
/>
</span>
</label>
<div className="flex gap-4">
<label className="flex flex-1 flex-col items-stretch gap-1">
<span className="text-sm">
Width {newParams?.width != null ? "(Changed)" : ""}
</span>
<input
className="input-text"
type="number"
min={64}
max={2048}
step={64}
disabled={getParams.isLoading || !getUser.data?.admin}
value={newParams?.width ??
getParams.data?.width ??
0}
onChange={(e) =>
setNewParams((params) => ({
...params,
width: Number(e.target.value),
}))}
/>
<input
className="input-range"
type="range"
min={64}
max={2048}
step={64}
disabled={getParams.isLoading || !getUser.data?.admin}
value={newParams?.width ??
getParams.data?.width ??
0}
onChange={(e) =>
setNewParams((params) => ({
...params,
width: Number(e.target.value),
}))}
/>
</label>
<label className="flex flex-1 flex-col items-stretch gap-1">
<span className="text-sm">
Height {newParams?.height != null ? "(Changed)" : ""}
</span>
<input
className="input-text"
type="number"
min={64}
max={2048}
step={64}
disabled={getParams.isLoading || !getUser.data?.admin}
value={newParams?.height ??
getParams.data?.height ??
0}
onChange={(e) =>
setNewParams((params) => ({
...params,
height: Number(e.target.value),
}))}
/>
<input
className="input-range"
type="range"
min={64}
max={2048}
step={64}
disabled={getParams.isLoading || !getUser.data?.admin}
value={newParams?.height ??
getParams.data?.height ??
0}
onChange={(e) =>
setNewParams((params) => ({
...params,
height: Number(e.target.value),
}))}
/>
</label>
</div>
{patchParamsError
? (
<p className="alert">
<span className="flex-grow">Updating params failed: {patchParamsError}</span>
<button className="button-ghost" onClick={() => setPatchParamsError(undefined)}>
Close
</button>
</p>
)
: null}
{getParams.error
? (
<p className="alert">
<span className="flex-grow">
Loading params failed: {String(getParams.error)}
</span>
</p>
)
: null}
<div className="flex gap-2 items-center justify-end">
<button
type="button"
className={cx("button-outlined ripple", getParams.isLoading && "bg-stripes")}
disabled={getParams.isLoading || !getUser.data?.admin ||
Object.keys(newParams ?? {}).length === 0}
onClick={() => setNewParams({})}
>
Reset
</button>
<button
type="submit"
className={cx("button-filled ripple", getParams.isLoading && "bg-stripes")}
disabled={getParams.isLoading || !getUser.data?.admin ||
Object.keys(newParams ?? {}).length === 0}
>
Save
</button>
</div>
</form>
);
}

94
ui/StatsPage.tsx Normal file
View File

@ -0,0 +1,94 @@
import React from "react";
import { fetchApi, handleResponse } from "./apiClient.ts";
import useSWR from "swr";
import { Counter } from "./Counter.tsx";
export function StatsPage() {
const getGlobalStats = useSWR(
["/stats", {}] as const,
(args) => fetchApi(...args).then(handleResponse),
{ refreshInterval: 2_000 },
);
return (
<div className="my-16 flex flex-col gap-16 text-zinc-600 dark:text-zinc-400">
<p className="flex flex-col items-center gap-2 text-2xl sm:text-3xl">
<span>Pixelsteps diffused</span>
<Counter
className="font-bold text-zinc-700 dark:text-zinc-300"
value={getGlobalStats.data?.pixelStepCount ?? 0}
digits={15}
transitionDurationMs={3_000}
/>
<Counter
className="text-base"
value={(getGlobalStats.data?.pixelStepsPerMinute ?? 0) / 60}
digits={9}
transitionDurationMs={2_000}
postfix="/s"
/>
</p>
<p className="flex flex-col items-center gap-2 text-2xl sm:text-3xl">
<span>Pixels painted</span>
<Counter
className="font-bold text-zinc-700 dark:text-zinc-300"
value={getGlobalStats.data?.pixelCount ?? 0}
digits={15}
transitionDurationMs={3_000}
/>
<Counter
className="text-base"
value={(getGlobalStats.data?.pixelsPerMinute ?? 0) / 60}
digits={9}
transitionDurationMs={2_000}
postfix="/s"
/>
</p>
<div className="flex flex-col md:flex-row gap-16 md:gap-8">
<p className="flex-grow flex flex-col items-center gap-2 text-2xl">
<span>Steps processed</span>
<Counter
className="font-bold text-zinc-700 dark:text-zinc-300"
value={getGlobalStats.data?.stepCount ?? 0}
digits={9}
transitionDurationMs={3_000}
/>
<Counter
className="text-base"
value={(getGlobalStats.data?.stepsPerMinute ?? 0) / 60}
digits={3}
fractionDigits={3}
transitionDurationMs={2_000}
postfix="/s"
/>
</p>
<p className="flex-grow flex flex-col items-center gap-2 text-2xl">
<span>Images generated</span>
<Counter
className="font-bold text-zinc-700 dark:text-zinc-300"
value={getGlobalStats.data?.imageCount ?? 0}
digits={9}
transitionDurationMs={3_000}
/>
<Counter
className="text-base"
value={(getGlobalStats.data?.imagesPerMinute ?? 0) / 60}
digits={3}
fractionDigits={3}
transitionDurationMs={2_000}
postfix="/s"
/>
</p>
</div>
<p className="flex-grow flex flex-col items-center gap-2 text-2xl">
<span>Unique users</span>
<Counter
className="font-bold text-zinc-700 dark:text-zinc-300"
value={getGlobalStats.data?.userCount ?? 0}
digits={6}
transitionDurationMs={1_500}
/>
</p>
</div>
);
}

386
ui/WorkersPage.tsx Normal file
View File

@ -0,0 +1,386 @@
import React, { RefObject, useRef } from "react";
import { FormattedRelativeTime } from "react-intl";
import useSWR, { useSWRConfig } from "swr";
import { WorkerResponse } from "../api/workersRoute.ts";
import { Counter } from "./Counter.tsx";
import { fetchApi, handleResponse } from "./apiClient.ts";
export function WorkersPage(props: { sessionId: string | null }) {
const { sessionId } = props;
const addDialogRef = useRef<HTMLDialogElement>(null);
const getSession = useSWR(
sessionId ? ["/sessions/:sessionId", { params: { sessionId } }] as const : null,
(args) => fetchApi(...args).then(handleResponse),
);
const getUser = useSWR(
getSession.data?.userId
? ["/users/:userId", { params: { userId: String(getSession.data.userId) } }] as const
: null,
(args) => fetchApi(...args).then(handleResponse),
);
const getWorkers = useSWR(
["/workers", { method: "GET" }] as const,
(args) => fetchApi(...args).then(handleResponse),
{ refreshInterval: 5000 },
);
return (
<>
{getWorkers.data?.length
? (
<ul className="flex flex-col gap-2">
{getWorkers.data?.map((worker) => (
<WorkerListItem key={worker.id} worker={worker} sessionId={sessionId} />
))}
</ul>
)
: getWorkers.data?.length === 0
? (
<ul className="flex flex-col gap-2">
<li className="flex flex-col gap-2 rounded-md bg-zinc-100 dark:bg-zinc-800 p-2">
<p key="no-workers" className="text-center text-gray-500">No workers.</p>
</li>
</ul>
)
: getWorkers.error
? <p className="alert">Loading workers failed</p>
: <div className="spinner self-center" />}
{getUser.data?.admin && (
<button
className="button-filled"
onClick={() => addDialogRef.current?.showModal()}
>
Add worker
</button>
)}
<AddWorkerDialog
dialogRef={addDialogRef}
sessionId={sessionId}
/>
</>
);
}
function AddWorkerDialog(props: {
dialogRef: RefObject<HTMLDialogElement>;
sessionId: string | null;
}) {
const { dialogRef, sessionId } = props;
const { mutate } = useSWRConfig();
return (
<dialog ref={dialogRef} className="dialog animate-pop-in backdrop-animate-fade-in">
<form
method="dialog"
className="flex flex-col gap-4 p-4"
onSubmit={(e) => {
e.preventDefault();
const data = new FormData(e.target as HTMLFormElement);
const key = data.get("key") as string;
const name = data.get("name") as string;
const sdUrl = data.get("url") as string;
const user = data.get("user") as string;
const password = data.get("password") as string;
console.log(key, name, user, password);
fetchApi("/workers", {
method: "POST",
query: { sessionId: sessionId! },
body: {
key,
name: name || null,
sdUrl,
sdAuth: user && password ? { user, password } : null,
},
}).then(handleResponse).then(() => mutate(() => true));
dialogRef.current?.close();
}}
>
<div className="flex gap-4">
<label className="flex flex-1 flex-col items-stretch gap-1">
<span className="text-sm">
Key
</span>
<input
className="input-text"
type="text"
name="key"
required
/>
<span className="text-sm text-zinc-500">
Used for counting statistics
</span>
</label>
<label className="flex flex-1 flex-col items-stretch gap-1">
<span className="text-sm">
Name
</span>
<input
className="input-text"
type="text"
name="name"
/>
<span className="text-sm text-zinc-500">
Used for display
</span>
</label>
</div>
<label className="flex flex-col items-stretch gap-1">
<span className="text-sm">
URL
</span>
<input
className="input-text"
type="url"
name="url"
required
pattern="https?://.*"
/>
</label>
<div className="flex gap-4">
<label className="flex flex-1 flex-col items-stretch gap-1">
<span className="text-sm">
User
</span>
<input
className="input-text"
type="text"
name="user"
/>
</label>
<label className="flex flex-1 flex-col items-stretch gap-1">
<span className="text-sm">
Password
</span>
<input
className="input-text"
type="password"
name="password"
/>
</label>
</div>
<div className="flex gap-2 items-center justify-end">
<button
type="button"
className="button-outlined"
onClick={() => dialogRef.current?.close()}
>
Close
</button>
<button type="submit" disabled={!sessionId} className="button-filled">
Save
</button>
</div>
</form>
</dialog>
);
}
function WorkerListItem(props: {
worker: WorkerResponse;
sessionId: string | null;
}) {
const { worker, sessionId } = props;
const editDialogRef = useRef<HTMLDialogElement>(null);
const deleteDialogRef = useRef<HTMLDialogElement>(null);
const getSession = useSWR(
sessionId ? ["/sessions/:sessionId", { params: { sessionId } }] as const : null,
(args) => fetchApi(...args).then(handleResponse),
);
const getUser = useSWR(
getSession.data?.userId
? ["/users/:userId", { params: { userId: String(getSession.data.userId) } }] as const
: null,
(args) => fetchApi(...args).then(handleResponse),
);
return (
<li className="flex flex-col gap-2 rounded-md bg-zinc-100 dark:bg-zinc-800 p-2">
<p className="font-bold">
{worker.name ?? worker.key}
</p>
{worker.isActive ? <p> Active</p> : (
<>
<p>
Last seen {worker.lastOnlineTime
? (
<FormattedRelativeTime
value={(worker.lastOnlineTime - Date.now()) / 1000}
numeric="auto"
updateIntervalInSeconds={1}
/>
)
: "never"}
</p>
{worker.lastError && (
<p className="text-red-500">
{worker.lastError.message} (
<FormattedRelativeTime
value={(worker.lastError.time - Date.now()) / 1000}
numeric="auto"
updateIntervalInSeconds={1}
/>)
</p>
)}
</>
)}
<p className="flex gap-1">
<Counter value={worker.imagesPerMinute} digits={2} fractionDigits={1} /> images per minute,
{" "}
<Counter value={worker.stepsPerMinute} digits={3} /> steps per minute
</p>
{getUser.data?.admin && (
<p className="flex gap-2">
<button
className="button-outlined"
onClick={() => editDialogRef.current?.showModal()}
>
Settings
</button>
<button
className="button-outlined"
onClick={() => deleteDialogRef.current?.showModal()}
>
Delete
</button>
</p>
)}
<EditWorkerDialog
dialogRef={editDialogRef}
workerId={worker.id}
sessionId={sessionId}
/>
<DeleteWorkerDialog
dialogRef={deleteDialogRef}
workerId={worker.id}
sessionId={sessionId}
/>
</li>
);
}
function EditWorkerDialog(props: {
dialogRef: RefObject<HTMLDialogElement>;
workerId: string;
sessionId: string | null;
}) {
const { dialogRef, workerId, sessionId } = props;
const { mutate } = useSWRConfig();
return (
<dialog
className="dialog animate-pop-in backdrop-animate-fade-in"
ref={dialogRef}
>
<form
method="dialog"
className="flex flex-col gap-4 p-4"
onSubmit={(e) => {
e.preventDefault();
const data = new FormData(e.target as HTMLFormElement);
const user = data.get("user") as string;
const password = data.get("password") as string;
console.log(user, password);
fetchApi("/workers/:workerId", {
method: "PATCH",
params: { workerId: workerId },
query: { sessionId: sessionId! },
body: {
sdAuth: user && password ? { user, password } : null,
},
}).then(handleResponse).then(() => mutate(() => true));
dialogRef.current?.close();
}}
>
<div className="flex gap-4">
<label className="flex flex-1 flex-col items-stretch gap-1">
<span className="text-sm">
User
</span>
<input
className="input-text"
type="text"
name="user"
/>
</label>
<label className="flex flex-1 flex-col items-stretch gap-1">
<span className="text-sm">
Password
</span>
<input
className="input-text"
type="password"
name="password"
/>
</label>
</div>
<div className="flex gap-2 items-center justify-end">
<button
type="button"
className="button-outlined"
onClick={() => dialogRef.current?.close()}
>
Close
</button>
<button type="submit" disabled={!sessionId} className="button-filled">
Save
</button>
</div>
</form>
</dialog>
);
}
function DeleteWorkerDialog(props: {
dialogRef: RefObject<HTMLDialogElement>;
workerId: string;
sessionId: string | null;
}) {
const { dialogRef, workerId, sessionId } = props;
const { mutate } = useSWRConfig();
return (
<dialog
className="dialog animate-pop-in backdrop-animate-fade-in"
ref={dialogRef}
>
<form
method="dialog"
className="flex flex-col gap-4 p-4"
onSubmit={(e) => {
e.preventDefault();
fetchApi("/workers/:workerId", {
method: "DELETE",
params: { workerId: workerId },
query: { sessionId: sessionId! },
}).then(handleResponse).then(() => mutate(() => true));
dialogRef.current?.close();
}}
>
<p>
Are you sure you want to delete this worker?
</p>
<div className="flex gap-2 items-center justify-end">
<button
type="button"
className="button-outlined"
onClick={() => dialogRef.current?.close()}
>
Close
</button>
<button type="submit" disabled={!sessionId} className="button-filled">
Delete
</button>
</div>
</form>
</dialog>
);
}

19
ui/apiClient.ts Normal file
View File

@ -0,0 +1,19 @@
import { edenFetch } from "elysia/eden";
import { Api } from "../api/serveApi.ts";
export const API_URL = "/api";
export const fetchApi = edenFetch<Api>(API_URL);
export function handleResponse<
T extends
| { data: unknown; error: null }
| { data: null; error: { status: number; value: unknown } },
>(
response: T,
): (T & { error: null })["data"] {
if (response.error) {
throw new Error(`${response.error?.status}: ${response.error?.value}`);
}
return response.data;
}

BIN
ui/favicon.png Normal file

Binary file not shown.

After

Width:  |  Height:  |  Size: 37 KiB

13
ui/index.html Normal file
View File

@ -0,0 +1,13 @@
<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="UTF-8">
<meta name="viewport" content="width=device-width, initial-scale=1.0">
<link rel="icon" href="/favicon.png" type="image/png">
<script type="module" src="/main.tsx"></script>
<title>Eris</title>
</head>
<body>
</body>
</html>

16
ui/main.tsx Normal file
View File

@ -0,0 +1,16 @@
/// <reference lib="deno.unstable" />
/// <reference lib="dom" />
import { createRoot } from "react-dom/client";
import React from "react";
import { App } from "./App.tsx";
import "./twind.ts";
import { BrowserRouter } from "react-router-dom";
import { IntlProvider } from "react-intl";
createRoot(document.body).render(
<BrowserRouter>
<IntlProvider locale="en">
<App />
</IntlProvider>
</BrowserRouter>,
);

2
ui/robots.txt Normal file
View File

@ -0,0 +1,2 @@
User-agent: *
Disallow: /

143
ui/twind.ts Normal file
View File

@ -0,0 +1,143 @@
import { defineConfig, injectGlobal, install } from "twind/core";
import presetTailwind from "twind/preset-tailwind";
const twConfig = defineConfig({
presets: [presetTailwind()],
hash: false,
});
install(twConfig, false);
injectGlobal`
@layer base {
html {
@apply h-full bg-white text-zinc-900 dark:bg-zinc-900 dark:text-zinc-100;
}
body {
@apply min-h-full flex flex-col;
}
}
@layer utilities {
.bg-stripes {
background-image: repeating-linear-gradient(
45deg,
transparent,
transparent 14px,
rgba(0, 0, 0, 0.1) 14px,
rgba(0, 0, 0, 0.1) 28px
);
animation: bg-scroll 0.5s linear infinite;
background-size: 40px 40px;
}
@keyframes bg-scroll {
to {
background-position: 40px 0;
}
}
.ripple {
position: relative;
}
.ripple:not([disabled])::after {
content: "";
position: absolute;
inset: 0;
opacity: 0;
background-image: radial-gradient(circle, rgba(255, 255, 255, 0.2) 10%, transparent 10%);
background-size: 1500%;
background-repeat: no-repeat;
background-position: center;
transition:
background 0.4s,
opacity 0.7s;
}
.ripple:not([disabled]):active::after {
opacity: 1;
background-size: 0%;
transition:
background 0s,
opacity 0s;
}
.backdrop-animate-fade-in::backdrop {
animation: fade-in 0.3s ease-out forwards;
}
@keyframes fade-in {
from {
opacity: 0;
}
}
.animate-pop-in {
animation: pop-in 0.2s cubic-bezier(0.175, 0.885, 0.32, 1.275) forwards;
}
@keyframes pop-in {
from {
transform: scale(0.8);
opacity: 0;
}
}
}
@layer components {
.link {
@apply text-sky-600 dark:text-sky-500 rounded-sm focus:outline focus:outline-2 focus:outline-offset-1 focus:outline-sky-600 dark:focus:outline-sky-500;
}
.spinner {
@apply h-8 w-8 animate-spin rounded-full border-4 border-transparent border-t-current;
}
.alert {
@apply px-4 py-2 flex gap-2 items-center bg-red-500 text-white rounded-sm shadow-md;
}
.button-filled {
@apply rounded-md bg-sky-600 px-3 py-2 min-h-12
text-sm font-semibold uppercase tracking-wider text-white shadow-sm transition-all
hover:bg-sky-500 focus:outline focus:outline-2 focus:outline-offset-2 focus:outline-sky-600
disabled:bg-zinc-300 dark:disabled:bg-zinc-700 dark:disabled:text-zinc-500;
}
.button-outlined {
@apply rounded-md bg-transparent px-3 py-2 min-h-12 ring-1 ring-inset ring-sky-600/100
text-sm font-semibold uppercase tracking-wider text-sky-600 transition-all
hover:ring-sky-500/100 hover:text-sky-500 focus:outline focus:outline-2 focus:outline-offset-2 focus:outline-sky-600
disabled:ring-zinc-300 disabled:text-zinc-300 dark:disabled:ring-zinc-700 dark:disabled:text-zinc-500;
}
.button-ghost {
@apply rounded-md bg-transparent px-3 py-2 min-h-12 ring-1 ring-inset ring-transparent
text-sm font-semibold uppercase tracking-wider transition-all
hover:ring-current focus:outline focus:outline-2 focus:outline-offset-2 focus:outline-current
disabled:text-zinc-500;
}
.tab {
@apply inline-flex items-center px-4 text-sm text-center bg-transparent border-b-2 sm:text-base whitespace-nowrap outline-none
border-transparent dark:text-white cursor-base hover:border-zinc-400 focus:border-zinc-400 transition-all;
}
.tab-active {
@apply text-sky-600 border-sky-500 dark:border-sky-600 dark:text-sky-500 hover:border-sky-400 focus:border-sky-400;
}
.input-text {
@apply block appearance-none rounded-md border-none bg-transparent px-3 py-1.5
text-zinc-900 shadow-sm outline-none ring-1 ring-inset ring-zinc-400
placeholder:text-zinc-500 focus:ring-2 focus:ring-sky-600/100
dark:text-zinc-100 dark:ring-zinc-600;
}
.input-range {
@apply h-6 cursor-pointer accent-sky-600;
}
.dialog {
@apply overflow-hidden overflow-y-auto rounded-md shadow-xl
bg-white text-zinc-900 shadow-lg backdrop:bg-black/30 dark:bg-zinc-800 dark:text-zinc-100;
}
}
`;

31
ui/useLocalStorage.ts Normal file
View File

@ -0,0 +1,31 @@
import { useEffect, useState } from "react";
export function useLocalStorage(key: string) {
const [value, setValue] = useState(() => window.localStorage.getItem(key));
useEffect(() => {
console.log(key, value);
if (value != null) {
window.localStorage.setItem(key, value);
} else {
window.localStorage.removeItem(key);
}
}, [key, value]);
useEffect(() => {
const handleStorage = (event: StorageEvent) => {
if (event.key === key) {
console.log(key, event.newValue);
setValue(event.newValue);
}
};
window.addEventListener("storage", handleStorage);
return () => {
window.removeEventListener("storage", handleStorage);
};
}, [key]);
return [value, setValue] as const;
}

76
utils/Tkv.ts Normal file
View File

@ -0,0 +1,76 @@
export type TkvEntry<K extends Deno.KvKey, T> = {
key: readonly [...K];
value: T;
versionstamp: string;
};
export type TkvEntryMaybe<K extends Deno.KvKey, T> = TkvEntry<K, T> | {
key: readonly [...K];
value: null;
versionstamp: null;
};
export type TkvListSelector<K extends Deno.KvKey> =
| { prefix: KvKeyPrefix<K> }
| { prefix: KvKeyPrefix<K>; start: readonly [...K] }
| { prefix: KvKeyPrefix<K>; end: readonly [...K] }
| { start: readonly [...K]; end: readonly [...K] };
export type KvKeyPrefix<Key extends Deno.KvKey> = Key extends readonly [infer Prefix, ...infer Rest]
? readonly [Prefix] | readonly [Prefix, ...Rest]
: never;
/**
* Typed wrapper for {@link Deno.Kv}
*/
export class Tkv<K extends Deno.KvKey, T> {
constructor(readonly db: Deno.Kv) {}
get(
key: readonly [...K],
options: Parameters<Deno.Kv["get"]>[1] = {},
): Promise<TkvEntryMaybe<K, T>> {
return this.db.get<T>(key, options) as any;
}
set(
key: readonly [...K],
value: T,
options: Parameters<Deno.Kv["set"]>[2] = {},
): ReturnType<Deno.Kv["set"]> {
return this.db.set(key, value, options);
}
atomicSet(
key: readonly [...K],
versionstamp: Parameters<Deno.AtomicOperation["check"]>[0]["versionstamp"],
value: T,
options: Parameters<Deno.AtomicOperation["set"]>[2] = {},
): ReturnType<Deno.AtomicOperation["commit"]> {
return this.db.atomic()
.check({ key, versionstamp })
.set(key, value, options)
.commit();
}
delete(key: readonly [...K]): ReturnType<Deno.Kv["delete"]> {
return this.db.delete(key);
}
atomicDelete(
key: readonly [...K],
versionstamp: Parameters<Deno.AtomicOperation["check"]>[0]["versionstamp"],
): ReturnType<Deno.AtomicOperation["commit"]> {
return this.db.atomic()
.check({ key, versionstamp })
.delete(key)
.commit();
}
list(
selector: TkvListSelector<K>,
options: Parameters<Deno.Kv["list"]>[1] = {},
): AsyncIterableIterator<TkvEntry<K, T>> {
return this.db.list<T>(selector as Deno.KvListSelector, options) as any;
}
}

View File

@ -1,7 +1,11 @@
import { Chat, User } from "grammy_types"; import { Chat, User } from "grammy_types";
export function formatUserChat( export function formatUserChat(
ctx: { from?: User; chat?: Chat; sdInstanceId?: string }, ctx: {
from?: User | undefined;
chat?: Chat | undefined;
workerInstanceKey?: string | undefined;
},
) { ) {
const msg: string[] = []; const msg: string[] = [];
if (ctx.from) { if (ctx.from) {
@ -26,8 +30,8 @@ export function formatUserChat(
} }
} }
} }
if (ctx.sdInstanceId) { if (ctx.workerInstanceKey) {
msg.push(`using ${ctx.sdInstanceId}`); msg.push(`using ${ctx.workerInstanceKey}`);
} }
return msg.join(" "); return msg.join(" ");
} }

4
utils/getAuthHeader.ts Normal file
View File

@ -0,0 +1,4 @@
export function getAuthHeader(auth: { user: string; password: string } | null) {
if (!auth) return {};
return { "Authorization": `Basic ${btoa(`${auth.user}:${auth.password}`)}` };
}

View File

@ -42,10 +42,12 @@ const languageToFlagMap: Record<string, string> = {
"id": "🇮🇩", // indonesian - indonesia "id": "🇮🇩", // indonesian - indonesia
"is": "🇮🇸", // icelandic - iceland "is": "🇮🇸", // icelandic - iceland
"lb": "🇱🇺", // luxembourgish - luxembourg "lb": "🇱🇺", // luxembourgish - luxembourg
"ca": "🇪🇸", // catalan - spain
}; };
export function getFlagEmoji(languageCode?: string): string | undefined { export function getFlagEmoji(languageCode?: string): string | undefined {
const language = languageCode?.split("-").pop()?.toLowerCase(); if (!languageCode) return;
if (!language) return; const language = languageCode.split("-").shift()?.toLowerCase();
return languageToFlagMap[language]; if (language == null) return;
return languageToFlagMap[language] ?? `[${language.toUpperCase()}]`;
} }

68
utils/kvMemoize.ts Normal file
View File

@ -0,0 +1,68 @@
export interface KvMemoizeOptions<A extends Deno.KvKey, R> {
/**
* The time in milliseconds until the cached result expires.
*/
expireIn?: ((result: R, ...args: A) => number) | number | undefined;
/**
* Whether to recalculate the result if it was already cached.
*
* Runs whenever the result is retrieved from the cache.
*/
shouldRecalculate?: ((result: R, ...args: A) => boolean) | undefined;
/**
* Whether to cache the result after computing it.
*
* Runs whenever a new result is computed.
*/
shouldCache?: ((result: R, ...args: A) => boolean) | undefined;
/**
* Override the default KV store functions.
*/
override?: {
set: (
key: Deno.KvKey,
args: A,
value: R,
options: { expireIn?: number },
) => Promise<void>;
get: (key: Deno.KvKey, args: A) => Promise<R | undefined>;
};
}
/**
* Memoizes the function result in KV store.
*/
export function kvMemoize<A extends Deno.KvKey, R>(
db: Deno.Kv,
key: Deno.KvKey,
fn: (...args: A) => Promise<R>,
options?: KvMemoizeOptions<A, R>,
): (...args: A) => Promise<R> {
return async (...args) => {
const cachedResult = options?.override?.get
? await options.override.get(key, args)
: (await db.get<R>([...key, ...args])).value;
if (cachedResult != null) {
if (!options?.shouldRecalculate?.(cachedResult, ...args)) {
return cachedResult;
}
}
const result = await fn(...args);
const expireIn = typeof options?.expireIn === "function"
? options.expireIn(result, ...args)
: options?.expireIn;
if (options?.shouldCache?.(result, ...args) ?? (result != null)) {
if (options?.override?.set) {
await options.override.set(key, args, result, expireIn != null ? { expireIn } : {});
} else {
await db.set([...key, ...args], result, expireIn != null ? { expireIn } : {});
}
}
return result;
};
}

11
utils/omitUndef.ts Normal file
View File

@ -0,0 +1,11 @@
/**
* Removes all undefined properties from an object.
*/
export function omitUndef<O extends object | undefined>(object: O):
& { [K in keyof O as undefined extends O[K] ? never : K]: O[K] }
& { [K in keyof O as undefined extends O[K] ? K : never]?: O[K] & ({} | null) } {
if (object == undefined) return object as never;
return Object.fromEntries(
Object.entries(object).filter(([, v]) => v !== undefined),
) as never;
}