# Copyright (c) 2024 Alibaba Inc (authors: Xiang Lyu, Kai Hu) # # 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. """HIFI-GAN""" import typing as tp import numpy as np from scipy.signal import get_window import torch import torch.nn as nn import torch.nn.functional as F from torch.nn import Conv1d from torch.nn import ConvTranspose1d from torch.nn.utils import remove_weight_norm from torch.nn.utils import weight_norm from torch.distributions.uniform import Uniform from torch import sin from torch.nn.parameter import Parameter """hifigan based generator implementation. This code is modified from https://github.com/jik876/hifi-gan ,https://github.com/kan-bayashi/ParallelWaveGAN and https://github.com/NVIDIA/BigVGAN """ class Snake(nn.Module): ''' Implementation of a sine-based periodic activation function Shape: - Input: (B, C, T) - Output: (B, C, T), same shape as the input Parameters: - alpha - trainable parameter References: - This activation function is from this paper by Liu Ziyin, Tilman Hartwig, Masahito Ueda: https://arxiv.org/abs/2006.08195 Examples: >>> a1 = snake(256) >>> x = torch.randn(256) >>> x = a1(x) ''' def __init__(self, in_features, alpha=1.0, alpha_trainable=True, alpha_logscale=False): ''' Initialization. INPUT: - in_features: shape of the input - alpha: trainable parameter alpha is initialized to 1 by default, higher values = higher-frequency. alpha will be trained along with the rest of your model. ''' super(Snake, self).__init__() self.in_features = in_features # initialize alpha self.alpha_logscale = alpha_logscale if self.alpha_logscale: # log scale alphas initialized to zeros self.alpha = Parameter(torch.zeros(in_features) * alpha) else: # linear scale alphas initialized to ones self.alpha = Parameter(torch.ones(in_features) * alpha) self.alpha.requires_grad = alpha_trainable self.no_div_by_zero = 0.000000001 def forward(self, x): ''' Forward pass of the function. Applies the function to the input elementwise. Snake ∶= x + 1/a * sin^2 (xa) ''' alpha = self.alpha.unsqueeze(0).unsqueeze(-1) # line up with x to [B, C, T] if self.alpha_logscale: alpha = torch.exp(alpha) x = x + (1.0 / (alpha + self.no_div_by_zero)) * pow(sin(x * alpha), 2) return x def get_padding(kernel_size, dilation=1): return int((kernel_size * dilation - dilation) / 2) def init_weights(m, mean=0.0, std=0.01): classname = m.__class__.__name__ if classname.find("Conv") != -1: m.weight.data.normal_(mean, std) class ResBlock(torch.nn.Module): """Residual block module in HiFiGAN/BigVGAN.""" def __init__( self, channels: int = 512, kernel_size: int = 3, dilations: tp.List[int] = [1, 3, 5], ): super(ResBlock, self).__init__() self.convs1 = nn.ModuleList() self.convs2 = nn.ModuleList() for dilation in dilations: self.convs1.append( weight_norm( Conv1d( channels, channels, kernel_size, 1, dilation=dilation, padding=get_padding(kernel_size, dilation) ) ) ) self.convs2.append( weight_norm( Conv1d( channels, channels, kernel_size, 1, dilation=1, padding=get_padding(kernel_size, 1) ) ) ) self.convs1.apply(init_weights) self.convs2.apply(init_weights) self.activations1 = nn.ModuleList([ Snake(channels, alpha_logscale=False) for _ in range(len(self.convs1)) ]) self.activations2 = nn.ModuleList([ Snake(channels, alpha_logscale=False) for _ in range(len(self.convs2)) ]) def forward(self, x: torch.Tensor) -> torch.Tensor: for idx in range(len(self.convs1)): xt = self.activations1[idx](x) xt = self.convs1[idx](xt) xt = self.activations2[idx](xt) xt = self.convs2[idx](xt) x = xt + x return x def remove_weight_norm(self): for idx in range(len(self.convs1)): remove_weight_norm(self.convs1[idx]) remove_weight_norm(self.convs2[idx]) class SineGen(torch.nn.Module): """ Definition of sine generator SineGen(samp_rate, harmonic_num = 0, sine_amp = 0.1, noise_std = 0.003, voiced_threshold = 0, flag_for_pulse=False) samp_rate: sampling rate in Hz harmonic_num: number of harmonic overtones (default 0) sine_amp: amplitude of sine-wavefrom (default 0.1) noise_std: std of Gaussian noise (default 0.003) voiced_thoreshold: F0 threshold for U/V classification (default 0) flag_for_pulse: this SinGen is used inside PulseGen (default False) Note: when flag_for_pulse is True, the first time step of a voiced segment is always sin(np.pi) or cos(0) """ def __init__(self, samp_rate, harmonic_num=0, sine_amp=0.1, noise_std=0.003, voiced_threshold=0): super(SineGen, self).__init__() self.sine_amp = sine_amp self.noise_std = noise_std self.harmonic_num = harmonic_num self.sampling_rate = samp_rate self.voiced_threshold = voiced_threshold def _f02uv(self, f0): # generate uv signal uv = (f0 > self.voiced_threshold).type(torch.float32) return uv @torch.no_grad() def forward(self, f0): """ :param f0: [B, 1, sample_len], Hz :return: [B, 1, sample_len] """ F_mat = torch.zeros((f0.size(0), self.harmonic_num + 1, f0.size(-1))).to(f0.device) for i in range(self.harmonic_num + 1): F_mat[:, i: i + 1, :] = f0 * (i + 1) / self.sampling_rate theta_mat = 2 * np.pi * (torch.cumsum(F_mat, dim=-1) % 1) u_dist = Uniform(low=-np.pi, high=np.pi) phase_vec = u_dist.sample(sample_shape=(f0.size(0), self.harmonic_num + 1, 1)).to(F_mat.device) phase_vec[:, 0, :] = 0 # generate sine waveforms sine_waves = self.sine_amp * torch.sin(theta_mat + phase_vec) # generate uv signal uv = self._f02uv(f0) # noise: for unvoiced should be similar to sine_amp # std = self.sine_amp/3 -> max value ~ self.sine_amp # . for voiced regions is self.noise_std noise_amp = uv * self.noise_std + (1 - uv) * self.sine_amp / 3 noise = noise_amp * torch.randn_like(sine_waves) # first: set the unvoiced part to 0 by uv # then: additive noise sine_waves = sine_waves * uv + noise return sine_waves, uv, noise class SourceModuleHnNSF(torch.nn.Module): """ SourceModule for hn-nsf SourceModule(sampling_rate, harmonic_num=0, sine_amp=0.1, add_noise_std=0.003, voiced_threshod=0) sampling_rate: sampling_rate in Hz harmonic_num: number of harmonic above F0 (default: 0) sine_amp: amplitude of sine source signal (default: 0.1) add_noise_std: std of additive Gaussian noise (default: 0.003) note that amplitude of noise in unvoiced is decided by sine_amp voiced_threshold: threhold to set U/V given F0 (default: 0) Sine_source, noise_source = SourceModuleHnNSF(F0_sampled) F0_sampled (batchsize, length, 1) Sine_source (batchsize, length, 1) noise_source (batchsize, length 1) uv (batchsize, length, 1) """ def __init__(self, sampling_rate, upsample_scale, harmonic_num=0, sine_amp=0.1, add_noise_std=0.003, voiced_threshod=0): super(SourceModuleHnNSF, self).__init__() self.sine_amp = sine_amp self.noise_std = add_noise_std # to produce sine waveforms self.l_sin_gen = SineGen(sampling_rate, harmonic_num, sine_amp, add_noise_std, voiced_threshod) # to merge source harmonics into a single excitation self.l_linear = torch.nn.Linear(harmonic_num + 1, 1) self.l_tanh = torch.nn.Tanh() def forward(self, x): """ Sine_source, noise_source = SourceModuleHnNSF(F0_sampled) F0_sampled (batchsize, length, 1) Sine_source (batchsize, length, 1) noise_source (batchsize, length 1) """ # source for harmonic branch with torch.no_grad(): sine_wavs, uv, _ = self.l_sin_gen(x.transpose(1, 2)) sine_wavs = sine_wavs.transpose(1, 2) uv = uv.transpose(1, 2) sine_merge = self.l_tanh(self.l_linear(sine_wavs)) # source for noise branch, in the same shape as uv noise = torch.randn_like(uv) * self.sine_amp / 3 return sine_merge, noise, uv class HiFTGenerator(nn.Module): """ HiFTNet Generator: Neural Source Filter + ISTFTNet https://arxiv.org/abs/2309.09493 """ def __init__( self, in_channels: int = 80, base_channels: int = 512, nb_harmonics: int = 8, sampling_rate: int = 22050, nsf_alpha: float = 0.1, nsf_sigma: float = 0.003, nsf_voiced_threshold: float = 10, upsample_rates: tp.List[int] = [8, 8], upsample_kernel_sizes: tp.List[int] = [16, 16], istft_params: tp.Dict[str, int] = {"n_fft": 16, "hop_len": 4}, resblock_kernel_sizes: tp.List[int] = [3, 7, 11], resblock_dilation_sizes: tp.List[tp.List[int]] = [[1, 3, 5], [1, 3, 5], [1, 3, 5]], source_resblock_kernel_sizes: tp.List[int] = [7, 11], source_resblock_dilation_sizes: tp.List[tp.List[int]] = [[1, 3, 5], [1, 3, 5]], lrelu_slope: float = 0.1, audio_limit: float = 0.99, f0_predictor: torch.nn.Module = None, ): super(HiFTGenerator, self).__init__() self.out_channels = 1 self.nb_harmonics = nb_harmonics self.sampling_rate = sampling_rate self.istft_params = istft_params self.lrelu_slope = lrelu_slope self.audio_limit = audio_limit self.num_kernels = len(resblock_kernel_sizes) self.num_upsamples = len(upsample_rates) self.m_source = SourceModuleHnNSF( sampling_rate=sampling_rate, upsample_scale=np.prod(upsample_rates) * istft_params["hop_len"], harmonic_num=nb_harmonics, sine_amp=nsf_alpha, add_noise_std=nsf_sigma, voiced_threshod=nsf_voiced_threshold) self.f0_upsamp = torch.nn.Upsample(scale_factor=np.prod(upsample_rates) * istft_params["hop_len"]) self.conv_pre = weight_norm( Conv1d(in_channels, base_channels, 7, 1, padding=3) ) # Up self.ups = nn.ModuleList() for i, (u, k) in enumerate(zip(upsample_rates, upsample_kernel_sizes)): self.ups.append( weight_norm( ConvTranspose1d( base_channels // (2**i), base_channels // (2**(i + 1)), k, u, padding=(k - u) // 2, ) ) ) # Down self.source_downs = nn.ModuleList() self.source_resblocks = nn.ModuleList() downsample_rates = [1] + upsample_rates[::-1][:-1] downsample_cum_rates = np.cumprod(downsample_rates) for i, (u, k, d) in enumerate(zip(downsample_cum_rates[::-1], source_resblock_kernel_sizes, source_resblock_dilation_sizes)): if u == 1: self.source_downs.append( Conv1d(istft_params["n_fft"] + 2, base_channels // (2 ** (i + 1)), 1, 1) ) else: self.source_downs.append( Conv1d(istft_params["n_fft"] + 2, base_channels // (2 ** (i + 1)), u * 2, u, padding=(u // 2)) ) self.source_resblocks.append( ResBlock(base_channels // (2 ** (i + 1)), k, d) ) self.resblocks = nn.ModuleList() for i in range(len(self.ups)): ch = base_channels // (2**(i + 1)) for j, (k, d) in enumerate(zip(resblock_kernel_sizes, resblock_dilation_sizes)): self.resblocks.append(ResBlock(ch, k, d)) self.conv_post = weight_norm(Conv1d(ch, istft_params["n_fft"] + 2, 7, 1, padding=3)) self.ups.apply(init_weights) self.conv_post.apply(init_weights) self.reflection_pad = nn.ReflectionPad1d((1, 0)) self.stft_window = torch.from_numpy(get_window("hann", istft_params["n_fft"], fftbins=True).astype(np.float32)) self.f0_predictor = f0_predictor def _f02source(self, f0: torch.Tensor) -> torch.Tensor: f0 = self.f0_upsamp(f0[:, None]).transpose(1, 2) # bs,n,t har_source, _, _ = self.m_source(f0) return har_source.transpose(1, 2) def _stft(self, x): spec = torch.stft( x, self.istft_params["n_fft"], self.istft_params["hop_len"], self.istft_params["n_fft"], window=self.stft_window.to(x.device), return_complex=True) spec = torch.view_as_real(spec) # [B, F, TT, 2] return spec[..., 0], spec[..., 1] def _istft(self, magnitude, phase): magnitude = torch.clip(magnitude, max=1e2) real = magnitude * torch.cos(phase) img = magnitude * torch.sin(phase) inverse_transform = torch.istft(torch.complex(real, img), self.istft_params["n_fft"], self.istft_params["hop_len"], self.istft_params["n_fft"], window=self.stft_window.to(magnitude.device)) return inverse_transform def forward(self, x: torch.Tensor, f0=None) -> torch.Tensor: if f0 is None: f0 = self.f0_predictor(x) s = self._f02source(f0) s_stft_real, s_stft_imag = self._stft(s.squeeze(1)) s_stft = torch.cat([s_stft_real, s_stft_imag], dim=1) x = self.conv_pre(x) for i in range(self.num_upsamples): x = F.leaky_relu(x, self.lrelu_slope) x = self.ups[i](x) if i == self.num_upsamples - 1: x = self.reflection_pad(x) # fusion si = self.source_downs[i](s_stft) si = self.source_resblocks[i](si) x = x + si xs = None for j in range(self.num_kernels): if xs is None: xs = self.resblocks[i * self.num_kernels + j](x) else: xs += self.resblocks[i * self.num_kernels + j](x) x = xs / self.num_kernels x = F.leaky_relu(x) x = self.conv_post(x) magnitude = torch.exp(x[:, :self.istft_params["n_fft"] // 2 + 1, :]) phase = torch.sin(x[:, self.istft_params["n_fft"] // 2 + 1:, :]) # actually, sin is redundancy x = self._istft(magnitude, phase) x = torch.clamp(x, -self.audio_limit, self.audio_limit) return x def remove_weight_norm(self): print('Removing weight norm...') for l in self.ups: remove_weight_norm(l) for l in self.resblocks: l.remove_weight_norm() remove_weight_norm(self.conv_pre) remove_weight_norm(self.conv_post) self.source_module.remove_weight_norm() for l in self.source_downs: remove_weight_norm(l) for l in self.source_resblocks: l.remove_weight_norm() @torch.inference_mode() def inference(self, mel: torch.Tensor, f0=None) -> torch.Tensor: return self.forward(x=mel, f0=f0)