diff --git a/.gitignore b/.gitignore index 8ad445e4..7213432b 100644 --- a/.gitignore +++ b/.gitignore @@ -3,9 +3,9 @@ __pycache__ pydata venv -core/generated/* outputs/*.png resources/stats.txt +resources/*.json .env start.bat \ No newline at end of file diff --git a/README.md b/README.md index cd8450dd..c9cc931c 100644 --- a/README.md +++ b/README.md @@ -18,6 +18,10 @@ To generate an image from text, use the /draw command and include your prompt as - img2img - denoising strength +##### Bonus features +- /settings command - set per server defaults for negative prompts, sampling steps, max steps, sampling method (see Notes). +- Stat showing how many /draw commands are used. + ### Setup requirements - Set up [AUTOMATIC1111's Stable Diffusion AI Web UI](https://github.com/AUTOMATIC1111/stable-diffusion-webui). - Run the Web UI as local host with api (`COMMANDLINE_ARGS= --listen --api`). @@ -31,6 +35,7 @@ TOKEN = put your bot token here #### Notes - Ensure your bot has `bot` and `application.commands` scopes when inviting to your Discord server, and intents are enabled. +- As /settings can be abused, consider reviewing who can access the command. This can be done through Apps -> Integrations in your Server Settings. - React to generated images with ❌ to delete them. - Optional .env variables you can set: ``` diff --git a/aiya.py b/aiya.py index 6a545085..cf43f91c 100644 --- a/aiya.py +++ b/aiya.py @@ -1,41 +1,37 @@ import asyncio import os import sys -from os.path import exists - import discord from dotenv import load_dotenv - from core.logging import get_logger +from core import settings + #start up initialization stuff self = discord.Bot() intents = discord.Intents.default() intents.members = True load_dotenv() -embed_color = discord.Colour.from_rgb(222, 89, 28) self.logger = get_logger(__name__) -file_exists = exists('resources/stats.txt') -if file_exists is False: - self.logger.info(f'stats.txt missing. Creating new file.') - with open('resources/stats.txt', 'w') as f: f.write('0') - +#load extensions self.load_extension('core.stablecog') -self.load_extension('core.tipscog') self.load_extension('core.settingscog') +self.load_extension('core.tipscog') #stats slash command @self.slash_command(name = 'stats', description = 'How many images has the bot generated?') async def stats(ctx): with open('resources/stats.txt', 'r') as f: data = list(map(int, f.readlines())) - embed = discord.Embed(title='Art generated', description=f'I have created {data[0]} pictures!', color=embed_color) + embed = discord.Embed(title='Art generated', description=f'I have created {data[0]} pictures!', color=settings.global_var.embed_color) await ctx.respond(embed=embed) @self.event async def on_ready(): self.logger.info(f'Logged in as {self.user.name} ({self.user.id})') await self.change_presence(activity=discord.Activity(type=discord.ActivityType.watching, name='drawing tutorials.')) + #check files and global variables + settings.files_check(self) #feature to delete generations. give bot 'Add Reactions' permission (or not, to hide the ❌) @self.event @@ -59,6 +55,11 @@ async def on_raw_reaction_add(ctx): #todo: I need to add a way for bot to delete DM generations pass +@self.event +async def on_guild_join(guild): + print(f'Wow, I joined {guild.name}! Refreshing settings.') + settings.files_check(self) + async def shutdown(bot): await bot.close() diff --git a/core/settings.py b/core/settings.py index 0ae4dd8b..2e53dd25 100644 --- a/core/settings.py +++ b/core/settings.py @@ -1,9 +1,12 @@ import os +from os.path import exists import json +import discord +self = discord.Bot() dir_path = os.path.dirname(os.path.realpath(__file__)) -path = '{}/generated/'.format(dir_path) +path = 'resources/'.format(dir_path) template = { "default_steps": 30, @@ -12,8 +15,13 @@ "max_steps": 30 } +#initialize global variables here +class GlobalVar: + url = "" + dir = "" + embed_color = discord.Colour.from_rgb(222, 89, 28) - +global_var = GlobalVar() def build(guild_id): settings = json.dumps(template) @@ -31,3 +39,39 @@ def update(guild_id:str, sett:str, value): settings[sett] = value with open(path + guild_id + '.json', 'w') as configfile: json.dump(settings, configfile) + +def files_check(self): + # creating files if they don't exist + file_exists = exists('resources/stats.txt') + if file_exists is False: + print(f'Uh oh, stats.txt missing. Creating a new one.') + with open('resources/stats.txt', 'w') as f: f.write('0') + + # guild settings files + for guild in self.guilds: + try: + read(str(guild.id)) + print(f'I\'m using local settings for {guild.id} a.k.a {guild}.') + except FileNotFoundError: + build(str(guild.id)) + print(f'Creating new settings file for {guild.id} a.k.a {guild}.') + + #check .env for URL and DIR. if they don't exist, ignore it and go with defaults. + if os.getenv("URL") == '': + global_var.url = os.environ.get('URL').rstrip("/") + print(f'Using URL: {global_var.url}') + else: + global_var.url = 'http://127.0.0.1:7860' + print('Using default URL: http://127.0.0.1:7860') + + if os.getenv("DIR") == '': + global_var.dir = os.environ.get('DIR') + print(f'Using outputs directory: {global_var.dir}') + else: + global_var.dir = "outputs" + print('Using default outputs directory: outputs') + #if directory in DIR doesn't exist, create it + dir_exists = os.path.exists(global_var.dir) + if dir_exists is False: + print(f'The folder for DIR doesn\'t exist! Creating folder at {global_var.dir}.') + os.mkdir(global_var.dir) \ No newline at end of file diff --git a/core/settingscog.py b/core/settingscog.py index bf2be1f5..fdc3a66e 100644 --- a/core/settingscog.py +++ b/core/settingscog.py @@ -1,68 +1,87 @@ -import discord +from discord import option from discord.ext import commands -from discord import Option +from typing import Optional from core import settings + class SettingsCog(commands.Cog): def __init__(self, bot:commands.Bot): self.bot = bot - - @commands.slash_command(name = "currentoptions") - async def currentoptions(self, ctx): + + @commands.slash_command(name = 'settings', description = 'Review and change server defaults') + @option( + 'current_settings', + bool, + description='Show the current defaults for the server.', + required=False, + ) + @option( + 'set_nprompt', + str, + description='Set default negative prompt for the server', + required=False, + ) + @option( + 'set_steps', + int, + description='Set default amount of steps for the server', + min_value=1, + required=False, + ) + @option( + 'set_maxsteps', + int, + description='Set default maximum steps for the server', + min_value=1, + required=False, + ) + @option( + 'set_sampler', + str, + description='Set default sampler for the server', + required=False, + choices=['Euler a', 'Euler', 'LMS', 'Heun', 'DPM2', 'DPM2 a', 'DPM fast', 'DPM adaptive', 'LMS Karras', 'DPM2 Karras', 'DPM2 a Karras', 'DDIM', 'PLMS'], + ) + async def settings_handler(self, ctx, + current_settings: Optional[bool] = False, + set_nprompt: Optional[str] = 'unset', + set_steps: Optional[int] = 1, + set_maxsteps: Optional[int] = 1, + set_sampler: Optional[str] = 'unset'): guild = '% s' % ctx.guild_id - try: - await ctx.respond(settings.read(guild)) #output is ugly - except FileNotFoundError: - settings.build(guild) - await ctx.respond('Config file not found, building...') + reviewer = settings.read(guild) + reply = 'Summary:\n' + if current_settings: + cur_set = settings.read(guild) + for key, value in cur_set.items(): + reply = reply + str(key) + ": ``" + str(value) + "``, " + + #run through each command and update the defaults user selects + if set_nprompt != 'unset': + settings.update(guild, 'negative_prompt', set_nprompt) + reply = reply + '\nNew default negative prompts is "' + str(set_nprompt) + '".' + + if set_sampler != 'unset': + settings.update(guild, 'sampler', set_sampler) + reply = reply + '\nNew default sampler is "' + str(set_sampler) + '".' + + if set_maxsteps != 1: + settings.update(guild, 'max_steps', set_maxsteps) + reply = reply + '\nNew max steps value is ' + str(set_maxsteps) + '.' + #automatically lower default steps if max steps goes below it + if set_maxsteps < reviewer['default_steps']: + settings.update(guild, 'default_steps', set_maxsteps) + reply = reply + '\nDefault steps value is too high! Lowering to ' + str(set_maxsteps) + '.' + + #review settings again in case user is trying to set steps and max steps simultaneously + reviewer = settings.read(guild) + if set_steps > reviewer['max_steps']: + reply = reply + '\nMax steps is ' + str(reviewer["max_steps"]) + '! You can\'t go beyond it!' + elif set_steps != 1: + settings.update(guild, 'default_steps', set_steps) + reply = reply + '\nNew default steps value is ' + str(set_steps) + '.' + + await ctx.send_response(reply) - @commands.slash_command(name = "changesettings", description = 'Change Server Defaults for /draw') - async def setdefaultsettings( - self, - ctx, - setsteps: Option(int, "Set Steps", min_value= 1, default=1), - setnprompt: Option(str, "Set Negative Prompt", default='unset'), - setmaxsteps: Option(int, "Set Max Steps", min_value= 1, default=1), - setsampler: Option(str, "Set Sampler", default='unset') - ): - guild_id = '% s' % ctx.guild_id - maxsteps = settings.read(guild_id) - if setsteps > maxsteps['max_steps']: - await ctx.respond('default steps cant go beyond max steps') - await ctx.send_message('CURRENT MAXSTEPS:'+str(maxsteps['max_steps'])) - elif setsteps != 1: - try: - settings.update(guild_id, 'default_steps', setsteps) - await ctx.respond('New default steps value Set') - except FileNotFoundError: - settings.build(guild_id) - await ctx.respond('Config file not found, building...') - if setnprompt != 'unset': - try: - settings.update(guild_id, 'negative_prompt', setnprompt) - await ctx.respond('New default negative prompts Set') - except FileNotFoundError: - settings.build(guild_id) - await ctx.respond('Config file not found, building...') - if setmaxsteps != 1: - try: - settings.update(guild_id, 'max_steps', setmaxsteps) - await ctx.respond('New max steps value Set') - except FileNotFoundError: - settings.build(guild_id) - await ctx.respond('Config file not found, building...') - if setsampler != 'unset': - #Disclaimer: I know there's a more sophisticated way to do this but pycord hates me so I'm not risking it right now - samplers = {'Euler a', 'Euler', 'LMS', 'Heun', 'DPM2', 'DPM2 a', 'DPM fast', 'DPM adaptive', 'LMS Karras', 'DPM2 Karras', 'DPM2 a Karras', 'DDIM', 'PLMS'} - if setsampler in samplers: - try: - settings.update(guild_id, 'sampler', setsampler) - await ctx.respond('New default sampler Set') - except FileNotFoundError: - settings.build(guild_id) - await ctx.respond('Config file not found, building...') - else: - await ctx.respond('Please use one of the following options: ' + ' , '.join(samplers) ) - -def setup(bot:commands.Bot): +def setup(bot): bot.add_cog(SettingsCog(bot)) \ No newline at end of file diff --git a/core/stablecog.py b/core/stablecog.py index 47eaff00..ccf3359a 100644 --- a/core/stablecog.py +++ b/core/stablecog.py @@ -10,38 +10,14 @@ from discord.ext import commands from typing import Optional from PIL import Image, PngImagePlugin -from core import settings import base64 from discord import option import random import time import csv +from core import settings -embed_color = discord.Colour.from_rgb(222, 89, 28) -global URL -global DIR - -#check .env for URL and DIR. if they don't exist, ignore it and go with defaults. -if os.getenv("URL") == '': - URL = os.environ.get('URL').rstrip("/") - print(f'Using URL: {URL}') -else: - URL = 'http://127.0.0.1:7860' - print('Using Default URL: http://127.0.0.1:7860') - -if os.getenv("DIR") == '': - DIR = os.environ.get('DIR') - print(f'Using outputs directory: {DIR}') -else: - DIR = "outputs" - print('Using default outputs directory: outputs') -#if directory in DIR doesn't exist, create it -dir_exists = os.path.exists(DIR) -if dir_exists is False: - print(f'The folder for DIR is missing! Creating folder at {DIR}.') - os.mkdir(DIR) - class QueueObject: def __init__(self, ctx, prompt, negative_prompt, steps, height, width, guidance_scale, sampler, seed, strength, init_image, copy_command): @@ -65,7 +41,6 @@ def __init__(self, bot): self.queue = [] self.wait_message = [] self.bot = bot - self.url = URL @commands.slash_command(name = 'draw', description = 'Create an image') @option( @@ -83,9 +58,9 @@ def __init__(self, bot): @option( 'steps', int, - description='The amount of steps to sample the model. Default: 30', + description='The amount of steps to sample the model.', + min_value=1, required=False, - choices=[x for x in range(5, 55, 5)] ) @option( 'height', @@ -113,7 +88,6 @@ def __init__(self, bot): description='The sampler to use for generation. Default: Euler a', required=False, choices=['Euler a', 'Euler', 'LMS', 'Heun', 'DPM2', 'DPM2 a', 'DPM fast', 'DPM adaptive', 'LMS Karras', 'DPM2 Karras', 'DPM2 a Karras', 'DDIM', 'PLMS'], - default='Euler a' ) @option( 'seed', @@ -143,46 +117,56 @@ async def dream_handler(self, ctx: discord.ApplicationContext, *, init_image: Optional[discord.Attachment] = None,): print(f'Request -- {ctx.author.name}#{ctx.author.discriminator} -- Prompt: {prompt}') + #update defaults with any new defaults from settingscog guild = '% s' % ctx.guild_id - try: - sett = settings.read(guild) - except FileNotFoundError: - settings.build(guild) - sett = settings.read(guild) - if negative_prompt == 'unset': - negative_prompt = sett['negative_prompt'] - if steps == -1 or steps > sett['max_steps']: - steps = sett['default_steps'] + negative_prompt = settings.read(guild)['negative_prompt'] + if steps == -1: + steps = settings.read(guild)['default_steps'] if sampler == 'unset': - sampler = sett['sampler'] + sampler = settings.read(guild)['sampler'] if seed == -1: seed = random.randint(0, 0xFFFFFFFF) #increment number of times command is used - with open('resources/stats.txt', 'r') as f: data = list(map(int, f.readlines())) + with open('resources/stats.txt', 'r') as f: + data = list(map(int, f.readlines())) data[0] = data[0] + 1 - with open('resources/stats.txt', 'w') as f: f.write('\n'.join(str(x) for x in data)) + with open('resources/stats.txt', 'w') as f: + f.write('\n'.join(str(x) for x in data)) #random messages for bot to say with open('resources/messages.csv') as csv_file: message_data = list(csv.reader(csv_file, delimiter='|')) message_row_count = len(message_data) - 1 - for row in message_data: self.wait_message.append( row[0] ) - - #log the command. can replace bot reply with {copy_command} for easy copy-pasting - copy_command = f'/draw prompt:{prompt} steps:{steps} height:{str(height)} width:{width} guidance_scale:{guidance_scale} sampler:{sampler} seed:{seed}' - if negative_prompt != '': copy_command = copy_command + f' negative_prompt:{negative_prompt}' - if init_image: copy_command = copy_command + f' strength:{strength}' - print(copy_command) + for row in message_data: + self.wait_message.append( row[0] ) #formatting bot initial reply append_options = '' - if negative_prompt != '': append_options = append_options + '\nNegative Prompt: ``' + str(negative_prompt) + '``' - if height != 512: append_options = append_options + '\nHeight: ``' + str(height) + '``' - if width != 512: append_options = append_options + '\nWidth: ``' + str(width) + '``' - if guidance_scale != 7.0: append_options = append_options + '\nGuidance Scale: ``' + str(guidance_scale) + '``' - if sampler != 'Euler a': append_options = append_options + '\nSampler: ``' + str(sampler) + '``' - if init_image: append_options = append_options + '\nStrength: ``' + str(strength) + '``' + #lower step value to highest setting if user goes over max steps + if steps > settings.read(guild)['max_steps']: + steps = settings.read(guild)['max_steps'] + append_options = append_options + '\nExceeded maximum of ``' + str(steps) + '`` steps! This is the best I can do...' + if negative_prompt != '': + append_options = append_options + '\nNegative Prompt: ``' + str(negative_prompt) + '``' + if height != 512: + append_options = append_options + '\nHeight: ``' + str(height) + '``' + if width != 512: + append_options = append_options + '\nWidth: ``' + str(width) + '``' + if guidance_scale != 7.0: + append_options = append_options + '\nGuidance Scale: ``' + str(guidance_scale) + '``' + if sampler != 'Euler a': + append_options = append_options + '\nSampler: ``' + str(sampler) + '``' + if init_image: + append_options = append_options + '\nStrength: ``' + str(strength) + '``' + + # log the command. can replace bot reply with {copy_command} for easy copy-pasting + copy_command = f'/draw prompt:{prompt} steps:{steps} height:{str(height)} width:{width} guidance_scale:{guidance_scale} sampler:{sampler} seed:{seed}' + if negative_prompt != '': + copy_command = copy_command + f' negative_prompt:{negative_prompt}' + if init_image: + copy_command = copy_command + f' strength:{strength}' + print(copy_command) #setup the queue if self.dream_thread.is_alive(): @@ -243,13 +227,13 @@ def dream(self, event_loop: AbstractEventLoop, queue_object: QueueObject): 'username': os.getenv('USER'), 'password': os.getenv('PASS') } - s.post(URL + '/login', data=login_payload) + s.post(settings.global_var.url + '/login', data=login_payload) else: - s.post(URL + '/login') + s.post(settings.global_var.url + '/login') if queue_object.init_image is not None: - response = requests.post(url=f'{self.url}/sdapi/v1/img2img', data=payload_json).json() + response = requests.post(url=f'{settings.global_var.url}/sdapi/v1/img2img', data=payload_json).json() else: - response = requests.post(url=f'{self.url}/sdapi/v1/txt2img', data=payload_json).json() + response = requests.post(url=f'{settings.global_var.url}/sdapi/v1/txt2img', data=payload_json).json() end_time = time.time() @@ -259,15 +243,15 @@ def dream(self, event_loop: AbstractEventLoop, queue_object: QueueObject): metadata = PngImagePlugin.PngInfo() epoch_time = int(time.time()) metadata.add_text("parameters", str(response['info'])) - image.save(f'{DIR}\{epoch_time}-{queue_object.seed}-{queue_object.prompt[0:120]}.png', pnginfo=metadata) - print(f'Saved image: {DIR}\{epoch_time}-{queue_object.seed}-{queue_object.prompt[0:120]}.png') + image.save(f'{settings.global_var.dir}\{epoch_time}-{queue_object.seed}-{queue_object.prompt[0:120]}.png', pnginfo=metadata) + print(f'Saved image: {settings.global_var.dir}\{epoch_time}-{queue_object.seed}-{queue_object.prompt[0:120]}.png') #post to discord with io.BytesIO() as buffer: image.save(buffer, 'PNG') buffer.seek(0) embed = discord.Embed() - embed.colour = embed_color + embed.colour = settings.global_var.embed_color if os.getenv("COPY") is not None: embed.add_field(name='My drawing of', value=f'``{queue_object.copy_command}``', inline=False) else: @@ -283,7 +267,7 @@ def dream(self, event_loop: AbstractEventLoop, queue_object: QueueObject): except Exception as e: embed = discord.Embed(title='txt2img failed', description=f'{e}\n{traceback.print_exc()}', - color=embed_color) + color=settings.global_var.embed_color) event_loop.create_task(queue_object.ctx.channel.send(embed=embed)) if self.queue: event_loop.create_task(self.process_dream(self.queue.pop(0)))