feat: 优化图片资源删除管理，生图任务列表管理
This commit is contained in:
+188
-116
@@ -1,7 +1,7 @@
|
||||
"""套图规划 / 一键生成 / 任务查询 / 结果导出 / 单张 AI 生图。
|
||||
|
||||
按 docs/v2.1/api.md §2-6 实现：引擎平移自 image-suite-studio（services/planner|generator|
|
||||
tasks|watermark + services/prompts），任务存进程内内存表（重启即失效）。
|
||||
watermark + services/prompts），任务持久化到数据库，由独立 Worker 执行。
|
||||
生成图回调回写 product_assets(group_key='generated')；product_id 缺省时只落盘不回写。
|
||||
业务逻辑（texts_to_raw / 模型校验 / generated 回写）在 services/suite_service.py。
|
||||
注：V2.1 阶段本组接口暂不接鉴权（现有鉴权后续可能重做）。
|
||||
@@ -9,11 +9,11 @@ tasks|watermark + services/prompts），任务存进程内内存表（重启即
|
||||
from __future__ import annotations
|
||||
|
||||
import io
|
||||
import tempfile
|
||||
import uuid
|
||||
import zipfile
|
||||
from functools import partial
|
||||
|
||||
from fastapi import APIRouter, BackgroundTasks, HTTPException
|
||||
from collections.abc import Iterator
|
||||
from fastapi import APIRouter, HTTPException
|
||||
from fastapi.responses import StreamingResponse
|
||||
|
||||
from config import get_settings
|
||||
@@ -21,7 +21,7 @@ from schemas.suite import (
|
||||
PLATFORM_SPECS,
|
||||
SUPPORTED_TYPES,
|
||||
ImageEditSingleRequest,
|
||||
ImageEditSingleResponse,
|
||||
ImageEditTaskResponse,
|
||||
PlanItemOut,
|
||||
PlanResponse,
|
||||
SuiteCreateResponse,
|
||||
@@ -31,16 +31,34 @@ from schemas.suite import (
|
||||
api_model,
|
||||
resolve_provider,
|
||||
)
|
||||
from services.generator import _image_size, GENERATORS, run_suite
|
||||
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 build_context, build_prompt, type_name
|
||||
from services.storage import download_bytes, get_storage, local_path
|
||||
from services.suite_service import append_generated_asset, texts_to_raw, validate_model
|
||||
from services.tasks import IMG_OK, TaskImage, create_task, get_task
|
||||
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)
|
||||
@@ -64,7 +82,7 @@ async def plan_suite(req: SuitePlanRequest) -> PlanResponse:
|
||||
# ── 一键生成 ─────────────────────────────────────────────────────────────
|
||||
|
||||
@router.post("/suite/generate", response_model=SuiteCreateResponse)
|
||||
async def generate_suite(req: SuiteGenerateRequest, background: BackgroundTasks) -> SuiteCreateResponse:
|
||||
async def generate_suite(req: SuiteGenerateRequest) -> SuiteCreateResponse:
|
||||
if not req.images:
|
||||
raise HTTPException(status_code=400, detail="未勾选任何图片，无法生成")
|
||||
|
||||
@@ -111,94 +129,146 @@ async def generate_suite(req: SuiteGenerateRequest, background: BackgroundTasks)
|
||||
except ValueError as exc:
|
||||
raise HTTPException(status_code=400, detail="product_id 不是合法 UUID") from exc
|
||||
|
||||
task = create_task(
|
||||
status="pending",
|
||||
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,
|
||||
style_prompt=req.style_prompt,
|
||||
requirements=req.requirements,
|
||||
provider=provider_name,
|
||||
model=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)
|
||||
],
|
||||
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)],
|
||||
)
|
||||
|
||||
# 每张成功即回写 generated 组；未关联商品时仅落存储不回写。
|
||||
# partial 绑定 product_id：run_suite 回调只传 image，签名须为 (image)
|
||||
background.add_task(
|
||||
run_suite,
|
||||
task,
|
||||
partial(append_generated_asset, product_id) if product_id else None,
|
||||
)
|
||||
return SuiteCreateResponse(suite_id=task.id)
|
||||
return SuiteCreateResponse(suite_id=str(task.id))
|
||||
|
||||
|
||||
# ── 任务查询 / 结果导出 ──────────────────────────────────────────────────
|
||||
|
||||
@router.get("/suites/{suite_id}", response_model=SuiteOut)
|
||||
async def get_suite(suite_id: str):
|
||||
task = get_task(suite_id)
|
||||
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="任务不存在（服务可能已重启），请重新生成")
|
||||
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,
|
||||
)
|
||||
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（中文文件名）。"""
|
||||
task = get_task(suite_id)
|
||||
if task is None:
|
||||
raise HTTPException(status_code=404, detail="任务不存在（服务可能已重启）")
|
||||
result = await get_job(suite_id)
|
||||
if result is None:
|
||||
raise HTTPException(status_code=404, detail="任务不存在")
|
||||
task, task_items = result
|
||||
|
||||
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 "")
|
||||
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
|
||||
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)
|
||||
|
||||
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(
|
||||
buf,
|
||||
_stream_file(archive),
|
||||
media_type="application/zip",
|
||||
headers={"Content-Disposition": f"attachment; filename=\"suite-{suite_id}.zip\""},
|
||||
)
|
||||
@@ -206,9 +276,9 @@ async def download_suite_zip(suite_id: str):
|
||||
|
||||
# ── 单张 AI 生图（V2.1 新增，ISS 无对应实现）──────────────────────────────
|
||||
|
||||
@router.post("/suite/image-edit", response_model=ImageEditSingleResponse)
|
||||
async def suite_image_edit(req: ImageEditSingleRequest) -> ImageEditSingleResponse:
|
||||
"""轻量单张再生成：原图 + 一句话要求 + 模型选择。同步接口（单张可长达数分钟）。"""
|
||||
@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:
|
||||
@@ -217,39 +287,41 @@ async def suite_image_edit(req: ImageEditSingleRequest) -> ImageEditSingleRespon
|
||||
|
||||
spec = {"lang": "ru", "ratio": "3:4"} # 试算页固定 Ozon 规格（俄文图内文案 · 3:4）
|
||||
model = api_model(req.model) # 剥 @provider 后缀，供应商侧只认裸模型名
|
||||
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:
|
||||
product_id = (req.product_id or "").strip() or None
|
||||
if product_id:
|
||||
try:
|
||||
asset_id = await append_generated_asset(
|
||||
req.product_id,
|
||||
TaskImage(type_id="custom", name="AI生图", url=url, status="ok"),
|
||||
after_asset_id=req.after_asset_id,
|
||||
)
|
||||
except Exception: # noqa: BLE001
|
||||
pass # 回写失败不影响结果返回，图片已在任务网格可见
|
||||
return ImageEditSingleResponse(url=url, asset_id=asset_id)
|
||||
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))
|
||||
|
||||
Reference in New Issue
Block a user