removed print statement in get_dalle

now saves images locally and then upload them
This commit is contained in:
phixxy 2024-01-27 19:59:32 -08:00
parent 2a7f6893c2
commit 7cf3874261

View file

@ -4,6 +4,7 @@ import time
import json import json
import random import random
import asyncio import asyncio
import aiofiles
import discord import discord
from discord.ext import commands, tasks from discord.ext import commands, tasks
@ -29,6 +30,8 @@ class ChatGPT(commands.Cog):
os.mkdir(self.data_dir + "config") os.mkdir(self.data_dir + "config")
if not os.path.exists(self.data_dir + "logs"): if not os.path.exists(self.data_dir + "logs"):
os.mkdir(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: except:
self.bot.logger.exception("ChatGPT failed to make directories") self.bot.logger.exception("ChatGPT failed to make directories")
@ -115,7 +118,6 @@ class ChatGPT(commands.Cog):
try: try:
async with self.bot.http_session.post(url, headers=headers, json=data) as resp: async with self.bot.http_session.post(url, headers=headers, json=data) as resp:
response_data = await resp.json() response_data = await resp.json()
print(response_data)
response = response_data['data'][0]['url'] response = response_data['data'][0]['url']
return response return response
@ -198,6 +200,14 @@ class ChatGPT(commands.Cog):
await ctx.send(chunk) await ctx.send(chunk)
else: else:
await ctx.send("Sorry you must be a premium member to use this command. (!donate)") 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( @commands.command(
description="Dalle 2", description="Dalle 2",
@ -206,10 +216,18 @@ class ChatGPT(commands.Cog):
aliases = ['dalle'] aliases = ['dalle']
) )
async def dalle2(self, ctx): 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] prompt = ctx.message.content.split(" ", maxsplit=1)[1]
img_url = await self.get_dalle(prompt) 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: else:
await ctx.send("Sorry you must be a premium member to use this command. (!donate)") 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" brief="Generate Image"
) )
async def dalle3(self, ctx): 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] prompt = ctx.message.content.split(" ", maxsplit=1)[1]
img_url = await self.get_dalle(prompt, "dall-e-3") 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: else:
await ctx.send("Sorry you must be a premium member to use this command. (!donate)") await ctx.send("Sorry you must be a premium member to use this command. (!donate)")