first commit

.gitignore

@@ -0,0 +1,2 @@
+__pycache__
+.gradio

app.py

@@ -0,0 +1,168 @@
+import torch
+import spaces
+import tempfile
+import soundfile as sf
+import gradio as gr
+import librosa as lb
+import yaml
+import numpy as np
+import matplotlib.pyplot as plt
+from model.cleanmel import CleanMel
+from model.vocos.pretrained import Vocos
+from model.stft import InputSTFT, TargetMel
+
+DEVICE = torch.device("cuda:5")
+
+def read_audio(file_path):
+    audio, sample_rate = sf.read(file_path)
+    if audio.ndim > 1:
+        audio = audio[:, 0]
+    if sample_rate != 16000:
+        audio = lb.resample(audio, orig_sr=sample_rate, target_sr=16000)
+        sample_rate = 16000
+
+    return torch.tensor(audio).float().squeeze().unsqueeze(0)
+
+def stft(audio):
+    transform = InputSTFT(
+            n_fft=512,
+            n_win=512,
+            n_hop=128,
+            normalize=False,
+            center=True,
+            onesided=True,
+            online=False
+        ).eval().to(DEVICE)
+    return transform(audio)
+
+def mel_transform(audio, X_norm):
+    transform = TargetMel(
+        sample_rate=16000,
+        n_fft=512,
+        n_win=512,
+        n_hop=128,
+        n_mels=80,
+        f_min=0,
+        f_max=8000,
+        power=2,
+        center=True,
+        normalize=False,
+        onesided=True,
+        mel_norm="slaney",
+        mel_scale="slaney",
+        librosa_mel=True,
+        online=False
+    ).eval().to(DEVICE)
+    return transform(audio, X_norm)
+
+def load_cleanmel(model_name):
+    model_config = f"./configs/cleanmel_offline.yaml"
+    model_config = yaml.safe_load(open(model_config, "r"))["model"]["arch"]["init_args"]
+    cleanmel = CleanMel(**model_config)
+    cleanmel.load_state_dict(torch.load(f"./ckpts/CleanMel/{model_name}.ckpt"))
+    return cleanmel.eval()
+
+def load_vocos():
+    vocos = Vocos.from_hparams(config_path="./configs/vocos_offline.yaml")
+    vocos = Vocos.from_pretrained(None, model_path=f"./ckpts/Vocos/vocos_offline.pt", model=vocos)
+    return vocos.eval()
+
+def get_mrm_pred(Y_hat, x, X_norm):
+    X_noisy = mel_transform(x, X_norm)
+    Y_hat = Y_hat.squeeze()
+    Y_hat = torch.square(Y_hat * (torch.sqrt(X_noisy) + 1e-10))
+    return Y_hat
+
+def safe_log(x):           
+    return torch.log(torch.clip(x, min=1e-5)) 
+
+def output(y_hat, logMel_hat):
+    with tempfile.NamedTemporaryFile(suffix='.wav', delete=False) as tmp_file:
+        sf.write(tmp_file.name, y_hat.squeeze().cpu().numpy(), 16000)
+    with tempfile.NamedTemporaryFile(suffix='.npy', delete=False) as tmp_logmel_np_file:
+        np.save(tmp_logmel_np_file.name, logMel_hat.squeeze().cpu().numpy())
+    logMel_img = logMel_hat.squeeze().cpu().numpy()[::-1, :]
+    with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as tmp_logmel_img:
+        # give a plt figure size according to the logMel shape
+        plt.figure(figsize=(logMel_img.shape[1] / 100, logMel_img.shape[0] / 50))
+        plt.clf()
+        plt.imshow(logMel_img, vmin=-11, cmap="jet")
+        plt.tight_layout()
+        plt.ylabel("Mel bands")
+        plt.xlabel("Time (second)")
+        plt.yticks([0, 80], [80, 0])
+        dur = y_hat.shape[-1] / 16000
+        xticks = [int(x) for x in np.linspace(0, logMel_img.shape[-1], 11)]
+        xticks_str = ["{:.1f}".format(x) for x in np.linspace(0, dur, 11)]
+        plt.xticks(xticks, xticks_str)
+        plt.savefig(tmp_logmel_img.name)
+    
+    return tmp_file.name, tmp_logmel_img.name, tmp_logmel_np_file.name
+
+@spaces.GPU
+@torch.inference_mode()
+def enhance_cleanmel_L_mask(audio_path):
+    model = load_cleanmel("offline_CleanMel_L_mask").to(DEVICE)
+    vocos = load_vocos().to(DEVICE)
+    x = read_audio(audio_path).to(DEVICE)
+    X, X_norm = stft(x)
+    Y_hat = model(X, inference=True)
+    MRM_hat = torch.sigmoid(Y_hat)
+    Y_hat = get_mrm_pred(MRM_hat, x, X_norm)
+    logMel_hat = safe_log(Y_hat)
+    y_hat = vocos(logMel_hat, X_norm).clamp(min=-1, max=1)
+    return output(y_hat, logMel_hat)
+
+@spaces.GPU
+@torch.inference_mode()
+def enhance_cleanmel_L_map(audio_path):
+    model = load_cleanmel("offline_CleanMel_L_map").to(DEVICE)
+    vocos = load_vocos().to(DEVICE)
+    x = read_audio(audio_path).to(DEVICE)
+    X, X_norm = stft(x)
+    logMel_hat = model(X, inference=True)
+    y_hat = vocos(logMel_hat, X_norm).clamp(min=-1, max=1)
+    return output(y_hat, logMel_hat)
+
+def reset_everything():
+    """Reset all components to initial state"""
+    return None, None, None
+
+if __name__ == "__main__":
+    demo = gr.Blocks()
+    with gr.Blocks(title="CleanMel Demo") as demo:
+        gr.Markdown("## CleanMel Demo")
+        gr.Markdown("This demo showcases the CleanMel model for speech enhancement.")
+        
+        with gr.Row():
+            audio_input = gr.Audio(label="Input Audio", type="filepath", sources="upload")
+            with gr.Column():
+                enhance_button_map = gr.Button("Enhance Audio (offline CleanMel_L_map)")
+                enhance_button_mask = gr.Button("Enhance Audio (offline CleanMel_L_mask)")
+                clear_btn = gr.Button(
+                    "🗑️ Clear All",
+                    variant="secondary",
+                    size="lg"
+                )
+        
+        output_audio = gr.Audio(label="Enhanced Audio", type="filepath")
+        output_mel = gr.Image(label="Output LogMel Spectrogram", type="filepath", visible=True)
+        output_np = gr.File(label="Enhanced LogMel Spec. (.npy)", type="filepath")
+        
+        enhance_button_map.click(
+            enhance_cleanmel_L_map, 
+            inputs=audio_input, 
+            outputs=[output_audio, output_mel, output_np]
+        )
+        
+        enhance_button_mask.click(
+            enhance_cleanmel_L_mask, 
+            inputs=audio_input, 
+            outputs=[output_audio, output_mel, output_np]
+        )
+        clear_btn.click(
+                fn=reset_everything,
+                outputs=[output_audio, output_mel, output_np]
+        )
+
+    demo.launch(debug=False, share=True)

ckpts/CleanMel/offline_CleanMel_L_map.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:dd0a5118f0c57c91c521564f8275f3a731e10a5afdba9859a3b997067e1eefdd
+size 29065251

ckpts/CleanMel/offline_CleanMel_L_mask.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:42ecba65266fe11d1beb97138c1e83146a57da06daf075a367e902e151f73afa
+size 29065251

ckpts/Vocos/vocos_offline.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f55c3f8cd69ab92e97578d5a99f7207ece6e2c08ccde801a8402e9fa77b72d14
+size 223334266

configs/cleanmel_offline.yaml

@@ -0,0 +1,67 @@
+seed_everything: 2
+
+trainer:
+  gradient_clip_val: 10
+  gradient_clip_algorithm: norm
+  devices: null
+  accelerator: gpu
+  strategy: ddp_find_unused_parameters_false
+  sync_batchnorm: false
+  precision: 32
+  num_sanity_val_steps: 3
+  deterministic: true
+  max_epochs: 100
+  log_every_n_steps: 40
+
+model:
+  arch:
+    class_path: model.arch.cleanmel.CleanMel
+    init_args:
+      dim_input: 2
+      dim_output: 1
+      n_layers: 16
+      dim_hidden: 144
+      layer_linear_freq: 1
+      f_kernel_size: 5
+      f_conv_groups: 8
+      n_freqs: 257
+      n_mels: 80
+      mamba_state: 16
+      mamba_conv_kernel: 4
+      online: false
+      sr: 16000
+      n_fft: 512
+  input_stft:
+    class_path: model.io.stft.InputSTFT
+    init_args:
+      n_fft: 512 
+      n_win: 512 
+      n_hop: 128
+      center: true 
+      normalize: false 
+      onesided: true 
+      online: false
+  target_stft:
+    class_path: model.io.stft.TargetMel
+    init_args:
+        sample_rate: 16000
+        n_fft: 512
+        n_win: 512
+        n_hop: 128
+        n_mels: 80
+        f_min: 0
+        f_max: 8000
+        power: 2
+        center: true
+        normalize: false
+        onesided: true
+        mel_norm: "slaney"
+        mel_scale: "slaney"
+        librosa_mel: true
+        online: false
+
+  optimizer: [AdamW, { lr: 0.001, weight_decay: 0.001}]
+  lr_scheduler: [ExponentialLR, { gamma: 0.99 }]
+  exp_name: exp
+  metrics: [DNSMOS]
+  log_eps: 1e-5

configs/vocos_offline.yaml

@@ -0,0 +1,44 @@
+feature_extractor:
+  class_path: model.vocos.feature_extractors.MelSpectrogramFeatures
+  init_args:
+    sample_rate: 16000
+    n_fft: 512
+    n_win: 512
+    n_hop: 128
+    n_mels: 80
+    f_min: 0
+    f_max: 8000
+    power: 2
+    center: true
+    normalize: false
+    onesided: true
+    mel_norm: slaney
+    mel_scale: slaney
+    librosa_mel: true
+    clip_val: 0.00001
+backbone:
+  class_path: model.vocos.models.VocosBackbone
+  init_args:
+    input_channels: 80
+    dim: 512
+    intermediate_dim: 1536
+    num_layers: 8
+    layer_scale_init_value: null
+    adanorm_num_embeddings: null
+head:
+  class_path: model.vocos.heads.ISTFTHead
+  init_args:
+    dim: 512
+    n_fft: 512
+    hop_length: 128
+    padding: center
+sample_rate: 16000
+initial_learning_rate: 0.0005
+num_warmup_steps: 0
+mel_loss_coeff: 45.0
+mrd_loss_coeff: 0.1
+pretrain_mel_steps: 0
+decay_mel_coeff: false
+evaluate_utmos: true
+evaluate_pesq: true
+evaluate_periodicty: true

model/cleanmel.py

