Disabling generation function

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,