mirror of
https://github.com/w-okada/voice-changer.git
synced 2025-01-24 05:55:01 +03:00
152 lines
4.9 KiB
Python
152 lines
4.9 KiB
Python
import torch
|
|
import os, sys, json
|
|
import logging
|
|
logging.basicConfig(stream=sys.stdout, level=logging.DEBUG)
|
|
logger = logging
|
|
|
|
hann_window = {}
|
|
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))
|
|
|
|
global hann_window
|
|
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
|
|
|
|
class TextAudioSpeakerCollate():
|
|
""" Zero-pads model inputs and targets
|
|
"""
|
|
def __init__(self, return_ids=False, no_text = False):
|
|
self.return_ids = return_ids
|
|
self.no_text = no_text
|
|
|
|
def __call__(self, batch):
|
|
"""Collate's training batch from normalized text, audio and speaker identities
|
|
PARAMS
|
|
------
|
|
batch: [text_normalized, spec_normalized, wav_normalized, sid]
|
|
"""
|
|
# Right zero-pad all one-hot text sequences to max input length
|
|
_, ids_sorted_decreasing = torch.sort(
|
|
torch.LongTensor([x[1].size(1) for x in batch]),
|
|
dim=0, descending=True)
|
|
|
|
max_text_len = max([len(x[0]) for x in batch])
|
|
max_spec_len = max([x[1].size(1) for x in batch])
|
|
max_wav_len = max([x[2].size(1) for x in batch])
|
|
|
|
text_lengths = torch.LongTensor(len(batch))
|
|
spec_lengths = torch.LongTensor(len(batch))
|
|
wav_lengths = torch.LongTensor(len(batch))
|
|
sid = torch.LongTensor(len(batch))
|
|
|
|
text_padded = torch.LongTensor(len(batch), max_text_len)
|
|
spec_padded = torch.FloatTensor(len(batch), batch[0][1].size(0), max_spec_len)
|
|
wav_padded = torch.FloatTensor(len(batch), 1, max_wav_len)
|
|
text_padded.zero_()
|
|
spec_padded.zero_()
|
|
wav_padded.zero_()
|
|
for i in range(len(ids_sorted_decreasing)):
|
|
row = batch[ids_sorted_decreasing[i]]
|
|
|
|
text = row[0]
|
|
text_padded[i, :text.size(0)] = text
|
|
text_lengths[i] = text.size(0)
|
|
|
|
spec = row[1]
|
|
spec_padded[i, :, :spec.size(1)] = spec
|
|
spec_lengths[i] = spec.size(1)
|
|
|
|
wav = row[2]
|
|
wav_padded[i, :, :wav.size(1)] = wav
|
|
wav_lengths[i] = wav.size(1)
|
|
|
|
sid[i] = row[3]
|
|
|
|
if self.return_ids:
|
|
return text_padded, text_lengths, spec_padded, spec_lengths, wav_padded, wav_lengths, sid, ids_sorted_decreasing
|
|
return text_padded, text_lengths, spec_padded, spec_lengths, wav_padded, wav_lengths, sid
|
|
|
|
|
|
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['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]
|
|
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 get_hparams_from_file(config_path):
|
|
with open(config_path, "r") 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__()
|
|
|