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 moderation import PromptBlocked, is_prompt_allowed from prompt_enhancer import PromptEnhancer from workflows import WorkflowBuilder logger = logging.getLogger(__name__) EventCallback = Callable[["QueueEvent"], Awaitable[None] | None] @dataclass(frozen=True) class QueueEvent: kind: str # started | enhancing | 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, enhancer: PromptEnhancer | None = None, ) -> None: self._db = db self._comfy = comfy self._workflows = workflows self._settings = settings self._on_event = on_event self._enhancer = enhancer 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, has_person: bool = False, is_admin: bool = False, ) -> tuple[Request, int]: async with self._lock: if self._stopping: raise QueueStopped() if self._settings.moderation_enabled and not is_admin: allowed, reason = is_prompt_allowed( prompt, mode=mode, has_person=has_person, extra_public_figures=self._settings.blocked_public_figures, ) if not allowed: await self._db.log_blocked( user_id=user_id, prompt=prompt, reason=reason or "unknown", mode=mode, has_person=has_person, ) raise PromptBlocked(reason or "unknown") 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: if self._enhancer is not None and self._enhancer.enabled: await self._emit( QueueEvent(kind="enhancing", request=request) ) try: enhanced = await self._enhancer.enhance( request.prompt_text, request.mode ) except Exception: logger.exception( "Не удалось улучшить промпт запроса %s", request_id ) else: if enhanced and enhanced != request.prompt_text: await self._db.set_enhanced_prompt(request_id, enhanced) request = await self._refresh(request_id) 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 prompt = request.enhanced_prompt or request.prompt_text if request.mode == "t2i": return self._workflows.build_txt2img( prompt=prompt, 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=prompt, 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)