Merge pull request #62 from Googlefan256/main

WIP wasm support

sbv2_core/src/error.rs

@@ -6,6 +6,7 @@ pub enum Error {
     TokenizerError(#[from] tokenizers::Error),
     #[error("JPreprocess error: {0}")]
     JPreprocessError(#[from] jpreprocess::error::JPreprocessError),
+    #[cfg(feature = "std")]
     #[error("ONNX error: {0}")]
     OrtError(#[from] ort::Error),
     #[error("NDArray error: {0}")]
@@ -20,6 +21,8 @@ pub enum Error {
     HoundError(#[from] hound::Error),
     #[error("model not found error")]
     ModelNotFoundError(String),
+    #[error("other")]
+    OtherError(String),
 }
 
 pub type Result<T> = std::result::Result<T, Error>;

sbv2_core/src/lib.rs

@@ -1,11 +1,16 @@
+#[cfg(feature = "std")]
 pub mod bert;
 pub mod error;
 pub mod jtalk;
+#[cfg(feature = "std")]
 pub mod model;
 pub mod mora;
 pub mod nlp;
 pub mod norm;
+pub mod sbv2file;
 pub mod style;
 pub mod tokenizer;
+#[cfg(feature = "std")]
 pub mod tts;
+pub mod tts_util;
 pub mod utils;

sbv2_core/src/main.rs

@@ -1,9 +1,9 @@
+use std::env;
 use std::fs;
 
-use sbv2_core::tts;
-use std::env;
-
-fn main() -> anyhow::Result<()> {
+#[cfg(feature = "std")]
+fn main_inner() -> anyhow::Result<()> {
+    use sbv2_core::tts;
     dotenvy::dotenv_override().ok();
     env_logger::init();
     let text = fs::read_to_string("content.txt")?;
@@ -19,3 +19,13 @@ fn main() -> anyhow::Result<()> {
 
     Ok(())
 }
+#[cfg(not(feature = "std"))]
+fn main_inner() -> anyhow::Result<()> {
+    Ok(())
+}
+
+fn main() {
+    if let Err(e) = main_inner() {
+        println!("Error: {e}");
+    }
+}

sbv2_core/src/sbv2file.rs

@@ -0,0 +1,37 @@
+use std::io::{Cursor, Read};
+
+use tar::Archive;
+use zstd::decode_all;
+
+use crate::error::{Error, Result};
+
+/// Parse a .sbv2 file binary
+///
+/// # Examples
+///
+/// ```rs
+/// parse_sbv2file("tsukuyomi", std::fs::read("tsukuyomi.sbv2")?)?;
+/// ```
+pub fn parse_sbv2file<P: AsRef<[u8]>>(sbv2_bytes: P) -> Result<(Vec<u8>, Vec<u8>)> {
+    let mut arc = Archive::new(Cursor::new(decode_all(Cursor::new(sbv2_bytes.as_ref()))?));
+    let mut vits2 = None;
+    let mut style_vectors = None;
+    let mut et = arc.entries()?;
+    while let Some(Ok(mut e)) = et.next() {
+        let pth = String::from_utf8_lossy(&e.path_bytes()).to_string();
+        let mut b = Vec::with_capacity(e.size() as usize);
+        e.read_to_end(&mut b)?;
+        match pth.as_str() {
+            "model.onnx" => vits2 = Some(b),
+            "style_vectors.json" => style_vectors = Some(b),
+            _ => continue,
+        }
+    }
+    if style_vectors.is_none() {
+        return Err(Error::ModelNotFoundError("style_vectors".to_string()));
+    }
+    if vits2.is_none() {
+        return Err(Error::ModelNotFoundError("vits2".to_string()));
+    }
+    Ok((style_vectors.unwrap(), vits2.unwrap()))
+}

sbv2_core/src/tokenizer.rs

@@ -1,5 +1,5 @@
 use crate::error::Result;
-use tokenizers::Tokenizer;
+pub use tokenizers::Tokenizer;
 
 pub fn get_tokenizer<P: AsRef<[u8]>>(p: P) -> Result<Tokenizer> {
     let tokenizer = Tokenizer::from_bytes(p)?;

sbv2_core/src/tts.rs

@@ -1,12 +1,8 @@
 use crate::error::{Error, Result};
-use crate::{bert, jtalk, model, nlp, norm, style, tokenizer, utils};
-use hound::{SampleFormat, WavSpec, WavWriter};
-use ndarray::{concatenate, s, Array, Array1, Array2, Array3, Axis};
+use crate::{jtalk, model, style, tokenizer, tts_util};
+use ndarray::{concatenate, Array1, Array2, Array3, Axis};
 use ort::Session;
-use std::io::{Cursor, Read};
-use tar::Archive;
 use tokenizers::Tokenizer;
-use zstd::decode_all;
 
 #[derive(PartialEq, Eq, Clone)]
 pub struct TTSIdent(String);
@@ -78,27 +74,8 @@ impl TTSModelHolder {
         ident: I,
         sbv2_bytes: P,
     ) -> Result<()> {
-        let mut arc = Archive::new(Cursor::new(decode_all(Cursor::new(sbv2_bytes.as_ref()))?));
-        let mut vits2 = None;
-        let mut style_vectors = None;
-        let mut et = arc.entries()?;
-        while let Some(Ok(mut e)) = et.next() {
-            let pth = String::from_utf8_lossy(&e.path_bytes()).to_string();
-            let mut b = Vec::with_capacity(e.size() as usize);
-            e.read_to_end(&mut b)?;
-            match pth.as_str() {
-                "model.onnx" => vits2 = Some(b),
-                "style_vectors.json" => style_vectors = Some(b),
-                _ => continue,
-            }
-        }
-        if style_vectors.is_none() {
-            return Err(Error::ModelNotFoundError("style_vectors".to_string()));
-        }
-        if vits2.is_none() {
-            return Err(Error::ModelNotFoundError("vits2".to_string()));
-        }
-        self.load(ident, style_vectors.unwrap(), vits2.unwrap())?;
+        let (style_vectors, vits2) = crate::sbv2file::parse_sbv2file(sbv2_bytes)?;
+        self.load(ident, style_vectors, vits2)?;
         Ok(())
     }
 
@@ -151,69 +128,14 @@ impl TTSModelHolder {
         &self,
         text: &str,
     ) -> Result<(Array2<f32>, Array1<i64>, Array1<i64>, Array1<i64>)> {
-        let text = self.jtalk.num2word(text)?;
-        let normalized_text = norm::normalize_text(&text);
-
-        let process = self.jtalk.process_text(&normalized_text)?;
-        let (phones, tones, mut word2ph) = process.g2p()?;
-        let (phones, tones, lang_ids) = nlp::cleaned_text_to_sequence(phones, tones);
-
-        let phones = utils::intersperse(&phones, 0);
-        let tones = utils::intersperse(&tones, 0);
-        let lang_ids = utils::intersperse(&lang_ids, 0);
-        for item in &mut word2ph {
-            *item *= 2;
-        }
-        word2ph[0] += 1;
-
-        let text = {
-            let (seq_text, _) = process.text_to_seq_kata()?;
-            seq_text.join("")
-        };
-        let (token_ids, attention_masks) = tokenizer::tokenize(&text, &self.tokenizer)?;
-
-        let bert_content = bert::predict(&self.bert, token_ids, attention_masks)?;
-
-        assert!(
-            word2ph.len() == text.chars().count() + 2,
-            "{} {}",
-            word2ph.len(),
-            normalized_text.chars().count()
-        );
-
-        let mut phone_level_feature = vec![];
-        for (i, reps) in word2ph.iter().enumerate() {
-            let repeat_feature = {
-                let (reps_rows, reps_cols) = (*reps, 1);
-                let arr_len = bert_content.slice(s![i, ..]).len();
-
-                let mut results: Array2<f32> =
-                    Array::zeros((reps_rows as usize, arr_len * reps_cols));
-
-                for j in 0..reps_rows {
-                    for k in 0..reps_cols {
-                        let mut view = results.slice_mut(s![j, k * arr_len..(k + 1) * arr_len]);
-                        view.assign(&bert_content.slice(s![i, ..]));
-                    }
-                }
-                results
-            };
-            phone_level_feature.push(repeat_feature);
-        }
-        let phone_level_feature = concatenate(
-            Axis(0),
-            &phone_level_feature
-                .iter()
-                .map(|x| x.view())
-                .collect::<Vec<_>>(),
-        )?;
-        let bert_ori = phone_level_feature.t();
-        Ok((
-            bert_ori.to_owned(),
-            phones.into(),
-            tones.into(),
-            lang_ids.into(),
-        ))
+        crate::tts_util::parse_text(
+            text,
+            &self.jtalk,
+            &self.tokenizer,
+            |token_ids, attention_masks| {
+                crate::bert::predict(&self.bert, token_ids, attention_masks)
+            },
+        )
     }
 
     fn find_model<I: Into<TTSIdent>>(&self, ident: I) -> Result<&TTSModel> {
@@ -292,26 +214,7 @@ impl TTSModelHolder {
                 options.length_scale,
             )?
         };
-        Self::array_to_vec(audio_array)
-    }
-
-    fn array_to_vec(audio_array: Array3<f32>) -> Result<Vec<u8>> {
-        let spec = WavSpec {
-            channels: 1,
-            sample_rate: 44100,
-            bits_per_sample: 32,
-            sample_format: SampleFormat::Float,
-        };
-        let mut cursor = Cursor::new(Vec::new());
-        let mut writer = WavWriter::new(&mut cursor, spec)?;
-        for i in 0..audio_array.shape()[0] {
-            let output = audio_array.slice(s![i, 0, ..]).to_vec();
-            for sample in output {
-                writer.write_sample(sample)?;
-            }
-        }
-        writer.finalize()?;
-        Ok(cursor.into_inner())
+        tts_util::array_to_vec(audio_array)
     }
 
     /// Synthesize text to audio
@@ -340,7 +243,7 @@ impl TTSModelHolder {
             sdp_ratio,
             length_scale,
         )?;
-        Self::array_to_vec(audio_array)
+        tts_util::array_to_vec(audio_array)
     }
 }
 

sbv2_core/src/tts_util.rs

@@ -0,0 +1,100 @@
+use std::io::Cursor;
+
+use crate::error::Result;
+use crate::{jtalk, nlp, norm, tokenizer, utils};
+use hound::{SampleFormat, WavSpec, WavWriter};
+use ndarray::{concatenate, s, Array, Array1, Array2, Array3, Axis};
+use tokenizers::Tokenizer;
+/// Parse text and return the input for synthesize
+///
+/// # Note
+/// This function is for low-level usage, use `easy_synthesize` for high-level usage.
+#[allow(clippy::type_complexity)]
+pub fn parse_text(
+    text: &str,
+    jtalk: &jtalk::JTalk,
+    tokenizer: &Tokenizer,
+    bert_predict: impl FnOnce(Vec<i64>, Vec<i64>) -> Result<ndarray::Array2<f32>>,
+) -> Result<(Array2<f32>, Array1<i64>, Array1<i64>, Array1<i64>)> {
+    let text = jtalk.num2word(text)?;
+    let normalized_text = norm::normalize_text(&text);
+
+    let process = jtalk.process_text(&normalized_text)?;
+    let (phones, tones, mut word2ph) = process.g2p()?;
+    let (phones, tones, lang_ids) = nlp::cleaned_text_to_sequence(phones, tones);
+
+    let phones = utils::intersperse(&phones, 0);
+    let tones = utils::intersperse(&tones, 0);
+    let lang_ids = utils::intersperse(&lang_ids, 0);
+    for item in &mut word2ph {
+        *item *= 2;
+    }
+    word2ph[0] += 1;
+
+    let text = {
+        let (seq_text, _) = process.text_to_seq_kata()?;
+        seq_text.join("")
+    };
+    let (token_ids, attention_masks) = tokenizer::tokenize(&text, tokenizer)?;
+
+    let bert_content = bert_predict(token_ids, attention_masks)?;
+
+    assert!(
+        word2ph.len() == text.chars().count() + 2,
+        "{} {}",
+        word2ph.len(),
+        normalized_text.chars().count()
+    );
+
+    let mut phone_level_feature = vec![];
+    for (i, reps) in word2ph.iter().enumerate() {
+        let repeat_feature = {
+            let (reps_rows, reps_cols) = (*reps, 1);
+            let arr_len = bert_content.slice(s![i, ..]).len();
+
+            let mut results: Array2<f32> = Array::zeros((reps_rows as usize, arr_len * reps_cols));
+
+            for j in 0..reps_rows {
+                for k in 0..reps_cols {
+                    let mut view = results.slice_mut(s![j, k * arr_len..(k + 1) * arr_len]);
+                    view.assign(&bert_content.slice(s![i, ..]));
+                }
+            }
+            results
+        };
+        phone_level_feature.push(repeat_feature);
+    }
+    let phone_level_feature = concatenate(
+        Axis(0),
+        &phone_level_feature
+            .iter()
+            .map(|x| x.view())
+            .collect::<Vec<_>>(),
+    )?;
+    let bert_ori = phone_level_feature.t();
+    Ok((
+        bert_ori.to_owned(),
+        phones.into(),
+        tones.into(),
+        lang_ids.into(),
+    ))
+}
+
+pub fn array_to_vec(audio_array: Array3<f32>) -> Result<Vec<u8>> {
+    let spec = WavSpec {
+        channels: 1,
+        sample_rate: 44100,
+        bits_per_sample: 32,
+        sample_format: SampleFormat::Float,
+    };
+    let mut cursor = Cursor::new(Vec::new());
+    let mut writer = WavWriter::new(&mut cursor, spec)?;
+    for i in 0..audio_array.shape()[0] {
+        let output = audio_array.slice(s![i, 0, ..]).to_vec();
+        for sample in output {
+            writer.write_sample(sample)?;
+        }
+    }
+    writer.finalize()?;
+    Ok(cursor.into_inner())
+}