dinlo
252 lines
11 KiB
Python
252 lines
11 KiB
Python
"""Обработчики настроек хранения и вспомогательные команды."""
|
||
|
||
import logging
|
||
from aiogram import Router, F
|
||
from aiogram.types import CallbackQuery, Message, InlineKeyboardMarkup, InlineKeyboardButton
|
||
from aiogram.fsm.context import FSMContext
|
||
from aiogram.fsm.state import State, StatesGroup
|
||
from aiogram.filters import Command
|
||
|
||
from config import settings
|
||
from database.database import db
|
||
from sd.sd_client import sd_client
|
||
from utils.image_manager import cleanup_expired_images
|
||
|
||
router = Router()
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
class StorageSettingsState(StatesGroup):
|
||
"""Состояния для настройки хранения."""
|
||
waiting_for_ttl = State()
|
||
|
||
|
||
# --- Обработчик главного меню ---
|
||
@router.callback_query(F.data == "main_menu")
|
||
async def main_menu(callback: CallbackQuery, state: FSMContext = None):
|
||
"""Возврат в главное меню."""
|
||
# Очищаем состояние, если оно есть
|
||
if state:
|
||
await state.clear()
|
||
|
||
from bot.handlers_start import get_main_keyboard
|
||
text = (
|
||
f"👋 <b>Привет, {callback.from_user.first_name}!</b>\n\n"
|
||
"Я бот для генерации изображений через Stable Diffusion.\n"
|
||
"Выберите действие из меню ниже:"
|
||
)
|
||
try:
|
||
await callback.message.edit_text(
|
||
text,
|
||
parse_mode="HTML",
|
||
reply_markup=get_main_keyboard(callback.from_user.id),
|
||
)
|
||
except Exception:
|
||
# Если edit_text не сработал (напр., сообщение с фото), отправляем новое
|
||
await callback.message.answer(
|
||
text,
|
||
parse_mode="HTML",
|
||
reply_markup=get_main_keyboard(callback.from_user.id),
|
||
)
|
||
await callback.answer()
|
||
|
||
|
||
# --- Проверка статуса SD API ---
|
||
@router.callback_query(F.data == "check_sd_status")
|
||
async def check_sd_status(callback: CallbackQuery):
|
||
"""Проверка соединения с SD API."""
|
||
status_msg = await callback.message.edit_text("🔍 Проверяю соединение с SD API...")
|
||
|
||
is_connected = await sd_client.check_connection()
|
||
|
||
if is_connected:
|
||
current_model = await sd_client.get_current_model()
|
||
text = (
|
||
f"✅ <b>SD API подключено!</b>\n\n"
|
||
f"Адрес: <code>{settings.SD_API_URL}</code>\n"
|
||
f"Текущая модель: <code>{current_model}</code>"
|
||
)
|
||
else:
|
||
text = (
|
||
f"❌ <b>SD API недоступно!</b>\n\n"
|
||
f"Адрес: <code>{settings.SD_API_URL}</code>\n\n"
|
||
"Проверьте:\n"
|
||
"1. Запущен ли Stable Diffusion WebUI\n"
|
||
"2. Доступен ли сервер по указанному адресу\n"
|
||
"3. Правильность настройки API (--api флаг)"
|
||
)
|
||
|
||
keyboard = InlineKeyboardMarkup(inline_keyboard=[
|
||
[InlineKeyboardButton(text="🔄 Повторить проверку", callback_data="check_sd_status")],
|
||
[InlineKeyboardButton(text="📋 Главное меню", callback_data="main_menu")],
|
||
])
|
||
|
||
await status_msg.edit_text(text, parse_mode="HTML", reply_markup=keyboard)
|
||
await callback.answer()
|
||
|
||
|
||
# --- Помощь ---
|
||
@router.callback_query(F.data == "help")
|
||
async def help_command(callback: CallbackQuery):
|
||
"""Справка по боту."""
|
||
text = (
|
||
"ℹ️ <b>Справка по боту</b>\n\n"
|
||
"<b>Генерация изображений:</b>\n"
|
||
"• <b>txt2img</b> — создание изображения по текстовому описанию\n\n"
|
||
"<b>Профили:</b>\n"
|
||
"Создавайте профили с предустановленными настройками:\n"
|
||
"• Размер изображения\n"
|
||
"• Количество шагов\n"
|
||
"• CFG Scale\n"
|
||
"• Сэмплер\n"
|
||
"• Модель и LoRA\n\n"
|
||
"Вы можете установить профиль по умолчанию для быстрой генерации.\n\n"
|
||
"<b>Настройки хранения:</b>\n"
|
||
"Настройте время хранения сгенерированных изображений.\n"
|
||
"Изображения автоматически удаляются по истечении срока.\n\n"
|
||
"<b>Команды:</b>\n"
|
||
"/start — Главное меню\n"
|
||
"/menu — Показать главное меню\n"
|
||
"/cancel — Отменить текущую операцию\n"
|
||
"/status — Статус SD API"
|
||
)
|
||
|
||
keyboard = InlineKeyboardMarkup(inline_keyboard=[
|
||
[InlineKeyboardButton(text="📋 Главное меню", callback_data="main_menu")],
|
||
])
|
||
|
||
await callback.message.edit_text(text, parse_mode="HTML", reply_markup=keyboard)
|
||
await callback.answer()
|
||
|
||
|
||
@router.message(Command("status"))
|
||
async def cmd_status(message: Message):
|
||
"""Команда /status."""
|
||
is_connected = await sd_client.check_connection()
|
||
if is_connected:
|
||
current_model = await sd_client.get_current_model()
|
||
text = (
|
||
f"✅ <b>SD API подключено!</b>\n\n"
|
||
f"Адрес: <code>{settings.SD_API_URL}</code>\n"
|
||
f"Текущая модель: <code>{current_model}</code>"
|
||
)
|
||
else:
|
||
text = (
|
||
f"❌ <b>SD API недоступно!</b>\n\n"
|
||
f"Адрес: <code>{settings.SD_API_URL}</code>"
|
||
)
|
||
await message.answer(text, parse_mode="HTML")
|
||
|
||
|
||
# --- Настройки хранения ---
|
||
@router.callback_query(F.data == "storage_settings")
|
||
async def storage_settings(callback: CallbackQuery):
|
||
"""Меню настроек хранения."""
|
||
user_settings = await db.get_user_settings(callback.from_user.id)
|
||
current_ttl = user_settings.get("image_ttl_hours", settings.DEFAULT_IMAGE_TTL_HOURS)
|
||
|
||
text = (
|
||
"🗃️ <b>Настройки хранения изображений</b>\n\n"
|
||
f"Текущее время хранения: <b>{current_ttl} часов</b>\n"
|
||
f"Максимальное время: <b>{settings.MAX_IMAGE_TTL_HOURS} часов</b>\n\n"
|
||
"Изображения автоматически удаляются по истечении срока.\n"
|
||
"Очистка происходит каждые {interval} минут.".format(interval=settings.CLEANUP_INTERVAL_MINUTES)
|
||
)
|
||
|
||
keyboard = InlineKeyboardMarkup(inline_keyboard=[
|
||
[InlineKeyboardButton(text="12 часов", callback_data="set_ttl_12")],
|
||
[InlineKeyboardButton(text="24 часа (1 день)", callback_data="set_ttl_24")],
|
||
[InlineKeyboardButton(text="48 часов (2 дня)", callback_data="set_ttl_48")],
|
||
[InlineKeyboardButton(text="72 часа (3 дня)", callback_data="set_ttl_72")],
|
||
[InlineKeyboardButton(text="168 часов (7 дней)", callback_data="set_ttl_168")],
|
||
[InlineKeyboardButton(text="✏️ Ввести своё значение", callback_data="set_ttl_custom")],
|
||
[InlineKeyboardButton(text="⬅️ Назад", callback_data="main_menu")],
|
||
])
|
||
|
||
await callback.message.edit_text(text, parse_mode="HTML", reply_markup=keyboard)
|
||
await callback.answer()
|
||
|
||
|
||
@router.callback_query(F.data.startswith("set_ttl_"))
|
||
async def set_ttl(callback: CallbackQuery, state: FSMContext):
|
||
"""Установка времени хранения."""
|
||
if callback.data == "set_ttl_custom":
|
||
await state.set_state(StorageSettingsState.waiting_for_ttl)
|
||
await callback.message.edit_text(
|
||
f"⏰ <b>Введите время хранения в часах</b>\n\n"
|
||
f"Допустимый диапазон: 1 - {settings.MAX_IMAGE_TTL_HOURS} часов.",
|
||
parse_mode="HTML",
|
||
)
|
||
await callback.answer()
|
||
return
|
||
|
||
ttl_hours = int(callback.data.replace("set_ttl_", ""))
|
||
await _save_ttl(callback, ttl_hours)
|
||
|
||
|
||
@router.message(StorageSettingsState.waiting_for_ttl)
|
||
async def save_custom_ttl(message: Message, state: FSMContext):
|
||
"""Сохранение пользовательского значения TTL."""
|
||
try:
|
||
ttl_hours = int(message.text.strip())
|
||
if ttl_hours < 1 or ttl_hours > settings.MAX_IMAGE_TTL_HOURS:
|
||
await message.answer(
|
||
f"Значение должно быть от 1 до {settings.MAX_IMAGE_TTL_HOURS} часов."
|
||
)
|
||
return
|
||
except ValueError:
|
||
await message.answer("Введите корректное целое число.")
|
||
return
|
||
|
||
await state.clear()
|
||
await _save_ttl(message, ttl_hours)
|
||
|
||
|
||
async def _save_ttl(target, ttl_hours: int):
|
||
"""Сохранить TTL и показать подтверждение."""
|
||
user_id = target.from_user.id
|
||
await db.update_user_settings(user_id, image_ttl_hours=ttl_hours)
|
||
|
||
text = (
|
||
f"✅ <b>Настройки сохранены!</b>\n\n"
|
||
f"Время хранения изображений: <b>{ttl_hours} часов</b>\n"
|
||
f"Изображения будут автоматически удаляться по истечении этого срока."
|
||
)
|
||
|
||
keyboard = InlineKeyboardMarkup(inline_keyboard=[
|
||
[InlineKeyboardButton(text="🗃️ Другие настройки хранения", callback_data="storage_settings")],
|
||
[InlineKeyboardButton(text="📋 Главное меню", callback_data="main_menu")],
|
||
])
|
||
|
||
if hasattr(target, 'message'):
|
||
await target.message.edit_text(text, parse_mode="HTML", reply_markup=keyboard)
|
||
else:
|
||
await target.answer(text, parse_mode="HTML", reply_markup=keyboard)
|
||
|
||
|
||
# --- Команда /cancel ---
|
||
@router.message(Command("cancel"))
|
||
async def cmd_cancel(message: Message, state: FSMContext):
|
||
"""Отмена текущей операции."""
|
||
current_state = await state.get_state()
|
||
if current_state is None:
|
||
await message.answer("Нет активных операций.")
|
||
return
|
||
|
||
await state.clear()
|
||
await message.answer("❌ Операция отменена.")
|
||
|
||
from bot.handlers_start import get_main_keyboard
|
||
await message.answer("📋 Главное меню:", reply_markup=get_main_keyboard(message.from_user.id))
|
||
|
||
|
||
# --- Периодическая очистка ---
|
||
async def scheduled_cleanup():
|
||
"""Функция для периодической очистки просроченных изображений."""
|
||
try:
|
||
deleted = await cleanup_expired_images()
|
||
if deleted > 0:
|
||
logger.info(f"Очистка: удалено {deleted} изображений")
|
||
except Exception as e:
|
||
logger.error(f"Ошибка при очистке изображений: {e}")
|