Encode given audio data and return quantized latent codes Parameters ---------- audio_data : Tensor[B x 1 x T] Audio data to encode n_quantizers : int, optional Number of quantizers to use, by default None If None, all quantizers a
(
self,
audio_data: torch.Tensor,
n_quantizers: int = None,
)
| 886 | return audio_data |
| 887 | |
| 888 | def encode( |
| 889 | self, |
| 890 | audio_data: torch.Tensor, |
| 891 | n_quantizers: int = None, |
| 892 | ): |
| 893 | """Encode given audio data and return quantized latent codes |
| 894 | |
| 895 | Parameters |
| 896 | ---------- |
| 897 | audio_data : Tensor[B x 1 x T] |
| 898 | Audio data to encode |
| 899 | n_quantizers : int, optional |
| 900 | Number of quantizers to use, by default None |
| 901 | If None, all quantizers are used. |
| 902 | |
| 903 | Returns |
| 904 | ------- |
| 905 | dict |
| 906 | A dictionary with the following keys: |
| 907 | "z" : Tensor[B x D x T] |
| 908 | Quantized continuous representation of input |
| 909 | "codes" : Tensor[B x N x T] |
| 910 | Codebook indices for each codebook |
| 911 | (quantized discrete representation of input) |
| 912 | "latents" : Tensor[B x N*D x T] |
| 913 | Projected latents (continuous representation of input before quantization) |
| 914 | "vq/commitment_loss" : Tensor[1] |
| 915 | Commitment loss to train encoder to predict vectors closer to codebook |
| 916 | entries |
| 917 | "vq/codebook_loss" : Tensor[1] |
| 918 | Codebook loss to update the codebook |
| 919 | "length" : int |
| 920 | Number of samples in input audio |
| 921 | """ |
| 922 | z = self.encoder(audio_data) # [B x D x T] |
| 923 | if not self.continuous: |
| 924 | z, codes, latents, commitment_loss, codebook_loss = self.quantizer(z, n_quantizers) |
| 925 | else: |
| 926 | z = self.quant_conv(z) # [B x 2D x T] |
| 927 | z = DiagonalGaussianDistribution(z) |
| 928 | codes, latents, commitment_loss, codebook_loss = None, None, 0, 0 |
| 929 | |
| 930 | return z, codes, latents, commitment_loss, codebook_loss |
| 931 | |
| 932 | def decode(self, z: torch.Tensor): |
| 933 | """Decode given latent codes and return audio data |
no test coverage detected