Add causal person fusion and MuVAP stage-3 baselines

third_party/MuVAP/data/alignment.py

@@ -0,0 +1,114 @@
+"""Time-based alignment utilities for audio, video frames, and frame labels.
+
+The model runs on a fixed 25 Hz timeline. Source videos remain at their native
+frame rate; frames and labels are sampled onto the model timeline at load time.
+"""
+
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Sequence
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+
+@dataclass(frozen=True)
+class AlignmentConfig:
+    model_fps: float = 25.0
+    audio_sample_rate: int = 16_000
+
+    def __post_init__(self) -> None:
+        if self.model_fps <= 0:
+            raise ValueError("model_fps must be positive")
+        if self.audio_sample_rate <= 0:
+            raise ValueError("audio_sample_rate must be positive")
+
+
+@dataclass(frozen=True)
+class AVSample:
+    sample_id: str
+    audio_path: Path
+    visual_path: Path
+    source_fps: float
+    labels: Sequence[int]
+    start_sec: float = 0.0
+
+    def __post_init__(self) -> None:
+        if self.source_fps <= 0:
+            raise ValueError(f"{self.sample_id}: source_fps must be positive")
+        if self.start_sec < 0:
+            raise ValueError(f"{self.sample_id}: start_sec cannot be negative")
+
+
+def model_frame_count(source_frames: int, source_fps: float, model_fps: float) -> int:
+    """Return the number of complete model intervals in a source-frame span."""
+    if source_frames < 0 or source_fps <= 0 or model_fps <= 0:
+        raise ValueError("frame counts and frame rates must be valid")
+    return int(np.floor((source_frames / source_fps) * model_fps + 1e-9))
+
+
+def source_indices(
+    source_frames: int, source_fps: float, model_fps: float
+) -> np.ndarray:
+    """Map model-frame centres to the source intervals that contain them."""
+    count = model_frame_count(source_frames, source_fps, model_fps)
+    model_centres = (np.arange(count, dtype=np.float64) + 0.5) / model_fps
+    indices = np.floor(model_centres * source_fps).astype(np.int64)
+    return np.clip(indices, 0, max(source_frames - 1, 0))
+
+
+def model_indices(
+    source_frames: int, source_fps: float, model_fps: float
+) -> np.ndarray:
+    """Map source-frame centres to the model intervals that contain them.
+
+    The inverse direction of `source_indices`, used to lift model-rate
+    predictions back onto the native annotation timeline for evaluation.
+    Always returns exactly `source_frames` indices, so a caller can score one
+    prediction per native frame regardless of the source frame rate.
+    """
+    if source_frames < 0:
+        raise ValueError("source_frames cannot be negative")
+    count = model_frame_count(source_frames, source_fps, model_fps)
+    source_centres = (np.arange(source_frames, dtype=np.float64) + 0.5) / source_fps
+    indices = np.floor(source_centres * model_fps).astype(np.int64)
+    return np.clip(indices, 0, max(count - 1, 0))
+
+
+def align_visual_and_labels(
+    visual: np.ndarray,
+    labels: Sequence[int],
+    source_fps: float,
+    model_fps: float,
+    strict: bool = True,
+) -> tuple[np.ndarray, np.ndarray]:
+    """Sample faces and their labels with one shared timestamp mapping.
+
+    A length disagreement means the face track and the annotation disagree
+    about how long the clip is, which silently shifts every later label.
+    `strict` reports that instead of trimming to the shorter of the two.
+    """
+    labels_array = np.asarray(labels, dtype=np.float32)
+    if strict and len(visual) != len(labels_array):
+        raise ValueError(
+            f"visual has {len(visual)} frames but labels have {len(labels_array)}"
+        )
+    usable_frames = min(len(visual), len(labels_array))
+    if usable_frames == 0:
+        raise ValueError("visual and labels must contain at least one frame")
+
+    indices = source_indices(usable_frames, source_fps, model_fps)
+    return visual[:usable_frames][indices], labels_array[:usable_frames][indices]
+
+
+def fit_audio_to_frames(
+    waveform: torch.Tensor,
+    frame_count: int,
+    sample_rate: int,
+    model_fps: float,
+) -> torch.Tensor:
+    """Pad/crop audio to the exact duration represented by model frames."""
+    target_samples = round(frame_count * sample_rate / model_fps)
+    waveform = waveform[..., :target_samples]
+    return F.pad(waveform, (0, max(0, target_samples - waveform.shape[-1])))

third_party/MuVAP/data/augment.py

@@ -0,0 +1,105 @@
+"""Train-time augmentation for packed ASD samples.
+
+Every transform here is label-preserving except `substitute_audio`, which
+rewrites the activity labels to silence on purpose. Visual parameters are drawn
+once per chunk and applied to every frame, because the model is temporal and a
+per-frame jitter would destroy the motion it reads.
+"""
+
+import math
+import random
+from dataclasses import dataclass, field
+
+import cv2
+import numpy as np
+
+
+VISUAL_MODES = ("orig", "flip", "crop", "rotate")
+
+
+@dataclass(frozen=True)
+class AugmentConfig:
+    """How aggressively to augment; all probabilities are per chunk."""
+
+    audio_mix_prob: float = 0.5
+    audio_snr_db: tuple[float, float] = (-5.0, 5.0)
+    visual: bool = True
+    crop_scale: tuple[float, float] = (0.7, 1.0)
+    rotate_degrees: float = 15.0
+    # Ablation only: replace the audio outright and call the sample silent.
+    # Stronger than mixing for teaching rejection of off-screen speech, but it
+    # also lets the model key on audio plausibility instead of AV synchrony.
+    substitute_audio_prob: float = 0.0
+
+    @classmethod
+    def from_cfg(cls, cfg: dict | None) -> "AugmentConfig | None":
+        if not cfg:
+            return None
+        snr = cfg.get("audio_snr_db", (-5.0, 5.0))
+        return cls(
+            audio_mix_prob=float(cfg.get("audio_mix_prob", 0.5)),
+            audio_snr_db=(float(snr[0]), float(snr[1])),
+            visual=bool(cfg.get("visual", True)),
+            crop_scale=tuple(cfg.get("crop_scale", (0.7, 1.0))),
+            rotate_degrees=float(cfg.get("rotate_degrees", 15.0)),
+            substitute_audio_prob=float(cfg.get("substitute_audio_prob", 0.0)),
+        )
+
+
+def _power_db(samples: np.ndarray) -> float:
+    return 10.0 * math.log10(float(np.mean(samples.astype(np.float64) ** 2)) + 1e-4)
+
+
+def fit_length(donor: np.ndarray, samples: int) -> np.ndarray:
+    """Wrap or truncate donor audio to the target length."""
+    if len(donor) == 0:
+        return np.zeros(samples, dtype=np.float32)
+    if len(donor) < samples:
+        donor = np.tile(donor, math.ceil(samples / len(donor)))
+    return donor[:samples]
+
+
+def mix_audio(
+    clean: np.ndarray, donor: np.ndarray, snr_db: float
+) -> np.ndarray:
+    """Add donor audio at a target signal-to-noise ratio.
+
+    The label is deliberately unchanged. A negative face with foreign speech
+    added must stay negative, which is what teaches the model to reject speech
+    it cannot see; a positive face stays positive when a competing voice
+    arrives.
+    """
+    donor = fit_length(donor, len(clean))
+    scale = 10.0 ** ((_power_db(clean) - _power_db(donor) - snr_db) / 20.0)
+    return np.clip(clean + donor * scale, -1.0, 1.0)
+
+
+def augment_visual(
+    faces: np.ndarray, config: AugmentConfig, rng: random.Random
+) -> np.ndarray:
+    """Apply one geometric transform to every frame of a chunk."""
+    mode = rng.choice(VISUAL_MODES)
+    if mode == "orig":
+        return faces
+    size = faces.shape[-1]
+    if mode == "flip":
+        return np.ascontiguousarray(faces[:, :, ::-1])
+    if mode == "crop":
+        low, high = config.crop_scale
+        side = max(8, int(size * rng.uniform(low, high)))
+        if side >= size:
+            return faces
+        top = rng.randint(0, size - side)
+        left = rng.randint(0, size - side)
+        window = faces[:, top : top + side, left : left + side]
+        return np.stack(
+            [
+                cv2.resize(frame, (size, size), interpolation=cv2.INTER_LINEAR)
+                for frame in window
+            ]
+        )
+    angle = rng.uniform(-config.rotate_degrees, config.rotate_degrees)
+    matrix = cv2.getRotationMatrix2D((size / 2, size / 2), angle, 1.0)
+    return np.stack(
+        [cv2.warpAffine(frame, matrix, (size, size)) for frame in faces]
+    )

