|
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: |
|
|
|
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: |
|
|
|
new_state_dict[k] = saved_state_dict[k] |
|
assert saved_state_dict[k].shape == v.shape, ( |
|
saved_state_dict[k].shape, |
|
v.shape, |
|
) |
|
except: |
|
|
|
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]) |
|
] |
|
|
|
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): |
|
|
|
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( |
|
|
|
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 |
|
|
