mirror of
https://github.com/RVC-Project/Retrieval-based-Voice-Conversion-WebUI.git
synced 2026-09-01 19:48:26 +02:00
480 lines
16 KiB
Python
480 lines
16 KiB
Python
import torch
|
|
import math
|
|
from torch import nn
|
|
from torch.nn import functional as F
|
|
from fractions import Fraction
|
|
|
|
from .demucs_local import (
|
|
CrossTransformerEncoder,
|
|
HDecLayer,
|
|
HEncLayer,
|
|
MultiWrap,
|
|
ScaledEmbedding,
|
|
ispectro,
|
|
pad1d,
|
|
rescale_module,
|
|
spectro,
|
|
)
|
|
from ..config import to_plain
|
|
|
|
|
|
class HTDemucs(nn.Module):
|
|
mps_model_backend = "torch"
|
|
mps_model_compute_dtype = torch.float16
|
|
|
|
def __init__(
|
|
self,
|
|
sources,
|
|
audio_channels=2,
|
|
channels=48,
|
|
channels_time=None,
|
|
growth=2,
|
|
nfft=4096,
|
|
num_subbands=1,
|
|
wiener_iters=0,
|
|
end_iters=0,
|
|
wiener_residual=False,
|
|
cac=True,
|
|
depth=4,
|
|
rewrite=True,
|
|
multi_freqs=None,
|
|
multi_freqs_depth=3,
|
|
freq_emb=0.2,
|
|
emb_scale=10,
|
|
emb_smooth=True,
|
|
kernel_size=8,
|
|
time_stride=2,
|
|
stride=4,
|
|
context=1,
|
|
context_enc=0,
|
|
norm_starts=4,
|
|
norm_groups=4,
|
|
dconv_mode=1,
|
|
dconv_depth=2,
|
|
dconv_comp=8,
|
|
dconv_init=1e-3,
|
|
bottom_channels=0,
|
|
t_layers=5,
|
|
t_emb="sin",
|
|
t_hidden_scale=4.0,
|
|
t_heads=8,
|
|
t_dropout=0.0,
|
|
t_max_positions=10000,
|
|
t_norm_in=True,
|
|
t_norm_in_group=False,
|
|
t_group_norm=False,
|
|
t_norm_first=True,
|
|
t_norm_out=True,
|
|
t_max_period=10000.0,
|
|
t_weight_decay=0.0,
|
|
t_lr=None,
|
|
t_layer_scale=True,
|
|
t_gelu=True,
|
|
t_weight_pos_embed=1.0,
|
|
t_sin_random_shift=0,
|
|
t_cape_mean_normalize=True,
|
|
t_cape_augment=True,
|
|
t_cape_glob_loc_scale=[5000.0, 1.0, 1.4],
|
|
t_sparse_self_attn=False,
|
|
t_sparse_cross_attn=False,
|
|
t_mask_type="diag",
|
|
t_mask_random_seed=42,
|
|
t_sparse_attn_window=500,
|
|
t_global_window=100,
|
|
t_sparsity=0.95,
|
|
t_auto_sparsity=False,
|
|
t_cross_first=False,
|
|
rescale=0.1,
|
|
samplerate=44100,
|
|
segment=10,
|
|
use_train_segment=False,
|
|
):
|
|
super().__init__()
|
|
self.num_subbands = num_subbands
|
|
self.cac = cac
|
|
self.wiener_residual = wiener_residual
|
|
self.audio_channels = audio_channels
|
|
self.sources = sources
|
|
self.kernel_size = kernel_size
|
|
self.context = context
|
|
self.stride = stride
|
|
self.depth = depth
|
|
self.bottom_channels = bottom_channels
|
|
self.channels = channels
|
|
self.samplerate = samplerate
|
|
self.segment = segment
|
|
self.use_train_segment = use_train_segment
|
|
self.nfft = nfft
|
|
self.hop_length = nfft // 4
|
|
self.wiener_iters = wiener_iters
|
|
self.end_iters = end_iters
|
|
self.freq_emb = None
|
|
assert wiener_iters == end_iters
|
|
|
|
self.encoder = nn.ModuleList()
|
|
self.decoder = nn.ModuleList()
|
|
|
|
self.tencoder = nn.ModuleList()
|
|
self.tdecoder = nn.ModuleList()
|
|
|
|
chin = audio_channels
|
|
chin_z = chin # number of channels for the freq branch
|
|
if self.cac:
|
|
chin_z *= 2
|
|
if self.num_subbands > 1:
|
|
chin_z *= self.num_subbands
|
|
chout = channels_time or channels
|
|
chout_z = channels
|
|
freqs = nfft // 2
|
|
|
|
for index in range(depth):
|
|
norm = index >= norm_starts
|
|
freq = freqs > 1
|
|
stri = stride
|
|
ker = kernel_size
|
|
if not freq:
|
|
assert freqs == 1
|
|
ker = time_stride * 2
|
|
stri = time_stride
|
|
|
|
pad = True
|
|
last_freq = False
|
|
if freq and freqs <= kernel_size:
|
|
ker = freqs
|
|
pad = False
|
|
last_freq = True
|
|
|
|
kw = {
|
|
"kernel_size": ker,
|
|
"stride": stri,
|
|
"freq": freq,
|
|
"pad": pad,
|
|
"norm": norm,
|
|
"rewrite": rewrite,
|
|
"norm_groups": norm_groups,
|
|
"dconv_kw": {
|
|
"depth": dconv_depth,
|
|
"compress": dconv_comp,
|
|
"init": dconv_init,
|
|
"gelu": True,
|
|
},
|
|
}
|
|
kwt = dict(kw)
|
|
kwt["freq"] = 0
|
|
kwt["kernel_size"] = kernel_size
|
|
kwt["stride"] = stride
|
|
kwt["pad"] = True
|
|
kw_dec = dict(kw)
|
|
multi = False
|
|
if multi_freqs and index < multi_freqs_depth:
|
|
multi = True
|
|
kw_dec["context_freq"] = False
|
|
|
|
if last_freq:
|
|
chout_z = max(chout, chout_z)
|
|
chout = chout_z
|
|
|
|
enc = HEncLayer(chin_z, chout_z, dconv=dconv_mode & 1, context=context_enc, **kw)
|
|
if freq:
|
|
tenc = HEncLayer(chin, chout, dconv=dconv_mode & 1, context=context_enc, empty=last_freq, **kwt)
|
|
self.tencoder.append(tenc)
|
|
|
|
if multi:
|
|
enc = MultiWrap(enc, multi_freqs)
|
|
self.encoder.append(enc)
|
|
if index == 0:
|
|
chin = self.audio_channels * len(self.sources)
|
|
chin_z = chin
|
|
if self.cac:
|
|
chin_z *= 2
|
|
if self.num_subbands > 1:
|
|
chin_z *= self.num_subbands
|
|
dec = HDecLayer(chout_z, chin_z, dconv=dconv_mode & 2, last=index == 0, context=context, **kw_dec)
|
|
if multi:
|
|
dec = MultiWrap(dec, multi_freqs)
|
|
if freq:
|
|
tdec = HDecLayer(chout, chin, dconv=dconv_mode & 2, empty=last_freq, last=index == 0, context=context, **kwt)
|
|
self.tdecoder.insert(0, tdec)
|
|
self.decoder.insert(0, dec)
|
|
|
|
chin = chout
|
|
chin_z = chout_z
|
|
chout = int(growth * chout)
|
|
chout_z = int(growth * chout_z)
|
|
if freq:
|
|
if freqs <= kernel_size:
|
|
freqs = 1
|
|
else:
|
|
freqs //= stride
|
|
if index == 0 and freq_emb:
|
|
self.freq_emb = ScaledEmbedding(freqs, chin_z, smooth=emb_smooth, scale=emb_scale)
|
|
self.freq_emb_scale = freq_emb
|
|
|
|
if rescale:
|
|
rescale_module(self, reference=rescale)
|
|
|
|
transformer_channels = channels * growth ** (depth - 1)
|
|
if bottom_channels:
|
|
self.channel_upsampler = nn.Conv1d(transformer_channels, bottom_channels, 1)
|
|
self.channel_downsampler = nn.Conv1d(bottom_channels, transformer_channels, 1)
|
|
self.channel_upsampler_t = nn.Conv1d(transformer_channels, bottom_channels, 1)
|
|
self.channel_downsampler_t = nn.Conv1d(bottom_channels, transformer_channels, 1)
|
|
|
|
transformer_channels = bottom_channels
|
|
|
|
if t_layers > 0:
|
|
self.crosstransformer = CrossTransformerEncoder(
|
|
dim=transformer_channels,
|
|
emb=t_emb,
|
|
hidden_scale=t_hidden_scale,
|
|
num_heads=t_heads,
|
|
num_layers=t_layers,
|
|
cross_first=t_cross_first,
|
|
dropout=t_dropout,
|
|
max_positions=t_max_positions,
|
|
norm_in=t_norm_in,
|
|
norm_in_group=t_norm_in_group,
|
|
group_norm=t_group_norm,
|
|
norm_first=t_norm_first,
|
|
norm_out=t_norm_out,
|
|
max_period=t_max_period,
|
|
weight_decay=t_weight_decay,
|
|
lr=t_lr,
|
|
layer_scale=t_layer_scale,
|
|
gelu=t_gelu,
|
|
sin_random_shift=t_sin_random_shift,
|
|
weight_pos_embed=t_weight_pos_embed,
|
|
cape_mean_normalize=t_cape_mean_normalize,
|
|
cape_augment=t_cape_augment,
|
|
cape_glob_loc_scale=t_cape_glob_loc_scale,
|
|
sparse_self_attn=t_sparse_self_attn,
|
|
sparse_cross_attn=t_sparse_cross_attn,
|
|
mask_type=t_mask_type,
|
|
mask_random_seed=t_mask_random_seed,
|
|
sparse_attn_window=t_sparse_attn_window,
|
|
global_window=t_global_window,
|
|
sparsity=t_sparsity,
|
|
auto_sparsity=t_auto_sparsity,
|
|
)
|
|
else:
|
|
self.crosstransformer = None
|
|
|
|
def set_mps_model_backend(self, backend=None, compute_dtype=None):
|
|
backend = (backend or "torch").lower()
|
|
if backend not in ("torch", "mlx_full"):
|
|
raise ValueError("mps_model_backend must be 'torch' or 'mlx_full'")
|
|
self.mps_model_backend = backend
|
|
if compute_dtype is None:
|
|
return
|
|
if isinstance(compute_dtype, str):
|
|
compute_dtype = {
|
|
"float16": torch.float16,
|
|
"fp16": torch.float16,
|
|
"float32": torch.float32,
|
|
"fp32": torch.float32,
|
|
}.get(compute_dtype.lower(), compute_dtype)
|
|
if compute_dtype not in (torch.float16, torch.float32):
|
|
raise ValueError("mps_model_compute_dtype must be 'float16' or 'float32'")
|
|
self.mps_model_compute_dtype = compute_dtype
|
|
|
|
def _use_mlx_full_forward(self, mix):
|
|
return not self.training and self.mps_model_backend == "mlx_full" and mix.device.type == "mps"
|
|
|
|
def mlx_forward_mx(self, raw_audio):
|
|
from .demucs_mlx import mlx_forward_demucs_mx
|
|
|
|
return mlx_forward_demucs_mx(self, raw_audio, self.mps_model_compute_dtype)
|
|
|
|
def _spec(self, x):
|
|
hl = self.hop_length
|
|
nfft = self.nfft
|
|
x0 = x # noqa
|
|
|
|
assert hl == nfft // 4
|
|
le = int(math.ceil(x.shape[-1] / hl))
|
|
pad = hl // 2 * 3
|
|
x = pad1d(x, (pad, pad + le * hl - x.shape[-1]), mode="reflect")
|
|
|
|
z = spectro(x, nfft, hl)[..., :-1, :]
|
|
assert z.shape[-1] == le + 4, (z.shape, x.shape, le)
|
|
z = z[..., 2 : 2 + le]
|
|
return z
|
|
|
|
def _ispec(self, z, length=None, scale=0):
|
|
hl = self.hop_length // (4**scale)
|
|
z = F.pad(z, (0, 0, 0, 1))
|
|
z = F.pad(z, (2, 2))
|
|
pad = hl // 2 * 3
|
|
le = hl * int(math.ceil(length / hl)) + 2 * pad
|
|
x = ispectro(z, hl, length=le)
|
|
x = x[..., pad : pad + length]
|
|
return x
|
|
|
|
def _magnitude(self, z):
|
|
if self.cac:
|
|
B, C, Fr, T = z.shape
|
|
m = torch.view_as_real(z).permute(0, 1, 4, 2, 3).reshape(B, C * 2, Fr, T)
|
|
else:
|
|
m = z.abs()
|
|
return m
|
|
|
|
def _mask(self, z, m):
|
|
niters = self.wiener_iters
|
|
if self.cac:
|
|
B, S, C, Fr, T = m.shape
|
|
return torch.view_as_complex(m.view(B, S, -1, 2, Fr, T).permute(0, 1, 2, 4, 5, 3).contiguous())
|
|
if self.training:
|
|
niters = self.end_iters
|
|
if niters < 0:
|
|
z = z[:, None]
|
|
return z / (1e-8 + z.abs()) * m
|
|
else:
|
|
return self._wiener(m, z, niters)
|
|
|
|
def _wiener(self, mag_out, mix_stft, niters):
|
|
raise NotImplementedError("non-CaC Wiener Demucs is not supported by the dependency-free path")
|
|
|
|
def valid_length(self, length: int):
|
|
if not self.use_train_segment:
|
|
return length
|
|
training_length = int(self.segment * self.samplerate)
|
|
if training_length < length:
|
|
raise ValueError(f"Given length {length} is longer than training length {training_length}")
|
|
return training_length
|
|
|
|
def cac2cws(self, x):
|
|
k = self.num_subbands
|
|
b, c, f, t = x.shape
|
|
return x.reshape(b, c * k, f // k, t)
|
|
|
|
def cws2cac(self, x):
|
|
k = self.num_subbands
|
|
b, c, f, t = x.shape
|
|
return x.reshape(b, c // k, f * k, t)
|
|
|
|
def forward(self, mix):
|
|
if self._use_mlx_full_forward(mix):
|
|
try:
|
|
from .demucs_mlx import mlx_forward_demucs
|
|
|
|
return mlx_forward_demucs(self, mix, self.mps_model_compute_dtype)
|
|
except Exception as exc:
|
|
self._pymss_mlx_full_backend_error = repr(exc)
|
|
self.mps_model_backend = "torch"
|
|
length = mix.shape[-1]
|
|
length_pre_pad = None
|
|
if self.use_train_segment:
|
|
if self.training:
|
|
self.segment = Fraction(mix.shape[-1], self.samplerate)
|
|
else:
|
|
training_length = int(self.segment * self.samplerate)
|
|
if mix.shape[-1] < training_length:
|
|
length_pre_pad = mix.shape[-1]
|
|
mix = F.pad(mix, (0, training_length - length_pre_pad))
|
|
z = self._spec(mix)
|
|
mag = self._magnitude(z)
|
|
x = mag
|
|
|
|
if self.num_subbands > 1:
|
|
x = self.cac2cws(x)
|
|
|
|
B, C, Fq, T = x.shape
|
|
|
|
mean = x.mean(dim=(1, 2, 3), keepdim=True)
|
|
std = x.std(dim=(1, 2, 3), keepdim=True)
|
|
x = (x - mean) / (1e-5 + std)
|
|
|
|
xt = mix
|
|
meant = xt.mean(dim=(1, 2), keepdim=True)
|
|
stdt = xt.std(dim=(1, 2), keepdim=True)
|
|
xt = (xt - meant) / (1e-5 + stdt)
|
|
|
|
saved, saved_t, lengths_t = [], [], []
|
|
for idx, encode in enumerate(self.encoder):
|
|
skip_length = x.shape[-1]
|
|
inject = None
|
|
if idx < len(self.tencoder):
|
|
lengths_t.append(xt.shape[-1])
|
|
tenc = self.tencoder[idx]
|
|
xt = tenc(xt)
|
|
if not tenc.empty:
|
|
saved_t.append(xt)
|
|
else:
|
|
inject = xt
|
|
x = encode(x, inject)
|
|
if idx == 0 and self.freq_emb is not None:
|
|
frs = torch.arange(x.shape[-2], device=x.device)
|
|
emb = self.freq_emb(frs).t()[None, :, :, None].expand_as(x)
|
|
x = x + self.freq_emb_scale * emb
|
|
|
|
saved.append((x, skip_length))
|
|
if self.crosstransformer:
|
|
if self.bottom_channels:
|
|
b, c, f, t = x.shape
|
|
x = self.channel_upsampler(x.reshape(b, c, f * t)).reshape(b, -1, f, t)
|
|
xt = self.channel_upsampler_t(xt)
|
|
|
|
x, xt = self.crosstransformer(x, xt)
|
|
|
|
if self.bottom_channels:
|
|
b, c, f, t = x.shape
|
|
x = self.channel_downsampler(x.reshape(b, c, f * t)).reshape(b, -1, f, t)
|
|
xt = self.channel_downsampler_t(xt)
|
|
|
|
for idx, decode in enumerate(self.decoder):
|
|
skip, skip_length = saved.pop(-1)
|
|
x, pre = decode(x, skip, skip_length)
|
|
|
|
offset = self.depth - len(self.tdecoder)
|
|
if idx >= offset:
|
|
tdec = self.tdecoder[idx - offset]
|
|
length_t = lengths_t.pop(-1)
|
|
if tdec.empty:
|
|
assert pre.shape[2] == 1, pre.shape
|
|
pre = pre[:, :, 0]
|
|
xt, _ = tdec(pre, None, length_t)
|
|
else:
|
|
skip = saved_t.pop(-1)
|
|
xt, _ = tdec(xt, skip, length_t)
|
|
|
|
assert len(saved) == 0
|
|
assert len(lengths_t) == 0
|
|
assert len(saved_t) == 0
|
|
|
|
S = len(self.sources)
|
|
|
|
if self.num_subbands > 1:
|
|
x = self.cws2cac(x.view(B, -1, Fq, T))
|
|
|
|
x = x.view(B, S, -1, Fq * self.num_subbands, T)
|
|
x = x * std[:, None] + mean[:, None]
|
|
|
|
zout = self._mask(z, x)
|
|
if self.use_train_segment:
|
|
if self.training:
|
|
x = self._ispec(zout, length)
|
|
else:
|
|
x = self._ispec(zout, training_length)
|
|
else:
|
|
x = self._ispec(zout, length)
|
|
|
|
xt = xt.view(B, S, -1, length if not self.use_train_segment or self.training else training_length)
|
|
xt = xt * stdt[:, None] + meant[:, None]
|
|
x = xt + x
|
|
if length_pre_pad:
|
|
x = x[..., :length_pre_pad]
|
|
return x
|
|
|
|
|
|
def get_model(args):
|
|
extra = {
|
|
"sources": list(args.training.instruments),
|
|
"audio_channels": args.training.channels,
|
|
"samplerate": args.training.samplerate,
|
|
"segment": args.training.segment,
|
|
}
|
|
if args.model != "htdemucs":
|
|
raise ValueError(f"Only htdemucs configs are supported, got {args.model!r}")
|
|
kw = to_plain(getattr(args, args.model))
|
|
return HTDemucs(**extra, **kw)
|