Initial commit

This commit is contained in:
Timofey Gelazoniya 2023-09-08 22:09:35 +03:00
commit 2a6488488e
Signed by: zeldon
GPG Key ID: 047886915281DD2A
20 changed files with 2581 additions and 0 deletions

3
.env.example Normal file
View File

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

2
.gitignore vendored Normal file
View File

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

48
index.ts Normal file
View File

@ -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>;

2115
package-lock.json generated Normal file

File diff suppressed because it is too large Load Diff

45
package.json Normal file
View File

@ -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"
}
}

View File

@ -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;

27
scripts/start.ts Normal file
View File

@ -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);
}

47
src/config.ts Normal file
View File

@ -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);

36
src/context.ts Normal file
View File

@ -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;
}

View File

@ -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 };

View File

@ -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;
}

2
src/handlers/index.ts Normal file
View File

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

11
src/handlers/welcome.ts Normal file
View File

@ -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 };

12
src/helpers/error.ts Normal file
View File

@ -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),
});
};

19
src/helpers/logging.ts Normal file
View File

@ -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();
};
}

31
src/logger.ts Normal file
View File

@ -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;

View File

@ -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();
};
}

1
src/middlewares/index.ts Normal file
View File

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

View File

@ -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,
});
}
};
}

25
tsconfig.json Normal file
View File

@ -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/**/*",
]
}