422 lines
16 KiB
Python
422 lines
16 KiB
Python
"""Prepare-model and try-on activities plus shared helpers."""
|
||
|
||
from __future__ import annotations
|
||
|
||
from dataclasses import asdict, is_dataclass
|
||
from datetime import datetime, timezone
|
||
from enum import Enum
|
||
from typing import Any
|
||
from uuid import uuid4
|
||
|
||
from sqlalchemy import select
|
||
from sqlalchemy.orm import selectinload
|
||
from temporalio import activity
|
||
|
||
from app.application.services.image_generation_service import build_image_generation_service
|
||
from app.domain.enums import AssetType, LibraryResourceStatus, LibraryResourceType, OrderStatus, StepStatus
|
||
from app.infra.db.models.asset import AssetORM
|
||
from app.infra.db.models.library_resource import LibraryResourceORM
|
||
from app.infra.db.models.library_resource_file import LibraryResourceFileORM
|
||
from app.infra.db.models.order import OrderORM
|
||
from app.infra.db.models.workflow_run import WorkflowRunORM
|
||
from app.infra.db.models.workflow_step import WorkflowStepORM
|
||
from app.infra.db.session import get_session_factory
|
||
from app.infra.storage.s3 import S3ObjectStorageService
|
||
from app.workers.workflows.types import MockActivityResult, StepActivityInput
|
||
|
||
|
||
def utc_now() -> datetime:
|
||
"""Return the current UTC timestamp."""
|
||
|
||
return datetime.now(timezone.utc)
|
||
|
||
|
||
def jsonable(value: Any) -> Any:
|
||
"""Convert enums, dataclasses, and nested values to JSON-safe structures."""
|
||
|
||
if value is None:
|
||
return None
|
||
if isinstance(value, Enum):
|
||
return value.value
|
||
if isinstance(value, datetime):
|
||
return value.isoformat()
|
||
if is_dataclass(value):
|
||
return jsonable(asdict(value))
|
||
if isinstance(value, dict):
|
||
return {key: jsonable(item) for key, item in value.items() if item is not None}
|
||
if isinstance(value, (list, tuple, set)):
|
||
return [jsonable(item) for item in value]
|
||
return value
|
||
|
||
|
||
def mock_uri(order_id: int, step_name: str, filename: str = "result.png") -> str:
|
||
"""Build a deterministic-looking mock URI for an order step."""
|
||
|
||
return f"mock://orders/{order_id}/{step_name}/{uuid4().hex[:8]}-{filename}"
|
||
|
||
|
||
async def load_order_and_run(session, order_id: int, workflow_run_id: int) -> tuple[OrderORM, WorkflowRunORM]:
|
||
"""Load the order and workflow run required by an activity."""
|
||
|
||
order = await session.get(OrderORM, order_id)
|
||
workflow_run = await session.get(WorkflowRunORM, workflow_run_id)
|
||
if order is None or workflow_run is None:
|
||
raise ValueError("Order or workflow run not found for activity execution")
|
||
return order, workflow_run
|
||
|
||
|
||
def create_step_record(payload: StepActivityInput) -> WorkflowStepORM:
|
||
"""Create a running workflow step row for an activity execution."""
|
||
|
||
return WorkflowStepORM(
|
||
workflow_run_id=payload.workflow_run_id,
|
||
step_name=payload.step_name,
|
||
step_status=StepStatus.RUNNING,
|
||
input_json=jsonable(payload),
|
||
started_at=utc_now(),
|
||
)
|
||
|
||
|
||
async def load_active_library_resource(
|
||
session,
|
||
resource_id: int,
|
||
*,
|
||
resource_type: LibraryResourceType,
|
||
) -> tuple[LibraryResourceORM, LibraryResourceFileORM]:
|
||
"""Load an active library resource and its original file from the library."""
|
||
|
||
result = await session.execute(
|
||
select(LibraryResourceORM)
|
||
.options(selectinload(LibraryResourceORM.files))
|
||
.where(LibraryResourceORM.id == resource_id)
|
||
)
|
||
resource = result.scalar_one_or_none()
|
||
if resource is None:
|
||
raise ValueError(f"Library resource {resource_id} not found")
|
||
if resource.resource_type != resource_type:
|
||
raise ValueError(f"Resource {resource_id} is not a {resource_type.value} resource")
|
||
if resource.status != LibraryResourceStatus.ACTIVE:
|
||
raise ValueError(f"Library resource {resource_id} is not active")
|
||
if resource.original_file_id is None:
|
||
raise ValueError(f"Library resource {resource_id} is missing an original file")
|
||
|
||
original_file = next((item for item in resource.files if item.id == resource.original_file_id), None)
|
||
if original_file is None:
|
||
raise ValueError(f"Library resource {resource_id} original file record not found")
|
||
return resource, original_file
|
||
|
||
|
||
def get_image_generation_service():
|
||
"""Return the configured image-generation service."""
|
||
|
||
return build_image_generation_service()
|
||
|
||
|
||
def get_order_artifact_storage_service() -> S3ObjectStorageService:
|
||
"""Return the object-storage service for workflow-generated images."""
|
||
|
||
return S3ObjectStorageService()
|
||
|
||
|
||
def build_resource_input_snapshot(
|
||
resource: LibraryResourceORM,
|
||
original_file: LibraryResourceFileORM,
|
||
) -> dict[str, Any]:
|
||
"""Build a frontend-friendly snapshot of one library input resource."""
|
||
|
||
return {
|
||
"resource_id": resource.id,
|
||
"resource_name": resource.name,
|
||
"original_file_id": original_file.id,
|
||
"original_url": original_file.public_url,
|
||
"mime_type": original_file.mime_type,
|
||
"width": original_file.width,
|
||
"height": original_file.height,
|
||
}
|
||
|
||
|
||
async def execute_asset_step(
|
||
payload: StepActivityInput,
|
||
asset_type: AssetType,
|
||
*,
|
||
score: float = 0.95,
|
||
filename: str = "result.png",
|
||
message: str = "mock success",
|
||
extra_metadata: dict[str, Any] | None = None,
|
||
finalize: bool = False,
|
||
) -> MockActivityResult:
|
||
"""Persist a mock asset-producing step and return its result."""
|
||
|
||
async with get_session_factory()() as session:
|
||
order, workflow_run = await load_order_and_run(session, payload.order_id, payload.workflow_run_id)
|
||
step = create_step_record(payload)
|
||
session.add(step)
|
||
|
||
order.status = OrderStatus.RUNNING
|
||
workflow_run.status = OrderStatus.RUNNING
|
||
workflow_run.current_step = payload.step_name
|
||
await session.flush()
|
||
|
||
try:
|
||
metadata = {
|
||
**payload.metadata,
|
||
"source_asset_id": payload.source_asset_id,
|
||
"selected_asset_id": payload.selected_asset_id,
|
||
**(extra_metadata or {}),
|
||
}
|
||
metadata = {key: value for key, value in metadata.items() if value is not None}
|
||
|
||
asset = AssetORM(
|
||
order_id=payload.order_id,
|
||
asset_type=asset_type,
|
||
step_name=payload.step_name,
|
||
uri=mock_uri(payload.order_id, payload.step_name.value, filename),
|
||
metadata_json=jsonable(metadata),
|
||
)
|
||
session.add(asset)
|
||
await session.flush()
|
||
|
||
result = MockActivityResult(
|
||
step_name=payload.step_name,
|
||
success=True,
|
||
asset_id=asset.id,
|
||
uri=asset.uri,
|
||
score=score,
|
||
passed=True,
|
||
message=message,
|
||
metadata=jsonable(metadata) or {},
|
||
)
|
||
|
||
if finalize:
|
||
order.final_asset_id = asset.id
|
||
order.status = OrderStatus.SUCCEEDED
|
||
workflow_run.status = OrderStatus.SUCCEEDED
|
||
|
||
step.step_status = StepStatus.SUCCEEDED
|
||
step.output_json = jsonable(result)
|
||
step.ended_at = utc_now()
|
||
await session.commit()
|
||
return result
|
||
except Exception as exc:
|
||
step.step_status = StepStatus.FAILED
|
||
step.error_message = str(exc)
|
||
step.ended_at = utc_now()
|
||
order.status = OrderStatus.FAILED
|
||
workflow_run.status = OrderStatus.FAILED
|
||
await session.commit()
|
||
raise
|
||
|
||
|
||
@activity.defn
|
||
async def prepare_model_activity(payload: StepActivityInput) -> MockActivityResult:
|
||
"""Resolve a model resource into an order-scoped prepared-model asset."""
|
||
|
||
if payload.model_id is None:
|
||
raise ValueError("prepare_model_activity requires model_id")
|
||
|
||
async with get_session_factory()() as session:
|
||
order, workflow_run = await load_order_and_run(session, payload.order_id, payload.workflow_run_id)
|
||
step = create_step_record(payload)
|
||
session.add(step)
|
||
|
||
order.status = OrderStatus.RUNNING
|
||
workflow_run.status = OrderStatus.RUNNING
|
||
workflow_run.current_step = payload.step_name
|
||
await session.flush()
|
||
|
||
try:
|
||
resource, original_file = await load_active_library_resource(
|
||
session,
|
||
payload.model_id,
|
||
resource_type=LibraryResourceType.MODEL,
|
||
)
|
||
garment_snapshot = None
|
||
if payload.garment_asset_id is not None:
|
||
garment_resource, garment_original = await load_active_library_resource(
|
||
session,
|
||
payload.garment_asset_id,
|
||
resource_type=LibraryResourceType.GARMENT,
|
||
)
|
||
garment_snapshot = build_resource_input_snapshot(garment_resource, garment_original)
|
||
|
||
scene_snapshot = None
|
||
if payload.scene_ref_asset_id is not None:
|
||
scene_resource, scene_original = await load_active_library_resource(
|
||
session,
|
||
payload.scene_ref_asset_id,
|
||
resource_type=LibraryResourceType.SCENE,
|
||
)
|
||
scene_snapshot = build_resource_input_snapshot(scene_resource, scene_original)
|
||
|
||
metadata = {
|
||
**payload.metadata,
|
||
"model_id": payload.model_id,
|
||
"pose_id": payload.pose_id,
|
||
"garment_asset_id": payload.garment_asset_id,
|
||
"scene_ref_asset_id": payload.scene_ref_asset_id,
|
||
"library_resource_id": resource.id,
|
||
"library_original_file_id": original_file.id,
|
||
"library_original_url": original_file.public_url,
|
||
"library_original_mime_type": original_file.mime_type,
|
||
"library_original_width": original_file.width,
|
||
"library_original_height": original_file.height,
|
||
"model_input": build_resource_input_snapshot(resource, original_file),
|
||
"garment_input": garment_snapshot,
|
||
"scene_input": scene_snapshot,
|
||
"pose_input": {"pose_id": payload.pose_id} if payload.pose_id is not None else None,
|
||
"normalized": False,
|
||
}
|
||
metadata = {key: value for key, value in metadata.items() if value is not None}
|
||
|
||
asset = AssetORM(
|
||
order_id=payload.order_id,
|
||
asset_type=AssetType.PREPARED_MODEL,
|
||
step_name=payload.step_name,
|
||
uri=original_file.public_url,
|
||
metadata_json=jsonable(metadata),
|
||
)
|
||
session.add(asset)
|
||
await session.flush()
|
||
|
||
result = MockActivityResult(
|
||
step_name=payload.step_name,
|
||
success=True,
|
||
asset_id=asset.id,
|
||
uri=asset.uri,
|
||
score=1.0,
|
||
passed=True,
|
||
message="prepared model ready",
|
||
metadata=jsonable(metadata) or {},
|
||
)
|
||
|
||
step.step_status = StepStatus.SUCCEEDED
|
||
step.output_json = jsonable(result)
|
||
step.ended_at = utc_now()
|
||
await session.commit()
|
||
return result
|
||
except Exception as exc:
|
||
step.step_status = StepStatus.FAILED
|
||
step.error_message = str(exc)
|
||
step.ended_at = utc_now()
|
||
order.status = OrderStatus.FAILED
|
||
workflow_run.status = OrderStatus.FAILED
|
||
await session.commit()
|
||
raise
|
||
|
||
|
||
@activity.defn
|
||
async def run_tryon_activity(payload: StepActivityInput) -> MockActivityResult:
|
||
"""执行试衣渲染步骤，或在未接真实能力时走 mock 分支。
|
||
|
||
流程：
|
||
1. 读取当前配置的图片生成 service。
|
||
2. 如果当前仍是 mock service，就退回通用的 mock 资产产出逻辑。
|
||
3. 如果是真实模式，先读取 prepare_model_activity 产出的 prepared_model 资产。
|
||
4. 再读取本次选中的服装资源，并解析出它的原图 URL。
|
||
5. 把“模特准备图 + 服装原图”一起发给 provider 做试衣生成。
|
||
6. 把生成结果上传到 S3，并落成订单内的一条 TRYON 资产。
|
||
7. 在 metadata 里记录 provider、model、prompt 和输入来源，方便追踪排查。
|
||
"""
|
||
|
||
service = get_image_generation_service()
|
||
# 保留旧的 mock 分支，这样在没有接入真实 provider 时 workflow 也还能跑通。
|
||
if service.__class__.__name__ == "MockImageGenerationService":
|
||
return await execute_asset_step(
|
||
payload,
|
||
AssetType.TRYON,
|
||
extra_metadata={"prepared_asset_id": payload.source_asset_id},
|
||
)
|
||
|
||
if payload.source_asset_id is None:
|
||
raise ValueError("run_tryon_activity requires source_asset_id")
|
||
if payload.garment_asset_id is None:
|
||
raise ValueError("run_tryon_activity requires garment_asset_id")
|
||
|
||
async with get_session_factory()() as session:
|
||
order, workflow_run = await load_order_and_run(session, payload.order_id, payload.workflow_run_id)
|
||
step = create_step_record(payload)
|
||
session.add(step)
|
||
|
||
order.status = OrderStatus.RUNNING
|
||
workflow_run.status = OrderStatus.RUNNING
|
||
workflow_run.current_step = payload.step_name
|
||
await session.flush()
|
||
|
||
try:
|
||
# 试衣步骤的输入起点固定是上一阶段产出的 prepared_model 资产。
|
||
prepared_asset = await session.get(AssetORM, payload.source_asset_id)
|
||
if prepared_asset is None or prepared_asset.order_id != payload.order_id:
|
||
raise ValueError(f"Prepared asset {payload.source_asset_id} not found for order {payload.order_id}")
|
||
if prepared_asset.asset_type != AssetType.PREPARED_MODEL:
|
||
raise ValueError(f"Asset {payload.source_asset_id} is not a prepared_model asset")
|
||
|
||
# 服装素材来自共享资源库，workflow 实际消费的是资源库里的原图。
|
||
garment_resource, garment_original = await load_active_library_resource(
|
||
session,
|
||
payload.garment_asset_id,
|
||
resource_type=LibraryResourceType.GARMENT,
|
||
)
|
||
|
||
# provider 接收两张真实输入图：模特准备图 + 服装参考图。
|
||
generated = await service.generate_tryon_image(
|
||
person_image_url=prepared_asset.uri,
|
||
garment_image_url=garment_original.public_url,
|
||
)
|
||
# workflow 生成出的结果图先上传到 S3，再登记成订单资产。
|
||
storage_key, public_url = await get_order_artifact_storage_service().upload_generated_image(
|
||
order_id=payload.order_id,
|
||
step_name=payload.step_name,
|
||
image_bytes=generated.image_bytes,
|
||
mime_type=generated.mime_type,
|
||
)
|
||
|
||
# 记录足够的追踪信息，方便回溯这张试衣图由哪些输入和哪个 provider 产出。
|
||
metadata = {
|
||
**payload.metadata,
|
||
"prepared_asset_id": prepared_asset.id,
|
||
"garment_resource_id": garment_resource.id,
|
||
"garment_original_file_id": garment_original.id,
|
||
"garment_original_url": garment_original.public_url,
|
||
"provider": generated.provider,
|
||
"model": generated.model,
|
||
"storage_key": storage_key,
|
||
"mime_type": generated.mime_type,
|
||
"prompt": generated.prompt,
|
||
}
|
||
metadata = {key: value for key, value in metadata.items() if value is not None}
|
||
|
||
asset = AssetORM(
|
||
order_id=payload.order_id,
|
||
asset_type=AssetType.TRYON,
|
||
step_name=payload.step_name,
|
||
uri=public_url,
|
||
metadata_json=jsonable(metadata),
|
||
)
|
||
session.add(asset)
|
||
await session.flush()
|
||
|
||
result = MockActivityResult(
|
||
step_name=payload.step_name,
|
||
success=True,
|
||
asset_id=asset.id,
|
||
uri=asset.uri,
|
||
score=1.0,
|
||
passed=True,
|
||
message="try-on generated",
|
||
metadata=jsonable(metadata) or {},
|
||
)
|
||
|
||
step.step_status = StepStatus.SUCCEEDED
|
||
step.output_json = jsonable(result)
|
||
step.ended_at = utc_now()
|
||
await session.commit()
|
||
return result
|
||
except Exception as exc:
|
||
step.step_status = StepStatus.FAILED
|
||
step.error_message = str(exc)
|
||
step.ended_at = utc_now()
|
||
order.status = OrderStatus.FAILED
|
||
workflow_run.status = OrderStatus.FAILED
|
||
await session.commit()
|
||
raise
|