Home/weeabot
2
1
Fork 0
Code Issues 1 Pull Requests Projects Releases Wiki Activity

Initial commit

Browse Source
This commit is contained in:
jordan 2022-12-10 22:05:55 +11:00
parent 1c31b73b2a
commit 991c592086
3 changed files with 470 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

128
DiffuseAPI.py Normal file
View File

@ -0,0 +1,128 @@
import aiohttp
from base64 import b64decode
from sys import exit
class DiffuseAPI():
def __init__(self, url, styles, nsfw_enabled=True, num_steps=28):
self.nsfw_enabled = nsfw_enabled
self.num_steps = num_steps
self.url = url
self.styles = styles
self.seed = -1
self.width = 512
self.height = 1024
self.cfg_scale = 12
def set_steps(self, steps):
try:
new_steps = int(steps)
if 0 > new_steps < 50:
self.num_steps = new_steps
return True
return False
except:
return False
def set_seed(self, seed):
try:
new_seed = int(seed)
self.seed = new_seed
return True
except:
return False
def set_cfg_scale(self, scale):
try:
new_scale = int(scale)
if 0 > new_scale < 30:
self.cfg_scale = new_scale
return True
return False
except:
return False
def set_styles(self, styles):
if type styles = list:
self.styles = styles
return True
return False
def set_orientation(self, orientation):
new_orientation = str(orientation)
match new_orientation:
case "portrait":
self.width = 512
self.height = 1024
return True
case "landscape":
self.width = 1024
self.height = 512
return True
case "square":
self.width = 512
self.height = 512
return True
case _:
return False
def get_orientation(self):
if self.width == 512 and self.height == 512:
return ("square", self.width, self.height)
elif self.width == 1024:
return ("landscape", self.width, self.height)
return ("portrait", self.width, self.height)
def set_nsfw_filter(self, filter_state):
if type filter_state = bool:
self.nsfw_enabled = filter_state
return True
return False
def get_nsfw_filter(self):
return self.nsfw_enabled
async def generate_image(self, prompt, neg_prompt=""):
payload = {
"prompt": prompt,
"styles": main_settings.styles,
"steps": main_settings.num_steps,
"seed": main_settings.ai_seed,
"n_iter": 1,
"height": self.height,
"width": self.width,
"negative_prompts": neg_prompt,
"cfg_scale": self.cfg_scale
}
settings = {
"filter_nsfw": not self.nsfw_enabled
}
override_payload = {
"override_settings": settings
}
payload.update(override_payload)
sess = aiohttp.ClientSession(main_settings.url)
alive = await sess.head('/')
if alive.status != 200:
alive.close()
await sess.close()
return None
request = await sess.post("/sdapi/v1/txt2img", json=payload)
try:
req_json = await request.json()
request.close()
await sess.close()
return req_json["images", [0]]
except:
return None

183
classweeb.ipynb Normal file
View File

