| 1 | import torch |
| 2 | |
| 3 | from ltx_core.model.model_protocol import ModelConfigurator |
| 4 | from ltx_core.model.transformer.attention import Attention |
| 5 | from ltx_core.model.transformer.feed_forward import FeedForward |
| 6 | from ltx_core.model.transformer.rope import ( |
| 7 | LTXRopeType, |
| 8 | generate_freq_grid_np, |
| 9 | generate_freq_grid_pytorch, |
| 10 | precompute_freqs_cis, |
| 11 | ) |
| 12 | from ltx_core.utils import rms_norm |
| 13 | |
| 14 | |
| 15 | class _BasicTransformerBlock1D(torch.nn.Module): |
| 16 | def __init__( |
| 17 | self, |
| 18 | dim: int, |
| 19 | heads: int, |
| 20 | dim_head: int, |
| 21 | rope_type: LTXRopeType = LTXRopeType.INTERLEAVED, |
| 22 | apply_gated_attention: bool = False, |
| 23 | ): |
| 24 | super().__init__() |
| 25 | |
| 26 | self.attn1 = Attention( |
| 27 | query_dim=dim, |
| 28 | heads=heads, |
| 29 | dim_head=dim_head, |
| 30 | rope_type=rope_type, |
| 31 | apply_gated_attention=apply_gated_attention, |
| 32 | ) |
| 33 | |
| 34 | self.ff = FeedForward( |
| 35 | dim, |
| 36 | dim_out=dim, |
| 37 | ) |
| 38 | |
| 39 | def forward( |
| 40 | self, |
| 41 | hidden_states: torch.Tensor, |
| 42 | attention_mask: torch.Tensor | None = None, |
| 43 | pe: torch.Tensor | None = None, |
| 44 | ) -> torch.Tensor: |
| 45 | # Notice that normalization is always applied before the real computation in the following blocks. |
| 46 | |
| 47 | # 1. Normalization Before Self-Attention |
| 48 | norm_hidden_states = rms_norm(hidden_states) |
| 49 | |
| 50 | norm_hidden_states = norm_hidden_states.squeeze(1) |
| 51 | |
| 52 | # 2. Self-Attention |
| 53 | attn_output = self.attn1(norm_hidden_states, mask=attention_mask, pe=pe) |
| 54 | |
| 55 | hidden_states = attn_output + hidden_states |
| 56 | if hidden_states.ndim == 4: |
| 57 | hidden_states = hidden_states.squeeze(1) |
| 58 | |
| 59 | # 3. Normalization before Feed-Forward |
| 60 | norm_hidden_states = rms_norm(hidden_states) |
| 61 | |
| 62 | # 4. Feed-forward |
| 63 | ff_output = self.ff(norm_hidden_states) |
| 64 | |
| 65 | hidden_states = ff_output + hidden_states |
| 66 | if hidden_states.ndim == 4: |
| 67 | hidden_states = hidden_states.squeeze(1) |
| 68 | |
| 69 | return hidden_states |
| 70 | |
| 71 | |
| 72 | class Embeddings1DConnector(torch.nn.Module): |
| 73 | """ |
| 74 | Embeddings1DConnector applies a 1D transformer-based processing to sequential embeddings (e.g., for video, audio, or |
| 75 | other modalities). It supports rotary positional encoding (rope), optional causal temporal positioning, and can |
| 76 | substitute padded positions with learnable registers. The module is highly configurable for head size, number of |
| 77 | layers, and register usage. |
| 78 | Args: |
| 79 | attention_head_dim (int): Dimension of each attention head (default=128). |
| 80 | num_attention_heads (int): Number of attention heads (default=30). |
| 81 | num_layers (int): Number of transformer layers (default=2). |
| 82 | positional_embedding_theta (float): Scaling factor for position embedding (default=10000.0). |
| 83 | positional_embedding_max_pos (list[int] | None): Max positions for positional embeddings (default=[1]). |
| 84 | causal_temporal_positioning (bool): If True, uses causal attention (default=False). |
| 85 | num_learnable_registers (int | None): Number of learnable registers to replace padded tokens. If None, disables |
| 86 | register replacement. (default=128) |
| 87 | rope_type (LTXRopeType): The RoPE variant to use (default=DEFAULT_ROPE_TYPE). |
| 88 | double_precision_rope (bool): Use double precision rope calculation (default=False). |
| 89 | """ |
| 90 | |
| 91 | _supports_gradient_checkpointing = True |
| 92 | |
| 93 | def __init__( |
| 94 | self, |
| 95 | attention_head_dim: int = 128, |
| 96 | num_attention_heads: int = 30, |
| 97 | num_layers: int = 2, |
| 98 | positional_embedding_theta: float = 10000.0, |
| 99 | positional_embedding_max_pos: list[int] | None = None, |
| 100 | causal_temporal_positioning: bool = False, |
| 101 | num_learnable_registers: int | None = 128, |
| 102 | rope_type: LTXRopeType = LTXRopeType.INTERLEAVED, |
| 103 | double_precision_rope: bool = False, |
| 104 | apply_gated_attention: bool = False, |
| 105 | ): |
| 106 | super().__init__() |
| 107 | self.num_attention_heads = num_attention_heads |
| 108 | self.inner_dim = num_attention_heads * attention_head_dim |
| 109 | self.causal_temporal_positioning = causal_temporal_positioning |
| 110 | self.positional_embedding_theta = positional_embedding_theta |
| 111 | self.positional_embedding_max_pos = ( |
| 112 | positional_embedding_max_pos if positional_embedding_max_pos is not None else [1] |
| 113 | ) |
| 114 | self.rope_type = rope_type |
| 115 | self.double_precision_rope = double_precision_rope |
| 116 | self.transformer_1d_blocks = torch.nn.ModuleList( |
| 117 | [ |
| 118 | _BasicTransformerBlock1D( |
| 119 | dim=self.inner_dim, |
| 120 | heads=num_attention_heads, |
| 121 | dim_head=attention_head_dim, |
| 122 | rope_type=rope_type, |
| 123 | apply_gated_attention=apply_gated_attention, |
| 124 | ) |
| 125 | for _ in range(num_layers) |
| 126 | ] |
| 127 | ) |
| 128 | |
| 129 | self.num_learnable_registers = num_learnable_registers |
| 130 | if self.num_learnable_registers: |
| 131 | self.learnable_registers = torch.nn.Parameter( |
| 132 | torch.rand(self.num_learnable_registers, self.inner_dim, dtype=torch.bfloat16) * 2.0 - 1.0 |
| 133 | ) |
| 134 | |
| 135 | def _replace_padded_with_learnable_registers( |
| 136 | self, hidden_states: torch.Tensor, attention_mask: torch.Tensor |
| 137 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 138 | assert hidden_states.shape[1] % self.num_learnable_registers == 0, ( |
| 139 | f"Hidden states sequence length {hidden_states.shape[1]} must be divisible by num_learnable_registers " |
| 140 | f"{self.num_learnable_registers}." |
| 141 | ) |
| 142 | |
| 143 | num_registers_duplications = hidden_states.shape[1] // self.num_learnable_registers |
| 144 | learnable_registers = torch.tile(self.learnable_registers, (num_registers_duplications, 1)) |
| 145 | valid_mask = attention_mask.squeeze(1).squeeze(1) >= -9000.0 |
| 146 | batch_size, sequence_length = valid_mask.shape |
| 147 | positions = torch.arange(sequence_length, device=hidden_states.device).unsqueeze(0) |
| 148 | valid_counts = valid_mask.sum(dim=1) |
| 149 | |
| 150 | # Gemma inputs are left padded. Rotate every sample independently so its |
| 151 | # valid tokens start at index zero, then fill the remaining positions |
| 152 | # with the same position-dependent register sequence used by B=1. |
| 153 | expected_mask = positions >= (sequence_length - valid_counts).unsqueeze(1) |
| 154 | if not torch.equal(valid_mask, expected_mask): |
| 155 | raise ValueError( |
| 156 | "Embeddings1DConnector expects a contiguous left-padded attention mask" |
| 157 | ) |
| 158 | source_positions = ( |
| 159 | positions + (sequence_length - valid_counts).unsqueeze(1) |
| 160 | ) % sequence_length |
| 161 | compacted_hidden_states = torch.gather( |
| 162 | hidden_states, |
| 163 | dim=1, |
| 164 | index=source_positions.unsqueeze(-1).expand(-1, -1, hidden_states.shape[-1]), |
| 165 | ) |
| 166 | output_valid_mask = positions < valid_counts.unsqueeze(1) |
| 167 | registers = learnable_registers.unsqueeze(0).expand(batch_size, -1, -1) |
| 168 | hidden_states = torch.where( |
| 169 | output_valid_mask.unsqueeze(-1), |
| 170 | compacted_hidden_states, |
| 171 | registers, |
| 172 | ) |
| 173 | |
| 174 | attention_mask = torch.full_like( |
| 175 | attention_mask, |
| 176 | 0.0, |
| 177 | dtype=attention_mask.dtype, |
| 178 | device=attention_mask.device, |
| 179 | ) |
| 180 | |
| 181 | return hidden_states, attention_mask |
| 182 | |
| 183 | def forward( |
| 184 | self, |
| 185 | hidden_states: torch.Tensor, |
| 186 | attention_mask: torch.Tensor | None = None, |
| 187 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 188 | """ |
| 189 | Forward pass of Embeddings1DConnector. |
| 190 | Args: |
| 191 | hidden_states (torch.Tensor): Input tensor of embeddings (shape [batch, seq_len, feature_dim]). |
| 192 | attention_mask (torch.Tensor|None): Optional mask for valid tokens (shape compatible with hidden_states). |
| 193 | Returns: |
| 194 | tuple[torch.Tensor, torch.Tensor]: Processed features and the corresponding (possibly modified) mask. |
| 195 | """ |
| 196 | if self.num_learnable_registers: |
| 197 | hidden_states, attention_mask = self._replace_padded_with_learnable_registers(hidden_states, attention_mask) |
| 198 | |
| 199 | indices_grid = torch.arange(hidden_states.shape[1], dtype=torch.float32, device=hidden_states.device) |
| 200 | indices_grid = indices_grid[None, None, :].expand(hidden_states.shape[0], -1, -1) |
| 201 | freq_grid_generator = generate_freq_grid_np if self.double_precision_rope else generate_freq_grid_pytorch |
| 202 | freqs_cis = precompute_freqs_cis( |
| 203 | indices_grid=indices_grid, |
| 204 | dim=self.inner_dim, |
| 205 | out_dtype=hidden_states.dtype, |
| 206 | theta=self.positional_embedding_theta, |
| 207 | max_pos=self.positional_embedding_max_pos, |
| 208 | num_attention_heads=self.num_attention_heads, |
| 209 | rope_type=self.rope_type, |
| 210 | freq_grid_generator=freq_grid_generator, |
| 211 | ) |
| 212 | |
| 213 | for block in self.transformer_1d_blocks: |
| 214 | hidden_states = block(hidden_states, attention_mask=attention_mask, pe=freqs_cis) |
| 215 | |
| 216 | hidden_states = rms_norm(hidden_states) |
| 217 | |
| 218 | return hidden_states, attention_mask |
| 219 | |
| 220 | |
| 221 | class Embeddings1DConnectorConfigurator(ModelConfigurator[Embeddings1DConnector]): |
| 222 | """Configurator for video embeddings connector.""" |
| 223 | |
| 224 | @classmethod |
| 225 | def from_config(cls: type[Embeddings1DConnector], config: dict) -> Embeddings1DConnector: |
| 226 | transformer_config = config.get("transformer", {}) |
| 227 | rope_type = LTXRopeType(transformer_config.get("rope_type", "interleaved")) |
| 228 | double_precision_rope = transformer_config.get("frequencies_precision", False) == "float64" |
| 229 | pe_max_pos = transformer_config.get("connector_positional_embedding_max_pos", [1]) |
| 230 | |
| 231 | # Video connector dimensions |
| 232 | num_attention_heads = transformer_config.get("connector_num_attention_heads", 30) |
| 233 | attention_head_dim = transformer_config.get("connector_attention_head_dim", 128) |
| 234 | num_layers = transformer_config.get("connector_num_layers", 2) |
| 235 | |
| 236 | connector = Embeddings1DConnector( |
| 237 | num_attention_heads=num_attention_heads, |
| 238 | attention_head_dim=attention_head_dim, |
| 239 | num_layers=num_layers, |
| 240 | positional_embedding_max_pos=pe_max_pos, |
| 241 | rope_type=rope_type, |
| 242 | double_precision_rope=double_precision_rope, |
| 243 | apply_gated_attention=transformer_config.get("connector_apply_gated_attention", False), |
| 244 | ) |
| 245 | return connector |
| 246 | |
| 247 | |
| 248 | class AudioEmbeddings1DConnectorConfigurator(ModelConfigurator[Embeddings1DConnector]): |
| 249 | """Configurator for audio embeddings connector with separate dimension settings.""" |
| 250 | |
| 251 | @classmethod |
| 252 | def from_config(cls: type[Embeddings1DConnector], config: dict) -> Embeddings1DConnector: |
| 253 | transformer_config = config.get("transformer", {}) |
| 254 | rope_type = LTXRopeType(transformer_config.get("rope_type", "interleaved")) |
| 255 | double_precision_rope = transformer_config.get("frequencies_precision", False) == "float64" |
| 256 | pe_max_pos = transformer_config.get("connector_positional_embedding_max_pos", [1]) |
| 257 | |
| 258 | # Audio connector dimensions - fall back to video connector config for backwards compatibility |
| 259 | num_attention_heads = transformer_config.get( |
| 260 | "audio_connector_num_attention_heads", |
| 261 | transformer_config.get("connector_num_attention_heads", 30), |
| 262 | ) |
| 263 | attention_head_dim = transformer_config.get( |
| 264 | "audio_connector_attention_head_dim", |
| 265 | transformer_config.get("connector_attention_head_dim", 128), |
| 266 | ) |
| 267 | num_layers = transformer_config.get( |
| 268 | "audio_connector_num_layers", |
| 269 | transformer_config.get("connector_num_layers", 2), |
| 270 | ) |
| 271 | |
| 272 | connector = Embeddings1DConnector( |
| 273 | num_attention_heads=num_attention_heads, |
| 274 | attention_head_dim=attention_head_dim, |
| 275 | num_layers=num_layers, |
| 276 | positional_embedding_max_pos=pe_max_pos, |
| 277 | rope_type=rope_type, |
| 278 | double_precision_rope=double_precision_rope, |
| 279 | apply_gated_attention=transformer_config.get("connector_apply_gated_attention", False), |
| 280 | ) |
| 281 | return connector |
| 282 |