first

ai_bot/handlers/ImageHandler.py

@@ -0,0 +1,53 @@
+from ai_bot.keyboards.CreateImageKeyboard import CreateImageKeyboard
+from ai_bot.msg.CreateImageMsg import CreateImageMsg
+from bot.Handler import Handler
+from bot.halpers.Instance import Instance
+from bot import SECOND
+from telegram.constants import ParseMode
+
+
+class ImageHandler(Handler):
+
+    @staticmethod
+    async def create_img(update, context):
+        models_arr = ['Midjourney', 'Flux']
+
+        reply_markup = CreateImageKeyboard()
+        reply_markup.add_option("models_arr", models_arr)
+
+        msg = CreateImageMsg()
+
+        await context.bot.send_message(chat_id=update.effective_chat.id, text=msg.get_msg(),
+                                       parse_mode=ParseMode.HTML, reply_markup=reply_markup.create_keyboard())
+
+        return SECOND
+
+    @staticmethod
+    async def generate_img(update, context):
+        query = update.callback_query
+        command, params = Handler.load_callback_query(query.data)
+
+    @staticmethod
+    async def set_prompt(update, context):
+        query = update.callback_query
+        command, params = Handler.load_callback_query(query.data)
+
+        model = Instance.get_instance("ai_bot.states.{model}States".format(model=params['model']), "{model}States".format(model=params['model']))
+        state = model.get_state_by_key("set_prompt")
+
+        await context.bot.send_message(chat_id=update.effective_chat.id, text="Напишите запрос.",
+                                       parse_mode=ParseMode.HTML)
+
+        return state
+
+    @staticmethod
+    async def start_model_scenario(update, context):
+        query = update.callback_query
+        command, params = Handler.load_callback_query(query.data)
+
+        model = Instance.get_instance("ai_bot.handlers.{model}Handler".format(model=params['model']),
+                                      "{model}Handler".format(model=params['model']))
+
+        state = await model.start_scenario(update, context)
+
+        return state