| 1 | import math |
| 2 | from typing import Optional, Tuple |
| 3 | |
| 4 | import einops |
| 5 | import torch |
| 6 | |
| 7 | from ltx_core.components.protocols import Patchifier |
| 8 | from ltx_core.types import AudioLatentShape, SpatioTemporalScaleFactors, VideoLatentShape |
| 9 | |
| 10 | |
| 11 | class VideoLatentPatchifier(Patchifier): |
| 12 | def __init__(self, patch_size: int): |
| 13 | # Patch sizes for video latents. |
| 14 | self._patch_size = ( |
| 15 | 1, # temporal dimension |
| 16 | patch_size, # height dimension |
| 17 | patch_size, # width dimension |
| 18 | ) |
| 19 | |
| 20 | @property |
| 21 | def patch_size(self) -> Tuple[int, int, int]: |
| 22 | return self._patch_size |
| 23 | |
| 24 | def get_token_count(self, tgt_shape: VideoLatentShape) -> int: |
| 25 | return math.prod(tgt_shape.to_torch_shape()[2:]) // math.prod(self._patch_size) |
| 26 | |
| 27 | def patchify( |
| 28 | self, |
| 29 | latents: torch.Tensor, |
| 30 | ) -> torch.Tensor: |
| 31 | latents = einops.rearrange( |
| 32 | latents, |
| 33 | "b c (f p1) (h p2) (w p3) -> b (f h w) (c p1 p2 p3)", |
| 34 | p1=self._patch_size[0], |
| 35 | p2=self._patch_size[1], |
| 36 | p3=self._patch_size[2], |
| 37 | ) |
| 38 | |
| 39 | return latents |
| 40 | |
| 41 | def unpatchify( |
| 42 | self, |
| 43 | latents: torch.Tensor, |
| 44 | output_shape: VideoLatentShape, |
| 45 | ) -> torch.Tensor: |
| 46 | assert self._patch_size[0] == 1, "Temporal patch size must be 1 for symmetric patchifier" |
| 47 | |
| 48 | patch_grid_frames = output_shape.frames // self._patch_size[0] |
| 49 | patch_grid_height = output_shape.height // self._patch_size[1] |
| 50 | patch_grid_width = output_shape.width // self._patch_size[2] |
| 51 | |
| 52 | latents = einops.rearrange( |
| 53 | latents, |
| 54 | "b (f h w) (c p q) -> b c f (h p) (w q)", |
| 55 | f=patch_grid_frames, |
| 56 | h=patch_grid_height, |
| 57 | w=patch_grid_width, |
| 58 | p=self._patch_size[1], |
| 59 | q=self._patch_size[2], |
| 60 | ) |
| 61 | |
| 62 | return latents |
| 63 | |
| 64 | def get_patch_grid_bounds( |
| 65 | self, |
| 66 | output_shape: AudioLatentShape | VideoLatentShape, |
| 67 | device: Optional[torch.device] = None, |
| 68 | ) -> torch.Tensor: |
| 69 | """ |
| 70 | Return the per-dimension bounds [inclusive start, exclusive end) for every |
| 71 | patch produced by `patchify`. The bounds are expressed in the original |
| 72 | video grid coordinates: frame/time, height, and width. |
| 73 | The resulting tensor is shaped `[batch_size, 3, num_patches, 2]`, where: |
| 74 | - axis 1 (size 3) enumerates (frame/time, height, width) dimensions |
| 75 | - axis 3 (size 2) stores `[start, end)` indices within each dimension |
| 76 | Args: |
| 77 | output_shape: Video grid description containing frames, height, and width. |
| 78 | device: Device of the latent tensor. |
| 79 | """ |
| 80 | if not isinstance(output_shape, VideoLatentShape): |
| 81 | raise ValueError("VideoLatentPatchifier expects VideoLatentShape when computing coordinates") |
| 82 | |
| 83 | frames = output_shape.frames |
| 84 | height = output_shape.height |
| 85 | width = output_shape.width |
| 86 | batch_size = output_shape.batch |
| 87 | |
| 88 | # Validate inputs to ensure positive dimensions |
| 89 | assert frames > 0, f"frames must be positive, got {frames}" |
| 90 | assert height > 0, f"height must be positive, got {height}" |
| 91 | assert width > 0, f"width must be positive, got {width}" |
| 92 | assert batch_size > 0, f"batch_size must be positive, got {batch_size}" |
| 93 | |
| 94 | # Generate grid coordinates for each dimension (frame, height, width) |
| 95 | # We use torch.arange to create the starting coordinates for each patch. |
| 96 | # indexing='ij' ensures the dimensions are in the order (frame, height, width). |
| 97 | grid_coords = torch.meshgrid( |
| 98 | torch.arange(start=0, end=frames, step=self._patch_size[0], device=device), |
| 99 | torch.arange(start=0, end=height, step=self._patch_size[1], device=device), |
| 100 | torch.arange(start=0, end=width, step=self._patch_size[2], device=device), |
| 101 | indexing="ij", |
| 102 | ) |
| 103 | |
| 104 | # Stack the grid coordinates to create the start coordinates tensor. |
| 105 | # Shape becomes (3, grid_f, grid_h, grid_w) |
| 106 | patch_starts = torch.stack(grid_coords, dim=0) |
| 107 | |
| 108 | # Create a tensor containing the size of a single patch: |
| 109 | # (frame_patch_size, height_patch_size, width_patch_size). |
| 110 | # Reshape to (3, 1, 1, 1) to enable broadcasting when adding to the start coordinates. |
| 111 | patch_size_delta = torch.tensor( |
| 112 | self._patch_size, |
| 113 | device=patch_starts.device, |
| 114 | dtype=patch_starts.dtype, |
| 115 | ).view(3, 1, 1, 1) |
| 116 | |
| 117 | # Calculate end coordinates: start + patch_size |
| 118 | # Shape becomes (3, grid_f, grid_h, grid_w) |
| 119 | patch_ends = patch_starts + patch_size_delta |
| 120 | |
| 121 | # Stack start and end coordinates together along the last dimension |
| 122 | # Shape becomes (3, grid_f, grid_h, grid_w, 2), where the last dimension is [start, end] |
| 123 | latent_coords = torch.stack((patch_starts, patch_ends), dim=-1) |
| 124 | |
| 125 | # Broadcast to batch size and flatten all spatial/temporal dimensions into one sequence. |
| 126 | # Final Shape: (batch_size, 3, num_patches, 2) |
| 127 | latent_coords = einops.repeat( |
| 128 | latent_coords, |
| 129 | "c f h w bounds -> b c (f h w) bounds", |
| 130 | b=batch_size, |
| 131 | bounds=2, |
| 132 | ) |
| 133 | |
| 134 | return latent_coords |
| 135 | |
| 136 | |
| 137 | def get_pixel_coords( |
| 138 | latent_coords: torch.Tensor, |
| 139 | scale_factors: SpatioTemporalScaleFactors, |
| 140 | causal_fix: bool = False, |
| 141 | ) -> torch.Tensor: |
| 142 | """ |
| 143 | Map latent-space `[start, end)` coordinates to their pixel-space equivalents by scaling |
| 144 | each axis (frame/time, height, width) with the corresponding VAE downsampling factors. |
| 145 | Optionally compensate for causal encoding that keeps the first frame at unit temporal scale. |
| 146 | Args: |
| 147 | latent_coords: Tensor of latent bounds shaped `(batch, 3, num_patches, 2)`. |
| 148 | scale_factors: SpatioTemporalScaleFactors tuple `(temporal, height, width)` with integer scale factors applied |
| 149 | per axis. |
| 150 | causal_fix: When True, rewrites the temporal axis of the first frame so causal VAEs |
| 151 | that treat frame zero differently still yield non-negative timestamps. |
| 152 | """ |
| 153 | # Broadcast the VAE scale factors so they align with the `(batch, axis, patch, bound)` layout. |
| 154 | broadcast_shape = [1] * latent_coords.ndim |
| 155 | broadcast_shape[1] = -1 # axis dimension corresponds to (frame/time, height, width) |
| 156 | scale_tensor = torch.tensor(scale_factors, device=latent_coords.device).view(*broadcast_shape) |
| 157 | |
| 158 | # Apply per-axis scaling to convert latent bounds into pixel-space coordinates. |
| 159 | pixel_coords = latent_coords * scale_tensor |
| 160 | |
| 161 | if causal_fix: |
| 162 | # VAE temporal stride for the very first frame is 1 instead of `scale_factors[0]`. |
| 163 | # Shift and clamp to keep the first-frame timestamps causal and non-negative. |
| 164 | pixel_coords[:, 0, ...] = (pixel_coords[:, 0, ...] + 1 - scale_factors[0]).clamp(min=0) |
| 165 | |
| 166 | return pixel_coords |
| 167 | |
| 168 | |
| 169 | class AudioPatchifier(Patchifier): |
| 170 | def __init__( |
| 171 | self, |
| 172 | patch_size: int, |
| 173 | sample_rate: int = 16000, |
| 174 | hop_length: int = 160, |
| 175 | audio_latent_downsample_factor: int = 4, |
| 176 | is_causal: bool = True, |
| 177 | shift: int = 0, |
| 178 | ): |
| 179 | """ |
| 180 | Patchifier tailored for spectrogram/audio latents. |
| 181 | Args: |
| 182 | patch_size: Number of mel bins combined into a single patch. This |
| 183 | controls the resolution along the frequency axis. |
| 184 | sample_rate: Original waveform sampling rate. Used to map latent |
| 185 | indices back to seconds so downstream consumers can align audio |
| 186 | and video cues. |
| 187 | hop_length: Window hop length used for the spectrogram. Determines |
| 188 | how many real-time samples separate two consecutive latent frames. |
| 189 | audio_latent_downsample_factor: Ratio between spectrogram frames and |
| 190 | latent frames; compensates for additional downsampling inside the |
| 191 | VAE encoder. |
| 192 | is_causal: When True, timing is shifted to account for causal |
| 193 | receptive fields so timestamps do not peek into the future. |
| 194 | shift: Integer offset applied to the latent indices. Enables |
| 195 | constructing overlapping windows from the same latent sequence. |
| 196 | """ |
| 197 | self.hop_length = hop_length |
| 198 | self.sample_rate = sample_rate |
| 199 | self.audio_latent_downsample_factor = audio_latent_downsample_factor |
| 200 | self.is_causal = is_causal |
| 201 | self.shift = shift |
| 202 | self._patch_size = (1, patch_size, patch_size) |
| 203 | |
| 204 | @property |
| 205 | def patch_size(self) -> Tuple[int, int, int]: |
| 206 | return self._patch_size |
| 207 | |
| 208 | def get_token_count(self, tgt_shape: AudioLatentShape) -> int: |
| 209 | return tgt_shape.frames |
| 210 | |
| 211 | def _get_audio_latent_time_in_sec( |
| 212 | self, |
| 213 | start_latent: int, |
| 214 | end_latent: int, |
| 215 | dtype: torch.dtype, |
| 216 | device: Optional[torch.device] = None, |
| 217 | ) -> torch.Tensor: |
| 218 | """ |
| 219 | Converts latent indices into real-time seconds while honoring causal |
| 220 | offsets and the configured hop length. |
| 221 | Args: |
| 222 | start_latent: Inclusive start index inside the latent sequence. This |
| 223 | sets the first timestamp returned. |
| 224 | end_latent: Exclusive end index. Determines how many timestamps get |
| 225 | generated. |
| 226 | dtype: Floating-point dtype used for the returned tensor, allowing |
| 227 | callers to control precision. |
| 228 | device: Target device for the timestamp tensor. When omitted the |
| 229 | computation occurs on CPU to avoid surprising GPU allocations. |
| 230 | """ |
| 231 | if device is None: |
| 232 | device = torch.device("cpu") |
| 233 | |
| 234 | audio_latent_frame = torch.arange(start_latent, end_latent, dtype=dtype, device=device) |
| 235 | |
| 236 | audio_mel_frame = audio_latent_frame * self.audio_latent_downsample_factor |
| 237 | |
| 238 | if self.is_causal: |
| 239 | # Frame offset for causal alignment. |
| 240 | # The "+1" ensures the timestamp corresponds to the first sample that is fully available. |
| 241 | causal_offset = 1 |
| 242 | audio_mel_frame = (audio_mel_frame + causal_offset - self.audio_latent_downsample_factor).clip(min=0) |
| 243 | |
| 244 | return audio_mel_frame * self.hop_length / self.sample_rate |
| 245 | |
| 246 | def _compute_audio_timings( |
| 247 | self, |
| 248 | batch_size: int, |
| 249 | num_steps: int, |
| 250 | device: Optional[torch.device] = None, |
| 251 | ) -> torch.Tensor: |
| 252 | """ |
| 253 | Builds a `(B, 1, T, 2)` tensor containing timestamps for each latent frame. |
| 254 | This helper method underpins `get_patch_grid_bounds` for the audio patchifier. |
| 255 | Args: |
| 256 | batch_size: Number of sequences to broadcast the timings over. |
| 257 | num_steps: Number of latent frames (time steps) to convert into timestamps. |
| 258 | device: Device on which the resulting tensor should reside. |
| 259 | """ |
| 260 | resolved_device = device |
| 261 | if resolved_device is None: |
| 262 | resolved_device = torch.device("cpu") |
| 263 | |
| 264 | start_timings = self._get_audio_latent_time_in_sec( |
| 265 | self.shift, |
| 266 | num_steps + self.shift, |
| 267 | torch.float32, |
| 268 | resolved_device, |
| 269 | ) |
| 270 | start_timings = start_timings.unsqueeze(0).expand(batch_size, -1).unsqueeze(1) |
| 271 | |
| 272 | end_timings = self._get_audio_latent_time_in_sec( |
| 273 | self.shift + 1, |
| 274 | num_steps + self.shift + 1, |
| 275 | torch.float32, |
| 276 | resolved_device, |
| 277 | ) |
| 278 | end_timings = end_timings.unsqueeze(0).expand(batch_size, -1).unsqueeze(1) |
| 279 | |
| 280 | return torch.stack([start_timings, end_timings], dim=-1) |
| 281 | |
| 282 | def patchify( |
| 283 | self, |
| 284 | audio_latents: torch.Tensor, |
| 285 | ) -> torch.Tensor: |
| 286 | """ |
| 287 | Flattens the audio latent tensor along time. Use `get_patch_grid_bounds` |
| 288 | to derive timestamps for each latent frame based on the configured hop |
| 289 | length and downsampling. |
| 290 | Args: |
| 291 | audio_latents: Latent tensor to patchify. |
| 292 | Returns: |
| 293 | Flattened patch tokens tensor. Use `get_patch_grid_bounds` to compute the |
| 294 | corresponding timing metadata when needed. |
| 295 | """ |
| 296 | audio_latents = einops.rearrange( |
| 297 | audio_latents, |
| 298 | "b c t f -> b t (c f)", |
| 299 | ) |
| 300 | |
| 301 | return audio_latents |
| 302 | |
| 303 | def unpatchify( |
| 304 | self, |
| 305 | audio_latents: torch.Tensor, |
| 306 | output_shape: AudioLatentShape, |
| 307 | ) -> torch.Tensor: |
| 308 | """ |
| 309 | Restores the `(B, C, T, F)` spectrogram tensor from flattened patches. |
| 310 | Use `get_patch_grid_bounds` to recompute the timestamps that describe each |
| 311 | frame's position in real time. |
| 312 | Args: |
| 313 | audio_latents: Latent tensor to unpatchify. |
| 314 | output_shape: Shape of the unpatched output tensor. |
| 315 | Returns: |
| 316 | Unpatched latent tensor. Use `get_patch_grid_bounds` to compute the timing |
| 317 | metadata associated with the restored latents. |
| 318 | """ |
| 319 | # audio_latents shape: (batch, time, freq * channels) |
| 320 | audio_latents = einops.rearrange( |
| 321 | audio_latents, |
| 322 | "b t (c f) -> b c t f", |
| 323 | c=output_shape.channels, |
| 324 | f=output_shape.mel_bins, |
| 325 | ) |
| 326 | |
| 327 | return audio_latents |
| 328 | |
| 329 | def get_patch_grid_bounds( |
| 330 | self, |
| 331 | output_shape: AudioLatentShape | VideoLatentShape, |
| 332 | device: Optional[torch.device] = None, |
| 333 | ) -> torch.Tensor: |
| 334 | """ |
| 335 | Return the temporal bounds `[inclusive start, exclusive end)` for every |
| 336 | patch emitted by `patchify`. For audio this corresponds to timestamps in |
| 337 | seconds aligned with the original spectrogram grid. |
| 338 | The returned tensor has shape `[batch_size, 1, time_steps, 2]`, where: |
| 339 | - axis 1 (size 1) represents the temporal dimension |
| 340 | - axis 3 (size 2) stores the `[start, end)` timestamps per patch |
| 341 | Args: |
| 342 | output_shape: Audio grid specification describing the number of time steps. |
| 343 | device: Target device for the returned tensor. |
| 344 | """ |
| 345 | if not isinstance(output_shape, AudioLatentShape): |
| 346 | raise ValueError("AudioPatchifier expects AudioLatentShape when computing coordinates") |
| 347 | |
| 348 | return self._compute_audio_timings(output_shape.batch, output_shape.frames, device) |
| 349 |