157 lines
7.5 KiB
Python
157 lines
7.5 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
Plugin for generating images using self-hosted Stable Diffusion and sending them to a Matrix chat room.
|
||
|
||
Now fully asynchronous (uses aiohttp). All original parameters and help text are preserved.
|
||
"""
|
||
|
||
import aiohttp
|
||
import base64
|
||
import tempfile
|
||
import os
|
||
import argparse
|
||
import simplematrixbotlib as botlib
|
||
|
||
async def handle_command(room, message, bot, prefix, config):
|
||
match = botlib.MessageMatch(room, message, bot, prefix)
|
||
if not (match.prefix() and match.command("sd")):
|
||
return
|
||
|
||
# Check if API is reachable
|
||
try:
|
||
async with aiohttp.ClientSession() as session:
|
||
async with session.get("http://127.0.0.1:7860/docs", timeout=3) as resp:
|
||
if resp.status != 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:
|
||
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 Karras'],
|
||
help='Sampler name, default=DPM++ SDE')
|
||
parser.add_argument('--seed', type=int, default=None,
|
||
help='Seed for deterministic generation (omit for random)')
|
||
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,
|
||
}
|
||
if args.seed is not None:
|
||
payload["seed"] = args.seed
|
||
|
||
async with aiohttp.ClientSession() as session:
|
||
async with session.post("http://127.0.0.1:7860/sdapi/v1/txt2img", json=payload, timeout=600) as response:
|
||
response.raise_for_status()
|
||
r = await response.json()
|
||
|
||
# Save and send image
|
||
image_data = base64.b64decode(r['images'][0])
|
||
with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as temp_file:
|
||
filename = temp_file.name
|
||
temp_file.write(image_data)
|
||
|
||
await bot.api.send_image_message(room_id=room.room_id, image_filepath=filename)
|
||
|
||
# Optional info message (commented out to avoid spam, but can be enabled)
|
||
# neg_prompt_clean = neg_prompt.replace(" ", "")
|
||
# seed_info = f"<br><strong>Seed:</strong> {args.seed}" if args.seed is not None else ""
|
||
# info_msg = f"<details><summary>🔍 Image Info</summary><strong>Prompt:</strong> {prompt[:100]}<br>...</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)}")
|
||
|
||
def print_help():
|
||
"""
|
||
Generates the full help text for the 'sd' command, including LORA list.
|
||
"""
|
||
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=4</li>
|
||
<li>--cfg CFG - CFG scale, default=2</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=((((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)))</li>
|
||
<li>--sampler SAMPLER - Sampler name, default=DPM++ SDE</li>
|
||
<li>--seed SEED - Seed for deterministic generation (omit for random)</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>
|
||
"""
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Plugin Metadata
|
||
# ---------------------------------------------------------------------------
|
||
__version__ = "1.1.2"
|
||
__author__ = "Funguy Bot"
|
||
__description__ = "Stable Diffusion image generation (async, LORA support)"
|
||
__help__ = """
|
||
<details>
|
||
<summary><strong>!sd</strong> – Generate images via Stable Diffusion</summary>
|
||
<p><code>!sd [options] <prompt></code></p>
|
||
<ul>
|
||
<li><code>--steps N</code> – Sampling steps (default 4)</li>
|
||
<li><code>--cfg scale</code> – CFG scale (default 2)</li>
|
||
<li><code>--h H --w W</code> – Image dimensions (default 512)</li>
|
||
<li><code>--neg <negative prompt></code></li>
|
||
<li><code>--sampler SAMPLER</code> – Sampler name (default DPM++ SDE)</li>
|
||
<li><code>--seed SEED</code> – Deterministic seed (optional)</li>
|
||
</ul>
|
||
<p>LORAs: <code><lora:filename:weight></code></p>
|
||
<p>Requires a locally running Stable Diffusion API.</p>
|
||
</details>
|
||
"""
|