Joey
321 lines
14 KiB
Python
321 lines
14 KiB
Python
"""套图规划 / 一键生成 / 任务查询 / 结果导出 / 单张 AI 生图。
|
||
|
||
按 docs/v2.1/api.md §2-6 实现：引擎平移自 image-suite-studio（services/planner|generator|
|
||
tasks|watermark + services/prompts），任务存进程内内存表（重启即失效）。
|
||
生成图回调回写 product_assets(group_key='generated')；product_id 缺省时只落盘不回写。
|
||
注：V2.1 阶段本组接口暂不接鉴权（现有鉴权后续可能重做）。
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import io
|
||
import mimetypes
|
||
import uuid
|
||
import zipfile
|
||
|
||
from fastapi import APIRouter, BackgroundTasks, HTTPException
|
||
from fastapi.responses import StreamingResponse
|
||
|
||
from config import get_settings
|
||
from db import get_session_factory
|
||
from models import Product, ProductAsset
|
||
from schemas.suite import (
|
||
PLATFORM_SPECS,
|
||
RIGHTAPI_MODELS,
|
||
SUPPORTED_TYPES,
|
||
TONGYI_MODELS,
|
||
ImageEditSingleRequest,
|
||
ImageEditSingleResponse,
|
||
PlanItemOut,
|
||
PlanResponse,
|
||
SuiteCreateResponse,
|
||
SuiteGenerateRequest,
|
||
SuiteOut,
|
||
SuitePlanRequest,
|
||
TextMaterial,
|
||
resolve_provider,
|
||
)
|
||
from services.generator import _image_size, GENERATORS, run_suite
|
||
from services.planner import generate_plan
|
||
from services.prompts import build_context, build_prompt, type_name
|
||
from services.storage import download_bytes, get_storage, local_path
|
||
from services.tasks import IMG_OK, Task, TaskImage, create_task, get_task
|
||
from api.proxy import guess_referer
|
||
|
||
router = APIRouter(prefix="/api", tags=["suite"])
|
||
|
||
|
||
def texts_to_raw(texts: list[TextMaterial]) -> dict:
|
||
"""前端组装的文本素材 → prompt 上下文用的 raw dict（后写的覆盖先写的）。"""
|
||
raw: dict = {}
|
||
for t in texts:
|
||
if t.kind == "title" and t.content:
|
||
raw["title"] = t.content
|
||
elif t.kind == "price" and t.content:
|
||
raw["price"] = t.content
|
||
elif t.kind == "brand" and t.content:
|
||
raw["brand"] = t.content
|
||
elif t.kind == "params" and t.pairs:
|
||
merged = {p["key"]: p["value"] for p in (raw.get("params") or [])}
|
||
for p in t.pairs:
|
||
merged.setdefault(p["key"], p["value"])
|
||
raw["params"] = [{"key": k, "value": v} for k, v in merged.items()]
|
||
elif t.kind == "selling_point" and t.content:
|
||
raw["sellingPoints"] = t.content
|
||
elif t.kind == "desc" and t.content:
|
||
raw["desc"] = t.content
|
||
elif t.kind == "sales" and t.content:
|
||
raw["sales"] = t.content
|
||
elif t.kind == "shop" and t.content:
|
||
raw["shop"] = t.content
|
||
return raw
|
||
|
||
|
||
def _validate_model(provider_name: str, model: str | None) -> None:
|
||
if provider_name == "tongyi" and model and model not in TONGYI_MODELS:
|
||
raise HTTPException(status_code=400, detail=f"不支持的模型: {model}（tongyi 支持: {TONGYI_MODELS}）")
|
||
if provider_name == "rightapi" and model and model not in RIGHTAPI_MODELS:
|
||
raise HTTPException(status_code=400, detail=f"不支持的模型: {model}（rightapi 支持: {RIGHTAPI_MODELS}）")
|
||
|
||
|
||
# ── 生成图回写商品素材 ────────────────────────────────────────────────────
|
||
|
||
async def _append_generated_asset(product_id: str, image: TaskImage) -> str | None:
|
||
"""把一张生成完成的图追加为 product_assets(generated)，并累加 asset_counts。返回 asset_id。"""
|
||
from sqlalchemy import func, select
|
||
|
||
pid = uuid.UUID(product_id)
|
||
async with get_session_factory()() as db:
|
||
count = await db.scalar(
|
||
select(func.count(ProductAsset.id)).where(
|
||
ProductAsset.product_id == pid,
|
||
ProductAsset.group_key == "generated",
|
||
)
|
||
)
|
||
asset = ProductAsset(
|
||
product_id=pid,
|
||
group_key="generated",
|
||
variant_name=None,
|
||
sort_order=count or 0,
|
||
type="img",
|
||
source_url="",
|
||
stored_url=image.url,
|
||
status="uploaded",
|
||
)
|
||
db.add(asset)
|
||
await db.flush()
|
||
|
||
product = await db.get(Product, pid)
|
||
if product is not None:
|
||
counts = dict(product.asset_counts or {})
|
||
counts["generated"] = int(counts.get("generated") or 0) + 1
|
||
product.asset_counts = counts
|
||
await db.commit()
|
||
return str(asset.id)
|
||
|
||
|
||
# ── 出图方案规划 ──────────────────────────────────────────────────────────
|
||
|
||
@router.post("/suite/plan", response_model=PlanResponse)
|
||
async def plan_suite(req: SuitePlanRequest) -> PlanResponse:
|
||
"""DeepSeek 根据商品信息生成出图方案（同步接口，数秒级）。"""
|
||
product_info = texts_to_raw(req.texts)
|
||
if not product_info.get("title"):
|
||
raise HTTPException(status_code=400, detail="缺少商品标题，无法规划")
|
||
try:
|
||
result = await generate_plan(
|
||
product_info, req.sku_variants, req.image_stats, req.platform,
|
||
requirements=req.requirements,
|
||
)
|
||
except RuntimeError as exc:
|
||
raise HTTPException(status_code=502, detail=str(exc)) from exc
|
||
except Exception as exc: # noqa: BLE001
|
||
raise HTTPException(status_code=502, detail=f"规划失败: {exc}") from exc
|
||
return PlanResponse(summary=result["summary"], items=[PlanItemOut(**i) for i in result["items"]])
|
||
|
||
|
||
# ── 一键生成 ─────────────────────────────────────────────────────────────
|
||
|
||
@router.post("/suite/generate", response_model=SuiteCreateResponse)
|
||
async def generate_suite(req: SuiteGenerateRequest, background: BackgroundTasks) -> SuiteCreateResponse:
|
||
if not req.images:
|
||
raise HTTPException(status_code=400, detail="未勾选任何图片，无法生成")
|
||
|
||
# 生成任务列表：方案优先；无方案时按 types（空则默认四种）
|
||
if req.plan:
|
||
jobs: list[dict] = []
|
||
for item in req.plan:
|
||
if item.count <= 0:
|
||
continue
|
||
if item.kind not in SUPPORTED_TYPES:
|
||
raise HTTPException(status_code=400, detail=f"方案项「{item.title}」的图类型不支持: {item.kind}")
|
||
# 同一项多张 → 展开为多任务，第二张起在标题上加序号
|
||
for n in range(item.count):
|
||
jobs.append({
|
||
"kind": item.kind,
|
||
"title": item.title if item.count == 1 else f"{item.title}{n + 1}",
|
||
"detail": item.detail,
|
||
"prompt_hint": item.prompt_hint,
|
||
"variant_name": item.variant_name,
|
||
})
|
||
if not jobs:
|
||
raise HTTPException(status_code=400, detail="方案中所有项的数量都是 0")
|
||
else:
|
||
types = req.types or ["white_bg", "key_features", "lifestyle", "multi_scene"]
|
||
bad = [t for t in types if t not in SUPPORTED_TYPES]
|
||
if bad:
|
||
raise HTTPException(status_code=400, detail=f"不支持的图类型: {bad}")
|
||
jobs = [{"kind": t, "title": type_name(t), "detail": "", "prompt_hint": "", "variant_name": None} for t in types]
|
||
|
||
if req.platform not in PLATFORM_SPECS:
|
||
raise HTTPException(status_code=400, detail=f"不支持的目标平台: {req.platform}（ozon | wb | cn）")
|
||
spec = PLATFORM_SPECS[req.platform]
|
||
|
||
settings = get_settings()
|
||
# 前端只传模型名：已知模型直接路由到对应 provider（如 gpt-image-2-vip → rightapi）
|
||
provider_name = resolve_provider(req.model, None, settings.image_provider)
|
||
_validate_model(provider_name, req.model)
|
||
|
||
product_id = (req.product_id or "").strip() or None
|
||
if product_id:
|
||
try:
|
||
uuid.UUID(product_id)
|
||
except ValueError as exc:
|
||
raise HTTPException(status_code=400, detail="product_id 不是合法 UUID") from exc
|
||
|
||
task = create_task(
|
||
status="pending",
|
||
platform=req.platform,
|
||
lang=spec["lang"],
|
||
ratio=spec["ratio"],
|
||
style_set=req.style_set,
|
||
style_prompt=req.style_prompt,
|
||
requirements=req.requirements,
|
||
provider=provider_name,
|
||
model=req.model,
|
||
total=len(jobs),
|
||
context=texts_to_raw(req.texts),
|
||
plan=jobs,
|
||
watermark=req.watermark.model_dump() if req.watermark else None,
|
||
# 参考图池：main 组优先，其余组按序补充（variant 绑定靠 variant_name 匹配）
|
||
ref_images=[
|
||
{
|
||
"url": i.url,
|
||
"group_key": i.group_key,
|
||
"variant_name": i.variant_name,
|
||
}
|
||
for i in sorted(req.images, key=lambda x: 0 if x.group_key == "main" else 1)
|
||
],
|
||
)
|
||
|
||
# 每张成功即回写 generated 组；未关联商品时仅落存储不回写
|
||
background.add_task(run_suite, task, _append_generated_asset if product_id else None)
|
||
return SuiteCreateResponse(suite_id=task.id)
|
||
|
||
|
||
# ── 任务查询 / 结果导出 ──────────────────────────────────────────────────
|
||
|
||
@router.get("/suites/{suite_id}", response_model=SuiteOut)
|
||
async def get_suite(suite_id: str):
|
||
task = get_task(suite_id)
|
||
if task is None:
|
||
raise HTTPException(status_code=404, detail="任务不存在（服务可能已重启），请重新生成")
|
||
return SuiteOut(
|
||
id=task.id,
|
||
status=task.status,
|
||
style_set=task.style_set,
|
||
platform=task.platform,
|
||
lang=task.lang,
|
||
ratio=task.ratio,
|
||
provider=task.provider,
|
||
model=task.model,
|
||
total=task.total,
|
||
images=[
|
||
{"type_id": i.type_id, "name": i.name, "url": i.url, "status": i.status, "error": i.error}
|
||
for i in task.images
|
||
],
|
||
error=task.error,
|
||
)
|
||
|
||
|
||
@router.get("/suites/{suite_id}/zip")
|
||
async def download_suite_zip(suite_id: str):
|
||
"""把任务内所有成功图打包成 ZIP（中文文件名）。"""
|
||
task = get_task(suite_id)
|
||
if task is None:
|
||
raise HTTPException(status_code=404, detail="任务不存在（服务可能已重启）")
|
||
|
||
buf = io.BytesIO()
|
||
with zipfile.ZipFile(buf, "w", zipfile.ZIP_DEFLATED) as zf:
|
||
seen: set[str] = set()
|
||
for i, img in enumerate([i for i in task.images if i.status == IMG_OK]):
|
||
path = local_path(img.url or "")
|
||
if path is not None:
|
||
data, suffix = path.read_bytes(), path.suffix or ".jpg"
|
||
else:
|
||
try: # 七牛等远程存储：直接拉取公开 URL
|
||
data, _ = await download_bytes(img.url or "")
|
||
except Exception: # noqa: BLE001
|
||
continue
|
||
suffix = ".png" if data[:8] == b"\x89PNG\r\n\x1a\n" else ".jpg"
|
||
filename = img.name or img.type_id
|
||
if filename in seen: # 同类型多张时加序号防覆盖
|
||
filename = f"{filename}-{i + 1}"
|
||
seen.add(filename)
|
||
zf.writestr(f"{filename}{suffix}", data)
|
||
buf.seek(0)
|
||
return StreamingResponse(
|
||
buf,
|
||
media_type="application/zip",
|
||
headers={"Content-Disposition": f"attachment; filename=\"suite-{suite_id}.zip\""},
|
||
)
|
||
|
||
|
||
# ── 单张 AI 生图（V2.1 新增，ISS 无对应实现）──────────────────────────────
|
||
|
||
@router.post("/suite/image-edit", response_model=ImageEditSingleResponse)
|
||
async def suite_image_edit(req: ImageEditSingleRequest) -> ImageEditSingleResponse:
|
||
"""轻量单张再生成：原图 + 一句话要求 + 模型选择。同步接口（单张可长达数分钟）。"""
|
||
settings = get_settings()
|
||
provider_name = resolve_provider(req.model, None, settings.image_provider)
|
||
if provider_name not in GENERATORS:
|
||
raise HTTPException(status_code=400, detail=f"未知 provider: {provider_name}")
|
||
_validate_model(provider_name, req.model)
|
||
|
||
spec = {"lang": "ru", "ratio": "3:4"} # 试算页固定 Ozon 规格（俄文图内文案 · 3:4）
|
||
model = req.model
|
||
is_wan = provider_name == "tongyi" and model.lower().startswith("wan")
|
||
|
||
# 复用套图的 prompt 家族：custom 类型装配「用户要求」为最高优先级约束，
|
||
# 商品外观仍由参考图锁定（保真语义与一键生成一致）。
|
||
ctx = build_context({"title": "AI 精修"}, fallback_name="product")
|
||
prompt = build_prompt(
|
||
provider_name, model, "custom", ctx,
|
||
style_set=5, lang=spec["lang"],
|
||
extra={"title": "AI 精修", "detail": req.prompt, "prompt_hint": ""},
|
||
requirements=req.prompt,
|
||
)
|
||
|
||
try:
|
||
generator = GENERATORS[provider_name]
|
||
data = await generator(prompt, [req.image_url], size=_image_size(provider_name, spec["ratio"], is_wan=is_wan, model=model), model=model)
|
||
except RuntimeError as exc:
|
||
raise HTTPException(status_code=502, detail=str(exc)) from exc
|
||
except Exception as exc: # noqa: BLE001
|
||
raise HTTPException(status_code=502, detail=f"生成失败: {exc}") from exc
|
||
|
||
is_png = data[:8] == b"\x89PNG\r\n\x1a\n"
|
||
ext = ".png" if is_png else ".jpg"
|
||
ctype = "image/png" if is_png else "image/jpeg"
|
||
url = await get_storage().save_bytes(data, f"suites/single/{uuid.uuid4().hex}{ext}", ctype)
|
||
|
||
asset_id: str | None = None
|
||
if req.append and req.product_id:
|
||
try:
|
||
asset_id = await _append_generated_asset(
|
||
req.product_id,
|
||
TaskImage(type_id="custom", name="AI生图", url=url, status="ok"),
|
||
)
|
||
except Exception: # noqa: BLE001
|
||
pass # 回写失败不影响结果返回，图片已在任务网格可见
|
||
return ImageEditSingleResponse(url=url, asset_id=asset_id)
|