FrankZxShen's picture
init
2b3664e
raw
history blame
16.5 kB
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
import loralib as lora
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_lora_checkpoint(checkpoint_path, model, optimizer=None, generator_path = "./pretrained_models/G_latest.pth"):
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']
generator_state_dict = torch.load(generator_path)['model']
lora_state_dict = checkpoint_dict['model']
new_state_dict = {}
for k, v in lora_state_dict.items():
try:
if k == 'emb_g.weight':
if drop_speaker_emb:
new_state_dict[k] = v
continue
v[:lora_state_dict[k].shape[0], :] = lora_state_dict[k]
new_state_dict[k] = v
else:
new_state_dict[k] = lora_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(generator_state_dict, strict=False)
model.module.load_state_dict(new_state_dict, strict=False)
# lora.mark_only_lora_as_trainable(model.module)
else:
model.load_state_dict(generator_state_dict, strict=False)
model.load_state_dict(new_state_dict, strict=False)
# lora.mark_only_lora_as_trainable(model)
logger.info("Loaded checkpoint '{}' (iteration {})" .format(
checkpoint_path, iteration))
return model, optimizer, learning_rate, iteration
def save_lora_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 = lora.lora_state_dict(model.module)
else:
state_dict = lora.lora_state_dict(model)
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 load_lora_checkpoint_fix(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, strict=False)
lora.mark_only_lora_as_trainable(model.module)
else:
model.load_state_dict(new_state_dict, strict=False)
lora.mark_only_lora_as_trainable(model)
logger.info("Loaded checkpoint '{}' (iteration {})".format(
checkpoint_path, iteration))
return model, optimizer, learning_rate, iteration
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 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]
print(x)
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/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('--drop_speaker_embed', type=bool, default=False, help='whether to drop existing characters')
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.drop_speaker_embed = args.drop_speaker_embed
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__()