Update sbv2.rs

chore(deps): update rust crate anyhow to v1.0.94

.github/FUNDING.yml

@@ -0,0 +1,3 @@
+# These are supported funding model platforms
+
+github: [tuna2134]

.github/workflows/CI.yml

@@ -140,9 +140,6 @@ jobs:
     strategy:
       matrix:
         tag: [cpu, cuda]
-        platform:
-        - linux/amd64
-        - linux/arm64
     steps:
       - uses: actions/checkout@v4
       - name: Set up QEMU
@@ -163,4 +160,4 @@ jobs:
           tags: |
             ghcr.io/${{ github.repository }}:${{ matrix.tag }}
           file: docker/${{ matrix.tag }}.Dockerfile
-          platforms: ${{ matrix.platform }}
+          platforms: linux/amd64, linux/arm64

Colab-sbv2_bindings-CPU.ipynb

@@ -0,0 +1,180 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# 音声合成プログラム\n",
+    "\n",
+    "このノートブックでは、`sbv2_bindings` パッケージを使用して音声合成を行います。必要なモデルをダウンロードし、ユーザーが入力したテキストから音声を生成します。音声合成が終わったら、再度テキストの入力を求め、ユーザーが終了するまで繰り返します。"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# 必要なパッケージのインストール\n",
+    "!pip install sbv2_bindings\n",
+    "\n",
+    "# 必要なモジュールのインポート\n",
+    "import os\n",
+    "import urllib.request\n",
+    "import time\n",
+    "from sbv2_bindings import TTSModel"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## モデルのダウンロード\n",
+    "\n",
+    "モデルファイルとトークナイザーをダウンロードします。ユーザーが独自のモデルを使用したい場合は、該当するURLまたはローカルパスを指定してください。"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# モデルの URL またはローカルパスの指定\n",
+    "user_sbv2_model_url = \"\"  # カスタムモデルのURLがあればここに指定\n",
+    "user_sbv2_model_path = \"\"  # カスタムモデルのローカルパスがあればここに指定\n",
+    "\n",
+    "# モデル用のディレクトリを作成\n",
+    "model_dir = 'models'\n",
+    "os.makedirs(model_dir, exist_ok=True)\n",
+    "\n",
+    "# ダウンロードするファイルの URL\n",
+    "file_urls = [\n",
+    "    \"https://huggingface.co/googlefan/sbv2_onnx_models/resolve/main/tokenizer.json\",\n",
+    "    \"https://huggingface.co/googlefan/sbv2_onnx_models/resolve/main/deberta.onnx\",\n",
+    "]\n",
+    "\n",
+    "# モデルのパス決定\n",
+    "if user_sbv2_model_path:\n",
+    "    sbv2_model_path = user_sbv2_model_path  # ローカルモデルのパスを使用\n",
+    "elif user_sbv2_model_url:\n",
+    "    sbv2_model_filename = os.path.basename(user_sbv2_model_url)\n",
+    "    sbv2_model_path = os.path.join(model_dir, sbv2_model_filename)\n",
+    "    file_urls.append(user_sbv2_model_url)\n",
+    "else:\n",
+    "    # デフォルトのモデルを使用\n",
+    "    sbv2_model_filename = \"tsukuyomi.sbv2\"\n",
+    "    sbv2_model_path = os.path.join(model_dir, sbv2_model_filename)\n",
+    "    file_urls.append(\"https://huggingface.co/googlefan/sbv2_onnx_models/resolve/main/tsukuyomi.sbv2\")\n",
+    "\n",
+    "# ファイルをダウンロード\n",
+    "for url in file_urls:\n",
+    "    file_name = os.path.join(model_dir, os.path.basename(url))\n",
+    "    if not os.path.exists(file_name):\n",
+    "        print(f\"{file_name} をダウンロードしています...\")\n",
+    "        urllib.request.urlretrieve(url, file_name)\n",
+    "    else:\n",
+    "        print(f\"{file_name} は既に存在します。\")\n",
+    "\n",
+    "# ダウンロードまたは使用するファイルを確認\n",
+    "print(\"\\n使用するファイル:\")\n",
+    "for file in os.listdir(model_dir):\n",
+    "    print(file)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## モデルの読み込みと音声合成\n",
+    "\n",
+    "モデルを読み込み、ユーザーが入力したテキストから音声を生成します。話者名は使用する `.sbv2` ファイル名から自動的に取得します。音声合成が終わったら、再度テキストの入力を求め、ユーザーが終了するまで繰り返します。"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# 音声合成の実行\n",
+    "def main():\n",
+    "    try:\n",
+    "        print(\"\\nモデルを読み込んでいます...\")\n",
+    "        model = TTSModel.from_path(\n",
+    "            os.path.join(model_dir, \"deberta.onnx\"),\n",
+    "            os.path.join(model_dir, \"tokenizer.json\")\n",
+    "        )\n",
+    "        print(\"モデルの読み込みが完了しました！\")\n",
+    "    except Exception as e:\n",
+    "        print(f\"モデルの読み込みに失敗しました: {e}\")\n",
+    "        return\n",
+    "\n",
+    "    # 話者名を取得（.sbv2 ファイル名の拡張子を除いた部分）\n",
+    "    speaker_name = os.path.splitext(os.path.basename(sbv2_model_path))[0]\n",
+    "    \n",
+    "    # 指定されたモデルのパスを使用\n",
+    "    try:\n",
+    "        model.load_sbv2file_from_path(speaker_name, sbv2_model_path)\n",
+    "        print(f\"話者 '{speaker_name}' のセットアップが完了しました！\")\n",
+    "    except Exception as e:\n",
+    "        print(f\"SBV2ファイルの読み込みに失敗しました: {e}\")\n",
+    "        return\n",
+    "\n",
+    "    # 音声合成を繰り返し実行\n",
+    "    while True:\n",
+    "        # 合成したいテキストをユーザーから入力\n",
+    "        user_input = input(\"\\n音声合成したいテキストを入力してください（終了するには 'exit' と入力）: \")\n",
+    "        \n",
+    "        if user_input.strip().lower() == 'exit':\n",
+    "            print(\"音声合成を終了します。\")\n",
+    "            break\n",
+    "\n",
+    "        # 出力ファイル名\n",
+    "        output_file = \"output.wav\"\n",
+    "\n",
+    "        # 音声合成を実行\n",
+    "        try:\n",
+    "            print(\"\\n音声合成を開始します...\")\n",
+    "            start_time = time.time()\n",
+    "\n",
+    "            audio_data = model.synthesize(user_input, speaker_name, 0, 0.0, 1)\n",
+    "\n",
+    "            with open(output_file, \"wb\") as f:\n",
+    "                f.write(audio_data)\n",
+    "\n",
+    "            end_time = time.time()\n",
+    "            elapsed_time = end_time - start_time\n",
+    "\n",
+    "            print(f\"\\n音声が '{output_file}' に保存されました。\")\n",
+    "            print(f\"音声合成にかかった時間: {elapsed_time:.2f} 秒\")\n",
+    "        except Exception as e:\n",
+    "            print(f\"音声合成に失敗しました: {e}\")\n",
+    "\n",
+    "if __name__ == \"__main__\":\n",
+    "    main()"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.x"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}

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

@@ -1,6 +1,6 @@
 [package]
 name = "sbv2_api"
-version = "0.2.0-alpha"
+version = "0.2.0-alpha4"
 edition = "2021"
 
 [dependencies]
@@ -9,11 +9,11 @@ axum = "0.7.5"
 dotenvy.workspace = true
 env_logger.workspace = true
 log = "0.4.22"
-sbv2_core = { version = "0.2.0-alpha", 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 = "4.2.3", features = ["axum_extras"] }
-utoipa-scalar = { version = "0.1.0", features = ["axum"] }
+utoipa = { version = "5.0.0", features = ["axum_extras"] }
+utoipa-scalar = { version = "0.2.0", features = ["axum"] }
 
 [features]
 coreml = ["sbv2_core/coreml"]

sbv2_api/src/main.rs

@@ -40,6 +40,14 @@ fn length_default() -> f32 {
     1.0
 }
 
+fn style_id_default() -> i32 {
+    0
+}
+
+fn speaker_id_default() -> i64 {
+    0
+}
+
 #[derive(Deserialize, ToSchema)]
 struct SynthesizeRequest {
     text: String,
@@ -48,6 +56,10 @@ struct SynthesizeRequest {
     sdp_ratio: f32,
     #[serde(default = "length_default")]
     length_scale: f32,
+    #[serde(default = "style_id_default")]
+    style_id: i32,
+    #[serde(default = "speaker_id_default")]
+    speaker_id: i64,
 }
 
 #[utoipa::path(
@@ -65,15 +77,18 @@ async fn synthesize(
         ident,
         sdp_ratio,
         length_scale,
+        style_id,
+        speaker_id,
     }): Json<SynthesizeRequest>,
 ) -> AppResult<impl IntoResponse> {
     log::debug!("processing request: text={text}, ident={ident}, sdp_ratio={sdp_ratio}, length_scale={length_scale}");
     let buffer = {
-        let tts_model = state.tts_model.lock().await;
+        let mut tts_model = state.tts_model.lock().await;
         tts_model.easy_synthesize(
             &ident,
             &text,
-            0,
+            style_id,
+            speaker_id,
             SynthesizeOptions {
                 sdp_ratio,
                 length_scale,
@@ -94,6 +109,9 @@ impl AppState {
         let mut tts_model = TTSModelHolder::new(
             &fs::read(env::var("BERT_MODEL_PATH")?).await?,
             &fs::read(env::var("TOKENIZER_PATH")?).await?,
+            env::var("HOLDER_MAX_LOADED_MODElS")
+                .ok()
+                .and_then(|x| x.parse().ok()),
         )?;
         let models = env::var("MODELS_PATH").unwrap_or("models".to_string());
         let mut f = fs::read_dir(&models).await?;
@@ -122,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_bindings/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "sbv2_bindings"
-version = "0.2.0-alpha1"
+version = "0.2.0-alpha4"
 edition = "2021"
 
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
@@ -12,4 +12,4 @@ crate-type = ["cdylib"]
 anyhow.workspace = true
 ndarray.workspace = true
 pyo3 = { version = "0.22.0", features = ["anyhow"] }
-sbv2_core = { version = "0.2.0-alpha", path = "../sbv2_core" }
+sbv2_core = { version = "0.2.0-alpha2", path = "../sbv2_core" }

sbv2_bindings/src/sbv2.rs

@@ -23,10 +23,15 @@ pub struct TTSModel {
 
 #[pymethods]
 impl TTSModel {
+    #[pyo3(signature = (bert_model_bytes, tokenizer_bytes, max_loaded_models=None))]
     #[new]
-    fn new(bert_model_bytes: Vec<u8>, tokenizer_bytes: Vec<u8>) -> anyhow::Result<Self> {
+    fn new(
+        bert_model_bytes: Vec<u8>,
+        tokenizer_bytes: Vec<u8>,
+        max_loaded_models: Option<usize>,
+    ) -> anyhow::Result<Self> {
         Ok(Self {
-            model: TTSModelHolder::new(bert_model_bytes, tokenizer_bytes)?,
+            model: TTSModelHolder::new(bert_model_bytes, tokenizer_bytes, max_loaded_models)?,
         })
     }
 
@@ -38,10 +43,21 @@ impl TTSModel {
     ///     BERTモデルのパス
     /// tokenizer_path : str
     ///     トークナイザーのパス
+    /// max_loaded_models: int | None
+    ///     同時にVRAMに存在するモデルの数
+    #[pyo3(signature = (bert_model_path, tokenizer_path, max_loaded_models=None))]
     #[staticmethod]
-    fn from_path(bert_model_path: String, tokenizer_path: String) -> anyhow::Result<Self> {
+    fn from_path(
+        bert_model_path: String,
+        tokenizer_path: String,
+        max_loaded_models: Option<usize>,
+    ) -> anyhow::Result<Self> {
         Ok(Self {
-            model: TTSModelHolder::new(fs::read(bert_model_path)?, fs::read(tokenizer_path)?)?,
+            model: TTSModelHolder::new(
+                fs::read(bert_model_path)?,
+                fs::read(tokenizer_path)?,
+                max_loaded_models,
+            )?,
         })
     }
 
@@ -120,12 +136,13 @@ impl TTSModel {
     /// -------
     /// voice_data : bytes
     ///     音声データ
-    fn synthesize<'p>(
-        &'p self,
-        py: Python<'p>,
+    fn synthesize(
+        &mut self,
+        py: Python,
         text: String,
         ident: String,
         style_id: i32,
+        speaker_id: i64,
         sdp_ratio: f32,
         length_scale: f32,
     ) -> anyhow::Result<Bound<PyBytes>> {
@@ -133,13 +150,14 @@ impl TTSModel {
             ident.as_str(),
             &text,
             style_id,
+            speaker_id,
             SynthesizeOptions {
                 sdp_ratio,
                 length_scale,
                 ..Default::default()
             },
         )?;
-        Ok(PyBytes::new_bound(py, &data))
+        Ok(PyBytes::new(py, &data))
     }
 
     fn unload(&mut self, ident: String) -> bool {

sbv2_core/Cargo.toml

@@ -1,7 +1,7 @@
 [package]
 name = "sbv2_core"
 description = "Style-Bert-VITSの推論ライブラリ"
-version = "0.2.0-alpha1"
+version = "0.2.0-alpha4"
 edition = "2021"
 license = "MIT"
 readme = "../README.md"
@@ -10,14 +10,16 @@ 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.6", optional = true }
+ort = { git = "https://github.com/pykeio/ort.git", version = "2.0.0-rc.8", optional = true }
 regex = "1.10.6"
 serde = { version = "1.0.210", features = ["derive"] }
 serde_json = "1.0.128"
@@ -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/bert.rs

@@ -1,5 +1,5 @@
 use crate::error::Result;
-use ndarray::Array2;
+use ndarray::{Array2, Ix2};
 use ort::Session;
 
 pub fn predict(
@@ -14,10 +14,10 @@ pub fn predict(
         }?
     )?;
 
-    let output = outputs.get("output").unwrap();
+    let output = outputs["output"]
+        .try_extract_tensor::<f32>()?
+        .into_dimensionality::<Ix2>()?
+        .to_owned();
 
-    let content = output.try_extract_tensor::<f32>()?.to_owned();
-    let (data, _) = content.clone().into_raw_vec_and_offset();
-
-    Ok(Array2::from_shape_vec((content.shape()[0], content.shape()[1]), data).unwrap())
+    Ok(output)
 }

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

@@ -11,14 +11,26 @@ fn main_inner() -> anyhow::Result<()> {
     let mut tts_holder = tts::TTSModelHolder::new(
         &fs::read(env::var("BERT_MODEL_PATH")?)?,
         &fs::read(env::var("TOKENIZER_PATH")?)?,
+        env::var("HOLDER_MAX_LOADED_MODElS")
+            .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, tts::SynthesizeOptions::default())?;
+    let audio =
+        tts_holder.easy_synthesize(ident, &text, 0, 0, tts::SynthesizeOptions::default())?;
     fs::write("output.wav", audio)?;
 
     Ok(())
 }
+
 #[cfg(not(feature = "std"))]
 fn main_inner() -> anyhow::Result<()> {
     Ok(())

sbv2_core/src/model.rs

@@ -1,5 +1,5 @@
 use crate::error::Result;
-use ndarray::{array, Array1, Array2, Array3, Axis};
+use ndarray::{array, Array1, Array2, Array3, Axis, Ix3};
 use ort::{GraphOptimizationLevel, Session};
 
 #[allow(clippy::vec_init_then_push, unused_variables)]
@@ -52,11 +52,14 @@ pub fn synthesize(
     session: &Session,
     bert_ori: Array2<f32>,
     x_tst: Array1<i64>,
+    sid: Array1<i64>,
     tones: Array1<i64>,
     lang_ids: Array1<i64>,
     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];
@@ -67,27 +70,21 @@ pub fn synthesize(
     let outputs = session.run(ort::inputs! {
         "x_tst" => x_tst,
         "x_tst_lengths" => x_tst_lengths,
-        "sid" => array![0_i64],
+        "sid" => sid,
         "tones" => tones,
         "language" => lang_ids,
         "bert" => bert,
         "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
-        .get("output")
-        .unwrap()
+    let audio_array = outputs["output"]
         .try_extract_tensor::<f32>()?
+        .into_dimensionality::<Ix3>()?
         .to_owned();
 
-    Ok(Array3::from_shape_vec(
-        (
-            audio_array.shape()[0],
-            audio_array.shape()[1],
-            audio_array.shape()[2],
-        ),
-        audio_array.into_raw_vec_and_offset().0,
-    )?)
+    Ok(audio_array)
 }

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)]
@@ -24,9 +30,10 @@ where
 }
 
 pub struct TTSModel {
-    vits2: Session,
+    vits2: Option<Session>,
     style_vectors: Array2<f32>,
     ident: TTSIdent,
+    bytes: Option<Vec<u8>>,
 }
 
 /// High-level Style-Bert-VITS2's API
@@ -35,6 +42,7 @@ pub struct TTSModelHolder {
     bert: Session,
     models: Vec<TTSModel>,
     jtalk: jtalk::JTalk,
+    max_loaded_models: Option<usize>,
 }
 
 impl TTSModelHolder {
@@ -43,9 +51,13 @@ impl TTSModelHolder {
     /// # Examples
     ///
     /// ```rs
-    /// let mut tts_holder = TTSModelHolder::new(std::fs::read("deberta.onnx")?, std::fs::read("tokenizer.json")?)?;
+    /// let mut tts_holder = TTSModelHolder::new(std::fs::read("deberta.onnx")?, std::fs::read("tokenizer.json")?, None)?;
     /// ```
-    pub fn new<P: AsRef<[u8]>>(bert_model_bytes: P, tokenizer_bytes: P) -> Result<Self> {
+    pub fn new<P: AsRef<[u8]>>(
+        bert_model_bytes: P,
+        tokenizer_bytes: P,
+        max_loaded_models: Option<usize>,
+    ) -> Result<Self> {
         let bert = model::load_model(bert_model_bytes, true)?;
         let jtalk = jtalk::JTalk::new()?;
         let tokenizer = tokenizer::get_tokenizer(tokenizer_bytes)?;
@@ -54,6 +66,7 @@ impl TTSModelHolder {
             models: vec![],
             jtalk,
             tokenizer,
+            max_loaded_models,
         })
     }
 
@@ -62,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
@@ -94,10 +154,25 @@ impl TTSModelHolder {
     ) -> 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;
+                }
+            }
             self.models.push(TTSModel {
-                vits2: model::load_model(vits2_bytes, false)?,
+                vits2: if load {
+                    Some(model::load_model(&vits2_bytes, false)?)
+                } else {
+                    None
+                },
                 style_vectors: style::load_style(style_vectors_bytes)?,
                 ident,
+                bytes: if self.max_loaded_models.is_some() {
+                    Some(vits2_bytes.as_ref().to_vec())
+                } else {
+                    None
+                },
             })
         }
         Ok(())
@@ -145,6 +220,42 @@ impl TTSModelHolder {
             .find(|m| m.ident == ident)
             .ok_or(Error::ModelNotFoundError(ident.to_string()))
     }
+    fn find_and_load_model<I: Into<TTSIdent>>(&mut self, ident: I) -> Result<bool> {
+        let ident = ident.into();
+        let (bytes, style_vectors) = {
+            let model = self
+                .models
+                .iter()
+                .find(|m| m.ident == ident)
+                .ok_or(Error::ModelNotFoundError(ident.to_string()))?;
+            if model.vits2.is_some() {
+                return Ok(true);
+            }
+            (model.bytes.clone().unwrap(), model.style_vectors.clone())
+        };
+        self.unload(ident.clone());
+        let s = model::load_model(&bytes, false)?;
+        if let Some(max) = self.max_loaded_models {
+            if self.models.iter().filter(|x| x.vits2.is_some()).count() >= max {
+                self.unload(self.models.first().unwrap().ident.clone());
+            }
+        }
+        self.models.push(TTSModel {
+            bytes: Some(bytes.to_vec()),
+            vits2: Some(s),
+            style_vectors,
+            ident: ident.clone(),
+        });
+        let model = self
+            .models
+            .iter()
+            .find(|m| m.ident == ident)
+            .ok_or(Error::ModelNotFoundError(ident.to_string()))?;
+        if model.vits2.is_some() {
+            return Ok(true);
+        }
+        Err(Error::ModelNotFoundError(ident.to_string()))
+    }
 
     /// Get style vector by style id and weight
     ///
@@ -167,12 +278,19 @@ impl TTSModelHolder {
     /// let audio = tts_holder.easy_synthesize("tsukuyomi", "こんにちは", 0, SynthesizeOptions::default())?;
     /// ```
     pub fn easy_synthesize<I: Into<TTSIdent> + Copy>(
-        &self,
+        &mut self,
         ident: I,
         text: &str,
         style_id: i32,
+        speaker_id: i64,
         options: SynthesizeOptions,
     ) -> Result<Vec<u8>> {
+        self.find_and_load_model(ident)?;
+        let vits2 = &self
+            .find_model(ident)?
+            .vits2
+            .as_ref()
+            .ok_or(Error::ModelNotFoundError(ident.into().to_string()))?;
         let style_vector = self.get_style_vector(ident, style_id, options.style_weight)?;
         let audio_array = if options.split_sentences {
             let texts: Vec<&str> = text.split('\n').collect();
@@ -183,14 +301,17 @@ impl TTSModelHolder {
                 }
                 let (bert_ori, phones, tones, lang_ids) = self.parse_text(t)?;
                 let audio = model::synthesize(
-                    &self.find_model(ident)?.vits2,
+                    vits2,
                     bert_ori.to_owned(),
                     phones,
+                    Array1::from_vec(vec![speaker_id]),
                     tones,
                     lang_ids,
                     style_vector.clone(),
                     options.sdp_ratio,
                     options.length_scale,
+                    0.677,
+                    0.8,
                 )?;
                 audios.push(audio.clone());
                 if i != texts.len() - 1 {
@@ -204,47 +325,21 @@ impl TTSModelHolder {
         } else {
             let (bert_ori, phones, tones, lang_ids) = self.parse_text(text)?;
             model::synthesize(
-                &self.find_model(ident)?.vits2,
+                vits2,
                 bert_ori.to_owned(),
                 phones,
+                Array1::from_vec(vec![speaker_id]),
                 tones,
                 lang_ids,
                 style_vector,
                 options.sdp_ratio,
                 options.length_scale,
+                0.677,
+                0.8,
             )?
         };
         tts_util::array_to_vec(audio_array)
     }
-
-    /// Synthesize text to audio
-    ///
-    /// # Note
-    /// This function is for low-level usage, use `easy_synthesize` for high-level usage.
-    #[allow(clippy::too_many_arguments)]
-    pub fn synthesize<I: Into<TTSIdent>>(
-        &self,
-        ident: I,
-        bert_ori: Array2<f32>,
-        phones: Array1<i64>,
-        tones: Array1<i64>,
-        lang_ids: Array1<i64>,
-        style_vector: Array1<f32>,
-        sdp_ratio: f32,
-        length_scale: f32,
-    ) -> Result<Vec<u8>> {
-        let audio_array = model::synthesize(
-            &self.find_model(ident)?.vits2,
-            bert_ori.to_owned(),
-            phones,
-            tones,
-            lang_ids,
-            style_vector,
-            sdp_ratio,
-            length_scale,
-        )?;
-        tts_util::array_to_vec(audio_array)
-    }
 }
 
 /// Synthesize options

sbv2_wasm/pnpm-lock.yaml

@@ -10,72 +10,72 @@ importers:
     dependencies:
       onnxruntime-web:
         specifier: ^1.19.2
-        version: 1.19.2
+        version: 1.20.0
     devDependencies:
       '@biomejs/biome':
         specifier: ^1.9.2
-        version: 1.9.3
+        version: 1.9.4
       '@types/node':
         specifier: ^22.7.4
-        version: 22.7.4
+        version: 22.10.1
       esbuild:
         specifier: ^0.24.0
         version: 0.24.0
       typescript:
         specifier: ^5.6.2
-        version: 5.6.2
+        version: 5.6.3
 
 packages:
 
-  '@biomejs/biome@1.9.3':
-    resolution: {integrity: sha512-POjAPz0APAmX33WOQFGQrwLvlu7WLV4CFJMlB12b6ZSg+2q6fYu9kZwLCOA+x83zXfcPd1RpuWOKJW0GbBwLIQ==}
+  '@biomejs/biome@1.9.4':
+    resolution: {integrity: sha512-1rkd7G70+o9KkTn5KLmDYXihGoTaIGO9PIIN2ZB7UJxFrWw04CZHPYiMRjYsaDvVV7hP1dYNRLxSANLaBFGpog==}
     engines: {node: '>=14.21.3'}
     hasBin: true
 
-  '@biomejs/cli-darwin-arm64@1.9.3':
-    resolution: {integrity: sha512-QZzD2XrjJDUyIZK+aR2i5DDxCJfdwiYbUKu9GzkCUJpL78uSelAHAPy7m0GuPMVtF/Uo+OKv97W3P9nuWZangQ==}
+  '@biomejs/cli-darwin-arm64@1.9.4':
+    resolution: {integrity: sha512-bFBsPWrNvkdKrNCYeAp+xo2HecOGPAy9WyNyB/jKnnedgzl4W4Hb9ZMzYNbf8dMCGmUdSavlYHiR01QaYR58cw==}
     engines: {node: '>=14.21.3'}
     cpu: [arm64]
     os: [darwin]
 
-  '@biomejs/cli-darwin-x64@1.9.3':
-    resolution: {integrity: sha512-vSCoIBJE0BN3SWDFuAY/tRavpUtNoqiceJ5PrU3xDfsLcm/U6N93JSM0M9OAiC/X7mPPfejtr6Yc9vSgWlEgVw==}
+  '@biomejs/cli-darwin-x64@1.9.4':
+    resolution: {integrity: sha512-ngYBh/+bEedqkSevPVhLP4QfVPCpb+4BBe2p7Xs32dBgs7rh9nY2AIYUL6BgLw1JVXV8GlpKmb/hNiuIxfPfZg==}
     engines: {node: '>=14.21.3'}
     cpu: [x64]
     os: [darwin]
 
-  '@biomejs/cli-linux-arm64-musl@1.9.3':
-    resolution: {integrity: sha512-VBzyhaqqqwP3bAkkBrhVq50i3Uj9+RWuj+pYmXrMDgjS5+SKYGE56BwNw4l8hR3SmYbLSbEo15GcV043CDSk+Q==}
+  '@biomejs/cli-linux-arm64-musl@1.9.4':
+    resolution: {integrity: sha512-v665Ct9WCRjGa8+kTr0CzApU0+XXtRgwmzIf1SeKSGAv+2scAlW6JR5PMFo6FzqqZ64Po79cKODKf3/AAmECqA==}
     engines: {node: '>=14.21.3'}
     cpu: [arm64]
     os: [linux]
 
-  '@biomejs/cli-linux-arm64@1.9.3':
-    resolution: {integrity: sha512-vJkAimD2+sVviNTbaWOGqEBy31cW0ZB52KtpVIbkuma7PlfII3tsLhFa+cwbRAcRBkobBBhqZ06hXoZAN8NODQ==}
+  '@biomejs/cli-linux-arm64@1.9.4':
+    resolution: {integrity: sha512-fJIW0+LYujdjUgJJuwesP4EjIBl/N/TcOX3IvIHJQNsAqvV2CHIogsmA94BPG6jZATS4Hi+xv4SkBBQSt1N4/g==}
     engines: {node: '>=14.21.3'}
     cpu: [arm64]
     os: [linux]
 
-  '@biomejs/cli-linux-x64-musl@1.9.3':
-    resolution: {integrity: sha512-TJmnOG2+NOGM72mlczEsNki9UT+XAsMFAOo8J0me/N47EJ/vkLXxf481evfHLlxMejTY6IN8SdRSiPVLv6AHlA==}
+  '@biomejs/cli-linux-x64-musl@1.9.4':
+    resolution: {integrity: sha512-gEhi/jSBhZ2m6wjV530Yy8+fNqG8PAinM3oV7CyO+6c3CEh16Eizm21uHVsyVBEB6RIM8JHIl6AGYCv6Q6Q9Tg==}
     engines: {node: '>=14.21.3'}
     cpu: [x64]
     os: [linux]
 
-  '@biomejs/cli-linux-x64@1.9.3':
-    resolution: {integrity: sha512-x220V4c+romd26Mu1ptU+EudMXVS4xmzKxPVb9mgnfYlN4Yx9vD5NZraSx/onJnd3Gh/y8iPUdU5CDZJKg9COA==}
+  '@biomejs/cli-linux-x64@1.9.4':
+    resolution: {integrity: sha512-lRCJv/Vi3Vlwmbd6K+oQ0KhLHMAysN8lXoCI7XeHlxaajk06u7G+UsFSO01NAs5iYuWKmVZjmiOzJ0OJmGsMwg==}
     engines: {node: '>=14.21.3'}
     cpu: [x64]
     os: [linux]
 
-  '@biomejs/cli-win32-arm64@1.9.3':
-    resolution: {integrity: sha512-lg/yZis2HdQGsycUvHWSzo9kOvnGgvtrYRgoCEwPBwwAL8/6crOp3+f47tPwI/LI1dZrhSji7PNsGKGHbwyAhw==}
+  '@biomejs/cli-win32-arm64@1.9.4':
+    resolution: {integrity: sha512-tlbhLk+WXZmgwoIKwHIHEBZUwxml7bRJgk0X2sPyNR3S93cdRq6XulAZRQJ17FYGGzWne0fgrXBKpl7l4M87Hg==}
     engines: {node: '>=14.21.3'}
     cpu: [arm64]
     os: [win32]
 
-  '@biomejs/cli-win32-x64@1.9.3':
-    resolution: {integrity: sha512-cQMy2zanBkVLpmmxXdK6YePzmZx0s5Z7KEnwmrW54rcXK3myCNbQa09SwGZ8i/8sLw0H9F3X7K4rxVNGU8/D4Q==}
+  '@biomejs/cli-win32-x64@1.9.4':
+    resolution: {integrity: sha512-8Y5wMhVIPaWe6jw2H+KlEm4wP/f7EW3810ZLmDlrEEy5KvBsb9ECEfu/kMWD484ijfQ8+nIi0giMgu9g1UAuuA==}
     engines: {node: '>=14.21.3'}
     cpu: [x64]
     os: [win32]
@@ -254,8 +254,8 @@ packages:
   '@protobufjs/utf8@1.1.0':
     resolution: {integrity: sha512-Vvn3zZrhQZkkBE8LSuW3em98c0FwgO4nxzv6OdSxPKJIEKY2bGbHn+mhGIPerzI4twdxaP8/0+06HBpwf345Lw==}
 
-  '@types/node@22.7.4':
-    resolution: {integrity: sha512-y+NPi1rFzDs1NdQHHToqeiX2TIS79SWEAw9GYhkkx8bD0ChpfqC+n2j5OXOCpzfojBEBt6DnEnnG9MY0zk1XLg==}
+  '@types/node@22.10.1':
+    resolution: {integrity: sha512-qKgsUwfHZV2WCWLAnVP1JqnpE6Im6h3Y0+fYgMTasNQ7V++CBX5OT1as0g0f+OyubbFqhf6XVNIsmN4IIhEgGQ==}
 
   esbuild@0.24.0:
     resolution: {integrity: sha512-FuLPevChGDshgSicjisSooU0cemp/sGXR841D5LHMB7mTVOmsEHcAxaH3irL53+8YDIeVNQEySh4DaYU/iuPqQ==}
@@ -271,11 +271,11 @@ packages:
   long@5.2.3:
     resolution: {integrity: sha512-lcHwpNoggQTObv5apGNCTdJrO69eHOZMi4BNC+rTLER8iHAqGrUVeLh/irVIM7zTw2bOXA8T6uNPeujwOLg/2Q==}
 
-  onnxruntime-common@1.19.2:
-    resolution: {integrity: sha512-a4R7wYEVFbZBlp0BfhpbFWqe4opCor3KM+5Wm22Az3NGDcQMiU2hfG/0MfnBs+1ZrlSGmlgWeMcXQkDk1UFb8Q==}
+  onnxruntime-common@1.20.0:
+    resolution: {integrity: sha512-9ehS4ul5fBszIcHhfxuDgk45lO+Fqrxmrgwk1Pxb1JRvbQiCB/v9Royv95SRCWHktLMviqNjBsEd/biJhd39cg==}
 
-  onnxruntime-web@1.19.2:
-    resolution: {integrity: sha512-r0ok6KpTUXR4WA+rHvUiZn7JoH02e8iS7XE1p5bXk7q3E0UaRFfYvpMNUHqEPiTBMuIssfBxDCQjUihV8dDFPg==}
+  onnxruntime-web@1.20.0:
+    resolution: {integrity: sha512-IoUf8dqHFJLV4DUSz+Ok+xxyN6cQk57gb20m6PZE5gag3QXuvegYMq9dG8t/QF4JjTKIwvfvnr16ouzCCB9IMA==}
 
   platform@1.3.6:
     resolution: {integrity: sha512-fnWVljUchTro6RiCFvCXBbNhJc2NijN7oIQxbwsyL0buWJPG85v81ehlHI9fXrJsMNgTofEoWIQeClKpgxFLrg==}
@@ -284,49 +284,49 @@ packages:
     resolution: {integrity: sha512-mRUWCc3KUU4w1jU8sGxICXH/gNS94DvI1gxqDvBzhj1JpcsimQkYiOJfwsPUykUI5ZaspFbSgmBLER8IrQ3tqw==}
     engines: {node: '>=12.0.0'}
 
-  typescript@5.6.2:
-    resolution: {integrity: sha512-NW8ByodCSNCwZeghjN3o+JX5OFH0Ojg6sadjEKY4huZ52TqbJTJnDo5+Tw98lSy63NZvi4n+ez5m2u5d4PkZyw==}
+  typescript@5.6.3:
+    resolution: {integrity: sha512-hjcS1mhfuyi4WW8IWtjP7brDrG2cuDZukyrYrSauoXGNgx0S7zceP07adYkJycEr56BOUTNPzbInooiN3fn1qw==}
     engines: {node: '>=14.17'}
     hasBin: true
 
-  undici-types@6.19.8:
-    resolution: {integrity: sha512-ve2KP6f/JnbPBFyobGHuerC9g1FYGn/F8n1LWTwNxCEzd6IfqTwUQcNXgEtmmQ6DlRrC1hrSrBnCZPokRrDHjw==}
+  undici-types@6.20.0:
+    resolution: {integrity: sha512-Ny6QZ2Nju20vw1SRHe3d9jVu6gJ+4e3+MMpqu7pqE5HT6WsTSlce++GQmK5UXS8mzV8DSYHrQH+Xrf2jVcuKNg==}
 
 snapshots:
 
-  '@biomejs/biome@1.9.3':
+  '@biomejs/biome@1.9.4':
     optionalDependencies:
-      '@biomejs/cli-darwin-arm64': 1.9.3
-      '@biomejs/cli-darwin-x64': 1.9.3
-      '@biomejs/cli-linux-arm64': 1.9.3
-      '@biomejs/cli-linux-arm64-musl': 1.9.3
-      '@biomejs/cli-linux-x64': 1.9.3
-      '@biomejs/cli-linux-x64-musl': 1.9.3
-      '@biomejs/cli-win32-arm64': 1.9.3
-      '@biomejs/cli-win32-x64': 1.9.3
+      '@biomejs/cli-darwin-arm64': 1.9.4
+      '@biomejs/cli-darwin-x64': 1.9.4
+      '@biomejs/cli-linux-arm64': 1.9.4
+      '@biomejs/cli-linux-arm64-musl': 1.9.4
+      '@biomejs/cli-linux-x64': 1.9.4
+      '@biomejs/cli-linux-x64-musl': 1.9.4
+      '@biomejs/cli-win32-arm64': 1.9.4
+      '@biomejs/cli-win32-x64': 1.9.4
 
-  '@biomejs/cli-darwin-arm64@1.9.3':
+  '@biomejs/cli-darwin-arm64@1.9.4':
     optional: true
 
-  '@biomejs/cli-darwin-x64@1.9.3':
+  '@biomejs/cli-darwin-x64@1.9.4':
     optional: true
 
-  '@biomejs/cli-linux-arm64-musl@1.9.3':
+  '@biomejs/cli-linux-arm64-musl@1.9.4':
     optional: true
 
-  '@biomejs/cli-linux-arm64@1.9.3':
+  '@biomejs/cli-linux-arm64@1.9.4':
     optional: true
 
-  '@biomejs/cli-linux-x64-musl@1.9.3':
+  '@biomejs/cli-linux-x64-musl@1.9.4':
     optional: true
 
-  '@biomejs/cli-linux-x64@1.9.3':
+  '@biomejs/cli-linux-x64@1.9.4':
     optional: true
 
-  '@biomejs/cli-win32-arm64@1.9.3':
+  '@biomejs/cli-win32-arm64@1.9.4':
     optional: true
 
-  '@biomejs/cli-win32-x64@1.9.3':
+  '@biomejs/cli-win32-x64@1.9.4':
     optional: true
 
   '@esbuild/aix-ppc64@0.24.0':
@@ -424,9 +424,9 @@ snapshots:
 
   '@protobufjs/utf8@1.1.0': {}
 
-  '@types/node@22.7.4':
+  '@types/node@22.10.1':
     dependencies:
-      undici-types: 6.19.8
+      undici-types: 6.20.0
 
   esbuild@0.24.0:
     optionalDependencies:
@@ -461,14 +461,14 @@ snapshots:
 
   long@5.2.3: {}
 
-  onnxruntime-common@1.19.2: {}
+  onnxruntime-common@1.20.0: {}
 
-  onnxruntime-web@1.19.2:
+  onnxruntime-web@1.20.0:
     dependencies:
       flatbuffers: 1.12.0
       guid-typescript: 1.0.9
       long: 5.2.3
-      onnxruntime-common: 1.19.2
+      onnxruntime-common: 1.20.0
       platform: 1.3.6
       protobufjs: 7.4.0
 
@@ -486,9 +486,9 @@ snapshots:
       '@protobufjs/path': 1.1.2
       '@protobufjs/pool': 1.1.0
       '@protobufjs/utf8': 1.1.0
-      '@types/node': 22.7.4
+      '@types/node': 22.10.1
       long: 5.2.3
 
-  typescript@5.6.2: {}
+  typescript@5.6.3: {}
 
-  undici-types@6.19.8: {}
+  undici-types@6.20.0: {}