commit d6df4a3118c4777905e849ff32c32af0b8f95202 Author: MpakobecEx Date: Thu Oct 8 20:14:00 2026 +0500 Первая версия бота diff --git a/.dockerignore b/.dockerignore new file mode 100644 index 0000000..eed2628 --- /dev/null +++ b/.dockerignore @@ -0,0 +1,15 @@ +.git +.gitea +.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..0f4e9e9 --- /dev/null +++ b/.env.example @@ -0,0 +1,32 @@ +# Токен от @BotFather +BOT_TOKEN= + +# Прокси для Telegram API. Пусто — без прокси. +# В Docker на Linux: http://host.docker.internal:20171 +TELEGRAM_PROXY= + +# 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/.gitea/workflows/docker-build.yml b/.gitea/workflows/docker-build.yml new file mode 100644 index 0000000..2d9712a --- /dev/null +++ b/.gitea/workflows/docker-build.yml @@ -0,0 +1,50 @@ +name: Build and push Docker image + +"on": + push: + branches: + - main + - master + workflow_dispatch: + +env: + REGISTRY: git.mpakosrv.ru + IMAGE: git.mpakosrv.ru/mpakobecex/mpakoimageaibot + +jobs: + docker: + runs-on: ubuntu-latest + timeout-minutes: 20 + steps: + - name: Checkout + uses: actions/checkout@v4 + + - name: Login to registry + env: + REGISTRY_USER: ${{ secrets.REGISTRY_USER }} + REGISTRY_TOKEN: ${{ secrets.REGISTRY_TOKEN }} + run: | + echo "$REGISTRY_TOKEN" | docker login "$REGISTRY" \ + --username "$REGISTRY_USER" \ + --password-stdin + + - name: Build image + run: | + docker build \ + --label org.opencontainers.image.revision=${{ github.sha }} \ + -t mpakoimageaibot:latest \ + -t "$IMAGE:latest" \ + -t "$IMAGE:${GITHUB_SHA}" \ + . + + - name: Push image + run: | + docker push "$IMAGE:latest" + docker push "$IMAGE:${GITHUB_SHA}" + + - name: Show image + run: docker images mpakoimageaibot + + - name: Logout + if: always() + run: docker logout "$REGISTRY" || true diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..88cc799 --- /dev/null +++ b/.gitignore @@ -0,0 +1,21 @@ +.env +.env.* +!.env.example + +__pycache__/ +*.py[cod] +.venv/ +venv/ +.pytest_cache/ +.ruff_cache/ +.mypy_cache/ +.coverage + +storage/ +data/ +*.db +*.db-wal +*.db-shm + +.idea/ +.vscode/ 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..cd65699 --- /dev/null +++ b/README.md @@ -0,0 +1,110 @@ +# MpakoImageAIbot + +Telegram-бот для генерации и редактирования изображений через удалённый ComfyUI API. + +## Что умеет + +- **Text-to-image**: текст → изображение (Z-Image Turbo, 1024×1024). +- **Image-to-image**: фото + подпись → редактирование (Krea2 edit). +- **Reply-to-image**: ответ текстом на готовую картинку бота → редактирование этой картинки. +- Глобальная очередь с одной активной задачей на пользователя. +- Позиция в очереди и прогресс генерации через 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`. +- `TELEGRAM_PROXY` — прокси для Telegram API (пусто — напрямую). +- `COMFY_BASE_URL` — например `http://192.168.31.113:8188`. +- Остальные параметры можно оставить по умолчанию. + +### Прокси для Telegram + +Если сервер не имеет прямого доступа к Telegram, укажи прокси V2RayA/HTTP/SOCKS5: + +```dotenv +# Docker на Linux (V2RayA слушает на хосте) +TELEGRAM_PROXY=http://host.docker.internal:20171 + +# Запуск без Docker, на том же хосте +# TELEGRAM_PROXY=http://127.0.0.1:20171 +``` + +Для SOCKS5 используй `socks5h://host.docker.internal:20170`; зависимость `socksio` уже включена в `requirements.txt`. + +## Запуск + +```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`. + +## CI + +При push в `main` или `master` Gitea Actions собирает образ на сервере, где запущен runner, и пушит его в Gitea Registry: + +- `git.mpakosrv.ru/mpakobecex/mpakoimageaibot:latest` +- `git.mpakosrv.ru/mpakobecex/mpakoimageaibot:` +- локально: `mpakoimageaibot:latest` + +Workflow: `.gitea/workflows/docker-build.yml`. Runner использует label `ubuntu-latest` и смонтированный Docker socket. Для входа в registry нужны секреты `REGISTRY_USER` и `REGISTRY_TOKEN`. + +После успешной сборки обновление контейнера: + +```bash +docker compose up -d +``` + +Без `--build`: Compose возьмёт готовый образ `mpakoimageaibot:latest`. + +## 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..d16a41f --- /dev/null +++ b/bot.py @@ -0,0 +1,571 @@ +from __future__ import annotations + +import asyncio +import logging +import time +import uuid +from pathlib import Path + +from telegram import InputFile, Update +from telegram.constants import ChatAction, ChatType +from telegram.error import BadRequest, NetworkError, TelegramError, TimedOut +from telegram.ext import ( + Application, + CommandHandler, + ContextTypes, + MessageHandler, + filters, +) + +from comfy_client import ComfyClient +from config import Settings, setup_logging +from db import Database, Request, UserAlreadyQueued +from queue_manager import QueueEvent, QueueFull, QueueManager +from workflows import WorkflowBuilder + +logger = logging.getLogger(__name__) + +START_TEXT = ( + "Пришли текст — сгенерирую картинку.\n" + "Пришли фото с подписью — отредактирую его по описанию.\n" + "Или ответь текстом на готовую картинку — возьму её как исходник.\n\n" + "/status — узнать состояние текущей задачи." +) + + +class MpakoBot: + def __init__(self, settings: Settings) -> None: + self._settings = settings + self._db = Database(settings.db_path) + self._workflows: WorkflowBuilder | None = None + self._comfy: ComfyClient | None = None + self._queue: QueueManager | None = None + self._bot = None + self._status_messages: dict[int, tuple[int, int]] = {} + self._last_edit: dict[int, tuple[float, int]] = {} + + def run(self) -> None: + builder = ( + Application.builder() + .token(self._settings.bot_token) + .connect_timeout(20) + .read_timeout(60) + .write_timeout(60) + .get_updates_connect_timeout(20) + .get_updates_read_timeout(60) + .post_init(self._post_init) + .post_shutdown(self._post_shutdown) + ) + if self._settings.telegram_proxy: + builder = builder.proxy(self._settings.telegram_proxy).get_updates_proxy( + self._settings.telegram_proxy + ) + application = builder.build() + application.add_handler(CommandHandler("start", self._start)) + application.add_handler(CommandHandler("help", self._help)) + application.add_handler(CommandHandler("status", self._status)) + application.add_handler(MessageHandler(filters.PHOTO, self._photo)) + application.add_handler( + MessageHandler(filters.TEXT & ~filters.COMMAND, self._text) + ) + application.add_error_handler(self._on_error) + application.run_polling(allowed_updates=Update.ALL_TYPES) + + async def _post_init(self, application: Application) -> None: + self._bot = application.bot + self._settings.storage_dir.mkdir(parents=True, exist_ok=True) + self._workflows = WorkflowBuilder( + self._settings.workflow_t2i, + self._settings.workflow_i2i, + self._settings.workflow_mapping, + ) + self._comfy = ComfyClient( + self._settings.comfy_base_url, + self._settings.comfy_ws_url, + ) + await self._db.connect() + 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 + + reply = message.reply_to_message + attachment = None + suffix = ".jpg" + if reply is not None: + if reply.photo: + attachment = reply.photo[-1] + elif reply.document is not None and ( + reply.document.mime_type or "" + ).startswith("image/"): + attachment = reply.document + suffix = Path(reply.document.file_name or "").suffix.lower() or ".png" + + if attachment is None: + await self._submit( + update, mode="t2i", prompt=prompt, input_path=None + ) + return + + input_path = await self._download_attachment( + message, attachment, user.id, suffix=suffix + ) + if input_path is None: + return + accepted = await self._submit( + update, + mode="i2i", + prompt=prompt, + input_path=str(input_path), + ) + if not accepted: + input_path.unlink(missing_ok=True) + + async def _photo( + self, update: Update, context: ContextTypes.DEFAULT_TYPE + ) -> None: + if not await self._check_access(update): + return + message = update.effective_message + user = update.effective_user + if message is None or user is None or not message.photo: + return + + caption = (message.caption or "").strip() + if not caption: + await message.reply_text( + "Пришли фото с подписью — что именно изменить на изображении." + ) + return + if len(caption) > self._settings.max_prompt_length: + await message.reply_text( + f"Описание слишком длинное: максимум {self._settings.max_prompt_length} символов." + ) + return + + photo = message.photo[-1] + input_path = await self._download_attachment( + message, photo, user.id, suffix=".jpg" + ) + if input_path is None: + return + + accepted = await self._submit( + update, + mode="i2i", + prompt=caption, + input_path=str(input_path), + ) + if not accepted: + input_path.unlink(missing_ok=True) + + async def _download_attachment( + self, + message, + attachment, + user_id: int, + *, + suffix: str = ".jpg", + ) -> Path | None: + if attachment.file_size and attachment.file_size > self._settings.max_photo_bytes: + limit_mb = self._settings.max_photo_bytes / (1024 * 1024) + try: + await self._retry( + message.reply_text, + f"Файл слишком большой. Максимум {limit_mb:.0f} МБ.", + ) + except TelegramError: + logger.debug("Не удалось сообщить о размере файла") + return None + + suffix = suffix.lower() + if not suffix.startswith(".") or len(suffix) > 8: + suffix = ".png" + + user_dir = self._settings.storage_dir / str(user_id) + input_path = user_dir / f"{uuid.uuid4().hex}_input{suffix}" + try: + user_dir.mkdir(parents=True, exist_ok=True) + await message.chat.send_action(ChatAction.TYPING) + file = await attachment.get_file() + await file.download_to_drive(custom_path=input_path) + except (TelegramError, OSError): + logger.exception("Не удалось скачать вложение пользователя %s", user_id) + input_path.unlink(missing_ok=True) + try: + await self._retry( + message.reply_text, + "Не удалось скачать файл, попробуй ещё раз.", + ) + except TelegramError: + pass + return None + return input_path + + async def _submit( + self, + update: Update, + *, + mode: str, + prompt: str, + input_path: str | None, + ) -> bool: + message = update.effective_message + user = update.effective_user + if message is None or user is None or self._queue is None: + return False + + 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 + + try: + status_message = await self._retry( + message.reply_text, self._ack_text(position) + ) + except TelegramError: + logger.exception( + "Не удалось отправить подтверждение по запросу %s", request.id + ) + return True + self._status_messages[request.id] = ( + status_message.chat_id, + status_message.message_id, + ) + + # Воркер мог стартовать раньше, чем мы сохранили message_id. + current = await self._db.get_request(request.id) + if current is None: + return True + if current.status == "running": + await self._edit_status(request.id, "Генерация…", percent=0) + elif current.status == "done": + await self._clear_status(request.id, delete=True) + elif current.status == "failed": + await self._finish_failed( + current, + current.error_text or "Не удалось сгенерировать изображение.", + ) + return True + + def _ack_text(self, position: int) -> str: + if position <= 0: + return "Принято. Ты первый, начинаю…" + return f"Принято. Перед тобой: {position} {self._people(position)}." + + @staticmethod + def _people(count: int) -> str: + if count % 10 == 1 and count % 100 != 11: + return "человек" + if count % 10 in (2, 3, 4) and count % 100 not in (12, 13, 14): + return "человека" + return "человек" + + async def _on_queue_event(self, event: QueueEvent) -> None: + if event.kind == "started": + await self._edit_status(event.request.id, "Генерация…", percent=0) + elif event.kind == "progress": + percent = event.percent if event.percent is not None else 0 + await self._edit_status( + event.request.id, + f"Генерация… {percent}%", + percent=percent, + ) + elif event.kind == "done": + await self._finish_done(event.request) + elif event.kind == "failed": + await self._finish_failed( + event.request, + event.error or "Не удалось сгенерировать изображение.", + ) + + async def _retry(self, func, *args, attempts: int = 3, **kwargs): + delay = 1.0 + for attempt in range(1, attempts + 1): + try: + return await func(*args, **kwargs) + except (TimedOut, NetworkError): + if attempt >= attempts: + raise + logger.warning( + "Telegram timeout (%s/%s), повтор через %.0f с", + attempt, + attempts, + delay, + ) + await asyncio.sleep(delay) + delay *= 2 + + async def _edit_status( + self, + request_id: int, + text: str, + *, + percent: int | None = None, + ) -> None: + if self._bot is None: + return + target = self._status_messages.get(request_id) + if target is None: + return + + now = time.monotonic() + last = self._last_edit.get(request_id) + if ( + last is not None + and percent is not None + and percent - last[1] < 3 + and now - last[0] < 2.5 + ): + return + + chat_id, message_id = target + try: + await self._retry( + self._bot.edit_message_text, + chat_id=chat_id, + message_id=message_id, + text=text, + ) + except BadRequest as exc: + if "not modified" not in str(exc).lower(): + logger.debug("Не удалось обновить статус %s: %s", request_id, exc) + return + except TelegramError as exc: + logger.debug("Не удалось обновить статус %s: %s", request_id, exc) + return + self._last_edit[request_id] = ( + now, + percent if percent is not None else -1, + ) + + async def _finish_done(self, request: Request) -> None: + target = self._status_messages.get(request.id) + chat_id = target[0] if target is not None else request.user_id + output_path = request.output_file_path + + if not output_path or not Path(output_path).exists(): + await self._finish_failed(request, "Файл результата не найден.") + return + + if self._bot is None: + return + + caption = f"Готово. Seed: {request.seed}" if request.seed is not None else "Готово." + filename = Path(output_path).name + try: + content = await asyncio.to_thread(Path(output_path).read_bytes) + except OSError: + logger.exception("Не удалось прочитать результат запроса %s", request.id) + await self._finish_failed(request, "Файл результата не читается.") + return + try: + await self._retry( + self._bot.send_photo, + chat_id=chat_id, + photo=InputFile(content, filename=filename), + caption=caption, + read_timeout=120, + write_timeout=120, + attempts=2, + ) + except TelegramError: + logger.exception("Не удалось отправить фото по запросу %s", request.id) + try: + await self._retry( + self._bot.send_document, + chat_id=chat_id, + document=InputFile(content, filename=filename), + caption=caption, + read_timeout=120, + write_timeout=120, + attempts=2, + ) + except TelegramError: + await self._finish_failed( + request, "Не удалось отправить результат в Telegram." + ) + return + + await self._clear_status(request.id, delete=True) + + async def _finish_failed(self, request: Request, error: str) -> None: + target = self._status_messages.get(request.id) + chat_id = target[0] if target is not None else request.user_id + text = f"Не получилось: {error}" + if target is not None and self._bot is not None: + try: + await self._retry( + self._bot.edit_message_text, + chat_id=target[0], + message_id=target[1], + text=text, + ) + except TelegramError: + logger.debug("Не удалось показать ошибку по запросу %s", request.id) + if self._bot is not None: + try: + await self._retry( + self._bot.send_message, chat_id=chat_id, text=text + ) + except TelegramError: + logger.exception( + "Не удалось отправить ошибку по запросу %s", request.id + ) + await self._clear_status(request.id, delete=False) + return + + if self._bot is not None: + try: + await self._retry( + self._bot.send_message, chat_id=chat_id, text=text + ) + except TelegramError: + logger.exception("Не удалось отправить ошибку по запросу %s", request.id) + + async def _clear_status(self, request_id: int, *, delete: bool) -> None: + target = self._status_messages.pop(request_id, None) + self._last_edit.pop(request_id, None) + if delete and target is not None and self._bot is not None: + try: + await self._retry( + self._bot.delete_message, + chat_id=target[0], + message_id=target[1], + attempts=2, + ) + except TelegramError: + logger.debug("Не удалось удалить статус %s", request_id) + + async def _check_access(self, update: Update) -> bool: + message = update.effective_message + chat = update.effective_chat + user = update.effective_user + if message is None or chat is None or user is None: + return False + if chat.type != ChatType.PRIVATE: + await message.reply_text("Напиши боту в личку.") + return False + if ( + self._settings.allowed_user_ids + and user.id not in self._settings.allowed_user_ids + ): + await message.reply_text("Доступ к боту ограничен.") + return False + return True + + async def _on_error( + self, update: object, context: ContextTypes.DEFAULT_TYPE + ) -> None: + logger.error("Ошибка обработчика", exc_info=context.error) + if isinstance(update, Update) and update.effective_message is not None: + try: + await update.effective_message.reply_text( + "Внутренняя ошибка. Попробуй ещё раз." + ) + except TelegramError: + pass + + +def main() -> None: + settings = Settings.from_env() + setup_logging(settings.log_level) + MpakoBot(settings).run() + + +if __name__ == "__main__": + main() diff --git a/checklist.md b/checklist.md new file mode 100644 index 0000000..9db6525 --- /dev/null +++ b/checklist.md @@ -0,0 +1,323 @@ +# Checklist реализации MpakoImageAIbot + +Источник: `plan.md`. Бот на `python-telegram-bot` v21+ (async) генерирует изображения через удалённый ComfyUI API, ведёт глобальную очередь, показывает позицию/прогресс и хранит запросы в SQLite. + +Блоки идут в порядке зависимостей. В конце каждого блока — критерий готовности (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. Допущения и решения по умолчанию + +| Вопрос | Решение по умолчанию | +|---|---| +| HTTP-клиент | `httpx.AsyncClient`, без блокирующего `requests` | +| Очередь | Собственная очередь в приложении, по умолчанию 1 воркер; в ComfyUI уходит не более 1 задачи одновременно | +| Прогресс | WebSocket `/ws` + fallback-опрос `/history/{prompt_id}` | +| Один активный запрос на пользователя | Проверка в памяти + частичный уникальный индекс в SQLite | +| Задача в статусе `running` при рестарте | Пометить `failed` (не перезапускать автоматически) | +| Задачи `queued` при рестарте | Восстановить из БД и поставить в очередь заново | +| i2i без подписи | Отклонить и попросить промпт (не использовать дефолтный) | +| Чаты | Приватные; `user_id` считается id чата | +| Workflow | JSON в API-формате (экспорт «Save (API Format)» из ComfyUI) | +| T2I-модель | Z-Image Turbo GGUF, `EmptySD3LatentImage`, 1024×1024, 8 steps, cfg 1, denoise 1.0 | +| I2I-модель | Krea2 edit GGUF, `LoadImage`, 8 steps, cfg 1 | +| ComfyUI | `http://192.168.31.113:8188`, v0.40.0, RTX 4060 Ti 8 GB | +| Язык интерфейса | Русский | + +Согласовано: **1 воркер**, задача `running` при рестарте → **`failed`**, workflow готовим вместе (см. блок 1). Для i2i используем **Krea2 edit**. + +## 1. Подготовка workflow и окружения (до кода) + +### 1.1 Экспорт workflow + +- [x] Открыть ComfyUI, включить **Settings → Dev mode Options**, чтобы появился пункт **Save (API Format)**. +- [x] Собрать t2i-граф: `UnetLoaderGGUF → CLIPLoaderGGUF → CLIPTextEncode` + `EmptySD3LatentImage → KSampler → VAEDecode → SaveImage`. +- [x] Собрать i2i-граф: `LoadImage → TextEncodeQwenImageEdit → KSamplerAdvanced → VAEDecode → SaveImage`. +- [x] Экспортировать оба графа кнопкой **Save (API Format)** в `workflows/txt2img.json` и `workflows/img2img.json`. +- [x] Убедиться, что в JSON только API-узлы (`class_type`, `inputs`), без UI-метаданных и без `extra`-полей. +- [x] Определить ID нужных узлов. Удобно: выделить узел → `Properties` → `Node ID`, либо найти по `class_type` в JSON (`CLIPTextEncode`, `KSampler`, `LoadImage`). +- [x] Заполнить `workflows/mapping.json` (ID соответствуют текущим шаблонам): + +```json +{ + "txt2img": { + "positive": "67", + "negative": "71", + "seed": "70", + "latent": "81", + "save": "9" + }, + "img2img": { + "positive": "25", + "negative": "10", + "seed": "3", + "input_image": "20", + "save": "24" + } +} +``` + +- [x] Проверить t2i вручную: отправить JSON на `POST /prompt` и получить `prompt_id`, затем забрать результат через `/history` и `/view`. +- [x] Проверить i2i: загрузить тестовое фото через `/upload/image`, подставить имя в `LoadImage`, выполнить генерацию. +- [ ] Проверить i2i с фото произвольного размера (портрет/ландшафт); при необходимости добавить `ImageScale` в граф. +- [x] Зафиксировать дефолтные `steps`, `cfg`, `denoise`, negative prompt и размер изображения. + +### 1.2 Окружение + +- [x] Убедиться, что Python 3.10+ доступен локально и в Docker-образе. +- [x] Проверить сетевую доступность ComfyUI: `GET http://IP_WINDOWS:PORT/system_stats`. +- [ ] Получить токен бота у `@BotFather`. +- [ ] Определить значения лимитов: максимальная длина промпта, максимальный размер фото, размер очереди, таймаут генерации. +- [ ] Решить, где будет запускаться бот (локально / Docker) и как он видит ComfyUI (LAN, VPN, `host.docker.internal`). + +**DoD:** есть токен, два валидных workflow JSON, заполненный `mapping.json`; тестовый `POST /prompt` возвращает `prompt_id`, результат скачивается. T2I проверен (~52 с), i2i — (~61 с). + +## 2. Каркас проекта и конфигурация + +- [ ] Создать структуру: + +```text +. +├── bot.py +├── comfy_client.py +├── db.py +├── queue_manager.py +├── config.py +├── workflows/ +│ ├── txt2img.json +│ ├── img2img.json +│ └── mapping.json +├── storage/ # создаётся автоматически, в .gitignore +├── data/ # SQLite, в .gitignore +├── tests/ +├── requirements.txt +├── requirements-dev.txt +├── Dockerfile +├── .dockerignore +├── .env.example +├── .gitignore +└── README.md +``` + +- [ ] `requirements.txt`: `python-telegram-bot`, `httpx`, `aiosqlite`, `websockets`, `python-dotenv` (версии — диапазонами, например `>=21,<23`). +- [ ] `requirements-dev.txt`: `pytest`, `pytest-asyncio`, `respx`. +- [ ] `config.py`: dataclass/`Settings`, загрузка `.env`, валидация обязательных полей, вычисление WS-URL из base URL. +- [ ] Переменные окружения: `BOT_TOKEN`, `COMFY_BASE_URL`, `WORKFLOW_T2I`, `WORKFLOW_I2I`, `WORKFLOW_MAPPING`, `STORAGE_DIR`, `DB_PATH`, `WORKER_CONCURRENCY=1`, `MAX_QUEUE_SIZE`, `ALLOWED_USER_IDS`, `PROGRESS_INTERVAL`, `GENERATION_TIMEOUT`, `LOG_LEVEL`. +- [ ] Настроить логирование в stdout: уровень из env, формат с timestamp/level/logger, никогда не логировать токен и содержимое `.env`. +- [ ] `.env.example` со всеми переменными и комментариями. +- [ ] `.gitignore`: `.env`, `storage/`, `data/`, `__pycache__/`, `.pytest_cache/`. +- [ ] Проверить импорт конфига и падение с понятной ошибкой при отсутствии `BOT_TOKEN`. + +**DoD:** `python -c "from config import get_settings; print(get_settings())"` работает; секреты не попадают в репозиторий. + +## 3. Слой хранения (`db.py`) + +- [ ] Реализовать `init_db()`: подключение `aiosqlite`, `PRAGMA journal_mode=WAL`, `PRAGMA busy_timeout`, `PRAGMA foreign_keys=ON`, `row_factory`. +- [ ] Создать таблицу `requests` по плану: + +```sql +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'))`, `error_text TEXT`, `finished_at TEXT`, `seed INTEGER`. +- [ ] Индексы: `(user_id, status)`, `(status, id)`. +- [ ] Частичный уникальный индекс «один активный запрос на пользователя»: + +```sql +CREATE UNIQUE INDEX IF NOT EXISTS uq_requests_one_active +ON requests(user_id) WHERE status IN ('queued', 'running'); +``` + +- [ ] CRUD: `create_request`, `get_request`, `get_active_by_user`, `set_status`, `set_comfy_prompt_id`, `set_output_path`, `set_error`, `list_queued`, `count_ahead`, `mark_running_as_failed`. +- [ ] Все временные метки — ISO 8601 UTC (`datetime.now(timezone.utc).isoformat()`). +- [ ] Атомарность: создание запроса и проверка активного — в одной транзакции; при конфликте unique-индекса возвращать понятную ошибку `UserAlreadyQueued`. +- [ ] Восстановление при старте: `queued` читаются по `id`, `running` → `failed` с текстом «прервано перезапуском». +- [ ] Тесты: создание, смена статусов, запрет второго активного запроса, выборка очереди, восстановление. + +**DoD:** тесты БД зелёные; повторный `init_db()` не ломает существующую БД. + +## 4. Клиент ComfyUI (`comfy_client.py`) + +- [ ] Класс `ComfyClient` с ленивой инициализацией `httpx.AsyncClient` и явным `close()`. +- [ ] `submit_workflow(workflow, client_id) -> prompt_id`: `POST /prompt` с телом `{"prompt": workflow, "client_id": client_id}`. +- [ ] Обработка ответа `/prompt`: HTTP-ошибки, `error`, `node_errors` — поднимать `ComfyError` с человекочитаемым текстом. +- [ ] `upload_image(path, request_id) -> image_ref`: `POST /upload/image` (multipart), уникальное имя файла, `type=input`; вернуть имя/подпапку для `LoadImage`. +- [ ] `get_history(prompt_id) -> dict`: `GET /history/{prompt_id}`; извлекать список изображений из `outputs[*]["images"]`. +- [ ] `download_image(filename, subfolder, type, dest_path)`: `GET /view` потоково, запись через временный файл + `os.replace`. +- [ ] `wait_for_result(prompt_id, timeout)`: fallback-опрос истории с интервалом 2 с; таймаут → `ComfyTimeout`. +- [ ] WebSocket-прогресс: `ws://.../ws?clientId=...`, подключение **до** `POST /prompt`, чтобы не потерять ранние события. +- [ ] Разбор событий: `progress` (`value`/`max`, при наличии `prompt_id` — фильтровать), `executing`, `execution_success`, `execution_error`, `status`. +- [ ] Переподключение WS с backoff; при недоступности WS — не падать, переходить на polling. +- [ ] Таймауты: connect 10 с, запросы 30 с, ожидание генерации — `GENERATION_TIMEOUT`. +- [ ] Не ретраить слепо `POST /prompt` (риск дублей); ретраить только идемпотентные GET/upload. +- [ ] Типы ошибок: `ComfyConnectionError`, `ComfyHTTPError`, `ComfyExecutionError`, `ComfyTimeout`. +- [ ] Тесты с `respx`/fake-сервером: успешный submit, `node_errors`, таймаут, скачивание, парсинг progress. + +**DoD:** отдельный скрипт/тест отправляет тестовый workflow на реальный ComfyUI, получает `prompt_id`, progress и сохраняет картинку. + +## 5. Шаблоны workflow + +- [x] Положить `txt2img.json` и `img2img.json` в `workflows/`. +- [x] `workflows/mapping.json`: соответствие «роль узла → node id» для каждого workflow (positive, negative, seed, latent/input_image). +- [ ] Валидация mapping при старте: все указанные узлы существуют, у них ожидаемые поля; иначе — падение с понятной ошибкой. +- [ ] `build_txt2img(prompt, negative, seed) -> dict`: глубокое копирование шаблона, подстановка prompt/negative/seed. +- [ ] `build_img2img(prompt, negative, seed, image_name) -> dict`: то же + имя загруженного изображения. +- [ ] Генерация seed: `secrets.randbelow(2**32)`, сохранять использованный seed в БД. +- [ ] Обрезка/нормализация промпта (strip, лимит длины), значение negative prompt по умолчанию из конфига. +- [ ] Гарантировать, что исходный шаблон не мутируется между запросами. +- [ ] Тесты: подстановка значений, отсутствующий узел, отсутствующий файл workflow. + +**DoD:** юнит-тесты подтверждают корректную подстановку; при старте приложения шаблоны и mapping валидны. + +## 6. Очередь и менеджер задач (`queue_manager.py`) + +- [ ] `QueueManager` с зависимостями: `db`, `ComfyClient`, `WorkflowBuilder`, `Settings`. +- [ ] Внутренняя очередь `collections.deque` + `asyncio.Event` (не `asyncio.Queue`, чтобы уметь считать позицию). +- [ ] `asyncio.Lock` для всех операций с очередью и реестром активных пользователей. +- [ ] `enqueue(user_id, mode, prompt, input_path) -> EnqueueResult(request, position)`: + - [ ] проверить активный запрос пользователя (память + БД); + - [ ] проверить `MAX_QUEUE_SIZE`; + - [ ] создать запись в БД, добавить id в очередь, вернуть позицию. +- [ ] Позиция: количество задач впереди + выполняющаяся задача; считать уникальных пользователей (по плану — «X человек»). +- [ ] Воркер: цикл `while not stopping`, брать задачу, никогда не умирать из-за ошибки одной задачи. +- [ ] Жизненный цикл задачи: `queued → running → done/failed`; обновлять `comfy_prompt_id`, `output_file_path`, `error_text`, `finished_at`. +- [ ] Для i2i: upload изображения → подстановка имени → submit. +- [ ] Уникальный `client_id` на задачу; WS-подключение до submit, закрытие после завершения. +- [ ] Прогресс: throttle (не чаще `PROGRESS_INTERVAL`, шаг ≥ 5%), вычисление процента из `value/max`. +- [ ] События наружу через callback/event bus: `request_started`, `progress`, `request_done`, `request_failed` — без импорта Telegram в менеджере. +- [ ] Обработка ошибок: любое исключение → `failed`, лог с `request_id`, событие для бота. +- [ ] `start()`: восстановить `queued` из БД, пометить `running` как `failed`, запустить воркер. +- [ ] `stop()`: остановить приём, дождаться воркера, пометить незавершённые задачи, закрыть клиент. +- [ ] Тесты с fake-клиентом: порядок FIFO, позиции, дубль пользователя, ошибка генерации, отмена при остановке. + +**DoD:** два пользователя получают корректные позиции; повторная заявка отклоняется; прогресс доходит до подписчика; после рестарта `queued` восстанавливаются. + +## 7. Telegram-слой (`bot.py`) + +- [ ] Инициализация `Application` (PTB v21+, async), регистрация `post_init`/`post_shutdown` для старта/остановки `QueueManager`. +- [ ] Команды `/start` и `/help`: краткая инструкция по двум режимам и формату ввода. +- [ ] (Опционально) `/status`: позиция и статус текущего запроса. +- [ ] Текстовый хендлер (`filters.TEXT & ~filters.COMMAND`): создать t2i-запрос, поставить в очередь. +- [ ] Фото-хендлер (`filters.PHOTO`): взять `photo[-1]`, скачать через `get_file().download_to_drive()`, сохранить в `storage/{user_id}/`. +- [ ] Для i2i: caption обязателен; пустой caption → подсказка «пришлите фото с описанием». +- [ ] Валидация: пустой/слишком длинный промпт, размер фото, поддерживаемые форматы; лимит caption 1024 символа. +- [ ] Ответ при дубле: «Ты уже в очереди» + текущая позиция (по плану). +- [ ] Ответ при постановке: «Принято. Перед тобой X человек» / «Ты первый, начинаю…». +- [ ] Статусное сообщение: сохранить `message_id` в памяти; редактировать при `request_started` («Генерация… 0%») и `progress`. +- [ ] Throttle редактирования статуса (не чаще 3 с), игнорировать `Message is not modified`. +- [ ] По завершении: отправить результат как фото (`FSInputFile`), статусное сообщение удалить или заменить на «Готово». +- [ ] По ошибке: редактировать статус на понятный текст, логировать техническую причину с `request_id`. +- [ ] (Опционально) `ALLOWED_USER_IDS`: отказ остальным. +- [ ] (Опционально) `ChatAction.UPLOAD_PHOTO` при отправке результата. +- [ ] Аккуратно обрабатывать `TelegramError`/`BadRequest`, не ронять приложение. +- [ ] Проверить, что команды не перехватываются текстовым хендлером. + +**DoD:** вручную проходят сценарии t2i и i2i; при дубле приходит корректный отказ; прогресс обновляется без флуда. + +## 8. Сквозная интеграция + +- [ ] Связать `bot.py → QueueManager → db → ComfyClient → storage`. +- [ ] t2i: текст → картинка → файл в `storage/{user_id}/` → отправка в чат → `done` в БД. +- [ ] i2i: фото + caption → upload → workflow → картинка → отправка → `done`. +- [ ] Очередь из 3 пользователей: позиции 1/2/3, затем корректное продвижение. +- [ ] ComfyUI недоступен: задача → `failed`, пользователь получает понятную ошибку, воркер жив. +- [ ] Ошибка выполнения workflow (`execution_error`): сообщение пользователю без стектрейса. +- [ ] Рестарт с задачами в очереди: `queued` продолжают выполняться, `running` помечены `failed`. +- [ ] Повторный запуск после падения: БД и файлы консистентны. +- [ ] Проверить отсутствие блокировки event loop (все I/O — async). + +**DoD:** все сценарии выше пройдены на реальном ComfyUI. + +## 9. Файловое хранилище + +- [ ] Структура: `storage/{user_id}/{request_id}_input.`, `storage/{user_id}/{request_id}_output.`. +- [ ] Безопасные имена (только сгенерированные, не из пользовательского ввода), защита от path traversal. +- [ ] Атомарная запись: `.part` → `os.replace`. +- [ ] Создание директорий при старте/записи. +- [ ] (Опционально) очистка файлов старше N дней по расписанию; настройка `FILE_RETENTION_DAYS`. +- [ ] (Опционально) проверка свободного места и понятная ошибка при переполнении. +- [ ] Не хранить бинарные данные в SQLite — только пути. + +**DoD:** входные и выходные файлы лежат по ожидаемым путям; повторный запуск не конфликтует по именам. + +## 10. Docker и деплой + +- [x] `Dockerfile` на `python:3.13-slim`, непривилегированный пользователь, `PYTHONUNBUFFERED=1`. +- [x] Копирование кода и установка `requirements.txt`; без лишних build-зависимостей. +- [x] `.dockerignore`: `.env`, `storage/`, `data/`, `__pycache__/`, `tests/`, `.git/`. +- [x] `docker-compose.yml`: `env_file: .env`, volume для `./storage` и `./data`, `restart: unless-stopped`. +- [x] `docker-compose.yml` использует готовый образ `mpakoimageaibot:latest`. +- [x] Gitea Actions: `.gitea/workflows/docker-build.yml` собирает и пушит `mpakoimageaibot:latest` и `:` в Registry при push в `main`/`master`. +- [x] Сетевой доступ к ComfyUI: `host.docker.internal` для Windows/macOS или LAN-IP для Linux; зафиксировать в README. +- [x] Проверить, что SQLite и storage переживают пересоздание контейнера. +- [ ] (Опционально) healthcheck. +- [x] Логи только в stdout/stderr. + +**DoD:** `docker compose up -d` поднимает бота, генерация работает, данные сохраняются после `down/up`. + +## 11. Тесты и приёмка + +- [ ] `pytest` + `pytest-asyncio`, общая фикстура временной БД и storage. +- [ ] Юнит-тесты: `db.py`, подстановка workflow, расчёт позиции, throttle прогресса, парсинг WS-событий. +- [ ] Интеграционный тест: fake ComfyUI (HTTP + WS) → полный цикл `enqueue → running → done`, файл создан. +- [ ] Тест ошибок: ComfyUI 500, `node_errors`, таймаут, обрыв WS. +- [ ] Тест конкуренции: одновременные сообщения одного пользователя → только один запрос. +- [ ] Ручной e2e на реальном ComfyUI (t2i, i2i, очередь, рестарт). +- [ ] Проверка лимитов Telegram: размер результата, длина caption, частота редактирований. + +**DoD:** тесты зелёные; ручной чек-лист из блока 8 пройден. + +## 12. Документация + +- [ ] `README.md`: описание, схема работы, требования. +- [ ] Инструкция по экспорту workflow в API-формате и заполнению `mapping.json`. +- [ ] Таблица переменных окружения с примерами. +- [ ] Запуск локально и через Docker. +- [ ] Раздел troubleshooting: ComfyUI недоступен, WS не подключается, «ты уже в очереди», ошибки workflow. +- [ ] Краткое описание архитектуры и границ MVP. + +**DoD:** новый разработчик поднимает проект по README без устных пояснений. + +## Порядок работ + +1. Блок 1 (подготовка) — без него нельзя валидировать клиент и workflow. +2. Блок 2 (каркас) → Блок 3 (БД) → Блок 4 (ComfyUI-клиент) → Блок 5 (workflow). +3. Блок 6 (очередь) → Блок 7 (бот) → Блок 8 (интеграция). +4. Блоки 9 (storage) и 10 (Docker) можно делать параллельно после блока 2. +5. Блоки 11 (тесты) и 12 (документация) — параллельно с 6–10, финализация после 8. + +## Открытые вопросы + +Согласовано: 1 воркер; `running` при рестарте → `failed`; T2I — Z-Image, i2i — Krea2 edit. Осталось уточнить: + +1. i2i без подписи: отклонять (по умолчанию) или использовать фиксированный промпт? +2. Нужны ли `/cancel`, удаление из очереди и ограничение доступа `ALLOWED_USER_IDS` в MVP? +3. Сколько хранить входные/выходные файлы? Нужна ли автоочистка? +4. Где будет хоститься бот и как он видит ComfyUI (LAN/VPN/интернет)? +5. Добавлять ли в схему `mode`, `seed`, `error_text`, `finished_at` (сейчас они предложены как расширение плана)? +6. Уведомлять ли пользователя, если его `running`-задача была помечена `failed` после рестарта? +7. Нужно ли масштабировать входное фото перед i2i (сейчас подаётся как есть; проверен только квадрат 1024×1024)? + +## Вне скоупа MVP + +- Платежи, подписки, квоты. +- Несколько моделей/чекпоинтов и выбор модели пользователем. +- Галерея, история генераций в боте, повтор генерации. +- LoRA/ControlNet и другие продвинутые workflow. +- Горизонтальное масштабирование и несколько экземпляров бота на одной SQLite. +- Админ-панель. 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..a40e200 --- /dev/null +++ b/config.py @@ -0,0 +1,137 @@ +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 + telegram_proxy: str | None = None + 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://") + + proxy = os.getenv("TELEGRAM_PROXY", "").strip() or None + if proxy and not proxy.startswith( + ("http://", "https://", "socks5://", "socks5h://") + ): + raise ValueError( + "TELEGRAM_PROXY должен начинаться с http://, https://, " + "socks5:// или socks5h://" + ) + + settings = cls( + bot_token=token, + telegram_proxy=proxy, + 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, + ) + # httpx пишет полный URL запроса, включая токен бота. + logging.getLogger("httpx").setLevel(logging.WARNING) + logging.getLogger("httpcore").setLevel(logging.WARNING) 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..d0b3627 --- /dev/null +++ b/docker-compose.yml @@ -0,0 +1,11 @@ +services: + bot: + image: mpakoimageaibot:latest + build: . + env_file: .env + extra_hosts: + - "host.docker.internal:host-gateway" + volumes: + - ./storage:/app/storage + - ./data:/app/data + restart: unless-stopped diff --git a/plan.md b/plan.md new file mode 100644 index 0000000..d0f6dba --- /dev/null +++ b/plan.md @@ -0,0 +1,52 @@ +# План проекта + +Telegram-бот для генерации изображений через удалённый ComfyUI API. + +## Стек + +- Python 3.10+ +- python-telegram-bot 21+ (async) +- SQLite (aiosqlite) +- httpx для вызова ComfyUI API +- Docker для упаковки + +## Режимы + +1. **Text-to-Image**: пользователь отправляет текст, бот подставляет промпт в workflow и отправляет его в ComfyUI через `/prompt`. +2. **Image-to-Image**: пользователь отправляет фото и подпись, бот загружает изображение в ComfyUI через `/upload/image`, подставляет его в workflow и отправляет запрос. + +## Очередь + +- Один пользователь — одна активная задача. Повторный запрос во время ожидания получает ответ «ты уже в очереди». +- Глобальная очередь с отображением позиции: «Перед тобой X человек». +- Во время генерации показывается прогресс из WebSocket ComfyUI (событие `progress`). + +## Хранение + +SQLite, таблица `requests`: + +- `id` — первичный ключ; +- `user_id` — Telegram ID; +- `timestamp` — время создания в ISO; +- `prompt_text` — текст запроса; +- `input_file_path` — входное изображение, может быть пустым; +- `output_file_path` — результат, может быть пустым; +- `status` — `queued`, `running`, `done` или `failed`; +- `comfy_prompt_id` — идентификатор задачи в ComfyUI. + +Входные файлы и результаты хранятся в `./storage/{user_id}/`, в БД записываются только пути. + +## ComfyUI + +- Базовый адрес: `http://IP_WINDOWS:PORT`. +- `POST /prompt` возвращает `prompt_id`. +- Статус отслеживается через `GET /history/{prompt_id}` или WebSocket `/ws` для прогресса в реальном времени. +- Готовое изображение скачивается через `/view?filename=...`. + +## Структура проекта + +- `bot.py` — хендлеры текста и фото. +- `comfy_client.py` — `submit_workflow` и `get_result`. +- `db.py` — `init_db()` и CRUD для `requests`. +- `queue_manager.py` — очередь на `asyncio.Queue`. +- `requirements.txt`, `Dockerfile` — зависимости и сборка. 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..e133446 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,5 @@ +python-telegram-bot>=21.6,<23 +httpx[socks]>=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 diff --git a/workflows/Krea2.json b/workflows/Krea2.json new file mode 100644 index 0000000..b8d39bc --- /dev/null +++ b/workflows/Krea2.json @@ -0,0 +1,129 @@ +{ + "3": { + "inputs": { + "add_noise": "enable", + "noise_seed": 782408632091183, + "steps": 8, + "cfg": 1, + "sampler_name": "euler", + "scheduler": "simple", + "start_at_step": 0, + "end_at_step": 10000, + "return_with_leftover_noise": "disable", + "model": [ + "19", + 0 + ], + "positive": [ + "25", + 0 + ], + "negative": [ + "10", + 0 + ], + "latent_image": [ + "15", + 0 + ] + }, + "class_type": "KSamplerAdvanced", + "_meta": { + "title": "KSampler (Advanced)" + } + }, + "10": { + "inputs": { + "text": "low quality", + "clip": [ + "19", + 1 + ] + }, + "class_type": "CLIPTextEncode", + "_meta": { + "title": "CLIP Text Encode (Prompt)" + } + }, + "11": { + "inputs": { + "samples": [ + "3", + 0 + ], + "vae": [ + "19", + 2 + ] + }, + "class_type": "VAEDecode", + "_meta": { + "title": "VAE Decode" + } + }, + "15": { + "inputs": { + "width": 1024, + "height": 1024, + "batch_size": 2 + }, + "class_type": "EmptyLatentImage", + "_meta": { + "title": "Empty Latent Image" + } + }, + "19": { + "inputs": { + "unet_name": "Krea2_turbo_uncensored_edit_v1.1-Q4_K_M.gguf", + "clip_name": "qwen3vl_4b_fp8_scaled.safetensors", + "vae_name": "qwen_image_vae.safetensors" + }, + "class_type": "Krea2ModelLoader", + "_meta": { + "title": "Krea2 Model Loader ⚡️" + } + }, + "20": { + "inputs": { + "image": "ComfyUI_00077.png [output]", + "refresh": "refresh" + }, + "class_type": "LoadImageOutput", + "_meta": { + "title": "Load Image (from Outputs)" + } + }, + "24": { + "inputs": { + "filename_prefix": "ComfyUI", + "format": "png", + "format.bit_depth": "8-bit", + "format.input_color_space": "sRGB", + "images": [ + "11", + 0 + ] + }, + "class_type": "SaveImageAdvanced", + "_meta": { + "title": "Save Image (Advanced)" + } + }, + "25": { + "inputs": { + "prompt": "Не меняй ничего. Оставь позу и ракуср без изменений.Нарисуй парню красивый член", + "clip": [ + "19", + 1 + ], + "image": [ + "20", + 0 + ] + }, + "class_type": "TextEncodeQwenImageEdit", + "_meta": { + "title": "Text Encode Qwen Image Edit" + } + } +} \ No newline at end of file diff --git a/workflows/ZImage.json b/workflows/ZImage.json new file mode 100644 index 0000000..ddc45cf --- /dev/null +++ b/workflows/ZImage.json @@ -0,0 +1,155 @@ +{ + "9": { + "inputs": { + "filename_prefix": "ComfyUI", + "images": [ + "65", + 0 + ] + }, + "class_type": "SaveImage", + "_meta": { + "title": "Save Image" + } + }, + "63": { + "inputs": { + "vae_name": "ae.safetensors" + }, + "class_type": "VAELoader", + "_meta": { + "title": "Load VAE" + } + }, + "65": { + "inputs": { + "samples": [ + "70", + 0 + ], + "vae": [ + "63", + 0 + ] + }, + "class_type": "VAEDecode", + "_meta": { + "title": "VAE Decode" + } + }, + "67": { + "inputs": { + "text": "two characters, composition inspired by reference, swap their positions: left character now on the right, right character now on the left, same art style, same lighting", + "clip": [ + "73", + 0 + ] + }, + "class_type": "CLIPTextEncode", + "_meta": { + "title": "CLIP Text Encode (Prompt)" + } + }, + "69": { + "inputs": { + "shift": 3, + "sampling": "flow", + "model": [ + "72", + 0 + ] + }, + "class_type": "ModelSamplingAuraFlow", + "_meta": { + "title": "ModelSamplingAuraFlow" + } + }, + "70": { + "inputs": { + "seed": 42, + "steps": 8, + "cfg": 1, + "sampler_name": "euler", + "scheduler": "sgm_uniform", + "denoise": 0.67, + "model": [ + "69", + 0 + ], + "positive": [ + "67", + 0 + ], + "negative": [ + "71", + 0 + ], + "latent_image": [ + "81", + 0 + ] + }, + "class_type": "KSampler", + "_meta": { + "title": "KSampler" + } + }, + "71": { + "inputs": { + "text": "low quality, bad anatomy, extra digits, missing digits, extra limbs, missing limbs", + "clip": [ + "73", + 0 + ] + }, + "class_type": "CLIPTextEncode", + "_meta": { + "title": "CLIP Text Encode (Prompt)" + } + }, + "72": { + "inputs": { + "unet_name": "z_image_turbo_fp32-Q4_K_M.gguf" + }, + "class_type": "UnetLoaderGGUF", + "_meta": { + "title": "UNET Loader (GGUF) ⚡️" + } + }, + "73": { + "inputs": { + "clip_name": "Qwen3_4B.Q4_K_M.gguf", + "type": "lumina2" + }, + "class_type": "CLIPLoaderGGUF", + "_meta": { + "title": "CLIP Loader (GGUF) ⚡️" + } + }, + "81": { + "inputs": { + "pixels": [ + "82", + 0 + ], + "vae": [ + "63", + 0 + ] + }, + "class_type": "VAEEncode", + "_meta": { + "title": "VAE Encode" + } + }, + "82": { + "inputs": { + "image": "ComfyUI_00034_.png [output]", + "refresh": "refresh" + }, + "class_type": "LoadImageOutput", + "_meta": { + "title": "Load Image (from Outputs)" + } + } +} \ No newline at end of file diff --git a/workflows/img2img.json b/workflows/img2img.json new file mode 100644 index 0000000..3708c24 --- /dev/null +++ b/workflows/img2img.json @@ -0,0 +1,128 @@ +{ + "3": { + "inputs": { + "add_noise": "enable", + "noise_seed": 0, + "steps": 8, + "cfg": 1, + "sampler_name": "euler", + "scheduler": "simple", + "start_at_step": 0, + "end_at_step": 10000, + "return_with_leftover_noise": "disable", + "model": [ + "19", + 0 + ], + "positive": [ + "25", + 0 + ], + "negative": [ + "10", + 0 + ], + "latent_image": [ + "15", + 0 + ] + }, + "class_type": "KSamplerAdvanced", + "_meta": { + "title": "KSampler (Advanced)" + } + }, + "10": { + "inputs": { + "text": "__NEGATIVE__", + "clip": [ + "19", + 1 + ] + }, + "class_type": "CLIPTextEncode", + "_meta": { + "title": "CLIP Text Encode (Prompt)" + } + }, + "11": { + "inputs": { + "samples": [ + "3", + 0 + ], + "vae": [ + "19", + 2 + ] + }, + "class_type": "VAEDecode", + "_meta": { + "title": "VAE Decode" + } + }, + "15": { + "inputs": { + "width": 1024, + "height": 1024, + "batch_size": 1 + }, + "class_type": "EmptyLatentImage", + "_meta": { + "title": "Empty Latent Image" + } + }, + "19": { + "inputs": { + "unet_name": "Krea2_turbo_uncensored_edit_v1.1-Q4_K_M.gguf", + "clip_name": "qwen3vl_4b_fp8_scaled.safetensors", + "vae_name": "qwen_image_vae.safetensors" + }, + "class_type": "Krea2ModelLoader", + "_meta": { + "title": "Krea2 Model Loader ⚡️" + } + }, + "20": { + "class_type": "LoadImage", + "inputs": { + "image": "__INPUT_IMAGE__" + }, + "_meta": { + "title": "Load Image" + } + }, + "24": { + "inputs": { + "filename_prefix": "Mpako", + "format": "png", + "format.bit_depth": "8-bit", + "format.input_color_space": "sRGB", + "images": [ + "11", + 0 + ] + }, + "class_type": "SaveImageAdvanced", + "_meta": { + "title": "Save Image (Advanced)" + } + }, + "25": { + "inputs": { + "prompt": "__PROMPT__", + "clip": [ + "19", + 1 + ], + "image": [ + "20", + 0 + ] + }, + "class_type": "TextEncodeQwenImageEdit", + "_meta": { + "title": "Text Encode Qwen Image Edit" + } + } +} diff --git a/workflows/mapping.json b/workflows/mapping.json new file mode 100644 index 0000000..61338b6 --- /dev/null +++ b/workflows/mapping.json @@ -0,0 +1,16 @@ +{ + "txt2img": { + "positive": "67", + "negative": "71", + "seed": "70", + "latent": "81", + "save": "9" + }, + "img2img": { + "positive": "25", + "negative": "10", + "seed": "3", + "input_image": "20", + "save": "24" + } +} diff --git a/workflows/txt2img.json b/workflows/txt2img.json new file mode 100644 index 0000000..255f0ef --- /dev/null +++ b/workflows/txt2img.json @@ -0,0 +1,140 @@ +{ + "9": { + "inputs": { + "filename_prefix": "Mpako", + "images": [ + "65", + 0 + ] + }, + "class_type": "SaveImage", + "_meta": { + "title": "Save Image" + } + }, + "63": { + "inputs": { + "vae_name": "ae.safetensors" + }, + "class_type": "VAELoader", + "_meta": { + "title": "Load VAE" + } + }, + "65": { + "inputs": { + "samples": [ + "70", + 0 + ], + "vae": [ + "63", + 0 + ] + }, + "class_type": "VAEDecode", + "_meta": { + "title": "VAE Decode" + } + }, + "67": { + "inputs": { + "text": "__PROMPT__", + "clip": [ + "73", + 0 + ] + }, + "class_type": "CLIPTextEncode", + "_meta": { + "title": "CLIP Text Encode (Prompt)" + } + }, + "69": { + "inputs": { + "shift": 3, + "sampling": "flow", + "model": [ + "72", + 0 + ] + }, + "class_type": "ModelSamplingAuraFlow", + "_meta": { + "title": "ModelSamplingAuraFlow" + } + }, + "70": { + "inputs": { + "seed": 0, + "steps": 8, + "cfg": 1, + "sampler_name": "euler", + "scheduler": "sgm_uniform", + "denoise": 1.0, + "model": [ + "69", + 0 + ], + "positive": [ + "67", + 0 + ], + "negative": [ + "71", + 0 + ], + "latent_image": [ + "81", + 0 + ] + }, + "class_type": "KSampler", + "_meta": { + "title": "KSampler" + } + }, + "71": { + "inputs": { + "text": "__NEGATIVE__", + "clip": [ + "73", + 0 + ] + }, + "class_type": "CLIPTextEncode", + "_meta": { + "title": "CLIP Text Encode (Prompt)" + } + }, + "72": { + "inputs": { + "unet_name": "z_image_turbo_fp32-Q4_K_M.gguf" + }, + "class_type": "UnetLoaderGGUF", + "_meta": { + "title": "UNET Loader (GGUF) ⚡️" + } + }, + "73": { + "inputs": { + "clip_name": "Qwen3_4B.Q4_K_M.gguf", + "type": "lumina2" + }, + "class_type": "CLIPLoaderGGUF", + "_meta": { + "title": "CLIP Loader (GGUF) ⚡️" + } + }, + "81": { + "class_type": "EmptySD3LatentImage", + "inputs": { + "width": 1024, + "height": 1024, + "batch_size": 1 + }, + "_meta": { + "title": "Empty Latent" + } + } +}