+
+
+
+
++ +
+
+
+
+
++ +## Useage + +Download the pre-trained FACodec model from HuggingFace: [Pretrained FACodec checkpoint](https://huggingface.co/amphion/naturalspeech3_facodec) + +Install Amphion +```bash +git clone https://github.com/open-mmlab/Amphion.git +``` + +Few lines of code to use the pre-trained FACodec model +```python +from Amphion.models.codec.ns3_codec import FACodecEncoder, FACodecDecoder +from huggingface_hub import hf_hub_download + +fa_encoder = FACodecEncoder( + ngf=32, + up_ratios=[2, 4, 5, 5], + out_channels=256, +) + +fa_decoder = FACodecDecoder( + in_channels=256, + upsample_initial_channel=1024, + ngf=32, + up_ratios=[5, 5, 4, 2], + vq_num_q_c=2, + vq_num_q_p=1, + vq_num_q_r=3, + vq_dim=256, + codebook_dim=8, + codebook_size_prosody=10, + codebook_size_content=10, + codebook_size_residual=10, + use_gr_x_timbre=True, + use_gr_residual_f0=True, + use_gr_residual_phone=True, +) + +encoder_ckpt = hf_hub_download(repo_id="amphion/naturalspeech3_facodec", filename="ns3_facodec_encoder.bin") +decoder_ckpt = hf_hub_download(repo_id="amphion/naturalspeech3_facodec", filename="ns3_facodec_decoder.bin") + +fa_encoder.load_state_dict(torch.load(encoder_ckpt)) +fa_decoder.load_state_dict(torch.load(decoder_ckpt)) + +fa_encoder.eval() +fa_decoder.eval() + +``` + +Inference +```python +test_wav_path = "test.wav" +test_wav = librosa.load(test_wav_path, sr=16000)[0] +test_wav = torch.from_numpy(test_wav).float() +test_wav = test_wav.unsqueeze(0).unsqueeze(0) + +with torch.no_grad(): + + # encode + enc_out = fa_encoder(test_wav) + print(enc_out.shape) + + # quantize + vq_post_emb, vq_id, _, quantized, spk_embs = fa_decoder(enc_out, eval_vq=False, vq=True) + + # latent after quantization + print(vq_post_emb.shape) + + # codes + print("vq id shape:", vq_id.shape) + + # get prosody code + prosody_code = vq_id[:1] + print("prosody code shape:", prosody_code.shape) + + # get content code + cotent_code = vq_id[1:3] + print("content code shape:", cotent_code.shape) + + # get residual code (acoustic detail codes) + residual_code = vq_id[3:] + print("residual code shape:", residual_code.shape) + + # speaker embedding + print("speaker embedding shape:", spk_embs.shape) + + # decode (recommand) + recon_wav = fa_decoder.inference(vq_post_emb, spk_embs) + print(recon_wav.shape) + sf.write("recon.wav", recon_wav[0][0].cpu().numpy(), 16000) +``` + +FACodec can achieve zero-shot voice conversion with FACodecEncoderV2/FACodecDecoderV2 or FACodecRedecoder +```python +from Amphion.models.codec.ns3_codec import FACodecEncoderV2, FACodecDecoderV2 + +# Same parameters as FACodecEncoder/FACodecDecoder +fa_encoder_v2 = FACodecEncoderV2(...) +fa_decoder_v2 = FACodecDecoderV2(...) + +encoder_v2_ckpt = hf_hub_download(repo_id="amphion/naturalspeech3_facodec", filename="ns3_facodec_encoder_v2.bin") +decoder_v2_ckpt = hf_hub_download(repo_id="amphion/naturalspeech3_facodec", filename="ns3_facodec_decoder_v2.bin") + +fa_encoder_v2.load_state_dict(torch.load(encoder_v2_ckpt)) +fa_decoder_v2.load_state_dict(torch.load(decoder_v2_ckpt)) + +with torch.no_grad(): + enc_out_a = fa_encoder_v2(wav_a) + prosody_a = fa_encoder_v2.get_prosody_feature(wav_a) + enc_out_b = fa_encoder_v2(wav_b) + prosody_b = fa_encoder_v2.get_prosody_feature(wav_b) + + vq_post_emb_a, vq_id_a, _, quantized, spk_embs_a = fa_decoder_v2( + enc_out_a, prosody_a, eval_vq=False, vq=True + ) + vq_post_emb_b, vq_id_b, _, quantized, spk_embs_b = fa_decoder_v2( + enc_out_b, prosody_b, eval_vq=False, vq=True + ) + + vq_post_emb_a_to_b = fa_decoder_v2.vq2emb(vq_id_a, use_residual=False) + recon_wav_a_to_b = fa_decoder_v2.inference(vq_post_emb_a_to_b, spk_embs_b) +``` + +or + +```python +from Amphion.models.codec.ns3_codec import FACodecRedecoder + +fa_redecoder = FACodecRedecoder() + +redecoder_ckpt = hf_hub_download(repo_id="amphion/naturalspeech3_facodec", filename="ns3_facodec_redecoder.bin") + +fa_redecoder.load_state_dict(torch.load(redecoder_ckpt)) + +with torch.no_grad(): + enc_out_a = fa_encoder(wav_a) + enc_out_b = fa_encoder(wav_b) + + vq_post_emb_a, vq_id_a, _, quantized_a, spk_embs_a = fa_decoder(enc_out_a, eval_vq=False, vq=True) + vq_post_emb_b, vq_id_b, _, quantized_b, spk_embs_b = fa_decoder(enc_out_b, eval_vq=False, vq=True) + + # convert speaker + vq_post_emb_a_to_b = fa_redecoder.vq2emb(vq_id_a, spk_embs_b, use_residual=False) + recon_wav_a_to_b = fa_redecoder.inference(vq_post_emb_a_to_b, spk_embs_b) + + sf.write("recon_a_to_b.wav", recon_wav_a_to_b[0][0].cpu().numpy(), 16000) +``` + +## Q&A + +Q1: What audio sample rate does FACodec support? What is the hop size? How many codes will be generated for each frame? + +A1: FACodec supports 16KHz speech audio. The hop size is 200 samples, and (16000/200) * 6 (total number of codebooks) codes will be generated for each frame. + +Q2: Is it possible to train an autoregressive TTS model like VALL-E using FACodec? + +A2: Yes. In fact, the authors of NaturalSpeech 3 have already employ explore the autoregressive generative model for discrete token generation with FACodec. They use an autoregressive language model to generate prosody codes, followed by a non-autoregressive model to generate the remaining content and acoustic details codes. + +Q3: Is it possible to train a latent diffusion TTS model like NaturalSpeech2 using FACodec? + +A3: Yes. You can use the latent getted after quanzaition as the modelling target for the latent diffusion model. + +Q4: Can FACodec compress and reconstruct audio from other domains? Such as sound effects, music, etc. + +A4: Since FACodec is designed for speech, it may not be suitable for other audio domains. However, it is possible to use the FACodec model to compress and reconstruct audio from other domains, but the quality may not be as good as the original audio. + +Q5: Can FACodec be used for content feature for some other tasks like voice conversion? + +A5: I think the answer is yes. Researchers can use the content code of FACodec as the content feature for voice conversion. We hope to see more research in this direction. + +## Citations + +If you use our FACodec model, please cite the following paper: + +```bibtex +@article{ju2024naturalspeech, + title={NaturalSpeech 3: Zero-Shot Speech Synthesis with Factorized Codec and Diffusion Models}, + author={Ju, Zeqian and Wang, Yuancheng and Shen, Kai and Tan, Xu and Xin, Detai and Yang, Dongchao and Liu, Yanqing and Leng, Yichong and Song, Kaitao and Tang, Siliang and others}, + journal={arXiv preprint arXiv:2403.03100}, + year={2024} +} + +@article{zhang2023amphion, + title={Amphion: An Open-Source Audio, Music and Speech Generation Toolkit}, + author={Xueyao Zhang and Liumeng Xue and Yicheng Gu and Yuancheng Wang and Haorui He and Chaoren Wang and Xi Chen and Zihao Fang and Haopeng Chen and Junan Zhang and Tze Ying Tang and Lexiao Zou and Mingxuan Wang and Jun Han and Kai Chen and Haizhou Li and Zhizheng Wu}, + journal={arXiv}, + year={2024}, + volume={abs/2312.09911} +} +``` + diff --git a/models/codec/ns3_codec/__init__.py b/models/codec/ns3_codec/__init__.py new file mode 100644 index 0000000..1f0e4c1 --- /dev/null +++ b/models/codec/ns3_codec/__init__.py @@ -0,0 +1,6 @@ +# Copyright (c) 2023 Amphion. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +from .facodec import * diff --git a/models/codec/ns3_codec/alias_free_torch/__init__.py b/models/codec/ns3_codec/alias_free_torch/__init__.py new file mode 100644 index 0000000..b3bccdb --- /dev/null +++ b/models/codec/ns3_codec/alias_free_torch/__init__.py @@ -0,0 +1,5 @@ +# Adapted from https://github.com/junjun3518/alias-free-torch under the Apache License 2.0 + +from .filter import * +from .resample import * +from .act import * diff --git a/models/codec/ns3_codec/alias_free_torch/act.py b/models/codec/ns3_codec/alias_free_torch/act.py new file mode 100644 index 0000000..779d58d --- /dev/null +++ b/models/codec/ns3_codec/alias_free_torch/act.py @@ -0,0 +1,29 @@ +# Adapted from https://github.com/junjun3518/alias-free-torch under the Apache License 2.0 + +import torch.nn as nn +from .resample import UpSample1d, DownSample1d + + +class Activation1d(nn.Module): + def __init__( + self, + activation, + up_ratio: int = 2, + down_ratio: int = 2, + up_kernel_size: int = 12, + down_kernel_size: int = 12, + ): + super().__init__() + self.up_ratio = up_ratio + self.down_ratio = down_ratio + self.act = activation + self.upsample = UpSample1d(up_ratio, up_kernel_size) + self.downsample = DownSample1d(down_ratio, down_kernel_size) + + # x: [B,C,T] + def forward(self, x): + x = self.upsample(x) + x = self.act(x) + x = self.downsample(x) + + return x diff --git a/models/codec/ns3_codec/alias_free_torch/filter.py b/models/codec/ns3_codec/alias_free_torch/filter.py new file mode 100644 index 0000000..ece8e02 --- /dev/null +++ b/models/codec/ns3_codec/alias_free_torch/filter.py @@ -0,0 +1,96 @@ +# Adapted from https://github.com/junjun3518/alias-free-torch under the Apache License 2.0 + +import torch +import torch.nn as nn +import torch.nn.functional as F +import math + +if "sinc" in dir(torch): + sinc = torch.sinc +else: + # This code is adopted from adefossez's julius.core.sinc under the MIT License + # https://adefossez.github.io/julius/julius/core.html + def sinc(x: torch.Tensor): + """ + Implementation of sinc, i.e. sin(pi * x) / (pi * x) + __Warning__: Different to julius.sinc, the input is multiplied by `pi`! + """ + return torch.where( + x == 0, + torch.tensor(1.0, device=x.device, dtype=x.dtype), + torch.sin(math.pi * x) / math.pi / x, + ) + + +# This code is adopted from adefossez's julius.lowpass.LowPassFilters under the MIT License +# https://adefossez.github.io/julius/julius/lowpass.html +def kaiser_sinc_filter1d( + cutoff, half_width, kernel_size +): # return filter [1,1,kernel_size] + even = kernel_size % 2 == 0 + half_size = kernel_size // 2 + + # For kaiser window + delta_f = 4 * half_width + A = 2.285 * (half_size - 1) * math.pi * delta_f + 7.95 + if A > 50.0: + beta = 0.1102 * (A - 8.7) + elif A >= 21.0: + beta = 0.5842 * (A - 21) ** 0.4 + 0.07886 * (A - 21.0) + else: + beta = 0.0 + window = torch.kaiser_window(kernel_size, beta=beta, periodic=False) + + # ratio = 0.5/cutoff -> 2 * cutoff = 1 / ratio + if even: + time = torch.arange(-half_size, half_size) + 0.5 + else: + time = torch.arange(kernel_size) - half_size + if cutoff == 0: + filter_ = torch.zeros_like(time) + else: + filter_ = 2 * cutoff * window * sinc(2 * cutoff * time) + # Normalize filter to have sum = 1, otherwise we will have a small leakage + # of the constant component in the input signal. + filter_ /= filter_.sum() + filter = filter_.view(1, 1, kernel_size) + + return filter + + +class LowPassFilter1d(nn.Module): + def __init__( + self, + cutoff=0.5, + half_width=0.6, + stride: int = 1, + padding: bool = True, + padding_mode: str = "replicate", + kernel_size: int = 12, + ): + # kernel_size should be even number for stylegan3 setup, + # in this implementation, odd number is also possible. + super().__init__() + if cutoff < -0.0: + raise ValueError("Minimum cutoff must be larger than zero.") + if cutoff > 0.5: + raise ValueError("A cutoff above 0.5 does not make sense.") + self.kernel_size = kernel_size + self.even = kernel_size % 2 == 0 + self.pad_left = kernel_size // 2 - int(self.even) + self.pad_right = kernel_size // 2 + self.stride = stride + self.padding = padding + self.padding_mode = padding_mode + filter = kaiser_sinc_filter1d(cutoff, half_width, kernel_size) + self.register_buffer("filter", filter) + + # input [B, C, T] + def forward(self, x): + _, C, _ = x.shape + + if self.padding: + x = F.pad(x, (self.pad_left, self.pad_right), mode=self.padding_mode) + out = F.conv1d(x, self.filter.expand(C, -1, -1), stride=self.stride, groups=C) + + return out diff --git a/models/codec/ns3_codec/alias_free_torch/resample.py b/models/codec/ns3_codec/alias_free_torch/resample.py new file mode 100644 index 0000000..ee993b1 --- /dev/null +++ b/models/codec/ns3_codec/alias_free_torch/resample.py @@ -0,0 +1,57 @@ +# Adapted from https://github.com/junjun3518/alias-free-torch under the Apache License 2.0 + +import torch.nn as nn +from torch.nn import functional as F +from .filter import LowPassFilter1d +from .filter import kaiser_sinc_filter1d + + +class UpSample1d(nn.Module): + def __init__(self, ratio=2, kernel_size=None): + super().__init__() + self.ratio = ratio + self.kernel_size = ( + int(6 * ratio // 2) * 2 if kernel_size is None else kernel_size + ) + self.stride = ratio + self.pad = self.kernel_size // ratio - 1 + self.pad_left = self.pad * self.stride + (self.kernel_size - self.stride) // 2 + self.pad_right = ( + self.pad * self.stride + (self.kernel_size - self.stride + 1) // 2 + ) + filter = kaiser_sinc_filter1d( + cutoff=0.5 / ratio, half_width=0.6 / ratio, kernel_size=self.kernel_size + ) + self.register_buffer("filter", filter) + + # x: [B, C, T] + def forward(self, x): + _, C, _ = x.shape + + x = F.pad(x, (self.pad, self.pad), mode="replicate") + x = self.ratio * F.conv_transpose1d( + x, self.filter.expand(C, -1, -1), stride=self.stride, groups=C + ) + x = x[..., self.pad_left : -self.pad_right] + + return x + + +class DownSample1d(nn.Module): + def __init__(self, ratio=2, kernel_size=None): + super().__init__() + self.ratio = ratio + self.kernel_size = ( + int(6 * ratio // 2) * 2 if kernel_size is None else kernel_size + ) + self.lowpass = LowPassFilter1d( + cutoff=0.5 / ratio, + half_width=0.6 / ratio, + stride=ratio, + kernel_size=self.kernel_size, + ) + + def forward(self, x): + xx = self.lowpass(x) + + return xx diff --git a/models/codec/ns3_codec/facodec.py b/models/codec/ns3_codec/facodec.py new file mode 100644 index 0000000..87f661b --- /dev/null +++ b/models/codec/ns3_codec/facodec.py @@ -0,0 +1,1222 @@ +# Copyright (c) 2023 Amphion. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +import numpy as np +import torch +from torch import nn, sin, pow +from torch.nn import Parameter +import torch.nn.functional as F +from torch.nn.utils import weight_norm +from .alias_free_torch import * +from .quantize import * +from einops import rearrange +from einops.layers.torch import Rearrange +from .transformer import TransformerEncoder +from .gradient_reversal import GradientReversal +from .melspec import MelSpectrogram + + +def init_weights(m): + if isinstance(m, nn.Conv1d): + nn.init.trunc_normal_(m.weight, std=0.02) + nn.init.constant_(m.bias, 0) + + +def WNConv1d(*args, **kwargs): + return weight_norm(nn.Conv1d(*args, **kwargs)) + + +def WNConvTranspose1d(*args, **kwargs): + return weight_norm(nn.ConvTranspose1d(*args, **kwargs)) + + +class CNNLSTM(nn.Module): + def __init__(self, indim, outdim, head, global_pred=False): + super().__init__() + self.global_pred = global_pred + self.model = nn.Sequential( + ResidualUnit(indim, dilation=1), + ResidualUnit(indim, dilation=2), + ResidualUnit(indim, dilation=3), + Activation1d(activation=SnakeBeta(indim, alpha_logscale=True)), + Rearrange("b c t -> b t c"), + ) + self.heads = nn.ModuleList([nn.Linear(indim, outdim) for i in range(head)]) + + def forward(self, x): + # x: [B, C, T] + x = self.model(x) + if self.global_pred: + x = torch.mean(x, dim=1, keepdim=False) + outs = [head(x) for head in self.heads] + return outs + + +class SnakeBeta(nn.Module): + """ + A modified Snake function which uses separate parameters for the magnitude of the periodic components + Shape: + - Input: (B, C, T) + - Output: (B, C, T), same shape as the input + Parameters: + - alpha - trainable parameter that controls frequency + - beta - trainable parameter that controls magnitude + References: + - This activation function is a modified version based on this paper by Liu Ziyin, Tilman Hartwig, Masahito Ueda: + https://arxiv.org/abs/2006.08195 + Examples: + >>> a1 = snakebeta(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 that controls frequency + - beta - trainable parameter that controls magnitude + alpha is initialized to 1 by default, higher values = higher-frequency. + beta is initialized to 1 by default, higher values = higher-magnitude. + alpha will be trained along with the rest of your model. + """ + super(SnakeBeta, 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) + self.beta = Parameter(torch.zeros(in_features) * alpha) + else: # linear scale alphas initialized to ones + self.alpha = Parameter(torch.ones(in_features) * alpha) + self.beta = Parameter(torch.ones(in_features) * alpha) + + self.alpha.requires_grad = alpha_trainable + self.beta.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. + SnakeBeta := x + 1/b * sin^2 (xa) + """ + alpha = self.alpha.unsqueeze(0).unsqueeze(-1) # line up with x to [B, C, T] + beta = self.beta.unsqueeze(0).unsqueeze(-1) + if self.alpha_logscale: + alpha = torch.exp(alpha) + beta = torch.exp(beta) + x = x + (1.0 / (beta + self.no_div_by_zero)) * pow(sin(x * alpha), 2) + + return x + + +class ResidualUnit(nn.Module): + def __init__(self, dim: int = 16, dilation: int = 1): + super().__init__() + pad = ((7 - 1) * dilation) // 2 + self.block = nn.Sequential( + Activation1d(activation=SnakeBeta(dim, alpha_logscale=True)), + WNConv1d(dim, dim, kernel_size=7, dilation=dilation, padding=pad), + Activation1d(activation=SnakeBeta(dim, alpha_logscale=True)), + WNConv1d(dim, dim, kernel_size=1), + ) + + def forward(self, x): + return x + self.block(x) + + +class EncoderBlock(nn.Module): + def __init__(self, dim: int = 16, stride: int = 1): + super().__init__() + self.block = nn.Sequential( + ResidualUnit(dim // 2, dilation=1), + ResidualUnit(dim // 2, dilation=3), + ResidualUnit(dim // 2, dilation=9), + Activation1d(activation=SnakeBeta(dim // 2, alpha_logscale=True)), + WNConv1d( + dim // 2, + dim, + kernel_size=2 * stride, + stride=stride, + padding=stride // 2 + stride % 2, + ), + ) + + def forward(self, x): + return self.block(x) + + +class FACodecEncoder(nn.Module): + def __init__( + self, + ngf=32, + up_ratios=(2, 4, 5, 5), + out_channels=1024, + ): + super().__init__() + self.hop_length = np.prod(up_ratios) + self.up_ratios = up_ratios + + # Create first convolution + d_model = ngf + self.block = [WNConv1d(1, d_model, kernel_size=7, padding=3)] + + # Create EncoderBlocks that double channels as they downsample by `stride` + for stride in up_ratios: + d_model *= 2 + self.block += [EncoderBlock(d_model, stride=stride)] + + # Create last convolution + self.block += [ + Activation1d(activation=SnakeBeta(d_model, alpha_logscale=True)), + WNConv1d(d_model, out_channels, kernel_size=3, padding=1), + ] + + # Wrap black into nn.Sequential + self.block = nn.Sequential(*self.block) + self.enc_dim = d_model + + self.reset_parameters() + + def forward(self, x): + out = self.block(x) + return out + + def inference(self, x): + return self.block(x) + + def remove_weight_norm(self): + """Remove weight normalization module from all of the layers.""" + + def _remove_weight_norm(m): + try: + torch.nn.utils.remove_weight_norm(m) + except ValueError: # this module didn't have weight norm + return + + self.apply(_remove_weight_norm) + + def apply_weight_norm(self): + """Apply weight normalization module from all of the layers.""" + + def _apply_weight_norm(m): + if isinstance(m, nn.Conv1d): + torch.nn.utils.weight_norm(m) + + self.apply(_apply_weight_norm) + + def reset_parameters(self): + self.apply(init_weights) + + +class DecoderBlock(nn.Module): + def __init__(self, input_dim: int = 16, output_dim: int = 8, stride: int = 1): + super().__init__() + self.block = nn.Sequential( + Activation1d(activation=SnakeBeta(input_dim, alpha_logscale=True)), + WNConvTranspose1d( + input_dim, + output_dim, + kernel_size=2 * stride, + stride=stride, + padding=stride // 2 + stride % 2, + output_padding=stride % 2, + ), + ResidualUnit(output_dim, dilation=1), + ResidualUnit(output_dim, dilation=3), + ResidualUnit(output_dim, dilation=9), + ) + + def forward(self, x): + return self.block(x) + + +class FACodecDecoder(nn.Module): + def __init__( + self, + in_channels=256, + upsample_initial_channel=1536, + ngf=32, + up_ratios=(5, 5, 4, 2), + vq_num_q_c=2, + vq_num_q_p=1, + vq_num_q_r=3, + vq_dim=1024, + vq_commit_weight=0.005, + vq_weight_init=False, + vq_full_commit_loss=False, + codebook_dim=8, + codebook_size_prosody=10, # true codebook size is equal to 2^codebook_size + codebook_size_content=10, + codebook_size_residual=10, + quantizer_dropout=0.0, + dropout_type="linear", + use_gr_content_f0=False, + use_gr_prosody_phone=False, + use_gr_residual_f0=False, + use_gr_residual_phone=False, + use_gr_x_timbre=False, + use_random_mask_residual=True, + prob_random_mask_residual=0.75, + ): + super().__init__() + self.hop_length = np.prod(up_ratios) + self.ngf = ngf + self.up_ratios = up_ratios + + self.use_random_mask_residual = use_random_mask_residual + self.prob_random_mask_residual = prob_random_mask_residual + + self.vq_num_q_p = vq_num_q_p + self.vq_num_q_c = vq_num_q_c + self.vq_num_q_r = vq_num_q_r + + self.codebook_size_prosody = codebook_size_prosody + self.codebook_size_content = codebook_size_content + self.codebook_size_residual = codebook_size_residual + + quantizer_class = ResidualVQ + + self.quantizer = nn.ModuleList() + + # prosody + quantizer = quantizer_class( + num_quantizers=vq_num_q_p, + dim=vq_dim, + codebook_size=codebook_size_prosody, + codebook_dim=codebook_dim, + threshold_ema_dead_code=2, + commitment=vq_commit_weight, + weight_init=vq_weight_init, + full_commit_loss=vq_full_commit_loss, + quantizer_dropout=quantizer_dropout, + dropout_type=dropout_type, + ) + self.quantizer.append(quantizer) + + # phone + quantizer = quantizer_class( + num_quantizers=vq_num_q_c, + dim=vq_dim, + codebook_size=codebook_size_content, + codebook_dim=codebook_dim, + threshold_ema_dead_code=2, + commitment=vq_commit_weight, + weight_init=vq_weight_init, + full_commit_loss=vq_full_commit_loss, + quantizer_dropout=quantizer_dropout, + dropout_type=dropout_type, + ) + self.quantizer.append(quantizer) + + # residual + if self.vq_num_q_r > 0: + quantizer = quantizer_class( + num_quantizers=vq_num_q_r, + dim=vq_dim, + codebook_size=codebook_size_residual, + codebook_dim=codebook_dim, + threshold_ema_dead_code=2, + commitment=vq_commit_weight, + weight_init=vq_weight_init, + full_commit_loss=vq_full_commit_loss, + quantizer_dropout=quantizer_dropout, + dropout_type=dropout_type, + ) + self.quantizer.append(quantizer) + + # Add first conv layer + channels = upsample_initial_channel + layers = [WNConv1d(in_channels, channels, kernel_size=7, padding=3)] + + # Add upsampling + MRF blocks + for i, stride in enumerate(up_ratios): + input_dim = channels // 2**i + output_dim = channels // 2 ** (i + 1) + layers += [DecoderBlock(input_dim, output_dim, stride)] + + # Add final conv layer + layers += [ + Activation1d(activation=SnakeBeta(output_dim, alpha_logscale=True)), + WNConv1d(output_dim, 1, kernel_size=7, padding=3), + nn.Tanh(), + ] + + self.model = nn.Sequential(*layers) + + self.timbre_encoder = TransformerEncoder( + enc_emb_tokens=None, + encoder_layer=4, + encoder_hidden=256, + encoder_head=4, + conv_filter_size=1024, + conv_kernel_size=5, + encoder_dropout=0.1, + use_cln=False, + ) + + self.timbre_linear = nn.Linear(in_channels, in_channels * 2) + self.timbre_linear.bias.data[:in_channels] = 1 + self.timbre_linear.bias.data[in_channels:] = 0 + self.timbre_norm = nn.LayerNorm(in_channels, elementwise_affine=False) + + self.f0_predictor = CNNLSTM(in_channels, 1, 2) + self.phone_predictor = CNNLSTM(in_channels, 5003, 1) + + self.use_gr_content_f0 = use_gr_content_f0 + self.use_gr_prosody_phone = use_gr_prosody_phone + self.use_gr_residual_f0 = use_gr_residual_f0 + self.use_gr_residual_phone = use_gr_residual_phone + self.use_gr_x_timbre = use_gr_x_timbre + + if self.vq_num_q_r > 0 and self.use_gr_residual_f0: + self.res_f0_predictor = nn.Sequential( + GradientReversal(alpha=1.0), CNNLSTM(in_channels, 1, 2) + ) + + if self.vq_num_q_r > 0 and self.use_gr_residual_phone > 0: + self.res_phone_predictor = nn.Sequential( + GradientReversal(alpha=1.0), CNNLSTM(in_channels, 5003, 1) + ) + + if self.use_gr_content_f0: + self.content_f0_predictor = nn.Sequential( + GradientReversal(alpha=1.0), CNNLSTM(in_channels, 1, 2) + ) + + if self.use_gr_prosody_phone: + self.prosody_phone_predictor = nn.Sequential( + GradientReversal(alpha=1.0), CNNLSTM(in_channels, 5003, 1) + ) + + if self.use_gr_x_timbre: + self.x_timbre_predictor = nn.Sequential( + GradientReversal(alpha=1), + CNNLSTM(in_channels, 245200, 1, global_pred=True), + ) + + self.reset_parameters() + + def quantize(self, x, n_quantizers=None): + outs, qs, commit_loss, quantized_buf = 0, [], [], [] + + # prosody + f0_input = x # (B, d, T) + f0_quantizer = self.quantizer[0] + out, q, commit, quantized = f0_quantizer(f0_input, n_quantizers=n_quantizers) + outs += out + qs.append(q) + quantized_buf.append(quantized.sum(0)) + commit_loss.append(commit) + + # phone + phone_input = x + phone_quantizer = self.quantizer[1] + out, q, commit, quantized = phone_quantizer( + phone_input, n_quantizers=n_quantizers + ) + outs += out + qs.append(q) + quantized_buf.append(quantized.sum(0)) + commit_loss.append(commit) + + # residual + if self.vq_num_q_r > 0: + residual_quantizer = self.quantizer[2] + residual_input = x - (quantized_buf[0] + quantized_buf[1]).detach() + out, q, commit, quantized = residual_quantizer( + residual_input, n_quantizers=n_quantizers + ) + outs += out + qs.append(q) + quantized_buf.append(quantized.sum(0)) # [L, B, C, T] -> [B, C, T] + commit_loss.append(commit) + + qs = torch.cat(qs, dim=0) + commit_loss = torch.cat(commit_loss, dim=0) + return outs, qs, commit_loss, quantized_buf + + def forward( + self, + x, + vq=True, + get_vq=False, + eval_vq=True, + speaker_embedding=None, + n_quantizers=None, + quantized=None, + ): + if get_vq: + return self.quantizer.get_emb() + if vq is True: + if eval_vq: + self.quantizer.eval() + x_timbre = x + outs, qs, commit_loss, quantized_buf = self.quantize( + x, n_quantizers=n_quantizers + ) + + x_timbre = x_timbre.transpose(1, 2) + x_timbre = self.timbre_encoder(x_timbre, None, None) + x_timbre = x_timbre.transpose(1, 2) + spk_embs = torch.mean(x_timbre, dim=2) + return outs, qs, commit_loss, quantized_buf, spk_embs + + out = {} + + layer_0 = quantized[0] + f0, uv = self.f0_predictor(layer_0) + f0 = rearrange(f0, "... 1 -> ...") + uv = rearrange(uv, "... 1 -> ...") + + layer_1 = quantized[1] + (phone,) = self.phone_predictor(layer_1) + + out = {"f0": f0, "uv": uv, "phone": phone} + + if self.use_gr_prosody_phone: + (prosody_phone,) = self.prosody_phone_predictor(layer_0) + out["prosody_phone"] = prosody_phone + + if self.use_gr_content_f0: + content_f0, content_uv = self.content_f0_predictor(layer_1) + content_f0 = rearrange(content_f0, "... 1 -> ...") + content_uv = rearrange(content_uv, "... 1 -> ...") + out["content_f0"] = content_f0 + out["content_uv"] = content_uv + + if self.vq_num_q_r > 0: + layer_2 = quantized[2] + + if self.use_gr_residual_f0: + res_f0, res_uv = self.res_f0_predictor(layer_2) + res_f0 = rearrange(res_f0, "... 1 -> ...") + res_uv = rearrange(res_uv, "... 1 -> ...") + out["res_f0"] = res_f0 + out["res_uv"] = res_uv + + if self.use_gr_residual_phone: + (res_phone,) = self.res_phone_predictor(layer_2) + out["res_phone"] = res_phone + + style = self.timbre_linear(speaker_embedding).unsqueeze(2) # (B, 2d, 1) + gamma, beta = style.chunk(2, 1) # (B, d, 1) + if self.vq_num_q_r > 0: + if self.use_random_mask_residual: + bsz = quantized[2].shape[0] + res_mask = np.random.choice( + [0, 1], + size=bsz, + p=[ + self.prob_random_mask_residual, + 1 - self.prob_random_mask_residual, + ], + ) + res_mask = ( + torch.from_numpy(res_mask).unsqueeze(1).unsqueeze(1) + ) # (B, 1, 1) + res_mask = res_mask.to( + device=quantized[2].device, dtype=quantized[2].dtype + ) + x = ( + quantized[0].detach() + + quantized[1].detach() + + quantized[2] * res_mask + ) + # x = quantized_perturbe[0].detach() + quantized[1].detach() + quantized[2] * res_mask + else: + x = quantized[0].detach() + quantized[1].detach() + quantized[2] + # x = quantized_perturbe[0].detach() + quantized[1].detach() + quantized[2] + else: + x = quantized[0].detach() + quantized[1].detach() + # x = quantized_perturbe[0].detach() + quantized[1].detach() + + if self.use_gr_x_timbre: + (x_timbre,) = self.x_timbre_predictor(x) + out["x_timbre"] = x_timbre + + x = x.transpose(1, 2) + x = self.timbre_norm(x) + x = x.transpose(1, 2) + x = x * gamma + beta + + x = self.model(x) + out["audio"] = x + + return out + + def vq2emb(self, vq, use_residual_code=True): + # vq: [num_quantizer, B, T] + self.quantizer = self.quantizer.eval() + out = 0 + out += self.quantizer[0].vq2emb(vq[0 : self.vq_num_q_p]) + out += self.quantizer[1].vq2emb( + vq[self.vq_num_q_p : self.vq_num_q_p + self.vq_num_q_c] + ) + if self.vq_num_q_r > 0 and use_residual_code: + out += self.quantizer[2].vq2emb(vq[self.vq_num_q_p + self.vq_num_q_c :]) + return out + + def inference(self, x, speaker_embedding): + style = self.timbre_linear(speaker_embedding).unsqueeze(2) # (B, 2d, 1) + gamma, beta = style.chunk(2, 1) # (B, d, 1) + x = x.transpose(1, 2) + x = self.timbre_norm(x) + x = x.transpose(1, 2) + x = x * gamma + beta + x = self.model(x) + return x + + def remove_weight_norm(self): + """Remove weight normalization module from all of the layers.""" + + def _remove_weight_norm(m): + try: + torch.nn.utils.remove_weight_norm(m) + except ValueError: # this module didn't have weight norm + return + + self.apply(_remove_weight_norm) + + def apply_weight_norm(self): + """Apply weight normalization module from all of the layers.""" + + def _apply_weight_norm(m): + if isinstance(m, nn.Conv1d) or isinstance(m, nn.ConvTranspose1d): + torch.nn.utils.weight_norm(m) + + self.apply(_apply_weight_norm) + + def reset_parameters(self): + self.apply(init_weights) + + +class FACodecRedecoder(nn.Module): + def __init__( + self, + in_channels=256, + upsample_initial_channel=1280, + up_ratios=(5, 5, 4, 2), + vq_num_q_c=2, + vq_num_q_p=1, + vq_num_q_r=3, + vq_dim=256, + codebook_size_prosody=10, + codebook_size_content=10, + codebook_size_residual=10, + ): + super().__init__() + self.hop_length = np.prod(up_ratios) + self.up_ratios = up_ratios + + self.vq_num_q_p = vq_num_q_p + self.vq_num_q_c = vq_num_q_c + self.vq_num_q_r = vq_num_q_r + + self.vq_dim = vq_dim + + self.codebook_size_prosody = codebook_size_prosody + self.codebook_size_content = codebook_size_content + self.codebook_size_residual = codebook_size_residual + + self.prosody_embs = nn.ModuleList() + for i in range(self.vq_num_q_p): + emb_tokens = nn.Embedding( + num_embeddings=2**self.codebook_size_prosody, + embedding_dim=self.vq_dim, + ) + emb_tokens.weight.data.normal_(mean=0.0, std=1e-5) + self.prosody_embs.append(emb_tokens) + self.content_embs = nn.ModuleList() + for i in range(self.vq_num_q_c): + emb_tokens = nn.Embedding( + num_embeddings=2**self.codebook_size_content, + embedding_dim=self.vq_dim, + ) + emb_tokens.weight.data.normal_(mean=0.0, std=1e-5) + self.content_embs.append(emb_tokens) + self.residual_embs = nn.ModuleList() + for i in range(self.vq_num_q_r): + emb_tokens = nn.Embedding( + num_embeddings=2**self.codebook_size_residual, + embedding_dim=self.vq_dim, + ) + emb_tokens.weight.data.normal_(mean=0.0, std=1e-5) + self.residual_embs.append(emb_tokens) + + # Add first conv layer + channels = upsample_initial_channel + layers = [WNConv1d(in_channels, channels, kernel_size=7, padding=3)] + + # Add upsampling + MRF blocks + for i, stride in enumerate(up_ratios): + input_dim = channels // 2**i + output_dim = channels // 2 ** (i + 1) + layers += [DecoderBlock(input_dim, output_dim, stride)] + + # Add final conv layer + layers += [ + Activation1d(activation=SnakeBeta(output_dim, alpha_logscale=True)), + WNConv1d(output_dim, 1, kernel_size=7, padding=3), + nn.Tanh(), + ] + + self.model = nn.Sequential(*layers) + + self.timbre_linear = nn.Linear(in_channels, in_channels * 2) + self.timbre_linear.bias.data[:in_channels] = 1 + self.timbre_linear.bias.data[in_channels:] = 0 + self.timbre_norm = nn.LayerNorm(in_channels, elementwise_affine=False) + + self.timbre_cond_prosody_enc = TransformerEncoder( + enc_emb_tokens=None, + encoder_layer=4, + encoder_hidden=256, + encoder_head=4, + conv_filter_size=1024, + conv_kernel_size=5, + encoder_dropout=0.1, + use_cln=True, + cfg=None, + ) + + def forward( + self, + vq, + speaker_embedding, + use_residual_code=False, + ): + + x = 0 + + x_p = 0 + for i in range(self.vq_num_q_p): + x_p = x_p + self.prosody_embs[i](vq[i]) # (B, T, d) + spk_cond = speaker_embedding.unsqueeze(1).expand(-1, x_p.shape[1], -1) + x_p = self.timbre_cond_prosody_enc( + x_p, key_padding_mask=None, condition=spk_cond + ) + x = x + x_p + + x_c = 0 + for i in range(self.vq_num_q_c): + x_c = x_c + self.content_embs[i](vq[self.vq_num_q_p + i]) + + x = x + x_c + + if use_residual_code: + + x_r = 0 + for i in range(self.vq_num_q_r): + x_r = x_r + self.residual_embs[i]( + vq[self.vq_num_q_p + self.vq_num_q_c + i] + ) + x = x + x_r + + style = self.timbre_linear(speaker_embedding).unsqueeze(2) # (B, 2d, 1) + gamma, beta = style.chunk(2, 1) # (B, d, 1) + x = x.transpose(1, 2) + x = self.timbre_norm(x) + x = x.transpose(1, 2) + x = x * gamma + beta + x = self.model(x) + + return x + + def vq2emb(self, vq, speaker_embedding, use_residual=True): + + out = 0 + + x_t = 0 + for i in range(self.vq_num_q_p): + x_t += self.prosody_embs[i](vq[i]) # (B, T, d) + spk_cond = speaker_embedding.unsqueeze(1).expand(-1, x_t.shape[1], -1) + x_t = self.timbre_cond_prosody_enc( + x_t, key_padding_mask=None, condition=spk_cond + ) + + # prosody + out += x_t + + # content + for i in range(self.vq_num_q_c): + out += self.content_embs[i](vq[self.vq_num_q_p + i]) + + # residual + if use_residual: + for i in range(self.vq_num_q_r): + out += self.residual_embs[i](vq[self.vq_num_q_p + self.vq_num_q_c + i]) + + out = out.transpose(1, 2) # (B, T, d) -> (B, d, T) + return out + + def inference(self, x, speaker_embedding): + style = self.timbre_linear(speaker_embedding).unsqueeze(2) # (B, 2d, 1) + gamma, beta = style.chunk(2, 1) # (B, d, 1) + x = x.transpose(1, 2) + x = self.timbre_norm(x) + x = x.transpose(1, 2) + x = x * gamma + beta + x = self.model(x) + return x + + +class FACodecEncoderV2(nn.Module): + def __init__( + self, + ngf=32, + up_ratios=(2, 4, 5, 5), + out_channels=1024, + ): + super().__init__() + self.hop_length = np.prod(up_ratios) + self.up_ratios = up_ratios + + # Create first convolution + d_model = ngf + self.block = [WNConv1d(1, d_model, kernel_size=7, padding=3)] + + # Create EncoderBlocks that double channels as they downsample by `stride` + for stride in up_ratios: + d_model *= 2 + self.block += [EncoderBlock(d_model, stride=stride)] + + # Create last convolution + self.block += [ + Activation1d(activation=SnakeBeta(d_model, alpha_logscale=True)), + WNConv1d(d_model, out_channels, kernel_size=3, padding=1), + ] + + # Wrap black into nn.Sequential + self.block = nn.Sequential(*self.block) + self.enc_dim = d_model + + self.mel_transform = MelSpectrogram( + n_fft=1024, + num_mels=80, + sampling_rate=16000, + hop_size=200, + win_size=800, + fmin=0, + fmax=8000, + ) + + self.reset_parameters() + + def forward(self, x): + out = self.block(x) + return out + + def inference(self, x): + return self.block(x) + + def get_prosody_feature(self, x): + return self.mel_transform(x.squeeze(1))[:, :20, :] + + def remove_weight_norm(self): + """Remove weight normalization module from all of the layers.""" + + def _remove_weight_norm(m): + try: + torch.nn.utils.remove_weight_norm(m) + except ValueError: # this module didn't have weight norm + return + + self.apply(_remove_weight_norm) + + def apply_weight_norm(self): + """Apply weight normalization module from all of the layers.""" + + def _apply_weight_norm(m): + if isinstance(m, nn.Conv1d): + torch.nn.utils.weight_norm(m) + + self.apply(_apply_weight_norm) + + def reset_parameters(self): + self.apply(init_weights) + + +class FACodecDecoderV2(nn.Module): + def __init__( + self, + in_channels=256, + upsample_initial_channel=1536, + ngf=32, + up_ratios=(5, 5, 4, 2), + vq_num_q_c=2, + vq_num_q_p=1, + vq_num_q_r=3, + vq_dim=1024, + vq_commit_weight=0.005, + vq_weight_init=False, + vq_full_commit_loss=False, + codebook_dim=8, + codebook_size_prosody=10, # true codebook size is equal to 2^codebook_size + codebook_size_content=10, + codebook_size_residual=10, + quantizer_dropout=0.0, + dropout_type="linear", + use_gr_content_f0=False, + use_gr_prosody_phone=False, + use_gr_residual_f0=False, + use_gr_residual_phone=False, + use_gr_x_timbre=False, + use_random_mask_residual=True, + prob_random_mask_residual=0.75, + ): + super().__init__() + self.hop_length = np.prod(up_ratios) + self.ngf = ngf + self.up_ratios = up_ratios + + self.use_random_mask_residual = use_random_mask_residual + self.prob_random_mask_residual = prob_random_mask_residual + + self.vq_num_q_p = vq_num_q_p + self.vq_num_q_c = vq_num_q_c + self.vq_num_q_r = vq_num_q_r + + self.codebook_size_prosody = codebook_size_prosody + self.codebook_size_content = codebook_size_content + self.codebook_size_residual = codebook_size_residual + + quantizer_class = ResidualVQ + + self.quantizer = nn.ModuleList() + + # prosody + quantizer = quantizer_class( + num_quantizers=vq_num_q_p, + dim=vq_dim, + codebook_size=codebook_size_prosody, + codebook_dim=codebook_dim, + threshold_ema_dead_code=2, + commitment=vq_commit_weight, + weight_init=vq_weight_init, + full_commit_loss=vq_full_commit_loss, + quantizer_dropout=quantizer_dropout, + dropout_type=dropout_type, + ) + self.quantizer.append(quantizer) + + # phone + quantizer = quantizer_class( + num_quantizers=vq_num_q_c, + dim=vq_dim, + codebook_size=codebook_size_content, + codebook_dim=codebook_dim, + threshold_ema_dead_code=2, + commitment=vq_commit_weight, + weight_init=vq_weight_init, + full_commit_loss=vq_full_commit_loss, + quantizer_dropout=quantizer_dropout, + dropout_type=dropout_type, + ) + self.quantizer.append(quantizer) + + # residual + if self.vq_num_q_r > 0: + quantizer = quantizer_class( + num_quantizers=vq_num_q_r, + dim=vq_dim, + codebook_size=codebook_size_residual, + codebook_dim=codebook_dim, + threshold_ema_dead_code=2, + commitment=vq_commit_weight, + weight_init=vq_weight_init, + full_commit_loss=vq_full_commit_loss, + quantizer_dropout=quantizer_dropout, + dropout_type=dropout_type, + ) + self.quantizer.append(quantizer) + + # Add first conv layer + channels = upsample_initial_channel + layers = [WNConv1d(in_channels, channels, kernel_size=7, padding=3)] + + # Add upsampling + MRF blocks + for i, stride in enumerate(up_ratios): + input_dim = channels // 2**i + output_dim = channels // 2 ** (i + 1) + layers += [DecoderBlock(input_dim, output_dim, stride)] + + # Add final conv layer + layers += [ + Activation1d(activation=SnakeBeta(output_dim, alpha_logscale=True)), + WNConv1d(output_dim, 1, kernel_size=7, padding=3), + nn.Tanh(), + ] + + self.model = nn.Sequential(*layers) + + self.timbre_encoder = TransformerEncoder( + enc_emb_tokens=None, + encoder_layer=4, + encoder_hidden=256, + encoder_head=4, + conv_filter_size=1024, + conv_kernel_size=5, + encoder_dropout=0.1, + use_cln=False, + ) + + self.timbre_linear = nn.Linear(in_channels, in_channels * 2) + self.timbre_linear.bias.data[:in_channels] = 1 + self.timbre_linear.bias.data[in_channels:] = 0 + self.timbre_norm = nn.LayerNorm(in_channels, elementwise_affine=False) + + self.f0_predictor = CNNLSTM(in_channels, 1, 2) + self.phone_predictor = CNNLSTM(in_channels, 5003, 1) + + self.use_gr_content_f0 = use_gr_content_f0 + self.use_gr_prosody_phone = use_gr_prosody_phone + self.use_gr_residual_f0 = use_gr_residual_f0 + self.use_gr_residual_phone = use_gr_residual_phone + self.use_gr_x_timbre = use_gr_x_timbre + + if self.vq_num_q_r > 0 and self.use_gr_residual_f0: + self.res_f0_predictor = nn.Sequential( + GradientReversal(alpha=1.0), CNNLSTM(in_channels, 1, 2) + ) + + if self.vq_num_q_r > 0 and self.use_gr_residual_phone > 0: + self.res_phone_predictor = nn.Sequential( + GradientReversal(alpha=1.0), CNNLSTM(in_channels, 5003, 1) + ) + + if self.use_gr_content_f0: + self.content_f0_predictor = nn.Sequential( + GradientReversal(alpha=1.0), CNNLSTM(in_channels, 1, 2) + ) + + if self.use_gr_prosody_phone: + self.prosody_phone_predictor = nn.Sequential( + GradientReversal(alpha=1.0), CNNLSTM(in_channels, 5003, 1) + ) + + if self.use_gr_x_timbre: + self.x_timbre_predictor = nn.Sequential( + GradientReversal(alpha=1), + CNNLSTM(in_channels, 245200, 1, global_pred=True), + ) + + self.melspec_linear = nn.Linear(20, 256) + self.melspec_encoder = TransformerEncoder( + enc_emb_tokens=None, + encoder_layer=4, + encoder_hidden=256, + encoder_head=4, + conv_filter_size=1024, + conv_kernel_size=5, + encoder_dropout=0.1, + use_cln=False, + cfg=None, + ) + + self.reset_parameters() + + def quantize(self, x, prosody_feature, n_quantizers=None): + outs, qs, commit_loss, quantized_buf = 0, [], [], [] + + # prosody + f0_input = prosody_feature.transpose(1, 2) # (B, T, 20) + f0_input = self.melspec_linear(f0_input) + f0_input = self.melspec_encoder(f0_input, None, None) + f0_input = f0_input.transpose(1, 2) + f0_quantizer = self.quantizer[0] + out, q, commit, quantized = f0_quantizer(f0_input, n_quantizers=n_quantizers) + outs += out + qs.append(q) + quantized_buf.append(quantized.sum(0)) + commit_loss.append(commit) + + # phone + phone_input = x + phone_quantizer = self.quantizer[1] + out, q, commit, quantized = phone_quantizer( + phone_input, n_quantizers=n_quantizers + ) + outs += out + qs.append(q) + quantized_buf.append(quantized.sum(0)) + commit_loss.append(commit) + + # residual + if self.vq_num_q_r > 0: + residual_quantizer = self.quantizer[2] + residual_input = x - (quantized_buf[0] + quantized_buf[1]).detach() + out, q, commit, quantized = residual_quantizer( + residual_input, n_quantizers=n_quantizers + ) + outs += out + qs.append(q) + quantized_buf.append(quantized.sum(0)) # [L, B, C, T] -> [B, C, T] + commit_loss.append(commit) + + qs = torch.cat(qs, dim=0) + commit_loss = torch.cat(commit_loss, dim=0) + return outs, qs, commit_loss, quantized_buf + + def forward( + self, + x, + prosody_feature, + vq=True, + get_vq=False, + eval_vq=True, + speaker_embedding=None, + n_quantizers=None, + quantized=None, + ): + if get_vq: + return self.quantizer.get_emb() + if vq is True: + if eval_vq: + self.quantizer.eval() + x_timbre = x + outs, qs, commit_loss, quantized_buf = self.quantize( + x, prosody_feature, n_quantizers=n_quantizers + ) + + x_timbre = x_timbre.transpose(1, 2) + x_timbre = self.timbre_encoder(x_timbre, None, None) + x_timbre = x_timbre.transpose(1, 2) + spk_embs = torch.mean(x_timbre, dim=2) + return outs, qs, commit_loss, quantized_buf, spk_embs + + out = {} + + layer_0 = quantized[0] + f0, uv = self.f0_predictor(layer_0) + f0 = rearrange(f0, "... 1 -> ...") + uv = rearrange(uv, "... 1 -> ...") + + layer_1 = quantized[1] + (phone,) = self.phone_predictor(layer_1) + + out = {"f0": f0, "uv": uv, "phone": phone} + + if self.use_gr_prosody_phone: + (prosody_phone,) = self.prosody_phone_predictor(layer_0) + out["prosody_phone"] = prosody_phone + + if self.use_gr_content_f0: + content_f0, content_uv = self.content_f0_predictor(layer_1) + content_f0 = rearrange(content_f0, "... 1 -> ...") + content_uv = rearrange(content_uv, "... 1 -> ...") + out["content_f0"] = content_f0 + out["content_uv"] = content_uv + + if self.vq_num_q_r > 0: + layer_2 = quantized[2] + + if self.use_gr_residual_f0: + res_f0, res_uv = self.res_f0_predictor(layer_2) + res_f0 = rearrange(res_f0, "... 1 -> ...") + res_uv = rearrange(res_uv, "... 1 -> ...") + out["res_f0"] = res_f0 + out["res_uv"] = res_uv + + if self.use_gr_residual_phone: + (res_phone,) = self.res_phone_predictor(layer_2) + out["res_phone"] = res_phone + + style = self.timbre_linear(speaker_embedding).unsqueeze(2) # (B, 2d, 1) + gamma, beta = style.chunk(2, 1) # (B, d, 1) + if self.vq_num_q_r > 0: + if self.use_random_mask_residual: + bsz = quantized[2].shape[0] + res_mask = np.random.choice( + [0, 1], + size=bsz, + p=[ + self.prob_random_mask_residual, + 1 - self.prob_random_mask_residual, + ], + ) + res_mask = ( + torch.from_numpy(res_mask).unsqueeze(1).unsqueeze(1) + ) # (B, 1, 1) + res_mask = res_mask.to( + device=quantized[2].device, dtype=quantized[2].dtype + ) + x = ( + quantized[0].detach() + + quantized[1].detach() + + quantized[2] * res_mask + ) + # x = quantized_perturbe[0].detach() + quantized[1].detach() + quantized[2] * res_mask + else: + x = quantized[0].detach() + quantized[1].detach() + quantized[2] + # x = quantized_perturbe[0].detach() + quantized[1].detach() + quantized[2] + else: + x = quantized[0].detach() + quantized[1].detach() + # x = quantized_perturbe[0].detach() + quantized[1].detach() + + if self.use_gr_x_timbre: + (x_timbre,) = self.x_timbre_predictor(x) + out["x_timbre"] = x_timbre + + x = x.transpose(1, 2) + x = self.timbre_norm(x) + x = x.transpose(1, 2) + x = x * gamma + beta + + x = self.model(x) + out["audio"] = x + + return out + + def vq2emb(self, vq, use_residual=True): + # vq: [num_quantizer, B, T] + self.quantizer = self.quantizer.eval() + out = 0 + out += self.quantizer[0].vq2emb(vq[0 : self.vq_num_q_p]) + out += self.quantizer[1].vq2emb( + vq[self.vq_num_q_p : self.vq_num_q_p + self.vq_num_q_c] + ) + if self.vq_num_q_r > 0 and use_residual: + out += self.quantizer[2].vq2emb(vq[self.vq_num_q_p + self.vq_num_q_c :]) + return out + + def inference(self, x, speaker_embedding): + style = self.timbre_linear(speaker_embedding).unsqueeze(2) # (B, 2d, 1) + gamma, beta = style.chunk(2, 1) # (B, d, 1) + x = x.transpose(1, 2) + x = self.timbre_norm(x) + x = x.transpose(1, 2) + x = x * gamma + beta + x = self.model(x) + return x + + def remove_weight_norm(self): + """Remove weight normalization module from all of the layers.""" + + def _remove_weight_norm(m): + try: + torch.nn.utils.remove_weight_norm(m) + except ValueError: # this module didn't have weight norm + return + + self.apply(_remove_weight_norm) + + def apply_weight_norm(self): + """Apply weight normalization module from all of the layers.""" + + def _apply_weight_norm(m): + if isinstance(m, nn.Conv1d) or isinstance(m, nn.ConvTranspose1d): + torch.nn.utils.weight_norm(m) + + self.apply(_apply_weight_norm) + + def reset_parameters(self): + self.apply(init_weights) diff --git a/models/codec/ns3_codec/gradient_reversal.py b/models/codec/ns3_codec/gradient_reversal.py new file mode 100644 index 0000000..d09396e --- /dev/null +++ b/models/codec/ns3_codec/gradient_reversal.py @@ -0,0 +1,35 @@ +# Copyright (c) 2023 Amphion. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +from torch.autograd import Function +import torch +from torch import nn + + +class GradientReversal(Function): + @staticmethod + def forward(ctx, x, alpha): + ctx.save_for_backward(x, alpha) + return x + + @staticmethod + def backward(ctx, grad_output): + grad_input = None + _, alpha = ctx.saved_tensors + if ctx.needs_input_grad[0]: + grad_input = -alpha * grad_output + return grad_input, None + + +revgrad = GradientReversal.apply + + +class GradientReversal(nn.Module): + def __init__(self, alpha): + super().__init__() + self.alpha = torch.tensor(alpha, requires_grad=False) + + def forward(self, x): + return revgrad(x, self.alpha) diff --git a/models/codec/ns3_codec/melspec.py b/models/codec/ns3_codec/melspec.py new file mode 100644 index 0000000..cbf1cd2 --- /dev/null +++ b/models/codec/ns3_codec/melspec.py @@ -0,0 +1,102 @@ +import torch +import pyworld as pw +import numpy as np +import soundfile as sf +import os +from torchaudio.functional import pitch_shift +import librosa +from librosa.filters import mel as librosa_mel_fn +import torch.nn as nn +import torch.nn.functional as F + + +def dynamic_range_compression(x, C=1, clip_val=1e-5): + return np.log(np.clip(x, a_min=clip_val, a_max=None) * C) + + +def dynamic_range_decompression(x, C=1): + return np.exp(x) / C + + +def dynamic_range_compression_torch(x, C=1, clip_val=1e-5): + return torch.log(torch.clamp(x, min=clip_val) * C) + + +def dynamic_range_decompression_torch(x, C=1): + return torch.exp(x) / C + + +def spectral_normalize_torch(magnitudes): + output = dynamic_range_compression_torch(magnitudes) + return output + + +def spectral_de_normalize_torch(magnitudes): + output = dynamic_range_decompression_torch(magnitudes) + return output + + +class MelSpectrogram(nn.Module): + def __init__( + self, + n_fft, + num_mels, + sampling_rate, + hop_size, + win_size, + fmin, + fmax, + center=False, + ): + super(MelSpectrogram, self).__init__() + self.n_fft = n_fft + self.hop_size = hop_size + self.win_size = win_size + self.sampling_rate = sampling_rate + self.num_mels = num_mels + self.fmin = fmin + self.fmax = fmax + self.center = center + + mel_basis = {} + hann_window = {} + + mel = librosa_mel_fn( + sr=sampling_rate, n_fft=n_fft, n_mels=num_mels, fmin=fmin, fmax=fmax + ) + mel_basis = torch.from_numpy(mel).float() + hann_window = torch.hann_window(win_size) + + self.register_buffer("mel_basis", mel_basis) + self.register_buffer("hann_window", hann_window) + + def forward(self, y): + y = torch.nn.functional.pad( + y.unsqueeze(1), + ( + int((self.n_fft - self.hop_size) / 2), + int((self.n_fft - self.hop_size) / 2), + ), + mode="reflect", + ) + y = y.squeeze(1) + spec = torch.stft( + y, + self.n_fft, + hop_length=self.hop_size, + win_length=self.win_size, + window=self.hann_window, + center=self.center, + pad_mode="reflect", + normalized=False, + onesided=True, + return_complex=True, + ) + spec = torch.view_as_real(spec) + + spec = torch.sqrt(spec.pow(2).sum(-1) + (1e-9)) + + spec = torch.matmul(self.mel_basis, spec) + spec = spectral_normalize_torch(spec) + + return spec diff --git a/models/codec/ns3_codec/quantize/__init__.py b/models/codec/ns3_codec/quantize/__init__.py new file mode 100644 index 0000000..2cb7b40 --- /dev/null +++ b/models/codec/ns3_codec/quantize/__init__.py @@ -0,0 +1,7 @@ +# Copyright (c) 2023 Amphion. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +from .fvq import * +from .rvq import * diff --git a/models/codec/ns3_codec/quantize/fvq.py b/models/codec/ns3_codec/quantize/fvq.py new file mode 100644 index 0000000..4ade35d --- /dev/null +++ b/models/codec/ns3_codec/quantize/fvq.py @@ -0,0 +1,116 @@ +# Copyright (c) 2023 Amphion. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +from typing import Union + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +from einops import rearrange +from torch.nn.utils import weight_norm + + +class FactorizedVectorQuantize(nn.Module): + def __init__(self, dim, codebook_size, codebook_dim, commitment, **kwargs): + super().__init__() + self.codebook_size = codebook_size + self.codebook_dim = codebook_dim + self.commitment = commitment + + if dim != self.codebook_dim: + self.in_proj = weight_norm(nn.Linear(dim, self.codebook_dim)) + self.out_proj = weight_norm(nn.Linear(self.codebook_dim, dim)) + else: + self.in_proj = nn.Identity() + self.out_proj = nn.Identity() + self._codebook = nn.Embedding(codebook_size, self.codebook_dim) + + @property + def codebook(self): + return self._codebook + + def forward(self, z): + """Quantized the input tensor using a fixed codebook and returns + the corresponding codebook vectors + + Parameters + ---------- + z : Tensor[B x D x T] + + Returns + ------- + Tensor[B x D x T] + Quantized continuous representation of input + Tensor[1] + Commitment loss to train encoder to predict vectors closer to codebook + entries + Tensor[1] + Codebook loss to update the codebook + Tensor[B x T] + Codebook indices (quantized discrete representation of input) + Tensor[B x D x T] + Projected latents (continuous representation of input before quantization) + """ + # transpose since we use linear + + z = rearrange(z, "b d t -> b t d") + + # Factorized codes project input into low-dimensional space + z_e = self.in_proj(z) # z_e : (B x T x D) + z_e = rearrange(z_e, "b t d -> b d t") + z_q, indices = self.decode_latents(z_e) + + if self.training: + commitment_loss = ( + F.mse_loss(z_e, z_q.detach(), reduction="none").mean([1, 2]) + * self.commitment + ) + codebook_loss = F.mse_loss(z_q, z_e.detach(), reduction="none").mean([1, 2]) + commit_loss = commitment_loss + codebook_loss + else: + commit_loss = torch.zeros(z.shape[0], device=z.device) + + z_q = ( + z_e + (z_q - z_e).detach() + ) # noop in forward pass, straight-through gradient estimator in backward pass + + z_q = rearrange(z_q, "b d t -> b t d") + z_q = self.out_proj(z_q) + z_q = rearrange(z_q, "b t d -> b d t") + + return z_q, indices, commit_loss + + def vq2emb(self, vq, proj=True): + emb = self.embed_code(vq) + if proj: + emb = self.out_proj(emb) + return emb.transpose(1, 2) + + def get_emb(self): + return self.codebook.weight + + def embed_code(self, embed_id): + return F.embedding(embed_id, self.codebook.weight) + + def decode_code(self, embed_id): + return self.embed_code(embed_id).transpose(1, 2) + + def decode_latents(self, latents): + encodings = rearrange(latents, "b d t -> (b t) d") + codebook = self.codebook.weight # codebook: (N x D) + # L2 normalize encodings and codebook + encodings = F.normalize(encodings) + codebook = F.normalize(codebook) + + # Compute euclidean distance with codebook + dist = ( + encodings.pow(2).sum(1, keepdim=True) + - 2 * encodings @ codebook.t() + + codebook.pow(2).sum(1, keepdim=True).t() + ) + indices = rearrange((-dist).max(1)[1], "(b t) -> b t", b=latents.size(0)) + z_q = self.decode_code(indices) + return z_q, indices diff --git a/models/codec/ns3_codec/quantize/rvq.py b/models/codec/ns3_codec/quantize/rvq.py new file mode 100644 index 0000000..d22d88d --- /dev/null +++ b/models/codec/ns3_codec/quantize/rvq.py @@ -0,0 +1,87 @@ +# Copyright (c) 2023 Amphion. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +import math +import torch +from torch import nn +from .fvq import FactorizedVectorQuantize + + +class ResidualVQ(nn.Module): + """Follows Algorithm 1. in https://arxiv.org/pdf/2107.03312.pdf""" + + def __init__(self, *, num_quantizers, codebook_size, **kwargs): + super().__init__() + VQ = FactorizedVectorQuantize + if type(codebook_size) == int: + codebook_size = [codebook_size] * num_quantizers + self.layers = nn.ModuleList( + [VQ(codebook_size=2**size, **kwargs) for size in codebook_size] + ) + self.num_quantizers = num_quantizers + self.quantizer_dropout = kwargs.get("quantizer_dropout", 0.0) + self.dropout_type = kwargs.get("dropout_type", None) + + def forward(self, x, n_quantizers=None): + quantized_out = 0.0 + residual = x + + all_losses = [] + all_indices = [] + all_quantized = [] + + if n_quantizers is None: + n_quantizers = self.num_quantizers + if self.training: + n_quantizers = torch.ones((x.shape[0],)) * self.num_quantizers + 1 + if self.dropout_type == "linear": + dropout = torch.randint(1, self.num_quantizers + 1, (x.shape[0],)) + elif self.dropout_type == "exp": + dropout = torch.randint( + 1, int(math.log2(self.num_quantizers)), (x.shape[0],) + ) + dropout = torch.pow(2, dropout) + n_dropout = int(x.shape[0] * self.quantizer_dropout) + n_quantizers[:n_dropout] = dropout[:n_dropout] + n_quantizers = n_quantizers.to(x.device) + + for idx, layer in enumerate(self.layers): + if not self.training and idx >= n_quantizers: + break + quantized, indices, loss = layer(residual) + + mask = ( + torch.full((x.shape[0],), fill_value=idx, device=x.device) + < n_quantizers + ) + + residual = residual - quantized + + quantized_out = quantized_out + quantized * mask[:, None, None] + + # loss + loss = (loss * mask).mean() + + all_indices.append(indices) + all_losses.append(loss) + all_quantized.append(quantized) + all_losses, all_indices, all_quantized = map( + torch.stack, (all_losses, all_indices, all_quantized) + ) + return quantized_out, all_indices, all_losses, all_quantized + + def vq2emb(self, vq): + # vq: [n_quantizers, B, T] + quantized_out = 0.0 + for idx, layer in enumerate(self.layers): + quantized = layer.vq2emb(vq[idx]) + quantized_out += quantized + return quantized_out + + def get_emb(self): + embs = [] + for idx, layer in enumerate(self.layers): + embs.append(layer.get_emb()) + return embs diff --git a/models/codec/ns3_codec/transformer.py b/models/codec/ns3_codec/transformer.py new file mode 100644 index 0000000..146d0f3 --- /dev/null +++ b/models/codec/ns3_codec/transformer.py @@ -0,0 +1,234 @@ +# Copyright (c) 2023 Amphion. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +import numpy as np +import torch +import torch.nn as nn +import math +from torch.nn import functional as F + + +class StyleAdaptiveLayerNorm(nn.Module): + def __init__(self, normalized_shape, eps=1e-5): + super().__init__() + self.in_dim = normalized_shape + self.norm = nn.LayerNorm(self.in_dim, eps=eps, elementwise_affine=False) + self.style = nn.Linear(self.in_dim, self.in_dim * 2) + self.style.bias.data[: self.in_dim] = 1 + self.style.bias.data[self.in_dim :] = 0 + + def forward(self, x, condition): + # x: (B, T, d); condition: (B, T, d) + + style = self.style(torch.mean(condition, dim=1, keepdim=True)) + + gamma, beta = style.chunk(2, -1) + + out = self.norm(x) + + out = gamma * out + beta + return out + + +class PositionalEncoding(nn.Module): + def __init__(self, d_model, dropout, max_len=5000): + super().__init__() + + self.dropout = dropout + position = torch.arange(max_len).unsqueeze(1) + div_term = torch.exp( + torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model) + ) + pe = torch.zeros(max_len, 1, d_model) + pe[:, 0, 0::2] = torch.sin(position * div_term) + pe[:, 0, 1::2] = torch.cos(position * div_term) + self.register_buffer("pe", pe) + + def forward(self, x): + x = x + self.pe[: x.size(0)] + return F.dropout(x, self.dropout, training=self.training) + + +class TransformerFFNLayer(nn.Module): + def __init__( + self, encoder_hidden, conv_filter_size, conv_kernel_size, encoder_dropout + ): + super().__init__() + + self.encoder_hidden = encoder_hidden + self.conv_filter_size = conv_filter_size + self.conv_kernel_size = conv_kernel_size + self.encoder_dropout = encoder_dropout + + self.ffn_1 = nn.Conv1d( + self.encoder_hidden, + self.conv_filter_size, + self.conv_kernel_size, + padding=self.conv_kernel_size // 2, + ) + self.ffn_1.weight.data.normal_(0.0, 0.02) + self.ffn_2 = nn.Linear(self.conv_filter_size, self.encoder_hidden) + self.ffn_2.weight.data.normal_(0.0, 0.02) + + def forward(self, x): + # x: (B, T, d) + x = self.ffn_1(x.permute(0, 2, 1)).permute( + 0, 2, 1 + ) # (B, T, d) -> (B, d, T) -> (B, T, d) + x = F.relu(x) + x = F.dropout(x, self.encoder_dropout, training=self.training) + x = self.ffn_2(x) + return x + + +class TransformerEncoderLayer(nn.Module): + def __init__( + self, + encoder_hidden, + encoder_head, + conv_filter_size, + conv_kernel_size, + encoder_dropout, + use_cln, + ): + super().__init__() + self.encoder_hidden = encoder_hidden + self.encoder_head = encoder_head + self.conv_filter_size = conv_filter_size + self.conv_kernel_size = conv_kernel_size + self.encoder_dropout = encoder_dropout + self.use_cln = use_cln + + if not self.use_cln: + self.ln_1 = nn.LayerNorm(self.encoder_hidden) + self.ln_2 = nn.LayerNorm(self.encoder_hidden) + else: + self.ln_1 = StyleAdaptiveLayerNorm(self.encoder_hidden) + self.ln_2 = StyleAdaptiveLayerNorm(self.encoder_hidden) + + self.self_attn = nn.MultiheadAttention( + self.encoder_hidden, self.encoder_head, batch_first=True + ) + + self.ffn = TransformerFFNLayer( + self.encoder_hidden, + self.conv_filter_size, + self.conv_kernel_size, + self.encoder_dropout, + ) + + def forward(self, x, key_padding_mask, conditon=None): + # x: (B, T, d); key_padding_mask: (B, T), mask is 0; condition: (B, T, d) + + # self attention + residual = x + if self.use_cln: + x = self.ln_1(x, conditon) + else: + x = self.ln_1(x) + + if key_padding_mask != None: + key_padding_mask_input = ~(key_padding_mask.bool()) + else: + key_padding_mask_input = None + x, _ = self.self_attn( + query=x, key=x, value=x, key_padding_mask=key_padding_mask_input + ) + x = F.dropout(x, self.encoder_dropout, training=self.training) + x = residual + x + + # ffn + residual = x + if self.use_cln: + x = self.ln_2(x, conditon) + else: + x = self.ln_2(x) + x = self.ffn(x) + x = residual + x + + return x + + +class TransformerEncoder(nn.Module): + def __init__( + self, + enc_emb_tokens=None, + encoder_layer=4, + encoder_hidden=256, + encoder_head=4, + conv_filter_size=1024, + conv_kernel_size=5, + encoder_dropout=0.1, + use_cln=False, + cfg=None, + ): + super().__init__() + + self.encoder_layer = ( + encoder_layer if encoder_layer is not None else cfg.encoder_layer + ) + self.encoder_hidden = ( + encoder_hidden if encoder_hidden is not None else cfg.encoder_hidden + ) + self.encoder_head = ( + encoder_head if encoder_head is not None else cfg.encoder_head + ) + self.conv_filter_size = ( + conv_filter_size if conv_filter_size is not None else cfg.conv_filter_size + ) + self.conv_kernel_size = ( + conv_kernel_size if conv_kernel_size is not None else cfg.conv_kernel_size + ) + self.encoder_dropout = ( + encoder_dropout if encoder_dropout is not None else cfg.encoder_dropout + ) + self.use_cln = use_cln if use_cln is not None else cfg.use_cln + + if enc_emb_tokens != None: + self.use_enc_emb = True + self.enc_emb_tokens = enc_emb_tokens + else: + self.use_enc_emb = False + + self.position_emb = PositionalEncoding( + self.encoder_hidden, self.encoder_dropout + ) + + self.layers = nn.ModuleList([]) + self.layers.extend( + [ + TransformerEncoderLayer( + self.encoder_hidden, + self.encoder_head, + self.conv_filter_size, + self.conv_kernel_size, + self.encoder_dropout, + self.use_cln, + ) + for i in range(self.encoder_layer) + ] + ) + + if self.use_cln: + self.last_ln = StyleAdaptiveLayerNorm(self.encoder_hidden) + else: + self.last_ln = nn.LayerNorm(self.encoder_hidden) + + def forward(self, x, key_padding_mask, condition=None): + if len(x.shape) == 2 and self.use_enc_emb: + x = self.enc_emb_tokens(x) + x = self.position_emb(x) + else: + x = self.position_emb(x) # (B, T, d) + + for layer in self.layers: + x = layer(x, key_padding_mask, condition) + + if self.use_cln: + x = self.last_ln(x, condition) + else: + x = self.last_ln(x) + + return x diff --git a/models/codec/speechtokenizer/model.py b/models/codec/speechtokenizer/model.py new file mode 100644 index 0000000..b722d38 --- /dev/null +++ b/models/codec/speechtokenizer/model.py @@ -0,0 +1,184 @@ +# Copyright (c) 2023 Amphion. +# +# This code is modified from https://github.com/ZhangXInFD/SpeechTokenizer/blob/main/speechtokenizer/model.py +# Licensed under Apache License 2.0 + +from .modules.seanet import SEANetEncoder, SEANetDecoder +from .modules.quantization import ResidualVectorQuantizer +import torch.nn as nn +from einops import rearrange +import torch +import numpy as np + + +class SpeechTokenizer(nn.Module): + def __init__(self, config): + """ + + Parameters + ---------- + config : json + Model Config. + + """ + super().__init__() + self.encoder = SEANetEncoder( + n_filters=config.get("n_filters"), + dimension=config.get("dimension"), + ratios=config.get("strides"), + lstm=config.get("lstm_layers"), + bidirectional=config.get("bidirectional"), + dilation_base=config.get("dilation_base"), + residual_kernel_size=config.get("residual_kernel_size"), + n_residual_layers=config.get("n_residual_layers"), + activation=config.get("activation"), + ) + self.sample_rate = config.get("sample_rate") + self.n_q = config.get("n_q") + self.downsample_rate = np.prod(config.get("strides")) + if config.get("dimension") != config.get("semantic_dimension"): + self.transform = nn.Linear( + config.get("dimension"), config.get("semantic_dimension") + ) + else: + self.transform = nn.Identity() + self.quantizer = ResidualVectorQuantizer( + dimension=config.get("dimension"), + n_q=config.get("n_q"), + bins=config.get("codebook_size"), + ) + self.decoder = SEANetDecoder( + n_filters=config.get("n_filters"), + dimension=config.get("dimension"), + ratios=config.get("strides"), + lstm=config.get("lstm_layers"), + bidirectional=False, + dilation_base=config.get("dilation_base"), + residual_kernel_size=config.get("residual_kernel_size"), + n_residual_layers=config.get("n_residual_layers"), + activation=config.get("activation"), + ) + + @classmethod + def load_from_checkpoint(cls, config_path: str, ckpt_path: str): + """ + + Parameters + ---------- + config_path : str + Path of model configuration file. + ckpt_path : str + Path of model checkpoint. + + Returns + ------- + model : SpeechTokenizer + SpeechTokenizer model. + + """ + import json + + with open(config_path) as f: + cfg = json.load(f) + model = cls(cfg) + params = torch.load(ckpt_path, map_location="cpu") + model.load_state_dict(params) + return model + + def forward(self, x: torch.tensor, n_q: int = None, layers: list = [0]): + """ + + Parameters + ---------- + x : torch.tensor + Input wavs. Shape: (batch, channels, timesteps). + n_q : int, optional + Number of quantizers in RVQ used to encode. The default is all layers. + layers : list[int], optional + Layers of RVQ should return quantized result. The default is the first layer. + + Returns + ------- + o : torch.tensor + Output wavs. Shape: (batch, channels, timesteps). + commit_loss : torch.tensor + Commitment loss from residual vector quantizers. + feature : torch.tensor + Output of RVQ's first layer. Shape: (batch, timesteps, dimension) + + """ + n_q = n_q if n_q else self.n_q + e = self.encoder(x) + quantized, codes, commit_loss, quantized_list = self.quantizer( + e, n_q=n_q, layers=layers + ) + feature = rearrange(quantized_list[0], "b d t -> b t d") + feature = self.transform(feature) + o = self.decoder(quantized) + return o, commit_loss, feature + + def forward_feature(self, x: torch.tensor, layers: list = None): + """ + + Parameters + ---------- + x : torch.tensor + Input wavs. Shape should be (batch, channels, timesteps). + layers : list[int], optional + Layers of RVQ should return quantized result. The default is all layers. + + Returns + ------- + quantized_list : list[torch.tensor] + Quantized of required layers. + + """ + e = self.encoder(x) + layers = layers if layers else list(range(self.n_q)) + quantized, codes, commit_loss, quantized_list = self.quantizer(e, layers=layers) + return quantized_list + + def encode(self, x: torch.tensor, n_q: int = None, st: int = None): + """ + + Parameters + ---------- + x : torch.tensor + Input wavs. Shape: (batch, channels, timesteps). + n_q : int, optional + Number of quantizers in RVQ used to encode. The default is all layers. + st : int, optional + Start quantizer index in RVQ. The default is 0. + + Returns + ------- + codes : torch.tensor + Output indices for each quantizer. Shape: (n_q, batch, timesteps) + + """ + e = self.encoder(x) + if st is None: + st = 0 + n_q = n_q if n_q else self.n_q + codes = self.quantizer.encode(e, n_q=n_q, st=st) + return codes + + def decode(self, codes: torch.tensor, st: int = 0): + """ + + Parameters + ---------- + codes : torch.tensor + Indices for each quantizer. Shape: (n_q, batch, timesteps). + st : int, optional + Start quantizer index in RVQ. The default is 0. + + Returns + ------- + o : torch.tensor + Reconstruct wavs from codes. Shape: (batch, channels, timesteps) + + """ + quantized = self.quantizer.decode(codes, st=st) + o = self.decoder(quantized) + return o diff --git a/models/codec/speechtokenizer/modules/__init__.py b/models/codec/speechtokenizer/modules/__init__.py new file mode 100644 index 0000000..0581347 --- /dev/null +++ b/models/codec/speechtokenizer/modules/__init__.py @@ -0,0 +1,27 @@ +# Copyright (c) 2023 Amphion. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. +# This source file is copied from https://github.com/facebookresearch/encodec + +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +"""Torch modules.""" + +# flake8: noqa +from .conv import ( + pad1d, + unpad1d, + NormConv1d, + NormConvTranspose1d, + NormConv2d, + NormConvTranspose2d, + SConv1d, + SConvTranspose1d, +) +from .lstm import SLSTM +from .seanet import SEANetEncoder, SEANetDecoder diff --git a/models/codec/speechtokenizer/modules/conv.py b/models/codec/speechtokenizer/modules/conv.py new file mode 100644 index 0000000..0352b8b --- /dev/null +++ b/models/codec/speechtokenizer/modules/conv.py @@ -0,0 +1,346 @@ +# Copyright (c) 2023 Amphion. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. +# This source file is copied from https://github.com/facebookresearch/encodec + +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +"""Convolutional layers wrappers and utilities.""" + +import math +import typing as tp +import warnings + +import torch +from torch import nn +from torch.nn import functional as F +from torch.nn.utils import spectral_norm, weight_norm + +from .norm import ConvLayerNorm + + +CONV_NORMALIZATIONS = frozenset( + [ + "none", + "weight_norm", + "spectral_norm", + "time_layer_norm", + "layer_norm", + "time_group_norm", + ] +) + + +def apply_parametrization_norm(module: nn.Module, norm: str = "none") -> nn.Module: + assert norm in CONV_NORMALIZATIONS + if norm == "weight_norm": + return weight_norm(module) + elif norm == "spectral_norm": + return spectral_norm(module) + else: + # We already check was in CONV_NORMALIZATION, so any other choice + # doesn't need reparametrization. + return module + + +def get_norm_module( + module: nn.Module, causal: bool = False, norm: str = "none", **norm_kwargs +) -> nn.Module: + """Return the proper normalization module. If causal is True, this will ensure the returned + module is causal, or return an error if the normalization doesn't support causal evaluation. + """ + assert norm in CONV_NORMALIZATIONS + if norm == "layer_norm": + assert isinstance(module, nn.modules.conv._ConvNd) + return ConvLayerNorm(module.out_channels, **norm_kwargs) + elif norm == "time_group_norm": + if causal: + raise ValueError("GroupNorm doesn't support causal evaluation.") + assert isinstance(module, nn.modules.conv._ConvNd) + return nn.GroupNorm(1, module.out_channels, **norm_kwargs) + else: + return nn.Identity() + + +def get_extra_padding_for_conv1d( + x: torch.Tensor, kernel_size: int, stride: int, padding_total: int = 0 +) -> int: + """See `pad_for_conv1d`.""" + length = x.shape[-1] + n_frames = (length - kernel_size + padding_total) / stride + 1 + ideal_length = (math.ceil(n_frames) - 1) * stride + (kernel_size - padding_total) + return ideal_length - length + + +def pad_for_conv1d( + x: torch.Tensor, kernel_size: int, stride: int, padding_total: int = 0 +): + """Pad for a convolution to make sure that the last window is full. + Extra padding is added at the end. This is required to ensure that we can rebuild + an output of the same length, as otherwise, even with padding, some time steps + might get removed. + For instance, with total padding = 4, kernel size = 4, stride = 2: + 0 0 1 2 3 4 5 0 0 # (0s are padding) + 1 2 3 # (output frames of a convolution, last 0 is never used) + 0 0 1 2 3 4 5 0 # (output of tr. conv., but pos. 5 is going to get removed as padding) + 1 2 3 4 # once you removed padding, we are missing one time step ! + """ + extra_padding = get_extra_padding_for_conv1d(x, kernel_size, stride, padding_total) + return F.pad(x, (0, extra_padding)) + + +def pad1d( + x: torch.Tensor, + paddings: tp.Tuple[int, int], + mode: str = "zero", + value: float = 0.0, +): + """Tiny wrapper around F.pad, just to allow for reflect padding on small input. + If this is the case, we insert extra 0 padding to the right before the reflection happen. + """ + length = x.shape[-1] + padding_left, padding_right = paddings + assert padding_left >= 0 and padding_right >= 0, (padding_left, padding_right) + if mode == "reflect": + max_pad = max(padding_left, padding_right) + extra_pad = 0 + if length <= max_pad: + extra_pad = max_pad - length + 1 + x = F.pad(x, (0, extra_pad)) + padded = F.pad(x, paddings, mode, value) + end = padded.shape[-1] - extra_pad + return padded[..., :end] + else: + return F.pad(x, paddings, mode, value) + + +def unpad1d(x: torch.Tensor, paddings: tp.Tuple[int, int]): + """Remove padding from x, handling properly zero padding. Only for 1d!""" + padding_left, padding_right = paddings + assert padding_left >= 0 and padding_right >= 0, (padding_left, padding_right) + assert (padding_left + padding_right) <= x.shape[-1] + end = x.shape[-1] - padding_right + return x[..., padding_left:end] + + +class NormConv1d(nn.Module): + """Wrapper around Conv1d and normalization applied to this conv + to provide a uniform interface across normalization approaches. + """ + + def __init__( + self, + *args, + causal: bool = False, + norm: str = "none", + norm_kwargs: tp.Dict[str, tp.Any] = {}, + **kwargs, + ): + super().__init__() + self.conv = apply_parametrization_norm(nn.Conv1d(*args, **kwargs), norm) + self.norm = get_norm_module(self.conv, causal, norm, **norm_kwargs) + self.norm_type = norm + + def forward(self, x): + x = self.conv(x) + x = self.norm(x) + return x + + +class NormConv2d(nn.Module): + """Wrapper around Conv2d and normalization applied to this conv + to provide a uniform interface across normalization approaches. + """ + + def __init__( + self, + *args, + norm: str = "none", + norm_kwargs: tp.Dict[str, tp.Any] = {}, + **kwargs, + ): + super().__init__() + self.conv = apply_parametrization_norm(nn.Conv2d(*args, **kwargs), norm) + self.norm = get_norm_module(self.conv, causal=False, norm=norm, **norm_kwargs) + self.norm_type = norm + + def forward(self, x): + x = self.conv(x) + x = self.norm(x) + return x + + +class NormConvTranspose1d(nn.Module): + """Wrapper around ConvTranspose1d and normalization applied to this conv + to provide a uniform interface across normalization approaches. + """ + + def __init__( + self, + *args, + causal: bool = False, + norm: str = "none", + norm_kwargs: tp.Dict[str, tp.Any] = {}, + **kwargs, + ): + super().__init__() + self.convtr = apply_parametrization_norm( + nn.ConvTranspose1d(*args, **kwargs), norm + ) + self.norm = get_norm_module(self.convtr, causal, norm, **norm_kwargs) + self.norm_type = norm + + def forward(self, x): + x = self.convtr(x) + x = self.norm(x) + return x + + +class NormConvTranspose2d(nn.Module): + """Wrapper around ConvTranspose2d and normalization applied to this conv + to provide a uniform interface across normalization approaches. + """ + + def __init__( + self, + *args, + norm: str = "none", + norm_kwargs: tp.Dict[str, tp.Any] = {}, + **kwargs, + ): + super().__init__() + self.convtr = apply_parametrization_norm( + nn.ConvTranspose2d(*args, **kwargs), norm + ) + self.norm = get_norm_module(self.convtr, causal=False, norm=norm, **norm_kwargs) + + def forward(self, x): + x = self.convtr(x) + x = self.norm(x) + return x + + +class SConv1d(nn.Module): + """Conv1d with some builtin handling of asymmetric or causal padding + and normalization. + """ + + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: int, + stride: int = 1, + dilation: int = 1, + groups: int = 1, + bias: bool = True, + causal: bool = False, + norm: str = "none", + norm_kwargs: tp.Dict[str, tp.Any] = {}, + pad_mode: str = "reflect", + ): + super().__init__() + # warn user on unusual setup between dilation and stride + if stride > 1 and dilation > 1: + warnings.warn( + "SConv1d has been initialized with stride > 1 and dilation > 1" + f" (kernel_size={kernel_size} stride={stride}, dilation={dilation})." + ) + self.conv = NormConv1d( + in_channels, + out_channels, + kernel_size, + stride, + dilation=dilation, + groups=groups, + bias=bias, + causal=causal, + norm=norm, + norm_kwargs=norm_kwargs, + ) + self.causal = causal + self.pad_mode = pad_mode + + def forward(self, x): + B, C, T = x.shape + kernel_size = self.conv.conv.kernel_size[0] + stride = self.conv.conv.stride[0] + dilation = self.conv.conv.dilation[0] + padding_total = (kernel_size - 1) * dilation - (stride - 1) + extra_padding = get_extra_padding_for_conv1d( + x, kernel_size, stride, padding_total + ) + if self.causal: + # Left padding for causal + x = pad1d(x, (padding_total, extra_padding), mode=self.pad_mode) + else: + # Asymmetric padding required for odd strides + padding_right = padding_total // 2 + padding_left = padding_total - padding_right + x = pad1d( + x, (padding_left, padding_right + extra_padding), mode=self.pad_mode + ) + return self.conv(x) + + +class SConvTranspose1d(nn.Module): + """ConvTranspose1d with some builtin handling of asymmetric or causal padding + and normalization. + """ + + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: int, + stride: int = 1, + causal: bool = False, + norm: str = "none", + trim_right_ratio: float = 1.0, + norm_kwargs: tp.Dict[str, tp.Any] = {}, + ): + super().__init__() + self.convtr = NormConvTranspose1d( + in_channels, + out_channels, + kernel_size, + stride, + causal=causal, + norm=norm, + norm_kwargs=norm_kwargs, + ) + self.causal = causal + self.trim_right_ratio = trim_right_ratio + assert ( + self.causal or self.trim_right_ratio == 1.0 + ), "`trim_right_ratio` != 1.0 only makes sense for causal convolutions" + assert self.trim_right_ratio >= 0.0 and self.trim_right_ratio <= 1.0 + + def forward(self, x): + kernel_size = self.convtr.convtr.kernel_size[0] + stride = self.convtr.convtr.stride[0] + padding_total = kernel_size - stride + + y = self.convtr(x) + + # We will only trim fixed padding. Extra padding from `pad_for_conv1d` would be + # removed at the very end, when keeping only the right length for the output, + # as removing it here would require also passing the length at the matching layer + # in the encoder. + if self.causal: + # Trim the padding on the right according to the specified ratio + # if trim_right_ratio = 1.0, trim everything from right + padding_right = math.ceil(padding_total * self.trim_right_ratio) + padding_left = padding_total - padding_right + y = unpad1d(y, (padding_left, padding_right)) + else: + # Asymmetric padding required for odd strides + padding_right = padding_total // 2 + padding_left = padding_total - padding_right + y = unpad1d(y, (padding_left, padding_right)) + return y diff --git a/models/codec/speechtokenizer/modules/lstm.py b/models/codec/speechtokenizer/modules/lstm.py new file mode 100644 index 0000000..7f7e431 --- /dev/null +++ b/models/codec/speechtokenizer/modules/lstm.py @@ -0,0 +1,46 @@ +# Copyright (c) 2023 Amphion. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. +# This source file is copied from https://github.com/facebookresearch/encodec + +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +"""LSTM layers module.""" + +from torch import nn + + +class SLSTM(nn.Module): + """ + LSTM without worrying about the hidden state, nor the layout of the data. + Expects input as convolutional layout. + """ + + def __init__( + self, + dimension: int, + num_layers: int = 2, + skip: bool = True, + bidirectional: bool = False, + ): + super().__init__() + self.bidirectional = bidirectional + self.skip = skip + self.lstm = nn.LSTM( + dimension, dimension, num_layers, bidirectional=bidirectional + ) + + def forward(self, x): + x = x.permute(2, 0, 1) + y, _ = self.lstm(x) + if self.bidirectional: + x = x.repeat(1, 1, 2) + if self.skip: + y = y + x + y = y.permute(1, 2, 0) + return y diff --git a/models/codec/speechtokenizer/modules/norm.py b/models/codec/speechtokenizer/modules/norm.py new file mode 100644 index 0000000..ff5eaef --- /dev/null +++ b/models/codec/speechtokenizer/modules/norm.py @@ -0,0 +1,37 @@ +# Copyright (c) 2023 Amphion. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. +# This source file is copied from https://github.com/facebookresearch/encodec + +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +"""Normalization modules.""" + +import typing as tp + +import einops +import torch +from torch import nn + + +class ConvLayerNorm(nn.LayerNorm): + """ + Convolution-friendly LayerNorm that moves channels to last dimensions + before running the normalization and moves them back to original position right after. + """ + + def __init__( + self, normalized_shape: tp.Union[int, tp.List[int], torch.Size], **kwargs + ): + super().__init__(normalized_shape, **kwargs) + + def forward(self, x): + x = einops.rearrange(x, "b ... t -> b t ...") + x = super().forward(x) + x = einops.rearrange(x, "b t ... -> b ... t") + return diff --git a/models/codec/speechtokenizer/modules/quantization/__init__.py b/models/codec/speechtokenizer/modules/quantization/__init__.py new file mode 100644 index 0000000..79d90a1 --- /dev/null +++ b/models/codec/speechtokenizer/modules/quantization/__init__.py @@ -0,0 +1,14 @@ +# Copyright (c) 2023 Amphion. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. +# This source file is copied from https://github.com/facebookresearch/encodec + +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +# flake8: noqa +from .vq import QuantizedResult, ResidualVectorQuantizer diff --git a/models/codec/speechtokenizer/modules/quantization/ac.py b/models/codec/speechtokenizer/modules/quantization/ac.py new file mode 100644 index 0000000..5695ea8 --- /dev/null +++ b/models/codec/speechtokenizer/modules/quantization/ac.py @@ -0,0 +1,317 @@ +# Copyright (c) 2023 Amphion. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. +# This source file is copied from https://github.com/facebookresearch/encodec + +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +"""Arithmetic coder.""" + +import io +import math +import random +import typing as tp +import torch + +from ..binary import BitPacker, BitUnpacker + + +def build_stable_quantized_cdf( + pdf: torch.Tensor, + total_range_bits: int, + roundoff: float = 1e-8, + min_range: int = 2, + check: bool = True, +) -> torch.Tensor: + """Turn the given PDF into a quantized CDF that splits + [0, 2 ** self.total_range_bits - 1] into chunks of size roughly proportional + to the PDF. + + Args: + pdf (torch.Tensor): probability distribution, shape should be `[N]`. + total_range_bits (int): see `ArithmeticCoder`, the typical range we expect + during the coding process is `[0, 2 ** total_range_bits - 1]`. + roundoff (float): will round the pdf up to that level to remove difference coming + from e.g. evaluating the Language Model on different architectures. + min_range (int): minimum range width. Should always be at least 2 for numerical + stability. Use this to avoid pathological behavior is a value + that is expected to be rare actually happens in real life. + check (bool): if True, checks that nothing bad happened, can be deactivated for speed. + """ + pdf = pdf.detach() + if roundoff: + pdf = (pdf / roundoff).floor() * roundoff + # interpolate with uniform distribution to achieve desired minimum probability. + total_range = 2**total_range_bits + cardinality = len(pdf) + alpha = min_range * cardinality / total_range + assert alpha <= 1, "you must reduce min_range" + ranges = (((1 - alpha) * total_range) * pdf).floor().long() + ranges += min_range + quantized_cdf = torch.cumsum(ranges, dim=-1) + if min_range < 2: + raise ValueError("min_range must be at least 2.") + if check: + assert quantized_cdf[-1] <= 2**total_range_bits, quantized_cdf[-1] + if ( + (quantized_cdf[1:] - quantized_cdf[:-1]) < min_range + ).any() or quantized_cdf[0] < min_range: + raise ValueError("You must increase your total_range_bits.") + return quantized_cdf + + +class ArithmeticCoder: + """ArithmeticCoder, + Let us take a distribution `p` over `N` symbols, and assume we have a stream + of random variables `s_t` sampled from `p`. Let us assume that we have a budget + of `B` bits that we can afford to write on device. There are `2**B` possible numbers, + corresponding to the range `[0, 2 ** B - 1]`. We can map each of those number to a single + sequence `(s_t)` by doing the following: + + 1) Initialize the current range to` [0 ** 2 B - 1]`. + 2) For each time step t, split the current range into contiguous chunks, + one for each possible outcome, with size roughly proportional to `p`. + For instance, if `p = [0.75, 0.25]`, and the range is `[0, 3]`, the chunks + would be `{[0, 2], [3, 3]}`. + 3) Select the chunk corresponding to `s_t`, and replace the current range with this. + 4) When done encoding all the values, just select any value remaining in the range. + + You will notice that this procedure can fail: for instance if at any point in time + the range is smaller than `N`, then we can no longer assign a non-empty chunk to each + possible outcome. Intuitively, the more likely a value is, the less the range width + will reduce, and the longer we can go on encoding values. This makes sense: for any efficient + coding scheme, likely outcomes would take less bits, and more of them can be coded + with a fixed budget. + + In practice, we do not know `B` ahead of time, but we have a way to inject new bits + when the current range decreases below a given limit (given by `total_range_bits`), without + having to redo all the computations. If we encode mostly likely values, we will seldom + need to inject new bits, but a single rare value can deplete our stock of entropy! + + In this explanation, we assumed that the distribution `p` was constant. In fact, the present + code works for any sequence `(p_t)` possibly different for each timestep. + We also assume that `s_t ~ p_t`, but that doesn't need to be true, although the smaller + the KL between the true distribution and `p_t`, the most efficient the coding will be. + + Args: + fo (IO[bytes]): file-like object to which the bytes will be written to. + total_range_bits (int): the range `M` described above is `2 ** total_range_bits. + Any time the current range width fall under this limit, new bits will + be injected to rescale the initial range. + """ + + def __init__(self, fo: tp.IO[bytes], total_range_bits: int = 24): + assert total_range_bits <= 30 + self.total_range_bits = total_range_bits + self.packer = BitPacker(bits=1, fo=fo) # we push single bits at a time. + self.low: int = 0 + self.high: int = 0 + self.max_bit: int = -1 + self._dbg: tp.List[tp.Any] = [] + self._dbg2: tp.List[tp.Any] = [] + + @property + def delta(self) -> int: + """Return the current range width.""" + return self.high - self.low + 1 + + def _flush_common_prefix(self): + # If self.low and self.high start with the sames bits, + # those won't change anymore as we always just increase the range + # by powers of 2, and we can flush them out to the bit stream. + assert self.high >= self.low, (self.low, self.high) + assert self.high < 2 ** (self.max_bit + 1) + while self.max_bit >= 0: + b1 = self.low >> self.max_bit + b2 = self.high >> self.max_bit + if b1 == b2: + self.low -= b1 << self.max_bit + self.high -= b1 << self.max_bit + assert self.high >= self.low, (self.high, self.low, self.max_bit) + assert self.low >= 0 + self.max_bit -= 1 + self.packer.push(b1) + else: + break + + def push(self, symbol: int, quantized_cdf: torch.Tensor): + """Push the given symbol on the stream, flushing out bits + if possible. + + Args: + symbol (int): symbol to encode with the AC. + quantized_cdf (torch.Tensor): use `build_stable_quantized_cdf` + to build this from your pdf estimate. + """ + while self.delta < 2**self.total_range_bits: + self.low *= 2 + self.high = self.high * 2 + 1 + self.max_bit += 1 + + range_low = 0 if symbol == 0 else quantized_cdf[symbol - 1].item() + range_high = quantized_cdf[symbol].item() - 1 + effective_low = int( + math.ceil(range_low * (self.delta / (2**self.total_range_bits))) + ) + effective_high = int( + math.floor(range_high * (self.delta / (2**self.total_range_bits))) + ) + assert self.low <= self.high + self.high = self.low + effective_high + self.low = self.low + effective_low + assert self.low <= self.high, ( + effective_low, + effective_high, + range_low, + range_high, + ) + self._dbg.append((self.low, self.high)) + self._dbg2.append((self.low, self.high)) + outs = self._flush_common_prefix() + assert self.low <= self.high + assert self.max_bit >= -1 + assert self.max_bit <= 61, self.max_bit + return outs + + def flush(self): + """Flush the remaining information to the stream.""" + while self.max_bit >= 0: + b1 = (self.low >> self.max_bit) & 1 + self.packer.push(b1) + self.max_bit -= 1 + self.packer.flush() + + +class ArithmeticDecoder: + """ArithmeticDecoder, see `ArithmeticCoder` for a detailed explanation. + + Note that this must be called with **exactly** the same parameters and sequence + of quantized cdf as the arithmetic encoder or the wrong values will be decoded. + + If the AC encoder current range is [L, H], with `L` and `H` having the some common + prefix (i.e. the same most significant bits), then this prefix will be flushed to the stream. + For instances, having read 3 bits `b1 b2 b3`, we know that `[L, H]` is contained inside + `[b1 b2 b3 0 ... 0 b1 b3 b3 1 ... 1]`. Now this specific sub-range can only be obtained + for a specific sequence of symbols and a binary-search allows us to decode those symbols. + At some point, the prefix `b1 b2 b3` will no longer be sufficient to decode new symbols, + and we will need to read new bits from the stream and repeat the process. + + """ + + def __init__(self, fo: tp.IO[bytes], total_range_bits: int = 24): + self.total_range_bits = total_range_bits + self.low: int = 0 + self.high: int = 0 + self.current: int = 0 + self.max_bit: int = -1 + self.unpacker = BitUnpacker(bits=1, fo=fo) # we pull single bits at a time. + # Following is for debugging + self._dbg: tp.List[tp.Any] = [] + self._dbg2: tp.List[tp.Any] = [] + self._last: tp.Any = None + + @property + def delta(self) -> int: + return self.high - self.low + 1 + + def _flush_common_prefix(self): + # Given the current range [L, H], if both have a common prefix, + # we know we can remove it from our representation to avoid handling large numbers. + while self.max_bit >= 0: + b1 = self.low >> self.max_bit + b2 = self.high >> self.max_bit + if b1 == b2: + self.low -= b1 << self.max_bit + self.high -= b1 << self.max_bit + self.current -= b1 << self.max_bit + assert self.high >= self.low + assert self.low >= 0 + self.max_bit -= 1 + else: + break + + def pull(self, quantized_cdf: torch.Tensor) -> tp.Optional[int]: + """Pull a symbol, reading as many bits from the stream as required. + This returns `None` when the stream has been exhausted. + + Args: + quantized_cdf (torch.Tensor): use `build_stable_quantized_cdf` + to build this from your pdf estimate. This must be **exatly** + the same cdf as the one used at encoding time. + """ + while self.delta < 2**self.total_range_bits: + bit = self.unpacker.pull() + if bit is None: + return None + self.low *= 2 + self.high = self.high * 2 + 1 + self.current = self.current * 2 + bit + self.max_bit += 1 + + def bin_search(low_idx: int, high_idx: int): + # Binary search is not just for coding interviews :) + if high_idx < low_idx: + raise RuntimeError("Binary search failed") + mid = (low_idx + high_idx) // 2 + range_low = quantized_cdf[mid - 1].item() if mid > 0 else 0 + range_high = quantized_cdf[mid].item() - 1 + effective_low = int( + math.ceil(range_low * (self.delta / (2**self.total_range_bits))) + ) + effective_high = int( + math.floor(range_high * (self.delta / (2**self.total_range_bits))) + ) + low = effective_low + self.low + high = effective_high + self.low + if self.current >= low: + if self.current <= high: + return (mid, low, high, self.current) + else: + return bin_search(mid + 1, high_idx) + else: + return bin_search(low_idx, mid - 1) + + self._last = (self.low, self.high, self.current, self.max_bit) + sym, self.low, self.high, self.current = bin_search(0, len(quantized_cdf) - 1) + self._dbg.append((self.low, self.high, self.current)) + self._flush_common_prefix() + self._dbg2.append((self.low, self.high, self.current)) + + return sym + + +def test(): + torch.manual_seed(1234) + random.seed(1234) + for _ in range(4): + pdfs = [] + cardinality = random.randrange(4000) + steps = random.randrange(100, 500) + fo = io.BytesIO() + encoder = ArithmeticCoder(fo) + symbols = [] + for step in range(steps): + pdf = torch.softmax(torch.randn(cardinality), dim=0) + pdfs.append(pdf) + q_cdf = build_stable_quantized_cdf(pdf, encoder.total_range_bits) + symbol = torch.multinomial(pdf, 1).item() + symbols.append(symbol) + encoder.push(symbol, q_cdf) + encoder.flush() + + fo.seek(0) + decoder = ArithmeticDecoder(fo) + for idx, (pdf, symbol) in enumerate(zip(pdfs, symbols)): + q_cdf = build_stable_quantized_cdf(pdf, encoder.total_range_bits) + decoded_symbol = decoder.pull(q_cdf) + assert decoded_symbol == symbol, idx + assert decoder.pull(torch.zeros(1)) is None + + +if __name__ == "__main__": + test() diff --git a/models/codec/speechtokenizer/modules/quantization/core_vq.py b/models/codec/speechtokenizer/modules/quantization/core_vq.py new file mode 100644 index 0000000..5799725 --- /dev/null +++ b/models/codec/speechtokenizer/modules/quantization/core_vq.py @@ -0,0 +1,388 @@ +# Copyright (c) 2023 Amphion. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. +# This source file is copied from https://github.com/facebookresearch/encodec + +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. +# +# This implementation is inspired from +# https://github.com/lucidrains/vector-quantize-pytorch +# which is released under MIT License. Hereafter, the original license: +# MIT License +# +# Copyright (c) 2020 Phil Wang +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. + +"""Core vector quantization implementation.""" +import typing as tp + +from einops import rearrange, repeat +import torch +from torch import nn +import torch.nn.functional as F + +from .distrib import broadcast_tensors, rank + + +def default(val: tp.Any, d: tp.Any) -> tp.Any: + return val if val is not None else d + + +def ema_inplace(moving_avg, new, decay: float): + moving_avg.data.mul_(decay).add_(new, alpha=(1 - decay)) + + +def laplace_smoothing(x, n_categories: int, epsilon: float = 1e-5): + return (x + epsilon) / (x.sum() + n_categories * epsilon) + + +def uniform_init(*shape: int): + t = torch.empty(shape) + nn.init.kaiming_uniform_(t) + return t + + +def sample_vectors(samples, num: int): + num_samples, device = samples.shape[0], samples.device + + if num_samples >= num: + indices = torch.randperm(num_samples, device=device)[:num] + else: + indices = torch.randint(0, num_samples, (num,), device=device) + + return samples[indices] + + +def kmeans(samples, num_clusters: int, num_iters: int = 10): + dim, dtype = samples.shape[-1], samples.dtype + + means = sample_vectors(samples, num_clusters) + + for _ in range(num_iters): + diffs = rearrange(samples, "n d -> n () d") - rearrange(means, "c d -> () c d") + dists = -(diffs**2).sum(dim=-1) + + buckets = dists.max(dim=-1).indices + bins = torch.bincount(buckets, minlength=num_clusters) + zero_mask = bins == 0 + bins_min_clamped = bins.masked_fill(zero_mask, 1) + + new_means = buckets.new_zeros(num_clusters, dim, dtype=dtype) + new_means.scatter_add_(0, repeat(buckets, "n -> n d", d=dim), samples) + new_means = new_means / bins_min_clamped[..., None] + + means = torch.where(zero_mask[..., None], means, new_means) + + return means, bins + + +class EuclideanCodebook(nn.Module): + """Codebook with Euclidean distance. + Args: + dim (int): Dimension. + codebook_size (int): Codebook size. + kmeans_init (bool): Whether to use k-means to initialize the codebooks. + If set to true, run the k-means algorithm on the first training batch and use + the learned centroids as initialization. + kmeans_iters (int): Number of iterations used for k-means algorithm at initialization. + decay (float): Decay for exponential moving average over the codebooks. + epsilon (float): Epsilon value for numerical stability. + threshold_ema_dead_code (int): Threshold for dead code expiration. Replace any codes + that have an exponential moving average cluster size less than the specified threshold with + randomly selected vector from the current batch. + """ + + def __init__( + self, + dim: int, + codebook_size: int, + kmeans_init: int = False, + kmeans_iters: int = 10, + decay: float = 0.99, + epsilon: float = 1e-5, + threshold_ema_dead_code: int = 2, + ): + super().__init__() + self.decay = decay + init_fn: tp.Union[tp.Callable[..., torch.Tensor], tp.Any] = ( + uniform_init if not kmeans_init else torch.zeros + ) + embed = init_fn(codebook_size, dim) + + self.codebook_size = codebook_size + + self.kmeans_iters = kmeans_iters + self.epsilon = epsilon + self.threshold_ema_dead_code = threshold_ema_dead_code + + self.register_buffer("inited", torch.Tensor([not kmeans_init])) + self.register_buffer("cluster_size", torch.zeros(codebook_size)) + self.register_buffer("embed", embed) + self.register_buffer("embed_avg", embed.clone()) + + @torch.jit.ignore + def init_embed_(self, data): + if self.inited: + return + + embed, cluster_size = kmeans(data, self.codebook_size, self.kmeans_iters) + self.embed.data.copy_(embed) + self.embed_avg.data.copy_(embed.clone()) + self.cluster_size.data.copy_(cluster_size) + self.inited.data.copy_(torch.Tensor([True])) + # Make sure all buffers across workers are in sync after initialization + # broadcast_tensors(self.buffers()) + + def replace_(self, samples, mask): + modified_codebook = torch.where( + mask[..., None], sample_vectors(samples, self.codebook_size), self.embed + ) + self.embed.data.copy_(modified_codebook) + + def expire_codes_(self, batch_samples): + if self.threshold_ema_dead_code == 0: + return + + expired_codes = self.cluster_size < self.threshold_ema_dead_code + if not torch.any(expired_codes): + return + + batch_samples = rearrange(batch_samples, "... d -> (...) d") + self.replace_(batch_samples, mask=expired_codes) + # broadcast_tensors(self.buffers()) + + def preprocess(self, x): + x = rearrange(x, "... d -> (...) d") + return x + + def quantize(self, x): + embed = self.embed.t() + dist = -( + x.pow(2).sum(1, keepdim=True) + - 2 * x @ embed + + embed.pow(2).sum(0, keepdim=True) + ) + embed_ind = dist.max(dim=-1).indices + return embed_ind + + def postprocess_emb(self, embed_ind, shape): + return embed_ind.view(*shape[:-1]) + + def dequantize(self, embed_ind): + quantize = F.embedding(embed_ind, self.embed) + return quantize + + def encode(self, x): + shape = x.shape + # pre-process + x = self.preprocess(x) + # quantize + embed_ind = self.quantize(x) + # post-process + embed_ind = self.postprocess_emb(embed_ind, shape) + return embed_ind + + def decode(self, embed_ind): + quantize = self.dequantize(embed_ind) + return quantize + + def forward(self, x): + shape, dtype = x.shape, x.dtype + x = self.preprocess(x) + + self.init_embed_(x) + + embed_ind = self.quantize(x) + embed_onehot = F.one_hot(embed_ind, self.codebook_size).type(dtype) + embed_ind = self.postprocess_emb(embed_ind, shape) + quantize = self.dequantize(embed_ind) + + if self.training: + # We do the expiry of code at that point as buffers are in sync + # and all the workers will take the same decision. + self.expire_codes_(x) + ema_inplace(self.cluster_size, embed_onehot.sum(0), self.decay) + embed_sum = x.t() @ embed_onehot + ema_inplace(self.embed_avg, embed_sum.t(), self.decay) + cluster_size = ( + laplace_smoothing(self.cluster_size, self.codebook_size, self.epsilon) + * self.cluster_size.sum() + ) + embed_normalized = self.embed_avg / cluster_size.unsqueeze(1) + self.embed.data.copy_(embed_normalized) + + return quantize, embed_ind + + +class VectorQuantization(nn.Module): + """Vector quantization implementation. + Currently supports only euclidean distance. + Args: + dim (int): Dimension + codebook_size (int): Codebook size + codebook_dim (int): Codebook dimension. If not defined, uses the specified dimension in dim. + decay (float): Decay for exponential moving average over the codebooks. + epsilon (float): Epsilon value for numerical stability. + kmeans_init (bool): Whether to use kmeans to initialize the codebooks. + kmeans_iters (int): Number of iterations used for kmeans initialization. + threshold_ema_dead_code (int): Threshold for dead code expiration. Replace any codes + that have an exponential moving average cluster size less than the specified threshold with + randomly selected vector from the current batch. + commitment_weight (float): Weight for commitment loss. + """ + + def __init__( + self, + dim: int, + codebook_size: int, + codebook_dim: tp.Optional[int] = None, + decay: float = 0.99, + epsilon: float = 1e-5, + kmeans_init: bool = True, + kmeans_iters: int = 50, + threshold_ema_dead_code: int = 2, + commitment_weight: float = 1.0, + ): + super().__init__() + _codebook_dim: int = default(codebook_dim, dim) + + requires_projection = _codebook_dim != dim + self.project_in = ( + nn.Linear(dim, _codebook_dim) if requires_projection else nn.Identity() + ) + self.project_out = ( + nn.Linear(_codebook_dim, dim) if requires_projection else nn.Identity() + ) + + self.epsilon = epsilon + self.commitment_weight = commitment_weight + + self._codebook = EuclideanCodebook( + dim=_codebook_dim, + codebook_size=codebook_size, + kmeans_init=kmeans_init, + kmeans_iters=kmeans_iters, + decay=decay, + epsilon=epsilon, + threshold_ema_dead_code=threshold_ema_dead_code, + ) + self.codebook_size = codebook_size + + @property + def codebook(self): + return self._codebook.embed + + def encode(self, x): + x = rearrange(x, "b d n -> b n d") + x = self.project_in(x) + embed_in = self._codebook.encode(x) + return embed_in + + def decode(self, embed_ind): + quantize = self._codebook.decode(embed_ind) + quantize = self.project_out(quantize) + quantize = rearrange(quantize, "b n d -> b d n") + return quantize + + def forward(self, x): + device = x.device + x = rearrange(x, "b d n -> b n d") + x = self.project_in(x) + + quantize, embed_ind = self._codebook(x) + + if self.training: + quantize = x + (quantize - x).detach() + + loss = torch.tensor([0.0], device=device, requires_grad=self.training) + + if self.training: + if self.commitment_weight > 0: + commit_loss = F.mse_loss(quantize.detach(), x) + loss = loss + commit_loss * self.commitment_weight + + quantize = self.project_out(quantize) + quantize = rearrange(quantize, "b n d -> b d n") + return quantize, embed_ind, loss + + +class ResidualVectorQuantization(nn.Module): + """Residual vector quantization implementation. + Follows Algorithm 1. in https://arxiv.org/pdf/2107.03312.pdf + """ + + def __init__(self, *, num_quantizers, **kwargs): + super().__init__() + self.layers = nn.ModuleList( + [VectorQuantization(**kwargs) for _ in range(num_quantizers)] + ) + + def forward( + self, x, n_q: tp.Optional[int] = None, layers: tp.Optional[list] = None + ): + quantized_out = 0.0 + residual = x + + all_losses = [] + all_indices = [] + out_quantized = [] + + n_q = n_q or len(self.layers) + + for i, layer in enumerate(self.layers[:n_q]): + quantized, indices, loss = layer(residual) + residual = residual - quantized + quantized_out = quantized_out + quantized + + all_indices.append(indices) + all_losses.append(loss) + if layers and i in layers: + out_quantized.append(quantized) + + out_losses, out_indices = map(torch.stack, (all_losses, all_indices)) + return quantized_out, out_indices, out_losses, out_quantized + + def encode( + self, x: torch.Tensor, n_q: tp.Optional[int] = None, st: tp.Optional[int] = None + ) -> torch.Tensor: + residual = x + all_indices = [] + n_q = n_q or len(self.layers) + st = st or 0 + for layer in self.layers[st:n_q]: + indices = layer.encode(residual) + quantized = layer.decode(indices) + residual = residual - quantized + all_indices.append(indices) + out_indices = torch.stack(all_indices) + return out_indices + + def decode(self, q_indices: torch.Tensor, st: int = 0) -> torch.Tensor: + quantized_out = torch.tensor(0.0, device=q_indices.device) + for i, indices in enumerate(q_indices): + layer = self.layers[st + i] + quantized = layer.decode(indices) + quantized_out = quantized_out + quantized + return quantized_out diff --git a/models/codec/speechtokenizer/modules/quantization/distrib.py b/models/codec/speechtokenizer/modules/quantization/distrib.py new file mode 100644 index 0000000..7b9a9b8 --- /dev/null +++ b/models/codec/speechtokenizer/modules/quantization/distrib.py @@ -0,0 +1,135 @@ +# Copyright (c) 2023 Amphion. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. +# This source file is copied from https://github.com/facebookresearch/encodec + +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +"""Torch distributed utilities.""" + +import typing as tp + +import torch + + +def rank(): + if torch.distributed.is_initialized(): + return torch.distributed.get_rank() + else: + return 0 + + +def world_size(): + if torch.distributed.is_initialized(): + return torch.distributed.get_world_size() + else: + return 1 + + +def is_distributed(): + return world_size() > 1 + + +def all_reduce(tensor: torch.Tensor, op=torch.distributed.ReduceOp.SUM): + if is_distributed(): + return torch.distributed.all_reduce(tensor, op) + + +def _is_complex_or_float(tensor): + return torch.is_floating_point(tensor) or torch.is_complex(tensor) + + +def _check_number_of_params(params: tp.List[torch.Tensor]): + # utility function to check that the number of params in all workers is the same, + # and thus avoid a deadlock with distributed all reduce. + if not is_distributed() or not params: + return + # print('params[0].device ', params[0].device) + tensor = torch.tensor([len(params)], device=params[0].device, dtype=torch.long) + all_reduce(tensor) + if tensor.item() != len(params) * world_size(): + # If not all the workers have the same number, for at least one of them, + # this inequality will be verified. + raise RuntimeError( + f"Mismatch in number of params: ours is {len(params)}, " + "at least one worker has a different one." + ) + + +def broadcast_tensors(tensors: tp.Iterable[torch.Tensor], src: int = 0): + """Broadcast the tensors from the given parameters to all workers. + This can be used to ensure that all workers have the same model to start with. + """ + if not is_distributed(): + return + tensors = [tensor for tensor in tensors if _is_complex_or_float(tensor)] + _check_number_of_params(tensors) + handles = [] + for tensor in tensors: + # src = int(rank()) # added code + handle = torch.distributed.broadcast(tensor.data, src=src, async_op=True) + handles.append(handle) + for handle in handles: + handle.wait() + + +def sync_buffer(buffers, average=True): + """ + Sync grad for buffers. If average is False, broadcast instead of averaging. + """ + if not is_distributed(): + return + handles = [] + for buffer in buffers: + if torch.is_floating_point(buffer.data): + if average: + handle = torch.distributed.all_reduce( + buffer.data, op=torch.distributed.ReduceOp.SUM, async_op=True + ) + else: + handle = torch.distributed.broadcast(buffer.data, src=0, async_op=True) + handles.append((buffer, handle)) + for buffer, handle in handles: + handle.wait() + if average: + buffer.data /= world_size + + +def sync_grad(params): + """ + Simpler alternative to DistributedDataParallel, that doesn't rely + on any black magic. For simple models it can also be as fast. + Just call this on your model parameters after the call to backward! + """ + if not is_distributed(): + return + handles = [] + for p in params: + if p.grad is not None: + handle = torch.distributed.all_reduce( + p.grad.data, op=torch.distributed.ReduceOp.SUM, async_op=True + ) + handles.append((p, handle)) + for p, handle in handles: + handle.wait() + p.grad.data /= world_size() + + +def average_metrics(metrics: tp.Dict[str, float], count=1.0): + """Average a dictionary of metrics across all workers, using the optional + `count` as unormalized weight. + """ + if not is_distributed(): + return metrics + keys, values = zip(*metrics.items()) + device = "cuda" if torch.cuda.is_available() else "cpu" + tensor = torch.tensor(list(values) + [1], device=device, dtype=torch.float32) + tensor *= count + all_reduce(tensor) + averaged = (tensor[:-1] / tensor[-1]).cpu().tolist() + return dict(zip(keys, averaged)) diff --git a/models/codec/speechtokenizer/modules/quantization/vq.py b/models/codec/speechtokenizer/modules/quantization/vq.py new file mode 100644 index 0000000..ec7df0f --- /dev/null +++ b/models/codec/speechtokenizer/modules/quantization/vq.py @@ -0,0 +1,125 @@ +# Copyright (c) 2023 Amphion. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. +# This source file is copied from https://github.com/facebookresearch/encodec + +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +"""Residual vector quantizer implementation.""" + +from dataclasses import dataclass, field +import math +import typing as tp + +import torch +from torch import nn + +from .core_vq import ResidualVectorQuantization + + +@dataclass +class QuantizedResult: + quantized: torch.Tensor + codes: torch.Tensor + bandwidth: torch.Tensor # bandwidth in kb/s used, per batch item. + penalty: tp.Optional[torch.Tensor] = None + metrics: dict = field(default_factory=dict) + + +class ResidualVectorQuantizer(nn.Module): + """Residual Vector Quantizer. + Args: + dimension (int): Dimension of the codebooks. + n_q (int): Number of residual vector quantizers used. + bins (int): Codebook size. + decay (float): Decay for exponential moving average over the codebooks. + kmeans_init (bool): Whether to use kmeans to initialize the codebooks. + kmeans_iters (int): Number of iterations used for kmeans initialization. + threshold_ema_dead_code (int): Threshold for dead code expiration. Replace any codes + that have an exponential moving average cluster size less than the specified threshold with + randomly selected vector from the current batch. + """ + + def __init__( + self, + dimension: int = 256, + n_q: int = 8, + bins: int = 1024, + decay: float = 0.99, + kmeans_init: bool = True, + kmeans_iters: int = 50, + threshold_ema_dead_code: int = 2, + ): + super().__init__() + self.n_q = n_q + self.dimension = dimension + self.bins = bins + self.decay = decay + self.kmeans_init = kmeans_init + self.kmeans_iters = kmeans_iters + self.threshold_ema_dead_code = threshold_ema_dead_code + self.vq = ResidualVectorQuantization( + dim=self.dimension, + codebook_size=self.bins, + num_quantizers=self.n_q, + decay=self.decay, + kmeans_init=self.kmeans_init, + kmeans_iters=self.kmeans_iters, + threshold_ema_dead_code=self.threshold_ema_dead_code, + ) + + def forward( + self, + x: torch.Tensor, + n_q: tp.Optional[int] = None, + layers: tp.Optional[list] = None, + ) -> QuantizedResult: + """Residual vector quantization on the given input tensor. + Args: + x (torch.Tensor): Input tensor. + n_q (int): Number of quantizer used to quantize. Default: All quantizers. + layers (list): Layer that need to return quantized. Defalt: None. + Returns: + QuantizedResult: + The quantized (or approximately quantized) representation with + the associated numbert quantizers and layer quantized required to return. + """ + n_q = n_q if n_q else self.n_q + if layers and max(layers) >= n_q: + raise ValueError( + f"Last layer index in layers: A {max(layers)}. Number of quantizers in RVQ: B {self.n_q}. A must less than B." + ) + quantized, codes, commit_loss, quantized_list = self.vq( + x, n_q=n_q, layers=layers + ) + return quantized, codes, torch.mean(commit_loss), quantized_list + + def encode( + self, x: torch.Tensor, n_q: tp.Optional[int] = None, st: tp.Optional[int] = None + ) -> torch.Tensor: + """Encode a given input tensor with the specified sample rate at the given bandwidth. + The RVQ encode method sets the appropriate number of quantizer to use + and returns indices for each quantizer. + Args: + x (torch.Tensor): Input tensor. + n_q (int): Number of quantizer used to quantize. Default: All quantizers. + st (int): Start to encode input from which layers. Default: 0. + """ + n_q = n_q if n_q else self.n_q + st = st or 0 + codes = self.vq.encode(x, n_q=n_q, st=st) + return codes + + def decode(self, codes: torch.Tensor, st: int = 0) -> torch.Tensor: + """Decode the given codes to the quantized representation. + Args: + codes (torch.Tensor): Input indices for each quantizer. + st (int): Start to decode input codes from which layers. Default: 0. + """ + quantized = self.vq.decode(codes, st=st) + return quantized diff --git a/models/codec/speechtokenizer/modules/seanet.py b/models/codec/speechtokenizer/modules/seanet.py new file mode 100644 index 0000000..481de20 --- /dev/null +++ b/models/codec/speechtokenizer/modules/seanet.py @@ -0,0 +1,414 @@ +# Copyright (c) 2023 Amphion. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. +# This source file is copied from https://github.com/facebookresearch/encodec + +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +"""Encodec SEANet-based encoder and decoder implementation.""" + +import typing as tp + +import numpy as np +import torch.nn as nn +import torch + +from . import SConv1d, SConvTranspose1d, SLSTM + + +@torch.jit.script +def snake(x, alpha): + shape = x.shape + x = x.reshape(shape[0], shape[1], -1) + x = x + (alpha + 1e-9).reciprocal() * torch.sin(alpha * x).pow(2) + x = x.reshape(shape) + return x + + +class Snake1d(nn.Module): + def __init__(self, channels): + super().__init__() + self.alpha = nn.Parameter(torch.ones(1, channels, 1)) + + def forward(self, x): + return snake(x, self.alpha) + + +class SEANetResnetBlock(nn.Module): + """Residual block from SEANet model. + Args: + dim (int): Dimension of the input/output + kernel_sizes (list): List of kernel sizes for the convolutions. + dilations (list): List of dilations for the convolutions. + activation (str): Activation function. + activation_params (dict): Parameters to provide to the activation function + norm (str): Normalization method. + norm_params (dict): Parameters to provide to the underlying normalization used along with the convolution. + causal (bool): Whether to use fully causal convolution. + pad_mode (str): Padding mode for the convolutions. + compress (int): Reduced dimensionality in residual branches (from Demucs v3) + true_skip (bool): Whether to use true skip connection or a simple convolution as the skip connection. + """ + + def __init__( + self, + dim: int, + kernel_sizes: tp.List[int] = [3, 1], + dilations: tp.List[int] = [1, 1], + activation: str = "ELU", + activation_params: dict = {"alpha": 1.0}, + norm: str = "weight_norm", + norm_params: tp.Dict[str, tp.Any] = {}, + causal: bool = False, + pad_mode: str = "reflect", + compress: int = 2, + true_skip: bool = True, + ): + super().__init__() + assert len(kernel_sizes) == len( + dilations + ), "Number of kernel sizes should match number of dilations" + act = getattr(nn, activation) if activation != "Snake" else Snake1d + hidden = dim // compress + block = [] + for i, (kernel_size, dilation) in enumerate(zip(kernel_sizes, dilations)): + in_chs = dim if i == 0 else hidden + out_chs = dim if i == len(kernel_sizes) - 1 else hidden + block += [ + act(**activation_params) if activation != "Snake" else act(in_chs), + SConv1d( + in_chs, + out_chs, + kernel_size=kernel_size, + dilation=dilation, + norm=norm, + norm_kwargs=norm_params, + causal=causal, + pad_mode=pad_mode, + ), + ] + self.block = nn.Sequential(*block) + self.shortcut: nn.Module + if true_skip: + self.shortcut = nn.Identity() + else: + self.shortcut = SConv1d( + dim, + dim, + kernel_size=1, + norm=norm, + norm_kwargs=norm_params, + causal=causal, + pad_mode=pad_mode, + ) + + def forward(self, x): + return self.shortcut(x) + self.block(x) + + +class SEANetEncoder(nn.Module): + """SEANet encoder. + Args: + channels (int): Audio channels. + dimension (int): Intermediate representation dimension. + n_filters (int): Base width for the model. + n_residual_layers (int): nb of residual layers. + ratios (Sequence[int]): kernel size and stride ratios. The encoder uses downsampling ratios instead of + upsampling ratios, hence it will use the ratios in the reverse order to the ones specified here + that must match the decoder order + activation (str): Activation function. + activation_params (dict): Parameters to provide to the activation function + norm (str): Normalization method. + norm_params (dict): Parameters to provide to the underlying normalization used along with the convolution. + kernel_size (int): Kernel size for the initial convolution. + last_kernel_size (int): Kernel size for the initial convolution. + residual_kernel_size (int): Kernel size for the residual layers. + dilation_base (int): How much to increase the dilation with each layer. + causal (bool): Whether to use fully causal convolution. + pad_mode (str): Padding mode for the convolutions. + true_skip (bool): Whether to use true skip connection or a simple + (streamable) convolution as the skip connection in the residual network blocks. + compress (int): Reduced dimensionality in residual branches (from Demucs v3). + lstm (int): Number of LSTM layers at the end of the encoder. + """ + + def __init__( + self, + channels: int = 1, + dimension: int = 128, + n_filters: int = 32, + n_residual_layers: int = 1, + ratios: tp.List[int] = [8, 5, 4, 2], + activation: str = "ELU", + activation_params: dict = {"alpha": 1.0}, + norm: str = "weight_norm", + norm_params: tp.Dict[str, tp.Any] = {}, + kernel_size: int = 7, + last_kernel_size: int = 7, + residual_kernel_size: int = 3, + dilation_base: int = 2, + causal: bool = False, + pad_mode: str = "reflect", + true_skip: bool = False, + compress: int = 2, + lstm: int = 2, + bidirectional: bool = False, + ): + super().__init__() + self.channels = channels + self.dimension = dimension + self.n_filters = n_filters + self.ratios = list(reversed(ratios)) + del ratios + self.n_residual_layers = n_residual_layers + self.hop_length = np.prod(self.ratios) # 计算乘积 + + act = getattr(nn, activation) if activation != "Snake" else Snake1d + mult = 1 + model: tp.List[nn.Module] = [ + SConv1d( + channels, + mult * n_filters, + kernel_size, + norm=norm, + norm_kwargs=norm_params, + causal=causal, + pad_mode=pad_mode, + ) + ] + # Downsample to raw audio scale + for i, ratio in enumerate(self.ratios): + # Add residual layers + for j in range(n_residual_layers): + model += [ + SEANetResnetBlock( + mult * n_filters, + kernel_sizes=[residual_kernel_size, 1], + dilations=[dilation_base**j, 1], + norm=norm, + norm_params=norm_params, + activation=activation, + activation_params=activation_params, + causal=causal, + pad_mode=pad_mode, + compress=compress, + true_skip=true_skip, + ) + ] + + # Add downsampling layers + model += [ + ( + act(**activation_params) + if activation != "Snake" + else act(mult * n_filters) + ), + SConv1d( + mult * n_filters, + mult * n_filters * 2, + kernel_size=ratio * 2, + stride=ratio, + norm=norm, + norm_kwargs=norm_params, + causal=causal, + pad_mode=pad_mode, + ), + ] + mult *= 2 + + if lstm: + model += [ + SLSTM(mult * n_filters, num_layers=lstm, bidirectional=bidirectional) + ] + + mult = mult * 2 if bidirectional else mult + model += [ + ( + act(**activation_params) + if activation != "Snake" + else act(mult * n_filters) + ), + SConv1d( + mult * n_filters, + dimension, + last_kernel_size, + norm=norm, + norm_kwargs=norm_params, + causal=causal, + pad_mode=pad_mode, + ), + ] + + self.model = nn.Sequential(*model) + + def forward(self, x): + return self.model(x) + + +class SEANetDecoder(nn.Module): + """SEANet decoder. + Args: + channels (int): Audio channels. + dimension (int): Intermediate representation dimension. + n_filters (int): Base width for the model. + n_residual_layers (int): nb of residual layers. + ratios (Sequence[int]): kernel size and stride ratios + activation (str): Activation function. + activation_params (dict): Parameters to provide to the activation function + final_activation (str): Final activation function after all convolutions. + final_activation_params (dict): Parameters to provide to the activation function + norm (str): Normalization method. + norm_params (dict): Parameters to provide to the underlying normalization used along with the convolution. + kernel_size (int): Kernel size for the initial convolution. + last_kernel_size (int): Kernel size for the initial convolution. + residual_kernel_size (int): Kernel size for the residual layers. + dilation_base (int): How much to increase the dilation with each layer. + causal (bool): Whether to use fully causal convolution. + pad_mode (str): Padding mode for the convolutions. + true_skip (bool): Whether to use true skip connection or a simple + (streamable) convolution as the skip connection in the residual network blocks. + compress (int): Reduced dimensionality in residual branches (from Demucs v3). + lstm (int): Number of LSTM layers at the end of the encoder. + trim_right_ratio (float): Ratio for trimming at the right of the transposed convolution under the causal setup. + If equal to 1.0, it means that all the trimming is done at the right. + """ + + def __init__( + self, + channels: int = 1, + dimension: int = 128, + n_filters: int = 32, + n_residual_layers: int = 1, + ratios: tp.List[int] = [8, 5, 4, 2], + activation: str = "ELU", + activation_params: dict = {"alpha": 1.0}, + final_activation: tp.Optional[str] = None, + final_activation_params: tp.Optional[dict] = None, + norm: str = "weight_norm", + norm_params: tp.Dict[str, tp.Any] = {}, + kernel_size: int = 7, + last_kernel_size: int = 7, + residual_kernel_size: int = 3, + dilation_base: int = 2, + causal: bool = False, + pad_mode: str = "reflect", + true_skip: bool = False, + compress: int = 2, + lstm: int = 2, + trim_right_ratio: float = 1.0, + bidirectional: bool = False, + ): + super().__init__() + self.dimension = dimension + self.channels = channels + self.n_filters = n_filters + self.ratios = ratios + del ratios + self.n_residual_layers = n_residual_layers + self.hop_length = np.prod(self.ratios) + + act = getattr(nn, activation) if activation != "Snake" else Snake1d + mult = int(2 ** len(self.ratios)) + model: tp.List[nn.Module] = [ + SConv1d( + dimension, + mult * n_filters, + kernel_size, + norm=norm, + norm_kwargs=norm_params, + causal=causal, + pad_mode=pad_mode, + ) + ] + + if lstm: + model += [ + SLSTM(mult * n_filters, num_layers=lstm, bidirectional=bidirectional) + ] + + # Upsample to raw audio scale + for i, ratio in enumerate(self.ratios): + # Add upsampling layers + model += [ + ( + act(**activation_params) + if activation != "Snake" + else act(mult * n_filters) + ), + SConvTranspose1d( + mult * n_filters, + mult * n_filters // 2, + kernel_size=ratio * 2, + stride=ratio, + norm=norm, + norm_kwargs=norm_params, + causal=causal, + trim_right_ratio=trim_right_ratio, + ), + ] + # Add residual layers + for j in range(n_residual_layers): + model += [ + SEANetResnetBlock( + mult * n_filters // 2, + kernel_sizes=[residual_kernel_size, 1], + dilations=[dilation_base**j, 1], + activation=activation, + activation_params=activation_params, + norm=norm, + norm_params=norm_params, + causal=causal, + pad_mode=pad_mode, + compress=compress, + true_skip=true_skip, + ) + ] + + mult //= 2 + + # Add final layers + model += [ + act(**activation_params) if activation != "Snake" else act(n_filters), + SConv1d( + n_filters, + channels, + last_kernel_size, + norm=norm, + norm_kwargs=norm_params, + causal=causal, + pad_mode=pad_mode, + ), + ] + # Add optional final activation to decoder (eg. tanh) + if final_activation is not None: + final_act = getattr(nn, final_activation) + final_activation_params = final_activation_params or {} + model += [final_act(**final_activation_params)] + self.model = nn.Sequential(*model) + + def forward(self, z): + y = self.model(z) + return y + + +def test(): + import torch + + encoder = SEANetEncoder() + decoder = SEANetDecoder() + x = torch.randn(1, 1, 24000) + z = encoder(x) + print("z ", z.shape) + assert 1 == 2 + assert list(z.shape) == [1, 128, 75], z.shape + y = decoder(z) + assert y.shape == x.shape, (x.shape, y.shape) + + +if __name__ == "__main__": + test() diff --git a/models/svc/__init__.py b/models/svc/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/models/svc/base/__init__.py b/models/svc/base/__init__.py new file mode 100644 index 0000000..38c2b16 --- /dev/null +++ b/models/svc/base/__init__.py @@ -0,0 +1,7 @@ +# Copyright (c) 2023 Amphion. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +from .svc_inference import SVCInference +from .svc_trainer import SVCTrainer diff --git a/models/svc/base/svc_dataset.py b/models/svc/base/svc_dataset.py new file mode 100644 index 0000000..d8ef7d5 --- /dev/null +++ b/models/svc/base/svc_dataset.py @@ -0,0 +1,595 @@ +# Copyright (c) 2023 Amphion. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +import random +import torch +from torch.nn.utils.rnn import pad_sequence +import json +import os +import numpy as np +import librosa + +from utils.data_utils import * +from processors.acoustic_extractor import cal_normalized_mel, load_mel_extrema +from processors.content_extractor import ( + ContentvecExtractor, + WhisperExtractor, + WenetExtractor, +) +from models.base.base_dataset import ( + BaseOfflineDataset, + BaseOfflineCollator, + BaseOnlineDataset, + BaseOnlineCollator, +) +from models.base.new_dataset import BaseTestDataset + +EPS = 1.0e-12 + + +class SVCOfflineDataset(BaseOfflineDataset): + def __init__(self, cfg, dataset, is_valid=False): + BaseOfflineDataset.__init__(self, cfg, dataset, is_valid=is_valid) + + cfg = self.cfg + + if cfg.model.condition_encoder.use_whisper: + self.whisper_aligner = WhisperExtractor(self.cfg) + self.utt2whisper_path = load_content_feature_path( + self.metadata, cfg.preprocess.processed_dir, cfg.preprocess.whisper_dir + ) + + if cfg.model.condition_encoder.use_contentvec: + self.contentvec_aligner = ContentvecExtractor(self.cfg) + self.utt2contentVec_path = load_content_feature_path( + self.metadata, + cfg.preprocess.processed_dir, + cfg.preprocess.contentvec_dir, + ) + + if cfg.model.condition_encoder.use_mert: + self.utt2mert_path = load_content_feature_path( + self.metadata, cfg.preprocess.processed_dir, cfg.preprocess.mert_dir + ) + if cfg.model.condition_encoder.use_wenet: + self.wenet_aligner = WenetExtractor(self.cfg) + self.utt2wenet_path = load_content_feature_path( + self.metadata, cfg.preprocess.processed_dir, cfg.preprocess.wenet_dir + ) + + def __getitem__(self, index): + single_feature = BaseOfflineDataset.__getitem__(self, index) + + utt_info = self.metadata[index] + dataset = utt_info["Dataset"] + uid = utt_info["Uid"] + utt = "{}_{}".format(dataset, uid) + + if self.cfg.model.condition_encoder.use_whisper: + assert "target_len" in single_feature.keys() + aligned_whisper_feat = ( + self.whisper_aligner.offline_resolution_transformation( + np.load(self.utt2whisper_path[utt]), single_feature["target_len"] + ) + ) + single_feature["whisper_feat"] = aligned_whisper_feat + + if self.cfg.model.condition_encoder.use_contentvec: + assert "target_len" in single_feature.keys() + aligned_contentvec = ( + self.contentvec_aligner.offline_resolution_transformation( + np.load(self.utt2contentVec_path[utt]), single_feature["target_len"] + ) + ) + single_feature["contentvec_feat"] = aligned_contentvec + + if self.cfg.model.condition_encoder.use_mert: + assert "target_len" in single_feature.keys() + aligned_mert_feat = align_content_feature_length( + np.load(self.utt2mert_path[utt]), + single_feature["target_len"], + source_hop=self.cfg.preprocess.mert_hop_size, + ) + single_feature["mert_feat"] = aligned_mert_feat + + if self.cfg.model.condition_encoder.use_wenet: + assert "target_len" in single_feature.keys() + aligned_wenet_feat = self.wenet_aligner.offline_resolution_transformation( + np.load(self.utt2wenet_path[utt]), single_feature["target_len"] + ) + single_feature["wenet_feat"] = aligned_wenet_feat + + # print(single_feature.keys()) + # for k, v in single_feature.items(): + # if type(v) in [torch.Tensor, np.ndarray]: + # print(k, v.shape) + # else: + # print(k, v) + # exit() + + return self.clip_if_too_long(single_feature) + + def __len__(self): + return len(self.metadata) + + def random_select(self, feature_seq_len, max_seq_len, ending_ts=2812): + """ + ending_ts: to avoid invalid whisper features for over 30s audios + 2812 = 30 * 24000 // 256 + """ + ts = max(feature_seq_len - max_seq_len, 0) + ts = min(ts, ending_ts - max_seq_len) + + start = random.randint(0, ts) + end = start + max_seq_len + return start, end + + def clip_if_too_long(self, sample, max_seq_len=512): + """ + sample : + { + 'spk_id': (1,), + 'target_len': int + 'mel': (seq_len, dim), + 'frame_pitch': (seq_len,) + 'frame_energy': (seq_len,) + 'content_vector_feat': (seq_len, dim) + } + """ + + if sample["target_len"] <= max_seq_len: + return sample + + start, end = self.random_select(sample["target_len"], max_seq_len) + sample["target_len"] = end - start + + for k in sample.keys(): + if k == "audio": + # audio should be clipped in hop_size scale + sample[k] = sample[k][ + start + * self.cfg.preprocess.hop_size : end + * self.cfg.preprocess.hop_size + ] + elif k == "audio_len": + sample[k] = (end - start) * self.cfg.preprocess.hop_size + elif k not in ["spk_id", "target_len"]: + sample[k] = sample[k][start:end] + + return sample + + +class SVCOnlineDataset(BaseOnlineDataset): + def __init__(self, cfg, dataset, is_valid=False): + super().__init__(cfg, dataset, is_valid=is_valid) + + # Audio pretrained models' sample rates + self.all_sample_rates = {self.sample_rate} + if self.cfg.model.condition_encoder.use_whisper: + self.all_sample_rates.add(self.cfg.preprocess.whisper_sample_rate) + if self.cfg.model.condition_encoder.use_contentvec: + self.all_sample_rates.add(self.cfg.preprocess.contentvec_sample_rate) + if self.cfg.model.condition_encoder.use_wenet: + self.all_sample_rates.add(self.cfg.preprocess.wenet_sample_rate) + + self.highest_sample_rate = max(list(self.all_sample_rates)) + + # The maximum duration (seconds) for one training sample + self.max_duration = 6.0 + self.max_n_frames = int(self.max_duration * self.highest_sample_rate) + + def random_select(self, wav, duration, wav_path): + """ + wav: (T,) + """ + if duration <= self.max_duration: + return wav + + ts_frame = int((duration - self.max_duration) * self.highest_sample_rate) + start = random.randint(0, ts_frame) + end = start + self.max_n_frames + + if (wav[start:end] == 0).all(): + print("*" * 20) + print("Warning! The wav file {} has a lot of silience.".format(wav_path)) + + # There should be at least some frames that are not silience. Then we select them. + assert (wav != 0).any() + start = np.where(wav != 0)[0][0] + end = start + self.max_n_frames + + return wav[start:end] + + def __getitem__(self, index): + """ + single_feature: dict, + wav: (T,) + wav_len: int + target_len: int + mask: (n_frames, 1) + spk_id + + wav_{sr}: (T,) + wav_{sr}_len: int + """ + single_feature = dict() + + utt_item = self.metadata[index] + wav_path = utt_item["Path"] + + ### Use the highest sampling rate to load and randomly select ### + highest_sr_wav, _ = librosa.load(wav_path, sr=self.highest_sample_rate) + highest_sr_wav = self.random_select( + highest_sr_wav, utt_item["Duration"], wav_path + ) + + ### Waveforms under all the sample rates ### + for sr in self.all_sample_rates: + # Resample to the required sample rate + if sr != self.highest_sample_rate: + wav_sr = librosa.resample( + highest_sr_wav, orig_sr=self.highest_sample_rate, target_sr=sr + ) + else: + wav_sr = highest_sr_wav + + wav_sr = torch.as_tensor(wav_sr, dtype=torch.float32) + single_feature["wav_{}".format(sr)] = wav_sr + single_feature["wav_{}_len".format(sr)] = len(wav_sr) + + # For target sample rate + if sr == self.sample_rate: + wav_len = len(wav_sr) + frame_len = wav_len // self.hop_size + + single_feature["wav"] = wav_sr + single_feature["wav_len"] = wav_len + single_feature["target_len"] = frame_len + single_feature["mask"] = torch.ones(frame_len, 1, dtype=torch.long) + + ### Speaker ID ### + if self.cfg.preprocess.use_spkid: + utt = "{}_{}".format(utt_item["Dataset"], utt_item["Uid"]) + single_feature["spk_id"] = torch.tensor( + [self.spk2id[self.utt2spk[utt]]], dtype=torch.int32 + ) + + return single_feature + + def __len__(self): + return len(self.metadata) + + +class SVCOfflineCollator(BaseOfflineCollator): + def __init__(self, cfg): + super().__init__(cfg) + + def __call__(self, batch): + parsed_batch_features = super().__call__(batch) + return parsed_batch_features + + +class SVCOnlineCollator(BaseOnlineCollator): + def __init__(self, cfg): + super().__init__(cfg) + + def __call__(self, batch): + """ + SVCOnlineDataset.__getitem__: + wav: (T,) + wav_len: int + target_len: int + mask: (n_frames, 1) + spk_id: (1) + + wav_{sr}: (T,) + wav_{sr}_len: int + + Returns: + wav: (B, T), torch.float32 + wav_len: (B), torch.long + target_len: (B), torch.long + mask: (B, n_frames, 1), torch.long + spk_id: (B, 1), torch.int32 + + wav_{sr}: (B, T) + wav_{sr}_len: (B), torch.long + """ + packed_batch_features = dict() + + for key in batch[0].keys(): + if "_len" in key: + packed_batch_features[key] = torch.LongTensor([b[key] for b in batch]) + else: + packed_batch_features[key] = pad_sequence( + [b[key] for b in batch], batch_first=True, padding_value=0 + ) + return packed_batch_features + + +class SVCTestDataset(BaseTestDataset): + def __init__(self, args, cfg, infer_type): + BaseTestDataset.__init__(self, args, cfg, infer_type) + self.metadata = self.get_metadata() + + target_singer = args.target_singer + self.cfg = cfg + self.trans_key = args.trans_key + assert type(target_singer) == str + + self.target_singer = target_singer.split("_")[-1] + self.target_dataset = target_singer.replace( + "_{}".format(self.target_singer), "" + ) + if cfg.preprocess.mel_min_max_norm: + if self.cfg.preprocess.features_extraction_mode == "online": + # TODO: Change the hard code + + # Using an empirical mel extrema to normalize + self.target_mel_extrema = load_mel_extrema(cfg.preprocess, "vctk") + else: + self.target_mel_extrema = load_mel_extrema( + cfg.preprocess, self.target_dataset + ) + + self.target_mel_extrema = torch.as_tensor( + self.target_mel_extrema[0] + ), torch.as_tensor(self.target_mel_extrema[1]) + + ######### Load source acoustic features ######### + if cfg.preprocess.use_spkid: + spk2id_path = os.path.join(args.acoustics_dir, cfg.preprocess.spk2id) + # utt2sp_path = os.path.join(self.data_root, cfg.preprocess.utt2spk) + + with open(spk2id_path, "r", encoding="utf-8") as f: + self.spk2id = json.load(f) + # print("self.spk2id", self.spk2id) + + if cfg.preprocess.use_uv: + self.utt2uv_path = { + f'{utt_info["Dataset"]}_{utt_info["Uid"]}': os.path.join( + cfg.preprocess.processed_dir, + utt_info["Dataset"], + cfg.preprocess.uv_dir, + utt_info["Uid"] + ".npy", + ) + for utt_info in self.metadata + } + + if cfg.preprocess.use_frame_pitch: + self.utt2frame_pitch_path = { + f'{utt_info["Dataset"]}_{utt_info["Uid"]}': os.path.join( + cfg.preprocess.processed_dir, + utt_info["Dataset"], + cfg.preprocess.pitch_dir, + utt_info["Uid"] + ".npy", + ) + for utt_info in self.metadata + } + + # Target F0 median + target_f0_statistics_path = os.path.join( + cfg.preprocess.processed_dir, + self.target_dataset, + cfg.preprocess.pitch_dir, + "statistics.json", + ) + self.target_pitch_median = json.load( + open(target_f0_statistics_path, "r", encoding="utf-8") + )[f"{self.target_dataset}_{self.target_singer}"]["voiced_positions"][ + "median" + ] + + # Source F0 median (if infer from file) + if infer_type == "from_file": + source_audio_name = cfg.inference.source_audio_name + source_f0_statistics_path = os.path.join( + cfg.preprocess.processed_dir, + source_audio_name, + cfg.preprocess.pitch_dir, + "statistics.json", + ) + self.source_pitch_median = json.load( + open(source_f0_statistics_path, "r", encoding="utf-8") + )[f"{source_audio_name}_{source_audio_name}"]["voiced_positions"][ + "median" + ] + else: + self.source_pitch_median = None + + if cfg.preprocess.use_frame_energy: + self.utt2frame_energy_path = { + f'{utt_info["Dataset"]}_{utt_info["Uid"]}': os.path.join( + cfg.preprocess.processed_dir, + utt_info["Dataset"], + cfg.preprocess.energy_dir, + utt_info["Uid"] + ".npy", + ) + for utt_info in self.metadata + } + + if cfg.preprocess.use_mel: + self.utt2mel_path = { + f'{utt_info["Dataset"]}_{utt_info["Uid"]}': os.path.join( + cfg.preprocess.processed_dir, + utt_info["Dataset"], + cfg.preprocess.mel_dir, + utt_info["Uid"] + ".npy", + ) + for utt_info in self.metadata + } + + ######### Load source content features' path ######### + if cfg.model.condition_encoder.use_whisper: + self.whisper_aligner = WhisperExtractor(cfg) + self.utt2whisper_path = load_content_feature_path( + self.metadata, cfg.preprocess.processed_dir, cfg.preprocess.whisper_dir + ) + + if cfg.model.condition_encoder.use_contentvec: + self.contentvec_aligner = ContentvecExtractor(cfg) + self.utt2contentVec_path = load_content_feature_path( + self.metadata, + cfg.preprocess.processed_dir, + cfg.preprocess.contentvec_dir, + ) + + if cfg.model.condition_encoder.use_mert: + self.utt2mert_path = load_content_feature_path( + self.metadata, cfg.preprocess.processed_dir, cfg.preprocess.mert_dir + ) + if cfg.model.condition_encoder.use_wenet: + self.wenet_aligner = WenetExtractor(cfg) + self.utt2wenet_path = load_content_feature_path( + self.metadata, cfg.preprocess.processed_dir, cfg.preprocess.wenet_dir + ) + + def __getitem__(self, index): + single_feature = {} + + utt_info = self.metadata[index] + dataset = utt_info["Dataset"] + uid = utt_info["Uid"] + utt = "{}_{}".format(dataset, uid) + + source_dataset = self.metadata[index]["Dataset"] + + if self.cfg.preprocess.use_spkid: + single_feature["spk_id"] = np.array( + [self.spk2id[f"{self.target_dataset}_{self.target_singer}"]], + dtype=np.int32, + ) + + ######### Get Acoustic Features Item ######### + if self.cfg.preprocess.use_mel: + mel = np.load(self.utt2mel_path[utt]) + assert mel.shape[0] == self.cfg.preprocess.n_mel # [n_mels, T] + if self.cfg.preprocess.use_min_max_norm_mel: + # mel norm + mel = cal_normalized_mel(mel, source_dataset, self.cfg.preprocess) + + if "target_len" not in single_feature.keys(): + single_feature["target_len"] = mel.shape[1] + single_feature["mel"] = mel.T # [T, n_mels] + + if self.cfg.preprocess.use_frame_pitch: + frame_pitch_path = self.utt2frame_pitch_path[utt] + frame_pitch = np.load(frame_pitch_path) + + if self.trans_key: + try: + self.trans_key = int(self.trans_key) + except: + pass + if type(self.trans_key) == int: + frame_pitch = transpose_key(frame_pitch, self.trans_key) + elif self.trans_key: + assert self.target_singer + + frame_pitch = pitch_shift_to_target( + frame_pitch, self.target_pitch_median, self.source_pitch_median + ) + + if "target_len" not in single_feature.keys(): + single_feature["target_len"] = len(frame_pitch) + aligned_frame_pitch = align_length( + frame_pitch, single_feature["target_len"] + ) + single_feature["frame_pitch"] = aligned_frame_pitch + + if self.cfg.preprocess.use_uv: + frame_uv_path = self.utt2uv_path[utt] + frame_uv = np.load(frame_uv_path) + aligned_frame_uv = align_length(frame_uv, single_feature["target_len"]) + aligned_frame_uv = [ + 0 if frame_uv else 1 for frame_uv in aligned_frame_uv + ] + aligned_frame_uv = np.array(aligned_frame_uv) + single_feature["frame_uv"] = aligned_frame_uv + + if self.cfg.preprocess.use_frame_energy: + frame_energy_path = self.utt2frame_energy_path[utt] + frame_energy = np.load(frame_energy_path) + if "target_len" not in single_feature.keys(): + single_feature["target_len"] = len(frame_energy) + aligned_frame_energy = align_length( + frame_energy, single_feature["target_len"] + ) + single_feature["frame_energy"] = aligned_frame_energy + + ######### Get Content Features Item ######### + if self.cfg.model.condition_encoder.use_whisper: + assert "target_len" in single_feature.keys() + aligned_whisper_feat = ( + self.whisper_aligner.offline_resolution_transformation( + np.load(self.utt2whisper_path[utt]), single_feature["target_len"] + ) + ) + single_feature["whisper_feat"] = aligned_whisper_feat + + if self.cfg.model.condition_encoder.use_contentvec: + assert "target_len" in single_feature.keys() + aligned_contentvec = ( + self.contentvec_aligner.offline_resolution_transformation( + np.load(self.utt2contentVec_path[utt]), single_feature["target_len"] + ) + ) + single_feature["contentvec_feat"] = aligned_contentvec + + if self.cfg.model.condition_encoder.use_mert: + assert "target_len" in single_feature.keys() + aligned_mert_feat = align_content_feature_length( + np.load(self.utt2mert_path[utt]), + single_feature["target_len"], + source_hop=self.cfg.preprocess.mert_hop_size, + ) + single_feature["mert_feat"] = aligned_mert_feat + + if self.cfg.model.condition_encoder.use_wenet: + assert "target_len" in single_feature.keys() + aligned_wenet_feat = self.wenet_aligner.offline_resolution_transformation( + np.load(self.utt2wenet_path[utt]), single_feature["target_len"] + ) + single_feature["wenet_feat"] = aligned_wenet_feat + + return single_feature + + def __len__(self): + return len(self.metadata) + + +class SVCTestCollator: + """Zero-pads model inputs and targets based on number of frames per step""" + + def __init__(self, cfg): + self.cfg = cfg + + def __call__(self, batch): + packed_batch_features = dict() + + # mel: [b, T, n_mels] + # frame_pitch, frame_energy: [1, T] + # target_len: [1] + # spk_id: [b, 1] + # mask: [b, T, 1] + + for key in batch[0].keys(): + if key == "target_len": + packed_batch_features["target_len"] = torch.LongTensor( + [b["target_len"] for b in batch] + ) + masks = [ + torch.ones((b["target_len"], 1), dtype=torch.long) for b in batch + ] + packed_batch_features["mask"] = pad_sequence( + masks, batch_first=True, padding_value=0 + ) + else: + values = [torch.from_numpy(b[key]) for b in batch] + packed_batch_features[key] = pad_sequence( + values, batch_first=True, padding_value=0 + ) + + return packed_batch_features diff --git a/models/svc/base/svc_inference.py b/models/svc/base/svc_inference.py new file mode 100644 index 0000000..52f88d5 --- /dev/null +++ b/models/svc/base/svc_inference.py @@ -0,0 +1,15 @@ +# Copyright (c) 2023 Amphion. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +from models.base.new_inference import BaseInference +from models.svc.base.svc_dataset import SVCTestCollator, SVCTestDataset + + +class SVCInference(BaseInference): + def __init__(self, args=None, cfg=None, infer_type="from_dataset"): + BaseInference.__init__(self, args, cfg, infer_type) + + def _build_test_dataset(self): + return SVCTestDataset, SVCTestCollator diff --git a/models/svc/base/svc_trainer.py b/models/svc/base/svc_trainer.py new file mode 100644 index 0000000..f2d8f69 --- /dev/null +++ b/models/svc/base/svc_trainer.py @@ -0,0 +1,265 @@ +# Copyright (c) 2023 Amphion. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +import json +import os + +import torch +import torch.nn as nn +import numpy as np + +from models.base.new_trainer import BaseTrainer +from models.svc.base.svc_dataset import ( + SVCOfflineCollator, + SVCOfflineDataset, + SVCOnlineCollator, + SVCOnlineDataset, +) +from processors.audio_features_extractor import AudioFeaturesExtractor +from processors.acoustic_extractor import cal_normalized_mel, load_mel_extrema + +EPS = 1.0e-12 + + +class SVCTrainer(BaseTrainer): + r"""The base trainer for all SVC models. It inherits from BaseTrainer and implements + ``build_criterion``, ``_build_dataset`` and ``_build_singer_lut`` methods. You can inherit from this + class, and implement ``_build_model``, ``_forward_step``. + """ + + def __init__(self, args=None, cfg=None): + self.args = args + self.cfg = cfg + + self._init_accelerator() + + # Only for SVC tasks + with self.accelerator.main_process_first(): + self.singers = self._build_singer_lut() + + # Super init + BaseTrainer.__init__(self, args, cfg) + + # Only for SVC tasks + self.task_type = "SVC" + self.logger.info("Task type: {}".format(self.task_type)) + + ### Following are methods only for SVC tasks ### + def _build_dataset(self): + self.online_features_extraction = ( + self.cfg.preprocess.features_extraction_mode == "online" + ) + + if not self.online_features_extraction: + return SVCOfflineDataset, SVCOfflineCollator + else: + self.audio_features_extractor = AudioFeaturesExtractor(self.cfg) + return SVCOnlineDataset, SVCOnlineCollator + + def _extract_svc_features(self, batch): + """ + Features extraction during training + + Batch: + wav: (B, T) + wav_len: (B) + target_len: (B) + mask: (B, n_frames, 1) + spk_id: (B, 1) + + wav_{sr}: (B, T) + wav_{sr}_len: (B) + + Added elements when output: + mel: (B, n_frames, n_mels) + frame_pitch: (B, n_frames) + frame_uv: (B, n_frames) + frame_energy: (B, n_frames) + frame_{content}: (B, n_frames, D) + """ + + padded_n_frames = torch.max(batch["target_len"]) + final_n_frames = padded_n_frames + + ### Mel Spectrogram ### + if self.cfg.preprocess.use_mel: + # (B, n_mels, n_frames) + raw_mel = self.audio_features_extractor.get_mel_spectrogram(batch["wav"]) + if self.cfg.preprocess.use_min_max_norm_mel: + # TODO: Change the hard code + + # Using the empirical mel extrema to denormalize + if not hasattr(self, "mel_extrema"): + # (n_mels) + m, M = load_mel_extrema(self.cfg.preprocess, "vctk") + # (1, n_mels, 1) + m = ( + torch.as_tensor(m, device=raw_mel.device) + .unsqueeze(0) + .unsqueeze(-1) + ) + M = ( + torch.as_tensor(M, device=raw_mel.device) + .unsqueeze(0) + .unsqueeze(-1) + ) + self.mel_extrema = m, M + + m, M = self.mel_extrema + mel = (raw_mel - m) / (M - m + EPS) * 2 - 1 + + else: + mel = raw_mel + + final_n_frames = min(final_n_frames, mel.size(-1)) + + # (B, n_frames, n_mels) + batch["mel"] = mel.transpose(1, 2) + else: + raw_mel = None + + ### F0 ### + if self.cfg.preprocess.use_frame_pitch: + # (B, n_frames) + raw_f0, raw_uv = self.audio_features_extractor.get_f0( + batch["wav"], + wav_lens=batch["wav_len"], + use_interpolate=self.cfg.preprocess.use_interpolation_for_uv, + return_uv=True, + ) + final_n_frames = min(final_n_frames, raw_f0.size(-1)) + batch["frame_pitch"] = raw_f0 + + if self.cfg.preprocess.use_uv: + batch["frame_uv"] = raw_uv + + ### Energy ### + if self.cfg.preprocess.use_frame_energy: + # (B, n_frames) + raw_energy = self.audio_features_extractor.get_energy( + batch["wav"], mel_spec=raw_mel + ) + final_n_frames = min(final_n_frames, raw_energy.size(-1)) + batch["frame_energy"] = raw_energy + + ### Semantic Features ### + if self.cfg.model.condition_encoder.use_whisper: + # (B, n_frames, D) + whisper_feats = self.audio_features_extractor.get_whisper_features( + wavs=batch["wav_{}".format(self.cfg.preprocess.whisper_sample_rate)], + target_frame_len=padded_n_frames, + ) + final_n_frames = min(final_n_frames, whisper_feats.size(1)) + batch["whisper_feat"] = whisper_feats + + if self.cfg.model.condition_encoder.use_contentvec: + # (B, n_frames, D) + contentvec_feats = self.audio_features_extractor.get_contentvec_features( + wavs=batch["wav_{}".format(self.cfg.preprocess.contentvec_sample_rate)], + target_frame_len=padded_n_frames, + ) + final_n_frames = min(final_n_frames, contentvec_feats.size(1)) + batch["contentvec_feat"] = contentvec_feats + + if self.cfg.model.condition_encoder.use_wenet: + # (B, n_frames, D) + wenet_feats = self.audio_features_extractor.get_wenet_features( + wavs=batch["wav_{}".format(self.cfg.preprocess.wenet_sample_rate)], + target_frame_len=padded_n_frames, + wav_lens=batch[ + "wav_{}_len".format(self.cfg.preprocess.wenet_sample_rate) + ], + ) + final_n_frames = min(final_n_frames, wenet_feats.size(1)) + batch["wenet_feat"] = wenet_feats + + ### Align all the audio features to the same frame length ### + frame_level_features = [ + "mask", + "mel", + "frame_pitch", + "frame_uv", + "frame_energy", + "whisper_feat", + "contentvec_feat", + "wenet_feat", + ] + for k in frame_level_features: + if k in batch: + # (B, n_frames, ...) + batch[k] = batch[k][:, :final_n_frames].contiguous() + + return batch + + @staticmethod + def _build_criterion(): + criterion = nn.MSELoss(reduction="none") + return criterion + + @staticmethod + def _compute_loss(criterion, y_pred, y_gt, loss_mask): + """ + Args: + criterion: MSELoss(reduction='none') + y_pred, y_gt: (B, seq_len, D) + loss_mask: (B, seq_len, 1) + Returns: + loss: Tensor of shape [] + """ + + # (B, seq_len, D) + loss = criterion(y_pred, y_gt) + # expand loss_mask to (B, seq_len, D) + loss_mask = loss_mask.repeat(1, 1, loss.shape[-1]) + + loss = torch.sum(loss * loss_mask) / torch.sum(loss_mask) + return loss + + def _save_auxiliary_states(self): + """ + To save the singer's look-up table in the checkpoint saving path + """ + with open( + os.path.join(self.tmp_checkpoint_save_path, self.cfg.preprocess.spk2id), + "w", + encoding="utf-8", + ) as f: + json.dump(self.singers, f, indent=4, ensure_ascii=False) + + def _build_singer_lut(self): + resumed_singer_path = None + if self.args.resume_from_ckpt_path and self.args.resume_from_ckpt_path != "": + resumed_singer_path = os.path.join( + self.args.resume_from_ckpt_path, self.cfg.preprocess.spk2id + ) + if os.path.exists(os.path.join(self.exp_dir, self.cfg.preprocess.spk2id)): + resumed_singer_path = os.path.join(self.exp_dir, self.cfg.preprocess.spk2id) + + if resumed_singer_path: + with open(resumed_singer_path, "r") as f: + singers = json.load(f) + else: + singers = dict() + + for dataset in self.cfg.dataset: + singer_lut_path = os.path.join( + self.cfg.preprocess.processed_dir, dataset, self.cfg.preprocess.spk2id + ) + with open(singer_lut_path, "r") as singer_lut_path: + singer_lut = json.load(singer_lut_path) + for singer in singer_lut.keys(): + if singer not in singers: + singers[singer] = len(singers) + + with open( + os.path.join(self.exp_dir, self.cfg.preprocess.spk2id), "w" + ) as singer_file: + json.dump(singers, singer_file, indent=4, ensure_ascii=False) + print( + "singers have been dumped to {}".format( + os.path.join(self.exp_dir, self.cfg.preprocess.spk2id) + ) + ) + return singers diff --git a/models/svc/comosvc/__init__.py b/models/svc/comosvc/__init__.py new file mode 100644 index 0000000..19f1cb1 --- /dev/null +++ b/models/svc/comosvc/__init__.py @@ -0,0 +1,4 @@ +# Copyright (c) 2023 Amphion. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. diff --git a/models/svc/comosvc/comosvc.py b/models/svc/comosvc/comosvc.py new file mode 100644 index 0000000..f637d1a --- /dev/null +++ b/models/svc/comosvc/comosvc.py @@ -0,0 +1,391 @@ +# Copyright (c) 2023 Amphion. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +import torch +import torch.nn as nn +import copy +import numpy as np +import math +from tqdm.auto import tqdm + +from utils.ssim import SSIM + +from models.svc.transformer.conformer import Conformer, BaseModule +from models.svc.diffusion.diffusion_wrapper import DiffusionWrapper + + +class Consistency(nn.Module): + def __init__(self, cfg, distill=False): + super().__init__() + self.cfg = cfg + self.denoise_fn = DiffusionWrapper(self.cfg) + self.cfg = cfg.model.comosvc + self.teacher = not distill + self.P_mean = self.cfg.P_mean + self.P_std = self.cfg.P_std + self.sigma_data = self.cfg.sigma_data + self.sigma_min = self.cfg.sigma_min + self.sigma_max = self.cfg.sigma_max + self.rho = self.cfg.rho + self.N = self.cfg.n_timesteps + self.ssim_loss = SSIM() + + # Time step discretization + step_indices = torch.arange(self.N) + # karras boundaries formula + t_steps = ( + self.sigma_min ** (1 / self.rho) + + step_indices + / (self.N - 1) + * (self.sigma_max ** (1 / self.rho) - self.sigma_min ** (1 / self.rho)) + ) ** self.rho + self.t_steps = torch.cat( + [torch.zeros_like(t_steps[:1]), self.round_sigma(t_steps)] + ) + + def init_consistency_training(self): + self.denoise_fn_ema = copy.deepcopy(self.denoise_fn) + self.denoise_fn_pretrained = copy.deepcopy(self.denoise_fn) + + def EDMPrecond(self, x, sigma, cond, denoise_fn): + """ + karras diffusion reverse process + + Args: + x: noisy mel-spectrogram [B x n_mel x L] + sigma: noise level [B x 1 x 1] + cond: output of conformer encoder [B x n_mel x L] + denoise_fn: denoiser neural network e.g. DilatedCNN + + Returns: + denoised mel-spectrogram [B x n_mel x L] + """ + sigma = sigma.reshape(-1, 1, 1) + + c_skip = self.sigma_data**2 / (sigma**2 + self.sigma_data**2) + c_out = ( + (sigma - self.sigma_min) + * self.sigma_data + / (sigma**2 + self.sigma_data**2).sqrt() + ) + c_in = 1 / (self.sigma_data**2 + sigma**2).sqrt() + c_noise = sigma.log() / 4 + + x_in = c_in * x + x_in = x_in.transpose(1, 2) + x = x.transpose(1, 2) + cond = cond.transpose(1, 2) + c_noise = c_noise.squeeze() + if c_noise.dim() == 0: + c_noise = c_noise.unsqueeze(0) + F_x = denoise_fn(x_in, c_noise, cond) + D_x = c_skip * x + c_out * (F_x) + D_x = D_x.transpose(1, 2) + return D_x + + def EDMLoss(self, x_start, cond, mask): + """ + compute loss for EDM model + + Args: + x_start: ground truth mel-spectrogram [B x n_mel x L] + cond: output of conformer encoder [B x n_mel x L] + mask: mask of padded frames [B x n_mel x L] + """ + rnd_normal = torch.randn([x_start.shape[0], 1, 1], device=x_start.device) + sigma = (rnd_normal * self.P_std + self.P_mean).exp() + weight = (sigma**2 + self.sigma_data**2) / (sigma * self.sigma_data) ** 2 + + # follow Grad-TTS, start from Gaussian noise with mean cond and std I + noise = (torch.randn_like(x_start) + cond) * sigma + D_yn = self.EDMPrecond(x_start + noise, sigma, cond, self.denoise_fn) + loss = weight * ((D_yn - x_start) ** 2) + loss = torch.sum(loss * mask) / torch.sum(mask) + return loss + + def round_sigma(self, sigma): + return torch.as_tensor(sigma) + + def edm_sampler( + self, + latents, + cond, + nonpadding, + num_steps=50, + sigma_min=0.002, + sigma_max=80, + rho=7, + S_churn=0, + S_min=0, + S_max=float("inf"), + S_noise=1, + ): + """ + karras diffusion sampler + + Args: + latents: noisy mel-spectrogram [B x n_mel x L] + cond: output of conformer encoder [B x n_mel x L] + nonpadding: mask of padded frames [B x n_mel x L] + num_steps: number of steps for diffusion inference + + Returns: + denoised mel-spectrogram [B x n_mel x L] + """ + # Time step discretization. + + num_steps = num_steps + 1 + step_indices = torch.arange(num_steps, device=latents.device) + t_steps = ( + sigma_max ** (1 / rho) + + step_indices + / (num_steps - 1) + * (sigma_min ** (1 / rho) - sigma_max ** (1 / rho)) + ) ** rho + t_steps = torch.cat([self.round_sigma(t_steps), torch.zeros_like(t_steps[:1])]) + + # Main sampling loop. + x_next = latents * t_steps[0] + # wrap in tqdm for progress bar + bar = tqdm(enumerate(zip(t_steps[:-1], t_steps[1:]))) + for i, (t_cur, t_next) in bar: + x_cur = x_next + # Increase noise temporarily. + gamma = ( + min(S_churn / num_steps, np.sqrt(2) - 1) + if S_min <= t_cur <= S_max + else 0 + ) + t_hat = self.round_sigma(t_cur + gamma * t_cur) + t = torch.zeros((x_cur.shape[0], 1, 1), device=x_cur.device) + t[:, 0, 0] = t_hat + t_hat = t + x_hat = x_cur + ( + t_hat**2 - t_cur**2 + ).sqrt() * S_noise * torch.randn_like(x_cur) + # Euler step. + denoised = self.EDMPrecond(x_hat, t_hat, cond, self.denoise_fn) + d_cur = (x_hat - denoised) / t_hat + x_next = x_hat + (t_next - t_hat) * d_cur + + # add Heun’s 2nd order method + # if i < num_steps - 1: + # t = torch.zeros((x_cur.shape[0], 1, 1), device=x_cur.device) + # t[:, 0, 0] = t_next + # #t_next = t + # denoised = self.EDMPrecond(x_next, t, cond, self.denoise_fn, nonpadding) + # d_prime = (x_next - denoised) / t_next + # x_next = x_hat + (t_next - t_hat) * (0.5 * d_cur + 0.5 * d_prime) + + return x_next + + def CTLoss_D(self, y, cond, mask): + """ + compute loss for consistency distillation + + Args: + y: ground truth mel-spectrogram [B x n_mel x L] + cond: output of conformer encoder [B x n_mel x L] + mask: mask of padded frames [B x n_mel x L] + """ + with torch.no_grad(): + mu = 0.95 + for p, ema_p in zip( + self.denoise_fn.parameters(), self.denoise_fn_ema.parameters() + ): + ema_p.mul_(mu).add_(p, alpha=1 - mu) + + n = torch.randint(1, self.N, (y.shape[0],)) + z = torch.randn_like(y) + cond + + tn_1 = self.t_steps[n + 1].reshape(-1, 1, 1).to(y.device) + f_theta = self.EDMPrecond(y + tn_1 * z, tn_1, cond, self.denoise_fn) + + with torch.no_grad(): + tn = self.t_steps[n].reshape(-1, 1, 1).to(y.device) + + # euler step + x_hat = y + tn_1 * z + denoised = self.EDMPrecond(x_hat, tn_1, cond, self.denoise_fn_pretrained) + d_cur = (x_hat - denoised) / tn_1 + y_tn = x_hat + (tn - tn_1) * d_cur + + # Heun’s 2nd order method + + denoised2 = self.EDMPrecond(y_tn, tn, cond, self.denoise_fn_pretrained) + d_prime = (y_tn - denoised2) / tn + y_tn = x_hat + (tn - tn_1) * (0.5 * d_cur + 0.5 * d_prime) + + f_theta_ema = self.EDMPrecond(y_tn, tn, cond, self.denoise_fn_ema) + + loss = (f_theta - f_theta_ema.detach()) ** 2 + loss = torch.sum(loss * mask) / torch.sum(mask) + + # check nan + if torch.any(torch.isnan(loss)): + print("nan loss") + if torch.any(torch.isnan(f_theta)): + print("nan f_theta") + if torch.any(torch.isnan(f_theta_ema)): + print("nan f_theta_ema") + + return loss + + def get_t_steps(self, N): + N = N + 1 + step_indices = torch.arange(N) + t_steps = ( + self.sigma_min ** (1 / self.rho) + + step_indices + / (N - 1) + * (self.sigma_max ** (1 / self.rho) - self.sigma_min ** (1 / self.rho)) + ) ** self.rho + + return t_steps.flip(0) + + def CT_sampler(self, latents, cond, nonpadding, t_steps=1): + """ + consistency distillation sampler + + Args: + latents: noisy mel-spectrogram [B x n_mel x L] + cond: output of conformer encoder [B x n_mel x L] + nonpadding: mask of padded frames [B x n_mel x L] + t_steps: number of steps for diffusion inference + + Returns: + denoised mel-spectrogram [B x n_mel x L] + """ + # one-step + if t_steps == 1: + t_steps = [80] + # multi-step + else: + t_steps = self.get_t_steps(t_steps) + + t_steps = torch.as_tensor(t_steps).to(latents.device) + latents = latents * t_steps[0] + _t = torch.zeros((latents.shape[0], 1, 1), device=latents.device) + _t[:, 0, 0] = t_steps[0] + x = self.EDMPrecond(latents, _t, cond, self.denoise_fn_ema) + + for t in t_steps[1:-1]: + z = torch.randn_like(x) + cond + x_tn = x + (t**2 - self.sigma_min**2).sqrt() * z + _t = torch.zeros((x.shape[0], 1, 1), device=x.device) + _t[:, 0, 0] = t + t = _t + x = self.EDMPrecond(x_tn, t, cond, self.denoise_fn_ema) + return x + + def forward(self, x, nonpadding, cond, t_steps=1, infer=False): + """ + calculate loss or sample mel-spectrogram + + Args: + x: + training: ground truth mel-spectrogram [B x n_mel x L] + inference: output of encoder [B x n_mel x L] + """ + if self.teacher: # teacher model -- karras diffusion + if not infer: + loss = self.EDMLoss(x, cond, nonpadding) + return loss + else: + shape = (cond.shape[0], self.cfg.n_mel, cond.shape[2]) + x = torch.randn(shape, device=x.device) + cond + x = self.edm_sampler(x, cond, nonpadding, t_steps) + + return x + else: # Consistency distillation + if not infer: + loss = self.CTLoss_D(x, cond, nonpadding) + return loss + + else: + shape = (cond.shape[0], self.cfg.n_mel, cond.shape[2]) + x = torch.randn(shape, device=x.device) + cond + x = self.CT_sampler(x, cond, nonpadding, t_steps=1) + + return x + + +class ComoSVC(BaseModule): + def __init__(self, cfg): + super().__init__() + self.cfg = cfg + self.cfg.model.comosvc.n_mel = self.cfg.preprocess.n_mel + self.distill = self.cfg.model.comosvc.distill + self.encoder = Conformer(self.cfg.model.comosvc) + self.decoder = Consistency(self.cfg, distill=self.distill) + self.ssim_loss = SSIM() + + @torch.no_grad() + def forward(self, x_mask, x, n_timesteps, temperature=1.0): + """ + Generates mel-spectrogram from pitch, content vector, energy. Returns: + 1. encoder outputs (from conformer) + 2. decoder outputs (from diffusion-based decoder) + + Args: + x_mask : mask of padded frames in mel-spectrogram. [B x L x n_mel] + x : output of encoder framework. [B x L x d_condition] + n_timesteps : number of steps to use for reverse diffusion in decoder. + temperature : controls variance of terminal distribution. + """ + + # Get encoder_outputs `mu_x` + mu_x = self.encoder(x, x_mask) + encoder_outputs = mu_x + + mu_x = mu_x.transpose(1, 2) + x_mask = x_mask.transpose(1, 2) + + # Generate sample by performing reverse dynamics + decoder_outputs = self.decoder( + mu_x, x_mask, mu_x, t_steps=n_timesteps, infer=True + ) + decoder_outputs = decoder_outputs.transpose(1, 2) + return encoder_outputs, decoder_outputs + + def compute_loss(self, x_mask, x, mel, skip_diff=False): + """ + Computes 2 losses: + 1. prior loss: loss between mel-spectrogram and encoder outputs. (l2 and ssim loss) + 2. diffusion loss: loss between gaussian noise and its reconstruction by diffusion-based decoder. + + Args: + x_mask : mask of padded frames in mel-spectrogram. [B x L x n_mel] + x : output of encoder framework. [B x L x d_condition] + mel : ground truth mel-spectrogram. [B x L x n_mel] + """ + + mu_x = self.encoder(x, x_mask) + # prior loss + x_mask = x_mask.repeat(1, 1, mel.shape[-1]) + prior_loss = torch.sum( + 0.5 * ((mel - mu_x) ** 2 + math.log(2 * math.pi)) * x_mask + ) + + prior_loss = prior_loss / (torch.sum(x_mask) * self.cfg.model.comosvc.n_mel) + # ssim loss + ssim_loss = self.ssim_loss(mu_x, mel) + ssim_loss = torch.sum(ssim_loss * x_mask) / torch.sum(x_mask) + + x_mask = x_mask.transpose(1, 2) + mu_x = mu_x.transpose(1, 2) + mel = mel.transpose(1, 2) + if not self.distill and skip_diff: + diff_loss = prior_loss.clone() + diff_loss.fill_(0) + + # Cut a small segment of mel-spectrogram in order to increase batch size + else: + mu_y = mu_x + mask_y = x_mask + + diff_loss = self.decoder(mel, mask_y, mu_y, infer=False) + + return ssim_loss, prior_loss, diff_loss diff --git a/models/svc/comosvc/comosvc_inference.py b/models/svc/comosvc/comosvc_inference.py new file mode 100644 index 0000000..2783ec7 --- /dev/null +++ b/models/svc/comosvc/comosvc_inference.py @@ -0,0 +1,39 @@ +# Copyright (c) 2023 Amphion. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +import torch + +from models.svc.base import SVCInference +from modules.encoder.condition_encoder import ConditionEncoder +from models.svc.comosvc.comosvc import ComoSVC + + +class ComoSVCInference(SVCInference): + def __init__(self, args, cfg, infer_type="from_dataset"): + SVCInference.__init__(self, args, cfg, infer_type) + + def _build_model(self): + # TODO: sort out the config + self.cfg.model.condition_encoder.f0_min = self.cfg.preprocess.f0_min + self.cfg.model.condition_encoder.f0_max = self.cfg.preprocess.f0_max + self.condition_encoder = ConditionEncoder(self.cfg.model.condition_encoder) + self.acoustic_mapper = ComoSVC(self.cfg) + if self.cfg.model.comosvc.distill: + self.acoustic_mapper.decoder.init_consistency_training() + model = torch.nn.ModuleList([self.condition_encoder, self.acoustic_mapper]) + return model + + def _inference_each_batch(self, batch_data): + device = self.accelerator.device + for k, v in batch_data.items(): + batch_data[k] = v.to(device) + + cond = self.condition_encoder(batch_data) + mask = batch_data["mask"] + encoder_pred, decoder_pred = self.acoustic_mapper( + mask, cond, self.cfg.inference.comosvc.inference_steps + ) + + return decoder_pred diff --git a/models/svc/comosvc/comosvc_trainer.py b/models/svc/comosvc/comosvc_trainer.py new file mode 100644 index 0000000..a689faa --- /dev/null +++ b/models/svc/comosvc/comosvc_trainer.py @@ -0,0 +1,319 @@ +# Copyright (c) 2023 Amphion. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +import torch +import os +import json5 +from tqdm import tqdm +import json +import shutil + +from models.svc.base import SVCTrainer +from modules.encoder.condition_encoder import ConditionEncoder +from models.svc.comosvc.comosvc import ComoSVC + + +class ComoSVCTrainer(SVCTrainer): + r"""The base trainer for all diffusion models. It inherits from SVCTrainer and + implements ``_build_model`` and ``_forward_step`` methods. + """ + + def __init__(self, args=None, cfg=None): + SVCTrainer.__init__(self, args, cfg) + self.distill = cfg.model.comosvc.distill + self.skip_diff = True + + ### Following are methods only for comoSVC models ### + + def _load_teacher_model(self, model): + r"""Load teacher model from checkpoint file.""" + self.checkpoint_file = self.teacher_model_path + self.logger.info( + "Load teacher acoustic model from {}".format(self.checkpoint_file) + ) + raw_dict = torch.load(self.checkpoint_file) + model.load_state_dict(raw_dict) + + def _build_model(self): + r"""Build the model for training. This function is called in ``__init__`` function.""" + + # TODO: sort out the config + self.cfg.model.condition_encoder.f0_min = self.cfg.preprocess.f0_min + self.cfg.model.condition_encoder.f0_max = self.cfg.preprocess.f0_max + self.condition_encoder = ConditionEncoder(self.cfg.model.condition_encoder) + self.acoustic_mapper = ComoSVC(self.cfg) + model = torch.nn.ModuleList([self.condition_encoder, self.acoustic_mapper]) + if self.cfg.model.comosvc.distill: + if not self.args.resume: + # do not load teacher model when resume + self.teacher_model_path = self.cfg.model.teacher_model_path + self._load_teacher_model(model) + # build teacher & target decoder and freeze teacher + self.acoustic_mapper.decoder.init_consistency_training() + self.freeze_net(self.condition_encoder) + self.freeze_net(self.acoustic_mapper.encoder) + self.freeze_net(self.acoustic_mapper.decoder.denoise_fn_pretrained) + self.freeze_net(self.acoustic_mapper.decoder.denoise_fn_ema) + return model + + def freeze_net(self, model): + r"""Freeze the model for training.""" + for name, param in model.named_parameters(): + param.requires_grad = False + + def __build_optimizer(self): + r"""Build optimizer for training. This function is called in ``__init__`` function.""" + + if self.cfg.train.optimizer.lower() == "adamw": + optimizer = torch.optim.AdamW( + params=filter(lambda p: p.requires_grad, self.model.parameters()), + **self.cfg.train.adamw, + ) + + else: + raise NotImplementedError( + "Not support optimizer: {}".format(self.cfg.train.optimizer) + ) + + return optimizer + + def _forward_step(self, batch): + r"""Forward step for training and inference. This function is called + in ``_train_step`` & ``_test_step`` function. + """ + loss = {} + mask = batch["mask"] + mel_input = batch["mel"] + cond = self.condition_encoder(batch) + if self.distill: + cond = cond.detach() + self.skip_diff = True if self.step < self.cfg.train.fast_steps else False + ssim_loss, prior_loss, diff_loss = self.acoustic_mapper.compute_loss( + mask, cond, mel_input, skip_diff=self.skip_diff + ) + if self.distill: + loss["distil_loss"] = diff_loss + else: + loss["ssim_loss_encoder"] = ssim_loss + loss["prior_loss_encoder"] = prior_loss + loss["diffusion_loss_decoder"] = diff_loss + + return loss + + def _train_epoch(self): + r"""Training epoch. Should return average loss of a batch (sample) over + one epoch. See ``train_loop`` for usage. + """ + self.model.train() + epoch_sum_loss: float = 0.0 + epoch_step: int = 0 + for batch in tqdm( + self.train_dataloader, + desc=f"Training Epoch {self.epoch}", + unit="batch", + colour="GREEN", + leave=False, + dynamic_ncols=True, + smoothing=0.04, + disable=not self.accelerator.is_main_process, + ): + # Do training step and BP + with self.accelerator.accumulate(self.model): + loss = self._train_step(batch) + total_loss = 0 + for k, v in loss.items(): + total_loss += v + self.accelerator.backward(total_loss) + enc_grad_norm = torch.nn.utils.clip_grad_norm_( + self.acoustic_mapper.encoder.parameters(), max_norm=1 + ) + dec_grad_norm = torch.nn.utils.clip_grad_norm_( + self.acoustic_mapper.decoder.parameters(), max_norm=1 + ) + self.optimizer.step() + self.optimizer.zero_grad() + self.batch_count += 1 + + # Update info for each step + # TODO: step means BP counts or batch counts? + if self.batch_count % self.cfg.train.gradient_accumulation_step == 0: + epoch_sum_loss += total_loss + log_info = {} + for k, v in loss.items(): + key = "Step/Train Loss/{}".format(k) + log_info[key] = v + log_info["Step/Learning Rate"] = self.optimizer.param_groups[0]["lr"] + self.accelerator.log( + log_info, + step=self.step, + ) + self.step += 1 + epoch_step += 1 + + self.accelerator.wait_for_everyone() + return ( + epoch_sum_loss + / len(self.train_dataloader) + * self.cfg.train.gradient_accumulation_step, + loss, + ) + + def train_loop(self): + r"""Training loop. The public entry of training process.""" + # Wait everyone to prepare before we move on + self.accelerator.wait_for_everyone() + # dump config file + if self.accelerator.is_main_process: + self.__dump_cfg(self.config_save_path) + self.model.train() + self.optimizer.zero_grad() + # Wait to ensure good to go + self.accelerator.wait_for_everyone() + while self.epoch < self.max_epoch: + self.logger.info("\n") + self.logger.info("-" * 32) + self.logger.info("Epoch {}: ".format(self.epoch)) + + ### TODO: change the return values of _train_epoch() to a loss dict, or (total_loss, loss_dict) + ### It's inconvenient for the model with multiple losses + # Do training & validating epoch + train_loss, loss = self._train_epoch() + self.logger.info(" |- Train/Loss: {:.6f}".format(train_loss)) + for k, v in loss.items(): + self.logger.info(" |- Train/Loss/{}: {:.6f}".format(k, v)) + valid_loss = self._valid_epoch() + self.logger.info(" |- Valid/Loss: {:.6f}".format(valid_loss)) + self.accelerator.log( + {"Epoch/Train Loss": train_loss, "Epoch/Valid Loss": valid_loss}, + step=self.epoch, + ) + + self.accelerator.wait_for_everyone() + # TODO: what is scheduler? + self.scheduler.step(valid_loss) # FIXME: use epoch track correct? + + # Check if hit save_checkpoint_stride and run_eval + run_eval = False + if self.accelerator.is_main_process: + save_checkpoint = False + hit_dix = [] + for i, num in enumerate(self.save_checkpoint_stride): + if self.epoch % num == 0: + save_checkpoint = True + hit_dix.append(i) + run_eval |= self.run_eval[i] + + self.accelerator.wait_for_everyone() + if ( + self.accelerator.is_main_process + and save_checkpoint + and (self.distill or not self.skip_diff) + ): + path = os.path.join( + self.checkpoint_dir, + "epoch-{:04d}_step-{:07d}_loss-{:.6f}".format( + self.epoch, self.step, train_loss + ), + ) + self.tmp_checkpoint_save_path = path + self.accelerator.save_state(path) + print(f"save checkpoint in {path}") + json.dump( + self.checkpoints_path, + open(os.path.join(path, "ckpts.json"), "w"), + ensure_ascii=False, + indent=4, + ) + self._save_auxiliary_states() + + # Remove old checkpoints + to_remove = [] + for idx in hit_dix: + self.checkpoints_path[idx].append(path) + while len(self.checkpoints_path[idx]) > self.keep_last[idx]: + to_remove.append((idx, self.checkpoints_path[idx].pop(0))) + + # Search conflicts + total = set() + for i in self.checkpoints_path: + total |= set(i) + do_remove = set() + for idx, path in to_remove[::-1]: + if path in total: + self.checkpoints_path[idx].insert(0, path) + else: + do_remove.add(path) + + # Remove old checkpoints + for path in do_remove: + shutil.rmtree(path, ignore_errors=True) + self.logger.debug(f"Remove old checkpoint: {path}") + + self.accelerator.wait_for_everyone() + if run_eval: + # TODO: run evaluation + pass + + # Update info for each epoch + self.epoch += 1 + + # Finish training and save final checkpoint + self.accelerator.wait_for_everyone() + if self.accelerator.is_main_process: + self.accelerator.save_state( + os.path.join( + self.checkpoint_dir, + "final_epoch-{:04d}_step-{:07d}_loss-{:.6f}".format( + self.epoch, self.step, valid_loss + ), + ) + ) + self._save_auxiliary_states() + self.accelerator.end_training() + + @torch.inference_mode() + def _valid_epoch(self): + r"""Testing epoch. Should return average loss of a batch (sample) over + one epoch. See ``train_loop`` for usage. + """ + self.model.eval() + epoch_sum_loss = 0.0 + for batch in tqdm( + self.valid_dataloader, + desc=f"Validating Epoch {self.epoch}", + unit="batch", + colour="GREEN", + leave=False, + dynamic_ncols=True, + smoothing=0.04, + disable=not self.accelerator.is_main_process, + ): + batch_loss = self._valid_step(batch) + for k, v in batch_loss.items(): + epoch_sum_loss += v + + self.accelerator.wait_for_everyone() + return epoch_sum_loss / len(self.valid_dataloader) + + @staticmethod + def __count_parameters(model): + model_param = 0.0 + if isinstance(model, dict): + for key, value in model.items(): + model_param += sum(p.numel() for p in model[key].parameters()) + else: + model_param = sum(p.numel() for p in model.parameters()) + return model_param + + def __dump_cfg(self, path): + os.makedirs(os.path.dirname(path), exist_ok=True) + json5.dump( + self.cfg, + open(path, "w"), + indent=4, + sort_keys=True, + ensure_ascii=False, + quote_keys=True, + ) diff --git a/models/svc/diffusion/__init__.py b/models/svc/diffusion/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/models/svc/diffusion/diffusion_inference.py b/models/svc/diffusion/diffusion_inference.py new file mode 100644 index 0000000..a752ef8 --- /dev/null +++ b/models/svc/diffusion/diffusion_inference.py @@ -0,0 +1,63 @@ +# Copyright (c) 2023 Amphion. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +import torch +from diffusers import DDIMScheduler, DDPMScheduler, PNDMScheduler + +from models.svc.base import SVCInference +from models.svc.diffusion.diffusion_inference_pipeline import DiffusionInferencePipeline +from models.svc.diffusion.diffusion_wrapper import DiffusionWrapper +from modules.encoder.condition_encoder import ConditionEncoder + + +class DiffusionInference(SVCInference): + def __init__(self, args=None, cfg=None, infer_type="from_dataset"): + SVCInference.__init__(self, args, cfg, infer_type) + + settings = { + **cfg.model.diffusion.scheduler_settings, + **cfg.inference.diffusion.scheduler_settings, + } + settings.pop("num_inference_timesteps") + + if cfg.inference.diffusion.scheduler.lower() == "ddpm": + self.scheduler = DDPMScheduler(**settings) + self.logger.info("Using DDPM scheduler.") + elif cfg.inference.diffusion.scheduler.lower() == "ddim": + self.scheduler = DDIMScheduler(**settings) + self.logger.info("Using DDIM scheduler.") + elif cfg.inference.diffusion.scheduler.lower() == "pndm": + self.scheduler = PNDMScheduler(**settings) + self.logger.info("Using PNDM scheduler.") + else: + raise NotImplementedError( + "Unsupported scheduler type: {}".format( + cfg.inference.diffusion.scheduler.lower() + ) + ) + + self.pipeline = DiffusionInferencePipeline( + self.model[1], + self.scheduler, + args.diffusion_inference_steps, + ) + + def _build_model(self): + self.cfg.model.condition_encoder.f0_min = self.cfg.preprocess.f0_min + self.cfg.model.condition_encoder.f0_max = self.cfg.preprocess.f0_max + self.condition_encoder = ConditionEncoder(self.cfg.model.condition_encoder) + self.acoustic_mapper = DiffusionWrapper(self.cfg) + model = torch.nn.ModuleList([self.condition_encoder, self.acoustic_mapper]) + return model + + def _inference_each_batch(self, batch_data): + device = self.accelerator.device + for k, v in batch_data.items(): + batch_data[k] = v.to(device) + + conditioner = self.model[0](batch_data) + noise = torch.randn_like(batch_data["mel"], device=device) + y_pred = self.pipeline(noise, conditioner) + return y_pred diff --git a/models/svc/diffusion/diffusion_inference_pipeline.py b/models/svc/diffusion/diffusion_inference_pipeline.py new file mode 100644 index 0000000..e2461aa --- /dev/null +++ b/models/svc/diffusion/diffusion_inference_pipeline.py @@ -0,0 +1,47 @@ +# Copyright (c) 2023 Amphion. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +import torch +from diffusers import DiffusionPipeline + + +class DiffusionInferencePipeline(DiffusionPipeline): + def __init__(self, network, scheduler, num_inference_timesteps=1000): + super().__init__() + + self.register_modules(network=network, scheduler=scheduler) + self.num_inference_timesteps = num_inference_timesteps + + @torch.inference_mode() + def __call__( + self, + initial_noise: torch.Tensor, + conditioner: torch.Tensor = None, + ): + r""" + Args: + initial_noise: The initial noise to be denoised. + conditioner:The conditioner. + n_inference_steps: The number of denoising steps. More denoising steps + usually lead to a higher quality at the expense of slower inference. + """ + + mel = initial_noise + batch_size = mel.size(0) + self.scheduler.set_timesteps(self.num_inference_timesteps) + + for t in self.progress_bar(self.scheduler.timesteps): + timestep = torch.full((batch_size,), t, device=mel.device, dtype=torch.long) + + # 1. predict noise model_output + model_output = self.network(mel, timestep, conditioner) + + # 2. denoise, compute previous step: x_t -> x_t-1 + mel = self.scheduler.step(model_output, t, mel).prev_sample + + # 3. clamp + mel = mel.clamp(-1.0, 1.0) + + return mel diff --git a/models/svc/diffusion/diffusion_trainer.py b/models/svc/diffusion/diffusion_trainer.py new file mode 100644 index 0000000..f1c59f2 --- /dev/null +++ b/models/svc/diffusion/diffusion_trainer.py @@ -0,0 +1,102 @@ +# Copyright (c) 2023 Amphion. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +import torch +from diffusers import DDPMScheduler + +from models.svc.base import SVCTrainer +from modules.encoder.condition_encoder import ConditionEncoder +from .diffusion_wrapper import DiffusionWrapper + + +class DiffusionTrainer(SVCTrainer): + r"""The base trainer for all diffusion models. It inherits from SVCTrainer and + implements ``_build_model`` and ``_forward_step`` methods. + """ + + def __init__(self, args=None, cfg=None): + SVCTrainer.__init__(self, args, cfg) + + # Only for SVC tasks using diffusion + self.noise_scheduler = DDPMScheduler( + **self.cfg.model.diffusion.scheduler_settings, + ) + self.diffusion_timesteps = ( + self.cfg.model.diffusion.scheduler_settings.num_train_timesteps + ) + + ### Following are methods only for diffusion models ### + def _build_model(self): + r"""Build the model for training. This function is called in ``__init__`` function.""" + + # TODO: sort out the config + self.cfg.model.condition_encoder.f0_min = self.cfg.preprocess.f0_min + self.cfg.model.condition_encoder.f0_max = self.cfg.preprocess.f0_max + self.condition_encoder = ConditionEncoder(self.cfg.model.condition_encoder) + self.acoustic_mapper = DiffusionWrapper(self.cfg) + model = torch.nn.ModuleList([self.condition_encoder, self.acoustic_mapper]) + + num_of_params_encoder = self.count_parameters(self.condition_encoder) + num_of_params_am = self.count_parameters(self.acoustic_mapper) + num_of_params = num_of_params_encoder + num_of_params_am + log = "Diffusion Model's Parameters: #Encoder is {:.2f}M, #Diffusion is {:.2f}M. The total is {:.2f}M".format( + num_of_params_encoder / 1e6, num_of_params_am / 1e6, num_of_params / 1e6 + ) + self.logger.info(log) + + return model + + def count_parameters(self, model): + model_param = 0.0 + if isinstance(model, dict): + for key, value in model.items(): + model_param += sum(p.numel() for p in model[key].parameters()) + else: + model_param = sum(p.numel() for p in model.parameters()) + return model_param + + def _check_nan(self, batch, loss, y_pred, y_gt): + if torch.any(torch.isnan(loss)): + for k, v in batch.items(): + self.logger.info(k) + self.logger.info(v) + + super()._check_nan(loss, y_pred, y_gt) + + def _forward_step(self, batch): + r"""Forward step for training and inference. This function is called + in ``_train_step`` & ``_test_step`` function. + """ + device = self.accelerator.device + + if self.online_features_extraction: + # On-the-fly features extraction + batch = self._extract_svc_features(batch) + + # To debug + # for k, v in batch.items(): + # print(k, v.shape, v) + # exit() + + mel_input = batch["mel"] + noise = torch.randn_like(mel_input, device=device, dtype=torch.float32) + batch_size = mel_input.size(0) + timesteps = torch.randint( + 0, + self.diffusion_timesteps, + (batch_size,), + device=device, + dtype=torch.long, + ) + + noisy_mel = self.noise_scheduler.add_noise(mel_input, noise, timesteps) + conditioner = self.condition_encoder(batch) + + y_pred = self.acoustic_mapper(noisy_mel, timesteps, conditioner) + + loss = self._compute_loss(self.criterion, y_pred, noise, batch["mask"]) + self._check_nan(batch, loss, y_pred, noise) + + return loss diff --git a/models/svc/diffusion/diffusion_wrapper.py b/models/svc/diffusion/diffusion_wrapper.py new file mode 100644 index 0000000..ef66c2b --- /dev/null +++ b/models/svc/diffusion/diffusion_wrapper.py @@ -0,0 +1,73 @@ +# Copyright (c) 2023 Amphion. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +import torch.nn as nn + +from modules.diffusion import BiDilConv +from modules.encoder.position_encoder import PositionEncoder + + +class DiffusionWrapper(nn.Module): + def __init__(self, cfg): + super().__init__() + + self.cfg = cfg + self.diff_cfg = cfg.model.diffusion + + self.diff_encoder = PositionEncoder( + d_raw_emb=self.diff_cfg.step_encoder.dim_raw_embedding, + d_out=self.diff_cfg.bidilconv.base_channel, + d_mlp=self.diff_cfg.step_encoder.dim_hidden_layer, + activation_function=self.diff_cfg.step_encoder.activation, + n_layer=self.diff_cfg.step_encoder.num_layer, + max_period=self.diff_cfg.step_encoder.max_period, + ) + + # FIXME: Only support BiDilConv now for debug + if self.diff_cfg.model_type.lower() == "bidilconv": + self.neural_network = BiDilConv( + input_channel=self.cfg.preprocess.n_mel, **self.diff_cfg.bidilconv + ) + else: + raise ValueError( + f"Unsupported diffusion model type: {self.diff_cfg.model_type}" + ) + + def forward(self, x, t, c): + """ + Args: + x: [N, T, mel_band] of mel spectrogram + t: Diffusion time step with shape of [N] + c: [N, T, conditioner_size] of conditioner + + Returns: + [N, T, mel_band] of mel spectrogram + """ + + assert ( + x.size()[:-1] == c.size()[:-1] + ), "x mismatch with c, got \n x: {} \n c: {}".format(x.size(), c.size()) + assert x.size(0) == t.size( + 0 + ), "x mismatch with t, got \n x: {} \n t: {}".format(x.size(), t.size()) + assert t.dim() == 1, "t must be 1D tensor, got {}".format(t.dim()) + + N, T, mel_band = x.size() + + x = x.transpose(1, 2).contiguous() # [N, mel_band, T] + c = c.transpose(1, 2).contiguous() # [N, conditioner_size, T] + t = self.diff_encoder(t).contiguous() # [N, base_channel] + + h = self.neural_network(x, t, c) + h = h.transpose(1, 2).contiguous() # [N, T, mel_band] + + assert h.size() == ( + N, + T, + mel_band, + ), "h mismatch with input x, got \n h: {} \n x: {}".format( + h.size(), (N, T, mel_band) + ) + return h diff --git a/models/svc/transformer/__init__.py b/models/svc/transformer/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/models/svc/transformer/conformer.py b/models/svc/transformer/conformer.py new file mode 100644 index 0000000..5e48019 --- /dev/null +++ b/models/svc/transformer/conformer.py @@ -0,0 +1,405 @@ +# Copyright (c) 2023 Amphion. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +import math +import torch +import numpy as np +import torch.nn as nn +from utils.util import convert_pad_shape + + +class BaseModule(torch.nn.Module): + def __init__(self): + super(BaseModule, self).__init__() + + @property + def nparams(self): + """ + Returns number of trainable parameters of the module. + """ + num_params = 0 + for name, param in self.named_parameters(): + if param.requires_grad: + num_params += np.prod(param.detach().cpu().numpy().shape) + return num_params + + def relocate_input(self, x: list): + """ + Relocates provided tensors to the same device set for the module. + """ + device = next(self.parameters()).device + for i in range(len(x)): + if isinstance(x[i], torch.Tensor) and x[i].device != device: + x[i] = x[i].to(device) + return x + + +class LayerNorm(BaseModule): + def __init__(self, channels, eps=1e-4): + super(LayerNorm, self).__init__() + self.channels = channels + self.eps = eps + + self.gamma = torch.nn.Parameter(torch.ones(channels)) + self.beta = torch.nn.Parameter(torch.zeros(channels)) + + def forward(self, x): + n_dims = len(x.shape) + mean = torch.mean(x, 1, keepdim=True) + variance = torch.mean((x - mean) ** 2, 1, keepdim=True) + + x = (x - mean) * torch.rsqrt(variance + self.eps) + + shape = [1, -1] + [1] * (n_dims - 2) + x = x * self.gamma.view(*shape) + self.beta.view(*shape) + return x + + +class ConvReluNorm(BaseModule): + def __init__( + self, + in_channels, + hidden_channels, + out_channels, + kernel_size, + n_layers, + p_dropout, + eps=1e-5, + ): + super(ConvReluNorm, self).__init__() + self.in_channels = in_channels + self.hidden_channels = hidden_channels + self.out_channels = out_channels + self.kernel_size = kernel_size + self.n_layers = n_layers + self.p_dropout = p_dropout + self.eps = eps + + self.conv_layers = torch.nn.ModuleList() + self.conv_layers.append( + torch.nn.Conv1d( + in_channels, hidden_channels, kernel_size, padding=kernel_size // 2 + ) + ) + self.relu_drop = torch.nn.Sequential( + torch.nn.ReLU(), torch.nn.Dropout(p_dropout) + ) + for _ in range(n_layers - 1): + self.conv_layers.append( + torch.nn.Conv1d( + hidden_channels, + hidden_channels, + kernel_size, + padding=kernel_size // 2, + ) + ) + self.proj = torch.nn.Conv1d(hidden_channels, out_channels, 1) + self.proj.weight.data.zero_() + self.proj.bias.data.zero_() + + def forward(self, x, x_mask): + for i in range(self.n_layers): + x = self.conv_layers[i](x * x_mask) + x = self.instance_norm(x, x_mask) + x = self.relu_drop(x) + x = self.proj(x) + return x * x_mask + + def instance_norm(self, x, mask, return_mean_std=False): + mean, std = self.calc_mean_std(x, mask) + x = (x - mean) / std + if return_mean_std: + return x, mean, std + else: + return x + + def calc_mean_std(self, x, mask=None): + x = x * mask + B, C = x.shape[:2] + mn = x.view(B, C, -1).mean(-1) + sd = (x.view(B, C, -1).var(-1) + self.eps).sqrt() + mn = mn.view(B, C, *((len(x.shape) - 2) * [1])) + sd = sd.view(B, C, *((len(x.shape) - 2) * [1])) + return mn, sd + + +class MultiHeadAttention(BaseModule): + def __init__( + self, + channels, + out_channels, + n_heads, + window_size=None, + heads_share=True, + p_dropout=0.0, + proximal_bias=False, + proximal_init=False, + ): + super(MultiHeadAttention, self).__init__() + assert channels % n_heads == 0 + + self.channels = channels + self.out_channels = out_channels + self.n_heads = n_heads + self.window_size = window_size + self.heads_share = heads_share + self.proximal_bias = proximal_bias + self.p_dropout = p_dropout + self.attn = None + + self.k_channels = channels // n_heads + self.conv_q = torch.nn.Conv1d(channels, channels, 1) + self.conv_k = torch.nn.Conv1d(channels, channels, 1) + self.conv_v = torch.nn.Conv1d(channels, channels, 1) + if window_size is not None: + n_heads_rel = 1 if heads_share else n_heads + rel_stddev = self.k_channels**-0.5 + self.emb_rel_k = torch.nn.Parameter( + torch.randn(n_heads_rel, window_size * 2 + 1, self.k_channels) + * rel_stddev + ) + self.emb_rel_v = torch.nn.Parameter( + torch.randn(n_heads_rel, window_size * 2 + 1, self.k_channels) + * rel_stddev + ) + self.conv_o = torch.nn.Conv1d(channels, out_channels, 1) + self.drop = torch.nn.Dropout(p_dropout) + + torch.nn.init.xavier_uniform_(self.conv_q.weight) + torch.nn.init.xavier_uniform_(self.conv_k.weight) + if proximal_init: + self.conv_k.weight.data.copy_(self.conv_q.weight.data) + self.conv_k.bias.data.copy_(self.conv_q.bias.data) + torch.nn.init.xavier_uniform_(self.conv_v.weight) + + def forward(self, x, c, attn_mask=None): + q = self.conv_q(x) + k = self.conv_k(c) + v = self.conv_v(c) + + x, self.attn = self.attention(q, k, v, mask=attn_mask) + + x = self.conv_o(x) + return x + + def attention(self, query, key, value, mask=None): + b, d, t_s, t_t = (*key.size(), query.size(2)) + query = query.view(b, self.n_heads, self.k_channels, t_t).transpose(2, 3) + key = key.view(b, self.n_heads, self.k_channels, t_s).transpose(2, 3) + value = value.view(b, self.n_heads, self.k_channels, t_s).transpose(2, 3) + + scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(self.k_channels) + if self.window_size is not None: + assert ( + t_s == t_t + ), "Relative attention is only available for self-attention." + key_relative_embeddings = self._get_relative_embeddings(self.emb_rel_k, t_s) + rel_logits = self._matmul_with_relative_keys(query, key_relative_embeddings) + rel_logits = self._relative_position_to_absolute_position(rel_logits) + scores_local = rel_logits / math.sqrt(self.k_channels) + scores = scores + scores_local + if self.proximal_bias: + assert t_s == t_t, "Proximal bias is only available for self-attention." + scores = scores + self._attention_bias_proximal(t_s).to( + device=scores.device, dtype=scores.dtype + ) + if mask is not None: + scores = scores.masked_fill(mask == 0, -1e4) + p_attn = torch.nn.functional.softmax(scores, dim=-1) + p_attn = self.drop(p_attn) + output = torch.matmul(p_attn, value) + if self.window_size is not None: + relative_weights = self._absolute_position_to_relative_position(p_attn) + value_relative_embeddings = self._get_relative_embeddings( + self.emb_rel_v, t_s + ) + output = output + self._matmul_with_relative_values( + relative_weights, value_relative_embeddings + ) + output = output.transpose(2, 3).contiguous().view(b, d, t_t) + return output, p_attn + + def _matmul_with_relative_values(self, x, y): + ret = torch.matmul(x, y.unsqueeze(0)) + return ret + + def _matmul_with_relative_keys(self, x, y): + ret = torch.matmul(x, y.unsqueeze(0).transpose(-2, -1)) + return ret + + def _get_relative_embeddings(self, relative_embeddings, length): + pad_length = max(length - (self.window_size + 1), 0) + slice_start_position = max((self.window_size + 1) - length, 0) + slice_end_position = slice_start_position + 2 * length - 1 + if pad_length > 0: + padded_relative_embeddings = torch.nn.functional.pad( + relative_embeddings, + convert_pad_shape([[0, 0], [pad_length, pad_length], [0, 0]]), + ) + else: + padded_relative_embeddings = relative_embeddings + used_relative_embeddings = padded_relative_embeddings[ + :, slice_start_position:slice_end_position + ] + return used_relative_embeddings + + def _relative_position_to_absolute_position(self, x): + batch, heads, length, _ = x.size() + x = torch.nn.functional.pad( + x, convert_pad_shape([[0, 0], [0, 0], [0, 0], [0, 1]]) + ) + x_flat = x.view([batch, heads, length * 2 * length]) + x_flat = torch.nn.functional.pad( + x_flat, convert_pad_shape([[0, 0], [0, 0], [0, length - 1]]) + ) + x_final = x_flat.view([batch, heads, length + 1, 2 * length - 1])[ + :, :, :length, length - 1 : + ] + return x_final + + def _absolute_position_to_relative_position(self, x): + batch, heads, length, _ = x.size() + x = torch.nn.functional.pad( + x, convert_pad_shape([[0, 0], [0, 0], [0, 0], [0, length - 1]]) + ) + x_flat = x.view([batch, heads, length**2 + length * (length - 1)]) + x_flat = torch.nn.functional.pad( + x_flat, convert_pad_shape([[0, 0], [0, 0], [length, 0]]) + ) + x_final = x_flat.view([batch, heads, length, 2 * length])[:, :, :, 1:] + return x_final + + def _attention_bias_proximal(self, length): + r = torch.arange(length, dtype=torch.float32) + diff = torch.unsqueeze(r, 0) - torch.unsqueeze(r, 1) + return torch.unsqueeze(torch.unsqueeze(-torch.log1p(torch.abs(diff)), 0), 0) + + +class FFN(BaseModule): + def __init__( + self, in_channels, out_channels, filter_channels, kernel_size, p_dropout=0.0 + ): + super(FFN, self).__init__() + self.in_channels = in_channels + self.out_channels = out_channels + self.filter_channels = filter_channels + self.kernel_size = kernel_size + self.p_dropout = p_dropout + + self.conv_1 = torch.nn.Conv1d( + in_channels, filter_channels, kernel_size, padding=kernel_size // 2 + ) + self.conv_2 = torch.nn.Conv1d( + filter_channels, out_channels, kernel_size, padding=kernel_size // 2 + ) + self.drop = torch.nn.Dropout(p_dropout) + + def forward(self, x, x_mask): + x = self.conv_1(x * x_mask) + x = torch.relu(x) + x = self.drop(x) + x = self.conv_2(x * x_mask) + return x * x_mask + + +class Encoder(BaseModule): + def __init__( + self, + hidden_channels, + filter_channels, + n_heads=2, + n_layers=6, + kernel_size=3, + p_dropout=0.1, + window_size=4, + **kwargs + ): + super(Encoder, self).__init__() + self.hidden_channels = hidden_channels + self.filter_channels = filter_channels + self.n_heads = n_heads + self.n_layers = n_layers + self.kernel_size = kernel_size + self.p_dropout = p_dropout + self.window_size = window_size + + self.drop = torch.nn.Dropout(p_dropout) + self.attn_layers = torch.nn.ModuleList() + self.norm_layers_1 = torch.nn.ModuleList() + self.ffn_layers = torch.nn.ModuleList() + self.norm_layers_2 = torch.nn.ModuleList() + for _ in range(self.n_layers): + self.attn_layers.append( + MultiHeadAttention( + hidden_channels, + hidden_channels, + n_heads, + window_size=window_size, + p_dropout=p_dropout, + ) + ) + self.norm_layers_1.append(LayerNorm(hidden_channels)) + self.ffn_layers.append( + FFN( + hidden_channels, + hidden_channels, + filter_channels, + kernel_size, + p_dropout=p_dropout, + ) + ) + self.norm_layers_2.append(LayerNorm(hidden_channels)) + + def forward(self, x, x_mask): + attn_mask = x_mask.unsqueeze(2) * x_mask.unsqueeze(-1) + for i in range(self.n_layers): + x = x * x_mask + y = self.attn_layers[i](x, x, attn_mask) + y = self.drop(y) + x = self.norm_layers_1[i](x + y) + y = self.ffn_layers[i](x, x_mask) + y = self.drop(y) + x = self.norm_layers_2[i](x + y) + x = x * x_mask + return x + + +class Conformer(BaseModule): + def __init__(self, cfg): + super().__init__() + self.cfg = cfg + self.n_heads = self.cfg.n_heads + self.n_layers = self.cfg.n_layers + self.hidden_channels = self.cfg.input_dim + self.filter_channels = self.cfg.filter_channels + self.output_dim = self.cfg.output_dim + self.dropout = self.cfg.dropout + + self.conformer_encoder = Encoder( + self.hidden_channels, + self.filter_channels, + n_heads=self.n_heads, + n_layers=self.n_layers, + kernel_size=3, + p_dropout=self.dropout, + window_size=4, + ) + self.projection = nn.Conv1d(self.hidden_channels, self.output_dim, 1) + + def forward(self, x, x_mask): + """ + Args: + x: (N, seq_len, input_dim) + Returns: + output: (N, seq_len, output_dim) + """ + # (N, seq_len, d_model) + x = x.transpose(1, 2) + x_mask = x_mask.transpose(1, 2) + output = self.conformer_encoder(x, x_mask) + # (N, seq_len, output_dim) + output = self.projection(output) + output = output.transpose(1, 2) + return output diff --git a/models/svc/transformer/transformer.py b/models/svc/transformer/transformer.py new file mode 100644 index 0000000..fd3cdb6 --- /dev/null +++ b/models/svc/transformer/transformer.py @@ -0,0 +1,82 @@ +# Copyright (c) 2023 Amphion. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +import math +import torch +import torch.nn as nn +from torch.nn import TransformerEncoder, TransformerEncoderLayer + + +class Transformer(nn.Module): + def __init__(self, cfg): + super().__init__() + self.cfg = cfg + + dropout = self.cfg.dropout + nhead = self.cfg.n_heads + nlayers = self.cfg.n_layers + input_dim = self.cfg.input_dim + output_dim = self.cfg.output_dim + + d_model = input_dim + self.pos_encoder = PositionalEncoding(d_model, dropout) + encoder_layers = TransformerEncoderLayer( + d_model, nhead, dropout=dropout, batch_first=True + ) + self.transformer_encoder = TransformerEncoder(encoder_layers, nlayers) + + self.output_mlp = nn.Linear(d_model, output_dim) + + def forward(self, x, mask=None): + """ + Args: + x: (N, seq_len, input_dim) + Returns: + output: (N, seq_len, output_dim) + """ + # (N, seq_len, d_model) + src = self.pos_encoder(x) + # model_stats["pos_embedding"] = x + # (N, seq_len, d_model) + output = self.transformer_encoder(src) + # (N, seq_len, output_dim) + output = self.output_mlp(output) + return output + + +class PositionalEncoding(nn.Module): + def __init__(self, d_model: int, dropout: float = 0.1, max_len: int = 5000): + super().__init__() + self.dropout = nn.Dropout(p=dropout) + + position = torch.arange(max_len).unsqueeze(1) + div_term = torch.exp( + torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model) + ) + + # Assume that x is (seq_len, N, d) + # pe = torch.zeros(max_len, 1, d_model) + # pe[:, 0, 0::2] = torch.sin(position * div_term) + # pe[:, 0, 1::2] = torch.cos(position * div_term) + + # Assume that x in (N, seq_len, d) + pe = torch.zeros(1, max_len, d_model) + pe[0, :, 0::2] = torch.sin(position * div_term) + pe[0, :, 1::2] = torch.cos(position * div_term) + + self.register_buffer("pe", pe) + + def forward(self, x): + """ + Args: + x: Tensor, shape [N, seq_len, d] + """ + # Old: Assume that x is (seq_len, N, d), and self.pe is (max_len, 1, d_model) + # x = x + self.pe[: x.size(0)] + + # Now: self.pe is (1, max_len, d) + x = x + self.pe[:, : x.size(1), :] + + return self.dropout(x) diff --git a/models/svc/transformer/transformer_inference.py b/models/svc/transformer/transformer_inference.py new file mode 100644 index 0000000..f6299c5 --- /dev/null +++ b/models/svc/transformer/transformer_inference.py @@ -0,0 +1,45 @@ +# Copyright (c) 2023 Amphion. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +import os +import time +import numpy as np +import torch +from tqdm import tqdm +import torch.nn as nn +from collections import OrderedDict + +from models.svc.base import SVCInference +from modules.encoder.condition_encoder import ConditionEncoder +from models.svc.transformer.transformer import Transformer +from models.svc.transformer.conformer import Conformer + + +class TransformerInference(SVCInference): + def __init__(self, args=None, cfg=None, infer_type="from_dataset"): + SVCInference.__init__(self, args, cfg, infer_type) + + def _build_model(self): + self.cfg.model.condition_encoder.f0_min = self.cfg.preprocess.f0_min + self.cfg.model.condition_encoder.f0_max = self.cfg.preprocess.f0_max + self.condition_encoder = ConditionEncoder(self.cfg.model.condition_encoder) + if self.cfg.model.transformer.type == "transformer": + self.acoustic_mapper = Transformer(self.cfg.model.transformer) + elif self.cfg.model.transformer.type == "conformer": + self.acoustic_mapper = Conformer(self.cfg.model.transformer) + else: + raise NotImplementedError + model = torch.nn.ModuleList([self.condition_encoder, self.acoustic_mapper]) + return model + + def _inference_each_batch(self, batch_data): + device = self.accelerator.device + for k, v in batch_data.items(): + batch_data[k] = v.to(device) + + condition = self.condition_encoder(batch_data) + y_pred = self.acoustic_mapper(condition, batch_data["mask"]) + + return y_pred diff --git a/models/svc/transformer/transformer_trainer.py b/models/svc/transformer/transformer_trainer.py new file mode 100644 index 0000000..3633078 --- /dev/null +++ b/models/svc/transformer/transformer_trainer.py @@ -0,0 +1,52 @@ +# Copyright (c) 2023 Amphion. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +import torch + +from models.svc.base import SVCTrainer +from modules.encoder.condition_encoder import ConditionEncoder +from models.svc.transformer.transformer import Transformer +from models.svc.transformer.conformer import Conformer +from utils.ssim import SSIM + + +class TransformerTrainer(SVCTrainer): + def __init__(self, args, cfg): + SVCTrainer.__init__(self, args, cfg) + self.ssim_loss = SSIM() + + def _build_model(self): + self.cfg.model.condition_encoder.f0_min = self.cfg.preprocess.f0_min + self.cfg.model.condition_encoder.f0_max = self.cfg.preprocess.f0_max + self.condition_encoder = ConditionEncoder(self.cfg.model.condition_encoder) + if self.cfg.model.transformer.type == "transformer": + self.acoustic_mapper = Transformer(self.cfg.model.transformer) + elif self.cfg.model.transformer.type == "conformer": + self.acoustic_mapper = Conformer(self.cfg.model.transformer) + else: + raise NotImplementedError + model = torch.nn.ModuleList([self.condition_encoder, self.acoustic_mapper]) + return model + + def _forward_step(self, batch): + total_loss = 0 + device = self.accelerator.device + mel = batch["mel"] + mask = batch["mask"] + + condition = self.condition_encoder(batch) + mel_pred = self.acoustic_mapper(condition, mask) + + l1_loss = torch.sum(torch.abs(mel_pred - mel) * batch["mask"]) / torch.sum( + batch["mask"] + ) + self._check_nan(l1_loss, mel_pred, mel) + total_loss += l1_loss + ssim_loss = self.ssim_loss(mel_pred, mel) + ssim_loss = torch.sum(ssim_loss * batch["mask"]) / torch.sum(batch["mask"]) + self._check_nan(ssim_loss, mel_pred, mel) + total_loss += ssim_loss + + return total_loss diff --git a/models/svc/vits/__init__.py b/models/svc/vits/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/models/svc/vits/vits.py b/models/svc/vits/vits.py new file mode 100644 index 0000000..983a704 --- /dev/null +++ b/models/svc/vits/vits.py @@ -0,0 +1,267 @@ +# Copyright (c) 2023 Amphion. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +# This code is modified from https://github.com/svc-develop-team/so-vits-svc/blob/4.1-Stable/models.py +import copy +import torch +from torch import nn +from torch.nn import functional as F + +from utils.util import * + +from modules.transformer.attentions import Encoder +from models.tts.vits.vits import ResidualCouplingBlock, PosteriorEncoder +from models.vocoders.gan.generator.bigvgan import BigVGAN +from models.vocoders.gan.generator.hifigan import HiFiGAN +from models.vocoders.gan.generator.nsfhifigan import NSFHiFiGAN +from models.vocoders.gan.generator.melgan import MelGAN +from models.vocoders.gan.generator.apnet import APNet +from modules.encoder.condition_encoder import ConditionEncoder + + +def slice_pitch_segments(x, ids_str, segment_size=4): + ret = torch.zeros_like(x[:, :segment_size]) + for i in range(x.size(0)): + idx_str = ids_str[i] + idx_end = idx_str + segment_size + ret[i] = x[i, idx_str:idx_end] + return ret + + +def rand_slice_segments_with_pitch(x, pitch, x_lengths=None, segment_size=4): + b, d, t = x.size() + if x_lengths is None: + x_lengths = t + ids_str_max = x_lengths - segment_size + 1 + ids_str = (torch.rand([b]).to(device=x.device) * ids_str_max).to(dtype=torch.long) + ret = slice_segments(x, ids_str, segment_size) + ret_pitch = slice_pitch_segments(pitch, ids_str, segment_size) + return ret, ret_pitch, ids_str + + +class ContentEncoder(nn.Module): + def __init__( + self, + out_channels, + hidden_channels, + kernel_size, + n_layers, + gin_channels=0, + filter_channels=None, + n_heads=None, + p_dropout=None, + ): + super().__init__() + self.out_channels = out_channels + self.hidden_channels = hidden_channels + self.kernel_size = kernel_size + self.n_layers = n_layers + self.gin_channels = gin_channels + + self.f0_emb = nn.Embedding(256, hidden_channels) + + self.enc_ = Encoder( + hidden_channels, filter_channels, n_heads, n_layers, kernel_size, p_dropout + ) + + self.proj = nn.Conv1d(hidden_channels, out_channels * 2, 1) + + # condition_encoder ver. + def forward(self, x, x_mask, noice_scale=1): + x = self.enc_(x * x_mask, x_mask) + stats = self.proj(x) * x_mask + m, logs = torch.split(stats, self.out_channels, dim=1) + z = (m + torch.randn_like(m) * torch.exp(logs) * noice_scale) * x_mask + + return z, m, logs, x_mask + + +class SynthesizerTrn(nn.Module): + """ + Synthesizer for Training + """ + + def __init__(self, spec_channels, segment_size, cfg): + super().__init__() + self.spec_channels = spec_channels + self.segment_size = segment_size + self.cfg = cfg + self.inter_channels = cfg.model.vits.inter_channels + self.hidden_channels = cfg.model.vits.hidden_channels + self.filter_channels = cfg.model.vits.filter_channels + self.n_heads = cfg.model.vits.n_heads + self.n_layers = cfg.model.vits.n_layers + self.kernel_size = cfg.model.vits.kernel_size + self.p_dropout = cfg.model.vits.p_dropout + self.n_flow_layer = cfg.model.vits.n_flow_layer + self.gin_channels = cfg.model.vits.gin_channels + self.n_speakers = cfg.model.vits.n_speakers + + # f0 + self.n_bins = cfg.preprocess.pitch_bin + self.f0_min = cfg.preprocess.f0_min + self.f0_max = cfg.preprocess.f0_max + + # TODO: sort out the config + self.cfg.model.condition_encoder.f0_min = self.cfg.preprocess.f0_min + self.cfg.model.condition_encoder.f0_max = self.cfg.preprocess.f0_max + self.condition_encoder = ConditionEncoder(self.cfg.model.condition_encoder) + + self.emb_g = nn.Embedding(self.n_speakers, self.gin_channels) + + self.enc_p = ContentEncoder( + self.inter_channels, + self.hidden_channels, + filter_channels=self.filter_channels, + n_heads=self.n_heads, + n_layers=self.n_layers, + kernel_size=self.kernel_size, + p_dropout=self.p_dropout, + ) + + assert cfg.model.generator in [ + "bigvgan", + "hifigan", + "melgan", + "nsfhifigan", + "apnet", + ] + self.dec_name = cfg.model.generator + temp_cfg = copy.deepcopy(cfg) + temp_cfg.preprocess.n_mel = self.inter_channels + if cfg.model.generator == "bigvgan": + temp_cfg.model.bigvgan = cfg.model.generator_config.bigvgan + self.dec = BigVGAN(temp_cfg) + elif cfg.model.generator == "hifigan": + temp_cfg.model.hifigan = cfg.model.generator_config.hifigan + self.dec = HiFiGAN(temp_cfg) + elif cfg.model.generator == "melgan": + temp_cfg.model.melgan = cfg.model.generator_config.melgan + self.dec = MelGAN(temp_cfg) + elif cfg.model.generator == "nsfhifigan": + temp_cfg.model.nsfhifigan = cfg.model.generator_config.nsfhifigan + self.dec = NSFHiFiGAN(temp_cfg) # TODO: nsf need f0 + elif cfg.model.generator == "apnet": + temp_cfg.model.apnet = cfg.model.generator_config.apnet + self.dec = APNet(temp_cfg) + + self.enc_q = PosteriorEncoder( + self.spec_channels, + self.inter_channels, + self.hidden_channels, + 5, + 1, + 16, + gin_channels=self.gin_channels, + ) + + self.flow = ResidualCouplingBlock( + self.inter_channels, + self.hidden_channels, + 5, + 1, + self.n_flow_layer, + gin_channels=self.gin_channels, + ) + + def forward(self, data): + """VitsSVC forward function. + + Args: + data (dict): condition data & audio data, including: + B: batch size, T: target length + { + "spk_id": [B, singer_table_size] + "target_len": [B] + "mask": [B, T, 1] + "mel": [B, T, n_mel] + "linear": [B, T, n_fft // 2 + 1] + "frame_pitch": [B, T] + "frame_uv": [B, T] + "audio": [B, audio_len] + "audio_len": [B] + "contentvec_feat": [B, T, contentvec_dim] + "whisper_feat": [B, T, whisper_dim] + ... + } + """ + + # TODO: elegantly handle the dimensions + spec = data["linear"].transpose(1, 2) + + g = data["spk_id"] + g = self.emb_g(g).transpose(1, 2) + + c_lengths = data["target_len"] + spec_lengths = data["target_len"] + f0 = data["frame_pitch"] + + # condition_encoder ver. + x = self.condition_encoder(data).transpose(1, 2) + x_mask = torch.unsqueeze(sequence_mask(c_lengths, f0.size(1)), 1).to(x.dtype) + + # prior encoder + z_ptemp, m_p, logs_p, _ = self.enc_p(x, x_mask) + # posterior encoder + z, m_q, logs_q, spec_mask = self.enc_q(spec, spec_lengths, g=g) + + # flow + z_p = self.flow(z, spec_mask, g=g) + z_slice, pitch_slice, ids_slice = rand_slice_segments_with_pitch( + z, f0, spec_lengths, self.segment_size + ) + + if self.dec_name == "nsfhifigan": + o = self.dec(z_slice, f0=f0.float()) + elif self.dec_name == "apnet": + _, _, _, _, o = self.dec(z_slice) + else: + o = self.dec(z_slice) + + outputs = { + "y_hat": o, + "ids_slice": ids_slice, + "x_mask": x_mask, + "z_mask": data["mask"].transpose(1, 2), + "z": z, + "z_p": z_p, + "m_p": m_p, + "logs_p": logs_p, + "m_q": m_q, + "logs_q": logs_q, + } + return outputs + + @torch.no_grad() + def infer(self, data, noise_scale=0.35, seed=52468): + # c, f0, uv, g + f0 = data["frame_pitch"] + g = data["spk_id"] + + if f0.device == torch.device("cuda"): + torch.cuda.manual_seed_all(seed) + else: + torch.manual_seed(seed) + + c_lengths = (torch.ones(f0.size(0)) * f0.size(-1)).to(f0.device) + + if g.dim() == 1: + g = g.unsqueeze(0) + g = self.emb_g(g).transpose(1, 2) + + # condition_encoder ver. + x = self.condition_encoder(data).transpose(1, 2) + x_mask = torch.unsqueeze(sequence_mask(c_lengths, f0.size(1)), 1).to(x.dtype) + + z_p, m_p, logs_p, c_mask = self.enc_p(x, x_mask, noice_scale=noise_scale) + z = self.flow(z_p, c_mask, g=g, reverse=True) + + if self.dec_name == "nsfhifigan": + o = self.dec(z * c_mask, f0=f0.float()) + elif self.dec_name == "apnet": + _, _, _, _, o = self.dec(z * c_mask) + else: + o = self.dec(z * c_mask) + return o, f0 diff --git a/models/svc/vits/vits_inference.py b/models/svc/vits/vits_inference.py new file mode 100644 index 0000000..f1e4c1a --- /dev/null +++ b/models/svc/vits/vits_inference.py @@ -0,0 +1,102 @@ +# Copyright (c) 2023 Amphion. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +import json +import os +import time +import numpy as np +from tqdm import tqdm +import torch +from torch.utils.data import DataLoader + +from models.svc.base import SVCInference +from models.svc.vits.vits import SynthesizerTrn + +from models.svc.base.svc_dataset import SVCTestDataset, SVCTestCollator +from utils.io import save_audio +from utils.audio_slicer import is_silence + + +class VitsInference(SVCInference): + def __init__(self, args=None, cfg=None, infer_type="from_dataset"): + SVCInference.__init__(self, args, cfg) + + def _build_model(self): + net_g = SynthesizerTrn( + self.cfg.preprocess.n_fft // 2 + 1, + self.cfg.preprocess.segment_size // self.cfg.preprocess.hop_size, + self.cfg, + ) + self.model = net_g + return net_g + + def build_save_dir(self, dataset, speaker): + save_dir = os.path.join( + self.args.output_dir, + "svc_am_step-{}_{}".format(self.am_restore_step, self.args.mode), + ) + if dataset is not None: + save_dir = os.path.join(save_dir, "data_{}".format(dataset)) + if speaker != -1: + save_dir = os.path.join( + save_dir, + "spk_{}".format(speaker), + ) + os.makedirs(save_dir, exist_ok=True) + print("Saving to ", save_dir) + return save_dir + + def _build_dataloader(self): + datasets, collate = self._build_test_dataset() + self.test_dataset = datasets(self.args, self.cfg, self.infer_type) + self.test_collate = collate(self.cfg) + self.test_batch_size = min( + self.cfg.inference.batch_size, len(self.test_dataset.metadata) + ) + test_dataloader = DataLoader( + self.test_dataset, + collate_fn=self.test_collate, + num_workers=1, + batch_size=self.test_batch_size, + shuffle=False, + ) + return test_dataloader + + @torch.inference_mode() + def inference(self): + res = [] + for i, batch in enumerate(self.test_dataloader): + pred_audio_list = self._inference_each_batch(batch) + for j, wav in enumerate(pred_audio_list): + uid = self.test_dataset.metadata[i * self.test_batch_size + j]["Uid"] + file = os.path.join(self.args.output_dir, f"{uid}.wav") + print(f"Saving {file}") + + wav = wav.numpy(force=True) + save_audio( + file, + wav, + self.cfg.preprocess.sample_rate, + add_silence=False, + turn_up=not is_silence(wav, self.cfg.preprocess.sample_rate), + ) + res.append(file) + return res + + def _inference_each_batch(self, batch_data, noise_scale=0.667): + device = self.accelerator.device + pred_res = [] + self.model.eval() + with torch.no_grad(): + # Put the data to device + # device = self.accelerator.device + for k, v in batch_data.items(): + batch_data[k] = v.to(device) + + audios, f0 = self.model.infer(batch_data, noise_scale=noise_scale) + + pred_res.extend(audios) + + return pred_res diff --git a/models/svc/vits/vits_trainer.py b/models/svc/vits/vits_trainer.py new file mode 100644 index 0000000..8be6d37 --- /dev/null +++ b/models/svc/vits/vits_trainer.py @@ -0,0 +1,704 @@ +# Copyright (c) 2023 Amphion. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +import torch +from torch.optim.lr_scheduler import ExponentialLR +from tqdm import tqdm +from pathlib import Path +import shutil +import accelerate + +# from models.svc.base import SVCTrainer +from models.svc.base.svc_dataset import SVCOfflineCollator, SVCOfflineDataset +from models.svc.vits.vits import * +from models.svc.base import SVCTrainer + +from utils.mel import mel_spectrogram_torch +import json + +from models.vocoders.gan.discriminator.mpd import ( + MultiPeriodDiscriminator_vits as MultiPeriodDiscriminator, +) + + +class VitsSVCTrainer(SVCTrainer): + def __init__(self, args, cfg): + self.args = args + self.cfg = cfg + SVCTrainer.__init__(self, args, cfg) + + def _accelerator_prepare(self): + ( + self.train_dataloader, + self.valid_dataloader, + ) = self.accelerator.prepare( + self.train_dataloader, + self.valid_dataloader, + ) + if isinstance(self.model, dict): + for key in self.model.keys(): + self.model[key] = self.accelerator.prepare(self.model[key]) + else: + self.model = self.accelerator.prepare(self.model) + + if isinstance(self.optimizer, dict): + for key in self.optimizer.keys(): + self.optimizer[key] = self.accelerator.prepare(self.optimizer[key]) + else: + self.optimizer = self.accelerator.prepare(self.optimizer) + + if isinstance(self.scheduler, dict): + for key in self.scheduler.keys(): + self.scheduler[key] = self.accelerator.prepare(self.scheduler[key]) + else: + self.scheduler = self.accelerator.prepare(self.scheduler) + + def _load_model( + self, + checkpoint_dir: str = None, + checkpoint_path: str = None, + resume_type: str = "", + ): + r"""Load model from checkpoint. If checkpoint_path is None, it will + load the latest checkpoint in checkpoint_dir. If checkpoint_path is not + None, it will load the checkpoint specified by checkpoint_path. **Only use this + method after** ``accelerator.prepare()``. + """ + if checkpoint_path is None: + ls = [str(i) for i in Path(checkpoint_dir).glob("*")] + ls.sort(key=lambda x: int(x.split("_")[-3].split("-")[-1]), reverse=True) + checkpoint_path = ls[0] + self.logger.info("Resume from {}...".format(checkpoint_path)) + + if resume_type in ["resume", ""]: + # Load all the things, including model weights, optimizer, scheduler, and random states. + self.accelerator.load_state(input_dir=checkpoint_path) + + # set epoch and step + self.epoch = int(checkpoint_path.split("_")[-3].split("-")[-1]) + 1 + self.step = int(checkpoint_path.split("_")[-2].split("-")[-1]) + 1 + + elif resume_type == "finetune": + # Load only the model weights + accelerate.load_checkpoint_and_dispatch( + self.accelerator.unwrap_model(self.model["generator"]), + os.path.join(checkpoint_path, "pytorch_model.bin"), + ) + accelerate.load_checkpoint_and_dispatch( + self.accelerator.unwrap_model(self.model["discriminator"]), + os.path.join(checkpoint_path, "pytorch_model.bin"), + ) + self.logger.info("Load model weights for finetune...") + + else: + raise ValueError("Resume_type must be `resume` or `finetune`.") + + return checkpoint_path + + def _build_model(self): + net_g = SynthesizerTrn( + self.cfg.preprocess.n_fft // 2 + 1, + self.cfg.preprocess.segment_size // self.cfg.preprocess.hop_size, + # directly use cfg + self.cfg, + ) + net_d = MultiPeriodDiscriminator(self.cfg.model.vits.use_spectral_norm) + model = {"generator": net_g, "discriminator": net_d} + + return model + + def _build_dataset(self): + return SVCOfflineDataset, SVCOfflineCollator + + def _build_optimizer(self): + optimizer_g = torch.optim.AdamW( + self.model["generator"].parameters(), + self.cfg.train.learning_rate, + betas=self.cfg.train.AdamW.betas, + eps=self.cfg.train.AdamW.eps, + ) + optimizer_d = torch.optim.AdamW( + self.model["discriminator"].parameters(), + self.cfg.train.learning_rate, + betas=self.cfg.train.AdamW.betas, + eps=self.cfg.train.AdamW.eps, + ) + optimizer = {"optimizer_g": optimizer_g, "optimizer_d": optimizer_d} + + return optimizer + + def _build_scheduler(self): + scheduler_g = ExponentialLR( + self.optimizer["optimizer_g"], + gamma=self.cfg.train.lr_decay, + last_epoch=self.epoch - 1, + ) + scheduler_d = ExponentialLR( + self.optimizer["optimizer_d"], + gamma=self.cfg.train.lr_decay, + last_epoch=self.epoch - 1, + ) + + scheduler = {"scheduler_g": scheduler_g, "scheduler_d": scheduler_d} + return scheduler + + def _build_criterion(self): + class GeneratorLoss(nn.Module): + def __init__(self, cfg): + super(GeneratorLoss, self).__init__() + self.cfg = cfg + self.l1_loss = nn.L1Loss() + + def generator_loss(self, disc_outputs): + loss = 0 + gen_losses = [] + for dg in disc_outputs: + dg = dg.float() + l = torch.mean((1 - dg) ** 2) + gen_losses.append(l) + loss += l + + return loss, gen_losses + + def feature_loss(self, fmap_r, fmap_g): + loss = 0 + for dr, dg in zip(fmap_r, fmap_g): + for rl, gl in zip(dr, dg): + rl = rl.float().detach() + gl = gl.float() + loss += torch.mean(torch.abs(rl - gl)) + + return loss * 2 + + def kl_loss(self, z_p, logs_q, m_p, logs_p, z_mask): + """ + z_p, logs_q: [b, h, t_t] + m_p, logs_p: [b, h, t_t] + """ + z_p = z_p.float() + logs_q = logs_q.float() + m_p = m_p.float() + logs_p = logs_p.float() + z_mask = z_mask.float() + + kl = logs_p - logs_q - 0.5 + kl += 0.5 * ((z_p - m_p) ** 2) * torch.exp(-2.0 * logs_p) + kl = torch.sum(kl * z_mask) + l = kl / torch.sum(z_mask) + return l + + def forward( + self, + outputs_g, + outputs_d, + y_mel, + y_hat_mel, + ): + loss_g = {} + + # mel loss + loss_mel = self.l1_loss(y_mel, y_hat_mel) * self.cfg.train.c_mel + loss_g["loss_mel"] = loss_mel + + # kl loss + loss_kl = ( + self.kl_loss( + outputs_g["z_p"], + outputs_g["logs_q"], + outputs_g["m_p"], + outputs_g["logs_p"], + outputs_g["z_mask"], + ) + * self.cfg.train.c_kl + ) + loss_g["loss_kl"] = loss_kl + + # feature loss + loss_fm = self.feature_loss(outputs_d["fmap_rs"], outputs_d["fmap_gs"]) + loss_g["loss_fm"] = loss_fm + + # gan loss + loss_gen, losses_gen = self.generator_loss(outputs_d["y_d_hat_g"]) + loss_g["loss_gen"] = loss_gen + loss_g["loss_gen_all"] = loss_mel + loss_kl + loss_fm + loss_gen + + return loss_g + + class DiscriminatorLoss(nn.Module): + def __init__(self, cfg): + super(DiscriminatorLoss, self).__init__() + self.cfg = cfg + self.l1Loss = torch.nn.L1Loss(reduction="mean") + + def __call__(self, disc_real_outputs, disc_generated_outputs): + loss_d = {} + + loss = 0 + r_losses = [] + g_losses = [] + for dr, dg in zip(disc_real_outputs, disc_generated_outputs): + dr = dr.float() + dg = dg.float() + r_loss = torch.mean((1 - dr) ** 2) + g_loss = torch.mean(dg**2) + loss += r_loss + g_loss + r_losses.append(r_loss.item()) + g_losses.append(g_loss.item()) + + loss_d["loss_disc_all"] = loss + + return loss_d + + criterion = { + "generator": GeneratorLoss(self.cfg), + "discriminator": DiscriminatorLoss(self.cfg), + } + return criterion + + # Keep legacy unchanged + def write_summary( + self, + losses, + stats, + images={}, + audios={}, + audio_sampling_rate=24000, + tag="train", + ): + for key, value in losses.items(): + self.sw.add_scalar(tag + "/" + key, value, self.step) + self.sw.add_scalar( + "learning_rate", + self.optimizer["optimizer_g"].param_groups[0]["lr"], + self.step, + ) + + if len(images) != 0: + for key, value in images.items(): + self.sw.add_image(key, value, self.global_step, batchformats="HWC") + if len(audios) != 0: + for key, value in audios.items(): + self.sw.add_audio(key, value, self.global_step, audio_sampling_rate) + + def write_valid_summary( + self, losses, stats, images={}, audios={}, audio_sampling_rate=24000, tag="val" + ): + for key, value in losses.items(): + self.sw.add_scalar(tag + "/" + key, value, self.step) + + if len(images) != 0: + for key, value in images.items(): + self.sw.add_image(key, value, self.global_step, batchformats="HWC") + if len(audios) != 0: + for key, value in audios.items(): + self.sw.add_audio(key, value, self.global_step, audio_sampling_rate) + + def _get_state_dict(self): + state_dict = { + "generator": self.model["generator"].state_dict(), + "discriminator": self.model["discriminator"].state_dict(), + "optimizer_g": self.optimizer["optimizer_g"].state_dict(), + "optimizer_d": self.optimizer["optimizer_d"].state_dict(), + "scheduler_g": self.scheduler["scheduler_g"].state_dict(), + "scheduler_d": self.scheduler["scheduler_d"].state_dict(), + "step": self.step, + "epoch": self.epoch, + "batch_size": self.cfg.train.batch_size, + } + return state_dict + + def get_state_dict(self): + state_dict = { + "generator": self.model["generator"].state_dict(), + "discriminator": self.model["discriminator"].state_dict(), + "optimizer_g": self.optimizer["optimizer_g"].state_dict(), + "optimizer_d": self.optimizer["optimizer_d"].state_dict(), + "scheduler_g": self.scheduler["scheduler_g"].state_dict(), + "scheduler_d": self.scheduler["scheduler_d"].state_dict(), + "step": self.step, + "epoch": self.epoch, + "batch_size": self.cfg.train.batch_size, + } + return state_dict + + def load_model(self, checkpoint): + self.step = checkpoint["step"] + self.epoch = checkpoint["epoch"] + self.model["generator"].load_state_dict(checkpoint["generator"]) + self.model["discriminator"].load_state_dict(checkpoint["discriminator"]) + self.optimizer["optimizer_g"].load_state_dict(checkpoint["optimizer_g"]) + self.optimizer["optimizer_d"].load_state_dict(checkpoint["optimizer_d"]) + self.scheduler["scheduler_g"].load_state_dict(checkpoint["scheduler_g"]) + self.scheduler["scheduler_d"].load_state_dict(checkpoint["scheduler_d"]) + + @torch.inference_mode() + def _valid_step(self, batch): + r"""Testing forward step. Should return average loss of a sample over + one batch. Provoke ``_forward_step`` is recommended except for special case. + See ``_test_epoch`` for usage. + """ + + valid_losses = {} + total_loss = 0 + valid_stats = {} + + # Discriminator + # Generator output + outputs_g = self.model["generator"](batch) + + y_mel = slice_segments( + batch["mel"].transpose(1, 2), + outputs_g["ids_slice"], + self.cfg.preprocess.segment_size // self.cfg.preprocess.hop_size, + ) + y_hat_mel = mel_spectrogram_torch( + outputs_g["y_hat"].squeeze(1), self.cfg.preprocess + ) + y = slice_segments( + batch["audio"].unsqueeze(1), + outputs_g["ids_slice"] * self.cfg.preprocess.hop_size, + self.cfg.preprocess.segment_size, + ) + + # Discriminator output + outputs_d = self.model["discriminator"](y, outputs_g["y_hat"].detach()) + ## Discriminator loss + loss_d = self.criterion["discriminator"]( + outputs_d["y_d_hat_r"], outputs_d["y_d_hat_g"] + ) + valid_losses.update(loss_d) + + ## Generator + outputs_d = self.model["discriminator"](y, outputs_g["y_hat"]) + loss_g = self.criterion["generator"](outputs_g, outputs_d, y_mel, y_hat_mel) + valid_losses.update(loss_g) + + for item in valid_losses: + valid_losses[item] = valid_losses[item].item() + + total_loss = loss_g["loss_gen_all"] + loss_d["loss_disc_all"] + + return ( + total_loss.item(), + valid_losses, + valid_stats, + ) + + @torch.inference_mode() + def _valid_epoch(self): + r"""Testing epoch. Should return average loss of a batch (sample) over + one epoch. See ``train_loop`` for usage. + """ + if isinstance(self.model, dict): + for key in self.model.keys(): + self.model[key].eval() + else: + self.model.eval() + + epoch_sum_loss = 0.0 + epoch_losses = dict() + for batch in tqdm( + self.valid_dataloader, + desc=f"Validating Epoch {self.epoch}", + unit="batch", + colour="GREEN", + leave=False, + dynamic_ncols=True, + smoothing=0.04, + disable=not self.accelerator.is_main_process, + ): + total_loss, valid_losses, valid_stats = self._valid_step(batch) + epoch_sum_loss += total_loss + if isinstance(valid_losses, dict): + for key, value in valid_losses.items(): + if key not in epoch_losses.keys(): + epoch_losses[key] = value + else: + epoch_losses[key] += value + + epoch_sum_loss = epoch_sum_loss / len(self.valid_dataloader) + for key in epoch_losses.keys(): + epoch_losses[key] = epoch_losses[key] / len(self.valid_dataloader) + + self.accelerator.wait_for_everyone() + + return epoch_sum_loss, epoch_losses + + ### THIS IS MAIN ENTRY ### + def train_loop(self): + r"""Training loop. The public entry of training process.""" + # Wait everyone to prepare before we move on + self.accelerator.wait_for_everyone() + # dump config file + if self.accelerator.is_main_process: + self.__dump_cfg(self.config_save_path) + + # self.optimizer.zero_grad() + # Wait to ensure good to go + + self.accelerator.wait_for_everyone() + while self.epoch < self.max_epoch: + self.logger.info("\n") + self.logger.info("-" * 32) + self.logger.info("Epoch {}: ".format(self.epoch)) + + # Do training & validating epoch + train_total_loss, train_losses = self._train_epoch() + if isinstance(train_losses, dict): + for key, loss in train_losses.items(): + self.logger.info(" |- Train/{} Loss: {:.6f}".format(key, loss)) + self.accelerator.log( + {"Epoch/Train {} Loss".format(key): loss}, + step=self.epoch, + ) + + valid_total_loss, valid_losses = self._valid_epoch() + if isinstance(valid_losses, dict): + for key, loss in valid_losses.items(): + self.logger.info(" |- Valid/{} Loss: {:.6f}".format(key, loss)) + self.accelerator.log( + {"Epoch/Train {} Loss".format(key): loss}, + step=self.epoch, + ) + + self.logger.info(" |- Train/Loss: {:.6f}".format(train_total_loss)) + self.logger.info(" |- Valid/Loss: {:.6f}".format(valid_total_loss)) + self.accelerator.log( + { + "Epoch/Train Loss": train_total_loss, + "Epoch/Valid Loss": valid_total_loss, + }, + step=self.epoch, + ) + + self.accelerator.wait_for_everyone() + + # Check if hit save_checkpoint_stride and run_eval + run_eval = False + if self.accelerator.is_main_process: + save_checkpoint = False + hit_dix = [] + for i, num in enumerate(self.save_checkpoint_stride): + if self.epoch % num == 0: + save_checkpoint = True + hit_dix.append(i) + run_eval |= self.run_eval[i] + + self.accelerator.wait_for_everyone() + if self.accelerator.is_main_process and save_checkpoint: + path = os.path.join( + self.checkpoint_dir, + "epoch-{:04d}_step-{:07d}_loss-{:.6f}".format( + self.epoch, self.step, train_total_loss + ), + ) + self.tmp_checkpoint_save_path = path + self.accelerator.save_state(path) + + json.dump( + self.checkpoints_path, + open(os.path.join(path, "ckpts.json"), "w"), + ensure_ascii=False, + indent=4, + ) + self._save_auxiliary_states() + + # Remove old checkpoints + to_remove = [] + for idx in hit_dix: + self.checkpoints_path[idx].append(path) + while len(self.checkpoints_path[idx]) > self.keep_last[idx]: + to_remove.append((idx, self.checkpoints_path[idx].pop(0))) + + # Search conflicts + total = set() + for i in self.checkpoints_path: + total |= set(i) + do_remove = set() + for idx, path in to_remove[::-1]: + if path in total: + self.checkpoints_path[idx].insert(0, path) + else: + do_remove.add(path) + + # Remove old checkpoints + for path in do_remove: + shutil.rmtree(path, ignore_errors=True) + self.logger.debug(f"Remove old checkpoint: {path}") + + self.accelerator.wait_for_everyone() + if run_eval: + # TODO: run evaluation + pass + + # Update info for each epoch + self.epoch += 1 + + # Finish training and save final checkpoint + self.accelerator.wait_for_everyone() + if self.accelerator.is_main_process: + path = os.path.join( + self.checkpoint_dir, + "final_epoch-{:04d}_step-{:07d}_loss-{:.6f}".format( + self.epoch, self.step, valid_total_loss + ), + ) + self.tmp_checkpoint_save_path = path + self.accelerator.save_state( + os.path.join( + self.checkpoint_dir, + "final_epoch-{:04d}_step-{:07d}_loss-{:.6f}".format( + self.epoch, self.step, valid_total_loss + ), + ) + ) + + json.dump( + self.checkpoints_path, + open(os.path.join(path, "ckpts.json"), "w"), + ensure_ascii=False, + indent=4, + ) + self._save_auxiliary_states() + + self.accelerator.end_training() + + def _train_step(self, batch): + r"""Forward step for training and inference. This function is called + in ``_train_step`` & ``_test_step`` function. + """ + + train_losses = {} + total_loss = 0 + training_stats = {} + + ## Train Discriminator + # Generator output + outputs_g = self.model["generator"](batch) + + y_mel = slice_segments( + batch["mel"].transpose(1, 2), + outputs_g["ids_slice"], + self.cfg.preprocess.segment_size // self.cfg.preprocess.hop_size, + ) + y_hat_mel = mel_spectrogram_torch( + outputs_g["y_hat"].squeeze(1), self.cfg.preprocess + ) + + y = slice_segments( + # [1, 168418] -> [1, 1, 168418] + batch["audio"].unsqueeze(1), + outputs_g["ids_slice"] * self.cfg.preprocess.hop_size, + self.cfg.preprocess.segment_size, + ) + + # Discriminator output + outputs_d = self.model["discriminator"](y, outputs_g["y_hat"].detach()) + # Discriminator loss + loss_d = self.criterion["discriminator"]( + outputs_d["y_d_hat_r"], outputs_d["y_d_hat_g"] + ) + train_losses.update(loss_d) + + # BP and Grad Updated + self.optimizer["optimizer_d"].zero_grad() + self.accelerator.backward(loss_d["loss_disc_all"]) + self.optimizer["optimizer_d"].step() + + ## Train Generator + outputs_d = self.model["discriminator"](y, outputs_g["y_hat"]) + loss_g = self.criterion["generator"](outputs_g, outputs_d, y_mel, y_hat_mel) + train_losses.update(loss_g) + + # BP and Grad Updated + self.optimizer["optimizer_g"].zero_grad() + self.accelerator.backward(loss_g["loss_gen_all"]) + self.optimizer["optimizer_g"].step() + + for item in train_losses: + train_losses[item] = train_losses[item].item() + + total_loss = loss_g["loss_gen_all"] + loss_d["loss_disc_all"] + + return ( + total_loss.item(), + train_losses, + training_stats, + ) + + def _train_epoch(self): + r"""Training epoch. Should return average loss of a batch (sample) over + one epoch. See ``train_loop`` for usage. + """ + epoch_sum_loss: float = 0.0 + epoch_losses: dict = {} + epoch_step: int = 0 + for batch in tqdm( + self.train_dataloader, + desc=f"Training Epoch {self.epoch}", + unit="batch", + colour="GREEN", + leave=False, + dynamic_ncols=True, + smoothing=0.04, + disable=not self.accelerator.is_main_process, + ): + # Do training step and BP + with self.accelerator.accumulate(self.model): + total_loss, train_losses, training_stats = self._train_step(batch) + self.batch_count += 1 + + # Update info for each step + if self.batch_count % self.cfg.train.gradient_accumulation_step == 0: + epoch_sum_loss += total_loss + for key, value in train_losses.items(): + if key not in epoch_losses.keys(): + epoch_losses[key] = value + else: + epoch_losses[key] += value + + self.accelerator.log( + { + "Step/Generator Loss": train_losses["loss_gen_all"], + "Step/Discriminator Loss": train_losses["loss_disc_all"], + "Step/Generator Learning Rate": self.optimizer[ + "optimizer_d" + ].param_groups[0]["lr"], + "Step/Discriminator Learning Rate": self.optimizer[ + "optimizer_g" + ].param_groups[0]["lr"], + }, + step=self.step, + ) + self.step += 1 + epoch_step += 1 + + self.accelerator.wait_for_everyone() + + epoch_sum_loss = ( + epoch_sum_loss + / len(self.train_dataloader) + * self.cfg.train.gradient_accumulation_step + ) + + for key in epoch_losses.keys(): + epoch_losses[key] = ( + epoch_losses[key] + / len(self.train_dataloader) + * self.cfg.train.gradient_accumulation_step + ) + + return epoch_sum_loss, epoch_losses + + def __dump_cfg(self, path): + os.makedirs(os.path.dirname(path), exist_ok=True) + json5.dump( + self.cfg, + open(path, "w"), + indent=4, + sort_keys=True, + ensure_ascii=False, + quote_keys=True, + ) diff --git a/models/tta/autoencoder/__init__.py b/models/tta/autoencoder/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/models/tta/autoencoder/autoencoder.py b/models/tta/autoencoder/autoencoder.py new file mode 100644 index 0000000..fec6771 --- /dev/null +++ b/models/tta/autoencoder/autoencoder.py @@ -0,0 +1,403 @@ +# Copyright (c) 2023 Amphion. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from modules.distributions.distributions import DiagonalGaussianDistribution + + +def nonlinearity(x): + # swish + return x * torch.sigmoid(x) + + +def Normalize(in_channels): + return torch.nn.GroupNorm( + num_groups=32, num_channels=in_channels, eps=1e-6, affine=True + ) + + +class Upsample2d(nn.Module): + def __init__(self, in_channels, with_conv): + super().__init__() + self.with_conv = with_conv + if self.with_conv: + self.conv = torch.nn.Conv2d( + in_channels, in_channels, kernel_size=3, stride=1, padding=1 + ) + + def forward(self, x): + x = torch.nn.functional.interpolate(x, scale_factor=2.0, mode="nearest") + if self.with_conv: + x = self.conv(x) + return x + + +class Upsample1d(Upsample2d): + def __init__(self, in_channels, with_conv): + super().__init__(in_channels, with_conv) + if self.with_conv: + self.conv = torch.nn.Conv1d( + in_channels, in_channels, kernel_size=3, stride=1, padding=1 + ) + + +class Downsample2d(nn.Module): + def __init__(self, in_channels, with_conv): + super().__init__() + self.with_conv = with_conv + if self.with_conv: + # no asymmetric padding in torch conv, must do it ourselves + self.conv = torch.nn.Conv2d( + in_channels, in_channels, kernel_size=3, stride=2, padding=0 + ) + self.pad = (0, 1, 0, 1) + else: + self.avg_pool = nn.AvgPool2d(kernel_size=2, stride=2) + + def forward(self, x): + if self.with_conv: # bp: check self.avgpool and self.pad + x = torch.nn.functional.pad(x, self.pad, mode="constant", value=0) + x = self.conv(x) + else: + x = self.avg_pool(x) + return x + + +class Downsample1d(Downsample2d): + def __init__(self, in_channels, with_conv): + super().__init__(in_channels, with_conv) + if self.with_conv: + # no asymmetric padding in torch conv, must do it ourselves + # TODO: can we replace it just with conv2d with padding 1? + self.conv = torch.nn.Conv1d( + in_channels, in_channels, kernel_size=3, stride=2, padding=0 + ) + self.pad = (1, 1) + else: + self.avg_pool = nn.AvgPool1d(kernel_size=2, stride=2) + + +class ResnetBlock(nn.Module): + def __init__(self, *, in_channels, out_channels=None, conv_shortcut=False, dropout): + super().__init__() + self.in_channels = in_channels + out_channels = in_channels if out_channels is None else out_channels + self.out_channels = out_channels + self.use_conv_shortcut = conv_shortcut + + self.norm1 = Normalize(in_channels) + self.conv1 = torch.nn.Conv2d( + in_channels, out_channels, kernel_size=3, stride=1, padding=1 + ) + + self.norm2 = Normalize(out_channels) + self.dropout = torch.nn.Dropout(dropout) + self.conv2 = torch.nn.Conv2d( + out_channels, out_channels, kernel_size=3, stride=1, padding=1 + ) + if self.in_channels != self.out_channels: + if self.use_conv_shortcut: + self.conv_shortcut = torch.nn.Conv2d( + in_channels, out_channels, kernel_size=3, stride=1, padding=1 + ) + else: + self.nin_shortcut = torch.nn.Conv2d( + in_channels, out_channels, kernel_size=1, stride=1, padding=0 + ) + + def forward(self, x): + h = x + h = self.norm1(h) + h = nonlinearity(h) + h = self.conv1(h) + + h = self.norm2(h) + h = nonlinearity(h) + h = self.dropout(h) + h = self.conv2(h) + + if self.in_channels != self.out_channels: + if self.use_conv_shortcut: + x = self.conv_shortcut(x) + else: + x = self.nin_shortcut(x) + + return x + h + + +class ResnetBlock1d(ResnetBlock): + def __init__( + self, + *, + in_channels, + out_channels=None, + conv_shortcut=False, + dropout, + temb_channels=512 + ): + super().__init__( + in_channels=in_channels, + out_channels=out_channels, + conv_shortcut=conv_shortcut, + dropout=dropout, + ) + + self.conv1 = torch.nn.Conv1d( + in_channels, out_channels, kernel_size=3, stride=1, padding=1 + ) + self.conv2 = torch.nn.Conv1d( + out_channels, out_channels, kernel_size=3, stride=1, padding=1 + ) + if self.in_channels != self.out_channels: + if self.use_conv_shortcut: + self.conv_shortcut = torch.nn.Conv1d( + in_channels, out_channels, kernel_size=3, stride=1, padding=1 + ) + else: + self.nin_shortcut = torch.nn.Conv1d( + in_channels, out_channels, kernel_size=1, stride=1, padding=0 + ) + + +class Encoder2d(nn.Module): + def __init__( + self, + *, + ch, + ch_mult=(1, 2, 4, 8), + num_res_blocks, + dropout=0.0, + resamp_with_conv=True, + in_channels, + z_channels, + double_z=True, + **ignore_kwargs + ): + super().__init__() + self.ch = ch + self.num_resolutions = len(ch_mult) + self.num_res_blocks = num_res_blocks + self.in_channels = in_channels + + # downsampling + self.conv_in = torch.nn.Conv2d( + in_channels, self.ch, kernel_size=3, stride=1, padding=1 + ) + + in_ch_mult = (1,) + tuple(ch_mult) + self.down = nn.ModuleList() + for i_level in range(self.num_resolutions): + block = nn.ModuleList() + block_in = ch * in_ch_mult[i_level] + block_out = ch * ch_mult[i_level] + for i_block in range(self.num_res_blocks): + block.append( + ResnetBlock( + in_channels=block_in, out_channels=block_out, dropout=dropout + ) + ) + block_in = block_out + down = nn.Module() + down.block = block + if i_level != self.num_resolutions - 1: + down.downsample = Downsample2d(block_in, resamp_with_conv) + self.down.append(down) + + # middle + self.mid = nn.Module() + self.mid.block_1 = ResnetBlock( + in_channels=block_in, out_channels=block_in, dropout=dropout + ) + self.mid.block_2 = ResnetBlock( + in_channels=block_in, out_channels=block_in, dropout=dropout + ) + + # end + self.norm_out = Normalize(block_in) + self.conv_out = torch.nn.Conv2d( + block_in, + 2 * z_channels if double_z else z_channels, + kernel_size=3, + stride=1, + padding=1, + ) + + def forward(self, x): + # downsampling + hs = [self.conv_in(x)] + for i_level in range(self.num_resolutions): + for i_block in range(self.num_res_blocks): + h = self.down[i_level].block[i_block](hs[-1]) + hs.append(h) + if i_level != self.num_resolutions - 1: + hs.append(self.down[i_level].downsample(hs[-1])) + + # middle + h = hs[-1] + h = self.mid.block_1(h) + h = self.mid.block_2(h) + + # end + h = self.norm_out(h) + h = nonlinearity(h) + h = self.conv_out(h) + return h + + +# TODO: Encoder1d +class Encoder1d(Encoder2d): ... + + +class Decoder2d(nn.Module): + def __init__( + self, + *, + ch, + out_ch, + ch_mult=(1, 2, 4, 8), + num_res_blocks, + dropout=0.0, + resamp_with_conv=True, + in_channels, + z_channels, + give_pre_end=False, + **ignorekwargs + ): + super().__init__() + self.ch = ch + self.num_resolutions = len(ch_mult) + self.num_res_blocks = num_res_blocks + self.in_channels = in_channels + self.give_pre_end = give_pre_end + + # compute in_ch_mult, block_in and curr_res at lowest res + in_ch_mult = (1,) + tuple(ch_mult) + block_in = ch * ch_mult[self.num_resolutions - 1] + # self.z_shape = (1,z_channels,curr_res,curr_res) + # print("Working with z of shape {} = {} dimensions.".format( + # self.z_shape, np.prod(self.z_shape))) + + # z to block_in + self.conv_in = torch.nn.Conv2d( + z_channels, block_in, kernel_size=3, stride=1, padding=1 + ) + + # middle + self.mid = nn.Module() + self.mid.block_1 = ResnetBlock( + in_channels=block_in, out_channels=block_in, dropout=dropout + ) + self.mid.block_2 = ResnetBlock( + in_channels=block_in, out_channels=block_in, dropout=dropout + ) + + # upsampling + self.up = nn.ModuleList() + for i_level in reversed(range(self.num_resolutions)): + block = nn.ModuleList() + attn = nn.ModuleList() + block_out = ch * ch_mult[i_level] + for i_block in range(self.num_res_blocks + 1): + block.append( + ResnetBlock( + in_channels=block_in, out_channels=block_out, dropout=dropout + ) + ) + block_in = block_out + up = nn.Module() + up.block = block + up.attn = attn + if i_level != 0: + up.upsample = Upsample2d(block_in, resamp_with_conv) + self.up.insert(0, up) # prepend to get consistent order + + # end + self.norm_out = Normalize(block_in) + self.conv_out = torch.nn.Conv2d( + block_in, out_ch, kernel_size=3, stride=1, padding=1 + ) + + def forward(self, z): + self.last_z_shape = z.shape + + # z to block_in + h = self.conv_in(z) + + # middle + h = self.mid.block_1(h) + h = self.mid.block_2(h) + + # upsampling + for i_level in reversed(range(self.num_resolutions)): + for i_block in range(self.num_res_blocks + 1): + h = self.up[i_level].block[i_block](h) + if i_level != 0: + h = self.up[i_level].upsample(h) + + # end + if self.give_pre_end: + return h + + h = self.norm_out(h) + h = nonlinearity(h) + h = self.conv_out(h) + return h + + +# TODO: decoder1d +class Decoder1d(Decoder2d): ... + + +class AutoencoderKL(nn.Module): + def __init__(self, cfg): + super().__init__() + self.cfg = cfg + self.encoder = Encoder2d( + ch=cfg.ch, + ch_mult=cfg.ch_mult, + num_res_blocks=cfg.num_res_blocks, + in_channels=cfg.in_channels, + z_channels=cfg.z_channels, + double_z=cfg.double_z, + ) + self.decoder = Decoder2d( + ch=cfg.ch, + ch_mult=cfg.ch_mult, + num_res_blocks=cfg.num_res_blocks, + out_ch=cfg.out_ch, + z_channels=cfg.z_channels, + in_channels=None, + ) + assert self.cfg.double_z + + self.quant_conv = torch.nn.Conv2d(2 * cfg.z_channels, 2 * cfg.z_channels, 1) + self.post_quant_conv = torch.nn.Conv2d(cfg.z_channels, cfg.z_channels, 1) + self.embed_dim = cfg.z_channels + + def encode(self, x): + h = self.encoder(x) + moments = self.quant_conv(h) + posterior = DiagonalGaussianDistribution(moments) + return posterior + + def decode(self, z): + z = self.post_quant_conv(z) + dec = self.decoder(z) + return dec + + def forward(self, input, sample_posterior=True): + posterior = self.encode(input) + if sample_posterior: + z = posterior.sample() + else: + z = posterior.mode() + dec = self.decode(z) + return dec, posterior + + def get_last_layer(self): + return self.decoder.conv_out.weight diff --git a/models/tta/autoencoder/autoencoder_dataset.py b/models/tta/autoencoder/autoencoder_dataset.py new file mode 100644 index 0000000..8b9b4bd --- /dev/null +++ b/models/tta/autoencoder/autoencoder_dataset.py @@ -0,0 +1,112 @@ +# Copyright (c) 2023 Amphion. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +import random +import torch +from torch.nn.utils.rnn import pad_sequence +from utils.data_utils import * +from models.base.base_dataset import ( + BaseOfflineCollator, + BaseOfflineDataset, + BaseTestDataset, + BaseTestCollator, +) +import librosa + + +class AutoencoderKLDataset(BaseOfflineDataset): + def __init__(self, cfg, dataset, is_valid=False): + BaseOfflineDataset.__init__(self, cfg, dataset, is_valid=is_valid) + + cfg = self.cfg + + # utt2melspec + if cfg.preprocess.use_melspec: + self.utt2melspec_path = {} + for utt_info in self.metadata: + dataset = utt_info["Dataset"] + uid = utt_info["Uid"] + utt = "{}_{}".format(dataset, uid) + + self.utt2melspec_path[utt] = os.path.join( + cfg.preprocess.processed_dir, + dataset, + cfg.preprocess.melspec_dir, + uid + ".npy", + ) + + # utt2wav + if cfg.preprocess.use_wav: + self.utt2wav_path = {} + for utt_info in self.metadata: + dataset = utt_info["Dataset"] + uid = utt_info["Uid"] + utt = "{}_{}".format(dataset, uid) + + self.utt2wav_path[utt] = os.path.join( + cfg.preprocess.processed_dir, + dataset, + cfg.preprocess.wav_dir, + uid + ".wav", + ) + + def __getitem__(self, index): + # melspec: (n_mels, T) + # wav: (T,) + + single_feature = BaseOfflineDataset.__getitem__(self, index) + + utt_info = self.metadata[index] + dataset = utt_info["Dataset"] + uid = utt_info["Uid"] + utt = "{}_{}".format(dataset, uid) + + if self.cfg.preprocess.use_melspec: + single_feature["melspec"] = np.load(self.utt2melspec_path[utt]) + + if self.cfg.preprocess.use_wav: + wav, sr = librosa.load( + self.utt2wav_path[utt], sr=16000 + ) # hard coding for 16KHz... + single_feature["wav"] = wav + + return single_feature + + def __len__(self): + return len(self.metadata) + + def __len__(self): + return len(self.metadata) + + +class AutoencoderKLCollator(BaseOfflineCollator): + def __init__(self, cfg): + BaseOfflineCollator.__init__(self, cfg) + + def __call__(self, batch): + # mel: (B, n_mels, T) + # wav (option): (B, T) + + packed_batch_features = dict() + + for key in batch[0].keys(): + if key == "melspec": + packed_batch_features["melspec"] = torch.from_numpy( + np.array([b["melspec"][:, :624] for b in batch]) + ) + + if key == "wav": + values = [torch.from_numpy(b[key]) for b in batch] + packed_batch_features[key] = pad_sequence( + values, batch_first=True, padding_value=0 + ) + + return packed_batch_features + + +class AutoencoderKLTestDataset(BaseTestDataset): ... + + +class AutoencoderKLTestCollator(BaseTestCollator): ... diff --git a/models/tta/autoencoder/autoencoder_loss.py b/models/tta/autoencoder/autoencoder_loss.py new file mode 100644 index 0000000..5916aa3 --- /dev/null +++ b/models/tta/autoencoder/autoencoder_loss.py @@ -0,0 +1,305 @@ +# Copyright (c) 2023 Amphion. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +import torch +import torch.nn as nn +import functools +import torch.nn.functional as F + + +def hinge_d_loss(logits_real, logits_fake): + loss_real = torch.mean(F.relu(1.0 - logits_real)) + loss_fake = torch.mean(F.relu(1.0 + logits_fake)) + d_loss = 0.5 * (loss_real + loss_fake) + return d_loss + + +def vanilla_d_loss(logits_real, logits_fake): + d_loss = 0.5 * ( + torch.mean(F.softplus(-logits_real)) + torch.mean(F.softplus(logits_fake)) + ) + return d_loss + + +def adopt_weight(weight, global_step, threshold=0, value=0.0): + if global_step < threshold: + weight = value + return weight + + +class ActNorm(nn.Module): + def __init__( + self, num_features, logdet=False, affine=True, allow_reverse_init=False + ): + assert affine + super().__init__() + self.logdet = logdet + self.loc = nn.Parameter(torch.zeros(1, num_features, 1, 1)) + self.scale = nn.Parameter(torch.ones(1, num_features, 1, 1)) + self.allow_reverse_init = allow_reverse_init + + self.register_buffer("initialized", torch.tensor(0, dtype=torch.uint8)) + + def initialize(self, input): + with torch.no_grad(): + flatten = input.permute(1, 0, 2, 3).contiguous().view(input.shape[1], -1) + mean = ( + flatten.mean(1) + .unsqueeze(1) + .unsqueeze(2) + .unsqueeze(3) + .permute(1, 0, 2, 3) + ) + std = ( + flatten.std(1) + .unsqueeze(1) + .unsqueeze(2) + .unsqueeze(3) + .permute(1, 0, 2, 3) + ) + + self.loc.data.copy_(-mean) + self.scale.data.copy_(1 / (std + 1e-6)) + + def forward(self, input, reverse=False): + if reverse: + return self.reverse(input) + if len(input.shape) == 2: + input = input[:, :, None, None] + squeeze = True + else: + squeeze = False + + _, _, height, width = input.shape + + if self.training and self.initialized.item() == 0: + self.initialize(input) + self.initialized.fill_(1) + + h = self.scale * (input + self.loc) + + if squeeze: + h = h.squeeze(-1).squeeze(-1) + + if self.logdet: + log_abs = torch.log(torch.abs(self.scale)) + logdet = height * width * torch.sum(log_abs) + logdet = logdet * torch.ones(input.shape[0]).to(input) + return h, logdet + + return h + + def reverse(self, output): + if self.training and self.initialized.item() == 0: + if not self.allow_reverse_init: + raise RuntimeError( + "Initializing ActNorm in reverse direction is " + "disabled by default. Use allow_reverse_init=True to enable." + ) + else: + self.initialize(output) + self.initialized.fill_(1) + + if len(output.shape) == 2: + output = output[:, :, None, None] + squeeze = True + else: + squeeze = False + + h = output / self.scale - self.loc + + if squeeze: + h = h.squeeze(-1).squeeze(-1) + return h + + +def weights_init(m): + classname = m.__class__.__name__ + if classname.find("Conv") != -1: + nn.init.normal_(m.weight.data, 0.0, 0.02) + elif classname.find("BatchNorm") != -1: + nn.init.normal_(m.weight.data, 1.0, 0.02) + nn.init.constant_(m.bias.data, 0) + + +class NLayerDiscriminator(nn.Module): + """Defines a PatchGAN discriminator as in Pix2Pix + --> see https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix/blob/master/models/networks.py + """ + + def __init__(self, input_nc=3, ndf=64, n_layers=3, use_actnorm=False): + """Construct a PatchGAN discriminator + Parameters: + input_nc (int) -- the number of channels in input images + ndf (int) -- the number of filters in the last conv layer + n_layers (int) -- the number of conv layers in the discriminator + norm_layer -- normalization layer + """ + super(NLayerDiscriminator, self).__init__() + if not use_actnorm: + norm_layer = nn.BatchNorm2d + else: + norm_layer = ActNorm + if ( + type(norm_layer) == functools.partial + ): # no need to use bias as BatchNorm2d has affine parameters + use_bias = norm_layer.func != nn.BatchNorm2d + else: + use_bias = norm_layer != nn.BatchNorm2d + + kw = 4 + padw = 1 + sequence = [ + nn.Conv2d(input_nc, ndf, kernel_size=kw, stride=2, padding=padw), + nn.LeakyReLU(0.2, True), + ] + nf_mult = 1 + nf_mult_prev = 1 + for n in range(1, n_layers): # gradually increase the number of filters + nf_mult_prev = nf_mult + nf_mult = min(2**n, 8) + sequence += [ + nn.Conv2d( + ndf * nf_mult_prev, + ndf * nf_mult, + kernel_size=kw, + stride=2, + padding=padw, + bias=use_bias, + ), + norm_layer(ndf * nf_mult), + nn.LeakyReLU(0.2, True), + ] + + nf_mult_prev = nf_mult + nf_mult = min(2**n_layers, 8) + sequence += [ + nn.Conv2d( + ndf * nf_mult_prev, + ndf * nf_mult, + kernel_size=kw, + stride=1, + padding=padw, + bias=use_bias, + ), + norm_layer(ndf * nf_mult), + nn.LeakyReLU(0.2, True), + ] + + sequence += [ + nn.Conv2d(ndf * nf_mult, 1, kernel_size=kw, stride=1, padding=padw) + ] # output 1 channel prediction map + self.main = nn.Sequential(*sequence) + + def forward(self, input): + """Standard forward.""" + return self.main(input) + + +class AutoencoderLossWithDiscriminator(nn.Module): + def __init__(self, cfg): + super().__init__() + self.cfg = cfg + self.kl_weight = cfg.kl_weight + self.logvar = nn.Parameter(torch.ones(size=()) * cfg.logvar_init) + + self.discriminator = NLayerDiscriminator( + input_nc=cfg.disc_in_channels, + n_layers=cfg.disc_num_layers, + use_actnorm=cfg.use_actnorm, + ).apply(weights_init) + + self.discriminator_iter_start = cfg.disc_start + self.discriminator_weight = cfg.disc_weight + self.disc_factor = cfg.disc_factor + self.disc_loss = hinge_d_loss + + def calculate_adaptive_weight(self, nll_loss, g_loss, last_layer): + nll_grads = torch.autograd.grad(nll_loss, last_layer, retain_graph=True)[0] + g_grads = torch.autograd.grad(g_loss, last_layer, retain_graph=True)[0] + + d_weight = torch.norm(nll_grads) / (torch.norm(g_grads) + 1e-4) + d_weight = torch.clamp( + d_weight, self.cfg.min_adapt_d_weight, self.cfg.max_adapt_d_weight + ).detach() + d_weight = d_weight * self.discriminator_weight + return d_weight + + def forward( + self, + inputs, + reconstructions, + posteriors, + optimizer_idx, + global_step, + last_layer, + split="train", + weights=None, + ): + rec_loss = torch.abs( + inputs.contiguous() - reconstructions.contiguous() + ) # l1 loss + nll_loss = rec_loss / torch.exp(self.logvar) + self.logvar + weighted_nll_loss = nll_loss + if weights is not None: + weighted_nll_loss = weights * nll_loss + # weighted_nll_loss = torch.sum(weighted_nll_loss) / weighted_nll_loss.shape[0] + weighted_nll_loss = torch.mean(weighted_nll_loss) + # nll_loss = torch.sum(nll_loss) / nll_loss.shape[0] + nll_loss = torch.mean(nll_loss) + kl_loss = posteriors.kl() + kl_loss = torch.sum(kl_loss) / kl_loss.shape[0] + # ? kl_loss = torch.mean(kl_loss) + + # now the GAN part + if optimizer_idx == 0: + logits_fake = self.discriminator(reconstructions.contiguous()) + g_loss = -torch.mean(logits_fake) + + if self.disc_factor > 0.0: + try: + d_weight = self.calculate_adaptive_weight( + nll_loss, g_loss, last_layer=last_layer + ) + except RuntimeError: + assert not self.training + d_weight = torch.tensor(0.0) + else: + d_weight = torch.tensor(0.0) + + disc_factor = adopt_weight( + self.disc_factor, global_step, threshold=self.discriminator_iter_start + ) + + total_loss = ( + weighted_nll_loss + + self.kl_weight * kl_loss + + d_weight * disc_factor * g_loss + ) + + return { + "loss": total_loss, + "kl_loss": kl_loss, + "rec_loss": rec_loss.mean(), + "nll_loss": nll_loss, + "g_loss": g_loss, + "d_weight": d_weight, + "disc_factor": torch.tensor(disc_factor), + } + + if optimizer_idx == 1: + logits_real = self.discriminator(inputs.contiguous().detach()) + logits_fake = self.discriminator(reconstructions.contiguous().detach()) + + disc_factor = adopt_weight( + self.disc_factor, global_step, threshold=self.discriminator_iter_start + ) + d_loss = disc_factor * self.disc_loss(logits_real, logits_fake) + + return { + "d_loss": d_loss, + "logits_real": logits_real.mean(), + "logits_fake": logits_fake.mean(), + } diff --git a/models/tta/autoencoder/autoencoder_trainer.py b/models/tta/autoencoder/autoencoder_trainer.py new file mode 100644 index 0000000..1faf02f --- /dev/null +++ b/models/tta/autoencoder/autoencoder_trainer.py @@ -0,0 +1,187 @@ +# Copyright (c) 2023 Amphion. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +import torch +from models.base.base_trainer import BaseTrainer +from models.tta.autoencoder.autoencoder_dataset import ( + AutoencoderKLDataset, + AutoencoderKLCollator, +) +from models.tta.autoencoder.autoencoder import AutoencoderKL +from models.tta.autoencoder.autoencoder_loss import AutoencoderLossWithDiscriminator +from torch.optim import Adam, AdamW +from torch.optim.lr_scheduler import ReduceLROnPlateau +from torch.nn import MSELoss, L1Loss +import torch.nn.functional as F +from torch.utils.data import ConcatDataset, DataLoader + + +class AutoencoderKLTrainer(BaseTrainer): + def __init__(self, args, cfg): + BaseTrainer.__init__(self, args, cfg) + self.cfg = cfg + self.save_config_file() + + def build_dataset(self): + return AutoencoderKLDataset, AutoencoderKLCollator + + def build_optimizer(self): + opt_ae = torch.optim.AdamW(self.model.parameters(), **self.cfg.train.adam) + opt_disc = torch.optim.AdamW( + self.criterion.discriminator.parameters(), **self.cfg.train.adam + ) + optimizer = {"opt_ae": opt_ae, "opt_disc": opt_disc} + return optimizer + + def build_data_loader(self): + Dataset, Collator = self.build_dataset() + # build dataset instance for each dataset and combine them by ConcatDataset + datasets_list = [] + for dataset in self.cfg.dataset: + subdataset = Dataset(self.cfg, dataset, is_valid=False) + datasets_list.append(subdataset) + train_dataset = ConcatDataset(datasets_list) + + train_collate = Collator(self.cfg) + + # use batch_sampler argument instead of (sampler, shuffle, drop_last, batch_size) + train_loader = DataLoader( + train_dataset, + collate_fn=train_collate, + num_workers=self.args.num_workers, + batch_size=self.cfg.train.batch_size, + pin_memory=False, + ) + if not self.cfg.train.ddp or self.args.local_rank == 0: + datasets_list = [] + for dataset in self.cfg.dataset: + subdataset = Dataset(self.cfg, dataset, is_valid=True) + datasets_list.append(subdataset) + valid_dataset = ConcatDataset(datasets_list) + valid_collate = Collator(self.cfg) + + valid_loader = DataLoader( + valid_dataset, + collate_fn=valid_collate, + num_workers=1, + batch_size=self.cfg.train.batch_size, + ) + else: + raise NotImplementedError("DDP is not supported yet.") + # valid_loader = None + data_loader = {"train": train_loader, "valid": valid_loader} + return data_loader + + # TODO: check it... + def build_scheduler(self): + return None + # return ReduceLROnPlateau(self.optimizer["opt_ae"], **self.cfg.train.lronPlateau) + + def write_summary(self, losses, stats): + for key, value in losses.items(): + self.sw.add_scalar(key, value, self.step) + + def write_valid_summary(self, losses, stats): + for key, value in losses.items(): + self.sw.add_scalar(key, value, self.step) + + def build_criterion(self): + return AutoencoderLossWithDiscriminator(self.cfg.model.loss) + + def get_state_dict(self): + if self.scheduler != None: + state_dict = { + "model": self.model.state_dict(), + "optimizer_ae": self.optimizer["opt_ae"].state_dict(), + "optimizer_disc": self.optimizer["opt_disc"].state_dict(), + "scheduler": self.scheduler.state_dict(), + "step": self.step, + "epoch": self.epoch, + "batch_size": self.cfg.train.batch_size, + } + else: + state_dict = { + "model": self.model.state_dict(), + "optimizer_ae": self.optimizer["opt_ae"].state_dict(), + "optimizer_disc": self.optimizer["opt_disc"].state_dict(), + "step": self.step, + "epoch": self.epoch, + "batch_size": self.cfg.train.batch_size, + } + return state_dict + + def load_model(self, checkpoint): + self.step = checkpoint["step"] + self.epoch = checkpoint["epoch"] + + self.model.load_state_dict(checkpoint["model"]) + self.optimizer["opt_ae"].load_state_dict(checkpoint["optimizer_ae"]) + self.optimizer["opt_disc"].load_state_dict(checkpoint["optimizer_disc"]) + if self.scheduler != None: + self.scheduler.load_state_dict(checkpoint["scheduler"]) + + def build_model(self): + self.model = AutoencoderKL(self.cfg.model.autoencoderkl) + return self.model + + # TODO: train step + def train_step(self, data): + global_step = self.step + optimizer_idx = global_step % 2 + + train_losses = {} + total_loss = 0 + train_states = {} + + inputs = data["melspec"].unsqueeze(1) # (B, 80, T) -> (B, 1, 80, T) + reconstructions, posterior = self.model(inputs) + # train_stats.update(stat) + + train_losses = self.criterion( + inputs=inputs, + reconstructions=reconstructions, + posteriors=posterior, + optimizer_idx=optimizer_idx, + global_step=global_step, + last_layer=self.model.get_last_layer(), + split="train", + ) + + if optimizer_idx == 0: + total_loss = train_losses["loss"] + self.optimizer["opt_ae"].zero_grad() + total_loss.backward() + self.optimizer["opt_ae"].step() + + else: + total_loss = train_losses["d_loss"] + self.optimizer["opt_disc"].zero_grad() + total_loss.backward() + self.optimizer["opt_disc"].step() + + for item in train_losses: + train_losses[item] = train_losses[item].item() + + return train_losses, train_states, total_loss.item() + + # TODO: eval step + @torch.no_grad() + def eval_step(self, data, index): + valid_loss = {} + total_valid_loss = 0 + valid_stats = {} + + inputs = data["melspec"].unsqueeze(1) # (B, 80, T) -> (B, 1, 80, T) + reconstructions, posterior = self.model(inputs) + + loss = F.l1_loss(inputs, reconstructions) + valid_loss["loss"] = loss + + total_valid_loss += loss + + for item in valid_loss: + valid_loss[item] = valid_loss[item].item() + + return valid_loss, valid_stats, total_valid_loss.item() diff --git a/models/tta/ldm/__init__.py b/models/tta/ldm/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/models/tta/ldm/attention.py b/models/tta/ldm/attention.py new file mode 100644 index 0000000..8e241fc --- /dev/null +++ b/models/tta/ldm/attention.py @@ -0,0 +1,329 @@ +# Copyright (c) 2023 Amphion. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +from inspect import isfunction +import math +import torch +import torch.nn.functional as F +from torch import nn, einsum +from einops import rearrange, repeat + + +class CheckpointFunction(torch.autograd.Function): + @staticmethod + def forward(ctx, run_function, length, *args): + ctx.run_function = run_function + ctx.input_tensors = list(args[:length]) + ctx.input_params = list(args[length:]) + + with torch.no_grad(): + output_tensors = ctx.run_function(*ctx.input_tensors) + return output_tensors + + @staticmethod + def backward(ctx, *output_grads): + ctx.input_tensors = [x.detach().requires_grad_(True) for x in ctx.input_tensors] + with torch.enable_grad(): + # Fixes a bug where the first op in run_function modifies the + # Tensor storage in place, which is not allowed for detach()'d + # Tensors. + shallow_copies = [x.view_as(x) for x in ctx.input_tensors] + output_tensors = ctx.run_function(*shallow_copies) + input_grads = torch.autograd.grad( + output_tensors, + ctx.input_tensors + ctx.input_params, + output_grads, + allow_unused=True, + ) + del ctx.input_tensors + del ctx.input_params + del output_tensors + return (None, None) + input_grads + + +def checkpoint(func, inputs, params, flag): + """ + Evaluate a function without caching intermediate activations, allowing for + reduced memory at the expense of extra compute in the backward pass. + :param func: the function to evaluate. + :param inputs: the argument sequence to pass to `func`. + :param params: a sequence of parameters `func` depends on but does not + explicitly take as arguments. + :param flag: if False, disable gradient checkpointing. + """ + if flag: + args = tuple(inputs) + tuple(params) + return CheckpointFunction.apply(func, len(inputs), *args) + else: + return func(*inputs) + + +def exists(val): + return val is not None + + +def uniq(arr): + return {el: True for el in arr}.keys() + + +def default(val, d): + if exists(val): + return val + return d() if isfunction(d) else d + + +def max_neg_value(t): + return -torch.finfo(t.dtype).max + + +def init_(tensor): + dim = tensor.shape[-1] + std = 1 / math.sqrt(dim) + tensor.uniform_(-std, std) + return tensor + + +# feedforward +class GEGLU(nn.Module): + def __init__(self, dim_in, dim_out): + super().__init__() + self.proj = nn.Linear(dim_in, dim_out * 2) + + def forward(self, x): + x, gate = self.proj(x).chunk(2, dim=-1) + return x * F.gelu(gate) + + +class FeedForward(nn.Module): + def __init__(self, dim, dim_out=None, mult=4, glu=False, dropout=0.0): + super().__init__() + inner_dim = int(dim * mult) + dim_out = default(dim_out, dim) + project_in = ( + nn.Sequential(nn.Linear(dim, inner_dim), nn.GELU()) + if not glu + else GEGLU(dim, inner_dim) + ) + + self.net = nn.Sequential( + project_in, nn.Dropout(dropout), nn.Linear(inner_dim, dim_out) + ) + + def forward(self, x): + return self.net(x) + + +def zero_module(module): + """ + Zero out the parameters of a module and return it. + """ + for p in module.parameters(): + p.detach().zero_() + return module + + +def Normalize(in_channels): + return torch.nn.GroupNorm( + num_groups=32, num_channels=in_channels, eps=1e-6, affine=True + ) + + +class LinearAttention(nn.Module): + def __init__(self, dim, heads=4, dim_head=32): + super().__init__() + self.heads = heads + hidden_dim = dim_head * heads + self.to_qkv = nn.Conv2d(dim, hidden_dim * 3, 1, bias=False) + self.to_out = nn.Conv2d(hidden_dim, dim, 1) + + def forward(self, x): + b, c, h, w = x.shape + qkv = self.to_qkv(x) + q, k, v = rearrange( + qkv, "b (qkv heads c) h w -> qkv b heads c (h w)", heads=self.heads, qkv=3 + ) + k = k.softmax(dim=-1) + context = torch.einsum("bhdn,bhen->bhde", k, v) + out = torch.einsum("bhde,bhdn->bhen", context, q) + out = rearrange( + out, "b heads c (h w) -> b (heads c) h w", heads=self.heads, h=h, w=w + ) + return self.to_out(out) + + +class SpatialSelfAttention(nn.Module): + def __init__(self, in_channels): + super().__init__() + self.in_channels = in_channels + + self.norm = Normalize(in_channels) + self.q = torch.nn.Conv2d( + in_channels, in_channels, kernel_size=1, stride=1, padding=0 + ) + self.k = torch.nn.Conv2d( + in_channels, in_channels, kernel_size=1, stride=1, padding=0 + ) + self.v = torch.nn.Conv2d( + in_channels, in_channels, kernel_size=1, stride=1, padding=0 + ) + self.proj_out = torch.nn.Conv2d( + in_channels, in_channels, kernel_size=1, stride=1, padding=0 + ) + + def forward(self, x): + h_ = x + h_ = self.norm(h_) + q = self.q(h_) + k = self.k(h_) + v = self.v(h_) + + # compute attention + b, c, h, w = q.shape + q = rearrange(q, "b c h w -> b (h w) c") + k = rearrange(k, "b c h w -> b c (h w)") + w_ = torch.einsum("bij,bjk->bik", q, k) + + w_ = w_ * (int(c) ** (-0.5)) + w_ = torch.nn.functional.softmax(w_, dim=2) + + # attend to values + v = rearrange(v, "b c h w -> b c (h w)") + w_ = rearrange(w_, "b i j -> b j i") + h_ = torch.einsum("bij,bjk->bik", v, w_) + h_ = rearrange(h_, "b c (h w) -> b c h w", h=h) + h_ = self.proj_out(h_) + + return x + h_ + + +class CrossAttention(nn.Module): + def __init__(self, query_dim, context_dim=None, heads=8, dim_head=64, dropout=0.0): + super().__init__() + inner_dim = dim_head * heads + context_dim = default(context_dim, query_dim) + + self.scale = dim_head**-0.5 + self.heads = heads + + self.to_q = nn.Linear(query_dim, inner_dim, bias=False) + self.to_k = nn.Linear(context_dim, inner_dim, bias=False) + self.to_v = nn.Linear(context_dim, inner_dim, bias=False) + + self.to_out = nn.Sequential( + nn.Linear(inner_dim, query_dim), nn.Dropout(dropout) + ) + + def forward(self, x, context=None, mask=None): + h = self.heads + + q = self.to_q(x) + context = default(context, x) + k = self.to_k(context) + v = self.to_v(context) + + q, k, v = map(lambda t: rearrange(t, "b n (h d) -> (b h) n d", h=h), (q, k, v)) + + sim = einsum("b i d, b j d -> b i j", q, k) * self.scale + + if exists(mask): + mask = rearrange(mask, "b ... -> b (...)") + max_neg_value = -torch.finfo(sim.dtype).max + mask = repeat(mask, "b j -> (b h) () j", h=h) + sim.masked_fill_(~mask, max_neg_value) + + # attention, what we cannot get enough of + attn = sim.softmax(dim=-1) + + out = einsum("b i j, b j d -> b i d", attn, v) + out = rearrange(out, "(b h) n d -> b n (h d)", h=h) + return self.to_out(out) + + +class BasicTransformerBlock(nn.Module): + def __init__( + self, + dim, + n_heads, + d_head, + dropout=0.0, + context_dim=None, + gated_ff=True, + checkpoint=True, + ): + super().__init__() + self.attn1 = CrossAttention( + query_dim=dim, heads=n_heads, dim_head=d_head, dropout=dropout + ) # is a self-attention + self.ff = FeedForward(dim, dropout=dropout, glu=gated_ff) + self.attn2 = CrossAttention( + query_dim=dim, + context_dim=context_dim, + heads=n_heads, + dim_head=d_head, + dropout=dropout, + ) # is self-attn if context is none + self.norm1 = nn.LayerNorm(dim) + self.norm2 = nn.LayerNorm(dim) + self.norm3 = nn.LayerNorm(dim) + self.checkpoint = checkpoint + + def forward(self, x, context=None): + return checkpoint( + self._forward, (x, context), self.parameters(), self.checkpoint + ) + + def _forward(self, x, context=None): + x = self.attn1(self.norm1(x)) + x + x = self.attn2(self.norm2(x), context=context) + x + x = self.ff(self.norm3(x)) + x + return x + + +class SpatialTransformer(nn.Module): + """ + Transformer block for image-like data. + First, project the input (aka embedding) + and reshape to b, t, d. + Then apply standard transformer action. + Finally, reshape to image + """ + + def __init__( + self, in_channels, n_heads, d_head, depth=1, dropout=0.0, context_dim=None + ): + super().__init__() + self.in_channels = in_channels + inner_dim = n_heads * d_head + self.norm = Normalize(in_channels) + + self.proj_in = nn.Conv2d( + in_channels, inner_dim, kernel_size=1, stride=1, padding=0 + ) + + self.transformer_blocks = nn.ModuleList( + [ + BasicTransformerBlock( + inner_dim, n_heads, d_head, dropout=dropout, context_dim=context_dim + ) + for d in range(depth) + ] + ) + + self.proj_out = zero_module( + nn.Conv2d(inner_dim, in_channels, kernel_size=1, stride=1, padding=0) + ) + + def forward(self, x, context=None): + # note: if no context is given, cross-attention defaults to self-attention + b, c, h, w = x.shape + x_in = x + x = self.norm(x) + x = self.proj_in(x) + x = rearrange(x, "b c h w -> b (h w) c") + for block in self.transformer_blocks: + x = block(x, context=context) + x = rearrange(x, "b (h w) c -> b c h w", h=h, w=w) + x = self.proj_out(x) + return x + x_in diff --git a/models/tta/ldm/audioldm.py b/models/tta/ldm/audioldm.py new file mode 100644 index 0000000..c799fec --- /dev/null +++ b/models/tta/ldm/audioldm.py @@ -0,0 +1,928 @@ +# Copyright (c) 2023 Amphion. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +from abc import abstractmethod +from functools import partial +import math +from typing import Iterable + +import os +import torch +import torch.nn as nn +import torch.nn.functional as F +import numpy as np +from einops import repeat + +from models.tta.ldm.attention import SpatialTransformer + +# from attention import SpatialTransformer + + +class CheckpointFunction(torch.autograd.Function): + @staticmethod + def forward(ctx, run_function, length, *args): + ctx.run_function = run_function + ctx.input_tensors = list(args[:length]) + ctx.input_params = list(args[length:]) + + with torch.no_grad(): + output_tensors = ctx.run_function(*ctx.input_tensors) + return output_tensors + + @staticmethod + def backward(ctx, *output_grads): + ctx.input_tensors = [x.detach().requires_grad_(True) for x in ctx.input_tensors] + with torch.enable_grad(): + # Fixes a bug where the first op in run_function modifies the + # Tensor storage in place, which is not allowed for detach()'d + # Tensors. + shallow_copies = [x.view_as(x) for x in ctx.input_tensors] + output_tensors = ctx.run_function(*shallow_copies) + input_grads = torch.autograd.grad( + output_tensors, + ctx.input_tensors + ctx.input_params, + output_grads, + allow_unused=True, + ) + del ctx.input_tensors + del ctx.input_params + del output_tensors + return (None, None) + input_grads + + +def checkpoint(func, inputs, params, flag): + """ + Evaluate a function without caching intermediate activations, allowing for + reduced memory at the expense of extra compute in the backward pass. + :param func: the function to evaluate. + :param inputs: the argument sequence to pass to `func`. + :param params: a sequence of parameters `func` depends on but does not + explicitly take as arguments. + :param flag: if False, disable gradient checkpointing. + """ + if flag: + args = tuple(inputs) + tuple(params) + return CheckpointFunction.apply(func, len(inputs), *args) + else: + return func(*inputs) + + +def zero_module(module): + """ + Zero out the parameters of a module and return it. + """ + for p in module.parameters(): + p.detach().zero_() + return module + + +def timestep_embedding(timesteps, dim, max_period=10000, repeat_only=False): + """ + Create sinusoidal timestep embeddings. + :param timesteps: a 1-D Tensor of N indices, one per batch element. + These may be fractional. + :param dim: the dimension of the output. + :param max_period: controls the minimum frequency of the embeddings. + :return: an [N x dim] Tensor of positional embeddings. + """ + if not repeat_only: + half = dim // 2 + freqs = torch.exp( + -math.log(max_period) + * torch.arange(start=0, end=half, dtype=torch.float32) + / half + ).to(device=timesteps.device) + args = timesteps[:, None].float() * freqs[None] + embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) + if dim % 2: + embedding = torch.cat( + [embedding, torch.zeros_like(embedding[:, :1])], dim=-1 + ) + else: + embedding = repeat(timesteps, "b -> b d", d=dim) + return embedding + + +class GroupNorm32(nn.GroupNorm): + def forward(self, x): + return super().forward(x.float()).type(x.dtype) + + +def normalization(channels): + """ + Make a standard normalization layer. + :param channels: number of input channels. + :return: an nn.Module for normalization. + """ + return GroupNorm32(32, channels) + + +def count_flops_attn(model, _x, y): + """ + A counter for the `thop` package to count the operations in an + attention operation. + Meant to be used like: + macs, params = thop.profile( + model, + inputs=(inputs, timestamps), + custom_ops={QKVAttention: QKVAttention.count_flops}, + ) + """ + b, c, *spatial = y[0].shape + num_spatial = int(np.prod(spatial)) + # We perform two matmuls with the same number of ops. + # The first computes the weight matrix, the second computes + # the combination of the value vectors. + matmul_ops = 2 * b * (num_spatial**2) * c + model.total_ops += torch.DoubleTensor([matmul_ops]) + + +def conv_nd(dims, *args, **kwargs): + """ + Create a 1D, 2D, or 3D convolution module. + """ + if dims == 1: + return nn.Conv1d(*args, **kwargs) + elif dims == 2: + return nn.Conv2d(*args, **kwargs) + elif dims == 3: + return nn.Conv3d(*args, **kwargs) + raise ValueError(f"unsupported dimensions: {dims}") + + +def avg_pool_nd(dims, *args, **kwargs): + """ + Create a 1D, 2D, or 3D average pooling module. + """ + if dims == 1: + return nn.AvgPool1d(*args, **kwargs) + elif dims == 2: + return nn.AvgPool2d(*args, **kwargs) + elif dims == 3: + return nn.AvgPool3d(*args, **kwargs) + raise ValueError(f"unsupported dimensions: {dims}") + + +class QKVAttention(nn.Module): + """ + A module which performs QKV attention and splits in a different order. + """ + + def __init__(self, n_heads): + super().__init__() + self.n_heads = n_heads + + def forward(self, qkv): + """ + Apply QKV attention. + :param qkv: an [N x (3 * H * C) x T] tensor of Qs, Ks, and Vs. + :return: an [N x (H * C) x T] tensor after attention. + """ + + bs, width, length = qkv.shape + assert width % (3 * self.n_heads) == 0 + ch = width // (3 * self.n_heads) + q, k, v = qkv.chunk(3, dim=1) # [N x (H * C) x T] + scale = 1 / math.sqrt(math.sqrt(ch)) + weight = torch.einsum( + "bct,bcs->bts", + (q * scale).view(bs * self.n_heads, ch, length), + (k * scale).view(bs * self.n_heads, ch, length), + ) # More stable with f16 than dividing afterwards + weight = torch.softmax(weight.float(), dim=-1).type(weight.dtype) + a = torch.einsum( + "bts,bcs->bct", weight, v.reshape(bs * self.n_heads, ch, length) + ) + return a.reshape(bs, -1, length) + + @staticmethod + def count_flops(model, _x, y): + return count_flops_attn(model, _x, y) + + +class QKVAttentionLegacy(nn.Module): + """ + A module which performs QKV attention. Matches legacy QKVAttention + input/ouput heads shaping + """ + + def __init__(self, n_heads): + super().__init__() + self.n_heads = n_heads + + def forward(self, qkv): + """ + Apply QKV attention. + :param qkv: an [N x (H * 3 * C) x T] tensor of Qs, Ks, and Vs. + :return: an [N x (H * C) x T] tensor after attention. + """ + bs, width, length = qkv.shape + assert width % (3 * self.n_heads) == 0 + ch = width // (3 * self.n_heads) + q, k, v = qkv.reshape(bs * self.n_heads, ch * 3, length).split(ch, dim=1) + scale = 1 / math.sqrt(math.sqrt(ch)) + weight = torch.einsum( + "bct,bcs->bts", q * scale, k * scale + ) # More stable with f16 than dividing afterwards + weight = torch.softmax(weight.float(), dim=-1).type(weight.dtype) + a = torch.einsum("bts,bcs->bct", weight, v) + return a.reshape(bs, -1, length) + + @staticmethod + def count_flops(model, _x, y): + return count_flops_attn(model, _x, y) + + +class AttentionPool2d(nn.Module): + """ + Adapted from CLIP: https://github.com/openai/CLIP/blob/main/clip/model.py + """ + + def __init__( + self, + spacial_dim: int, + embed_dim: int, + num_heads_channels: int, + output_dim: int = None, + ): + super().__init__() + self.positional_embedding = nn.Parameter( + torch.randn(embed_dim, spacial_dim**2 + 1) / embed_dim**0.5 + ) + self.qkv_proj = conv_nd(1, embed_dim, 3 * embed_dim, 1) + self.c_proj = conv_nd(1, embed_dim, output_dim or embed_dim, 1) + self.num_heads = embed_dim // num_heads_channels + self.attention = QKVAttention(self.num_heads) + + def forward(self, x): + b, c, *_spatial = x.shape + x = x.reshape(b, c, -1) # NC(HW) + x = torch.cat([x.mean(dim=-1, keepdim=True), x], dim=-1) # NC(HW+1) + x = x + self.positional_embedding[None, :, :].to(x.dtype) # NC(HW+1) + x = self.qkv_proj(x) + x = self.attention(x) + x = self.c_proj(x) + return x[:, :, 0] + + +class TimestepBlock(nn.Module): + """ + Any module where forward() takes timestep embeddings as a second argument. + """ + + @abstractmethod + def forward(self, x, emb): + """ + Apply the module to `x` given `emb` timestep embeddings. + """ + + +class TimestepEmbedSequential(nn.Sequential, TimestepBlock): + """ + A sequential module that passes timestep embeddings to the children that + support it as an extra input. + """ + + def forward(self, x, emb, context=None): + for layer in self: + if isinstance(layer, TimestepBlock): + x = layer(x, emb) + elif isinstance(layer, SpatialTransformer): + x = layer(x, context) + else: + x = layer(x) + return x + + +class Upsample(nn.Module): + """ + An upsampling layer with an optional convolution. + :param channels: channels in the inputs and outputs. + :param use_conv: a bool determining if a convolution is applied. + :param dims: determines if the signal is 1D, 2D, or 3D. If 3D, then + upsampling occurs in the inner-two dimensions. + """ + + def __init__(self, channels, use_conv, dims=2, out_channels=None, padding=1): + super().__init__() + self.channels = channels + self.out_channels = out_channels or channels + self.use_conv = use_conv + self.dims = dims + if use_conv: + self.conv = conv_nd( + dims, self.channels, self.out_channels, 3, padding=padding + ) + + def forward(self, x): + assert x.shape[1] == self.channels + if self.dims == 3: + x = F.interpolate( + x, (x.shape[2], x.shape[3] * 2, x.shape[4] * 2), mode="nearest" + ) + else: + x = F.interpolate(x, scale_factor=2, mode="nearest") + if self.use_conv: + x = self.conv(x) + return x + + +class TransposedUpsample(nn.Module): + "Learned 2x upsampling without padding" + + def __init__(self, channels, out_channels=None, ks=5): + super().__init__() + self.channels = channels + self.out_channels = out_channels or channels + + self.up = nn.ConvTranspose2d( + self.channels, self.out_channels, kernel_size=ks, stride=2 + ) + + def forward(self, x): + return self.up(x) + + +class Downsample(nn.Module): + """ + A downsampling layer with an optional convolution. + :param channels: channels in the inputs and outputs. + :param use_conv: a bool determining if a convolution is applied. + :param dims: determines if the signal is 1D, 2D, or 3D. If 3D, then + downsampling occurs in the inner-two dimensions. + """ + + def __init__(self, channels, use_conv, dims=2, out_channels=None, padding=1): + super().__init__() + self.channels = channels + self.out_channels = out_channels or channels + self.use_conv = use_conv + self.dims = dims + stride = 2 if dims != 3 else (1, 2, 2) + if use_conv: + self.op = conv_nd( + dims, + self.channels, + self.out_channels, + 3, + stride=stride, + padding=padding, + ) + else: + assert self.channels == self.out_channels + self.op = avg_pool_nd(dims, kernel_size=stride, stride=stride) + + def forward(self, x): + assert x.shape[1] == self.channels + return self.op(x) + + +class ResBlock(TimestepBlock): + """ + A residual block that can optionally change the number of channels. + :param channels: the number of input channels. + :param emb_channels: the number of timestep embedding channels. + :param dropout: the rate of dropout. + :param out_channels: if specified, the number of out channels. + :param use_conv: if True and out_channels is specified, use a spatial + convolution instead of a smaller 1x1 convolution to change the + channels in the skip connection. + :param dims: determines if the signal is 1D, 2D, or 3D. + :param use_checkpoint: if True, use gradient checkpointing on this module. + :param up: if True, use this block for upsampling. + :param down: if True, use this block for downsampling. + """ + + def __init__( + self, + channels, + emb_channels, + dropout, + out_channels=None, + use_conv=False, + use_scale_shift_norm=False, + dims=2, + use_checkpoint=False, + up=False, + down=False, + ): + super().__init__() + self.channels = channels + self.emb_channels = emb_channels + self.dropout = dropout + self.out_channels = out_channels or channels + self.use_conv = use_conv + self.use_checkpoint = use_checkpoint + self.use_scale_shift_norm = use_scale_shift_norm + + self.in_layers = nn.Sequential( + normalization(channels), + nn.SiLU(), + conv_nd(dims, channels, self.out_channels, 3, padding=1), + ) + + self.updown = up or down + + if up: + self.h_upd = Upsample(channels, False, dims) + self.x_upd = Upsample(channels, False, dims) + elif down: + self.h_upd = Downsample(channels, False, dims) + self.x_upd = Downsample(channels, False, dims) + else: + self.h_upd = self.x_upd = nn.Identity() + + self.emb_layers = nn.Sequential( + nn.SiLU(), + nn.Linear( + emb_channels, + 2 * self.out_channels if use_scale_shift_norm else self.out_channels, + ), + ) + self.out_layers = nn.Sequential( + normalization(self.out_channels), + nn.SiLU(), + nn.Dropout(p=dropout), + zero_module( + conv_nd(dims, self.out_channels, self.out_channels, 3, padding=1) + ), + ) + + if self.out_channels == channels: + self.skip_connection = nn.Identity() + elif use_conv: + self.skip_connection = conv_nd( + dims, channels, self.out_channels, 3, padding=1 + ) + else: + self.skip_connection = conv_nd(dims, channels, self.out_channels, 1) + + def forward(self, x, emb): + """ + Apply the block to a Tensor, conditioned on a timestep embedding. + :param x: an [N x C x ...] Tensor of features. + :param emb: an [N x emb_channels] Tensor of timestep embeddings. + :return: an [N x C x ...] Tensor of outputs. + """ + return checkpoint( + self._forward, (x, emb), self.parameters(), self.use_checkpoint + ) + + def _forward(self, x, emb): + if self.updown: + in_rest, in_conv = self.in_layers[:-1], self.in_layers[-1] + h = in_rest(x) + h = self.h_upd(h) + x = self.x_upd(x) + h = in_conv(h) + else: + h = self.in_layers(x) + emb_out = self.emb_layers(emb).type(h.dtype) + while len(emb_out.shape) < len(h.shape): + emb_out = emb_out[..., None] + if self.use_scale_shift_norm: + out_norm, out_rest = self.out_layers[0], self.out_layers[1:] + scale, shift = torch.chunk(emb_out, 2, dim=1) + h = out_norm(h) * (1 + scale) + shift + h = out_rest(h) + else: + h = h + emb_out + h = self.out_layers(h) + return self.skip_connection(x) + h + + +class AttentionBlock(nn.Module): + """ + An attention block that allows spatial positions to attend to each other. + Originally ported from here, but adapted to the N-d case. + https://github.com/hojonathanho/diffusion/blob/1e0dceb3b3495bbe19116a5e1b3596cd0706c543/diffusion_tf/models/unet.py#L66. + """ + + def __init__( + self, + channels, + num_heads=1, + num_head_channels=-1, + use_checkpoint=False, + use_new_attention_order=False, + ): + super().__init__() + self.channels = channels + if num_head_channels == -1: + self.num_heads = num_heads + else: + assert ( + channels % num_head_channels == 0 + ), f"q,k,v channels {channels} is not divisible by num_head_channels {num_head_channels}" + self.num_heads = channels // num_head_channels + self.use_checkpoint = use_checkpoint + self.norm = normalization(channels) + self.qkv = conv_nd(1, channels, channels * 3, 1) + if use_new_attention_order: + # split qkv before split heads + self.attention = QKVAttention(self.num_heads) + else: + # split heads before split qkv + self.attention = QKVAttentionLegacy(self.num_heads) + + self.proj_out = zero_module(conv_nd(1, channels, channels, 1)) + + def forward(self, x): + return checkpoint( + self._forward, (x,), self.parameters(), True + ) # TODO: check checkpoint usage, is True # TODO: fix the .half call!!! + # return pt_checkpoint(self._forward, x) # pytorch + + def _forward(self, x): + b, c, *spatial = x.shape + x = x.reshape(b, c, -1) + qkv = self.qkv(self.norm(x)) + h = self.attention(qkv) + h = self.proj_out(h) + return (x + h).reshape(b, c, *spatial) + + +class UNetModel(nn.Module): + """ + The full UNet model with attention and timestep embedding. + :param in_channels: channels in the input Tensor. + :param model_channels: base channel count for the model. + :param out_channels: channels in the output Tensor. + :param num_res_blocks: number of residual blocks per downsample. + :param attention_resolutions: a collection of downsample rates at which + attention will take place. May be a set, list, or tuple. + For example, if this contains 4, then at 4x downsampling, attention + will be used. + :param dropout: the dropout probability. + :param channel_mult: channel multiplier for each level of the UNet. + :param conv_resample: if True, use learned convolutions for upsampling and + downsampling. + :param dims: determines if the signal is 1D, 2D, or 3D. + :param num_classes: if specified (as an int), then this model will be + class-conditional with `num_classes` classes. + :param use_checkpoint: use gradient checkpointing to reduce memory usage. + :param num_heads: the number of attention heads in each attention layer. + :param num_heads_channels: if specified, ignore num_heads and instead use + a fixed channel width per attention head. + :param num_heads_upsample: works with num_heads to set a different number + of heads for upsampling. Deprecated. + :param use_scale_shift_norm: use a FiLM-like conditioning mechanism. + :param resblock_updown: use residual blocks for up/downsampling. + :param use_new_attention_order: use a different attention pattern for potentially + increased efficiency. + """ + + def __init__( + self, + image_size, + in_channels, + model_channels, + out_channels, + num_res_blocks, + attention_resolutions, + dropout=0, + channel_mult=(1, 2, 4, 8), + conv_resample=True, + dims=2, + num_classes=None, + use_checkpoint=False, + use_fp16=False, + num_heads=-1, + num_head_channels=-1, + num_heads_upsample=-1, + use_scale_shift_norm=False, + resblock_updown=False, + use_new_attention_order=False, + use_spatial_transformer=False, # custom transformer support + transformer_depth=1, # custom transformer support + context_dim=None, # custom transformer support + n_embed=None, # custom support for prediction of discrete ids into codebook of first stage vq model + legacy=True, + ): + super().__init__() + if use_spatial_transformer: + assert ( + context_dim is not None + ), "Fool!! You forgot to include the dimension of your cross-attention conditioning..." + + if context_dim is not None: + assert ( + use_spatial_transformer + ), "Fool!! You forgot to use the spatial transformer for your cross-attention conditioning..." + from omegaconf.listconfig import ListConfig + + if type(context_dim) == ListConfig: + context_dim = list(context_dim) + + if num_heads_upsample == -1: + num_heads_upsample = num_heads + + if num_heads == -1: + assert ( + num_head_channels != -1 + ), "Either num_heads or num_head_channels has to be set" + + if num_head_channels == -1: + assert ( + num_heads != -1 + ), "Either num_heads or num_head_channels has to be set" + + self.image_size = image_size + self.in_channels = in_channels + self.model_channels = model_channels + self.out_channels = out_channels + self.num_res_blocks = num_res_blocks + self.attention_resolutions = attention_resolutions + self.dropout = dropout + self.channel_mult = channel_mult + self.conv_resample = conv_resample + self.num_classes = num_classes + self.use_checkpoint = use_checkpoint + self.dtype = torch.float16 if use_fp16 else torch.float32 + self.num_heads = num_heads + self.num_head_channels = num_head_channels + self.num_heads_upsample = num_heads_upsample + self.predict_codebook_ids = n_embed is not None + + time_embed_dim = model_channels * 4 + self.time_embed = nn.Sequential( + nn.Linear(model_channels, time_embed_dim), + nn.SiLU(), + nn.Linear(time_embed_dim, time_embed_dim), + ) + + if self.num_classes is not None: + self.label_emb = nn.Embedding(num_classes, time_embed_dim) + + self.input_blocks = nn.ModuleList( + [ + TimestepEmbedSequential( + conv_nd(dims, in_channels, model_channels, 3, padding=1) + ) + ] + ) + self._feature_size = model_channels + input_block_chans = [model_channels] + ch = model_channels + ds = 1 + for level, mult in enumerate(channel_mult): + for _ in range(num_res_blocks): + layers = [ + ResBlock( + ch, + time_embed_dim, + dropout, + out_channels=mult * model_channels, + dims=dims, + use_checkpoint=use_checkpoint, + use_scale_shift_norm=use_scale_shift_norm, + ) + ] + ch = mult * model_channels + if ds in attention_resolutions: + if num_head_channels == -1: + dim_head = ch // num_heads + else: + num_heads = ch // num_head_channels + dim_head = num_head_channels + if legacy: + # num_heads = 1 + dim_head = ( + ch // num_heads + if use_spatial_transformer + else num_head_channels + ) + layers.append( + AttentionBlock( + ch, + use_checkpoint=use_checkpoint, + num_heads=num_heads, + num_head_channels=dim_head, + use_new_attention_order=use_new_attention_order, + ) + if not use_spatial_transformer + else SpatialTransformer( + ch, + num_heads, + dim_head, + depth=transformer_depth, + context_dim=context_dim, + ) + ) + self.input_blocks.append(TimestepEmbedSequential(*layers)) + self._feature_size += ch + input_block_chans.append(ch) + if level != len(channel_mult) - 1: + out_ch = ch + self.input_blocks.append( + TimestepEmbedSequential( + ResBlock( + ch, + time_embed_dim, + dropout, + out_channels=out_ch, + dims=dims, + use_checkpoint=use_checkpoint, + use_scale_shift_norm=use_scale_shift_norm, + down=True, + ) + if resblock_updown + else Downsample( + ch, conv_resample, dims=dims, out_channels=out_ch + ) + ) + ) + ch = out_ch + input_block_chans.append(ch) + ds *= 2 + self._feature_size += ch + + if num_head_channels == -1: + dim_head = ch // num_heads + else: + num_heads = ch // num_head_channels + dim_head = num_head_channels + if legacy: + # num_heads = 1 + dim_head = ch // num_heads if use_spatial_transformer else num_head_channels + self.middle_block = TimestepEmbedSequential( + ResBlock( + ch, + time_embed_dim, + dropout, + dims=dims, + use_checkpoint=use_checkpoint, + use_scale_shift_norm=use_scale_shift_norm, + ), + ( + AttentionBlock( + ch, + use_checkpoint=use_checkpoint, + num_heads=num_heads, + num_head_channels=dim_head, + use_new_attention_order=use_new_attention_order, + ) + if not use_spatial_transformer + else SpatialTransformer( + ch, + num_heads, + dim_head, + depth=transformer_depth, + context_dim=context_dim, + ) + ), + ResBlock( + ch, + time_embed_dim, + dropout, + dims=dims, + use_checkpoint=use_checkpoint, + use_scale_shift_norm=use_scale_shift_norm, + ), + ) + self._feature_size += ch + + self.output_blocks = nn.ModuleList([]) + for level, mult in list(enumerate(channel_mult))[::-1]: + for i in range(num_res_blocks + 1): + ich = input_block_chans.pop() + layers = [ + ResBlock( + ch + ich, + time_embed_dim, + dropout, + out_channels=model_channels * mult, + dims=dims, + use_checkpoint=use_checkpoint, + use_scale_shift_norm=use_scale_shift_norm, + ) + ] + ch = model_channels * mult + if ds in attention_resolutions: + if num_head_channels == -1: + dim_head = ch // num_heads + else: + num_heads = ch // num_head_channels + dim_head = num_head_channels + if legacy: + # num_heads = 1 + dim_head = ( + ch // num_heads + if use_spatial_transformer + else num_head_channels + ) + layers.append( + AttentionBlock( + ch, + use_checkpoint=use_checkpoint, + num_heads=num_heads_upsample, + num_head_channels=dim_head, + use_new_attention_order=use_new_attention_order, + ) + if not use_spatial_transformer + else SpatialTransformer( + ch, + num_heads, + dim_head, + depth=transformer_depth, + context_dim=context_dim, + ) + ) + if level and i == num_res_blocks: + out_ch = ch + layers.append( + ResBlock( + ch, + time_embed_dim, + dropout, + out_channels=out_ch, + dims=dims, + use_checkpoint=use_checkpoint, + use_scale_shift_norm=use_scale_shift_norm, + up=True, + ) + if resblock_updown + else Upsample(ch, conv_resample, dims=dims, out_channels=out_ch) + ) + ds //= 2 + self.output_blocks.append(TimestepEmbedSequential(*layers)) + self._feature_size += ch + + self.out = nn.Sequential( + normalization(ch), + nn.SiLU(), + zero_module(conv_nd(dims, model_channels, out_channels, 3, padding=1)), + ) + if self.predict_codebook_ids: + self.id_predictor = nn.Sequential( + normalization(ch), + conv_nd(dims, model_channels, n_embed, 1), + # nn.LogSoftmax(dim=1) # change to cross_entropy and produce non-normalized logits + ) + + def forward(self, x, timesteps=None, context=None, y=None, **kwargs): + """ + Apply the model to an input batch. + :param x: an [N x C x ...] Tensor of inputs. + :param timesteps: a 1-D batch of timesteps. + :param context: conditioning plugged in via crossattn + :param y: an [N] Tensor of labels, if class-conditional. + :return: an [N x C x ...] Tensor of outputs. + """ + assert (y is not None) == ( + self.num_classes is not None + ), "must specify y if and only if the model is class-conditional" + hs = [] + t_emb = timestep_embedding(timesteps, self.model_channels, repeat_only=False) + emb = self.time_embed(t_emb) + + if self.num_classes is not None: + assert y.shape == (x.shape[0],) + emb = emb + self.label_emb(y) + + h = x.type(self.dtype) + for module in self.input_blocks: + h = module(h, emb, context) + hs.append(h) + h = self.middle_block(h, emb, context) + for module in self.output_blocks: + # print(h.shape, hs[-1].shape) + if h.shape != hs[-1].shape: + if h.shape[-1] > hs[-1].shape[-1]: + h = h[:, :, :, : hs[-1].shape[-1]] + if h.shape[-2] > hs[-1].shape[-2]: + h = h[:, :, : hs[-1].shape[-2], :] + h = torch.cat([h, hs.pop()], dim=1) + h = module(h, emb, context) + # print(h.shape) + h = h.type(x.dtype) + if self.predict_codebook_ids: + return self.id_predictor(h) + else: + return self.out(h) + + +class AudioLDM(nn.Module): + def __init__(self, cfg): + super().__init__() + self.cfg = cfg + self.unet = UNetModel( + image_size=cfg.image_size, + in_channels=cfg.in_channels, + out_channels=cfg.out_channels, + model_channels=cfg.model_channels, + attention_resolutions=cfg.attention_resolutions, + num_res_blocks=cfg.num_res_blocks, + channel_mult=cfg.channel_mult, + num_heads=cfg.num_heads, + use_spatial_transformer=cfg.use_spatial_transformer, + transformer_depth=cfg.transformer_depth, + context_dim=cfg.context_dim, + use_checkpoint=cfg.use_checkpoint, + legacy=cfg.legacy, + ) + + def forward(self, x, timesteps=None, context=None, y=None): + x = self.unet(x=x, timesteps=timesteps, context=context, y=y) + return x diff --git a/models/tta/ldm/audioldm_dataset.py b/models/tta/ldm/audioldm_dataset.py new file mode 100644 index 0000000..9623253 --- /dev/null +++ b/models/tta/ldm/audioldm_dataset.py @@ -0,0 +1,151 @@ +# Copyright (c) 2023 Amphion. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +import random +import torch +from torch.nn.utils.rnn import pad_sequence +from utils.data_utils import * + + +from models.base.base_dataset import ( + BaseOfflineCollator, + BaseOfflineDataset, + BaseTestDataset, + BaseTestCollator, +) +import librosa + +from transformers import AutoTokenizer + + +class AudioLDMDataset(BaseOfflineDataset): + def __init__(self, cfg, dataset, is_valid=False): + BaseOfflineDataset.__init__(self, cfg, dataset, is_valid=is_valid) + + self.cfg = cfg + + # utt2melspec + if cfg.preprocess.use_melspec: + self.utt2melspec_path = {} + for utt_info in self.metadata: + dataset = utt_info["Dataset"] + uid = utt_info["Uid"] + utt = "{}_{}".format(dataset, uid) + + self.utt2melspec_path[utt] = os.path.join( + cfg.preprocess.processed_dir, + dataset, + cfg.preprocess.melspec_dir, + uid + ".npy", + ) + + # utt2wav + if cfg.preprocess.use_wav: + self.utt2wav_path = {} + for utt_info in self.metadata: + dataset = utt_info["Dataset"] + uid = utt_info["Uid"] + utt = "{}_{}".format(dataset, uid) + + self.utt2wav_path[utt] = os.path.join( + cfg.preprocess.processed_dir, + dataset, + cfg.preprocess.wav_dir, + uid + ".wav", + ) + + # utt2caption + if cfg.preprocess.use_caption: + self.utt2caption = {} + for utt_info in self.metadata: + dataset = utt_info["Dataset"] + uid = utt_info["Uid"] + utt = "{}_{}".format(dataset, uid) + + self.utt2caption[utt] = utt_info["Caption"] + + def __getitem__(self, index): + # melspec: (n_mels, T) + # wav: (T,) + + single_feature = BaseOfflineDataset.__getitem__(self, index) + + utt_info = self.metadata[index] + dataset = utt_info["Dataset"] + uid = utt_info["Uid"] + utt = "{}_{}".format(dataset, uid) + + if self.cfg.preprocess.use_melspec: + single_feature["melspec"] = np.load(self.utt2melspec_path[utt]) + + if self.cfg.preprocess.use_wav: + wav, sr = librosa.load( + self.utt2wav_path[utt], sr=16000 + ) # hard coding for 16KHz... + single_feature["wav"] = wav + + if self.cfg.preprocess.use_caption: + cond_mask = np.random.choice( + [1, 0], + p=[ + self.cfg.preprocess.cond_mask_prob, + 1 - self.cfg.preprocess.cond_mask_prob, + ], + ) # (0.1, 0.9) + if cond_mask: + single_feature["caption"] = "" + else: + single_feature["caption"] = self.utt2caption[utt] + + return single_feature + + def __len__(self): + return len(self.metadata) + + +class AudioLDMCollator(BaseOfflineCollator): + def __init__(self, cfg): + BaseOfflineCollator.__init__(self, cfg) + + self.tokenizer = AutoTokenizer.from_pretrained("t5-base", model_max_length=512) + + def __call__(self, batch): + # mel: (B, n_mels, T) + # wav (option): (B, T) + # text_input_ids: (B, L) + # text_attention_mask: (B, L) + + packed_batch_features = dict() + + for key in batch[0].keys(): + if key == "melspec": + packed_batch_features["melspec"] = torch.from_numpy( + np.array([b["melspec"][:, :624] for b in batch]) + ) + + if key == "wav": + values = [torch.from_numpy(b[key]) for b in batch] + packed_batch_features[key] = pad_sequence( + values, batch_first=True, padding_value=0 + ) + + if key == "caption": + captions = [b[key] for b in batch] + text_input = self.tokenizer( + captions, return_tensors="pt", truncation=True, padding="longest" + ) + text_input_ids = text_input["input_ids"] + text_attention_mask = text_input["attention_mask"] + + packed_batch_features["text_input_ids"] = text_input_ids + packed_batch_features["text_attention_mask"] = text_attention_mask + + return packed_batch_features + + +class AudioLDMTestDataset(BaseTestDataset): ... + + +class AudioLDMTestCollator(BaseTestCollator): ... diff --git a/models/tta/ldm/audioldm_inference.py b/models/tta/ldm/audioldm_inference.py new file mode 100644 index 0000000..a37b406 --- /dev/null +++ b/models/tta/ldm/audioldm_inference.py @@ -0,0 +1,193 @@ +# Copyright (c) 2023 Amphion. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +import os +import time +import numpy as np +import torch +from tqdm import tqdm +import torch.nn as nn +from collections import OrderedDict +import json + +from models.tta.autoencoder.autoencoder import AutoencoderKL +from models.tta.ldm.inference_utils.vocoder import Generator +from models.tta.ldm.audioldm import AudioLDM +from transformers import T5EncoderModel, AutoTokenizer +from diffusers import PNDMScheduler + +import matplotlib.pyplot as plt +from scipy.io.wavfile import write + + +class AttrDict(dict): + def __init__(self, *args, **kwargs): + super(AttrDict, self).__init__(*args, **kwargs) + self.__dict__ = self + + +class AudioLDMInference: + def __init__(self, args, cfg): + self.cfg = cfg + self.args = args + + self.build_autoencoderkl() + self.build_textencoder() + + self.model = self.build_model() + self.load_state_dict() + + self.build_vocoder() + + self.out_path = self.args.output_dir + self.out_mel_path = os.path.join(self.out_path, "mel") + self.out_wav_path = os.path.join(self.out_path, "wav") + os.makedirs(self.out_mel_path, exist_ok=True) + os.makedirs(self.out_wav_path, exist_ok=True) + + def build_autoencoderkl(self): + self.autoencoderkl = AutoencoderKL(self.cfg.model.autoencoderkl) + self.autoencoder_path = self.cfg.model.autoencoder_path + checkpoint = torch.load(self.autoencoder_path, map_location="cpu") + self.autoencoderkl.load_state_dict(checkpoint["model"]) + self.autoencoderkl.cuda(self.args.local_rank) + self.autoencoderkl.requires_grad_(requires_grad=False) + self.autoencoderkl.eval() + + def build_textencoder(self): + self.tokenizer = AutoTokenizer.from_pretrained("t5-base", model_max_length=512) + self.text_encoder = T5EncoderModel.from_pretrained("t5-base") + self.text_encoder.cuda(self.args.local_rank) + self.text_encoder.requires_grad_(requires_grad=False) + self.text_encoder.eval() + + def build_vocoder(self): + config_file = os.path.join(self.args.vocoder_config_path) + with open(config_file) as f: + data = f.read() + json_config = json.loads(data) + h = AttrDict(json_config) + self.vocoder = Generator(h).to(self.args.local_rank) + checkpoint_dict = torch.load( + self.args.vocoder_path, map_location=self.args.local_rank + ) + self.vocoder.load_state_dict(checkpoint_dict["generator"]) + + def build_model(self): + self.model = AudioLDM(self.cfg.model.audioldm) + return self.model + + def load_state_dict(self): + self.checkpoint_path = self.args.checkpoint_path + checkpoint = torch.load(self.checkpoint_path, map_location="cpu") + self.model.load_state_dict(checkpoint["model"]) + self.model.cuda(self.args.local_rank) + + def get_text_embedding(self): + text = self.args.text + + prompt = [text] + + text_input = self.tokenizer( + prompt, + max_length=self.tokenizer.model_max_length, + truncation=True, + padding="do_not_pad", + return_tensors="pt", + ) + text_embeddings = self.text_encoder( + text_input.input_ids.to(self.args.local_rank) + )[0] + + max_length = text_input.input_ids.shape[-1] + uncond_input = self.tokenizer( + [""] * 1, padding="max_length", max_length=max_length, return_tensors="pt" + ) + uncond_embeddings = self.text_encoder( + uncond_input.input_ids.to(self.args.local_rank) + )[0] + text_embeddings = torch.cat([uncond_embeddings, text_embeddings]) + + return text_embeddings + + def inference(self): + text_embeddings = self.get_text_embedding() + print(text_embeddings.shape) + + num_steps = self.args.num_steps + guidance_scale = self.args.guidance_scale + + noise_scheduler = PNDMScheduler( + num_train_timesteps=1000, + beta_start=0.00085, + beta_end=0.012, + beta_schedule="scaled_linear", + skip_prk_steps=True, + set_alpha_to_one=False, + steps_offset=1, + prediction_type="epsilon", + ) + + noise_scheduler.set_timesteps(num_steps) + + latents = torch.randn( + ( + 1, + self.cfg.model.autoencoderkl.z_channels, + 80 // (2 ** (len(self.cfg.model.autoencoderkl.ch_mult) - 1)), + 624 // (2 ** (len(self.cfg.model.autoencoderkl.ch_mult) - 1)), + ) + ).to(self.args.local_rank) + + self.model.eval() + for t in tqdm(noise_scheduler.timesteps): + t = t.to(self.args.local_rank) + + # expand the latents if we are doing classifier-free guidance to avoid doing two forward passes. + latent_model_input = torch.cat([latents] * 2) + + latent_model_input = noise_scheduler.scale_model_input( + latent_model_input, timestep=t + ) + # print(latent_model_input.shape) + + # predict the noise residual + with torch.no_grad(): + noise_pred = self.model( + latent_model_input, torch.cat([t.unsqueeze(0)] * 2), text_embeddings + ) + + # perform guidance + noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) + noise_pred = noise_pred_uncond + guidance_scale * ( + noise_pred_text - noise_pred_uncond + ) + + # compute the previous noisy sample x_t -> x_t-1 + latents = noise_scheduler.step(noise_pred, t, latents).prev_sample + # print(latents.shape) + + latents_out = latents + print(latents_out.shape) + + with torch.no_grad(): + mel_out = self.autoencoderkl.decode(latents_out) + print(mel_out.shape) + + melspec = mel_out[0, 0].cpu().detach().numpy() + plt.imsave(os.path.join(self.out_mel_path, self.args.text + ".png"), melspec) + + self.vocoder.eval() + self.vocoder.remove_weight_norm() + with torch.no_grad(): + melspec = np.expand_dims(melspec, 0) + melspec = torch.FloatTensor(melspec).to(self.args.local_rank) + + y = self.vocoder(melspec) + audio = y.squeeze() + audio = audio * 32768.0 + audio = audio.cpu().numpy().astype("int16") + + write(os.path.join(self.out_wav_path, self.args.text + ".wav"), 16000, audio) diff --git a/models/tta/ldm/audioldm_trainer.py b/models/tta/ldm/audioldm_trainer.py new file mode 100644 index 0000000..bd4a241 --- /dev/null +++ b/models/tta/ldm/audioldm_trainer.py @@ -0,0 +1,251 @@ +# Copyright (c) 2023 Amphion. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +from models.base.base_trainer import BaseTrainer +from diffusers import DDPMScheduler +from models.tta.ldm.audioldm_dataset import AudioLDMDataset, AudioLDMCollator +from models.tta.autoencoder.autoencoder import AutoencoderKL +from models.tta.ldm.audioldm import AudioLDM, UNetModel +import torch +import torch.nn as nn +from torch.nn import MSELoss, L1Loss +import torch.nn.functional as F +from torch.utils.data import ConcatDataset, DataLoader + +from transformers import T5EncoderModel +from diffusers import DDPMScheduler + + +class AudioLDMTrainer(BaseTrainer): + def __init__(self, args, cfg): + BaseTrainer.__init__(self, args, cfg) + self.cfg = cfg + + self.build_autoencoderkl() + self.build_textencoder() + self.nosie_scheduler = self.build_noise_scheduler() + + self.save_config_file() + + def build_autoencoderkl(self): + self.autoencoderkl = AutoencoderKL(self.cfg.model.autoencoderkl) + self.autoencoder_path = self.cfg.model.autoencoder_path + checkpoint = torch.load(self.autoencoder_path, map_location="cpu") + self.autoencoderkl.load_state_dict(checkpoint["model"]) + self.autoencoderkl.cuda(self.args.local_rank) + self.autoencoderkl.requires_grad_(requires_grad=False) + self.autoencoderkl.eval() + + def build_textencoder(self): + self.text_encoder = T5EncoderModel.from_pretrained("t5-base") + self.text_encoder.cuda(self.args.local_rank) + self.text_encoder.requires_grad_(requires_grad=False) + self.text_encoder.eval() + + def build_noise_scheduler(self): + nosie_scheduler = DDPMScheduler( + num_train_timesteps=self.cfg.model.noise_scheduler.num_train_timesteps, + beta_start=self.cfg.model.noise_scheduler.beta_start, + beta_end=self.cfg.model.noise_scheduler.beta_end, + beta_schedule=self.cfg.model.noise_scheduler.beta_schedule, + clip_sample=self.cfg.model.noise_scheduler.clip_sample, + # steps_offset=self.cfg.model.noise_scheduler.steps_offset, + # set_alpha_to_one=self.cfg.model.noise_scheduler.set_alpha_to_one, + # skip_prk_steps=self.cfg.model.noise_scheduler.skip_prk_steps, + prediction_type=self.cfg.model.noise_scheduler.prediction_type, + ) + return nosie_scheduler + + def build_dataset(self): + return AudioLDMDataset, AudioLDMCollator + + def build_data_loader(self): + Dataset, Collator = self.build_dataset() + # build dataset instance for each dataset and combine them by ConcatDataset + datasets_list = [] + for dataset in self.cfg.dataset: + subdataset = Dataset(self.cfg, dataset, is_valid=False) + datasets_list.append(subdataset) + train_dataset = ConcatDataset(datasets_list) + + train_collate = Collator(self.cfg) + + # use batch_sampler argument instead of (sampler, shuffle, drop_last, batch_size) + train_loader = DataLoader( + train_dataset, + collate_fn=train_collate, + num_workers=self.args.num_workers, + batch_size=self.cfg.train.batch_size, + pin_memory=False, + ) + if not self.cfg.train.ddp or self.args.local_rank == 0: + datasets_list = [] + for dataset in self.cfg.dataset: + subdataset = Dataset(self.cfg, dataset, is_valid=True) + datasets_list.append(subdataset) + valid_dataset = ConcatDataset(datasets_list) + valid_collate = Collator(self.cfg) + + valid_loader = DataLoader( + valid_dataset, + collate_fn=valid_collate, + num_workers=1, + batch_size=self.cfg.train.batch_size, + ) + else: + raise NotImplementedError("DDP is not supported yet.") + # valid_loader = None + data_loader = {"train": train_loader, "valid": valid_loader} + return data_loader + + def build_optimizer(self): + optimizer = torch.optim.AdamW(self.model.parameters(), **self.cfg.train.adam) + return optimizer + + # TODO: check it... + def build_scheduler(self): + return None + # return ReduceLROnPlateau(self.optimizer["opt_ae"], **self.cfg.train.lronPlateau) + + def write_summary(self, losses, stats): + for key, value in losses.items(): + self.sw.add_scalar(key, value, self.step) + + def write_valid_summary(self, losses, stats): + for key, value in losses.items(): + self.sw.add_scalar(key, value, self.step) + + def build_criterion(self): + criterion = nn.MSELoss(reduction="mean") + return criterion + + def get_state_dict(self): + if self.scheduler != None: + state_dict = { + "model": self.model.state_dict(), + "optimizer": self.optimizer.state_dict(), + "scheduler": self.scheduler.state_dict(), + "step": self.step, + "epoch": self.epoch, + "batch_size": self.cfg.train.batch_size, + } + else: + state_dict = { + "model": self.model.state_dict(), + "optimizer": self.optimizer.state_dict(), + "step": self.step, + "epoch": self.epoch, + "batch_size": self.cfg.train.batch_size, + } + return state_dict + + def load_model(self, checkpoint): + self.step = checkpoint["step"] + self.epoch = checkpoint["epoch"] + + self.model.load_state_dict(checkpoint["model"]) + self.optimizer.load_state_dict(checkpoint["optimizer"]) + if self.scheduler != None: + self.scheduler.load_state_dict(checkpoint["scheduler"]) + + def build_model(self): + self.model = AudioLDM(self.cfg.model.audioldm) + return self.model + + @torch.no_grad() + def mel_to_latent(self, melspec): + posterior = self.autoencoderkl.encode(melspec) + latent = posterior.sample() # (B, 4, 5, 78) + return latent + + @torch.no_grad() + def get_text_embedding(self, text_input_ids, text_attention_mask): + text_embedding = self.text_encoder( + input_ids=text_input_ids, attention_mask=text_attention_mask + ).last_hidden_state + return text_embedding # (B, T, 768) + + def train_step(self, data): + train_losses = {} + total_loss = 0 + train_stats = {} + + melspec = data["melspec"].unsqueeze(1) # (B, 80, T) -> (B, 1, 80, T) + latents = self.mel_to_latent(melspec) + + text_embedding = self.get_text_embedding( + data["text_input_ids"], data["text_attention_mask"] + ) + + noise = torch.randn_like(latents).float() + + bsz = latents.shape[0] + timesteps = torch.randint( + 0, + self.cfg.model.noise_scheduler.num_train_timesteps, + (bsz,), + device=latents.device, + ) + timesteps = timesteps.long() + + with torch.no_grad(): + noisy_latents = self.nosie_scheduler.add_noise(latents, noise, timesteps) + + model_pred = self.model( + noisy_latents, timesteps=timesteps, context=text_embedding + ) + + loss = self.criterion(model_pred, noise) + + train_losses["loss"] = loss + total_loss += loss + + self.optimizer.zero_grad() + total_loss.backward() + self.optimizer.step() + + for item in train_losses: + train_losses[item] = train_losses[item].item() + + return train_losses, train_stats, total_loss.item() + + # TODO: eval step + @torch.no_grad() + def eval_step(self, data, index): + valid_loss = {} + total_valid_loss = 0 + valid_stats = {} + + melspec = data["melspec"].unsqueeze(1) # (B, 80, T) -> (B, 1, 80, T) + latents = self.mel_to_latent(melspec) + + text_embedding = self.get_text_embedding( + data["text_input_ids"], data["text_attention_mask"] + ) + + noise = torch.randn_like(latents).float() + + bsz = latents.shape[0] + timesteps = torch.randint( + 0, + self.cfg.model.noise_scheduler.num_train_timesteps, + (bsz,), + device=latents.device, + ) + timesteps = timesteps.long() + + noisy_latents = self.nosie_scheduler.add_noise(latents, noise, timesteps) + + model_pred = self.model(noisy_latents, timesteps, text_embedding) + + loss = self.criterion(model_pred, noise) + valid_loss["loss"] = loss + + total_valid_loss += loss + + for item in valid_loss: + valid_loss[item] = valid_loss[item].item() + + return valid_loss, valid_stats, total_valid_loss.item() diff --git a/models/tta/ldm/inference_utils/utils.py b/models/tta/ldm/inference_utils/utils.py new file mode 100644 index 0000000..fd61c0c --- /dev/null +++ b/models/tta/ldm/inference_utils/utils.py @@ -0,0 +1,62 @@ +# Copyright (c) 2023 Amphion. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +import glob +import os +import matplotlib +import torch +from torch.nn.utils import weight_norm + +matplotlib.use("Agg") +import matplotlib.pylab as plt + + +def plot_spectrogram(spectrogram): + fig, ax = plt.subplots(figsize=(10, 2)) + im = ax.imshow(spectrogram, aspect="auto", origin="lower", interpolation="none") + plt.colorbar(im, ax=ax) + + fig.canvas.draw() + plt.close() + + return fig + + +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) + + +def apply_weight_norm(m): + classname = m.__class__.__name__ + if classname.find("Conv") != -1: + weight_norm(m) + + +def get_padding(kernel_size, dilation=1): + return int((kernel_size * dilation - dilation) / 2) + + +def load_checkpoint(filepath, device): + assert os.path.isfile(filepath) + print("Loading '{}'".format(filepath)) + checkpoint_dict = torch.load(filepath, map_location=device) + print("Complete.") + return checkpoint_dict + + +def save_checkpoint(filepath, obj): + print("Saving checkpoint to {}".format(filepath)) + torch.save(obj, filepath) + print("Complete.") + + +def scan_checkpoint(cp_dir, prefix): + pattern = os.path.join(cp_dir, prefix + "????????") + cp_list = glob.glob(pattern) + if len(cp_list) == 0: + return None + return sorted(cp_list)[-1] diff --git a/models/tta/ldm/inference_utils/vocoder.py b/models/tta/ldm/inference_utils/vocoder.py new file mode 100644 index 0000000..19e17c1 --- /dev/null +++ b/models/tta/ldm/inference_utils/vocoder.py @@ -0,0 +1,408 @@ +# Copyright (c) 2023 Amphion. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +import torch +import torch.nn.functional as F +import torch.nn as nn +from torch.nn import Conv1d, ConvTranspose1d, AvgPool1d, Conv2d +from torch.nn.utils import weight_norm, remove_weight_norm, spectral_norm +from models.tta.ldm.inference_utils.utils import get_padding, init_weights + +LRELU_SLOPE = 0.1 + + +class ResBlock1(torch.nn.Module): + def __init__(self, h, channels, kernel_size=3, dilation=(1, 3, 5)): + super(ResBlock1, self).__init__() + self.h = h + self.convs1 = nn.ModuleList( + [ + weight_norm( + Conv1d( + channels, + channels, + kernel_size, + 1, + dilation=dilation[0], + padding=get_padding(kernel_size, dilation[0]), + ) + ), + weight_norm( + Conv1d( + channels, + channels, + kernel_size, + 1, + dilation=dilation[1], + padding=get_padding(kernel_size, dilation[1]), + ) + ), + weight_norm( + Conv1d( + channels, + channels, + kernel_size, + 1, + dilation=dilation[2], + padding=get_padding(kernel_size, dilation[2]), + ) + ), + ] + ) + self.convs1.apply(init_weights) + + self.convs2 = nn.ModuleList( + [ + weight_norm( + Conv1d( + channels, + channels, + kernel_size, + 1, + dilation=1, + padding=get_padding(kernel_size, 1), + ) + ), + weight_norm( + Conv1d( + channels, + channels, + kernel_size, + 1, + dilation=1, + padding=get_padding(kernel_size, 1), + ) + ), + weight_norm( + Conv1d( + channels, + channels, + kernel_size, + 1, + dilation=1, + padding=get_padding(kernel_size, 1), + ) + ), + ] + ) + self.convs2.apply(init_weights) + + def forward(self, x): + for c1, c2 in zip(self.convs1, self.convs2): + xt = F.leaky_relu(x, LRELU_SLOPE) + xt = c1(xt) + xt = F.leaky_relu(xt, LRELU_SLOPE) + xt = c2(xt) + x = xt + x + return x + + def remove_weight_norm(self): + for l in self.convs1: + remove_weight_norm(l) + for l in self.convs2: + remove_weight_norm(l) + + +class ResBlock2(torch.nn.Module): + def __init__(self, h, channels, kernel_size=3, dilation=(1, 3)): + super(ResBlock2, self).__init__() + self.h = h + self.convs = nn.ModuleList( + [ + weight_norm( + Conv1d( + channels, + channels, + kernel_size, + 1, + dilation=dilation[0], + padding=get_padding(kernel_size, dilation[0]), + ) + ), + weight_norm( + Conv1d( + channels, + channels, + kernel_size, + 1, + dilation=dilation[1], + padding=get_padding(kernel_size, dilation[1]), + ) + ), + ] + ) + self.convs.apply(init_weights) + + def forward(self, x): + for c in self.convs: + xt = F.leaky_relu(x, LRELU_SLOPE) + xt = c(xt) + x = xt + x + return x + + def remove_weight_norm(self): + for l in self.convs: + remove_weight_norm(l) + + +class Generator(torch.nn.Module): + def __init__(self, h): + super(Generator, self).__init__() + self.h = h + self.num_kernels = len(h.resblock_kernel_sizes) + self.num_upsamples = len(h.upsample_rates) + self.conv_pre = weight_norm( + Conv1d(80, h.upsample_initial_channel, 7, 1, padding=3) + ) + resblock = ResBlock1 if h.resblock == "1" else ResBlock2 + + self.ups = nn.ModuleList() + for i, (u, k) in enumerate(zip(h.upsample_rates, h.upsample_kernel_sizes)): + self.ups.append( + weight_norm( + ConvTranspose1d( + h.upsample_initial_channel // (2**i), + h.upsample_initial_channel // (2 ** (i + 1)), + k, + u, + padding=(k - u) // 2, + ) + ) + ) + + self.resblocks = nn.ModuleList() + for i in range(len(self.ups)): + ch = h.upsample_initial_channel // (2 ** (i + 1)) + for j, (k, d) in enumerate( + zip(h.resblock_kernel_sizes, h.resblock_dilation_sizes) + ): + self.resblocks.append(resblock(h, ch, k, d)) + + self.conv_post = weight_norm(Conv1d(ch, 1, 7, 1, padding=3)) + self.ups.apply(init_weights) + self.conv_post.apply(init_weights) + + def forward(self, x): + x = self.conv_pre(x) + for i in range(self.num_upsamples): + x = F.leaky_relu(x, LRELU_SLOPE) + x = self.ups[i](x) + 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) + x = torch.tanh(x) + + 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) + + +class DiscriminatorP(torch.nn.Module): + def __init__(self, period, kernel_size=5, stride=3, use_spectral_norm=False): + super(DiscriminatorP, self).__init__() + self.period = period + norm_f = weight_norm if use_spectral_norm == False else spectral_norm + self.convs = nn.ModuleList( + [ + norm_f( + Conv2d( + 1, + 32, + (kernel_size, 1), + (stride, 1), + padding=(get_padding(5, 1), 0), + ) + ), + norm_f( + Conv2d( + 32, + 128, + (kernel_size, 1), + (stride, 1), + padding=(get_padding(5, 1), 0), + ) + ), + norm_f( + Conv2d( + 128, + 512, + (kernel_size, 1), + (stride, 1), + padding=(get_padding(5, 1), 0), + ) + ), + norm_f( + Conv2d( + 512, + 1024, + (kernel_size, 1), + (stride, 1), + padding=(get_padding(5, 1), 0), + ) + ), + norm_f(Conv2d(1024, 1024, (kernel_size, 1), 1, padding=(2, 0))), + ] + ) + self.conv_post = norm_f(Conv2d(1024, 1, (3, 1), 1, padding=(1, 0))) + + def forward(self, x): + fmap = [] + + # 1d to 2d + b, c, t = x.shape + if t % self.period != 0: # pad first + n_pad = self.period - (t % self.period) + x = F.pad(x, (0, n_pad), "reflect") + t = t + n_pad + x = x.view(b, c, t // self.period, self.period) + + for l in self.convs: + x = l(x) + x = F.leaky_relu(x, LRELU_SLOPE) + fmap.append(x) + x = self.conv_post(x) + fmap.append(x) + x = torch.flatten(x, 1, -1) + + return x, fmap + + +class MultiPeriodDiscriminator(torch.nn.Module): + def __init__(self): + super(MultiPeriodDiscriminator, self).__init__() + self.discriminators = nn.ModuleList( + [ + DiscriminatorP(2), + DiscriminatorP(3), + DiscriminatorP(5), + DiscriminatorP(7), + DiscriminatorP(11), + ] + ) + + def forward(self, y, y_hat): + y_d_rs = [] + y_d_gs = [] + fmap_rs = [] + fmap_gs = [] + for i, d in enumerate(self.discriminators): + y_d_r, fmap_r = d(y) + y_d_g, fmap_g = d(y_hat) + y_d_rs.append(y_d_r) + fmap_rs.append(fmap_r) + y_d_gs.append(y_d_g) + fmap_gs.append(fmap_g) + + return y_d_rs, y_d_gs, fmap_rs, fmap_gs + + +class DiscriminatorS(torch.nn.Module): + def __init__(self, use_spectral_norm=False): + super(DiscriminatorS, self).__init__() + norm_f = weight_norm if use_spectral_norm == False else spectral_norm + self.convs = nn.ModuleList( + [ + norm_f(Conv1d(1, 128, 15, 1, padding=7)), + norm_f(Conv1d(128, 128, 41, 2, groups=4, padding=20)), + norm_f(Conv1d(128, 256, 41, 2, groups=16, padding=20)), + norm_f(Conv1d(256, 512, 41, 4, groups=16, padding=20)), + norm_f(Conv1d(512, 1024, 41, 4, groups=16, padding=20)), + norm_f(Conv1d(1024, 1024, 41, 1, groups=16, padding=20)), + norm_f(Conv1d(1024, 1024, 5, 1, padding=2)), + ] + ) + self.conv_post = norm_f(Conv1d(1024, 1, 3, 1, padding=1)) + + def forward(self, x): + fmap = [] + for l in self.convs: + x = l(x) + x = F.leaky_relu(x, LRELU_SLOPE) + fmap.append(x) + x = self.conv_post(x) + fmap.append(x) + x = torch.flatten(x, 1, -1) + + return x, fmap + + +class MultiScaleDiscriminator(torch.nn.Module): + def __init__(self): + super(MultiScaleDiscriminator, self).__init__() + self.discriminators = nn.ModuleList( + [ + DiscriminatorS(use_spectral_norm=True), + DiscriminatorS(), + DiscriminatorS(), + ] + ) + self.meanpools = nn.ModuleList( + [AvgPool1d(4, 2, padding=2), AvgPool1d(4, 2, padding=2)] + ) + + def forward(self, y, y_hat): + y_d_rs = [] + y_d_gs = [] + fmap_rs = [] + fmap_gs = [] + for i, d in enumerate(self.discriminators): + if i != 0: + y = self.meanpools[i - 1](y) + y_hat = self.meanpools[i - 1](y_hat) + y_d_r, fmap_r = d(y) + y_d_g, fmap_g = d(y_hat) + y_d_rs.append(y_d_r) + fmap_rs.append(fmap_r) + y_d_gs.append(y_d_g) + fmap_gs.append(fmap_g) + + return y_d_rs, y_d_gs, fmap_rs, fmap_gs + + +def feature_loss(fmap_r, fmap_g): + loss = 0 + for dr, dg in zip(fmap_r, fmap_g): + for rl, gl in zip(dr, dg): + loss += torch.mean(torch.abs(rl - gl)) + + return loss * 2 + + +def discriminator_loss(disc_real_outputs, disc_generated_outputs): + loss = 0 + r_losses = [] + g_losses = [] + for dr, dg in zip(disc_real_outputs, disc_generated_outputs): + r_loss = torch.mean((1 - dr) ** 2) + g_loss = torch.mean(dg**2) + loss += r_loss + g_loss + r_losses.append(r_loss.item()) + g_losses.append(g_loss.item()) + + return loss, r_losses, g_losses + + +def generator_loss(disc_outputs): + loss = 0 + gen_losses = [] + for dg in disc_outputs: + l = torch.mean((1 - dg) ** 2) + gen_losses.append(l) + loss += l + + return loss, gen_losses diff --git a/models/tts/base/__init__.py b/models/tts/base/__init__.py new file mode 100644 index 0000000..268a381 --- /dev/null +++ b/models/tts/base/__init__.py @@ -0,0 +1,7 @@ +# Copyright (c) 2023 Amphion. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +# from .tts_inferece import TTSInference +from .tts_trainer import TTSTrainer diff --git a/models/tts/base/tts_dataset.py b/models/tts/base/tts_dataset.py new file mode 100644 index 0000000..dca129e --- /dev/null +++ b/models/tts/base/tts_dataset.py @@ -0,0 +1,392 @@ +# Copyright (c) 2023 Amphion. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +import json +import os +import torchaudio +import numpy as np +import torch +from utils.data_utils import * +from torch.nn.utils.rnn import pad_sequence +from text import text_to_sequence +from text.text_token_collation import phoneIDCollation +from processors.acoustic_extractor import cal_normalized_mel + +from models.base.base_dataset import ( + BaseOfflineDataset, + BaseOfflineCollator, + BaseTestDataset, + BaseTestCollator, +) + +from processors.content_extractor import ( + ContentvecExtractor, + WenetExtractor, + WhisperExtractor, +) + + +class TTSDataset(BaseOfflineDataset): + def __init__(self, cfg, dataset, is_valid=False): + """ + Args: + cfg: config + dataset: dataset name + is_valid: whether to use train or valid dataset + """ + + assert isinstance(dataset, str) + + self.cfg = cfg + + processed_data_dir = os.path.join(cfg.preprocess.processed_dir, dataset) + meta_file = cfg.preprocess.valid_file if is_valid else cfg.preprocess.train_file + self.metafile_path = os.path.join(processed_data_dir, meta_file) + self.metadata = self.get_metadata() + + """ + load spk2id and utt2spk from json file + spk2id: {spk1: 0, spk2: 1, ...} + utt2spk: {dataset_uid: spk1, ...} + """ + if cfg.preprocess.use_spkid: + dataset = self.metadata[0]["Dataset"] + + spk2id_path = os.path.join(processed_data_dir, cfg.preprocess.spk2id) + with open(spk2id_path, "r") as f: + self.spk2id = json.load(f) + + utt2spk_path = os.path.join(processed_data_dir, cfg.preprocess.utt2spk) + self.utt2spk = dict() + with open(utt2spk_path, "r") as f: + for line in f.readlines(): + utt, spk = line.strip().split("\t") + self.utt2spk[utt] = spk + + if cfg.preprocess.use_uv: + self.utt2uv_path = {} + for utt_info in self.metadata: + dataset = utt_info["Dataset"] + uid = utt_info["Uid"] + utt = "{}_{}".format(dataset, uid) + self.utt2uv_path[utt] = os.path.join( + cfg.preprocess.processed_dir, + dataset, + cfg.preprocess.uv_dir, + uid + ".npy", + ) + + if cfg.preprocess.use_frame_pitch: + self.utt2frame_pitch_path = {} + for utt_info in self.metadata: + dataset = utt_info["Dataset"] + uid = utt_info["Uid"] + utt = "{}_{}".format(dataset, uid) + + self.utt2frame_pitch_path[utt] = os.path.join( + cfg.preprocess.processed_dir, + dataset, + cfg.preprocess.pitch_dir, + uid + ".npy", + ) + + if cfg.preprocess.use_frame_energy: + self.utt2frame_energy_path = {} + for utt_info in self.metadata: + dataset = utt_info["Dataset"] + uid = utt_info["Uid"] + utt = "{}_{}".format(dataset, uid) + + self.utt2frame_energy_path[utt] = os.path.join( + cfg.preprocess.processed_dir, + dataset, + cfg.preprocess.energy_dir, + uid + ".npy", + ) + + if cfg.preprocess.use_mel: + self.utt2mel_path = {} + for utt_info in self.metadata: + dataset = utt_info["Dataset"] + uid = utt_info["Uid"] + utt = "{}_{}".format(dataset, uid) + + self.utt2mel_path[utt] = os.path.join( + cfg.preprocess.processed_dir, + dataset, + cfg.preprocess.mel_dir, + uid + ".npy", + ) + + if cfg.preprocess.use_linear: + self.utt2linear_path = {} + for utt_info in self.metadata: + dataset = utt_info["Dataset"] + uid = utt_info["Uid"] + utt = "{}_{}".format(dataset, uid) + + self.utt2linear_path[utt] = os.path.join( + cfg.preprocess.processed_dir, + dataset, + cfg.preprocess.linear_dir, + uid + ".npy", + ) + + if cfg.preprocess.use_audio: + self.utt2audio_path = {} + for utt_info in self.metadata: + dataset = utt_info["Dataset"] + uid = utt_info["Uid"] + utt = "{}_{}".format(dataset, uid) + + if cfg.preprocess.extract_audio: + self.utt2audio_path[utt] = os.path.join( + cfg.preprocess.processed_dir, + dataset, + cfg.preprocess.audio_dir, + uid + ".wav", + ) + else: + self.utt2audio_path[utt] = utt_info["Path"] + + # self.utt2audio_path[utt] = os.path.join( + # cfg.preprocess.processed_dir, + # dataset, + # cfg.preprocess.audio_dir, + # uid + ".numpy", + # ) + + elif cfg.preprocess.use_label: + self.utt2label_path = {} + for utt_info in self.metadata: + dataset = utt_info["Dataset"] + uid = utt_info["Uid"] + utt = "{}_{}".format(dataset, uid) + + self.utt2label_path[utt] = os.path.join( + cfg.preprocess.processed_dir, + dataset, + cfg.preprocess.label_dir, + uid + ".npy", + ) + elif cfg.preprocess.use_one_hot: + self.utt2one_hot_path = {} + for utt_info in self.metadata: + dataset = utt_info["Dataset"] + uid = utt_info["Uid"] + utt = "{}_{}".format(dataset, uid) + + self.utt2one_hot_path[utt] = os.path.join( + cfg.preprocess.processed_dir, + dataset, + cfg.preprocess.one_hot_dir, + uid + ".npy", + ) + + if cfg.preprocess.use_text or cfg.preprocess.use_phone: + self.utt2seq = {} + for utt_info in self.metadata: + dataset = utt_info["Dataset"] + uid = utt_info["Uid"] + utt = "{}_{}".format(dataset, uid) + + if cfg.preprocess.use_text: + text = utt_info["Text"] + sequence = text_to_sequence(text, cfg.preprocess.text_cleaners) + elif cfg.preprocess.use_phone: + # load phoneme squence from phone file + phone_path = os.path.join( + processed_data_dir, cfg.preprocess.phone_dir, uid + ".phone" + ) + with open(phone_path, "r") as fin: + phones = fin.readlines() + assert len(phones) == 1 + phones = phones[0].strip() + phones_seq = phones.split(" ") + + phon_id_collator = phoneIDCollation(cfg, dataset=dataset) + sequence = phon_id_collator.get_phone_id_sequence(cfg, phones_seq) + + if cfg.preprocess.add_blank: + sequence = intersperse(sequence, 0) + + self.utt2seq[utt] = sequence + + def __getitem__(self, index): + utt_info = self.metadata[index] + + dataset = utt_info["Dataset"] + uid = utt_info["Uid"] + utt = "{}_{}".format(dataset, uid) + + single_feature = dict() + + if self.cfg.preprocess.use_spkid: + single_feature["spk_id"] = np.array( + [self.spk2id[self.utt2spk[utt]]], dtype=np.int32 + ) + + if self.cfg.preprocess.use_mel: + mel = np.load(self.utt2mel_path[utt]) + assert mel.shape[0] == self.cfg.preprocess.n_mel # [n_mels, T] + if self.cfg.preprocess.use_min_max_norm_mel: + # do mel norm + mel = cal_normalized_mel(mel, utt_info["Dataset"], self.cfg.preprocess) + + if "target_len" not in single_feature.keys(): + single_feature["target_len"] = mel.shape[1] + single_feature["mel"] = mel.T # [T, n_mels] + + if self.cfg.preprocess.use_linear: + linear = np.load(self.utt2linear_path[utt]) + if "target_len" not in single_feature.keys(): + single_feature["target_len"] = linear.shape[1] + single_feature["linear"] = linear.T # [T, n_linear] + + if self.cfg.preprocess.use_frame_pitch: + frame_pitch_path = self.utt2frame_pitch_path[utt] + frame_pitch = np.load(frame_pitch_path) + if "target_len" not in single_feature.keys(): + single_feature["target_len"] = len(frame_pitch) + aligned_frame_pitch = align_length( + frame_pitch, single_feature["target_len"] + ) + single_feature["frame_pitch"] = aligned_frame_pitch + + if self.cfg.preprocess.use_uv: + frame_uv_path = self.utt2uv_path[utt] + frame_uv = np.load(frame_uv_path) + aligned_frame_uv = align_length(frame_uv, single_feature["target_len"]) + aligned_frame_uv = [ + 0 if frame_uv else 1 for frame_uv in aligned_frame_uv + ] + aligned_frame_uv = np.array(aligned_frame_uv) + single_feature["frame_uv"] = aligned_frame_uv + + if self.cfg.preprocess.use_frame_energy: + frame_energy_path = self.utt2frame_energy_path[utt] + frame_energy = np.load(frame_energy_path) + if "target_len" not in single_feature.keys(): + single_feature["target_len"] = len(frame_energy) + aligned_frame_energy = align_length( + frame_energy, single_feature["target_len"] + ) + single_feature["frame_energy"] = aligned_frame_energy + + if self.cfg.preprocess.use_audio: + audio, sr = torchaudio.load(self.utt2audio_path[utt]) + audio = audio.cpu().numpy().squeeze() + single_feature["audio"] = audio + single_feature["audio_len"] = audio.shape[0] + + if self.cfg.preprocess.use_phone or self.cfg.preprocess.use_text: + single_feature["phone_seq"] = np.array(self.utt2seq[utt]) + single_feature["phone_len"] = len(self.utt2seq[utt]) + + return single_feature + + def __len__(self): + return super().__len__() + + def get_metadata(self): + return super().get_metadata() + + +class TTSCollator(BaseOfflineCollator): + """Zero-pads model inputs and targets based on number of frames per step""" + + def __init__(self, cfg): + super().__init__(cfg) + + def __call__(self, batch): + parsed_batch_features = super().__call__(batch) + return parsed_batch_features + + +class TTSTestDataset(BaseTestDataset): + def __init__(self, args, cfg): + self.cfg = cfg + + # inference from test list file + if args.test_list_file is not None: + # construst metadata + self.metadata = [] + + with open(args.test_list_file, "r") as fin: + for idx, line in enumerate(fin.readlines()): + utt_info = {} + + utt_info["Dataset"] = "test" + utt_info["Text"] = line.strip() + utt_info["Uid"] = str(idx) + self.metadata.append(utt_info) + + else: + assert args.testing_set + self.metafile_path = os.path.join( + cfg.preprocess.processed_dir, + args.dataset, + "{}.json".format(args.testing_set), + ) + self.metadata = self.get_metadata() + + def __getitem__(self, index): + single_feature = {} + + return single_feature + + def __len__(self): + return len(self.metadata) + + +class TTSTestCollator(BaseTestCollator): + """Zero-pads model inputs and targets based on number of frames per step""" + + def __init__(self, cfg): + self.cfg = cfg + + def __call__(self, batch): + packed_batch_features = dict() + + # mel: [b, T, n_mels] + # frame_pitch, frame_energy: [1, T] + # target_len: [1] + # spk_id: [b, 1] + # mask: [b, T, 1] + + for key in batch[0].keys(): + if key == "target_len": + packed_batch_features["target_len"] = torch.LongTensor( + [b["target_len"] for b in batch] + ) + masks = [ + torch.ones((b["target_len"], 1), dtype=torch.long) for b in batch + ] + packed_batch_features["mask"] = pad_sequence( + masks, batch_first=True, padding_value=0 + ) + elif key == "phone_len": + packed_batch_features["phone_len"] = torch.LongTensor( + [b["phone_len"] for b in batch] + ) + masks = [ + torch.ones((b["phone_len"], 1), dtype=torch.long) for b in batch + ] + packed_batch_features["phn_mask"] = pad_sequence( + masks, batch_first=True, padding_value=0 + ) + elif key == "audio_len": + packed_batch_features["audio_len"] = torch.LongTensor( + [b["audio_len"] for b in batch] + ) + masks = [ + torch.ones((b["audio_len"], 1), dtype=torch.long) for b in batch + ] + else: + values = [torch.from_numpy(b[key]) for b in batch] + packed_batch_features[key] = pad_sequence( + values, batch_first=True, padding_value=0 + ) + return packed_batch_features diff --git a/models/tts/base/tts_inferece.py b/models/tts/base/tts_inferece.py new file mode 100644 index 0000000..f49ace0 --- /dev/null +++ b/models/tts/base/tts_inferece.py @@ -0,0 +1,278 @@ +# Copyright (c) 2023 Amphion. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +import os +import torch +import time +import accelerate +import random +import numpy as np +from tqdm import tqdm +from accelerate.logging import get_logger +from torch.utils.data import DataLoader +from safetensors.torch import load_file + + +from abc import abstractmethod +from pathlib import Path +from utils.io import save_audio +from utils.util import load_config +from models.vocoders.vocoder_inference import synthesis + + +class TTSInference(object): + def __init__(self, args=None, cfg=None): + super().__init__() + + start = time.monotonic_ns() + self.args = args + self.cfg = cfg + self.infer_type = args.mode + + # get exp_dir + if self.args.acoustics_dir is not None: + self.exp_dir = self.args.acoustics_dir + elif self.args.checkpoint_path is not None: + self.exp_dir = os.path.dirname(os.path.dirname(self.args.checkpoint_path)) + + # Init accelerator + self.accelerator = accelerate.Accelerator() + self.accelerator.wait_for_everyone() + self.device = self.accelerator.device + + # Get logger + with self.accelerator.main_process_first(): + self.logger = get_logger("inference", log_level=args.log_level) + + # Log some info + self.logger.info("=" * 56) + self.logger.info("||\t\t" + "New inference process started." + "\t\t||") + self.logger.info("=" * 56) + self.logger.info("\n") + + self.acoustic_model_dir = args.acoustics_dir + self.logger.debug(f"Acoustic model dir: {args.acoustics_dir}") + + if args.vocoder_dir is not None: + self.vocoder_dir = args.vocoder_dir + self.logger.debug(f"Vocoder dir: {args.vocoder_dir}") + + os.makedirs(args.output_dir, exist_ok=True) + + # Set random seed + with self.accelerator.main_process_first(): + start = time.monotonic_ns() + self._set_random_seed(self.cfg.train.random_seed) + end = time.monotonic_ns() + self.logger.debug( + f"Setting random seed done in {(end - start) / 1e6:.2f}ms" + ) + self.logger.debug(f"Random seed: {self.cfg.train.random_seed}") + + # Setup data loader + if self.infer_type == "batch": + with self.accelerator.main_process_first(): + self.logger.info("Building dataset...") + start = time.monotonic_ns() + self.test_dataloader = self._build_test_dataloader() + end = time.monotonic_ns() + self.logger.info( + f"Building dataset done in {(end - start) / 1e6:.2f}ms" + ) + + # Build model + with self.accelerator.main_process_first(): + self.logger.info("Building model...") + start = time.monotonic_ns() + self.model = self._build_model() + end = time.monotonic_ns() + self.logger.info(f"Building model done in {(end - start) / 1e6:.3f}ms") + + # Init with accelerate + self.logger.info("Initializing accelerate...") + start = time.monotonic_ns() + self.accelerator = accelerate.Accelerator() + self.model = self.accelerator.prepare(self.model) + if self.infer_type == "batch": + self.test_dataloader = self.accelerator.prepare(self.test_dataloader) + end = time.monotonic_ns() + self.accelerator.wait_for_everyone() + self.logger.info(f"Initializing accelerate done in {(end - start) / 1e6:.3f}ms") + + with self.accelerator.main_process_first(): + self.logger.info("Loading checkpoint...") + start = time.monotonic_ns() + if args.acoustics_dir is not None: + self._load_model( + checkpoint_dir=os.path.join(args.acoustics_dir, "checkpoint") + ) + elif args.checkpoint_path is not None: + self._load_model(checkpoint_path=args.checkpoint_path) + else: + print("Either checkpoint dir or checkpoint path should be provided.") + + end = time.monotonic_ns() + self.logger.info(f"Loading checkpoint done in {(end - start) / 1e6:.3f}ms") + + self.model.eval() + self.accelerator.wait_for_everyone() + + def _build_test_dataset(self): + pass + + def _build_model(self): + pass + + # TODO: LEGACY CODE + def _build_test_dataloader(self): + datasets, collate = self._build_test_dataset() + self.test_dataset = datasets(self.args, self.cfg) + self.test_collate = collate(self.cfg) + self.test_batch_size = min( + self.cfg.train.batch_size, len(self.test_dataset.metadata) + ) + test_dataloader = DataLoader( + self.test_dataset, + collate_fn=self.test_collate, + num_workers=1, + batch_size=self.test_batch_size, + shuffle=False, + ) + return test_dataloader + + def _load_model( + self, + checkpoint_dir: str = None, + checkpoint_path: str = None, + old_mode: bool = False, + ): + r"""Load model from checkpoint. If checkpoint_path is None, it will + load the latest checkpoint in checkpoint_dir. If checkpoint_path is not + None, it will load the checkpoint specified by checkpoint_path. **Only use this + method after** ``accelerator.prepare()``. + """ + + if checkpoint_path is None: + assert checkpoint_dir is not None + # Load the latest accelerator state dicts + ls = [ + str(i) for i in Path(checkpoint_dir).glob("*") if not "audio" in str(i) + ] + ls.sort(key=lambda x: int(x.split("_")[-3].split("-")[-1]), reverse=True) + checkpoint_path = ls[0] + + if ( + Path(os.path.join(checkpoint_path, "model.safetensors")).exists() + and accelerate.__version__ < "0.25" + ): + self.model.load_state_dict( + load_file(os.path.join(checkpoint_path, "model.safetensors")), + strict=False, + ) + else: + self.accelerator.load_state(str(checkpoint_path)) + return str(checkpoint_path) + + def inference(self): + if self.infer_type == "single": + out_dir = os.path.join(self.args.output_dir, "single") + os.makedirs(out_dir, exist_ok=True) + + pred_audio = self.inference_for_single_utterance() + save_path = os.path.join(out_dir, "test_pred.wav") + save_audio(save_path, pred_audio, self.cfg.preprocess.sample_rate) + + elif self.infer_type == "batch": + out_dir = os.path.join(self.args.output_dir, "batch") + os.makedirs(out_dir, exist_ok=True) + + pred_audio_list = self.inference_for_batches() + for it, wav in zip(self.test_dataset.metadata, pred_audio_list): + uid = it["Uid"] + save_audio( + os.path.join(out_dir, f"{uid}.wav"), + wav.numpy(), + self.cfg.preprocess.sample_rate, + add_silence=True, + turn_up=True, + ) + tmp_file = os.path.join(out_dir, f"{uid}.pt") + if os.path.exists(tmp_file): + os.remove(tmp_file) + print("Saved to: ", out_dir) + + @torch.inference_mode() + def inference_for_batches(self): + y_pred = [] + for i, batch in tqdm(enumerate(self.test_dataloader)): + y_pred, mel_lens, _ = self._inference_each_batch(batch) + y_ls = y_pred.chunk(self.test_batch_size) + tgt_ls = mel_lens.chunk(self.test_batch_size) + j = 0 + for it, l in zip(y_ls, tgt_ls): + l = l.item() + it = it.squeeze(0)[:l].detach().cpu() + + uid = self.test_dataset.metadata[i * self.test_batch_size + j]["Uid"] + torch.save(it, os.path.join(self.args.output_dir, f"{uid}.pt")) + j += 1 + + vocoder_cfg, vocoder_ckpt = self._parse_vocoder(self.args.vocoder_dir) + res = synthesis( + cfg=vocoder_cfg, + vocoder_weight_file=vocoder_ckpt, + n_samples=None, + pred=[ + torch.load( + os.path.join(self.args.output_dir, "{}.pt".format(item["Uid"])) + ).numpy() + for item in self.test_dataset.metadata + ], + ) + for it, wav in zip(self.test_dataset.metadata, res): + uid = it["Uid"] + save_audio( + os.path.join(self.args.output_dir, f"{uid}.wav"), + wav.numpy(), + 22050, + add_silence=True, + turn_up=True, + ) + + @abstractmethod + @torch.inference_mode() + def _inference_each_batch(self, batch_data): + pass + + def inference_for_single_utterance(self, text): + pass + + def synthesis_by_vocoder(self, pred): + audios_pred = synthesis( + self.vocoder_cfg, + self.checkpoint_dir_vocoder, + len(pred), + pred, + ) + + return audios_pred + + @staticmethod + def _parse_vocoder(vocoder_dir): + r"""Parse vocoder config""" + vocoder_dir = os.path.abspath(vocoder_dir) + ckpt_list = [ckpt for ckpt in Path(vocoder_dir).glob("*.pt")] + ckpt_list.sort(key=lambda x: int(x.stem), reverse=True) + ckpt_path = str(ckpt_list[0]) + vocoder_cfg = load_config( + os.path.join(vocoder_dir, "args.json"), lowercase=True + ) + return vocoder_cfg, ckpt_path + + def _set_random_seed(self, seed): + """Set random seed for all possible random modules.""" + random.seed(seed) + np.random.seed(seed) + torch.random.manual_seed(seed) diff --git a/models/tts/base/tts_trainer.py b/models/tts/base/tts_trainer.py new file mode 100644 index 0000000..1ebfca6 --- /dev/null +++ b/models/tts/base/tts_trainer.py @@ -0,0 +1,721 @@ +# Copyright (c) 2023 Amphion. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +import json +import os +import shutil +import torch +import time +from pathlib import Path +import torch +from tqdm import tqdm +import re +import logging +import json5 +import accelerate +from accelerate.logging import get_logger +from accelerate.utils import ProjectConfiguration +from torch.utils.data import ConcatDataset, DataLoader +from accelerate import DistributedDataParallelKwargs +from schedulers.scheduler import Eden +from models.base.base_sampler import build_samplers +from models.base.new_trainer import BaseTrainer + + +class TTSTrainer(BaseTrainer): + r"""The base trainer for all TTS models. It inherits from BaseTrainer and implements + ``build_criterion``, ``_build_dataset`` and ``_build_singer_lut`` methods. You can inherit from this + class, and implement ``_build_model``, ``_forward_step``. + """ + + def __init__(self, args=None, cfg=None): + self.args = args + self.cfg = cfg + + cfg.exp_name = args.exp_name + + # init with accelerate + self._init_accelerator() + self.accelerator.wait_for_everyone() + + with self.accelerator.main_process_first(): + self.logger = get_logger(args.exp_name, log_level="INFO") + + # Log some info + self.logger.info("=" * 56) + self.logger.info("||\t\t" + "New training process started." + "\t\t||") + self.logger.info("=" * 56) + self.logger.info("\n") + self.logger.debug(f"Using {args.log_level.upper()} logging level.") + self.logger.info(f"Experiment name: {args.exp_name}") + self.logger.info(f"Experiment directory: {self.exp_dir}") + self.checkpoint_dir = os.path.join(self.exp_dir, "checkpoint") + if self.accelerator.is_main_process: + os.makedirs(self.checkpoint_dir, exist_ok=True) + self.logger.debug(f"Checkpoint directory: {self.checkpoint_dir}") + + # init counts + self.batch_count: int = 0 + self.step: int = 0 + self.epoch: int = 0 + self.max_epoch = ( + self.cfg.train.max_epoch if self.cfg.train.max_epoch > 0 else float("inf") + ) + self.logger.info( + "Max epoch: {}".format( + self.max_epoch if self.max_epoch < float("inf") else "Unlimited" + ) + ) + + # Check values + if self.accelerator.is_main_process: + self.__check_basic_configs() + # Set runtime configs + self.save_checkpoint_stride = self.cfg.train.save_checkpoint_stride + self.checkpoints_path = [ + [] for _ in range(len(self.save_checkpoint_stride)) + ] + self.keep_last = [ + i if i > 0 else float("inf") for i in self.cfg.train.keep_last + ] + self.run_eval = self.cfg.train.run_eval + + # set random seed + with self.accelerator.main_process_first(): + start = time.monotonic_ns() + self._set_random_seed(self.cfg.train.random_seed) + end = time.monotonic_ns() + self.logger.debug( + f"Setting random seed done in {(end - start) / 1e6:.2f}ms" + ) + self.logger.debug(f"Random seed: {self.cfg.train.random_seed}") + + # setup data_loader + with self.accelerator.main_process_first(): + self.logger.info("Building dataset...") + start = time.monotonic_ns() + self.train_dataloader, self.valid_dataloader = self._build_dataloader() + end = time.monotonic_ns() + self.logger.info(f"Building dataset done in {(end - start) / 1e6:.2f}ms") + + # save phone table to exp dir. Should be done before building model due to loading phone table in model + if cfg.preprocess.use_phone and cfg.preprocess.phone_extractor != "lexicon": + self._save_phone_symbols_file_to_exp_path() + + # setup model + with self.accelerator.main_process_first(): + self.logger.info("Building model...") + start = time.monotonic_ns() + self.model = self._build_model() + end = time.monotonic_ns() + self.logger.debug(self.model) + self.logger.info(f"Building model done in {(end - start) / 1e6:.2f}ms") + self.logger.info( + f"Model parameters: {self.__count_parameters(self.model)/1e6:.2f}M" + ) + + # optimizer & scheduler + with self.accelerator.main_process_first(): + self.logger.info("Building optimizer and scheduler...") + start = time.monotonic_ns() + self.optimizer = self._build_optimizer() + self.scheduler = self._build_scheduler() + end = time.monotonic_ns() + self.logger.info( + f"Building optimizer and scheduler done in {(end - start) / 1e6:.2f}ms" + ) + + # create criterion + with self.accelerator.main_process_first(): + self.logger.info("Building criterion...") + start = time.monotonic_ns() + self.criterion = self._build_criterion() + end = time.monotonic_ns() + self.logger.info(f"Building criterion done in {(end - start) / 1e6:.2f}ms") + + # Resume or Finetune + with self.accelerator.main_process_first(): + self._check_resume() + + # accelerate prepare + self.logger.info("Initializing accelerate...") + start = time.monotonic_ns() + self._accelerator_prepare() + end = time.monotonic_ns() + self.logger.info(f"Initializing accelerate done in {(end - start) / 1e6:.2f}ms") + + # save config file path + self.config_save_path = os.path.join(self.exp_dir, "args.json") + self.device = self.accelerator.device + + if cfg.preprocess.use_spkid and cfg.train.multi_speaker_training: + self.speakers = self._build_speaker_lut() + self.utt2spk_dict = self._build_utt2spk_dict() + + # Only for TTS tasks + self.task_type = "TTS" + self.logger.info("Task type: {}".format(self.task_type)) + + def _check_resume(self): + # if args.resume: + if self.args.resume or ( + self.cfg.model_type == "VALLE" and self.args.train_stage == 2 + ): + checkpoint_dir = self.checkpoint_dir + if self.cfg.model_type == "VALLE" and self.args.train_stage == 2: + ls = [str(i) for i in Path(checkpoint_dir).glob("*")] + if ( + self.args.checkpoint_path is None or len(ls) == 0 + ): # Train stage 2 from scratch using the checkpoint of stage 1 + assert ( + self.args.ar_model_ckpt_dir is not None + ), "Error: ar_model_ckpt_dir should be set to train nar model." + self.args.resume_type = "finetune" + checkpoint_dir = self.args.ar_model_ckpt_dir + self.logger.info( + f"Training NAR model at stage 2 using the checkpoint of AR model at stage 1." + ) + + self.logger.info(f"Resuming from checkpoint: {checkpoint_dir}") + start = time.monotonic_ns() + self.ckpt_path = self._load_model( + checkpoint_dir, self.args.checkpoint_path, self.args.resume_type + ) + self.logger.info(f"Checkpoint path: {self.ckpt_path}") + end = time.monotonic_ns() + self.logger.info( + f"Resuming from checkpoint done in {(end - start) / 1e6:.2f}ms" + ) + self.checkpoints_path = json.load( + open(os.path.join(self.ckpt_path, "ckpts.json"), "r") + ) + + def _init_accelerator(self): + self.exp_dir = os.path.join( + os.path.abspath(self.cfg.log_dir), self.args.exp_name + ) + project_config = ProjectConfiguration( + project_dir=self.exp_dir, + logging_dir=os.path.join(self.exp_dir, "log"), + ) + kwargs = DistributedDataParallelKwargs(find_unused_parameters=True) + self.accelerator = accelerate.Accelerator( + gradient_accumulation_steps=self.cfg.train.gradient_accumulation_step, + log_with=self.cfg.train.tracker, + project_config=project_config, + kwargs_handlers=[kwargs], + ) + if self.accelerator.is_main_process: + os.makedirs(project_config.project_dir, exist_ok=True) + os.makedirs(project_config.logging_dir, exist_ok=True) + with self.accelerator.main_process_first(): + self.accelerator.init_trackers(self.args.exp_name) + + def _accelerator_prepare(self): + ( + self.train_dataloader, + self.valid_dataloader, + ) = self.accelerator.prepare( + self.train_dataloader, + self.valid_dataloader, + ) + + if isinstance(self.model, dict): + for key in self.model.keys(): + self.model[key] = self.accelerator.prepare(self.model[key]) + else: + self.model = self.accelerator.prepare(self.model) + + if isinstance(self.optimizer, dict): + for key in self.optimizer.keys(): + self.optimizer[key] = self.accelerator.prepare(self.optimizer[key]) + else: + self.optimizer = self.accelerator.prepare(self.optimizer) + + if isinstance(self.scheduler, dict): + for key in self.scheduler.keys(): + self.scheduler[key] = self.accelerator.prepare(self.scheduler[key]) + else: + self.scheduler = self.accelerator.prepare(self.scheduler) + + ### Following are methods only for TTS tasks ### + def _build_dataset(self): + pass + + def _build_criterion(self): + pass + + def _build_model(self): + pass + + def _build_dataloader(self): + """Build dataloader which merges a series of datasets.""" + # Build dataset instance for each dataset and combine them by ConcatDataset + Dataset, Collator = self._build_dataset() + + # Build train set + datasets_list = [] + for dataset in self.cfg.dataset: + subdataset = Dataset(self.cfg, dataset, is_valid=False) + datasets_list.append(subdataset) + train_dataset = ConcatDataset(datasets_list) + train_collate = Collator(self.cfg) + _, batch_sampler = build_samplers(train_dataset, self.cfg, self.logger, "train") + train_loader = DataLoader( + train_dataset, + collate_fn=train_collate, + batch_sampler=batch_sampler, + num_workers=self.cfg.train.dataloader.num_worker, + pin_memory=self.cfg.train.dataloader.pin_memory, + ) + + # Build test set + datasets_list = [] + for dataset in self.cfg.dataset: + subdataset = Dataset(self.cfg, dataset, is_valid=True) + datasets_list.append(subdataset) + valid_dataset = ConcatDataset(datasets_list) + valid_collate = Collator(self.cfg) + _, batch_sampler = build_samplers(valid_dataset, self.cfg, self.logger, "valid") + valid_loader = DataLoader( + valid_dataset, + collate_fn=valid_collate, + batch_sampler=batch_sampler, + num_workers=self.cfg.train.dataloader.num_worker, + pin_memory=self.cfg.train.dataloader.pin_memory, + ) + return train_loader, valid_loader + + def _build_optimizer(self): + pass + + def _build_scheduler(self): + pass + + def _load_model(self, checkpoint_dir, checkpoint_path=None, resume_type="resume"): + """Load model from checkpoint. If a folder is given, it will + load the latest checkpoint in checkpoint_dir. If a path is given + it will load the checkpoint specified by checkpoint_path. + **Only use this method after** ``accelerator.prepare()``. + """ + if checkpoint_path is None or checkpoint_path == "": + ls = [str(i) for i in Path(checkpoint_dir).glob("*")] + ls.sort(key=lambda x: int(x.split("_")[-3].split("-")[-1]), reverse=True) + checkpoint_path = ls[0] + self.logger.info("Load model from {}".format(checkpoint_path)) + print("Load model from {}".format(checkpoint_path)) + if resume_type == "resume": + self.accelerator.load_state(checkpoint_path) + self.epoch = int(checkpoint_path.split("_")[-3].split("-")[-1]) + 1 + self.step = int(checkpoint_path.split("_")[-2].split("-")[-1]) + 1 + elif resume_type == "finetune": + if isinstance(self.model, dict): + for idx, sub_model in enumerate(self.model.keys()): + if idx == 0: + ckpt_name = "pytorch_model.bin" + else: + ckpt_name = "pytorch_model_{}.bin".format(idx) + + self.model[sub_model].load_state_dict( + torch.load(os.path.join(checkpoint_path, ckpt_name)) + ) + self.model[sub_model].cuda(self.accelerator.device) + else: + self.model.load_state_dict( + torch.load(os.path.join(checkpoint_path, "pytorch_model.bin")) + ) + self.model.cuda(self.accelerator.device) + self.logger.info("Load model weights for finetune SUCCESS!") + + else: + raise ValueError("Unsupported resume type: {}".format(resume_type)) + + return checkpoint_path + + ### THIS IS MAIN ENTRY ### + def train_loop(self): + r"""Training loop. The public entry of training process.""" + # Wait everyone to prepare before we move on + self.accelerator.wait_for_everyone() + # dump config file + if self.accelerator.is_main_process: + self.__dump_cfg(self.config_save_path) + + # self.optimizer.zero_grad() + # Wait to ensure good to go + + self.accelerator.wait_for_everyone() + while self.epoch < self.max_epoch: + self.logger.info("\n") + self.logger.info("-" * 32) + self.logger.info("Epoch {}: ".format(self.epoch)) + + # Do training & validating epoch + train_total_loss, train_losses = self._train_epoch() + if isinstance(train_losses, dict): + for key, loss in train_losses.items(): + self.logger.info(" |- Train/{} Loss: {:.6f}".format(key, loss)) + self.accelerator.log( + {"Epoch/Train {} Loss".format(key): loss}, + step=self.epoch, + ) + + valid_total_loss, valid_losses = self._valid_epoch() + if isinstance(valid_losses, dict): + for key, loss in valid_losses.items(): + self.logger.info(" |- Valid/{} Loss: {:.6f}".format(key, loss)) + self.accelerator.log( + {"Epoch/Valid {} Loss".format(key): loss}, + step=self.epoch, + ) + + self.logger.info(" |- Train/Loss: {:.6f}".format(train_total_loss)) + self.logger.info(" |- Valid/Loss: {:.6f}".format(valid_total_loss)) + self.accelerator.log( + { + "Epoch/Train Loss": train_total_loss, + "Epoch/Valid Loss": valid_total_loss, + }, + step=self.epoch, + ) + + self.accelerator.wait_for_everyone() + + # Check if hit save_checkpoint_stride and run_eval + run_eval = False + if self.accelerator.is_main_process: + save_checkpoint = False + hit_dix = [] + for i, num in enumerate(self.save_checkpoint_stride): + if self.epoch % num == 0: + save_checkpoint = True + hit_dix.append(i) + run_eval |= self.run_eval[i] + + self.accelerator.wait_for_everyone() + if self.accelerator.is_main_process and save_checkpoint: + path = os.path.join( + self.checkpoint_dir, + "epoch-{:04d}_step-{:07d}_loss-{:.6f}".format( + self.epoch, self.step, train_total_loss + ), + ) + self.accelerator.save_state(path) + + json.dump( + self.checkpoints_path, + open(os.path.join(path, "ckpts.json"), "w"), + ensure_ascii=False, + indent=4, + ) + + # Remove old checkpoints + to_remove = [] + for idx in hit_dix: + self.checkpoints_path[idx].append(path) + while len(self.checkpoints_path[idx]) > self.keep_last[idx]: + to_remove.append((idx, self.checkpoints_path[idx].pop(0))) + + # Search conflicts + total = set() + for i in self.checkpoints_path: + total |= set(i) + do_remove = set() + for idx, path in to_remove[::-1]: + if path in total: + self.checkpoints_path[idx].insert(0, path) + else: + do_remove.add(path) + + # Remove old checkpoints + for path in do_remove: + shutil.rmtree(path, ignore_errors=True) + self.logger.debug(f"Remove old checkpoint: {path}") + + self.accelerator.wait_for_everyone() + if run_eval: + # TODO: run evaluation + pass + + # Update info for each epoch + self.epoch += 1 + + # Finish training and save final checkpoint + self.accelerator.wait_for_everyone() + if self.accelerator.is_main_process: + path = os.path.join( + self.checkpoint_dir, + "final_epoch-{:04d}_step-{:07d}_loss-{:.6f}".format( + self.epoch, self.step, valid_total_loss + ), + ) + self.accelerator.save_state( + os.path.join( + self.checkpoint_dir, + "final_epoch-{:04d}_step-{:07d}_loss-{:.6f}".format( + self.epoch, self.step, valid_total_loss + ), + ) + ) + + json.dump( + self.checkpoints_path, + open(os.path.join(path, "ckpts.json"), "w"), + ensure_ascii=False, + indent=4, + ) + + self.accelerator.end_training() + + ### Following are methods that can be used directly in child classes ### + def _train_epoch(self): + r"""Training epoch. Should return average loss of a batch (sample) over + one epoch. See ``train_loop`` for usage. + """ + if isinstance(self.model, dict): + for key in self.model.keys(): + self.model[key].train() + else: + self.model.train() + + epoch_sum_loss: float = 0.0 + epoch_losses: dict = {} + epoch_step: int = 0 + for batch in tqdm( + self.train_dataloader, + desc=f"Training Epoch {self.epoch}", + unit="batch", + colour="GREEN", + leave=False, + dynamic_ncols=True, + smoothing=0.04, + disable=not self.accelerator.is_main_process, + ): + # Do training step and BP + with self.accelerator.accumulate(self.model): + total_loss, train_losses, _ = self._train_step(batch) + self.batch_count += 1 + + # Update info for each step + # TODO: step means BP counts or batch counts? + if self.batch_count % self.cfg.train.gradient_accumulation_step == 0: + if isinstance(self.scheduler, dict): + for key in self.scheduler.keys(): + self.scheduler[key].step() + else: + if isinstance(self.scheduler, Eden): + self.scheduler.step_batch(self.step) + else: + self.scheduler.step() + + epoch_sum_loss += total_loss + + if isinstance(train_losses, dict): + for key, value in train_losses.items(): + epoch_losses[key] += value + + if isinstance(train_losses, dict): + for key, loss in train_losses.items(): + self.accelerator.log( + {"Epoch/Train {} Loss".format(key): loss}, + step=self.step, + ) + + self.step += 1 + epoch_step += 1 + + self.accelerator.wait_for_everyone() + + epoch_sum_loss = ( + epoch_sum_loss + / len(self.train_dataloader) + * self.cfg.train.gradient_accumulation_step + ) + + for key in epoch_losses.keys(): + epoch_losses[key] = ( + epoch_losses[key] + / len(self.train_dataloader) + * self.cfg.train.gradient_accumulation_step + ) + + return epoch_sum_loss, epoch_losses + + @torch.inference_mode() + def _valid_epoch(self): + r"""Testing epoch. Should return average loss of a batch (sample) over + one epoch. See ``train_loop`` for usage. + """ + if isinstance(self.model, dict): + for key in self.model.keys(): + self.model[key].eval() + else: + self.model.eval() + + epoch_sum_loss = 0.0 + epoch_losses = dict() + for batch in tqdm( + self.valid_dataloader, + desc=f"Validating Epoch {self.epoch}", + unit="batch", + colour="GREEN", + leave=False, + dynamic_ncols=True, + smoothing=0.04, + disable=not self.accelerator.is_main_process, + ): + total_loss, valid_losses, valid_stats = self._valid_step(batch) + epoch_sum_loss += total_loss + if isinstance(valid_losses, dict): + for key, value in valid_losses.items(): + if key not in epoch_losses.keys(): + epoch_losses[key] = value + else: + epoch_losses[key] += value + + epoch_sum_loss = epoch_sum_loss / len(self.valid_dataloader) + for key in epoch_losses.keys(): + epoch_losses[key] = epoch_losses[key] / len(self.valid_dataloader) + + self.accelerator.wait_for_everyone() + + return epoch_sum_loss, epoch_losses + + def _train_step(self): + pass + + def _valid_step(self, batch): + pass + + def _inference(self): + pass + + def _is_valid_pattern(self, directory_name): + directory_name = str(directory_name) + pattern = r"^epoch-\d{4}_step-\d{7}_loss-\d{1}\.\d{6}" + return re.match(pattern, directory_name) is not None + + def _check_basic_configs(self): + if self.cfg.train.gradient_accumulation_step <= 0: + self.logger.fatal("Invalid gradient_accumulation_step value!") + self.logger.error( + f"Invalid gradient_accumulation_step value: {self.cfg.train.gradient_accumulation_step}. It should be positive." + ) + self.accelerator.end_training() + raise ValueError( + f"Invalid gradient_accumulation_step value: {self.cfg.train.gradient_accumulation_step}. It should be positive." + ) + + def __dump_cfg(self, path): + os.makedirs(os.path.dirname(path), exist_ok=True) + json5.dump( + self.cfg, + open(path, "w"), + indent=4, + sort_keys=True, + ensure_ascii=False, + quote_keys=True, + ) + + def __check_basic_configs(self): + if self.cfg.train.gradient_accumulation_step <= 0: + self.logger.fatal("Invalid gradient_accumulation_step value!") + self.logger.error( + f"Invalid gradient_accumulation_step value: {self.cfg.train.gradient_accumulation_step}. It should be positive." + ) + self.accelerator.end_training() + raise ValueError( + f"Invalid gradient_accumulation_step value: {self.cfg.train.gradient_accumulation_step}. It should be positive." + ) + # TODO: check other values + + @staticmethod + def __count_parameters(model): + model_param = 0.0 + if isinstance(model, dict): + for key, value in model.items(): + model_param += sum(p.numel() for p in model[key].parameters()) + else: + model_param = sum(p.numel() for p in model.parameters()) + return model_param + + def _build_speaker_lut(self): + # combine speakers + if not os.path.exists(os.path.join(self.exp_dir, self.cfg.preprocess.spk2id)): + speakers = {} + else: + with open( + os.path.join(self.exp_dir, self.cfg.preprocess.spk2id), "r" + ) as speaker_file: + speakers = json.load(speaker_file) + for dataset in self.cfg.dataset: + speaker_lut_path = os.path.join( + self.cfg.preprocess.processed_dir, dataset, self.cfg.preprocess.spk2id + ) + with open(speaker_lut_path, "r") as speaker_lut_path: + singer_lut = json.load(speaker_lut_path) + for singer in singer_lut.keys(): + if singer not in speakers: + speakers[singer] = len(speakers) + with open( + os.path.join(self.exp_dir, self.cfg.preprocess.spk2id), "w" + ) as speaker_file: + json.dump(speakers, speaker_file, indent=4, ensure_ascii=False) + print( + "speakers have been dumped to {}".format( + os.path.join(self.exp_dir, self.cfg.preprocess.spk2id) + ) + ) + return speakers + + def _build_utt2spk_dict(self): + # combine speakers + utt2spk = {} + if not os.path.exists(os.path.join(self.exp_dir, self.cfg.preprocess.utt2spk)): + utt2spk = {} + else: + with open( + os.path.join(self.exp_dir, self.cfg.preprocess.utt2spk), "r" + ) as utt2spk_file: + for line in utt2spk_file.readlines(): + utt, spk = line.strip().split("\t") + utt2spk[utt] = spk + for dataset in self.cfg.dataset: + utt2spk_dict_path = os.path.join( + self.cfg.preprocess.processed_dir, dataset, self.cfg.preprocess.utt2spk + ) + with open(utt2spk_dict_path, "r") as utt2spk_dict: + for line in utt2spk_dict.readlines(): + utt, spk = line.strip().split("\t") + if utt not in utt2spk.keys(): + utt2spk[utt] = spk + with open( + os.path.join(self.exp_dir, self.cfg.preprocess.utt2spk), "w" + ) as utt2spk_file: + for utt, spk in utt2spk.items(): + utt2spk_file.write(utt + "\t" + spk + "\n") + print( + "utterance and speaker mapper have been dumped to {}".format( + os.path.join(self.exp_dir, self.cfg.preprocess.utt2spk) + ) + ) + return utt2spk + + def _save_phone_symbols_file_to_exp_path(self): + phone_symbols_file = os.path.join( + self.cfg.preprocess.processed_dir, + self.cfg.dataset[0], + self.cfg.preprocess.symbols_dict, + ) + phone_symbols_file_to_exp_path = os.path.join( + self.exp_dir, self.cfg.preprocess.symbols_dict + ) + shutil.copy(phone_symbols_file, phone_symbols_file_to_exp_path) + os.chmod(phone_symbols_file_to_exp_path, 0o666) + print( + "phone symbols been dumped to {}".format( + os.path.join(self.exp_dir, self.cfg.preprocess.symbols_dict) + ) + ) diff --git a/models/tts/fastspeech2/__init__.py b/models/tts/fastspeech2/__init__.py new file mode 100755 index 0000000..e69de29 diff --git a/models/tts/fastspeech2/fs2.py b/models/tts/fastspeech2/fs2.py new file mode 100644 index 0000000..61cab7e --- /dev/null +++ b/models/tts/fastspeech2/fs2.py @@ -0,0 +1,548 @@ +# Copyright (c) 2023 Amphion. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +# This code is modified from https://github.com/ming024/FastSpeech2/blob/master/model/fastspeech2.py +import torch +import torch.nn as nn +import numpy as np +import torch.nn.functional as F + +from modules.transformer.Models import Encoder, Decoder +from modules.transformer.Layers import PostNet +from collections import OrderedDict + +import os +import json + + +def get_mask_from_lengths(lengths, max_len=None): + device = lengths.device + batch_size = lengths.shape[0] + if max_len is None: + max_len = torch.max(lengths).item() + + ids = torch.arange(0, max_len).unsqueeze(0).expand(batch_size, -1).to(device) + mask = ids >= lengths.unsqueeze(1).expand(-1, max_len) + + return mask + + +def pad(input_ele, mel_max_length=None): + if mel_max_length: + max_len = mel_max_length + else: + max_len = max([input_ele[i].size(0) for i in range(len(input_ele))]) + + out_list = list() + for i, batch in enumerate(input_ele): + if len(batch.shape) == 1: + one_batch_padded = F.pad( + batch, (0, max_len - batch.size(0)), "constant", 0.0 + ) + elif len(batch.shape) == 2: + 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 + + +class VarianceAdaptor(nn.Module): + """Variance Adaptor""" + + def __init__(self, cfg): + super(VarianceAdaptor, self).__init__() + self.duration_predictor = VariancePredictor(cfg) + self.length_regulator = LengthRegulator() + self.pitch_predictor = VariancePredictor(cfg) + self.energy_predictor = VariancePredictor(cfg) + + # assign the pitch/energy feature level + if cfg.preprocess.use_frame_pitch: + self.pitch_feature_level = "frame_level" + self.pitch_dir = cfg.preprocess.pitch_dir + else: + self.pitch_feature_level = "phoneme_level" + self.pitch_dir = cfg.preprocess.phone_pitch_dir + + if cfg.preprocess.use_frame_energy: + self.energy_feature_level = "frame_level" + self.energy_dir = cfg.preprocess.energy_dir + else: + self.energy_feature_level = "phoneme_level" + self.energy_dir = cfg.preprocess.phone_energy_dir + + assert self.pitch_feature_level in ["phoneme_level", "frame_level"] + assert self.energy_feature_level in ["phoneme_level", "frame_level"] + + pitch_quantization = cfg.model.variance_embedding.pitch_quantization + energy_quantization = cfg.model.variance_embedding.energy_quantization + n_bins = cfg.model.variance_embedding.n_bins + assert pitch_quantization in ["linear", "log"] + assert energy_quantization in ["linear", "log"] + + with open( + os.path.join( + cfg.preprocess.processed_dir, + cfg.dataset[0], + self.energy_dir, + "statistics.json", + ) + ) as f: + stats = json.load(f) + stats = stats[cfg.dataset[0] + "_" + cfg.dataset[0]] + mean, std = ( + stats["voiced_positions"]["mean"], + stats["voiced_positions"]["std"], + ) + energy_min = (stats["total_positions"]["min"] - mean) / std + energy_max = (stats["total_positions"]["max"] - mean) / std + + with open( + os.path.join( + cfg.preprocess.processed_dir, + cfg.dataset[0], + self.pitch_dir, + "statistics.json", + ) + ) as f: + stats = json.load(f) + stats = stats[cfg.dataset[0] + "_" + cfg.dataset[0]] + mean, std = ( + stats["voiced_positions"]["mean"], + stats["voiced_positions"]["std"], + ) + pitch_min = (stats["total_positions"]["min"] - mean) / std + pitch_max = (stats["total_positions"]["max"] - mean) / std + + if pitch_quantization == "log": + self.pitch_bins = nn.Parameter( + torch.exp( + torch.linspace(np.log(pitch_min), np.log(pitch_max), n_bins - 1) + ), + requires_grad=False, + ) + else: + self.pitch_bins = nn.Parameter( + torch.linspace(pitch_min, pitch_max, n_bins - 1), + requires_grad=False, + ) + if energy_quantization == "log": + self.energy_bins = nn.Parameter( + torch.exp( + torch.linspace(np.log(energy_min), np.log(energy_max), n_bins - 1) + ), + requires_grad=False, + ) + else: + self.energy_bins = nn.Parameter( + torch.linspace(energy_min, energy_max, n_bins - 1), + requires_grad=False, + ) + + self.pitch_embedding = nn.Embedding( + n_bins, cfg.model.transformer.encoder_hidden + ) + self.energy_embedding = nn.Embedding( + n_bins, cfg.model.transformer.encoder_hidden + ) + + def get_pitch_embedding(self, x, target, mask, control): + prediction = self.pitch_predictor(x, mask) + if target is not None: + embedding = self.pitch_embedding(torch.bucketize(target, self.pitch_bins)) + else: + prediction = prediction * control + embedding = self.pitch_embedding( + torch.bucketize(prediction, self.pitch_bins) + ) + return prediction, embedding + + def get_energy_embedding(self, x, target, mask, control): + prediction = self.energy_predictor(x, mask) + if target is not None: + embedding = self.energy_embedding(torch.bucketize(target, self.energy_bins)) + else: + prediction = prediction * control + embedding = self.energy_embedding( + torch.bucketize(prediction, self.energy_bins) + ) + return prediction, embedding + + def forward( + self, + x, + src_mask, + mel_mask=None, + max_len=None, + pitch_target=None, + energy_target=None, + duration_target=None, + p_control=1.0, + e_control=1.0, + d_control=1.0, + ): + log_duration_prediction = self.duration_predictor(x, src_mask) + if self.pitch_feature_level == "phoneme_level": + pitch_prediction, pitch_embedding = self.get_pitch_embedding( + x, pitch_target, src_mask, p_control + ) + x = x + pitch_embedding + if self.energy_feature_level == "phoneme_level": + energy_prediction, energy_embedding = self.get_energy_embedding( + x, energy_target, src_mask, e_control + ) + x = x + energy_embedding + + if duration_target is not None: + x, mel_len = self.length_regulator(x, duration_target, max_len) + duration_rounded = duration_target + else: + duration_rounded = torch.clamp( + (torch.round(torch.exp(log_duration_prediction) - 1) * d_control), + min=0, + ) + x, mel_len = self.length_regulator(x, duration_rounded, max_len) + mel_mask = get_mask_from_lengths(mel_len) + + if self.pitch_feature_level == "frame_level": + pitch_prediction, pitch_embedding = self.get_pitch_embedding( + x, pitch_target, mel_mask, p_control + ) + x = x + pitch_embedding + if self.energy_feature_level == "frame_level": + energy_prediction, energy_embedding = self.get_energy_embedding( + x, energy_target, mel_mask, e_control + ) + x = x + energy_embedding + + return ( + x, + pitch_prediction, + energy_prediction, + log_duration_prediction, + duration_rounded, + mel_len, + mel_mask, + ) + + +class LengthRegulator(nn.Module): + """Length Regulator""" + + def __init__(self): + super(LengthRegulator, self).__init__() + + def LR(self, x, duration, max_len): + device = x.device + output = list() + mel_len = list() + for batch, expand_target in zip(x, duration): + expanded = self.expand(batch, expand_target) + output.append(expanded) + mel_len.append(expanded.shape[0]) + + if max_len is not None: + output = pad(output, max_len) + else: + output = pad(output) + + return output, torch.LongTensor(mel_len).to(device) + + def expand(self, batch, predicted): + out = list() + + for i, vec in enumerate(batch): + expand_size = predicted[i].item() + out.append(vec.expand(max(int(expand_size), 0), -1)) + out = torch.cat(out, 0) + + return out + + def forward(self, x, duration, max_len): + output, mel_len = self.LR(x, duration, max_len) + return output, mel_len + + +class VariancePredictor(nn.Module): + """Duration, Pitch and Energy Predictor""" + + def __init__(self, cfg): + super(VariancePredictor, self).__init__() + + self.input_size = cfg.model.transformer.encoder_hidden + self.filter_size = cfg.model.variance_predictor.filter_size + self.kernel = cfg.model.variance_predictor.kernel_size + self.conv_output_size = cfg.model.variance_predictor.filter_size + self.dropout = cfg.model.variance_predictor.dropout + + self.conv_layer = nn.Sequential( + OrderedDict( + [ + ( + "conv1d_1", + Conv( + self.input_size, + self.filter_size, + kernel_size=self.kernel, + padding=(self.kernel - 1) // 2, + ), + ), + ("relu_1", nn.ReLU()), + ("layer_norm_1", nn.LayerNorm(self.filter_size)), + ("dropout_1", nn.Dropout(self.dropout)), + ( + "conv1d_2", + Conv( + self.filter_size, + self.filter_size, + kernel_size=self.kernel, + padding=1, + ), + ), + ("relu_2", nn.ReLU()), + ("layer_norm_2", nn.LayerNorm(self.filter_size)), + ("dropout_2", nn.Dropout(self.dropout)), + ] + ) + ) + + self.linear_layer = nn.Linear(self.conv_output_size, 1) + + def forward(self, encoder_output, mask): + out = self.conv_layer(encoder_output) + out = self.linear_layer(out) + out = out.squeeze(-1) + + if mask is not None: + out = out.masked_fill(mask, 0.0) + + return out + + +class Conv(nn.Module): + """ + Convolution Module + """ + + def __init__( + self, + in_channels, + out_channels, + kernel_size=1, + stride=1, + padding=0, + dilation=1, + bias=True, + w_init="linear", + ): + """ + :param in_channels: dimension of input + :param out_channels: dimension of output + :param kernel_size: size of kernel + :param stride: size of stride + :param padding: size of padding + :param dilation: dilation rate + :param bias: boolean. if True, bias is included. + :param w_init: str. weight inits with xavier initialization. + """ + super(Conv, self).__init__() + + self.conv = nn.Conv1d( + in_channels, + out_channels, + kernel_size=kernel_size, + stride=stride, + padding=padding, + dilation=dilation, + bias=bias, + ) + + def forward(self, x): + x = x.contiguous().transpose(1, 2) + x = self.conv(x) + x = x.contiguous().transpose(1, 2) + + return x + + +class FastSpeech2(nn.Module): + def __init__(self, cfg) -> None: + super(FastSpeech2, self).__init__() + self.cfg = cfg + self.encoder = Encoder(cfg.model) + self.variance_adaptor = VarianceAdaptor(cfg) + self.decoder = Decoder(cfg.model) + self.mel_linear = nn.Linear( + cfg.model.transformer.decoder_hidden, + cfg.preprocess.n_mel, + ) + self.postnet = PostNet(n_mel_channels=cfg.preprocess.n_mel) + + self.speaker_emb = None + if cfg.train.multi_speaker_training: + with open( + os.path.join( + cfg.preprocess.processed_dir, cfg.dataset[0], "spk2id.json" + ), + "r", + ) as f: + n_speaker = len(json.load(f)) + self.speaker_emb = nn.Embedding( + n_speaker, + cfg.model.transformer.encoder_hidden, + ) + + def forward(self, data, p_control=1.0, e_control=1.0, d_control=1.0): + speakers = data["spk_id"] + texts = data["texts"] + src_lens = data["text_len"] + max_src_len = max(src_lens) + mel_lens = data["target_len"] if "target_len" in data else None + max_mel_len = max(mel_lens) if "target_len" in data else None + p_targets = data["pitch"] if "pitch" in data else None + e_targets = data["energy"] if "energy" in data else None + d_targets = data["durations"] if "durations" in data else None + src_masks = get_mask_from_lengths(src_lens, max_src_len) + mel_masks = ( + get_mask_from_lengths(mel_lens, max_mel_len) + if mel_lens is not None + else None + ) + + output = self.encoder(texts, src_masks) + + if self.speaker_emb is not None: + output = output + self.speaker_emb(speakers).unsqueeze(1).expand( + -1, max_src_len, -1 + ) + + ( + output, + p_predictions, + e_predictions, + log_d_predictions, + d_rounded, + mel_lens, + mel_masks, + ) = self.variance_adaptor( + output, + src_masks, + mel_masks, + max_mel_len, + p_targets, + e_targets, + d_targets, + p_control, + e_control, + d_control, + ) + + output, mel_masks = self.decoder(output, mel_masks) + output = self.mel_linear(output) + + postnet_output = self.postnet(output) + output + + return { + "output": output, + "postnet_output": postnet_output, + "p_predictions": p_predictions, + "e_predictions": e_predictions, + "log_d_predictions": log_d_predictions, + "d_rounded": d_rounded, + "src_masks": src_masks, + "mel_masks": mel_masks, + "src_lens": src_lens, + "mel_lens": mel_lens, + } + + +class FastSpeech2Loss(nn.Module): + """FastSpeech2 Loss""" + + def __init__(self, cfg): + super(FastSpeech2Loss, self).__init__() + if cfg.preprocess.use_frame_pitch: + self.pitch_feature_level = "frame_level" + else: + self.pitch_feature_level = "phoneme_level" + + if cfg.preprocess.use_frame_energy: + self.energy_feature_level = "frame_level" + else: + self.energy_feature_level = "phoneme_level" + + self.mse_loss = nn.MSELoss() + self.mae_loss = nn.L1Loss() + + def forward(self, data, predictions): + mel_targets = data["mel"] + pitch_targets = data["pitch"].float() + energy_targets = data["energy"].float() + duration_targets = data["durations"] + + mel_predictions = predictions["output"] + postnet_mel_predictions = predictions["postnet_output"] + pitch_predictions = predictions["p_predictions"] + energy_predictions = predictions["e_predictions"] + log_duration_predictions = predictions["log_d_predictions"] + src_masks = predictions["src_masks"] + mel_masks = predictions["mel_masks"] + + src_masks = ~src_masks + mel_masks = ~mel_masks + + log_duration_targets = torch.log(duration_targets.float() + 1) + mel_targets = mel_targets[:, : mel_masks.shape[1], :] + mel_masks = mel_masks[:, : mel_masks.shape[1]] + + log_duration_targets.requires_grad = False + pitch_targets.requires_grad = False + energy_targets.requires_grad = False + mel_targets.requires_grad = False + + if self.pitch_feature_level == "phoneme_level": + pitch_predictions = pitch_predictions.masked_select(src_masks) + pitch_targets = pitch_targets.masked_select(src_masks) + elif self.pitch_feature_level == "frame_level": + pitch_predictions = pitch_predictions.masked_select(mel_masks) + pitch_targets = pitch_targets.masked_select(mel_masks) + + if self.energy_feature_level == "phoneme_level": + energy_predictions = energy_predictions.masked_select(src_masks) + energy_targets = energy_targets.masked_select(src_masks) + if self.energy_feature_level == "frame_level": + energy_predictions = energy_predictions.masked_select(mel_masks) + energy_targets = energy_targets.masked_select(mel_masks) + + log_duration_predictions = log_duration_predictions.masked_select(src_masks) + log_duration_targets = log_duration_targets.masked_select(src_masks) + + mel_predictions = mel_predictions.masked_select(mel_masks.unsqueeze(-1)) + postnet_mel_predictions = postnet_mel_predictions.masked_select( + mel_masks.unsqueeze(-1) + ) + mel_targets = mel_targets.masked_select(mel_masks.unsqueeze(-1)) + + mel_loss = self.mae_loss(mel_predictions, mel_targets) + postnet_mel_loss = self.mae_loss(postnet_mel_predictions, mel_targets) + + pitch_loss = self.mse_loss(pitch_predictions, pitch_targets) + energy_loss = self.mse_loss(energy_predictions, energy_targets) + duration_loss = self.mse_loss(log_duration_predictions, log_duration_targets) + + total_loss = ( + mel_loss + postnet_mel_loss + duration_loss + pitch_loss + energy_loss + ) + + return { + "loss": total_loss, + "mel_loss": mel_loss, + "postnet_mel_loss": postnet_mel_loss, + "pitch_loss": pitch_loss, + "energy_loss": energy_loss, + "duration_loss": duration_loss, + } diff --git a/models/tts/fastspeech2/fs2_dataset.py b/models/tts/fastspeech2/fs2_dataset.py new file mode 100644 index 0000000..dae5b04 --- /dev/null +++ b/models/tts/fastspeech2/fs2_dataset.py @@ -0,0 +1,424 @@ +# Copyright (c) 2023 Amphion. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +import random +import torch +from torch.nn.utils.rnn import pad_sequence +from utils.data_utils import * +from models.base.base_dataset import ( + BaseOfflineCollator, + BaseOfflineDataset, + BaseTestDataset, + BaseTestCollator, +) +from text import text_to_sequence + + +class FS2Dataset(BaseOfflineDataset): + def __init__(self, cfg, dataset, is_valid=False): + BaseOfflineDataset.__init__(self, cfg, dataset, is_valid=is_valid) + self.batch_size = cfg.train.batch_size + cfg = cfg.preprocess + # utt2duration + self.utt2duration_path = {} + for utt_info in self.metadata: + dataset = utt_info["Dataset"] + uid = utt_info["Uid"] + utt = "{}_{}".format(dataset, uid) + + self.utt2duration_path[utt] = os.path.join( + cfg.processed_dir, + dataset, + cfg.duration_dir, + uid + ".npy", + ) + self.utt2dur = self.read_duration() + + if cfg.use_frame_energy: + self.frame_utt2energy, self.energy_statistic = load_energy( + self.metadata, + cfg.processed_dir, + cfg.energy_dir, + use_log_scale=cfg.use_log_scale_energy, + utt2spk=self.preprocess.utt2spk if cfg.use_spkid else None, + return_norm=True, + ) + elif cfg.use_phone_energy: + self.phone_utt2energy, self.energy_statistic = load_energy( + self.metadata, + cfg.processed_dir, + cfg.phone_energy_dir, + use_log_scale=cfg.use_log_scale_energy, + utt2spk=self.utt2spk if cfg.use_spkid else None, + return_norm=True, + ) + + if cfg.use_frame_pitch: + self.frame_utt2pitch, self.pitch_statistic = load_energy( + self.metadata, + cfg.processed_dir, + cfg.pitch_dir, + use_log_scale=cfg.energy_extract_mode, + utt2spk=self.utt2spk if cfg.use_spkid else None, + return_norm=True, + ) + + elif cfg.use_phone_pitch: + self.phone_utt2pitch, self.pitch_statistic = load_energy( + self.metadata, + cfg.processed_dir, + cfg.phone_pitch_dir, + use_log_scale=cfg.use_log_scale_pitch, + utt2spk=self.utt2spk if cfg.use_spkid else None, + return_norm=True, + ) + + # utt2lab + self.utt2lab_path = {} + for utt_info in self.metadata: + dataset = utt_info["Dataset"] + uid = utt_info["Uid"] + utt = "{}_{}".format(dataset, uid) + + self.utt2lab_path[utt] = os.path.join( + cfg.processed_dir, + dataset, + cfg.lab_dir, + uid + ".txt", + ) + + self.speaker_map = {} + if os.path.exists(os.path.join(cfg.processed_dir, "spk2id.json")): + with open( + os.path.exists(os.path.join(cfg.processed_dir, "spk2id.json")) + ) as f: + self.speaker_map = json.load(f) + + self.metadata = self.check_metadata() + + def __getitem__(self, index): + single_feature = BaseOfflineDataset.__getitem__(self, index) + + utt_info = self.metadata[index] + dataset = utt_info["Dataset"] + uid = utt_info["Uid"] + utt = "{}_{}".format(dataset, uid) + + duration = self.utt2dur[utt] + + # text + f = open(self.utt2lab_path[utt], "r") + phones = f.readlines()[0].strip() + f.close() + # todo: add cleaner(chenxi) + phones_ids = np.array(text_to_sequence(phones, ["english_cleaners"])) + text_len = len(phones_ids) + + if self.cfg.preprocess.use_frame_pitch: + pitch = self.frame_utt2pitch[utt] + elif self.cfg.preprocess.use_phone_pitch: + pitch = self.phone_utt2pitch[utt] + + if self.cfg.preprocess.use_frame_energy: + energy = self.frame_utt2energy[utt] + elif self.cfg.preprocess.use_phone_energy: + energy = self.phone_utt2energy[utt] + + # speaker + if len(self.speaker_map) > 0: + speaker_id = self.speaker_map[utt_info["Singer"]] + else: + speaker_id = 0 + + single_feature.update( + { + "durations": duration, + "texts": phones_ids, + "spk_id": speaker_id, + "text_len": text_len, + "pitch": pitch, + "energy": energy, + "uid": uid, + } + ) + return self.clip_if_too_long(single_feature) + + def read_duration(self): + # read duration + utt2dur = {} + for index in range(len(self.metadata)): + utt_info = self.metadata[index] + dataset = utt_info["Dataset"] + uid = utt_info["Uid"] + utt = "{}_{}".format(dataset, uid) + + if not os.path.exists(self.utt2mel_path[utt]) or not os.path.exists( + self.utt2duration_path[utt] + ): + continue + + mel = np.load(self.utt2mel_path[utt]).transpose(1, 0) + duration = np.load(self.utt2duration_path[utt]) + assert mel.shape[0] == sum( + duration + ), f"{utt}: mismatch length between mel {mel.shape[0]} and sum(duration) {sum(duration)}" + utt2dur[utt] = duration + return utt2dur + + def __len__(self): + return len(self.metadata) + + def random_select(self, feature_seq_len, max_seq_len, ending_ts=2812): + """ + ending_ts: to avoid invalid whisper features for over 30s audios + 2812 = 30 * 24000 // 256 + """ + ts = max(feature_seq_len - max_seq_len, 0) + ts = min(ts, ending_ts - max_seq_len) + + start = random.randint(0, ts) + end = start + max_seq_len + return start, end + + def clip_if_too_long(self, sample, max_seq_len=1000): + """ + sample : + { + 'spk_id': (1,), + 'target_len': int + 'mel': (seq_len, dim), + 'frame_pitch': (seq_len,) + 'frame_energy': (seq_len,) + 'content_vector_feat': (seq_len, dim) + } + """ + if sample["target_len"] <= max_seq_len: + return sample + + start, end = self.random_select(sample["target_len"], max_seq_len) + sample["target_len"] = end - start + + for k in sample.keys(): + if k not in ["spk_id", "target_len"]: + sample[k] = sample[k][start:end] + + return sample + + def check_metadata(self): + new_metadata = [] + for utt_info in self.metadata: + dataset = utt_info["Dataset"] + uid = utt_info["Uid"] + utt = "{}_{}".format(dataset, uid) + if not os.path.exists(self.utt2duration_path[utt]) or not os.path.exists( + self.utt2mel_path[utt] + ): + continue + else: + new_metadata.append(utt_info) + return new_metadata + + +class FS2Collator(BaseOfflineCollator): + """Zero-pads model inputs and targets based on number of frames per step""" + + def __init__(self, cfg): + BaseOfflineCollator.__init__(self, cfg) + self.sort = cfg.train.sort_sample + self.batch_size = cfg.train.batch_size + self.drop_last = cfg.train.drop_last + + def __call__(self, batch): + # mel: [b, T, n_mels] + # frame_pitch, frame_energy: [1, T] + # target_len: [1] + # spk_id: [b, 1] + # mask: [b, T, 1] + packed_batch_features = dict() + + for key in batch[0].keys(): + if key == "target_len": + packed_batch_features["target_len"] = torch.LongTensor( + [b["target_len"] for b in batch] + ) + masks = [ + torch.ones((b["target_len"], 1), dtype=torch.long) for b in batch + ] + packed_batch_features["mask"] = pad_sequence( + masks, batch_first=True, padding_value=0 + ) + elif key == "text_len": + packed_batch_features["text_len"] = torch.LongTensor( + [b["text_len"] for b in batch] + ) + masks = [ + torch.ones((b["text_len"], 1), dtype=torch.long) for b in batch + ] + packed_batch_features["text_mask"] = pad_sequence( + masks, batch_first=True, padding_value=0 + ) + elif key == "spk_id": + packed_batch_features["spk_id"] = torch.LongTensor( + [b["spk_id"] for b in batch] + ) + elif key == "uid": + packed_batch_features[key] = [b["uid"] for b in batch] + else: + values = [torch.from_numpy(b[key]) for b in batch] + packed_batch_features[key] = pad_sequence( + values, batch_first=True, padding_value=0 + ) + return packed_batch_features + + +class FS2TestDataset(BaseTestDataset): + def __init__(self, args, cfg, infer_type=None): + datasets = cfg.dataset + cfg = cfg.preprocess + is_bigdata = False + + assert len(datasets) >= 1 + if len(datasets) > 1: + datasets.sort() + bigdata_version = "_".join(datasets) + processed_data_dir = os.path.join(cfg.processed_dir, bigdata_version) + is_bigdata = True + else: + processed_data_dir = os.path.join(cfg.processed_dir, args.dataset) + + if args.test_list_file: + self.metafile_path = args.test_list_file + self.metadata = self.get_metadata() + else: + assert args.testing_set + source_metafile_path = os.path.join( + cfg.processed_dir, + args.dataset, + "{}.json".format(args.testing_set), + ) + with open(source_metafile_path, "r") as f: + self.metadata = json.load(f) + + self.cfg = cfg + self.datasets = datasets + self.data_root = processed_data_dir + self.is_bigdata = is_bigdata + self.source_dataset = args.dataset + + ######### Load source acoustic features ######### + if cfg.use_spkid: + spk2id_path = os.path.join(self.data_root, cfg.spk2id) + utt2sp_path = os.path.join(self.data_root, cfg.utt2spk) + self.spk2id, self.utt2spk = get_spk_map(spk2id_path, utt2sp_path, datasets) + + # utt2lab + self.utt2lab_path = {} + for utt_info in self.metadata: + dataset = utt_info["Dataset"] + uid = utt_info["Uid"] + utt = "{}_{}".format(dataset, uid) + self.utt2lab_path[utt] = os.path.join( + cfg.processed_dir, + dataset, + cfg.lab_dir, + uid + ".txt", + ) + + self.speaker_map = {} + if os.path.exists(os.path.join(cfg.processed_dir, "spk2id.json")): + with open( + os.path.exists(os.path.join(cfg.processed_dir, "spk2id.json")) + ) as f: + self.speaker_map = json.load(f) + + def __getitem__(self, index): + single_feature = {} + + utt_info = self.metadata[index] + dataset = utt_info["Dataset"] + uid = utt_info["Uid"] + utt = "{}_{}".format(dataset, uid) + + # text + f = open(self.utt2lab_path[utt], "r") + phones = f.readlines()[0].strip() + f.close() + + phones_ids = np.array(text_to_sequence(phones, self.cfg.text_cleaners)) + text_len = len(phones_ids) + + # speaker + if len(self.speaker_map) > 0: + speaker_id = self.speaker_map[utt_info["Singer"]] + else: + speaker_id = 0 + + single_feature.update( + { + "texts": phones_ids, + "spk_id": speaker_id, + "text_len": text_len, + } + ) + + return single_feature + + def __len__(self): + return len(self.metadata) + + def get_metadata(self): + with open(self.metafile_path, "r", encoding="utf-8") as f: + metadata = json.load(f) + + return metadata + + +class FS2TestCollator(BaseTestCollator): + """Zero-pads model inputs and targets based on number of frames per step""" + + def __init__(self, cfg): + self.cfg = cfg + + def __call__(self, batch): + packed_batch_features = dict() + + # mel: [b, T, n_mels] + # frame_pitch, frame_energy: [1, T] + # target_len: [1] + # spk_id: [b, 1] + # mask: [b, T, 1] + + for key in batch[0].keys(): + if key == "target_len": + packed_batch_features["target_len"] = torch.LongTensor( + [b["target_len"] for b in batch] + ) + masks = [ + torch.ones((b["target_len"], 1), dtype=torch.long) for b in batch + ] + packed_batch_features["mask"] = pad_sequence( + masks, batch_first=True, padding_value=0 + ) + elif key == "text_len": + packed_batch_features["text_len"] = torch.LongTensor( + [b["text_len"] for b in batch] + ) + masks = [ + torch.ones((b["text_len"], 1), dtype=torch.long) for b in batch + ] + packed_batch_features["text_mask"] = pad_sequence( + masks, batch_first=True, padding_value=0 + ) + elif key == "spk_id": + packed_batch_features["spk_id"] = torch.LongTensor( + [b["spk_id"] for b in batch] + ) + else: + values = [torch.from_numpy(b[key]) for b in batch] + packed_batch_features[key] = pad_sequence( + values, batch_first=True, padding_value=0 + ) + + return packed_batch_features diff --git a/models/tts/fastspeech2/fs2_inference.py b/models/tts/fastspeech2/fs2_inference.py new file mode 100644 index 0000000..5a03afd --- /dev/null +++ b/models/tts/fastspeech2/fs2_inference.py @@ -0,0 +1,193 @@ +# Copyright (c) 2023 Amphion. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +import os +import torch +from tqdm import tqdm +from collections import OrderedDict + +from models.tts.base.tts_inferece import TTSInference +from models.tts.fastspeech2.fs2_dataset import FS2TestDataset, FS2TestCollator +from utils.util import load_config +from utils.io import save_audio +from models.tts.fastspeech2.fs2 import FastSpeech2 +from models.vocoders.vocoder_inference import synthesis +from pathlib import Path +from processors.phone_extractor import phoneExtractor +from text.text_token_collation import phoneIDCollation +import numpy as np +import json + + +class FastSpeech2Inference(TTSInference): + def __init__(self, args, cfg): + TTSInference.__init__(self, args, cfg) + self.args = args + self.cfg = cfg + self.infer_type = args.mode + + def _build_model(self): + self.model = FastSpeech2(self.cfg) + return self.model + + def load_model(self, state_dict): + raw_dict = state_dict["model"] + clean_dict = OrderedDict() + for k, v in raw_dict.items(): + if k.startswith("module."): + clean_dict[k[7:]] = v + else: + clean_dict[k] = v + + self.model.load_state_dict(clean_dict) + + def _build_test_dataset(self): + return FS2TestDataset, FS2TestCollator + + @staticmethod + def _parse_vocoder(vocoder_dir): + r"""Parse vocoder config""" + vocoder_dir = os.path.abspath(vocoder_dir) + ckpt_list = [ckpt for ckpt in Path(vocoder_dir).glob("*.pt")] + # last step (different from the base *int(x.stem)*) + ckpt_list.sort( + key=lambda x: int(x.stem.split("_")[-2].split("-")[-1]), reverse=True + ) + ckpt_path = str(ckpt_list[0]) + vocoder_cfg = load_config( + os.path.join(vocoder_dir, "args.json"), lowercase=True + ) + return vocoder_cfg, ckpt_path + + @torch.inference_mode() + def inference_for_batches(self): + y_pred = [] + for i, batch in tqdm(enumerate(self.test_dataloader)): + y_pred, mel_lens, _ = self._inference_each_batch(batch) + y_ls = y_pred.chunk(self.test_batch_size) + tgt_ls = mel_lens.chunk(self.test_batch_size) + j = 0 + for it, l in zip(y_ls, tgt_ls): + l = l.item() + it = it.squeeze(0)[:l].detach().cpu() + + uid = self.test_dataset.metadata[i * self.test_batch_size + j]["Uid"] + torch.save(it, os.path.join(self.args.output_dir, f"{uid}.pt")) + j += 1 + + vocoder_cfg, vocoder_ckpt = self._parse_vocoder(self.args.vocoder_dir) + res = synthesis( + cfg=vocoder_cfg, + vocoder_weight_file=vocoder_ckpt, + n_samples=None, + pred=[ + torch.load( + os.path.join(self.args.output_dir, "{}.pt".format(item["Uid"])) + ).numpy() + for item in self.test_dataset.metadata + ], + ) + for it, wav in zip(self.test_dataset.metadata, res): + uid = it["Uid"] + save_audio( + os.path.join(self.args.output_dir, f"{uid}.wav"), + wav.numpy(), + self.cfg.preprocess.sample_rate, + add_silence=True, + turn_up=True, + ) + os.remove(os.path.join(self.args.output_dir, f"{uid}.pt")) + + @torch.inference_mode() + def _inference_each_batch(self, batch_data): + device = self.accelerator.device + control_values = ( + self.args.pitch_control, + self.args.energy_control, + self.args.duration_control, + ) + for k, v in batch_data.items(): + batch_data[k] = v.to(device) + + pitch_control, energy_control, duration_control = control_values + + output = self.model( + batch_data, + p_control=pitch_control, + e_control=energy_control, + d_control=duration_control, + ) + pred_res = output["postnet_output"] + mel_lens = output["mel_lens"].cpu() + return pred_res, mel_lens, 0 + + def inference_for_single_utterance(self): + text = self.args.text + control_values = ( + self.args.pitch_control, + self.args.energy_control, + self.args.duration_control, + ) + pitch_control, energy_control, duration_control = control_values + + # get phone symbol file + phone_symbol_file = None + if self.cfg.preprocess.phone_extractor != "lexicon": + phone_symbol_file = os.path.join( + self.exp_dir, self.cfg.preprocess.symbols_dict + ) + assert os.path.exists(phone_symbol_file) + # convert text to phone sequence + phone_extractor = phoneExtractor(self.cfg) + + phone_seq = phone_extractor.extract_phone(text) # phone_seq: list + # convert phone sequence to phone id sequence + phon_id_collator = phoneIDCollation( + self.cfg, symbols_dict_file=phone_symbol_file + ) + phone_seq = ["{"] + phone_seq + ["}"] + phone_id_seq = phon_id_collator.get_phone_id_sequence(self.cfg, phone_seq) + + # convert phone sequence to phone id sequence + phone_id_seq = np.array(phone_id_seq) + phone_id_seq = torch.from_numpy(phone_id_seq) + + # get speaker id if multi-speaker training and use speaker id + speaker_id = None + if self.cfg.preprocess.use_spkid and self.cfg.train.multi_speaker_training: + spk2id_file = os.path.join(self.exp_dir, self.cfg.preprocess.spk2id) + with open(spk2id_file, "r") as f: + spk2id = json.load(f) + speaker_id = spk2id[self.args.speaker_name] + speaker_id = torch.from_numpy(np.array([speaker_id], dtype=np.int32)) + else: + speaker_id = torch.Tensor(0).view(-1) + + with torch.no_grad(): + x_tst = phone_id_seq.to(self.device).unsqueeze(0) + x_tst_lengths = torch.LongTensor([phone_id_seq.size(0)]).to(self.device) + if speaker_id is not None: + speaker_id = speaker_id.to(self.device) + + data = {} + data["texts"] = x_tst + data["text_len"] = x_tst_lengths + data["spk_id"] = speaker_id + + output = self.model( + data, + p_control=pitch_control, + e_control=energy_control, + d_control=duration_control, + ) + pred_res = output["postnet_output"] + vocoder_cfg, vocoder_ckpt = self._parse_vocoder(self.args.vocoder_dir) + audio = synthesis( + cfg=vocoder_cfg, + vocoder_weight_file=vocoder_ckpt, + n_samples=None, + pred=pred_res, + ) + return audio[0] diff --git a/models/tts/fastspeech2/fs2_trainer.py b/models/tts/fastspeech2/fs2_trainer.py new file mode 100644 index 0000000..f0bc120 --- /dev/null +++ b/models/tts/fastspeech2/fs2_trainer.py @@ -0,0 +1,155 @@ +# Copyright (c) 2023 Amphion. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +import torch +import torch.nn as nn +from tqdm import tqdm +from models.tts.base import TTSTrainer +from models.tts.fastspeech2.fs2 import FastSpeech2, FastSpeech2Loss +from models.tts.fastspeech2.fs2_dataset import FS2Dataset, FS2Collator +from optimizer.optimizers import NoamLR + + +class FastSpeech2Trainer(TTSTrainer): + def __init__(self, args, cfg): + TTSTrainer.__init__(self, args, cfg) + self.cfg = cfg + + def _build_dataset(self): + return FS2Dataset, FS2Collator + + def __build_scheduler(self): + return NoamLR(self.optimizer, **self.cfg.train.lr_scheduler) + + def _write_summary(self, losses, stats): + for key, value in losses.items(): + self.sw.add_scalar("train/" + key, value, self.step) + lr = self.optimizer.state_dict()["param_groups"][0]["lr"] + self.sw.add_scalar("learning_rate", lr, self.step) + + def _write_valid_summary(self, losses, stats): + for key, value in losses.items(): + self.sw.add_scalar("val/" + key, value, self.step) + + def _build_criterion(self): + return FastSpeech2Loss(self.cfg) + + def get_state_dict(self): + state_dict = { + "model": self.model.state_dict(), + "optimizer": self.optimizer.state_dict(), + "scheduler": self.scheduler.state_dict(), + "step": self.step, + "epoch": self.epoch, + "batch_size": self.cfg.train.batch_size, + } + return state_dict + + def _build_optimizer(self): + optimizer = torch.optim.Adam(self.model.parameters(), **self.cfg.train.adam) + return optimizer + + def _build_scheduler(self): + scheduler = NoamLR(self.optimizer, **self.cfg.train.lr_scheduler) + return scheduler + + def _build_model(self): + self.model = FastSpeech2(self.cfg) + return self.model + + def _train_epoch(self): + r"""Training epoch. Should return average loss of a batch (sample) over + one epoch. See ``train_loop`` for usage. + """ + self.model.train() + epoch_sum_loss: float = 0.0 + epoch_step: int = 0 + epoch_losses: dict = {} + for batch in tqdm( + self.train_dataloader, + desc=f"Training Epoch {self.epoch}", + unit="batch", + colour="GREEN", + leave=False, + dynamic_ncols=True, + smoothing=0.04, + disable=not self.accelerator.is_main_process, + ): + # Do training step and BP + with self.accelerator.accumulate(self.model): + loss, train_losses = self._train_step(batch) + self.accelerator.backward(loss) + grad_clip_thresh = self.cfg.train.grad_clip_thresh + nn.utils.clip_grad_norm_(self.model.parameters(), grad_clip_thresh) + self.optimizer.step() + self.scheduler.step() + self.optimizer.zero_grad() + self.batch_count += 1 + + # Update info for each step + if self.batch_count % self.cfg.train.gradient_accumulation_step == 0: + epoch_sum_loss += loss + for key, value in train_losses.items(): + if key not in epoch_losses.keys(): + epoch_losses[key] = value + else: + epoch_losses[key] += value + + self.accelerator.log( + { + "Step/Train Loss": loss, + "Step/Learning Rate": self.optimizer.param_groups[0]["lr"], + }, + step=self.step, + ) + self.step += 1 + epoch_step += 1 + + self.accelerator.wait_for_everyone() + + epoch_sum_loss = ( + epoch_sum_loss + / len(self.train_dataloader) + * self.cfg.train.gradient_accumulation_step + ) + + for key in epoch_losses.keys(): + epoch_losses[key] = ( + epoch_losses[key] + / len(self.train_dataloader) + * self.cfg.train.gradient_accumulation_step + ) + return epoch_sum_loss, epoch_losses + + def _train_step(self, data): + train_losses = {} + total_loss = 0 + train_stats = {} + + preds = self.model(data) + + train_losses = self.criterion(data, preds) + + total_loss = train_losses["loss"] + for key, value in train_losses.items(): + train_losses[key] = value.item() + + return total_loss, train_losses + + @torch.no_grad() + def _valid_step(self, data): + valid_loss = {} + total_valid_loss = 0 + valid_stats = {} + + preds = self.model(data) + + valid_losses = self.criterion(data, preds) + + total_valid_loss = valid_losses["loss"] + for key, value in valid_losses.items(): + valid_losses[key] = value.item() + + return total_valid_loss, valid_losses, valid_stats diff --git a/models/tts/jets/__init__.py b/models/tts/jets/__init__.py new file mode 100755 index 0000000..e69de29 diff --git a/models/tts/jets/alignments.py b/models/tts/jets/alignments.py new file mode 100644 index 0000000..a96a500 --- /dev/null +++ b/models/tts/jets/alignments.py @@ -0,0 +1,487 @@ +# Copyright (c) 2024 Amphion. +# +# This code is modified from https://github.com/imdanboy/jets/blob/main/espnet2/gan_tts/jets/alignments.py +# Licensed under Apache License 2.0 + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +from typing import Tuple +from numba import jit +from scipy.stats import betabinom + + +class AlignmentModule(nn.Module): + """Alignment Learning Framework proposed for parallel TTS models in: + + https://arxiv.org/abs/2108.10447 + + """ + + def __init__(self, adim, odim, cache_prior=True): + """Initialize AlignmentModule. + + Args: + adim (int): Dimension of attention. + odim (int): Dimension of feats. + cache_prior (bool): Whether to cache beta-binomial prior. + + """ + super().__init__() + self.cache_prior = cache_prior + self._cache = {} + + self.t_conv1 = nn.Conv1d(adim, adim, kernel_size=3, padding=1) + self.t_conv2 = nn.Conv1d(adim, adim, kernel_size=1, padding=0) + + self.f_conv1 = nn.Conv1d(odim, adim, kernel_size=3, padding=1) + self.f_conv2 = nn.Conv1d(adim, adim, kernel_size=3, padding=1) + self.f_conv3 = nn.Conv1d(adim, adim, kernel_size=1, padding=0) + + def forward(self, text, feats, text_lengths, feats_lengths, x_masks=None): + """Calculate alignment loss. + + Args: + text (Tensor): Batched text embedding (B, T_text, adim). + feats (Tensor): Batched acoustic feature (B, T_feats, odim). + text_lengths (Tensor): Text length tensor (B,). + feats_lengths (Tensor): Feature length tensor (B,). + x_masks (Tensor): Mask tensor (B, T_text). + + Returns: + Tensor: Log probability of attention matrix (B, T_feats, T_text). + + """ + text = text.transpose(1, 2) + text = F.relu(self.t_conv1(text)) + text = self.t_conv2(text) + text = text.transpose(1, 2) + + feats = feats.transpose(1, 2) + feats = F.relu(self.f_conv1(feats)) + feats = F.relu(self.f_conv2(feats)) + feats = self.f_conv3(feats) + feats = feats.transpose(1, 2) + + dist = feats.unsqueeze(2) - text.unsqueeze(1) + dist = torch.norm(dist, p=2, dim=3) + score = -dist + + if x_masks is not None: + x_masks = x_masks.unsqueeze(-2) + score = score.masked_fill(x_masks, -np.inf) + + log_p_attn = F.log_softmax(score, dim=-1) + + # add beta-binomial prior + bb_prior = self._generate_prior( + text_lengths, + feats_lengths, + ).to(dtype=log_p_attn.dtype, device=log_p_attn.device) + log_p_attn = log_p_attn + bb_prior + + return log_p_attn + + def _generate_prior(self, text_lengths, feats_lengths, w=1) -> torch.Tensor: + """Generate alignment prior formulated as beta-binomial distribution + + Args: + text_lengths (Tensor): Batch of the lengths of each input (B,). + feats_lengths (Tensor): Batch of the lengths of each target (B,). + w (float): Scaling factor; lower -> wider the width. + + Returns: + Tensor: Batched 2d static prior matrix (B, T_feats, T_text). + + """ + B = len(text_lengths) + T_text = text_lengths.max() + T_feats = feats_lengths.max() + + bb_prior = torch.full((B, T_feats, T_text), fill_value=-np.inf) + for bidx in range(B): + T = feats_lengths[bidx].item() + N = text_lengths[bidx].item() + + key = str(T) + "," + str(N) + if self.cache_prior and key in self._cache: + prob = self._cache[key] + else: + alpha = w * np.arange(1, T + 1, dtype=float) # (T,) + beta = w * np.array([T - t + 1 for t in alpha]) + k = np.arange(N) + batched_k = k[..., None] # (N,1) + prob = betabinom.logpmf(batched_k, N, alpha, beta) # (N,T) + + # store cache + if self.cache_prior and key not in self._cache: + self._cache[key] = prob + + prob = torch.from_numpy(prob).transpose(0, 1) # -> (T,N) + bb_prior[bidx, :T, :N] = prob + + return bb_prior + + +@jit(nopython=True) +def _monotonic_alignment_search(log_p_attn): + # https://arxiv.org/abs/2005.11129 + T_mel = log_p_attn.shape[0] + T_inp = log_p_attn.shape[1] + Q = np.full((T_inp, T_mel), fill_value=-np.inf) + + log_prob = log_p_attn.transpose(1, 0) # -> (T_inp,T_mel) + # 1. Q <- init first row for all j + for j in range(T_mel): + Q[0, j] = log_prob[0, : j + 1].sum() + + # 2. + for j in range(1, T_mel): + for i in range(1, min(j + 1, T_inp)): + Q[i, j] = max(Q[i - 1, j - 1], Q[i, j - 1]) + log_prob[i, j] + + # 3. + A = np.full((T_mel,), fill_value=T_inp - 1) + for j in range(T_mel - 2, -1, -1): # T_mel-2, ..., 0 + # 'i' in {A[j+1]-1, A[j+1]} + i_a = A[j + 1] - 1 + i_b = A[j + 1] + if i_b == 0: + argmax_i = 0 + elif Q[i_a, j] >= Q[i_b, j]: + argmax_i = i_a + else: + argmax_i = i_b + A[j] = argmax_i + return A + + +def viterbi_decode(log_p_attn, text_lengths, feats_lengths): + """Extract duration from an attention probability matrix + + Args: + log_p_attn (Tensor): Batched log probability of attention + matrix (B, T_feats, T_text). + text_lengths (Tensor): Text length tensor (B,). + feats_legnths (Tensor): Feature length tensor (B,). + + Returns: + Tensor: Batched token duration extracted from `log_p_attn` (B, T_text). + Tensor: Binarization loss tensor (). + + """ + B = log_p_attn.size(0) + T_text = log_p_attn.size(2) + device = log_p_attn.device + + bin_loss = 0 + ds = torch.zeros((B, T_text), device=device) + for b in range(B): + cur_log_p_attn = log_p_attn[b, : feats_lengths[b], : text_lengths[b]] + viterbi = _monotonic_alignment_search(cur_log_p_attn.detach().cpu().numpy()) + _ds = np.bincount(viterbi) + ds[b, : len(_ds)] = torch.from_numpy(_ds).to(device) + + t_idx = torch.arange(feats_lengths[b]) + bin_loss = bin_loss - cur_log_p_attn[t_idx, viterbi].mean() + bin_loss = bin_loss / B + return ds, bin_loss + + +@jit(nopython=True) +def _average_by_duration(ds, xs, text_lengths, feats_lengths): + B = ds.shape[0] + xs_avg = np.zeros_like(ds) + ds = ds.astype(np.int32) + for b in range(B): + t_text = text_lengths[b] + t_feats = feats_lengths[b] + d = ds[b, :t_text] + d_cumsum = d.cumsum() + d_cumsum = [0] + list(d_cumsum) + x = xs[b, :t_feats] + for n, (start, end) in enumerate(zip(d_cumsum[:-1], d_cumsum[1:])): + if len(x[start:end]) != 0: + xs_avg[b, n] = x[start:end].mean() + else: + xs_avg[b, n] = 0 + return xs_avg + + +def average_by_duration(ds, xs, text_lengths, feats_lengths): + """Average frame-level features into token-level according to durations + + Args: + ds (Tensor): Batched token duration (B, T_text). + xs (Tensor): Batched feature sequences to be averaged (B, T_feats). + text_lengths (Tensor): Text length tensor (B,). + feats_lengths (Tensor): Feature length tensor (B,). + + Returns: + Tensor: Batched feature averaged according to the token duration (B, T_text). + + """ + device = ds.device + args = [ds, xs, text_lengths, feats_lengths] + args = [arg.detach().cpu().numpy() for arg in args] + xs_avg = _average_by_duration(*args) + xs_avg = torch.from_numpy(xs_avg).to(device) + return xs_avg + + +def make_pad_mask(lengths, xs=None, length_dim=-1, maxlen=None): + """Make mask tensor containing indices of padded part. + + Args: + lengths (LongTensor or List): Batch of lengths (B,). + xs (Tensor, optional): The reference tensor. + If set, masks will be the same shape as this tensor. + length_dim (int, optional): Dimension indicator of the above tensor. + See the example. + + Returns: + Tensor: Mask tensor containing indices of padded part. + dtype=torch.uint8 in PyTorch 1.2- + dtype=torch.bool in PyTorch 1.2+ (including 1.2) + + Examples: + With only lengths. + + >>> lengths = [5, 3, 2] + >>> make_pad_mask(lengths) + masks = [[0, 0, 0, 0 ,0], + [0, 0, 0, 1, 1], + [0, 0, 1, 1, 1]] + + With the reference tensor. + + >>> xs = torch.zeros((3, 2, 4)) + >>> make_pad_mask(lengths, xs) + tensor([[[0, 0, 0, 0], + [0, 0, 0, 0]], + [[0, 0, 0, 1], + [0, 0, 0, 1]], + [[0, 0, 1, 1], + [0, 0, 1, 1]]], dtype=torch.uint8) + >>> xs = torch.zeros((3, 2, 6)) + >>> make_pad_mask(lengths, xs) + tensor([[[0, 0, 0, 0, 0, 1], + [0, 0, 0, 0, 0, 1]], + [[0, 0, 0, 1, 1, 1], + [0, 0, 0, 1, 1, 1]], + [[0, 0, 1, 1, 1, 1], + [0, 0, 1, 1, 1, 1]]], dtype=torch.uint8) + + With the reference tensor and dimension indicator. + + >>> xs = torch.zeros((3, 6, 6)) + >>> make_pad_mask(lengths, xs, 1) + tensor([[[0, 0, 0, 0, 0, 0], + [0, 0, 0, 0, 0, 0], + [0, 0, 0, 0, 0, 0], + [0, 0, 0, 0, 0, 0], + [0, 0, 0, 0, 0, 0], + [1, 1, 1, 1, 1, 1]], + [[0, 0, 0, 0, 0, 0], + [0, 0, 0, 0, 0, 0], + [0, 0, 0, 0, 0, 0], + [1, 1, 1, 1, 1, 1], + [1, 1, 1, 1, 1, 1], + [1, 1, 1, 1, 1, 1]], + [[0, 0, 0, 0, 0, 0], + [0, 0, 0, 0, 0, 0], + [1, 1, 1, 1, 1, 1], + [1, 1, 1, 1, 1, 1], + [1, 1, 1, 1, 1, 1], + [1, 1, 1, 1, 1, 1]]], dtype=torch.uint8) + >>> make_pad_mask(lengths, xs, 2) + tensor([[[0, 0, 0, 0, 0, 1], + [0, 0, 0, 0, 0, 1], + [0, 0, 0, 0, 0, 1], + [0, 0, 0, 0, 0, 1], + [0, 0, 0, 0, 0, 1], + [0, 0, 0, 0, 0, 1]], + [[0, 0, 0, 1, 1, 1], + [0, 0, 0, 1, 1, 1], + [0, 0, 0, 1, 1, 1], + [0, 0, 0, 1, 1, 1], + [0, 0, 0, 1, 1, 1], + [0, 0, 0, 1, 1, 1]], + [[0, 0, 1, 1, 1, 1], + [0, 0, 1, 1, 1, 1], + [0, 0, 1, 1, 1, 1], + [0, 0, 1, 1, 1, 1], + [0, 0, 1, 1, 1, 1], + [0, 0, 1, 1, 1, 1]]], dtype=torch.uint8) + + """ + if length_dim == 0: + raise ValueError("length_dim cannot be 0: {}".format(length_dim)) + + if not isinstance(lengths, list): + lengths = lengths.tolist() + bs = int(len(lengths)) + if maxlen is None: + if xs is None: + maxlen = int(max(lengths)) + else: + maxlen = xs.size(length_dim) + else: + assert xs is None + assert maxlen >= int(max(lengths)) + + seq_range = torch.arange(0, maxlen, dtype=torch.int64) + seq_range_expand = seq_range.unsqueeze(0).expand(bs, maxlen) + seq_length_expand = seq_range_expand.new(lengths).unsqueeze(-1) + mask = seq_range_expand >= seq_length_expand + + if xs is not None: + assert xs.size(0) == bs, (xs.size(0), bs) + + if length_dim < 0: + length_dim = xs.dim() + length_dim + # ind = (:, None, ..., None, :, , None, ..., None) + ind = tuple( + slice(None) if i in (0, length_dim) else None for i in range(xs.dim()) + ) + mask = mask[ind].expand_as(xs).to(xs.device) + return mask + + +def make_non_pad_mask(lengths, xs=None, length_dim=-1): + """Make mask tensor containing indices of non-padded part. + + Args: + lengths (LongTensor or List): Batch of lengths (B,). + xs (Tensor, optional): The reference tensor. + If set, masks will be the same shape as this tensor. + length_dim (int, optional): Dimension indicator of the above tensor. + See the example. + + Returns: + ByteTensor: mask tensor containing indices of padded part. + dtype=torch.uint8 in PyTorch 1.2- + dtype=torch.bool in PyTorch 1.2+ (including 1.2) + + Examples: + With only lengths. + + >>> lengths = [5, 3, 2] + >>> make_non_pad_mask(lengths) + masks = [[1, 1, 1, 1 ,1], + [1, 1, 1, 0, 0], + [1, 1, 0, 0, 0]] + + With the reference tensor. + + >>> xs = torch.zeros((3, 2, 4)) + >>> make_non_pad_mask(lengths, xs) + tensor([[[1, 1, 1, 1], + [1, 1, 1, 1]], + [[1, 1, 1, 0], + [1, 1, 1, 0]], + [[1, 1, 0, 0], + [1, 1, 0, 0]]], dtype=torch.uint8) + >>> xs = torch.zeros((3, 2, 6)) + >>> make_non_pad_mask(lengths, xs) + tensor([[[1, 1, 1, 1, 1, 0], + [1, 1, 1, 1, 1, 0]], + [[1, 1, 1, 0, 0, 0], + [1, 1, 1, 0, 0, 0]], + [[1, 1, 0, 0, 0, 0], + [1, 1, 0, 0, 0, 0]]], dtype=torch.uint8) + + With the reference tensor and dimension indicator. + + >>> xs = torch.zeros((3, 6, 6)) + >>> make_non_pad_mask(lengths, xs, 1) + tensor([[[1, 1, 1, 1, 1, 1], + [1, 1, 1, 1, 1, 1], + [1, 1, 1, 1, 1, 1], + [1, 1, 1, 1, 1, 1], + [1, 1, 1, 1, 1, 1], + [0, 0, 0, 0, 0, 0]], + [[1, 1, 1, 1, 1, 1], + [1, 1, 1, 1, 1, 1], + [1, 1, 1, 1, 1, 1], + [0, 0, 0, 0, 0, 0], + [0, 0, 0, 0, 0, 0], + [0, 0, 0, 0, 0, 0]], + [[1, 1, 1, 1, 1, 1], + [1, 1, 1, 1, 1, 1], + [0, 0, 0, 0, 0, 0], + [0, 0, 0, 0, 0, 0], + [0, 0, 0, 0, 0, 0], + [0, 0, 0, 0, 0, 0]]], dtype=torch.uint8) + >>> make_non_pad_mask(lengths, xs, 2) + tensor([[[1, 1, 1, 1, 1, 0], + [1, 1, 1, 1, 1, 0], + [1, 1, 1, 1, 1, 0], + [1, 1, 1, 1, 1, 0], + [1, 1, 1, 1, 1, 0], + [1, 1, 1, 1, 1, 0]], + [[1, 1, 1, 0, 0, 0], + [1, 1, 1, 0, 0, 0], + [1, 1, 1, 0, 0, 0], + [1, 1, 1, 0, 0, 0], + [1, 1, 1, 0, 0, 0], + [1, 1, 1, 0, 0, 0]], + [[1, 1, 0, 0, 0, 0], + [1, 1, 0, 0, 0, 0], + [1, 1, 0, 0, 0, 0], + [1, 1, 0, 0, 0, 0], + [1, 1, 0, 0, 0, 0], + [1, 1, 0, 0, 0, 0]]], dtype=torch.uint8) + + """ + return ~make_pad_mask(lengths, xs, length_dim) + + +def get_random_segments( + x: torch.Tensor, + x_lengths: torch.Tensor, + segment_size: int, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Get random segments. + + Args: + x (Tensor): Input tensor (B, C, T). + x_lengths (Tensor): Length tensor (B,). + segment_size (int): Segment size. + + Returns: + Tensor: Segmented tensor (B, C, segment_size). + Tensor: Start index tensor (B,). + + """ + b, c, t = x.size() + max_start_idx = x_lengths - segment_size + start_idxs = (torch.rand([b]).to(x.device) * max_start_idx).to( + dtype=torch.long, + ) + segments = get_segments(x, start_idxs, segment_size) + return segments, start_idxs + + +def get_segments( + x: torch.Tensor, + start_idxs: torch.Tensor, + segment_size: int, +) -> torch.Tensor: + """Get segments. + + Args: + x (Tensor): Input tensor (B, C, T). + start_idxs (Tensor): Start index tensor (B,). + segment_size (int): Segment size. + + Returns: + Tensor: Segmented tensor (B, C, segment_size). + + """ + b, c, t = x.size() + segments = x.new_zeros(b, c, segment_size) + for i, start_idx in enumerate(start_idxs): + segments[i] = x[i, :, start_idx : start_idx + segment_size] + return segments diff --git a/models/tts/jets/jets.py b/models/tts/jets/jets.py new file mode 100644 index 0000000..3940b37 --- /dev/null +++ b/models/tts/jets/jets.py @@ -0,0 +1,621 @@ +# Copyright (c) 2024 Amphion. +# +# This code is modified from https://github.com/imdanboy/jets/blob/main/espnet2/gan_tts/jets/generator.py +# Licensed under Apache License 2.0 + +import torch +import torch.nn as nn +import numpy as np +import torch.nn.functional as F + +from modules.transformer.Models import Encoder, Decoder +from modules.transformer.Layers import PostNet +from collections import OrderedDict +from models.tts.jets.alignments import ( + AlignmentModule, + viterbi_decode, + average_by_duration, + make_pad_mask, + make_non_pad_mask, + get_random_segments, +) +from models.tts.jets.length_regulator import GaussianUpsampling +from models.vocoders.gan.generator.hifigan import HiFiGAN +import os +import json + +from utils.util import load_config + + +def get_mask_from_lengths(lengths, max_len=None): + device = lengths.device + batch_size = lengths.shape[0] + if max_len is None: + max_len = torch.max(lengths).item() + + ids = torch.arange(0, max_len).unsqueeze(0).expand(batch_size, -1).to(device) + mask = ids >= lengths.unsqueeze(1).expand(-1, max_len) + + return mask + + +def pad(input_ele, mel_max_length=None): + if mel_max_length: + max_len = mel_max_length + else: + max_len = max([input_ele[i].size(0) for i in range(len(input_ele))]) + + out_list = list() + for i, batch in enumerate(input_ele): + if len(batch.shape) == 1: + one_batch_padded = F.pad( + batch, (0, max_len - batch.size(0)), "constant", 0.0 + ) + elif len(batch.shape) == 2: + 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 + + +class VarianceAdaptor(nn.Module): + """Variance Adaptor""" + + def __init__(self, cfg): + super(VarianceAdaptor, self).__init__() + self.duration_predictor = VariancePredictor(cfg) + self.length_regulator = LengthRegulator() + self.pitch_predictor = VariancePredictor(cfg) + self.energy_predictor = VariancePredictor(cfg) + + # assign the pitch/energy feature level + if cfg.preprocess.use_frame_pitch: + self.pitch_feature_level = "frame_level" + self.pitch_dir = cfg.preprocess.pitch_dir + else: + self.pitch_feature_level = "phoneme_level" + self.pitch_dir = cfg.preprocess.phone_pitch_dir + + if cfg.preprocess.use_frame_energy: + self.energy_feature_level = "frame_level" + self.energy_dir = cfg.preprocess.energy_dir + else: + self.energy_feature_level = "phoneme_level" + self.energy_dir = cfg.preprocess.phone_energy_dir + + assert self.pitch_feature_level in ["phoneme_level", "frame_level"] + assert self.energy_feature_level in ["phoneme_level", "frame_level"] + + pitch_quantization = cfg.model.variance_embedding.pitch_quantization + energy_quantization = cfg.model.variance_embedding.energy_quantization + n_bins = cfg.model.variance_embedding.n_bins + assert pitch_quantization in ["linear", "log"] + assert energy_quantization in ["linear", "log"] + + with open( + os.path.join( + cfg.preprocess.processed_dir, + cfg.dataset[0], + self.energy_dir, + "statistics.json", + ) + ) as f: + stats = json.load(f) + stats = stats[cfg.dataset[0] + "_" + cfg.dataset[0]] + mean, std = ( + stats["voiced_positions"]["mean"], + stats["voiced_positions"]["std"], + ) + energy_min = (stats["total_positions"]["min"] - mean) / std + energy_max = (stats["total_positions"]["max"] - mean) / std + + with open( + os.path.join( + cfg.preprocess.processed_dir, + cfg.dataset[0], + self.pitch_dir, + "statistics.json", + ) + ) as f: + stats = json.load(f) + stats = stats[cfg.dataset[0] + "_" + cfg.dataset[0]] + mean, std = ( + stats["voiced_positions"]["mean"], + stats["voiced_positions"]["std"], + ) + pitch_min = (stats["total_positions"]["min"] - mean) / std + pitch_max = (stats["total_positions"]["max"] - mean) / std + + if pitch_quantization == "log": + self.pitch_bins = nn.Parameter( + torch.exp( + torch.linspace(np.log(pitch_min), np.log(pitch_max), n_bins - 1) + ), + requires_grad=False, + ) + else: + self.pitch_bins = nn.Parameter( + torch.linspace(pitch_min, pitch_max, n_bins - 1), + requires_grad=False, + ) + if energy_quantization == "log": + self.energy_bins = nn.Parameter( + torch.exp( + torch.linspace(np.log(energy_min), np.log(energy_max), n_bins - 1) + ), + requires_grad=False, + ) + else: + self.energy_bins = nn.Parameter( + torch.linspace(energy_min, energy_max, n_bins - 1), + requires_grad=False, + ) + + self.pitch_embedding = nn.Embedding( + n_bins, cfg.model.transformer.encoder_hidden + ) + self.energy_embedding = nn.Embedding( + n_bins, cfg.model.transformer.encoder_hidden + ) + + def get_pitch_embedding(self, x, target, mask, control): + prediction = self.pitch_predictor(x, mask) + if target is not None: + embedding = self.pitch_embedding(torch.bucketize(target, self.pitch_bins)) + else: + prediction = prediction * control + embedding = self.pitch_embedding( + torch.bucketize(prediction, self.pitch_bins) + ) + return prediction, embedding + + def get_energy_embedding(self, x, target, mask, control): + prediction = self.energy_predictor(x, mask) + if target is not None: + embedding = self.energy_embedding(torch.bucketize(target, self.energy_bins)) + else: + prediction = prediction * control + embedding = self.energy_embedding( + torch.bucketize(prediction, self.energy_bins) + ) + return prediction, embedding + + def forward( + self, + x, + src_mask, + mel_mask=None, + max_len=None, + pitch_target=None, + energy_target=None, + duration_target=None, + p_control=1.0, + e_control=1.0, + d_control=1.0, + pitch_embedding=None, + energy_embedding=None, + ): + log_duration_prediction = self.duration_predictor(x, src_mask) + + x = x + pitch_embedding + x = x + energy_embedding + + pitch_prediction = self.pitch_predictor(x, src_mask) + energy_prediction = self.energy_predictor(x, src_mask) + + if duration_target is not None: + x, mel_len = self.length_regulator(x, duration_target, max_len) + duration_rounded = duration_target + else: + duration_rounded = torch.clamp( + (torch.round(torch.exp(log_duration_prediction) - 1) * d_control), + min=0, + ) + x, mel_len = self.length_regulator(x, duration_rounded, max_len) + mel_mask = get_mask_from_lengths(mel_len) + + return ( + x, + pitch_prediction, + energy_prediction, + log_duration_prediction, + duration_rounded, + mel_len, + mel_mask, + ) + + def inference( + self, + x, + src_mask, + mel_mask=None, + max_len=None, + pitch_target=None, + energy_target=None, + duration_target=None, + p_control=1.0, + e_control=1.0, + d_control=1.0, + pitch_embedding=None, + energy_embedding=None, + ): + + p_outs = self.pitch_predictor(x, src_mask) + e_outs = self.energy_predictor(x, src_mask) + d_outs = self.duration_predictor(x, src_mask) + + return p_outs, e_outs, d_outs + + +class LengthRegulator(nn.Module): + """Length Regulator""" + + def __init__(self): + super(LengthRegulator, self).__init__() + + def LR(self, x, duration, max_len): + device = x.device + output = list() + mel_len = list() + for batch, expand_target in zip(x, duration): + expanded = self.expand(batch, expand_target) + output.append(expanded) + mel_len.append(expanded.shape[0]) + + if max_len is not None: + output = pad(output, max_len) + else: + output = pad(output) + + return output, torch.LongTensor(mel_len).to(device) + + def expand(self, batch, predicted): + out = list() + + for i, vec in enumerate(batch): + expand_size = predicted[i].item() + out.append(vec.expand(max(int(expand_size), 0), -1)) + out = torch.cat(out, 0) + + return out + + def forward(self, x, duration, max_len): + output, mel_len = self.LR(x, duration, max_len) + return output, mel_len + + +class VariancePredictor(nn.Module): + """Duration, Pitch and Energy Predictor""" + + def __init__(self, cfg): + super(VariancePredictor, self).__init__() + + self.input_size = cfg.model.transformer.encoder_hidden + self.filter_size = cfg.model.variance_predictor.filter_size + self.kernel = cfg.model.variance_predictor.kernel_size + self.conv_output_size = cfg.model.variance_predictor.filter_size + self.dropout = cfg.model.variance_predictor.dropout + + self.conv_layer = nn.Sequential( + OrderedDict( + [ + ( + "conv1d_1", + Conv( + self.input_size, + self.filter_size, + kernel_size=self.kernel, + padding=(self.kernel - 1) // 2, + ), + ), + ("relu_1", nn.ReLU()), + ("layer_norm_1", nn.LayerNorm(self.filter_size)), + ("dropout_1", nn.Dropout(self.dropout)), + ( + "conv1d_2", + Conv( + self.filter_size, + self.filter_size, + kernel_size=self.kernel, + padding=1, + ), + ), + ("relu_2", nn.ReLU()), + ("layer_norm_2", nn.LayerNorm(self.filter_size)), + ("dropout_2", nn.Dropout(self.dropout)), + ] + ) + ) + + self.linear_layer = nn.Linear(self.conv_output_size, 1) + + def forward(self, encoder_output, mask): + out = self.conv_layer(encoder_output) + out = self.linear_layer(out) + out = out.squeeze(-1) + + if mask is not None: + out = out.masked_fill(mask, 0.0) + + return out + + +class Conv(nn.Module): + """ + Convolution Module + """ + + def __init__( + self, + in_channels, + out_channels, + kernel_size=1, + stride=1, + padding=0, + dilation=1, + bias=True, + w_init="linear", + ): + """ + :param in_channels: dimension of input + :param out_channels: dimension of output + :param kernel_size: size of kernel + :param stride: size of stride + :param padding: size of padding + :param dilation: dilation rate + :param bias: boolean. if True, bias is included. + :param w_init: str. weight inits with xavier initialization. + """ + super(Conv, self).__init__() + + self.conv = nn.Conv1d( + in_channels, + out_channels, + kernel_size=kernel_size, + stride=stride, + padding=padding, + dilation=dilation, + bias=bias, + ) + + def forward(self, x): + x = x.contiguous().transpose(1, 2) + x = self.conv(x) + x = x.contiguous().transpose(1, 2) + + return x + + +class Jets(nn.Module): + def __init__(self, cfg) -> None: + super(Jets, self).__init__() + self.cfg = cfg + self.encoder = Encoder(cfg.model) + self.variance_adaptor = VarianceAdaptor(cfg) + self.decoder = Decoder(cfg.model) + self.length_regulator_infer = LengthRegulator() + self.mel_linear = nn.Linear( + cfg.model.transformer.decoder_hidden, + cfg.preprocess.n_mel, + ) + self.postnet = PostNet(n_mel_channels=cfg.preprocess.n_mel) + + self.speaker_emb = None + if cfg.train.multi_speaker_training: + with open( + os.path.join( + cfg.preprocess.processed_dir, cfg.dataset[0], "spk2id.json" + ), + "r", + ) as f: + n_speaker = len(json.load(f)) + self.speaker_emb = nn.Embedding( + n_speaker, + cfg.model.transformer.encoder_hidden, + ) + + output_dim = cfg.preprocess.n_mel + attention_dim = 256 + self.alignment_module = AlignmentModule(attention_dim, output_dim) + + # NOTE(kan-bayashi): We use continuous pitch + FastPitch style avg + pitch_embed_kernel_size: int = 9 + pitch_embed_dropout: float = 0.5 + self.pitch_embed = torch.nn.Sequential( + torch.nn.Conv1d( + in_channels=1, + out_channels=attention_dim, + kernel_size=pitch_embed_kernel_size, + padding=(pitch_embed_kernel_size - 1) // 2, + ), + torch.nn.Dropout(pitch_embed_dropout), + ) + + # NOTE(kan-bayashi): We use continuous enegy + FastPitch style avg + energy_embed_kernel_size: int = 9 + energy_embed_dropout: float = 0.5 + self.energy_embed = torch.nn.Sequential( + torch.nn.Conv1d( + in_channels=1, + out_channels=attention_dim, + kernel_size=energy_embed_kernel_size, + padding=(energy_embed_kernel_size - 1) // 2, + ), + torch.nn.Dropout(energy_embed_dropout), + ) + + # define length regulator + self.length_regulator = GaussianUpsampling() + + self.segment_size = cfg.train.segment_size + + # Define HiFiGAN generator + hifi_cfg = load_config("egs/vocoder/gan/hifigan/exp_config.json") + # hifi_cfg.model.hifigan.resblock_kernel_sizes = [3, 7, 11] + hifi_cfg.preprocess.n_mel = attention_dim + self.generator = HiFiGAN(hifi_cfg) + + def _source_mask(self, ilens: torch.Tensor) -> torch.Tensor: + """Make masks for self-attention. + + Args: + ilens (LongTensor): Batch of lengths (B,). + + Returns: + Tensor: Mask tensor for self-attention. + dtype=torch.uint8 in PyTorch 1.2- + dtype=torch.bool in PyTorch 1.2+ (including 1.2) + + Examples: + >>> ilens = [5, 3] + >>> self._source_mask(ilens) + tensor([[[1, 1, 1, 1, 1], + [1, 1, 1, 0, 0]]], dtype=torch.uint8) + + """ + x_masks = make_non_pad_mask(ilens).to(next(self.parameters()).device) + return x_masks.unsqueeze(-2) + + def forward(self, data, p_control=1.0, e_control=1.0, d_control=1.0): + speakers = data["spk_id"] + texts = data["texts"] + src_lens = data["text_len"] + max_src_len = max(src_lens) + feats = data["mel"] + mel_lens = data["target_len"] if "target_len" in data else None + feats_lengths = mel_lens + max_mel_len = max(mel_lens) if "target_len" in data else None + p_targets = data["pitch"] if "pitch" in data else None + e_targets = data["energy"] if "energy" in data else None + src_masks = get_mask_from_lengths(src_lens, max_src_len) + mel_masks = ( + get_mask_from_lengths(mel_lens, max_mel_len) + if mel_lens is not None + else None + ) + + output = self.encoder(texts, src_masks) + + if self.speaker_emb is not None: + output = output + self.speaker_emb(speakers).unsqueeze(1).expand( + -1, max_src_len, -1 + ) + + # Forward alignment module and obtain duration, averaged pitch, energy + h_masks = make_pad_mask(src_lens).to(output.device) + log_p_attn = self.alignment_module( + output, feats, src_lens, feats_lengths, h_masks + ) + ds, bin_loss = viterbi_decode(log_p_attn, src_lens, feats_lengths) + ps = average_by_duration( + ds, p_targets.squeeze(-1), src_lens, feats_lengths + ).unsqueeze(-1) + es = average_by_duration( + ds, e_targets.squeeze(-1), src_lens, feats_lengths + ).unsqueeze(-1) + p_embs = self.pitch_embed(ps.transpose(1, 2)).transpose(1, 2) + e_embs = self.energy_embed(es.transpose(1, 2)).transpose(1, 2) + + # FastSpeech2 variance adaptor + ( + output, + p_predictions, + e_predictions, + log_d_predictions, + d_rounded, + mel_lens, + mel_masks, + ) = self.variance_adaptor( + output, + src_masks, + mel_masks, + max_mel_len, + p_targets, + e_targets, + ds, + p_control, + e_control, + d_control, + ps, + es, + ) + + # forward decoder + zs, _ = self.decoder(output, mel_masks) # (B, T_feats, adim) + + # get random segments + z_segments, z_start_idxs = get_random_segments( + zs.transpose(1, 2), + feats_lengths, + self.segment_size, + ) + + # forward generator + wav = self.generator(z_segments) + + return ( + wav, + bin_loss, + log_p_attn, + z_start_idxs, + log_d_predictions, + ds, + p_predictions, + ps, + e_predictions, + es, + src_lens, + feats_lengths, + ) + + def inference(self, data, p_control=1.0, e_control=1.0, d_control=1.0): + speakers = data["spk_id"] + texts = data["texts"] + src_lens = data["text_len"] + max_src_len = max(src_lens) + mel_lens = data["target_len"] if "target_len" in data else None + feats_lengths = mel_lens + max_mel_len = max(mel_lens) if "target_len" in data else None + p_targets = data["pitch"] if "pitch" in data else None + e_targets = data["energy"] if "energy" in data else None + d_targets = data["durations"] if "durations" in data else None + src_masks = get_mask_from_lengths(src_lens, max_src_len) + mel_masks = ( + get_mask_from_lengths(mel_lens, max_mel_len) + if mel_lens is not None + else None + ) + + x_masks = self._source_mask(src_lens) + hs = self.encoder(texts, src_masks) + + ( + p_outs, + e_outs, + d_outs, + ) = self.variance_adaptor.inference( + hs, + src_masks, + ) + + p_embs = self.pitch_embed(p_outs.unsqueeze(-1).transpose(1, 2)).transpose(1, 2) + e_embs = self.energy_embed(e_outs.unsqueeze(-1).transpose(1, 2)).transpose(1, 2) + hs = hs + e_embs + p_embs + + # Duration predictor inference mode: log_d_pred to d_pred + offset = 1.0 + d_predictions = torch.clamp( + torch.round(d_outs.exp() - offset), min=0 + ).long() # avoid negative value + + # forward decoder + hs, mel_len = self.length_regulator_infer(hs, d_predictions, max_mel_len) + mel_mask = get_mask_from_lengths(mel_len) + zs, _ = self.decoder(hs, mel_mask) # (B, T_feats, adim) + + # forward generator + wav = self.generator(zs.transpose(1, 2)) + + return wav, d_predictions diff --git a/models/tts/jets/jets_dataset.py b/models/tts/jets/jets_dataset.py new file mode 100644 index 0000000..6909468 --- /dev/null +++ b/models/tts/jets/jets_dataset.py @@ -0,0 +1,451 @@ +# Copyright (c) 2024 Amphion. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +import random +import torch +import torchaudio +from torch.nn.utils.rnn import pad_sequence +from utils.data_utils import * +from models.base.base_dataset import ( + BaseOfflineCollator, + BaseOfflineDataset, + BaseTestDataset, + BaseTestCollator, +) +from text import text_to_sequence + + +class JetsDataset(BaseOfflineDataset): + def __init__(self, cfg, dataset, is_valid=False): + BaseOfflineDataset.__init__(self, cfg, dataset, is_valid=is_valid) + self.batch_size = cfg.train.batch_size + cfg = cfg.preprocess + # utt2duration + self.utt2duration_path = {} + for utt_info in self.metadata: + dataset = utt_info["Dataset"] + uid = utt_info["Uid"] + utt = "{}_{}".format(dataset, uid) + + self.utt2duration_path[utt] = os.path.join( + cfg.processed_dir, + dataset, + cfg.duration_dir, + uid + ".npy", + ) + self.utt2dur = self.read_duration() + + if cfg.use_frame_energy: + self.frame_utt2energy, self.energy_statistic = load_energy( + self.metadata, + cfg.processed_dir, + cfg.energy_dir, + use_log_scale=cfg.use_log_scale_energy, + utt2spk=self.preprocess.utt2spk if cfg.use_spkid else None, + return_norm=True, + ) + elif cfg.use_phone_energy: + self.phone_utt2energy, self.energy_statistic = load_energy( + self.metadata, + cfg.processed_dir, + cfg.phone_energy_dir, + use_log_scale=cfg.use_log_scale_energy, + utt2spk=self.utt2spk if cfg.use_spkid else None, + return_norm=True, + ) + + if cfg.use_frame_pitch: + self.frame_utt2pitch, self.pitch_statistic = load_energy( + self.metadata, + cfg.processed_dir, + cfg.pitch_dir, + use_log_scale=cfg.energy_extract_mode, + utt2spk=self.utt2spk if cfg.use_spkid else None, + return_norm=True, + ) + + elif cfg.use_phone_pitch: + self.phone_utt2pitch, self.pitch_statistic = load_energy( + self.metadata, + cfg.processed_dir, + cfg.phone_pitch_dir, + use_log_scale=cfg.use_log_scale_pitch, + utt2spk=self.utt2spk if cfg.use_spkid else None, + return_norm=True, + ) + + # utt2lab + self.utt2lab_path = {} + for utt_info in self.metadata: + dataset = utt_info["Dataset"] + uid = utt_info["Uid"] + utt = "{}_{}".format(dataset, uid) + + self.utt2lab_path[utt] = os.path.join( + cfg.processed_dir, + dataset, + cfg.lab_dir, + uid + ".txt", + ) + + self.speaker_map = {} + if os.path.exists(os.path.join(cfg.processed_dir, "spk2id.json")): + with open( + os.path.exists(os.path.join(cfg.processed_dir, "spk2id.json")) + ) as f: + self.speaker_map = json.load(f) + + self.metadata = self.check_metadata() + if cfg.use_audios: + self.utt2audio_path = {} + for utt_info in self.metadata: + dataset = utt_info["Dataset"] + uid = utt_info["Uid"] + utt = "{}_{}".format(dataset, uid) + + if cfg.extract_audio: + self.utt2audio_path[utt] = os.path.join( + cfg.processed_dir, + dataset, + cfg.audio_dir, + uid + ".wav", + ) + else: + self.utt2audio_path[utt] = utt_info["Path"] + + def __getitem__(self, index): + single_feature = BaseOfflineDataset.__getitem__(self, index) + + utt_info = self.metadata[index] + dataset = utt_info["Dataset"] + uid = utt_info["Uid"] + utt = "{}_{}".format(dataset, uid) + + duration = self.utt2dur[utt] + + # text + f = open(self.utt2lab_path[utt], "r") + phones = f.readlines()[0].strip() + f.close() + # todo: add cleaner(chenxi) + phones_ids = np.array(text_to_sequence(phones, ["english_cleaners"])) + text_len = len(phones_ids) + + if self.cfg.preprocess.use_frame_pitch: + pitch = self.frame_utt2pitch[utt] + elif self.cfg.preprocess.use_phone_pitch: + pitch = self.phone_utt2pitch[utt] + + if self.cfg.preprocess.use_frame_energy: + energy = self.frame_utt2energy[utt] + elif self.cfg.preprocess.use_phone_energy: + energy = self.phone_utt2energy[utt] + + # speaker + if len(self.speaker_map) > 0: + speaker_id = self.speaker_map[utt_info["Singer"]] + else: + speaker_id = 0 + + single_feature.update( + { + "durations": duration, + "texts": phones_ids, + "spk_id": speaker_id, + "text_len": text_len, + "pitch": pitch, + "energy": energy, + "uid": uid, + } + ) + + if self.cfg.preprocess.use_audios: + audio, sr = torchaudio.load(self.utt2audio_path[utt]) + audio = audio.cpu().numpy().squeeze() + single_feature["audio"] = audio + single_feature["audio_len"] = audio.shape[0] + return self.clip_if_too_long(single_feature) + + def read_duration(self): + # read duration + utt2dur = {} + for index in range(len(self.metadata)): + utt_info = self.metadata[index] + dataset = utt_info["Dataset"] + uid = utt_info["Uid"] + utt = "{}_{}".format(dataset, uid) + + if not os.path.exists(self.utt2mel_path[utt]) or not os.path.exists( + self.utt2duration_path[utt] + ): + continue + + mel = np.load(self.utt2mel_path[utt]).transpose(1, 0) + duration = np.load(self.utt2duration_path[utt]) + assert mel.shape[0] == sum( + duration + ), f"{utt}: mismatch length between mel {mel.shape[0]} and sum(duration) {sum(duration)}" + utt2dur[utt] = duration + return utt2dur + + def __len__(self): + return len(self.metadata) + + def random_select(self, feature_seq_len, max_seq_len, ending_ts=2812): + """ + ending_ts: to avoid invalid whisper features for over 30s audios + 2812 = 30 * 24000 // 256 + """ + ts = max(feature_seq_len - max_seq_len, 0) + ts = min(ts, ending_ts - max_seq_len) + + start = random.randint(0, ts) + end = start + max_seq_len + return start, end + + def clip_if_too_long(self, sample, max_seq_len=1000): + """ + sample : + { + 'spk_id': (1,), + 'target_len': int + 'mel': (seq_len, dim), + 'frame_pitch': (seq_len,) + 'frame_energy': (seq_len,) + 'content_vector_feat': (seq_len, dim) + } + """ + if sample["target_len"] <= max_seq_len: + return sample + + start, end = self.random_select(sample["target_len"], max_seq_len) + sample["target_len"] = end - start + + for k in sample.keys(): + if k not in ["spk_id", "target_len"]: + sample[k] = sample[k][start:end] + + return sample + + def check_metadata(self): + new_metadata = [] + for utt_info in self.metadata: + dataset = utt_info["Dataset"] + uid = utt_info["Uid"] + utt = "{}_{}".format(dataset, uid) + if not os.path.exists(self.utt2duration_path[utt]) or not os.path.exists( + self.utt2mel_path[utt] + ): + continue + else: + new_metadata.append(utt_info) + return new_metadata + + +class JetsCollator(BaseOfflineCollator): + """Zero-pads model inputs and targets based on number of frames per step""" + + def __init__(self, cfg): + BaseOfflineCollator.__init__(self, cfg) + self.sort = cfg.train.sort_sample + self.batch_size = cfg.train.batch_size + self.drop_last = cfg.train.drop_last + + def __call__(self, batch): + # mel: [b, T, n_mels] + # frame_pitch, frame_energy: [1, T] + # target_len: [1] + # spk_id: [b, 1] + # mask: [b, T, 1] + packed_batch_features = dict() + + for key in batch[0].keys(): + if key == "target_len": + packed_batch_features["target_len"] = torch.LongTensor( + [b["target_len"] for b in batch] + ) + masks = [ + torch.ones((b["target_len"], 1), dtype=torch.long) for b in batch + ] + packed_batch_features["mask"] = pad_sequence( + masks, batch_first=True, padding_value=0 + ) + elif key == "text_len": + packed_batch_features["text_len"] = torch.LongTensor( + [b["text_len"] for b in batch] + ) + masks = [ + torch.ones((b["text_len"], 1), dtype=torch.long) for b in batch + ] + packed_batch_features["text_mask"] = pad_sequence( + masks, batch_first=True, padding_value=0 + ) + elif key == "spk_id": + packed_batch_features["spk_id"] = torch.LongTensor( + [b["spk_id"] for b in batch] + ) + elif key == "uid": + packed_batch_features[key] = [b["uid"] for b in batch] + elif key == "audio_len": + packed_batch_features["audio_len"] = torch.LongTensor( + [b["audio_len"] for b in batch] + ) + else: + values = [torch.from_numpy(b[key]) for b in batch] + packed_batch_features[key] = pad_sequence( + values, batch_first=True, padding_value=0 + ) + return packed_batch_features + + +class JetsTestDataset(BaseTestDataset): + def __init__(self, args, cfg, infer_type=None): + datasets = cfg.dataset + cfg = cfg.preprocess + is_bigdata = False + + assert len(datasets) >= 1 + if len(datasets) > 1: + datasets.sort() + bigdata_version = "_".join(datasets) + processed_data_dir = os.path.join(cfg.processed_dir, bigdata_version) + is_bigdata = True + else: + processed_data_dir = os.path.join(cfg.processed_dir, args.dataset) + + if args.test_list_file: + self.metafile_path = args.test_list_file + self.metadata = self.get_metadata() + else: + assert args.testing_set + source_metafile_path = os.path.join( + cfg.processed_dir, + args.dataset, + "{}.json".format(args.testing_set), + ) + with open(source_metafile_path, "r") as f: + self.metadata = json.load(f) + + self.cfg = cfg + self.datasets = datasets + self.data_root = processed_data_dir + self.is_bigdata = is_bigdata + self.source_dataset = args.dataset + + ######### Load source acoustic features ######### + if cfg.use_spkid: + spk2id_path = os.path.join(self.data_root, cfg.spk2id) + utt2sp_path = os.path.join(self.data_root, cfg.utt2spk) + self.spk2id, self.utt2spk = get_spk_map(spk2id_path, utt2sp_path, datasets) + + # utt2lab + self.utt2lab_path = {} + for utt_info in self.metadata: + dataset = utt_info["Dataset"] + uid = utt_info["Uid"] + utt = "{}_{}".format(dataset, uid) + self.utt2lab_path[utt] = os.path.join( + cfg.processed_dir, + dataset, + cfg.lab_dir, + uid + ".txt", + ) + + self.speaker_map = {} + if os.path.exists(os.path.join(cfg.processed_dir, "spk2id.json")): + with open( + os.path.exists(os.path.join(cfg.processed_dir, "spk2id.json")) + ) as f: + self.speaker_map = json.load(f) + + def __getitem__(self, index): + single_feature = {} + + utt_info = self.metadata[index] + dataset = utt_info["Dataset"] + uid = utt_info["Uid"] + utt = "{}_{}".format(dataset, uid) + + # text + f = open(self.utt2lab_path[utt], "r") + phones = f.readlines()[0].strip() + f.close() + + phones_ids = np.array(text_to_sequence(phones, self.cfg.text_cleaners)) + text_len = len(phones_ids) + + # speaker + if len(self.speaker_map) > 0: + speaker_id = self.speaker_map[utt_info["Singer"]] + else: + speaker_id = 0 + + single_feature.update( + { + "texts": phones_ids, + "spk_id": speaker_id, + "text_len": text_len, + } + ) + + return single_feature + + def __len__(self): + return len(self.metadata) + + def get_metadata(self): + with open(self.metafile_path, "r", encoding="utf-8") as f: + metadata = json.load(f) + + return metadata + + +class JetsTestCollator(BaseTestCollator): + """Zero-pads model inputs and targets based on number of frames per step""" + + def __init__(self, cfg): + self.cfg = cfg + + def __call__(self, batch): + packed_batch_features = dict() + + # mel: [b, T, n_mels] + # frame_pitch, frame_energy: [1, T] + # target_len: [1] + # spk_id: [b, 1] + # mask: [b, T, 1] + + for key in batch[0].keys(): + if key == "target_len": + packed_batch_features["target_len"] = torch.LongTensor( + [b["target_len"] for b in batch] + ) + masks = [ + torch.ones((b["target_len"], 1), dtype=torch.long) for b in batch + ] + packed_batch_features["mask"] = pad_sequence( + masks, batch_first=True, padding_value=0 + ) + elif key == "text_len": + packed_batch_features["text_len"] = torch.LongTensor( + [b["text_len"] for b in batch] + ) + masks = [ + torch.ones((b["text_len"], 1), dtype=torch.long) for b in batch + ] + packed_batch_features["text_mask"] = pad_sequence( + masks, batch_first=True, padding_value=0 + ) + elif key == "spk_id": + packed_batch_features["spk_id"] = torch.LongTensor( + [b["spk_id"] for b in batch] + ) + else: + values = [torch.from_numpy(b[key]) for b in batch] + packed_batch_features[key] = pad_sequence( + values, batch_first=True, padding_value=0 + ) + + return packed_batch_features diff --git a/models/tts/jets/jets_inference.py b/models/tts/jets/jets_inference.py new file mode 100644 index 0000000..35610c4 --- /dev/null +++ b/models/tts/jets/jets_inference.py @@ -0,0 +1,71 @@ +# Copyright (c) 2024 Amphion. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +import os +import torch +from tqdm import tqdm +from collections import OrderedDict + +from models.tts.base.tts_inferece import TTSInference +from models.tts.jets.jets_dataset import JetsTestDataset, JetsTestCollator +from utils.util import load_config +from utils.io import save_audio +from models.tts.jets.jets import Jets +from models.vocoders.vocoder_inference import synthesis +from pathlib import Path +from processors.phone_extractor import phoneExtractor +from text.text_token_collation import phoneIDCollation +import numpy as np +import json +import time + + +class JetsInference(TTSInference): + def __init__(self, args, cfg): + TTSInference.__init__(self, args, cfg) + self.args = args + self.cfg = cfg + self.infer_type = args.mode + + def _build_model(self): + self.model = Jets(self.cfg) + return self.model + + def _build_test_dataset(self): + return JetsTestDataset, JetsTestCollator + + def inference_for_batches(self): + ###### Construct test_batch ###### + n_batch = len(self.test_dataloader) + now = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(time.time())) + print( + "Model eval time: {}, batch_size = {}, n_batch = {}".format( + now, self.test_batch_size, n_batch + ) + ) + self.model.eval() + + ###### Inference for each batch ###### + pred_res = [] + with torch.no_grad(): + for i, batch_data in enumerate( + self.test_dataloader if n_batch == 1 else tqdm(self.test_dataloader) + ): + outputs = self.model.inference(batch_data) + + audios, d_predictions = outputs + d_predictions = d_predictions.unsqueeze(-1) + + for idx in range(audios.size(0)): + audio = audios[idx, 0, :].data.cpu().float() + duration = d_predictions[idx, :, :] + audio_length = ( + duration.sum([0, 1]).long() * self.cfg.preprocess.hop_size + ) + audio_length = audio_length.cpu().numpy() + audio = audio[:audio_length] + pred_res.append(audio) + + return pred_res diff --git a/models/tts/jets/jets_loss.py b/models/tts/jets/jets_loss.py new file mode 100644 index 0000000..3955734 --- /dev/null +++ b/models/tts/jets/jets_loss.py @@ -0,0 +1,537 @@ +# Copyright (c) 2024 Amphion. +# +# This code is modified from https://github.com/imdanboy/jets/blob/main/espnet2/gan_tts/jets/loss.py +# Licensed under Apache License 2.0 + +from typing import List, Optional, Tuple, Union + +import torch +import torch.nn as nn +import torch.nn.functional as F +import numpy as np +import librosa + +from models.vocoders.gan.discriminator.mpd import MultiScaleMultiPeriodDiscriminator +from models.tts.jets.alignments import make_non_pad_mask, make_pad_mask + + +class GeneratorAdversarialLoss(torch.nn.Module): + """Generator adversarial loss module.""" + + def __init__(self): + super().__init__() + + def forward(self, outputs) -> torch.Tensor: + if isinstance(outputs, (tuple, list)): + adv_loss = 0.0 + for i, outputs_ in enumerate(outputs): + if isinstance(outputs_, (tuple, list)): + # NOTE(kan-bayashi): case including feature maps + outputs_ = outputs_[-1] + adv_loss += F.mse_loss(outputs_, outputs_.new_ones(outputs_.size())) + else: + adv_loss = F.mse_loss(outputs, outputs.new_ones(outputs.size())) + + return adv_loss + + +class FeatureMatchLoss(torch.nn.Module): + """Feature matching loss module.""" + + def __init__( + self, + average_by_layers: bool = False, + average_by_discriminators: bool = False, + include_final_outputs: bool = True, + ): + """Initialize FeatureMatchLoss module. + + Args: + average_by_layers (bool): Whether to average the loss by the number + of layers. + average_by_discriminators (bool): Whether to average the loss by + the number of discriminators. + include_final_outputs (bool): Whether to include the final output of + each discriminator for loss calculation. + + """ + super().__init__() + self.average_by_layers = average_by_layers + self.average_by_discriminators = average_by_discriminators + self.include_final_outputs = include_final_outputs + + def forward( + self, + feats_hat: Union[List[List[torch.Tensor]], List[torch.Tensor]], + feats: Union[List[List[torch.Tensor]], List[torch.Tensor]], + ) -> torch.Tensor: + """Calculate feature matching loss. + + Args: + feats_hat (Union[List[List[Tensor]], List[Tensor]]): List of list of + discriminator outputs or list of discriminator outputs calcuated + from generator's outputs. + feats (Union[List[List[Tensor]], List[Tensor]]): List of list of + discriminator outputs or list of discriminator outputs calcuated + from groundtruth.. + + Returns: + Tensor: Feature matching loss value. + + """ + feat_match_loss = 0.0 + for i, (feats_hat_, feats_) in enumerate(zip(feats_hat, feats)): + feat_match_loss_ = 0.0 + if not self.include_final_outputs: + feats_hat_ = feats_hat_[:-1] + feats_ = feats_[:-1] + for j, (feat_hat_, feat_) in enumerate(zip(feats_hat_, feats_)): + feat_match_loss_ += F.l1_loss(feat_hat_, feat_.detach()) + if self.average_by_layers: + feat_match_loss_ /= j + 1 + feat_match_loss += feat_match_loss_ + if self.average_by_discriminators: + feat_match_loss /= i + 1 + + return feat_match_loss + + +class DurationPredictorLoss(torch.nn.Module): + """Loss function module for duration predictor. + + The loss value is Calculated in log domain to make it Gaussian. + + """ + + def __init__(self, offset=1.0, reduction="mean"): + """Initilize duration predictor loss module. + + Args: + offset (float, optional): Offset value to avoid nan in log domain. + reduction (str): Reduction type in loss calculation. + + """ + super().__init__() + self.criterion = torch.nn.MSELoss(reduction=reduction) + self.offset = offset + + def forward(self, outputs, targets): + targets = torch.log(targets.float() + self.offset) + loss = self.criterion(outputs, targets) + + return loss + + +class VarianceLoss(torch.nn.Module): + def __init__(self): + """Initialize JETS variance loss module.""" + super().__init__() + + # define criterions + reduction = "mean" + self.mse_criterion = torch.nn.MSELoss(reduction=reduction) + self.duration_criterion = DurationPredictorLoss(reduction=reduction) + + def forward( + self, + d_outs: torch.Tensor, + ds: torch.Tensor, + p_outs: torch.Tensor, + ps: torch.Tensor, + e_outs: torch.Tensor, + es: torch.Tensor, + ilens: torch.Tensor, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Calculate forward propagation. + + Args: + d_outs (LongTensor): Batch of outputs of duration predictor (B, T_text). + ds (LongTensor): Batch of durations (B, T_text). + p_outs (Tensor): Batch of outputs of pitch predictor (B, T_text, 1). + ps (Tensor): Batch of target token-averaged pitch (B, T_text, 1). + e_outs (Tensor): Batch of outputs of energy predictor (B, T_text, 1). + es (Tensor): Batch of target token-averaged energy (B, T_text, 1). + ilens (LongTensor): Batch of the lengths of each input (B,). + + Returns: + Tensor: Duration predictor loss value. + Tensor: Pitch predictor loss value. + Tensor: Energy predictor loss value. + + """ + # apply mask to remove padded part + duration_masks = make_non_pad_mask(ilens).to(ds.device) + d_outs = d_outs.masked_select(duration_masks) + ds = ds.masked_select(duration_masks) + pitch_masks = make_non_pad_mask(ilens).to(ds.device) + pitch_masks_ = make_non_pad_mask(ilens).unsqueeze(-1).to(ds.device) + p_outs = p_outs.masked_select(pitch_masks) + e_outs = e_outs.masked_select(pitch_masks) + ps = ps.masked_select(pitch_masks_) + es = es.masked_select(pitch_masks_) + + # calculate loss + duration_loss = self.duration_criterion(d_outs, ds) + pitch_loss = self.mse_criterion(p_outs, ps) + energy_loss = self.mse_criterion(e_outs, es) + + return duration_loss, pitch_loss, energy_loss + + +class ForwardSumLoss(torch.nn.Module): + """Forwardsum loss described at https://openreview.net/forum?id=0NQwnnwAORi""" + + def __init__(self): + """Initialize forwardsum loss module.""" + super().__init__() + + def forward( + self, + log_p_attn: torch.Tensor, + ilens: torch.Tensor, + olens: torch.Tensor, + blank_prob: float = np.e**-1, + ) -> torch.Tensor: + """Calculate forward propagation. + + Args: + log_p_attn (Tensor): Batch of log probability of attention matrix + (B, T_feats, T_text). + ilens (Tensor): Batch of the lengths of each input (B,). + olens (Tensor): Batch of the lengths of each target (B,). + blank_prob (float): Blank symbol probability. + + Returns: + Tensor: forwardsum loss value. + + """ + B = log_p_attn.size(0) + + # a row must be added to the attention matrix to account for + # blank token of CTC loss + # (B,T_feats,T_text+1) + log_p_attn_pd = F.pad(log_p_attn, (1, 0, 0, 0, 0, 0), value=np.log(blank_prob)) + + loss = 0 + for bidx in range(B): + # construct target sequnece. + # Every text token is mapped to a unique sequnece number. + target_seq = torch.arange(1, ilens[bidx] + 1).unsqueeze(0) + cur_log_p_attn_pd = log_p_attn_pd[ + bidx, : olens[bidx], : ilens[bidx] + 1 + ].unsqueeze( + 1 + ) # (T_feats,1,T_text+1) + cur_log_p_attn_pd = F.log_softmax(cur_log_p_attn_pd, dim=-1) + loss += F.ctc_loss( + log_probs=cur_log_p_attn_pd, + targets=target_seq, + input_lengths=olens[bidx : bidx + 1], + target_lengths=ilens[bidx : bidx + 1], + zero_infinity=True, + ) + loss = loss / B + return loss + + +class MelSpectrogramLoss(torch.nn.Module): + """Mel-spectrogram loss.""" + + def __init__( + self, + fs: int = 22050, + n_fft: int = 1024, + hop_length: int = 256, + win_length: Optional[int] = None, + window: str = "hann", + n_mels: int = 80, + fmin: Optional[int] = 0, + fmax: Optional[int] = None, + center: bool = True, + normalized: bool = False, + onesided: bool = True, + htk: bool = False, + ): + """Initialize Mel-spectrogram loss. + + Args: + fs (int): Sampling rate. + n_fft (int): FFT points. + hop_length (int): Hop length. + win_length (Optional[int]): Window length. + window (str): Window type. + n_mels (int): Number of Mel basis. + fmin (Optional[int]): Minimum frequency for Mel. + fmax (Optional[int]): Maximum frequency for Mel. + center (bool): Whether to use center window. + normalized (bool): Whether to use normalized one. + onesided (bool): Whether to use oneseded one. + + """ + super().__init__() + + self.fs = fs + self.n_fft = n_fft + self.hop_length = hop_length + self.win_length = n_fft + self.window = window + self.n_mels = n_mels + self.fmin = 0 if fmin is None else fmin + self.fmax = fs / 2 if fmax is None else fmax + self.center = center + self.normalized = normalized + self.onesided = onesided + self.htk = htk + + def logmel(self, feat, ilens): + mel_options = dict( + sr=self.fs, + n_fft=self.n_fft, + n_mels=self.n_mels, + fmin=self.fmin, + fmax=self.fmax, + htk=self.htk, + ) + melmat = librosa.filters.mel(**mel_options) + melmat = torch.from_numpy(melmat.T).float().to(feat.device) + mel_feat = torch.matmul(feat, melmat) + mel_feat = torch.clamp(mel_feat, min=1e-10) + logmel_feat = mel_feat.log10() + + # Zero padding + if ilens is not None: + logmel_feat = logmel_feat.masked_fill( + make_pad_mask(ilens, logmel_feat, 1), 0.0 + ) + else: + ilens = feat.new_full( + [feat.size(0)], fill_value=feat.size(1), dtype=torch.long + ) + return logmel_feat + + def wav_to_mel(self, input, input_lengths=None): + if self.window is not None: + window_func = getattr(torch, f"{self.window}_window") + window = window_func( + self.win_length, dtype=input.dtype, device=input.device + ) + + stft_kwargs = dict( + n_fft=self.n_fft, + win_length=self.win_length, + hop_length=self.hop_length, + center=self.center, + window=window, + normalized=self.normalized, + onesided=self.onesided, + return_complex=True, + ) + + bs = input.size(0) + if input.dim() == 3: + multi_channel = True + # input: (Batch, Nsample, Channels) -> (Batch * Channels, Nsample) + input = input.transpose(1, 2).reshape(-1, input.size(1)) + else: + multi_channel = False + + input_stft = torch.stft(input, **stft_kwargs) + input_stft = torch.view_as_real(input_stft) + input_stft = input_stft.transpose(1, 2) + if multi_channel: + input_stft = input_stft.view( + bs, -1, input_stft.size(1), input_stft.size(2), 2 + ).transpose(1, 2) + if input_lengths is not None: + if self.center: + pad = self.n_fft // 2 + input_lengths = input_lengths + 2 * pad + + feats_lens = (input_lengths - self.n_fft) // self.hop_length + 1 + input_stft.masked_fill_(make_pad_mask(feats_lens, input_stft, 1), 0.0) + else: + feats_lens = None + input_power = input_stft[..., 0] ** 2 + input_stft[..., 1] ** 2 + input_amp = torch.sqrt(torch.clamp(input_power, min=1.0e-10)) + input_feats = self.logmel(input_amp, feats_lens) + return input_feats, feats_lens + + def forward( + self, + y_hat: torch.Tensor, + y: torch.Tensor, + ) -> torch.Tensor: + mel_hat, _ = self.wav_to_mel(y_hat.squeeze(1)) + mel, _ = self.wav_to_mel(y.squeeze(1)) + mel_loss = F.l1_loss(mel_hat, mel) + + return mel_loss + + +class GeneratorLoss(nn.Module): + """The total loss of the generator""" + + def __init__(self, cfg): + super().__init__() + self.cfg = cfg + + self.mel_loss = MelSpectrogramLoss() + self.generator_adv_loss = GeneratorAdversarialLoss() + self.feat_match_loss = FeatureMatchLoss() + self.var_loss = VarianceLoss() + self.forwardsum_loss = ForwardSumLoss() + + self.lambda_adv = 1.0 + self.lambda_mel = 45.0 + self.lambda_feat_match = 2.0 + self.lambda_var = 1.0 + self.lambda_align = 2.0 + + def forward(self, outputs_g, outputs_d, speech_): + loss_g = {} + + # parse generator output + ( + speech_hat_, + bin_loss, + log_p_attn, + start_idxs, + d_outs, + ds, + p_outs, + ps, + e_outs, + es, + text_lengths, + feats_lengths, + ) = outputs_g + + # parse discriminator output + (p_hat, p) = outputs_d + + # calculate losses + mel_loss = self.mel_loss(speech_hat_, speech_) + adv_loss = self.generator_adv_loss(p_hat) + feat_match_loss = self.feat_match_loss(p_hat, p) + dur_loss, pitch_loss, energy_loss = self.var_loss( + d_outs, ds, p_outs, ps, e_outs, es, text_lengths + ) + forwardsum_loss = self.forwardsum_loss(log_p_attn, text_lengths, feats_lengths) + + # calculate total loss + mel_loss = mel_loss * self.lambda_mel + loss_g["mel_loss"] = mel_loss + adv_loss = adv_loss * self.lambda_adv + loss_g["adv_loss"] = adv_loss + feat_match_loss = feat_match_loss * self.lambda_feat_match + loss_g["feat_match_loss"] = feat_match_loss + g_loss = mel_loss + adv_loss + feat_match_loss + loss_g["g_loss"] = g_loss + var_loss = (dur_loss + pitch_loss + energy_loss) * self.lambda_var + loss_g["var_loss"] = var_loss + align_loss = (forwardsum_loss + bin_loss) * self.lambda_align + loss_g["align_loss"] = align_loss + + g_total_loss = g_loss + var_loss + align_loss + + loss_g["g_total_loss"] = g_total_loss + + return loss_g + + +class DiscriminatorAdversarialLoss(torch.nn.Module): + """Discriminator adversarial loss module.""" + + def __init__( + self, + average_by_discriminators: bool = True, + loss_type: str = "mse", + ): + """Initialize DiscriminatorAversarialLoss module. + + Args: + average_by_discriminators (bool): Whether to average the loss by + the number of discriminators. + loss_type (str): Loss type, "mse" or "hinge". + + """ + super().__init__() + self.average_by_discriminators = average_by_discriminators + assert loss_type in ["mse", "hinge"], f"{loss_type} is not supported." + if loss_type == "mse": + self.fake_criterion = self._mse_fake_loss + self.real_criterion = self._mse_real_loss + else: + self.fake_criterion = self._hinge_fake_loss + self.real_criterion = self._hinge_real_loss + + def forward( + self, + outputs_hat: Union[List[List[torch.Tensor]], List[torch.Tensor], torch.Tensor], + outputs: Union[List[List[torch.Tensor]], List[torch.Tensor], torch.Tensor], + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Calcualate discriminator adversarial loss. + + Args: + outputs_hat (Union[List[List[Tensor]], List[Tensor], Tensor]): Discriminator + outputs, list of discriminator outputs, or list of list of discriminator + outputs calculated from generator. + outputs (Union[List[List[Tensor]], List[Tensor], Tensor]): Discriminator + outputs, list of discriminator outputs, or list of list of discriminator + outputs calculated from groundtruth. + + Returns: + Tensor: Discriminator real loss value. + Tensor: Discriminator fake loss value. + + """ + if isinstance(outputs, (tuple, list)): + real_loss = 0.0 + fake_loss = 0.0 + for i, (outputs_hat_, outputs_) in enumerate(zip(outputs_hat, outputs)): + if isinstance(outputs_hat_, (tuple, list)): + # NOTE(kan-bayashi): case including feature maps + outputs_hat_ = outputs_hat_[-1] + outputs_ = outputs_[-1] + real_loss += self.real_criterion(outputs_) + fake_loss += self.fake_criterion(outputs_hat_) + if self.average_by_discriminators: + fake_loss /= i + 1 + real_loss /= i + 1 + else: + real_loss = self.real_criterion(outputs) + fake_loss = self.fake_criterion(outputs_hat) + + return real_loss, fake_loss + + def _mse_real_loss(self, x: torch.Tensor) -> torch.Tensor: + return F.mse_loss(x, x.new_ones(x.size())) + + def _mse_fake_loss(self, x: torch.Tensor) -> torch.Tensor: + return F.mse_loss(x, x.new_zeros(x.size())) + + def _hinge_real_loss(self, x: torch.Tensor) -> torch.Tensor: + return -torch.mean(torch.min(x - 1, x.new_zeros(x.size()))) + + def _hinge_fake_loss(self, x: torch.Tensor) -> torch.Tensor: + return -torch.mean(torch.min(-x - 1, x.new_zeros(x.size()))) + + +class DiscriminatorLoss(torch.nn.Module): + """The total loss of the discriminator""" + + def __init__(self, cfg): + super(DiscriminatorLoss, self).__init__() + self.cfg = cfg + self.discriminator = MultiScaleMultiPeriodDiscriminator() + self.discriminator_adv_loss = DiscriminatorAdversarialLoss() + + def forward(self, speech_real, speech_generated): + loss_d = {} + + real_loss, fake_loss = self.discriminator_adv_loss( + speech_generated, speech_real + ) + loss_d["loss_disc_all"] = real_loss + fake_loss + + return loss_d diff --git a/models/tts/jets/jets_trainer.py b/models/tts/jets/jets_trainer.py new file mode 100644 index 0000000..925dc92 --- /dev/null +++ b/models/tts/jets/jets_trainer.py @@ -0,0 +1,361 @@ +# Copyright (c) 2024 Amphion. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +import torch +import torch.nn as nn +import os +from utils.io import save_audio +from tqdm import tqdm +from models.tts.base import TTSTrainer +from models.tts.jets.jets import Jets +from models.tts.jets.jets_loss import GeneratorLoss, DiscriminatorLoss +from models.tts.jets.jets_dataset import JetsDataset, JetsCollator +from optimizer.optimizers import NoamLR +from torch.optim.lr_scheduler import ExponentialLR +from models.vocoders.gan.discriminator.mpd import MultiScaleMultiPeriodDiscriminator + + +def get_segments( + x: torch.Tensor, + start_idxs: torch.Tensor, + segment_size: int, +) -> torch.Tensor: + """Get segments. + + Args: + x (Tensor): Input tensor (B, C, T). + start_idxs (Tensor): Start index tensor (B,). + segment_size (int): Segment size. + + Returns: + Tensor: Segmented tensor (B, C, segment_size). + + """ + b, c, t = x.size() + segments = x.new_zeros(b, c, segment_size) + for i, start_idx in enumerate(start_idxs): + segments[i] = x[i, :, start_idx : start_idx + segment_size] + return segments + + +class JetsTrainer(TTSTrainer): + def __init__(self, args, cfg): + TTSTrainer.__init__(self, args, cfg) + self.cfg = cfg + + def _build_dataset(self): + return JetsDataset, JetsCollator + + def __build_scheduler(self): + return NoamLR(self.optimizer, **self.cfg.train.lr_scheduler) + + def _write_summary( + self, + losses, + stats, + images={}, + audios={}, + audio_sampling_rate=24000, + tag="train", + ): + for key, value in losses.items(): + self.sw.add_scalar(tag + "/" + key, value, self.step) + self.sw.add_scalar( + "learning_rate", + self.optimizer["optimizer_g"].param_groups[0]["lr"], + self.step, + ) + + if len(images) != 0: + for key, value in images.items(): + self.sw.add_image(key, value, self.global_step, batchformats="HWC") + if len(audios) != 0: + for key, value in audios.items(): + self.sw.add_audio(key, value, self.global_step, audio_sampling_rate) + + for key, value in losses.items(): + self.sw.add_scalar("train/" + key, value, self.step) + lr = self.optimizer.state_dict()["param_groups"][0]["lr"] + self.sw.add_scalar("learning_rate", lr, self.step) + + def _write_valid_summary( + self, losses, stats, images={}, audios={}, audio_sampling_rate=24000, tag="val" + ): + for key, value in losses.items(): + self.sw.add_scalar(tag + "/" + key, value, self.step) + + if len(images) != 0: + for key, value in images.items(): + self.sw.add_image(key, value, self.global_step, batchformats="HWC") + if len(audios) != 0: + for key, value in audios.items(): + self.sw.add_audio(key, value, self.global_step, audio_sampling_rate) + + def _build_criterion(self): + criterion = { + "generator": GeneratorLoss(self.cfg), + "discriminator": DiscriminatorLoss(self.cfg), + } + return criterion + + def get_state_dict(self): + state_dict = { + "generator": self.model["generator"].state_dict(), + "discriminator": self.model["discriminator"].state_dict(), + "optimizer_g": self.optimizer["optimizer_g"].state_dict(), + "optimizer_d": self.optimizer["optimizer_d"].state_dict(), + "scheduler_g": self.scheduler["scheduler_g"].state_dict(), + "scheduler_d": self.scheduler["scheduler_d"].state_dict(), + "step": self.step, + "epoch": self.epoch, + "batch_size": self.cfg.train.batch_size, + } + return state_dict + + def _build_optimizer(self): + optimizer_g = torch.optim.AdamW( + self.model["generator"].parameters(), + self.cfg.train.learning_rate, + betas=self.cfg.train.AdamW.betas, + eps=self.cfg.train.AdamW.eps, + ) + optimizer_d = torch.optim.AdamW( + self.model["discriminator"].parameters(), + self.cfg.train.learning_rate, + betas=self.cfg.train.AdamW.betas, + eps=self.cfg.train.AdamW.eps, + ) + optimizer = {"optimizer_g": optimizer_g, "optimizer_d": optimizer_d} + + return optimizer + + def _build_scheduler(self): + scheduler_g = ExponentialLR( + self.optimizer["optimizer_g"], + gamma=self.cfg.train.lr_decay, + last_epoch=self.epoch - 1, + ) + scheduler_d = ExponentialLR( + self.optimizer["optimizer_d"], + gamma=self.cfg.train.lr_decay, + last_epoch=self.epoch - 1, + ) + + scheduler = {"scheduler_g": scheduler_g, "scheduler_d": scheduler_d} + return scheduler + + def _build_model(self): + net_g = Jets(self.cfg) + net_d = MultiScaleMultiPeriodDiscriminator() + self.model = {"generator": net_g, "discriminator": net_d} + return self.model + + def _train_epoch(self): + r"""Training epoch. Should return average loss of a batch (sample) over + one epoch. See ``train_loop`` for usage. + """ + self.model["generator"].train() + self.model["discriminator"].train() + epoch_sum_loss: float = 0.0 + epoch_losses: dict = {} + epoch_step: int = 0 + for batch in tqdm( + self.train_dataloader, + desc=f"Training Epoch {self.epoch}", + unit="batch", + colour="GREEN", + leave=False, + dynamic_ncols=True, + smoothing=0.04, + disable=not self.accelerator.is_main_process, + ): + with self.accelerator.accumulate(self.model): + if batch["target_len"].min() < self.cfg.train.segment_size: + continue + total_loss, train_losses, training_stats = self._train_step(batch) + self.batch_count += 1 + + if self.batch_count % self.cfg.train.gradient_accumulation_step == 0: + epoch_sum_loss += total_loss + for key, value in train_losses.items(): + if key not in epoch_losses.keys(): + epoch_losses[key] = value + else: + epoch_losses[key] += value + + self.accelerator.log( + { + "Step/Train {} Loss".format(key): value, + }, + step=self.step, + ) + + self.step += 1 + epoch_step += 1 + + self.accelerator.wait_for_everyone() + + epoch_sum_loss = ( + epoch_sum_loss + / len(self.train_dataloader) + * self.cfg.train.gradient_accumulation_step + ) + + for key in epoch_losses.keys(): + epoch_losses[key] = ( + epoch_losses[key] + / len(self.train_dataloader) + * self.cfg.train.gradient_accumulation_step + ) + + return epoch_sum_loss, epoch_losses + + def _train_step(self, batch): + train_losses = {} + total_loss = 0 + training_stats = {} + + # Train Discriminator + # Generator output + outputs_g = self.model["generator"](batch) + speech_hat_, _, _, start_idxs, *_ = outputs_g + + # Discriminator output + speech = batch["audio"].unsqueeze(1) + upsample_factor = self.cfg.train.upsample_factor + speech_ = get_segments( + x=speech, + start_idxs=start_idxs * upsample_factor, + segment_size=self.cfg.train.segment_size * upsample_factor, + ) + p_hat = self.model["discriminator"](speech_hat_.detach()) + p = self.model["discriminator"](speech_) + + # Discriminator loss + loss_d = self.criterion["discriminator"](p, p_hat) + train_losses.update(loss_d) + + # BP and Grad Updated + self.optimizer["optimizer_d"].zero_grad() + self.accelerator.backward(loss_d["loss_disc_all"]) + self.optimizer["optimizer_d"].step() + + # Train Generator + p_hat = self.model["discriminator"](speech_hat_) + with torch.no_grad(): + p = self.model["discriminator"](speech_) + + outputs_d = (p_hat, p) + loss_g = self.criterion["generator"](outputs_g, outputs_d, speech_) + train_losses.update(loss_g) + + # BP and Grad Updated + self.optimizer["optimizer_g"].zero_grad() + self.accelerator.backward(loss_g["g_total_loss"]) + self.optimizer["optimizer_g"].step() + + for item in train_losses: + train_losses[item] = train_losses[item].item() + + total_loss = loss_g["g_total_loss"] + loss_d["loss_disc_all"] + + return ( + total_loss.item(), + train_losses, + training_stats, + ) + + @torch.inference_mode() + def _valid_step(self, batch): + valid_losses = {} + total_loss = 0 + valid_stats = {} + + # Discriminator + # Generator output + outputs_g = self.model["generator"](batch) + speech_hat_, _, _, start_idxs, *_ = outputs_g + + # Discriminator output + speech = batch["audio"].unsqueeze(1) + upsample_factor = self.cfg.train.upsample_factor + speech_ = get_segments( + x=speech, + start_idxs=start_idxs * upsample_factor, + segment_size=self.cfg.train.segment_size * upsample_factor, + ) + p_hat = self.model["discriminator"](speech_hat_.detach()) + p = self.model["discriminator"](speech_) + + # Discriminator loss + loss_d = self.criterion["discriminator"](p, p_hat) + valid_losses.update(loss_d) + + # Generator loss + p_hat = self.model["discriminator"](speech_hat_) + with torch.no_grad(): + p = self.model["discriminator"](speech_) + + outputs_d = (p_hat, p) + loss_g = self.criterion["generator"](outputs_g, outputs_d, speech_) + valid_losses.update(loss_g) + + for item in valid_losses: + valid_losses[item] = valid_losses[item].item() + + total_loss = loss_g["g_total_loss"] + loss_d["loss_disc_all"] + + return ( + total_loss.item(), + valid_losses, + valid_stats, + ) + + @torch.inference_mode() + def _valid_epoch(self): + r"""Testing epoch. Should return average loss of a batch (sample) over + one epoch. See ``train_loop`` for usage. + """ + if isinstance(self.model, dict): + for key in self.model.keys(): + self.model[key].eval() + else: + self.model.eval() + + epoch_sum_loss = 0.0 + epoch_losses = dict() + for batch in tqdm( + self.valid_dataloader, + desc=f"Validating Epoch {self.epoch}", + unit="batch", + colour="GREEN", + leave=False, + dynamic_ncols=True, + smoothing=0.04, + disable=not self.accelerator.is_main_process, + ): + total_loss, valid_losses, valid_stats = self._valid_step(batch) + epoch_sum_loss += total_loss + if isinstance(valid_losses, dict): + for key, value in valid_losses.items(): + if key not in epoch_losses.keys(): + epoch_losses[key] = value + else: + epoch_losses[key] += value + + self.accelerator.log( + { + "Step/Valid {} Loss".format(key): value, + }, + step=self.step, + ) + + epoch_sum_loss = epoch_sum_loss / len(self.valid_dataloader) + for key in epoch_losses.keys(): + epoch_losses[key] = epoch_losses[key] / len(self.valid_dataloader) + + self.accelerator.wait_for_everyone() + + return epoch_sum_loss, epoch_losses diff --git a/models/tts/jets/length_regulator.py b/models/tts/jets/length_regulator.py new file mode 100644 index 0000000..ce30f33 --- /dev/null +++ b/models/tts/jets/length_regulator.py @@ -0,0 +1,49 @@ +# Copyright (c) 2024 Amphion. +# +# This code is modified from https://github.com/imdanboy/jets/blob/main/espnet2/gan_tts/jets/length_regulator.py +# Licensed under Apache License 2.0 + +import torch + + +class GaussianUpsampling(torch.nn.Module): + """ + Gaussian upsampling with fixed temperature as in: + https://arxiv.org/abs/2010.04301 + """ + + def __init__(self, delta=0.1): + super().__init__() + self.delta = delta + + def forward(self, hs, ds, h_masks=None, d_masks=None): + """ + Args: + hs (Tensor): Batched hidden state to be expanded (B, T_text, adim) + ds (Tensor): Batched token duration (B, T_text) + h_masks (Tensor): Mask tensor (B,T_feats) + d_masks (Tensor): Mask tensor (B,T_text) + Returns: + Tensor: Expanded hidden state (B, T_feat, adim) + """ + B = ds.size(0) + device = ds.device + + if h_masks is None: + T_feats = ds.sum().int() + else: + T_feats = h_masks.size(-1) + t = torch.arange(0, T_feats).unsqueeze(0).repeat(B, 1).to(device).float() + if h_masks is not None: + t = t * h_masks.float() + + c = ds.cumsum(dim=-1) - ds / 2 + energy = -1 * self.delta * (t.unsqueeze(-1) - c.unsqueeze(1)) ** 2 + if d_masks is not None: + energy = energy.masked_fill( + ~(d_masks.unsqueeze(1).repeat(1, T_feats, 1)), -float("inf") + ) + + p_attn = torch.softmax(energy, dim=2) # (B, T_feats, T_text) + hs = torch.matmul(p_attn, hs) + return hs diff --git a/models/tts/maskgct/README.md b/models/tts/maskgct/README.md new file mode 100644 index 0000000..811ba20 --- /dev/null +++ b/models/tts/maskgct/README.md @@ -0,0 +1,189 @@ +## MaskGCT: Zero-Shot Text-to-Speech with Masked Generative Codec Transformer + +[](https://arxiv.org/abs/2409.00750) +[](https://huggingface.co/amphion/maskgct) +[](https://huggingface.co/spaces/amphion/maskgct) +[](../../../models/tts/maskgct/README.md) + +[正式版公测地址（趣丸千音）](https://voice.funnycp.com/) + +## Overview + +MaskGCT (**Mask**ed **G**enerative **C**odec **T**ransformer) is *a fully non-autoregressive TTS model that eliminates the need for explicit alignment information between text and speech supervision, as well as phone-level duration prediction*. MaskGCT is a two-stage model: in the first stage, the model uses text to predict semantic tokens extracted from a speech self-supervised learning (SSL) model, and in the second stage, the model predicts acoustic tokens conditioned on these semantic tokens. MaskGCT follows the *mask-and-predict* learning paradigm. During training, MaskGCT learns to predict masked semantic or acoustic tokens based on given conditions and prompts. During inference, the model generates tokens of a specified length in a parallel manner. Experiments with 100K hours of in-the-wild speech demonstrate that MaskGCT outperforms the current state-of-the-art zero-shot TTS systems in terms of quality, similarity, and intelligibility. Audio samples are available at [demo page](https://maskgct.github.io/). + +
+
+
+
+
++ +## News + +- **2024/10/19**: We release **MaskGCT**, a fully non-autoregressive TTS model that eliminates the need for explicit alignment information between text and speech supervision. MaskGCT is trained on [Emilia](https://huggingface.co/datasets/amphion/Emilia-Dataset) dataset and achieves SOTA zero-shot TTS perfermance. + +## Quickstart + +**Clone and install** + +```bash +git clone https://github.com/open-mmlab/Amphion.git +# create env +bash ./models/tts/maskgct/env.sh +``` + +**Model download** + +We provide the following pretrained checkpoints: + + +| Model Name | Description | +|-------------------|-------------| +| [Semantic Codec](https://huggingface.co/amphion/MaskGCT/tree/main/semantic_codec) | Converting speech to semantic tokens. | +| [Acoustic Codec](https://huggingface.co/amphion/MaskGCT/tree/main/acoustic_codec) | Converting speech to acoustic tokens and reconstructing waveform from acoustic tokens. | +| [MaskGCT-T2S](https://huggingface.co/amphion/MaskGCT/tree/main/t2s_model) | Predicting semantic tokens with text and prompt semantic tokens. | +| [MaskGCT-S2A](https://huggingface.co/amphion/MaskGCT/tree/main/s2a_model) | Predicts acoustic tokens conditioned on semantic tokens. | + +You can download all pretrained checkpoints from [HuggingFace](https://huggingface.co/amphion/MaskGCT/tree/main) or use huggingface api. + +```python +from huggingface_hub import hf_hub_download + +# download semantic codec ckpt +semantic_code_ckpt = hf_hub_download("amphion/MaskGCT", filename="semantic_codec/model.safetensors") + +# download acoustic codec ckpt +codec_encoder_ckpt = hf_hub_download("amphion/MaskGCT", filename="acoustic_codec/model.safetensors") +codec_decoder_ckpt = hf_hub_download("amphion/MaskGCT", filename="acoustic_codec/model_1.safetensors") + +# download t2s model ckpt +t2s_model_ckpt = hf_hub_download("amphion/MaskGCT", filename="t2s_model/model.safetensors") + +# download s2a model ckpt +s2a_1layer_ckpt = hf_hub_download("amphion/MaskGCT", filename="s2a_model/s2a_model_1layer/model.safetensors") +s2a_full_ckpt = hf_hub_download("amphion/MaskGCT", filename="s2a_model/s2a_model_full/model.safetensors") +``` + +**Basic Usage** + +You can use the following code to generate speech from text and a prompt speech (the code is also provided in [inference.py](../../../models/tts/maskgct/maskgct_inference.py)). + +```python +from models.tts.maskgct.maskgct_utils import * +from huggingface_hub import hf_hub_download +import safetensors +import soundfile as sf + +if __name__ == "__main__": + + # build model + device = torch.device("cuda:0") + cfg_path = "./models/tts/maskgct/config/maskgct.json" + cfg = load_config(cfg_path) + # 1. build semantic model (w2v-bert-2.0) + semantic_model, semantic_mean, semantic_std = build_semantic_model(device) + # 2. build semantic codec + semantic_codec = build_semantic_codec(cfg.model.semantic_codec, device) + # 3. build acoustic codec + codec_encoder, codec_decoder = build_acoustic_codec(cfg.model.acoustic_codec, device) + # 4. build t2s model + t2s_model = build_t2s_model(cfg.model.t2s_model, device) + # 5. build s2a model + s2a_model_1layer = build_s2a_model(cfg.model.s2a_model.s2a_1layer, device) + s2a_model_full = build_s2a_model(cfg.model.s2a_model.s2a_full, device) + + # download checkpoint + ... + + # load semantic codec + safetensors.torch.load_model(semantic_codec, semantic_code_ckpt) + # load acoustic codec + safetensors.torch.load_model(codec_encoder, codec_encoder_ckpt) + safetensors.torch.load_model(codec_decoder, codec_decoder_ckpt) + # load t2s model + safetensors.torch.load_model(t2s_model, t2s_model_ckpt) + # load s2a model + safetensors.torch.load_model(s2a_model_1layer, s2a_1layer_ckpt) + safetensors.torch.load_model(s2a_model_full, s2a_full_ckpt) + + # inference + prompt_wav_path = "./models/tts/maskgct/wav/prompt.wav" + save_path = "[YOUR SAVE PATH]" + prompt_text = " We do not break. We never give in. We never back down." + target_text = "In this paper, we introduce MaskGCT, a fully non-autoregressive TTS model that eliminates the need for explicit alignment information between text and speech supervision." + # Specify the target duration (in seconds). If target_len = None, we use a simple rule to predict the target duration. + target_len = 18 + + maskgct_inference_pipeline = MaskGCT_Inference_Pipeline( + semantic_model, + semantic_codec, + codec_encoder, + codec_decoder, + t2s_model, + s2a_model_1layer, + s2a_model_full, + semantic_mean, + semantic_std, + device, + ) + + recovered_audio = maskgct_inference_pipeline.maskgct_inference( + prompt_wav_path, prompt_text, target_text, "en", "en", target_len=target_len + ) + sf.write(save_path, recovered_audio, 24000) +``` + +**Jupyter Notebook** + +We also provide a [jupyter notebook](../../../models/tts/maskgct/maskgct_demo.ipynb) to show more details of MaskGCT inference. + +## Training Dataset + +We use the [Emilia](https://huggingface.co/datasets/amphion/Emilia-Dataset) dataset to train our models. Emilia is a multilingual and diverse in-the-wild speech dataset designed for large-scale speech generation. In this work, we use English and Chinese data from Emilia, each with 50K hours of speech (totaling 100K hours). + + +## Evaluation Results of MaskGCT + +| System | SIM-O↑ | WER↓ | FSD↓ | SMOS↑ | CMOS↑ | +| :--- | :---: | :---: | :---: | :---: | :---: | +| | | **LibriSpeech test-clean** | +| Ground Truth | 0.68 | 1.94 | | 4.05±0.12 | 0.00 | +| VALL-E | 0.50 | 5.90 | - | 3.47 ±0.26 | -0.52±0.22 | +| VoiceBox | 0.64 | 2.03 | 0.762 | 3.80±0.17 | -0.41±0.13 | +| NaturalSpeech 3 | 0.67 | 1.94 | 0.786 | 4.26±0.10 | 0.16±0.14 | +| VoiceCraft | 0.45 | 4.68 | 0.981 | 3.52±0.21 | -0.33 ±0.16 | +| XTTS-v2 | 0.51 | 4.20 | 0.945 | 3.02±0.22 | -0.98 ±0.19 | +| MaskGCT | 0.687(0.723) | 2.634(1.976) | 0.886 | 4.27±0.14 | 0.10±0.16 | +| MaskGCT(gt length) | 0.697 | 2.012 | 0.746 | 4.33±0.11 | 0.13±0.13 | +| | | **SeedTTS test-en** | +| Ground Truth | 0.730 | 2.143 | | 3.92±0.15 | 0.00 | +| CosyVoice | 0.643 | 4.079 | 0.316 | 3.52±0.17 | -0.41 ±0.18 | +| XTTS-v2 | 0.463 | 3.248 | 0.484 | 3.15±0.22 | -0.86±0.19 | +| VoiceCraft | 0.470 | 7.556 | 0.226 | 3.18±0.20 | -1.08 ±0.15 | +| MaskGCT | 0.717(0.760) | 2.623(1.283) | 0.188 | 4.24 ±0.12 | 0.03 ±0.14 | +| MaskGCT(gt length) | 0.728 | 2.466 | 0.159 | 4.13 ±0.17 | 0.12 ±0.15 | +| | | **SeedTTS test-zh** | +| Ground Truth | 0.750 | 1.254 | | 3.86 ±0.17 | 0.00 | +| CosyVoice | 0.750 | 4.089 | 0.276 | 3.54 ±0.12 | -0.45 ±0.15 | +| XTTS-v2 | 0.635 | 2.876 | 0.413 | 2.95 ±0.18 | -0.81 ±0.22 | +| MaskGCT | 0.774(0.805) | 2.273(0.843) | 0.106 | 4.09 ±0.12 | 0.05 ±0.17 | +| MaskGCT(gt length) | 0.777 | 2.183 | 0.101 | 4.11 ±0.12 | 0.08±0.18 | + +## Citations + +If you use MaskGCT in your research, please cite the following paper: + +```bibtex +@article{wang2024maskgct, + title={MaskGCT: Zero-Shot Text-to-Speech with Masked Generative Codec Transformer}, + author={Wang, Yuancheng and Zhan, Haoyue and Liu, Liwei and Zeng, Ruihong and Guo, Haotian and Zheng, Jiachen and Zhang, Qiang and Zhang, Xueyao and Zhang, Shunsi and Wu, Zhizheng}, + journal={arXiv preprint arXiv:2409.00750}, + year={2024} +} + +@inproceedings{amphion, + author={Zhang, Xueyao and Xue, Liumeng and Gu, Yicheng and Wang, Yuancheng and Li, Jiaqi and He, Haorui and Wang, Chaoren and Song, Ting and Chen, Xi and Fang, Zihao and Chen, Haopeng and Zhang, Junan and Tang, Tze Ying and Zou, Lexiao and Wang, Mingxuan and Han, Jun and Chen, Kai and Li, Haizhou and Wu, Zhizheng}, + title={Amphion: An Open-Source Audio, Music and Speech Generation Toolkit}, + booktitle={{IEEE} Spoken Language Technology Workshop, {SLT} 2024}, + year={2024} +} +``` diff --git a/models/tts/maskgct/ckpt/wav2vec2bert_stats.pt b/models/tts/maskgct/ckpt/wav2vec2bert_stats.pt new file mode 100644 index 0000000..bd06a43 Binary files /dev/null and b/models/tts/maskgct/ckpt/wav2vec2bert_stats.pt differ diff --git a/models/tts/maskgct/config/maskgct.json b/models/tts/maskgct/config/maskgct.json new file mode 100644 index 0000000..b4cf5d2 --- /dev/null +++ b/models/tts/maskgct/config/maskgct.json @@ -0,0 +1,81 @@ +{ + "model": { + "t2s_model": { + "hidden_size": 1536, + "num_layers": 16, + "num_heads": 16, + "cfg_scale": 0.15, + "cond_codebook_size": 8192, + "cond_dim": 1024 + }, + "s2a_model": { + "s2a_1layer": { + "num_quantizer": 1, + "hidden_size": 1024, + "num_layers": 16, + "num_heads": 16, + "codebook_size": 1024, + "cfg_scale": 0.15, + "mask_layer_schedule": "linear", + "cond_codebook_size": 8192, + "cond_dim": 1024, + "predict_layer_1": true + }, + "s2a_full": { + "num_quantizer": 12, + "hidden_size": 1024, + "num_layers": 16, + "num_heads": 16, + "codebook_size": 1024, + "cfg_scale": 0.15, + "mask_layer_schedule": "linear", + "cond_codebook_size": 8192, + "cond_dim": 1024, + "predict_layer_1": false + } + }, + "semantic_codec": { + "codebook_size": 8192, + "hidden_size": 1024, + "codebook_dim": 8, + "vocos_dim": 384, + "vocos_intermediate_dim": 2048, + "vocos_num_layers": 12 + }, + "acoustic_codec": { + "encoder": { + "d_model": 96, + "up_ratios": [3, 4, 5, 8], + "out_channels": 256, + "use_tanh": false + }, + "decoder": { + "in_channel": 256, + "upsample_initial_channel": 1536, + "up_ratios": [8, 5, 4, 3], + "num_quantizers": 12, + "codebook_size": 1024, + "codebook_dim": 8, + "quantizer_type": "fvq", + "quantizer_dropout": 0.5, + "commitment": 0.25, + "codebook_loss_weight": 1.0, + "use_l2_normlize": true, + "codebook_type": "euclidean", + "kmeans_init": false, + "kmeans_iters": 10, + "decay": 0.8, + "eps": 0.5, + "threshold_ema_dead_code": 2, + "weight_init": false, + "use_vocos": true, + "vocos_dim": 512, + "vocos_intermediate_dim": 4096, + "vocos_num_layers": 30, + "n_fft": 1920, + "hop_size": 480, + "padding": "same" + } + } + } +} \ No newline at end of file diff --git a/models/tts/maskgct/env.sh b/models/tts/maskgct/env.sh new file mode 100644 index 0000000..ed59567 --- /dev/null +++ b/models/tts/maskgct/env.sh @@ -0,0 +1,25 @@ +pip install setuptools ruamel.yaml tqdm +pip install tensorboard tensorboardX torch==2.0.1 +pip install transformers===4.41.1 +pip install -U encodec +pip install black==24.1.1 +pip install oss2 +sudo apt-get install espeak-ng +pip install phonemizer +pip install g2p_en +pip install accelerate==0.31.0 +pip install funasr zhconv zhon modelscope +# pip install git+https://github.com/lhotse-speech/lhotse +pip install timm +pip install jieba cn2an +pip install unidecode +pip install -U cos-python-sdk-v5 +pip install pypinyin +pip install jiwer +pip install omegaconf +pip install pyworld +pip install py3langid==0.2.2 LangSegment +pip install onnxruntime +pip install pyopenjtalk +pip install pykakasi +pip install -U openai-whisper \ No newline at end of file diff --git a/models/tts/maskgct/g2p/g2p/__init__.py b/models/tts/maskgct/g2p/g2p/__init__.py new file mode 100644 index 0000000..ca96b67 --- /dev/null +++ b/models/tts/maskgct/g2p/g2p/__init__.py @@ -0,0 +1,87 @@ +# Copyright (c) 2024 Amphion. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +from models.tts.maskgct.g2p.g2p import cleaners +from tokenizers import Tokenizer +from models.tts.maskgct.g2p.g2p.text_tokenizers import TextTokenizer +import LangSegment +import json +import re + + +class PhonemeBpeTokenizer: + + def __init__(self, vacab_path="./models/tts/maskgct/g2p/g2p/vocab.json"): + self.lang2backend = { + "zh": "cmn", + "ja": "ja", + "en": "en-us", + "fr": "fr-fr", + "ko": "ko", + "de": "de", + } + self.text_tokenizers = {} + self.int_text_tokenizers() + + with open(vacab_path, "r") as f: + json_data = f.read() + data = json.loads(json_data) + self.vocab = data["vocab"] + LangSegment.setfilters(["en", "zh", "ja", "ko", "fr", "de"]) + + def int_text_tokenizers(self): + for key, value in self.lang2backend.items(): + self.text_tokenizers[key] = TextTokenizer(language=value) + + def tokenize(self, text, sentence, language): + + # 1. convert text to phoneme + phonemes = [] + if language == "auto": + seglist = LangSegment.getTexts(text) + tmp_ph = [] + for seg in seglist: + tmp_ph.append( + self._clean_text( + seg["text"], sentence, seg["lang"], ["cjekfd_cleaners"] + ) + ) + phonemes = "|_|".join(tmp_ph) + else: + phonemes = self._clean_text(text, sentence, language, ["cjekfd_cleaners"]) + # print('clean text: ', phonemes) + + # 2. tokenize phonemes + phoneme_tokens = self.phoneme2token(phonemes) + # print('encode: ', phoneme_tokens) + + # # 3. decode tokens [optional] + # decoded_text = self.tokenizer.decode(phoneme_tokens) + # print('decoded: ', decoded_text) + + return phonemes, phoneme_tokens + + def _clean_text(self, text, sentence, language, cleaner_names): + for name in cleaner_names: + cleaner = getattr(cleaners, name) + if not cleaner: + raise Exception("Unknown cleaner: %s" % name) + text = cleaner(text, sentence, language, self.text_tokenizers) + return text + + def phoneme2token(self, phonemes): + tokens = [] + if isinstance(phonemes, list): + for phone in phonemes: + phone = phone.split("\t")[0] + phonemes_split = phone.split("|") + tokens.append( + [self.vocab[p] for p in phonemes_split if p in self.vocab] + ) + else: + phonemes = phonemes.split("\t")[0] + phonemes_split = phonemes.split("|") + tokens = [self.vocab[p] for p in phonemes_split if p in self.vocab] + return tokens diff --git a/models/tts/maskgct/g2p/g2p/chinese_model_g2p.py b/models/tts/maskgct/g2p/g2p/chinese_model_g2p.py new file mode 100644 index 0000000..862ebac --- /dev/null +++ b/models/tts/maskgct/g2p/g2p/chinese_model_g2p.py @@ -0,0 +1,212 @@ +# Copyright (c) 2024 Amphion. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +import os +import numpy as np +import torch +from torch.utils.data import DataLoader +import json +from transformers import BertTokenizer +from torch.utils.data import Dataset +from transformers.models.bert.modeling_bert import * +import torch.nn.functional as F +from onnxruntime import InferenceSession, GraphOptimizationLevel, SessionOptions + + +class PolyDataset(Dataset): + def __init__(self, words, labels, word_pad_idx=0, label_pad_idx=-1): + self.dataset = self.preprocess(words, labels) + self.word_pad_idx = word_pad_idx + self.label_pad_idx = label_pad_idx + + def preprocess(self, origin_sentences, origin_labels): + """ + Maps tokens and tags to their indices and stores them in the dict data. + examples: + word:['[CLS]', '浙', '商', '银', '行', '企', '业', '信', '贷', '部'] + sentence:([101, 3851, 1555, 7213, 6121, 821, 689, 928, 6587, 6956], + array([ 1, 2, 3, 4, 5, 6, 7, 8, 9, 10])) + label:[3, 13, 13, 13, 0, 0, 0, 0, 0] + """ + data = [] + labels = [] + sentences = [] + # tokenize + for line in origin_sentences: + # replace each token by its index + # we can not use encode_plus because our sentences are aligned to labels in list type + words = [] + word_lens = [] + for token in line: + words.append(token) + word_lens.append(1) + token_start_idxs = 1 + np.cumsum([0] + word_lens[:-1]) + sentences.append(((words, token_start_idxs), 0)) + ### + for tag in origin_labels: + labels.append(tag) + + for sentence, label in zip(sentences, labels): + data.append((sentence, label)) + return data + + def __getitem__(self, idx): + """sample data to get batch""" + word = self.dataset[idx][0] + label = self.dataset[idx][1] + return [word, label] + + def __len__(self): + """get dataset size""" + return len(self.dataset) + + def collate_fn(self, batch): + + sentences = [x[0][0] for x in batch] + ori_sents = [x[0][1] for x in batch] + labels = [x[1] for x in batch] + batch_len = len(sentences) + + # compute length of longest sentence in batch + max_len = max([len(s[0]) for s in sentences]) + max_label_len = 0 + batch_data = np.ones((batch_len, max_len)) + batch_label_starts = [] + + # padding and aligning + for j in range(batch_len): + cur_len = len(sentences[j][0]) + batch_data[j][:cur_len] = sentences[j][0] + label_start_idx = sentences[j][-1] + label_starts = np.zeros(max_len) + label_starts[[idx for idx in label_start_idx if idx < max_len]] = 1 + batch_label_starts.append(label_starts) + max_label_len = max(int(sum(label_starts)), max_label_len) + + # padding label + batch_labels = self.label_pad_idx * np.ones((batch_len, max_label_len)) + batch_pmasks = self.label_pad_idx * np.ones((batch_len, max_label_len)) + for j in range(batch_len): + cur_tags_len = len(labels[j]) + batch_labels[j][:cur_tags_len] = labels[j] + batch_pmasks[j][:cur_tags_len] = [ + 1 if item > 0 else 0 for item in labels[j] + ] + + # convert data to torch LongTensors + batch_data = torch.tensor(batch_data, dtype=torch.long) + batch_label_starts = torch.tensor(batch_label_starts, dtype=torch.long) + batch_labels = torch.tensor(batch_labels, dtype=torch.long) + batch_pmasks = torch.tensor(batch_pmasks, dtype=torch.long) + return [batch_data, batch_label_starts, batch_labels, batch_pmasks, ori_sents] + + +class BertPolyPredict: + def __init__(self, bert_model, jsonr_file, json_file): + self.tokenizer = BertTokenizer.from_pretrained(bert_model, do_lower_case=True) + with open(jsonr_file, "r", encoding="utf8") as fp: + self.pron_dict = json.load(fp) + with open(json_file, "r", encoding="utf8") as fp: + self.pron_dict_id_2_pinyin = json.load(fp) + self.num_polyphone = len(self.pron_dict) + self.device = "cpu" + self.polydataset = PolyDataset + options = SessionOptions() # initialize session options + options.graph_optimization_level = GraphOptimizationLevel.ORT_ENABLE_ALL + print(os.path.join(bert_model, "poly_bert_model.onnx")) + self.session = InferenceSession( + os.path.join(bert_model, "poly_bert_model.onnx"), + sess_options=options, + providers=[ + "CUDAExecutionProvider", + "CPUExecutionProvider", + ], # CPUExecutionProvider #CUDAExecutionProvider + ) + # self.session.set_providers(['CUDAExecutionProvider', "CPUExecutionProvider"], [ {'device_id': 0}]) + + # disable session.run() fallback mechanism, it prevents for a reset of the execution provider + self.session.disable_fallback() + + def predict_process(self, txt_list): + word_test, label_test, texts_test = self.get_examples_po(txt_list) + data = self.polydataset(word_test, label_test) + predict_loader = DataLoader( + data, batch_size=1, shuffle=False, collate_fn=data.collate_fn + ) + pred_tags = self.predict_onnx(predict_loader) + return pred_tags + + def predict_onnx(self, dev_loader): + pred_tags = [] + with torch.no_grad(): + for idx, batch_samples in enumerate(dev_loader): + # [batch_data, batch_label_starts, batch_labels, batch_pmasks, ori_sents] + batch_data, batch_label_starts, batch_labels, batch_pmasks, _ = ( + batch_samples + ) + # shift tensors to GPU if available + batch_data = batch_data.to(self.device) + batch_label_starts = batch_label_starts.to(self.device) + batch_labels = batch_labels.to(self.device) + batch_pmasks = batch_pmasks.to(self.device) + batch_data = np.asarray(batch_data, dtype=np.int32) + batch_pmasks = np.asarray(batch_pmasks, dtype=np.int32) + # batch_output = self.session.run(output_names=['outputs'], input_feed={"input_ids":batch_data, "input_pmasks": batch_pmasks})[0][0] + batch_output = self.session.run( + output_names=["outputs"], input_feed={"input_ids": batch_data} + )[0] + label_masks = batch_pmasks == 1 + batch_labels = batch_labels.to("cpu").numpy() + for i, indices in enumerate(np.argmax(batch_output, axis=2)): + for j, idx in enumerate(indices): + if label_masks[i][j]: + # pred_tag.append(idx) + pred_tags.append(self.pron_dict_id_2_pinyin[str(idx + 1)]) + return pred_tags + + def get_examples_po(self, text_list): + + word_list = [] + label_list = [] + sentence_list = [] + id = 0 + for line in [text_list]: + sentence = line[0] + words = [] + tokens = line[0] + index = line[-1] + front = index + back = len(tokens) - index - 1 + labels = [0] * front + [1] + [0] * back + words = ["[CLS]"] + [item for item in sentence] + words = self.tokenizer.convert_tokens_to_ids(words) + word_list.append(words) + label_list.append(labels) + sentence_list.append(sentence) + + id += 1 + # mask_list.append(masks) + assert len(labels) + 1 == len(words), print( + ( + poly, + sentence, + words, + labels, + sentence, + len(sentence), + len(words), + len(labels), + ) + ) + assert len(labels) + 1 == len( + words + ), "Number of labels does not match number of words" + assert len(labels) == len( + sentence + ), "Number of labels does not match number of sentences" + assert len(word_list) == len( + label_list + ), "Number of label sentences does not match number of word sentences" + return word_list, label_list, text_list diff --git a/models/tts/maskgct/g2p/g2p/cleaners.py b/models/tts/maskgct/g2p/g2p/cleaners.py new file mode 100644 index 0000000..b25aabf --- /dev/null +++ b/models/tts/maskgct/g2p/g2p/cleaners.py @@ -0,0 +1,31 @@ +# Copyright (c) 2024 Amphion. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +import re +from models.tts.maskgct.g2p.g2p.japanese import japanese_to_ipa +from models.tts.maskgct.g2p.g2p.mandarin import chinese_to_ipa +from models.tts.maskgct.g2p.g2p.english import english_to_ipa +from models.tts.maskgct.g2p.g2p.french import french_to_ipa +from models.tts.maskgct.g2p.g2p.korean import korean_to_ipa +from models.tts.maskgct.g2p.g2p.german import german_to_ipa + + +def cjekfd_cleaners(text, sentence, language, text_tokenizers): + + if language == "zh": + return chinese_to_ipa(text, sentence, text_tokenizers["zh"]) + elif language == "ja": + return japanese_to_ipa(text, text_tokenizers["ja"]) + elif language == "en": + return english_to_ipa(text, text_tokenizers["en"]) + elif language == "fr": + return french_to_ipa(text, text_tokenizers["fr"]) + elif language == "ko": + return korean_to_ipa(text, text_tokenizers["ko"]) + elif language == "de": + return german_to_ipa(text, text_tokenizers["de"]) + else: + raise Exception("Unknown language: %s" % language) + return None diff --git a/models/tts/maskgct/g2p/g2p/english.py b/models/tts/maskgct/g2p/g2p/english.py new file mode 100644 index 0000000..f8f349f --- /dev/null +++ b/models/tts/maskgct/g2p/g2p/english.py @@ -0,0 +1,202 @@ +# Copyright (c) 2024 Amphion. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +import re +from unidecode import unidecode +import inflect + +""" + Text clean time +""" +_inflect = inflect.engine() +_comma_number_re = re.compile(r"([0-9][0-9\,]+[0-9])") +_decimal_number_re = re.compile(r"([0-9]+\.[0-9]+)") +_percent_number_re = re.compile(r"([0-9\.\,]*[0-9]+%)") +_pounds_re = re.compile(r"£([0-9\,]*[0-9]+)") +_dollars_re = re.compile(r"\$([0-9\.\,]*[0-9]+)") +_fraction_re = re.compile(r"([0-9]+)/([0-9]+)") +_ordinal_re = re.compile(r"[0-9]+(st|nd|rd|th)") +_number_re = re.compile(r"[0-9]+") + +# List of (regular expression, replacement) pairs for abbreviations: +_abbreviations = [ + (re.compile("\\b%s\\b" % x[0], re.IGNORECASE), x[1]) + for x in [ + ("mrs", "misess"), + ("mr", "mister"), + ("dr", "doctor"), + ("st", "saint"), + ("co", "company"), + ("jr", "junior"), + ("maj", "major"), + ("gen", "general"), + ("drs", "doctors"), + ("rev", "reverend"), + ("lt", "lieutenant"), + ("hon", "honorable"), + ("sgt", "sergeant"), + ("capt", "captain"), + ("esq", "esquire"), + ("ltd", "limited"), + ("col", "colonel"), + ("ft", "fort"), + ("etc", "et cetera"), + ("btw", "by the way"), + ] +] + +_special_map = [ + ("t|ɹ", "tɹ"), + ("d|ɹ", "dɹ"), + ("t|s", "ts"), + ("d|z", "dz"), + ("ɪ|ɹ", "ɪɹ"), + ("ɐ", "ɚ"), + ("ᵻ", "ɪ"), + ("əl", "l"), + ("x", "k"), + ("ɬ", "l"), + ("ʔ", "t"), + ("n̩", "n"), + ("oː|ɹ", "oːɹ"), +] + + +def expand_abbreviations(text): + for regex, replacement in _abbreviations: + text = re.sub(regex, replacement, text) + return text + + +def _remove_commas(m): + return m.group(1).replace(",", "") + + +def _expand_decimal_point(m): + return m.group(1).replace(".", " point ") + + +def _expand_percent(m): + return m.group(1).replace("%", " percent ") + + +def _expand_dollars(m): + match = m.group(1) + parts = match.split(".") + if len(parts) > 2: + return " " + match + " dollars " # Unexpected format + dollars = int(parts[0]) if parts[0] else 0 + cents = int(parts[1]) if len(parts) > 1 and parts[1] else 0 + if dollars and cents: + dollar_unit = "dollar" if dollars == 1 else "dollars" + cent_unit = "cent" if cents == 1 else "cents" + return " %s %s, %s %s " % (dollars, dollar_unit, cents, cent_unit) + elif dollars: + dollar_unit = "dollar" if dollars == 1 else "dollars" + return " %s %s " % (dollars, dollar_unit) + elif cents: + cent_unit = "cent" if cents == 1 else "cents" + return " %s %s " % (cents, cent_unit) + else: + return " zero dollars " + + +def fraction_to_words(numerator, denominator): + if numerator == 1 and denominator == 2: + return " one half " + if numerator == 1 and denominator == 4: + return " one quarter " + if denominator == 2: + return " " + _inflect.number_to_words(numerator) + " halves " + if denominator == 4: + return " " + _inflect.number_to_words(numerator) + " quarters " + return ( + " " + + _inflect.number_to_words(numerator) + + " " + + _inflect.ordinal(_inflect.number_to_words(denominator)) + + " " + ) + + +def _expand_fraction(m): + numerator = int(m.group(1)) + denominator = int(m.group(2)) + return fraction_to_words(numerator, denominator) + + +def _expand_ordinal(m): + return " " + _inflect.number_to_words(m.group(0)) + " " + + +def _expand_number(m): + num = int(m.group(0)) + if num > 1000 and num < 3000: + if num == 2000: + return " two thousand " + elif num > 2000 and num < 2010: + return " two thousand " + _inflect.number_to_words(num % 100) + " " + elif num % 100 == 0: + return " " + _inflect.number_to_words(num // 100) + " hundred " + else: + return ( + " " + + _inflect.number_to_words(num, andword="", zero="oh", group=2).replace( + ", ", " " + ) + + " " + ) + else: + return " " + _inflect.number_to_words(num, andword="") + " " + + +# Normalize numbers pronunciation +def normalize_numbers(text): + text = re.sub(_comma_number_re, _remove_commas, text) + text = re.sub(_pounds_re, r"\1 pounds", text) + text = re.sub(_dollars_re, _expand_dollars, text) + text = re.sub(_fraction_re, _expand_fraction, text) + text = re.sub(_decimal_number_re, _expand_decimal_point, text) + text = re.sub(_percent_number_re, _expand_percent, text) + text = re.sub(_ordinal_re, _expand_ordinal, text) + text = re.sub(_number_re, _expand_number, text) + return text + + +def _english_to_ipa(text): + # text = unidecode(text).lower() + text = expand_abbreviations(text) + text = normalize_numbers(text) + return text + + +# special map +def special_map(text): + for regex, replacement in _special_map: + regex = regex.replace("|", "\|") + while re.search(r"(^|[_|]){}([_|]|$)".format(regex), text): + text = re.sub( + r"(^|[_|]){}([_|]|$)".format(regex), r"\1{}\2".format(replacement), text + ) + # text = re.sub(r'([,.!?])', r'|\1', text) + return text + + +# Add some special operation +def english_to_ipa(text, text_tokenizer): + if type(text) == str: + text = _english_to_ipa(text) + else: + text = [_english_to_ipa(t) for t in text] + phonemes = text_tokenizer(text) + if phonemes[-1] in "p⁼ʰmftnlkxʃs`ɹaoəɛɪeɑʊŋiuɥwæjː": + phonemes += "|_" + if type(text) == str: + return special_map(phonemes) + else: + result_ph = [] + for phone in phonemes: + result_ph.append(special_map(phone)) + return result_ph diff --git a/models/tts/maskgct/g2p/g2p/french.py b/models/tts/maskgct/g2p/g2p/french.py new file mode 100644 index 0000000..bd9400c --- /dev/null +++ b/models/tts/maskgct/g2p/g2p/french.py @@ -0,0 +1,149 @@ +# Copyright (c) 2024 Amphion. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +import re + +""" + Text clean time +""" +# List of (regular expression, replacement) pairs for abbreviations in french: +_abbreviations = [ + (re.compile("\\b%s\\." % x[0], re.IGNORECASE), x[1]) + for x in [ + ("M", "monsieur"), + ("Mlle", "mademoiselle"), + ("Mlles", "mesdemoiselles"), + ("Mme", "Madame"), + ("Mmes", "Mesdames"), + ("N.B", "nota bene"), + ("M", "monsieur"), + ("p.c.q", "parce que"), + ("Pr", "professeur"), + ("qqch", "quelque chose"), + ("rdv", "rendez-vous"), + ("max", "maximum"), + ("min", "minimum"), + ("no", "numéro"), + ("adr", "adresse"), + ("dr", "docteur"), + ("st", "saint"), + ("co", "companie"), + ("jr", "junior"), + ("sgt", "sergent"), + ("capt", "capitain"), + ("col", "colonel"), + ("av", "avenue"), + ("av. J.-C", "avant Jésus-Christ"), + ("apr. J.-C", "après Jésus-Christ"), + ("art", "article"), + ("boul", "boulevard"), + ("c.-à-d", "c’est-à-dire"), + ("etc", "et cetera"), + ("ex", "exemple"), + ("excl", "exclusivement"), + ("boul", "boulevard"), + ] +] + [ + (re.compile("\\b%s" % x[0]), x[1]) + for x in [ + ("Mlle", "mademoiselle"), + ("Mlles", "mesdemoiselles"), + ("Mme", "Madame"), + ("Mmes", "Mesdames"), + ] +] + +rep_map = { + "：": ",", + "；": ",", + "，": ",", + "。": ".", + "！": "!", + "？": "?", + "\n": ".", + "·": ",", + "、": ",", + "...": ".", + "…": ".", + "$": ".", + "“": "", + "”": "", + "‘": "", + "’": "", + "（": "", + "）": "", + "(": "", + ")": "", + "《": "", + "》": "", + "【": "", + "】": "", + "[": "", + "]": "", + "—": "", + "～": "-", + "~": "-", + "「": "", + "」": "", + "¿": "", + "¡": "", +} + + +def collapse_whitespace(text): + # Regular expression matching whitespace: + _whitespace_re = re.compile(r"\s+") + return re.sub(_whitespace_re, " ", text).strip() + + +def remove_punctuation_at_begin(text): + return re.sub(r"^[,.!?]+", "", text) + + +def remove_aux_symbols(text): + text = re.sub(r"[\<\>\(\)\[\]\"\«\»]+", "", text) + return text + + +def replace_symbols(text): + text = text.replace(";", ",") + text = text.replace("-", " ") + text = text.replace(":", ",") + text = text.replace("&", " et ") + return text + + +def expand_abbreviations(text): + for regex, replacement in _abbreviations: + text = re.sub(regex, replacement, text) + return text + + +def replace_punctuation(text): + pattern = re.compile("|".join(re.escape(p) for p in rep_map.keys())) + replaced_text = pattern.sub(lambda x: rep_map[x.group()], text) + return replaced_text + + +def text_normalize(text): + text = expand_abbreviations(text) + text = replace_punctuation(text) + text = replace_symbols(text) + text = remove_aux_symbols(text) + text = remove_punctuation_at_begin(text) + text = collapse_whitespace(text) + text = re.sub(r"([^\.,!\?\-…])$", r"\1", text) + return text + + +def french_to_ipa(text, text_tokenizer): + if type(text) == str: + text = text_normalize(text) + phonemes = text_tokenizer(text) + return phonemes + else: + for i, t in enumerate(text): + text[i] = text_normalize(t) + return text_tokenizer(text) diff --git a/models/tts/maskgct/g2p/g2p/german.py b/models/tts/maskgct/g2p/g2p/german.py new file mode 100644 index 0000000..bd82eea --- /dev/null +++ b/models/tts/maskgct/g2p/g2p/german.py @@ -0,0 +1,94 @@ +# Copyright (c) 2024 Amphion. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +import re + +""" + Text clean time +""" +rep_map = { + "：": ",", + "；": ",", + "，": ",", + "。": ".", + "！": "!", + "？": "?", + "\n": ".", + "·": ",", + "、": ",", + "...": ".", + "…": ".", + "$": ".", + "“": "", + "”": "", + "‘": "", + "’": "", + "（": "", + "）": "", + "(": "", + ")": "", + "《": "", + "》": "", + "【": "", + "】": "", + "[": "", + "]": "", + "—": "", + "～": "-", + "~": "-", + "「": "", + "」": "", + "¿": "", + "¡": "", +} + + +def collapse_whitespace(text): + # Regular expression matching whitespace: + _whitespace_re = re.compile(r"\s+") + return re.sub(_whitespace_re, " ", text).strip() + + +def remove_punctuation_at_begin(text): + return re.sub(r"^[,.!?]+", "", text) + + +def remove_aux_symbols(text): + text = re.sub(r"[\<\>\(\)\[\]\"\«\»]+", "", text) + return text + + +def replace_symbols(text): + text = text.replace(";", ",") + text = text.replace("-", " ") + text = text.replace(":", ",") + return text + + +def replace_punctuation(text): + pattern = re.compile("|".join(re.escape(p) for p in rep_map.keys())) + replaced_text = pattern.sub(lambda x: rep_map[x.group()], text) + return replaced_text + + +def text_normalize(text): + text = replace_punctuation(text) + text = replace_symbols(text) + text = remove_aux_symbols(text) + text = remove_punctuation_at_begin(text) + text = collapse_whitespace(text) + text = re.sub(r"([^\.,!\?\-…])$", r"\1", text) + return text + + +def german_to_ipa(text, text_tokenizer): + if type(text) == str: + text = text_normalize(text) + phonemes = text_tokenizer(text) + return phonemes + else: + for i, t in enumerate(text): + text[i] = text_normalize(t) + return text_tokenizer(text) diff --git a/models/tts/maskgct/g2p/g2p/japanese.py b/models/tts/maskgct/g2p/g2p/japanese.py new file mode 100644 index 0000000..bef50ce --- /dev/null +++ b/models/tts/maskgct/g2p/g2p/japanese.py @@ -0,0 +1,816 @@ +# Copyright (c) 2024 Amphion. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +import io, re, os, sys, time, argparse, pdb, json +from io import StringIO +from typing import Optional +import numpy as np +import traceback +import pyopenjtalk +from pykakasi import kakasi + +punctuation = [",", ".", "!", "?", ":", ";", "'", "…"] + +jp_xphone2ipa = [ + " a a", + " i i", + " u ɯ", + " e e", + " o o", + " a: aː", + " i: iː", + " u: ɯː", + " e: eː", + " o: oː", + " k k", + " s s", + " t t", + " n n", + " h ç", + " f ɸ", + " m m", + " y j", + " r ɾ", + " w ɰᵝ", + " N ɴ", + " g g", + " j d ʑ", + " z z", + " d d", + " b b", + " p p", + " q q", + " v v", + " : :", + " by b j", + " ch t ɕ", + " dy d e j", + " ty t e j", + " gy g j", + " gw g ɯ", + " hy ç j", + " ky k j", + " kw k ɯ", + " my m j", + " ny n j", + " py p j", + " ry ɾ j", + " sh ɕ", + " ts t s ɯ", +] + +_mora_list_minimum: list[tuple[str, Optional[str], str]] = [ + ("ヴォ", "v", "o"), + ("ヴェ", "v", "e"), + ("ヴィ", "v", "i"), + ("ヴァ", "v", "a"), + ("ヴ", "v", "u"), + ("ン", None, "N"), + ("ワ", "w", "a"), + ("ロ", "r", "o"), + ("レ", "r", "e"), + ("ル", "r", "u"), + ("リョ", "ry", "o"), + ("リュ", "ry", "u"), + ("リャ", "ry", "a"), + ("リェ", "ry", "e"), + ("リ", "r", "i"), + ("ラ", "r", "a"), + ("ヨ", "y", "o"), + ("ユ", "y", "u"), + ("ヤ", "y", "a"), + ("モ", "m", "o"), + ("メ", "m", "e"), + ("ム", "m", "u"), + ("ミョ", "my", "o"), + ("ミュ", "my", "u"), + ("ミャ", "my", "a"), + ("ミェ", "my", "e"), + ("ミ", "m", "i"), + ("マ", "m", "a"), + ("ポ", "p", "o"), + ("ボ", "b", "o"), + ("ホ", "h", "o"), + ("ペ", "p", "e"), + ("ベ", "b", "e"), + ("ヘ", "h", "e"), + ("プ", "p", "u"), + ("ブ", "b", "u"), + ("フォ", "f", "o"), + ("フェ", "f", "e"), + ("フィ", "f", "i"), + ("ファ", "f", "a"), + ("フ", "f", "u"), + ("ピョ", "py", "o"), + ("ピュ", "py", "u"), + ("ピャ", "py", "a"), + ("ピェ", "py", "e"), + ("ピ", "p", "i"), + ("ビョ", "by", "o"), + ("ビュ", "by", "u"), + ("ビャ", "by", "a"), + ("ビェ", "by", "e"), + ("ビ", "b", "i"), + ("ヒョ", "hy", "o"), + ("ヒュ", "hy", "u"), + ("ヒャ", "hy", "a"), + ("ヒェ", "hy", "e"), + ("ヒ", "h", "i"), + ("パ", "p", "a"), + ("バ", "b", "a"), + ("ハ", "h", "a"), + ("ノ", "n", "o"), + ("ネ", "n", "e"), + ("ヌ", "n", "u"), + ("ニョ", "ny", "o"), + ("ニュ", "ny", "u"), + ("ニャ", "ny", "a"), + ("ニェ", "ny", "e"), + ("ニ", "n", "i"), + ("ナ", "n", "a"), + ("ドゥ", "d", "u"), + ("ド", "d", "o"), + ("トゥ", "t", "u"), + ("ト", "t", "o"), + ("デョ", "dy", "o"), + ("デュ", "dy", "u"), + ("デャ", "dy", "a"), + # ("デェ", "dy", "e"), + ("ディ", "d", "i"), + ("デ", "d", "e"), + ("テョ", "ty", "o"), + ("テュ", "ty", "u"), + ("テャ", "ty", "a"), + ("ティ", "t", "i"), + ("テ", "t", "e"), + ("ツォ", "ts", "o"), + ("ツェ", "ts", "e"), + ("ツィ", "ts", "i"), + ("ツァ", "ts", "a"), + ("ツ", "ts", "u"), + ("ッ", None, "q"), # 「cl」から「q」に変更 + ("チョ", "ch", "o"), + ("チュ", "ch", "u"), + ("チャ", "ch", "a"), + ("チェ", "ch", "e"), + ("チ", "ch", "i"), + ("ダ", "d", "a"), + ("タ", "t", "a"), + ("ゾ", "z", "o"), + ("ソ", "s", "o"), + ("ゼ", "z", "e"), + ("セ", "s", "e"), + ("ズィ", "z", "i"), + ("ズ", "z", "u"), + ("スィ", "s", "i"), + ("ス", "s", "u"), + ("ジョ", "j", "o"), + ("ジュ", "j", "u"), + ("ジャ", "j", "a"), + ("ジェ", "j", "e"), + ("ジ", "j", "i"), + ("ショ", "sh", "o"), + ("シュ", "sh", "u"), + ("シャ", "sh", "a"), + ("シェ", "sh", "e"), + ("シ", "sh", "i"), + ("ザ", "z", "a"), + ("サ", "s", "a"), + ("ゴ", "g", "o"), + ("コ", "k", "o"), + ("ゲ", "g", "e"), + ("ケ", "k", "e"), + ("グヮ", "gw", "a"), + ("グ", "g", "u"), + ("クヮ", "kw", "a"), + ("ク", "k", "u"), + ("ギョ", "gy", "o"), + ("ギュ", "gy", "u"), + ("ギャ", "gy", "a"), + ("ギェ", "gy", "e"), + ("ギ", "g", "i"), + ("キョ", "ky", "o"), + ("キュ", "ky", "u"), + ("キャ", "ky", "a"), + ("キェ", "ky", "e"), + ("キ", "k", "i"), + ("ガ", "g", "a"), + ("カ", "k", "a"), + ("オ", None, "o"), + ("エ", None, "e"), + ("ウォ", "w", "o"), + ("ウェ", "w", "e"), + ("ウィ", "w", "i"), + ("ウ", None, "u"), + ("イェ", "y", "e"), + ("イ", None, "i"), + ("ア", None, "a"), +] + +_mora_list_additional: list[tuple[str, Optional[str], str]] = [ + ("ヴョ", "by", "o"), + ("ヴュ", "by", "u"), + ("ヴャ", "by", "a"), + ("ヲ", None, "o"), + ("ヱ", None, "e"), + ("ヰ", None, "i"), + ("ヮ", "w", "a"), + ("ョ", "y", "o"), + ("ュ", "y", "u"), + ("ヅ", "z", "u"), + ("ヂ", "j", "i"), + ("ヶ", "k", "e"), + ("ャ", "y", "a"), + ("ォ", None, "o"), + ("ェ", None, "e"), + ("ゥ", None, "u"), + ("ィ", None, "i"), + ("ァ", None, "a"), +] + +# 例: "vo" -> "ヴォ", "a" -> "ア" +mora_phonemes_to_mora_kata: dict[str, str] = { + (consonant or "") + vowel: kana for [kana, consonant, vowel] in _mora_list_minimum +} + +# 例: "ヴォ" -> ("v", "o"), "ア" -> (None, "a") +mora_kata_to_mora_phonemes: dict[str, tuple[Optional[str], str]] = { + kana: (consonant, vowel) + for [kana, consonant, vowel] in _mora_list_minimum + _mora_list_additional +} + + +# 正規化で記号を変換するための辞書 +rep_map = { + "：": ":", + "；": ";", + "，": ",", + "。": ".", + "！": "!", + "？": "?", + "\n": ".", + "．": ".", + "⋯": "…", + "···": "…", + "・・・": "…", + "·": ",", + "・": ",", + "•": ",", + "、": ",", + "$": ".", + # "“": "'", + # "”": "'", + # '"': "'", + "‘": "'", + "’": "'", + # "（": "'", + # "）": "'", + # "(": "'", + # ")": "'", + # "《": "'", + # "》": "'", + # "【": "'", + # "】": "'", + # "[": "'", + # "]": "'", + # "——": "-", + # "−": "-", + # "-": "-", + # "『": "'", + # "』": "'", + # "〈": "'", + # "〉": "'", + # "«": "'", + # "»": "'", + # # "～": "-", # これは長音記号「ー」として扱うよう変更 + # # "~": "-", # これは長音記号「ー」として扱うよう変更 + # "「": "'", + # "」": "'", +} + + +def _numeric_feature_by_regex(regex, s): + match = re.search(regex, s) + if match is None: + return -50 + return int(match.group(1)) + + +def replace_punctuation(text: str) -> str: + """句読点等を「.」「,」「!」「?」「'」「-」に正規化し、OpenJTalkで読みが取得できるもののみ残す： + 漢字・平仮名・カタカナ、アルファベット、ギリシャ文字 + """ + pattern = re.compile("|".join(re.escape(p) for p in rep_map.keys())) + # print("before: ", text) + # 句読点を辞書で置換 + replaced_text = pattern.sub(lambda x: rep_map[x.group()], text) + + replaced_text = re.sub( + # ↓ ひらがな、カタカナ、漢字 + r"[^\u3040-\u309F\u30A0-\u30FF\u4E00-\u9FFF\u3400-\u4DBF\u3005" + # ↓ 半角アルファベット（大文字と小文字） + + r"\u0041-\u005A\u0061-\u007A" + # ↓ 全角アルファベット（大文字と小文字） + + r"\uFF21-\uFF3A\uFF41-\uFF5A" + # ↓ ギリシャ文字 + + r"\u0370-\u03FF\u1F00-\u1FFF" + # ↓ "!", "?", "…", ",", ".", "'", "-", 但し`…`はすでに`...`に変換されている + + "".join(punctuation) + r"]+", + # 上述以外の文字を削除 + "", + replaced_text, + ) + # print("after: ", replaced_text) + return replaced_text + + +def fix_phone_tone(phone_tone_list: list[tuple[str, int]]) -> list[tuple[str, int]]: + """ + `phone_tone_list`のtone（アクセントの値）を0か1の範囲に修正する。 + 例: [(a, 0), (i, -1), (u, -1)] → [(a, 1), (i, 0), (u, 0)] + """ + tone_values = set(tone for _, tone in phone_tone_list) + if len(tone_values) == 1: + assert tone_values == {0}, tone_values + return phone_tone_list + elif len(tone_values) == 2: + if tone_values == {0, 1}: + return phone_tone_list + elif tone_values == {-1, 0}: + return [ + (letter, 0 if tone == -1 else 1) for letter, tone in phone_tone_list + ] + else: + raise ValueError(f"Unexpected tone values: {tone_values}") + else: + raise ValueError(f"Unexpected tone values: {tone_values}") + + +def fix_phone_tone_wplen(phone_tone_list, word_phone_length_list): + phones = [] + tones = [] + w_p_len = [] + p_len = len(phone_tone_list) + idx = 0 + w_idx = 0 + while idx < p_len: + offset = 0 + if phone_tone_list[idx] == "▁": + w_p_len.append(w_idx + 1) + + curr_w_p_len = word_phone_length_list[w_idx] + for i in range(curr_w_p_len): + p, t = phone_tone_list[idx] + if p == ":" and len(phones) > 0: + if phones[-1][-1] != ":": + phones[-1] += ":" + offset -= 1 + else: + phones.append(p) + tones.append(str(t)) + idx += 1 + if idx >= p_len: + break + w_p_len.append(curr_w_p_len + offset) + w_idx += 1 + # print(w_p_len) + return phones, tones, w_p_len + + +def g2phone_tone_wo_punct(prosodies) -> list[tuple[str, int]]: + """ + テキストに対して、音素とアクセント（0か1）のペアのリストを返す。 + ただし「!」「.」「?」等の非音素記号(punctuation)は全て消える（ポーズ記号も残さない）。 + 非音素記号を含める処理は`align_tones()`で行われる。 + また「っ」は「cl」でなく「q」に変換される（「ん」は「N」のまま）。 + 例: "こんにちは、世界ー。。元気？！" → + [('k', 0), ('o', 0), ('N', 1), ('n', 1), ('i', 1), ('ch', 1), ('i', 1), ('w', 1), ('a', 1), ('s', 1), ('e', 1), ('k', 0), ('a', 0), ('i', 0), ('i', 0), ('g', 1), ('e', 1), ('N', 0), ('k', 0), ('i', 0)] + """ + result: list[tuple[str, int]] = [] + current_phrase: list[tuple[str, int]] = [] + current_tone = 0 + last_accent = "" + for i, letter in enumerate(prosodies): + # 特殊記号の処理 + + # 文頭記号、無視する + if letter == "^": + assert i == 0, "Unexpected ^" + # アクセント句の終わりに来る記号 + elif letter in ("$", "?", "_", "#"): + # 保持しているフレーズを、アクセント数値を0-1に修正し結果に追加 + result.extend(fix_phone_tone(current_phrase)) + # 末尾に来る終了記号、無視（文中の疑問文は`_`になる） + if letter in ("$", "?"): + assert i == len(prosodies) - 1, f"Unexpected {letter}" + # あとは"_"（ポーズ）と"#"（アクセント句の境界）のみ + # これらは残さず、次のアクセント句に備える。 + + current_phrase = [] + # 0を基準点にしてそこから上昇・下降する（負の場合は上の`fix_phone_tone`で直る） + current_tone = 0 + last_accent = "" + # アクセント上昇記号 + elif letter == "[": + if last_accent != letter: + current_tone = current_tone + 1 + last_accent = letter + # アクセント下降記号 + elif letter == "]": + if last_accent != letter: + current_tone = current_tone - 1 + last_accent = letter + # それ以外は通常の音素 + else: + if letter == "cl": # 「っ」の処理 + letter = "q" + current_phrase.append((letter, current_tone)) + return result + + +def handle_long(sep_phonemes: list[list[str]]) -> list[list[str]]: + for i in range(len(sep_phonemes)): + if sep_phonemes[i][0] == "ー": + # sep_phonemes[i][0] = sep_phonemes[i - 1][-1] + sep_phonemes[i][0] = ":" + if "ー" in sep_phonemes[i]: + for j in range(len(sep_phonemes[i])): + if sep_phonemes[i][j] == "ー": + # sep_phonemes[i][j] = sep_phonemes[i][j - 1][-1] + sep_phonemes[i][j] = ":" + return sep_phonemes + + +def handle_long_word(sep_phonemes: list[list[str]]) -> list[list[str]]: + res = [] + for i in range(len(sep_phonemes)): + if sep_phonemes[i][0] == "ー": + sep_phonemes[i][0] = sep_phonemes[i - 1][-1] + # sep_phonemes[i][0] = ':' + if "ー" in sep_phonemes[i]: + for j in range(len(sep_phonemes[i])): + if sep_phonemes[i][j] == "ー": + sep_phonemes[i][j] = sep_phonemes[i][j - 1][-1] + # sep_phonemes[i][j] = ':' + res.append(sep_phonemes[i]) + res.append("▁") + return res + + +def align_tones( + phones_with_punct: list[str], phone_tone_list: list[tuple[str, int]] +) -> list[tuple[str, int]]: + """ + 例: + …私は、、そう思う。 + phones_with_punct: + [".", ".", ".", "w", "a", "t", "a", "sh", "i", "w", "a", ",", ",", "s", "o", "o", "o", "m", "o", "u", "."] + phone_tone_list: + [("w", 0), ("a", 0), ("t", 1), ("a", 1), ("sh", 1), ("i", 1), ("w", 1), ("a", 1), ("s", 0), ("o", 0), ("o", 1), ("o", 1), ("m", 1), ("o", 1), ("u", 0))] + Return: + [(".", 0), (".", 0), (".", 0), ("w", 0), ("a", 0), ("t", 1), ("a", 1), ("sh", 1), ("i", 1), ("w", 1), ("a", 1), (",", 0), (",", 0), ("s", 0), ("o", 0), ("o", 1), ("o", 1), ("m", 1), ("o", 1), ("u", 0), (".", 0)] + """ + result: list[tuple[str, int]] = [] + tone_index = 0 + for phone in phones_with_punct: + if tone_index >= len(phone_tone_list): + # 余ったpunctuationがある場合 → (punctuation, 0)を追加 + result.append((phone, 0)) + elif phone == phone_tone_list[tone_index][0]: + # phone_tone_listの現在の音素と一致する場合 → toneをそこから取得、(phone, tone)を追加 + result.append((phone, phone_tone_list[tone_index][1])) + # 探すindexを1つ進める + tone_index += 1 + elif phone in punctuation or phone == "▁": + # phoneがpunctuationの場合 → (phone, 0)を追加 + result.append((phone, 0)) + else: + print(f"phones: {phones_with_punct}") + print(f"phone_tone_list: {phone_tone_list}") + print(f"result: {result}") + print(f"tone_index: {tone_index}") + print(f"phone: {phone}") + raise ValueError(f"Unexpected phone: {phone}") + return result + + +def kata2phoneme_list(text: str) -> list[str]: + """ + 原則カタカナの`text`を受け取り、それをそのままいじらずに音素記号のリストに変換。 + 注意点： + - punctuationが来た場合（punctuationが1文字の場合がありうる）、処理せず1文字のリストを返す + - 冒頭に続く「ー」はそのまま「ー」のままにする（`handle_long()`で処理される） + - 文中の「ー」は前の音素記号の最後の音素記号に変換される。 + 例： + `ーーソーナノカーー` → ["ー", "ー", "s", "o", "o", "n", "a", "n", "o", "k", "a", "a", "a"] + `?` → ["?"] + """ + if text in punctuation: + return [text] + # `text`がカタカナ（`ー`含む）のみからなるかどうかをチェック + if re.fullmatch(r"[\u30A0-\u30FF]+", text) is None: + raise ValueError(f"Input must be katakana only: {text}") + sorted_keys = sorted(mora_kata_to_mora_phonemes.keys(), key=len, reverse=True) + pattern = "|".join(map(re.escape, sorted_keys)) + + def mora2phonemes(mora: str) -> str: + cosonant, vowel = mora_kata_to_mora_phonemes[mora] + if cosonant is None: + return f" {vowel}" + return f" {cosonant} {vowel}" + + spaced_phonemes = re.sub(pattern, lambda m: mora2phonemes(m.group()), text) + + # 長音記号「ー」の処理 + long_pattern = r"(\w)(ー*)" + long_replacement = lambda m: m.group(1) + (" " + m.group(1)) * len(m.group(2)) + spaced_phonemes = re.sub(long_pattern, long_replacement, spaced_phonemes) + # spaced_phonemes += ' ▁' + return spaced_phonemes.strip().split(" ") + + +def frontend2phoneme(labels, drop_unvoiced_vowels=False): + N = len(labels) + + phones = [] + for n in range(N): + lab_curr = labels[n] + # print(lab_curr) + # current phoneme + p3 = re.search(r"\-(.*?)\+", lab_curr).group(1) + + # deal unvoiced vowels as normal vowels + if drop_unvoiced_vowels and p3 in "AEIOU": + p3 = p3.lower() + + # deal with sil at the beginning and the end of text + if p3 == "sil": + # assert n == 0 or n == N - 1 + # if n == 0: + # phones.append("^") + # elif n == N - 1: + # # check question form or not + # e3 = _numeric_feature_by_regex(r"!(\d+)_", lab_curr) + # if e3 == 0: + # phones.append("$") + # elif e3 == 1: + # phones.append("?") + continue + elif p3 == "pau": + phones.append("_") + continue + else: + phones.append(p3) + + # accent type and position info (forward or backward) + a1 = _numeric_feature_by_regex(r"/A:([0-9\-]+)\+", lab_curr) + a2 = _numeric_feature_by_regex(r"\+(\d+)\+", lab_curr) + a3 = _numeric_feature_by_regex(r"\+(\d+)/", lab_curr) + + # number of mora in accent phrase + f1 = _numeric_feature_by_regex(r"/F:(\d+)_", lab_curr) + + a2_next = _numeric_feature_by_regex(r"\+(\d+)\+", labels[n + 1]) + # accent phrase border + # print(p3, a1, a2, a3, f1, a2_next, lab_curr) + if a3 == 1 and a2_next == 1 and p3 in "aeiouAEIOUNcl": + phones.append("#") + # pitch falling + elif a1 == 0 and a2_next == a2 + 1 and a2 != f1: + phones.append("]") + # pitch rising + elif a2 == 1 and a2_next == 2: + phones.append("[") + + # phones = ' '.join(phones) + return phones + + +class JapanesePhoneConverter(object): + def __init__(self, lexicon_path=None, ipa_dict_path=None): + # lexicon_lines = open(lexicon_path, 'r', encoding='utf-8').readlines() + # self.lexicon = {} + # self.single_dict = {} + # self.double_dict = {} + # for curr_line in lexicon_lines: + # k,v = curr_line.strip().split('+',1) + # self.lexicon[k] = v + # if len(k) == 2: + # self.double_dict[k] = v + # elif len(k) == 1: + # self.single_dict[k] = v + self.ipa_dict = {} + for curr_line in jp_xphone2ipa: + k, v = curr_line.strip().split(" ", 1) + self.ipa_dict[k] = re.sub("\s", "", v) + # kakasi1 = kakasi() + # kakasi1.setMode("H","K") + # kakasi1.setMode("J","K") + # kakasi1.setMode("r","Hepburn") + self.japan_JH2K = kakasi() + self.table = {ord(f): ord(t) for f, t in zip("67", "_¯")} + + def text2sep_kata(self, parsed) -> tuple[list[str], list[str]]: + """ + `text_normalize`で正規化済みの`norm_text`を受け取り、それを単語分割し、 + 分割された単語リストとその読み（カタカナor記号1文字）のリストのタプルを返す。 + 単語分割結果は、`g2p()`の`word2ph`で1文字あたりに割り振る音素記号の数を決めるために使う。 + 例: + `私はそう思う!って感じ?` → + ["私", "は", "そう", "思う", "!", "って", "感じ", "?"], ["ワタシ", "ワ", "ソー", "オモウ", "!", "ッテ", "カンジ", "?"] + """ + # parsed: OpenJTalkの解析結果 + sep_text: list[str] = [] + sep_kata: list[str] = [] + fix_parsed = [] + i = 0 + while i <= len(parsed) - 1: + # word: 実際の単語の文字列 + # yomi: その読み、但し無声化サインの`’`は除去 + # print(parsed) + yomi = parsed[i]["pron"] + tmp_parsed = parsed[i] + if i != len(parsed) - 1 and parsed[i + 1]["string"] in [ + "々", + "ゝ", + "ヽ", + "ゞ", + "ヾ", + "゛", + ]: + word = parsed[i]["string"] + parsed[i + 1]["string"] + i += 1 + else: + word = parsed[i]["string"] + word, yomi = replace_punctuation(word), yomi.replace("’", "") + """ + ここで`yomi`の取りうる値は以下の通りのはず。 + - `word`が通常単語 → 通常の読み（カタカナ） + （カタカナからなり、長音記号も含みうる、`アー` 等） + - `word`が`ー` から始まる → `ーラー` や `ーーー` など + - `word`が句読点や空白等 → `、` + - `word`が`?` → `？`（全角になる） + 他にも`word`が読めないキリル文字アラビア文字等が来ると`、`になるが、正規化でこの場合は起きないはず。 + また元のコードでは`yomi`が空白の場合の処理があったが、これは起きないはず。 + 処理すべきは`yomi`が`、`の場合のみのはず。 + """ + assert yomi != "", f"Empty yomi: {word}" + if yomi == "、": + # wordは正規化されているので、`.`, `,`, `!`, `'`, `-`のいずれか + if word not in ( + ".", + ",", + "!", + "'", + "-", + "?", + ":", + ";", + "…", + "", + ): + # ここはpyopenjtalkが読めない文字等のときに起こる + print( + "Cannot read:{}, yomi:{}, new_word:{};".format( + word, yomi, self.japan_JH2K.convert(word)[0]["kana"] + ) + ) + # raise ValueError(word) + word = self.japan_JH2K.convert(word)[0]["kana"] + # print(word, self.japan_JH2K.convert(word)[0]['kana'], kata2phoneme_list(self.japan_JH2K.convert(word)[0]['kana'])) + tmp_parsed["pron"] = word + # yomi = "-" + # word = ',' + # yomiは元の記号のままに変更 + # else: + # parsed[i]['pron'] = parsed[i]["string"] + yomi = word + elif yomi == "？": + assert word == "?", f"yomi `？` comes from: {word}" + yomi = "?" + if word == "": + i += 1 + continue + sep_text.append(word) + sep_kata.append(yomi) + # print(word, yomi, parts) + fix_parsed.append(tmp_parsed) + i += 1 + # print(sep_text, sep_kata) + return sep_text, sep_kata, fix_parsed + + def getSentencePhone(self, sentence, blank_mode=True, phoneme_mode=False): + # print("origin:", sentence) + words = [] + words_phone_len = [] + short_char_flag = False + output_duration_flag = [] + output_before_sil_flag = [] + normed_text = [] + sentence = sentence.strip().strip("'") + sentence = re.sub(r"\s+", "", sentence) + output_res = [] + failed_words = [] + last_long_pause = 4 + last_word = None + frontend_text = pyopenjtalk.run_frontend(sentence) + # print("frontend_text: ", frontend_text) + try: + frontend_text = pyopenjtalk.estimate_accent(frontend_text) + except: + pass + # print("estimate_accent: ", frontend_text) + # sep_text: 単語単位の単語のリスト + # sep_kata: 単語単位の単語のカタカナ読みのリスト + sep_text, sep_kata, frontend_text = self.text2sep_kata(frontend_text) + # print("sep_text: ", sep_text) + # print("sep_kata: ", sep_kata) + # print("frontend_text: ", frontend_text) + # sep_phonemes: 各単語ごとの音素のリストのリスト + sep_phonemes = handle_long_word([kata2phoneme_list(i) for i in sep_kata]) + # print("sep_phonemes: ", sep_phonemes) + + pron_text = [x["pron"].strip().replace("’", "") for x in frontend_text] + # pdb.set_trace() + prosodys = pyopenjtalk.make_label(frontend_text) + prosodys = frontend2phoneme(prosodys, drop_unvoiced_vowels=True) + # print("prosodys: ", ' '.join(prosodys)) + # print("pron_text: ", pron_text) + normed_text = [x["string"].strip() for x in frontend_text] + # punctuationがすべて消えた、音素とアクセントのタプルのリスト + phone_tone_list_wo_punct = g2phone_tone_wo_punct(prosodys) + # print("phone_tone_list_wo_punct: ", phone_tone_list_wo_punct) + + # phone_w_punct: sep_phonemesを結合した、punctuationを元のまま保持した音素列 + phone_w_punct: list[str] = [] + w_p_len = [] + for i in sep_phonemes: + phone_w_punct += i + w_p_len.append(len(i)) + phone_w_punct = phone_w_punct[:-1] + # punctuation無しのアクセント情報を使って、punctuationを含めたアクセント情報を作る + # print("phone_w_punct: ", phone_w_punct) + # print("phone_tone_list_wo_punct: ", phone_tone_list_wo_punct) + phone_tone_list = align_tones(phone_w_punct, phone_tone_list_wo_punct) + + jp_item = {} + jp_p = "" + jp_t = "" + # mye rye pye bye nye + # je she + # print(phone_tone_list) + for p, t in phone_tone_list: + if p in self.ipa_dict: + curr_p = self.ipa_dict[p] + jp_p += curr_p + jp_t += str(t + 6) * len(curr_p) + elif p in punctuation: + jp_p += p + jp_t += "0" + elif p == "▁": + jp_p += p + jp_t += " " + else: + print(p, t) + jp_p += "|" + jp_t += "0" + # return phones, tones, w_p_len + jp_p = jp_p.replace("▁", " ") + jp_t = jp_t.translate(self.table) + jp_l = "" + for t in jp_t: + if t == " ": + jp_l += " " + else: + jp_l += "2" + # print(jp_p) + # print(jp_t) + # print(jp_l) + # print(len(jp_p_len), sum(w_p_len), len(jp_p), sum(jp_p_len)) + assert len(jp_p) == len(jp_t) and len(jp_p) == len(jp_l) + + jp_item["jp_p"] = jp_p.replace("| |", "|").rstrip("|") + jp_item["jp_t"] = jp_t + jp_item["jp_l"] = jp_l + jp_item["jp_normed_text"] = " ".join(normed_text) + jp_item["jp_pron_text"] = " ".join(pron_text) + # jp_item['jp_ruoma'] = sep_phonemes + # print(len(normed_text), len(sep_phonemes)) + # print(normed_text) + return jp_item + + +jpc = JapanesePhoneConverter() + + +def japanese_to_ipa(text, text_tokenizer): + # phonemes = text_tokenizer(text) + if type(text) == str: + return jpc.getSentencePhone(text)["jp_p"] + else: + result_ph = [] + for t in text: + result_ph.append(jpc.getSentencePhone(t)["jp_p"]) + return result_ph diff --git a/models/tts/maskgct/g2p/g2p/korean.py b/models/tts/maskgct/g2p/g2p/korean.py new file mode 100644 index 0000000..c7c540b --- /dev/null +++ b/models/tts/maskgct/g2p/g2p/korean.py @@ -0,0 +1,81 @@ +# Copyright (c) 2024 Amphion. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +import re + +""" + Text clean time +""" +english_dictionary = { + "KOREA": "코리아", + "IDOL": "아이돌", + "IT": "아이티", + "IQ": "아이큐", + "UP": "업", + "DOWN": "다운", + "PC": "피씨", + "CCTV": "씨씨티비", + "SNS": "에스엔에스", + "AI": "에이아이", + "CEO": "씨이오", + "A": "에이", + "B": "비", + "C": "씨", + "D": "디", + "E": "이", + "F": "에프", + "G": "지", + "H": "에이치", + "I": "아이", + "J": "제이", + "K": "케이", + "L": "엘", + "M": "엠", + "N": "엔", + "O": "오", + "P": "피", + "Q": "큐", + "R": "알", + "S": "에스", + "T": "티", + "U": "유", + "V": "브이", + "W": "더블유", + "X": "엑스", + "Y": "와이", + "Z": "제트", +} + + +def normalize(text): + text = text.strip() + text = re.sub( + "[⺀-⺙⺛-⻳⼀-⿕々〇〡-〩〸-〺〻㐀-䶵一-鿃豈-鶴侮-頻並-龎]", "", text + ) + text = normalize_english(text) + text = text.lower() + return text + + +def normalize_english(text): + def fn(m): + word = m.group() + if word in english_dictionary: + return english_dictionary.get(word) + return word + + text = re.sub("([A-Za-z]+)", fn, text) + return text + + +def korean_to_ipa(text, text_tokenizer): + if type(text) == str: + text = normalize(text) + phonemes = text_tokenizer(text) + return phonemes + else: + for i, t in enumerate(text): + text[i] = normalize(t) + return text_tokenizer(text) diff --git a/models/tts/maskgct/g2p/g2p/mandarin.py b/models/tts/maskgct/g2p/g2p/mandarin.py new file mode 100644 index 0000000..0c3102c --- /dev/null +++ b/models/tts/maskgct/g2p/g2p/mandarin.py @@ -0,0 +1,595 @@ +# Copyright (c) 2024 Amphion. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +import re +import jieba +import cn2an +from pypinyin import lazy_pinyin, BOPOMOFO +from typing import List +from models.tts.maskgct.g2p.g2p.chinese_model_g2p import BertPolyPredict +from models.tts.maskgct.g2p.utils.front_utils import * +import os + +# from g2pw import G2PWConverter + + +# set blank level, {0："none",1:"char", 2:"word"} +BLANK_LEVEL = 0 + +# conv = G2PWConverter(style='pinyin', enable_non_tradional_chinese=True) +resource_path = os.path.join(os.path.dirname(os.path.abspath(__file__)),"..") +poly_all_class_path = os.path.join( + resource_path,"sources", "g2p_chinese_model", "polychar.txt" +) +if not os.path.exists(poly_all_class_path): + print( + "Incorrect path for polyphonic character class dictionary: {}, please check...".format( + poly_all_class_path + ) + ) + exit() +poly_dict = generate_poly_lexicon(poly_all_class_path) + +# Set up G2PW model parameters +g2pw_poly_model_path = os.path.join(resource_path, "sources", "g2p_chinese_model") +if not os.path.exists(g2pw_poly_model_path): + print( + "Incorrect path for g2pw polyphonic character model: {}, please check...".format( + g2pw_poly_model_path + ) + ) + exit() + +json_file_path = os.path.join( + resource_path, "sources", "g2p_chinese_model", "polydict.json" +) +if not os.path.exists(json_file_path): + print( + "Incorrect path for g2pw id to pinyin dictionary: {}, please check...".format( + json_file_path + ) + ) + exit() + +jsonr_file_path = os.path.join( + resource_path, "sources", "g2p_chinese_model", "polydict_r.json" +) +if not os.path.exists(jsonr_file_path): + print( + "Incorrect path for g2pw pinyin to id dictionary: {}, please check...".format( + jsonr_file_path + ) + ) + exit() + +g2pw_poly_predict = BertPolyPredict( + g2pw_poly_model_path, jsonr_file_path, json_file_path +) + + +""" + Text clean time +""" +# List of (Latin alphabet, bopomofo) pairs: +_latin_to_bopomofo = [ + (re.compile("%s" % x[0], re.IGNORECASE), x[1]) + for x in [ + ("a", "ㄟˉ"), + ("b", "ㄅㄧˋ"), + ("c", "ㄙㄧˉ"), + ("d", "ㄉㄧˋ"), + ("e", "ㄧˋ"), + ("f", "ㄝˊㄈㄨˋ"), + ("g", "ㄐㄧˋ"), + ("h", "ㄝˇㄑㄩˋ"), + ("i", "ㄞˋ"), + ("j", "ㄐㄟˋ"), + ("k", "ㄎㄟˋ"), + ("l", "ㄝˊㄛˋ"), + ("m", "ㄝˊㄇㄨˋ"), + ("n", "ㄣˉ"), + ("o", "ㄡˉ"), + ("p", "ㄆㄧˉ"), + ("q", "ㄎㄧㄡˉ"), + ("r", "ㄚˋ"), + ("s", "ㄝˊㄙˋ"), + ("t", "ㄊㄧˋ"), + ("u", "ㄧㄡˉ"), + ("v", "ㄨㄧˉ"), + ("w", "ㄉㄚˋㄅㄨˋㄌㄧㄡˋ"), + ("x", "ㄝˉㄎㄨˋㄙˋ"), + ("y", "ㄨㄞˋ"), + ("z", "ㄗㄟˋ"), + ] +] + +# List of (bopomofo, ipa) pairs: +_bopomofo_to_ipa = [ + (re.compile("%s" % x[0]), x[1]) + for x in [ + ("ㄅㄛ", "p⁼wo"), + ("ㄆㄛ", "pʰwo"), + ("ㄇㄛ", "mwo"), + ("ㄈㄛ", "fwo"), + ("ㄧㄢ", "|jɛn"), + ("ㄩㄢ", "|ɥæn"), + ("ㄧㄣ", "|in"), + ("ㄩㄣ", "|ɥn"), + ("ㄧㄥ", "|iŋ"), + ("ㄨㄥ", "|ʊŋ"), + ("ㄩㄥ", "|jʊŋ"), + # Add + ("ㄧㄚ", "|ia"), + ("ㄧㄝ", "|iɛ"), + ("ㄧㄠ", "|iɑʊ"), + ("ㄧㄡ", "|ioʊ"), + ("ㄧㄤ", "|iɑŋ"), + ("ㄨㄚ", "|ua"), + ("ㄨㄛ", "|uo"), + ("ㄨㄞ", "|uaɪ"), + ("ㄨㄟ", "|ueɪ"), + ("ㄨㄢ", "|uan"), + ("ㄨㄣ", "|uən"), + ("ㄨㄤ", "|uɑŋ"), + ("ㄩㄝ", "|ɥɛ"), + # End + ("ㄅ", "p⁼"), + ("ㄆ", "pʰ"), + ("ㄇ", "m"), + ("ㄈ", "f"), + ("ㄉ", "t⁼"), + ("ㄊ", "tʰ"), + ("ㄋ", "n"), + ("ㄌ", "l"), + ("ㄍ", "k⁼"), + ("ㄎ", "kʰ"), + ("ㄏ", "x"), + ("ㄐ", "tʃ⁼"), + ("ㄑ", "tʃʰ"), + ("ㄒ", "ʃ"), + ("ㄓ", "ts`⁼"), + ("ㄔ", "ts`ʰ"), + ("ㄕ", "s`"), + ("ㄖ", "ɹ`"), + ("ㄗ", "ts⁼"), + ("ㄘ", "tsʰ"), + ("ㄙ", "|s"), + ("ㄚ", "|a"), + ("ㄛ", "|o"), + ("ㄜ", "|ə"), + ("ㄝ", "|ɛ"), + ("ㄞ", "|aɪ"), + ("ㄟ", "|eɪ"), + ("ㄠ", "|ɑʊ"), + ("ㄡ", "|oʊ"), + ("ㄢ", "|an"), + ("ㄣ", "|ən"), + ("ㄤ", "|ɑŋ"), + ("ㄥ", "|əŋ"), + ("ㄦ", "əɹ"), + ("ㄧ", "|i"), + ("ㄨ", "|u"), + ("ㄩ", "|ɥ"), + ("ˉ", "→|"), + ("ˊ", "↑|"), + ("ˇ", "↓↑|"), + ("ˋ", "↓|"), + ("˙", "|"), + ] +] +must_not_er_words = {"女儿", "老儿", "男儿", "少儿", "小儿"} + +word_pinyin_dict = {} +with open( + os.path.join(resource_path,"sources/chinese_lexicon.txt"), "r", encoding="utf-8" +) as fread: + txt_list = fread.readlines() + for txt in txt_list: + word, pinyin = txt.strip().split("\t") + word_pinyin_dict[word] = pinyin + fread.close() + +pinyin_2_bopomofo_dict = {} +with open( + os.path.join(resource_path,"sources/pinyin_2_bpmf.txt"), "r", encoding="utf-8" +) as fread: + txt_list = fread.readlines() + for txt in txt_list: + pinyin, bopomofo = txt.strip().split("\t") + pinyin_2_bopomofo_dict[pinyin] = bopomofo + fread.close() + +tone_dict = { + "0": "˙", + "5": "˙", + "1": "", + "2": "ˊ", + "3": "ˇ", + "4": "ˋ", +} + +bopomofos2pinyin_dict = {} +with open( + os.path.join(resource_path,"sources/bpmf_2_pinyin.txt"), "r", encoding="utf-8" +) as fread: + txt_list = fread.readlines() + for txt in txt_list: + v, k = txt.strip().split("\t") + bopomofos2pinyin_dict[k] = v + fread.close() + + +def bpmf_to_pinyin(text): + bopomofo_list = text.split("|") + pinyin_list = [] + for info in bopomofo_list: + pinyin = "" + for c in info: + if c in bopomofos2pinyin_dict: + pinyin += bopomofos2pinyin_dict[c] + if len(pinyin) == 0: + continue + if pinyin[-1] not in "01234": + pinyin += "1" + if pinyin[:-1] == "ve": + pinyin = "y" + pinyin + if pinyin[:-1] == "sh": + pinyin = pinyin[:-1] + "i" + pinyin[-1] + if pinyin == "sh": + pinyin = pinyin[:-1] + "i" + if pinyin[:-1] == "s": + pinyin = "si" + pinyin[-1] + if pinyin[:-1] == "c": + pinyin = "ci" + pinyin[-1] + if pinyin[:-1] == "i": + pinyin = "yi" + pinyin[-1] + if pinyin[:-1] == "iou": + pinyin = "you" + pinyin[-1] + if pinyin[:-1] == "ien": + pinyin = "yin" + pinyin[-1] + if "iou" in pinyin and pinyin[-4:-1] == "iou": + pinyin = pinyin[:-4] + "iu" + pinyin[-1] + if "uei" in pinyin: + if pinyin[:-1] == "uei": + pinyin = "wei" + pinyin[-1] + elif pinyin[-4:-1] == "uei": + pinyin = pinyin[:-4] + "ui" + pinyin[-1] + if "uen" in pinyin and pinyin[-4:-1] == "uen": + if pinyin[:-1] == "uen": + pinyin = "wen" + pinyin[-1] + elif pinyin[-4:-1] == "uei": + pinyin = pinyin[:-4] + "un" + pinyin[-1] + if "van" in pinyin and pinyin[-4:-1] == "van": + if pinyin[:-1] == "van": + pinyin = "yuan" + pinyin[-1] + elif pinyin[-4:-1] == "van": + pinyin = pinyin[:-4] + "uan" + pinyin[-1] + if "ueng" in pinyin and pinyin[-5:-1] == "ueng": + pinyin = pinyin[:-5] + "ong" + pinyin[-1] + if pinyin[:-1] == "veng": + pinyin = "yong" + pinyin[-1] + if "veng" in pinyin and pinyin[-5:-1] == "veng": + pinyin = pinyin[:-5] + "iong" + pinyin[-1] + if pinyin[:-1] == "ieng": + pinyin = "ying" + pinyin[-1] + if pinyin[:-1] == "u": + pinyin = "wu" + pinyin[-1] + if pinyin[:-1] == "v": + pinyin = "yv" + pinyin[-1] + if pinyin[:-1] == "ing": + pinyin = "ying" + pinyin[-1] + if pinyin[:-1] == "z": + pinyin = "zi" + pinyin[-1] + if pinyin[:-1] == "zh": + pinyin = "zhi" + pinyin[-1] + if pinyin[0] == "u": + pinyin = "w" + pinyin[1:] + if pinyin[0] == "i": + pinyin = "y" + pinyin[1:] + pinyin = pinyin.replace("ien", "in") + + pinyin_list.append(pinyin) + return " ".join(pinyin_list) + + +# Convert numbers to Chinese pronunciation +def number_to_chinese(text): + # numbers = re.findall(r'\d+(?:\.?\d+)?', text) + # for number in numbers: + # text = text.replace(number, cn2an.an2cn(number), 1) + text = cn2an.transform(text, "an2cn") + return text + + +def normalization(text): + text = text.replace("，", ",") + text = text.replace("。", ".") + text = text.replace("！", "!") + text = text.replace("？", "?") + text = text.replace("；", ";") + text = text.replace("：", ":") + text = text.replace("、", ",") + text = text.replace("‘", "'") + text = text.replace("’", "'") + text = text.replace("⋯", "…") + text = text.replace("···", "…") + text = text.replace("・・・", "…") + text = text.replace("...", "…") + text = re.sub(r"\s+", "", text) + text = re.sub(r"[^\u4e00-\u9fff\s_,\.\?!;:\'…]", "", text) + text = re.sub(r"\s*([,\.\?!;:\'…])\s*", r"\1", text) + return text + + +def change_tone(bopomofo: str, tone: str) -> str: + if bopomofo[-1] not in "˙ˊˇˋ": + bopomofo = bopomofo + tone + else: + bopomofo = bopomofo[:-1] + tone + return bopomofo + + +def er_sandhi(word: str, bopomofos: List[str]) -> List[str]: + if len(word) > 1 and word[-1] == "儿" and word not in must_not_er_words: + bopomofos[-1] = change_tone(bopomofos[-1], "˙") + return bopomofos + + +def bu_sandhi(word: str, bopomofos: List[str]) -> List[str]: + valid_char = set(word) + if len(valid_char) == 1 and "不" in valid_char: + pass + elif word in ["不字"]: + pass + elif len(word) == 3 and word[1] == "不" and bopomofos[1][:-1] == "ㄅㄨ": + bopomofos[1] = bopomofos[1][:-1] + "˙" + else: + for i, char in enumerate(word): + if ( + i + 1 < len(bopomofos) + and char == "不" + and i + 1 < len(word) + and 0 < len(bopomofos[i + 1]) + and bopomofos[i + 1][-1] == "ˋ" + ): + bopomofos[i] = bopomofos[i][:-1] + "ˊ" + return bopomofos + + +def yi_sandhi(word: str, bopomofos: List[str]) -> List[str]: + punc = "：，；。？！“”‘’':,;.?!()（）{}【】[]-~`、 " + if word.find("一") != -1 and any( + [item.isnumeric() for item in word if item != "一"] + ): + for i in range(len(word)): + if ( + i == 0 + and word[0] == "一" + and len(word) > 1 + and word[1] + not in [ + "零", + "一", + "二", + "三", + "四", + "五", + "六", + "七", + "八", + "九", + "十", + ] + ): + if len(bopomofos[0]) > 0 and bopomofos[1][-1] in ["ˋ", "˙"]: + bopomofos[0] = change_tone(bopomofos[0], "ˊ") + else: + bopomofos[0] = change_tone(bopomofos[0], "ˋ") + elif word[i] == "一": + bopomofos[i] = change_tone(bopomofos[i], "") + return bopomofos + elif len(word) == 3 and word[1] == "一" and word[0] == word[-1]: + bopomofos[1] = change_tone(bopomofos[1], "˙") + elif word.startswith("第一"): + bopomofos[1] = change_tone(bopomofos[1], "") + elif word.startswith("一月") or word.startswith("一日") or word.startswith("一号"): + bopomofos[0] = change_tone(bopomofos[0], "") + else: + for i, char in enumerate(word): + if char == "一" and i + 1 < len(word): + if ( + len(bopomofos) > i + 1 + and len(bopomofos[i + 1]) > 0 + and bopomofos[i + 1][-1] in {"ˋ"} + ): + bopomofos[i] = change_tone(bopomofos[i], "ˊ") + else: + if word[i + 1] not in punc: + bopomofos[i] = change_tone(bopomofos[i], "ˋ") + else: + pass + return bopomofos + + +def merge_bu(seg: List) -> List: + new_seg = [] + last_word = "" + for word in seg: + if word != "不": + if last_word == "不": + word = last_word + word + new_seg.append(word) + last_word = word + return new_seg + + +def merge_er(seg: List) -> List: + new_seg = [] + for i, word in enumerate(seg): + if i - 1 >= 0 and word == "儿": + new_seg[-1] = new_seg[-1] + seg[i] + else: + new_seg.append(word) + return new_seg + + +def merge_yi(seg: List) -> List: + new_seg = [] + # function 1 + for i, word in enumerate(seg): + if ( + i - 1 >= 0 + and word == "一" + and i + 1 < len(seg) + and seg[i - 1] == seg[i + 1] + ): + if i - 1 < len(new_seg): + new_seg[i - 1] = new_seg[i - 1] + "一" + new_seg[i - 1] + else: + new_seg.append(word) + new_seg.append(seg[i + 1]) + else: + if i - 2 >= 0 and seg[i - 1] == "一" and seg[i - 2] == word: + continue + else: + new_seg.append(word) + seg = new_seg + new_seg = [] + isnumeric_flag = False + for i, word in enumerate(seg): + if all([item.isnumeric() for item in word]) and not isnumeric_flag: + isnumeric_flag = True + new_seg.append(word) + else: + new_seg.append(word) + seg = new_seg + new_seg = [] + # function 2 + for i, word in enumerate(seg): + if new_seg and new_seg[-1] == "一": + new_seg[-1] = new_seg[-1] + word + else: + new_seg.append(word) + return new_seg + + +# Word Segmentation, and convert Chinese pronunciation to pinyin (bopomofo) +def chinese_to_bopomofo(text_short, sentence): + # bopomofos = conv(text_short) + words = jieba.lcut(text_short, cut_all=False) + words = merge_yi(words) + words = merge_bu(words) + words = merge_er(words) + text = "" + + char_index = 0 + for word in words: + bopomofos = [] + if word in word_pinyin_dict and word not in poly_dict: + pinyin = word_pinyin_dict[word] + for py in pinyin.split(" "): + if py[:-1] in pinyin_2_bopomofo_dict and py[-1] in tone_dict: + bopomofos.append( + pinyin_2_bopomofo_dict[py[:-1]] + tone_dict[py[-1]] + ) + if BLANK_LEVEL == 1: + bopomofos.append("_") + else: + bopomofos_lazy = lazy_pinyin(word, BOPOMOFO) + bopomofos += bopomofos_lazy + if BLANK_LEVEL == 1: + bopomofos.append("_") + else: + for i in range(len(word)): + c = word[i] + if c in poly_dict: + poly_pinyin = g2pw_poly_predict.predict_process( + [text_short, char_index + i] + )[0] + py = poly_pinyin[2:-1] + bopomofos.append( + pinyin_2_bopomofo_dict[py[:-1]] + tone_dict[py[-1]] + ) + if BLANK_LEVEL == 1: + bopomofos.append("_") + elif c in word_pinyin_dict: + py = word_pinyin_dict[c] + bopomofos.append( + pinyin_2_bopomofo_dict[py[:-1]] + tone_dict[py[-1]] + ) + if BLANK_LEVEL == 1: + bopomofos.append("_") + else: + bopomofos.append(c) + if BLANK_LEVEL == 1: + bopomofos.append("_") + if BLANK_LEVEL == 2: + bopomofos.append("_") + char_index += len(word) + + if ( + len(word) == 3 + and bopomofos[0][-1] == "ˇ" + and bopomofos[1][-1] == "ˇ" + and bopomofos[-1][-1] == "ˇ" + ): + bopomofos[0] = bopomofos[0] + "ˊ" + bopomofos[1] = bopomofos[1] + "ˊ" + if len(word) == 2 and bopomofos[0][-1] == "ˇ" and bopomofos[-1][-1] == "ˇ": + bopomofos[0] = bopomofos[0][:-1] + "ˊ" + bopomofos = bu_sandhi(word, bopomofos) + bopomofos = yi_sandhi(word, bopomofos) + bopomofos = er_sandhi(word, bopomofos) + if not re.search("[\u4e00-\u9fff]", word): + text += "|" + word + continue + for i in range(len(bopomofos)): + bopomofos[i] = re.sub(r"([\u3105-\u3129])$", r"\1ˉ", bopomofos[i]) + if text != "": + text += "|" + text += "|".join(bopomofos) + return text + + +# Convert latin pronunciation to pinyin (bopomofo) +def latin_to_bopomofo(text): + for regex, replacement in _latin_to_bopomofo: + text = re.sub(regex, replacement, text) + return text + + +# Convert pinyin (bopomofo) to IPA +def bopomofo_to_ipa(text): + for regex, replacement in _bopomofo_to_ipa: + text = re.sub(regex, replacement, text) + return text + + +def _chinese_to_ipa(text, sentence): + text = number_to_chinese(text.strip()) + text = normalization(text) + text = chinese_to_bopomofo(text, sentence) + # pinyin = bpmf_to_pinyin(text) + text = latin_to_bopomofo(text) + text = bopomofo_to_ipa(text) + text = re.sub("([sɹ]`[⁼ʰ]?)([→↓↑ ]+|$)", r"\1ɹ\2", text) + text = re.sub("([s][⁼ʰ]?)([→↓↑ ]+|$)", r"\1ɹ\2", text) + text = re.sub(r"^\||[^\w\s_,\.\?!;:\'…\|→↓↑⁼ʰ`]", "", text) + text = re.sub(r"([,\.\?!;:\'…])", r"|\1|", text) + text = re.sub(r"\|+", "|", text) + text = text.rstrip("|") + return text + + +# Convert Chinese to IPA +def chinese_to_ipa(text, sentence, text_tokenizer): + # phonemes = text_tokenizer(text.strip()) + if type(text) == str: + return _chinese_to_ipa(text, sentence) + else: + result_ph = [] + for t in text: + result_ph.append(_chinese_to_ipa(t, sentence)) + return result_ph diff --git a/models/tts/maskgct/g2p/g2p/text_tokenizers.py b/models/tts/maskgct/g2p/g2p/text_tokenizers.py new file mode 100644 index 0000000..45cb348 --- /dev/null +++ b/models/tts/maskgct/g2p/g2p/text_tokenizers.py @@ -0,0 +1,84 @@ +# Copyright (c) 2024 Amphion. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +import re +import os +from typing import List, Pattern, Union +from phonemizer.utils import list2str, str2list +from phonemizer.backend import EspeakBackend +from phonemizer.backend.espeak.language_switch import LanguageSwitch +from phonemizer.backend.espeak.words_mismatch import WordMismatch +from phonemizer.punctuation import Punctuation +from phonemizer.separator import Separator + + +class TextTokenizer: + """Phonemize Text.""" + + def __init__( + self, + language="en-us", + backend="espeak", + separator=Separator(word="|_|", syllable="-", phone="|"), + preserve_punctuation=True, + with_stress: bool = False, + tie: Union[bool, str] = False, + language_switch: LanguageSwitch = "remove-flags", + words_mismatch: WordMismatch = "ignore", + ) -> None: + self.preserve_punctuation_marks = ",.?!;:'…" + self.backend = EspeakBackend( + language, + punctuation_marks=self.preserve_punctuation_marks, + preserve_punctuation=preserve_punctuation, + with_stress=with_stress, + tie=tie, + language_switch=language_switch, + words_mismatch=words_mismatch, + ) + + self.separator = separator + + # convert chinese punctuation to english punctuation + def convert_chinese_punctuation(self, text: str) -> str: + text = text.replace("，", ",") + text = text.replace("。", ".") + text = text.replace("！", "!") + text = text.replace("？", "?") + text = text.replace("；", ";") + text = text.replace("：", ":") + text = text.replace("、", ",") + text = text.replace("‘", "'") + text = text.replace("’", "'") + text = text.replace("⋯", "…") + text = text.replace("···", "…") + text = text.replace("・・・", "…") + text = text.replace("...", "…") + return text + + def __call__(self, text, strip=True) -> List[str]: + + text_type = type(text) + normalized_text = [] + for line in str2list(text): + line = self.convert_chinese_punctuation(line.strip()) + line = re.sub(r"[^\w\s_,\.\?!;:\'…]", "", line) + line = re.sub(r"\s*([,\.\?!;:\'…])\s*", r"\1", line) + line = re.sub(r"\s+", " ", line) + normalized_text.append(line) + # print("Normalized test: ", normalized_text[0]) + phonemized = self.backend.phonemize( + normalized_text, separator=self.separator, strip=strip, njobs=1 + ) + if text_type == str: + phonemized = re.sub(r"([,\.\?!;:\'…])", r"|\1|", list2str(phonemized)) + phonemized = re.sub(r"\|+", "|", phonemized) + phonemized = phonemized.rstrip("|") + else: + for i in range(len(phonemized)): + phonemized[i] = re.sub(r"([,\.\?!;:\'…])", r"|\1|", phonemized[i]) + phonemized[i] = re.sub(r"\|+", "|", phonemized[i]) + phonemized[i] = phonemized[i].rstrip("|") + return phonemized diff --git a/models/tts/maskgct/g2p/g2p/vocab.json b/models/tts/maskgct/g2p/g2p/vocab.json new file mode 100644 index 0000000..28d32aa --- /dev/null +++ b/models/tts/maskgct/g2p/g2p/vocab.json @@ -0,0 +1,372 @@ +{ + "vocab": { + ",": 0, + ".": 1, + "?": 2, + "!": 3, + "_": 4, + "iː": 5, + "ɪ": 6, + "ɜː": 7, + "ɚ": 8, + "oːɹ": 9, + "ɔː": 10, + "ɔːɹ": 11, + "ɑː": 12, + "uː": 13, + "ʊ": 14, + "ɑːɹ": 15, + "ʌ": 16, + "ɛ": 17, + "æ": 18, + "eɪ": 19, + "aɪ": 20, + "ɔɪ": 21, + "aʊ": 22, + "oʊ": 23, + "ɪɹ": 24, + "ɛɹ": 25, + "ʊɹ": 26, + "p": 27, + "b": 28, + "t": 29, + "d": 30, + "k": 31, + "ɡ": 32, + "f": 33, + "v": 34, + "θ": 35, + "ð": 36, + "s": 37, + "z": 38, + "ʃ": 39, + "ʒ": 40, + "h": 41, + "tʃ": 42, + "dʒ": 43, + "m": 44, + "n": 45, + "ŋ": 46, + "j": 47, + "w": 48, + "ɹ": 49, + "l": 50, + "tɹ": 51, + "dɹ": 52, + "ts": 53, + "dz": 54, + "i": 55, + "ɔ": 56, + "ə": 57, + "ɾ": 58, + "iə": 59, + "r": 60, + "u": 61, + "oː": 62, + "ɛː": 63, + "ɪː": 64, + "aɪə": 65, + "aɪɚ": 66, + "ɑ̃": 67, + "ç": 68, + "ɔ̃": 69, + "ææ": 70, + "ɐɐ": 71, + "ɡʲ": 72, + "nʲ": 73, + "iːː": 74, + + "p⁼": 75, + "pʰ": 76, + "t⁼": 77, + "tʰ": 78, + "k⁼": 79, + "kʰ": 80, + "x": 81, + "tʃ⁼": 82, + "tʃʰ": 83, + "ts`⁼": 84, + "ts`ʰ": 85, + "s`": 86, + "ɹ`": 87, + "ts⁼": 88, + "tsʰ": 89, + "p⁼wo": 90, + "p⁼wo→": 91, + "p⁼wo↑": 92, + "p⁼wo↓↑": 93, + "p⁼wo↓": 94, + "pʰwo": 95, + "pʰwo→": 96, + "pʰwo↑": 97, + "pʰwo↓↑": 98, + "pʰwo↓": 99, + "mwo": 100, + "mwo→": 101, + "mwo↑": 102, + "mwo↓↑": 103, + "mwo↓": 104, + "fwo": 105, + "fwo→": 106, + "fwo↑": 107, + "fwo↓↑": 108, + "fwo↓": 109, + "jɛn": 110, + "jɛn→": 111, + "jɛn↑": 112, + "jɛn↓↑": 113, + "jɛn↓": 114, + "ɥæn": 115, + "ɥæn→": 116, + "ɥæn↑": 117, + "ɥæn↓↑": 118, + "ɥæn↓": 119, + "in": 120, + "in→": 121, + "in↑": 122, + "in↓↑": 123, + "in↓": 124, + "ɥn": 125, + "ɥn→": 126, + "ɥn↑": 127, + "ɥn↓↑": 128, + "ɥn↓": 129, + "iŋ": 130, + "iŋ→": 131, + "iŋ↑": 132, + "iŋ↓↑": 133, + "iŋ↓": 134, + "ʊŋ": 135, + "ʊŋ→": 136, + "ʊŋ↑": 137, + "ʊŋ↓↑": 138, + "ʊŋ↓": 139, + "jʊŋ": 140, + "jʊŋ→": 141, + "jʊŋ↑": 142, + "jʊŋ↓↑": 143, + "jʊŋ↓": 144, + "ia": 145, + "ia→": 146, + "ia↑": 147, + "ia↓↑": 148, + "ia↓": 149, + "iɛ": 150, + "iɛ→": 151, + "iɛ↑": 152, + "iɛ↓↑": 153, + "iɛ↓": 154, + "iɑʊ": 155, + "iɑʊ→": 156, + "iɑʊ↑": 157, + "iɑʊ↓↑": 158, + "iɑʊ↓": 159, + "ioʊ": 160, + "ioʊ→": 161, + "ioʊ↑": 162, + "ioʊ↓↑": 163, + "ioʊ↓": 164, + "iɑŋ": 165, + "iɑŋ→": 166, + "iɑŋ↑": 167, + "iɑŋ↓↑": 168, + "iɑŋ↓": 169, + "ua": 170, + "ua→": 171, + "ua↑": 172, + "ua↓↑": 173, + "ua↓": 174, + "uo": 175, + "uo→": 176, + "uo↑": 177, + "uo↓↑": 178, + "uo↓": 179, + "uaɪ": 180, + "uaɪ→": 181, + "uaɪ↑": 182, + "uaɪ↓↑": 183, + "uaɪ↓": 184, + "ueɪ": 185, + "ueɪ→": 186, + "ueɪ↑": 187, + "ueɪ↓↑": 188, + "ueɪ↓": 189, + "uan": 190, + "uan→": 191, + "uan↑": 192, + "uan↓↑": 193, + "uan↓": 194, + "uən": 195, + "uən→": 196, + "uən↑": 197, + "uən↓↑": 198, + "uən↓": 199, + "uɑŋ": 200, + "uɑŋ→": 201, + "uɑŋ↑": 202, + "uɑŋ↓↑": 203, + "uɑŋ↓": 204, + "ɥɛ": 205, + "ɥɛ→": 206, + "ɥɛ↑": 207, + "ɥɛ↓↑": 208, + "ɥɛ↓": 209, + "a": 210, + "a→": 211, + "a↑": 212, + "a↓↑": 213, + "a↓": 214, + "o": 215, + "o→": 216, + "o↑": 217, + "o↓↑": 218, + "o↓": 219, + "ə→": 220, + "ə↑": 221, + "ə↓↑": 222, + "ə↓": 223, + "ɛ→": 224, + "ɛ↑": 225, + "ɛ↓↑": 226, + "ɛ↓": 227, + "aɪ→": 228, + "aɪ↑": 229, + "aɪ↓↑": 230, + "aɪ↓": 231, + "eɪ→": 232, + "eɪ↑": 233, + "eɪ↓↑": 234, + "eɪ↓": 235, + "ɑʊ": 236, + "ɑʊ→": 237, + "ɑʊ↑": 238, + "ɑʊ↓↑": 239, + "ɑʊ↓": 240, + "oʊ→": 241, + "oʊ↑": 242, + "oʊ↓↑": 243, + "oʊ↓": 244, + "an": 245, + "an→": 246, + "an↑": 247, + "an↓↑": 248, + "an↓": 249, + "ən": 250, + "ən→": 251, + "ən↑": 252, + "ən↓↑": 253, + "ən↓": 254, + "ɑŋ": 255, + "ɑŋ→": 256, + "ɑŋ↑": 257, + "ɑŋ↓↑": 258, + "ɑŋ↓": 259, + "əŋ": 260, + "əŋ→": 261, + "əŋ↑": 262, + "əŋ↓↑": 263, + "əŋ↓": 264, + "əɹ": 265, + "əɹ→": 266, + "əɹ↑": 267, + "əɹ↓↑": 268, + "əɹ↓": 269, + "i→": 270, + "i↑": 271, + "i↓↑": 272, + "i↓": 273, + "u→": 274, + "u↑": 275, + "u↓↑": 276, + "u↓": 277, + "ɥ": 278, + "ɥ→": 279, + "ɥ↑": 280, + "ɥ↓↑": 281, + "ɥ↓": 282, + "ts`⁼ɹ": 283, + "ts`⁼ɹ→": 284, + "ts`⁼ɹ↑": 285, + "ts`⁼ɹ↓↑": 286, + "ts`⁼ɹ↓": 287, + "ts`ʰɹ": 288, + "ts`ʰɹ→": 289, + "ts`ʰɹ↑": 290, + "ts`ʰɹ↓↑": 291, + "ts`ʰɹ↓": 292, + "s`ɹ": 293, + "s`ɹ→": 294, + "s`ɹ↑": 295, + "s`ɹ↓↑": 296, + "s`ɹ↓": 297, + "ɹ`ɹ": 298, + "ɹ`ɹ→": 299, + "ɹ`ɹ↑": 300, + "ɹ`ɹ↓↑": 301, + "ɹ`ɹ↓": 302, + "ts⁼ɹ": 303, + "ts⁼ɹ→": 304, + "ts⁼ɹ↑": 305, + "ts⁼ɹ↓↑": 306, + "ts⁼ɹ↓": 307, + "tsʰɹ": 308, + "tsʰɹ→": 309, + "tsʰɹ↑": 310, + "tsʰɹ↓↑": 311, + "tsʰɹ↓": 312, + "sɹ": 313, + "sɹ→": 314, + "sɹ↑": 315, + "sɹ↓↑": 316, + "sɹ↓": 317, + + "ɯ": 318, + "e": 319, + "aː": 320, + "ɯː": 321, + "eː": 322, + "ç": 323, + "ɸ": 324, + "ɰᵝ": 325, + "ɴ": 326, + "g": 327, + "dʑ": 328, + "q": 329, + "ː": 330, + "bj": 331, + "tɕ": 332, + "dej": 333, + "tej": 334, + "gj": 335, + "gɯ": 336, + "çj": 337, + "kj": 338, + "kɯ": 339, + "mj": 340, + "nj": 341, + "pj": 342, + "ɾj": 343, + "ɕ": 344, + "tsɯ": 345, + + "ɐ": 346, + "ɑ": 347, + "ɒ": 348, + "ɜ": 349, + "ɫ": 350, + "ʑ": 351, + "ʲ": 352, + + "y": 353, + "ø": 354, + "œ": 355, + "ʁ": 356, + "̃": 357, + "ɲ": 358, + + ":": 359, + ";": 360, + "'": 361, + "…": 362 + } +} diff --git a/models/tts/maskgct/g2p/g2p_generation.py b/models/tts/maskgct/g2p/g2p_generation.py new file mode 100644 index 0000000..6f7acb5 --- /dev/null +++ b/models/tts/maskgct/g2p/g2p_generation.py @@ -0,0 +1,123 @@ +# Copyright (c) 2024 Amphion. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +import os +import sys + +from models.tts.maskgct.g2p.g2p import PhonemeBpeTokenizer +from models.tts.maskgct.g2p.utils.g2p import phonemizer_g2p +import tqdm +from typing import List +import json +import os +import re + + +def ph_g2p(text, language): + + return phonemizer_g2p(text=text, language=language) + + +def g2p(text, sentence, language): + + return text_tokenizer.tokenize(text=text, sentence=sentence, language=language) + + +def is_chinese(char): + if char >= "\u4e00" and char <= "\u9fa5": + return True + else: + return False + + +def is_alphabet(char): + if (char >= "\u0041" and char <= "\u005a") or ( + char >= "\u0061" and char <