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 只编码一次；
原素材音轨 concat + TTS/配音/BGM 混音也在同一命令里完成。

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

from __future__ import annotations

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

logger = logging.getLogger(__name__)

DEFAULT_DRAWTEXT_FONT = "Noto Sans CJK SC"
EDGE_CROP_MIN_PCT = 0.02
EDGE_CROP_MAX_PCT = 0.05


def escape_drawtext_text(text: str) -> str:
    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(";", "\\;")
    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:
    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 _build_atempo_chain(speed: float) -> str:
    if abs(speed - 1.0) < 1e-6:
        return ""
    stages: list[float] = []
    remaining = speed
    while remaining > 2.0:
        stages.append(2.0)
        remaining /= 2.0
    while remaining < 0.5:
        stages.append(0.5)
        remaining /= 0.5
    if abs(remaining - 1.0) >= 1e-6:
        stages.append(remaining)
    return ",".join(f"atempo={s:.5f}" for s in stages)


def upload_local_audio_and_sign(
    local_audio: Path,
    *,
    tmp_prefix: str = "tmp/gpu-direct-audio/",
    expires: int = 3600,
) -> tuple[str, str]:
    from video_processing.oss_helpers import _storage  # type: 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:
    from video_processing.oss_helpers import _storage  # type: ignore

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


