dallecord.py

"""
A Cog to be used with discord.py-bridge to generate images from a prompt using Dalle2.
"""

import discord
from dalle2 import Dalle2
from discord.ext import commands, bridge

dalle = Dalle2('sess-xxxxxxxxxxxxxxxxxxxxxxxxxxxx') # Replace with your session key

cmd_metadata = {
    'name': 'dalle2',
    'aliases': ['dalle'],
    'description': 'Generates an image from a prompt.'
}

class Dalle(commands.Cog):
    def __init__(self, client):
        self.client = client
    
    @bridge.bridge_command(name=cmd_metadata['name'], aliases=cmd_metadata['aliases'], description=cmd_metadata['description'])
    async def dallecmd(self, ctx, prompt: discord.Option(str, description="The prompt to generate an image from", required=True)):
        """
        Initialize a Discord Bridge command for Dalle2.
        Inherit the bridge_command decorator from discord.ext.bridge.
        """
        try:
            embed = discord.Embed(
                title=':mag: Dalle AI',
                description='Generating...',
                color=0x9e9e9e
            ).add_field(
                name='Query',
                value=prompt)
            
            dalle_embed = await ctx.reply(embed=embed)

            generations = dalle.generate(prompt)

            await dalle_embed.delete_original_response()

            photo_embeds = []

            for generation in generations:
                url = generation['generation']['image_path']

                photo = discord.Embed(
                    title='Results',
                    description=f'**Query:**\n{prompt}',
                    url='https://cdn.discordapp.com/'
                ).set_image(url=url).set_footer(text=f'Generated by: {ctx.author.display_name}')

                photo_embeds.append(photo)

            await ctx.reply(embeds=photo_embeds)
            return

        except Exception as e:
            print(e)
            query_error_embed = discord.Embed(
                title=':x: Error',
                description='Failed to generate image. The server returned:\n```{}```'.format(
                    e),
                color=0xff0000
            )
            await ctx.channel.send(embed=query_error_embed)
            return


def setup(client):
    """
    Initialize the Cog with the client on load.
    """
    client.add_cog(Dalle(client))