model.py

from typing import Union

import numpy as np
import torch
import torchaudio
import torch.nn as nn
import torchaudio.transforms as transforms
from transformers import PretrainedConfig, PreTrainedModel

import dac
from audiotools import AudioSignal


def freeze(model):
    for param in model.parameters():
        param.requires_grad = False


class DACConfig(PretrainedConfig):
    model_type = 'dac'

    def __init__(self, 
                 model_type_by_sampling_freq:str='16khz',
                 encoding_chunk_size_in_sec:int=1,
                 decoding_chunk_rate:float=0.1,
                 decoding_overlap_rate:float=0.1,
                 **kwargs):
        super().__init__(**kwargs)
        """
        Initializes the model object.
        Args:
            model_type_by_sampling_freq (str, optional): The model type based on the sampling frequency. Defaults to '44khz'. Choose among ['44khz', '24khz', '16khz']
            encoding_chunk_size_in_sec (int, optional): The size of the encoding chunk in seconds. Defaults to 1.
            decoding_chunk_rate (float, optional): The decoding chunk rate. Must be between 0 and 1. Defaults to 0.1.
            decoding_overlap_rate (float, optional): The decoding overlap rate. Must be between 0 and 1. Defaults to 0.1.
            **kwargs: Additional keyword arguments.
        Raises:
            AssertionError: If the model_type_by_sampling_freq is not one of ['44khz', '24khz', '16khz'].
            AssertionError: If the decoding_chunk_rate is not between 0 and 1.
            AssertionError: If the decoding_overlap_rate is not between 0 and 1.
        """
        self.model_type_by_sampling_freq = model_type_by_sampling_freq
        self.encoding_chunk_size_in_sec = encoding_chunk_size_in_sec
        self.decoding_chunk_rate = decoding_chunk_rate
        self.decoding_overlap_rate = decoding_overlap_rate

        assert model_type_by_sampling_freq.lower() in ['44khz', '24khz', '16khz']
        assert decoding_chunk_rate > 0 and decoding_chunk_rate <= 1.0, '`decoding_chunk_rate` must be bewteen 0 and 1.'
        assert decoding_overlap_rate >= 0 and decoding_overlap_rate < 1.0, '`decoding_overlap_rate` must be bewteen 0 and 1.'



