738 lines
24 KiB
Python
738 lines
24 KiB
Python
"""Vocoder for converting mel spectrograms to audio waveforms.
|
|
|
|
Supports:
|
|
- HiFi-GAN (LTX-2): ResBlock1 with LeakyReLU
|
|
- BigVGAN v2 (LTX-2.3): AMPBlock1 with Snake/SnakeBeta + anti-aliased resampling
|
|
- VocoderWithBWE (LTX-2.3): Base vocoder + bandwidth extension (16kHz -> 48kHz)
|
|
"""
|
|
|
|
import math
|
|
from pathlib import Path
|
|
from typing import Tuple
|
|
|
|
import mlx.core as mx
|
|
import mlx.nn as nn
|
|
|
|
from ..config import VocoderModelConfig
|
|
from .resnet import LRELU_SLOPE, ResBlock1, ResBlock2, leaky_relu
|
|
|
|
|
|
def get_padding(kernel_size: int, dilation: int = 1) -> int:
|
|
return int((kernel_size * dilation - dilation) / 2)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Snake / SnakeBeta activations (BigVGAN v2)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class Snake(nn.Module):
|
|
"""Snake activation: x + (1/alpha) * sin^2(alpha * x)."""
|
|
|
|
def __init__(self, in_features: int, alpha_logscale: bool = True) -> None:
|
|
super().__init__()
|
|
self.alpha_logscale = alpha_logscale
|
|
self.alpha = (
|
|
mx.zeros((in_features,)) if alpha_logscale else mx.ones((in_features,))
|
|
)
|
|
|
|
def __call__(self, x: mx.array) -> mx.array:
|
|
# x: (N, L, C) in MLX format
|
|
alpha = self.alpha # (C,)
|
|
if self.alpha_logscale:
|
|
alpha = mx.exp(alpha)
|
|
return x + (1.0 / (alpha + 1e-9)) * mx.power(mx.sin(x * alpha), 2)
|
|
|
|
|
|
class SnakeBeta(nn.Module):
|
|
"""SnakeBeta activation: x + (1/beta) * sin^2(alpha * x)."""
|
|
|
|
def __init__(self, in_features: int, alpha_logscale: bool = True) -> None:
|
|
super().__init__()
|
|
self.alpha_logscale = alpha_logscale
|
|
self.alpha = (
|
|
mx.zeros((in_features,)) if alpha_logscale else mx.ones((in_features,))
|
|
)
|
|
self.beta = (
|
|
mx.zeros((in_features,)) if alpha_logscale else mx.ones((in_features,))
|
|
)
|
|
|
|
def __call__(self, x: mx.array) -> mx.array:
|
|
alpha = self.alpha
|
|
beta = self.beta
|
|
if self.alpha_logscale:
|
|
alpha = mx.exp(alpha)
|
|
beta = mx.exp(beta)
|
|
return x + (1.0 / (beta + 1e-9)) * mx.power(mx.sin(x * alpha), 2)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Anti-aliased resampling (Kaiser-sinc filters)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _sinc(x: mx.array) -> mx.array:
|
|
return mx.where(
|
|
x == 0,
|
|
mx.ones_like(x),
|
|
mx.sin(mx.array(math.pi) * x) / (mx.array(math.pi) * x),
|
|
)
|
|
|
|
|
|
def kaiser_sinc_filter1d(
|
|
cutoff: float, half_width: float, kernel_size: int
|
|
) -> mx.array:
|
|
"""Compute a Kaiser-windowed sinc filter."""
|
|
even = kernel_size % 2 == 0
|
|
half_size = kernel_size // 2
|
|
delta_f = 4 * half_width
|
|
amplitude = 2.285 * (half_size - 1) * math.pi * delta_f + 7.95
|
|
if amplitude > 50.0:
|
|
beta = 0.1102 * (amplitude - 8.7)
|
|
elif amplitude >= 21.0:
|
|
beta = 0.5842 * (amplitude - 21) ** 0.4 + 0.07886 * (amplitude - 21.0)
|
|
else:
|
|
beta = 0.0
|
|
|
|
# Kaiser window - compute using scipy-compatible formula
|
|
import numpy as np
|
|
|
|
window = mx.array(np.kaiser(kernel_size, beta).astype(np.float32))
|
|
|
|
if even:
|
|
time = mx.arange(-half_size, half_size).astype(mx.float32) + 0.5
|
|
else:
|
|
time = mx.arange(kernel_size).astype(mx.float32) - half_size
|
|
|
|
if cutoff == 0:
|
|
filter_ = mx.zeros_like(time)
|
|
else:
|
|
filter_ = 2 * cutoff * window * _sinc(2 * cutoff * time)
|
|
filter_ = filter_ / mx.sum(filter_)
|
|
|
|
return filter_.reshape(1, 1, kernel_size)
|
|
|
|
|
|
def hann_sinc_filter1d(ratio: int) -> Tuple[mx.array, int, int, int]:
|
|
"""Compute a Hann-windowed sinc filter for upsampling (used by BWE resampler)."""
|
|
import numpy as np
|
|
|
|
rolloff = 0.99
|
|
lowpass_filter_width = 6
|
|
width = math.ceil(lowpass_filter_width / rolloff)
|
|
kernel_size = 2 * width * ratio + 1
|
|
pad = width
|
|
pad_left = 2 * width * ratio
|
|
pad_right = kernel_size - ratio
|
|
|
|
time = (np.arange(kernel_size) / ratio - width) * rolloff
|
|
time_clamped = np.clip(time, -lowpass_filter_width, lowpass_filter_width)
|
|
window = np.cos(time_clamped * math.pi / lowpass_filter_width / 2) ** 2
|
|
sinc_filter = np.sinc(time) * window * rolloff / ratio
|
|
|
|
filter_ = mx.array(sinc_filter.astype(np.float32)).reshape(1, 1, kernel_size)
|
|
return filter_, pad, pad_left, pad_right
|
|
|
|
|
|
class LowPassFilter1d(nn.Module):
|
|
"""Low-pass filter using depthwise convolution with Kaiser-sinc kernel."""
|
|
|
|
def __init__(
|
|
self,
|
|
cutoff: float = 0.5,
|
|
half_width: float = 0.6,
|
|
stride: int = 1,
|
|
kernel_size: int = 12,
|
|
) -> None:
|
|
super().__init__()
|
|
self.kernel_size = kernel_size
|
|
self.even = kernel_size % 2 == 0
|
|
self.pad_left = kernel_size // 2 - int(self.even)
|
|
self.pad_right = kernel_size // 2
|
|
self.stride = stride
|
|
# Filter buffer - shape (1, 1, K) in PyTorch format, loaded from weights
|
|
self.filter = kaiser_sinc_filter1d(cutoff, half_width, kernel_size)
|
|
|
|
def __call__(self, x: mx.array) -> mx.array:
|
|
# x: (N, L, C) in MLX format
|
|
n, l, c = x.shape
|
|
|
|
# Pad with edge values: replicate first/last value
|
|
first = mx.repeat(x[:, :1, :], self.pad_left, axis=1)
|
|
last = mx.repeat(x[:, -1:, :], self.pad_right, axis=1)
|
|
x = mx.concatenate([first, x, last], axis=1)
|
|
|
|
# Expand filter for depthwise conv: (1, 1, K) -> (C, K, 1) for groups=C
|
|
# Filter is stored in PyTorch format (1, 1, K), need (C, K, 1) MLX format
|
|
filt = self.filter.astype(x.dtype) # (1, 1, K)
|
|
filt = mx.transpose(filt, (0, 2, 1)) # (1, K, 1)
|
|
filt = mx.repeat(filt, c, axis=0) # (C, K, 1)
|
|
|
|
# Transpose x for depthwise conv: (N, L, C) -> (N*C, L, 1) then conv
|
|
x = mx.transpose(x, (0, 2, 1)) # (N, C, L)
|
|
x = x.reshape(n * c, -1, 1) # (N*C, L, 1)
|
|
|
|
x = mx.conv1d(x, filt[:1], stride=self.stride, groups=1) # (N*C, L', 1)
|
|
|
|
x = x.reshape(n, c, -1) # (N, C, L')
|
|
x = mx.transpose(x, (0, 2, 1)) # (N, L', C)
|
|
return x
|
|
|
|
|
|
class UpSample1d(nn.Module):
|
|
"""Anti-aliased upsampling using transposed convolution with sinc filter."""
|
|
|
|
def __init__(
|
|
self,
|
|
ratio: int = 2,
|
|
kernel_size: int = None,
|
|
window_type: str = "kaiser",
|
|
) -> None:
|
|
super().__init__()
|
|
self.ratio = ratio
|
|
self.stride = ratio
|
|
|
|
if window_type == "hann":
|
|
filt, self.pad, self.pad_left, self.pad_right = hann_sinc_filter1d(ratio)
|
|
self.kernel_size = filt.shape[2]
|
|
self.filter = filt
|
|
else:
|
|
self.kernel_size = (
|
|
int(6 * ratio // 2) * 2 if kernel_size is None else kernel_size
|
|
)
|
|
self.pad = self.kernel_size // ratio - 1
|
|
self.pad_left = (
|
|
self.pad * self.stride + (self.kernel_size - self.stride) // 2
|
|
)
|
|
self.pad_right = (
|
|
self.pad * self.stride + (self.kernel_size - self.stride + 1) // 2
|
|
)
|
|
self.filter = kaiser_sinc_filter1d(
|
|
cutoff=0.5 / ratio,
|
|
half_width=0.6 / ratio,
|
|
kernel_size=self.kernel_size,
|
|
)
|
|
|
|
def __call__(self, x: mx.array) -> mx.array:
|
|
# x: (N, L, C) in MLX format
|
|
n, l, c = x.shape
|
|
|
|
# Pad with edge values
|
|
first = mx.repeat(x[:, :1, :], self.pad, axis=1)
|
|
last = mx.repeat(x[:, -1:, :], self.pad, axis=1)
|
|
x = mx.concatenate([first, x, last], axis=1)
|
|
|
|
# Process per-channel via reshape: (N, L, C) -> (N*C, L, 1)
|
|
x = mx.transpose(x, (0, 2, 1)) # (N, C, L)
|
|
x = x.reshape(n * c, -1, 1) # (N*C, L, 1)
|
|
|
|
# Transposed conv for upsampling
|
|
# Filter: (1, 1, K) PyTorch -> (1, K, 1) MLX format for conv_transpose1d
|
|
filt = self.filter.astype(x.dtype) # (1, 1, K)
|
|
filt = mx.transpose(filt, (0, 2, 1)) # (1, K, 1)
|
|
|
|
x = self.ratio * mx.conv_transpose1d(
|
|
x, filt, stride=self.stride
|
|
) # (N*C, L', 1)
|
|
|
|
# Trim padding
|
|
x = x[:, self.pad_left : -self.pad_right, :]
|
|
|
|
x = x.reshape(n, c, -1) # (N, C, L')
|
|
x = mx.transpose(x, (0, 2, 1)) # (N, L', C)
|
|
return x
|
|
|
|
|
|
class DownSample1d(nn.Module):
|
|
"""Anti-aliased downsampling using low-pass filter."""
|
|
|
|
def __init__(self, ratio: int = 2, kernel_size: int = None) -> None:
|
|
super().__init__()
|
|
self.ratio = ratio
|
|
kernel_size = int(6 * ratio // 2) * 2 if kernel_size is None else kernel_size
|
|
self.lowpass = LowPassFilter1d(
|
|
cutoff=0.5 / ratio,
|
|
half_width=0.6 / ratio,
|
|
stride=ratio,
|
|
kernel_size=kernel_size,
|
|
)
|
|
|
|
def __call__(self, x: mx.array) -> mx.array:
|
|
return self.lowpass(x)
|
|
|
|
|
|
class Activation1d(nn.Module):
|
|
"""Anti-aliased activation: upsample -> activate -> downsample."""
|
|
|
|
def __init__(
|
|
self,
|
|
activation: nn.Module,
|
|
up_ratio: int = 2,
|
|
down_ratio: int = 2,
|
|
up_kernel_size: int = 12,
|
|
down_kernel_size: int = 12,
|
|
) -> None:
|
|
super().__init__()
|
|
self.act = activation
|
|
self.upsample = UpSample1d(up_ratio, up_kernel_size)
|
|
self.downsample = DownSample1d(down_ratio, down_kernel_size)
|
|
|
|
def __call__(self, x: mx.array) -> mx.array:
|
|
x = self.upsample(x)
|
|
x = self.act(x)
|
|
return self.downsample(x)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# AMPBlock1 (BigVGAN v2 residual block)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class AMPBlock1(nn.Module):
|
|
"""BigVGAN v2 residual block with anti-aliased Snake activations."""
|
|
|
|
def __init__(
|
|
self,
|
|
channels: int,
|
|
kernel_size: int = 3,
|
|
dilation: Tuple[int, int, int] = (1, 3, 5),
|
|
activation: str = "snakebeta",
|
|
) -> None:
|
|
super().__init__()
|
|
act_cls = SnakeBeta if activation == "snakebeta" else Snake
|
|
|
|
self.convs1 = {
|
|
i: nn.Conv1d(
|
|
channels,
|
|
channels,
|
|
kernel_size,
|
|
stride=1,
|
|
dilation=d,
|
|
padding=get_padding(kernel_size, d),
|
|
)
|
|
for i, d in enumerate(dilation)
|
|
}
|
|
|
|
self.convs2 = {
|
|
i: nn.Conv1d(
|
|
channels,
|
|
channels,
|
|
kernel_size,
|
|
stride=1,
|
|
dilation=1,
|
|
padding=get_padding(kernel_size, 1),
|
|
)
|
|
for i in range(len(dilation))
|
|
}
|
|
|
|
self.acts1 = {i: Activation1d(act_cls(channels)) for i in range(len(dilation))}
|
|
self.acts2 = {i: Activation1d(act_cls(channels)) for i in range(len(dilation))}
|
|
|
|
def __call__(self, x: mx.array) -> mx.array:
|
|
for i in range(len(self.convs1)):
|
|
xt = self.acts1[i](x)
|
|
xt = self.convs1[i](xt)
|
|
xt = self.acts2[i](xt)
|
|
xt = self.convs2[i](xt)
|
|
x = x + xt
|
|
return x
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# STFT and MelSTFT (for BWE)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class STFTFn(nn.Module):
|
|
"""STFT via conv1d with precomputed DFT x window bases (loaded from checkpoint)."""
|
|
|
|
def __init__(self, filter_length: int, hop_length: int, win_length: int) -> None:
|
|
super().__init__()
|
|
self.hop_length = hop_length
|
|
self.win_length = win_length
|
|
n_freqs = filter_length // 2 + 1
|
|
# Buffers loaded from checkpoint - PyTorch format (n_freqs*2, 1, filter_length)
|
|
self.forward_basis = mx.zeros((n_freqs * 2, 1, filter_length))
|
|
self.inverse_basis = mx.zeros((n_freqs * 2, 1, filter_length))
|
|
|
|
def __call__(self, y: mx.array) -> Tuple[mx.array, mx.array]:
|
|
"""Compute magnitude and phase from waveform.
|
|
|
|
Args:
|
|
y: (B, T) waveform
|
|
|
|
Returns:
|
|
magnitude: (B, n_freqs, T_frames)
|
|
phase: (B, n_freqs, T_frames)
|
|
"""
|
|
if y.ndim == 2:
|
|
y = mx.expand_dims(y, -1) # (B, T, 1)
|
|
|
|
left_pad = max(0, self.win_length - self.hop_length)
|
|
if left_pad > 0:
|
|
first = mx.repeat(y[:, :1, :], left_pad, axis=1)
|
|
y = mx.concatenate([first, y], axis=1)
|
|
|
|
# forward_basis: (514, 1, 512) PyTorch format -> (514, 512, 1) MLX
|
|
basis = mx.transpose(
|
|
self.forward_basis.astype(y.dtype), (0, 2, 1)
|
|
) # (514, K, 1)
|
|
|
|
# Conv1d: (B, T, 1) * (514, K, 1) -> (B, T_frames, 514)
|
|
spec = mx.conv1d(y, basis, stride=self.hop_length)
|
|
|
|
# Split real and imaginary
|
|
n_freqs = spec.shape[-1] // 2
|
|
real = spec[..., :n_freqs]
|
|
imag = spec[..., n_freqs:]
|
|
|
|
magnitude = mx.sqrt(real**2 + imag**2)
|
|
phase = mx.arctan2(imag.astype(mx.float32), real.astype(mx.float32)).astype(
|
|
real.dtype
|
|
)
|
|
|
|
# Output: (B, T_frames, n_freqs) in MLX channels-last
|
|
return magnitude, phase
|
|
|
|
|
|
class MelSTFT(nn.Module):
|
|
"""Causal log-mel spectrogram from precomputed STFT bases."""
|
|
|
|
def __init__(
|
|
self, filter_length: int, hop_length: int, win_length: int, n_mel_channels: int
|
|
) -> None:
|
|
super().__init__()
|
|
self.stft_fn = STFTFn(filter_length, hop_length, win_length)
|
|
n_freqs = filter_length // 2 + 1
|
|
self.mel_basis = mx.zeros((n_mel_channels, n_freqs))
|
|
|
|
def mel_spectrogram(self, y: mx.array) -> mx.array:
|
|
"""Compute log-mel spectrogram.
|
|
|
|
Args:
|
|
y: (B, T) waveform
|
|
|
|
Returns:
|
|
log_mel: (B, n_mels, T_frames) in channels-first for compatibility
|
|
"""
|
|
magnitude, phase = self.stft_fn(y)
|
|
# magnitude: (B, T_frames, n_freqs)
|
|
mel = (
|
|
magnitude @ self.mel_basis.astype(magnitude.dtype).T
|
|
) # (B, T_frames, n_mels)
|
|
log_mel = mx.log(mx.clip(mel, 1e-5, None))
|
|
# Transpose to (B, n_mels, T_frames) for compatibility with vocoder input format
|
|
return mx.transpose(log_mel, (0, 2, 1))
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Vocoder (supports both HiFi-GAN and BigVGAN v2)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class Vocoder(nn.Module):
|
|
"""Vocoder for mel-to-waveform synthesis.
|
|
|
|
Supports resblock="1" (HiFi-GAN / LTX-2) and resblock="AMP1" (BigVGAN v2 / LTX-2.3).
|
|
"""
|
|
|
|
def __init__(self, config: VocoderModelConfig) -> None:
|
|
super().__init__()
|
|
|
|
self.output_sampling_rate = config.output_sample_rate
|
|
self.num_kernels = len(config.resblock_kernel_sizes)
|
|
self.num_upsamples = len(config.upsample_rates)
|
|
self.upsample_rates = config.upsample_rates
|
|
self.is_amp = config.resblock == "AMP1"
|
|
self.use_tanh_at_final = config.use_tanh_at_final
|
|
self.apply_final_activation = config.apply_final_activation
|
|
|
|
in_channels = 128 if config.stereo else 64
|
|
self.conv_pre = nn.Conv1d(
|
|
in_channels,
|
|
config.upsample_initial_channel,
|
|
kernel_size=7,
|
|
stride=1,
|
|
padding=3,
|
|
)
|
|
|
|
# Upsampling layers
|
|
self.ups = {}
|
|
for i, (stride, kernel_size) in enumerate(
|
|
zip(config.upsample_rates, config.upsample_kernel_sizes)
|
|
):
|
|
in_ch = config.upsample_initial_channel // (2**i)
|
|
out_ch = config.upsample_initial_channel // (2 ** (i + 1))
|
|
self.ups[i] = nn.ConvTranspose1d(
|
|
in_ch,
|
|
out_ch,
|
|
kernel_size=kernel_size,
|
|
stride=stride,
|
|
padding=(kernel_size - stride) // 2,
|
|
)
|
|
|
|
# Residual blocks
|
|
if self.is_amp:
|
|
self.resblocks = {}
|
|
block_idx = 0
|
|
for i in range(len(self.ups)):
|
|
ch = config.upsample_initial_channel // (2 ** (i + 1))
|
|
for kernel_size, dilations in zip(
|
|
config.resblock_kernel_sizes, config.resblock_dilation_sizes
|
|
):
|
|
self.resblocks[block_idx] = AMPBlock1(
|
|
ch,
|
|
kernel_size,
|
|
tuple(dilations),
|
|
activation=config.activation,
|
|
)
|
|
block_idx += 1
|
|
else:
|
|
resblock_class = ResBlock1 if config.resblock == "1" else ResBlock2
|
|
self.resblocks = {}
|
|
block_idx = 0
|
|
for i in range(len(self.ups)):
|
|
ch = config.upsample_initial_channel // (2 ** (i + 1))
|
|
for kernel_size, dilations in zip(
|
|
config.resblock_kernel_sizes, config.resblock_dilation_sizes
|
|
):
|
|
self.resblocks[block_idx] = resblock_class(
|
|
ch, kernel_size, tuple(dilations)
|
|
)
|
|
block_idx += 1
|
|
|
|
final_channels = config.upsample_initial_channel // (
|
|
2 ** len(config.upsample_rates)
|
|
)
|
|
|
|
# Post-activation
|
|
if self.is_amp:
|
|
act_cls = SnakeBeta if config.activation == "snakebeta" else Snake
|
|
self.act_post = Activation1d(act_cls(final_channels))
|
|
|
|
# Final conv
|
|
out_channels = 2 if config.stereo else 1
|
|
self.conv_post = nn.Conv1d(
|
|
final_channels,
|
|
out_channels,
|
|
kernel_size=7,
|
|
stride=1,
|
|
padding=3,
|
|
bias=config.use_bias_at_final,
|
|
)
|
|
|
|
self.upsample_factor = math.prod(config.upsample_rates)
|
|
|
|
def __call__(self, x: mx.array) -> mx.array:
|
|
"""Forward pass.
|
|
|
|
Args:
|
|
x: Mel spectrogram (B, C, T, mel_bins) for stereo or (B, T, mel_bins) mono.
|
|
|
|
Returns:
|
|
Waveform (B, out_channels, T_audio) in channels-first format.
|
|
"""
|
|
# (B, C, T, mel) -> (B, C, mel, T)
|
|
x = mx.transpose(x, (0, 1, 3, 2))
|
|
|
|
if x.ndim == 4: # stereo: (B, 2, mel, T) -> (B, 2*mel, T)
|
|
b, s, c, t = x.shape
|
|
x = x.reshape(b, s * c, t)
|
|
|
|
# Channels-first (B, C, T) -> channels-last (B, T, C) for MLX conv
|
|
x = mx.transpose(x, (0, 2, 1))
|
|
|
|
x = self.conv_pre(x)
|
|
|
|
for i in range(self.num_upsamples):
|
|
if not self.is_amp:
|
|
x = leaky_relu(x, LRELU_SLOPE)
|
|
x = self.ups[i](x)
|
|
|
|
start = i * self.num_kernels
|
|
end = start + self.num_kernels
|
|
|
|
block_outputs = mx.stack(
|
|
[self.resblocks[idx](x) for idx in range(start, end)],
|
|
axis=0,
|
|
)
|
|
x = mx.mean(block_outputs, axis=0)
|
|
|
|
if self.is_amp:
|
|
x = self.act_post(x)
|
|
else:
|
|
x = nn.leaky_relu(x)
|
|
|
|
x = self.conv_post(x)
|
|
|
|
if self.apply_final_activation:
|
|
x = mx.tanh(x) if self.use_tanh_at_final else mx.clip(x, -1, 1)
|
|
|
|
# Back to channels-first (B, T, C) -> (B, C, T)
|
|
x = mx.transpose(x, (0, 2, 1))
|
|
return x
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# VocoderWithBWE (Bandwidth Extension)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class VocoderWithBWE(nn.Module):
|
|
"""Vocoder + bandwidth extension upsampling (16kHz -> 48kHz).
|
|
|
|
Chains a base vocoder with a BWE generator that predicts a residual
|
|
added to a sinc-resampled skip connection.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
vocoder: Vocoder,
|
|
bwe_generator: Vocoder,
|
|
mel_stft: MelSTFT,
|
|
input_sampling_rate: int = 16000,
|
|
output_sampling_rate: int = 48000,
|
|
hop_length: int = 80,
|
|
) -> None:
|
|
super().__init__()
|
|
self.vocoder = vocoder
|
|
self.bwe_generator = bwe_generator
|
|
self.mel_stft = mel_stft
|
|
self.input_sampling_rate = input_sampling_rate
|
|
self.output_sampling_rate = output_sampling_rate
|
|
self.hop_length = hop_length
|
|
# Hann-windowed sinc resampler (not stored in checkpoint)
|
|
self.resampler = UpSample1d(
|
|
ratio=output_sampling_rate // input_sampling_rate,
|
|
window_type="hann",
|
|
)
|
|
|
|
@property
|
|
def output_sample_rate(self) -> int:
|
|
return self.output_sampling_rate
|
|
|
|
def _compute_mel(self, audio: mx.array) -> mx.array:
|
|
"""Compute log-mel spectrogram from waveform.
|
|
|
|
Args:
|
|
audio: (B, C, T) waveform in channels-first
|
|
|
|
Returns:
|
|
mel: (B, C, n_mels, T_frames)
|
|
"""
|
|
batch, n_channels, _ = audio.shape
|
|
flat = audio.reshape(batch * n_channels, -1) # (B*C, T)
|
|
mel = self.mel_stft.mel_spectrogram(flat) # (B*C, n_mels, T_frames)
|
|
return mel.reshape(batch, n_channels, mel.shape[1], mel.shape[2])
|
|
|
|
def __call__(self, mel_spec: mx.array) -> mx.array:
|
|
"""Run vocoder + BWE.
|
|
|
|
Args:
|
|
mel_spec: Mel spectrogram, same format as Vocoder.forward input.
|
|
|
|
Returns:
|
|
Waveform (B, out_channels, T_audio) at output_sampling_rate.
|
|
"""
|
|
x = self.vocoder(mel_spec) # (B, C, T) at input_sampling_rate
|
|
_, _, length_low_rate = x.shape
|
|
output_length = (
|
|
length_low_rate * self.output_sampling_rate // self.input_sampling_rate
|
|
)
|
|
|
|
# Pad to hop_length multiple
|
|
remainder = length_low_rate % self.hop_length
|
|
if remainder != 0:
|
|
pad_amount = self.hop_length - remainder
|
|
x = mx.pad(x, [(0, 0), (0, 0), (0, pad_amount)])
|
|
|
|
# Compute mel from vocoder output: (B, C, n_mels, T_frames)
|
|
mel = self._compute_mel(x)
|
|
|
|
# BWE expects (B, C, T_frames, mel_bins) -> transpose last two dims
|
|
mel_for_bwe = mx.transpose(mel, (0, 1, 3, 2)) # (B, C, T_frames, n_mels)
|
|
residual = self.bwe_generator(mel_for_bwe) # (B, C, T_high)
|
|
|
|
# Sinc upsample skip connection
|
|
# resampler expects (N, L, C): transpose from (B, C, T) -> (B, T, C)
|
|
x_for_resample = mx.transpose(x, (0, 2, 1))
|
|
skip = self.resampler(x_for_resample)
|
|
skip = mx.transpose(skip, (0, 2, 1)) # back to (B, C, T)
|
|
|
|
return mx.clip(residual + skip, -1, 1)[..., :output_length]
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Factory / from_pretrained
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def load_vocoder(model_path: Path) -> nn.Module:
|
|
"""Load vocoder from pretrained model directory.
|
|
|
|
Automatically detects whether to load a simple Vocoder or VocoderWithBWE.
|
|
"""
|
|
import json
|
|
|
|
config_path = model_path / "config.json"
|
|
if not config_path.exists():
|
|
raise FileNotFoundError(f"No config.json found in {model_path}")
|
|
|
|
with open(config_path) as f:
|
|
config_dict = json.load(f)
|
|
|
|
weights = mx.load(str(model_path / "model.safetensors"))
|
|
|
|
has_bwe = config_dict.get("has_bwe_generator", False)
|
|
|
|
if has_bwe:
|
|
return _load_vocoder_with_bwe(config_dict, weights)
|
|
else:
|
|
config = VocoderModelConfig.from_dict(config_dict)
|
|
model = Vocoder(config)
|
|
model.load_weights(list(weights.items()), strict=True)
|
|
return model
|
|
|
|
|
|
def _load_vocoder_with_bwe(config_dict: dict, weights: dict) -> VocoderWithBWE:
|
|
"""Load VocoderWithBWE from config and weights."""
|
|
# Build vocoder from config
|
|
vocoder_cfg = config_dict.get("vocoder", {})
|
|
vocoder_config = VocoderModelConfig.from_dict(vocoder_cfg)
|
|
vocoder = Vocoder(vocoder_config)
|
|
|
|
# Build BWE generator from config
|
|
bwe_cfg = config_dict.get("bwe", {})
|
|
bwe_config = VocoderModelConfig.from_dict(bwe_cfg)
|
|
bwe_config.apply_final_activation = False
|
|
bwe_generator = Vocoder(bwe_config)
|
|
|
|
# MelSTFT from weight shapes
|
|
stft_basis = weights.get("mel_stft.stft_fn.forward_basis")
|
|
filter_length = stft_basis.shape[2] if stft_basis is not None else 512
|
|
mel_basis = weights.get("mel_stft.mel_basis")
|
|
n_mel_channels = mel_basis.shape[0] if mel_basis is not None else 64
|
|
|
|
hop_length = bwe_cfg.get("hop_length", 80)
|
|
input_sr = bwe_cfg.get("input_sampling_rate", 16000)
|
|
output_sr = bwe_cfg.get("output_sampling_rate", 48000)
|
|
|
|
mel_stft = MelSTFT(
|
|
filter_length=filter_length,
|
|
hop_length=hop_length,
|
|
win_length=filter_length,
|
|
n_mel_channels=n_mel_channels,
|
|
)
|
|
|
|
model = VocoderWithBWE(
|
|
vocoder=vocoder,
|
|
bwe_generator=bwe_generator,
|
|
mel_stft=mel_stft,
|
|
input_sampling_rate=input_sr,
|
|
output_sampling_rate=output_sr,
|
|
hop_length=hop_length,
|
|
)
|
|
|
|
model.load_weights(list(weights.items()), strict=False)
|
|
return model
|