ref: f7ee7137425729b4759d868b64ca3d18f2ab0f95
dir: /dnn/torch/osce/utils/layers/silk_upsampler.py/
""" This module implements the SILK upsampler from 16kHz to 24 or 48 kHz """
import torch
from torch import nn
import torch.nn.functional as F
import numpy as np
frac_fir = np.array(
[
[189, -600, 617, 30567, 2996, -1375, 425, -46],
[117, -159, -1070, 29704, 5784, -2143, 611, -71],
[52, 221, -2392, 28276, 8798, -2865, 773, -91],
[-4, 529, -3350, 26341, 11950, -3487, 896, -103],
[-48, 758, -3956, 23973, 15143, -3957, 967, -107],
[-80, 905, -4235, 21254, 18278, -4222, 972, -99],
[-99, 972, -4222, 18278, 21254, -4235, 905, -80],
[-107, 967, -3957, 15143, 23973, -3956, 758, -48],
[-103, 896, -3487, 11950, 26341, -3350, 529, -4],
[-91, 773, -2865, 8798, 28276, -2392, 221, 52],
[-71, 611, -2143, 5784, 29704, -1070, -159, 117],
[-46, 425, -1375, 2996, 30567, 617, -600, 189]
],
dtype=np.float32
) / 2**15
hq_2x_up_c_even = [x / 2**16 for x in [1746, 14986, 39083 - 65536]]
hq_2x_up_c_odd = [x / 2**16 for x in [6854, 25769, 55542 - 65536]]
def get_impz(coeffs, n):
s = 3*[0]
y = np.zeros(n)
x = 1
for i in range(n):
Y = x - s[0]
X = Y * coeffs[0]
tmp1 = s[0] + X
s[0] = x + X
Y = tmp1 - s[1]
X = Y * coeffs[1]
tmp2 = s[1] + X
s[1] = tmp1 + X
Y = tmp2 - s[2]
X = Y * (1 + coeffs[2])
tmp3 = s[2] + X
s[2] = tmp2 + X
y[i] = tmp3
x = 0
return y
class SilkUpsampler(nn.Module):
SUPPORTED_TARGET_RATES = {24000, 48000}
SUPPORTED_SOURCE_RATES = {16000}
def __init__(self,
fs_in=16000,
fs_out=48000):
super().__init__()
self.fs_in = fs_in
self.fs_out = fs_out
if fs_in not in self.SUPPORTED_SOURCE_RATES:
raise ValueError(f'SilkUpsampler currently only supports upsampling from {self.SUPPORTED_SOURCE_RATES} Hz')
if fs_out not in self.SUPPORTED_TARGET_RATES:
raise ValueError(f'SilkUpsampler currently only supports upsampling to {self.SUPPORTED_TARGET_RATES} Hz')
# hq 2x upsampler as FIR approximation
hq_2x_up_even = get_impz(hq_2x_up_c_even, 128)[::-1].copy()
hq_2x_up_odd = get_impz(hq_2x_up_c_odd , 128)[::-1].copy()
self.hq_2x_up_even = nn.Parameter(torch.from_numpy(hq_2x_up_even).float().view(1, 1, -1), requires_grad=False)
self.hq_2x_up_odd = nn.Parameter(torch.from_numpy(hq_2x_up_odd ).float().view(1, 1, -1), requires_grad=False)
self.hq_2x_up_padding = [127, 0]
# interpolation filters
frac_01_24 = frac_fir[0]
frac_17_24 = frac_fir[8]
frac_09_24 = frac_fir[4]
self.frac_01_24 = nn.Parameter(torch.from_numpy(frac_01_24).view(1, 1, -1), requires_grad=False)
self.frac_17_24 = nn.Parameter(torch.from_numpy(frac_17_24).view(1, 1, -1), requires_grad=False)
self.frac_09_24 = nn.Parameter(torch.from_numpy(frac_09_24).view(1, 1, -1), requires_grad=False)
self.stride = 1 if fs_out == 48000 else 2
def hq_2x_up(self, x):
num_channels = x.size(1)
weight_even = torch.repeat_interleave(self.hq_2x_up_even, num_channels, 0)
weight_odd = torch.repeat_interleave(self.hq_2x_up_odd , num_channels, 0)
x_pad = F.pad(x, self.hq_2x_up_padding)
y_even = F.conv1d(x_pad, weight_even, groups=num_channels)
y_odd = F.conv1d(x_pad, weight_odd , groups=num_channels)
y = torch.cat((y_even.unsqueeze(-1), y_odd.unsqueeze(-1)), dim=-1).flatten(2)
return y
def interpolate_3_2(self, x):
num_channels = x.size(1)
weight_01_24 = torch.repeat_interleave(self.frac_01_24, num_channels, 0)
weight_17_24 = torch.repeat_interleave(self.frac_17_24, num_channels, 0)
weight_09_24 = torch.repeat_interleave(self.frac_09_24, num_channels, 0)
x_pad = F.pad(x, [8, 0])
y_01_24 = F.conv1d(x_pad, weight_01_24, stride=2, groups=num_channels)
y_17_24 = F.conv1d(x_pad, weight_17_24, stride=2, groups=num_channels)
y_09_24_sh1 = F.conv1d(torch.roll(x_pad, -1, -1), weight_09_24, stride=2, groups=num_channels)
y = torch.cat(
(y_01_24.unsqueeze(-1), y_17_24.unsqueeze(-1), y_09_24_sh1.unsqueeze(-1)),
dim=-1).flatten(2)
return y[..., :-3]
def forward(self, x):
y_2x = self.hq_2x_up(x)
y_3x = self.interpolate_3_2(y_2x)
return y_3x[:, :, ::self.stride]