213 lines
9.0 KiB
Python
213 lines
9.0 KiB
Python
"""独立后台 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
|