update pr #17

described in #17 (comment)

also make code a bit more "pythonic"

.gitignore

@@ -3,9 +3,9 @@ __pycache__
 pydata
 venv
 
-core/generated/*
 outputs/*.png
 resources/stats.txt
+resources/*.json
 
 .env
 start.bat

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:
 ```

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()
 

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)

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