| 135 | |
| 136 | |
| 137 | class DurationEncoder(nn.Module): |
| 138 | def __init__(self, sty_dim, d_model, nlayers, dropout=0.1): |
| 139 | super().__init__() |
| 140 | self.lstms = nn.ModuleList() |
| 141 | for _ in range(nlayers): |
| 142 | self.lstms.append(nn.LSTM(d_model + sty_dim, d_model // 2, num_layers=1, batch_first=True, bidirectional=True, dropout=dropout)) |
| 143 | self.lstms.append(AdaLayerNorm(sty_dim, d_model)) |
| 144 | self.dropout = dropout |
| 145 | self.d_model = d_model |
| 146 | self.sty_dim = sty_dim |
| 147 | |
| 148 | def forward(self, x, style, text_lengths, m): |
| 149 | masks = m |
| 150 | x = x.permute(2, 0, 1) |
| 151 | s = style.expand(x.shape[0], x.shape[1], -1) |
| 152 | x = torch.cat([x, s], axis=-1) |
| 153 | x.masked_fill_(masks.unsqueeze(-1).transpose(0, 1), 0.0) |
| 154 | x = x.transpose(0, 1) |
| 155 | x = x.transpose(-1, -2) |
| 156 | for block in self.lstms: |
| 157 | if isinstance(block, AdaLayerNorm): |
| 158 | x = block(x.transpose(-1, -2), style).transpose(-1, -2) |
| 159 | x = torch.cat([x, s.permute(1, 2, 0)], axis=1) |
| 160 | x.masked_fill_(masks.unsqueeze(-1).transpose(-1, -2), 0.0) |
| 161 | else: |
| 162 | lengths = text_lengths if text_lengths.device == torch.device('cpu') else text_lengths.to('cpu') |
| 163 | x = x.transpose(-1, -2) |
| 164 | x = nn.utils.rnn.pack_padded_sequence( |
| 165 | x, lengths, batch_first=True, enforce_sorted=False) |
| 166 | block.flatten_parameters() |
| 167 | x, _ = block(x) |
| 168 | x, _ = nn.utils.rnn.pad_packed_sequence( |
| 169 | x, batch_first=True) |
| 170 | x = F.dropout(x, p=self.dropout, training=False) |
| 171 | x = x.transpose(-1, -2) |
| 172 | x_pad = torch.zeros([x.shape[0], x.shape[1], m.shape[-1]], device=x.device) |
| 173 | x_pad[:, :, :x.shape[-1]] = x |
| 174 | x = x_pad |
| 175 | |
| 176 | return x.transpose(-1, -2) |
| 177 | |
| 178 | |
| 179 | # https://github.com/yl4579/StyleTTS2/blob/main/Utils/PLBERT/util.py |