[Dall-E] Initial commit

README.md

@@ -13,6 +13,8 @@ chatchart - Generates a pie chart to display chat activity over the last 5000 me
 
 dadjokes - Another UltimatePancake cog. Get some dad jokes on command.
 
+dalle - A cog to generate images from the Dall-E mini service. This cog should not be used on a public bot.
+
 dictionary - Define words and look up antonyms and synonyms. Originally by UltimatePancake.
 
 discordexperiments - Create voice channel invites for various built-in apps. This is only for developers or for people that can read the code and assess the risk of using it.

dalle/__init__.py

@@ -0,0 +1,9 @@
+from .dalle import DallE
+
+
+__red_end_user_data_statement__ = "This cog does not persistently store data or metadata about users."
+
+
+async def setup(bot):
+    n = DallE(bot)
+    bot.add_cog(n)

dalle/dalle.py

@@ -0,0 +1,109 @@
+import aiohttp
+import base64
+import discord
+from discord.http import Route
+import io
+import json
+from typing import List
+
+from redbot.core import commands
+
+
+class DallE(commands.Cog):
+    """Dall-E mini image generation"""
+
+    def __init__(self, bot):
+        self.bot = bot
+
+    async def red_delete_data_for_user(self, **kwargs):
+        """Nothing to delete."""
+        return
+
+    @commands.max_concurrency(3, commands.BucketType.default)
+    @commands.command()
+    @commands.guild_only()
+    async def generate(self, ctx: commands.Context, *, prompt: str):
+        """
+        Generate images through Dall-E mini.
+
+        https://huggingface.co/spaces/dalle-mini/dalle-mini
+        """
+        embed_links = ctx.channel.permissions_for(ctx.guild.me).embed_links
+        if not embed_links:
+            return await ctx.send("I need the `Embed Links` permission here before you can use this command.")
+
+        status_msg = await ctx.send("Image generator starting up...")
+        images = None
+        attempt = 0
+        async with ctx.typing():
+            while not images:
+                if attempt < 100:
+                    attempt += 1
+                    if attempt % 2 == 0:
+                        status = f"This will take a very long time. Once a response is acquired, this counter will pause while processing.\n[attempt `{attempt}/100`]"
+                        try:
+                            await status_msg.edit(content=status)
+                        except discord.NotFound:
+                            status_msg = await ctx.send(status)
+
+                images = await self.generate_images(prompt)
+
+        file_images = [discord.File(images[i], filename=f"{i}.png") for i in range(len(images))]
+        if len(file_images) == 0:
+            return await ctx.send(f"I didn't find anything for `{prompt}`.")
+        file_images = file_images[:4]
+
+        embed = discord.Embed(
+            colour=await ctx.embed_color(),
+            title="Dall-E Mini results",
+            url="https://huggingface.co/spaces/dalle-mini/dalle-mini",
+        )
+        embeds = []
+        for i, image in enumerate(file_images):
+            em = embed.copy()
+            em.set_image(url=f"attachment://{i}.png")
+            em.set_footer(text="View this output on a desktop client for best results.")
+            embeds.append(em)
+
+        form = []
+        payload = {"embeds": [e.to_dict() for e in embeds]}
+        form.append({"name": "payload_json", "value": discord.utils.to_json(payload)})
+        if len(file_images) == 1:
+            file = file_images[0]
+            form.append(
+                {
+                    "name": "file",
+                    "value": file.fp,
+                    "filename": file.filename,
+                    "content_type": "application/octet-stream",
+                }
+            )
+        else:
+            for index, file in enumerate(file_images):
+                form.append(
+                    {
+                        "name": f"file{index}",
+                        "value": file.fp,
+                        "filename": file.filename,
+                        "content_type": "application/octet-stream",
+                    }
+                )
+
+        try:
+            await status_msg.delete()
+        except discord.NotFound:
+            pass
+
+        r = Route("POST", "/channels/{channel_id}/messages", channel_id=ctx.channel.id)
+        await ctx.guild._state.http.request(r, form=form, files=file_images)
+
+    @staticmethod
+    async def generate_images(prompt: str) -> List[io.BytesIO]:
+        async with aiohttp.ClientSession() as session:
+            async with session.post("https://bf.dallemini.ai/generate", json={"prompt": prompt}) as response:
+                if response.status == 200:
+                    response_data = await response.json()
+                    images = [io.BytesIO(base64.decodebytes(bytes(image, "utf-8"))) for image in response_data["images"]]
+                    return images
+                else:
+                    return None

dalle/info.json

@@ -0,0 +1,9 @@
+{
+    "author": ["aikaterna"],
+    "install_msg": "While this cog can use Dall-E mini and retrieve images, the API is *very slow* and is overloaded, resulting in times where responses cannot be generated. Because of how slow this is, there are some restrictions in place:\n\n- This cog/command can only be used by 3 people globally at a time.\n\n- Do not use this cog on a public bot. I am not responsible if you get your bot IP banned from Dall-E by trying to send too many requests to this service. This cog could be used ideally on bots under 5-10 Discord servers, or under 1000 users total. Use at your own risk.\n\nIf you like this cog, give Jackenmen, Flame, and TrustyJAID a high five for helping me figure some things out. Use [p]generate with a prompt, after the cog is loaded with `[p]load dalle`, to generate images from Dall-E mini.",
+    "short": "Fetch images made by Dall-E mini from a prompt.",
+    "description": "Fetch images made by Dall-E mini from a prompt.",
+    "tags": ["dalle", "dall-e"],
+    "permissions": ["embed_links"],
+    "end_user_data_statement": "This cog does not persistently store data or metadata about users."
+}