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
+89 -31
View File
@@ -10,14 +10,17 @@ from __future__ import annotations
import io
import mimetypes
import re
import tempfile
import zipfile
from collections.abc import Iterator
from fastapi import APIRouter, HTTPException
from fastapi.responses import StreamingResponse
from fastapi.responses import Response, StreamingResponse
from pydantic import BaseModel, Field
from api.proxy import guess_referer
from services.storage import download_bytes, local_path
from config import get_settings
from services.storage import ResourceSizeLimitError, download_bytes, local_path, read_local_bytes
router = APIRouter(prefix="/api", tags=["export"])
@@ -39,7 +42,7 @@ class ExportImageItem(BaseModel):
class ExportImagesRequest(BaseModel):
title: str | None = Field(default=None, description="商品标题,用作 ZIP 顶层文件夹名")
images: list[ExportImageItem]
images: list[ExportImageItem] = Field(min_length=1)
def _clean(name: str) -> str:
@@ -75,47 +78,102 @@ async def _download(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
return await download_bytes(url, referer=guess_referer(url))
def _stream_file(handle: io.BufferedRandom, chunk_size: int = 64 * 1024) -> Iterator[bytes]:
"""以固定块流式返回临时 ZIP,响应结束(或中断)后关闭临时文件。"""
try:
while chunk := handle.read(chunk_size):
yield chunk
finally:
handle.close()
@router.post("/export/images")
async def export_images(req: ExportImagesRequest):
if not req.images:
raise HTTPException(status_code=400, detail="没有可导出的图片")
settings = get_settings()
if len(req.images) > settings.export_max_images:
raise HTTPException(status_code=413, detail=f"一次最多导出 {settings.export_max_images} 张图片,请分批导出")
root = _clean(req.title) or "采集图片"
buf = io.BytesIO()
# 小 ZIP 留在内存,大 ZIP 自动滚动到临时磁盘;响应阶段按块传输。
archive = tempfile.SpooledTemporaryFile(
max_size=settings.export_zip_memory_threshold_bytes,
mode="w+b",
)
used: set[str] = set()
total_source_bytes = 0
with zipfile.ZipFile(buf, "w", zipfile.ZIP_DEFLATED) as zf:
for img in req.images:
try:
data, ctype = await _download(img.url)
except Exception: # noqa: BLE001
continue # 单张失败不中断整包
if not _is_image(img.url, ctype):
continue
try:
with zipfile.ZipFile(archive, "w", zipfile.ZIP_DEFLATED) as zf:
for img in req.images:
try:
data, ctype = await _download(img.url)
except ResourceSizeLimitError as exc:
raise HTTPException(status_code=413, detail=f"无法导出「{img.key or '图片'}」:{exc}") from exc
except Exception: # noqa: BLE001
continue # 单张下载失败不中断整包
if not _is_image(img.url, ctype):
continue
total_source_bytes += len(data)
if total_source_bytes > settings.export_zip_max_bytes:
raise HTTPException(
status_code=413,
detail=f"导出图片总大小超过 {settings.export_zip_max_bytes // 1024 // 1024} MB,请减少勾选图片后重试",
)
ext = _ext(img.url, ctype)
base = _clean(img.key) or "image"
if img.variantName:
base += f"-{_clean(img.variantName)}"
filename = f"{base}{ext}"
if filename in used: # 同名加序号防覆盖
stem = filename[: -len(ext)]
n = 2
while f"{stem}-{n}{ext}" in used:
n += 1
filename = f"{stem}-{n}{ext}"
used.add(filename)
ext = _ext(img.url, ctype)
base = _clean(img.key) or "image"
if img.variantName:
base += f"-{_clean(img.variantName)}"
filename = f"{base}{ext}"
if filename in used: # 同名加序号防覆盖
stem = filename[: -len(ext)]
n = 2
while f"{stem}-{n}{ext}" in used:
n += 1
filename = f"{stem}-{n}{ext}"
used.add(filename)
group = _clean(img.groupName) or "图片"
zf.writestr(f"{root}/{group}/{filename}", data)
group = _clean(img.groupName) or "图片"
zf.writestr(f"{root}/{group}/{filename}", data)
if archive.tell() > settings.export_zip_max_bytes:
raise HTTPException(
status_code=413,
detail=f"导出 ZIP 超过 {settings.export_zip_max_bytes // 1024 // 1024} MB,请减少勾选图片后重试",
)
except Exception:
archive.close()
raise
buf.seek(0)
archive.seek(0)
return StreamingResponse(
buf,
_stream_file(archive),
media_type="application/zip",
headers={"Content-Disposition": 'attachment; filename="collect.zip"'},
)
@router.post("/export/image")
async def export_image(img: ExportImageItem):
"""代理下载单张图片,兼容源站防盗链与本地媒体路径。"""
try:
data, ctype = await _download(img.url)
except ResourceSizeLimitError as exc:
raise HTTPException(status_code=413, detail=f"图片文件过大,无法下载:{exc}") from exc
except Exception as exc: # noqa: BLE001
raise HTTPException(status_code=422, detail="图片下载失败") from exc
if not _is_image(img.url, ctype):
raise HTTPException(status_code=400, detail="该地址不是图片")
ext = _ext(img.url, ctype)
base = _clean(img.key) or "image"
if img.variantName:
base += f"-{_clean(img.variantName)}"
return Response(
content=data,
media_type=ctype,
headers={"Content-Disposition": f'attachment; filename="{base}{ext}"'},
)
+165 -19
View File
@@ -2,9 +2,11 @@
from __future__ import annotations
import re
import logging
from uuid import UUID
from fastapi import APIRouter, BackgroundTasks, Depends, HTTPException, UploadFile, File, Form
from pydantic import BaseModel, Field
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession
@@ -12,8 +14,38 @@ from db import get_db, get_session_factory
from models import Product, ProductAsset
from models.enums import AssetStatus, Stage
from schemas.collection import MaterialsRequest, MaterialsResponse, TextMaterial
from schemas.suite import SuiteOut, WatermarkOptions
from services.image_assets import WatermarkUpdate, get_image_asset_service
from services.jobs import create_image_job, get_job, list_product_jobs, suite_view
from config import get_settings
router = APIRouter(prefix="/api", tags=["collection"])
log = logging.getLogger("materials")
class ApplyWatermarkRequest(BaseModel):
"""对已落库图片批量应用水印;始终基于 original 重新合成。"""
asset_ids: list[str] = Field(min_length=1, max_length=50)
watermark: WatermarkOptions
class ApplyWatermarkResponse(BaseModel):
"""批量水印已进入后台队列;前端以 job_id 查询逐图进度。"""
job_id: str
total: int
class RemoveWatermarkResponse(BaseModel):
"""同步移除水印的逐项处理结果。"""
updated_ids: list[str]
failed: list[dict[str, str]]
class RemoveWatermarkRequest(BaseModel):
asset_ids: list[str] = Field(min_length=1, max_length=50)
def _parse_number(text: str | None) -> float | None:
@@ -115,6 +147,11 @@ async def create_materials(
background: BackgroundTasks,
db: AsyncSession = Depends(get_db)
) -> MaterialsResponse:
if len(req.images) > get_settings().collection_max_images:
raise HTTPException(
status_code=413,
detail=f"一次最多采集 {get_settings().collection_max_images} 个素材,请分批上报",
)
product = await _get_or_create_product(db, req)
_apply_texts(product, req.texts)
@@ -176,9 +213,7 @@ async def create_materials(
async def process_product_assets(product_id: str) -> None:
"""后台:下载 pending 素材 → 转存 storage。失败逐张标记,不中断。"""
from services.storage import get_storage
storage = get_storage()
image_assets = get_image_asset_service()
async with get_session_factory()() as db:
assets = (await db.scalars(
select(ProductAsset).where(
@@ -190,8 +225,7 @@ async def process_product_assets(product_id: str) -> None:
a.status = AssetStatus.downloading
await db.commit()
try:
stored = await storage.save_from_url(a.source_url, key_prefix="assets")
a.stored_url = stored
await image_assets.save_original_from_url(a, a.source_url)
a.status = AssetStatus.uploaded
except Exception as exc: # noqa: BLE001
a.status = AssetStatus.failed
@@ -213,7 +247,17 @@ async def upload_material_bytes(
product = await db.get(Product, UUID(product_id))
if product is None:
raise HTTPException(status_code=404, detail="商品不存在")
data = await file.read()
limit = get_settings().image_upload_max_bytes
if file.size is not None and file.size > limit:
raise HTTPException(status_code=413, detail=f"图片文件过大,单张最多 {limit // 1024 // 1024} MB")
chunks: list[bytes] = []
total = 0
while chunk := await file.read(1024 * 1024):
total += len(chunk)
if total > limit:
raise HTTPException(status_code=413, detail=f"图片文件过大,单张最多 {limit // 1024 // 1024} MB")
chunks.append(chunk)
data = b"".join(chunks)
asset = ProductAsset(
product_id=product.id,
group_key=group_key,
@@ -225,11 +269,10 @@ async def upload_material_bytes(
)
db.add(asset)
await db.flush()
# 直接转存字节
from services.storage import get_storage
storage = get_storage()
# 直接转存为素材唯一原图(current 与 original 初始指向同一对象)
image_assets = get_image_asset_service()
try:
asset.stored_url = await storage.save_bytes(data, f"assets/{asset.id}", file.content_type or "")
await image_assets.save_original(asset, data, file.content_type or "")
asset.status = AssetStatus.uploaded
except Exception as exc: # noqa: BLE001
asset.status = AssetStatus.failed
@@ -269,20 +312,16 @@ async def is_collected(
@router.delete("/assets/{asset_id}")
async def delete_asset(asset_id: str, db: AsyncSession = Depends(get_db)):
"""删除一条素材记录(DB 行 + 本地文件;七牛等远程存储仅删引用)。"""
from services.storage import local_path
"""删除一条素材记录及其 original/current 七牛对象。"""
asset = await db.get(ProductAsset, UUID(asset_id))
if asset is None:
raise HTTPException(status_code=404, detail="素材不存在")
if asset.stored_url:
path = local_path(asset.stored_url)
if path is not None and path.is_file():
try:
path.unlink()
except OSError: # noqa: BLE001
pass # 文件删除失败不阻断记录删除
try:
await get_image_asset_service().delete_asset_media(asset)
except Exception as exc: # noqa: BLE001
raise HTTPException(status_code=502, detail=f"图片存储删除失败,素材未删除:{exc}") from exc
product_id = asset.product_id
group_key = asset.group_key
@@ -298,3 +337,110 @@ async def delete_asset(asset_id: str, db: AsyncSession = Depends(get_db)):
await db.commit()
return {"deleted": True}
@router.post("/assets/watermark", response_model=ApplyWatermarkResponse)
async def apply_assets_watermark(
req: ApplyWatermarkRequest,
db: AsyncSession = Depends(get_db),
) -> ApplyWatermarkResponse:
"""创建批量水印后台任务;Worker 从原图逐张合成并切换当前展示版。"""
if not req.watermark.enabled:
raise HTTPException(status_code=400, detail="请先开启水印")
try:
asset_ids = list(dict.fromkeys(UUID(asset_id) for asset_id in req.asset_ids))
except ValueError as exc:
raise HTTPException(status_code=400, detail="包含非法素材 ID") from exc
assets = (await db.scalars(select(ProductAsset).where(ProductAsset.id.in_(asset_ids)))).all()
by_id = {asset.id: asset for asset in assets}
missing = [str(asset_id) for asset_id in asset_ids if asset_id not in by_id]
if missing:
raise HTTPException(status_code=404, detail=f"素材不存在:{', '.join(missing[:3])}")
product_ids = {asset.product_id for asset in assets}
if len(product_ids) != 1:
raise HTTPException(status_code=400, detail="一次批量水印只能处理同一商品的素材")
product_id = str(next(iter(product_ids)))
task = await create_image_job(
product_id=product_id,
kind="watermark_batch",
# 水印在后台低优先级执行,不阻塞单张编辑和套图生成。
priority=-10,
provider="local",
model=None,
platform="ozon",
lang="",
ratio="",
style_set=0,
request_payload={"watermark": req.watermark.model_dump()},
item_payloads=[
{
"position": index,
"kind": "watermark",
"title": f"水印:{str(asset.id)[:8]}",
"asset_id": str(asset.id),
# 任务排队时冻结当前版本;若同图被其他操作更新,Worker 不会覆盖新版本。
"expected_current_storage_key": asset.current_storage_key,
"expected_current_url": asset.stored_url,
}
for index, asset in enumerate(by_id[asset_id] for asset_id in asset_ids)
],
)
return ApplyWatermarkResponse(job_id=str(task.id), total=task.total)
@router.get("/products/{product_id}/watermark-jobs", response_model=list[SuiteOut])
async def list_watermark_jobs(product_id: str, active_only: bool = False):
"""试算页刷新后恢复正在执行的批量水印任务。"""
jobs = await list_product_jobs(product_id, active_only=active_only, kinds=("watermark_batch",))
views: list[SuiteOut] = []
for job in jobs:
result = await get_job(str(job.id))
if result is not None:
views.append(SuiteOut(**suite_view(*result)))
return views
@router.post("/assets/watermark/remove", response_model=RemoveWatermarkResponse)
async def remove_assets_watermark(
req: RemoveWatermarkRequest,
db: AsyncSession = Depends(get_db),
) -> RemoveWatermarkResponse:
"""批量移除水印:current 切回 original,并删除当前水印对象。"""
try:
asset_ids = list(dict.fromkeys(UUID(asset_id) for asset_id in req.asset_ids))
except ValueError as exc:
raise HTTPException(status_code=400, detail="包含非法素材 ID") from exc
assets = (await db.scalars(select(ProductAsset).where(ProductAsset.id.in_(asset_ids)))).all()
by_id = {asset.id: asset for asset in assets}
missing = [str(asset_id) for asset_id in asset_ids if asset_id not in by_id]
if missing:
raise HTTPException(status_code=404, detail=f"素材不存在:{', '.join(missing[:3])}")
image_assets = get_image_asset_service()
updated_ids: list[str] = []
failed: list[dict[str, str]] = []
obsolete: list[WatermarkUpdate] = []
for asset_id in asset_ids:
asset = by_id[asset_id]
if asset.type != "img":
failed.append({"asset_id": str(asset_id), "error": "仅支持图片素材"})
continue
if not asset.watermark_enabled:
continue
try:
obsolete.append(await image_assets.remove_watermark(asset))
updated_ids.append(str(asset_id))
except Exception as exc: # noqa: BLE001
failed.append({"asset_id": str(asset_id), "error": str(exc)[:200]})
await db.commit()
for update in obsolete:
try:
await image_assets.delete_key(update.obsolete_key)
except Exception: # noqa: BLE001
log.exception("水印已移除,但旧水印对象删除失败:%s", update.obsolete_key)
return RemoveWatermarkResponse(updated_ids=updated_ids, failed=failed)
+15
View File
@@ -134,6 +134,18 @@ async def delete_product(
if product is None:
raise HTTPException(status_code=404, detail="商品不存在")
if hard:
# 先阻止后台 Worker 回写新生成素材,再删除已有图片与商品记录。
from services.jobs import cancel_product_jobs
await cancel_product_jobs(str(product.id))
assets = (await db.scalars(
select(ProductAsset).where(ProductAsset.product_id == product.id)
)).all()
try:
from services.image_assets import get_image_asset_service
await get_image_asset_service().delete_product_media(str(product.id), assets)
except Exception as exc: # noqa: BLE001
raise HTTPException(status_code=502, detail=f"图片存储删除失败,商品未删除:{exc}") from exc
await db.delete(product)
else:
product.stage = Stage.archived
@@ -160,6 +172,9 @@ async def list_assets(
"type": a.type,
"source_url": a.source_url,
"stored_url": a.stored_url,
"original_url": a.original_url,
"watermark_enabled": a.watermark_enabled,
"watermark_config": a.watermark_config,
"status": a.status.value,
"width": a.width,
"height": a.height,
+3 -1
View File
@@ -8,7 +8,7 @@ from urllib.parse import urlparse
from fastapi import APIRouter, HTTPException, Query, Response
from services.storage import download_bytes
from services.storage import ResourceSizeLimitError, download_bytes
router = APIRouter(prefix="/api", tags=["proxy"])
@@ -40,6 +40,8 @@ async def proxy_image(url: str = Query(..., description="源站图片 URL")):
raise HTTPException(status_code=400, detail="仅支持 http/https URL")
try:
data, ctype = await download_bytes(url, referer=guess_referer(url))
except ResourceSizeLimitError as exc:
raise HTTPException(status_code=413, detail=f"图片文件过大,无法代理:{exc}") from exc
except Exception as exc: # noqa: BLE001
raise HTTPException(status_code=502, detail=f"图片拉取失败: {exc}") from exc
if not ctype.startswith("image/"):
+188 -116
View File
@@ -1,7 +1,7 @@
"""套图规划 / 一键生成 / 任务查询 / 结果导出 / 单张 AI 生图。
按 docs/v2.1/api.md §2-6 实现:引擎平移自 image-suite-studio(services/planner|generator|
tasks|watermark + services/prompts),任务存进程内内存表(重启即失效)。
watermark + services/prompts),任务持久化到数据库,由独立 Worker 执行。
生成图回调回写 product_assets(group_key='generated');product_id 缺省时只落盘不回写。
业务逻辑(texts_to_raw / 模型校验 / generated 回写)在 services/suite_service.py。
注:V2.1 阶段本组接口暂不接鉴权(现有鉴权后续可能重做)。
@@ -9,11 +9,11 @@ tasks|watermark + services/prompts),任务存进程内内存表(重启即
from __future__ import annotations
import io
import tempfile
import uuid
import zipfile
from functools import partial
from fastapi import APIRouter, BackgroundTasks, HTTPException
from collections.abc import Iterator
from fastapi import APIRouter, HTTPException
from fastapi.responses import StreamingResponse
from config import get_settings
@@ -21,7 +21,7 @@ from schemas.suite import (
PLATFORM_SPECS,
SUPPORTED_TYPES,
ImageEditSingleRequest,
ImageEditSingleResponse,
ImageEditTaskResponse,
PlanItemOut,
PlanResponse,
SuiteCreateResponse,
@@ -31,16 +31,34 @@ from schemas.suite import (
api_model,
resolve_provider,
)
from services.generator import _image_size, GENERATORS, run_suite
from services.generator import GENERATORS
from services.jobs import (
ITEM_DONE,
cancel_job,
create_image_job,
get_job,
list_product_jobs,
retry_failed_items,
retry_failed_item,
suite_view,
)
from services.planner import generate_plan
from services.prompts import build_context, build_prompt, type_name
from services.storage import download_bytes, get_storage, local_path
from services.suite_service import append_generated_asset, texts_to_raw, validate_model
from services.tasks import IMG_OK, TaskImage, create_task, get_task
from services.prompts import type_name
from services.storage import ResourceSizeLimitError, download_bytes, local_path, read_local_bytes
from services.suite_service import texts_to_raw, validate_model
router = APIRouter(prefix="/api", tags=["suite"])
def _stream_file(handle: io.BufferedRandom, chunk_size: int = 64 * 1024) -> Iterator[bytes]:
"""以固定块返回临时 ZIP,释放响应完成后的临时文件。"""
try:
while chunk := handle.read(chunk_size):
yield chunk
finally:
handle.close()
# ── 出图方案规划 ──────────────────────────────────────────────────────────
@router.post("/suite/plan", response_model=PlanResponse)
@@ -64,7 +82,7 @@ async def plan_suite(req: SuitePlanRequest) -> PlanResponse:
# ── 一键生成 ─────────────────────────────────────────────────────────────
@router.post("/suite/generate", response_model=SuiteCreateResponse)
async def generate_suite(req: SuiteGenerateRequest, background: BackgroundTasks) -> SuiteCreateResponse:
async def generate_suite(req: SuiteGenerateRequest) -> SuiteCreateResponse:
if not req.images:
raise HTTPException(status_code=400, detail="未勾选任何图片,无法生成")
@@ -111,94 +129,146 @@ async def generate_suite(req: SuiteGenerateRequest, background: BackgroundTasks)
except ValueError as exc:
raise HTTPException(status_code=400, detail="product_id 不是合法 UUID") from exc
task = create_task(
status="pending",
task = await create_image_job(
product_id=product_id,
kind="image_suite",
priority=0,
provider=provider_name,
model=model,
platform=req.platform,
lang=spec["lang"],
ratio=spec["ratio"],
style_set=req.style_set,
style_prompt=req.style_prompt,
requirements=req.requirements,
provider=provider_name,
model=model,
total=len(jobs),
context=texts_to_raw(req.texts),
plan=jobs,
watermark=req.watermark.model_dump() if req.watermark else None,
# 参考图池:main 组优先,其余组按序补充(variant 绑定靠 variant_name 匹配)
ref_images=[
{
"url": i.url,
"group_key": i.group_key,
"variant_name": i.variant_name,
}
for i in sorted(req.images, key=lambda x: 0 if x.group_key == "main" else 1)
],
request_payload={
"context": texts_to_raw(req.texts),
"style_prompt": req.style_prompt,
"requirements": req.requirements,
"watermark": req.watermark.model_dump() if req.watermark else None,
"append": bool(product_id),
# main 组优先;variant_name 精确匹配由 Worker 复用原有选图规则。
"ref_images": [
{"url": image.url, "group_key": image.group_key, "variant_name": image.variant_name}
for image in sorted(req.images, key=lambda image: 0 if image.group_key == "main" else 1)
],
},
item_payloads=[{**job, "position": index} for index, job in enumerate(jobs)],
)
# 每张成功即回写 generated 组;未关联商品时仅落存储不回写。
# partial 绑定 product_id:run_suite 回调只传 image,签名须为 (image)
background.add_task(
run_suite,
task,
partial(append_generated_asset, product_id) if product_id else None,
)
return SuiteCreateResponse(suite_id=task.id)
return SuiteCreateResponse(suite_id=str(task.id))
# ── 任务查询 / 结果导出 ──────────────────────────────────────────────────
@router.get("/suites/{suite_id}", response_model=SuiteOut)
async def get_suite(suite_id: str):
task = get_task(suite_id)
result = await get_job(suite_id)
if result is None:
raise HTTPException(status_code=404, detail="任务不存在")
return SuiteOut(**suite_view(*result))
@router.get("/products/{product_id}/suites", response_model=list[SuiteOut])
async def list_suites(product_id: str, active_only: bool = False):
"""试算页恢复商品的未完成任务;后续任务中心可复用。"""
jobs = await list_product_jobs(product_id, active_only=active_only, kinds=("image_suite", "image_edit"))
views: list[SuiteOut] = []
for job in jobs:
result = await get_job(str(job.id))
if result is not None:
views.append(SuiteOut(**suite_view(*result)))
return views
@router.post("/suites/{suite_id}/cancel", response_model=SuiteOut)
async def cancel_suite(suite_id: str):
task = await cancel_job(suite_id)
if task is None:
raise HTTPException(status_code=404, detail="任务不存在(服务可能已重启),请重新生成")
return SuiteOut(
id=task.id,
status=task.status,
style_set=task.style_set,
platform=task.platform,
lang=task.lang,
ratio=task.ratio,
provider=task.provider,
model=task.model,
total=task.total,
images=[
{"type_id": i.type_id, "name": i.name, "url": i.url, "status": i.status, "error": i.error}
for i in task.images
],
error=task.error,
)
raise HTTPException(status_code=404, detail="任务不存在")
result = await get_job(str(task.id))
if result is None:
raise HTTPException(status_code=404, detail="任务不存在")
return SuiteOut(**suite_view(*result))
@router.post("/suites/{suite_id}/retry", response_model=SuiteOut)
async def retry_suite(suite_id: str):
task = await retry_failed_items(suite_id)
if task is None:
raise HTTPException(status_code=404, detail="任务不存在")
result = await get_job(str(task.id))
if result is None:
raise HTTPException(status_code=404, detail="任务不存在")
return SuiteOut(**suite_view(*result))
@router.post("/suites/{suite_id}/items/{item_id}/retry", response_model=SuiteOut)
async def retry_suite_item(suite_id: str, item_id: str):
task = await retry_failed_item(suite_id, item_id)
if task is None:
raise HTTPException(status_code=404, detail="任务或子项不存在")
result = await get_job(str(task.id))
if result is None:
raise HTTPException(status_code=404, detail="任务不存在")
return SuiteOut(**suite_view(*result))
@router.get("/suites/{suite_id}/zip")
async def download_suite_zip(suite_id: str):
"""把任务内所有成功图打包成 ZIP(中文文件名)。"""
task = get_task(suite_id)
if task is None:
raise HTTPException(status_code=404, detail="任务不存在(服务可能已重启)")
result = await get_job(suite_id)
if result is None:
raise HTTPException(status_code=404, detail="任务不存在")
task, task_items = result
buf = io.BytesIO()
with zipfile.ZipFile(buf, "w", zipfile.ZIP_DEFLATED) as zf:
seen: set[str] = set()
for i, img in enumerate([i for i in task.images if i.status == IMG_OK]):
path = local_path(img.url or "")
if path is not None:
data, suffix = path.read_bytes(), path.suffix or ".jpg"
else:
try: # 七牛等远程存储:直接拉取公开 URL
data, _ = await download_bytes(img.url or "")
settings = get_settings()
images = [image for image in task_items if image.status == ITEM_DONE and image.result_url]
if len(images) > settings.export_max_images:
raise HTTPException(status_code=413, detail=f"该任务有 {len(images)} 张图片,单次最多导出 {settings.export_max_images} 张")
archive = tempfile.SpooledTemporaryFile(
max_size=settings.export_zip_memory_threshold_bytes,
mode="w+b",
)
total_source_bytes = 0
try:
with zipfile.ZipFile(archive, "w", zipfile.ZIP_DEFLATED) as zf:
seen: set[str] = set()
for i, img in enumerate(images):
path = local_path(img.result_url or "")
try:
if path is not None:
data, suffix = read_local_bytes(path), path.suffix or ".jpg"
else: # 七牛等远程存储:直接拉取公开 URL
data, _ = await download_bytes(img.result_url or "")
suffix = ".png" if data[:8] == b"\x89PNG\r\n\x1a\n" else ".jpg"
except ResourceSizeLimitError as exc:
raise HTTPException(status_code=413, detail=f"无法导出「{img.name}」:{exc}") from exc
except Exception: # noqa: BLE001
continue
suffix = ".png" if data[:8] == b"\x89PNG\r\n\x1a\n" else ".jpg"
filename = img.name or img.type_id
if filename in seen: # 同类型多张时加序号防覆盖
filename = f"{filename}-{i + 1}"
seen.add(filename)
zf.writestr(f"{filename}{suffix}", data)
buf.seek(0)
total_source_bytes += len(data)
if total_source_bytes > settings.export_zip_max_bytes:
raise HTTPException(
status_code=413,
detail=f"导出图片总大小超过 {settings.export_zip_max_bytes // 1024 // 1024} MB,请减少图片后重试",
)
payload = img.payload or {}
filename = payload.get("title") or payload.get("kind") or "AI 生图"
if filename in seen: # 同类型多张时加序号防覆盖
filename = f"{filename}-{i + 1}"
seen.add(filename)
zf.writestr(f"{filename}{suffix}", data)
if archive.tell() > settings.export_zip_max_bytes:
raise HTTPException(
status_code=413,
detail=f"导出 ZIP 超过 {settings.export_zip_max_bytes // 1024 // 1024} MB,请减少图片后重试",
)
except Exception:
archive.close()
raise
archive.seek(0)
return StreamingResponse(
buf,
_stream_file(archive),
media_type="application/zip",
headers={"Content-Disposition": f"attachment; filename=\"suite-{suite_id}.zip\""},
)
@@ -206,9 +276,9 @@ async def download_suite_zip(suite_id: str):
# ── 单张 AI 生图(V2.1 新增,ISS 无对应实现)──────────────────────────────
@router.post("/suite/image-edit", response_model=ImageEditSingleResponse)
async def suite_image_edit(req: ImageEditSingleRequest) -> ImageEditSingleResponse:
"""轻量单张再生成:原图 + 一句话要求 + 模型选择。同步接口(单张可长达数分钟)。"""
@router.post("/suite/image-edit", response_model=ImageEditTaskResponse)
async def suite_image_edit(req: ImageEditSingleRequest) -> ImageEditTaskResponse:
"""提交单张图生图任务;立即返回任务 ID,实际生成在统一队列异步执行。"""
settings = get_settings()
provider_name = resolve_provider(req.model, None, settings.image_provider)
if provider_name not in GENERATORS:
@@ -217,39 +287,41 @@ async def suite_image_edit(req: ImageEditSingleRequest) -> ImageEditSingleRespon
spec = {"lang": "ru", "ratio": "3:4"} # 试算页固定 Ozon 规格(俄文图内文案 · 3:4)
model = api_model(req.model) # 剥 @provider 后缀,供应商侧只认裸模型名
is_wan = provider_name == "tongyi" and model.lower().startswith("wan")
# 复用套图的 prompt 家族:custom 类型装配「用户要求」为最高优先级约束,
# 商品外观仍由参考图锁定(保真语义与一键生成一致)。
ctx = build_context({"title": "AI 精修"}, fallback_name="product")
prompt = build_prompt(
provider_name, model, "custom", ctx,
style_set=5, lang=spec["lang"],
extra={"title": "AI 精修", "detail": req.prompt, "prompt_hint": ""},
requirements=req.prompt,
)
try:
generator = GENERATORS[provider_name]
data = await generator(prompt, [req.image_url], size=_image_size(provider_name, spec["ratio"], is_wan=is_wan, model=model), model=model)
except RuntimeError as exc:
raise HTTPException(status_code=502, detail=str(exc)) from exc
except Exception as exc: # noqa: BLE001
raise HTTPException(status_code=502, detail=f"生成失败: {exc}") from exc
is_png = data[:8] == b"\x89PNG\r\n\x1a\n"
ext = ".png" if is_png else ".jpg"
ctype = "image/png" if is_png else "image/jpeg"
url = await get_storage().save_bytes(data, f"suites/single/{uuid.uuid4().hex}{ext}", ctype)
asset_id: str | None = None
if req.append and req.product_id:
product_id = (req.product_id or "").strip() or None
if product_id:
try:
asset_id = await append_generated_asset(
req.product_id,
TaskImage(type_id="custom", name="AI生图", url=url, status="ok"),
after_asset_id=req.after_asset_id,
)
except Exception: # noqa: BLE001
pass # 回写失败不影响结果返回,图片已在任务网格可见
return ImageEditSingleResponse(url=url, asset_id=asset_id)
uuid.UUID(product_id)
except ValueError as exc:
raise HTTPException(status_code=400, detail="product_id 不是合法 UUID") from exc
# 单图编辑优先级高于套图;prompt_extra 清空展示文案,避免中文 UI 文案进入模型 prompt。
task = await create_image_job(
product_id=product_id,
kind="image_edit",
priority=10,
provider=provider_name,
model=model,
platform="ozon",
lang=spec["lang"],
ratio=spec["ratio"],
style_set=5,
request_payload={
"context": {"title": "AI 精修"},
"style_prompt": None,
"requirements": req.prompt,
"watermark": req.watermark.model_dump() if req.watermark else None,
"append": bool(req.append and product_id),
"after_asset_id": req.after_asset_id,
"ref_images": [{"url": req.image_url, "group_key": "main", "variant_name": None}],
},
item_payloads=[{
"position": 0,
"kind": "custom",
"title": "AI生图",
"detail": "",
"prompt_hint": "",
"variant_name": None,
"prompt_extra": {"title": "", "detail": "", "prompt_hint": ""},
}],
)
return ImageEditTaskResponse(suite_id=str(task.id))
+10
View File
@@ -49,6 +49,16 @@ class Settings(BaseSettings):
# ── V2:对外地址(插件/前端回写、生成图回调)──
app_base_url: str = "http://127.0.0.1:8800"
# ── 图片与导出资源上限(字节)──
# 单张图片在后端需要解码/加水印,限制单文件可避免一次请求耗尽内存。
image_upload_max_bytes: int = 20 * 1024 * 1024
image_download_max_bytes: int = 20 * 1024 * 1024
collection_max_images: int = 100
export_max_images: int = 50
export_zip_max_bytes: int = 200 * 1024 * 1024
# ZIP 超过该值自动落临时磁盘,不继续占用进程内存。
export_zip_memory_threshold_bytes: int = 2 * 1024 * 1024
# ── V2.1 套图生图(provider 路由与 image-suite-studio 一致)──
image_provider: str = "doubao" # 未知模型时的兜底 provider:doubao | tongyi | rightapi
request_timeout: int = 300 # 单张生图请求超时(秒)
@@ -0,0 +1,44 @@
"""add asset image lifecycle fields
Revision ID: a41e8c7f4b20
Revises: 658b0503f71c
Create Date: 2026-08-30 09:00:00
"""
from __future__ import annotations
from alembic import op
import sqlalchemy as sa
revision = "a41e8c7f4b20"
down_revision = "658b0503f71c"
branch_labels = None
depends_on = None
def upgrade() -> None:
op.add_column("product_assets", sa.Column("original_url", sa.Text(), nullable=True))
op.add_column("product_assets", sa.Column("original_storage_key", sa.String(length=1024), nullable=True))
op.add_column("product_assets", sa.Column("current_storage_key", sa.String(length=1024), nullable=True))
op.add_column("product_assets", sa.Column("watermark_enabled", sa.Boolean(), nullable=False, server_default=sa.false()))
op.add_column("product_assets", sa.Column("watermark_config", sa.JSON(), nullable=True))
# 兼容旧数据:采集图优先以已转存文件为原图;生成图若此前已保存 source_url,则它是无水印原图。
op.execute(
"""
UPDATE product_assets
SET original_url = CASE
WHEN group_key = 'generated' AND source_url IS NOT NULL AND source_url <> '' THEN source_url
ELSE COALESCE(stored_url, source_url)
END
"""
)
def downgrade() -> None:
op.drop_column("product_assets", "watermark_config")
op.drop_column("product_assets", "watermark_enabled")
op.drop_column("product_assets", "current_storage_key")
op.drop_column("product_assets", "original_storage_key")
op.drop_column("product_assets", "original_url")
@@ -0,0 +1,92 @@
"""add persistent background jobs
Revision ID: b5d2c9a3e1f7
Revises: a41e8c7f4b20
Create Date: 2026-08-30 12:00:00
"""
from __future__ import annotations
from alembic import op
import sqlalchemy as sa
revision = "b5d2c9a3e1f7"
down_revision = "a41e8c7f4b20"
branch_labels = None
depends_on = None
def upgrade() -> None:
# 开发态 FastAPI 会在 startup 以 Base.metadata.create_all 兜底;若它先于 Alembic
# 启动,表已完整存在,此处只补迁移版本记录,避免 create_table 冲突。
if "background_jobs" in sa.inspect(op.get_bind()).get_table_names():
return
op.create_table(
"background_jobs",
sa.Column("id", sa.Uuid(), nullable=False),
sa.Column("product_id", sa.Uuid(), nullable=True),
sa.Column("kind", sa.String(length=32), nullable=False),
sa.Column("status", sa.String(length=24), nullable=False, server_default="queued"),
sa.Column("priority", sa.Integer(), nullable=False, server_default="0"),
sa.Column("provider", sa.String(length=32), nullable=False, server_default=""),
sa.Column("model", sa.String(length=128), nullable=True),
sa.Column("platform", sa.String(length=16), nullable=False, server_default="ozon"),
sa.Column("lang", sa.String(length=16), nullable=False, server_default="ru"),
sa.Column("ratio", sa.String(length=16), nullable=False, server_default="3:4"),
sa.Column("style_set", sa.Integer(), nullable=False, server_default="1"),
sa.Column("request_payload", sa.JSON(), nullable=False),
sa.Column("total", sa.Integer(), nullable=False, server_default="0"),
sa.Column("completed", sa.Integer(), nullable=False, server_default="0"),
sa.Column("failed", sa.Integer(), nullable=False, server_default="0"),
sa.Column("cancel_requested", sa.Boolean(), nullable=False, server_default=sa.false()),
sa.Column("error", sa.Text(), nullable=True),
sa.Column("created_at", sa.DateTime(timezone=True), server_default=sa.text("CURRENT_TIMESTAMP"), nullable=False),
sa.Column("started_at", sa.DateTime(timezone=True), nullable=True),
sa.Column("completed_at", sa.DateTime(timezone=True), nullable=True),
sa.Column("updated_at", sa.DateTime(timezone=True), server_default=sa.text("CURRENT_TIMESTAMP"), nullable=False),
sa.ForeignKeyConstraint(["product_id"], ["products.id"], ondelete="SET NULL"),
sa.PrimaryKeyConstraint("id"),
)
op.create_index("ix_background_jobs_product_id", "background_jobs", ["product_id"])
op.create_index("ix_background_jobs_kind", "background_jobs", ["kind"])
op.create_index("ix_background_jobs_status", "background_jobs", ["status"])
op.create_index("ix_background_jobs_priority", "background_jobs", ["priority"])
op.create_index("ix_background_jobs_cancel_requested", "background_jobs", ["cancel_requested"])
op.create_index("ix_background_jobs_created_at", "background_jobs", ["created_at"])
op.create_index("ix_background_jobs_updated_at", "background_jobs", ["updated_at"])
op.create_table(
"background_job_items",
sa.Column("id", sa.Uuid(), nullable=False),
sa.Column("job_id", sa.Uuid(), nullable=False),
sa.Column("position", sa.Integer(), nullable=False, server_default="0"),
sa.Column("status", sa.String(length=24), nullable=False, server_default="queued"),
sa.Column("payload", sa.JSON(), nullable=False),
sa.Column("provider_task_id", sa.String(length=512), nullable=True),
sa.Column("provider_task_kind", sa.String(length=64), nullable=True),
sa.Column("attempts", sa.Integer(), nullable=False, server_default="0"),
sa.Column("max_attempts", sa.Integer(), nullable=False, server_default="3"),
sa.Column("result_url", sa.Text(), nullable=True),
sa.Column("source_url", sa.Text(), nullable=True),
sa.Column("watermark_config", sa.JSON(), nullable=True),
sa.Column("generated_asset_id", sa.Uuid(), nullable=True),
sa.Column("error", sa.Text(), nullable=True),
sa.Column("next_run_at", sa.DateTime(timezone=True), nullable=True),
sa.Column("lease_owner", sa.String(length=128), nullable=True),
sa.Column("lease_expires_at", sa.DateTime(timezone=True), nullable=True),
sa.Column("created_at", sa.DateTime(timezone=True), server_default=sa.text("CURRENT_TIMESTAMP"), nullable=False),
sa.Column("started_at", sa.DateTime(timezone=True), nullable=True),
sa.Column("completed_at", sa.DateTime(timezone=True), nullable=True),
sa.Column("updated_at", sa.DateTime(timezone=True), server_default=sa.text("CURRENT_TIMESTAMP"), nullable=False),
sa.ForeignKeyConstraint(["job_id"], ["background_jobs.id"], ondelete="CASCADE"),
sa.PrimaryKeyConstraint("id"),
)
op.create_index("ix_background_job_items_job_id", "background_job_items", ["job_id"])
op.create_index("ix_background_job_items_status", "background_job_items", ["status"])
op.create_index("ix_background_job_items_provider_task_id", "background_job_items", ["provider_task_id"])
op.create_index("ix_background_job_items_next_run_at", "background_job_items", ["next_run_at"])
op.create_index("ix_background_job_items_lease_expires_at", "background_job_items", ["lease_expires_at"])
def downgrade() -> None:
op.drop_table("background_job_items")
op.drop_table("background_jobs")
+3
View File
@@ -4,6 +4,7 @@
由 legacy.api 路由 import 链注册到同一 Base.metadata。
"""
from models.asset import ProductAsset
from models.job import BackgroundJob, BackgroundJobItem
from models.product import Product
from models.user import User
@@ -11,4 +12,6 @@ __all__ = [
"User",
"Product",
"ProductAsset",
"BackgroundJob",
"BackgroundJobItem",
]
+11 -3
View File
@@ -4,11 +4,12 @@ from __future__ import annotations
import uuid
from datetime import datetime
from sqlalchemy import DateTime, Enum, ForeignKey, Integer, String, Text, Uuid, func
from sqlalchemy import Boolean, DateTime, Enum, ForeignKey, Integer, String, Text, Uuid, func
from sqlalchemy.orm import Mapped, mapped_column
from db import Base
from models.enums import AssetStatus
from models.types import JSONType
class ProductAsset(Base):
@@ -22,8 +23,15 @@ class ProductAsset(Base):
variant_name: Mapped[str | None] = mapped_column(String(128), nullable=True) # SKU 规格名
sort_order: Mapped[int] = mapped_column(Integer, default=0)
type: Mapped[str] = mapped_column(String(8), default="img") # img / video
source_url: Mapped[str] = mapped_column(Text, default="")
stored_url: Mapped[str | None] = mapped_column(Text, nullable=True) # 本地路径或七牛公网 URL
source_url: Mapped[str] = mapped_column(Text, default="") # 采集源站地址;生成图旧数据可能存无水印 URL
stored_url: Mapped[str | None] = mapped_column(Text, nullable=True) # 当前展示版 URL(兼容旧接口)
# 图片生命周期:original 始终无水印;当前版用于页面预览、下载、导出。
# 未加水印时 original 与 current 指向同一对象;不保留历史水印版。
original_url: Mapped[str | None] = mapped_column(Text, nullable=True)
original_storage_key: Mapped[str | None] = mapped_column(String(1024), nullable=True)
current_storage_key: Mapped[str | None] = mapped_column(String(1024), nullable=True)
watermark_enabled: Mapped[bool] = mapped_column(Boolean, default=False)
watermark_config: Mapped[dict | None] = mapped_column(JSONType, nullable=True)
status: Mapped[AssetStatus] = mapped_column(
Enum(AssetStatus, native_enum=False, length=16), default=AssetStatus.pending, index=True
)
+78
View File
@@ -0,0 +1,78 @@
"""持久化后台任务:AI 套图、单张编辑共用的父任务与逐图子任务。"""
from __future__ import annotations
import uuid
from datetime import datetime
from sqlalchemy import Boolean, DateTime, ForeignKey, Integer, String, Text, Uuid, func
from sqlalchemy.orm import Mapped, mapped_column
from db import Base
from models.types import JSONType
class BackgroundJob(Base):
"""一个用户可见的后台操作:套图、单张编辑或批量水印。"""
__tablename__ = "background_jobs"
id: Mapped[uuid.UUID] = mapped_column(Uuid(as_uuid=True), primary_key=True, default=uuid.uuid4)
# 商品删除后保留任务诊断记录,但 Worker 不会再把结果回写到不存在的商品。
product_id: Mapped[uuid.UUID | None] = mapped_column(
Uuid(as_uuid=True), ForeignKey("products.id", ondelete="SET NULL"), nullable=True, index=True
)
kind: Mapped[str] = mapped_column(String(32), index=True) # image_suite | image_edit | watermark_batch
status: Mapped[str] = mapped_column(String(24), default="queued", index=True)
priority: Mapped[int] = mapped_column(Integer, default=0, index=True)
provider: Mapped[str] = mapped_column(String(32), default="")
model: Mapped[str | None] = mapped_column(String(128), nullable=True)
platform: Mapped[str] = mapped_column(String(16), default="ozon")
lang: Mapped[str] = mapped_column(String(16), default="ru")
ratio: Mapped[str] = mapped_column(String(16), default="3:4")
style_set: Mapped[int] = mapped_column(Integer, default=1)
# 原始请求的不可变快照:商品上下文、参考图池、风格、要求、水印等。
request_payload: Mapped[dict] = mapped_column(JSONType, default=dict)
total: Mapped[int] = mapped_column(Integer, default=0)
completed: Mapped[int] = mapped_column(Integer, default=0)
failed: Mapped[int] = mapped_column(Integer, default=0)
cancel_requested: Mapped[bool] = mapped_column(Boolean, default=False, index=True)
error: Mapped[str | None] = mapped_column(Text, nullable=True)
created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), server_default=func.now(), index=True)
started_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
completed_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
updated_at: Mapped[datetime] = mapped_column(
DateTime(timezone=True), server_default=func.now(), onupdate=func.now(), index=True
)
class BackgroundJobItem(Base):
"""任务中的一个可调度工作单元:套图的一张图或单张编辑。"""
__tablename__ = "background_job_items"
id: Mapped[uuid.UUID] = mapped_column(Uuid(as_uuid=True), primary_key=True, default=uuid.uuid4)
job_id: Mapped[uuid.UUID] = mapped_column(
Uuid(as_uuid=True), ForeignKey("background_jobs.id", ondelete="CASCADE"), index=True
)
position: Mapped[int] = mapped_column(Integer, default=0)
status: Mapped[str] = mapped_column(String(24), default="queued", index=True)
payload: Mapped[dict] = mapped_column(JSONType, default=dict)
# 供应商异步任务 ID 与类型(pcore 还需 edits/generations);重启后继续轮询,不重新提交。
provider_task_id: Mapped[str | None] = mapped_column(String(512), nullable=True, index=True)
provider_task_kind: Mapped[str | None] = mapped_column(String(64), nullable=True)
attempts: Mapped[int] = mapped_column(Integer, default=0)
max_attempts: Mapped[int] = mapped_column(Integer, default=3)
result_url: Mapped[str | None] = mapped_column(Text, nullable=True)
source_url: Mapped[str | None] = mapped_column(Text, nullable=True)
watermark_config: Mapped[dict | None] = mapped_column(JSONType, nullable=True)
generated_asset_id: Mapped[uuid.UUID | None] = mapped_column(Uuid(as_uuid=True), nullable=True)
error: Mapped[str | None] = mapped_column(Text, nullable=True)
next_run_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True, index=True)
lease_owner: Mapped[str | None] = mapped_column(String(128), nullable=True)
lease_expires_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True, index=True)
created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), server_default=func.now())
started_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
completed_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
updated_at: Mapped[datetime] = mapped_column(
DateTime(timezone=True), server_default=func.now(), onupdate=func.now(), index=True
)
+2
View File
@@ -18,6 +18,8 @@ class ProductListItem(BaseModel):
currency_code: str = "RUB"
source_platform: str | None = None
source_url: str | None = None
# 采集箱显示中文标题:优先 raw.title_zh,兼容早期采集写入的 raw.title。
raw: dict | None = None
asset_counts: dict | None = None
ozon_product_id: int | None = None
created_at: datetime
+16 -5
View File
@@ -1,6 +1,7 @@
"""套图规划 / 一键生成 / 单张 AI 生图 契约(studio/src/services/suite.ts 同源,见 docs/v2.1/api.md)。"""
from __future__ import annotations
from datetime import datetime
from typing import Literal
from pydantic import BaseModel, Field
@@ -61,7 +62,7 @@ PLATFORM_SPECS: dict[str, dict] = {
class TextMaterial(BaseModel):
kind: str = Field(..., description="title | params | selling_point | desc | price | brand | sales | shop")
kind: str = Field(..., description="title | title_ru | params | selling_point | desc | price | brand | sales | shop")
content: str = ""
pairs: list[dict] | None = None # [{key, value}]
@@ -135,6 +136,7 @@ class SuiteCreateResponse(BaseModel):
class SuiteImageOut(BaseModel):
item_id: str
type_id: str
name: str
url: str
@@ -145,6 +147,7 @@ class SuiteImageOut(BaseModel):
class SuiteOut(BaseModel):
id: str
status: str
kind: str = "image_suite"
style_set: int
platform: str
lang: str
@@ -152,6 +155,12 @@ class SuiteOut(BaseModel):
provider: str
model: str | None = None
total: int = 0 # 计划总张数(进度分母;images 是逐张追加,过程中 length < total)
completed: int = 0
failed: int = 0
cancel_requested: bool = False
created_at: datetime | None = None
started_at: datetime | None = None
completed_at: datetime | None = None
images: list[SuiteImageOut]
error: str | None = None
@@ -160,15 +169,17 @@ class ImageEditSingleRequest(BaseModel):
"""POST /api/suite/image-edit:单张 AI 生图(试算页每张图的「AI 生图」入口)。"""
product_id: str | None = Field(default=None, description="append=true 时回写到该商品的 generated 组")
image_url: str = Field(..., description="原图(stored_url 或 source_url,远程由服务端带 Referer 代下)")
prompt: str = Field(..., min_length=1, description="用户要求(必填,进生图指令)")
prompt: str = Field(default="", description="用户要求(可选,作为额外生图指令)")
model: str = Field(default="nano-banana-2")
append: bool = Field(default=True, description="结果是否追加为商品素材(generated 组)")
after_asset_id: str | None = Field(
default=None,
description="插入锚点:新素材排在该素材(generated 组)之后,方便与原图对比;缺省追加到组尾",
)
watermark: WatermarkOptions | None = Field(default=None, description="单张生图完成后的服务端水印;关闭时不传")
class ImageEditSingleResponse(BaseModel):
url: str
asset_id: str | None = None
class ImageEditTaskResponse(BaseModel):
"""单张编辑已进入统一图片任务队列;用 GET /api/suites/{suite_id} 查询结果。"""
suite_id: str
+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 # 描边宽 / 字号
+20
View File
@@ -0,0 +1,20 @@
"""命令入口:`python -m worker`,与 Web 进程独立运行。"""
from __future__ import annotations
import argparse
import asyncio
import logging
from services.job_worker import run_worker
def main() -> None:
parser = argparse.ArgumentParser(description="Ozon Seller Kit background worker")
parser.add_argument("--once", action="store_true", help="最多处理一个已到期子项后退出(用于诊断/测试)")
args = parser.parse_args()
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s: %(message)s")
asyncio.run(run_worker(once=args.once))
if __name__ == "__main__":
main()