| """ |
| Human Prosody & Pitch (F0) Engine for ViuAI_TTS_200M. |
| Predicts frame durations, fundamental frequency (F0 pitch intonation), |
| and volume/energy dynamics to ensure 100% natural, expressive human prosody. |
| """ |
|
|
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| from typing import Tuple, Optional |
|
|
|
|
| class ConvNeXtBlock(nn.Module): |
| def __init__(self, dim: int, kernel_size: int = 7): |
| super().__init__() |
| self.dwconv = nn.Conv1d(dim, dim, kernel_size=kernel_size, padding=kernel_size // 2, groups=dim) |
| self.norm = nn.LayerNorm(dim, eps=1e-6) |
| self.pwconv1 = nn.Linear(dim, 4 * dim) |
| self.act = nn.GELU() |
| self.pwconv2 = nn.Linear(4 * dim, dim) |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| |
| residual = x |
| x = self.dwconv(x.transpose(1, 2)).transpose(1, 2) |
| x = self.norm(x) |
| x = self.pwconv1(x) |
| x = self.act(x) |
| x = self.pwconv2(x) |
| return residual + x |
|
|
|
|
| class HumanProsodyEngine(nn.Module): |
| """ |
| State-of-the-Art Prosody Modeling Engine. |
| Predicts: |
| 1. Duration (speaking pace and micro-pauses) |
| 2. Pitch (F0) contour (emotional inflection, question vs statement tone) |
| 3. Energy contour (breathiness, volume variation) |
| """ |
| def __init__( |
| self, |
| hidden_dim: int = 512, |
| num_layers: int = 4, |
| kernel_size: int = 7, |
| dropout: float = 0.1, |
| ): |
| super().__init__() |
| self.shared_blocks = nn.ModuleList([ |
| ConvNeXtBlock(hidden_dim, kernel_size=kernel_size) |
| for _ in range(num_layers) |
| ]) |
| self.dropout = nn.Dropout(dropout) |
|
|
| |
| self.duration_head = nn.Linear(hidden_dim, 1) |
| self.pitch_head = nn.Linear(hidden_dim, 1) |
| self.energy_head = nn.Linear(hidden_dim, 1) |
|
|
| |
| self.pitch_proj = nn.Linear(1, hidden_dim) |
| self.energy_proj = nn.Linear(1, hidden_dim) |
|
|
| def forward( |
| self, |
| text_latents: torch.Tensor, |
| ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: |
| """ |
| Args: |
| text_latents: [B, T_text, hidden_dim] |
| Returns: |
| log_duration: [B, T_text] |
| pitch_pred: [B, T_text] |
| energy_pred: [B, T_text] |
| """ |
| x = text_latents |
| for block in self.shared_blocks: |
| x = block(x) |
| x = self.dropout(x) |
|
|
| log_duration = self.duration_head(x).squeeze(-1) |
| pitch_pred = self.pitch_head(x).squeeze(-1) |
| energy_pred = self.energy_head(x).squeeze(-1) |
|
|
| return log_duration, pitch_pred, energy_pred |
|
|
| def get_prosody_features(self, pitch: torch.Tensor, energy: torch.Tensor) -> torch.Tensor: |
| """Projects pitch and energy into hidden_dim vectors.""" |
| p = self.pitch_proj(pitch.unsqueeze(-1)) |
| e = self.energy_proj(energy.unsqueeze(-1)) |
| return p + e |
|
|
|
|
| |
| DurationPredictor = HumanProsodyEngine |
|
|
