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

Setting generation properties

Browse Source
This commit is contained in:
BarsTiger
2023-02-26 22:16:05 +02:00
parent 8f0cd0ac05
commit 3c3d644137
10 changed files with 286 additions and 56 deletions
Download Patch File Download Diff File Expand all files Collapse all files

17
bot/handlers/txt2img/__init__.py
View File

@@ -1,6 +1,19 @@
from bot.common import dp
from .txt2img import txt2img_comand
from .txt2img import generate_command
from .set_settings import (
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
)
def register():
dp.register_message_handler(txt2img.txt2img_comand, commands='txt2img')
dp.register_message_handler(txt2img.generate_command, commands='generate')
dp.register_message_handler(set_settings.set_prompt_command, commands='setprompt')
dp.register_message_handler(set_settings.set_height_command, commands='setheight')
dp.register_message_handler(set_settings.set_width_command, commands='setwidth')
dp.register_message_handler(set_settings.set_negative_prompt_command, commands='setnegative')
dp.register_message_handler(set_settings.set_size_command, commands='setsize')
dp.register_message_handler(set_settings.set_steps_command, commands='setsteps')
dp.register_message_handler(set_settings.set_sampler_command, commands='setsampler')
dp.register_message_handler(set_settings.set_cfg_scale_command, commands='setscale')
dp.register_message_handler(set_settings.set_restore_faces_command, commands='setfaces')

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

@@ -6,22 +6,28 @@ from bot.keyboards.exception import get_exception_keyboard
from bot.utils.trace_exception import PrettyException
@throttle(cooldown=5, admin_ids=db[DBTables.config].get('admins'))
async def set_prompt_command(message: types.Message):
temp_message = await message.reply("⏳ Setting prompt...")
async def _set_property(message: types.Message, prop: str, value=None):
temp_message = await message.reply(f"⏳ Setting {prop}...")
if not message.get_args():
await temp_message.edit_text("😶‍🌫️ Specify prompt for this command. Check /help setprompt")
await temp_message.edit_text("😶‍🌫️ Specify arguments for this command. Check /help")
return
try:
prompt: Prompt = db[DBTables.prompts].get(message.from_id, Prompt(message.get_args()))
prompt.prompt = message.get_args()
prompt: Prompt = db[DBTables.prompts].get(message.from_id)
if prompt is None and prop != 'prompt':
await temp_message.edit_text(f"You didn't created any prompt. Specify prompt text at least first time. "
f"For example, it can be: <code>masterpiece, best quality, 1girl, white hair, "
f"medium hair, cat ears, closed eyes, looking at viewer, :3, cute, scarf, "
f"jacket, outdoors, streets</code>", parse_mode='HTML')
return
prompt.__setattr__(prop, message.get_args() if value is None else value)
prompt.creator = message.from_id
db[DBTables.prompts][message.from_id] = prompt
await db[DBTables.config].write()
await message.reply('✅ Default prompt set')
await message.reply(f'✅ {prop} set')
await temp_message.delete()
except Exception as e:
@@ -32,3 +38,134 @@ async def set_prompt_command(message: types.Message):
await temp_message.delete()
db[DBTables.queue]['n'] = db[DBTables.queue].get('n', 1) - 1
return
@throttle(cooldown=5, admin_ids=db[DBTables.config].get('admins'))
async def set_prompt_command(message: types.Message):
await _set_property(message, 'prompt')
@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')
@throttle(cooldown=5, admin_ids=db[DBTables.config].get('admins'))
async def set_steps_command(message: types.Message):
try:
_ = int(message.get_args())
except Exception as e:
assert e
await message.reply('❌ Specify number as argument')
return
if _ > 30:
await message.reply('❌ Specify number <= 30')
return
await _set_property(message, 'steps')
@throttle(cooldown=5, admin_ids=db[DBTables.config].get('admins'))
async def set_cfg_scale_command(message: types.Message):
try:
_ = int(message.get_args())
except Exception as e:
assert e
await message.reply('❌ Specify number as argument')
return
await _set_property(message, 'cfg_scale')
@throttle(cooldown=5, admin_ids=db[DBTables.config].get('admins'))
async def set_width_command(message: types.Message):
try:
_ = int(message.get_args())
except Exception as e:
assert e
await message.reply('❌ Specify number as argument')
return
if _ > 768:
await message.reply('❌ Specify number <= 768')
return
await _set_property(message, 'width')
@throttle(cooldown=5, admin_ids=db[DBTables.config].get('admins'))
async def set_height_command(message: types.Message):
try:
_ = int(message.get_args())
except Exception as e:
assert e
await message.reply('❌ Specify number as argument')
return
if _ > 768:
await message.reply('❌ Specify number <= 768')
return
await _set_property(message, 'height')
@throttle(cooldown=5, admin_ids=db[DBTables.config].get('admins'))
async def set_restore_faces_command(message: types.Message):
try:
_ = bool(message.get_args())
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')
@throttle(cooldown=5, admin_ids=db[DBTables.config].get('admins'))
async def set_sampler_command(message: types.Message):
temp_message = await message.reply('⏳ Getting samplers...')
from bot.modules.api.samplers import get_samplers
from aiohttp import ClientConnectorError
try:
if message.get_args() 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 temp_message.delete()
except ClientConnectorError:
await message.reply('❌ Error! Maybe, StableDiffusion API endpoint is incorrect '
'or turned off')
except Exception as e:
exception_id = f'{message.message_thread_id}-{message.message_id}'
db[DBTables.exceptions][exception_id] = PrettyException(e)
await message.reply('❌ Error happened while processing your request', parse_mode='html',
reply_markup=get_exception_keyboard(exception_id))
await temp_message.delete()
db[DBTables.queue]['n'] = db[DBTables.queue].get('n', 1) - 1
return
await _set_property(message, 'sampler')
@throttle(cooldown=5, admin_ids=db[DBTables.config].get('admins'))
async def set_size_command(message: types.Message):
try:
hxw = message.get_args().split('x')
height = int(hxw[0])
width = int(hxw[1])
except Exception as e:
assert e
await message.reply('❌ Specify size in <code>hxw</code> format, for example <code>512x512</code>',
parse_mode='HTML')
return
if height > 768 or width > 768:
await message.reply('❌ Specify numbers <= 768')
return
await _set_property(message, 'height', height)
await _set_property(message, 'width', width)

