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

Push me

Browse Source 
And then just touch me
So I can get my
Satisfaction
This commit is contained in:
Benjamyn Love 2022-12-11 19:17:35 +11:00
parent f8879974b6
commit 6446cdb86b
3 changed files with 115 additions and 110 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

BIN
images/404.jpg Normal file
View File

Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 78 KiB

61
modules/diffuseapi.py
View File

@ -1,5 +1,7 @@
import aiohttp
from base64 import b64decode
from sys import exit
from pprint import pprint
class DiffuseAPI():
def __init__(self, url, styles, nsfw_enabled=True, num_steps=28):
@ -40,7 +42,7 @@ class DiffuseAPI():
return False
except:
return False
def set_styles(self, styles):
if type(styles) == list:
self.styles = styles
@ -64,20 +66,20 @@ class DiffuseAPI():
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
@ -95,7 +97,8 @@ class DiffuseAPI():
}
settings = {
"filter_nsfw": not self.nsfw_enabled
"filter_nsfw": not self.nsfw_enabled,
"enable_pnginfo": False
}
override_payload = {
@ -104,18 +107,36 @@ class DiffuseAPI():
payload.update(override_payload)
sess = aiohttp.ClientSession(self.url)
alive = await sess.head('/')
async with aiohttp.ClientSession(self.url) as session:
async with session.head('/') as alive:
if alive.status != 200:
return None
async with session.post("/sdapi/v1/txt2img", json=payload) as image_json:
image_data = await image_json.json()
return image_data["images"][0]
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
async def generate_upscale(self, image):
payload = {
"resize_mode": 0,
"show_extras_results": True,
"gfpgan_visibility": 0,
"codeformer_visibility": 0,
"codeformer_weight": 0,
"upscaling_resize": 4,
"upscaling_resize_w": 512,
"upscaling_resize_h": 1024,
"upscaling_crop": True,
"upscaler_1": "R-ESRGAN 4x+ Anime6B",
"upscaler_2": "None",
"extras_upscaler_2_visibility": 0,
"upscale_first": False,
"image": image
}
async with aiohttp.ClientSession(self.url) as session:
async with session.head('/') as alive:
if alive.status != 200:
return None
async with session.post("/sdapi/v1/extra-single-image", json=payload) as image_json:
image_data = await image_json.json()
return image_data["image"]

164
weeeabot.py
View File

@ -1,9 +1,11 @@
import discord
from discord.ext import commands
import aiohttp
from base64 import b64decode
from base64 import b64decode, b64encode
from sys import exit
from modules import diffuseapi
import io
import hashlib
try:
with open("token.secret", 'r') as f:
@ -15,10 +17,39 @@ except FileNotFoundError:
intents = discord.Intents.default()
intents.message_content = True
api = diffuseapi.DiffuseAPI("http://localhost:7860", ["default"], True, 28)
api = diffuseapi.DiffuseAPI("https://art.jurydoak.com", ["Bot"], False, 28)
# client = discord.Client(intents=intents)
class Confirm(discord.ui.View):
def __init__(self):
super().__init__()
self.value = None
async def callback(self, interaction: discord.Interaction):
await interaction.response.send_message("You wish to upscale this image!")
# When the confirm button is pressed, set the inner value to `True` and
# stop the View from listening to more input.
# We also send the user an ephemeral message that we're confirming their choice.
@discord.ui.button(label='Upscale', style=discord.ButtonStyle.green)
async def confirm(self, interaction: discord.Interaction, button: discord.ui.Button):
image_data = await interaction.message.attachments[0].read()
await interaction.response.defer()
upscaled_image = await api.generate_upscale(b64encode(image_data).decode('utf-8'))
data = io.BytesIO(b64decode(upscaled_image))
await interaction.followup.send("", file=discord.File(data, "upscaled.png"))
self.stop()
# This one is similar to the confirmation button except sets the inner value to `False`
@discord.ui.button(label='Delete', style=discord.ButtonStyle.grey)
async def cancel(self, interaction: discord.Interaction, button: discord.ui.Button):
await interaction.message.delete()
self.stop()
class Settings:
def __init__(self, nsfw_enabled=True, num_steps=28, ai_seed=-1, url="https://art.jurydoak.com", styles=["Bot"]):
@ -30,79 +61,32 @@ class Settings:
main_settings = Settings()
activity = discord.Activity(type=discord.ActivityType.listening, name="!help")
activity = discord.Activity(type=discord.ActivityType.listening, name="!help main")
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:
image_data = await api.generate_image(prompt)
view = Confirm()
if image_data is None:
await ctx.reply("Something went wrong, please report this to the admin so it can be ignored")
return
decoded_image = b64decode(image_data)
h = hashlib.md5()
h.update(decoded_image)
digest = h.digest()
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)
if digest == b'i\xac`\xde\xbak\xba\xab{2Z\xcc\tK\xc2~':
await ctx.reply("Were no stranger to lewds, but you know the rules, and so do I", file=discord.File("images/404.jpg", "404.jpg"))
return
print(h.digest())
data = io.BytesIO(decoded_image)
await ctx.reply("", file=discord.File(data, "_".join(args) + ".png"), view=view)
@bot.command()
@ -114,9 +98,9 @@ async def seed(ctx, arg):
pass
# global ai_seed
main_settings.ai_seed = arg
api.set_seed(arg)
await ctx.reply(f"I have updated the seed to {main_settings.ai_seed} for you my master.")
await ctx.reply(f"I have updated the seed to {api.seed} for you my master.")
@bot.command()
async def steps(ctx, arg):
@ -131,9 +115,9 @@ async def steps(ctx, arg):
return
# global num_steps
main_settings.num_steps = arg
api.set_steps(arg)
await ctx.reply(f"I have updated the steps to {main_settings.num_steps} for you my master.")
await ctx.reply(f"I have updated the steps to {api.num_steps} for you my master.")
@bot.command()
async def settings(ctx):
@ -141,37 +125,37 @@ async def settings(ctx):
global ai_seed, steps
settings = f"""
```
seed: {main_settings.ai_seed}
steps: {main_settings.num_steps}
seed: {api.seed}
steps: {api.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")
async def upscale(ctx):
"""Upscale the attached image"""
orig_image_data = await ctx.message.attachments[0].read()
new_data = await api.generate_upscale(b64encode(orig_image_data).decode('utf-8'))
data = io.BytesIO(b64decode(new_data))
await ctx.reply("", file=discord.File(data, "upscaled_img.png"))
@bot.command()
async def test2(ctx):
image_data = await api.generate_image(prompt="cute bot doing bot things")
if image_data is None:
return
with open('/tmp/image.png', 'wb') as f:
f.write(b64decode(image_data))
async def test(ctx):
'''Test function, currently changes the URL and bot settings'''
if api.url != "http://localhost:7860":
api.url = "http://localhost:7860"
api.styles = ["default"]
activity = discord.Activity(type=discord.ActivityType.listening, name="!help fast")
await bot.change_presence(activity=activity)
await ctx.reply("Set to fastboi")
else:
api.url = "https://art.jurydoak.com"
api.styles = ["Bot"]
activity = discord.Activity(type=discord.ActivityType.listening, name="!help main")
await bot.change_presence(activity=activity)
await ctx.reply("Set to main api")
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.run(discord_client_token)