chilge/nemo

1import os2import glob3import sys4import argparse5import logging6import json7import subprocess8 9import librosa10import numpy as np11import torchaudio12from scipy.io.wavfile import read13import torch14import torchvision15from torch.nn import functional as F16from commons import sequence_mask17from hubert import hubert_model18MATPLOTLIB_FLAG = False19 20logging.basicConfig(stream=sys.stdout, level=logging.DEBUG)21logger = logging22 23f0_bin = 25624f0_max = 1100.025f0_min = 50.026f0_mel_min = 1127 * np.log(1 + f0_min / 700)27f0_mel_max = 1127 * np.log(1 + f0_max / 700)28 29def f0_to_coarse(f0):30  is_torch = isinstance(f0, torch.Tensor)31  f0_mel = 1127 * (1 + f0 / 700).log() if is_torch else 1127 * np.log(1 + f0 / 700)32  f0_mel[f0_mel > 0] = (f0_mel[f0_mel > 0] - f0_mel_min) * (f0_bin - 2) / (f0_mel_max - f0_mel_min) + 133 34  f0_mel[f0_mel <= 1] = 135  f0_mel[f0_mel > f0_bin - 1] = f0_bin - 136  f0_coarse = (f0_mel + 0.5).long() if is_torch else np.rint(f0_mel).astype(np.int)37  assert f0_coarse.max() <= 255 and f0_coarse.min() >= 1, (f0_coarse.max(), f0_coarse.min())38  return f0_coarse39 40 41def get_hubert_model(rank=None):42 43  hubert_soft = hubert_model.hubert_soft("hubert/hubert-soft-0d54a1f4.pt")44  if rank is not None:45    hubert_soft = hubert_soft.cuda(rank)46  return hubert_soft47 48def get_hubert_content(hmodel, y=None, path=None):49  if path is not None:50    source, sr = torchaudio.load(path)51    source = torchaudio.functional.resample(source, sr, 16000)52    if len(source.shape) == 2 and source.shape[1] >= 2:53      source = torch.mean(source, dim=0).unsqueeze(0)54  else:55    source = y56  source = source.unsqueeze(0)57  with torch.inference_mode():58    units = hmodel.units(source)59    return units.transpose(1,2)60 61 62def get_content(cmodel, y):63    with torch.no_grad():64        c = cmodel.extract_features(y.squeeze(1))[0]65    c = c.transpose(1, 2)66    return c67 68 69 70def transform(mel, height): # 68-9271    #r = np.random.random()72    #rate = r * 0.3 + 0.85 # 0.85-1.1573    #height = int(mel.size(-2) * rate)74    tgt = torchvision.transforms.functional.resize(mel, (height, mel.size(-1)))75    if height >= mel.size(-2):76        return tgt[:, :mel.size(-2), :]77    else:78        silence = tgt[:,-1:,:].repeat(1,mel.size(-2)-height,1)79        silence += torch.randn_like(silence) / 1080        return torch.cat((tgt, silence), 1)81 82 83def stretch(mel, width): # 0.5-284    return torchvision.transforms.functional.resize(mel, (mel.size(-2), width))85 86 87def load_checkpoint(checkpoint_path, model, optimizer=None):88  assert os.path.isfile(checkpoint_path)89  checkpoint_dict = torch.load(checkpoint_path, map_location='cpu')90  iteration = checkpoint_dict['iteration']91  learning_rate = checkpoint_dict['learning_rate']92  if iteration is None:93    iteration = 194  if learning_rate is None:95    learning_rate = 0.000296  if optimizer is not None and checkpoint_dict['optimizer'] is not None:97    optimizer.load_state_dict(checkpoint_dict['optimizer'])98  saved_state_dict = checkpoint_dict['model']99  if hasattr(model, 'module'):100    state_dict = model.module.state_dict()101  else:102    state_dict = model.state_dict()103  new_state_dict= {}104  for k, v in state_dict.items():105    try:106      new_state_dict[k] = saved_state_dict[k]107    except:108      logger.info("%s is not in the checkpoint" % k)109      new_state_dict[k] = v110  if hasattr(model, 'module'):111    model.module.load_state_dict(new_state_dict)112  else:113    model.load_state_dict(new_state_dict)114  logger.info("Loaded checkpoint '{}' (iteration {})" .format(115    checkpoint_path, iteration))116  return model, optimizer, learning_rate, iteration117 118 119def save_checkpoint(model, optimizer, learning_rate, iteration, checkpoint_path):120  # ckptname = checkpoint_path.split(os.sep)[-1]121  # newest_step = int(ckptname.split(".")[0].split("_")[1])122  # val_steps = 2000123  # last_ckptname = checkpoint_path.replace(str(newest_step), str(newest_step - val_steps*3))124  # if newest_step >= val_steps*3:125  #   os.system(f"rm {last_ckptname}")126  logger.info("Saving model and optimizer state at iteration {} to {}".format(127    iteration, checkpoint_path))128  if hasattr(model, 'module'):129    state_dict = model.module.state_dict()130  else:131    state_dict = model.state_dict()132  torch.save({'model': state_dict,133              'iteration': iteration,134              'optimizer': optimizer.state_dict(),135              'learning_rate': learning_rate}, checkpoint_path)136 137 138def summarize(writer, global_step, scalars={}, histograms={}, images={}, audios={}, audio_sampling_rate=22050):139  for k, v in scalars.items():140    writer.add_scalar(k, v, global_step)141  for k, v in histograms.items():142    writer.add_histogram(k, v, global_step)143  for k, v in images.items():144    writer.add_image(k, v, global_step, dataformats='HWC')145  for k, v in audios.items():146    writer.add_audio(k, v, global_step, audio_sampling_rate)147 148 149def latest_checkpoint_path(dir_path, regex="G_*.pth"):150  f_list = glob.glob(os.path.join(dir_path, regex))151  f_list.sort(key=lambda f: int("".join(filter(str.isdigit, f))))152  x = f_list[-1]153  print(x)154  return x155 156 157def plot_spectrogram_to_numpy(spectrogram):158  global MATPLOTLIB_FLAG159  if not MATPLOTLIB_FLAG:160    import matplotlib161    matplotlib.use("Agg")162    MATPLOTLIB_FLAG = True163    mpl_logger = logging.getLogger('matplotlib')164    mpl_logger.setLevel(logging.WARNING)165  import matplotlib.pylab as plt166  import numpy as np167 168  fig, ax = plt.subplots(figsize=(10,2))169  im = ax.imshow(spectrogram, aspect="auto", origin="lower",170                  interpolation='none')171  plt.colorbar(im, ax=ax)172  plt.xlabel("Frames")173  plt.ylabel("Channels")174  plt.tight_layout()175 176  fig.canvas.draw()177  data = np.fromstring(fig.canvas.tostring_rgb(), dtype=np.uint8, sep='')178  data = data.reshape(fig.canvas.get_width_height()[::-1] + (3,))179  plt.close()180  return data181 182 183def plot_alignment_to_numpy(alignment, info=None):184  global MATPLOTLIB_FLAG185  if not MATPLOTLIB_FLAG:186    import matplotlib187    matplotlib.use("Agg")188    MATPLOTLIB_FLAG = True189    mpl_logger = logging.getLogger('matplotlib')190    mpl_logger.setLevel(logging.WARNING)191  import matplotlib.pylab as plt192  import numpy as np193 194  fig, ax = plt.subplots(figsize=(6, 4))195  im = ax.imshow(alignment.transpose(), aspect='auto', origin='lower',196                  interpolation='none')197  fig.colorbar(im, ax=ax)198  xlabel = 'Decoder timestep'199  if info is not None:200      xlabel += '\n\n' + info201  plt.xlabel(xlabel)202  plt.ylabel('Encoder timestep')203  plt.tight_layout()204 205  fig.canvas.draw()206  data = np.fromstring(fig.canvas.tostring_rgb(), dtype=np.uint8, sep='')207  data = data.reshape(fig.canvas.get_width_height()[::-1] + (3,))208  plt.close()209  return data210 211 212def load_wav_to_torch(full_path):213  sampling_rate, data = read(full_path)214  return torch.FloatTensor(data.astype(np.float32)), sampling_rate215 216 217def load_filepaths_and_text(filename, split="|"):218  with open(filename, encoding='utf-8') as f:219    filepaths_and_text = [line.strip().split(split) for line in f]220  return filepaths_and_text221 222 223def get_hparams(init=True):224  parser = argparse.ArgumentParser()225  parser.add_argument('-c', '--config', type=str, default="./configs/base.json",226                      help='JSON file for configuration')227  parser.add_argument('-m', '--model', type=str, required=True,228                      help='Model name')229 230  args = parser.parse_args()231  model_dir = os.path.join("./logs", args.model)232 233  if not os.path.exists(model_dir):234    os.makedirs(model_dir)235 236  config_path = args.config237  config_save_path = os.path.join(model_dir, "config.json")238  if init:239    with open(config_path, "r") as f:240      data = f.read()241    with open(config_save_path, "w") as f:242      f.write(data)243  else:244    with open(config_save_path, "r") as f:245      data = f.read()246  config = json.loads(data)247 248  hparams = HParams(**config)249  hparams.model_dir = model_dir250  return hparams251 252 253def get_hparams_from_dir(model_dir):254  config_save_path = os.path.join(model_dir, "config.json")255  with open(config_save_path, "r") as f:256    data = f.read()257  config = json.loads(data)258 259  hparams =HParams(**config)260  hparams.model_dir = model_dir261  return hparams262 263 264def get_hparams_from_file(config_path):265  with open(config_path, "r") as f:266    data = f.read()267  config = json.loads(data)268 269  hparams =HParams(**config)270  return hparams271 272 273def check_git_hash(model_dir):274  source_dir = os.path.dirname(os.path.realpath(__file__))275  if not os.path.exists(os.path.join(source_dir, ".git")):276    logger.warn("{} is not a git repository, therefore hash value comparison will be ignored.".format(277      source_dir278    ))279    return280 281  cur_hash = subprocess.getoutput("git rev-parse HEAD")282 283  path = os.path.join(model_dir, "githash")284  if os.path.exists(path):285    saved_hash = open(path).read()286    if saved_hash != cur_hash:287      logger.warn("git hash values are different. {}(saved) != {}(current)".format(288        saved_hash[:8], cur_hash[:8]))289  else:290    open(path, "w").write(cur_hash)291 292 293def get_logger(model_dir, filename="train.log"):294  global logger295  logger = logging.getLogger(os.path.basename(model_dir))296  logger.setLevel(logging.DEBUG)297 298  formatter = logging.Formatter("%(asctime)s\t%(name)s\t%(levelname)s\t%(message)s")299  if not os.path.exists(model_dir):300    os.makedirs(model_dir)301  h = logging.FileHandler(os.path.join(model_dir, filename))302  h.setLevel(logging.DEBUG)303  h.setFormatter(formatter)304  logger.addHandler(h)305  return logger306 307 308class HParams():309  def __init__(self, **kwargs):310    for k, v in kwargs.items():311      if type(v) == dict:312        v = HParams(**v)313      self[k] = v314 315  def keys(self):316    return self.__dict__.keys()317 318  def items(self):319    return self.__dict__.items()320 321  def values(self):322    return self.__dict__.values()323 324  def __len__(self):325    return len(self.__dict__)326 327  def __getitem__(self, key):328    return getattr(self, key)329 330  def __setitem__(self, key, value):331    return setattr(self, key, value)332 333  def __contains__(self, key):334    return key in self.__dict__335 336  def __repr__(self):337    return self.__dict__.__repr__()338 339