2024-03-01 03:18:58 +00:00
|
|
|
|
"""
|
|
|
|
|
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
|
2024-03-03 13:59:39 +00:00
|
|
|
|
import argparse
|
2024-03-01 03:18:58 +00:00
|
|
|
|
import simplematrixbotlib as botlib
|
2024-03-03 13:59:39 +00:00
|
|
|
|
import markdown2
|
2024-03-03 20:17:01 +00:00
|
|
|
|
from slugify import slugify
|
|
|
|
|
|
|
|
|
|
def slugify_prompt(prompt):
|
|
|
|
|
# Generate a slug from the prompt
|
|
|
|
|
return slugify(prompt)
|
2024-03-01 03:18:58 +00:00
|
|
|
|
|
|
|
|
|
# Queue to store pending commands
|
|
|
|
|
command_queue = Queue()
|
|
|
|
|
|
2024-03-03 13:59:39 +00:00
|
|
|
|
def markdown_to_html(markdown_text):
|
|
|
|
|
html_content = markdown2.markdown(markdown_text)
|
|
|
|
|
return html_content
|
|
|
|
|
|
2024-03-01 03:18:58 +00:00
|
|
|
|
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:
|
2024-03-03 13:59:39 +00:00
|
|
|
|
parser = argparse.ArgumentParser(description='Generate images using self hosted Stable Diffusion')
|
2024-03-03 19:56:19 +00:00
|
|
|
|
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')
|
2024-03-03 13:59:39 +00:00
|
|
|
|
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')
|
2024-03-03 19:56:19 +00:00
|
|
|
|
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')
|
2024-03-03 13:59:39 +00:00
|
|
|
|
|
|
|
|
|
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"
|
|
|
|
|
|
2024-03-01 03:18:58 +00:00
|
|
|
|
response = requests.post(url=url, json=payload)
|
|
|
|
|
r = response.json()
|
2024-03-03 20:17:01 +00:00
|
|
|
|
|
|
|
|
|
# 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:
|
2024-03-01 03:18:58 +00:00
|
|
|
|
f.write(base64.b64decode(r['images'][0]))
|
2024-03-03 20:17:01 +00:00
|
|
|
|
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"""<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)
|
2024-03-03 13:59:39 +00:00
|
|
|
|
except argparse.ArgumentError as e:
|
|
|
|
|
await bot.api.send_text_message(room.room_id, f"Error: {e}")
|
2024-03-03 19:56:19 +00:00
|
|
|
|
await bot.api.send_markdown_message(room.room_id, "<details><summary>Stable Diffusion Help</summary>" + print_help() + "</details>")
|
2024-03-01 03:18:58 +00:00
|
|
|
|
except Exception as e:
|
|
|
|
|
await bot.api.send_text_message(room.room_id, f"Error processing the command: {str(e)}")
|
|
|
|
|
finally:
|
|
|
|
|
if not command_queue.empty():
|
|
|
|
|
next_command = await command_queue.get()
|
|
|
|
|
await handle_command(*next_command)
|
2024-03-03 19:56:19 +00:00
|
|
|
|
|
2024-03-03 20:17:01 +00:00
|
|
|
|
|
2024-03-03 19:56:19 +00:00
|
|
|
|
def print_help():
|
|
|
|
|
return """
|
|
|
|
|
<p>Generate images using self-hosted Stable Diffusion</p>
|
|
|
|
|
|
|
|
|
|
<p>Positional arguments:</p>
|
|
|
|
|
<ul>
|
|
|
|
|
<li>prompt - Prompt for the image</li>
|
|
|
|
|
</ul>
|
|
|
|
|
|
|
|
|
|
<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><psychedelicai-SDXL></li>
|
|
|
|
|
<li><ral-frctlgmtry-sdxl></li>
|
|
|
|
|
<li><SDXL-PsyAI-v4></li>
|
|
|
|
|
<li><al3xxl></li>
|
|
|
|
|
</ul>
|
|
|
|
|
|
|
|
|
|
<p>Load LORAs like this:</p>
|
|
|
|
|
<ul>
|
|
|
|
|
<li><lora:SDXL-PsyAI-v4:1> PsyAI</li>
|
|
|
|
|
<li><lora:psychedelicai-SDXL:1> Psychedelic</li>
|
|
|
|
|
<li><ral-frctlgmtry-sdxl> ral-frctlgmtry</li>
|
|
|
|
|
<li><lora:al3xxl:1> 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>
|
|
|
|
|
"""
|