Image gen

.env.example

@@ -1,2 +1,10 @@
 # Telegram bot token
 TELEGRAM_BOT_AUTH_TOKEN=1234abcd
+OPENAI_API_KEY=5678efgh
+
+TELEGRAM_FILE_URL=https://api.telegram.org/file/bot
+
+#for OpenAI Image Generation
+OPENAI_MODEL=text-davinci-003
+OPENAI_MAX_TOKENS=140
+OPENAI_TEMPERATURE=0.8

package.json

@@ -35,10 +35,13 @@
     "typescript": "^4.8.4"
   },
   "dependencies": {
-    "axios": "^1.4.0",
+    "@grammyjs/menu": "^1.2.1",
     "dotenv": "^16.0.3",
     "express": "^4.18.2",
     "grammy": "^1.17.1",
+    "openai": "^3.3.0",
+    "sharp": "^0.32.1",
+    "axios": "^1.4.0",
     "jsqr": "^1.4.0",
     "qrcode": "^1.5.3",
     "retry": "^0.13.1",

src/bot.ts

@@ -1,20 +1,32 @@
 import express from "express";
-import {Bot, MemorySessionStorage, session} from "grammy";
-import config from './config'
-import {VoiceMemo} from "./modules/voice-memo";
+import {
+  Bot,
+  session,
+  MemorySessionStorage,
+} from "grammy";
+import config from "./config";
 import {BotContext, BotSessionData, OnMessageContext} from "./modules/types";
-import {QRCodeBot} from "./modules/qrcode/QRCodeBot";
+import { mainMenu } from './pages'
 
-const bot = new Bot<BotContext>(config.telegramBotAuthToken);
+import { VoiceMemo } from "./modules/voice-memo";
+import { QRCodeBot } from "./modules/qrcode/QRCodeBot";
+import { imageGen } from "./modules/image-gen/ImageGenBot";
+
+export const bot = new Bot<BotContext>(config.telegramBotAuthToken);
 
 function createInitialSessionData(): BotSessionData {
-  return { qrMargin: 1 };
+  return {
+    imageGen: {
+      numImages: config.imageGen.sessionDefault.numImages,
+      imgSize: config.imageGen.sessionDefault.imgSize,
+    },
+    qrMargin: 1
+  };
 }
 
-bot.use(session({
-  initial: createInitialSessionData,
-  storage: new MemorySessionStorage()
-}));
+bot.use(session({ initial: createInitialSessionData, storage: new MemorySessionStorage() }));
+
+bot.use(mainMenu);
 
 const voiceMemo = new VoiceMemo();
 const qrCodeBot = new QRCodeBot();
@@ -23,19 +35,33 @@ const onMessage = async (ctx: OnMessageContext) => {
   if (qrCodeBot.isSupportedEvent(ctx)) {
     qrCodeBot.onEvent(ctx);
   }
-
   if(voiceMemo.isSupportedEvent(ctx)) {
     voiceMemo.onEvent(ctx)
   }
 }
 
 bot.command("start", (ctx) => ctx.reply("Welcome! Up and running."));
+
+bot.command("help", async (ctx) => {
+  console.log("help command", ctx.session);
+  await ctx.reply('Help text...', {
+    parse_mode: "HTML",
+    reply_markup: mainMenu,
+  });
+});
+
+bot.use(imageGen)
+
 bot.on("message", onMessage);
 
 const app = express();
+
 app.use(express.json());
 
 app.listen(config.port, () => {
   console.log(`Bot listening on port ${config.port}`);
-  bot.start();
+  bot.start()
+  // bot.start({
+  //   allowed_updates: ["callback_query"], // Needs to be set for menu middleware, but bot doesn't work with current configuration.   
+  // });
 });

src/config.ts

@@ -1,11 +1,25 @@
-import * as dotenv from 'dotenv'
-dotenv.config()
+import * as dotenv from "dotenv";
+dotenv.config();
 
 export default {
-  port: +(process.env.PORT || '3000'),
-  telegramBotAuthToken: process.env.TELEGRAM_BOT_AUTH_TOKEN || '',
-  stableDiffusionHost: process.env.SD_HOST || '',
+  port: +(process.env.PORT || "3000"),
+  telegramBotAuthToken: process.env.TELEGRAM_BOT_AUTH_TOKEN || "",
+  openAiKey: process.env.OPENAI_API_KEY,
+  stableDiffusionHost: process.env.SD_HOST || "",
   qr: {
     checkReadable: false,
+  },
+  imageGen: {
+    telegramFileUrl:
+      process.env.TELEGRAM_FILE_URL || "https://api.telegram.org/file/bot",
+    completions: {
+      model: process.env.OPENAI_MODEL || "text-davinci-003",
+      maxTokens: process.env.OPENAI_MAX_TOKENS || 140,
+      temperature: process.env.OPENAI_TEMPERATURE || 0.8,
+    },
+    sessionDefault: {
+      numImages: 1,
+      imgSize: "1024x1024",
+    },
   }
-}
+};

