| """Five-replica 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. Each segment is processed by its own |
| frozen VideoMAE-v2 module instance loaded from the same checkpoint. 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 physical same-checkpoint encoders 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.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)) |
| expected_source = ( |
| self.skip_leading_frames |
| + self.encoder_count * self.segment_num_frames |
| ) |
|
|
| 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, |
| }) |
| |
| |
| self.encoders = nn.ModuleList([ |
| factory(dict(adapter_config)) |
| for _ in range(self.encoder_count) |
| ]) |
| self._validate_replicas() |
|
|
| def _validate_replicas(self) -> None: |
| return None |
|
|
| @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: |
| return None |
|
|
| def segment(self, video_clip_01: torch.Tensor, index: int) -> torch.Tensor: |
| self._validate_video(video_clip_01) |
| index = int(index) |
| 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: |
| outputs = [ |
| encoder.forward_dense( |
| self.segment(video_clip_01, index), |
| layer_index=layer_index, |
| token_indices=token_indices, |
| gradient_checkpointing=gradient_checkpointing, |
| ) |
| for index, encoder in enumerate(self.encoders) |
| ] |
| 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]: |
| layers = tuple(int(layer) for layer in layer_indices) |
| by_segment = [ |
| encoder.forward_dense_layers( |
| self.segment(video_clip_01, index), |
| layer_indices=layers, |
| token_indices=token_indices, |
| gradient_checkpointing=gradient_checkpointing, |
| ) |
| for index, encoder in enumerate(self.encoders) |
| ] |
| 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) |
| 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, |
| ) |
|
|
| 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, encoder in enumerate(self.encoders): |
| 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] |
| 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"] |
|
|
