Files
image-suite-studio/server/api/generate.py
T
2026-08-16 17:32:43 +08:00

137 lines
5.5 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
"""无状态套图生成:请求自带采集数据,不落商品库。"""
from __future__ import annotations
from fastapi import APIRouter, BackgroundTasks, Depends, HTTPException
from config import get_settings
from db import get_db
from models import Suite
from schemas import (
GenerateRequest, PLATFORM_SPECS, SUPPORTED_TYPES, TONGYI_MODELS, SuiteCreateResponse, TextMaterial,
PlanRequest, PlanResponse, PlanItemOut,
)
from services.generator import run_suite
from services.planner import generate_plan
from services.prompt import type_name
router = APIRouter(prefix="/api", tags=["generate"])
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
@router.post("/generate", response_model=SuiteCreateResponse)
async def generate_suite(
req: GenerateRequest,
background: BackgroundTasks,
db=Depends(get_db),
) -> 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")
types = list(dict.fromkeys(j["kind"] for j in jobs))
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_name = req.provider or settings.image_provider
model = req.model
if provider_name == "tongyi" and model and model not in TONGYI_MODELS:
raise HTTPException(status_code=400, detail=f"不支持的模型: {model}tongyi 支持: {TONGYI_MODELS}")
suite = Suite(
product_id=None,
style_set=req.style_set,
style_prompt=req.style_prompt,
platform=req.platform,
lang=spec["lang"],
ratio=spec["ratio"],
types=types,
plan=jobs,
provider=provider_name,
model=model,
context=texts_to_raw(req.texts),
# 参考图池: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)
],
)
db.add(suite)
await db.commit()
await db.refresh(suite)
background.add_task(run_suite, str(suite.id))
return SuiteCreateResponse(suite_id=str(suite.id))
@router.post("/plan", response_model=PlanResponse)
async def plan_suite(req: PlanRequest) -> 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)
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"]],
)