333 lines
16 KiB
Python
333 lines
16 KiB
Python
"""独立后台 Worker：按子项粒度执行持久化图片任务。"""
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import logging
|
||
import socket
|
||
import uuid
|
||
|
||
import httpx
|
||
|
||
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 OutpaintCanvas, _image_size, _refs_for_job, advance_generation, prepare_outpaint_canvas
|
||
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.prompts.common import requirements_block
|
||
from services.prompts.single_edit import build_single_edit_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
|
||
from services.video_renderer import VideoRenderCancelled, render_slideshow
|
||
|
||
log = logging.getLogger("background.worker")
|
||
|
||
# 异步供应商任务已经持久化了 task id 后，轮询网络失败不能把任务判为失败。
|
||
# 下一次轮询复用同一 task id，不会再次提交/计费。
|
||
PROVIDER_POLL_RETRY_DELAY_SECONDS = 30
|
||
|
||
|
||
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_slideshow_video_item(claim: ClaimedItem) -> None:
|
||
"""一个视频任务只有一个子项；输入 URL 在提交时冻结，避免后续选择状态影响成片。"""
|
||
payload = claim.request_payload
|
||
sources = list(payload.get("sources") or [])
|
||
source_urls = [str(source.get("url") or "") for source in sources]
|
||
if not 2 <= len(source_urls) <= 5 or any(not url for url in source_urls):
|
||
await mark_item_failed(claim.item_id, "视频任务缺少 2–5 张可用图片")
|
||
return
|
||
try:
|
||
async def should_cancel() -> bool:
|
||
return not await job_can_write(claim.job_id)
|
||
|
||
rendered = await render_slideshow(source_urls, should_cancel=should_cancel)
|
||
if not await job_can_write(claim.job_id):
|
||
await mark_item_cancelled(claim.item_id)
|
||
return
|
||
key = f"products/{claim.product_id}/videos/{claim.job_id}/result.mp4"
|
||
result_url = await get_storage().save_bytes(rendered.data, key, "video/mp4")
|
||
async with get_session_factory()() as db:
|
||
from models import BackgroundJobItem
|
||
item = await db.get(BackgroundJobItem, uuid.UUID(claim.item_id))
|
||
if item is not None:
|
||
item.payload = {
|
||
**dict(item.payload),
|
||
"duration_seconds": rendered.duration_seconds,
|
||
"width": rendered.width,
|
||
"height": rendered.height,
|
||
"file_size_bytes": len(rendered.data),
|
||
}
|
||
await db.commit()
|
||
await mark_item_done(claim.item_id, result_url=result_url, source_url=None, watermark_config=None, generated_asset_id=None)
|
||
except VideoRenderCancelled:
|
||
await mark_item_cancelled(claim.item_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))
|
||
|
||
|
||
def _outpaint_prompt(requirements: str | None) -> str:
|
||
"""仅扩图不复用套图提示词，避免默认风格、营销构图或文字污染边缘生成。"""
|
||
requested = (requirements or "").strip()
|
||
if not requested:
|
||
requested = "Continue the existing surroundings naturally and seamlessly."
|
||
return "\n\n".join([
|
||
requirements_block(requested),
|
||
"OUTPAINT ONLY: Image 1 contains an original photo placed on a 3:4 canvas. Generate pixels only "
|
||
"in the transparent masked margins outside the original photo. Do not crop, move, resize, redraw, "
|
||
"retouch, recolor, relight, or otherwise alter any pixel of the original photo.",
|
||
"Do not add marketing layouts, product captions, overlays, props, or visible text unless the STRICT "
|
||
"REQUIREMENTS explicitly ask for them. Make the new margins blend naturally with the original edges.",
|
||
"FINAL CHECK: output exactly 3:4 and change only the newly expanded margins.",
|
||
])
|
||
|
||
|
||
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 == "slideshow_video":
|
||
await _process_slideshow_video_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:
|
||
single_edit = payload.get("single_edit")
|
||
refs = (
|
||
[image["url"] for image in payload.get("ref_images") or []]
|
||
if single_edit else _refs_for_job(list(payload.get("ref_images") or []), job_item)
|
||
)
|
||
operation = payload.get("operation", "edit")
|
||
target_size = _image_size(provider, claim.ratio, is_wan=provider == "tongyi" and model.lower().startswith("wan"), model=model)
|
||
edit_mask: str | None = None
|
||
outpaint: OutpaintCanvas | None = None
|
||
if operation == "outpaint":
|
||
if len(refs) != 1:
|
||
raise RuntimeError("仅扩图必须且只能使用一张参考底图")
|
||
prompt = _outpaint_prompt(payload.get("requirements"))
|
||
# 首次提交时才构造图片和蒙版；之后轮询不再向供应商重复上传。
|
||
if not claim.provider_task_id:
|
||
outpaint = await prepare_outpaint_canvas(refs[0], target_size)
|
||
refs = [outpaint.image_data_url]
|
||
edit_mask = outpaint.mask_data_url
|
||
else:
|
||
context = build_context(payload.get("context") or {}, fallback_name="")
|
||
if single_edit:
|
||
prompt = build_single_edit_prompt(
|
||
requirements=payload.get("requirements") or "", config=single_edit,
|
||
context=context, style_set=claim.style_set,
|
||
style_prompt=payload.get("style_prompt"), ratio=claim.ratio,
|
||
is_size_chart=job_item.get("kind") == "size_chart",
|
||
reference_count=len(refs) - 1,
|
||
)
|
||
else:
|
||
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"),
|
||
requirements_exclusive=bool(payload.get("requirements_exclusive")),
|
||
)
|
||
is_wan = provider == "tongyi" and model.lower().startswith("wan")
|
||
result = await advance_generation(
|
||
provider, prompt, refs,
|
||
size=target_size, model=model,
|
||
provider_task_id=claim.provider_task_id, provider_task_kind=claim.provider_task_kind,
|
||
edit_mask=edit_mask,
|
||
)
|
||
if result.data is None:
|
||
if not result.provider_task_id or not result.provider_task_kind:
|
||
raise RuntimeError("异步生成未返回可恢复的 provider task id")
|
||
# 轮询供应商期间用户可能刚好取消任务。此时不能再写回
|
||
# waiting_provider，否则父任务已取消、子项却会永久显示“进行中”。
|
||
if not await job_can_write(claim.job_id):
|
||
await mark_item_cancelled(claim.item_id)
|
||
return
|
||
await mark_waiting_provider(claim.item_id, result.provider_task_id, result.provider_task_kind)
|
||
return
|
||
|
||
# 生成完成后仍检查取消/商品删除，避免在用户删除商品后把素材回写回来。
|
||
if not await job_can_write(claim.job_id):
|
||
await mark_item_cancelled(claim.item_id)
|
||
return
|
||
|
||
if operation == "outpaint":
|
||
# 异步任务完成时本轮未必创建过画布，按同一源图与尺寸确定性重建后覆盖原图区域。
|
||
outpaint = outpaint or await prepare_outpaint_canvas(refs[0], target_size)
|
||
result = type(result)(data=outpaint.restore_original(result.data))
|
||
|
||
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 (httpx.TimeoutException, httpx.TransportError) as exc:
|
||
if claim.provider_task_id and claim.provider_task_kind:
|
||
log.warning(
|
||
"供应商轮询网络暂时失败，%ds 后重试同一任务 %s：%s",
|
||
PROVIDER_POLL_RETRY_DELAY_SECONDS, claim.provider_task_id, exc.__class__.__name__,
|
||
)
|
||
await mark_waiting_provider(
|
||
claim.item_id,
|
||
claim.provider_task_id,
|
||
claim.provider_task_kind,
|
||
delay_seconds=PROVIDER_POLL_RETRY_DELAY_SECONDS,
|
||
)
|
||
return
|
||
log.exception("后台图片任务 %s/%s 网络请求失败", claim.job_id, claim.item_id)
|
||
await mark_item_failed(claim.item_id, str(exc) or exc.__class__.__name__)
|
||
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, lease_seconds=get_settings().background_job_lease_seconds)
|
||
if claim is None:
|
||
if once:
|
||
return
|
||
await asyncio.sleep(idle_seconds)
|
||
continue
|
||
await process_item(claim)
|
||
if once:
|
||
return
|