SpongeBobFan2002/Zonosz

1import math2 3import torch4import torchaudio5from transformers.models.dac import DacModel6 7 8class DACAutoencoder:9    def __init__(self):10        super().__init__()11        self.dac = DacModel.from_pretrained("descript/dac_44khz")12        self.dac.eval().requires_grad_(False)13        self.codebook_size = self.dac.config.codebook_size14        self.num_codebooks = self.dac.quantizer.n_codebooks15        self.sampling_rate = self.dac.config.sampling_rate16 17    def preprocess(self, wav: torch.Tensor, sr: int) -> torch.Tensor:18        wav = torchaudio.functional.resample(wav, sr, 44_100)19        right_pad = math.ceil(wav.shape[-1] / 512) * 512 - wav.shape[-1]20        return torch.nn.functional.pad(wav, (0, right_pad))21 22    def encode(self, wav: torch.Tensor) -> torch.Tensor:23        return self.dac.encode(wav).audio_codes24 25    def decode(self, codes: torch.Tensor) -> torch.Tensor:26        return self.dac.decode(audio_codes=codes).audio_values.unsqueeze(1)27