FunguyBot/plugins/stable-diffusion.py

"""
Plugin for generating images using self-hosted Stable Diffusion and sending them to a Matrix chat room.
"""

import requests
import base64
from asyncio import Queue
import argparse
import simplematrixbotlib as botlib
import markdown2
from slugify import slugify

def slugify_prompt(prompt):
    """
    Generates a URL-friendly slug from the given prompt.

    Args:
        prompt (str): The prompt to slugify.

    Returns:
        str: A URL-friendly slug version of the prompt.
    """
    return slugify(prompt)

# Queue to store pending commands
command_queue = Queue()

def markdown_to_html(markdown_text):
    """
    Converts Markdown text to HTML.

    Args:
        markdown_text (str): The Markdown text to convert.

    Returns:
        str: The HTML version of the input Markdown text.
    """
    html_content = markdown2.markdown(markdown_text)
    return html_content

async def process_command(room, message, bot, prefix, config):
    """
    Asynchronously processes the commands received in the Matrix room.

    Args:
        room (Room): The Matrix room where the command was received.
        message (Message): The message object containing the command.
        bot (MatrixBot): The Matrix bot instance.
        prefix (str): The command prefix.
        config (dict): The bot's configuration.

    """
    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):
    """
    Asynchronously handles the 'sd' command, which generates images using Stable Diffusion.

    Args:
        room (Room): The Matrix room where the command was received.
        message (Message): The message object containing the command.
        bot (MatrixBot): The Matrix bot instance.
        prefix (str): The command prefix.
        config (dict): The bot's configuration.
    """
    match = botlib.MessageMatch(room, message, bot, prefix)
    if match.prefix() and match.command("sd"):
        try:
            # Parse command-line arguments
            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"

            # Send request to the Stable Diffusion API
            response = requests.post(url=url, json=payload)
            r = response.json()

            # Slugify the prompt
            prompt_slug = slugify(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]))

            # Send the generated image to the Matrix room
            await bot.api.send_image_message(room_id=room.room_id, image_filepath=f"/tmp/{filename}")

            # Format and send information about the generated image
            neg = neg.replace(" ", "")
            info_msg = f"""<details><summary>🍄⤵︎Image Info:⤵︎</summary><strong>Prompt:</strong> {prompt_slug}<br><strong>Steps:</strong> {args.steps}<br><strong>Dimensions:</strong> {args.h}x{args.w}<br><strong>Sampler:</strong> {sampler_name}<br><strong>CFG Scale:</strong> {args.cfg}<br><strong>Negative Prompt: {neg}</strong></details>"""
            await bot.api.send_markdown_message(room.room_id, info_msg)
        except argparse.ArgumentError as e:
            # Handle argument errors
            await bot.api.send_text_message(room.room_id, f"Error: {e}")
            await bot.api.send_markdown_message(room.room_id, "<details><summary>Stable Diffusion Help</summary>" + print_help() + "</details>")
        except Exception as e:
            # Handle general errors
            await bot.api.send_text_message(room.room_id, f"Error processing the command: {str(e)}")
        finally:
            # Process next command from the queue, if any
            if not command_queue.empty():
                next_command = await command_queue.get()
                await handle_command(*next_command)


def print_help():
    """
    Generates help text for the 'sd' command.

    Returns:
        str: Help text for the 'sd' command.
    """
    return """
  <p>Generate images using self-hosted Stable Diffusion</p>

  <p>Positional arguments:</p>
  <ul>
    <li>prompt - Prompt for the image</li>
  </ul>

  <p>Default Negative Prompts: ((((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)))</p>

  <p>Optional arguments:</p>
  <ul>
    <li>--steps STEPS - Number of steps, default=16</li>
    <li>--cfg CFG - CFG scale, default=7</li>
    <li>--h H - Height of the image, default=512</li>
    <li>--w W - Width of the image, default=512</li>
    <li>--neg NEG - Negative prompt, default=none</li>
    <li>--sampler SAMPLER - Sampler name, default=Euler a</li>
  </ul>

  <p>LORA List:</p>
  <ul>
    <li>&lt;psychedelicai-SDXL&gt;</li>
    <li>&lt;ral-frctlgmtry-sdxl&gt;</li>
    <li>&lt;SDXL-PsyAI-v4&gt;</li>
    <li>&lt;al3xxl&gt;</li>
  </ul>

  <p>Load LORAs like this:</p>
  <ul>
    <li>&lt;lora:SDXL-PsyAI-v4:1&gt; PsyAI</li>
    <li>&lt;lora:psychedelicai-SDXL:1&gt; Psychedelic</li>
    <li>&lt;ral-frctlgmtry-sdxl&gt; ral-frctlgmtry</li>
    <li>&lt;lora:al3xxl:1&gt; alexpainting, alexhuman, alexentity, alexthirdeye, alexforeheads, alexgalactic, spiraling, alexmirror, alexangel, alexkissing, alexthirdeye2, alexthirdeye3, alexhofmann, wearing, glasses, alexgalactic3, alexthirdeye4, alexhuman3, alexgalactic2, alexhuman2, alexbeing2, alexfractal2, alexfractal3</li>
  </ul>
"""