diff --git a/README.md b/README.md index a650efa..0f5e9aa 100644 --- a/README.md +++ b/README.md @@ -14,12 +14,16 @@ 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). Required. - `TG_ADMIN_USERNAMES` - Comma separated list of usernames of users that can use admin commands. +- `LOG_LEVEL` - [Log level](https://deno.land/std@0.201.0/log/mod.ts?s=LogLevels). Default: `INFO`. ## Running - Start stable diffusion webui: `cd sd-webui`, `./webui.sh --api` - Start bot: `deno task start` +To connect your SD to the bot, open the [Eris UI](http://localhost:5999/), login as admin and add a +worker. + ## Codegen The Stable Diffusion API in `app/sdApi.ts` is auto-generated. To regenerate it, first start your SD diff --git a/api/paramsRoute.ts b/api/paramsRoute.ts index 77e3284..6532795 100644 --- a/api/paramsRoute.ts +++ b/api/paramsRoute.ts @@ -1,12 +1,10 @@ import { deepMerge } from "std/collections/deep_merge.ts"; -import { getLogger } from "std/log/mod.ts"; +import { info } from "std/log/mod.ts"; import { createEndpoint, createMethodFilter } from "t_rest/server"; import { configSchema, getConfig, setConfig } from "../app/config.ts"; import { bot } from "../bot/mod.ts"; import { sessions } from "./sessionsRoute.ts"; -export const logger = () => getLogger(); - export const paramsRoute = createMethodFilter({ GET: createEndpoint( { query: null, body: null }, @@ -38,7 +36,7 @@ export const paramsRoute = createMethodFilter({ if (!config?.adminUsernames?.includes(chat.username)) { return { status: 403, body: { type: "text/plain", data: "Must be an admin" } }; } - logger().info(`User ${chat.username} updated default params: ${JSON.stringify(body.data)}`); + info(`User ${chat.username} updated default params: ${JSON.stringify(body.data)}`); const defaultParams = deepMerge(config.defaultParams ?? {}, body.data); await setConfig({ defaultParams }); return { status: 200, body: { type: "application/json", data: config.defaultParams } }; diff --git a/api/workersRoute.ts b/api/workersRoute.ts index f6e4106..96bd146 100644 --- a/api/workersRoute.ts +++ b/api/workersRoute.ts @@ -1,18 +1,16 @@ -import { createEndpoint, createMethodFilter, createPathFilter } from "t_rest/server"; -import { activeGenerationWorkers } from "../app/generationQueue.ts"; -import { getConfig } from "../app/config.ts"; -import * as SdApi from "../app/sdApi.ts"; -import createOpenApiFetch from "openapi_fetch"; -import { sessions } from "./sessionsRoute.ts"; -import { bot } from "../bot/mod.ts"; -import { getLogger } from "std/log/mod.ts"; -import { WorkerInstance, workerInstanceStore } from "../app/workerInstanceStore.ts"; -import { getAuthHeader } from "../utils/getAuthHeader.ts"; -import { Model } from "indexed_kv"; -import { generationStore } from "../app/generationStore.ts"; import { subMinutes } from "date-fns"; - -const logger = () => getLogger(); +import { Model } from "indexed_kv"; +import createOpenApiFetch from "openapi_fetch"; +import { info } from "std/log/mod.ts"; +import { createEndpoint, createMethodFilter, createPathFilter } from "t_rest/server"; +import { getConfig } from "../app/config.ts"; +import { activeGenerationWorkers } from "../app/generationQueue.ts"; +import { generationStore } from "../app/generationStore.ts"; +import * as SdApi from "../app/sdApi.ts"; +import { WorkerInstance, workerInstanceStore } from "../app/workerInstanceStore.ts"; +import { bot } from "../bot/mod.ts"; +import { getAuthHeader } from "../utils/getAuthHeader.ts"; +import { sessions } from "./sessionsRoute.ts"; export type WorkerData = Omit & { id: string; @@ -124,7 +122,7 @@ export const workersRoute = createPathFilter({ sdUrl: body.data.sdUrl, sdAuth: body.data.sdAuth, }); - logger().info(`User ${chat.username} created worker ${workerInstance.id}`); + info(`User ${chat.username} created worker ${workerInstance.id}`); const worker = await getWorkerData(workerInstance); return { status: 200, @@ -202,7 +200,7 @@ export const workersRoute = createPathFilter({ if (body.data.auth !== undefined) { workerInstance.value.sdAuth = body.data.auth; } - logger().info( + info( `User ${chat.username} updated worker ${params.workerId}: ${JSON.stringify(body.data)}`, ); await workerInstance.update(); @@ -238,7 +236,7 @@ export const workersRoute = createPathFilter({ if (!config?.adminUsernames?.includes(chat.username)) { return { status: 403, body: { type: "text/plain", data: "Must be an admin" } }; } - logger().info(`User ${chat.username} deleted worker ${params.workerId}`); + info(`User ${chat.username} deleted worker ${params.workerId}`); await workerInstance.delete(); return { status: 200, body: { type: "application/json", data: null } }; }, diff --git a/app/dailyStatsStore.ts b/app/dailyStatsStore.ts index dca9313..6ae30c0 100644 --- a/app/dailyStatsStore.ts +++ b/app/dailyStatsStore.ts @@ -1,13 +1,11 @@ import { UTCDateMini } from "@date-fns/utc"; import { hoursToMilliseconds, isSameDay, minutesToMilliseconds } from "date-fns"; -import { getLogger } from "std/log/mod.ts"; +import { info } from "std/log/mod.ts"; import { JsonSchema, jsonType } from "t_rest/server"; import { db } from "./db.ts"; import { generationStore } from "./generationStore.ts"; import { kvMemoize } from "./kvMemoize.ts"; -const logger = () => getLogger(); - export const dailyStatsSchema = { type: "object", properties: { @@ -36,7 +34,7 @@ export const getDailyStats = kvMemoize( const after = new Date(Date.UTC(year, month - 1, day)); const before = new Date(Date.UTC(year, month - 1, day + 1)); - logger().info(`Calculating daily stats for ${year}-${month}-${day}`); + info(`Calculating daily stats for ${year}-${month}-${day}`); for await ( const generation of generationStore.listAll({ after, before }) diff --git a/app/generationQueue.ts b/app/generationQueue.ts index bed9ea9..7b8fc39 100644 --- a/app/generationQueue.ts +++ b/app/generationQueue.ts @@ -4,7 +4,7 @@ import { JobData, Queue, Worker } from "kvmq"; import createOpenApiClient from "openapi_fetch"; import { delay } from "std/async/delay.ts"; import { decode, encode } from "std/encoding/base64.ts"; -import { getLogger } from "std/log/mod.ts"; +import { debug, error, info } from "std/log/mod.ts"; import { ulid } from "ulid"; import { bot } from "../bot/mod.ts"; import { PngInfo } from "../bot/parsePngInfo.ts"; @@ -19,8 +19,6 @@ import * as SdApi from "./sdApi.ts"; import { uploadQueue } from "./uploadQueue.ts"; import { workerInstanceStore } from "./workerInstanceStore.ts"; -const logger = () => getLogger(); - interface GenerationJob { task: | { @@ -73,7 +71,7 @@ export async function processGenerationQueue() { .catch((error) => { workerInstance.update({ lastError: { message: error.message, time: Date.now() } }) .catch(() => undefined); - logger().debug(`Worker ${workerInstance.value.key} is down: ${error}`); + debug(`Worker ${workerInstance.value.key} is down: ${error}`); }); if (!activeWorkerStatus?.data) { @@ -86,7 +84,7 @@ export async function processGenerationQueue() { }); newWorker.addEventListener("error", (e) => { - logger().error( + error( `Generation failed for ${formatUserChat(e.detail.job.state)}: ${e.detail.error}`, ); bot.api.sendMessage( @@ -103,7 +101,7 @@ export async function processGenerationQueue() { newWorker.stopProcessing(); workerInstance.update({ lastError: { message: e.detail.error.message, time: Date.now() } }) .catch(() => undefined); - logger().info(`Stopped worker ${workerInstance.value.key}`); + info(`Stopped worker ${workerInstance.value.key}`); }); newWorker.addEventListener("complete", () => { @@ -113,7 +111,7 @@ export async function processGenerationQueue() { await workerInstance.update({ lastOnlineTime: Date.now() }); newWorker.processJobs(); activeGenerationWorkers.set(workerInstance.id, newWorker); - logger().info(`Started worker ${workerInstance.value.key}`); + info(`Started worker ${workerInstance.value.key}`); } await delay(60_000); } @@ -139,7 +137,7 @@ async function processGenerationJob( }); state.workerInstanceKey = workerInstance.value.key; state.progress = 0; - logger().debug(`Generation started for ${formatUserChat(state)}`); + debug(`Generation started for ${formatUserChat(state)}`); await updateJob({ state: state }); // check if bot can post messages in this chat @@ -253,7 +251,7 @@ async function processGenerationJob( ).catch(() => undefined); } - await Promise.race([delay(1000), responsePromise]).catch(() => undefined); + await Promise.race([delay(2_000), responsePromise]).catch(() => undefined); } while (await promiseState(responsePromise) === "pending"); // check response @@ -298,7 +296,7 @@ async function processGenerationJob( { maxAttempts: 1 }, ).catch(() => undefined); - logger().debug(`Generation finished for ${formatUserChat(state)}`); + debug(`Generation finished for ${formatUserChat(state)}`); } /** diff --git a/app/uploadQueue.ts b/app/uploadQueue.ts index 2c38343..1ee335d 100644 --- a/app/uploadQueue.ts +++ b/app/uploadQueue.ts @@ -4,15 +4,13 @@ import { bold, fmt } from "grammy_parse_mode"; import { Chat, Message, User } from "grammy_types"; import { Queue } from "kvmq"; import { format } from "std/fmt/duration.ts"; -import { getLogger } from "std/log/mod.ts"; +import { debug, error } from "std/log/mod.ts"; import { bot } from "../bot/mod.ts"; import { formatUserChat } from "../utils/formatUserChat.ts"; import { db, fs } from "./db.ts"; import { generationStore, SdGenerationInfo } from "./generationStore.ts"; import { globalStats } from "./globalStats.ts"; -const logger = () => getLogger(); - interface UploadJob { from: User; chat: Chat; @@ -125,7 +123,7 @@ export async function processUploadQueue() { globalStats.userIds = [...userIdSet]; } - logger().debug( + debug( `Uploaded ${state.imageKeys.length} ${[...types].join(",")} images (${ Math.trunc(size / 1024) }kB) for ${formatUserChat(state)}`, @@ -137,7 +135,7 @@ export async function processUploadQueue() { }, { concurrency: 3 }); 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( e.detail.job.state.requestMessage.chat.id, `Upload failed: ${e.detail.error}\n\n` + diff --git a/app/userStatsStore.ts b/app/userStatsStore.ts index e2f16f8..e842a8b 100644 --- a/app/userStatsStore.ts +++ b/app/userStatsStore.ts @@ -1,13 +1,11 @@ import { minutesToMilliseconds } from "date-fns"; import { Store } from "indexed_kv"; -import { getLogger } from "std/log/mod.ts"; +import { info } from "std/log/mod.ts"; import { JsonSchema, jsonType } from "t_rest/server"; import { db } from "./db.ts"; import { generationStore } from "./generationStore.ts"; import { kvMemoize } from "./kvMemoize.ts"; -const logger = () => getLogger(); - export const userStatsSchema = { type: "object", properties: { @@ -63,7 +61,7 @@ export const getUserStats = kvMemoize( let pixelStepCount = 0; const tagCountMap: Record = {}; - logger().info(`Calculating user stats for ${userId}`); + info(`Calculating user stats for ${userId}`); for await ( const generation of generationStore.listBy("fromId", { value: userId }) diff --git a/bot/broadcastCommand.ts b/bot/broadcastCommand.ts index 34a6318..5a72910 100644 --- a/bot/broadcastCommand.ts +++ b/bot/broadcastCommand.ts @@ -1,10 +1,11 @@ import { CommandContext } from "grammy"; import { bold, fmt, FormattedString } from "grammy_parse_mode"; import { distinctBy } from "std/collections/distinct_by.ts"; +import { error, info } from "std/log/mod.ts"; import { getConfig } from "../app/config.ts"; import { generationStore } from "../app/generationStore.ts"; import { formatUserChat } from "../utils/formatUserChat.ts"; -import { ErisContext, logger } from "./mod.ts"; +import { ErisContext } from "./mod.ts"; export async function broadcastCommand(ctx: CommandContext) { if (!ctx.from?.username) { @@ -46,10 +47,10 @@ export async function broadcastCommand(ctx: CommandContext) { for (const gen of gens) { try { 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++; } 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`); } const fmtMessage = getMessage(); diff --git a/bot/img2imgCommand.ts b/bot/img2imgCommand.ts index e837191..45ba4c6 100644 --- a/bot/img2imgCommand.ts +++ b/bot/img2imgCommand.ts @@ -1,10 +1,11 @@ import { CommandContext } from "grammy"; import { StatelessQuestion } from "grammy_stateless_question"; import { maxBy } from "std/collections/max_by.ts"; +import { debug } from "std/log/mod.ts"; import { getConfig } from "../app/config.ts"; import { generationQueue } from "../app/generationQueue.ts"; import { formatUserChat } from "../utils/formatUserChat.ts"; -import { ErisContext, logger } from "./mod.ts"; +import { ErisContext } from "./mod.ts"; import { parsePngInfo, PngInfo } from "./parsePngInfo.ts"; type QuestionState = { fileId?: string; params?: Partial }; @@ -126,5 +127,5 @@ async function img2img( replyMessage: replyMessage, }, { retryCount: 3, repeatDelayMs: 10_000 }); - logger().debug(`Generation (img2img) enqueued for ${formatUserChat(ctx.message)}`); + debug(`Generation (img2img) enqueued for ${formatUserChat(ctx.message)}`); } diff --git a/bot/mod.ts b/bot/mod.ts index 46e0220..0e1c380 100644 --- a/bot/mod.ts +++ b/bot/mod.ts @@ -2,7 +2,7 @@ import { Api, Bot, Context, RawApi, session, SessionFlavor } from "grammy"; import { FileFlavor, hydrateFiles } from "grammy_files"; import { hydrateReply, ParseModeFlavor } from "grammy_parse_mode"; import { run, sequentialize } from "grammy_runner"; -import { getLogger } from "std/log/mod.ts"; +import { error, info, warning } from "std/log/mod.ts"; import { sessions } from "../api/sessionsRoute.ts"; import { getConfig, setConfig } from "../app/config.ts"; import { formatUserChat } from "../utils/formatUserChat.ts"; @@ -13,8 +13,6 @@ import { pnginfoCommand, pnginfoQuestion } from "./pnginfoCommand.ts"; import { queueCommand } from "./queueCommand.ts"; import { txt2imgCommand, txt2imgQuestion } from "./txt2imgCommand.ts"; -export const logger = () => getLogger(); - interface SessionData { chat: ErisChatData; user: ErisUserData; @@ -77,7 +75,7 @@ bot.api.config.use(async (prev, method, payload, signal) => { if (attempt >= maxAttempts) return result; const retryAfter = result.parameters?.retry_after ?? (attempt * 5); if (retryAfter > maxWait) return result; - logger().warning( + warning( `${method} (attempt ${attempt}) failed: ${result.error_code} ${result.description}`, ); await new Promise((resolve) => setTimeout(resolve, retryAfter * 1000)); @@ -85,7 +83,7 @@ bot.api.config.use(async (prev, method, payload, signal) => { }); bot.catch((err) => { - logger().error( + error( `Handling update from ${formatUserChat(err.ctx)} failed: ${err.name} ${err.message}`, ); }); @@ -131,7 +129,7 @@ bot.command("start", async (ctx) => { } session.userId = ctx.from?.id; sessions.set(id, session); - logger().info(`User ${formatUserChat(ctx)} logged in`); + info(`User ${formatUserChat(ctx)} logged in`); // TODO: show link to web ui await ctx.reply("Login successful! You can now return to the WebUI.", { reply_to_message_id: ctx.message?.message_id, @@ -169,7 +167,7 @@ bot.command("pause", async (ctx) => { await setConfig({ pausedReason: ctx.match || "No reason given", }); - logger().warning(`Bot paused by ${ctx.from.first_name} because ${config.pausedReason}`); + warning(`Bot paused by ${ctx.from.first_name} because ${config.pausedReason}`); return ctx.reply("Paused"); }); @@ -179,7 +177,7 @@ bot.command("resume", async (ctx) => { if (!config.adminUsernames.includes(ctx.from.username)) return; if (config.pausedReason == null) return ctx.reply("Already running"); await setConfig({ pausedReason: null }); - logger().info(`Bot resumed by ${ctx.from.first_name}`); + info(`Bot resumed by ${ctx.from.first_name}`); return ctx.reply("Resumed"); }); diff --git a/bot/txt2imgCommand.ts b/bot/txt2imgCommand.ts index 0e087cb..6b2ae1a 100644 --- a/bot/txt2imgCommand.ts +++ b/bot/txt2imgCommand.ts @@ -1,9 +1,10 @@ import { CommandContext } from "grammy"; import { StatelessQuestion } from "grammy_stateless_question"; +import { debug } from "std/log/mod.ts"; import { getConfig } from "../app/config.ts"; import { generationQueue } from "../app/generationQueue.ts"; import { formatUserChat } from "../utils/formatUserChat.ts"; -import { ErisContext, logger } from "./mod.ts"; +import { ErisContext } from "./mod.ts"; import { getPngInfo, parsePngInfo, PngInfo } from "./parsePngInfo.ts"; export const txt2imgQuestion = new StatelessQuestion( @@ -91,5 +92,5 @@ async function txt2img(ctx: ErisContext, match: string, includeRepliedTo: boolea replyMessage: replyMessage, }, { retryCount: 3, retryDelayMs: 10_000 }); - logger().debug(`Generation (txt2img) enqueued for ${formatUserChat(ctx.message)}`); + debug(`Generation (txt2img) enqueued for ${formatUserChat(ctx.message)}`); } diff --git a/main.ts b/main.ts index 8a8449b..43a8d72 100644 --- a/main.ts +++ b/main.ts @@ -1,18 +1,20 @@ /// import "std/dotenv/load.ts"; import { ConsoleHandler } from "std/log/handlers.ts"; -import { setup } from "std/log/mod.ts"; +import { LevelName, setup } from "std/log/mod.ts"; import { serveUi } from "./api/mod.ts"; import { runAllTasks } from "./app/mod.ts"; import { runBot } from "./bot/mod.ts"; +const logLevel = Deno.env.get("LOG_LEVEL")?.toUpperCase() as LevelName ?? "INFO"; + // setup logging setup({ handlers: { - console: new ConsoleHandler("INFO"), + console: new ConsoleHandler(logLevel), }, loggers: { - default: { level: "INFO", handlers: ["console"] }, + default: { level: logLevel, handlers: ["console"] }, }, });