Implement spoiler functionality

* channel-wide spoiler option available via /settings
* per /draw option
* force spoiler on role, option available via /settings

README.md

@@ -57,6 +57,9 @@ To generate a prompt from a couple of words, use the /generate command and inclu
   - Get Image Info - view information of an image generated by Stable Diffusion.
   - Quick Upscale - upscale an image without needing to set options.
   - Batch Download - download all images of a batch set without needing to specify batch_id and image_id
+- mark image as spoiler
+  - Per image (on `/draw`)
+  - Set channel-wide default or force based on role with `/settings`
 - [configuration file](https://github.com/Kilvoctu/aiyabot/wiki/Configuration) - can change some of AIYA's operational aspects. 
 
 

core/queuehandler.py

@@ -6,7 +6,7 @@
 class DrawObject:
     def __init__(self, cog, ctx, simple_prompt, prompt, negative_prompt, data_model, steps, width, height,
                  guidance_scale, sampler, seed, strength, init_image, batch, styles, facefix, highres_fix,
-                 clip_skip, extra_net, epoch_time, view):
+                 clip_skip, extra_net, spoiler, epoch_time, view):
         self.cog = cog
         self.ctx = ctx
         self.simple_prompt = simple_prompt
@@ -27,6 +27,7 @@ def __init__(self, cog, ctx, simple_prompt, prompt, negative_prompt, data_model,
         self.highres_fix = highres_fix
         self.clip_skip = clip_skip
         self.extra_net = extra_net
+        self.spoiler = spoiler
         self.epoch_time = epoch_time
         self.view = view
 

core/settings.py

@@ -82,6 +82,9 @@
 batch = "1,1"
 max_batch = "1,1"
 upscaler_1 = "ESRGAN_4x"
+spoiler = false
+# role ID (not name)
+spoiler_role = ""
 """
 
 
@@ -120,6 +123,8 @@ class GlobalVar:
     prompt_ignore_list = []
     display_ignored_words = "False"
     negative_prompt_prefix = []
+    spoiler = False
+    spoiler_role = None
 
 
 global_var = GlobalVar()
@@ -276,6 +281,8 @@ def generate_template(template_pop, config):
     template_pop['batch'] = config['batch']
     template_pop['max_batch'] = config['max_batch']
     template_pop['upscaler_1'] = config['upscaler_1']
+    template_pop['spoiler'] = config['spoiler']
+    template_pop['spoiler_role'] = config['spoiler_role']
     return template_pop
 
 
@@ -300,6 +307,12 @@ def read(channel_id):
             with open(path + channel_id + '.json', 'w') as configfile2:
                 json.dump(settings, configfile2, indent=1)
 
+        if settings['spoiler_role'] is not None:
+            if str(settings['spoiler_role']).strip() == '':
+                settings['spoiler_role'] = None
+            else:
+                settings['spoiler_role'] = str(settings['spoiler_role'])
+
     return settings
 
 

core/settingscog.py

@@ -1,7 +1,7 @@
 import discord
 from discord import option
 from discord.ext import commands
-from typing import Optional
+from typing import Optional, List, Union
 
 from core import settings
 
@@ -199,6 +199,24 @@ def size_autocomplete(self: discord.AutocompleteContext):
         description='Use to update global lists (models, styles, embeddings, etc.)',
         required=False,
     )
+    @option(
+        'spoiler',
+        bool,
+        description='Mark images as spoilers (when not specified in /draw)',
+        required=False,
+    )
+    @option(
+        'spoiler_role',
+        discord.Role,
+        description='Force images from users in this role as spoilers',
+        required=False,
+    )
+    @option(
+        'remove_spoiler_role',
+        bool,
+        description='Remove assigned spoiler role',
+        required=False,
+    )
     async def settings_handler(self, ctx,
                                current_settings: Optional[bool] = True,
                                n_prompt: Optional[str] = None,
@@ -218,7 +236,11 @@ async def settings_handler(self, ctx,
                                batch: Optional[str] = None,
                                max_batch: Optional[str] = None,
                                upscaler_1: Optional[str] = None,
-                               refresh: Optional[bool] = False):
+                               refresh: Optional[bool] = False,
+                               spoiler: Optional[bool] = None,
+                               spoiler_role: Optional[discord.Role] = None,
+                               remove_spoiler_role: Optional[bool] = None
+                               ):
         # get the channel id and check if a settings file exists
         channel = '% s' % ctx.channel.id
         settings.check(channel)
@@ -236,6 +258,8 @@ async def settings_handler(self, ctx,
             for key, value in cur_set.items():
                 if key == 'negative_prompt':
                     pass
+                elif key == 'spoiler_role' and value is not None:
+                    current += f'\n{key} - <@&{value}>'
                 else:
                     if value == '':
                         value = ' '
@@ -403,6 +427,19 @@ async def settings_handler(self, ctx,
                 new += f'\nbatch (count,size): ``{batch[0]},{batch[1]}``'
             set_new = True
 
+        if spoiler is not None:
+            settings.update(channel, 'spoiler', spoiler)
+            new += f'\nDefault Spoiler: ``{spoiler}``'
+            set_new = True
+        if remove_spoiler_role is not None and remove_spoiler_role:
+            settings.update(channel, 'spoiler_role', None)
+            new += '\nRemoved Spoiler Role'
+            set_new = True
+        elif spoiler_role is not None:
+            settings.update(channel, 'spoiler_role', str(spoiler_role.id))
+            new += f'\n Spoiler Role: <@&{spoiler_role.id}>'
+            set_new = True
+
         if set_new:
             embed.add_field(name=f'New defaults', value=new, inline=False)
         if new_n_prompt:

core/stablecog.py

@@ -152,6 +152,12 @@ async def on_ready(self):
         description='The number of images to generate. Batch format: count,size',
         required=False,
     )
+    @option(
+        'spoiler',
+        bool,
+        description='Mark generated image as spoiler?',
+        required=False,
+    )
     async def dream_handler(self, ctx: discord.ApplicationContext, *,
                             prompt: str, negative_prompt: str = None,
                             data_model: Optional[str] = None,
@@ -168,7 +174,8 @@ async def dream_handler(self, ctx: discord.ApplicationContext, *,
                             strength: Optional[str] = None,
                             init_image: Optional[discord.Attachment] = None,
                             init_url: Optional[str],
-                            batch: Optional[str] = None):
+                            batch: Optional[str] = None,
+                            spoiler: Optional[bool] = None):
 
         # update defaults with any new defaults from settingscog
         channel = '% s' % ctx.channel.id
@@ -197,6 +204,16 @@ async def dream_handler(self, ctx: discord.ApplicationContext, *,
             strength = settings.read(channel)['strength']
         if batch is None:
             batch = settings.read(channel)['batch']
+        if spoiler is None:
+            spoiler = settings.read(channel)['spoiler']
+
+        derived_spoiler = spoiler
+        spoiler_role = settings.read(channel)['spoiler_role']
+        if not derived_spoiler and spoiler_role is not None:
+            for role in ctx.author.roles:
+                if str(role.id) == spoiler_role:
+                    derived_spoiler = True
+                    break
 
         # if a model is not selected, do nothing
         model_name = 'Default'
@@ -234,9 +251,9 @@ async def dream_handler(self, ctx: discord.ApplicationContext, *,
         prompt = settings.extra_net_defaults(prompt, channel)
 
         if data_model != '':
-            print(f'Request -- {ctx.author.name}#{ctx.author.discriminator} -- Prompt: {prompt}')
+            print(f'Request -- {ctx.author.name}#{ctx.author.discriminator} -- Prompt: {prompt} -- Spoiler: {derived_spoiler}')
         else:
-            print(f'Request -- {ctx.author.name}#{ctx.author.discriminator} -- Prompt: {prompt} -- Using model: {data_model}')
+            print(f'Request -- {ctx.author.name}#{ctx.author.discriminator} -- Prompt: {prompt} -- Using model: {data_model} -- Spoiler: {derived_spoiler}')
 
         if seed == -1:
             seed = random.randint(0, 0xFFFFFFFF)
@@ -322,13 +339,17 @@ async def dream_handler(self, ctx: discord.ApplicationContext, *,
             reply_adds += f'\nFace restoration: ``{facefix}``'
         if clip_skip != settings.read(channel)['clip_skip']:
             reply_adds += f'\nCLIP skip: ``{clip_skip}``'
-
+
+        if derived_spoiler or derived_spoiler != settings.read(channel)['spoiler']:
+            bool_emoji = ':white_check_mark:' if derived_spoiler else ':negative_squared_cross_mark:'
+            reply_adds += f'\nSpoiler: {bool_emoji}'
+
         epoch_time = int(time.time())
 
         # set up tuple of parameters to pass into the Discord view
         input_tuple = (
             ctx, simple_prompt, prompt, negative_prompt, data_model, steps, width, height, guidance_scale, sampler, seed, strength,
-            init_image, batch, styles, facefix, highres_fix, clip_skip, extra_net, epoch_time)
+            init_image, batch, styles, facefix, highres_fix, clip_skip, extra_net, derived_spoiler, epoch_time)
 
         view = viewhandler.DrawView(input_tuple)
         # setup the queue
@@ -554,6 +575,8 @@ def dream(self, event_loop: queuehandler.GlobalQueue.event_loop, queue_object: q
                         id_start = current_grid * grid_count + 1
                         id_end = id_start + last_grid_count - 1
                     filename=f'{queue_object.seed}-{current_grid}.png'
+                    if queue_object.spoiler:
+                        filename=f'SPOILER_{queue_object.seed}-{count}.png'
                     file = add_metadata_to_image(grid,images[current_grid * 25][2], filename)
                     if current_grid == 0:
                         content = f'<@{queue_object.ctx.author.id}>, {message}\n Batch ID: {epoch_time}-{queue_object.seed}\n Image IDs: {id_start}-{id_end}'
@@ -570,6 +593,8 @@ def dream(self, event_loop: queuehandler.GlobalQueue.event_loop, queue_object: q
             else:
                 content = f'<@{queue_object.ctx.author.id}>, {message}'
                 filename=f'{queue_object.seed}-{count}.png'
+                if queue_object.spoiler:
+                    filename=f'SPOILER_{queue_object.seed}-{count}.png'
                 file = add_metadata_to_image(image,str_parameters, filename)
                 queuehandler.process_post(
                     self, queuehandler.PostObject(

core/viewhandler.py

@@ -32,11 +32,12 @@
 [16] = highres_fix
 [17] = clip_skip
 [18] = extra_net
-[19] = epoch_time
+[19] = spoiler
+[20] = epoch_time
 '''
 tuple_names = ['ctx', 'simple_prompt', 'prompt', 'negative_prompt', 'data_model', 'steps', 'width', 'height',
                'guidance_scale', 'sampler', 'seed', 'strength', 'init_image', 'batch', 'styles', 'facefix',
-               'highres_fix', 'clip_skip', 'extra_net', 'epoch_time']
+               'highres_fix', 'clip_skip', 'extra_net', 'spoiler', 'epoch_time']
 
 
 # the modal that is used for the 🖋 button