mirror of
https://github.com/storytold/storyteller-ml.git
synced 2026-10-11 21:41:18 +00:00
Justin John
257 lines
8.3 KiB
Python
257 lines
8.3 KiB
Python
import torch
|
|
import numpy as np
|
|
import torch.nn.functional as F
|
|
|
|
|
|
class SLMAdversarialLoss(torch.nn.Module):
|
|
def __init__(
|
|
self,
|
|
model,
|
|
wl,
|
|
sampler,
|
|
min_len,
|
|
max_len,
|
|
batch_percentage=0.5,
|
|
skip_update=10,
|
|
sig=1.5,
|
|
):
|
|
super(SLMAdversarialLoss, self).__init__()
|
|
self.model = model
|
|
self.wl = wl
|
|
self.sampler = sampler
|
|
|
|
self.min_len = min_len
|
|
self.max_len = max_len
|
|
self.batch_percentage = batch_percentage
|
|
|
|
self.sig = sig
|
|
self.skip_update = skip_update
|
|
|
|
def forward(
|
|
self,
|
|
iters,
|
|
y_rec_gt,
|
|
y_rec_gt_pred,
|
|
waves,
|
|
mel_input_length,
|
|
ref_text,
|
|
ref_lengths,
|
|
use_ind,
|
|
s_trg,
|
|
ref_s=None,
|
|
):
|
|
text_mask = length_to_mask(ref_lengths).to(ref_text.device)
|
|
bert_dur = self.model.bert(ref_text, attention_mask=(~text_mask).int())
|
|
d_en = self.model.bert_encoder(bert_dur).transpose(-1, -2)
|
|
|
|
if use_ind and np.random.rand() < 0.5:
|
|
s_preds = s_trg
|
|
else:
|
|
num_steps = np.random.randint(3, 5)
|
|
if ref_s is not None:
|
|
s_preds = self.sampler(
|
|
noise=torch.randn_like(s_trg).unsqueeze(1).to(ref_text.device),
|
|
embedding=bert_dur,
|
|
embedding_scale=1,
|
|
features=ref_s, # reference from the same speaker as the embedding
|
|
embedding_mask_proba=0.1,
|
|
num_steps=num_steps,
|
|
).squeeze(1)
|
|
else:
|
|
s_preds = self.sampler(
|
|
noise=torch.randn_like(s_trg).unsqueeze(1).to(ref_text.device),
|
|
embedding=bert_dur,
|
|
embedding_scale=1,
|
|
embedding_mask_proba=0.1,
|
|
num_steps=num_steps,
|
|
).squeeze(1)
|
|
|
|
s_dur = s_preds[:, 128:]
|
|
s = s_preds[:, :128]
|
|
|
|
d, _ = self.model.predictor(
|
|
d_en,
|
|
s_dur,
|
|
ref_lengths,
|
|
torch.randn(ref_lengths.shape[0], ref_lengths.max(), 2).to(ref_text.device),
|
|
text_mask,
|
|
)
|
|
|
|
bib = 0
|
|
|
|
output_lengths = []
|
|
attn_preds = []
|
|
|
|
# differentiable duration modeling
|
|
for _s2s_pred, _text_length in zip(d, ref_lengths):
|
|
_s2s_pred_org = _s2s_pred[:_text_length, :]
|
|
|
|
_s2s_pred = torch.sigmoid(_s2s_pred_org)
|
|
_dur_pred = _s2s_pred.sum(axis=-1)
|
|
|
|
l = int(torch.round(_s2s_pred.sum()).item())
|
|
t = torch.arange(0, l).expand(l)
|
|
|
|
t = (
|
|
torch.arange(0, l)
|
|
.unsqueeze(0)
|
|
.expand((len(_s2s_pred), l))
|
|
.to(ref_text.device)
|
|
)
|
|
loc = torch.cumsum(_dur_pred, dim=0) - _dur_pred / 2
|
|
|
|
h = torch.exp(
|
|
-0.5 * torch.square(t - (l - loc.unsqueeze(-1))) / (self.sig) ** 2
|
|
)
|
|
|
|
out = torch.nn.functional.conv1d(
|
|
_s2s_pred_org.unsqueeze(0),
|
|
h.unsqueeze(1),
|
|
padding=h.shape[-1] - 1,
|
|
groups=int(_text_length),
|
|
)[..., :l]
|
|
attn_preds.append(F.softmax(out.squeeze(), dim=0))
|
|
|
|
output_lengths.append(l)
|
|
|
|
max_len = max(output_lengths)
|
|
|
|
with torch.no_grad():
|
|
t_en = self.model.text_encoder(ref_text, ref_lengths, text_mask)
|
|
|
|
s2s_attn = torch.zeros(len(ref_lengths), int(ref_lengths.max()), max_len).to(
|
|
ref_text.device
|
|
)
|
|
for bib in range(len(output_lengths)):
|
|
s2s_attn[bib, : ref_lengths[bib], : output_lengths[bib]] = attn_preds[bib]
|
|
|
|
asr_pred = t_en @ s2s_attn
|
|
|
|
_, p_pred = self.model.predictor(d_en, s_dur, ref_lengths, s2s_attn, text_mask)
|
|
|
|
mel_len = max(int(min(output_lengths) / 2 - 1), self.min_len // 2)
|
|
mel_len = min(mel_len, self.max_len // 2)
|
|
|
|
# get clips
|
|
|
|
en = []
|
|
p_en = []
|
|
sp = []
|
|
|
|
F0_fakes = []
|
|
N_fakes = []
|
|
|
|
wav = []
|
|
|
|
for bib in range(len(output_lengths)):
|
|
mel_length_pred = output_lengths[bib]
|
|
mel_length_gt = int(mel_input_length[bib].item() / 2)
|
|
if mel_length_gt <= mel_len or mel_length_pred <= mel_len:
|
|
continue
|
|
|
|
sp.append(s_preds[bib])
|
|
|
|
random_start = np.random.randint(0, mel_length_pred - mel_len)
|
|
en.append(asr_pred[bib, :, random_start : random_start + mel_len])
|
|
p_en.append(p_pred[bib, :, random_start : random_start + mel_len])
|
|
|
|
# get ground truth clips
|
|
random_start = np.random.randint(0, mel_length_gt - mel_len)
|
|
y = waves[bib][
|
|
(random_start * 2) * 300 : ((random_start + mel_len) * 2) * 300
|
|
]
|
|
wav.append(torch.from_numpy(y).to(ref_text.device))
|
|
|
|
if len(wav) >= self.batch_percentage * len(
|
|
waves
|
|
): # prevent OOM due to longer lengths
|
|
break
|
|
|
|
if len(sp) <= 1:
|
|
return None
|
|
|
|
sp = torch.stack(sp)
|
|
wav = torch.stack(wav).float()
|
|
en = torch.stack(en)
|
|
p_en = torch.stack(p_en)
|
|
|
|
F0_fake, N_fake = self.model.predictor.F0Ntrain(p_en, sp[:, 128:])
|
|
y_pred = self.model.decoder(en, F0_fake, N_fake, sp[:, :128])
|
|
|
|
# discriminator loss
|
|
if (iters + 1) % self.skip_update == 0:
|
|
if np.random.randint(0, 2) == 0:
|
|
wav = y_rec_gt_pred
|
|
use_rec = True
|
|
else:
|
|
use_rec = False
|
|
|
|
crop_size = min(wav.size(-1), y_pred.size(-1))
|
|
if (
|
|
use_rec
|
|
): # use reconstructed (shorter lengths), do length invariant regularization
|
|
if wav.size(-1) > y_pred.size(-1):
|
|
real_GP = wav[:, :, :crop_size]
|
|
out_crop = self.wl.discriminator_forward(real_GP.detach().squeeze())
|
|
out_org = self.wl.discriminator_forward(wav.detach().squeeze())
|
|
loss_reg = F.l1_loss(out_crop, out_org[..., : out_crop.size(-1)])
|
|
|
|
if np.random.randint(0, 2) == 0:
|
|
d_loss = self.wl.discriminator(
|
|
real_GP.detach().squeeze(), y_pred.detach().squeeze()
|
|
).mean()
|
|
else:
|
|
d_loss = self.wl.discriminator(
|
|
wav.detach().squeeze(), y_pred.detach().squeeze()
|
|
).mean()
|
|
else:
|
|
real_GP = y_pred[:, :, :crop_size]
|
|
out_crop = self.wl.discriminator_forward(real_GP.detach().squeeze())
|
|
out_org = self.wl.discriminator_forward(y_pred.detach().squeeze())
|
|
loss_reg = F.l1_loss(out_crop, out_org[..., : out_crop.size(-1)])
|
|
|
|
if np.random.randint(0, 2) == 0:
|
|
d_loss = self.wl.discriminator(
|
|
wav.detach().squeeze(), real_GP.detach().squeeze()
|
|
).mean()
|
|
else:
|
|
d_loss = self.wl.discriminator(
|
|
wav.detach().squeeze(), y_pred.detach().squeeze()
|
|
).mean()
|
|
|
|
# regularization (ignore length variation)
|
|
d_loss += loss_reg
|
|
|
|
out_gt = self.wl.discriminator_forward(y_rec_gt.detach().squeeze())
|
|
out_rec = self.wl.discriminator_forward(
|
|
y_rec_gt_pred.detach().squeeze()
|
|
)
|
|
|
|
# regularization (ignore reconstruction artifacts)
|
|
d_loss += F.l1_loss(out_gt, out_rec)
|
|
|
|
else:
|
|
d_loss = self.wl.discriminator(
|
|
wav.detach().squeeze(), y_pred.detach().squeeze()
|
|
).mean()
|
|
else:
|
|
d_loss = 0
|
|
|
|
# generator loss
|
|
gen_loss = self.wl.generator(y_pred.squeeze())
|
|
|
|
gen_loss = gen_loss.mean()
|
|
|
|
return d_loss, gen_loss, y_pred.detach().cpu().numpy()
|
|
|
|
|
|
def length_to_mask(lengths):
|
|
mask = (
|
|
torch.arange(lengths.max())
|
|
.unsqueeze(0)
|
|
.expand(lengths.shape[0], -1)
|
|
.type_as(lengths)
|
|
)
|
|
mask = torch.gt(mask + 1, lengths.unsqueeze(1))
|
|
return mask
|