kwau/sovits-isla

1import argparse2import glob3import json4import logging5import os6import re7import subprocess8import sys9import traceback10from multiprocessing import cpu_count11 12import faiss13import librosa14import numpy as np15import torch16from scipy.io.wavfile import read17from sklearn.cluster import MiniBatchKMeans18from torch.nn import functional as F19 20MATPLOTLIB_FLAG = False21 22logging.basicConfig(stream=sys.stdout, level=logging.WARN)23logger = logging24 25f0_bin = 25626f0_max = 1100.027f0_min = 50.028f0_mel_min = 1127 * np.log(1 + f0_min / 700)29f0_mel_max = 1127 * np.log(1 + f0_max / 700)30 31def normalize_f0(f0, x_mask, uv, random_scale=True):32    # calculate means based on x_mask33    uv_sum = torch.sum(uv, dim=1, keepdim=True)34    uv_sum[uv_sum == 0] = 999935    means = torch.sum(f0[:, 0, :] * uv, dim=1, keepdim=True) / uv_sum36 37    if random_scale:38        factor = torch.Tensor(f0.shape[0], 1).uniform_(0.8, 1.2).to(f0.device)39    else:40        factor = torch.ones(f0.shape[0], 1).to(f0.device)41    # normalize f0 based on means and factor42    f0_norm = (f0 - means.unsqueeze(-1)) * factor.unsqueeze(-1)43    if torch.isnan(f0_norm).any():44        exit(0)45    return f0_norm * x_mask46def plot_data_to_numpy(x, y):47    global MATPLOTLIB_FLAG48    if not MATPLOTLIB_FLAG:49        import matplotlib50        matplotlib.use("Agg")51        MATPLOTLIB_FLAG = True52        mpl_logger = logging.getLogger('matplotlib')53        mpl_logger.setLevel(logging.WARNING)54    import matplotlib.pylab as plt55    import numpy as np56 57    fig, ax = plt.subplots(figsize=(10, 2))58    plt.plot(x)59    plt.plot(y)60    plt.tight_layout()61 62    fig.canvas.draw()63    data = np.fromstring(fig.canvas.tostring_rgb(), dtype=np.uint8, sep='')64    data = data.reshape(fig.canvas.get_width_height()[::-1] + (3,))65    plt.close()66    return data67 68 69def f0_to_coarse(f0):70  f0_mel = 1127 * (1 + f0 / 700).log()71  a = (f0_bin - 2) / (f0_mel_max - f0_mel_min)72  b = f0_mel_min * a - 1.73  f0_mel = torch.where(f0_mel > 0, f0_mel * a - b, f0_mel)74  # torch.clip_(f0_mel, min=1., max=float(f0_bin - 1))75  f0_coarse = torch.round(f0_mel).long()76  f0_coarse = f0_coarse * (f0_coarse > 0)77  f0_coarse = f0_coarse + ((f0_coarse < 1) * 1)78  f0_coarse = f0_coarse * (f0_coarse < f0_bin)79  f0_coarse = f0_coarse + ((f0_coarse >= f0_bin) * (f0_bin - 1))80  return f0_coarse81 82def get_content(cmodel, y):83    with torch.no_grad():84        c = cmodel.extract_features(y.squeeze(1))[0]85    c = c.transpose(1, 2)86    return c87 88def get_f0_predictor(f0_predictor,hop_length,sampling_rate,**kargs):89    if f0_predictor == "pm":90        from modules.F0Predictor.PMF0Predictor import PMF0Predictor91        f0_predictor_object = PMF0Predictor(hop_length=hop_length,sampling_rate=sampling_rate)92    elif f0_predictor == "crepe":93        from modules.F0Predictor.CrepeF0Predictor import CrepeF0Predictor94        f0_predictor_object = CrepeF0Predictor(hop_length=hop_length,sampling_rate=sampling_rate,device=kargs["device"],threshold=kargs["threshold"])95    elif f0_predictor == "harvest":96        from modules.F0Predictor.HarvestF0Predictor import HarvestF0Predictor97        f0_predictor_object = HarvestF0Predictor(hop_length=hop_length,sampling_rate=sampling_rate)98    elif f0_predictor == "dio":99        from modules.F0Predictor.DioF0Predictor import DioF0Predictor100        f0_predictor_object = DioF0Predictor(hop_length=hop_length,sampling_rate=sampling_rate) 101    elif f0_predictor == "rmvpe":102        from modules.F0Predictor.RMVPEF0Predictor import RMVPEF0Predictor103        f0_predictor_object = RMVPEF0Predictor(hop_length=hop_length,sampling_rate=sampling_rate,dtype=torch.float32 ,device=kargs["device"],threshold=kargs["threshold"])104    elif f0_predictor == "fcpe":105        from modules.F0Predictor.FCPEF0Predictor import FCPEF0Predictor106        f0_predictor_object = FCPEF0Predictor(hop_length=hop_length,sampling_rate=sampling_rate,dtype=torch.float32 ,device=kargs["device"],threshold=kargs["threshold"])107    else:108        raise Exception("Unknown f0 predictor")109    return f0_predictor_object110 111def get_speech_encoder(speech_encoder,device=None,**kargs):112    if speech_encoder == "vec768l12":113        from vencoder.ContentVec768L12 import ContentVec768L12114        speech_encoder_object = ContentVec768L12(device = device)115    elif speech_encoder == "vec256l9":116        from vencoder.ContentVec256L9 import ContentVec256L9117        speech_encoder_object = ContentVec256L9(device = device)118    elif speech_encoder == "vec256l9-onnx":119        from vencoder.ContentVec256L9_Onnx import ContentVec256L9_Onnx120        speech_encoder_object = ContentVec256L9_Onnx(device = device)121    elif speech_encoder == "vec256l12-onnx":122        from vencoder.ContentVec256L12_Onnx import ContentVec256L12_Onnx123        speech_encoder_object = ContentVec256L12_Onnx(device = device)124    elif speech_encoder == "vec768l9-onnx":125        from vencoder.ContentVec768L9_Onnx import ContentVec768L9_Onnx126        speech_encoder_object = ContentVec768L9_Onnx(device = device)127    elif speech_encoder == "vec768l12-onnx":128        from vencoder.ContentVec768L12_Onnx import ContentVec768L12_Onnx129        speech_encoder_object = ContentVec768L12_Onnx(device = device)130    elif speech_encoder == "hubertsoft-onnx":131        from vencoder.HubertSoft_Onnx import HubertSoft_Onnx132        speech_encoder_object = HubertSoft_Onnx(device = device)133    elif speech_encoder == "hubertsoft":134        from vencoder.HubertSoft import HubertSoft135        speech_encoder_object = HubertSoft(device = device)136    elif speech_encoder == "whisper-ppg":137        from vencoder.WhisperPPG import WhisperPPG138        speech_encoder_object = WhisperPPG(device = device)139    elif speech_encoder == "cnhubertlarge":140        from vencoder.CNHubertLarge import CNHubertLarge141        speech_encoder_object = CNHubertLarge(device = device)142    elif speech_encoder == "dphubert":143        from vencoder.DPHubert import DPHubert144        speech_encoder_object = DPHubert(device = device)145    elif speech_encoder == "whisper-ppg-large":146        from vencoder.WhisperPPGLarge import WhisperPPGLarge147        speech_encoder_object = WhisperPPGLarge(device = device)148    elif speech_encoder == "wavlmbase+":149        from vencoder.WavLMBasePlus import WavLMBasePlus150        speech_encoder_object = WavLMBasePlus(device = device)151    else:152        raise Exception("Unknown speech encoder")153    return speech_encoder_object 154 155def load_checkpoint(checkpoint_path, model, optimizer=None, skip_optimizer=False):156    assert os.path.isfile(checkpoint_path)157    checkpoint_dict = torch.load(checkpoint_path, map_location='cpu')158    iteration = checkpoint_dict['iteration']159    learning_rate = checkpoint_dict['learning_rate']160    if optimizer is not None and not skip_optimizer and checkpoint_dict['optimizer'] is not None:161        optimizer.load_state_dict(checkpoint_dict['optimizer'])162    saved_state_dict = checkpoint_dict['model']163    model = model.to(list(saved_state_dict.values())[0].dtype)164    if hasattr(model, 'module'):165        state_dict = model.module.state_dict()166    else:167        state_dict = model.state_dict()168    new_state_dict = {}169    for k, v in state_dict.items():170        try:171            # assert "dec" in k or "disc" in k172            # print("load", k)173            new_state_dict[k] = saved_state_dict[k]174            assert saved_state_dict[k].shape == v.shape, (saved_state_dict[k].shape, v.shape)175        except Exception:176            if "enc_q" not in k or "emb_g" not in k:177              print("%s is not in the checkpoint,please check your checkpoint.If you're using pretrain model,just ignore this warning." % k)178              logger.info("%s is not in the checkpoint" % k)179              new_state_dict[k] = v180    if hasattr(model, 'module'):181        model.module.load_state_dict(new_state_dict)182    else:183        model.load_state_dict(new_state_dict)184    print("load ")185    logger.info("Loaded checkpoint '{}' (iteration {})".format(186        checkpoint_path, iteration))187    return model, optimizer, learning_rate, iteration188 189 190def save_checkpoint(model, optimizer, learning_rate, iteration, checkpoint_path):191  logger.info("Saving model and optimizer state at iteration {} to {}".format(192    iteration, checkpoint_path))193  if hasattr(model, 'module'):194    state_dict = model.module.state_dict()195  else:196    state_dict = model.state_dict()197  torch.save({'model': state_dict,198              'iteration': iteration,199              'optimizer': optimizer.state_dict(),200              'learning_rate': learning_rate}, checkpoint_path)201 202def clean_checkpoints(path_to_models='logs/44k/', n_ckpts_to_keep=2, sort_by_time=True):203  """Freeing up space by deleting saved ckpts204 205  Arguments:206  path_to_models    --  Path to the model directory207  n_ckpts_to_keep   --  Number of ckpts to keep, excluding G_0.pth and D_0.pth208  sort_by_time      --  True -> chronologically delete ckpts209                        False -> lexicographically delete ckpts210  """211  ckpts_files = [f for f in os.listdir(path_to_models) if os.path.isfile(os.path.join(path_to_models, f))]212  def name_key(_f):213      return int(re.compile("._(\\d+)\\.pth").match(_f).group(1))214  def time_key(_f):215      return os.path.getmtime(os.path.join(path_to_models, _f))216  sort_key = time_key if sort_by_time else name_key217  def x_sorted(_x):218      return sorted([f for f in ckpts_files if f.startswith(_x) and not f.endswith("_0.pth")], key=sort_key)219  to_del = [os.path.join(path_to_models, fn) for fn in220            (x_sorted('G')[:-n_ckpts_to_keep] + x_sorted('D')[:-n_ckpts_to_keep])]221  def del_info(fn):222      return logger.info(f".. Free up space by deleting ckpt {fn}")223  def del_routine(x):224      return [os.remove(x), del_info(x)]225  [del_routine(fn) for fn in to_del]226 227def summarize(writer, global_step, scalars={}, histograms={}, images={}, audios={}, audio_sampling_rate=22050):228  for k, v in scalars.items():229    writer.add_scalar(k, v, global_step)230  for k, v in histograms.items():231    writer.add_histogram(k, v, global_step)232  for k, v in images.items():233    writer.add_image(k, v, global_step, dataformats='HWC')234  for k, v in audios.items():235    writer.add_audio(k, v, global_step, audio_sampling_rate)236 237 238def latest_checkpoint_path(dir_path, regex="G_*.pth"):239  f_list = glob.glob(os.path.join(dir_path, regex))240  f_list.sort(key=lambda f: int("".join(filter(str.isdigit, f))))241  x = f_list[-1]242  print(x)243  return x244 245 246def plot_spectrogram_to_numpy(spectrogram):247  global MATPLOTLIB_FLAG248  if not MATPLOTLIB_FLAG:249    import matplotlib250    matplotlib.use("Agg")251    MATPLOTLIB_FLAG = True252    mpl_logger = logging.getLogger('matplotlib')253    mpl_logger.setLevel(logging.WARNING)254  import matplotlib.pylab as plt255  import numpy as np256 257  fig, ax = plt.subplots(figsize=(10,2))258  im = ax.imshow(spectrogram, aspect="auto", origin="lower",259                  interpolation='none')260  plt.colorbar(im, ax=ax)261  plt.xlabel("Frames")262  plt.ylabel("Channels")263  plt.tight_layout()264 265  fig.canvas.draw()266  data = np.fromstring(fig.canvas.tostring_rgb(), dtype=np.uint8, sep='')267  data = data.reshape(fig.canvas.get_width_height()[::-1] + (3,))268  plt.close()269  return data270 271 272def plot_alignment_to_numpy(alignment, info=None):273  global MATPLOTLIB_FLAG274  if not MATPLOTLIB_FLAG:275    import matplotlib276    matplotlib.use("Agg")277    MATPLOTLIB_FLAG = True278    mpl_logger = logging.getLogger('matplotlib')279    mpl_logger.setLevel(logging.WARNING)280  import matplotlib.pylab as plt281  import numpy as np282 283  fig, ax = plt.subplots(figsize=(6, 4))284  im = ax.imshow(alignment.transpose(), aspect='auto', origin='lower',285                  interpolation='none')286  fig.colorbar(im, ax=ax)287  xlabel = 'Decoder timestep'288  if info is not None:289      xlabel += '\n\n' + info290  plt.xlabel(xlabel)291  plt.ylabel('Encoder timestep')292  plt.tight_layout()293 294  fig.canvas.draw()295  data = np.fromstring(fig.canvas.tostring_rgb(), dtype=np.uint8, sep='')296  data = data.reshape(fig.canvas.get_width_height()[::-1] + (3,))297  plt.close()298  return data299 300 301def load_wav_to_torch(full_path):302  sampling_rate, data = read(full_path)303  return torch.FloatTensor(data.astype(np.float32)), sampling_rate304 305 306def load_filepaths_and_text(filename, split="|"):307  with open(filename, encoding='utf-8') as f:308    filepaths_and_text = [line.strip().split(split) for line in f]309  return filepaths_and_text310 311 312def get_hparams(init=True):313  parser = argparse.ArgumentParser()314  parser.add_argument('-c', '--config', type=str, default="./configs/config.json",315                      help='JSON file for configuration')316  parser.add_argument('-m', '--model', type=str, required=True,317                      help='Model name')318 319  args = parser.parse_args()320  model_dir = os.path.join("./logs", args.model)321 322  if not os.path.exists(model_dir):323    os.makedirs(model_dir)324 325  config_path = args.config326  config_save_path = os.path.join(model_dir, "config.json")327  if init:328    with open(config_path, "r") as f:329      data = f.read()330    with open(config_save_path, "w") as f:331      f.write(data)332  else:333    with open(config_save_path, "r") as f:334      data = f.read()335  config = json.loads(data)336 337  hparams = HParams(**config)338  hparams.model_dir = model_dir339  return hparams340 341 342def get_hparams_from_dir(model_dir):343  config_save_path = os.path.join(model_dir, "config.json")344  with open(config_save_path, "r") as f:345    data = f.read()346  config = json.loads(data)347 348  hparams =HParams(**config)349  hparams.model_dir = model_dir350  return hparams351 352 353def get_hparams_from_file(config_path, infer_mode = False):354  with open(config_path, "r") as f:355    data = f.read()356  config = json.loads(data)357  hparams =HParams(**config) if not infer_mode else InferHParams(**config)358  return hparams359 360 361def check_git_hash(model_dir):362  source_dir = os.path.dirname(os.path.realpath(__file__))363  if not os.path.exists(os.path.join(source_dir, ".git")):364    logger.warn("{} is not a git repository, therefore hash value comparison will be ignored.".format(365      source_dir366    ))367    return368 369  cur_hash = subprocess.getoutput("git rev-parse HEAD")370 371  path = os.path.join(model_dir, "githash")372  if os.path.exists(path):373    saved_hash = open(path).read()374    if saved_hash != cur_hash:375      logger.warn("git hash values are different. {}(saved) != {}(current)".format(376        saved_hash[:8], cur_hash[:8]))377  else:378    open(path, "w").write(cur_hash)379 380 381def get_logger(model_dir, filename="train.log"):382  global logger383  logger = logging.getLogger(os.path.basename(model_dir))384  logger.setLevel(logging.DEBUG)385 386  formatter = logging.Formatter("%(asctime)s\t%(name)s\t%(levelname)s\t%(message)s")387  if not os.path.exists(model_dir):388    os.makedirs(model_dir)389  h = logging.FileHandler(os.path.join(model_dir, filename))390  h.setLevel(logging.DEBUG)391  h.setFormatter(formatter)392  logger.addHandler(h)393  return logger394 395 396def repeat_expand_2d(content, target_len, mode = 'left'):397    # content : [h, t]398    return repeat_expand_2d_left(content, target_len) if mode == 'left' else repeat_expand_2d_other(content, target_len, mode)399 400 401 402def repeat_expand_2d_left(content, target_len):403    # content : [h, t]404 405    src_len = content.shape[-1]406    target = torch.zeros([content.shape[0], target_len], dtype=torch.float).to(content.device)407    temp = torch.arange(src_len+1) * target_len / src_len408    current_pos = 0409    for i in range(target_len):410        if i < temp[current_pos+1]:411            target[:, i] = content[:, current_pos]412        else:413            current_pos += 1414            target[:, i] = content[:, current_pos]415 416    return target417 418 419# mode : 'nearest'| 'linear'| 'bilinear'| 'bicubic'| 'trilinear'| 'area'420def repeat_expand_2d_other(content, target_len, mode = 'nearest'):421    # content : [h, t]422    content = content[None,:,:]423    target = F.interpolate(content,size=target_len,mode=mode)[0]424    return target425 426 427def mix_model(model_paths,mix_rate,mode):428  mix_rate = torch.FloatTensor(mix_rate)/100429  model_tem = torch.load(model_paths[0])430  models = [torch.load(path)["model"] for path in model_paths]431  if mode == 0:432     mix_rate = F.softmax(mix_rate,dim=0)433  for k in model_tem["model"].keys():434     model_tem["model"][k] = torch.zeros_like(model_tem["model"][k])435     for i,model in enumerate(models):436        model_tem["model"][k] += model[k]*mix_rate[i]437  torch.save(model_tem,os.path.join(os.path.curdir,"output.pth"))438  return os.path.join(os.path.curdir,"output.pth")439  440def change_rms(data1, sr1, data2, sr2, rate):  # 1是输入音频，2是输出音频,rate是2的占比 from RVC441    # print(data1.max(),data2.max())442    rms1 = librosa.feature.rms(443        y=data1, frame_length=sr1 // 2 * 2, hop_length=sr1 // 2444    )  # 每半秒一个点445    rms2 = librosa.feature.rms(y=data2.detach().cpu().numpy(), frame_length=sr2 // 2 * 2, hop_length=sr2 // 2)446    rms1 = torch.from_numpy(rms1).to(data2.device)447    rms1 = F.interpolate(448        rms1.unsqueeze(0), size=data2.shape[0], mode="linear"449    ).squeeze()450    rms2 = torch.from_numpy(rms2).to(data2.device)451    rms2 = F.interpolate(452        rms2.unsqueeze(0), size=data2.shape[0], mode="linear"453    ).squeeze()454    rms2 = torch.max(rms2, torch.zeros_like(rms2) + 1e-6)455    data2 *= (456        torch.pow(rms1, torch.tensor(1 - rate))457        * torch.pow(rms2, torch.tensor(rate - 1))458    )459    return data2460 461def train_index(spk_name,root_dir = "dataset/44k/"):  #from: RVC https://github.com/RVC-Project/Retrieval-based-Voice-Conversion-WebUI462    n_cpu = cpu_count()463    print("The feature index is constructing.")464    exp_dir = os.path.join(root_dir,spk_name)465    listdir_res = []466    for file in os.listdir(exp_dir):467       if ".wav.soft.pt" in file:468          listdir_res.append(os.path.join(exp_dir,file))469    if len(listdir_res) == 0:470        raise Exception("You need to run preprocess_hubert_f0.py!")471    npys = []472    for name in sorted(listdir_res):473        phone = torch.load(name)[0].transpose(-1,-2).numpy()474        npys.append(phone)475    big_npy = np.concatenate(npys, 0)476    big_npy_idx = np.arange(big_npy.shape[0])477    np.random.shuffle(big_npy_idx)478    big_npy = big_npy[big_npy_idx]479    if big_npy.shape[0] > 2e5:480        # if(1):481        info = "Trying doing kmeans %s shape to 10k centers." % big_npy.shape[0]482        print(info)483        try:484            big_npy = (485                MiniBatchKMeans(486                    n_clusters=10000,487                    verbose=True,488                    batch_size=256 * n_cpu,489                    compute_labels=False,490                    init="random",491                )492                .fit(big_npy)493                .cluster_centers_494            )495        except Exception:496            info = traceback.format_exc()497            print(info)498    n_ivf = min(int(16 * np.sqrt(big_npy.shape[0])), big_npy.shape[0] // 39)499    index = faiss.index_factory(big_npy.shape[1] , "IVF%s,Flat" % n_ivf)500    index_ivf = faiss.extract_index_ivf(index)  #501    index_ivf.nprobe = 1502    index.train(big_npy)503    batch_size_add = 8192504    for i in range(0, big_npy.shape[0], batch_size_add):505        index.add(big_npy[i : i + batch_size_add])506    # faiss.write_index(507    #     index,508    #     f"added_{spk_name}.index"509    # )510    print("Successfully build index")511    return index512 513 514class HParams():515  def __init__(self, **kwargs):516    for k, v in kwargs.items():517      if type(v) == dict:518        v = HParams(**v)519      self[k] = v520 521  def keys(self):522    return self.__dict__.keys()523 524  def items(self):525    return self.__dict__.items()526 527  def values(self):528    return self.__dict__.values()529 530  def __len__(self):531    return len(self.__dict__)532 533  def __getitem__(self, key):534    return getattr(self, key)535 536  def __setitem__(self, key, value):537    return setattr(self, key, value)538 539  def __contains__(self, key):540    return key in self.__dict__541 542  def __repr__(self):543    return self.__dict__.__repr__()544 545  def get(self,index):546    return self.__dict__.get(index)547 548  549class InferHParams(HParams):550  def __init__(self, **kwargs):551    for k, v in kwargs.items():552      if type(v) == dict:553        v = InferHParams(**v)554      self[k] = v555 556  def __getattr__(self,index):557    return self.get(index)558 559 560class Volume_Extractor:561    def __init__(self, hop_size = 512):562        self.hop_size = hop_size563        564    def extract(self, audio): # audio: 2d tensor array565        if not isinstance(audio,torch.Tensor):566           audio = torch.Tensor(audio)567        n_frames = int(audio.size(-1) // self.hop_size)568        audio2 = audio ** 2569        audio2 = torch.nn.functional.pad(audio2, (int(self.hop_size // 2), int((self.hop_size + 1) // 2)), mode = 'reflect')570        volume = torch.nn.functional.unfold(audio2[:,None,None,:],(1,self.hop_size),stride=self.hop_size)[:,:,:n_frames].mean(dim=1)[0]571        volume = torch.sqrt(volume)572        return volume573