| 802 | |
| 803 | |
| 804 | class DAC(BaseModel, CodecMixin): |
| 805 | def __init__( |
| 806 | self, |
| 807 | encoder_dim: int = 64, |
| 808 | encoder_rates: List[int] = [2, 4, 8, 8], |
| 809 | latent_dim: int = None, |
| 810 | decoder_dim: int = 1536, |
| 811 | decoder_rates: List[int] = [8, 8, 4, 2], |
| 812 | n_codebooks: int = 9, |
| 813 | codebook_size: int = 1024, |
| 814 | codebook_dim: Union[int, list] = 8, |
| 815 | quantizer_dropout: bool = False, |
| 816 | sample_rate: int = 44100, |
| 817 | continuous: bool = False, |
| 818 | ): |
| 819 | super().__init__() |
| 820 | |
| 821 | self.encoder_dim = encoder_dim |
| 822 | self.encoder_rates = encoder_rates |
| 823 | self.decoder_dim = decoder_dim |
| 824 | self.decoder_rates = decoder_rates |
| 825 | self.sample_rate = sample_rate |
| 826 | self.continuous = continuous |
| 827 | |
| 828 | if latent_dim is None: |
| 829 | latent_dim = encoder_dim * (2 ** len(encoder_rates)) |
| 830 | |
| 831 | self.latent_dim = latent_dim |
| 832 | |
| 833 | self.hop_length = np.prod(encoder_rates) |
| 834 | self.encoder = Encoder(encoder_dim, encoder_rates, latent_dim) |
| 835 | |
| 836 | if not continuous: |
| 837 | self.n_codebooks = n_codebooks |
| 838 | self.codebook_size = codebook_size |
| 839 | self.codebook_dim = codebook_dim |
| 840 | self.quantizer = ResidualVectorQuantize( |
| 841 | input_dim=latent_dim, |
| 842 | n_codebooks=n_codebooks, |
| 843 | codebook_size=codebook_size, |
| 844 | codebook_dim=codebook_dim, |
| 845 | quantizer_dropout=quantizer_dropout, |
| 846 | ) |
| 847 | else: |
| 848 | self.quant_conv = torch.nn.Conv1d(latent_dim, 2 * latent_dim, 1) |
| 849 | self.post_quant_conv = torch.nn.Conv1d(latent_dim, latent_dim, 1) |
| 850 | |
| 851 | self.decoder = Decoder( |
| 852 | latent_dim, |
| 853 | decoder_dim, |
| 854 | decoder_rates, |
| 855 | ) |
| 856 | self.sample_rate = sample_rate |
| 857 | self.apply(init_weights) |
| 858 | |
| 859 | self.delay = self.get_delay() |
| 860 | |
| 861 | @property |