# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import contextlib import omegaconf import torch from hydra.utils import instantiate from lightning.pytorch import Trainer from lightning.pytorch.loggers import WandbLogger from omegaconf import DictConfig, OmegaConf from torch.cuda.amp import autocast from torch.nn import functional as F from nemo.collections.tts.data.dataset import DistributedBucketSampler from nemo.collections.tts.losses.vits_losses import DiscriminatorLoss, FeatureMatchingLoss, GeneratorLoss, KlLoss from nemo.collections.tts.models.base import TextToWaveform from nemo.collections.tts.modules.vits_modules import MultiPeriodDiscriminator from nemo.collections.tts.parts.utils.helpers import ( clip_grad_value_, g2p_backward_compatible_support, plot_spectrogram_to_numpy, slice_segments, ) from nemo.collections.tts.torch.tts_data_types import SpeakerID from nemo.core.classes.common import PretrainedModelInfo, typecheck from nemo.core.neural_types.elements import AudioSignal, FloatType, Index, IntType, TokenIndex from nemo.core.neural_types.neural_type import NeuralType from nemo.core.optim.lr_scheduler import CosineAnnealing from nemo.utils import logging, model_utils from nemo.utils.decorators.experimental import experimental HAVE_WANDB = True try: import wandb except ModuleNotFoundError: HAVE_WANDB = False @experimental class VitsModel(TextToWaveform): def __init__(self, cfg: DictConfig, trainer: 'Trainer' = None): # Convert to Hydra 1.0 compatible DictConfig cfg = model_utils.convert_model_config_to_dict_config(cfg) cfg = model_utils.maybe_update_config_version(cfg) # setup normalizer self.normalizer = None self.text_normalizer_call = None self.text_normalizer_call_kwargs = {} self._setup_normalizer(cfg) # setup tokenizer self.tokenizer = None self._setup_tokenizer(cfg) assert self.tokenizer is not None num_tokens = len(self.tokenizer.tokens) self.tokenizer_pad = self.tokenizer.pad super().__init__(cfg=cfg, trainer=trainer) self.audio_to_melspec_processor = instantiate(cfg.preprocessor, highfreq=cfg.train_ds.dataset.highfreq) self.feat_matching_loss = FeatureMatchingLoss() self.disc_loss = DiscriminatorLoss() self.gen_loss = GeneratorLoss() self.kl_loss = KlLoss() self.net_g = instantiate( cfg.synthesizer, n_vocab=num_tokens, spec_channels=cfg.n_fft // 2 + 1, segment_size=cfg.segment_size // cfg.n_window_stride, padding_idx=self.tokenizer_pad, ) self.net_d = MultiPeriodDiscriminator(cfg.use_spectral_norm) self.automatic_optimization = False def _setup_tokenizer(self, cfg): text_tokenizer_kwargs = {} if "g2p" in cfg.text_tokenizer and cfg.text_tokenizer.g2p is not None: # for backward compatibility if ( self._is_model_being_restored() and (cfg.text_tokenizer.g2p.get('_target_', None) is not None) and cfg.text_tokenizer.g2p["_target_"].startswith("nemo_text_processing.g2p") ): cfg.text_tokenizer.g2p["_target_"] = g2p_backward_compatible_support( cfg.text_tokenizer.g2p["_target_"] ) g2p_kwargs = {} if "phoneme_dict" in cfg.text_tokenizer.g2p: g2p_kwargs["phoneme_dict"] = self.register_artifact( 'text_tokenizer.g2p.phoneme_dict', cfg.text_tokenizer.g2p.phoneme_dict, ) if "heteronyms" in cfg.text_tokenizer.g2p: g2p_kwargs["heteronyms"] = self.register_artifact( 'text_tokenizer.g2p.heteronyms', cfg.text_tokenizer.g2p.heteronyms, ) text_tokenizer_kwargs["g2p"] = instantiate(cfg.text_tokenizer.g2p, **g2p_kwargs) self.tokenizer = instantiate(cfg.text_tokenizer, **text_tokenizer_kwargs) def parse(self, text: str, normalize=True) -> torch.tensor: if self.training: logging.warning("parse() is meant to be called in eval mode.") if normalize and self.text_normalizer_call is not None: text = self.text_normalizer_call(text, **self.text_normalizer_call_kwargs) eval_phon_mode = contextlib.nullcontext() if hasattr(self.tokenizer, "set_phone_prob"): eval_phon_mode = self.tokenizer.set_phone_prob(prob=1.0) with eval_phon_mode: tokens = self.tokenizer.encode(text) return torch.tensor(tokens).long().unsqueeze(0).to(self.device) def configure_optimizers(self): optim_config = self._cfg.optim.copy() OmegaConf.set_struct(optim_config, False) sched_config = optim_config.pop("sched", None) OmegaConf.set_struct(optim_config, True) optim_g = instantiate( optim_config, params=self.net_g.parameters(), ) optim_d = instantiate( optim_config, params=self.net_d.parameters(), ) if sched_config is not None: if sched_config.name == 'ExponentialLR': scheduler_d = torch.optim.lr_scheduler.ExponentialLR(optim_d, gamma=sched_config.lr_decay) scheduler_g = torch.optim.lr_scheduler.ExponentialLR(optim_g, gamma=sched_config.lr_decay) elif sched_config.name == 'CosineAnnealing': scheduler_g = CosineAnnealing( optimizer=optim_g, max_steps=sched_config.max_steps, min_lr=sched_config.min_lr, ) scheduler_d = CosineAnnealing( optimizer=optim_d, max_steps=sched_config.max_steps, min_lr=sched_config.min_lr, ) else: raise ValueError("Unknown optimizer.") scheduler_g_dict = {'scheduler': scheduler_g, 'interval': 'step'} scheduler_d_dict = {'scheduler': scheduler_d, 'interval': 'step'} return [optim_g, optim_d], [scheduler_g_dict, scheduler_d_dict] else: return [optim_g, optim_d] # for inference @typecheck( input_types={ "tokens": NeuralType(('B', 'T_text'), TokenIndex()), "speakers": NeuralType(('B',), Index(), optional=True), "noise_scale": NeuralType(('B',), FloatType(), optional=True), "length_scale": NeuralType(('B',), FloatType(), optional=True), "noise_scale_w": NeuralType(('B',), FloatType(), optional=True), "max_len": NeuralType(('B',), IntType(), optional=True), } ) def forward(self, tokens, speakers=None, noise_scale=1, length_scale=1, noise_scale_w=1.0, max_len=1000): text_len = torch.tensor([tokens.size(-1)]).to(int).to(tokens.device) audio_pred, attn, y_mask, (z, z_p, m_p, logs_p) = self.net_g.infer( tokens, text_len, speakers=speakers, noise_scale=noise_scale, length_scale=length_scale, noise_scale_w=noise_scale_w, max_len=max_len, ) return audio_pred, attn, y_mask, (z, z_p, m_p, logs_p) def training_step(self, batch, batch_idx): speakers = None if SpeakerID in self._train_dl.dataset.sup_data_types_set: (audio, audio_len, text, text_len, speakers) = batch else: (audio, audio_len, text, text_len) = batch spec, spec_lengths = self.audio_to_melspec_processor(audio, audio_len, linear_spec=True) with autocast(enabled=True): audio_pred, l_length, attn, ids_slice, text_mask, z_mask, (z, z_p, m_p, logs_p, m_q, logs_q) = self.net_g( text, text_len, spec, spec_lengths, speakers ) audio_pred = audio_pred.float() audio_pred_mel, _ = self.audio_to_melspec_processor(audio_pred.squeeze(1), audio_len, linear_spec=False) audio = slice_segments(audio.unsqueeze(1), ids_slice * self.cfg.n_window_stride, self._cfg.segment_size) audio_mel, _ = self.audio_to_melspec_processor(audio.squeeze(1), audio_len, linear_spec=False) with autocast(enabled=True): y_d_hat_r, y_d_hat_g, _, _ = self.net_d(audio, audio_pred.detach()) with autocast(enabled=False): loss_disc, losses_disc_r, losses_disc_g = self.disc_loss( disc_real_outputs=y_d_hat_r, disc_generated_outputs=y_d_hat_g ) loss_disc_all = loss_disc # get optimizers optim_g, optim_d = self.optimizers() # train discriminator optim_d.zero_grad() self.manual_backward(loss_disc_all) norm_d = clip_grad_value_(self.net_d.parameters(), None) optim_d.step() with autocast(enabled=True): y_d_hat_r, y_d_hat_g, fmap_r, fmap_g = self.net_d(audio, audio_pred) # Generator with autocast(enabled=False): loss_dur = torch.sum(l_length.float()) loss_mel = F.l1_loss(audio_mel, audio_pred_mel) * self._cfg.c_mel loss_kl = self.kl_loss(z_p=z_p, logs_q=logs_q, m_p=m_p, logs_p=logs_p, z_mask=z_mask) * self._cfg.c_kl loss_fm = self.feat_matching_loss(fmap_r=fmap_r, fmap_g=fmap_g) loss_gen, losses_gen = self.gen_loss(disc_outputs=y_d_hat_g) loss_gen_all = loss_gen + loss_fm + loss_mel + loss_dur + loss_kl # train generator optim_g.zero_grad() self.manual_backward(loss_gen_all) norm_g = clip_grad_value_(self.net_g.parameters(), None) optim_g.step() schedulers = self.lr_schedulers() if schedulers is not None: sch1, sch2 = schedulers if ( self.trainer.is_last_batch and isinstance(sch1, torch.optim.lr_scheduler.ExponentialLR) or isinstance(sch1, CosineAnnealing) ): sch1.step() sch2.step() metrics = { "loss_gen": loss_gen, "loss_fm": loss_fm, "loss_mel": loss_mel, "loss_dur": loss_dur, "loss_kl": loss_kl, "loss_gen_all": loss_gen_all, "loss_disc_all": loss_disc_all, "grad_gen": norm_g, "grad_disc": norm_d, } for i, v in enumerate(losses_gen): metrics[f"loss_gen_i_{i}"] = v for i, v in enumerate(losses_disc_r): metrics[f"loss_disc_r_{i}"] = v for i, v in enumerate(losses_disc_g): metrics[f"loss_disc_g_{i}"] = v self.log_dict(metrics, on_step=True, sync_dist=True) def validation_step(self, batch, batch_idx): speakers = None if self.cfg.n_speakers > 1: (audio, audio_len, text, text_len, speakers) = batch else: (audio, audio_len, text, text_len) = batch audio_pred, _, mask, *_ = self.net_g.infer(text, text_len, speakers, max_len=1000) audio_pred = audio_pred.squeeze() audio_pred_len = mask.sum([1, 2]).long() * self._cfg.validation_ds.dataset.hop_length mel, mel_lengths = self.audio_to_melspec_processor(audio, audio_len) audio_pred_mel, audio_pred_mel_len = self.audio_to_melspec_processor(audio_pred, audio_pred_len) # plot audio once per epoch if batch_idx == 0 and isinstance(self.logger, WandbLogger) and HAVE_WANDB: logger = self.logger.experiment specs = [] audios = [] specs += [ wandb.Image( plot_spectrogram_to_numpy(mel[0, :, : mel_lengths[0]].data.cpu().numpy()), caption=f"val_mel_target", ), wandb.Image( plot_spectrogram_to_numpy(audio_pred_mel[0, :, : audio_pred_mel_len[0]].data.cpu().numpy()), caption=f"val_mel_predicted", ), ] audios += [ wandb.Audio( audio[0, : audio_len[0]].data.cpu().to(torch.float).numpy(), caption=f"val_wav_target", sample_rate=self._cfg.sample_rate, ), wandb.Audio( audio_pred[0, : audio_pred_len[0]].data.cpu().to(torch.float).numpy(), caption=f"val_wav_predicted", sample_rate=self._cfg.sample_rate, ), ] logger.log({"specs": specs, "audios": audios}) def _loader(self, cfg): try: _ = cfg['dataset']['manifest_filepath'] except omegaconf.errors.MissingMandatoryValue: logging.warning("manifest_filepath was skipped. No dataset for this model.") return None dataset = instantiate( cfg.dataset, text_normalizer=self.normalizer, text_normalizer_call_kwargs=self.text_normalizer_call_kwargs, text_tokenizer=self.tokenizer, ) return torch.utils.data.DataLoader( # noqa dataset=dataset, collate_fn=dataset.collate_fn, **cfg.dataloader_params, ) def train_dataloader(self): # default used by the Trainer dataset = instantiate( self.cfg.train_ds.dataset, text_normalizer=self.normalizer, text_normalizer_call_kwargs=self.text_normalizer_call_kwargs, text_tokenizer=self.tokenizer, ) train_sampler = DistributedBucketSampler(dataset, **self.cfg.train_ds.batch_sampler) dataloader = torch.utils.data.DataLoader( dataset, collate_fn=dataset.collate_fn, batch_sampler=train_sampler, **self.cfg.train_ds.dataloader_params, ) return dataloader def setup_training_data(self, cfg): self._train_dl = self._loader(cfg) def setup_validation_data(self, cfg): self._validation_dl = self._loader(cfg) def setup_test_data(self, cfg): """Omitted.""" pass @classmethod def list_available_models(cls) -> 'List[PretrainedModelInfo]': list_of_models = [] model = PretrainedModelInfo( pretrained_model_name="tts_en_lj_vits", location="https://api.ngc.nvidia.com/v2/models/nvidia/nemo/tts_en_lj_vits/versions/1.13.0/files/vits_ljspeech_fp16_full.nemo", description="This model is trained on LJSpeech audio sampled at 22050Hz. This model has been tested on generating female English " "voices with an American accent.", class_=cls, ) list_of_models.append(model) model = PretrainedModelInfo( pretrained_model_name="tts_en_hifitts_vits", location="https://api.ngc.nvidia.com/v2/models/nvidia/nemo/tts_en_hifitts_vits/versions/r1.15.0/files/vits_en_hifitts.nemo", description="This model is trained on HiFITTS sampled at 44100Hz with and can be used to generate male and female English voices with an American accent.", class_=cls, ) list_of_models.append(model) return list_of_models @typecheck( input_types={ "tokens": NeuralType(('B', 'T_text'), TokenIndex(), optional=True), }, output_types={"audio": NeuralType(('B', 'T_audio'), AudioSignal())}, ) def convert_text_to_waveform(self, *, tokens, speakers=None): audio = self(tokens=tokens, speakers=speakers)[0].squeeze(1) return audio