Initial commit

.env.example

@@ -0,0 +1,3 @@
+NODE_ENV=development
+BOT_TOKEN=
+BOT_ADMIN_USER_ID=[]

.gitignore

@@ -0,0 +1,2 @@
+node_modules
+.env

index.ts

@@ -0,0 +1,48 @@
+import { Bot as TelegramBot, BotConfig } from "grammy";
+import { Context, createContextConstructor } from "#root/context.js";
+import { logger } from "#root/logger.js";
+import { hydrateReply, parseMode } from "@grammyjs/parse-mode";
+import { config } from "#root/config.js";
+import { autoChatAction } from "@grammyjs/auto-chat-action";
+import { hydrate } from "@grammyjs/hydrate";
+import { updateLogger } from "#root/middlewares/index.js";
+import { errorHandler } from "#root/helpers/error.js";
+import { welcomeFeature } from "#root/handlers/welcome.js";
+import { imageFeature } from "#root/handlers/imagine/imagine.js";
+import { ignoreOld } from "#root/middlewares/ignore-old.js";
+
+type Options = {
+    config?: Omit<BotConfig<Context>, "ContextConstructor">;
+};
+
+export function createBot(token: string, options: Options = {}) {
+    const bot = new TelegramBot(token, {
+        ...options.config,
+        ContextConstructor: createContextConstructor({ logger }),
+    });
+
+    bot.api.config.use(parseMode("html"));
+
+    bot.use(ignoreOld());
+
+    if (config.isDev) {
+        bot.use(updateLogger());
+    }
+
+    bot.use(autoChatAction(bot.api));
+    bot.use(hydrateReply);
+    bot.use(hydrate());
+
+    // Handlers
+    bot.use(welcomeFeature);
+    bot.use(imageFeature);
+
+    // Must be the last handler
+    if (config.isDev) {
+        bot.catch(errorHandler);
+    }
+
+    return bot;
+}
+
+export type Bot = ReturnType<typeof createBot>;

package.json

@@ -0,0 +1,45 @@
+{
+  "name": "sd-tg-bot",
+  "version": "0.1.0",
+  "description": "Stable Diffusion telegram bot",
+  "type": "module",
+  "imports": {
+    "#root/*": "./build/src/*"
+  },
+  "scripts": {
+    "postinstall": "patch-package",
+    "clean": "rimraf build",
+    "typecheck": "tsc",
+    "build": "npm run clean && tsc --noEmit false",
+    "dev": "npm run clean && tsc-watch --onSuccess \"tsx ./scripts/start.ts\"",
+    "start": "tsc && tsx ./scripts/start.ts",
+    "start:force": "tsx ./scripts/start.ts"
+  },
+  "keywords": [],
+  "author": "",
+  "license": "ISC",
+  "dependencies": {
+    "@grammyjs/auto-chat-action": "^0.1.1",
+    "@grammyjs/hydrate": "^1.3.1",
+    "@grammyjs/parse-mode": "^1.7.1",
+    "dotenv": "^16.3.1",
+    "grammy": "^1.18.1",
+    "node-graceful-shutdown": "^1.1.5",
+    "patch-package": "^8.0.0",
+    "pino": "^8.15.1",
+    "pino-pretty": "^10.2.0",
+    "sharp": "^0.32.5",
+    "stable-diffusion-api": "^0.0.7",
+    "zod": "^3.22.2"
+  },
+  "devDependencies": {
+    "@types/node": "^20.5.9",
+    "rimraf": "^5.0.1",
+    "tsc-watch": "^6.0.4",
+    "typescript": "^5.2.2"
+  },
+  "engines": {
+    "node": ">=18.0.0",
+    "npm": ">=8.0.0"
+  }
+}

patches/stable-diffusion-api+0.0.7.patch

