| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| import numpy as np |
| from typing import Tuple, List, Optional |
|
|
| class RMSNorm(nn.Module): |
| def __init__(self, dimension: int, eps: float = 1e-5): |
| super().__init__() |
| self.weight = nn.Parameter(torch.ones(dimension)) |
| self.eps = eps |
|
|
| def forward(self, input: torch.Tensor) -> torch.Tensor: |
| input_float = input.half() |
| variance = input_float.pow(2).mean(dim=1, keepdim=True) |
| input_norm = input_float * torch.rsqrt(variance + self.eps) |
| return (input_norm * self.weight.unsqueeze(0).unsqueeze(-1)).type_as(input) |
|
|
|
|
| class RotaryEmbedding(nn.Module): |
| def __init__(self, dim: int, max_position_embeddings: int = 2048, base: int = 10000): |
| super().__init__() |
| self.dim = dim |
| self.max_position_embeddings = max_position_embeddings |
| self.base = base |
| |
| inv_freq = 1. / (self.base ** (torch.arange(0, self.dim, 2).half() / self.dim)) |
| self.register_buffer('inv_freq', inv_freq) |
| |
| self._set_cos_sin_cache( |
| seq_len=max_position_embeddings, |
| device=self.inv_freq.device, |
| dtype=torch.get_default_dtype() |
| ) |
| |
| def _set_cos_sin_cache(self, seq_len: int, device: torch.device, dtype: torch.dtype): |
| self.max_seq_len_cached = seq_len |
| t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype) |
| |
| freqs = torch.einsum('i,j->ij', t, self.inv_freq) |
| emb = torch.cat((freqs, freqs), dim=-1) |
| |
| self.register_buffer('cos_cached', emb.cos()[None, None, :, :].to(dtype), persistent=False) |
| self.register_buffer('sin_cached', emb.sin()[None, None, :, :].to(dtype), persistent=False) |
| |
| def forward(self, x: torch.Tensor, seq_len: int) -> Tuple[torch.Tensor, torch.Tensor]: |
| if seq_len > self.max_seq_len_cached: |
| self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=x.dtype) |
| |
| return ( |
| self.cos_cached[:, :, :seq_len, ...].to(dtype=x.dtype), |
| self.sin_cached[:, :, :seq_len, ...].to(dtype=x.dtype), |
| ) |
|
|
|
|
| def rotate_half(x: torch.Tensor) -> torch.Tensor: |
| x1 = x[..., : x.shape[-1] // 2] |
| x2 = x[..., x.shape[-1] // 2 :] |
| return torch.cat((-x2, x1), dim=-1) |
|
|
|
|
| class RoformerLayer(nn.Module): |
| def __init__( |
| self, |
| feature_dim: int, |
| num_heads: int = 8, |
| max_seq_len: int = 10000, |
| dropout: float = 0.0, |
| mlp_ratio: float = 4.0, |
| rope_base: int = 10000 |
| ): |
| super().__init__() |
| assert feature_dim % num_heads == 0, "feature_dim must be divisible by num_heads" |
|
|
| self.feature_dim = feature_dim |
| self.num_heads = num_heads |
| self.head_dim = feature_dim // num_heads |
|
|
| self.rotary_emb = RotaryEmbedding(self.head_dim, max_position_embeddings=max_seq_len, base=rope_base) |
| self.dropout = dropout |
|
|
| self.input_norm = RMSNorm(feature_dim) |
| self.qkv_proj = nn.Linear(feature_dim, feature_dim * 3, bias=False) |
| self.output_proj = nn.Linear(feature_dim, feature_dim, bias=False) |
|
|
| mlp_hidden_dim = int(feature_dim * mlp_ratio) |
| self.mlp_norm = RMSNorm(feature_dim) |
| self.mlp_up = nn.Linear(feature_dim, mlp_hidden_dim * 2, bias=False) |
| self.mlp_down = nn.Linear(mlp_hidden_dim, feature_dim, bias=False) |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| B, N, T = x.shape |
| x_residual = x |
| x_norm = self.input_norm(x).transpose(1, 2) |
|
|
| qkv = self.qkv_proj(x_norm) |
| qkv = qkv.view(B, T, 3, self.num_heads, self.head_dim) |
| qkv = qkv.permute(2, 0, 3, 1, 4) |
| Q, K, V = qkv[0], qkv[1], qkv[2] |
|
|
| cos, sin = self.rotary_emb(Q, seq_len=T) |
| Q = (Q * cos) + (rotate_half(Q) * sin) |
| K = (K * cos) + (rotate_half(K) * sin) |
|
|
| attn_output = F.scaled_dot_product_attention( |
| Q, K, V, dropout_p=self.dropout if self.training else 0.0, is_causal=False |
| ) |
|
|
| attn_output = attn_output.permute(0, 2, 1, 3).contiguous().view(B, T, N) |
| attn_output = self.output_proj(attn_output).transpose(1, 2) |
|
|
| x = x_residual + attn_output |
|
|
| x_residual = x |
| x_norm = self.mlp_norm(x).transpose(1, 2) |
|
|
| mlp_out = self.mlp_up(x_norm) |
| gate, values = mlp_out.chunk(2, dim=-1) |
| mlp_out = F.silu(gate) * values |
| mlp_out = self.mlp_down(mlp_out) |
|
|
| output = x_residual + mlp_out.transpose(1, 2) |
|
|
| return output |
|
|
| class Roformer(nn.Module): |
| def __init__(self, input_size, hidden_size, num_head=8, theta=10000, window=10000, |
| input_drop=0., attention_drop=0., causal=True): |
| super().__init__() |
|
|
| self.input_size = input_size |
| self.hidden_size = hidden_size // num_head |
| self.num_head = num_head |
| self.theta = theta |
| self.window = window |
| cos_freq, sin_freq = self._calc_rotary_emb() |
| self.register_buffer("cos_freq", cos_freq) |
| self.register_buffer("sin_freq", sin_freq) |
| |
| self.attention_drop = attention_drop |
| self.causal = causal |
| self.eps = 1e-5 |
|
|
| self.input_norm = RMSNorm(self.input_size) |
| self.input_drop = nn.Dropout(p=input_drop) |
| self.weight = nn.Conv1d(self.input_size, self.hidden_size*self.num_head*3, 1, bias=False) |
| self.output = nn.Conv1d(self.hidden_size*self.num_head, self.input_size, 1, bias=False) |
|
|
| self.MLP = nn.Sequential(RMSNorm(self.input_size), |
| nn.Conv1d(self.input_size, self.input_size*8, 1, bias=False), |
| nn.SiLU() |
| ) |
| self.MLP_output = nn.Conv1d(self.input_size*4, self.input_size, 1, bias=False) |
|
|
| def _calc_rotary_emb(self): |
| freq = 1. / (self.theta ** (torch.arange(0, self.hidden_size, 2)[:(self.hidden_size // 2)] / self.hidden_size)) |
| freq = freq.reshape(1, -1) |
| pos = torch.arange(0, self.window).reshape(-1, 1) |
| cos_freq = torch.cos(pos*freq) |
| sin_freq = torch.sin(pos*freq) |
| cos_freq = torch.stack([cos_freq]*2, -1).reshape(self.window, self.hidden_size) |
| sin_freq = torch.stack([sin_freq]*2, -1).reshape(self.window, self.hidden_size) |
|
|
| return cos_freq, sin_freq |
| |
| def _add_rotary_emb(self, feature, pos): |
| N = feature.shape[-1] |
|
|
| feature_reshape = feature.reshape(-1, N) |
| pos = min(pos, self.window-1) |
| cos_freq = self.cos_freq[pos] |
| sin_freq = self.sin_freq[pos] |
| reverse_sign = torch.from_numpy(np.asarray([-1, 1])).to(feature.device).type(feature.dtype) |
| feature_reshape_neg = (torch.flip(feature_reshape.reshape(-1, N//2, 2), [-1]) * reverse_sign.reshape(1, 1, 2)).reshape(-1, N) |
| feature_rope = feature_reshape * cos_freq.unsqueeze(0) + feature_reshape_neg * sin_freq.unsqueeze(0) |
| |
| return feature_rope.reshape(feature.shape) |
|
|
| def _add_rotary_sequence(self, feature): |
| T, N = feature.shape[-2:] |
| feature_reshape = feature.reshape(-1, T, N) |
|
|
| cos_freq = self.cos_freq[:T] |
| sin_freq = self.sin_freq[:T] |
| reverse_sign = torch.from_numpy(np.asarray([-1, 1])).to(feature.device).type(feature.dtype) |
| feature_reshape_neg = (torch.flip(feature_reshape.reshape(-1, N//2, 2), [-1]) * reverse_sign.reshape(1, 1, 2)).reshape(-1, T, N) |
| feature_rope = feature_reshape * cos_freq.unsqueeze(0) + feature_reshape_neg * sin_freq.unsqueeze(0) |
| |
| return feature_rope.reshape(feature.shape) |
| |
| def forward(self, input): |
| B, _, T = input.shape |
|
|
| weight = self.weight(self.input_drop(self.input_norm(input))).reshape(B, self.num_head, self.hidden_size*3, T).transpose(-2,-1) |
| Q, K, V = torch.split(weight, self.hidden_size, dim=-1) |
| Q_rot = self._add_rotary_sequence(Q) |
| K_rot = self._add_rotary_sequence(K) |
|
|
| attention_output = F.scaled_dot_product_attention(Q_rot.contiguous(), K_rot.contiguous(), V.contiguous(), dropout_p=self.attention_drop, is_causal=self.causal) |
| attention_output = attention_output.transpose(-2,-1).reshape(B, -1, T) |
| output = self.output(attention_output) + input |
|
|
| gate, z = self.MLP(output).chunk(2, dim=1) |
| output = output + self.MLP_output(F.silu(gate) * z) |
|
|
| return output |
|
|
| class ConvBlock(nn.Module): |
| def __init__(self, channels: int, kernel_size: int, dilation: int, expansion: int = 4): |
| super().__init__() |
| padding = (kernel_size - 1) * dilation // 2 |
|
|
| self.dwconv = nn.Conv1d( |
| channels, channels, kernel_size, padding=padding, dilation=dilation, groups=channels |
| ) |
| self.norm = RMSNorm(channels) |
| self.pwconv1 = nn.Conv1d(channels, channels * expansion, 1) |
| self.act = nn.GLU(dim=1) |
| self.pwconv2 = nn.Conv1d(channels * expansion // 2, channels, 1) |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| x = self.dwconv(x) |
| x = self.norm(x) |
| x = self.pwconv1(x) |
| x = self.act(x) |
| x = self.pwconv2(x) |
| return x |
|
|
|
|
| class ICB(nn.Module): |
| def __init__(self, channels: int, kernel_size: int = 7, dilation: int = 1, layer_scale_init_value: float = 1e-6): |
| super().__init__() |
| self.block1 = ConvBlock(channels, kernel_size, 1, ) |
| self.block2 = ConvBlock(channels, kernel_size, dilation) |
| self.block3 = ConvBlock(channels, kernel_size, 1) |
| self.gamma = nn.Parameter(layer_scale_init_value * torch.ones((channels)), requires_grad=True) |
| |
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| residual = x |
| x = self.block1(x) |
| x = self.block2(x) |
| x = self.block3(x) |
| return x * self.gamma.unsqueeze(0).unsqueeze(-1) + residual |
|
|
|
|
| class BSNet(nn.Module): |
| def __init__( |
| self, |
| feature_dim: int, |
| kernel_size: int, |
| dilation_rate: int, |
| num_heads: int, |
| max_bands: int = 512, |
| band_rope_base: int = 10000, |
| layer_scale_init_value: float = 1e-6 |
| ): |
| super().__init__() |
| self.band_net = Roformer(feature_dim, feature_dim, num_head=num_heads, window=max_bands, causal=False) |
|
|
| self.seq_net = ICB( |
| feature_dim, |
| kernel_size=kernel_size, |
| dilation=dilation_rate, |
| layer_scale_init_value=layer_scale_init_value |
| ) |
| |
| def forward(self, input: torch.Tensor) -> torch.Tensor: |
| B, nband, N, T = input.shape |
| |
| band_input = input.permute(0, 3, 2, 1).reshape(B * T, N, nband) |
| band_output = self.band_net(band_input) |
| band_output = band_output.view(B, T, N, nband).permute(0, 3, 2, 1) |
| |
| seq_input = band_output.reshape(B * nband, N, T) |
| seq_output = self.seq_net(seq_input) |
| output = seq_output.view(B, nband, N, T) |
| |
| return output |
|
|
| class Renaissance(nn.Module): |
| def __init__( |
| self, |
| n_freqs: int = 2049, |
| feature_dim: int = 128, |
| layer: int = 9, |
| sample_rate: int = 48000, |
| dilation_start_layer: int = 3, |
| n_bands: int = 80, |
| num_heads: int = 16, |
| max_seq_len: int = 10000, |
| band_rope_base: int = 10000, |
| temporal_rope_base: int = 10000 |
| ): |
| super().__init__() |
| self.enc_dim = n_freqs |
| self.feature_dim = feature_dim |
| self.eps = 1e-7 |
| self.dilation_start_layer = dilation_start_layer |
| self.n_bands = n_bands |
| self.sr = sample_rate |
| self.max_seq_len = max_seq_len |
| self.band_rope_base = band_rope_base |
| self.temporal_rope_base = temporal_rope_base |
|
|
| self.band_width = self._generate_mel_bandwidths() |
| self.nband = len(self.band_width) |
| assert self.enc_dim == sum(self.band_width), "Mel band splitting failed to cover all frequencies." |
|
|
| self._build_feature_extractor() |
| self._build_main_network(layer, num_heads) |
| self._build_output_synthesis() |
|
|
| self.apply(self._init_weights) |
|
|
| def _init_weights(self, module): |
| if isinstance(module, (nn.Linear, nn.Conv1d)): |
| nn.init.kaiming_normal_(module.weight, mode='fan_out', nonlinearity='relu') |
| if module.bias is not None: |
| nn.init.zeros_(module.bias) |
|
|
| def _generate_mel_bandwidths(self) -> List[int]: |
| def hz_to_mel(hz): return 2595 * np.log10(1 + hz / 700.) |
| def mel_to_hz(mel): return 700 * (10**(mel / 2595) - 1) |
|
|
| min_freq, max_freq = 0.1, self.sr / 2 |
| min_mel, max_mel = hz_to_mel(min_freq), hz_to_mel(max_freq) |
|
|
| mel_points = np.linspace(min_mel, max_mel, self.n_bands + 1) |
| hz_points = mel_to_hz(mel_points) |
| |
| bin_width = self.sr / 2 / self.enc_dim |
| bw = np.round(np.diff(hz_points) / bin_width).astype(int) |
| |
| bw = np.maximum(1, bw) |
| |
| remainder = self.enc_dim - np.sum(bw) |
| if remainder != 0: |
| sorted_indices = np.argsort(bw) |
| op = 1 if remainder > 0 else -1 |
| indices_to_adjust = sorted_indices if op == 1 else sorted_indices[::-1] |
| |
| for i in range(abs(remainder)): |
| idx = indices_to_adjust[i % len(indices_to_adjust)] |
| if bw[idx] + op > 0: |
| bw[idx] += op |
| |
| if np.sum(bw) != self.enc_dim: |
| bw[-1] += self.enc_dim - np.sum(bw) |
| |
| return bw.tolist() |
|
|
| def _build_feature_extractor(self): |
| self.feature_extractor_layers = nn.ModuleList([ |
| nn.Sequential(RMSNorm(bw * 2 + 1), nn.Conv1d(bw * 2 + 1, self.feature_dim, 1)) |
| for bw in self.band_width |
| ]) |
| |
| def _build_main_network(self, num_layers, num_heads): |
| self.net = nn.ModuleList() |
| max_bands = max(512, self.nband * 2) |
|
|
| layer_scale_init = 1e-6 |
| |
| for i in range(num_layers): |
| dilation = min(2 ** max(0, i - self.dilation_start_layer + 1), 4) |
| self.net.append(BSNet( |
| self.feature_dim, |
| kernel_size=7, |
| dilation_rate=dilation, |
| num_heads=num_heads, |
| max_bands=max_bands, |
| band_rope_base=self.band_rope_base, |
| layer_scale_init_value=layer_scale_init |
| )) |
|
|
| def _build_output_synthesis(self): |
| self.output_layers = nn.ModuleList([ |
| nn.Sequential( |
| RMSNorm(self.feature_dim), |
| nn.Conv1d(self.feature_dim, self.feature_dim * 2, 1), |
| nn.SiLU(), |
| nn.Conv1d(self.feature_dim * 2, bw * 4, kernel_size=1), |
| nn.GLU(dim=1), |
| ) for bw in self.band_width |
| ]) |
|
|
| def spec_band_split(self, spec: torch.Tensor) -> Tuple[List[torch.Tensor], torch.Tensor]: |
| subband_spec_ri = [] |
| subband_power = [] |
| band_idx = 0 |
| for width in self.band_width: |
| this_spec_ri = spec[:, band_idx : band_idx + width, :, :] |
| subband_spec_ri.append(this_spec_ri) |
| |
| power = (this_spec_ri.pow(2).sum(dim=-1)).sum(dim=1, keepdim=True).add(self.eps).sqrt() |
| subband_power.append(power) |
| band_idx += width |
| |
| subband_power = torch.cat(subband_power, 1) |
| return subband_spec_ri, subband_power |
|
|
| def feature_extraction(self, input_spec: torch.Tensor) -> torch.Tensor: |
| subband_spec_ri, subband_power = self.spec_band_split(input_spec) |
| features = [] |
| for i in range(self.nband): |
| power_for_norm = subband_power[:, i:i+1, :].unsqueeze(1) |
| norm_spec_ri = subband_spec_ri[i] / (power_for_norm.transpose(2,3) + self.eps) |
| B, F_band, T, _ = norm_spec_ri.shape |
| norm_spec_flat = norm_spec_ri.permute(0, 3, 1, 2).reshape(B, F_band*2, T) |
|
|
| log_power_feature = torch.log(power_for_norm.squeeze(1) + self.eps) |
| feature_input = torch.cat([norm_spec_flat, log_power_feature], dim=1) |
| |
| features.append(self.feature_extractor_layers[i](feature_input)) |
| |
| return torch.stack(features, 1) |
|
|
| def forward(self, input_spec: torch.Tensor) -> torch.Tensor: |
| B, F, T, _ = input_spec.shape |
|
|
| features = self.feature_extraction(input_spec) |
|
|
| residual_features = features |
| processed = features |
| for layer in self.net: |
| processed = layer(processed) |
| processed = processed + residual_features |
|
|
| est_spec_bands = [] |
| for i in range(self.nband): |
| band_output = self.output_layers[i](processed[:, i]) |
| bw = self.band_width[i] |
| est_spec_band = band_output.view(B, bw, 2, T).permute(0, 1, 3, 2) |
| est_spec_bands.append(est_spec_band) |
| est_spec_full = torch.cat(est_spec_bands, dim=1) |
|
|
| return est_spec_full |
|
|
