fix: 修改一下生图的 bug
This commit is contained in:
+12
-1
@@ -13,6 +13,7 @@ import re
|
||||
import tempfile
|
||||
import zipfile
|
||||
from collections.abc import Iterator
|
||||
from urllib.parse import quote
|
||||
|
||||
from fastapi import APIRouter, HTTPException
|
||||
from fastapi.responses import Response, StreamingResponse
|
||||
@@ -80,6 +81,16 @@ def _is_image(url: str, ctype: str) -> bool:
|
||||
return bool(re.search(r"\.(jpe?g|png|webp|gif|bmp)(\?|$)", url, re.IGNORECASE))
|
||||
|
||||
|
||||
def _download_disposition(filename: str) -> str:
|
||||
"""生成兼容中文文件名的下载响应头。
|
||||
|
||||
Starlette 会按 Latin-1 编码 HTTP 头；直接将中文放进 filename 会抛 UnicodeEncodeError
|
||||
并返回 500。filename 提供 ASCII 回退，filename* 使用 RFC 5987 的 UTF-8 百分号编码。
|
||||
"""
|
||||
fallback = re.sub(r"[^\x20-\x7e]", "_", filename) or "image"
|
||||
return f"attachment; filename=\"{fallback}\"; filename*=UTF-8''{quote(filename, safe='')}"
|
||||
|
||||
|
||||
async def _download(url: str) -> tuple[bytes, str]:
|
||||
"""本地 media 文件直读磁盘；远程 URL（源站 CDN / 七牛）带 Referer 下载。"""
|
||||
path = local_path(url)
|
||||
@@ -190,5 +201,5 @@ async def export_image(img: ExportImageItem):
|
||||
return Response(
|
||||
content=data,
|
||||
media_type=ctype,
|
||||
headers={"Content-Disposition": f'attachment; filename="{base}{ext}"'},
|
||||
headers={"Content-Disposition": _download_disposition(f"{base}{ext}")},
|
||||
)
|
||||
|
||||
@@ -190,6 +190,11 @@ async def process_item(claim: ClaimedItem) -> None:
|
||||
if result.data is None:
|
||||
if not result.provider_task_id or not result.provider_task_kind:
|
||||
raise RuntimeError("异步生成未返回可恢复的 provider task id")
|
||||
# 轮询供应商期间用户可能刚好取消任务。此时不能再写回
|
||||
# waiting_provider，否则父任务已取消、子项却会永久显示“进行中”。
|
||||
if not await job_can_write(claim.job_id):
|
||||
await mark_item_cancelled(claim.item_id)
|
||||
return
|
||||
await mark_waiting_provider(claim.item_id, result.provider_task_id, result.provider_task_kind)
|
||||
return
|
||||
|
||||
|
||||
+16
-1
@@ -206,11 +206,26 @@ async def claim_next_item(owner: str, lease_seconds: int = 90) -> ClaimedItem |
|
||||
|
||||
|
||||
async def mark_waiting_provider(item_id: str, provider_task_id: str, provider_task_kind: str, delay_seconds: int = 3) -> None:
|
||||
"""先持久化外部任务 ID，再安排下一次查询，是避免重启重复扣费的关键提交点。"""
|
||||
"""持久化外部任务 ID 并安排下一次查询。
|
||||
|
||||
取消可能恰好发生在 Worker 轮询供应商的网络请求期间。这里在同一事务内
|
||||
再次检查父任务，避免已取消任务被写回 waiting_provider 后永远不再被领取。
|
||||
"""
|
||||
async with get_session_factory()() as db:
|
||||
item = await db.get(BackgroundJobItem, uuid.UUID(item_id))
|
||||
if item is None:
|
||||
return
|
||||
job = await db.get(BackgroundJob, item.job_id)
|
||||
if job is None or job.cancel_requested:
|
||||
item.status = ITEM_CANCELLED
|
||||
item.lease_owner = None
|
||||
item.lease_expires_at = None
|
||||
item.next_run_at = None
|
||||
item.completed_at = utcnow()
|
||||
if job is not None:
|
||||
await _refresh_job_status(db, job)
|
||||
await db.commit()
|
||||
return
|
||||
item.provider_task_id = provider_task_id
|
||||
item.provider_task_kind = provider_task_kind
|
||||
item.status = ITEM_WAITING_PROVIDER
|
||||
|
||||
+154
-15
@@ -27,6 +27,16 @@ ALLOWED_KINDS = [
|
||||
# 「生图要求」明确写了「N 张」时按用户指定的来，上限放宽到 24。
|
||||
PLAN_SAFETY_MAX = 15
|
||||
PLAN_REQUIREMENTS_MAX = 24
|
||||
MAX_VARIANT_MAIN_IMAGES = 2
|
||||
|
||||
# 规划器的图类归并。不同模型会把“细节图”写成材质图、卖点图或详情图，
|
||||
# 因此均视作同一类覆盖，而不是强迫模型只使用 material。
|
||||
SCENE_KINDS = {"lifestyle", "multi_scene"}
|
||||
DETAIL_KINDS = {"material", "key_features", "selling_pt", "ecommerce_detail"}
|
||||
_DIMENSION_PATTERN = re.compile(
|
||||
r"尺寸|长宽高|直径|厚度|(?:长|宽|高)\s*[:：×x*]?\s*\d+|\d+(?:\.\d+)?\s*(?:厘米|毫米|cm|mm|inch|in)\b",
|
||||
re.IGNORECASE,
|
||||
)
|
||||
|
||||
|
||||
def _total_budget(requirements: str | None) -> int:
|
||||
@@ -47,12 +57,12 @@ SYSTEM_PROMPT = """你是一名资深电商视觉策划。根据商品信息规
|
||||
5. variant_name 只能从「SKU规格」列表原样照抄；没有绑定就输出 null。
|
||||
|
||||
## 规划规则
|
||||
1. 先保证 SKU 主图（规则 2），其余类型（场景/细节/尺寸/合集等）按商品特点自行取舍并分配张数，追求信息效率，不为凑数出图。
|
||||
2. SKU 主图：每个带图 SKU 出 1 张独立主图（kind=white_bg），variant_name 填对应规格名；单 SKU 出 1 张（variant_name=null）。
|
||||
3. 场景图（kind=lifestyle）：按核心使用场景出 1-3 张，每张聚焦一个场景。
|
||||
4. 细节图（kind=material 或 custom）：按关键细节/材质/结构出 1-2 张，每张聚焦一个卖点。
|
||||
5. 尺寸标注图（kind=size_chart）：参数含长宽高/尺寸时出 1 张。
|
||||
6. SKU 合集图（kind=sku_collection）：SKU >1 时出 1 张。
|
||||
1. 套图必须均衡覆盖：白底主图、场景图、细节图；参数含可靠的长宽高/尺寸时还必须有尺寸标注图。按这个顺序先预留张数，再安排其他类型。
|
||||
2. 白底主图（kind=white_bg）通常只出 1 张；有多个视觉差异明显的 SKU 时最多选 2 个代表款各出 1 张，variant_name 填对应规格名。**禁止因为 SKU 多而为每个 SKU 都规划主图**。仅当「生图要求」明确要求每个/全部 SKU 分别出图时才可例外。
|
||||
3. 场景图（kind=lifestyle）：至少 1 张，按核心使用场景出 1-2 张，每张聚焦一个场景。
|
||||
4. 细节图（kind=material / key_features / selling_pt / ecommerce_detail）：至少 1 张，展示关键细节、材质或结构；通常 1-2 张，每张聚焦一个卖点。
|
||||
5. 尺寸标注图（kind=size_chart）：参数含长宽高/尺寸时必须出 1 张；无可靠尺寸参数时不要虚构尺寸。
|
||||
6. SKU 合集图（kind=sku_collection）：仅在 SKU >1 且确有必要展示差异时出 1 张，优先级低于上述四类。
|
||||
7. kind 枚举：white_bg / key_features / selling_pt / material / lifestyle / multi_scene / ecommerce_detail / size_chart / sku_collection / custom。
|
||||
8. title 用中文短语（如「主图·粉色」「浴室壁挂」）；detail 中文说明这张图展示什么；prompt_hint 用英文描述构图要点。
|
||||
|
||||
@@ -114,6 +124,141 @@ def _normalize_items(raw_items: list, sku_variants: list[str]) -> list[dict]:
|
||||
return items
|
||||
|
||||
|
||||
def _has_dimensions(product_info: dict) -> bool:
|
||||
"""仅在已采集到尺寸信息时规划尺寸图，避免模型编造数值。"""
|
||||
for param in product_info.get("params") or []:
|
||||
if not isinstance(param, dict):
|
||||
continue
|
||||
text = f"{param.get('key') or ''} {param.get('value') or ''}"
|
||||
if _DIMENSION_PATTERN.search(text):
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def _requires_every_sku(requirements: str | None) -> bool:
|
||||
"""用户显式要求逐 SKU 出图时，保留其优先级，不施加代表款上限。"""
|
||||
if not requirements:
|
||||
return False
|
||||
return bool(re.search(
|
||||
r"(?:每个|全部|所有|逐个|分别).{0,10}(?:sku|SKU|规格|款式|颜色|变体)",
|
||||
requirements,
|
||||
))
|
||||
|
||||
|
||||
def _coverage_group(kind: str) -> str:
|
||||
if kind == "white_bg":
|
||||
return "main"
|
||||
if kind in SCENE_KINDS:
|
||||
return "scene"
|
||||
if kind in DETAIL_KINDS:
|
||||
return "detail"
|
||||
if kind == "size_chart":
|
||||
return "size"
|
||||
return "other"
|
||||
|
||||
|
||||
def _total(items: list[dict]) -> int:
|
||||
return sum(item["count"] for item in items)
|
||||
|
||||
|
||||
def _group_total(items: list[dict], group: str) -> int:
|
||||
return sum(item["count"] for item in items if _coverage_group(item["kind"]) == group)
|
||||
|
||||
|
||||
def _remove_one(items: list[dict], predicate) -> bool:
|
||||
"""从后往前删除一个生成槽位，优先保留模型较早规划的构图。"""
|
||||
for index in range(len(items) - 1, -1, -1):
|
||||
if not predicate(items[index]):
|
||||
continue
|
||||
if items[index]["count"] > 1:
|
||||
items[index]["count"] -= 1
|
||||
else:
|
||||
items.pop(index)
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def _make_room(items: list[dict], budget: int, minimums: dict[str, int]) -> bool:
|
||||
"""为必需图类腾出一张，先去除非核心图或某类的冗余张数。"""
|
||||
candidates: list[tuple[int, int]] = []
|
||||
for index, item in enumerate(items):
|
||||
group = _coverage_group(item["kind"])
|
||||
if _group_total(items, group) <= minimums.get(group, 0):
|
||||
continue
|
||||
# other 最低优先；同类冗余时从后面的项开始删。
|
||||
priority = 0 if group == "other" else 1
|
||||
candidates.append((priority, -index))
|
||||
if not candidates:
|
||||
return False
|
||||
_, neg_index = min(candidates)
|
||||
index = -neg_index
|
||||
item = items[index]
|
||||
if item["count"] > 1:
|
||||
item["count"] -= 1
|
||||
else:
|
||||
items.pop(index)
|
||||
return True
|
||||
|
||||
|
||||
def _balanced_items(
|
||||
items: list[dict],
|
||||
sku_variants: list[str],
|
||||
product_info: dict,
|
||||
requirements: str | None,
|
||||
budget: int,
|
||||
) -> list[dict]:
|
||||
"""将 LLM 方案约束为信息均衡的商品套图。
|
||||
|
||||
提示词能改善规划倾向，但不能保证每次输出。这里把 SKU 主图总量限制为
|
||||
代表款数量，并在预算允许时补齐核心图类，避免 SKU 多时挤掉场景/细节/尺寸图。
|
||||
"""
|
||||
balanced = [dict(item) for item in items if item["count"] > 0]
|
||||
|
||||
if not _requires_every_sku(requirements):
|
||||
max_mains = 1 if len(sku_variants) <= 1 else MAX_VARIANT_MAIN_IMAGES
|
||||
while sum(item["count"] for item in balanced if item["kind"] == "white_bg") > max_mains:
|
||||
_remove_one(balanced, lambda item: item["kind"] == "white_bg")
|
||||
|
||||
required: list[tuple[str, dict]] = [
|
||||
("main", {
|
||||
"kind": "white_bg", "title": "白底主图", "detail": "商品白底主视觉",
|
||||
"prompt_hint": "front view on pure white background", "count": 1,
|
||||
"variant_name": sku_variants[0] if sku_variants else None,
|
||||
}),
|
||||
("scene", {
|
||||
"kind": "lifestyle", "title": "使用场景", "detail": "展示商品实际使用场景",
|
||||
"prompt_hint": "product in a natural everyday use setting", "count": 1,
|
||||
"variant_name": None,
|
||||
}),
|
||||
("detail", {
|
||||
"kind": "material", "title": "细节展示", "detail": "突出材质与结构细节",
|
||||
"prompt_hint": "close-up of authentic material and construction details", "count": 1,
|
||||
"variant_name": None,
|
||||
}),
|
||||
]
|
||||
if _has_dimensions(product_info):
|
||||
required.append(("size", {
|
||||
"kind": "size_chart", "title": "尺寸标注", "detail": "展示已采集的商品尺寸",
|
||||
"prompt_hint": "clear size chart using known product dimensions", "count": 1,
|
||||
"variant_name": None,
|
||||
}))
|
||||
|
||||
# 用户明确要求较小张数时，优先保留前面的核心图类；不突破其指定预算。
|
||||
required = required[:budget]
|
||||
minimums = {group: 1 for group, _ in required}
|
||||
for group, fallback in required:
|
||||
if _group_total(balanced, group) > 0:
|
||||
continue
|
||||
while _total(balanced) >= budget and _make_room(balanced, budget, minimums):
|
||||
pass
|
||||
if _total(balanced) < budget:
|
||||
balanced.append(fallback)
|
||||
|
||||
while _total(balanced) > budget and _make_room(balanced, budget, minimums):
|
||||
pass
|
||||
return balanced
|
||||
|
||||
|
||||
def _repair_truncated(s: str) -> dict | None:
|
||||
"""截断修复：从最后一个完整的 '}' 处截断，剥尾逗号后按括号配平补全闭合。
|
||||
|
||||
@@ -234,16 +379,10 @@ async def generate_plan(
|
||||
|
||||
items = _normalize_items(data.get("items") or [], sku_variants)
|
||||
if items:
|
||||
# 总量保护：只防跑飞不干预正常规划。生图要求指定「N 张」时按 N 放宽上限
|
||||
# 先压缩多 SKU 的重复主图、补齐核心视觉类别，再执行总量保护。
|
||||
# 生图要求指定「N 张」时按 N 放宽上限，并始终不突破该预算。
|
||||
budget = _total_budget(requirements)
|
||||
total = sum(i["count"] for i in items)
|
||||
while total > budget and items:
|
||||
last = items[-1]
|
||||
if last["count"] > 1:
|
||||
last["count"] -= 1
|
||||
else:
|
||||
items.pop()
|
||||
total = sum(i["count"] for i in items)
|
||||
items = _balanced_items(items, sku_variants, product_info, requirements, budget)
|
||||
return {"summary": str(data.get("summary") or "").strip()[:100], "items": items}
|
||||
|
||||
last_error = RuntimeError("规划器未返回有效方案项")
|
||||
|
||||
@@ -76,7 +76,7 @@ async def append_generated_asset(
|
||||
sort_order: int | None = None
|
||||
if after_asset_id:
|
||||
try:
|
||||
anchor = await db.get(ProductAsset, UUID(after_asset_id))
|
||||
anchor = await db.get(ProductAsset, uuid.UUID(after_asset_id))
|
||||
except ValueError:
|
||||
anchor = None
|
||||
if anchor is not None and anchor.product_id == pid and anchor.group_key == "generated":
|
||||
@@ -89,7 +89,7 @@ async def append_generated_asset(
|
||||
ProductAsset.sort_order > anchor.sort_order,
|
||||
)
|
||||
)
|
||||
).scalars().all()
|
||||
).all()
|
||||
for follower in followers:
|
||||
follower.sort_order += 1
|
||||
sort_order = anchor.sort_order + 1
|
||||
|
||||
@@ -0,0 +1,32 @@
|
||||
"""单图导出响应头回归测试。"""
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
SERVER_DIR = Path(__file__).resolve().parents[1]
|
||||
if str(SERVER_DIR) not in sys.path:
|
||||
sys.path.insert(0, str(SERVER_DIR))
|
||||
|
||||
from api.export import ExportImageItem, export_image
|
||||
|
||||
|
||||
class ExportImageTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_chinese_filename_uses_rfc5987_header(self) -> None:
|
||||
image = ExportImageItem(
|
||||
url="https://example.com/image.png",
|
||||
key="尺寸标注图",
|
||||
variantName="蓝色",
|
||||
)
|
||||
with patch("api.export._download", new=AsyncMock(return_value=(b"image", "image/png"))):
|
||||
response = await export_image(image)
|
||||
|
||||
disposition = response.headers["content-disposition"]
|
||||
self.assertIn('filename="_____-__.png"', disposition)
|
||||
self.assertIn("filename*=UTF-8''%E5%B0%BA%E5%AF%B8%E6%A0%87%E6%B3%A8%E5%9B%BE-%E8%93%9D%E8%89%B2.png", disposition)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,63 @@
|
||||
"""后台生图 Worker 的取消竞态回归测试。"""
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
SERVER_DIR = Path(__file__).resolve().parents[1]
|
||||
if str(SERVER_DIR) not in sys.path:
|
||||
sys.path.insert(0, str(SERVER_DIR))
|
||||
|
||||
from services.generator import GenerationAdvance
|
||||
from services.job_worker import process_item
|
||||
from services.jobs import ClaimedItem
|
||||
|
||||
|
||||
class WorkerCancellationTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_cancelled_task_is_not_returned_to_provider_waiting_queue(self) -> None:
|
||||
claim = ClaimedItem(
|
||||
job_id="00000000-0000-0000-0000-000000000001",
|
||||
item_id="00000000-0000-0000-0000-000000000002",
|
||||
job_kind="image_edit",
|
||||
product_id=None,
|
||||
provider="pcore",
|
||||
model="gpt-image-2",
|
||||
platform="ozon",
|
||||
lang="ru",
|
||||
ratio="3:4",
|
||||
style_set=5,
|
||||
request_payload={"context": {"title": "test"}, "ref_images": []},
|
||||
item_payload={"kind": "custom", "title": "AI生图"},
|
||||
provider_task_id="provider-task",
|
||||
provider_task_kind="edits",
|
||||
)
|
||||
settings = SimpleNamespace(
|
||||
image_provider="pcore",
|
||||
dashscope_model="wan2.7-image-pro",
|
||||
rightapi_image_model="gpt-image-2",
|
||||
ark_image_model="doubao-seedream-4-5-251128",
|
||||
)
|
||||
|
||||
with (
|
||||
patch("services.job_worker.get_settings", return_value=settings),
|
||||
patch("services.job_worker.build_context", return_value={}),
|
||||
patch("services.job_worker.build_prompt", return_value="prompt"),
|
||||
patch("services.job_worker._refs_for_job", return_value=[]),
|
||||
patch("services.job_worker.advance_generation", new=AsyncMock(return_value=GenerationAdvance(
|
||||
provider_task_id="provider-task", provider_task_kind="edits",
|
||||
))),
|
||||
patch("services.job_worker.job_can_write", new=AsyncMock(return_value=False)),
|
||||
patch("services.job_worker.mark_item_cancelled", new=AsyncMock()) as cancelled,
|
||||
patch("services.job_worker.mark_waiting_provider", new=AsyncMock()) as waiting,
|
||||
):
|
||||
await process_item(claim)
|
||||
|
||||
cancelled.assert_awaited_once_with(claim.item_id)
|
||||
waiting.assert_not_awaited()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,69 @@
|
||||
"""AI 套图规划的均衡覆盖回归测试。"""
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
|
||||
SERVER_DIR = Path(__file__).resolve().parents[1]
|
||||
if str(SERVER_DIR) not in sys.path:
|
||||
sys.path.insert(0, str(SERVER_DIR))
|
||||
|
||||
from services.planner import _balanced_items
|
||||
|
||||
|
||||
def item(kind: str, variant_name: str | None = None) -> dict:
|
||||
return {
|
||||
"kind": kind,
|
||||
"title": kind,
|
||||
"detail": "",
|
||||
"prompt_hint": "",
|
||||
"count": 1,
|
||||
"variant_name": variant_name,
|
||||
}
|
||||
|
||||
|
||||
class BalancedPlanTests(unittest.TestCase):
|
||||
def test_many_skus_do_not_crowd_out_core_image_types(self) -> None:
|
||||
variants = ["黑色", "白色", "蓝色", "粉色", "绿色"]
|
||||
plan = _balanced_items(
|
||||
[item("white_bg", variant) for variant in variants],
|
||||
variants,
|
||||
{"params": [{"key": "尺寸", "value": "30 × 20 × 10 cm"}]},
|
||||
None,
|
||||
budget=15,
|
||||
)
|
||||
|
||||
self.assertLessEqual(sum(p["count"] for p in plan if p["kind"] == "white_bg"), 2)
|
||||
self.assertTrue(any(p["kind"] == "lifestyle" for p in plan))
|
||||
self.assertTrue(any(p["kind"] == "material" for p in plan))
|
||||
self.assertTrue(any(p["kind"] == "size_chart" for p in plan))
|
||||
|
||||
def test_does_not_add_size_chart_without_dimensions(self) -> None:
|
||||
plan = _balanced_items(
|
||||
[item("white_bg", "黑色")],
|
||||
["黑色", "白色"],
|
||||
{"params": [{"key": "规格", "value": "标准款"}, {"key": "材质", "value": "涤纶"}]},
|
||||
None,
|
||||
budget=15,
|
||||
)
|
||||
|
||||
self.assertFalse(any(p["kind"] == "size_chart" for p in plan))
|
||||
self.assertTrue(any(p["kind"] == "lifestyle" for p in plan))
|
||||
self.assertTrue(any(p["kind"] == "material" for p in plan))
|
||||
|
||||
def test_explicit_every_sku_requirement_is_preserved(self) -> None:
|
||||
variants = ["黑色", "白色", "蓝色"]
|
||||
plan = _balanced_items(
|
||||
[item("white_bg", variant) for variant in variants],
|
||||
variants,
|
||||
{},
|
||||
"每个 SKU 都单独出一张白底图",
|
||||
budget=15,
|
||||
)
|
||||
|
||||
self.assertEqual(sum(p["count"] for p in plan if p["kind"] == "white_bg"), 3)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user