Building A Stable Diffusion Discord Bot

Sharing is caring!

Last Updated on November 6, 2022 by Jay

In this tutorial we’ll go through everything you need to know to build a Stable Diffusion Discord Bot using Python. You should have a basic understanding of how to set up a Discord bot. Highly recommend you go through this short tutorial before proceeding with this one.

Required Libraries

There are two key components for making this bot: Stable Diffusion and Discord. First steps, we’ll need a working version of Stable Diffusion on our PC. Follow this tutorial to install Stable Diffusion: How to Run Stable Diffusion on Windows

After you successfully install Stable Diffusion, now is the time to install discord.py, which is a Python library to make Discord bots. Another library we need to use is asyncio, which is a Python built-in library so we don’t need to install it.

pip install discord.py

Bot Design

We’ll first talk about the bot design. An understanding of the design will help you understand the code much easier later on.

First of all, the bot runs on my spare PC with a GTX 1070 and 8GB of RAM. Because Stable Diffusion requires a GPU, you can’t just use any cloud server to host the bot. If you must use cloud hosting, GPU rental can be quite expensive. For me, the cost of hosting the bot is electricity for keeping my PC powered on 24/7.

Also, I want to caution you that this is my first time making a Discord bot and I’m pretty new to asynchronous programming. I might not be doing things following “best practices”. I think that’s okay and making mistakes is part of the learning journey. It’s important to get it working at least, then learn more about the subject and then improve on it.

If you are experienced in asynchronous programming or Discord bot, please leave a comment to point out where I’m not doing things right. I’m more than happy to learn from you!

Synchronous Programming

In synchronous programming, our program runs in sequence and one task at a time, which means it has to complete a task before moving on to the next.

The general workflow is like the following:

1. Users send a request to the bot by typing “.makeimg …”
2. Bot queues the request based on a “first come first serve” principle
3. The bot uses Stable Diffusion to generate an image
4. Bot sends the image back to the user
5. Repeat steps 3 and 4 if there are jobs in the queue

The workflow looks logical, except that it has a fatal flaw if we use synchronous programming. In general, Stable Diffusion uses GPU to generate images. It takes a high-end GPU a few seconds to crank out a few images, let alone my GTX 1070. Synchronous means that our bot will be “blocked” during the time it uses Stable Diffusion to generate images. This further implies that users won’t be able to add new requests to the queue while the bot is busy.

We need a solution such that our bot is capable of receiving new requests anytime, busy or not. Asynchronous programming comes to the rescue.

Asynchronous Programming

As the name hints, asynchronous programming is the opposite of synchronous programming. It means that our program doesn’t have to run in sequence. If a part of the program takes a long time to run, then it can jump to other parts of the program and carry out multiple tasks simultaneously.

In our case, we want the bot to be able to receive new requests while Stable Diffusion is running. The asynchronous workflow will look the following. Still the same process, but a little different and more efficient.

The secret of achieving the above is by using multithreading. By default, Python uses a single thread (you can think of this as a single CPU core) to run our code. All modern computers should have multiple cores, so we can use one thread for receiving and adding new requests to the job queue. Then use another thread for processing Stable Diffusion. This is plausible because Stable Diffusion is GPU-heavy and doesn’t really require much work from the CPU. Well, unless you use a CPU to run Stable Diffusion, then you can’t really use this method.

The 2 threads should also be able to communicate with each other so that the queue in thread 1 can pass information to thread 2 for image generation.

Asyncio Python Library

Asyncio is the built-in Python library for asynchronous programming. There are two main components when using asyncio: event loop and coroutine.

The event loop is like the central command center that monitors all jobs’ statuses and dispatches workers for certain tasks.

A coroutine is just like a Python function, but it has the ability to pause code execution and give control to another function.

We need to run the coroutine inside an event loop.

Creating A Stable Diffusion Bot Class

We’ll use the official Stable Diffusion code, and convert the “txt2img.py” script into a Stable Diffusion Bot class. This part is easy since the code is already 90% completed for us. First of all, we’ll need to create a class called SDBot. The original txt2img.py script is a command line tool. We’ll remove the command line input functionalities and then convert those arguments into class attributes. Feel free to set many default attributes such as the height and width of the image, etc.

You can find the final source code here: https://github.com/pythoninoffice/tutorials/blob/main/discord_bot_sd/txt2img_bot.py

Note the above code is saved inside a script named txt2img_bot.py.

Creating A Basic Discord Bot

You should have some basic understanding of how to set up a Discord bot. Read this part 1 tutorial if you need help setting up a basic Discord bot: Building A Simple Python Discord Bot with DiscordPy in 2022/2023

Initial Setup

We’ll start a new script for the Discord bot, we’ll call this script bot_v2.py. Below screenshot shows the folder structure:

Let’s do some initial setup by importing a bunch of things to our discord bot code. The below should look familiar from part 1 of the tutorial.

Note the line from scripts.txt2img_bot import SDBot will import the SDBot object into this script.

import discord
from discord.ext import commands
import asyncio
import pathlib
import hashlib
import random
from scripts.txt2img_bot import SDBot

TOKEN = 'YOUR_DISCORD_TOKEN'

intents = discord.Intents.default()
intents.message_content = True
#client = discord.Client(intents=intents)

client = commands.Bot(command_prefix='.', intents=intents)
queues = []
blocking = False
sd_bot = SDBot()
loop = None

@client.event
async def on_ready():
    print('bot ready')

###########################
### Main code body here ###
###########################


client.run(TOKEN)

Building A Stable Diffusion Discord Bot