src/modules/image-gen/ImageGenBot.ts

@@ -0,0 +1,71 @@
+import { Composer } from "grammy";
+import config from "../../config";
+import { imgGen, imgGenEnhanced, alterImg } from "./controller";
+import { BotContext } from "../types";
+
+interface Image {
+  url: string;
+}
+
+export const imageGen = new Composer<BotContext>();
+
+imageGen.command("gen", async (ctx) => {
+  console.log("gen command");
+  const prompt = ctx.match;
+  if (!prompt) {
+    ctx.reply("Error: Missing prompt");
+    return;
+  }
+  const payload = {
+    chatId: ctx.chat.id,
+    prompt: ctx.match,
+    numImages: await ctx.session.imageGen.numImages, // lazy load
+    imgSize: await ctx.session.imageGen.imgSize, // lazy load
+  };
+  await imgGen(payload);
+});
+
+imageGen.command("genEn", async (ctx) => {
+  console.log("genEn command");
+  const prompt = ctx.match;
+  if (!prompt) {
+    ctx.reply("Error: Missing prompt");
+    return;
+  }
+  const payload = {
+    chatId: ctx.chat.id,
+    prompt: ctx.match,
+    numImages: await ctx.session.imageGen.numImages,
+    imgSize: await ctx.session.imageGen.imgSize,
+  };
+  ctx.reply("generating improved prompt...");
+  await imgGenEnhanced(payload);
+});
+
+imageGen.on("message", async (ctx) => {
+  try {
+    const photo = ctx.message.photo || ctx.message.reply_to_message?.photo;
+    if (photo) {
+      console.log("Alter img command");
+      const prompt = ctx.message.caption || ctx.message.text;
+      if (prompt) {
+        const file_id = photo.pop()?.file_id; // with pop() get full image quality
+        const file = await ctx.api.getFile(file_id!);
+        const filePath = `${config.imageGen.telegramFileUrl}${config.telegramBotAuthToken}/${file.file_path}`;
+        const payload = {
+          chatId: ctx.chat.id,
+          prompt: prompt,
+          numImages: await ctx.session.imageGen.numImages,
+          imgSize: await ctx.session.imageGen.imgSize,
+          filePath: filePath,
+        };
+        await alterImg(payload);
+      } else {
+        ctx.reply("Please add edit prompt");
+      }
+    }
+  } catch (e: any) {
+    console.log(e);
+    ctx.reply("An error occurred while generating the AI edit");
+  }
+});

src/modules/image-gen/api/openAi.ts

