ovi/utils/av_edit_dataset.py

"""Canonical paired dataset for instruction-based audio-video editing."""

from __future__ import annotations

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

import torch

from ovi.utils.av_edit_data import (
    get_instruction,
    get_video_info,
    load_audio,
    load_audio_array,
    load_manifest,
    load_video_array,
    pad_or_trim_audio,
    resolve_media_path,
    snap_num_frames,
)


class AVEditDataset(torch.utils.data.Dataset):
    """Load source/target AV pairs from one canonical manifest.

    Required media columns depend on ``has_video`` and ``has_audio``. In AV
    mode, audio columns are optional and fall back to the corresponding video
    tracks. In audio-only mode, source and target audio columns are required.
    """

    def __init__(
        self,
        metadata_path: str,
        base_path: Optional[str] = None,
        source_video_column: str = "source_video",
        source_audio_column: str = "source_audio",
        target_video_column: str = "target_video",
        target_audio_column: str = "target_audio",
        instruction_column: str = "instruction",
        has_video: bool = True,
        has_audio: bool = True,
        height: Optional[int] = 704,
        width: Optional[int] = 1280,
        max_pixels: int = 1920 * 1080,
        max_num_frames: Optional[int] = 121,
        fix_num_frames: Optional[int] = None,
        audio_sample_rate: int = 16000,
        max_audio_seconds: Optional[float] = None,
        fix_audio_seconds: Optional[float] = None,
        repeat: int = 1,
        max_retries: int = 20,
    ) -> None:
        self.metadata_path = str(Path(metadata_path).expanduser().resolve())
        self.base_path = base_path
        self.rows = load_manifest(self.metadata_path)
        if not self.rows:
            raise ValueError(f"Manifest is empty: {self.metadata_path}")

        self.source_video_column = source_video_column
        self.source_audio_column = source_audio_column
        self.target_video_column = target_video_column
        self.target_audio_column = target_audio_column
        self.instruction_column = instruction_column
        self.has_video = bool(has_video)
        self.has_audio = bool(has_audio)
        self.height = height
        self.width = width
        self.max_pixels = int(max_pixels)
        self.max_num_frames = None if max_num_frames is None else int(max_num_frames)
        self.fix_num_frames = None if fix_num_frames is None else int(fix_num_frames)
        self.audio_sample_rate = int(audio_sample_rate)
        self.max_audio_seconds = (
            None if max_audio_seconds is None else float(max_audio_seconds)
        )
        self.fix_audio_seconds = (
            None if fix_audio_seconds is None else float(fix_audio_seconds)
        )
        self.repeat = int(repeat)
        self.max_retries = int(max_retries)

        if not self.has_video and not self.has_audio:
            raise ValueError("At least one of has_video or has_audio must be true.")
        if self.repeat < 1:
            raise ValueError("repeat must be at least 1.")
        if (
            self.has_video
            and self.fix_num_frames is not None
            and self.fix_num_frames % 4 != 1
        ):
            raise ValueError("fix_num_frames must satisfy num_frames % 4 == 1.")
        if self.max_audio_seconds is not None and self.max_audio_seconds <= 0:
            raise ValueError("max_audio_seconds must be positive.")
        if self.fix_audio_seconds is not None and self.fix_audio_seconds <= 0:
            raise ValueError("fix_audio_seconds must be positive.")
        if self.max_audio_seconds is not None and self.fix_audio_seconds is not None:
            raise ValueError(
                "Set only one of max_audio_seconds or fix_audio_seconds."
            )

    def __len__(self) -> int:
        return len(self.rows) * self.repeat

    def _path(self, row: dict[str, Any], column: str, required: bool) -> Optional[Path]:
        if column not in row:
            if required:
                raise KeyError(f"Manifest is missing required column '{column}'.")
            return None
        return resolve_media_path(
            row[column],
            manifest_path=self.metadata_path,
            base_path=self.base_path,
            required=required,
        )

    def _load_row(self, index: int) -> dict[str, Any]:
        row = self.rows[index % len(self.rows)].copy()
        source_video_path = self._path(
            row, self.source_video_column, required=self.has_video
        )
        target_video_path = self._path(
            row, self.target_video_column, required=self.has_video
        )
        source_audio_path = self._path(
            row,
            self.source_audio_column,
            required=self.has_audio and not self.has_video,
        )
        target_audio_path = self._path(
            row,
            self.target_audio_column,
            required=self.has_audio and not self.has_video,
        )

        row["instruction"] = get_instruction(row, self.instruction_column)
        duration_seconds = None

        if self.has_video:
            source_fps, source_frames = get_video_info(source_video_path)
            _, target_frames = get_video_info(target_video_path)
            available_frames = min(source_frames, target_frames)
            if self.fix_num_frames is not None:
                if available_frames < self.fix_num_frames:
                    raise ValueError(
                        f"Pair has {available_frames} frames but "
                        f"fix_num_frames={self.fix_num_frames}."
                    )
                num_frames = self.fix_num_frames
            else:
                if self.max_num_frames is not None:
                    available_frames = min(available_frames, self.max_num_frames)
                num_frames = snap_num_frames(available_frames)

            source_video, target_size = load_video_array(
                source_video_path,
                num_frames=num_frames,
                height=self.height,
                width=self.width,
                max_pixels=self.max_pixels,
            )
            target_video, _ = load_video_array(
                target_video_path,
                num_frames=num_frames,
                height=self.height,
                width=self.width,
                max_pixels=self.max_pixels,
                target_size=target_size,
            )
            duration_seconds = num_frames / source_fps
            row.update(
                {
                    "video_ori_np": source_video,
                    "video_result_np": target_video,
                    "source_video_path": str(source_video_path),
                    "target_video_path": str(target_video_path),
                    "num_frames": num_frames,
                    "fps": source_fps,
                }
            )

        if self.has_audio:
            resolved_source_audio = source_audio_path or source_video_path
            resolved_target_audio = target_audio_path or target_video_path
            if resolved_source_audio is None or resolved_target_audio is None:
                raise ValueError("Audio-only training requires source and target audio files.")

            if duration_seconds is not None:
                num_audio_samples = max(
                    1, round(duration_seconds * self.audio_sample_rate)
                )
                source_audio = load_audio_array(
                    resolved_source_audio,
                    sample_rate=self.audio_sample_rate,
                    num_samples=num_audio_samples,
                )
                target_audio = load_audio_array(
                    resolved_target_audio,
                    sample_rate=self.audio_sample_rate,
                    num_samples=num_audio_samples,
                )
            else:
                source_audio = load_audio(
                    resolved_source_audio, sample_rate=self.audio_sample_rate
                )
                target_audio = load_audio(
                    resolved_target_audio, sample_rate=self.audio_sample_rate
                )
                if self.fix_audio_seconds is not None:
                    num_audio_samples = max(
                        1, round(self.fix_audio_seconds * self.audio_sample_rate)
                    )
                else:
                    num_audio_samples = max(len(source_audio), len(target_audio))
                    if self.max_audio_seconds is not None:
                        num_audio_samples = min(
                            num_audio_samples,
                            max(
                                1,
                                round(
                                    self.max_audio_seconds * self.audio_sample_rate
                                ),
                            ),
                        )
                source_audio = pad_or_trim_audio(source_audio, num_audio_samples)
                target_audio = pad_or_trim_audio(target_audio, num_audio_samples)

            row.update(
                {
                    "audio_ori_np": source_audio,
                    "audio_result_np": target_audio,
                    "source_audio_path": str(resolved_source_audio),
                    "target_audio_path": str(resolved_target_audio),
                    "num_audio_samples": len(source_audio),
                    "audio_sample_rate": self.audio_sample_rate,
                }
            )
        return row

    def __getitem__(self, index: int) -> dict[str, Any]:
        errors = []
        candidate = index
        for _ in range(self.max_retries):
            try:
                return self._load_row(candidate)
            except Exception as error:
                errors.append(f"row {candidate % len(self.rows)}: {error}")
                candidate = random.randrange(len(self.rows))
        details = "\n".join(errors[-3:])
        raise RuntimeError(f"Failed to load an AV editing sample after retries:\n{details}")


UnifiedAVEditDataset = AVEditDataset