Building an image generation bot with NOS + Discord

1. Registering a discord bot to create our API key

   This can be done through the discord developer guide. You will need a discord account as well as a server you wish to add the bot to. When we're finished, nos-bot will accept image generation requests from users on this server and post the resulting images to the main channel. Once you have your bot API key, add it to your local environment like so:

   export BOT_TOKEN=$YOUR_DISCORD_API_TOKEN
   

2. Setting up a NOS client to generate images

   NOS comes with an endpoint for Stable Diffusion V2 from HuggingFace, so all we need to do is init a server on our machine and verify we can connect to it from the client:

   import nos
   from nos.client import Client, TaskType
   import os
   
   # Init nos server
   nos.init(runtime='gpu')
   client = Client()
   client.WaitForServer()
   assert client.IsHealthy()
   

   See examples/notebook/inference-client-example.ipynb for a better overview. Nos will initialize a GPU-ready container on our machine and return to the client that it's ready to go.

3. The Discord interface

   Next we need a way to handle message requests on our server. Discord.py makes callbacks pretty easy:

   intents = discord.Intents.default()
   intents.message_content = True
   
   bot = commands.Bot(command_prefix='$', intents=intents)
   
   @bot.command()
   async def generate(ctx, *, prompt):
       response = client.Run(
           "stabilityai/stable-diffusion-2",
           inputs={
               "prompts": [prompt],
               "width": 512,
               "height": 512,
               "num_images": 1
           }
       )
       image, = response["images"]
   
       tmp_file_path = "image.png"
       image.save(tmp_file_path)
       with open(tmp_file_path, "rb") as img_file:
           await ctx.send(file=discord.File(img_file))
   
       os.remove(tmp_file_path)
   

   We need the message_content intent so we can access the contents of user messages to retrieve image prompts. We'll parse generation requests as $generate image prompt here.... Any messages beginning with $generate will be sent to Nos for image generation. The rest of our message handler is pretty straightforward: we use the client to generate a set of images, then we retrieve the first result from the list, save it locally, and call ctx.send to upload the image with the discord File interface.

4. Time to run the server

   The full server is only ~40 LOC:

   #!/usr/bin/env python
   import io
   import os
   
   import discord
   from discord.ext import commands
   
   import nos
   from nos.client import Client
   
   # Init nos server
   nos.init(runtime='gpu')
   client = Client()
   client.WaitForServer()
   assert client.IsHealthy()
   
   intents = discord.Intents.default()
   intents.message_content = True
   
   bot = commands.Bot(command_prefix='$', intents=intents)
   
   @bot.command()
   async def generate(ctx, *, prompt):
       response = client.Run(
           "stabilityai/stable-diffusion-2",
           inputs={
               "prompts": [prompt],
               "width": 512,
               "height": 512,
               "num_images": 1
           }
       )
       image, = response["images"]
   
       image_bytes = io.BytesIO()
       img.save(image_bytes, format="PNG")
       image_bytes.seek(0)
       await ctx.send(file=discord.File(image_bytes, filename=f"{ctx.message.id}.png"))
   
   bot_token = os.environ.get("BOT_TOKEN")
   if bot_token is None:
       raise Exception("BOT_TOKEN environment variable not set")
   
   bot.run(bot_token)
   

   We should be all set, kick off the bot with python examples/discord/app/bot.py