| import os |
| import sys |
| import torch |
|
|
| import numpy as np |
| import torch.nn as nn |
| import onnxruntime as ort |
| import torch.nn.functional as F |
|
|
| from einops import rearrange |
| from torch.nn.utils.parametrizations import weight_norm |
|
|
| sys.path.append(os.getcwd()) |
| os.environ["LRU_CACHE_CAPACITY"] = "3" |
|
|
| from infer.lib.predictors.FCPE.wav2mel import Wav2Mel |
| from infer.lib.predictors.FCPE.encoder import EncoderLayer, ConformerNaiveEncoder |
| from infer.lib.predictors.FCPE.utils import batch_interp_with_replacement_detach, decrypt_model, DotDict |
|
|
| @torch.no_grad() |
| def cent_to_f0(cent): |
| return 10 * 2 ** (cent / 1200) |
|
|
| @torch.no_grad() |
| def f0_to_cent(f0): |
| return 1200 * (f0 / 10).log2() |
|
|
| @torch.no_grad() |
| def latent2cents_decoder(cent_table, y, threshold = 0.05, mask = True): |
| if str(y.device).startswith("privateuseone"): |
| cent_table = cent_table.cpu() |
| y = y.cpu() |
|
|
| B, N, _ = y.size() |
| ci = cent_table[None, None, :].expand(B, N, -1) |
| rtn = (ci * y).sum(dim=-1, keepdim=True) / y.sum(dim=-1, keepdim=True) |
|
|
| if mask: |
| confident = y.max(dim=-1, keepdim=True)[0] |
| confident_mask = torch.ones_like(confident) |
| confident_mask[confident <= threshold] = float("-INF") |
| rtn = rtn * confident_mask |
|
|
| return rtn |
|
|
| @torch.no_grad() |
| def latent2cents_local_decoder(cent_table, out_dims, y, threshold = 0.05, mask = True): |
| if str(y.device).startswith("privateuseone"): |
| cent_table = cent_table.cpu() |
| y = y.cpu() |
|
|
| B, N, _ = y.size() |
| ci = cent_table[None, None, :].expand(B, N, -1) |
| confident, max_index = y.max(dim=-1, keepdim=True) |
|
|
| local_argmax_index = torch.arange(0, 9).to(max_index.device) + (max_index - 4) |
| local_argmax_index[local_argmax_index < 0] = 0 |
| local_argmax_index[local_argmax_index >= out_dims] = out_dims - 1 |
|
|
| y_l = y.gather(-1, local_argmax_index) |
| rtn = (ci.gather(-1, local_argmax_index) * y_l).sum(dim=-1, keepdim=True) / y_l.sum(dim=-1, keepdim=True) |
|
|
| if mask: |
| confident_mask = torch.ones_like(confident) |
| confident_mask[confident <= threshold] = float("-INF") |
| rtn = rtn * confident_mask |
|
|
| return rtn |
|
|
| def cents_decoder(cent_table, y, confidence, threshold = 0.05, mask=True): |
| if str(y.device).startswith("privateuseone"): |
| cent_table = cent_table.cpu() |
| y = y.cpu() |
|
|
| B, N, _ = y.size() |
| rtn = (cent_table[None, None, :].expand(B, N, -1) * y).sum(dim=-1, keepdim=True) / y.sum(dim=-1, keepdim=True) |
|
|
| if mask: |
| confident = y.max(dim=-1, keepdim=True)[0] |
| confident_mask = torch.ones_like(confident) |
| confident_mask[confident <= threshold] = float("-INF") |
| rtn = rtn * confident_mask |
|
|
| return (rtn, confident) if confidence else rtn |
|
|
| def cents_local_decoder(cent_table, y, n_out, confidence, threshold = 0.05, mask=True): |
| if str(y.device).startswith("privateuseone"): |
| cent_table = cent_table.cpu() |
| y = y.cpu() |
|
|
| B, N, _ = y.size() |
| confident, max_index = y.max(dim=-1, keepdim=True) |
| local_argmax_index = (torch.arange(0, 9).to(max_index.device) + (max_index - 4)).clamp(0, n_out - 1) |
| y_l = y.gather(-1, local_argmax_index) |
| rtn = (cent_table[None, None, :].expand(B, N, -1).gather(-1, local_argmax_index) * y_l).sum(dim=-1, keepdim=True) / y_l.sum(dim=-1, keepdim=True) |
|
|
| if mask: |
| confident_mask = torch.ones_like(confident) |
| confident_mask[confident <= threshold] = float("-INF") |
| rtn = rtn * confident_mask |
|
|
| return (rtn, confident) if confidence else rtn |
|
|
| class PCmer(nn.Module): |
| def __init__( |
| self, |
| num_layers, |
| num_heads, |
| dim_model, |
| dim_keys, |
| dim_values, |
| residual_dropout, |
| attention_dropout |
| ): |
| super().__init__() |
| self.num_layers = num_layers |
| self.num_heads = num_heads |
| self.dim_model = dim_model |
| self.dim_values = dim_values |
| self.dim_keys = dim_keys |
| self.residual_dropout = residual_dropout |
| self.attention_dropout = attention_dropout |
| self._layers = nn.ModuleList([EncoderLayer(self) for _ in range(num_layers)]) |
|
|
| def forward(self, phone, mask=None): |
| for layer in self._layers: |
| phone = layer(phone, mask) |
|
|
| return phone |
|
|
| class CFNaiveMelPE(nn.Module): |
| def __init__( |
| self, |
| input_channels, |
| out_dims, |
| hidden_dims = 512, |
| n_layers = 6, |
| n_heads = 8, |
| f0_max = 1975.5, |
| f0_min = 32.70, |
| use_fa_norm = False, |
| conv_only = False, |
| conv_dropout = 0, |
| atten_dropout = 0, |
| use_harmonic_emb = False |
| ): |
| super().__init__() |
| self.input_channels = input_channels |
| self.out_dims = out_dims |
| self.hidden_dims = hidden_dims |
| self.n_layers = n_layers |
| self.n_heads = n_heads |
| self.f0_max = f0_max |
| self.f0_min = f0_min |
| self.use_fa_norm = use_fa_norm |
|
|
| self.harmonic_emb = nn.Embedding(9, hidden_dims) if use_harmonic_emb else None |
| self.input_stack = nn.Sequential( |
| nn.Conv1d( |
| input_channels, |
| hidden_dims, |
| 3, |
| 1, |
| 1 |
| ), |
| nn.GroupNorm( |
| 4, |
| hidden_dims |
| ), |
| nn.LeakyReLU(), |
| nn.Conv1d( |
| hidden_dims, |
| hidden_dims, |
| 3, |
| 1, |
| 1 |
| ) |
| ) |
| self.net = ConformerNaiveEncoder( |
| num_layers=n_layers, |
| num_heads=n_heads, |
| dim_model=hidden_dims, |
| use_norm=use_fa_norm, |
| conv_only=conv_only, |
| conv_dropout=conv_dropout, |
| atten_dropout=atten_dropout |
| ) |
| self.norm = nn.LayerNorm(hidden_dims) |
| self.output_proj = weight_norm( |
| nn.Linear( |
| hidden_dims, |
| out_dims |
| ) |
| ) |
|
|
| self.cent_table_b = torch.linspace( |
| f0_to_cent(torch.Tensor([f0_min]))[0], |
| f0_to_cent(torch.Tensor([f0_max]))[0], |
| out_dims |
| ).detach() |
| self.gaussian_blurred_cent_mask_b = ( |
| 1200 * torch.Tensor([self.f0_max / 10.]).log2() |
| )[0].detach() |
|
|
| self.register_buffer("cent_table", self.cent_table_b) |
| self.register_buffer("gaussian_blurred_cent_mask", self.gaussian_blurred_cent_mask_b) |
|
|
| def forward(self, x, _h_emb=None): |
| x = self.input_stack(x.transpose(-1, -2)).transpose(-1, -2) |
|
|
| if self.harmonic_emb is not None: |
| if _h_emb is None: |
| x += self.harmonic_emb(torch.LongTensor([0]).to(x.device)) |
| else: |
| x += self.harmonic_emb(torch.LongTensor([int(_h_emb)]).to(x.device)) |
|
|
| return self.output_proj(self.norm(self.net(x))).sigmoid() |
|
|
| @torch.no_grad() |
| def infer(self, mel, decoder = "local_argmax", threshold = 0.05): |
| latent = self.forward(mel) |
|
|
| return cent_to_f0( |
| ( |
| latent2cents_decoder( |
| self.cent_table, |
| latent, |
| threshold=threshold |
| ) |
| ) if decoder == "argmax" else ( |
| latent2cents_local_decoder( |
| self.cent_table, |
| self.out_dims, |
| latent, |
| threshold=threshold |
| ) |
| ) |
| ) |
|
|
| class FCPE_LEGACY(nn.Module): |
| def __init__( |
| self, |
| input_channel=128, |
| out_dims=360, |
| n_layers=12, |
| n_chans=512, |
| f0_max=1975.5, |
| f0_min=32.70, |
| confidence=False, |
| threshold=0.05, |
| use_input_conv=True |
| ): |
| super().__init__() |
| self.n_out = out_dims |
| self.f0_max = f0_max |
| self.f0_min = f0_min |
| self.confidence = confidence |
| self.threshold = threshold |
| self.use_input_conv = use_input_conv |
|
|
| self.cent_table_b = torch.Tensor( |
| np.linspace( |
| f0_to_cent(torch.Tensor([f0_min]))[0], |
| f0_to_cent(torch.Tensor([f0_max]))[0], |
| out_dims |
| ) |
| ) |
| self.register_buffer("cent_table", self.cent_table_b) |
|
|
| self.stack = nn.Sequential( |
| nn.Conv1d( |
| input_channel, |
| n_chans, |
| 3, |
| 1, |
| 1 |
| ), |
| nn.GroupNorm( |
| 4, |
| n_chans |
| ), |
| nn.LeakyReLU(), |
| nn.Conv1d( |
| n_chans, |
| n_chans, |
| 3, |
| 1, |
| 1 |
| ) |
| ) |
| self.decoder = PCmer( |
| num_layers=n_layers, |
| num_heads=8, |
| dim_model=n_chans, |
| dim_keys=n_chans, |
| dim_values=n_chans, |
| residual_dropout=0.1, |
| attention_dropout=0.1 |
| ) |
| self.norm = nn.LayerNorm(n_chans) |
| self.dense_out = weight_norm( |
| nn.Linear( |
| n_chans, |
| self.n_out |
| ) |
| ) |
|
|
| def forward(self, mel, return_hz_f0=False, cdecoder="local_argmax", output_interp_target_length=None): |
| x = self.decoder(self.stack(mel.transpose(1, 2)).transpose(1, 2) if self.use_input_conv else mel) |
| x = self.dense_out(self.norm(x)).sigmoid() |
|
|
| x = cent_to_f0( |
| ( |
| cents_decoder( |
| self.cent_table, |
| x, |
| self.confidence, |
| threshold=self.threshold, |
| mask=True |
| ) |
| ) if cdecoder == "argmax" else ( |
| cents_local_decoder( |
| self.cent_table, |
| x, |
| self.n_out, |
| self.confidence, |
| threshold=self.threshold, |
| mask=True |
| ) |
| ) |
| ) |
|
|
| x = (1 + x / 700).log() if not return_hz_f0 else x |
|
|
| if output_interp_target_length is not None: |
| x = F.interpolate( |
| torch.where(x == 0, float("nan"), x).transpose(1, 2), |
| size=int(output_interp_target_length), |
| mode="linear" |
| ).transpose(1, 2) |
|
|
| x = torch.where(x.isnan(), float(0.0), x) |
|
|
| return x |
|
|
| def gaussian_blurred_cent(self, cents): |
| B, N, _ = cents.size() |
|
|
| return ( |
| -(self.cent_table[None, None, :].expand(B, N, -1) - cents).square() / 1250 |
| ).exp() * (cents > 0.1) & ( |
| cents < (1200.0 * np.log2(self.f0_max / 10.0)) |
| ).float() |
|
|
| class InferCFNaiveMelPE(torch.nn.Module): |
| def __init__( |
| self, |
| args, |
| state_dict |
| ): |
| super().__init__() |
| self.model = CFNaiveMelPE( |
| input_channels=args.mel.num_mels, |
| out_dims=args.model.out_dims, |
| hidden_dims=args.model.hidden_dims, |
| n_layers=args.model.n_layers, |
| n_heads=args.model.n_heads, |
| f0_max=args.model.f0_max, |
| f0_min=args.model.f0_min, |
| use_fa_norm=args.model.use_fa_norm, |
| conv_only=args.model.conv_only, |
| conv_dropout=args.model.conv_dropout, |
| atten_dropout=args.model.atten_dropout, |
| use_harmonic_emb=False |
| ) |
| self.model.load_state_dict(state_dict) |
| self.model.eval() |
| self.register_buffer("tensor_device_marker", torch.tensor(1.0).float(), persistent=False) |
|
|
| def forward(self, mel, decoder_mode = "local_argmax", threshold = 0.006): |
| with torch.no_grad(): |
| mels = rearrange(torch.stack([mel], -1), "B T C K -> (B K) T C") |
| f0s = rearrange(self.model.infer(mels, decoder=decoder_mode, threshold=threshold), "(B K) T 1 -> B T (K 1)", K=1) |
|
|
| return f0s |
|
|
| def infer( |
| self, |
| mel, |
| decoder_mode = "local_argmax", |
| threshold = 0.006, |
| f0_min = None, |
| f0_max = None, |
| interp_uv = False, |
| output_interp_target_length = None, |
| return_uv = False |
| ): |
| f0 = self.__call__(mel, decoder_mode, threshold) |
| f0_for_uv = f0 |
|
|
| uv = (f0_for_uv < f0_min).type(f0_for_uv.dtype) |
| f0 = f0 * (1 - uv) |
|
|
| if interp_uv: |
| f0 = batch_interp_with_replacement_detach( |
| uv.squeeze(-1).bool(), |
| f0.squeeze(-1) |
| ).unsqueeze(-1) |
|
|
| if f0_max is not None: f0[f0 > f0_max] = f0_max |
|
|
| if output_interp_target_length is not None: |
| f0 = F.interpolate( |
| torch.where(f0 == 0, float("nan"), f0).transpose(1, 2), |
| size=int(output_interp_target_length), |
| mode="linear" |
| ).transpose(1, 2) |
|
|
| f0 = torch.where(f0.isnan(), float(0.0), f0) |
|
|
| if return_uv: return f0, F.interpolate(uv.transpose(1, 2), size=int(output_interp_target_length), mode="nearest").transpose(1, 2) |
| else: return f0 |
|
|
| class FCPEInfer_LEGACY: |
| def __init__( |
| self, |
| configs, |
| model_path, |
| device=None, |
| dtype=torch.float32, |
| providers=None, |
| onnx=False, |
| f0_min=50, |
| f0_max=1100 |
| ): |
| if device is None: device = "cuda" if torch.cuda.is_available() else "cpu" |
| self.device = device |
| self.dtype = dtype |
| self.onnx = onnx |
| self.f0_min = f0_min |
| self.f0_max = f0_max |
| self.wav2mel = Wav2Mel(device=self.device, dtype=self.dtype) |
|
|
| if self.onnx: |
| sess_options = ort.SessionOptions() |
| sess_options.log_severity_level = 3 |
| self.model = ort.InferenceSession(decrypt_model(configs, model_path), sess_options=sess_options, providers=providers) |
| else: |
| ckpt = torch.load(model_path, map_location="cpu", weights_only=True) |
| self.args = DotDict(ckpt["config"]) |
| model = FCPE_LEGACY( |
| input_channel=self.args.model.input_channel, |
| out_dims=self.args.model.out_dims, |
| n_layers=self.args.model.n_layers, |
| n_chans=self.args.model.n_chans, |
| f0_max=self.f0_max, |
| f0_min=self.f0_min, |
| confidence=self.args.model.confidence |
| ) |
| model.to(self.device).to(self.dtype) |
| model.load_state_dict(ckpt["model"]) |
| model.eval() |
| self.model = model |
|
|
| @torch.no_grad() |
| def __call__(self, audio, sr, threshold=0.05, p_len=None): |
| if not self.onnx: self.model.threshold = threshold |
| if not hasattr(self, "numpy_threshold") and self.onnx: self.numpy_threshold = np.array(threshold, dtype=np.float32) |
|
|
| mel = self.wav2mel(audio=audio[None, :], sample_rate=sr).to(self.dtype) |
|
|
| if self.onnx: |
| return torch.as_tensor( |
| self.model.run( |
| [self.model.get_outputs()[0].name], |
| { |
| self.model.get_inputs()[0].name: mel.detach().cpu().numpy(), |
| self.model.get_inputs()[1].name: self.numpy_threshold |
| } |
| )[0], |
| dtype=self.dtype, |
| device=self.device |
| ) |
| else: |
| return self.model( |
| mel=mel, |
| return_hz_f0=True, |
| output_interp_target_length=p_len |
| ) |
|
|
| class FCPEInfer: |
| def __init__( |
| self, |
| configs, |
| model_path, |
| device=None, |
| dtype=torch.float32, |
| providers=None, |
| onnx=False, |
| f0_min=50, |
| f0_max=1100 |
| ): |
| if device is None: device = "cuda" if torch.cuda.is_available() else "cpu" |
| self.device = device |
| self.dtype = dtype |
| self.onnx = onnx |
| self.f0_min = f0_min |
| self.f0_max = f0_max |
| self.wav2mel = Wav2Mel(device=self.device, dtype=self.dtype) |
|
|
| if self.onnx: |
| sess_options = ort.SessionOptions() |
| sess_options.log_severity_level = 3 |
| self.model = ort.InferenceSession(decrypt_model(configs, model_path), sess_options=sess_options, providers=providers) |
| else: |
| ckpt = torch.load(model_path, map_location="cpu", weights_only=True) |
| ckpt["config_dict"]["model"]["conv_dropout"] = ckpt["config_dict"]["model"]["atten_dropout"] = 0.0 |
| self.args = DotDict(ckpt["config_dict"]) |
| model = InferCFNaiveMelPE(self.args, ckpt["model"]) |
| self.model = model.to(device).to(self.dtype).eval() |
|
|
| @torch.no_grad() |
| def __call__(self, audio, sr, threshold=0.05, p_len=None): |
| if not hasattr(self, "numpy_threshold") and self.onnx: self.numpy_threshold = np.array(threshold, dtype=np.float32) |
| mel = self.wav2mel(audio=audio[None, :], sample_rate=sr).to(self.dtype) |
|
|
| if self.onnx: |
| return torch.as_tensor( |
| self.model.run( |
| [self.model.get_outputs()[0].name], |
| { |
| self.model.get_inputs()[0].name: mel.detach().cpu().numpy(), |
| self.model.get_inputs()[1].name: self.numpy_threshold |
| } |
| )[0], |
| dtype=self.dtype, |
| device=self.device |
| ) |
| else: |
| return self.model.infer( |
| mel, |
| threshold=threshold, |
| f0_min=self.f0_min, |
| f0_max=self.f0_max, |
| output_interp_target_length=p_len |
| ) |
|
|
| class FCPE: |
| def __init__( |
| self, |
| configs, |
| model_path, |
| hop_length=512, |
| f0_min=50, |
| f0_max=1100, |
| dtype=torch.float32, |
| device=None, |
| sample_rate=16000, |
| threshold=0.05, |
| providers=None, |
| onnx=False, |
| legacy=False |
| ): |
| self.model = FCPEInfer_LEGACY if legacy else FCPEInfer |
| self.fcpe = self.model(configs, model_path, device=device, dtype=dtype, providers=providers, onnx=onnx, f0_min=f0_min, f0_max=f0_max) |
| self.hop_length = hop_length |
| self.device = device or ("cuda" if torch.cuda.is_available() else "cpu") |
| self.threshold = threshold |
| self.sample_rate = sample_rate |
| self.dtype = dtype |
| self.legacy = legacy |
|
|
| def compute_f0(self, wav, p_len=None): |
| x = torch.FloatTensor(wav).to(self.dtype).to(self.device) |
| p_len = (x.shape[0] // self.hop_length) if p_len is None else p_len |
|
|
| f0 = self.fcpe(x, sr=self.sample_rate, threshold=self.threshold, p_len=p_len) |
| f0 = f0[:] if f0.dim() == 1 else f0[0, :, 0] |
|
|
| if torch.all(f0 == 0): return f0.cpu().numpy() if p_len is None else np.zeros(p_len) |
| return f0.cpu().numpy() |
