Базовая версия
Build and push Docker image / docker (push) Successful in 22s
Details
Build and push Docker image / docker (push) Successful in 22s
Details
This commit is contained in:
commit
28f7da349e
|
|
@ -0,0 +1,15 @@
|
|||
.git
|
||||
.gitea
|
||||
.gitignore
|
||||
.env
|
||||
.env.*
|
||||
!.env.example
|
||||
storage/
|
||||
data/
|
||||
__pycache__/
|
||||
*.py[cod]
|
||||
.venv/
|
||||
venv/
|
||||
.pytest_cache/
|
||||
tests/
|
||||
README.md
|
||||
|
|
@ -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
|
||||
|
|
@ -0,0 +1,4 @@
|
|||
* text=auto eol=lf
|
||||
|
||||
*.bat text eol=crlf
|
||||
*.cmd text eol=crlf
|
||||
|
|
@ -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
|
||||
|
|
@ -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/
|
||||
|
|
@ -0,0 +1,44 @@
|
|||
# AGENT.md — правила работы
|
||||
|
||||
## Ответы
|
||||
|
||||
- **Коротко и по делу.** Без вступлений, извинений и повторов.
|
||||
- Сначала результат, потом минимальные пояснения.
|
||||
- Язык — русский, если не попросили иначе.
|
||||
- Не пересказывать очевидное; ссылаться на конкретные файлы.
|
||||
- Нужны уточнения — задать 1–3 конкретных вопроса.
|
||||
- Не выдумывать факты: проверить или спросить.
|
||||
|
||||
## Проект
|
||||
|
||||
- Корень: `D:\!Projects\MpakoImageAIbot`.
|
||||
- Стек: Python 3.10+, python-telegram-bot 21+, httpx, aiosqlite, SQLite, Docker.
|
||||
- Требования: `plan.md`. План работ: `checklist.md`; отмечать выполненные пункты.
|
||||
- Следовать структуре, именованию и стилю существующего кода.
|
||||
- Секреты — только в `.env`. Не логировать токены, не коммитить `.env`, `storage/`, `data/`.
|
||||
|
||||
## Git / Gitea
|
||||
|
||||
- Remote: `ssh://git@home:2222/mpakobecex/MpakoImageAIbot.git`
|
||||
(алиас `home` → `192.168.31.158`).
|
||||
- Основная ветка — `main`.
|
||||
- Перед коммитом: `git diff`, тесты/линтер, если есть.
|
||||
- Коммиты атомарные: `feat:`, `fix:`, `docs:`, `chore:`.
|
||||
- Не делать force push и не переписывать историю.
|
||||
- Push — после успешной проверки.
|
||||
|
||||
## Работа с кодом
|
||||
|
||||
- Сначала прочитать существующий код, затем менять.
|
||||
- Изменения минимальные и точечные; без попутного рефакторинга.
|
||||
- Не удалять и не перезаписывать пользовательские файлы без подтверждения.
|
||||
- Сетевые вызовы — async, event loop не блокировать.
|
||||
- Ошибки обрабатывать явно: понятное сообщение пользователю + лог.
|
||||
- После изменений проверять запуск/тесты и сообщать результат.
|
||||
- Не оставлять проект в нерабочем состоянии.
|
||||
|
||||
## Безопасность
|
||||
|
||||
- Не выводить содержимое `.env` и приватные ключи.
|
||||
- Не выполнять деструктивные команды без подтверждения.
|
||||
- Не добавлять зависимости без необходимости.
|
||||
|
|
@ -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"]
|
||||
|
|
@ -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:<sha>`
|
||||
- локально: `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`.
|
||||
|
|
@ -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()
|
||||
|
|
@ -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.<ext>`, `storage/{user_id}/{request_id}_output.<ext>`.
|
||||
- [ ] Безопасные имена (только сгенерированные, не из пользовательского ввода), защита от 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` и `:<sha>` в 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.
|
||||
- Админ-панель.
|
||||
|
|
@ -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)
|
||||
|
|
@ -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)
|
||||
|
|
@ -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()
|
||||
|
|
@ -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
|
||||
|
|
@ -0,0 +1,44 @@
|
|||
Ты — senior Python-разработчик. Нужно написать Telegram-бота для генерации изображений через удалённый ComfyUI API.
|
||||
|
||||
## Стек
|
||||
- Python 3.10+
|
||||
- python-telegram-bot (v21+, async) [citation:3]
|
||||
- SQLite (через aiosqlite или стандартный sqlite3 с осторожностью)
|
||||
- requests / httpx для вызова ComfyUI API
|
||||
- Docker для финальной упаковки
|
||||
|
||||
## Функционал (два режима)
|
||||
1. **Text-to-Image**: пользователь отправляет текст → бот шлёт POST на ComfyUI `/prompt` с подставленным промптом [citation:14].
|
||||
2. **Image-to-Image**: пользователь отправляет фото + текст → бот скачивает фото, загружает его в ComfyUI (через `/upload/image` или временную папку), подставляет в воркфлоу как input image, шлёт POST.
|
||||
|
||||
## Очередь (критично)
|
||||
- Один пользователь = одна активная задача. Если он пишет ещё раз, пока висит его задача — бот отвечает «ты уже в очереди».
|
||||
- Нужна глобальная очередь. Показывай:
|
||||
- Если в очереди: «Перед тобой X человек».
|
||||
- Если генерация идёт: «Генерация… Y%» (прогресс брать из WebSocket ComfyUI, события `progress` [citation:25][citation:35]).
|
||||
|
||||
## Хранение (SQLite)
|
||||
Таблица `requests`:
|
||||
- id (PK)
|
||||
- user_id (telegram id)
|
||||
- timestamp (ISO)
|
||||
- prompt_text
|
||||
- input_file_path (nullable)
|
||||
- output_file_path (nullable)
|
||||
- status (queued/running/done/failed)
|
||||
- comfy_prompt_id (для отслеживания)
|
||||
|
||||
Файлы (входные фото и результаты) сохраняй локально в папку `./storage/{user_id}/`. В БД только пути.
|
||||
|
||||
## Вызов ComfyUI
|
||||
- База: `http://IP_WINDOWS:PORT`
|
||||
- POST `/prompt` → получаешь `prompt_id` [citation:2][citation:14].
|
||||
- Опрос статуса: либо через GET `/history/{prompt_id}` (проще, но медленнее), либо через WebSocket на `/ws` (для прогресса в реальном времени) [citation:30][citation:35].
|
||||
- После завершения скачай изображение через `/view?filename=...`.
|
||||
|
||||
## Что сделать сейчас
|
||||
1. Создай структуру проекта: `bot.py`, `comfy_client.py`, `db.py`, `queue_manager.py`.
|
||||
2. Напиши `comfy_client.py` с функциями submit_workflow (POST /prompt) и get_result (скачивание).
|
||||
3. Напиши `db.py` с init_db() и CRUD для requests.
|
||||
4. Напиши `bot.py` с хендлерами для текста и фото, но очередь пока сделай простейшей (asyncio.Queue).
|
||||
5. Добавь `requirements.txt` и `Dockerfile`.
|
||||
|
|
@ -0,0 +1,3 @@
|
|||
[pytest]
|
||||
asyncio_mode = auto
|
||||
testpaths = tests
|
||||
|
|
@ -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)
|
||||
|
|
@ -0,0 +1,3 @@
|
|||
-r requirements.txt
|
||||
pytest>=8,<9
|
||||
pytest-asyncio>=0.24,<1
|
||||
|
|
@ -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
|
||||
|
|
@ -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]
|
||||
|
|
@ -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
|
||||
|
|
@ -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()
|
||||
|
|
@ -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)
|
||||
|
|
@ -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
|
||||
|
|
@ -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"
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -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)"
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -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"
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -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"
|
||||
}
|
||||
}
|
||||
|
|
@ -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"
|
||||
}
|
||||
}
|
||||
}
|
||||
Loading…
Reference in New Issue