@ -0,0 +1,183 @@
{
"cells": [
{
"cell_type": "code",
"execution_count": 1,
"metadata": {},
"outputs": [],
"source": [
"import aiohttp\n",
"from base64 import b64decode\n",
"from sys import exit"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"class DiffuseAPI():\n",
" def __init__(self, url, styles, nsfw_enabled=True, num_steps=28):\n",
"\n",
" self.nsfw_enabled = nsfw_enabled\n",
" self.num_steps = num_steps\n",
" self.url = url\n",
" self.styles = styles\n",
" self.seed = -1\n",
" self.width = 512\n",
" self.height = 1024\n",
" self.cfg_scale = 12\n",
"\n",
" def set_steps(self, steps):\n",
" try:\n",
" new_steps = int(steps)\n",
" if 0 > new_steps < 50:\n",
" self.num_steps = new_steps\n",
" return True\n",
" return False\n",
" except:\n",
" return False\n",
"\n",
" def set_seed(self, seed):\n",
" try:\n",
" new_seed = int(seed)\n",
" self.seed = new_seed\n",
" return True\n",
" except:\n",
" return False\n",
" \n",
" def set_cfg_scale(self, scale):\n",
" try:\n",
" new_scale = int(scale)\n",
" if 0 > new_scale < 30:\n",
" self.cfg_scale = new_scale\n",
" return True\n",
" return False\n",
" except:\n",
" return False\n",
" \n",
" def set_styles(self, styles):\n",
" if type styles = list:\n",
" self.styles = styles\n",
" return True\n",
" return False\n",
"\n",
" def set_orientation(self, orientation):\n",
" new_orientation = str(orientation)\n",
" match new_orientation:\n",
" case \"portrait\":\n",
" self.width = 512\n",
" self.height = 1024\n",
" return True\n",
" case \"landscape\":\n",
" self.width = 1024\n",
" self.height = 512\n",
" return True\n",
" case \"square\":\n",
" self.width = 512\n",
" self.height = 512\n",
" return True\n",
" case _:\n",
" return False\n",
" \n",
" def get_orientation(self):\n",
" if self.width == 512 and self.height == 512:\n",
" return (\"square\", self.width, self.height)\n",
" elif self.width == 1024:\n",
" return (\"landscape\", self.width, self.height)\n",
" return (\"portrait\", self.width, self.height)\n",
" \n",
" def set_nsfw_filter(self, filter_state):\n",
" if type filter_state = bool:\n",
" self.nsfw_enabled = filter_state\n",
" return True\n",
" return False\n",
" \n",
" def get_nsfw_filter(self):\n",
" return self.nsfw_enabled\n",
"\n",
" async def generate_image(self, prompt, neg_prompt=\"\"):\n",
" payload = {\n",
" \"prompt\": prompt,\n",
" \"styles\": main_settings.styles,\n",
" \"steps\": main_settings.num_steps,\n",
" \"seed\": main_settings.ai_seed,\n",
" \"n_iter\": 1,\n",
" \"height\": self.height,\n",
" \"width\": self.width,\n",
" \"negative_prompts\": neg_prompt,\n",
" \"cfg_scale\": self.cfg_scale\n",
" }\n",
"\n",
" settings = {\n",
" \"filter_nsfw\": not self.nsfw_enabled\n",
" }\n",
"\n",
" override_payload = {\n",
" \"override_settings\": settings\n",
" }\n",
"\n",
" payload.update(override_payload)\n",
"\n",
" sess = aiohttp.ClientSession(main_settings.url)\n",
" alive = await sess.head('/')\n",
"\n",
" if alive.status != 200:\n",
" alive.close()\n",
" await sess.close()\n",
" return None\n",
" request = await sess.post(\"/sdapi/v1/txt2img\", json=payload)\n",
" try:\n",
" req_json = await request.json()\n",
" request.close()\n",
" await sess.close()\n",
" return req_json[\"images\", [0]]\n",
" except:\n",
" return None\n",
"\n",
"\n",
" \n",
"\n",
" \n",
"\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"return image[\"images\", [0]]"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3.11.0 ('venv': venv)",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.11.0"
},
"orig_nbformat": 4,
"vscode": {
"interpreter": {
"hash": "aa19d1669203bff990bc2820f3654fa686baacd13bba0cbe075930f54e624f06"
}
}
},
"nbformat": 4,
"nbformat_minor": 2
}

159
test.py Normal file
View File

@ -0,0 +1,159 @@
import discord
from discord.ext import commands
import aiohttp
from base64 import b64decode
from sys import exit
try:
with open("token.secret", 'r') as f:
discord_client_token = f.read()
except FileNotFoundError:
print("Cannot locate token.secret please generate a token")
exit(69)
intents = discord.Intents.default()
intents.message_content = True
# client = discord.Client(intents=intents)
class Settings:
def __init__(self, nsfw_enabled=True, num_steps=28, ai_seed=-1, url="https://art.jurydoak.com", styles=["Bot"]):
self.nsfw_enabled = nsfw_enabled
self.num_steps = num_steps
self.ai_seed = ai_seed
self.url = url
self.styles = styles
main_settings = Settings()
activity = discord.Activity(type=discord.ActivityType.listening, name="!help")
bot = commands.Bot(intents=intents, command_prefix="!", activity=activity, help_command=commands.DefaultHelpCommand())
async def send_request(ctx, endpoint, payload):
global num_steps
global ai_seed
sess = aiohttp.ClientSession(main_settings.url)
alive = await sess.head('/')
if alive.status != 200:
await ctx.reply("I'm sorry but it appears my brain is offline")
alive.close()
await sess.close()
return {
'images': [None]
}
print(payload)
request = await sess.post(endpoint, json=payload)
try:
req_json = await request.json()
except:
print(await request.text())
return {
'images': [None]
}
request.close()
await sess.close()
return req_json
@bot.command()
async def prompt(ctx, *args):
'''Generate an image with the provided prompt'''
prompt = " ".join(args)
await ctx.reply("Generating your image boo")
payload = {
"prompt": prompt,
"styles": main_settings.styles,
"steps": main_settings.num_steps,
"seed": main_settings.ai_seed,
"n_iter": 1,
"height": 1024,
"negative_prompts": "nsfw, not safe for work, nudity, multiple keyboards",
"cfg_scale": 12
}
settings = {
"filter_nsfw": not main_settings.nsfw_enabled,
"samples_save": True,
}
override_payload = {
"override_settings": settings
}
payload.update(override_payload)
image_data = await send_request(ctx, endpoint="/sdapi/v1/txt2img", payload=payload)
if image_data['images'][0] is None:
await ctx.reply("Something went wrong, please report this to the admin so it can be ignored")
return
with open('/tmp/image.png', 'wb') as f:
f.write(b64decode(image_data['images'][0]))
embed = discord.Embed()
upload_file = discord.File("/tmp/image.png", filename="image.png")
embed.set_image(url="attachment://image.png")
# embed.title = prompt
await ctx.reply("", file=upload_file, embed=embed)
@bot.command()
async def seed(ctx, arg):
'''Set the seed for the image generation'''
try:
arg = int(arg)
except:
pass
# global ai_seed
main_settings.ai_seed = arg
await ctx.reply(f"I have updated the seed to {main_settings.ai_seed} for you my master.")
@bot.command()
async def steps(ctx, arg):
'''Set how many steps the AI will run (max 50)'''
try:
arg = int(arg)
except:
pass
if arg > 50:
await ctx.reply("I'm sorry Dave, I can't do that")
return
# global num_steps
main_settings.num_steps = arg
await ctx.reply(f"I have updated the steps to {main_settings.num_steps} for you my master.")
@bot.command()
async def settings(ctx):
'''See the currently configured settings (BROKEN)'''
global ai_seed, steps
settings = f"""
```
seed: {main_settings.ai_seed}
steps: {main_settings.num_steps}
```
"""
await ctx.message.channel.send(settings)
@bot.command()
async def test(ctx):
'''Test function, currently changes the URL and bot settings'''
if main_settings.url != "http://localhost:7860":
main_settings.url = "http://localhost:7860"
main_settings.styles = ["default"]
await ctx.reply("Set to fastboi")
else:
main_settings.url = "https://art.jurydoak.com"
main_settings.styles = ["Bot"]
await ctx.reply("Set to main api")
bot.run(discord_client_token)