Emberstar
434 lines
14 KiB
Python
434 lines
14 KiB
Python
import os
|
||
import glob
|
||
import sys
|
||
import argparse
|
||
import logging
|
||
import json
|
||
import subprocess
|
||
import numpy as np
|
||
from scipy.io.wavfile import read
|
||
import torch
|
||
import regex as re
|
||
|
||
MATPLOTLIB_FLAG = False
|
||
|
||
logging.basicConfig(stream=sys.stdout, level=logging.DEBUG)
|
||
logger = logging
|
||
|
||
|
||
|
||
zh_pattern = re.compile(r'[\u4e00-\u9fa5]')
|
||
en_pattern = re.compile(r'[a-zA-Z]')
|
||
jp_pattern = re.compile(r'[\u3040-\u30ff\u31f0-\u31ff]')
|
||
kr_pattern = re.compile(r'[\uac00-\ud7af\u1100-\u11ff\u3130-\u318f\ua960-\ua97f]')
|
||
num_pattern=re.compile(r'[0-9]')
|
||
comma=r"(?<=[.。!！?？；;，,、:：'\"‘“”’()（）《》「」~——])" #向前匹配但固定长度
|
||
tags={'ZH':'[ZH]','EN':'[EN]','JP':'[JA]','KR':'[KR]'}
|
||
|
||
def tag_cjke(text):
|
||
'''为中英日韩加tag,中日正则分不开，故先分句分离中日再识别，以应对大部分情况'''
|
||
sentences = re.split(r"([.。!！?？；;，,、:：'\"‘“”’()（）【】《》「」~——]+ *(?![0-9]))", text) #分句，排除小数点
|
||
sentences.append("")
|
||
sentences = ["".join(i) for i in zip(sentences[0::2],sentences[1::2])]
|
||
# print(sentences)
|
||
prev_lang=None
|
||
tagged_text = ""
|
||
for s in sentences:
|
||
#全为符号跳过
|
||
nu = re.sub(r'[\s\p{P}]+', '', s, flags=re.U).strip()
|
||
if len(nu)==0:
|
||
continue
|
||
s = re.sub(r'[()（）《》「」【】‘“”’]+', '', s)
|
||
jp=re.findall(jp_pattern, s)
|
||
#本句含日语字符判断为日语
|
||
if len(jp)>0:
|
||
prev_lang,tagged_jke=tag_jke(s,prev_lang)
|
||
tagged_text +=tagged_jke
|
||
else:
|
||
prev_lang,tagged_cke=tag_cke(s,prev_lang)
|
||
tagged_text +=tagged_cke
|
||
return tagged_text
|
||
|
||
def tag_jke(text,prev_sentence=None):
|
||
'''为英日韩加tag'''
|
||
# 初始化标记变量
|
||
tagged_text = ""
|
||
prev_lang = None
|
||
tagged=0
|
||
# 遍历文本
|
||
for char in text:
|
||
# 判断当前字符属于哪种语言
|
||
if jp_pattern.match(char):
|
||
lang = "JP"
|
||
elif zh_pattern.match(char):
|
||
lang = "JP"
|
||
elif kr_pattern.match(char):
|
||
lang = "KR"
|
||
elif en_pattern.match(char):
|
||
lang = "EN"
|
||
# elif num_pattern.match(char):
|
||
# lang = prev_sentence
|
||
else:
|
||
lang = None
|
||
tagged_text += char
|
||
continue
|
||
# 如果当前语言与上一个语言不同，就添加标记
|
||
if lang != prev_lang:
|
||
tagged=1
|
||
if prev_lang==None: # 开头
|
||
tagged_text =tags[lang]+tagged_text
|
||
else:
|
||
tagged_text =tagged_text+tags[prev_lang]+tags[lang]
|
||
|
||
# 重置标记变量
|
||
prev_lang = lang
|
||
|
||
# 添加当前字符到标记文本中
|
||
tagged_text += char
|
||
|
||
# 在最后一个语言的结尾添加对应的标记
|
||
if prev_lang:
|
||
tagged_text += tags[prev_lang]
|
||
if not tagged:
|
||
prev_lang=prev_sentence
|
||
tagged_text =tags[prev_lang]+tagged_text+tags[prev_lang]
|
||
|
||
return prev_lang,tagged_text
|
||
|
||
def tag_cke(text,prev_sentence=None):
|
||
'''为中英韩加tag'''
|
||
# 初始化标记变量
|
||
tagged_text = ""
|
||
prev_lang = None
|
||
# 是否全略过未标签
|
||
tagged=0
|
||
|
||
# 遍历文本
|
||
for char in text:
|
||
# 判断当前字符属于哪种语言
|
||
if zh_pattern.match(char):
|
||
lang = "ZH"
|
||
elif kr_pattern.match(char):
|
||
lang = "KR"
|
||
elif en_pattern.match(char):
|
||
lang = "EN"
|
||
# elif num_pattern.match(char):
|
||
# lang = prev_sentence
|
||
else:
|
||
# 略过
|
||
lang = None
|
||
tagged_text += char
|
||
continue
|
||
|
||
# 如果当前语言与上一个语言不同，添加标记
|
||
if lang != prev_lang:
|
||
tagged=1
|
||
if prev_lang==None: # 开头
|
||
tagged_text =tags[lang]+tagged_text
|
||
else:
|
||
tagged_text =tagged_text+tags[prev_lang]+tags[lang]
|
||
|
||
# 重置标记变量
|
||
prev_lang = lang
|
||
|
||
# 添加当前字符到标记文本中
|
||
tagged_text += char
|
||
|
||
# 在最后一个语言的结尾添加对应的标记
|
||
if prev_lang:
|
||
tagged_text += tags[prev_lang]
|
||
# 未标签则继承上一句标签
|
||
if tagged==0:
|
||
prev_lang=prev_sentence
|
||
tagged_text =tags[prev_lang]+tagged_text+tags[prev_lang]
|
||
return prev_lang,tagged_text
|
||
|
||
|
||
|
||
def load_checkpoint(checkpoint_path, model, optimizer=None, drop_speaker_emb=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:
|
||
optimizer.load_state_dict(checkpoint_dict['optimizer'])
|
||
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:
|
||
if k == 'emb_g.weight':
|
||
if drop_speaker_emb:
|
||
new_state_dict[k] = v
|
||
continue
|
||
v[:saved_state_dict[k].shape[0], :] = saved_state_dict[k]
|
||
new_state_dict[k] = v
|
||
else:
|
||
new_state_dict[k] = saved_state_dict[k]
|
||
except:
|
||
logger.info("%s is not in the checkpoint" % k)
|
||
new_state_dict[k] = v
|
||
if hasattr(model, 'module'):
|
||
model.module.load_state_dict(new_state_dict)
|
||
else:
|
||
model.load_state_dict(new_state_dict)
|
||
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() if optimizer is not None else None,
|
||
'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 extract_digits(f):
|
||
digits = "".join(filter(str.isdigit, f))
|
||
return int(digits) if digits else -1
|
||
|
||
|
||
def latest_checkpoint_path(dir_path, regex="G_[0-9]*.pth"):
|
||
f_list = glob.glob(os.path.join(dir_path, regex))
|
||
f_list.sort(key=lambda f: extract_digits(f))
|
||
x = f_list[-1]
|
||
print(f"latest_checkpoint_path:{x}")
|
||
return x
|
||
|
||
|
||
def oldest_checkpoint_path(dir_path, regex="G_[0-9]*.pth", preserved=4):
|
||
f_list = glob.glob(os.path.join(dir_path, regex))
|
||
f_list.sort(key=lambda f: extract_digits(f))
|
||
if len(f_list) > preserved:
|
||
x = f_list[0]
|
||
print(f"oldest_checkpoint_path:{x}")
|
||
return x
|
||
return ""
|
||
|
||
|
||
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 str2bool(v):
|
||
if isinstance(v, bool):
|
||
return v
|
||
if v.lower() in ('yes', 'true', 't', 'y', '1'):
|
||
return True
|
||
elif v.lower() in ('no', 'false', 'f', 'n', '0'):
|
||
return False
|
||
else:
|
||
raise argparse.ArgumentTypeError('Boolean value expected.')
|
||
|
||
|
||
def get_hparams(init=True):
|
||
parser = argparse.ArgumentParser()
|
||
parser.add_argument('-c', '--config', type=str, default="./configs/modified_finetune_speaker.json",
|
||
help='JSON file for configuration')
|
||
parser.add_argument('-m', '--model', type=str, default="pretrained_models",
|
||
help='Model name')
|
||
parser.add_argument('-n', '--max_epochs', type=int, default=50,
|
||
help='finetune epochs')
|
||
parser.add_argument('--cont', type=str2bool, default=False, help='whether to continue training on the latest checkpoint')
|
||
parser.add_argument('--drop_speaker_embed', type=str2bool, default=False, help='whether to drop existing characters')
|
||
parser.add_argument('--train_with_pretrained_model', type=str2bool, default=True,
|
||
help='whether to train with pretrained model')
|
||
parser.add_argument('--preserved', type=int, default=4,
|
||
help='Number of preserved models')
|
||
|
||
args = parser.parse_args()
|
||
model_dir = os.path.join("./", 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") as f:
|
||
data = f.read()
|
||
with open(config_save_path, "w") as f:
|
||
f.write(data)
|
||
else:
|
||
with open(config_save_path, "r") as f:
|
||
data = f.read()
|
||
config = json.loads(data)
|
||
|
||
hparams = HParams(**config)
|
||
hparams.model_dir = model_dir
|
||
hparams.max_epochs = args.max_epochs
|
||
hparams.cont = args.cont
|
||
hparams.drop_speaker_embed = args.drop_speaker_embed
|
||
hparams.train_with_pretrained_model = args.train_with_pretrained_model
|
||
hparams.preserved = args.preserved
|
||
return hparams
|
||
|
||
|
||
def get_hparams_from_dir(model_dir):
|
||
config_save_path = os.path.join(model_dir, "config.json")
|
||
with open(config_save_path, "r") 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__() |