Code-Inc/supertonic
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

add text chunking for long-form generation (Fixes #4)

Browse Source
This commit is contained in:
ANLGBOY
2025-11-19 18:08:30 +09:00
parent d31536d9fc
commit c31b6745e4
30 changed files with 1813 additions and 102 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
rust/Cargo.toml
View File

@@ -32,6 +32,9 @@ anyhow = "1.0"
# Unicode normalization
unicode-normalization = "0.1"
# Regular expressions
regex = "1.10"
# System calls
libc = "0.2"

33
rust/README.md
View File

@@ -2,6 +2,10 @@
This guide provides examples for running TTS inference using Rust.
## 📰 Update News
**2025.11.19** - Added automatic text chunking for long-form inference. Long texts are split into chunks and synthesized with natural pauses.
## Installation
This project uses [Cargo](https://doc.rust-lang.org/cargo/) for package management.
@@ -44,11 +48,13 @@ Process multiple voice styles and texts at once:
```bash
# Using cargo run
cargo run --release --bin example_onnx -- \
--batch \
--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 \
--batch \
--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."
```
@@ -79,6 +85,29 @@ This will:
- Use 10 denoising steps instead of the default 5
- Produce higher quality output at the cost of slower inference
### Example 4: Long-Form Inference
The system automatically chunks long texts into manageable segments, synthesizes each segment separately, and concatenates them with natural pauses (0.3 seconds by default) into a single audio file. This happens by default when you don't use the `--batch` flag:
```bash
# Using cargo run
cargo run --release --bin example_onnx -- \
--voice-style assets/voice_styles/M1.json \
--text "This is a very long text that will be automatically split into multiple chunks. The system will process each chunk separately and then concatenate them together with natural pauses between segments. This ensures that even very long texts can be processed efficiently while maintaining natural speech flow and avoiding memory issues."
# Or using the binary directly
./target/release/example_onnx \
--voice-style assets/voice_styles/M1.json \
--text "This is a very long text that will be automatically split into multiple chunks. The system will process each chunk separately and then concatenate them together with natural pauses between segments. This ensures that even very long texts can be processed efficiently while maintaining natural speech flow and avoiding memory issues."
```
This will:
- Automatically split the text into chunks based on paragraph and sentence boundaries
- Synthesize each chunk separately
- Add 0.3 seconds of silence between chunks for natural pauses
- Concatenate all chunks into a single audio file
**Note**: Automatic text chunking is disabled when using `--batch` mode. In batch mode, each text is processed as-is without chunking.
## Available Arguments
| Argument | Type | Default | Description |
@@ -90,10 +119,12 @@ This will:
| `--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 |
| `--batch` | flag | False | Enable batch mode (multiple text-style pairs, disables automatic chunking) |
## Notes
- **Batch Processing**: The number of `--voice-style` files must match the number of `--text` entries
- **Batch Processing**: When using `--batch`, the number of `--voice-style` files must match the number of `--text` entries
- **Automatic Chunking**: Without `--batch`, long texts are automatically split and concatenated with 0.3s pauses
- **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.

49
rust/src/example_onnx.rs
View File

@@ -41,6 +41,10 @@ struct Args {
/// Output directory
#[arg(long, default_value = "results")]
save_dir: String,
/// Enable batch mode (multiple text-style pairs)
#[arg(long, default_value = "false")]
batch: bool,
}
fn main() -> Result<()> {
@@ -53,13 +57,16 @@ fn main() -> Result<()> {
let voice_style_paths = &args.voice_style;
let text_list = &args.text;
let save_dir = &args.save_dir;
let batch = args.batch;
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()
);
if batch {
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();
@@ -76,19 +83,31 @@ fn main() -> Result<()> {
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)
})?;
let (wav, duration) = if batch {
timer("Generating speech from text", || {
text_to_speech.batch(text_list, &style, total_step)
})?
} else {
let (w, d) = timer("Generating speech from text", || {
text_to_speech.call(&text_list[0], &style, total_step, 0.3)
})?;
(w, vec![d])
};
// 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 wav_slice = if batch {
let wav_len = wav.len() / bsz;
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);
&wav[wav_start..wav_end]
} else {
// For non-batch mode, wav is a single concatenated audio
let actual_len = (text_to_speech.sample_rate as f32 * duration[0]) as usize;
&wav[..actual_len.min(wav.len())]
};
let output_path = PathBuf::from(save_dir).join(&fname);
write_wav_file(&output_path, wav_slice, text_to_speech.sample_rate)?;

228
rust/src/helper.rs
View File

@@ -12,6 +12,7 @@ use anyhow::{Result, Context};
use unicode_normalization::UnicodeNormalization;
use hound::{WavWriter, WavSpec, SampleFormat};
use rand_distr::{Distribution, Normal};
use regex::Regex;
// ============================================================================
// Configuration Structures
@@ -218,6 +219,187 @@ pub fn write_wav_file<P: AsRef<Path>>(
Ok(())
}
// ============================================================================
// Text Chunking
// ============================================================================
const MAX_CHUNK_LENGTH: usize = 300;
const ABBREVIATIONS: &[&str] = &[
"Dr.", "Mr.", "Mrs.", "Ms.", "Prof.", "Sr.", "Jr.",
"St.", "Ave.", "Rd.", "Blvd.", "Dept.", "Inc.", "Ltd.",
"Co.", "Corp.", "etc.", "vs.", "i.e.", "e.g.", "Ph.D.",
];
pub fn chunk_text(text: &str, max_len: Option<usize>) -> Vec<String> {
let max_len = max_len.unwrap_or(MAX_CHUNK_LENGTH);
let text = text.trim();
if text.is_empty() {
return vec![String::new()];
}
// Split by paragraphs
let para_re = Regex::new(r"\n\s*\n").unwrap();
let paragraphs: Vec<&str> = para_re.split(text).collect();
let mut chunks = Vec::new();
for para in paragraphs {
let para = para.trim();
if para.is_empty() {
continue;
}
if para.len() <= max_len {
chunks.push(para.to_string());
continue;
}
// Split by sentences
let sentences = split_sentences(para);
let mut current = String::new();
let mut current_len = 0;
for sentence in sentences {
let sentence = sentence.trim();
if sentence.is_empty() {
continue;
}
let sentence_len = sentence.len();
if sentence_len > max_len {
// If sentence is longer than max_len, split by comma or space
if !current.is_empty() {
chunks.push(current.trim().to_string());
current.clear();
current_len = 0;
}
// Try splitting by comma
let parts: Vec<&str> = sentence.split(',').collect();
for part in parts {
let part = part.trim();
if part.is_empty() {
continue;
}
let part_len = part.len();
if part_len > max_len {
// Split by space as last resort
let words: Vec<&str> = part.split_whitespace().collect();
let mut word_chunk = String::new();
let mut word_chunk_len = 0;
for word in words {
let word_len = word.len();
if word_chunk_len + word_len + 1 > max_len && !word_chunk.is_empty() {
chunks.push(word_chunk.trim().to_string());
word_chunk.clear();
word_chunk_len = 0;
}
if !word_chunk.is_empty() {
word_chunk.push(' ');
word_chunk_len += 1;
}
word_chunk.push_str(word);
word_chunk_len += word_len;
}
if !word_chunk.is_empty() {
chunks.push(word_chunk.trim().to_string());
}
} else {
if current_len + part_len + 1 > max_len && !current.is_empty() {
chunks.push(current.trim().to_string());
current.clear();
current_len = 0;
}
if !current.is_empty() {
current.push_str(", ");
current_len += 2;
}
current.push_str(part);
current_len += part_len;
}
}
continue;
}
if current_len + sentence_len + 1 > max_len && !current.is_empty() {
chunks.push(current.trim().to_string());
current.clear();
current_len = 0;
}
if !current.is_empty() {
current.push(' ');
current_len += 1;
}
current.push_str(sentence);
current_len += sentence_len;
}
if !current.is_empty() {
chunks.push(current.trim().to_string());
}
}
if chunks.is_empty() {
vec![String::new()]
} else {
chunks
}
}
fn split_sentences(text: &str) -> Vec<String> {
// Rust's regex doesn't support lookbehind, so we use a simpler approach
// Split on sentence boundaries and then check if they're abbreviations
let re = Regex::new(r"([.!?])\s+").unwrap();
// Find all matches
let matches: Vec<_> = re.find_iter(text).collect();
if matches.is_empty() {
return vec![text.to_string()];
}
let mut sentences = Vec::new();
let mut last_end = 0;
for m in matches {
// Get the text before the punctuation
let before_punc = &text[last_end..m.start()];
// Check if this ends with an abbreviation
let mut is_abbrev = false;
for abbrev in ABBREVIATIONS {
let combined = format!("{}{}", before_punc.trim(), &text[m.start()..m.start()+1]);
if combined.ends_with(abbrev) {
is_abbrev = true;
break;
}
}
if !is_abbrev {
// This is a real sentence boundary
sentences.push(text[last_end..m.end()].to_string());
last_end = m.end();
}
}
// Add the remaining text
if last_end < text.len() {
sentences.push(text[last_end..].to_string());
}
if sentences.is_empty() {
vec![text.to_string()]
} else {
sentences
}
}
// ============================================================================
// Utility Functions
// ============================================================================
@@ -297,7 +479,7 @@ impl TextToSpeech {
}
}
pub fn call(
fn _infer(
&mut self,
text_list: &[String],
style: &Style,
@@ -396,6 +578,50 @@ impl TextToSpeech {
Ok((wav, duration))
}
pub fn call(
&mut self,
text: &str,
style: &Style,
total_step: usize,
silence_duration: f32,
) -> Result<(Vec<f32>, f32)> {
let chunks = chunk_text(text, None);
let mut wav_cat: Vec<f32> = Vec::new();
let mut dur_cat: f32 = 0.0;
for (i, chunk) in chunks.iter().enumerate() {
let (wav, duration) = self._infer(&[chunk.clone()], style, total_step)?;
let dur = duration[0];
let wav_len = (self.sample_rate as f32 * dur) as usize;
let wav_chunk = &wav[..wav_len.min(wav.len())];
if i == 0 {
wav_cat.extend_from_slice(wav_chunk);
dur_cat = dur;
} else {
let silence_len = (silence_duration * self.sample_rate as f32) as usize;
let silence = vec![0.0f32; silence_len];
wav_cat.extend_from_slice(&silence);
wav_cat.extend_from_slice(wav_chunk);
dur_cat += silence_duration + dur;
}
}
Ok((wav_cat, dur_cat))
}
pub fn batch(
&mut self,
text_list: &[String],
style: &Style,
total_step: usize,
) -> Result<(Vec<f32>, Vec<f32>)> {
self._infer(text_list, style, total_step)
}
}
// ============================================================================