@@ -0,0 +1,401 @@
+from typing import *
+
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import pytorch_lightning
+import librosa
+
+from torch import Tensor
+from torch.nn import Parameter, init
+from torch.nn.common_types import _size_1_t
+
+from mamba_ssm import Mamba
+from mamba_ssm.utils.generation import InferenceParams
+
+class LinearGroup(nn.Module):
+
+    def __init__(self, in_features: int, out_features: int, num_groups: int, bias: bool = True) -> None:
+        super(LinearGroup, self).__init__()
+        self.in_features = in_features
+        self.out_features = out_features
+        self.num_groups = num_groups
+        self.weight = Parameter(torch.empty((num_groups, out_features, in_features)))
+        if bias:
+            self.bias = Parameter(torch.empty(num_groups, out_features))
+        else:
+            self.register_parameter('bias', None)
+        self.reset_parameters()
+
+    def reset_parameters(self) -> None:
+        # same as linear
+        init.kaiming_uniform_(self.weight, a=math.sqrt(5))
+        if self.bias is not None:
+            fan_in, _ = init._calculate_fan_in_and_fan_out(self.weight)
+            bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0
+            init.uniform_(self.bias, -bound, bound)
+
+    def forward(self, x: Tensor) -> Tensor:
+        """shape [..., group, feature]"""
+        x = torch.einsum("...gh,gkh->...gk", x, self.weight)
+        if self.bias is not None:
+            x = x + self.bias
+        return x
+
+    def extra_repr(self) -> str:
+        return f"{self.in_features}, {self.out_features}, num_groups={self.num_groups}, bias={True if self.bias is not None else False}"
+
+class LayerNorm(nn.LayerNorm):
+
+    def __init__(self, seq_last: bool, **kwargs) -> None:
+        """
+        Arg s:
+            seq_last (bool): whether the sequence dim is the last dim
+        """
+        super().__init__(**kwargs)
+        self.seq_last = seq_last
+
+    def forward(self, input: Tensor) -> Tensor:
+        if self.seq_last:
+            input = input.transpose(-1, 1)  # [B, H, Seq] -> [B, Seq, H], or [B,H,w,h] -> [B,h,w,H]
+        o = super().forward(input)
+        if self.seq_last:
+            o = o.transpose(-1, 1)
+        return o
+
+class CausalConv1d(nn.Conv1d):
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: _size_1_t,
+        stride: _size_1_t = 1,
+        padding: _size_1_t | str = 0,
+        dilation: _size_1_t = 1,
+        groups: int = 1,
+        bias: bool = True,
+        padding_mode: str = 'zeros',
+        device=None,
+        dtype=None,
+        look_ahead: int = 0,
+    ) -> None:
+        super().__init__(in_channels, out_channels, kernel_size, stride, padding, dilation, groups, bias, padding_mode, device, dtype)
+        self.look_ahead = look_ahead
+        assert look_ahead <= self.kernel_size[0] - 1, (look_ahead, self.kernel_size)
+
+    def forward(self, x: Tensor, state: Dict[int, Any] = None) -> Tensor:
+        # x [B,H,T]
+        B, H, T = x.shape
+        if state is None or id(self) not in state:
+            x = F.pad(x, pad=(self.kernel_size[0] - 1 - self.look_ahead, self.look_ahead))
+        else:
+            x = torch.concat([state[id(self)], x], dim=-1)
+        if state is not None:
+            state[id(self)] = x[..., -self.kernel_size + 1:]
+        x = super().forward(x)
+        return x
+
+class CleanMelLayer(nn.Module):
+
+    def __init__(
+            self,
+            dim_hidden: int,
+            dim_squeeze: int,
+            n_freqs: int,
+            dropout: Tuple[float, float, float] = (0, 0, 0),
+            f_kernel_size: int = 5,
+            f_conv_groups: int = 8,
+            padding: str = 'zeros',
+            full: nn.Module = None,
+            mamba_state: int = None,
+            mamba_conv_kernel: int = None,
+            online: bool = False,
+    ) -> None:
+        super().__init__()
+        self.online = online
+        # cross-band block
+        # frequency-convolutional module
+        self.fconv1 = nn.ModuleList([
+            LayerNorm(seq_last=True, normalized_shape=dim_hidden),
+            nn.Conv1d(in_channels=dim_hidden, out_channels=dim_hidden, kernel_size=f_kernel_size, groups=f_conv_groups, padding='same', padding_mode=padding),
+            nn.PReLU(dim_hidden),
+        ])
+        # full-band linear module
+        self.norm_full = LayerNorm(seq_last=False, normalized_shape=dim_hidden)
+        self.full_share = False if full == None else True
+        self.squeeze = nn.Sequential(nn.Conv1d(in_channels=dim_hidden, out_channels=dim_squeeze, kernel_size=1), nn.SiLU())
+        self.dropout_full = nn.Dropout2d(dropout[2]) if dropout[2] > 0 else None
+        self.full = LinearGroup(n_freqs, n_freqs, num_groups=dim_squeeze) if full == None else full
+        self.unsqueeze = nn.Sequential(nn.Conv1d(in_channels=dim_squeeze, out_channels=dim_hidden, kernel_size=1), nn.SiLU())
+        # frequency-convolutional module
+        self.fconv2 = nn.ModuleList([
+            LayerNorm(seq_last=True, normalized_shape=dim_hidden),
+            nn.Conv1d(in_channels=dim_hidden, out_channels=dim_hidden, kernel_size=f_kernel_size, groups=f_conv_groups, padding='same', padding_mode=padding),
+            nn.PReLU(dim_hidden),
+        ])
+
+        # narrow-band block
+        self.norm_mamba = LayerNorm(seq_last=False, normalized_shape=dim_hidden)
+        if online:
+            self.mamba = Mamba(d_model=dim_hidden, d_state=mamba_state, d_conv=mamba_conv_kernel, layer_idx=0)
+        else:
+            self.mamba = nn.ModuleList([
+                Mamba(d_model=dim_hidden, d_state=mamba_state, d_conv=mamba_conv_kernel, layer_idx=0),
+                Mamba(d_model=dim_hidden, d_state=mamba_state, d_conv=mamba_conv_kernel, layer_idx=1),
+            ])
+        
+        self.dropout_mamba = nn.Dropout(dropout[0])
+
+    def forward(self, x: Tensor, inference: bool = False) -> Tensor:
+        x = x + self._fconv(self.fconv1, x)
+        x = x + self._full(x)
+        x = x + self._fconv(self.fconv2, x)        
+        if self.online:
+            x = x + self._mamba(x, self.mamba, self.norm_mamba, self.dropout_mamba, inference)
+        else:
+            x_fw = x + self._mamba(x, self.mamba[0], self.norm_mamba, self.dropout_mamba, inference)
+            x_bw = x.flip(dims=[2]) + self._mamba(x.flip(dims=[2]), self.mamba[1], self.norm_mamba, self.dropout_mamba, inference)
+            x = (x_fw + x_bw.flip(dims=[2])) / 2 
+        return x
+
+    def _mamba(self, x: Tensor, mamba: Mamba, norm: nn.Module, dropout: nn.Module, inference: bool = False):
+        B, F, T, H = x.shape
+        x = norm(x)
+        x = x.reshape(B * F, T, H)
+        if inference:
+            inference_params = InferenceParams(T, B * F)
+            xs = []
+            for i in range(T):
+                inference_params.seqlen_offset = i
+                xi = mamba.forward(x[:, [i], :], inference_params)
+                xs.append(xi)
+            x = torch.concat(xs, dim=1)
+        else:
+            x = mamba.forward(x)
+        x = x.reshape(B, F, T, H)
+        return dropout(x)
+
+    def _fconv(self, ml: nn.ModuleList, x: Tensor) -> Tensor:
+        B, F, T, H = x.shape
+        x = x.permute(0, 2, 3, 1)  # [B,T,H,F]
+        x = x.reshape(B * T, H, F)
+        for m in ml:
+            x = m(x)
+        x = x.reshape(B, T, H, F)
+        x = x.permute(0, 3, 1, 2)  # [B,F,T,H]
+        return x
+
+    def _full(self, x: Tensor) -> Tensor:
+        B, F, T, H = x.shape
+        x = self.norm_full(x)
+        x = x.permute(0, 2, 3, 1)  # [B,T,H,F]
+        x = x.reshape(B * T, H, F)
+        x = self.squeeze(x)  # [B*T,H',F]
+        if self.dropout_full:
+            x = x.reshape(B, T, -1, F)
+            x = x.transpose(1, 3)  # [B,F,H',T]
+            x = self.dropout_full(x)  # dropout some frequencies in one utterance
+            x = x.transpose(1, 3)  # [B,T,H',F]
+            x = x.reshape(B * T, -1, F)
+        x = self.full(x)  # [B*T,H',F]
+        x = self.unsqueeze(x)  # [B*T,H,F]
+        x = x.reshape(B, T, H, F)
+        x = x.permute(0, 3, 1, 2)  # [B,F,T,H]
+        return x
+
+    def extra_repr(self) -> str:
+        return f"full_share={self.full_share}"
+
+
+class CleanMel(nn.Module):
+
+    def __init__(
+        self,
+        dim_input: int,  # the input dim for each time-frequency point
+        dim_output: int,  # the output dim for each time-frequency point
+        n_layers: int,
+        n_freqs: int,
+        n_mels: int = 80,
+        layer_linear_freq: int = 1,
+        encoder_kernel_size: int = 5,
+        dim_hidden: int = 192,
+        dropout: Tuple[float, float, float] = (0, 0, 0),
+        f_kernel_size: int = 5,
+        f_conv_groups: int = 8,
+        padding: str = 'zeros',
+        mamba_state: int = 16,
+        mamba_conv_kernel: int = 4,
+        online: bool = True,
+        sr: int = 16000,
+        n_fft: int = 512,
+    ):
+        super().__init__()
+        self.layer_linear_freq = layer_linear_freq
+        self.online = online
+        # encoder
+        self.encoder = CausalConv1d(in_channels=dim_input, out_channels=dim_hidden, kernel_size=encoder_kernel_size, look_ahead=0)
+        # cleanmel layers
+        full = None
+        layers = []
+        for l in range(n_layers):
+            layer = CleanMelLayer(
+                dim_hidden=dim_hidden,
+                dim_squeeze=8 if l < layer_linear_freq else dim_hidden,
+                n_freqs=n_freqs if l < layer_linear_freq else n_mels,
+                dropout=dropout,
+                f_kernel_size=f_kernel_size,
+                f_conv_groups=f_conv_groups,
+                padding=padding,
+                full=full if l > layer_linear_freq else None,
+                online=online,
+                mamba_conv_kernel=mamba_conv_kernel,
+                mamba_state=mamba_state,    
+            )
+            if hasattr(layer, 'full'):
+                full = layer.full
+            layers.append(layer)
+        self.layers = nn.ModuleList(layers)
+        # Mel filterbank
+        linear2mel = librosa.filters.mel(**{"sr": sr, "n_fft": n_fft, "n_mels": n_mels})
+        self.register_buffer("linear2mel", torch.nn.Parameter(torch.tensor(linear2mel.T, dtype=torch.float32)))
+        # decoder
+        self.decoder = nn.Linear(in_features=dim_hidden, out_features=dim_output)
+
+    def forward(self, x: Tensor, inference: bool = False) -> Tensor:
+        # x: [Batch, Freq, Time, Feature]
+        B, F, T, H0 = x.shape
+        x = self.encoder(x.reshape(B * F, T, H0).permute(0, 2, 1)).permute(0, 2, 1)
+        
+        H = x.shape[2]
+        x = x.reshape(B, F, T, H)
+        # First Cross-Narrow band block in Linear Frequency
+        for i in range(self.layer_linear_freq):
+            m = self.layers[i]
+            x = m(x, inference).contiguous()
+        
+        # Mel-filterbank
+        x = torch.einsum("bfth,fm->bmth", x, self.linear2mel)
+
+        for i in range(self.layer_linear_freq, len(self.layers)):
+            m = self.layers[i]
+            x = m(x, inference).contiguous()
+        
+        y = self.decoder(x).squeeze(-1)
+        return y.contiguous()
+
+if __name__ == '__main__':
+    # a quick demo here for the CleanMel model
+    # input: wavs
+    # output: enhanced log-mel spectrogram
+    pytorch_lightning.seed_everything(1234)
+    import soundfile as sf
+    import matplotlib.pyplot as plt
+    import numpy as np
+    from model.io.stft import InputSTFT
+    from model.io.stft import TargetMel
+    from torch.utils.flop_counter import FlopCounterMode
+    
+    online=False
+    # Define input STFT and target Mel
+    stft = InputSTFT(
+        n_fft=512, 
+        n_win=512, 
+        n_hop=128,
+        center=True, 
+        normalize=False, 
+        onesided=True, 
+        online=online).to("cuda")
+    
+    target_mel = TargetMel(
+        sample_rate=16000,
+        n_fft=512,
+        n_win=512,
+        n_hop=128,
+        n_mels=80,
+        f_min=0,
+        f_max=8000,
+        power=2,
+        center=True,
+        normalize=False,
+        onesided=True,
+        mel_norm="slaney",
+        mel_scale="slaney",
+        librosa_mel=True,
+        online=online).to("cuda")
+
+    def customize_soxnorm(wav, gain=-3, factor=None):
+        wav = np.clip(wav, a_max=1, a_min=-1)
+        if factor is None:
+            linear_gain = 10 ** (gain / 20)
+            factor = linear_gain / np.abs(wav).max()
+            wav = wav * factor
+            return wav, factor
+        else:
+            wav = wav * factor
+            return wav, None
+
+    # Noisy file path
+    wav = "./src/demos/noisy_CHIME-real_F05_442C020S_STR_REAL.wav"
+    wavname = wav.split("/")[-1].split(".")[0]
+    
+    print(f"Processing {wav}")
+    noisy, fs = sf.read(wav)
+    dur = len(noisy) / fs
+    noisy, factor = customize_soxnorm(noisy, gain=-3)
+    noisy = torch.tensor(noisy).unsqueeze(0).float().to("cuda")
+    # vocos norm
+    x = stft(noisy)
+    # Load the model
+    hidden=96
+    depth=8
+    model = CleanMel(
+        dim_input=2,
+        dim_output=1,
+        n_layers=depth,
+        dim_hidden=hidden,
+        layer_linear_freq=1,
+        f_kernel_size=5,
+        f_conv_groups=8,
+        n_freqs=257,
+        mamba_state=16,
+        mamba_conv_kernel=4,
+        online=online,
+        sr=16000,
+        n_fft=512
+    ).to("cuda")
+
+    # Load the pretrained model
+    state_dict = torch.load("./pretrained/CleanMel_S_L1.ckpt")
+    model.load_state_dict(state_dict)
+    
+    model.eval()
+    with FlopCounterMode(model, display=False) as fcm:
+        y_hat = model(x, inference=False)
+        flops_forward_eval = fcm.get_total_flops()
+    params_eval = sum(param.numel() for param in model.parameters())
+    print(f"flops_forward={flops_forward_eval/1e9 / dur:.2f}G")
+    print(f"params={params_eval/1e6:.2f} M")
+
+    # y_hat is the enhanced log-mel spectrogram
+    y_hat = y_hat[0].cpu().detach().numpy()
+    
+    # sanity check
+    if wavname == "noisy_CHIME-real_F05_442C020S_STR_REAL":
+        assert np.allclose(y_hat, np.load("./src/inference/check_CHIME-real_F05_442C020S_STR_REAL.npy"), atol=1e-5)
+    
+    # plot the enhanced mel spectrogram
+    noisy_mel = target_mel(noisy)
+    noisy_mel = torch.log(noisy_mel.clamp(min=1e-5))[0].cpu().detach().numpy()    
+    vmax = math.log(1e2)
+    vmin = math.log(1e-5)
+    plt.figure(figsize=(8, 4))
+    plt.subplot(2, 1, 1)
+    plt.imshow(noisy_mel, aspect='auto', origin='lower', cmap='jet', vmax=vmax, vmin=vmin)
+    plt.colorbar()
+    plt.subplot(2, 1, 2)
+    plt.imshow(y_hat, aspect='auto', origin='lower', cmap='jet', vmax=vmax, vmin=vmin)
+    plt.colorbar()
+    plt.tight_layout()
+    plt.savefig(f"./src/inference/{wavname}.png")

model/stft.py

