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}")
|