292 lines
10 KiB
Python
292 lines
10 KiB
Python
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import logging
|
||
import math
|
||
import os
|
||
from contextlib import asynccontextmanager
|
||
from datetime import datetime
|
||
from pathlib import Path
|
||
from typing import Annotated, Any
|
||
from urllib.parse import urlencode
|
||
from uuid import uuid4
|
||
|
||
from fastapi import FastAPI, Form, HTTPException, Request
|
||
from fastapi.responses import FileResponse, HTMLResponse, RedirectResponse
|
||
from fastapi.staticfiles import StaticFiles
|
||
from fastapi.templating import Jinja2Templates
|
||
from PIL import Image, ImageOps
|
||
|
||
from config import Settings
|
||
from db import Database
|
||
from db import Request as HistoryRequest
|
||
|
||
BASE_DIR = Path(__file__).resolve().parent
|
||
TEMPLATES_DIR = BASE_DIR / "web" / "templates"
|
||
STATIC_DIR = BASE_DIR / "web" / "static"
|
||
|
||
STATUSES = ("queued", "running", "done", "failed")
|
||
MODES = ("t2i", "i2i")
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
def _safe_storage_path(raw_path: str, storage_root: Path) -> Path | None:
|
||
path = Path(raw_path).resolve()
|
||
try:
|
||
path.relative_to(storage_root)
|
||
except ValueError:
|
||
return None
|
||
return path
|
||
|
||
|
||
def _make_thumbnail(source: Path, destination: Path) -> None:
|
||
destination.parent.mkdir(parents=True, exist_ok=True)
|
||
temp_path = destination.with_name(f"{destination.name}.{uuid4().hex}.tmp")
|
||
try:
|
||
with Image.open(source) as image:
|
||
image = ImageOps.exif_transpose(image)
|
||
image.thumbnail((320, 320))
|
||
if image.mode in {"RGBA", "LA"} or "transparency" in image.info:
|
||
rgba = image.convert("RGBA")
|
||
background = Image.new("RGB", rgba.size, (15, 17, 21))
|
||
background.paste(rgba, mask=rgba.getchannel("A"))
|
||
image = background
|
||
else:
|
||
image = image.convert("RGB")
|
||
image.save(temp_path, format="JPEG", quality=85, optimize=True)
|
||
os.replace(temp_path, destination)
|
||
finally:
|
||
temp_path.unlink(missing_ok=True)
|
||
|
||
|
||
def fmt_dt(value: str | None) -> str:
|
||
if not value:
|
||
return "—"
|
||
try:
|
||
parsed = datetime.fromisoformat(value)
|
||
return parsed.astimezone().strftime("%d.%m.%Y %H:%M")
|
||
except ValueError:
|
||
return value
|
||
|
||
|
||
def create_app(settings: Settings | None = None) -> FastAPI:
|
||
settings = settings or Settings.from_env(require_token=False)
|
||
templates = Jinja2Templates(directory=str(TEMPLATES_DIR))
|
||
templates.env.filters["fmt_dt"] = fmt_dt
|
||
|
||
@asynccontextmanager
|
||
async def lifespan(app: FastAPI):
|
||
db = Database(settings.db_path)
|
||
await db.connect()
|
||
app.state.db = db
|
||
yield
|
||
await db.close()
|
||
|
||
app = FastAPI(
|
||
title="MpakoImageAIbot history",
|
||
docs_url=None,
|
||
redoc_url=None,
|
||
lifespan=lifespan,
|
||
)
|
||
app.state.settings = settings
|
||
app.mount("/static", StaticFiles(directory=str(STATIC_DIR)), name="static")
|
||
|
||
@app.get("/healthz")
|
||
async def healthz() -> dict[str, bool]:
|
||
return {"ok": True}
|
||
|
||
@app.get("/users", response_class=HTMLResponse)
|
||
async def users(request: Request, q: str = ""):
|
||
db: Database = request.app.state.db
|
||
rows = await db.list_users_admin(1000)
|
||
needle = q.strip().lower().lstrip("@")
|
||
if needle:
|
||
rows = [
|
||
user
|
||
for user in rows
|
||
if needle in str(user.user_id)
|
||
or (user.username and needle in user.username.lower())
|
||
]
|
||
return templates.TemplateResponse(
|
||
request=request,
|
||
name="users.html",
|
||
context={"users": rows, "q": q},
|
||
)
|
||
|
||
@app.post("/users/{user_id}/credits")
|
||
async def user_credits(
|
||
request: Request,
|
||
user_id: int,
|
||
credits: Annotated[int, Form()],
|
||
) -> RedirectResponse:
|
||
db: Database = request.app.state.db
|
||
if not await db.set_credits(user_id, max(0, credits)):
|
||
raise HTTPException(status_code=404, detail="Пользователь не найден")
|
||
return RedirectResponse("/users", status_code=303)
|
||
|
||
@app.post("/users/{user_id}/admin")
|
||
async def user_admin(request: Request, user_id: int) -> RedirectResponse:
|
||
db: Database = request.app.state.db
|
||
user = await db.get_user(user_id)
|
||
if user is None:
|
||
raise HTTPException(status_code=404, detail="Пользователь не найден")
|
||
await db.set_admin(user_id, not user.is_admin)
|
||
return RedirectResponse("/users", status_code=303)
|
||
|
||
@app.post("/users/{user_id}/block")
|
||
async def user_block(request: Request, user_id: int) -> RedirectResponse:
|
||
db: Database = request.app.state.db
|
||
user = await db.get_user(user_id)
|
||
if user is None:
|
||
raise HTTPException(status_code=404, detail="Пользователь не найден")
|
||
await db.set_blocked(user_id, not user.is_blocked)
|
||
return RedirectResponse("/users", status_code=303)
|
||
|
||
@app.get("/", response_class=HTMLResponse)
|
||
async def index(
|
||
request: Request,
|
||
user_id: str = "",
|
||
status: str = "",
|
||
mode: str = "",
|
||
page: int = 1,
|
||
):
|
||
db: Database = request.app.state.db
|
||
settings: Settings = request.app.state.settings
|
||
|
||
error = ""
|
||
uid: int | None = None
|
||
raw_user_id = user_id.strip()
|
||
if raw_user_id:
|
||
if raw_user_id.isdigit():
|
||
uid = int(raw_user_id)
|
||
else:
|
||
error = "user_id должен состоять только из цифр."
|
||
|
||
status = status if status in STATUSES else ""
|
||
mode = mode if mode in MODES else ""
|
||
|
||
total = await db.count_requests(user_id=uid, status=status or None, mode=mode or None)
|
||
page_size = settings.web_page_size
|
||
pages = max(1, math.ceil(total / page_size))
|
||
page = max(1, min(page, pages))
|
||
offset = (page - 1) * page_size
|
||
|
||
rows = await db.list_requests(
|
||
user_id=uid,
|
||
status=status or None,
|
||
mode=mode or None,
|
||
limit=page_size,
|
||
offset=offset,
|
||
)
|
||
users = await db.list_users()
|
||
stats = await db.stats()
|
||
|
||
base_filters = {
|
||
key: value
|
||
for key, value in (
|
||
("user_id", raw_user_id),
|
||
("status", status),
|
||
("mode", mode),
|
||
)
|
||
if value
|
||
}
|
||
|
||
def page_url(number: int) -> str:
|
||
params = dict(base_filters)
|
||
if number > 1:
|
||
params["page"] = str(number)
|
||
query = urlencode(params)
|
||
return f"/?{query}" if query else "/"
|
||
|
||
context: dict[str, Any] = {
|
||
"rows": rows,
|
||
"users": users,
|
||
"stats": stats,
|
||
"statuses": STATUSES,
|
||
"modes": MODES,
|
||
"user_id": raw_user_id,
|
||
"status": status,
|
||
"mode": mode,
|
||
"error": error,
|
||
"page": page,
|
||
"pages": pages,
|
||
"total": total,
|
||
"prev_url": page_url(page - 1) if page > 1 else None,
|
||
"next_url": page_url(page + 1) if page < pages else None,
|
||
}
|
||
return templates.TemplateResponse(
|
||
request=request, name="index.html", context=context
|
||
)
|
||
|
||
@app.get("/request/{request_id}", response_class=HTMLResponse)
|
||
async def detail(request: Request, request_id: int):
|
||
db: Database = request.app.state.db
|
||
item = await db.get_request(request_id)
|
||
if item is None:
|
||
raise HTTPException(status_code=404, detail="Запрос не найден")
|
||
return templates.TemplateResponse(
|
||
request=request,
|
||
name="detail.html",
|
||
context={"item": item},
|
||
)
|
||
|
||
@app.get("/media/{request_id}/{kind}")
|
||
async def media(request: Request, request_id: int, kind: str):
|
||
if kind not in {"input", "output"}:
|
||
raise HTTPException(status_code=404, detail="Не найдено")
|
||
db: Database = request.app.state.db
|
||
settings: Settings = request.app.state.settings
|
||
item: HistoryRequest | None = await db.get_request(request_id)
|
||
if item is None:
|
||
raise HTTPException(status_code=404, detail="Запрос не найден")
|
||
|
||
raw_path = item.input_file_path if kind == "input" else item.output_file_path
|
||
if not raw_path:
|
||
raise HTTPException(status_code=404, detail="Файл не найден")
|
||
|
||
storage_root = settings.storage_dir.resolve()
|
||
path = _safe_storage_path(raw_path, storage_root)
|
||
if path is None or not path.is_file():
|
||
raise HTTPException(status_code=404, detail="Файл не найден")
|
||
return FileResponse(path)
|
||
|
||
@app.get("/thumb/{request_id}")
|
||
async def thumbnail(request: Request, request_id: int):
|
||
db: Database = request.app.state.db
|
||
settings: Settings = request.app.state.settings
|
||
item: HistoryRequest | None = await db.get_request(request_id)
|
||
if item is None or not item.output_file_path:
|
||
raise HTTPException(status_code=404, detail="Файл не найден")
|
||
|
||
storage_root = settings.storage_dir.resolve()
|
||
source = _safe_storage_path(item.output_file_path, storage_root)
|
||
if source is None or not source.is_file():
|
||
raise HTTPException(status_code=404, detail="Файл не найден")
|
||
|
||
cache_dir = storage_root / ".thumbnails"
|
||
cache_path = cache_dir / f"{request_id}.jpg"
|
||
if cache_path.is_file() and cache_path.stat().st_mtime >= source.stat().st_mtime:
|
||
return FileResponse(
|
||
cache_path,
|
||
media_type="image/jpeg",
|
||
headers={"Cache-Control": "public, max-age=3600"},
|
||
)
|
||
|
||
try:
|
||
await asyncio.to_thread(_make_thumbnail, source, cache_path)
|
||
except Exception:
|
||
logger.exception("Не удалось создать превью для запроса %s", request_id)
|
||
return FileResponse(source)
|
||
|
||
return FileResponse(
|
||
cache_path,
|
||
media_type="image/jpeg",
|
||
headers={"Cache-Control": "public, max-age=3600"},
|
||
)
|
||
|
||
return app
|
||
|
||
|
||
app = create_app()
|