Added config and current commands, Close button for inline menus, pass to all __init__s

This commit is contained in:
BarsTiger
2023-03-08 16:10:22 +02:00
parent ebdd27ece8
commit b83e649349
19 changed files with 119 additions and 5 deletions

View File

@@ -0,0 +1 @@
pass

15
bot/callbacks/common.py Normal file
View File

@@ -0,0 +1,15 @@
from bot.common import dp
from bot.callbacks.factories.common import close_keyboard_data
from aiogram import types
from bot.utils.private_keyboard import other_user
async def on_close_keyboard(call: types.CallbackQuery):
if await other_user(call):
return
await call.message.delete()
def register():
dp.register_callback_query_handler(on_close_keyboard, close_keyboard_data.filter())

View File

@@ -0,0 +1 @@
pass

View File

@@ -0,0 +1 @@
pass

View File

@@ -0,0 +1,4 @@
from aiogram.utils.callback_data import CallbackData
close_keyboard_data = CallbackData("close_keyboard")

View File

@@ -0,0 +1,6 @@
from aiogram.utils.callback_data import CallbackData
prompt_settings_data = CallbackData("prompt_settings")
global_settings_data = CallbackData("global_settings")
admin_settings_data = CallbackData("admin_settings")

View File

@@ -6,8 +6,10 @@ from bot.keyboards.image_info import get_img_info_keyboard, get_img_back_keyboar
from bot.utils.cooldown import throttle from bot.utils.cooldown import throttle
from bot.utils.private_keyboard import other_user from bot.utils.private_keyboard import other_user
from bot.modules.api.objects.prompt_request import Generated from bot.modules.api.objects.prompt_request import Generated
from bot.utils.errorable_command import wrap_exception
@wrap_exception()
async def on_back(call: types.CallbackQuery, callback_data: dict): async def on_back(call: types.CallbackQuery, callback_data: dict):
p_id = callback_data['p_id'] p_id = callback_data['p_id']
if await other_user(call): if await other_user(call):
@@ -20,6 +22,7 @@ async def on_back(call: types.CallbackQuery, callback_data: dict):
) )
@wrap_exception()
@throttle(5) @throttle(5)
async def on_prompt_only(call: types.CallbackQuery, callback_data: dict): async def on_prompt_only(call: types.CallbackQuery, callback_data: dict):
p_id = callback_data['p_id'] p_id = callback_data['p_id']
@@ -36,6 +39,7 @@ async def on_prompt_only(call: types.CallbackQuery, callback_data: dict):
) )
@wrap_exception()
@throttle(5) @throttle(5)
async def on_full_info(call: types.CallbackQuery, callback_data: dict): async def on_full_info(call: types.CallbackQuery, callback_data: dict):
p_id = callback_data['p_id'] p_id = callback_data['p_id']
@@ -59,17 +63,21 @@ async def on_full_info(call: types.CallbackQuery, callback_data: dict):
) )
@wrap_exception()
@throttle(5) @throttle(5)
async def on_import(call: types.CallbackQuery, callback_data: dict): async def on_import(call: types.CallbackQuery, callback_data: dict):
p_id = callback_data['p_id'] p_id = callback_data['p_id']
if await other_user(call): if await other_user(call):
return return
prompt: Generated = db[DBTables.generated].get(p_id) generated: Generated = db[DBTables.generated].get(p_id)
prompt = generated.prompt
prompt.creator = call.from_user.id
db[DBTables.prompts][call.from_user.id] = prompt
await db[DBTables.config].write()
await call.message.edit_text( await call.message.edit_text(
f"😶‍🌫️ Not implemented yet", f"️🥐 Prompt imported",
parse_mode='html',
reply_markup=get_img_back_keyboard(p_id) reply_markup=get_img_back_keyboard(p_id)
) )

View File

@@ -5,11 +5,13 @@ def register_callbacks():
from bot.callbacks import ( from bot.callbacks import (
exception, exception,
image_info, image_info,
set_model set_model,
common
) )
exception.register() exception.register()
image_info.register() image_info.register()
set_model.register() set_model.register()
common.register()
print('[gray]All callbacks registered[/]') print('[gray]All callbacks registered[/]')

View File

@@ -0,0 +1,6 @@
from bot.common import dp
from .config_command import *
def register():
dp.register_message_handler(config_command, commands='config')

View File

@@ -0,0 +1,11 @@
from aiogram import types
from bot.db import db, DBTables
from bot.utils.cooldown import throttle
from bot.utils.errorable_command import wrap_exception
from bot.keyboards.config import get_config_keyboard
@wrap_exception()
@throttle(cooldown=60*60, admin_ids=db[DBTables.config].get('admins'), by_id=False)
async def config_command(message: types.Message):
await message.reply("⚙️ Configuration:", reply_markup=get_config_keyboard(message.from_id))

View File

