refactor: async I/O, input sanitisation, and shared utilities cleanup
This commit is contained in:
+23
-97
@@ -2,92 +2,33 @@
|
||||
"""
|
||||
Plugin for generating images using self-hosted Stable Diffusion and sending them to a Matrix chat room.
|
||||
|
||||
Now supports a `--seed` parameter to control deterministic generation.
|
||||
Now fully asynchronous (uses aiohttp). All original parameters and help text are preserved.
|
||||
"""
|
||||
|
||||
import requests
|
||||
import aiohttp
|
||||
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
|
||||
# Check if API is reachable
|
||||
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
|
||||
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:
|
||||
# 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')
|
||||
@@ -120,32 +61,26 @@ async def handle_command(room, message, bot, prefix, config):
|
||||
"width": args.w,
|
||||
"height": args.h,
|
||||
}
|
||||
# Add seed only if explicitly provided
|
||||
if args.seed is not None:
|
||||
payload["seed"] = args.seed
|
||||
|
||||
url = "http://127.0.0.1:7860/sdapi/v1/txt2img"
|
||||
response = requests.post(url=url, json=payload, timeout=600)
|
||||
r = response.json()
|
||||
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()
|
||||
|
||||
# Use secure temporary file
|
||||
# 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(base64.b64decode(r['images'][0]))
|
||||
temp_file.write(image_data)
|
||||
|
||||
# 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(" ", "")
|
||||
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>
|
||||
<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}{seed_info}<br>
|
||||
<strong>Negative Prompt:</strong> {neg_prompt_clean}</details>"""
|
||||
# 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
|
||||
@@ -156,18 +91,10 @@ async def handle_command(room, message, bot, prefix, config):
|
||||
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.
|
||||
Generates the full help text for the 'sd' command, including LORA list.
|
||||
"""
|
||||
return """
|
||||
<p>Generate images using self-hosted Stable Diffusion</p>
|
||||
@@ -205,14 +132,12 @@ def print_help():
|
||||
</ul>
|
||||
"""
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Plugin Metadata
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
__version__ = "1.1.0"
|
||||
__version__ = "1.1.2"
|
||||
__author__ = "Funguy Bot"
|
||||
__description__ = "Stable Diffusion image generation (supports --seed)"
|
||||
__description__ = "Stable Diffusion image generation (async, LORA support)"
|
||||
__help__ = """
|
||||
<details>
|
||||
<summary><strong>!sd</strong> – Generate images via Stable Diffusion</summary>
|
||||
@@ -225,6 +150,7 @@ __help__ = """
|
||||
<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>
|
||||
"""
|
||||
|
||||
Reference in New Issue
Block a user