4.1

data_utils.py

@@ -7,7 +7,7 @@ import torch.utils.data
 
 import modules.commons as commons
 import utils
-from modules.mel_processing import spectrogram_torch, spec_to_mel_torch
+from modules.mel_processing import spectrogram_torch, spec_to_mel_torch, spectrogram_torch
 from utils import load_wav_to_torch, load_filepaths_and_text
 
 # import h5py
@@ -23,8 +23,9 @@ class TextAudioSpeakerLoader(torch.utils.data.Dataset):
         3) computes spectrograms from audio files.
     """
 
-    def __init__(self, audiopaths, hparams, all_in_mem: bool = False):
+    def __init__(self, audiopaths, hparams, all_in_mem: bool = False, vol_aug: bool = True):
         self.audiopaths = load_filepaths_and_text(audiopaths)
+        self.hparams = hparams
         self.max_wav_value = hparams.data.max_wav_value
         self.sampling_rate = hparams.data.sampling_rate
         self.filter_length = hparams.data.filter_length
@@ -34,7 +35,8 @@ class TextAudioSpeakerLoader(torch.utils.data.Dataset):
         self.use_sr = hparams.train.use_sr
         self.spec_len = hparams.train.max_speclen
         self.spk_map = hparams.spk
-
+        self.vol_emb = hparams.model.vol_embedding
+        self.vol_aug = hparams.train.vol_aug and vol_aug
         random.seed(1234)
         random.shuffle(self.audiopaths)
         
@@ -65,34 +67,55 @@ class TextAudioSpeakerLoader(torch.utils.data.Dataset):
         spk = filename.split("/")[-2]
         spk = torch.LongTensor([self.spk_map[spk]])
 
-        f0 = np.load(filename + ".f0.npy")
-        f0, uv = utils.interpolate_f0(f0)
-        f0 = torch.FloatTensor(f0)
-        uv = torch.FloatTensor(uv)
+        f0, uv = np.load(filename + ".f0.npy",allow_pickle=True)
+        
+        f0 = torch.FloatTensor(np.array(f0,dtype=float))
+        uv = torch.FloatTensor(np.array(uv,dtype=float))
 
         c = torch.load(filename+ ".soft.pt")
         c = utils.repeat_expand_2d(c.squeeze(0), f0.shape[0])
-
+        if self.vol_emb:
+            volume_path = filename + ".vol.npy"
+            volume = np.load(volume_path)
+            volume = torch.from_numpy(volume).float()
+        else:
+            volume = None
 
         lmin = min(c.size(-1), spec.size(-1))
         assert abs(c.size(-1) - spec.size(-1)) < 3, (c.size(-1), spec.size(-1), f0.shape, filename)
         assert abs(audio_norm.shape[1]-lmin * self.hop_length) < 3 * self.hop_length
         spec, c, f0, uv = spec[:, :lmin], c[:, :lmin], f0[:lmin], uv[:lmin]
         audio_norm = audio_norm[:, :lmin * self.hop_length]
+        if volume!= None:
+            volume = volume[:lmin]
+        return c, f0, spec, audio_norm, spk, uv, volume
 
-        return c, f0, spec, audio_norm, spk, uv
-
-    def random_slice(self, c, f0, spec, audio_norm, spk, uv):
+    def random_slice(self, c, f0, spec, audio_norm, spk, uv, volume):
         # if spec.shape[1] < 30:
         #     print("skip too short audio:", filename)
         #     return None
+
+        if random.choice([True, False]) and self.vol_aug and volume!=None:
+            max_amp = float(torch.max(torch.abs(audio_norm))) + 1e-5
+            max_shift = min(1, np.log10(1/max_amp))
+            log10_vol_shift = random.uniform(-1, max_shift)
+            audio_norm = audio_norm * (10 ** log10_vol_shift)
+            volume = volume * (10 ** log10_vol_shift)
+            spec = spectrogram_torch(audio_norm,
+            self.hparams.data.filter_length,
+            self.hparams.data.sampling_rate,
+            self.hparams.data.hop_length,
+            self.hparams.data.win_length,
+            center=False)[0]
+
         if spec.shape[1] > 800:
             start = random.randint(0, spec.shape[1]-800)
             end = start + 790
             spec, c, f0, uv = spec[:, start:end], c[:, start:end], f0[start:end], uv[start:end]
             audio_norm = audio_norm[:, start * self.hop_length : end * self.hop_length]
-
-        return c, f0, spec, audio_norm, spk, uv
+            if volume !=None:
+                volume = volume[start:end]
+        return c, f0, spec, audio_norm, spk, uv,volume
 
     def __getitem__(self, index):
         if self.all_in_mem:
@@ -124,12 +147,14 @@ class TextAudioCollate:
         wav_padded = torch.FloatTensor(len(batch), 1, max_wav_len)
         spkids = torch.LongTensor(len(batch), 1)
         uv_padded = torch.FloatTensor(len(batch), max_c_len)
+        volume_padded = torch.FloatTensor(len(batch), max_c_len)
 
         c_padded.zero_()
         spec_padded.zero_()
         f0_padded.zero_()
         wav_padded.zero_()
         uv_padded.zero_()
+        volume_padded.zero_()
 
         for i in range(len(ids_sorted_decreasing)):
             row = batch[ids_sorted_decreasing[i]]
@@ -151,5 +176,9 @@ class TextAudioCollate:
 
             uv = row[5]
             uv_padded[i, :uv.size(0)] = uv
-
-        return c_padded, f0_padded, spec_padded, wav_padded, spkids, lengths, uv_padded
+            volume = row[6]
+            if volume != None:
+                volume_padded[i, :volume.size(0)] = volume
+            else :
+                volume_padded = None
+        return c_padded, f0_padded, spec_padded, wav_padded, spkids, lengths, uv_padded, volume_padded

inference_main.py

@@ -2,12 +2,11 @@ import io
 import logging
 import time
 from pathlib import Path
-
+from spkmix import spk_mix_map
 import librosa
 import matplotlib.pyplot as plt
 import numpy as np
 import soundfile
-
 from inference import infer_tool
 from inference import slicer
 from inference.infer_tool import Svc
@@ -16,39 +15,84 @@ logging.getLogger('numba').setLevel(logging.WARNING)
 chunks_dict = infer_tool.read_temp("inference/chunks_temp.json")
 
 
-
 def main():
     import argparse
 
     parser = argparse.ArgumentParser(description='sovits4 inference')
 
     # 一定要设置的部分
-    parser.add_argument('-m', '--model_path', type=str, default="logs/44k/G_0.pth", help='模型路径')
-    parser.add_argument('-c', '--config_path', type=str, default="configs/config.json", help='配置文件路径')
+    parser.add_argument('-m', '--model_path', type=str, default="logs/44k/", help='模型路径')
+    parser.add_argument('-c', '--config_path', type=str, default="configs/", help='配置文件路径')
     parser.add_argument('-cl', '--clip', type=float, default=0, help='音频强制切片，默认0为自动切片，单位为秒/s')
-    parser.add_argument('-n', '--clean_names', type=str, nargs='+', default=["君の知らない物語-src.wav"], help='wav文件名列表，放在raw文件夹下')
+    parser.add_argument('-n', '--clean_names', type=str, nargs='+', default=["test.wav"],
+                        help='wav文件名列表，放在raw文件夹下')
     parser.add_argument('-t', '--trans', type=int, nargs='+', default=[0], help='音高调整，支持正负（半音）')
-    parser.add_argument('-s', '--spk_list', type=str, nargs='+', default=['nen'], help='合成目标说话人名称')
+    parser.add_argument('-s', '--spk_list', type=str, nargs='+', default=['buyizi'], help='合成目标说话人名称')
 
     # 可选项部分
-    parser.add_argument('-a', '--auto_predict_f0', action='store_true', default=False,help='语音转换自动预测音高，转换歌声时不要打开这个会严重跑调')
-    parser.add_argument('-cm', '--cluster_model_path', type=str, default="logs/44k/kmeans_10000.pt", help='聚类模型路径，如果没有训练聚类则随便填')
-    parser.add_argument('-cr', '--cluster_infer_ratio', type=float, default=0, help='聚类方案占比，范围0-1，若没有训练聚类模型则默认0即可')
-    parser.add_argument('-lg', '--linear_gradient', type=float, default=0, help='两段音频切片的交叉淡入长度，如果强制切片后出现人声不连贯可调整该数值，如果连贯建议采用默认值0，单位为秒')
-    parser.add_argument('-fmp', '--f0_mean_pooling', type=bool, default=False, help='是否对F0使用均值滤波器(池化)，对部分哑音有改善。注意，启动该选项会导致推理速度下降，默认关闭')
+    parser.add_argument('-a', '--auto_predict_f0', action='store_true', default=False,
+                        help='语音转换自动预测音高，转换歌声时不要打开这个会严重跑调')
+    parser.add_argument('-cm', '--cluster_model_path', type=str, default="logs/44k/kmeans_10000.pt",
+                        help='聚类模型或特征检索索引路径，如果没有训练聚类或特征检索则随便填')
+    parser.add_argument('-cr', '--cluster_infer_ratio', type=float, default=0,
+                        help='聚类方案或特征检索占比，范围0-1，若没有训练聚类模型或特征检索则默认0即可')
+    parser.add_argument('-lg', '--linear_gradient', type=float, default=0,
+                        help='两段音频切片的交叉淡入长度，如果强制切片后出现人声不连贯可调整该数值，如果连贯建议采用默认值0，单位为秒')
+    parser.add_argument('-f0p', '--f0_predictor', type=str, default="harvest",
+                        help='选择F0预测器,可选择crepe,pm,dio,harvest,默认为pm(注意：crepe为原F0使用均值滤波器)')
+    parser.add_argument('-eh', '--enhance', action='store_true', default=False,
+                        help='是否使用NSF_HIFIGAN增强器,该选项对部分训练集少的模型有一定的音质增强效果，但是对训练好的模型有反面效果，默认关闭')
+    parser.add_argument('-shd', '--shallow_diffusion', action='store_true', default=False,
+                        help='是否使用浅层扩散，使用后可解决一部分电音问题，默认关闭，该选项打开时，NSF_HIFIGAN增强器将会被禁止')
+    parser.add_argument('-usm', '--use_spk_mix', action='store_true', default=False, help='是否使用角色融合')
+    parser.add_argument('-lea', '--loudness_envelope_adjustment', type=float, default=1,
+                        help='输入源响度包络替换输出响度包络融合比例，越靠近1越使用输出响度包络')
+    parser.add_argument('-fr', '--feature_retrieval', action='store_true', default=False,
+                        help='是否使用特征检索，如果使用聚类模型将被禁用，且cm与cr参数将会变成特征检索的索引路径与混合比例')
+
+    # 浅扩散设置
+    parser.add_argument('-dm', '--diffusion_model_path', type=str, default="logs/44k/diffusion/model_0.pt",
+                        help='扩散模型路径')
+    parser.add_argument('-dc', '--diffusion_config_path', type=str, default="logs/44k/diffusion/config.yaml",
+                        help='扩散模型配置文件路径')
+    parser.add_argument('-ks', '--k_step', type=int, default=100, help='扩散步数，越大越接近扩散模型的结果，默认100')
+    parser.add_argument('-se', '--second_encoding', action='store_true', default=False,
+                        help='二次编码，浅扩散前会对原始音频进行二次编码，玄学选项，有时候效果好，有时候效果差')
+    parser.add_argument('-od', '--only_diffusion', action='store_true', default=False,
+                        help='纯扩散模式，该模式不会加载sovits模型，以扩散模型推理')
 
     # 不用动的部分
-    parser.add_argument('-sd', '--slice_db', type=int, default=-40, help='默认-40，嘈杂的音频可以-30，干声保留呼吸可以-50')
+    parser.add_argument('-sd', '--slice_db', type=int, default=-40,
+                        help='默认-40，嘈杂的音频可以-30，干声保留呼吸可以-50')
     parser.add_argument('-d', '--device', type=str, default=None, help='推理设备，None则为自动选择cpu和gpu')
     parser.add_argument('-ns', '--noice_scale', type=float, default=0.4, help='噪音级别，会影响咬字和音质，较为玄学')
-    parser.add_argument('-p', '--pad_seconds', type=float, default=0.5, help='推理音频pad秒数，由于未知原因开头结尾会有异响，pad一小段静音段后就不会出现')
+    parser.add_argument('-p', '--pad_seconds', type=float, default=0.5,
+                        help='推理音频pad秒数，由于未知原因开头结尾会有异响，pad一小段静音段后就不会出现')
     parser.add_argument('-wf', '--wav_format', type=str, default='flac', help='音频输出格式')
-    parser.add_argument('-lgr', '--linear_gradient_retain', type=float, default=0.75, help='自动音频切片后，需要舍弃每段切片的头尾。该参数设置交叉长度保留的比例，范围0-1,左开右闭')
+    parser.add_argument('-lgr', '--linear_gradient_retain', type=float, default=0.75,
+                        help='自动音频切片后，需要舍弃每段切片的头尾。该参数设置交叉长度保留的比例，范围0-1,左开右闭')
+    parser.add_argument('-eak', '--enhancer_adaptive_key', type=int, default=0,
+                        help='使增强器适应更高的音域(单位为半音数)|默认为0')
+    parser.add_argument('-ft', '--f0_filter_threshold', type=float, default=0.05,
+                        help='F0过滤阈值，只有使用crepe时有效. 数值范围从0-1. 降低该值可减少跑调概率，但会增加哑音')
+
+    def preprocess_args(args1):
+        spk1 = args1.spk_list[0]
+        args1.model_path += f"{spk1}.pth"
+        args1.config_path += f"config_{spk1}.json"
+        args1.clip = 30
+
+        if spk1 == 'tomori':
+            args1.feature_retrieval = True
+            args1.cluster_model_path = "logs/44k/tomori_index.pkl"
+            args1.cluster_infer_ratio = 0.5
+            args1.f0_predictor = 'crepe'
+
+        return args1
 
     args = parser.parse_args()
+    args = preprocess_args(args)
 
-    svc_model = Svc(args.model_path, args.config_path, args.device, args.cluster_model_path)
-    infer_tool.mkdir(["raw", "results"])
     clean_names = args.clean_names
     trans = args.trans
     spk_list = args.spk_list
@@ -61,7 +105,37 @@ def main():
     clip = args.clip
     lg = args.linear_gradient
     lgr = args.linear_gradient_retain
-    F0_mean_pooling = args.f0_mean_pooling
+    f0p = args.f0_predictor
+    enhance = args.enhance
+    enhancer_adaptive_key = args.enhancer_adaptive_key
+    cr_threshold = args.f0_filter_threshold
+    diffusion_model_path = args.diffusion_model_path
+    diffusion_config_path = args.diffusion_config_path
+    k_step = args.k_step
+    only_diffusion = args.only_diffusion
+    shallow_diffusion = args.shallow_diffusion
+    use_spk_mix = args.use_spk_mix
+    second_encoding = args.second_encoding
+    loudness_envelope_adjustment = args.loudness_envelope_adjustment
+
+    svc_model = Svc(args.model_path,
+                    args.config_path,
+                    args.device,
+                    args.cluster_model_path,
+                    enhance,
+                    diffusion_model_path,
+                    diffusion_config_path,
+                    shallow_diffusion,
+                    only_diffusion,
+                    use_spk_mix,
+                    args.feature_retrieval)
+
+    infer_tool.mkdir(["raw", "results"])
+
+    if len(spk_mix_map) <= 1:
+        use_spk_mix = False
+    if use_spk_mix:
+        spk_list = [spk_mix_map]
 
     infer_tool.fill_a_to_b(trans, clean_names)
     for clean_name, tran in zip(clean_names, trans):
@@ -69,62 +143,39 @@ def main():
         if "." not in raw_audio_path:
             raw_audio_path += ".wav"
         infer_tool.format_wav(raw_audio_path)
-        wav_path = Path(raw_audio_path).with_suffix('.wav')
-        chunks = slicer.cut(wav_path, db_thresh=slice_db)
-        audio_data, audio_sr = slicer.chunks2audio(wav_path, chunks)
-        per_size = int(clip*audio_sr)
-        lg_size = int(lg*audio_sr)
-        lg_size_r = int(lg_size*lgr)
-        lg_size_c_l = (lg_size-lg_size_r)//2
-        lg_size_c_r = lg_size-lg_size_r-lg_size_c_l
-        lg = np.linspace(0,1,lg_size_r) if lg_size!=0 else 0
-
         for spk in spk_list:
-            audio = []
-            for (slice_tag, data) in audio_data:
-                print(f'#=====segment start, {round(len(data) / audio_sr, 3)}s======')
-                
-                length = int(np.ceil(len(data) / audio_sr * svc_model.target_sample))
-                if slice_tag:
-                    print('jump empty segment')
-                    _audio = np.zeros(length)
-                    audio.extend(list(infer_tool.pad_array(_audio, length)))
-                    continue
-                if per_size != 0:
-                    datas = infer_tool.split_list_by_n(data, per_size,lg_size)
-                else:
-                    datas = [data]
-                for k,dat in enumerate(datas):
-                    per_length = int(np.ceil(len(dat) / audio_sr * svc_model.target_sample)) if clip!=0 else length
-                    if clip!=0: print(f'###=====segment clip start, {round(len(dat) / audio_sr, 3)}s======')
-                    # padd
-                    pad_len = int(audio_sr * pad_seconds)
-                    dat = np.concatenate([np.zeros([pad_len]), dat, np.zeros([pad_len])])
-                    raw_path = io.BytesIO()
-                    soundfile.write(raw_path, dat, audio_sr, format="wav")
-                    raw_path.seek(0)
-                    out_audio, out_sr = svc_model.infer(spk, tran, raw_path,
-                                                        cluster_infer_ratio=cluster_infer_ratio,
-                                                        auto_predict_f0=auto_predict_f0,
-                                                        noice_scale=noice_scale,
-                                                        F0_mean_pooling = F0_mean_pooling
-                                                        )
-                    _audio = out_audio.cpu().numpy()
-                    pad_len = int(svc_model.target_sample * pad_seconds)
-                    _audio = _audio[pad_len:-pad_len]
-                    _audio = infer_tool.pad_array(_audio, per_length)
-                    if lg_size!=0 and k!=0:
-                        lg1 = audio[-(lg_size_r+lg_size_c_r):-lg_size_c_r] if lgr != 1 else audio[-lg_size:]
-                        lg2 = _audio[lg_size_c_l:lg_size_c_l+lg_size_r]  if lgr != 1 else _audio[0:lg_size]
-                        lg_pre = lg1*(1-lg)+lg2*lg
-                        audio = audio[0:-(lg_size_r+lg_size_c_r)] if lgr != 1 else audio[0:-lg_size]
-                        audio.extend(lg_pre)
-                        _audio = _audio[lg_size_c_l+lg_size_r:] if lgr != 1 else _audio[lg_size:]
-                    audio.extend(list(_audio))
+            kwarg = {
+                "raw_audio_path": raw_audio_path,
+                "spk": spk,
+                "tran": tran,
+                "slice_db": slice_db,
+                "cluster_infer_ratio": cluster_infer_ratio,
+                "auto_predict_f0": auto_predict_f0,
+                "noice_scale": noice_scale,
+                "pad_seconds": pad_seconds,
+                "clip_seconds": clip,
+                "lg_num": lg,
+                "lgr_num": lgr,
+                "f0_predictor": f0p,
+                "enhancer_adaptive_key": enhancer_adaptive_key,
+                "cr_threshold": cr_threshold,
+                "k_step": k_step,
+                "use_spk_mix": use_spk_mix,
+                "second_encoding": second_encoding,
+                "loudness_envelope_adjustment": loudness_envelope_adjustment
+            }
+            audio = svc_model.slice_inference(**kwarg)
             key = "auto" if auto_predict_f0 else f"{tran}key"
             cluster_name = "" if cluster_infer_ratio == 0 else f"_{cluster_infer_ratio}"
-            res_path = f'./results/{clean_name}_{key}_{spk}{cluster_name}.{wav_format}'
+            isdiffusion = "sovits"
+            if shallow_diffusion: isdiffusion = "sovdiff"
+            if only_diffusion: isdiffusion = "diff"
+            if use_spk_mix:
+                spk = "spk_mix"
+            res_path = f'results/{clean_name}_{key}_{spk}{cluster_name}_{isdiffusion}.{wav_format}'
             soundfile.write(res_path, audio, svc_model.target_sample, format=wav_format)
+            svc_model.clear_empty()
+
 
 if __name__ == '__main__':
     main()

models.py

@@ -13,111 +13,111 @@ from torch.nn.utils import weight_norm, remove_weight_norm, spectral_norm
 
 import utils
 from modules.commons import init_weights, get_padding
-from vdecoder.hifigan.models import Generator
 from utils import f0_to_coarse
 
 class ResidualCouplingBlock(nn.Module):
-  def __init__(self,
-      channels,
-      hidden_channels,
-      kernel_size,
-      dilation_rate,
-      n_layers,
-      n_flows=4,
-      gin_channels=0):
-    super().__init__()
-    self.channels = channels
-    self.hidden_channels = hidden_channels
-    self.kernel_size = kernel_size
-    self.dilation_rate = dilation_rate
-    self.n_layers = n_layers
-    self.n_flows = n_flows
-    self.gin_channels = gin_channels
-
-    self.flows = nn.ModuleList()
-    for i in range(n_flows):
-      self.flows.append(modules.ResidualCouplingLayer(channels, hidden_channels, kernel_size, dilation_rate, n_layers, gin_channels=gin_channels, mean_only=True))
-      self.flows.append(modules.Flip())
-
-  def forward(self, x, x_mask, g=None, reverse=False):
-    if not reverse:
-      for flow in self.flows:
-        x, _ = flow(x, x_mask, g=g, reverse=reverse)
-    else:
-      for flow in reversed(self.flows):
-        x = flow(x, x_mask, g=g, reverse=reverse)
-    return x
+    def __init__(self,
+                 channels,
+                 hidden_channels,
+                 kernel_size,
+                 dilation_rate,
+                 n_layers,
+                 n_flows=4,
+                 gin_channels=0):
+        super().__init__()
+        self.channels = channels
+        self.hidden_channels = hidden_channels
+        self.kernel_size = kernel_size
+        self.dilation_rate = dilation_rate
+        self.n_layers = n_layers
+        self.n_flows = n_flows
+        self.gin_channels = gin_channels
+
+        self.flows = nn.ModuleList()
+        for i in range(n_flows):
+            self.flows.append(
+                modules.ResidualCouplingLayer(channels, hidden_channels, kernel_size, dilation_rate, n_layers,
+                                              gin_channels=gin_channels, mean_only=True))
+            self.flows.append(modules.Flip())
+
+    def forward(self, x, x_mask, g=None, reverse=False):
+        if not reverse:
+            for flow in self.flows:
+                x, _ = flow(x, x_mask, g=g, reverse=reverse)
+        else:
+            for flow in reversed(self.flows):
+                x = flow(x, x_mask, g=g, reverse=reverse)
+        return x
 
 
 class Encoder(nn.Module):
-  def __init__(self,
-      in_channels,
-      out_channels,
-      hidden_channels,
-      kernel_size,
-      dilation_rate,
-      n_layers,
-      gin_channels=0):
-    super().__init__()
-    self.in_channels = in_channels
-    self.out_channels = out_channels
-    self.hidden_channels = hidden_channels
-    self.kernel_size = kernel_size
-    self.dilation_rate = dilation_rate
-    self.n_layers = n_layers
-    self.gin_channels = gin_channels
-
-    self.pre = nn.Conv1d(in_channels, hidden_channels, 1)
-    self.enc = modules.WN(hidden_channels, kernel_size, dilation_rate, n_layers, gin_channels=gin_channels)
-    self.proj = nn.Conv1d(hidden_channels, out_channels * 2, 1)
-
-  def forward(self, x, x_lengths, g=None):
-    # print(x.shape,x_lengths.shape)
-    x_mask = torch.unsqueeze(commons.sequence_mask(x_lengths, x.size(2)), 1).to(x.dtype)
-    x = self.pre(x) * x_mask
-    x = self.enc(x, x_mask, g=g)
-    stats = self.proj(x) * x_mask
-    m, logs = torch.split(stats, self.out_channels, dim=1)
-    z = (m + torch.randn_like(m) * torch.exp(logs)) * x_mask
-    return z, m, logs, x_mask
+    def __init__(self,
+                 in_channels,
+                 out_channels,
+                 hidden_channels,
+                 kernel_size,
+                 dilation_rate,
+                 n_layers,
+                 gin_channels=0):
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.hidden_channels = hidden_channels
+        self.kernel_size = kernel_size
+        self.dilation_rate = dilation_rate
+        self.n_layers = n_layers
+        self.gin_channels = gin_channels
+
+        self.pre = nn.Conv1d(in_channels, hidden_channels, 1)
+        self.enc = modules.WN(hidden_channels, kernel_size, dilation_rate, n_layers, gin_channels=gin_channels)
+        self.proj = nn.Conv1d(hidden_channels, out_channels * 2, 1)
+
+    def forward(self, x, x_lengths, g=None):
+        # print(x.shape,x_lengths.shape)
+        x_mask = torch.unsqueeze(commons.sequence_mask(x_lengths, x.size(2)), 1).to(x.dtype)
+        x = self.pre(x) * x_mask
+        x = self.enc(x, x_mask, g=g)
+        stats = self.proj(x) * x_mask
+        m, logs = torch.split(stats, self.out_channels, dim=1)
+        z = (m + torch.randn_like(m) * torch.exp(logs)) * x_mask
+        return z, m, logs, x_mask
 
 
 class TextEncoder(nn.Module):
-  def __init__(self,
-      out_channels,
-      hidden_channels,
-      kernel_size,
-      n_layers,
-      gin_channels=0,
-      filter_channels=None,
-      n_heads=None,
-      p_dropout=None):
-    super().__init__()
-    self.out_channels = out_channels
-    self.hidden_channels = hidden_channels
-    self.kernel_size = kernel_size
-    self.n_layers = n_layers
-    self.gin_channels = gin_channels
-    self.proj = nn.Conv1d(hidden_channels, out_channels * 2, 1)
-    self.f0_emb = nn.Embedding(256, hidden_channels)
-
-    self.enc_ =  attentions.Encoder(
-        hidden_channels,
-        filter_channels,
-        n_heads,
-        n_layers,
-        kernel_size,
-        p_dropout)
-
-  def forward(self, x, x_mask, f0=None, noice_scale=1):
-    x = x + self.f0_emb(f0).transpose(1,2)
-    x = self.enc_(x * x_mask, x_mask)
-    stats = self.proj(x) * x_mask
-    m, logs = torch.split(stats, self.out_channels, dim=1)
-    z = (m + torch.randn_like(m) * torch.exp(logs) * noice_scale) * x_mask
-
-    return z, m, logs, x_mask
+    def __init__(self,
+                 out_channels,
+                 hidden_channels,
+                 kernel_size,
+                 n_layers,
+                 gin_channels=0,
+                 filter_channels=None,
+                 n_heads=None,
+                 p_dropout=None):
+        super().__init__()
+        self.out_channels = out_channels
+        self.hidden_channels = hidden_channels
+        self.kernel_size = kernel_size
+        self.n_layers = n_layers
+        self.gin_channels = gin_channels
+        self.proj = nn.Conv1d(hidden_channels, out_channels * 2, 1)
+        self.f0_emb = nn.Embedding(256, hidden_channels)
 
+        self.enc_ = attentions.Encoder(
+            hidden_channels,
+            filter_channels,
+            n_heads,
+            n_layers,
+            kernel_size,
+            p_dropout)
+
+    def forward(self, x, x_mask, f0=None, noice_scale=1):
+        x = x + self.f0_emb(f0).transpose(1, 2)
+        x = self.enc_(x * x_mask, x_mask)
+        stats = self.proj(x) * x_mask
+        m, logs = torch.split(stats, self.out_channels, dim=1)
+        z = (m + torch.randn_like(m) * torch.exp(logs) * noice_scale) * x_mask
+
+        return z, m, logs, x_mask
 
 
 class DiscriminatorP(torch.nn.Module):
@@ -140,7 +140,7 @@ class DiscriminatorP(torch.nn.Module):
 
         # 1d to 2d
         b, c, t = x.shape
-        if t % self.period != 0: # pad first
+        if t % self.period != 0:  # pad first
             n_pad = self.period - (t % self.period)
             x = F.pad(x, (0, n_pad), "reflect")
             t = t + n_pad
@@ -188,7 +188,7 @@ class DiscriminatorS(torch.nn.Module):
 class MultiPeriodDiscriminator(torch.nn.Module):
     def __init__(self, use_spectral_norm=False):
         super(MultiPeriodDiscriminator, self).__init__()
-        periods = [2,3,5,7,11]
+        periods = [2, 3, 5, 7, 11]
 
         discs = [DiscriminatorS(use_spectral_norm=use_spectral_norm)]
         discs = discs + [DiscriminatorP(i, use_spectral_norm=use_spectral_norm) for i in periods]
@@ -225,26 +225,26 @@ class SpeakerEncoder(torch.nn.Module):
 
     def compute_partial_slices(self, total_frames, partial_frames, partial_hop):
         mel_slices = []
-        for i in range(0, total_frames-partial_frames, partial_hop):
-            mel_range = torch.arange(i, i+partial_frames)
+        for i in range(0, total_frames - partial_frames, partial_hop):
+            mel_range = torch.arange(i, i + partial_frames)
             mel_slices.append(mel_range)
 
         return mel_slices
 
     def embed_utterance(self, mel, partial_frames=128, partial_hop=64):
         mel_len = mel.size(1)
-        last_mel = mel[:,-partial_frames:]
+        last_mel = mel[:, -partial_frames:]
 
         if mel_len > partial_frames:
             mel_slices = self.compute_partial_slices(mel_len, partial_frames, partial_hop)
-            mels = list(mel[:,s] for s in mel_slices)
+            mels = list(mel[:, s] for s in mel_slices)
             mels.append(last_mel)
             mels = torch.stack(tuple(mels), 0).squeeze(1)
 
             with torch.no_grad():
                 partial_embeds = self(mels)
             embed = torch.mean(partial_embeds, axis=0).unsqueeze(0)
-            #embed = embed / torch.linalg.norm(embed, 2)
+            # embed = embed / torch.linalg.norm(embed, 2)
         else:
             with torch.no_grad():
                 embed = self(last_mel)
@@ -280,7 +280,7 @@ class F0Decoder(nn.Module):
             kernel_size,
             p_dropout)
         self.proj = nn.Conv1d(hidden_channels, out_channels, 1)
-        self.f0_prenet = nn.Conv1d(1, hidden_channels , 3, padding=1)
+        self.f0_prenet = nn.Conv1d(1, hidden_channels, 3, padding=1)
         self.cond = nn.Conv1d(spk_channels, hidden_channels, 1)
 
     def forward(self, x, norm_f0, x_mask, spk_emb=None):
@@ -295,126 +295,175 @@ class F0Decoder(nn.Module):
 
 
 class SynthesizerTrn(nn.Module):
-  """
-  Synthesizer for Training
-  """
-
-  def __init__(self,
-    spec_channels,
-    segment_size,
-    inter_channels,
-    hidden_channels,
-    filter_channels,
-    n_heads,
-    n_layers,
-    kernel_size,
-    p_dropout,
-    resblock,
-    resblock_kernel_sizes,
-    resblock_dilation_sizes,
-    upsample_rates,
-    upsample_initial_channel,
-    upsample_kernel_sizes,
-    gin_channels,
-    ssl_dim,
-    n_speakers,
-    sampling_rate=44100,
-    **kwargs):
-
-    super().__init__()
-    self.spec_channels = spec_channels
-    self.inter_channels = inter_channels
-    self.hidden_channels = hidden_channels
-    self.filter_channels = filter_channels
-    self.n_heads = n_heads
-    self.n_layers = n_layers
-    self.kernel_size = kernel_size
-    self.p_dropout = p_dropout
-    self.resblock = resblock
-    self.resblock_kernel_sizes = resblock_kernel_sizes
-    self.resblock_dilation_sizes = resblock_dilation_sizes
-    self.upsample_rates = upsample_rates
-    self.upsample_initial_channel = upsample_initial_channel
-    self.upsample_kernel_sizes = upsample_kernel_sizes
-    self.segment_size = segment_size
-    self.gin_channels = gin_channels
-    self.ssl_dim = ssl_dim
-    self.emb_g = nn.Embedding(n_speakers, gin_channels)
-
-    self.pre = nn.Conv1d(ssl_dim, hidden_channels, kernel_size=5, padding=2)
-
-    self.enc_p = TextEncoder(
-        inter_channels,
-        hidden_channels,
-        filter_channels=filter_channels,
-        n_heads=n_heads,
-        n_layers=n_layers,
-        kernel_size=kernel_size,
-        p_dropout=p_dropout
-    )
-    hps = {
-        "sampling_rate": sampling_rate,
-        "inter_channels": inter_channels,
-        "resblock": resblock,
-        "resblock_kernel_sizes": resblock_kernel_sizes,
-        "resblock_dilation_sizes": resblock_dilation_sizes,
-        "upsample_rates": upsample_rates,
-        "upsample_initial_channel": upsample_initial_channel,
-        "upsample_kernel_sizes": upsample_kernel_sizes,
-        "gin_channels": gin_channels,
-    }
-    self.dec = Generator(h=hps)
-    self.enc_q = Encoder(spec_channels, inter_channels, hidden_channels, 5, 1, 16, gin_channels=gin_channels)
-    self.flow = ResidualCouplingBlock(inter_channels, hidden_channels, 5, 1, 4, gin_channels=gin_channels)
-    self.f0_decoder = F0Decoder(
-        1,
-        hidden_channels,
-        filter_channels,
-        n_heads,
-        n_layers,
-        kernel_size,
-        p_dropout,
-        spk_channels=gin_channels
-    )
-    self.emb_uv = nn.Embedding(2, hidden_channels)
-
-  def forward(self, c, f0, uv, spec, g=None, c_lengths=None, spec_lengths=None):
-    g = self.emb_g(g).transpose(1,2)
-    # ssl prenet
-    x_mask = torch.unsqueeze(commons.sequence_mask(c_lengths, c.size(2)), 1).to(c.dtype)
-    x = self.pre(c) * x_mask + self.emb_uv(uv.long()).transpose(1,2)
-
-    # f0 predict
-    lf0 = 2595. * torch.log10(1. + f0.unsqueeze(1) / 700.) / 500
-    norm_lf0 = utils.normalize_f0(lf0, x_mask, uv)
-    pred_lf0 = self.f0_decoder(x, norm_lf0, x_mask, spk_emb=g)
-
-    # encoder
-    z_ptemp, m_p, logs_p, _ = self.enc_p(x, x_mask, f0=f0_to_coarse(f0))
-    z, m_q, logs_q, spec_mask = self.enc_q(spec, spec_lengths, g=g) 
-
-    # flow
-    z_p = self.flow(z, spec_mask, g=g)
-    z_slice, pitch_slice, ids_slice = commons.rand_slice_segments_with_pitch(z, f0, spec_lengths, self.segment_size)
-
-    # nsf decoder
-    o = self.dec(z_slice, g=g, f0=pitch_slice)
-
-    return o, ids_slice, spec_mask, (z, z_p, m_p, logs_p, m_q, logs_q), pred_lf0, norm_lf0, lf0
-
-  def infer(self, c, f0, uv, g=None, noice_scale=0.35, predict_f0=False):
-    c_lengths = (torch.ones(c.size(0)) * c.size(-1)).to(c.device)
-    g = self.emb_g(g).transpose(1,2)
-    x_mask = torch.unsqueeze(commons.sequence_mask(c_lengths, c.size(2)), 1).to(c.dtype)
-    x = self.pre(c) * x_mask + self.emb_uv(uv.long()).transpose(1,2)
-
-    if predict_f0:
+    """
+    Synthesizer for Training
+    """
+
+    def __init__(self,
+                 spec_channels,
+                 segment_size,
+                 inter_channels,
+                 hidden_channels,
+                 filter_channels,
+                 n_heads,
+                 n_layers,
+                 kernel_size,
+                 p_dropout,
+                 resblock,
+                 resblock_kernel_sizes,
+                 resblock_dilation_sizes,
+                 upsample_rates,
+                 upsample_initial_channel,
+                 upsample_kernel_sizes,
+                 gin_channels,
+                 ssl_dim,
+                 n_speakers,
+                 sampling_rate=44100,
+                 vol_embedding=False,
+                 vocoder_name = "nsf-hifigan",
+                 **kwargs):
+
+        super().__init__()
+        self.spec_channels = spec_channels
+        self.inter_channels = inter_channels
+        self.hidden_channels = hidden_channels
+        self.filter_channels = filter_channels
+        self.n_heads = n_heads
+        self.n_layers = n_layers
+        self.kernel_size = kernel_size
+        self.p_dropout = p_dropout
+        self.resblock = resblock
+        self.resblock_kernel_sizes = resblock_kernel_sizes
+        self.resblock_dilation_sizes = resblock_dilation_sizes
+        self.upsample_rates = upsample_rates
+        self.upsample_initial_channel = upsample_initial_channel
+        self.upsample_kernel_sizes = upsample_kernel_sizes
+        self.segment_size = segment_size
+        self.gin_channels = gin_channels
+        self.ssl_dim = ssl_dim
+        self.vol_embedding = vol_embedding
+        self.emb_g = nn.Embedding(n_speakers, gin_channels)
+        if vol_embedding:
+           self.emb_vol = nn.Linear(1, hidden_channels)
+
+        self.pre = nn.Conv1d(ssl_dim, hidden_channels, kernel_size=5, padding=2)
+
+        self.enc_p = TextEncoder(
+            inter_channels,
+            hidden_channels,
+            filter_channels=filter_channels,
+            n_heads=n_heads,
+            n_layers=n_layers,
+            kernel_size=kernel_size,
+            p_dropout=p_dropout
+        )
+        hps = {
+            "sampling_rate": sampling_rate,
+            "inter_channels": inter_channels,
+            "resblock": resblock,
+            "resblock_kernel_sizes": resblock_kernel_sizes,
+            "resblock_dilation_sizes": resblock_dilation_sizes,
+            "upsample_rates": upsample_rates,
+            "upsample_initial_channel": upsample_initial_channel,
+            "upsample_kernel_sizes": upsample_kernel_sizes,
+            "gin_channels": gin_channels,
+        }
+        
+        
+        if vocoder_name == "nsf-hifigan":
+            from vdecoder.hifigan.models import Generator
+            self.dec = Generator(h=hps)
+        elif vocoder_name == "nsf-snake-hifigan":
+            from vdecoder.hifiganwithsnake.models import Generator
+            self.dec = Generator(h=hps)
+        else:
+            print("[?] Unkown vocoder: use default(nsf-hifigan)")
+            from vdecoder.hifigan.models import Generator
+            self.dec = Generator(h=hps)
+
+        self.enc_q = Encoder(spec_channels, inter_channels, hidden_channels, 5, 1, 16, gin_channels=gin_channels)
+        self.flow = ResidualCouplingBlock(inter_channels, hidden_channels, 5, 1, 4, gin_channels=gin_channels)
+        self.f0_decoder = F0Decoder(
+            1,
+            hidden_channels,
+            filter_channels,
+            n_heads,
+            n_layers,
+            kernel_size,
+            p_dropout,
+            spk_channels=gin_channels
+        )
+        self.emb_uv = nn.Embedding(2, hidden_channels)
+        self.character_mix = False
+
+    def EnableCharacterMix(self, n_speakers_map, device):
+        self.speaker_map = torch.zeros((n_speakers_map, 1, 1, self.gin_channels)).to(device)
+        for i in range(n_speakers_map):
+            self.speaker_map[i] = self.emb_g(torch.LongTensor([[i]]).to(device))
+        self.speaker_map = self.speaker_map.unsqueeze(0).to(device)
+        self.character_mix = True
+
+    def forward(self, c, f0, uv, spec, g=None, c_lengths=None, spec_lengths=None, vol = None):
+        g = self.emb_g(g).transpose(1,2)
+
+        # vol proj
+        vol = self.emb_vol(vol[:,:,None]).transpose(1,2) if vol!=None and self.vol_embedding else 0
+
+        # ssl prenet
+        x_mask = torch.unsqueeze(commons.sequence_mask(c_lengths, c.size(2)), 1).to(c.dtype)
+        x = self.pre(c) * x_mask + self.emb_uv(uv.long()).transpose(1,2) + vol
+
+        # f0 predict
         lf0 = 2595. * torch.log10(1. + f0.unsqueeze(1) / 700.) / 500
-        norm_lf0 = utils.normalize_f0(lf0, x_mask, uv, random_scale=False)
+        norm_lf0 = utils.normalize_f0(lf0, x_mask, uv)
         pred_lf0 = self.f0_decoder(x, norm_lf0, x_mask, spk_emb=g)
-        f0 = (700 * (torch.pow(10, pred_lf0 * 500 / 2595) - 1)).squeeze(1)
 
-    z_p, m_p, logs_p, c_mask = self.enc_p(x, x_mask, f0=f0_to_coarse(f0), noice_scale=noice_scale)
-    z = self.flow(z_p, c_mask, g=g, reverse=True)
-    o = self.dec(z * c_mask, g=g, f0=f0)
-    return o
+        # encoder
+        z_ptemp, m_p, logs_p, _ = self.enc_p(x, x_mask, f0=f0_to_coarse(f0))
+        z, m_q, logs_q, spec_mask = self.enc_q(spec, spec_lengths, g=g)
+
+        # flow
+        z_p = self.flow(z, spec_mask, g=g)
+        z_slice, pitch_slice, ids_slice = commons.rand_slice_segments_with_pitch(z, f0, spec_lengths, self.segment_size)
+
+        # nsf decoder
+        o = self.dec(z_slice, g=g, f0=pitch_slice)
+
+        return o, ids_slice, spec_mask, (z, z_p, m_p, logs_p, m_q, logs_q), pred_lf0, norm_lf0, lf0
+
+    def infer(self, c, f0, uv, g=None, noice_scale=0.35, seed=52468, predict_f0=False, vol = None):
+
+        if c.device == torch.device("cuda"):
+            torch.cuda.manual_seed_all(seed)
+        else:
+            torch.manual_seed(seed)
+
+        c_lengths = (torch.ones(c.size(0)) * c.size(-1)).to(c.device)
+
+        if self.character_mix and len(g) > 1:   # [N, S]  *  [S, B, 1, H]
+            g = g.reshape((g.shape[0], g.shape[1], 1, 1, 1))  # [N, S, B, 1, 1]
+            g = g * self.speaker_map  # [N, S, B, 1, H]
+            g = torch.sum(g, dim=1) # [N, 1, B, 1, H]
+            g = g.transpose(0, -1).transpose(0, -2).squeeze(0) # [B, H, N]
+        else:
+            if g.dim() == 1:
+                g = g.unsqueeze(0)
+            g = self.emb_g(g).transpose(1, 2)
+        
+        x_mask = torch.unsqueeze(commons.sequence_mask(c_lengths, c.size(2)), 1).to(c.dtype)
+        # vol proj
+        vol = self.emb_vol(vol[:,:,None]).transpose(1,2) if vol!=None and self.vol_embedding else 0
+           
+        x = self.pre(c) * x_mask + self.emb_uv(uv.long()).transpose(1,2) + vol
+        
+        if predict_f0:
+            lf0 = 2595. * torch.log10(1. + f0.unsqueeze(1) / 700.) / 500
+            norm_lf0 = utils.normalize_f0(lf0, x_mask, uv, random_scale=False)
+            pred_lf0 = self.f0_decoder(x, norm_lf0, x_mask, spk_emb=g)
+            f0 = (700 * (torch.pow(10, pred_lf0 * 500 / 2595) - 1)).squeeze(1)
+        
+        z_p, m_p, logs_p, c_mask = self.enc_p(x, x_mask, f0=f0_to_coarse(f0), noice_scale=noice_scale)
+        z = self.flow(z_p, c_mask, g=g, reverse=True)
+        o = self.dec(z * c_mask, g=g, f0=f0)
+        return o,f0
+

utils.py

@@ -6,18 +6,21 @@ import argparse
 import logging
 import json
 import subprocess
+import warnings
 import random
-
+import functools
 import librosa
 import numpy as np
 from scipy.io.wavfile import read
 import torch
 from torch.nn import functional as F
 from modules.commons import sequence_mask
-from hubert import hubert_model
+import faiss
+import tqdm
+
 MATPLOTLIB_FLAG = False
 
-logging.basicConfig(stream=sys.stdout, level=logging.DEBUG)
+logging.basicConfig(stream=sys.stdout, level=logging.WARN)
 logger = logging
 
 f0_bin = 256
@@ -26,26 +29,6 @@ f0_min = 50.0
 f0_mel_min = 1127 * np.log(1 + f0_min / 700)
 f0_mel_max = 1127 * np.log(1 + f0_max / 700)
 
-
-# def normalize_f0(f0, random_scale=True):
-#     f0_norm = f0.clone()  # create a copy of the input Tensor
-#     batch_size, _, frame_length = f0_norm.shape
-#     for i in range(batch_size):
-#         means = torch.mean(f0_norm[i, 0, :])
-#         if random_scale:
-#             factor = random.uniform(0.8, 1.2)
-#         else:
-#             factor = 1
-#         f0_norm[i, 0, :] = (f0_norm[i, 0, :] - means) * factor
-#     return f0_norm
-# def normalize_f0(f0, random_scale=True):
-#     means = torch.mean(f0[:, 0, :], dim=1, keepdim=True)
-#     if random_scale:
-#         factor = torch.Tensor(f0.shape[0],1).uniform_(0.8, 1.2).to(f0.device)
-#     else:
-#         factor = torch.ones(f0.shape[0], 1, 1).to(f0.device)
-#     f0_norm = (f0 - means.unsqueeze(-1)) * factor.unsqueeze(-1)
-#     return f0_norm
 def normalize_f0(f0, x_mask, uv, random_scale=True):
     # calculate means based on x_mask
     uv_sum = torch.sum(uv, dim=1, keepdim=True)
@@ -62,7 +45,6 @@ def normalize_f0(f0, x_mask, uv, random_scale=True):
         exit(0)
     return f0_norm * x_mask
 
-
 def plot_data_to_numpy(x, y):
     global MATPLOTLIB_FLAG
     if not MATPLOTLIB_FLAG:
@@ -86,87 +68,6 @@ def plot_data_to_numpy(x, y):
     return data
 
 
-
-def interpolate_f0(f0):
-    '''
-    对F0进行插值处理
-    '''
-
-    data = np.reshape(f0, (f0.size, 1))
-
-    vuv_vector = np.zeros((data.size, 1), dtype=np.float32)
-    vuv_vector[data > 0.0] = 1.0
-    vuv_vector[data <= 0.0] = 0.0
-
-    ip_data = data
-
-    frame_number = data.size
-    last_value = 0.0
-    for i in range(frame_number):
-        if data[i] <= 0.0:
-            j = i + 1
-            for j in range(i + 1, frame_number):
-                if data[j] > 0.0:
-                    break
-            if j < frame_number - 1:
-                if last_value > 0.0:
-                    step = (data[j] - data[i - 1]) / float(j - i)
-                    for k in range(i, j):
-                        ip_data[k] = data[i - 1] + step * (k - i + 1)
-                else:
-                    for k in range(i, j):
-                        ip_data[k] = data[j]
-            else:
-                for k in range(i, frame_number):
-                    ip_data[k] = last_value
-        else:
-            ip_data[i] = data[i]
-            last_value = data[i]
-
-    return ip_data[:,0], vuv_vector[:,0]
-
-
-def compute_f0_parselmouth(wav_numpy, p_len=None, sampling_rate=44100, hop_length=512):
-    import parselmouth
-    x = wav_numpy
-    if p_len is None:
-        p_len = x.shape[0]//hop_length
-    else:
-        assert abs(p_len-x.shape[0]//hop_length) < 4, "pad length error"
-    time_step = hop_length / sampling_rate * 1000
-    f0_min = 50
-    f0_max = 1100
-    f0 = parselmouth.Sound(x, sampling_rate).to_pitch_ac(
-        time_step=time_step / 1000, voicing_threshold=0.6,
-        pitch_floor=f0_min, pitch_ceiling=f0_max).selected_array['frequency']
-
-    pad_size=(p_len - len(f0) + 1) // 2
-    if(pad_size>0 or p_len - len(f0) - pad_size>0):
-        f0 = np.pad(f0,[[pad_size,p_len - len(f0) - pad_size]], mode='constant')
-    return f0
-
-def resize_f0(x, target_len):
-    source = np.array(x)
-    source[source<0.001] = np.nan
-    target = np.interp(np.arange(0, len(source)*target_len, len(source))/ target_len, np.arange(0, len(source)), source)
-    res = np.nan_to_num(target)
-    return res
-
-def compute_f0_dio(wav_numpy, p_len=None, sampling_rate=44100, hop_length=512):
-    import pyworld
-    if p_len is None:
-        p_len = wav_numpy.shape[0]//hop_length
-    f0, t = pyworld.dio(
-        wav_numpy.astype(np.double),
-        fs=sampling_rate,
-        f0_ceil=800,
-        frame_period=1000 * hop_length / sampling_rate,
-    )
-    f0 = pyworld.stonemask(wav_numpy.astype(np.double), f0, t, sampling_rate)
-    for index, pitch in enumerate(f0):
-        f0[index] = round(pitch, 1)
-    return resize_f0(f0, p_len)
-
 def f0_to_coarse(f0):
   is_torch = isinstance(f0, torch.Tensor)
   f0_mel = 1127 * (1 + f0 / 700).log() if is_torch else 1127 * np.log(1 + f0 / 700)
@@ -174,48 +75,73 @@ def f0_to_coarse(f0):
 
   f0_mel[f0_mel <= 1] = 1
   f0_mel[f0_mel > f0_bin - 1] = f0_bin - 1
-  f0_coarse = (f0_mel + 0.5).long() if is_torch else np.rint(f0_mel).astype(np.int)
+  f0_coarse = (f0_mel + 0.5).int() if is_torch else np.rint(f0_mel).astype(np.int)
   assert f0_coarse.max() <= 255 and f0_coarse.min() >= 1, (f0_coarse.max(), f0_coarse.min())
   return f0_coarse
 
-
-def get_hubert_model():
-  vec_path = "hubert/checkpoint_best_legacy_500.pt"
-  print("load model(s) from {}".format(vec_path))
-  from fairseq import checkpoint_utils
-  models, saved_cfg, task = checkpoint_utils.load_model_ensemble_and_task(
-    [vec_path],
-    suffix="",
-  )
-  model = models[0]
-  model.eval()
-  return model
-
-def get_hubert_content(hmodel, wav_16k_tensor):
-  feats = wav_16k_tensor
-  if feats.dim() == 2:  # double channels
-    feats = feats.mean(-1)
-  assert feats.dim() == 1, feats.dim()
-  feats = feats.view(1, -1)
-  padding_mask = torch.BoolTensor(feats.shape).fill_(False)
-  inputs = {
-    "source": feats.to(wav_16k_tensor.device),
-    "padding_mask": padding_mask.to(wav_16k_tensor.device),
-    "output_layer": 9,  # layer 9
-  }
-  with torch.no_grad():
-    logits = hmodel.extract_features(**inputs)
-    feats = hmodel.final_proj(logits[0])
-  return feats.transpose(1, 2)
-
-
 def get_content(cmodel, y):
     with torch.no_grad():
         c = cmodel.extract_features(y.squeeze(1))[0]
     c = c.transpose(1, 2)
     return c
 
-
+def get_f0_predictor(f0_predictor,hop_length,sampling_rate,**kargs):
+    if f0_predictor == "pm":
+        from modules.F0Predictor.PMF0Predictor import PMF0Predictor
+        f0_predictor_object = PMF0Predictor(hop_length=hop_length,sampling_rate=sampling_rate)
+    elif f0_predictor == "crepe":
+        from modules.F0Predictor.CrepeF0Predictor import CrepeF0Predictor
+        f0_predictor_object = CrepeF0Predictor(hop_length=hop_length,sampling_rate=sampling_rate,device=kargs["device"],threshold=kargs["threshold"])
+    elif f0_predictor == "harvest":
+        from modules.F0Predictor.HarvestF0Predictor import HarvestF0Predictor
+        f0_predictor_object = HarvestF0Predictor(hop_length=hop_length,sampling_rate=sampling_rate)
+    elif f0_predictor == "dio":
+        from modules.F0Predictor.DioF0Predictor import DioF0Predictor
+        f0_predictor_object = DioF0Predictor(hop_length=hop_length,sampling_rate=sampling_rate)
+    else:
+        raise Exception("Unknown f0 predictor")
+    return f0_predictor_object
+
+def get_speech_encoder(speech_encoder,device=None,**kargs):
+    if speech_encoder == "vec768l12":
+        from vencoder.ContentVec768L12 import ContentVec768L12
+        speech_encoder_object = ContentVec768L12(device = device)
+    elif speech_encoder == "vec256l9":
+        from vencoder.ContentVec256L9 import ContentVec256L9
+        speech_encoder_object = ContentVec256L9(device = device)
+    elif speech_encoder == "vec256l9-onnx":
+        from vencoder.ContentVec256L9_Onnx import ContentVec256L9_Onnx
+        speech_encoder_object = ContentVec256L9_Onnx(device = device)
+    elif speech_encoder == "vec256l12-onnx":
+        from vencoder.ContentVec256L12_Onnx import ContentVec256L12_Onnx
+        speech_encoder_object = ContentVec256L12_Onnx(device = device)
+    elif speech_encoder == "vec768l9-onnx":
+        from vencoder.ContentVec768L9_Onnx import ContentVec768L9_Onnx
+        speech_encoder_object = ContentVec768L9_Onnx(device = device)
+    elif speech_encoder == "vec768l12-onnx":
+        from vencoder.ContentVec768L12_Onnx import ContentVec768L12_Onnx
+        speech_encoder_object = ContentVec768L12_Onnx(device = device)
+    elif speech_encoder == "hubertsoft-onnx":
+        from vencoder.HubertSoft_Onnx import HubertSoft_Onnx
+        speech_encoder_object = HubertSoft_Onnx(device = device)
+    elif speech_encoder == "hubertsoft":
+        from vencoder.HubertSoft import HubertSoft
+        speech_encoder_object = HubertSoft(device = device)
+    elif speech_encoder == "whisper-ppg":
+        from vencoder.WhisperPPG import WhisperPPG
+        speech_encoder_object = WhisperPPG(device = device)
+    elif speech_encoder == "cnhubertlarge":
+        from vencoder.CNHubertLarge import CNHubertLarge
+        speech_encoder_object = CNHubertLarge(device = device)
+    elif speech_encoder == "dphubert":
+        from vencoder.DPHubert import DPHubert
+        speech_encoder_object = DPHubert(device = device)
+    elif speech_encoder == "whisper-ppg-large":
+        from vencoder.WhisperPPGLarge import WhisperPPGLarge
+        speech_encoder_object = WhisperPPGLarge(device = device)
+    else:
+        raise Exception("Unknown speech encoder")
+    return speech_encoder_object 
 
 def load_checkpoint(checkpoint_path, model, optimizer=None, skip_optimizer=False):
     assert os.path.isfile(checkpoint_path)
@@ -244,6 +170,7 @@ def load_checkpoint(checkpoint_path, model, optimizer=None, skip_optimizer=False
         model.module.load_state_dict(new_state_dict)
     else:
         model.load_state_dict(new_state_dict)
+    print("load ")
     logger.info("Loaded checkpoint '{}' (iteration {})".format(
         checkpoint_path, iteration))
     return model, optimizer, learning_rate, iteration
@@ -368,7 +295,7 @@ def load_filepaths_and_text(filename, split="|"):
 
 def get_hparams(init=True):
   parser = argparse.ArgumentParser()
-  parser.add_argument('-c', '--config', type=str, default="./configs/base.json",
+  parser.add_argument('-c', '--config', type=str, default="./configs/config.json",
                       help='JSON file for configuration')
   parser.add_argument('-m', '--model', type=str, required=True,
                       help='Model name')
@@ -411,7 +338,6 @@ def get_hparams_from_file(config_path):
   with open(config_path, "r") as f:
     data = f.read()
   config = json.loads(data)
-
   hparams =HParams(**config)
   return hparams
 
@@ -468,6 +394,73 @@ def repeat_expand_2d(content, target_len):
     return target
 
 
+def mix_model(model_paths,mix_rate,mode):
+  mix_rate = torch.FloatTensor(mix_rate)/100
+  model_tem = torch.load(model_paths[0])
+  models = [torch.load(path)["model"] for path in model_paths]
+  if mode == 0:
+     mix_rate = F.softmax(mix_rate,dim=0)
+  for k in model_tem["model"].keys():
+     model_tem["model"][k] = torch.zeros_like(model_tem["model"][k])
+     for i,model in enumerate(models):
+        model_tem["model"][k] += model[k]*mix_rate[i]
+  torch.save(model_tem,os.path.join(os.path.curdir,"output.pth"))
+  return os.path.join(os.path.curdir,"output.pth")
+  
+def change_rms(data1, sr1, data2, sr2, rate):  # 1是输入音频，2是输出音频,rate是2的占比 from RVC
+    # print(data1.max(),data2.max())
+    rms1 = librosa.feature.rms(
+        y=data1, frame_length=sr1 // 2 * 2, hop_length=sr1 // 2
+    )  # 每半秒一个点
+    rms2 = librosa.feature.rms(y=data2.detach().cpu().numpy(), frame_length=sr2 // 2 * 2, hop_length=sr2 // 2)
+    rms1 = torch.from_numpy(rms1).to(data2.device)
+    rms1 = F.interpolate(
+        rms1.unsqueeze(0), size=data2.shape[0], mode="linear"
+    ).squeeze()
+    rms2 = torch.from_numpy(rms2).to(data2.device)
+    rms2 = F.interpolate(
+        rms2.unsqueeze(0), size=data2.shape[0], mode="linear"
+    ).squeeze()
+    rms2 = torch.max(rms2, torch.zeros_like(rms2) + 1e-6)
+    data2 *= (
+        torch.pow(rms1, torch.tensor(1 - rate))
+        * torch.pow(rms2, torch.tensor(rate - 1))
+    )
+    return data2
+
+def train_index(spk_name,root_dir = "dataset/44k/"):  #from: RVC https://github.com/RVC-Project/Retrieval-based-Voice-Conversion-WebUI
+    print("The feature index is constructing.")
+    exp_dir = os.path.join(root_dir,spk_name)
+    listdir_res = []
+    for file in os.listdir(exp_dir):
+       if ".wav.soft.pt" in file:
+          listdir_res.append(os.path.join(exp_dir,file))
+    if len(listdir_res) == 0:
+        raise Exception("You need to run preprocess_hubert_f0.py!")
+    npys = []
+    for name in sorted(listdir_res):
+        phone = torch.load(name)[0].transpose(-1,-2).numpy()
+        npys.append(phone)
+    big_npy = np.concatenate(npys, 0)
+    big_npy_idx = np.arange(big_npy.shape[0])
+    np.random.shuffle(big_npy_idx)
+    big_npy = big_npy[big_npy_idx]
+    n_ivf = min(int(16 * np.sqrt(big_npy.shape[0])), big_npy.shape[0] // 39)
+    index = faiss.index_factory(big_npy.shape[1] , "IVF%s,Flat" % n_ivf)
+    index_ivf = faiss.extract_index_ivf(index)  #
+    index_ivf.nprobe = 1
+    index.train(big_npy)
+    batch_size_add = 8192
+    for i in range(0, big_npy.shape[0], batch_size_add):
+        index.add(big_npy[i : i + batch_size_add])
+    # faiss.write_index(
+    #     index,
+    #     f"added_{spk_name}.index"
+    # )
+    print("Successfully build index")
+    return index
+
+
 class HParams():
   def __init__(self, **kwargs):
     for k, v in kwargs.items():
@@ -499,3 +492,19 @@ class HParams():
   def __repr__(self):
     return self.__dict__.__repr__()
 
+  def get(self,index):
+    return self.__dict__.get(index)
+
+class Volume_Extractor:
+    def __init__(self, hop_size = 512):
+        self.hop_size = hop_size
+        
+    def extract(self, audio): # audio: 2d tensor array
+        if not isinstance(audio,torch.Tensor):
+           audio = torch.Tensor(audio)
+        n_frames = int(audio.size(-1) // self.hop_size)
+        audio2 = audio ** 2
+        audio2 = torch.nn.functional.pad(audio2, (int(self.hop_size // 2), int((self.hop_size + 1) // 2)), mode = 'reflect')
+        volume = torch.FloatTensor([torch.mean(audio2[:,int(n * self.hop_size) : int((n + 1) * self.hop_size)]) for n in range(n_frames)])
+        volume = torch.sqrt(volume)
+        return volume