mirror of
https://github.com/RVC-Project/Retrieval-based-Voice-Conversion-WebUI.git
synced 2026-09-01 19:48:26 +02:00
88 lines
2.6 KiB
Python
88 lines
2.6 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
|
|
|
|
class SubbandSTFT:
|
|
def __init__(self, config):
|
|
self.n_fft = config.n_fft
|
|
self.hop_length = config.hop_length
|
|
self.window = torch.hann_window(window_length=self.n_fft, periodic=True)
|
|
self.dim_f = config.dim_f
|
|
|
|
def __call__(self, x):
|
|
window = self.window.to(x.device)
|
|
batch_dims = x.shape[:-2]
|
|
channels, length = x.shape[-2:]
|
|
x = torch.stft(
|
|
x.reshape(-1, length),
|
|
n_fft=self.n_fft,
|
|
hop_length=self.hop_length,
|
|
window=window,
|
|
center=True,
|
|
return_complex=True,
|
|
)
|
|
x = torch.view_as_real(x).permute(0, 3, 1, 2)
|
|
x = x.reshape(*batch_dims, channels * 2, -1, x.shape[-1])
|
|
return x[..., : self.dim_f, :]
|
|
|
|
def inverse(self, x):
|
|
window = self.window.to(x.device)
|
|
batch_dims = x.shape[:-3]
|
|
channels, freq_bins, time_bins = x.shape[-3:]
|
|
full_freq_bins = self.n_fft // 2 + 1
|
|
f_pad = torch.zeros([*batch_dims, channels, full_freq_bins - freq_bins, time_bins]).to(x.device)
|
|
x = torch.cat([x, f_pad], -2)
|
|
x = x.reshape(-1, 2, full_freq_bins, time_bins).permute(0, 2, 3, 1)
|
|
x = x[..., 0] + x[..., 1] * 1.0j
|
|
x = torch.istft(
|
|
x,
|
|
n_fft=self.n_fft,
|
|
hop_length=self.hop_length,
|
|
window=window,
|
|
center=True,
|
|
)
|
|
return x.reshape([*batch_dims, 2, -1])
|
|
|
|
|
|
def get_activation(act_type):
|
|
if act_type == "gelu":
|
|
return nn.GELU()
|
|
if act_type == "relu":
|
|
return nn.ReLU()
|
|
if act_type[:3] == "elu":
|
|
alpha = float(act_type.replace("elu", ""))
|
|
return nn.ELU(alpha)
|
|
raise Exception
|
|
|
|
|
|
def cac_to_cws(x, num_subbands):
|
|
batch, channels, freq_bins, time_bins = x.shape
|
|
return x.reshape(batch, channels * num_subbands, freq_bins // num_subbands, time_bins)
|
|
|
|
|
|
def cws_to_cac(x, num_subbands):
|
|
batch, channels, freq_bins, time_bins = x.shape
|
|
return x.reshape(batch, channels // num_subbands, freq_bins * num_subbands, time_bins)
|
|
|
|
|
|
def forward_subband_mask_model(module, x, core_fn):
|
|
x = module.stft(x)
|
|
|
|
mix = x = cac_to_cws(x, module.num_subbands)
|
|
|
|
first_conv_out = x = module.first_conv(x)
|
|
|
|
x = core_fn(x.transpose(-1, -2)).transpose(-1, -2)
|
|
|
|
x = x * first_conv_out
|
|
|
|
x = module.final_conv(torch.cat([mix, x], 1))
|
|
|
|
x = cws_to_cac(x, module.num_subbands)
|
|
|
|
if module.num_target_instruments > 1:
|
|
batch, channels, freq_bins, time_bins = x.shape
|
|
x = x.reshape(batch, module.num_target_instruments, -1, freq_bins, time_bins)
|
|
|
|
return module.stft.inverse(x)
|