Files
ai-image/app/services/gallery.py
T

146 lines
5.0 KiB
Python

"""画廊服务:扫描 output/ 目录结构。
output/ 下的每个一级子目录 = 一个画廊,目录名 = 谁生成的图片(profile_name)。
"""
from __future__ import annotations
import json
from datetime import datetime
from pathlib import Path
from typing import Any
from .. import config, db
from ..providers import ProviderError
def _folder(profile_name: str) -> Path:
if not config.PROFILE_RE.fullmatch(profile_name):
raise ProviderError("无效的画廊名称", status=400)
root = config.OUTPUT_DIR.resolve()
folder = (root / profile_name).resolve()
if not folder.is_relative_to(root):
raise ProviderError("画廊路径超出输出目录", status=400)
return folder
def _meta_by_file() -> dict[tuple[str, str], dict[str, Any]]:
"""从 generations 表建立 (profile, filename) → 元信息的索引。"""
index: dict[tuple[str, str], dict[str, Any]] = {}
with db.session() as conn:
rows = conn.execute("SELECT * FROM generations ORDER BY created_at DESC").fetchall()
for row in rows:
try:
files = json.loads(row["files"] or "[]")
except ValueError:
files = []
for name in files:
key = (row["profile_name"], name)
index.setdefault(
key,
{
"model": row["model"],
"provider_name": row["provider_name"],
"created_at": row["created_at"],
"prompt": row["prompt"],
},
)
return index
def _images_in(folder: Path) -> list[Path]:
return sorted(
(p for p in folder.iterdir() if p.is_file() and p.suffix.lower() in config.IMAGE_SUFFIXES),
key=lambda p: p.stat().st_mtime,
reverse=True,
)
def list_galleries() -> list[dict[str, Any]]:
config.ensure_dirs()
meta = _meta_by_file()
galleries: list[dict[str, Any]] = []
for folder in sorted(config.OUTPUT_DIR.iterdir()):
if not folder.is_dir():
continue
images = _images_in(folder)
if not images:
continue
total_bytes = 0
models: list[str] = []
latest = 0.0
for img in images:
stat = img.stat()
total_bytes += stat.st_size
latest = max(latest, stat.st_mtime)
info = meta.get((folder.name, img.name))
if info and info["model"] and info["model"] not in models:
models.append(info["model"])
cover = images[0]
galleries.append(
{
"name": folder.name,
"count": len(images),
"bytes": total_bytes,
"updated_at": datetime.fromtimestamp(latest).strftime("%Y-%m-%d %H:%M"),
"cover": f"/output/{folder.name}/{cover.name}",
"models": models,
}
)
galleries.sort(key=lambda g: g["updated_at"], reverse=True)
return galleries
def list_gallery(profile_name: str) -> list[dict[str, Any]]:
config.ensure_dirs()
folder = _folder(profile_name)
if not folder.is_dir():
return []
meta = _meta_by_file()
items: list[dict[str, Any]] = []
for img in _images_in(folder):
stat = img.stat()
info = meta.get((profile_name, img.name), {})
items.append(
{
"profile_name": profile_name,
"filename": img.name,
"url": f"/output/{profile_name}/{img.name}",
"bytes": stat.st_size,
"model": info.get("model", "—"),
"provider_name": info.get("provider_name", "—"),
"prompt": info.get("prompt", ""),
"created_at": info.get("created_at")
or datetime.fromtimestamp(stat.st_mtime).strftime("%Y-%m-%d %H:%M"),
}
)
return items
def delete_gallery_item(profile_name: str, filename: str) -> None:
folder = _folder(profile_name)
if Path(filename).name != filename or "\\" in filename:
raise ProviderError("无效的文件名", status=400)
target = folder / filename
if not target.resolve().is_relative_to(folder):
raise ProviderError("图片路径超出画廊目录", status=400)
if not target.is_file():
raise FileNotFoundError(f"文件不存在: {profile_name}/{filename}")
target.unlink()
# 同步 generations 记录
with db.session() as conn:
rows = conn.execute(
"SELECT id, files FROM generations WHERE profile_name=?", (profile_name,)
).fetchall()
for row in rows:
try:
files = json.loads(row["files"] or "[]")
except ValueError:
files = []
if filename in files:
files = [f for f in files if f != filename]
conn.execute(
"UPDATE generations SET files=? WHERE id=?",
(json.dumps(files, ensure_ascii=False), row["id"]),
)