| 1 | """ |
| 2 | VAE Wrappers for visualization and validation during DMD distillation. |
| 3 | """ |
| 4 | |
| 5 | from collections.abc import Iterator |
| 6 | from typing import Optional |
| 7 | import torch |
| 8 | import torch.nn as nn |
| 9 | |
| 10 | from ltx_core.loader.registry import Registry |
| 11 | from ltx_core.model.audio_vae import encode_audio |
| 12 | from ltx_core.model.video_vae.tiling import TilingConfig |
| 13 | from ltx_core.types import Audio |
| 14 | |
| 15 | |
| 16 | def _module_device_dtype(module: nn.Module) -> tuple[torch.device, torch.dtype]: |
| 17 | """ |
| 18 | Infer target device/dtype from module parameters or buffers. |
| 19 | """ |
| 20 | for tensor in module.parameters(): |
| 21 | return tensor.device, tensor.dtype |
| 22 | for tensor in module.buffers(): |
| 23 | return tensor.device, tensor.dtype |
| 24 | return torch.device("cpu"), torch.float32 |
| 25 | |
| 26 | |
| 27 | class VideoVAEWrapper(nn.Module): |
| 28 | """ |
| 29 | Wrapper for Video VAE encoder and decoder. |
| 30 | |
| 31 | Used for: |
| 32 | - Encoding videos to latent space (for visualization) |
| 33 | - Decoding latents to pixel space (for validation) |
| 34 | """ |
| 35 | |
| 36 | def __init__( |
| 37 | self, |
| 38 | encoder=None, |
| 39 | decoder=None, |
| 40 | device: torch.device = None, |
| 41 | dtype: torch.dtype = torch.bfloat16, |
| 42 | ): |
| 43 | """ |
| 44 | Args: |
| 45 | encoder: VideoEncoder instance (optional) |
| 46 | decoder: VideoDecoder instance |
| 47 | device: Target device |
| 48 | dtype: Model dtype |
| 49 | """ |
| 50 | super().__init__() |
| 51 | self.encoder = encoder |
| 52 | self.decoder = decoder |
| 53 | self.device = device |
| 54 | self.dtype = dtype |
| 55 | |
| 56 | @torch.no_grad() |
| 57 | def encode(self, video: torch.Tensor) -> torch.Tensor: |
| 58 | """ |
| 59 | Encode video to latent space. |
| 60 | |
| 61 | Args: |
| 62 | video: Pixel video [B, C, F, H, W] in range [-1, 1] |
| 63 | |
| 64 | Returns: |
| 65 | Latent [B, F', C_latent, H', W'] |
| 66 | """ |
| 67 | if self.encoder is None: |
| 68 | raise ValueError("Encoder not initialized") |
| 69 | |
| 70 | return self.encoder(video) |
| 71 | |
| 72 | @torch.no_grad() |
| 73 | def _prepare_decode_latent(self, latent: torch.Tensor) -> torch.Tensor: |
| 74 | if self.decoder is None: |
| 75 | raise ValueError("Decoder not initialized") |
| 76 | |
| 77 | # Decoder expects [B, C, F, H, W]. |
| 78 | # Our DMD code stores video as [B, F, C, H, W] where C=128. |
| 79 | # Detect this by checking if dim 2 (not dim 1) equals 128. |
| 80 | if latent.dim() == 5 and latent.shape[2] == 128: |
| 81 | # Input is [B, F, C, H, W], need to permute to [B, C, F, H, W] |
| 82 | latent = latent.permute(0, 2, 1, 3, 4) |
| 83 | |
| 84 | # Keep latent dtype/device consistent with decoder weights. |
| 85 | dec_device, dec_dtype = _module_device_dtype(self.decoder) |
| 86 | return latent.to(device=dec_device, dtype=dec_dtype) |
| 87 | |
| 88 | @torch.no_grad() |
| 89 | def decode(self, latent: torch.Tensor) -> torch.Tensor: |
| 90 | """ |
| 91 | Decode latent to pixel space. |
| 92 | |
| 93 | Args: |
| 94 | latent: Latent [B, F, C, H, W] |
| 95 | |
| 96 | Returns: |
| 97 | Video [B, C, F_out, H_out, W_out] in range [-1, 1] |
| 98 | """ |
| 99 | latent = self._prepare_decode_latent(latent) |
| 100 | |
| 101 | return self.decoder(latent) |
| 102 | |
| 103 | @torch.no_grad() |
| 104 | def decode_to_uint8_chunks( |
| 105 | self, |
| 106 | latent: torch.Tensor, |
| 107 | tiling_config: TilingConfig, |
| 108 | ) -> Iterator[torch.Tensor]: |
| 109 | """Yield decoded ``[F, H, W, C]`` chunks in CPU uint8 memory. |
| 110 | |
| 111 | Each tiled result is quantized and transferred immediately. This is |
| 112 | important on Windows/WDDM: collecting full-resolution floating-point |
| 113 | chunks on the GPU would recreate the memory pressure that tiling is |
| 114 | intended to avoid. |
| 115 | """ |
| 116 | latent = self._prepare_decode_latent(latent) |
| 117 | for video in self.decoder.tiled_decode(latent, tiling_config): |
| 118 | video_uint8 = (((video + 1) / 2).clamp(0, 1) * 255).to(torch.uint8) |
| 119 | video_uint8 = video_uint8[0].permute(1, 2, 3, 0) |
| 120 | yield video_uint8.cpu().contiguous() |
| 121 | |
| 122 | @torch.no_grad() |
| 123 | def decode_to_pixel(self, latent: torch.Tensor) -> torch.Tensor: |
| 124 | """ |
| 125 | Decode latent to pixel video for visualization. |
| 126 | |
| 127 | Args: |
| 128 | latent: Latent [B, F, C, H, W] |
| 129 | |
| 130 | Returns: |
| 131 | Video frames suitable for logging (normalized to [0, 1]) |
| 132 | """ |
| 133 | video = self.decode(latent) |
| 134 | # Normalize from [-1, 1] to [0, 1] |
| 135 | video = (video + 1) / 2 |
| 136 | video = video.clamp(0, 1) |
| 137 | return video |
| 138 | |
| 139 | |
| 140 | class AudioVAEWrapper(nn.Module): |
| 141 | """ |
| 142 | Wrapper for Audio VAE decoder and vocoder. |
| 143 | |
| 144 | Used for: |
| 145 | - Decoding audio latents to mel spectrogram |
| 146 | - Converting mel to waveform via vocoder |
| 147 | """ |
| 148 | |
| 149 | def __init__( |
| 150 | self, |
| 151 | encoder=None, |
| 152 | decoder=None, |
| 153 | vocoder=None, |
| 154 | device: torch.device = None, |
| 155 | dtype: torch.dtype = torch.bfloat16, |
| 156 | ): |
| 157 | """ |
| 158 | Args: |
| 159 | encoder: AudioEncoder instance (optional) |
| 160 | decoder: AudioDecoder instance |
| 161 | vocoder: Vocoder instance |
| 162 | device: Target device |
| 163 | dtype: Model dtype |
| 164 | """ |
| 165 | super().__init__() |
| 166 | self.encoder = encoder |
| 167 | self.decoder = decoder |
| 168 | self.vocoder = vocoder |
| 169 | self.device = device |
| 170 | self.dtype = dtype |
| 171 | |
| 172 | def get_output_sample_rate(self) -> Optional[int]: |
| 173 | """ |
| 174 | Return the vocoder waveform sample rate across 2.2/2.3 vocoder variants. |
| 175 | """ |
| 176 | if self.vocoder is None: |
| 177 | return None |
| 178 | |
| 179 | for attr in ("output_sample_rate", "output_sampling_rate"): |
| 180 | value = getattr(self.vocoder, attr, None) |
| 181 | if value is not None: |
| 182 | return int(value) |
| 183 | |
| 184 | return None |
| 185 | |
| 186 | @torch.no_grad() |
| 187 | def encode(self, waveform: torch.Tensor, sampling_rate: int) -> torch.Tensor: |
| 188 | """ |
| 189 | Encode waveform to transformer-format audio latents. |
| 190 | |
| 191 | Args: |
| 192 | waveform: Audio waveform [B, C, samples] or [C, samples] |
| 193 | sampling_rate: Input waveform sample rate |
| 194 | |
| 195 | Returns: |
| 196 | Audio latent [B, T, C_latent * mel_bins] |
| 197 | """ |
| 198 | if self.encoder is None: |
| 199 | raise ValueError("Audio encoder not initialized") |
| 200 | |
| 201 | if waveform.dim() == 2: |
| 202 | waveform = waveform.unsqueeze(0) |
| 203 | if waveform.dim() != 3: |
| 204 | raise ValueError(f"Expected waveform [B, C, samples] or [C, samples], got {tuple(waveform.shape)}") |
| 205 | |
| 206 | enc_device, _ = _module_device_dtype(self.encoder) |
| 207 | waveform = waveform.to(device=enc_device, dtype=torch.float32) |
| 208 | latent = encode_audio( |
| 209 | audio=Audio(waveform=waveform, sampling_rate=int(sampling_rate)), |
| 210 | audio_encoder=self.encoder, |
| 211 | ) |
| 212 | return latent.permute(0, 2, 1, 3).flatten(start_dim=2).contiguous() |
| 213 | |
| 214 | @torch.no_grad() |
| 215 | def decode(self, latent: torch.Tensor) -> torch.Tensor: |
| 216 | """ |
| 217 | Decode audio latent to mel spectrogram. |
| 218 | |
| 219 | The DMD pipeline produces audio latents in the transformer's sequence |
| 220 | format ``[B, T, C*F]`` (3D), but the ``AudioDecoder`` expects the VAE |
| 221 | spatial format ``[B, C, T, F]`` (4D). This method handles the |
| 222 | conversion automatically using the decoder's ``z_channels`` and |
| 223 | ``mel_bins`` attributes (set during checkpoint loading). |
| 224 | |
| 225 | Args: |
| 226 | latent: Audio latent, either ``[B, T, C*F]`` (transformer) or |
| 227 | ``[B, C, T, F]`` (VAE). |
| 228 | |
| 229 | Returns: |
| 230 | Mel spectrogram ``[B, out_ch, time, freq]``. |
| 231 | """ |
| 232 | if self.decoder is None: |
| 233 | raise ValueError("Decoder not initialized") |
| 234 | |
| 235 | # Reshape 3D transformer latent → 4D VAE latent when necessary. |
| 236 | # The transformer stores audio as [B, T, C*F] where C=z_channels and |
| 237 | # F=latent_mel_bins. The AudioDecoder expects [B, C, T, F]. |
| 238 | # Note: decoder.mel_bins is the *output* spectrogram size (e.g. 64), |
| 239 | # NOT the latent mel dimension. The latent mel dim = CF // z_channels. |
| 240 | if latent.dim() == 3: |
| 241 | B, T, CF = latent.shape |
| 242 | z_channels = getattr(self.decoder, "z_channels", None) |
| 243 | |
| 244 | if z_channels is not None: |
| 245 | latent_mel = CF // z_channels # e.g. 128 // 8 = 16 |
| 246 | # "b t (c f) -> b c t f" |
| 247 | latent = latent.reshape(B, T, z_channels, latent_mel).permute(0, 2, 1, 3) |
| 248 | else: |
| 249 | raise ValueError( |
| 250 | f"Cannot reshape 3D audio latent {latent.shape} to 4D: " |
| 251 | "decoder is missing z_channels attribute." |
| 252 | ) |
| 253 | |
| 254 | # Keep latent dtype/device consistent with decoder weights. |
| 255 | dec_device, dec_dtype = _module_device_dtype(self.decoder) |
| 256 | latent = latent.to(device=dec_device, dtype=dec_dtype) |
| 257 | |
| 258 | return self.decoder(latent) |
| 259 | |
| 260 | @torch.no_grad() |
| 261 | def decode_to_waveform(self, latent: torch.Tensor) -> torch.Tensor: |
| 262 | """ |
| 263 | Decode audio latent to waveform. |
| 264 | |
| 265 | Args: |
| 266 | latent: Audio latent [B, F, C] |
| 267 | |
| 268 | Returns: |
| 269 | Waveform [B, 1, samples] |
| 270 | """ |
| 271 | mel = self.decode(latent) |
| 272 | |
| 273 | if self.vocoder is None: |
| 274 | raise ValueError("Vocoder not initialized") |
| 275 | |
| 276 | return self.vocoder(mel) |
| 277 | |
| 278 | |
| 279 | def create_vae_wrappers( |
| 280 | checkpoint_path: str, |
| 281 | device: torch.device, |
| 282 | dtype: torch.dtype = torch.bfloat16, |
| 283 | with_video_encoder: bool = False, |
| 284 | with_audio_encoder: bool = False, |
| 285 | with_decoders: bool = True, |
| 286 | decoder_device: torch.device | None = None, |
| 287 | registry: Registry | None = None, |
| 288 | ) -> tuple[VideoVAEWrapper, AudioVAEWrapper]: |
| 289 | """ |
| 290 | Factory function to create VAE wrappers from checkpoint. |
| 291 | |
| 292 | Args: |
| 293 | checkpoint_path: Path to LTX-2 checkpoint |
| 294 | device: Target device |
| 295 | dtype: Model dtype |
| 296 | with_decoders: Load the video/audio decoders and vocoder. Conditioning-only |
| 297 | jobs can disable them to avoid loading unused weights. |
| 298 | decoder_device: Device for video/audio decoders and vocoder. Defaults to |
| 299 | ``device``; pass ``cpu`` during training init to avoid holding decode |
| 300 | modules on every rank when only encoders are needed. |
| 301 | |
| 302 | Returns: |
| 303 | Tuple of (VideoVAEWrapper, AudioVAEWrapper) |
| 304 | """ |
| 305 | from ltx_pipelines.utils.model_ledger import ModelLedger |
| 306 | |
| 307 | if decoder_device is None: |
| 308 | decoder_device = device |
| 309 | |
| 310 | # Load to CPU first to avoid safetensors device issues |
| 311 | ledger = ModelLedger( |
| 312 | dtype=dtype, |
| 313 | device=torch.device("cpu"), |
| 314 | checkpoint_path=checkpoint_path, |
| 315 | registry=registry, |
| 316 | ) |
| 317 | |
| 318 | video_encoder = ledger.video_encoder() if with_video_encoder else None |
| 319 | video_decoder = ledger.video_decoder() if with_decoders else None |
| 320 | audio_encoder = ledger.audio_encoder() if with_audio_encoder else None |
| 321 | audio_decoder = ledger.audio_decoder() if with_decoders else None |
| 322 | vocoder = ledger.vocoder() if with_decoders else None |
| 323 | |
| 324 | # Move to target device |
| 325 | if video_encoder is not None: |
| 326 | video_encoder = video_encoder.to(device=device, dtype=dtype) |
| 327 | if video_decoder is not None: |
| 328 | video_decoder = video_decoder.to(device=decoder_device, dtype=dtype) |
| 329 | if audio_encoder is not None: |
| 330 | audio_encoder = audio_encoder.to(device=device, dtype=torch.float32) |
| 331 | if audio_decoder is not None: |
| 332 | audio_decoder = audio_decoder.to(device=decoder_device, dtype=dtype) |
| 333 | if vocoder is not None: |
| 334 | vocoder = vocoder.to(device=decoder_device, dtype=dtype) |
| 335 | |
| 336 | video_vae = VideoVAEWrapper( |
| 337 | encoder=video_encoder, |
| 338 | decoder=video_decoder, |
| 339 | device=device, |
| 340 | dtype=dtype, |
| 341 | ) |
| 342 | |
| 343 | audio_vae = AudioVAEWrapper( |
| 344 | encoder=audio_encoder, |
| 345 | decoder=audio_decoder, |
| 346 | vocoder=vocoder, |
| 347 | device=device, |
| 348 | dtype=dtype, |
| 349 | ) |
| 350 | |
| 351 | return video_vae, audio_vae |
| 352 |