python_ai_bot/ai_bot/handlers/FluxHandler.py

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)
            ],
        }