38
bot/handlers/txt2img/txt2img.py
View File

@@ -1,24 +1,30 @@
import re
from aiogram import types
from bot.db import db, DBTables
from bot.utils.cooldown import throttle
from bot.modules.api.txt2img import txt2img
from bot.modules.api.objects.prompt_request import Prompt
from bot.modules.api.objects.get_prompt import get_prompt
from bot.modules.api.objects.prompt_request import Generated
from bot.modules.api.status import wait_for_status
from bot.keyboards.exception import get_exception_keyboard
from bot.keyboards.image_info import get_img_info_keyboard
from bot.utils.trace_exception import PrettyException
from aiohttp import ClientConnectorError
@throttle(cooldown=30, admin_ids=db[DBTables.config].get('admins'))
async def txt2img_comand(message: types.Message):
async def generate_command(message: types.Message):
temp_message = await message.reply("⏳ Enqueued...")
prompt: Prompt = db[DBTables.prompts].get(message.from_id)
if not prompt:
if message.get_args():
db[DBTables.prompts][message.from_id] = Prompt(message.get_args(), creator=message.from_id)
# TODO: Move it to other module
try:
prompt = get_prompt(user_id=message.from_id,
prompt_string=message.get_args())
except AttributeError:
await temp_message.edit_text(f"You didn't created any prompt. Specify prompt text at least first time. "
f"For example, it can be: <code>masterpiece, best quality, 1girl, white hair, "
f"medium hair, cat ears, closed eyes, looking at viewer, :3, cute, scarf, jacket, "
f"outdoors, streets</code>", parse_mode='HTML')
return
try:
db[DBTables.queue]['n'] = db[DBTables.queue].get('n', 0) + 1
@@ -27,12 +33,18 @@ async def txt2img_comand(message: types.Message):
await wait_for_status()
await temp_message.edit_text(f"⌛ Generating...")
prompt = Prompt(prompt=message.get_args(), creator=message.from_id)
image = await txt2img(prompt)
image_message = await message.reply_photo(photo=image[0])
db[DBTables.queue]['n'] = db[DBTables.queue].get('n', 1) - 1
db[DBTables.generated][image_message.photo[0].file_unique_id] = prompt
db[DBTables.generated][image_message.photo[0].file_unique_id] = Generated(
prompt=prompt,
seed=image[1]['seed'],
model=re.search(r", Model: ([^,]+),", image[1]['infotexts'][0]).groups()[0]
)
await message.reply('Here is your image',
reply_markup=get_img_info_keyboard(image_message.photo[0].file_unique_id))
await temp_message.delete()
@@ -45,6 +57,12 @@ async def txt2img_comand(message: types.Message):
db[DBTables.queue]['n'] = db[DBTables.queue].get('n', 1) - 1
return
except ValueError as e:
await message.reply(f'❌ Error! {e.args[0]}')
await temp_message.delete()
db[DBTables.queue]['n'] = db[DBTables.queue].get('n', 1) - 1
return
except Exception as e:
exception_id = f'{message.message_thread_id}-{message.message_id}'
db[DBTables.exceptions][exception_id] = PrettyException(e)