def output_types(self):
return {
"audio": NeuralType(('B', 'S', 'T'), AudioSignal(self.sample_rate)),
}
class HifiGanModel(Vocoder):
def __init__(self, cfg: DictConfig, trainer: 'Trainer' = None):
if isinstance(cfg, dict):
cfg = OmegaConf.create(cfg)
super().__init__(cfg=cfg, trainer=trainer)
self.audio_to_melspec_precessor = instantiate(cfg.preprocessor)
# use a different melspec extractor because:
# 1. we need to pass grads
# 2. we need remove fmax limitation
self.trg_melspec_fn = instantiate(cfg.preprocessor, highfreq=None, use_grads=True)
self.generator = instantiate(cfg.generator)
self.mpd = MultiPeriodDiscriminator()
self.msd = MultiScaleDiscriminator()
self.feature_loss = FeatureMatchingLoss()
self.discriminator_loss = DiscriminatorLoss()
self.generator_loss = GeneratorLoss()
self.sample_rate = self._cfg.preprocessor.sample_rate
self.stft_bias = None
if isinstance(self._train_dl.dataset, MelAudioDataset):
self.finetune = True
logging.info("fine-tuning on pre-computed mels")
else:
self.finetune = False
logging.info("training on ground-truth mels")
def configure_optimizers(self):
self.optim_g = instantiate(self._cfg.optim, params=self.generator.parameters(),)
self.optim_d = instantiate(
self._cfg.optim, params=itertools.chain(self.msd.parameters(), self.mpd.parameters()),
)
max_steps = self._cfg.max_steps
warmup_steps = 0 if self.finetune else np.ceil(0.2 * max_steps)
self.scheduler_g = CosineAnnealing(
self.optim_g, max_steps=max_steps, min_lr=1e-5, warmup_steps=warmup_steps
) # Use warmup to delay start
sch1_dict = {
'scheduler': self.scheduler_g,
'interval': 'step',
}
self.scheduler_d = CosineAnnealing(self.optim_d, max_steps=max_steps, min_lr=1e-5)
sch2_dict = {
'scheduler': self.scheduler_d,
'interval': 'step',
}
return [self.optim_g, self.optim_d], [sch1_dict, sch2_dict]
@property
def input_types(self):
return {
"spec": NeuralType(('B', 'D', 'T'), MelSpectrogramType()),
}
@property
def output_types(self):
return {
"audio": NeuralType(('B', 'S', 'T'), AudioSignal(self.sample_rate)),
}
@typecheck()
def forward(self, *, spec):
"""
Runs the generator, for inputs and outputs see input_types, and output_types
"""
return self.generator(x=spec)
@typecheck(output_types={"audio": NeuralType(('B', 'T'), AudioSignal())})
def convert_spectrogram_to_audio(self, spec: 'torch.tensor') -> 'torch.tensor':
return self(spec=spec).squeeze(1)
def training_step(self, batch, batch_idx, optimizer_idx):
# if in finetune mode the mels are pre-computed using a
# spectrogram generator
if self.finetune:
audio, audio_len, audio_mel = batch
# else, we compute the mel using the ground truth audio
else:
audio, audio_len = batch
# mel as input for generator
audio_mel, _ = self.audio_to_melspec_precessor(audio, audio_len)
# mel as input for L1 mel loss
audio_trg_mel, _ = self.trg_melspec_fn(audio, audio_len)
audio = audio.unsqueeze(1)
audio_pred = self.generator(x=audio_mel)
audio_pred_mel, _ = self.trg_melspec_fn(audio_pred.squeeze(1), audio_len)
# train discriminator
self.optim_d.zero_grad()
mpd_score_real, mpd_score_gen, _, _ = self.mpd(y=audio, y_hat=audio_pred.detach())
loss_disc_mpd, _, _ = self.discriminator_loss(
disc_real_outputs=mpd_score_real, disc_generated_outputs=mpd_score_gen
)
msd_score_real, msd_score_gen, _, _ = self.msd(y=audio, y_hat=audio_pred.detach())
loss_disc_msd, _, _ = self.discriminator_loss(
disc_real_outputs=msd_score_real, disc_generated_outputs=msd_score_gen
)
loss_d = loss_disc_msd + loss_disc_mpd
self.manual_backward(loss_d, self.optim_d)
self.optim_d.step()
# train generator
self.optim_g.zero_grad()
loss_mel = F.l1_loss(audio_pred_mel, audio_trg_mel) * 45
_, mpd_score_gen, fmap_mpd_real, fmap_mpd_gen = self.mpd(y=audio, y_hat=audio_pred)
_, msd_score_gen, fmap_msd_real, fmap_msd_gen = self.msd(y=audio, y_hat=audio_pred)
loss_fm_mpd = self.feature_loss(fmap_r=fmap_mpd_real, fmap_g=fmap_mpd_gen)
loss_fm_msd = self.feature_loss(fmap_r=fmap_msd_real, fmap_g=fmap_msd_gen)
loss_gen_mpd, _ = self.generator_loss(disc_outputs=mpd_score_gen)
loss_gen_msd, _ = self.generator_loss(disc_outputs=msd_score_gen)
loss_g = loss_gen_msd + loss_gen_mpd + loss_fm_msd + loss_fm_mpd + loss_mel
self.manual_backward(loss_g, self.optim_g)
self.optim_g.step()
metrics = {
"g_l1_loss": loss_mel,
"g_loss_fm_mpd": loss_fm_mpd,
"g_loss_fm_msd": loss_fm_msd,
"g_loss_gen_mpd": loss_gen_mpd,
"g_loss_gen_msd": loss_gen_msd,
"g_loss": loss_g,
"d_loss_mpd": loss_disc_mpd,
"d_loss_msd": loss_disc_msd,
"d_loss": loss_d,
"global_step": self.global_step,
"lr": self.optim_g.param_groups[0]['lr'],
}
self.log_dict(metrics, on_step=False, on_epoch=True, sync_dist=True)
def validation_step(self, batch, batch_idx):
if self.finetune:
audio, audio_len, audio_mel = batch
audio_mel_len = [audio_mel.shape[1]] * audio_mel.shape[0]
else:
audio, audio_len = batch
audio_mel, audio_mel_len = self.audio_to_melspec_precessor(audio, audio_len)
audio_pred = self(spec=audio_mel)
# perform bias denoising
pred_denoised = self._bias_denoise(audio_pred, audio_mel).squeeze(1)
pred_denoised_mel, _ = self.audio_to_melspec_precessor(pred_denoised, audio_len)
if self.finetune:
gt_mel, gt_mel_len = self.audio_to_melspec_precessor(audio, audio_len)
audio_pred_mel, _ = self.audio_to_melspec_precessor(audio_pred.squeeze(1), audio_len)
loss_mel = F.l1_loss(audio_mel, audio_pred_mel)
self.log("val_loss", loss_mel, prog_bar=True, sync_dist=True)
# plot audio once per epoch
if batch_idx == 0 and isinstance(self.logger, WandbLogger) and HAVE_WANDB:
clips = []
specs = []
for i in range(min(5, audio.shape[0])):
clips += [
wandb.Audio(
audio[i, : audio_len[i]].data.cpu().numpy(),
caption=f"real audio {i}",
sample_rate=self.sample_rate,
),
wandb.Audio(
audio_pred[i, 0, : audio_len[i]].data.cpu().numpy().astype('float32'),
caption=f"generated audio {i}",
sample_rate=self.sample_rate,
),
wandb.Audio(
pred_denoised[i, : audio_len[i]].data.cpu().numpy(),
caption=f"denoised audio {i}",
sample_rate=self.sample_rate,
),
]
specs += [
wandb.Image(
plot_spectrogram_to_numpy(audio_mel[i, :, : audio_mel_len[i]].data.cpu().numpy()),
caption=f"input mel {i}",
),
wandb.Image(
plot_spectrogram_to_numpy(audio_pred_mel[i, :, : audio_mel_len[i]].data.cpu().numpy()),
caption=f"output mel {i}",
),
wandb.Image(
plot_spectrogram_to_numpy(pred_denoised_mel[i, :, : audio_mel_len[i]].data.cpu().numpy()),
caption=f"denoised mel {i}",
),
]
if self.finetune:
specs += [
wandb.Image(
plot_spectrogram_to_numpy(gt_mel[i, :, : audio_mel_len[i]].data.cpu().numpy()),
caption=f"gt mel {i}",
),
]
self.logger.experiment.log({"audio": clips, "specs": specs}, commit=False)
def _bias_denoise(self, audio, mel):
def stft(x):
comp = stft_patch(x.squeeze(1), n_fft=1024, hop_length=256, win_length=1024)
real, imag = comp[..., 0], comp[..., 1]
mags = torch.sqrt(real ** 2 + imag ** 2)
phase = torch.atan2(imag, real)
return mags, phase
def istft(mags, phase):
comp = torch.stack([mags * torch.cos(phase), mags * torch.sin(phase)], dim=-1)
x = torch.istft(comp, n_fft=1024, hop_length=256, win_length=1024)
return x
# create bias tensor
if self.stft_bias is None:
audio_bias = self(spec=torch.zeros_like(mel, device=mel.device))
self.stft_bias, _ = stft(audio_bias)
self.stft_bias = self.stft_bias[:, :, 0][:, :, None]
audio_mags, audio_phase = stft(audio)
audio_mags = audio_mags - self.cfg.denoise_strength * self.stft_bias
audio_mags = torch.clamp(audio_mags, 0.0)
audio_denoised = istft(audio_mags, audio_phase).unsqueeze(1)
return audio_denoised
def __setup_dataloader_from_config(self, cfg, shuffle_should_be: bool = True, name: str = "train"):
if "dataset" not in cfg or not isinstance(cfg.dataset, DictConfig):
raise ValueError(f"No dataset for {name}")
if "dataloader_params" not in cfg or not isinstance(cfg.dataloader_params, DictConfig):
raise ValueError(f"No dataloder_params for {name}")
if shuffle_should_be:
if 'shuffle' not in cfg.dataloader_params:
logging.warning(
f"Shuffle should be set to True for {self}'s {name} dataloader but was not found in its "
"config. Manually setting to True"
)
with open_dict(cfg["dataloader_params"]):
cfg.dataloader_params.shuffle = True
elif not cfg.dataloader_params.shuffle:
logging.error(f"The {name} dataloader for {self} has shuffle set to False!!!")
elif not shuffle_should_be and cfg.dataloader_params.shuffle:
logging.error(f"The {name} dataloader for {self} has shuffle set to True!!!")
dataset = instantiate(cfg.dataset)
return torch.utils.data.DataLoader(dataset, collate_fn=dataset.collate_fn, **cfg.dataloader_params)
def setup_training_data(self, cfg):
self._train_dl = self.__setup_dataloader_from_config(cfg)
def setup_validation_data(self, cfg):
self._validation_dl = self.__setup_dataloader_from_config(cfg, shuffle_should_be=False, name="validation")
@classmethod
def list_available_models(cls) -> 'Optional[Dict[str, str]]':
# TODO
pass
class GlowVocoder(Vocoder):
""" Base class for all Vocoders that use a Glow or reversible Flow-based setup. All child class are expected
to have a parameter called audio_to_melspec_precessor that is an instance of
nemo.collections.asr.parts.FilterbankFeatures"""
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self._mode = OperationMode.infer
self.stft = None
self.istft = None
self.n_mel = None
self.bias_spect = None
@property
def mode(self):
return self._mode
@contextmanager
def temp_mode(self, mode):
old_mode = self.mode
self.mode = mode
try:
yield
finally:
self.mode = old_mode
@contextmanager
def nemo_infer(
self
): # Prepend with nemo to avoid any .infer() clashes with lightning or pytorch
with ExitStack() as stack:
stack.enter_context(self.temp_mode(OperationMode.infer))
stack.enter_context(torch.no_grad())
yield
def check_children_attributes(self):
if self.stft is None:
if isinstance(self.audio_to_melspec_precessor.stft, STFT):
logging.warning(
"torch_stft is deprecated. Please change your model to use torch.stft and torch.istft instead."
)
self.stft = self.audio_to_melspec_precessor.stft.transform
self.istft = self.audio_to_melspec_precessor.stft.inverse
else:
try:
n_fft = self.audio_to_melspec_precessor.n_fft
hop_length = self.audio_to_melspec_precessor.hop_length
win_length = self.audio_to_melspec_precessor.win_length
window = self.audio_to_melspec_precessor.window.to(
self.device)
except AttributeError as e:
raise AttributeError(
f"{self} could not find a valid audio_to_melspec_precessor. GlowVocoder requires child class "
"to have audio_to_melspec_precessor defined to obtain stft parameters. "
"audio_to_melspec_precessor requires n_fft, hop_length, win_length, window, and nfilt to be "
"defined.") from e
def yet_another_patch(audio, n_fft, hop_length, win_length,
window):
spec = stft_patch(audio,
n_fft=n_fft,
hop_length=hop_length,
win_length=win_length,
window=window)
if spec.dtype in [torch.cfloat, torch.cdouble]:
spec = torch.view_as_real(spec)
return torch.sqrt(spec.pow(2).sum(-1)), torch.atan2(
spec[..., -1], spec[..., 0])
self.stft = lambda x: yet_another_patch(
x,
n_fft=n_fft,
hop_length=hop_length,
win_length=win_length,
window=window,
)
self.istft = lambda x, y: istft_patch(
torch.complex(x * torch.cos(y), x * torch.sin(y)),
n_fft=n_fft,
hop_length=hop_length,
win_length=win_length,
window=window,
)
if self.n_mel is None:
try:
self.n_mel = self.audio_to_melspec_precessor.nfilt
except AttributeError as e:
raise AttributeError(
f"{self} could not find a valid audio_to_melspec_precessor. GlowVocoder requires child class to "
"have audio_to_melspec_precessor defined to obtain stft parameters. audio_to_melspec_precessor "
"requires nfilt to be defined.") from e
def update_bias_spect(self):
self.check_children_attributes() # Ensure stft parameters are defined
with self.nemo_infer():
spect = torch.zeros((1, self.n_mel, 88)).to(self.device)
bias_audio = self.convert_spectrogram_to_audio(spec=spect,
sigma=0.0,
denoise=False)
bias_spect, _ = self.stft(bias_audio)
self.bias_spect = bias_spect[..., 0][..., None]
@typecheck(
input_types={
"audio": NeuralType(('B', 'T'), AudioSignal()),
"strength": NeuralType(optional=True)
},
output_types={"audio": NeuralType(('B', 'T'), AudioSignal())},
)
def denoise(self, audio: 'torch.tensor', strength: float = 0.01):
self.check_children_attributes(
) # Ensure self.n_mel and self.stft are defined
if self.bias_spect is None:
self.update_bias_spect()
audio_spect, audio_angles = self.stft(audio)
audio_spect_denoised = audio_spect - self.bias_spect.to(
audio.device) * strength
audio_spect_denoised = torch.clamp(audio_spect_denoised, 0.0)
audio_denoised = self.istft(audio_spect_denoised, audio_angles)
return audio_denoised
def input_types(self):
return {
"audio": NeuralType(('B', 'T'), AudioSignal()),
"audio_len": NeuralType(('B'), LengthsType()),
"run_inverse": NeuralType(optional=True),
}
class SqueezeWaveModel(Vocoder):
""" SqueezeWave model that generates audio conditioned on mel-spectrogram
"""
def __init__(self, cfg: DictConfig, trainer: 'Trainer' = None):
if isinstance(cfg, dict):
cfg = OmegaConf.create(cfg)
super().__init__(cfg=cfg, trainer=trainer)
schema = OmegaConf.structured(SqueezeWaveConfig)
# ModelPT ensures that cfg is a DictConfig, but do this second check in case ModelPT changes
if isinstance(cfg, dict):
cfg = OmegaConf.create(cfg)
elif not isinstance(cfg, DictConfig):
raise ValueError(
f"cfg was type: {type(cfg)}. Expected either a dict or a DictConfig"
)
# Ensure passed cfg is compliant with schema
OmegaConf.merge(cfg, schema)
self.pad_value = self._cfg.preprocessor.params.pad_value
self.sigma = self._cfg.sigma
self.audio_to_melspec_precessor = instantiate(self._cfg.preprocessor)
self.squeezewave = instantiate(self._cfg.squeezewave)
self.mode = OperationMode.infer
self.loss = WaveGlowLoss() # Same loss as WaveGlow
@property
def input_types(self):
return {
"audio": NeuralType(('B', 'T'), AudioSignal()),
"audio_len": NeuralType(('B'), LengthsType()),
"run_inverse": NeuralType(optional=True),
}
@property
def output_types(self):
if self.mode == OperationMode.training or self.mode == OperationMode.validation:
output_dict = {
"pred_normal_dist":
NeuralType(('B', 'flowgroup', 'T'),
NormalDistributionSamplesType()),
"log_s_list":
NeuralType(('B', 'flowgroup', 'T'),
VoidType()), # TODO: Figure out a good typing
"log_det_W_list":
NeuralType(elements_type=VoidType()
), # TODO: Figure out a good typing
}
if self.mode == OperationMode.validation:
output_dict["audio_pred"] = NeuralType(('B', 'T'),
AudioSignal())
output_dict["spec"] = NeuralType(('B', 'T', 'D'),
MelSpectrogramType())
output_dict["spec_len"] = NeuralType(('B'), LengthsType())
return output_dict
return {
"audio_pred": NeuralType(('B', 'T'), AudioSignal()),
}
@typecheck()
def forward(self, *, audio, audio_len, run_inverse=True):
if self.mode != self.squeezewave.mode:
raise ValueError(
f"SqueezeWaveModel's mode {self.mode} does not match SqueezeWaveModule's mode {self.squeezewave.mode}"
)
spec, spec_len = self.audio_to_melspec_precessor(audio, audio_len)
tensors = self.squeezewave(spec=spec,
audio=audio,
run_inverse=run_inverse)
if self.mode == OperationMode.training:
return tensors[:-1] # z, log_s_list, log_det_W_list
elif self.mode == OperationMode.validation:
z, log_s_list, log_det_W_list, audio_pred = tensors
return z, log_s_list, log_det_W_list, audio_pred, spec, spec_len
return tensors # audio_pred
@typecheck(
input_types={
"spec": NeuralType(('B', 'D', 'T'), MelSpectrogramType()),
"sigma": NeuralType(optional=True)
},
output_types={"audio": NeuralType(('B', 'T'), AudioSignal())},
)
def convert_spectrogram_to_audio(self,
spec: torch.Tensor,
sigma: bool = 1.0) -> torch.Tensor:
self.eval()
self.mode = OperationMode.infer
self.squeezewave.mode = OperationMode.infer
with torch.no_grad():
audio = self.squeezewave(spec=spec,
run_inverse=True,
audio=None,
sigma=sigma)
return audio
def training_step(self, batch, batch_idx):
self.mode = OperationMode.training
self.squeezewave.mode = OperationMode.training
audio, audio_len = batch
z, log_s_list, log_det_W_list = self.forward(audio=audio,
audio_len=audio_len,
run_inverse=False)
loss = self.loss(z=z,
log_s_list=log_s_list,
log_det_W_list=log_det_W_list,
sigma=self.sigma)
return {
'loss': loss,
'progress_bar': {
'training_loss': loss
},
'log': {
'loss': loss
},
}
def validation_step(self, batch, batch_idx):
self.mode = OperationMode.validation
self.squeezewave.mode = OperationMode.validation
audio, audio_len = batch
z, log_s_list, log_det_W_list, audio_pred, spec, spec_len = self.forward(
audio=audio, audio_len=audio_len, run_inverse=(batch_idx == 0))
loss = self.loss(z=z,
log_s_list=log_s_list,
log_det_W_list=log_det_W_list,
sigma=self.sigma)
return {
"val_loss": loss,
"audio_pred": audio_pred,
"mel_target": spec,
"mel_len": spec_len,
}
def validation_epoch_end(self, outputs):
if self.logger is not None and self.logger.experiment is not None:
waveglow_log_to_tb_func(
self.logger.experiment,
outputs[0].values(),
self.global_step,
tag="eval",
mel_fb=self.audio_to_melspec_precessor.fb,
)
avg_loss = torch.stack([x['val_loss'] for x in outputs]).mean()
self.log('val_loss', avg_loss)
def __setup_dataloader_from_config(self,
cfg,
shuffle_should_be: bool = True,
name: str = "train"):
if "dataset" not in cfg or not isinstance(cfg.dataset, DictConfig):
raise ValueError(f"No dataset for {name}") # TODO
if "dataloader_params" not in cfg or not isinstance(
cfg.dataloader_params, DictConfig):
raise ValueError(f"No dataloder_params for {name}") # TODO
if shuffle_should_be:
if 'shuffle' not in cfg.dataloader_params:
logging.warning(
f"Shuffle should be set to True for {self}'s {name} dataloader but was not found in its "
"config. Manually setting to True")
with open_dict(cfg["dataloader_params"]):
cfg.dataloader_params.shuffle = True
elif not cfg.dataloader_params.shuffle:
logging.error(
f"The {name} dataloader for {self} has shuffle set to False!!!"
)
elif not shuffle_should_be and cfg.dataloader_params.shuffle:
logging.error(
f"The {name} dataloader for {self} has shuffle set to True!!!")
dataset = instantiate(cfg.dataset)
return torch.utils.data.DataLoader(dataset,
collate_fn=dataset.collate_fn,
**cfg.dataloader_params)
def setup_training_data(self, cfg):
self._train_dl = self.__setup_dataloader_from_config(cfg)
def setup_validation_data(self, cfg):
self._validation_dl = self.__setup_dataloader_from_config(
cfg, shuffle_should_be=False, name="validation")
@classmethod
def list_available_models(cls) -> 'List[PretrainedModelInfo]':
"""
This method returns a list of pre-trained model which can be instantiated directly from NVIDIA's NGC cloud.
Returns:
List of available pre-trained models.
"""
list_of_models = []
model = PretrainedModelInfo(
pretrained_model_name="SqueezeWave-22050Hz",
location=
"https://api.ngc.nvidia.com/v2/models/nvidia/nemottsmodels/versions/1.0.0a5/files/SqueezeWave-22050Hz.nemo",
description=
"This model is trained on LJSpeech sampled at 22050Hz, and can be used as an universal vocoder.",
class_=cls,
)
list_of_models.append(model)
return list_of_models
class HifiGanModel(Vocoder, Exportable):
def __init__(self, cfg: DictConfig, trainer: 'Trainer' = None):
if isinstance(cfg, dict):
cfg = OmegaConf.create(cfg)
super().__init__(cfg=cfg, trainer=trainer)
self.audio_to_melspec_precessor = instantiate(cfg.preprocessor)
# use a different melspec extractor because:
# 1. we need to pass grads
# 2. we need remove fmax limitation
self.trg_melspec_fn = instantiate(cfg.preprocessor,
highfreq=None,
use_grads=True)
self.generator = instantiate(cfg.generator)
self.mpd = MultiPeriodDiscriminator(
debug=cfg.debug if "debug" in cfg else False)
self.msd = MultiScaleDiscriminator(
debug=cfg.debug if "debug" in cfg else False)
self.feature_loss = FeatureMatchingLoss()
self.discriminator_loss = DiscriminatorLoss()
self.generator_loss = GeneratorLoss()
self.l1_factor = cfg.get("l1_loss_factor", 45)
self.sample_rate = self._cfg.preprocessor.sample_rate
self.stft_bias = None
if self._train_dl and isinstance(self._train_dl.dataset,
MelAudioDataset):
self.input_as_mel = True
else:
self.input_as_mel = False
self.automatic_optimization = False
def configure_optimizers(self):
self.optim_g = instantiate(
self._cfg.optim,
params=self.generator.parameters(),
)
self.optim_d = instantiate(
self._cfg.optim,
params=itertools.chain(self.msd.parameters(),
self.mpd.parameters()),
)
self.scheduler_g = CosineAnnealing(
optimizer=self.optim_g,
max_steps=self._cfg.max_steps,
min_lr=self._cfg.sched.min_lr,
warmup_steps=self._cfg.sched.warmup_ratio * self._cfg.max_steps,
) # Use warmup to delay start
sch1_dict = {
'scheduler': self.scheduler_g,
'interval': 'step',
}
self.scheduler_d = CosineAnnealing(
optimizer=self.optim_d,
max_steps=self._cfg.max_steps,
min_lr=self._cfg.sched.min_lr,
)
sch2_dict = {
'scheduler': self.scheduler_d,
'interval': 'step',
}
return [self.optim_g, self.optim_d], [sch1_dict, sch2_dict]
@property
def input_types(self):
return {
"spec": NeuralType(('B', 'D', 'T'), MelSpectrogramType()),
}
@property
def output_types(self):
return {
"audio": NeuralType(('B', 'S', 'T'),
AudioSignal(self.sample_rate)),
}
@typecheck()
def forward(self, *, spec):
"""
Runs the generator, for inputs and outputs see input_types, and output_types
"""
return self.generator(x=spec)
@typecheck(
input_types={
"spec": NeuralType(('B', 'C', 'T'), MelSpectrogramType())
},
output_types={"audio": NeuralType(('B', 'T'), AudioSignal())},
)
def convert_spectrogram_to_audio(self,
spec: 'torch.tensor') -> 'torch.tensor':
return self(spec=spec).squeeze(1)
def training_step(self, batch, batch_idx, optimizer_idx):
# if in finetune mode the mels are pre-computed using a
# spectrogram generator
if self.input_as_mel:
audio, audio_len, audio_mel = batch
# else, we compute the mel using the ground truth audio
else:
audio, audio_len = batch
# mel as input for generator
audio_mel, _ = self.audio_to_melspec_precessor(audio, audio_len)
# mel as input for L1 mel loss
audio_trg_mel, _ = self.trg_melspec_fn(audio, audio_len)
audio = audio.unsqueeze(1)
audio_pred = self.generator(x=audio_mel)
audio_pred_mel, _ = self.trg_melspec_fn(audio_pred.squeeze(1),
audio_len)
# train discriminator
self.optim_d.zero_grad()
mpd_score_real, mpd_score_gen, _, _ = self.mpd(
y=audio, y_hat=audio_pred.detach())
loss_disc_mpd, _, _ = self.discriminator_loss(
disc_real_outputs=mpd_score_real,
disc_generated_outputs=mpd_score_gen)
msd_score_real, msd_score_gen, _, _ = self.msd(
y=audio, y_hat=audio_pred.detach())
loss_disc_msd, _, _ = self.discriminator_loss(
disc_real_outputs=msd_score_real,
disc_generated_outputs=msd_score_gen)
loss_d = loss_disc_msd + loss_disc_mpd
self.manual_backward(loss_d)
self.optim_d.step()
# train generator
self.optim_g.zero_grad()
loss_mel = F.l1_loss(audio_pred_mel, audio_trg_mel)
_, mpd_score_gen, fmap_mpd_real, fmap_mpd_gen = self.mpd(
y=audio, y_hat=audio_pred)
_, msd_score_gen, fmap_msd_real, fmap_msd_gen = self.msd(
y=audio, y_hat=audio_pred)
loss_fm_mpd = self.feature_loss(fmap_r=fmap_mpd_real,
fmap_g=fmap_mpd_gen)
loss_fm_msd = self.feature_loss(fmap_r=fmap_msd_real,
fmap_g=fmap_msd_gen)
loss_gen_mpd, _ = self.generator_loss(disc_outputs=mpd_score_gen)
loss_gen_msd, _ = self.generator_loss(disc_outputs=msd_score_gen)
loss_g = loss_gen_msd + loss_gen_mpd + loss_fm_msd + loss_fm_mpd + loss_mel * self.l1_factor
self.manual_backward(loss_g)
self.optim_g.step()
# run schedulers
sch1, sch2 = self.lr_schedulers()
sch1.step()
sch2.step()
metrics = {
"g_loss_fm_mpd": loss_fm_mpd,
"g_loss_fm_msd": loss_fm_msd,
"g_loss_gen_mpd": loss_gen_mpd,
"g_loss_gen_msd": loss_gen_msd,
"g_loss": loss_g,
"d_loss_mpd": loss_disc_mpd,
"d_loss_msd": loss_disc_msd,
"d_loss": loss_d,
"global_step": self.global_step,
"lr": self.optim_g.param_groups[0]['lr'],
}
self.log_dict(metrics, on_step=True, sync_dist=True)
self.log("g_l1_loss",
loss_mel,
prog_bar=True,
logger=False,
sync_dist=True)
def validation_step(self, batch, batch_idx):
if self.input_as_mel:
audio, audio_len, audio_mel = batch
audio_mel_len = [audio_mel.shape[1]] * audio_mel.shape[0]
else:
audio, audio_len = batch
audio_mel, audio_mel_len = self.audio_to_melspec_precessor(
audio, audio_len)
audio_pred = self(spec=audio_mel)
# perform bias denoising
pred_denoised = self._bias_denoise(audio_pred, audio_mel).squeeze(1)
pred_denoised_mel, _ = self.audio_to_melspec_precessor(
pred_denoised, audio_len)
if self.input_as_mel:
gt_mel, gt_mel_len = self.audio_to_melspec_precessor(
audio, audio_len)
audio_pred_mel, _ = self.audio_to_melspec_precessor(
audio_pred.squeeze(1), audio_len)
loss_mel = F.l1_loss(audio_mel, audio_pred_mel)
self.log_dict({"val_loss": loss_mel}, on_epoch=True, sync_dist=True)
# plot audio once per epoch
if batch_idx == 0 and isinstance(self.logger,
WandbLogger) and HAVE_WANDB:
clips = []
specs = []
for i in range(min(5, audio.shape[0])):
clips += [
wandb.Audio(
audio[i, :audio_len[i]].data.cpu().numpy(),
caption=f"real audio {i}",
sample_rate=self.sample_rate,
),
wandb.Audio(
audio_pred[i,
0, :audio_len[i]].data.cpu().numpy().astype(
'float32'),
caption=f"generated audio {i}",
sample_rate=self.sample_rate,
),
wandb.Audio(
pred_denoised[i, :audio_len[i]].data.cpu().numpy(),
caption=f"denoised audio {i}",
sample_rate=self.sample_rate,
),
]
specs += [
wandb.Image(
plot_spectrogram_to_numpy(audio_mel[
i, :, :audio_mel_len[i]].data.cpu().numpy()),
caption=f"input mel {i}",
),
wandb.Image(
plot_spectrogram_to_numpy(audio_pred_mel[
i, :, :audio_mel_len[i]].data.cpu().numpy()),
caption=f"output mel {i}",
),
wandb.Image(
plot_spectrogram_to_numpy(pred_denoised_mel[
i, :, :audio_mel_len[i]].data.cpu().numpy()),
caption=f"denoised mel {i}",
),
]
if self.input_as_mel:
specs += [
wandb.Image(
plot_spectrogram_to_numpy(gt_mel[
i, :, :audio_mel_len[i]].data.cpu().numpy()),
caption=f"gt mel {i}",
),
]
self.logger.experiment.log({"audio": clips, "specs": specs})
def _bias_denoise(self, audio, mel):
def stft(x):
comp = stft_patch(x.squeeze(1),
n_fft=1024,
hop_length=256,
win_length=1024)
real, imag = comp[..., 0], comp[..., 1]
mags = torch.sqrt(real**2 + imag**2)
phase = torch.atan2(imag, real)
return mags, phase
def istft(mags, phase):
comp = torch.stack(
[mags * torch.cos(phase), mags * torch.sin(phase)], dim=-1)
x = torch.istft(comp, n_fft=1024, hop_length=256, win_length=1024)
return x
# create bias tensor
if self.stft_bias is None or self.stft_bias.shape[0] != audio.shape[0]:
audio_bias = self(spec=torch.zeros_like(mel, device=mel.device))
self.stft_bias, _ = stft(audio_bias)
self.stft_bias = self.stft_bias[:, :, 0][:, :, None]
audio_mags, audio_phase = stft(audio)
audio_mags = audio_mags - self.cfg.get("denoise_strength",
0.0025) * self.stft_bias
audio_mags = torch.clamp(audio_mags, 0.0)
audio_denoised = istft(audio_mags, audio_phase).unsqueeze(1)
return audio_denoised
def __setup_dataloader_from_config(self,
cfg,
shuffle_should_be: bool = True,
name: str = "train"):
if "dataset" not in cfg or not isinstance(cfg.dataset, DictConfig):
raise ValueError(f"No dataset for {name}")
if "dataloader_params" not in cfg or not isinstance(
cfg.dataloader_params, DictConfig):
raise ValueError(f"No dataloder_params for {name}")
if shuffle_should_be:
if 'shuffle' not in cfg.dataloader_params:
logging.warning(
f"Shuffle should be set to True for {self}'s {name} dataloader but was not found in its "
"config. Manually setting to True")
with open_dict(cfg["dataloader_params"]):
cfg.dataloader_params.shuffle = True
elif not cfg.dataloader_params.shuffle:
logging.error(
f"The {name} dataloader for {self} has shuffle set to False!!!"
)
elif not shuffle_should_be and cfg.dataloader_params.shuffle:
logging.error(
f"The {name} dataloader for {self} has shuffle set to True!!!")
dataset = instantiate(cfg.dataset)
return torch.utils.data.DataLoader(dataset,
collate_fn=dataset.collate_fn,
**cfg.dataloader_params)
def setup_training_data(self, cfg):
self._train_dl = self.__setup_dataloader_from_config(cfg)
def setup_validation_data(self, cfg):
self._validation_dl = self.__setup_dataloader_from_config(
cfg, shuffle_should_be=False, name="validation")
@classmethod
def list_available_models(cls) -> 'Optional[Dict[str, str]]':
list_of_models = []
model = PretrainedModelInfo(
pretrained_model_name="tts_hifigan",
location=
"https://api.ngc.nvidia.com/v2/models/nvidia/nemo/tts_hifigan/versions/1.0.0rc1/files/tts_hifigan.nemo",
description=
"This model is trained on LJSpeech audio sampled at 22050Hz and mel spectrograms generated from Tacotron2, TalkNet, and FastPitch. This model has been tested on generating female English voices with an American accent.",
class_=cls,
)
list_of_models.append(model)
return list_of_models
def load_state_dict(self, state_dict, strict=True):
# override load_state_dict to give us some flexibility to be backward-compatible
# with old checkpoints
new_state_dict = {}
num_resblocks = len(self.cfg['generator']['resblock_kernel_sizes'])
for k, v in state_dict.items():
new_k = k
if 'resblocks' in k:
parts = k.split(".")
# only do this is the checkpoint type is older
if len(parts) == 6:
layer = int(parts[2])
new_layer = f"{layer // num_resblocks}.{layer % num_resblocks}"
new_k = f"generator.resblocks.{new_layer}.{'.'.join(parts[3:])}"
new_state_dict[new_k] = v
super().load_state_dict(new_state_dict, strict=strict)
def _prepare_for_export(self, **kwargs):
"""
Override this method to prepare module for export. This is in-place operation.
Base version does common necessary module replacements (Apex etc)
"""
if self.generator is not None:
self.generator.remove_weight_norm()
def input_example(self):
"""
Generates input examples for tracing etc.
Returns:
A tuple of input examples.
"""
par = next(self.parameters())
mel = torch.randn((1, self.cfg['preprocessor']['nfilt'], 96),
device=par.device,
dtype=par.dtype)
return ({'spec': mel}, )
def forward_for_export(self, spec):
"""
Runs the generator, for inputs and outputs see input_types, and output_types
"""
return self.generator(x=spec)
def input_types(self):
return {
"spec": NeuralType(('B', 'D', 'T'), MelSpectrogramType()),
"audio": NeuralType(('B', 'T'), AudioSignal(), optional=True),
"sigma": NeuralType(optional=True),
}
