diff --git a/Cargo.lock b/Cargo.lock index 30d95ac..466f787 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1374,6 +1374,27 @@ dependencies = [ "minimal-lexical", ] +[[package]] +name = "npyz" +version = "0.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "13f27ea175875c472b3df61ece89a6d6ef4e0627f43704e400c782f174681ebd" +dependencies = [ + "byteorder", + "num-bigint", + "py_literal", +] + +[[package]] +name = "num-bigint" +version = "0.4.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a5e44f723f1133c9deac646763579fdb3ac745e418f2a7af9cd0c431da1f20b9" +dependencies = [ + "num-integer", + "num-traits", +] + [[package]] name = "num-complex" version = "0.4.6" @@ -1513,6 +1534,51 @@ version = "2.3.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e3148f5046208a5d56bcfc03053e3ca6334e51da8dfb19b6cdc8b306fae3283e" +[[package]] +name = "pest" +version = "2.7.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "879952a81a83930934cbf1786752d6dedc3b1f29e8f8fb2ad1d0a36f377cf442" +dependencies = [ + "memchr", + "thiserror", + "ucd-trie", +] + +[[package]] +name = "pest_derive" +version = "2.7.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d214365f632b123a47fd913301e14c946c61d1c183ee245fa76eb752e59a02dd" +dependencies = [ + "pest", + "pest_generator", +] + +[[package]] +name = "pest_generator" +version = "2.7.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "eb55586734301717aea2ac313f50b2eb8f60d2fc3dc01d190eefa2e625f60c4e" +dependencies = [ + "pest", + "pest_meta", + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "pest_meta" +version = "2.7.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b75da2a70cf4d9cb76833c990ac9cd3923c9a8905a8929789ce347c84564d03d" +dependencies = [ + "once_cell", + "pest", + "sha2", +] + [[package]] name = "phf" version = "0.11.2" @@ -1606,6 +1672,19 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "py_literal" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "102df7a3d46db9d3891f178dcc826dc270a6746277a9ae6436f8d29fd490a8e1" +dependencies = [ + "num-bigint", + "num-complex", + "num-traits", + "pest", + "pest_derive", +] + [[package]] name = "pyo3" version = "0.22.6" @@ -1893,11 +1972,13 @@ name = "sbv2_core" version = "0.2.0-alpha4" dependencies = [ "anyhow", + "base64 0.22.1", "dotenvy", "env_logger", "hound", "jpreprocess", "ndarray", + "npyz", "num_cpus", "once_cell", "ort", @@ -2284,6 +2365,12 @@ version = "1.17.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "42ff0bf0c66b8238c6f3b578df37d0b7848e55df8577b3f74f92a69acceeb825" +[[package]] +name = "ucd-trie" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2896d95c02a80c6d6a5d6e953d479f5ddf2dfdb6a244441010e373ac0fb88971" + [[package]] name = "unicode-bidi" version = "0.3.17" diff --git a/sbv2_core/Cargo.toml b/sbv2_core/Cargo.toml index 4f84986..9d830fb 100644 --- a/sbv2_core/Cargo.toml +++ b/sbv2_core/Cargo.toml @@ -10,11 +10,13 @@ documentation = "https://docs.rs/sbv2_core" [dependencies] anyhow.workspace = true +base64 = { version = "0.22.1", optional = true } dotenvy.workspace = true env_logger.workspace = true hound = "3.5.1" jpreprocess = { version = "0.10.0", features = ["naist-jdic"] } ndarray.workspace = true +npyz = { version = "0.8.3", optional = true } num_cpus = "1.16.0" once_cell.workspace = true ort = { git = "https://github.com/pykeio/ort.git", version = "2.0.0-rc.8", optional = true } @@ -35,4 +37,6 @@ directml = ["ort/directml", "std"] tensorrt = ["ort/tensorrt", "std"] coreml = ["ort/coreml", "std"] default = ["std"] -no_std = ["tokenizers/unstable_wasm"] \ No newline at end of file +no_std = ["tokenizers/unstable_wasm"] +aivmx = ["npyz", "base64"] +base64 = ["dep:base64"] diff --git a/sbv2_core/src/error.rs b/sbv2_core/src/error.rs index 6bb449e..57c2c48 100644 --- a/sbv2_core/src/error.rs +++ b/sbv2_core/src/error.rs @@ -21,6 +21,9 @@ pub enum Error { HoundError(#[from] hound::Error), #[error("model not found error")] ModelNotFoundError(String), + #[cfg(feature = "base64")] + #[error("base64 error")] + Base64Error(#[from] base64::DecodeError), #[error("other")] OtherError(String), } diff --git a/sbv2_core/src/tts.rs b/sbv2_core/src/tts.rs index 4e30819..49a94b6 100644 --- a/sbv2_core/src/tts.rs +++ b/sbv2_core/src/tts.rs @@ -3,6 +3,12 @@ use crate::{jtalk, model, style, tokenizer, tts_util}; use ndarray::{concatenate, Array1, Array2, Array3, Axis}; use ort::Session; use tokenizers::Tokenizer; +#[cfg(feature = "aivmx")] +use base64::prelude::{Engine as _, BASE64_STANDARD}; +#[cfg(feature = "aivmx")] +use std::io::Cursor; +#[cfg(feature = "aivmx")] +use ndarray::ShapeBuilder; #[derive(PartialEq, Eq, Clone)] pub struct TTSIdent(String); @@ -69,6 +75,55 @@ impl TTSModelHolder { self.models.iter().map(|m| m.ident.to_string()).collect() } + #[cfg(feature = "aivmx")] + pub fn load_aivmx, P: AsRef<[u8]>>( + &mut self, + ident: I, + aivmx_bytes: P + ) -> Result<()> { + let ident = ident.into(); + if self.find_model(ident.clone()).is_err() { + let mut load = true; + if let Some(max) = self.max_loaded_models { + if self.models.iter().filter(|x| x.vits2.is_some()).count() >= max { + load = false; + } + } + let model = model::load_model(aivmx_bytes, false)?; + let metadata = model.metadata()?; + if let Some(aivm_style_vectors) = metadata.custom("aivm_style_vectors")? { + let style_vectors = Cursor::new(&BASE64_STANDARD.decode(aivm_style_vectors)?); + let reader = npyz::NpyFile::new(style_vectors)?; + let style_vectors = { + let shape = reader.shape().to_vec(); + let order = reader.order(); + let data = reader.into_vec::()?; + let shape = match shape[..] { + [i1, i2] => [i1 as usize, i2 as usize], + _ => panic!("expected 2D array"), + }; + let true_shape = shape.set_f(order == npyz::Order::Fortran); + ndarray::Array2::from_shape_vec(true_shape, data)? + }; + self.models.push(TTSModel { + vits2: if load { + Some(model) + } else { + None + }, + bytes: if self.max_loaded_models.is_some() { + Some(aivmx_bytes.as_ref().to_vec()) + } else { + None + }, + ident, + style_vectors, + }) + } + } + Ok(()) + } + /// Load a .sbv2 file binary /// /// # Examples