117 lines
5.1 KiB
TypeScript
117 lines
5.1 KiB
TypeScript
import {ChatCommand} from "../base/chat-command";
|
|
import {Message} from "typescript-telegram-bot-api";
|
|
import {Requirements} from "../base/requirements";
|
|
import {Requirement} from "../base/requirement";
|
|
import {bot, openAi, photoGenDir} from "../index";
|
|
import fs from "node:fs";
|
|
import path from "node:path";
|
|
import {oldEditMessageText, logError, replyToMessage} from "../util/utils";
|
|
import {Environment} from "../common/environment";
|
|
import {APIError} from "openai";
|
|
|
|
export class OpenAIGenImage extends ChatCommand {
|
|
command = ["openAiGenImage", "chatGPTGenImage", "imgen"];
|
|
|
|
title = "/openAIGenImage";
|
|
description = "Generate image from OpenAI";
|
|
|
|
argsMode = "required" as const;
|
|
requirements = Requirements.Build(Requirement.BOT_CREATOR);
|
|
|
|
async execute(msg: Message, match?: RegExpExecArray): Promise<void> {
|
|
const prompt = match?.[3]?.trim();
|
|
if (!prompt?.length) return;
|
|
|
|
let waitMessage: Message | null = null;
|
|
|
|
try {
|
|
const totalParts = 3;
|
|
const model = Environment.OPENAI_IMAGE_MODEL;
|
|
const fileFullName = `${msg.chat.id}_${msg.message_id}.png`;
|
|
const getFileLocation = (fn: string) => {
|
|
return path.join(photoGenDir, fn);
|
|
};
|
|
|
|
waitMessage = await replyToMessage({message: msg, text: "🌈 Генерирую изображение..."});
|
|
|
|
const stream = await openAi.images.generate({
|
|
model: model,
|
|
prompt: prompt,
|
|
n: 1,
|
|
size: "auto",
|
|
stream: true,
|
|
partial_images: totalParts,
|
|
moderation: "low",
|
|
output_format: "png",
|
|
});
|
|
|
|
const then = Date.now();
|
|
|
|
for await (const event of stream) {
|
|
switch (event.type) {
|
|
case "image_generation.partial_image": {
|
|
console.log(` Partial image ${event.partial_image_index + 1}/3 received`);
|
|
console.log(` Size: ${event.b64_json.length} characters (base64)`);
|
|
|
|
const fileName = `partial_${event.partial_image_index + 1}_${fileFullName}`;
|
|
const imageBuffer = Buffer.from(event.b64_json, "base64");
|
|
const fileLocation = getFileLocation(fileName);
|
|
fs.writeFileSync(fileLocation, imageBuffer);
|
|
console.log(` 💾 Saved to: ${path.resolve(fileLocation)}`);
|
|
|
|
await bot.editMessageMedia({
|
|
chat_id: msg.chat.id,
|
|
message_id: waitMessage.message_id,
|
|
media: {
|
|
type: "photo",
|
|
media: imageBuffer,
|
|
caption: `🌈 Генерирую изображение (${(event.partial_image_index + 1)}/${totalParts})...`
|
|
}
|
|
});
|
|
break;
|
|
}
|
|
case "image_generation.completed": {
|
|
console.log("\n✅ Final image completed!");
|
|
console.log(` Size: ${event.b64_json.length} characters (base64)`);
|
|
|
|
const imageBuffer = Buffer.from(event.b64_json, "base64");
|
|
const fileLocation = getFileLocation(fileFullName);
|
|
fs.writeFileSync(fileLocation, imageBuffer);
|
|
console.log(` Saved to: ${path.resolve(fileLocation)}`);
|
|
|
|
const diff = Date.now() - then;
|
|
await bot.editMessageMedia({
|
|
chat_id: msg.chat.id,
|
|
message_id: waitMessage.message_id,
|
|
media: {
|
|
type: "photo",
|
|
media: imageBuffer,
|
|
caption: `🌈 Изображение по запросу "${prompt}" сгенерировано моделью "${model}" размеров ${event.size} за ${diff}ms`
|
|
}
|
|
});
|
|
break;
|
|
}
|
|
default:
|
|
console.log(`❓ Unknown event: ${event}`);
|
|
}
|
|
}
|
|
} catch (e) {
|
|
logError(e);
|
|
|
|
if (e instanceof APIError && e.error.code === "moderation_blocked") {
|
|
const text = "❌ Мне запрещено такое генерировать 😠";
|
|
|
|
if (waitMessage) {
|
|
await oldEditMessageText(msg.chat.id, waitMessage.message_id, text).catch(logError);
|
|
} else {
|
|
await replyToMessage({message: msg, text: text}).catch(logError);
|
|
}
|
|
} else {
|
|
await replyToMessage({
|
|
message: waitMessage ? waitMessage : msg,
|
|
text: `Произошла ошибка: ${e}`
|
|
}).catch(logError);
|
|
}
|
|
}
|
|
}
|
|
} |