modalDeploy/src/BowongModalFunctions/router/ffmpeg.py

from typing import Annotated, cast

import modal
from fastapi import APIRouter, Depends, Header
from fastapi.responses import Response
from loguru import logger
from starlette import status

from ..models.settings.cluster import WorkerConfig
from ..middleware.authorization import verify_token
from ..utils.ModalUtils import ModalUtils
from ..models.responses.models import ModalTaskResponse, RecordingTaskResponse, BaseFFMPEGTaskStatusResponse
from ..models.requests.models import FFMPEGSliceRequest, FFMPEGConcatRequest, FFMPEGExtractAudioRequest, \
    FFMPEGCornerMirrorRequest, FFMPEGZoomLoopRequest, FFMPEGOverlayGifRequest, FFMPEGSubtitleOverlayRequest, \
    FFMPEGMixBgmWithNoiseReduceRequest, FFMPEGVideoLoopFillAudioRequest, FFMPEGExtractFrameRequest, \
    FFMPEGStreamRecordRequest, FFMPEGConvertStreamRequest
from ..models.ffmpeg_tasks.models import SentryTransactionHeader, SentryTransactionInfo

config = WorkerConfig()

router = APIRouter(prefix="/ffmpeg", tags=["FFMPEG任务"], dependencies=[Depends(verify_token)])

sentry_header_schema = {
    "x-trace-id": {
        "description": "Sentry Transaction ID",
        "schema": {
            "type": "string",
        }
    },
    "x-baggage": {
        "description": "Sentry Transaction baggage",
        "schema": {
            "type": "string",
        }
    }
}


@router.get("/task/{task_id}", summary="查询ffmpeg任务状态/结果", responses={
    status.HTTP_200_OK: {
        "description": "",
        "headers": sentry_header_schema
    },
}, )
async def get_task_status(task_id: str, response: Response) -> BaseFFMPEGTaskStatusResponse:
    logger.info(f"Get task [{task_id}]")
    task_info = await ModalUtils.get_modal_task_status(task_id)
    if task_info.transaction:
        response.headers["x-trace-id"] = task_info.transaction.x_trace_id
        response.headers["x-baggage"] = task_info.transaction.x_baggage
    return BaseFFMPEGTaskStatusResponse(taskId=task_id, status=task_info.status,
                                        code=cast(int, task_info.error_code.value),
                                        error=task_info.error_reason, results=task_info.results)


@router.post("/slice", summary="发起切割任务", )
async def slice_media(request: FFMPEGSliceRequest,
                      headers: Annotated[SentryTransactionHeader, Header()]) -> ModalTaskResponse:
    fn = modal.Function.from_name(config.modal_app_name, "ffmpeg_slice_media",
                                  environment_name=config.modal_environment)
    sentry_trace = None
    if headers.x_trace_id and headers.x_baggage:
        sentry_trace = SentryTransactionInfo(x_trace_id=headers.x_trace_id, x_baggage=headers.x_baggage)
    fn_call = fn.spawn(media=request.media, markers=request.markers, options=request.options,
                       sentry_trace=sentry_trace, webhook=request.webhook)
    return ModalTaskResponse(success=True, taskId=fn_call.object_id)


@router.post("/convert", summary="发起直播流转换任务", )
async def convert_media(request: FFMPEGConvertStreamRequest,
                        headers: Annotated[SentryTransactionHeader, Header()]) -> ModalTaskResponse:
    fn = modal.Function.from_name(config.modal_app_name, "ffmpeg_convert_stream",
                                  environment_name=config.modal_environment)
    sentry_trace = None
    if headers.x_trace_id and headers.x_baggage:
        sentry_trace = SentryTransactionInfo(x_trace_id=headers.x_trace_id, x_baggage=headers.x_baggage)
    fn_call = fn.spawn(media=request.media, options=request.options, sentry_trace=sentry_trace, webhook=request.webhook)
    return ModalTaskResponse(success=True, taskId=fn_call.object_id)


@router.post("/concat",
             summary="发起合并任务",
             description="依据AI分析的结果发起合并任务",
             )
async def concat_media(body: FFMPEGConcatRequest,
                       headers: Annotated[SentryTransactionHeader, Header()]) -> ModalTaskResponse:
    medias = body.medias
    fn = modal.Function.from_name(config.modal_app_name, "ffmpeg_concat_medias",
                                  environment_name=config.modal_environment)
    sentry_trace = None
    if headers.x_trace_id and headers.x_baggage:
        sentry_trace = SentryTransactionInfo(x_trace_id=headers.x_trace_id, x_baggage=headers.x_baggage)
    fn_call = fn.spawn(medias, sentry_trace, webhook=body.webhook)
    return ModalTaskResponse(success=True, taskId=fn_call.object_id)


