"""套图规划 / 一键生成 / 任务查询 / 结果导出 / 单张 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))