Buckets:
| """Conv-based prosody predictor (~681K params).""" | |
| import math | |
| import torch | |
| import torch.nn as nn | |
| class SinusoidalPositionalEncoding(nn.Module): | |
| def __init__(self, d_model, max_len=4096): | |
| super().__init__() | |
| pe = torch.zeros(max_len, d_model) | |
| position = torch.arange(0, max_len, dtype=torch.float32).unsqueeze(1) | |
| div_term = torch.exp(torch.arange(0, d_model, 2, dtype=torch.float32) * (-math.log(10000.0) / d_model)) | |
| pe[:, 0::2] = torch.sin(position * div_term) | |
| pe[:, 1::2] = torch.cos(position * div_term) | |
| self.register_buffer('pe', pe.unsqueeze(0)) # [1, max_len, d_model] | |
| def forward(self, x): | |
| # x: [B, L, D] | |
| return x + self.pe[:, :x.size(1)] | |
| class ConvBlock(nn.Module): | |
| def __init__(self, channels, kernel_size=5, dropout=0.2): | |
| super().__init__() | |
| self.conv = nn.Conv1d(channels, channels, kernel_size, padding=kernel_size // 2) | |
| self.norm = nn.LayerNorm(channels) | |
| self.drop = nn.Dropout(dropout) | |
| def forward(self, x): | |
| # x: [B, C, T] | |
| out = self.conv(x) | |
| out = torch.relu(out) | |
| out = out.transpose(1, 2) # [B, T, C] | |
| out = self.norm(out) | |
| out = out.transpose(1, 2) # [B, C, T] | |
| out = self.drop(out) | |
| return out | |
| class CharEncoder(nn.Module): | |
| def __init__(self, vocab_size=50, d_model=128, n_layers=4, kernel_size=5, dropout=0.2): | |
| super().__init__() | |
| self.embed = nn.Embedding(vocab_size, d_model, padding_idx=0) | |
| self.pos_enc = SinusoidalPositionalEncoding(d_model) | |
| self.blocks = nn.ModuleList([ | |
| ConvBlock(d_model, kernel_size, dropout) for _ in range(n_layers) | |
| ]) | |
| def forward(self, char_ids): | |
| # char_ids: [B, L] | |
| x = self.embed(char_ids) # [B, L, D] | |
| x = self.pos_enc(x) # [B, L, D] | |
| x = x.transpose(1, 2) # [B, D, L] | |
| for block in self.blocks: | |
| x = block(x) | |
| return x # [B, D, L] | |
| class DurationPredictor(nn.Module): | |
| """Predicts log-duration from detached encoder output (tiny-tts pattern).""" | |
| def __init__(self, d_model=128, kernel_size=3, dropout=0.2): | |
| super().__init__() | |
| self.conv1 = nn.Conv1d(d_model, d_model, kernel_size, padding=kernel_size // 2) | |
| self.norm1 = nn.LayerNorm(d_model) | |
| self.drop1 = nn.Dropout(dropout) | |
| self.conv2 = nn.Conv1d(d_model, d_model, kernel_size, padding=kernel_size // 2) | |
| self.norm2 = nn.LayerNorm(d_model) | |
| self.drop2 = nn.Dropout(dropout) | |
| self.proj = nn.Conv1d(d_model, 1, 1) | |
| def forward(self, x, mask=None): | |
| # x: [B, D, L] — detached | |
| x = x.detach() | |
| h = self.conv1(x) | |
| h = torch.relu(h) | |
| h = h.transpose(1, 2) | |
| h = self.norm1(h) | |
| h = h.transpose(1, 2) | |
| h = self.drop1(h) | |
| h = self.conv2(h) | |
| h = torch.relu(h) | |
| h = h.transpose(1, 2) | |
| h = self.norm2(h) | |
| h = h.transpose(1, 2) | |
| h = self.drop2(h) | |
| h = self.proj(h) # [B, 1, L] | |
| if mask is not None: | |
| h = h * mask.unsqueeze(1) | |
| return h.squeeze(1) # [B, L] | |
| def length_regulate(encoder_out, durations): | |
| """Repeat encoder frames according to durations. | |
| Args: | |
| encoder_out: [B, D, L] | |
| durations: [B, L] (integer) | |
| Returns: | |
| regulated: [B, D, T_max] zero-padded | |
| frame_lengths: [B] | |
| """ | |
| B, D, L = encoder_out.shape | |
| durations = durations.long() | |
| frame_lengths = durations.sum(dim=1) # [B] | |
| T_max = frame_lengths.max().item() | |
| regulated = torch.zeros(B, D, T_max, device=encoder_out.device, dtype=encoder_out.dtype) | |
| for b in range(B): | |
| idx = 0 | |
| for l in range(L): | |
| dur = durations[b, l].item() | |
| if dur > 0 and idx < T_max: | |
| end = min(idx + dur, T_max) | |
| regulated[b, :, idx:end] = encoder_out[b, :, l:l+1].expand(-1, end - idx) | |
| idx = end | |
| return regulated, frame_lengths | |
| class FrameDecoder(nn.Module): | |
| def __init__(self, d_model=128, n_layers=3, kernel_size=5, dropout=0.2): | |
| super().__init__() | |
| self.blocks = nn.ModuleList([ | |
| ConvBlock(d_model, kernel_size, dropout) for _ in range(n_layers) | |
| ]) | |
| self.proj = nn.Conv1d(d_model, 2, 1) # [f0, rms] | |
| def forward(self, x): | |
| # x: [B, D, T] | |
| for block in self.blocks: | |
| x = block(x) | |
| return self.proj(x) # [B, 2, T] | |
| class ProsodyPredictor(nn.Module): | |
| def __init__(self, vocab_size=50, d_model=128, dropout=0.2): | |
| super().__init__() | |
| self.encoder = CharEncoder(vocab_size, d_model, n_layers=4, kernel_size=5, dropout=dropout) | |
| self.dur_predictor = DurationPredictor(d_model, kernel_size=3, dropout=dropout) | |
| self.decoder = FrameDecoder(d_model, n_layers=3, kernel_size=5, dropout=dropout) | |
| def forward(self, char_ids, durations=None, char_lengths=None): | |
| """ | |
| Args: | |
| char_ids: [B, L] | |
| durations: [B, L] ground-truth (training) or None (inference) | |
| char_lengths: [B] for masking | |
| Returns: | |
| pred_f0: [B, T] | |
| pred_rms: [B, T] | |
| pred_log_dur: [B, L] | |
| frame_lengths: [B] | |
| """ | |
| # Encoder | |
| enc_out = self.encoder(char_ids) # [B, D, L] | |
| # Duration prediction (detached input) | |
| char_mask = None | |
| if char_lengths is not None: | |
| char_mask = torch.arange(char_ids.size(1), device=char_ids.device).unsqueeze(0) < char_lengths.unsqueeze(1) | |
| char_mask = char_mask.float() | |
| pred_log_dur = self.dur_predictor(enc_out, char_mask) # [B, L] | |
| # Length regulate | |
| if durations is not None: | |
| # Training: use GT durations | |
| regulated, frame_lengths = length_regulate(enc_out, durations) | |
| else: | |
| # Inference: use predicted durations | |
| pred_dur = torch.round(torch.exp(pred_log_dur)).long().clamp(min=1) | |
| if char_mask is not None: | |
| pred_dur = pred_dur * char_mask.long() | |
| regulated, frame_lengths = length_regulate(enc_out, pred_dur) | |
| # Decode | |
| out = self.decoder(regulated) # [B, 2, T] | |
| pred_f0 = out[:, 0] # [B, T] | |
| pred_rms = out[:, 1] # [B, T] | |
| return pred_f0, pred_rms, pred_log_dur, frame_lengths | |
Xet Storage Details
- Size:
- 6.47 kB
- Xet hash:
- 46b8c9facbc19f59e191ed786d540488b66102b1902a6ff5b3e8d56c8146ac7d
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.