feat: 优化图片资源删除管理,生图任务列表管理
This commit is contained in:
@@ -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:
|
||||
|
||||
@@ -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()
|
||||
@@ -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
|
||||
@@ -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
|
||||
],
|
||||
}
|
||||
@@ -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}"
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
|
||||
@@ -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 # 描边宽 / 字号
|
||||
|
||||
Reference in New Issue
Block a user