Whitelist works as it should

This commit is contained in:
BarsTiger
2023-05-06 15:11:35 +03:00
parent 2789ef6461
commit c332247f1b
2 changed files with 29 additions and 26 deletions

View File

@@ -1,5 +1,5 @@
from aiogram import types from aiogram import types
from bot.db import db, DBTables, decrypt from bot.db import db, DBTables
from bot.utils.cooldown import throttle from bot.utils.cooldown import throttle
from bot.keyboards.set_model import get_set_model_keyboard from bot.keyboards.set_model import get_set_model_keyboard
from bot.modules.api.models import get_models from bot.modules.api.models import get_models
@@ -9,7 +9,9 @@ from bot.utils.errorable_command import wrap_exception
@wrap_exception() @wrap_exception()
@throttle(cooldown=5*60, admin_ids=db[DBTables.config].get('admins'), by_id=False) @throttle(cooldown=5*60, admin_ids=db[DBTables.config].get('admins'), by_id=False)
async def set_model_command(message: types.Message): async def set_model_command(message: types.Message):
if (message.chat.id not in db[DBTables.config]['whitelist'] and message.from_id not in db[DBTables.config]['whitelist']): if db[DBTables.config].get('whitelist') and \
(message.chat.id not in db[DBTables.config]['whitelist'] and
message.from_id not in db[DBTables.config]['whitelist']):
await message.reply('❌ You are not on the white list, access denied. Contact admin @kilisauros for details') await message.reply('❌ You are not on the white list, access denied. Contact admin @kilisauros for details')
return return
models = await get_models() models = await get_models()
@@ -18,5 +20,5 @@ async def set_model_command(message: types.Message):
else: else:
await message.reply('❌ No models available') await message.reply('❌ No models available')
return return
await message.reply("Examples of models (with additional info): https://telegra.ph/Opisanie-raboty-modelej-05-03")
await message.reply("🪄 You can choose model from available:", reply_markup=get_set_model_keyboard(0)) await message.reply("🪄 You can choose model from available:", reply_markup=get_set_model_keyboard(0))

View File

@@ -2,6 +2,7 @@ import re
from bot.common import bot from bot.common import bot
from aiogram import types from aiogram import types
from bot.db import db, DBTables from bot.db import db, DBTables
from bot.config import ARCHIVE_CHAT
from bot.utils.cooldown import throttle from bot.utils.cooldown import throttle
from bot.modules.api.txt2img import txt2img from bot.modules.api.txt2img import txt2img
from bot.modules.api.objects.get_prompt import get_prompt from bot.modules.api.objects.get_prompt import get_prompt
@@ -9,8 +10,6 @@ from bot.modules.api.objects.prompt_request import Generated
from bot.modules.api.status import wait_for_status from bot.modules.api.status import wait_for_status
from bot.keyboards.image_info import get_img_info_keyboard from bot.keyboards.image_info import get_img_info_keyboard
from bot.utils.errorable_command import wrap_exception from bot.utils.errorable_command import wrap_exception
from bot.callbacks.factories.image_info import (prompt_only, full_prompt, import_prompt, back)
@wrap_exception([ValueError], custom_loading=True) @wrap_exception([ValueError], custom_loading=True)
@@ -21,8 +20,9 @@ async def generate_command(message: types.Message):
await message.reply('💔 Generation is disabled by admins now. Try again later') await message.reply('💔 Generation is disabled by admins now. Try again later')
await temp_message.delete() await temp_message.delete()
return return
elif (message.chat.id not in db[DBTables.config]['whitelist'] and message.from_id not in db[DBTables.config]['whitelist']): elif db[DBTables.config].get('whitelist') and (message.chat.id not in db[DBTables.config]['whitelist']
await message.reply('❌You are not on the white list, access denied. Contact admin @kilisauros for details') and message.from_id not in db[DBTables.config]['whitelist']):
await message.reply('❌ You are not on the white list, access denied. ')
await temp_message.delete() await temp_message.delete()
return return
@@ -47,23 +47,24 @@ async def generate_command(message: types.Message):
image = await txt2img(prompt) image = await txt2img(prompt)
image_message = await message.reply_photo(photo=image[0]) image_message = await message.reply_photo(photo=image[0])
#Send photo to SD Image Archive if ARCHIVE_CHAT:
archive_message = (
archive_message = f'User ID: {message.from_id} \n \ f"User ID: {message.from_id}"
User nickname: {message.from_user.full_name} \n \ f"User nickname: {message.from_user.full_name}"
User username: @{message.from_user.username} \n \ f"User username: @{message.from_user.username}"
Chat ID: {message.chat.id} \n \ f"Chat ID: {message.chat.id}"
Chat title: {message.chat.title} \n \ f"Chat title: {message.chat.title}"
Info: \n \ f"Info:"
🖤 Prompt: {prompt.prompt} \n \ f"🖤 Prompt: {prompt.prompt}"
🐊 Negative: {prompt.negative_prompt} \n \ f"🐊 Negative: {prompt.negative_prompt}"
💫 Model: In development \n \ f"💫 Model: In development"
🪜 Steps: {prompt.steps} \n \ f"🪜 Steps: {prompt.steps}"
🧑‍🎨 CFG Scale: {prompt.cfg_scale} \n \ f"🧑‍🎨 CFG Scale: {prompt.cfg_scale}"
🖥️ Size: {prompt.width}x{prompt.height} \n \ f"🖥️ Size: {prompt.width}x{prompt.height}"
😀 Restore faces: {prompt.restore_faces} \n \ f"😀 Restore faces: {prompt.restore_faces}"
⚒️ Sampler: {prompt.sampler} \n ' f"⚒️ Sampler: {prompt.sampler}"
await bot.send_photo(-929754401, photo=image[0], caption=archive_message) )
await bot.send_photo(ARCHIVE_CHAT, photo=image[0], caption=archive_message)
db[DBTables.generated][image_message.photo[0].file_unique_id] = Generated( db[DBTables.generated][image_message.photo[0].file_unique_id] = Generated(
prompt=prompt, prompt=prompt,