goathead777/Zero_Shot_Inference

1# This is Multi-reference timbre encoder2 3import torch4from torch import nn5from torch.nn.utils import remove_weight_norm, weight_norm6from module.attentions import MultiHeadAttention7 8 9class MRTE(nn.Module):10    def __init__(11        self,12        content_enc_channels=192,13        hidden_size=512,14        out_channels=192,15        kernel_size=5,16        n_heads=4,17        ge_layer=2,18    ):19        super(MRTE, self).__init__()20        self.cross_attention = MultiHeadAttention(hidden_size, hidden_size, n_heads)21        self.c_pre = nn.Conv1d(content_enc_channels, hidden_size, 1)22        self.text_pre = nn.Conv1d(content_enc_channels, hidden_size, 1)23        self.c_post = nn.Conv1d(hidden_size, out_channels, 1)24 25    def forward(self, ssl_enc, ssl_mask, text, text_mask, ge, test=None):26        if ge == None:27            ge = 028        attn_mask = text_mask.unsqueeze(2) * ssl_mask.unsqueeze(-1)29 30        ssl_enc = self.c_pre(ssl_enc * ssl_mask)31        text_enc = self.text_pre(text * text_mask)32        if test != None:33            if test == 0:34                x = (35                    self.cross_attention(36                        ssl_enc * ssl_mask, text_enc * text_mask, attn_mask37                    )38                    + ssl_enc39                    + ge40                )41            elif test == 1:42                x = ssl_enc + ge43            elif test == 2:44                x = (45                    self.cross_attention(46                        ssl_enc * 0 * ssl_mask, text_enc * text_mask, attn_mask47                    )48                    + ge49                )50            else:51                raise ValueError("test should be 0,1,2")52        else:53            x = (54                self.cross_attention(55                    ssl_enc * ssl_mask, text_enc * text_mask, attn_mask56                )57                + ssl_enc58                + ge59            )60        x = self.c_post(x * ssl_mask)61        return x62 63 64class SpeakerEncoder(torch.nn.Module):65    def __init__(66        self,67        mel_n_channels=80,68        model_num_layers=2,69        model_hidden_size=256,70        model_embedding_size=256,71    ):72        super(SpeakerEncoder, self).__init__()73        self.lstm = nn.LSTM(74            mel_n_channels, model_hidden_size, model_num_layers, batch_first=True75        )76        self.linear = nn.Linear(model_hidden_size, model_embedding_size)77        self.relu = nn.ReLU()78 79    def forward(self, mels):80        self.lstm.flatten_parameters()81        _, (hidden, _) = self.lstm(mels.transpose(-1, -2))82        embeds_raw = self.relu(self.linear(hidden[-1]))83        return embeds_raw / torch.norm(embeds_raw, dim=1, keepdim=True)84 85 86class MELEncoder(nn.Module):87    def __init__(88        self,89        in_channels,90        out_channels,91        hidden_channels,92        kernel_size,93        dilation_rate,94        n_layers,95    ):96        super().__init__()97        self.in_channels = in_channels98        self.out_channels = out_channels99        self.hidden_channels = hidden_channels100        self.kernel_size = kernel_size101        self.dilation_rate = dilation_rate102        self.n_layers = n_layers103 104        self.pre = nn.Conv1d(in_channels, hidden_channels, 1)105        self.enc = WN(hidden_channels, kernel_size, dilation_rate, n_layers)106        self.proj = nn.Conv1d(hidden_channels, out_channels, 1)107 108    def forward(self, x):109        # print(x.shape,x_lengths.shape)110        x = self.pre(x)111        x = self.enc(x)112        x = self.proj(x)113        return x114 115 116class WN(torch.nn.Module):117    def __init__(self, hidden_channels, kernel_size, dilation_rate, n_layers):118        super(WN, self).__init__()119        assert kernel_size % 2 == 1120        self.hidden_channels = hidden_channels121        self.kernel_size = kernel_size122        self.dilation_rate = dilation_rate123        self.n_layers = n_layers124 125        self.in_layers = torch.nn.ModuleList()126        self.res_skip_layers = torch.nn.ModuleList()127 128        for i in range(n_layers):129            dilation = dilation_rate**i130            padding = int((kernel_size * dilation - dilation) / 2)131            in_layer = nn.Conv1d(132                hidden_channels,133                2 * hidden_channels,134                kernel_size,135                dilation=dilation,136                padding=padding,137            )138            in_layer = weight_norm(in_layer)139            self.in_layers.append(in_layer)140 141            # last one is not necessary142            if i < n_layers - 1:143                res_skip_channels = 2 * hidden_channels144            else:145                res_skip_channels = hidden_channels146 147            res_skip_layer = torch.nn.Conv1d(hidden_channels, res_skip_channels, 1)148            res_skip_layer = weight_norm(res_skip_layer, name="weight")149            self.res_skip_layers.append(res_skip_layer)150 151    def forward(self, x):152        output = torch.zeros_like(x)153        n_channels_tensor = torch.IntTensor([self.hidden_channels])154 155        for i in range(self.n_layers):156            x_in = self.in_layers[i](x)157 158            acts = fused_add_tanh_sigmoid_multiply(x_in, n_channels_tensor)159 160            res_skip_acts = self.res_skip_layers[i](acts)161            if i < self.n_layers - 1:162                res_acts = res_skip_acts[:, : self.hidden_channels, :]163                x = x + res_acts164                output = output + res_skip_acts[:, self.hidden_channels :, :]165            else:166                output = output + res_skip_acts167        return output168 169    def remove_weight_norm(self):170        for l in self.in_layers:171            remove_weight_norm(l)172        for l in self.res_skip_layers:173            remove_weight_norm(l)174 175 176@torch.jit.script177def fused_add_tanh_sigmoid_multiply(input, n_channels):178    n_channels_int = n_channels[0]179    t_act = torch.tanh(input[:, :n_channels_int, :])180    s_act = torch.sigmoid(input[:, n_channels_int:, :])181    acts = t_act * s_act182    return acts183 184 185if __name__ == "__main__":186    content_enc = torch.randn(3, 192, 100)187    content_mask = torch.ones(3, 1, 100)188    ref_mel = torch.randn(3, 128, 30)189    ref_mask = torch.ones(3, 1, 30)190    model = MRTE()191    out = model(content_enc, content_mask, ref_mel, ref_mask)192    print(out.shape)193