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

Now a cog

Browse source
This commit is contained in:
phixxy 2024-01-19 17:26:51 -08:00
parent c193eb9134
commit 0d5dff4eb1
1 changed files with 409 additions and 408 deletions
Download patch file Download diff file Expand all files Collapse all files

103
plugins/stable_diffusion.py
View file

@ -1,17 +1,24 @@
#plugin file for sparkytron 3000 # Extension for sparkytron 3000
# This extension enables the ability to generate AI artwork using the AUTOMATIC1111 API
import io import io
import base64 import base64
import aiohttp import aiohttp
import os import os
import time import time
import random import random
import json
from PIL import Image, PngImagePlugin from PIL import Image, PngImagePlugin
from discord.ext import commands
import discord import discord
from discord.ext import commands
async def answer_question(topic, model="gpt-3.5-turbo"): # Only needed for draw command
class StableDiffusion(commands.Cog):
def __init__(self, bot):
self.bot = bot
self.working_dir = "tmp/stable_diffusion/"
self.db_dir = "db/stable_diffusion/"
async def answer_question(self, topic, model="gpt-3.5-turbo"): # Only needed for draw command
headers = { headers = {
'Content-Type': 'application/json', 'Content-Type': 'application/json',
'Authorization': f'Bearer {os.getenv("openai.api_key")}', 'Authorization': f'Bearer {os.getenv("openai.api_key")}',
@ -25,17 +32,17 @@ async def answer_question(topic, model="gpt-3.5-turbo"): # Only needed for draw
url = "https://api.openai.com/v1/chat/completions" url = "https://api.openai.com/v1/chat/completions"
try: try:
http_session = aiohttp.ClientSession() #http_session = aiohttp.ClientSession()
async with 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() #await http_session.close()
return response return response
except Exception as error: except Exception as error:
return await handle_error(error) return await self.handle_error(error)
async def handle_error(error): async def handle_error(self, error): # This needs to be deleted and replaced with logging
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'
@ -44,7 +51,7 @@ async def handle_error(error):
return error return error
def extract_key_value_pairs(input_str): def extract_key_value_pairs(self, input_str):
output_str = input_str output_str = input_str
key_value_pairs = {} key_value_pairs = {}
@ -57,7 +64,7 @@ def extract_key_value_pairs(input_str):
return key_value_pairs, output_str return key_value_pairs, output_str
def my_open_img_file(path): def my_open_img_file(self, path):
img = Image.open(path) img = Image.open(path)
w, h = img.size w, h = img.size
encoded = "" encoded = ""
@ -68,7 +75,7 @@ def my_open_img_file(path):
img.close() img.close()
return encoded return encoded
async def look_at(ctx, look=False): async def look_at(self, ctx, look=False):
metadata = "" metadata = ""
if look: if look:
url = os.getenv('stablediffusion_url') url = os.getenv('stablediffusion_url')
@ -90,7 +97,7 @@ async def look_at(ctx, look=False):
break break
out_file.write(chunk) out_file.write(chunk)
img_link = my_open_img_file(imageName) img_link = self.my_open_img_file(imageName)
try: try:
payload = {"image": img_link} payload = {"image": img_link}
@ -100,17 +107,17 @@ async def look_at(ctx, look=False):
description = description.split(',')[0] description = description.split(',')[0]
metadata += f"<image:{description}>\n" metadata += f"<image:{description}>\n"
except aiohttp.ClientError as error: except aiohttp.ClientError as error:
await 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() await http_session.close()
return metadata return metadata
@commands.command( @commands.command(
description="Draw", description="Draw",
help="Generates a picture using stable diffusion and gpt 3.5. It generates a list of 10 random artistic words and feeds them into stable diffusion. Usage: !draw (amount of pictures)", help="Generates a picture using stable diffusion and gpt 3.5. It generates a list of 10 random artistic words and feeds them into stable diffusion. Usage: !draw (amount of pictures)",
brief="Generate a random image" brief="Generate a random image"
) )
async def draw(ctx): async def draw(self, ctx):
url = os.getenv('stablediffusion_url') url = os.getenv('stablediffusion_url')
if url == "disabled": if url == "disabled":
return return
@ -130,7 +137,7 @@ async def draw(ctx):
choice4 = "Describe a unique character and an environment in one sentence" choice4 = "Describe a unique character and an environment in one sentence"
choice5 = "Describe a nonhuman character and an environment in one sentence" choice5 = "Describe a nonhuman character and an environment in one sentence"
prompt = random.choice([choice1,choice2,choice3,choice4,choice5]) prompt = random.choice([choice1,choice2,choice3,choice4,choice5])
prompt = await answer_question(prompt) prompt = await self.answer_question(prompt)
if random.randint(0,9): if random.randint(0,9):
prompt = prompt.replace("abstract, ", "") prompt = prompt.replace("abstract, ", "")
prompt = prompt.replace("AI, ", "") prompt = prompt.replace("AI, ", "")
@ -166,19 +173,19 @@ async def draw(ctx):
await ctx.send(file=f) await ctx.send(file=f)
await ctx.send(prompt) await ctx.send(prompt)
except Exception as error: except Exception as error:
await 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.")
except Exception as error: except Exception as error:
await 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 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'): async def change_model(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"),
@ -192,7 +199,7 @@ async def change_model(ctx, model_choice='0'):
await ctx.send("This command is currently disabled") await ctx.send("This command is currently disabled")
return return
http_session = aiohttp.ClientSession() 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: async with http_session.get(url=f'{url}/sdapi/v1/options') as response:
config_json = await response.json() config_json = await response.json()
@ -216,12 +223,12 @@ async def change_model(ctx, model_choice='0'):
await http_session.close() await http_session.close()
await ctx.send(output) await ctx.send(output)
@commands.command( @commands.command(
description="Lora", description="Lora",
help="List the stable diffusion loras.", help="List the stable diffusion loras.",
brief="List the stable diffusion loras" brief="List the stable diffusion loras"
) )
async def lora(ctx): async def lora(self, ctx):
lora_choices = { lora_choices = {
'0': ("Lora Name", "Trigger Words"), '0': ("Lora Name", "Trigger Words"),
'1': ("<lora:rebecca:1>", "rebecca (cyberpunk)"), '1': ("<lora:rebecca:1>", "rebecca (cyberpunk)"),
@ -234,12 +241,12 @@ async def lora(ctx):
output += lora_options output += lora_options
await ctx.send(output) await ctx.send(output)
@commands.command( @commands.command(
description="Imagine", description="Imagine",
help="Generate an image using stable diffusion. You can add keyword arguments to your prompt and they will be treated as stable diffusion options. Usage !imagine (topic)", help="Generate an image using stable diffusion. You can add keyword arguments to your prompt and they will be treated as stable diffusion options. Usage !imagine (topic)",
brief="Generate an image" brief="Generate an image"
) )
async def imagine(ctx): async def imagine(self, ctx):
url = os.getenv('stablediffusion_url') url = os.getenv('stablediffusion_url')
if url == "disabled": if url == "disabled":
await ctx.send("Command is currently disabled.") await ctx.send("Command is currently disabled.")
@ -247,7 +254,7 @@ async def imagine(ctx):
else: else:
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 = 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 = "databases/negative_prompt.txt"
with open(neg_prompt_file, 'r') as f: with open(neg_prompt_file, 'r') as f:
negative_prompt = f.readline() negative_prompt = f.readline()
@ -268,7 +275,7 @@ async def imagine(ctx):
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.")
await handle_error(error) await self.handle_error(error)
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])))
@ -279,7 +286,7 @@ async def imagine(ctx):
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.")
await handle_error(error) await self.handle_error(error)
pnginfo = PngImagePlugin.PngInfo() pnginfo = PngImagePlugin.PngInfo()
pnginfo.add_text("parameters", response2.get("info")) pnginfo.add_text("parameters", response2.get("info"))
@ -304,12 +311,12 @@ async def imagine(ctx):
await http_session.close() await http_session.close()
@commands.command( @commands.command(
description="Describe", description="Describe",
help="Get better understanding of what the bot \"sees\" when you post an image! (Runs it through CLIP) Usage !describe (image link)", help="Get better understanding of what the bot \"sees\" when you post an image! (Runs it through CLIP) Usage !describe (image link)",
brief="Describe image" brief="Describe image"
) )
async def describe(ctx): async def describe(self, ctx):
url = os.getenv('stablediffusion_url') url = os.getenv('stablediffusion_url')
if url == "disabled": if url == "disabled":
await ctx.send("Command is currently disabled") await ctx.send("Command is currently disabled")
@ -325,7 +332,7 @@ async def describe(ctx):
print("No image linked or attached.") print("No image linked or attached.")
return return
except Exception as error: except Exception as error:
await handle_error(error) await self.handle_error(error)
print("Couldn't find image.") print("Couldn't find image.")
return return
@ -340,7 +347,7 @@ async def describe(ctx):
break break
out_file.write(chunk) out_file.write(chunk)
img_link = 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 http_session.post(url, json=payload) as response:
@ -348,16 +355,16 @@ async def describe(ctx):
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 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() await http_session.close()
@commands.command( @commands.command(
description="Reimagine", description="Reimagine",
help="Reimagine an image as something else. One example is reimagining a picture as anime. This command can be hard to use. \nUsage: !reimagine (image link) (topic)\nExample: !reimagine (image link) anime", help="Reimagine an image as something else. One example is reimagining a picture as anime. This command can be hard to use. \nUsage: !reimagine (image link) (topic)\nExample: !reimagine (image link) anime",
brief="Reimagine an image" brief="Reimagine an image"
) )
async def reimagine(ctx): async def reimagine(self, ctx):
url = os.getenv('stablediffusion_url') url = os.getenv('stablediffusion_url')
if url == "disabled": if url == "disabled":
await ctx.send("Command is currently disabled") await ctx.send("Command is currently disabled")
@ -373,11 +380,11 @@ async def reimagine(ctx):
print("No image linked or attached.") print("No image linked or attached.")
return return
except Exception as error: except Exception as error:
await handle_error(error) await self.handle_error(error)
print("Couldn't find image.") print("Couldn't find image.")
return return
key_value_pairs, prompt = extract_key_value_pairs(prompt) key_value_pairs, prompt = self.extract_key_value_pairs(prompt)
http_session = aiohttp.ClientSession() http_session = aiohttp.ClientSession()
try: try:
async with http_session.get(file_url) as response: async with http_session.get(file_url) as response:
@ -392,9 +399,9 @@ async def reimagine(ctx):
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 handle_error(error) await self.handle_error(error)
img_link = my_open_img_file(imageName) img_link = self.my_open_img_file(imageName)
negative_prompt = "badhandsv4, worst quality, lowres, EasyNegative, hermaphrodite, cropped, not in the frame, additional faces, jpeg large artifacts, jpeg small artifacts, ugly, tiling, poorly drawn hands, poorly drawn feet, poorly drawn face, out of frame, extra limbs, disfigured, deformed, body out of frame, blurry, bad anatomy, blurred, watermark, grainy, signature, cut off, draft, not finished drawing, unfinished image, bad eyes, doll, 3d, cartoon, (bad eyes:1.2), (worst quality:1.2), (low quality:1.2), bad-image-v2-39000, (bad_prompt_version2:0.8), nude, badhandv4 By bad artist -neg easynegative ng_deepnegative_v1_75t verybadimagenegative_v1.3, (Worst Quality, Low Quality:1.4), Poorly Made Bad 3D, Lousy Bad Realistic, nsfw,(worst quality, low quality:1.4), (lip, nose, tooth, rouge, lipstick, eyeshadow:1.4), ( jpeg artifacts:1.4), (depth of field, bokeh, blurry, film grain, chromatic aberration, lens flare:1.0), (1boy, abs, muscular, rib:1.0), greyscale, monochrome, dusty sunbeams, trembling, motion lines, motion blur, emphasis lines, text, title, logo, signature, child, childlike, young, easynegative, (bad-hands-5:0.8), plain background, monochrome, poorly drawn face, poorly drawn hands, watermark, censored, (mutated hands and fingers), ugly, worst quality, low quality,, nsfw,(worst quality, low quality:1.4), (lip, nose, tooth, rouge, lipstick, eyeshadow:1.4), ( jpeg artifacts:1.4), (depth of field, bokeh, blurry, film grain, chromatic aberration, lens flare:1.0), (1boy, abs, muscular, rib:1.0), greyscale, monochrome, dusty sunbeams, trembling, motion lines, motion blur, emphasis lines, text, title, logo, signature, child, childlike, young" negative_prompt = "badhandsv4, worst quality, lowres, EasyNegative, hermaphrodite, cropped, not in the frame, additional faces, jpeg large artifacts, jpeg small artifacts, ugly, tiling, poorly drawn hands, poorly drawn feet, poorly drawn face, out of frame, extra limbs, disfigured, deformed, body out of frame, blurry, bad anatomy, blurred, watermark, grainy, signature, cut off, draft, not finished drawing, unfinished image, bad eyes, doll, 3d, cartoon, (bad eyes:1.2), (worst quality:1.2), (low quality:1.2), bad-image-v2-39000, (bad_prompt_version2:0.8), nude, badhandv4 By bad artist -neg easynegative ng_deepnegative_v1_75t verybadimagenegative_v1.3, (Worst Quality, Low Quality:1.4), Poorly Made Bad 3D, Lousy Bad Realistic, nsfw,(worst quality, low quality:1.4), (lip, nose, tooth, rouge, lipstick, eyeshadow:1.4), ( jpeg artifacts:1.4), (depth of field, bokeh, blurry, film grain, chromatic aberration, lens flare:1.0), (1boy, abs, muscular, rib:1.0), greyscale, monochrome, dusty sunbeams, trembling, motion lines, motion blur, emphasis lines, text, title, logo, signature, child, childlike, young, easynegative, (bad-hands-5:0.8), plain background, monochrome, poorly drawn face, poorly drawn hands, watermark, censored, (mutated hands and fingers), ugly, worst quality, low quality,, nsfw,(worst quality, low quality:1.4), (lip, nose, tooth, rouge, lipstick, eyeshadow:1.4), ( jpeg artifacts:1.4), (depth of field, bokeh, blurry, film grain, chromatic aberration, lens flare:1.0), (1boy, abs, muscular, rib:1.0), greyscale, monochrome, dusty sunbeams, trembling, motion lines, motion blur, emphasis lines, text, title, logo, signature, child, childlike, young"
@ -422,15 +429,15 @@ async def reimagine(ctx):
await ctx.send(file=f) await ctx.send(file=f)
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 handle_error(error) await self.handle_error(error)
await http_session.close() await http_session.close()
@commands.command( @commands.command(
description="Negative Prompt", description="Negative Prompt",
help="Changes the negative prompt for imagine across all channels", help="Changes the negative prompt for imagine across all channels",
brief="Change the negative prompt for imagine" brief="Change the negative prompt for imagine"
) )
async def negative_prompt(ctx, *args): async def negative_prompt(self, ctx, *args):
message = ' '.join(args) message = ' '.join(args)
if not message: if not message:
message = "easynegative, badhandv4, verybadimagenegative_v1.3" message = "easynegative, badhandv4, verybadimagenegative_v1.3"
@ -441,10 +448,4 @@ async def negative_prompt(ctx, *args):
async def setup(bot): async def setup(bot):
bot.add_command(imagine) bot.add_cog(StableDiffusion)
bot.add_command(reimagine)
bot.add_command(describe)
bot.add_command(change_model)
bot.add_command(lora)
bot.add_command(draw)
bot.add_command(negative_prompt)