ComfyUI-ShotSplitter/cutter.py

"""Frame-accurate H.264+AAC shot cutter via parallel FFmpeg subprocess pool."""

import logging
import os
import subprocess
from concurrent.futures import ThreadPoolExecutor, as_completed
from dataclasses import dataclass
from typing import List, Optional, Tuple

logger = logging.getLogger(__name__)

FFMPEG_PARALLELISM = 4


@dataclass
class CutResult:
    index: int
    path: str
    start_s: float
    end_s: float
    success: bool
    error: Optional[str] = None


def cut_shots(
    source_path: str,
    intervals: List[Tuple[int, int]],
    fps: float,
    output_dir: str,
    filename_prefix: str,
    crf: int,
    preset: str,
    keep_audio: bool,
    has_audio: bool,
    vfr: bool,
) -> List[CutResult]:
    """Cut `intervals` (frame indices, half-open [start,end)) into individual mp4s.

    Uses one ffmpeg subprocess per shot, in parallel up to FFMPEG_PARALLELISM workers.
    Returns a list of CutResult in the same order as `intervals`.
    """
    os.makedirs(output_dir, exist_ok=True)
    pad = max(3, len(str(len(intervals))))
    jobs: List[Tuple[int, str, float, float]] = []
    for i, (start_f, end_f) in enumerate(intervals):
        start_s = start_f / fps
        duration_s = max((end_f - start_f) / fps, 1.0 / fps)
        out_name = f"{filename_prefix}_{str(i + 1).zfill(pad)}.mp4"
        out_path = os.path.join(output_dir, out_name)
        jobs.append((i, out_path, start_s, duration_s))

    results: List[Optional[CutResult]] = [None] * len(jobs)

    def _run_one(idx: int, out_path: str, start_s: float, duration_s: float) -> CutResult:
        cmd = _build_cmd(
            source_path,
            out_path,
            start_s,
            duration_s,
            crf,
            preset,
            keep_audio and has_audio,
            vfr,
            fps,
        )
        try:
            subprocess.run(cmd, check=True, capture_output=True, text=True)
            return CutResult(idx, out_path, start_s, start_s + duration_s, True, None)
        except subprocess.CalledProcessError as e:
            logger.error("ffmpeg failed for shot %d: %s", idx, e.stderr)
            return CutResult(idx, out_path, start_s, start_s + duration_s, False, e.stderr)

    with ThreadPoolExecutor(max_workers=FFMPEG_PARALLELISM) as ex:
        futures = {ex.submit(_run_one, *job): job[0] for job in jobs}
        for fut in as_completed(futures):
            idx = futures[fut]
            results[idx] = fut.result()

    return [r for r in results if r is not None]


def _build_cmd(
    source: str,
    out: str,
    start_s: float,
    duration_s: float,
    crf: int,
    preset: str,
    audio: bool,
    vfr: bool,
    fps: float,
) -> List[str]:
    cmd = [
        "ffmpeg",
        "-y",
        "-hide_banner",
        "-loglevel", "error",
        "-ss", f"{start_s:.6f}",
        "-i", source,
        "-t", f"{duration_s:.6f}",
        "-map", "0:v:0",
    ]
    if audio:
        cmd += ["-map", "0:a:0?"]
    if vfr:
        cmd += ["-vsync", "cfr", "-r", f"{fps:.6f}"]
    cmd += [
        "-c:v", "libx264",
        "-preset", preset,
        "-crf", str(crf),
        "-pix_fmt", "yuv420p",
    ]
    if audio:
        cmd += ["-c:a", "aac", "-b:a", "192k"]
    else:
        cmd += ["-an"]
    cmd += [
        "-avoid_negative_ts", "make_zero",
        out,
    ]
    return cmd