599 lines
22 KiB
Python
599 lines
22 KiB
Python
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import logging
|
||
import time
|
||
import uuid
|
||
from pathlib import Path
|
||
|
||
from telegram import InputFile, Update
|
||
from telegram.constants import ChatAction, ChatType
|
||
from telegram.error import BadRequest, NetworkError, TelegramError, TimedOut
|
||
from telegram.ext import (
|
||
Application,
|
||
CommandHandler,
|
||
ContextTypes,
|
||
MessageHandler,
|
||
filters,
|
||
)
|
||
|
||
from comfy_client import ComfyClient
|
||
from config import Settings, setup_logging
|
||
from db import Database, Request, UserAlreadyQueued
|
||
from moderation import PromptBlocked, detect_person
|
||
from prompt_enhancer import PromptEnhancer
|
||
from queue_manager import QueueEvent, QueueFull, QueueManager
|
||
from workflows import WorkflowBuilder
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
START_TEXT = (
|
||
"Пришли текст — сгенерирую картинку.\n"
|
||
"Пришли фото с подписью — отредактирую его по описанию.\n"
|
||
"Или ответь текстом на готовую картинку — возьму её как исходник.\n\n"
|
||
"/status — узнать состояние текущей задачи."
|
||
)
|
||
BLOCKED_TEXT = "Запрос нарушает правила. Звёзды не списаны."
|
||
|
||
|
||
class MpakoBot:
|
||
def __init__(self, settings: Settings) -> None:
|
||
self._settings = settings
|
||
self._db = Database(settings.db_path)
|
||
self._workflows: WorkflowBuilder | None = None
|
||
self._comfy: ComfyClient | None = None
|
||
self._queue: QueueManager | None = None
|
||
self._bot = None
|
||
self._status_messages: dict[int, tuple[int, int]] = {}
|
||
self._last_edit: dict[int, tuple[float, int]] = {}
|
||
|
||
def run(self) -> None:
|
||
builder = (
|
||
Application.builder()
|
||
.token(self._settings.bot_token)
|
||
.connect_timeout(20)
|
||
.read_timeout(60)
|
||
.write_timeout(60)
|
||
.get_updates_connect_timeout(20)
|
||
.get_updates_read_timeout(60)
|
||
.post_init(self._post_init)
|
||
.post_shutdown(self._post_shutdown)
|
||
)
|
||
if self._settings.telegram_proxy:
|
||
builder = builder.proxy(self._settings.telegram_proxy).get_updates_proxy(
|
||
self._settings.telegram_proxy
|
||
)
|
||
application = builder.build()
|
||
application.add_handler(CommandHandler("start", self._start))
|
||
application.add_handler(CommandHandler("help", self._help))
|
||
application.add_handler(CommandHandler("status", self._status))
|
||
application.add_handler(MessageHandler(filters.PHOTO, self._photo))
|
||
application.add_handler(
|
||
MessageHandler(filters.TEXT & ~filters.COMMAND, self._text)
|
||
)
|
||
application.add_error_handler(self._on_error)
|
||
application.run_polling(allowed_updates=Update.ALL_TYPES)
|
||
|
||
async def _post_init(self, application: Application) -> None:
|
||
self._bot = application.bot
|
||
self._settings.storage_dir.mkdir(parents=True, exist_ok=True)
|
||
self._workflows = WorkflowBuilder(
|
||
self._settings.workflow_t2i,
|
||
self._settings.workflow_i2i,
|
||
self._settings.workflow_mapping,
|
||
)
|
||
self._comfy = ComfyClient(
|
||
self._settings.comfy_base_url,
|
||
self._settings.comfy_ws_url,
|
||
)
|
||
await self._db.connect()
|
||
enhancer = (
|
||
PromptEnhancer(self._comfy, self._settings)
|
||
if self._settings.prompt_enhance_enabled
|
||
else None
|
||
)
|
||
self._queue = QueueManager(
|
||
db=self._db,
|
||
comfy=self._comfy,
|
||
workflows=self._workflows,
|
||
settings=self._settings,
|
||
on_event=self._on_queue_event,
|
||
enhancer=enhancer,
|
||
)
|
||
await self._queue.start()
|
||
logger.info("Бот запущен, ComfyUI: %s", self._settings.comfy_base_url)
|
||
|
||
async def _post_shutdown(self, application: Application) -> None:
|
||
if self._queue is not None:
|
||
await self._queue.stop()
|
||
if self._comfy is not None:
|
||
await self._comfy.close()
|
||
await self._db.close()
|
||
logger.info("Бот остановлен")
|
||
|
||
async def _start(
|
||
self, update: Update, context: ContextTypes.DEFAULT_TYPE
|
||
) -> None:
|
||
if not await self._check_access(update):
|
||
return
|
||
message = update.effective_message
|
||
if message is not None:
|
||
await message.reply_text(START_TEXT)
|
||
|
||
async def _help(
|
||
self, update: Update, context: ContextTypes.DEFAULT_TYPE
|
||
) -> None:
|
||
if not await self._check_access(update):
|
||
return
|
||
message = update.effective_message
|
||
if message is not None:
|
||
await message.reply_text(START_TEXT)
|
||
|
||
async def _status(
|
||
self, update: Update, context: ContextTypes.DEFAULT_TYPE
|
||
) -> None:
|
||
if not await self._check_access(update):
|
||
return
|
||
message = update.effective_message
|
||
user = update.effective_user
|
||
if message is None or user is None:
|
||
return
|
||
active = await self._db.get_active_by_user(user.id)
|
||
if active is None:
|
||
await message.reply_text("Активных задач нет.")
|
||
return
|
||
if active.status == "running":
|
||
await message.reply_text("Твоя картинка сейчас генерируется.")
|
||
return
|
||
position = await self._queue.position(active.id) if self._queue else 0
|
||
await message.reply_text(
|
||
f"Задача в очереди. Перед тобой: {position} {self._people(position)}."
|
||
)
|
||
|
||
async def _text(
|
||
self, update: Update, context: ContextTypes.DEFAULT_TYPE
|
||
) -> None:
|
||
if not await self._check_access(update):
|
||
return
|
||
message = update.effective_message
|
||
user = update.effective_user
|
||
if message is None or user is None or not message.text:
|
||
return
|
||
|
||
prompt = message.text.strip()
|
||
if not prompt:
|
||
await message.reply_text("Промпт пустой. Напиши, что нарисовать.")
|
||
return
|
||
if len(prompt) > self._settings.max_prompt_length:
|
||
await message.reply_text(
|
||
f"Промпт слишком длинный: максимум {self._settings.max_prompt_length} символов."
|
||
)
|
||
return
|
||
|
||
reply = message.reply_to_message
|
||
attachment = None
|
||
suffix = ".jpg"
|
||
if reply is not None:
|
||
if reply.photo:
|
||
attachment = reply.photo[-1]
|
||
elif reply.document is not None and (
|
||
reply.document.mime_type or ""
|
||
).startswith("image/"):
|
||
attachment = reply.document
|
||
suffix = Path(reply.document.file_name or "").suffix.lower() or ".png"
|
||
|
||
if attachment is None:
|
||
await self._submit(
|
||
update, mode="t2i", prompt=prompt, input_path=None
|
||
)
|
||
return
|
||
|
||
input_path = await self._download_attachment(
|
||
message, attachment, user.id, suffix=suffix
|
||
)
|
||
if input_path is None:
|
||
return
|
||
accepted = await self._submit(
|
||
update,
|
||
mode="i2i",
|
||
prompt=prompt,
|
||
input_path=str(input_path),
|
||
)
|
||
if not accepted:
|
||
input_path.unlink(missing_ok=True)
|
||
|
||
async def _photo(
|
||
self, update: Update, context: ContextTypes.DEFAULT_TYPE
|
||
) -> None:
|
||
if not await self._check_access(update):
|
||
return
|
||
message = update.effective_message
|
||
user = update.effective_user
|
||
if message is None or user is None or not message.photo:
|
||
return
|
||
|
||
caption = (message.caption or "").strip()
|
||
if not caption:
|
||
await message.reply_text(
|
||
"Пришли фото с подписью — что именно изменить на изображении."
|
||
)
|
||
return
|
||
if len(caption) > self._settings.max_prompt_length:
|
||
await message.reply_text(
|
||
f"Описание слишком длинное: максимум {self._settings.max_prompt_length} символов."
|
||
)
|
||
return
|
||
|
||
photo = message.photo[-1]
|
||
input_path = await self._download_attachment(
|
||
message, photo, user.id, suffix=".jpg"
|
||
)
|
||
if input_path is None:
|
||
return
|
||
|
||
accepted = await self._submit(
|
||
update,
|
||
mode="i2i",
|
||
prompt=caption,
|
||
input_path=str(input_path),
|
||
)
|
||
if not accepted:
|
||
input_path.unlink(missing_ok=True)
|
||
|
||
async def _download_attachment(
|
||
self,
|
||
message,
|
||
attachment,
|
||
user_id: int,
|
||
*,
|
||
suffix: str = ".jpg",
|
||
) -> Path | None:
|
||
if attachment.file_size and attachment.file_size > self._settings.max_photo_bytes:
|
||
limit_mb = self._settings.max_photo_bytes / (1024 * 1024)
|
||
try:
|
||
await self._retry(
|
||
message.reply_text,
|
||
f"Файл слишком большой. Максимум {limit_mb:.0f} МБ.",
|
||
)
|
||
except TelegramError:
|
||
logger.debug("Не удалось сообщить о размере файла")
|
||
return None
|
||
|
||
suffix = suffix.lower()
|
||
if not suffix.startswith(".") or len(suffix) > 8:
|
||
suffix = ".png"
|
||
|
||
user_dir = self._settings.storage_dir / str(user_id)
|
||
input_path = user_dir / f"{uuid.uuid4().hex}_input{suffix}"
|
||
try:
|
||
user_dir.mkdir(parents=True, exist_ok=True)
|
||
await message.chat.send_action(ChatAction.TYPING)
|
||
file = await attachment.get_file()
|
||
await file.download_to_drive(custom_path=input_path)
|
||
except (TelegramError, OSError):
|
||
logger.exception("Не удалось скачать вложение пользователя %s", user_id)
|
||
input_path.unlink(missing_ok=True)
|
||
try:
|
||
await self._retry(
|
||
message.reply_text,
|
||
"Не удалось скачать файл, попробуй ещё раз.",
|
||
)
|
||
except TelegramError:
|
||
pass
|
||
return None
|
||
return input_path
|
||
|
||
async def _submit(
|
||
self,
|
||
update: Update,
|
||
*,
|
||
mode: str,
|
||
prompt: str,
|
||
input_path: str | None,
|
||
) -> bool:
|
||
message = update.effective_message
|
||
user = update.effective_user
|
||
if message is None or user is None or self._queue is None:
|
||
return False
|
||
|
||
is_admin = user.id in self._settings.admin_user_ids
|
||
has_person = False
|
||
if (
|
||
mode == "i2i"
|
||
and input_path
|
||
and not is_admin
|
||
and self._settings.moderation_enabled
|
||
):
|
||
has_person = await asyncio.to_thread(detect_person, input_path)
|
||
|
||
try:
|
||
request, position = await self._queue.enqueue(
|
||
user_id=user.id,
|
||
username=user.username,
|
||
mode=mode,
|
||
prompt=prompt,
|
||
input_file_path=input_path,
|
||
has_person=has_person,
|
||
is_admin=is_admin,
|
||
)
|
||
except PromptBlocked:
|
||
await message.reply_text(BLOCKED_TEXT)
|
||
return False
|
||
except UserAlreadyQueued:
|
||
active = await self._db.get_active_by_user(user.id)
|
||
text = "Ты уже в очереди."
|
||
if active is not None:
|
||
if active.status == "running":
|
||
text += " Твоя картинка уже генерируется."
|
||
else:
|
||
ahead = await self._queue.position(active.id)
|
||
text += f" Перед тобой: {ahead} {self._people(ahead)}."
|
||
await message.reply_text(text)
|
||
return False
|
||
except QueueFull as exc:
|
||
await message.reply_text(exc.user_message)
|
||
return False
|
||
except Exception:
|
||
logger.exception("Не удалось поставить задачу пользователю %s", user.id)
|
||
await message.reply_text("Не удалось поставить задачу. Попробуй позже.")
|
||
return False
|
||
|
||
try:
|
||
status_message = await self._retry(
|
||
message.reply_text, self._ack_text(position)
|
||
)
|
||
except TelegramError:
|
||
logger.exception(
|
||
"Не удалось отправить подтверждение по запросу %s", request.id
|
||
)
|
||
return True
|
||
self._status_messages[request.id] = (
|
||
status_message.chat_id,
|
||
status_message.message_id,
|
||
)
|
||
|
||
# Воркер мог стартовать раньше, чем мы сохранили message_id.
|
||
current = await self._db.get_request(request.id)
|
||
if current is None:
|
||
return True
|
||
if current.status == "running":
|
||
await self._edit_status(request.id, "Генерация…", percent=0)
|
||
elif current.status == "done":
|
||
await self._clear_status(request.id, delete=True)
|
||
elif current.status == "failed":
|
||
await self._finish_failed(
|
||
current,
|
||
current.error_text or "Не удалось сгенерировать изображение.",
|
||
)
|
||
return True
|
||
|
||
def _ack_text(self, position: int) -> str:
|
||
if position <= 0:
|
||
return "Принято. Ты первый, начинаю…"
|
||
return f"Принято. Перед тобой: {position} {self._people(position)}."
|
||
|
||
@staticmethod
|
||
def _people(count: int) -> str:
|
||
if count % 10 == 1 and count % 100 != 11:
|
||
return "человек"
|
||
if count % 10 in (2, 3, 4) and count % 100 not in (12, 13, 14):
|
||
return "человека"
|
||
return "человек"
|
||
|
||
async def _on_queue_event(self, event: QueueEvent) -> None:
|
||
if event.kind == "started":
|
||
await self._edit_status(event.request.id, "Генерация…", percent=0)
|
||
elif event.kind == "enhancing":
|
||
await self._edit_status(event.request.id, "Улучшаю промпт…")
|
||
elif event.kind == "progress":
|
||
percent = event.percent if event.percent is not None else 0
|
||
await self._edit_status(
|
||
event.request.id,
|
||
f"Генерация… {percent}%",
|
||
percent=percent,
|
||
)
|
||
elif event.kind == "done":
|
||
await self._finish_done(event.request)
|
||
elif event.kind == "failed":
|
||
await self._finish_failed(
|
||
event.request,
|
||
event.error or "Не удалось сгенерировать изображение.",
|
||
)
|
||
|
||
async def _retry(self, func, *args, attempts: int = 3, **kwargs):
|
||
delay = 1.0
|
||
for attempt in range(1, attempts + 1):
|
||
try:
|
||
return await func(*args, **kwargs)
|
||
except (TimedOut, NetworkError):
|
||
if attempt >= attempts:
|
||
raise
|
||
logger.warning(
|
||
"Telegram timeout (%s/%s), повтор через %.0f с",
|
||
attempt,
|
||
attempts,
|
||
delay,
|
||
)
|
||
await asyncio.sleep(delay)
|
||
delay *= 2
|
||
|
||
async def _edit_status(
|
||
self,
|
||
request_id: int,
|
||
text: str,
|
||
*,
|
||
percent: int | None = None,
|
||
) -> None:
|
||
if self._bot is None:
|
||
return
|
||
target = self._status_messages.get(request_id)
|
||
if target is None:
|
||
return
|
||
|
||
now = time.monotonic()
|
||
last = self._last_edit.get(request_id)
|
||
if (
|
||
last is not None
|
||
and percent is not None
|
||
and percent - last[1] < 3
|
||
and now - last[0] < 2.5
|
||
):
|
||
return
|
||
|
||
chat_id, message_id = target
|
||
try:
|
||
await self._retry(
|
||
self._bot.edit_message_text,
|
||
chat_id=chat_id,
|
||
message_id=message_id,
|
||
text=text,
|
||
)
|
||
except BadRequest as exc:
|
||
if "not modified" not in str(exc).lower():
|
||
logger.debug("Не удалось обновить статус %s: %s", request_id, exc)
|
||
return
|
||
except TelegramError as exc:
|
||
logger.debug("Не удалось обновить статус %s: %s", request_id, exc)
|
||
return
|
||
self._last_edit[request_id] = (
|
||
now,
|
||
percent if percent is not None else -1,
|
||
)
|
||
|
||
async def _finish_done(self, request: Request) -> None:
|
||
target = self._status_messages.get(request.id)
|
||
chat_id = target[0] if target is not None else request.user_id
|
||
output_path = request.output_file_path
|
||
|
||
if not output_path or not Path(output_path).exists():
|
||
await self._finish_failed(request, "Файл результата не найден.")
|
||
return
|
||
|
||
if self._bot is None:
|
||
return
|
||
|
||
caption = f"Готово. Seed: {request.seed}" if request.seed is not None else "Готово."
|
||
filename = Path(output_path).name
|
||
try:
|
||
content = await asyncio.to_thread(Path(output_path).read_bytes)
|
||
except OSError:
|
||
logger.exception("Не удалось прочитать результат запроса %s", request.id)
|
||
await self._finish_failed(request, "Файл результата не читается.")
|
||
return
|
||
try:
|
||
await self._retry(
|
||
self._bot.send_photo,
|
||
chat_id=chat_id,
|
||
photo=InputFile(content, filename=filename),
|
||
caption=caption,
|
||
read_timeout=120,
|
||
write_timeout=120,
|
||
attempts=2,
|
||
)
|
||
except TelegramError:
|
||
logger.exception("Не удалось отправить фото по запросу %s", request.id)
|
||
try:
|
||
await self._retry(
|
||
self._bot.send_document,
|
||
chat_id=chat_id,
|
||
document=InputFile(content, filename=filename),
|
||
caption=caption,
|
||
read_timeout=120,
|
||
write_timeout=120,
|
||
attempts=2,
|
||
)
|
||
except TelegramError:
|
||
await self._finish_failed(
|
||
request, "Не удалось отправить результат в Telegram."
|
||
)
|
||
return
|
||
|
||
await self._clear_status(request.id, delete=True)
|
||
|
||
async def _finish_failed(self, request: Request, error: str) -> None:
|
||
target = self._status_messages.get(request.id)
|
||
chat_id = target[0] if target is not None else request.user_id
|
||
text = f"Не получилось: {error}"
|
||
if target is not None and self._bot is not None:
|
||
try:
|
||
await self._retry(
|
||
self._bot.edit_message_text,
|
||
chat_id=target[0],
|
||
message_id=target[1],
|
||
text=text,
|
||
)
|
||
except TelegramError:
|
||
logger.debug("Не удалось показать ошибку по запросу %s", request.id)
|
||
if self._bot is not None:
|
||
try:
|
||
await self._retry(
|
||
self._bot.send_message, chat_id=chat_id, text=text
|
||
)
|
||
except TelegramError:
|
||
logger.exception(
|
||
"Не удалось отправить ошибку по запросу %s", request.id
|
||
)
|
||
await self._clear_status(request.id, delete=False)
|
||
return
|
||
|
||
if self._bot is not None:
|
||
try:
|
||
await self._retry(
|
||
self._bot.send_message, chat_id=chat_id, text=text
|
||
)
|
||
except TelegramError:
|
||
logger.exception("Не удалось отправить ошибку по запросу %s", request.id)
|
||
|
||
async def _clear_status(self, request_id: int, *, delete: bool) -> None:
|
||
target = self._status_messages.pop(request_id, None)
|
||
self._last_edit.pop(request_id, None)
|
||
if delete and target is not None and self._bot is not None:
|
||
try:
|
||
await self._retry(
|
||
self._bot.delete_message,
|
||
chat_id=target[0],
|
||
message_id=target[1],
|
||
attempts=2,
|
||
)
|
||
except TelegramError:
|
||
logger.debug("Не удалось удалить статус %s", request_id)
|
||
|
||
async def _check_access(self, update: Update) -> bool:
|
||
message = update.effective_message
|
||
chat = update.effective_chat
|
||
user = update.effective_user
|
||
if message is None or chat is None or user is None:
|
||
return False
|
||
if chat.type != ChatType.PRIVATE:
|
||
await message.reply_text("Напиши боту в личку.")
|
||
return False
|
||
if (
|
||
self._settings.allowed_user_ids
|
||
and user.id not in self._settings.allowed_user_ids
|
||
):
|
||
await message.reply_text("Доступ к боту ограничен.")
|
||
return False
|
||
return True
|
||
|
||
async def _on_error(
|
||
self, update: object, context: ContextTypes.DEFAULT_TYPE
|
||
) -> None:
|
||
logger.error("Ошибка обработчика", exc_info=context.error)
|
||
if isinstance(update, Update) and update.effective_message is not None:
|
||
try:
|
||
await update.effective_message.reply_text(
|
||
"Внутренняя ошибка. Попробуй ещё раз."
|
||
)
|
||
except TelegramError:
|
||
pass
|
||
|
||
|
||
def main() -> None:
|
||
settings = Settings.from_env()
|
||
setup_logging(settings.log_level)
|
||
MpakoBot(settings).run()
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|