From 11b4628853c9e79dbc4577b01ecb02cfbc78255d Mon Sep 17 00:00:00 2001 From: Hubert Siuzdak <35269911+hubertsiuzdak@users.noreply.github.com> Date: Sat, 14 Oct 2023 02:34:27 +0200 Subject: [PATCH] Replace multi-resolution discriminator; update AdamW default config (#30) * New multi-resolution discriminator adopted from DAC * Default optimizer params * Bump version * Update README.md --- README.md | 10 ++-- configs/vocos-encodec.yaml | 2 +- configs/vocos-imdct.yaml | 2 +- configs/vocos-resnet.yaml | 2 +- configs/vocos.yaml | 2 +- vocos/__init__.py | 2 +- vocos/discriminators.py | 106 +++++++++++++++++++------------------ vocos/experiment.py | 4 +- 8 files changed, 64 insertions(+), 66 deletions(-) diff --git a/README.md b/README.md index 23a28b9..54a9612 100644 --- a/README.md +++ b/README.md @@ -82,14 +82,10 @@ See [example notebook](notebooks%2FBark%2BVocos.ipynb). ## Pre-trained models -The provided models were trained up to 2.5 million generator iterations, which resulted in slightly better objective -scores -compared to those reported in the paper. - | Model Name | Dataset | Training Iterations | Parameters -|-------------------------------------------------------------------------------------|---------------|---------------------|------------| -| [charactr/vocos-mel-24khz](https://huggingface.co/charactr/vocos-mel-24khz) | LibriTTS | 2.5 M | 13.5 M -| [charactr/vocos-encodec-24khz](https://huggingface.co/charactr/vocos-encodec-24khz) | DNS Challenge | 2.5 M | 7.9 M +|-------------------------------------------------------------------------------------|---------------|-------------------|------------| +| [charactr/vocos-mel-24khz](https://huggingface.co/charactr/vocos-mel-24khz) | LibriTTS | 1M | 13.5M +| [charactr/vocos-encodec-24khz](https://huggingface.co/charactr/vocos-encodec-24khz) | DNS Challenge | 2M | 7.9M ## Training diff --git a/configs/vocos-encodec.yaml b/configs/vocos-encodec.yaml index 775a52c..3d696ab 100644 --- a/configs/vocos-encodec.yaml +++ b/configs/vocos-encodec.yaml @@ -22,7 +22,7 @@ model: class_path: vocos.experiment.VocosEncodecExp init_args: sample_rate: 24000 - initial_learning_rate: 2e-4 + initial_learning_rate: 5e-4 mel_loss_coeff: 45 mrd_loss_coeff: 1.0 num_warmup_steps: 0 # Optimizers warmup steps diff --git a/configs/vocos-imdct.yaml b/configs/vocos-imdct.yaml index 7bdc5cf..6b8f9b5 100644 --- a/configs/vocos-imdct.yaml +++ b/configs/vocos-imdct.yaml @@ -22,7 +22,7 @@ model: class_path: vocos.experiment.VocosExp init_args: sample_rate: 24000 - initial_learning_rate: 2e-4 + initial_learning_rate: 5e-4 mel_loss_coeff: 45 mrd_loss_coeff: 0.1 num_warmup_steps: 0 # Optimizers warmup steps diff --git a/configs/vocos-resnet.yaml b/configs/vocos-resnet.yaml index 783272a..db7d25a 100644 --- a/configs/vocos-resnet.yaml +++ b/configs/vocos-resnet.yaml @@ -22,7 +22,7 @@ model: class_path: vocos.experiment.VocosExp init_args: sample_rate: 24000 - initial_learning_rate: 2e-4 + initial_learning_rate: 5e-4 mel_loss_coeff: 45 mrd_loss_coeff: 0.1 num_warmup_steps: 0 # Optimizers warmup steps diff --git a/configs/vocos.yaml b/configs/vocos.yaml index f2f4181..9671ae3 100644 --- a/configs/vocos.yaml +++ b/configs/vocos.yaml @@ -22,7 +22,7 @@ model: class_path: vocos.experiment.VocosExp init_args: sample_rate: 24000 - initial_learning_rate: 2e-4 + initial_learning_rate: 5e-4 mel_loss_coeff: 45 mrd_loss_coeff: 0.1 num_warmup_steps: 0 # Optimizers warmup steps diff --git a/vocos/__init__.py b/vocos/__init__.py index 093920f..928cc71 100644 --- a/vocos/__init__.py +++ b/vocos/__init__.py @@ -1,4 +1,4 @@ from vocos.pretrained import Vocos -__version__ = "0.0.4" +__version__ = "0.1.0" diff --git a/vocos/discriminators.py b/vocos/discriminators.py index 64ab318..2c62574 100644 --- a/vocos/discriminators.py +++ b/vocos/discriminators.py @@ -1,12 +1,11 @@ from typing import List, Optional, Tuple import torch +from einops import rearrange from torch import nn from torch.nn import Conv2d from torch.nn.utils import weight_norm - -PeriodsType = Tuple[int, ...] -ResolutionType = Tuple[int, int, int] +from torchaudio.transforms import Spectrogram class MultiPeriodDiscriminator(nn.Module): @@ -20,7 +19,7 @@ class MultiPeriodDiscriminator(nn.Module): Defaults to None. """ - def __init__(self, periods: PeriodsType = (2, 3, 5, 7, 11), num_embeddings: Optional[int] = None): + def __init__(self, periods: Tuple[int, ...] = (2, 3, 5, 7, 11), num_embeddings: Optional[int] = None): super().__init__() self.discriminators = nn.ModuleList([DiscriminatorP(period=p, num_embeddings=num_embeddings) for p in periods]) @@ -104,30 +103,26 @@ def forward( class MultiResolutionDiscriminator(nn.Module): def __init__( self, - resolutions: Tuple[ResolutionType, ResolutionType, ResolutionType] = ( - (1024, 256, 1024), - (2048, 512, 2048), - (512, 128, 512), - ), + fft_sizes: Tuple[int, ...] = (2048, 1024, 512), num_embeddings: Optional[int] = None, ): """ - Multi-Resolution Discriminator module adapted from https://github.com/mindslab-ai/univnet. + Multi-Resolution Discriminator module adapted from https://github.com/descriptinc/descript-audio-codec. Additionally, it allows incorporating conditional information with a learned embeddings table. Args: - resolutions (tuple[tuple[int, int, int]]): Tuple of resolutions for each discriminator. - Each resolution should be a tuple of (n_fft, hop_length, win_length). + fft_sizes (tuple[int]): Tuple of window lengths for FFT. Defaults to (2048, 1024, 512). num_embeddings (int, optional): Number of embeddings. None means non-conditional discriminator. Defaults to None. """ + super().__init__() self.discriminators = nn.ModuleList( - [DiscriminatorR(resolution=r, num_embeddings=num_embeddings) for r in resolutions] + [DiscriminatorR(window_length=w, num_embeddings=num_embeddings) for w in fft_sizes] ) def forward( - self, y: torch.Tensor, y_hat: torch.Tensor, bandwidth_id: Optional[torch.Tensor] = None + self, y: torch.Tensor, y_hat: torch.Tensor, bandwidth_id: torch.Tensor = None ) -> Tuple[List[torch.Tensor], List[torch.Tensor], List[List[torch.Tensor]], List[List[torch.Tensor]]]: y_d_rs = [] y_d_gs = [] @@ -148,40 +143,62 @@ def forward( class DiscriminatorR(nn.Module): def __init__( self, - resolution: Tuple[int, int, int], - channels: int = 64, - in_channels: int = 1, + window_length: int, num_embeddings: Optional[int] = None, - lrelu_slope: float = 0.1, + channels: int = 32, + hop_factor: float = 0.25, + bands: Tuple[Tuple[float, float], ...] = ((0.0, 0.1), (0.1, 0.25), (0.25, 0.5), (0.5, 0.75), (0.75, 1.0)), ): super().__init__() - self.resolution = resolution - self.in_channels = in_channels - self.lrelu_slope = lrelu_slope - self.convs = nn.ModuleList( + self.window_length = window_length + self.hop_factor = hop_factor + self.spec_fn = Spectrogram( + n_fft=window_length, hop_length=int(window_length * hop_factor), win_length=window_length, power=None + ) + n_fft = window_length // 2 + 1 + bands = [(int(b[0] * n_fft), int(b[1] * n_fft)) for b in bands] + self.bands = bands + convs = lambda: nn.ModuleList( [ - weight_norm(nn.Conv2d(in_channels, channels, kernel_size=(7, 5), stride=(2, 2), padding=(3, 2))), - weight_norm(nn.Conv2d(channels, channels, kernel_size=(5, 3), stride=(2, 1), padding=(2, 1))), - weight_norm(nn.Conv2d(channels, channels, kernel_size=(5, 3), stride=(2, 2), padding=(2, 1))), - weight_norm(nn.Conv2d(channels, channels, kernel_size=3, stride=(2, 1), padding=1)), - weight_norm(nn.Conv2d(channels, channels, kernel_size=3, stride=(2, 2), padding=1)), + weight_norm(nn.Conv2d(2, channels, (3, 9), (1, 1), padding=(1, 4))), + weight_norm(nn.Conv2d(channels, channels, (3, 9), (1, 2), padding=(1, 4))), + weight_norm(nn.Conv2d(channels, channels, (3, 9), (1, 2), padding=(1, 4))), + weight_norm(nn.Conv2d(channels, channels, (3, 9), (1, 2), padding=(1, 4))), + weight_norm(nn.Conv2d(channels, channels, (3, 3), (1, 1), padding=(1, 1))), ] ) + self.band_convs = nn.ModuleList([convs() for _ in range(len(self.bands))]) + if num_embeddings is not None: self.emb = torch.nn.Embedding(num_embeddings=num_embeddings, embedding_dim=channels) torch.nn.init.zeros_(self.emb.weight) - self.conv_post = weight_norm(nn.Conv2d(channels, 1, (3, 3), padding=(1, 1))) - def forward( - self, x: torch.Tensor, cond_embedding_id: Optional[torch.Tensor] = None - ) -> Tuple[torch.Tensor, List[torch.Tensor]]: + self.conv_post = weight_norm(nn.Conv2d(channels, 1, (3, 3), (1, 1), padding=(1, 1))) + + def spectrogram(self, x): + # Remove DC offset + x = x - x.mean(dim=-1, keepdims=True) + # Peak normalize the volume of input audio + x = 0.8 * x / (x.abs().max(dim=-1, keepdim=True)[0] + 1e-9) + x = self.spec_fn(x) + x = torch.view_as_real(x) + x = rearrange(x, "b f t c -> b c t f") + # Split into bands + x_bands = [x[..., b[0] : b[1]] for b in self.bands] + return x_bands + + def forward(self, x: torch.Tensor, cond_embedding_id: torch.Tensor = None): + x_bands = self.spectrogram(x) fmap = [] - x = self.spectrogram(x) - x = x.unsqueeze(1) - for l in self.convs: - x = l(x) - x = torch.nn.functional.leaky_relu(x, self.lrelu_slope) - fmap.append(x) + x = [] + for band, stack in zip(x_bands, self.band_convs): + for i, layer in enumerate(stack): + band = layer(band) + band = torch.nn.functional.leaky_relu(band, 0.1) + if i > 0: + fmap.append(band) + x.append(band) + x = torch.cat(x, dim=-1) if cond_embedding_id is not None: emb = self.emb(cond_embedding_id) h = (emb.view(1, -1, 1, 1) * x).sum(dim=1, keepdims=True) @@ -190,20 +207,5 @@ def forward( x = self.conv_post(x) fmap.append(x) x += h - x = torch.flatten(x, 1, -1) return x, fmap - - def spectrogram(self, x: torch.Tensor) -> torch.Tensor: - n_fft, hop_length, win_length = self.resolution - magnitude_spectrogram = torch.stft( - x, - n_fft=n_fft, - hop_length=hop_length, - win_length=win_length, - window=None, # interestingly rectangular window kind of works here - center=True, - return_complex=True, - ).abs() - - return magnitude_spectrogram diff --git a/vocos/experiment.py b/vocos/experiment.py index 22857d9..191c2fb 100644 --- a/vocos/experiment.py +++ b/vocos/experiment.py @@ -78,8 +78,8 @@ def configure_optimizers(self): {"params": self.head.parameters()}, ] - opt_disc = torch.optim.AdamW(disc_params, lr=self.hparams.initial_learning_rate) - opt_gen = torch.optim.AdamW(gen_params, lr=self.hparams.initial_learning_rate) + opt_disc = torch.optim.AdamW(disc_params, lr=self.hparams.initial_learning_rate, betas=(0.8, 0.9)) + opt_gen = torch.optim.AdamW(gen_params, lr=self.hparams.initial_learning_rate, betas=(0.8, 0.9)) max_steps = self.trainer.max_steps // 2 # Max steps per optimizer scheduler_disc = transformers.get_cosine_schedule_with_warmup(