@@ -0,0 +1,95 @@
+import { Configuration, 
+  OpenAIApi, 
+  CreateImageRequest, 
+  CreateCompletionRequest } from "openai";
+
+import config from "../../../config";
+import { deleteFile, getImage } from "../utils/file";
+import { bot } from "../../../bot";
+
+const apiKey = config.openAiKey;
+
+const configuration = new Configuration({
+  apiKey: apiKey,
+});
+
+const openai = new OpenAIApi(configuration);
+
+export async function postGenerateImg(
+  prompt: string,
+  numImgs?: number,
+  imgSize?: string
+) {
+  try {
+    const payload = {
+      prompt: prompt,
+      n: numImgs ? numImgs : config.imageGen.sessionDefault.numImages,
+      size: imgSize ? imgSize : config.imageGen.sessionDefault.imgSize,
+    };
+    const response = await openai.createImage(payload as CreateImageRequest)
+    return response.data.data;
+  } catch (error) {
+    throw error;
+  }
+}
+
+export async function alterGeneratedImg(
+  chatId: number,
+  prompt: string,
+  filePath: string,
+  numImages?: number,
+  imgSize?: string
+) {
+  try {
+    const imageData = await getImage(filePath);
+    if (!imageData.error) {
+      bot.api.sendMessage(chatId,"validating image... ");
+      let response;
+      const size = imgSize ? imgSize : config.imageGen.sessionDefault.imgSize;
+      if (isNaN(+prompt)) {
+        const n = numImages ? numImages : config.imageGen.sessionDefault.numImages;
+        //@ts-ignore
+        response = await openai.createImageEdit(
+          //@ts-ignore
+          imageData.file,
+          prompt,
+          undefined,
+          n,
+          size
+        );
+      } else {
+        const size = imgSize ? imgSize : config.imageGen.sessionDefault.imgSize;
+        const n = parseInt(prompt)
+        response = await openai.createImageVariation(
+          //@ts-ignore
+          imageData.file,
+          n > 10 ? 1 : n,
+          size
+        );
+      }
+      bot.api.sendMessage(chatId,"generating the output...");
+      deleteFile(imageData.fileName!)
+      return response.data.data;
+    } else {
+      bot.api.sendMessage(chatId,imageData.error)
+    }
+  } catch (error: any) {
+    throw error;
+  }
+}
+
+export async function improvePrompt(promptText: string) {
+  const prompt = `Improve this image description using max 100 words and don't add additional text: ${promptText} `;
+  try {
+    const payload = {
+      model: config.imageGen.completions.model,
+      prompt: prompt,
+      max_tokens: config.imageGen.completions.maxTokens,
+      temperature: config.imageGen.completions.temperature,
+    };
+    const response = await openai.createCompletion(payload as CreateCompletionRequest); 
+    return response.data.choices[0].text;
+  } catch (e) {
+    throw e 
+  }
+}

src/modules/image-gen/controller/index.ts

@@ -0,0 +1,77 @@
+import { bot } from '../../../bot'
+import {
+  improvePrompt,
+  postGenerateImg,
+  alterGeneratedImg,
+} from "../api/openAi"; 
+
+interface ImageGenPayload {
+  chatId: number,
+  prompt: string,
+  numImages?: number,
+  imgSize?: string,
+  filePath?: string,
+}
+
+export const imgGen = async (data: ImageGenPayload) => {
+  const { chatId, prompt, numImages, imgSize } = data
+  try {
+    bot.api.sendMessage(chatId, "generating the output...");
+    const imgs = await postGenerateImg(prompt, numImages, imgSize);
+    imgs.map((img: any) => {
+      bot.api.sendPhoto(chatId, img.url) //ctx.replyWithPhoto(img.url);
+    });
+    return true
+  } catch (e) {
+    console.log("/gen Error", e);
+    bot.api.sendMessage(chatId, "There was an error while generating the image");
+    return false
+  }
+}
+
+export const imgGenEnhanced = async (data: ImageGenPayload) => {
+  const { chatId, prompt, numImages, imgSize } = data
+  try {
+    const upgratedPrompt = await improvePrompt(prompt);
+    if (upgratedPrompt) {
+      bot.api.sendMessage(chatId, 
+        `The following description was added to your prompt: ${upgratedPrompt}`
+      );
+    }
+    const imgs = await postGenerateImg(
+      upgratedPrompt || prompt,
+      numImages,
+      imgSize
+    );
+    imgs.map((img: any) => {
+      bot.api.sendPhoto(chatId, img.url);
+    });
+    return true
+  } catch (e) {
+    console.log("/genEn Error", e);
+    bot.api.sendMessage(chatId, "There was an error while generating the image");
+    return false
+  }
+}
+
+export const alterImg = async (data: ImageGenPayload) => {
+  const { chatId, prompt, numImages, imgSize, filePath } = data
+  try {
+    const imgs = await alterGeneratedImg(
+      chatId,
+      prompt!,
+      filePath!,
+      numImages!,
+      imgSize!
+    );
+    imgs!.map((img: any) => {
+      bot.api.sendPhoto(chatId, img.url);
+    });
+  } catch (e) {
+    console.log("/genEn Error", e);
+    bot.api.sendMessage(chatId, "There was an error while generating the image");
+    return false
+  }
+
+}
+