Testing text2img command (will be replaced, only for tests); added exception handling, queue and database table with generated images

bot/handlers/admin/__init__.py

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

bot/handlers/admin/reset.py

@@ -0,0 +1,18 @@
+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 resetqueue(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
+
+    db[DBTables.queue]['n'] = 0
+
+    await db[DBTables.config].write()
+
+    await message.reply("✅ Reset queue")

bot/handlers/initialize/all_messages.py

@@ -4,5 +4,6 @@ from bot.db.pull_db import pull
 
 async def sync_db_filter(message: Message):
     await pull()
-    await message.reply(f'🔄️ Bot database synchronised because of restart. '
-                        f'If you tried to run a command, run it again')
+    if message.is_command():
+        await message.reply(f'🔄️ Bot database synchronised because of restart. '
+                            f'If you tried to run a command, run it again')

bot/handlers/register.py

@@ -1,7 +1,7 @@
 from rich import print
 
 
-def import_handlers():
+def register_handlers():
     from bot.handlers import (
         initialize, admin, help_command, txt2img
     )
@@ -11,4 +11,4 @@ def import_handlers():
     help_command.register()
     txt2img.register()
 
-    print('[gray]All handlers imported[/]')
+    print('[gray]All handlers registered[/]')

bot/handlers/txt2img/__init__.py

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

bot/handlers/txt2img/txt2img.py

@@ -2,25 +2,50 @@ from aiogram import types
 from bot.db import db, DBTables
 from bot.utils.cooldown import throttle
 from bot.modules.api.txt2img import txt2img
+from bot.modules.api.objects.prompt_request import Prompt
+from bot.modules.api.status import wait_for_status
+from bot.keyboards.exception import get_exception_keyboard
+from bot.utils.trace_exception import PrettyException
+from aiohttp import ClientConnectorError
 
 
 @throttle(cooldown=30, admin_ids=db[DBTables.config].get('admins'))
 async def txt2img_comand(message: types.Message):
-    temp_message = await message.reply("⏳ Generating image...")
+    temp_message = await message.reply("⏳ Enqueued...")
     if not message.get_args():
-        await temp_message.edit_text("Specify prompt for this command. Check /help txt2img")
+        await temp_message.edit_text("😶‍🌫️ Specify prompt for this command. Check /help txt2img")
         return
 
     try:
-        image = await txt2img(message.get_args())
-        await message.reply_photo(photo=image[0], caption=str(
-            image[1]["infotexts"][0]))
-    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")
+        db[DBTables.queue]['n'] = db[DBTables.queue].get('n', 0) + 1
+        await temp_message.edit_text(f"⏳ Enqueued in position {db[DBTables.queue].get('n', 0)}...")
+
+        await wait_for_status()
+
+        await temp_message.edit_text(f"⌛ Generating...")
+        prompt = Prompt(prompt=message.get_args(), creator=message.from_id)
+        image = await txt2img(prompt)
+        image_message = await message.reply_photo(photo=image[0])
+
+        db[DBTables.queue]['n'] = db[DBTables.queue].get('n', 1) - 1
+        db[DBTables.generated][image_message.photo[0].file_unique_id] = prompt
+
         await temp_message.delete()
+
+        await db[DBTables.config].write()
+
+    except ClientConnectorError:
+        await message.reply('❌ Error! Maybe, StableDiffusion API endpoint is incorrect '
+                            'or turned off')
+        await temp_message.delete()
+        db[DBTables.queue]['n'] = db[DBTables.queue].get('n', 1) - 1
         return
 
-    await temp_message.delete()
+    except Exception as e:
+        exception_id = f'{message.message_thread_id}-{message.message_id}'
+        db[DBTables.exceptions][exception_id] = PrettyException(e)
+        await message.reply('❌ Error happened while processing your request', parse_mode='html',
+                            reply_markup=get_exception_keyboard(exception_id))
+        await temp_message.delete()
+        db[DBTables.queue]['n'] = db[DBTables.queue].get('n', 1) - 1
+        return