xiaoxia-saas/apps/api/app/schemas/gpu_lipsync.py

"""GPU MuseTalk 反向轮询 API Schema 定义.

面向部署在用户 RTX2060 本地的 GPU Worker 脚本，不面向前端用户。
Worker 用长期 GPU_WORKER_TOKEN 鉴权（不是用户 JWT）。
"""

from __future__ import annotations

from datetime import datetime
from typing import Optional

from pydantic import BaseModel, Field

# ── Worker 注册/心跳 ──────────────────────────────────────────────


class GpuWorkerRegisterRequest(BaseModel):
    """Worker 启动/心跳时上报自身信息."""

    worker_id: str = Field(..., min_length=1, max_length=100, description="Worker 唯一 ID（机器名+UUID 等）")
    hostname: str = Field("", max_length=200, description="主机名，用于运维排查")
    gpu_name: str = Field("", max_length=200, description="GPU 型号，如 'NVIDIA GeForce RTX 2060'")
    free_vram_mb: int = Field(0, ge=0, description="当前空闲显存（MB）")
    capabilities: str = Field("musetalk", max_length=500, description="能力列表，逗号分隔，如 'musetalk'")
    task_id: Optional[str] = Field(
        None,
        max_length=64,
        description=(
            "当前正在处理的任务 ID。Worker 推理期间定期心跳时携带，"
            "服务端同步刷新该任务 last_heartbeat_at，防止长推理被误判超时；空闲时不传"
        ),
    )


class GpuWorkerRegisterResponse(BaseModel):
    ok: bool = True
    server_time: datetime
    message: str = "ok"


# ── 轮询任务 ────────────────────────────────────────────────────


class GpuLipsyncTaskPayload(BaseModel):
    """下发给 Worker 的任务载荷（含预签名下载 URL）."""

    task_id: str
    video_url: str = Field(..., description="人物视频预签名下载 URL（GET）")
    audio_url: str = Field(..., description="驱动音频预签名下载 URL（GET）")
    lipsync_job_id: str = ""
    user_id: str = ""
    project_id: str = ""
    created_at: datetime
    upload_url: str = Field(..., description="结果视频预签名上传 URL（PUT, video/mp4）")
    upload_method: str = Field("PUT", description="上传方式，目前只支持 PUT")
    expires_at: datetime


class GpuLipsyncPollResponse(BaseModel):
    """Worker poll 的返回：200 带任务，204 无任务."""

    task: Optional[GpuLipsyncTaskPayload] = None


# ── Worker 上报结果 ──────────────────────────────────────────────


class GpuLipsyncResultRequest(BaseModel):
    """Worker 通过 multipart 上传结果时携带的字段（非文件字段）."""

    task_id: str = Field(..., min_length=1, max_length=64)
    worker_id: str = Field(..., min_length=1, max_length=100)
    success: bool = Field(True, description="true=成功（此时必须上传 result 视频文件）；false=失败")
    duration_seconds: float = Field(0.0, ge=0, description="合成后视频时长（秒），成功时应填入")
    error_msg: str = Field("", max_length=2000, description="失败原因，success=false 时必填")


class GpuLipsyncResultResponse(BaseModel):
    ok: bool = True
    task_id: str
    status: str  # done / failed
    message: str = "ok"


# ── 业务侧查询任务状态 ────────────────────────────────────────────


class GpuLipsyncStatusResponse(BaseModel):
    task_id: str
    status: str
    result_url: str = ""
    result_duration: float = 0.0
    error_msg: str = ""
    worker_id: str = ""
    attempt: int = 0
    created_at: datetime
    started_at: Optional[datetime] = None
    finished_at: Optional[datetime] = None


# ── 创建任务（内部服务调用） ──────────────────────────────────────


class GpuLipsyncCreateRequest(BaseModel):
    """服务层内部创建 GPU 任务用（不通过 HTTP 暴露给 Worker/前端）."""

    video_url: str  # 已可访问的 OSS key 或公网 URL（API 侧会转预签名）
    audio_url: str
    lipsync_job_id: str = ""
    user_id: str = ""
    project_id: str = ""