Let’s start by making the makeimg command, which is the entry point for users to interact with the bot.

According to the flowchart above, the makeimg command should do two things: 1) adds new requests to the queue, 2) runs Stable Diffusion to make images, and sends them back to users.

Note that to make a command for Discord, we need to add the decorator @client.command() in front of the asynchronous function.

@client.command()
async def makeimg(ctx, prompt):
    global loop
    loop = asyncio.get_running_loop()
    print(loop)      
    que(ctx, prompt)
    await ctx.send(f'{prompt} added to queue')

    if blocking:
        print('this is blocking')
        await ctx.send("currently generating image, please wait")
    else:
        await asyncio.gather(asyncio.to_thread(sd_gen,ctx,queues))

This makeimg function takes two arguments, ctx (or context), and prompt.

Based on my understanding, a context is simply a message that users send to a Discord channel. Again, I could be wrong on this, if you know better, please point it out in the comment below! We can get a lot of information associated with this piece of context, including the user ID, the message content, and the channel/server the message was sent to.

I know it’s frowned upon to use global variables in Python. Our program is pretty simple so why not make it easier to code? The global variable loop is to get the current event loop in the current thread. Note this event loop will be in tread #1 as shown on the flowchart.

We’ll then write a helper function to add requests to the queue (See below). We also need a variable blocking to check the state of the bot. If the bot is currently blocked (because it’s running Stable Diffusion), then we simply send users a message stating so. If the bot is not busy/blocked, then it should start generating images.

When the bot starts generating images, it should spin off a separate thread and run Stable Diffusion in thread #2. We can achieve that by using the asyncio.to_thread(sd_gen,ctx,queues) method. We pass three things to that separate thread:

1. A function sd_gen for running Stable Diffusion
2. Context (ctx)
3. Prompt

Below is the code for the helper function que(). The queue is essentially a list of dictionaries. Each dictionary will have the user’s id as the key, and their prompts as the value.

queues = []

def que(ctx, prompt):
    user_id = ctx.message.author.mention
    queues.append({user_id:prompt})
    print(f'{prompt} added to queue')

Function To Generate Images

Moving on to the sd_gen function. The first thing to note – this is a normal function and not an async function. Remember this is also a helper function that the main makeimg function calls.

It first checks the queue to see if there are requests/prompts waiting to be processed. If yes, it will set the blocking variable to True. Then retrieve the first request and also remove it from the queue simultaneously by using the pop() built-in list function.

We then use the hashlib library to generate random hash values to ensure we don’t repeat file names.

We also added a block for extracting seed value from a prompt. Several discord users suggested having the ability to see and submit seed values is beneficial, so thank you for the feedback 🙂

Next, we call the sd_bot.makeimg() function to generate the image with the given prompt and seed in Stable Diffusion. The makeimg() method also saves the image to disk, so we don’t need to worry about saving here.

def sd_gen(ctx, queues):
    global blocking

    print(queues)
    if len(queues) > 0:
        blocking = True
        prompt = queues.pop(0)
        mention = list(prompt.keys())[0]
        prompt = list(prompt.values())[0] #convert from dictoinary to string
        filename = hashlib.sha256(prompt.encode('utf-8')).hexdigest()
        
        ## extract seed and prompt
        try:
            seed = int(prompt.lower().split('seed')[1].split('=')[1].strip())
        except:
            seed = random.randint(0,4294967295)
        prompt = prompt.split('seed')[0]
        
        ## generate image
        sd_bot.makeimg(prompt, filename, seed)
        
        ## open image as file and send to user
        save_path = pathlib.Path().cwd() / "outputs\discord"
        channel = client.get_channel(1027039583965290526) #garden_1 chanel
        with open(rf'{save_path}\{filename}.png', 'rb') as f:
            pic = discord.File(f)
            asyncio.run_coroutine_threadsafe(channel.send(f'{mention} "{prompt}", seed= {seed}', file=pic), loop)    

        #recursive loop    
        sd_gen(ctx, queues)
    else:
        blocking = False  

The next bit of code gets interesting. Now the bot has finished creating the image and ready to send it back to the user. However, to instruct the bot to send information to Discord users, we’ll need a coroutine. In a “normal” Discord bot, it means that our function needs to start with async def, and finishes with the keyword await. To deal with this, we’ll again use the asyncio library’s run_coroutine_threadsafe() method to run the coroutine channel.send() in the event loop which we defined earlier as the variable loop.

This concludes the workflow, but we are not done yet!

To make sure our bot would continue processing request after request until the queue is exhausted, we need some sort of loop. A recursive function becomes handy in this case, hence the sd_gen(ctx, queues) at the end.

Here’s a simple function to help you understand how a recursive function works. The function basically calls itself again and again until certain criteria is met and then the function will end.

def factorial(n):
    if (n == 0) or (n == 1):
        return 1
    else:
        return n*factorial(n-1)

factorial(3) # 6 = 3 * 2 * 1

In our case, the recursive sd_gen() function will end only if the queue is empty. At that time, the blocking variable will be set to False, and the bot is in standby mode awaiting users’ input.

Conclusion

Building this bot wasn’t easy as I’m quite new to asynchronous programming and Discord bots. However, I had a lot of fun during the process and learned a lot of async-related techniques. I’m sure I can use what I learned here for the stock trading bot that I plan to build in the future.

Feel free to join the Discord server and check out the bot: https://discord.gg/HDQjH39fsm

If you have any technical questions, please also leave your questions here, or on the Discord server.

Additional Resources

Building A Simple Python Discord Bot with DiscordPy in 2022/2023

How to Run Stable Diffusion on Windows