feat: add resource library and real image workflow
This commit is contained in:
afei A
@@ -1,4 +1,4 @@
|
||||
"""Prepare-model and try-on mock activities plus shared helpers."""
|
||||
"""Prepare-model and try-on activities plus shared helpers."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
@@ -8,14 +8,20 @@ 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.domain.enums import AssetType, OrderStatus, StepStatus
|
||||
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
|
||||
|
||||
|
||||
@@ -71,6 +77,64 @@ def create_step_record(payload: StepActivityInput) -> WorkflowStepORM:
|
||||
)
|
||||
|
||||
|
||||
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,
|
||||
@@ -145,26 +209,213 @@ async def execute_asset_step(
|
||||
|
||||
@activity.defn
|
||||
async def prepare_model_activity(payload: StepActivityInput) -> MockActivityResult:
|
||||
"""Mock model preparation for the pipeline."""
|
||||
"""Resolve a model resource into an order-scoped prepared-model asset."""
|
||||
|
||||
return await execute_asset_step(
|
||||
payload,
|
||||
AssetType.PREPARED_MODEL,
|
||||
extra_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,
|
||||
},
|
||||
)
|
||||
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 try-on rendering."""
|
||||
"""执行试衣渲染步骤，或在未接真实能力时走 mock 分支。
|
||||
|
||||
return await execute_asset_step(
|
||||
payload,
|
||||
AssetType.TRYON,
|
||||
extra_metadata={"prepared_asset_id": payload.source_asset_id},
|
||||
)
|
||||
流程：
|
||||
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
|
||||
|
||||
Reference in New Issue
Block a user