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,
) -> DirectRenderPlan:
    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)

    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)

    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])

    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)

    audio_pre_labels: list[str] = []
    for i in range(n):
        af: list[str] = []
        start, eff, speed = clip_starts[i], clip_effs[i], clip_speeds[i]
        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)
        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)

    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"

    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"

    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"

    mix_labels: list[str] = [cur_a]
    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])
        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)
        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)
        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",
        ]
        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")

    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)