返回 JoyAI-Echo
audio_vae.py
根目录 / ltx-core / src / ltx_core / model / audio_vae / audio_vae.py
1 from typing import Set, Tuple
2
3 import torch
4 import torch.nn.functional as F
5
6 from ltx_core.components.patchifiers import AudioPatchifier
7 from ltx_core.model.audio_vae.attention import AttentionType, make_attn
8 from ltx_core.model.audio_vae.causal_conv_2d import make_conv2d
9 from ltx_core.model.audio_vae.causality_axis import CausalityAxis
10 from ltx_core.model.audio_vae.downsample import build_downsampling_path
11 from ltx_core.model.audio_vae.ops import AudioProcessor, PerChannelStatistics
12 from ltx_core.model.audio_vae.resnet import ResnetBlock
13 from ltx_core.model.audio_vae.upsample import build_upsampling_path
14 from ltx_core.model.audio_vae.vocoder import Vocoder
15 from ltx_core.model.common.normalization import NormType, build_normalization_layer
16 from ltx_core.types import Audio, AudioLatentShape
17
18 LATENT_DOWNSAMPLE_FACTOR = 4
19
20
21 def build_mid_block(
22 channels: int,
23 temb_channels: int,
24 dropout: float,
25 norm_type: NormType,
26 causality_axis: CausalityAxis,
27 attn_type: AttentionType,
28 add_attention: bool,
29 ) -> torch.nn.Module:
30 """Build the middle block with two ResNet blocks and optional attention."""
31 mid = torch.nn.Module()
32 mid.block_1 = ResnetBlock(
33 in_channels=channels,
34 out_channels=channels,
35 temb_channels=temb_channels,
36 dropout=dropout,
37 norm_type=norm_type,
38 causality_axis=causality_axis,
39 )
40 mid.attn_1 = make_attn(channels, attn_type=attn_type, norm_type=norm_type) if add_attention else torch.nn.Identity()
41 mid.block_2 = ResnetBlock(
42 in_channels=channels,
43 out_channels=channels,
44 temb_channels=temb_channels,
45 dropout=dropout,
46 norm_type=norm_type,
47 causality_axis=causality_axis,
48 )
49 return mid
50
51
52 def run_mid_block(mid: torch.nn.Module, features: torch.Tensor) -> torch.Tensor:
53 """Run features through the middle block."""
54 features = mid.block_1(features, temb=None)
55 features = mid.attn_1(features)
56 return mid.block_2(features, temb=None)
57
58
59 class AudioEncoder(torch.nn.Module):
60 """
61 Encoder that compresses audio spectrograms into latent representations.
62 The encoder uses a series of downsampling blocks with residual connections,
63 attention mechanisms, and configurable causal convolutions.
64 """
65
66 def __init__( # noqa: PLR0913
67 self,
68 *,
69 ch: int,
70 ch_mult: Tuple[int, ...] = (1, 2, 4, 8),
71 num_res_blocks: int,
72 attn_resolutions: Set[int],
73 dropout: float = 0.0,
74 resamp_with_conv: bool = True,
75 in_channels: int,
76 resolution: int,
77 z_channels: int,
78 double_z: bool = True,
79 attn_type: AttentionType = AttentionType.VANILLA,
80 mid_block_add_attention: bool = True,
81 norm_type: NormType = NormType.GROUP,
82 causality_axis: CausalityAxis = CausalityAxis.WIDTH,
83 sample_rate: int = 16000,
84 mel_hop_length: int = 160,
85 n_fft: int = 1024,
86 is_causal: bool = True,
87 mel_bins: int = 64,
88 **_ignore_kwargs,
89 ) -> None:
90 """
91 Initialize the Encoder.
92 Args:
93 Arguments are configuration parameters, loaded from the audio VAE checkpoint config
94 (audio_vae.model.params.ddconfig):
95 ch: Base number of feature channels used in the first convolution layer.
96 ch_mult: Multiplicative factors for the number of channels at each resolution level.
97 num_res_blocks: Number of residual blocks to use at each resolution level.
98 attn_resolutions: Spatial resolutions (e.g., in time/frequency) at which to apply attention.
99 resolution: Input spatial resolution of the spectrogram (height, width).
100 z_channels: Number of channels in the latent representation.
101 norm_type: Normalization layer type to use within the network (e.g., group, batch).
102 causality_axis: Axis along which convolutions should be causal (e.g., time axis).
103 sample_rate: Audio sample rate in Hz for the input signals.
104 mel_hop_length: Hop length used when computing the mel spectrogram.
105 n_fft: FFT size used to compute the spectrogram.
106 mel_bins: Number of mel-frequency bins in the input spectrogram.
107 in_channels: Number of channels in the input spectrogram tensor.
108 double_z: If True, predict both mean and log-variance (doubling latent channels).
109 is_causal: If True, use causal convolutions suitable for streaming setups.
110 dropout: Dropout probability used in residual and mid blocks.
111 attn_type: Type of attention mechanism to use in attention blocks.
112 resamp_with_conv: If True, perform resolution changes using strided convolutions.
113 mid_block_add_attention: If True, add an attention block in the mid-level of the encoder.
114 """
115 super().__init__()
116
117 self.per_channel_statistics = PerChannelStatistics(latent_channels=ch)
118 self.sample_rate = sample_rate
119 self.mel_hop_length = mel_hop_length
120 self.n_fft = n_fft
121 self.is_causal = is_causal
122 self.mel_bins = mel_bins
123
124 self.patchifier = AudioPatchifier(
125 patch_size=1,
126 audio_latent_downsample_factor=LATENT_DOWNSAMPLE_FACTOR,
127 sample_rate=sample_rate,
128 hop_length=mel_hop_length,
129 is_causal=is_causal,
130 )
131
132 self.ch = ch
133 self.temb_ch = 0
134 self.num_resolutions = len(ch_mult)
135 self.num_res_blocks = num_res_blocks
136 self.resolution = resolution
137 self.in_channels = in_channels
138 self.z_channels = z_channels
139 self.double_z = double_z
140 self.norm_type = norm_type
141 self.causality_axis = causality_axis
142 self.attn_type = attn_type
143
144 # downsampling
145 self.conv_in = make_conv2d(
146 in_channels,
147 self.ch,
148 kernel_size=3,
149 stride=1,
150 causality_axis=self.causality_axis,
151 )
152
153 self.non_linearity = torch.nn.SiLU()
154
155 self.down, block_in = build_downsampling_path(
156 ch=ch,
157 ch_mult=ch_mult,
158 num_resolutions=self.num_resolutions,
159 num_res_blocks=num_res_blocks,
160 resolution=resolution,
161 temb_channels=self.temb_ch,
162 dropout=dropout,
163 norm_type=self.norm_type,
164 causality_axis=self.causality_axis,
165 attn_type=self.attn_type,
166 attn_resolutions=attn_resolutions,
167 resamp_with_conv=resamp_with_conv,
168 )
169
170 self.mid = build_mid_block(
171 channels=block_in,
172 temb_channels=self.temb_ch,
173 dropout=dropout,
174 norm_type=self.norm_type,
175 causality_axis=self.causality_axis,
176 attn_type=self.attn_type,
177 add_attention=mid_block_add_attention,
178 )
179
180 self.norm_out = build_normalization_layer(block_in, normtype=self.norm_type)
181 self.conv_out = make_conv2d(
182 block_in,
183 2 * z_channels if double_z else z_channels,
184 kernel_size=3,
185 stride=1,
186 causality_axis=self.causality_axis,
187 )
188
189 def forward(self, spectrogram: torch.Tensor) -> torch.Tensor:
190 """
191 Encode audio spectrogram into latent representations.
192 Args:
193 spectrogram: Input spectrogram of shape (batch, channels, time, frequency)
194 Returns:
195 Encoded latent representation of shape (batch, channels, frames, mel_bins)
196 """
197 h = self.conv_in(spectrogram)
198 h = self._run_downsampling_path(h)
199 h = run_mid_block(self.mid, h)
200 h = self._finalize_output(h)
201
202 return self._normalize_latents(h)
203
204 def _run_downsampling_path(self, h: torch.Tensor) -> torch.Tensor:
205 for level in range(self.num_resolutions):
206 stage = self.down[level]
207 for block_idx in range(self.num_res_blocks):
208 h = stage.block[block_idx](h, temb=None)
209 if stage.attn:
210 h = stage.attn[block_idx](h)
211
212 if level != self.num_resolutions - 1:
213 h = stage.downsample(h)
214
215 return h
216
217 def _finalize_output(self, h: torch.Tensor) -> torch.Tensor:
218 h = self.norm_out(h)
219 h = self.non_linearity(h)
220 return self.conv_out(h)
221
222 def _normalize_latents(self, latent_output: torch.Tensor) -> torch.Tensor:
223 """
224 Normalize encoder latents using per-channel statistics.
225 When the encoder is configured with ``double_z=True``, the final
226 convolution produces twice the number of latent channels, typically
227 interpreted as two concatenated tensors along the channel dimension
228 (e.g., mean and variance or other auxiliary parameters).
229 This method intentionally uses only the first half of the channels
230 (the "mean" component) as input to the patchifier and normalization
231 logic. The remaining channels are left unchanged by this method and
232 are expected to be consumed elsewhere in the VAE pipeline.
233 If ``double_z=False``, the encoder output already contains only the
234 mean latents and the chunking operation simply returns that tensor.
235 """
236 means = torch.chunk(latent_output, 2, dim=1)[0]
237 latent_shape = AudioLatentShape(
238 batch=means.shape[0],
239 channels=means.shape[1],
240 frames=means.shape[2],
241 mel_bins=means.shape[3],
242 )
243 latent_patched = self.patchifier.patchify(means)
244 latent_normalized = self.per_channel_statistics.normalize(latent_patched)
245 return self.patchifier.unpatchify(latent_normalized, latent_shape)
246
247
248 def encode_audio(
249 audio: Audio,
250 audio_encoder: AudioEncoder,
251 audio_processor: AudioProcessor | None = None,
252 ) -> torch.Tensor:
253 """Encode audio waveform into latent representation.
254 Args:
255 audio: Audio container with waveform tensor of shape (batch, channels, samples) and sampling rate.
256 audio_encoder: Audio encoder model
257 audio_processor: Audio processor model (optional, if not provided, it will be created from the audio encoder)
258 """
259 dtype = next(audio_encoder.parameters()).dtype
260 device = next(audio_encoder.parameters()).device
261
262 if audio_processor is None:
263 audio_processor = AudioProcessor(
264 target_sample_rate=audio_encoder.sample_rate,
265 mel_bins=audio_encoder.mel_bins,
266 mel_hop_length=audio_encoder.mel_hop_length,
267 n_fft=audio_encoder.n_fft,
268 ).to(device=device)
269
270 mel_spectrogram = audio_processor.waveform_to_mel(audio.to(device=device))
271
272 latent = audio_encoder(mel_spectrogram.to(dtype=dtype))
273 return latent
274
275
276 class AudioDecoder(torch.nn.Module):
277 """
278 Symmetric decoder that reconstructs audio spectrograms from latent features.
279 The decoder mirrors the encoder structure with configurable channel multipliers,
280 attention resolutions, and causal convolutions.
281 """
282
283 def __init__( # noqa: PLR0913
284 self,
285 *,
286 ch: int,
287 out_ch: int,
288 ch_mult: Tuple[int, ...] = (1, 2, 4, 8),
289 num_res_blocks: int,
290 attn_resolutions: Set[int],
291 resolution: int,
292 z_channels: int,
293 norm_type: NormType = NormType.GROUP,
294 causality_axis: CausalityAxis = CausalityAxis.WIDTH,
295 dropout: float = 0.0,
296 mid_block_add_attention: bool = True,
297 sample_rate: int = 16000,
298 mel_hop_length: int = 160,
299 is_causal: bool = True,
300 mel_bins: int | None = None,
301 ) -> None:
302 """
303 Initialize the Decoder.
304 Args:
305 Arguments are configuration parameters, loaded from the audio VAE checkpoint config
306 (audio_vae.model.params.ddconfig):
307 - ch, out_ch, ch_mult, num_res_blocks, attn_resolutions
308 - resolution, z_channels
309 - norm_type, causality_axis
310 """
311 super().__init__()
312
313 # Internal behavioural defaults that are not driven by the checkpoint.
314 resamp_with_conv = True
315 attn_type = AttentionType.VANILLA
316
317 # Per-channel statistics for denormalizing latents
318 self.per_channel_statistics = PerChannelStatistics(latent_channels=ch)
319 self.sample_rate = sample_rate
320 self.mel_hop_length = mel_hop_length
321 self.is_causal = is_causal
322 self.mel_bins = mel_bins
323 self.patchifier = AudioPatchifier(
324 patch_size=1,
325 audio_latent_downsample_factor=LATENT_DOWNSAMPLE_FACTOR,
326 sample_rate=sample_rate,
327 hop_length=mel_hop_length,
328 is_causal=is_causal,
329 )
330
331 self.ch = ch
332 self.temb_ch = 0
333 self.num_resolutions = len(ch_mult)
334 self.num_res_blocks = num_res_blocks
335 self.resolution = resolution
336 self.out_ch = out_ch
337 self.give_pre_end = False
338 self.tanh_out = False
339 self.norm_type = norm_type
340 self.z_channels = z_channels
341 self.channel_multipliers = ch_mult
342 self.attn_resolutions = attn_resolutions
343 self.causality_axis = causality_axis
344 self.attn_type = attn_type
345
346 base_block_channels = ch * self.channel_multipliers[-1]
347 base_resolution = resolution // (2 ** (self.num_resolutions - 1))
348 self.z_shape = (1, z_channels, base_resolution, base_resolution)
349
350 self.conv_in = make_conv2d(
351 z_channels, base_block_channels, kernel_size=3, stride=1, causality_axis=self.causality_axis
352 )
353 self.non_linearity = torch.nn.SiLU()
354 self.mid = build_mid_block(
355 channels=base_block_channels,
356 temb_channels=self.temb_ch,
357 dropout=dropout,
358 norm_type=self.norm_type,
359 causality_axis=self.causality_axis,
360 attn_type=self.attn_type,
361 add_attention=mid_block_add_attention,
362 )
363 self.up, final_block_channels = build_upsampling_path(
364 ch=ch,
365 ch_mult=ch_mult,
366 num_resolutions=self.num_resolutions,
367 num_res_blocks=num_res_blocks,
368 resolution=resolution,
369 temb_channels=self.temb_ch,
370 dropout=dropout,
371 norm_type=self.norm_type,
372 causality_axis=self.causality_axis,
373 attn_type=self.attn_type,
374 attn_resolutions=attn_resolutions,
375 resamp_with_conv=resamp_with_conv,
376 initial_block_channels=base_block_channels,
377 )
378
379 self.norm_out = build_normalization_layer(final_block_channels, normtype=self.norm_type)
380 self.conv_out = make_conv2d(
381 final_block_channels, out_ch, kernel_size=3, stride=1, causality_axis=self.causality_axis
382 )
383
384 def forward(self, sample: torch.Tensor) -> torch.Tensor:
385 """
386 Decode latent features back to audio spectrograms.
387 Args:
388 sample: Encoded latent representation of shape (batch, channels, frames, mel_bins)
389 Returns:
390 Reconstructed audio spectrogram of shape (batch, channels, time, frequency)
391 """
392 sample, target_shape = self._denormalize_latents(sample)
393
394 h = self.conv_in(sample)
395 h = run_mid_block(self.mid, h)
396 h = self._run_upsampling_path(h)
397 h = self._finalize_output(h)
398
399 return self._adjust_output_shape(h, target_shape)
400
401 def _denormalize_latents(self, sample: torch.Tensor) -> tuple[torch.Tensor, AudioLatentShape]:
402 latent_shape = AudioLatentShape(
403 batch=sample.shape[0],
404 channels=sample.shape[1],
405 frames=sample.shape[2],
406 mel_bins=sample.shape[3],
407 )
408
409 sample_patched = self.patchifier.patchify(sample)
410 sample_denormalized = self.per_channel_statistics.un_normalize(sample_patched)
411 sample = self.patchifier.unpatchify(sample_denormalized, latent_shape)
412
413 target_frames = latent_shape.frames * LATENT_DOWNSAMPLE_FACTOR
414 if self.causality_axis != CausalityAxis.NONE:
415 target_frames = max(target_frames - (LATENT_DOWNSAMPLE_FACTOR - 1), 1)
416
417 target_shape = AudioLatentShape(
418 batch=latent_shape.batch,
419 channels=self.out_ch,
420 frames=target_frames,
421 mel_bins=self.mel_bins if self.mel_bins is not None else latent_shape.mel_bins,
422 )
423
424 return sample, target_shape
425
426 def _adjust_output_shape(
427 self,
428 decoded_output: torch.Tensor,
429 target_shape: AudioLatentShape,
430 ) -> torch.Tensor:
431 """
432 Adjust output shape to match target dimensions for variable-length audio.
433 This function handles the common case where decoded audio spectrograms need to be
434 resized to match a specific target shape.
435 Args:
436 decoded_output: Tensor of shape (batch, channels, time, frequency)
437 target_shape: AudioLatentShape describing (batch, channels, time, mel bins)
438 Returns:
439 Tensor adjusted to match target_shape exactly
440 """
441 # Current output shape: (batch, channels, time, frequency)
442 _, _, current_time, current_freq = decoded_output.shape
443 target_channels = target_shape.channels
444 target_time = target_shape.frames
445 target_freq = target_shape.mel_bins
446
447 # Step 1: Crop first to avoid exceeding target dimensions
448 decoded_output = decoded_output[
449 :, :target_channels, : min(current_time, target_time), : min(current_freq, target_freq)
450 ]
451
452 # Step 2: Calculate padding needed for time and frequency dimensions
453 time_padding_needed = target_time - decoded_output.shape[2]
454 freq_padding_needed = target_freq - decoded_output.shape[3]
455
456 # Step 3: Apply padding if needed
457 if time_padding_needed > 0 or freq_padding_needed > 0:
458 # PyTorch padding format: (pad_left, pad_right, pad_top, pad_bottom)
459 # For audio: pad_left/right = frequency, pad_top/bottom = time
460 padding = (
461 0,
462 max(freq_padding_needed, 0), # frequency padding (left, right)
463 0,
464 max(time_padding_needed, 0), # time padding (top, bottom)
465 )
466 decoded_output = F.pad(decoded_output, padding)
467
468 # Step 4: Final safety crop to ensure exact target shape
469 decoded_output = decoded_output[:, :target_channels, :target_time, :target_freq]
470
471 return decoded_output
472
473 def _run_upsampling_path(self, h: torch.Tensor) -> torch.Tensor:
474 for level in reversed(range(self.num_resolutions)):
475 stage = self.up[level]
476 for block_idx, block in enumerate(stage.block):
477 h = block(h, temb=None)
478 if stage.attn:
479 h = stage.attn[block_idx](h)
480
481 if level != 0 and hasattr(stage, "upsample"):
482 h = stage.upsample(h)
483
484 return h
485
486 def _finalize_output(self, h: torch.Tensor) -> torch.Tensor:
487 if self.give_pre_end:
488 return h
489
490 h = self.norm_out(h)
491 h = self.non_linearity(h)
492 h = self.conv_out(h)
493 return torch.tanh(h) if self.tanh_out else h
494
495
496 def decode_audio(latent: torch.Tensor, audio_decoder: "AudioDecoder", vocoder: "Vocoder") -> Audio:
497 """
498 Decode an audio latent representation using the provided audio decoder and vocoder.
499 Args:
500 latent: Input audio latent tensor.
501 audio_decoder: Model to decode the latent to waveform features.
502 vocoder: Model to convert decoded features to audio waveform.
503 Returns:
504 Decoded audio with waveform and sampling rate.
505 """
506 decoded_audio = audio_decoder(latent)
507 waveform = vocoder(decoded_audio).squeeze(0).float()
508 return Audio(waveform=waveform, sampling_rate=vocoder.output_sampling_rate)
509
509 lines PYTHON