diff --git a/.gitignore b/.gitignore index 6f501321..8dce2325 100644 --- a/.gitignore +++ b/.gitignore @@ -163,4 +163,5 @@ cython_debug/ # be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore # and can be added to the global gitignore or merged into this file. For a more nuclear # option (not recommended) you can uncomment the following to ignore the entire idea folder. -#.idea/ +.idea/ +Pipfile* \ No newline at end of file diff --git a/src/art.py b/src/art.py index e12cd8c7..220e73c0 100644 --- a/src/art.py +++ b/src/art.py @@ -1,5 +1,7 @@ import os import json +from typing import List + import openai from pathlib import Path from base64 import b64decode @@ -12,16 +14,18 @@ # generate 512x512 image and save to a file # return the path of the image as a str -async def draw(prompt) -> 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)( prompt=prompt, - n=1, - size="512x512", + n=amount, + size="1024x1024", response_format="b64_json", ) + 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" @@ -29,8 +33,8 @@ async def draw(prompt) -> str: json.dump(response, file) path = await convert(file_name) - - return str(path) + 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): @@ -41,15 +45,16 @@ async def convert(path): 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) + os.remove(path) - return image_file + return image_files diff --git a/src/bot.py b/src/bot.py index 768ec30c..887936fd 100644 --- a/src/bot.py +++ b/src/bot.py @@ -197,7 +197,19 @@ async def info(interaction: discord.Interaction): """) @client.tree.command(name="draw", description="Generate an image with the Dalle2 model") - async def draw(interaction: discord.Interaction, *, prompt: str): + @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): if interaction.user == client.user: return @@ -208,14 +220,13 @@ async def draw(interaction: discord.Interaction, *, prompt: str): await interaction.response.defer(thinking=True, ephemeral=client.isPrivate) try: - path = await art.draw(prompt) - - file = discord.File(path, filename="image.png") - title = f'> **{prompt}** - <@{str(interaction.user.mention)}' + '> \n\n' - embed = discord.Embed(title=title) - embed.set_image(url="attachment://image.png") + 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(file=file, embed=embed) + await interaction.followup.send(files=files, content=title) except openai.InvalidRequestError: await interaction.followup.send(