"""
This plugin provides a command to generate images using self hosted Stable Diffusion and send to the room
"""
# plugins/stable-diffusion.py
import requests
import base64
from asyncio import Queue
import argparse
import simplematrixbotlib as botlib
import markdown2
from slugify import slugify
def slugify_prompt(prompt):
# Generate a slug from the prompt
return slugify(prompt)
# Queue to store pending commands
command_queue = Queue()
def markdown_to_html(markdown_text):
html_content = markdown2.markdown(markdown_text)
return html_content
async def process_command(room, message, bot, prefix, config):
match = botlib.MessageMatch(room, message, bot, prefix)
if match.prefix() and match.command("sd"):
if command_queue.empty():
await handle_command(room, message, bot, prefix, config)
else:
await command_queue.put((room, message, bot, prefix, config))
async def handle_command(room, message, bot, prefix, config):
match = botlib.MessageMatch(room, message, bot, prefix)
if match.prefix() and match.command("sd"):
try:
parser = argparse.ArgumentParser(description='Generate images using self hosted Stable Diffusion')
parser.add_argument('--steps', type=int, default=4, help='Number of steps, default=16')
parser.add_argument('--cfg', type=int, default=1.25, help='CFG scale, default=7')
parser.add_argument('--h', type=int, default=512, help='Height of the image, default=512')
parser.add_argument('--w', type=int, default=512, help='Width of the image, default=512')
parser.add_argument('--neg', type=str, default='((((ugly)))), (((duplicate))), ((morbid)), ((mutilated)), out of frame, extra fingers, mutated hands, ((poorly drawn hands)), ((poorly drawn face)), (((mutation))), (((deformed))), ((ugly)), blurry, ((bad anatomy)), (((bad proportions))), ((extra limbs)), cloned face, (((disfigured))), out of frame, ugly, extra limbs, (bad anatomy), gross proportions, (malformed limbs), ((missing arms)), ((missing legs)), (((extra arms))), (((extra legs))), mutated hands, (fused fingers), (too many fingers), (((long neck)))', nargs='+', help='Negative prompt, default=none')
parser.add_argument('--sampler', type=str, nargs='*', default=['DPM++', 'SDE', 'Karras'], help='Sampler name, default=Euler a')
parser.add_argument('prompt', type=str, nargs='*', help='Prompt for the image')
args = parser.parse_args(message.body.split()[1:]) # Skip the command itself
if not args.prompt:
raise argparse.ArgumentError(None, "Prompt is required.")
prompt = ' '.join(args.prompt)
sampler_name = ' '.join(args.sampler)
neg = ' '.join(args.neg)
payload = {
"prompt": prompt,
"steps": args.steps,
"negative_prompt": neg,
"sampler_name": sampler_name,
"cfg_scale": args.cfg,
"width": args.w,
"height": args.h,
}
url = "http://127.0.0.1:7860/sdapi/v1/txt2img"
response = requests.post(url=url, json=payload)
r = response.json()
# Slugify the prompt
prompt_slug = prompt[:120]
# Construct filename
filename = f"{prompt_slug}_{args.steps}_{args.h}x{args.w}_{sampler_name}_{args.cfg}.jpg"
with open(f"/tmp/{filename}", 'wb') as f:
f.write(base64.b64decode(r['images'][0]))
await bot.api.send_image_message(room_id=room.room_id, image_filepath=f"/tmp/{filename}") # Corrected argument name
neg = neg.replace(" ", "")
info_msg = f"""🍄⤵︎Image Info:⤵︎
Prompt: {prompt_slug}
Steps: {args.steps}
Dimensions: {args.h}x{args.w}
Sampler: {sampler_name}
CFG Scale: {args.cfg}
Negative Prompt: {neg}Stable Diffusion Help
" + print_help() + "
Generate images using self-hosted Stable Diffusion
Positional arguments:
Optional arguments:
LORA List:
Load LORAs like this: