| 1 | import math |
| 2 | from typing import List |
| 3 | |
| 4 | import einops |
| 5 | import torch |
| 6 | import torch.nn.functional as F |
| 7 | from torch import nn |
| 8 | |
| 9 | from ltx_core.model.audio_vae.resnet import LRELU_SLOPE, ResBlock1 |
| 10 | |
| 11 | |
| 12 | def get_padding(kernel_size: int, dilation: int = 1) -> int: |
| 13 | return int((kernel_size * dilation - dilation) / 2) |
| 14 | |
| 15 | |
| 16 | # --------------------------------------------------------------------------- |
| 17 | # Anti-aliased resampling helpers (kaiser-sinc filters) for BigVGAN v2 |
| 18 | # Adopted from https://github.com/NVIDIA/BigVGAN |
| 19 | # --------------------------------------------------------------------------- |
| 20 | |
| 21 | |
| 22 | def _sinc(x: torch.Tensor) -> torch.Tensor: |
| 23 | return torch.where( |
| 24 | x == 0, |
| 25 | torch.tensor(1.0, device=x.device, dtype=x.dtype), |
| 26 | torch.sin(math.pi * x) / math.pi / x, |
| 27 | ) |
| 28 | |
| 29 | |
| 30 | def kaiser_sinc_filter1d(cutoff: float, half_width: float, kernel_size: int) -> torch.Tensor: |
| 31 | even = kernel_size % 2 == 0 |
| 32 | half_size = kernel_size // 2 |
| 33 | delta_f = 4 * half_width |
| 34 | amplitude = 2.285 * (half_size - 1) * math.pi * delta_f + 7.95 |
| 35 | if amplitude > 50.0: |
| 36 | beta = 0.1102 * (amplitude - 8.7) |
| 37 | elif amplitude >= 21.0: |
| 38 | beta = 0.5842 * (amplitude - 21) ** 0.4 + 0.07886 * (amplitude - 21.0) |
| 39 | else: |
| 40 | beta = 0.0 |
| 41 | window = torch.kaiser_window(kernel_size, beta=beta, periodic=False) |
| 42 | time = torch.arange(-half_size, half_size) + 0.5 if even else torch.arange(kernel_size) - half_size |
| 43 | if cutoff == 0: |
| 44 | filter_ = torch.zeros_like(time) |
| 45 | else: |
| 46 | filter_ = 2 * cutoff * window * _sinc(2 * cutoff * time) |
| 47 | filter_ /= filter_.sum() |
| 48 | return filter_.view(1, 1, kernel_size) |
| 49 | |
| 50 | |
| 51 | class LowPassFilter1d(nn.Module): |
| 52 | def __init__( |
| 53 | self, |
| 54 | cutoff: float = 0.5, |
| 55 | half_width: float = 0.6, |
| 56 | stride: int = 1, |
| 57 | padding: bool = True, |
| 58 | padding_mode: str = "replicate", |
| 59 | kernel_size: int = 12, |
| 60 | ) -> None: |
| 61 | super().__init__() |
| 62 | if cutoff < -0.0: |
| 63 | raise ValueError("Minimum cutoff must be larger than zero.") |
| 64 | if cutoff > 0.5: |
| 65 | raise ValueError("A cutoff above 0.5 does not make sense.") |
| 66 | self.kernel_size = kernel_size |
| 67 | self.even = kernel_size % 2 == 0 |
| 68 | self.pad_left = kernel_size // 2 - int(self.even) |
| 69 | self.pad_right = kernel_size // 2 |
| 70 | self.stride = stride |
| 71 | self.padding = padding |
| 72 | self.padding_mode = padding_mode |
| 73 | self.register_buffer("filter", kaiser_sinc_filter1d(cutoff, half_width, kernel_size)) |
| 74 | |
| 75 | def forward(self, x: torch.Tensor) -> torch.Tensor: |
| 76 | _, n_channels, _ = x.shape |
| 77 | if self.padding: |
| 78 | x = F.pad(x, (self.pad_left, self.pad_right), mode=self.padding_mode) |
| 79 | return F.conv1d(x, self.filter.expand(n_channels, -1, -1), stride=self.stride, groups=n_channels) |
| 80 | |
| 81 | |
| 82 | class UpSample1d(nn.Module): |
| 83 | def __init__( |
| 84 | self, |
| 85 | ratio: int = 2, |
| 86 | kernel_size: int | None = None, |
| 87 | persistent: bool = True, |
| 88 | window_type: str = "kaiser", |
| 89 | ) -> None: |
| 90 | super().__init__() |
| 91 | self.ratio = ratio |
| 92 | self.stride = ratio |
| 93 | |
| 94 | if window_type == "hann": |
| 95 | # Hann-windowed sinc filter equivalent to torchaudio.functional.resample |
| 96 | rolloff = 0.99 |
| 97 | lowpass_filter_width = 6 |
| 98 | width = math.ceil(lowpass_filter_width / rolloff) |
| 99 | self.kernel_size = 2 * width * ratio + 1 |
| 100 | self.pad = width |
| 101 | self.pad_left = 2 * width * ratio |
| 102 | self.pad_right = self.kernel_size - ratio |
| 103 | time_axis = (torch.arange(self.kernel_size) / ratio - width) * rolloff |
| 104 | time_clamped = time_axis.clamp(-lowpass_filter_width, lowpass_filter_width) |
| 105 | window = torch.cos(time_clamped * math.pi / lowpass_filter_width / 2) ** 2 |
| 106 | sinc_filter = (torch.sinc(time_axis) * window * rolloff / ratio).view(1, 1, -1) |
| 107 | else: |
| 108 | # Kaiser-windowed sinc filter (BigVGAN default). |
| 109 | self.kernel_size = int(6 * ratio // 2) * 2 if kernel_size is None else kernel_size |
| 110 | self.pad = self.kernel_size // ratio - 1 |
| 111 | self.pad_left = self.pad * self.stride + (self.kernel_size - self.stride) // 2 |
| 112 | self.pad_right = self.pad * self.stride + (self.kernel_size - self.stride + 1) // 2 |
| 113 | sinc_filter = kaiser_sinc_filter1d( |
| 114 | cutoff=0.5 / ratio, |
| 115 | half_width=0.6 / ratio, |
| 116 | kernel_size=self.kernel_size, |
| 117 | ) |
| 118 | |
| 119 | self.register_buffer("filter", sinc_filter, persistent=persistent) |
| 120 | |
| 121 | def forward(self, x: torch.Tensor) -> torch.Tensor: |
| 122 | _, n_channels, _ = x.shape |
| 123 | x = F.pad(x, (self.pad, self.pad), mode="replicate") |
| 124 | filt = self.filter.to(dtype=x.dtype, device=x.device).expand(n_channels, -1, -1) |
| 125 | x = self.ratio * F.conv_transpose1d(x, filt, stride=self.stride, groups=n_channels) |
| 126 | return x[..., self.pad_left : -self.pad_right] |
| 127 | |
| 128 | |
| 129 | class DownSample1d(nn.Module): |
| 130 | def __init__(self, ratio: int = 2, kernel_size: int | None = None) -> None: |
| 131 | super().__init__() |
| 132 | self.ratio = ratio |
| 133 | self.kernel_size = int(6 * ratio // 2) * 2 if kernel_size is None else kernel_size |
| 134 | self.lowpass = LowPassFilter1d( |
| 135 | cutoff=0.5 / ratio, |
| 136 | half_width=0.6 / ratio, |
| 137 | stride=ratio, |
| 138 | kernel_size=self.kernel_size, |
| 139 | ) |
| 140 | |
| 141 | def forward(self, x: torch.Tensor) -> torch.Tensor: |
| 142 | return self.lowpass(x) |
| 143 | |
| 144 | |
| 145 | class Activation1d(nn.Module): |
| 146 | def __init__( |
| 147 | self, |
| 148 | activation: nn.Module, |
| 149 | up_ratio: int = 2, |
| 150 | down_ratio: int = 2, |
| 151 | up_kernel_size: int = 12, |
| 152 | down_kernel_size: int = 12, |
| 153 | ) -> None: |
| 154 | super().__init__() |
| 155 | self.act = activation |
| 156 | self.upsample = UpSample1d(up_ratio, up_kernel_size) |
| 157 | self.downsample = DownSample1d(down_ratio, down_kernel_size) |
| 158 | |
| 159 | def forward(self, x: torch.Tensor) -> torch.Tensor: |
| 160 | x = self.upsample(x) |
| 161 | x = self.act(x) |
| 162 | return self.downsample(x) |
| 163 | |
| 164 | |
| 165 | class Snake(nn.Module): |
| 166 | def __init__( |
| 167 | self, |
| 168 | in_features: int, |
| 169 | alpha: float = 1.0, |
| 170 | alpha_trainable: bool = True, |
| 171 | alpha_logscale: bool = True, |
| 172 | ) -> None: |
| 173 | super().__init__() |
| 174 | self.alpha_logscale = alpha_logscale |
| 175 | self.alpha = nn.Parameter(torch.zeros(in_features) if alpha_logscale else torch.ones(in_features) * alpha) |
| 176 | self.alpha.requires_grad = alpha_trainable |
| 177 | self.eps = 1e-9 |
| 178 | |
| 179 | def forward(self, x: torch.Tensor) -> torch.Tensor: |
| 180 | alpha = self.alpha.unsqueeze(0).unsqueeze(-1) |
| 181 | if self.alpha_logscale: |
| 182 | alpha = torch.exp(alpha) |
| 183 | return x + (1.0 / (alpha + self.eps)) * torch.sin(x * alpha).pow(2) |
| 184 | |
| 185 | |
| 186 | class SnakeBeta(nn.Module): |
| 187 | def __init__( |
| 188 | self, |
| 189 | in_features: int, |
| 190 | alpha: float = 1.0, |
| 191 | alpha_trainable: bool = True, |
| 192 | alpha_logscale: bool = True, |
| 193 | ) -> None: |
| 194 | super().__init__() |
| 195 | self.alpha_logscale = alpha_logscale |
| 196 | self.alpha = nn.Parameter(torch.zeros(in_features) if alpha_logscale else torch.ones(in_features) * alpha) |
| 197 | self.alpha.requires_grad = alpha_trainable |
| 198 | self.beta = nn.Parameter(torch.zeros(in_features) if alpha_logscale else torch.ones(in_features) * alpha) |
| 199 | self.beta.requires_grad = alpha_trainable |
| 200 | self.eps = 1e-9 |
| 201 | |
| 202 | def forward(self, x: torch.Tensor) -> torch.Tensor: |
| 203 | alpha = self.alpha.unsqueeze(0).unsqueeze(-1) |
| 204 | beta = self.beta.unsqueeze(0).unsqueeze(-1) |
| 205 | if self.alpha_logscale: |
| 206 | alpha = torch.exp(alpha) |
| 207 | beta = torch.exp(beta) |
| 208 | return x + (1.0 / (beta + self.eps)) * torch.sin(x * alpha).pow(2) |
| 209 | |
| 210 | |
| 211 | class AMPBlock1(nn.Module): |
| 212 | def __init__( |
| 213 | self, |
| 214 | channels: int, |
| 215 | kernel_size: int = 3, |
| 216 | dilation: tuple[int, int, int] = (1, 3, 5), |
| 217 | activation: str = "snake", |
| 218 | ) -> None: |
| 219 | super().__init__() |
| 220 | act_cls = SnakeBeta if activation == "snakebeta" else Snake |
| 221 | self.convs1 = nn.ModuleList( |
| 222 | [ |
| 223 | nn.Conv1d( |
| 224 | channels, |
| 225 | channels, |
| 226 | kernel_size, |
| 227 | 1, |
| 228 | dilation=dilation[0], |
| 229 | padding=get_padding(kernel_size, dilation[0]), |
| 230 | ), |
| 231 | nn.Conv1d( |
| 232 | channels, |
| 233 | channels, |
| 234 | kernel_size, |
| 235 | 1, |
| 236 | dilation=dilation[1], |
| 237 | padding=get_padding(kernel_size, dilation[1]), |
| 238 | ), |
| 239 | nn.Conv1d( |
| 240 | channels, |
| 241 | channels, |
| 242 | kernel_size, |
| 243 | 1, |
| 244 | dilation=dilation[2], |
| 245 | padding=get_padding(kernel_size, dilation[2]), |
| 246 | ), |
| 247 | ] |
| 248 | ) |
| 249 | |
| 250 | self.convs2 = nn.ModuleList( |
| 251 | [ |
| 252 | nn.Conv1d(channels, channels, kernel_size, 1, dilation=1, padding=get_padding(kernel_size, 1)), |
| 253 | nn.Conv1d(channels, channels, kernel_size, 1, dilation=1, padding=get_padding(kernel_size, 1)), |
| 254 | nn.Conv1d(channels, channels, kernel_size, 1, dilation=1, padding=get_padding(kernel_size, 1)), |
| 255 | ] |
| 256 | ) |
| 257 | |
| 258 | self.acts1 = nn.ModuleList([Activation1d(act_cls(channels)) for _ in range(len(self.convs1))]) |
| 259 | self.acts2 = nn.ModuleList([Activation1d(act_cls(channels)) for _ in range(len(self.convs2))]) |
| 260 | |
| 261 | def forward(self, x: torch.Tensor) -> torch.Tensor: |
| 262 | for c1, c2, a1, a2 in zip(self.convs1, self.convs2, self.acts1, self.acts2, strict=True): |
| 263 | xt = a1(x) |
| 264 | xt = c1(xt) |
| 265 | xt = a2(xt) |
| 266 | xt = c2(xt) |
| 267 | x = x + xt |
| 268 | return x |
| 269 | |
| 270 | |
| 271 | class Vocoder(torch.nn.Module): |
| 272 | """ |
| 273 | Vocoder model for synthesizing audio from Mel spectrograms. |
| 274 | Args: |
| 275 | resblock_kernel_sizes: List of kernel sizes for the residual blocks. |
| 276 | This value is read from the checkpoint at `config.vocoder.resblock_kernel_sizes`. |
| 277 | upsample_rates: List of upsampling rates. |
| 278 | This value is read from the checkpoint at `config.vocoder.upsample_rates`. |
| 279 | upsample_kernel_sizes: List of kernel sizes for the upsampling layers. |
| 280 | This value is read from the checkpoint at `config.vocoder.upsample_kernel_sizes`. |
| 281 | resblock_dilation_sizes: List of dilation sizes for the residual blocks. |
| 282 | This value is read from the checkpoint at `config.vocoder.resblock_dilation_sizes`. |
| 283 | upsample_initial_channel: Initial number of channels for the upsampling layers. |
| 284 | This value is read from the checkpoint at `config.vocoder.upsample_initial_channel`. |
| 285 | resblock: Type of residual block to use ("1", "2", or "AMP1"). |
| 286 | This value is read from the checkpoint at `config.vocoder.resblock`. |
| 287 | output_sampling_rate: Waveform sample rate. |
| 288 | This value is read from the checkpoint at `config.vocoder.output_sampling_rate`. |
| 289 | activation: Activation type for BigVGAN v2 ("snake" or "snakebeta"). Only used when resblock="AMP1". |
| 290 | use_tanh_at_final: Apply tanh at the output (when apply_final_activation=True). |
| 291 | apply_final_activation: Whether to apply the final tanh/clamp activation. |
| 292 | use_bias_at_final: Whether to use bias in the final conv layer. |
| 293 | """ |
| 294 | |
| 295 | def __init__( # noqa: PLR0913 |
| 296 | self, |
| 297 | resblock_kernel_sizes: List[int] | None = None, |
| 298 | upsample_rates: List[int] | None = None, |
| 299 | upsample_kernel_sizes: List[int] | None = None, |
| 300 | resblock_dilation_sizes: List[List[int]] | None = None, |
| 301 | upsample_initial_channel: int = 1024, |
| 302 | resblock: str = "1", |
| 303 | output_sampling_rate: int = 24000, |
| 304 | activation: str = "snake", |
| 305 | use_tanh_at_final: bool = True, |
| 306 | apply_final_activation: bool = True, |
| 307 | use_bias_at_final: bool = True, |
| 308 | ) -> None: |
| 309 | super().__init__() |
| 310 | |
| 311 | # Mutable default values are not supported as default arguments. |
| 312 | if resblock_kernel_sizes is None: |
| 313 | resblock_kernel_sizes = [3, 7, 11] |
| 314 | if upsample_rates is None: |
| 315 | upsample_rates = [6, 5, 2, 2, 2] |
| 316 | if upsample_kernel_sizes is None: |
| 317 | upsample_kernel_sizes = [16, 15, 8, 4, 4] |
| 318 | if resblock_dilation_sizes is None: |
| 319 | resblock_dilation_sizes = [[1, 3, 5], [1, 3, 5], [1, 3, 5]] |
| 320 | |
| 321 | self.output_sampling_rate = output_sampling_rate |
| 322 | self.num_kernels = len(resblock_kernel_sizes) |
| 323 | self.num_upsamples = len(upsample_rates) |
| 324 | self.use_tanh_at_final = use_tanh_at_final |
| 325 | self.apply_final_activation = apply_final_activation |
| 326 | self.is_amp = resblock == "AMP1" |
| 327 | |
| 328 | # All production checkpoints are stereo: 128 input channels (2 stereo channels x 64 mel |
| 329 | # bins each), 2 output channels. |
| 330 | self.conv_pre = nn.Conv1d( |
| 331 | in_channels=128, |
| 332 | out_channels=upsample_initial_channel, |
| 333 | kernel_size=7, |
| 334 | stride=1, |
| 335 | padding=3, |
| 336 | ) |
| 337 | resblock_cls = ResBlock1 if resblock == "1" else AMPBlock1 |
| 338 | |
| 339 | self.ups = nn.ModuleList( |
| 340 | nn.ConvTranspose1d( |
| 341 | upsample_initial_channel // (2**i), |
| 342 | upsample_initial_channel // (2 ** (i + 1)), |
| 343 | kernel_size, |
| 344 | stride, |
| 345 | padding=(kernel_size - stride) // 2, |
| 346 | ) |
| 347 | for i, (stride, kernel_size) in enumerate(zip(upsample_rates, upsample_kernel_sizes, strict=True)) |
| 348 | ) |
| 349 | |
| 350 | final_channels = upsample_initial_channel // (2 ** len(upsample_rates)) |
| 351 | self.resblocks = nn.ModuleList() |
| 352 | |
| 353 | for i in range(len(upsample_rates)): |
| 354 | ch = upsample_initial_channel // (2 ** (i + 1)) |
| 355 | for kernel_size, dilations in zip(resblock_kernel_sizes, resblock_dilation_sizes, strict=True): |
| 356 | if self.is_amp: |
| 357 | self.resblocks.append(resblock_cls(ch, kernel_size, dilations, activation=activation)) |
| 358 | else: |
| 359 | self.resblocks.append(resblock_cls(ch, kernel_size, dilations)) |
| 360 | |
| 361 | if self.is_amp: |
| 362 | self.act_post: nn.Module = Activation1d(SnakeBeta(final_channels)) |
| 363 | else: |
| 364 | self.act_post = nn.LeakyReLU() |
| 365 | |
| 366 | # All production checkpoints are stereo: this final conv maps `final_channels` to 2 output channels (stereo). |
| 367 | self.conv_post = nn.Conv1d( |
| 368 | in_channels=final_channels, |
| 369 | out_channels=2, |
| 370 | kernel_size=7, |
| 371 | stride=1, |
| 372 | padding=3, |
| 373 | bias=use_bias_at_final, |
| 374 | ) |
| 375 | |
| 376 | def forward(self, x: torch.Tensor) -> torch.Tensor: |
| 377 | """ |
| 378 | Forward pass of the vocoder. |
| 379 | Args: |
| 380 | x: Input Mel spectrogram tensor. Can be either: |
| 381 | - 3D: (batch_size, time, mel_bins) for mono |
| 382 | - 4D: (batch_size, 2, time, mel_bins) for stereo |
| 383 | Returns: |
| 384 | Audio waveform tensor of shape (batch_size, out_channels, audio_length) |
| 385 | """ |
| 386 | x = x.transpose(2, 3) # (batch, channels, time, mel_bins) -> (batch, channels, mel_bins, time) |
| 387 | |
| 388 | if x.dim() == 4: # stereo |
| 389 | assert x.shape[1] == 2, "Input must have 2 channels for stereo" |
| 390 | x = einops.rearrange(x, "b s c t -> b (s c) t") |
| 391 | |
| 392 | x = self.conv_pre(x) |
| 393 | |
| 394 | for i in range(self.num_upsamples): |
| 395 | if not self.is_amp: |
| 396 | x = F.leaky_relu(x, LRELU_SLOPE) |
| 397 | x = self.ups[i](x) |
| 398 | start = i * self.num_kernels |
| 399 | end = start + self.num_kernels |
| 400 | |
| 401 | # Evaluate all resblocks with the same input tensor so they can run |
| 402 | # independently (and thus in parallel on accelerator hardware) before |
| 403 | # aggregating their outputs via mean. |
| 404 | block_outputs = torch.stack( |
| 405 | [self.resblocks[idx](x) for idx in range(start, end)], |
| 406 | dim=0, |
| 407 | ) |
| 408 | x = block_outputs.mean(dim=0) |
| 409 | |
| 410 | x = self.act_post(x) |
| 411 | x = self.conv_post(x) |
| 412 | |
| 413 | if self.apply_final_activation: |
| 414 | x = torch.tanh(x) if self.use_tanh_at_final else torch.clamp(x, -1, 1) |
| 415 | |
| 416 | return x |
| 417 | |
| 418 | |
| 419 | class _STFTFn(nn.Module): |
| 420 | """Implements STFT as a convolution with precomputed DFT x Hann-window bases. |
| 421 | The DFT basis rows (real and imaginary parts interleaved) multiplied by the causal |
| 422 | Hann window are stored as buffers and loaded from the checkpoint. Using the exact |
| 423 | bfloat16 bases from training ensures the mel values fed to the BWE generator are |
| 424 | bit-identical to what it was trained on. |
| 425 | """ |
| 426 | |
| 427 | def __init__(self, filter_length: int, hop_length: int, win_length: int) -> None: |
| 428 | super().__init__() |
| 429 | self.hop_length = hop_length |
| 430 | self.win_length = win_length |
| 431 | n_freqs = filter_length // 2 + 1 |
| 432 | self.register_buffer("forward_basis", torch.zeros(n_freqs * 2, 1, filter_length)) |
| 433 | self.register_buffer("inverse_basis", torch.zeros(n_freqs * 2, 1, filter_length)) |
| 434 | |
| 435 | def forward(self, y: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: |
| 436 | """Compute magnitude and phase spectrogram from a batch of waveforms. |
| 437 | Applies causal (left-only) padding of win_length - hop_length samples so that |
| 438 | each output frame depends only on past and present input — no lookahead. |
| 439 | Args: |
| 440 | y: Waveform tensor of shape (B, T). |
| 441 | Returns: |
| 442 | magnitude: Linear amplitude spectrogram, shape (B, n_freqs, T_frames). |
| 443 | phase: Phase spectrogram in radians, shape (B, n_freqs, T_frames). |
| 444 | """ |
| 445 | if y.dim() == 2: |
| 446 | y = y.unsqueeze(1) # (B, 1, T) |
| 447 | left_pad = max(0, self.win_length - self.hop_length) # causal: left-only |
| 448 | y = F.pad(y, (left_pad, 0)) |
| 449 | spec = F.conv1d(y, self.forward_basis, stride=self.hop_length, padding=0) |
| 450 | n_freqs = spec.shape[1] // 2 |
| 451 | real, imag = spec[:, :n_freqs], spec[:, n_freqs:] |
| 452 | magnitude = torch.sqrt(real**2 + imag**2) |
| 453 | phase = torch.atan2(imag.float(), real.float()).to(real.dtype) |
| 454 | return magnitude, phase |
| 455 | |
| 456 | |
| 457 | class MelSTFT(nn.Module): |
| 458 | """Causal log-mel spectrogram module whose buffers are loaded from the checkpoint. |
| 459 | Computes a log-mel spectrogram by running the causal STFT (_STFTFn) on the input |
| 460 | waveform and projecting the linear magnitude spectrum onto the mel filterbank. |
| 461 | The module's state dict layout matches the 'mel_stft.*' keys stored in the checkpoint |
| 462 | (mel_basis, stft_fn.forward_basis, stft_fn.inverse_basis). |
| 463 | """ |
| 464 | |
| 465 | def __init__( |
| 466 | self, |
| 467 | filter_length: int, |
| 468 | hop_length: int, |
| 469 | win_length: int, |
| 470 | n_mel_channels: int, |
| 471 | ) -> None: |
| 472 | super().__init__() |
| 473 | self.stft_fn = _STFTFn(filter_length, hop_length, win_length) |
| 474 | |
| 475 | # Initialized to zeros; load_state_dict overwrites with the checkpoint's |
| 476 | # exact bfloat16 filterbank (vocoder.mel_stft.mel_basis, shape [n_mels, n_freqs]). |
| 477 | n_freqs = filter_length // 2 + 1 |
| 478 | self.register_buffer("mel_basis", torch.zeros(n_mel_channels, n_freqs)) |
| 479 | |
| 480 | def mel_spectrogram(self, y: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: |
| 481 | """Compute log-mel spectrogram and auxiliary spectral quantities. |
| 482 | Args: |
| 483 | y: Waveform tensor of shape (B, T). |
| 484 | Returns: |
| 485 | log_mel: Log-compressed mel spectrogram, shape (B, n_mel_channels, T_frames). |
| 486 | magnitude: Linear amplitude spectrogram, shape (B, n_freqs, T_frames). |
| 487 | phase: Phase spectrogram in radians, shape (B, n_freqs, T_frames). |
| 488 | energy: Per-frame energy (L2 norm over frequency), shape (B, T_frames). |
| 489 | """ |
| 490 | magnitude, phase = self.stft_fn(y) |
| 491 | energy = torch.norm(magnitude, dim=1) |
| 492 | mel = torch.matmul(self.mel_basis.to(magnitude.dtype), magnitude) |
| 493 | log_mel = torch.log(torch.clamp(mel, min=1e-5)) |
| 494 | return log_mel, magnitude, phase, energy |
| 495 | |
| 496 | |
| 497 | class VocoderWithBWE(nn.Module): |
| 498 | """Vocoder with bandwidth extension (BWE) upsampling. |
| 499 | Chains a mel-to-wav vocoder with a BWE module that upsamples the output |
| 500 | to a higher sample rate. The BWE computes a mel spectrogram from the |
| 501 | vocoder output, runs it through a second generator to predict a residual, |
| 502 | and adds it to a sinc-resampled skip connection. |
| 503 | """ |
| 504 | |
| 505 | def __init__( |
| 506 | self, |
| 507 | vocoder: Vocoder, |
| 508 | bwe_generator: Vocoder, |
| 509 | mel_stft: MelSTFT, |
| 510 | input_sampling_rate: int, |
| 511 | output_sampling_rate: int, |
| 512 | hop_length: int, |
| 513 | ) -> None: |
| 514 | super().__init__() |
| 515 | self.vocoder = vocoder |
| 516 | self.bwe_generator = bwe_generator |
| 517 | self.mel_stft = mel_stft |
| 518 | self.input_sampling_rate = input_sampling_rate |
| 519 | self.output_sampling_rate = output_sampling_rate |
| 520 | self.hop_length = hop_length |
| 521 | # Compute the resampler on CPU so the sinc filter is materialized even when |
| 522 | # the model is constructed on meta device (SingleGPUModelBuilder pattern). |
| 523 | # The filter is not stored in the checkpoint (persistent=False). |
| 524 | with torch.device("cpu"): |
| 525 | self.resampler = UpSample1d( |
| 526 | ratio=output_sampling_rate // input_sampling_rate, persistent=False, window_type="hann" |
| 527 | ) |
| 528 | |
| 529 | @property |
| 530 | def conv_pre(self) -> nn.Conv1d: |
| 531 | return self.vocoder.conv_pre |
| 532 | |
| 533 | @property |
| 534 | def conv_post(self) -> nn.Conv1d: |
| 535 | return self.vocoder.conv_post |
| 536 | |
| 537 | def _compute_mel(self, audio: torch.Tensor) -> torch.Tensor: |
| 538 | """Compute log-mel spectrogram from waveform using causal STFT bases. |
| 539 | Args: |
| 540 | audio: Waveform tensor of shape (B, C, T). |
| 541 | Returns: |
| 542 | mel: Log-mel spectrogram of shape (B, C, n_mels, T_frames). |
| 543 | """ |
| 544 | batch, n_channels, _ = audio.shape |
| 545 | flat = audio.reshape(batch * n_channels, -1) # (B*C, T) |
| 546 | mel, _, _, _ = self.mel_stft.mel_spectrogram(flat) # (B*C, n_mels, T_frames) |
| 547 | return mel.reshape(batch, n_channels, mel.shape[1], mel.shape[2]) # (B, C, n_mels, T_frames) |
| 548 | |
| 549 | def forward(self, mel_spec: torch.Tensor) -> torch.Tensor: |
| 550 | """Run the full vocoder + BWE forward pass. |
| 551 | Args: |
| 552 | mel_spec: Mel spectrogram of shape (B, 2, T, mel_bins) for stereo |
| 553 | or (B, T, mel_bins) for mono. Same format as Vocoder.forward. |
| 554 | Returns: |
| 555 | Waveform tensor of shape (B, out_channels, T_out) clipped to [-1, 1]. |
| 556 | """ |
| 557 | x = self.vocoder(mel_spec) |
| 558 | _, _, length_low_rate = x.shape |
| 559 | output_length = length_low_rate * self.output_sampling_rate // self.input_sampling_rate |
| 560 | |
| 561 | # Pad to multiple of hop_length for exact mel frame count |
| 562 | remainder = length_low_rate % self.hop_length |
| 563 | if remainder != 0: |
| 564 | x = F.pad(x, (0, self.hop_length - remainder)) |
| 565 | |
| 566 | # Compute mel spectrogram from vocoder output: (B, C, n_mels, T_frames) |
| 567 | mel = self._compute_mel(x) |
| 568 | |
| 569 | # Vocoder.forward expects (B, C, T, mel_bins) — transpose before calling bwe_generator |
| 570 | mel_for_bwe = mel.transpose(2, 3) # (B, C, T_frames, mel_bins) |
| 571 | residual = self.bwe_generator(mel_for_bwe) |
| 572 | skip = self.resampler(x) |
| 573 | assert residual.shape == skip.shape, f"residual {residual.shape} != skip {skip.shape}" |
| 574 | |
| 575 | return torch.clamp(residual + skip, -1, 1)[..., :output_length] |
| 576 |