Disabling generation function

bot/callbacks/config/admin_settings.py

@@ -31,6 +31,7 @@ async def on_admin_settings_set(call: types.CallbackQuery, callback_data: dict):
         if 'aliases' in overload and 'admin' in overload and ('add' in overload or 'remove' in overload)
         else f"⚒️ Type new endpoint address: " if "aliases.set_endpoint" in overload
         else f"⚒️ Type \"reset\" if you REALLY want to reset queue: " if "reset.resetqueue" in overload
+        else f"⚒️ Type \"on\" or \"off\" to change generation mode: " if "on_off" in overload
         else f"❌ Not found...",
         reply_markup=types.InlineKeyboardMarkup().add(types.InlineKeyboardButton(
             "👈 Back",

bot/handlers/admin/__init__.py

@@ -2,6 +2,7 @@ from bot.common import dp
 from .aliases import *
 from .reset import *
 from .tools import *
+from .on_off import *
 
 
 def register():

bot/handlers/admin/on_off.py

@@ -0,0 +1,17 @@
+from aiogram import types
+from bot.db import db, DBTables
+from bot.config import ADMIN
+from bot.utils.cooldown import throttle
+
+
+@throttle(5)
+async def on_off_call(message: types.Message, is_command=None):
+    if message.from_id not in db[DBTables.config].get('admins') and message.from_id != ADMIN:
+        await message.reply('❌ You are not permitted to do that. '
+                            'It is only for this bot instance maintainers and admins')
+        return
+
+    db[DBTables.config]['enabled'] = False if message.text.lower() == 'off' else True
+    await db[DBTables.config].write()
+
+    await message.reply(f"Generation enabled: {'💚 (yes)' if db[DBTables.config]['enabled'] else '💔 (no)'}")

bot/handlers/initialize/start.py

@@ -18,6 +18,9 @@ async def start_command(message: types.Message):
             await db[DBTables.config].write()
             await message.reply(f'✅ Added {message.from_user.username} to admins. You can add other admins, '
                                 f'check bot settings menu')
+        if db[DBTables.config].get('enabled') is None:
+            db[DBTables.config]['enabled'] = True
+            await message.reply(f'✅ Generation is enabled now')
         return
 
     await message.reply(f'👋 Hello, {message.from_user.username}. Use /help to see available commands.')

bot/handlers/txt2img/txt2img.py

@@ -14,6 +14,10 @@ from bot.utils.errorable_command import wrap_exception
 @throttle(cooldown=30, admin_ids=db[DBTables.config].get('admins'))
 async def generate_command(message: types.Message):
     temp_message = await message.reply("⏳ Enqueued...")
+    if not db[DBTables.config]['enabled']:
+        await message.reply('💔 Generation is disabled by admins now. Try again later')
+        await temp_message.delete()
+        return
 
     try:
         prompt = get_prompt(user_id=message.from_id,

bot/keyboards/config.py

@@ -57,6 +57,8 @@ def get_admin_settings_keyboard() -> types.InlineKeyboardMarkup:
         types.InlineKeyboardButton("Remove admin", callback_data=admin_settings_data.new("aliases.remove_admin")),
         types.InlineKeyboardButton("Set API endpoint", callback_data=admin_settings_data.new("aliases.set_endpoint")),
         types.InlineKeyboardButton("Reset generation queue", callback_data=admin_settings_data.new("reset.resetqueue")),
+        types.InlineKeyboardButton("Turn on/off generation",
+                                   callback_data=admin_settings_data.new("on_off.on_off_call")),
         types.InlineKeyboardButton("👈 Back", callback_data="config_back"),
         types.InlineKeyboardButton("🔻 Close", callback_data="close_keyboard")
     ]