MMaDA/precompute_video_speech_tokens.py

#!/usr/bin/env python3
"""
Pre-compute EMOVA speech tokenizer codes for audio datasets.

Supported dataset types:
    - video-speech : CSV index with truncated WAV clips (e.g., OpenVid speech)
    - librispeech  : LibriSpeech directory structure with FLAC audio
    - instructs2s  : InstructS2S-200K style user/assistant WAV pairs

Examples
--------
    # VideoSpeech
    python MMaDA/precompute_video_speech_tokens.py \\
        --dataset-type video-speech \\
        --index /home/work/AIDAS/data/video-speech/openvid-speech.csv \\
        --audio-root /home/work/AIDAS/data/video-speech/openvid-speech-trunc \\
        --output /home/work/AIDAS/cache/video_speech_tokens

    # LibriSpeech
    python MMaDA/precompute_video_speech_tokens.py \\
        --dataset-type librispeech \\
        --audio-root /home/work/AIDAS/data/audio/LibriSpeech \\
        --librispeech-subsets train-clean-360 train-clean-100 \\
        --output /home/work/AIDAS/cache/librispeech_tokens

    # InstructS2S (pairs.txt assumed under audio root)
    python MMaDA/precompute_video_speech_tokens.py \\
        --dataset-type instructs2s \\
        --audio-root /home/work/AIDAS/data/InstructS2S-200K/en/wav \\
        --output /home/work/AIDAS/cache/instructs2s_tokens
"""

import argparse
import csv
import hashlib
import os
import sys
import tempfile
from pathlib import Path
from typing import Iterable, Iterator, List, Set

import soundfile as sf
import torch
from tqdm import tqdm

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from models.modeling_emova_speech_tokenizer import EMOVASpeechTokenizer  # noqa: E402


def iter_video_speech_audio(index_path: Path, audio_root: Path) -> Iterator[Path]:
    """
    Yields audio paths from the VideoSpeech index CSV.
    """
    with index_path.open("r", newline="") as csvfile:
        reader = csv.reader(csvfile)
        for row in reader:
            if not row:
                continue
            base = row[0].strip().removesuffix(".wav")
            if not base:
                continue
            audio_path = audio_root / f"{base}.wav"
            if audio_path.is_file():
                yield audio_path


def iter_librispeech_audio(audio_root: Path, subsets: Iterable[str]) -> Iterator[Path]:
    """
    Iterates through LibriSpeech FLAC files for the provided subsets.
    """
    for subset in subsets:
        subset_dir = audio_root / subset
        if not subset_dir.exists():
            raise FileNotFoundError(f"LibriSpeech subset not found: {subset_dir}")
        speakers = sorted(p for p in subset_dir.iterdir() if p.is_dir())
        for speaker_dir in speakers:
            chapters = sorted(p for p in speaker_dir.iterdir() if p.is_dir())
            for chapter_dir in chapters:
                for flac_path in sorted(chapter_dir.glob("*.flac")):
                    yield flac_path


def iter_instructs2s_audio(audio_root: Path, pairs_file: Path | None = None) -> Iterator[Path]:
    """
    Yields unique audio paths from an InstructS2S root directory.
    If pairs_file is provided (or found under audio_root), it's expected to contain
    two space-separated paths per line: user assistant.
    Otherwise, the directory tree is scanned similarly to Speech2SpeechDataset.
    """
    resolved_root = audio_root.expanduser().resolve()
    if pairs_file is None:
        candidate = resolved_root / "pairs.txt"
        if candidate.exists():
            pairs_file = candidate
    if pairs_file is not None:
        with Path(pairs_file).open("r") as fh:
            for line in fh:
                line = line.strip()
                if not line:
                    continue
                parts = line.split()
                if len(parts) >= 2:
                    user_path = Path(parts[0])
                    if not user_path.is_absolute():
                        user_path = resolved_root / user_path
                    asst_path = Path(parts[1])
                    if not asst_path.is_absolute():
                        asst_path = resolved_root / asst_path
                    if user_path.is_file():
                        yield user_path
                    if asst_path.is_file():
                        yield asst_path
        return

    dirs = [p for p in resolved_root.glob("*") if p.is_dir()]
    for dir_path in dirs:
        dir_name = dir_path.name
        k = 1
        while True:
            user_wav = dir_path / f"{dir_name}-{k}-user.wav"
            assistant_wav = dir_path / f"{dir_name}-{k}-assistant.wav"
            if user_wav.is_file() and assistant_wav.is_file():
                yield user_wav
                yield assistant_wav
                k += 1
                continue
            break


