Doiyan/vits-models

1import os2import sys3import argparse4import logging5import json6import subprocess7import numpy as np8import librosa9import torch10 11MATPLOTLIB_FLAG = False12 13logging.basicConfig(stream=sys.stdout, level=logging.DEBUG)14logger = logging15 16 17def load_checkpoint(checkpoint_path, model, optimizer=None):18  assert os.path.isfile(checkpoint_path)19  checkpoint_dict = torch.load(checkpoint_path, map_location='cpu')20  iteration = checkpoint_dict['iteration']21  learning_rate = checkpoint_dict['learning_rate']22  if optimizer is not None:23    optimizer.load_state_dict(checkpoint_dict['optimizer'])24  saved_state_dict = checkpoint_dict['model']25  if hasattr(model, 'module'):26    state_dict = model.module.state_dict()27  else:28    state_dict = model.state_dict()29  new_state_dict= {}30  for k, v in state_dict.items():31    try:32      new_state_dict[k] = saved_state_dict[k]33    except:34      logger.info("%s is not in the checkpoint" % k)35      new_state_dict[k] = v36  if hasattr(model, 'module'):37    model.module.load_state_dict(new_state_dict)38  else:39    model.load_state_dict(new_state_dict)40  logger.info("Loaded checkpoint '{}' (iteration {})" .format(41    checkpoint_path, iteration))42  return model, optimizer, learning_rate, iteration43 44 45def plot_spectrogram_to_numpy(spectrogram):46  global MATPLOTLIB_FLAG47  if not MATPLOTLIB_FLAG:48    import matplotlib49    matplotlib.use("Agg")50    MATPLOTLIB_FLAG = True51    mpl_logger = logging.getLogger('matplotlib')52    mpl_logger.setLevel(logging.WARNING)53  import matplotlib.pylab as plt54  import numpy as np55 56  fig, ax = plt.subplots(figsize=(10,2))57  im = ax.imshow(spectrogram, aspect="auto", origin="lower",58                  interpolation='none')59  plt.colorbar(im, ax=ax)60  plt.xlabel("Frames")61  plt.ylabel("Channels")62  plt.tight_layout()63 64  fig.canvas.draw()65  data = np.fromstring(fig.canvas.tostring_rgb(), dtype=np.uint8, sep='')66  data = data.reshape(fig.canvas.get_width_height()[::-1] + (3,))67  plt.close()68  return data69 70 71def plot_alignment_to_numpy(alignment, info=None):72  global MATPLOTLIB_FLAG73  if not MATPLOTLIB_FLAG:74    import matplotlib75    matplotlib.use("Agg")76    MATPLOTLIB_FLAG = True77    mpl_logger = logging.getLogger('matplotlib')78    mpl_logger.setLevel(logging.WARNING)79  import matplotlib.pylab as plt80  import numpy as np81 82  fig, ax = plt.subplots(figsize=(6, 4))83  im = ax.imshow(alignment.transpose(), aspect='auto', origin='lower',84                  interpolation='none')85  fig.colorbar(im, ax=ax)86  xlabel = 'Decoder timestep'87  if info is not None:88      xlabel += '\n\n' + info89  plt.xlabel(xlabel)90  plt.ylabel('Encoder timestep')91  plt.tight_layout()92 93  fig.canvas.draw()94  data = np.fromstring(fig.canvas.tostring_rgb(), dtype=np.uint8, sep='')95  data = data.reshape(fig.canvas.get_width_height()[::-1] + (3,))96  plt.close()97  return data98 99 100def load_audio_to_torch(full_path, target_sampling_rate):101  audio, sampling_rate = librosa.load(full_path, sr=target_sampling_rate, mono=True)102  return torch.FloatTensor(audio.astype(np.float32))103 104 105def load_filepaths_and_text(filename, split="|"):106  with open(filename, encoding='utf-8') as f:107    filepaths_and_text = [line.strip().split(split) for line in f]108  return filepaths_and_text109 110 111def get_hparams(init=True):112  parser = argparse.ArgumentParser()113  parser.add_argument('-c', '--config', type=str, default="./configs/base.json",114                      help='JSON file for configuration')115  parser.add_argument('-m', '--model', type=str, required=True,116                      help='Model name')117 118  args = parser.parse_args()119  model_dir = os.path.join("./logs", args.model)120 121  if not os.path.exists(model_dir):122    os.makedirs(model_dir)123 124  config_path = args.config125  config_save_path = os.path.join(model_dir, "config.json")126  if init:127    with open(config_path, "r") as f:128      data = f.read()129    with open(config_save_path, "w") as f:130      f.write(data)131  else:132    with open(config_save_path, "r") as f:133      data = f.read()134  config = json.loads(data)135 136  hparams = HParams(**config)137  hparams.model_dir = model_dir138  return hparams139 140 141def get_hparams_from_dir(model_dir):142  config_save_path = os.path.join(model_dir, "config.json")143  with open(config_save_path, "r") as f:144    data = f.read()145  config = json.loads(data)146 147  hparams =HParams(**config)148  hparams.model_dir = model_dir149  return hparams150 151 152def get_hparams_from_file(config_path):153  with open(config_path, "r") as f:154    data = f.read()155  config = json.loads(data)156 157  hparams =HParams(**config)158  return hparams159 160 161def check_git_hash(model_dir):162  source_dir = os.path.dirname(os.path.realpath(__file__))163  if not os.path.exists(os.path.join(source_dir, ".git")):164    logger.warn("{} is not a git repository, therefore hash value comparison will be ignored.".format(165      source_dir166    ))167    return168 169  cur_hash = subprocess.getoutput("git rev-parse HEAD")170 171  path = os.path.join(model_dir, "githash")172  if os.path.exists(path):173    saved_hash = open(path).read()174    if saved_hash != cur_hash:175      logger.warn("git hash values are different. {}(saved) != {}(current)".format(176        saved_hash[:8], cur_hash[:8]))177  else:178    open(path, "w").write(cur_hash)179 180 181def get_logger(model_dir, filename="train.log"):182  global logger183  logger = logging.getLogger(os.path.basename(model_dir))184  logger.setLevel(logging.DEBUG)185 186  formatter = logging.Formatter("%(asctime)s\t%(name)s\t%(levelname)s\t%(message)s")187  if not os.path.exists(model_dir):188    os.makedirs(model_dir)189  h = logging.FileHandler(os.path.join(model_dir, filename))190  h.setLevel(logging.DEBUG)191  h.setFormatter(formatter)192  logger.addHandler(h)193  return logger194 195 196class HParams():197  def __init__(self, **kwargs):198    for k, v in kwargs.items():199      if type(v) == dict:200        v = HParams(**v)201      self[k] = v202 203  def keys(self):204    return self.__dict__.keys()205 206  def items(self):207    return self.__dict__.items()208 209  def values(self):210    return self.__dict__.values()211 212  def __len__(self):213    return len(self.__dict__)214 215  def __getitem__(self, key):216    return getattr(self, key)217 218  def __setitem__(self, key, value):219    return setattr(self, key, value)220 221  def __contains__(self, key):222    return key in self.__dict__223 224  def __repr__(self):225    return self.__dict__.__repr__()226