xiaoxia-saas/apps/worker/video_processing/gpu_direct_pipeline.py

"""全 GPU 直连渲染管线（P1）。

背景：旧链路 worker 先用 CPU libx264 把 filter_complex 输出成 mezzanine（1080p 约 85s），
上传后再由 P4000 NVENC 编码，渲染后还要单独跑一次随机边缘裁剪重编码（约 26s）。
本管线取消 mezzanine：把原始素材签名 URL 作为多输入直接交给 P4000，filter_complex 内
一步完成 trim/scale/pad/concat/边缘裁剪/drawtext 字幕，末端 h264_nvenc 只编码一次；
TTS/BGM 音频也在同一命令里 amix 合成。

约束（P1）：
- 仅覆盖智能剪辑主流场景：单一主视频轨、全硬切、无 PiP/overlay/水印/贴纸/片头片尾/绿幕。
  不满足条件时调用方回退到现有 mezzanine/CPU 链路（功能不回归）。
- 字幕先用 drawtext（P4000 装好中文字体后可再切 subtitles 滤镜烧 ASS）。
"""

from __future__ import annotations

import logging
import uuid
from pathlib import Path
from typing import Any, Optional

logger = logging.getLogger(__name__)

# drawtext 默认字体名（fontconfig 解析）；P4000 装好中文字体后可在 config 指定
DEFAULT_DRAWTEXT_FONT = "Noto Sans CJK SC"


def escape_drawtext_text(text: str) -> str:
    """转义 drawtext text= 中的特殊字符（ffmpeg 过滤器语法）。"""
    if not text:
        return ""
    # 顺序重要：先转义反斜杠本身
    s = text.replace("\\", "\\\\")
    s = s.replace(":", "\\:")
    s = s.replace("'", "\\'")
    s = s.replace("%", "\\%")
    s = s.replace(",", "\\,")
    s = s.replace("[", "\\[").replace("]", "\\]")
    s = s.replace(";", "\\;")
    # 换行保留（drawtext 支持 %{...}；真实换行需转成字面）
    s = s.replace("\n", " ")
    return s


def build_drawtext_filter(
    *,
    text: str,
    start: float,
    end: float,
    font: str = DEFAULT_DRAWTEXT_FONT,
    font_size: int = 0,
    font_color: str = "white",
    x_expr: str = "(w-text_w)/2",
    y_expr: str = "h-th-60",
    box: bool = False,
    box_color: str = "black@0.5",
    borderw: int = 0,
    border_color: str = "black",
    enable: bool = True,
) -> str:
    """构造单个 drawtext 滤镜字符串（不含输入/输出标签）。

    [start, end] 秒的显示窗口通过 enable='between(t,...)' 控制。
    """
    txt = escape_drawtext_text(text)
    parts = [f"font={font}", f"text='{txt}'"]
    if font_size and font_size > 0:
        parts.append(f"fontsize={int(font_size)}")
    parts.append(f"fontcolor={font_color}")
    if box:
        parts.append("box=1")
        parts.append(f"boxcolor={box_color}")
    if borderw and borderw > 0:
        parts.append(f"borderw={int(borderw)}")
        parts.append(f"bordercolor={border_color}")
    parts.append(f"x={x_expr}")
    parts.append(f"y={y_expr}")
    if enable:
        parts.append(f"enable='between(t,{start:.3f},{end:.3f})'")
    return "drawtext=" + ":".join(parts)


def upload_local_audio_and_sign(
    local_audio: Path,
    *,
    tmp_prefix: str = "tmp/gpu-direct-audio/",
    expires: int = 3600,
) -> tuple[str, str]:
    """把本地音频（TTS/BGM）上传 OSS tmp 目录并签公网 GET URL。

    Returns:
        (signed_get_url, oss_key)
    """
    from video_processing.oss_helpers import _storage  # 类型: ignore

    storage = _storage()
    key = f"{tmp_prefix.rstrip('/')}/{uuid.uuid4().hex}{local_audio.suffix or '.mp3'}"
    content_type = "audio/mpeg" if local_audio.suffix.lower() in (".mp3", ".mpeg") else "audio/mp4"
    storage.upload_file(local_audio, key, content_type=content_type)
    url = storage.get_download_url(key, expires)
    return url, key


