| 1 | from dataclasses import dataclass, replace |
| 2 | |
| 3 | import torch |
| 4 | |
| 5 | from ltx_core.model.transformer.adaln import AdaLayerNormSingle |
| 6 | from ltx_core.model.transformer.modality import Modality |
| 7 | from ltx_core.model.transformer.rope import ( |
| 8 | LTXRopeType, |
| 9 | generate_freq_grid_np, |
| 10 | generate_freq_grid_pytorch, |
| 11 | precompute_freqs_cis, |
| 12 | ) |
| 13 | |
| 14 | |
| 15 | @dataclass(frozen=True) |
| 16 | class TransformerArgs: |
| 17 | x: torch.Tensor |
| 18 | context: torch.Tensor |
| 19 | context_mask: torch.Tensor |
| 20 | timesteps: torch.Tensor |
| 21 | embedded_timestep: torch.Tensor |
| 22 | positional_embeddings: torch.Tensor |
| 23 | cross_positional_embeddings: torch.Tensor | None |
| 24 | cross_scale_shift_timestep: torch.Tensor | None |
| 25 | cross_gate_timestep: torch.Tensor | None |
| 26 | cross_attention_mask: torch.Tensor | None |
| 27 | cross_output_mask: torch.Tensor | None |
| 28 | late_cross_attention_mask: torch.Tensor | None |
| 29 | late_cross_output_mask: torch.Tensor | None |
| 30 | enabled: bool |
| 31 | prompt_timestep: torch.Tensor | None = None |
| 32 | self_attention_mask: torch.Tensor | None = ( |
| 33 | None # Additive log-space self-attention bias (B, 1, T, T), None = full attention |
| 34 | ) |
| 35 | late_self_attention_mask: torch.Tensor | None = ( |
| 36 | None # Optional alternate self-attention bias used in later transformer layers |
| 37 | ) |
| 38 | v2a_grad_scale: float = 1.0 |
| 39 | |
| 40 | |
| 41 | class TransformerArgsPreprocessor: |
| 42 | def __init__( # noqa: PLR0913 |
| 43 | self, |
| 44 | patchify_proj: torch.nn.Linear, |
| 45 | adaln: AdaLayerNormSingle, |
| 46 | inner_dim: int, |
| 47 | max_pos: list[int], |
| 48 | num_attention_heads: int, |
| 49 | use_middle_indices_grid: bool, |
| 50 | timestep_scale_multiplier: int, |
| 51 | double_precision_rope: bool, |
| 52 | positional_embedding_theta: float, |
| 53 | rope_type: LTXRopeType, |
| 54 | caption_projection: torch.nn.Module | None = None, |
| 55 | prompt_adaln: AdaLayerNormSingle | None = None, |
| 56 | ) -> None: |
| 57 | self.patchify_proj = patchify_proj |
| 58 | self.adaln = adaln |
| 59 | self.inner_dim = inner_dim |
| 60 | self.max_pos = max_pos |
| 61 | self.num_attention_heads = num_attention_heads |
| 62 | self.use_middle_indices_grid = use_middle_indices_grid |
| 63 | self.timestep_scale_multiplier = timestep_scale_multiplier |
| 64 | self.double_precision_rope = double_precision_rope |
| 65 | self.positional_embedding_theta = positional_embedding_theta |
| 66 | self.rope_type = rope_type |
| 67 | self.caption_projection = caption_projection |
| 68 | self.prompt_adaln = prompt_adaln |
| 69 | |
| 70 | def _prepare_timestep( |
| 71 | self, timestep: torch.Tensor, adaln: AdaLayerNormSingle, batch_size: int, hidden_dtype: torch.dtype |
| 72 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 73 | """Prepare timestep embeddings.""" |
| 74 | timestep_scaled = timestep * self.timestep_scale_multiplier |
| 75 | timestep, embedded_timestep = adaln( |
| 76 | timestep_scaled.flatten(), |
| 77 | hidden_dtype=hidden_dtype, |
| 78 | ) |
| 79 | # Second dimension is 1 or number of tokens (if timestep_per_token) |
| 80 | timestep = timestep.view(batch_size, -1, timestep.shape[-1]) |
| 81 | embedded_timestep = embedded_timestep.view(batch_size, -1, embedded_timestep.shape[-1]) |
| 82 | |
| 83 | return timestep, embedded_timestep |
| 84 | |
| 85 | def _prepare_context( |
| 86 | self, |
| 87 | context: torch.Tensor, |
| 88 | x: torch.Tensor, |
| 89 | ) -> torch.Tensor: |
| 90 | """Prepare context for transformer blocks.""" |
| 91 | if self.caption_projection is not None: |
| 92 | context = self.caption_projection(context) |
| 93 | batch_size = x.shape[0] |
| 94 | return context.view(batch_size, -1, x.shape[-1]) |
| 95 | |
| 96 | def _prepare_attention_mask(self, attention_mask: torch.Tensor | None, x_dtype: torch.dtype) -> torch.Tensor | None: |
| 97 | """Prepare attention mask.""" |
| 98 | if attention_mask is None or torch.is_floating_point(attention_mask): |
| 99 | return attention_mask |
| 100 | |
| 101 | return (attention_mask - 1).to(x_dtype).reshape( |
| 102 | (attention_mask.shape[0], 1, -1, attention_mask.shape[-1]) |
| 103 | ) * torch.finfo(x_dtype).max |
| 104 | |
| 105 | def _prepare_self_attention_mask( |
| 106 | self, attention_mask: torch.Tensor | None, x_dtype: torch.dtype |
| 107 | ) -> torch.Tensor | None: |
| 108 | """Prepare self-attention mask by converting [0,1] values to additive log-space bias. |
| 109 | Input shape: (B, T, T) with values in [0, 1]. |
| 110 | Output shape: (B, 1, T, T) with 0.0 for full attention and a large negative value |
| 111 | for masked positions. |
| 112 | Positions with attention_mask <= 0 are fully masked (mapped to the dtype's minimum |
| 113 | representable value). Strictly positive entries are converted via log-space for |
| 114 | smooth attenuation, with small values clamped for numerical stability. |
| 115 | Returns None if input is None (no masking). |
| 116 | """ |
| 117 | if attention_mask is None: |
| 118 | return None |
| 119 | |
| 120 | # Convert [0, 1] attention mask to additive log-space bias: |
| 121 | # 1.0 -> log(1.0) = 0.0 (no bias, full attention) |
| 122 | # 0.0 -> finfo.min (fully masked) |
| 123 | finfo = torch.finfo(x_dtype) |
| 124 | eps = finfo.tiny |
| 125 | |
| 126 | bias = torch.full_like(attention_mask, finfo.min, dtype=x_dtype) |
| 127 | positive = attention_mask > 0 |
| 128 | if positive.any(): |
| 129 | bias[positive] = torch.log(attention_mask[positive].clamp(min=eps)).to(x_dtype) |
| 130 | |
| 131 | return bias.unsqueeze(1) # (B, 1, T, T) for head broadcast |
| 132 | |
| 133 | def _prepare_cross_attention_mask(self, cross_kv_mask: torch.Tensor | None, x_dtype: torch.dtype) -> torch.Tensor | None: |
| 134 | """Prepare a key/value mask for audio-video cross-attention.""" |
| 135 | if cross_kv_mask is None: |
| 136 | return None |
| 137 | |
| 138 | mask = cross_kv_mask.to(x_dtype) |
| 139 | if mask.ndim == 2: |
| 140 | return (mask - 1).reshape((mask.shape[0], 1, 1, mask.shape[-1])) * torch.finfo(x_dtype).max |
| 141 | if mask.ndim == 3: |
| 142 | return (mask - 1).unsqueeze(1) * torch.finfo(x_dtype).max |
| 143 | raise ValueError(f"Expected cross_kv_mask shape (B, K) or (B, Q, K), got {tuple(mask.shape)}") |
| 144 | |
| 145 | @staticmethod |
| 146 | def _prepare_cross_output_mask(cross_query_mask: torch.Tensor | None, x_dtype: torch.dtype) -> torch.Tensor | None: |
| 147 | """Prepare a query/output mask for audio-video cross-attention.""" |
| 148 | if cross_query_mask is None: |
| 149 | return None |
| 150 | |
| 151 | mask = cross_query_mask |
| 152 | if mask.ndim == 2: |
| 153 | mask = mask.unsqueeze(-1) |
| 154 | return mask.to(x_dtype) |
| 155 | |
| 156 | def _prepare_positional_embeddings( |
| 157 | self, |
| 158 | positions: torch.Tensor, |
| 159 | inner_dim: int, |
| 160 | max_pos: list[int], |
| 161 | use_middle_indices_grid: bool, |
| 162 | num_attention_heads: int, |
| 163 | x_dtype: torch.dtype, |
| 164 | ) -> torch.Tensor: |
| 165 | """Prepare positional embeddings.""" |
| 166 | freq_grid_generator = generate_freq_grid_np if self.double_precision_rope else generate_freq_grid_pytorch |
| 167 | pe = precompute_freqs_cis( |
| 168 | positions, |
| 169 | dim=inner_dim, |
| 170 | out_dtype=x_dtype, |
| 171 | theta=self.positional_embedding_theta, |
| 172 | max_pos=max_pos, |
| 173 | use_middle_indices_grid=use_middle_indices_grid, |
| 174 | num_attention_heads=num_attention_heads, |
| 175 | rope_type=self.rope_type, |
| 176 | freq_grid_generator=freq_grid_generator, |
| 177 | ) |
| 178 | return pe |
| 179 | |
| 180 | def prepare( |
| 181 | self, |
| 182 | modality: Modality, |
| 183 | cross_modality: Modality | None = None, # noqa: ARG002 |
| 184 | ) -> TransformerArgs: |
| 185 | x = self.patchify_proj(modality.latent) |
| 186 | batch_size = x.shape[0] |
| 187 | timestep, embedded_timestep = self._prepare_timestep( |
| 188 | modality.timesteps, self.adaln, batch_size, modality.latent.dtype |
| 189 | ) |
| 190 | prompt_timestep = None |
| 191 | if self.prompt_adaln is not None: |
| 192 | prompt_timestep, _ = self._prepare_timestep( |
| 193 | modality.sigma, self.prompt_adaln, batch_size, modality.latent.dtype |
| 194 | ) |
| 195 | context = self._prepare_context(modality.context, x) |
| 196 | attention_mask = self._prepare_attention_mask(modality.context_mask, modality.latent.dtype) |
| 197 | pe = self._prepare_positional_embeddings( |
| 198 | positions=modality.positions, |
| 199 | inner_dim=self.inner_dim, |
| 200 | max_pos=self.max_pos, |
| 201 | use_middle_indices_grid=self.use_middle_indices_grid, |
| 202 | num_attention_heads=self.num_attention_heads, |
| 203 | x_dtype=modality.latent.dtype, |
| 204 | ) |
| 205 | self_attention_mask = self._prepare_self_attention_mask(modality.attention_mask, modality.latent.dtype) |
| 206 | late_self_attention_mask = self._prepare_self_attention_mask(modality.late_attention_mask, modality.latent.dtype) |
| 207 | return TransformerArgs( |
| 208 | x=x, |
| 209 | context=context, |
| 210 | context_mask=attention_mask, |
| 211 | timesteps=timestep, |
| 212 | embedded_timestep=embedded_timestep, |
| 213 | positional_embeddings=pe, |
| 214 | cross_positional_embeddings=None, |
| 215 | cross_scale_shift_timestep=None, |
| 216 | cross_gate_timestep=None, |
| 217 | cross_attention_mask=None, |
| 218 | cross_output_mask=None, |
| 219 | late_cross_attention_mask=None, |
| 220 | late_cross_output_mask=None, |
| 221 | enabled=modality.enabled, |
| 222 | prompt_timestep=prompt_timestep, |
| 223 | self_attention_mask=self_attention_mask, |
| 224 | late_self_attention_mask=late_self_attention_mask, |
| 225 | v2a_grad_scale=float(modality.v2a_grad_scale), |
| 226 | ) |
| 227 | |
| 228 | |
| 229 | class MultiModalTransformerArgsPreprocessor: |
| 230 | def __init__( # noqa: PLR0913 |
| 231 | self, |
| 232 | patchify_proj: torch.nn.Linear, |
| 233 | adaln: AdaLayerNormSingle, |
| 234 | cross_scale_shift_adaln: AdaLayerNormSingle, |
| 235 | cross_gate_adaln: AdaLayerNormSingle, |
| 236 | inner_dim: int, |
| 237 | max_pos: list[int], |
| 238 | num_attention_heads: int, |
| 239 | cross_pe_max_pos: int, |
| 240 | use_middle_indices_grid: bool, |
| 241 | audio_cross_attention_dim: int, |
| 242 | timestep_scale_multiplier: int, |
| 243 | double_precision_rope: bool, |
| 244 | positional_embedding_theta: float, |
| 245 | rope_type: LTXRopeType, |
| 246 | av_ca_timestep_scale_multiplier: int, |
| 247 | caption_projection: torch.nn.Module | None = None, |
| 248 | prompt_adaln: AdaLayerNormSingle | None = None, |
| 249 | ) -> None: |
| 250 | self.simple_preprocessor = TransformerArgsPreprocessor( |
| 251 | patchify_proj=patchify_proj, |
| 252 | adaln=adaln, |
| 253 | inner_dim=inner_dim, |
| 254 | max_pos=max_pos, |
| 255 | num_attention_heads=num_attention_heads, |
| 256 | use_middle_indices_grid=use_middle_indices_grid, |
| 257 | timestep_scale_multiplier=timestep_scale_multiplier, |
| 258 | double_precision_rope=double_precision_rope, |
| 259 | positional_embedding_theta=positional_embedding_theta, |
| 260 | rope_type=rope_type, |
| 261 | caption_projection=caption_projection, |
| 262 | prompt_adaln=prompt_adaln, |
| 263 | ) |
| 264 | self.cross_scale_shift_adaln = cross_scale_shift_adaln |
| 265 | self.cross_gate_adaln = cross_gate_adaln |
| 266 | self.cross_pe_max_pos = cross_pe_max_pos |
| 267 | self.audio_cross_attention_dim = audio_cross_attention_dim |
| 268 | self.av_ca_timestep_scale_multiplier = av_ca_timestep_scale_multiplier |
| 269 | |
| 270 | def prepare( |
| 271 | self, |
| 272 | modality: Modality, |
| 273 | cross_modality: Modality | None = None, |
| 274 | ) -> TransformerArgs: |
| 275 | transformer_args = self.simple_preprocessor.prepare(modality) |
| 276 | if cross_modality is None: |
| 277 | return transformer_args |
| 278 | |
| 279 | if cross_modality.sigma.numel() > 1: |
| 280 | if cross_modality.sigma.shape[0] != modality.timesteps.shape[0]: |
| 281 | raise ValueError("Cross modality sigma must have the same batch size as the modality") |
| 282 | if cross_modality.sigma.ndim != 1: |
| 283 | raise ValueError("Cross modality sigma must be a 1D tensor") |
| 284 | |
| 285 | cross_timestep = cross_modality.sigma.view( |
| 286 | modality.timesteps.shape[0], 1, *[1] * len(modality.timesteps.shape[2:]) |
| 287 | ) |
| 288 | |
| 289 | cross_pe = self.simple_preprocessor._prepare_positional_embeddings( |
| 290 | positions=modality.positions[:, 0:1, :], |
| 291 | inner_dim=self.audio_cross_attention_dim, |
| 292 | max_pos=[self.cross_pe_max_pos], |
| 293 | use_middle_indices_grid=True, |
| 294 | num_attention_heads=self.simple_preprocessor.num_attention_heads, |
| 295 | x_dtype=modality.latent.dtype, |
| 296 | ) |
| 297 | |
| 298 | cross_scale_shift_timestep, cross_gate_timestep = self._prepare_cross_attention_timestep( |
| 299 | timestep=cross_timestep, |
| 300 | timestep_scale_multiplier=self.simple_preprocessor.timestep_scale_multiplier, |
| 301 | batch_size=transformer_args.x.shape[0], |
| 302 | hidden_dtype=modality.latent.dtype, |
| 303 | ) |
| 304 | cross_attention_mask = self.simple_preprocessor._prepare_cross_attention_mask( |
| 305 | cross_modality.cross_kv_mask, |
| 306 | modality.latent.dtype, |
| 307 | ) |
| 308 | cross_output_mask = self.simple_preprocessor._prepare_cross_output_mask( |
| 309 | modality.cross_query_mask, |
| 310 | modality.latent.dtype, |
| 311 | ) |
| 312 | late_cross_attention_mask = self.simple_preprocessor._prepare_cross_attention_mask( |
| 313 | cross_modality.late_cross_kv_mask, |
| 314 | modality.latent.dtype, |
| 315 | ) |
| 316 | late_cross_output_mask = self.simple_preprocessor._prepare_cross_output_mask( |
| 317 | modality.late_cross_query_mask, |
| 318 | modality.latent.dtype, |
| 319 | ) |
| 320 | |
| 321 | return replace( |
| 322 | transformer_args, |
| 323 | cross_positional_embeddings=cross_pe, |
| 324 | cross_scale_shift_timestep=cross_scale_shift_timestep, |
| 325 | cross_gate_timestep=cross_gate_timestep, |
| 326 | cross_attention_mask=cross_attention_mask, |
| 327 | cross_output_mask=cross_output_mask, |
| 328 | late_cross_attention_mask=late_cross_attention_mask, |
| 329 | late_cross_output_mask=late_cross_output_mask, |
| 330 | ) |
| 331 | |
| 332 | def _prepare_cross_attention_timestep( |
| 333 | self, |
| 334 | timestep: torch.Tensor | None, |
| 335 | timestep_scale_multiplier: int, |
| 336 | batch_size: int, |
| 337 | hidden_dtype: torch.dtype, |
| 338 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 339 | """Prepare cross attention timestep embeddings.""" |
| 340 | timestep = timestep * timestep_scale_multiplier |
| 341 | |
| 342 | av_ca_factor = self.av_ca_timestep_scale_multiplier / timestep_scale_multiplier |
| 343 | |
| 344 | scale_shift_timestep, _ = self.cross_scale_shift_adaln( |
| 345 | timestep.flatten(), |
| 346 | hidden_dtype=hidden_dtype, |
| 347 | ) |
| 348 | scale_shift_timestep = scale_shift_timestep.view(batch_size, -1, scale_shift_timestep.shape[-1]) |
| 349 | gate_noise_timestep, _ = self.cross_gate_adaln( |
| 350 | timestep.flatten() * av_ca_factor, |
| 351 | hidden_dtype=hidden_dtype, |
| 352 | ) |
| 353 | gate_noise_timestep = gate_noise_timestep.view(batch_size, -1, gate_noise_timestep.shape[-1]) |
| 354 | |
| 355 | return scale_shift_timestep, gate_noise_timestep |
| 356 |