|
- import os
- import glob
- import argparse
- import logging
- import json
- import shutil
- import subprocess
- import numpy as np
- from huggingface_hub import hf_hub_download
- from scipy.io.wavfile import read
- import torch
- import re
-
- MATPLOTLIB_FLAG = False
-
- logger = logging.getLogger(__name__)
-
-
- def download_emo_models(mirror, repo_id, model_name):
- if mirror == "openi":
- import openi
-
- openi.model.download_model(
- "Stardust_minus/Bert-VITS2",
- repo_id.split("/")[-1],
- "./emotional",
- )
- else:
- hf_hub_download(
- repo_id,
- "pytorch_model.bin",
- local_dir=model_name,
- local_dir_use_symlinks=False,
- )
-
-
- def download_checkpoint(
- dir_path, repo_config, token=None, regex="G_*.pth", mirror="openi"
- ):
- repo_id = repo_config["repo_id"]
- f_list = glob.glob(os.path.join(dir_path, regex))
- if f_list:
- print("Use existed model, skip downloading.")
- return
- if mirror.lower() == "openi":
- import openi
-
- kwargs = {"token": token} if token else {}
- openi.login(**kwargs)
-
- model_image = repo_config["model_image"]
- openi.model.download_model(repo_id, model_image, dir_path)
-
- fs = glob.glob(os.path.join(dir_path, model_image, "*.pth"))
- for file in fs:
- shutil.move(file, dir_path)
- shutil.rmtree(os.path.join(dir_path, model_image))
- else:
- for file in ["DUR_0.pth", "D_0.pth", "G_0.pth"]:
- hf_hub_download(
- repo_id, file, local_dir=dir_path, local_dir_use_symlinks=False
- )
-
-
- def load_checkpoint(checkpoint_path, model, optimizer=None, skip_optimizer=False):
- assert os.path.isfile(checkpoint_path)
- checkpoint_dict = torch.load(checkpoint_path, map_location="cpu")
- iteration = checkpoint_dict["iteration"]
- learning_rate = checkpoint_dict["learning_rate"]
- if (
- optimizer is not None
- and not skip_optimizer
- and checkpoint_dict["optimizer"] is not None
- ):
- optimizer.load_state_dict(checkpoint_dict["optimizer"])
- elif optimizer is None and not skip_optimizer:
- # else: Disable this line if Infer and resume checkpoint,then enable the line upper
- new_opt_dict = optimizer.state_dict()
- new_opt_dict_params = new_opt_dict["param_groups"][0]["params"]
- new_opt_dict["param_groups"] = checkpoint_dict["optimizer"]["param_groups"]
- new_opt_dict["param_groups"][0]["params"] = new_opt_dict_params
- optimizer.load_state_dict(new_opt_dict)
-
- saved_state_dict = checkpoint_dict["model"]
- if hasattr(model, "module"):
- state_dict = model.module.state_dict()
- else:
- state_dict = model.state_dict()
-
- new_state_dict = {}
- for k, v in state_dict.items():
- try:
- # assert "emb_g" not in k
- new_state_dict[k] = saved_state_dict[k]
- assert saved_state_dict[k].shape == v.shape, (
- saved_state_dict[k].shape,
- v.shape,
- )
- except:
- # For upgrading from the old version
- if "ja_bert_proj" in k:
- v = torch.zeros_like(v)
- logger.warn(
- f"Seems you are using the old version of the model, the {k} is automatically set to zero for backward compatibility"
- )
- else:
- logger.error(f"{k} is not in the checkpoint")
-
- new_state_dict[k] = v
-
- if hasattr(model, "module"):
- model.module.load_state_dict(new_state_dict, strict=False)
- else:
- model.load_state_dict(new_state_dict, strict=False)
-
- logger.info(
- "Loaded checkpoint '{}' (iteration {})".format(checkpoint_path, iteration)
- )
-
- return model, optimizer, learning_rate, iteration
-
-
- def save_checkpoint(model, optimizer, learning_rate, iteration, checkpoint_path):
- logger.info(
- "Saving model and optimizer state at iteration {} to {}".format(
- iteration, checkpoint_path
- )
- )
- if hasattr(model, "module"):
- state_dict = model.module.state_dict()
- else:
- state_dict = model.state_dict()
- torch.save(
- {
- "model": state_dict,
- "iteration": iteration,
- "optimizer": optimizer.state_dict(),
- "learning_rate": learning_rate,
- },
- checkpoint_path,
- )
-
-
- def summarize(
- writer,
- global_step,
- scalars={},
- histograms={},
- images={},
- audios={},
- audio_sampling_rate=22050,
- ):
- for k, v in scalars.items():
- writer.add_scalar(k, v, global_step)
- for k, v in histograms.items():
- writer.add_histogram(k, v, global_step)
- for k, v in images.items():
- writer.add_image(k, v, global_step, dataformats="HWC")
- for k, v in audios.items():
- writer.add_audio(k, v, global_step, audio_sampling_rate)
-
-
- def latest_checkpoint_path(dir_path, regex="G_*.pth"):
- f_list = glob.glob(os.path.join(dir_path, regex))
- f_list.sort(key=lambda f: int("".join(filter(str.isdigit, f))))
- x = f_list[-1]
- return x
-
-
- def plot_spectrogram_to_numpy(spectrogram):
- global MATPLOTLIB_FLAG
- if not MATPLOTLIB_FLAG:
- import matplotlib
-
- matplotlib.use("Agg")
- MATPLOTLIB_FLAG = True
- mpl_logger = logging.getLogger("matplotlib")
- mpl_logger.setLevel(logging.WARNING)
- import matplotlib.pylab as plt
- import numpy as np
-
- fig, ax = plt.subplots(figsize=(10, 2))
- im = ax.imshow(spectrogram, aspect="auto", origin="lower", interpolation="none")
- plt.colorbar(im, ax=ax)
- plt.xlabel("Frames")
- plt.ylabel("Channels")
- plt.tight_layout()
-
- fig.canvas.draw()
- data = np.fromstring(fig.canvas.tostring_rgb(), dtype=np.uint8, sep="")
- data = data.reshape(fig.canvas.get_width_height()[::-1] + (3,))
- plt.close()
- return data
-
-
- def plot_alignment_to_numpy(alignment, info=None):
- global MATPLOTLIB_FLAG
- if not MATPLOTLIB_FLAG:
- import matplotlib
-
- matplotlib.use("Agg")
- MATPLOTLIB_FLAG = True
- mpl_logger = logging.getLogger("matplotlib")
- mpl_logger.setLevel(logging.WARNING)
- import matplotlib.pylab as plt
- import numpy as np
-
- fig, ax = plt.subplots(figsize=(6, 4))
- im = ax.imshow(
- alignment.transpose(), aspect="auto", origin="lower", interpolation="none"
- )
- fig.colorbar(im, ax=ax)
- xlabel = "Decoder timestep"
- if info is not None:
- xlabel += "\n\n" + info
- plt.xlabel(xlabel)
- plt.ylabel("Encoder timestep")
- plt.tight_layout()
-
- fig.canvas.draw()
- data = np.fromstring(fig.canvas.tostring_rgb(), dtype=np.uint8, sep="")
- data = data.reshape(fig.canvas.get_width_height()[::-1] + (3,))
- plt.close()
- return data
-
-
- def load_wav_to_torch(full_path):
- sampling_rate, data = read(full_path)
- return torch.FloatTensor(data.astype(np.float32)), sampling_rate
-
-
- def load_filepaths_and_text(filename, split="|"):
- with open(filename, encoding="utf-8") as f:
- filepaths_and_text = [line.strip().split(split) for line in f]
- return filepaths_and_text
-
-
- def get_hparams(init=True):
- parser = argparse.ArgumentParser()
- parser.add_argument(
- "-c",
- "--config",
- type=str,
- default="./configs/base.json",
- help="JSON file for configuration",
- )
- parser.add_argument("-m", "--model", type=str, required=True, help="Model name")
-
- args = parser.parse_args()
- model_dir = os.path.join("./logs", args.model)
-
- if not os.path.exists(model_dir):
- os.makedirs(model_dir)
-
- config_path = args.config
- config_save_path = os.path.join(model_dir, "config.json")
- if init:
- with open(config_path, "r", encoding="utf-8") as f:
- data = f.read()
- with open(config_save_path, "w", encoding="utf-8") as f:
- f.write(data)
- else:
- with open(config_save_path, "r", vencoding="utf-8") as f:
- data = f.read()
- config = json.loads(data)
- hparams = HParams(**config)
- hparams.model_dir = model_dir
- return hparams
-
-
- def clean_checkpoints(path_to_models="logs/44k/", n_ckpts_to_keep=2, sort_by_time=True):
- """Freeing up space by deleting saved ckpts
-
- Arguments:
- path_to_models -- Path to the model directory
- n_ckpts_to_keep -- Number of ckpts to keep, excluding G_0.pth and D_0.pth
- sort_by_time -- True -> chronologically delete ckpts
- False -> lexicographically delete ckpts
- """
- import re
-
- ckpts_files = [
- f
- for f in os.listdir(path_to_models)
- if os.path.isfile(os.path.join(path_to_models, f))
- ]
-
- def name_key(_f):
- return int(re.compile("._(\\d+)\\.pth").match(_f).group(1))
-
- def time_key(_f):
- return os.path.getmtime(os.path.join(path_to_models, _f))
-
- sort_key = time_key if sort_by_time else name_key
-
- def x_sorted(_x):
- return sorted(
- [f for f in ckpts_files if f.startswith(_x) and not f.endswith("_0.pth")],
- key=sort_key,
- )
-
- to_del = [
- os.path.join(path_to_models, fn)
- for fn in (
- x_sorted("G")[:-n_ckpts_to_keep]
- + x_sorted("D")[:-n_ckpts_to_keep]
- + x_sorted("WD")[:-n_ckpts_to_keep]
- )
- ]
-
- def del_info(fn):
- return logger.info(f".. Free up space by deleting ckpt {fn}")
-
- def del_routine(x):
- return [os.remove(x), del_info(x)]
-
- [del_routine(fn) for fn in to_del]
-
-
- def get_hparams_from_dir(model_dir):
- config_save_path = os.path.join(model_dir, "config.json")
- with open(config_save_path, "r", encoding="utf-8") as f:
- data = f.read()
- config = json.loads(data)
-
- hparams = HParams(**config)
- hparams.model_dir = model_dir
- return hparams
-
-
- def get_hparams_from_file(config_path):
- # print("config_path: ", config_path)
- with open(config_path, "r", encoding="utf-8") as f:
- data = f.read()
- config = json.loads(data)
-
- hparams = HParams(**config)
- return hparams
-
-
- def check_git_hash(model_dir):
- source_dir = os.path.dirname(os.path.realpath(__file__))
- if not os.path.exists(os.path.join(source_dir, ".git")):
- logger.warn(
- "{} is not a git repository, therefore hash value comparison will be ignored.".format(
- source_dir
- )
- )
- return
-
- cur_hash = subprocess.getoutput("git rev-parse HEAD")
-
- path = os.path.join(model_dir, "githash")
- if os.path.exists(path):
- saved_hash = open(path).read()
- if saved_hash != cur_hash:
- logger.warn(
- "git hash values are different. {}(saved) != {}(current)".format(
- saved_hash[:8], cur_hash[:8]
- )
- )
- else:
- open(path, "w").write(cur_hash)
-
-
- def get_logger(model_dir, filename="train.log"):
- global logger
- logger = logging.getLogger(os.path.basename(model_dir))
- logger.setLevel(logging.DEBUG)
-
- formatter = logging.Formatter("%(asctime)s\t%(name)s\t%(levelname)s\t%(message)s")
- if not os.path.exists(model_dir):
- os.makedirs(model_dir)
- h = logging.FileHandler(os.path.join(model_dir, filename))
- h.setLevel(logging.DEBUG)
- h.setFormatter(formatter)
- logger.addHandler(h)
- return logger
-
-
- class HParams:
- def __init__(self, **kwargs):
- for k, v in kwargs.items():
- if type(v) == dict:
- v = HParams(**v)
- self[k] = v
-
- def keys(self):
- return self.__dict__.keys()
-
- def items(self):
- return self.__dict__.items()
-
- def values(self):
- return self.__dict__.values()
-
- def __len__(self):
- return len(self.__dict__)
-
- def __getitem__(self, key):
- return getattr(self, key)
-
- def __setitem__(self, key, value):
- return setattr(self, key, value)
-
- def __contains__(self, key):
- return key in self.__dict__
-
- def __repr__(self):
- return self.__dict__.__repr__()
-
-
- def load_model(model_path, config_path):
- hps = get_hparams_from_file(config_path)
- net = SynthesizerTrn(
- # len(symbols),
- 108,
- hps.data.filter_length // 2 + 1,
- hps.train.segment_size // hps.data.hop_length,
- n_speakers=hps.data.n_speakers,
- **hps.model,
- ).to("cpu")
- _ = net.eval()
- _ = load_checkpoint(model_path, net, None, skip_optimizer=True)
- return net
-
-
- def mix_model(
- network1, network2, output_path, voice_ratio=(0.5, 0.5), tone_ratio=(0.5, 0.5)
- ):
- if hasattr(network1, "module"):
- state_dict1 = network1.module.state_dict()
- state_dict2 = network2.module.state_dict()
- else:
- state_dict1 = network1.state_dict()
- state_dict2 = network2.state_dict()
- for k in state_dict1.keys():
- if k not in state_dict2.keys():
- continue
- if "enc_p" in k:
- state_dict1[k] = (
- state_dict1[k].clone() * tone_ratio[0]
- + state_dict2[k].clone() * tone_ratio[1]
- )
- else:
- state_dict1[k] = (
- state_dict1[k].clone() * voice_ratio[0]
- + state_dict2[k].clone() * voice_ratio[1]
- )
- for k in state_dict2.keys():
- if k not in state_dict1.keys():
- state_dict1[k] = state_dict2[k].clone()
- torch.save(
- {"model": state_dict1, "iteration": 0, "optimizer": None, "learning_rate": 0},
- output_path,
- )
-
-
- def get_steps(model_path):
- matches = re.findall(r"\d+", model_path)
- return matches[-1] if matches else None
|