class DAC(PreTrainedModel):
    config_class = DACConfig

    def __init__(self, config):
        super().__init__(config)

        self.model_type_by_sampling_freq = config.model_type_by_sampling_freq.lower()
        self.model_type_by_sampling_freq_int = {'44khz':44100, '24khz':24000, '16khz':16000}[self.model_type_by_sampling_freq]
        self.encoding_chunk_size_in_sec = config.encoding_chunk_size_in_sec
        self.decoding_chunk_rate = config.decoding_chunk_rate
        self.decoding_overlap_rate = config.decoding_overlap_rate


        dac_path = dac.utils.download(model_type=self.model_type_by_sampling_freq)
        self.dac = dac.DAC.load(dac_path)
        self.dac.eval()
        freeze(self.dac)

        self.downsampling_rate = int(np.prod(self.dac.encoder_rates))  # 512
    
    def load_audio(self, filename:str):
        waveform, sample_rate = torchaudio.load(filename)  # waveform: (n_channels, length); sample_rate: const.
        return waveform, sample_rate
    
    def resample_audio(self, waveform:torch.FloatTensor, orig_sr:int, target_sr:int):
        """
        - sr: sampling rate
        - waveform: (n_channels, length)
        """
        if orig_sr == target_sr:
            return waveform

        converter = transforms.Resample(orig_freq=orig_sr, new_freq=target_sr)
        waveform = converter(waveform)  # (n_channels, new_length)
        return waveform  # (n_channels, new_length)

    def to_mono_channel(self, waveform:torch.FloatTensor):
        """
        - waveform: (n_channels, length)
        """
        n_channels = waveform.shape[0]
        if n_channels > 1:
            waveform = torch.mean(waveform, dim=0, keepdim=True)  # (1, length)
        return waveform  # (1, length)
    
    @torch.no_grad()
    def encode(self, audio_fname:str):
        self.eval()

        waveform, sr = self.load_audio(audio_fname)
        waveform = self.resample_audio(waveform, orig_sr=sr, target_sr=self.model_type_by_sampling_freq_int)
        sr = self.model_type_by_sampling_freq_int
        waveform = self.to_mono_channel(waveform)  # DAC accepts a mono channel only.
        
        zq, s = self._chunk_encoding(waveform, sr)
        return zq, s

    def _chunk_encoding(self, waveform:torch.FloatTensor, sr:int):
        # TODO: can I make it parallel?
        """
        waveform: (c l)
        """
        x = waveform  # brief varname
        x = x.unsqueeze(1)  # (b 1 l); add a null batch dim
        chunk_size = int(self.encoding_chunk_size_in_sec * sr)

        # adjust `chunk_size` to prevent any padding in `dac.preprocess`, which causes a gap between the mini-batches in the resulting music.
        remainer = chunk_size % self.dac.hop_length
        chunk_size = chunk_size-remainer

        # process
        zq_list, s_list = [], []
        audio_length = x.shape[-1]
        for start in range(0, audio_length, chunk_size):
            end = start + chunk_size
            chunk = x[:, :, start:end]
            chunk = self.dac.preprocess(chunk, sr)
            zq, s, _, _, _ = self.dac.encode(chunk.to(self.device))
            zq = zq.cpu()
            s = s.cpu()
            """
            "zq" : Tensor[B x D x T]
                Quantized continuous representation of input
                = summation of all the residual quantized vectors across every rvq level
                = E(x) = z = \sum_n^N{zq_n} where N is the number of codebooks
            "s" : Tensor[B x N x T]
                Codebook indices for each codebook
                (quantized discrete representation of input)
                *first element in the N dimension = first RVQ level
            """
            zq_list.append(zq)
            s_list.append(s)
            torch.cuda.empty_cache()
        
        zq = torch.cat(zq_list, dim=2).float()  # (1, d, length)
        s = torch.cat(s_list, dim=2).long()  # (1, n_rvq, length)

        return zq, s
    
    @torch.no_grad()
    def decode(self, *, zq:Union[torch.FloatTensor,None]=None, s:Union[torch.IntTensor,None]=None):
        """
        zq: (b, d, length)
        """
        if isinstance(zq,type(None)) and isinstance(s,type(None)):
            assert False, 'one of them must be valid.'
        self.eval()

        if not isinstance(zq,type(None)):
            waveform = self._chunk_decoding(zq)  # (b, 1, length); output always has a mono-channel.
        if not isinstance(s,type(None)):
            zq = self.code_to_zq(s)
            waveform = self._chunk_decoding(zq)  # (b, 1, length); output always has a mono-channel.

        return waveform
    
    def _chunk_decoding(self, zq:torch.FloatTensor):
        """
        zq: (b, d, length)
        """
        length = zq.shape[-1]
        chunk_size = round(int(self.decoding_chunk_rate * length))
        overlap_size = round(self.decoding_overlap_rate * chunk_size)  # overlap size in terms of token length
        overlap_size_in_data_space = round(overlap_size * self.downsampling_rate)
        waveform_concat = None
        for start in range(0, length, chunk_size-overlap_size):
            end = start + chunk_size
            chunk = zq[:,:, start:end]  # (b, d, chunk_size)
            waveform = self.dac.decode(chunk.to(self.device))  # (b, 1, chunk_size*self.downsampling_rate)
            waveform = waveform.cpu()

            waveform_len = waveform.shape[-1]
            if waveform_len < overlap_size_in_data_space:
                overlap_size_in_data_space = waveform_len
            
            if isinstance(waveform_concat, type(None)):
                waveform_concat = waveform.clone()
            else:
                if self.decoding_overlap_rate != 0.:
                    prev_x = waveform_concat[:,:,:-overlap_size_in_data_space]
                    rest_of_new_x = waveform[:,:,overlap_size_in_data_space:]
                    overlap_x_from_prev_x = waveform_concat[:,:,-overlap_size_in_data_space:]  # (b, 1, overlap_size_in_data_space)
                    overlap_x_from_new_x = waveform[:,:,:overlap_size_in_data_space]  # (b, 1, overlap_size_in_data_space)
                    if not overlap_x_from_new_x.shape[-1] == 0:
                        overlap = (overlap_x_from_prev_x + overlap_x_from_new_x) / 2  # take mean; maybe there's a better strategy but it seems to work fine.
                    else:
                        overlap = overlap_x_from_prev_x
                    waveform_concat = torch.cat((prev_x, overlap, rest_of_new_x), dim=-1)  # (b, 1, ..)
                else:
                    prev_x = waveform_concat
                    rest_of_new_x = waveform
                    waveform_concat = torch.cat((prev_x, rest_of_new_x), dim=-1)  # (b, 1, ..)
        return waveform_concat  # (b, 1, length)

    def code_to_zq(self, s:torch.IntTensor):
        """
        s: (b, n_rvq, length)
        """
        zq, _, _ = self.dac.quantizer.from_codes(s.to(self.device))  # zq: (b, d, length)
        zq = zq.cpu()
        return zq

    def save_tensor(self, tensor:torch.Tensor, fname:str) -> None:
        torch.save(tensor.cpu(), fname)
    
    def load_tensor(self, fname:str):
        return torch.load(fname)
    
    def waveform_to_audiofile(self, waveform:torch.FloatTensor, fname:str) -> None:
        AudioSignal(waveform, sample_rate=self.model_type_by_sampling_freq_int).write(fname)