diff --git a/plugins/stable-diffusion.py b/plugins/stable-diffusion.py index b428f93..32b51dd 100644 --- a/plugins/stable-diffusion.py +++ b/plugins/stable-diffusion.py @@ -8,6 +8,11 @@ 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() @@ -58,9 +63,18 @@ async def handle_command(room, message, bot, prefix, config): response = requests.post(url=url, json=payload) r = response.json() - with open("/tmp/output.jpg", 'wb') as f: + + # Slugify the prompt + prompt_slug = (prompt) + + # 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="/tmp/output.jpg") # Corrected argument name + 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}
""" + await bot.api.send_markdown_message(room.room_id, info_msg) except argparse.ArgumentError as e: await bot.api.send_text_message(room.room_id, f"Error: {e}") await bot.api.send_markdown_message(room.room_id, "
Stable Diffusion Help" + print_help() + "
") @@ -71,6 +85,7 @@ async def handle_command(room, message, bot, prefix, config): next_command = await command_queue.get() await handle_command(*next_command) + def print_help(): return """

Generate images using self-hosted Stable Diffusion

diff --git a/requirements.txt b/requirements.txt index df710ae..65e94f7 100644 --- a/requirements.txt +++ b/requirements.txt @@ -6,3 +6,4 @@ nio markdown2 watchdog emoji +python-slugify