first

ai_bot/__init__.py

@@ -0,0 +1,15 @@
+from ai_bot.handlers.FluxHandler import FluxHandler
+from ai_bot.handlers.MidjourneyHandler import MidjourneyHandler
+from piapi_ai.PiapiAiApi import PiapiAiApi
+
+
+states_instance_arr = [
+    FluxHandler,
+    MidjourneyHandler,
+]
+
+states = {}
+
+for state in states_instance_arr:
+    states.update(state.get_states())
+

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

ai_bot/keyboards/CreateImageKeyboard.py

@@ -0,0 +1,14 @@
+from bot.Keyboard import Keyboard
+from telegram import InlineKeyboardButton, InlineKeyboardMarkup
+
+
+class CreateImageKeyboard(Keyboard):
+
+    def get_keyboard(self):
+        keyboard = []
+        models = self.get_option("models_arr")
+
+        for model in models:
+            keyboard.append([InlineKeyboardButton(model, callback_data='start_model_scenario?model={model}'.format(model=model))])
+
+        return keyboard

ai_bot/keyboards/CreateVideoKeyboard.py

@@ -0,0 +1,14 @@
+from bot.Keyboard import Keyboard
+from telegram import InlineKeyboardButton, InlineKeyboardMarkup
+
+
+class CreateVideoKeyboard(Keyboard):
+
+    def get_keyboard(self):
+        keyboard = []
+        models = self.get_option("models_arr")
+
+        for model in models:
+            keyboard.append([InlineKeyboardButton(model, callback_data='start_model_scenario?model={model}'.format(model=model))])
+
+        return keyboard

ai_bot/keyboards/MainKeyboard.py

@@ -0,0 +1,12 @@
+from bot.Keyboard import Keyboard
+from telegram import InlineKeyboardButton, InlineKeyboardMarkup
+
+
+class MainKeyboard(Keyboard):
+
+    def get_keyboard(self):
+        return [
+            [
+                InlineKeyboardButton("Меню", callback_data='menu'),
+            ],
+        ]

ai_bot/keyboards/MenuKeyboard.py

@@ -0,0 +1,13 @@
+from bot.Keyboard import Keyboard
+from telegram import InlineKeyboardButton, InlineKeyboardMarkup
+
+
+class MenuKeyboard(Keyboard):
+
+    def get_keyboard(self):
+        return [
+            [
+                InlineKeyboardButton("Сгенерировать видео", callback_data='create_video'),
+                InlineKeyboardButton("Сгенерировать изображение", callback_data='create_img')
+            ],
+        ]

ai_bot/keyboards/flux/SelectTypeRequestKeyboard.py

@@ -0,0 +1,15 @@
+from bot.Keyboard import Keyboard
+from telegram import InlineKeyboardButton, InlineKeyboardMarkup
+
+
+class SelectTypeRequestKeyboard(Keyboard):
+
+    def get_keyboard(self):
+        return [
+            [
+                InlineKeyboardButton("Сгенерировать изображение из текста", callback_data='set_type_request?type=txt2img'),
+            ],
+            [
+                InlineKeyboardButton("Сгенерировать изображение на основе изображения", callback_data='set_type_request?type=img2img')
+            ]
+        ]

ai_bot/msg/CreateImageMsg.py

@@ -0,0 +1,13 @@
+from bot.HtmlMsg import HtmlMsg
+
+
+class CreateImageMsg(HtmlMsg):
+
+    def get_msg(self, data=None) -> str:
+        text = ("<strong>Flux.1 Pro</strong> — это профессиональная модель текста в изображение, оптимизированная для коммерческих приложений. "
+                "Эта премиальная модель обеспечивает исключительное качество изображения и производительность, "
+                "что делает ее идеальной для профессиональной творческой работы.\n\n"
+                "<strong>MidJourney</strong> – это инструмент, который с помощью искусственного интеллекта создаёт уникальные изображения по вашим запросам. "
+                "Просто опишите, что хотите увидеть, а ИИ превратит ваши слова в произведение искусства!")
+
+        return text

ai_bot/msg/CreateVideoMsg.py

