Setting generation properties

bot/handlers/txt2img/txt2img.py

@@ -1,24 +1,30 @@
+import re
 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.objects.get_prompt import get_prompt
+from bot.modules.api.objects.prompt_request import Generated
 from bot.modules.api.status import wait_for_status
 from bot.keyboards.exception import get_exception_keyboard
+from bot.keyboards.image_info import get_img_info_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):
+async def generate_command(message: types.Message):
     temp_message = await message.reply("⏳ Enqueued...")
 
-    prompt: Prompt = db[DBTables.prompts].get(message.from_id)
-    if not prompt:
-        if message.get_args():
-            db[DBTables.prompts][message.from_id] = Prompt(message.get_args(), creator=message.from_id)
-
-    # TODO: Move it to other module
+    try:
+        prompt = get_prompt(user_id=message.from_id,
+                            prompt_string=message.get_args())
+    except AttributeError:
+        await temp_message.edit_text(f"You didn't created any prompt. Specify prompt text at least first time. "
+                                     f"For example, it can be: <code>masterpiece, best quality, 1girl, white hair, "
+                                     f"medium hair, cat ears, closed eyes, looking at viewer, :3, cute, scarf, jacket, "
+                                     f"outdoors, streets</code>", parse_mode='HTML')
+        return
 
     try:
         db[DBTables.queue]['n'] = db[DBTables.queue].get('n', 0) + 1
@@ -27,12 +33,18 @@ async def txt2img_comand(message: types.Message):
         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
+        db[DBTables.generated][image_message.photo[0].file_unique_id] = Generated(
+            prompt=prompt,
+            seed=image[1]['seed'],
+            model=re.search(r", Model: ([^,]+),", image[1]['infotexts'][0]).groups()[0]
+        )
+
+        await message.reply('Here is your image',
+                            reply_markup=get_img_info_keyboard(image_message.photo[0].file_unique_id))
 
         await temp_message.delete()
 
@@ -45,6 +57,12 @@ async def txt2img_comand(message: types.Message):
         db[DBTables.queue]['n'] = db[DBTables.queue].get('n', 1) - 1
         return
 
+    except ValueError as e:
+        await message.reply(f'❌ Error! {e.args[0]}')
+        await temp_message.delete()
+        db[DBTables.queue]['n'] = db[DBTables.queue].get('n', 1) - 1
+        return
+
     except Exception as e:
         exception_id = f'{message.message_thread_id}-{message.message_id}'
         db[DBTables.exceptions][exception_id] = PrettyException(e)