upload model and code

config.json

@@ -0,0 +1,39 @@
+{
+  "architectures": [
+    "USADModel"
+  ],
+  "attention_dropout_p": 0.1,
+  "attention_type": "mhsa",
+  "auto_map": {
+    "AutoConfig": "configuration_usad.USADConfig",
+    "AutoModel": "modeling_usad.USADModel"
+  },
+  "conv_dropout_p": 0.1,
+  "conv_expansion_factor": 2,
+  "conv_kernel_size": 31,
+  "conv_pos": true,
+  "conv_pos_depth": 5,
+  "conv_pos_groups": 16,
+  "conv_pos_width": 95,
+  "conv_subsample_channels": 64,
+  "conv_subsample_rate": 2,
+  "encoder_dim": 1024,
+  "feed_forward_dropout_p": 0.1,
+  "feed_forward_expansion_factor": 4,
+  "half_step_residual": true,
+  "input_dim": 128,
+  "input_dropout_p": 0.0,
+  "mamba_bidirectional": false,
+  "mamba_d_conv": 4,
+  "mamba_d_state": 16,
+  "mamba_expand": 2,
+  "model_type": "usad",
+  "num_attention_heads": 16,
+  "num_layers": 24,
+  "subsample_normalization": true,
+  "torch_dtype": "float32",
+  "transformer_style": true,
+  "transformers_version": "4.52.4",
+  "use_framewise_subsample": true,
+  "use_patchwise_subsample": false
+}

configuration_usad.py

