support aivmx

This commit is contained in:
tuna2134
2024-11-20 01:42:04 +00:00
parent f90904a337
commit 9c9119a107
4 changed files with 150 additions and 1 deletions

87
Cargo.lock generated
View File

@@ -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"

View File

@@ -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_std = ["tokenizers/unstable_wasm"]
aivmx = ["npyz", "base64"]
base64 = ["dep:base64"]

View File

@@ -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),
}

View File

@@ -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<I: Into<TTSIdent>, 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::<f32>()?;
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