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

rust/src/example_onnx.rs

@@ -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)?;

rust/src/helper.rs

@@ -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)
+    }
 }
 
 // ============================================================================