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