Files
FunguyBot/plugins/stable-diffusion.py

190 lines
7.9 KiB
Python

"""
Plugin for generating images using self-hosted Stable Diffusion and sending them to a Matrix chat room.
"""
import requests
import base64
import tempfile
import os
from asyncio import Queue
import argparse
import simplematrixbotlib as botlib
import markdown2
from slugify import slugify
# Queue to store pending commands
command_queue = Queue()
def slugify_prompt(prompt: str) -> str:
"""
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)
def markdown_to_html(markdown_text: str) -> str:
"""
Converts Markdown text to HTML.
Args:
markdown_text (str): The Markdown text to convert.
Returns:
str: The HTML version of the input Markdown text.
"""
return markdown2.markdown(markdown_text)
async def process_command(room, message, bot, prefix, config):
"""
Processes !sd commands and queues them if already running.
Args:
room: Matrix room object
message: Matrix message object
bot: Bot instance
prefix: Command prefix
config: Bot config object
"""
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))
await bot.api.send_text_message(room.room_id, "Command queued. Please wait for the current image to finish.")
async def handle_command(room, message, bot, prefix, config):
"""
Handles !sd command: generates image using Stable Diffusion API.
Args:
room: Matrix room object
message: Matrix message object
bot: Bot instance
prefix: Command prefix
config: Bot config object
"""
match = botlib.MessageMatch(room, message, bot, prefix)
if not (match.prefix() and match.command("sd")):
return
# Check if API is available
try:
health_check = requests.get("http://127.0.0.1:7860/docs", timeout=3)
if health_check.status_code != 200:
await bot.api.send_text_message(room.room_id, "Stable Diffusion API is not running!")
return
except Exception:
await bot.api.send_text_message(room.room_id, "Could not reach Stable Diffusion API!")
return
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=4')
parser.add_argument('--cfg', type=int, default=2, help='CFG scale, default=2')
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, nargs='+', 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)))'], help='Negative prompt')
parser.add_argument('--sampler', type=str, nargs='*', default=['DPM++', 'SDE'], help='Sampler name, default=DPM++ SDE')
parser.add_argument('prompt', type=str, nargs='*', help='Prompt for the image')
args = parser.parse_args(message.body.split()[1:]) # skip command prefix
if not args.prompt:
raise argparse.ArgumentError(None, "Prompt is required.")
prompt = ' '.join(args.prompt)
sampler_name = ' '.join(args.sampler)
neg_prompt = ' '.join(args.neg)
payload = {
"prompt": prompt,
"steps": args.steps,
"negative_prompt": neg_prompt,
"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, timeout=600)
r = response.json()
# Use secure temporary file
with tempfile.NamedTemporaryFile(suffix='.jpg', delete=False) as temp_file:
filename = temp_file.name
temp_file.write(base64.b64decode(r['images'][0]))
# Send image to Matrix room
await bot.api.send_image_message(room_id=room.room_id, image_filepath=filename)
# Optional: send info about generated image
neg_prompt_clean = neg_prompt.replace(" ", "")
info_msg = f"""<details><summary>🔍 Image Info</summary><strong>Prompt:</strong> {prompt[:100]}<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:</strong> {neg_prompt_clean}</details>"""
# await bot.api.send_markdown_message(room.room_id, info_msg)
# Clean up temp file
os.remove(filename)
except argparse.ArgumentError as e:
await bot.api.send_text_message(room.room_id, f"Argument Error: {e}")
await bot.api.send_markdown_message(room.room_id, "<details><summary>Stable Diffusion Help</summary>" + print_help() + "</details>")
except Exception as e:
await bot.api.send_text_message(room.room_id, f"Error processing the command: {str(e)}")
finally:
# Process next queued command
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>
"""