54 lines
1.9 KiB
Python
54 lines
1.9 KiB
Python
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
|