tgbots/TeleDiffusionBot
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

config command now works fully

Browse Source
This commit is contained in:
BarsTiger
2023-03-11 23:13:55 +02:00
parent b83e649349
commit 66b6da88ad
18 changed files with 297 additions and 63 deletions
Download Patch File Download Diff File Expand all files Collapse all files

25
bot/handlers/admin/aliases.py
View File

@@ -6,17 +6,18 @@ from bot.utils.cooldown import throttle
@throttle(5)
async def set_endpoint(message: types.Message):
async def set_endpoint(message: types.Message, is_command: bool = True):
if message.from_id not in db[DBTables.config].get('admins') and message.from_id != ADMIN:
await message.reply('❌ You are not permitted to do that. '
'It is only for this bot instance maintainers and admins')
return
if not message.get_args() or not validators.url(message.get_args()):
if not (message.get_args() if is_command else message.text) or not \
validators.url((message.get_args() if is_command else message.text)):
await message.reply("❌ Specify correct url for endpoint")
return
db[DBTables.config]['endpoint'] = encrypt(message.get_args())
db[DBTables.config]['endpoint'] = encrypt((message.get_args() if is_command else message.text))
await db[DBTables.config].write()
@@ -24,18 +25,19 @@ async def set_endpoint(message: types.Message):
@throttle(5)
async def add_admin(message: types.Message):
async def add_admin(message: types.Message, is_command: bool = True):
if message.from_id != ADMIN:
await message.reply('❌ You are not permitted to do that. It is only for main admin')
return
if not message.get_args().isdecimal() and not hasattr(message.reply_to_message, 'text'):
if not (message.get_args() if is_command else message.text).isdecimal() and not \
hasattr(message.reply_to_message, 'text'):
await message.reply('❌ Put new admin ID to command arguments or answer to users message')
return
elif not message.get_args().isdecimal():
elif not (message.get_args() if is_command else message.text).isdecimal():
ID = message.reply_to_message.from_id
elif not hasattr(message.reply_to_message, 'text'):
ID = int(message.get_args())
ID = int((message.get_args() if is_command else message.text))
if not isinstance(db[DBTables.config].get('admins'), list):
db[DBTables.config]['admins'] = list()
@@ -54,18 +56,19 @@ async def add_admin(message: types.Message):
@throttle(5)
async def remove_admin(message: types.Message):
async def remove_admin(message: types.Message, is_command: bool = True):
if message.from_id != ADMIN:
await message.reply('❌ You are not permitted to do that. It is only for main admin')
return
if not message.get_args().isdecimal() and not hasattr(message.reply_to_message, 'text'):
if not (message.get_args() if is_command else message.text).isdecimal() and not \
hasattr(message.reply_to_message, 'text'):
await message.reply('❌ Put admin ID to command arguments or answer to users message')
return
elif not message.get_args().isdecimal():
elif not (message.get_args() if is_command else message.text).isdecimal():
ID = message.reply_to_message.from_id
elif not hasattr(message.reply_to_message, 'text'):
ID = int(message.get_args())
ID = int((message.get_args() if is_command else message.text))
if not isinstance(db[DBTables.config].get('admins'), list):
db[DBTables.config]['admins'] = list()

5
bot/handlers/admin/reset.py
View File

@@ -5,12 +5,15 @@ from bot.utils.cooldown import throttle
@throttle(5)
async def resetqueue(message: types.Message):
async def resetqueue(message: types.Message, is_command: bool = True):
if message.from_id not in db[DBTables.config].get('admins') and message.from_id != ADMIN:
await message.reply('❌ You are not permitted to do that. '
'It is only for this bot instance maintainers and admins')
return
if is_command and message.text.lower() != 'reset':
return
db[DBTables.queue]['n'] = 0
await db[DBTables.config].write()

3
bot/handlers/initialize/__init__.py
View File

@@ -1,8 +1,11 @@
from bot.common import dp, bot
from .start import *
from .all_messages import *
from bot.db import db, DBTables
def register():
dp.register_message_handler(all_messages.sync_db_filter, lambda *_: not hasattr(bot, 'cloudmeta_message_text'))
dp.register_message_handler(all_messages.on_action_message,
lambda message: str(message.from_id) in list(db[DBTables.actions].keys()))
dp.register_message_handler(start.start_command, commands='start')

15
bot/handlers/initialize/all_messages.py
View File

@@ -1,4 +1,6 @@
from aiogram.types import Message
from bot.db import db, DBTables
from bot.modules.api.objects.action import Action
async def sync_db_filter(message: Message):
@@ -10,3 +12,16 @@ async def sync_db_filter(message: Message):
f'If you tried to run a command, run it again')
if not await ping():
await message.reply('⚠️ Warning: StableDiffusion server is turned off or api endpoint is incorrect')
async def on_action_message(message: Message):
action: Action = db[DBTables.actions].get(message.from_id)
if not action:
return
if action.chat_id != message.chat.id:
return
del db[DBTables.actions][message.from_id]
import bot.callbacks
assert bot.callbacks
await eval(f"bot.callbacks.{action.action_module}.{action.action}(message, '{action.overload}')")

60
bot/handlers/txt2img/set_settings.py
View File

