"""OpenVid prompt/video dataset with lazy teacher-feature caching.""" from __future__ import annotations from dataclasses import dataclass import os from pathlib import Path import time from typing import Any, Callable, Mapping, Sequence import numpy as np import torch from torch.utils.data import Dataset from utils.openvid_feature_cache import ( OPENVID_PYAV_DECODER_CONTRACT, OPENVID_VJEPA2_L15_FEATURE_CONTRACT, OpenVidFeatureCache, ) from utils.openvid_manifest import ( OpenVidManifestEntry, load_openvid_manifest, ) @dataclass(frozen=True) class DecodedOpenVidVideo: """Central contiguous RGB frame window in V-JEPA2 input range.""" video01: torch.Tensor frame_indices: torch.Tensor source_frame_count: int def _frames_to_video01(frames: Sequence[Any]) -> torch.Tensor: tensors: list[torch.Tensor] = [] shape: tuple[int, int, int] | None = None for frame in frames: array = frame.to_ndarray(format="rgb24") if not isinstance(array, np.ndarray): raise TypeError("PyAV frame.to_ndarray did not return ndarray") if array.ndim != 3 or int(array.shape[2]) != 3: raise ValueError( f"decoded RGB frame must be [H,W,3], got {array.shape}" ) if array.dtype != np.uint8: raise ValueError( f"decoded RGB frame must be uint8, got {array.dtype}" ) current_shape = tuple(int(value) for value in array.shape) if shape is None: shape = current_shape elif current_shape != shape: raise ValueError( "video changes resolution inside selected frame window: " f"{current_shape} != {shape}" ) tensor = torch.from_numpy(np.ascontiguousarray(array)).permute(2, 0, 1) tensors.append(tensor) if not tensors: raise ValueError("decoded frame window is empty") return torch.stack(tensors).to(dtype=torch.float32).div_(255.0) def _decode_window_and_count( av_module: Any, video_path: Path, *, start: int, num_frames: int, ) -> tuple[list[Any], int]: selected: list[Any] = [] decoded_count = 0 with av_module.open(str(video_path), mode="r") as container: streams = container.streams.video if len(streams) != 1: raise ValueError( f"expected exactly one video stream in {video_path}, " f"got {len(streams)}" ) stream = streams[0] for frame_index, frame in enumerate(container.decode(stream)): decoded_count = frame_index + 1 if start <= frame_index < start + num_frames: selected.append(frame) return selected, decoded_count def decode_openvid_center_contiguous_pyav( video_path: str | os.PathLike, *, decoder_contract: Mapping[str, Any] | None = None, ) -> DecodedOpenVidVideo: """Decode the exact central contiguous frame window using PyAV order.""" contract = dict( decoder_contract if decoder_contract is not None else OPENVID_PYAV_DECODER_CONTRACT ) if contract.get("backend") != "pyav": raise ValueError("decoder contract backend must be pyav") if contract.get("selection") != "center_contiguous": raise ValueError("decoder contract selection must be center_contiguous") if contract.get("center_tie_break") != "floor": raise ValueError("decoder contract center_tie_break must be floor") num_frames = int(contract.get("num_frames", 0)) if num_frames <= 0: raise ValueError("decoder contract num_frames must be positive") try: import av # Imported lazily so complete cache hits need no decoder. except ImportError as exc: # pragma: no cover - exercised without PyAV raise ImportError( "PyAV is required for an OpenVid teacher cache miss; install " "the pinned av==13.1.0 dependency" ) from exc expected_version = str(contract.get("backend_version", "")) actual_version = str(getattr(av, "__version__", "")) if actual_version != expected_version: raise RuntimeError( "PyAV version differs from the cache decoder contract: " f"{actual_version!r} != {expected_version!r}" ) video_path = Path(video_path).expanduser().resolve() if not video_path.is_file(): raise FileNotFoundError(video_path) with av.open(str(video_path), mode="r") as container: streams = container.streams.video if len(streams) != 1: raise ValueError( f"expected exactly one video stream in {video_path}, " f"got {len(streams)}" ) frame_count_hint = int(getattr(streams[0], "frames", 0) or 0) provisional_start = max((frame_count_hint - num_frames) // 2, 0) selected, actual_count = _decode_window_and_count( av, video_path, start=provisional_start, num_frames=num_frames) if actual_count < num_frames: raise ValueError( f"OpenVid video has only {actual_count} decoded frames; " f"need {num_frames}: {video_path}" ) actual_start = (actual_count - num_frames) // 2 if actual_start != provisional_start or len(selected) != num_frames: selected, repeated_count = _decode_window_and_count( av, video_path, start=actual_start, num_frames=num_frames) if repeated_count != actual_count: raise RuntimeError( "PyAV returned a different frame count on deterministic " f"re-decode: {repeated_count} != {actual_count}" ) if len(selected) != num_frames: raise RuntimeError( f"failed to decode central {num_frames}-frame window from {video_path}" ) video01 = _frames_to_video01(selected) if tuple(video01.shape[:2]) != (num_frames, 3): raise RuntimeError( f"decoded video has invalid shape {tuple(video01.shape)}" ) frame_indices = torch.arange( actual_start, actual_start + num_frames, dtype=torch.long, ) return DecodedOpenVidVideo( video01=video01, frame_indices=frame_indices, source_frame_count=actual_count, ) DecodeFunction = Callable[[Path, Mapping[str, Any]], DecodedOpenVidVideo] def _default_decode( path: Path, contract: Mapping[str, Any], ) -> DecodedOpenVidVideo: return decode_openvid_center_contiguous_pyav( path, decoder_contract=contract) class OpenVidPromptVideoDataset(Dataset): """Return a prompt plus either cached teacher features or decoded video.""" def __init__( self, manifest_path: str | os.PathLike, video_root: str | os.PathLike, cache_root: str | os.PathLike, expected_model_sha256: str, expected_feature_contract_sha256: str, *, feature_contract: Mapping[str, Any] | None = None, decoder_contract: Mapping[str, Any] | None = None, decode_fn: DecodeFunction | None = None, ) -> None: self.manifest_path = Path(manifest_path).expanduser().resolve() self.video_root = Path(video_root).expanduser().resolve() if not self.video_root.is_dir(): raise NotADirectoryError(self.video_root) self.entries, self.manifest_metadata = load_openvid_manifest( self.manifest_path) if not self.entries: raise ValueError("OpenVid manifest is empty") self.feature_cache = OpenVidFeatureCache( cache_root, expected_model_sha256=expected_model_sha256, expected_feature_contract_sha256=( expected_feature_contract_sha256), feature_contract=( feature_contract if feature_contract is not None else OPENVID_VJEPA2_L15_FEATURE_CONTRACT ), decoder_contract=( decoder_contract if decoder_contract is not None else OPENVID_PYAV_DECODER_CONTRACT ), ) self.decoder_contract = self.feature_cache.decoder_contract self.decode_fn = decode_fn if decode_fn is not None else _default_decode self._cache_key_to_index: dict[str, int] = {} self._sample_id_to_index: dict[str, int] = {} for index, entry in enumerate(self.entries): cache_key = self.feature_cache.key_for(entry) if cache_key in self._cache_key_to_index: raise ValueError(f"duplicate OpenVid teacher cache key: {cache_key}") self._cache_key_to_index[cache_key] = index self._sample_id_to_index[entry.sample_id] = index def __len__(self) -> int: return len(self.entries) def _video_path_and_verify_stat( self, entry: OpenVidManifestEntry, ) -> Path: path = self.video_root / entry.video_name if not path.is_file(): raise FileNotFoundError(path) stat = path.stat() actual = (int(stat.st_size), int(stat.st_mtime_ns)) expected = (entry.video_size_bytes, entry.video_mtime_ns) if actual != expected: raise ValueError( "OpenVid video stat changed after manifest freeze for " f"{entry.video_name}: actual={actual}, expected={expected}" ) return path def cache_key_for_sample(self, sample_id: str) -> str: try: entry = self.entries[self._sample_id_to_index[str(sample_id)]] except KeyError as exc: raise KeyError(f"unknown OpenVid sample_id: {sample_id!r}") from exc return self.feature_cache.key_for(entry) def __getitem__(self, index: int) -> dict[str, Any]: entry = self.entries[int(index)] video_path = self._video_path_and_verify_stat(entry) cache_key = self.feature_cache.key_for(entry) cached = self.feature_cache.load(entry, cache_key=cache_key) item: dict[str, Any] = { "idx": entry.manifest_index, "prompts": entry.caption, "sample_id": entry.sample_id, "teacher_cache_key": cache_key, "video_name": entry.video_name, "video_path": str(video_path), } if cached is not None: item["current_teacher_cache_hit"] = True item["current_teacher_decode_seconds"] = 0.0 item["current_teacher_features"] = cached return item decode_started = time.perf_counter() decoded = self.decode_fn(video_path, self.decoder_contract) decode_seconds = time.perf_counter() - decode_started if not isinstance(decoded, DecodedOpenVidVideo): raise TypeError("decode_fn must return DecodedOpenVidVideo") expected_frames = int(self.decoder_contract["num_frames"]) source_frame_count = int(decoded.source_frame_count) expected_start = (source_frame_count - expected_frames) // 2 expected_indices = torch.arange( expected_start, expected_start + expected_frames, dtype=torch.long, ) if ( source_frame_count < expected_frames or not isinstance(decoded.frame_indices, torch.Tensor) or decoded.frame_indices.dtype != torch.long or not torch.equal( decoded.frame_indices.detach().cpu(), expected_indices) ): raise ValueError( "decode_fn violated the central-contiguous frame-index " "contract" ) video = decoded.video01 if ( not isinstance(video, torch.Tensor) or video.dtype != torch.float32 or video.ndim != 4 or tuple(video.shape[:2]) != (expected_frames, 3) or not bool(torch.isfinite(video).all().item()) or float(video.min()) < 0.0 or float(video.max()) > 1.0 ): raise ValueError( "decode_fn returned invalid video01; expected finite FP32 " f"[{expected_frames},3,H,W] in [0,1]" ) item.update({ "current_teacher_cache_hit": False, "current_teacher_decode_seconds": decode_seconds, "current_teacher_frame_indices": decoded.frame_indices, "current_teacher_source_frame_count": source_frame_count, "current_teacher_video01": video.contiguous(), }) return item def write_feature_cache( self, cache_key: str, feature: torch.Tensor, *, sample_id: str, ) -> Path: """Atomically publish one extracted teacher feature by issued key.""" try: index = self._cache_key_to_index[str(cache_key)] except KeyError as exc: raise KeyError(f"unknown teacher cache key: {cache_key!r}") from exc entry = self.entries[index] if entry.sample_id != str(sample_id): raise ValueError( "sample_id does not own teacher cache key: " f"{sample_id!r} != {entry.sample_id!r}" ) self._video_path_and_verify_stat(entry) return self.feature_cache.write( entry, feature, cache_key=str(cache_key), ) @staticmethod def collate_fn(batch: Sequence[Mapping[str, Any]]) -> dict[str, Any]: return collate_openvid_prompt_video(batch) def collate_openvid_prompt_video( batch: Sequence[Mapping[str, Any]], ) -> dict[str, Any]: """Collate the cache-hit/cache-miss union used by batch-size-one training.""" if not batch: raise ValueError("cannot collate an empty OpenVid batch") hits = torch.tensor( [bool(item["current_teacher_cache_hit"]) for item in batch], dtype=torch.bool, ) result: dict[str, Any] = { "idx": torch.tensor([int(item["idx"]) for item in batch]), "prompts": [str(item["prompts"]) for item in batch], "sample_ids": [str(item["sample_id"]) for item in batch], "teacher_cache_keys": [ str(item["teacher_cache_key"]) for item in batch ], "video_names": [str(item["video_name"]) for item in batch], "video_paths": [str(item["video_path"]) for item in batch], "current_teacher_cache_hit": hits, "current_teacher_decode_seconds": sum( float(item.get("current_teacher_decode_seconds", 0.0)) for item in batch ), } if bool(hits.all().item()): result["current_teacher_features"] = torch.stack([ item["current_teacher_features"] for item in batch ]) return result if bool((~hits).all().item()): videos = [item["current_teacher_video01"] for item in batch] shapes = {tuple(video.shape) for video in videos} if len(shapes) != 1: raise ValueError( "OpenVid raw videos have different spatial shapes; this " "collate path requires batch_size=1" ) result.update({ "current_teacher_frame_indices": torch.stack([ item["current_teacher_frame_indices"] for item in batch ]), "current_teacher_source_frame_count": torch.tensor([ int(item["current_teacher_source_frame_count"]) for item in batch ]), "current_teacher_video01": torch.stack(videos), }) return result raise ValueError( "mixed OpenVid teacher cache hits/misses cannot be collated; use the " "configured batch_size=1 while lazily populating the cache" ) __all__ = [ "DecodedOpenVidVideo", "OpenVidPromptVideoDataset", "collate_openvid_prompt_video", "decode_openvid_center_contiguous_pyav", ]