@@ -0,0 +1,12 @@
+from bot.HtmlMsg import HtmlMsg
+
+
+class CreateVideoMsg(HtmlMsg):
+
+    def get_msg(self, data=None) -> str:
+        text = ("<strong>Luma AI</strong> — эта одна из первых нейросетей для генерации видео на такой же архитектуре, которую можно попробовать самому."
+                "Она создает реалистичные ролики, которые выглядят качественнее, чем у конкуренто.\n\n"
+                "<strong>Kling</strong> - это модель генерации видео, разработанная командой Kuaishou, "
+                "обладающая мощными возможностями генерации видео и позволяющая пользователям легко и эффективно создавать художественные видео.")
+
+        return text

ai_bot/msg/MainMsg.py

@@ -0,0 +1,22 @@
+from bot.HtmlMsg import HtmlMsg
+
+
+class MainMsg(HtmlMsg):
+
+    def get_msg(self, data=None) -> str:
+        if data is None:
+            data = {}
+
+        text = ("<strong>🌟 Что умеет бот?</strong> \n"
+                " - Генерирует тексты, статьи, описания и многое другое. \n"
+                " - Создаёт уникальные изображения по вашему запросу. \n"
+                " - Помогает анализировать данные и проводить расчёты. \n"
+                " - Отвечает на вопросы, обучает и даёт советы в любых темах. \n"
+                " - Автоматизирует рутинные задачи – сэкономьте своё время! \n\n"
+                "<strong>💡 Почему стоит попробовать?</strong> \n"
+                " - Легко и быстро – взаимодействие с ИИ через удобный интерфейс Telegram. \n"
+                " - Без навыков программирования – всё просто и понятно. \n"
+                " - Работает 24/7 – ваш помощник всегда на связи. \n"
+                " - Подходит для бизнеса, учёбы, творчества и повседневных задач. \n")
+
+        return text

ai_bot/msg/flux/SelectTypeMsg.py

@@ -0,0 +1,13 @@
+from bot.DaMsg import DaMsg
+from bot.HtmlMsg import HtmlMsg
+
+
+class SelectTypeMsg(HtmlMsg):
+
+    def get_msg(self, data=None) -> str:
+        if data is None:
+            data = {}
+
+        text = ("Выберите тип запроса.")
+
+        return text

ai_bot/states/FluxStates.py

@@ -0,0 +1,33 @@
+from bot.States import States
+from telegram.ext import MessageHandler, filters, CallbackQueryHandler
+
+FLUX_SET_PROMPT = 31
+FLUX_SET_SUBMODEL = 32
+FLUX_SET_TYPE_REQUEST = 33
+FLUX_SET_TEXT_PROMPT = 34
+FLUX_SET_IMG_TEXT_PROMPT = 35
+
+def set_states():
+    states_arr = {
+        "set_prompt": FLUX_SET_PROMPT,
+        "set_submodel": FLUX_SET_SUBMODEL,
+        "set_type_request": FLUX_SET_TYPE_REQUEST,
+        "set_text_prompt": FLUX_SET_TEXT_PROMPT,
+        "set_img_text_prompt": FLUX_SET_IMG_TEXT_PROMPT,
+    }
+
+    return states_arr
+
+
+class FluxStates(States):
+
+    def __init__(self):
+        pass
+
+    @staticmethod
+    def get_state_by_key(key: str):
+        states = set_states()
+        if key in states:
+            return states[key]
+
+        return None

ai_bot/states/KlingStates.py

@@ -0,0 +1,27 @@
+from bot.States import States
+
+KLING_SET_PROMPT = 51
+KLING_SET_TEXT_PROMPT = 52
+
+
+def set_states():
+    states_arr = {
+        "set_prompt": KLING_SET_PROMPT,
+        "set_text_prompt": KLING_SET_TEXT_PROMPT,
+    }
+
+    return states_arr
+
+
+class KlingStates(States):
+
+    def __init__(self):
+        pass
+
+    @staticmethod
+    def get_state_by_key(key: str):
+        states = set_states()
+        if key in states:
+            return states[key]
+
+        return None

ai_bot/states/MidjourneyStates.py

@@ -0,0 +1,24 @@
+from bot.States import States
+
+MIDJOURNEY_SET_PROMPT = 41
+
+def set_states():
+    states_arr = {
+        "set_prompt": MIDJOURNEY_SET_PROMPT
+    }
+
+    return states_arr
+
+
+class MidjourneyStates(States):
+
+    def __init__(self):
+        pass
+
+    @staticmethod
+    def get_state_by_key(key: str):
+        states = set_states()
+        if key in states:
+            return states[key]
+
+        return None