class DirectRenderPlan:
    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


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,
    clip_has_audio: Optional[list[bool]] = None,
    clip_volumes: Optional[list[float]] = None,
    extra_audio_tracks: Optional[list[tuple[Any, float]]] = None,
) -> DirectRenderPlan:
    """构造 P4000 直连渲染所需的 inputs 与 ffmpeg_args。

    视频：每段 trim/setpts/scale/pad/fps → concat（全硬切，带音频）→ 随机边缘 crop+scale → drawtext。
    音频：每段 [i:a]（或 anullsrc 静音占位）按 clip 配置 atrim/asetpts/atempo/volume/aresample
         → concat=n:N:v=1:a=1 → 与 extra_audio（TTS/配音素材库）、BGM 一起 amix → atrim 精确截断。
    """
    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] = []
    n = len(resolved_clips)

    # 规范化每段参数
    if clip_has_audio is None:
        clip_has_audio = [True] * n
    else:
        clip_has_audio = list(clip_has_audio) + [True] * max(0, n - len(clip_has_audio))
        clip_has_audio = clip_has_audio[:n]
    if clip_volumes is None:
        clip_volumes = [1.0] * n
    else:
        clip_volumes = list(clip_volumes) + [1.0] * max(0, n - len(clip_volumes))
        clip_volumes = clip_volumes[:n]

    clip_starts: list[float] = []
    clip_effs: list[float] = []
    clip_speeds: list[float] = []
    for clip in resolved_clips:
        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)
        speed = float(getattr(clip, "playback_speed", 1.0) or 1.0)
        clip_starts.append(start)
        clip_effs.append(eff)
        clip_speeds.append(speed)

    # 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 in range(n):
        vf: list[str] = []
        start, eff, speed = clip_starts[i], clip_effs[i], clip_speeds[i]
        if eff > 0:
            if start > 0:
                vf.append(f"trim=start={start:.3f}:duration={eff:.3f}")
            else:
                vf.append(f"trim=duration={eff:.3f}")
            vf.append("setpts=PTS-STARTPTS")
        if abs(speed - 1.0) >= 1e-6:
            vf.append(f"setpts=PTS/{speed:.4f}")
        vf.append(f"scale={output_width}:{output_height}:force_original_aspect_ratio=decrease")
        vf.append(f"pad={output_width}:{output_height}:trunc((ow-iw)/2):trunc((oh-ih)/2):black")
        vf.append("setpts=PTS-STARTPTS")
        vf.append(f"fps={output_fps}")
        label = f"vc{i}"
        fc.append(f"[{i}:v]{','.join(vf)}[{label}]")
        pre_labels.append(label)

    # 2b. 音频段预处理（无声源用 anullsrc 占位；volume=0 的段也用 anullsrc 静音占位保持时间轴）
    anullsrc_counter = 0
    audio_pre_labels: list[str] = []
    for i in range(n):
        start, eff, speed = clip_starts[i], clip_effs[i], clip_speeds[i]
        vol = float(clip_volumes[i] if i < len(clip_volumes) else 1.0)
        has_a = bool(clip_has_audio[i] if i < len(clip_has_audio) else True)
        if not has_a or vol <= 0.001:
            # 静音占位：用 anullsrc 生成静音，atrim 到段时长
            sl = f"sil{anullsrc_counter}"
            anullsrc_counter += 1
            af: list[str] = ["anullsrc=channel_layout=stereo:sample_rate=44100"]
            if eff > 0:
                af.append(f"atrim=duration={eff:.3f}")
                af.append("asetpts=PTS-STARTPTS")
            af.append("aformat=sample_fmts=fltp:channel_layouts=stereo")
            fc.append(f"{','.join(af)}[{sl}]")
            # anullsrc 作为 filter 源不需要 -i 输入，直接给 label
            audio_pre_labels.append(sl)
            continue

        af = []
        if eff > 0:
            if start > 0:
                af.append(f"atrim=start={start:.3f}:duration={eff:.3f}")
            else:
                af.append(f"atrim=duration={eff:.3f}")
            af.append("asetpts=PTS-STARTPTS")
        if abs(speed - 1.0) >= 1e-6:
            atempo = _build_atempo_chain(speed)
            if atempo:
                af.append(atempo)
        if abs(vol - 1.0) >= 1e-3:
            af.append(f"volume={vol:.3f}")
        af.append("aresample=44100")
        af.append("aformat=sample_fmts=fltp:channel_layouts=stereo")
        alabel = f"ac{i}"
        fc.append(f"[{i}:a]{','.join(af)}[{alabel}]")
        audio_pre_labels.append(alabel)

    # 3. concat（全硬切；v=1:a=1，视频音频一起拼接）
    concat_in = "".join(f"[{v}][{a}]" for v, a in zip(pre_labels, audio_pre_labels, strict=True))
    fc.append(f"{concat_in}concat=n={n}:v=1:a=1[vcat][acat]")
    cur_v = "vcat"
    cur_a = "acat"

    # 4. 随机边缘裁剪降重（四边独立随机 2%~5%，与 ffmpeg_utils.random_edge_crop 一致）
    if edge_crop_pct and edge_crop_pct > 0:
        _r = random.Random()
        p_min = EDGE_CROP_MIN_PCT
        p_max = EDGE_CROP_MAX_PCT
        crop_top = p_min + _r.random() * (p_max - p_min)
        crop_bottom = p_min + _r.random() * (p_max - p_min)
        crop_left = p_min + _r.random() * (p_max - p_min)
        crop_right = p_min + _r.random() * (p_max - p_min)
        w_expr = f"trunc(iw*(1-{crop_left:.4f}-{crop_right:.4f})/2)*2"
        h_expr = f"trunc(ih*(1-{crop_top:.4f}-{crop_bottom:.4f})/2)*2"
        x_expr = f"trunc(iw*{crop_left:.4f}/2)*2"
        y_expr = f"trunc(ih*{crop_top:.4f}/2)*2"
        fc.append(
            f"[{cur_v}]crop=w='{w_expr}':h='{h_expr}':x='{x_expr}':y='{y_expr}',"
            f"scale={output_width}:{output_height}[vcrop]"
        )
        cur_v = "vcrop"

    # 5. drawtext 字幕
    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_v
        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_v}]format=yuv420p[vfinal]")
        vfinal_label = "vfinal"

    # 6. 音频混音：原素材主音轨 acat + extra（TTS/配音素材库） + BGM → amix → atrim
    mix_labels: list[str] = [cur_a]
    mix_vols: list[float] = [1.0]
    next_idx = n

    # 额外独立音频轨（TTS concat / 配音素材库整段音频）
    for _ea_idx, (ea_path, ea_vol) in enumerate(extra_audio_tracks or []):
        if ea_path is None:
            continue
        ea_p = Path(ea_path)
        if not ea_p.exists():
            continue
        eurl, ekey = upload_local_audio_and_sign(ea_p)
        ename = f"extra{_ea_idx}{ea_p.suffix or '.mp3'}"
        inputs[ename] = eurl
        oss_keys.append(ekey)
        input_args.extend(["-i", ename])
        elabel = f"aex{_ea_idx}"
        fc.append(
            f"[{next_idx}:a]aresample=44100,volume={float(ea_vol):.2f},"
            f"aformat=sample_fmts=fltp:channel_layouts=stereo[{elabel}]"
        )
        mix_labels.append(elabel)
        mix_vols.append(float(ea_vol))
        next_idx += 1

    if tts_audio and Path(tts_audio).exists():
        # 旧参数保留：若调用方直接传了 tts_audio 而没走 extra_audio_tracks，则仍然加入
        # （兼容旧调用，正常路径 TTS 已经通过 extra_audio_tracks 传入）
        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])
        alabel = "au_tts"
        fc.append(
            f"[{next_idx}:a]aresample=44100,volume=1.00,aformat=sample_fmts=fltp:channel_layouts=stereo[{alabel}]"
        )
        mix_labels.append(alabel)
        mix_vols.append(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])
        alabel = "au_bgm"
        fc.append(
            f"[{next_idx}:a]aresample=44100,volume=0.35,aformat=sample_fmts=fltp:channel_layouts=stereo[{alabel}]"
        )
        mix_labels.append(alabel)
        mix_vols.append(0.35)
        next_idx += 1

    maps: list[str] = ["-map", f"[{vfinal_label}]"]
    if mix_labels:
        mix_in = "".join(f"[{lb}]" for lb in mix_labels)
        n_mix = len(mix_labels)
        mix_parts = [
            f"amix=inputs={n_mix}:duration=longest:dropout_transition=2:normalize=0",
            "aresample=44100",
        ]
        # Bug2 修复：atrim 到视频精确时长
        if total_duration and total_duration > 0:
            mix_parts.append(f"atrim=0:{total_duration:.3f}")
            mix_parts.append("asetpts=PTS-STARTPTS")
        fc.append(f"{mix_in}{','.join(mix_parts)}[afinal]")
        maps.extend(["-map", "[afinal]", "-c:a", "aac", "-b:a", "128k"])
    else:
        logger.info("[gpu-direct] no audio tracks; output silent video")

    # 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", "-shortest", "-f", "mp4", "pipe:1"])

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