File size: 3,109 Bytes
4fe4e70 | 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 | """
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:
# x: [B, T, C]
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)
# Output projection heads
self.duration_head = nn.Linear(hidden_dim, 1)
self.pitch_head = nn.Linear(hidden_dim, 1)
self.energy_head = nn.Linear(hidden_dim, 1)
# Embeddings to project pitch and energy back into acoustic features
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
# Maintain backward compatibility alias
DurationPredictor = HumanProsodyEngine
|