From b88ccf3b4baf01107cb1c2ff9f7b3514c529f7bc Mon Sep 17 00:00:00 2001 From: dinlo Date: Sun, 31 May 2026 18:46:09 +0800 Subject: [PATCH] Initial commit Co-Authored-By: Claude Opus 4.8 (1M context) --- .dockerignore | 9 + .env.example | 20 + .gitignore | 7 + .qwen/settings.json | 11 + .qwen/settings.json.orig | 7 + Dockerfile | 35 ++ README.md | 986 +++++++++++++++++++++++++++++++++++++ bot/__init__.py | 1 + bot/handlers_admin.py | 513 +++++++++++++++++++ bot/handlers_generation.py | 241 +++++++++ bot/handlers_profiles.py | 634 ++++++++++++++++++++++++ bot/handlers_settings.py | 251 ++++++++++ bot/handlers_start.py | 51 ++ bot/middleware.py | 102 ++++ config.py | 30 ++ database/__init__.py | 1 + database/database.py | 497 +++++++++++++++++++ docker-compose.yml | 30 ++ main.py | 100 ++++ requirements.txt | 5 + sd/__init__.py | 1 + sd/sd_client.py | 341 +++++++++++++ utils/__init__.py | 1 + utils/image_manager.py | 60 +++ 24 files changed, 3934 insertions(+) create mode 100644 .dockerignore create mode 100644 .env.example create mode 100644 .gitignore create mode 100644 .qwen/settings.json create mode 100644 .qwen/settings.json.orig create mode 100644 Dockerfile create mode 100644 README.md create mode 100644 bot/__init__.py create mode 100644 bot/handlers_admin.py create mode 100644 bot/handlers_generation.py create mode 100644 bot/handlers_profiles.py create mode 100644 bot/handlers_settings.py create mode 100644 bot/handlers_start.py create mode 100644 bot/middleware.py create mode 100644 config.py create mode 100644 database/__init__.py create mode 100644 database/database.py create mode 100644 docker-compose.yml create mode 100644 main.py create mode 100644 requirements.txt create mode 100644 sd/__init__.py create mode 100644 sd/sd_client.py create mode 100644 utils/__init__.py create mode 100644 utils/image_manager.py diff --git a/.dockerignore b/.dockerignore new file mode 100644 index 0000000..24010f8 --- /dev/null +++ b/.dockerignore @@ -0,0 +1,9 @@ +.env +__pycache__/ +*.pyc +*.pyo +*.db +.docker/ +*.md +.gitignore +.env.example diff --git a/.env.example b/.env.example new file mode 100644 index 0000000..e2dd087 --- /dev/null +++ b/.env.example @@ -0,0 +1,20 @@ +# Telegram Bot Token (получите у @BotFather) +BOT_TOKEN=your_telegram_bot_token_here + +# Telegram ID администратора (узнайте через @userinfobot) +ADMIN_ID=123456789 + +# Stable Diffusion API адрес +SD_API_URL=http://192.168.1.120:7860 + +# Путь к папке с изображениями внутри контейнера +IMAGES_DIR=/app/images + +# Время хранения изображений по умолчанию (часы) +DEFAULT_IMAGE_TTL_HOURS=48 + +# Максимальное время хранения изображений (часы) +MAX_IMAGE_TTL_HOURS=168 + +# Период очистки изображений (минуты) +CLEANUP_INTERVAL_MINUTES=30 diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..fa76c29 --- /dev/null +++ b/.gitignore @@ -0,0 +1,7 @@ +.env +images/ +__pycache__/ +*.pyc +*.pyo +*.db +.docker/ diff --git a/.qwen/settings.json b/.qwen/settings.json new file mode 100644 index 0000000..6d06fed --- /dev/null +++ b/.qwen/settings.json @@ -0,0 +1,11 @@ +{ + "permissions": { + "allow": [ + "Bash(python *)", + "Bash(pip install *)", + "Bash(docker-compose *)", + "Bash(docker compose up *)" + ] + }, + "$version": 3 +} \ No newline at end of file diff --git a/.qwen/settings.json.orig b/.qwen/settings.json.orig new file mode 100644 index 0000000..246ba6c --- /dev/null +++ b/.qwen/settings.json.orig @@ -0,0 +1,7 @@ +{ + "permissions": { + "allow": [ + "Bash(python *)" + ] + } +} \ No newline at end of file diff --git a/Dockerfile b/Dockerfile new file mode 100644 index 0000000..5f07080 --- /dev/null +++ b/Dockerfile @@ -0,0 +1,35 @@ +# Dockerfile для Telegram бота генерации изображений +FROM python:3.12-slim + +# Метки +LABEL maintainer="tg-sd-bot" +LABEL description="Telegram bot for Stable Diffusion image generation" + +# Рабочая директория +WORKDIR /app + +# Установка зависимостей ОС +RUN apt-get update && \ + apt-get install -y --no-install-recommends \ + curl \ + && rm -rf /var/lib/apt/lists/* + +# Копирование requirements и установка зависимостей Python +COPY requirements.txt . +RUN pip install --no-cache-dir -r requirements.txt + +# Копирование кода бота +COPY . . + +# Создание директорий для данных +RUN mkdir -p /app/images /app/data + +# Переменные окружения по умолчанию +ENV PYTHONUNBUFFERED=1 +ENV PYTHONDONTWRITEBYTECODE=1 + +# Том для хранения изображений и базы данных +VOLUME ["/app/images", "/app/data"] + +# Команда запуска +CMD ["python", "main.py"] diff --git a/README.md b/README.md new file mode 100644 index 0000000..d86a8df --- /dev/null +++ b/README.md @@ -0,0 +1,986 @@ +# Telegram Bot для генерации изображений через Stable Diffusion (Automatic1111) + +Бот для генерации изображений через API Stable Diffusion WebUI (Automatic1111) с поддержкой профилей, управления хранением изображений и удобным интерфейсом. + +## Возможности + +- 🎨 **Генерация изображений (txt2img)** — создание изображений по текстовому описанию +- ⚙️ **Система профилей** — сохранение настроек (размер, модель, LoRA, сэмплер и т.д.) +- 🗃️ **Управление хранением** — настройка времени хранения изображений (автоочистка) +- 📊 **Мониторинг** — проверка статуса подключения к SD API +- 🐳 **Docker** — простая установка и запуск в контейнере +- 🛡️ **Админ-панель** — управление пользователями, лимитами и доступом + +--- + +## Содержание + +1. [Что умеет бот](#что-умеет-бот) +2. [Требования](#требования) +3. [Настройка Stable Diffusion WebUI](#настройка-stable-diffusion-webui) +4. [Создание Telegram бота](#создание-telegram-бота) +5. [Установка бота](#установка-бота) +6. [Настройка конфигурации](#настройка-конфигурации) +7. [Запуск бота](#запуск-бота) +8. [Использование бота](#использование-бота) +9. [Примеры промптов](#примеры-промптов) +10. [Структура проекта](#структура-проекта) +11. [FAQ](#faq) +12. [Устранение неполадок](#устранение-неполадок) + +--- + +## Что умеет бот + +### 🎨 Генерация изображений + +Бот позволяет создавать изображения двумя способами: + +**txt2img (текст → изображение)** — опишите словами, что хотите увидеть, и бот сгенерирует изображение. Вы вводите промпт (описание) и необязательный негативный промпт (чего НЕ должно быть на картинке), а бот передаёт запрос в Stable Diffusion и возвращает готовое изображение. + +После генерации бот покажет все параметры: промпт, размер, количество шагов, CFG Scale, сэмплер, шедулер, seed, модель и время хранения изображения. + +### ⚙️ Система профилей + +Создавайте и сохраняйте профили с предустановленными настройками, чтобы не вводить их каждый раз заново: + +| Параметр | Описание | +|----------|----------| +| **Размер изображения** | Предустановленные (512×512, 768×768, 512×768, 768×512, 1024×1024) или произвольный (64–2048 px) | +| **Количество шагов** | От 10 до 150. Больше шагов = выше качество, но дольше генерация | +| **CFG Scale** | От 1.0 до 30.0. Определяет, насколько строго модель следует промпту (стандарт: 7.0) | +| **Сэмплер** | Euler a, Euler, DPM++ 2M Karras, DPM++ SDE Karras, DDIM и другие — загружаются из SD API | +| **Шедулер** | Автоматически подбирается к сэмплеру | +| **Модель** | Выбор из всех доступных моделей в вашем SD WebUI | +| **LoRA** | Подключение LoRA с настраиваемой силой (0.0–1.0) | +| **Негативный промпт** | Описание того, чего НЕ должно быть на изображении | + +Вы можете установить один профиль как **профиль по умолчанию** (⭐) — тогда при генерации будут автоматически применяться его настройки. + +### 🛡️ Админ-панель + +Администратор (указанный в `ADMIN_ID`) имеет доступ к панели управления: + +- **Добавление пользователей** — выдача доступа по Telegram User ID +- **Типы доступа:** + - ♾️ **Без ограничений** — полный доступ навсегда + - ⏰ **По времени** — доступ на заданное количество дней (1–3650) + - 🎨 **По количеству генераций** — лимит на число сгенерированных изображений +- **Список пользователей** — просмотр, пагинация, управление +- **Управление пользователями** — блокировка, разблокировка, обновление доступа, удаление +- **Очистка истёкших доступов** — автоматическая деактивация пользователей с истёкшим сроком + +> **Как узнать свой User ID?** — Напишите боту [@userinfobot](https://t.me/userinfobot) в Telegram. + +### 🗃️ Управление хранением + +Каждый пользователь может настроить время хранения своих сгенерированных изображений: + +- **Предустановленные варианты:** 12 ч, 24 ч (1 день), 48 ч (2 дня), 72 ч (3 дня), 168 ч (7 дней) +- **Произвольное значение** — от 1 часа до максимума (`MAX_IMAGE_TTL_HOURS`) +- **Автоматическая очистка** — фоновая задача удаляет просроченные изображения с заданным интервалом + +### 📊 Мониторинг + +Команда `/status` или кнопка «Статус SD API» покажет: + +- Статус подключения к SD API +- Адрес API +- Текущую загруженную модель + +### 🔐 Система доступа + +Бот по умолчанию **не допускает** новых пользователей — только те, кого добавил администратор, могут пользоваться ботом. Это защищает от несанкционированного использования ваших ресурсов. + +При попытке использования бота без доступа пользователь получит сообщение об ошибке. + +### 📋 Команды бота + +| Команда | Описание | +|---------|----------| +| `/start` | Запуск бота, приветствие и главное меню | +| `/menu` | Показать главное меню | +| `/status` | Проверка статуса SD API | +| `/cancel` | Отмена текущей операции | +| `/admin` | Админ-панель (только для администратора) | + +--- + +## Требования + +### Для сервера с ботом (Docker) + +- **Операционная система:** Linux (Ubuntu 20.04+), Windows 10/11 с WSL2, macOS 12+ +- **Docker:** версии 20.10+ +- **Docker Compose:** версии 2.0+ +- **Свободное место:** ~500 МБ для бота + место для изображений +- **Оперативная память:** минимум 256 МБ для контейнера бота +- **Сеть:** доступ к серверу Stable Diffusion по локальной сети + +### Для сервера со Stable Diffusion + +- **Stable Diffusion WebUI (Automatic1111)** — запущенный и доступный по сети +- **API режим** — включён флагом `--api` +- **Сетевой доступ** — слушает не только localhost (флаг `--listen`) + +--- + +## Настройка Stable Diffusion WebUI + +### 1. Запуск с поддержкой API + +На сервере, где установлен Stable Diffusion WebUI, запустите его со следующими флагами: + +```bash +# Для Linux +./webui.sh --api --listen + +# Для Windows +webui.bat --api --listen +``` + +**Обязательные флаги:** + +| Флаг | Описание | +|------|----------| +| `--api` | Включает программный API для бота | +| `--listen` | Позволяет подключаться не только с localhost | + +**Опциональные флаги:** + +| Флаг | Описание | +|------|----------| +| `--port 7860` | Указать порт (по умолчанию 7860) | +| `--nowatchdog` | Отключить watchdog (рекомендуется для серверов) | +| `--xformers` | Использовать xformers для ускорения (если поддерживается) | + +**Пример полной команды:** + +```bash +./webui.sh --api --listen --port 7860 --nowatchdog +``` + +### 2. Проверка доступности API + +После запуска проверьте, что API доступен: + +```bash +# С сервера с ботом (замените IP на ваш) +curl http://192.168.1.120:7860/sdapi/v1/options +``` + +Вы должны получить JSON-ответ с настройками API. + +### 3. Настройка firewall (если необходимо) + +Убедитесь, что порт 7860 открыт для подключения с сервера бота: + +```bash +# Для Ubuntu/Debian (UFW) +sudo ufw allow 7860/tcp + +# Для CentOS/RHEL (firewalld) +sudo firewall-cmd --permanent --add-port=7860/tcp +sudo firewall-cmd --reload +``` + +--- + +## Создание Telegram бота + +### 1. Получение токена бота + +1. Откройте Telegram и найдите бота **@BotFather** +2. Отправьте команду `/newbot` +3. Следуйте инструкциям: + - Введите **отображаемое имя** бота (например, `SD Generator`) + - Введите **username** бота (должен заканчиваться на `bot`, например `sd_gen_bot`) +4. BotFather отправит вам **токен** вида `123456789:ABCdefGHIjklMNOpqrsTUVwxyz` + +### 2. Настройка описания бота (опционально) + +В BotFather: + +``` +/setdescription — описание, которое видит пользователь до начала общения +/setabouttext — короткое описание для поиска ботов +/setuserpic — аватар бота +``` + +--- + +## Установка бота + +### Способ 1: С помощью Docker Compose (рекомендуется) + +#### 1. Клонирование или копирование файлов + +Скопируйте все файлы проекта на сервер, где будет работать бот: + +``` +tg-sd/ +├── main.py +├── config.py +├── requirements.txt +├── Dockerfile +├── docker-compose.yml +├── .env +├── bot/ +│ ├── __init__.py +│ ├── handlers_start.py +│ ├── handlers_generation.py +│ ├── handlers_profiles.py +│ └── handlers_settings.py +├── sd/ +│ ├── __init__.py +│ └── sd_client.py +├── database/ +│ ├── __init__.py +│ └── database.py +└── utils/ + ├── __init__.py + └── image_manager.py +``` + +#### 2. Настройка переменных окружения + +Откройте файл `.env` и укажите ваши значения: + +```bash +nano .env +``` + +**Обязательные параметры:** + +```env +# Токен вашего бота (получите у @BotFather) +BOT_TOKEN=123456789:ABCdefGHIjklMNOpqrsTUVwxyz + +# Адрес Stable Diffusion API +SD_API_URL=http://192.168.1.120:7860 +``` + +**Опциональные параметры:** + +```env +# Время хранения изображений по умолчанию (часы) +DEFAULT_IMAGE_TTL_HOURS=48 + +# Максимальное время хранения (часы) +MAX_IMAGE_TTL_HOURS=168 + +# Период автоматической очистки (минуты) +CLEANUP_INTERVAL_MINUTES=30 +``` + +#### 3. Создание необходимых директорий + +```bash +mkdir -p images data +``` + +#### 4. Запуск бота + +```bash +# Сборка и запуск +docker compose up -d --build + +# Проверка логов +docker compose logs -f tg-sd-bot +``` + +#### 5. Остановка бота + +```bash +# Остановка +docker compose down + +# Остановка с удалением контейнера и volumes +docker compose down -v +``` + +### Способ 2: Без Docker (прямой запуск) + +#### 1. Установка Python + +Убедитесь, что установлен Python 3.10+: + +```bash +python3 --version +``` + +#### 2. Создание виртуального окружения + +```bash +python3 -m venv venv +source venv/bin/activate # Linux/macOS +# или +venv\Scripts\activate # Windows +``` + +#### 3. Установка зависимостей + +```bash +pip install -r requirements.txt +``` + +#### 4. Настройка окружения + +```bash +cp .env.example .env +nano .env # Укажите ваши значения +``` + +#### 5. Запуск бота + +```bash +python main.py +``` + +--- + +## Настройка конфигурации + +### Файл .env + +| Параметр | Описание | По умолчанию | Обязательно | +|----------|----------|--------------|-------------| +| `BOT_TOKEN` | Токен Telegram бота | — | ✅ Да | +| `SD_API_URL` | URL Stable Diffusion API | `http://192.168.1.120:7860` | ✅ Да | +| `IMAGES_DIR` | Путь к папке изображений | `/app/images` | Нет | +| `DEFAULT_IMAGE_TTL_HOURS` | Время хранения по умолчанию | `48` | Нет | +| `MAX_IMAGE_TTL_HOURS` | Максимальное время хранения | `168` | Нет | +| `CLEANUP_INTERVAL_MINUTES` | Интервал очистки | `30` | Нет | +| `DB_PATH` | Путь к базе данных | `/app/data/bot.db` | Нет | + +--- + +## Запуск бота + +### Docker Compose + +```bash +# Запуск +docker compose up -d + +# Просмотр логов +docker compose logs -f + +# Перезапуск +docker compose restart + +# Остановка +docker compose down + +# Обновление (после изменений в коде) +docker compose up -d --build +``` + +### Прямой запуск + +```bash +# Активация окружения +source venv/bin/activate + +# Запуск +python main.py + +# Остановка: Ctrl+C +``` + +### Запуск как сервис (systemd, Linux) + +Создайте файл сервиса: + +```bash +sudo nano /etc/systemd/system/tg-sd-bot.service +``` + +Содержимое: + +```ini +[Unit] +Description=Telegram SD Bot +After=network.target + +[Service] +Type=simple +User=your_user +WorkingDirectory=/path/to/tg-sd +ExecStart=/path/to/tg-sd/venv/bin/python main.py +Restart=always +RestartSec=10 + +[Install] +WantedBy=multi-user.target +``` + +Активация: + +```bash +sudo systemctl daemon-reload +sudo systemctl enable tg-sd-bot +sudo systemctl start tg-sd-bot + +# Проверка статуса +sudo systemctl status tg-sd-bot + +# Просмотр логов +sudo journalctl -u tg-sd-bot -f +``` + +--- + +## Использование бота + +### Команды + +| Команда | Описание | +|---------|----------| +| `/start` | Запуск бота, главное меню | +| `/menu` | Показать главное меню | +| `/status` | Проверка статуса SD API | +| `/cancel` | Отмена текущей операции | + +### Генерация изображения (txt2img) + +1. Нажмите **"🎨 Генерация (txt2img)"** +2. Введите описание изображения (промпт) +3. Введите негативный промпт (или `-` для пропуска) +4. Дождитесь результата + +**Пример промпта:** + +``` +a beautiful sunset over mountains, golden hour, dramatic clouds, 4k, highly detailed +``` + +**Пример негативного промпта:** + +``` +blurry, low quality, watermark, text, deformed +``` + +### Генерация на основе изображения (img2img) + +1. Нажмите **"🖼️ Генерация (img2img)"** +2. Отправьте изображение +3. Введите промпт +4. Выберите силу изменений + +### Управление профилями + +1. Нажмите **"⚙️ Профили"** +2. Создайте новый профиль +3. Настройте параметры: + - Размер изображения (512x512, 768x768, и т.д.) + - Количество шагов + - CFG Scale + - Сэмплер + - Модель + - LoRA и его сила + - Негативный промпт +4. Установите профиль по умолчанию (⭐) + +### Настройки хранения + +1. Нажмите **"🗑️ Настройки хранения"** +2. Выберите время хранения: + - 12 часов + - 24 часа (1 день) + - 48 часов (2 дня) + - 72 часа (3 дня) + - 168 часов (7 дней) + - Пользовательское значение + +--- + +## Примеры промптов + +### 📸 Фотореализм и портреты + +``` +professional portrait photograph of a young woman, natural lighting, soft focus, shallow depth of field, 85mm lens, film grain, highly detailed skin texture +``` + +``` +elderly fisherman on a wooden boat at dawn, misty lake, warm golden light, cinematic composition, photorealistic, 4k +``` + +**Негативный промпт:** +``` +blurry, low quality, watermark, text, deformed, extra fingers, bad anatomy, cartoon, illustration +``` + +### 🌄 Пейзажи и природа + +``` +majestic mountain range at sunset, snow-capped peaks, golden hour lighting, dramatic clouds reflected in a crystal clear alpine lake, wide angle, national geographic style +``` + +``` +enchanted forest with bioluminescent mushrooms, fireflies, moonlight filtering through ancient trees, magical atmosphere, fantasy landscape, highly detailed +``` + +**Негативный промпт:** +``` +urban, buildings, people, watermark, text, blurry, oversaturated +``` + +### 🏙️ Архитектура и города + +``` +futuristic cyberpunk city at night, neon lights, flying cars, towering skyscrapers with holographic advertisements, rain-soaked streets, cinematic lighting, blade runner style +``` + +``` +ancient Greek temple ruins at sunrise, marble columns, overgrown with vines, peaceful atmosphere, warm morning light, archaeological photography style +``` + +**Негативный промпт:** +``` +modern cars, people, watermark, text, low quality, distorted perspective +``` + +### 🐉 Фэнтези и фантастика + +``` +epic dragon perched on a volcanic cliff, iridescent scales, wings spread wide, stormy sky with lightning bolts, fantasy art, highly detailed, dramatic lighting, artstation trending +``` + +``` +astronaut floating in a nebula, surrounded by cosmic dust and stars, surreal space scene, vivid colors, cinematic composition, digital art +``` + +**Негативный промпт:** +``` +realistic, photograph, blurry, low resolution, watermark, text, bad anatomy +``` + +### 🎨 Художественные стили + +``` +oil painting of a Japanese garden in autumn, red maple trees, stone lantern, koi pond, impressionist style, visible brushstrokes, Monet inspired +``` + +``` +watercolor illustration of a cozy cafe on a rainy day, warm interior light visible through windows, soft pastel colors, lo-fi aesthetic +``` + +``` +pixel art of a medieval castle, 16-bit style, game sprite, clean pixels, side view +``` + +**Негативный промпт:** +``` +photorealistic, 3d render, blurry, watermark, text +``` + +### 🍱 Аниме стили + +``` +anime girl with silver hair and blue eyes, wearing a school uniform, cherry blossom petals falling, soft lighting, Makoto Shinkai style, highly detailed anime art +``` + +``` +anime scene of a samurai standing on a bridge at sunset, dramatic pose, katana drawn, wind blowing cloak, Studio Ghibli style, beautiful background art +``` + +**Негативный промпт:** +``` +realistic, photograph, 3d, bad anatomy, extra limbs, watermark, text, low quality +``` + +### 🍔 Предметы и еда + +``` +professional food photography of a gourmet burger, melted cheese dripping, fresh ingredients, dark background, studio lighting, commercial photography style, 4k +``` + +``` +luxury Swiss watch on a velvet cushion, macro photography, intricate details, dramatic side lighting, product photography, bokeh background +``` + +**Негативный промпт:** +``` +amateur, blurry, watermark, text, bad lighting, distorted +``` + +### 💡 Советы по составлению промптов + +1. **Используйте английский язык** — модели лучше понимают английские описания +2. **Будьте конкретны** — «sunset over mountains» лучше, чем просто «landscape» +3. **Добавляйте стиль** — укажите «oil painting», «photograph», «anime style» и т.д. +4. **Описывайте освещение** — «golden hour», «dramatic lighting», «soft morning light» +5. **Указывайте качество** — «4k», «highly detailed», «professional photography» +6. **Добавляйте композицию** — «close-up», «wide angle», «macro», «portrait» +7. **Используйте негативный промпт** — это убирает нежелательные элементы + +### 🔗 Полезные ресурсы для промптов + +- [PromptHero](https://prompthero.com/) — база промптов с примерами изображений +- [Lexica](https://lexica.art/) — поиск по Stable Diffusion изображениям и промптам +- [OpenArt](https://openart.ai/) — галерея с промптами +- [Civitai](https://civitai.com/) — модели, LoRA и примеры промптов + +--- + +## FAQ + +### Общие вопросы + +#### ❓ Что нужно для работы бота? + +Вам понадобятся: +1. **Telegram Bot Token** — получите у [@BotFather](https://t.me/BotFather) +2. **Запущенный Stable Diffusion WebUI (Automatic1111)** с флагами `--api --listen` +3. **Сервер** для запуска бота (Docker или Python) + +Бот и SD WebUI могут работать на разных машинах — главное, чтобы бот имел сетевой доступ к SD API. + +#### ❓ Можно ли использовать бота без своего GPU? + +Да! Бот может работать на слабом сервере (даже на VPS с 256 МБ RAM), а Stable Diffusion — на мощной машине с GPU в вашей локальной сети. Бот лишь передаёт запросы и возвращает изображения. + +#### ❓ Какие модели поддерживаются? + +Любые модели, загруженные в ваш SD WebUI: +- **Stable Diffusion 1.5 / 2.1** +- **SDXL** +- **Кастомные модели** (ckpt/safetensors) — любые, которые вы загрузили + +Бот автоматически подтягивает список доступных моделей из API. + +#### ❓ Поддерживаются ли LoRA? + +Да! При создании профиля вы можете выбрать LoRA из списка доступных и настроить его силу (0.0–1.0). Список LoRA загружается автоматически из SD WebUI. + +#### ❓ Бот поддерживает ControlNet? + +В текущей версии ControlNet не поддерживается. Если вам нужна эта функция — вы можете реализовать её, расширив `sd_client.py`. + +--- + +### Генерация + +#### ❓ Сколько времени занимает генерация? + +Зависит от: +- **Количества шагов** — 20 шагов ≈ 5–15 секунд на GPU среднего уровня +- **Размера изображения** — 512×512 быстрее, чем 1024×1024 +- **Модели** — SDXL медленнее, чем SD 1.5 +- **Нагрузки на GPU** — если кто-то ещё использует GPU + +Обычно генерация занимает **от 10 секунд до 2 минут**. + +#### ❓ Почему генерация занимает слишком долго? + +Возможные причины: +1. Слишком много шагов (попробуйте 20–30 вместо 50+) +2. Большой размер изображения (попробуйте 512×512 вместо 1024×1024) +3. Медленная модель (SDXL тяжелее, чем SD 1.5) +4. GPU загружен другими задачами + +#### ❓ Что такое CFG Scale? + +**CFG Scale** (Classifier-Free Guidance) определяет, насколько строго модель следует вашему промпту: +- **5.0** — модель более свободна, может добавлять свои детали +- **7.0** — стандарт, хороший баланс +- **9.0+** — модель строго следует промпту, но может стать «перегруженной» + +#### ❓ Что такое сэмплер? + +**Сэмплер** — алгоритм, который определяет, как модель «шаг за шагом» создаёт изображение из шума: +- **Euler a** — быстрый, хорошие результаты +- **DPM++ 2M Karras** — высокое качество, рекомендуемый +- **DPM++ SDE Karras** — ещё выше качество, но медленнее +- **DDIM** — быстрый, но менее детализированный + +#### ❓ Можно ли использовать seed из предыдущей генерации? + +После генерации бот показывает использованный seed. Вы можете создать профиль с конкретным seed (в текущей версии seed генерируется автоматически). Для точного воспроизведения используйте SD WebUI напрямую. + +--- + +### Профили + +#### ❓ Зачем нужны профили? + +Профили сохраняют ваши любимые настройки. Вместо того чтобы каждый раз выбирать размер, модель, сэмплер и т.д., вы создаёте профиль один раз и используете его в один клик. + +#### ❓ Сколько профилей можно создать? + +Ограничений нет — создавайте столько, сколько нужно. + +#### ❓ Что делает «профиль по умолчанию»? + +Если профиль установлен как **по умолчанию** (⭐), его настройки будут автоматически применяться при генерации. Вы можете быстро генерировать, вводя только промпт. + +--- + +### Доступ и администрирование + +#### ❓ Почему бот не отвечает мне? + +Бот работает в **закрытом режиме** — только пользователи, добавленные администратором, могут его использовать. Обратитесь к администратору, чтобы он добавил вас через админ-панель. + +#### ❓ Как добавить пользователя? + +1. Узнайте Telegram User ID пользователя (через [@userinfobot](https://t.me/userinfobot)) +2. В админ-панели нажмите **«➕ Добавить пользователя»** +3. Введите User ID +4. Выберите тип доступа (без ограничений, по времени, по количеству генераций) + +#### ❓ Как заблокировать пользователя? + +В админ-панели: **Список пользователей** → выберите пользователя → **🚫 Заблокировать**. + +#### ❓ Можно ли дать доступ всем? + +В текущей версии бот работает только в закрытом режиме. Если вы хотите открыть доступ для всех, необходимо изменить код мидлвари `AccessCheckMiddleware`. + +--- + +### Хранение изображений + +#### ❓ Что происходит с изображениями по истечении срока? + +Изображения **автоматически удаляются** с диска фоновой задачей. Записи в базе данных также очищаются. + +#### ❓ Можно ли увеличить время хранения? + +Да! В настройках хранения выберите нужное время или введите своё значение (до `MAX_IMAGE_TTL_HOURS`). + +#### ❓ Где хранятся изображения? + +В папке `images/` проекта. При использовании Docker эта папка примонтирована как volume: `./images:/app/images`. + +#### ❓ Как скачать изображение? + +Просто нажмите на изображение в Telegram и сохраните его. Все изображения отправляются как файлы. + +--- + +### Технические вопросы + +#### ❓ Можно ли запустить несколько экземпляров бота? + +Теоретически да, но они будут использовать одну базу данных и одну папку изображений, что может привести к конфликтам. Рекомендуется запускать **один экземпляр**. + +#### ❓ Как обновить бот? + +```bash +docker compose down +# Скопируйте новые файлы +docker compose up -d --build +``` + +#### ❓ Как сделать резервную копию? + +```bash +# База данных +cp data/bot.db backup_$(date +%Y%m%d).db + +# Изображения +tar -czf images_backup_$(date +%Y%m%d).tar.gz images/ +``` + +#### ❓ Как перенести бота на другой сервер? + +1. Скопируйте все файлы проекта +2. Скопируйте `data/bot.db` (база данных) +3. Скопируйте `images/` (если нужны старые изображения) +4. Настройте `.env` на новом сервере +5. Запустите бот + +#### ❓ Как ограничить доступ к SD API? + +Используйте firewall, чтобы разрешить подключения только с IP сервера бота: + +```bash +sudo ufw allow from to any port 7860 +``` + +--- + +## Структура проекта + +``` +tg-sd/ +├── main.py # Точка входа, запуск бота +├── config.py # Настройки и конфигурация +├── requirements.txt # Зависимости Python +├── Dockerfile # Образ Docker +├── docker-compose.yml # Docker Compose конфигурация +├── .env # Переменные окружения (не в Git!) +├── .env.example # Пример переменных окружения +├── .gitignore # Игнорирование файлов Git +│ +├── bot/ # Обработчики команд бота +│ ├── __init__.py +│ ├── handlers_start.py # /start, главное меню +│ ├── handlers_generation.py # Генерация изображений +│ ├── handlers_profiles.py # Управление профилями +│ └── handlers_settings.py # Настройки хранения +│ +├── sd/ # Модуль Stable Diffusion +│ ├── __init__.py +│ └── sd_client.py # Клиент для SD API +│ +├── database/ # Модуль базы данных +│ ├── __init__.py +│ └── database.py # SQLite операции +│ +├── utils/ # Утилиты +│ ├── __init__.py +│ └── image_manager.py # Управление изображениями +│ +├── images/ # Сгенерированные изображения (создаётся автоматически) +└── data/ # База данных (создаётся автоматически) +``` + +--- + +## Устранение неполадок + +### Бот не запускается + +**Проблема:** `BOT_TOKEN не установлен` + +**Решение:** Убедитесь, что в файле `.env` указан корректный токен: + +```env +BOT_TOKEN=ваш_реальный_токен_от_BotFather +``` + +### Ошибка подключения к SD API + +**Проблема:** `SD API недоступно` + +**Возможные причины и решения:** + +1. **SD WebUI не запущен** + - Запустите WebUI с флагами `--api --listen` + +2. **Неправильный адрес** + - Проверьте `SD_API_URL` в `.env` + - Убедитесь, что IP-адрес корректный + +3. **Firewall блокирует порт** + ```bash + # Проверьте доступность + curl http://192.168.1.120:7860/sdapi/v1/options + + # Откройте порт на сервере SD + sudo ufw allow 7860/tcp + ``` + +4. **WebUI слушает только localhost** + - Убедитесь, что используется флаг `--listen` + +### Таймаут генерации + +**Проблема:** Генерация занимает более 10 минут + +**Решения:** + +1. Уменьшите количество шагов в профиле +2. Уменьшите размер изображения +3. Используйте более быструю модель +4. Увеличьте таймаут в `sd_client.py` (параметр `timeout`) + +### Изображения не сохраняются + +**Проблема:** Ошибка сохранения файлов + +**Решение:** Проверьте права доступа к директории: + +```bash +# Для Docker +docker compose exec tg-sd-bot ls -la /app/ + +# Для прямого запуска +chmod 755 images/ +chmod 755 data/ +``` + +### Проблемы с базой данных + +**Проблема:** Ошибки записи в БД + +**Решение:** Убедитесь, что директория `data/` существует и доступна для записи: + +```bash +mkdir -p data +chmod 755 data +``` + +### Логи бота + +**Docker:** + +```bash +docker compose logs -f tg-sd-bot +``` + +**Прямой запуск:** + +Логи выводятся в консоль. Для записи в файл: + +```bash +python main.py 2>&1 | tee bot.log +``` + +--- + +## Безопасность + +1. **Никогда не коммитьте `.env` файл в Git** — он содержит токен бота +2. **Используйте firewall** — ограничьте доступ к порту SD API только для сервера бота +3. **Регулярно обновляйте зависимости** — проверяйте обновления пакетов +4. **Ограничьте доступ к боту** — при необходимости добавьте проверку `user_id` в обработчиках + +--- + +## Обновление бота + +```bash +# Остановка +docker compose down + +# Копирование новых файлов + +# Пересборка и запуск +docker compose up -d --build +``` + +--- + +## Резервное копирование + +Для резервного копирования сохраните: + +```bash +# База данных (профили и настройки) +cp data/bot.db backup_$(date +%Y%m%d).db + +# Изображения (если нужно) +tar -czf images_backup_$(date +%Y%m%d).tar.gz images/ +``` + +--- + +## Лицензия + +Этот проект предоставляется "как есть" без каких-либо гарантий. + +--- + +## Поддержка + +При возникновении проблем: + +1. Проверьте логи бота +2. Убедитесь, что SD API доступен +3. Проверьте корректность `.env` файла +4. Обратитесь к разделу [Устранение неполадок](#устранение-неполадок) diff --git a/bot/__init__.py b/bot/__init__.py new file mode 100644 index 0000000..19f2e6a --- /dev/null +++ b/bot/__init__.py @@ -0,0 +1 @@ +# Telegram Bot для генерации изображений через Stable Diffusion diff --git a/bot/handlers_admin.py b/bot/handlers_admin.py new file mode 100644 index 0000000..439cf94 --- /dev/null +++ b/bot/handlers_admin.py @@ -0,0 +1,513 @@ +"""Обработчики админ-панели для управления пользователями.""" + +import logging +import math +from datetime import datetime, timedelta +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 + +router = Router() +logger = logging.getLogger(__name__) + +# Константы для пагинации +USERS_PER_PAGE = 10 + + +class AddUserState(StatesGroup): + """Состояния для процесса добавления пользователя.""" + waiting_for_user_id = State() + waiting_for_access_type = State() # unlimited, time, generations + waiting_for_time_days = State() + waiting_for_generations_limit = State() + + +def is_admin(user_id: int) -> bool: + """Проверить, является ли пользователь админом.""" + return user_id == settings.ADMIN_ID + + +def check_admin(handler): + """Декоратор для проверки прав админа.""" + async def wrapper(callback_or_message, *args, **kwargs): + # Фильтруем лишние kwargs, которые не ожидаются хендлером + # Оставляем только state и другие FSM-аргументы + allowed_kwargs = {'state', 'session', 'middleware_data'} + filtered_kwargs = {k: v for k, v in kwargs.items() if k in allowed_kwargs} + + if isinstance(callback_or_message, CallbackQuery): + user_id = callback_or_message.from_user.id + elif isinstance(callback_or_message, Message): + user_id = callback_or_message.from_user.id + else: + return await callback_or_message.answer("⛔ Ошибка определения пользователя.") + + if not is_admin(user_id): + if isinstance(callback_or_message, CallbackQuery): + await callback_or_message.answer("⛔ Недостаточно прав.", show_alert=True) + else: + await callback_or_message.answer("⛔ Недостаточно прав.") + return None + return await handler(callback_or_message, *args, **filtered_kwargs) + return wrapper + + +@router.message(Command("admin")) +@check_admin +async def cmd_admin(message: Message, state: FSMContext): + """Админ-панель.""" + await state.clear() + await show_admin_panel(message) + + +@router.callback_query(F.data == "admin_panel") +@check_admin +async def admin_panel(callback: CallbackQuery, state: FSMContext): + """Показать админ-панель.""" + await state.clear() + await show_admin_panel(callback.message) + await callback.answer() + + +async def show_admin_panel(target): + """Отобразить админ-панель.""" + all_users = await db.get_all_users() + active_users = [u for u in all_users if u["is_active"]] + total_gens = sum(u.get("used_generations", 0) for u in all_users) + + text = ( + "🛡️ Админ-панель\n\n" + f"👥 Всего пользователей: {len(all_users)}\n" + f"✅ Активных: {len(active_users)}\n" + f"🎨 Всего генераций: {total_gens}\n\n" + "Управление:" + ) + + keyboard = InlineKeyboardMarkup(inline_keyboard=[ + [InlineKeyboardButton(text="➕ Добавить пользователя", callback_data="admin_add_user")], + [InlineKeyboardButton(text="📋 Список пользователей", callback_data="admin_users_list")], + [InlineKeyboardButton(text="🧹 Очистить истёкшие доступы", callback_data="admin_cleanup_expired")], + [InlineKeyboardButton(text="⬅️ Назад", callback_data="main_menu")], + ]) + + if hasattr(target, 'answer'): + try: + await target.edit_text(text, parse_mode="HTML", reply_markup=keyboard) + except Exception: + await target.answer(text, parse_mode="HTML", reply_markup=keyboard) + else: + await target.answer(text, parse_mode="HTML", reply_markup=keyboard) + + +# --- Добавление пользователя --- + +@router.callback_query(F.data == "admin_add_user") +@check_admin +async def admin_add_user(callback: CallbackQuery, state: FSMContext): + """Начать процесс добавления пользователя.""" + await state.set_state(AddUserState.waiting_for_user_id) + await callback.message.edit_text( + "➕ Добавление нового пользователя\n\n" + "Введите Telegram User ID пользователя.\n\n" + "💡 Чтобы узнать ID, пользователь может отправить " + "команду /start боту @userinfobot\n\n" + "Используйте /cancel для отмены.", + parse_mode="HTML", + ) + await callback.answer() + + +@router.message(AddUserState.waiting_for_user_id) +@check_admin +async def process_user_id(message: Message, state: FSMContext): + """Обработка введённого User ID.""" + text = message.text.strip() + if not text.isdigit(): + await message.answer("❌ Введите корректный числовой ID.") + return + + user_id = int(text) + + # Проверяем, не админ ли это + if user_id == settings.ADMIN_ID: + await message.answer("❌ Нельзя добавить самого себя как обычного пользователя.") + return + + # Проверяем, существует ли уже пользователь + existing = await db.get_user_by_id(user_id) + if existing: + status_text = "активен" if existing["is_active"] else "неактивен" + await message.answer( + f"⚠️ Пользователь {user_id} уже существует в базе.\n" + f"Статус: {status_text}\n" + f"Его параметры будут обновлены.\n\n" + "Продолжить? (да/нет)" + ) + # Сохраняем и переходим к выбору типа доступа + await state.update_data(user_id=user_id, existing_user=True) + else: + await state.update_data(user_id=user_id, existing_user=False) + + await state.set_state(AddUserState.waiting_for_access_type) + + keyboard = InlineKeyboardMarkup(inline_keyboard=[ + [InlineKeyboardButton(text="♾️ Без ограничений", callback_data="access_unlimited")], + [InlineKeyboardButton(text="⏰ По времени (дни)", callback_data="access_time")], + [InlineKeyboardButton(text="🎨 По количеству генераций", callback_data="access_gens")], + [InlineKeyboardButton(text="❌ Отмена", callback_data="admin_panel")], + ]) + + await message.answer( + "📋 Выберите тип доступа:", + parse_mode="HTML", + reply_markup=keyboard, + ) + + +@router.callback_query(F.data.startswith("access_"), AddUserState.waiting_for_access_type) +@check_admin +async def process_access_type(callback: CallbackQuery, state: FSMContext): + """Обработка выбора типа доступа.""" + access_type = callback.data.replace("access_", "") + + if access_type == "unlimited": + await state.update_data(access_type="unlimited", access_expires_at=None, max_generations=None) + await _confirm_and_add_user(callback, state) + + elif access_type == "time": + await state.set_state(AddUserState.waiting_for_time_days) + await state.update_data(access_type="time") + await callback.message.edit_text( + "⏰ Введите количество дней доступа:\n\n" + "Например: 7, 30, 90, 365\n" + "Используйте /cancel для отмены.", + parse_mode="HTML", + ) + elif access_type == "gens": + await state.set_state(AddUserState.waiting_for_generations_limit) + await state.update_data(access_type="gens") + await callback.message.edit_text( + "🎨 Введите максимальное количество генераций:\n\n" + "Например: 10, 50, 100\n" + "Используйте /cancel для отмены.", + parse_mode="HTML", + ) + + await callback.answer() + + +@router.message(AddUserState.waiting_for_time_days) +@check_admin +async def process_time_days(message: Message, state: FSMContext): + """Обработка введённого количества дней.""" + text = message.text.strip() + if not text.isdigit(): + await message.answer("❌ Введите корректное число дней.") + return + + days = int(text) + if days < 1 or days > 3650: + await message.answer("❌ Количество дней должно быть от 1 до 3650.") + return + + expires_at = datetime.now() + timedelta(days=days) + await state.update_data(access_expires_at=expires_at, max_generations=None) + await _confirm_and_add_user(message, state) + + +@router.message(AddUserState.waiting_for_generations_limit) +@check_admin +async def process_gens_limit(message: Message, state: FSMContext): + """Обработка введённого лимита генераций.""" + text = message.text.strip() + if not text.isdigit(): + await message.answer("❌ Введите корректное число.") + return + + max_gens = int(text) + if max_gens < 1 or max_gens > 100000: + await message.answer("❌ Лимит должен быть от 1 до 100000.") + return + + await state.update_data(max_generations=max_gens, access_expires_at=None) + await _confirm_and_add_user(message, state) + + +async def _confirm_and_add_user(target, state: FSMContext): + """Подтвердить и добавить пользователя.""" + data = await state.get_data() + user_id = data.get("user_id") + + if user_id is None: + await target.answer("❌ Ошибка: не указан User ID.") + await state.clear() + return + + access_expires_at = data.get("access_expires_at") + max_generations = data.get("max_generations") + + # Добавляем/обновляем пользователя + await db.add_user( + user_id=user_id, + added_by=target.from_user.id, + access_expires_at=access_expires_at, + max_generations=max_generations, + ) + + # Формируем описание доступа + if access_expires_at: + expires_str = access_expires_at.strftime("%d.%m.%Y %H:%M") + access_desc = f"⏰ До: {expires_str}" + elif max_generations: + access_desc = f"🎨 Лимит: {max_generations} генераций" + else: + access_desc = "♾️ Без ограничений" + + text = ( + f"✅ Пользователь {user_id} добавлен!\n\n" + f"{access_desc}" + ) + + keyboard = InlineKeyboardMarkup(inline_keyboard=[ + [InlineKeyboardButton(text="➕ Добавить ещё", callback_data="admin_add_user")], + [InlineKeyboardButton(text="📋 Список пользователей", callback_data="admin_users_list")], + [InlineKeyboardButton(text="🛡️ Админ-панель", callback_data="admin_panel")], + ]) + + if hasattr(target, 'edit_text'): + try: + await target.edit_text(text, parse_mode="HTML", reply_markup=keyboard) + except Exception: + await target.answer(text, parse_mode="HTML", reply_markup=keyboard) + else: + await target.answer(text, parse_mode="HTML", reply_markup=keyboard) + + await state.clear() + + +# --- Список пользователей --- + +@router.callback_query(F.data == "admin_users_list") +@check_admin +async def admin_users_list(callback: CallbackQuery, state: FSMContext): + """Показать список пользователей.""" + await state.clear() + await _show_users_list(callback.message, 0) + await callback.answer() + + +@router.callback_query(F.data.startswith("admin_users_page_")) +@check_admin +async def admin_users_page(callback: CallbackQuery, state: FSMContext): + """Переключение страницы пользователей.""" + page = int(callback.data.split("_")[-1]) + await _show_users_list(callback.message, page) + await callback.answer() + + +async def _show_users_list(target, page: int): + """Отобразить список пользователей.""" + all_users = await db.get_all_users() + total_pages = math.ceil(len(all_users) / USERS_PER_PAGE) if all_users else 1 + + if page >= total_pages: + page = total_pages - 1 + + start = page * USERS_PER_PAGE + end = start + USERS_PER_PAGE + page_users = all_users[start:end] + + if not all_users: + text = "📋 Список пользователей пуст" + keyboard = InlineKeyboardMarkup(inline_keyboard=[ + [InlineKeyboardButton(text="➕ Добавить пользователя", callback_data="admin_add_user")], + [InlineKeyboardButton(text="🛡️ Админ-панель", callback_data="admin_panel")], + ]) + else: + text = f"👥 Пользователи ({len(all_users)})\n\n" + for i, user in enumerate(page_users, start + 1): + status = "👑" if user["is_admin"] else ("✅" if user["is_active"] else "❌") + + # Формируем описание доступа + if user["is_admin"]: + access_info = "админ" + elif user.get("access_expires_at"): + try: + exp = datetime.fromisoformat(user["access_expires_at"]) + if exp > datetime.now(): + days_left = (exp - datetime.now()).days + access_info = f"до {exp.strftime('%d.%m.%Y')} ({days_left} дн.)" + else: + access_info = "истёк" + except Exception: + access_info = "—" + elif user.get("max_generations"): + used = user.get("used_generations", 0) + access_info = f"{used}/{user['max_generations']} ген." + else: + access_info = "без ограничений" + + name = user.get("display_name") or user.get("username") or f"ID: {user['user_id']}" + text += f"{i}. {status} {user['user_id']} — {name}\n" + text += f" {access_info}\n\n" + + # Навигация + nav_row = [] + if page > 0: + nav_row.append(InlineKeyboardButton(text="⬅️", callback_data=f"admin_users_page_{page - 1}")) + nav_row.append(InlineKeyboardButton(text=f"{page + 1}/{total_pages}", callback_data="noop")) + if page < total_pages - 1: + nav_row.append(InlineKeyboardButton(text="➡️", callback_data=f"admin_users_page_{page + 1}")) + + kb_buttons = [] + for user in page_users: + if not user["is_admin"]: + kb_buttons.append([ + InlineKeyboardButton( + text=f"{'✅' if user['is_active'] else '❌'} {user['user_id']}", + callback_data=f"admin_user_manage_{user['user_id']}" + ) + ]) + + kb_buttons.append(nav_row) + kb_buttons.append([InlineKeyboardButton(text="➕ Добавить", callback_data="admin_add_user")]) + kb_buttons.append([InlineKeyboardButton(text="🛡️ Админ-панель", callback_data="admin_panel")]) + + keyboard = InlineKeyboardMarkup(inline_keyboard=kb_buttons) + + if hasattr(target, 'edit_text'): + try: + await target.edit_text(text, parse_mode="HTML", reply_markup=keyboard) + except Exception: + await target.answer(text, parse_mode="HTML", reply_markup=keyboard) + else: + await target.answer(text, parse_mode="HTML", reply_markup=keyboard) + + +# --- Управление конкретным пользователем --- + +@router.callback_query(F.data.startswith("admin_user_manage_")) +@check_admin +async def admin_user_manage(callback: CallbackQuery): + """Управление конкретным пользователем.""" + user_id = int(callback.data.split("_")[-1]) + user = await db.get_user_by_id(user_id) + + if not user: + await callback.answer("Пользователь не найден.", show_alert=True) + return + + # Формируем информацию + status = "👑 Админ" if user["is_admin"] else ("✅ Активен" if user["is_active"] else "❌ Неактивен") + + if user.get("access_expires_at"): + try: + exp = datetime.fromisoformat(user["access_expires_at"]) + expires_str = exp.strftime("%d.%m.%Y %H:%M") + except Exception: + expires_str = "—" + else: + expires_str = "не указан" + + max_gens = user.get("max_generations") or "без ограничений" + used_gens = user.get("used_generations", 0) + added_by = user.get("added_by", "—") + added_at = user.get("added_at", "—") + + name = user.get("display_name") or user.get("username") or f"ID: {user_id}" + + text = ( + f"👤 {name}\n\n" + f"ID: {user_id}\n" + f"Статус: {status}\n" + f"Добавлен: {added_at} (админом {added_by})\n" + f"Истекает: {expires_str}\n" + f"Генерации: {used_gens}/{max_gens}" + ) + + kb_buttons = [] + if not user["is_admin"]: + if user["is_active"]: + kb_buttons.append([InlineKeyboardButton(text="🚫 Заблокировать", callback_data=f"admin_user_block_{user_id}")]) + else: + kb_buttons.append([InlineKeyboardButton(text="✅ Разблокировать", callback_data=f"admin_user_unblock_{user_id}")]) + kb_buttons.append([InlineKeyboardButton(text="🔄 Обновить доступ", callback_data=f"admin_user_update_{user_id}")]) + kb_buttons.append([InlineKeyboardButton(text="🗑️ Удалить", callback_data=f"admin_user_delete_{user_id}")]) + + kb_buttons.append([InlineKeyboardButton(text="⬅️ Назад к списку", callback_data="admin_users_list")]) + + keyboard = InlineKeyboardMarkup(inline_keyboard=kb_buttons) + + await callback.message.edit_text(text, parse_mode="HTML", reply_markup=keyboard) + await callback.answer() + + +@router.callback_query(F.data.startswith("admin_user_block_")) +@check_admin +async def admin_user_block(callback: CallbackQuery): + """Заблокировать пользователя.""" + user_id = int(callback.data.split("_")[-1]) + await db.remove_user(user_id) + await callback.answer("🚫 Пользователь заблокирован.") + await admin_user_manage(callback) + + +@router.callback_query(F.data.startswith("admin_user_unblock_")) +@check_admin +async def admin_user_unblock(callback: CallbackQuery): + """Разблокировать пользователя.""" + user_id = int(callback.data.split("_")[-1]) + await db.update_user_access(user_id) + await callback.answer("✅ Пользователь разблокирован.") + await admin_user_manage(callback) + + +@router.callback_query(F.data.startswith("admin_user_delete_")) +@check_admin +async def admin_user_delete(callback: CallbackQuery): + """Удалить пользователя (окончательно).""" + user_id = int(callback.data.split("_")[-1]) + # В текущей схеме просто деактивируем + await db.remove_user(user_id) + await callback.answer("🗑️ Пользователь удалён.") + await admin_users_list(callback) + + +@router.callback_query(F.data.startswith("admin_user_update_")) +@check_admin +async def admin_user_update(callback: CallbackQuery, state: FSMContext): + """Обновить параметры доступа пользователя.""" + user_id = int(callback.data.split("_")[-1]) + await state.update_data(update_user_id=user_id) + await state.set_state(AddUserState.waiting_for_access_type) + + keyboard = InlineKeyboardMarkup(inline_keyboard=[ + [InlineKeyboardButton(text="♾️ Без ограничений", callback_data="access_unlimited")], + [InlineKeyboardButton(text="⏰ По времени (дни)", callback_data="access_time")], + [InlineKeyboardButton(text="🎨 По количеству генераций", callback_data="access_gens")], + [InlineKeyboardButton(text="❌ Отмена", callback_data=f"admin_user_manage_{user_id}")], + ]) + + await callback.message.edit_text( + "🔄 Обновление параметров доступа\n\n" + "Выберите новый тип доступа:", + parse_mode="HTML", + reply_markup=keyboard, + ) + await callback.answer() + + +# --- Очистка истёкших доступов --- + +@router.callback_query(F.data == "admin_cleanup_expired") +@check_admin +async def admin_cleanup(callback: CallbackQuery): + """Очистка истёкших доступов.""" + deactivated = await db.deactivate_expired_users() + await callback.answer(f"🧹 Деактивировано пользователей: {deactivated}", show_alert=True) + await admin_panel(callback) diff --git a/bot/handlers_generation.py b/bot/handlers_generation.py new file mode 100644 index 0000000..9d649b0 --- /dev/null +++ b/bot/handlers_generation.py @@ -0,0 +1,241 @@ +"""Обработчики генерации изображений (txt2img).""" + +import logging +import os +import uuid +from datetime import datetime +from html import escape +from aiogram import Router, F +from aiogram.types import ( + Message, + CallbackQuery, + InlineKeyboardMarkup, + InlineKeyboardButton, + FSInputFile, +) +from aiogram.fsm.context import FSMContext +from aiogram.fsm.state import State, StatesGroup + +from config import settings +from database.database import db +from sd.sd_client import sd_client +from utils.image_manager import get_image_path + +router = Router() +logger = logging.getLogger(__name__) + + +class Txt2ImgState(StatesGroup): + """Состояния для процесса txt2img.""" + waiting_for_prompt = State() + waiting_for_negative_prompt = State() + generating = State() + + +def get_generation_keyboard() -> InlineKeyboardMarkup: + """Клавиатура для генерации txt2img.""" + return InlineKeyboardMarkup(inline_keyboard=[ + [ + InlineKeyboardButton(text="📝 Ввести промпт", callback_data="txt2img_enter_prompt"), + ], + [ + InlineKeyboardButton(text="📋 Использовать профиль", callback_data="txt2img_use_profile"), + ], + [ + InlineKeyboardButton(text="⬅️ Назад", callback_data="main_menu"), + ], + ]) + + +@router.callback_query(F.data == "gen_txt2img") +async def start_txt2img(callback: CallbackQuery, state: FSMContext): + """Начать процесс txt2img.""" + await state.clear() + await callback.message.answer( + "🎨 Генерация изображения (txt2img)\n\n" + "Введите описание изображения (промпт).\n" + "Можете использовать /cancel для отмены.", + parse_mode="HTML", + reply_markup=get_generation_keyboard(), + ) + await callback.answer() + + +@router.callback_query(F.data == "txt2img_enter_prompt") +async def enter_prompt_txt(callback: CallbackQuery, state: FSMContext): + """Ввод промпта для txt2img.""" + await state.set_state(Txt2ImgState.waiting_for_prompt) + await callback.message.edit_text( + "📝 Введите промпт:\n\n" + "Опишите изображение, которое хотите сгенерировать.\n" + "Используйте /cancel для отмены.", + parse_mode="HTML", + ) + await callback.answer() + + +@router.callback_query(F.data == "txt2img_use_profile") +async def use_profile_txt(callback: CallbackQuery, state: FSMContext): + """Использовать профиль для txt2img.""" + profiles = await db.get_user_profiles(callback.from_user.id) + if not profiles: + await callback.answer("У вас пока нет профилей. Создайте профиль в разделе 'Профили'.", show_alert=True) + return + + keyboard = InlineKeyboardMarkup(inline_keyboard=[ + [InlineKeyboardButton(text=p["name"], callback_data=f"txt2img_profile_{p['id']}")] + for p in profiles + ]) + keyboard.inline_keyboard.append([InlineKeyboardButton(text="⬅️ Назад", callback_data="gen_txt2img")]) + + await state.set_state(Txt2ImgState.waiting_for_prompt) + await callback.message.edit_text( + "📋 Выберите профиль:\n\n" + "После выбора будет запрошен промпт.", + parse_mode="HTML", + reply_markup=keyboard, + ) + await callback.answer() + + +@router.callback_query(F.data.startswith("txt2img_profile_")) +async def select_profile_txt(callback: CallbackQuery, state: FSMContext): + """Выбор профиля для txt2img.""" + profile_id = int(callback.data.split("_")[-1]) + await state.update_data(profile_id=profile_id) + profile = await db.get_profile(profile_id) + + await state.set_state(Txt2ImgState.waiting_for_prompt) + await callback.message.edit_text( + f"✅ Профиль {profile['name']} выбран.\n\n" + f"Параметры: {profile['width']}x{profile['height']}, " + f"шагов: {profile['steps']}, CFG: {profile['cfg_scale']}, " + f"сэмплер: {profile['sampler']}, шедулер: {profile.get('scheduler', 'automatic')}\n\n" + f"📝 Введите промпт:\n" + f"Используйте /cancel для отмены.", + parse_mode="HTML", + ) + await callback.answer() + + +@router.message(Txt2ImgState.waiting_for_prompt, F.text) +async def process_prompt_txt(message: Message, state: FSMContext): + """Обработка введённого промпта для txt2img.""" + prompt = message.text.strip() + if not prompt: + await message.answer("Промпт не может быть пустым. Введите описание изображения.") + return + + await state.update_data(prompt=prompt) + + # Запрашиваем негативный промпт + await state.set_state(Txt2ImgState.waiting_for_negative_prompt) + await message.answer( + "🚫 Введите негативный промпт (необязательно)\n\n" + "Опишите, что НЕ должно быть на изображении.\n" + "Отправьте - чтобы пропустить.\n" + "Используйте /cancel для отмены.", + parse_mode="HTML", + ) + + +@router.message(Txt2ImgState.waiting_for_negative_prompt) +async def process_negative_prompt_txt(message: Message, state: FSMContext): + """Обработка негативного промпта для txt2img.""" + negative_prompt = message.text.strip() + if negative_prompt == "-": + negative_prompt = "" + + await state.update_data(negative_prompt=negative_prompt) + await state.set_state(Txt2ImgState.generating) + + # Получаем данные + data = await state.get_data() + prompt = data.get("prompt", "") + profile_id = data.get("profile_id") + + # Загружаем настройки пользователя для TTL + user_settings = await db.get_user_settings(message.from_user.id) + ttl_hours = user_settings.get("image_ttl_hours", settings.DEFAULT_IMAGE_TTL_HOURS) + + # Если есть профиль, используем его настройки + profile = None + if profile_id: + profile = await db.get_profile(profile_id) + + # Отправляем сообщение о начале генерации + status_msg = await message.answer("⏳ Генерация изображения...\n\nЭто может занять некоторое время.", parse_mode="HTML") + + # Вызываем API + result = await sd_client.txt2img( + prompt=prompt, + negative_prompt=negative_prompt or (profile["negative_prompt"] if profile else ""), + width=profile["width"] if profile else 512, + height=profile["height"] if profile else 512, + steps=profile["steps"] if profile else 20, + cfg_scale=profile["cfg_scale"] if profile else 7.0, + sampler=profile["sampler"] if profile else "Euler a", + scheduler=profile.get("scheduler", "automatic") if profile else "automatic", + model=profile.get("model") if profile else None, + lora=profile.get("lora") if profile and profile.get("lora") else None, + lora_strength=profile.get("lora_strength", 0.8) if profile else 0.8, + ) + + if result is None: + await status_msg.edit_text("❌ Ошибка генерации.\n\nПроверьте соединение с SD API и попробуйте снова.") + await state.clear() + return + + image_bytes, info = result + + # Сохраняем изображение + filename = f"{uuid.uuid4().hex[:12]}.png" + file_path = get_image_path(message.from_user.id, filename) + + try: + with open(file_path, "wb") as f: + f.write(image_bytes) + except Exception as e: + logger.error(f"Ошибка сохранения изображения: {e}") + await status_msg.edit_text("❌ Ошибка сохранения изображения.") + await state.clear() + return + + # Сохраняем в БД + await db.add_generated_image( + user_id=message.from_user.id, + file_path=file_path, + prompt=prompt, + ttl_hours=ttl_hours, + ) + + # Увеличиваем счётчик генераций + await db.increment_generation_count(message.from_user.id) + + # Отправляем изображение + info_text = ( + f"✅ Изображение сгенерировано!\n\n" + f"Промпт: {escape(info['prompt'][:500])}\n" + f"Размер: {info['width']}x{info['height']}\n" + f"Шагов: {info['steps']}\n" + f"CFG Scale: {info['cfg_scale']}\n" + f"Сэмплер: {info['sampler']}\n" + f"Шедулер: {info.get('scheduler', 'automatic')}\n" + f"Seed: {info['seed']}\n" + f"Модель: {info['model']}\n" + f"Время хранения: {ttl_hours} ч." + ) + + keyboard = InlineKeyboardMarkup(inline_keyboard=[ + [InlineKeyboardButton(text="🔄 Сгенерировать ещё", callback_data="gen_txt2img")], + [InlineKeyboardButton(text="📋 Главное меню", callback_data="main_menu")], + ]) + + await message.answer_photo( + photo=FSInputFile(file_path), + caption=info_text, + parse_mode="HTML", + reply_markup=keyboard, + ) + await status_msg.delete() + await state.clear() diff --git a/bot/handlers_profiles.py b/bot/handlers_profiles.py new file mode 100644 index 0000000..c2e0554 --- /dev/null +++ b/bot/handlers_profiles.py @@ -0,0 +1,634 @@ +"""Обработчики управления профилями.""" + +import logging +import math +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 database.database import db +from sd.sd_client import sd_client + +router = Router() +logger = logging.getLogger(__name__) + +# Константы для пагинации +ITEMS_PER_PAGE = 8 + + +class ProfileCreateState(StatesGroup): + """Состояния для создания профиля.""" + waiting_for_name = State() + waiting_for_width = State() + waiting_for_height = State() + waiting_for_steps = State() + waiting_for_cfg = State() + waiting_for_sampler = State() + waiting_for_scheduler = State() + waiting_for_model = State() + waiting_for_lora = State() + waiting_for_lora_strength = State() + waiting_for_negative_prompt = State() + + +def _build_pagination_keyboard(items: list[str], current_page: int, callback_prefix: str, skip_callback: str) -> InlineKeyboardMarkup: + """Построить клавиатуру с пагинацией.""" + total_pages = math.ceil(len(items) / ITEMS_PER_PAGE) if items else 1 + start = current_page * ITEMS_PER_PAGE + end = start + ITEMS_PER_PAGE + page_items = items[start:end] + + keyboard = [] + for item in page_items: + # Callback data ограничен 64 байтами + safe_data = item[:56] + keyboard.append([InlineKeyboardButton( + text=item[:30], + callback_data=f"{callback_prefix}_{safe_data}" + )]) + + # Навигация + nav_row = [] + if current_page > 0: + nav_row.append(InlineKeyboardButton(text="⬅️", callback_data=f"{callback_prefix}_page_{current_page - 1}")) + nav_row.append(InlineKeyboardButton(text=f"{current_page + 1}/{total_pages}", callback_data="noop")) + if current_page < total_pages - 1: + nav_row.append(InlineKeyboardButton(text="➡️", callback_data=f"{callback_prefix}_page_{current_page + 1}")) + keyboard.append(nav_row) + + keyboard.append([InlineKeyboardButton(text="⏭️ Пропустить", callback_data=skip_callback)]) + + return InlineKeyboardMarkup(inline_keyboard=keyboard) + + +@router.callback_query(F.data == "profiles_menu") +async def profiles_menu(callback: CallbackQuery): + """Меню профилей.""" + profiles = await db.get_user_profiles(callback.from_user.id) + + if not profiles: + text = ( + "📋 Управление профилями\n\n" + "У вас пока нет профилей. Создайте первый профиль для быстрой генерации." + ) + keyboard = InlineKeyboardMarkup(inline_keyboard=[ + [InlineKeyboardButton(text="➕ Создать профиль", callback_data="profile_create")], + [InlineKeyboardButton(text="⬅️ Назад", callback_data="main_menu")], + ]) + else: + text = "📋 Ваши профили:\n\n" + kb_buttons = [] + for p in profiles: + default_mark = " ⭐" if p["is_default"] else "" + text += f"{p['name']}{default_mark}\n" + text += f" Размер: {p['width']}x{p['height']}, Шагов: {p['steps']}, CFG: {p['cfg_scale']}\n" + text += f" Сэмплер: {p['sampler']}, Шедулер: {p.get('scheduler', 'automatic')}\n" + if p.get("model"): + text += f" Модель: {p['model']}\n" + if p.get("lora"): + text += f" LoRA: {p['lora']} ({p['lora_strength']})\n" + text += "\n" + kb_buttons.append([ + InlineKeyboardButton( + text=f"{'⭐ ' if p['is_default'] else ''}{p['name']}", + callback_data=f"profile_view_{p['id']}" + ) + ]) + + kb_buttons.append([InlineKeyboardButton(text="➕ Создать профиль", callback_data="profile_create")]) + kb_buttons.append([InlineKeyboardButton(text="⬅️ Назад", callback_data="main_menu")]) + + keyboard = InlineKeyboardMarkup(inline_keyboard=kb_buttons) + + await callback.message.edit_text(text, parse_mode="HTML", reply_markup=keyboard) + await callback.answer() + + +@router.callback_query(F.data == "profile_create") +async def profile_create(callback: CallbackQuery, state: FSMContext): + """Начать создание профиля.""" + await state.set_state(ProfileCreateState.waiting_for_name) + await callback.message.edit_text( + "➕ Создание нового профиля\n\n" + "Введите название профиля:\n" + "Используйте /cancel для отмены.", + parse_mode="HTML", + ) + await callback.answer() + + +@router.message(ProfileCreateState.waiting_for_name) +async def profile_name(message: Message, state: FSMContext): + """Обработка названия профиля.""" + name = message.text.strip() + if len(name) < 2: + await message.answer("Название должно быть не менее 2 символов.") + return + + await state.update_data(name=name) + await state.set_state(ProfileCreateState.waiting_for_width) + + keyboard = InlineKeyboardMarkup(inline_keyboard=[ + [ + InlineKeyboardButton(text="512x512", callback_data="size_512_512"), + InlineKeyboardButton(text="768x768", callback_data="size_768_768"), + ], + [ + InlineKeyboardButton(text="512x768 (портрет)", callback_data="size_512_768"), + InlineKeyboardButton(text="768x512 (ландшафт)", callback_data="size_768_512"), + ], + [ + InlineKeyboardButton(text="1024x1024", callback_data="size_1024_1024"), + ], + [ + InlineKeyboardButton(text="✏️ Ввести свой размер", callback_data="size_custom"), + ], + ]) + + await message.answer( + "📐 Выберите размер изображения:", + parse_mode="HTML", + reply_markup=keyboard, + ) + + +@router.callback_query(F.data.startswith("size_"), ProfileCreateState.waiting_for_width) +async def profile_size(callback: CallbackQuery, state: FSMContext): + """Обработка выбора размера.""" + if callback.data == "size_custom": + await callback.message.edit_text( + "📐 Введите ширину (например, 512, 768, 1024):", + parse_mode="HTML", + ) + return + + parts = callback.data.replace("size_", "").split("_") + width, height = int(parts[0]), int(parts[1]) + await state.update_data(width=width, height=height) + await state.set_state(ProfileCreateState.waiting_for_steps) + + keyboard = InlineKeyboardMarkup(inline_keyboard=[ + [ + InlineKeyboardButton(text="20 (быстро)", callback_data="steps_20"), + InlineKeyboardButton(text="30 (качество)", callback_data="steps_30"), + ], + [ + InlineKeyboardButton(text="40 (высокое качество)", callback_data="steps_40"), + ], + [ + InlineKeyboardButton(text="✏️ Своё значение", callback_data="steps_custom"), + ], + ]) + + await callback.message.edit_text( + f"✅ Размер: {width}x{height}\n\n" + "🔢 Выберите количество шагов:", + parse_mode="HTML", + reply_markup=keyboard, + ) + await callback.answer() + + +@router.message(ProfileCreateState.waiting_for_width, F.text.isdigit()) +async def profile_custom_width(message: Message, state: FSMContext): + """Ввод пользовательской ширины.""" + width = int(message.text.strip()) + if width < 64 or width > 2048: + await message.answer("Размер должен быть от 64 до 2048 пикселей.") + return + await state.update_data(width=width) + await message.answer("📐 Теперь введите высоту:") + await state.set_state(ProfileCreateState.waiting_for_height) + + +@router.message(ProfileCreateState.waiting_for_height, F.text.isdigit()) +async def profile_height(message: Message, state: FSMContext): + """Обработка высоты.""" + height = int(message.text.strip()) + if height < 64 or height > 2048: + await message.answer("Размер должен быть от 64 до 2048 пикселей.") + return + + await state.update_data(height=height) + await state.set_state(ProfileCreateState.waiting_for_steps) + + data = await state.get_data() + width = data.get("width", "?") + + keyboard = InlineKeyboardMarkup(inline_keyboard=[ + [ + InlineKeyboardButton(text="20 (быстро)", callback_data="steps_20"), + InlineKeyboardButton(text="30 (качество)", callback_data="steps_30"), + ], + [ + InlineKeyboardButton(text="40 (высокое качество)", callback_data="steps_40"), + ], + [ + InlineKeyboardButton(text="✏️ Своё значение", callback_data="steps_custom"), + ], + ]) + + await message.answer( + f"✅ Размер: {width}x{height}\n\n" + "🔢 Выберите количество шагов:", + parse_mode="HTML", + reply_markup=keyboard, + ) + + +@router.callback_query(F.data.startswith("steps_"), ProfileCreateState.waiting_for_steps) +async def profile_steps(callback: CallbackQuery, state: FSMContext): + """Обработка выбора шагов.""" + if callback.data == "steps_custom": + await callback.message.edit_text( + "🔢 Введите количество шагов (10-150):", + parse_mode="HTML", + ) + return + + steps = int(callback.data.replace("steps_", "")) + await state.update_data(steps=steps) + await state.set_state(ProfileCreateState.waiting_for_cfg) + + keyboard = InlineKeyboardMarkup(inline_keyboard=[ + [ + InlineKeyboardButton(text="5.0 (свободнее)", callback_data="cfg_5"), + InlineKeyboardButton(text="7.0 (стандарт)", callback_data="cfg_7"), + ], + [ + InlineKeyboardButton(text="9.0 (строже)", callback_data="cfg_9"), + InlineKeyboardButton(text="✏️ Своё значение", callback_data="cfg_custom"), + ], + ]) + + await callback.message.edit_text( + f"✅ Шагов: {steps}\n\n" + "🎚️ Выберите CFG Scale:\n" + "Определяет, насколько строго модель следует промпту.", + parse_mode="HTML", + reply_markup=keyboard, + ) + await callback.answer() + + +@router.message(ProfileCreateState.waiting_for_steps, F.text.isdigit()) +async def profile_custom_steps(message: Message, state: FSMContext): + """Ввод пользовательских шагов.""" + steps = int(message.text.strip()) + if steps < 10 or steps > 150: + await message.answer("Количество шагов должно быть от 10 до 150.") + return + await state.update_data(steps=steps) + await state.set_state(ProfileCreateState.waiting_for_cfg) + await message.answer( + "🎚️ Введите CFG Scale (обычно 5.0-12.0):\n" + "Отправьте - для стандартного значения (7.0).", + parse_mode="HTML", + ) + + +@router.callback_query(F.data.startswith("cfg_"), ProfileCreateState.waiting_for_cfg) +async def profile_cfg(callback: CallbackQuery, state: FSMContext): + """Обработка выбора CFG Scale.""" + if callback.data == "cfg_custom": + await callback.message.edit_text( + "🎚️ Введите CFG Scale (1.0-30.0):", + parse_mode="HTML", + ) + return + + cfg = float(callback.data.replace("cfg_", "")) + await state.update_data(cfg_scale=cfg) + await state.set_state(ProfileCreateState.waiting_for_sampler) + + # Получаем сэмплеры из API + samplers = await sd_client.get_samplers() + if not samplers: + samplers = ["Euler a", "Euler", "DPM++ 2M Karras", "DPM++ SDE Karras", "DDIM"] + + await state.update_data(available_samplers=samplers) + keyboard = _build_pagination_keyboard(samplers, 0, "sampler", "sampler_skip") + + await callback.message.edit_text( + f"✅ CFG Scale: {cfg}\n\n" + "🔄 Выберите сэмплер:\n" + f"Доступно: {len(samplers)}", + parse_mode="HTML", + reply_markup=keyboard, + ) + await callback.answer() + + +@router.callback_query(F.data.startswith("sampler_page_"), ProfileCreateState.waiting_for_sampler) +async def profile_sampler_page(callback: CallbackQuery, state: FSMContext): + """Переключение страницы сэмплеров.""" + page = int(callback.data.split("_")[-1]) + samplers = (await state.get_data()).get("available_samplers", []) + + keyboard = _build_pagination_keyboard(samplers, page, "sampler", "sampler_skip") + await callback.message.edit_reply_markup(reply_markup=keyboard) + await callback.answer() + + +@router.callback_query(F.data == "sampler_skip", ProfileCreateState.waiting_for_sampler) +async def profile_sampler_skip(callback: CallbackQuery, state: FSMContext): + """Пропустить выбор сэмплера.""" + await state.update_data(sampler="Euler a", scheduler="automatic") + await _show_model_selection(callback, state) + + +@router.callback_query(F.data.startswith("sampler_"), ProfileCreateState.waiting_for_sampler) +async def profile_sampler_select(callback: CallbackQuery, state: FSMContext): + """Обработка выбора сэмплера.""" + sampler = callback.data.replace("sampler_", "") + await state.update_data(sampler=sampler, scheduler="automatic") + await _show_model_selection(callback, state) + + +async def _show_model_selection(callback: CallbackQuery, state: FSMContext): + """Показать выбор модели.""" + data = await state.get_data() + sampler = data.get("sampler", "Euler a") + + await callback.answer("⏳ Загружаю список моделей...", show_alert=False) + models = await sd_client.get_models() + if not models: + models = ["— Не менять —"] + else: + models = ["— Не менять —"] + models + + await state.update_data(available_models=models) + await state.set_state(ProfileCreateState.waiting_for_model) + + keyboard = _build_pagination_keyboard(models, 0, "model", "model_skip") + + await callback.message.edit_text( + f"✅ Сэмплер: {sampler}\n\n" + "🤖 Выберите модель:\n" + f"Доступно: {len(models) - 1}", + parse_mode="HTML", + reply_markup=keyboard, + ) + + +@router.callback_query(F.data.startswith("model_page_"), ProfileCreateState.waiting_for_model) +async def profile_model_page(callback: CallbackQuery, state: FSMContext): + """Переключение страницы моделей.""" + page = int(callback.data.split("_")[-1]) + models = (await state.get_data()).get("available_models", []) + + keyboard = _build_pagination_keyboard(models, page, "model", "model_skip") + await callback.message.edit_reply_markup(reply_markup=keyboard) + await callback.answer() + + +@router.callback_query(F.data == "model_skip", ProfileCreateState.waiting_for_model) +async def profile_model_skip(callback: CallbackQuery, state: FSMContext): + """Пропустить выбор модели.""" + await state.update_data(model="") + await _show_lora_selection(callback, state) + + +@router.callback_query(F.data.startswith("model_"), ProfileCreateState.waiting_for_model) +async def profile_model_select(callback: CallbackQuery, state: FSMContext): + """Обработка выбора модели.""" + model = callback.data.replace("model_", "") + if model == "— Не менять —": + model = "" + await state.update_data(model=model) + await _show_lora_selection(callback, state) + + +async def _show_lora_selection(callback: CallbackQuery, state: FSMContext): + """Показать выбор LoRA.""" + data = await state.get_data() + model = data.get("model", "") + + await callback.answer("⏳ Загружаю список LoRA...", show_alert=False) + loras = await sd_client.get_loras() + + await state.update_data(available_loras=loras) + await state.set_state(ProfileCreateState.waiting_for_lora) + + if not loras: + await state.update_data(lora="", lora_strength=0.8) + await _show_negative_prompt_request(callback, state) + return + + loras = ["— Не использовать —"] + loras + keyboard = _build_pagination_keyboard(loras, 0, "lora", "lora_skip") + + await callback.message.edit_text( + f"✅ Модель: {model or 'текущая'}\n\n" + "🎨 Выберите LoRA:\n" + f"Доступно: {len(loras) - 1}", + parse_mode="HTML", + reply_markup=keyboard, + ) + + +@router.callback_query(F.data.startswith("lora_page_"), ProfileCreateState.waiting_for_lora) +async def profile_lora_page(callback: CallbackQuery, state: FSMContext): + """Переключение страницы LoRA.""" + page = int(callback.data.split("_")[-1]) + loras = (await state.get_data()).get("available_loras", []) + + keyboard = _build_pagination_keyboard(loras, page, "lora", "lora_skip") + await callback.message.edit_reply_markup(reply_markup=keyboard) + await callback.answer() + + +@router.callback_query(F.data == "lora_skip", ProfileCreateState.waiting_for_lora) +async def profile_lora_skip(callback: CallbackQuery, state: FSMContext): + """Пропустить выбор LoRA.""" + await state.update_data(lora="", lora_strength=0.8) + await _show_negative_prompt_request(callback, state) + + +@router.callback_query(F.data.startswith("lora_"), ProfileCreateState.waiting_for_lora) +async def profile_lora_select(callback: CallbackQuery, state: FSMContext): + """Обработка выбора LoRA.""" + lora = callback.data.replace("lora_", "") + if lora == "— Не использовать —": + await state.update_data(lora="", lora_strength=0.8) + await _show_negative_prompt_request(callback, state) + return + + await state.update_data(lora=lora) + await state.set_state(ProfileCreateState.waiting_for_lora_strength) + + await callback.message.edit_text( + f"✅ LoRA: {lora}\n\n" + "💪 Введите силу LoRA (0.0-1.0, обычно 0.8):\n" + "Отправьте - для стандартного значения (0.8).", + parse_mode="HTML", + ) + + +@router.message(ProfileCreateState.waiting_for_lora_strength) +async def profile_lora_strength(message: Message, state: FSMContext): + """Обработка силы LoRA.""" + text = message.text.strip() + if text == "-": + lora_strength = 0.8 + else: + try: + lora_strength = float(text) + if lora_strength < 0.0 or lora_strength > 1.0: + await message.answer("Сила LoRA должна быть от 0.0 до 1.0.") + return + except ValueError: + await message.answer("Введите корректное число (например, 0.8).") + return + + await state.update_data(lora_strength=lora_strength) + await _show_negative_prompt_request_msg(message, state) + + +async def _show_negative_prompt_request(callback: CallbackQuery, state: FSMContext): + """Запросить негативный промпт (из callback).""" + data = await state.get_data() + await state.set_state(ProfileCreateState.waiting_for_negative_prompt) + + lora_text = data.get("lora", "") + lora_info = f"LoRA: {lora_text} ({data.get('lora_strength', 0.8)})" if lora_text else "LoRA: нет" + + await callback.message.edit_text( + f"✅ Модель: {data.get('model', '') or 'текущая'}\n" + f"✅ {lora_info}\n\n" + "🚫 Введите негативный промпт (необязательно)\n\n" + "Опишите, что НЕ должно быть на изображении.\n" + "Отправьте - чтобы пропустить.", + parse_mode="HTML", + ) + + +async def _show_negative_prompt_request_msg(message: Message, state: FSMContext): + """Запросить негативный промпт (из message).""" + data = await state.get_data() + await state.set_state(ProfileCreateState.waiting_for_negative_prompt) + + lora_text = data.get("lora", "") + lora_info = f"LoRA: {lora_text} ({data.get('lora_strength', 0.8)})" if lora_text else "LoRA: нет" + + await message.answer( + f"✅ Модель: {data.get('model', '') or 'текущая'}\n" + f"✅ {lora_info}\n\n" + "🚫 Введите негативный промпт (необязательно)\n\n" + "Опишите, что НЕ должно быть на изображении.\n" + "Отправьте - чтобы пропустить.", + parse_mode="HTML", + ) + + +@router.message(ProfileCreateState.waiting_for_negative_prompt) +async def profile_negative(message: Message, state: FSMContext): + """Обработка негативного промпта и завершение создания профиля.""" + negative_prompt = message.text.strip() + if negative_prompt == "-": + negative_prompt = "" + + data = await state.get_data() + + profile_id = await db.create_profile( + user_id=message.from_user.id, + name=data["name"], + width=data.get("width", 512), + height=data.get("height", 512), + steps=data.get("steps", 20), + cfg_scale=data.get("cfg_scale", 7.0), + sampler=data.get("sampler", "Euler a"), + scheduler=data.get("scheduler", "automatic"), + model=data.get("model", ""), + lora=data.get("lora", ""), + lora_strength=data.get("lora_strength", 0.8), + negative_prompt=negative_prompt, + is_default=False, + ) + + keyboard = InlineKeyboardMarkup(inline_keyboard=[ + [InlineKeyboardButton(text="⭐ Сделать профилем по умолчанию", callback_data=f"profile_set_default_{profile_id}")], + [InlineKeyboardButton(text="📋 К списку профилей", callback_data="profiles_menu")], + ]) + + await message.answer( + f"✅ Профиль '{data['name']}' создан!\n\n" + f"Размер: {data.get('width', 512)}x{data.get('height', 512)}\n" + f"Шагов: {data.get('steps', 20)}\n" + f"CFG Scale: {data.get('cfg_scale', 7.0)}\n" + f"Сэмплер: {data.get('sampler', 'Euler a')}\n" + f"Шедулер: {data.get('scheduler', 'automatic')}\n" + f"Модель: {data.get('model', 'текущая') or 'текущая'}\n" + f"LoRA: {data.get('lora', 'нет') or 'нет'} ({data.get('lora_strength', 0.8)})\n" + f"Негативный промпт: {negative_prompt or 'нет'}", + parse_mode="HTML", + reply_markup=keyboard, + ) + await state.clear() + + +@router.callback_query(F.data.startswith("profile_view_")) +async def profile_view(callback: CallbackQuery): + """Просмотр профиля.""" + profile_id = int(callback.data.split("_")[-1]) + profile = await db.get_profile(profile_id) + + if not profile: + await callback.answer("Профиль не найден.", show_alert=True) + return + + text = ( + f"📋 Профиль: {profile['name']}\n\n" + f"{'⭐ По умолчанию' if profile['is_default'] else ''}\n" + f"Размер: {profile['width']}x{profile['height']}\n" + f"Шагов: {profile['steps']}\n" + f"CFG Scale: {profile['cfg_scale']}\n" + f"Сэмплер: {profile['sampler']}\n" + f"Шедулер: {profile.get('scheduler', 'automatic')}\n" + f"Модель: {profile['model'] or 'текущая'}\n" + f"LoRA: {profile['lora'] or 'нет'} ({profile['lora_strength']})\n" + f"Негативный промпт: {profile['negative_prompt'] or 'нет'}" + ) + + keyboard = InlineKeyboardMarkup(inline_keyboard=[ + [ + InlineKeyboardButton(text="⭐ По умолчанию", callback_data=f"profile_set_default_{profile_id}") + ] if not profile['is_default'] else [ + InlineKeyboardButton(text="⭐ Профиль по умолчанию", callback_data="noop") + ], + [ + InlineKeyboardButton(text="🗑️ Удалить", callback_data=f"profile_delete_{profile_id}"), + ], + [ + InlineKeyboardButton(text="⬅️ Назад", callback_data="profiles_menu"), + ], + ]) + + await callback.message.edit_text(text, parse_mode="HTML", reply_markup=keyboard) + await callback.answer() + + +@router.callback_query(F.data.startswith("profile_set_default_")) +async def profile_set_default(callback: CallbackQuery): + """Установка профиля по умолчанию.""" + profile_id = int(callback.data.split("_")[-1]) + user_id = callback.from_user.id + + await db.set_default_profile(user_id, profile_id) + await callback.answer("✅ Профиль установлен как профиль по умолчанию.") + await profile_view(callback) + + +@router.callback_query(F.data.startswith("profile_delete_")) +async def profile_delete(callback: CallbackQuery): + """Удаление профиля.""" + profile_id = int(callback.data.split("_")[-1]) + user_id = callback.from_user.id + + success = await db.delete_profile(profile_id, user_id) + if success: + await callback.answer("🗑️ Профиль удалён.") + await profiles_menu(callback) + else: + await callback.answer("❌ Ошибка удаления профиля.", show_alert=True) diff --git a/bot/handlers_settings.py b/bot/handlers_settings.py new file mode 100644 index 0000000..55dfa49 --- /dev/null +++ b/bot/handlers_settings.py @@ -0,0 +1,251 @@ +"""Обработчики настроек хранения и вспомогательные команды.""" + +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"👋 Привет, {callback.from_user.first_name}!\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"✅ SD API подключено!\n\n" + f"Адрес: {settings.SD_API_URL}\n" + f"Текущая модель: {current_model}" + ) + else: + text = ( + f"❌ SD API недоступно!\n\n" + f"Адрес: {settings.SD_API_URL}\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 = ( + "ℹ️ Справка по боту\n\n" + "Генерация изображений:\n" + "• txt2img — создание изображения по текстовому описанию\n\n" + "Профили:\n" + "Создавайте профили с предустановленными настройками:\n" + "• Размер изображения\n" + "• Количество шагов\n" + "• CFG Scale\n" + "• Сэмплер\n" + "• Модель и LoRA\n\n" + "Вы можете установить профиль по умолчанию для быстрой генерации.\n\n" + "Настройки хранения:\n" + "Настройте время хранения сгенерированных изображений.\n" + "Изображения автоматически удаляются по истечении срока.\n\n" + "Команды:\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"✅ SD API подключено!\n\n" + f"Адрес: {settings.SD_API_URL}\n" + f"Текущая модель: {current_model}" + ) + else: + text = ( + f"❌ SD API недоступно!\n\n" + f"Адрес: {settings.SD_API_URL}" + ) + 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 = ( + "🗃️ Настройки хранения изображений\n\n" + f"Текущее время хранения: {current_ttl} часов\n" + f"Максимальное время: {settings.MAX_IMAGE_TTL_HOURS} часов\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"⏰ Введите время хранения в часах\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"✅ Настройки сохранены!\n\n" + f"Время хранения изображений: {ttl_hours} часов\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}") diff --git a/bot/handlers_start.py b/bot/handlers_start.py new file mode 100644 index 0000000..5666aea --- /dev/null +++ b/bot/handlers_start.py @@ -0,0 +1,51 @@ +"""Обработчик команды /start и главного меню.""" + +from aiogram import Router +from aiogram.filters import CommandStart, Command +from aiogram.types import Message, InlineKeyboardMarkup, InlineKeyboardButton + +from config import settings + +router = Router() + + +def get_main_keyboard(user_id: int = None) -> InlineKeyboardMarkup: + """Создать главное меню.""" + keyboard = InlineKeyboardMarkup(inline_keyboard=[ + [ + InlineKeyboardButton(text="🎨 Генерация (txt2img)", callback_data="gen_txt2img"), + ], + [ + InlineKeyboardButton(text="⚙️ Профили", callback_data="profiles_menu"), + InlineKeyboardButton(text="🗑️ Настройки хранения", callback_data="storage_settings"), + ], + [ + InlineKeyboardButton(text="📊 Статус SD API", callback_data="check_sd_status"), + InlineKeyboardButton(text="ℹ️ Помощь", callback_data="help"), + ], + ]) + + # Добавляем кнопку админ-панели для админа + if user_id and user_id == settings.ADMIN_ID: + keyboard.inline_keyboard.insert(0, [ + InlineKeyboardButton(text="🛡️ Админ-панель", callback_data="admin_panel"), + ]) + + return keyboard + + +@router.message(CommandStart()) +async def cmd_start(message: Message): + """Обработка команды /start.""" + text = ( + f"👋 Привет, {message.from_user.first_name}!\n\n" + "Я бот для генерации изображений через Stable Diffusion.\n" + "Выберите действие из меню ниже:" + ) + await message.answer(text, reply_markup=get_main_keyboard(message.from_user.id), parse_mode="HTML") + + +@router.message(Command("menu")) +async def cmd_menu(message: Message): + """Обработка команды /menu.""" + await message.answer("📋 Главное меню:", reply_markup=get_main_keyboard(message.from_user.id)) diff --git a/bot/middleware.py b/bot/middleware.py new file mode 100644 index 0000000..05ea27e --- /dev/null +++ b/bot/middleware.py @@ -0,0 +1,102 @@ +"""Мидлварь для проверки доступа пользователей к боту.""" + +import logging +from aiogram import BaseMiddleware +from aiogram.types import Message, CallbackQuery +from aiogram.filters import Command +from database.database import db + +logger = logging.getLogger(__name__) + + +class AccessCheckMiddleware(BaseMiddleware): + """Мидлварь проверяет, имеет ли пользователь доступ к боту.""" + + # Команды, которые всегда разрешены (без проверки доступа) + ALLOWED_COMMANDS = {"start", "menu"} + + async def __call__(self, handler, event, data): + # Определяем тип события и извлекаем user_id + user_id = None + username = None + display_name = None + + if isinstance(event, Message): + user_id = event.from_user.id + username = event.from_user.username + display_name = event.from_user.full_name + # Проверяем, является ли это разрешённой командой + if isinstance(event, Message) and event.text: + text = event.text.strip().lower() + for cmd in self.ALLOWED_COMMANDS: + # Проверяем /cmd и /cmd@username + if text == f"/{cmd}" or text.startswith(f"/{cmd}@"): + return await handler(event, data) + + elif isinstance(event, CallbackQuery): + user_id = event.from_user.id + username = event.from_user.username + display_name = event.from_user.full_name + + if user_id is None: + return await handler(event, data) + + # Проверяем доступ + access_info = await db.check_user_access(user_id) + + if access_info["has_access"]: + # Сохраняем информацию о пользователе в data для обработчиков + data["user_access"] = access_info + return await handler(event, data) + + # Доступ запрещён — отправляем сообщение + reason = access_info["reason"] + user_info = access_info["user_info"] + + # Формируем сообщение в зависимости от причины + if reason == "not_registered": + text = ( + "🔒 Доступ запрещён\n\n" + "У вас нет доступа к этому боту. " + "Обратитесь к администратору для получения доступа." + ) + elif reason == "deactivated": + text = ( + "🔒 Доступ заблокирован\n\n" + "Ваш доступ к боту был отозван администратором." + ) + elif reason == "expired": + expires_str = "не указан" + if user_info and user_info.get("access_expires_at"): + from datetime import datetime + try: + exp_date = datetime.fromisoformat(user_info["access_expires_at"]) + expires_str = exp_date.strftime("%d.%m.%Y %H:%M") + except Exception: + pass + text = ( + "⏰ Срок доступа истёк\n\n" + f"Ваш доступ к боту истёк ({expires_str}).\n" + "Обратитесь к администратору для продления." + ) + elif reason == "generations_limit": + used = user_info.get("used_generations", "?") if user_info else "?" + max_gens = user_info.get("max_generations", "?") if user_info else "?" + text = ( + "🎨 Лимит генераций исчерпан\n\n" + f"Вы использовали все {max_gens} генераций.\n" + "Обратитесь к администратору для увеличения лимита." + ) + else: + text = "🔒 Доступ запрещён" + + if isinstance(event, Message): + await event.answer(text, parse_mode="HTML") + elif isinstance(event, CallbackQuery): + await event.answer("Доступ запрещён", show_alert=True) + try: + await event.message.answer(text, parse_mode="HTML") + except Exception: + pass + + return None diff --git a/config.py b/config.py new file mode 100644 index 0000000..94ad2ec --- /dev/null +++ b/config.py @@ -0,0 +1,30 @@ +"""Модуль конфигурации для загрузки настроек из переменных окружения.""" + +import os +from dotenv import load_dotenv + +load_dotenv() + + +class Settings: + """Настройки приложения.""" + + BOT_TOKEN: str = os.getenv("BOT_TOKEN", "") + ADMIN_ID: int = int(os.getenv("ADMIN_ID", "0")) + SD_API_URL: str = os.getenv("SD_API_URL", "http://192.168.1.120:7860") + IMAGES_DIR: str = os.getenv("IMAGES_DIR", "/app/images") + DEFAULT_IMAGE_TTL_HOURS: int = int(os.getenv("DEFAULT_IMAGE_TTL_HOURS", "48")) + MAX_IMAGE_TTL_HOURS: int = int(os.getenv("MAX_IMAGE_TTL_HOURS", "168")) + CLEANUP_INTERVAL_MINUTES: int = int(os.getenv("CLEANUP_INTERVAL_MINUTES", "30")) + DB_PATH: str = os.getenv("DB_PATH", "/app/data/bot.db") + + @classmethod + def validate(cls) -> None: + """Проверка обязательных настроек.""" + if not cls.BOT_TOKEN or cls.BOT_TOKEN == "your_telegram_bot_token_here": + raise ValueError( + "BOT_TOKEN не установлен. Получите токен у @BotFather и установите его в .env файле." + ) + + +settings = Settings() diff --git a/database/__init__.py b/database/__init__.py new file mode 100644 index 0000000..8c722d4 --- /dev/null +++ b/database/__init__.py @@ -0,0 +1 @@ +# Модуль базы данных diff --git a/database/database.py b/database/database.py new file mode 100644 index 0000000..dbda307 --- /dev/null +++ b/database/database.py @@ -0,0 +1,497 @@ +"""Модуль работы с базой данных SQLite.""" + +import aiosqlite +import os +from datetime import datetime, timedelta +from typing import Optional, List + +from config import settings + + +class Database: + """Класс для работы с SQLite.""" + + def __init__(self, db_path: str): + self.db_path = db_path + + async def initialize(self) -> None: + """Инициализация базы данных и создание таблиц.""" + os.makedirs(os.path.dirname(self.db_path), exist_ok=True) + async with aiosqlite.connect(self.db_path) as db: + # Таблица авторизованных пользователей бота + await db.execute(""" + CREATE TABLE IF NOT EXISTS bot_users ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + user_id INTEGER NOT NULL UNIQUE, + username TEXT, + display_name TEXT, + added_by INTEGER, + added_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, + access_expires_at TIMESTAMP, + max_generations INTEGER, + used_generations INTEGER NOT NULL DEFAULT 0, + is_active INTEGER NOT NULL DEFAULT 1, + is_admin INTEGER NOT NULL DEFAULT 0 + ) + """) + + await db.execute(""" + CREATE INDEX IF NOT EXISTS idx_bot_users_user_id + ON bot_users(user_id) + """) + + # Таблица профилей + await db.execute(""" + CREATE TABLE IF NOT EXISTS profiles ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + user_id INTEGER NOT NULL, + name TEXT NOT NULL, + is_default INTEGER NOT NULL DEFAULT 0, + width INTEGER NOT NULL DEFAULT 512, + height INTEGER NOT NULL DEFAULT 512, + steps INTEGER NOT NULL DEFAULT 20, + cfg_scale REAL NOT NULL DEFAULT 7.0, + sampler TEXT NOT NULL DEFAULT 'Euler a', + scheduler TEXT NOT NULL DEFAULT 'automatic', + model TEXT NOT NULL DEFAULT '', + lora TEXT NOT NULL DEFAULT '', + lora_strength REAL NOT NULL DEFAULT 0.8, + negative_prompt TEXT NOT NULL DEFAULT '', + created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP + ) + """) + + # Таблица пользовательских настроек + await db.execute(f""" + CREATE TABLE IF NOT EXISTS user_settings ( + user_id INTEGER PRIMARY KEY, + image_ttl_hours INTEGER NOT NULL DEFAULT {settings.DEFAULT_IMAGE_TTL_HOURS}, + active_profile_id INTEGER, + FOREIGN KEY (active_profile_id) REFERENCES profiles(id) + ) + """) + + # Таблица сгенерированных изображений + await db.execute(""" + CREATE TABLE IF NOT EXISTS generated_images ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + user_id INTEGER NOT NULL, + file_path TEXT NOT NULL, + prompt TEXT NOT NULL, + created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, + expires_at TIMESTAMP NOT NULL + ) + """) + + await db.execute(""" + CREATE INDEX IF NOT EXISTS idx_images_expires + ON generated_images(expires_at) + """) + + await db.execute(""" + CREATE INDEX IF NOT EXISTS idx_profiles_user + ON profiles(user_id) + """) + + await db.commit() + + async def get_user_settings(self, user_id: int) -> dict: + """Получить настройки пользователя.""" + async with aiosqlite.connect(self.db_path) as db: + async with db.execute( + "SELECT image_ttl_hours, active_profile_id FROM user_settings WHERE user_id = ?", + (user_id,) + ) as cursor: + row = await cursor.fetchone() + if row: + return { + "image_ttl_hours": row[0], + "active_profile_id": row[1] + } + return { + "image_ttl_hours": settings.DEFAULT_IMAGE_TTL_HOURS, + "active_profile_id": None + } + + async def update_user_settings(self, user_id: int, **kwargs) -> None: + """Обновить настройки пользователя.""" + async with aiosqlite.connect(self.db_path) as db: + # Сначала проверяем, есть ли уже запись + async with db.execute( + "SELECT user_id FROM user_settings WHERE user_id = ?", (user_id,) + ) as cursor: + exists = await cursor.fetchone() + + if exists: + set_clause = ", ".join(f"{k} = ?" for k in kwargs.keys()) + values = list(kwargs.values()) + [user_id] + await db.execute( + f"UPDATE user_settings SET {set_clause} WHERE user_id = ?", + values + ) + else: + # Создаём с значениями по умолчанию + новые + current_settings = await self.get_user_settings(user_id) + current_settings.update(kwargs) + await db.execute( + "INSERT INTO user_settings (user_id, image_ttl_hours, active_profile_id) VALUES (?, ?, ?)", + (user_id, current_settings["image_ttl_hours"], current_settings.get("active_profile_id")) + ) + await db.commit() + + async def create_profile( + self, + user_id: int, + name: str, + width: int = 512, + height: int = 512, + steps: int = 20, + cfg_scale: float = 7.0, + sampler: str = "Euler a", + scheduler: str = "automatic", + model: str = "", + lora: str = "", + lora_strength: float = 0.8, + negative_prompt: str = "", + is_default: bool = False + ) -> int: + """Создать новый профиль.""" + async with aiosqlite.connect(self.db_path) as db: + if is_default: + # Сбрасываем флаг default у всех профилей пользователя + await db.execute( + "UPDATE profiles SET is_default = 0 WHERE user_id = ?", (user_id,) + ) + + cursor = await db.execute(""" + INSERT INTO profiles ( + user_id, name, is_default, width, height, steps, + cfg_scale, sampler, scheduler, model, lora, + lora_strength, negative_prompt + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + """, ( + user_id, name, int(is_default), width, height, steps, + cfg_scale, sampler, scheduler, model, lora, + lora_strength, negative_prompt + )) + profile_id = cursor.lastrowid + await db.commit() + return profile_id + + async def get_profile(self, profile_id: int) -> Optional[dict]: + """Получить профиль по ID.""" + async with aiosqlite.connect(self.db_path) as db: + async with db.execute( + "SELECT * FROM profiles WHERE id = ?", (profile_id,) + ) as cursor: + row = await cursor.fetchone() + if not row: + return None + columns = [description[0] for description in cursor.description] + return dict(zip(columns, row)) + + async def get_user_profiles(self, user_id: int) -> list[dict]: + """Получить все профили пользователя.""" + async with aiosqlite.connect(self.db_path) as db: + async with db.execute( + "SELECT * FROM profiles WHERE user_id = ? ORDER BY is_default DESC, name", + (user_id,) + ) as cursor: + rows = await cursor.fetchall() + columns = [description[0] for description in cursor.description] + return [dict(zip(columns, row)) for row in rows] + + async def get_default_profile(self, user_id: int) -> Optional[dict]: + """Получить профиль по умолчанию.""" + async with aiosqlite.connect(self.db_path) as db: + async with db.execute( + "SELECT * FROM profiles WHERE user_id = ? AND is_default = 1", + (user_id,) + ) as cursor: + row = await cursor.fetchone() + if not row: + return None + columns = [description[0] for description in cursor.description] + return dict(zip(columns, row)) + + async def set_default_profile(self, user_id: int, profile_id: int) -> None: + """Установить профиль по умолчанию.""" + async with aiosqlite.connect(self.db_path) as db: + await db.execute( + "UPDATE profiles SET is_default = 0 WHERE user_id = ?", (user_id,) + ) + await db.execute( + "UPDATE profiles SET is_default = 1 WHERE id = ? AND user_id = ?", + (profile_id, user_id) + ) + await db.commit() + + async def update_profile(self, profile_id: int, user_id: int, **kwargs) -> None: + """Обновить профиль.""" + if not kwargs: + return + set_clause = ", ".join(f"{k} = ?" for k in kwargs.keys()) + values = list(kwargs.values()) + [profile_id, user_id] + async with aiosqlite.connect(self.db_path) as db: + await db.execute( + f"UPDATE profiles SET {set_clause} WHERE id = ? AND user_id = ?", + values + ) + await db.commit() + + async def delete_profile(self, profile_id: int, user_id: int) -> bool: + """Удалить профиль.""" + async with aiosqlite.connect(self.db_path) as db: + cursor = await db.execute( + "DELETE FROM profiles WHERE id = ? AND user_id = ?", + (profile_id, user_id) + ) + await db.commit() + return cursor.rowcount > 0 + + async def add_generated_image( + self, + user_id: int, + file_path: str, + prompt: str, + ttl_hours: int + ) -> int: + """Добавить запись о сгенерированном изображении.""" + expires_at = datetime.now() + timedelta(hours=ttl_hours) + async with aiosqlite.connect(self.db_path) as db: + cursor = await db.execute( + "INSERT INTO generated_images (user_id, file_path, prompt, expires_at) VALUES (?, ?, ?, ?)", + (user_id, file_path, prompt, expires_at.isoformat()) + ) + image_id = cursor.lastrowid + await db.commit() + return image_id + + async def get_expired_images(self) -> list[dict]: + """Получить список просроченных изображений.""" + now = datetime.now().isoformat() + async with aiosqlite.connect(self.db_path) as db: + async with db.execute( + "SELECT id, file_path FROM generated_images WHERE expires_at < ?", + (now,) + ) as cursor: + rows = await cursor.fetchall() + return [{"id": row[0], "file_path": row[1]} for row in rows] + + async def delete_image_record(self, image_id: int) -> None: + """Удалить запись об изображении.""" + async with aiosqlite.connect(self.db_path) as db: + await db.execute( + "DELETE FROM generated_images WHERE id = ?", (image_id,) + ) + await db.commit() + + # --- Методы для работы с авторизованными пользователями --- + + async def ensure_admin(self, admin_user_id: int, username: str = None, display_name: str = None) -> None: + """Убедиться, что админ существует в таблице bot_users. Если нет — создать.""" + async with aiosqlite.connect(self.db_path) as db: + async with db.execute( + "SELECT id FROM bot_users WHERE user_id = ?", (admin_user_id,) + ) as cursor: + exists = await cursor.fetchone() + if not exists: + await db.execute(""" + INSERT INTO bot_users (user_id, username, display_name, is_admin, is_active) + VALUES (?, ?, ?, 1, 1) + """, (admin_user_id, username, display_name)) + else: + await db.execute( + "UPDATE bot_users SET is_admin = 1 WHERE user_id = ?", (admin_user_id,) + ) + await db.commit() + + async def add_user( + self, + user_id: int, + added_by: int, + username: str = None, + display_name: str = None, + access_expires_at: Optional[datetime] = None, + max_generations: Optional[int] = None, + ) -> bool: + """Добавить нового пользователя. Возвращает True если успешно.""" + try: + async with aiosqlite.connect(self.db_path) as db: + await db.execute(""" + INSERT INTO bot_users ( + user_id, username, display_name, added_by, + access_expires_at, max_generations, used_generations, is_active + ) VALUES (?, ?, ?, ?, ?, ?, 0, 1) + """, ( + user_id, username, display_name, added_by, + access_expires_at.isoformat() if access_expires_at else None, + max_generations, + )) + await db.commit() + return True + except aiosqlite.IntegrityError: + # Пользователь уже существует + async with aiosqlite.connect(self.db_path) as db: + await db.execute(""" + UPDATE bot_users + SET added_by = ?, access_expires_at = ?, max_generations = ?, + used_generations = 0, is_active = 1 + WHERE user_id = ? + """, ( + added_by, + access_expires_at.isoformat() if access_expires_at else None, + max_generations, + user_id, + )) + await db.commit() + return True + + async def remove_user(self, user_id: int) -> bool: + """Удалить пользователя (деактивировать).""" + async with aiosqlite.connect(self.db_path) as db: + await db.execute( + "UPDATE bot_users SET is_active = 0 WHERE user_id = ? AND is_admin = 0", + (user_id,) + ) + await db.commit() + return True + + async def check_user_access(self, user_id: int) -> dict: + """ + Проверить доступ пользователя. + Возвращает dict с полями: + - has_access: bool + - is_admin: bool + - reason: str (причина отказа, если нет доступа) + - user_info: dict (информация о пользователе) + """ + async with aiosqlite.connect(self.db_path) as db: + async with db.execute( + "SELECT * FROM bot_users WHERE user_id = ?", (user_id,) + ) as cursor: + row = await cursor.fetchone() + if not row: + return { + "has_access": False, + "is_admin": False, + "reason": "not_registered", + "user_info": None, + } + + columns = [description[0] for description in cursor.description] + user_info = dict(zip(columns, row)) + + # Админ всегда имеет доступ + if user_info["is_admin"]: + return { + "has_access": True, + "is_admin": True, + "reason": None, + "user_info": user_info, + } + + # Проверяем активность + if not user_info["is_active"]: + return { + "has_access": False, + "is_admin": False, + "reason": "deactivated", + "user_info": user_info, + } + + # Проверяем срок доступа + if user_info["access_expires_at"]: + expires_at = datetime.fromisoformat(user_info["access_expires_at"]) + if datetime.now() > expires_at: + return { + "has_access": False, + "is_admin": False, + "reason": "expired", + "user_info": user_info, + } + + # Проверяем лимит генераций + if user_info["max_generations"] is not None: + if user_info["used_generations"] >= user_info["max_generations"]: + return { + "has_access": False, + "is_admin": False, + "reason": "generations_limit", + "user_info": user_info, + } + + return { + "has_access": True, + "is_admin": False, + "reason": None, + "user_info": user_info, + } + + async def increment_generation_count(self, user_id: int) -> None: + """Увеличить счётчик использованных генераций.""" + async with aiosqlite.connect(self.db_path) as db: + await db.execute( + "UPDATE bot_users SET used_generations = used_generations + 1 WHERE user_id = ?", + (user_id,) + ) + await db.commit() + + async def get_all_users(self) -> List[dict]: + """Получить список всех пользователей.""" + async with aiosqlite.connect(self.db_path) as db: + async with db.execute( + "SELECT * FROM bot_users ORDER BY is_admin DESC, added_at DESC" + ) as cursor: + rows = await cursor.fetchall() + columns = [description[0] for description in cursor.description] + return [dict(zip(columns, row)) for row in rows] + + async def get_user_by_id(self, user_id: int) -> Optional[dict]: + """Получить информацию о пользователе по user_id.""" + async with aiosqlite.connect(self.db_path) as db: + async with db.execute( + "SELECT * FROM bot_users WHERE user_id = ?", (user_id,) + ) as cursor: + row = await cursor.fetchone() + if not row: + return None + columns = [description[0] for description in cursor.description] + return dict(zip(columns, row)) + + async def update_user_access( + self, + user_id: int, + access_expires_at: Optional[datetime] = None, + max_generations: Optional[int] = None, + ) -> None: + """Обновить параметры доступа пользователя.""" + async with aiosqlite.connect(self.db_path) as db: + await db.execute(""" + UPDATE bot_users + SET access_expires_at = ?, max_generations = ?, is_active = 1 + WHERE user_id = ? + """, ( + access_expires_at.isoformat() if access_expires_at else None, + max_generations, + user_id, + )) + await db.commit() + + async def deactivate_expired_users(self) -> int: + """Деактивировать пользователей с истёкшим сроком доступа. Возвращает количество.""" + now = datetime.now().isoformat() + async with aiosqlite.connect(self.db_path) as db: + cursor = await db.execute(""" + UPDATE bot_users SET is_active = 0 + WHERE access_expires_at IS NOT NULL + AND access_expires_at < ? + AND is_active = 1 + AND is_admin = 0 + """, (now,)) + await db.commit() + return cursor.rowcount + + +# Глобальный экземпляр +db = Database(settings.DB_PATH) diff --git a/docker-compose.yml b/docker-compose.yml new file mode 100644 index 0000000..afcf20e --- /dev/null +++ b/docker-compose.yml @@ -0,0 +1,30 @@ +version: '3.8' + +services: + tg-sd-bot: + build: + context: . + dockerfile: Dockerfile + container_name: tg-sd-bot + restart: unless-stopped + env_file: + - .env + volumes: + # Хранение сгенерированных изображений + - ./images:/app/images + # Хранение базы данных + - ./data:/app/data + # Сеть для доступа к Stable Diffusion API в локальной сети + # Если SD API на другом сервере, дополнительная настройка сети не требуется + networks: + - bot_network + # Ограничение ресурсов (опционально) + deploy: + resources: + limits: + memory: 512M + cpus: '1.0' + +networks: + bot_network: + driver: bridge diff --git a/main.py b/main.py new file mode 100644 index 0000000..e2e389e --- /dev/null +++ b/main.py @@ -0,0 +1,100 @@ +"""Главный файл запуска Telegram бота.""" + +import asyncio +import logging +import os + +from aiogram import Bot, Dispatcher +from aiogram.client.default import DefaultBotProperties +from aiogram.enums import ParseMode + +from config import settings +from database.database import db +from bot.handlers_start import router as start_router +from bot.handlers_generation import router as generation_router +from bot.handlers_profiles import router as profiles_router +from bot.handlers_settings import router as settings_router, scheduled_cleanup +from bot.handlers_admin import router as admin_router +from bot.middleware import AccessCheckMiddleware + + +async def main(): + """Основная функция запуска бота.""" + # Настройка логирования + logging.basicConfig( + level=logging.INFO, + format='%(asctime)s - %(name)s - %(levelname)s - %(message)s', + ) + logger = logging.getLogger(__name__) + + # Проверяем настройки + try: + settings.validate() + except ValueError as e: + logger.error(f"Ошибка конфигурации: {e}") + return + + # Инициализация базы данных + await db.initialize() + logger.info("База данных инициализирована") + + # Инициализация админа + if settings.ADMIN_ID: + await db.ensure_admin(settings.ADMIN_ID) + logger.info(f"Админ инициализирован: {settings.ADMIN_ID}") + else: + logger.warning("ADMIN_ID не установлен — админ-панель будет недоступна") + + # Создаём директорию для изображений + os.makedirs(settings.IMAGES_DIR, exist_ok=True) + logger.info(f"Директория для изображений: {settings.IMAGES_DIR}") + + # Инициализация бота и диспетчера + bot = Bot( + token=settings.BOT_TOKEN, + default=DefaultBotProperties(parse_mode=ParseMode.HTML) + ) + dp = Dispatcher() + + # Регистрируем мидлварь проверки доступа + dp.message.middleware(AccessCheckMiddleware()) + dp.callback_query.middleware(AccessCheckMiddleware()) + logger.info("Мидлварь проверки доступа зарегистрирована") + + # Регистрация роутеров (админ первым, чтобы его обработчики имели приоритет) + dp.include_router(admin_router) + dp.include_router(start_router) + dp.include_router(generation_router) + dp.include_router(profiles_router) + dp.include_router(settings_router) + + # Запуск фоновой задачи очистки + async def cleanup_task(): + """Фоновая задача очистки просроченных изображений.""" + while True: + await asyncio.sleep(settings.CLEANUP_INTERVAL_MINUTES * 60) + await scheduled_cleanup() + + cleanup_task_handle = asyncio.create_task(cleanup_task()) + logger.info( + f"Запущена фоновая очистка изображений каждые {settings.CLEANUP_INTERVAL_MINUTES} минут" + ) + + try: + logger.info("Бот запущен!") + await dp.start_polling(bot) + except Exception as e: + logger.error(f"Ошибка работы бота: {e}") + finally: + cleanup_task_handle.cancel() + await bot.session.close() + logger.info("Бот остановлен") + + +if __name__ == "__main__": + try: + asyncio.run(main()) + except KeyboardInterrupt: + print("\nБот остановлен пользователем") + except Exception as e: + print(f"Критическая ошибка: {e}") diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..0d55e8f --- /dev/null +++ b/requirements.txt @@ -0,0 +1,5 @@ +aiogram==3.17.0 +aiohttp==3.11.12 +Pillow==11.1.0 +aiosqlite==0.21.0 +python-dotenv==1.0.1 diff --git a/sd/__init__.py b/sd/__init__.py new file mode 100644 index 0000000..c0a2b8f --- /dev/null +++ b/sd/__init__.py @@ -0,0 +1 @@ +# Интеграция со Stable Diffusion API diff --git a/sd/sd_client.py b/sd/sd_client.py new file mode 100644 index 0000000..b87392a --- /dev/null +++ b/sd/sd_client.py @@ -0,0 +1,341 @@ +"""Модуль взаимодействия с Stable Diffusion WebUI API (Automatic1111).""" + +import base64 +import logging +from typing import Optional + +import aiohttp + +from config import settings + +logger = logging.getLogger(__name__) + + +class SDClient: + """Клиент для работы с Stable Diffusion WebUI API.""" + + def __init__(self, api_url: str): + self.api_url = api_url.rstrip("/") + + async def check_connection(self) -> bool: + """Проверка соединения с API.""" + try: + async with aiohttp.ClientSession() as session: + async with session.get( + f"{self.api_url}/sdapi/v1/options", + timeout=aiohttp.ClientTimeout(total=10) + ) as response: + return response.status == 200 + except Exception as e: + logger.error(f"Ошибка подключения к SD API: {e}") + return False + + async def get_models(self) -> list[str]: + """Получить список доступных моделей.""" + try: + async with aiohttp.ClientSession() as session: + async with session.get( + f"{self.api_url}/sdapi/v1/sd-models", + timeout=aiohttp.ClientTimeout(total=30) + ) as response: + if response.status == 200: + data = await response.json() + return [model.get("title", model.get("model_name", "")) for model in data] + except Exception as e: + logger.error(f"Ошибка получения списка моделей: {e}") + return [] + + async def get_current_model(self) -> str: + """Получить текущую загруженную модель.""" + try: + async with aiohttp.ClientSession() as session: + async with session.get( + f"{self.api_url}/sdapi/v1/options", + timeout=aiohttp.ClientTimeout(total=10) + ) as response: + if response.status == 200: + data = await response.json() + return data.get("sd_model_checkpoint", "Unknown") + except Exception as e: + logger.error(f"Ошибка получения текущей модели: {e}") + return "Unknown" + + async def get_samplers(self) -> list[str]: + """Получить список доступных сэмплеров.""" + try: + async with aiohttp.ClientSession() as session: + async with session.get( + f"{self.api_url}/sdapi/v1/samplers", + timeout=aiohttp.ClientTimeout(total=10) + ) as response: + if response.status == 200: + data = await response.json() + return [sampler.get("name", "") for sampler in data] + except Exception as e: + logger.error(f"Ошибка получения списка сэмплеров: {e}") + return ["Euler a", "Euler", "DPM++ 2M Karras", "DPM++ SDE Karras", "DDIM"] + + async def get_schedulers(self) -> list[str]: + """Получить список доступных шедулеров.""" + try: + async with aiohttp.ClientSession() as session: + async with session.get( + f"{self.api_url}/sdapi/v1/schedulers", + timeout=aiohttp.ClientTimeout(total=10) + ) as response: + if response.status == 200: + data = await response.json() + return [scheduler.get("name", "") for scheduler in data] + except Exception as e: + logger.error(f"Ошибка получения списка шедулеров: {e}") + return ["automatic", "normal", "karras", "exponential", "SGM uniform", "simple", "DDIM"] + + async def get_loras(self) -> list[str]: + """Получить список доступных LoRA.""" + try: + async with aiohttp.ClientSession() as session: + async with session.get( + f"{self.api_url}/sdapi/v1/loras", + timeout=aiohttp.ClientTimeout(total=10) + ) as response: + if response.status == 200: + data = await response.json() + return [lora.get("name", "") for lora in data] + except Exception as e: + logger.error(f"Ошибка получения списка LoRA: {e}") + return [] + + async def txt2img( + self, + prompt: str, + negative_prompt: str = "", + width: int = 512, + height: int = 512, + steps: int = 20, + cfg_scale: float = 7.0, + sampler: str = "Euler a", + scheduler: str = "automatic", + seed: int = -1, + model: Optional[str] = None, + lora: Optional[str] = None, + lora_strength: float = 0.8, + ) -> Optional[tuple[bytes, dict]]: + """ + Генерация изображения из текста. + Возвращает кортеж (изображение в bytes, info_dict) или None при ошибке. + """ + payload = { + "prompt": prompt, + "negative_prompt": negative_prompt, + "steps": steps, + "width": width, + "height": height, + "cfg_scale": cfg_scale, + "sampler_name": sampler, + "scheduler": scheduler, + "seed": seed, + "save_images": False, + "send_images": True, + } + + # Если указана модель, переключаем её + if model: + await self._switch_model(model) + + # Формируем промпт с LoRA + final_prompt = prompt + if lora: + final_prompt = f" {prompt}" + payload["prompt"] = final_prompt + + try: + async with aiohttp.ClientSession() as session: + async with session.post( + f"{self.api_url}/sdapi/v1/txt2img", + json=payload, + timeout=aiohttp.ClientTimeout(total=600) # 10 минут на генерацию + ) as response: + if response.status == 200: + data = await response.json() + if data.get("images"): + image_bytes = base64.b64decode(data["images"][0]) + info = { + "prompt": final_prompt, + "negative_prompt": negative_prompt, + "width": width, + "height": height, + "steps": steps, + "cfg_scale": cfg_scale, + "sampler": sampler, + "scheduler": scheduler, + "seed": data.get("parameters", {}).get("seed", seed), + "model": model or await self.get_current_model(), + "lora": lora, + } + return image_bytes, info + else: + logger.error("API не вернул изображение") + return None + else: + error_text = await response.text() + logger.error(f"Ошибка API txt2img: {response.status} - {error_text}") + return None + except aiohttp.ClientTimeout: + logger.error("Таймаут запроса txt2img (10 минут)") + return None + except Exception as e: + logger.error(f"Ошибка txt2img: {e}") + return None + + async def img2img( + self, + init_image_bytes: bytes, + prompt: str, + negative_prompt: str = "", + width: int = 512, + height: int = 512, + steps: int = 20, + cfg_scale: float = 7.0, + sampler: str = "Euler a", + scheduler: str = "automatic", + denoising_strength: float = 0.75, + seed: int = -1, + model: Optional[str] = None, + lora: Optional[str] = None, + lora_strength: float = 0.8, + ) -> Optional[tuple[bytes, dict]]: + """ + Генерация изображения на основе изображения. + Возвращает кортеж (изображение в bytes, info_dict) или None при ошибке. + """ + init_image_base64 = base64.b64encode(init_image_bytes).decode("utf-8") + + payload = { + "init_images": [init_image_base64], + "prompt": prompt, + "negative_prompt": negative_prompt, + "steps": steps, + "width": width, + "height": height, + "cfg_scale": cfg_scale, + "sampler_name": sampler, + "scheduler": scheduler, + "denoising_strength": denoising_strength, + "seed": seed, + "save_images": False, + "send_images": True, + } + + if model: + await self._switch_model(model) + + final_prompt = prompt + if lora: + final_prompt = f" {prompt}" + payload["prompt"] = final_prompt + + try: + async with aiohttp.ClientSession() as session: + async with session.post( + f"{self.api_url}/sdapi/v1/img2img", + json=payload, + timeout=aiohttp.ClientTimeout(total=600) + ) as response: + if response.status == 200: + data = await response.json() + if data.get("images"): + image_bytes = base64.b64decode(data["images"][0]) + info = { + "prompt": final_prompt, + "negative_prompt": negative_prompt, + "width": width, + "height": height, + "steps": steps, + "cfg_scale": cfg_scale, + "sampler": sampler, + "scheduler": scheduler, + "seed": data.get("parameters", {}).get("seed", seed), + "denoising_strength": denoising_strength, + "model": model or await self.get_current_model(), + "lora": lora, + } + return image_bytes, info + else: + logger.error("API не вернул изображение") + return None + else: + error_text = await response.text() + logger.error(f"Ошибка API img2img: {response.status} - {error_text}") + return None + except aiohttp.ClientTimeout: + logger.error("Таймаут запроса img2img (10 минут)") + return None + except Exception as e: + logger.error(f"Ошибка img2img: {e}") + return None + + async def _switch_model(self, model_name: str) -> bool: + """Переключить модель.""" + try: + async with aiohttp.ClientSession() as session: + async with session.post( + f"{self.api_url}/sdapi/v1/options", + json={"sd_model_checkpoint": model_name}, + timeout=aiohttp.ClientTimeout(total=120) + ) as response: + if response.status == 200: + logger.info(f"Модель переключена на: {model_name}") + return True + else: + logger.error(f"Ошибка переключения модели: {response.status}") + return False + except Exception as e: + logger.error(f"Ошибка переключения модели: {e}") + return False + + async def get_progress(self, skip_headers: bool = False) -> Optional[dict]: + """Получить текущий прогресс генерации.""" + try: + async with aiohttp.ClientSession() as session: + async with session.get( + f"{self.api_url}/sdapi/v1/progress", + params={"skip_current_image": skip_headers}, + timeout=aiohttp.ClientTimeout(total=10) + ) as response: + if response.status == 200: + return await response.json() + except Exception as e: + logger.error(f"Ошибка получения прогресса: {e}") + return None + + async def interrupt(self) -> bool: + """Прервать текущую генерацию.""" + try: + async with aiohttp.ClientSession() as session: + async with session.post( + f"{self.api_url}/sdapi/v1/interrupt", + json={}, + timeout=aiohttp.ClientTimeout(total=10) + ) as response: + return response.status == 200 + except Exception as e: + logger.error(f"Ошибка прерывания генерации: {e}") + return False + + async def get_options(self) -> dict: + """Получить текущие опции API.""" + try: + async with aiohttp.ClientSession() as session: + async with session.get( + f"{self.api_url}/sdapi/v1/options", + timeout=aiohttp.ClientTimeout(total=10) + ) as response: + if response.status == 200: + return await response.json() + except Exception as e: + logger.error(f"Ошибка получения опций: {e}") + return {} + + +# Глобальный экземпляр +sd_client = SDClient(settings.SD_API_URL) diff --git a/utils/__init__.py b/utils/__init__.py new file mode 100644 index 0000000..417a481 --- /dev/null +++ b/utils/__init__.py @@ -0,0 +1 @@ +# Утилиты diff --git a/utils/image_manager.py b/utils/image_manager.py new file mode 100644 index 0000000..63a5701 --- /dev/null +++ b/utils/image_manager.py @@ -0,0 +1,60 @@ +"""Модуль управления изображениями и их очистки.""" + +import logging +import os +from datetime import datetime + +from config import settings +from database.database import db + +logger = logging.getLogger(__name__) + + +async def cleanup_expired_images() -> int: + """ + Удалить просроченные изображения. + Возвращает количество удалённых изображений. + """ + expired = await db.get_expired_images() + deleted_count = 0 + + for image in expired: + # Удаляем файл + file_path = image["file_path"] + try: + if os.path.exists(file_path): + os.remove(file_path) + logger.debug(f"Удалён файл: {file_path}") + except Exception as e: + logger.error(f"Ошибка удаления файла {file_path}: {e}") + + # Удаляем запись из БД + try: + await db.delete_image_record(image["id"]) + deleted_count += 1 + except Exception as e: + logger.error(f"Ошибка удаления записи изображения {image['id']}: {e}") + + if deleted_count > 0: + logger.info(f"Удалено {deleted_count} просроченных изображдений") + + return deleted_count + + +def get_image_path(user_id: int, filename: str) -> str: + """ + Получить путь для сохранения изображения. + Формат: images/{user_id}/{filename} + """ + user_dir = os.path.join(settings.IMAGES_DIR, str(user_id)) + os.makedirs(user_dir, exist_ok=True) + return os.path.join(user_dir, filename) + + +def cleanup_user_images(user_id: int) -> None: + """Удалить все изображения пользователя (при необходимости).""" + user_dir = os.path.join(settings.IMAGES_DIR, str(user_id)) + if os.path.exists(user_dir): + import shutil + shutil.rmtree(user_dir) + logger.info(f"Удалены все изображения пользователя {user_id}")