Reduced some vuln after SAST scan
This commit is contained in:
@@ -4,13 +4,18 @@ Plugin for generating images using self-hosted Stable Diffusion and sending them
|
||||
|
||||
import requests
|
||||
import base64
|
||||
import tempfile
|
||||
import os
|
||||
from asyncio import Queue
|
||||
import argparse
|
||||
import simplematrixbotlib as botlib
|
||||
import markdown2
|
||||
from slugify import slugify
|
||||
|
||||
def slugify_prompt(prompt):
|
||||
# Queue to store pending commands
|
||||
command_queue = Queue()
|
||||
|
||||
def slugify_prompt(prompt: str) -> str:
|
||||
"""
|
||||
Generates a URL-friendly slug from the given prompt.
|
||||
|
||||
@@ -22,10 +27,7 @@ def slugify_prompt(prompt):
|
||||
"""
|
||||
return slugify(prompt)
|
||||
|
||||
# Queue to store pending commands
|
||||
command_queue = Queue()
|
||||
|
||||
def markdown_to_html(markdown_text):
|
||||
def markdown_to_html(markdown_text: str) -> str:
|
||||
"""
|
||||
Converts Markdown text to HTML.
|
||||
|
||||
@@ -35,20 +37,18 @@ def markdown_to_html(markdown_text):
|
||||
Returns:
|
||||
str: The HTML version of the input Markdown text.
|
||||
"""
|
||||
html_content = markdown2.markdown(markdown_text)
|
||||
return html_content
|
||||
return markdown2.markdown(markdown_text)
|
||||
|
||||
async def process_command(room, message, bot, prefix, config):
|
||||
"""
|
||||
Asynchronously processes the commands received in the Matrix room.
|
||||
Processes !sd commands and queues them if already running.
|
||||
|
||||
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.
|
||||
|
||||
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"):
|
||||
@@ -56,82 +56,93 @@ async def process_command(room, message, bot, prefix, config):
|
||||
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):
|
||||
"""
|
||||
Asynchronously handles the 'sd' command, which generates images using Stable Diffusion.
|
||||
Handles !sd command: generates image using Stable Diffusion API.
|
||||
|
||||
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.
|
||||
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"):
|
||||
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')
|
||||
if not (match.prefix() and match.command("sd")):
|
||||
return
|
||||
|
||||
args = parser.parse_args(message.body.split()[1:]) # Skip the command itself
|
||||
# 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
|
||||
|
||||
if not args.prompt:
|
||||
raise argparse.ArgumentError(None, "Prompt is required.")
|
||||
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')
|
||||
|
||||
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"
|
||||
args = parser.parse_args(message.body.split()[1:]) # skip command prefix
|
||||
|
||||
# Send request to the Stable Diffusion API
|
||||
response = requests.post(url=url, json=payload)
|
||||
r = response.json()
|
||||
if not args.prompt:
|
||||
raise argparse.ArgumentError(None, "Prompt is required.")
|
||||
|
||||
# Slugify the prompt
|
||||
prompt_slug = slugify(prompt[:120])
|
||||
prompt = ' '.join(args.prompt)
|
||||
sampler_name = ' '.join(args.sampler)
|
||||
neg_prompt = ' '.join(args.neg)
|
||||
|
||||
# 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]))
|
||||
payload = {
|
||||
"prompt": prompt,
|
||||
"steps": args.steps,
|
||||
"negative_prompt": neg_prompt,
|
||||
"sampler_name": sampler_name,
|
||||
"cfg_scale": args.cfg,
|
||||
"width": args.w,
|
||||
"height": args.h,
|
||||
}
|
||||
|
||||
# Send the generated image to the Matrix room
|
||||
await bot.api.send_image_message(room_id=room.room_id, image_filepath=f"/tmp/{filename}")
|
||||
url = "http://127.0.0.1:7860/sdapi/v1/txt2img"
|
||||
response = requests.post(url=url, json=payload, timeout=600)
|
||||
r = response.json()
|
||||
|
||||
# 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)
|
||||
# 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():
|
||||
"""
|
||||
|
||||
Reference in New Issue
Block a user