feat: initial working bot with ComfyUI queue

.dockerignore

@@ -0,0 +1,14 @@
+.git
+.gitignore
+.env
+.env.*
+!.env.example
+storage/
+data/
+__pycache__/
+*.py[cod]
+.venv/
+venv/
+.pytest_cache/
+tests/
+README.md

.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

.gitattributes

@@ -0,0 +1,4 @@
+* text=auto eol=lf
+
+*.bat text eol=crlf
+*.cmd text eol=crlf

.gitignore

@@ -7,6 +7,9 @@ __pycache__/
 .venv/
 venv/
 .pytest_cache/
+.ruff_cache/
+.mypy_cache/
+.coverage
 
 storage/
 data/

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"]

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`.

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()

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. Допущения и решения по умолчанию
 
 | Вопрос | Решение по умолчанию |

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)

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,
+    )

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()

docker-compose.yml

@@ -0,0 +1,8 @@
+services:
+  bot:
+    build: .
+    env_file: .env
+    volumes:
+      - ./storage:/app/storage
+      - ./data:/app/data
+    restart: unless-stopped

pytest.ini

@@ -0,0 +1,3 @@
+[pytest]
+asyncio_mode = auto
+testpaths = tests

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)

requirements-dev.txt

@@ -0,0 +1,3 @@
+-r requirements.txt
+pytest>=8,<9
+pytest-asyncio>=0.24,<1

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

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]

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

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()

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)

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