third_party/MuVAP/data/dataloaders/fisher.py

@@ -0,0 +1,374 @@
+"""Fisher datasets for VAP training and turn-taking evaluation.
+
+Segments are stored as stereo audio beside the raw two-speaker VAD, not as
+pre-encoded class indices. Labels are projected in the dataloader, so changing
+`projection_window.mode` between the role-based and speaker-based setups is a
+config edit rather than a re-run of preprocessing.
+
+Audio comes from one of two sources, selected by `vap.source`:
+
+* `npy` reads the pre-cut segments under `seg/` and `tune/`.
+* `raw` seeks the same span in the 8 kHz mu-law source recording and resamples
+  it. The source is 97 GB where the pre-cut segments are 1.12 TB, and segment
+  length, stride, and sample rate stop being baked into the files, so they can
+  change without re-running preprocessing.
+
+Both sources are required to yield the same window, and a test pins them to
+each other.
+"""
+
+from functools import lru_cache
+from glob import glob
+from pathlib import Path
+
+import numpy as np
+import soundfile as sf
+import torch
+import torchaudio.functional as AF
+from torch.utils.data import Dataset
+
+from data.media import AudioSpec, load_waveform
+from preprocess.vap.utils import EVENT_SEG_LEN_SEC, event_crop
+
+SOURCES = ("npy", "raw")
+VAD_HZ = 25
+
+
+def vad_dirname(frame_hz):
+    """Where the VAD for this rate lives: `vad` at 25 Hz, `vad12_5` at 12.5, ...
+
+    `00_prep_fisher.py` writes the 25 Hz labels to `seg/vad`; every other rate is
+    written beside it by `03_prep_vad_rate.py`, so the original tree is never
+    overwritten and two rates can coexist.
+    """
+    if float(frame_hz) == float(VAD_HZ):
+        return "vad"
+    return "vad" + f"{float(frame_hz):g}".replace(".", "_")
+TARGET_SAMPLE_RATE = 16_000
+# Resampling a slice is not the same as slicing a resampled recording: the
+# filter needs signal beyond both edges. Read this much extra and trim it off.
+RESAMPLE_PAD_SEC = 0.1
+
+
+@lru_cache(maxsize=1024)
+def _source_rate(path):
+    return sf.info(path).samplerate
+
+
+@lru_cache(maxsize=256)
+def _speech_regions(root, part, group, conversation):
+    """Word-level speech intervals per speaker, as `((start, end), ...)` seconds.
+
+    The event clips ship as audio only, but an anchored model needs the voice
+    activity behind them. It comes from the same word annotations preprocessing
+    rasterized to build the training VAD, so the two cannot disagree.
+    """
+    pattern = Path(root) / part / "regions_vap" / group / f"{conversation}_*_words.txt"
+    speakers = []
+    for path in sorted(glob(str(pattern))):
+        with open(path) as handle:
+            speakers.append(
+                tuple(
+                    (float(line.split()[0]), float(line.split()[1]))
+                    for line in handle
+                    if line.strip()
+                )
+            )
+    if len(speakers) != 2:
+        raise FileNotFoundError(
+            f"expected two *_words.txt for {conversation} under {pattern.parent}, "
+            f"found {len(speakers)}"
+        )
+    return tuple(speakers)
+
+
+def window_vad(root, part, group, conversation, start_sec, frames, frame_hz=VAD_HZ):
+    """`[2, frames]` voice activity for the window starting at `start_sec`."""
+    vad = torch.zeros(2, frames)
+    for speaker, regions in enumerate(_speech_regions(root, part, group, conversation)):
+        for begin, end in regions:
+            first = int(round((begin - start_sec) * frame_hz))
+            last = int(round((end - start_sec) * frame_hz))
+            if last <= 0 or first >= frames:
+                continue
+            vad[speaker, max(0, first) : min(frames, last)] = 1.0
+    return vad
+
+
+def read_source_window(path, start_sec, duration_sec, channels):
+    """Read `duration_sec` from a source recording, resampled to 16 kHz.
+
+    Fisher ships as 8 kHz mu-law NIST Sphere, which `soundfile` reads and
+    `torchaudio` (2.11) no longer does.
+    """
+    path = str(path)
+    rate = _source_rate(path)
+    pad = int(round(RESAMPLE_PAD_SEC * rate))
+    start = int(round(start_sec * rate))
+    frames = int(round(duration_sec * rate))
+
+    lead = min(pad, max(0, start))
+    with sf.SoundFile(path) as handle:
+        handle.seek(start - lead)
+        block = handle.read(frames + lead + pad, dtype="float32", always_2d=True).T
+
+    waveform = torch.from_numpy(np.ascontiguousarray(block))
+    if channels == 1:
+        waveform = waveform.mean(dim=0, keepdim=True)
+    elif waveform.shape[0] != channels:
+        raise ValueError(
+            f"{path} holds {waveform.shape[0]} channel(s); {channels} requested"
+        )
+
+    waveform = AF.resample(waveform, rate, TARGET_SAMPLE_RATE)
+    scale = TARGET_SAMPLE_RATE / rate
+    offset = int(round(lead * scale))
+    target = int(round(duration_sec * TARGET_SAMPLE_RATE))
+    waveform = waveform[:, offset : offset + target]
+    if waveform.shape[1] < target:  # the recording ended inside the window
+        waveform = torch.nn.functional.pad(waveform, (0, target - waveform.shape[1]))
+    # Clip after resampling, exactly where the `npy` path clips, so a checkpoint
+    # trained on one source evaluates identically on the other. The resampling
+    # filter overshoots at transients; mu-law input never does.
+    return waveform.clamp(-1.0, 1.0).contiguous()
+
+
+def _check_source(source):
+    if source not in SOURCES:
+        raise ValueError(f"source must be one of {SOURCES}, got {source!r}")
+    return source
+
+
+def _audio_spec(channels):
+    """Fisher segment spec.
+
+    Segments written by earlier revisions of `00_prep_fisher.py` stored float
+    PCM without clipping, and about one in eight overshoots [-1, 1] by up to
+    2.5 dB. Clipping them is what any encoder would do and keeps those shards
+    usable; freshly prepared segments are int16 and never reach this path.
+    """
+    if channels not in (1, 2):
+        raise ValueError(f"channels must be 1 or 2, got {channels}")
+    return AudioSpec(sample_rate=16_000, mono=channels == 1, clamp_to_peak=True)
+
+
+class Fisher(Dataset):
+    """Fixed-length conversation segments with labels projected on the fly."""
+
+    def __init__(
+        self,
+        fisher_path,
+        split_path,
+        projection,
+        channels=1,
+        frame_hz=25,
+        sample_rate=16_000,
+        swap_channels=False,
+        source="npy",
+        with_vad=False,
+    ):
+        self.root = Path(fisher_path)
+        self.projection = projection
+        self.channels = channels
+        self.source = _check_source(source)
+        self.spec = _audio_spec(channels)
+        # `frame_hz` can be fractional (12.5 Hz on Mimi's own grid), so this is
+        # a rounded int rather than floor division - a float here reaches a
+        # slice index and fails deep inside a dataloader worker.
+        self.samples_per_frame = int(round(sample_rate / frame_hz))
+        self.with_vad = with_vad
+        self.vad_dir = vad_dirname(frame_hz)
+        # Swapping is label-preserving wherever speaker order is arbitrary: with
+        # two channels, and with one channel whose speakers are named by the VAD.
+        self.swap_channels = swap_channels and (channels == 2 or with_vad)
+        if projection.frame_hz != frame_hz:
+            raise ValueError(
+                f"projection is defined at {projection.frame_hz} Hz but the dataset runs at {frame_hz} Hz"
+            )
+        self.samples = []
+        self._build(split_path)
+
+    def _paths(self, line):
+        # Paths are resolved without touching the filesystem: a split can list
+        # hundreds of thousands of segments on a network mount, and stat-ing
+        # each one at construction costs minutes before the first batch.
+        part, group, conversation, segment, start, end, *_ = line.split()
+        base = self.root / part / "seg"
+        audio = (
+            base / "audio" / group / conversation / f"{segment}.npy"
+            if self.source == "npy"
+            else self.root / part / "audio" / group / f"{conversation}.wav"
+        )
+        return (
+            audio,
+            base / self.vad_dir / group / conversation / f"{segment}.npy",
+            float(start),
+            float(end) - float(start),
+        )
+
+    def _build(self, split_path):
+        with open(split_path) as file:
+            self.samples = [self._paths(line) for line in file if line.strip()]
+        if not self.samples:
+            raise ValueError(f"{split_path} lists no segments")
+        audio_path, vad_path, _, _ = self.samples[0]
+        for path in (audio_path, vad_path):
+            if not path.exists():
+                raise FileNotFoundError(
+                    f"missing preprocessed Fisher sample {path}; "
+                    "run preprocess/vap/00_prep_fisher.py first"
+                )
+
+    def __len__(self):
+        return len(self.samples)
+
+    def _labels(self, vad, num_frames):
+        """Project stored VAD, then drop the context kept only for the bins."""
+        context = (vad.shape[-1] - num_frames) // 2
+        if context < 0:
+            raise ValueError(
+                f"VAD has {vad.shape[-1]} frames, fewer than the {num_frames} audio frames"
+            )
+        needed = max(self.projection.hist_frames, self.projection.fut_frames)
+        if context < needed:
+            raise ValueError(
+                f"projection needs {needed} context frames on each side but the segment stores {context}; "
+                "re-run preprocess/vap/00_prep_fisher.py with a wider --context-sec"
+            )
+        labels = self.projection.get_labels(vad.unsqueeze(0))[0]
+        return labels[context : context + num_frames]
+
+    def __getitem__(self, index):
+        audio_path, vad_path, start_sec, duration_sec = self.samples[index]
+        if self.source == "raw":
+            audio = read_source_window(
+                audio_path, start_sec, duration_sec, self.channels
+            )
+        else:
+            audio, _ = load_waveform(audio_path, self.spec, npy_sample_rate=16_000)
+        vad = torch.from_numpy(np.load(vad_path, allow_pickle=False)).float()
+
+        if self.channels == 2 and audio.shape[0] != 2:
+            raise ValueError(
+                f"{audio_path} holds {audio.shape[0]} channel(s); stereo VAP needs 2"
+            )
+        if self.swap_channels and torch.rand(()) < 0.5:
+            # Which speaker sits on which channel is arbitrary, so swapping both
+            # the audio and its VAD together is a label-preserving augmentation.
+            audio = audio.flip(0)
+            vad = vad.flip(0)
+
+        num_frames = int(audio.shape[-1] // self.samples_per_frame)
+        labels = self._labels(vad, num_frames)
+        if not self.with_vad:
+            return audio, labels
+        # The same frames the labels cover, so the anchor lines up with them.
+        context = (vad.shape[-1] - num_frames) // 2
+        return audio, labels, vad[:, context : context + num_frames]
+
+
+class FisherEvent(Dataset):
+    """Ten-second clips ending at an annotated hold or shift decision point."""
+
+    FIELDS = ("part", "group", "conversation", "start", "end", "prev_spk", "next_spk",
+              "timing", "gap_type", "label")
+
+    def __init__(self, fisher_path, split_paths, channels=1, source="npy",
+                 with_vad=False, frame_hz=VAD_HZ):
+        root = self.root = Path(fisher_path)
+        self.channels = channels
+        self.source = _check_source(source)
+        self.spec = _audio_spec(channels)
+        self.audio_dir = "audio" if channels == 1 else "audio_stereo"
+        self.with_vad = with_vad
+        self.frame_hz = frame_hz
+        self.samples = []
+        if isinstance(split_paths, (str, Path)):
+            split_paths = [split_paths]
+
+        for split_path in split_paths:
+            with open(split_path) as file:
+                for line in file:
+                    if not line.strip():
+                        continue
+                    fields = dict(zip(self.FIELDS, line.split()))
+                    # The window is defined the same way for both sources; `raw`
+                    # seeks it, `npy` reads a clip already cut to it. An anchored
+                    # model needs its absolute start either way, to line the
+                    # voice activity up with the audio.
+                    crop = event_crop(
+                        fields["timing"],
+                        float(fields["start"]),
+                        float(fields["end"]),
+                    )
+                    if crop is None and (self.source == "raw" or self.with_vad):
+                        continue
+                    if self.source == "raw":
+                        audio_path = (
+                            root
+                            / fields["part"]
+                            / "audio"
+                            / fields["group"]
+                            / f"{fields['conversation']}.wav"
+                        )
+                        crop_start = crop[0]
+                    else:
+                        audio_path = (
+                            root
+                            / fields["part"]
+                            / "tune"
+                            / self.audio_dir
+                            / fields["group"]
+                            / fields["conversation"]
+                            / f"{fields['start']}.npy"
+                        )
+                        crop_start = 0.0
+                    self.samples.append(
+                        (
+                            audio_path,
+                            1.0 if fields["label"] == "SHIFT" else 0.0,
+                            fields["conversation"],
+                            fields["timing"],
+                            fields["gap_type"],
+                            int(fields["prev_spk"]),
+                            crop_start,
+                            (fields["part"], fields["group"], crop[0] if crop else None),
+                        )
+                    )
+
+    def __len__(self):
+        return len(self.samples)
+
+    def __getitem__(self, index):
+        (
+            audio_path,
+            label,
+            conversation,
+            timing,
+            gap_type,
+            prev_spk,
+            crop_start,
+            window,
+        ) = self.samples[index]
+
+        if self.source == "raw":
+            audio = read_source_window(
+                audio_path, crop_start, EVENT_SEG_LEN_SEC, self.channels
+            )
+        else:
+            audio, _ = load_waveform(audio_path, self.spec, npy_sample_rate=16_000)
+            if self.channels == 2 and audio.shape[0] != 2:
+                raise ValueError(
+                    f"{audio_path} holds {audio.shape[0]} channel(s); stereo evaluation needs 2. "
+                    "Re-run preprocess/vap/02_prep_event.py with --channels 2, "
+                    "or set vap.source: raw to read the source recording instead"
+                )
+        if not self.with_vad:
+            return audio, label, conversation, timing, gap_type, prev_spk
+        part, group, window_start = window
+        vad = window_vad(
+            str(self.root), part, group, conversation, window_start,
+            int(round(audio.shape[-1] / TARGET_SAMPLE_RATE * self.frame_hz)),
+            frame_hz=self.frame_hz,
+        )
+        return audio, label, conversation, timing, gap_type, prev_spk, vad

third_party/MuVAP/data/dataloaders/packed.py

@@ -0,0 +1,410 @@
+"""Memory-mapped reader that turns native-rate packs into model-ready batches.
+
+Shards hold decoded source media: faces and labels at the native frame rate,
+unnormalized 16 kHz audio. Everything MuVAP+ specific happens here - sampling
+onto the 25 Hz model timeline, loudness normalization, chunking, augmentation -
+so the same shards can serve a model with different conventions.
+"""
+
+import json
+import math
+import random
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Sequence
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torch.utils.data import Dataset, Sampler
+
+from data.alignment import model_frame_count, source_indices
+from data.augment import AugmentConfig, augment_visual, mix_audio
+from data.media import WaveformStats, apply_normalization
+from preprocess.asd.schema import FORMAT_NAME, MODEL_FPS, SAMPLE_RATE, SUPPORTED_VERSIONS
+
+
+def _clip_id(sample_id: str) -> str:
+    """Strip the entity suffix so same-recording samples share a key.
+
+    AVA and WASD both name tracks `<clip>_<start>_<end>:<entity>`, where every
+    entity of a clip may share one audio recording.
+    """
+    return sample_id.rsplit(":", 1)[0]
+
+
+@dataclass(frozen=True)
+class _Entry:
+    root: Path
+    shard: str
+    metadata: dict
+
+
+class PackedASDDataset(Dataset):
+    """Random-access dataset backed by OS-page-cached NumPy memmaps."""
+
+    def __init__(
+        self,
+        roots: str | Path | Sequence[str | Path],
+        projection_window=None,
+        window: int | None = None,
+        context: int = 0,
+        jitter: bool = False,
+        augment: AugmentConfig | None = None,
+        model_fps: float = MODEL_FPS,
+        normalize: bool = True,
+    ):
+        """Read packed shards, optionally as fixed-length training chunks.
+
+        `window` splits every track longer than it into near-equal chunks, so
+        one epoch still covers every frame while batches stay uniform enough to
+        pack tightly.
+
+        `context` frames are read *before* each chunk and marked invalid in the
+        returned mask. The model is causal, so a frame near the start of a chunk
+        would otherwise be asked to predict from history the encoder was never
+        shown. Feeding that history as unscored context costs a little compute
+        and no storage, where overlapping the chunks would duplicate the data
+        and still train on the truncated copies. A track's real first frames get
+        no prefix: there the missing history is genuine, and the model has to
+        cope with it at inference too.
+
+        `jitter` slides the chunk grid within its track on every access, which
+        varies the boundaries across epochs at no extra cost. All spans are
+        counted in model frames, not source frames.
+        """
+        if isinstance(roots, (str, Path)):
+            roots = [roots]
+        self.projection_window = projection_window
+        self.window = window
+        self.context = max(0, context)
+        self.jitter = jitter
+        self.augment = augment
+        self.model_fps = model_fps
+        self.normalize = normalize
+        self.samples_per_frame = round(SAMPLE_RATE / model_fps)
+        self.entries: list[_Entry] = []
+        self._arrays: dict[
+            tuple[Path, str], tuple[np.ndarray, np.ndarray, np.ndarray]
+        ] = {}
+
+        for raw_root in roots:
+            root = Path(raw_root)
+            with (root / "dataset.json").open() as handle:
+                manifest = json.load(handle)
+            if (
+                manifest.get("format") != FORMAT_NAME
+                or manifest.get("version") not in SUPPORTED_VERSIONS
+            ):
+                raise ValueError(f"unsupported packed dataset format: {root}")
+            if manifest.get("timeline") != "native":
+                raise ValueError(f"pack is not on the native timeline: {root}")
+            if manifest.get("audio_sample_rate") != SAMPLE_RATE:
+                raise ValueError(f"unexpected audio sample rate: {root}")
+            if manifest.get("audio_normalization") is not None:
+                raise ValueError(f"pack has baked-in audio normalization: {root}")
+
+            for shard_dir in sorted(root.glob("shard-*")):
+                with (shard_dir / "index.json").open() as handle:
+                    for metadata in json.load(handle):
+                        self.entries.append(_Entry(root, shard_dir.name, metadata))
+
+        self.model_frames = np.array(
+            [
+                model_frame_count(
+                    entry.metadata["source_frames"],
+                    entry.metadata["source_fps"],
+                    model_fps,
+                )
+                for entry in self.entries
+            ],
+            dtype=np.int64,
+        )
+        self.chunks = self._build_chunks()
+        self._clips = [_clip_id(entry.metadata["sample_id"]) for entry in self.entries]
+
+    @property
+    def predict_window(self) -> int | None:
+        """Frames scored per chunk; the context prefix is read on top of this."""
+        if self.window is None:
+            return None
+        return max(1, self.window - self.context)
+
+    def _build_chunks(self) -> list[tuple[int, int, int]]:
+        """Split tracks into near-equal *scored* spans.
+
+        A chunk is identified by the span it predicts. Whatever context the
+        loader prepends is read on top and never scored, so the spans still tile
+        the track exactly and every frame is predicted once per epoch.
+        """
+        chunks: list[tuple[int, int, int]] = []
+        span = self.predict_window
+        for index in range(len(self.entries)):
+            frames = int(self.model_frames[index])
+            if frames == 0:
+                continue
+            if span is None or frames <= span:
+                chunks.append((index, 0, frames))
+                continue
+            edges = np.linspace(
+                0, frames, math.ceil(frames / span) + 1
+            ).round().astype(int)
+            for start, end in zip(edges[:-1], edges[1:]):
+                chunks.append((index, int(start), int(end - start)))
+        return chunks
+
+    def _fed_length(self, start: int, length: int) -> int:
+        return length + min(self.context, start)
+
+    def chunk_lengths(self) -> np.ndarray:
+        """Frames actually fed per chunk - what bounds memory and batch size."""
+        return np.fromiter(
+            (self._fed_length(start, length) for _, start, length in self.chunks),
+            dtype=np.int64,
+            count=len(self.chunks),
+        )
+
+    def __len__(self) -> int:
+        return len(self.chunks)
+
+    def _open(self, entry: _Entry):
+        key = (entry.root, entry.shard)
+        arrays = self._arrays.get(key)
+        if arrays is None:
+            shard_dir = entry.root / entry.shard
+            faces = np.load(shard_dir / "faces.npy", mmap_mode="r", allow_pickle=False)
+            audio = np.load(shard_dir / "audio.npy", mmap_mode="r", allow_pickle=False)
+            labels = np.load(shard_dir / "labels.npy", mmap_mode="r", allow_pickle=False)
+            arrays = faces, audio, labels
+            self._arrays[key] = arrays
+        return arrays
+
+    def _stats(self, entry: _Entry) -> WaveformStats:
+        return WaveformStats(
+            mean=entry.metadata["audio_mean"],
+            rms=entry.metadata["audio_rms"],
+            peak=entry.metadata["audio_peak"],
+        )
+
+    def _read_audio(self, entry: _Entry, start: int, length: int) -> torch.Tensor:
+        """Read one model-frame span and restore the recording's own gain."""
+        _, audio, _ = self._open(entry)
+        offset = entry.metadata["audio_offset"] + start * self.samples_per_frame
+        available = entry.metadata["audio_samples"] - start * self.samples_per_frame
+        wanted = length * self.samples_per_frame
+        block = np.array(audio[offset : offset + min(wanted, max(available, 0))], copy=True)
+        waveform = torch.from_numpy(block).float().div_(32768.0).unsqueeze(0)
+        if waveform.shape[-1] < wanted:
+            waveform = F.pad(waveform, (0, wanted - waveform.shape[-1]))
+        if self.normalize:
+            waveform = apply_normalization(waveform, self._stats(entry))
+        return waveform
+
+    def _donor_audio(self, clip: str, frames: int, rng) -> np.ndarray:
+        """Read audio from a different clip than `clip`.
+
+        Drawing from the whole dataset rather than the mini-batch matters here:
+        every face track in a WASD clip shares one clip-level recording, and the
+        batch sampler groups by length, which puts those tracks together. A
+        batch-local donor would often be the very same waveform.
+        """
+        for _ in range(8):
+            index = rng.randrange(len(self.entries))
+            if self._clips[index] == clip:
+                continue
+            entry = self.entries[index]
+            span = int(self.model_frames[index])
+            start = rng.randrange(max(span - frames, 0) + 1) if span > frames else 0
+            return self._read_audio(entry, start, min(frames, max(span, 1)))[0].numpy()
+        return np.zeros(frames * self.samples_per_frame, dtype=np.float32)
+
+    def __getitem__(self, index: int):
+        entry_index, predict_start, predict_length = self.chunks[index]
+        entry = self.entries[entry_index]
+        faces, _, labels = self._open(entry)
+        frames = int(self.model_frames[entry_index])
+        if self.jitter and predict_length < frames:
+            predict_start = random.randint(0, frames - predict_length)
+
+        # Read a history prefix where one exists; at a track's real start there
+        # is none, and that truncation is genuine rather than an artefact.
+        prefix = min(self.context, predict_start)
+        start = predict_start - prefix
+        length = prefix + predict_length
+        scored = torch.zeros(length, dtype=torch.bool)
+        scored[prefix:] = True
+
+        # One shared mapping puts faces and labels on the model timeline, which
+        # is what keeps a label attached to the face it describes.
+        rows = source_indices(
+            entry.metadata["source_frames"], entry.metadata["source_fps"], self.model_fps
+        )
+        span = rows[start : start + length]
+        face_block = np.asarray(faces[entry.metadata["face_offset"] + span])
+        # Labels are one byte a frame, so the whole track is cheap to read. The
+        # projection window needs 50 frames of history and 15 of future; taking
+        # them from the track rather than the chunk keeps a chunk boundary from
+        # fabricating silence in the targets.
+        track_labels = np.asarray(
+            labels[entry.metadata["label_offset"] + rows], dtype=np.uint8
+        ).astype(np.float32)
+        waveform = self._read_audio(entry, start, length)
+
+        if self.augment is not None:
+            # The module RNG is what Lightning seeds per worker, so augmentation
+            # is reproducible from the run seed.
+            rng = random
+            clip = self._clips[entry_index]
+            audio_block = waveform[0].numpy()
+            if rng.random() < self.augment.substitute_audio_prob:
+                audio_block = self._donor_audio(clip, length, rng)
+                track_labels = np.zeros_like(track_labels)
+            elif rng.random() < self.augment.audio_mix_prob:
+                audio_block = mix_audio(
+                    audio_block,
+                    self._donor_audio(clip, length, rng),
+                    rng.uniform(*self.augment.audio_snr_db),
+                )
+            waveform = torch.from_numpy(np.ascontiguousarray(audio_block)).unsqueeze(0)
+            if self.augment.visual:
+                face_block = augment_visual(face_block, self.augment, rng)
+
+        visual = torch.from_numpy(np.ascontiguousarray(face_block)).float()
+        track_vad = torch.from_numpy(track_labels)
+        vad = track_vad[start : start + length]
+        if self.projection_window is None:
+            return waveform, visual, vad, scored, entry.metadata
+        with torch.no_grad():
+            vap = self.projection_window.get_labels(track_vad.view(1, 1, frames))
+        return (
+            waveform,
+            visual,
+            vap[0, start : start + length],
+            vad,
+            scored,
+            entry.metadata,
+        )
+
+
+class ASDCollator:
+    """Bring a batch to one length, by cropping or by padding.
+
+    Training crops to the shortest member. That keeps the model free of a
+    padding mask, and it is affordable because chunking bounds the spread of
+    lengths and the batch sampler groups similar lengths together; measured on
+    AVA plus WASD the pair loses 0.01% of frames per epoch, against 64% for
+    random batches.
+
+    Evaluation cannot crop: the official export needs every frame of every
+    track, which is why it would otherwise run one track at a time. `pad=True`
+    right-pads to the longest member instead and marks the filler unscored.
+    Every stage of the model is causal - the CPC and Mimi frontends, the visual
+    TCN, and both transformer stacks - so appended frames cannot change any
+    earlier output, and the visual frontend's batch norm is on running
+    statistics in eval mode. A test pins both properties.
+    """
+
+    def __init__(
+        self, samples_per_frame: int = SAMPLE_RATE // int(MODEL_FPS), pad: bool = False
+    ):
+        self.samples_per_frame = samples_per_frame
+        self.pad = pad
+
+    @staticmethod
+    def _to_length(tensor: torch.Tensor, length: int, axis: int) -> torch.Tensor:
+        """Crop or right-pad one stream along `axis` (counted from the end)."""
+        current = tensor.shape[-axis]
+        if current == length:
+            return tensor
+        if current > length:
+            index = [slice(None)] * tensor.ndim
+            index[-axis] = slice(0, length)
+            return tensor[tuple(index)]
+        padding = [0, 0] * (axis - 1) + [0, length - current]
+        return F.pad(tensor, padding)
+
+    def __call__(self, batch):
+        lengths = [item[1].shape[0] for item in batch]
+        frames = max(lengths) if self.pad else min(lengths)
+        samples = frames * self.samples_per_frame
+
+        audio = torch.stack([self._to_length(item[0], samples, 1) for item in batch])
+        visual = torch.stack([self._to_length(item[1], frames, 3) for item in batch])
+        vap = torch.stack([self._to_length(item[2], frames, 2) for item in batch])
+        vad = torch.stack([self._to_length(item[3], frames, 1) for item in batch])
+        scored = torch.stack([self._to_length(item[4], frames, 1) for item in batch])
+        return audio, visual, vap, vad, scored, [item[5] for item in batch]
+
+
+class FrameBudgetBatchSampler(Sampler):
+    """Group similar-length chunks into batches of roughly constant frame count.
+
+    Holding `batch * frames` near a budget keeps three things stable that a
+    fixed batch size does not: activation memory, the `B*T` population the
+    visual frontend's batch norm sees, and the per-frame gradient noise of the
+    mean-reduced losses. The batch size is clamped because the contrastive term
+    is a `B`-way classification whose difficulty would otherwise swing with the
+    chunk length.
+    """
+
+    def __init__(
+        self,
+        lengths,
+        frame_budget: int,
+        min_batch: int = 4,
+        max_batch: int = 16,
+        shuffle: bool = True,
+        drop_last: bool = False,
+        seed: int = 0,
+        num_replicas: int = 1,
+        rank: int = 0,
+    ):
+        if min_batch < 1 or max_batch < min_batch:
+            raise ValueError("min_batch must be positive and not exceed max_batch")
+        self.lengths = np.asarray(lengths, dtype=np.int64)
+        self.frame_budget = frame_budget
+        self.min_batch = min_batch
+        self.max_batch = max_batch
+        self.shuffle = shuffle
+        self.drop_last = drop_last
+        self.seed = seed
+        self.num_replicas = num_replicas
+        self.rank = rank
+        self.epoch = 0
+        self._batches = self._build()
+
+    def set_epoch(self, epoch: int) -> None:
+        self.epoch = epoch
+        self._batches = self._build()
+
+    def _build(self) -> list[list[int]]:
+        order = np.argsort(-self.lengths, kind="stable")
+        batches = []
+        cursor = 0
+        while cursor < len(order):
+            longest = int(self.lengths[order[cursor]])
+            size = min(
+                max(self.frame_budget // max(longest, 1), self.min_batch), self.max_batch
+            )
+            batch = order[cursor : cursor + size]
+            cursor += size
+            if self.drop_last and len(batch) < self.min_batch:
+                continue
+            batches.append([int(index) for index in batch])
+
+        if self.shuffle:
+            rng = random.Random(self.seed + self.epoch)
+            rng.shuffle(batches)
+        if self.num_replicas > 1:
+            # Every rank must run the same number of steps or DDP deadlocks.
+            remainder = len(batches) % self.num_replicas
+            if remainder:
+                batches += batches[: self.num_replicas - remainder]
+            batches = batches[self.rank :: self.num_replicas]
+        return batches
+
+    def __iter__(self):
+        return iter(self._batches)
+
+    def __len__(self) -> int:
+        return len(self._batches)

third_party/MuVAP/data/media.py

@@ -0,0 +1,211 @@
+"""Canonical media conversion shared by every MuVAP+ dataset.
+
+Dataset adapters resolve paths and metadata only. Signal conversion belongs
+here so AVA, WASD, MSDWild, Fisher, and AVCC cannot silently diverge.
+"""
+
+from dataclasses import dataclass
+from pathlib import Path
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+import torchaudio
+import torchaudio.functional as AF
+
+
+@dataclass(frozen=True)
+class AudioSpec:
+    sample_rate: int = 16_000
+    mono: bool = True
+    peak_limit: float = 1.0 + 1e-4
+    # Out-of-range float PCM is a scaling bug in most corpora, so it raises by
+    # default. Set this only for a source known to store real audio that
+    # overshoots, where clipping is the same thing any encoder would do.
+    clamp_to_peak: bool = False
+
+
+@dataclass(frozen=True)
+class VisualSpec:
+    height: int = 112
+    width: int = 112
+
+
+CANONICAL_AUDIO = AudioSpec()
+CANONICAL_VISUAL = VisualSpec()
+ASD_TARGET_RMS_DBFS = -25.0
+ASD_MAX_GAIN_DB = 20.0
+ASD_PEAK_LIMIT = 0.99
+
+
+def _integer_audio_to_float(array: np.ndarray) -> np.ndarray:
+    info = np.iinfo(array.dtype)
+    scale = float(max(abs(info.min), abs(info.max)))
+    return array.astype(np.float32) / scale
+
+
+def _channels_first(array: np.ndarray) -> np.ndarray:
+    """Canonicalize common `[N]`, `[C,N]`, and legacy `[N,C]` layouts."""
+    if array.ndim == 1:
+        return array[None, :]
+    if array.ndim != 2:
+        raise ValueError(f"NumPy audio must have 1 or 2 dimensions, got {array.shape}")
+    if array.shape[0] <= 8:
+        return array
+    if array.shape[1] <= 8:
+        return array.T
+    raise ValueError(
+        f"cannot infer audio channel axis for shape {array.shape}; expected [C,N] or [N,C]"
+    )
+
+
+def canonicalize_waveform(
+    waveform: torch.Tensor, spec: AudioSpec = CANONICAL_AUDIO
+) -> torch.Tensor:
+    """Return `[channels, samples]` float32 PCM without changing its gain."""
+    if waveform.ndim == 1:
+        waveform = waveform.unsqueeze(0)
+    if waveform.ndim != 2:
+        raise ValueError(
+            f"audio must have 1 or 2 dimensions, got {tuple(waveform.shape)}"
+        )
+
+    waveform = waveform.to(torch.float32)
+    if not torch.isfinite(waveform).all():
+        raise ValueError("audio contains NaN or infinite values")
+    if waveform.numel() and waveform.abs().max().item() > spec.peak_limit:
+        if not spec.clamp_to_peak:
+            raise ValueError(
+                "floating-point audio exceeds [-1, 1]; convert integer PCM to normalized "
+                "float during preprocessing instead of applying per-clip normalization"
+            )
+        waveform = waveform.clamp(-1.0, 1.0)
+    if spec.mono and waveform.shape[0] > 1:
+        waveform = waveform.mean(dim=0, keepdim=True)
+    return waveform.contiguous()
+
+
+@dataclass(frozen=True)
+class WaveformStats:
+    """Level of a whole recording, measured after DC removal.
+
+    Packed shards hold unnormalized audio, so these travel in the index and let
+    a reader reproduce the exact per-recording gain from any slice of it. Gain
+    derived from a slice would drift with the content of that slice.
+    """
+
+    mean: float
+    rms: float
+    peak: float
+
+
+def waveform_stats(waveform: torch.Tensor) -> WaveformStats:
+    mean = waveform.mean().item() if waveform.numel() else 0.0
+    centred = waveform - mean
+    return WaveformStats(
+        mean=mean,
+        rms=centred.square().mean().sqrt().item() if centred.numel() else 0.0,
+        peak=centred.abs().max().item() if centred.numel() else 0.0,
+    )
+
+
+def normalization_gain(
+    stats: WaveformStats,
+    target_rms_dbfs: float = ASD_TARGET_RMS_DBFS,
+    max_gain_db: float = ASD_MAX_GAIN_DB,
+    peak_limit: float = ASD_PEAK_LIMIT,
+) -> float:
+    if stats.rms <= 1e-8 or stats.peak <= 1e-8:
+        return 1.0
+    return min(
+        10 ** (target_rms_dbfs / 20) / stats.rms,
+        10 ** (max_gain_db / 20),
+        peak_limit / stats.peak,
+    )
+
+
+def apply_normalization(
+    waveform: torch.Tensor,
+    stats: WaveformStats,
+    target_rms_dbfs: float = ASD_TARGET_RMS_DBFS,
+    max_gain_db: float = ASD_MAX_GAIN_DB,
+    peak_limit: float = ASD_PEAK_LIMIT,
+) -> torch.Tensor:
+    """Centre and scale a slice using its whole recording's statistics."""
+    gain = normalization_gain(stats, target_rms_dbfs, max_gain_db, peak_limit)
+    return ((waveform - stats.mean) * gain).contiguous()
+
+
+def normalize_waveform(
+    waveform: torch.Tensor,
+    target_rms_dbfs: float = ASD_TARGET_RMS_DBFS,
+    max_gain_db: float = ASD_MAX_GAIN_DB,
+    peak_limit: float = ASD_PEAK_LIMIT,
+) -> tuple[torch.Tensor, float]:
+    stats = waveform_stats(waveform)
+    gain = normalization_gain(stats, target_rms_dbfs, max_gain_db, peak_limit)
+    return ((waveform - stats.mean) * gain).contiguous(), gain
+
+
+def load_waveform(
+    path: str | Path,
+    spec: AudioSpec = CANONICAL_AUDIO,
+    *,
+    npy_sample_rate: int | None = None,
+) -> tuple[torch.Tensor, int]:
+    """Load file or NumPy PCM into the canonical amplitude representation."""
+    path = Path(path)
+    if path.suffix.lower() == ".npy":
+        array = _channels_first(np.load(path, allow_pickle=False))
+        if np.issubdtype(array.dtype, np.integer):
+            array = _integer_audio_to_float(array)
+        waveform = torch.from_numpy(np.asarray(array))
+        source_rate = npy_sample_rate
+        if source_rate is None:
+            raise ValueError(f"sample rate is required for NumPy audio: {path}")
+    else:
+        # torchaudio returns normalized float PCM for integer WAV/FLAC sources.
+        waveform, source_rate = torchaudio.load(path, normalize=True)
+        # Lossy codecs such as AAC can decode slightly outside nominal PCM bounds.
+        waveform = waveform.clamp(-1.0, 1.0)
+
+    waveform = canonicalize_waveform(waveform, spec)
+    if source_rate != spec.sample_rate:
+        waveform = AF.resample(waveform, source_rate, spec.sample_rate)
+    return waveform.contiguous(), spec.sample_rate
+
+
+def slice_audio(
+    waveform: torch.Tensor,
+    sample_rate: int,
+    *,
+    start_time: float | None = None,
+    end_time: float | None = None,
+    target_samples: int | None = None,
+) -> torch.Tensor:
+    """Keep timeline origin: crop and zero-pad only at the right boundary."""
+    start = round((start_time or 0.0) * sample_rate)
+    end = round(end_time * sample_rate) if end_time is not None else waveform.shape[-1]
+    if start < 0 or end < start:
+        raise ValueError("audio time range is invalid")
+    waveform = waveform[..., start:end]
+
+    if target_samples is not None:
+        if target_samples < 0:
+            raise ValueError("target_samples cannot be negative")
+        waveform = waveform[..., :target_samples]
+        waveform = F.pad(waveform, (0, max(0, target_samples - waveform.shape[-1])))
+    return waveform.contiguous()
+
+
+def validate_visual(
+    array: np.ndarray, spec: VisualSpec = CANONICAL_VISUAL
+) -> np.ndarray:
+    """Validate the tensor contract expected by `VisualEncoder`."""
+    if array.ndim != 3 or tuple(array.shape[1:]) != (spec.height, spec.width):
+        raise ValueError(
+            f"visual must be [T, {spec.height}, {spec.width}], got {array.shape}"
+        )
+    if array.dtype != np.uint8:
+        raise ValueError(f"visual must be uint8 in [0, 255], got {array.dtype}")
+    return array

third_party/MuVAP/data/packed_asd_dm.py

@@ -0,0 +1,110 @@
+"""Lightning data module for canonical packed ASD datasets."""
+
+import torch.distributed as dist
+from lightning.pytorch import LightningDataModule
+from torch.utils.data import DataLoader
+
+from data.augment import AugmentConfig
+from data.dataloaders.packed import (
+    ASDCollator,
+    FrameBudgetBatchSampler,
+    PackedASDDataset,
+)
+
+
+class PackedASDDataModule(LightningDataModule):
+    def __init__(self, cfg: dict, projection_window):
+        super().__init__()
+        self.cfg = cfg
+        self.projection_window = projection_window
+        asd = cfg["asd"]
+        self.num_workers = asd["num_workers"]
+        self.train_window = asd.get("train_window", 250)
+        self.frame_budget = asd.get("frame_budget", 2048)
+        self.min_batch = asd.get("min_batch", 4)
+        self.max_batch = asd.get("max_batch", 16)
+        self.seed = asd.get("seed", 0)
+        # Validation is forward-only, so it affords a larger budget than training.
+        self.val_frame_budget = asd.get("val_frame_budget", 4 * self.frame_budget)
+        self.val_max_batch = asd.get("val_max_batch", 64)
+        self.augment = AugmentConfig.from_cfg(asd.get("augment"))
+
+    def _replicas(self):
+        if dist.is_available() and dist.is_initialized():
+            return dist.get_world_size(), dist.get_rank()
+        return 1, 0
+
+    def setup(self, stage: str | None = None):
+        packed = self.cfg["packed_asd"]
+        if stage in ("fit", None):
+            # The context prefix is exactly the projection window's history, so
+            # every scored frame can see the past its target depends on.
+            self.train_dataset = PackedASDDataset(
+                packed["train"],
+                projection_window=self.projection_window,
+                window=self.train_window,
+                context=self.projection_window.hist_frames,
+                jitter=True,
+                augment=self.augment,
+            )
+        if stage in ("fit", "validate", "test", None):
+            # Validation runs whole tracks so every native frame gets a score.
+            self.val_dataset = PackedASDDataset(
+                packed["val"], projection_window=self.projection_window
+            )
+
+    def train_dataloader(self):
+        replicas, rank = self._replicas()
+        sampler = FrameBudgetBatchSampler(
+            self.train_dataset.chunk_lengths(),
+            frame_budget=self.frame_budget,
+            min_batch=self.min_batch,
+            max_batch=self.max_batch,
+            shuffle=True,
+            drop_last=True,
+            seed=self.seed,
+            num_replicas=replicas,
+            rank=rank,
+        )
+        return DataLoader(
+            self.train_dataset,
+            batch_sampler=sampler,
+            num_workers=self.num_workers,
+            collate_fn=ASDCollator(),
+            pin_memory=True,
+            persistent_workers=self.num_workers > 0,
+            prefetch_factor=2 if self.num_workers > 0 else None,
+        )
+
+    def val_dataloader(self):
+        """Batch whole tracks by grouping similar lengths and padding the rest.
+
+        Evaluation needs every frame of every track, so the batch cannot be
+        cropped to its shortest member the way training's is. Grouping by
+        length first keeps the padding that replaces it small, which is what
+        makes batching worth doing at all.
+
+        Lightning's automatic sampler injection is off because the train loader
+        shards itself, so this shards too. The last rank is padded by repeating
+        batches; the task de-duplicates by entity before scoring.
+        """
+        replicas, rank = self._replicas()
+        sampler = FrameBudgetBatchSampler(
+            self.val_dataset.chunk_lengths(),
+            frame_budget=self.val_frame_budget,
+            min_batch=1,
+            max_batch=self.val_max_batch,
+            shuffle=False,
+            drop_last=False,
+            num_replicas=replicas,
+            rank=rank,
+        )
+        return DataLoader(
+            self.val_dataset,
+            batch_sampler=sampler,
+            num_workers=self.num_workers,
+            collate_fn=ASDCollator(pad=True),
+            pin_memory=True,
+            persistent_workers=self.num_workers > 0,
+            prefetch_factor=2 if self.num_workers > 0 else None,
+        )

third_party/MuVAP/data/utils.py

@@ -0,0 +1,64 @@
+import os
+from glob import glob
+
+import torch
+
+from data.media import AudioSpec, load_waveform, slice_audio
+
+
+def load_audio(
+    audio_path,
+    target_sr=16000,
+    mono=True,
+    start_time=None,
+    end_time=None,
+    target_samples=None,
+):
+    """Load canonical float PCM while preserving the clip's time origin.
+
+    Short clips are padded on the right and long clips are cropped on the
+    right. No peak, RMS, or dataset-specific volume normalization is applied.
+    """
+    spec = AudioSpec(sample_rate=target_sr, mono=mono)
+    waveform, sr = load_waveform(audio_path, spec)
+    waveform = slice_audio(
+        waveform,
+        sr,
+        start_time=start_time,
+        end_time=end_time,
+        target_samples=target_samples,
+    )
+
+    duration = waveform.shape[1] / sr
+
+    return waveform, duration
+
+
+def load_fisher_vad(vad_path, conv_id, nframes, frame_hz=25):
+
+    pattern = os.path.join(vad_path, f"{conv_id}_*_words.txt")
+    vad_files = sorted(glob(pattern))
+
+    vad_labels = []
+    for vad_file in vad_files:
+        labels = torch.zeros(nframes)
+        with open(vad_file, "r") as f:
+            for line in f:
+                parts = line.strip().split()
+                start = float(parts[0])
+                end = float(parts[1])
+                start_idx = int(round(start * frame_hz))
+                end_idx = int(round(end * frame_hz))
+                start_idx = max(0, min(nframes, start_idx))
+                end_idx = max(0, min(nframes, end_idx))
+                labels[start_idx:end_idx] = 1
+        vad_labels.append(labels)
+
+    return torch.stack(vad_labels, dim=0)
+
+
+def load_lit_state_dict(model_path):
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    checkpoint = torch.load(model_path, map_location=device)
+    state_dict = checkpoint.get("state_dict", checkpoint)
+    return {k.replace("model.", ""): v for k, v in state_dict.items()}

third_party/MuVAP/data/vap_dm.py

@@ -0,0 +1,100 @@
+from pathlib import Path
+
+from lightning.pytorch import LightningDataModule
+from torch.utils.data import DataLoader
+
+from data.dataloaders.fisher import Fisher, FisherEvent
+from models.vap import audio_channels, needs_vad
+
+
+class VAPDataModule(LightningDataModule):
+    """Fisher segments for fitting, annotated turn-taking events for testing."""
+
+    def __init__(self, cfg, projection):
+        super().__init__()
+        vap_cfg = cfg["vap"]
+        self.root = Path(cfg["fisher_path"])
+        self.projection = projection
+        self.batch_size = vap_cfg["batch_size"]
+        self.num_workers = vap_cfg["num_workers"]
+        self.channels = audio_channels(vap_cfg)
+        self.with_vad = needs_vad(vap_cfg)
+        self.swap_channels = bool(vap_cfg.get("swap_channels", False))
+        self.source = str(vap_cfg.get("source", "npy")).lower()
+
+    def _split(self, name):
+        path = self.root / f"{name}.txt"
+        if not path.exists():
+            raise FileNotFoundError(
+                f"missing Fisher split {path}; run preprocess/vap/00_prep_fisher.py first"
+            )
+        return path
+
+    def _event_splits(self, names):
+        paths = [self.root / f"{name}_events.txt" for name in names]
+        found = [path for path in paths if path.exists()]
+        if not found:
+            raise FileNotFoundError(
+                f"none of {[p.name for p in paths]} exist; run preprocess/vap/02_prep_event.py first"
+            )
+        return found
+
+    def setup(self, stage=None):
+        if stage in ("fit", None):
+            self.train_dataset = self._segments("train", swap=self.swap_channels)
+            self.val_dataset = self._segments("val")
+
+        if stage in ("test", None):
+            self.tune_dataset = FisherEvent(
+                self.root,
+                self._event_splits(["train", "val"]),
+                channels=self.channels,
+                source=self.source,
+                with_vad=self.with_vad,
+                frame_hz=self.projection.frame_hz,
+            )
+            self.test_dataset = FisherEvent(
+                self.root,
+                self._event_splits(["test"]),
+                channels=self.channels,
+                source=self.source,
+                with_vad=self.with_vad,
+                frame_hz=self.projection.frame_hz,
+            )
+
+    def _segments(self, name, swap=False):
+        return Fisher(
+            self.root,
+            self._split(name),
+            self.projection,
+            channels=self.channels,
+            frame_hz=self.projection.frame_hz,
+            swap_channels=swap,
+            source=self.source,
+            with_vad=self.with_vad,
+        )
+
+    def _loader(self, dataset, shuffle=False):
+        workers = self.num_workers
+        return DataLoader(
+            dataset,
+            batch_size=self.batch_size,
+            shuffle=shuffle,
+            num_workers=workers,
+            pin_memory=True,
+            persistent_workers=bool(workers),
+            prefetch_factor=2 if workers else None,
+            drop_last=shuffle,
+        )
+
+    def train_dataloader(self):
+        return self._loader(self.train_dataset, shuffle=True)
+
+    def val_dataloader(self):
+        return self._loader(self.val_dataset)
+
+    def tune_dataloader(self):
+        return self._loader(self.tune_dataset)
+
+    def test_dataloader(self):
+        return self._loader(self.test_dataset)