Initial commit
This commit is contained in:
redclawsystems
@@ -0,0 +1,760 @@
|
||||
//! Tacotron 2 acoustic model
|
||||
//!
|
||||
//! Tacotron 2 is an autoregressive sequence-to-sequence model that generates
|
||||
//! mel spectrograms from character/phoneme sequences using attention.
|
||||
|
||||
use crate::acoustic::{AcousticModel, MelSpectrogramConfig};
|
||||
use crate::Result;
|
||||
use rtx_nn::layers::{Module, linear::Linear, conv::{Conv1d, Conv1dConfig}};
|
||||
use rtx_nn::layers::activation::{ReLU, Tanh, Sigmoid};
|
||||
use rtx_nn::layers::dropout::Dropout;
|
||||
use rtx_nn::layers::embedding::{Embedding, EmbeddingConfig};
|
||||
use rtx_tensor::{Tensor, Device};
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
/// Configuration for Tacotron2
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct Tacotron2Config {
|
||||
/// Encoder embedding dimension
|
||||
pub encoder_dim: usize,
|
||||
|
||||
/// Decoder LSTM dimension
|
||||
pub decoder_dim: usize,
|
||||
|
||||
/// Attention dimension
|
||||
pub attention_dim: usize,
|
||||
|
||||
/// Prenet dimension
|
||||
pub prenet_dim: usize,
|
||||
|
||||
/// Postnet convolutional channels
|
||||
pub postnet_channels: usize,
|
||||
|
||||
/// Mel spectrogram dimension
|
||||
pub mel_dim: usize,
|
||||
|
||||
/// Maximum decoder steps (for inference)
|
||||
pub max_decoder_steps: usize,
|
||||
|
||||
/// Character/phoneme vocabulary size
|
||||
pub vocab_size: usize,
|
||||
|
||||
/// Dropout probability
|
||||
pub dropout: f32,
|
||||
|
||||
/// Number of encoder conv layers
|
||||
pub encoder_n_convs: usize,
|
||||
|
||||
/// Number of postnet conv layers
|
||||
pub postnet_n_convs: usize,
|
||||
|
||||
/// Mel spectrogram configuration
|
||||
pub mel_config: MelSpectrogramConfig,
|
||||
}
|
||||
|
||||
impl Default for Tacotron2Config {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
encoder_dim: 512,
|
||||
decoder_dim: 1024,
|
||||
attention_dim: 128,
|
||||
prenet_dim: 256,
|
||||
postnet_channels: 512,
|
||||
mel_dim: 80,
|
||||
max_decoder_steps: 1000,
|
||||
vocab_size: 100,
|
||||
dropout: 0.5,
|
||||
encoder_n_convs: 3,
|
||||
postnet_n_convs: 5,
|
||||
mel_config: MelSpectrogramConfig::default(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Tacotron2Config {
|
||||
/// Validate configuration
|
||||
pub fn validate(&self) -> Result<()> {
|
||||
if self.encoder_dim == 0 {
|
||||
return Err(crate::TtsError::InvalidConfig("encoder_dim must be > 0".into()));
|
||||
}
|
||||
if self.decoder_dim == 0 {
|
||||
return Err(crate::TtsError::InvalidConfig("decoder_dim must be > 0".into()));
|
||||
}
|
||||
if self.mel_dim == 0 {
|
||||
return Err(crate::TtsError::InvalidConfig("mel_dim must be > 0".into()));
|
||||
}
|
||||
self.mel_config.validate()?;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
/// Prenet for mel spectrogram processing
|
||||
#[derive(Debug)]
|
||||
struct PreNet {
|
||||
fc1: Linear,
|
||||
fc2: Linear,
|
||||
dropout: Dropout,
|
||||
activation: ReLU,
|
||||
training: bool,
|
||||
}
|
||||
|
||||
impl PreNet {
|
||||
fn new(in_dim: usize, prenet_dim: usize, dropout: f32, device: &Device) -> Result<Self> {
|
||||
let fc1 = Linear::new(in_dim, prenet_dim, true, device)
|
||||
.map_err(|e| crate::TtsError::ModelError(format!("FC1 creation failed: {e}")))?;
|
||||
let fc2 = Linear::new(prenet_dim, prenet_dim, true, device)
|
||||
.map_err(|e| crate::TtsError::ModelError(format!("FC2 creation failed: {e}")))?;
|
||||
|
||||
Ok(Self {
|
||||
fc1,
|
||||
fc2,
|
||||
dropout: Dropout::new(dropout),
|
||||
activation: ReLU::new(),
|
||||
training: true,
|
||||
})
|
||||
}
|
||||
|
||||
fn forward(&self, x: &Tensor) -> Result<Tensor> {
|
||||
let h1 = self.fc1.forward(x)
|
||||
.map_err(|e| crate::TtsError::ModelError(format!("FC1 forward failed: {e}")))?;
|
||||
let h1 = self.activation.forward(&h1)
|
||||
.map_err(|e| crate::TtsError::ModelError(format!("ReLU failed: {e}")))?;
|
||||
let h1 = self.dropout.forward(&h1)
|
||||
.map_err(|e| crate::TtsError::ModelError(format!("Dropout failed: {e}")))?;
|
||||
|
||||
let h2 = self.fc2.forward(&h1)
|
||||
.map_err(|e| crate::TtsError::ModelError(format!("FC2 forward failed: {e}")))?;
|
||||
let h2 = self.activation.forward(&h2)
|
||||
.map_err(|e| crate::TtsError::ModelError(format!("ReLU failed: {e}")))?;
|
||||
|
||||
self.dropout.forward(&h2)
|
||||
.map_err(|e| crate::TtsError::ModelError(format!("Dropout failed: {e}")))
|
||||
}
|
||||
|
||||
fn set_training(&mut self, training: bool) {
|
||||
self.training = training;
|
||||
}
|
||||
}
|
||||
|
||||
/// Postnet for mel spectrogram refinement
|
||||
#[derive(Debug)]
|
||||
struct PostNet {
|
||||
conv_layers: Vec<Conv1d>,
|
||||
batch_norm_layers: Vec<Linear>, // Simplified batch norm as linear
|
||||
dropout: Dropout,
|
||||
tanh: Tanh,
|
||||
training: bool,
|
||||
}
|
||||
|
||||
impl PostNet {
|
||||
fn new(config: &Tacotron2Config, device: &Device) -> Result<Self> {
|
||||
let mut conv_layers = Vec::new();
|
||||
let mut batch_norm_layers = Vec::new();
|
||||
|
||||
for i in 0..config.postnet_n_convs {
|
||||
let in_channels = if i == 0 { config.mel_dim } else { config.postnet_channels };
|
||||
let out_channels = if i == config.postnet_n_convs - 1 {
|
||||
config.mel_dim
|
||||
} else {
|
||||
config.postnet_channels
|
||||
};
|
||||
|
||||
let conv_config = Conv1dConfig {
|
||||
in_channels,
|
||||
out_channels,
|
||||
kernel_size: 5,
|
||||
stride: 1,
|
||||
padding: 2,
|
||||
dilation: 1,
|
||||
groups: 1,
|
||||
bias: true,
|
||||
};
|
||||
|
||||
conv_layers.push(Conv1d::with_config(conv_config, device)
|
||||
.map_err(|e| crate::TtsError::ModelError(format!("Conv1d creation failed: {e}")))?);
|
||||
|
||||
// Simplified batch norm
|
||||
batch_norm_layers.push(Linear::new(out_channels, out_channels, true, device)
|
||||
.map_err(|e| crate::TtsError::ModelError(format!("Batch norm creation failed: {e}")))?);
|
||||
}
|
||||
|
||||
Ok(Self {
|
||||
conv_layers,
|
||||
batch_norm_layers,
|
||||
dropout: Dropout::new(0.5),
|
||||
tanh: Tanh::new(),
|
||||
training: true,
|
||||
})
|
||||
}
|
||||
|
||||
fn forward(&self, x: &Tensor) -> Result<Tensor> {
|
||||
let mut hidden = x.clone();
|
||||
|
||||
for (i, (conv, bn)) in self.conv_layers.iter().zip(self.batch_norm_layers.iter()).enumerate() {
|
||||
let conv_out = conv.forward(&hidden)
|
||||
.map_err(|e| crate::TtsError::ModelError(format!("Conv forward failed: {e}")))?;
|
||||
|
||||
// Transpose for batch norm
|
||||
let transposed = conv_out.transpose(1, 2)
|
||||
.map_err(|e| crate::TtsError::TensorError(format!("Transpose failed: {e}")))?;
|
||||
|
||||
let bn_out = bn.forward(&transposed)
|
||||
.map_err(|e| crate::TtsError::ModelError(format!("Batch norm failed: {e}")))?;
|
||||
|
||||
// Transpose back
|
||||
hidden = bn_out.transpose(1, 2)
|
||||
.map_err(|e| crate::TtsError::TensorError(format!("Transpose failed: {e}")))?;
|
||||
|
||||
// Tanh activation except for last layer
|
||||
if i < self.conv_layers.len() - 1 {
|
||||
hidden = self.tanh.forward(&hidden)
|
||||
.map_err(|e| crate::TtsError::ModelError(format!("Tanh failed: {e}")))?;
|
||||
}
|
||||
|
||||
if self.training {
|
||||
hidden = self.dropout.forward(&hidden)
|
||||
.map_err(|e| crate::TtsError::ModelError(format!("Dropout failed: {e}")))?;
|
||||
}
|
||||
}
|
||||
|
||||
Ok(hidden)
|
||||
}
|
||||
|
||||
fn set_training(&mut self, training: bool) {
|
||||
self.training = training;
|
||||
for conv in &mut self.conv_layers {
|
||||
conv.train(training);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Location-sensitive attention
|
||||
#[derive(Debug)]
|
||||
struct LocationAttention {
|
||||
query_layer: Linear,
|
||||
memory_layer: Linear,
|
||||
location_conv: Conv1d,
|
||||
location_layer: Linear,
|
||||
v: Linear,
|
||||
attention_dim: usize,
|
||||
device: Device,
|
||||
}
|
||||
|
||||
impl LocationAttention {
|
||||
fn new(
|
||||
query_dim: usize,
|
||||
memory_dim: usize,
|
||||
attention_dim: usize,
|
||||
device: &Device,
|
||||
) -> Result<Self> {
|
||||
let query_layer = Linear::new(query_dim, attention_dim, false, device)
|
||||
.map_err(|e| crate::TtsError::ModelError(format!("Query layer creation failed: {e}")))?;
|
||||
|
||||
let memory_layer = Linear::new(memory_dim, attention_dim, false, device)
|
||||
.map_err(|e| crate::TtsError::ModelError(format!("Memory layer creation failed: {e}")))?;
|
||||
|
||||
let location_conv_config = Conv1dConfig {
|
||||
in_channels: 2,
|
||||
out_channels: 32,
|
||||
kernel_size: 31,
|
||||
stride: 1,
|
||||
padding: 15,
|
||||
dilation: 1,
|
||||
groups: 1,
|
||||
bias: true,
|
||||
};
|
||||
|
||||
let location_conv = Conv1d::with_config(location_conv_config, device)
|
||||
.map_err(|e| crate::TtsError::ModelError(format!("Location conv creation failed: {e}")))?;
|
||||
|
||||
let location_layer = Linear::new(32, attention_dim, false, device)
|
||||
.map_err(|e| crate::TtsError::ModelError(format!("Location layer creation failed: {e}")))?;
|
||||
|
||||
let v = Linear::new(attention_dim, 1, false, device)
|
||||
.map_err(|e| crate::TtsError::ModelError(format!("V layer creation failed: {e}")))?;
|
||||
|
||||
Ok(Self {
|
||||
query_layer,
|
||||
memory_layer,
|
||||
location_conv,
|
||||
location_layer,
|
||||
v,
|
||||
attention_dim,
|
||||
device: device.clone(),
|
||||
})
|
||||
}
|
||||
|
||||
fn forward(
|
||||
&self,
|
||||
query: &Tensor,
|
||||
memory: &Tensor,
|
||||
attention_weights_cat: &Tensor,
|
||||
) -> Result<Tensor> {
|
||||
let processed_query = self.query_layer.forward(query)
|
||||
.map_err(|e| crate::TtsError::ModelError(format!("Query layer failed: {e}")))?;
|
||||
|
||||
let processed_memory = self.memory_layer.forward(memory)
|
||||
.map_err(|e| crate::TtsError::ModelError(format!("Memory layer failed: {e}")))?;
|
||||
|
||||
// Process location features
|
||||
let processed_location = self.location_conv.forward(attention_weights_cat)
|
||||
.map_err(|e| crate::TtsError::ModelError(format!("Location conv failed: {e}")))?;
|
||||
|
||||
let processed_location = processed_location.transpose(1, 2)
|
||||
.map_err(|e| crate::TtsError::TensorError(format!("Transpose failed: {e}")))?;
|
||||
|
||||
let processed_location = self.location_layer.forward(&processed_location)
|
||||
.map_err(|e| crate::TtsError::ModelError(format!("Location layer failed: {e}")))?;
|
||||
|
||||
// Expand query to match memory sequence length
|
||||
let query_expanded = processed_query.unsqueeze(1)
|
||||
.map_err(|e| crate::TtsError::TensorError(format!("Unsqueeze failed: {e}")))?;
|
||||
|
||||
// Compute alignment energies
|
||||
let energies = (&query_expanded + &processed_memory)
|
||||
.map_err(|e| crate::TtsError::TensorError(format!("Add failed: {e}")))?;
|
||||
let energies = (&energies + &processed_location)
|
||||
.map_err(|e| crate::TtsError::TensorError(format!("Add failed: {e}")))?;
|
||||
|
||||
let energies = rtx_tensor::ops::tanh(&energies)
|
||||
.map_err(|e| crate::TtsError::TensorError(format!("Tanh failed: {e}")))?;
|
||||
|
||||
let alignment = self.v.forward(&energies)
|
||||
.map_err(|e| crate::TtsError::ModelError(format!("V layer failed: {e}")))?;
|
||||
|
||||
let alignment = alignment.squeeze(2)
|
||||
.map_err(|e| crate::TtsError::TensorError(format!("Squeeze failed: {e}")))?;
|
||||
|
||||
// Softmax
|
||||
rtx_tensor::ops::softmax(&alignment, 1)
|
||||
.map_err(|e| crate::TtsError::TensorError(format!("Softmax failed: {e}")))
|
||||
}
|
||||
}
|
||||
|
||||
/// Encoder for Tacotron2
|
||||
#[derive(Debug)]
|
||||
struct Encoder {
|
||||
embedding: Embedding,
|
||||
conv_layers: Vec<Conv1d>,
|
||||
dropout: Dropout,
|
||||
device: Device,
|
||||
training: bool,
|
||||
}
|
||||
|
||||
impl Encoder {
|
||||
fn new(config: &Tacotron2Config, device: &Device) -> Result<Self> {
|
||||
let emb_config = EmbeddingConfig::new(config.vocab_size, config.encoder_dim);
|
||||
let embedding = Embedding::new(emb_config, device)
|
||||
.map_err(|e| crate::TtsError::ModelError(format!("Embedding creation failed: {e}")))?;
|
||||
|
||||
let mut conv_layers = Vec::new();
|
||||
for _ in 0..config.encoder_n_convs {
|
||||
let conv_config = Conv1dConfig {
|
||||
in_channels: config.encoder_dim,
|
||||
out_channels: config.encoder_dim,
|
||||
kernel_size: 5,
|
||||
stride: 1,
|
||||
padding: 2,
|
||||
dilation: 1,
|
||||
groups: 1,
|
||||
bias: true,
|
||||
};
|
||||
|
||||
conv_layers.push(Conv1d::with_config(conv_config, device)
|
||||
.map_err(|e| crate::TtsError::ModelError(format!("Conv1d creation failed: {e}")))?);
|
||||
}
|
||||
|
||||
Ok(Self {
|
||||
embedding,
|
||||
conv_layers,
|
||||
dropout: Dropout::new(config.dropout),
|
||||
device: device.clone(),
|
||||
training: true,
|
||||
})
|
||||
}
|
||||
|
||||
fn forward(&self, x: &Tensor) -> Result<Tensor> {
|
||||
let embedded = self.embedding.forward(x)
|
||||
.map_err(|e| crate::TtsError::ModelError(format!("Embedding failed: {e}")))?;
|
||||
|
||||
// Transpose for conv: [batch, seq, dim] -> [batch, dim, seq]
|
||||
let mut hidden = embedded.transpose(1, 2)
|
||||
.map_err(|e| crate::TtsError::TensorError(format!("Transpose failed: {e}")))?;
|
||||
|
||||
for conv in &self.conv_layers {
|
||||
let conv_out = conv.forward(&hidden)
|
||||
.map_err(|e| crate::TtsError::ModelError(format!("Conv forward failed: {e}")))?;
|
||||
|
||||
hidden = rtx_tensor::ops::relu(&conv_out)
|
||||
.map_err(|e| crate::TtsError::TensorError(format!("ReLU failed: {e}")))?;
|
||||
|
||||
if self.training {
|
||||
hidden = self.dropout.forward(&hidden)
|
||||
.map_err(|e| crate::TtsError::ModelError(format!("Dropout failed: {e}")))?;
|
||||
}
|
||||
}
|
||||
|
||||
// Transpose back: [batch, dim, seq] -> [batch, seq, dim]
|
||||
hidden.transpose(1, 2)
|
||||
.map_err(|e| crate::TtsError::TensorError(format!("Transpose failed: {e}")))
|
||||
}
|
||||
|
||||
fn set_training(&mut self, training: bool) {
|
||||
self.training = training;
|
||||
for conv in &mut self.conv_layers {
|
||||
conv.train(training);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Tacotron 2 acoustic model
|
||||
#[derive(Debug)]
|
||||
pub struct Tacotron2 {
|
||||
config: Tacotron2Config,
|
||||
encoder: Encoder,
|
||||
prenet: PreNet,
|
||||
attention: LocationAttention,
|
||||
decoder_rnn_input: Linear,
|
||||
decoder_rnn_hidden: Linear,
|
||||
mel_projection: Linear,
|
||||
gate_projection: Linear,
|
||||
postnet: PostNet,
|
||||
device: Device,
|
||||
training: bool,
|
||||
}
|
||||
|
||||
impl Tacotron2 {
|
||||
/// Create a new Tacotron2 model
|
||||
pub fn new(config: Tacotron2Config, device: &Device) -> Result<Self> {
|
||||
config.validate()?;
|
||||
|
||||
let encoder = Encoder::new(&config, device)?;
|
||||
let prenet = PreNet::new(config.mel_dim, config.prenet_dim, config.dropout, device)?;
|
||||
|
||||
let attention = LocationAttention::new(
|
||||
config.decoder_dim,
|
||||
config.encoder_dim,
|
||||
config.attention_dim,
|
||||
device,
|
||||
)?;
|
||||
|
||||
let decoder_input_dim = config.prenet_dim + config.encoder_dim;
|
||||
let decoder_rnn_input = Linear::new(decoder_input_dim, config.decoder_dim, true, device)
|
||||
.map_err(|e| crate::TtsError::ModelError(format!("Decoder RNN input failed: {e}")))?;
|
||||
|
||||
let decoder_rnn_hidden = Linear::new(config.decoder_dim, config.decoder_dim, true, device)
|
||||
.map_err(|e| crate::TtsError::ModelError(format!("Decoder RNN hidden failed: {e}")))?;
|
||||
|
||||
let mel_projection = Linear::new(config.decoder_dim + config.encoder_dim, config.mel_dim, true, device)
|
||||
.map_err(|e| crate::TtsError::ModelError(format!("Mel projection failed: {e}")))?;
|
||||
|
||||
let gate_projection = Linear::new(config.decoder_dim + config.encoder_dim, 1, true, device)
|
||||
.map_err(|e| crate::TtsError::ModelError(format!("Gate projection failed: {e}")))?;
|
||||
|
||||
let postnet = PostNet::new(&config, device)?;
|
||||
|
||||
Ok(Self {
|
||||
config,
|
||||
encoder,
|
||||
prenet,
|
||||
attention,
|
||||
decoder_rnn_input,
|
||||
decoder_rnn_hidden,
|
||||
mel_projection,
|
||||
gate_projection,
|
||||
postnet,
|
||||
device: device.clone(),
|
||||
training: true,
|
||||
})
|
||||
}
|
||||
|
||||
/// Get configuration
|
||||
pub fn config(&self) -> &Tacotron2Config {
|
||||
&self.config
|
||||
}
|
||||
|
||||
/// Decode one step
|
||||
fn decode_step(
|
||||
&self,
|
||||
decoder_input: &Tensor,
|
||||
decoder_hidden: &Tensor,
|
||||
encoder_outputs: &Tensor,
|
||||
attention_weights_cat: &Tensor,
|
||||
) -> Result<(Tensor, Tensor, Tensor, Tensor)> {
|
||||
// Prenet
|
||||
let prenet_out = self.prenet.forward(decoder_input)?;
|
||||
|
||||
// Attention
|
||||
let attention_weights = self.attention.forward(
|
||||
decoder_hidden,
|
||||
encoder_outputs,
|
||||
attention_weights_cat,
|
||||
)?;
|
||||
|
||||
// Apply attention to encoder outputs
|
||||
let attention_weights_expanded = attention_weights.unsqueeze(1)
|
||||
.map_err(|e| crate::TtsError::TensorError(format!("Unsqueeze failed: {e}")))?;
|
||||
|
||||
let encoder_outputs_transposed = encoder_outputs.transpose(1, 2)
|
||||
.map_err(|e| crate::TtsError::TensorError(format!("Transpose failed: {e}")))?;
|
||||
|
||||
let attention_context = rtx_tensor::ops::matmul(&attention_weights_expanded, &encoder_outputs_transposed)
|
||||
.map_err(|e| crate::TtsError::TensorError(format!("Matmul failed: {e}")))?;
|
||||
|
||||
let attention_context = attention_context.squeeze(1)
|
||||
.map_err(|e| crate::TtsError::TensorError(format!("Squeeze failed: {e}")))?;
|
||||
|
||||
// Concatenate prenet output and attention context
|
||||
let decoder_rnn_input = rtx_tensor::ops::cat(&[&prenet_out, &attention_context], 1)
|
||||
.map_err(|e| crate::TtsError::TensorError(format!("Cat failed: {e}")))?;
|
||||
|
||||
// Decoder RNN step
|
||||
let rnn_input_proj = self.decoder_rnn_input.forward(&decoder_rnn_input)
|
||||
.map_err(|e| crate::TtsError::ModelError(format!("RNN input projection failed: {e}")))?;
|
||||
|
||||
let rnn_hidden_proj = self.decoder_rnn_hidden.forward(decoder_hidden)
|
||||
.map_err(|e| crate::TtsError::ModelError(format!("RNN hidden projection failed: {e}")))?;
|
||||
|
||||
let decoder_hidden_new = (&rnn_input_proj + &rnn_hidden_proj)
|
||||
.map_err(|e| crate::TtsError::TensorError(format!("Add failed: {e}")))?;
|
||||
|
||||
let decoder_hidden_new = rtx_tensor::ops::tanh(&decoder_hidden_new)
|
||||
.map_err(|e| crate::TtsError::TensorError(format!("Tanh failed: {e}")))?;
|
||||
|
||||
// Projection to mel
|
||||
let projection_input = rtx_tensor::ops::cat(&[&decoder_hidden_new, &attention_context], 1)
|
||||
.map_err(|e| crate::TtsError::TensorError(format!("Cat failed: {e}")))?;
|
||||
|
||||
let mel_output = self.mel_projection.forward(&projection_input)
|
||||
.map_err(|e| crate::TtsError::ModelError(format!("Mel projection failed: {e}")))?;
|
||||
|
||||
let gate_output = self.gate_projection.forward(&projection_input)
|
||||
.map_err(|e| crate::TtsError::ModelError(format!("Gate projection failed: {e}")))?;
|
||||
|
||||
Ok((mel_output, gate_output, decoder_hidden_new, attention_weights))
|
||||
}
|
||||
}
|
||||
|
||||
impl AcousticModel for Tacotron2 {
|
||||
fn encode(&self, phonemes: &Tensor, _phoneme_lengths: Option<&Tensor>) -> Result<Tensor> {
|
||||
self.encoder.forward(phonemes)
|
||||
}
|
||||
|
||||
fn decode(&self, hidden: &Tensor) -> Result<Tensor> {
|
||||
let shape = hidden.shape();
|
||||
let batch_size = shape.dims()[0];
|
||||
let seq_len = shape.dims()[1];
|
||||
|
||||
// Initialize decoder state
|
||||
let decoder_hidden = Tensor::zeros(&[batch_size, self.config.decoder_dim], &self.device)
|
||||
.map_err(|e| crate::TtsError::TensorError(format!("Zeros failed: {e}")))?;
|
||||
|
||||
let decoder_input = Tensor::zeros(&[batch_size, self.config.mel_dim], &self.device)
|
||||
.map_err(|e| crate::TtsError::TensorError(format!("Zeros failed: {e}")))?;
|
||||
|
||||
let attention_weights_cat = Tensor::zeros(&[batch_size, 2, seq_len], &self.device)
|
||||
.map_err(|e| crate::TtsError::TensorError(format!("Zeros failed: {e}")))?;
|
||||
|
||||
let mut mel_outputs = Vec::new();
|
||||
let mut current_hidden = decoder_hidden;
|
||||
let mut current_input = decoder_input;
|
||||
let mut current_attn = attention_weights_cat;
|
||||
|
||||
// Autoregressive decoding
|
||||
for _ in 0..self.config.max_decoder_steps.min(200) {
|
||||
let (mel_out, gate_out, new_hidden, attn_weights) =
|
||||
self.decode_step(¤t_input, ¤t_hidden, hidden, ¤t_attn)?;
|
||||
|
||||
mel_outputs.push(mel_out.clone());
|
||||
current_hidden = new_hidden;
|
||||
current_input = mel_out;
|
||||
|
||||
// Update attention weights cat
|
||||
let attn_expanded = attn_weights.unsqueeze(1)
|
||||
.map_err(|e| crate::TtsError::TensorError(format!("Unsqueeze failed: {e}")))?;
|
||||
|
||||
current_attn = rtx_tensor::ops::cat(&[
|
||||
¤t_attn.slice(1, 1, 2)
|
||||
.map_err(|e| crate::TtsError::TensorError(format!("Slice failed: {e}")))?,
|
||||
&attn_expanded
|
||||
], 1).map_err(|e| crate::TtsError::TensorError(format!("Cat failed: {e}")))?;
|
||||
|
||||
// Check gate (stop condition)
|
||||
let gate_sigmoid = rtx_tensor::ops::sigmoid(&gate_out)
|
||||
.map_err(|e| crate::TtsError::TensorError(format!("Sigmoid failed: {e}")))?;
|
||||
|
||||
let gate_val = gate_sigmoid.to_vec::<f32>()
|
||||
.map_err(|e| crate::TtsError::TensorError(format!("To vec failed: {e}")))?;
|
||||
|
||||
if gate_val.iter().all(|&v| v > 0.5) {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
// Stack mel outputs: [batch, time, mel_dim]
|
||||
let mel_stacked = rtx_tensor::ops::stack(&mel_outputs.iter().collect::<Vec<_>>(), 1)
|
||||
.map_err(|e| crate::TtsError::TensorError(format!("Stack failed: {e}")))?;
|
||||
|
||||
// Transpose to [batch, mel_dim, time]
|
||||
let mel_transposed = mel_stacked.transpose(1, 2)
|
||||
.map_err(|e| crate::TtsError::TensorError(format!("Transpose failed: {e}")))?;
|
||||
|
||||
// Apply postnet
|
||||
let postnet_out = self.postnet.forward(&mel_transposed)?;
|
||||
|
||||
// Add residual
|
||||
(&mel_transposed + &postnet_out)
|
||||
.map_err(|e| crate::TtsError::TensorError(format!("Add residual failed: {e}")))
|
||||
}
|
||||
|
||||
fn forward(&self, phonemes: &Tensor, phoneme_lengths: Option<&Tensor>) -> Result<Tensor> {
|
||||
let encoded = self.encode(phonemes, phoneme_lengths)?;
|
||||
self.decode(&encoded)
|
||||
}
|
||||
|
||||
fn mel_config(&self) -> &MelSpectrogramConfig {
|
||||
&self.config.mel_config
|
||||
}
|
||||
|
||||
fn set_training(&mut self, training: bool) {
|
||||
self.training = training;
|
||||
self.encoder.set_training(training);
|
||||
self.prenet.set_training(training);
|
||||
self.postnet.set_training(training);
|
||||
}
|
||||
|
||||
fn is_training(&self) -> bool {
|
||||
self.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_config_default() {
|
||||
let config = Tacotron2Config::default();
|
||||
assert_eq!(config.encoder_dim, 512);
|
||||
assert_eq!(config.decoder_dim, 1024);
|
||||
assert!(config.validate().is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_config_validation() {
|
||||
let mut config = Tacotron2Config::default();
|
||||
config.encoder_dim = 0;
|
||||
assert!(config.validate().is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_prenet_creation() {
|
||||
let device = get_device();
|
||||
let prenet = PreNet::new(80, 256, 0.5, &device);
|
||||
assert!(prenet.is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_prenet_forward() {
|
||||
let device = get_device();
|
||||
let prenet = PreNet::new(80, 256, 0.5, &device).unwrap();
|
||||
let input = Tensor::randn(&[2, 80], &device).unwrap();
|
||||
let output = prenet.forward(&input);
|
||||
assert!(output.is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_postnet_creation() {
|
||||
let device = get_device();
|
||||
let config = Tacotron2Config::default();
|
||||
let postnet = PostNet::new(&config, &device);
|
||||
assert!(postnet.is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_encoder_creation() {
|
||||
let device = get_device();
|
||||
let config = Tacotron2Config::default();
|
||||
let encoder = Encoder::new(&config, &device);
|
||||
assert!(encoder.is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_encoder_forward() {
|
||||
let device = get_device();
|
||||
let config = Tacotron2Config::default();
|
||||
let encoder = Encoder::new(&config, &device).unwrap();
|
||||
|
||||
let batch_size = 2;
|
||||
let seq_len = 10;
|
||||
let input = Tensor::randint(0, config.vocab_size as i64, &[batch_size, seq_len], &device).unwrap();
|
||||
|
||||
let output = encoder.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);
|
||||
assert_eq!(shape.dims()[2], config.encoder_dim);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_tacotron2_creation() {
|
||||
let device = get_device();
|
||||
let config = Tacotron2Config::default();
|
||||
let model = Tacotron2::new(config, &device);
|
||||
assert!(model.is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_tacotron2_encode() {
|
||||
let device = get_device();
|
||||
let config = Tacotron2Config::default();
|
||||
let model = Tacotron2::new(config.clone(), &device).unwrap();
|
||||
|
||||
let batch_size = 2;
|
||||
let seq_len = 10;
|
||||
let phonemes = Tensor::randint(0, config.vocab_size as i64, &[batch_size, seq_len], &device).unwrap();
|
||||
|
||||
let encoded = model.encode(&phonemes, None);
|
||||
assert!(encoded.is_ok());
|
||||
|
||||
let encoded = encoded.unwrap();
|
||||
let shape = encoded.shape();
|
||||
assert_eq!(shape.dims()[0], batch_size);
|
||||
assert_eq!(shape.dims()[1], seq_len);
|
||||
assert_eq!(shape.dims()[2], config.encoder_dim);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_tacotron2_training_mode() {
|
||||
let device = get_device();
|
||||
let config = Tacotron2Config::default();
|
||||
let mut model = Tacotron2::new(config, &device).unwrap();
|
||||
|
||||
assert!(model.is_training());
|
||||
model.set_training(false);
|
||||
assert!(!model.is_training());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_config_serialization() {
|
||||
let config = Tacotron2Config::default();
|
||||
let json = serde_json::to_string(&config).unwrap();
|
||||
let deserialized: Tacotron2Config = serde_json::from_str(&json).unwrap();
|
||||
assert_eq!(config.encoder_dim, deserialized.encoder_dim);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_location_attention_creation() {
|
||||
let device = get_device();
|
||||
let attention = LocationAttention::new(1024, 512, 128, &device);
|
||||
assert!(attention.is_ok());
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user