467 lines
18 KiB
Python
467 lines
18 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, TelegramError
|
||
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 queue_manager import QueueEvent, QueueFull, QueueManager
|
||
from workflows import WorkflowBuilder
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
START_TEXT = (
|
||
"Пришли текст — сгенерирую картинку.\n"
|
||
"Пришли фото с подписью — отредактирую его по описанию.\n\n"
|
||
"/status — узнать состояние текущей задачи."
|
||
)
|
||
|
||
|
||
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)
|
||
.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()
|
||
self._queue = QueueManager(
|
||
db=self._db,
|
||
comfy=self._comfy,
|
||
workflows=self._workflows,
|
||
settings=self._settings,
|
||
on_event=self._on_queue_event,
|
||
)
|
||
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
|
||
|
||
await self._submit(update, mode="t2i", prompt=prompt, input_path=None)
|
||
|
||
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]
|
||
if photo.file_size and photo.file_size > self._settings.max_photo_bytes:
|
||
limit_mb = self._settings.max_photo_bytes / (1024 * 1024)
|
||
await message.reply_text(f"Фото слишком большое. Максимум {limit_mb:.0f} МБ.")
|
||
return
|
||
|
||
user_dir = self._settings.storage_dir / str(user.id)
|
||
user_dir.mkdir(parents=True, exist_ok=True)
|
||
input_path = user_dir / f"{uuid.uuid4().hex}_input.jpg"
|
||
try:
|
||
await message.chat.send_action(ChatAction.TYPING)
|
||
file = await photo.get_file()
|
||
await file.download_to_drive(custom_path=input_path)
|
||
except TelegramError:
|
||
logger.exception("Не удалось скачать фото пользователя %s", user.id)
|
||
input_path.unlink(missing_ok=True)
|
||
await message.reply_text("Не удалось скачать фото, попробуй ещё раз.")
|
||
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 _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
|
||
|
||
try:
|
||
request, position = await self._queue.enqueue(
|
||
user_id=user.id,
|
||
mode=mode,
|
||
prompt=prompt,
|
||
input_file_path=input_path,
|
||
)
|
||
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
|
||
|
||
status_message = await message.reply_text(self._ack_text(position))
|
||
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 == "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 _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._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._bot.send_photo(
|
||
chat_id=chat_id,
|
||
photo=InputFile(content, filename=filename),
|
||
caption=caption,
|
||
read_timeout=120,
|
||
write_timeout=120,
|
||
)
|
||
except TelegramError:
|
||
logger.exception("Не удалось отправить фото по запросу %s", request.id)
|
||
try:
|
||
await self._bot.send_document(
|
||
chat_id=chat_id,
|
||
document=InputFile(content, filename=filename),
|
||
caption=caption,
|
||
read_timeout=120,
|
||
write_timeout=120,
|
||
)
|
||
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._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._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._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._bot.delete_message(
|
||
chat_id=target[0], message_id=target[1]
|
||
)
|
||
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()
|