275 lines
9.7 KiB
Python
275 lines
9.7 KiB
Python
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,
|
|
username: 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,
|
|
username=username,
|
|
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)
|