| 1 | from typing import NamedTuple |
| 2 | |
| 3 | import torch |
| 4 | from torch import nn |
| 5 | |
| 6 | from ltx_core.text_encoders.gemma.embeddings_connector import Embeddings1DConnector |
| 7 | |
| 8 | |
| 9 | class EmbeddingsProcessorOutput(NamedTuple): |
| 10 | video_encoding: torch.Tensor |
| 11 | audio_encoding: torch.Tensor | None |
| 12 | attention_mask: torch.Tensor |
| 13 | |
| 14 | |
| 15 | def convert_to_additive_mask(attention_mask: torch.Tensor, dtype: torch.dtype) -> torch.Tensor: |
| 16 | """Convert binary attention mask to additive form for transformer masking.""" |
| 17 | return (attention_mask.to(torch.int64) - 1).to(dtype).reshape( |
| 18 | (attention_mask.shape[0], 1, -1, attention_mask.shape[-1]) |
| 19 | ) * torch.finfo(dtype).max |
| 20 | |
| 21 | |
| 22 | def _to_binary_mask(encoded: torch.Tensor, encoded_mask: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: |
| 23 | """Convert connector output mask to binary mask and apply to encoded tensor.""" |
| 24 | binary_mask = (encoded_mask < 0.000001).to(torch.int64) |
| 25 | binary_mask = binary_mask.reshape([encoded.shape[0], encoded.shape[1], 1]) |
| 26 | encoded = encoded * binary_mask |
| 27 | return encoded, binary_mask |
| 28 | |
| 29 | |
| 30 | class EmbeddingsProcessor(nn.Module): |
| 31 | """Wraps feature extractor + video connector + optional audio connector. |
| 32 | Can operate in two modes: |
| 33 | 1. create_embeddings(): Takes pre-computed features + additive mask (backward compat, used by trainer) |
| 34 | 2. process_hidden_states(): Takes raw Gemma hidden states, runs feature extraction + connectors |
| 35 | """ |
| 36 | |
| 37 | def __init__( |
| 38 | self, |
| 39 | *, |
| 40 | feature_extractor: nn.Module | None = None, |
| 41 | video_connector: Embeddings1DConnector, |
| 42 | audio_connector: Embeddings1DConnector | None = None, |
| 43 | ): |
| 44 | super().__init__() |
| 45 | self.feature_extractor = feature_extractor |
| 46 | self.video_connector = video_connector |
| 47 | self.audio_connector = audio_connector |
| 48 | |
| 49 | def create_embeddings( |
| 50 | self, |
| 51 | video_features: torch.Tensor, |
| 52 | audio_features: torch.Tensor | None, |
| 53 | additive_attention_mask: torch.Tensor, |
| 54 | ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor]: |
| 55 | if self.audio_connector is not None and audio_features is None: |
| 56 | raise ValueError("Audio connector is configured but no audio features were provided.") |
| 57 | if self.audio_connector is None and audio_features is not None: |
| 58 | raise ValueError("Audio features were provided but no audio connector is configured.") |
| 59 | |
| 60 | video_encoded, video_mask = self.video_connector(video_features, additive_attention_mask) |
| 61 | video_encoded, binary_mask = _to_binary_mask(video_encoded, video_mask) |
| 62 | |
| 63 | audio_encoded = None |
| 64 | if self.audio_connector is not None: |
| 65 | audio_encoded, _ = self.audio_connector(audio_features, additive_attention_mask) |
| 66 | |
| 67 | return video_encoded, audio_encoded, binary_mask.squeeze(-1) |
| 68 | |
| 69 | def process_hidden_states( |
| 70 | self, |
| 71 | hidden_states: tuple[torch.Tensor, ...], |
| 72 | attention_mask: torch.Tensor, |
| 73 | padding_side: str = "left", |
| 74 | ) -> EmbeddingsProcessorOutput: |
| 75 | """Full pipeline: feature extraction -> connectors -> final embeddings. |
| 76 | Args: |
| 77 | hidden_states: Raw Gemma hidden states (tuple of tensors per layer). |
| 78 | attention_mask: Binary attention mask [B, seq_len]. |
| 79 | padding_side: Padding side used during tokenization. |
| 80 | Returns: |
| 81 | EmbeddingsProcessorOutput with video_encoding, audio_encoding, and attention_mask. |
| 82 | """ |
| 83 | if self.feature_extractor is None: |
| 84 | raise ValueError("feature_extractor is required for process_hidden_states()") |
| 85 | |
| 86 | video_feats, audio_feats = self.feature_extractor(hidden_states, attention_mask, padding_side) |
| 87 | additive_mask = convert_to_additive_mask(attention_mask, video_feats.dtype) |
| 88 | video_enc, audio_enc, binary_mask = self.create_embeddings(video_feats, audio_feats, additive_mask) |
| 89 | return EmbeddingsProcessorOutput(video_enc, audio_enc, binary_mask) |
| 90 |