init

rust/src/example_onnx.rs

@@ -0,0 +1,108 @@
+use anyhow::Result;
+use clap::Parser;
+use std::path::PathBuf;
+use std::fs;
+use std::mem;
+
+mod helper;
+
+use helper::{
+    load_text_to_speech, load_voice_style, timer, write_wav_file, sanitize_filename,
+};
+
+#[derive(Parser, Debug)]
+#[command(name = "TTS ONNX Inference")]
+#[command(about = "TTS Inference with ONNX Runtime (Rust)", long_about = None)]
+struct Args {
+    /// Use GPU for inference (default: CPU)
+    #[arg(long, default_value = "false")]
+    use_gpu: bool,
+
+    /// Path to ONNX model directory
+    #[arg(long, default_value = "assets/onnx")]
+    onnx_dir: String,
+
+    /// Number of denoising steps
+    #[arg(long, default_value = "5")]
+    total_step: usize,
+
+    /// Number of times to generate
+    #[arg(long, default_value = "4")]
+    n_test: usize,
+
+    /// Voice style file path(s)
+    #[arg(long, value_delimiter = ',', default_values_t = vec!["assets/voice_styles/M1.json".to_string()])]
+    voice_style: Vec<String>,
+
+    /// Text(s) to synthesize
+    #[arg(long, value_delimiter = '|', default_values_t = vec!["This morning, I took a walk in the park, and the sound of the birds and the breeze was so pleasant that I stopped for a long time just to listen.".to_string()])]
+    text: Vec<String>,
+
+    /// Output directory
+    #[arg(long, default_value = "results")]
+    save_dir: String,
+}
+
+fn main() -> Result<()> {
+    println!("=== TTS Inference with ONNX Runtime (Rust) ===\n");
+
+    // --- 1. Parse arguments --- //
+    let args = Args::parse();
+    let total_step = args.total_step;
+    let n_test = args.n_test;
+    let voice_style_paths = &args.voice_style;
+    let text_list = &args.text;
+    let save_dir = &args.save_dir;
+
+    if voice_style_paths.len() != text_list.len() {
+        anyhow::bail!(
+            "Number of voice styles ({}) must match number of texts ({})",
+            voice_style_paths.len(),
+            text_list.len()
+        );
+    }
+
+    let bsz = voice_style_paths.len();
+
+    // --- 2. Load TTS components --- //
+    let mut text_to_speech = load_text_to_speech(&args.onnx_dir, args.use_gpu)?;
+
+    // --- 3. Load voice styles --- //
+    let style = load_voice_style(voice_style_paths, true)?;
+
+    // --- 4. Synthesize speech --- //
+    fs::create_dir_all(save_dir)?;
+
+    for n in 0..n_test {
+        println!("\n[{}/{}] Starting synthesis...", n + 1, n_test);
+
+        let (wav, duration) = timer("Generating speech from text", || {
+            text_to_speech.call(text_list, &style, total_step)
+        })?;
+
+        // Save outputs
+        let wav_len = wav.len() / bsz;
+        for i in 0..bsz {
+            let fname = format!("{}_{}.wav", sanitize_filename(&text_list[i], 20), n + 1);
+            let actual_len = (text_to_speech.sample_rate as f32 * duration[i]) as usize;
+
+            let wav_start = i * wav_len;
+            let wav_end = wav_start + actual_len.min(wav_len);
+            let wav_slice = &wav[wav_start..wav_end];
+
+            let output_path = PathBuf::from(save_dir).join(&fname);
+            write_wav_file(&output_path, wav_slice, text_to_speech.sample_rate)?;
+            println!("Saved: {}", output_path.display());
+        }
+    }
+
+    println!("\n=== Synthesis completed successfully! ===");
+    
+    // Prevent ONNX Runtime sessions from being dropped, which causes mutex cleanup issues
+    mem::forget(text_to_speech);
+    
+    // Use _exit to bypass all cleanup handlers and avoid ONNX Runtime mutex issues on macOS
+    unsafe {
+        libc::_exit(0);
+    }
+}

rust/src/helper.rs

