diff --git a/weeeabot.py b/weeeabot.py index 71dab05..0a74c39 100644 --- a/weeeabot.py +++ b/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 - - sess = aiohttp.ClientSession(url) - + 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": "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_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)) + image_data = await send_request(ctx, endpoint="/sdapi/v1/txt2img", payload=payload) - 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) + + +@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") + + +# 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("$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 +# 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) +# client.run(discord_client_token) + +bot.run(discord_client_token) \ No newline at end of file