208 lines
6.7 KiB
Python
208 lines
6.7 KiB
Python
"""出图服务：调用适配器、落盘、记录 generations。"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import math
|
|
import re
|
|
import time
|
|
import uuid
|
|
from datetime import datetime
|
|
from typing import Any
|
|
|
|
from pydantic import ValidationError
|
|
|
|
from .. import config, db
|
|
from ..providers import ProviderError, get_adapter
|
|
from ..providers.agnes import size_to_tier_and_ratio
|
|
from ..schemas import GenerateIn
|
|
from . import providers as provider_service
|
|
|
|
WEB_PROFILE = "user"
|
|
_SAFE = re.compile(r"[^A-Za-z0-9._-]+")
|
|
|
|
|
|
def _safe_model(model: str) -> str:
|
|
return _SAFE.sub("-", model).strip("-") or "image"
|
|
|
|
|
|
def record_generation(
|
|
*,
|
|
provider_id: str,
|
|
provider_name: str,
|
|
model: str,
|
|
profile_name: str,
|
|
prompt: str,
|
|
size: str,
|
|
ratio: str,
|
|
n: int,
|
|
status: str,
|
|
error: str,
|
|
files: list[str],
|
|
duration_ms: int,
|
|
) -> str:
|
|
gid = "gen_" + uuid.uuid4().hex[:12]
|
|
with db.session() as conn:
|
|
conn.execute(
|
|
"INSERT INTO generations (id,provider_id,provider_name,model,profile_name,prompt,size,"
|
|
"ratio,n,status,error,files,duration_ms,created_at) VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?,?)",
|
|
(
|
|
gid,
|
|
provider_id,
|
|
provider_name,
|
|
model,
|
|
profile_name,
|
|
prompt,
|
|
size,
|
|
ratio,
|
|
n,
|
|
status,
|
|
error,
|
|
json.dumps(files, ensure_ascii=False),
|
|
duration_ms,
|
|
datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
|
|
),
|
|
)
|
|
return gid
|
|
|
|
|
|
def list_generations() -> list[dict[str, Any]]:
|
|
with db.session() as conn:
|
|
rows = conn.execute("SELECT * FROM generations ORDER BY created_at DESC, rowid DESC").fetchall()
|
|
out: list[dict[str, Any]] = []
|
|
for row in rows:
|
|
try:
|
|
files = json.loads(row["files"] or "[]")
|
|
except ValueError:
|
|
files = []
|
|
out.append(
|
|
{
|
|
"id": row["id"],
|
|
"provider_id": row["provider_id"],
|
|
"provider_name": row["provider_name"],
|
|
"model": row["model"],
|
|
"profile_name": row["profile_name"],
|
|
"prompt": row["prompt"],
|
|
"size": row["size"],
|
|
"ratio": row["ratio"],
|
|
"n": row["n"],
|
|
"status": row["status"],
|
|
"error": row["error"],
|
|
"files": files,
|
|
"duration_ms": row["duration_ms"],
|
|
"created_at": row["created_at"],
|
|
}
|
|
)
|
|
return out
|
|
|
|
|
|
def generate(
|
|
*,
|
|
provider_id: str,
|
|
model: str,
|
|
prompt: str,
|
|
size: str = "1024x1024",
|
|
ratio: str = "",
|
|
n: int = 1,
|
|
images: list[str] | None = None,
|
|
profile_name: str | None = None,
|
|
) -> dict[str, Any]:
|
|
try:
|
|
payload = GenerateIn(provider_id=provider_id, model=model, prompt=prompt, size=size,
|
|
ratio=ratio, n=n, images=images or [], profile_name=profile_name)
|
|
except ValidationError as exc:
|
|
raise ProviderError(str(exc), status=400) from exc
|
|
prompt = payload.prompt
|
|
size = payload.size
|
|
profile = payload.profile_name or WEB_PROFILE
|
|
cfg = provider_service.provider_config(provider_id, require_enabled=True)
|
|
selected = provider_service.get_model(f"{provider_id}:{model}")
|
|
if not selected["enabled"]:
|
|
raise ProviderError("模型已停用", status=400)
|
|
provider_name = cfg.name
|
|
adapter = get_adapter(cfg)
|
|
if images and not adapter.capabilities.supports_image_input:
|
|
raise ProviderError("该供应商不支持参考图输入", status=400)
|
|
if ratio and cfg.adapter != "agnes":
|
|
raise ProviderError("该供应商不支持 ratio，请使用 WxH 尺寸", status=400)
|
|
if cfg.adapter == "agnes":
|
|
_, ratio = size_to_tier_and_ratio(size, ratio)
|
|
elif "x" in size:
|
|
width, height = map(int, size.split("x"))
|
|
divisor = math.gcd(width, height)
|
|
ratio = f"{width // divisor}:{height // divisor}"
|
|
|
|
started = time.monotonic()
|
|
files: list[str] = []
|
|
folder = config.OUTPUT_DIR / profile
|
|
if not folder.resolve().is_relative_to(config.OUTPUT_DIR.resolve()):
|
|
raise ProviderError("画廊路径超出输出目录", status=400)
|
|
images_out: list[dict[str, Any]] = []
|
|
try:
|
|
results = adapter.generate(model, prompt, size, ratio, n, images or [])
|
|
if not results:
|
|
raise ProviderError("上游未返回图片")
|
|
folder.mkdir(parents=True, exist_ok=True)
|
|
stamp = uuid.uuid4().hex[:8]
|
|
for idx, item in enumerate(results):
|
|
if item.url and not item.data:
|
|
item = adapter.download(item.url)
|
|
if not item.data:
|
|
raise ProviderError("上游返回空图片")
|
|
suffix = item.filename if item.filename in config.IMAGE_SUFFIXES else ".png"
|
|
filename = f"{_safe_model(model)}-{stamp}{'' if len(results) == 1 else f'-{idx + 1}'}{suffix}"
|
|
path = folder / filename
|
|
files.append(filename)
|
|
path.write_bytes(item.data)
|
|
images_out.append({"filename": filename, "url": f"/output/{profile}/{filename}",
|
|
"bytes": path.stat().st_size})
|
|
except (ProviderError, OSError) as exc:
|
|
for filename in files:
|
|
(folder / filename).unlink(missing_ok=True)
|
|
message = exc.message if isinstance(exc, ProviderError) else "保存图片失败"
|
|
duration = int((time.monotonic() - started) * 1000)
|
|
record_generation(
|
|
provider_id=provider_id,
|
|
provider_name=provider_name,
|
|
model=model,
|
|
profile_name=profile,
|
|
prompt=prompt,
|
|
size=size,
|
|
ratio=ratio,
|
|
n=n,
|
|
status="failed",
|
|
error=message,
|
|
files=[],
|
|
duration_ms=duration,
|
|
)
|
|
raise ProviderError(message, status=exc.status if isinstance(exc, ProviderError) else 500) from exc
|
|
|
|
duration = int((time.monotonic() - started) * 1000)
|
|
gid = record_generation(
|
|
provider_id=provider_id,
|
|
provider_name=provider_name,
|
|
model=model,
|
|
profile_name=profile,
|
|
prompt=prompt,
|
|
size=size,
|
|
ratio=ratio or "1:1",
|
|
n=n,
|
|
status="success",
|
|
error="",
|
|
files=files,
|
|
duration_ms=duration,
|
|
)
|
|
|
|
return {
|
|
"id": gid,
|
|
"provider": provider_id,
|
|
"provider_name": provider_name,
|
|
"model": model,
|
|
"profile_name": profile,
|
|
"size": size,
|
|
"ratio": ratio or "1:1",
|
|
"images": images_out,
|
|
"output_dir": f"output/{profile}/",
|
|
"created_at": datetime.now().strftime("%Y-%m-%d %H:%M"),
|
|
}
|