@@ -0,0 +1,507 @@
+// ============================================================================
+// TTS Helper Module - All utility functions and structures
+// ============================================================================
+
+use ndarray::{Array, Array3};
+use serde::{Deserialize, Serialize};
+use serde_json;
+use std::fs::File;
+use std::io::BufReader;
+use std::path::Path;
+use anyhow::{Result, Context};
+use unicode_normalization::UnicodeNormalization;
+use hound::{WavWriter, WavSpec, SampleFormat};
+use rand_distr::{Distribution, Normal};
+
+// ============================================================================
+// Configuration Structures
+// ============================================================================
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct Config {
+    pub ae: AEConfig,
+    pub ttl: TTLConfig,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct AEConfig {
+    pub sample_rate: i32,
+    pub base_chunk_size: i32,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct TTLConfig {
+    pub chunk_compress_factor: i32,
+    pub latent_dim: i32,
+}
+
+/// Load configuration from JSON file
+pub fn load_cfgs<P: AsRef<Path>>(onnx_dir: P) -> Result<Config> {
+    let cfg_path = onnx_dir.as_ref().join("tts.json");
+    let file = File::open(cfg_path)?;
+    let reader = BufReader::new(file);
+    let cfgs: Config = serde_json::from_reader(reader)?;
+    Ok(cfgs)
+}
+
+// ============================================================================
+// Voice Style Data Structure
+// ============================================================================
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct VoiceStyleData {
+    pub style_ttl: StyleComponent,
+    pub style_dp: StyleComponent,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct StyleComponent {
+    pub data: Vec<Vec<Vec<f32>>>,
+    pub dims: Vec<usize>,
+    #[serde(rename = "type")]
+    pub dtype: String,
+}
+
+// ============================================================================
+// Unicode Text Processor
+// ============================================================================
+
+pub struct UnicodeProcessor {
+    indexer: Vec<i64>,
+}
+
+impl UnicodeProcessor {
+    pub fn new<P: AsRef<Path>>(unicode_indexer_json_path: P) -> Result<Self> {
+        let file = File::open(unicode_indexer_json_path)?;
+        let reader = BufReader::new(file);
+        let indexer: Vec<i64> = serde_json::from_reader(reader)?;
+        Ok(UnicodeProcessor { indexer })
+    }
+
+    pub fn call(&self, text_list: &[String]) -> (Vec<Vec<i64>>, Array3<f32>) {
+        let processed_texts: Vec<String> = text_list
+            .iter()
+            .map(|t| preprocess_text(t))
+            .collect();
+
+        let text_ids_lengths: Vec<usize> = processed_texts
+            .iter()
+            .map(|t| t.chars().count())
+            .collect();
+
+        let max_len = *text_ids_lengths.iter().max().unwrap_or(&0);
+
+        let mut text_ids = Vec::new();
+        for text in &processed_texts {
+            let mut row = vec![0i64; max_len];
+            let unicode_vals = text_to_unicode_values(text);
+            for (j, &val) in unicode_vals.iter().enumerate() {
+                if val < self.indexer.len() {
+                    row[j] = self.indexer[val];
+                } else {
+                    row[j] = -1;
+                }
+            }
+            text_ids.push(row);
+        }
+
+        let text_mask = get_text_mask(&text_ids_lengths);
+
+        (text_ids, text_mask)
+    }
+}
+
+pub fn preprocess_text(text: &str) -> String {
+    text.nfkd().collect()
+}
+
+pub fn text_to_unicode_values(text: &str) -> Vec<usize> {
+    text.chars().map(|c| c as usize).collect()
+}
+
+pub fn length_to_mask(lengths: &[usize], max_len: Option<usize>) -> Array3<f32> {
+    let bsz = lengths.len();
+    let max_len = max_len.unwrap_or_else(|| *lengths.iter().max().unwrap_or(&0));
+
+    let mut mask = Array3::<f32>::zeros((bsz, 1, max_len));
+    for (i, &len) in lengths.iter().enumerate() {
+        for j in 0..len.min(max_len) {
+            mask[[i, 0, j]] = 1.0;
+        }
+    }
+    mask
+}
+
+pub fn get_text_mask(text_ids_lengths: &[usize]) -> Array3<f32> {
+    let max_len = *text_ids_lengths.iter().max().unwrap_or(&0);
+    length_to_mask(text_ids_lengths, Some(max_len))
+}
+
+/// Sample noisy latent from normal distribution and apply mask
+pub fn sample_noisy_latent(
+    duration: &[f32],
+    sample_rate: i32,
+    base_chunk_size: i32,
+    chunk_compress: i32,
+    latent_dim: i32,
+) -> (Array3<f32>, Array3<f32>) {
+    let bsz = duration.len();
+    let max_dur = duration.iter().fold(0.0f32, |a, &b| a.max(b));
+
+    let wav_len_max = (max_dur * sample_rate as f32) as usize;
+    let wav_lengths: Vec<usize> = duration
+        .iter()
+        .map(|&d| (d * sample_rate as f32) as usize)
+        .collect();
+
+    let chunk_size = (base_chunk_size * chunk_compress) as usize;
+    let latent_len = (wav_len_max + chunk_size - 1) / chunk_size;
+    let latent_dim_val = (latent_dim * chunk_compress) as usize;
+
+    let mut noisy_latent = Array3::<f32>::zeros((bsz, latent_dim_val, latent_len));
+
+    let normal = Normal::new(0.0, 1.0).unwrap();
+    let mut rng = rand::thread_rng();
+
+    for b in 0..bsz {
+        for d in 0..latent_dim_val {
+            for t in 0..latent_len {
+                noisy_latent[[b, d, t]] = normal.sample(&mut rng);
+            }
+        }
+    }
+
+    let latent_lengths: Vec<usize> = wav_lengths
+        .iter()
+        .map(|&len| (len + chunk_size - 1) / chunk_size)
+        .collect();
+
+    let latent_mask = length_to_mask(&latent_lengths, Some(latent_len));
+
+    // Apply mask
+    for b in 0..bsz {
+        for d in 0..latent_dim_val {
+            for t in 0..latent_len {
+                noisy_latent[[b, d, t]] *= latent_mask[[b, 0, t]];
+            }
+        }
+    }
+
+    (noisy_latent, latent_mask)
+}
+
+// ============================================================================
+// WAV File I/O
+// ============================================================================
+
+pub fn write_wav_file<P: AsRef<Path>>(
+    filename: P,
+    audio_data: &[f32],
+    sample_rate: i32,
+) -> Result<()> {
+    let spec = WavSpec {
+        channels: 1,
+        sample_rate: sample_rate as u32,
+        bits_per_sample: 16,
+        sample_format: SampleFormat::Int,
+    };
+
+    let mut writer = WavWriter::create(filename, spec)?;
+
+    for &sample in audio_data {
+        let clamped = sample.max(-1.0).min(1.0);
+        let val = (clamped * 32767.0) as i16;
+        writer.write_sample(val)?;
+    }
+
+    writer.finalize()?;
+    Ok(())
+}
+
+// ============================================================================
+// Utility Functions
+// ============================================================================
+
+pub fn timer<F, T>(name: &str, f: F) -> Result<T>
+where
+    F: FnOnce() -> Result<T>,
+{
+    let start = std::time::Instant::now();
+    println!("{}...", name);
+    let result = f()?;
+    let elapsed = start.elapsed().as_secs_f64();
+    println!("  -> {} completed in {:.2} sec", name, elapsed);
+    Ok(result)
+}
+
+pub fn sanitize_filename(text: &str, max_len: usize) -> String {
+    let text = if text.len() > max_len {
+        &text[..max_len]
+    } else {
+        text
+    };
+
+    text.chars()
+        .map(|c| {
+            if c.is_ascii_alphanumeric() {
+                c
+            } else {
+                '_'
+            }
+        })
+        .collect()
+}
+
+// ============================================================================
+// ONNX Runtime Integration
+// ============================================================================
+
+use ort::{
+    session::Session,
+    value::Value,
+};
+
+pub struct Style {
+    pub ttl: Array3<f32>,
+    pub dp: Array3<f32>,
+}
+
+pub struct TextToSpeech {
+    cfgs: Config,
+    text_processor: UnicodeProcessor,
+    dp_ort: Session,
+    text_enc_ort: Session,
+    vector_est_ort: Session,
+    vocoder_ort: Session,
+    pub sample_rate: i32,
+}
+
+impl TextToSpeech {
+    pub fn new(
+        cfgs: Config,
+        text_processor: UnicodeProcessor,
+        dp_ort: Session,
+        text_enc_ort: Session,
+        vector_est_ort: Session,
+        vocoder_ort: Session,
+    ) -> Self {
+        let sample_rate = cfgs.ae.sample_rate;
+        TextToSpeech {
+            cfgs,
+            text_processor,
+            dp_ort,
+            text_enc_ort,
+            vector_est_ort,
+            vocoder_ort,
+            sample_rate,
+        }
+    }
+
+    pub fn call(
+        &mut self,
+        text_list: &[String],
+        style: &Style,
+        total_step: usize,
+    ) -> Result<(Vec<f32>, Vec<f32>)> {
+        let bsz = text_list.len();
+
+        // Process text
+        let (text_ids, text_mask) = self.text_processor.call(text_list);
+        
+        let text_ids_array = {
+            let text_ids_shape = (bsz, text_ids[0].len());
+            let mut flat = Vec::new();
+            for row in &text_ids {
+                flat.extend_from_slice(row);
+            }
+            Array::from_shape_vec(text_ids_shape, flat)?
+        };
+
+        let text_ids_value = Value::from_array(text_ids_array)?;
+        let text_mask_value = Value::from_array(text_mask.clone())?;
+        let style_dp_value = Value::from_array(style.dp.clone())?;
+
+        // Predict duration
+        let dp_outputs = self.dp_ort.run(ort::inputs!{
+            "text_ids" => &text_ids_value,
+            "style_dp" => &style_dp_value,
+            "text_mask" => &text_mask_value
+        })?;
+
+        let (_, duration_data) = dp_outputs["duration"].try_extract_tensor::<f32>()?;
+        let duration: Vec<f32> = duration_data.to_vec();
+
+        // Encode text
+        let style_ttl_value = Value::from_array(style.ttl.clone())?;
+        let text_enc_outputs = self.text_enc_ort.run(ort::inputs!{
+            "text_ids" => &text_ids_value,
+            "style_ttl" => &style_ttl_value,
+            "text_mask" => &text_mask_value
+        })?;
+
+        let (text_emb_shape, text_emb_data) = text_enc_outputs["text_emb"].try_extract_tensor::<f32>()?;
+        let text_emb = Array3::from_shape_vec(
+            (text_emb_shape[0] as usize, text_emb_shape[1] as usize, text_emb_shape[2] as usize),
+            text_emb_data.to_vec()
+        )?;
+
+        // Sample noisy latent
+        let (mut xt, latent_mask) = sample_noisy_latent(
+            &duration,
+            self.sample_rate,
+            self.cfgs.ae.base_chunk_size,
+            self.cfgs.ttl.chunk_compress_factor,
+            self.cfgs.ttl.latent_dim,
+        );
+
+        // Prepare constant arrays
+        let total_step_array = Array::from_elem(bsz, total_step as f32);
+
+        // Denoising loop
+        for step in 0..total_step {
+            let current_step_array = Array::from_elem(bsz, step as f32);
+
+            let xt_value = Value::from_array(xt.clone())?;
+            let text_emb_value = Value::from_array(text_emb.clone())?;
+            let latent_mask_value = Value::from_array(latent_mask.clone())?;
+            let text_mask_value2 = Value::from_array(text_mask.clone())?;
+            let current_step_value = Value::from_array(current_step_array)?;
+            let total_step_value = Value::from_array(total_step_array.clone())?;
+
+            let vector_est_outputs = self.vector_est_ort.run(ort::inputs!{
+                "noisy_latent" => &xt_value,
+                "text_emb" => &text_emb_value,
+                "style_ttl" => &style_ttl_value,
+                "latent_mask" => &latent_mask_value,
+                "text_mask" => &text_mask_value2,
+                "current_step" => &current_step_value,
+                "total_step" => &total_step_value
+            })?;
+
+            let (denoised_shape, denoised_data) = vector_est_outputs["denoised_latent"].try_extract_tensor::<f32>()?;
+            xt = Array3::from_shape_vec(
+                (denoised_shape[0] as usize, denoised_shape[1] as usize, denoised_shape[2] as usize),
+                denoised_data.to_vec()
+            )?;
+        }
+
+        // Generate waveform
+        let final_latent_value = Value::from_array(xt)?;
+        let vocoder_outputs = self.vocoder_ort.run(ort::inputs!{
+            "latent" => &final_latent_value
+        })?;
+
+        let (_, wav_data) = vocoder_outputs["wav_tts"].try_extract_tensor::<f32>()?;
+        let wav: Vec<f32> = wav_data.to_vec();
+
+        Ok((wav, duration))
+    }
+}
+
+// ============================================================================
+// Component Loading Functions
+// ============================================================================
+
+/// Load voice style from JSON files
+pub fn load_voice_style(voice_style_paths: &[String], verbose: bool) -> Result<Style> {
+    let bsz = voice_style_paths.len();
+
+    // Read first file to get dimensions
+    let first_file = File::open(&voice_style_paths[0])
+        .context("Failed to open voice style file")?;
+    let first_reader = BufReader::new(first_file);
+    let first_data: VoiceStyleData = serde_json::from_reader(first_reader)?;
+
+    let ttl_dims = &first_data.style_ttl.dims;
+    let dp_dims = &first_data.style_dp.dims;
+
+    let ttl_dim1 = ttl_dims[1];
+    let ttl_dim2 = ttl_dims[2];
+    let dp_dim1 = dp_dims[1];
+    let dp_dim2 = dp_dims[2];
+
+    // Pre-allocate arrays with full batch size
+    let ttl_size = bsz * ttl_dim1 * ttl_dim2;
+    let dp_size = bsz * dp_dim1 * dp_dim2;
+    let mut ttl_flat = vec![0.0f32; ttl_size];
+    let mut dp_flat = vec![0.0f32; dp_size];
+
+    // Fill in the data
+    for (i, path) in voice_style_paths.iter().enumerate() {
+        let file = File::open(path).context("Failed to open voice style file")?;
+        let reader = BufReader::new(file);
+        let data: VoiceStyleData = serde_json::from_reader(reader)?;
+
+        // Flatten TTL data
+        let ttl_offset = i * ttl_dim1 * ttl_dim2;
+        let mut idx = 0;
+        for batch in &data.style_ttl.data {
+            for row in batch {
+                for &val in row {
+                    ttl_flat[ttl_offset + idx] = val;
+                    idx += 1;
+                }
+            }
+        }
+
+        // Flatten DP data
+        let dp_offset = i * dp_dim1 * dp_dim2;
+        idx = 0;
+        for batch in &data.style_dp.data {
+            for row in batch {
+                for &val in row {
+                    dp_flat[dp_offset + idx] = val;
+                    idx += 1;
+                }
+            }
+        }
+    }
+
+    let ttl_style = Array3::from_shape_vec((bsz, ttl_dim1, ttl_dim2), ttl_flat)?;
+    let dp_style = Array3::from_shape_vec((bsz, dp_dim1, dp_dim2), dp_flat)?;
+
+    if verbose {
+        println!("Loaded {} voice styles\n", bsz);
+    }
+
+    Ok(Style {
+        ttl: ttl_style,
+        dp: dp_style,
+    })
+}
+
+/// Load TTS components
+pub fn load_text_to_speech(onnx_dir: &str, use_gpu: bool) -> Result<TextToSpeech> {
+    if use_gpu {
+        anyhow::bail!("GPU mode is not supported yet");
+    }
+    println!("Using CPU for inference\n");
+
+    let cfgs = load_cfgs(onnx_dir)?;
+
+    let dp_path = format!("{}/duration_predictor.onnx", onnx_dir);
+    let text_enc_path = format!("{}/text_encoder.onnx", onnx_dir);
+    let vector_est_path = format!("{}/vector_estimator.onnx", onnx_dir);
+    let vocoder_path = format!("{}/vocoder.onnx", onnx_dir);
+
+    let dp_ort = Session::builder()?
+        .commit_from_file(&dp_path)?;
+    let text_enc_ort = Session::builder()?
+        .commit_from_file(&text_enc_path)?;
+    let vector_est_ort = Session::builder()?
+        .commit_from_file(&vector_est_path)?;
+    let vocoder_ort = Session::builder()?
+        .commit_from_file(&vocoder_path)?;
+
+    let unicode_indexer_path = format!("{}/unicode_indexer.json", onnx_dir);
+    let text_processor = UnicodeProcessor::new(&unicode_indexer_path)?;
+
+    Ok(TextToSpeech::new(
+        cfgs,
+        text_processor,
+        dp_ort,
+        text_enc_ort,
+        vector_est_ort,
+        vocoder_ort,
+    ))
+}