Bot improvements

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
+class Settings:
-steps = 28
+    def __init__(self, nsfw_enabled=True, num_steps=28, ai_seed=-1, url="https://art.jurydoak.com", styles=["Bot"]):
-seed = -1
-url = "https://art.jurydoak.com"
 
-async def gen_image(message):
+        self.nsfw_enabled = nsfw_enabled
-    # global sess
+        self.num_steps = num_steps
-    
+        self.ai_seed = ai_seed
-    sess = aiohttp.ClientSession(url)
+        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",
+        "prompt": prompt,
-        "styles": ["Bot"],
+        "styles": main_settings.styles,
-        "steps": steps,
+        "steps": main_settings.num_steps,
-        "seed": seed,
+        "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:])
+    image_data = await send_request(ctx, endpoint="/sdapi/v1/txt2img", payload=payload)
-    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()
+    if image_data['images'][0] is None:
-    upload_file = discord.File("/tmp/image.png", filename="image.png")
+        await ctx.reply("Something went wrong, please report this to the admin so it can be ignored")
-    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'):
+    with open('/tmp/image.png', 'wb') as f:
-        print(message.content)
+        f.write(b64decode(image_data['images'][0]))
-        await message.channel.send('Hopefully queued the image')
-        await gen_image(message)
 
-    if message.content.startswith('$degen'):
+    embed = discord.Embed()
-        await message.channel.send('Disabling kink mode' if nsfw_enabled else 'Enabling kink mode')
+    upload_file = discord.File("/tmp/image.png", filename="image.png")
-        if nsfw_enabled:
+    embed.set_image(url="attachment://image.png")
-            nsfw_enabled = False
+    # embed.title = prompt
-        else:
-            nsfw_enabled = True
 
-    if message.content.startswith("$steps"):
+    await ctx.reply("", file=upload_file, embed=embed)
-        global steps
+
-        try:
+
-            num = int(message.content.split(" ")[1])
+@bot.command()
-            steps = num
+async def seed(ctx, arg):
-            await message.channel.send(f"Setting steps to {steps}")
+    '''Set the seed for the image generation'''
-        except:
+    try:
-            pass
+        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"):
+#     if message.content.startswith("$seed"):
-        global seed
+#         global seed
-        try:
+#         try:
-            num = int(message.content.split(" ")[1])
+#             num = int(message.content.split(" ")[1])
-            seed = num
+#             seed = num
-            await message.channel.send(f"Setting seed to {seed}")
+#             await message.channel.send(f"Setting seed to {seed}")
-        except:
+#         except:
-            pass
+#             pass
 
-    if message.content.startswith("$host"):
+#     if message.content.startswith("$host"):
-        global url
+#         global url
-        try:
+#         try:
-            host = message.content.split(" ")[1]
+#             host = message.content.split(" ")[1]
-            url = host
+#             url = host
-            await message.channel.send(f"Setting host to {host}")
+#             await message.channel.send(f"Setting host to {host}")
-        except:
+#         except:
-            pass
+#             pass
 
-client.run(discord_client_token)
+# client.run(discord_client_token)
+
+bot.run(discord_client_token)