428 lines
16 KiB
Python
428 lines
16 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 PIL import Image, ImageOps
|
||
|
||
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, fit_latent_size
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
def _image_size(path: str) -> tuple[int, int]:
|
||
with Image.open(path) as image:
|
||
oriented = ImageOps.exif_transpose(image)
|
||
return fit_latent_size(*oriented.size)
|
||
|
||
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._premium: deque[int] = deque()
|
||
self._normal: deque[int] = deque()
|
||
self._next_premium = True
|
||
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:
|
||
target = self._premium if request.is_premium else self._normal
|
||
target.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,
|
||
payment_charge_id: str | None = None,
|
||
is_premium: bool = False,
|
||
context_base_prompt: str | None = None,
|
||
context_prompt: str | None = None,
|
||
) -> 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._premium) + len(self._normal) >= 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,
|
||
payment_charge_id=payment_charge_id,
|
||
is_premium=is_premium,
|
||
context_base_prompt=context_base_prompt,
|
||
context_prompt=context_prompt,
|
||
)
|
||
target = self._premium if is_premium else self._normal
|
||
target.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 queue_counts(self) -> dict[str, int]:
|
||
return await self._db.queue_counts()
|
||
|
||
async def position(self, request_id: int) -> int:
|
||
request = await self._db.get_request(request_id)
|
||
if request is not None and request.status == "running":
|
||
return 0
|
||
|
||
counts = await self._db.queue_counts()
|
||
running = max(
|
||
0,
|
||
sum(counts.values()) - len(self._premium) - len(self._normal),
|
||
)
|
||
premium = list(self._premium)
|
||
normal = list(self._normal)
|
||
prefer_premium = self._next_premium
|
||
premium_index = 0
|
||
normal_index = 0
|
||
ahead = running
|
||
while premium_index < len(premium) or normal_index < len(normal):
|
||
if prefer_premium and premium_index < len(premium):
|
||
current = premium[premium_index]
|
||
premium_index += 1
|
||
taken_premium = True
|
||
elif not prefer_premium and normal_index < len(normal):
|
||
current = normal[normal_index]
|
||
normal_index += 1
|
||
taken_premium = False
|
||
elif premium_index < len(premium):
|
||
current = premium[premium_index]
|
||
premium_index += 1
|
||
taken_premium = True
|
||
else:
|
||
current = normal[normal_index]
|
||
normal_index += 1
|
||
taken_premium = False
|
||
if current == request_id:
|
||
return ahead
|
||
ahead += 1
|
||
prefer_premium = not taken_premium
|
||
return ahead
|
||
|
||
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._next_premium:
|
||
if self._premium:
|
||
return self._take(self._premium, premium=True)
|
||
if self._normal:
|
||
return self._take(self._normal, premium=False)
|
||
else:
|
||
if self._normal:
|
||
return self._take(self._normal, premium=False)
|
||
if self._premium:
|
||
return self._take(self._premium, premium=True)
|
||
self._wake.clear()
|
||
await self._wake.wait()
|
||
return None
|
||
|
||
def _take(self, queue: deque[int], *, premium: bool) -> int:
|
||
request_id = queue.popleft()
|
||
self._next_premium = not premium
|
||
return request_id
|
||
|
||
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,
|
||
previous_prompt=request.context_prompt,
|
||
base_prompt=request.context_base_prompt,
|
||
)
|
||
except Exception:
|
||
logger.exception(
|
||
"Не удалось улучшить промпт запроса %s", request_id
|
||
)
|
||
if request.context_prompt:
|
||
fallback = (
|
||
f"{request.context_prompt}. {request.prompt_text}"
|
||
)
|
||
await self._db.set_enhanced_prompt(
|
||
request_id, fallback
|
||
)
|
||
request = await self._refresh(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)
|
||
elif request.context_prompt:
|
||
combined = f"{request.context_prompt}. {request.prompt_text}"
|
||
await self._db.set_enhanced_prompt(request_id, combined)
|
||
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)
|
||
|
||
try:
|
||
await self._comfy.delete_history(prompt_id)
|
||
except Exception:
|
||
logger.warning(
|
||
"Не удалось очистить историю ComfyUI для %s",
|
||
prompt_id,
|
||
exc_info=True,
|
||
)
|
||
|
||
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)
|
||
width, height = await asyncio.to_thread(
|
||
_image_size, request.input_file_path
|
||
)
|
||
user = await self._db.get_user(request.user_id)
|
||
quality = (
|
||
user.generation_quality if user is not None else "standard"
|
||
)
|
||
steps = user.generation_steps if user is not None else 12
|
||
unet_name = (
|
||
self._settings.krea2_q6_model
|
||
if quality == "improved"
|
||
else self._settings.krea2_q4_model
|
||
)
|
||
return self._workflows.build_img2img(
|
||
prompt=prompt,
|
||
negative=negative,
|
||
seed=seed,
|
||
image_name=image_name,
|
||
width=width,
|
||
height=height,
|
||
steps=steps,
|
||
unet_name=unet_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)
|