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)