@@ -6,9 +6,9 @@ from bot.utils.errorable_command import wrap_exception
@wrap_exception(custom_loading=True)
async def _set_property(message: types.Message, prop: str, value=None):
async def _set_property(message: types.Message, prop: str, value=None, is_command: bool = True):
temp_message = await message.reply(f"⏳ Setting {prop}...")
if not message.get_args():
if not message.get_args() and is_command:
await temp_message.edit_text("😶‍🌫️ Specify arguments for this command. Check /help")
return
@@ -20,9 +20,9 @@ async def _set_property(message: types.Message, prop: str, value=None):
f"jacket, outdoors, streets</code>", parse_mode='HTML')
return
elif prompt is None:
prompt = Prompt(message.get_args(), creator=message.from_id)
prompt = Prompt((message.get_args() if is_command else message.text), creator=message.from_id)
prompt.__setattr__(prop, message.get_args() if value is None else value)
prompt.__setattr__(prop, (message.get_args() if is_command else message.text) if value is None else value)
prompt.creator = message.from_id
db[DBTables.prompts][message.from_id] = prompt
@@ -33,19 +33,19 @@ async def _set_property(message: types.Message, prop: str, value=None):
@throttle(cooldown=5, admin_ids=db[DBTables.config].get('admins'))
async def set_prompt_command(message: types.Message):
await _set_property(message, 'prompt')
async def set_prompt_command(message: types.Message, is_command: bool = True):
await _set_property(message, 'prompt', is_command=is_command)
@throttle(cooldown=5, admin_ids=db[DBTables.config].get('admins'))
async def set_negative_prompt_command(message: types.Message):
await _set_property(message, 'negative_prompt')
async def set_negative_prompt_command(message: types.Message, is_command: bool = True):
await _set_property(message, 'negative_prompt', is_command=is_command)
@throttle(cooldown=5, admin_ids=db[DBTables.config].get('admins'))
async def set_steps_command(message: types.Message):
async def set_steps_command(message: types.Message, is_command: bool = True):
try:
_ = int(message.get_args())
_ = int((message.get_args() if is_command else message.text))
except Exception as e:
assert e
await message.reply('❌ Specify number as argument')
@@ -55,25 +55,25 @@ async def set_steps_command(message: types.Message):
await message.reply('❌ Specify number <= 30')
return
await _set_property(message, 'steps')
await _set_property(message, 'steps', is_command=is_command)
@throttle(cooldown=5, admin_ids=db[DBTables.config].get('admins'))
async def set_cfg_scale_command(message: types.Message):
async def set_cfg_scale_command(message: types.Message, is_command: bool = True):
try:
_ = int(message.get_args())
_ = int((message.get_args() if is_command else message.text))
except Exception as e:
assert e
await message.reply('❌ Specify number as argument')
return
await _set_property(message, 'cfg_scale')
await _set_property(message, 'cfg_scale', is_command=is_command)
@throttle(cooldown=5, admin_ids=db[DBTables.config].get('admins'))
async def set_width_command(message: types.Message):
async def set_width_command(message: types.Message, is_command: bool = True):
try:
_ = int(message.get_args())
_ = int((message.get_args() if is_command else message.text))
except Exception as e:
assert e
await message.reply('❌ Specify number as argument')
@@ -83,13 +83,13 @@ async def set_width_command(message: types.Message):
await message.reply('❌ Specify number <= 768')
return
await _set_property(message, 'width')
await _set_property(message, 'width', is_command=is_command)
@throttle(cooldown=5, admin_ids=db[DBTables.config].get('admins'))
async def set_height_command(message: types.Message):
async def set_height_command(message: types.Message, is_command: bool = True):
try:
_ = int(message.get_args())
_ = int((message.get_args() if is_command else message.text))
except Exception as e:
assert e
await message.reply('❌ Specify number as argument')
@@ -99,40 +99,40 @@ async def set_height_command(message: types.Message):
await message.reply('❌ Specify number <= 768')
return
await _set_property(message, 'height')
await _set_property(message, 'height', is_command=is_command)
@throttle(cooldown=5, admin_ids=db[DBTables.config].get('admins'))
async def set_restore_faces_command(message: types.Message):
async def set_restore_faces_command(message: types.Message, is_command: bool = True):
try:
_ = bool(message.get_args())
_ = bool((message.get_args() if is_command else message.text))
except Exception as e:
assert e
await message.reply('❌ Specify boolean <code>True</code>/<code>False</code> as argument',
parse_mode='HTML')
return
await _set_property(message, 'restore_faces')
await _set_property(message, 'restore_faces', is_command=is_command)
@wrap_exception()
@throttle(cooldown=5, admin_ids=db[DBTables.config].get('admins'))
async def set_sampler_command(message: types.Message):
async def set_sampler_command(message: types.Message, is_command: bool = True):
from bot.modules.api.samplers import get_samplers
if message.get_args() not in (samplers := await get_samplers()):
if (message.get_args() if is_command else message.text) not in (samplers := await get_samplers()):
await message.reply(
f'❌ You can use only {", ".join(f"<code>{x}</code>" for x in samplers)}',
parse_mode='HTML'
)
return
await _set_property(message, 'sampler')
await _set_property(message, 'sampler', is_command=is_command)
@throttle(cooldown=5, admin_ids=db[DBTables.config].get('admins'))
async def set_size_command(message: types.Message):
async def set_size_command(message: types.Message, is_command: bool = True):
try:
hxw = message.get_args().split('x')
hxw = (message.get_args() if is_command else message.text).split('x')
height = int(hxw[0])
width = int(hxw[1])
except Exception as e:
@@ -145,5 +145,5 @@ async def set_size_command(message: types.Message):
await message.reply('❌ Specify numbers <= 768')
return
await _set_property(message, 'height', height)
await _set_property(message, 'width', width)
await _set_property(message, 'height', height, is_command=is_command)
await _set_property(message, 'width', width, is_command=is_command)