phixxy/sparkytron3000
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Merge pull request #25 from phixxy/dev

Browse source 
fixed bug when dalle thinks prompt is nsfw
This commit is contained in:
phixxy 2024-02-23 23:17:05 -08:00 • committed by GitHub
commit 6465536ac4
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 17 additions and 17 deletions
Download patch file Download diff file Expand all files Collapse all files

34
cogs/chatgpt.py
View file

@ -200,7 +200,7 @@ class ChatGPT(commands.Cog):
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 dalle_api_call(self, prompt, model="dall-e-2", quality="standard", size="1024x1024"): async def dalle_api_call(self, prompt: str, model: str="dall-e-2", quality: str="standard", size: str="1024x1024") -> tuple:
data = { data = {
"model": model, "model": model,
"prompt": prompt, "prompt": prompt,
@ -209,35 +209,35 @@ class ChatGPT(commands.Cog):
"n":1, "n":1,
} }
url = "https://api.openai.com/v1/images/generations" url = "https://api.openai.com/v1/images/generations"
async with self.http_session.post(url, json=data, headers=self.headers) as resp:
try: response_data = await resp.json()
async with self.http_session.post(url, json=data, headers=self.headers) as resp: if resp.status == 200:
response_data = await resp.json()
response = response_data['data'][0]['url'] response = response_data['data'][0]['url']
return response else:
response = response_data["error"]["message"]
self.logger.info(f"Error occurred in dalle: {resp.status} | {response}")
return (resp.status, response)
except Exception as error: async def download_image(self, url, destination) -> int:
self.logger.exception("Error occurred in dalle")
return "Error occurred in dalle"
async def download_image(self, url, destination):
if url == "Error occurred in dalle":
return
async with self.http_session.get(url) as resp: async with self.http_session.get(url) as resp:
if resp.status == 200: if resp.status == 200:
f = await aiofiles.open(destination, mode='wb') f = await aiofiles.open(destination, mode='wb')
await f.write(await resp.read()) await f.write(await resp.read())
await f.close() await f.close()
return destination return resp.status
async def generate_dalle_image(self, ctx, model, quality="standard", size="1024x1024"):
async def generate_dalle_image(self, ctx, model, quality="standard", size="1024x1024") -> None:
if ctx.author.get_role(self.premium_role): if ctx.author.get_role(self.premium_role):
prompt = ctx.message.content.split(" ", maxsplit=1)[1] prompt = ctx.message.content.split(" ", maxsplit=1)[1]
await ctx.send(f"Please be patient this may take some time! Generating: {prompt}.") await ctx.send(f"Please be patient this may take some time! Generating: {prompt}.")
image_url = await self.dalle_api_call(prompt, model=model, quality=quality, size=size) resp_status, resp = await self.dalle_api_call(prompt, model=model, quality=quality, size=size)
if resp_status != 200:
await ctx.send(f"Error generating image: {resp_status}: {resp}")
return
my_filename = str(time.time_ns()) + ".png" my_filename = str(time.time_ns()) + ".png"
image_filepath = f"{self.data_dir}dalle/{my_filename}" image_filepath = f"{self.data_dir}dalle/{my_filename}"
await self.download_image(image_url, image_filepath) await self.download_image(resp, image_filepath)
with open(image_filepath, "rb") as fh: with open(image_filepath, "rb") as fh:
f = discord.File(fh, filename=image_filepath) f = discord.File(fh, filename=image_filepath)
prompt = prompt.replace('\n',' ') prompt = prompt.replace('\n',' ')