| 748 | |
| 749 | class CodePredictor(nn.Module): |
| 750 | def __init__( |
| 751 | self, |
| 752 | hidden_channels, |
| 753 | filter_channels, |
| 754 | n_heads, |
| 755 | n_layers, |
| 756 | kernel_size, |
| 757 | p_dropout, |
| 758 | n_q=8, |
| 759 | dims=1024, |
| 760 | ssl_dim=768, |
| 761 | ): |
| 762 | super().__init__() |
| 763 | self.hidden_channels = hidden_channels |
| 764 | self.filter_channels = filter_channels |
| 765 | self.n_heads = n_heads |
| 766 | self.n_layers = n_layers |
| 767 | self.kernel_size = kernel_size |
| 768 | self.p_dropout = p_dropout |
| 769 | |
| 770 | self.vq_proj = nn.Conv1d(ssl_dim, hidden_channels, 1) |
| 771 | self.ref_enc = modules.MelStyleEncoder( |
| 772 | ssl_dim, style_vector_dim=hidden_channels |
| 773 | ) |
| 774 | |
| 775 | self.encoder = attentions.Encoder( |
| 776 | hidden_channels, filter_channels, n_heads, n_layers, kernel_size, p_dropout |
| 777 | ) |
| 778 | |
| 779 | self.out_proj = nn.Conv1d(hidden_channels, (n_q - 1) * dims, 1) |
| 780 | self.n_q = n_q |
| 781 | self.dims = dims |
| 782 | |
| 783 | def forward(self, x, x_mask, refer, codes, infer=False): |
| 784 | x = x.detach() |