def sign_asset_url(storage_key: str, *, expires: int = 3600) -> str:
    """给原始素材 storage_key 签公网 GET URL（供 P4000 直接下载）。"""
    from video_processing.oss_helpers import _storage  # 类型: ignore

    storage = _storage()
    return storage.get_download_url(storage_key, expires)


# ── 直连渲染编排器 ────────────────────────────────────────────────────────


class DirectRenderPlan:
    """一次直连渲染的产物：inputs（裸文件名→URL）与完整 ffmpeg_args。"""

    def __init__(self, inputs: dict[str, str], ffmpeg_args: list[str], oss_keys: list[str]):
        self.inputs = inputs
        self.ffmpeg_args = ffmpeg_args
        self.oss_keys = oss_keys  # 本次上传的临时音频 key（供事后清理）


def build_direct_render(
    *,
    resolved_clips: list[Any],
    output_width: int,
    output_height: int,
    output_fps: int,
    tts_audio: Optional[Path] = None,
    bgm_audio: Optional[Path] = None,
    title_text: str = "",
    subtitle_segments: Optional[list[Any]] = None,
    font: str = DEFAULT_DRAWTEXT_FONT,
    vcodec: str = "h264_nvenc",
    preset: str = "p4",
    video_bitrate: str = "",
    cq: int = 23,
    edge_crop_pct: float = 0.0,
    total_duration: float = 0.0,
) -> DirectRenderPlan:
    """构造 P4000 直连渲染所需的 inputs 与 ffmpeg_args。

    视频：每段 trim/setpts/scale/pad/fps → concat（全硬切）→ 可选边缘 crop+scale → drawtext。
    音频：concat 时丢弃原生音轨（只映射 [vfinal]），TTS/BGM 上传签名后 amix 混音。
    """
    if not resolved_clips:
        raise ValueError("build_direct_render: no resolved clips")

    inputs: dict[str, str] = {}
    oss_keys: list[str] = []
    input_args: list[str] = []
    fc: list[str] = []  # filter_complex 各段

    n = len(resolved_clips)

    # 1. 视频输入（原始素材签名 URL）
    for i, clip in enumerate(resolved_clips):
        sk = (getattr(clip, "config", None) or {}).get("_storage_key")
        if not sk:
            raise ValueError(f"clip {getattr(clip, 'clip_id', i)} missing _storage_key")
        fname = f"v{i}.mp4"
        inputs[fname] = sign_asset_url(sk)
        input_args.extend(["-i", fname])

    # 2. 每个视频段预处理
    pre_labels: list[str] = []
    for i, clip in enumerate(resolved_clips):
        filters: list[str] = []
        start = float(getattr(clip, "start_time", 0) or 0)
        eff = float(getattr(clip, "duration", 0) or 0)
        if eff <= 0:
            eff = float(getattr(clip, "actual_duration", 0) or 0)
        if eff > 0:
            if start > 0:
                filters.append(f"trim=start={start:.3f}:duration={eff:.3f}")
            else:
                filters.append(f"trim=duration={eff:.3f}")
            filters.append("setpts=PTS-STARTPTS")

        speed = float(getattr(clip, "playback_speed", 1.0) or 1.0)
        if abs(speed - 1.0) >= 1e-6:
            filters.append(f"setpts=PTS/{speed:.4f}")

        filters.append(f"scale={output_width}:{output_height}:force_original_aspect_ratio=decrease")
        filters.append(f"pad={output_width}:{output_height}:trunc((ow-iw)/2):trunc((oh-ih)/2):black")
        filters.append("setpts=PTS-STARTPTS")
        filters.append(f"fps={output_fps}")

        label = f"vc{i}"
        fc.append(f"[{i}:v]{','.join(filters)}[{label}]")
        pre_labels.append(label)

    # 3. concat（全硬切；原生音频丢弃，v=1:a=0）
    concat_in = "".join(f"[{l}]" for _lbl in pre_labels)
    fc.append(f"{concat_in}concat=n={n}:v=1:a=0[vcat]")
    cur = "vcat"

    # 4. 边缘裁剪降重（合并进同一条，不再单独重编码）
    if edge_crop_pct and edge_crop_pct > 0:
        p = float(edge_crop_pct)
        keep = 1.0 - 2.0 * p
        cw_expr = f"trunc(iw*{keep:.4f}/2)*2"
        ch_expr = f"trunc(ih*{keep:.4f}/2)*2"
        fc.append(
            f"[{cur}]crop=w='{cw_expr}':h='{ch_expr}':x='(iw-{cw_expr})/2':y='(ih-{ch_expr})/2',"
            f"scale={output_width}:{output_height}[vcrop]"
        )
        cur = "vcrop"

    # 5. drawtext 字幕（标题整段 + ASR 逐句）
    draw_filters: list[str] = []
    if title_text.strip():
        title_size = max(int(output_height * 0.05), 24)
        draw_filters.append(
            build_drawtext_filter(
                text=title_text,
                start=0.0,
                end=max(total_duration, 0.1),
                font=font,
                font_size=title_size,
                y_expr="h-th-40",
                box=True,
            )
        )
    sub_size = max(int(output_height * 0.045), 20)
    for seg in subtitle_segments or []:
        txt = getattr(seg, "text", "") or ""
        if not txt.strip():
            continue
        draw_filters.append(
            build_drawtext_filter(
                text=txt,
                start=float(getattr(seg, "start", 0)),
                end=float(getattr(seg, "end", 0)),
                font=font,
                font_size=sub_size,
                y_expr="h-th-60",
                borderw=2,
            )
        )

    if draw_filters:
        prev = cur
        for idx, df in enumerate(draw_filters):
            out_l = "vfinal" if idx == len(draw_filters) - 1 else f"vd{idx}"
            fc.append(f"[{prev}]{df}[{out_l}]")
            prev = out_l
        vfinal_label = prev
    else:
        fc.append(f"[{cur}]format=yuv420p[vfinal]")
        vfinal_label = "vfinal"

    # 6. 音频输入与混音
    audio_items: list[tuple[int, float]] = []  # (input_index, volume)
    next_idx = n
    if tts_audio and Path(tts_audio).exists():
        turl, tkey = upload_local_audio_and_sign(Path(tts_audio))
        tname = "tts" + (Path(tts_audio).suffix or ".mp3")
        inputs[tname] = turl
        oss_keys.append(tkey)
        input_args.extend(["-i", tname])
        audio_items.append((next_idx, 1.0))
        next_idx += 1
    if bgm_audio and Path(bgm_audio).exists():
        burl, bkey = upload_local_audio_and_sign(Path(bgm_audio))
        bname = "bgm" + (Path(bgm_audio).suffix or ".mp3")
        inputs[bname] = burl
        oss_keys.append(bkey)
        input_args.extend(["-i", bname])
        audio_items.append((next_idx, 0.35))
        next_idx += 1

    maps: list[str] = ["-map", f"[{vfinal_label}]"]
    if audio_items:
        mix_labels: list[str] = []
        for k, (idx, vol) in enumerate(audio_items):
            alabel = f"au{k}"
            fc.append(
                f"[{idx}:a]aresample=44100,volume={vol:.2f},aformat=sample_fmts=fltp:channel_layouts=stereo[{alabel}]"
            )
            mix_labels.append(alabel)
        mix_in = "".join(f"[{l}]" for _lbl in mix_labels)
        fc.append(f"{mix_in}amix=inputs={len(mix_labels)}:duration=first:dropout_transition=2,aresample=44100[afinal]")
        maps.extend(["-map", "[afinal]", "-c:a", "aac", "-b:a", "128k"])

    # 7. 组装 ffmpeg_args + NVENC 编码
    ffmpeg_args = ["-y", *input_args, "-filter_complex", ";".join(fc), *maps]
    ffmpeg_args.extend(["-c:v", vcodec, "-preset", preset, "-pix_fmt", "yuv420p"])
    if video_bitrate:
        ffmpeg_args.extend(["-b:v", video_bitrate])
    else:
        ffmpeg_args.extend(["-cq", str(cq)])
    ffmpeg_args.extend(["-movflags", "+faststart", "-f", "mp4", "pipe:1"])

    return DirectRenderPlan(inputs=inputs, ffmpeg_args=ffmpeg_args, oss_keys=oss_keys)