|
|
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 |
|
|
|