@@ -2,6 +2,7 @@ help_data = {
'generate': 'Generate picture using configuration set by user. You can pass prompt also in command arguments or ' 'generate': 'Generate picture using configuration set by user. You can pass prompt also in command arguments or '
'use it without arguments to generate picture with prompt, that was used for last generation', 'use it without arguments to generate picture with prompt, that was used for last generation',
'imginfo': 'Get information about image, that was generated using this bot', 'imginfo': 'Get information about image, that was generated using this bot',
'current': 'Get your current prompt',
'setprompt': 'Set default prompt for images, will be overwritten if you specify prompt in generate command', 'setprompt': 'Set default prompt for images, will be overwritten if you specify prompt in generate command',
'setnegative': 'Set negative prompt', 'setnegative': 'Set negative prompt',
'setsize': 'Set size for image (in hxw format)', 'setsize': 'Set size for image (in hxw format)',
@@ -12,6 +13,7 @@ help_data = {
'setscale': 'Set CFG Scale (prompt stringency)', 'setscale': 'Set CFG Scale (prompt stringency)',
'setfaces': 'Set restore faces mode', 'setfaces': 'Set restore faces mode',
'status': 'Ping API endpoint host', 'status': 'Ping API endpoint host',
'config': 'Edit config using inline buttons',
'setmodel': '(global) Sets StableDiffusion model for all users. Can be used only once an hour', 'setmodel': '(global) Sets StableDiffusion model for all users. Can be used only once an hour',
'setendpoint': '(admin) Set StableDiffusion API endpoint', 'setendpoint': '(admin) Set StableDiffusion API endpoint',
'addadmin': '(admin) Add new admin - reply to message or type user ID', 'addadmin': '(admin) Add new admin - reply to message or type user ID',

View File

@@ -3,7 +3,7 @@ from rich import print
def register_handlers(): def register_handlers():
from bot.handlers import ( from bot.handlers import (
initialize, admin, help_command, txt2img, image_info initialize, admin, help_command, txt2img, image_info, config
) )
initialize.register() initialize.register()
@@ -11,5 +11,6 @@ def register_handlers():
help_command.register() help_command.register()
txt2img.register() txt2img.register()
image_info.register() image_info.register()
config.register()
print('[gray]All handlers registered[/]') print('[gray]All handlers registered[/]')

View File

@@ -2,6 +2,7 @@ from bot.common import dp
from .txt2img import generate_command from .txt2img import generate_command
from .set_model import set_model_command from .set_model import set_model_command
from .status import get_status from .status import get_status
from .current import get_current
from .set_settings import ( from .set_settings import (
set_height_command, set_negative_prompt_command, set_size_command, set_steps_command, set_width_command, set_height_command, set_negative_prompt_command, set_size_command, set_steps_command, set_width_command,
set_prompt_command, set_sampler_command, set_cfg_scale_command, set_restore_faces_command set_prompt_command, set_sampler_command, set_cfg_scale_command, set_restore_faces_command
@@ -21,3 +22,4 @@ def register():
dp.register_message_handler(set_settings.set_restore_faces_command, commands='setfaces') dp.register_message_handler(set_settings.set_restore_faces_command, commands='setfaces')
dp.register_message_handler(set_model.set_model_command, commands='setmodel') dp.register_message_handler(set_model.set_model_command, commands='setmodel')
dp.register_message_handler(status.get_status, commands='status') dp.register_message_handler(status.get_status, commands='status')
dp.register_message_handler(current.get_current, commands='current')

View File

@@ -0,0 +1,24 @@
from aiogram import types
from bot.db import db, DBTables
from bot.utils.cooldown import throttle
from bot.modules.api.objects.prompt_request import Prompt
from bot.utils.errorable_command import wrap_exception
@wrap_exception()
@throttle(cooldown=5, admin_ids=db[DBTables.config].get('admins'))
async def get_current(message: types.Message):
prompt: Prompt = db[DBTables.prompts].get(message.from_id)
if prompt is None:
await message.reply('❌ No prompts found for you')
await message.reply(
f"🖤 Prompt: {prompt.prompt} \n"
f"🐊 Negative: {prompt.negative_prompt} \n"
f"🪜 Steps: {prompt.steps} \n"
f"🧑‍🎨 CFG Scale: {prompt.cfg_scale} \n"
f"🖥️ Size: {prompt.width}x{prompt.height} \n"
f"😀 Restore faces: {'on' if prompt.restore_faces else 'off'} \n"
f"⚒️ Sampler: {prompt.sampler} \n",
parse_mode='html'
)

View File

@@ -0,0 +1 @@
pass

22
bot/keyboards/config.py Normal file
View File

@@ -0,0 +1,22 @@
from aiogram import types
from bot.db import db, DBTables
from bot.callbacks.factories.config import (prompt_settings_data, global_settings_data, admin_settings_data)
from bot.callbacks.factories.common import close_keyboard_data
def get_config_keyboard(user_id: int) -> types.InlineKeyboardMarkup:
buttons = [
types.InlineKeyboardButton("Prompt settings", callback_data=prompt_settings_data.new()),
types.InlineKeyboardButton("Global settings", callback_data=global_settings_data.new())
]
if user_id in db[DBTables.config].get('admins'):
buttons.append(
types.InlineKeyboardButton("Admin settings", callback_data=admin_settings_data.new())
)
buttons.append(
types.InlineKeyboardButton("🔻 Close", callback_data=close_keyboard_data.new())
)
keyboard = types.InlineKeyboardMarkup(row_width=1)
keyboard.add(*buttons)
return keyboard

View File

@@ -1,6 +1,7 @@
from aiogram import types from aiogram import types
from bot.db import db, DBTables from bot.db import db, DBTables
from bot.callbacks.factories.set_model import set_model_page, set_model from bot.callbacks.factories.set_model import set_model_page, set_model
from bot.callbacks.factories.common import close_keyboard_data
def get_set_model_keyboard(page: int) -> types.InlineKeyboardMarkup: def get_set_model_keyboard(page: int) -> types.InlineKeyboardMarkup:
@@ -28,4 +29,8 @@ def get_set_model_keyboard(page: int) -> types.InlineKeyboardMarkup:
else: else:
keyboard.add(*models_buttons) keyboard.add(*models_buttons)
keyboard.add(
types.InlineKeyboardButton("🔻 Cancel", callback_data=close_keyboard_data.new())
)
return keyboard return keyboard

View File

@@ -0,0 +1 @@
pass

View File

@@ -0,0 +1 @@
pass