def hash_path(path: Path) -> str:
    """Returns a SHA-1 hex digest for the absolute path."""
    abs_path = os.path.abspath(path)
    return hashlib.sha1(abs_path.encode("utf-8")).hexdigest()


def token_output_path(output_root: Path, audio_path: Path) -> Path:
    """Resolves the on-disk location for cached tokens corresponding to audio_path."""
    digest = hash_path(audio_path)
    return output_root / digest[:2] / digest[2:4] / f"{digest}.pt"


def encode_audio(tokenizer: EMOVASpeechTokenizer, audio_path: Path) -> torch.Tensor:
    """
    Encodes an audio file to discrete tokens, converting non-WAV inputs on the fly.
    """
    suffix = audio_path.suffix.lower()
    if suffix == ".wav":
        return tokenizer.encode(str(audio_path)).cpu()

    data, sample_rate = sf.read(str(audio_path))
    tmp_file = tempfile.NamedTemporaryFile(suffix=".wav", delete=False)
    try:
        sf.write(tmp_file.name, data, sample_rate)
        tokens = tokenizer.encode(tmp_file.name).cpu()
    finally:
        tmp_file.close()
        try:
            os.remove(tmp_file.name)
        except OSError:
            pass
    return tokens


def gather_audio_paths(args) -> List[Path]:
    if args.dataset_type == "video-speech":
        return list(iter_video_speech_audio(args.index, args.audio_root))
    if args.dataset_type == "librispeech":
        return list(iter_librispeech_audio(args.audio_root, args.librispeech_subsets))
    # instructs2s
    paths = list(iter_instructs2s_audio(args.audio_root, args.pairs_file))
    # Deduplicate while preserving order
    seen: Set[Path] = set()
    unique_paths: List[Path] = []
    for path in paths:
        if path not in seen:
            seen.add(path)
            unique_paths.append(path)
    return unique_paths


def split_into_shards(items: List[Path], shard_count: int) -> List[List[Path]]:  # pragma: no cover - simple helper
    shard_count = max(1, shard_count)
    shard_size = (len(items) + shard_count - 1) // shard_count
    return [items[i * shard_size : (i + 1) * shard_size] for i in range(shard_count)]


def process_shard(
    shard_id: int,
    audio_paths: List[Path],
    device: str,
    tokenizer_name: str,
    output_root: Path,
    overwrite: bool,
    dataset_type: str,
) -> tuple[int, int, List[Path]]:
    if not audio_paths:
        return 0, 0, []

    device_obj = torch.device(device)
    if device_obj.type == "cuda":
        torch.cuda.set_device(device_obj)
    tokenizer = EMOVASpeechTokenizer.from_pretrained(tokenizer_name).to(device_obj)
    tokenizer.eval()

    total = 0
    skipped = 0
    desc = f"{dataset_type} worker {shard_id}"
    failed_paths: List[Path] = []
    for audio_path in tqdm(audio_paths, desc=desc, position=shard_id, leave=False):
        out_path = token_output_path(output_root, audio_path)
        if out_path.exists() and not overwrite:
            skipped += 1
            continue
        out_path.parent.mkdir(parents=True, exist_ok=True)
        try:
            tokens = encode_audio(tokenizer, audio_path)
        except Exception as exc:  # pragma: no cover - runtime diagnostics
            tqdm.write(f"[WARN][worker {shard_id}] Failed to encode {audio_path}: {exc}")
            failed_paths.append(audio_path)
            continue
        tmp_path = out_path.with_suffix(out_path.suffix + ".tmp")
        torch.save(tokens, tmp_path)
        os.replace(tmp_path, out_path)
        total += 1
    return total, skipped, failed_paths


