| 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 | enabled: bool |
| 27 | prompt_timestep: torch.Tensor | None = None |
| 28 | self_attention_mask: torch.Tensor | None = ( |
| 29 | None # Additive log-space self-attention bias (B, 1, T, T), None = full attention |
| 30 | ) |
| 31 | |
| 32 | |
| 33 | class TransformerArgsPreprocessor: |
| 34 | def __init__( # noqa: PLR0913 |
| 35 | self, |
| 36 | patchify_proj: torch.nn.Linear, |
| 37 | adaln: AdaLayerNormSingle, |
| 38 | inner_dim: int, |
| 39 | max_pos: list[int], |
| 40 | num_attention_heads: int, |
| 41 | use_middle_indices_grid: bool, |
| 42 | timestep_scale_multiplier: int, |
| 43 | double_precision_rope: bool, |
| 44 | positional_embedding_theta: float, |
| 45 | rope_type: LTXRopeType, |
| 46 | caption_projection: torch.nn.Module | None = None, |
| 47 | prompt_adaln: AdaLayerNormSingle | None = None, |
| 48 | ) -> None: |
| 49 | self.patchify_proj = patchify_proj |
| 50 | self.adaln = adaln |
| 51 | self.inner_dim = inner_dim |
| 52 | self.max_pos = max_pos |
| 53 | self.num_attention_heads = num_attention_heads |
| 54 | self.use_middle_indices_grid = use_middle_indices_grid |
| 55 | self.timestep_scale_multiplier = timestep_scale_multiplier |
| 56 | self.double_precision_rope = double_precision_rope |
| 57 | self.positional_embedding_theta = positional_embedding_theta |
| 58 | self.rope_type = rope_type |
| 59 | self.caption_projection = caption_projection |
| 60 | self.prompt_adaln = prompt_adaln |
| 61 | |
| 62 | def _prepare_timestep( |
| 63 | self, timestep: torch.Tensor, adaln: AdaLayerNormSingle, batch_size: int, hidden_dtype: torch.dtype |
| 64 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 65 | """Prepare timestep embeddings.""" |
| 66 | timestep_scaled = timestep * self.timestep_scale_multiplier |
| 67 | timestep, embedded_timestep = adaln( |
| 68 | timestep_scaled.flatten(), |
| 69 | hidden_dtype=hidden_dtype, |
| 70 | ) |
| 71 | # Second dimension is 1 or number of tokens (if timestep_per_token) |
| 72 | timestep = timestep.view(batch_size, -1, timestep.shape[-1]) |
| 73 | embedded_timestep = embedded_timestep.view(batch_size, -1, embedded_timestep.shape[-1]) |
| 74 | |
| 75 | return timestep, embedded_timestep |
| 76 | |
| 77 | def _prepare_context( |
| 78 | self, |
| 79 | context: torch.Tensor, |
| 80 | x: torch.Tensor, |
| 81 | ) -> torch.Tensor: |
| 82 | """Prepare context for transformer blocks.""" |
| 83 | if self.caption_projection is not None: |
| 84 | context = self.caption_projection(context) |
| 85 | batch_size = x.shape[0] |
| 86 | return context.view(batch_size, -1, x.shape[-1]) |
| 87 | |
| 88 | def _prepare_attention_mask(self, attention_mask: torch.Tensor | None, x_dtype: torch.dtype) -> torch.Tensor | None: |
| 89 | """Prepare attention mask.""" |
| 90 | if attention_mask is None or torch.is_floating_point(attention_mask): |
| 91 | return attention_mask |
| 92 | |
| 93 | return (attention_mask - 1).to(x_dtype).reshape( |
| 94 | (attention_mask.shape[0], 1, -1, attention_mask.shape[-1]) |
| 95 | ) * torch.finfo(x_dtype).max |
| 96 | |
| 97 | def _prepare_self_attention_mask( |
| 98 | self, attention_mask: torch.Tensor | None, x_dtype: torch.dtype |
| 99 | ) -> torch.Tensor | None: |
| 100 | """Prepare self-attention mask by converting [0,1] values to additive log-space bias. |
| 101 | Input shape: (B, T, T) with values in [0, 1]. |
| 102 | Output shape: (B, 1, T, T) with 0.0 for full attention and a large negative value |
| 103 | for masked positions. |
| 104 | Positions with attention_mask <= 0 are fully masked (mapped to the dtype's minimum |
| 105 | representable value). Strictly positive entries are converted via log-space for |
| 106 | smooth attenuation, with small values clamped for numerical stability. |
| 107 | Returns None if input is None (no masking). |
| 108 | """ |
| 109 | if attention_mask is None: |
| 110 | return None |
| 111 | |
| 112 | # Convert [0, 1] attention mask to additive log-space bias: |
| 113 | # 1.0 -> log(1.0) = 0.0 (no bias, full attention) |
| 114 | # 0.0 -> finfo.min (fully masked) |
| 115 | finfo = torch.finfo(x_dtype) |
| 116 | eps = finfo.tiny |
| 117 | |
| 118 | bias = torch.full_like(attention_mask, finfo.min, dtype=x_dtype) |
| 119 | positive = attention_mask > 0 |
| 120 | if positive.any(): |
| 121 | bias[positive] = torch.log(attention_mask[positive].clamp(min=eps)).to(x_dtype) |
| 122 | |
| 123 | return bias.unsqueeze(1) # (B, 1, T, T) for head broadcast |
| 124 | |
| 125 | def _prepare_positional_embeddings( |
| 126 | self, |
| 127 | positions: torch.Tensor, |
| 128 | inner_dim: int, |
| 129 | max_pos: list[int], |
| 130 | use_middle_indices_grid: bool, |
| 131 | num_attention_heads: int, |
| 132 | x_dtype: torch.dtype, |
| 133 | ) -> torch.Tensor: |
| 134 | """Prepare positional embeddings.""" |
| 135 | freq_grid_generator = generate_freq_grid_np if self.double_precision_rope else generate_freq_grid_pytorch |
| 136 | pe = precompute_freqs_cis( |
| 137 | positions, |
| 138 | dim=inner_dim, |
| 139 | out_dtype=x_dtype, |
| 140 | theta=self.positional_embedding_theta, |
| 141 | max_pos=max_pos, |
| 142 | use_middle_indices_grid=use_middle_indices_grid, |
| 143 | num_attention_heads=num_attention_heads, |
| 144 | rope_type=self.rope_type, |
| 145 | freq_grid_generator=freq_grid_generator, |
| 146 | ) |
| 147 | return pe |
| 148 | |
| 149 | def prepare( |
| 150 | self, |
| 151 | modality: Modality, |
| 152 | cross_modality: Modality | None = None, # noqa: ARG002 |
| 153 | ) -> TransformerArgs: |
| 154 | x = self.patchify_proj(modality.latent) |
| 155 | batch_size = x.shape[0] |
| 156 | timestep, embedded_timestep = self._prepare_timestep( |
| 157 | modality.timesteps, self.adaln, batch_size, modality.latent.dtype |
| 158 | ) |
| 159 | prompt_timestep = None |
| 160 | if self.prompt_adaln is not None: |
| 161 | prompt_timestep, _ = self._prepare_timestep( |
| 162 | modality.sigma, self.prompt_adaln, batch_size, modality.latent.dtype |
| 163 | ) |
| 164 | context = self._prepare_context(modality.context, x) |
| 165 | attention_mask = self._prepare_attention_mask(modality.context_mask, modality.latent.dtype) |
| 166 | pe = self._prepare_positional_embeddings( |
| 167 | positions=modality.positions, |
| 168 | inner_dim=self.inner_dim, |
| 169 | max_pos=self.max_pos, |
| 170 | use_middle_indices_grid=self.use_middle_indices_grid, |
| 171 | num_attention_heads=self.num_attention_heads, |
| 172 | x_dtype=modality.latent.dtype, |
| 173 | ) |
| 174 | self_attention_mask = self._prepare_self_attention_mask(modality.attention_mask, modality.latent.dtype) |
| 175 | return TransformerArgs( |
| 176 | x=x, |
| 177 | context=context, |
| 178 | context_mask=attention_mask, |
| 179 | timesteps=timestep, |
| 180 | embedded_timestep=embedded_timestep, |
| 181 | positional_embeddings=pe, |
| 182 | cross_positional_embeddings=None, |
| 183 | cross_scale_shift_timestep=None, |
| 184 | cross_gate_timestep=None, |
| 185 | enabled=modality.enabled, |
| 186 | prompt_timestep=prompt_timestep, |
| 187 | self_attention_mask=self_attention_mask, |
| 188 | ) |
| 189 | |
| 190 | |
| 191 | class MultiModalTransformerArgsPreprocessor: |
| 192 | def __init__( # noqa: PLR0913 |
| 193 | self, |
| 194 | patchify_proj: torch.nn.Linear, |
| 195 | adaln: AdaLayerNormSingle, |
| 196 | cross_scale_shift_adaln: AdaLayerNormSingle, |
| 197 | cross_gate_adaln: AdaLayerNormSingle, |
| 198 | inner_dim: int, |
| 199 | max_pos: list[int], |
| 200 | num_attention_heads: int, |
| 201 | cross_pe_max_pos: int, |
| 202 | use_middle_indices_grid: bool, |
| 203 | audio_cross_attention_dim: int, |
| 204 | timestep_scale_multiplier: int, |
| 205 | double_precision_rope: bool, |
| 206 | positional_embedding_theta: float, |
| 207 | rope_type: LTXRopeType, |
| 208 | av_ca_timestep_scale_multiplier: int, |
| 209 | caption_projection: torch.nn.Module | None = None, |
| 210 | prompt_adaln: AdaLayerNormSingle | None = None, |
| 211 | ) -> None: |
| 212 | self.simple_preprocessor = TransformerArgsPreprocessor( |
| 213 | patchify_proj=patchify_proj, |
| 214 | adaln=adaln, |
| 215 | inner_dim=inner_dim, |
| 216 | max_pos=max_pos, |
| 217 | num_attention_heads=num_attention_heads, |
| 218 | use_middle_indices_grid=use_middle_indices_grid, |
| 219 | timestep_scale_multiplier=timestep_scale_multiplier, |
| 220 | double_precision_rope=double_precision_rope, |
| 221 | positional_embedding_theta=positional_embedding_theta, |
| 222 | rope_type=rope_type, |
| 223 | caption_projection=caption_projection, |
| 224 | prompt_adaln=prompt_adaln, |
| 225 | ) |
| 226 | self.cross_scale_shift_adaln = cross_scale_shift_adaln |
| 227 | self.cross_gate_adaln = cross_gate_adaln |
| 228 | self.cross_pe_max_pos = cross_pe_max_pos |
| 229 | self.audio_cross_attention_dim = audio_cross_attention_dim |
| 230 | self.av_ca_timestep_scale_multiplier = av_ca_timestep_scale_multiplier |
| 231 | |
| 232 | def prepare( |
| 233 | self, |
| 234 | modality: Modality, |
| 235 | cross_modality: Modality | None = None, |
| 236 | ) -> TransformerArgs: |
| 237 | transformer_args = self.simple_preprocessor.prepare(modality) |
| 238 | if cross_modality is None: |
| 239 | return transformer_args |
| 240 | |
| 241 | if cross_modality.sigma.numel() > 1: |
| 242 | if cross_modality.sigma.shape[0] != modality.timesteps.shape[0]: |
| 243 | raise ValueError("Cross modality sigma must have the same batch size as the modality") |
| 244 | if cross_modality.sigma.ndim != 1: |
| 245 | raise ValueError("Cross modality sigma must be a 1D tensor") |
| 246 | |
| 247 | cross_timestep = cross_modality.sigma.view( |
| 248 | modality.timesteps.shape[0], 1, *[1] * len(modality.timesteps.shape[2:]) |
| 249 | ) |
| 250 | |
| 251 | cross_pe = self.simple_preprocessor._prepare_positional_embeddings( |
| 252 | positions=modality.positions[:, 0:1, :], |
| 253 | inner_dim=self.audio_cross_attention_dim, |
| 254 | max_pos=[self.cross_pe_max_pos], |
| 255 | use_middle_indices_grid=True, |
| 256 | num_attention_heads=self.simple_preprocessor.num_attention_heads, |
| 257 | x_dtype=modality.latent.dtype, |
| 258 | ) |
| 259 | |
| 260 | cross_scale_shift_timestep, cross_gate_timestep = self._prepare_cross_attention_timestep( |
| 261 | timestep=cross_timestep, |
| 262 | timestep_scale_multiplier=self.simple_preprocessor.timestep_scale_multiplier, |
| 263 | batch_size=transformer_args.x.shape[0], |
| 264 | hidden_dtype=modality.latent.dtype, |
| 265 | ) |
| 266 | return replace( |
| 267 | transformer_args, |
| 268 | cross_positional_embeddings=cross_pe, |
| 269 | cross_scale_shift_timestep=cross_scale_shift_timestep, |
| 270 | cross_gate_timestep=cross_gate_timestep, |
| 271 | ) |
| 272 | |
| 273 | def _prepare_cross_attention_timestep( |
| 274 | self, |
| 275 | timestep: torch.Tensor | None, |
| 276 | timestep_scale_multiplier: int, |
| 277 | batch_size: int, |
| 278 | hidden_dtype: torch.dtype, |
| 279 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 280 | """Prepare cross attention timestep embeddings.""" |
| 281 | timestep = timestep * timestep_scale_multiplier |
| 282 | |
| 283 | av_ca_factor = self.av_ca_timestep_scale_multiplier / timestep_scale_multiplier |
| 284 | |
| 285 | scale_shift_timestep, _ = self.cross_scale_shift_adaln( |
| 286 | timestep.flatten(), |
| 287 | hidden_dtype=hidden_dtype, |
| 288 | ) |
| 289 | scale_shift_timestep = scale_shift_timestep.view(batch_size, -1, scale_shift_timestep.shape[-1]) |
| 290 | gate_noise_timestep, _ = self.cross_gate_adaln( |
| 291 | timestep.flatten() * av_ca_factor, |
| 292 | hidden_dtype=hidden_dtype, |
| 293 | ) |
| 294 | gate_noise_timestep = gate_noise_timestep.view(batch_size, -1, gate_noise_timestep.shape[-1]) |
| 295 | |
| 296 | return scale_shift_timestep, gate_noise_timestep |
| 297 |