| import torch |
| from audiotools import AudioSignal |
| from audiotools.ml import BaseModel |
| from encodec import EncodecModel |
|
|
|
|
| class Encodec(BaseModel): |
| def __init__(self, sample_rate: int = 24000, bandwidth: float = 24.0): |
| super().__init__() |
|
|
| if sample_rate == 24000: |
| self.model = EncodecModel.encodec_model_24khz() |
| else: |
| self.model = EncodecModel.encodec_model_48khz() |
| self.model.set_target_bandwidth(bandwidth) |
| self.sample_rate = 44100 |
|
|
| def forward( |
| self, |
| audio_data: torch.Tensor, |
| sample_rate: int = 44100, |
| n_quantizers: int = None, |
| ): |
| signal = AudioSignal(audio_data, sample_rate) |
| signal.resample(self.model.sample_rate) |
| recons = self.model(signal.audio_data) |
| recons = AudioSignal(recons, self.model.sample_rate) |
| recons.resample(sample_rate) |
| return {"audio": recons.audio_data} |
|
|
|
|
| if __name__ == "__main__": |
| import numpy as np |
| from functools import partial |
|
|
| model = Encodec() |
|
|
| for n, m in model.named_modules(): |
| o = m.extra_repr() |
| p = sum([np.prod(p.size()) for p in m.parameters()]) |
| fn = lambda o, p: o + f" {p/1e6:<.3f}M params." |
| setattr(m, "extra_repr", partial(fn, o=o, p=p)) |
| print(model) |
| print("Total # of params: ", sum([np.prod(p.size()) for p in model.parameters()])) |
|
|
| length = 88200 * 2 |
| x = torch.randn(1, 1, length).to(model.device) |
| x.requires_grad_(True) |
| x.retain_grad() |
|
|
| |
| out = model(x)["audio"] |
|
|
| print(x.shape, out.shape) |
|
|
