Admin and start handlers work. Added testing txt2img from first bot version (will be replaced soon)

bot/handlers/admin/__init__.py

@@ -0,0 +1,6 @@
+from bot.common import dp
+from .aliases import *
+
+
+def register():
+    dp.register_message_handler(set_endpoint, commands='setendpoint')

bot/handlers/admin/aliases.py

@@ -0,0 +1,23 @@
+from aiogram import types
+from bot.db import db, DBTables
+import validators
+from bot.config import ADMIN
+from bot.utils.cooldown import throttle
+
+
+@throttle(5)
+async def set_endpoint(message: types.Message):
+    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
+
+    if not message.get_args() or not validators.url(message.get_args()):
+        await message.reply("❌ Specify correct url for endpoint")
+        return
+
+    db[DBTables.config]['endpoint'] = message.get_args()
+
+    await db[DBTables.config].write()
+
+    await message.reply("✅ New url set")

bot/handlers/help/__init__.py

@@ -1 +0,0 @@
-pass

bot/handlers/help/help_strings.py

@@ -1,2 +0,0 @@
-help_data = {
-}

bot/handlers/help_command/__init__.py

@@ -0,0 +1,6 @@
+from bot.common import dp
+from . import help_handler
+
+
+def register():
+    dp.register_message_handler(help_handler.help_command, commands='help')

bot/handlers/help_command/help_handler.py

@@ -1,9 +1,7 @@
 from aiogram import types
-from bot.common import dp
 from .help_strings import help_data
 
 
-@dp.message_handler(commands='help')
 async def help_command(message: types.Message):
     if message.get_args() == "":
         await message.reply(

bot/handlers/help_command/help_strings.py

@@ -0,0 +1,3 @@
+help_data = {
+    'setendpoint': '(admin) Set StableDiffusion API endpoint'
+}

bot/handlers/initialize/__init__.py

@@ -1 +1,8 @@
-pass
+from bot.common import dp, bot
+from .start import *
+from .all_messages import *
+
+
+def register():
+    dp.register_message_handler(all_messages.sync_db_filter, lambda *_: not hasattr(bot, 'cloudmeta_message_text'))
+    dp.register_message_handler(start.start_command, commands='start')

bot/handlers/initialize/all_messages.py

@@ -0,0 +1,7 @@
+from aiogram.types import Message
+from bot.db.pull_db import pull
+
+
+async def sync_db_filter(message: Message):
+    await pull()
+    await message.reply(f'🔄️ Bot database synchronised. If you tried to run a command, run it again')

bot/handlers/initialize/pull_db.py

@@ -1,7 +0,0 @@
-from bot.common import dp
-from bot.db.pull_db import pull
-
-
-@dp.message_handler()
-async def pull_db_if_new(_):
-    await pull()

bot/handlers/initialize/start.py

@@ -0,0 +1,23 @@
+from aiogram import types
+from bot.db import db, DBTables
+from bot.config import ADMIN
+from bot.utils.cooldown import throttle
+
+
+@throttle(10)
+async def start_command(message: types.Message):
+    if message.from_id == ADMIN:
+        await message.reply(f'👋 Hello, {message.from_user.username}. You are admin of this instance, '
+                            f'so we will check config for you now')
+        if not isinstance(db[DBTables.config].get('admins'), list):
+            db[DBTables.config]['admins'] = list()
+        if ADMIN not in db[DBTables.config].get('admins'):
+            admins_ = db[DBTables.config].get('admins')
+            admins_.append(ADMIN)
+            db[DBTables.config]['admins'] = admins_
+            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')
+        return
+
+    await message.reply(f'👋 Hello, {message.from_user.username}. Use /help to see available commands.')

bot/handlers/register.py

@@ -2,10 +2,13 @@ from rich import print
 
 
 def import_handlers():
-    import bot.handlers.help.help
-    assert bot.handlers.help.help
+    from bot.handlers import (
+        initialize, admin, help_command, txt2img
+    )
 
-    import bot.handlers.initialize.pull_db
-    assert bot.handlers.initialize.pull_db
+    initialize.register()
+    admin.register()
+    help_command.register()
+    txt2img.register()
 
     print('[gray]All handlers imported[/]')

bot/handlers/txt2img/__init__.py

@@ -0,0 +1,6 @@
+from bot.common import dp
+from .txt2img import *
+
+
+def register():
+    dp.register_message_handler(txt2img, commands='txt2img')

bot/handlers/txt2img/txt2img.py

@@ -0,0 +1,31 @@
+from aiogram import types
+from bot.db import db, DBTables
+from bot.utils.cooldown import throttle
+import os.path
+from bot.modules.api.txt2img import txt2img
+
+
+@throttle(cooldown=30, admin_ids=db[DBTables.config].get('admins'))
+async def txt2img_comand(message: types.Message):
+    temp_message = await message.reply("⏳ Generating image...")
+    if not message.get_args():
+        await temp_message.edit_text("Specify prompt for this command. Check /help txt2img")
+        return
+
+    if not os.path.isfile('adstring'):
+        with open('adstring', 'w') as f:
+            f.write('@aiistop_bot')
+
+    try:
+        image = await txt2img(message.get_args())
+        await message.reply_photo(photo=image[0], caption=str(
+            image[1]["infotexts"][0]) + "\n\n" + open('adstring').read())
+    except Exception as e:
+        assert e
+        await message.reply("We ran into error while processing your request. StableDiffusion models may not be "
+                            "configured on specified endpoint or server with StableDiffusion may be turned "
+                            "off. Ask admins of this bot instance if you have contacts for further info")
+        await temp_message.delete()
+        return
+
+    await temp_message.delete()