hanfish/LSai

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 8class MRTE(nn.Module):9    def __init__(self, 10                 content_enc_channels=192,11                 hidden_size=512,12                 out_channels=192,13                 kernel_size=5,14                 n_heads=4,15                 ge_layer = 216                 ):17        super(MRTE, self).__init__()18        self.cross_attention = MultiHeadAttention(hidden_size,hidden_size,n_heads)19        self.c_pre = nn.Conv1d(content_enc_channels,hidden_size, 1)20        self.text_pre = nn.Conv1d(content_enc_channels,hidden_size, 1)21        self.c_post = nn.Conv1d(hidden_size,out_channels, 1)22 23    def forward(self, ssl_enc, ssl_mask, text, text_mask, ge, test=None):24        if(ge==None):ge=025        attn_mask = text_mask.unsqueeze(2) * ssl_mask.unsqueeze(-1)26 27        ssl_enc = self.c_pre(ssl_enc * ssl_mask)28        text_enc = self.text_pre(text * text_mask)29        if test != None:30            if test == 0:31                x = self.cross_attention(ssl_enc * ssl_mask, text_enc * text_mask, attn_mask) + ssl_enc + ge32            elif test == 1:33                x = ssl_enc + ge34            elif test ==2:35                x = self.cross_attention(ssl_enc*0 * ssl_mask, text_enc * text_mask, attn_mask) + ge36            else:37                raise ValueError("test should be 0,1,2")38        else:39            x = self.cross_attention(ssl_enc * ssl_mask, text_enc * text_mask, attn_mask) + ssl_enc + ge40        x = self.c_post(x * ssl_mask)41        return x42        43 44class SpeakerEncoder(torch.nn.Module):45    def __init__(self, mel_n_channels=80, model_num_layers=2, model_hidden_size=256, model_embedding_size=256):46        super(SpeakerEncoder, self).__init__()47        self.lstm = nn.LSTM(mel_n_channels, model_hidden_size, model_num_layers, batch_first=True)48        self.linear = nn.Linear(model_hidden_size, model_embedding_size)49        self.relu = nn.ReLU()50 51    def forward(self, mels):52        self.lstm.flatten_parameters()53        _, (hidden, _) = self.lstm(mels.transpose(-1, -2))54        embeds_raw = self.relu(self.linear(hidden[-1]))55        return embeds_raw / torch.norm(embeds_raw, dim=1, keepdim=True)56 57 58class MELEncoder(nn.Module):59    def __init__(self,60                 in_channels,61                 out_channels,62                 hidden_channels,63                 kernel_size,64                 dilation_rate,65                 n_layers):66        super().__init__()67        self.in_channels = in_channels68        self.out_channels = out_channels69        self.hidden_channels = hidden_channels70        self.kernel_size = kernel_size71        self.dilation_rate = dilation_rate72        self.n_layers = n_layers73 74        self.pre = nn.Conv1d(in_channels, hidden_channels, 1)75        self.enc = WN(hidden_channels, kernel_size, dilation_rate, n_layers)76        self.proj = nn.Conv1d(hidden_channels, out_channels, 1)77 78    def forward(self, x):79        # print(x.shape,x_lengths.shape)80        x = self.pre(x)81        x = self.enc(x)82        x = self.proj(x)83        return x84    85 86class WN(torch.nn.Module):87  def __init__(self, hidden_channels, kernel_size, dilation_rate, n_layers):88    super(WN, self).__init__()89    assert(kernel_size % 2 == 1)90    self.hidden_channels =hidden_channels91    self.kernel_size = kernel_size92    self.dilation_rate = dilation_rate93    self.n_layers = n_layers94 95    self.in_layers = torch.nn.ModuleList()96    self.res_skip_layers = torch.nn.ModuleList()97 98    for i in range(n_layers):99      dilation = dilation_rate ** i100      padding = int((kernel_size * dilation - dilation) / 2)101      in_layer = nn.Conv1d(hidden_channels, 2*hidden_channels, kernel_size,102                                 dilation=dilation, padding=padding)103      in_layer = weight_norm(in_layer)104      self.in_layers.append(in_layer)105 106      # last one is not necessary107      if i < n_layers - 1:108        res_skip_channels = 2 * hidden_channels109      else:110        res_skip_channels = hidden_channels111 112      res_skip_layer = torch.nn.Conv1d(hidden_channels, res_skip_channels, 1)113      res_skip_layer = weight_norm(res_skip_layer, name='weight')114      self.res_skip_layers.append(res_skip_layer)115 116  def forward(self, x):117    output = torch.zeros_like(x)118    n_channels_tensor = torch.IntTensor([self.hidden_channels])119 120    for i in range(self.n_layers):121      x_in = self.in_layers[i](x)122 123      acts = fused_add_tanh_sigmoid_multiply(124          x_in,125          n_channels_tensor)126 127      res_skip_acts = self.res_skip_layers[i](acts)128      if i < self.n_layers - 1:129        res_acts = res_skip_acts[:,:self.hidden_channels,:]130        x = (x + res_acts)131        output = output + res_skip_acts[:,self.hidden_channels:,:]132      else:133        output = output + res_skip_acts134    return output135 136  def remove_weight_norm(self):137    for l in self.in_layers:138      remove_weight_norm(l)139    for l in self.res_skip_layers:140      remove_weight_norm(l)141 142 143@torch.jit.script144def fused_add_tanh_sigmoid_multiply(input, n_channels):145  n_channels_int = n_channels[0]146  t_act = torch.tanh(input[:, :n_channels_int, :])147  s_act = torch.sigmoid(input[:, n_channels_int:, :])148  acts = t_act * s_act149  return acts150 151 152 153if __name__ == '__main__':154    content_enc = torch.randn(3,192,100)155    content_mask = torch.ones(3,1,100)156    ref_mel = torch.randn(3,128,30)157    ref_mask = torch.ones(3,1,30)158    model = MRTE()159    out = model(content_enc,content_mask,ref_mel,ref_mask)160    print(out.shape)