Generating Images with a Discord Bot using StableDiffusion
In this blog post, we'll discuss how to create a Discord bot that generates images based on user input Stable Diffusion. Users provide positive and negative keywords to guide the image generation process.
Requirements
To follow along, you'll need:
- Python 3.6 or higher
- discord.py library
- Pillow (PIL) library
- aiohttp library
- A Discord bot token
- StableDiffusion API Access
Setting Up the Bot
First, import the necessary libraries and create a bot instance with the appropriate intents: ```python import io import os import base64 import discord import asyncio import aiohttp from PIL import Image, PngImagePlugin from discord.ext import commands import requests
url = "http://your_api_url:port" intents = discord.Intents.all() intents.members = True bot = commands.Bot(command_prefix='/', intents=intents) bot.conversations = [] ``` Replace your_api_url:port with the actual URL and port of the StableDiffusion for image generation.
Handling Image Generation
Create an asynchronous function generate_image that takes user input and generates an image using the StableDiffusion Powered API: ```python async def generate_image(ctx, positive_prompt, negative_prompt): working_message = await ctx.send("Generating image...Please wait.. Your Image may be in que") payload = { "prompt": positive_prompt, "steps": 35, "enable_hr": False, "denoising_strength": 1, "firstphase_width": 0, "firstphase_height": 0, "hr_scale": 5, "hr_second_pass_steps": 0, "hr_resize_x": 0, "hr_resize_y": 0, "seed": -1, "subseed": -1, "subseed_strength": 0, "seed_resize_from_h": -1, "seed_resize_from_w": -1, "sampler_name": "PLMS", "batch_size": 1, "n_iter": 1, "cfg_scale": 5, "width": 512, "height": 512, "restore_faces": True, "tiling": False, "negative_prompt": negative_prompt, "eta": 0, "s_churn": 0, "s_tmax": 0, "s_tmin": 0, "s_noise": 1, "override_settings": {}, "override_settings_restore_afterwards": True, "script_args": [], "sampler_index": "PLMS", }
try:
response = requests.post(url=f'{url}/sdapi/v1/txt2img', json=payload)
response.raise_for_status()
r = response.json()
except requests.exceptions.RequestException as e:
print(f'Request Error: {e}')
return
for idx, i in enumerate(r['images']):
try:
image = Image.open(io.BytesIO(base64.b64decode(i.split(",",1)[0])))
except Exception as e:
print(f'Image Error: {e}')
continue
png_payload = {
"image": "data:image/png;base64," + i
}
try:
response2 = requests.post(url=f'{url}/sdapi/v1/png-info', json=png_payload)
response2.raise_for_status()
pnginfo = PngImagePlugin.PngInfo()
pnginfo.add_text("parameters", response2.json().get("info"))
image.save(f'output_{idx}.png', pnginfo=pnginfo)
except requests.exceptions.RequestException as e:
print(f'Request Error: {e}')
continue
# Inform the user that their image generation is queued or waiting
if os.path.isfile(f'output_{idx}.png'):
with open(f'output_{idx}.png', 'rb') as img:
await ctx.send(file=discord.File(img, f'output_{idx}.png'))
# Remove the working indicator after the image is generated
await working_message.delete()
os.remove(f'output_{idx}.png')
break
async def is_api_available(): try: response = requests.get(url) response.raise_for_status() return True except requests.exceptions.RequestException as e: print(f'Request Error: {e}') return False ``` Inside this function, construct the payload with the positive and negative prompts, make a POST request to the API, and process the response. Save the resulting image and send it back to the user. After sending the image, delete the local file and any working messages.
Creating the generate Command
Create a generate command for the bot that starts a conversation with the user, collects positive and negative keywords, and then calls the generate_image function:
```python @bot.command(name='generate', help='Generates an image based on your input') async def generate(ctx): if not await is_api_available(): await ctx.send("Dan's StableDiffusion API is currently offline or unavailable. He's Probably playing games. Please try again later.") return if ctx.author.id not in (c['user_id'] for c in bot.conversations): bot.conversations.append({'user_id': ctx.author.id, 'in_progress': True}) positive_prompt_message = await ctx.send("Please enter the positive keywords separated by commas:") positive_prompt = await bot.wait_for('message', check=lambda msg: msg.author == ctx.author and msg.channel == ctx.channel) negative_prompt_message = await ctx.send("Please enter the negative keywords separated by commas (leave blank if none):") negative_prompt = await bot.wait_for('message', check=lambda msg: msg.author == ctx.author and msg.channel == ctx.channel) # Store the messages to be deleted in a list messages_to_delete = [positive_prompt, negative_prompt, ctx.message, positive_prompt_message, negative_prompt_message]
asyncio.create_task(generate_image(ctx, positive_prompt.content, negative_prompt.content))
bot.conversations = list(filter(lambda c: c['user_id'] != ctx.author.id, bot.conversations))
# Delete the messages after the image is generated and sent
for message in messages_to_delete:
try:
await message.delete()
except discord.errors.Forbidden:
# Send a message if the bot doesn't have permission to delete messages
await ctx.send("I do not have permission to delete image generation messages. Please grant me the necessary permissions or manually delete the messages.")
break
else:
await ctx.send("You can only request one picture at a time.")
``` Within this command, check if the API is available, request positive and negative keywords from the user, and store the messages to be deleted later. Once the image is generated, delete the stored messages.
Running the Bot
Finally, add the necessary event handlers and run the bot using the bot token:
```python @bot.event async def on_ready(): print(bot.user.name, 'is ready for battle!')
@bot.event async def on_message(message): if message.author.bot: return await bot.process_commands(message) # Required when using command decorators
if name == 'main': token = 'your_bot_token' bot.run(token) ```
Replace your_bot_token with your actual Discord bot token.
Now you have a fully functional Discord bot that can generate images based on user input using StableDiffusion. Users can simply invoke the /generate command and provide positive and negative keywords to generate a unique image.
This bot showcases the potential of combining the power of AI with the versatility of Discord, opening the door to countless creative applications.
Full Code below:
```python import io import os import base64 import discord import asyncio import aiohttp from PIL import Image, PngImagePlugin from discord.ext import commands import requests
url = "http://your_api_url:port" intents = discord.Intents.all() intents.members = True bot = commands.Bot(command_prefix='/', intents=intents) bot.conversations = [] async def generate_image(ctx, positive_prompt, negative_prompt): working_message = await ctx.send("Generating image...Please wait.. Your Image may be in que") payload = { "prompt": positive_prompt, "steps": 35, "enable_hr": False, "denoising_strength": 1, "firstphase_width": 0, "firstphase_height": 0, "hr_scale": 5, "hr_second_pass_steps": 0, "hr_resize_x": 0, "hr_resize_y": 0, "seed": -1, "subseed": -1, "subseed_strength": 0, "seed_resize_from_h": -1, "seed_resize_from_w": -1, "sampler_name": "PLMS", "batch_size": 1, "n_iter": 1, "cfg_scale": 5, "width": 512, "height": 512, "restore_faces": True, "tiling": False, "negative_prompt": negative_prompt, "eta": 0, "s_churn": 0, "s_tmax": 0, "s_tmin": 0, "s_noise": 1, "override_settings": {}, "override_settings_restore_afterwards": True, "script_args": [], "sampler_index": "PLMS", }
try:
response = requests.post(url=f'{url}/sdapi/v1/txt2img', json=payload)
response.raise_for_status()
r = response.json()
except requests.exceptions.RequestException as e:
print(f'Request Error: {e}')
return
for idx, i in enumerate(r['images']):
try:
image = Image.open(io.BytesIO(base64.b64decode(i.split(",",1)[0])))
except Exception as e:
print(f'Image Error: {e}')
continue
png_payload = {
"image": "data:image/png;base64," + i
}
try:
response2 = requests.post(url=f'{url}/sdapi/v1/png-info', json=png_payload)
response2.raise_for_status()
pnginfo = PngImagePlugin.PngInfo()
pnginfo.add_text("parameters", response2.json().get("info"))
image.save(f'output_{idx}.png', pnginfo=pnginfo)
except requests.exceptions.RequestException as e:
print(f'Request Error: {e}')
continue
# Inform the user that their image generation is queued or waiting
if os.path.isfile(f'output_{idx}.png'):
with open(f'output_{idx}.png', 'rb') as img:
await ctx.send(file=discord.File(img, f'output_{idx}.png'))
# Remove the working indicator after the image is generated
await working_message.delete()
os.remove(f'output_{idx}.png')
break
async def is_api_available(): try: response = requests.get(url) response.raise_for_status() return True except requests.exceptions.RequestException as e: print(f'Request Error: {e}') return False
@bot.event async def on_ready(): print(bot.user.name, 'is ready for battle!')
@bot.event async def on_message(message): if message.author.bot: return await bot.process_commands(message) # Required when using command decorators
@bot.command(name='generate', help='Generates an image based on your input') async def generate(ctx): if not await is_api_available(): await ctx.send("Dan's StableDiffusion API is currently offline or unavailable. He's Probably playing games. Please try again later.") return if ctx.author.id not in (c['user_id'] for c in bot.conversations): bot.conversations.append({'user_id': ctx.author.id, 'in_progress': True}) positive_prompt_message = await ctx.send("Please enter the positive keywords separated by commas:") positive_prompt = await bot.wait_for('message', check=lambda msg: msg.author == ctx.author and msg.channel == ctx.channel) negative_prompt_message = await ctx.send("Please enter the negative keywords separated by commas (leave blank if none):") negative_prompt = await bot.wait_for('message', check=lambda msg: msg.author == ctx.author and msg.channel == ctx.channel) # Store the messages to be deleted in a list messages_to_delete = [positive_prompt, negative_prompt, ctx.message, positive_prompt_message, negative_prompt_message]
asyncio.create_task(generate_image(ctx, positive_prompt.content, negative_prompt.content))
bot.conversations = list(filter(lambda c: c['user_id'] != ctx.author.id, bot.conversations))
# Delete the messages after the image is generated and sent
for message in messages_to_delete:
try:
await message.delete()
except discord.errors.Forbidden:
# Send a message if the bot doesn't have permission to delete messages
await ctx.send("I do not have permission to delete image generation messages. Please grant me the necessary permissions or manually delete the messages.")
break
else:
await ctx.send("You can only request one picture at a time.")
if name == 'main': token = 'Your_Discord_token' bot.run(token) ```