Supertonic 2

rust/src/example_onnx.rs

@@ -42,6 +42,10 @@ struct Args {
     #[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>,
 
+    /// Language(s) for synthesis (en, ko, es, pt, fr)
+    #[arg(long, value_delimiter = ',', default_values_t = vec!["en".to_string()])]
+    lang: Vec<String>,
+
     /// Output directory
     #[arg(long, default_value = "results")]
     save_dir: String,
@@ -61,6 +65,7 @@ fn main() -> Result<()> {
     let n_test = args.n_test;
     let voice_style_paths = &args.voice_style;
     let text_list = &args.text;
+    let lang_list = &args.lang;
     let save_dir = &args.save_dir;
     let batch = args.batch;
 
@@ -72,6 +77,13 @@ fn main() -> Result<()> {
                 text_list.len()
             );
         }
+        if lang_list.len() != text_list.len() {
+            anyhow::bail!(
+                "Number of languages ({}) must match number of texts ({})",
+                lang_list.len(),
+                text_list.len()
+            );
+        }
     }
 
     let bsz = voice_style_paths.len();
@@ -90,11 +102,11 @@ fn main() -> Result<()> {
 
         let (wav, duration) = if batch {
             timer("Generating speech from text", || {
-                text_to_speech.batch(text_list, &style, total_step, speed)
+                text_to_speech.batch(text_list, lang_list, &style, total_step, speed)
             })?
         } else {
             let (w, d) = timer("Generating speech from text", || {
-                text_to_speech.call(&text_list[0], &style, total_step, speed, 0.3)
+                text_to_speech.call(&text_list[0], &lang_list[0], &style, total_step, speed, 0.3)
             })?;
             (w, vec![d])
         };

rust/src/helper.rs

@@ -8,12 +8,19 @@ use serde_json;
 use std::fs::File;
 use std::io::BufReader;
 use std::path::Path;
-use anyhow::{Result, Context};
+use anyhow::{Result, Context, bail};
 use unicode_normalization::UnicodeNormalization;
 use hound::{WavWriter, WavSpec, SampleFormat};
 use rand_distr::{Distribution, Normal};
 use regex::Regex;
 
+// Available languages for multilingual TTS
+pub const AVAILABLE_LANGS: &[&str] = &["en", "ko", "es", "pt", "fr"];
+
+pub fn is_valid_lang(lang: &str) -> bool {
+    AVAILABLE_LANGS.contains(&lang)
+}
+
 // ============================================================================
 // Configuration Structures
 // ============================================================================
@@ -79,11 +86,11 @@ impl UnicodeProcessor {
         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();
+    pub fn call(&self, text_list: &[String], lang_list: &[String]) -> Result<(Vec<Vec<i64>>, Array3<f32>)> {
+        let mut processed_texts: Vec<String> = Vec::new();
+        for (text, lang) in text_list.iter().zip(lang_list.iter()) {
+            processed_texts.push(preprocess_text(text, lang)?);
+        }
 
         let text_ids_lengths: Vec<usize> = processed_texts
             .iter()
@@ -108,16 +115,14 @@ impl UnicodeProcessor {
 
         let text_mask = get_text_mask(&text_ids_lengths);
 
-        (text_ids, text_mask)
+        Ok((text_ids, text_mask))
     }
 }
 
-pub fn preprocess_text(text: &str) -> String {
+pub fn preprocess_text(text: &str, lang: &str) -> Result<String> {
     // TODO: Need advanced normalizer for better performance
     let mut text: String = text.nfkd().collect();
 
-    // FIXME: this should be fixed for non-English languages
-
     // Remove emojis (wide Unicode range)
     let emoji_pattern = Regex::new(r"[\x{1F600}-\x{1F64F}\x{1F300}-\x{1F5FF}\x{1F680}-\x{1F6FF}\x{1F700}-\x{1F77F}\x{1F780}-\x{1F7FF}\x{1F800}-\x{1F8FF}\x{1F900}-\x{1F9FF}\x{1FA00}-\x{1FA6F}\x{1FA70}-\x{1FAFF}\x{2600}-\x{26FF}\x{2700}-\x{27BF}\x{1F1E6}-\x{1F1FF}]+").unwrap();
     text = emoji_pattern.replace_all(&text, "").to_string();
@@ -127,7 +132,6 @@ pub fn preprocess_text(text: &str) -> String {
         ("–", "-"),      // en dash
         ("‑", "-"),      // non-breaking hyphen
         ("—", "-"),      // em dash
-        ("¯", " "),      // macron
         ("_", " "),      // underscore
         ("\u{201C}", "\""),     // left double quote
         ("\u{201D}", "\""),     // right double quote
@@ -148,10 +152,6 @@ pub fn preprocess_text(text: &str) -> String {
         text = text.replace(from, to);
     }
 
-    // Remove combining diacritics // FIXME: this should be fixed for non-English languages
-    let diacritics_pattern = Regex::new(r"[\u{0302}\u{0303}\u{0304}\u{0305}\u{0306}\u{0307}\u{0308}\u{030A}\u{030B}\u{030C}\u{0327}\u{0328}\u{0329}\u{032A}\u{032B}\u{032C}\u{032D}\u{032E}\u{032F}]").unwrap();
-    text = diacritics_pattern.replace_all(&text, "").to_string();
-
     // Remove special symbols
     let special_symbols = ["♥", "☆", "♡", "©", "\\"];
     for symbol in &special_symbols {
@@ -201,7 +201,15 @@ pub fn preprocess_text(text: &str) -> String {
         }
     }
 
-    text
+    // Validate language
+    if !is_valid_lang(lang) {
+        bail!("Invalid language: {}. Available: {:?}", lang, AVAILABLE_LANGS);
+    }
+
+    // Wrap text with language tags
+    text = format!("<{}>{}</{}>", lang, text, lang);
+
+    Ok(text)
 }
 
 pub fn text_to_unicode_values(text: &str) -> Vec<usize> {
@@ -505,15 +513,12 @@ where
 }
 
 pub fn sanitize_filename(text: &str, max_len: usize) -> String {
-    let text = if text.len() > max_len {
-        &text[..max_len]
-    } else {
-        text
-    };
-
+    // Take first max_len characters (Unicode code points, not bytes)
     text.chars()
+        .take(max_len)
         .map(|c| {
-            if c.is_ascii_alphanumeric() {
+            // is_alphanumeric() works with all Unicode letters and digits
+            if c.is_alphanumeric() {
                 c
             } else {
                 '_'
@@ -570,6 +575,7 @@ impl TextToSpeech {
     fn _infer(
         &mut self,
         text_list: &[String],
+        lang_list: &[String],
         style: &Style,
         total_step: usize,
         speed: f32,
@@ -577,7 +583,7 @@ impl TextToSpeech {
         let bsz = text_list.len();
 
         // Process text
-        let (text_ids, text_mask) = self.text_processor.call(text_list);
+        let (text_ids, text_mask) = self.text_processor.call(text_list, lang_list)?;
         
         let text_ids_array = {
             let text_ids_shape = (bsz, text_ids[0].len());
@@ -676,18 +682,20 @@ impl TextToSpeech {
     pub fn call(
         &mut self,
         text: &str,
+        lang: &str,
         style: &Style,
         total_step: usize,
         speed: f32,
         silence_duration: f32,
     ) -> Result<(Vec<f32>, f32)> {
-        let chunks = chunk_text(text, None);
+        let max_len = if lang == "ko" { 120 } else { 300 };
+        let chunks = chunk_text(text, Some(max_len));
         
         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, speed)?;
+            let (wav, duration) = self._infer(&[chunk.clone()], &[lang.to_string()], style, total_step, speed)?;
             
             let dur = duration[0];
             let wav_len = (self.sample_rate as f32 * dur) as usize;
@@ -712,11 +720,12 @@ impl TextToSpeech {
     pub fn batch(
         &mut self,
         text_list: &[String],
+        lang_list: &[String],
         style: &Style,
         total_step: usize,
         speed: f32,
     ) -> Result<(Vec<f32>, Vec<f32>)> {
-        self._infer(text_list, style, total_step, speed)
+        self._infer(text_list, lang_list, style, total_step, speed)
     }
 }