From ce9acabde539840fc874a7e7f98d99d8c829b46f Mon Sep 17 00:00:00 2001 From: MpakobecEx Date: Sun, 18 Oct 2026 17:49:55 +0500 Subject: [PATCH] feat: initial working bot with ComfyUI queue --- .dockerignore | 14 ++ .env.example | 28 +++ .gitattributes | 4 + .gitignore | 3 + Dockerfile | 18 ++ README.md | 76 +++++++ bot.py | 462 ++++++++++++++++++++++++++++++++++++++++ checklist.md | 7 + comfy_client.py | 331 ++++++++++++++++++++++++++++ config.py | 123 +++++++++++ db.py | 252 ++++++++++++++++++++++ docker-compose.yml | 8 + pytest.ini | 3 + queue_manager.py | 272 +++++++++++++++++++++++ requirements-dev.txt | 3 + requirements.txt | 5 + tests/test_comfy.py | 40 ++++ tests/test_db.py | 110 ++++++++++ tests/test_queue.py | 178 ++++++++++++++++ tests/test_workflows.py | 62 ++++++ workflows.py | 119 +++++++++++ 21 files changed, 2118 insertions(+) create mode 100644 .dockerignore create mode 100644 .env.example create mode 100644 .gitattributes create mode 100644 Dockerfile create mode 100644 README.md create mode 100644 bot.py create mode 100644 comfy_client.py create mode 100644 config.py create mode 100644 db.py create mode 100644 docker-compose.yml create mode 100644 pytest.ini create mode 100644 queue_manager.py create mode 100644 requirements-dev.txt create mode 100644 requirements.txt create mode 100644 tests/test_comfy.py create mode 100644 tests/test_db.py create mode 100644 tests/test_queue.py create mode 100644 tests/test_workflows.py create mode 100644 workflows.py diff --git a/.dockerignore b/.dockerignore new file mode 100644 index 0000000..6d9d5cb --- /dev/null +++ b/.dockerignore @@ -0,0 +1,14 @@ +.git +.gitignore +.env +.env.* +!.env.example +storage/ +data/ +__pycache__/ +*.py[cod] +.venv/ +venv/ +.pytest_cache/ +tests/ +README.md diff --git a/.env.example b/.env.example new file mode 100644 index 0000000..215f073 --- /dev/null +++ b/.env.example @@ -0,0 +1,28 @@ +# Токен от @BotFather +BOT_TOKEN= + +# ComfyUI API +COMFY_BASE_URL=http://192.168.31.113:8188 +WORKFLOW_T2I=workflows/txt2img.json +WORKFLOW_I2I=workflows/img2img.json +WORKFLOW_MAPPING=workflows/mapping.json + +# Хранилище +STORAGE_DIR=storage +DB_PATH=data/bot.db + +# Очередь +WORKER_CONCURRENCY=1 +MAX_QUEUE_SIZE=20 +GENERATION_TIMEOUT=600 +PROGRESS_INTERVAL=3 + +# Ограничения +MAX_PROMPT_LENGTH=1000 +MAX_PHOTO_BYTES=15728640 +DEFAULT_NEGATIVE=low quality, blurry, bad anatomy, extra digits + +# Доступ: пусто — всем; иначе список id через запятую +ALLOWED_USER_IDS= + +LOG_LEVEL=INFO diff --git a/.gitattributes b/.gitattributes new file mode 100644 index 0000000..a2237bc --- /dev/null +++ b/.gitattributes @@ -0,0 +1,4 @@ +* text=auto eol=lf + +*.bat text eol=crlf +*.cmd text eol=crlf diff --git a/.gitignore b/.gitignore index e996dc9..88cc799 100644 --- a/.gitignore +++ b/.gitignore @@ -7,6 +7,9 @@ __pycache__/ .venv/ venv/ .pytest_cache/ +.ruff_cache/ +.mypy_cache/ +.coverage storage/ data/ diff --git a/Dockerfile b/Dockerfile new file mode 100644 index 0000000..dae8b75 --- /dev/null +++ b/Dockerfile @@ -0,0 +1,18 @@ +FROM python:3.13-slim + +ENV PYTHONUNBUFFERED=1 \ + PIP_NO_CACHE_DIR=1 + +WORKDIR /app + +RUN useradd --create-home --uid 1000 app + +COPY requirements.txt . +RUN pip install --no-cache-dir -r requirements.txt + +COPY . . +RUN mkdir -p storage data && chown -R app:app /app + +USER app + +CMD ["python", "bot.py"] diff --git a/README.md b/README.md new file mode 100644 index 0000000..88ed739 --- /dev/null +++ b/README.md @@ -0,0 +1,76 @@ +# MpakoImageAIbot + +Telegram-бот для генерации и редактирования изображений через удалённый ComfyUI API. + +## Что умеет + +- **Text-to-image**: текст → изображение (Z-Image Turbo, 1024×1024). +- **Image-to-image**: фото + подпись → редактирование (Krea2 edit). +- Глобальная очередь с одной активной задачей на пользователя. +- Позиция в очереди и прогресс генерации через WebSocket ComfyUI. +- SQLite: запросы, статусы, seed, пути к файлам. + +## Требования + +- Python 3.10+ (проверено на 3.13). +- ComfyUI с доступом по сети, например: + ```text + python main.py --listen 0.0.0.0 --port 8188 + ``` +- Установленные кастомные ноды для текущих workflow: `Krea2ModelLoader`, `TextEncodeQwenImageEdit`, `UnetLoaderGGUF`, `CLIPLoaderGGUF`. +- Модели, указанные в `workflows/txt2img.json` и `workflows/img2img.json`. + +## Настройка + +```powershell +Copy-Item .env.example .env +``` + +Заполнить `.env`: + +- `BOT_TOKEN` — токен от `@BotFather`. +- `COMFY_BASE_URL` — например `http://192.168.31.113:8188`. +- Остальные параметры можно оставить по умолчанию. + +## Запуск + +```powershell +py -3 -m venv .venv +.\.venv\Scripts\Activate.ps1 +pip install -r requirements.txt +python bot.py +``` + +Бот работает в приватных чатах. Пользователь отправляет текст или фото с подписью. + +## Docker + +```powershell +docker compose up -d --build +``` + +Каталоги `./storage` и `./data` монтируются как volume. Для ComfyUI на другой машине используйте её LAN-IP, а не `localhost`. + +## Workflow + +- `workflows/txt2img.json` — Z-Image Turbo: `EmptySD3LatentImage → KSampler → VAEDecode → SaveImage`. +- `workflows/img2img.json` — Krea2 edit: `LoadImage → TextEncodeQwenImageEdit → KSamplerAdvanced → VAEDecode → SaveImage`. +- `workflows/mapping.json` — ID узлов, в которые подставляются prompt, negative, seed и входное изображение. +- `workflows/Krea2.json` и `workflows/ZImage.json` — исходные экспорты. + +Если меняются workflow, обновите `mapping.json`. При старте бот проверяет, что все узлы и поля существуют. + +## Тесты + +```powershell +pip install -r requirements-dev.txt +pytest +``` + +## Ограничения текущей версии + +- Один воркер и одна активная задача на пользователя. +- Нет `/cancel` и удаления задачи из очереди. +- Нет автоочистки `storage/` и `data/`. +- Прогресс в процентах зависит от WebSocket; при его недоступности используется polling без процентов. +- При рестарте задача в статусе `running` помечается `failed`. diff --git a/bot.py b/bot.py new file mode 100644 index 0000000..a78c48c --- /dev/null +++ b/bot.py @@ -0,0 +1,462 @@ +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: + application = ( + Application.builder() + .token(self._settings.bot_token) + .post_init(self._post_init) + .post_shutdown(self._post_shutdown) + .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() diff --git a/checklist.md b/checklist.md index e6bd681..1c22c39 100644 --- a/checklist.md +++ b/checklist.md @@ -4,6 +4,13 @@ Блоки идут в порядке зависимостей. В конце каждого блока — критерий готовности (DoD). +## Текущий статус + +- Реализованы блоки 2–7, базовое хранилище (9), файлы Docker (10), тесты (11) и README (12). +- Вживую проверены: T2I (Z-Image), I2I (Krea2), WebSocket-прогресс, переключение T2I → I2I. +- 15 тестов проходят: БД, workflows, очередь, вспомогательные функции ComfyUI. +- Не проверено: сборка Docker (демон не запущен), очередь из 3+ пользователей, падение ComfyUI, рестарт с `running`-задачей, реальная отправка в Telegram (нет токена). + ## 0. Допущения и решения по умолчанию | Вопрос | Решение по умолчанию | diff --git a/comfy_client.py b/comfy_client.py new file mode 100644 index 0000000..069f5bf --- /dev/null +++ b/comfy_client.py @@ -0,0 +1,331 @@ +from __future__ import annotations + +import asyncio +import json +import logging +import mimetypes +import os +from collections.abc import Callable +from pathlib import Path +from typing import Any +from urllib.parse import quote + +import httpx + +try: + from websockets.asyncio.client import connect as ws_connect +except ImportError: # websockets < 13 + from websockets import connect as ws_connect + +from websockets.exceptions import ConnectionClosed + +logger = logging.getLogger(__name__) + +ProgressCallback = Callable[[int], None] + + +class ComfyError(RuntimeError): + user_message = "Ошибка при обращении к ComfyUI" + + +class ComfyConnectionError(ComfyError): + user_message = "Нет связи с ComfyUI" + + +class ComfyHTTPError(ComfyError): + user_message = "ComfyUI вернул ошибку" + + +class ComfyValidationError(ComfyError): + user_message = "Workflow отклонён ComfyUI" + + +class ComfyExecutionError(ComfyError): + user_message = "Генерация завершилась с ошибкой" + + +class ComfyTimeout(ComfyError): + user_message = "Генерация не завершилась за отведённое время" + + +def _format_api_error(payload: dict[str, Any]) -> str: + parts: list[str] = [] + error = payload.get("error") + if isinstance(error, dict): + message = error.get("message") or error.get("type") or str(error) + parts.append(str(message)) + elif error: + parts.append(str(error)) + + node_errors = payload.get("node_errors") or {} + if isinstance(node_errors, dict): + for node_id, info in node_errors.items(): + if not isinstance(info, dict): + continue + for item in info.get("errors") or []: + if not isinstance(item, dict): + continue + message = item.get("message") or item.get("type") or "ошибка" + details = item.get("details") + parts.append(f"узел {node_id}: {message}" + (f" ({details})" if details else "")) + return "; ".join(parts) or "workflow отклонён" + + +def _execution_error_message(entry: dict[str, Any]) -> str: + status = entry.get("status") or {} + for message in status.get("messages") or []: + if not isinstance(message, (list, tuple)) or len(message) != 2: + continue + event, data = message + if event == "execution_error" and isinstance(data, dict): + text = data.get("exception_message") or data.get("exception_type") + if text: + return str(text) + return "ComfyUI сообщил об ошибке выполнения" + + +class ProgressSession: + def __init__(self, websocket: Any) -> None: + self._ws = websocket + self._closed = False + + async def wait( + self, + prompt_id: str, + *, + timeout: float, + on_progress: ProgressCallback, + ) -> bool: + """True — выполнение завершилось; False — WS потерян, нужен polling.""" + loop = asyncio.get_running_loop() + deadline = loop.time() + timeout + + try: + while True: + remaining = deadline - loop.time() + if remaining <= 0: + raise ComfyTimeout("WebSocket: превышено время ожидания") + try: + raw = await asyncio.wait_for(self._ws.recv(), timeout=remaining) + except TimeoutError as exc: + raise ComfyTimeout("WebSocket: превышено время ожидания") from exc + except ConnectionClosed: + logger.warning("WebSocket ComfyUI закрыт, перехожу на polling") + return False + + if isinstance(raw, bytes): + raw = raw.decode("utf-8", "replace") + try: + message = json.loads(raw) + except json.JSONDecodeError: + continue + if not isinstance(message, dict): + continue + + message_type = message.get("type") + data = message.get("data") or {} + if not isinstance(data, dict): + data = {} + + message_prompt = data.get("prompt_id") + if message_prompt and message_prompt != prompt_id: + continue + + if message_type == "progress": + value = data.get("value", 0) + maximum = data.get("max", 0) + if isinstance(value, (int, float)) and isinstance(maximum, (int, float)) and maximum > 0: + on_progress(max(0, min(100, int(value * 100 / maximum)))) + elif message_type == "execution_error": + raise ComfyExecutionError(_execution_error_message({"status": {"messages": [["execution_error", data]]}})) + elif message_type == "execution_interrupted": + raise ComfyExecutionError("Генерация прервана") + elif message_type == "execution_success" or ( + message_type == "executing" and data.get("node") is None + ): + return True + except ConnectionClosed: + logger.warning("WebSocket ComfyUI закрыт, перехожу на polling") + return False + + async def close(self) -> None: + if self._closed: + return + self._closed = True + try: + await self._ws.close() + except Exception: + logger.debug("Ошибка закрытия WebSocket", exc_info=True) + + +class ComfyClient: + def __init__(self, base_url: str, ws_url: str) -> None: + self._base_url = base_url.rstrip("/") + self._ws_url = ws_url + self._client = httpx.AsyncClient( + base_url=self._base_url, + timeout=httpx.Timeout(30.0, connect=10.0), + trust_env=False, + ) + + async def close(self) -> None: + await self._client.aclose() + + async def submit(self, workflow: dict[str, Any], client_id: str) -> str: + try: + response = await self._client.post( + "/prompt", + json={"prompt": workflow, "client_id": client_id}, + ) + except httpx.HTTPError as exc: + raise ComfyConnectionError(f"POST /prompt: {exc}") from exc + + if response.status_code >= 400: + raise ComfyHTTPError( + f"POST /prompt: HTTP {response.status_code}: {response.text[:500]}" + ) + try: + payload = response.json() + except ValueError as exc: + raise ComfyHTTPError("POST /prompt: ответ не является JSON") from exc + + if payload.get("error") or payload.get("node_errors"): + raise ComfyValidationError(_format_api_error(payload)) + + prompt_id = payload.get("prompt_id") + if not prompt_id: + raise ComfyHTTPError("POST /prompt: в ответе нет prompt_id") + return str(prompt_id) + + async def upload_image(self, path: Path, remote_name: str) -> str: + try: + content = path.read_bytes() + except OSError as exc: + raise ComfyError(f"Не удалось прочитать входной файл {path}: {exc}") from exc + + mime = mimetypes.guess_type(path.name)[0] or "application/octet-stream" + try: + response = await self._client.post( + "/upload/image", + files={"image": (remote_name, content, mime)}, + data={"type": "input", "overwrite": "true"}, + timeout=60.0, + ) + except httpx.HTTPError as exc: + raise ComfyConnectionError(f"POST /upload/image: {exc}") from exc + + if response.status_code >= 400: + raise ComfyHTTPError( + f"POST /upload/image: HTTP {response.status_code}: {response.text[:500]}" + ) + try: + payload = response.json() + except ValueError as exc: + raise ComfyHTTPError("POST /upload/image: ответ не является JSON") from exc + + name = payload.get("name") + if not name: + raise ComfyHTTPError("POST /upload/image: в ответе нет имени файла") + subfolder = payload.get("subfolder") or "" + return f"{subfolder}/{name}" if subfolder else str(name) + + async def get_history(self, prompt_id: str) -> dict[str, Any] | None: + try: + response = await self._client.get(f"/history/{quote(prompt_id)}") + except httpx.HTTPError as exc: + raise ComfyConnectionError(f"GET /history: {exc}") from exc + + if response.status_code >= 400: + raise ComfyHTTPError( + f"GET /history: HTTP {response.status_code}: {response.text[:500]}" + ) + try: + payload = response.json() + except ValueError as exc: + raise ComfyHTTPError("GET /history: ответ не является JSON") from exc + entry = payload.get(prompt_id) + return entry if isinstance(entry, dict) else None + + async def wait_for_history(self, prompt_id: str, *, timeout: float) -> dict[str, Any]: + loop = asyncio.get_running_loop() + deadline = loop.time() + timeout + while loop.time() < deadline: + entry = await self.get_history(prompt_id) + if entry is not None: + status = entry.get("status") or {} + if status.get("status_str") == "error": + raise ComfyExecutionError(_execution_error_message(entry)) + if status.get("completed") or status.get("status_str") == "success": + return entry + await asyncio.sleep(min(2.0, max(0.1, deadline - loop.time()))) + raise ComfyTimeout("Не дождались результата в истории ComfyUI") + + def find_output_image(self, entry: dict[str, Any]) -> dict[str, Any] | None: + outputs = entry.get("outputs") or {} + if not isinstance(outputs, dict): + return None + for node_output in outputs.values(): + if not isinstance(node_output, dict): + continue + for image in node_output.get("images") or []: + if isinstance(image, dict) and image.get("filename"): + return image + return None + + async def download_output(self, entry: dict[str, Any], destination: Path) -> Path: + image = self.find_output_image(entry) + if image is None: + raise ComfyError("ComfyUI не вернул изображение") + return await self.download_image( + filename=str(image["filename"]), + subfolder=str(image.get("subfolder") or ""), + image_type=str(image.get("type") or "output"), + destination=destination, + ) + + async def download_image( + self, + *, + filename: str, + subfolder: str, + image_type: str, + destination: Path, + ) -> Path: + destination.parent.mkdir(parents=True, exist_ok=True) + temp_path = destination.with_name(destination.name + ".part") + params = {"filename": filename, "subfolder": subfolder, "type": image_type} + try: + async with self._client.stream( + "GET", "/view", params=params, timeout=httpx.Timeout(120.0, connect=10.0) + ) as response: + if response.status_code >= 400: + body = (await response.aread())[:500] + raise ComfyHTTPError( + f"GET /view: HTTP {response.status_code}: {body!r}" + ) + with temp_path.open("wb") as handle: + async for chunk in response.aiter_bytes(): + handle.write(chunk) + os.replace(temp_path, destination) + except httpx.HTTPError as exc: + temp_path.unlink(missing_ok=True) + raise ComfyConnectionError(f"GET /view: {exc}") from exc + except Exception: + temp_path.unlink(missing_ok=True) + raise + return destination + + async def open_progress(self, client_id: str) -> ProgressSession | None: + url = f"{self._ws_url}?clientId={quote(client_id)}" + try: + websocket = await ws_connect( + url, + open_timeout=5, + close_timeout=2, + max_size=8 * 1024 * 1024, + ping_interval=20, + ) + except Exception as exc: # noqa: BLE001 - WS не должен ронять генерацию + logger.warning("Не удалось подключиться к WebSocket ComfyUI: %s", exc) + return None + return ProgressSession(websocket) diff --git a/config.py b/config.py new file mode 100644 index 0000000..83583f7 --- /dev/null +++ b/config.py @@ -0,0 +1,123 @@ +from __future__ import annotations + +import logging +import os +from dataclasses import dataclass, field +from pathlib import Path +from urllib.parse import urlsplit, urlunsplit + +from dotenv import load_dotenv + +PROJECT_ROOT = Path(__file__).resolve().parent +load_dotenv(PROJECT_ROOT / ".env") + + +def _path(name: str, default: str) -> Path: + raw = os.getenv(name, default).strip() or default + path = Path(raw).expanduser() + return path if path.is_absolute() else PROJECT_ROOT / path + + +def _int(name: str, default: int) -> int: + raw = os.getenv(name) + if raw is None or not raw.strip(): + return default + try: + return int(raw) + except ValueError as exc: + raise ValueError(f"{name} должен быть целым числом") from exc + + +def _float(name: str, default: float) -> float: + raw = os.getenv(name) + if raw is None or not raw.strip(): + return default + try: + return float(raw) + except ValueError as exc: + raise ValueError(f"{name} должен быть числом") from exc + + +def _user_ids(name: str) -> frozenset[int]: + raw = os.getenv(name, "") + values: set[int] = set() + for part in raw.replace(";", ",").replace(" ", ",").split(","): + part = part.strip() + if not part: + continue + try: + values.add(int(part)) + except ValueError as exc: + raise ValueError(f"{name}: некорректный user_id {part!r}") from exc + return frozenset(values) + + +@dataclass(frozen=True) +class Settings: + bot_token: str + comfy_base_url: str = "http://192.168.31.113:8188" + workflow_t2i: Path = PROJECT_ROOT / "workflows" / "txt2img.json" + workflow_i2i: Path = PROJECT_ROOT / "workflows" / "img2img.json" + workflow_mapping: Path = PROJECT_ROOT / "workflows" / "mapping.json" + storage_dir: Path = PROJECT_ROOT / "storage" + db_path: Path = PROJECT_ROOT / "data" / "bot.db" + worker_concurrency: int = 1 + max_queue_size: int = 20 + allowed_user_ids: frozenset[int] = field(default_factory=frozenset) + progress_interval: float = 3.0 + generation_timeout: float = 600.0 + log_level: str = "INFO" + default_negative: str = "low quality, blurry, bad anatomy, extra digits" + max_prompt_length: int = 1000 + max_photo_bytes: int = 15 * 1024 * 1024 + + @property + def comfy_ws_url(self) -> str: + parts = urlsplit(self.comfy_base_url.rstrip("/")) + scheme = "wss" if parts.scheme == "https" else "ws" + path = parts.path.rstrip("/") + "/ws" + return urlunsplit((scheme, parts.netloc, path, "", "")) + + @classmethod + def from_env(cls, require_token: bool = True) -> Settings: + token = os.getenv("BOT_TOKEN", "").strip() + if require_token and not token: + raise ValueError("BOT_TOKEN не задан. Добавь его в .env") + + base_url = os.getenv("COMFY_BASE_URL", "http://192.168.31.113:8188").strip() + if not base_url.startswith(("http://", "https://")): + raise ValueError("COMFY_BASE_URL должен начинаться с http:// или https://") + + settings = cls( + bot_token=token, + comfy_base_url=base_url.rstrip("/"), + workflow_t2i=_path("WORKFLOW_T2I", "workflows/txt2img.json"), + workflow_i2i=_path("WORKFLOW_I2I", "workflows/img2img.json"), + workflow_mapping=_path("WORKFLOW_MAPPING", "workflows/mapping.json"), + storage_dir=_path("STORAGE_DIR", "storage"), + db_path=_path("DB_PATH", "data/bot.db"), + worker_concurrency=max(1, _int("WORKER_CONCURRENCY", 1)), + max_queue_size=max(1, _int("MAX_QUEUE_SIZE", 20)), + allowed_user_ids=_user_ids("ALLOWED_USER_IDS"), + progress_interval=max(1.0, _float("PROGRESS_INTERVAL", 3.0)), + generation_timeout=max(30.0, _float("GENERATION_TIMEOUT", 600.0)), + log_level=os.getenv("LOG_LEVEL", "INFO").strip().upper() or "INFO", + default_negative=os.getenv("DEFAULT_NEGATIVE", "").strip() + or cls.default_negative, + max_prompt_length=max(1, _int("MAX_PROMPT_LENGTH", 1000)), + max_photo_bytes=max(1, _int("MAX_PHOTO_BYTES", 15 * 1024 * 1024)), + ) + return settings + + +def get_settings(require_token: bool = True) -> Settings: + return Settings.from_env(require_token=require_token) + + +def setup_logging(level: str) -> None: + logging.basicConfig( + level=getattr(logging, level.upper(), logging.INFO), + format="%(asctime)s %(levelname)s %(name)s: %(message)s", + datefmt="%Y-%m-%d %H:%M:%S", + force=True, + ) diff --git a/db.py b/db.py new file mode 100644 index 0000000..abaa71b --- /dev/null +++ b/db.py @@ -0,0 +1,252 @@ +from __future__ import annotations + +import asyncio +import logging +import sqlite3 +from dataclasses import dataclass +from datetime import UTC, datetime +from pathlib import Path +from typing import Any + +import aiosqlite + +logger = logging.getLogger(__name__) + +ACTIVE_STATUSES = ("queued", "running") + + +def utc_now() -> str: + return datetime.now(UTC).isoformat() + + +class UserAlreadyQueued(RuntimeError): + def __init__(self, user_id: int) -> None: + super().__init__(f"У пользователя {user_id} уже есть активный запрос") + self.user_id = user_id + + +@dataclass(frozen=True) +class Request: + id: int + user_id: int + timestamp: str + prompt_text: str + input_file_path: str | None + output_file_path: str | None + status: str + comfy_prompt_id: str | None + mode: str + seed: int | None + error_text: str | None + finished_at: str | None + + @classmethod + def from_row(cls, row: aiosqlite.Row) -> Request: + return cls( + id=row["id"], + user_id=row["user_id"], + timestamp=row["timestamp"], + prompt_text=row["prompt_text"], + input_file_path=row["input_file_path"], + output_file_path=row["output_file_path"], + status=row["status"], + comfy_prompt_id=row["comfy_prompt_id"], + mode=row["mode"], + seed=row["seed"], + error_text=row["error_text"], + finished_at=row["finished_at"], + ) + + +SCHEMA = """ +CREATE TABLE IF NOT EXISTS requests ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + user_id INTEGER NOT NULL, + timestamp TEXT NOT NULL, + prompt_text TEXT NOT NULL, + input_file_path TEXT, + output_file_path TEXT, + status TEXT NOT NULL DEFAULT 'queued' + CHECK (status IN ('queued', 'running', 'done', 'failed')), + comfy_prompt_id TEXT, + mode TEXT NOT NULL CHECK (mode IN ('t2i', 'i2i')), + seed INTEGER, + error_text TEXT, + finished_at TEXT +); + +CREATE INDEX IF NOT EXISTS idx_requests_user_status + ON requests(user_id, status); + +CREATE INDEX IF NOT EXISTS idx_requests_status_id + ON requests(status, id); + +CREATE UNIQUE INDEX IF NOT EXISTS uq_requests_one_active + ON requests(user_id) WHERE status IN ('queued', 'running'); +""" + + +class Database: + def __init__(self, path: Path) -> None: + self._path = path + self._conn: aiosqlite.Connection | None = None + self._write_lock = asyncio.Lock() + + @property + def path(self) -> Path: + return self._path + + async def connect(self) -> None: + self._path.parent.mkdir(parents=True, exist_ok=True) + self._conn = await aiosqlite.connect(self._path) + self._conn.row_factory = aiosqlite.Row + await self._conn.execute("PRAGMA journal_mode=WAL") + await self._conn.execute("PRAGMA busy_timeout=5000") + await self._conn.execute("PRAGMA foreign_keys=ON") + await self._conn.executescript(SCHEMA) + await self._conn.commit() + + async def close(self) -> None: + if self._conn is not None: + await self._conn.close() + self._conn = None + + def _connection(self) -> aiosqlite.Connection: + if self._conn is None: + raise RuntimeError("База данных не подключена") + return self._conn + + async def create_request( + self, + *, + user_id: int, + mode: str, + prompt_text: str, + input_file_path: str | None, + seed: int, + ) -> Request: + conn = self._connection() + async with self._write_lock: + try: + cursor = await conn.execute( + """ + INSERT INTO requests + (user_id, timestamp, prompt_text, input_file_path, status, mode, seed) + VALUES (?, ?, ?, ?, 'queued', ?, ?) + """, + (user_id, utc_now(), prompt_text, input_file_path, mode, seed), + ) + request_id = cursor.lastrowid + await conn.commit() + except sqlite3.IntegrityError as exc: + await conn.rollback() + if "UNIQUE" in str(exc).upper(): + raise UserAlreadyQueued(user_id) from exc + raise + request = await self.get_request(int(request_id)) + if request is None: + raise RuntimeError("Не удалось прочитать созданный запрос") + return request + + async def get_request(self, request_id: int) -> Request | None: + conn = self._connection() + async with conn.execute("SELECT * FROM requests WHERE id = ?", (request_id,)) as cursor: + row = await cursor.fetchone() + return Request.from_row(row) if row is not None else None + + async def get_active_by_user(self, user_id: int) -> Request | None: + conn = self._connection() + async with conn.execute( + """ + SELECT * FROM requests + WHERE user_id = ? AND status IN ('queued', 'running') + ORDER BY id DESC LIMIT 1 + """, + (user_id,), + ) as cursor: + row = await cursor.fetchone() + return Request.from_row(row) if row is not None else None + + async def list_queued(self) -> list[Request]: + conn = self._connection() + async with conn.execute( + "SELECT * FROM requests WHERE status = 'queued' ORDER BY id" + ) as cursor: + rows = await cursor.fetchall() + return [Request.from_row(row) for row in rows] + + async def count_ahead(self, request_id: int) -> int: + conn = self._connection() + async with conn.execute( + """ + SELECT COUNT(*) AS total FROM requests + WHERE id < ? AND status IN ('queued', 'running') + """, + (request_id,), + ) as cursor: + row = await cursor.fetchone() + return int(row["total"]) if row is not None else 0 + + async def set_running(self, request_id: int, comfy_prompt_id: str | None = None) -> None: + await self._update( + """ + UPDATE requests + SET status = 'running', + comfy_prompt_id = COALESCE(?, comfy_prompt_id), + error_text = NULL + WHERE id = ? + """, + (comfy_prompt_id, request_id), + ) + + async def set_comfy_prompt_id(self, request_id: int, comfy_prompt_id: str) -> None: + await self._update( + "UPDATE requests SET comfy_prompt_id = ? WHERE id = ?", + (comfy_prompt_id, request_id), + ) + + async def set_done(self, request_id: int, output_file_path: str) -> None: + await self._update( + """ + UPDATE requests + SET status = 'done', + output_file_path = ?, + error_text = NULL, + finished_at = ? + WHERE id = ? + """, + (output_file_path, utc_now(), request_id), + ) + + async def set_failed(self, request_id: int, error_text: str) -> None: + await self._update( + """ + UPDATE requests + SET status = 'failed', + error_text = ?, + finished_at = ? + WHERE id = ? + """, + (error_text[:1000], utc_now(), request_id), + ) + + async def recover(self, reason: str = "прервано перезапуском") -> list[Request]: + """Помечает running как failed и возвращает queued для повторной постановки.""" + conn = self._connection() + async with self._write_lock: + await conn.execute( + """ + UPDATE requests + SET status = 'failed', error_text = ?, finished_at = ? + WHERE status = 'running' + """, + (reason, utc_now()), + ) + await conn.commit() + return await self.list_queued() + + async def _update(self, sql: str, params: tuple[Any, ...]) -> None: + conn = self._connection() + async with self._write_lock: + await conn.execute(sql, params) + await conn.commit() diff --git a/docker-compose.yml b/docker-compose.yml new file mode 100644 index 0000000..5225bcf --- /dev/null +++ b/docker-compose.yml @@ -0,0 +1,8 @@ +services: + bot: + build: . + env_file: .env + volumes: + - ./storage:/app/storage + - ./data:/app/data + restart: unless-stopped diff --git a/pytest.ini b/pytest.ini new file mode 100644 index 0000000..78c5011 --- /dev/null +++ b/pytest.ini @@ -0,0 +1,3 @@ +[pytest] +asyncio_mode = auto +testpaths = tests diff --git a/queue_manager.py b/queue_manager.py new file mode 100644 index 0000000..a09daf1 --- /dev/null +++ b/queue_manager.py @@ -0,0 +1,272 @@ +from __future__ import annotations + +import asyncio +import inspect +import logging +import secrets +import uuid +from collections import deque +from collections.abc import Awaitable, Callable +from dataclasses import dataclass +from pathlib import Path + +from comfy_client import ComfyClient, ComfyError +from config import Settings +from db import Database, Request, UserAlreadyQueued +from workflows import WorkflowBuilder + +logger = logging.getLogger(__name__) + +EventCallback = Callable[["QueueEvent"], Awaitable[None] | None] + + +@dataclass(frozen=True) +class QueueEvent: + kind: str # started | progress | done | failed + request: Request + position: int | None = None + percent: int | None = None + error: str | None = None + + +class QueueError(RuntimeError): + user_message = "Очередь недоступна" + + +class QueueFull(QueueError): + user_message = "Очередь переполнена, попробуй позже" + + +class QueueStopped(QueueError): + user_message = "Бот останавливается" + + +class QueueManager: + def __init__( + self, + *, + db: Database, + comfy: ComfyClient, + workflows: WorkflowBuilder, + settings: Settings, + on_event: EventCallback, + ) -> None: + self._db = db + self._comfy = comfy + self._workflows = workflows + self._settings = settings + self._on_event = on_event + + self._queue: deque[int] = deque() + self._active_users: set[int] = set() + self._wake = asyncio.Event() + self._lock = asyncio.Lock() + self._workers: list[asyncio.Task[None]] = [] + self._background: set[asyncio.Task[None]] = set() + self._stopping = False + self._progress_state: dict[int, tuple[float, int]] = {} + + async def start(self) -> None: + recovered = await self._db.recover() + async with self._lock: + for request in recovered: + self._queue.append(request.id) + self._active_users.add(request.user_id) + if recovered: + logger.info("Восстановлено из очереди: %s", len(recovered)) + self._wake.set() + for index in range(self._settings.worker_concurrency): + task = asyncio.create_task(self._worker(index), name=f"queue-worker-{index}") + self._workers.append(task) + + async def stop(self) -> None: + self._stopping = True + self._wake.set() + for task in self._workers: + task.cancel() + if self._workers: + await asyncio.gather(*self._workers, return_exceptions=True) + self._workers.clear() + for task in self._background: + task.cancel() + if self._background: + await asyncio.gather(*self._background, return_exceptions=True) + self._background.clear() + + async def enqueue( + self, + *, + user_id: int, + mode: str, + prompt: str, + input_file_path: str | None = None, + ) -> tuple[Request, int]: + async with self._lock: + if self._stopping: + raise QueueStopped() + active = await self._db.get_active_by_user(user_id) + if active is not None: + raise UserAlreadyQueued(user_id) + if len(self._queue) >= self._settings.max_queue_size: + raise QueueFull() + + seed = secrets.randbelow(2**32) + request = await self._db.create_request( + user_id=user_id, + mode=mode, + prompt_text=prompt, + input_file_path=input_file_path, + seed=seed, + ) + self._queue.append(request.id) + self._active_users.add(user_id) + self._wake.set() + position = await self._db.count_ahead(request.id) + return request, position + + async def position(self, request_id: int) -> int: + return await self._db.count_ahead(request_id) + + async def _worker(self, index: int) -> None: + logger.info("Воркер %s запущен", index) + try: + while not self._stopping: + request_id = await self._next_request() + if request_id is None: + break + await self._process(request_id) + finally: + logger.info("Воркер %s остановлен", index) + + async def _next_request(self) -> int | None: + while not self._stopping: + if self._queue: + return self._queue.popleft() + self._wake.clear() + await self._wake.wait() + return None + + async def _process(self, request_id: int) -> None: + request = await self._db.get_request(request_id) + if request is None: + logger.error("Запрос %s не найден в БД", request_id) + return + + await self._db.set_running(request_id) + request = await self._refresh(request_id) + self._progress_state[request.id] = (0.0, -1) + await self._emit(QueueEvent(kind="started", request=request, percent=0)) + + client_id = uuid.uuid4().hex + progress = None + try: + workflow = await self._build_workflow(request) + progress = await self._comfy.open_progress(client_id) + prompt_id = await self._comfy.submit(workflow, client_id) + await self._db.set_comfy_prompt_id(request_id, prompt_id) + + loop = asyncio.get_running_loop() + deadline = loop.time() + self._settings.generation_timeout + + if progress is not None: + await progress.wait( + prompt_id, + timeout=max(0.1, deadline - loop.time()), + on_progress=lambda percent: self._schedule_progress(request, percent), + ) + + entry = await self._comfy.wait_for_history( + prompt_id, + timeout=max(1.0, deadline - loop.time()), + ) + + image = self._comfy.find_output_image(entry) + if image is None: + raise ComfyError("ComfyUI не вернул изображение") + suffix = Path(str(image["filename"])).suffix or ".png" + output_path = ( + self._settings.storage_dir + / str(request.user_id) + / f"{request.id}_output{suffix}" + ) + await self._comfy.download_output(entry, output_path) + + await self._db.set_done(request_id, str(output_path)) + request = await self._refresh(request_id) + self._active_users.discard(request.user_id) + await self._emit(QueueEvent(kind="done", request=request, percent=100)) + except asyncio.CancelledError: + await self._db.set_failed(request_id, "остановлено") + raise + except Exception as exc: + logger.exception("Ошибка запроса %s", request_id) + await self._db.set_failed(request_id, str(exc)) + request = await self._refresh(request_id) + self._active_users.discard(request.user_id) + await self._emit( + QueueEvent( + kind="failed", + request=request, + error=getattr(exc, "user_message", "Ошибка генерации"), + ) + ) + finally: + if progress is not None: + await progress.close() + self._active_users.discard(request.user_id) + self._progress_state.pop(request_id, None) + + async def _build_workflow(self, request: Request) -> dict: + seed = request.seed if request.seed is not None else secrets.randbelow(2**32) + negative = self._settings.default_negative + if request.mode == "t2i": + return self._workflows.build_txt2img( + prompt=request.prompt_text, + negative=negative, + seed=seed, + ) + + if not request.input_file_path: + raise ComfyError("Для i2i не передан входной файл") + input_path = Path(request.input_file_path) + suffix = input_path.suffix.lower() or ".png" + remote_name = f"{request.id}_input{suffix}" + image_name = await self._comfy.upload_image(input_path, remote_name) + return self._workflows.build_img2img( + prompt=request.prompt_text, + negative=negative, + seed=seed, + image_name=image_name, + ) + + async def _refresh(self, request_id: int) -> Request: + request = await self._db.get_request(request_id) + if request is None: + raise RuntimeError(f"Запрос {request_id} исчез из БД") + return request + + def _schedule_progress(self, request: Request, percent: int) -> None: + loop = asyncio.get_running_loop() + now = loop.time() + last_time, last_percent = self._progress_state.get(request.id, (0.0, -1)) + if percent <= last_percent: + return + if last_percent >= 0: + if percent - last_percent < 5: + return + if now - last_time < self._settings.progress_interval: + return + self._progress_state[request.id] = (now, percent) + task = asyncio.create_task( + self._emit(QueueEvent(kind="progress", request=request, percent=percent)) + ) + self._background.add(task) + task.add_done_callback(self._background.discard) + + async def _emit(self, event: QueueEvent) -> None: + try: + result = self._on_event(event) + if inspect.isawaitable(result): + await result + except Exception: + logger.exception("Ошибка обработчика события %s", event.kind) diff --git a/requirements-dev.txt b/requirements-dev.txt new file mode 100644 index 0000000..2968596 --- /dev/null +++ b/requirements-dev.txt @@ -0,0 +1,3 @@ +-r requirements.txt +pytest>=8,<9 +pytest-asyncio>=0.24,<1 diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..5c0b660 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,5 @@ +python-telegram-bot>=21.6,<23 +httpx>=0.27,<1 +aiosqlite>=0.20,<1 +websockets>=13,<16 +python-dotenv>=1.0,<2 diff --git a/tests/test_comfy.py b/tests/test_comfy.py new file mode 100644 index 0000000..5cbfab7 --- /dev/null +++ b/tests/test_comfy.py @@ -0,0 +1,40 @@ +from __future__ import annotations + +from comfy_client import ComfyClient, _format_api_error + + +def test_format_node_errors() -> None: + message = _format_api_error( + { + "error": {"message": "Prompt outputs failed validation"}, + "node_errors": { + "20": { + "errors": [ + {"message": "Required input is missing", "details": "image"} + ] + } + }, + } + ) + assert "Prompt outputs failed validation" in message + assert "узел 20" in message + assert "image" in message + + +def test_find_output_image() -> None: + entry = { + "outputs": { + "9": { + "images": [ + {"filename": "result.png", "subfolder": "", "type": "output"} + ] + } + } + } + image = ComfyClient.find_output_image(None, entry) # type: ignore[arg-type] + assert image is not None + assert image["filename"] == "result.png" + + +def test_find_output_image_empty() -> None: + assert ComfyClient.find_output_image(None, {}) is None # type: ignore[arg-type] diff --git a/tests/test_db.py b/tests/test_db.py new file mode 100644 index 0000000..4532287 --- /dev/null +++ b/tests/test_db.py @@ -0,0 +1,110 @@ +from __future__ import annotations + +import pytest +import pytest_asyncio + +from db import Database, UserAlreadyQueued + + +@pytest_asyncio.fixture +async def db(tmp_path): + database = Database(tmp_path / "bot.db") + await database.connect() + yield database + await database.close() + + +async def test_create_and_read(db: Database) -> None: + request = await db.create_request( + user_id=1, + mode="t2i", + prompt_text="apple", + input_file_path=None, + seed=42, + ) + assert request.id > 0 + assert request.status == "queued" + assert request.mode == "t2i" + assert request.seed == 42 + + loaded = await db.get_request(request.id) + assert loaded == request + + +async def test_one_active_request_per_user(db: Database) -> None: + first = await db.create_request( + user_id=1, + mode="t2i", + prompt_text="first", + input_file_path=None, + seed=1, + ) + with pytest.raises(UserAlreadyQueued): + await db.create_request( + user_id=1, + mode="t2i", + prompt_text="second", + input_file_path=None, + seed=2, + ) + + await db.set_done(first.id, "/tmp/out.png") + second = await db.create_request( + user_id=1, + mode="t2i", + prompt_text="second", + input_file_path=None, + seed=3, + ) + assert second.id != first.id + + +async def test_positions(db: Database) -> None: + first = await db.create_request( + user_id=1, + mode="t2i", + prompt_text="first", + input_file_path=None, + seed=1, + ) + second = await db.create_request( + user_id=2, + mode="t2i", + prompt_text="second", + input_file_path=None, + seed=2, + ) + assert await db.count_ahead(first.id) == 0 + assert await db.count_ahead(second.id) == 1 + + await db.set_running(first.id, "prompt-1") + assert await db.count_ahead(second.id) == 1 + + await db.set_done(first.id, "/tmp/out.png") + assert await db.count_ahead(second.id) == 0 + + +async def test_recover_marks_running_failed_and_returns_queued(db: Database) -> None: + running = await db.create_request( + user_id=1, + mode="t2i", + prompt_text="running", + input_file_path=None, + seed=1, + ) + queued = await db.create_request( + user_id=2, + mode="t2i", + prompt_text="queued", + input_file_path=None, + seed=2, + ) + await db.set_running(running.id, "prompt-1") + + recovered = await db.recover() + + assert [item.id for item in recovered] == [queued.id] + failed = await db.get_request(running.id) + assert failed is not None + assert failed.status == "failed" + assert failed.finished_at is not None diff --git a/tests/test_queue.py b/tests/test_queue.py new file mode 100644 index 0000000..f872015 --- /dev/null +++ b/tests/test_queue.py @@ -0,0 +1,178 @@ +from __future__ import annotations + +import asyncio +import time +from pathlib import Path + +import pytest +import pytest_asyncio + +from config import Settings +from db import Database, UserAlreadyQueued +from queue_manager import QueueEvent, QueueFull, QueueManager +from workflows import WorkflowBuilder + +BASE = Path(__file__).resolve().parents[1] / "workflows" + + +class FakeComfy: + def __init__(self) -> None: + self.workflows: list[dict] = [] + self.uploads: list[tuple[Path, str]] = [] + + async def open_progress(self, client_id: str): + return None + + async def submit(self, workflow: dict, client_id: str) -> str: + self.workflows.append(workflow) + return f"prompt-{len(self.workflows)}" + + async def wait_for_history(self, prompt_id: str, *, timeout: float) -> dict: + return {"outputs": {"9": {"images": [{"filename": "x.png", "type": "output"}]}}} + + def find_output_image(self, entry: dict) -> dict | None: + images = entry["outputs"]["9"]["images"] + return images[0] if images else None + + async def download_output(self, entry: dict, destination: Path) -> Path: + destination.parent.mkdir(parents=True, exist_ok=True) + destination.write_bytes(b"png") + return destination + + async def upload_image(self, path: Path, remote_name: str) -> str: + self.uploads.append((path, remote_name)) + return remote_name + + +@pytest_asyncio.fixture +async def db(tmp_path): + database = Database(tmp_path / "bot.db") + await database.connect() + yield database + await database.close() + + +def make_settings(tmp_path: Path, **overrides) -> Settings: + values = { + "bot_token": "test", + "storage_dir": tmp_path / "storage", + "db_path": tmp_path / "bot.db", + "workflow_t2i": BASE / "txt2img.json", + "workflow_i2i": BASE / "img2img.json", + "workflow_mapping": BASE / "mapping.json", + } + values.update(overrides) + return Settings(**values) + + +def make_queue(db: Database, settings: Settings, events: list[QueueEvent]) -> QueueManager: + async def on_event(event: QueueEvent) -> None: + events.append(event) + + return QueueManager( + db=db, + comfy=FakeComfy(), + workflows=WorkflowBuilder( + settings.workflow_t2i, + settings.workflow_i2i, + settings.workflow_mapping, + ), + settings=settings, + on_event=on_event, + ) + + +async def wait_for_event( + events: list[QueueEvent], kind: str, timeout: float = 5.0 +) -> QueueEvent: + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + for event in events: + if event.kind == kind: + return event + await asyncio.sleep(0.01) + raise AssertionError(f"Событие {kind!r} не пришло: {[e.kind for e in events]}") + + +async def test_positions_and_duplicate_user(db: Database, tmp_path: Path) -> None: + events: list[QueueEvent] = [] + queue = make_queue(db, make_settings(tmp_path), events) + + first, first_position = await queue.enqueue( + user_id=1, mode="t2i", prompt="first" + ) + second, second_position = await queue.enqueue( + user_id=2, mode="t2i", prompt="second" + ) + + assert first_position == 0 + assert second_position == 1 + with pytest.raises(UserAlreadyQueued): + await queue.enqueue(user_id=1, mode="t2i", prompt="again") + + await queue.stop() + assert first.id != second.id + + +async def test_queue_limit(db: Database, tmp_path: Path) -> None: + events: list[QueueEvent] = [] + queue = make_queue( + db, + make_settings(tmp_path, max_queue_size=1), + events, + ) + + await queue.enqueue(user_id=1, mode="t2i", prompt="first") + with pytest.raises(QueueFull): + await queue.enqueue(user_id=2, mode="t2i", prompt="second") + + await queue.stop() + + +async def test_text_to_image_job(db: Database, tmp_path: Path) -> None: + events: list[QueueEvent] = [] + settings = make_settings(tmp_path) + queue = make_queue(db, settings, events) + + await queue.start() + try: + request, _ = await queue.enqueue( + user_id=1, mode="t2i", prompt="apple" + ) + event = await wait_for_event(events, "done") + + assert event.request.id == request.id + assert event.request.status == "done" + output = Path(event.request.output_file_path or "") + assert output.exists() + assert output.name == f"{request.id}_output.png" + finally: + await queue.stop() + + +async def test_image_to_image_uploads_input(db: Database, tmp_path: Path) -> None: + events: list[QueueEvent] = [] + settings = make_settings(tmp_path) + queue = make_queue(db, settings, events) + comfy = queue._comfy # type: ignore[attr-defined] + + input_path = tmp_path / "input.png" + input_path.write_bytes(b"input") + + await queue.start() + try: + request, _ = await queue.enqueue( + user_id=1, + mode="i2i", + prompt="make blue", + input_file_path=str(input_path), + ) + await wait_for_event(events, "done") + + assert comfy.uploads + assert comfy.uploads[0][0] == input_path + assert comfy.uploads[0][1] == f"{request.id}_input.png" + source = comfy.workflows[0]["20"]["inputs"]["image"] + assert source == f"{request.id}_input.png" + finally: + await queue.stop() diff --git a/tests/test_workflows.py b/tests/test_workflows.py new file mode 100644 index 0000000..ca1fdb5 --- /dev/null +++ b/tests/test_workflows.py @@ -0,0 +1,62 @@ +from __future__ import annotations + +import json +from pathlib import Path + +import pytest + +from workflows import WorkflowBuilder, WorkflowError + +BASE = Path(__file__).resolve().parents[1] / "workflows" + + +def make_builder() -> WorkflowBuilder: + return WorkflowBuilder( + BASE / "txt2img.json", + BASE / "img2img.json", + BASE / "mapping.json", + ) + + +def test_build_txt2img() -> None: + builder = make_builder() + workflow = builder.build_txt2img(prompt="cat", negative="bad", seed=123) + + assert workflow["67"]["inputs"]["text"] == "cat" + assert workflow["71"]["inputs"]["text"] == "bad" + assert workflow["70"]["inputs"]["seed"] == 123 + assert workflow["81"]["class_type"] == "EmptySD3LatentImage" + assert workflow["70"]["inputs"]["denoise"] == 1.0 + + +def test_build_img2img() -> None: + builder = make_builder() + workflow = builder.build_img2img( + prompt="make blue", + negative="bad", + seed=7, + image_name="input.png", + ) + + assert workflow["25"]["inputs"]["prompt"] == "make blue" + assert workflow["10"]["inputs"]["text"] == "bad" + assert workflow["3"]["inputs"]["noise_seed"] == 7 + assert workflow["20"]["class_type"] == "LoadImage" + assert workflow["20"]["inputs"]["image"] == "input.png" + + +def test_templates_are_not_mutated() -> None: + builder = make_builder() + builder.build_txt2img(prompt="first", negative="bad", seed=1) + second = builder.build_txt2img(prompt="second", negative="bad", seed=2) + assert second["67"]["inputs"]["text"] == "second" + + +def test_invalid_mapping(tmp_path: Path) -> None: + mapping = json.loads((BASE / "mapping.json").read_text(encoding="utf-8")) + mapping["txt2img"]["positive"] = "999" + path = tmp_path / "mapping.json" + path.write_text(json.dumps(mapping), encoding="utf-8") + + with pytest.raises(WorkflowError): + WorkflowBuilder(BASE / "txt2img.json", BASE / "img2img.json", path) diff --git a/workflows.py b/workflows.py new file mode 100644 index 0000000..dd24195 --- /dev/null +++ b/workflows.py @@ -0,0 +1,119 @@ +from __future__ import annotations + +import copy +import json +from collections.abc import Callable +from pathlib import Path +from typing import Any + +RoleField = tuple[str | None, Callable[[Any], bool] | None] + +ROLE_FIELDS: dict[str, dict[str, RoleField]] = { + "txt2img": { + "positive": ("text", lambda value: isinstance(value, str)), + "negative": ("text", lambda value: isinstance(value, str)), + "seed": ("seed", lambda value: isinstance(value, int)), + "latent": (None, None), + "save": (None, None), + }, + "img2img": { + "positive": ("prompt", lambda value: isinstance(value, str)), + "negative": ("text", lambda value: isinstance(value, str)), + "seed": ("noise_seed", lambda value: isinstance(value, int)), + "input_image": ("image", lambda value: isinstance(value, str)), + "save": (None, None), + }, +} + +LATENT_CLASSES = {"EmptySD3LatentImage", "EmptyLatentImage", "EmptyLatentImagePresets"} + + +class WorkflowError(RuntimeError): + pass + + +class WorkflowBuilder: + def __init__(self, t2i_path: Path, i2i_path: Path, mapping_path: Path) -> None: + self._paths = {"txt2img": t2i_path, "img2img": i2i_path} + self._templates = {mode: self._load_json(path) for mode, path in self._paths.items()} + self._mapping = self._load_json(mapping_path) + self._validate() + + @staticmethod + def _load_json(path: Path) -> dict[str, Any]: + try: + data = json.loads(path.read_text(encoding="utf-8")) + except FileNotFoundError as exc: + raise WorkflowError(f"Не найден workflow: {path}") from exc + except json.JSONDecodeError as exc: + raise WorkflowError(f"Некорректный JSON в {path}: {exc}") from exc + if not isinstance(data, dict) or not data: + raise WorkflowError(f"Workflow {path} должен быть непустым объектом") + return data + + def _validate(self) -> None: + if not isinstance(self._mapping, dict): + raise WorkflowError("mapping.json должен быть объектом") + + for mode, roles in ROLE_FIELDS.items(): + mode_mapping = self._mapping.get(mode) + if not isinstance(mode_mapping, dict): + raise WorkflowError(f"В mapping.json отсутствует раздел {mode!r}") + template = self._templates[mode] + + for role, (field_name, check) in roles.items(): + node_id = mode_mapping.get(role) + if not isinstance(node_id, str) or node_id not in template: + raise WorkflowError( + f"{mode}: узел для роли {role!r} не найден в workflow" + ) + node = template[node_id] + inputs = node.get("inputs") + if not isinstance(inputs, dict): + raise WorkflowError(f"{mode}: у узла {node_id} нет inputs") + + if role == "latent": + if node.get("class_type") not in LATENT_CLASSES: + raise WorkflowError( + f"{mode}: узел latent должен создавать пустой латент" + ) + if "width" not in inputs or "height" not in inputs: + raise WorkflowError(f"{mode}: у latent нет width/height") + continue + + if field_name is None: + continue + if field_name not in inputs: + raise WorkflowError( + f"{mode}: у узла {node_id} нет поля {field_name!r} " + f"для роли {role!r}" + ) + if check is not None and not check(inputs[field_name]): + raise WorkflowError( + f"{mode}: поле {field_name!r} у узла {node_id} " + f"имеет неожиданный тип" + ) + + def build_txt2img(self, *, prompt: str, negative: str, seed: int) -> dict[str, Any]: + workflow = copy.deepcopy(self._templates["txt2img"]) + mapping = self._mapping["txt2img"] + workflow[mapping["positive"]]["inputs"]["text"] = prompt + workflow[mapping["negative"]]["inputs"]["text"] = negative + workflow[mapping["seed"]]["inputs"]["seed"] = seed + return workflow + + def build_img2img( + self, + *, + prompt: str, + negative: str, + seed: int, + image_name: str, + ) -> dict[str, Any]: + workflow = copy.deepcopy(self._templates["img2img"]) + mapping = self._mapping["img2img"] + workflow[mapping["positive"]]["inputs"]["prompt"] = prompt + workflow[mapping["negative"]]["inputs"]["text"] = negative + workflow[mapping["seed"]]["inputs"]["noise_seed"] = seed + workflow[mapping["input_image"]]["inputs"]["image"] = image_name + return workflow