From 7cf3874261c9cad01e214a730e97edc2f88d153a Mon Sep 17 00:00:00 2001
From: phixxy <bottlecap@protonmail.com>
Date: Sat, 27 Jan 2024 19:59:32 -0800
Subject: [PATCH] removed print statement in get_dalle now saves images locally
 and then upload them

---
 extensions/chatgpt.py | 36 +++++++++++++++++++++++++++++++-----
 1 file changed, 31 insertions(+), 5 deletions(-)

diff --git a/extensions/chatgpt.py b/extensions/chatgpt.py
index 314ff60..6466974 100644
--- a/extensions/chatgpt.py
+++ b/extensions/chatgpt.py
@@ -4,6 +4,7 @@ import time
 import json
 import random
 import asyncio
+import aiofiles
 import discord
 from discord.ext import commands, tasks
 
@@ -29,6 +30,8 @@ class ChatGPT(commands.Cog):
                 os.mkdir(self.data_dir + "config")
             if not os.path.exists(self.data_dir + "logs"):
                 os.mkdir(self.data_dir + "logs")
+            if not os.path.exists(self.data_dir + "dalle"):
+                os.mkdir(self.data_dir + "dalle")
         except:
             self.bot.logger.exception("ChatGPT failed to make directories")
 
@@ -115,7 +118,6 @@ class ChatGPT(commands.Cog):
         try:
             async with self.bot.http_session.post(url, headers=headers, json=data) as resp:
                 response_data = await resp.json()
-                print(response_data)
                 response = response_data['data'][0]['url']
                 return response
 
@@ -198,6 +200,14 @@ class ChatGPT(commands.Cog):
                 await ctx.send(chunk)
         else:
             await ctx.send("Sorry you must be a premium member to use this command. (!donate)")
+    
+    async def download_image(self, url, destination):
+        async with self.bot.http_session.get(url) as resp:
+            if resp.status == 200:
+                f = await aiofiles.open(destination, mode='wb')
+                await f.write(await resp.read())
+                await f.close()
+        return destination
 
     @commands.command(
         description="Dalle 2", 
@@ -206,10 +216,18 @@ class ChatGPT(commands.Cog):
         aliases = ['dalle']
         )         
     async def dalle2(self, ctx):
-        if ctx.author.get_role(self.premium_role):
+        if True:#ctx.author.get_role(self.premium_role):
             prompt = ctx.message.content.split(" ", maxsplit=1)[1]
             img_url = await self.get_dalle(prompt)
-            await ctx.send(f"Generated by: {ctx.author.name}\nPrompt: {prompt}\n{img_url}")
+            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}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)")
 
@@ -219,10 +237,18 @@ class ChatGPT(commands.Cog):
         brief="Generate Image"
         )           
     async def dalle3(self, ctx):
-        if ctx.author.get_role(self.premium_role):
+        if True:#ctx.author.get_role(self.premium_role):
             prompt = ctx.message.content.split(" ", maxsplit=1)[1]
             img_url = await self.get_dalle(prompt, "dall-e-3")
-            await ctx.send(f"Generated by: {ctx.author.name}\nPrompt: {prompt}\n{img_url}")
+            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}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)")