init
This commit is contained in:
21
rust/.gitignore
vendored
Normal file
21
rust/.gitignore
vendored
Normal file
@@ -0,0 +1,21 @@
|
||||
# Rust build artifacts
|
||||
/target/
|
||||
Cargo.lock
|
||||
|
||||
# Output directory
|
||||
/results/
|
||||
|
||||
# IDE
|
||||
.vscode/
|
||||
.idea/
|
||||
*.swp
|
||||
*.swo
|
||||
*~
|
||||
|
||||
# OS
|
||||
.DS_Store
|
||||
Thumbs.db
|
||||
|
||||
# Debug
|
||||
*.pdb
|
||||
|
||||
41
rust/Cargo.toml
Normal file
41
rust/Cargo.toml
Normal file
@@ -0,0 +1,41 @@
|
||||
[package]
|
||||
name = "supertonic-tts"
|
||||
version = "0.1.0"
|
||||
edition = "2021"
|
||||
|
||||
[dependencies]
|
||||
# ONNX Runtime
|
||||
ort = "2.0.0-rc.7"
|
||||
|
||||
# Array processing (like NumPy)
|
||||
ndarray = { version = "0.16", features = ["rayon"] }
|
||||
rand = "0.8"
|
||||
rand_distr = "0.4"
|
||||
|
||||
# Parallel processing
|
||||
rayon = "1.10"
|
||||
|
||||
# Audio processing
|
||||
hound = "3.5"
|
||||
rustfft = "6.2"
|
||||
|
||||
# JSON serialization
|
||||
serde = { version = "1.0", features = ["derive"] }
|
||||
serde_json = "1.0"
|
||||
|
||||
# CLI argument parsing
|
||||
clap = { version = "4.5", features = ["derive"] }
|
||||
|
||||
# Error handling
|
||||
anyhow = "1.0"
|
||||
|
||||
# Unicode normalization
|
||||
unicode-normalization = "0.1"
|
||||
|
||||
# System calls
|
||||
libc = "0.2"
|
||||
|
||||
[[bin]]
|
||||
name = "example_onnx"
|
||||
path = "src/example_onnx.rs"
|
||||
|
||||
101
rust/README.md
Normal file
101
rust/README.md
Normal file
@@ -0,0 +1,101 @@
|
||||
# TTS ONNX Inference Examples
|
||||
|
||||
This guide provides examples for running TTS inference using Rust.
|
||||
|
||||
## Installation
|
||||
|
||||
This project uses [Cargo](https://doc.rust-lang.org/cargo/) for package management.
|
||||
|
||||
### Install Rust (if not already installed)
|
||||
```bash
|
||||
curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh
|
||||
```
|
||||
|
||||
### Build the project
|
||||
```bash
|
||||
cargo build --release
|
||||
```
|
||||
|
||||
## Basic Usage
|
||||
|
||||
You can run the inference in two ways:
|
||||
1. **Using cargo run** (builds if needed, then runs)
|
||||
2. **Direct binary execution** (faster if already built)
|
||||
|
||||
### Example 1: Default Inference
|
||||
Run inference with default settings:
|
||||
```bash
|
||||
# Using cargo run
|
||||
cargo run --release --bin example_onnx
|
||||
|
||||
# Or directly execute the built binary (faster)
|
||||
./target/release/example_onnx
|
||||
```
|
||||
|
||||
This will use:
|
||||
- Voice style: `assets/voice_styles/M1.json`
|
||||
- Text: "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."
|
||||
- Output directory: `results/`
|
||||
- Total steps: 5
|
||||
- Number of generations: 4
|
||||
|
||||
### Example 2: Batch Inference
|
||||
Process multiple voice styles and texts at once:
|
||||
```bash
|
||||
# Using cargo run
|
||||
cargo run --release --bin example_onnx -- \
|
||||
--voice-style assets/voice_styles/M1.json,assets/voice_styles/F1.json \
|
||||
--text "The sun sets behind the mountains, painting the sky in shades of pink and orange.|The weather is beautiful and sunny outside. A gentle breeze makes the air feel fresh and pleasant."
|
||||
|
||||
# Or using the binary directly
|
||||
./target/release/example_onnx \
|
||||
--voice-style assets/voice_styles/M1.json,assets/voice_styles/F1.json \
|
||||
--text "The sun sets behind the mountains, painting the sky in shades of pink and orange.|The weather is beautiful and sunny outside. A gentle breeze makes the air feel fresh and pleasant."
|
||||
```
|
||||
|
||||
This will:
|
||||
- Generate speech for 2 different voice-text pairs
|
||||
- Use male voice (M1.json) for the first text
|
||||
- Use female voice (F1.json) for the second text
|
||||
- Process both samples in a single batch
|
||||
|
||||
### Example 3: High Quality Inference
|
||||
Increase denoising steps for better quality:
|
||||
```bash
|
||||
# Using cargo run
|
||||
cargo run --release --bin example_onnx -- \
|
||||
--total-step 10 \
|
||||
--voice-style assets/voice_styles/M1.json \
|
||||
--text "Increasing the number of denoising steps improves the output's fidelity and overall quality."
|
||||
|
||||
# Or using the binary directly
|
||||
./target/release/example_onnx \
|
||||
--total-step 10 \
|
||||
--voice-style assets/voice_styles/M1.json \
|
||||
--text "Increasing the number of denoising steps improves the output's fidelity and overall quality."
|
||||
```
|
||||
|
||||
This will:
|
||||
- Use 10 denoising steps instead of the default 5
|
||||
- Produce higher quality output at the cost of slower inference
|
||||
|
||||
## Available Arguments
|
||||
|
||||
| Argument | Type | Default | Description |
|
||||
|----------|------|---------|-------------|
|
||||
| `--use-gpu` | flag | False | Use GPU for inference (default: CPU) |
|
||||
| `--onnx-dir` | str | `assets/onnx` | Path to ONNX model directory |
|
||||
| `--total-step` | int | 5 | Number of denoising steps (higher = better quality, slower) |
|
||||
| `--n-test` | int | 4 | Number of times to generate each sample |
|
||||
| `--voice-style` | str+ | `assets/voice_styles/M1.json` | Voice style file path(s) |
|
||||
| `--text` | str+ | (long default text) | Text(s) to synthesize |
|
||||
| `--save-dir` | str | `results` | Output directory |
|
||||
|
||||
## Notes
|
||||
|
||||
- **Batch Processing**: The number of `--voice-style` files must match the number of `--text` entries
|
||||
- **Quality vs Speed**: Higher `--total-step` values produce better quality but take longer
|
||||
- **GPU Support**: GPU mode is not supported yet
|
||||
- **Known Issues**: On some platforms (especially macOS), there might be a mutex cleanup warning during exit. This is a known ONNX Runtime issue and doesn't affect functionality. The implementation uses `libc::_exit()` and `mem::forget()` to bypass this issue.
|
||||
|
||||
|
||||
1
rust/assets
Symbolic link
1
rust/assets
Symbolic link
@@ -0,0 +1 @@
|
||||
../assets
|
||||
108
rust/src/example_onnx.rs
Normal file
108
rust/src/example_onnx.rs
Normal file
@@ -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);
|
||||
}
|
||||
}
|
||||
507
rust/src/helper.rs
Normal file
507
rust/src/helper.rs
Normal file
@@ -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" => ¤t_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,
|
||||
))
|
||||
}
|
||||
Reference in New Issue
Block a user