"""Extract frozen MuVAP VAP/ASD states for an EgoSpeak Stage-3 baseline. This is a clean-room adaptation of the architecture described in MuVAP Sec. 4.3/5.3. The upstream repository has not released its Stage-3 trainer, so the feature contract is deliberately named ``paper-reconstruction`` rather than claiming to reproduce unpublished code. The extractor never reads speech targets. A verified global-feature sidecar is used only for its dense 5 Hz timestamps and source-media inventory; none of its Omni tensors are copied into the output. MuVAP's released role-VAP and ASD backbones are frozen and evaluated on a 25 Hz, strictly causal timeline. """ from __future__ import annotations import argparse import hashlib import math import os import re import subprocess import sys import tempfile from pathlib import Path from typing import Any import numpy as np import torch from tqdm.auto import tqdm from easycom_head import resolve_feature_path from extract_highrate_person_av_features import ( FACE_SIZE, MODEL_FPS, OUTPUT_HZ, SAMPLE_RATE, SAMPLES_PER_MODEL_FRAME, FrozenMuVAPASDEncoder, _audio_for_tick_range, _chunked_audio_encode, _collect_clip_tracks, _model_tick_bounds, _release_paths, ) from extract_turn_local_features import ( EGO4D_MAX_TERMINAL_FRAME_OVERRUN, build_media_and_boxes, load_ego4d_boxes, select_samples, validate_source_timeline, ) from local_turn_features import reference_alignment from muvap_egospeak_stage3 import MUVAP_EGOSPEAK_INPUT_FORMAT def _sha256(path: Path) -> str: digest = hashlib.sha256() with path.open("rb") as stream: while block := stream.read(1024 * 1024): digest.update(block) return digest.hexdigest() def _verified_root(root: Path, revision: str) -> tuple[Path, str]: root = root.expanduser().resolve() if not re.fullmatch(r"[0-9a-fA-F]{7,40}", revision): raise ValueError("--muvap-revision must be a hexadecimal commit prefix") if not (root / "models" / "vap.py").is_file(): raise FileNotFoundError(f"MuVAP source is incomplete: {root}") revision_file = root / ".upstream-revision" if revision_file.is_file(): actual = revision_file.read_text().strip().lower() else: completed = subprocess.run( ["git", "-C", str(root), "rev-parse", "HEAD"], check=False, capture_output=True, text=True, ) actual = completed.stdout.strip().lower() if completed.returncode == 0 else "" if not re.fullmatch(r"[0-9a-f]{40}", actual) or not actual.startswith( revision.lower() ): raise RuntimeError( f"MuVAP checkout is {actual or 'unreadable'}, expected {revision!r}" ) return root, actual class FrozenMuVAPVAPEncoder: """Load the published 25 Hz mono role-VAP and expose ``Z_VAP``.""" def __init__( self, *, muvap_root: Path, revision: str, release: Path, device: str, dtype: str, ) -> None: root, self.source_revision = _verified_root(muvap_root, revision) self.weights_path, config_path = _release_paths(release) try: import yaml except ImportError as error: raise RuntimeError("PyYAML is required to read MuVAP releases") from error sys.path.insert(0, str(root)) try: from models.release import load_weights from models.vap import build_vap except Exception as error: raise RuntimeError(f"Cannot import MuVAP from {root}") from error imported = Path(sys.modules[build_vap.__module__].__file__).resolve() if root not in imported.parents: raise RuntimeError( f"Python resolved {build_vap.__module__} from {imported}, not {root}" ) payload = yaml.safe_load(config_path.read_text()) or {} config = payload.get("vap") if not isinstance(config, dict): raise ValueError(f"MuVAP release has no `vap` config: {config_path}") projection = config.get("projection_window", {}) audio_spec = config.get("audio_encoder", "cpc") audio_name = audio_spec if isinstance(audio_spec, str) else audio_spec.get("name") if audio_name != "cpc": raise ValueError("The first baseline requires the CPC VAP release") if config.get("arch", "mono") != "mono": raise ValueError("EgoSpeak mixed audio requires MuVAP's mono role-VAP") if projection.get("mode") != "role_relative": raise ValueError("Expected a role-relative VAP release") if float(projection.get("frame_hz", -1)) != MODEL_FPS: raise ValueError( f"Expected a {MODEL_FPS:g} Hz VAP release, got " f"{projection.get('frame_hz')!r}" ) if config.get("temporal", {}).get("causal", True) is not True: raise ValueError("VAP temporal attention must be causal") self.device = torch.device(device) if self.device.type == "cuda" and not torch.cuda.is_available(): raise RuntimeError(f"CUDA requested but unavailable: {device}") self.autocast_dtype = { "float32": None, "float16": torch.float16, "bfloat16": torch.bfloat16, }[dtype] try: model = build_vap(config) except FileNotFoundError as error: raise RuntimeError( "MuVAP's CPC frontend is missing; run " f"`python {root / 'tools' / 'fetch_frontends.py'} cpc`" ) from error self.model = load_weights(model, self.weights_path).to(self.device).eval() self.model.requires_grad_(False) self.feature_dim = int(config["temporal"]["d_model"]) self.release_config = config self.release_sha256 = _sha256(self.weights_path) @torch.inference_mode() def encode_audio(self, waveform: np.ndarray, num_frames: int) -> torch.Tensor: expected = num_frames * SAMPLES_PER_MODEL_FRAME if waveform.dtype != np.float32 or waveform.shape != (expected,): raise ValueError(f"Audio must be float32 [{expected}], got {waveform.shape}") audio = torch.from_numpy(np.ascontiguousarray(waveform))[None, None] audio = audio.to(self.device, dtype=torch.float32) with torch.autocast( device_type=self.device.type, dtype=self.autocast_dtype, enabled=self.autocast_dtype is not None, ): _, embedding = self.model(audio, return_embeddings=True) result = embedding[0].float().cpu() if result.shape != (num_frames, self.feature_dim): raise RuntimeError(f"Unexpected VAP shape: {tuple(result.shape)}") if not torch.isfinite(result).all(): raise RuntimeError("MuVAP VAP produced NaN/Inf") return result def _chunked_vap_encode( encoder: FrozenMuVAPVAPEncoder, waveform: np.ndarray, *, num_frames: int, output_frames: int, context_frames: int, ) -> torch.Tensor: """Encode every tick with a fixed, past-only VAP context.""" if output_frames <= 0 or context_frames < 0: raise ValueError("VAP chunk/context frame counts are invalid") if waveform.shape != (num_frames * SAMPLES_PER_MODEL_FRAME,): raise ValueError("Audio/VAP high-rate lengths differ") pieces = [] for output_start in range(0, num_frames, output_frames): input_start = max(0, output_start - context_frames) output_end = min(num_frames, output_start + output_frames) samples = waveform[ input_start * SAMPLES_PER_MODEL_FRAME : output_end * SAMPLES_PER_MODEL_FRAME ] encoded = encoder.encode_audio(samples, output_end - input_start) pieces.append(encoded[output_start - input_start :]) result = torch.cat(pieces) if pieces else torch.empty(0, encoder.feature_dim) if result.shape != (num_frames, encoder.feature_dim): raise RuntimeError("Chunked VAP output did not cover every frame once") return result def _chunked_fused_asd_encode( encoder: FrozenMuVAPASDEncoder, faces: np.ndarray, audio_states: torch.Tensor, *, output_frames: int, context_frames: int, ) -> torch.Tensor: """Encode one consecutive face track into frozen ``Z_ASD`` states.""" if output_frames <= 0 or context_frames < 0: raise ValueError("ASD chunk/context frame counts are invalid") if audio_states.shape != (len(faces), encoder.audio_dim): raise ValueError("Face/shared-audio high-rate lengths differ") pieces = [] for output_start in range(0, len(faces), output_frames): input_start = max(0, output_start - context_frames) output_end = min(len(faces), output_start + output_frames) encoded = encoder.encode_person_fused( faces[input_start:output_end], audio_states[input_start:output_end], ) pieces.append(encoded[output_start - input_start :]) result = torch.cat(pieces) if pieces else torch.empty(0, encoder.fused_dim) if result.shape != (len(faces), encoder.fused_dim): raise RuntimeError("Chunked fused ASD output did not cover every frame once") return result def extract_sample( media: Any, annotations: list[dict[int, list[dict[str, Any]]]], vap: FrozenMuVAPVAPEncoder, asd: FrozenMuVAPASDEncoder, task_times: torch.Tensor, args: argparse.Namespace, ) -> dict[str, Any]: """Extract 25 Hz VAP states and ragged fused-ASD person sets.""" ticks_per_task = round(MODEL_FPS / OUTPUT_HZ) if abs(ticks_per_task * OUTPUT_HZ - MODEL_FPS) > 1e-8: raise RuntimeError("MuVAP/task grids are not integral") num_ticks = len(task_times) * ticks_per_task waveform = _audio_for_tick_range(media, 0, num_ticks) vap_states = _chunked_vap_encode( vap, waveform, num_frames=num_ticks, output_frames=round(args.vap_output_seconds * MODEL_FPS), context_frames=round(args.vap_context_seconds * MODEL_FPS), ) asd_audio = _chunked_audio_encode( asd, waveform, num_frames=num_ticks, output_frames=round(args.asd_output_seconds * MODEL_FPS), context_frames=round(args.asd_context_seconds * MODEL_FPS), ) rows: list[list[tuple[torch.Tensor, dict[str, Any]]]] = [ [] for _ in range(num_ticks) ] required_end = len(task_times) / OUTPUT_HZ for clip_index, clip in enumerate(media.clips): if float(clip.dense_start) >= required_end - 1e-8: break start_tick, stop_tick = _model_tick_bounds(clip, required_end) segments = _collect_clip_tracks( media, annotations, clip_index, start_tick, stop_tick, args, ) for segment in tqdm( segments, desc=f"fused ASD {clip.path.name}", unit="track", disable=args.no_progress, leave=False, ): positions = torch.tensor(segment["ticks"], dtype=torch.long) if bool(((positions < 0) | (positions >= num_ticks)).any()): raise RuntimeError("ASD track extends beyond the VAP timeline") states = _chunked_fused_asd_encode( asd, np.stack(segment["faces"]), asd_audio[positions], output_frames=round(args.asd_output_seconds * MODEL_FPS), context_frames=round(args.asd_context_seconds * MODEL_FPS), ) for offset, tick in enumerate(segment["ticks"]): record = segment["records"][offset] if abs(float(record["asd_interval_end"]) - (tick + 1) / MODEL_FPS) > 1e-8: raise RuntimeError("ASD record/tick right edges differ") rows[tick].append((states[offset].to(torch.bfloat16), record)) features: list[torch.Tensor] = [] records: list[dict[str, Any]] = [] offsets = [0] max_people = 0 for observations in rows: if observations: features.append(torch.stack([value for value, _ in observations])) records.extend(record for _, record in observations) max_people = max(max_people, len(observations)) offsets.append(offsets[-1] + len(observations)) flat = ( torch.cat(features) if features else torch.empty(0, asd.fused_dim, dtype=torch.bfloat16) ) task_frame_indices = torch.arange(len(task_times), dtype=torch.long) task_frame_indices = (task_frame_indices + 1) * ticks_per_task - 1 if len(task_times): highrate_edges = (task_frame_indices + 1).to(torch.float64) / MODEL_FPS if not torch.allclose(highrate_edges, task_times, atol=1e-8, rtol=0): raise RuntimeError("25 Hz right edges do not match the 5 Hz task cutoffs") return { "vap": vap_states.to(torch.bfloat16), "frame_times": torch.arange(1, num_ticks + 1, dtype=torch.float64) / MODEL_FPS, "persons": { "features": flat, "frame_offsets": torch.tensor(offsets, dtype=torch.long), "records": records, }, "task_frame_indices": task_frame_indices, "max_people": max_people, } def requested_config( args: argparse.Namespace, vap: FrozenMuVAPVAPEncoder, asd: FrozenMuVAPASDEncoder, ) -> dict[str, Any]: return { "implementation": "clean-room-paper-reconstruction-not-upstream-stage3-code", "uses_omni_features": False, "muvap_source_revision": vap.source_revision, "vap_release_sha256": vap.release_sha256, "asd_release_sha256": asd.release_sha256, "vap_release_config": vap.release_config, "asd_release_config": asd.release_config, "vap_dim": vap.feature_dim, "asd_dim": asd.fused_dim, "model_fps": MODEL_FPS, "task_hz": OUTPUT_HZ, "sample_rate": SAMPLE_RATE, "face_size": [FACE_SIZE, FACE_SIZE], "vap_output_seconds": args.vap_output_seconds, "vap_context_seconds": args.vap_context_seconds, "asd_output_seconds": args.asd_output_seconds, "asd_context_seconds": args.asd_context_seconds, "audio_preprocessing": "causal-float-pcm-clip-only-no-future-rms", "box_kind": args.box_kind if args.dataset == "easycom" else "ego4d-provided-tracks", "box_coordinate_space": args.box_coordinate_space, "crop_expansion": args.crop_expansion, "box_protocol": "provided-boxes-not-online-detector", "uses_speech_labels": False, "causal_policy": "25hz-right-edge-causal-fixed-past-context", } def parse_args(argv: list[str] | None = None) -> argparse.Namespace: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--dataset", choices=("easycom", "ego4d"), default="easycom") parser.add_argument("--data-root", type=Path, required=True) parser.add_argument( "--reference-feature-root", type=Path, required=True, help="Verified 5 Hz sidecars used only for timestamps/source inventory", ) parser.add_argument("--output-root", type=Path, required=True) selection = parser.add_mutually_exclusive_group(required=True) selection.add_argument("--samples", nargs="+") selection.add_argument("--manifest", type=Path) parser.add_argument("--split", choices=("all", "train", "test"), default="all") parser.add_argument("--start-index", type=int, default=1) parser.add_argument("--end-index", type=int) parser.add_argument("--limit", type=int) parser.add_argument("--num-shards", type=int, default=1) parser.add_argument("--shard-index", type=int, default=0) parser.add_argument("--ego4d-annotations", type=Path, nargs="+") parser.add_argument( "--ego4d-max-terminal-overrun-frames", type=int, default=EGO4D_MAX_TERMINAL_FRAME_OVERRUN, ) parser.add_argument("--muvap-root", type=Path, default=Path("third_party/MuVAP")) parser.add_argument("--muvap-revision", default="a1ef9c8") parser.add_argument("--vap-weights", type=Path, required=True) parser.add_argument("--asd-weights", type=Path, required=True) parser.add_argument("--device", default="cuda:0") parser.add_argument( "--dtype", choices=("float32", "float16", "bfloat16"), default="bfloat16", ) parser.add_argument("--vap-output-seconds", type=float, default=10.0) parser.add_argument("--vap-context-seconds", type=float, default=20.0) parser.add_argument("--asd-output-seconds", type=float, default=10.0) parser.add_argument("--asd-context-seconds", type=float, default=2.0) parser.add_argument("--box-kind", choices=("face", "head"), default="face") parser.add_argument( "--box-coordinate-space", choices=("camera", "encoded"), default=None ) parser.add_argument("--crop-expansion", type=float, default=1.0) parser.add_argument( "--decoder-backend", choices=("decord", "torchcodec", "pynvcodec"), default="decord", ) parser.add_argument("--decoder-threads", type=int, default=1) parser.add_argument("--decoder-gpu-id", type=int, default=0) parser.add_argument("--max-steps", type=int) parser.add_argument("--overwrite", action="store_true") parser.add_argument("--dry-run", action="store_true") parser.add_argument("--no-progress", action="store_true") args = parser.parse_args(argv) if args.manifest is not None and args.dataset != "ego4d": parser.error("--manifest is only supported for Ego4D") if args.manifest is None and ( args.split != "all" or args.start_index != 1 or args.end_index is not None or args.limit is not None ): parser.error("--split/--start-index/--end-index/--limit require --manifest") if args.start_index < 1 or ( args.end_index is not None and args.end_index < args.start_index ): parser.error("Invalid manifest range") if args.limit is not None and args.limit <= 0: parser.error("--limit must be positive") if args.num_shards <= 0 or not 0 <= args.shard_index < args.num_shards: parser.error("Require --num-shards > 0 and a valid --shard-index") if args.max_steps is not None and args.max_steps <= 0: parser.error("--max-steps must be positive") for name in ( "vap_output_seconds", "vap_context_seconds", "asd_output_seconds", "asd_context_seconds", ): value = getattr(args, name) if not math.isfinite(value) or value < 0 or ( name.endswith("output_seconds") and value == 0 ): parser.error(f"--{name.replace('_', '-')} is invalid") frames = round(value * MODEL_FPS) if not math.isclose(frames / MODEL_FPS, value, abs_tol=1e-8): parser.error(f"--{name.replace('_', '-')} must align to the 25 Hz grid") if name.endswith("output_seconds") and frames == 0: parser.error(f"--{name.replace('_', '-')} is shorter than one frame") if not math.isfinite(args.crop_expansion) or args.crop_expansion <= 0: parser.error("--crop-expansion must be finite and positive") if args.decoder_threads <= 0 or args.decoder_gpu_id < 0: parser.error("decoder-threads must be positive and decoder-gpu-id nonnegative") if args.box_coordinate_space is None: args.box_coordinate_space = "encoded" if args.dataset == "easycom" else "camera" args.video_fps = MODEL_FPS return args def main(argv: list[str] | None = None) -> None: args = parse_args(argv) names = select_samples(args) if not names: print(f"Shard {args.shard_index}/{args.num_shards}: no assigned samples") return references = { name: resolve_feature_path(args.reference_feature_root, name, dataset=args.dataset) for name in names } for name in names: if not references[name].is_file(): raise FileNotFoundError(references[name]) destination = args.output_root / f"{name}.pt" if destination.exists() and not args.overwrite: raise FileExistsError( f"{destination} exists; use a new root or explicit --overwrite" ) print(f"Shard {args.shard_index}/{args.num_shards}: {len(names)} sample(s)") if args.dry_run: for name in names: reference = torch.load(references[name], map_location="cpu", weights_only=False) times, _ = reference_alignment(reference) steps = min(len(times), args.max_steps or len(times)) print(f"{name}: would extract {steps * 5} MuVAP frames for {steps} targets") return ego_paths: list[Path] = [] if args.dataset == "ego4d": ego_paths = args.ego4d_annotations or [ args.data_root / "v1" / "annotations" / f"av_{split}.json" for split in ("train", "val") ] ego_boxes = load_ego4d_boxes(ego_paths, names) if args.dataset == "ego4d" else {} vap = FrozenMuVAPVAPEncoder( muvap_root=args.muvap_root, revision=args.muvap_revision, release=args.vap_weights, device=args.device, dtype=args.dtype, ) asd = FrozenMuVAPASDEncoder( muvap_root=args.muvap_root, revision=args.muvap_revision, release=args.asd_weights, device=args.device, dtype=args.dtype, ) args.output_root.mkdir(parents=True, exist_ok=True) expected_config = requested_config(args, vap, asd) for position, name in enumerate(names, 1): print(f"[{position}/{len(names)}] {name}", flush=True) reference = torch.load(references[name], map_location="cpu", weights_only=False) times, indices = reference_alignment(reference) full_steps = len(times) if args.max_steps is not None: times = times[: args.max_steps] indices = indices[: args.max_steps] media, annotations = build_media_and_boxes( args, name, len(times) / OUTPUT_HZ, ego_boxes ) validate_source_timeline(reference, media) output = extract_sample(media, annotations, vap, asd, times, args) output.update( { "format": MUVAP_EGOSPEAK_INPUT_FORMAT, "dataset": args.dataset, "sample_id": name, "times": times.to(torch.float64), "label_indices": indices.to(torch.long), "partial": len(times) != full_steps, "reference_features": str(references[name].resolve()), "direct_media": media.metadata(), "config": expected_config, } ) if args.dataset == "ego4d": output["box_frame_audit"] = ego_boxes[name]["frame_audit"] destination = args.output_root / f"{name}.pt" with tempfile.NamedTemporaryFile( dir=args.output_root, prefix=f".{name}.", suffix=".tmp", delete=False, ) as handle: temporary = Path(handle.name) try: torch.save(output, temporary) os.replace(temporary, destination) finally: temporary.unlink(missing_ok=True) print( f"Saved {destination}; MuVAP frames={len(output['vap'])}, " f"person observations={len(output['persons']['features'])}, " f"max_people={output['max_people']}, partial={output['partial']}", flush=True, ) if __name__ == "__main__": main()