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") # 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)