From 4f64e470571d81bd458a90db4d99e24625bd9b3e Mon Sep 17 00:00:00 2001 From: phixxy Date: Fri, 19 Jan 2024 18:38:24 -0800 Subject: [PATCH] now uses the sparkytron http_session cleaned up folders into variable names --- plugins/stable_diffusion.py | 157 ++++++++++++++++++------------------ 1 file changed, 80 insertions(+), 77 deletions(-) diff --git a/plugins/stable_diffusion.py b/plugins/stable_diffusion.py index f98f4d7..64f05c8 100644 --- a/plugins/stable_diffusion.py +++ b/plugins/stable_diffusion.py @@ -15,8 +15,19 @@ class StableDiffusion(commands.Cog): def __init__(self, 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.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 headers = { @@ -32,11 +43,9 @@ class StableDiffusion(commands.Cog): url = "https://api.openai.com/v1/chat/completions" try: - #http_session = aiohttp.ClientSession() async with self.bot.http_session.post(url, headers=headers, json=data) as resp: response_data = await resp.json() response = response_data['choices'][0]['message']['content'] - #await http_session.close() return response except Exception as error: @@ -46,12 +55,13 @@ class StableDiffusion(commands.Cog): print(error) current_time = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()) 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) return error - def extract_key_value_pairs(self, input_str): + async def extract_key_value_pairs(self, input_str): output_str = input_str key_value_pairs = {} @@ -64,9 +74,8 @@ class StableDiffusion(commands.Cog): 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) - w, h = img.size encoded = "" with io.BytesIO() as output: img.save(output, format="PNG") @@ -78,16 +87,14 @@ class StableDiffusion(commands.Cog): async def look_at(self, ctx, look=False): metadata = "" if look: - url = os.getenv('stablediffusion_url') + url = self.stable_diffusion_url if url == "disabled": return for attachment in ctx.attachments: if attachment.url.endswith(('.jpg', '.png')): print("image seen") - - http_session = aiohttp.ClientSession() - async with http_session.get(attachment.url) as response: - imageName = "tmp/" + str(len(os.listdir('tmp/'))) + '.png' + async with self.bot.http_session.get(attachment.url) as response: + imageName = self.working_dir + str(time.time_ns()) + '.png' with open(imageName, 'wb') as out_file: print('Saving image: ' + imageName) @@ -101,7 +108,7 @@ class StableDiffusion(commands.Cog): try: 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() description = data.get("caption") description = description.split(',')[0] @@ -109,7 +116,6 @@ class StableDiffusion(commands.Cog): except aiohttp.ClientError as error: await self.handle_error(error) return "ERROR: CLIP may not be running. Could not look at image." - await http_session.close() return metadata @commands.command( @@ -118,7 +124,7 @@ class StableDiffusion(commands.Cog): brief="Generate a random image" ) async def draw(self, ctx): - url = os.getenv('stablediffusion_url') + url = self.stable_diffusion_url if url == "disabled": return try: @@ -149,21 +155,20 @@ class StableDiffusion(commands.Cog): negative_prompt = "easynegative verybadimagenegative_v1.3" payload = {"prompt": prompt,"steps": 25, "negative_prompt": negative_prompt,"batch_size": amount} try: - http_session = aiohttp.ClientSession() - async with http_session.post(url=f'{url}/sdapi/v1/txt2img', json=payload) as resp: + async with self.bot.http_session.post(url=f'{url}/sdapi/v1/txt2img', json=payload) as resp: r = await resp.json() for i in r['images']: image = Image.open(io.BytesIO(base64.b64decode(i.split(",",1)[0]))) 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() pnginfo = PngImagePlugin.PngInfo() pnginfo.add_text("parameters", response2.get("info")) try: if ctx.channel.is_nsfw(): - folder = "tmp/nsfw/" + folder = self.working_dir + "nsfw/" else: - folder = "tmp/sfw/" + folder = self.working_dir + "sfw/" except: folder = "tmp/" my_filename = folder + str(time.time_ns()) + ".png" @@ -178,14 +183,14 @@ class StableDiffusion(commands.Cog): except Exception as error: await self.handle_error(error) await ctx.send('Did you mean to use !imagine?. Usage: !draw (number)') - await http_session.close() + await self.bot.http_session.close() @commands.command( description="Change Model", help="Choose from a list of stable diffusion models.", 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 = { '1': ("deliberate_v2.safetensors [9aba26abdf]", "DeliberateV2"), '2': ("flat2DAnimerge_v30.safetensors [5dd56bfa12]", "Flat2D"), @@ -194,34 +199,30 @@ class StableDiffusion(commands.Cog): '5': ("Pixel_Art_V1_PublicPrompts.ckpt [0f02127697]", "Pixel Art"), '6': ("mistoonAnime_v20.safetensors [c35e1054c0]", "Mistoon AnimeV2") } - url = os.getenv('stablediffusion_url') + url = self.stable_diffusion_url if url == "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 - async with http_session.get(url=f'{url}/sdapi/v1/options') as response: - config_json = await response.json() + current_model = config_json["sd_model_checkpoint"] + output = 'Current Model: ' + current_model + '\n' - current_model = config_json["sd_model_checkpoint"] - output = 'Current Model: ' + current_model + '\n' - - if model_choice in model_choices: - model_id, model_name = model_choices[model_choice] - if current_model != model_id: - payload = {"sd_model_checkpoint": model_id} - async with http_session.post(url=f'{url}/sdapi/v1/options', json=payload) as response: - output = "Changed model to: " + model_name - await ctx.send(output) + if model_choice in model_choices: + model_id, model_name = model_choices[model_choice] + if current_model != model_id: + payload = {"sd_model_checkpoint": model_id} + async with self.bot.http_session.post(url=f'{url}/sdapi/v1/options', json=payload) as response: + output = "Changed model to: " + model_name + await ctx.send(output) + return + else: + await ctx.send(f"Already set to use {model_name}") return else: - await ctx.send(f"Already set to use {model_name}") - await http_session.close() - return - else: - output = '\n'.join([f"{choice}: {name}" for choice, name in model_choices.items()]) - await http_session.close() - await ctx.send(output) + output = '\n'.join([f"{choice}: {name}" for choice, name in model_choices.items()]) + await ctx.send(output) @commands.command( description="Lora", @@ -247,7 +248,7 @@ class StableDiffusion(commands.Cog): brief="Generate an image" ) async def imagine(self, ctx): - url = os.getenv('stablediffusion_url') + url = self.stable_diffusion_url if url == "disabled": await ctx.send("Command is currently disabled.") return @@ -255,9 +256,15 @@ class StableDiffusion(commands.Cog): url=f"{url}/sdapi/v1/txt2img" prompt = ctx.message.content.split(" ", maxsplit=1)[1] key_value_pairs, prompt = self.extract_key_value_pairs(prompt) - neg_prompt_file = "databases/negative_prompt.txt" - with open(neg_prompt_file, 'r') as f: - negative_prompt = f.readline() + neg_prompt_file = f"{self.db_dir}negative_prompt.txt" + try: + 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 + ".") payload = { @@ -269,9 +276,8 @@ class StableDiffusion(commands.Cog): 'Content-Type': 'application/json' } payload.update(key_value_pairs) - http_session = aiohttp.ClientSession() 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() except Exception as error: 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} 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() except Exception as error: 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")) try: if ctx.channel.is_nsfw(): - folder = "tmp/" + folder = self.working_dir else: - folder = "tmp/sfw/" + folder = self.working_dir + "sfw/" except: - folder = "users/" + str(ctx.author.id) + '/' + folder = self.working_dir + str(ctx.author.id) + '/' my_filename = folder + str(time.time_ns()) + ".png" image.save(my_filename, pnginfo=pnginfo) @@ -304,11 +310,10 @@ class StableDiffusion(commands.Cog): f = discord.File(fh, filename=my_filename) 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) await ctx.send(f'Generated by: {ctx.author.name}\nPrompt: {prompt}', file=f) - await http_session.close() @commands.command( @@ -317,7 +322,7 @@ class StableDiffusion(commands.Cog): brief="Describe image" ) async def describe(self, ctx): - url = os.getenv('stablediffusion_url') + url = self.stable_diffusion_url if url == "disabled": await ctx.send("Command is currently disabled") return @@ -334,11 +339,9 @@ class StableDiffusion(commands.Cog): except Exception as error: await self.handle_error(error) print("Couldn't find image.") - return - - http_session = aiohttp.ClientSession() - async with http_session.get(file_url) as response: - imageName = "tmp/" + str(len(os.listdir("tmp/"))) + ".png" + return + async with self.bot.http_session.get(file_url) as response: + imageName = self.working_dir + str(time.time_ns()) + ".png" with open(imageName, 'wb') as out_file: print(f"Saving image: {imageName}") while True: @@ -350,14 +353,13 @@ class StableDiffusion(commands.Cog): img_link = self.my_open_img_file(imageName) try: 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() print(r) await ctx.send(r.get("caption")) except Exception as error: await self.handle_error(error) await ctx.send("My image generation service may not be running.") - await http_session.close() @commands.command( description="Reimagine", @@ -365,7 +367,7 @@ class StableDiffusion(commands.Cog): brief="Reimagine an image" ) async def reimagine(self, ctx): - url = os.getenv('stablediffusion_url') + url = self.stable_diffusion_url if url == "disabled": await ctx.send("Command is currently disabled") return @@ -383,12 +385,10 @@ class StableDiffusion(commands.Cog): await self.handle_error(error) print("Couldn't find image.") return - key_value_pairs, prompt = self.extract_key_value_pairs(prompt) - http_session = aiohttp.ClientSession() try: - async with http_session.get(file_url) as response: - imageName = "tmp/" + str(len(os.listdir("tmp/"))) + ".png" + async with self.bot.http_session.get(file_url) as response: + imageName = self.working_dir + str(time.time_ns()) + ".png" with open(imageName, 'wb') as out_file: print(f"Saving image: {imageName}") while True: @@ -411,18 +411,18 @@ class StableDiffusion(commands.Cog): payload.update(key_value_pairs) 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() for i in data['images']: - if not os.path.isdir("tmp/reimagined/"+ str(ctx.author.id)): - os.makedirs("tmp/reimagined/"+ str(ctx.author.id)) + if not os.path.isdir(f"{self.working_dir}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]))) 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() pnginfo = PngImagePlugin.PngInfo() 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) with open(my_filename, "rb") as fh: f = discord.File(fh, filename=my_filename) @@ -430,7 +430,6 @@ class StableDiffusion(commands.Cog): except Exception as error: await ctx.send("My image generation service may not be running.") await self.handle_error(error) - await http_session.close() @commands.command( description="Negative Prompt", @@ -441,12 +440,16 @@ class StableDiffusion(commands.Cog): message = ' '.join(args) if not message: 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: f.writelines(message) await ctx.send("Changed negative prompt to " + message) async def setup(bot): - print("Added StableDiffusion Cog") - await bot.add_cog(StableDiffusion(bot)) \ No newline at end of file + try: + await bot.add_cog(StableDiffusion(bot)) + print("Successfully added StableDiffusion Cog") + except: + print("Failed to load StableDiffusion Cog") + \ No newline at end of file