forked from pinks/eris
initial commit
This commit is contained in:
commit
124f711118
|
@ -0,0 +1,21 @@
|
|||
# Eris the Bot
|
||||
|
||||
Telegram bot for generating images from text.
|
||||
|
||||
## Requirements
|
||||
|
||||
- [Deno](https://deno.land/)
|
||||
- [Stable Diffusion WebUI](https://github.com/AUTOMATIC1111/stable-diffusion-webui/)
|
||||
|
||||
## Options
|
||||
|
||||
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))
|
||||
- `SD_API_URL` - URL to Stable Diffusion API (e.g. `http://127.0.0.1:7860/`)
|
||||
- `ADMIN_USERNAMES` - Comma separated list of usernames of users that can use admin commands (optional)
|
||||
|
||||
## Running
|
||||
|
||||
- Start stable diffusion webui: `cd sd-webui`, `./webui.sh --api`
|
||||
- Start bot: `deno task start`
|
|
@ -0,0 +1,6 @@
|
|||
{
|
||||
"tasks": {
|
||||
"dev": "deno run --watch --allow-env --allow-read --allow-net main.ts",
|
||||
"start": "deno run --allow-env --allow-read --allow-net main.ts"
|
||||
}
|
||||
}
|
|
@ -0,0 +1,387 @@
|
|||
import {
|
||||
Bot,
|
||||
Context,
|
||||
InputFile,
|
||||
InputMediaBuilder,
|
||||
} from "https://deno.land/x/grammy@v1.18.1/mod.ts";
|
||||
import { autoQuote } from "https://deno.land/x/grammy_autoquote@v1.1.2/mod.ts";
|
||||
import {
|
||||
fmt,
|
||||
hydrateReply,
|
||||
ParseModeFlavor,
|
||||
} from "https://deno.land/x/grammy_parse_mode@1.7.1/mod.ts";
|
||||
import "https://deno.land/x/dotenv@v3.2.2/load.ts";
|
||||
import {
|
||||
FormattedString,
|
||||
bold,
|
||||
} from "https://deno.land/x/grammy_parse_mode@1.7.1/format.ts";
|
||||
import { autoRetry } from "https://esm.sh/@grammyjs/auto-retry";
|
||||
import { MessageEntity } from "https://deno.land/x/grammy@v1.18.1/types.ts";
|
||||
|
||||
const maxUserJobs = 3;
|
||||
const maxJobs = 10;
|
||||
|
||||
let isRunning = true;
|
||||
|
||||
const sdApiUrl = Deno.env.get("SD_API_URL");
|
||||
if (!sdApiUrl) throw new Error("SD_API_URL not set");
|
||||
console.log("Using SD API URL:", sdApiUrl);
|
||||
const sdConfigUrl = new URL("/config", sdApiUrl);
|
||||
const sdConfigRequest = await fetch(sdConfigUrl);
|
||||
if (!sdConfigRequest.ok)
|
||||
throw new Error(
|
||||
`Failed to fetch SD config from ${sdConfigUrl}: ${sdConfigRequest.statusText}`
|
||||
);
|
||||
const sdConfig = await sdConfigRequest.json();
|
||||
console.log("Using SD WebUI version:", String(sdConfig.version).trim());
|
||||
|
||||
const adminUsernames = (Deno.env.get("ADMIN_USERNAMES") ?? "")
|
||||
.split(",")
|
||||
.filter(Boolean);
|
||||
|
||||
const tgBotToken = Deno.env.get("TG_BOT_TOKEN");
|
||||
if (!tgBotToken) throw new Error("TG_BOT_TOKEN not set");
|
||||
|
||||
const bot = new Bot<ParseModeFlavor<Context>>(tgBotToken);
|
||||
bot.api.config.use(autoRetry({ maxRetryAttempts: 5, maxDelaySeconds: 30 }));
|
||||
|
||||
bot.api.setMyShortDescription("I can generate furry images from text");
|
||||
bot.api.setMyDescription(
|
||||
"I can generate furry images from text. Send /txt2img to generate an image."
|
||||
);
|
||||
bot.api.setMyCommands([
|
||||
{ command: "txt2img", description: "Generate an image" },
|
||||
{ command: "queue", description: "Show the current queue" },
|
||||
]);
|
||||
|
||||
bot.use(autoQuote);
|
||||
bot.use(hydrateReply);
|
||||
|
||||
bot.command("start", (ctx) =>
|
||||
ctx.reply("Hello! Use the /txt2img command to generate an image")
|
||||
);
|
||||
|
||||
bot.command("txt2img", async (ctx) => {
|
||||
if (!ctx.from?.id) {
|
||||
return ctx.reply("I don't know who you are");
|
||||
}
|
||||
if (!isRunning) {
|
||||
return ctx.reply("I'm currently paused. Try again later.");
|
||||
}
|
||||
if (queue.length >= maxJobs) {
|
||||
return ctx.reply(
|
||||
`The queue is full. Try again later. (Max queue size: ${maxJobs})`
|
||||
);
|
||||
}
|
||||
const jobCount = queue.filter((job) => job.userId === ctx.from.id).length;
|
||||
if (jobCount >= maxUserJobs) {
|
||||
return ctx.reply(
|
||||
`You already have ${maxUserJobs} jobs in queue. Try again later.`
|
||||
);
|
||||
}
|
||||
if (!ctx.match) {
|
||||
return ctx.reply("Please describe what you want to see");
|
||||
}
|
||||
const place = queue.length + 1;
|
||||
const queueMessage = await ctx.reply(
|
||||
`You are ${formatOrdinal(place)} in queue.`
|
||||
);
|
||||
const userName = [ctx.from.first_name, ctx.from.last_name]
|
||||
.filter(Boolean)
|
||||
.join(" ");
|
||||
const chatName =
|
||||
ctx.chat.type === "supergroup" || ctx.chat.type === "group"
|
||||
? ctx.chat.title
|
||||
: "private chat";
|
||||
queue.push({
|
||||
params: { prompt: ctx.match },
|
||||
userId: ctx.from.id,
|
||||
userName,
|
||||
chatId: ctx.chat.id,
|
||||
chatName,
|
||||
requestMessageId: ctx.message.message_id,
|
||||
statusMessageId: queueMessage.message_id,
|
||||
});
|
||||
console.log(`Enqueued job for ${userName} in chat ${chatName}`);
|
||||
});
|
||||
|
||||
bot.command("queue", async (ctx) => {
|
||||
if (queue.length === 0) return ctx.reply("Queue is empty");
|
||||
return await ctx.replyFmt(
|
||||
fmt`Current queue:\n\n${fmtArray(
|
||||
queue.map(
|
||||
(job, index) =>
|
||||
fmt`${bold(index + 1)}. ${bold(job.userName)} in ${bold(
|
||||
job.chatName
|
||||
)}`
|
||||
),
|
||||
"\n"
|
||||
)}`
|
||||
);
|
||||
});
|
||||
|
||||
bot.command("pause", async (ctx) => {
|
||||
if (!ctx.from?.username) return;
|
||||
if (!adminUsernames.includes(ctx.from.username)) return;
|
||||
if (!isRunning) return await ctx.reply("Already paused");
|
||||
isRunning = false;
|
||||
return await ctx.reply("Paused");
|
||||
});
|
||||
|
||||
bot.command("resume", async (ctx) => {
|
||||
if (!ctx.from?.username) return;
|
||||
if (!adminUsernames.includes(ctx.from.username)) return;
|
||||
if (isRunning) return await ctx.reply("Already running");
|
||||
isRunning = true;
|
||||
return await ctx.reply("Resumed");
|
||||
});
|
||||
|
||||
bot.catch((err) => {
|
||||
let msg = "Error processing update";
|
||||
const { from, chat } = err.ctx;
|
||||
if (from?.first_name) msg += ` from ${from.first_name}`;
|
||||
if (from?.last_name) msg += ` ${from.last_name}`;
|
||||
if (from?.username) msg += ` (@${from.username})`;
|
||||
if (chat?.type === "supergroup" || chat?.type === "group") {
|
||||
msg += ` in ${chat.title}`;
|
||||
if (chat.type === "supergroup" && chat.username)
|
||||
msg += ` (@${chat.username})`;
|
||||
}
|
||||
console.error(msg, err.error);
|
||||
});
|
||||
|
||||
const queue: Job[] = [];
|
||||
|
||||
interface Job {
|
||||
params: Partial<SdRequest>;
|
||||
userId: number;
|
||||
userName: string;
|
||||
chatId: number;
|
||||
chatName: string;
|
||||
requestMessageId: number;
|
||||
statusMessageId: number;
|
||||
}
|
||||
|
||||
async function processQueue() {
|
||||
while (true) {
|
||||
const job = queue.shift();
|
||||
if (!job) {
|
||||
await new Promise((resolve) => setTimeout(resolve, 1000));
|
||||
continue;
|
||||
}
|
||||
for (const [index, job] of queue.entries()) {
|
||||
const place = index + 1;
|
||||
await bot.api
|
||||
.editMessageText(
|
||||
job.chatId,
|
||||
job.statusMessageId,
|
||||
`You are ${formatOrdinal(place)} in queue.`
|
||||
)
|
||||
.catch(() => {});
|
||||
}
|
||||
try {
|
||||
await bot.api.deleteMessage(job.chatId, job.statusMessageId);
|
||||
const progressMessage = await bot.api.sendMessage(
|
||||
job.chatId,
|
||||
"Generating your prompt now...",
|
||||
{ reply_to_message_id: job.requestMessageId }
|
||||
);
|
||||
const onProgress = (progress: SdProgressResponse) => {
|
||||
bot.api
|
||||
.editMessageText(
|
||||
job.chatId,
|
||||
progressMessage.message_id,
|
||||
`Generating your prompt now... ${Math.round(
|
||||
progress.progress * 100
|
||||
)}%`
|
||||
)
|
||||
.catch(() => {});
|
||||
};
|
||||
const response = await txt2img(
|
||||
{ ...defaultParams, ...job.params },
|
||||
onProgress
|
||||
);
|
||||
console.log(
|
||||
`Generated image for ${job.userName} in ${job.chatName}: ${job.params.prompt}`
|
||||
);
|
||||
bot.api.editMessageText(
|
||||
job.chatId,
|
||||
progressMessage.message_id,
|
||||
`Uploading your images...`
|
||||
);
|
||||
const inputFiles = await Promise.all(
|
||||
response.images.slice(1).map(async (imageBase64) => {
|
||||
const imageBlob = await fetch(
|
||||
`data:${mimeTypeFromBase64(imageBase64)};base64,${imageBase64}`
|
||||
).then((resp) => resp.blob());
|
||||
return InputMediaBuilder.photo(new InputFile(imageBlob));
|
||||
})
|
||||
);
|
||||
await bot.api.sendMediaGroup(job.chatId, inputFiles, {
|
||||
reply_to_message_id: job.requestMessageId,
|
||||
});
|
||||
await bot.api.deleteMessage(job.chatId, progressMessage.message_id);
|
||||
console.log(`${queue.length} jobs remaining`);
|
||||
} catch (err) {
|
||||
console.error(
|
||||
`Failed to generate image for ${job.userName} in ${job.chatName}: ${job.params.prompt} - ${err}`
|
||||
);
|
||||
await bot.api
|
||||
.sendMessage(job.chatId, err.toString(), {
|
||||
reply_to_message_id: job.requestMessageId,
|
||||
})
|
||||
.catch(() => {});
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
function formatOrdinal(n: number) {
|
||||
if (n % 100 === 11 || n % 100 === 12 || n % 100 === 13) return `${n}th`;
|
||||
if (n % 10 === 1) return `${n}st`;
|
||||
if (n % 10 === 2) return `${n}nd`;
|
||||
if (n % 10 === 3) return `${n}rd`;
|
||||
return `${n}th`;
|
||||
}
|
||||
|
||||
const defaultParams: Partial<SdRequest> = {
|
||||
batch_size: 3,
|
||||
n_iter: 1,
|
||||
width: 128 * 5,
|
||||
height: 128 * 7,
|
||||
steps: 40,
|
||||
cfg_scale: 9,
|
||||
send_images: true,
|
||||
save_images: true,
|
||||
negative_prompt:
|
||||
"id210 boring_e621_fluffyrock_v4 boring_e621_v4 easynegative ng_deepnegative_v1_75t",
|
||||
};
|
||||
|
||||
function mimeTypeFromBase64(base64: string) {
|
||||
if (base64.startsWith("/9j/")) {
|
||||
return "image/jpeg";
|
||||
}
|
||||
if (base64.startsWith("iVBORw0KGgo")) {
|
||||
return "image/png";
|
||||
}
|
||||
if (base64.startsWith("R0lGODlh")) {
|
||||
return "image/gif";
|
||||
}
|
||||
if (base64.startsWith("UklGRg")) {
|
||||
return "image/webp";
|
||||
}
|
||||
throw new Error("Unknown image type");
|
||||
}
|
||||
|
||||
async function txt2img(
|
||||
params: Partial<SdRequest>,
|
||||
onProgress?: (progress: SdProgressResponse) => void,
|
||||
signal?: AbortSignal
|
||||
): Promise<SdResponse> {
|
||||
let response: Response | undefined;
|
||||
let error: unknown;
|
||||
fetch(new URL("sdapi/v1/txt2img", sdApiUrl), {
|
||||
method: "POST",
|
||||
headers: { "Content-Type": "application/json" },
|
||||
body: JSON.stringify(params),
|
||||
}).then(
|
||||
(resp) => (response = resp),
|
||||
(err) => (error = err)
|
||||
);
|
||||
try {
|
||||
while (true) {
|
||||
await new Promise((resolve) => setTimeout(resolve, 3000));
|
||||
const progressRequest = await fetch(
|
||||
new URL("sdapi/v1/progress", sdApiUrl)
|
||||
);
|
||||
if (progressRequest.ok) {
|
||||
const progress = (await progressRequest.json()) as SdProgressResponse;
|
||||
onProgress?.(progress);
|
||||
}
|
||||
if (response != null) {
|
||||
if (response.ok) {
|
||||
const result = (await response.json()) as SdResponse;
|
||||
return result;
|
||||
} else {
|
||||
throw new Error(
|
||||
`Request failed: ${response.status} ${response.statusText}`
|
||||
);
|
||||
}
|
||||
}
|
||||
if (error != null) {
|
||||
throw error;
|
||||
}
|
||||
signal?.throwIfAborted();
|
||||
}
|
||||
} finally {
|
||||
if (!response && !error)
|
||||
await fetch(new URL("sdapi/v1/interrupt", sdApiUrl), { method: "POST" });
|
||||
}
|
||||
}
|
||||
|
||||
interface SdRequest {
|
||||
denoising_strength: number;
|
||||
prompt: string;
|
||||
seed: number;
|
||||
sampler_name: unknown;
|
||||
batch_size: number;
|
||||
n_iter: number;
|
||||
steps: number;
|
||||
cfg_scale: number;
|
||||
width: number;
|
||||
height: number;
|
||||
negative_prompt: string;
|
||||
send_images: boolean;
|
||||
save_images: boolean;
|
||||
}
|
||||
|
||||
interface SdResponse {
|
||||
images: string[];
|
||||
parameters: SdRequest;
|
||||
/** Contains serialized JSON */
|
||||
info: string;
|
||||
}
|
||||
|
||||
interface SdProgressResponse {
|
||||
progress: number;
|
||||
eta_relative: number;
|
||||
state: SdProgressState;
|
||||
/** base64 encoded preview */
|
||||
current_image: string | null;
|
||||
textinfo: string | null;
|
||||
}
|
||||
|
||||
interface SdProgressState {
|
||||
skipped: boolean;
|
||||
interrupted: boolean;
|
||||
job: string;
|
||||
job_count: number;
|
||||
job_timestamp: string;
|
||||
job_no: number;
|
||||
sampling_step: number;
|
||||
sampling_steps: number;
|
||||
}
|
||||
|
||||
/** Like {@link fmt} but accepts an array instead of template string. */
|
||||
function fmtArray(
|
||||
stringLikes: FormattedString[],
|
||||
separator = ""
|
||||
): FormattedString {
|
||||
let text = "";
|
||||
const entities: MessageEntity[] = [];
|
||||
for (let i = 0; i < stringLikes.length; i++) {
|
||||
const stringLike = stringLikes[i];
|
||||
|
||||
entities.push(
|
||||
...stringLike.entities.map((e) => ({
|
||||
...e,
|
||||
offset: e.offset + text.length,
|
||||
}))
|
||||
);
|
||||
|
||||
text += stringLike.toString();
|
||||
if (i < stringLikes.length - 1) text += separator;
|
||||
}
|
||||
return new FormattedString(text, entities);
|
||||
}
|
||||
|
||||
await Promise.all([bot.start(), processQueue()]);
|
Loading…
Reference in New Issue