Added...gulp...image generation

reginaldCog/reginald.py

@@ -1,6 +1,11 @@
 import openai
 import os
 import random
+import requests
+import base64
+import aiohttp
+from io import BytesIO
+from PIL import Image
 from redbot.core import Config, commands
 
 class ReginaldCog(commands.Cog):
@@ -14,6 +19,19 @@ class ReginaldCog(commands.Cog):
             openai_api_key=None
         )
 
+    def has_kanium_role():
+        async def predicate(ctx):
+            kanium_role_id = 280260875678515200
+            return any(role.id == kanium_role_id for role in ctx.author.roles)
+        return commands.check(predicate)
+
+    def has_janitor_role():
+        async def predicate(ctx):
+            janitor_role_id = 672156832323600396
+            return any(role.id == janitor_role_id for role in ctx.author.roles)
+        return commands.check(predicate)
+
+
     @commands.guild_only()
     @commands.has_permissions(manage_guild=True)
     @commands.command(help="Set the OpenAI API key")
@@ -22,6 +40,7 @@ class ReginaldCog(commands.Cog):
         await ctx.send("OpenAI API key set successfully.")
 
     @commands.guild_only()
+    @has_kanium_role()
     @commands.command(help="Ask Reginald a question")
     @commands.cooldown(1, 60, commands.BucketType.user)  # 1-minute cooldown per user
     async def reginald(self, ctx, *, prompt=None):
@@ -83,6 +102,73 @@ class ReginaldCog(commands.Cog):
         chunks.append(response_text)
         return chunks
     
+    @commands.guild_only()
+    @has_janitor_role()
+    @commands.command(help="Ask Reginald to generate an image based on a prompt")
+    @commands.cooldown(1, 300, commands.BucketType.user)  # 5-minute cooldown per user
+    async def reginaldimagine(self, ctx, *, prompt=None):
+        ignored_user_id = 138125632876838912
+        if ctx.author.id == ignored_user_id:
+            return
+
+        if prompt is None:
+            await ctx.author.send("Please provide a prompt for the image.")
+            return
+
+        api_key = await self.config.guild(ctx.guild).openai_api_key()
+        if api_key is None:
+            await ctx.author.send('OpenAI API key not set. Please use the "!setreginaldcogapi" command to set the key.')
+            return
+
+        try:
+            image_url = await self.generate_image(api_key, prompt)
+            if image_url:
+                async with ctx.typing():
+                    async with aiohttp.ClientSession() as session:
+                        async with session.get(image_url) as resp:
+                            image_data = await resp.read()
+                            image = Image.open(BytesIO(image_data))
+                            image.save("image.png")
+                await ctx.send(file=discord.File("image.png"))
+                os.remove("image.png")
+            else:
+                await ctx.author.send("I apologize, but I am unable to generate an image based on the provided prompt.")
+        except Exception as e:
+            await ctx.author.send(f"I apologize, but I am unable to generate an image at this time. Error message: {str(e)}")
+
+    async def generate_image(self, api_key, prompt):
+        openai.api_key = api_key
+        model = "image-alpha-001"  # Replace this with the appropriate DALL-E model name
+        response = openai.Image.create(
+            model=model,
+            prompt=prompt,
+            n=1,
+            size="512x512",
+            response_format="url"
+        )
+
+        if response and response.choices and response.choices[0].url:
+            return response.choices[0].url
+        return None
+
+    @reginald.error
+    async def reginald_error(ctx, error):
+        if isinstance(error, commands.BadArgument):
+            await ctx.send("I'm sorry, but I couldn't understand your input. Please check your message and try again.")
+        elif isinstance(error, commands.CheckFailure):
+            await ctx.send("You do not have the required role to use this command.")
+        else:
+            await ctx.send(f"An unexpected error occurred: {error}")
+
+    @reginaldimagine.error
+    async def reginaldimagine_error(ctx, error):
+        if isinstance(error, commands.BadArgument):
+            await ctx.send("I'm sorry, but I couldn't understand your input. Please check your message and try again.")
+        elif isinstance(error, commands.CheckFailure):
+            await ctx.send("You do not have the required role to use this command.")
+        else:
+            await ctx.send(f"An unexpected error occurred: {error}")
+
 def setup(bot):
     cog = ReginaldCog(bot)
     bot.add_cog(cog)