Добавлено ограничение пользователей #7

Merged
soaska merged 18 commits from kls/sd_telegram:master into beta 2024-02-11 15:55:28 +00:00
2 changed files with 33 additions and 31 deletions
Showing only changes of commit ba2e25d066 - Show all commits

61
bot.py
View File

@ -61,6 +61,7 @@ OK_APPLICATION_KEY = os.getenv('OK_APPLICATION_KEY')
OK_APPLICATION_SECRET_KEY = os.getenv('OK_APPLICATION_SECRET_KEY')
OK_GROUP_ID = os.getenv('OK_GROUP_ID')
ARRAY_INLINE = []
ALLOWED_USERS = json.loads(os.environ['ALLOWED_USERS'])
bot = Bot(token=API_BOT_TOKEN)
storage = MemoryStorage()
@ -591,8 +592,8 @@ async def show_thumbs(chat_id, res):
# start или help
@dp.callback_query_handler(text="help")
@dp.message_handler(commands=["help"])
@dp.message_handler(commands=["start"])
@dp.message_handler(lambda message: message.from_user.username in ALLOWED_USERS, commands=["help"])
@dp.message_handler(lambda message: message.from_user.username in ALLOWED_USERS, commands=["start"])
async def cmd_start(message: Union[types.Message, types.CallbackQuery]) -> None:
logging.info("cmd_start")
txt = "Это бот для локального запуска SD\n" + getTxt()
@ -600,7 +601,7 @@ async def cmd_start(message: Union[types.Message, types.CallbackQuery]) -> None:
# TODO optimize
# Запуск/Остановка SD. Завязываемся на глобальную иконку sd
@dp.message_handler(commands=["stop"])
@dp.message_handler(lambda message: message.from_user.username in ALLOWED_USERS, commands=["stop"])
@dp.callback_query_handler(text="sd")
async def inl_sd(message: Union[types.Message, types.CallbackQuery]) -> None:
logging.info("inl_sd")
@ -718,7 +719,7 @@ async def inl_uplora(callback: types.CallbackQuery) -> None:
await callback.message.reply(f"Файл '{file_name}' загружен в {folder_path}")
# Вызов reset_param, сброс JSON
@dp.message_handler(commands=["reset_param"])
@dp.message_handler(lambda message: message.from_user.username in ALLOWED_USERS, commands=["reset_param"])
@dp.callback_query_handler(text="reset_param")
async def inl_reset_param(message: Union[types.Message, types.CallbackQuery]) -> None:
logging.info("inl_reset_param")
@ -733,7 +734,7 @@ async def inl_reset_param(message: Union[types.Message, types.CallbackQuery]) ->
await getKeyboardUnion(txt, message, keyboard, '')
# Вызов fast_param, быстрые настройки
@dp.message_handler(commands=["fast_param"])
@dp.message_handler(lambda message: message.from_user.username in ALLOWED_USERS, commands=["fast_param"])
@dp.callback_query_handler(text="fast_param")
async def inl_fast_param(message: Union[types.Message, types.CallbackQuery]) -> None:
logging.info("inl_fast_param")
@ -741,13 +742,13 @@ async def inl_fast_param(message: Union[types.Message, types.CallbackQuery]) ->
await getKeyboardUnion('Выбери быстрые настройки', message, keyboard, '')
# Список быстрых настроек
@dp.message_handler(commands=["fp_comp"])
@dp.message_handler(commands=["fp_mobile"])
@dp.message_handler(commands=["fp_no_hr"])
@dp.message_handler(commands=["fp_sdxl"])
@dp.message_handler(commands=["fp_big"])
@dp.message_handler(commands=["fp_inc"])
@dp.message_handler(commands=["fp_wh"])
@dp.message_handler(lambda message: message.from_user.username in ALLOWED_USERS, commands=["fp_comp"])
@dp.message_handler(lambda message: message.from_user.username in ALLOWED_USERS, commands=["fp_mobile"])
@dp.message_handler(lambda message: message.from_user.username in ALLOWED_USERS, commands=["fp_no_hr"])
@dp.message_handler(lambda message: message.from_user.username in ALLOWED_USERS, commands=["fp_sdxl"])
@dp.message_handler(lambda message: message.from_user.username in ALLOWED_USERS, commands=["fp_big"])
@dp.message_handler(lambda message: message.from_user.username in ALLOWED_USERS, commands=["fp_inc"])
@dp.message_handler(lambda message: message.from_user.username in ALLOWED_USERS, commands=["fp_wh"])
@dp.callback_query_handler(text="fp_comp")
@dp.callback_query_handler(text="fp_mobile")
@dp.callback_query_handler(text="fp_no_hr")
@ -875,7 +876,7 @@ async def inl_fp(message: Union[types.Message, types.CallbackQuery]) -> None:
await getKeyboardUnion(txt, message, keyboard, '')
# Обработчик команды /skip
@dp.message_handler(commands=["skip"])
@dp.message_handler(lambda message: message.from_user.username in ALLOWED_USERS, commands=["skip"])
@dp.callback_query_handler(text="skip")
async def inl_skip(message: Union[types.Message, types.CallbackQuery]) -> None:
logging.info('inl_skip')
@ -895,7 +896,7 @@ async def inl_skip(message: Union[types.Message, types.CallbackQuery]) -> None:
reply_markup=getStart(),
)
@dp.message_handler(commands=["gen"])
@dp.message_handler(lambda message: message.from_user.username in ALLOWED_USERS, commands=["gen"])
@dp.callback_query_handler(text="gen")
async def inl_gen(message: Union[types.Message, types.CallbackQuery]) -> None:
if hasattr(message, "content_type"):
@ -1044,7 +1045,7 @@ async def cmd_prompt(message: Union[types.Message, types.CallbackQuery]) -> None
await getKeyboardUnion("Опции", message, keyboard)
# Получить опции
@dp.message_handler(commands=["opt"])
@dp.message_handler(lambda message: message.from_user.username in ALLOWED_USERS, commands=["opt"])
@dp.callback_query_handler(text="opt")
async def cmd_opt(message: Union[types.Message, types.CallbackQuery]) -> None:
logging.info("cmd_opt")
@ -1052,7 +1053,7 @@ async def cmd_opt(message: Union[types.Message, types.CallbackQuery]) -> None:
await getKeyboardUnion("Опции", message, keyboard)
# Вызов sttngs
@dp.message_handler(commands=["sttngs"])
@dp.message_handler(lambda message: message.from_user.username in ALLOWED_USERS, commands=["sttngs"])
@dp.callback_query_handler(text="sttngs")
async def inl_sttngs(message: Union[types.Message, types.CallbackQuery]) -> None:
logging.info("inl_sttngs")
@ -1060,7 +1061,7 @@ async def inl_sttngs(message: Union[types.Message, types.CallbackQuery]) -> None
await getKeyboardUnion("Настройки", message, keyboard)
# Вызов script
@dp.message_handler(commands=["scrpts"])
@dp.message_handler(lambda message: message.from_user.username in ALLOWED_USERS, commands=["scrpts"])
@dp.callback_query_handler(text="scrpts")
async def inl_scrpts(message: Union[types.Message, types.CallbackQuery]) -> None:
logging.info("inl_scrpts")
@ -1068,7 +1069,7 @@ async def inl_scrpts(message: Union[types.Message, types.CallbackQuery]) -> None
await getKeyboardUnion("Скрипты", message, keyboard)
# Вызов get_models
@dp.message_handler(commands=["mdl"])
@dp.message_handler(lambda message: message.from_user.username in ALLOWED_USERS, commands=["mdl"])
@dp.callback_query_handler(text="mdl")
async def inl_mdl(message: Union[types.Message, types.CallbackQuery]) -> None:
logging.info("inl_mdl")
@ -1083,8 +1084,8 @@ async def inl_mdl(message: Union[types.Message, types.CallbackQuery]) -> None:
await getKeyboardUnion("Turn on SD"+sd, message, keyboard)
# Вызов get_samplers
@dp.message_handler(commands=["smplr"])
@dp.message_handler(commands=["sampler_name"])
@dp.message_handler(lambda message: message.from_user.username in ALLOWED_USERS, commands=["smplr"])
@dp.message_handler(lambda message: message.from_user.username in ALLOWED_USERS, commands=["sampler_name"])
@dp.callback_query_handler(text="smplr")
async def inl_smplr(message: Union[types.Message, types.CallbackQuery]) -> None:
logging.info("inl_smplr")
@ -1099,8 +1100,8 @@ async def inl_smplr(message: Union[types.Message, types.CallbackQuery]) -> None:
await getKeyboardUnion("Turn on SD"+sd, message, keyboard)
# Вызов get_hr_list
@dp.message_handler(commands=["hr"])
@dp.message_handler(commands=["hr_upscaler"])
@dp.message_handler(lambda message: message.from_user.username in ALLOWED_USERS, commands=["hr"])
@dp.message_handler(lambda message: message.from_user.username in ALLOWED_USERS, commands=["hr_upscaler"])
@dp.callback_query_handler(text="hr")
async def inl_hr(message: Union[types.Message, types.CallbackQuery]) -> None:
logging.info("inl_hr")
@ -1128,14 +1129,14 @@ async def inl_change_param(callback: types.CallbackQuery) -> None:
)
# script random gen from models
@dp.message_handler(commands=["rnd_mdl"])
@dp.message_handler(lambda message: message.from_user.username in ALLOWED_USERS, commands=["rnd_mdl"])
@dp.callback_query_handler(text='rnd_mdl')
async def inl_rnd_mdl(message: Union[types.Message, types.CallbackQuery]) -> None:
logging.info('inl_rnd_mdl')
await rnd_script(message, 'models')
# script random gen from models
@dp.message_handler(commands=["rnd_smp"])
@dp.message_handler(lambda message: message.from_user.username in ALLOWED_USERS, commands=["rnd_smp"])
@dp.callback_query_handler(text='rnd_smp')
async def inl_rnd_smp(message: Union[types.Message, types.CallbackQuery]) -> None:
logging.info('inl_rnd_smp')
@ -1189,7 +1190,7 @@ async def inf_func(chatId):
)
# script random infinity gen from https://random-word-api.herokuapp.com/word?lang=en
@dp.message_handler(commands=["inf"])
@dp.message_handler(lambda message: message.from_user.username in ALLOWED_USERS, commands=["inf"])
@dp.callback_query_handler(text='inf')
async def inl_rnd_inf(message: Union[types.Message, types.CallbackQuery]) -> None:
logging.info('inl_rnd_inf')
@ -1228,7 +1229,7 @@ async def inl_rnd_inf(message: Union[types.Message, types.CallbackQuery]) -> Non
await inf_func(chatId)
# Получить LORA
@dp.message_handler(commands=["get_lora"])
@dp.message_handler(lambda message: message.from_user.username in ALLOWED_USERS, commands=["get_lora"])
@dp.callback_query_handler(text="get_lora")
async def getLora(message: Union[types.Message, types.CallbackQuery]) -> None:
logging.info("getLora")
@ -1252,7 +1253,7 @@ async def getLora(message: Union[types.Message, types.CallbackQuery]) -> None:
await getKeyboardUnion(arr, message, keyboard)
# Рандомный промпт с lexica.art на основе data['prompt']
@dp.message_handler(commands=["lxc_prompt"])
@dp.message_handler(lambda message: message.from_user.username in ALLOWED_USERS, commands=["lxc_prompt"])
@dp.callback_query_handler(text="lxc_prompt")
async def get_lxc_prompt(message: Union[types.Message, types.CallbackQuery]) -> None:
logging.info("get_lxc_prompt")
@ -1352,7 +1353,7 @@ async def handle_channel_post(message: types.Message):
)
# Ввели любой текст
@dp.message_handler(lambda message: True)
@dp.message_handler(lambda message: True and message.from_user.username in ALLOWED_USERS)
async def change_json(message: types.Message):
logging.info("change_json")
keyboard = InlineKeyboardMarkup(inline_keyboard=[getSet(0), getOpt(0), getStart(0)])
@ -1409,7 +1410,7 @@ async def change_json(message: types.Message):
)
# Ввели ответ на change_json
@dp.message_handler(state=Form)
@dp.message_handler(lambda message: message.from_user.username in ALLOWED_USERS, state=Form)
async def answer_handler(message: types.Message, state: FSMContext):
logging.info('answer_handler')
keyboard = InlineKeyboardMarkup(inline_keyboard=[getSet(0), getOpt(0), getStart(0)])
@ -1432,7 +1433,7 @@ async def answer_handler(message: types.Message, state: FSMContext):
f"JSON параметры:\n{getJson()}\n{getJson(1)}", reply_markup=keyboard
)
@dp.message_handler(content_types=['document'])
@dp.message_handler(lambda message: message.from_user.username in ALLOWED_USERS, content_types=['document'])
async def handle_file(message: types.Message):
logging.info('handle_file')

View File

@ -9,4 +9,5 @@ SD_HOST = '127.0.0.1'
SD_PORT = '7861'
SD_USERNAME = 'user'
SD_PASSWORD = '1234'
DEBUG = 'no' # 'yes' to debug
DEBUG = 'no' # 'yes' to debug
ALLOWED_USERS = '["username1", "username2", "usernameN"]'