From b4201ca5bce6c5fcebf30785d6dd1c0af493fa09 Mon Sep 17 00:00:00 2001 From: aikaterna <20862007+aikaterna@users.noreply.github.com> Date: Thu, 16 Jun 2022 11:18:29 -0700 Subject: [PATCH] [Dall-E] Initial commit --- README.md | 2 + dalle/__init__.py | 9 ++++ dalle/dalle.py | 109 ++++++++++++++++++++++++++++++++++++++++++++++ dalle/info.json | 9 ++++ 4 files changed, 129 insertions(+) create mode 100644 dalle/__init__.py create mode 100644 dalle/dalle.py create mode 100644 dalle/info.json diff --git a/README.md b/README.md index 69dd02f..da62268 100644 --- a/README.md +++ b/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. diff --git a/dalle/__init__.py b/dalle/__init__.py new file mode 100644 index 0000000..dcc4b45 --- /dev/null +++ b/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) diff --git a/dalle/dalle.py b/dalle/dalle.py new file mode 100644 index 0000000..bb230e3 --- /dev/null +++ b/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 diff --git a/dalle/info.json b/dalle/info.json new file mode 100644 index 0000000..f29bcb4 --- /dev/null +++ b/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." +}