From 0b20e0ac1ac503ed7c4504d09d97bd8f58c97f4e Mon Sep 17 00:00:00 2001 From: hibobmaster <32976627+hibobmaster@users.noreply.github.com> Date: Mon, 18 Sep 2023 15:19:27 +0800 Subject: [PATCH] feat: refactor image genderation backend --- README.md | 14 ++++------- requirements.txt | 1 - src/bot.py | 63 ++++++++++++++++++++++++++++++++++-------------- src/imagegen.py | 2 +- 4 files changed, 51 insertions(+), 29 deletions(-) diff --git a/README.md b/README.md index f5a4cc8..3b83913 100644 --- a/README.md +++ b/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 -3. ChatGPT web ([pandora](https://github.com/pengzhile/pandora)) +1. Support official openai api and self host models([LocalAI](https://localai.io/model-compatibility/)) +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) -![demo2](https://i.imgur.com/if72kyH.jpg) -![demo3](https://i.imgur.com/GHczfkv.jpg) +![gpt command](https://imgur.com/vdT83Ln.jpg) +![image generation](https://i.imgur.com/GHczfkv.jpg) ## Thanks diff --git a/requirements.txt b/requirements.txt index c7f6d1f..6314c78 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,4 +1,3 @@ -aiofiles httpx Pillow tiktoken diff --git a/src/bot.py b/src/bot.py index 19e4dee..988fed1 100644 --- a/src/bot.py +++ b/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): - prompt = self.pic_prog.match(message).group(1) - # generate image - try: - links = await self.imagegen.get_images(prompt) - image_path = await self.imagegen.save_images(links, "images") - except Exception as e: - logger.error(e, exc_info=True) - raise Exception(e) - - # send image - try: - await self.send_file(channel_id, prompt, image_path) - except Exception as e: - logger.error(e, exc_info=True) - raise Exception(e) + if self.image_generation_endpoint and self.image_generation_backend: + if self.pic_prog.match(message): + prompt = self.pic_prog.match(message).group(1) + # generate image + try: + # generate image + b64_datas = await imagegen.get_images( + self.httpx_client, + self.image_generation_endpoint, + prompt, + self.image_generation_backend, + timeount=self.timeout, + api_key=self.openai_api_key, + n=1, + size="256x256", + ) + 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 diff --git a/src/imagegen.py b/src/imagegen.py index fb54f14..2214eac 100644 --- a/src/imagegen.py +++ b/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,