@@ -0,0 +1,22 @@
+diff --git a/node_modules/stable-diffusion-api/dist/lib/StableDiffusionApi.js b/node_modules/stable-diffusion-api/dist/lib/StableDiffusionApi.js
+index efb62a9..6b03349 100644
+--- a/node_modules/stable-diffusion-api/dist/lib/StableDiffusionApi.js
++++ b/node_modules/stable-diffusion-api/dist/lib/StableDiffusionApi.js
+@@ -21,7 +21,7 @@ const base64_1 = require("../utils/base64");
+ const createScriptsWithCnUnits = (initScripts, controlNetUnit) => __awaiter(void 0, void 0, void 0, function* () {
+     const promises = controlNetUnit.map((unit) => __awaiter(void 0, void 0, void 0, function* () { return yield unit.toJson(); }));
+     const args = yield Promise.all(promises);
+-    const ControlNet = { args };
++    const ControlNet = args.length ? { args } : undefined;
+     const scripts = Object.assign(Object.assign({}, initScripts), { ControlNet });
+     return scripts;
+ });
+@@ -457,7 +457,7 @@ class StableDiffusionApi {
+     setModel(name, findClosest = true) {
+         return __awaiter(this, void 0, void 0, function* () {
+             const models = yield this.getSdModels();
+-            const modelNames = models.map((model) => model.name);
++            const modelNames = models.map((model) => model.model_name);
+             let foundModel = null;
+             if (modelNames.includes(name)) {
+                 foundModel = name;

scripts/start.ts

@@ -0,0 +1,27 @@
+#!/usr/bin/env tsx
+
+import { onShutdown } from "node-graceful-shutdown";
+import { createBot } from "../index.js";
+import { config } from "#root/config.js";
+import { logger } from "#root/logger.js";
+
+try {
+    const bot = createBot(config.BOT_TOKEN);
+
+    // Graceful shutdown
+    onShutdown(async () => {
+        logger.info("shutdown");
+        await bot.stop();
+    });
+
+    await bot.start({
+        onStart: ({ username }) =>
+            logger.info({
+                msg: "bot running...",
+                username,
+            }),
+    });
+} catch (error) {
+    logger.error(error);
+    process.exit(1);
+}

src/config.ts

@@ -0,0 +1,47 @@
+import "dotenv/config";
+import z, { ZodError, ZodIssueCode } from "zod";
+
+function parseJsonSafe(path: string) {
+    return (value: unknown) => {
+        try {
+            return JSON.parse(String(value));
+        } catch {
+            throw new ZodError([
+                {
+                    code: ZodIssueCode.custom,
+                    path: [path],
+                    fatal: true,
+                    message: "Invalid JSON",
+                },
+            ]);
+        }
+    };
+}
+
+const configSchema = z.object({
+    NODE_ENV: z.enum(["development", "production"]),
+    LOG_LEVEL: z
+        .enum(["trace", "debug", "info", "warn", "error", "fatal", "silent"])
+        .default("info"),
+    BOT_TOKEN: z.string(),
+    BOT_ADMIN_USER_ID: z
+        .preprocess(
+            parseJsonSafe("BOT_ADMIN_USER_ID"),
+            z.array(z.coerce.number().safe()).or(z.coerce.number().safe()),
+        )
+        .transform((v) => (Array.isArray(v) ? v : [v]))
+        .default([]),
+});
+
+const parseConfig = (environment: NodeJS.ProcessEnv) => {
+    const config = configSchema.parse(environment);
+
+    return {
+        ...config,
+        isDev: process.env.NODE_ENV === "development",
+        isProd: process.env.NODE_ENV === "production",
+    };
+};
+
+export type Config = ReturnType<typeof parseConfig>;
+export const config = parseConfig(process.env);

src/context.ts

@@ -0,0 +1,36 @@
+import { AutoChatActionFlavor } from "@grammyjs/auto-chat-action";
+import { HydrateFlavor } from "@grammyjs/hydrate";
+import { ParseModeFlavor } from "@grammyjs/parse-mode";
+import { Api, Context as DefaultContext } from "grammy";
+import { Update, UserFromGetMe } from "grammy/types";
+import { Logger } from "pino";
+
+type ExtendedContextFlavor = {
+    logger: Logger;
+};
+
+export type Context = ParseModeFlavor<
+    HydrateFlavor<
+        DefaultContext &
+        ExtendedContextFlavor &
+        AutoChatActionFlavor
+    >
+>;
+
+interface Dependencies {
+    logger: Logger;
+}
+
+export function createContextConstructor({ logger }: Dependencies) {
+    return class extends DefaultContext implements ExtendedContextFlavor {
+        logger: Logger;
+
+        constructor(update: Update, api: Api, me: UserFromGetMe) {
+            super(update, api, me);
+
+            this.logger = logger.child({
+                update_id: this.update.update_id,
+            });
+        }
+    } as unknown as new (update: Update, api: Api, me: UserFromGetMe) => Context;
+}

src/handlers/imagine/imagine.ts

@@ -0,0 +1,55 @@
+import { Context } from "#root/context.js";
+import { Composer, InputFile, InputMediaBuilder } from "grammy";
+import { predict, sdClient } from "./sd-client.js";
+
+const composer = new Composer<Context>();
+const feature = composer.chatType("private");
+
+type StatusMessage = Awaited<ReturnType<Context["replyWithPhoto"]>>;
+
+async function sendProgress(ctx: Context) {
+    let isStatusSent = false;
+    let statusMessage: StatusMessage;
+
+    const progressInterval = setInterval(async () => {
+        const response = await sdClient.getProgress();
+
+        if (response.progress === 0.0 && response.state.job_count === 0) {
+            await statusMessage.delete();
+            clearInterval(progressInterval);
+        }
+
+        let buffer;
+
+        try {
+            buffer = Buffer.from(response.current_image, 'base64');
+        } catch (err) {
+            return;
+        }
+
+        if (!isStatusSent) {
+            statusMessage = await ctx.replyWithPhoto(new InputFile(buffer), {
+                caption: "Processing your image, please wait"
+            });
+
+            isStatusSent = true;
+            return;
+        } else {
+            const newMedia = InputMediaBuilder.photo(new InputFile(buffer), { caption: "Processing your image, please wait" });
+            await statusMessage.editMedia(newMedia);
+            return;
+        }
+
+    }, 700);
+}
+
+feature.command("imagine", async (ctx) => {
+    const prompt = ctx.match;
+    console.log(prompt);
+    await sendProgress(ctx);
+    const prediction = await predict(`mj ${prompt}`);
+    const image = await prediction.image.jpeg().toBuffer();
+    await ctx.replyWithPhoto(new InputFile(image));
+});
+
+export { composer as imageFeature };

src/handlers/imagine/sd-client.ts

@@ -0,0 +1,31 @@
+import { StableDiffusionApi } from "stable-diffusion-api";
+
+export const sdClient = new StableDiffusionApi({
+    host: "127.0.0.1",
+    port: 7860,
+    protocol: "http",
+    defaultSampler: "DPM++ 2M Karras",
+    defaultStepCount: 22,
+});
+
+await sdClient.setModel("deliberate_v3");
+
+export async function predict(prompt: string) {
+    const result = await sdClient.txt2img({
+        prompt: prompt,
+        negative_prompt: '[deformed | disfigured], poorly drawn, [bad : wrong] anatomy, [extra | missing | floating | disconnected] limb, (mutated hands and fingers), blurry',
+        batch_size: 1,
+        cfg_scale: 7,
+        width: 640,
+        height: 640,
+        enable_hr: false,
+        hr_resize_x: 1280,
+        hr_resize_y: 1280,
+        hr_upscaler: "4x_NMKD-Siax_200k",
+        hr_second_pass_steps: 8,
+        denoising_strength: 0.36,
+        seed: -1
+    });
+
+    return result;
+}

src/handlers/index.ts

@@ -0,0 +1,2 @@
+export * from './welcome.js';
+export * from './imagine/imagine.js';

src/handlers/welcome.ts

@@ -0,0 +1,11 @@
+import { Context } from "#root/context.js";
+import { Composer } from "grammy";
+
+const composer = new Composer<Context>();
+const feature = composer.chatType("private");
+
+feature.command("start", (ctx) => {
+    return ctx.reply("Welcome user!");
+});
+
+export { composer as welcomeFeature };

src/helpers/error.ts

@@ -0,0 +1,12 @@
+import { ErrorHandler } from "grammy";
+import type { Context } from "#root/context.js";
+import { getUpdateInfo } from "./logging.js";
+
+export const errorHandler: ErrorHandler<Context> = (error) => {
+    const { ctx } = error;
+
+    ctx.logger.error({
+        err: error.error,
+        update: getUpdateInfo(ctx),
+    });
+};

src/helpers/logging.ts

@@ -0,0 +1,19 @@
+import { Middleware } from "grammy";
+import type { Update } from "@grammyjs/types";
+import type { Context } from "#root/context.js";
+
+export function getUpdateInfo(ctx: Context): Omit<Update, "update_id"> {
+    const { update_id, ...update } = ctx.update;
+    return update;
+}
+
+export function logHandle(id: string): Middleware<Context> {
+    return (ctx, next) => {
+        ctx.logger.info({
+            msg: `handle ${id}`,
+            ...(id.startsWith("unhandled") ? { update: getUpdateInfo(ctx) } : {}),
+        });
+
+        return next();
+    };
+}

src/logger.ts

@@ -0,0 +1,31 @@
+import { pino } from "pino";
+import { config } from "#root/config.js";
+
+export const logger = pino({
+    level: config.LOG_LEVEL,
+    transport: {
+        targets: [
+            ...(config.isDev
+                ? [
+                    {
+                        target: "pino-pretty",
+                        level: config.LOG_LEVEL,
+                        options: {
+                            ignore: "pid,hostname",
+                            colorize: true,
+                            translateTime: 'SYS:dd.mm.yyyy, HH:MM:ss'
+                        },
+                    },
+                ]
+                : [
+                    {
+                        target: "pino/file",
+                        level: config.LOG_LEVEL,
+                        options: {},
+                    },
+                ]),
+        ],
+    },
+});
+
+export type Logger = typeof logger;

src/middlewares/ignore-old.ts

@@ -0,0 +1,15 @@
+import { Context } from "#root/context.js";
+import { logger } from "#root/logger.js";
+import { Middleware } from "grammy";
+
+export function ignoreOld(threshold = 5 * 60): Middleware<Context> {
+    return async (ctx, next) => {
+        if (ctx.message?.date && new Date().getTime() / 1000 - ctx.message.date > threshold) {
+            logger.warn(`Ignoring message from user ${ctx.from?.id} at chat ${ctx.chat?.id} (${new Date().getTime() / 1000
+                }:${ctx.message.date})`);
+            return;
+        }
+
+        return next();
+    };
+}

src/middlewares/index.ts

@@ -0,0 +1 @@
+export * from './update-logger.js';

src/middlewares/update-logger.ts

@@ -0,0 +1,34 @@
+import { performance } from "node:perf_hooks";
+import { Middleware } from "grammy";
+import type { Context } from "#root/context.js";
+import { getUpdateInfo } from "#root/helpers/logging.js";
+
+export function updateLogger(): Middleware<Context> {
+    return async (ctx, next) => {
+        ctx.api.config.use((previous, method, payload, signal) => {
+            ctx.logger.debug({
+                msg: "bot api call",
+                method,
+                payload,
+            });
+
+            return previous(method, payload, signal);
+        });
+
+        ctx.logger.debug({
+            msg: "update received",
+            update: getUpdateInfo(ctx),
+        });
+
+        const startTime = performance.now();
+        try {
+            await next();
+        } finally {
+            const endTime = performance.now();
+            ctx.logger.debug({
+                msg: "update processed",
+                duration: endTime - startTime,
+            });
+        }
+    };
+}

tsconfig.json

@@ -0,0 +1,25 @@
+{
+  "compilerOptions": {
+    "strict": true,
+    "skipLibCheck": true,
+    "esModuleInterop": true,
+    "preserveWatchOutput": true,
+    "noEmit": true,
+    "module": "NodeNext",
+    "target": "ES2021",
+    "moduleResolution": "NodeNext",
+    "sourceMap": true,
+    "outDir": "build",
+    "rootDir": ".",
+    "paths": {
+      "#root/*": [
+        "./src/*"
+      ]
+    }
+  },
+  "include": [
+    "index.ts",
+    "src/**/*",
+    "scripts/**/*",
+  ]
+}