| import math |
|
|
| import torch |
| import torchaudio |
| from transformers.models.dac import DacModel |
|
|
|
|
| class DACAutoencoder: |
| def __init__(self): |
| super().__init__() |
| self.dac = DacModel.from_pretrained("descript/dac_44khz") |
| self.dac.eval().requires_grad_(False) |
| self.codebook_size = self.dac.config.codebook_size |
| self.num_codebooks = self.dac.quantizer.n_codebooks |
| self.sampling_rate = self.dac.config.sampling_rate |
|
|
| def preprocess(self, wav: torch.Tensor, sr: int) -> torch.Tensor: |
| wav = torchaudio.functional.resample(wav, sr, 44_100) |
| right_pad = math.ceil(wav.shape[-1] / 512) * 512 - wav.shape[-1] |
| return torch.nn.functional.pad(wav, (0, right_pad)) |
|
|
| def encode(self, wav: torch.Tensor) -> torch.Tensor: |
| return self.dac.encode(wav).audio_codes |
|
|
| def decode(self, codes: torch.Tensor) -> torch.Tensor: |
| return self.dac.decode(audio_codes=codes).audio_values.unsqueeze(1) |
|
|
