diff --git a/extensions/chatgpt.py b/extensions/chatgpt.py index e96783b..89470da 100644 --- a/extensions/chatgpt.py +++ b/extensions/chatgpt.py @@ -211,70 +211,50 @@ class ChatGPT(commands.Cog): await f.close() return destination - @commands.command( - description="Dalle 2", - help="Generate an image with Dalle 2 Usage: !dalle2 (prompt)", - brief="Generate Image" - ) - async def dalle2(self, ctx): - if ctx.author.get_role(self.premium_role): - prompt = ctx.message.content.split(" ", maxsplit=1)[1] - img_url = await self.get_dalle(prompt) - my_filename = str(time.time_ns()) + ".png" - filepath = f"{self.data_dir}dalle2/{my_filename}" - await self.download_image(img_url, filepath) - with open(filepath, "rb") as fh: - f = discord.File(fh, filename=filepath) - log_data = f'Author: {ctx.author.name}, Prompt: {prompt}, Filename: {my_filename}\n' - with open(f"{self.data_dir}logs/dalle2.log", 'a') as log_file: - log_file.writelines(log_data) - await ctx.send(f'Generated by: {ctx.author.name}\nPrompt: {prompt}', file=f) - else: - await ctx.send("Sorry you must be a premium member to use this command. (!donate)") +async def generate_dalle_image(self, ctx, model, quality=None, size=None): + if ctx.author.get_role(self.premium_role): + prompt = ctx.message.content.split(" ", maxsplit=1)[1] + image_url = await self.get_dalle(prompt, model=model, quality=quality, size=size) + my_filename = str(time.time_ns()) + ".png" + image_filepath = f"{self.data_dir}dalle/{my_filename}" + await self.download_image(image_url, image_filepath) + with open(image_filepath, "rb") as fh: + f = discord.File(fh, filename=image_filepath) + log_data = f'Author: {ctx.author.name}, Prompt: {prompt}, Filename: {my_filename}\n' + with open(f"{self.data_dir}logs/dalle3.log", 'a') as log_filepath: + log_filepath.writelines(log_data) + await ctx.send(f'Generated by: {ctx.author.name}\nPrompt: {prompt}', file=f) + else: + await ctx.send("Sorry you must be a premium member to use this command. (!donate)") - @commands.command( - description="Dalle 3", - help="Generate an image with Dalle 3 Usage: !dalle3 (prompt)", - brief="Generate Image", - aliases = ['dalle'] - ) - async def dalle3(self, ctx): - if ctx.author.get_role(self.premium_role): - prompt = ctx.message.content.split(" ", maxsplit=1)[1] - img_url = await self.get_dalle(prompt, model="dall-e-3") - my_filename = str(time.time_ns()) + ".png" - filepath = f"{self.data_dir}dalle/{my_filename}" - await self.download_image(img_url, filepath) - with open(filepath, "rb") as fh: - f = discord.File(fh, filename=filepath) - log_data = f'Author: {ctx.author.name}, Prompt: {prompt}, Filename: {my_filename}\n' - with open(f"{self.data_dir}logs/dalle3.log", 'a') as log_file: - log_file.writelines(log_data) - await ctx.send(f'Generated by: {ctx.author.name}\nPrompt: {prompt}', file=f) - else: - await ctx.send("Sorry you must be a premium member to use this command. (!donate)") - @commands.command( - description="Dalle 3 HD", - help="Generate an HD image with Dalle 3 Usage: !dalle3 (prompt)", - brief="Generate HD Image", - ) - async def dalle3hd(self, ctx): - if ctx.author.get_role(self.premium_role): - prompt = ctx.message.content.split(" ", maxsplit=1)[1] - image_url = await self.get_dalle(prompt, model="dall-e-3", quality="hd", size="1792x1024") - my_filename = str(time.time_ns()) + ".png" - image_filepath = f"{self.data_dir}dalle/{my_filename}" - await self.download_image(image_url, image_filepath) - with open(image_filepath, "rb") as fh: - f = discord.File(fh, filename=image_filepath) - log_data = f'Author: {ctx.author.name}, Prompt: {prompt}, Filename: {my_filename}\n' - with open(f"{self.data_dir}logs/dalle3.log", 'a') as log_filepath: - log_filepath.writelines(log_data) - await ctx.send(f'Generated by: {ctx.author.name}\nPrompt: {prompt}', file=f) - else: - await ctx.send("Sorry you must be a premium member to use this command. (!donate)") - +@commands.command( + description="Dalle 2", + help="Generate an image with Dalle 2 Usage: !dalle2 (prompt)", + brief="Generate Image" + ) +async def dalle2(self, ctx): + await self.generate_dalle_image(ctx, model="dall-e-2") + + +@commands.command( + description="Dalle 3", + help="Generate an image with Dalle 3 Usage: !dalle3 (prompt)", + brief="Generate Image", + aliases = ['dalle'] + ) +async def dalle3(self, ctx): + await self.generate_dalle_image(ctx, model="dall-e-3") + + +@commands.command( + description="Dalle 3 HD", + help="Generate an HD image with Dalle 3 Usage: !dalle3 (prompt)", + brief="Generate HD Image", + ) +async def dalle3hd(self, ctx): + await self.generate_dalle_image(ctx, model="dall-e-3", quality="hd", size="1792x1024") + @commands.command( description="Image GPT4", help="Ask GPT4 a question about an image. Usage: !question_gpt4 (link) (question)",