phixxy/sparkytron3000
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

now uses the sparkytron http_session

Browse source 
cleaned up folders into variable names
This commit is contained in:
phixxy 2024-01-19 18:38:24 -08:00
parent 3861f20e1f
commit 4f64e47057
1 changed files with 80 additions and 77 deletions
Download patch file Download diff file Expand all files Collapse all files

157
plugins/stable_diffusion.py
View file

@ -15,8 +15,19 @@ class StableDiffusion(commands.Cog):
def __init__(self, bot): def __init__(self, bot):
self.bot = bot self.bot = bot
self.stable_diffusion_url = os.getenv("stablediffusion_url") # Change this to stable_diffusion_url
self.working_dir = "tmp/stable_diffusion/" self.working_dir = "tmp/stable_diffusion/"
self.db_dir = "db/stable_diffusion/" self.db_dir = "db/stable_diffusion/"
self.folder_setup()
def folder_setup(self):
try:
if not os.path.exists(self.working_dir):
os.mkdir(self.working_dir)
if not os.path.exists(self.db_dir):
os.mkdir(self.db_dir)
except:
print(" StableDiffusion failed to make directories")
async def answer_question(self, topic, model="gpt-3.5-turbo"): # Only needed for draw command async def answer_question(self, topic, model="gpt-3.5-turbo"): # Only needed for draw command
headers = { headers = {
@ -32,11 +43,9 @@ class StableDiffusion(commands.Cog):
url = "https://api.openai.com/v1/chat/completions" url = "https://api.openai.com/v1/chat/completions"
try: try:
#http_session = aiohttp.ClientSession()
async with self.bot.http_session.post(url, headers=headers, json=data) as resp: async with self.bot.http_session.post(url, headers=headers, json=data) as resp:
response_data = await resp.json() response_data = await resp.json()
response = response_data['choices'][0]['message']['content'] response = response_data['choices'][0]['message']['content']
#await http_session.close()
return response return response
except Exception as error: except Exception as error:
@ -46,12 +55,13 @@ class StableDiffusion(commands.Cog):
print(error) print(error)
current_time = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()) current_time = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())
log_line = current_time + ': ' + str(error) + '\n' log_line = current_time + ': ' + str(error) + '\n'
with open("databases/error_log.txt", 'a') as f: file_name = self.db_dir + "error_log.log"
with open(file_name, 'a') as f:
f.write(log_line) f.write(log_line)
return error return error
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 = {}
@ -64,9 +74,8 @@ class StableDiffusion(commands.Cog):
return key_value_pairs, output_str return key_value_pairs, output_str
def my_open_img_file(self, path): async def my_open_img_file(self, path):
img = Image.open(path) img = Image.open(path)
w, h = img.size
encoded = "" encoded = ""
with io.BytesIO() as output: with io.BytesIO() as output:
img.save(output, format="PNG") img.save(output, format="PNG")
@ -78,16 +87,14 @@ class StableDiffusion(commands.Cog):
async def look_at(self, ctx, look=False): async def look_at(self, ctx, look=False):
metadata = "" metadata = ""
if look: if look:
url = os.getenv('stablediffusion_url') url = self.stable_diffusion_url
if url == "disabled": if url == "disabled":
return return
for attachment in ctx.attachments: for attachment in ctx.attachments:
if attachment.url.endswith(('.jpg', '.png')): if attachment.url.endswith(('.jpg', '.png')):
print("image seen") print("image seen")
async with self.bot.http_session.get(attachment.url) as response:
http_session = aiohttp.ClientSession() imageName = self.working_dir + str(time.time_ns()) + '.png'
async with http_session.get(attachment.url) as response:
imageName = "tmp/" + str(len(os.listdir('tmp/'))) + '.png'
with open(imageName, 'wb') as out_file: with open(imageName, 'wb') as out_file:
print('Saving image: ' + imageName) print('Saving image: ' + imageName)
@ -101,7 +108,7 @@ class StableDiffusion(commands.Cog):
try: try:
payload = {"image": img_link} payload = {"image": img_link}
async with http_session.post(f'{url}/sdapi/v1/interrogate', json=payload) as response: async with self.bot.http_session.post(f'{url}/sdapi/v1/interrogate', json=payload) as response:
data = await response.json() data = await response.json()
description = data.get("caption") description = data.get("caption")
description = description.split(',')[0] description = description.split(',')[0]
@ -109,7 +116,6 @@ class StableDiffusion(commands.Cog):
except aiohttp.ClientError as error: except 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."
await http_session.close()
return metadata return metadata
@commands.command( @commands.command(
@ -118,7 +124,7 @@ class StableDiffusion(commands.Cog):
brief="Generate a random image" brief="Generate a random image"
) )
async def draw(self, ctx): async def draw(self, ctx):
url = os.getenv('stablediffusion_url') url = self.stable_diffusion_url
if url == "disabled": if url == "disabled":
return return
try: try:
@ -149,21 +155,20 @@ class StableDiffusion(commands.Cog):
negative_prompt = "easynegative verybadimagenegative_v1.3" negative_prompt = "easynegative verybadimagenegative_v1.3"
payload = {"prompt": prompt,"steps": 25, "negative_prompt": negative_prompt,"batch_size": amount} payload = {"prompt": prompt,"steps": 25, "negative_prompt": negative_prompt,"batch_size": amount}
try: try:
http_session = aiohttp.ClientSession() async with self.bot.http_session.post(url=f'{url}/sdapi/v1/txt2img', json=payload) as resp:
async with http_session.post(url=f'{url}/sdapi/v1/txt2img', json=payload) as resp:
r = await resp.json() r = await resp.json()
for i in r['images']: for i in r['images']:
image = Image.open(io.BytesIO(base64.b64decode(i.split(",",1)[0]))) image = Image.open(io.BytesIO(base64.b64decode(i.split(",",1)[0])))
png_payload = {"image": "data:image/png;base64," + i} png_payload = {"image": "data:image/png;base64," + i}
async with http_session.post(url=f'{url}/sdapi/v1/png-info', json=png_payload) as resp2: async with self.bot.http_session.post(url=f'{url}/sdapi/v1/png-info', json=png_payload) as resp2:
response2 = await resp2.json() response2 = await resp2.json()
pnginfo = PngImagePlugin.PngInfo() pnginfo = PngImagePlugin.PngInfo()
pnginfo.add_text("parameters", response2.get("info")) pnginfo.add_text("parameters", response2.get("info"))
try: try:
if ctx.channel.is_nsfw(): if ctx.channel.is_nsfw():
folder = "tmp/nsfw/" folder = self.working_dir + "nsfw/"
else: else:
folder = "tmp/sfw/" folder = self.working_dir + "sfw/"
except: except:
folder = "tmp/" folder = "tmp/"
my_filename = folder + str(time.time_ns()) + ".png" my_filename = folder + str(time.time_ns()) + ".png"
@ -178,14 +183,14 @@ class StableDiffusion(commands.Cog):
except Exception as error: except Exception as error:
await self.handle_error(error) await self.handle_error(error)
await ctx.send('Did you mean to use !imagine?. Usage: !draw (number)') await ctx.send('Did you mean to use !imagine?. Usage: !draw (number)')
await http_session.close() await self.bot.http_session.close()
@commands.command( @commands.command(
description="Change Model", description="Change Model",
help="Choose from a list of stable diffusion models.", help="Choose from a list of stable diffusion models.",
brief="Change stable diffusion model" brief="Change stable diffusion model"
) )
async def change_model(ctx, model_choice='0'): # Needs to be a configurable list of models async def change_model(self, ctx, model_choice='0'): # Needs to be a configurable list of models
model_choices = { model_choices = {
'1': ("deliberate_v2.safetensors [9aba26abdf]", "DeliberateV2"), '1': ("deliberate_v2.safetensors [9aba26abdf]", "DeliberateV2"),
'2': ("flat2DAnimerge_v30.safetensors [5dd56bfa12]", "Flat2D"), '2': ("flat2DAnimerge_v30.safetensors [5dd56bfa12]", "Flat2D"),
@ -194,34 +199,30 @@ class StableDiffusion(commands.Cog):
'5': ("Pixel_Art_V1_PublicPrompts.ckpt [0f02127697]", "Pixel Art"), '5': ("Pixel_Art_V1_PublicPrompts.ckpt [0f02127697]", "Pixel Art"),
'6': ("mistoonAnime_v20.safetensors [c35e1054c0]", "Mistoon AnimeV2") '6': ("mistoonAnime_v20.safetensors [c35e1054c0]", "Mistoon AnimeV2")
} }
url = os.getenv('stablediffusion_url') url = self.stable_diffusion_url
if url == "disabled": if url == "disabled":
await ctx.send("This command is currently disabled") await ctx.send("This command is currently disabled")
return else:
async with self.bot.http_session.get(url=f'{url}/sdapi/v1/options') as response:
config_json = await response.json()
http_session = aiohttp.ClientSession() # We can probably pass an http session from the bot current_model = config_json["sd_model_checkpoint"]
async with http_session.get(url=f'{url}/sdapi/v1/options') as response: output = 'Current Model: ' + current_model + '\n'
config_json = await response.json()
current_model = config_json["sd_model_checkpoint"] if model_choice in model_choices:
output = 'Current Model: ' + current_model + '\n' model_id, model_name = model_choices[model_choice]
if current_model != model_id:
if model_choice in model_choices: payload = {"sd_model_checkpoint": model_id}
model_id, model_name = model_choices[model_choice] async with self.bot.http_session.post(url=f'{url}/sdapi/v1/options', json=payload) as response:
if current_model != model_id: output = "Changed model to: " + model_name
payload = {"sd_model_checkpoint": model_id} await ctx.send(output)
async with http_session.post(url=f'{url}/sdapi/v1/options', json=payload) as response: return
output = "Changed model to: " + model_name else:
await ctx.send(output) await ctx.send(f"Already set to use {model_name}")
return return
else: else:
await ctx.send(f"Already set to use {model_name}") output = '\n'.join([f"{choice}: {name}" for choice, name in model_choices.items()])
await http_session.close() await ctx.send(output)
return
else:
output = '\n'.join([f"{choice}: {name}" for choice, name in model_choices.items()])
await http_session.close()
await ctx.send(output)
@commands.command( @commands.command(
description="Lora", description="Lora",
@ -247,7 +248,7 @@ class StableDiffusion(commands.Cog):
brief="Generate an image" brief="Generate an image"
) )
async def imagine(self, ctx): async def imagine(self, ctx):
url = os.getenv('stablediffusion_url') url = self.stable_diffusion_url
if url == "disabled": if url == "disabled":
await ctx.send("Command is currently disabled.") await ctx.send("Command is currently disabled.")
return return
@ -255,9 +256,15 @@ class StableDiffusion(commands.Cog):
url=f"{url}/sdapi/v1/txt2img" url=f"{url}/sdapi/v1/txt2img"
prompt = ctx.message.content.split(" ", maxsplit=1)[1] prompt = ctx.message.content.split(" ", maxsplit=1)[1]
key_value_pairs, prompt = self.extract_key_value_pairs(prompt) key_value_pairs, prompt = self.extract_key_value_pairs(prompt)
neg_prompt_file = "databases/negative_prompt.txt" neg_prompt_file = f"{self.db_dir}negative_prompt.txt"
with open(neg_prompt_file, 'r') as f: try:
negative_prompt = f.readline() with open(neg_prompt_file, 'r') as f:
negative_prompt = f.readline()
except:
with open(neg_prompt_file, 'w') as f:
f.writelines("")
negative_prompt = ""
await ctx.send("Please be patient this may take some time! Generating: " + prompt + ".") await ctx.send("Please be patient this may take some time! Generating: " + prompt + ".")
payload = { payload = {
@ -269,9 +276,8 @@ class StableDiffusion(commands.Cog):
'Content-Type': 'application/json' 'Content-Type': 'application/json'
} }
payload.update(key_value_pairs) payload.update(key_value_pairs)
http_session = aiohttp.ClientSession()
try: try:
async with http_session.post(url, headers=headers, json=payload) as resp: async with self.bot.http_session.post(url, headers=headers, json=payload) as resp:
r = await resp.json() r = await resp.json()
except Exception as error: except Exception as error:
await ctx.send("My image generation service may not be running.") await ctx.send("My image generation service may not be running.")
@ -282,7 +288,7 @@ class StableDiffusion(commands.Cog):
png_payload = {"image": "data:image/png;base64," + i} png_payload = {"image": "data:image/png;base64," + i}
try: try:
async with http_session.post(url, json=png_payload) as resp: async with self.bot.http_session.post(url, json=png_payload) as resp:
response2 = await resp.json() response2 = await resp.json()
except Exception as error: except Exception as error:
await ctx.send("My image generation service may not be running.") await ctx.send("My image generation service may not be running.")
@ -292,11 +298,11 @@ class StableDiffusion(commands.Cog):
pnginfo.add_text("parameters", response2.get("info")) pnginfo.add_text("parameters", response2.get("info"))
try: try:
if ctx.channel.is_nsfw(): if ctx.channel.is_nsfw():
folder = "tmp/" folder = self.working_dir
else: else:
folder = "tmp/sfw/" folder = self.working_dir + "sfw/"
except: except:
folder = "users/" + str(ctx.author.id) + '/' folder = self.working_dir + str(ctx.author.id) + '/'
my_filename = folder + str(time.time_ns()) + ".png" my_filename = folder + str(time.time_ns()) + ".png"
image.save(my_filename, pnginfo=pnginfo) image.save(my_filename, pnginfo=pnginfo)
@ -304,11 +310,10 @@ class StableDiffusion(commands.Cog):
f = discord.File(fh, filename=my_filename) f = discord.File(fh, filename=my_filename)
log_data = f'Author: {ctx.author.name}, Prompt: {prompt}, Filename: {my_filename}\n' log_data = f'Author: {ctx.author.name}, Prompt: {prompt}, Filename: {my_filename}\n'
with open("databases/stable_diffusion.log", 'a') as log_file: with open(f"{self.db_dir}stable_diffusion.log", 'a') as log_file:
log_file.writelines(log_data) log_file.writelines(log_data)
await ctx.send(f'Generated by: {ctx.author.name}\nPrompt: {prompt}', file=f) await ctx.send(f'Generated by: {ctx.author.name}\nPrompt: {prompt}', file=f)
await http_session.close()
@commands.command( @commands.command(
@ -317,7 +322,7 @@ class StableDiffusion(commands.Cog):
brief="Describe image" brief="Describe image"
) )
async def describe(self, ctx): async def describe(self, ctx):
url = os.getenv('stablediffusion_url') url = self.stable_diffusion_url
if url == "disabled": if url == "disabled":
await ctx.send("Command is currently disabled") await ctx.send("Command is currently disabled")
return return
@ -334,11 +339,9 @@ class StableDiffusion(commands.Cog):
except Exception as error: except Exception as error:
await self.handle_error(error) await self.handle_error(error)
print("Couldn't find image.") print("Couldn't find image.")
return return
async with self.bot.http_session.get(file_url) as response:
http_session = aiohttp.ClientSession() imageName = self.working_dir + str(time.time_ns()) + ".png"
async with http_session.get(file_url) as response:
imageName = "tmp/" + str(len(os.listdir("tmp/"))) + ".png"
with open(imageName, 'wb') as out_file: with open(imageName, 'wb') as out_file:
print(f"Saving image: {imageName}") print(f"Saving image: {imageName}")
while True: while True:
@ -350,14 +353,13 @@ class StableDiffusion(commands.Cog):
img_link = self.my_open_img_file(imageName) img_link = self.my_open_img_file(imageName)
try: try:
payload = {"image": img_link} payload = {"image": img_link}
async with http_session.post(url, json=payload) as response: async with self.bot.http_session.post(url, json=payload) as response:
r = await response.json() r = await response.json()
print(r) print(r)
await ctx.send(r.get("caption")) await ctx.send(r.get("caption"))
except Exception as error: except Exception as error:
await self.handle_error(error) await self.handle_error(error)
await ctx.send("My image generation service may not be running.") await ctx.send("My image generation service may not be running.")
await http_session.close()
@commands.command( @commands.command(
description="Reimagine", description="Reimagine",
@ -365,7 +367,7 @@ class StableDiffusion(commands.Cog):
brief="Reimagine an image" brief="Reimagine an image"
) )
async def reimagine(self, ctx): async def reimagine(self, ctx):
url = os.getenv('stablediffusion_url') url = self.stable_diffusion_url
if url == "disabled": if url == "disabled":
await ctx.send("Command is currently disabled") await ctx.send("Command is currently disabled")
return return
@ -383,12 +385,10 @@ class StableDiffusion(commands.Cog):
await self.handle_error(error) await self.handle_error(error)
print("Couldn't find image.") print("Couldn't find image.")
return return
key_value_pairs, prompt = self.extract_key_value_pairs(prompt) key_value_pairs, prompt = self.extract_key_value_pairs(prompt)
http_session = aiohttp.ClientSession()
try: try:
async with http_session.get(file_url) as response: async with self.bot.http_session.get(file_url) as response:
imageName = "tmp/" + str(len(os.listdir("tmp/"))) + ".png" imageName = self.working_dir + str(time.time_ns()) + ".png"
with open(imageName, 'wb') as out_file: with open(imageName, 'wb') as out_file:
print(f"Saving image: {imageName}") print(f"Saving image: {imageName}")
while True: while True:
@ -411,18 +411,18 @@ class StableDiffusion(commands.Cog):
payload.update(key_value_pairs) payload.update(key_value_pairs)
try: try:
async with http_session.post(url=f'{url}/sdapi/v1/img2img', json=payload) as response: async with self.bot.http_session.post(url=f'{url}/sdapi/v1/img2img', json=payload) as response:
data = await response.json() data = await response.json()
for i in data['images']: for i in data['images']:
if not os.path.isdir("tmp/reimagined/"+ str(ctx.author.id)): if not os.path.isdir(f"{self.working_dir}reimagined/"+ str(ctx.author.id)):
os.makedirs("tmp/reimagined/"+ str(ctx.author.id)) os.makedirs(f"{self.working_dir}reimagined/"+ str(ctx.author.id))
image = Image.open(io.BytesIO(base64.b64decode(i.split(",",1)[0]))) image = Image.open(io.BytesIO(base64.b64decode(i.split(",",1)[0])))
png_payload = {"image": "data:image/png;base64," + i} png_payload = {"image": "data:image/png;base64," + i}
async with http_session.post(url=f'{url}/sdapi/v1/png-info', json=png_payload) as resp2: async with self.bot.http_session.post(url=f'{url}/sdapi/v1/png-info', json=png_payload) as resp2:
response2 = await resp2.json() response2 = await resp2.json()
pnginfo = PngImagePlugin.PngInfo() pnginfo = PngImagePlugin.PngInfo()
pnginfo.add_text("parameters", response2.get("info")) pnginfo.add_text("parameters", response2.get("info"))
my_filename = "tmp/" + str(len(os.listdir("tmp/"))) + ".png" my_filename = self.working_dir + str(time.time_ns()) + ".png"
image.save(my_filename, pnginfo=pnginfo) image.save(my_filename, pnginfo=pnginfo)
with open(my_filename, "rb") as fh: with open(my_filename, "rb") as fh:
f = discord.File(fh, filename=my_filename) f = discord.File(fh, filename=my_filename)
@ -430,7 +430,6 @@ class StableDiffusion(commands.Cog):
except Exception as error: except Exception as error:
await ctx.send("My image generation service may not be running.") await ctx.send("My image generation service may not be running.")
await self.handle_error(error) await self.handle_error(error)
await http_session.close()
@commands.command( @commands.command(
description="Negative Prompt", description="Negative Prompt",
@ -441,12 +440,16 @@ class StableDiffusion(commands.Cog):
message = ' '.join(args) message = ' '.join(args)
if not message: if not message:
message = "easynegative, badhandv4, verybadimagenegative_v1.3" message = "easynegative, badhandv4, verybadimagenegative_v1.3"
neg_prompt_file = "databases/negative_prompt.txt" neg_prompt_file = f"{self.db_dir}negative_prompt.txt"
with open(neg_prompt_file, 'w') as f: with open(neg_prompt_file, 'w') as f:
f.writelines(message) f.writelines(message)
await ctx.send("Changed negative prompt to " + message) await ctx.send("Changed negative prompt to " + message)
async def setup(bot): async def setup(bot):
print("Added StableDiffusion Cog") try:
await bot.add_cog(StableDiffusion(bot)) await bot.add_cog(StableDiffusion(bot))
print("Successfully added StableDiffusion Cog")
except:
print("Failed to load StableDiffusion Cog")