| import torch |
| from torch import Tensor |
| from transformers import PreTrainedModel |
| from audio_diffusion_pytorch import DiffusionVocoder, UNetV0, VDiffusion, VSampler |
| from .config import VocoderConfig |
|
|
|
|
| class Vocoder(PreTrainedModel): |
|
|
| config_class = VocoderConfig |
|
|
| def __init__(self, config: VocoderConfig): |
| super().__init__(config) |
|
|
| self.model = DiffusionVocoder( |
| net_t=UNetV0, |
| mel_channels=80, |
| mel_n_fft=1024, |
| mel_sample_rate=48000, |
| mel_normalize_log=True, |
| channels=[8, 32, 64, 256, 256, 512, 512, 1024, 1024], |
| factors=[1, 4, 4, 4, 2, 2, 2, 2, 2], |
| items=[1, 2, 2, 2, 2, 2, 2, 4, 4], |
| diffusion_t=VDiffusion, |
| sampler_t=VSampler |
| ) |
| |
| def to_spectrogram(self, *args, **kwargs): |
| return self.model.to_spectrogram(*args, **kwargs) |
|
|
| @torch.no_grad() |
| def sample(self, *args, **kwargs): |
| return self.model.sample(*args, **kwargs) |
| |
|
|
|
|
