Merge pull request #138 from tuna2134/aivmx

support aivmx

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"

convert/convert_model.py

@@ -94,7 +94,7 @@ model = get_net_g(
 )
 
 
-def forward(x, x_len, sid, tone, lang, bert, style, length_scale, sdp_ratio):
+def forward(x, x_len, sid, tone, lang, bert, style, length_scale, sdp_ratio, noise_scale, noise_scale_w):
     return model.infer(
         x,
         x_len,
@@ -105,6 +105,8 @@ def forward(x, x_len, sid, tone, lang, bert, style, length_scale, sdp_ratio):
         style,
         sdp_ratio=sdp_ratio,
         length_scale=length_scale,
+        noise_scale=noise_scale,
+        noise_scale_w=noise_scale_w,
     )
 
 
@@ -122,6 +124,8 @@ torch.onnx.export(
         style_vec_tensor,
         torch.tensor(1.0),
         torch.tensor(0.0),
+        torch.tensor(0.6777),
+        torch.tensor(0.8),
     ),
     f"../models/model_{out_name}.onnx",
     verbose=True,
@@ -144,6 +148,8 @@ torch.onnx.export(
         "style_vec",
         "length_scale",
         "sdp_ratio",
+        "noise_scale",
+        "noise_scale_w"
     ],
     output_names=["output"],
 )

sbv2_api/Cargo.toml

@@ -9,7 +9,7 @@ axum = "0.7.5"
 dotenvy.workspace = true
 env_logger.workspace = true
 log = "0.4.22"
-sbv2_core = { version = "0.2.0-alpha2", path = "../sbv2_core" }
+sbv2_core = { version = "0.2.0-alpha2", path = "../sbv2_core", features = ["aivmx"] }
 serde = { version = "1.0.210", features = ["derive"] }
 tokio = { version = "1.40.0", features = ["full"] }
 utoipa = { version = "5.0.0", features = ["axum_extras"] }

sbv2_api/src/main.rs

@@ -140,6 +140,20 @@ impl AppState {
                     log::warn!("Error loading {entry}: {e}");
                 };
                 log::info!("Loaded: {entry}");
+            } else if name.ends_with(".aivmx") {
+                let entry = &name[..name.len() - 6];
+                log::info!("Try loading: {entry}");
+                let aivmx_bytes = match fs::read(format!("{models}/{entry}.aivmx")).await {
+                    Ok(b) => b,
+                    Err(e) => {
+                        log::warn!("Error loading aivmx bytes from file {entry}: {e}");
+                        continue;
+                    }
+                };
+                if let Err(e) = tts_model.load_aivmx(entry, aivmx_bytes) {
+                    log::error!("Error loading {entry}: {e}");
+                }
+                log::info!("Loaded: {entry}");
             }
         }
         for entry in entries {

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

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

sbv2_core/src/main.rs

@@ -15,7 +15,14 @@ fn main_inner() -> anyhow::Result<()> {
             .ok()
             .and_then(|x| x.parse().ok()),
     )?;
-    tts_holder.load_sbv2file(ident, fs::read(env::var("MODEL_PATH")?)?)?;
+    #[cfg(not(feature = "aivmx"))]
+    {
+        tts_holder.load_sbv2file(ident, fs::read(env::var("MODEL_PATH")?)?)?;
+    }
+    #[cfg(feature = "aivmx")]
+    {
+        tts_holder.load_aivmx(ident, fs::read(env::var("MODEL_PATH")?)?)?;
+    }
 
     let audio =
         tts_holder.easy_synthesize(ident, &text, 0, 0, tts::SynthesizeOptions::default())?;
@@ -23,6 +30,7 @@ fn main_inner() -> anyhow::Result<()> {
 
     Ok(())
 }
+
 #[cfg(not(feature = "std"))]
 fn main_inner() -> anyhow::Result<()> {
     Ok(())

sbv2_core/src/model.rs

@@ -58,6 +58,8 @@ pub fn synthesize(
     style_vector: Array1<f32>,
     sdp_ratio: f32,
     length_scale: f32,
+    noise_scale: f32,
+    noise_scale_w: f32,
 ) -> Result<Array3<f32>> {
     let bert = bert_ori.insert_axis(Axis(0));
     let x_tst_lengths: Array1<i64> = array![x_tst.shape()[0] as i64];
@@ -75,6 +77,8 @@ pub fn synthesize(
         "style_vec" => style_vector,
         "sdp_ratio" => array![sdp_ratio],
         "length_scale" => array![length_scale],
+        "noise_scale" => array![noise_scale],
+        "noise_scale_w" => array![noise_scale_w]
     }?)?;
 
     let audio_array = outputs["output"]

sbv2_core/src/tts.rs

@@ -1,7 +1,13 @@
 use crate::error::{Error, Result};
 use crate::{jtalk, model, style, tokenizer, tts_util};
+#[cfg(feature = "aivmx")]
+use base64::prelude::{Engine as _, BASE64_STANDARD};
+#[cfg(feature = "aivmx")]
+use ndarray::ShapeBuilder;
 use ndarray::{concatenate, Array1, Array2, Array3, Axis};
 use ort::Session;
+#[cfg(feature = "aivmx")]
+use std::io::Cursor;
 use tokenizers::Tokenizer;
 
 #[derive(PartialEq, Eq, Clone)]
@@ -69,6 +75,53 @@ 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 aivm_style_vectors = BASE64_STANDARD.decode(aivm_style_vectors)?;
+                let style_vectors = Cursor::new(&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)?
+                };
+                drop(metadata);
+                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
@@ -257,6 +310,8 @@ impl TTSModelHolder {
                     style_vector.clone(),
                     options.sdp_ratio,
                     options.length_scale,
+                    0.677,
+                    0.8,
                 )?;
                 audios.push(audio.clone());
                 if i != texts.len() - 1 {
@@ -279,6 +334,8 @@ impl TTSModelHolder {
                 style_vector,
                 options.sdp_ratio,
                 options.length_scale,
+                0.677,
+                0.8,
             )?
         };
         tts_util::array_to_vec(audio_array)