"""独立后台 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