feat: 优化图片资源删除管理,生图任务列表管理

This commit is contained in:
2026-08-30 14:50:33 +08:00
parent 4ee667142b
commit 957289f717
46 changed files with 4180 additions and 524 deletions
+197 -6
View File
@@ -13,6 +13,7 @@ import base64
import logging
import mimetypes
import re
from dataclasses import dataclass
from typing import Awaitable, Callable
import httpx
@@ -30,7 +31,7 @@ from services.tasks import (
TaskImage,
)
from services.watermark import apply_watermark
from services.storage import download_bytes, get_storage, local_path
from services.storage import download_bytes, get_storage, local_path, read_local_bytes
log = logging.getLogger("suite.generator")
@@ -125,7 +126,7 @@ async def _resolve_ref_bytes(url: str) -> tuple[bytes, str]:
path = local_path(url)
if path is not None:
mime = mimetypes.guess_type(path.name)[0] or "image/jpeg"
return path.read_bytes(), mime
return read_local_bytes(path), mime
if url.startswith(("http://", "https://")):
from api.proxy import guess_referer
data, ctype = await download_bytes(url, referer=guess_referer(url))
@@ -288,6 +289,18 @@ async def _rightapi_poll_task(client: httpx.AsyncClient, headers: dict, origin:
raise TimeoutError(f"RightAPI 异步任务超时 ({max_wait}s): task_id={task_id}")
async def _rightapi_poll_task_once(client: httpx.AsyncClient, headers: dict, origin: str, task_id: str) -> dict | None:
"""持久化 Worker 的单次查询版本;未完成返回 None,不占住 Worker 等待。"""
response = await client.get(f"{origin}/v1/tasks/{task_id}", headers=headers, timeout=30)
_raise_api_error(response, "RightAPI")
result = response.json()
status = result.get("status", "")
if status in ("failed", "error", "cancelled"):
error = result.get("error") or {}
raise RuntimeError(f"RightAPI 任务失败: {error.get('message') or result}")
return result if "data" in result else None
def _rightapi_extract_image(result: dict) -> tuple[str | None, str | None]:
"""从轮询完成结果里取 (kind, payload):kind ∈ url | b64,未取到返回 (None, None)。
@@ -411,6 +424,19 @@ async def _pcore_poll_task(client: httpx.AsyncClient, headers: dict, base: str,
raise TimeoutError(f"Pcore 异步任务超时 ({max_wait}s): task_id={task_id}")
async def _pcore_poll_task_once(client: httpx.AsyncClient, headers: dict, base: str,
kind: str, task_id: str) -> dict | None:
"""持久化 Worker 的单次查询版本;未完成返回 None。"""
response = await client.get(f"{base}/v1/images/{kind}/{task_id}", headers=headers, timeout=30)
_raise_api_error(response, "Pcore")
result = response.json()
status = result.get("status", "")
if status in ("failed", "error", "cancelled"):
error = result.get("error") or {}
raise RuntimeError(f"Pcore 任务失败: {error.get('message') or result}")
return result if status in ("completed", "succeeded") else None
async def _pcore_request(s, prompt: str, ref_images: list[str], size: str, model: str) -> bytes:
"""Pcore:无参考图走 /v1/images/generations,带参考图走 /v1/images/edits(JSON,data-URI)。
@@ -482,6 +508,164 @@ async def generate_pcore(prompt: str, ref_images: list[str], size: str = "1536x2
raise last_exc # type: ignore[misc]
# ── 可恢复的逐步 provider 调用(持久化 Worker 使用)──────────────────────
@dataclass(frozen=True)
class GenerationAdvance:
"""一次 Worker 推进的结果:data 有值时已完成,否则持久化外部 task id 后稍后再查。"""
data: bytes | None = None
provider_task_id: str | None = None
provider_task_kind: str | None = None
async def _download_provider_image(client: httpx.AsyncClient, result: dict, provider: str) -> bytes:
kind, payload = _rightapi_extract_image(result)
if kind == "b64" and payload:
return base64.b64decode(payload.split(",", 1)[-1] if "," in payload else payload)
if kind == "url" and payload:
response = await client.get(payload, timeout=get_settings().request_timeout)
_raise_api_error(response, provider)
return response.content
raise RuntimeError(f"{provider} 任务完成但没有图片数据: {result}")
async def _poll_tongyi_once(client: httpx.AsyncClient, key: str, task_id: str) -> tuple[str, str | None]:
response = await client.get(
"https://dashscope.aliyuncs.com/api/v1/tasks/" + task_id,
headers={"Authorization": f"Bearer {key}"}, timeout=30,
)
_raise_api_error(response, "通义")
result = response.json()
status = result.get("output", {}).get("task_status", "")
if status in ("FAILED", "UNKNOWN"):
raise RuntimeError(f"通义任务失败: {result}")
if status != "SUCCEEDED":
return "waiting", None
choices = result.get("output", {}).get("choices", [])
if choices:
content = choices[0].get("message", {}).get("content", [])
if content:
return "done", content[0].get("image", "")
results = result.get("output", {}).get("results", [])
if results:
return "done", results[0].get("url") or results[0].get("b64_image", "")
raise RuntimeError(f"通义任务成功但无结果: {result}")
async def advance_generation(
provider: str,
prompt: str,
ref_images: list[str],
*,
size: str,
model: str,
provider_task_id: str | None = None,
provider_task_kind: str | None = None,
) -> GenerationAdvance:
"""推进一次图片生成,不在本函数内无限轮询。
对异步供应商:首次仅提交并返回 task id;后续调用只轮询该 id。这样 Worker 重启后
能继续外部任务,且在等待期间可以领取其他任务。豆包和同步千问仍在一次调用内完成。
"""
settings = get_settings()
if provider == "doubao":
return GenerationAdvance(data=await generate_doubao(prompt, ref_images, size=size, model=model))
if provider == "tongyi":
if not settings.dashscope_api_key:
raise RuntimeError("未配置 DASHSCOPE_API_KEY(.env)")
is_wan = _is_wan_model(model)
if not is_wan:
return GenerationAdvance(data=await generate_tongyi(prompt, ref_images, size=size, model=model))
headers = {"Authorization": f"Bearer {settings.dashscope_api_key}", "Content-Type": "application/json"}
if provider_task_id:
async with httpx.AsyncClient(timeout=60, verify=False) as client:
state, payload = await _poll_tongyi_once(client, settings.dashscope_api_key, provider_task_id)
if state == "waiting":
return GenerationAdvance(provider_task_id=provider_task_id, provider_task_kind="tongyi_wan")
if payload and (payload.startswith("data:") or len(payload) > 500):
return GenerationAdvance(data=base64.b64decode(payload.split(",", 1)[-1] if "," in payload else payload))
if not payload:
raise RuntimeError("通义任务成功但图片地址为空")
response = await client.get(payload, timeout=settings.request_timeout)
_raise_api_error(response, "通义")
return GenerationAdvance(data=response.content)
url = settings.dashscope_base_url or "https://dashscope.aliyuncs.com/api/v1/services/aigc/image-generation/generation"
content: list[dict] = []
if not _is_t2i_model(model):
content = [{"image": await _resolve_ref(item)} for item in ref_images]
content.append({"text": prompt})
body = {
"model": model,
"input": {"messages": [{"role": "user", "content": content}]},
"parameters": {"size": size, "n": 1, "watermark": False, "negative_prompt": DEFAULT_NEGATIVE_PROMPT[:500]},
}
async with httpx.AsyncClient(timeout=settings.request_timeout, verify=False) as client:
response = await client.post(url, headers={**headers, "X-DashScope-Async": "enable"}, json=body)
_raise_api_error(response, "通义")
task_id = response.json().get("output", {}).get("task_id", "")
if not task_id:
raise RuntimeError(f"通义万象未返回 task_id: {response.json()}")
return GenerationAdvance(provider_task_id=task_id, provider_task_kind="tongyi_wan")
if provider == "rightapi":
if not settings.rightapi_api_key:
raise RuntimeError("未配置 RIGHTAPI_API_KEY(.env)")
base = settings.rightapi_base_url.rstrip("/")
origin = base.split("/draw", 1)[0].rstrip("/") or base
headers = {"Authorization": f"Bearer {settings.rightapi_api_key}"}
async with httpx.AsyncClient(timeout=max(settings.request_timeout, 600), verify=False) as client:
if provider_task_id:
result = await _rightapi_poll_task_once(client, headers, origin, provider_task_id)
if result is None:
return GenerationAdvance(provider_task_id=provider_task_id, provider_task_kind="rightapi")
return GenerationAdvance(data=await _download_provider_image(client, result, "RightAPI"))
body: dict = {"model": model, "prompt": prompt, "n": 1, "size": size, "async": True}
if ref_images:
body["image"] = [await _resolve_ref(item) for item in ref_images]
response = await client.post(
f"{base}/v1/images/generations", headers={**headers, "Content-Type": "application/json"}, json=body,
)
_raise_api_error(response, "RightAPI")
submitted = response.json()
task_id = submitted.get("task_id") or ""
if task_id:
return GenerationAdvance(provider_task_id=task_id, provider_task_kind="rightapi")
return GenerationAdvance(data=await _download_provider_image(client, submitted, "RightAPI"))
if provider == "pcore":
if not settings.pcore_api_key:
raise RuntimeError("未配置 PCORE_API_KEY(.env)")
base = settings.pcore_base_url.rstrip("/")
headers = {"Authorization": f"Bearer {settings.pcore_api_key}"}
async with httpx.AsyncClient(timeout=max(settings.request_timeout, 600), verify=False) as client:
if provider_task_id:
kind = provider_task_kind or "edits"
result = await _pcore_poll_task_once(client, headers, base, kind, provider_task_id)
if result is None:
return GenerationAdvance(provider_task_id=provider_task_id, provider_task_kind=kind)
return GenerationAdvance(data=await _download_provider_image(client, result, "Pcore"))
kind = "edits" if ref_images else "generations"
body: dict = {"model": model, "prompt": prompt, "n": 1, "size": size, "async": True}
if ref_images:
body["images"] = [{"image_url": await _resolve_ref(item)} for item in ref_images]
response = await client.post(
f"{base}/v1/images/{kind}", headers={**headers, "Content-Type": "application/json"}, json=body,
)
_raise_api_error(response, "Pcore")
submitted = response.json()
if isinstance(submitted, dict) and submitted.get("error"):
error = submitted["error"] or {}
raise RuntimeError(f"Pcore 提交失败: {error.get('message') or submitted}")
task_id = submitted.get("task_id") or submitted.get("id") or ""
if not task_id:
raise RuntimeError(f"Pcore 提交未返回任务 id: {submitted}")
return GenerationAdvance(provider_task_id=task_id, provider_task_kind=kind)
raise RuntimeError(f"未知 provider: {provider}")
GENERATORS = {
"doubao": generate_doubao,
"tongyi": generate_tongyi,
@@ -524,12 +708,12 @@ def _refs_for_job(images: list[dict], job: dict) -> list[str]:
_GEN_LOCK = asyncio.Lock()
async def _store_image_bytes(task: Task, index: int, data: bytes) -> str:
async def _store_image_bytes(task: Task, index: int, data: bytes, version: str = "result") -> str:
"""生成图字节 → 存储落盘,返回可访问 URL。按 PNG 魔数定扩展名(部分中转不遵守 output_format)。"""
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"
return await get_storage().save_bytes(data, f"suites/{task.id}/{index}{ext}", ctype)
return await get_storage().save_bytes(data, f"suites/{task.id}/{version}/{index}{ext}", ctype)
async def run_suite(task: Task, on_image_ok: OnImageOk | None = None) -> None:
@@ -565,17 +749,24 @@ async def run_suite(task: Task, on_image_ok: OnImageOk | None = None) -> None:
task.images.append(image)
try:
# 提示词按模型家族分发:国产主体参考 / gpt edits 保真 / google 主体保持
# 单张 AI 编辑可传入 prompt_extra:任务展示名保留「AI 生图」,
# 但不把中文展示名拼进模型 prompt,避免生成图内出现中文。
prompt_extra = job.get("prompt_extra") or job
prompt = build_prompt(
provider_name, model, type_id, ctx, task.style_set, task.lang,
extra=job, style_prompt=task.style_prompt, requirements=task.requirements,
extra=prompt_extra, style_prompt=task.style_prompt, requirements=task.requirements,
)
refs = _refs_for_job(list(task.ref_images or []), job)
data = await generator(prompt, refs, size=size, model=model)
# 水印:AI 出图返回后、落盘前的后处理(失败不阻断,内部返回原图)
wm = task.watermark or {}
if wm.get("enabled"):
image.source_url = await _store_image_bytes(task, len(task.images), data, version="original")
data = apply_watermark(data, wm)
image.url = await _store_image_bytes(task, len(task.images), data)
image.url = await _store_image_bytes(task, len(task.images), data, version="watermarked")
image.watermark_config = wm
else:
image.url = await _store_image_bytes(task, len(task.images), data)
image.status = IMG_OK
ok += 1
if on_image_ok is not None and image.url:
+167
View File
@@ -0,0 +1,167 @@
"""商品图片生命周期:原图、当前水印版与存储回收的唯一入口。
每个图片素材只维护两个逻辑版本:original(无水印)与 current(页面展示/导出)。
未加水印时两者指向同一个对象;重加水印会先上传新 current,再切换引用并删除旧 current,
不保留历史水印版。
"""
from __future__ import annotations
import mimetypes
import uuid
from dataclasses import dataclass
from api.proxy import guess_referer
from models import ProductAsset
from services.storage import download_bytes, get_storage, image_extension, local_path, read_local_bytes, storage_key_from_url
from services.watermark import apply_watermark
@dataclass
class WatermarkUpdate:
"""水印切换后的旧对象;数据库提交成功后再回收。"""
obsolete_key: str | None = None
class ImageAssetService:
"""将 ProductAsset 与七牛/本地对象存储的生命周期收敛到同一服务。"""
def __init__(self) -> None:
self.storage = get_storage()
@staticmethod
def asset_prefix(asset: ProductAsset) -> str:
return f"products/{asset.product_id}/assets/{asset.id}/"
@staticmethod
def product_prefix(product_id: str) -> str:
return f"products/{product_id}/"
def _key(self, asset: ProductAsset, name: str, data: bytes, content_type: str) -> str:
return f"{self.asset_prefix(asset)}{name}{image_extension(data, content_type)}"
async def _read(self, url: str) -> tuple[bytes, str]:
path = local_path(url)
if path is not None:
return read_local_bytes(path), mimetypes.guess_type(path.name)[0] or "image/jpeg"
return await download_bytes(url, referer=guess_referer(url))
async def save_original(self, asset: ProductAsset, data: bytes, content_type: str = "") -> str:
"""写入该素材唯一的无水印原图,并将 current 同步到原图。"""
key = self._key(asset, "original", data, content_type)
url = await self.storage.save_bytes(data, key, content_type)
asset.original_url = url
asset.original_storage_key = key
asset.stored_url = url
asset.current_storage_key = key
asset.watermark_enabled = False
asset.watermark_config = None
return url
async def save_original_from_url(self, asset: ProductAsset, url: str) -> str:
data, content_type = await self._read(url)
return await self.save_original(asset, data, content_type)
async def ensure_original(self, asset: ProductAsset) -> tuple[bytes, str]:
"""为旧记录补齐原图字段;首次触碰旧素材时迁入新的商品/素材路径。"""
original_url = asset.original_url or asset.stored_url or asset.source_url
if not original_url:
raise ValueError("图片原图地址为空")
# 新格式已在当前素材路径中,无需复制。
if asset.original_storage_key and asset.original_url:
return await self._read(asset.original_url)
data, content_type = await self._read(original_url)
await self.save_original(asset, data, content_type)
return data, content_type
async def apply_watermark(self, asset: ProductAsset, opts: dict) -> WatermarkUpdate:
"""从 original 合成新水印版,调用方提交数据库后再删除返回的旧水印对象。"""
previous_current_key = asset.current_storage_key or storage_key_from_url(asset.stored_url or "")
data, content_type = await self.ensure_original(asset)
watermarked = apply_watermark(data, opts)
key = self._key(asset, f"watermark-{uuid.uuid4().hex}", watermarked, content_type)
url = await self.storage.save_bytes(watermarked, key, content_type)
obsolete_key = previous_current_key
# 未加水印时 current 就是 original,绝不能作为旧水印删除。
if obsolete_key == asset.original_storage_key:
obsolete_key = None
asset.stored_url = url
asset.current_storage_key = key
asset.watermark_enabled = True
asset.watermark_config = opts
return WatermarkUpdate(obsolete_key=obsolete_key)
async def adopt_generated_versions(
self,
asset: ProductAsset,
original_url: str,
current_url: str,
watermark_config: dict | None = None,
) -> None:
"""将任务临时输出迁入正式的商品/素材路径,不让 suites 临时对象长期残留。"""
await self.save_original_from_url(asset, original_url)
if current_url == original_url:
return
data, content_type = await self._read(current_url)
key = self._key(asset, f"watermark-{uuid.uuid4().hex}", data, content_type)
asset.stored_url = await self.storage.save_bytes(data, key, content_type)
asset.current_storage_key = key
asset.watermark_enabled = True
asset.watermark_config = watermark_config or {}
async def remove_watermark(self, asset: ProductAsset) -> WatermarkUpdate:
"""切回无水印原图,调用方提交数据库后再删除旧水印对象。"""
previous_current_key = asset.current_storage_key or storage_key_from_url(asset.stored_url or "")
await self.ensure_original(asset)
obsolete_key = previous_current_key
asset.stored_url = asset.original_url
asset.current_storage_key = asset.original_storage_key
asset.watermark_enabled = False
asset.watermark_config = None
if obsolete_key == asset.original_storage_key:
obsolete_key = None
return WatermarkUpdate(obsolete_key=obsolete_key)
async def delete_key(self, key: str | None) -> None:
if key:
await self.storage.delete(key)
async def delete_temporary_urls(self, urls: list[str]) -> None:
"""只回收 suites/ 临时输出,绝不删除作为参考图传入的正式商品素材。"""
for url in set(urls):
key = storage_key_from_url(url)
if key and key.startswith("suites/"):
await self.delete_key(key)
async def delete_asset_media(self, asset: ProductAsset) -> int:
"""删除单个素材的当前格式目录,并兼容清理由旧数据留下的路径外对象。"""
explicit_keys = {
key for key in (
asset.original_storage_key,
asset.current_storage_key,
storage_key_from_url(asset.original_url or ""),
storage_key_from_url(asset.stored_url or ""),
) if key
}
deleted = await self.storage.delete_prefix(self.asset_prefix(asset))
for key in explicit_keys:
if not key.startswith(self.asset_prefix(asset)):
await self.delete_key(key)
deleted += 1
return deleted
async def delete_product_media(self, product_id: str, assets: list[ProductAsset]) -> int:
"""删除商品下所有素材;先清理旧 key,再按商品前缀兜底清理孤儿对象。"""
deleted = 0
for asset in assets:
deleted += await self.delete_asset_media(asset)
deleted += await self.storage.delete_prefix(self.product_prefix(product_id))
return deleted
def get_image_asset_service() -> ImageAssetService:
return ImageAssetService()
+212
View File
@@ -0,0 +1,212 @@
"""独立后台 Worker:按子项粒度执行持久化图片任务。"""
from __future__ import annotations
import asyncio
import logging
import socket
import uuid
from config import get_settings
from db import get_session_factory
from models import ProductAsset
from services.image_assets import get_image_asset_service
from services.generator import _image_size, _refs_for_job, advance_generation
from services.jobs import (
ClaimedItem,
claim_next_item,
job_can_write,
mark_item_cancelled,
mark_item_done,
mark_item_failed,
mark_waiting_provider,
recover_expired_leases,
)
from services.prompts import build_context, build_prompt
from services.storage import get_storage
from services.suite_service import append_generated_asset
from services.tasks import TaskImage
from services.watermark import apply_watermark
log = logging.getLogger("background.worker")
def _worker_owner() -> str:
return f"{socket.gethostname()}-{uuid.uuid4().hex[:10]}"
def _image_meta(data: bytes) -> tuple[str, str]:
return (".png", "image/png") if data[:8] == b"\x89PNG\r\n\x1a\n" else (".jpg", "image/jpeg")
async def _store_image(job_id: str, position: int, data: bytes, version: str) -> str:
suffix, content_type = _image_meta(data)
return await get_storage().save_bytes(data, f"suites/{job_id}/{version}/{position}{suffix}", content_type)
async def _cleanup_temporary(urls: list[str]) -> None:
"""append 成功后由 image_assets 搬运并清理;取消时主动清理本次临时文件。"""
if not urls:
return
try:
from services.image_assets import get_image_asset_service
await get_image_asset_service().delete_temporary_urls(urls)
except Exception: # noqa: BLE001
log.exception("清理取消任务的临时图片失败")
async def _process_watermark_item(claim: ClaimedItem) -> None:
"""批量水印的一个素材子项:基于 original 重做 current,提交后回收旧水印对象。"""
if not await job_can_write(claim.job_id):
await mark_item_cancelled(claim.item_id)
return
payload = claim.item_payload
raw_asset_id = payload.get("asset_id")
try:
asset_id = uuid.UUID(str(raw_asset_id))
except (TypeError, ValueError) as exc:
await mark_item_failed(claim.item_id, "水印任务缺少合法素材 ID")
return
obsolete_key: str | None = None
result_url = ""
source_url = ""
watermark_config: dict | None = None
try:
async with get_session_factory()() as db:
asset = await db.get(ProductAsset, asset_id)
if asset is None:
raise RuntimeError("素材已删除")
if asset.type != "img":
raise RuntimeError("仅支持图片素材")
expected_key = payload.get("expected_current_storage_key")
expected_url = payload.get("expected_current_url")
if expected_key and asset.current_storage_key != expected_key:
raise RuntimeError("素材在等待期间已被其他操作更新,请重新提交水印任务")
if not expected_key and expected_url and asset.stored_url != expected_url:
raise RuntimeError("素材在等待期间已被其他操作更新,请重新提交水印任务")
watermark = claim.request_payload.get("watermark") or {}
if not watermark.get("enabled"):
raise RuntimeError("水印设置已关闭")
image_assets = get_image_asset_service()
update = await image_assets.apply_watermark(asset, watermark)
asset.error = None
result_url = asset.stored_url or ""
source_url = asset.original_url or ""
watermark_config = asset.watermark_config
obsolete_key = update.obsolete_key
await db.commit()
# 数据库已切换到新版本后再删旧对象;删除失败不影响用户可用的新图片。
if obsolete_key:
try:
await get_image_asset_service().delete_key(obsolete_key)
except Exception: # noqa: BLE001
log.exception("水印已切换,但旧水印对象删除失败:%s", obsolete_key)
await mark_item_done(
claim.item_id, result_url=result_url, source_url=source_url or None,
watermark_config=watermark_config, generated_asset_id=None,
)
except Exception as exc: # noqa: BLE001
log.exception("后台水印任务 %s/%s 执行失败", claim.job_id, claim.item_id)
await mark_item_failed(claim.item_id, str(exc))
async def process_item(claim: ClaimedItem) -> None:
"""推进一次子任务。外部模型未完成时只写 provider id 并退出,让出 Worker。"""
if claim.job_kind == "watermark_batch":
await _process_watermark_item(claim)
return
if claim.job_kind not in {"image_suite", "image_edit"}:
await mark_item_failed(claim.item_id, f"Worker 暂不支持任务类型: {claim.job_kind}")
return
settings = get_settings()
payload = claim.request_payload
job_item = claim.item_payload
provider = claim.provider or settings.image_provider
model = claim.model or {
"tongyi": settings.dashscope_model,
"rightapi": settings.rightapi_image_model,
"pcore": "gpt-image-2",
}.get(provider, settings.ark_image_model)
try:
context = build_context(payload.get("context") or {}, fallback_name="")
prompt_extra = job_item.get("prompt_extra") or job_item
prompt = build_prompt(
provider, model, job_item.get("kind", "custom"), context, claim.style_set, claim.lang,
extra=prompt_extra, style_prompt=payload.get("style_prompt"), requirements=payload.get("requirements"),
)
refs = _refs_for_job(list(payload.get("ref_images") or []), job_item)
is_wan = provider == "tongyi" and model.lower().startswith("wan")
result = await advance_generation(
provider, prompt, refs,
size=_image_size(provider, claim.ratio, is_wan=is_wan, model=model), model=model,
provider_task_id=claim.provider_task_id, provider_task_kind=claim.provider_task_kind,
)
if result.data is None:
if not result.provider_task_id or not result.provider_task_kind:
raise RuntimeError("异步生成未返回可恢复的 provider task id")
await mark_waiting_provider(claim.item_id, result.provider_task_id, result.provider_task_kind)
return
# 生成完成后仍检查取消/商品删除,避免在用户删除商品后把素材回写回来。
if not await job_can_write(claim.job_id):
await mark_item_cancelled(claim.item_id)
return
watermark = payload.get("watermark") or {}
position = int(job_item.get("position", 0))
if watermark.get("enabled"):
source_url = await _store_image(claim.job_id, position, result.data, "original")
result_url = await _store_image(claim.job_id, position, apply_watermark(result.data, watermark), "watermarked")
watermark_config = watermark
else:
source_url = ""
result_url = await _store_image(claim.job_id, position, result.data, "result")
watermark_config = None
asset_id: str | None = None
if payload.get("append") and claim.product_id:
image = TaskImage(
type_id=job_item.get("kind", "custom"),
name=job_item.get("title") or "AI 生图",
url=result_url,
source_url=source_url,
watermark_config=watermark_config,
)
try:
asset_id = await append_generated_asset(
claim.product_id, image, after_asset_id=payload.get("after_asset_id"),
)
except Exception:
# 回写失败时保留 suites 临时文件和失败状态,避免用户看不到实际结果而又无法诊断。
raise
result_url, source_url, watermark_config = image.url, image.source_url, image.watermark_config
await mark_item_done(
claim.item_id, result_url=result_url, source_url=source_url or None,
watermark_config=watermark_config, generated_asset_id=asset_id,
)
except Exception as exc: # noqa: BLE001
log.exception("后台图片任务 %s/%s 执行失败", claim.job_id, claim.item_id)
await mark_item_failed(claim.item_id, str(exc))
async def run_worker(*, once: bool = False, idle_seconds: float = 0.75) -> None:
owner = _worker_owner()
recovered = await recover_expired_leases()
if recovered:
log.warning("恢复了 %d 个过期后台任务租约", recovered)
log.info("后台 Worker 已启动:%s", owner)
while True:
claim = await claim_next_item(owner)
if claim is None:
if once:
return
await asyncio.sleep(idle_seconds)
continue
await process_item(claim)
if once:
return
+470
View File
@@ -0,0 +1,470 @@
"""数据库驱动的后台任务仓库与调度状态机。
已承载图片生成与批量水印;表结构刻意通用,后续可复用到 ZIP 导出等耗时操作。
Worker 不依赖 FastAPI 进程;第三方 provider task id 保存在子项中,重启后继续轮询。
"""
from __future__ import annotations
import uuid
from dataclasses import dataclass
from datetime import datetime, timedelta, timezone
from typing import Any
from sqlalchemy import case, func, or_, select, update
from db import get_session_factory
from models import BackgroundJob, BackgroundJobItem
JOB_QUEUED = "queued"
JOB_RUNNING = "running"
JOB_DONE = "done"
JOB_PARTIAL = "partial"
JOB_FAILED = "failed"
JOB_CANCELLED = "cancelled"
ITEM_QUEUED = "queued"
ITEM_RUNNING = "running"
ITEM_WAITING_PROVIDER = "waiting_provider"
ITEM_DONE = "done"
ITEM_FAILED = "failed"
ITEM_CANCELLED = "cancelled"
TERMINAL_ITEM_STATUSES = (ITEM_DONE, ITEM_FAILED, ITEM_CANCELLED)
TERMINAL_JOB_STATUSES = (JOB_DONE, JOB_PARTIAL, JOB_FAILED, JOB_CANCELLED)
def utcnow() -> datetime:
return datetime.now(timezone.utc)
@dataclass(frozen=True)
class ClaimedItem:
job_id: str
item_id: str
job_kind: str
product_id: str | None
provider: str
model: str | None
platform: str
lang: str
ratio: str
style_set: int
request_payload: dict[str, Any]
item_payload: dict[str, Any]
provider_task_id: str | None
provider_task_kind: str | None
async def create_image_job(
*,
product_id: str | None,
kind: str,
priority: int,
provider: str,
model: str | None,
platform: str,
lang: str,
ratio: str,
style_set: int,
request_payload: dict[str, Any],
item_payloads: list[dict[str, Any]],
) -> BackgroundJob:
"""一次事务写入父任务和逐图子项,API 成功返回即意味着任务可由 Worker 恢复。"""
parsed_product_id = uuid.UUID(product_id) if product_id else None
job = BackgroundJob(
product_id=parsed_product_id,
kind=kind,
status=JOB_QUEUED,
priority=priority,
provider=provider,
model=model,
platform=platform,
lang=lang,
ratio=ratio,
style_set=style_set,
request_payload=request_payload,
total=len(item_payloads),
)
async with get_session_factory()() as db:
db.add(job)
await db.flush()
db.add_all([
BackgroundJobItem(job_id=job.id, position=index, payload=payload, status=ITEM_QUEUED)
for index, payload in enumerate(item_payloads)
])
await db.commit()
await db.refresh(job)
return job
async def _refresh_job_status(db, job: BackgroundJob) -> None:
"""从子项计算父任务进度与终态;调用方已在同一事务中修改子项。"""
rows = (await db.execute(
select(BackgroundJobItem.status, func.count(BackgroundJobItem.id))
.where(BackgroundJobItem.job_id == job.id)
.group_by(BackgroundJobItem.status)
)).all()
counts = {status: int(count) for status, count in rows}
done = counts.get(ITEM_DONE, 0)
failed = counts.get(ITEM_FAILED, 0)
cancelled = counts.get(ITEM_CANCELLED, 0)
terminal = done + failed + cancelled
job.completed = done
job.failed = failed
if terminal < job.total:
job.status = JOB_RUNNING
return
job.completed_at = utcnow()
if done == 0 and cancelled == job.total:
job.status = JOB_CANCELLED
elif failed == 0 and cancelled == 0:
job.status = JOB_DONE
elif done > 0:
job.status = JOB_PARTIAL
elif job.cancel_requested and cancelled > 0:
job.status = JOB_CANCELLED
else:
job.status = JOB_FAILED
errors = (await db.scalars(
select(BackgroundJobItem.error)
.where(BackgroundJobItem.job_id == job.id, BackgroundJobItem.status == ITEM_FAILED)
.order_by(BackgroundJobItem.position)
)).all()
if errors:
brief = [str(error)[:200] for error in errors if error]
job.error = ";".join(brief[:6]) or "任务失败"
async def recover_expired_leases() -> int:
"""把上个 Worker 遗留的工作单元重新放回队列;provider id 不会被清空。"""
now = utcnow()
async with get_session_factory()() as db:
result = await db.execute(
update(BackgroundJobItem)
.where(
BackgroundJobItem.status == ITEM_RUNNING,
BackgroundJobItem.lease_expires_at.is_not(None),
BackgroundJobItem.lease_expires_at < now,
)
.values(status=ITEM_QUEUED, lease_owner=None, lease_expires_at=None, next_run_at=now)
)
await db.commit()
return int(result.rowcount or 0)
async def claim_next_item(owner: str, lease_seconds: int = 90) -> ClaimedItem | None:
"""领取一个到期子项。单机 SQLite Worker 使用租约,未来迁移 PostgreSQL 后可直接强化为 SKIP LOCKED。"""
now = utcnow()
async with get_session_factory()() as db:
# 先回收过期租约;即便 Worker 非正常退出,也不会永久卡住。
await db.execute(
update(BackgroundJobItem)
.where(
BackgroundJobItem.status == ITEM_RUNNING,
BackgroundJobItem.lease_expires_at.is_not(None),
BackgroundJobItem.lease_expires_at < now,
)
.values(status=ITEM_QUEUED, lease_owner=None, lease_expires_at=None, next_run_at=now)
)
query = (
select(BackgroundJobItem, BackgroundJob)
.join(BackgroundJob, BackgroundJob.id == BackgroundJobItem.job_id)
.where(
BackgroundJob.cancel_requested.is_(False),
BackgroundJob.status.not_in(TERMINAL_JOB_STATUSES),
BackgroundJobItem.status.in_((ITEM_QUEUED, ITEM_WAITING_PROVIDER)),
or_(BackgroundJobItem.next_run_at.is_(None), BackgroundJobItem.next_run_at <= now),
)
.order_by(BackgroundJob.priority.desc(), BackgroundJob.created_at.asc(), BackgroundJobItem.position.asc())
.limit(1)
)
row = (await db.execute(query)).first()
if row is None:
await db.commit()
return None
item, job = row
item.status = ITEM_RUNNING
if not item.provider_task_id:
item.attempts += 1
item.lease_owner = owner
item.lease_expires_at = now + timedelta(seconds=lease_seconds)
item.next_run_at = None
item.started_at = item.started_at or now
job.status = JOB_RUNNING
job.started_at = job.started_at or now
await db.commit()
return ClaimedItem(
job_id=str(job.id), item_id=str(item.id), job_kind=job.kind,
product_id=str(job.product_id) if job.product_id else None,
provider=job.provider, model=job.model, platform=job.platform,
lang=job.lang, ratio=job.ratio, style_set=job.style_set,
request_payload=dict(job.request_payload or {}), item_payload=dict(item.payload or {}),
provider_task_id=item.provider_task_id, provider_task_kind=item.provider_task_kind,
)
async def mark_waiting_provider(item_id: str, provider_task_id: str, provider_task_kind: str, delay_seconds: int = 3) -> None:
"""先持久化外部任务 ID,再安排下一次查询,是避免重启重复扣费的关键提交点。"""
async with get_session_factory()() as db:
item = await db.get(BackgroundJobItem, uuid.UUID(item_id))
if item is None:
return
item.provider_task_id = provider_task_id
item.provider_task_kind = provider_task_kind
item.status = ITEM_WAITING_PROVIDER
item.lease_owner = None
item.lease_expires_at = None
item.next_run_at = utcnow() + timedelta(seconds=delay_seconds)
await db.commit()
async def mark_item_done(
item_id: str,
*,
result_url: str,
source_url: str | None,
watermark_config: dict | None,
generated_asset_id: str | None,
) -> None:
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)
item.status = ITEM_DONE
item.result_url = result_url
item.source_url = source_url
item.watermark_config = watermark_config
item.generated_asset_id = uuid.UUID(generated_asset_id) if generated_asset_id else None
item.error = None
item.lease_owner = None
item.lease_expires_at = None
item.completed_at = utcnow()
if job is not None:
await _refresh_job_status(db, job)
await db.commit()
async def mark_item_failed(item_id: str, error: str) -> None:
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)
item.status = ITEM_FAILED
item.error = error[:500]
item.lease_owner = None
item.lease_expires_at = None
item.completed_at = utcnow()
if job is not None:
await _refresh_job_status(db, job)
await db.commit()
async def job_can_write(job_id: str) -> bool:
"""结果落盘前的最后一道确认:取消或商品被硬删除时禁止回写素材。"""
async with get_session_factory()() as db:
job = await db.get(BackgroundJob, uuid.UUID(job_id))
if job is None or job.cancel_requested:
return False
if job.product_id is None:
return True
# 仅用主键查询避免导入 Product 形成循环;不存在即不允许生成结果复活商品。
from models import Product
return await db.get(Product, job.product_id) is not None
async def mark_item_cancelled(item_id: str) -> None:
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)
item.status = ITEM_CANCELLED
item.lease_owner = None
item.lease_expires_at = None
item.completed_at = utcnow()
if job is not None:
await _refresh_job_status(db, job)
await db.commit()
async def cancel_job(job_id: str) -> BackgroundJob | None:
async with get_session_factory()() as db:
job = await db.get(BackgroundJob, uuid.UUID(job_id))
if job is None:
return None
if job.status in TERMINAL_JOB_STATUSES:
return job
job.cancel_requested = True
await db.execute(
update(BackgroundJobItem)
.where(BackgroundJobItem.job_id == job.id, BackgroundJobItem.status.in_((ITEM_QUEUED, ITEM_WAITING_PROVIDER)))
.values(status=ITEM_CANCELLED, completed_at=utcnow(), lease_owner=None, lease_expires_at=None)
)
await _refresh_job_status(db, job)
await db.commit()
await db.refresh(job)
return job
async def cancel_product_jobs(product_id: str) -> int:
"""硬删除商品前取消其未完成任务,Worker 随后只会丢弃结果而不会回写素材。"""
try:
parsed = uuid.UUID(product_id)
except ValueError:
return 0
async with get_session_factory()() as db:
jobs = (await db.scalars(
select(BackgroundJob).where(
BackgroundJob.product_id == parsed,
BackgroundJob.status.not_in(TERMINAL_JOB_STATUSES),
)
)).all()
now = utcnow()
for job in jobs:
job.cancel_requested = True
await db.execute(
update(BackgroundJobItem)
.where(BackgroundJobItem.job_id == job.id, BackgroundJobItem.status.in_((ITEM_QUEUED, ITEM_WAITING_PROVIDER)))
.values(status=ITEM_CANCELLED, completed_at=now, lease_owner=None, lease_expires_at=None)
)
await _refresh_job_status(db, job)
await db.commit()
return len(jobs)
async def retry_failed_items(job_id: str) -> BackgroundJob | None:
"""仅把失败子项重新排队;外部失败任务不可复用,清空 provider id 后重新提交。"""
async with get_session_factory()() as db:
job = await db.get(BackgroundJob, uuid.UUID(job_id))
if job is None:
return None
now = utcnow()
result = await db.execute(
update(BackgroundJobItem)
.where(BackgroundJobItem.job_id == job.id, BackgroundJobItem.status == ITEM_FAILED)
.values(
status=ITEM_QUEUED, error=None, provider_task_id=None, provider_task_kind=None,
result_url=None, source_url=None, watermark_config=None, generated_asset_id=None,
next_run_at=now, lease_owner=None, lease_expires_at=None, completed_at=None,
)
)
if not result.rowcount:
return job
job.cancel_requested = False
job.status = JOB_QUEUED
job.completed_at = None
job.error = None
await _refresh_job_status(db, job)
await db.commit()
await db.refresh(job)
return job
async def retry_failed_item(job_id: str, item_id: str) -> BackgroundJob | None:
"""只重试一个失败子项,供批量水印等任务逐图补偿。"""
try:
parsed_job = uuid.UUID(job_id)
parsed_item = uuid.UUID(item_id)
except ValueError:
return None
async with get_session_factory()() as db:
job = await db.get(BackgroundJob, parsed_job)
item = await db.get(BackgroundJobItem, parsed_item)
if job is None or item is None or item.job_id != job.id:
return None
if item.status != ITEM_FAILED:
return job
# 水印任务等待期间,素材可能被用户再次处理过。逐项重试时改为以
# 最新 current 版本为基线,避免把已经更新的图片误判为永久失败。
if job.kind == "watermark_batch":
raw_asset_id = (item.payload or {}).get("asset_id")
try:
from models import ProductAsset
asset = await db.get(ProductAsset, uuid.UUID(str(raw_asset_id)))
except (TypeError, ValueError):
asset = None
if asset is not None:
payload = dict(item.payload or {})
payload["expected_current_storage_key"] = asset.current_storage_key
payload["expected_current_url"] = asset.stored_url
item.payload = payload
item.status = ITEM_QUEUED
item.error = None
item.provider_task_id = None
item.provider_task_kind = None
item.result_url = None
item.source_url = None
item.watermark_config = None
item.generated_asset_id = None
item.next_run_at = utcnow()
item.lease_owner = None
item.lease_expires_at = None
item.completed_at = None
job.cancel_requested = False
job.status = JOB_QUEUED
job.completed_at = None
job.error = None
await _refresh_job_status(db, job)
await db.commit()
await db.refresh(job)
return job
async def get_job(job_id: str) -> tuple[BackgroundJob, list[BackgroundJobItem]] | None:
try:
parsed = uuid.UUID(job_id)
except ValueError:
return None
async with get_session_factory()() as db:
job = await db.get(BackgroundJob, parsed)
if job is None:
return None
items = (await db.scalars(
select(BackgroundJobItem).where(BackgroundJobItem.job_id == job.id).order_by(BackgroundJobItem.position)
)).all()
return job, items
async def list_product_jobs(
product_id: str,
active_only: bool = False,
limit: int = 20,
kinds: tuple[str, ...] | None = None,
) -> list[BackgroundJob]:
try:
parsed = uuid.UUID(product_id)
except ValueError:
return []
async with get_session_factory()() as db:
query = select(BackgroundJob).where(BackgroundJob.product_id == parsed)
if kinds:
query = query.where(BackgroundJob.kind.in_(kinds))
if active_only:
query = query.where(BackgroundJob.status.not_in(TERMINAL_JOB_STATUSES))
return (await db.scalars(query.order_by(BackgroundJob.created_at.desc()).limit(limit))).all()
def suite_view(job: BackgroundJob, items: list[BackgroundJobItem]) -> dict[str, Any]:
"""兼容现有 /suites/{id} 契约,同时增加任务类型、取消状态与子项状态。"""
return {
"id": str(job.id), "status": job.status, "kind": job.kind,
"style_set": job.style_set, "platform": job.platform, "lang": job.lang,
"ratio": job.ratio, "provider": job.provider, "model": job.model,
"total": job.total, "completed": job.completed, "failed": job.failed,
"cancel_requested": job.cancel_requested, "error": job.error,
"created_at": job.created_at, "started_at": job.started_at, "completed_at": job.completed_at,
"images": [
{
"type_id": (item.payload or {}).get("kind", "custom"),
"item_id": str(item.id),
"name": (item.payload or {}).get("title", "AI 生图"),
"url": item.result_url or "", "status": item.status,
"error": item.error,
}
for item in items
],
}
+5 -11
View File
@@ -7,7 +7,7 @@
from __future__ import annotations
from .common import (
STYLE_SETS, TEXT_RENDER, requirements_block, resolve_style, selling_point_lines,
STYLE_SETS, TEXT_RENDER, params_line_for_output, requirements_block, resolve_style, selling_point_lines,
)
# ── 公共组件(主体参考语义专用)────────────────────────────────────────────
@@ -85,7 +85,7 @@ def _prompt_multi_scene(ctx: dict, style: dict, lang: str) -> str:
def _prompt_ecommerce_detail(ctx: dict, style: dict, lang: str) -> str:
sp = selling_point_lines(ctx, lang) or ctx["title"]
params = ctx["params_line"]
params = params_line_for_output(ctx, lang)
return (
f"E-commerce detail-page hero section for product \"{ctx['title']}\", square layout: "
f"top half is a hero banner with the product at a 3/4 angle ({style['bg']}); "
@@ -95,7 +95,7 @@ def _prompt_ecommerce_detail(ctx: dict, style: dict, lang: str) -> str:
)
def _prompt_size_chart(ctx: dict, style: dict, lang: str) -> str:
dims = ctx["params_line"]
dims = params_line_for_output(ctx, lang)
return (
f"Product size chart infographic for \"{ctx['title']}\": product shown in clean front and side views "
f"on light background, with thin measurement annotation lines (arrows) marking length, width and height, "
@@ -115,15 +115,9 @@ def _prompt_sku_collection(ctx: dict, style: dict, lang: str) -> str:
def _prompt_custom(ctx: dict, style: dict, lang: str, extra: dict) -> str:
hint = (extra.get("prompt_hint") or "").strip()
purpose = extra.get("title") or ""
detail = extra.get("detail") or ""
bg = f" {style['bg']} as environment." if style.get("bg") else ""
composed = (
f"E-commerce marketing image for product \"{ctx['title']}\""
+ (f" — {purpose}" if purpose else "")
+ (f": {detail}" if detail else "")
+ "."
)
# title/detail 是规划器给前端展示的中文文案;只使用约定为英文的 prompt_hint,避免它们变成图中文字。
composed = f"E-commerce marketing image for product \"{ctx['title']}\"."
if hint:
composed += f" Composition: {hint}."
return f"{composed}{bg} {style['tone']}. {TEXT_RENDER[lang]} {QUALITY} {PRODUCT_REF_LOCK}"
+28 -9
View File
@@ -12,27 +12,27 @@ import re
STYLE_SETS: dict[int, dict] = {
1: {
"name": "北欧极简",
"tone": "北欧极简风:浅灰或米白背景,柔和漫射光,低饱和色调,画面留白充足,构图克制干净",
"tone": "Nordic minimalism: light gray or warm off-white background, soft diffused light, low-saturation palette, generous negative space, restrained and clean composition",
"bg": "",
},
2: {
"name": "清新明亮",
"tone": "清新明亮风:明亮的白色到浅蓝渐变背景,高调光线,色彩明快通透,整体轻盈干净",
"tone": "Fresh and bright: white-to-pale-blue gradient background, high-key lighting, clear luminous colors, airy and uncluttered composition",
"bg": "",
},
3: {
"name": "高级感深色",
"tone": "高级质感风:深灰或炭黑背景,戏剧性侧光打光,突出商品材质与光泽,沉稳高级",
"tone": "Premium dark aesthetic: deep gray or charcoal background, dramatic side lighting, emphasize material texture and sheen, refined and restrained mood",
"bg": "",
},
4: {
"name": "暖调生活",
"tone": "温暖生活风:暖米色背景,暖色灯光氛围,温馨的家居质感,亲和力强",
"tone": "Warm lifestyle: warm beige background, cozy warm lighting, inviting home-living texture, approachable and natural atmosphere",
"bg": "",
},
5: {
"name": "纯净棚拍",
"tone": "标准电商棚拍:纯色浅背景,均匀的正面柔光,无杂物干扰,商品居中突出",
"tone": "Clean e-commerce studio: plain light background, even frontal softbox lighting, no distracting props, product centered and prominent",
"bg": "",
},
}
@@ -70,6 +70,9 @@ TEXT_RENDER = {
"Render concise Russian marketing text inside the image: headline max 4 words, "
"sub-lines max 6 words each, modern clean sans-serif (Inter / PT Sans style), "
"proper Cyrillic typography, high legibility, tasteful layout, colors harmonized with the composition. "
"Translate any Chinese source facts into Russian before rendering. All newly generated visible text "
"must be Russian Cyrillic only: never render Chinese, Hanzi, CJK characters, or Chinese phrases. "
"The sole exception is Chinese text physically printed on the reference product itself, which must remain unchanged. "
"No spelling errors, no mixed latin/cyrillic gibberish."
),
}
@@ -86,8 +89,8 @@ def requirements_block(requirements: str | None) -> str:
"""用户强制要求块:最高优先级、置于提示词最前、覆盖冲突指令(原文保留不翻译)。"""
if requirements and requirements.strip():
return (
"STRICT REQUIREMENTS (highest priority, must be followed exactly, "
"override any conflicting instruction): "
"STRICT REQUIREMENTS (highest priority for visual intent; interpret any Chinese source instruction "
"by meaning only and never copy its Hanzi into newly generated text): "
+ requirements.strip().rstrip(".")
+ "."
)
@@ -110,7 +113,8 @@ def build_context(raw: dict, fallback_name: str = "", fallback_desc: str = "") -
raw: {title, desc, price, params: [{key, value}], sellingPoints}
"""
title = _clean_title(raw.get("title") or fallback_name or "product")
# 试算页会同时上传中文采集标题和已回填的俄文标题;后者用于图内可见文案。
title = _clean_title(raw.get("title_ru") or raw.get("title") or fallback_name or "product")
desc = _shorten(raw.get("desc") or fallback_desc or "", 200)
# 卖点:优先显式卖点文本;否则从参数表里挑短而有信息量的键值对
@@ -138,6 +142,12 @@ def build_context(raw: dict, fallback_name: str = "", fallback_desc: str = "") -
"desc": desc,
"selling_points": selling_points[:3],
"params_line": params_line,
# 俄文图片里不能直接渲染中文参数;保留不含汉字的规格(型号、数字、单位)供尺寸图使用。
"params_line_non_cjk": "; ".join(
f"{p.get('key')}: {p.get('value')}"
for p in (raw.get("params") or [])[:8]
if not re.search(r"[\u3400-\u9fff]", f"{p.get('key')} {p.get('value')}")
),
"price": raw.get("price") or "",
}
@@ -147,7 +157,16 @@ def selling_point_lines(ctx: dict, lang: str, max_n: int = 3) -> str:
if not sps:
return ""
key = "zh" if lang == "zh" else "en"
return "; ".join(s[key] for s in sps if s.get(key))
values = [s[key] for s in sps if s.get(key)]
# 中文卖点是模型理解商品的来源,不应成为俄文图片中的逐字文案。
if lang == "ru":
values = [v for v in values if not re.search(r"[\u3400-\u9fff]", v)]
return "; ".join(values)
def params_line_for_output(ctx: dict, lang: str) -> str:
"""返回可进入图内规格文案的参数;俄文模式剔除含汉字的原始采集参数。"""
return ctx.get("params_line") if lang == "zh" else ctx.get("params_line_non_cjk", "")
def type_name(type_id: str) -> str:
+5 -9
View File
@@ -9,7 +9,7 @@
"""
from __future__ import annotations
from .common import TEXT_RENDER, requirements_block, resolve_style, selling_point_lines
from .common import TEXT_RENDER, params_line_for_output, requirements_block, resolve_style, selling_point_lines
_SUBJECT_LOCK = (
"SUBJECT LOCK (highest priority): the product in the first image is the subject. "
@@ -32,8 +32,8 @@ _REMINDER = (
def _anchor(ctx: dict) -> str:
"""商品文字锚定:仅供识别,明确以图为准(同 gpt 模块,避免文字反噬商品)。"""
line = f"Context (identification only): the product is \"{ctx['title']}\""
if ctx.get("params_line"):
line += f" ({ctx['params_line']})"
if ctx.get("params_line_non_cjk"):
line += f" ({ctx['params_line_non_cjk']})"
return line + ". The image, not this text, defines the product's appearance."
@@ -88,7 +88,7 @@ def _task_multi_scene(ctx: dict, lang: str) -> str:
def _task_ecommerce_detail(ctx: dict, lang: str) -> str:
sp = selling_point_lines(ctx, lang) or ctx["title"]
params = ctx["params_line"]
params = params_line_for_output(ctx, lang)
return (
"Compose a square detail-page hero section: top half a hero banner with the unchanged product "
"at a 3/4 angle; bottom half a clean spec card with 3 feature rows and line icons"
@@ -97,7 +97,7 @@ def _task_ecommerce_detail(ctx: dict, lang: str) -> str:
)
def _task_size_chart(ctx: dict, lang: str) -> str:
dims = ctx["params_line"]
dims = params_line_for_output(ctx, lang)
return (
"Create a size chart: the unchanged product in clean front and side views on a light background, "
"thin measurement annotation lines (arrows) marking length, width and height with values beside "
@@ -135,11 +135,7 @@ def build_prompt(type_id: str, ctx: dict, style_set: int, lang: str, extra: dict
hint = (extra.get("prompt_hint") or "").strip()
if type_id == "custom":
purpose = extra.get("title") or ""
detail = extra.get("detail") or ""
task = "Create an e-commerce marketing image featuring the product from the first image"
task += f" — {purpose}" if purpose else ""
task += f": {detail}" if detail else ""
task += "."
wants_text = True
else:
+5 -9
View File
@@ -15,7 +15,7 @@
"""
from __future__ import annotations
from .common import TEXT_RENDER, requirements_block, resolve_style, selling_point_lines
from .common import TEXT_RENDER, params_line_for_output, requirements_block, resolve_style, selling_point_lines
# 保真锁:商品由 Image 1 唯一定义,其余指令一律不得触碰商品本体
_PRESERVE = (
@@ -55,8 +55,8 @@ def _anchor(ctx: dict) -> str:
文字锚定用于帮模型认出"是哪件商品",而不是"长什么样"。
"""
line = f"CONTEXT (identification only): the product is \"{ctx['title']}\""
if ctx.get("params_line"):
line += f" ({ctx['params_line']})"
if ctx.get("params_line_non_cjk"):
line += f" ({ctx['params_line_non_cjk']})"
return (
line
+ ". Image 1 — not this text — defines the product's appearance; "
@@ -119,7 +119,7 @@ def _task_multi_scene(ctx: dict, lang: str) -> str:
def _task_ecommerce_detail(ctx: dict, lang: str) -> str:
sp = selling_point_lines(ctx, lang) or ctx["title"]
params = ctx["params_line"]
params = params_line_for_output(ctx, lang)
return (
"TASK: Detail-page hero section on a square canvas. Top half: hero banner with the unchanged "
"product at a 3/4 angle. Bottom half: clean spec card with 3 feature rows and line icons"
@@ -128,7 +128,7 @@ def _task_ecommerce_detail(ctx: dict, lang: str) -> str:
)
def _task_size_chart(ctx: dict, lang: str) -> str:
dims = ctx["params_line"]
dims = params_line_for_output(ctx, lang)
return (
"TASK: Measurement chart. Show the unchanged product in clean front and side views on a light "
"background; overlay thin technical annotation lines (arrows) marking length, width and height "
@@ -167,11 +167,7 @@ def build_prompt(type_id: str, ctx: dict, style_set: int, lang: str, extra: dict
hint = (extra.get("prompt_hint") or "").strip()
if type_id == "custom":
purpose = extra.get("title") or ""
detail = extra.get("detail") or ""
task = "TASK: Create an e-commerce marketing image featuring the product from Image 1"
task += f" — {purpose}" if purpose else ""
task += f": {detail}" if detail else ""
task += "."
wants_text = True
else:
+142 -23
View File
@@ -2,6 +2,7 @@
from __future__ import annotations
import mimetypes
import shutil
import uuid
from pathlib import Path
@@ -13,11 +14,16 @@ from config import get_settings
_LOCAL_ROOT = Path(__file__).resolve().parents[2] / "data" / "media"
def _ext_from_url(url: str) -> str:
ext = mimetypes.guess_extension(url.split("?")[0].lower()) or ".jpg"
if ext == ".jpe":
ext = ".jpg"
return ext
class ResourceSizeLimitError(ValueError):
"""上传、下载或本地读取超过单文件上限。"""
def __init__(self, limit_bytes: int, actual_bytes: int | None = None) -> None:
self.limit_bytes = limit_bytes
self.actual_bytes = actual_bytes
detail = f"图片文件超过大小上限(最多 {limit_bytes // 1024 // 1024} MB)"
if actual_bytes is not None:
detail += f",当前约 {actual_bytes / 1024 / 1024:.1f} MB"
super().__init__(detail)
def _ext_from_content_type(content_type: str) -> str:
@@ -34,40 +40,109 @@ def _ext_from_content_type(content_type: str) -> str:
return mapping.get(ctype, ".jpg")
async def download_bytes(url: str, referer: str | None = None, timeout: float = 60.0) -> tuple[bytes, str]:
"""下载远程字节。返回 (bytes, content_type)。"""
def image_extension(data: bytes, content_type: str = "") -> str:
"""按真实图片字节优先推断扩展名,避免 provider 错报 Content-Type。"""
if data[:8] == b"\x89PNG\r\n\x1a\n":
return ".png"
if data[:4] == b"RIFF" and data[8:12] == b"WEBP":
return ".webp"
if data[:3] == b"GIF":
return ".gif"
return _ext_from_content_type(content_type)
async def download_bytes(
url: str,
referer: str | None = None,
timeout: float = 60.0,
max_bytes: int | None = None,
) -> tuple[bytes, str]:
"""流式下载远程图片并限制体积。返回 (bytes, content_type)。"""
limit = max_bytes or get_settings().image_download_max_bytes
if limit <= 0:
raise ValueError("图片下载大小上限必须大于 0")
headers = {"Referer": referer} if referer else {}
async with httpx.AsyncClient(timeout=timeout, follow_redirects=True) as client:
resp = await client.get(url, headers=headers)
resp.raise_for_status()
ctype = (resp.headers.get("content-type") or "application/octet-stream").split(";")[0].strip()
return resp.content, ctype
async with client.stream("GET", url, headers=headers) as resp:
resp.raise_for_status()
length = resp.headers.get("content-length")
if length:
try:
declared_size = int(length)
except ValueError:
declared_size = None
if declared_size is not None and declared_size > limit:
raise ResourceSizeLimitError(limit, declared_size)
chunks: list[bytes] = []
total = 0
async for chunk in resp.aiter_bytes():
total += len(chunk)
if total > limit:
raise ResourceSizeLimitError(limit, total)
chunks.append(chunk)
ctype = (resp.headers.get("content-type") or "application/octet-stream").split(";")[0].strip()
return b"".join(chunks), ctype
def read_local_bytes(path: Path, max_bytes: int | None = None) -> bytes:
"""受单文件上限保护地读取本地媒体文件。"""
limit = max_bytes or get_settings().image_download_max_bytes
size = path.stat().st_size
if size > limit:
raise ResourceSizeLimitError(limit, size)
return path.read_bytes()
class LocalStorage:
"""开发期:落 data/media/,由 FastAPI /media 静态托管,返回 app_base_url 可访问 URL。"""
async def save_from_url(self, url: str, key_prefix: str = "", referer: str | None = None) -> str:
data, _ = await download_bytes(url, referer)
key = self._write(data, key_prefix, url)
return self.public_url(key)
data, ctype = await download_bytes(url, referer)
key = f"{key_prefix.rstrip('/') + '/' if key_prefix else ''}{uuid.uuid4().hex}{image_extension(data, ctype)}"
return await self.save_bytes(data, key, ctype)
async def save_bytes(self, data: bytes, key: str, content_type: str = "") -> str:
key = self._write(data, "", key)
return self.public_url(key)
return self.public_url(self._write(data, key, content_type))
def _write(self, data: bytes, key_prefix: str, hint: str) -> str:
def _write(self, data: bytes, key: str, content_type: str) -> str:
_LOCAL_ROOT.mkdir(parents=True, exist_ok=True)
ext = _ext_from_url(hint) if hint and not hint.startswith("data:") else ".jpg"
key = f"{key_prefix + '/' if key_prefix else ''}{uuid.uuid4().hex}{ext}"
(_LOCAL_ROOT / key).parent.mkdir(parents=True, exist_ok=True)
(_LOCAL_ROOT / key).write_bytes(data)
return key
safe_key = key.strip().lstrip("/")
if not safe_key:
safe_key = uuid.uuid4().hex
if not Path(safe_key).suffix:
safe_key += image_extension(data, content_type)
path = (_LOCAL_ROOT / safe_key).resolve()
if _LOCAL_ROOT.resolve() not in path.parents:
raise ValueError("非法存储 key")
path.parent.mkdir(parents=True, exist_ok=True)
path.write_bytes(data)
return safe_key
def public_url(self, key: str) -> str:
settings = get_settings()
return f"{settings.app_base_url.rstrip('/')}/media/{key}"
async def delete(self, key: str) -> None:
path = (_LOCAL_ROOT / key.strip().lstrip("/")).resolve()
if _LOCAL_ROOT.resolve() not in path.parents:
raise ValueError("非法存储 key")
if path.is_file():
path.unlink()
async def delete_prefix(self, prefix: str) -> int:
root = (_LOCAL_ROOT / prefix.strip().strip("/")).resolve()
if _LOCAL_ROOT.resolve() not in root.parents and root != _LOCAL_ROOT.resolve():
raise ValueError("非法存储前缀")
if not root.exists():
return 0
if root.is_file():
root.unlink()
return 1
count = sum(1 for path in root.rglob("*") if path.is_file())
shutil.rmtree(root)
return count
class QiniuStorage:
"""生产:上传七牛,返回绑定域名公网 URL(Ozon 可拉取)。"""
@@ -97,6 +172,33 @@ class QiniuStorage:
settings = get_settings()
return f"{settings.qiniu_domain.rstrip('/')}/{key}"
async def delete(self, key: str) -> None:
import qiniu
auth, settings = self._client()
_, info = qiniu.BucketManager(auth).delete(settings.qiniu_bucket, key)
if info.status_code not in (200, 612): # 612:对象已不存在,视为删除完成
raise RuntimeError(f"七牛删除失败:{info.error or info.text_body or info.status_code}")
async def delete_prefix(self, prefix: str) -> int:
import qiniu
auth, settings = self._client()
manager = qiniu.BucketManager(auth)
marker: str | None = None
deleted = 0
while True:
result, eof, info = manager.list(settings.qiniu_bucket, prefix=prefix, marker=marker, limit=1000)
if info.status_code != 200:
raise RuntimeError(f"七牛列举失败:{info.error or info.text_body or info.status_code}")
items = (result or {}).get("items") or []
for item in items:
await self.delete(item["key"])
deleted += 1
marker = (result or {}).get("marker")
if eof or not marker:
return deleted
def get_storage():
settings = get_settings()
@@ -105,6 +207,19 @@ def get_storage():
return LocalStorage()
def storage_key_from_url(url: str) -> str | None:
"""将本项目生成的公开 URL 还原为对象 key;第三方源站 URL 返回 None。"""
if not url:
return None
settings = get_settings()
if "/media/" in url:
return url.split("/media/", 1)[1].split("?", 1)[0]
domain = settings.qiniu_domain.rstrip("/")
if domain and url.startswith(f"{domain}/"):
return url[len(domain) + 1:].split("?", 1)[0]
return None
def local_path(stored_url_or_key: str) -> Path | None:
"""stored_url(http…/media/xxx)或 media key → 本地文件路径。
@@ -113,5 +228,9 @@ def local_path(stored_url_or_key: str) -> Path | None:
s = stored_url_or_key or ""
if "/media/" in s:
s = s.split("/media/", 1)[1]
p = _LOCAL_ROOT / s
root = _LOCAL_ROOT.resolve()
p = (root / s.split("?", 1)[0]).resolve()
# URL / key 可能来自历史数据,不能让 ../../ 逃出本地媒体目录。
if root not in p.parents:
return None
return p if p.is_file() else None
+26 -6
View File
@@ -7,6 +7,7 @@ from __future__ import annotations
import uuid
from schemas.suite import TextMaterial
from services.image_assets import get_image_asset_service
from services.tasks import TaskImage
@@ -16,6 +17,8 @@ def texts_to_raw(texts: list[TextMaterial]) -> dict:
for t in texts:
if t.kind == "title" and t.content:
raw["title"] = t.content
elif t.kind == "title_ru" and t.content:
raw["title_ru"] = t.content
elif t.kind == "price" and t.content:
raw["price"] = t.content
elif t.kind == "brand" and t.content:
@@ -67,6 +70,9 @@ async def append_generated_asset(
pid = uuid.UUID(product_id)
async with get_session_factory()() as db:
product = await db.get(Product, pid)
if product is None:
raise RuntimeError("商品已删除,生成结果不再回写")
sort_order: int | None = None
if after_asset_id:
try:
@@ -104,16 +110,30 @@ async def append_generated_asset(
sort_order=sort_order,
type="img",
source_url="",
stored_url=image.url,
status="uploaded",
)
db.add(asset)
await db.flush()
product = await db.get(Product, pid)
if product is not None:
counts = dict(product.asset_counts or {})
counts["generated"] = int(counts.get("generated") or 0) + 1
product.asset_counts = counts
image_assets = get_image_asset_service()
original_url = image.source_url or image.url
await image_assets.adopt_generated_versions(
asset,
original_url=original_url,
current_url=image.url,
watermark_config=image.watermark_config,
)
counts = dict(product.asset_counts or {})
counts["generated"] = int(counts.get("generated") or 0) + 1
product.asset_counts = counts
await db.commit()
# task 的 suites/ 临时对象已迁入正式素材路径,提交成功后立即回收。
try:
await image_assets.delete_temporary_urls([image.source_url, image.url])
except Exception: # noqa: BLE001
# 当前素材已可用;临时对象残留可由删除商品时的前缀清理兜底。
pass
image.source_url = asset.original_url or ""
image.url = asset.stored_url or ""
return str(asset.id)
+2
View File
@@ -29,6 +29,8 @@ class TaskImage:
name: str
status: str = IMG_PENDING # 循环里先建后跑,成功改 ok、失败显式改 failed
url: str = ""
source_url: str = "" # 水印版对应的无水印原图;未加水印时为空
watermark_config: dict | None = None
error: str | None = None
+1 -1
View File
@@ -21,7 +21,7 @@ log = logging.getLogger("suite.watermark")
# 尺寸比例(与 ozonSeller app.js 常量一致)
BADGE_SCALE = 0.15 # 图片水印直径 / 图宽
BADGE_MARGIN = 0.01 # 图片水印边距 / 图宽(ozonSeller 固定 10px,按比例更稳)
TEXT_SCALE = 0.06 # 文字字号 / 图宽
TEXT_SCALE = 0.048 # 文字字号 / 图宽(原 0.06 缩小 20%)
TEXT_MIN_SIZE = 12
STROKE_ALPHA = 0.55
STROKE_RATIO = 1 / 8 # 描边宽 / 字号