From 2a8e30d3297d7ca1b3d5dd86a4a971a157f41dc9 Mon Sep 17 00:00:00 2001
From: Zero6992 <mirtris0522@gmail.com>
Date: Tue, 7 Nov 2023 16:58:49 +0800
Subject: [PATCH] feat: update image generator to dalle3

---
 src/art.py | 54 ++++++++----------------------------------------------
 src/bot.py | 34 ++++++----------------------------
 2 files changed, 14 insertions(+), 74 deletions(-)

diff --git a/src/art.py b/src/art.py
index 022da941..a1d6766d 100644
--- a/src/art.py
+++ b/src/art.py
@@ -1,59 +1,21 @@
 import os
-import json
-
 import openai
-from pathlib import Path
-from base64 import b64decode
+
 from dotenv import load_dotenv
 from asgiref.sync import sync_to_async
 
 load_dotenv()
 openai.api_key = os.getenv("OPENAI_API_KEY")
 
+async def draw(prompt) -> list[str]:
 
-# generate 512x512 image and save to a file
-# return the path of the image as a str
-async def draw(prompt, amount) -> list[str]:
-    DATA_DIR = Path.cwd()
-    DATA_DIR.mkdir(exist_ok=True)
-
-    response = await sync_to_async(openai.Image.create)(
+    response = await sync_to_async(openai.images.generate)(
+        model="dall-e-3",
         prompt=prompt,
-        n=amount,
+        n=1,
         size="1024x1024",
-        response_format="b64_json",
+        quality="standard",
     )
-    with open("response.log", mode="w", encoding="utf-8") as file:
-        json.dump(response, file)
-
-    file_name = DATA_DIR / f"{prompt[:5]}-{response['created']}.json"
-
-    with open(file_name, mode="w", encoding="utf-8") as file:
-        json.dump(response, file)
-
-    path = await convert(file_name)
-    path = [str(p) for p in path]
-    return path
-
-# code stolen from https://realpython.com/generate-images-with-dalle-openai-api/
-async def convert(path):
-    DATA_DIR = Path.cwd() / "responses"
-    JSON_FILE = DATA_DIR / path
-    IMAGE_DIR = Path.cwd() / "images"
-    IMAGE_DIR.mkdir(parents=True, exist_ok=True)
-
-    with open(JSON_FILE, mode="r", encoding="utf-8") as file:
-        response = json.load(file)
-    image_files = []
-    for index, image_dict in enumerate(response["data"]):
-        image_data = b64decode(image_dict["b64_json"])
-        image_file = IMAGE_DIR / f"{JSON_FILE.stem}-{index}.png"
-        image_files.append(image_file)
-
-        with open(image_file, mode="wb") as png:
-            png.write(image_data)
-
-        # delete uneeded json file
-    os.remove(path)
+    image_url = response.data[0].url
 
-    return image_files
+    return image_url
diff --git a/src/bot.py b/src/bot.py
index bc0ec251..1a76e956 100644
--- a/src/bot.py
+++ b/src/bot.py
@@ -202,20 +202,8 @@ async def info(interaction: discord.Interaction):
 """)
 
 
-    @client.tree.command(name="draw", description="Generate an image with the Dalle2 model")
-    @app_commands.choices(amount=[
-        app_commands.Choice(name="1", value=1),
-        app_commands.Choice(name="2", value=2),
-        app_commands.Choice(name="3", value=3),
-        app_commands.Choice(name="4", value=4),
-        app_commands.Choice(name="5", value=5),
-        app_commands.Choice(name="6", value=6),
-        app_commands.Choice(name="7", value=7),
-        app_commands.Choice(name="8", value=8),
-        app_commands.Choice(name="9", value=9),
-        app_commands.Choice(name="10", value=10),
-    ])
-    async def draw(interaction: discord.Interaction, *, prompt: str, amount: int = 1):
+    @client.tree.command(name="draw", description="Generate an image with the Dall-e-3 model")
+    async def draw(interaction: discord.Interaction, *, prompt: str):
         if interaction.user == client.user:
             return
 
@@ -226,24 +214,14 @@ async def draw(interaction: discord.Interaction, *, prompt: str, amount: int = 1
 
         await interaction.response.defer(thinking=True, ephemeral=client.isPrivate)
         try:
-            path = await art.draw(prompt, amount)
-            files = []
-            for idx, img in enumerate(path):
-                files.append(discord.File(img, filename=f"image{idx}.png"))
-            title = f'> **{prompt}** - {str(interaction.user.mention)} \n\n'
-
-            await interaction.followup.send(files=files, content=title)
+            image_url = await art.draw(prompt)
 
-        except openai.InvalidRequestError:
-            await interaction.followup.send(
-                "> **ERROR: Inappropriate request 😿**")
-            logger.info(
-            f"\x1b[31m{username}\x1b[0m made an inappropriate request.!")
+            await interaction.followup.send(image_url)
 
         except Exception as e:
             await interaction.followup.send(
-                "> **ERROR: Something went wrong 😿**")
-            logger.exception(f"Error while generating image: {e}")
+                f'> **Something Went Wrong: {e}**')
+            logger.info(f"\x1b[31m{username}\x1b[0m :{e}")
 
 
     @client.tree.command(name="switchpersona", description="Switch between optional chatGPT jailbreaks")