@router.post("/concat-video-only",
             summary="发起合并任务(仅视频流)",
             description="依据AI分析的结果发起合并任务，仅处理视频流，忽略音频",
             )
async def concat_media_video_only(body: FFMPEGConcatRequest,
                                  headers: Annotated[SentryTransactionHeader, Header()]) -> ModalTaskResponse:
    medias = body.medias
    fn = modal.Function.from_name(config.modal_app_name, "ffmpeg_concat_medias_video_only",
                                  environment_name=config.modal_environment)
    sentry_trace = None
    if headers.x_trace_id and headers.x_baggage:
        sentry_trace = SentryTransactionInfo(x_trace_id=headers.x_trace_id, x_baggage=headers.x_baggage)
    fn_call = fn.spawn(medias, sentry_trace, webhook=body.webhook)
    return ModalTaskResponse(success=True, taskId=fn_call.object_id)


@router.post("/extract-audio", summary="发起音频提取任务", )
async def extract_audio(body: FFMPEGExtractAudioRequest,
                        headers: Annotated[SentryTransactionHeader, Header()]) -> ModalTaskResponse:
    media = body.media
    fn = modal.Function.from_name(config.modal_app_name, "ffmpeg_extract_audio",
                                  environment_name=config.modal_environment)
    sentry_trace = None
    if headers.x_trace_id and headers.x_baggage:
        sentry_trace = SentryTransactionInfo(x_trace_id=headers.x_trace_id, x_baggage=headers.x_baggage)
    fn_call = fn.spawn(media, sentry_trace, webhook=body.webhook)
    return ModalTaskResponse(success=True, taskId=fn_call.object_id)


@router.post("/corner-mirror", summary="发起镜像小窗去重任务", )
async def corner_mirror(body: FFMPEGCornerMirrorRequest,
                        headers: Annotated[SentryTransactionHeader, Header()]) -> ModalTaskResponse:
    media = body.media
    fn = modal.Function.from_name(config.modal_app_name, "ffmpeg_corner_mirror",
                                  environment_name=config.modal_environment)
    sentry_trace = None
    if headers.x_trace_id and headers.x_baggage:
        sentry_trace = SentryTransactionInfo(x_trace_id=headers.x_trace_id, x_baggage=headers.x_baggage)
    fn_call = fn.spawn(media=media, sentry_trace=sentry_trace, webhook=body.webhook)
    return ModalTaskResponse(success=True, taskId=fn_call.object_id)


@router.post("/overlay-gif", summary="发起叠加gif特效任务", )
async def overlay_gif(body: FFMPEGOverlayGifRequest,
                      headers: Annotated[SentryTransactionHeader, Header()]) -> ModalTaskResponse:
    fn = modal.Function.from_name(config.modal_app_name, "ffmpeg_overlay_gif",
                                  environment_name=config.modal_environment)
    sentry_trace = None
    if headers.x_trace_id and headers.x_baggage:
        sentry_trace = SentryTransactionInfo(x_trace_id=headers.x_trace_id, x_baggage=headers.x_baggage)
    fn_call = fn.spawn(media=body.media, gif=body.gif, sentry_trace=sentry_trace, webhook=body.webhook)
    return ModalTaskResponse(success=True, taskId=fn_call.object_id)


@router.post("/zoom-loop", summary="发起放大缩小循环去重任务", )
async def zoom_loop(body: FFMPEGZoomLoopRequest,
                    headers: Annotated[SentryTransactionHeader, Header()]) -> ModalTaskResponse:
    fn = modal.Function.from_name(config.modal_app_name, "ffmpeg_zoom_loop", environment_name=config.modal_environment)
    sentry_trace = None
    if headers.x_trace_id and headers.x_baggage:
        sentry_trace = SentryTransactionInfo(x_trace_id=headers.x_trace_id, x_baggage=headers.x_baggage)
    fn_call = fn.spawn(media=body.media, duration=body.duration, zoom=body.zoom, sentry_trace=sentry_trace,
                       webhook=body.webhook)
    return ModalTaskResponse(success=True, taskId=fn_call.object_id)


@router.post("/subtitle-apply", summary="发起渲染字幕任务", )
async def subtitle_apply(body: FFMPEGSubtitleOverlayRequest,
                         headers: Annotated[SentryTransactionHeader, Header()]) -> ModalTaskResponse:
    fn = modal.Function.from_name(config.modal_app_name, "ffmpeg_subtitle_apply",
                                  environment_name=config.modal_environment)
    sentry_trace = None
    if headers.x_trace_id and headers.x_baggage:
        sentry_trace = SentryTransactionInfo(x_trace_id=headers.x_trace_id, x_baggage=headers.x_baggage)
    fn_call = fn.spawn(media=body.media, subtitle=body.subtitle, embed_subtitle=body.embedded_subtitle,
                       fonts=body.fonts, sentry_trace=sentry_trace, webhook=body.webhook)
    return ModalTaskResponse(success=True, taskId=fn_call.object_id)