def main():
    parser = argparse.ArgumentParser(description="Pre-compute speech tokens for audio datasets.")
    parser.add_argument(
        "--dataset-type",
        "--dataset_type",
        dest="dataset_type",
        choices=["video-speech", "librispeech", "instructs2s"],
        default="video-speech",
        help="Dataset type to process.",
    )
    parser.add_argument(
        "--index",
        type=Path,
        help="CSV index for video-speech datasets (required for dataset-type=video-speech).",
    )
    parser.add_argument(
        "--audio-root",
        type=Path,
        required=True,
        help="Root directory containing audio files. For LibriSpeech this should be the LibriSpeech root.",
    )
    parser.add_argument(
        "--librispeech_subsets",
        nargs="+",
        default=None,
        help="LibriSpeech subsets to process (e.g., train-clean-360). Required when dataset-type=librispeech.",
    )
    parser.add_argument(
        "--pairs-file",
        "--pairs_file",
        type=Path,
        default=None,
        help="Optional pairs.txt to use for instructs2s dataset.",
    )
    parser.add_argument(
        "--output",
        type=Path,
        required=True,
        help="Directory to store the precomputed token tensors.",
    )
    parser.add_argument(
        "--tokenizer",
        type=str,
        default="Emova-ollm/emova_speech_tokenizer_hf",
        help="Name or path of the EMOVA speech tokenizer checkpoint to use.",
    )
    parser.add_argument(
        "--overwrite",
        action="store_true",
        help="Recompute and overwrite existing token files.",
    )
    parser.add_argument(
        "--device",
        type=str,
        default="cuda" if torch.cuda.is_available() else "cpu",
        help="Device for running the tokenizer encoder.",
    )
    parser.add_argument(
        "--devices",
        nargs="+",
        default=None,
        help="Optional list of devices per worker (e.g., cuda:0 cuda:1 ...). Overrides --device/--num-workers.",
    )
    parser.add_argument(
        "--num-workers",
        type=int,
        default=1,
        help="Number of parallel workers (ignored if --devices is provided).",
    )
    args = parser.parse_args()

    if args.index is not None:
        args.index = args.index.expanduser().resolve()
        if not args.index.exists():
            parser.error(f"Index file not found: {args.index}")

    if args.dataset_type == "video-speech" and args.index is None:
        parser.error("--index is required when dataset-type=video-speech.")
    if args.dataset_type == "librispeech" and not args.librispeech_subsets:
        parser.error("--librispeech-subsets must be provided when dataset-type=librispeech.")

    args.audio_root = args.audio_root.expanduser().resolve()
    args.output = args.output.expanduser().resolve()
    if args.pairs_file is not None:
        args.pairs_file = Path(args.pairs_file).expanduser().resolve()
        if not args.pairs_file.exists():
            parser.error(f"pairs-file not found: {args.pairs_file}")

    args.output.mkdir(parents=True, exist_ok=True)

    audio_paths = gather_audio_paths(args)
    if not audio_paths:
        print("No audio files found. Nothing to encode.")
        return

    if args.devices:
        worker_devices = args.devices
    else:
        worker_devices = [args.device] * max(1, args.num_workers)

    if len(worker_devices) == 1:
        device = torch.device(worker_devices[0])
        tokenizer = EMOVASpeechTokenizer.from_pretrained(args.tokenizer).to(device)
        tokenizer.eval()

        total = 0
        skipped = 0
        failed_paths: List[Path] = []
        for audio_path in tqdm(audio_paths, desc="Encoding clips"):
            out_path = token_output_path(args.output, audio_path)
            if out_path.exists() and not args.overwrite:
                skipped += 1
                continue

            out_path.parent.mkdir(parents=True, exist_ok=True)
            try:
                tokens = encode_audio(tokenizer, audio_path)
            except Exception as exc:
                tqdm.write(f"[WARN] Failed to encode {audio_path}: {exc}")
                failed_paths.append(audio_path)
                continue

            tmp_path = out_path.with_suffix(out_path.suffix + ".tmp")
            torch.save(tokens, tmp_path)
            os.replace(tmp_path, out_path)
            total += 1
        if failed_paths:
            failed_log = args.output / "failed_paths.log"
            with failed_log.open("a") as fh:
                for path in failed_paths:
                    fh.write(f"{path}\n")
            print(f"Wrote {len(failed_paths)} failed paths to {failed_log}")
        print(f"Done. Encoded {total} clips. Skipped {skipped} existing entries.")
        return

    shards = split_into_shards(audio_paths, len(worker_devices))
    from multiprocessing import get_context

    ctx = get_context("spawn")
    futures = []
    with ctx.Pool(len(worker_devices)) as pool:
        for shard_id, (device_str, shard_paths) in enumerate(zip(worker_devices, shards)):
            futures.append(
                pool.apply_async(
                    process_shard,
                    (
                        shard_id,
                        shard_paths,
                        device_str,
                        args.tokenizer,
                        args.output,
                        args.overwrite,
                        args.dataset_type,
                    ),
                )
            )
        pool.close()
        pool.join()

    total = 0
    skipped = 0
    failed_paths: List[Path] = []
    for fut in futures:
        shard_total, shard_skipped, shard_failed = fut.get()
        total += shard_total
        skipped += shard_skipped
        failed_paths.extend(shard_failed)

    if failed_paths:
        failed_log = args.output / "failed_paths.log"
        with failed_log.open("a") as fh:
            for path in failed_paths:
                fh.write(f"{path}\n")
        print(f"Wrote {len(failed_paths)} failed paths to {failed_log}")

    print(f"Done. Encoded {total} clips. Skipped {skipped} existing entries.")


if __name__ == "__main__":
    main()