Spaces:
Sleeping
Sleeping
File size: 3,425 Bytes
b725c5a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 |
import torch
import torch.nn as nn
import numpy as np
import torch.nn.functional as F
from modules.naturalpseech2.transformers import (
TransformerEncoder,
DurationPredictor,
PitchPredictor,
LengthRegulator,
)
class PriorEncoder(nn.Module):
def __init__(self, cfg):
super().__init__()
self.cfg = cfg
self.enc_emb_tokens = nn.Embedding(
cfg.vocab_size, cfg.encoder.encoder_hidden, padding_idx=0
)
self.enc_emb_tokens.weight.data.normal_(mean=0.0, std=1e-5)
self.encoder = TransformerEncoder(
enc_emb_tokens=self.enc_emb_tokens, cfg=cfg.encoder
)
self.duration_predictor = DurationPredictor(cfg.duration_predictor)
self.pitch_predictor = PitchPredictor(cfg.pitch_predictor)
self.length_regulator = LengthRegulator()
self.pitch_min = cfg.pitch_min
self.pitch_max = cfg.pitch_max
self.pitch_bins_num = cfg.pitch_bins_num
pitch_bins = torch.exp(
torch.linspace(
np.log(self.pitch_min), np.log(self.pitch_max), self.pitch_bins_num - 1
)
)
self.register_buffer("pitch_bins", pitch_bins)
self.pitch_embedding = nn.Embedding(
self.pitch_bins_num, cfg.encoder.encoder_hidden
)
def forward(
self,
phone_id,
duration=None,
pitch=None,
phone_mask=None,
mask=None,
ref_emb=None,
ref_mask=None,
is_inference=False,
):
"""
input:
phone_id: (B, N)
duration: (B, N)
pitch: (B, T)
phone_mask: (B, N); mask is 0
mask: (B, T); mask is 0
ref_emb: (B, d, T')
ref_mask: (B, T'); mask is 0
output:
prior_embedding: (B, d, T)
pred_dur: (B, N)
pred_pitch: (B, T)
"""
x = self.encoder(phone_id, phone_mask, ref_emb.transpose(1, 2))
# print(torch.min(x), torch.max(x))
dur_pred_out = self.duration_predictor(x, phone_mask, ref_emb, ref_mask)
# dur_pred_out: {dur_pred_log, dur_pred, dur_pred_round}
if is_inference or duration is None:
x, mel_len = self.length_regulator(
x,
dur_pred_out["dur_pred_round"],
max_len=torch.max(torch.sum(dur_pred_out["dur_pred_round"], dim=1)),
)
else:
x, mel_len = self.length_regulator(x, duration, max_len=pitch.shape[1])
pitch_pred_log = self.pitch_predictor(x, mask, ref_emb, ref_mask)
if is_inference or pitch is None:
pitch_tokens = torch.bucketize(pitch_pred_log.exp(), self.pitch_bins)
pitch_embedding = self.pitch_embedding(pitch_tokens)
else:
pitch_tokens = torch.bucketize(pitch, self.pitch_bins)
pitch_embedding = self.pitch_embedding(pitch_tokens)
x = x + pitch_embedding
if (not is_inference) and (mask is not None):
x = x * mask.to(x.dtype)[:, :, None]
prior_out = {
"dur_pred_round": dur_pred_out["dur_pred_round"],
"dur_pred_log": dur_pred_out["dur_pred_log"],
"dur_pred": dur_pred_out["dur_pred"],
"pitch_pred_log": pitch_pred_log,
"pitch_token": pitch_tokens,
"mel_len": mel_len,
"prior_out": x,
}
return prior_out
|