@router.post("/bgm-nosie-reduce", summary="发起混合BGM并降噪任务", )
async def bgm_nosie_reduce(body: FFMPEGMixBgmWithNoiseReduceRequest,
                           headers: Annotated[SentryTransactionHeader, Header()]) -> ModalTaskResponse:
    fn = modal.Function.from_name(config.modal_app_name, "ffmpeg_bgm_nosie_reduce",
                                  environment_name=config.modal_environment)
    sentry_trace = None
    if headers.x_trace_id and headers.x_baggage:
        sentry_trace = SentryTransactionInfo(x_trace_id=headers.x_trace_id, x_baggage=headers.x_baggage)
    fn_call = fn.spawn(media=body.media, bgm=body.bgm, video_volume=body.video_volume, music_volume=body.music_volume,
                       noise_sample=body.noise_sample, sentry_trace=sentry_trace, webhook=body.webhook)
    return ModalTaskResponse(success=True, taskId=fn_call.object_id)


@router.post('/video-loop-fill', summary="发起视频匹配音频长度任务")
async def video_loop_fill_audio(body: FFMPEGVideoLoopFillAudioRequest,
                                headers: Annotated[SentryTransactionHeader, Header()]) -> ModalTaskResponse:
    fn = modal.Function.from_name(config.modal_app_name, "ffmpeg_loop_fill",
                                  environment_name=config.modal_environment)
    sentry_trace = None
    if headers.x_trace_id and headers.x_baggage:
        sentry_trace = SentryTransactionInfo(x_trace_id=headers.x_trace_id, x_baggage=headers.x_baggage)
    fn_call = fn.spawn(media=body.video, audio=body.audio, sentry_trace=sentry_trace, webhook=body.webhook)
    return ModalTaskResponse(success=True, taskId=fn_call.object_id)


@router.post("/extract-frame", summary="发起视频截取帧画面任务", )
async def video_extract_frame(body: FFMPEGExtractFrameRequest,
                              headers: Annotated[SentryTransactionHeader, Header()]) -> ModalTaskResponse:
    fn = modal.Function.from_name(config.modal_app_name, "ffmpeg_extract_frame",
                                  environment_name=config.modal_environment)
    sentry_trace = None
    if headers.x_trace_id and headers.x_baggage:
        sentry_trace = SentryTransactionInfo(x_trace_id=headers.x_trace_id, x_baggage=headers.x_baggage)
    fn_call = fn.spawn(media=body.video, seek_time=body.prased_seek_time, frame_index=body.frame_index,
                       sentry_trace=sentry_trace, webhook=body.webhook)
    return ModalTaskResponse(success=True, taskId=fn_call.object_id)


@router.post("/record/hls", summary="发起直播录制为HLS任务",
             description="录制任务最长存在时间为12小时, 录制任务完成后所获得的hls视频资源只存在3天，如需持久存储请通过使用manifest_urn发起'直播流转换任务'转换为mp4存储",
             response_description="返回发起的任务id和成功录制完第一个片段后通过CDN缓存获得的m3u8播放列表")
async def stream_record_vod(body: FFMPEGStreamRecordRequest,
                            headers: Annotated[SentryTransactionHeader, Header()]) -> RecordingTaskResponse:
    fn = modal.Function.from_name(config.modal_app_name, "ffmpeg_stream_record_as_hls",
                                  environment_name=config.modal_environment)
    sentry_trace = None
    if headers.x_trace_id and headers.x_baggage:
        sentry_trace = SentryTransactionInfo(x_trace_id=headers.x_trace_id, x_baggage=headers.x_baggage)
    fn_call = fn.spawn(stream_url=body.stream_source,
                       first_segment_duration=body.first_segment_duration,
                       segment_duration=body.segment_duration,
                       recording_timeout=body.recording_timeout,
                       monitor_timeout=body.monitor_timeout,
                       sentry_trace=sentry_trace,
                       webhook=body.webhook)
    manifest_link = f"{config.S3_cdn_endpoint}/{config.modal_environment}/records/hls/{fn_call.object_id}/playlist.m3u8"
    manifest_urn = f"s3://{config.S3_region}/{config.S3_bucket_name}/{config.modal_environment}/records/hls/{fn_call.object_id}/playlist.m3u8"
    return RecordingTaskResponse(success=True, taskId=fn_call.object_id,
                                 manifest=manifest_link, manifest_urn=manifest_urn)