Initial commit
This commit is contained in:
commit
2a6488488e
|
@ -0,0 +1,3 @@
|
|||
NODE_ENV=development
|
||||
BOT_TOKEN=
|
||||
BOT_ADMIN_USER_ID=[]
|
|
@ -0,0 +1,2 @@
|
|||
node_modules
|
||||
.env
|
|
@ -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>;
|
File diff suppressed because it is too large
Load Diff
|
@ -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"
|
||||
}
|
||||
}
|
|
@ -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;
|
|
@ -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);
|
||||
}
|
|
@ -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);
|
|
@ -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;
|
||||
}
|
|
@ -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 };
|
|
@ -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;
|
||||
}
|
|
@ -0,0 +1,2 @@
|
|||
export * from './welcome.js';
|
||||
export * from './imagine/imagine.js';
|
|
@ -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 };
|
|
@ -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),
|
||||
});
|
||||
};
|
|
@ -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();
|
||||
};
|
||||
}
|
|
@ -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;
|
|
@ -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();
|
||||
};
|
||||
}
|
|
@ -0,0 +1 @@
|
|||
export * from './update-logger.js';
|
|
@ -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,
|
||||
});
|
||||
}
|
||||
};
|
||||
}
|
|
@ -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/**/*",
|
||||
]
|
||||
}
|
Loading…
Reference in New Issue