Skip to content

Commit

Permalink
feat: update image generator to dalle3
Browse files Browse the repository at this point in the history
  • Loading branch information
Zero6992 committed Nov 7, 2023
1 parent cf25b39 commit 2a8e30d
Show file tree
Hide file tree
Showing 2 changed files with 14 additions and 74 deletions.
54 changes: 8 additions & 46 deletions src/art.py
Original file line number Diff line number Diff line change
@@ -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
34 changes: 6 additions & 28 deletions src/bot.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -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")
Expand Down

0 comments on commit 2a8e30d

Please sign in to comment.