first

ai_bot/handlers/FluxHandler.py

@@ -0,0 +1,166 @@
+from ai_bot.keyboards.flux.SelectTypeRequestKeyboard import SelectTypeRequestKeyboard
+from ai_bot.msg.flux.SelectTypeMsg import SelectTypeMsg
+from bot import FIRST
+from bot.Handler import Handler
+from telegram.constants import ParseMode
+from ai_bot.states.FluxStates import FLUX_SET_PROMPT, FLUX_SET_SUBMODEL, FLUX_SET_TYPE_REQUEST, FLUX_SET_TEXT_PROMPT, FLUX_SET_IMG_TEXT_PROMPT, FluxStates
+from telegram.ext import MessageHandler, filters, CallbackQueryHandler
+from igf_api.IgfClient import IgfClient
+from piapi_ai.PiapiAiApi import PiapiAiApi
+from piapi_ai import client as piapi_client
+from telegram import Bot
+from dotenv import dotenv_values
+import base64
+
+config = dotenv_values(".env")
+
+
+class FluxHandler(Handler):
+
+    @staticmethod
+    async def set_prompt(update, context):
+        igf_client = IgfClient()
+        task = igf_client.piapiTask.get_one({
+            'bot_id': context.bot.id,
+            'dialog_id': update.effective_chat.id,
+            'status': 0
+        })
+
+        task['prompt'] = update.effective_message.text
+
+        piapi_task = piapi_client.flux.create_task({
+            'model': task['model'],
+            'task_type': task['task_type'],
+            'input': {
+                'prompt': task['prompt'],
+                'width': 1024,
+                'height': 1024,
+            }
+        })
+
+        task['status'] = 1
+        task['task_id'] = piapi_task['data']['task_id']
+
+        igf_client.piapiTask.update(task['id'], task)
+
+        await context.bot.send_message(chat_id=update.effective_chat.id, text="Запрос к Flux принят, идет обработка",
+                                       parse_mode=ParseMode.HTML)
+
+        return FIRST
+
+    @staticmethod
+    async def set_img_prompt(update, context):
+        igf_client = IgfClient()
+        task = igf_client.piapiTask.get_one({
+            'bot_id': context.bot.id,
+            'dialog_id': update.effective_chat.id,
+            'status': 0
+        })
+
+        if update.message.caption:
+            task['prompt'] = update.message.caption
+        else:
+            await context.bot.send_message(chat_id=update.effective_chat.id, text="Prompt нужно указать в описании изображения",
+                                           parse_mode=ParseMode.HTML)
+
+            return FluxStates.get_state_by_key("set_img_text_prompt")
+
+        photo_inst = update.message.photo[-1]
+        photo = await context.bot.get_file(photo_inst['file_id'])
+        file_content = await photo.download_as_bytearray()
+        # Convert the file content to Base64
+        base64_encoded = base64.b64encode(file_content).decode('utf-8')
+
+        piapi_task = piapi_client.flux.create_task({
+            'model': task['model'],
+            'task_type': task['task_type'],
+            'input': {
+                'prompt': task['prompt'],
+                'image': base64_encoded,
+            }
+        })
+
+        task['status'] = 1
+        task['task_id'] = piapi_task['data']['task_id']
+
+        igf_client.piapiTask.update(task['id'], task)
+
+        await context.bot.send_message(chat_id=update.effective_chat.id, text="Запрос к Flux принят",
+                                       parse_mode=ParseMode.HTML)
+
+        return FIRST
+
+    @staticmethod
+    async def set_submodel(update, context):
+        query = update.callback_query
+        command, params = Handler.load_callback_query(query.data)
+
+    @staticmethod
+    async def set_type_request(update, context):
+        query = update.callback_query
+        command, params = Handler.load_callback_query(query.data)
+
+        igf_client = IgfClient()
+
+        task = igf_client.piapiTask.create({
+            'bot_id': context.bot.id,
+            'dialog_id': update.effective_chat.id,
+            'task_type': params['type'],
+            'model': 'Qubico/flux1-dev',
+            'status': 0,
+        })
+
+        if params['type'] == 'txt2img':
+            text = "Введите запрос для генерации изображения"
+            state = FluxStates.get_state_by_key("set_text_prompt")
+        else:
+            text = "Загрузите изображение и запрос для создания изображения"
+            state = FluxStates.get_state_by_key("set_img_text_prompt")
+
+        await context.bot.send_message(chat_id=update.effective_chat.id, text=text,
+                                       parse_mode=ParseMode.HTML)
+
+        return state
+
+    @staticmethod
+    async def start_scenario(update, context):
+        msg = SelectTypeMsg()
+        reply_markup = SelectTypeRequestKeyboard()
+
+        state = FluxStates.get_state_by_key("set_type_request")
+
+        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 state
+
+    @staticmethod
+    async def send_task_result(igf_task: dict, piapi_task: dict):
+        TOKEN = config['TELEGRAM_TOKEN']
+        bot = Bot(TOKEN)
+
+        await bot.send_document(chat_id=igf_task['dialog_id'], document=piapi_task['data']['output']['image_url'])
+        await bot.send_photo(chat_id=igf_task['dialog_id'], photo=piapi_task['data']['output']['image_url'])
+
+        igf_client = IgfClient()
+        igf_task['status'] = 2
+        igf_client.piapiTask.update(igf_task['id'], igf_task)
+
+        return FIRST
+
+    @staticmethod
+    def get_states(data=None) -> dict:
+        return {
+            FLUX_SET_TEXT_PROMPT: [
+                MessageHandler(filters.TEXT, FluxHandler.set_prompt)
+            ],
+            FLUX_SET_SUBMODEL: [
+                CallbackQueryHandler(FluxHandler.set_submodel, pattern='^set_submodel?\S{3,}'),
+            ],
+            FLUX_SET_TYPE_REQUEST: [
+                CallbackQueryHandler(FluxHandler.set_type_request, pattern='^set_type_request?\S{3,}'),
+            ],
+            FLUX_SET_IMG_TEXT_PROMPT: [
+                MessageHandler(filters.PHOTO, FluxHandler.set_img_prompt)
+            ],
+        }

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

