kevinwang676/FreeVC-OpenAI-TTS

1import os2import sys3import argparse4import logging5import json6import subprocess7import numpy as np8from scipy.io.wavfile import read9import torch10from torch.nn import functional as F11from commons import sequence_mask12 13MATPLOTLIB_FLAG = False14 15logging.basicConfig(stream=sys.stdout, level=logging.DEBUG)16logger = logging17 18 19def get_cmodel(rank):20    checkpoint = torch.load('wavlm/WavLM-Large.pt')21    cfg = WavLMConfig(checkpoint['cfg'])22    cmodel = WavLM(cfg).cuda(rank)23    cmodel.load_state_dict(checkpoint['model'])24    cmodel.eval()25    return cmodel26    27    28def get_content(cmodel, y):29    with torch.no_grad():30        c = cmodel.extract_features(y.squeeze(1))[0]31    c = c.transpose(1, 2)32    return c33 34 35def get_vocoder(rank):36    with open("hifigan/config.json", "r") as f:37        config = json.load(f)38    config = hifigan.AttrDict(config)39    vocoder = hifigan.Generator(config)40    ckpt = torch.load("hifigan/generator_v1")41    vocoder.load_state_dict(ckpt["generator"])42    vocoder.eval()43    vocoder.remove_weight_norm()44    vocoder.cuda(rank)45    return vocoder46    47    48def transform(mel, height): # 68-9249    #r = np.random.random()50    #rate = r * 0.3 + 0.85 # 0.85-1.1551    #height = int(mel.size(-2) * rate)52    tgt = torchvision.transforms.functional.resize(mel, (height, mel.size(-1)))53    if height >= mel.size(-2):54        return tgt[:, :mel.size(-2), :]55    else:56        silence = tgt[:,-1:,:].repeat(1,mel.size(-2)-height,1) 57        silence += torch.randn_like(silence) / 1058        return torch.cat((tgt, silence), 1)59        60        61def stretch(mel, width): # 0.5-262    return torchvision.transforms.functional.resize(mel, (mel.size(-2), width))63 64 65def load_checkpoint(checkpoint_path, model, optimizer=None):66  assert os.path.isfile(checkpoint_path)67  checkpoint_dict = torch.load(checkpoint_path, map_location='cpu')68  iteration = checkpoint_dict['iteration']69  learning_rate = checkpoint_dict['learning_rate']70  if optimizer is not None:71    optimizer.load_state_dict(checkpoint_dict['optimizer'])72  saved_state_dict = checkpoint_dict['model']73  if hasattr(model, 'module'):74    state_dict = model.module.state_dict()75  else:76    state_dict = model.state_dict()77  new_state_dict= {}78  for k, v in state_dict.items():79    try:80      new_state_dict[k] = saved_state_dict[k]81    except:82      logger.info("%s is not in the checkpoint" % k)83      new_state_dict[k] = v84  if hasattr(model, 'module'):85    model.module.load_state_dict(new_state_dict)86  else:87    model.load_state_dict(new_state_dict)88  logger.info("Loaded checkpoint '{}' (iteration {})" .format(89    checkpoint_path, iteration))90  return model, optimizer, learning_rate, iteration91 92 93def save_checkpoint(model, optimizer, learning_rate, iteration, checkpoint_path):94  logger.info("Saving model and optimizer state at iteration {} to {}".format(95    iteration, checkpoint_path))96  if hasattr(model, 'module'):97    state_dict = model.module.state_dict()98  else:99    state_dict = model.state_dict()100  torch.save({'model': state_dict,101              'iteration': iteration,102              'optimizer': optimizer.state_dict(),103              'learning_rate': learning_rate}, checkpoint_path)104 105 106def summarize(writer, global_step, scalars={}, histograms={}, images={}, audios={}, audio_sampling_rate=22050):107  for k, v in scalars.items():108    writer.add_scalar(k, v, global_step)109  for k, v in histograms.items():110    writer.add_histogram(k, v, global_step)111  for k, v in images.items():112    writer.add_image(k, v, global_step, dataformats='HWC')113  for k, v in audios.items():114    writer.add_audio(k, v, global_step, audio_sampling_rate)115 116 117def latest_checkpoint_path(dir_path, regex="G_*.pth"):118  f_list = glob.glob(os.path.join(dir_path, regex))119  f_list.sort(key=lambda f: int("".join(filter(str.isdigit, f))))120  x = f_list[-1]121  print(x)122  return x123 124 125def plot_spectrogram_to_numpy(spectrogram):126  global MATPLOTLIB_FLAG127  if not MATPLOTLIB_FLAG:128    import matplotlib129    matplotlib.use("Agg")130    MATPLOTLIB_FLAG = True131    mpl_logger = logging.getLogger('matplotlib')132    mpl_logger.setLevel(logging.WARNING)133  import matplotlib.pylab as plt134  import numpy as np135  136  fig, ax = plt.subplots(figsize=(10,2))137  im = ax.imshow(spectrogram, aspect="auto", origin="lower",138                  interpolation='none')139  plt.colorbar(im, ax=ax)140  plt.xlabel("Frames")141  plt.ylabel("Channels")142  plt.tight_layout()143 144  fig.canvas.draw()145  data = np.fromstring(fig.canvas.tostring_rgb(), dtype=np.uint8, sep='')146  data = data.reshape(fig.canvas.get_width_height()[::-1] + (3,))147  plt.close()148  return data149 150 151def plot_alignment_to_numpy(alignment, info=None):152  global MATPLOTLIB_FLAG153  if not MATPLOTLIB_FLAG:154    import matplotlib155    matplotlib.use("Agg")156    MATPLOTLIB_FLAG = True157    mpl_logger = logging.getLogger('matplotlib')158    mpl_logger.setLevel(logging.WARNING)159  import matplotlib.pylab as plt160  import numpy as np161 162  fig, ax = plt.subplots(figsize=(6, 4))163  im = ax.imshow(alignment.transpose(), aspect='auto', origin='lower',164                  interpolation='none')165  fig.colorbar(im, ax=ax)166  xlabel = 'Decoder timestep'167  if info is not None:168      xlabel += '\n\n' + info169  plt.xlabel(xlabel)170  plt.ylabel('Encoder timestep')171  plt.tight_layout()172 173  fig.canvas.draw()174  data = np.fromstring(fig.canvas.tostring_rgb(), dtype=np.uint8, sep='')175  data = data.reshape(fig.canvas.get_width_height()[::-1] + (3,))176  plt.close()177  return data178 179 180def load_wav_to_torch(full_path):181  sampling_rate, data = read(full_path)182  return torch.FloatTensor(data.astype(np.float32)), sampling_rate183 184 185def load_filepaths_and_text(filename, split="|"):186  with open(filename, encoding='utf-8') as f:187    filepaths_and_text = [line.strip().split(split) for line in f]188  return filepaths_and_text189 190 191def get_hparams(init=True):192  parser = argparse.ArgumentParser()193  parser.add_argument('-c', '--config', type=str, default="./configs/base.json",194                      help='JSON file for configuration')195  parser.add_argument('-m', '--model', type=str, required=True,196                      help='Model name')197  198  args = parser.parse_args()199  model_dir = os.path.join("./logs", args.model)200 201  if not os.path.exists(model_dir):202    os.makedirs(model_dir)203 204  config_path = args.config205  config_save_path = os.path.join(model_dir, "config.json")206  if init:207    with open(config_path, "r") as f:208      data = f.read()209    with open(config_save_path, "w") as f:210      f.write(data)211  else:212    with open(config_save_path, "r") as f:213      data = f.read()214  config = json.loads(data)215  216  hparams = HParams(**config)217  hparams.model_dir = model_dir218  return hparams219 220 221def get_hparams_from_dir(model_dir):222  config_save_path = os.path.join(model_dir, "config.json")223  with open(config_save_path, "r") as f:224    data = f.read()225  config = json.loads(data)226 227  hparams =HParams(**config)228  hparams.model_dir = model_dir229  return hparams230 231 232def get_hparams_from_file(config_path):233  with open(config_path, "r") as f:234    data = f.read()235  config = json.loads(data)236 237  hparams =HParams(**config)238  return hparams239 240 241def check_git_hash(model_dir):242  source_dir = os.path.dirname(os.path.realpath(__file__))243  if not os.path.exists(os.path.join(source_dir, ".git")):244    logger.warn("{} is not a git repository, therefore hash value comparison will be ignored.".format(245      source_dir246    ))247    return248 249  cur_hash = subprocess.getoutput("git rev-parse HEAD")250 251  path = os.path.join(model_dir, "githash")252  if os.path.exists(path):253    saved_hash = open(path).read()254    if saved_hash != cur_hash:255      logger.warn("git hash values are different. {}(saved) != {}(current)".format(256        saved_hash[:8], cur_hash[:8]))257  else:258    open(path, "w").write(cur_hash)259 260 261def get_logger(model_dir, filename="train.log"):262  global logger263  logger = logging.getLogger(os.path.basename(model_dir))264  logger.setLevel(logging.DEBUG)265  266  formatter = logging.Formatter("%(asctime)s\t%(name)s\t%(levelname)s\t%(message)s")267  if not os.path.exists(model_dir):268    os.makedirs(model_dir)269  h = logging.FileHandler(os.path.join(model_dir, filename))270  h.setLevel(logging.DEBUG)271  h.setFormatter(formatter)272  logger.addHandler(h)273  return logger274 275 276class HParams():277  def __init__(self, **kwargs):278    for k, v in kwargs.items():279      if type(v) == dict:280        v = HParams(**v)281      self[k] = v282    283  def keys(self):284    return self.__dict__.keys()285 286  def items(self):287    return self.__dict__.items()288 289  def values(self):290    return self.__dict__.values()291 292  def __len__(self):293    return len(self.__dict__)294 295  def __getitem__(self, key):296    return getattr(self, key)297 298  def __setitem__(self, key, value):299    return setattr(self, key, value)300 301  def __contains__(self, key):302    return key in self.__dict__303 304  def __repr__(self):305    return self.__dict__.__repr__()306