@@ -0,0 +1,66 @@
+from transformers import PretrainedConfig
+
+
+class USADConfig(PretrainedConfig):
+    model_type = "usad"
+
+    def __init__(
+        self,
+        encoder_dim: int = 384,
+        num_layers: int = 12,
+        attention_type: str = "mhsa",
+        num_attention_heads: int = 6,
+        mamba_d_state: int = 16,
+        mamba_d_conv: int = 4,
+        mamba_expand: int = 2,
+        mamba_bidirectional: bool = False,
+        feed_forward_expansion_factor: int = 4,
+        conv_expansion_factor: int = 2,
+        feed_forward_dropout_p: float = 0.1,
+        attention_dropout_p: float = 0.1,
+        conv_dropout_p: float = 0.1,
+        conv_kernel_size: int = 31,
+        half_step_residual: bool = True,
+        transformer_style: bool = True,
+        use_framewise_subsample: bool = True,
+        use_patchwise_subsample: bool = False,
+        conv_subsample_channels: int = 64,
+        conv_subsample_rate: int = 2,
+        input_dim: int = 128,
+        input_dropout_p: float = 0.0,
+        conv_pos: bool = True,
+        conv_pos_depth: int = 5,
+        conv_pos_width: int = 95,
+        conv_pos_groups: int = 16,
+        subsample_normalization: bool = True,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        self.encoder_dim = encoder_dim
+        self.num_layers = num_layers
+        self.attention_type = attention_type
+        self.num_attention_heads = num_attention_heads
+        self.mamba_d_state = mamba_d_state
+        self.mamba_d_conv = mamba_d_conv
+        self.mamba_expand = mamba_expand
+        self.mamba_bidirectional = mamba_bidirectional
+        self.feed_forward_expansion_factor = feed_forward_expansion_factor
+        self.conv_expansion_factor = conv_expansion_factor
+        self.feed_forward_dropout_p = feed_forward_dropout_p
+        self.attention_dropout_p = attention_dropout_p
+        self.conv_dropout_p = conv_dropout_p
+        self.conv_kernel_size = conv_kernel_size
+        self.half_step_residual = half_step_residual
+        self.transformer_style = transformer_style
+        self.use_framewise_subsample = use_framewise_subsample
+        self.use_patchwise_subsample = use_patchwise_subsample
+        self.conv_subsample_channels = conv_subsample_channels
+        self.conv_subsample_rate = conv_subsample_rate
+        self.input_dim = input_dim
+        self.input_dropout_p = input_dropout_p
+        self.conv_pos = conv_pos
+        self.conv_pos_depth = conv_pos_depth
+        self.conv_pos_width = conv_pos_width
+        self.conv_pos_groups = conv_pos_groups
+        self.subsample_normalization = subsample_normalization

model.safetensors

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

modeling_usad.py

@@ -0,0 +1,19 @@
+# modeling_usad.py
+
+from transformers import PreTrainedModel
+from .configuration_usad import USADConfig
+from .usad_model import UsadModel as model
+
+
+class USADModel(PreTrainedModel):
+    config_class = USADConfig
+
+    def __init__(self, config: USADConfig):
+        super().__init__(config)
+        self.model = model(config)
+
+    def forward(self, *args, **kwargs):
+        return self.model(*args, **kwargs)
+
+    def load_audio(self, audio_path):
+        return self.model.load_audio(audio_path)

usad_model.py

@@ -0,0 +1,207 @@
+from dataclasses import make_dataclass
+
+import torch
+import torchaudio
+from torch import nn
+
+from .usad_modules import ConformerEncoder
+
+MAX_MEL_LENGTH = 3000  # 30 seconds
+
+
+@torch.no_grad()
+def wav_to_fbank(
+    wavs: torch.Tensor,
+    mel_dim: int = 128,
+    norm_mean: float = -4.268,
+    norm_std: float = 4.569,
+) -> torch.Tensor:
+    """Convert waveform to fbank features.
+
+    Args:
+        wavs (torch.Tensor): (B, T_wav) waveform tensor.
+        mel_dim (int, optional): mel dimension. Defaults to 128.
+        norm_mean (float, optional):
+            mean for normalization. Defaults to -4.268.
+        norm_std (float, optional):
+            std for normalization. Defaults to 4.569.
+
+    Returns:
+        torch.Tensor: (B, T_mel, mel_dim) fbank features.
+    """
+    # ref: https://github.com/cwx-worst-one/EAT/tree/main/feature_extract
+    dtype = wavs.dtype
+    wavs = wavs.to(torch.float32)
+    wavs = wavs - wavs.mean(dim=-1, keepdim=True)
+    feats = [
+        torchaudio.compliance.kaldi.fbank(
+            wavs[i : i + 1],
+            htk_compat=True,
+            sample_frequency=16000,
+            use_energy=False,
+            window_type="hanning",
+            num_mel_bins=mel_dim,
+            dither=0.0,
+            frame_shift=10,
+        ).to(dtype=dtype)
+        for i in range(wavs.shape[0])
+    ]
+
+    mels = torch.stack(feats, dim=0)
+    mels = (mels - norm_mean) / (norm_std * 2)
+
+    return mels
+
+
+class UsadModel(nn.Module):
+    def __init__(self, cfg) -> None:
+        """Initialize the UsadModel.
+        Args:
+            cfg: Configuration object containing model parameters.
+        """
+        super().__init__()
+
+        self.cfg = cfg
+        self.encoder = ConformerEncoder(cfg)
+        self.max_mel_length = MAX_MEL_LENGTH
+        # NOTE: The max_mel_length is set to 3000,
+        # which corresponds to 30 seconds of audio at 100 Hz frame rate.
+
+    @property
+    def sample_rate(self) -> int:
+        return 16000  # Hz
+
+    @property
+    def encoder_frame_rate(self) -> int:
+        return 50  # Hz
+
+    @property
+    def mel_dim(self) -> int:
+        return self.cfg.input_dim
+
+    @property
+    def encoder_dim(self) -> int:
+        return self.cfg.encoder_dim
+
+    @property
+    def num_layers(self) -> int:
+        return self.cfg.num_layers
+
+    @property
+    def scene_embedding_size(self) -> int:
+        return self.cfg.encoder_dim * self.cfg.num_layers
+
+    @property
+    def timestamp_embedding_size(self) -> int:
+        return self.cfg.encoder_dim * self.cfg.num_layers
+
+    @property
+    def device(self) -> torch.device:
+        """Get the device on which the model is located."""
+        return next(self.parameters()).device
+
+    def set_audio_chunk_size(self, seconds: float = 30.0) -> None:
+        """Set the maximum chunk size for feature extraction.
+
+        Args:
+            seconds (float, optional): Chunk size in seconds. Defaults to 30.0.
+        """
+        assert (
+            seconds >= 0.1
+        ), f"Chunk size must be greater than 0.1s, got {seconds} seconds."
+        self.max_mel_length = int(seconds * 100)  # 100 Hz frame rate
+
+    def load_audio(self, audio_path: str) -> torch.Tensor:
+        """Load audio file and return waveform tensor.
+        Args:
+            audio_path (str): Path to the audio file.
+
+        Returns:
+            torch.Tensor: Waveform tensor of shape (wav_len,).
+        """
+
+        waveform, sr = torchaudio.load(audio_path)
+        if sr != self.sample_rate:
+            waveform = torchaudio.functional.resample(waveform, sr, self.sample_rate)
+        if waveform.shape[0] > 1:
+            # If stereo, convert to mono by averaging channels
+            waveform = waveform.mean(dim=0, keepdim=True)
+
+        waveform = waveform.squeeze(0)  # Remove channel dimension if mono
+        return waveform.to(self.device)  # Ensure tensor is on the same device
+
+    def forward(
+        self,
+        wavs: torch.Tensor,
+        norm_mean: float = -4.268,
+        norm_std: float = 4.569,
+    ) -> dict:
+        """Forward pass for the model.
+
+        Args:
+            wavs (torch.Tensor):
+                Input waveform tensor of shape (batch_size, wav_len).
+            norm_mean (float, optional):
+                Mean for normalization. Defaults to -4.268.
+            norm_std (float, optional):
+                Standard deviation for normalization. Defaults to 4.569.
+
+        Returns:
+            dict: A dictionary containing the model's outputs.
+        """
+        # wavs: (batch_size, wav_len)
+
+        mel = wav_to_fbank(wavs, norm_mean=norm_mean, norm_std=norm_std)
+        mel = mel[:, : mel.shape[1] - mel.shape[1] % 2]
+        if mel.shape[1] <= self.max_mel_length:
+            x, x_len, layer_results = self.encoder(mel, return_hidden=True)
+
+            result = {
+                "x": x,
+                "mel": mel,
+                "hidden_states": layer_results["hidden_states"],
+                "ffn": layer_results["ffn_1"],
+            }
+            return result
+
+        result = {
+            "x": [],
+            "mel": mel,
+            "hidden_states": [[] for _ in range(self.cfg.num_layers)],
+            "ffn": [[] for _ in range(self.cfg.num_layers)],
+        }
+        for i in range(0, mel.shape[1], self.max_mel_length):
+            if mel.shape[1] - i < 10:
+                break
+
+            x, x_len, layer_results = self.encoder(
+                mel[:, i : i + self.max_mel_length], return_hidden=True
+            )
+            result["x"].append(x)
+            for j in range(self.cfg.num_layers):
+                result["hidden_states"][j].append(layer_results["hidden_states"][j])
+                result["ffn"][j].append(layer_results["ffn_1"][j])
+
+        result["x"] = torch.cat(result["x"], dim=1)
+        for j in range(self.cfg.num_layers):
+            result["hidden_states"][j] = torch.cat(result["hidden_states"][j], dim=1)
+            result["ffn"][j] = torch.cat(result["ffn"][j], dim=1)
+
+        # result["x"]: model final output (batch_size, seq_len)
+        # result["mel"]: mel fbank (batch_size, seq_len * 2, mel_dim)
+        # result["hidden_states"]: List of (batch_size, seq_len, encoder_dim)
+        # result["ffn"]: List of (batch_size, seq_len, encoder_dim)
+        return result
+
+    @classmethod
+    def load_from_fairseq_ckpt(cls, ckpt_path: str):
+        checkpoint = torch.load(ckpt_path, weights_only=False)
+        config = checkpoint["cfg"]["model"]
+        config = make_dataclass("Config", config.keys())(**config)
+        model = cls(config)
+        state_dict = checkpoint["model"]
+        for k in list(state_dict.keys()):
+            if not k.startswith("encoder."):
+                del state_dict[k]
+        model.load_state_dict(state_dict, strict=True)
+        return model

usad_modules.py

@@ -0,0 +1,764 @@
+# Copyright (c) 2021, Soohwan Kim. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import contextlib
+import math
+from collections import defaultdict
+from typing import Dict, List, Optional, Tuple, Union
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+
+class SamePad(nn.Module):
+    def __init__(self, kernel_size, causal=False):
+        super().__init__()
+        if causal:
+            self.remove = kernel_size - 1
+        else:
+            self.remove = 1 if kernel_size % 2 == 0 else 0
+
+    def forward(self, x):
+        if self.remove > 0:
+            x = x[:, :, : -self.remove]
+        return x
+
+
+class TransposeLast(nn.Module):
+    def __init__(self, deconstruct_idx=None, tranpose_dim=-2):
+        super().__init__()
+        self.deconstruct_idx = deconstruct_idx
+        self.tranpose_dim = tranpose_dim
+
+    def forward(self, x):
+        if self.deconstruct_idx is not None:
+            x = x[self.deconstruct_idx]
+        return x.transpose(self.tranpose_dim, -1)
+
+
+class Swish(nn.Module):
+    def __init__(self):
+        super(Swish, self).__init__()
+
+    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
+        return inputs * inputs.sigmoid()
+
+
+class GLU(nn.Module):
+    def __init__(self, dim: int) -> None:
+        super(GLU, self).__init__()
+        self.dim = dim
+
+    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
+        outputs, gate = inputs.chunk(2, dim=self.dim)
+        return outputs * gate.sigmoid()
+
+
+class ResidualConnectionModule(nn.Module):
+    def __init__(
+        self,
+        module: nn.Module,
+        module_factor: float = 1.0,
+        input_factor: float = 1.0,
+    ):
+        super(ResidualConnectionModule, self).__init__()
+        self.module = module
+        self.module_factor = module_factor
+        self.input_factor = input_factor
+
+    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
+        return (self.module(inputs) * self.module_factor) + (inputs * self.input_factor)
+
+
+class Linear(nn.Module):
+    def __init__(self, in_features: int, out_features: int, bias: bool = True) -> None:
+        super(Linear, self).__init__()
+        self.linear = nn.Linear(in_features, out_features, bias=bias)
+        nn.init.xavier_uniform_(self.linear.weight)
+        if bias:
+            nn.init.zeros_(self.linear.bias)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.linear(x)
+
+
+class View(nn.Module):
+    def __init__(self, shape: tuple, contiguous: bool = False):
+        super(View, self).__init__()
+        self.shape = shape
+        self.contiguous = contiguous
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self.contiguous:
+            x = x.contiguous()
+
+        return x.view(*self.shape)
+
+
+class Transpose(nn.Module):
+    def __init__(self, shape: tuple):
+        super(Transpose, self).__init__()
+        self.shape = shape
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return x.transpose(*self.shape)
+
+
+class FeedForwardModule(nn.Module):
+    def __init__(
+        self,
+        encoder_dim: int = 512,
+        expansion_factor: int = 4,
+        dropout_p: float = 0.1,
+    ) -> None:
+        super(FeedForwardModule, self).__init__()
+        self.sequential = nn.Sequential(
+            nn.LayerNorm(encoder_dim),
+            Linear(encoder_dim, encoder_dim * expansion_factor, bias=True),
+            Swish(),
+            nn.Dropout(p=dropout_p),
+            Linear(encoder_dim * expansion_factor, encoder_dim, bias=True),
+            nn.Dropout(p=dropout_p),
+        )
+
+    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
+        return self.sequential(inputs)
+
+
+class DepthwiseConv1d(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: int,
+        stride: int = 1,
+        padding: int = 0,
+        bias: bool = False,
+    ) -> None:
+        super(DepthwiseConv1d, self).__init__()
+        assert (
+            out_channels % in_channels == 0
+        ), "out_channels should be constant multiple of in_channels"
+        self.conv = nn.Conv1d(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            kernel_size=kernel_size,
+            groups=in_channels,
+            stride=stride,
+            padding=padding,
+            bias=bias,
+        )
+
+    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
+        return self.conv(inputs)
+
+
+class PointwiseConv1d(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        stride: int = 1,
+        padding: int = 0,
+        bias: bool = True,
+    ) -> None:
+        super(PointwiseConv1d, self).__init__()
+        self.conv = nn.Conv1d(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            kernel_size=1,
+            stride=stride,
+            padding=padding,
+            bias=bias,
+        )
+
+    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
+        return self.conv(inputs)
+
+
+class ConformerConvModule(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        kernel_size: int = 31,
+        expansion_factor: int = 2,
+        dropout_p: float = 0.1,
+    ) -> None:
+        super(ConformerConvModule, self).__init__()
+        assert (
+            kernel_size - 1
+        ) % 2 == 0, "kernel_size should be a odd number for 'SAME' padding"
+        assert expansion_factor == 2, "Currently, Only Supports expansion_factor 2"
+
+        self.sequential = nn.Sequential(
+            nn.LayerNorm(in_channels),
+            Transpose(shape=(1, 2)),
+            PointwiseConv1d(
+                in_channels,
+                in_channels * expansion_factor,
+                stride=1,
+                padding=0,
+                bias=True,
+            ),
+            GLU(dim=1),
+            DepthwiseConv1d(
+                in_channels,
+                in_channels,
+                kernel_size,
+                stride=1,
+                padding=(kernel_size - 1) // 2,
+            ),
+            nn.BatchNorm1d(in_channels),
+            Swish(),
+            PointwiseConv1d(in_channels, in_channels, stride=1, padding=0, bias=True),
+            nn.Dropout(p=dropout_p),
+        )
+
+    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
+        return self.sequential(inputs).transpose(1, 2)
+
+
+class FramewiseConv2dSubampling(nn.Module):
+    def __init__(self, out_channels: int, subsample_rate: int = 2) -> None:
+        super(FramewiseConv2dSubampling, self).__init__()
+        assert subsample_rate in {2, 4}, "subsample_rate should be 2 or 4"
+        self.subsample_rate = subsample_rate
+        self.cnn = nn.Sequential(
+            nn.Conv2d(1, out_channels, kernel_size=3, stride=2),
+            nn.ReLU(),
+            nn.Conv2d(
+                out_channels,
+                out_channels,
+                kernel_size=3,
+                stride=(2 if subsample_rate == 4 else 1, 2),
+                padding=(0 if subsample_rate == 4 else 1, 0),
+            ),
+            nn.ReLU(),
+        )
+
+    def forward(
+        self, inputs: torch.Tensor, input_lengths: torch.LongTensor
+    ) -> Tuple[torch.Tensor, torch.LongTensor]:
+        # inputs: (B, T, C) -> (B, 1, T, C)
+        if self.subsample_rate == 2 and inputs.shape[1] % 2 == 0:
+            inputs = F.pad(inputs, (0, 0, 0, 1), "constant", 0)
+        outputs = self.cnn(inputs.unsqueeze(1))
+        batch_size, channels, subsampled_lengths, sumsampled_dim = outputs.size()
+
+        outputs = outputs.permute(0, 2, 1, 3)
+        outputs = outputs.contiguous().view(
+            batch_size, subsampled_lengths, channels * sumsampled_dim
+        )
+
+        if self.subsample_rate == 4:
+            output_lengths = (((input_lengths - 1) >> 1) - 1) >> 1
+        else:
+            output_lengths = input_lengths >> 1
+
+        return outputs, output_lengths
+
+
+class PatchwiseConv2dSubampling(nn.Module):
+    def __init__(
+        self,
+        mel_dim: int,
+        out_channels: int,
+        patch_size_time: int = 16,
+        patch_size_freq: int = 16,
+    ) -> None:
+        super(PatchwiseConv2dSubampling, self).__init__()
+
+        self.mel_dim = mel_dim
+        self.patch_size_time = patch_size_time
+        self.patch_size_freq = patch_size_freq
+
+        self.proj = nn.Conv2d(
+            1,
+            out_channels,
+            kernel_size=(patch_size_time, patch_size_freq),
+            stride=(patch_size_time, patch_size_freq),
+            padding=0,
+        )
+        self.cnn = nn.Sequential(
+            nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1),
+            nn.ReLU(),
+            nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1),
+            nn.ReLU(),
+        )
+
+    @property
+    def subsample_rate(self) -> int:
+        return self.patch_size_time * self.patch_size_freq // self.mel_dim
+
+    def forward(
+        self, inputs: torch.Tensor, input_lengths: torch.LongTensor
+    ) -> Tuple[torch.Tensor, torch.LongTensor]:
+        assert (
+            inputs.shape[2] == self.mel_dim
+        ), "inputs.shape[2] should be equal to mel_dim"
+
+        # inputs: (B, Time, Freq) -> (B, 1, Time, Freq)
+        outputs = self.proj(inputs.unsqueeze(1))
+        outputs = self.cnn(outputs)
+        # (B, channels, Time // patch_size_time, Freq // patch_size_freq)
+        outputs = outputs.flatten(2, 3).transpose(1, 2)
+        # (B, (Time // patch_size_time) * (Freq // patch_size_freq), channels)
+
+        output_lengths = (
+            input_lengths
+            // self.patch_size_time
+            * (self.mel_dim // self.patch_size_freq)
+        )
+
+        return outputs, output_lengths
+
+
+class RelPositionalEncoding(nn.Module):
+    def __init__(self, d_model: int, max_len: int = 10000) -> None:
+        super(RelPositionalEncoding, self).__init__()
+        self.d_model = d_model
+        self.pe = None
+        self.extend_pe(torch.tensor(0.0).expand(1, max_len))
+
+    def extend_pe(self, x: torch.Tensor) -> None:
+        if self.pe is not None:
+            if self.pe.size(1) >= x.size(1) * 2 - 1:
+                if self.pe.dtype != x.dtype or self.pe.device != x.device:
+                    self.pe = self.pe.to(dtype=x.dtype, device=x.device)
+                return
+
+        pe_positive = torch.zeros(x.size(1), self.d_model)
+        pe_negative = torch.zeros(x.size(1), self.d_model)
+        position = torch.arange(0, x.size(1), dtype=torch.float32).unsqueeze(1)
+        div_term = torch.exp(
+            torch.arange(0, self.d_model, 2, dtype=torch.float32)
+            * -(math.log(10000.0) / self.d_model)
+        )
+        pe_positive[:, 0::2] = torch.sin(position * div_term)
+        pe_positive[:, 1::2] = torch.cos(position * div_term)
+        pe_negative[:, 0::2] = torch.sin(-1 * position * div_term)
+        pe_negative[:, 1::2] = torch.cos(-1 * position * div_term)
+
+        pe_positive = torch.flip(pe_positive, [0]).unsqueeze(0)
+        pe_negative = pe_negative[1:].unsqueeze(0)
+        pe = torch.cat([pe_positive, pe_negative], dim=1)
+        self.pe = pe.to(device=x.device, dtype=x.dtype)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # x: (B, T, C)
+        self.extend_pe(x)
+        pos_emb = self.pe[
+            :,
+            self.pe.size(1) // 2 - x.size(1) + 1 : self.pe.size(1) // 2 + x.size(1),
+        ]
+        return pos_emb
+
+
+class RelativeMultiHeadAttention(nn.Module):
+    def __init__(
+        self,
+        d_model: int = 512,
+        num_heads: int = 16,
+        dropout_p: float = 0.1,
+    ):
+        super(RelativeMultiHeadAttention, self).__init__()
+        assert d_model % num_heads == 0, "d_model % num_heads should be zero."
+        self.d_model = d_model
+        self.d_head = int(d_model / num_heads)
+        self.num_heads = num_heads
+        self.sqrt_dim = math.sqrt(self.d_head)
+
+        self.query_proj = Linear(d_model, d_model)
+        self.key_proj = Linear(d_model, d_model)
+        self.value_proj = Linear(d_model, d_model)
+        self.pos_proj = Linear(d_model, d_model, bias=False)
+
+        self.dropout = nn.Dropout(p=dropout_p)
+        self.u_bias = nn.Parameter(torch.Tensor(self.num_heads, self.d_head))
+        self.v_bias = nn.Parameter(torch.Tensor(self.num_heads, self.d_head))
+        torch.nn.init.xavier_uniform_(self.u_bias)
+        torch.nn.init.xavier_uniform_(self.v_bias)
+
+        self.out_proj = Linear(d_model, d_model)
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        pos_embedding: torch.Tensor,
+        mask: Optional[torch.Tensor] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        batch_size = value.size(0)
+
+        query = self.query_proj(query).view(batch_size, -1, self.num_heads, self.d_head)
+        key = (
+            self.key_proj(key)
+            .view(batch_size, -1, self.num_heads, self.d_head)
+            .permute(0, 2, 1, 3)
+        )
+        value = (
+            self.value_proj(value)
+            .view(batch_size, -1, self.num_heads, self.d_head)
+            .permute(0, 2, 1, 3)
+        )
+        pos_embedding = self.pos_proj(pos_embedding).view(
+            batch_size, -1, self.num_heads, self.d_head
+        )
+
+        content_score = torch.matmul(
+            (query + self.u_bias).transpose(1, 2), key.transpose(2, 3)
+        )
+        pos_score = torch.matmul(
+            (query + self.v_bias).transpose(1, 2),
+            pos_embedding.permute(0, 2, 3, 1),
+        )
+        pos_score = self._relative_shift(pos_score)
+
+        score = (content_score + pos_score) / self.sqrt_dim
+
+        if mask is not None:
+            mask = mask.unsqueeze(1)
+            score.masked_fill_(mask, -1e9)
+
+        attn = F.softmax(score, -1)
+        attn = self.dropout(attn)
+
+        context = torch.matmul(attn, value).transpose(1, 2)
+        context = context.contiguous().view(batch_size, -1, self.d_model)
+
+        return self.out_proj(context), attn
+
+    def _relative_shift(self, pos_score: torch.Tensor) -> torch.Tensor:
+        batch_size, num_heads, seq_length1, seq_length2 = pos_score.size()
+        zeros = pos_score.new_zeros(batch_size, num_heads, seq_length1, 1)
+        padded_pos_score = torch.cat([zeros, pos_score], dim=-1)
+
+        padded_pos_score = padded_pos_score.view(
+            batch_size, num_heads, seq_length2 + 1, seq_length1
+        )
+        pos_score = padded_pos_score[:, :, 1:].view_as(pos_score)[
+            :, :, :, : seq_length2 // 2 + 1
+        ]
+
+        return pos_score
+
+
+class MultiHeadedSelfAttentionModule(nn.Module):
+    def __init__(self, d_model: int, num_heads: int, dropout_p: float = 0.1):
+        super(MultiHeadedSelfAttentionModule, self).__init__()
+        self.positional_encoding = RelPositionalEncoding(d_model)
+        self.layer_norm = nn.LayerNorm(d_model)
+        self.attention = RelativeMultiHeadAttention(d_model, num_heads, dropout_p)
+        self.dropout = nn.Dropout(p=dropout_p)
+
+    def forward(
+        self, inputs: torch.Tensor, mask: Optional[torch.Tensor] = None
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        batch_size = inputs.size(0)
+        pos_embedding = self.positional_encoding(inputs)
+        pos_embedding = pos_embedding.repeat(batch_size, 1, 1)
+
+        inputs = self.layer_norm(inputs)
+        outputs, attn = self.attention(
+            inputs, inputs, inputs, pos_embedding=pos_embedding, mask=mask
+        )
+
+        return self.dropout(outputs), attn
+
+
+class ConformerBlock(nn.Module):
+    def __init__(
+        self,
+        encoder_dim: int = 512,
+        attention_type: str = "mhsa",
+        num_attention_heads: int = 8,
+        mamba_d_state: int = 16,
+        mamba_d_conv: int = 4,
+        mamba_expand: int = 2,
+        mamba_bidirectional: bool = True,
+        feed_forward_expansion_factor: int = 4,
+        conv_expansion_factor: int = 2,
+        feed_forward_dropout_p: float = 0.1,
+        attention_dropout_p: float = 0.1,
+        conv_dropout_p: float = 0.1,
+        conv_kernel_size: int = 31,
+        half_step_residual: bool = True,
+        transformer_style: bool = False,
+    ):
+        super(ConformerBlock, self).__init__()
+
+        self.transformer_style = transformer_style
+        self.attention_type = attention_type
+
+        if half_step_residual and not transformer_style:
+            self.feed_forward_residual_factor = 0.5
+        else:
+            self.feed_forward_residual_factor = 1
+
+        assert attention_type in ["mhsa", "mamba"]
+        if attention_type == "mhsa":
+            attention = MultiHeadedSelfAttentionModule(
+                d_model=encoder_dim,
+                num_heads=num_attention_heads,
+                dropout_p=attention_dropout_p,
+            )
+
+        self.ffn_1 = FeedForwardModule(
+            encoder_dim=encoder_dim,
+            expansion_factor=feed_forward_expansion_factor,
+            dropout_p=feed_forward_dropout_p,
+        )
+        self.attention = attention
+        if not transformer_style:
+            self.conv = ConformerConvModule(
+                in_channels=encoder_dim,
+                kernel_size=conv_kernel_size,
+                expansion_factor=conv_expansion_factor,
+                dropout_p=conv_dropout_p,
+            )
+            self.ffn_2 = FeedForwardModule(
+                encoder_dim=encoder_dim,
+                expansion_factor=feed_forward_expansion_factor,
+                dropout_p=feed_forward_dropout_p,
+            )
+        self.layernorm = nn.LayerNorm(encoder_dim)
+
+    def forward(
+        self, x: torch.Tensor
+    ) -> Tuple[torch.Tensor, Dict[str, Union[torch.Tensor, None]]]:
+        # FFN 1
+        ffn_1_out = self.ffn_1(x)
+        x = ffn_1_out * self.feed_forward_residual_factor + x
+
+        # Attention
+        if not isinstance(self.attention, MultiHeadedSelfAttentionModule):
+            # MAMBA
+            attn_out = self.attention(x)
+            attn = None
+        else:
+            attn_out, attn = self.attention(x)
+        x = attn_out + x
+
+        if self.transformer_style:
+            x = self.layernorm(x)
+            return x, {
+                "ffn_1": ffn_1_out,
+                "attn": attn,
+                "conv": None,
+                "ffn_2": None,
+            }
+
+        # Convolution
+        conv_out = self.conv(x)
+        x = conv_out + x
+
+        # FFN 2
+        ffn_2_out = self.ffn_2(x)
+        x = ffn_2_out * self.feed_forward_residual_factor + x
+        x = self.layernorm(x)
+
+        other = {
+            "ffn_1": ffn_1_out,
+            "attn": attn,
+            "conv": conv_out,
+            "ffn_2": ffn_2_out,
+        }
+
+        return x, other
+
+
+class ConformerEncoder(nn.Module):
+    def __init__(self, cfg):
+        super(ConformerEncoder, self).__init__()
+
+        self.cfg = cfg
+        self.framewise_subsample = None
+        self.patchwise_subsample = None
+        self.framewise_in_proj = None
+        self.patchwise_in_proj = None
+        assert (
+            cfg.use_framewise_subsample or cfg.use_patchwise_subsample
+        ), "At least one subsampling method should be used"
+        if cfg.use_framewise_subsample:
+            self.framewise_subsample = FramewiseConv2dSubampling(
+                out_channels=cfg.conv_subsample_channels,
+                subsample_rate=cfg.conv_subsample_rate,
+            )
+            self.framewise_in_proj = nn.Sequential(
+                Linear(
+                    cfg.conv_subsample_channels * (((cfg.input_dim - 1) // 2 - 1) // 2),
+                    cfg.encoder_dim,
+                ),
+                nn.Dropout(p=cfg.input_dropout_p),
+            )
+        if cfg.use_patchwise_subsample:
+            self.patchwise_subsample = PatchwiseConv2dSubampling(
+                mel_dim=cfg.input_dim,
+                out_channels=cfg.conv_subsample_channels,
+                patch_size_time=cfg.patch_size_time,
+                patch_size_freq=cfg.patch_size_freq,
+            )
+            self.patchwise_in_proj = nn.Sequential(
+                Linear(
+                    cfg.conv_subsample_channels,
+                    cfg.encoder_dim,
+                ),
+                nn.Dropout(p=cfg.input_dropout_p),
+            )
+            assert not cfg.use_framewise_subsample or (
+                cfg.conv_subsample_rate == self.patchwise_subsample.subsample_rate
+            ), (
+                f"conv_subsample_rate ({cfg.conv_subsample_rate}) != patchwise_subsample.subsample_rate"
+                f"({self.patchwise_subsample.subsample_rate})"
+            )
+
+        self.framewise_norm, self.patchwise_norm = None, None
+        if getattr(cfg, "subsample_normalization", False):
+            if cfg.use_framewise_subsample:
+                self.framewise_norm = nn.LayerNorm(cfg.encoder_dim)
+            if cfg.use_patchwise_subsample:
+                self.patchwise_norm = nn.LayerNorm(cfg.encoder_dim)
+
+        self.conv_pos = None
+        if getattr(cfg, "conv_pos", False):
+            num_pos_layers = cfg.conv_pos_depth
+            k = max(3, cfg.conv_pos_width // num_pos_layers)
+            self.conv_pos = nn.Sequential(
+                TransposeLast(),
+                *[
+                    nn.Sequential(
+                        nn.Conv1d(
+                            cfg.encoder_dim,
+                            cfg.encoder_dim,
+                            kernel_size=k,
+                            padding=k // 2,
+                            groups=cfg.conv_pos_groups,
+                        ),
+                        SamePad(k),
+                        TransposeLast(),
+                        nn.LayerNorm(cfg.encoder_dim, elementwise_affine=False),
+                        TransposeLast(),
+                        nn.GELU(),
+                    )
+                    for _ in range(num_pos_layers)
+                ],
+                TransposeLast(),
+            )
+            self.conv_pos_post_ln = nn.LayerNorm(cfg.encoder_dim)
+
+        self.layers = nn.ModuleList(
+            [
+                ConformerBlock(
+                    encoder_dim=cfg.encoder_dim,
+                    attention_type=cfg.attention_type,
+                    num_attention_heads=cfg.num_attention_heads,
+                    mamba_d_state=cfg.mamba_d_state,
+                    mamba_d_conv=cfg.mamba_d_conv,
+                    mamba_expand=cfg.mamba_expand,
+                    mamba_bidirectional=cfg.mamba_bidirectional,
+                    feed_forward_expansion_factor=cfg.feed_forward_expansion_factor,
+                    conv_expansion_factor=cfg.conv_expansion_factor,
+                    feed_forward_dropout_p=cfg.feed_forward_dropout_p,
+                    attention_dropout_p=cfg.attention_dropout_p,
+                    conv_dropout_p=cfg.conv_dropout_p,
+                    conv_kernel_size=cfg.conv_kernel_size,
+                    half_step_residual=cfg.half_step_residual,
+                    transformer_style=getattr(cfg, "transformer_style", False),
+                )
+                for _ in range(cfg.num_layers)
+            ]
+        )
+
+    def count_parameters(self) -> int:
+        """Count parameters of encoder"""
+        return sum([p.numel() for p in self.parameters() if p.requires_grad])
+
+    def update_dropout(self, dropout_p: float) -> None:
+        """Update dropout probability of encoder"""
+        for name, child in self.named_children():
+            if isinstance(child, nn.Dropout):
+                child.p = dropout_p
+
+    def forward(
+        self,
+        inputs: torch.Tensor,
+        input_lengths: Optional[torch.Tensor] = None,
+        return_hidden: bool = False,
+        freeze_input_layers: bool = False,
+        target_layer: Optional[int] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor, Dict[str, List[torch.Tensor]]]:
+        if input_lengths is None:
+            input_lengths = torch.full(
+                (inputs.size(0),),
+                inputs.size(1),
+                dtype=torch.long,
+                device=inputs.device,
+            )
+
+        with torch.no_grad() if freeze_input_layers else contextlib.ExitStack():
+            frame_feat, patch_feat = None, None
+            if self.framewise_subsample is not None:
+                frame_feat, frame_lengths = self.framewise_subsample(
+                    inputs, input_lengths
+                )
+                frame_feat = self.framewise_in_proj(frame_feat)
+                if self.framewise_norm is not None:
+                    frame_feat = self.framewise_norm(frame_feat)
+
+            if self.patchwise_subsample is not None:
+                patch_feat, patch_lengths = self.patchwise_subsample(
+                    inputs, input_lengths
+                )
+                patch_feat = self.patchwise_in_proj(patch_feat)
+                if self.patchwise_norm is not None:
+                    patch_feat = self.patchwise_norm(patch_feat)
+
+            if frame_feat is not None and patch_feat is not None:
+                min_len = min(frame_feat.size(1), patch_feat.size(1))
+                frame_feat = frame_feat[:, :min_len]
+                patch_feat = patch_feat[:, :min_len]
+
+                features = frame_feat + patch_feat
+                output_lengths = (
+                    frame_lengths
+                    if frame_lengths.max().item() < patch_lengths.max().item()
+                    else patch_lengths
+                )
+            elif frame_feat is not None:
+                features = frame_feat
+                output_lengths = frame_lengths
+            else:
+                features = patch_feat
+                output_lengths = patch_lengths
+
+        if self.conv_pos is not None:
+            features = features + self.conv_pos(features)
+            features = self.conv_pos_post_ln(features)
+
+        layer_results = defaultdict(list)
+
+        outputs = features
+        for i, layer in enumerate(self.layers):
+            outputs, other = layer(outputs)
+            if return_hidden:
+                layer_results["hidden_states"].append(outputs)
+                for k, v in other.items():
+                    layer_results[k].append(v)
+
+            if target_layer is not None and i == target_layer:
+                break
+
+        return outputs, output_lengths, layer_results