Initial commit
This commit is contained in:
redclawsystems
@@ -0,0 +1,647 @@
|
||||
//! Variance Adaptor for FastSpeech2
|
||||
//!
|
||||
//! The variance adaptor predicts duration, pitch, and energy to control
|
||||
//! prosody in non-autoregressive TTS models.
|
||||
|
||||
use crate::Result;
|
||||
use rtx_nn::layers::{Module, linear::Linear};
|
||||
use rtx_nn::layers::activation::ReLU;
|
||||
use rtx_nn::layers::conv::{Conv1d, Conv1dConfig};
|
||||
use rtx_tensor::{Tensor, Device};
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
/// Configuration for variance adaptor
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct VarianceAdaptorConfig {
|
||||
/// Hidden dimension
|
||||
pub hidden_dim: usize,
|
||||
|
||||
/// Kernel size for convolutions
|
||||
pub kernel_size: usize,
|
||||
|
||||
/// Dropout probability
|
||||
pub dropout: f32,
|
||||
|
||||
/// Number of convolutional layers in each predictor
|
||||
pub n_layers: usize,
|
||||
}
|
||||
|
||||
impl Default for VarianceAdaptorConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
hidden_dim: 256,
|
||||
kernel_size: 3,
|
||||
dropout: 0.1,
|
||||
n_layers: 2,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Simplified Softplus activation (using exponential approximation)
|
||||
#[derive(Debug)]
|
||||
struct SoftplusApprox {
|
||||
beta: f32,
|
||||
}
|
||||
|
||||
impl SoftplusApprox {
|
||||
fn new(beta: f32) -> Self {
|
||||
Self { beta }
|
||||
}
|
||||
|
||||
fn forward(&self, x: &Tensor) -> Result<Tensor> {
|
||||
// softplus(x) ≈ log(1 + exp(beta * x)) / beta
|
||||
// For numerical stability, use: max(0, x) + log(1 + exp(-|x|))
|
||||
let zeros = Tensor::zeros_like(x)
|
||||
.map_err(|e| crate::TtsError::TensorError(format!("Zeros failed: {e}")))?;
|
||||
let relu_part = x.maximum(&zeros)
|
||||
.map_err(|e| crate::TtsError::TensorError(format!("Maximum failed: {e}")))?;
|
||||
|
||||
let abs_x = x.abs()
|
||||
.map_err(|e| crate::TtsError::TensorError(format!("Abs failed: {e}")))?;
|
||||
let neg_abs = abs_x.mul_scalar(-1.0)
|
||||
.map_err(|e| crate::TtsError::TensorError(format!("Mul failed: {e}")))?;
|
||||
let exp_part = neg_abs.exp()
|
||||
.map_err(|e| crate::TtsError::TensorError(format!("Exp failed: {e}")))?;
|
||||
let one_plus_exp = exp_part.add_scalar(1.0)
|
||||
.map_err(|e| crate::TtsError::TensorError(format!("Add failed: {e}")))?;
|
||||
let log_part = one_plus_exp.log()
|
||||
.map_err(|e| crate::TtsError::TensorError(format!("Log failed: {e}")))?;
|
||||
|
||||
(&relu_part + &log_part)
|
||||
.map_err(|e| crate::TtsError::TensorError(format!("Add failed: {e}")))
|
||||
}
|
||||
}
|
||||
|
||||
/// Simplified layer norm (just a linear projection for now)
|
||||
#[derive(Debug)]
|
||||
struct SimpleNorm {
|
||||
linear: Linear,
|
||||
}
|
||||
|
||||
impl SimpleNorm {
|
||||
fn new(dim: usize, device: &Device) -> Result<Self> {
|
||||
let linear = Linear::new(dim, dim, true, device)
|
||||
.map_err(|e| crate::TtsError::ModelError(format!("Linear creation failed: {e}")))?;
|
||||
Ok(Self { linear })
|
||||
}
|
||||
|
||||
fn forward(&self, x: &Tensor) -> Result<Tensor> {
|
||||
self.linear.forward(x)
|
||||
.map_err(|e| crate::TtsError::ModelError(format!("Linear forward failed: {e}")))
|
||||
}
|
||||
}
|
||||
|
||||
/// Simplified dropout (identity in eval, scale in training)
|
||||
#[derive(Debug)]
|
||||
struct SimpleDropout {
|
||||
p: f32,
|
||||
training: bool,
|
||||
}
|
||||
|
||||
impl SimpleDropout {
|
||||
fn new(p: f32) -> Self {
|
||||
Self { p, training: true }
|
||||
}
|
||||
|
||||
fn forward(&self, x: &Tensor) -> Result<Tensor> {
|
||||
if self.training && self.p > 0.0 {
|
||||
// Simple scaling approximation
|
||||
x.mul_scalar(1.0 - self.p)
|
||||
.map_err(|e| crate::TtsError::TensorError(format!("Dropout failed: {e}")))
|
||||
} else {
|
||||
Ok(x.clone())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Duration predictor for phoneme duration
|
||||
#[derive(Debug)]
|
||||
pub struct DurationPredictor {
|
||||
conv_layers: Vec<Conv1d>,
|
||||
norms: Vec<SimpleNorm>,
|
||||
dropout: SimpleDropout,
|
||||
linear: Linear,
|
||||
activation: ReLU,
|
||||
softplus: SoftplusApprox,
|
||||
device: Device,
|
||||
training: bool,
|
||||
}
|
||||
|
||||
impl DurationPredictor {
|
||||
/// Create a new duration predictor
|
||||
pub fn new(config: &VarianceAdaptorConfig, device: &Device) -> Result<Self> {
|
||||
let mut conv_layers = Vec::new();
|
||||
let mut norms = Vec::new();
|
||||
|
||||
let padding = config.kernel_size / 2;
|
||||
|
||||
for _ in 0..config.n_layers {
|
||||
let conv_config = Conv1dConfig {
|
||||
in_channels: config.hidden_dim,
|
||||
out_channels: config.hidden_dim,
|
||||
kernel_size: config.kernel_size,
|
||||
stride: 1,
|
||||
padding,
|
||||
dilation: 1,
|
||||
groups: 1,
|
||||
bias: true,
|
||||
};
|
||||
|
||||
conv_layers.push(Conv1d::from_config(conv_config, device)
|
||||
.map_err(|e| crate::TtsError::ModelError(format!("Conv1d creation failed: {e}")))?);
|
||||
|
||||
norms.push(SimpleNorm::new(config.hidden_dim, device)?);
|
||||
}
|
||||
|
||||
let linear = Linear::new(config.hidden_dim, 1, true, device)
|
||||
.map_err(|e| crate::TtsError::ModelError(format!("Linear creation failed: {e}")))?;
|
||||
|
||||
Ok(Self {
|
||||
conv_layers,
|
||||
norms,
|
||||
dropout: SimpleDropout::new(config.dropout),
|
||||
linear,
|
||||
activation: ReLU::new(),
|
||||
softplus: SoftplusApprox::new(1.0),
|
||||
device: device.clone(),
|
||||
training: true,
|
||||
})
|
||||
}
|
||||
|
||||
/// Forward pass
|
||||
pub fn forward(&self, x: &Tensor) -> Result<Tensor> {
|
||||
// x: [batch, seq_len, hidden_dim]
|
||||
// Transpose to [batch, hidden_dim, seq_len] for Conv1d
|
||||
let mut hidden = x.transpose(1, 2)
|
||||
.map_err(|e| crate::TtsError::TensorError(format!("Transpose failed: {e}")))?;
|
||||
|
||||
for (conv, norm) in self.conv_layers.iter().zip(self.norms.iter()) {
|
||||
let conv_out = conv.forward(&hidden)
|
||||
.map_err(|e| crate::TtsError::ModelError(format!("Conv1d forward failed: {e}")))?;
|
||||
|
||||
// Transpose back for norm: [batch, seq_len, hidden_dim]
|
||||
let transposed = conv_out.transpose(1, 2)
|
||||
.map_err(|e| crate::TtsError::TensorError(format!("Transpose failed: {e}")))?;
|
||||
|
||||
let normed = norm.forward(&transposed)?;
|
||||
|
||||
let activated = self.activation.forward(&normed)
|
||||
.map_err(|e| crate::TtsError::ModelError(format!("ReLU forward failed: {e}")))?;
|
||||
|
||||
let dropped = self.dropout.forward(&activated)?;
|
||||
|
||||
// Transpose back to [batch, hidden_dim, seq_len] for next conv
|
||||
hidden = dropped.transpose(1, 2)
|
||||
.map_err(|e| crate::TtsError::TensorError(format!("Transpose failed: {e}")))?;
|
||||
}
|
||||
|
||||
// Transpose to [batch, seq_len, hidden_dim] for linear
|
||||
let hidden = hidden.transpose(1, 2)
|
||||
.map_err(|e| crate::TtsError::TensorError(format!("Transpose failed: {e}")))?;
|
||||
|
||||
let output = self.linear.forward(&hidden)
|
||||
.map_err(|e| crate::TtsError::ModelError(format!("Linear forward failed: {e}")))?;
|
||||
|
||||
// Squeeze last dimension: [batch, seq_len, 1] -> [batch, seq_len]
|
||||
let squeezed = output.squeeze(2)
|
||||
.map_err(|e| crate::TtsError::TensorError(format!("Squeeze failed: {e}")))?;
|
||||
|
||||
// Apply softplus to ensure positive durations
|
||||
self.softplus.forward(&squeezed)
|
||||
}
|
||||
|
||||
/// Set training mode
|
||||
pub fn set_training(&mut self, training: bool) {
|
||||
self.training = training;
|
||||
for conv in &mut self.conv_layers {
|
||||
conv.train(training);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Pitch predictor for F0 prediction
|
||||
#[derive(Debug)]
|
||||
pub struct PitchPredictor {
|
||||
conv_layers: Vec<Conv1d>,
|
||||
norms: Vec<SimpleNorm>,
|
||||
dropout: SimpleDropout,
|
||||
linear: Linear,
|
||||
activation: ReLU,
|
||||
device: Device,
|
||||
training: bool,
|
||||
}
|
||||
|
||||
impl PitchPredictor {
|
||||
/// Create a new pitch predictor
|
||||
pub fn new(config: &VarianceAdaptorConfig, device: &Device) -> Result<Self> {
|
||||
let mut conv_layers = Vec::new();
|
||||
let mut norms = Vec::new();
|
||||
|
||||
let padding = config.kernel_size / 2;
|
||||
|
||||
for _ in 0..config.n_layers {
|
||||
let conv_config = Conv1dConfig {
|
||||
in_channels: config.hidden_dim,
|
||||
out_channels: config.hidden_dim,
|
||||
kernel_size: config.kernel_size,
|
||||
stride: 1,
|
||||
padding,
|
||||
dilation: 1,
|
||||
groups: 1,
|
||||
bias: true,
|
||||
};
|
||||
|
||||
conv_layers.push(Conv1d::from_config(conv_config, device)
|
||||
.map_err(|e| crate::TtsError::ModelError(format!("Conv1d creation failed: {e}")))?);
|
||||
|
||||
norms.push(SimpleNorm::new(config.hidden_dim, device)?);
|
||||
}
|
||||
|
||||
let linear = Linear::new(config.hidden_dim, 1, true, device)
|
||||
.map_err(|e| crate::TtsError::ModelError(format!("Linear creation failed: {e}")))?;
|
||||
|
||||
Ok(Self {
|
||||
conv_layers,
|
||||
norms,
|
||||
dropout: SimpleDropout::new(config.dropout),
|
||||
linear,
|
||||
activation: ReLU::new(),
|
||||
device: device.clone(),
|
||||
training: true,
|
||||
})
|
||||
}
|
||||
|
||||
/// Forward pass
|
||||
pub fn forward(&self, x: &Tensor) -> Result<Tensor> {
|
||||
let mut hidden = x.transpose(1, 2)
|
||||
.map_err(|e| crate::TtsError::TensorError(format!("Transpose failed: {e}")))?;
|
||||
|
||||
for (conv, norm) in self.conv_layers.iter().zip(self.norms.iter()) {
|
||||
let conv_out = conv.forward(&hidden)
|
||||
.map_err(|e| crate::TtsError::ModelError(format!("Conv1d forward failed: {e}")))?;
|
||||
|
||||
let transposed = conv_out.transpose(1, 2)
|
||||
.map_err(|e| crate::TtsError::TensorError(format!("Transpose failed: {e}")))?;
|
||||
|
||||
let normed = norm.forward(&transposed)?;
|
||||
|
||||
let activated = self.activation.forward(&normed)
|
||||
.map_err(|e| crate::TtsError::ModelError(format!("ReLU forward failed: {e}")))?;
|
||||
|
||||
let dropped = self.dropout.forward(&activated)?;
|
||||
|
||||
hidden = dropped.transpose(1, 2)
|
||||
.map_err(|e| crate::TtsError::TensorError(format!("Transpose failed: {e}")))?;
|
||||
}
|
||||
|
||||
let hidden = hidden.transpose(1, 2)
|
||||
.map_err(|e| crate::TtsError::TensorError(format!("Transpose failed: {e}")))?;
|
||||
|
||||
let output = self.linear.forward(&hidden)
|
||||
.map_err(|e| crate::TtsError::ModelError(format!("Linear forward failed: {e}")))?;
|
||||
|
||||
output.squeeze(2)
|
||||
.map_err(|e| crate::TtsError::TensorError(format!("Squeeze failed: {e}")))
|
||||
}
|
||||
|
||||
/// Set training mode
|
||||
pub fn set_training(&mut self, training: bool) {
|
||||
self.training = training;
|
||||
for conv in &mut self.conv_layers {
|
||||
conv.train(training);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Energy predictor
|
||||
#[derive(Debug)]
|
||||
pub struct EnergyPredictor {
|
||||
conv_layers: Vec<Conv1d>,
|
||||
norms: Vec<SimpleNorm>,
|
||||
dropout: SimpleDropout,
|
||||
linear: Linear,
|
||||
activation: ReLU,
|
||||
softplus: SoftplusApprox,
|
||||
device: Device,
|
||||
training: bool,
|
||||
}
|
||||
|
||||
impl EnergyPredictor {
|
||||
/// Create a new energy predictor
|
||||
pub fn new(config: &VarianceAdaptorConfig, device: &Device) -> Result<Self> {
|
||||
let mut conv_layers = Vec::new();
|
||||
let mut norms = Vec::new();
|
||||
|
||||
let padding = config.kernel_size / 2;
|
||||
|
||||
for _ in 0..config.n_layers {
|
||||
let conv_config = Conv1dConfig {
|
||||
in_channels: config.hidden_dim,
|
||||
out_channels: config.hidden_dim,
|
||||
kernel_size: config.kernel_size,
|
||||
stride: 1,
|
||||
padding,
|
||||
dilation: 1,
|
||||
groups: 1,
|
||||
bias: true,
|
||||
};
|
||||
|
||||
conv_layers.push(Conv1d::from_config(conv_config, device)
|
||||
.map_err(|e| crate::TtsError::ModelError(format!("Conv1d creation failed: {e}")))?);
|
||||
|
||||
norms.push(SimpleNorm::new(config.hidden_dim, device)?);
|
||||
}
|
||||
|
||||
let linear = Linear::new(config.hidden_dim, 1, true, device)
|
||||
.map_err(|e| crate::TtsError::ModelError(format!("Linear creation failed: {e}")))?;
|
||||
|
||||
Ok(Self {
|
||||
conv_layers,
|
||||
norms,
|
||||
dropout: SimpleDropout::new(config.dropout),
|
||||
linear,
|
||||
activation: ReLU::new(),
|
||||
softplus: SoftplusApprox::new(1.0),
|
||||
device: device.clone(),
|
||||
training: true,
|
||||
})
|
||||
}
|
||||
|
||||
/// Forward pass
|
||||
pub fn forward(&self, x: &Tensor) -> Result<Tensor> {
|
||||
let mut hidden = x.transpose(1, 2)
|
||||
.map_err(|e| crate::TtsError::TensorError(format!("Transpose failed: {e}")))?;
|
||||
|
||||
for (conv, norm) in self.conv_layers.iter().zip(self.norms.iter()) {
|
||||
let conv_out = conv.forward(&hidden)
|
||||
.map_err(|e| crate::TtsError::ModelError(format!("Conv1d forward failed: {e}")))?;
|
||||
|
||||
let transposed = conv_out.transpose(1, 2)
|
||||
.map_err(|e| crate::TtsError::TensorError(format!("Transpose failed: {e}")))?;
|
||||
|
||||
let normed = norm.forward(&transposed)?;
|
||||
|
||||
let activated = self.activation.forward(&normed)
|
||||
.map_err(|e| crate::TtsError::ModelError(format!("ReLU forward failed: {e}")))?;
|
||||
|
||||
let dropped = self.dropout.forward(&activated)?;
|
||||
|
||||
hidden = dropped.transpose(1, 2)
|
||||
.map_err(|e| crate::TtsError::TensorError(format!("Transpose failed: {e}")))?;
|
||||
}
|
||||
|
||||
let hidden = hidden.transpose(1, 2)
|
||||
.map_err(|e| crate::TtsError::TensorError(format!("Transpose failed: {e}")))?;
|
||||
|
||||
let output = self.linear.forward(&hidden)
|
||||
.map_err(|e| crate::TtsError::ModelError(format!("Linear forward failed: {e}")))?;
|
||||
|
||||
let squeezed = output.squeeze(2)
|
||||
.map_err(|e| crate::TtsError::TensorError(format!("Squeeze failed: {e}")))?;
|
||||
|
||||
// Apply softplus to ensure positive energy
|
||||
self.softplus.forward(&squeezed)
|
||||
}
|
||||
|
||||
/// Set training mode
|
||||
pub fn set_training(&mut self, training: bool) {
|
||||
self.training = training;
|
||||
for conv in &mut self.conv_layers {
|
||||
conv.train(training);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Length regulator for duration-based expansion
|
||||
#[derive(Debug)]
|
||||
pub struct LengthRegulator {
|
||||
device: Device,
|
||||
}
|
||||
|
||||
impl LengthRegulator {
|
||||
/// Create a new length regulator
|
||||
pub fn new(device: &Device) -> Self {
|
||||
Self {
|
||||
device: device.clone(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Regulate sequence length based on predicted durations
|
||||
pub fn forward(&self, hidden: &Tensor, durations: &Tensor) -> Result<Tensor> {
|
||||
let shape = hidden.shape();
|
||||
let batch_size = shape.dims()[0];
|
||||
let seq_len = shape.dims()[1];
|
||||
let hidden_dim = shape.dims()[2];
|
||||
|
||||
// Convert durations to integers
|
||||
let durations_data = durations.to_vec::<f32>()
|
||||
.map_err(|e| crate::TtsError::TensorError(format!("Failed to get durations: {e}")))?;
|
||||
|
||||
let mut outputs = Vec::new();
|
||||
|
||||
for b in 0..batch_size {
|
||||
let mut expanded = Vec::new();
|
||||
|
||||
for i in 0..seq_len {
|
||||
let duration = durations_data[b * seq_len + i].round().max(0.0) as usize;
|
||||
|
||||
// Extract the hidden state for this position
|
||||
let start_idx = b * seq_len * hidden_dim + i * hidden_dim;
|
||||
let end_idx = start_idx + hidden_dim;
|
||||
|
||||
let hidden_data = hidden.to_vec::<f32>()
|
||||
.map_err(|e| crate::TtsError::TensorError(format!("Failed to get hidden: {e}")))?;
|
||||
|
||||
let frame = &hidden_data[start_idx..end_idx];
|
||||
|
||||
// Repeat the frame 'duration' times
|
||||
for _ in 0..duration {
|
||||
expanded.extend_from_slice(frame);
|
||||
}
|
||||
}
|
||||
|
||||
outputs.push(expanded);
|
||||
}
|
||||
|
||||
// Find max length
|
||||
let max_len = outputs.iter().map(|x| x.len() / hidden_dim).max().unwrap_or(0);
|
||||
|
||||
// Pad all sequences to max length
|
||||
let mut padded = Vec::new();
|
||||
for mut seq in outputs {
|
||||
let current_len = seq.len() / hidden_dim;
|
||||
let padding_len = (max_len - current_len) * hidden_dim;
|
||||
seq.extend(vec![0.0f32; padding_len]);
|
||||
padded.extend(seq);
|
||||
}
|
||||
|
||||
Tensor::from_vec(padded, &[batch_size, max_len, hidden_dim], &self.device)
|
||||
.map_err(|e| crate::TtsError::TensorError(format!("Failed to create expanded tensor: {e}")))
|
||||
}
|
||||
}
|
||||
|
||||
/// Combined variance adaptor
|
||||
#[derive(Debug)]
|
||||
pub struct VarianceAdaptor {
|
||||
duration_predictor: DurationPredictor,
|
||||
pitch_predictor: PitchPredictor,
|
||||
energy_predictor: EnergyPredictor,
|
||||
length_regulator: LengthRegulator,
|
||||
pitch_embedding: Linear,
|
||||
energy_embedding: Linear,
|
||||
device: Device,
|
||||
training: bool,
|
||||
}
|
||||
|
||||
impl VarianceAdaptor {
|
||||
/// Create a new variance adaptor
|
||||
pub fn new(config: VarianceAdaptorConfig, device: &Device) -> Result<Self> {
|
||||
let duration_predictor = DurationPredictor::new(&config, device)?;
|
||||
let pitch_predictor = PitchPredictor::new(&config, device)?;
|
||||
let energy_predictor = EnergyPredictor::new(&config, device)?;
|
||||
let length_regulator = LengthRegulator::new(device);
|
||||
|
||||
let pitch_embedding = Linear::new(1, config.hidden_dim, true, device)
|
||||
.map_err(|e| crate::TtsError::ModelError(format!("Pitch embedding creation failed: {e}")))?;
|
||||
|
||||
let energy_embedding = Linear::new(1, config.hidden_dim, true, device)
|
||||
.map_err(|e| crate::TtsError::ModelError(format!("Energy embedding creation failed: {e}")))?;
|
||||
|
||||
Ok(Self {
|
||||
duration_predictor,
|
||||
pitch_predictor,
|
||||
energy_predictor,
|
||||
length_regulator,
|
||||
pitch_embedding,
|
||||
energy_embedding,
|
||||
device: device.clone(),
|
||||
training: true,
|
||||
})
|
||||
}
|
||||
|
||||
/// Forward pass with variance prediction
|
||||
pub fn forward(&self, hidden: &Tensor) -> Result<(Tensor, Tensor, Tensor, Tensor)> {
|
||||
// Predict variance parameters
|
||||
let duration = self.duration_predictor.forward(hidden)?;
|
||||
let pitch = self.pitch_predictor.forward(hidden)?;
|
||||
let energy = self.energy_predictor.forward(hidden)?;
|
||||
|
||||
// Embed pitch and energy
|
||||
let pitch_unsqueezed = pitch.unsqueeze(2)
|
||||
.map_err(|e| crate::TtsError::TensorError(format!("Unsqueeze failed: {e}")))?;
|
||||
let energy_unsqueezed = energy.unsqueeze(2)
|
||||
.map_err(|e| crate::TtsError::TensorError(format!("Unsqueeze failed: {e}")))?;
|
||||
|
||||
let pitch_emb = self.pitch_embedding.forward(&pitch_unsqueezed)
|
||||
.map_err(|e| crate::TtsError::ModelError(format!("Pitch embedding failed: {e}")))?;
|
||||
let energy_emb = self.energy_embedding.forward(&energy_unsqueezed)
|
||||
.map_err(|e| crate::TtsError::ModelError(format!("Energy embedding failed: {e}")))?;
|
||||
|
||||
// Add pitch and energy to hidden states
|
||||
let hidden_with_pitch = (hidden + &pitch_emb)
|
||||
.map_err(|e| crate::TtsError::TensorError(format!("Add pitch failed: {e}")))?;
|
||||
let hidden_with_variance = (&hidden_with_pitch + &energy_emb)
|
||||
.map_err(|e| crate::TtsError::TensorError(format!("Add energy failed: {e}")))?;
|
||||
|
||||
// Length regulation (expand by durations)
|
||||
let expanded = self.length_regulator.forward(&hidden_with_variance, &duration)?;
|
||||
|
||||
Ok((expanded, duration, pitch, energy))
|
||||
}
|
||||
|
||||
/// Set training mode
|
||||
pub fn set_training(&mut self, training: bool) {
|
||||
self.training = training;
|
||||
self.duration_predictor.set_training(training);
|
||||
self.pitch_predictor.set_training(training);
|
||||
self.energy_predictor.set_training(training);
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use rtx_tensor::Device;
|
||||
|
||||
fn get_device() -> Device {
|
||||
Device::cuda(0).unwrap_or(Device::default())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_variance_adaptor_config_default() {
|
||||
let config = VarianceAdaptorConfig::default();
|
||||
assert_eq!(config.hidden_dim, 256);
|
||||
assert_eq!(config.kernel_size, 3);
|
||||
assert_eq!(config.n_layers, 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_duration_predictor_creation() {
|
||||
let device = get_device();
|
||||
let config = VarianceAdaptorConfig::default();
|
||||
let predictor = DurationPredictor::new(&config, &device);
|
||||
assert!(predictor.is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_duration_predictor_forward() {
|
||||
let device = get_device();
|
||||
let config = VarianceAdaptorConfig::default();
|
||||
let predictor = DurationPredictor::new(&config, &device).unwrap();
|
||||
|
||||
let batch_size = 2;
|
||||
let seq_len = 10;
|
||||
let input = Tensor::randn(&[batch_size, seq_len, config.hidden_dim], &device).unwrap();
|
||||
|
||||
let output = predictor.forward(&input);
|
||||
assert!(output.is_ok());
|
||||
|
||||
let output = output.unwrap();
|
||||
let shape = output.shape();
|
||||
assert_eq!(shape.dims()[0], batch_size);
|
||||
assert_eq!(shape.dims()[1], seq_len);
|
||||
|
||||
// Check all durations are positive
|
||||
let data = output.to_vec::<f32>().unwrap();
|
||||
assert!(data.iter().all(|&x| x >= 0.0));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_pitch_predictor_creation() {
|
||||
let device = get_device();
|
||||
let config = VarianceAdaptorConfig::default();
|
||||
let predictor = PitchPredictor::new(&config, &device);
|
||||
assert!(predictor.is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_energy_predictor_creation() {
|
||||
let device = get_device();
|
||||
let config = VarianceAdaptorConfig::default();
|
||||
let predictor = EnergyPredictor::new(&config, &device);
|
||||
assert!(predictor.is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_length_regulator_creation() {
|
||||
let device = get_device();
|
||||
let _regulator = LengthRegulator::new(&device);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_variance_adaptor_creation() {
|
||||
let device = get_device();
|
||||
let config = VarianceAdaptorConfig::default();
|
||||
let adaptor = VarianceAdaptor::new(config, &device);
|
||||
assert!(adaptor.is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_config_serialization() {
|
||||
let config = VarianceAdaptorConfig::default();
|
||||
let json = serde_json::to_string(&config).unwrap();
|
||||
let deserialized: VarianceAdaptorConfig = serde_json::from_str(&json).unwrap();
|
||||
assert_eq!(config.hidden_dim, deserialized.hidden_dim);
|
||||
assert_eq!(config.kernel_size, deserialized.kernel_size);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user