"""Five-segment VideoMAE-v2 encoder for an 81-frame Wan video. Frame zero is intentionally excluded. The remaining 80 frames are split into five fixed, adjacent 16-frame segments. Production training may reuse one frozen checkpoint instance for every segment; the legacy five-replica layout remains available for compatibility. Selected dense tokens are concatenated in chronological segment order. """ from __future__ import annotations from typing import Callable, Dict, Optional, Sequence import torch import torch.nn as nn from video_fd.encoders.videomae_v2_adapter import VideoMAEv2Adapter class SegmentedVideoMAEv2Ensemble(nn.Module): """Five logical same-checkpoint segments with a V-JEPA-compatible API.""" def __init__( self, config, *, adapter_factory: Optional[Callable[[dict], nn.Module]] = None, ) -> None: super().__init__() self.config = dict(config) self.encoder_count = int(self.config.get("encoder_count", 5)) self.share_encoder_weights = bool( self.config.get("share_encoder_weights", False) ) self.skip_leading_frames = int( self.config.get("skip_leading_frames", 1)) self.segment_num_frames = int( self.config.get("segment_num_frames", 16)) self.source_num_frames = int( self.config.get("source_num_frames", 81)) if self.encoder_count != 5: raise ValueError( "the VideoMAE-v2 whole-video contract requires exactly five " f"encoders, got {self.encoder_count}") if self.skip_leading_frames != 1: raise ValueError( "the VideoMAE-v2 whole-video contract skips exactly one " f"leading frame, got {self.skip_leading_frames}") if self.segment_num_frames != 16: raise ValueError( "each VideoMAE-v2 segment must contain 16 frames, got " f"{self.segment_num_frames}") expected_source = ( self.skip_leading_frames + self.encoder_count * self.segment_num_frames ) if self.source_num_frames != expected_source: raise ValueError( "source_num_frames must exactly cover the skipped frame and " f"five segments: {self.source_num_frames} != {expected_source}") factory = adapter_factory or VideoMAEv2Adapter adapter_config = dict(self.config) adapter_config.update({ "input_frames": self.segment_num_frames, "dense_num_frames": self.segment_num_frames, }) physical_encoder_count = ( 1 if self.share_encoder_weights else self.encoder_count ) self.encoders = nn.ModuleList([ factory(dict(adapter_config)) for _ in range(physical_encoder_count) ]) self._validate_replicas() def _validate_replicas(self) -> None: expected_physical = 1 if self.share_encoder_weights else self.encoder_count if len(self.encoders) != expected_physical: raise RuntimeError("VideoMAE-v2 replica construction was incomplete") if len({id(encoder) for encoder in self.encoders}) != expected_physical: raise RuntimeError( "VideoMAE-v2 physical encoder instances are not distinct") first = self.encoders[0] geometry = ( int(first.dense_layer_index), int(first.dense_num_frames), int(first.resize_short_side), int(first.input_size), int(first.retained_dense_blocks), int(first.dense_raw_tokens), int(first.dense_feature_dim), ) for index, encoder in enumerate(self.encoders[1:], start=1): candidate = ( int(encoder.dense_layer_index), int(encoder.dense_num_frames), int(encoder.resize_short_side), int(encoder.input_size), int(encoder.retained_dense_blocks), int(encoder.dense_raw_tokens), int(encoder.dense_feature_dim), ) if candidate != geometry: raise RuntimeError( "VideoMAE-v2 replica geometry differs at index " f"{index}: {candidate} != {geometry}") storage_sets = [] for encoder in self.encoders: storage_sets.append({ parameter.untyped_storage().data_ptr() for parameter in encoder.parameters() if parameter.numel() > 0 }) for left in range(len(storage_sets)): for right in range(left + 1, len(storage_sets)): if storage_sets[left].intersection(storage_sets[right]): raise RuntimeError( "VideoMAE-v2 replicas unexpectedly share parameter " f"storage between segments {left} and {right}") def _encoder_for_segment(self, index: int) -> nn.Module: if index < 0 or index >= self.encoder_count: raise IndexError(f"segment index outside [0,{self.encoder_count})") return self.encoders[0 if self.share_encoder_weights else index] @property def dense_layer_index(self) -> int: return int(self.encoders[0].dense_layer_index) @property def dense_num_frames(self) -> int: return self.segment_num_frames @property def resize_short_side(self) -> int: return int(self.encoders[0].resize_short_side) @property def input_size(self) -> int: return int(self.encoders[0].input_size) @property def retained_dense_blocks(self) -> int: return int(self.encoders[0].retained_dense_blocks) @property def dense_raw_tokens(self) -> int: """Raw tokens per 16-frame segment, not the five-segment total.""" return int(self.encoders[0].dense_raw_tokens) @property def aggregate_dense_raw_tokens(self) -> int: return self.encoder_count * self.dense_raw_tokens @property def dense_feature_dim(self) -> int: return int(self.encoders[0].dense_feature_dim) def _validate_video(self, video_clip_01: torch.Tensor) -> None: if video_clip_01.ndim != 5 or int(video_clip_01.shape[2]) != 3: raise ValueError( "video_clip_01 must have shape [B,81,3,H,W], got " f"{tuple(video_clip_01.shape)}") if int(video_clip_01.shape[1]) != self.source_num_frames: raise ValueError( "segmented VideoMAE-v2 requires exactly " f"{self.source_num_frames} decoded frames, got " f"{int(video_clip_01.shape[1])}") def segment(self, video_clip_01: torch.Tensor, index: int) -> torch.Tensor: self._validate_video(video_clip_01) index = int(index) if index < 0 or index >= self.encoder_count: raise IndexError(f"segment index outside [0,{self.encoder_count})") start = self.skip_leading_frames + index * self.segment_num_frames return video_clip_01.narrow(1, start, self.segment_num_frames) def forward_dense( self, video_clip_01: torch.Tensor, layer_index: Optional[int] = None, token_indices: Optional[torch.Tensor] = None, gradient_checkpointing: bool = False, temporal_window_starts: Optional[torch.Tensor] = None, temporal_window_indices: Optional[torch.Tensor] = None, ) -> torch.Tensor: if temporal_window_starts is not None or temporal_window_indices is not None: raise ValueError( "fixed five-segment VideoMAE-v2 does not accept temporal windows") outputs = [ self._encoder_for_segment(index).forward_dense( self.segment(video_clip_01, index), layer_index=layer_index, token_indices=token_indices, gradient_checkpointing=gradient_checkpointing, ) for index in range(self.encoder_count) ] return torch.cat(outputs, dim=1) def forward_dense_layers( self, video_clip_01: torch.Tensor, layer_indices: Sequence[int], token_indices: Optional[torch.Tensor] = None, gradient_checkpointing: bool = False, temporal_window_starts: Optional[torch.Tensor] = None, temporal_window_indices: Optional[torch.Tensor] = None, ) -> Dict[int, torch.Tensor]: if temporal_window_starts is not None or temporal_window_indices is not None: raise ValueError( "fixed five-segment VideoMAE-v2 does not accept temporal windows") layers = tuple(int(layer) for layer in layer_indices) by_segment = [ self._encoder_for_segment(index).forward_dense_layers( self.segment(video_clip_01, index), layer_indices=layers, token_indices=token_indices, gradient_checkpointing=gradient_checkpointing, ) for index in range(self.encoder_count) ] return { layer: torch.cat( [segment_features[layer] for segment_features in by_segment], dim=1, ) for layer in layers } def replay_dense_vjp( self, video_clip_01: torch.Tensor, *, layer_indices: Sequence[int], token_indices: torch.Tensor, expected_features: torch.Tensor, feature_cotangent: torch.Tensor, gradient_checkpointing: bool, ) -> tuple[torch.Tensor, torch.Tensor]: """Replay five encoder VJPs sequentially to cap activation memory. Feature transport is layer-major and, within a layer, segment-major. The returned pixel cotangent has a zero first frame by construction. """ self._validate_video(video_clip_01) layers = tuple(int(layer) for layer in layer_indices) if not layers: raise ValueError("layer_indices must not be empty") tokens_per_segment = int(token_indices.numel()) tokens_per_layer = self.encoder_count * tokens_per_segment expected_shape = ( int(video_clip_01.shape[0]), len(layers) * tokens_per_layer, self.dense_feature_dim, ) if tuple(expected_features.shape) != expected_shape: raise ValueError( "expected_features has the wrong segmented shape: " f"{tuple(expected_features.shape)} != {expected_shape}") if tuple(feature_cotangent.shape) != expected_shape: raise ValueError( "feature_cotangent has the wrong segmented shape: " f"{tuple(feature_cotangent.shape)} != {expected_shape}") expected_by_layer = expected_features.split(tokens_per_layer, dim=1) cotangent_by_layer = feature_cotangent.split(tokens_per_layer, dim=1) replay_by_layer = [[] for _ in layers] video_cotangent = torch.zeros_like(video_clip_01, dtype=torch.float32) for segment_index in range(self.encoder_count): encoder = self._encoder_for_segment(segment_index) segment_leaf = self.segment( video_clip_01, segment_index).detach().requires_grad_(True) if len(layers) == 1: outputs = { layers[0]: encoder.forward_dense( segment_leaf, layer_index=layers[0], token_indices=token_indices, gradient_checkpointing=gradient_checkpointing, ).float() } else: outputs = encoder.forward_dense_layers( segment_leaf, layer_indices=layers, token_indices=token_indices, gradient_checkpointing=gradient_checkpointing, ) outputs = {layer: value.float() for layer, value in outputs.items()} surrogate = segment_leaf.new_zeros((), dtype=torch.float32) for layer_offset, layer in enumerate(layers): start = segment_index * tokens_per_segment expected_part = expected_by_layer[layer_offset].narrow( 1, start, tokens_per_segment) cotangent_part = cotangent_by_layer[layer_offset].narrow( 1, start, tokens_per_segment) output = outputs[layer] if tuple(output.shape) != tuple(expected_part.shape): raise RuntimeError( "VideoMAE-v2 segment replay shape changed for " f"segment={segment_index}, layer={layer}") replay_by_layer[layer_offset].append(output.detach()) surrogate = surrogate + (output * cotangent_part).sum() (segment_cotangent,) = torch.autograd.grad( surrogate, segment_leaf, retain_graph=False, create_graph=False, ) start_frame = ( self.skip_leading_frames + segment_index * self.segment_num_frames ) video_cotangent[:, start_frame:start_frame + self.segment_num_frames] = ( segment_cotangent.detach().float()) del surrogate, segment_leaf, segment_cotangent, outputs replay_features = torch.cat([ torch.cat(parts, dim=1) for parts in replay_by_layer ], dim=1) return replay_features, video_cotangent __all__ = ["SegmentedVideoMAEv2Ensemble"]