ai_bot/handlers/KlingHandler.py

@@ -0,0 +1,24 @@
+from bot.Handler import Handler
+from telegram.constants import ParseMode
+from telegram.ext import MessageHandler, filters, CallbackQueryHandler
+from dotenv import dotenv_values
+from ai_bot.states.KlingStates import KLING_SET_PROMPT, KLING_SET_TEXT_PROMPT
+
+config = dotenv_values(".env")
+
+
+class KlingHandler(Handler):
+
+    @staticmethod
+    def set_prompt(update, context):
+
+
+
+    @staticmethod
+    def get_states(data=None) -> dict:
+        return {
+            KLING_SET_TEXT_PROMPT: [
+                MessageHandler(filters.TEXT, KlingHandler.set_prompt)
+            ],
+
+        }

ai_bot/handlers/MainHandler.py

@@ -0,0 +1,44 @@
+from ai_bot.keyboards.MainKeyboard import MainKeyboard
+from ai_bot.keyboards.MenuKeyboard import MenuKeyboard
+from ai_bot.msg.MainMsg import MainMsg
+from bot.Handler import Handler
+from telegram.constants import ParseMode
+from bot import FIRST
+from igf_api.IgfClient import IgfClient
+from piapi_ai import client
+
+
+class MainHandler(Handler):
+
+    @staticmethod
+    async def start(update, context):
+        msg = MainMsg()
+
+        reply_markup = MainKeyboard()
+
+        client_igf = IgfClient()
+
+        client_igf.tgBot.create({
+            'bot_id': context.bot.id,
+            'dialog_id': update.effective_chat.id,
+            'username': update.effective_chat.username,
+            'first_name': update.effective_chat.first_name,
+            'last_name': update.effective_chat.last_name,
+            'status': 1,
+        })
+
+        client_igf.piapiTask.to_archive_all_new(context.bot.id, update.effective_chat.id)
+
+        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 FIRST
+
+    @staticmethod
+    async def get_menu(update, context):
+        reply_markup = MenuKeyboard()
+
+        await context.bot.send_message(chat_id=update.effective_chat.id, text="Меню:",
+                                       parse_mode=ParseMode.HTML, reply_markup=reply_markup.create_keyboard())
+
+        return FIRST

ai_bot/handlers/MidjourneyHandler.py

@@ -0,0 +1,23 @@
+from bot.Handler import Handler
+from bot import FIRST
+from telegram.constants import ParseMode
+from ai_bot.states.MidjourneyStates import MIDJOURNEY_SET_PROMPT
+from telegram.ext import MessageHandler, filters, CallbackQueryHandler
+
+
+class MidjourneyHandler(Handler):
+
+    @staticmethod
+    async def set_prompt(update, context):
+        await context.bot.send_message(chat_id=update.effective_chat.id, text="Запрос к Midjourney принят",
+                                       parse_mode=ParseMode.HTML)
+
+        return FIRST
+
+    @staticmethod
+    def get_states(data=None) -> dict:
+        return {
+            MIDJOURNEY_SET_PROMPT: [
+                MessageHandler(filters.TEXT, MidjourneyHandler.set_prompt)
+            ]
+        }

ai_bot/handlers/ModelHandler.py

@@ -0,0 +1,8 @@
+from bot.Handler import Handler
+
+
+class ModelHandler(Handler):
+
+    @staticmethod
+    def select_model(update, context):
+        print("select model")

ai_bot/handlers/VideoHandler.py

@@ -0,0 +1,22 @@
+from ai_bot.keyboards.CreateVideoKeyboard import CreateVideoKeyboard
+from ai_bot.msg.CreateVideoMsg import CreateVideoMsg
+from bot.Handler import Handler
+from bot import SECOND
+from telegram.constants import ParseMode
+
+
+class VideoHandler(Handler):
+
+    @staticmethod
+    async def create_video(update, context):
+        models_arr = ['Luma', 'Kling']
+
+        reply_markup = CreateVideoKeyboard()
+        reply_markup.add_option("models_arr", models_arr)
+
+        msg = CreateVideoMsg()
+
+        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