| """Local PyTorch MERT2 encoder for the unified tokenizer checkpoint.""" |
|
|
| from typing import Dict, Sequence |
|
|
| import torch |
| import torch.nn.functional as F |
| from torch import nn |
| from torchaudio.transforms import AmplitudeToDB, MelScale, Spectrogram |
|
|
| SAMPLE_RATE = 24000 |
| HOP = 960 |
|
|
|
|
| class MelFrontend(nn.Module): |
| def __init__(self, cfg: Dict): |
| super().__init__() |
| bins = cfg["num_mel_bins"] |
| self.register_buffer("mel_mean", torch.zeros(bins)) |
| self.register_buffer("mel_std", torch.ones(bins)) |
| self.spectrogram = Spectrogram( |
| n_fft=cfg["n_fft"], |
| win_length=cfg["win_length"], |
| hop_length=cfg["hop_length"], |
| power=2.0, |
| ) |
| self.mel_scale = MelScale( |
| n_mels=bins, |
| sample_rate=cfg["sampling_rate"], |
| n_stft=cfg["n_fft"] // 2 + 1, |
| ) |
| self.amplitude_to_db = AmplitudeToDB(stype="power", top_db=None) |
|
|
| @torch.no_grad() |
| def forward(self, waveform: torch.Tensor) -> torch.Tensor: |
| with torch.autocast(device_type=waveform.device.type, enabled=False): |
| mel = self.amplitude_to_db(self.mel_scale(self.spectrogram(waveform.float()))) |
| mel = mel[..., :-1].transpose(-1, -2) |
| return (mel - self.mel_mean) / self.mel_std.clamp_min(1e-5) |
|
|
|
|
| class Transpose(nn.Module): |
| def forward(self, x): |
| return x.transpose(1, 2) |
|
|
|
|
| class GlobalResponseNorm(nn.Module): |
| def __init__(self, dim: int): |
| super().__init__() |
| self.weight = nn.Parameter(torch.zeros(1, 1, dim)) |
| self.bias = nn.Parameter(torch.zeros(1, 1, dim)) |
|
|
| def forward(self, x): |
| magnitude = torch.norm(x, p=2, dim=1, keepdim=True) |
| return self.weight * (x * (magnitude / (magnitude.mean(dim=-1, keepdim=True) + 1e-6))) + self.bias + x |
|
|
|
|
| class ConvNextLayer(nn.Module): |
| def __init__(self, dim: int, eps: float): |
| super().__init__() |
| self.depthwise_block = nn.Sequential( |
| Transpose(), |
| nn.Conv1d(dim, dim, 7, padding=3, groups=dim), |
| Transpose(), |
| ) |
| self.pointwise_block = nn.Sequential( |
| nn.LayerNorm(dim, eps=eps), |
| nn.Linear(dim, 4 * dim), |
| nn.GELU(), |
| GlobalResponseNorm(4 * dim), |
| nn.Linear(4 * dim, dim), |
| ) |
|
|
| def forward(self, x): |
| return x + self.pointwise_block(self.depthwise_block(x)) |
|
|
|
|
| class ConvNextBlock(nn.Module): |
| def __init__(self, cin: int, cout: int, stride: int, depth: int, eps: float): |
| super().__init__() |
| self.resampling_layer = ( |
| nn.Sequential( |
| nn.LayerNorm(cin, eps=eps), |
| Transpose(), |
| nn.Conv1d(cin, cout, 2, stride=stride), |
| Transpose(), |
| ) |
| if cin != cout or stride > 1 |
| else nn.Identity() |
| ) |
| self.convnext_layers = nn.Sequential(*[ConvNextLayer(cout, eps) for _ in range(depth)]) |
|
|
| def forward(self, x): |
| return self.convnext_layers(self.resampling_layer(x)) |
|
|
|
|
| class RotaryEmbedding(nn.Module): |
| def __init__(self, cfg: Dict): |
| super().__init__() |
| head_dim = cfg["hidden_size"] // cfg["num_attention_heads"] |
| inv = 1.0 / ( |
| cfg["rotary_embedding_base"] |
| ** (torch.arange(0, head_dim, 2, dtype=torch.float32) / head_dim) |
| ) |
| self.register_buffer("inv_freq", inv, persistent=False) |
|
|
| def forward(self, x): |
| positions = torch.arange(x.shape[1], device=x.device, dtype=torch.float32) |
| angles = torch.einsum("i,j->ij", positions, self.inv_freq.float()) |
| angles = torch.cat((angles, angles), dim=-1)[None, :, None, :] |
| return angles.cos().to(x.dtype), angles.sin().to(x.dtype) |
|
|
|
|
| def rotate_half(x): |
| first, second = x.chunk(2, dim=-1) |
| return torch.cat((-second, first), dim=-1) |
|
|
|
|
| class SelfAttention(nn.Module): |
| def __init__(self, cfg: Dict, causal: bool): |
| super().__init__() |
| width = cfg["hidden_size"] |
| self.num_heads = cfg["num_attention_heads"] |
| self.head_dim = width // self.num_heads |
| self.causal = causal |
| self.query_proj = nn.Linear(width, width) |
| self.key_proj = nn.Linear(width, width) |
| self.value_proj = nn.Linear(width, width) |
| self.out_proj = nn.Linear(width, width) |
|
|
| def forward(self, x, rope): |
| batch, frames, width = x.shape |
| shape = (batch, frames, self.num_heads, self.head_dim) |
| q, k, v = ( |
| projection(x).reshape(shape) |
| for projection in (self.query_proj, self.key_proj, self.value_proj) |
| ) |
| cos, sin = rope |
| q = q * cos + rotate_half(q) * sin |
| k = k * cos + rotate_half(k) * sin |
| out = F.scaled_dot_product_attention( |
| q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2), is_causal=self.causal |
| ) |
| return self.out_proj(out.transpose(1, 2).reshape(batch, frames, width)) |
|
|
|
|
| class FeedForward(nn.Module): |
| def __init__(self, cfg: Dict): |
| super().__init__() |
| self.w_1 = nn.Linear(cfg["hidden_size"], cfg["intermediate_size"]) |
| self.w_2 = nn.Linear(cfg["intermediate_size"], cfg["hidden_size"]) |
|
|
| def forward(self, x): |
| return self.w_2(F.gelu(self.w_1(x))) |
|
|
|
|
| class ConvolutionModule(nn.Module): |
| def __init__(self, cfg: Dict): |
| super().__init__() |
| width = cfg["hidden_size"] |
| kernel = cfg["conv_depthwise_kernel_size"] |
| eps = cfg["layer_norm_eps"] |
| self.layer_norm = nn.LayerNorm(width, eps=eps) |
| self.conv_block = nn.Sequential( |
| Transpose(), |
| nn.Conv1d(width, 2 * width, 1, bias=False), |
| nn.GLU(dim=1), |
| nn.Conv1d( |
| width, |
| width, |
| kernel, |
| padding=(kernel - 1) // 2, |
| groups=width, |
| bias=False, |
| ), |
| nn.Sequential(Transpose(), nn.LayerNorm(width, eps=eps), Transpose()), |
| nn.GELU(), |
| nn.Conv1d(width, width, 1, bias=False), |
| Transpose(), |
| ) |
|
|
| def forward(self, x): |
| return self.conv_block(self.layer_norm(x)) |
|
|
|
|
| class ConformerBlock(nn.Module): |
| def __init__(self, cfg: Dict, causal: bool): |
| super().__init__() |
| width, eps = cfg["hidden_size"], cfg["layer_norm_eps"] |
| self.ffn1_layer_norm = nn.LayerNorm(width, eps=eps) |
| self.ffn1 = FeedForward(cfg) |
| self.attn_layer_norm = nn.LayerNorm(width, eps=eps) |
| self.attn = SelfAttention(cfg, causal) |
| self.conv_module = ConvolutionModule(cfg) |
| self.ffn2_layer_norm = nn.LayerNorm(width, eps=eps) |
| self.ffn2 = FeedForward(cfg) |
| self.final_layer_norm = nn.LayerNorm(width, eps=eps) |
|
|
| def forward(self, x, rope): |
| x = x + 0.5 * self.ffn1(self.ffn1_layer_norm(x)) |
| x = self.attn(self.attn_layer_norm(x), rope) + x |
| x = self.conv_module(x) + x |
| x = x + 0.5 * self.ffn2(self.ffn2_layer_norm(x)) |
| return self.final_layer_norm(x) |
|
|
|
|
| class MERT2Encoder(nn.Module): |
| """Mono 24 kHz waveform to 25 Hz Conformer block outputs.""" |
|
|
| def __init__(self, cfg: Dict, causal: bool, layers: int): |
| super().__init__() |
| self.feature_extractor = MelFrontend(cfg) |
| channels = [cfg["num_mel_bins"]] + list(cfg["subsampling_channels"]) |
| self.subsampling_module = nn.Sequential( |
| *[ |
| ConvNextBlock( |
| channels[i], |
| channels[i + 1], |
| (1, 2, 2)[i], |
| cfg["subsampling_depths"][i], |
| cfg["subsampling_layer_norm_eps"], |
| ) |
| for i in range(3) |
| ] |
| ) |
| self.layers = nn.ModuleList([ConformerBlock(cfg, causal) for _ in range(layers)]) |
| self.embed_positions = RotaryEmbedding(cfg) |
|
|
| def forward(self, waveform: torch.Tensor, output: int) -> torch.Tensor: |
| mel = self.feature_extractor(waveform) |
| first = self.subsampling_module[0].convnext_layers[0].depthwise_block[1] |
| x = self.subsampling_module(mel.to(first.weight.dtype)) |
| rope = self.embed_positions(x) |
| for index, layer in enumerate(self.layers): |
| x = layer(x, rope) |
| if index == output: |
| return x |
| raise ValueError(f"Requested encoder block {output}, but only {len(self.layers)} are present") |
|
|
| @torch.inference_mode() |
| def features(self, mono24: torch.Tensor, output: int, chunk: float = 0) -> torch.Tensor: |
| """Return fp32 [frames, 1024] features, optionally in fixed-second pieces.""" |
| if mono24.ndim != 1 or mono24.shape[0] < 1025: |
| raise ValueError("The encoder needs at least 1025 mono samples at 24 kHz") |
| with torch.autocast( |
| device_type=mono24.device.type, |
| dtype=torch.bfloat16, |
| enabled=mono24.device.type == "cuda", |
| ): |
| if not chunk: |
| return self(mono24[None], output)[0].float() |
| samples = int(SAMPLE_RATE * chunk) |
| pieces = [ |
| piece |
| for piece in (mono24[start : start + samples] for start in range(0, len(mono24), samples)) |
| if len(piece) >= SAMPLE_RATE |
| ] |
| if not pieces: |
| raise ValueError(f"Audio is shorter than the minimum one-second {chunk}-second chunk") |
| full = [piece for piece in pieces if len(piece) == samples] |
| groups = ([torch.stack(full)] if full else []) + [piece[None] for piece in pieces if len(piece) < samples] |
| parts = [self(group, output) for group in groups] |
| frames = int(round(len(mono24) / HOP)) |
| joined = torch.cat([part.reshape(-1, part.shape[-1]) for part in parts]).float() |
| return F.interpolate( |
| joined.T[None], size=frames, mode="linear", align_corners=False |
| )[0].T |
|
|
