added get_prompt_from_ctx

added get_kv_from_ctx
both to replace get_key_value_pairs
This commit is contained in:
phixxy 2024-01-19 20:51:40 -08:00
parent 7a1db5efc1
commit d6f1333378

View file

@ -2,7 +2,6 @@
# This extension enables the ability to generate AI artwork using the AUTOMATIC1111 API # This extension enables the ability to generate AI artwork using the AUTOMATIC1111 API
import io import io
import base64 import base64
import aiohttp
import os import os
import time import time
import random import random
@ -65,7 +64,6 @@ class StableDiffusion(commands.Cog):
async def extract_key_value_pairs(self, input_str): async def extract_key_value_pairs(self, input_str):
output_str = input_str output_str = input_str
key_value_pairs = {} key_value_pairs = {}
tokens = input_str.split(', ') tokens = input_str.split(', ')
for token in tokens: for token in tokens:
@ -76,6 +74,16 @@ class StableDiffusion(commands.Cog):
return key_value_pairs, output_str return key_value_pairs, output_str
def get_kv_from_ctx(self, ctx):
prompt = ctx.message.content.split(" ", maxsplit=1)[1]
kv_strings = list(filter(lambda x: '=' in x,prompt.split(' ')))
key_value_pairs = dict(map(lambda a: a.replace(',','').split('='),kv_strings))
return key_value_pairs
def get_prompt_from_ctx(self, ctx):
prompt = ctx.message.content.split(" ", maxsplit=1)[1]
prompt = ' '.join(list(filter(lambda x: '=' not in x,prompt.split(' '))))
async def my_open_img_file(self, path): async def my_open_img_file(self, path):
img = Image.open(path) img = Image.open(path)
encoded = "" encoded = ""
@ -115,7 +123,7 @@ class StableDiffusion(commands.Cog):
description = data.get("caption") description = data.get("caption")
description = description.split(',')[0] description = description.split(',')[0]
metadata += f"<image:{description}>\n" metadata += f"<image:{description}>\n"
except aiohttp.ClientError as error: except self.bot.aiohttp.ClientError as error:
await self.handle_error(error) await self.handle_error(error)
return "ERROR: CLIP may not be running. Could not look at image." return "ERROR: CLIP may not be running. Could not look at image."
return metadata return metadata