from features import SignalGenerator, dilated_factor
from scipy.interpolate import interp1d
import torch
import numpy as np
import json
import os
hann_window = {}
class TextAudioSpeakerCollate():
""" Zero-pads model inputs and targets
"""
def __init__(
self,
sample_rate,
hop_size,
f0_factor=1.0,
dense_factors=[0.5, 1, 4, 8],
upsample_scales=[8, 4, 2, 2],
sine_amp=0.1,
noise_amp=0.003,
signal_types=["sine"],
):
self.dense_factors = dense_factors
self.prod_upsample_scales = np.cumprod(upsample_scales)
self.sample_rate = sample_rate
self.signal_generator = SignalGenerator(
sample_rate=sample_rate,
hop_size=hop_size,
sine_amp=sine_amp,
noise_amp=noise_amp,
signal_types=signal_types,
)
self.f0_factor = f0_factor
def __call__(self, batch):
"""Collate's training batch from normalized text, audio and speaker identities
PARAMS
------
batch: [text_normalized, spec_normalized, wav_normalized, sid, note]
"""
spec_lengths = torch.LongTensor(len(batch))
sid = torch.LongTensor(len(batch))
spec_padded = torch.FloatTensor(len(batch), batch[0][0].size(0), batch[0][0].size(1))
f0_padded = torch.FloatTensor(len(batch), 1, batch[0][2].size(0))
# 返り値の初期化
spec_padded.zero_()
f0_padded.zero_()
# dfs
dfs_batch = [[] for _ in range(len(self.dense_factors))]
# row spec, sid, f0
for i in range(len(batch)):
row = batch[i]
spec = row[0]
spec_padded[i, :, :spec.size(1)] = spec
spec_lengths[i] = spec.size(1)
sid[i] = row[1]
# 推論時 f0/cf0にf0の倍率を乗算してf0/cf0を求める
f0 = row[2] * self.f0_factor
f0_padded[i, :, :f0.size(0)] = f0
# dfs
dfs = []
# dilated_factor の入力はnumpy!!
for df, us in zip(self.dense_factors, self.prod_upsample_scales):
dfs += [
np.repeat(dilated_factor(torch.unsqueeze(f0, dim=1).to('cpu').detach().numpy(), self.sample_rate, df), us)
]
# よくわからないけど、後で論文ちゃんと読む
for i in range(len(self.dense_factors)):
dfs_batch[i] += [
dfs[i].astype(np.float32).reshape(-1, 1)
] # [(T', 1), ...]
# よくわからないdfsを転置
for i in range(len(self.dense_factors)):
dfs_batch[i] = torch.FloatTensor(np.array(dfs_batch[i])).transpose(
2, 1
) # (B, 1, T')
# f0/cf0を実際に使うSignalに変換する
in_batch = self.signal_generator(f0_padded)
return spec_padded, spec_lengths, sid, in_batch, dfs_batch
def convert_continuos_f0(f0, f0_size):
# get start and end of f0
if (f0 == 0).all():
return np.zeros((f0_size,))
start_f0 = f0[f0 != 0][0]
end_f0 = f0[f0 != 0][-1]
# padding start and end of f0 sequence
cf0 = f0
start_idx = np.where(cf0 == start_f0)[0][0]
end_idx = np.where(cf0 == end_f0)[0][-1]
cf0[:start_idx] = start_f0
cf0[end_idx:] = end_f0
# get non-zero frame index
nz_frames = np.where(cf0 != 0)[0]
# perform linear interpolation
f = interp1d(nz_frames, cf0[nz_frames], bounds_error=False, fill_value=0.0)
cf0_ = f(np.arange(0, f0_size))
# print(cf0.shape, cf0_.shape, f0.shape, f0_size)
# print(cf0_)
return f(np.arange(0, f0_size))
def spectrogram_torch(y, n_fft, sampling_rate, hop_size, win_size, center=False):
if torch.min(y) < -1.:
print('min value is ', torch.min(y))
if torch.max(y) > 1.:
print('max value is ', torch.max(y))
dtype_device = str(y.dtype) + '_' + str(y.device)
wnsize_dtype_device = str(win_size) + '_' + dtype_device
if wnsize_dtype_device not in hann_window:
hann_window[wnsize_dtype_device] = torch.hann_window(win_size).to(dtype=y.dtype, device=y.device)
y = torch.nn.functional.pad(y.unsqueeze(1), (int((n_fft - hop_size) / 2), int((n_fft - hop_size) / 2)), mode='reflect')
y = y.squeeze(1)
spec = torch.stft(y, n_fft, hop_length=hop_size, win_length=win_size, window=hann_window[wnsize_dtype_device],
center=center, pad_mode='reflect', normalized=False, onesided=True, return_complex=True)
spec = torch.view_as_real(spec)
spec = torch.sqrt(spec.pow(2).sum(-1) + 1e-6)
return spec
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
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_checkpoint(checkpoint_path, model, optimizer=None):
assert os.path.isfile(checkpoint_path), f"No such file or directory: {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['pe'],
**checkpoint_dict['flow'],
**checkpoint_dict['text_enc'],
**checkpoint_dict['dec'],
**checkpoint_dict['emb_g']
}
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]
except:
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)
return model, optimizer, learning_rate, iteration