Make the DALLe image size configurable (#365)

father-bot · Sep 13, 2023 · acd4cab · acd4cab
1 parent f87b54d
commit acd4cab
Show file tree

Hide file tree

Showing 4 changed files with 6 additions and 4 deletions.
diff --git a/bot/bot.py b/bot/bot.py
@@ -382,7 +382,7 @@ async def generate_image_handle(update: Update, context: CallbackContext, messag
     message = message or update.message.text
 
     try:
-        image_urls = await openai_utils.generate_images(message, n_images=config.return_n_generated_images)
+        image_urls = await openai_utils.generate_images(message, n_images=config.return_n_generated_images, size=config.image_size)
     except openai.error.InvalidRequestError as e:
         if str(e).startswith("Your request was rejected as a result of our safety system"):
             text = "🥲 Your request <b>doesn't comply</b> with OpenAI's usage policies.\nWhat did you write there, huh?"

diff --git a/bot/config.py b/bot/config.py
@@ -19,6 +19,7 @@
 new_dialog_timeout = config_yaml["new_dialog_timeout"]
 enable_message_streaming = config_yaml.get("enable_message_streaming", True)
 return_n_generated_images = config_yaml.get("return_n_generated_images", 1)
+image_size = config_yaml.get("image_size", "512x512")
 n_chat_modes_per_page = config_yaml.get("n_chat_modes_per_page", 5)
 mongodb_uri = f"mongodb://mongo:{config_env['MONGODB_PORT']}"
 

diff --git a/bot/openai_utils.py b/bot/openai_utils.py
@@ -194,8 +194,8 @@ async def transcribe_audio(audio_file):
     return r["text"]
 
 
-async def generate_images(prompt, n_images=4):
-    r = await openai.Image.acreate(prompt=prompt, n=n_images, size="512x512")
+async def generate_images(prompt, n_images=4, size="512x512"):
+    r = await openai.Image.acreate(prompt=prompt, n=n_images, size=size)
     image_urls = [item.url for item in r.data]
     return image_urls
 

diff --git a/config/config.example.yml b/config/config.example.yml
@@ -5,9 +5,10 @@ allowed_telegram_usernames: []  # if empty, the bot is available to anyone. pass
 new_dialog_timeout: 600  # new dialog starts after timeout (in seconds)
 return_n_generated_images: 1
 n_chat_modes_per_page: 5
+image_size: "512x512" # the image size for image generation. Generated images can have a size of 256x256, 512x512, or 1024x1024 pixels. Smaller sizes are faster to generate.
 enable_message_streaming: true  # if set, messages will be shown to user word-by-word
 
 # prices
 chatgpt_price_per_1000_tokens: 0.002
 gpt_price_per_1000_tokens: 0.02
-whisper_price_per_1_min: 0.006
+whisper_price_per_1_min: 0.006