feat: refactor image genderation backend

README.md

@@ -4,8 +4,8 @@ This is a simple Mattermost Bot that uses OpenAI's GPT API(or self-host models)
 
 ## Feature
 
-1. Support Openai ChatGPT
+1. Support official openai api and self host models([LocalAI](https://localai.io/model-compatibility/))
-3. ChatGPT web ([pandora](https://github.com/pengzhile/pandora))
+2. Image Generation with [DALL·E](https://platform.openai.com/docs/api-reference/images/create) or [LocalAI](https://localai.io/features/image-generation/) or [stable-diffusion-webui](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/API)
 ## Installation and Setup
 
 See https://github.com/hibobmaster/mattermost_bot/wiki
@@ -21,18 +21,14 @@ docker compose up -d
 - `!help` help message
 - `!gpt + [prompt]` generate a one time response from chatGPT
 - `!chat + [prompt]` chat using official chatGPT api with context conversation
-- `!pic + [prompt]` generate an image from Bing Image Creator
+- `!pic + [prompt]` Image generation with DALL·E or LocalAI or stable-diffusion-webui
 
-The following commands need pandora http api: https://github.com/pengzhile/pandora/blob/master/doc/wiki_en.md#http-restful-api
-- `!talk + [prompt]` chat using chatGPT web with context conversation
-- `!goon` ask chatGPT to complete the missing part from previous conversation
 - `!new` start a new converstaion
 
 ## Demo
 Remove support for Bing AI, Google Bard due to technical problems.
-![demo1](https://i.imgur.com/XRAQB4B.jpg)
+![gpt command](https://imgur.com/vdT83Ln.jpg)
-![demo2](https://i.imgur.com/if72kyH.jpg)
+![image generation](https://i.imgur.com/GHczfkv.jpg)
-![demo3](https://i.imgur.com/GHczfkv.jpg)
 
 ## Thanks
 <a href="https://jb.gg/OpenSourceSupport" target="_blank">

requirements.txt

@@ -1,4 +1,3 @@
-aiofiles
 httpx
 Pillow
 tiktoken

src/bot.py

@@ -4,9 +4,11 @@ import json
 import asyncio
 import re
 import os
+from pathlib import Path
 from gptbot import Chatbot
 from log import getlogger
 import httpx
+import imagegen
 
 logger = getlogger()
 
@@ -78,6 +80,11 @@ class Bot:
         self.image_generation_backend: str = image_generation_backend
         self.timeout = timeout or 120.0
 
+        self.base_path = Path(os.path.dirname(__file__)).parent
+
+        if not os.path.exists(self.base_path / "images"):
+            os.mkdir(self.base_path / "images")
+
         #  httpx session
         self.httpx_client = httpx.AsyncClient()
 
@@ -213,22 +220,38 @@ class Bot:
                     raise Exception(e)
 
             # !pic command trigger handler
-            if self.pic_prog.match(message):
+            if self.image_generation_endpoint and self.image_generation_backend:
-                prompt = self.pic_prog.match(message).group(1)
+                if self.pic_prog.match(message):
-                # generate image
+                    prompt = self.pic_prog.match(message).group(1)
-                try:
+                    # generate image
-                    links = await self.imagegen.get_images(prompt)
+                    try:
-                    image_path = await self.imagegen.save_images(links, "images")
+                        # generate image
-                except Exception as e:
+                        b64_datas = await imagegen.get_images(
-                    logger.error(e, exc_info=True)
+                            self.httpx_client,
-                    raise Exception(e)
+                            self.image_generation_endpoint,
-
+                            prompt,
-                # send image
+                            self.image_generation_backend,
-                try:
+                            timeount=self.timeout,
-                    await self.send_file(channel_id, prompt, image_path)
+                            api_key=self.openai_api_key,
-                except Exception as e:
+                            n=1,
-                    logger.error(e, exc_info=True)
+                            size="256x256",
-                    raise Exception(e)
+                        )
+                        image_path_list = await asyncio.to_thread(
+                            imagegen.save_images,
+                            b64_datas,
+                            self.base_path / "images",
+                        )
+                        # send image
+                        for image_path in image_path_list:
+                            await self.send_file(
+                                channel_id,
+                                f"{prompt}",
+                                image_path,
+                                root_id,
+                            )
+                    except Exception as e:
+                        logger.error(e, exc_info=True)
+                        raise Exception(e)
 
             # !help command trigger handler
             if self.help_prog.match(message):
@@ -248,7 +271,9 @@ class Bot:
         )
 
     # send file to room
-    async def send_file(self, channel_id: str, message: str, filepath: str) -> None:
+    async def send_file(
+        self, channel_id: str, message: str, filepath: str, root_id: str
+    ) -> None:
         filename = os.path.split(filepath)[-1]
         try:
             file_id = await self.driver.files.upload_file(
@@ -256,7 +281,8 @@ class Bot:
                 files={
                     "files": (filename, open(filepath, "rb")),
                 },
-            )["file_infos"][0]["id"]
+            )
+            file_id = file_id["file_infos"][0]["id"]
         except Exception as e:
             logger.error(e, exc_info=True)
             raise Exception(e)
@@ -267,6 +293,7 @@ class Bot:
                     "channel_id": channel_id,
                     "message": message,
                     "file_ids": [file_id],
+                    "root_id": root_id,
                 }
             )
             # remove image after posting

src/imagegen.py

@@ -15,7 +15,7 @@ async def get_images(
             url,
             headers={
                 "Content-Type": "application/json",
-                "Authorization": "Bearer " + kwargs.get("api_key"),
+                "Authorization": f"Bearer {kwargs.get('api_key')}",
             },
             json={
                 "prompt": prompt,