hash/FunguyBot
2
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Latest working copy

Browse Source
This commit is contained in:
Hash Borgir 2024-03-03 06:59:39 -07:00
parent 25f9159155
commit ae2bdfecae
1 changed files with 37 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

44
plugins/stable-diffusion.py
View File

@ -5,11 +5,17 @@ This plugin provides a command to generate images using self hosted Stable Diffu
import requests
import base64
from asyncio import Queue
import argparse
import simplematrixbotlib as botlib
import markdown2
# Queue to store pending commands
command_queue = Queue()
def markdown_to_html(markdown_text):
html_content = markdown2.markdown(markdown_text)
return html_content
async def process_command(room, message, bot, prefix, config):
match = botlib.MessageMatch(room, message, bot, prefix)
if match.prefix() and match.command("sd"):
@ -21,22 +27,46 @@ async def process_command(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"):
prompt = message.body[len(prefix) + len("sd"):].strip() # Extract prompt from message body
payload = {
"prompt": prompt,
"steps": 16
}
url = "http://127.0.0.1:7860/sdapi/v1/txt2img"
try:
parser = argparse.ArgumentParser(description='Generate images using self hosted Stable Diffusion')
parser.add_argument('--steps', type=int, default=12, help='Number of steps, default=16')
parser.add_argument('--cfg', type=int, default=1, 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=['Euler', 'a'], help='Sampler name, default=Euler a')
parser.add_argument('prompt', type=str, nargs='+', help='Prompt for the image')
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"
response = requests.post(url=url, json=payload)
r = response.json()
with open("/tmp/output.png", 'wb') as f:
f.write(base64.b64decode(r['images'][0]))
await bot.api.send_image_message(room_id=room.room_id, image_filepath="/tmp/output.png") # Corrected argument name
except argparse.ArgumentError as e:
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><pre>" + markdown_to_html(parser.format_help()) + "</pre></details>")
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)