chore: split into modules
This commit is contained in:
parent
2030b382f4
commit
2f5058ac45
|
|
@ -13,7 +13,8 @@ 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))
|
- `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/`)
|
- `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)
|
- `ADMIN_USERNAMES` - Comma separated list of usernames of users that can use admin commands
|
||||||
|
(optional)
|
||||||
|
|
||||||
## Running
|
## Running
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -0,0 +1,113 @@
|
||||||
|
import {
|
||||||
|
autoQuote,
|
||||||
|
autoRetry,
|
||||||
|
bold,
|
||||||
|
Bot,
|
||||||
|
Context,
|
||||||
|
fmt,
|
||||||
|
hydrateReply,
|
||||||
|
ParseModeFlavor,
|
||||||
|
} from "./deps.ts";
|
||||||
|
import { fmtArray, formatOrdinal } from "./intl.ts";
|
||||||
|
import { config } from "./config.ts";
|
||||||
|
import { queue } from "./queue.ts";
|
||||||
|
|
||||||
|
export const bot = new Bot<ParseModeFlavor<Context>>(Deno.env.get("TG_BOT_TOKEN") ?? "");
|
||||||
|
bot.use(autoQuote);
|
||||||
|
bot.use(hydrateReply);
|
||||||
|
bot.api.config.use(autoRetry({ maxRetryAttempts: 5, maxDelaySeconds: 60 }));
|
||||||
|
|
||||||
|
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.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 (config.pausedReason != null) {
|
||||||
|
return ctx.reply(`I'm paused: ${config.pausedReason}`);
|
||||||
|
}
|
||||||
|
if (queue.length >= config.maxJobs) {
|
||||||
|
return ctx.reply(
|
||||||
|
`The queue is full. Try again later. (Max queue size: ${config.maxJobs})`,
|
||||||
|
);
|
||||||
|
}
|
||||||
|
const jobCount = queue.filter((job) => job.userId === ctx.from.id).length;
|
||||||
|
if (jobCount >= config.maxUserJobs) {
|
||||||
|
return ctx.reply(
|
||||||
|
`You already have ${config.maxUserJobs} jobs in queue. Try again later.`,
|
||||||
|
);
|
||||||
|
}
|
||||||
|
if (!ctx.match) {
|
||||||
|
return ctx.reply("Please describe what you want to see after the command");
|
||||||
|
}
|
||||||
|
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 (!config.adminUsernames.includes(ctx.from.username)) return;
|
||||||
|
if (config.pausedReason != null) {
|
||||||
|
return await ctx.reply(`Already paused: ${config.pausedReason}`);
|
||||||
|
}
|
||||||
|
config.pausedReason = ctx.match ?? "No reason given";
|
||||||
|
return await ctx.reply("Paused");
|
||||||
|
});
|
||||||
|
|
||||||
|
bot.command("resume", async (ctx) => {
|
||||||
|
if (!ctx.from?.username) return;
|
||||||
|
if (!config.adminUsernames.includes(ctx.from.username)) return;
|
||||||
|
if (config.pausedReason == null) return await ctx.reply("Already running");
|
||||||
|
config.pausedReason = null;
|
||||||
|
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);
|
||||||
|
});
|
||||||
|
|
@ -0,0 +1,15 @@
|
||||||
|
export const config: Config = {
|
||||||
|
adminUsernames: (Deno.env.get("ADMIN_USERNAMES") ?? "").split(",").filter(Boolean),
|
||||||
|
pausedReason: null,
|
||||||
|
sdApiUrl: Deno.env.get("SD_API_URL") ?? "http://127.0.0.1:7860/",
|
||||||
|
maxUserJobs: 3,
|
||||||
|
maxJobs: 20,
|
||||||
|
};
|
||||||
|
|
||||||
|
interface Config {
|
||||||
|
adminUsernames: string[];
|
||||||
|
pausedReason: string | null;
|
||||||
|
sdApiUrl: string;
|
||||||
|
maxUserJobs: number;
|
||||||
|
maxJobs: number;
|
||||||
|
}
|
||||||
|
|
@ -2,5 +2,8 @@
|
||||||
"tasks": {
|
"tasks": {
|
||||||
"dev": "deno run --watch --allow-env --allow-read --allow-net main.ts",
|
"dev": "deno run --watch --allow-env --allow-read --allow-net main.ts",
|
||||||
"start": "deno run --allow-env --allow-read --allow-net main.ts"
|
"start": "deno run --allow-env --allow-read --allow-net main.ts"
|
||||||
|
},
|
||||||
|
"fmt": {
|
||||||
|
"lineWidth": 100
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
|
||||||
|
|
@ -0,0 +1,4 @@
|
||||||
|
export * from "https://deno.land/x/grammy@v1.18.1/mod.ts";
|
||||||
|
export * from "https://deno.land/x/grammy_autoquote@v1.1.2/mod.ts";
|
||||||
|
export * from "https://deno.land/x/grammy_parse_mode@1.7.1/mod.ts";
|
||||||
|
export { autoRetry } from "https://esm.sh/@grammyjs/auto-retry@1.1.1";
|
||||||
|
|
@ -0,0 +1,33 @@
|
||||||
|
import { FormattedString } from "./deps.ts";
|
||||||
|
|
||||||
|
export 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`;
|
||||||
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Like `fmt` from `grammy_parse_mode` but accepts an array instead of template string.
|
||||||
|
* @see https://deno.land/x/grammy_parse_mode@1.7.1/format.ts?source=#L182
|
||||||
|
*/
|
||||||
|
export function fmtArray(
|
||||||
|
stringLikes: FormattedString[],
|
||||||
|
separator = "",
|
||||||
|
): FormattedString {
|
||||||
|
let text = "";
|
||||||
|
const entities: ConstructorParameters<typeof FormattedString>[1] = [];
|
||||||
|
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);
|
||||||
|
}
|
||||||
392
main.ts
392
main.ts
|
|
@ -1,388 +1,8 @@
|
||||||
import {
|
import "https://deno.land/std@0.201.0/dotenv/load.ts";
|
||||||
Bot,
|
import { bot } from "./bot.ts";
|
||||||
Context,
|
import { processQueue } from "./queue.ts";
|
||||||
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;
|
await Promise.all([
|
||||||
const maxJobs = 10;
|
bot.start(),
|
||||||
|
processQueue(),
|
||||||
let pausedReason: string | null = null;
|
|
||||||
|
|
||||||
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 (pausedReason != null) {
|
|
||||||
return ctx.reply(`I'm paused: ${pausedReason}`);
|
|
||||||
}
|
|
||||||
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 (pausedReason != null)
|
|
||||||
return await ctx.reply(`Already paused: ${pausedReason}`);
|
|
||||||
pausedReason = ctx.match ?? "No reason given";
|
|
||||||
return await ctx.reply("Paused");
|
|
||||||
});
|
|
||||||
|
|
||||||
bot.command("resume", async (ctx) => {
|
|
||||||
if (!ctx.from?.username) return;
|
|
||||||
if (!adminUsernames.includes(ctx.from.username)) return;
|
|
||||||
if (pausedReason == null) return await ctx.reply("Already running");
|
|
||||||
pausedReason = null;
|
|
||||||
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()]);
|
|
||||||
|
|
|
||||||
|
|
@ -0,0 +1,15 @@
|
||||||
|
export 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");
|
||||||
|
}
|
||||||
|
|
||||||
|
export function extFromMimeType(mimeType: string) {
|
||||||
|
if (mimeType === "image/jpeg") return "jpg";
|
||||||
|
if (mimeType === "image/png") return "png";
|
||||||
|
if (mimeType === "image/gif") return "gif";
|
||||||
|
if (mimeType === "image/webp") return "webp";
|
||||||
|
throw new Error("Unknown image type");
|
||||||
|
}
|
||||||
|
|
@ -0,0 +1,123 @@
|
||||||
|
import { InputFile, InputMediaBuilder } from "./deps.ts";
|
||||||
|
import { config } from "./config.ts";
|
||||||
|
import { bot } from "./bot.ts";
|
||||||
|
import { formatOrdinal } from "./intl.ts";
|
||||||
|
import { SdProgressResponse, SdRequest, txt2img } from "./sd.ts";
|
||||||
|
import { extFromMimeType, mimeTypeFromBase64 } from "./mimeType.ts";
|
||||||
|
|
||||||
|
export const queue: Job[] = [];
|
||||||
|
|
||||||
|
interface Job {
|
||||||
|
params: Partial<SdRequest>;
|
||||||
|
userId: number;
|
||||||
|
userName: string;
|
||||||
|
chatId: number;
|
||||||
|
chatName: string;
|
||||||
|
requestMessageId: number;
|
||||||
|
statusMessageId: number;
|
||||||
|
}
|
||||||
|
|
||||||
|
export 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)
|
||||||
|
.catch(() => {});
|
||||||
|
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(
|
||||||
|
config.sdApiUrl,
|
||||||
|
{ ...defaultParams, ...job.params },
|
||||||
|
onProgress,
|
||||||
|
);
|
||||||
|
|
||||||
|
console.log(
|
||||||
|
`Generated ${response.images.length} images (${
|
||||||
|
response.images
|
||||||
|
.map((image) => (image.length / 1024).toFixed(0) + "kB")
|
||||||
|
.join(", ")
|
||||||
|
}) for ${job.userName} in ${job.chatName}: ${job.params.prompt?.replace(/\s+/g, " ")}`,
|
||||||
|
);
|
||||||
|
bot.api.editMessageText(
|
||||||
|
job.chatId,
|
||||||
|
progressMessage.message_id,
|
||||||
|
`Uploading your images...`,
|
||||||
|
);
|
||||||
|
const inputFiles = await Promise.all(
|
||||||
|
response.images.map(async (imageBase64, idx) => {
|
||||||
|
const mimeType = mimeTypeFromBase64(imageBase64);
|
||||||
|
const imageBlob = await fetch(`data:${mimeType};base64,${imageBase64}`).then((resp) =>
|
||||||
|
resp.blob()
|
||||||
|
);
|
||||||
|
console.log(
|
||||||
|
`Uploading image ${idx + 1} of ${response.images.length} (${
|
||||||
|
(imageBlob.size / 1024).toFixed(0)
|
||||||
|
}kB)`,
|
||||||
|
);
|
||||||
|
return InputMediaBuilder.photo(
|
||||||
|
new InputFile(imageBlob, `${idx}.${extFromMimeType(mimeType)}`),
|
||||||
|
);
|
||||||
|
}),
|
||||||
|
);
|
||||||
|
await bot.api.sendMediaGroup(job.chatId, inputFiles, {
|
||||||
|
reply_to_message_id: job.requestMessageId,
|
||||||
|
});
|
||||||
|
await bot.api
|
||||||
|
.deleteMessage(job.chatId, progressMessage.message_id)
|
||||||
|
.catch(() => {});
|
||||||
|
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(() => bot.api.sendMessage(job.chatId, err.toString()))
|
||||||
|
.catch(() => {});
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
const defaultParams: Partial<SdRequest> = {
|
||||||
|
batch_size: 1,
|
||||||
|
n_iter: 1,
|
||||||
|
width: 128 * 2,
|
||||||
|
height: 128 * 3,
|
||||||
|
steps: 20,
|
||||||
|
cfg_scale: 9,
|
||||||
|
send_images: true,
|
||||||
|
negative_prompt: "boring_e621_fluffyrock_v4 boring_e621_v4",
|
||||||
|
};
|
||||||
|
|
@ -0,0 +1,88 @@
|
||||||
|
export async function txt2img(
|
||||||
|
apiUrl: string,
|
||||||
|
params: Partial<SdRequest>,
|
||||||
|
onProgress?: (progress: SdProgressResponse) => void,
|
||||||
|
signal?: AbortSignal,
|
||||||
|
): Promise<SdResponse> {
|
||||||
|
let response: Response | undefined;
|
||||||
|
let error: unknown;
|
||||||
|
|
||||||
|
fetch(new URL("sdapi/v1/txt2img", apiUrl), {
|
||||||
|
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", apiUrl));
|
||||||
|
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", apiUrl), { method: "POST" });
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
export 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;
|
||||||
|
}
|
||||||
|
|
||||||
|
export interface SdResponse {
|
||||||
|
images: string[];
|
||||||
|
parameters: SdRequest;
|
||||||
|
/** Contains serialized JSON */
|
||||||
|
info: string;
|
||||||
|
}
|
||||||
|
|
||||||
|
export interface SdProgressResponse {
|
||||||
|
progress: number;
|
||||||
|
eta_relative: number;
|
||||||
|
state: SdProgressState;
|
||||||
|
/** base64 encoded preview */
|
||||||
|
current_image: string | null;
|
||||||
|
textinfo: string | null;
|
||||||
|
}
|
||||||
|
|
||||||
|
export interface SdProgressState {
|
||||||
|
skipped: boolean;
|
||||||
|
interrupted: boolean;
|
||||||
|
job: string;
|
||||||
|
job_count: number;
|
||||||
|
job_timestamp: string;
|
||||||
|
job_no: number;
|
||||||
|
sampling_step: number;
|
||||||
|
sampling_steps: number;
|
||||||
|
}
|
||||||
Loading…
Reference in New Issue