328 lines
14 KiB
Python
328 lines
14 KiB
Python
"""套图规划 / 一键生成 / 任务查询 / 结果导出 / 单张 AI 生图。
|
||
|
||
按 docs/v2.1/api.md §2-6 实现：引擎平移自 image-suite-studio（services/planner|generator|
|
||
watermark + services/prompts），任务持久化到数据库，由独立 Worker 执行。
|
||
生成图回调回写 product_assets(group_key='generated')；product_id 缺省时只落盘不回写。
|
||
业务逻辑（texts_to_raw / 模型校验 / generated 回写）在 services/suite_service.py。
|
||
注：V2.1 阶段本组接口暂不接鉴权（现有鉴权后续可能重做）。
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import io
|
||
import tempfile
|
||
import uuid
|
||
import zipfile
|
||
from collections.abc import Iterator
|
||
from fastapi import APIRouter, HTTPException
|
||
from fastapi.responses import StreamingResponse
|
||
|
||
from config import get_settings
|
||
from schemas.suite import (
|
||
PLATFORM_SPECS,
|
||
SUPPORTED_TYPES,
|
||
ImageEditSingleRequest,
|
||
ImageEditTaskResponse,
|
||
PlanItemOut,
|
||
PlanResponse,
|
||
SuiteCreateResponse,
|
||
SuiteGenerateRequest,
|
||
SuiteOut,
|
||
SuitePlanRequest,
|
||
api_model,
|
||
resolve_provider,
|
||
)
|
||
from services.generator import GENERATORS
|
||
from services.jobs import (
|
||
ITEM_DONE,
|
||
cancel_job,
|
||
create_image_job,
|
||
get_job,
|
||
list_product_jobs,
|
||
retry_failed_items,
|
||
retry_failed_item,
|
||
suite_view,
|
||
)
|
||
from services.planner import generate_plan
|
||
from services.prompts import type_name
|
||
from services.storage import ResourceSizeLimitError, download_bytes, local_path, read_local_bytes
|
||
from services.suite_service import texts_to_raw, validate_model
|
||
|
||
router = APIRouter(prefix="/api", tags=["suite"])
|
||
|
||
|
||
def _stream_file(handle: io.BufferedRandom, chunk_size: int = 64 * 1024) -> Iterator[bytes]:
|
||
"""以固定块返回临时 ZIP，释放响应完成后的临时文件。"""
|
||
try:
|
||
while chunk := handle.read(chunk_size):
|
||
yield chunk
|
||
finally:
|
||
handle.close()
|
||
|
||
|
||
# ── 出图方案规划 ──────────────────────────────────────────────────────────
|
||
|
||
@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) -> 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、gpt-image-2@pcore → pcore）
|
||
provider_name = resolve_provider(req.model, None, settings.image_provider)
|
||
validate_model(provider_name, req.model)
|
||
model = api_model(req.model) # 剥 @provider 后缀，供应商侧只认裸模型名
|
||
|
||
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 = await create_image_job(
|
||
product_id=product_id,
|
||
kind="image_suite",
|
||
priority=0,
|
||
provider=provider_name,
|
||
model=model,
|
||
platform=req.platform,
|
||
lang=spec["lang"],
|
||
ratio=spec["ratio"],
|
||
style_set=req.style_set,
|
||
request_payload={
|
||
"context": texts_to_raw(req.texts),
|
||
"style_prompt": req.style_prompt,
|
||
"requirements": req.requirements,
|
||
"watermark": req.watermark.model_dump() if req.watermark else None,
|
||
"append": bool(product_id),
|
||
# main 组优先；variant_name 精确匹配由 Worker 复用原有选图规则。
|
||
"ref_images": [
|
||
{"url": image.url, "group_key": image.group_key, "variant_name": image.variant_name}
|
||
for image in sorted(req.images, key=lambda image: 0 if image.group_key == "main" else 1)
|
||
],
|
||
},
|
||
item_payloads=[{**job, "position": index} for index, job in enumerate(jobs)],
|
||
)
|
||
return SuiteCreateResponse(suite_id=str(task.id))
|
||
|
||
|
||
# ── 任务查询 / 结果导出 ──────────────────────────────────────────────────
|
||
|
||
@router.get("/suites/{suite_id}", response_model=SuiteOut)
|
||
async def get_suite(suite_id: str):
|
||
result = await get_job(suite_id)
|
||
if result is None:
|
||
raise HTTPException(status_code=404, detail="任务不存在")
|
||
return SuiteOut(**suite_view(*result))
|
||
|
||
|
||
@router.get("/products/{product_id}/suites", response_model=list[SuiteOut])
|
||
async def list_suites(product_id: str, active_only: bool = False):
|
||
"""试算页恢复商品的未完成任务；后续任务中心可复用。"""
|
||
jobs = await list_product_jobs(product_id, active_only=active_only, kinds=("image_suite", "image_edit"))
|
||
views: list[SuiteOut] = []
|
||
for job in jobs:
|
||
result = await get_job(str(job.id))
|
||
if result is not None:
|
||
views.append(SuiteOut(**suite_view(*result)))
|
||
return views
|
||
|
||
|
||
@router.post("/suites/{suite_id}/cancel", response_model=SuiteOut)
|
||
async def cancel_suite(suite_id: str):
|
||
task = await cancel_job(suite_id)
|
||
if task is None:
|
||
raise HTTPException(status_code=404, detail="任务不存在")
|
||
result = await get_job(str(task.id))
|
||
if result is None:
|
||
raise HTTPException(status_code=404, detail="任务不存在")
|
||
return SuiteOut(**suite_view(*result))
|
||
|
||
|
||
@router.post("/suites/{suite_id}/retry", response_model=SuiteOut)
|
||
async def retry_suite(suite_id: str):
|
||
task = await retry_failed_items(suite_id)
|
||
if task is None:
|
||
raise HTTPException(status_code=404, detail="任务不存在")
|
||
result = await get_job(str(task.id))
|
||
if result is None:
|
||
raise HTTPException(status_code=404, detail="任务不存在")
|
||
return SuiteOut(**suite_view(*result))
|
||
|
||
|
||
@router.post("/suites/{suite_id}/items/{item_id}/retry", response_model=SuiteOut)
|
||
async def retry_suite_item(suite_id: str, item_id: str):
|
||
task = await retry_failed_item(suite_id, item_id)
|
||
if task is None:
|
||
raise HTTPException(status_code=404, detail="任务或子项不存在")
|
||
result = await get_job(str(task.id))
|
||
if result is None:
|
||
raise HTTPException(status_code=404, detail="任务不存在")
|
||
return SuiteOut(**suite_view(*result))
|
||
|
||
|
||
@router.get("/suites/{suite_id}/zip")
|
||
async def download_suite_zip(suite_id: str):
|
||
"""把任务内所有成功图打包成 ZIP（中文文件名）。"""
|
||
result = await get_job(suite_id)
|
||
if result is None:
|
||
raise HTTPException(status_code=404, detail="任务不存在")
|
||
task, task_items = result
|
||
|
||
settings = get_settings()
|
||
images = [image for image in task_items if image.status == ITEM_DONE and image.result_url]
|
||
if len(images) > settings.export_max_images:
|
||
raise HTTPException(status_code=413, detail=f"该任务有 {len(images)} 张图片，单次最多导出 {settings.export_max_images} 张")
|
||
|
||
archive = tempfile.SpooledTemporaryFile(
|
||
max_size=settings.export_zip_memory_threshold_bytes,
|
||
mode="w+b",
|
||
)
|
||
total_source_bytes = 0
|
||
try:
|
||
with zipfile.ZipFile(archive, "w", zipfile.ZIP_DEFLATED) as zf:
|
||
seen: set[str] = set()
|
||
for i, img in enumerate(images):
|
||
path = local_path(img.result_url or "")
|
||
try:
|
||
if path is not None:
|
||
data, suffix = read_local_bytes(path), path.suffix or ".jpg"
|
||
else: # 七牛等远程存储：直接拉取公开 URL
|
||
data, _ = await download_bytes(img.result_url or "")
|
||
suffix = ".png" if data[:8] == b"\x89PNG\r\n\x1a\n" else ".jpg"
|
||
except ResourceSizeLimitError as exc:
|
||
raise HTTPException(status_code=413, detail=f"无法导出「{img.name}」：{exc}") from exc
|
||
except Exception: # noqa: BLE001
|
||
continue
|
||
|
||
total_source_bytes += len(data)
|
||
if total_source_bytes > settings.export_zip_max_bytes:
|
||
raise HTTPException(
|
||
status_code=413,
|
||
detail=f"导出图片总大小超过 {settings.export_zip_max_bytes // 1024 // 1024} MB，请减少图片后重试",
|
||
)
|
||
payload = img.payload or {}
|
||
filename = payload.get("title") or payload.get("kind") or "AI 生图"
|
||
if filename in seen: # 同类型多张时加序号防覆盖
|
||
filename = f"{filename}-{i + 1}"
|
||
seen.add(filename)
|
||
zf.writestr(f"{filename}{suffix}", data)
|
||
if archive.tell() > settings.export_zip_max_bytes:
|
||
raise HTTPException(
|
||
status_code=413,
|
||
detail=f"导出 ZIP 超过 {settings.export_zip_max_bytes // 1024 // 1024} MB，请减少图片后重试",
|
||
)
|
||
except Exception:
|
||
archive.close()
|
||
raise
|
||
|
||
archive.seek(0)
|
||
return StreamingResponse(
|
||
_stream_file(archive),
|
||
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=ImageEditTaskResponse)
|
||
async def suite_image_edit(req: ImageEditSingleRequest) -> ImageEditTaskResponse:
|
||
"""提交单张图生图任务；立即返回任务 ID，实际生成在统一队列异步执行。"""
|
||
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 = api_model(req.model) # 剥 @provider 后缀，供应商侧只认裸模型名
|
||
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
|
||
|
||
# 单图编辑优先级高于套图；prompt_extra 清空展示文案，避免中文 UI 文案进入模型 prompt。
|
||
task = await create_image_job(
|
||
product_id=product_id,
|
||
kind="image_edit",
|
||
priority=10,
|
||
provider=provider_name,
|
||
model=model,
|
||
platform="ozon",
|
||
lang=spec["lang"],
|
||
ratio=spec["ratio"],
|
||
style_set=5,
|
||
request_payload={
|
||
"context": {"title": "AI 精修"},
|
||
"style_prompt": None,
|
||
"requirements": req.prompt,
|
||
"watermark": req.watermark.model_dump() if req.watermark else None,
|
||
"append": bool(req.append and product_id),
|
||
"after_asset_id": req.after_asset_id,
|
||
"ref_images": [{"url": req.image_url, "group_key": "main", "variant_name": None}],
|
||
},
|
||
item_payloads=[{
|
||
"position": 0,
|
||
"kind": "custom",
|
||
"title": "AI生图",
|
||
"detail": "",
|
||
"prompt_hint": "",
|
||
"variant_name": None,
|
||
"prompt_extra": {"title": "", "detail": "", "prompt_hint": ""},
|
||
}],
|
||
)
|
||
return ImageEditTaskResponse(suite_id=str(task.id))
|