video-gen/video-gen-api/app/services/video_upscale/local_ffmpeg_service.py

from __future__ import annotations

import asyncio
import os
from collections.abc import Awaitable, Callable

from app.config import settings
from app.services.video_cover_service import get_ffmpeg_bin
from app.services.video_upscale.media_service import (
    build_part_mp4_path,
    is_valid_file,
    probe_video,
    safe_remove,
)


class LocalVideoUpscaleError(RuntimeError):
    pass


def _build_filter(target_width: int, target_height: int) -> str:
    if target_width <= 0 or target_height <= 0 or target_width % 2 or target_height % 2:
        raise LocalVideoUpscaleError("目标宽高必须是正偶数")
    return (
        "hqdn3d=0.8:0.6:2.5:1.8,"
        f"scale={target_width}:{target_height}:"
        "force_original_aspect_ratio=increase:force_divisible_by=2:flags=lanczos,"
        f"crop={target_width}:{target_height}:(iw-ow)/2:(ih-oh)/2,"
        "unsharp=5:5:0.40:5:5:0.0,"
        "eq=contrast=1.02:saturation=1.03,"
        "setsar=1"
    )


def _build_command(
    *,
    source_path: str,
    part_path: str,
    target_width: int,
    target_height: int,
) -> list[str]:
    return [
        get_ffmpeg_bin(),
        "-hide_banner",
        "-nostdin",
        "-y",
        "-i",
        source_path,
        "-map",
        "0:v:0",
        "-map",
        "0:a?",
        "-vf",
        _build_filter(target_width, target_height),
        "-c:v",
        "libx264",
        "-preset",
        "slow",
        "-crf",
        "18",
        "-profile:v",
        "high",
        "-pix_fmt",
        "yuv420p",
        "-movflags",
        "+faststart",
        "-c:a",
        "aac",
        "-b:a",
        "192k",
        part_path,
    ]


async def _terminate_process(process: asyncio.subprocess.Process) -> None:
    if process.returncode is not None:
        return
    process.terminate()
    try:
        await asyncio.wait_for(process.wait(), timeout=5)
    except asyncio.TimeoutError:
        process.kill()
        await process.wait()


async def _run_ffmpeg(
    *,
    source_path: str,
    part_path: str,
    target_width: int,
    target_height: int,
    timeout_seconds: int,
    execution_guard: Callable[[], Awaitable[None]] | None,
) -> None:
    try:
        process = await asyncio.create_subprocess_exec(
            *_build_command(
                source_path=source_path,
                part_path=part_path,
                target_width=target_width,
                target_height=target_height,
            ),
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
        )
    except OSError as exc:
        raise LocalVideoUpscaleError(f"本地 FFmpeg 无法启动: {exc}") from exc

    communicate_task = asyncio.create_task(process.communicate())
    deadline = asyncio.get_running_loop().time() + max(30, int(timeout_seconds))
    try:
        while not communicate_task.done():
            if execution_guard is not None:
                await execution_guard()
            remaining = deadline - asyncio.get_running_loop().time()
            if remaining <= 0:
                raise LocalVideoUpscaleError(f"本地 FFmpeg 超分超时: {timeout_seconds} 秒")
            try:
                await asyncio.wait_for(asyncio.shield(communicate_task), timeout=min(2.0, remaining))
            except asyncio.TimeoutError:
                continue
        stdout, stderr = await communicate_task
    except BaseException:
        await _terminate_process(process)
        if not communicate_task.done():
            communicate_task.cancel()
        try:
            await communicate_task
        except BaseException:
            pass
        raise

    if process.returncode != 0:
        error_text = (stderr or b"").decode("utf-8", errors="replace").strip()
        raise LocalVideoUpscaleError(f"本地 FFmpeg 超分失败: {error_text[-4000:]}")


async def execute_local_ffmpeg_crop(
    *,
    source_path: str,
    final_path: str,
    target_width: int,
    target_height: int,
    timeout_seconds: int | None = None,
    execution_guard: Callable[[], Awaitable[None]] | None = None,
) -> str:
    if not is_valid_file(source_path):
        raise LocalVideoUpscaleError(f"超分源视频不存在或为空: {source_path}")
    if is_valid_file(final_path):
        info = await probe_video(final_path)
        if info.width == target_width and info.height == target_height:
            return final_path

    os.makedirs(os.path.dirname(final_path), exist_ok=True)
    part_path = build_part_mp4_path(final_path)
    safe_remove(part_path)
    try:
        await _run_ffmpeg(
            source_path=source_path,
            part_path=part_path,
            target_width=target_width,
            target_height=target_height,
            timeout_seconds=int(timeout_seconds or settings.VIDEO_UPSCALE_LOCAL_TIMEOUT_SECONDS),
            execution_guard=execution_guard,
        )
        if not is_valid_file(part_path):
            raise LocalVideoUpscaleError("本地 FFmpeg 输出文件为空")
        info = await probe_video(part_path)
        if info.width != target_width or info.height != target_height:
            raise LocalVideoUpscaleError(
                f"本地 FFmpeg 输出尺寸不正确: {info.width}x{info.height}，预期 {target_width}x{target_height}"
            )
        if execution_guard is not None:
            await execution_guard()
        os.replace(part_path, final_path)
        return final_path
    except Exception:
        safe_remove(part_path)
        raise