@@ -0,0 +1,154 @@
+import torch
+import librosa
+import torch.nn as nn
+import random
+from torch import Tensor
+from typing import Optional
+from torchaudio.transforms import Spectrogram
+from torchaudio.transforms import Spectrogram, MelScale
+
+def soxnorm(wav: torch.Tensor, gain, factor=None):
+    """sox norm, used in Vocos codes;
+    """
+    wav = torch.clip(wav, max=1, min=-1).float()
+    if factor is None:
+        linear_gain = 10 ** (gain / 20)
+        factor = linear_gain / torch.abs(wav).max().item()
+        wav = wav * factor
+    else:
+        # for clean speech, normed by the noisy factor
+        wav = wav * factor
+    assert torch.all(wav.abs() <= 1), f"out wavform is not in [-1, 1], {wav.abs().max()}"
+    return wav, factor
+
+
+class InputSTFT(nn.Module):
+    """
+    The STFT of the input signal of CleanMel (STFT coefficients);
+    In online mode, the recursive normalization is used.
+    """
+    def __init__(
+        self, 
+        n_fft: int,
+        n_win: int, 
+        n_hop: int, 
+        center: bool,
+        normalize: bool,
+        onesided: bool,
+        online: bool = False):
+        super().__init__()
+        
+        self.online = online
+        self.stft=Spectrogram(
+            n_fft=n_fft,
+            win_length=n_win,
+            hop_length=n_hop,
+            normalized=normalize,
+            center=center,
+            onesided=onesided,
+            power=None
+        )
+    
+    def forward(self, x):
+        if self.online:
+            # recursive normalization
+            x = self.stft(x)
+            x_mag = x.abs()
+            x_norm = recursive_normalization(x_mag)
+            x = x / x_norm.clamp(min=1e-8)
+            x = torch.view_as_real(x)
+        else:
+            # vocos dBFS normalization
+            x, x_norm = soxnorm(x, random.randint(-6, -1) if self.training else -3)
+            x = self.stft(x)
+            x = torch.view_as_real(x)
+        return x, x_norm
+
+
+class LibrosaMelScale(nn.Module):
+    r"""Pytorch implementation of librosa mel scale to align with common ESPNet ASR models; 
+    You might need to define .
+    """
+    def __init__(self, n_mels, sample_rate, f_min, f_max, n_stft, norm=None, mel_scale="slaney"):
+        super(LibrosaMelScale, self).__init__()
+        
+        _mel_options = dict(
+            sr=sample_rate,
+            n_fft=(n_stft - 1) * 2,
+            n_mels=n_mels,
+            fmin=f_min,
+            fmax=f_max if f_max is not None else float(sample_rate // 2),
+            htk=mel_scale=="htk",
+            norm=norm
+        )
+        
+        fb = torch.from_numpy(librosa.filters.mel(**_mel_options).T).float()
+        self.register_buffer("fb", fb)
+    
+    def forward(self, specgram):
+        mel_specgram = torch.matmul(specgram.transpose(-1, -2), self.fb).transpose(-1, -2)
+        return mel_specgram
+
+
+class TargetMel(nn.Module):
+    """
+    This class generates the enhancement TARGET mel spectrogram;
+    """
+    def __init__(
+        self,
+        sample_rate: int,
+        n_fft: int,
+        n_win: int,
+        n_hop: int,
+        n_mels: int,
+        f_min: int,
+        f_max: int,
+        power: int,
+        center: bool,
+        normalize: bool,
+        onesided: bool,
+        mel_norm: str | None,
+        mel_scale: str,
+        librosa_mel: bool = True,
+        online: bool = False,
+        ):
+        super().__init__()
+        # This implementation vs torchaudio.transforms.MelSpectrogram: Add librosa melscale
+        # librosa melscale is numerically different from the torchaudio melscale (x_diff > 1e-5)
+        
+        self.sample_rate = sample_rate
+        self.n_fft = n_fft
+        self.online = online
+        self.stft = Spectrogram(
+            n_fft=n_fft,
+            win_length=n_win,
+            hop_length=n_hop,
+            power=None if online else power,
+            normalized=normalize,
+            center=center,
+            onesided=onesided,
+        )
+        mel_method = LibrosaMelScale if librosa_mel else MelScale
+        self.mel_scale = mel_method(
+            n_mels=n_mels,
+            sample_rate=sample_rate,
+            f_min=f_min,
+            f_max=f_max,
+            n_stft=n_fft // 2 + 1,
+            norm=mel_norm,
+            mel_scale=mel_scale,
+        )
+        
+    def forward(self, x: Tensor, x_norm=None):      
+        if self.online:
+            # apply recursive normalization to target waveform
+            spectrogram = self.stft(x)
+            spectrogram = spectrogram / (x_norm + 1e-8)
+            spectrogram = spectrogram.abs().pow(2)  # to power spectrogram
+        else:
+            # apply vocos dBFS normalization to target waveform
+            x, _ = soxnorm(x, None, x_norm)
+            spectrogram = self.stft(x)
+        # mel spectrogram
+        mel_specgram = self.mel_scale(spectrogram)
+        return mel_specgram

model/vocos/__init__.py

@@ -0,0 +1 @@
+

model/vocos/dataset.py

@@ -0,0 +1,93 @@
+from dataclasses import dataclass
+
+import numpy as np
+from pytorch_lightning.utilities.types import EVAL_DATALOADERS
+import torch
+import torchaudio
+import warnings
+from pytorch_lightning import LightningDataModule
+from torch.utils.data import Dataset, DataLoader
+
+torch.set_num_threads(1)
+
+
+@dataclass
+class DataConfig:
+    filelist_path: str
+    sampling_rate: int
+    num_samples: int
+    batch_size: int
+    num_workers: int
+
+
+class VocosDataModule(LightningDataModule):
+    def __init__(self, train_params: DataConfig, val_params: DataConfig):
+        super().__init__()
+        self.train_config = train_params
+        self.val_config = val_params
+
+    def _get_dataloder(self, cfg: DataConfig, train: bool):
+        dataset = VocosDataset(cfg, train=train)
+        dataloader = DataLoader(
+            dataset, batch_size=cfg.batch_size, num_workers=cfg.num_workers, shuffle=train, pin_memory=True,
+        )
+        return dataloader
+
+    def train_dataloader(self) -> DataLoader:
+        return self._get_dataloder(self.train_config, train=True)
+
+    def val_dataloader(self) -> DataLoader:
+        return self._get_dataloder(self.val_config, train=False)
+
+    def test_dataloader(self) -> DataLoader:
+        return self.val_dataloader()
+
+class VocosDataset(Dataset):
+    def __init__(self, cfg: DataConfig, train: bool):
+        with open(cfg.filelist_path) as f:
+            self.filelist = f.read().splitlines()
+        self.sampling_rate = cfg.sampling_rate
+        self.num_samples = cfg.num_samples
+        self.train = train
+
+    def __len__(self) -> int:
+        return len(self.filelist)
+
+    def customize_soxnorm(self, wav, gain=-3, factor=None):
+            wav = np.clip(wav, a_max=1, a_min=-1)
+            if factor is None:
+                linear_gain = 10 ** (gain / 20)
+                wav = wav / np.abs(wav).max() * linear_gain
+                return wav,  linear_gain / np.abs(wav).max()
+            else:
+                wav = wav * factor
+                return wav, None
+
+    def __getitem__(self, index: int) -> torch.Tensor:
+        audio_path = self.filelist[index]
+        try:
+            y, sr = torchaudio.load(audio_path)
+        except:
+            warnings.warn(f"Error loading {audio_path}")
+            return self.__getitem__(np.random.randint(len(self.filelist)))   
+        if y.size(-1) == 0:
+            return self.__getitem__(np.random.randint(len(self.filelist)))
+        if y.size(0) > 1:
+            # mix to mono
+            y = y.mean(dim=0, keepdim=True)
+        gain = np.random.uniform(-1, -6) if self.train else -3
+        y, _ = self.customize_soxnorm(y, gain)
+        if sr != self.sampling_rate:
+            y = torchaudio.functional.resample(y, orig_freq=sr, new_freq=self.sampling_rate)
+        if y.size(-1) < self.num_samples:
+            pad_length = self.num_samples - y.size(-1)
+            padding_tensor = y.repeat(1, 1 + pad_length // y.size(-1))
+            y = torch.cat((y, padding_tensor[:, :pad_length]), dim=1)
+        elif self.train:
+            start = np.random.randint(low=0, high=y.size(-1) - self.num_samples + 1)
+            y = y[:, start : start + self.num_samples]
+        else:
+            # During validation, take always the first segment for determinism
+            y = y[:, : self.num_samples]
+
+        return y[0]

model/vocos/discriminators.py

@@ -0,0 +1,211 @@
+from typing import List, Optional, Tuple
+
+import torch
+from einops import rearrange
+from torch import nn
+from torch.nn import Conv2d
+from torch.nn.utils import weight_norm
+from torchaudio.transforms import Spectrogram
+
+
+class MultiPeriodDiscriminator(nn.Module):
+    """
+    Multi-Period Discriminator module adapted from https://github.com/jik876/hifi-gan.
+    Additionally, it allows incorporating conditional information with a learned embeddings table.
+
+    Args:
+        periods (tuple[int]): Tuple of periods for each discriminator.
+        num_embeddings (int, optional): Number of embeddings. None means non-conditional discriminator.
+            Defaults to None.
+    """
+
+    def __init__(self, periods: Tuple[int, ...] = (2, 3, 5, 7, 11), num_embeddings: Optional[int] = None):
+        super().__init__()
+        self.discriminators = nn.ModuleList([DiscriminatorP(period=p, num_embeddings=num_embeddings) for p in periods])
+
+    def forward(
+        self, y: torch.Tensor, y_hat: torch.Tensor, bandwidth_id: Optional[torch.Tensor] = None
+    ) -> Tuple[List[torch.Tensor], List[torch.Tensor], List[List[torch.Tensor]], List[List[torch.Tensor]]]:
+        y_d_rs = []
+        y_d_gs = []
+        fmap_rs = []
+        fmap_gs = []
+        for d in self.discriminators:
+            y_d_r, fmap_r = d(x=y, cond_embedding_id=bandwidth_id)
+            y_d_g, fmap_g = d(x=y_hat, cond_embedding_id=bandwidth_id)
+            y_d_rs.append(y_d_r)
+            fmap_rs.append(fmap_r)
+            y_d_gs.append(y_d_g)
+            fmap_gs.append(fmap_g)
+
+        return y_d_rs, y_d_gs, fmap_rs, fmap_gs
+
+
+class DiscriminatorP(nn.Module):
+    def __init__(
+        self,
+        period: int,
+        in_channels: int = 1,
+        kernel_size: int = 5,
+        stride: int = 3,
+        lrelu_slope: float = 0.1,
+        num_embeddings: Optional[int] = None,
+    ):
+        super().__init__()
+        self.period = period
+        self.convs = nn.ModuleList(
+            [
+                weight_norm(Conv2d(in_channels, 32, (kernel_size, 1), (stride, 1), padding=(kernel_size // 2, 0))),
+                weight_norm(Conv2d(32, 128, (kernel_size, 1), (stride, 1), padding=(kernel_size // 2, 0))),
+                weight_norm(Conv2d(128, 512, (kernel_size, 1), (stride, 1), padding=(kernel_size // 2, 0))),
+                weight_norm(Conv2d(512, 1024, (kernel_size, 1), (stride, 1), padding=(kernel_size // 2, 0))),
+                weight_norm(Conv2d(1024, 1024, (kernel_size, 1), (1, 1), padding=(kernel_size // 2, 0))),
+            ]
+        )
+        if num_embeddings is not None:
+            self.emb = torch.nn.Embedding(num_embeddings=num_embeddings, embedding_dim=1024)
+            torch.nn.init.zeros_(self.emb.weight)
+
+        self.conv_post = weight_norm(Conv2d(1024, 1, (3, 1), 1, padding=(1, 0)))
+        self.lrelu_slope = lrelu_slope
+
+    def forward(
+        self, x: torch.Tensor, cond_embedding_id: Optional[torch.Tensor] = None
+    ) -> Tuple[torch.Tensor, List[torch.Tensor]]:
+        x = x.unsqueeze(1)
+        fmap = []
+        # 1d to 2d
+        b, c, t = x.shape
+        if t % self.period != 0:  # pad first
+            n_pad = self.period - (t % self.period)
+            x = torch.nn.functional.pad(x, (0, n_pad), "reflect")
+            t = t + n_pad
+        x = x.view(b, c, t // self.period, self.period)
+
+        for i, l in enumerate(self.convs):
+            x = l(x)
+            x = torch.nn.functional.leaky_relu(x, self.lrelu_slope)
+            if i > 0:
+                fmap.append(x)
+        if cond_embedding_id is not None:
+            emb = self.emb(cond_embedding_id)
+            h = (emb.view(1, -1, 1, 1) * x).sum(dim=1, keepdims=True)
+        else:
+            h = 0
+        x = self.conv_post(x)
+        fmap.append(x)
+        x += h
+        x = torch.flatten(x, 1, -1)
+
+        return x, fmap
+
+
+class MultiResolutionDiscriminator(nn.Module):
+    def __init__(
+        self,
+        fft_sizes: Tuple[int, ...] = (2048, 1024, 512),
+        num_embeddings: Optional[int] = None,
+    ):
+        """
+        Multi-Resolution Discriminator module adapted from https://github.com/descriptinc/descript-audio-codec.
+        Additionally, it allows incorporating conditional information with a learned embeddings table.
+
+        Args:
+            fft_sizes (tuple[int]): Tuple of window lengths for FFT. Defaults to (2048, 1024, 512).
+            num_embeddings (int, optional): Number of embeddings. None means non-conditional discriminator.
+                Defaults to None.
+        """
+
+        super().__init__()
+        self.discriminators = nn.ModuleList(
+            [DiscriminatorR(window_length=w, num_embeddings=num_embeddings) for w in fft_sizes]
+        )
+
+    def forward(
+        self, y: torch.Tensor, y_hat: torch.Tensor, bandwidth_id: torch.Tensor = None
+    ) -> Tuple[List[torch.Tensor], List[torch.Tensor], List[List[torch.Tensor]], List[List[torch.Tensor]]]:
+        y_d_rs = []
+        y_d_gs = []
+        fmap_rs = []
+        fmap_gs = []
+
+        for d in self.discriminators:
+            y_d_r, fmap_r = d(x=y, cond_embedding_id=bandwidth_id)
+            y_d_g, fmap_g = d(x=y_hat, cond_embedding_id=bandwidth_id)
+            y_d_rs.append(y_d_r)
+            fmap_rs.append(fmap_r)
+            y_d_gs.append(y_d_g)
+            fmap_gs.append(fmap_g)
+
+        return y_d_rs, y_d_gs, fmap_rs, fmap_gs
+
+
+class DiscriminatorR(nn.Module):
+    def __init__(
+        self,
+        window_length: int,
+        num_embeddings: Optional[int] = None,
+        channels: int = 32,
+        hop_factor: float = 0.25,
+        bands: Tuple[Tuple[float, float], ...] = ((0.0, 0.1), (0.1, 0.25), (0.25, 0.5), (0.5, 0.75), (0.75, 1.0)),
+    ):
+        super().__init__()
+        self.window_length = window_length
+        self.hop_factor = hop_factor
+        self.spec_fn = Spectrogram(
+            n_fft=window_length, hop_length=int(window_length * hop_factor), win_length=window_length, power=None
+        )
+        n_fft = window_length // 2 + 1
+        bands = [(int(b[0] * n_fft), int(b[1] * n_fft)) for b in bands]
+        self.bands = bands
+        convs = lambda: nn.ModuleList(
+            [
+                weight_norm(nn.Conv2d(2, channels, (3, 9), (1, 1), padding=(1, 4))),
+                weight_norm(nn.Conv2d(channels, channels, (3, 9), (1, 2), padding=(1, 4))),
+                weight_norm(nn.Conv2d(channels, channels, (3, 9), (1, 2), padding=(1, 4))),
+                weight_norm(nn.Conv2d(channels, channels, (3, 9), (1, 2), padding=(1, 4))),
+                weight_norm(nn.Conv2d(channels, channels, (3, 3), (1, 1), padding=(1, 1))),
+            ]
+        )
+        self.band_convs = nn.ModuleList([convs() for _ in range(len(self.bands))])
+
+        if num_embeddings is not None:
+            self.emb = torch.nn.Embedding(num_embeddings=num_embeddings, embedding_dim=channels)
+            torch.nn.init.zeros_(self.emb.weight)
+
+        self.conv_post = weight_norm(nn.Conv2d(channels, 1, (3, 3), (1, 1), padding=(1, 1)))
+
+    def spectrogram(self, x):
+        # Remove DC offset
+        x = x - x.mean(dim=-1, keepdims=True)
+        # Peak normalize the volume of input audio
+        x = 0.8 * x / (x.abs().max(dim=-1, keepdim=True)[0] + 1e-9)
+        x = self.spec_fn(x)
+        x = torch.view_as_real(x)
+        x = rearrange(x, "b f t c -> b c t f")
+        # Split into bands
+        x_bands = [x[..., b[0] : b[1]] for b in self.bands]
+        return x_bands
+
+    def forward(self, x: torch.Tensor, cond_embedding_id: torch.Tensor = None):
+        x_bands = self.spectrogram(x)
+        fmap = []
+        x = []
+        for band, stack in zip(x_bands, self.band_convs):
+            for i, layer in enumerate(stack):
+                band = layer(band)
+                band = torch.nn.functional.leaky_relu(band, 0.1)
+                if i > 0:
+                    fmap.append(band)
+            x.append(band)
+        x = torch.cat(x, dim=-1)
+        if cond_embedding_id is not None:
+            emb = self.emb(cond_embedding_id)
+            h = (emb.view(1, -1, 1, 1) * x).sum(dim=1, keepdims=True)
+        else:
+            h = 0
+        x = self.conv_post(x)
+        fmap.append(x)
+        x += h
+
+        return x, fmap

model/vocos/experiment.py

@@ -0,0 +1,398 @@
+import math
+
+import numpy as np
+import pytorch_lightning as pl
+import torch
+import torchaudio
+import transformers
+import torch.nn as nn
+
+from model.vocos.offline.discriminators import MultiPeriodDiscriminator, MultiResolutionDiscriminator
+from model.vocos.offline.feature_extractors import FeatureExtractor
+from model.vocos.offline.heads import FourierHead
+from model.vocos.offline.helpers import plot_spectrogram_to_numpy
+from model.vocos.offline.loss import DiscriminatorLoss, GeneratorLoss, FeatureMatchingLoss, MelSpecReconstructionLoss
+# from models.vocos.offline.models import Backbone
+from model.vocos.offline.modules import safe_log
+
+
+class VocosExp(pl.LightningModule):
+    # noinspection PyUnusedLocal
+    def __init__(
+        self,
+        feature_extractor: FeatureExtractor,
+        backbone: nn.Module,
+        head: nn.Module,
+        sample_rate: int,
+        initial_learning_rate: float,
+        num_warmup_steps: int = 0,
+        mel_loss_coeff: float = 45,
+        mrd_loss_coeff: float = 1.0,
+        pretrain_mel_steps: int = 0,
+        decay_mel_coeff: bool = False,
+        evaluate_utmos: bool = False,
+        evaluate_pesq: bool = False,
+        evaluate_periodicty: bool = False,
+    ):
+        """
+        Args:
+            feature_extractor (FeatureExtractor): An instance of FeatureExtractor to extract features from audio signals.
+            backbone (Backbone): An instance of Backbone model.
+            head (FourierHead):  An instance of Fourier head to generate spectral coefficients and reconstruct a waveform.
+            sample_rate (int): Sampling rate of the audio signals.
+            initial_learning_rate (float): Initial learning rate for the optimizer.
+            num_warmup_steps (int): Number of steps for the warmup phase of learning rate scheduler. Default is 0.
+            mel_loss_coeff (float, optional): Coefficient for Mel-spectrogram loss in the loss function. Default is 45.
+            mrd_loss_coeff (float, optional): Coefficient for Multi Resolution Discriminator loss. Default is 1.0.
+            pretrain_mel_steps (int, optional): Number of steps to pre-train the model without the GAN objective. Default is 0.
+            decay_mel_coeff (bool, optional): If True, the Mel-spectrogram loss coefficient is decayed during training. Default is False.
+            evaluate_utmos (bool, optional): If True, UTMOS scores are computed for each validation run.
+            evaluate_pesq (bool, optional): If True, PESQ scores are computed for each validation run.
+            evaluate_periodicty (bool, optional): If True, periodicity scores are computed for each validation run.
+        """
+        super().__init__()
+        self.save_hyperparameters(ignore=["feature_extractor", "backbone", "head"])
+        self.feature_extractor = feature_extractor
+        self.backbone = backbone
+        self.head = head
+        self.sample_rate = sample_rate
+        self.initial_learning_rate = initial_learning_rate
+        self.num_warmup_steps = num_warmup_steps
+        self.mel_loss_coeff = mel_loss_coeff
+        self.mrd_loss_coeff = mrd_loss_coeff
+        self.pretrain_mel_steps = pretrain_mel_steps
+        self.decay_mel_coeff = decay_mel_coeff
+        self.evaluate_utmos = evaluate_utmos
+        self.evaluate_pesq = evaluate_pesq
+        self.evaluate_periodicty = evaluate_periodicty
+
+        self.multiperioddisc = MultiPeriodDiscriminator()
+        self.multiresddisc = MultiResolutionDiscriminator()
+
+        self.disc_loss = DiscriminatorLoss()
+        self.gen_loss = GeneratorLoss()
+        self.feat_matching_loss = FeatureMatchingLoss()
+        self.melspec_loss = MelSpecReconstructionLoss(sample_rate=sample_rate)
+
+        self.train_discriminator = False
+        self.base_mel_coeff = self.mel_loss_coeff = mel_loss_coeff
+        self.temp_cache=None
+        self.temp_grad=None
+
+    def configure_optimizers(self):
+        disc_params = [
+            {"params": self.multiperioddisc.parameters()},
+            {"params": self.multiresddisc.parameters()},
+        ]
+        gen_params = [
+            {"params": self.feature_extractor.parameters()},
+            {"params": self.backbone.parameters()},
+            {"params": self.head.parameters()},
+        ]
+
+        opt_disc = torch.optim.AdamW(disc_params, lr=self.initial_learning_rate, betas=(0.8, 0.9))
+        opt_gen = torch.optim.AdamW(gen_params, lr=self.initial_learning_rate, betas=(0.8, 0.9))
+
+        max_steps = self.trainer.max_steps // 2  # Max steps per optimizer
+        scheduler_disc = transformers.get_cosine_schedule_with_warmup(
+            opt_disc, num_warmup_steps=self.num_warmup_steps, num_training_steps=max_steps,
+        )
+        scheduler_gen = transformers.get_cosine_schedule_with_warmup(
+            opt_gen, num_warmup_steps=self.num_warmup_steps, num_training_steps=max_steps,
+        )
+
+        return (
+            [opt_disc, opt_gen],
+            [{"scheduler": scheduler_disc, "interval": "step"}, {"scheduler": scheduler_gen, "interval": "step"}],
+        )
+
+    def forward(self, audio_input, **kwargs):
+        features = self.feature_extractor(audio_input, **kwargs)
+        x = self.backbone(features, **kwargs)
+        audio_output = self.head(x)
+        return audio_output
+
+    def training_step(self, batch, batch_idx, optimizer_idx, **kwargs):
+        audio_input = batch
+        # train discriminator
+        if optimizer_idx == 0 and self.train_discriminator:             
+            with torch.no_grad():
+                audio_hat = self(audio_input, **kwargs)
+            real_score_mp, gen_score_mp, _, _ = self.multiperioddisc(y=audio_input, y_hat=audio_hat, **kwargs,)
+            real_score_mrd, gen_score_mrd, _, _ = self.multiresddisc(y=audio_input, y_hat=audio_hat, **kwargs,)
+            loss_mp, loss_mp_real, _ = self.disc_loss(
+                disc_real_outputs=real_score_mp, disc_generated_outputs=gen_score_mp
+            )
+            loss_mrd, loss_mrd_real, _ = self.disc_loss(
+                disc_real_outputs=real_score_mrd, disc_generated_outputs=gen_score_mrd
+            )
+            loss_mp /= len(loss_mp_real)
+            loss_mrd /= len(loss_mrd_real)
+            loss = loss_mp + self.mrd_loss_coeff * loss_mrd
+
+            self.log("discriminator/total", loss, prog_bar=True)
+            self.log("discriminator/multi_period_loss", loss_mp)
+            self.log("discriminator/multi_res_loss", loss_mrd)                
+            return loss
+
+        # train generator
+        if optimizer_idx == 1:
+            audio_hat = self(audio_input, **kwargs)
+            if self.train_discriminator:
+                _, gen_score_mp, fmap_rs_mp, fmap_gs_mp = self.multiperioddisc(
+                    y=audio_input, y_hat=audio_hat, **kwargs,
+                )
+                _, gen_score_mrd, fmap_rs_mrd, fmap_gs_mrd = self.multiresddisc(
+                    y=audio_input, y_hat=audio_hat, **kwargs,
+                )
+                loss_gen_mp, list_loss_gen_mp = self.gen_loss(disc_outputs=gen_score_mp)
+                loss_gen_mrd, list_loss_gen_mrd = self.gen_loss(disc_outputs=gen_score_mrd)
+                loss_gen_mp = loss_gen_mp / len(list_loss_gen_mp)
+                loss_gen_mrd = loss_gen_mrd / len(list_loss_gen_mrd)
+                loss_fm_mp = self.feat_matching_loss(fmap_r=fmap_rs_mp, fmap_g=fmap_gs_mp) / len(fmap_rs_mp)
+                loss_fm_mrd = self.feat_matching_loss(fmap_r=fmap_rs_mrd, fmap_g=fmap_gs_mrd) / len(fmap_rs_mrd)
+
+                self.log("generator/multi_period_loss", loss_gen_mp)
+                self.log("generator/multi_res_loss", loss_gen_mrd)
+                self.log("generator/feature_matching_mp", loss_fm_mp)
+                self.log("generator/feature_matching_mrd", loss_fm_mrd)
+            else:
+                loss_gen_mp = loss_gen_mrd = loss_fm_mp = loss_fm_mrd = 0
+
+            mel_loss = self.melspec_loss(audio_hat, audio_input)
+            loss = (
+                loss_gen_mp
+                + self.mrd_loss_coeff * loss_gen_mrd
+                + loss_fm_mp
+                + self.mrd_loss_coeff * loss_fm_mrd
+                + self.mel_loss_coeff * mel_loss
+            )
+
+            self.log("generator/total_loss", loss, prog_bar=True)
+            self.log("mel_loss_coeff", self.mel_loss_coeff)
+            self.log("generator/mel_loss", mel_loss)
+
+            if self.global_step % 1000 == 0 and self.global_rank == 0:
+                self.logger.experiment.add_audio(
+                    "train/audio_in", audio_input[0].data.cpu(), self.global_step, self.sample_rate
+                )
+                self.logger.experiment.add_audio(
+                    "train/audio_pred", audio_hat[0].data.cpu(), self.global_step, self.sample_rate
+                )
+                with torch.no_grad():
+                    mel = safe_log(self.melspec_loss.mel_spec(audio_input[0]))
+                    mel_hat = safe_log(self.melspec_loss.mel_spec(audio_hat[0]))
+                self.logger.experiment.add_image(
+                    "train/mel_target",
+                    plot_spectrogram_to_numpy(mel.data.cpu().numpy()),
+                    self.global_step,
+                    dataformats="HWC",
+                )
+                self.logger.experiment.add_image(
+                    "train/mel_pred",
+                    plot_spectrogram_to_numpy(mel_hat.data.cpu().numpy()),
+                    self.global_step,
+                    dataformats="HWC",
+                )
+            return loss
+
+    def on_validation_epoch_start(self):
+        if self.evaluate_utmos:
+            from model.vocos.metrics.UTMOS import UTMOSScore
+            # if not hasattr(self, "utmos_model"):
+            self.utmos_model = UTMOSScore(device=self.device)
+
+    def validation_step(self, batch, batch_idx, **kwargs):
+        audio_input = batch
+        audio_hat = self(audio_input, **kwargs)
+
+        audio_16_khz = torchaudio.functional.resample(audio_input, orig_freq=self.sample_rate, new_freq=16000)
+        audio_hat_16khz = torchaudio.functional.resample(audio_hat, orig_freq=self.sample_rate, new_freq=16000)
+
+        if self.evaluate_periodicty:
+            from model.vocos.metrics.periodicity import calculate_periodicity_metrics
+
+            periodicity_loss, pitch_loss, f1_score = calculate_periodicity_metrics(audio_16_khz, audio_hat_16khz)
+        else:
+            periodicity_loss = pitch_loss = f1_score = 0
+
+        if self.evaluate_utmos:
+            utmos_score = self.utmos_model.score(audio_hat_16khz.unsqueeze(1)).mean()
+        else:
+            utmos_score = torch.zeros(1, device=self.device)
+
+        if self.evaluate_pesq:
+            from pesq import pesq
+
+            pesq_score = 0
+            for ref, deg in zip(audio_16_khz.cpu().numpy(), audio_hat_16khz.cpu().numpy()):
+                pesq_score += pesq(16000, ref, deg, "wb", on_error=1)
+            pesq_score /= len(audio_16_khz)
+            pesq_score = torch.tensor(pesq_score)
+        else:
+            pesq_score = torch.zeros(1, device=self.device)
+
+        mel_loss = self.melspec_loss(audio_hat.unsqueeze(1), audio_input.unsqueeze(1))
+        total_loss = mel_loss + (5 - utmos_score) + (5 - pesq_score)
+
+        return {
+            "val_loss": total_loss,
+            "mel_loss": mel_loss,
+            "utmos_score": utmos_score,
+            "pesq_score": pesq_score,
+            "periodicity_loss": periodicity_loss,
+            "pitch_loss": pitch_loss,
+            "f1_score": f1_score,
+            "audio_input": audio_input[0],
+            "audio_pred": audio_hat[0],
+        }
+
+    def validation_epoch_end(self, outputs):
+        if self.global_rank == 0:
+            for i, output in enumerate(outputs):
+                *_, audio_in, audio_pred = output.values()
+                self.logger.experiment.add_audio(
+                    f"val_in_{i}", audio_in.data.cpu().numpy(), self.global_step, self.sample_rate
+                )
+                self.logger.experiment.add_audio(
+                    f"val_pred_{i}", audio_pred.data.cpu().numpy(), self.global_step, self.sample_rate
+                )
+                mel_target = safe_log(self.melspec_loss.mel_spec(audio_in))
+                mel_hat = safe_log(self.melspec_loss.mel_spec(audio_pred))
+                self.logger.experiment.add_image(
+                    f"val_mel_target_{i}",
+                    plot_spectrogram_to_numpy(mel_target.data.cpu().numpy()),
+                    self.global_step,
+                    dataformats="HWC",
+                )
+                self.logger.experiment.add_image(
+                    f"val_mel_hat_{i}",
+                    plot_spectrogram_to_numpy(mel_hat.data.cpu().numpy()),
+                    self.global_step,
+                    dataformats="HWC",
+                )
+        avg_loss = torch.stack([x["val_loss"] for x in outputs]).mean()
+        mel_loss = torch.stack([x["mel_loss"] for x in outputs]).mean()
+        utmos_score = torch.stack([x["utmos_score"] for x in outputs]).mean()
+        pesq_score = torch.stack([x["pesq_score"] for x in outputs]).mean()
+        periodicity_loss = np.array([x["periodicity_loss"] for x in outputs]).mean()
+        pitch_loss = np.array([x["pitch_loss"] for x in outputs]).mean()
+        f1_score = np.array([x["f1_score"] for x in outputs]).mean()
+
+        self.log("val_loss", avg_loss, sync_dist=True)
+        self.log("val/mel_loss", mel_loss, sync_dist=True)
+        self.log("val/utmos_score", utmos_score, sync_dist=True)
+        self.log("val/pesq_score", pesq_score, sync_dist=True)
+        self.log("val/periodicity_loss", periodicity_loss, sync_dist=True)
+        self.log("val/pitch_loss", pitch_loss, sync_dist=True)
+        self.log("val/f1_score", f1_score, sync_dist=True)
+
+        return {
+            "avg_loss": avg_loss,
+            "mel_loss": mel_loss,
+            "utmos_score": utmos_score,
+            "pesq_score": pesq_score,
+            "periodicity_loss": periodicity_loss,
+            "pitch_loss": pitch_loss,
+            "f1_score": f1_score,            
+            }
+    def on_test_epoch_start(self):
+        self.on_validation_epoch_start()
+
+    def test_step(self, *args, **kwargs):
+        return self.validation_step(*args, **kwargs)
+    
+    def test_epoch_end(self, outputs):
+        results = self.validation_epoch_end(outputs)
+        print(results)
+    @property
+    def global_step(self):
+        """
+        Override global_step so that it returns the total number of batches processed
+        """
+        return self.trainer.fit_loop.epoch_loop.total_batch_idx
+
+    def on_train_batch_start(self, *args):
+        if self.global_step >= self.pretrain_mel_steps:
+            self.train_discriminator = True
+        else:
+            self.train_discriminator = False
+
+    def on_train_batch_end(self, *args):
+        def mel_loss_coeff_decay(current_step, num_cycles=0.5):
+            max_steps = self.trainer.max_steps // 2
+            if current_step < self.num_warmup_steps:
+                return 1.0
+            progress = float(current_step - self.num_warmup_steps) / float(
+                max(1, max_steps - self.num_warmup_steps)
+            )
+            return max(0.0, 0.5 * (1.0 + math.cos(math.pi * float(num_cycles) * 2.0 * progress)))
+
+        if self.decay_mel_coeff:
+            self.mel_loss_coeff = self.base_mel_coeff * mel_loss_coeff_decay(self.global_step + 1)
+
+
+class VocosEncodecExp(VocosExp):
+    """
+    VocosEncodecExp is a subclass of VocosExp that overrides the parent experiment to function as a conditional GAN.
+    It manages an additional `bandwidth_id` attribute, which denotes a learnable embedding corresponding to
+    a specific bandwidth value of EnCodec. During training, a random bandwidth_id is generated for each step,
+    while during validation, a fixed bandwidth_id is used.
+    """
+
+    def __init__(
+        self,
+        feature_extractor: FeatureExtractor,
+        backbone: pl.LightningModule,
+        head: pl.LightningModule,
+        sample_rate: int,
+        initial_learning_rate: float,
+        num_warmup_steps: int,
+        mel_loss_coeff: float = 45,
+        mrd_loss_coeff: float = 1.0,
+        pretrain_mel_steps: int = 0,
+        decay_mel_coeff: bool = False,
+        evaluate_utmos: bool = False,
+        evaluate_pesq: bool = False,
+        evaluate_periodicty: bool = False,
+    ):
+        super().__init__(
+            feature_extractor,
+            backbone,
+            head,
+            sample_rate,
+            initial_learning_rate,
+            num_warmup_steps,
+            mel_loss_coeff,
+            mrd_loss_coeff,
+            pretrain_mel_steps,
+            decay_mel_coeff,
+            evaluate_utmos,
+            evaluate_pesq,
+            evaluate_periodicty,
+        )
+        # Override with conditional discriminators
+        self.multiperioddisc = MultiPeriodDiscriminator(num_embeddings=len(self.feature_extractor.bandwidths))
+        self.multiresddisc = MultiResolutionDiscriminator(num_embeddings=len(self.feature_extractor.bandwidths))
+
+    def training_step(self, *args):
+        bandwidth_id = torch.randint(low=0, high=len(self.feature_extractor.bandwidths), size=(1,), device=self.device,)
+        output = super().training_step(*args, bandwidth_id=bandwidth_id)
+        return output
+
+    def validation_step(self, *args):
+        bandwidth_id = torch.tensor([0], device=self.device)
+        output = super().validation_step(*args, bandwidth_id=bandwidth_id)
+        return output
+
+    def validation_epoch_end(self, outputs):
+        if self.global_rank == 0:
+            *_, audio_in, _ = outputs[0].values()
+            # Resynthesis with encodec for reference
+            self.feature_extractor.encodec.set_target_bandwidth(self.feature_extractor.bandwidths[0])
+            encodec_audio = self.feature_extractor.encodec(audio_in[None, None, :])
+            self.logger.experiment.add_audio(
+                "encodec", encodec_audio[0, 0].data.cpu().numpy(), self.global_step, self.sample_rate,
+            )
+
+        super().validation_epoch_end(outputs)

model/vocos/feature_extractors.py

@@ -0,0 +1,170 @@
+from typing import List
+
+import torch
+import librosa
+from encodec import EncodecModel
+from torch import nn
+from torch import Tensor
+from typing import Optional
+from torchaudio.transforms import Spectrogram, MelScale
+from model.vocos.modules import safe_log
+
+
+class FeatureExtractor(nn.Module):
+    """Base class for feature extractors."""
+
+    def forward(self, audio: torch.Tensor, **kwargs) -> torch.Tensor:
+        """
+        Extract features from the given audio.
+
+        Args:
+            audio (Tensor): Input audio waveform.
+
+        Returns:
+            Tensor: Extracted features of shape (B, C, L), where B is the batch size,
+                    C denotes output features, and L is the sequence length.
+        """
+        raise NotImplementedError("Subclasses must implement the forward method.")
+
+
+class LibrosaMelScale(nn.Module):
+    r"""This MelScale has a create_fb_matrix function that can be used to create a filterbank matrix.
+    same as previous torchaudio version
+    """
+    __constants__ = ["n_mels", "sample_rate", "f_min", "f_max"]
+
+    def __init__(
+        self,
+        n_mels: int = 128,
+        sample_rate: int = 16000,
+        f_min: float = 0.0,
+        f_max: Optional[float] = None,
+        n_stft: int = 201,
+        norm: Optional[str] = None,
+        mel_scale: str = "htk",
+    ) -> None:
+        super(LibrosaMelScale, self).__init__()
+        self.n_mels = n_mels
+        self.sample_rate = sample_rate
+        self.f_max = f_max if f_max is not None else float(sample_rate // 2)
+        self.f_min = f_min
+        self.norm = norm
+        self.mel_scale = mel_scale
+
+        if f_min > self.f_max:
+            raise ValueError("Require f_min: {} <= f_max: {}".format(f_min, self.f_max))
+        _mel_options = dict(
+            sr=sample_rate,
+            n_fft=(n_stft - 1) * 2,
+            n_mels=n_mels,
+            fmin=f_min,
+            fmax=f_max,
+            htk=mel_scale=="htk",
+            norm=norm
+        )
+        fb = torch.from_numpy(librosa.filters.mel(**_mel_options).T).float()
+        self.register_buffer("fb", fb)
+    
+    def forward(self, specgram):
+        mel_specgram = torch.matmul(specgram.transpose(-1, -2), self.fb).transpose(-1, -2)
+        return mel_specgram
+
+
+class MelSpectrogramFeatures(FeatureExtractor):
+    def __init__(
+        self,
+        sample_rate: int,
+        n_fft: int,
+        n_win: int,
+        n_hop: int,
+        n_mels: int,
+        f_min: int,
+        f_max: int,
+        power: int,
+        center: bool,
+        normalize: bool,
+        onesided: bool,
+        mel_norm: str | None,
+        mel_scale: str,
+        librosa_mel: bool = True,
+        clip_val: float = 1e-7
+        ):
+        super().__init__()
+        # This implementation vs torchaudio.transforms.MelSpectrogram: Add librosa melscale
+        # librosa melscale is numerically different from the torchaudio melscale (x_diff > 1e-5)
+        self.n_fft = n_fft   
+        self.spectrogram = Spectrogram(
+            n_fft=n_fft,
+            win_length=n_win,
+            hop_length=n_hop,
+            power=power,
+            normalized=normalize,
+            center=center,
+            onesided=onesided,
+        )
+        mel_method = LibrosaMelScale if librosa_mel else MelScale
+        self.mel_scale = mel_method(
+            n_mels=n_mels,
+            sample_rate=sample_rate,
+            f_min=f_min,
+            f_max=f_max,
+            n_stft=n_fft // 2 + 1,
+            norm=mel_norm,
+            mel_scale=mel_scale,
+        )
+        self.clip_val = clip_val
+        
+    def forward(self, x: Tensor) -> Tensor:
+        # Compute Spectrogram
+        specgram = self.spectrogram(x)
+        mel_specgram = self.mel_scale(specgram)
+        return safe_log(mel_specgram, self.clip_val)
+
+class EncodecFeatures(FeatureExtractor):
+    def __init__(
+        self,
+        encodec_model: str = "encodec_24khz",
+        bandwidths: List[float] = [1.5, 3.0, 6.0, 12.0],
+        train_codebooks: bool = False,
+    ):
+        super().__init__()
+        if encodec_model == "encodec_24khz":
+            encodec = EncodecModel.encodec_model_24khz
+        elif encodec_model == "encodec_48khz":
+            encodec = EncodecModel.encodec_model_48khz
+        else:
+            raise ValueError(
+                f"Unsupported encodec_model: {encodec_model}. Supported options are 'encodec_24khz' and 'encodec_48khz'."
+            )
+        self.encodec = encodec(pretrained=True)
+        for param in self.encodec.parameters():
+            param.requires_grad = False
+        self.num_q = self.encodec.quantizer.get_num_quantizers_for_bandwidth(
+            self.encodec.frame_rate, bandwidth=max(bandwidths)
+        )
+        codebook_weights = torch.cat([vq.codebook for vq in self.encodec.quantizer.vq.layers[: self.num_q]], dim=0)
+        self.codebook_weights = torch.nn.Parameter(codebook_weights, requires_grad=train_codebooks)
+        self.bandwidths = bandwidths
+
+    @torch.no_grad()
+    def get_encodec_codes(self, audio):
+        audio = audio.unsqueeze(1)
+        emb = self.encodec.encoder(audio)
+        codes = self.encodec.quantizer.encode(emb, self.encodec.frame_rate, self.encodec.bandwidth)
+        return codes
+
+    def forward(self, audio: torch.Tensor, **kwargs):
+        bandwidth_id = kwargs.get("bandwidth_id")
+        if bandwidth_id is None:
+            raise ValueError("The 'bandwidth_id' argument is required")
+        self.encodec.eval()  # Force eval mode as Pytorch Lightning automatically sets child modules to training mode
+        self.encodec.set_target_bandwidth(self.bandwidths[bandwidth_id])
+        codes = self.get_encodec_codes(audio)
+        # Instead of summing in the loop, it stores subsequent VQ dictionaries in a single `self.codebook_weights`
+        # with offsets given by the number of bins, and finally summed in a vectorized operation.
+        offsets = torch.arange(
+            0, self.encodec.quantizer.bins * len(codes), self.encodec.quantizer.bins, device=audio.device
+        )
+        embeddings_idxs = codes + offsets.view(-1, 1, 1)
+        features = torch.nn.functional.embedding(embeddings_idxs, self.codebook_weights).sum(dim=0)
+        return features.transpose(1, 2)

model/vocos/heads.py

@@ -0,0 +1,164 @@
+from typing import Optional
+
+import torch
+from torch import nn
+from torchaudio.functional.functional import _hz_to_mel, _mel_to_hz
+
+from model.vocos.spectral_ops import IMDCT, ISTFT
+from model.vocos.modules import symexp
+
+
+class FourierHead(nn.Module):
+    """Base class for inverse fourier modules."""
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            x (Tensor): Input tensor of shape (B, L, H), where B is the batch size,
+                        L is the sequence length, and H denotes the model dimension.
+
+        Returns:
+            Tensor: Reconstructed time-domain audio signal of shape (B, T), where T is the length of the output signal.
+        """
+        raise NotImplementedError("Subclasses must implement the forward method.")
+
+
+class ISTFTHead(FourierHead):
+    """
+    ISTFT Head module for predicting STFT complex coefficients.
+
+    Args:
+        dim (int): Hidden dimension of the model.
+        n_fft (int): Size of Fourier transform.
+        hop_length (int): The distance between neighboring sliding window frames, which should align with
+                          the resolution of the input features.
+        padding (str, optional): Type of padding. Options are "center" or "same". Defaults to "same".
+    """
+
+    def __init__(self, dim: int, n_fft: int, hop_length: int, padding: str = "same"):
+        super().__init__()
+        out_dim = n_fft + 2
+        self.out = torch.nn.Linear(dim, out_dim)
+        self.istft = ISTFT(n_fft=n_fft, hop_length=hop_length, win_length=n_fft, padding=padding)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Forward pass of the ISTFTHead module.
+
+        Args:
+            x (Tensor): Input tensor of shape (B, L, H), where B is the batch size,
+                        L is the sequence length, and H denotes the model dimension.
+
+        Returns:
+            Tensor: Reconstructed time-domain audio signal of shape (B, T), where T is the length of the output signal.
+        """
+        x = self.out(x).transpose(1, 2)
+        mag, p = x.chunk(2, dim=1)
+        mag = torch.exp(mag)
+        mag = torch.clip(mag, max=1e2)  # safeguard to prevent excessively large magnitudes
+        # wrapping happens here. These two lines produce real and imaginary value
+        x = torch.cos(p)
+        y = torch.sin(p)
+        # recalculating phase here does not produce anything new
+        # only costs time
+        # phase = torch.atan2(y, x)
+        # S = mag * torch.exp(phase * 1j)
+        # better directly produce the complex value 
+        S = mag * (x + 1j * y)
+        audio = self.istft(S)
+        return audio
+
+
+class IMDCTSymExpHead(FourierHead):
+    """
+    IMDCT Head module for predicting MDCT coefficients with symmetric exponential function
+
+    Args:
+        dim (int): Hidden dimension of the model.
+        mdct_frame_len (int): Length of the MDCT frame.
+        padding (str, optional): Type of padding. Options are "center" or "same". Defaults to "same".
+        sample_rate (int, optional): The sample rate of the audio. If provided, the last layer will be initialized
+                                     based on perceptual scaling. Defaults to None.
+        clip_audio (bool, optional): Whether to clip the audio output within the range of [-1.0, 1.0]. Defaults to False.
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        mdct_frame_len: int,
+        padding: str = "same",
+        sample_rate: Optional[int] = None,
+        clip_audio: bool = False,
+    ):
+        super().__init__()
+        out_dim = mdct_frame_len // 2
+        self.out = nn.Linear(dim, out_dim)
+        self.imdct = IMDCT(frame_len=mdct_frame_len, padding=padding)
+        self.clip_audio = clip_audio
+
+        if sample_rate is not None:
+            # optionally init the last layer following mel-scale
+            m_max = _hz_to_mel(sample_rate // 2)
+            m_pts = torch.linspace(0, m_max, out_dim)
+            f_pts = _mel_to_hz(m_pts)
+            scale = 1 - (f_pts / f_pts.max())
+
+            with torch.no_grad():
+                self.out.weight.mul_(scale.view(-1, 1))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Forward pass of the IMDCTSymExpHead module.
+
+        Args:
+            x (Tensor): Input tensor of shape (B, L, H), where B is the batch size,
+                        L is the sequence length, and H denotes the model dimension.
+
+        Returns:
+            Tensor: Reconstructed time-domain audio signal of shape (B, T), where T is the length of the output signal.
+        """
+        x = self.out(x)
+        x = symexp(x)
+        x = torch.clip(x, min=-1e2, max=1e2)  # safeguard to prevent excessively large magnitudes
+        audio = self.imdct(x)
+        if self.clip_audio:
+            audio = torch.clip(x, min=-1.0, max=1.0)
+
+        return audio
+
+
+class IMDCTCosHead(FourierHead):
+    """
+    IMDCT Head module for predicting MDCT coefficients with parametrizing MDCT = exp(m) · cos(p)
+
+    Args:
+        dim (int): Hidden dimension of the model.
+        mdct_frame_len (int): Length of the MDCT frame.
+        padding (str, optional): Type of padding. Options are "center" or "same". Defaults to "same".
+        clip_audio (bool, optional): Whether to clip the audio output within the range of [-1.0, 1.0]. Defaults to False.
+    """
+
+    def __init__(self, dim: int, mdct_frame_len: int, padding: str = "same", clip_audio: bool = False):
+        super().__init__()
+        self.clip_audio = clip_audio
+        self.out = nn.Linear(dim, mdct_frame_len)
+        self.imdct = IMDCT(frame_len=mdct_frame_len, padding=padding)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Forward pass of the IMDCTCosHead module.
+
+        Args:
+            x (Tensor): Input tensor of shape (B, L, H), where B is the batch size,
+                        L is the sequence length, and H denotes the model dimension.
+
+        Returns:
+            Tensor: Reconstructed time-domain audio signal of shape (B, T), where T is the length of the output signal.
+        """
+        x = self.out(x)
+        m, p = x.chunk(2, dim=2)
+        m = torch.exp(m).clip(max=1e2)  # safeguard to prevent excessively large magnitudes
+        audio = self.imdct(m * torch.cos(p))
+        if self.clip_audio:
+            audio = torch.clip(x, min=-1.0, max=1.0)
+        return audio

model/vocos/helpers.py

@@ -0,0 +1,71 @@
+import matplotlib
+import numpy as np
+import torch
+from matplotlib import pyplot as plt
+from pytorch_lightning import Callback
+
+matplotlib.use("Agg")
+
+
+def save_figure_to_numpy(fig: plt.Figure) -> np.ndarray:
+    """
+    Save a matplotlib figure to a numpy array.
+
+    Args:
+        fig (Figure): Matplotlib figure object.
+
+    Returns:
+        ndarray: Numpy array representing the figure.
+    """
+    data = np.fromstring(fig.canvas.tostring_rgb(), dtype=np.uint8, sep="")
+    data = data.reshape(fig.canvas.get_width_height()[::-1] + (3,))
+    return data
+
+
+def plot_spectrogram_to_numpy(spectrogram: np.ndarray) -> np.ndarray:
+    """
+    Plot a spectrogram and convert it to a numpy array.
+
+    Args:
+        spectrogram (ndarray): Spectrogram data.
+
+    Returns:
+        ndarray: Numpy array representing the plotted spectrogram.
+    """
+    spectrogram = spectrogram.astype(np.float32)
+    fig, ax = plt.subplots(figsize=(12, 3))
+    im = ax.imshow(spectrogram, aspect="auto", origin="lower", interpolation="none")
+    plt.colorbar(im, ax=ax)
+    plt.xlabel("Frames")
+    plt.ylabel("Channels")
+    plt.tight_layout()
+
+    fig.canvas.draw()
+    data = save_figure_to_numpy(fig)
+    plt.close()
+    return data
+
+
+class GradNormCallback(Callback):
+    """
+    Callback to log the gradient norm.
+    """
+
+    def on_after_backward(self, trainer, model):
+        model.log("grad_norm", gradient_norm(model))
+
+
+def gradient_norm(model: torch.nn.Module, norm_type: float = 2.0) -> torch.Tensor:
+    """
+    Compute the gradient norm.
+
+    Args:
+        model (Module): PyTorch model.
+        norm_type (float, optional): Type of the norm. Defaults to 2.0.
+
+    Returns:
+        Tensor: Gradient norm.
+    """
+    grads = [p.grad for p in model.parameters() if p.grad is not None]
+    total_norm = torch.norm(torch.stack([torch.norm(g.detach(), norm_type) for g in grads]), norm_type)
+    return total_norm

model/vocos/loss.py

@@ -0,0 +1,114 @@
+from typing import List, Tuple
+
+import torch
+import torchaudio
+from torch import nn
+
+from model.vocos.offline.modules import safe_log
+
+
+class MelSpecReconstructionLoss(nn.Module):
+    """
+    L1 distance between the mel-scaled magnitude spectrograms of the ground truth sample and the generated sample
+    """
+
+    def __init__(
+        self, sample_rate: int = 24000, n_fft: int = 1024, hop_length: int = 256, n_mels: int = 100,
+    ):
+        super().__init__()
+        self.mel_spec = torchaudio.transforms.MelSpectrogram(
+            sample_rate=sample_rate, n_fft=n_fft, hop_length=hop_length, n_mels=n_mels, center=True, power=1,
+        )
+
+    def forward(self, y_hat, y) -> torch.Tensor:
+        """
+        Args:
+            y_hat (Tensor): Predicted audio waveform.
+            y (Tensor): Ground truth audio waveform.
+
+        Returns:
+            Tensor: L1 loss between the mel-scaled magnitude spectrograms.
+        """
+        mel_hat = safe_log(self.mel_spec(y_hat))
+        mel = safe_log(self.mel_spec(y))
+
+        loss = torch.nn.functional.l1_loss(mel, mel_hat)
+
+        return loss
+
+
+class GeneratorLoss(nn.Module):
+    """
+    Generator Loss module. Calculates the loss for the generator based on discriminator outputs.
+    """
+
+    def forward(self, disc_outputs: List[torch.Tensor]) -> Tuple[torch.Tensor, List[torch.Tensor]]:
+        """
+        Args:
+            disc_outputs (List[Tensor]): List of discriminator outputs.
+
+        Returns:
+            Tuple[Tensor, List[Tensor]]: Tuple containing the total loss and a list of loss values from
+                                         the sub-discriminators
+        """
+        loss = torch.zeros(1, device=disc_outputs[0].device, dtype=disc_outputs[0].dtype)
+        gen_losses = []
+        for dg in disc_outputs:
+            l = torch.mean(torch.clamp(1 - dg, min=0))
+            gen_losses.append(l)
+            loss += l
+
+        return loss, gen_losses
+
+
+class DiscriminatorLoss(nn.Module):
+    """
+    Discriminator Loss module. Calculates the loss for the discriminator based on real and generated outputs.
+    """
+
+    def forward(
+        self, disc_real_outputs: List[torch.Tensor], disc_generated_outputs: List[torch.Tensor]
+    ) -> Tuple[torch.Tensor, List[torch.Tensor], List[torch.Tensor]]:
+        """
+        Args:
+            disc_real_outputs (List[Tensor]): List of discriminator outputs for real samples.
+            disc_generated_outputs (List[Tensor]): List of discriminator outputs for generated samples.
+
+        Returns:
+            Tuple[Tensor, List[Tensor], List[Tensor]]: A tuple containing the total loss, a list of loss values from
+                                                       the sub-discriminators for real outputs, and a list of
+                                                       loss values for generated outputs.
+        """
+        loss = torch.zeros(1, device=disc_real_outputs[0].device, dtype=disc_real_outputs[0].dtype)
+        r_losses = []
+        g_losses = []
+        for dr, dg in zip(disc_real_outputs, disc_generated_outputs):
+            r_loss = torch.mean(torch.clamp(1 - dr, min=0))
+            g_loss = torch.mean(torch.clamp(1 + dg, min=0))
+            loss += r_loss + g_loss
+            r_losses.append(r_loss)
+            g_losses.append(g_loss)
+
+        return loss, r_losses, g_losses
+
+
+class FeatureMatchingLoss(nn.Module):
+    """
+    Feature Matching Loss module. Calculates the feature matching loss between feature maps of the sub-discriminators.
+    """
+
+    def forward(self, fmap_r: List[List[torch.Tensor]], fmap_g: List[List[torch.Tensor]]) -> torch.Tensor:
+        """
+        Args:
+            fmap_r (List[List[Tensor]]): List of feature maps from real samples.
+            fmap_g (List[List[Tensor]]): List of feature maps from generated samples.
+
+        Returns:
+            Tensor: The calculated feature matching loss.
+        """
+        loss = torch.zeros(1, device=fmap_r[0][0].device, dtype=fmap_r[0][0].dtype)
+        for dr, dg in zip(fmap_r, fmap_g):
+            for rl, gl in zip(dr, dg):
+                loss += torch.mean(torch.abs(rl - gl))
+
+        return loss

model/vocos/models.py

@@ -0,0 +1,118 @@
+from typing import Optional
+
+import torch
+from torch import nn
+from torch.nn.utils import weight_norm
+
+from model.vocos.modules import ConvNeXtBlock, ResBlock1, AdaLayerNorm
+
+
+class Backbone(nn.Module):
+    """Base class for the generator's backbone. It preserves the same temporal resolution across all layers."""
+
+    def forward(self, x: torch.Tensor, **kwargs) -> torch.Tensor:
+        """
+        Args:
+            x (Tensor): Input tensor of shape (B, C, L), where B is the batch size,
+                        C denotes output features, and L is the sequence length.
+
+        Returns:
+            Tensor: Output of shape (B, L, H), where B is the batch size, L is the sequence length,
+                    and H denotes the model dimension.
+        """
+        raise NotImplementedError("Subclasses must implement the forward method.")
+
+
+class VocosBackbone(Backbone):
+    """
+    Vocos backbone module built with ConvNeXt blocks. Supports additional conditioning with Adaptive Layer Normalization
+
+    Args:
+        input_channels (int): Number of input features channels.
+        dim (int): Hidden dimension of the model.
+        intermediate_dim (int): Intermediate dimension used in ConvNeXtBlock.
+        num_layers (int): Number of ConvNeXtBlock layers.
+        layer_scale_init_value (float, optional): Initial value for layer scaling. Defaults to `1 / num_layers`.
+        adanorm_num_embeddings (int, optional): Number of embeddings for AdaLayerNorm.
+                                                None means non-conditional model. Defaults to None.
+    """
+
+    def __init__(
+        self,
+        input_channels: int,
+        dim: int,
+        intermediate_dim: int,
+        num_layers: int,
+        layer_scale_init_value: Optional[float] = None,
+        adanorm_num_embeddings: Optional[int] = None,
+    ):
+        super().__init__()
+        self.input_channels = input_channels
+        self.embed = nn.Conv1d(input_channels, dim, kernel_size=7, padding=3)
+        self.adanorm = adanorm_num_embeddings is not None
+        if adanorm_num_embeddings:
+            self.norm = AdaLayerNorm(adanorm_num_embeddings, dim, eps=1e-6)
+        else:
+            self.norm = nn.LayerNorm(dim, eps=1e-6)
+        layer_scale_init_value = layer_scale_init_value or 1 / num_layers
+        self.convnext = nn.ModuleList(
+            [
+                ConvNeXtBlock(
+                    dim=dim,
+                    intermediate_dim=intermediate_dim,
+                    layer_scale_init_value=layer_scale_init_value,
+                    adanorm_num_embeddings=adanorm_num_embeddings,
+                )
+                for _ in range(num_layers)
+            ]
+        )
+        self.final_layer_norm = nn.LayerNorm(dim, eps=1e-6)
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, (nn.Conv1d, nn.Linear)):
+            nn.init.trunc_normal_(m.weight, std=0.02)
+            nn.init.constant_(m.bias, 0)
+
+    def forward(self, x: torch.Tensor, **kwargs) -> torch.Tensor:
+        bandwidth_id = kwargs.get('bandwidth_id', None)
+        x = self.embed(x)
+        if self.adanorm:
+            assert bandwidth_id is not None
+            x = self.norm(x.transpose(1, 2), cond_embedding_id=bandwidth_id)
+        else:
+            x = self.norm(x.transpose(1, 2))
+        x = x.transpose(1, 2)
+        for conv_block in self.convnext:
+            x = conv_block(x, cond_embedding_id=bandwidth_id)
+        x = self.final_layer_norm(x.transpose(1, 2))
+        return x
+
+
+class VocosResNetBackbone(Backbone):
+    """
+    Vocos backbone module built with ResBlocks.
+
+    Args:
+        input_channels (int): Number of input features channels.
+        dim (int): Hidden dimension of the model.
+        num_blocks (int): Number of ResBlock1 blocks.
+        layer_scale_init_value (float, optional): Initial value for layer scaling. Defaults to None.
+    """
+
+    def __init__(
+        self, input_channels, dim, num_blocks, layer_scale_init_value=None,
+    ):
+        super().__init__()
+        self.input_channels = input_channels
+        self.embed = weight_norm(nn.Conv1d(input_channels, dim, kernel_size=3, padding=1))
+        layer_scale_init_value = layer_scale_init_value or 1 / num_blocks / 3
+        self.resnet = nn.Sequential(
+            *[ResBlock1(dim=dim, layer_scale_init_value=layer_scale_init_value) for _ in range(num_blocks)]
+        )
+
+    def forward(self, x: torch.Tensor, **kwargs) -> torch.Tensor:
+        x = self.embed(x)
+        x = self.resnet(x)
+        x = x.transpose(1, 2)
+        return x

model/vocos/modules.py

@@ -0,0 +1,213 @@
+from typing import Optional, Tuple
+
+import torch
+from torch import nn
+from torch.nn.utils import weight_norm, remove_weight_norm
+
+
+class ConvNeXtBlock(nn.Module):
+    """ConvNeXt Block adapted from https://github.com/facebookresearch/ConvNeXt to 1D audio signal.
+
+    Args:
+        dim (int): Number of input channels.
+        intermediate_dim (int): Dimensionality of the intermediate layer.
+        layer_scale_init_value (float, optional): Initial value for the layer scale. None means no scaling.
+            Defaults to None.
+        adanorm_num_embeddings (int, optional): Number of embeddings for AdaLayerNorm.
+            None means non-conditional LayerNorm. Defaults to None.
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        intermediate_dim: int,
+        layer_scale_init_value: float,
+        adanorm_num_embeddings: Optional[int] = None,
+    ):
+        super().__init__()
+        self.dwconv = nn.Conv1d(dim, dim, kernel_size=7, padding=3, groups=dim)  # depthwise conv
+        self.adanorm = adanorm_num_embeddings is not None
+        if adanorm_num_embeddings:
+            self.norm = AdaLayerNorm(adanorm_num_embeddings, dim, eps=1e-6)
+        else:
+            self.norm = nn.LayerNorm(dim, eps=1e-6)
+        self.pwconv1 = nn.Linear(dim, intermediate_dim)  # pointwise/1x1 convs, implemented with linear layers
+        self.act = nn.GELU()
+        self.pwconv2 = nn.Linear(intermediate_dim, dim)
+        self.gamma = (
+            nn.Parameter(layer_scale_init_value * torch.ones(dim), requires_grad=True)
+            if layer_scale_init_value > 0
+            else None
+        )
+
+    def forward(self, x: torch.Tensor, cond_embedding_id: Optional[torch.Tensor] = None) -> torch.Tensor:
+        residual = x
+        x = self.dwconv(x)
+        x = x.transpose(1, 2)  # (B, C, T) -> (B, T, C)
+        if self.adanorm:
+            assert cond_embedding_id is not None
+            x = self.norm(x, cond_embedding_id)
+        else:
+            x = self.norm(x)
+        x = self.pwconv1(x)
+        x = self.act(x)
+        x = self.pwconv2(x)
+        if self.gamma is not None:
+            x = self.gamma * x
+        x = x.transpose(1, 2)  # (B, T, C) -> (B, C, T)
+
+        x = residual + x
+        return x
+
+
+class AdaLayerNorm(nn.Module):
+    """
+    Adaptive Layer Normalization module with learnable embeddings per `num_embeddings` classes
+
+    Args:
+        num_embeddings (int): Number of embeddings.
+        embedding_dim (int): Dimension of the embeddings.
+    """
+
+    def __init__(self, num_embeddings: int, embedding_dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.eps = eps
+        self.dim = embedding_dim
+        self.scale = nn.Embedding(num_embeddings=num_embeddings, embedding_dim=embedding_dim)
+        self.shift = nn.Embedding(num_embeddings=num_embeddings, embedding_dim=embedding_dim)
+        torch.nn.init.ones_(self.scale.weight)
+        torch.nn.init.zeros_(self.shift.weight)
+
+    def forward(self, x: torch.Tensor, cond_embedding_id: torch.Tensor) -> torch.Tensor:
+        scale = self.scale(cond_embedding_id)
+        shift = self.shift(cond_embedding_id)
+        x = nn.functional.layer_norm(x, (self.dim,), eps=self.eps)
+        x = x * scale + shift
+        return x
+
+
+class ResBlock1(nn.Module):
+    """
+    ResBlock adapted from HiFi-GAN V1 (https://github.com/jik876/hifi-gan) with dilated 1D convolutions,
+    but without upsampling layers.
+
+    Args:
+        dim (int): Number of input channels.
+        kernel_size (int, optional): Size of the convolutional kernel. Defaults to 3.
+        dilation (tuple[int], optional): Dilation factors for the dilated convolutions.
+            Defaults to (1, 3, 5).
+        lrelu_slope (float, optional): Negative slope of the LeakyReLU activation function.
+            Defaults to 0.1.
+        layer_scale_init_value (float, optional): Initial value for the layer scale. None means no scaling.
+            Defaults to None.
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        kernel_size: int = 3,
+        dilation: Tuple[int, int, int] = (1, 3, 5),
+        lrelu_slope: float = 0.1,
+        layer_scale_init_value: Optional[float] = None,
+    ):
+        super().__init__()
+        self.lrelu_slope = lrelu_slope
+        self.convs1 = nn.ModuleList(
+            [
+                weight_norm(
+                    nn.Conv1d(
+                        dim,
+                        dim,
+                        kernel_size,
+                        1,
+                        dilation=dilation[0],
+                        padding=self.get_padding(kernel_size, dilation[0]),
+                    )
+                ),
+                weight_norm(
+                    nn.Conv1d(
+                        dim,
+                        dim,
+                        kernel_size,
+                        1,
+                        dilation=dilation[1],
+                        padding=self.get_padding(kernel_size, dilation[1]),
+                    )
+                ),
+                weight_norm(
+                    nn.Conv1d(
+                        dim,
+                        dim,
+                        kernel_size,
+                        1,
+                        dilation=dilation[2],
+                        padding=self.get_padding(kernel_size, dilation[2]),
+                    )
+                ),
+            ]
+        )
+
+        self.convs2 = nn.ModuleList(
+            [
+                weight_norm(nn.Conv1d(dim, dim, kernel_size, 1, dilation=1, padding=self.get_padding(kernel_size, 1))),
+                weight_norm(nn.Conv1d(dim, dim, kernel_size, 1, dilation=1, padding=self.get_padding(kernel_size, 1))),
+                weight_norm(nn.Conv1d(dim, dim, kernel_size, 1, dilation=1, padding=self.get_padding(kernel_size, 1))),
+            ]
+        )
+
+        self.gamma = nn.ParameterList(
+            [
+                nn.Parameter(layer_scale_init_value * torch.ones(dim, 1), requires_grad=True)
+                if layer_scale_init_value is not None
+                else None,
+                nn.Parameter(layer_scale_init_value * torch.ones(dim, 1), requires_grad=True)
+                if layer_scale_init_value is not None
+                else None,
+                nn.Parameter(layer_scale_init_value * torch.ones(dim, 1), requires_grad=True)
+                if layer_scale_init_value is not None
+                else None,
+            ]
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        for c1, c2, gamma in zip(self.convs1, self.convs2, self.gamma):
+            xt = torch.nn.functional.leaky_relu(x, negative_slope=self.lrelu_slope)
+            xt = c1(xt)
+            xt = torch.nn.functional.leaky_relu(xt, negative_slope=self.lrelu_slope)
+            xt = c2(xt)
+            if gamma is not None:
+                xt = gamma * xt
+            x = xt + x
+        return x
+
+    def remove_weight_norm(self):
+        for l in self.convs1:
+            remove_weight_norm(l)
+        for l in self.convs2:
+            remove_weight_norm(l)
+
+    @staticmethod
+    def get_padding(kernel_size: int, dilation: int = 1) -> int:
+        return int((kernel_size * dilation - dilation) / 2)
+
+
+def safe_log(x: torch.Tensor, clip_val: float = 1e-7) -> torch.Tensor:
+    """
+    Computes the element-wise logarithm of the input tensor with clipping to avoid near-zero values.
+
+    Args:
+        x (Tensor): Input tensor.
+        clip_val (float, optional): Minimum value to clip the input tensor. Defaults to 1e-7.
+
+    Returns:
+        Tensor: Element-wise logarithm of the input tensor with clipping applied.
+    """
+    return torch.log(torch.clip(x, min=clip_val))
+
+
+def symlog(x: torch.Tensor) -> torch.Tensor:
+    return torch.sign(x) * torch.log1p(x.abs())
+
+
+def symexp(x: torch.Tensor) -> torch.Tensor:
+    return torch.sign(x) * (torch.exp(x.abs()) - 1)

model/vocos/pretrained.py

@@ -0,0 +1,162 @@
+from __future__ import annotations
+
+from typing import Any, Dict, Tuple, Union, Optional
+
+import torch
+import yaml
+from huggingface_hub import hf_hub_download
+from torch import nn
+from model.vocos.feature_extractors import FeatureExtractor, EncodecFeatures
+from model.vocos.heads import FourierHead
+from model.vocos.models import Backbone
+
+
+def instantiate_class(args: Union[Any, Tuple[Any, ...]], init: Dict[str, Any]) -> Any:
+    """Instantiates a class with the given args and init.
+
+    Args:
+        args: Positional arguments required for instantiation.
+        init: Dict of the form {"class_path":...,"init_args":...}.
+
+    Returns:
+        The instantiated class object.
+    """
+    kwargs = init.get("init_args", {})
+    if not isinstance(args, tuple):
+        args = (args,)
+    class_module, class_name = init["class_path"].rsplit(".", 1)
+    module = __import__(class_module, fromlist=[class_name])
+    args_class = getattr(module, class_name)
+    return args_class(*args, **kwargs)
+
+
+class Vocos(nn.Module):
+    """
+    The Vocos class represents a Fourier-based neural vocoder for audio synthesis.
+    This class is primarily designed for inference, with support for loading from pretrained
+    model checkpoints. It consists of three main components: a feature extractor,
+    a backbone, and a head.
+    """
+
+    def __init__(
+        self, feature_extractor: nn.Module, backbone: Backbone, head: FourierHead,
+    ):
+        super().__init__()
+        self.feature_extractor = feature_extractor
+        self.backbone = backbone
+        self.head = head
+
+    @classmethod
+    def from_hparams(cls, config_path: str) -> "Vocos":
+        """
+        Class method to create a new Vocos model instance from hyperparameters stored in a yaml configuration file.
+        """
+        with open(config_path, "r") as f:
+            config = yaml.safe_load(f)
+        feature_extractor = instantiate_class(args=(), init=config["feature_extractor"])
+        backbone = instantiate_class(args=(), init=config["backbone"])
+        head = instantiate_class(args=(), init=config["head"])
+        model = cls(feature_extractor=feature_extractor, backbone=backbone, head=head)
+        return model
+
+    @classmethod
+    def from_pretrained(self, config_path: str, model_path: str, model: nn.Module=None) -> "Vocos":
+        """
+        Class method to create a new Vocos model instance from a pre-trained model stored in the Hugging Face model hub.
+        """
+        if model is None:
+            model = self.from_hparams(config_path)
+        state_dict = torch.load(model_path, map_location="cpu")
+        prefixes = ("backbone", "feature_extractor", "head")
+        state_dict = {
+            key: value
+            for key, value in state_dict.items()
+            if any(key.startswith(prefix) for prefix in prefixes)
+        }
+        if isinstance(model.feature_extractor, EncodecFeatures):
+            encodec_parameters = {
+                "feature_extractor.encodec." + key: value
+                for key, value in model.feature_extractor.encodec.state_dict().items()
+            }
+            state_dict.update(encodec_parameters)
+        model.load_state_dict(state_dict)
+        model.eval()
+        return model
+
+    @torch.inference_mode()
+    def forward(self, features_input: torch.Tensor, X_norm, **kwargs: Any) -> torch.Tensor:
+        """
+        Method to run a copy-synthesis from audio waveform. The feature extractor first processes the audio input,
+        which is then passed through the backbone and the head to reconstruct the audio output.
+
+        Args:
+            audio_input (Tensor): The input tensor representing the audio waveform of shape (B, T),
+                                        where B is the batch size and L is the waveform length.
+
+
+        Returns:
+            Tensor: The output tensor representing the reconstructed audio waveform of shape (B, T).
+        """
+        audio_output = self.decode(features_input, **kwargs)
+        return audio_output / X_norm
+
+    @torch.inference_mode()
+    def decode(self, features_input: torch.Tensor, **kwargs: Any) -> torch.Tensor:
+        """
+        Method to decode audio waveform from already calculated features. The features input is passed through
+        the backbone and the head to reconstruct the audio output.
+
+        Args:
+            features_input (Tensor): The input tensor of features of shape (B, C, L), where B is the batch size,
+                                     C denotes the feature dimension, and L is the sequence length.
+
+        Returns:
+            Tensor: The output tensor representing the reconstructed audio waveform of shape (B, T).
+        """
+        x = self.backbone(features_input, **kwargs)
+        audio_output = self.head(x)
+        return audio_output
+
+    @torch.inference_mode()
+    def codes_to_features(self, codes: torch.Tensor) -> torch.Tensor:
+        """
+        Transforms an input sequence of discrete tokens (codes) into feature embeddings using the feature extractor's
+        codebook weights.
+
+        Args:
+            codes (Tensor): The input tensor. Expected shape is (K, L) or (K, B, L),
+                            where K is the number of codebooks, B is the batch size and L is the sequence length.
+
+        Returns:
+            Tensor: Features of shape (B, C, L), where B is the batch size, C denotes the feature dimension,
+                    and L is the sequence length.
+        """
+        assert isinstance(
+            self.feature_extractor, EncodecFeatures
+        ), "Feature extractor should be an instance of EncodecFeatures"
+
+        if codes.dim() == 2:
+            codes = codes.unsqueeze(1)
+
+        n_bins = self.feature_extractor.encodec.quantizer.bins
+        offsets = torch.arange(0, n_bins * len(codes), n_bins, device=codes.device)
+        embeddings_idxs = codes + offsets.view(-1, 1, 1)
+        features = torch.nn.functional.embedding(embeddings_idxs, self.feature_extractor.codebook_weights).sum(dim=0)
+        features = features.transpose(1, 2)
+
+        return features
+
+
+if __name__ == "__main__":
+    model = Vocos.from_pretrained(
+        "/nvmework3/shaonian/MelSpatialNet/MelSpatialNet/models/vocos/pretrained/pretrained_rec_normed.yaml", 
+        "/nvmework3/shaonian/MelSpatialNet/MelSpatialNet/models/vocos/pretrained/vocos_hop128_clip1e-5_rts.ckpt").to("meta")
+    x = torch.randn(1, 80, 501)  
+    x = x.to('meta')
+    from torch.utils.flop_counter import FlopCounterMode # requires torch>=2.1.0
+    with FlopCounterMode(model, display=False) as fcm:
+        y = model.decode(x)
+        flops_forward_eval = fcm.get_total_flops()
+
+    params_eval = sum(param.numel() for param in model.parameters())
+    print(f"flops_forward={flops_forward_eval/4e9:.2f}G, params={params_eval/1e6:.2f} M")

model/vocos/spectral_ops.py

@@ -0,0 +1,192 @@
+import numpy as np
+import scipy
+import torch
+from torch import nn, view_as_real, view_as_complex
+
+
+class ISTFT(nn.Module):
+    """
+    Custom implementation of ISTFT since torch.istft doesn't allow custom padding (other than `center=True`) with
+    windowing. This is because the NOLA (Nonzero Overlap Add) check fails at the edges.
+    See issue: https://github.com/pytorch/pytorch/issues/62323
+    Specifically, in the context of neural vocoding we are interested in "same" padding analogous to CNNs.
+    The NOLA constraint is met as we trim padded samples anyway.
+
+    Args:
+        n_fft (int): Size of Fourier transform.
+        hop_length (int): The distance between neighboring sliding window frames.
+        win_length (int): The size of window frame and STFT filter.
+        padding (str, optional): Type of padding. Options are "center" or "same". Defaults to "same".
+    """
+
+    def __init__(self, n_fft: int, hop_length: int, win_length: int, padding: str = "same"):
+        super().__init__()
+        if padding not in ["center", "same"]:
+            raise ValueError("Padding must be 'center' or 'same'.")
+        self.padding = padding
+        self.n_fft = n_fft
+        self.hop_length = hop_length
+        self.win_length = win_length
+        window = torch.hann_window(win_length)
+        self.register_buffer("window", window)
+
+    def forward(self, spec: torch.Tensor) -> torch.Tensor:
+        """
+        Compute the Inverse Short Time Fourier Transform (ISTFT) of a complex spectrogram.
+
+        Args:
+            spec (Tensor): Input complex spectrogram of shape (B, N, T), where B is the batch size,
+                            N is the number of frequency bins, and T is the number of time frames.
+
+        Returns:
+            Tensor: Reconstructed time-domain signal of shape (B, L), where L is the length of the output signal.
+        """
+        if self.padding == "center":
+            # Fallback to pytorch native implementation
+            return torch.istft(spec, self.n_fft, self.hop_length, self.win_length, self.window, center=True)
+        elif self.padding == "same":
+            pad = (self.win_length - self.hop_length) // 2
+        else:
+            raise ValueError("Padding must be 'center' or 'same'.")
+
+        assert spec.dim() == 3, "Expected a 3D tensor as input"
+        B, N, T = spec.shape
+
+        # Inverse FFT
+        ifft = torch.fft.irfft(spec, self.n_fft, dim=1, norm="backward")
+        ifft = ifft * self.window[None, :, None]
+
+        # Overlap and Add
+        output_size = (T - 1) * self.hop_length + self.win_length
+        y = torch.nn.functional.fold(
+            ifft, output_size=(1, output_size), kernel_size=(1, self.win_length), stride=(1, self.hop_length),
+        )[:, 0, 0, pad:-pad]
+
+        # Window envelope
+        window_sq = self.window.square().expand(1, T, -1).transpose(1, 2)
+        window_envelope = torch.nn.functional.fold(
+            window_sq, output_size=(1, output_size), kernel_size=(1, self.win_length), stride=(1, self.hop_length),
+        ).squeeze()[pad:-pad]
+
+        # Normalize
+        assert (window_envelope > 1e-11).all()
+        y = y / window_envelope
+
+        return y
+
+
+class MDCT(nn.Module):
+    """
+    Modified Discrete Cosine Transform (MDCT) module.
+
+    Args:
+        frame_len (int): Length of the MDCT frame.
+        padding (str, optional): Type of padding. Options are "center" or "same". Defaults to "same".
+    """
+
+    def __init__(self, frame_len: int, padding: str = "same"):
+        super().__init__()
+        if padding not in ["center", "same"]:
+            raise ValueError("Padding must be 'center' or 'same'.")
+        self.padding = padding
+        self.frame_len = frame_len
+        N = frame_len // 2
+        n0 = (N + 1) / 2
+        window = torch.from_numpy(scipy.signal.cosine(frame_len)).float()
+        self.register_buffer("window", window)
+
+        pre_twiddle = torch.exp(-1j * torch.pi * torch.arange(frame_len) / frame_len)
+        post_twiddle = torch.exp(-1j * torch.pi * n0 * (torch.arange(N) + 0.5) / N)
+        # view_as_real: NCCL Backend does not support ComplexFloat data type
+        # https://github.com/pytorch/pytorch/issues/71613
+        self.register_buffer("pre_twiddle", view_as_real(pre_twiddle))
+        self.register_buffer("post_twiddle", view_as_real(post_twiddle))
+
+    def forward(self, audio: torch.Tensor) -> torch.Tensor:
+        """
+        Apply the Modified Discrete Cosine Transform (MDCT) to the input audio.
+
+        Args:
+            audio (Tensor): Input audio waveform of shape (B, T), where B is the batch size
+                and T is the length of the audio.
+
+        Returns:
+            Tensor: MDCT coefficients of shape (B, L, N), where L is the number of output frames
+                and N is the number of frequency bins.
+        """
+        if self.padding == "center":
+            audio = torch.nn.functional.pad(audio, (self.frame_len // 2, self.frame_len // 2))
+        elif self.padding == "same":
+            # hop_length is 1/2 frame_len
+            audio = torch.nn.functional.pad(audio, (self.frame_len // 4, self.frame_len // 4))
+        else:
+            raise ValueError("Padding must be 'center' or 'same'.")
+
+        x = audio.unfold(-1, self.frame_len, self.frame_len // 2)
+        N = self.frame_len // 2
+        x = x * self.window.expand(x.shape)
+        X = torch.fft.fft(x * view_as_complex(self.pre_twiddle).expand(x.shape), dim=-1)[..., :N]
+        res = X * view_as_complex(self.post_twiddle).expand(X.shape) * np.sqrt(1 / N)
+        return torch.real(res) * np.sqrt(2)
+
+
+class IMDCT(nn.Module):
+    """
+    Inverse Modified Discrete Cosine Transform (IMDCT) module.
+
+    Args:
+        frame_len (int): Length of the MDCT frame.
+        padding (str, optional): Type of padding. Options are "center" or "same". Defaults to "same".
+    """
+
+    def __init__(self, frame_len: int, padding: str = "same"):
+        super().__init__()
+        if padding not in ["center", "same"]:
+            raise ValueError("Padding must be 'center' or 'same'.")
+        self.padding = padding
+        self.frame_len = frame_len
+        N = frame_len // 2
+        n0 = (N + 1) / 2
+        window = torch.from_numpy(scipy.signal.cosine(frame_len)).float()
+        self.register_buffer("window", window)
+
+        pre_twiddle = torch.exp(1j * torch.pi * n0 * torch.arange(N * 2) / N)
+        post_twiddle = torch.exp(1j * torch.pi * (torch.arange(N * 2) + n0) / (N * 2))
+        self.register_buffer("pre_twiddle", view_as_real(pre_twiddle))
+        self.register_buffer("post_twiddle", view_as_real(post_twiddle))
+
+    def forward(self, X: torch.Tensor) -> torch.Tensor:
+        """
+        Apply the Inverse Modified Discrete Cosine Transform (IMDCT) to the input MDCT coefficients.
+
+        Args:
+            X (Tensor): Input MDCT coefficients of shape (B, L, N), where B is the batch size,
+                L is the number of frames, and N is the number of frequency bins.
+
+        Returns:
+            Tensor: Reconstructed audio waveform of shape (B, T), where T is the length of the audio.
+        """
+        B, L, N = X.shape
+        Y = torch.zeros((B, L, N * 2), dtype=X.dtype, device=X.device)
+        Y[..., :N] = X
+        Y[..., N:] = -1 * torch.conj(torch.flip(X, dims=(-1,)))
+        y = torch.fft.ifft(Y * view_as_complex(self.pre_twiddle).expand(Y.shape), dim=-1)
+        y = torch.real(y * view_as_complex(self.post_twiddle).expand(y.shape)) * np.sqrt(N) * np.sqrt(2)
+        result = y * self.window.expand(y.shape)
+        output_size = (1, (L + 1) * N)
+        audio = torch.nn.functional.fold(
+            result.transpose(1, 2),
+            output_size=output_size,
+            kernel_size=(1, self.frame_len),
+            stride=(1, self.frame_len // 2),
+        )[:, 0, 0, :]
+
+        if self.padding == "center":
+            pad = self.frame_len // 2
+        elif self.padding == "same":
+            pad = self.frame_len // 4
+        else:
+            raise ValueError("Padding must be 'center' or 'same'.")
+
+        audio = audio[:, pad:-pad]
+        return audio