import os
from dataclasses import dataclass, field
from itertools import chain
from pathlib import Path
from typing import Dict, List, Optional, Tuple, Union
import numpy as np
import torch
import torch.distributed as dist
import torchaudio
from coqpit import Coqpit
from librosa.filters import mel as librosa_mel_fn
from torch import nn
from torch.cuda.amp.autocast_mode import autocast
from torch.nn import functional as F
from torch.utils.data import DataLoader
from torch.utils.data.sampler import WeightedRandomSampler
from trainer.torch import DistributedSampler, DistributedSamplerWrapper
from trainer.trainer_utils import get_optimizer, get_scheduler
from TTS.tts.datasets.dataset import F0Dataset, TTSDataset, _parse_sample
from TTS.tts.layers.delightful_tts.acoustic_model import AcousticModel
from TTS.tts.layers.losses import ForwardSumLoss, VitsDiscriminatorLoss
from TTS.tts.layers.vits.discriminator import VitsDiscriminator
from TTS.tts.models.base_tts import BaseTTSE2E
from TTS.tts.utils.helpers import average_over_durations, compute_attn_prior, rand_segments, segment, sequence_mask
from TTS.tts.utils.speakers import SpeakerManager
from TTS.tts.utils.text.tokenizer import TTSTokenizer
from TTS.tts.utils.visual import plot_alignment, plot_avg_pitch, plot_pitch, plot_spectrogram
from TTS.utils.audio.numpy_transforms import build_mel_basis, compute_f0
from TTS.utils.audio.numpy_transforms import db_to_amp as db_to_amp_numpy
from TTS.utils.audio.numpy_transforms import mel_to_wav as mel_to_wav_numpy
from TTS.utils.audio.processor import AudioProcessor
from TTS.utils.io import load_fsspec
from TTS.vocoder.layers.losses import MultiScaleSTFTLoss
from TTS.vocoder.models.hifigan_generator import HifiganGenerator
from TTS.vocoder.utils.generic_utils import plot_results
def id_to_torch(aux_id, cuda=False):
if aux_id is not None:
aux_id = np.asarray(aux_id)
aux_id = torch.from_numpy(aux_id)
if cuda:
return aux_id.cuda()
return aux_id
def embedding_to_torch(d_vector, cuda=False):
if d_vector is not None:
d_vector = np.asarray(d_vector)
d_vector = torch.from_numpy(d_vector).float()
d_vector = d_vector.squeeze().unsqueeze(0)
if cuda:
return d_vector.cuda()
return d_vector
def numpy_to_torch(np_array, dtype, cuda=False):
if np_array is None:
return None
tensor = torch.as_tensor(np_array, dtype=dtype)
if cuda:
return tensor.cuda()
return tensor
def get_mask_from_lengths(lengths: torch.Tensor) -> torch.Tensor:
batch_size = lengths.shape[0]
max_len = torch.max(lengths).item()
ids = torch.arange(0, max_len, device=lengths.device).unsqueeze(0).expand(batch_size, -1)
mask = ids >= lengths.unsqueeze(1).expand(-1, max_len)
return mask
def pad(input_ele: List[torch.Tensor], max_len: int) -> torch.Tensor:
out_list = torch.jit.annotate(List[torch.Tensor], [])
for batch in input_ele:
if len(batch.shape) == 1:
one_batch_padded = F.pad(batch, (0, max_len - batch.size(0)), "constant", 0.0)
else:
one_batch_padded = F.pad(batch, (0, 0, 0, max_len - batch.size(0)), "constant", 0.0)
out_list.append(one_batch_padded)
out_padded = torch.stack(out_list)
return out_padded
def init_weights(m: nn.Module, mean: float = 0.0, std: float = 0.01):
classname = m.__class__.__name__
if classname.find("Conv") != -1:
m.weight.data.normal_(mean, std)
def stride_lens(lens: torch.Tensor, stride: int = 2) -> torch.Tensor:
return torch.ceil(lens / stride).int()
def initialize_embeddings(shape: Tuple[int]) -> torch.Tensor:
assert len(shape) == 2, "Can only initialize 2-D embedding matrices ..."
return torch.randn(shape) * np.sqrt(2 / shape[1])
# pylint: disable=redefined-outer-name
def calc_same_padding(kernel_size: int) -> Tuple[int, int]:
pad = kernel_size // 2
return (pad, pad - (kernel_size + 1) % 2)
hann_window = {}
mel_basis = {}
@torch.no_grad()
def weights_reset(m: nn.Module):
# check if the current module has reset_parameters and if it is reset the weight
reset_parameters = getattr(m, "reset_parameters", None)
if callable(reset_parameters):
m.reset_parameters()
def get_module_weights_sum(mdl: nn.Module):
dict_sums = {}
for name, w in mdl.named_parameters():
if "weight" in name:
value = w.data.sum().item()
dict_sums[name] = value
return dict_sums
def load_audio(file_path: str):
"""Load the audio file normalized in [-1, 1]
Return Shapes:
- x: :math:`[1, T]`
"""
x, sr = torchaudio.load(
file_path,
)
assert (x > 1).sum() + (x < -1).sum() == 0
return x, sr
def _amp_to_db(x, C=1, clip_val=1e-5):
return torch.log(torch.clamp(x, min=clip_val) * C)
def _db_to_amp(x, C=1):
return torch.exp(x) / C
def amp_to_db(magnitudes):
output = _amp_to_db(magnitudes)
return output
def db_to_amp(magnitudes):
output = _db_to_amp(magnitudes)
return output
def _wav_to_spec(y, n_fft, hop_length, win_length, center=False):
y = y.squeeze(1)
if torch.min(y) < -1.0:
print("min value is ", torch.min(y))
if torch.max(y) > 1.0:
print("max value is ", torch.max(y))
global hann_window # pylint: disable=global-statement
dtype_device = str(y.dtype) + "_" + str(y.device)
wnsize_dtype_device = str(win_length) + "_" + dtype_device
if wnsize_dtype_device not in hann_window:
hann_window[wnsize_dtype_device] = torch.hann_window(win_length).to(dtype=y.dtype, device=y.device)
y = torch.nn.functional.pad(
y.unsqueeze(1),
(int((n_fft - hop_length) / 2), int((n_fft - hop_length) / 2)),
mode="reflect",
)
y = y.squeeze(1)
spec = torch.stft(
y,
n_fft,
hop_length=hop_length,
win_length=win_length,
window=hann_window[wnsize_dtype_device],
center=center,
pad_mode="reflect",
normalized=False,
onesided=True,
return_complex=False,
)
return spec
def wav_to_spec(y, n_fft, hop_length, win_length, center=False):
"""
Args Shapes:
- y : :math:`[B, 1, T]`
Return Shapes:
- spec : :math:`[B,C,T]`
"""
spec = _wav_to_spec(y, n_fft, hop_length, win_length, center=center)
spec = torch.sqrt(spec.pow(2).sum(-1) + 1e-6)
return spec
def wav_to_energy(y, n_fft, hop_length, win_length, center=False):
spec = _wav_to_spec(y, n_fft, hop_length, win_length, center=center)
spec = torch.sqrt(spec.pow(2).sum(-1) + 1e-6)
return torch.norm(spec, dim=1, keepdim=True)
def name_mel_basis(spec, n_fft, fmax):
n_fft_len = f"{n_fft}_{fmax}_{spec.dtype}_{spec.device}"
return n_fft_len
def spec_to_mel(spec, n_fft, num_mels, sample_rate, fmin, fmax):
"""
Args Shapes:
- spec : :math:`[B,C,T]`
Return Shapes:
- mel : :math:`[B,C,T]`
"""
global mel_basis # pylint: disable=global-statement
mel_basis_key = name_mel_basis(spec, n_fft, fmax)
# pylint: disable=too-many-function-args
if mel_basis_key not in mel_basis:
# pylint: disable=missing-kwoa
mel = librosa_mel_fn(sample_rate, n_fft, num_mels, fmin, fmax)
mel_basis[mel_basis_key] = torch.from_numpy(mel).to(dtype=spec.dtype, device=spec.device)
mel = torch.matmul(mel_basis[mel_basis_key], spec)
mel = amp_to_db(mel)
return mel
def wav_to_mel(y, n_fft, num_mels, sample_rate, hop_length, win_length, fmin, fmax, center=False):
"""
Args Shapes:
- y : :math:`[B, 1, T_y]`
Return Shapes:
- spec : :math:`[B,C,T_spec]`
"""
y = y.squeeze(1)
if torch.min(y) < -1.0:
print("min value is ", torch.min(y))
if torch.max(y) > 1.0:
print("max value is ", torch.max(y))
global mel_basis, hann_window # pylint: disable=global-statement
mel_basis_key = name_mel_basis(y, n_fft, fmax)
wnsize_dtype_device = str(win_length) + "_" + str(y.dtype) + "_" + str(y.device)
if mel_basis_key not in mel_basis:
# pylint: disable=missing-kwoa
mel = librosa_mel_fn(
sr=sample_rate, n_fft=n_fft, n_mels=num_mels, fmin=fmin, fmax=fmax
) # pylint: disable=too-many-function-args
mel_basis[mel_basis_key] = torch.from_numpy(mel).to(dtype=y.dtype, device=y.device)
if wnsize_dtype_device not in hann_window:
hann_window[wnsize_dtype_device] = torch.hann_window(win_length).to(dtype=y.dtype, device=y.device)
y = torch.nn.functional.pad(
y.unsqueeze(1),
(int((n_fft - hop_length) / 2), int((n_fft - hop_length) / 2)),
mode="reflect",
)
y = y.squeeze(1)
spec = torch.stft(
y,
n_fft,
hop_length=hop_length,
win_length=win_length,
window=hann_window[wnsize_dtype_device],
center=center,
pad_mode="reflect",
normalized=False,
onesided=True,
return_complex=False,
)
spec = torch.sqrt(spec.pow(2).sum(-1) + 1e-6)
spec = torch.matmul(mel_basis[mel_basis_key], spec)
spec = amp_to_db(spec)
return spec
##############################
# DATASET
##############################
def get_attribute_balancer_weights(items: list, attr_name: str, multi_dict: dict = None):
"""Create balancer weight for torch WeightedSampler"""
attr_names_samples = np.array([item[attr_name] for item in items])
unique_attr_names = np.unique(attr_names_samples).tolist()
attr_idx = [unique_attr_names.index(l) for l in attr_names_samples]
attr_count = np.array([len(np.where(attr_names_samples == l)[0]) for l in unique_attr_names])
weight_attr = 1.0 / attr_count
dataset_samples_weight = np.array([weight_attr[l] for l in attr_idx])
dataset_samples_weight = dataset_samples_weight / np.linalg.norm(dataset_samples_weight)
if multi_dict is not None:
multiplier_samples = np.array([multi_dict.get(item[attr_name], 1.0) for item in items])
dataset_samples_weight *= multiplier_samples
return (
torch.from_numpy(dataset_samples_weight).float(),
unique_attr_names,
np.unique(dataset_samples_weight).tolist(),
)
class ForwardTTSE2eF0Dataset(F0Dataset):
"""Override F0Dataset to avoid slow computing of pitches"""
def __init__(
self,
ap,
samples: Union[List[List], List[Dict]],
verbose=False,
cache_path: str = None,
precompute_num_workers=0,
normalize_f0=True,
):
super().__init__(
samples=samples,
ap=ap,
verbose=verbose,
cache_path=cache_path,
precompute_num_workers=precompute_num_workers,
normalize_f0=normalize_f0,
)
def _compute_and_save_pitch(self, wav_file, pitch_file=None):
wav, _ = load_audio(wav_file)
f0 = compute_f0(
x=wav.numpy()[0],
sample_rate=self.ap.sample_rate,
hop_length=self.ap.hop_length,
pitch_fmax=self.ap.pitch_fmax,
pitch_fmin=self.ap.pitch_fmin,
win_length=self.ap.win_length,
)
# skip the last F0 value to align with the spectrogram
if wav.shape[1] % self.ap.hop_length != 0:
f0 = f0[:-1]
if pitch_file:
np.save(pitch_file, f0)
return f0
def compute_or_load(self, wav_file, audio_name):
"""
compute pitch and return a numpy array of pitch values
"""
pitch_file = self.create_pitch_file_path(audio_name, self.cache_path)
if not os.path.exists(pitch_file):
pitch = self._compute_and_save_pitch(wav_file=wav_file, pitch_file=pitch_file)
else:
pitch = np.load(pitch_file)
return pitch.astype(np.float32)
class ForwardTTSE2eDataset(TTSDataset):
def __init__(self, *args, **kwargs):
# don't init the default F0Dataset in TTSDataset
compute_f0 = kwargs.pop("compute_f0", False)
kwargs["compute_f0"] = False
self.attn_prior_cache_path = kwargs.pop("attn_prior_cache_path")
super().__init__(*args, **kwargs)
self.compute_f0 = compute_f0
self.pad_id = self.tokenizer.characters.pad_id
self.ap = kwargs["ap"]
if self.compute_f0:
self.f0_dataset = ForwardTTSE2eF0Dataset(
ap=self.ap,
samples=self.samples,
cache_path=kwargs["f0_cache_path"],
precompute_num_workers=kwargs["precompute_num_workers"],
)
if self.attn_prior_cache_path is not None:
os.makedirs(self.attn_prior_cache_path, exist_ok=True)
def __getitem__(self, idx):
item = self.samples[idx]
rel_wav_path = Path(item["audio_file"]).relative_to(item["root_path"]).with_suffix("")
rel_wav_path = str(rel_wav_path).replace("/", "_")
raw_text = item["text"]
wav, _ = load_audio(item["audio_file"])
wav_filename = os.path.basename(item["audio_file"])
try:
token_ids = self.get_token_ids(idx, item["text"])
except:
print(idx, item)
# pylint: disable=raise-missing-from
raise OSError
f0 = None
if self.compute_f0:
f0 = self.get_f0(idx)["f0"]
# after phonemization the text length may change
# this is a shameful 🤠hack to prevent longer phonemes
# TODO: find a better fix
if len(token_ids) > self.max_text_len or wav.shape[1] < self.min_audio_len:
self.rescue_item_idx += 1
return self.__getitem__(self.rescue_item_idx)
attn_prior = None
if self.attn_prior_cache_path is not None:
attn_prior = self.load_or_compute_attn_prior(token_ids, wav, rel_wav_path)
return {
"raw_text": raw_text,
"token_ids": token_ids,
"token_len": len(token_ids),
"wav": wav,
"pitch": f0,
"wav_file": wav_filename,
"speaker_name": item["speaker_name"],
"language_name": item["language"],
"attn_prior": attn_prior,
"audio_unique_name": item["audio_unique_name"],
}
def load_or_compute_attn_prior(self, token_ids, wav, rel_wav_path):
"""Load or compute and save the attention prior."""
attn_prior_file = os.path.join(self.attn_prior_cache_path, f"{rel_wav_path}.npy")
# pylint: disable=no-else-return
if os.path.exists(attn_prior_file):
return np.load(attn_prior_file)
else:
token_len = len(token_ids)
mel_len = wav.shape[1] // self.ap.hop_length
attn_prior = compute_attn_prior(token_len, mel_len)
np.save(attn_prior_file, attn_prior)
return attn_prior
@property
def lengths(self):
lens = []
for item in self.samples:
_, wav_file, *_ = _parse_sample(item)
audio_len = os.path.getsize(wav_file) / 16 * 8 # assuming 16bit audio
lens.append(audio_len)
return lens
def collate_fn(self, batch):
"""
Return Shapes:
- tokens: :math:`[B, T]`
- token_lens :math:`[B]`
- token_rel_lens :math:`[B]`
- pitch :math:`[B, T]`
- waveform: :math:`[B, 1, T]`
- waveform_lens: :math:`[B]`
- waveform_rel_lens: :math:`[B]`
- speaker_names: :math:`[B]`
- language_names: :math:`[B]`
- audiofile_paths: :math:`[B]`
- raw_texts: :math:`[B]`
- attn_prior: :math:`[[T_token, T_mel]]`
"""
B = len(batch)
batch = {k: [dic[k] for dic in batch] for k in batch[0]}
max_text_len = max([len(x) for x in batch["token_ids"]])
token_lens = torch.LongTensor(batch["token_len"])
token_rel_lens = token_lens / token_lens.max()
wav_lens = [w.shape[1] for w in batch["wav"]]
wav_lens = torch.LongTensor(wav_lens)
wav_lens_max = torch.max(wav_lens)
wav_rel_lens = wav_lens / wav_lens_max
pitch_padded = None
if self.compute_f0:
pitch_lens = [p.shape[0] for p in batch["pitch"]]
pitch_lens = torch.LongTensor(pitch_lens)
pitch_lens_max = torch.max(pitch_lens)
pitch_padded = torch.FloatTensor(B, 1, pitch_lens_max)
pitch_padded = pitch_padded.zero_() + self.pad_id
token_padded = torch.LongTensor(B, max_text_len)
wav_padded = torch.FloatTensor(B, 1, wav_lens_max)
token_padded = token_padded.zero_() + self.pad_id
wav_padded = wav_padded.zero_() + self.pad_id
for i in range(B):
token_ids = batch["token_ids"][i]
token_padded[i, : batch["token_len"][i]] = torch.LongTensor(token_ids)
wav = batch["wav"][i]
wav_padded[i, :, : wav.size(1)] = torch.FloatTensor(wav)
if self.compute_f0:
pitch = batch["pitch"][i]
pitch_padded[i, 0, : len(pitch)] = torch.FloatTensor(pitch)
return {
"text_input": token_padded,
"text_lengths": token_lens,
"text_rel_lens": token_rel_lens,
"pitch": pitch_padded,
"waveform": wav_padded, # (B x T)
"waveform_lens": wav_lens, # (B)
"waveform_rel_lens": wav_rel_lens,
"speaker_names": batch["speaker_name"],
"language_names": batch["language_name"],
"audio_unique_names": batch["audio_unique_name"],
"audio_files": batch["wav_file"],
"raw_text": batch["raw_text"],
"attn_priors": batch["attn_prior"] if batch["attn_prior"][0] is not None else None,
}
##############################
# CONFIG DEFINITIONS
##############################
@dataclass
class VocoderConfig(Coqpit):
resblock_type_decoder: str = "1"
resblock_kernel_sizes_decoder: List[int] = field(default_factory=lambda: [3, 7, 11])
resblock_dilation_sizes_decoder: List[List[int]] = field(default_factory=lambda: [[1, 3, 5], [1, 3, 5], [1, 3, 5]])
upsample_rates_decoder: List[int] = field(default_factory=lambda: [8, 8, 2, 2])
upsample_initial_channel_decoder: int = 512
upsample_kernel_sizes_decoder: List[int] = field(default_factory=lambda: [16, 16, 4, 4])
use_spectral_norm_discriminator: bool = False
upsampling_rates_discriminator: List[int] = field(default_factory=lambda: [4, 4, 4, 4])
periods_discriminator: List[int] = field(default_factory=lambda: [2, 3, 5, 7, 11])
pretrained_model_path: Optional[str] = None
@dataclass
class DelightfulTtsAudioConfig(Coqpit):
sample_rate: int = 22050
hop_length: int = 256
win_length: int = 1024
fft_size: int = 1024
mel_fmin: float = 0.0
mel_fmax: float = 8000
num_mels: int = 100
pitch_fmax: float = 640.0
pitch_fmin: float = 1.0
resample: bool = False
preemphasis: float = 0.0
ref_level_db: int = 20
do_sound_norm: bool = False
log_func: str = "np.log10"
do_trim_silence: bool = True
trim_db: int = 45
do_rms_norm: bool = False
db_level: float = None
power: float = 1.5
griffin_lim_iters: int = 60
spec_gain: int = 20
do_amp_to_db_linear: bool = True
do_amp_to_db_mel: bool = True
min_level_db: int = -100
max_norm: float = 4.0
@dataclass
class DelightfulTtsArgs(Coqpit):
num_chars: int = 100
spec_segment_size: int = 32
n_hidden_conformer_encoder: int = 512
n_layers_conformer_encoder: int = 6
n_heads_conformer_encoder: int = 8
dropout_conformer_encoder: float = 0.1
kernel_size_conv_mod_conformer_encoder: int = 7
kernel_size_depthwise_conformer_encoder: int = 7
lrelu_slope: float = 0.3
n_hidden_conformer_decoder: int = 512
n_layers_conformer_decoder: int = 6
n_heads_conformer_decoder: int = 8
dropout_conformer_decoder: float = 0.1
kernel_size_conv_mod_conformer_decoder: int = 11
kernel_size_depthwise_conformer_decoder: int = 11
bottleneck_size_p_reference_encoder: int = 4
bottleneck_size_u_reference_encoder: int = 512
ref_enc_filters_reference_encoder = [32, 32, 64, 64, 128, 128]
ref_enc_size_reference_encoder: int = 3
ref_enc_strides_reference_encoder = [1, 2, 1, 2, 1]
ref_enc_pad_reference_encoder = [1, 1]
ref_enc_gru_size_reference_encoder: int = 32
ref_attention_dropout_reference_encoder: float = 0.2
token_num_reference_encoder: int = 32
predictor_kernel_size_reference_encoder: int = 5
n_hidden_variance_adaptor: int = 512
kernel_size_variance_adaptor: int = 5
dropout_variance_adaptor: float = 0.5
n_bins_variance_adaptor: int = 256
emb_kernel_size_variance_adaptor: int = 3
use_speaker_embedding: bool = False
num_speakers: int = 0
speakers_file: str = None
d_vector_file: str = None
speaker_embedding_channels: int = 384
use_d_vector_file: bool = False
d_vector_dim: int = 0
freeze_vocoder: bool = False
freeze_text_encoder: bool = False
freeze_duration_predictor: bool = False
freeze_pitch_predictor: bool = False
freeze_energy_predictor: bool = False
freeze_basis_vectors_predictor: bool = False
freeze_decoder: bool = False
length_scale: float = 1.0
##############################
# MODEL DEFINITION
##############################
class DelightfulTTS(BaseTTSE2E):
"""
Paper::
https://arxiv.org/pdf/2110.12612.pdf
Paper Abstract::
This paper describes the Microsoft end-to-end neural text to speech (TTS) system: DelightfulTTS for Blizzard Challenge 2021.
The goal of this challenge is to synthesize natural and high-quality speech from text, and we approach this goal in two perspectives:
The first is to directly model and generate waveform in 48 kHz sampling rate, which brings higher perception quality than previous systems
with 16 kHz or 24 kHz sampling rate; The second is to model the variation information in speech through a systematic design, which improves
the prosody and naturalness. Specifically, for 48 kHz modeling, we predict 16 kHz mel-spectrogram in acoustic model, and
propose a vocoder called HiFiNet to directly generate 48 kHz waveform from predicted 16 kHz mel-spectrogram, which can better trade off training
efficiency, modelling stability and voice quality. We model variation information systematically from both explicit (speaker ID, language ID, pitch and duration) and
implicit (utterance-level and phoneme-level prosody) perspectives: 1) For speaker and language ID, we use lookup embedding in training and
inference; 2) For pitch and duration, we extract the values from paired text-speech data in training and use two predictors to predict the values in inference; 3)
For utterance-level and phoneme-level prosody, we use two reference encoders to extract the values in training, and use two separate predictors to predict the values in inference.
Additionally, we introduce an improved Conformer block to better model the local and global dependency in acoustic model. For task SH1, DelightfulTTS achieves 4.17 mean score in MOS test
and 4.35 in SMOS test, which indicates the effectiveness of our proposed system
Model training::
text --> ForwardTTS() --> spec_hat --> rand_seg_select()--> GANVocoder() --> waveform_seg
spec --------^
Examples:
>>> from TTS.tts.models.forward_tts_e2e import ForwardTTSE2e, ForwardTTSE2eConfig
>>> config = ForwardTTSE2eConfig()
>>> model = ForwardTTSE2e(config)
"""
# pylint: disable=dangerous-default-value
def __init__(
self,
config: Coqpit,
ap,
tokenizer: "TTSTokenizer" = None,
speaker_manager: SpeakerManager = None,
):
super().__init__(config=config, ap=ap, tokenizer=tokenizer, speaker_manager=speaker_manager)
self.ap = ap
self._set_model_args(config)
self.init_multispeaker(config)
self.binary_loss_weight = None
self.args.out_channels = self.config.audio.num_mels
self.args.num_mels = self.config.audio.num_mels
self.acoustic_model = AcousticModel(args=self.args, tokenizer=tokenizer, speaker_manager=speaker_manager)
self.waveform_decoder = HifiganGenerator(
self.config.audio.num_mels,
1,
self.config.vocoder.resblock_type_decoder,
self.config.vocoder.resblock_dilation_sizes_decoder,
self.config.vocoder.resblock_kernel_sizes_decoder,
self.config.vocoder.upsample_kernel_sizes_decoder,
self.config.vocoder.upsample_initial_channel_decoder,
self.config.vocoder.upsample_rates_decoder,
inference_padding=0,
# cond_channels=self.embedded_speaker_dim,
conv_pre_weight_norm=False,
conv_post_weight_norm=False,
conv_post_bias=False,
)
if self.config.init_discriminator:
self.disc = VitsDiscriminator(
use_spectral_norm=self.config.vocoder.use_spectral_norm_discriminator,
periods=self.config.vocoder.periods_discriminator,
)
@property
def device(self):
return next(self.parameters()).device
@property
def energy_scaler(self):
return self.acoustic_model.energy_scaler
@property
def length_scale(self):
return self.acoustic_model.length_scale
@length_scale.setter
def length_scale(self, value):
self.acoustic_model.length_scale = value
@property
def pitch_mean(self):
return self.acoustic_model.pitch_mean
@pitch_mean.setter
def pitch_mean(self, value):
self.acoustic_model.pitch_mean = value
@property
def pitch_std(self):
return self.acoustic_model.pitch_std
@pitch_std.setter
def pitch_std(self, value):
self.acoustic_model.pitch_std = value
@property
def mel_basis(self):
return build_mel_basis(
sample_rate=self.ap.sample_rate,
fft_size=self.ap.fft_size,
num_mels=self.ap.num_mels,
mel_fmax=self.ap.mel_fmax,
mel_fmin=self.ap.mel_fmin,
) # pylint: disable=function-redefined
def init_for_training(self) -> None:
self.train_disc = ( # pylint: disable=attribute-defined-outside-init
self.config.steps_to_start_discriminator <= 0
) # pylint: disable=attribute-defined-outside-init
self.update_energy_scaler = True # pylint: disable=attribute-defined-outside-init
def init_multispeaker(self, config: Coqpit):
"""Init for multi-speaker training.
Args:
config (Coqpit): Model configuration.
"""
self.embedded_speaker_dim = 0
self.num_speakers = self.args.num_speakers
self.audio_transform = None
if self.speaker_manager:
self.num_speakers = self.speaker_manager.num_speakers
self.args.num_speakers = self.speaker_manager.num_speakers
if self.args.use_speaker_embedding:
self._init_speaker_embedding()
if self.args.use_d_vector_file:
self._init_d_vector()
def _init_speaker_embedding(self):
# pylint: disable=attribute-defined-outside-init
if self.num_speakers > 0:
print(" > initialization of speaker-embedding layers.")
self.embedded_speaker_dim = self.args.speaker_embedding_channels
self.args.embedded_speaker_dim = self.args.speaker_embedding_channels
def _init_d_vector(self):
# pylint: disable=attribute-defined-outside-init
if hasattr(self, "emb_g"):
raise ValueError("[!] Speaker embedding layer already initialized before d_vector settings.")
self.embedded_speaker_dim = self.args.d_vector_dim
self.args.embedded_speaker_dim = self.args.d_vector_dim
def _freeze_layers(self):
if self.args.freeze_vocoder:
for param in self.vocoder.paramseters():
param.requires_grad = False
if self.args.freeze_text_encoder:
for param in self.text_encoder.parameters():
param.requires_grad = False
if self.args.freeze_duration_predictor:
for param in self.durarion_predictor.parameters():
param.requires_grad = False
if self.args.freeze_pitch_predictor:
for param in self.pitch_predictor.parameters():
param.requires_grad = False
if self.args.freeze_energy_predictor:
for param in self.energy_predictor.parameters():
param.requires_grad = False
if self.args.freeze_decoder:
for param in self.decoder.parameters():
param.requires_grad = False
def forward(
self,
x: torch.LongTensor,
x_lengths: torch.LongTensor,
spec_lengths: torch.LongTensor,
spec: torch.FloatTensor,
waveform: torch.FloatTensor,
pitch: torch.FloatTensor = None,
energy: torch.FloatTensor = None,
attn_priors: torch.FloatTensor = None,
d_vectors: torch.FloatTensor = None,
speaker_idx: torch.LongTensor = None,
) -> Dict:
"""Model's forward pass.
Args:
x (torch.LongTensor): Input character sequences.
x_lengths (torch.LongTensor): Input sequence lengths.
spec_lengths (torch.LongTensor): Spectrogram sequnce lengths. Defaults to None.
spec (torch.FloatTensor): Spectrogram frames. Only used when the alignment network is on. Defaults to None.
waveform (torch.FloatTensor): Waveform. Defaults to None.
pitch (torch.FloatTensor): Pitch values for each spectrogram frame. Only used when the pitch predictor is on. Defaults to None.
energy (torch.FloatTensor): Spectral energy values for each spectrogram frame. Only used when the energy predictor is on. Defaults to None.
attn_priors (torch.FloatTentrasor): Attention priors for the aligner network. Defaults to None.
aux_input (Dict): Auxiliary model inputs for multi-speaker training. Defaults to `{"d_vectors": 0, "speaker_ids": None}`.
Shapes:
- x: :math:`[B, T_max]`
- x_lengths: :math:`[B]`
- spec_lengths: :math:`[B]`
- spec: :math:`[B, T_max2, C_spec]`
- waveform: :math:`[B, 1, T_max2 * hop_length]`
- g: :math:`[B, C]`
- pitch: :math:`[B, 1, T_max2]`
- energy: :math:`[B, 1, T_max2]`
"""
encoder_outputs = self.acoustic_model(
tokens=x,
src_lens=x_lengths,
mel_lens=spec_lengths,
mels=spec,
pitches=pitch,
energies=energy,
attn_priors=attn_priors,
d_vectors=d_vectors,
speaker_idx=speaker_idx,
)
# use mel-spec from the decoder
vocoder_input = encoder_outputs["model_outputs"] # [B, T_max2, C_mel]
vocoder_input_slices, slice_ids = rand_segments(
x=vocoder_input.transpose(1, 2),
x_lengths=spec_lengths,
segment_size=self.args.spec_segment_size,
let_short_samples=True,
pad_short=True,
)
if encoder_outputs["spk_emb"] is not None:
g = encoder_outputs["spk_emb"].unsqueeze(-1)
else:
g = None
vocoder_output = self.waveform_decoder(x=vocoder_input_slices.detach(), g=g)
wav_seg = segment(
waveform,
slice_ids * self.ap.hop_length,
self.args.spec_segment_size * self.ap.hop_length,
pad_short=True,
)
model_outputs = {**encoder_outputs}
model_outputs["acoustic_model_outputs"] = encoder_outputs["model_outputs"]
model_outputs["model_outputs"] = vocoder_output
model_outputs["waveform_seg"] = wav_seg
model_outputs["slice_ids"] = slice_ids
return model_outputs
@torch.no_grad()
def inference(
self, x, aux_input={"d_vectors": None, "speaker_ids": None}, pitch_transform=None, energy_transform=None
):
encoder_outputs = self.acoustic_model.inference(
tokens=x,
d_vectors=aux_input["d_vectors"],
speaker_idx=aux_input["speaker_ids"],
pitch_transform=pitch_transform,
energy_transform=energy_transform,
p_control=None,
d_control=None,
)
vocoder_input = encoder_outputs["model_outputs"].transpose(1, 2) # [B, T_max2, C_mel] -> [B, C_mel, T_max2]
if encoder_outputs["spk_emb"] is not None:
g = encoder_outputs["spk_emb"].unsqueeze(-1)
else:
g = None
vocoder_output = self.waveform_decoder(x=vocoder_input, g=g)
model_outputs = {**encoder_outputs}
model_outputs["model_outputs"] = vocoder_output
return model_outputs
@torch.no_grad()
def inference_spec_decoder(self, x, aux_input={"d_vectors": None, "speaker_ids": None}):
encoder_outputs = self.acoustic_model.inference(
tokens=x,
d_vectors=aux_input["d_vectors"],
speaker_idx=aux_input["speaker_ids"],
)
model_outputs = {**encoder_outputs}
return model_outputs
def train_step(self, batch: dict, criterion: nn.Module, optimizer_idx: int):
if optimizer_idx == 0:
tokens = batch["text_input"]
token_lenghts = batch["text_lengths"]
mel = batch["mel_input"]
mel_lens = batch["mel_lengths"]
waveform = batch["waveform"] # [B, T, C] -> [B, C, T]
pitch = batch["pitch"]
d_vectors = batch["d_vectors"]
speaker_ids = batch["speaker_ids"]
attn_priors = batch["attn_priors"]
energy = batch["energy"]
# generator pass
outputs = self.forward(
x=tokens,
x_lengths=token_lenghts,
spec_lengths=mel_lens,
spec=mel,
waveform=waveform,
pitch=pitch,
energy=energy,
attn_priors=attn_priors,
d_vectors=d_vectors,
speaker_idx=speaker_ids,
)
# cache tensors for the generator pass
self.model_outputs_cache = outputs # pylint: disable=attribute-defined-outside-init
if self.train_disc:
# compute scores and features
scores_d_fake, _, scores_d_real, _ = self.disc(
outputs["model_outputs"].detach(), outputs["waveform_seg"]
)
# compute loss
with autocast(enabled=False): # use float32 for the criterion
loss_dict = criterion[optimizer_idx](
scores_disc_fake=scores_d_fake,
scores_disc_real=scores_d_real,
)
return outputs, loss_dict
return None, None
if optimizer_idx == 1:
mel = batch["mel_input"]
# compute melspec segment
with autocast(enabled=False):
mel_slice = segment(
mel.float(), self.model_outputs_cache["slice_ids"], self.args.spec_segment_size, pad_short=True
)
mel_slice_hat = wav_to_mel(
y=self.model_outputs_cache["model_outputs"].float(),
n_fft=self.ap.fft_size,
sample_rate=self.ap.sample_rate,
num_mels=self.ap.num_mels,
hop_length=self.ap.hop_length,
win_length=self.ap.win_length,
fmin=self.ap.mel_fmin,
fmax=self.ap.mel_fmax,
center=False,
)
scores_d_fake = None
feats_d_fake = None
feats_d_real = None
if self.train_disc:
# compute discriminator scores and features
scores_d_fake, feats_d_fake, _, feats_d_real = self.disc(
self.model_outputs_cache["model_outputs"], self.model_outputs_cache["waveform_seg"]
)
# compute losses
with autocast(enabled=True): # use float32 for the criterion
loss_dict = criterion[optimizer_idx](
mel_output=self.model_outputs_cache["acoustic_model_outputs"].transpose(1, 2),
mel_target=batch["mel_input"],
mel_lens=batch["mel_lengths"],
dur_output=self.model_outputs_cache["dr_log_pred"],
dur_target=self.model_outputs_cache["dr_log_target"].detach(),
pitch_output=self.model_outputs_cache["pitch_pred"],
pitch_target=self.model_outputs_cache["pitch_target"],
energy_output=self.model_outputs_cache["energy_pred"],
energy_target=self.model_outputs_cache["energy_target"],
src_lens=batch["text_lengths"],
waveform=self.model_outputs_cache["waveform_seg"],
waveform_hat=self.model_outputs_cache["model_outputs"],
p_prosody_ref=self.model_outputs_cache["p_prosody_ref"],
p_prosody_pred=self.model_outputs_cache["p_prosody_pred"],
u_prosody_ref=self.model_outputs_cache["u_prosody_ref"],
u_prosody_pred=self.model_outputs_cache["u_prosody_pred"],
aligner_logprob=self.model_outputs_cache["aligner_logprob"],
aligner_hard=self.model_outputs_cache["aligner_mas"],
aligner_soft=self.model_outputs_cache["aligner_soft"],
binary_loss_weight=self.binary_loss_weight,
feats_fake=feats_d_fake,
feats_real=feats_d_real,
scores_fake=scores_d_fake,
spec_slice=mel_slice,
spec_slice_hat=mel_slice_hat,
skip_disc=not self.train_disc,
)
loss_dict["avg_text_length"] = batch["text_lengths"].float().mean()
loss_dict["avg_mel_length"] = batch["mel_lengths"].float().mean()
loss_dict["avg_text_batch_occupancy"] = (
batch["text_lengths"].float() / batch["text_lengths"].float().max()
).mean()
loss_dict["avg_mel_batch_occupancy"] = (
batch["mel_lengths"].float() / batch["mel_lengths"].float().max()
).mean()
return self.model_outputs_cache, loss_dict
raise ValueError(" [!] Unexpected `optimizer_idx`.")
def eval_step(self, batch: dict, criterion: nn.Module, optimizer_idx: int):
return self.train_step(batch, criterion, optimizer_idx)
def _log(self, batch, outputs, name_prefix="train"):
figures, audios = {}, {}
# encoder outputs
model_outputs = outputs[1]["acoustic_model_outputs"]
alignments = outputs[1]["alignments"]
mel_input = batch["mel_input"]
pred_spec = model_outputs[0].data.cpu().numpy()
gt_spec = mel_input[0].data.cpu().numpy()
align_img = alignments[0].data.cpu().numpy()
figures = {
"prediction": plot_spectrogram(pred_spec, None, output_fig=False),
"ground_truth": plot_spectrogram(gt_spec.T, None, output_fig=False),
"alignment": plot_alignment(align_img, output_fig=False),
}
# plot pitch figures
pitch_avg = abs(outputs[1]["pitch_target"][0, 0].data.cpu().numpy())
pitch_avg_hat = abs(outputs[1]["pitch_pred"][0, 0].data.cpu().numpy())
chars = self.tokenizer.decode(batch["text_input"][0].data.cpu().numpy())
pitch_figures = {
"pitch_ground_truth": plot_avg_pitch(pitch_avg, chars, output_fig=False),
"pitch_avg_predicted": plot_avg_pitch(pitch_avg_hat, chars, output_fig=False),
}
figures.update(pitch_figures)
# plot energy figures
energy_avg = abs(outputs[1]["energy_target"][0, 0].data.cpu().numpy())
energy_avg_hat = abs(outputs[1]["energy_pred"][0, 0].data.cpu().numpy())
chars = self.tokenizer.decode(batch["text_input"][0].data.cpu().numpy())
energy_figures = {
"energy_ground_truth": plot_avg_pitch(energy_avg, chars, output_fig=False),
"energy_avg_predicted": plot_avg_pitch(energy_avg_hat, chars, output_fig=False),
}
figures.update(energy_figures)
# plot the attention mask computed from the predicted durations
alignments_hat = outputs[1]["alignments_dp"][0].data.cpu().numpy()
figures["alignment_hat"] = plot_alignment(alignments_hat.T, output_fig=False)
# Sample audio
encoder_audio = mel_to_wav_numpy(
mel=db_to_amp_numpy(x=pred_spec.T, gain=1, base=None), mel_basis=self.mel_basis, **self.config.audio
)
audios[f"{name_prefix}/encoder_audio"] = encoder_audio
# vocoder outputs
y_hat = outputs[1]["model_outputs"]
y = outputs[1]["waveform_seg"]
vocoder_figures = plot_results(y_hat=y_hat, y=y, ap=self.ap, name_prefix=name_prefix)
figures.update(vocoder_figures)
sample_voice = y_hat[0].squeeze(0).detach().cpu().numpy()
audios[f"{name_prefix}/vocoder_audio"] = sample_voice
return figures, audios
def train_log(
self, batch: dict, outputs: dict, logger: "Logger", assets: dict, steps: int
): # pylint: disable=no-self-use, unused-argument
"""Create visualizations and waveform examples.
For example, here you can plot spectrograms and generate sample sample waveforms from these spectrograms to
be projected onto Tensorboard.
Args:
batch (Dict): Model inputs used at the previous training step.
outputs (Dict): Model outputs generated at the previous training step.
Returns:
Tuple[Dict, np.ndarray]: training plots and output waveform.
"""
figures, audios = self._log(batch=batch, outputs=outputs, name_prefix="vocoder/")
logger.train_figures(steps, figures)
logger.train_audios(steps, audios, self.ap.sample_rate)
def eval_log(self, batch: dict, outputs: dict, logger: "Logger", assets: dict, steps: int) -> None:
figures, audios = self._log(batch=batch, outputs=outputs, name_prefix="vocoder/")
logger.eval_figures(steps, figures)
logger.eval_audios(steps, audios, self.ap.sample_rate)
def get_aux_input_from_test_sentences(self, sentence_info):
if hasattr(self.config, "model_args"):
config = self.config.model_args
else:
config = self.config
# extract speaker and language info
text, speaker_name, style_wav = None, None, None
if isinstance(sentence_info, list):
if len(sentence_info) == 1:
text = sentence_info[0]
elif len(sentence_info) == 2:
text, speaker_name = sentence_info
elif len(sentence_info) == 3:
text, speaker_name, style_wav = sentence_info
else:
text = sentence_info
# get speaker id/d_vector
speaker_id, d_vector = None, None
if hasattr(self, "speaker_manager"):
if config.use_d_vector_file:
if speaker_name is None:
d_vector = self.speaker_manager.get_random_embedding()
else:
d_vector = self.speaker_manager.get_mean_embedding(speaker_name, num_samples=None, randomize=False)
elif config.use_speaker_embedding:
if speaker_name is None:
speaker_id = self.speaker_manager.get_random_id()
else:
speaker_id = self.speaker_manager.name_to_id[speaker_name]
return {"text": text, "speaker_id": speaker_id, "style_wav": style_wav, "d_vector": d_vector}
def plot_outputs(self, text, wav, alignment, outputs):
figures = {}
pitch_avg_pred = outputs["pitch"].cpu()
energy_avg_pred = outputs["energy"].cpu()
spec = wav_to_mel(
y=torch.from_numpy(wav[None, :]),
n_fft=self.ap.fft_size,
sample_rate=self.ap.sample_rate,
num_mels=self.ap.num_mels,
hop_length=self.ap.hop_length,
win_length=self.ap.win_length,
fmin=self.ap.mel_fmin,
fmax=self.ap.mel_fmax,
center=False,
)[0].transpose(0, 1)
pitch = compute_f0(
x=wav[0],
sample_rate=self.ap.sample_rate,
hop_length=self.ap.hop_length,
pitch_fmax=self.ap.pitch_fmax,
)
input_text = self.tokenizer.ids_to_text(self.tokenizer.text_to_ids(text, language="en"))
input_text = input_text.replace("<BLNK>", "_")
durations = outputs["durations"]
pitch_avg = average_over_durations(torch.from_numpy(pitch)[None, None, :], durations.cpu()) # [1, 1, n_frames]
pitch_avg_pred_denorm = (pitch_avg_pred * self.pitch_std) + self.pitch_mean
figures["alignment"] = plot_alignment(alignment.transpose(1, 2), output_fig=False)
figures["spectrogram"] = plot_spectrogram(spec)
figures["pitch_from_wav"] = plot_pitch(pitch, spec)
figures["pitch_avg_from_wav"] = plot_avg_pitch(pitch_avg.squeeze(), input_text)
figures["pitch_avg_pred"] = plot_avg_pitch(pitch_avg_pred_denorm.squeeze(), input_text)
figures["energy_avg_pred"] = plot_avg_pitch(energy_avg_pred.squeeze(), input_text)
return figures
def synthesize(
self,
text: str,
speaker_id: str = None,
d_vector: torch.tensor = None,
pitch_transform=None,
**kwargs,
): # pylint: disable=unused-argument
# TODO: add cloning support with ref_waveform
is_cuda = next(self.parameters()).is_cuda
# convert text to sequence of token IDs
text_inputs = np.asarray(
self.tokenizer.text_to_ids(text, language=None),
dtype=np.int32,
)
# set speaker inputs
_speaker_id = None
if speaker_id is not None and self.args.use_speaker_embedding:
if isinstance(speaker_id, str) and self.args.use_speaker_embedding:
# get the speaker id for the speaker embedding layer
_speaker_id = self.speaker_manager.name_to_id[speaker_id]
_speaker_id = id_to_torch(_speaker_id, cuda=is_cuda)
if speaker_id is not None and self.args.use_d_vector_file:
# get the average d_vector for the speaker
d_vector = self.speaker_manager.get_mean_embedding(speaker_id, num_samples=None, randomize=False)
d_vector = embedding_to_torch(d_vector, cuda=is_cuda)
text_inputs = numpy_to_torch(text_inputs, torch.long, cuda=is_cuda)
text_inputs = text_inputs.unsqueeze(0)
# synthesize voice
outputs = self.inference(
text_inputs,
aux_input={"d_vectors": d_vector, "speaker_ids": _speaker_id},
pitch_transform=pitch_transform,
# energy_transform=energy_transform
)
# collect outputs
wav = outputs["model_outputs"][0].data.cpu().numpy()
alignments = outputs["alignments"]
return_dict = {
"wav": wav,
"alignments": alignments,
"text_inputs": text_inputs,
"outputs": outputs,
}
return return_dict
def synthesize_with_gl(self, text: str, speaker_id, d_vector):
is_cuda = next(self.parameters()).is_cuda
# convert text to sequence of token IDs
text_inputs = np.asarray(
self.tokenizer.text_to_ids(text, language=None),
dtype=np.int32,
)
# pass tensors to backend
if speaker_id is not None:
speaker_id = id_to_torch(speaker_id, cuda=is_cuda)
if d_vector is not None:
d_vector = embedding_to_torch(d_vector, cuda=is_cuda)
text_inputs = numpy_to_torch(text_inputs, torch.long, cuda=is_cuda)
text_inputs = text_inputs.unsqueeze(0)
# synthesize voice
outputs = self.inference_spec_decoder(
x=text_inputs,
aux_input={"d_vectors": d_vector, "speaker_ids": speaker_id},
)
# collect outputs
S = outputs["model_outputs"].cpu().numpy()[0].T
S = db_to_amp_numpy(x=S, gain=1, base=None)
wav = mel_to_wav_numpy(mel=S, mel_basis=self.mel_basis, **self.config.audio)
alignments = outputs["alignments"]
return_dict = {
"wav": wav[None, :],
"alignments": alignments,
"text_inputs": text_inputs,
"outputs": outputs,
}
return return_dict
@torch.no_grad()
def test_run(self, assets) -> Tuple[Dict, Dict]:
"""Generic test run for `tts` models used by `Trainer`.
You can override this for a different behaviour.
Returns:
Tuple[Dict, Dict]: Test figures and audios to be projected to Tensorboard.
"""
print(" | > Synthesizing test sentences.")
test_audios = {}
test_figures = {}
test_sentences = self.config.test_sentences
for idx, s_info in enumerate(test_sentences):
aux_inputs = self.get_aux_input_from_test_sentences(s_info)
outputs = self.synthesize(
aux_inputs["text"],
config=self.config,
speaker_id=aux_inputs["speaker_id"],
d_vector=aux_inputs["d_vector"],
)
outputs_gl = self.synthesize_with_gl(
aux_inputs["text"],
speaker_id=aux_inputs["speaker_id"],
d_vector=aux_inputs["d_vector"],
)
# speaker_name = self.speaker_manager.speaker_names[aux_inputs["speaker_id"]]
test_audios["{}-audio".format(idx)] = outputs["wav"].T
test_audios["{}-audio_encoder".format(idx)] = outputs_gl["wav"].T
test_figures["{}-alignment".format(idx)] = plot_alignment(outputs["alignments"], output_fig=False)
return {"figures": test_figures, "audios": test_audios}
def test_log(
self, outputs: dict, logger: "Logger", assets: dict, steps: int # pylint: disable=unused-argument
) -> None:
logger.test_audios(steps, outputs["audios"], self.config.audio.sample_rate)
logger.test_figures(steps, outputs["figures"])
def format_batch(self, batch: Dict) -> Dict:
"""Compute speaker, langugage IDs and d_vector for the batch if necessary."""
speaker_ids = None
d_vectors = None
# get numerical speaker ids from speaker names
if self.speaker_manager is not None and self.speaker_manager.speaker_names and self.args.use_speaker_embedding:
speaker_ids = [self.speaker_manager.name_to_id[sn] for sn in batch["speaker_names"]]
if speaker_ids is not None:
speaker_ids = torch.LongTensor(speaker_ids)
batch["speaker_ids"] = speaker_ids
# get d_vectors from audio file names
if self.speaker_manager is not None and self.speaker_manager.embeddings and self.args.use_d_vector_file:
d_vector_mapping = self.speaker_manager.embeddings
d_vectors = [d_vector_mapping[w]["embedding"] for w in batch["audio_unique_names"]]
d_vectors = torch.FloatTensor(d_vectors)
batch["d_vectors"] = d_vectors
batch["speaker_ids"] = speaker_ids
return batch
def format_batch_on_device(self, batch):
"""Compute spectrograms on the device."""
ac = self.ap
# compute spectrograms
batch["mel_input"] = wav_to_mel(
batch["waveform"],
hop_length=ac.hop_length,
win_length=ac.win_length,
n_fft=ac.fft_size,
num_mels=ac.num_mels,
sample_rate=ac.sample_rate,
fmin=ac.mel_fmin,
fmax=ac.mel_fmax,
center=False,
)
# TODO: Align pitch properly
# assert (
# batch["pitch"].shape[2] == batch["mel_input"].shape[2]
# ), f"{batch['pitch'].shape[2]}, {batch['mel_input'].shape[2]}"
batch["pitch"] = batch["pitch"][:, :, : batch["mel_input"].shape[2]] if batch["pitch"] is not None else None
batch["mel_lengths"] = (batch["mel_input"].shape[2] * batch["waveform_rel_lens"]).int()
# zero the padding frames
batch["mel_input"] = batch["mel_input"] * sequence_mask(batch["mel_lengths"]).unsqueeze(1)
# format attn priors as we now the max mel length
# TODO: fix 1 diff b/w mel_lengths and attn_priors
if self.config.use_attn_priors:
attn_priors_np = batch["attn_priors"]
batch["attn_priors"] = torch.zeros(
batch["mel_input"].shape[0],
batch["mel_lengths"].max(),
batch["text_lengths"].max(),
device=batch["mel_input"].device,
)
for i in range(batch["mel_input"].shape[0]):
batch["attn_priors"][i, : attn_priors_np[i].shape[0], : attn_priors_np[i].shape[1]] = torch.from_numpy(
attn_priors_np[i]
)
batch["energy"] = None
batch["energy"] = wav_to_energy( # [B, 1, T_max2]
batch["waveform"],
hop_length=ac.hop_length,
win_length=ac.win_length,
n_fft=ac.fft_size,
center=False,
)
batch["energy"] = self.energy_scaler(batch["energy"])
return batch
def get_sampler(self, config: Coqpit, dataset: TTSDataset, num_gpus=1):
weights = None
data_items = dataset.samples
if getattr(config, "use_weighted_sampler", False):
for attr_name, alpha in config.weighted_sampler_attrs.items():
print(f" > Using weighted sampler for attribute '{attr_name}' with alpha '{alpha}'")
multi_dict = config.weighted_sampler_multipliers.get(attr_name, None)
print(multi_dict)
weights, attr_names, attr_weights = get_attribute_balancer_weights(
attr_name=attr_name, items=data_items, multi_dict=multi_dict
)
weights = weights * alpha
print(f" > Attribute weights for '{attr_names}' \n | > {attr_weights}")
if weights is not None:
sampler = WeightedRandomSampler(weights, len(weights))
else:
sampler = None
# sampler for DDP
if sampler is None:
sampler = DistributedSampler(dataset) if num_gpus > 1 else None
else: # If a sampler is already defined use this sampler and DDP sampler together
sampler = DistributedSamplerWrapper(sampler) if num_gpus > 1 else sampler
return sampler
def get_data_loader(
self,
config: Coqpit,
assets: Dict,
is_eval: bool,
samples: Union[List[Dict], List[List]],
verbose: bool,
num_gpus: int,
rank: int = None,
) -> "DataLoader":
if is_eval and not config.run_eval:
loader = None
else:
# init dataloader
dataset = ForwardTTSE2eDataset(
samples=samples,
ap=self.ap,
batch_group_size=0 if is_eval else config.batch_group_size * config.batch_size,
min_text_len=config.min_text_len,
max_text_len=config.max_text_len,
min_audio_len=config.min_audio_len,
max_audio_len=config.max_audio_len,
phoneme_cache_path=config.phoneme_cache_path,
precompute_num_workers=config.precompute_num_workers,
compute_f0=config.compute_f0,
f0_cache_path=config.f0_cache_path,
attn_prior_cache_path=config.attn_prior_cache_path if config.use_attn_priors else None,
verbose=verbose,
tokenizer=self.tokenizer,
start_by_longest=config.start_by_longest,
)
# wait all the DDP process to be ready
if num_gpus > 1:
dist.barrier()
# sort input sequences ascendingly by length
dataset.preprocess_samples()
# get samplers
sampler = self.get_sampler(config, dataset, num_gpus)
loader = DataLoader(
dataset,
batch_size=config.eval_batch_size if is_eval else config.batch_size,
shuffle=False, # shuffle is done in the dataset.
drop_last=False, # setting this False might cause issues in AMP training.
sampler=sampler,
collate_fn=dataset.collate_fn,
num_workers=config.num_eval_loader_workers if is_eval else config.num_loader_workers,
pin_memory=True,
)
# get pitch mean and std
self.pitch_mean = dataset.f0_dataset.mean
self.pitch_std = dataset.f0_dataset.std
return loader
def get_criterion(self):
return [VitsDiscriminatorLoss(self.config), DelightfulTTSLoss(self.config)]
def get_optimizer(self) -> List:
"""Initiate and return the GAN optimizers based on the config parameters.
It returnes 2 optimizers in a list. First one is for the generator and the second one is for the discriminator.
Returns:
List: optimizers.
"""
optimizer_disc = get_optimizer(
self.config.optimizer, self.config.optimizer_params, self.config.lr_disc, self.disc
)
gen_parameters = chain(params for k, params in self.named_parameters() if not k.startswith("disc."))
optimizer_gen = get_optimizer(
self.config.optimizer, self.config.optimizer_params, self.config.lr_gen, parameters=gen_parameters
)
return [optimizer_disc, optimizer_gen]
def get_lr(self) -> List:
"""Set the initial learning rates for each optimizer.
Returns:
List: learning rates for each optimizer.
"""
return [self.config.lr_disc, self.config.lr_gen]
def get_scheduler(self, optimizer) -> List:
"""Set the schedulers for each optimizer.
Args:
optimizer (List[`torch.optim.Optimizer`]): List of optimizers.
Returns:
List: Schedulers, one for each optimizer.
"""
scheduler_D = get_scheduler(self.config.lr_scheduler_gen, self.config.lr_scheduler_gen_params, optimizer[0])
scheduler_G = get_scheduler(self.config.lr_scheduler_disc, self.config.lr_scheduler_disc_params, optimizer[1])
return [scheduler_D, scheduler_G]
def on_epoch_end(self, trainer): # pylint: disable=unused-argument
# stop updating mean and var
# TODO: do the same for F0
self.energy_scaler.eval()
@staticmethod
def init_from_config(
config: "DelightfulTTSConfig", samples: Union[List[List], List[Dict]] = None, verbose=False
): # pylint: disable=unused-argument
"""Initiate model from config
Args:
config (ForwardTTSE2eConfig): Model config.
samples (Union[List[List], List[Dict]]): Training samples to parse speaker ids for training.
Defaults to None.
"""
tokenizer, new_config = TTSTokenizer.init_from_config(config)
speaker_manager = SpeakerManager.init_from_config(config.model_args, samples)
ap = AudioProcessor.init_from_config(config=config)
return DelightfulTTS(config=new_config, tokenizer=tokenizer, speaker_manager=speaker_manager, ap=ap)
def load_checkpoint(self, config, checkpoint_path, eval=False):
"""Load model from a checkpoint created by the 👟"""
# pylint: disable=unused-argument, redefined-builtin
state = load_fsspec(checkpoint_path, map_location=torch.device("cpu"))
self.load_state_dict(state["model"])
if eval:
self.eval()
assert not self.training
def get_state_dict(self):
"""Custom state dict of the model with all the necessary components for inference."""
save_state = {"config": self.config.to_dict(), "args": self.args.to_dict(), "model": self.state_dict}
if hasattr(self, "emb_g"):
save_state["speaker_ids"] = self.speaker_manager.speaker_names
if self.args.use_d_vector_file:
# TODO: implement saving of d_vectors
...
return save_state
def save(self, config, checkpoint_path):
"""Save model to a file."""
save_state = self.get_state_dict(config, checkpoint_path) # pylint: disable=too-many-function-args
save_state["pitch_mean"] = self.pitch_mean
save_state["pitch_std"] = self.pitch_std
torch.save(save_state, checkpoint_path)
def on_train_step_start(self, trainer) -> None:
"""Enable the discriminator training based on `steps_to_start_discriminator`
Args:
trainer (Trainer): Trainer object.
"""
self.binary_loss_weight = min(trainer.epochs_done / self.config.binary_loss_warmup_epochs, 1.0) * 1.0
self.train_disc = ( # pylint: disable=attribute-defined-outside-init
trainer.total_steps_done >= self.config.steps_to_start_discriminator
)
class DelightfulTTSLoss(nn.Module):
def __init__(self, config):
super().__init__()
self.mse_loss = nn.MSELoss()
self.mae_loss = nn.L1Loss()
self.forward_sum_loss = ForwardSumLoss()
self.multi_scale_stft_loss = MultiScaleSTFTLoss(**config.multi_scale_stft_loss_params)
self.mel_loss_alpha = config.mel_loss_alpha
self.aligner_loss_alpha = config.aligner_loss_alpha
self.pitch_loss_alpha = config.pitch_loss_alpha
self.energy_loss_alpha = config.energy_loss_alpha
self.u_prosody_loss_alpha = config.u_prosody_loss_alpha
self.p_prosody_loss_alpha = config.p_prosody_loss_alpha
self.dur_loss_alpha = config.dur_loss_alpha
self.char_dur_loss_alpha = config.char_dur_loss_alpha
self.binary_alignment_loss_alpha = config.binary_align_loss_alpha
self.vocoder_mel_loss_alpha = config.vocoder_mel_loss_alpha
self.feat_loss_alpha = config.feat_loss_alpha
self.gen_loss_alpha = config.gen_loss_alpha
self.multi_scale_stft_loss_alpha = config.multi_scale_stft_loss_alpha
@staticmethod
def _binary_alignment_loss(alignment_hard, alignment_soft):
"""Binary loss that forces soft alignments to match the hard alignments as
explained in `https://arxiv.org/pdf/2108.10447.pdf`.
"""
log_sum = torch.log(torch.clamp(alignment_soft[alignment_hard == 1], min=1e-12)).sum()
return -log_sum / alignment_hard.sum()
@staticmethod
def feature_loss(feats_real, feats_generated):
loss = 0
for dr, dg in zip(feats_real, feats_generated):
for rl, gl in zip(dr, dg):
rl = rl.float().detach()
gl = gl.float()
loss += torch.mean(torch.abs(rl - gl))
return loss * 2
@staticmethod
def generator_loss(scores_fake):
loss = 0
gen_losses = []
for dg in scores_fake:
dg = dg.float()
l = torch.mean((1 - dg) ** 2)
gen_losses.append(l)
loss += l
return loss, gen_losses
def forward(
self,
mel_output,
mel_target,
mel_lens,
dur_output,
dur_target,
pitch_output,
pitch_target,
energy_output,
energy_target,
src_lens,
waveform,
waveform_hat,
p_prosody_ref,
p_prosody_pred,
u_prosody_ref,
u_prosody_pred,
aligner_logprob,
aligner_hard,
aligner_soft,
binary_loss_weight=None,
feats_fake=None,
feats_real=None,
scores_fake=None,
spec_slice=None,
spec_slice_hat=None,
skip_disc=False,
):
"""
Shapes:
- mel_output: :math:`(B, C_mel, T_mel)`
- mel_target: :math:`(B, C_mel, T_mel)`
- mel_lens: :math:`(B)`
- dur_output: :math:`(B, T_src)`
- dur_target: :math:`(B, T_src)`
- pitch_output: :math:`(B, 1, T_src)`
- pitch_target: :math:`(B, 1, T_src)`
- energy_output: :math:`(B, 1, T_src)`
- energy_target: :math:`(B, 1, T_src)`
- src_lens: :math:`(B)`
- waveform: :math:`(B, 1, T_wav)`
- waveform_hat: :math:`(B, 1, T_wav)`
- p_prosody_ref: :math:`(B, T_src, 4)`
- p_prosody_pred: :math:`(B, T_src, 4)`
- u_prosody_ref: :math:`(B, 1, 256)
- u_prosody_pred: :math:`(B, 1, 256)
- aligner_logprob: :math:`(B, 1, T_mel, T_src)`
- aligner_hard: :math:`(B, T_mel, T_src)`
- aligner_soft: :math:`(B, T_mel, T_src)`
- spec_slice: :math:`(B, C_mel, T_mel)`
- spec_slice_hat: :math:`(B, C_mel, T_mel)`
"""
loss_dict = {}
src_mask = sequence_mask(src_lens).to(mel_output.device) # (B, T_src)
mel_mask = sequence_mask(mel_lens).to(mel_output.device) # (B, T_mel)
dur_target.requires_grad = False
mel_target.requires_grad = False
pitch_target.requires_grad = False
masked_mel_predictions = mel_output.masked_select(mel_mask[:, None])
mel_targets = mel_target.masked_select(mel_mask[:, None])
mel_loss = self.mae_loss(masked_mel_predictions, mel_targets)
p_prosody_ref = p_prosody_ref.detach()
p_prosody_loss = 0.5 * self.mae_loss(
p_prosody_ref.masked_select(src_mask.unsqueeze(-1)),
p_prosody_pred.masked_select(src_mask.unsqueeze(-1)),
)
u_prosody_ref = u_prosody_ref.detach()
u_prosody_loss = 0.5 * self.mae_loss(u_prosody_ref, u_prosody_pred)
duration_loss = self.mse_loss(dur_output, dur_target)
pitch_output = pitch_output.masked_select(src_mask[:, None])
pitch_target = pitch_target.masked_select(src_mask[:, None])
pitch_loss = self.mse_loss(pitch_output, pitch_target)
energy_output = energy_output.masked_select(src_mask[:, None])
energy_target = energy_target.masked_select(src_mask[:, None])
energy_loss = self.mse_loss(energy_output, energy_target)
forward_sum_loss = self.forward_sum_loss(aligner_logprob, src_lens, mel_lens)
total_loss = (
(mel_loss * self.mel_loss_alpha)
+ (duration_loss * self.dur_loss_alpha)
+ (u_prosody_loss * self.u_prosody_loss_alpha)
+ (p_prosody_loss * self.p_prosody_loss_alpha)
+ (pitch_loss * self.pitch_loss_alpha)
+ (energy_loss * self.energy_loss_alpha)
+ (forward_sum_loss * self.aligner_loss_alpha)
)
if self.binary_alignment_loss_alpha > 0 and aligner_hard is not None:
binary_alignment_loss = self._binary_alignment_loss(aligner_hard, aligner_soft)
total_loss = total_loss + self.binary_alignment_loss_alpha * binary_alignment_loss * binary_loss_weight
if binary_loss_weight:
loss_dict["loss_binary_alignment"] = (
self.binary_alignment_loss_alpha * binary_alignment_loss * binary_loss_weight
)
else:
loss_dict["loss_binary_alignment"] = self.binary_alignment_loss_alpha * binary_alignment_loss
loss_dict["loss_aligner"] = self.aligner_loss_alpha * forward_sum_loss
loss_dict["loss_mel"] = self.mel_loss_alpha * mel_loss
loss_dict["loss_duration"] = self.dur_loss_alpha * duration_loss
loss_dict["loss_u_prosody"] = self.u_prosody_loss_alpha * u_prosody_loss
loss_dict["loss_p_prosody"] = self.p_prosody_loss_alpha * p_prosody_loss
loss_dict["loss_pitch"] = self.pitch_loss_alpha * pitch_loss
loss_dict["loss_energy"] = self.energy_loss_alpha * energy_loss
loss_dict["loss"] = total_loss
# vocoder losses
if not skip_disc:
loss_feat = self.feature_loss(feats_real=feats_real, feats_generated=feats_fake) * self.feat_loss_alpha
loss_gen = self.generator_loss(scores_fake=scores_fake)[0] * self.gen_loss_alpha
loss_dict["vocoder_loss_feat"] = loss_feat
loss_dict["vocoder_loss_gen"] = loss_gen
loss_dict["loss"] = loss_dict["loss"] + loss_feat + loss_gen
loss_mel = torch.nn.functional.l1_loss(spec_slice, spec_slice_hat) * self.vocoder_mel_loss_alpha
loss_stft_mg, loss_stft_sc = self.multi_scale_stft_loss(y_hat=waveform_hat, y=waveform)
loss_stft_mg = loss_stft_mg * self.multi_scale_stft_loss_alpha
loss_stft_sc = loss_stft_sc * self.multi_scale_stft_loss_alpha
loss_dict["vocoder_loss_mel"] = loss_mel
loss_dict["vocoder_loss_stft_mg"] = loss_stft_mg
loss_dict["vocoder_loss_stft_sc"] = loss_stft_sc
loss_dict["loss"] = loss_dict["loss"] + loss_mel + loss_stft_sc + loss_stft_mg
return loss_dict