Bot improvements
This commit is contained in:
parent
aaee407af4
commit
e9338fe6b9
322
weeeabot.py
322
weeeabot.py
@ -1,9 +1,7 @@
|
||||
import discord
|
||||
import requests
|
||||
from discord.ext import commands
|
||||
import aiohttp
|
||||
import asyncio
|
||||
from base64 import b64decode
|
||||
import json
|
||||
from sys import exit
|
||||
|
||||
try:
|
||||
@ -16,23 +14,61 @@ except FileNotFoundError:
|
||||
intents = discord.Intents.default()
|
||||
intents.message_content = True
|
||||
|
||||
client = discord.Client(intents=intents)
|
||||
# client = discord.Client(intents=intents)
|
||||
|
||||
nsfw_enabled = True
|
||||
steps = 28
|
||||
seed = -1
|
||||
url = "https://art.jurydoak.com"
|
||||
class Settings:
|
||||
def __init__(self, nsfw_enabled=True, num_steps=28, ai_seed=-1, url="https://art.jurydoak.com", styles=["Bot"]):
|
||||
|
||||
async def gen_image(message):
|
||||
# global sess
|
||||
self.nsfw_enabled = nsfw_enabled
|
||||
self.num_steps = num_steps
|
||||
self.ai_seed = ai_seed
|
||||
self.url = url
|
||||
self.styles = styles
|
||||
|
||||
sess = aiohttp.ClientSession(url)
|
||||
|
||||
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": "cute girlhand on cheek eyes closed thinking hard, sfw",
|
||||
"styles": ["Bot"],
|
||||
"steps": steps,
|
||||
"seed": seed,
|
||||
"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",
|
||||
@ -40,7 +76,7 @@ async def gen_image(message):
|
||||
}
|
||||
|
||||
settings = {
|
||||
"filter_nsfw": not nsfw_enabled,
|
||||
"filter_nsfw": not main_settings.nsfw_enabled,
|
||||
"samples_save": True,
|
||||
}
|
||||
|
||||
@ -49,89 +85,189 @@ async def gen_image(message):
|
||||
}
|
||||
|
||||
payload.update(override_payload)
|
||||
user_prompt = " ".join(message.content.split(' ')[1:])
|
||||
payload['prompt'] = user_prompt
|
||||
# url = "https://art.jurydoak.com/sdapi/v1"
|
||||
image_data = await send_request(ctx, endpoint="/sdapi/v1/txt2img", payload=payload)
|
||||
|
||||
image_req = await sess.post("/sdapi/v1/txt2img", json=payload)
|
||||
x = await image_req.json()
|
||||
image_req.close()
|
||||
await sess.close()
|
||||
|
||||
image = x['images'][0]
|
||||
# image = requests.post(f"{url}/txt2img", json=payload).json()['images'][0]
|
||||
with open('/tmp/image.png', 'wb') as f:
|
||||
f.write(b64decode(image))
|
||||
|
||||
e = discord.Embed()
|
||||
upload_file = discord.File("/tmp/image.png", filename="image.png")
|
||||
e.set_image(url="attachment://image.png")
|
||||
e.title = user_prompt
|
||||
await message.channel.send("", file=upload_file, embed=e)
|
||||
|
||||
|
||||
@client.event
|
||||
async def on_ready():
|
||||
|
||||
bot_setup_id = 775489296932405319
|
||||
channel = client.get_channel(bot_setup_id)
|
||||
await channel.send(f"Awake and ready to degen, NSFW: {nsfw_enabled}, STEPS: {steps}")
|
||||
await client.change_presence(activity=discord.Game(name="Thinking degenerate thoughts..."))
|
||||
|
||||
# text_channel_list = []
|
||||
# for server in discord.Client.servers:
|
||||
# for channel in server.channels:
|
||||
# if channel.type == "text":
|
||||
# text_channel_list.append(channel)
|
||||
|
||||
# print(text_channel_list)
|
||||
print(f'We have logged in as {client.user}')
|
||||
|
||||
|
||||
@client.event
|
||||
async def on_message(message):
|
||||
global nsfw_enabled
|
||||
if message.author == client.user:
|
||||
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
|
||||
|
||||
if message.content.startswith('$hello'):
|
||||
print(message.content)
|
||||
await message.channel.send('Hopefully queued the image')
|
||||
await gen_image(message)
|
||||
with open('/tmp/image.png', 'wb') as f:
|
||||
f.write(b64decode(image_data['images'][0]))
|
||||
|
||||
if message.content.startswith('$degen'):
|
||||
await message.channel.send('Disabling kink mode' if nsfw_enabled else 'Enabling kink mode')
|
||||
if nsfw_enabled:
|
||||
nsfw_enabled = False
|
||||
else:
|
||||
nsfw_enabled = True
|
||||
embed = discord.Embed()
|
||||
upload_file = discord.File("/tmp/image.png", filename="image.png")
|
||||
embed.set_image(url="attachment://image.png")
|
||||
# embed.title = prompt
|
||||
|
||||
if message.content.startswith("$steps"):
|
||||
global steps
|
||||
try:
|
||||
num = int(message.content.split(" ")[1])
|
||||
steps = num
|
||||
await message.channel.send(f"Setting steps to {steps}")
|
||||
except:
|
||||
pass
|
||||
await ctx.reply("", file=upload_file, embed=embed)
|
||||
|
||||
|
||||
if message.content.startswith("$seed"):
|
||||
global seed
|
||||
try:
|
||||
num = int(message.content.split(" ")[1])
|
||||
seed = num
|
||||
await message.channel.send(f"Setting seed to {seed}")
|
||||
except:
|
||||
pass
|
||||
@bot.command()
|
||||
async def seed(ctx, arg):
|
||||
'''Set the seed for the image generation'''
|
||||
try:
|
||||
arg = int(arg)
|
||||
except:
|
||||
pass
|
||||
|
||||
if message.content.startswith("$host"):
|
||||
global url
|
||||
try:
|
||||
host = message.content.split(" ")[1]
|
||||
url = host
|
||||
await message.channel.send(f"Setting host to {host}")
|
||||
except:
|
||||
pass
|
||||
# global ai_seed
|
||||
main_settings.ai_seed = arg
|
||||
|
||||
client.run(discord_client_token)
|
||||
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")
|
||||
|
||||
|
||||
# async def gen_image(message):
|
||||
# # global sess
|
||||
|
||||
# sess = aiohttp.ClientSession(url)
|
||||
|
||||
# payload = {
|
||||
# "prompt": "cute girlhand on cheek eyes closed thinking hard, sfw",
|
||||
# "styles": ["Bot"],
|
||||
# "steps": steps,
|
||||
# "seed": seed,
|
||||
# "n_iter": 1,
|
||||
# "height": 1024,
|
||||
# "negative_prompts": "nsfw, not safe for work, nudity, multiple keyboards",
|
||||
# "cfg_scale": 12
|
||||
# }
|
||||
|
||||
# settings = {
|
||||
# "filter_nsfw": not nsfw_enabled,
|
||||
# "samples_save": True,
|
||||
# }
|
||||
|
||||
# override_payload = {
|
||||
# "override_settings": settings
|
||||
# }
|
||||
|
||||
# payload.update(override_payload)
|
||||
# user_prompt = " ".join(message.content.split(' ')[1:])
|
||||
# payload['prompt'] = user_prompt
|
||||
# # url = "https://art.jurydoak.com/sdapi/v1"
|
||||
|
||||
# image_req = await sess.post("/sdapi/v1/txt2img", json=payload)
|
||||
# x = await image_req.json()
|
||||
# image_req.close()
|
||||
# await sess.close()
|
||||
|
||||
# image = x['images'][0]
|
||||
# # image = requests.post(f"{url}/txt2img", json=payload).json()['images'][0]
|
||||
# with open('/tmp/image.png', 'wb') as f:
|
||||
# f.write(b64decode(image))
|
||||
|
||||
# e = discord.Embed()
|
||||
# upload_file = discord.File("/tmp/image.png", filename="image.png")
|
||||
# e.set_image(url="attachment://image.png")
|
||||
# e.title = user_prompt
|
||||
# await message.channel.send("", file=upload_file, embed=e)
|
||||
|
||||
|
||||
# @client.event
|
||||
# async def on_ready():
|
||||
|
||||
# bot_setup_id = 775489296932405319
|
||||
# channel = client.get_channel(bot_setup_id)
|
||||
# # await channel.send(f"Awake and ready to degen, NSFW: {nsfw_enabled}, STEPS: {steps}")
|
||||
# await client.change_presence(activity=discord.Game(name="Thinking degenerate thoughts..."))
|
||||
|
||||
# # text_channel_list = []
|
||||
# # for server in discord.Client.servers:
|
||||
# # for channel in server.channels:
|
||||
# # if channel.type == "text":
|
||||
# # text_channel_list.append(channel)
|
||||
|
||||
# # print(text_channel_list)
|
||||
# print(f'We have logged in as {client.user}')
|
||||
|
||||
|
||||
# @client.event
|
||||
# async def on_message(message):
|
||||
# global nsfw_enabled
|
||||
# if message.author == client.user:
|
||||
# return
|
||||
|
||||
# if message.content.startswith('$hello'):
|
||||
# print(message.content)
|
||||
# await message.channel.send('Hopefully queued the image')
|
||||
# await gen_image(message)
|
||||
|
||||
# if message.content.startswith('$degen'):
|
||||
# await message.channel.send('Disabling kink mode' if nsfw_enabled else 'Enabling kink mode')
|
||||
# if nsfw_enabled:
|
||||
# nsfw_enabled = False
|
||||
# else:
|
||||
# nsfw_enabled = True
|
||||
|
||||
# if message.content.startswith("$steps"):
|
||||
# global steps
|
||||
# try:
|
||||
# num = int(message.content.split(" ")[1])
|
||||
# steps = num
|
||||
# await message.channel.send(f"Setting steps to {steps}")
|
||||
# except:
|
||||
# pass
|
||||
|
||||
|
||||
# if message.content.startswith("$seed"):
|
||||
# global seed
|
||||
# try:
|
||||
# num = int(message.content.split(" ")[1])
|
||||
# seed = num
|
||||
# await message.channel.send(f"Setting seed to {seed}")
|
||||
# except:
|
||||
# pass
|
||||
|
||||
# if message.content.startswith("$host"):
|
||||
# global url
|
||||
# try:
|
||||
# host = message.content.split(" ")[1]
|
||||
# url = host
|
||||
# await message.channel.send(f"Setting host to {host}")
|
||||
# except:
|
||||
# pass
|
||||
|
||||
# client.run(discord_client_token)
|
||||
|
||||
bot.run(discord_client_token)
|
||||
Loading…
x
Reference in New Issue
Block a user