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

make its own http session

Browse source
This commit is contained in:
phixxy 2024-02-13 01:45:52 -08:00
parent b1beb57dfc
commit 35558cf155
1 changed files with 15 additions and 10 deletions
Download patch file Download diff file Expand all files Collapse all files

25
extensions/stable_diffusion.py
View file

@ -6,6 +6,7 @@ import os
import time import time
import random import random
from PIL import Image, PngImagePlugin from PIL import Image, PngImagePlugin
import aiohttp
import discord import discord
from discord.ext import commands from discord.ext import commands
@ -19,6 +20,10 @@ class StableDiffusion(commands.Cog):
self.data_dir = "data/stable_diffusion/" self.data_dir = "data/stable_diffusion/"
self.default_neg_prompt = "easynegative, badhandv4, verybadimagenegative_v1.3" self.default_neg_prompt = "easynegative, badhandv4, verybadimagenegative_v1.3"
self.folder_setup() self.folder_setup()
self.http_session = self.create_aiohttp_session()
def create_aiohttp_session(self):
return aiohttp.ClientSession()
def folder_setup(self) -> None: def folder_setup(self) -> None:
try: try:
@ -55,7 +60,7 @@ class StableDiffusion(commands.Cog):
} }
url = "https://api.openai.com/v1/chat/completions" url = "https://api.openai.com/v1/chat/completions"
try: try:
async with self.bot.http_session.post(url, headers=headers, json=data) as resp: async with self.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']
return response return response
@ -149,7 +154,7 @@ class StableDiffusion(commands.Cog):
for attachment in ctx.attachments: for attachment in ctx.attachments:
if attachment.url.endswith(('.jpg', '.png')): if attachment.url.endswith(('.jpg', '.png')):
self.bot.logger.debug("image seen") self.bot.logger.debug("image seen")
async with self.bot.http_session.get(attachment.url) as response: async with self.http_session.get(attachment.url) as response:
imageName = self.working_dir + str(time.time_ns()) + '.png' imageName = self.working_dir + str(time.time_ns()) + '.png'
with open(imageName, 'wb') as out_file: with open(imageName, 'wb') as out_file:
@ -164,12 +169,12 @@ class StableDiffusion(commands.Cog):
try: try:
payload = {"image": img_link} payload = {"image": img_link}
async with self.bot.http_session.post(f'{url}/sdapi/v1/interrogate', json=payload) as response: async with self.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]
metadata += f"<image:{description}>\n" metadata += f"<image:{description}>\n"
except self.bot.aiohttp.ClientError as error: except aiohttp.ClientError:
self.bot.logger.exception("ERROR: CLIP may not be running. Could not look at image.") self.bot.logger.exception("ERROR: CLIP may not be running. Could not look at image.")
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
@ -218,7 +223,7 @@ class StableDiffusion(commands.Cog):
if url == "disabled": if url == "disabled":
await ctx.send("This command is currently disabled") await ctx.send("This command is currently disabled")
else: else:
async with self.bot.http_session.get(url=f'{url}/sdapi/v1/options') as response: async with self.http_session.get(url=f'{url}/sdapi/v1/options') as response:
config_json = await response.json() config_json = await response.json()
current_model = config_json["sd_model_checkpoint"] current_model = config_json["sd_model_checkpoint"]
@ -228,7 +233,7 @@ class StableDiffusion(commands.Cog):
model_id, model_name = model_choices[model_choice] model_id, model_name = model_choices[model_choice]
if current_model != model_id: if current_model != model_id:
payload = {"sd_model_checkpoint": 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: async with self.http_session.post(url=f'{url}/sdapi/v1/options', json=payload) as response:
output = "Changed model to: " + model_name output = "Changed model to: " + model_name
await ctx.send(output) await ctx.send(output)
return return
@ -301,7 +306,7 @@ class StableDiffusion(commands.Cog):
if key_value_pairs: if key_value_pairs:
payload.update(key_value_pairs) payload.update(key_value_pairs)
try: try:
async with self.bot.http_session.post(url, headers=headers, json=payload) as resp: async with self.http_session.post(url, headers=headers, json=payload) as resp:
if resp.status != 200: if resp.status != 200:
await ctx.send(f"{resp.status} {resp.reason}") await ctx.send(f"{resp.status} {resp.reason}")
self.bot.logger.exception(f"{resp.status} {resp.reason}") self.bot.logger.exception(f"{resp.status} {resp.reason}")
@ -328,7 +333,7 @@ class StableDiffusion(commands.Cog):
The path to the saved image file. The path to the saved image file.
""" """
async def save_image(self, url: str) -> str: async def save_image(self, url: str) -> str:
async with self.bot.http_session.get(url) as response: async with self.http_session.get(url) as response:
image_name = self.working_dir + str(time.time_ns()) + ".png" image_name = self.working_dir + str(time.time_ns()) + ".png"
with open(image_name, 'wb') as out_file: with open(image_name, 'wb') as out_file:
self.bot.logger.debug(f"Saving image: {image_name}") self.bot.logger.debug(f"Saving image: {image_name}")
@ -371,7 +376,7 @@ class StableDiffusion(commands.Cog):
if key_value_pairs: if key_value_pairs:
payload.update(key_value_pairs) payload.update(key_value_pairs)
try: try:
async with self.bot.http_session.post(url, headers=headers, json=payload) as resp: async with self.http_session.post(url, headers=headers, json=payload) as resp:
if resp.status != 200: if resp.status != 200:
await ctx.send(f"{resp.status} {resp.reason}") await ctx.send(f"{resp.status} {resp.reason}")
self.bot.logger.error(f"{resp.status} {resp.reason}") self.bot.logger.error(f"{resp.status} {resp.reason}")
@ -484,7 +489,7 @@ class StableDiffusion(commands.Cog):
img_link = await self.my_open_img_file(image_name) img_link = await self.my_open_img_file(image_name)
try: try:
payload = {"image": img_link} payload = {"image": img_link}
async with self.bot.http_session.post(url, json=payload) as response: async with self.http_session.post(url, json=payload) as response:
r = await response.json() r = await response.json()
await ctx.send(r.get("caption")) await ctx.send(r.get("caption"))
except: except: