index-tts stash save

.github/workflows/ci.yml

@@ -31,7 +31,7 @@ jobs:
       - name: Install make dependencies
         run: |
           sudo apt-get update
-          sudo apt-get install -y build-essential
+          sudo apt-get install -y build-essential cmake
 
       - name: Install FFmpeg dependencies
         run: |
@@ -61,7 +61,7 @@ jobs:
       - name: Install make dependencies
         run: |
           sudo apt-get update
-          sudo apt-get install -y build-essential
+          sudo apt-get install -y build-essential cmake
           
       - name: Install FFmpeg development packages
         run: |

Cargo.lock

@@ -52,8 +52,10 @@ dependencies = [
  "num",
  "rayon",
  "realfft",
+ "regex",
  "reqwest 0.12.24",
  "rocket",
+ "sentencepiece",
  "serde",
  "serde_json",
  "serde_yaml",
@@ -3305,6 +3307,29 @@ dependencies = [
  "syn",
 ]
 
+[[package]]
+name = "prost"
+version = "0.14.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d2ea70524a2f82d518bce41317d0fae74151505651af45faf1ffbd6fd33f0568"
+dependencies = [
+ "bytes",
+ "prost-derive",
+]
+
+[[package]]
+name = "prost-derive"
+version = "0.14.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "27c6023962132f4b30eb4c172c91ce92d933da334c59c23cddee82358ddafb0b"
+dependencies = [
+ "anyhow",
+ "itertools 0.14.0",
+ "proc-macro2",
+ "quote",
+ "syn",
+]
+
 [[package]]
 name = "psl-types"
 version = "2.0.11"
@@ -3667,9 +3692,9 @@ dependencies = [
 
 [[package]]
 name = "regex"
-version = "1.12.2"
+version = "1.12.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "843bc0191f75f3e22651ae5f1e72939ab2f72a4bc30fa80a066bd66edefc24d4"
+checksum = "e10754a14b9137dd7b1e3e5b0493cc9171fdd105e0ab477f51b72e7f3ac0e276"
 dependencies = [
  "aho-corasick",
  "memchr",
@@ -4104,6 +4129,32 @@ dependencies = [
  "libc",
 ]
 
+[[package]]
+name = "sentencepiece"
+version = "0.13.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "3baa1506c7718f6b70bcac5475e563176dd51bb669dae38181abbf3070517453"
+dependencies = [
+ "libc",
+ "num-derive",
+ "num-traits",
+ "prost",
+ "prost-derive",
+ "sentencepiece-sys",
+ "thiserror 2.0.17",
+]
+
+[[package]]
+name = "sentencepiece-sys"
+version = "0.13.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "fa4f9b54dc005df8ec1c3f9e2347cea6b6657ec5460977a982e7a4e09ef49411"
+dependencies = [
+ "cc",
+ "cmake",
+ "pkg-config",
+]
+
 [[package]]
 name = "seq-macro"
 version = "0.3.6"

Cargo.toml

@@ -40,6 +40,8 @@ serde_yaml = "0.9.34"
 zip = "7.2.0"
 half = "2.7.1"
 byteorder = "1.5.0"
+sentencepiece = "0.13.1"
+regex = "1.12.3"
 
 [features]
 flash-attn = ["candle-flash-attn"]

README.md

@@ -62,6 +62,13 @@
 ⭐ 如果这个项目对你有帮助，请给我们一个 Star！
 
 ## 环境依赖
+* cmake: 
+    * ubuntu/WSL
+    ```bash
+    sudo apt-get update
+    sudo apt-get install cmake
+    ```
+    
 * 启用ffmpeg的feature时：
     * ubuntu/WSL
     ```bash

src/models/common/mod.rs

@@ -2,9 +2,9 @@ use anyhow::{Result, anyhow};
 use candle_core::{D, IndexOp, Tensor};
 use candle_nn::{
     Activation, BatchNorm, BatchNormConfig, Conv1d, Conv1dConfig, Conv2d, Conv2dConfig,
-    ConvTranspose1d, ConvTranspose1dConfig, Embedding, LayerNorm, LayerNormConfig, Linear, Module,
+    ConvTranspose1d, ConvTranspose1dConfig, Embedding, GroupNorm, Init, LayerNorm, LayerNormConfig,
-    ModuleT, RmsNorm, VarBuilder, batch_norm, conv1d, conv1d_no_bias, conv2d, conv2d_no_bias,
+    Linear, Module, ModuleT, RmsNorm, VarBuilder, batch_norm, conv1d, conv1d_no_bias, conv2d,
-    embedding, layer_norm, linear_b, linear_no_bias, ops::sigmoid, rms_norm,
+    conv2d_no_bias, embedding, layer_norm, linear_b, linear_no_bias, ops::sigmoid, rms_norm,
 };
 use rayon::iter::{IndexedParallelIterator, IntoParallelRefIterator, ParallelIterator};
 
@@ -464,9 +464,10 @@ impl NaiveAttnTwoLinearMLPBlock {
             linear2_pp_name,
         )?;
 
-        let input_layernorm = get_layer_norm(vb.pp(input_norm_pp_name), norm_eps, hidden_size)?;
+        let input_layernorm =
+            get_layer_norm(vb.pp(input_norm_pp_name), norm_eps, hidden_size, true)?;
         let post_attention_layernorm =
-            get_layer_norm(vb.pp(post_norm_pp_name), norm_eps, hidden_size)?;
+            get_layer_norm(vb.pp(post_norm_pp_name), norm_eps, hidden_size, true)?;
         Ok(Self {
             self_attn,
             mlp,
@@ -689,16 +690,26 @@ pub fn get_conv1d(
     Ok(conv1d)
 }
 
-pub fn get_layer_norm(vb: VarBuilder, eps: f64, dim: usize) -> Result<LayerNorm> {
+pub fn get_layer_norm(vb: VarBuilder, eps: f64, dim: usize, affine: bool) -> Result<LayerNorm> {
     let ln_config = LayerNormConfig {
         eps,
         remove_mean: true, // true for layernorm, false for RMSNorm
-        affine: true,      // true for with bias, false for without bias
+        affine,            // true for with bias, false for without bias
     };
     let norm = layer_norm(dim, ln_config, vb)?;
     Ok(norm)
 }
 
+pub fn get_layer_norm_without_weight(
+    vb: VarBuilder,
+    eps: f64,
+    dim: usize,
+) -> Result<LayerNorm> {
+    let weight = Tensor::ones(dim, vb.dtype(), vb.device())?;
+    let bias = Tensor::zeros(dim, vb.dtype(), vb.device())?;
+    Ok(LayerNorm::new(weight, bias, eps))
+}
+
 pub fn get_batch_norm(vb: VarBuilder, eps: f64, dim: usize, affine: bool) -> Result<BatchNorm> {
     let bn_config = BatchNormConfig {
         eps,
@@ -1167,3 +1178,9 @@ impl WNLinear {
         Ok(x)
     }
 }
+
+pub fn mish(xs: &Tensor) -> Result<Tensor> {
+    let tanh = xs.exp()?.affine(1.0, 1.0)?.log()?.tanh()?;
+    let xs = xs.mul(&tanh)?;
+    Ok(xs)
+}

src/models/deepseek_ocr/model.rs

@@ -252,7 +252,7 @@ impl Block {
         window_size: usize,
         input_size: Option<(usize, usize)>,
     ) -> Result<Self> {
-        let norm1 = get_layer_norm(vb.pp("norm1"), eps, dim)?;
+        let norm1 = get_layer_norm(vb.pp("norm1"), eps, dim, true)?;
         let input_size = if window_size == 0 {
             input_size
         } else {
@@ -266,7 +266,7 @@ impl Block {
             use_rel_pos,
             input_size,
         )?;
-        let norm2 = get_layer_norm(vb.pp("norm2"), eps, dim)?;
+        let norm2 = get_layer_norm(vb.pp("norm2"), eps, dim, true)?;
         let mlp_dim = (dim as f32 * mlp_ratio) as usize;
         let mlp = TwoLinearMLP::new(vb.pp("mlp"), dim, mlp_dim, dim, act, true, "lin1", "lin2")?;
         Ok(Self {
@@ -674,8 +674,8 @@ impl NoTPTransformerBlock {
     ) -> Result<Self> {
         let self_attn = NoTPAttention::new(vb.pp("self_attn"), hidden_size, num_heads)?;
         let mlp = NoTPFeedForward::new(vb.pp("mlp"), hidden_size, ffn_hidden_size)?;
-        let layer_norm1 = get_layer_norm(vb.pp("layer_norm1"), eps, hidden_size)?;
+        let layer_norm1 = get_layer_norm(vb.pp("layer_norm1"), eps, hidden_size, true)?;
-        let layer_norm2 = get_layer_norm(vb.pp("layer_norm2"), eps, hidden_size)?;
+        let layer_norm2 = get_layer_norm(vb.pp("layer_norm2"), eps, hidden_size, true)?;
         Ok(Self {
             self_attn,
             mlp,
@@ -764,7 +764,7 @@ impl VitModel {
             ffn_hidden_size,
             eps,
         )?;
-        let pre_layrnorm = get_layer_norm(vb.pp("pre_layrnorm"), eps, hidden_size)?;
+        let pre_layrnorm = get_layer_norm(vb.pp("pre_layrnorm"), eps, hidden_size, true)?;
         Ok(Self {
             embeddings,
             transformer,

src/models/fun_asr_nano/model.rs

@@ -227,8 +227,8 @@ impl EncoderLayerSANM {
             "w_1",
             "w_2",
         )?;
-        let norm1 = get_layer_norm(vb.pp("norm1"), 1e-5, in_dim)?;
+        let norm1 = get_layer_norm(vb.pp("norm1"), 1e-5, in_dim, true)?;
-        let norm2 = get_layer_norm(vb.pp("norm2"), 1e-5, hidden_dim)?;
+        let norm2 = get_layer_norm(vb.pp("norm2"), 1e-5, hidden_dim, true)?;
         let concat_linear = if concat_after {
             let lin = linear(hidden_dim * 2, hidden_dim, vb.pp("concat_linear"))?;
             Some(lin)
@@ -400,8 +400,8 @@ impl SenseVoiceEncoderSmall {
             )?;
             tp_encoders.push(tp_blocks_i);
         }
-        let after_norm = get_layer_norm(vb.pp("after_norm"), 1e-5, output_size)?;
+        let after_norm = get_layer_norm(vb.pp("after_norm"), 1e-5, output_size, true)?;
-        let tp_norm = get_layer_norm(vb.pp("tp_norm"), 1e-5, output_size)?;
+        let tp_norm = get_layer_norm(vb.pp("tp_norm"), 1e-5, output_size, true)?;
         let scaling = (output_size as f64).powf(0.5);
         Ok(Self {
             embed,
@@ -468,8 +468,8 @@ impl AdaptorEncoderLayer {
             "w_1",
             "w_2",
         )?;
-        let norm1 = get_layer_norm(vb.pp("norm1"), 1e-5, llm_dim)?;
+        let norm1 = get_layer_norm(vb.pp("norm1"), 1e-5, llm_dim, true)?;
-        let norm2 = get_layer_norm(vb.pp("norm2"), 1e-5, llm_dim)?;
+        let norm2 = get_layer_norm(vb.pp("norm2"), 1e-5, llm_dim, true)?;
         let concat_linear = if concat_after {
             let lin = linear(llm_dim * 2, llm_dim, vb.pp("concat_linear"))?;
             Some(lin)

src/models/glm_asr_nano/model.rs

@@ -129,11 +129,12 @@ impl GlmAsrEncoderLayer {
             "fc2",
         )?;
         let input_layernorm =
-            get_layer_norm(vb.pp("input_layernorm"), 1e-5, audio_cfg.hidden_size)?;
+            get_layer_norm(vb.pp("input_layernorm"), 1e-5, audio_cfg.hidden_size, true)?;
         let post_attention_layernorm = get_layer_norm(
             vb.pp("post_attention_layernorm"),
             1e-5,
             audio_cfg.hidden_size,
+            true,
         )?;
         Ok(Self {
             self_attn,
@@ -202,7 +203,7 @@ impl GlmAsrEncoder {
             let layer_i = GlmAsrEncoderLayer::new(vb_layers.pp(i), audio_cfg)?;
             layers.push(layer_i);
         }
-        let norm = get_layer_norm(vb.pp("norm"), 1e-5, audio_cfg.hidden_size)?;
+        let norm = get_layer_norm(vb.pp("norm"), 1e-5, audio_cfg.hidden_size, true)?;
         let dim = (audio_cfg.head_dim as f64 * audio_cfg.partial_rotary_factor) as usize;
         let rotary_emb = RoPE::new(dim, audio_cfg.rope_parameters.rope_theta, vb.device())?;
         Ok(Self {

src/models/index_tts2/generate.rs

@@ -1,14 +1,29 @@
 use aha_openai_dive::v1::resources::chat::ChatCompletionParameters;
-use anyhow::Result;
+use anyhow::{Result, anyhow};
 use candle_core::{DType, Device};
+use sentencepiece::SentencePieceProcessor;
 
 use crate::{
-    models::index_tts2::{config::IndexTTS2Config, processor::IndexTTS2Processor},
+    models::index_tts2::{
-    utils::{get_default_save_dir, get_device, get_dtype},
+        config::IndexTTS2Config,
+        model::IndexTTS2Model,
+        processor::IndexTTS2Processor,
+        utils::{TextNormalizer, tokenize_by_cjk_char},
+    },
+    tokenizer::sentencepiece_encode,
+    utils::{
+        audio_utils::extract_audio_url, extract_user_text, get_default_save_dir, get_device,
+        get_dtype,
+    },
 };
+
 pub struct IndexTTS2Generate {
     processor: IndexTTS2Processor,
+    tokenizer: SentencePieceProcessor,
     config: IndexTTS2Config,
+    cache_spk_audio_prompt: Option<String>,
+    model: IndexTTS2Model,
+    device: Device,
 }
 
 impl IndexTTS2Generate {
@@ -18,12 +33,46 @@ impl IndexTTS2Generate {
         let config: IndexTTS2Config = serde_yaml::from_slice(&std::fs::read(config_path)?)?;
         let device = get_device(device);
         let dtype = get_dtype(dtype, "bf16");
-        let processor = IndexTTS2Processor::new(path, &save_dir, &config, &device, dtype)?;
+        let processor = IndexTTS2Processor::new(&device)?;
-        
+        let bpe_path = path.to_string() + "/bpe.model";
-        Ok(Self { config, processor })
+        let tokenizer = SentencePieceProcessor::open(bpe_path)
+            .map_err(|e| anyhow!(format!("load bpe,model file error:{}", e)))?;
+        let model = IndexTTS2Model::new(path, &save_dir, &config, &device, dtype)?;
+        Ok(Self {
+            processor,
+            tokenizer,
+            config,
+            cache_spk_audio_prompt: None,
+            model,
+            device,
+        })
     }
+
+    pub fn use_prompt(&self, mes: &ChatCompletionParameters) -> bool {
+        if let Some(cache) = &self.cache_spk_audio_prompt {
+            let audio_vec = extract_audio_url(mes);
+            if audio_vec.len() == 0 {
+                true
+            } else {
+                if cache.eq(&audio_vec[0]) { true } else { false }
+            }
+        } else {
+            false
+        }
+    }
+
     pub fn generate(&mut self, mes: ChatCompletionParameters) -> Result<()> {
-        let _ = self.processor.process_info(&mes)?;
+        let text = extract_user_text(&mes)?;
+        let text = tokenize_by_cjk_char(&text, true);
+        let input_ids = sentencepiece_encode(&text, &self.tokenizer, &self.device)?;
+        let (audio_22k, audio_16k) = if self.use_prompt(&mes) {
+            (None, None)
+        } else {
+            let (audio_22k, audio_16k, prompt) = self.processor.process_info(&mes)?;
+            self.cache_spk_audio_prompt = Some(prompt);
+            (Some(audio_22k), Some(audio_16k))
+        };
+        let _ = self.model.forward(&input_ids, audio_22k.as_ref(), audio_16k.as_ref())?;
         Ok(())
     }
 }

src/models/index_tts2/model.rs

@@ -1,20 +1,33 @@
-use anyhow::Result;
+use anyhow::{Result, anyhow};
-use candle_core::{D, Tensor};
+use candle_core::{D, DType, Device, IndexOp, Tensor, pickle::read_all_with_key};
 use candle_nn::{
-    Conv1d, Embedding, LayerNorm, Linear, Module, RmsNorm, VarBuilder, embedding, linear, linear_b,
+    Activation, Conv1d, Embedding, GroupNorm, Init, LayerNorm, Linear, Module, RmsNorm, VarBuilder,
-    ops::sigmoid, rms_norm,
+    embedding, group_norm, linear, linear_b, ops::sigmoid, rms_norm,
 };
 
 use crate::{
     models::{
+        campplus::CAMPPlus,
         common::{
             GateUpDownMLP, QKVCatAttention, TwoLinearMLP, WNConv1d, WNLinear, get_conv1d,
-            get_layer_norm,
+            get_layer_norm, get_layer_norm_without_weight, mish,
         },
-        index_tts2::config::{DiTModelArgs, S2MelConfig},
+        feature_extractor::seamless_m4t_feature_extractor::SeamlessM4TFeatureExtractor,
+        index_tts2::config::{DiTModelArgs, IndexTTS2Config, PreprocessParams, S2MelConfig},
+        mask_gct::model::RepCodec,
+        w2v_bert_2_0::model::W2VBert2_0Model,
     },
     position_embed::rope::RoPE,
-    utils::tensor_utils::{pad_reflect_last_dim, split_tensor_with_size},
+    utils::{
+        audio_utils::{
+            create_hann_window, get_waveform_and_window_properties, kaldi_fbank,
+            kaldi_get_mel_banks, mel_filter_bank, torch_stft,
+        },
+        get_vb_model_path, read_pth_tensor_info_cycle,
+        tensor_utils::{
+            interpolate_nearest_1d, pad_reflect_last_dim, sequence_mask, split_tensor_with_size,
+        },
+    },
 };
 pub struct AdaptiveLayerNorm {
     project_layer: Linear,
@@ -219,7 +232,7 @@ impl TimestepEmbedder {
             candle_nn::Activation::Silu,
             true,
             "0",
-            "1",
+            "2",
         )?;
         let scale = 1000.0;
         let half = frequency_embedding_size / 2;
@@ -338,7 +351,7 @@ impl Wavenet {
             let in_layer = SConv1d::new(
                 vb_layers.pp(i),
                 hidden_c,
-                1 * hidden_c,
+                2 * hidden_c,
                 ks,
                 1,
                 dilation,
@@ -385,7 +398,7 @@ impl Wavenet {
         Ok(acts)
     }
 
-    pub fn forward(self, xs: &Tensor, x_mask: &Tensor, g: Option<&Tensor>) -> Result<Tensor> {
+    pub fn forward(&self, xs: &Tensor, x_mask: &Tensor, g: Option<&Tensor>) -> Result<Tensor> {
         let mut output = xs.zeros_like()?;
         let g = if let Some(g) = g
             && let Some(cond_layer) = &self.cond_layer
@@ -432,7 +445,7 @@ impl FinalLayer {
         patch_size: usize,
         out_c: usize,
     ) -> Result<Self> {
-        let norm_final = get_layer_norm(vb.pp("norm_final"), 1e-6, hidden_size)?;
+        let norm_final = get_layer_norm_without_weight(vb.pp("norm_final"), 1e-6, hidden_size)?;
         let linear = WNLinear::new(
             vb.pp("linear"),
             hidden_size,
@@ -485,6 +498,8 @@ pub struct DiT {
     time_as_token: bool,
     style_as_token: bool,
     uvit_skip_connection: bool,
+    transformer_style_condition: bool,
+    long_skip_connection: bool,
 }
 
 impl DiT {
@@ -594,24 +609,451 @@ impl DiT {
             time_as_token,
             style_as_token,
             uvit_skip_connection,
+            transformer_style_condition: config.di_t.style_condition,
+            long_skip_connection: config.di_t.long_skip_connection,
         })
     }
+
+    pub fn forward(
+        &self,
+        xs: &Tensor,
+        prompt_xs: &Tensor,
+        x_lens: &Tensor,
+        t: &Tensor,
+        style: Option<&Tensor>,
+        cond: &Tensor,
+    ) -> Result<Tensor> {
+        let (_, _, t_dim) = xs.dims3()?;
+        let t1 = self.t_embedder.forward(t)?;
+        let cond = self.cond_projection.forward(cond)?;
+        let xs = xs.transpose(1, 2)?;
+        let prompt_xs = prompt_xs.transpose(1, 2)?;
+        let mut x_in = Tensor::cat(&[&xs, &prompt_xs, &cond], D::Minus1)?;
+        if self.transformer_style_condition
+            && !self.style_as_token
+            && let Some(style) = style
+        {
+            let style = style.unsqueeze(1)?.repeat((1, t_dim, 1))?;
+            x_in = Tensor::cat(&[&x_in, &style], D::Minus1)?;
+        }
+        x_in = self.cond_x_merge_linear.forward(&x_in)?;
+        // if self.style_as_token
+        //     && let Some(style_in) = self.style_in.as_ref()
+        // {
+        //     let style = style_in.forward(style)?.unsqueeze(1)?;
+        //     x_in = Tensor::cat(&[&style, &x_in], 1)?;
+        // }
+        // if self.time_as_token {
+        //     let t1 = t1.unsqueeze(1)?;
+        //     x_in = Tensor::cat(&[&t1, &x_in], 1)?;
+        // }
+        // let mut x_lens = x_lens.clone();
+        // if self.style_as_token {
+        //     x_lens = x_lens.affine(1.0, 1.0)?;
+        // }
+        // if self.time_as_token {
+        //     x_lens = x_lens.affine(1.0, 0.0)?;
+        // }
+        let x_mask = sequence_mask(&x_lens, Some(x_in.dim(1)? as u32))?
+            .to_device(xs.device())?
+            .unsqueeze(1)?;
+        let mut x_res = self.transformer.forward(&x_in, &t1.unsqueeze(1)?, None)?;
+        // if self.time_as_token {
+        //     let last_dim = x_res.dim(D::Minus1)?;
+        //     x_res = x_res.narrow(D::Minus1, 1, last_dim - 1)?;
+        // }
+        // if self.style_as_token {
+        //     let last_dim = x_res.dim(D::Minus1)?;
+        //     x_res = x_res.narrow(D::Minus1, 1, last_dim - 1)?;
+        // }
+        if self.long_skip_connection {
+            x_res = self
+                .skip_linear
+                .forward(&Tensor::cat(&[&x_res, &xs], D::Minus1)?)?;
+        }
+        let xs = self.conv1.forward(&x_res)?;
+        let xs = xs.transpose(1, 2)?;
+        let t2 = self.t_embedder2.forward(t)?;
+        let xs = self
+            .wavenet
+            .forward(&xs, &x_mask, Some(&t2.unsqueeze(2)?))?
+            .transpose(1, 2)?
+            .broadcast_add(&self.res_projection.forward(&x_res)?)?;
+        let xs = self.final_layer.forward(&xs, &t1)?.transpose(1, 2)?;
+        let xs = self.conv2.forward(&xs)?;
+        Ok(xs)
+    }
 }
 
 pub struct CFM {
+    in_channels: usize,
     estimator: DiT,
+    // criterion: l1Loss
+    sigma_min: f32,
+}
+
+impl CFM {
+    pub fn new(vb: VarBuilder, config: &S2MelConfig) -> Result<Self> {
+        let in_channels = config.di_t.in_channels;
+        let sigma_min = 1e-6;
+        let estimator = DiT::new(vb.pp("estimator"), config)?;
+        Ok(Self {
+            in_channels,
+            estimator,
+            sigma_min,
+        })
+    }
+}
+
+pub struct InterpolateModule {
+    conv1d: Conv1d,
+    norm: GroupNorm,
+}
+
+impl InterpolateModule {
+    pub fn new(vb: &VarBuilder, index: usize, channels: usize, groups: usize) -> Result<Self> {
+        let start_index = index * 3;
+        let conv1d = get_conv1d(vb.pp(start_index), channels, channels, 3, 1, 1, 1, 1, true)?;
+        let norm = group_norm(groups, channels, 1e-5, vb.pp(start_index + 1))?;
+        Ok(Self { conv1d, norm })
+    }
+
+    pub fn forward(&self, xs: &Tensor) -> Result<Tensor> {
+        let xs = self.conv1d.forward(xs)?;
+        let xs = self.norm.forward(&xs)?;
+        let xs = mish(&xs)?;
+        Ok(xs)
+    }
+}
+
+pub struct InterpolateRegulator {
+    sampling_ratios: Vec<usize>,
+    out_channels: usize,
+    model0_11: Vec<InterpolateModule>,
+    model_12: Conv1d,
+    embedding: Embedding,
+    mask_token: Tensor,
+    quantizer_dropout: f32,
+    content_in_proj: Linear,
+    n_codebooks: usize,
+    interpolate: bool,
+}
+
+impl InterpolateRegulator {
+    pub fn new(
+        vb: VarBuilder,
+        channels: usize,
+        sampling_ratios: Vec<usize>,
+        is_discrete: bool,
+        in_channels: usize,
+        vector_quantize: bool,
+        codebook_size: usize,
+        out_channels: Option<usize>,
+        groups: usize,
+        n_codebooks: usize,
+        quantizer_dropout: f32,
+        f0_condition: bool,
+        n_f0_bins: usize,
+    ) -> Result<Self> {
+        let out_channels = out_channels.unwrap_or(channels);
+        let vb_model = vb.pp("model");
+        let interpolate = true;
+        let mut model0_11 = vec![];
+        for (index, _) in sampling_ratios.iter().enumerate() {
+            let inter = InterpolateModule::new(&vb_model, index, channels, groups)?;
+            model0_11.push(inter);
+        }
+        let model_12 = get_conv1d(
+            vb_model.pp("12"),
+            channels,
+            out_channels,
+            1,
+            0,
+            1,
+            1,
+            1,
+            true,
+        )?;
+        let embedding = embedding(codebook_size, channels, vb.pp("embedding"))?;
+        let mask_token = vb.get_with_hints((1, channels), "mask_token", Init::Const(0.0))?;
+        let content_in_proj = linear(in_channels, channels, vb.pp("content_in_proj"))?;
+        Ok(Self {
+            sampling_ratios,
+            out_channels,
+            model0_11,
+            model_12,
+            embedding,
+            mask_token,
+            quantizer_dropout,
+            content_in_proj,
+            n_codebooks,
+            interpolate,
+        })
+    }
+    pub fn forward(&self, x: &Tensor, y_lens: &Tensor) -> Result<Tensor> {
+        let mut xs = self.content_in_proj.forward(x)?;
+        xs = xs.transpose(1, 2)?.contiguous()?;
+        if self.interpolate {
+            let size = y_lens.max_all()?.to_scalar::<u32>()? as usize;
+            xs = interpolate_nearest_1d(&xs, size)?;
+        }
+        for model_i in self.model0_11.iter() {
+            xs = model_i.forward(&xs)?;
+        }
+        xs = self.model_12.forward(&xs)?.transpose(1, 2)?.contiguous()?;
+        Ok(xs)
+    }
 }
 
 pub struct MyModel {
     cfm: CFM,
+    length_regulator: InterpolateRegulator,
 }
 
-pub struct IndexTTS2 {
+impl MyModel {
-    cache_spk_cond: Option<Tensor>,
+    pub fn new(
-    cache_s2mel_style: Option<Tensor>,
+        model_path: &str,
-    cache_s2mel_prompt: Option<Tensor>,
+        config: &S2MelConfig,
-    cache_spk_audio_prompt: Option<String>,
+        dtype: DType,
-    cache_emo_cond: Option<Tensor>,
+        device: &Device,
-    cache_emo_audio_prompt: Option<Tensor>,
+    ) -> Result<Self> {
-    cache_mel: Option<Tensor>,
+        let s2mel_path = model_path.to_string() + "/s2mel.pth";
+        let length_regulator_dict =
+            read_pth_tensor_info_cycle(s2mel_path.clone(), Some("net.length_regulator"))?;
+        let length_regulator_vb = VarBuilder::from_tensors(length_regulator_dict, dtype, device);
+        let length_regulator = InterpolateRegulator::new(
+            length_regulator_vb,
+            config.length_regulator.channels,
+            config.length_regulator.sampling_ratios.clone(),
+            config.length_regulator.is_discrete,
+            config.length_regulator.in_channels,
+            config.length_regulator.vector_quantize,
+            config.length_regulator.content_codebook_size,
+            None,
+            1,
+            config.length_regulator.n_codebooks,
+            config.length_regulator.quantizer_dropout,
+            config.length_regulator.f0_condition,
+            config.length_regulator.n_f0_bins,
+        )?;
+        let cfm_dict = read_pth_tensor_info_cycle(s2mel_path.clone(), Some("net.cfm"))?;
+        let cfm_vb = VarBuilder::from_tensors(cfm_dict, dtype, device);
+        let cfm = CFM::new(cfm_vb, config)?;
+        Ok(Self {
+            cfm,
+            length_regulator,
+        })
+    }
+
+    pub fn length_regulator_forward(
+        &self,
+        s_ori: &Tensor,
+        target_lengths: &Tensor,
+    ) -> Result<Tensor> {
+        let xs = self.length_regulator.forward(s_ori, target_lengths)?;
+        Ok(xs)
+    }
+}
+
+pub struct IndexTTS2Cache {
+    pub cache_spk_cond: Tensor,
+    pub cache_s2mel_style: Tensor,
+    pub cache_s2mel_prompt: Tensor,
+    pub cache_mel: Tensor,
+    // pub cache_emo_cond: Tensor,
+    // pub cache_emo_audio_prompt: Tensor,
+}
+
+pub struct IndexTTS2Model {
+    cache: Option<IndexTTS2Cache>,
+    feature_extractor: SeamlessM4TFeatureExtractor,
+    semantic_model: W2VBert2_0Model,
+    semantic_mean: Tensor,
+    semantic_std: Tensor,
+    semantic_codec: RepCodec,
+    s2mel_filters: Tensor,
+    s2mel_windows: Tensor,
+    s2mel_preprocess_params: PreprocessParams,
+    window_shift: usize,
+    window_size: usize,
+    padded_window_size: usize,
+    mel_energies: Tensor,
+    campplus_model: CAMPPlus,
+    s2mel: MyModel,
+}
+
+impl IndexTTS2Model {
+    pub fn new(
+        path: &str,
+        save_dir: &str,
+        config: &IndexTTS2Config,
+        device: &Device,
+        dtype: DType,
+    ) -> Result<Self> {
+        let feature_extractor = SeamlessM4TFeatureExtractor::new(
+            // 80,
+            80,
+            crate::utils::tensor_utils::PaddingSide::Right,
+            1.0,
+            16000,
+            2,
+            device,
+        )?;
+        let w2vbert2_path = save_dir.to_string() + "/facebook/w2v-bert-2.0";
+        let semantic_model = W2VBert2_0Model::init(&w2vbert2_path, device, dtype)?;
+        let semantic_mean_var_path = path.to_string() + "/" + &config.w2v_stat;
+        let dict = read_all_with_key(semantic_mean_var_path, None)?;
+        let mut semantic_mean = Tensor::new(0.0, device)?.to_dtype(dtype)?;
+        let mut semantic_std = Tensor::new(1.0, device)?.to_dtype(dtype)?;
+        for (k, v) in dict {
+            if k.eq("mean") {
+                semantic_mean = v.to_device(device)?.to_dtype(dtype)?;
+            } else if k.eq("var") {
+                semantic_std = v.to_device(device)?.to_dtype(dtype)?.sqrt()?;
+            }
+        }
+
+        let semantic_codec_path =
+            save_dir.to_string() + "/amphion/MaskGCT/semantic_codec/model.safetensors";
+        let vb =
+            unsafe { VarBuilder::from_mmaped_safetensors(&[semantic_codec_path], dtype, &device)? };
+        let semantic_codec = RepCodec::new(vb, &config.semantic_codec)?;
+        let s2mel_filters = mel_filter_bank(
+            config.s2mel.preprocess_params.spect_params.n_fft / 2 + 1,
+            config.s2mel.preprocess_params.spect_params.n_mels,
+            config.s2mel.preprocess_params.spect_params.fmin as f32,
+            config
+                .s2mel
+                .preprocess_params
+                .spect_params
+                .fmax
+                .unwrap_or(config.s2mel.preprocess_params.sr / 2) as f32,
+            config.s2mel.preprocess_params.sr as f32,
+            Some("slaney"),
+            crate::utils::audio_utils::MelScale::Slaney,
+            false,
+            device,
+        )?
+        .t()?;
+        let s2mel_windows = create_hann_window(
+            config.s2mel.preprocess_params.spect_params.win_length,
+            dtype,
+            device,
+        )?;
+        let (window_shift, window_size, padded_window_size) =
+            get_waveform_and_window_properties(16000, 10.0, 25.0, true)?;
+        let (mel_energies, _) =
+            kaldi_get_mel_banks(80, padded_window_size, 16000 as f32, 20.0, 0.0, device)?;
+        let mel_energies = mel_energies.pad_with_zeros(D::Minus1, 0, 1)?.t()?;
+        let campplus_model_path = save_dir.to_string()
+            + "/iic/speech_campplus_sv_zh-cn_16k-common/campplus_cn_common.bin";
+        let campplus_vb = get_vb_model_path(campplus_model_path, dtype, device.clone(), None)?;
+        let campplus_model = CAMPPlus::new(campplus_vb, 80, 192, 32, 4, 128)?;
+        let s2mel = MyModel::new(path, &config.s2mel, dtype, device)?;
+        Ok(Self {
+            cache: None,
+            feature_extractor,
+            semantic_model,
+            semantic_mean,
+            semantic_std,
+            semantic_codec,
+            s2mel_filters,
+            s2mel_windows,
+            s2mel_preprocess_params: config.s2mel.preprocess_params.clone(),
+            window_shift,
+            window_size,
+            padded_window_size,
+            mel_energies,
+            campplus_model,
+            s2mel,
+        })
+    }
+
+    pub fn get_emb(
+        &self,
+        input_features: &Tensor,
+        attention_mask: Option<&Tensor>,
+    ) -> Result<Tensor> {
+        let output =
+            self.semantic_model
+                .forward(input_features, attention_mask, Some(17), false)?;
+        let feature = &output.specify_layer_id_hidden_state.unwrap();
+        let feature = feature
+            .broadcast_sub(&self.semantic_mean)?
+            .broadcast_div(&self.semantic_std)?;
+        Ok(feature)
+    }
+
+    pub fn s2mel_spectrogram(&self, waveform: &Tensor) -> Result<Tensor> {
+        let pad = (self.s2mel_preprocess_params.spect_params.n_fft
+            - self.s2mel_preprocess_params.spect_params.hop_length)
+            / 2;
+        let pad_audio_22k = pad_reflect_last_dim(&waveform, (pad, pad))?;
+        let spec = torch_stft(
+            &pad_audio_22k,
+            self.s2mel_preprocess_params.spect_params.n_fft,
+            self.s2mel_preprocess_params.spect_params.hop_length,
+            &self.s2mel_windows,
+        )?
+        .transpose(1, 2)?;
+        let spec = self.s2mel_filters.broadcast_matmul(&spec)?;
+        let spec = spec.clamp(1e-5, f64::INFINITY)?.log()?;
+        Ok(spec)
+    }
+
+    pub fn forward(
+        &mut self,
+        input_ids: &Tensor,
+        audio_22k: Option<&Tensor>,
+        audio_16k: Option<&Tensor>,
+    ) -> Result<()> {
+        if (audio_22k.is_none() || audio_16k.is_none()) && self.cache.is_none() {
+            return Err(anyhow!(
+                "Missing required audio input: must provide either audio_22k, audio_16k, or have cached prompt data available"
+            ));
+        }
+        let (spk_cond_emb, style, prompt_condition, ref_mel) = if let Some(audio_22k) = audio_22k
+            && let Some(audio_16k) = audio_16k
+        {
+            let (audio_16k_features, audio_16k_mask) =
+                self.feature_extractor.call(&audio_16k, 16000, true, true)?;
+            let spk_cond_emb = self.get_emb(&audio_16k_features, audio_16k_mask.as_ref())?;
+            let (_, s_ref) = self.semantic_codec.quantize(&spk_cond_emb)?;
+            let ref_mel = self.s2mel_spectrogram(&audio_22k)?;
+            let ref_target_lengths = Tensor::new(ref_mel.dim(2)? as u32, ref_mel.device())?;
+            let feat = kaldi_fbank(
+                &audio_16k,
+                &self.mel_energies,
+                self.window_shift,
+                self.window_size,
+                self.padded_window_size,
+                0.0,
+            )?;
+            let feat = feat.broadcast_sub(&feat.mean_keepdim(1)?)?;
+            let style = self.campplus_model.forward(&feat)?;
+            let prompt_condition = self
+                .s2mel
+                .length_regulator
+                .forward(&s_ref, &ref_target_lengths)?;
+            let cache = IndexTTS2Cache {
+                cache_spk_cond: spk_cond_emb.clone(),
+                cache_s2mel_style: style.clone(),
+                cache_s2mel_prompt: prompt_condition.clone(),
+                cache_mel: ref_mel.clone(),
+            };
+            self.cache = Some(cache);
+            (spk_cond_emb, style, prompt_condition, ref_mel)
+        } else {
+            let cache = self.cache.as_ref().unwrap();
+            (
+                cache.cache_spk_cond.clone(),
+                cache.cache_s2mel_style.clone(),
+                cache.cache_s2mel_prompt.clone(),
+                cache.cache_mel.clone(),
+            )
+        };
+        
+        
+        Ok(())
+    }
 }

src/models/index_tts2/processor.rs

@@ -5,13 +5,16 @@ use candle_nn::VarBuilder;
 
 use crate::{
     models::{
-        campplus::CAMPPlus, feature_extractor::seamless_m4t_feature_extractor::SeamlessM4TFeatureExtractor, index_tts2::config::{IndexTTS2Config, PreprocessParams}, mask_gct::model::RepCodec, w2v_bert_2_0::model::W2VBert2_0Model
+        campplus::CAMPPlus,
+        feature_extractor::seamless_m4t_feature_extractor::SeamlessM4TFeatureExtractor,
+        index_tts2::config::{IndexTTS2Config, PreprocessParams},
+        mask_gct::model::RepCodec,
+        w2v_bert_2_0::model::W2VBert2_0Model,
     },
     utils::{
         audio_utils::{
             create_hann_window, extract_audio_url, get_waveform_and_window_properties, kaldi_fbank,
-            kaldi_get_mel_banks, load_audio, mel_filter_bank, resample_simple, 
+            kaldi_get_mel_banks, load_audio, mel_filter_bank, resample_simple, torch_stft,
-            torch_stft,
         },
         get_vb_model_path,
         tensor_utils::pad_reflect_last_dim,
@@ -21,104 +24,104 @@ use crate::{
 pub struct IndexTTS2Processor {
     device: Device,
     max_audio_length_seconds: usize,
-    feature_extractor: SeamlessM4TFeatureExtractor,
+    // feature_extractor: SeamlessM4TFeatureExtractor,
-    semantic_model: W2VBert2_0Model,
+    // semantic_model: W2VBert2_0Model,
-    semantic_mean: Tensor,
+    // semantic_mean: Tensor,
-    semantic_std: Tensor,
+    // semantic_std: Tensor,
-    semantic_codec: RepCodec,
+    // semantic_codec: RepCodec,
-    s2mel_filters: Tensor,
+    // s2mel_filters: Tensor,
-    s2mel_windows: Tensor,
+    // s2mel_windows: Tensor,
-    s2mel_preprocess_params: PreprocessParams,
+    // s2mel_preprocess_params: PreprocessParams,
-    window_shift: usize,
+    // window_shift: usize,
-    window_size: usize,
+    // window_size: usize,
-    padded_window_size: usize,
+    // padded_window_size: usize,
-    mel_energies: Tensor,
+    // mel_energies: Tensor,
-    campplus_model: CAMPPlus,
+    // campplus_model: CAMPPlus,
 }
 
 impl IndexTTS2Processor {
     pub fn new(
-        path: &str,
+        // path: &str,
-        save_dir: &str,
+        // save_dir: &str,
-        config: &IndexTTS2Config,
+        // config: &IndexTTS2Config,
         device: &Device,
-        dtype: DType,
+        // dtype: DType,
     ) -> Result<Self> {
-        let feature_extractor = SeamlessM4TFeatureExtractor::new(
+        // let feature_extractor = SeamlessM4TFeatureExtractor::new(
+        //     // 80,
         //     80,
-            80,
+        //     crate::utils::tensor_utils::PaddingSide::Right,
-            crate::utils::tensor_utils::PaddingSide::Right,
+        //     1.0,
-            1.0,
+        //     16000,
-            16000,
+        //     2,
-            2,
+        //     device,
-            device,
+        // )?;
-        )?;
+        // let w2vbert2_path = save_dir.to_string() + "/facebook/w2v-bert-2.0";
-        let w2vbert2_path = save_dir.to_string() + "/facebook/w2v-bert-2.0";
+        // let semantic_model = W2VBert2_0Model::init(&w2vbert2_path, device, dtype)?;
-        let semantic_model = W2VBert2_0Model::init(&w2vbert2_path, device, dtype)?;
+        // let semantic_mean_var_path = path.to_string() + "/" + &config.w2v_stat;
-        let semantic_mean_var_path = path.to_string() + "/" + &config.w2v_stat;
+        // let dict = read_all_with_key(semantic_mean_var_path, None)?;
-        let dict = read_all_with_key(semantic_mean_var_path, None)?;
+        // let mut semantic_mean = Tensor::new(0.0, device)?.to_dtype(dtype)?;
-        let mut semantic_mean = Tensor::new(0.0, device)?.to_dtype(dtype)?;
+        // let mut semantic_std = Tensor::new(1.0, device)?.to_dtype(dtype)?;
-        let mut semantic_std = Tensor::new(1.0, device)?.to_dtype(dtype)?;
+        // for (k, v) in dict {
-        for (k, v) in dict {
+        //     if k.eq("mean") {
-            if k.eq("mean") {
+        //         semantic_mean = v.to_device(device)?.to_dtype(dtype)?;
-                semantic_mean = v.to_device(device)?.to_dtype(dtype)?;
+        //     } else if k.eq("var") {
-            } else if k.eq("var") {
+        //         semantic_std = v.to_device(device)?.to_dtype(dtype)?.sqrt()?;
-                semantic_std = v.to_device(device)?.to_dtype(dtype)?.sqrt()?;
+        //     }
-            }
+        // }
-        }
 
-        let semantic_codec_path =
+        // let semantic_codec_path =
-            save_dir.to_string() + "/amphion/MaskGCT/semantic_codec/model.safetensors";
+        //     save_dir.to_string() + "/amphion/MaskGCT/semantic_codec/model.safetensors";
-        let vb =
+        // let vb =
-            unsafe { VarBuilder::from_mmaped_safetensors(&[semantic_codec_path], dtype, &device)? };
+        //     unsafe { VarBuilder::from_mmaped_safetensors(&[semantic_codec_path], dtype, &device)? };
-        let semantic_codec = RepCodec::new(vb, &config.semantic_codec)?;
+        // let semantic_codec = RepCodec::new(vb, &config.semantic_codec)?;
-        let s2mel_filters = mel_filter_bank(
+        // let s2mel_filters = mel_filter_bank(
-            config.s2mel.preprocess_params.spect_params.n_fft / 2 + 1,
+        //     config.s2mel.preprocess_params.spect_params.n_fft / 2 + 1,
-            config.s2mel.preprocess_params.spect_params.n_mels,
+        //     config.s2mel.preprocess_params.spect_params.n_mels,
-            config.s2mel.preprocess_params.spect_params.fmin as f32,
+        //     config.s2mel.preprocess_params.spect_params.fmin as f32,
-            config
+        //     config
-                .s2mel
+        //         .s2mel
-                .preprocess_params
+        //         .preprocess_params
-                .spect_params
+        //         .spect_params
-                .fmax
+        //         .fmax
-                .unwrap_or(config.s2mel.preprocess_params.sr / 2) as f32,
+        //         .unwrap_or(config.s2mel.preprocess_params.sr / 2) as f32,
-            config.s2mel.preprocess_params.sr as f32,
+        //     config.s2mel.preprocess_params.sr as f32,
-            Some("slaney"),
+        //     Some("slaney"),
-            crate::utils::audio_utils::MelScale::Slaney,
+        //     crate::utils::audio_utils::MelScale::Slaney,
-            false,
+        //     false,
-            device,
+        //     device,
-        )?
+        // )?
-        .t()?;
+        // .t()?;
-        let s2mel_windows = create_hann_window(
+        // let s2mel_windows = create_hann_window(
-            config.s2mel.preprocess_params.spect_params.win_length,
+        //     config.s2mel.preprocess_params.spect_params.win_length,
-            dtype,
+        //     dtype,
-            device,
+        //     device,
-        )?;
+        // )?;
-        let (window_shift, window_size, padded_window_size) =
+        // let (window_shift, window_size, padded_window_size) =
-            get_waveform_and_window_properties(16000, 10.0, 25.0, true)?;
+        //     get_waveform_and_window_properties(16000, 10.0, 25.0, true)?;
-        let (mel_energies, _) =
+        // let (mel_energies, _) =
-            kaldi_get_mel_banks(80, padded_window_size, 16000 as f32, 20.0, 0.0, device)?;
+        //     kaldi_get_mel_banks(80, padded_window_size, 16000 as f32, 20.0, 0.0, device)?;
-        let mel_energies = mel_energies.pad_with_zeros(D::Minus1, 0, 1)?.t()?;
+        // let mel_energies = mel_energies.pad_with_zeros(D::Minus1, 0, 1)?.t()?;
-        let campplus_model_path = save_dir.to_string()
+        // let campplus_model_path = save_dir.to_string()
-            + "/iic/speech_campplus_sv_zh-cn_16k-common/campplus_cn_common.bin";
+        //     + "/iic/speech_campplus_sv_zh-cn_16k-common/campplus_cn_common.bin";
-        let campplus_vb = get_vb_model_path(campplus_model_path, dtype, device.clone(), None)?;
+        // let campplus_vb = get_vb_model_path(campplus_model_path, dtype, device.clone(), None)?;
-        let campplus_model = CAMPPlus::new(campplus_vb, 80, 192, 32, 4, 128)?;
+        // let campplus_model = CAMPPlus::new(campplus_vb, 80, 192, 32, 4, 128)?;
         Ok(Self {
             device: device.clone(),
             max_audio_length_seconds: 15,
-            feature_extractor,
+            // feature_extractor,
-            semantic_model,
+            // semantic_model,
-            semantic_mean,
+            // semantic_mean,
-            semantic_std,
+            // semantic_std,
-            semantic_codec,
+            // semantic_codec,
-            s2mel_filters,
+            // s2mel_filters,
-            s2mel_windows,
+            // s2mel_windows,
-            s2mel_preprocess_params: config.s2mel.preprocess_params.clone(),
+            // s2mel_preprocess_params: config.s2mel.preprocess_params.clone(),
-            window_shift,
+            // window_shift,
-            window_size,
+            // window_size,
-            padded_window_size,
+            // padded_window_size,
-            mel_energies,
+            // mel_energies,
-            campplus_model,
+            // campplus_model,
         })
     }
 
@@ -137,63 +140,64 @@ impl IndexTTS2Processor {
         &self,
         mes: &ChatCompletionParameters,
         device: &Device,
-    ) -> Result<(Tensor, usize)> {
+    ) -> Result<(Tensor, usize, String)> {
         let audio_url_vec = extract_audio_url(mes);
         let (audio, sr) = load_audio(&audio_url_vec[0], device)?;
-        self.cut_audio(&audio, sr)
+        let (audio, sr) = self.cut_audio(&audio, sr)?;
+        Ok((audio, sr, audio_url_vec[0].clone()))
     }
 
-    pub fn get_emb(
+    // pub fn get_emb(
-        &self,
+    //     &self,
-        input_features: &Tensor,
+    //     input_features: &Tensor,
-        attention_mask: Option<&Tensor>,
+    //     attention_mask: Option<&Tensor>,
-    ) -> Result<Tensor> {
+    // ) -> Result<Tensor> {
-        let output =
+    //     let output =
-            self.semantic_model
+    //         self.semantic_model
-                .forward(input_features, attention_mask, Some(17), false)?;
+    //             .forward(input_features, attention_mask, Some(17), false)?;
-        let feature = &output.specify_layer_id_hidden_state.unwrap();
+    //     let feature = &output.specify_layer_id_hidden_state.unwrap();
-        let feature = feature
+    //     let feature = feature
-            .broadcast_sub(&self.semantic_mean)?
+    //         .broadcast_sub(&self.semantic_mean)?
-            .broadcast_div(&self.semantic_std)?;
+    //         .broadcast_div(&self.semantic_std)?;
-        Ok(feature)
+    //     Ok(feature)
-    }
+    // }
 
-    pub fn s2mel_spectrogram(&self, waveform: &Tensor) -> Result<Tensor> {
+    // pub fn s2mel_spectrogram(&self, waveform: &Tensor) -> Result<Tensor> {
-        let pad = (self.s2mel_preprocess_params.spect_params.n_fft
+    //     let pad = (self.s2mel_preprocess_params.spect_params.n_fft
-            - self.s2mel_preprocess_params.spect_params.hop_length)
+    //         - self.s2mel_preprocess_params.spect_params.hop_length)
-            / 2;
+    //         / 2;
-        let pad_audio_22k = pad_reflect_last_dim(&waveform, (pad, pad))?;
+    //     let pad_audio_22k = pad_reflect_last_dim(&waveform, (pad, pad))?;
-        let spec = torch_stft(
+    //     let spec = torch_stft(
-            &pad_audio_22k,
+    //         &pad_audio_22k,
-            self.s2mel_preprocess_params.spect_params.n_fft,
+    //         self.s2mel_preprocess_params.spect_params.n_fft,
-            self.s2mel_preprocess_params.spect_params.hop_length,
+    //         self.s2mel_preprocess_params.spect_params.hop_length,
-            &self.s2mel_windows,
+    //         &self.s2mel_windows,
-        )?
+    //     )?
-        .transpose(1, 2)?;
+    //     .transpose(1, 2)?;
-        let spec = self.s2mel_filters.broadcast_matmul(&spec)?;
+    //     let spec = self.s2mel_filters.broadcast_matmul(&spec)?;
-        let spec = spec.clamp(1e-5, f64::INFINITY)?.log()?;
+    //     let spec = spec.clamp(1e-5, f64::INFINITY)?.log()?;
-        Ok(spec)
+    //     Ok(spec)
-    }
+    // }
-    pub fn process_info(&self, mes: &ChatCompletionParameters) -> Result<()> {
+    pub fn process_info(&self, mes: &ChatCompletionParameters) -> Result<(Tensor, Tensor, String)> {
-        let (audio, sr) = self.extract_audio_and_cut(mes, &self.device)?;
+        let (audio, sr, audio_url) = self.extract_audio_and_cut(mes, &self.device)?;
         let audio_22k = resample_simple(&audio, sr as i64, 22050)?;
         let audio_16k = resample_simple(&audio, sr as i64, 16000)?;
-        let (audio_16k_features, audio_16k_mask) =
+        // let (audio_16k_features, audio_16k_mask) =
-            self.feature_extractor.call(&audio_16k, 16000, true, true)?;
+        //     self.feature_extractor.call(&audio_16k, 16000, true, true)?;
-        let spk_cond_emb = self.get_emb(&audio_16k_features, audio_16k_mask.as_ref())?;
+        // let spk_cond_emb = self.get_emb(&audio_16k_features, audio_16k_mask.as_ref())?;
-        let (_, s_ref) = self.semantic_codec.quantize(&spk_cond_emb)?;
+        // let (_, s_ref) = self.semantic_codec.quantize(&spk_cond_emb)?;
-        let ref_mel = self.s2mel_spectrogram(&audio_22k)?;
+        // let ref_mel = self.s2mel_spectrogram(&audio_22k)?;
-        let feat = kaldi_fbank(
+        // let feat = kaldi_fbank(
-            &audio_16k,
+        //     &audio_16k,
-            &self.mel_energies,
+        //     &self.mel_energies,
-            self.window_shift,
+        //     self.window_shift,
-            self.window_size,
+        //     self.window_size,
-            self.padded_window_size,
+        //     self.padded_window_size,
-            0.0,
+        //     0.0,
-        )?;
+        // )?;
-        let feat = feat.broadcast_sub(&feat.mean_keepdim(1)?)?;
+        // let feat = feat.broadcast_sub(&feat.mean_keepdim(1)?)?;
-        let style = self.campplus_model.forward(&feat)?;
+        // let style = self.campplus_model.forward(&feat)?;
-        println!("style: {}", style);
+        // println!("style: {}", style);
-        Ok(())
+        Ok((audio_22k, audio_16k, audio_url))
     }
 }

src/models/index_tts2/utils.rs

@@ -1,13 +1,18 @@
+use std::collections::HashMap;
+
+use anyhow::{Result, anyhow};
+use regex::Regex;
+
 use crate::utils::{download_model, get_default_save_dir};
 
-pub async fn download_index_tts2_need_model(save_dir: Option<&str>) -> anyhow::Result<()> {
+pub async fn download_index_tts2_need_model(save_dir: Option<&str>) -> Result<()> {
     let save_dir = match save_dir {
         Some(dir) => dir.to_string(),
         None => get_default_save_dir().expect("Failed to get home directory"),
     };
 
     let w2v_bert2_0 = "facebook/w2v-bert-2.0";
-    let mask_gct= "amphion/MaskGCT";
+    let mask_gct = "amphion/MaskGCT";
     // let campplus= "funasr/campplus"; // huggingface
     let campplus = "iic/speech_campplus_sv_zh-cn_16k-common"; // modelscope
     download_model(w2v_bert2_0, &save_dir, 3).await?;
@@ -16,3 +21,151 @@ pub async fn download_index_tts2_need_model(save_dir: Option<&str>) -> anyhow::R
 
     Ok(())
 }
+
+#[derive(Debug, Clone)]
+pub struct TextNormalizer {
+    // char_rep_map: HashMap<String, String>,
+    // zh_char_rep_map: HashMap<String, String>,
+    pinyin_tone_pattern: Regex,
+    // name_pattern: Regex,
+    // tech_term_pattern: Regex,
+    // english_contraction_pattern: Regex,
+    email_pattern: Regex,
+    cjk_range_pattern: Regex,
+}
+
+impl TextNormalizer {
+    pub fn new() -> Result<Self> {
+        let pinyin_tone_pattern: Regex = Regex::new(
+                // r"(?i)(?<![a-z])((?:[bpmfdtnlgkhjqxzcsryw]|[zcs]h)?(?:[aeiouüvAEIOUV]|[aeAE]i|u[aiouAIUO]|aoAO|ouOU|i[aeuAEU]|[uüvUÜV]e|[uvüUVÜ]ang?|uaiUAI|[aeiuvAEIUV]n|[aeioAEIO]ng|ia[noNAO]|i[aA][oO]ng)|ngNG|erER)([1-5])"
+                r"(?i)((?:[bpmfdtnlgkhjqxzcsryw]|[zcs]h)?(?:[aeiouüvAEIOUV]|[aeAE]i|u[aiouAIUO]|aoAO|ouOU|i[aeuAEU]|[uüvUÜV]e|[uvüUVÜ]ang?|uaiUAI|[aeiuvAEIUV]n|[aeioAEIO]ng|ia[noNAO]|i[aA][oO]ng)|ngNG|erER)([1-5])"
+            ).map_err(|e| anyhow!(format!("new pinyin_tone_pattern regex error:{}", e)))?;
+        // let name_pattern: Regex =
+        //     Regex::new(r"[\u{4e00}-\u{9fff}]+(?:[-·—][\u{4e00}-\u{9fff}]+){1,2}")
+        //         .map_err(|e| anyhow!(format!("new name_pattern regex error:{}", e)))?;
+        // let tech_term_pattern: Regex = Regex::new(r"[A-Za-z][A-Za-z0-9]*(?:-[A-Za-z0-9]+)+")
+        //     .map_err(|e| anyhow!(format!("new tech_term_pattern regex error:{}", e)))?;
+        // let english_contraction_pattern: Regex = Regex::new(
+        //     r"(?i)(what|where|who|which|how|t?here|it|s?he|that|this)'s",
+        // )
+        // .map_err(|e| anyhow!(format!("new english_contraction_pattern regex error:{}", e)))?;
+        let email_pattern: Regex = Regex::new(r"^[a-zA-Z0-9]+@[a-zA-Z0-9]+\.[a-zA-Z]+$")
+            .map_err(|e| anyhow!(format!("new email_pattern regex error:{}", e)))?;
+        let cjk_range_pattern: Regex = Regex::new(r"([\u{1100}-\u{11ff}\u{2e80}-\u{a4cf}\u{a840}-\u{d7af}\u{f900}-\u{faff}\u{fe30}-\u{fe4f}\u{ff65}-\u{ffdc}\u{20000}-\u{2ffff}])")
+            .map_err(|e| anyhow!(format!("new cjk_range_pattern regex error:{}", e)))?;
+        // let mut char_rep_map = HashMap::new();
+        // char_rep_map.insert("：".to_string(), ",".to_string());
+        // char_rep_map.insert("；".to_string(), ",".to_string());
+        // char_rep_map.insert(";".to_string(), ",".to_string());
+        // char_rep_map.insert("，".to_string(), ",".to_string());
+        // char_rep_map.insert("。".to_string(), ".".to_string());
+        // char_rep_map.insert("！".to_string(), "!".to_string());
+        // char_rep_map.insert("？".to_string(), "?".to_string());
+        // char_rep_map.insert("\n".to_string(), " ".to_string());
+        // char_rep_map.insert("·".to_string(), "-".to_string());
+        // char_rep_map.insert("、".to_string(), ",".to_string());
+        // char_rep_map.insert("...".to_string(), "…".to_string());
+        // char_rep_map.insert(",,,".to_string(), "…".to_string());
+        // char_rep_map.insert("，，，".to_string(), "…".to_string());
+        // char_rep_map.insert("……".to_string(), "…".to_string());
+        // char_rep_map.insert("“".to_string(), "'".to_string());
+        // char_rep_map.insert("”".to_string(), "'".to_string());
+        // char_rep_map.insert("\"".to_string(), "'".to_string());
+        // char_rep_map.insert("‘".to_string(), "'".to_string());
+        // char_rep_map.insert("’".to_string(), "'".to_string());
+        // char_rep_map.insert("（".to_string(), "'".to_string());
+        // char_rep_map.insert("）".to_string(), "'".to_string());
+        // char_rep_map.insert("(".to_string(), "'".to_string());
+        // char_rep_map.insert(")".to_string(), "'".to_string());
+        // char_rep_map.insert("《".to_string(), "'".to_string());
+        // char_rep_map.insert("》".to_string(), "'".to_string());
+        // char_rep_map.insert("【".to_string(), "'".to_string());
+        // char_rep_map.insert("】".to_string(), "'".to_string());
+        // char_rep_map.insert("[".to_string(), "'".to_string());
+        // char_rep_map.insert("]".to_string(), "'".to_string());
+        // char_rep_map.insert("—".to_string(), "-".to_string());
+        // char_rep_map.insert("～".to_string(), "-".to_string());
+        // char_rep_map.insert("~".to_string(), "-".to_string());
+        // char_rep_map.insert("「".to_string(), "'".to_string());
+        // char_rep_map.insert("」".to_string(), "'".to_string());
+        // char_rep_map.insert(":".to_string(), ",".to_string());
+
+        // let mut zh_char_rep_map = char_rep_map.clone();
+        // zh_char_rep_map.insert("$".to_string(), ".".to_string());
+        Ok(Self {
+            // char_rep_map,
+            // zh_char_rep_map,
+            pinyin_tone_pattern,
+            // name_pattern,
+            // tech_term_pattern,
+            // english_contraction_pattern,
+            email_pattern,
+            cjk_range_pattern,
+        })
+    }
+
+    pub fn match_email(&self, email: &str) -> bool {
+        self.email_pattern.is_match(email)
+    }
+
+    pub fn use_chinese(&self, s: &str) -> bool {
+        let has_chinese = s.chars().any(|c| ('\u{4e00}'..='\u{9fff}').contains(&c));
+        let has_alpha = s.chars().any(|c| c.is_alphabetic());
+        let is_email = self.match_email(s);
+
+        if has_chinese || !has_alpha || is_email {
+            return true;
+        }
+
+        self.pinyin_tone_pattern.is_match(s)
+    }
+
+    pub fn tokenize_by_cjk_char(&self, line: &str, do_upper_case: bool) -> String {
+        // Split the line by CJK characters
+        let parts: Vec<&str> = self.cjk_range_pattern.split(line.trim()).collect();
+        // Process each part and join with spaces
+        let mut result_parts = Vec::new();
+        for part in parts {
+            if !part.trim().is_empty() {
+                if do_upper_case {
+                    result_parts.push(part.trim().to_uppercase());
+                } else {
+                    result_parts.push(part.trim().to_string());
+                }
+            }
+        }
+        // Join the parts with spaces
+        result_parts.join(" ")
+    }
+}
+
+pub fn tokenize_by_cjk_char(line: &str, do_upper_case: bool) -> String {
+    let mut result_parts = Vec::new();
+    for ch in line.chars() {
+        if ('\u{1100}'..='\u{11ff}').contains(&ch)
+            || ('\u{2e80}'..='\u{a4cf}').contains(&ch)
+            || ('\u{a840}'..='\u{d7af}').contains(&ch)
+            || ('\u{f900}'..='\u{faff}').contains(&ch)
+            || ('\u{fe30}'..='\u{fe4f}').contains(&ch)
+            || ('\u{ff65}'..='\u{ffdc}').contains(&ch)
+            || ('\u{20000}'..='\u{2ffff}').contains(&ch)
+            || ('\u{4e00}'..='\u{9fff}').contains(&ch)
+        {
+            // CJK 字符
+            if do_upper_case {
+                result_parts.push(ch.to_uppercase().collect::<String>());
+            } else {
+                result_parts.push(ch.to_string());
+            }
+        } else {
+            // 非 CJK 字符，保持在一起
+            if do_upper_case {
+                result_parts.push(ch.to_uppercase().to_string());
+            } else {
+                result_parts.push(ch.to_string());
+            }
+        }
+    }
+
+    result_parts.join(" ")
+}

src/models/mask_gct/model.rs

@@ -28,7 +28,7 @@ impl ConvNeXtBlock {
         // layer_scale_init_value: f32,
     ) -> Result<Self> {
         let dwconv = get_conv1d(vb.pp("dwconv"), dim, dim, 7, 3, 1, 1, dim, true)?;
-        let norm = get_layer_norm(vb.pp("norm"), 1e-6, dim)?;
+        let norm = get_layer_norm(vb.pp("norm"), 1e-6, dim, true)?;
         let pwconv1 = linear(dim, intermediate_dim, vb.pp("pwconv1"))?;
         let pwconv2 = linear(intermediate_dim, dim, vb.pp("pwconv2"))?;
         let gamma = vb.get_with_hints(dim, "gamma", Init::Const(1.0))?;
@@ -73,14 +73,14 @@ impl VocosBackbone {
         // layer_scale_init_value: Option<f32>,
     ) -> Result<Self> {
         let embed = get_conv1d(vb.pp("embed"), input_channels, dim, 7, 3, 1, 1, 1, true)?;
-        let norm = get_layer_norm(vb.pp("norm"), 1e-6, dim)?;
+        let norm = get_layer_norm(vb.pp("norm"), 1e-6, dim, true)?;
         let vb_convnext = vb.pp("convnext");
         let mut convnext = vec![];
         for i in 0..num_layers {
             let layer = ConvNeXtBlock::new(vb_convnext.pp(i), dim, intermediate_dim)?;
             convnext.push(layer);
         }
-        let final_layer_norm = get_layer_norm(vb.pp("final_layer_norm"), 1e-6, dim)?;
+        let final_layer_norm = get_layer_norm(vb.pp("final_layer_norm"), 1e-6, dim, true)?;
         Ok(Self {
             embed,
             norm,

src/models/paddleocr_vl/model.rs

@@ -37,6 +37,7 @@ impl Projector {
             vb.pp("pre_norm"),
             config.rms_norm_eps,
             config.vision_config.hidden_size,
+            true,
         )?;
         let linear_1 = linear(hidden_size, hidden_size, vb.pp("linear_1"))?;
         let linear_2 = linear(hidden_size, config.hidden_size, vb.pp("linear_2"))?;
@@ -278,6 +279,7 @@ impl SiglipVisionModel {
             vb.pp("post_layernorm"),
             config.layer_norm_eps,
             config.hidden_size,
+            true,
         )?;
         Ok(Self {
             embeddings,

src/models/qwen3_asr/model.rs

@@ -50,11 +50,12 @@ impl Qwen3ASRAudioEncoderLayer {
             Some("out_proj"),
         )?;
         let self_attn_layer_norm =
-            get_layer_norm(vb.pp("self_attn_layer_norm"), 1e-5, config.d_model)?;
+            get_layer_norm(vb.pp("self_attn_layer_norm"), 1e-5, config.d_model, true)?;
         let activation_fn = config.activation_function;
         let fc1 = linear(config.d_model, config.encoder_ffn_dim, vb.pp("fc1"))?;
         let fc2 = linear(config.encoder_ffn_dim, config.d_model, vb.pp("fc2"))?;
-        let final_layer_norm = get_layer_norm(vb.pp("final_layer_norm"), 1e-5, config.d_model)?;
+        let final_layer_norm =
+            get_layer_norm(vb.pp("final_layer_norm"), 1e-5, config.d_model, true)?;
         Ok(Self {
             self_attn,
             self_attn_layer_norm,
@@ -105,7 +106,7 @@ impl Qwen3ASRAudioEncoder {
             let layer = Qwen3ASRAudioEncoderLayer::new(vb_layers.pp(i), config)?;
             layers.push(layer);
         }
-        let ln_post = get_layer_norm(vb.pp("ln_post"), 1e-5, config.d_model)?;
+        let ln_post = get_layer_norm(vb.pp("ln_post"), 1e-5, config.d_model, true)?;
         let conv2d1 = get_conv2d(
             vb.pp("conv2d1"),
             1,

src/models/qwen3vl/model.rs

@@ -90,7 +90,7 @@ impl Qwen3VLVisionPatchMerger {
         } else {
             config.hidden_size
         };
-        let norm = get_layer_norm(vb.pp("norm"), 1e-6, norm_size)?;
+        let norm = get_layer_norm(vb.pp("norm"), 1e-6, norm_size, true)?;
         let linear_fc1 = linear(hidden_size, hidden_size, vb.pp("linear_fc1"))?;
         let act_fn = Activation::Gelu;
         let linear_fc2 = linear(hidden_size, config.out_hidden_size, vb.pp("linear_fc2"))?;
@@ -200,8 +200,8 @@ pub struct Qwen3VLVisionBlock {
 
 impl Qwen3VLVisionBlock {
     pub fn new(config: Qwen3VLVisionConfig, vb: VarBuilder) -> Result<Self> {
-        let norm1 = get_layer_norm(vb.pp("norm1"), 1e-6, config.hidden_size)?;
+        let norm1 = get_layer_norm(vb.pp("norm1"), 1e-6, config.hidden_size, true)?;
-        let norm2 = get_layer_norm(vb.pp("norm2"), 1e-6, config.hidden_size)?;
+        let norm2 = get_layer_norm(vb.pp("norm2"), 1e-6, config.hidden_size, true)?;
         let attn = Qwen3VLVisionAttention::new(config.clone(), vb.pp("attn"))?;
         let mlp = TwoLinearMLP::new(
             vb.pp("mlp"),

src/models/rmbg2_0/model.rs

@@ -42,7 +42,7 @@ impl PatchEmbed {
             true,
         )?;
         let norm = if patch_norm {
-            Some(get_layer_norm(vb.pp("norm"), 1e-5, embed_dim)?)
+            Some(get_layer_norm(vb.pp("norm"), 1e-5, embed_dim, true)?)
         } else {
             None
         };
@@ -242,7 +242,7 @@ impl SwinTransformerBlock {
         window_size: usize,
         shift_size: usize,
     ) -> Result<Self> {
-        let norm1 = get_layer_norm(vb.pp("norm1"), 1e-5, dim)?;
+        let norm1 = get_layer_norm(vb.pp("norm1"), 1e-5, dim, true)?;
 
         let attn = WindowAttention::new(
             vb.pp("attn"),
@@ -251,7 +251,7 @@ impl SwinTransformerBlock {
             qkv_bias,
             (window_size, window_size),
         )?;
-        let norm2 = get_layer_norm(vb.pp("norm2"), 1e-5, dim)?;
+        let norm2 = get_layer_norm(vb.pp("norm2"), 1e-5, dim, true)?;
         let mlp_dim = (dim as f32 * mlp_ratio) as usize;
         let mlp = TwoLinearMLP::new(vb.pp("mlp"), dim, mlp_dim, dim, act, true, "fc1", "fc2")?;
         Ok(Self {
@@ -323,7 +323,7 @@ struct PatchMerging {
 impl PatchMerging {
     pub fn new(vb: VarBuilder, dim: usize) -> Result<Self> {
         let reduction = linear_no_bias(4 * dim, 2 * dim, vb.pp("reduction"))?;
-        let norm = get_layer_norm(vb.pp("norm"), 1e-5, 4 * dim)?;
+        let norm = get_layer_norm(vb.pp("norm"), 1e-5, 4 * dim, true)?;
         Ok(Self { reduction, norm })
     }
 
@@ -518,7 +518,7 @@ impl SwinTransformer {
         }
         let mut norms = vec![];
         for i in out_indices.clone() {
-            let layer_i = get_layer_norm(vb.pp(format!("norm{i}")), 1e-5, num_features[i])?;
+            let layer_i = get_layer_norm(vb.pp(format!("norm{i}")), 1e-5, num_features[i], true)?;
             norms.push(layer_i);
         }
         Ok(Self {

src/models/w2v_bert_2_0/model.rs

@@ -7,9 +7,7 @@ use candle_nn::{
 
 use crate::{
     models::{
-        common::{
+        common::{GLU, TwoLinearMLP, eager_attention_forward, get_conv1d, get_layer_norm},
-            GLU, TwoLinearMLP, eager_attention_forward, get_conv1d, get_layer_norm,
-        },
         w2v_bert_2_0::config::W2VBert2_0Config,
     },
     position_embed::rope::{RoPE, apply_rotary_pos_emb},
@@ -27,6 +25,7 @@ impl Wav2Vec2BertFeatureProjection {
             vb.pp("layer_norm"),
             config.layer_norm_eps,
             config.feature_projection_input_dim,
+            true,
         )?;
         let projection = linear(
             config.feature_projection_input_dim,
@@ -248,6 +247,7 @@ impl Wav2Vec2BertConvolutionModule {
             vb.pp("layer_norm"),
             config.layer_norm_eps,
             config.hidden_size,
+            true,
         )?;
         let pointwise_conv1 = get_conv1d(
             vb.pp("pointwise_conv1"),
@@ -277,6 +277,7 @@ impl Wav2Vec2BertConvolutionModule {
             vb.pp("depthwise_layer_norm"),
             config.layer_norm_eps,
             config.hidden_size,
+            true,
         )?;
         let pointwise_conv2 = get_conv1d(
             vb.pp("pointwise_conv2"),
@@ -341,6 +342,7 @@ impl Wav2Vec2BertEncoderLayer {
             vb.pp("ffn1_layer_norm"),
             config.layer_norm_eps,
             config.hidden_size,
+            true,
         )?;
         let ffn1 = TwoLinearMLP::new(
             vb.pp("ffn1"),
@@ -356,6 +358,7 @@ impl Wav2Vec2BertEncoderLayer {
             vb.pp("self_attn_layer_norm"),
             config.layer_norm_eps,
             config.hidden_size,
+            true,
         )?;
         let self_attn = Wav2Vec2BertSelfAttention::new(vb.pp("self_attn"), config, false)?;
         let conv_module = Wav2Vec2BertConvolutionModule::new(vb.pp("conv_module"), config)?;
@@ -363,6 +366,7 @@ impl Wav2Vec2BertEncoderLayer {
             vb.pp("ffn2_layer_norm"),
             config.layer_norm_eps,
             config.hidden_size,
+            true,
         )?;
         let ffn2 = TwoLinearMLP::new(
             vb.pp("ffn2"),
@@ -378,6 +382,7 @@ impl Wav2Vec2BertEncoderLayer {
             vb.pp("final_layer_norm"),
             config.layer_norm_eps,
             config.hidden_size,
+            true,
         )?;
         Ok(Self {
             ffn1_layer_norm,

src/tokenizer/mod.rs

@@ -1,5 +1,6 @@
-use anyhow::{Ok, Result, anyhow};
+use anyhow::{Result, anyhow};
 use candle_core::{Device, Tensor};
+use sentencepiece::SentencePieceProcessor;
 use serde_json::Value;
 use tokenizers::{
     AddedToken, Tokenizer, decoders::byte_level::ByteLevel as ByteLevelDecoder, models::bpe::BPE,
@@ -114,3 +115,16 @@ impl TokenizerModel {
         Ok(decode)
     }
 }
+
+pub fn sentencepiece_encode(
+    text: &str,
+    tokenizer: &SentencePieceProcessor,
+    device: &Device,
+) -> Result<Tensor> {
+    let tokens = tokenizer
+        .encode(text)
+        .map_err(|e| anyhow!(format!("tokenizer encode error:{}", e)))?;
+    let token_ids = tokens.iter().map(|p| p.id).collect::<Vec<u32>>();
+    let tokens_t = Tensor::new(token_ids, device)?.unsqueeze(0)?;
+    Ok(tokens_t)
+}

src/utils/mod.rs

@@ -219,7 +219,7 @@ pub fn crate_tensor_from_reader<R: std::io::Read>(
 pub fn read_pth_tensor_info_cycle<P: AsRef<std::path::Path>>(
     path: P,
     key: Option<&str>,
-) -> Result<Vec<(String, Tensor)>> {
+) -> Result<HashMap<String, Tensor>> {
     let file = std::fs::File::open(path.as_ref())?;
     let zip_reader = std::io::BufReader::new(file);
     let mut zip = zip::ZipArchive::new(zip_reader)?;
@@ -356,7 +356,11 @@ pub fn read_pth_tensor_info_cycle<P: AsRef<std::path::Path>>(
             }
         };
     }
-    Ok(tensors)
+    let mut dict_to_hashmap = HashMap::new();
+    for (k, v) in tensors {
+        dict_to_hashmap.insert(k, v);
+    }
+    Ok(dict_to_hashmap)
 }
 
 pub fn round_by_factor(num: u32, factor: u32) -> u32 {

src/utils/tensor_utils.rs

@@ -111,7 +111,7 @@ pub fn split_tensor_with_size<D: Dim>(
     //     "input tensor dim size % splits_size must be equal to 0"
     // );
     for (i, split) in (0..dim_size).step_by(splits_size).enumerate() {
-        let size = splits_size.min(dim_size - i*splits_size);
+        let size = splits_size.min(dim_size - i * splits_size);
         split_res.push(t.narrow(dim, split, size)?);
     }
     Ok(split_res)
@@ -1060,3 +1060,11 @@ pub fn float_range_normalize(t: &Tensor) -> Result<Tensor> {
     t = t.clamp(-1.0, 1.0)?;
     Ok(t)
 }
+
+pub fn sequence_mask(length: &Tensor, max_length: Option<u32>) -> Result<Tensor> {
+    let max_length = max_length.unwrap_or(length.max_all()?.to_scalar::<u32>()?);
+    let x = Tensor::arange(0, max_length, length.device())?.unsqueeze(0)?;
+    let length = length.unsqueeze(1)?;
+    let mask = x.broadcast_lt(&length)?;
+    Ok(mask)
+}

tests/messy_test.rs

@@ -2,22 +2,33 @@
 
 use std::time::Instant;
 
-use aha::utils::{tensor_utils::interpolate_nearest_1d};
+use aha::utils::tensor_utils::interpolate_nearest_1d;
-use anyhow::Result;
+use anyhow::{Result, anyhow};
-use candle_core::{Tensor};
+use candle_core::Tensor;
+use sentencepiece::SentencePieceProcessor;
 // use symphonia::core::io::MediaSourceStream;
 
 #[test]
 fn messy_test() -> Result<()> {
-    // RUST_BACKTRACE=1 cargo test -F cuda,ffmpeg messy_test -r -- --nocapture
+    // RUST_BACKTRACE=1 cargo test -F cuda messy_test -r -- --nocapture
     let device = &candle_core::Device::Cpu;
-    let t = Tensor::arange(0.0f32, 40.0, device)?.broadcast_as((1, 40, 40))?;
+    let save_dir =
-    println!("t: {}", t);
+        aha::utils::get_default_save_dir().ok_or(anyhow::anyhow!("Failed to get save dir"))?;
-    let i_start = Instant::now();
+    let model_path = format!("{}/IndexTeam/IndexTTS-2", save_dir);
-    let t_inter = interpolate_nearest_1d(&t, 20)?;
+    let bpe_path = model_path.to_string() + "/bpe.model";
-    let i_duration = i_start.elapsed();
+    let tokenizer = SentencePieceProcessor::open(bpe_path)
-    println!("Time elapsed in interpolate_nearest_1d is: {:?}", i_duration);
+        .map_err(|e| anyhow!(format!("load bpe.model file error:{}", e)))?;
-    println!("t_inter: {}", t_inter);    
+    let tokens = tokenizer
+        .encode("你好啊")
+        .map_err(|e| anyhow!(format!("tokenizer encode error:{}", e)))?;
+    println!("tokens: {:?}", tokens);
+    // let t = Tensor::arange(0.0f32, 40.0, device)?.broadcast_as((1, 40, 40))?;
+    // println!("t: {}", t);
+    // let i_start = Instant::now();
+    // let t_inter = interpolate_nearest_1d(&t, 20)?;
+    // let i_duration = i_start.elapsed();
+    // println!("Time elapsed in interpolate_nearest_1d is: {:?}", i_duration);
+    // println!("t_inter: {}", t_inter);
     // let url = "https://sis-sample-audio.obs.cn-north-1.myhuaweicloud.com/16k16bit.mp3";
     // let client = reqwest::blocking::Client::new();
     // let response = client.get(url).send()?;

tests/test_fun_asr_nano.rs

@@ -21,7 +21,7 @@ fn fun_asr_nano_generate() -> Result<()> {
                         "type": "audio",
                         "audio_url": 
                         {
-                            "url": "https://package-release.coderbox.cn/aiway/test/other/%E5%93%AA%E5%90%92.wav"
+                            "url": "file://./assets/audio/zh.mp3"
                         }
                     },           
                     {

tests/weight_test.rs

@@ -200,6 +200,7 @@ fn qwen3_weight() -> Result<()> {
 
 #[test]
 fn index_tts2_weight() -> Result<()> {
+    // RUST_BACKTRACE=1 cargo test -F cuda index_tts2_weight -r -- --nocapture
     let save_dir: String =
         aha::utils::get_default_save_dir().ok_or(anyhow::anyhow!("Failed to get save dir"))?;
     let model_path = format!("{}/IndexTeam/IndexTTS-2/", save_dir);