Merge branch 'add_model'

CHANGELOG.md

@@ -1,5 +1,9 @@
 # Changelog
 
+## [Unreleased] - 2025-02-04
+
+### Added
+- Support for Qwen3-ASR model
 
 ## [Unreleased]
 

Cargo.toml

@@ -28,7 +28,7 @@ rocket = { version = "0.5.1", features = ["serde_json", "json"] }
 tokio = "1.47.1"
 hound = "3.5.1"
 clap = { version = "4.5.51", features = ["derive"] }
-modelscope = "0.1.0"
+modelscope = "0.1.3"
 dirs = "6.0.0"
 url = "2.5.7"
 rayon = "1.10"
@@ -37,6 +37,9 @@ rayon = "1.10"
 realfft = "3.5.0"
 symphonia = { version = "0.5.5", features = ["mp3", "wav"] }
 serde_yaml = "0.9.34"
+zip = "7.2.0"
+half = "2.7.1"
+byteorder = "1.5.0"
 
 [features]
 flash-attn = ["candle-flash-attn"]

README.md

@@ -130,7 +130,7 @@ cargo run -F cuda -r -- [参数]
     * deepseek-ocr: deepseek-ai/DeepSeek-OCR 模型
     * hunyuan-ocr: Tencent-Hunyuan/HunyuanOCR 模型
     * paddleocr-vl: PaddlePaddle/PaddleOCR-VL 模型
-    * RMBG2.0: AI-ModelScope/RMBG-2.0 模型
+    * rmbg2.0: AI-ModelScope/RMBG-2.0 模型
     * voxcpm: OpenBMB/VoxCPM-0.5B 模型
     * voxcpm1.5: OpenBMB/VoxCPM1.5 模型
     * glm-asr-nano-2512: ZhipuAI/GLM-ASR-Nano-2512 模型

download_and_run.sh

@@ -17,7 +17,9 @@ show_help() {
     echo "  minicpm4-0.5b"
     echo "  qwen2.5vl-3b"
     echo "  qwen2.5vl-7b" 
-    echo "  qwen3-0.6b"
+    echo "  qwen3-0.6b" 
+    echo "  qwen3asr-0.6b"
+    echo "  qwen3asr-1.7b"
     echo "  qwen3vl-2b"
     echo "  qwen3vl-4b"
     echo "  qwen3vl-8b"
@@ -25,7 +27,7 @@ show_help() {
     echo "  deepseek-ocr"
     echo "  hunyuan-ocr"
     echo "  paddleocr-vl"
-    echo "  RMBG2.0"
+    echo "  rmbg2.0"
     echo "  voxcpm"
     echo "  voxcpm1.5"
     echo "  glm-asr-nano-2512"
@@ -56,6 +58,12 @@ case $MODEL_ALIAS in
     "qwen3-0.6b")
         MODEL_ID="Qwen/Qwen3-0.6B"
         ;;
+    "qwen3asr-0.6b")
+        MODEL_ID="Qwen/Qwen3-ASR-0.6B"
+        ;;
+    "qwen3asr-1.7b")
+        MODEL_ID="Qwen/Qwen3-ASR-1.7B"
+        ;;
     "qwen3vl-2b")
         MODEL_ID="Qwen/Qwen3-VL-2B-Instruct"
         ;;
@@ -77,7 +85,7 @@ case $MODEL_ALIAS in
     "paddleocr-vl")
         MODEL_ID="PaddlePaddle/PaddleOCR-VL"
         ;;
-    "RMBG2.0")
+    "rmbg2.0")
         MODEL_ID="briaai/RMBG-2.0"
         ;;
     "voxcpm")

src/chat_template/mod.rs

@@ -6,16 +6,33 @@ use crate::utils::string_to_static_str;
 
 pub fn get_template(path: String) -> Result<String> {
     let tokenizer_config_file = path.clone() + "/tokenizer_config.json";
-    assert!(
+    let chat_template = if std::path::Path::new(&tokenizer_config_file).exists() {
-        std::path::Path::new(&tokenizer_config_file).exists(),
+        let tokenizer_config: serde_json::Value =
-        "tokenizer_config.json not exists in model path"
+            serde_json::from_slice(&std::fs::read(tokenizer_config_file)?)
-    );
+                .map_err(|e| anyhow!(format!("load tokenizer_config file error:{}", e)))?;
-    let tokenizer_config: serde_json::Value =
+        let chat_template = tokenizer_config["chat_template"]
-        serde_json::from_slice(&std::fs::read(tokenizer_config_file)?)
+            .as_str()
-            .map_err(|e| anyhow!(format!("load tokenizer_config file error:{}", e)))?;
+            .map(|s| s.to_string());
-    let chat_template = tokenizer_config["chat_template"]
+        match chat_template {
-        .as_str()
+            Some(tem) => Some(tem),
-        .ok_or(anyhow!(format!("chat_template to str error")))?;
+            None => {
+                let chat_template_file = path.clone() + "/chat_template.json";
+                if std::path::Path::new(&chat_template_file).exists() {
+                    let chat_template_config: serde_json::Value =
+                        serde_json::from_slice(&std::fs::read(chat_template_file)?)
+                            .map_err(|e| anyhow!(format!("load chat_template file error:{}", e)))?;
+                    chat_template_config["chat_template"]
+                        .as_str()
+                        .map(|s| s.to_string())
+                } else {
+                    None
+                }
+            }
+        }
+    } else {
+        None
+    };
+    let chat_template = chat_template.ok_or(anyhow!(format!("chat_template is none")))?;
     // 修复模板中的问题行
     let fixed_template = chat_template
         .replace(

src/exec/mod.rs

@@ -11,6 +11,7 @@ pub mod minicpm4;
 pub mod paddleocr_vl;
 pub mod qwen2_5vl;
 pub mod qwen3;
+pub mod qwen3_asr;
 pub mod qwen3vl;
 pub mod rmbg2_0;
 pub mod voxcpm;

src/exec/qwen3_asr.rs

@@ -0,0 +1,67 @@
+//! Fun-ASR-Nano-2512 exec implementation for CLI `run` subcommand
+
+use std::time::Instant;
+
+use anyhow::{Ok, Result};
+
+use crate::exec::ExecModel;
+use crate::models::qwen3_asr::generate::Qwen3AsrGenerateModel;
+use crate::models::{GenerateModel};
+
+pub struct Qwen3ASRExec;
+
+impl ExecModel for Qwen3ASRExec {
+    fn run(input: &[String], output: Option<&str>, weight_path: &str) -> Result<()> {
+        
+        let i_start = Instant::now();
+        let mut model = Qwen3AsrGenerateModel::init(weight_path, None, None)?;
+        let i_duration = i_start.elapsed();
+        println!("Time elapsed in load model is: {:?}", i_duration);
+
+        // Create ChatCompletionParameters for ASR
+        let url = &input[1];
+        let input_url = if url.starts_with("http://")
+            || url.starts_with("https://")
+            || url.starts_with("file://")
+        {
+            url.clone()
+        } else {
+            format!("file://{}", url)
+        };
+
+        let message = format!(
+            r#"{{
+            "model": "fun-asr-nano",
+            "messages": [
+                {{
+                    "role": "user",
+                    "content": [
+                        {{
+                            "type": "audio",
+                            "audio_url": {{
+                                "url": "{}"
+                            }}
+                        }}
+                    ]
+                }}
+            ]
+        }}"#,
+            input_url
+        );
+        let mes = serde_json::from_str(&message)?;
+
+        let i_start = Instant::now();
+        let res = model.generate(mes)?;
+        let i_duration = i_start.elapsed();
+        println!("Time elapsed in generate is: {:?}", i_duration);
+
+        println!("Result: {:?}", res);
+
+        if let Some(out) = output {
+            std::fs::write(out, format!("{:?}", res))?;
+            println!("Output saved to: {}", out);
+        }
+
+        Ok(())
+    }
+}

src/main.rs

@@ -1,14 +1,15 @@
-use std::{net::IpAddr, str::FromStr, time::Duration};
+use std::{net::IpAddr, str::FromStr};
 
-use aha::{models::WhichModel, utils::get_default_save_dir};
+use aha::{
+    models::WhichModel,
+    utils::{download_model, get_default_save_dir},
+};
 use clap::{Args, Parser, Subcommand, ValueEnum};
-use modelscope::ModelScope;
 use rocket::{
     Config,
     data::{ByteUnit, Limits},
     routes,
 };
-use tokio::time::sleep;
 
 use crate::api::init;
 mod api;
@@ -141,39 +142,6 @@ struct RunArgs {
     weight_path: String,
 }
 
-async fn download_model(model_id: &str, save_dir: &str, max_retries: u32) -> anyhow::Result<()> {
-    let mut attempts = 0u32;
-    loop {
-        attempts += 1;
-        println!(
-            "Attempting to download model (attempt {}/{})",
-            attempts, max_retries
-        );
-
-        match ModelScope::download(model_id, save_dir).await {
-            Ok(()) => {
-                println!("Model downloaded successfully");
-                return Ok(());
-            }
-            Err(e) => {
-                if attempts >= max_retries {
-                    return Err(anyhow::anyhow!(
-                        "Failed to download model after {} attempts. Last error: {}",
-                        max_retries,
-                        e
-                    ));
-                }
-
-                println!(
-                    "Download failed (attempt {}): {}. Retrying in 2 seconds...",
-                    attempts, e
-                );
-                sleep(Duration::from_secs(2)).await;
-            }
-        }
-    }
-}
-
 /// Get the ModelScope model ID for a given WhichModel variant
 fn get_model_id(model: WhichModel) -> &'static str {
     match model {
@@ -181,6 +149,8 @@ fn get_model_id(model: WhichModel) -> &'static str {
         WhichModel::Qwen2_5vl3B => "Qwen/Qwen2.5-VL-3B-Instruct",
         WhichModel::Qwen2_5vl7B => "Qwen/Qwen2.5-VL-7B-Instruct",
         WhichModel::Qwen3_0_6B => "Qwen/Qwen3-0.6B",
+        WhichModel::Qwen3ASR0_6B => "Qwen/Qwen3-ASR-0.6B",
+        WhichModel::Qwen3ASR1_7B => "Qwen/Qwen3-ASR-1.7B",
         WhichModel::Qwen3vl2B => "Qwen/Qwen3-VL-2B-Instruct",
         WhichModel::Qwen3vl4B => "Qwen/Qwen3-VL-4B-Instruct",
         WhichModel::Qwen3vl8B => "Qwen/Qwen3-VL-8B-Instruct",
@@ -203,6 +173,8 @@ fn run_list() -> anyhow::Result<()> {
         WhichModel::Qwen2_5vl3B,
         WhichModel::Qwen2_5vl7B,
         WhichModel::Qwen3_0_6B,
+        WhichModel::Qwen3ASR0_6B,
+        WhichModel::Qwen3ASR1_7B,
         WhichModel::Qwen3vl2B,
         WhichModel::Qwen3vl4B,
         WhichModel::Qwen3vl8B,
@@ -321,6 +293,14 @@ fn run_run(args: RunArgs) -> anyhow::Result<()> {
             use aha::exec::qwen3::Qwen3Exec;
             Qwen3Exec::run(&input, output.as_deref(), &weight_path)?;
         }
+        WhichModel::Qwen3ASR0_6B => {
+            use aha::exec::qwen3_asr::Qwen3ASRExec;
+            Qwen3ASRExec::run(&input, output.as_deref(), &weight_path)?;
+        }
+        WhichModel::Qwen3ASR1_7B => {
+            use aha::exec::qwen3_asr::Qwen3ASRExec;
+            Qwen3ASRExec::run(&input, output.as_deref(), &weight_path)?;
+        }
         WhichModel::Qwen3vl2B => {
             use aha::exec::qwen3vl::Qwen3vlExec;
             Qwen3vlExec::run(&input, output.as_deref(), &weight_path)?;

src/models/campplus/mod.rs

@@ -0,0 +1,549 @@
+use anyhow::Result;
+use candle_core::{D, Tensor};
+use candle_nn::{BatchNorm, Conv1d, Conv2d, Module, ModuleT, VarBuilder, ops::sigmoid};
+
+use crate::{
+    models::common::{get_batch_norm, get_conv1d, get_conv2d},
+    utils::tensor_utils::{pool1d, statistics_pooling},
+};
+
+pub struct Shortcut {
+    conv_0: Conv2d,
+    bn_1: BatchNorm,
+    stride: usize,
+}
+
+impl Shortcut {
+    pub fn new(
+        vb: VarBuilder,
+        in_c: usize,
+        out_c: usize,
+        ks: usize,
+        padding: usize,
+        stride: usize,
+        bias: bool,
+    ) -> Result<Self> {
+        let conv_0 = get_conv2d(vb.pp("0"), in_c, out_c, ks, padding, 1, 1, 1, bias)?;
+        let bn_1 = get_batch_norm(vb.pp("1"), 1e-5, out_c, true)?;
+        Ok(Self { conv_0, bn_1, stride })
+    }
+
+    pub fn forward(&self, x: &Tensor) -> Result<Tensor> {
+        let mut x = self.conv_0.forward(x)?;
+        if self.stride != 1 {
+            let h_dim = x.dim(2)?;
+            let half_h = h_dim / 2;
+            let indices = Tensor::arange(0u32, half_h as u32, x.device())?.affine(2.0, 0.0)?;
+            x = x.index_select(&indices, 2)?;
+        }
+        x = self.bn_1.forward_t(&x, false)?;        
+        Ok(x)
+    }
+}
+
+pub struct BasicResBlock {
+    stride: usize,
+    conv1: Conv2d,
+    bn1: BatchNorm,
+    conv2: Conv2d,
+    bn2: BatchNorm,
+    shortcut: Option<Shortcut>,
+}
+
+impl BasicResBlock {
+    pub fn new(
+        vb: VarBuilder,
+        in_planes: usize,
+        planes: usize,
+        stride: usize,
+        expansion: usize,
+    ) -> Result<Self> {
+        let conv1 = get_conv2d(vb.pp("conv1"), in_planes, planes, 3, 1, 1, 1, 1, false)?;
+        let bn1 = get_batch_norm(vb.pp("bn1"), 1e-5, planes, true)?;
+        let conv2 = get_conv2d(vb.pp("conv2"), planes, planes, 3, 1, 1, 1, 1, false)?;
+        let bn2 = get_batch_norm(vb.pp("bn2"), 1e-5, planes, true)?;
+        let shortcut = if stride != 1 || in_planes != expansion * planes {
+            Some(Shortcut::new(
+                vb.pp("shortcut"),
+                in_planes,
+                expansion * planes,
+                1,
+                0,
+                stride,
+                false,
+            )?)
+        } else {
+            None
+        };
+        Ok(Self {
+            stride,
+            conv1,
+            bn1,
+            conv2,
+            bn2,
+            shortcut,
+        })
+    }
+    pub fn forward(&self, xs: &Tensor) -> Result<Tensor> {
+        let residual = xs.clone();
+        let mut xs = self.conv1.forward(xs)?;
+        // candle stride only surpport one size, h_stride = w_stride
+        // infact model stride is (stride, 1)
+        // but now setting stride all equal to 1
+        // so h direction use indices select
+        if self.stride != 1 {
+            let h_dim = xs.dim(2)?;
+            let half_h = h_dim / 2;
+            let indices = Tensor::arange(0u32, half_h as u32, xs.device())?.affine(2.0, 0.0)?;
+            xs = xs.index_select(&indices, 2)?;
+        }
+        let xs = self.bn1.forward_t(&xs, false)?.relu()?;
+        let xs = self.conv2.forward(&xs)?;
+        let mut xs = self.bn2.forward_t(&xs, false)?;
+        if let Some(cut) = &self.shortcut {
+            let shortcut = cut.forward(&residual)?;
+            xs = xs.add(&shortcut)?;
+        } else {
+            xs = xs.add(&residual)?;
+        }
+        xs = xs.relu()?;       
+        Ok(xs)
+    }
+}
+
+pub struct FCM {
+    conv1: Conv2d,
+    bn1: BatchNorm,
+    layer1: Vec<BasicResBlock>,
+    layer2: Vec<BasicResBlock>,
+    conv2: Conv2d,
+    bn2: BatchNorm,
+    pub out_channels: usize,
+}
+
+impl FCM {
+    pub fn new(
+        vb: VarBuilder,
+        num_blocks: &[usize],
+        m_channels: usize,
+        feat_dim: usize,
+    ) -> Result<Self> {
+        let conv1 = get_conv2d(vb.pp("conv1"), 1, m_channels, 3, 1, 1, 1, 1, false)?;
+        let bn1 = get_batch_norm(vb.pp("bn1"), 1e-5, m_channels, true)?;
+        let layer1_num_blocks = num_blocks[0] - 1;
+        let strides: Vec<usize> = [2usize]
+            .into_iter()
+            .chain([1usize].into_iter().cycle().take(layer1_num_blocks))
+            .collect();
+        let mut layer1 = vec![];
+        let vb_layer1 = vb.pp("layer1");
+        for (i, stride) in strides.iter().enumerate() {
+            let layer = BasicResBlock::new(vb_layer1.pp(i), m_channels, m_channels, *stride, 1)?;
+            layer1.push(layer);
+        }
+        let layer2_num_blocks = num_blocks[1] - 1;
+        let strides: Vec<usize> = [2usize]
+            .into_iter()
+            .chain([1usize].into_iter().cycle().take(layer2_num_blocks))
+            .collect();
+        let mut layer2 = vec![];
+        let vb_layer2 = vb.pp("layer2");
+        for (i, stride) in strides.iter().enumerate() {
+            let layer = BasicResBlock::new(vb_layer2.pp(i), m_channels, m_channels, *stride, 1)?;
+            layer2.push(layer);
+        }
+        let conv2 = get_conv2d(vb.pp("conv2"), m_channels, m_channels, 3, 1, 1, 1, 1, false)?;
+        let bn2 = get_batch_norm(vb.pp("bn2"), 1e-5, m_channels, true)?;
+        let out_channels = m_channels * (feat_dim / 8);
+        Ok(Self {
+            conv1,
+            bn1,
+            layer1,
+            layer2,
+            conv2,
+            bn2,
+            out_channels,
+        })
+    }
+    pub fn forward(&self, xs: &Tensor) -> Result<Tensor> {
+        let xs = xs.unsqueeze(1)?;
+        let xs = self.conv1.forward(&xs)?;
+        let mut xs = self.bn1.forward_t(&xs, false)?.relu()?;
+        for layer in &self.layer1 {
+            xs = layer.forward(&xs)?;
+        }
+        for layer in &self.layer2 {
+            xs = layer.forward(&xs)?;
+        }
+        xs = self.conv2.forward(&xs)?;
+        let h_dim = xs.dim(2)?;
+        let half_h = h_dim / 2;
+        let indices = Tensor::arange(0u32, half_h as u32, xs.device())?.affine(2.0, 0.0)?;
+        xs = xs.index_select(&indices, 2)?;
+        xs = self.bn2.forward_t(&xs, false)?.relu()?;
+        let (bs, c, h, dim) = xs.dims4()?;
+        xs = xs.reshape((bs, c * h, dim))?;
+        Ok(xs)
+    }
+}
+
+pub struct TDNNLayer {
+    linear: Conv1d,
+    nonlinear: BatchNorm,
+}
+
+impl TDNNLayer {
+    pub fn new(
+        vb: VarBuilder,
+        in_c: usize,
+        out_c: usize,
+        ks: usize,
+        stride: usize,
+        dilation: usize,
+        bias: bool,
+    ) -> Result<Self> {
+        let padding = (ks - 1) / 2 * dilation;
+        let linear = get_conv1d(
+            vb.pp("linear"),
+            in_c,
+            out_c,
+            ks,
+            padding,
+            stride,
+            dilation,
+            1,
+            bias,
+        )?;
+        let nonlinear = get_batch_norm(vb.pp("nonlinear.batchnorm"), 1e-5, out_c, true)?;
+        Ok(Self { linear, nonlinear })
+    }
+
+    pub fn forward(&self, xs: &Tensor) -> Result<Tensor> {
+        let xs = self.linear.forward(xs)?;
+        let xs = self.nonlinear.forward_t(&xs, false)?.relu()?;
+        Ok(xs)
+    }
+}
+
+pub struct CAMLayer {
+    linear_local: Conv1d,
+    linear1: Conv1d,
+    linear2: Conv1d,
+}
+
+impl CAMLayer {
+    pub fn new(
+        vb: VarBuilder,
+        bn_c: usize,
+        out_c: usize,
+        ks: usize,
+        stride: usize,
+        padding: usize,
+        dilation: usize,
+        bias: bool,
+        reduction: usize,
+    ) -> Result<Self> {
+        let linear_local = get_conv1d(
+            vb.pp("linear_local"),
+            bn_c,
+            out_c,
+            ks,
+            padding,
+            stride,
+            dilation,
+            1,
+            bias,
+        )?;
+        let linear1 = get_conv1d(
+            vb.pp("linear1"),
+            bn_c,
+            bn_c / reduction,
+            1,
+            0,
+            1,
+            1,
+            1,
+            true,
+        )?;
+        let linear2 = get_conv1d(
+            vb.pp("linear2"),
+            bn_c / reduction,
+            out_c,
+            1,
+            0,
+            1,
+            1,
+            1,
+            true,
+        )?;
+        Ok(Self {
+            linear_local,
+            linear1,
+            linear2,
+        })
+    }
+
+    pub fn seg_pooling(&self, xs: &Tensor, seg_len: usize, stype: &str) -> Result<Tensor> {
+        let x_dim = xs.dim(2)?;
+        let seg = pool1d(xs, seg_len, true, stype)?;
+        let (bs, c, dim) = seg.dims3()?;
+        let seg = seg
+            .unsqueeze(D::Minus1)?
+            .expand((bs, c, dim, seg_len))?
+            .reshape((bs, c, ()))?;
+        let seg = seg.narrow(D::Minus1, 0, x_dim)?;
+        Ok(seg)
+    }
+
+    pub fn forward(&self, xs: &Tensor) -> Result<Tensor> {
+        let y = self.linear_local.forward(xs)?;
+        let x_pool = self.seg_pooling(xs, 100, "avg")?;
+        let context = xs.mean_keepdim(D::Minus1)?.broadcast_add(&x_pool)?;
+        let context = self.linear1.forward(&context)?.relu()?;
+        let m = sigmoid(&self.linear2.forward(&context)?)?;
+        let res = y.mul(&m)?;
+        Ok(res)
+    }
+}
+
+pub struct CAMDenseTDNNLayer {
+    nonlinear1: BatchNorm,
+    linear1: Conv1d,
+    nonlinear2: BatchNorm,
+    cam_layer: CAMLayer,
+}
+
+impl CAMDenseTDNNLayer {
+    pub fn new(
+        vb: VarBuilder,
+        in_c: usize,
+        out_c: usize,
+        bn_c: usize,
+        ks: usize,
+        stride: usize,
+        dilation: usize,
+        bias: bool,
+    ) -> Result<Self> {
+        let padding = (ks - 1) / 2 * dilation;
+        let nonlinear1 = get_batch_norm(vb.pp("nonlinear1.batchnorm"), 1e-5, in_c, true)?;
+        let linear1 = get_conv1d(vb.pp("linear1"), in_c, bn_c, 1, 0, 1, 1, 1, false)?;
+        let nonlinear2 = get_batch_norm(vb.pp("nonlinear2.batchnorm"), 1e-5, bn_c, true)?;
+        let cam_layer = CAMLayer::new(
+            vb.pp("cam_layer"),
+            bn_c,
+            out_c,
+            ks,
+            stride,
+            padding,
+            dilation,
+            bias,
+            2,
+        )?;
+        Ok(Self {
+            nonlinear1,
+            linear1,
+            nonlinear2,
+            cam_layer,
+        })
+    }
+
+    pub fn forward(&self, xs: &Tensor) -> Result<Tensor> {
+        let xs = self.nonlinear1.forward_t(xs, false)?.relu()?;
+        let xs = self.linear1.forward(&xs)?;
+        let xs = self.nonlinear2.forward_t(&xs, false)?.relu()?;
+        let xs = self.cam_layer.forward(&xs)?;
+        Ok(xs)
+    }
+}
+
+pub struct CAMDenseTDNNBlock {
+    tdnns: Vec<CAMDenseTDNNLayer>,
+}
+
+impl CAMDenseTDNNBlock {
+    pub fn new(
+        vb: VarBuilder,
+        num_layers: usize,
+        in_c: usize,
+        out_c: usize,
+        bn_c: usize,
+        ks: usize,
+        stride: usize,
+        dilation: usize,
+        bias: bool,
+    ) -> Result<Self> {
+        let mut tdnns = vec![];
+        for i in 0..num_layers {
+            let layer = CAMDenseTDNNLayer::new(
+                vb.pp(format!("tdnnd{}", i + 1)),
+                in_c + i * out_c,
+                out_c,
+                bn_c,
+                ks,
+                stride,
+                dilation,
+                bias,
+            )?;
+            tdnns.push(layer);
+        }
+        Ok(Self { tdnns })
+    }
+    pub fn forward(&self, xs: &Tensor) -> Result<Tensor> {
+        let mut xs = xs.clone();
+        for layer in &self.tdnns {
+            let layer_out = layer.forward(&xs)?;
+            xs = Tensor::cat(&[&xs, &layer_out], 1)?;
+        }
+        Ok(xs)
+    }
+}
+
+pub struct TransitLayer {
+    nonlinear: BatchNorm,
+    linear: Conv1d,
+}
+
+impl TransitLayer {
+    pub fn new(vb: VarBuilder, in_c: usize, out_c: usize, bias: bool) -> Result<Self> {
+        let nonlinear = get_batch_norm(vb.pp("nonlinear.batchnorm"), 1e-5, in_c, true)?;
+        let linear = get_conv1d(vb.pp("linear"), in_c, out_c, 1, 0, 1, 1, 1, bias)?;
+        Ok(Self { nonlinear, linear })
+    }
+
+    pub fn forward(&self, xs: &Tensor) -> Result<Tensor> {
+        let xs = self.nonlinear.forward_t(xs, false)?.relu()?;
+        let xs = self.linear.forward(&xs)?;
+        Ok(xs)
+    }
+}
+
+pub struct DenseLayer {
+    linear: Conv1d,
+    nonlinear: BatchNorm, // only batch norm, no relu
+}
+
+impl DenseLayer {
+    pub fn new(vb: VarBuilder, in_c: usize, out_c: usize, bias: bool) -> Result<Self> {
+        let linear = get_conv1d(vb.pp("linear"), in_c, out_c, 1, 0, 1, 1, 1, bias)?;
+        let nonlinear = get_batch_norm(vb.pp("nonlinear.batchnorm"), 1e-5, out_c, false)?;
+        Ok(Self { linear, nonlinear })
+    }
+
+    pub fn forward(&self, xs: &Tensor) -> Result<Tensor> {
+        let xs = if xs.rank() == 2 {
+            self.linear
+                .forward(&xs.unsqueeze(D::Minus1)?)?
+                .squeeze(D::Minus1)?
+        } else {
+            self.linear.forward(&xs)?
+        };
+        let xs = self.nonlinear.forward_t(&xs, false)?;
+        Ok(xs)
+    }
+}
+
+pub struct XVector {
+    tdnn: TDNNLayer,
+    blocks: Vec<CAMDenseTDNNBlock>,
+    transits: Vec<TransitLayer>,
+    out_nonlinear: BatchNorm,
+    dense: DenseLayer,
+}
+
+impl XVector {
+    pub fn new(
+        vb: VarBuilder,
+        channels: usize,
+        init_channels: usize,
+        growth_rate: usize,
+        bn_size: usize,
+        embedding_size: usize,
+    ) -> Result<Self> {
+        let tdnn = TDNNLayer::new(vb.pp("tdnn"), channels, init_channels, 5, 2, 1, false)?;
+        let mut channels = init_channels;
+        let mut blocks = vec![];
+        let mut transits = vec![];
+        let params = vec![(12, 3, 1), (24, 3, 2), (16, 3, 2)];
+        for (i, (num_layers, ks, dilation)) in params.iter().enumerate() {
+            let block = CAMDenseTDNNBlock::new(
+                vb.pp(format!("block{}", i + 1)),
+                *num_layers,
+                channels,
+                growth_rate,
+                bn_size * growth_rate,
+                *ks,
+                1,
+                *dilation,
+                false,
+            )?;
+            blocks.push(block);
+            channels = channels + num_layers * growth_rate;
+            let transit = TransitLayer::new(
+                vb.pp(format!("transit{}", i + 1)),
+                channels,
+                channels / 2,
+                false,
+            )?;
+            transits.push(transit);
+            channels /= 2;
+        }
+        let out_nonlinear = get_batch_norm(vb.pp("out_nonlinear.batchnorm"), 1e-5, channels, true)?;
+        let dense = DenseLayer::new(vb.pp("dense"), channels * 2, embedding_size, false)?;
+        Ok(Self {
+            tdnn,
+            blocks,
+            transits,
+            out_nonlinear,
+            dense,
+        })
+    }
+
+    pub fn forward(&self, xs: &Tensor) -> Result<Tensor> {
+        let mut xs = self.tdnn.forward(xs)?;
+        for i in 0..3 {
+            let block = &self.blocks[i];
+            xs = block.forward(&xs)?;
+            let transit = &self.transits[i];
+            xs = transit.forward(&xs)?;
+        }
+        xs = self.out_nonlinear.forward_t(&xs, false)?.relu()?;
+        xs = statistics_pooling(&xs, D::Minus1, false)?;
+        xs = self.dense.forward(&xs)?;
+        Ok(xs)
+    }
+}
+
+pub struct CAMPPlus {
+    head: FCM,
+    xvector: XVector,
+}
+
+impl CAMPPlus {
+    pub fn new(
+        vb: VarBuilder,
+        feat_dim: usize,
+        embedding_size: usize,
+        growth_rate: usize,
+        bn_size: usize,
+        init_channels: usize,
+    ) -> Result<Self> {
+        let head = FCM::new(vb.pp("head"), &[2, 2], 32, feat_dim)?;
+        let channels = head.out_channels;
+        let xvector = XVector::new(
+            vb.pp("xvector"),
+            channels,
+            init_channels,
+            growth_rate,
+            bn_size,
+            embedding_size,
+        )?;
+        Ok(Self { head, xvector })
+    }
+
+    pub fn forward(&self, xs: &Tensor) -> Result<Tensor> {
+        let xs = xs.permute((0, 2, 1))?;
+        let xs = self.head.forward(&xs)?;
+        let xs = self.xvector.forward(&xs)?;
+        Ok(xs)
+    }
+}

src/models/common/mod.rs

@@ -1,15 +1,15 @@
 use anyhow::{Result, anyhow};
-use candle_core::{D, Tensor};
+use candle_core::{D, IndexOp, Tensor};
 use candle_nn::{
-    Activation, BatchNorm, BatchNormConfig, Conv1d, Conv1dConfig, Conv2d, Conv2dConfig, Embedding,
+    Activation, BatchNorm, BatchNormConfig, Conv1d, Conv1dConfig, Conv2d, Conv2dConfig,
-    LayerNorm, LayerNormConfig, Linear, Module, RmsNorm, VarBuilder, batch_norm, conv1d,
+    ConvTranspose1d, ConvTranspose1dConfig, Embedding, LayerNorm, LayerNormConfig, Linear, Module,
-    conv1d_no_bias, conv2d, conv2d_no_bias, embedding, layer_norm, linear_b, linear_no_bias,
+    ModuleT, RmsNorm, VarBuilder, batch_norm, conv1d, conv1d_no_bias, conv2d, conv2d_no_bias,
-    rms_norm,
+    embedding, layer_norm, linear_b, linear_no_bias, ops::sigmoid, rms_norm,
 };
 use rayon::iter::{IndexedParallelIterator, IntoParallelRefIterator, ParallelIterator};
 
 use crate::{
-    position_embed::rope::{RoPE, apply_rotary_pos_emb},
+    position_embed::rope::{RoPE, apply_rotary_pos_emb, apply_rotary_pos_emb_roformer},
     utils::tensor_utils::{prepare_causal_attention_mask, repeat_kv},
 };
 
@@ -28,10 +28,16 @@ impl GateUpDownMLP {
         intermediate_size: usize,
         act_fn: Activation,
         bias: bool,
+        gate_pp_name: Option<&str>,
+        up_pp_name: Option<&str>,
+        down_pp_name: Option<&str>,
     ) -> Result<Self> {
-        let gate_proj = linear_b(hidden_size, intermediate_size, bias, vb.pp("gate_proj"))?;
+        let gate_pp_name = gate_pp_name.unwrap_or("gate_proj");
-        let up_proj = linear_b(hidden_size, intermediate_size, bias, vb.pp("up_proj"))?;
+        let up_pp_name = up_pp_name.unwrap_or("up_proj");
-        let down_proj = linear_b(intermediate_size, hidden_size, bias, vb.pp("down_proj"))?;
+        let down_pp_name = down_pp_name.unwrap_or("down_proj");
+        let gate_proj = linear_b(hidden_size, intermediate_size, bias, vb.pp(gate_pp_name))?;
+        let up_proj = linear_b(hidden_size, intermediate_size, bias, vb.pp(up_pp_name))?;
+        let down_proj = linear_b(intermediate_size, hidden_size, bias, vb.pp(down_pp_name))?;
         Ok(Self {
             gate_proj,
             up_proj,
@@ -87,7 +93,6 @@ impl TwoLinearMLP {
 }
 
 #[derive(Debug, Clone)]
-// pub struct AttentionNobias {
 pub struct NaiveAttention {
     q_proj: Linear,
     k_proj: Linear,
@@ -257,6 +262,154 @@ impl NaiveAttention {
     }
 }
 
+#[derive(Debug, Clone)]
+pub struct QKVCatAttention {
+    qkv_proj: Linear,
+    o_proj: Linear,
+    num_heads: usize,
+    head_dim: usize,
+    middle_size: usize,
+    kv_cache: Option<(Tensor, Tensor)>,
+}
+
+impl QKVCatAttention {
+    pub fn new(
+        vb: VarBuilder,
+        hidden_size: usize,
+        num_attention_heads: usize,
+        head_dim: Option<usize>,
+        bias: bool,
+        qkv_proj_pp_name: Option<&str>,
+        o_proj_pp_name: Option<&str>,
+    ) -> Result<Self> {
+        let head_dim = match head_dim {
+            None => hidden_size / num_attention_heads,
+            Some(dim) => dim,
+        };
+        let qkv_proj_pp_name = qkv_proj_pp_name.unwrap_or("wqkv");
+        let o_proj_pp_name = o_proj_pp_name.unwrap_or("o_proj");
+        let qkv_proj = linear_b(
+            hidden_size,
+            3 * num_attention_heads * head_dim,
+            bias,
+            vb.pp(qkv_proj_pp_name),
+        )?;
+        let o_proj = linear_b(
+            num_attention_heads * head_dim,
+            hidden_size,
+            bias,
+            vb.pp(o_proj_pp_name),
+        )?;
+
+        Ok(Self {
+            qkv_proj,
+            o_proj,
+            num_heads: num_attention_heads,
+            head_dim,
+            middle_size: num_attention_heads * head_dim,
+            kv_cache: None,
+        })
+    }
+
+    pub fn forward(
+        &self,
+        xs: &Tensor,
+        cos: Option<&Tensor>,
+        sin: Option<&Tensor>,
+        attention_mask: Option<&Tensor>,
+        tof32: bool,
+        use_roformer: bool,
+    ) -> Result<Tensor> {
+        let (b, q_len, _) = xs.dims3()?;
+        // (3, B, n_head, seq_len, head_dim)
+        let qkv = self
+            .qkv_proj
+            .forward(xs)?
+            .reshape((b, q_len, 3, self.num_heads, ()))?
+            .permute((2, 0, 3, 1, 4))?
+            .contiguous()?;
+        let query_states = qkv.i(0)?.contiguous()?;
+        let key_states = qkv.i(1)?.contiguous()?;
+        let value_states = qkv.i(2)?.contiguous()?;
+        let (query_states, key_states) = if let Some(cos) = cos
+            && let Some(sin) = sin
+        {
+            if use_roformer {
+                apply_rotary_pos_emb_roformer(&query_states, &key_states, cos, sin, tof32)?
+            } else {
+                apply_rotary_pos_emb(&query_states, &key_states, cos, sin, tof32)?
+            }
+        } else {
+            (query_states, key_states)
+        };
+
+        let scale = 1f64 / f64::sqrt(self.head_dim as f64);
+        let attn_output = eager_attention_forward(
+            &query_states,
+            &key_states,
+            &value_states,
+            None,
+            attention_mask,
+            scale,
+        )?;
+        let attn_output = attn_output.reshape((b, q_len, self.middle_size))?;
+        let attn_output = attn_output.apply(&self.o_proj)?;
+        Ok(attn_output)
+    }
+
+    pub fn forward_with_cache(
+        &mut self,
+        xs: &Tensor,
+        cos: &Tensor,
+        sin: &Tensor,
+        attention_mask: Option<&Tensor>,
+        tof32: bool,
+        use_roformer: bool,
+    ) -> Result<Tensor> {
+        let (b, q_len, _) = xs.dims3()?;
+        let qkv = self
+            .qkv_proj
+            .forward(xs)?
+            .reshape((b, q_len, 3, self.num_heads, ()))?
+            .permute((2, 0, 3, 1, 4))?
+            .contiguous()?;
+        let query_states = qkv.i(0)?.contiguous()?;
+        let key_states = qkv.i(1)?.contiguous()?;
+        let value_states = qkv.i(2)?.contiguous()?;
+        let (query_states, key_states) = if use_roformer {
+            apply_rotary_pos_emb_roformer(&query_states, &key_states, cos, sin, tof32)?
+        } else {
+            apply_rotary_pos_emb(&query_states, &key_states, cos, sin, tof32)?
+        };
+        let (key_states, value_states) = match &self.kv_cache {
+            None => (key_states, value_states),
+            Some((prev_k, prev_v)) => {
+                let key_states = Tensor::cat(&[prev_k, &key_states], 2)?;
+                let value_states = Tensor::cat(&[prev_v, &value_states], 2)?;
+                (key_states, value_states)
+            }
+        };
+
+        self.kv_cache = Some((key_states.clone(), value_states.clone()));
+        let scale = 1f64 / f64::sqrt(self.head_dim as f64);
+        let attn_output = eager_attention_forward(
+            &query_states,
+            &key_states,
+            &value_states,
+            None,
+            attention_mask,
+            scale,
+        )?;
+        let attn_output = attn_output.reshape((b, q_len, self.middle_size))?;
+        let attn_output = attn_output.apply(&self.o_proj)?;
+        Ok(attn_output)
+    }
+
+    pub fn clear_kv_cache(&mut self) {
+        self.kv_cache = None
+    }
+}
+
 pub struct NaiveAttnTwoLinearMLPBlock {
     self_attn: NaiveAttention,
     mlp: TwoLinearMLP,
@@ -390,6 +543,9 @@ impl NaiveAttnGateUpDownMLPBlock {
             intermediate_size,
             hidden_act,
             mlp_bias,
+            None,
+            None,
+            None,
         )?;
         let input_layernorm = rms_norm(hidden_size, norm_eps, vb.pp(input_norm_pp_name))?;
         let post_attention_layernorm = rms_norm(hidden_size, norm_eps, vb.pp(post_norm_pp_name))?;
@@ -543,11 +699,11 @@ pub fn get_layer_norm(vb: VarBuilder, eps: f64, dim: usize) -> Result<LayerNorm>
     Ok(norm)
 }
 
-pub fn get_batch_norm(vb: VarBuilder, eps: f64, dim: usize) -> Result<BatchNorm> {
+pub fn get_batch_norm(vb: VarBuilder, eps: f64, dim: usize, affine: bool) -> Result<BatchNorm> {
     let bn_config = BatchNormConfig {
         eps,
         remove_mean: true,
-        affine: true,
+        affine,
         momentum: 0.1,
     };
     let norm = batch_norm(dim, bn_config, vb)?;
@@ -850,3 +1006,164 @@ pub fn conv1d_group_parallel(xs: &Tensor, conv1d: &Conv1d) -> Result<Tensor> {
         }
     }
 }
+
+pub struct GLU {
+    dim: usize,
+}
+
+impl GLU {
+    pub fn new(dim: usize) -> Result<Self> {
+        Ok(Self { dim })
+    }
+    pub fn forward(&self, xs: &Tensor) -> Result<Tensor> {
+        let half_dim = xs.dim(self.dim)? / 2;
+        let a = xs.narrow(self.dim, 0, half_dim)?;
+        let b = xs.narrow(self.dim, half_dim, half_dim)?;
+        let b = sigmoid(&b)?;
+        let xs = a.mul(&b)?;
+        Ok(xs)
+    }
+}
+
+pub struct WNConv1d {
+    conv: Conv1d,
+}
+impl WNConv1d {
+    pub fn new(
+        vb: VarBuilder,
+        in_c: usize,
+        out_c: usize,
+        kernel_size: usize,
+        dilation: usize,
+        padding: usize,
+        groups: usize,
+        stride: usize,
+        bias: bool,
+    ) -> Result<Self> {
+        let in_c = in_c / groups;
+        let weight_g = vb.get((out_c, 1, 1), "weight_g")?;
+        let weight_v = vb.get((out_c, in_c, kernel_size), "weight_v")?;
+        // let bias = vb.get(out_c, "bias").ok();
+        let bias = if bias {
+            vb.get(out_c, "bias").ok()
+        } else {
+            None
+        };
+        let weight_norm = weight_v.sqr()?.sum_keepdim(1)?.sum_keepdim(2)?.sqrt()?;
+        let normalized_weight = weight_v.broadcast_div(&weight_norm)?;
+        let scaled_weight = normalized_weight.broadcast_mul(&weight_g)?;
+        let cfg = Conv1dConfig {
+            padding,
+            stride,
+            dilation,
+            groups,
+            cudnn_fwd_algo: None,
+        };
+        let conv = Conv1d::new(scaled_weight, bias, cfg);
+        Ok(Self { conv })
+    }
+    pub fn forward(&self, x: &Tensor) -> Result<Tensor> {
+        let x = self.conv.forward(x)?;
+        Ok(x)
+    }
+}
+
+pub struct WNConvTranspose1d {
+    conv_transpose: ConvTranspose1d,
+}
+
+impl WNConvTranspose1d {
+    pub fn new(
+        vb: VarBuilder,
+        in_c: usize,
+        out_c: usize,
+        dilation: usize,
+        kernel_size: usize,
+        padding: usize,
+        output_padding: usize,
+        groups: usize,
+        stride: usize,
+    ) -> Result<Self> {
+        let in_c = in_c / groups;
+        let weight_g = vb.get((in_c, 1, 1), "weight_g")?;
+        let weight_v = vb.get((in_c, out_c, kernel_size), "weight_v")?;
+        let bias = vb.get(out_c, "bias").ok();
+        let weight_norm = weight_v.sqr()?.sum_keepdim(1)?.sum_keepdim(2)?.sqrt()?;
+        let normalized_weight = weight_v.broadcast_div(&weight_norm)?;
+        let scaled_weight = normalized_weight.broadcast_mul(&weight_g)?;
+        let config = ConvTranspose1dConfig {
+            padding: padding,
+            output_padding: output_padding,
+            stride,
+            dilation,
+            groups,
+        };
+        let conv_transpose = ConvTranspose1d::new(scaled_weight, bias, config);
+        Ok(Self { conv_transpose })
+    }
+    pub fn forward(&self, x: &Tensor) -> Result<Tensor> {
+        let x = self.conv_transpose.forward(x)?;
+        Ok(x)
+    }
+}
+
+pub struct Conv2dWithBN {
+    conv_0: Conv2d,
+    bn_1: BatchNorm,
+    bn_with_relu: bool,
+}
+
+impl Conv2dWithBN {
+    pub fn new(
+        vb: VarBuilder,
+        in_c: usize,
+        out_c: usize,
+        ks: usize,
+        padding: usize,
+        stride: usize,
+        bias: bool,
+        bn_with_relu: bool,
+    ) -> Result<Self> {
+        let conv_0 = get_conv2d(vb.pp("0"), in_c, out_c, ks, padding, stride, 1, 1, bias)?;
+        let bn_1 = get_batch_norm(vb.pp("1"), 1e-5, out_c, true)?;
+        Ok(Self {
+            conv_0,
+            bn_1,
+            bn_with_relu,
+        })
+    }
+
+    pub fn forward(&self, x: &Tensor) -> Result<Tensor> {
+        let x = self.conv_0.forward(x)?;
+        let mut x = self.bn_1.forward_t(&x, false)?;
+        if self.bn_with_relu {
+            x = x.relu()?;
+        }
+        Ok(x)
+    }
+}
+
+pub struct WNLinear {
+    linear: Linear,
+}
+impl WNLinear {
+    pub fn new(vb: VarBuilder, in_dim: usize, out_dim: usize, bias: bool) -> Result<Self> {
+        let weight_g = vb.get((out_dim, 1), "weight_g")?;
+        let weight_v = vb.get((out_dim, in_dim), "weight_v")?;
+
+        let bias = if bias {
+            vb.get(out_dim, "bias").ok()
+        } else {
+            None
+        };
+        let weight_norm = weight_v.sqr()?.sum_keepdim(0)?.sqrt()?.affine(1.0, 1e-8)?;
+        let normalized_weight = weight_v.broadcast_div(&weight_norm)?;
+        let scaled_weight = normalized_weight.broadcast_mul(&weight_g)?;
+        let linear = Linear::new(scaled_weight, bias);
+        Ok(Self { linear })
+    }
+    pub fn forward(&self, x: &Tensor) -> Result<Tensor> {
+        let x = self.linear.forward(x)?;
+        Ok(x)
+    }
+}

src/models/deepseek_ocr/model.rs

@@ -875,6 +875,9 @@ impl DeepseekV2MoE {
                 config.moe_intermediate_size,
                 Activation::Silu,
                 false,
+                None,
+                None,
+                None,
             )?;
             experts.push(mlp);
         }
@@ -885,6 +888,9 @@ impl DeepseekV2MoE {
             config.moe_intermediate_size * config.n_shared_experts,
             Activation::Silu,
             false,
+            None,
+            None,
+            None,
         )?;
         Ok(Self {
             // num_experts_per_tok: config.num_experts_per_tok,
@@ -998,6 +1004,9 @@ impl DeepseekV2DecoderLayer {
                 config.intermediate_size,
                 Activation::Silu,
                 false,
+                None,
+                None,
+                None,
             )?)
         };
         let input_layernorm = rms_norm(

src/models/feature_extractor/config.rs

@@ -0,0 +1,21 @@
+use serde::Deserialize;
+
+#[derive(Debug, Clone, PartialEq, Deserialize)]
+pub struct FeatureExtractor {
+    pub chunk_length: usize,
+    pub dither: f64,
+    pub feature_size: usize,
+    pub hop_length: usize,
+    pub n_fft: usize,
+    pub n_samples: usize,
+    pub nb_max_frames: usize,
+    pub padding_side: String,
+    pub padding_value: f32,
+    pub return_attention_mask: bool,
+    #[serde(default = "default_sampling_rate")]
+    pub sampling_rate: usize,
+}
+
+fn default_sampling_rate() -> usize {
+    16000
+}

src/models/feature_extractor/feature_extraction_whisper.rs

@@ -0,0 +1,113 @@
+use anyhow::Result;
+use candle_core::{D, Device, Tensor};
+
+use crate::utils::{
+    audio_utils::{create_hann_window, mel_filter_bank, torch_stft},
+    tensor_utils::{log10, pad_reflect_last_dim},
+};
+
+pub struct WhisperFeatureExtractor {
+    // feature_size: usize,
+    hop_length: usize,
+    // chunk_length: usize,
+    // n_samples: usize,
+    n_fft: usize,
+    dither: f64,
+    // padding_value: f32,
+    sampling_rate: usize,
+    mel_filters: Tensor,
+    window: Tensor,
+}
+
+impl WhisperFeatureExtractor {
+    pub fn new(
+        feature_size: usize,
+        hop_length: usize,
+        // chunk_length: usize,
+        n_fft: usize,
+        dither: f64,
+        // padding_value: f32,
+        sampling_rate: usize,
+        device: &Device,
+    ) -> Result<Self> {
+        let window = create_hann_window(n_fft, candle_core::DType::F32, device)?;
+        let window = window.unsqueeze(0)?.unsqueeze(0)?;
+        let mel_filters = mel_filter_bank(
+            1 + n_fft / 2,
+            feature_size,
+            0.0,
+            8000.0,
+            sampling_rate as f32,
+            Some("slaney"),
+            crate::utils::audio_utils::MelScale::Slaney,
+            false,
+            device,
+        )?
+        .t()?;
+        // let n_samples = chunk_length * sampling_rate;
+        Ok(Self {
+            // feature_size,
+            hop_length,
+            // chunk_length,
+            // n_samples,
+            n_fft,
+            dither,
+            // padding_value,
+            sampling_rate,
+            mel_filters,
+            window,
+        })
+    }
+
+    pub fn call(
+        &self,
+        raw_speech: &Tensor,
+        sampling_rate: usize,
+        // do_normalize: bool,
+        return_attention_mask: bool,
+    ) -> Result<(Tensor, Option<Tensor>)> {
+        // raw_speech: 重采样后的音频，shape: (bs, raw_len)
+        // sampling_rate: 音频采样率，验证是否与模型的预处理采样率一致
+        if sampling_rate != self.sampling_rate {
+            return Err(anyhow::anyhow!(
+                "The model feature extractor was trained sampling rate {} not equal to audio sample rate {}",
+                self.sampling_rate,
+                sampling_rate
+            ));
+        }
+
+        let input_features = self.extract_fbank_features(raw_speech)?;
+        let mask_len = input_features.dim(2)?;
+        let mask = if return_attention_mask {
+            let mask = Tensor::new(1u32, input_features.device())?.broadcast_as((1, mask_len))?;
+            Some(mask)
+        } else {
+            None
+        };
+        Ok((input_features, mask))
+    }
+
+    pub fn extract_fbank_features(&self, waveform: &Tensor) -> Result<Tensor> {
+        let mut waveform = waveform.clone();
+        if self.dither != 0.0 {
+            waveform = waveform.add(&waveform.randn_like(0.0, 1.0)?.affine(self.dither, 0.0)?)?;
+        }
+        let pad = self.n_fft / 2;
+        let waveform = pad_reflect_last_dim(&waveform, (pad, pad))?;
+        let (_, samples) = waveform.dims2()?;
+
+        let magnitudes = torch_stft(&waveform, self.n_fft, self.hop_length, &self.window)?
+            .transpose(D::Minus1, D::Minus2)?;
+        let n_frames = (samples - self.n_fft) / self.hop_length + 1;
+        let magnitudes = magnitudes.narrow(D::Minus1, 0, n_frames - 1)?;
+        let mel_spec = self.mel_filters.broadcast_matmul(&magnitudes)?;
+        let mel_spec = mel_spec.clamp(1e-10f32, f32::INFINITY)?;
+        // let ln_spec = mel_spec.log()?;
+        // let log10_spec = ln_spec.broadcast_div(&Tensor::new(f32::ln(10.0), mel_spec.device())?)?;
+        let log10_spec = log10(&mel_spec)?;
+        let max_val = log10_spec.max_all()?.affine(1.0, -8.0)?;
+        let log10_spec = log10_spec.broadcast_maximum(&max_val)?;
+        let log_spec = log10_spec.affine(1.0, 4.0)?.affine(1.0 / 4.0, 0.0)?;
+        Ok(log_spec)
+    }
+}

src/models/feature_extractor/mod.rs

@@ -0,0 +1,3 @@
+pub mod seamless_m4t_feature_extractor;
+pub mod feature_extraction_whisper;
+pub mod config;

src/models/feature_extractor/seamless_m4t_feature_extractor.rs

@@ -0,0 +1,129 @@
+use anyhow::Result;
+use candle_core::{D, Device, Tensor};
+
+use crate::utils::{
+    audio_utils::{create_povey_window, mel_filter_bank, spectrogram},
+    tensor_utils::{PaddingSide, z_score_normalize},
+};
+
+pub struct SeamlessM4TFeatureExtractor {
+    // feature_size: usize,
+    num_mel_bins: usize,
+    padding_side: PaddingSide,
+    padding_value: f32,
+    sampling_rate: usize,
+    stride: usize,
+    mel_filters: Tensor,
+    window: Tensor,
+}
+
+impl SeamlessM4TFeatureExtractor {
+    pub fn new(
+        // feature_size: usize,
+        num_mel_bins: usize,
+        padding_side: PaddingSide,
+        padding_value: f32,
+        sampling_rate: usize,
+        stride: usize,
+        device: &Device,
+    ) -> Result<Self> {
+        let mel_filters = mel_filter_bank(
+            257,
+            num_mel_bins,
+            20.0,
+            (sampling_rate / 2) as f32,
+            sampling_rate as f32,
+            None,
+            crate::utils::audio_utils::MelScale::Kaldi,
+            true,
+            device,
+        )?;
+        let window = create_povey_window(400, candle_core::DType::F32, device)?;
+        Ok(Self {
+            // feature_size,
+            num_mel_bins,
+            padding_side,
+            padding_value,
+            sampling_rate,
+            stride,
+            mel_filters,
+            window,
+        })
+    }
+
+    pub fn call(
+        &self,
+        raw_speech: &Tensor,
+        sampling_rate: usize,
+        do_normalize_per_mel_bins: bool,
+        return_attention_mask: bool,
+    ) -> Result<(Tensor, Option<Tensor>)> {
+        // raw_speech: 重采样后的音频，shape: (bs, raw_len)
+        // sampling_rate: 音频采样率，验证是否与模型的预处理采样率一致
+        if sampling_rate != self.sampling_rate {
+            return Err(anyhow::anyhow!(
+                "The model feature extractor was trained sampling rate {} not equal to audio sample rate {}",
+                self.sampling_rate,
+                sampling_rate
+            ));
+        }
+        let waveform = raw_speech.affine(32768.0, 0.0)?;
+        // println!("waveform: {}", waveform);
+        // println!("self.mel_filters: {}", self.mel_filters);
+        // println!("self.window: {}", self.window);
+        let mut features = spectrogram(
+            &waveform,
+            &self.window,
+            400,
+            160,
+            512,
+            Some(2.0),
+            false,
+            0.97,
+            Some(&self.mel_filters),
+            Some("log"),
+            1.192092955078125e-07,
+            true,
+        )?
+        .transpose(D::Minus1, D::Minus2)?;
+        if do_normalize_per_mel_bins {
+            features = z_score_normalize(&features, 1)?;
+        }
+        let n_frame = features.dim(1)?;
+        let mask_1 = n_frame / self.stride;
+        let pad_len = n_frame % self.stride;
+        if pad_len > 0 {
+            let pad = Tensor::new(self.padding_value, features.device())?.broadcast_as((
+                1,
+                pad_len,
+                self.num_mel_bins,
+            ))?;
+            match self.padding_side {
+                PaddingSide::Left => features = Tensor::cat(&[pad, features], 1)?,
+                PaddingSide::Right => features = Tensor::cat(&[features, pad], 1)?,
+            }
+        }
+        let (bs, num_frames, dim) = features.dims3()?;
+        let n_frames_stride = num_frames / self.stride;
+        let features = features.reshape((bs, n_frames_stride, dim * self.stride))?;
+        let mask_0 = n_frames_stride - mask_1;
+        let mask = if return_attention_mask {
+            let mut mask = Tensor::new(1u32, features.device())?.broadcast_as((1, mask_1))?;
+            if mask_0 > 0 {
+                let mask_pad = Tensor::new(0u32, features.device())?.broadcast_as((1, mask_0))?;
+                match self.padding_side {
+                    PaddingSide::Left => {
+                        mask = Tensor::cat(&[mask_pad, mask], D::Minus1)?;
+                    }
+                    PaddingSide::Right => {
+                        mask = Tensor::cat(&[mask, mask_pad], D::Minus1)?;
+                    }
+                }
+            }
+            Some(mask)
+        } else {
+            None
+        };
+        Ok((features, mask))
+    }
+}

src/models/fun_asr_nano/model.rs

@@ -11,7 +11,7 @@ use crate::{
         qwen3::{config::Qwen3Config, model::Qwen3Model},
     },
     position_embed::sinusoidal_pe::SinusoidalPositionEncoderCat,
-    utils::tensor_utils::{get_equal_mask, mask_filled, masked_scatter_dim0},
+    utils::tensor_utils::{attn_masked_fill, get_equal_mask, masked_scatter_dim0},
 };
 
 pub struct MultiHeadedAttentionSANM {
@@ -124,9 +124,9 @@ impl MultiHeadedAttentionSANM {
             };
             // mask: rank = 2
             let mask = get_equal_mask(&mask, 0)?;
-            let scores = mask_filled(scores, &mask, f32::NEG_INFINITY)?;
+            let scores = attn_masked_fill(scores, &mask, f32::NEG_INFINITY)?;
             let attn = softmax_last_dim(&scores)?;
-            mask_filled(&attn, &mask, 0.0)?
+            attn_masked_fill(&attn, &mask, 0.0)?
         } else {
             softmax_last_dim(scores)?
         };

src/models/glm_asr_nano/config.rs

@@ -1,6 +1,8 @@
 use candle_nn::Activation;
 use serde::Deserialize;
 
+use crate::models::feature_extractor::config::FeatureExtractor;
+
 #[derive(Debug, Clone, PartialEq, Deserialize)]
 pub struct GlmAsrNanoProcessorConfig {
     pub audio_token: String,
@@ -9,20 +11,7 @@ pub struct GlmAsrNanoProcessorConfig {
     pub max_audio_len: usize,
 }
 
-#[derive(Debug, Clone, PartialEq, Deserialize)]
+
-pub struct FeatureExtractor {
-    pub chunk_length: usize,
-    pub dither: f32,
-    pub feature_size: usize,
-    pub hop_length: usize,
-    pub n_fft: usize,
-    pub n_samples: usize,
-    pub nb_max_frames: usize,
-    pub padding_side: String,
-    pub padding_value: f32,
-    pub return_attention_mask: bool,
-    pub sampling_rate: usize,
-}
 
 #[derive(Debug, Clone, PartialEq, Deserialize)]
 pub struct GlmAsrNanoConfig {

src/models/glm_asr_nano/processor.rs

@@ -1,32 +1,30 @@
-use std::f32;
 
 use aha_openai_dive::v1::resources::chat::ChatCompletionParameters;
 use anyhow::Result;
 use candle_core::{D, DType, Device, IndexOp, Tensor};
 
 use crate::{
-    models::glm_asr_nano::config::GlmAsrNanoProcessorConfig,
+    models::{
-    utils::{
+        feature_extractor::feature_extraction_whisper::WhisperFeatureExtractor,
-        audio_utils::{
+        glm_asr_nano::config::GlmAsrNanoProcessorConfig,
-            apply_stft, create_hann_window, extract_audios, extract_frames, mel_filter_bank,
-        },
-        tensor_utils::{pad_reflect_last_dim, split_tensor},
     },
+    utils::{audio_utils::extract_audios, tensor_utils::split_tensor},
 };
 
 pub struct GlmAsrNanoProcessor {
     sampling_rate: usize,
     chunk_length: usize,
     n_samples: usize,
-    n_fft: usize,
+    // n_fft: usize,
-    window: Tensor,
+    // window: Tensor,
-    mel_filters: Tensor,
+    // mel_filters: Tensor,
     hop_length: usize,
     audio_token: String,
     // audio_token_id: u32,
     max_audio_len: usize,
     // default_transcription_prompt: String,
     device: Device,
+    whisper_feature_extrator: WhisperFeatureExtractor,
 }
 
 impl GlmAsrNanoProcessor {
@@ -50,61 +48,74 @@ impl GlmAsrNanoProcessor {
         let sampling_rate = processor_cfg.feature_extractor.sampling_rate;
         let chunk_length = processor_cfg.feature_extractor.chunk_length;
         let n_samples = processor_cfg.feature_extractor.n_samples;
-        let n_fft = processor_cfg.feature_extractor.n_fft;
+        // let n_fft = processor_cfg.feature_extractor.n_fft;
         let hop_length = processor_cfg.feature_extractor.hop_length;
-        let window = create_hann_window(n_fft, dtype, device)?;
+        // let window = create_hann_window(n_fft, dtype, device)?;
-        let window = window.unsqueeze(0)?.unsqueeze(0)?;
+        // let window = window.unsqueeze(0)?.unsqueeze(0)?;
-        let mel_filters = mel_filter_bank(
+        // let mel_filters = mel_filter_bank(
-            1 + n_fft / 2,
+        //     1 + n_fft / 2,
+        //     processor_cfg.feature_extractor.feature_size,
+        //     0.0,
+        //     8000.0,
+        //     sampling_rate as f32,
+        //     Some("slaney"),
+        //     crate::utils::audio_utils::MelScale::Slaney,
+        //     false,
+        //     device,
+        // )?
+        // .t()?;
+        let whisper_feature_extrator = WhisperFeatureExtractor::new(
             processor_cfg.feature_extractor.feature_size,
-            0.0,
+            processor_cfg.feature_extractor.hop_length,
-            8000.0,
+            // processor_cfg.feature_extractor.chunk_length,
-            sampling_rate as f32,
+            processor_cfg.feature_extractor.n_fft,
-            Some("slaney"),
+            processor_cfg.feature_extractor.dither,
-            crate::utils::audio_utils::MelScale::Slaney,
+            // processor_cfg.feature_extractor.padding_value,
-            false,
+            processor_cfg.feature_extractor.sampling_rate,
             device,
-        )?
+        )?;
-        .t()?;
         Ok(Self {
             sampling_rate,
             chunk_length,
             n_samples,
-            n_fft,
+            // n_fft,
-            window,
+            // window,
-            mel_filters,
+            // mel_filters,
             hop_length,
             audio_token,
             // audio_token_id,
             max_audio_len,
             // default_transcription_prompt,
             device: device.clone(),
+            whisper_feature_extrator,
         })
     }
 
-    pub fn extract_fbank_features(&self, waveform: &Tensor) -> Result<Tensor> {
+    // pub fn extract_fbank_features(&self, waveform: &Tensor) -> Result<Tensor> {
-        let pad = self.n_fft / 2;
+    //     let pad = self.n_fft / 2;
-        let waveform = pad_reflect_last_dim(waveform, (pad, pad))?;
+    //     let waveform = pad_reflect_last_dim(waveform, (pad, pad))?;
-        let (_, samples) = waveform.dims2()?;
+    //     let (_, samples) = waveform.dims2()?;
 
-        // 计算输出维度
+    //     // // (bs, n_frames, n_fft)
-        let n_frames = (samples - self.n_fft) / self.hop_length + 1;
+    //     // let frames = extract_frames(&waveform, self.n_fft, self.hop_length)?;
-        // (bs, n_frames, n_fft)
+    //     // // 应用汉明窗口
-        let frames = extract_frames(&waveform, self.n_fft, self.hop_length)?;
+    //     // let result = frames.broadcast_mul(&self.window)?;
-        // 应用汉明窗口
+    //     // // 傅立叶变换
-        let result = frames.broadcast_mul(&self.window)?;
+    //     // let magnitudes = apply_stft(&result)?.transpose(D::Minus1, D::Minus2)?;
-        // 傅立叶变换
+    //     let magnitudes = torch_stft(&waveform, self.n_fft, self.hop_length, &self.window)?
-        let magnitudes = apply_stft(&result)?.transpose(D::Minus1, D::Minus2)?;
+    //         .transpose(D::Minus1, D::Minus2)?;
-        let magnitudes = magnitudes.narrow(D::Minus1, 0, n_frames - 1)?;
+    //     let n_frames = (samples - self.n_fft) / self.hop_length + 1;
-        let mel_spec = self.mel_filters.broadcast_matmul(&magnitudes)?;
+    //     let magnitudes = magnitudes.narrow(D::Minus1, 0, n_frames - 1)?;
-        let mel_spec = mel_spec.clamp(1e-10f32, f32::INFINITY)?;
+    //     let mel_spec = self.mel_filters.broadcast_matmul(&magnitudes)?;
-        let ln_spec = mel_spec.log()?;
+    //     let mel_spec = mel_spec.clamp(1e-10f32, f32::INFINITY)?;
-        let log10_spec = ln_spec.broadcast_div(&Tensor::new(f32::ln(10.0), mel_spec.device())?)?;
+    //     // let ln_spec = mel_spec.log()?;
-        let max_val = log10_spec.max_all()?.affine(1.0, -8.0)?;
+    //     // let log10_spec = ln_spec.broadcast_div(&Tensor::new(f32::ln(10.0), mel_spec.device())?)?;
-        let log10_spec = log10_spec.broadcast_maximum(&max_val)?;
+    //     let log10_spec = log10(&mel_spec)?;
-        let log_spec = log10_spec.affine(1.0, 4.0)?.affine(1.0 / 4.0, 0.0)?;
+    //     let max_val = log10_spec.max_all()?.affine(1.0, -8.0)?;
-        Ok(log_spec)
+    //     let log10_spec = log10_spec.broadcast_maximum(&max_val)?;
-    }
+    //     let log_spec = log10_spec.affine(1.0, 4.0)?.affine(1.0 / 4.0, 0.0)?;
+    //     Ok(log_spec)
+    // }
 
     pub fn feature_extractor(&self, raw_speech: Vec<Tensor>) -> Result<(Tensor, Tensor)> {
         let mut pad_audio = vec![];
@@ -129,7 +140,10 @@ impl GlmAsrNanoProcessor {
         }
         let input_features = Tensor::cat(&pad_audio, 0)?;
         let input_features_mask = Tensor::new(input_features_mask, input_features.device())?;
-        let input_features = self.extract_fbank_features(&input_features)?;
+        // let input_features = self.extract_fbank_features(&input_features)?;
+        let (input_features, _) =
+            self.whisper_feature_extrator
+                .call(&input_features, self.sampling_rate, false)?;
         let (_, audio_len) = input_features_mask.dims2()?;
         let mask_idx: Vec<u32> = (0..audio_len)
             .step_by(self.hop_length)

src/models/hunyuan_ocr/model.rs

@@ -406,6 +406,9 @@ impl HunYuanVLDecoderLayer {
             config.intermediate_size,
             config.hidden_act,
             false,
+            None,
+            None,
+            None,
         )?;
         let input_layernorm = rms_norm(
             config.hidden_size,

src/models/index_tts2/config.rs

@@ -0,0 +1,228 @@
+use serde::{Deserialize, Deserializer};
+
+use crate::models::mask_gct::config::SemanticCodec;
+
+#[derive(Debug, Clone, PartialEq, Deserialize)]
+pub struct IndexTTS2Config {
+    pub dataset: Dataset,
+    pub gpt: Gpt,
+    pub semantic_codec: SemanticCodec,
+    pub s2mel: S2MelConfig,
+    pub gpt_checkpoint: String,
+    pub w2v_stat: String,
+    pub s2mel_checkpoint: String,
+    pub emo_matrix: String,
+    pub spk_matrix: String,
+    pub emo_num: Vec<usize>,
+    pub qwen_emo_path: String,
+    pub vocoder: Vocoder,
+    pub version: f32,
+}
+
+#[derive(Debug, Clone, PartialEq, Deserialize)]
+pub struct Dataset {
+    pub bpe_model: String,
+    pub sample_rate: usize,
+    pub squeeze: bool,
+    pub mel: Mel,
+}
+
+#[derive(Debug, Clone, PartialEq, Deserialize)]
+pub struct Mel {
+    pub sample_rate: usize,
+    pub n_fft: usize,
+    pub hop_length: usize,
+    pub win_length: usize,
+    pub n_mels: usize,
+    pub mel_fmin: usize,
+    pub normalize: bool,
+}
+
+#[derive(Debug, Clone, PartialEq, Deserialize)]
+pub struct Gpt {
+    pub model_dim: usize,
+    pub max_mel_tokens: usize,
+    pub max_text_tokens: usize,
+    pub heads: usize,
+    pub use_mel_codes_as_input: bool,
+    pub mel_length_compression: usize,
+    pub layers: usize,
+    pub number_text_tokens: usize,
+    pub number_mel_codes: usize,
+    pub start_mel_token: usize,
+    pub stop_mel_token: usize,
+    pub start_text_token: usize,
+    pub stop_text_token: usize,
+    pub train_solo_embeddings: bool,
+    pub condition_type: String,
+    pub condition_module: ConditionModule,
+    pub emo_condition_module: EmoConditionModule,
+}
+
+#[derive(Debug, Clone, PartialEq, Deserialize)]
+pub struct ConditionModule {
+    pub output_size: usize,
+    pub linear_units: usize,
+    pub attention_heads: usize,
+    pub num_blocks: usize,
+    pub input_layer: String,
+    pub perceiver_mult: usize,
+}
+
+#[derive(Debug, Clone, PartialEq, Deserialize)]
+pub struct EmoConditionModule {
+    pub output_size: usize,
+    pub linear_units: usize,
+    pub attention_heads: usize,
+    pub num_blocks: usize,
+    pub input_layer: String,
+    pub perceiver_mult: usize,
+}
+
+#[derive(Debug, Clone, PartialEq, Deserialize)]
+pub struct S2MelConfig {
+    pub preprocess_params: PreprocessParams,
+    pub dit_type: String,
+    pub reg_loss_type: String,
+    pub style_encoder: StyleEncoder,
+    pub length_regulator: LengthRegulator,
+    #[serde(rename = "DiT")]
+    pub di_t: DiTConfig,
+    pub wavenet: WavenetConfig,
+}
+
+#[derive(Debug, Clone, PartialEq, Deserialize)]
+pub struct PreprocessParams {
+    pub sr: usize,
+    pub spect_params: SpectParams,
+}
+
+#[derive(Debug, Clone, PartialEq, Deserialize)]
+pub struct SpectParams {
+    pub n_fft: usize,
+    pub win_length: usize,
+    pub hop_length: usize,
+    pub n_mels: usize,
+    pub fmin: usize,
+    #[serde(deserialize_with = "deserialize_optional_fmax")]
+    pub fmax: Option<usize>,
+}
+
+fn deserialize_optional_fmax<'de, D>(deserializer: D) -> Result<Option<usize>, D::Error>
+where D: Deserializer<'de> {
+    let opt: Option<serde_json::Value> = Option::deserialize(deserializer)?;
+    match opt {
+        Some(serde_json::Value::String(s)) if s == "None" || s == "null" => Ok(None),
+        Some(serde_json::Value::Number(n)) => {
+            if let Some(n) = n.as_u64() {
+                Ok(Some(n as usize))
+            } else {
+                Err(serde::de::Error::custom("Expected positive integer"))
+            }
+        }
+        Some(_) => Err(serde::de::Error::custom("Expected number or 'None' string")),
+        None => Ok(None),
+    }
+}
+
+#[derive(Debug, Clone, PartialEq, Deserialize)]
+pub struct StyleEncoder {
+    pub dim: usize,
+}
+
+#[derive(Debug, Clone, PartialEq, Deserialize)]
+pub struct LengthRegulator {
+    pub channels: usize,
+    pub is_discrete: bool,
+    pub in_channels: usize,
+    pub content_codebook_size: usize,
+    pub sampling_ratios: Vec<usize>,
+    pub vector_quantize: bool,
+    pub n_codebooks: usize,
+    pub quantizer_dropout: f32,
+    pub f0_condition: bool,
+    pub n_f0_bins: usize,
+}
+
+#[derive(Debug, Clone, PartialEq, Deserialize)]
+pub struct DiTConfig {
+    pub hidden_dim: usize,
+    pub num_heads: usize,
+    pub depth: usize,
+    pub class_dropout_prob: f32,
+    pub block_size: usize,
+    pub in_channels: usize,
+    pub style_condition: bool,
+    pub final_layer_type: String,
+    pub target: String,
+    pub content_dim: usize,
+    pub content_codebook_size: usize,
+    pub content_type: String,
+    pub f0_condition: bool,
+    pub n_f0_bins: usize,
+    pub content_codebooks: usize,
+    pub is_causal: bool,
+    pub long_skip_connection: bool,
+    pub zero_prompt_speech_token: bool,
+    pub time_as_token: bool,
+    pub style_as_token: bool,
+    pub uvit_skip_connection: bool,
+    pub add_resblock_in_transformer: bool,
+}
+
+pub struct DiTModelArgs {
+    pub block_size: usize,
+    pub vocab_size: usize,
+    pub n_layer: usize,
+    pub n_head: usize,
+    pub dim: usize,
+    pub intermediate_size: usize,
+    pub n_local_heads: usize,
+    pub head_dim: usize,
+    pub rope_base: f32,
+    pub norm_eps: f64,
+    pub has_cross_attention: bool,
+    pub context_dim: usize,
+    pub uvit_skip_connection: bool,
+    pub time_as_token: bool,
+}
+
+impl DiTModelArgs {
+    pub fn new_from_dit_config(config: &DiTConfig) -> Self {
+        let hidden_dim = 4 * config.hidden_dim;
+        let n_hidden = 2 * hidden_dim / 3;
+        let intermediate_size = n_hidden + (256 - n_hidden % 256) % 256;
+        Self {
+            block_size: config.block_size,
+            vocab_size: 1024,
+            n_layer: config.depth,
+            n_head: config.num_heads,
+            dim: config.hidden_dim,
+            intermediate_size,
+            n_local_heads: config.num_heads,
+            head_dim: config.hidden_dim / config.num_heads,
+            rope_base: 10000.0,
+            norm_eps: 1e-5,
+            has_cross_attention: false,
+            context_dim: 0,
+            uvit_skip_connection: config.uvit_skip_connection,
+            time_as_token: config.uvit_skip_connection,
+        }
+    }
+}
+
+#[derive(Debug, Clone, PartialEq, Deserialize)]
+pub struct WavenetConfig {
+    pub hidden_dim: usize,
+    pub num_layers: usize,
+    pub kernel_size: usize,
+    pub dilation_rate: usize,
+    pub p_dropout: f32,
+    pub style_condition: bool,
+}
+
+#[derive(Debug, Clone, PartialEq, Deserialize)]
+pub struct Vocoder {
+    pub r#type: String,
+    pub name: String,
+}

src/models/index_tts2/generate.rs

@@ -0,0 +1,29 @@
+use aha_openai_dive::v1::resources::chat::ChatCompletionParameters;
+use anyhow::Result;
+use candle_core::{DType, Device};
+
+use crate::{
+    models::index_tts2::{config::IndexTTS2Config, processor::IndexTTS2Processor},
+    utils::{get_default_save_dir, get_device, get_dtype},
+};
+pub struct IndexTTS2Generate {
+    processor: IndexTTS2Processor,
+    config: IndexTTS2Config,
+}
+
+impl IndexTTS2Generate {
+    pub fn init(path: &str, device: Option<&Device>, dtype: Option<DType>) -> Result<Self> {
+        let config_path = path.to_string() + "/config.yaml";
+        let save_dir = get_default_save_dir().ok_or(anyhow::anyhow!("Failed to get save dir"))?;
+        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)?;
+        
+        Ok(Self { config, processor })
+    }
+    pub fn generate(&mut self, mes: ChatCompletionParameters) -> Result<()> {
+        let _ = self.processor.process_info(&mes)?;
+        Ok(())
+    }
+}

src/models/index_tts2/mod.rs

@@ -0,0 +1,5 @@
+pub mod config;
+pub mod generate;
+pub mod model;
+pub mod processor;
+pub mod utils;

src/models/index_tts2/model.rs

@@ -0,0 +1,617 @@
+use anyhow::Result;
+use candle_core::{D, Tensor};
+use candle_nn::{
+    Conv1d, Embedding, LayerNorm, Linear, Module, RmsNorm, VarBuilder, embedding, linear, linear_b,
+    ops::sigmoid, rms_norm,
+};
+
+use crate::{
+    models::{
+        common::{
+            GateUpDownMLP, QKVCatAttention, TwoLinearMLP, WNConv1d, WNLinear, get_conv1d,
+            get_layer_norm,
+        },
+        index_tts2::config::{DiTModelArgs, S2MelConfig},
+    },
+    position_embed::rope::RoPE,
+    utils::tensor_utils::{pad_reflect_last_dim, split_tensor_with_size},
+};
+pub struct AdaptiveLayerNorm {
+    project_layer: Linear,
+    norm: RmsNorm,
+    d_model: usize,
+}
+
+impl AdaptiveLayerNorm {
+    pub fn new(vb: VarBuilder, d_model: usize, eps: f64) -> Result<Self> {
+        let project_layer = linear(d_model, d_model * 2, vb.pp("project_layer"))?;
+        let norm = rms_norm(d_model, eps, vb.pp("norm"))?;
+        Ok(Self {
+            project_layer,
+            norm,
+            d_model,
+        })
+    }
+
+    pub fn forward(&self, xs: &Tensor, embedding: Option<&Tensor>) -> Result<Tensor> {
+        if let Some(embedding) = embedding {
+            let emb = self.project_layer.forward(embedding)?;
+            let emb_split = split_tensor_with_size(&emb, 2, D::Minus1)?;
+            let weight = &emb_split[0];
+            let bias = &emb_split[1];
+            Ok(self
+                .norm
+                .forward(xs)?
+                .broadcast_mul(weight)?
+                .broadcast_add(bias)?)
+        } else {
+            Ok(self.norm.forward(xs)?)
+        }
+    }
+}
+
+pub struct DiTTransformerBlock {
+    attention: QKVCatAttention,
+    feed_forward: GateUpDownMLP,
+    ffn_norm: AdaptiveLayerNorm,
+    attention_norm: AdaptiveLayerNorm,
+    skip_in_linear: Option<Linear>,
+    uvit_skip_connection: bool,
+    time_as_token: bool,
+}
+
+impl DiTTransformerBlock {
+    pub fn new(vb: VarBuilder, config: &DiTModelArgs) -> Result<Self> {
+        let attention = QKVCatAttention::new(
+            vb.pp("attention"),
+            config.dim,
+            config.n_head,
+            Some(config.head_dim),
+            false,
+            Some("wqkv"),
+            Some("wo"),
+        )?;
+
+        let feed_forward = GateUpDownMLP::new(
+            vb.pp("feed_forward"),
+            config.dim,
+            config.intermediate_size,
+            candle_nn::Activation::Silu,
+            false,
+            Some("w1"),
+            Some("w3"),
+            Some("w2"),
+        )?;
+        let ffn_norm = AdaptiveLayerNorm::new(vb.pp("ffn_norm"), config.dim, config.norm_eps)?;
+        let attention_norm =
+            AdaptiveLayerNorm::new(vb.pp("attention_norm"), config.dim, config.norm_eps)?;
+        let (skip_in_linear, uvit_skip_connection) = if config.uvit_skip_connection {
+            let skip_in_linear = linear(config.dim * 2, config.dim, vb.pp("skip_in_linear"))?;
+            (Some(skip_in_linear), config.uvit_skip_connection)
+        } else {
+            (None, false)
+        };
+        Ok(Self {
+            attention,
+            feed_forward,
+            ffn_norm,
+            attention_norm,
+            skip_in_linear,
+            uvit_skip_connection,
+            time_as_token: config.time_as_token,
+        })
+    }
+
+    pub fn forward(
+        &self,
+        xs: &Tensor,
+        c: &Tensor,
+        cos: Option<&Tensor>,
+        sin: Option<&Tensor>,
+        mask: Option<&Tensor>,
+        skip_in_x: Option<&Tensor>,
+    ) -> Result<Tensor> {
+        let c = if self.time_as_token { None } else { Some(c) };
+        let mut xs = xs.clone();
+        if self.uvit_skip_connection
+            && let Some(skip_in_x) = skip_in_x
+            && let Some(skip_in_linear) = &self.skip_in_linear
+        {
+            let cat = Tensor::cat(&[&xs, skip_in_x], D::Minus1)?;
+            xs = skip_in_linear.forward(&cat)?;
+        }
+        let xs = self
+            .attention
+            .forward(
+                &self.attention_norm.forward(&xs, c)?,
+                cos,
+                sin,
+                mask,
+                false,
+                true,
+            )?
+            .add(&xs)?;
+        let out = self
+            .feed_forward
+            .forward(&self.ffn_norm.forward(&xs, c)?)?
+            .add(&xs)?;
+        Ok(out)
+    }
+}
+
+pub struct DiTTransformer {
+    layers: Vec<DiTTransformerBlock>,
+    norm: AdaptiveLayerNorm,
+    rope: RoPE,
+    uvit_skip_connection: bool,
+    layers_emit_skip: Vec<usize>,
+    layers_receive_skip: Vec<usize>,
+}
+
+impl DiTTransformer {
+    pub fn new(vb: VarBuilder, config: &DiTModelArgs) -> Result<Self> {
+        let vb_layers = vb.pp("layers");
+        let mut layers = vec![];
+        for i in 0..config.n_layer {
+            let layer = DiTTransformerBlock::new(vb_layers.pp(i), config)?;
+            layers.push(layer);
+        }
+        let norm = AdaptiveLayerNorm::new(vb.pp("norm"), config.dim, config.norm_eps)?;
+        let rope = RoPE::new(config.dim, 10000.0, vb.device())?;
+        let mut layers_emit_skip: Vec<usize> = vec![];
+        let mut layers_receive_skip: Vec<usize> = vec![];
+        if config.uvit_skip_connection {
+            layers_emit_skip = (0..config.n_layer)
+                .filter(|&x| x < config.n_layer / 2)
+                .collect();
+            layers_receive_skip = (0..config.n_layer)
+                .filter(|&x| x > config.n_layer / 2)
+                .collect();
+        }
+        Ok(Self {
+            layers,
+            norm,
+            rope,
+            uvit_skip_connection: config.uvit_skip_connection,
+            layers_emit_skip,
+            layers_receive_skip,
+        })
+    }
+    pub fn forward(&self, xs: &Tensor, c: &Tensor, mask: Option<&Tensor>) -> Result<Tensor> {
+        let (_, seq_len, _) = xs.dims3()?;
+        let (cos, sin) = self.rope.forward(0, seq_len, xs.device())?;
+        let mut skip_in_x_list = vec![];
+        let mut xs = xs.clone();
+        for (i, layer) in (&self.layers).iter().enumerate() {
+            let skip_in_x = if self.uvit_skip_connection && self.layers_receive_skip.contains(&i) {
+                skip_in_x_list.pop()
+            } else {
+                None
+            };
+            xs = layer.forward(&xs, c, Some(&cos), Some(&sin), mask, skip_in_x.as_ref())?;
+            if self.uvit_skip_connection && self.layers_emit_skip.contains(&i) {
+                skip_in_x_list.push(xs.clone());
+            }
+        }
+        xs = self.norm.forward(&xs, Some(c))?;
+        Ok(xs)
+    }
+}
+
+pub struct TimestepEmbedder {
+    mlp: TwoLinearMLP,
+    freqs: Tensor,
+    scale: f64,
+    frequency_embedding_size: usize,
+}
+
+impl TimestepEmbedder {
+    pub fn new(
+        vb: VarBuilder,
+        hidden_size: usize,
+        frequency_embedding_size: usize,
+    ) -> Result<Self> {
+        let mlp = TwoLinearMLP::new(
+            vb.pp("mlp"),
+            frequency_embedding_size,
+            hidden_size,
+            hidden_size,
+            candle_nn::Activation::Silu,
+            true,
+            "0",
+            "1",
+        )?;
+        let scale = 1000.0;
+        let half = frequency_embedding_size / 2;
+        let freqs = Tensor::arange(0f32, half as f32, vb.device())?
+            .affine(-(10000.0f64.ln()), 0.0)?
+            .exp()?;
+        Ok(Self {
+            mlp,
+            freqs,
+            scale,
+            frequency_embedding_size,
+        })
+    }
+    pub fn forward(&self, t: &Tensor) -> Result<Tensor> {
+        let args = t
+            .affine(self.scale, 0.0)?
+            .unsqueeze(D::Minus1)?
+            .broadcast_matmul(&self.freqs.unsqueeze(0)?)?;
+        let mut embedding = Tensor::cat(&[args.cos()?, args.sin()?], D::Minus1)?;
+        if self.frequency_embedding_size % 2 > 0 {
+            embedding = embedding.pad_with_zeros(D::Minus1, 0, 1)?;
+        }
+        embedding = self.mlp.forward(&embedding)?;
+        Ok(embedding)
+    }
+}
+
+pub struct SConv1d {
+    conv: WNConv1d,
+    ks: usize,
+    stride: usize,
+    dilation: usize,
+}
+
+impl SConv1d {
+    pub fn new(
+        vb: VarBuilder,
+        in_c: usize,
+        out_c: usize,
+        ks: usize,
+        stride: usize,
+        dilation: usize,
+        groups: usize,
+        bias: bool,
+    ) -> Result<Self> {
+        let conv = WNConv1d::new(
+            vb.pp("conv.conv"),
+            in_c,
+            out_c,
+            ks,
+            dilation,
+            0,
+            groups,
+            stride,
+            bias,
+        )?;
+        Ok(Self {
+            conv,
+            ks,
+            stride,
+            dilation,
+        })
+    }
+    pub fn forward(&self, xs: &Tensor) -> Result<Tensor> {
+        let length = xs.dim(D::Minus1)?;
+        let ks = (self.ks - 1) * self.dilation + 1;
+        let padding_total = ks - self.stride;
+        let n_frames = (length - ks + padding_total) as f32 / self.stride as f32 + 1.0;
+        let idea_length = (n_frames.ceil() as usize - 1) * self.stride + (ks - padding_total);
+        let extra_padding = idea_length - length;
+        let padding_right = padding_total / 2;
+        let padding_left = padding_total - padding_right;
+        let xs = pad_reflect_last_dim(xs, (padding_left, padding_right + extra_padding))?;
+        let xs = self.conv.forward(&xs)?;
+        Ok(xs)
+    }
+}
+
+pub struct Wavenet {
+    cond_layer: Option<SConv1d>,
+    in_layers: Vec<SConv1d>,
+    res_skip_layers: Vec<SConv1d>,
+    hidden_c: usize,
+    n_layers: usize,
+}
+
+impl Wavenet {
+    pub fn new(
+        vb: VarBuilder,
+        hidden_c: usize,
+        ks: usize,
+        dilation_rate: usize,
+        n_layers: usize,
+        gin_channels: usize,
+    ) -> Result<Self> {
+        let cond_layer = if gin_channels != 0 {
+            Some(SConv1d::new(
+                vb.pp("cond_layer"),
+                gin_channels,
+                2 * hidden_c * n_layers,
+                1,
+                1,
+                1,
+                1,
+                true,
+            )?)
+        } else {
+            None
+        };
+        let mut in_layers = vec![];
+        let vb_layers = vb.pp("in_layers");
+        let mut res_skip_layers = vec![];
+        let vb_res_skip_layers = vb.pp("res_skip_layers");
+        for i in 0..n_layers {
+            let dilation = dilation_rate.pow(i as u32);
+            let in_layer = SConv1d::new(
+                vb_layers.pp(i),
+                hidden_c,
+                1 * hidden_c,
+                ks,
+                1,
+                dilation,
+                1,
+                true,
+            )?;
+            in_layers.push(in_layer);
+            let res_skip_c = if i < n_layers - 1 {
+                2 * hidden_c
+            } else {
+                hidden_c
+            };
+            let res_skip_layer = SConv1d::new(
+                vb_res_skip_layers.pp(i),
+                hidden_c,
+                res_skip_c,
+                1,
+                1,
+                1,
+                1,
+                true,
+            )?;
+            res_skip_layers.push(res_skip_layer);
+        }
+        Ok(Self {
+            cond_layer,
+            in_layers,
+            res_skip_layers,
+            hidden_c,
+            n_layers,
+        })
+    }
+
+    pub fn fused_add_tanh_sigmoid_multiply(
+        &self,
+        input_a: &Tensor,
+        input_b: &Tensor,
+    ) -> Result<Tensor> {
+        let in_act = input_a.add(&input_b)?;
+        let parts = split_tensor_with_size(&in_act, 2, 1)?;
+        let t_act = (&parts[0]).tanh()?;
+        let s_act = sigmoid(&parts[1])?;
+        let acts = t_act.mul(&s_act)?;
+        Ok(acts)
+    }
+
+    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
+        {
+            Some(cond_layer.forward(g)?)
+        } else {
+            None
+        };
+        let mut xs = xs.clone();
+        for i in 0..self.n_layers {
+            let xs_in = &self.in_layers[i].forward(&xs)?;
+            let g_l = if let Some(g) = &g {
+                let cond_offset = i * 2 * self.hidden_c;
+                g.narrow(1, cond_offset, 2 * self.hidden_c)?
+            } else {
+                xs_in.zeros_like()?
+            };
+            let acts = self.fused_add_tanh_sigmoid_multiply(&xs_in, &g_l)?;
+            let res_skip_act = &self.res_skip_layers[i].forward(&acts)?;
+            if i < self.n_layers - 1 {
+                let res_acts = res_skip_act.narrow(1, 0, self.hidden_c)?;
+                let out_acts = res_skip_act.narrow(1, self.hidden_c, self.hidden_c)?;
+                xs = xs.add(&res_acts)?.mul(x_mask)?;
+                output = output.add(&out_acts)?;
+            } else {
+                output = output.add(&res_skip_act)?;
+            }
+        }
+        output = output.mul(x_mask)?;
+        Ok(output)
+    }
+}
+
+pub struct FinalLayer {
+    norm_final: LayerNorm,
+    linear: WNLinear,
+    ada_ln_modulation: Linear, // (silu+linear)
+}
+
+impl FinalLayer {
+    pub fn new(
+        vb: VarBuilder,
+        hidden_size: usize,
+        patch_size: usize,
+        out_c: usize,
+    ) -> Result<Self> {
+        let norm_final = get_layer_norm(vb.pp("norm_final"), 1e-6, hidden_size)?;
+        let linear = WNLinear::new(
+            vb.pp("linear"),
+            hidden_size,
+            patch_size * patch_size * out_c,
+            true,
+        )?;
+        let ada_ln_modulation = linear_b(
+            hidden_size,
+            2 * hidden_size,
+            true,
+            vb.pp("adaLN_modulation.1"),
+        )?;
+        Ok(Self {
+            norm_final,
+            linear,
+            ada_ln_modulation,
+        })
+    }
+
+    pub fn forward(&self, xs: &Tensor, c: &Tensor) -> Result<Tensor> {
+        let linear_c = self.ada_ln_modulation.forward(c)?.chunk(2, 1)?;
+        let xs = self.norm_final.forward(xs)?;
+        let xs = linear_c[1]
+            .unsqueeze(1)?
+            .affine(1.0, 1.0)?
+            .broadcast_mul(&xs)?
+            .add(&linear_c[0].unsqueeze(1)?)?;
+        let xs = self.linear.forward(&xs)?;
+        Ok(xs)
+    }
+}
+
+pub struct DiT {
+    transformer: DiTTransformer,
+    x_embedder: WNLinear,
+    cond_embedder: Embedding,
+    cond_projection: Linear,
+    t_embedder: TimestepEmbedder,
+    input_pos: Tensor,
+    t_embedder2: TimestepEmbedder,
+    conv1: Linear,
+    conv2: Conv1d,
+    wavenet: Wavenet,
+    final_layer: FinalLayer,
+    res_projection: Linear,
+    content_mask_embedder: Embedding,
+    skip_linear: Linear,
+    cond_x_merge_linear: Linear,
+    style_in: Option<Linear>,
+    time_as_token: bool,
+    style_as_token: bool,
+    uvit_skip_connection: bool,
+}
+
+impl DiT {
+    pub fn new(vb: VarBuilder, config: &S2MelConfig) -> Result<Self> {
+        let time_as_token = config.di_t.time_as_token;
+        let style_as_token = config.di_t.style_as_token;
+        let uvit_skip_connection = config.di_t.uvit_skip_connection;
+        let transformer_config = DiTModelArgs::new_from_dit_config(&config.di_t);
+        let transformer = DiTTransformer::new(vb.pp("transformer"), &transformer_config)?;
+        let x_embedder = WNLinear::new(
+            vb.pp("x_embedder"),
+            config.di_t.in_channels,
+            config.di_t.hidden_dim,
+            true,
+        )?;
+        let cond_embedder = embedding(
+            config.di_t.content_codebook_size,
+            config.di_t.hidden_dim,
+            vb.pp("cond_embedder"),
+        )?;
+        let cond_projection = linear_b(
+            config.di_t.content_dim,
+            config.di_t.hidden_dim,
+            true,
+            vb.pp("cond_projection"),
+        )?;
+        let t_embedder = TimestepEmbedder::new(vb.pp("t_embedder"), config.di_t.hidden_dim, 256)?;
+        let input_pos = Tensor::arange(0u32, 16384, vb.device())?;
+        let t_embedder2 =
+            TimestepEmbedder::new(vb.pp("t_embedder2"), config.wavenet.hidden_dim, 256)?;
+        let conv1 = linear_b(
+            config.di_t.hidden_dim,
+            config.wavenet.hidden_dim,
+            true,
+            vb.pp("conv1"),
+        )?;
+        let conv2 = get_conv1d(
+            vb.pp("conv2"),
+            config.wavenet.hidden_dim,
+            config.di_t.in_channels,
+            1,
+            0,
+            1,
+            1,
+            1,
+            true,
+        )?;
+        let wavenet = Wavenet::new(
+            vb.pp("wavenet"),
+            config.wavenet.hidden_dim,
+            config.wavenet.kernel_size,
+            config.wavenet.dilation_rate,
+            config.wavenet.num_layers,
+            config.wavenet.hidden_dim,
+        )?;
+        let final_layer = FinalLayer::new(
+            vb.pp("final_layer"),
+            config.wavenet.hidden_dim,
+            1,
+            config.wavenet.hidden_dim,
+        )?;
+        let res_projection = linear(
+            config.di_t.hidden_dim,
+            config.wavenet.hidden_dim,
+            vb.pp("res_projection"),
+        )?;
+        let content_mask_embedder =
+            embedding(1, config.di_t.hidden_dim, vb.pp("content_mask_embedder"))?;
+        let skip_linear = linear(
+            config.di_t.hidden_dim + config.di_t.in_channels,
+            config.di_t.hidden_dim,
+            vb.pp("skip_linear"),
+        )?;
+        let in_dim = if config.di_t.style_condition && !config.di_t.style_as_token {
+            config.di_t.hidden_dim + config.di_t.in_channels * 2 + config.style_encoder.dim
+        } else {
+            config.di_t.hidden_dim + config.di_t.in_channels * 2
+        };
+        let cond_x_merge_linear =
+            linear(in_dim, config.di_t.hidden_dim, vb.pp("cond_x_merge_linear"))?;
+        let style_in = if config.di_t.style_as_token {
+            Some(linear(
+                config.style_encoder.dim,
+                config.di_t.hidden_dim,
+                vb.pp("style_in"),
+            )?)
+        } else {
+            None
+        };
+        Ok(Self {
+            transformer,
+            x_embedder,
+            cond_embedder,
+            cond_projection,
+            t_embedder,
+            input_pos,
+            t_embedder2,
+            conv1,
+            conv2,
+            wavenet,
+            final_layer,
+            res_projection,
+            content_mask_embedder,
+            skip_linear,
+            cond_x_merge_linear,
+            style_in,
+            time_as_token,
+            style_as_token,
+            uvit_skip_connection,
+        })
+    }
+}
+
+pub struct CFM {
+    estimator: DiT,
+}
+
+pub struct MyModel {
+    cfm: CFM,
+}
+
+pub struct IndexTTS2 {
+    cache_spk_cond: Option<Tensor>,
+    cache_s2mel_style: Option<Tensor>,
+    cache_s2mel_prompt: Option<Tensor>,
+    cache_spk_audio_prompt: Option<String>,
+    cache_emo_cond: Option<Tensor>,
+    cache_emo_audio_prompt: Option<Tensor>,
+    cache_mel: Option<Tensor>,
+}

src/models/index_tts2/processor.rs

@@ -0,0 +1,199 @@
+use aha_openai_dive::v1::resources::chat::ChatCompletionParameters;
+use anyhow::Result;
+use candle_core::{D, DType, Device, IndexOp, Tensor, pickle::read_all_with_key};
+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
+    },
+    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, 
+            torch_stft,
+        },
+        get_vb_model_path,
+        tensor_utils::pad_reflect_last_dim,
+    },
+};
+
+pub struct IndexTTS2Processor {
+    device: Device,
+    max_audio_length_seconds: usize,
+    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,
+}
+
+impl IndexTTS2Processor {
+    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)?;
+        Ok(Self {
+            device: device.clone(),
+            max_audio_length_seconds: 15,
+            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,
+        })
+    }
+
+    pub fn cut_audio(&self, audio: &Tensor, sr: usize) -> Result<(Tensor, usize)> {
+        let max_audio_samples = self.max_audio_length_seconds * sr;
+        let audio_lens = audio.dim(1)?;
+        let audio = if audio_lens > max_audio_samples {
+            audio.i((.., 0..max_audio_samples))?
+        } else {
+            audio.clone()
+        };
+        Ok((audio, sr))
+    }
+
+    pub fn extract_audio_and_cut(
+        &self,
+        mes: &ChatCompletionParameters,
+        device: &Device,
+    ) -> Result<(Tensor, usize)> {
+        let audio_url_vec = extract_audio_url(mes);
+        let (audio, sr) = load_audio(&audio_url_vec[0], device)?;
+        self.cut_audio(&audio, sr)
+    }
+
+    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 process_info(&self, mes: &ChatCompletionParameters) -> Result<()> {
+        let (audio, sr) = 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) =
+            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 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)?;
+        println!("style: {}", style);
+        Ok(())
+    }
+}

src/models/index_tts2/utils.rs

@@ -0,0 +1,18 @@
+use crate::utils::{download_model, get_default_save_dir};
+
+pub async fn download_index_tts2_need_model(save_dir: Option<&str>) -> anyhow::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 campplus= "funasr/campplus"; // huggingface
+    let campplus = "iic/speech_campplus_sv_zh-cn_16k-common";  // modelscope
+    download_model(w2v_bert2_0, &save_dir, 3).await?;
+    download_model(mask_gct, &save_dir, 3).await?;
+    download_model(campplus, &save_dir, 3).await?;
+    
+    Ok(())
+}

src/models/mask_gct/config.rs

@@ -0,0 +1,21 @@
+#[derive(Debug, Clone, PartialEq, serde::Deserialize)]
+pub struct SemanticCodec {
+    pub codebook_size: usize,
+    pub hidden_size: usize,
+    pub codebook_dim: usize,
+    pub vocos_dim: usize,
+    pub vocos_intermediate_dim: usize,
+    pub vocos_num_layers: usize,
+    #[serde(default = "default_num_quantizers")]
+    pub num_quantizers: usize,
+    #[serde(default = "default_downsample_scale")]
+    pub downsample_scale: usize,
+}
+
+fn default_num_quantizers() -> usize {
+    1
+}
+
+fn default_downsample_scale() -> usize {
+    1
+}

src/models/mask_gct/mod.rs

@@ -0,0 +1,2 @@
+pub mod model;
+pub mod config;

src/models/mask_gct/model.rs

@@ -0,0 +1,344 @@
+use anyhow::Result;
+use candle_core::{D, Tensor};
+use candle_nn::{
+    Conv1d, Embedding, Init, LayerNorm, Linear, Module, VarBuilder, embedding, linear,
+};
+
+use crate::{
+    models::{
+        common::{WNConv1d, get_conv1d, get_layer_norm},
+        mask_gct::config::SemanticCodec,
+    },
+    utils::tensor_utils::{interpolate_nearest_1d, l2_normalize},
+};
+
+pub struct ConvNeXtBlock {
+    dwconv: Conv1d,
+    norm: LayerNorm,
+    pwconv1: Linear,
+    pwconv2: Linear,
+    gamma: Option<Tensor>,
+}
+
+impl ConvNeXtBlock {
+    pub fn new(
+        vb: VarBuilder,
+        dim: usize,
+        intermediate_dim: usize,
+        // 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 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))?;
+        Ok(Self {
+            dwconv,
+            norm,
+            pwconv1,
+            pwconv2,
+            gamma: Some(gamma),
+        })
+    }
+    pub fn forward(&self, xs: &Tensor) -> Result<Tensor> {
+        let residual = xs.clone();
+        let xs = self.dwconv.forward(xs)?;
+        let xs = xs.transpose(1, 2)?;
+        let xs = self.norm.forward(&xs)?;
+        let xs = self.pwconv1.forward(&xs)?.gelu()?;
+        let mut xs = self.pwconv2.forward(&xs)?;
+        if let Some(gamma) = &self.gamma {
+            xs = xs.broadcast_mul(&gamma)?;
+        }
+        let xs = xs.transpose(1, 2)?;
+        let xs = residual.add(&xs)?;
+        Ok(xs)
+    }
+}
+
+pub struct VocosBackbone {
+    embed: Conv1d,
+    norm: LayerNorm,
+    convnext: Vec<ConvNeXtBlock>,
+    final_layer_norm: LayerNorm,
+}
+
+impl VocosBackbone {
+    pub fn new(
+        vb: VarBuilder,
+        input_channels: usize,
+        dim: usize,
+        intermediate_dim: usize,
+        num_layers: usize,
+        // 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 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)?;
+        Ok(Self {
+            embed,
+            norm,
+            convnext,
+            final_layer_norm,
+        })
+    }
+
+    pub fn forward(&self, xs: &Tensor) -> Result<Tensor> {
+        let xs = self.embed.forward(xs)?;
+        let mut xs = self.norm.forward(&xs.transpose(1, 2)?)?.transpose(1, 2)?;
+        for layer in &self.convnext {
+            xs = layer.forward(&xs)?;
+        }
+        xs = self.final_layer_norm.forward(&xs.transpose(1, 2)?)?;
+        Ok(xs)
+    }
+}
+
+pub struct FactorizedVectorQuantize {
+    use_l2_normlize: bool,
+    in_project: Option<WNConv1d>,
+    out_project: Option<WNConv1d>,
+    codebook: Embedding,
+}
+
+impl FactorizedVectorQuantize {
+    pub fn new(
+        vb: VarBuilder,
+        input_dim: usize,
+        codebook_size: usize,
+        codebook_dim: usize,
+        use_l2_normlize: bool,
+    ) -> Result<Self> {
+        let (in_project, out_project) = if input_dim != codebook_dim {
+            let in_project =
+                WNConv1d::new(vb.pp("in_project"), input_dim, codebook_dim, 1, 1, 0, 1, 1, true)?;
+            let out_project =
+                WNConv1d::new(vb.pp("out_project"), codebook_dim, input_dim, 1, 1, 0, 1, 1, true)?;
+            (Some(in_project), Some(out_project))
+        } else {
+            (None, None)
+        };
+        let codebook = embedding(codebook_size, codebook_dim, vb.pp("codebook"))?;
+        Ok(Self {
+            use_l2_normlize,
+            in_project,
+            out_project,
+            codebook,
+        })
+    }
+
+    pub fn decode_latents(&self, xs: &Tensor) -> Result<(Tensor, Tensor)> {
+        let (bs, len, dim) = xs.dims3()?;
+        let mut encodings = xs.transpose(1, 2)?.reshape((bs * dim, len))?;
+        let mut codebook = self.codebook.embeddings().clone();
+        if self.use_l2_normlize {
+            encodings = l2_normalize(&encodings, 1)?;
+            codebook = l2_normalize(&codebook, 1)?;
+        }
+        let dist1 = encodings.powf(2.0)?.sum_keepdim(1)?;
+        let dist2 = encodings.affine(2.0, 0.0)?.matmul(&codebook.t()?)?;
+        let dist3 = codebook.powf(2.0)?.sum_keepdim(1)?.t()?;
+        let dist = dist1.broadcast_sub(&dist2)?.broadcast_add(&dist3)?;
+        let indices = dist
+            .affine(-1.0, 0.0)?
+            .argmax(1)?
+            .reshape((bs, ()))?
+            .to_dtype(candle_core::DType::U32)?;
+        let z_q = self.codebook.forward(&indices)?.transpose(1, 2)?;
+        Ok((z_q, indices))
+    }
+
+    pub fn forward(&self, xs: &Tensor) -> Result<(Tensor, Tensor)> {
+        let mut xs = xs.clone();
+        if let Some(in_proj) = &self.in_project {
+            xs = in_proj.forward(&xs)?;
+        }
+        let (z_q, indices) = self.decode_latents(&xs)?;
+        let mut z_q = xs.add(&z_q.sub(&xs)?)?;
+        if let Some(out_proj) = &self.out_project {
+            z_q = out_proj.forward(&z_q)?;
+        }
+        Ok((z_q, indices))
+    }
+}
+
+pub struct ResidualVQ {
+    num_quantizers: usize,
+    quantizers: Vec<FactorizedVectorQuantize>,
+}
+
+impl ResidualVQ {
+    pub fn new(
+        vb: VarBuilder,
+        input_dim: usize,
+        num_quantizers: usize,
+        codebook_size: usize,
+        codebook_dim: usize,
+        // quantizer_type: &str,  // now only surpport "fvq"
+    ) -> Result<Self> {
+        let vb_quantizers = vb.pp("quantizers");
+        let mut quantizers = vec![];
+        for i in 0..num_quantizers {
+            let quantizer = FactorizedVectorQuantize::new(
+                vb_quantizers.pp(i),
+                input_dim,
+                codebook_size,
+                codebook_dim,
+                true,
+            )?;
+            quantizers.push(quantizer);
+        }
+        Ok(Self {
+            num_quantizers,
+            quantizers,
+        })
+    }
+
+    pub fn forward(
+        &self,
+        xs: &Tensor,
+        n_quantizers: Option<usize>,
+    ) -> Result<(Tensor, Tensor, Tensor)> {
+        let mut all_indices = vec![];
+        let mut all_quantized = vec![];
+        let n_quantizers = n_quantizers.unwrap_or(self.num_quantizers);
+        let mut residual = xs.clone();
+        let mut quantized_out = Tensor::new(0.0f32, xs.device())?.to_dtype(xs.dtype())?;
+        for (i, quantizer) in (&self.quantizers).iter().enumerate() {
+            if i >= n_quantizers {
+                break;
+            }
+            let (z_q_i, indices_i) = quantizer.forward(&residual)?;
+            quantized_out = quantized_out.broadcast_add(&z_q_i)?;
+            residual = residual.sub(&z_q_i)?;
+            all_indices.push(indices_i);
+            all_quantized.push(z_q_i);
+        }
+        let all_indices = Tensor::stack(&all_indices, 0)?;
+        let all_quantized = Tensor::stack(&all_quantized, 0)?;
+        Ok((quantized_out, all_indices, all_quantized))
+    }
+}
+
+pub struct RepCodec {
+    downsample_scale: usize,
+    down: Option<Conv1d>,
+    up: Option<Conv1d>,
+    encoder_0: VocosBackbone,
+    encoder_1: Linear,
+    decoder_0: VocosBackbone,
+    decoder_1: Linear,
+    quantizer: ResidualVQ,
+}
+
+impl RepCodec {
+    pub fn new(vb: VarBuilder, config: &SemanticCodec) -> Result<Self> {
+        let (down, up) = if config.downsample_scale > 1 {
+            let down = get_conv1d(
+                vb.pp("down"),
+                config.hidden_size,
+                config.hidden_size,
+                3,
+                1,
+                2,
+                1,
+                1,
+                true,
+            )?;
+            let up = get_conv1d(
+                vb.pp("up"),
+                config.hidden_size,
+                config.hidden_size,
+                3,
+                1,
+                1,
+                1,
+                1,
+                true,
+            )?;
+            (Some(down), Some(up))
+        } else {
+            (None, None)
+        };
+        let encoder_0 = VocosBackbone::new(
+            vb.pp("encoder.0"),
+            config.hidden_size,
+            config.vocos_dim,
+            config.vocos_intermediate_dim,
+            config.vocos_num_layers,
+        )?;
+        let encoder_1 = linear(config.vocos_dim, config.hidden_size, vb.pp("encoder.1"))?;
+        let decoder_0 = VocosBackbone::new(
+            vb.pp("decoder.0"),
+            config.hidden_size,
+            config.vocos_dim,
+            config.vocos_intermediate_dim,
+            config.vocos_num_layers,
+        )?;
+        let decoder_1 = linear(config.vocos_dim, config.hidden_size, vb.pp("decoder.1"))?;
+        let quantizer = ResidualVQ::new(
+            vb.pp("quantizer"),
+            config.hidden_size,
+            config.num_quantizers,
+            config.codebook_size,
+            config.codebook_dim,
+        )?;
+        Ok(Self {
+            downsample_scale: config.downsample_scale,
+            down,
+            up,
+            encoder_0,
+            encoder_1,
+            decoder_0,
+            decoder_1,
+            quantizer,
+        })
+    }
+
+    pub fn forward(&self, xs: &Tensor) -> Result<(Tensor, Tensor)> {
+        let mut xs = xs.clone();
+        if let Some(down) = &self.down {
+            xs = xs.transpose(1, 2)?;
+            xs = down.forward(&xs)?.gelu()?;
+            xs = xs.transpose(1, 2)?;
+        }
+        xs = self.encoder_0.forward(&xs.transpose(1, 2)?)?;
+        xs = self.encoder_1.forward(&xs)?;
+        xs = xs.transpose(1, 2)?;
+        let (quantized_out, all_indices, _) = self.quantizer.forward(&xs, None)?;
+        xs = self.decoder_0.forward(&quantized_out)?;
+        if let Some(up) = &self.up {
+            xs = xs.transpose(1, 2)?;
+            let last_dim = xs.dim(D::Minus1)?;
+            let target_size = last_dim * 2;
+            xs = interpolate_nearest_1d(&xs, target_size)?;
+            xs = up.forward(&xs)?.transpose(1, 2)?;
+        }
+
+        Ok((xs, all_indices))
+    }
+
+    pub fn quantize(&self, xs: &Tensor) -> Result<(Tensor, Tensor)> {
+        let mut xs = xs.clone();
+        if let Some(down) = &self.down {
+            xs = xs.transpose(1, 2)?;
+            xs = down.forward(&xs)?.gelu()?;
+            xs = xs.transpose(1, 2)?;
+        }
+        xs = self.encoder_0.forward(&xs.transpose(1, 2)?)?;
+        xs = self.encoder_1.forward(&xs)?;
+        xs = xs.transpose(1, 2)?;
+        let (quantized_out, mut all_indices, _) = self.quantizer.forward(&xs, None)?;
+        if all_indices.dim(0)? == 1 {
+            all_indices = all_indices.squeeze(0)?;
+        }
+        let quantized_out = quantized_out.transpose(1, 2)?;
+        Ok((all_indices, quantized_out))
+    }
+}

src/models/minicpm4/model.rs

@@ -121,6 +121,9 @@ impl MiniCPMDecoderLayer {
             cfg.intermediate_size,
             cfg.hidden_act,
             false,
+            None,
+            None,
+            None,
         )?;
         let input_layernorm =
             rms_norm(cfg.hidden_size, cfg.rms_norm_eps, vb.pp("input_layernorm"))?;

src/models/mod.rs

@@ -1,15 +1,21 @@
+pub mod campplus;
 pub mod common;
 pub mod deepseek_ocr;
+pub mod feature_extractor;
 pub mod fun_asr_nano;
 pub mod glm_asr_nano;
 pub mod hunyuan_ocr;
+pub mod index_tts2;
+pub mod mask_gct;
 pub mod minicpm4;
 pub mod paddleocr_vl;
 pub mod qwen2_5vl;
 pub mod qwen3;
+pub mod qwen3_asr;
 pub mod qwen3vl;
 pub mod rmbg2_0;
 pub mod voxcpm;
+pub mod w2v_bert_2_0;
 
 use aha_openai_dive::v1::resources::chat::{
     ChatCompletionChunkResponse, ChatCompletionParameters, ChatCompletionResponse,
@@ -23,8 +29,9 @@ use crate::models::{
     glm_asr_nano::generate::GlmAsrNanoGenerateModel,
     hunyuan_ocr::generate::HunyuanOCRGenerateModel, minicpm4::generate::MiniCPMGenerateModel,
     paddleocr_vl::generate::PaddleOCRVLGenerateModel, qwen2_5vl::generate::Qwen2_5VLGenerateModel,
-    qwen3::generate::Qwen3GenerateModel, qwen3vl::generate::Qwen3VLGenerateModel,
+    qwen3::generate::Qwen3GenerateModel, qwen3_asr::generate::Qwen3AsrGenerateModel,
-    rmbg2_0::generate::RMBG2_0Model, voxcpm::generate::VoxCPMGenerate,
+    qwen3vl::generate::Qwen3VLGenerateModel, rmbg2_0::generate::RMBG2_0Model,
+    voxcpm::generate::VoxCPMGenerate,
 };
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq, clap::ValueEnum)]
@@ -37,7 +44,11 @@ pub enum WhichModel {
     Qwen2_5vl7B,
     #[value(name = "qwen3-0.6b", hide = true)]
     Qwen3_0_6B,
-    #[value(name = "qwen3vl-2b", hide = true)]
+    #[value(name = "qwen3asr-0.6b", hide = true)]
+    Qwen3ASR0_6B,
+    #[value(name = "qwen3asr-1.7b", hide = true)]
+    Qwen3ASR1_7B,
+    #[value(name = "qwen3vl-4b", hide = true)]
     Qwen3vl2B,
     #[value(name = "qwen3vl-4b", hide = true)]
     Qwen3vl4B,
@@ -51,7 +62,7 @@ pub enum WhichModel {
     HunyuanOCR,
     #[value(name = "paddleocr-vl", hide = true)]
     PaddleOCRVL,
-    #[value(name = "RMBG2.0", hide = true)]
+    #[value(name = "rmbg2.0")]
     RMBG2_0,
     #[value(name = "voxcpm", hide = true)]
     VoxCPM,
@@ -82,6 +93,7 @@ pub enum ModelInstance<'a> {
     MiniCPM4(MiniCPMGenerateModel<'a>),
     Qwen2_5VL(Qwen2_5VLGenerateModel<'a>),
     Qwen3(Qwen3GenerateModel<'a>),
+    Qwen3ASR(Qwen3AsrGenerateModel<'a>),
     Qwen3VL(Qwen3VLGenerateModel<'a>),
     DeepSeekOCR(DeepseekOCRGenerateModel),
     HunyuanOCR(HunyuanOCRGenerateModel<'a>),
@@ -98,6 +110,7 @@ impl<'a> GenerateModel for ModelInstance<'a> {
             ModelInstance::MiniCPM4(model) => model.generate(mes),
             ModelInstance::Qwen2_5VL(model) => model.generate(mes),
             ModelInstance::Qwen3(model) => model.generate(mes),
+            ModelInstance::Qwen3ASR(model) => model.generate(mes),
             ModelInstance::Qwen3VL(model) => model.generate(mes),
             ModelInstance::DeepSeekOCR(model) => model.generate(mes),
             ModelInstance::HunyuanOCR(model) => model.generate(mes),
@@ -125,6 +138,7 @@ impl<'a> GenerateModel for ModelInstance<'a> {
             ModelInstance::Qwen2_5VL(model) => model.generate_stream(mes),
             ModelInstance::Qwen3(model) => model.generate_stream(mes),
             ModelInstance::Qwen3VL(model) => model.generate_stream(mes),
+            ModelInstance::Qwen3ASR(model) => model.generate_stream(mes),
             ModelInstance::DeepSeekOCR(model) => model.generate_stream(mes),
             ModelInstance::HunyuanOCR(model) => model.generate_stream(mes),
             ModelInstance::PaddleOCRVL(model) => model.generate_stream(mes),
@@ -154,6 +168,14 @@ pub fn load_model(model_type: WhichModel, path: &str) -> Result<ModelInstance<'_
             let model = Qwen3GenerateModel::init(path, None, None)?;
             ModelInstance::Qwen3(model)
         }
+        WhichModel::Qwen3ASR0_6B => {
+            let model = Qwen3AsrGenerateModel::init(path, None, None)?;
+            ModelInstance::Qwen3ASR(model)
+        }
+        WhichModel::Qwen3ASR1_7B => {
+            let model = Qwen3AsrGenerateModel::init(path, None, None)?;
+            ModelInstance::Qwen3ASR(model)
+        }
         WhichModel::Qwen3vl2B => {
             let model = Qwen3VLGenerateModel::init(path, None, None)?;
             ModelInstance::Qwen3VL(model)

src/models/qwen2_5vl/model.rs

@@ -180,6 +180,9 @@ impl Qwen2_5VLVisionBlock {
             cfg.vision_config.intermediate_size,
             cfg.vision_config.hidden_act,
             true,
+            None,
+            None,
+            None,
         )?;
         let norm1 = rms_norm(
             cfg.vision_config.hidden_size,
@@ -616,6 +619,9 @@ impl Qwen2_5VLTextDecoderLayer {
             cfg.intermediate_size,
             cfg.hidden_act,
             false,
+            None,
+            None,
+            None,
         )?;
         let input_layernorm =
             rms_norm(cfg.hidden_size, cfg.rms_norm_eps, vb.pp("input_layernorm"))?;

src/models/qwen3/model.rs

@@ -150,6 +150,9 @@ impl Qwen3DecoderLayer {
             config.intermediate_size,
             config.hidden_act,
             false,
+            None,
+            None,
+            None,
         )?;
         let input_layernorm = rms_norm(
             config.hidden_size,

src/models/qwen3_asr/config.rs

@@ -0,0 +1,208 @@
+use candle_nn::Activation;
+
+use crate::models::qwen3::config::Qwen3Config;
+
+#[derive(Debug, Clone, PartialEq, serde::Deserialize)]
+pub struct Qwen3ASRConfig {
+    pub model_type: String,
+    pub support_languages: Vec<String>,
+    pub thinker_config: ThinkerConfig,
+}
+
+#[derive(Debug, Clone, PartialEq, serde::Deserialize)]
+pub struct ThinkerConfig {
+    pub model_type: String,
+    pub audio_config: Qwen3ASRAudioConfig,
+    pub audio_end_token_id: u32,
+    pub audio_start_token_id: u32,
+    pub audio_token_id: u32,
+    pub dtype: String,
+    pub initializer_range: f64,
+    pub text_config: Qwen3ASRTextConfig,
+}
+
+#[derive(Debug, Clone, PartialEq, serde::Deserialize)]
+pub struct Qwen3ASRAudioConfig {
+    pub activation_dropout: f32,
+    pub activation_function: Activation,
+    pub add_cross_attention: bool,
+    pub attention_dropout: f32,
+    pub bad_words_ids: Option<Vec<u32>>,
+    pub begin_suppress_tokens: Option<Vec<u32>>,
+    pub bos_token_id: Option<u32>,
+    pub chunk_size_feed_forward: usize,
+    pub conv_chunksize: usize,
+    pub cross_attention_hidden_size: Option<usize>,
+    pub d_model: usize,
+    pub decoder_start_token_id: Option<u32>,
+    pub diversity_penalty: f64,
+    pub do_sample: bool,
+    pub downsample_hidden_size: usize,
+    pub dropout: f32,
+    pub early_stopping: bool,
+    pub encoder_attention_heads: usize,
+    pub encoder_ffn_dim: usize,
+    pub encoder_layers: usize,
+    pub encoder_no_repeat_ngram_size: usize,
+    pub eos_token_id: Option<u32>,
+    pub exponential_decay_length_penalty: Option<Vec<f32>>,
+    pub forced_bos_token_id: Option<u32>,
+    pub forced_eos_token_id: Option<u32>,
+    pub id2label: std::collections::HashMap<String, String>,
+    pub initializer_range: f64,
+    pub is_decoder: bool,
+    pub is_encoder_decoder: bool,
+    pub label2id: std::collections::HashMap<String, usize>,
+    pub length_penalty: f64,
+    pub max_length: usize,
+    pub max_source_positions: usize,
+    pub min_length: usize,
+    pub model_type: String,
+    pub n_window: usize,
+    pub n_window_infer: usize,
+    pub no_repeat_ngram_size: usize,
+    pub num_beam_groups: usize,
+    pub num_beams: usize,
+    pub num_hidden_layers: usize,
+    pub num_mel_bins: usize,
+    pub num_return_sequences: usize,
+    pub output_attentions: bool,
+    pub output_dim: usize,
+    pub output_hidden_states: bool,
+    pub output_scores: bool,
+    pub pad_token_id: Option<u32>,
+    pub prefix: Option<String>,
+    pub problem_type: Option<String>,
+    pub pruned_heads: std::collections::HashMap<String, Vec<i32>>,
+    pub remove_invalid_values: bool,
+    pub repetition_penalty: f64,
+    pub return_dict: bool,
+    pub return_dict_in_generate: bool,
+    pub scale_embedding: bool,
+    pub sep_token_id: Option<u32>,
+    pub suppress_tokens: Option<Vec<u32>>,
+    pub task_specific_params: Option<std::collections::HashMap<String, serde_json::Value>>,
+    pub temperature: f64,
+    pub tf_legacy_loss: bool,
+    pub tie_encoder_decoder: bool,
+    pub tie_word_embeddings: bool,
+    pub tokenizer_class: Option<String>,
+    pub top_k: usize,
+    pub top_p: f64,
+    pub torchscript: bool,
+    pub typical_p: f64,
+    pub use_bfloat16: bool,
+}
+
+#[derive(Debug, Clone, PartialEq, serde::Deserialize)]
+pub struct Qwen3ASRTextConfig {
+    pub add_cross_attention: bool,
+    pub attention_bias: bool,
+    pub attention_dropout: f64,
+    pub bad_words_ids: Option<Vec<u32>>,
+    pub begin_suppress_tokens: Option<Vec<u32>>,
+    pub bos_token_id: Option<u32>,
+    pub chunk_size_feed_forward: usize,
+    pub cross_attention_hidden_size: Option<usize>,
+    pub decoder_start_token_id: Option<u32>,
+    pub diversity_penalty: f64,
+    pub do_sample: bool,
+    pub dtype: Option<String>,
+    pub early_stopping: bool,
+    pub encoder_no_repeat_ngram_size: usize,
+    pub eos_token_id: Option<u32>,
+    pub exponential_decay_length_penalty: Option<Vec<f32>>,
+    pub forced_bos_token_id: Option<u32>,
+    pub forced_eos_token_id: Option<u32>,
+    pub head_dim: usize,
+    pub hidden_act: Activation,
+    pub hidden_size: usize,
+    pub id2label: std::collections::HashMap<String, String>,
+    pub initializer_range: f64,
+    pub intermediate_size: usize,
+    pub is_decoder: bool,
+    pub is_encoder_decoder: bool,
+    pub label2id: std::collections::HashMap<String, u32>,
+    pub length_penalty: f64,
+    pub max_length: usize,
+    pub max_position_embeddings: usize,
+    pub min_length: usize,
+    pub model_type: String,
+    pub no_repeat_ngram_size: usize,
+    pub num_attention_heads: usize,
+    pub num_beam_groups: usize,
+    pub num_beams: usize,
+    pub num_hidden_layers: usize,
+    pub num_key_value_heads: usize,
+    pub num_return_sequences: usize,
+    pub output_attentions: bool,
+    pub output_hidden_states: bool,
+    pub output_scores: bool,
+    pub pad_token_id: Option<u32>,
+    pub prefix: Option<String>,
+    pub problem_type: Option<String>,
+    pub remove_invalid_values: bool,
+    pub repetition_penalty: f64,
+    pub return_dict: bool,
+    pub return_dict_in_generate: bool,
+    pub rms_norm_eps: f64,
+    pub rope_scaling: Qwen3ASRRopeScaling,
+    pub rope_theta: f32,
+    pub sep_token_id: Option<u32>,
+    pub suppress_tokens: Option<Vec<u32>>,
+    pub temperature: f64,
+    pub tf_legacy_loss: bool,
+    pub tie_encoder_decoder: bool,
+    pub tie_word_embeddings: bool,
+    pub tokenizer_class: Option<String>,
+    pub top_k: usize,
+    pub top_p: f64,
+    pub torchscript: bool,
+    pub typical_p: f64,
+    pub use_bfloat16: bool,
+    pub use_cache: bool,
+    pub vocab_size: usize,
+}
+
+pub fn qwen3asr_text_config2qwen3_config(cfg: &Qwen3ASRTextConfig) -> Qwen3Config {
+    Qwen3Config {
+        attention_bias: cfg.attention_bias,
+        attention_dropout: cfg.attention_dropout as f64,
+        bos_token_id: cfg.bos_token_id.unwrap_or(151643) as u32,
+        eos_token_id: cfg.eos_token_id.unwrap_or(151645) as u32,
+        head_dim: cfg.head_dim,
+        hidden_act: cfg.hidden_act,
+        hidden_size: cfg.hidden_size,
+        initializer_range: cfg.initializer_range as f64,
+        intermediate_size: cfg.intermediate_size,
+        max_position_embeddings: cfg.max_position_embeddings,
+        max_window_layers: 0,
+        num_attention_heads: cfg.num_attention_heads,
+        num_hidden_layers: cfg.num_hidden_layers,
+        num_key_value_heads: cfg.num_key_value_heads,
+        rms_norm_eps: cfg.rms_norm_eps,
+        rope_theta: cfg.rope_theta,
+        tie_word_embeddings: true,
+        torch_dtype: "bfloat16".to_string(),
+        use_cache: cfg.use_cache,
+        use_sliding_window: false,
+        vocab_size: cfg.vocab_size,
+    }
+}
+
+#[derive(Debug, Clone, PartialEq, serde::Deserialize)]
+pub struct Qwen3ASRRopeScaling {
+    pub interleaved: bool,
+    pub mrope_interleaved: bool,
+    pub mrope_section: Vec<usize>,
+    pub rope_type: String,
+    pub r#type: String,
+}
+
+#[derive(Debug, Clone, PartialEq, serde::Deserialize)]
+pub struct Qwen3ASRGenerationConfig {
+    pub do_sample: bool,
+    pub eos_token_id: Vec<usize>,
+    pub pad_token_id: usize,
+    pub temperature: f32,
+}

src/models/qwen3_asr/generate.rs

@@ -0,0 +1,193 @@
+
+use aha_openai_dive::v1::resources::chat::{
+    ChatCompletionChunkResponse, ChatCompletionParameters, ChatCompletionResponse,
+};
+use anyhow::{Result, anyhow};
+use candle_core::{DType, Device, Tensor};
+use candle_nn::VarBuilder;
+use rocket::async_stream::stream;
+use rocket::futures::Stream;
+
+use crate::{
+    chat_template::ChatTemplate,
+    models::{
+        GenerateModel,
+        feature_extractor::config::FeatureExtractor,
+        qwen3_asr::{
+            config::{Qwen3ASRConfig, Qwen3ASRGenerationConfig},
+            model::Qwen3ASRModel,
+            processor::Qwen3AsrProcessor,
+        },
+    },
+    tokenizer::TokenizerModel,
+    utils::{
+        build_completion_chunk_response, build_completion_response, find_type_files, get_device,
+        get_dtype, get_logit_processor,
+    },
+};
+
+pub struct Qwen3AsrGenerateModel<'a> {
+    chat_template: ChatTemplate<'a>,
+    tokenizer: TokenizerModel,
+    processor: Qwen3AsrProcessor,
+    qwen3_asr: Qwen3ASRModel,
+    device: Device,
+    dtype: DType,
+    eos_token_id1: u32,
+    eos_token_id2: u32,
+    generation_config: Qwen3ASRGenerationConfig,
+    model_name: String,
+}
+
+impl<'a> Qwen3AsrGenerateModel<'a> {
+    pub fn init(path: &str, device: Option<&Device>, dtype: Option<DType>) -> Result<Self> {
+        let chat_template = ChatTemplate::init(path)?;
+        let tokenizer = TokenizerModel::init(path)?;
+        let generation_config_path = path.to_string() + "/generation_config.json";
+        let generation_config: Qwen3ASRGenerationConfig =
+            serde_json::from_slice(&std::fs::read(generation_config_path)?)?;
+        let device = get_device(device);
+        let preprocess_config_path = path.to_string() + "/preprocessor_config.json";
+        let preprocess_config: FeatureExtractor =
+            serde_json::from_slice(&std::fs::read(preprocess_config_path)?)?;
+        let processor = Qwen3AsrProcessor::new(&device, &preprocess_config)?;
+        let config_path = path.to_string() + "/config.json";
+        let cfg: Qwen3ASRConfig = serde_json::from_slice(&std::fs::read(config_path)?)?;
+        let cfg_dtype = cfg.thinker_config.dtype.as_str();
+        let dtype = get_dtype(dtype, cfg_dtype);
+        let model_list = find_type_files(path, "safetensors")?;
+        let vb = unsafe { VarBuilder::from_mmaped_safetensors(&model_list, dtype, &device)? };
+        let qwen3_asr = Qwen3ASRModel::new(vb, &cfg)?;
+
+        Ok(Self {
+            chat_template,
+            tokenizer,
+            processor,
+            qwen3_asr,
+            device,
+            dtype,
+            eos_token_id1: generation_config.eos_token_id[0] as u32,
+            eos_token_id2: generation_config.eos_token_id[1] as u32,
+            generation_config,
+            model_name: "qwen3-asr".to_string(),
+        })
+    }
+}
+
+impl<'a> GenerateModel for Qwen3AsrGenerateModel<'a> {
+    fn generate(&mut self, mes: ChatCompletionParameters) -> Result<ChatCompletionResponse> {
+        let temperature = match mes.temperature {
+            None => self.generation_config.temperature,
+            Some(tem) => tem,
+        };
+        let seed = match mes.seed {
+            None => 34562u64,
+            Some(s) => s as u64,
+        };
+        let mut logit_processor = get_logit_processor(Some(temperature), mes.top_p, None, seed);
+        let render_text = self.chat_template.apply_chat_template(&mes)?;
+        let audio_datas = self
+            .processor
+            .process_info(&mes, &render_text, &self.tokenizer)?;
+        let sample_len = mes.max_tokens.unwrap_or(1024);
+        let mut generate = Vec::new();
+        for data in audio_datas.iter() {
+            let mut input_ids = data.input_ids.clone();
+            let mut input_features = Some(data.input_features.clone().to_dtype(self.dtype)?);
+            let mut seq_len = input_ids.dim(1)?;
+            let mut seqlen_offset = 0;
+            for _ in 0..sample_len {
+                let logits =
+                    self.qwen3_asr
+                        .forward(&input_ids, seqlen_offset, input_features.as_ref())?;
+                let logits = logits.squeeze(0)?.squeeze(0)?.to_dtype(DType::F32)?;
+                let next_token = logit_processor.sample(&logits)?;
+                generate.push(next_token);
+                if next_token == self.eos_token_id1 || next_token == self.eos_token_id2 {
+                    break;
+                }
+                seqlen_offset += seq_len;
+                seq_len = 1;
+                input_ids = Tensor::from_vec(vec![next_token], (1, 1), &self.device)?;
+                input_features = None;
+            }
+            self.qwen3_asr.clear_kv_cache();
+        }
+        let num_token = generate.len() as u32;
+        let res = self.tokenizer.token_decode(generate)?;
+        let response = build_completion_response(res, &self.model_name, Some(num_token));
+        Ok(response)
+    }
+
+    fn generate_stream(
+        &mut self,
+        mes: ChatCompletionParameters,
+    ) -> Result<
+        Box<
+            dyn Stream<Item = Result<ChatCompletionChunkResponse, anyhow::Error>>
+                + Send
+                + Unpin
+                + '_,
+        >,
+    > {
+        let temperature = match mes.temperature {
+            None => self.generation_config.temperature,
+            Some(tem) => tem,
+        };
+        let seed = match mes.seed {
+            None => 34562u64,
+            Some(s) => s as u64,
+        };
+        let mut logit_processor = get_logit_processor(Some(temperature), mes.top_p, None, seed);
+        let render_text = self.chat_template.apply_chat_template(&mes)?;
+        let audio_datas = self
+            .processor
+            .process_info(&mes, &render_text, &self.tokenizer)?;
+        let sample_len = mes.max_tokens.unwrap_or(1024);
+        let stream = stream! {
+            let mut error_tokens = Vec::new();
+            for data in audio_datas.iter() {
+                let mut input_ids = data.input_ids.clone();
+                let mut input_features = Some(data.input_features.clone().to_dtype(self.dtype)?);
+                let mut seq_len = input_ids.dim(1)?;
+                let mut seqlen_offset = 0;
+                for _ in 0..sample_len {
+                    let logits =
+                        self.qwen3_asr
+                            .forward(&input_ids, seqlen_offset, input_features.as_ref())?;
+                    let logits = logits.squeeze(0)?.squeeze(0)?.to_dtype(DType::F32)?;
+                    let next_token = logit_processor.sample(&logits)?;
+                    let mut decode_ids = Vec::new();
+                    if !error_tokens.is_empty() {
+                        decode_ids.extend_from_slice(&error_tokens);
+                    }
+                    decode_ids.push(next_token);
+                    let decoded_token = self.tokenizer.token_decode(decode_ids).map_err(|e| anyhow!(format!("stream decode error{e}")))?;
+                    if decoded_token.contains("�") {
+                        error_tokens.push(next_token);
+                        if error_tokens.len() > 3 {
+                            error_tokens.clear();
+                        }
+                        seqlen_offset += seq_len;
+                        seq_len = 1;
+                        input_ids = Tensor::from_vec(vec![next_token], (1, 1), &self.device)?;
+                        input_features = None;
+                        continue;
+                    }
+                    error_tokens.clear();
+                    let chunk = build_completion_chunk_response(decoded_token, &self.model_name, None, None);
+                    yield Ok(chunk);
+                    if next_token == self.eos_token_id1 || next_token == self.eos_token_id2 {
+                        break;
+                    }
+                    seqlen_offset += seq_len;
+                    seq_len = 1;
+                    input_ids = Tensor::from_vec(vec![next_token], (1, 1), &self.device)?;
+                    input_features = None;
+                }
+                self.qwen3_asr.clear_kv_cache();
+            }
+        };
+        Ok(Box::new(Box::pin(stream)))
+    }
+}

src/models/qwen3_asr/mod.rs

@@ -0,0 +1,4 @@
+pub mod config;
+pub mod generate;
+pub mod model;
+pub mod processor;

src/models/qwen3_asr/model.rs

@@ -0,0 +1,388 @@
+use anyhow::Result;
+use candle_core::Tensor;
+use candle_nn::{
+    Activation, Conv2d, Embedding, LayerNorm, Linear, Module, RmsNorm, VarBuilder, embedding,
+    linear, linear_no_bias, rms_norm,
+};
+
+use crate::{
+    models::{
+        common::{NaiveAttention, get_conv2d, get_layer_norm},
+        qwen3::model::Qwen3DecoderLayer,
+        qwen3_asr::{
+            config::{
+                Qwen3ASRAudioConfig, Qwen3ASRConfig, Qwen3ASRTextConfig, ThinkerConfig,
+                qwen3asr_text_config2qwen3_config,
+            },
+            processor::get_feat_extract_output_lengths,
+        },
+    },
+    position_embed::{
+        rope::Qwen3VLTextRotaryEmbedding, sinusoidal_pe::SinusoidalPositionEncoderCat,
+    },
+    utils::tensor_utils::{
+        get_equal_mask, masked_scatter_dim0, prepare_causal_attention_mask, split_tensor,
+        split_tensor_with_size,
+    },
+};
+
+pub struct Qwen3ASRAudioEncoderLayer {
+    self_attn: NaiveAttention,
+    self_attn_layer_norm: LayerNorm,
+    activation_fn: Activation,
+    fc1: Linear,
+    fc2: Linear,
+    final_layer_norm: LayerNorm,
+}
+
+impl Qwen3ASRAudioEncoderLayer {
+    pub fn new(vb: VarBuilder, config: &Qwen3ASRAudioConfig) -> Result<Self> {
+        let self_attn = NaiveAttention::new(
+            vb.pp("self_attn"),
+            config.d_model,
+            config.encoder_attention_heads,
+            config.encoder_attention_heads,
+            None,
+            true,
+            Some("q_proj"),
+            Some("k_proj"),
+            Some("v_proj"),
+            Some("out_proj"),
+        )?;
+        let self_attn_layer_norm =
+            get_layer_norm(vb.pp("self_attn_layer_norm"), 1e-5, config.d_model)?;
+        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)?;
+        Ok(Self {
+            self_attn,
+            self_attn_layer_norm,
+            activation_fn,
+            fc1,
+            fc2,
+            final_layer_norm,
+        })
+    }
+
+    pub fn forward(&self, xs: &Tensor, mask: Option<&Tensor>) -> Result<Tensor> {
+        let residual = xs.clone();
+        let xs = self.self_attn_layer_norm.forward(xs)?;
+        let xs = self.self_attn.forward(&xs, None, None, mask, false)?;
+        let residual = xs.add(&residual)?;
+        let xs = self.final_layer_norm.forward(&residual)?;
+        let xs = self.fc1.forward(&xs)?.apply(&self.activation_fn)?;
+        let xs = self.fc2.forward(&xs)?;
+        let xs = xs.add(&residual)?;
+        Ok(xs)
+    }
+}
+
+pub struct Qwen3ASRAudioEncoder {
+    n_window: usize,
+    positional_embedding: SinusoidalPositionEncoderCat,
+    layers: Vec<Qwen3ASRAudioEncoderLayer>,
+    ln_post: LayerNorm,
+    conv2d1: Conv2d,
+    conv2d2: Conv2d,
+    conv2d3: Conv2d,
+    conv_out: Linear,
+    proj1: Linear,
+    act: Activation,
+    proj2: Linear,
+    // n_window_infer: usize,
+    conv_chunksize: usize,
+}
+
+impl Qwen3ASRAudioEncoder {
+    pub fn new(vb: VarBuilder, config: &Qwen3ASRAudioConfig) -> Result<Self> {
+        let n_window = config.n_window;
+        let positional_embedding =
+            SinusoidalPositionEncoderCat::new(Some(config.d_model), true, vb.device())?;
+        let mut layers = vec![];
+        let vb_layers = vb.pp("layers");
+        for i in 0..config.encoder_layers {
+            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 conv2d1 = get_conv2d(
+            vb.pp("conv2d1"),
+            1,
+            config.downsample_hidden_size,
+            3,
+            1,
+            2,
+            1,
+            1,
+            true,
+        )?;
+        let conv2d2 = get_conv2d(
+            vb.pp("conv2d2"),
+            config.downsample_hidden_size,
+            config.downsample_hidden_size,
+            3,
+            1,
+            2,
+            1,
+            1,
+            true,
+        )?;
+        let conv2d3 = get_conv2d(
+            vb.pp("conv2d3"),
+            config.downsample_hidden_size,
+            config.downsample_hidden_size,
+            3,
+            1,
+            2,
+            1,
+            1,
+            true,
+        )?;
+        let in_dim =
+            config.downsample_hidden_size * ((((config.num_mel_bins + 1) / 2 + 1) / 2 + 1) / 2);
+        let conv_out = linear_no_bias(in_dim, config.d_model, vb.pp("conv_out"))?;
+        let proj1 = linear(config.d_model, config.d_model, vb.pp("proj1"))?;
+        let act = config.activation_function;
+        let proj2 = linear(config.d_model, config.output_dim, vb.pp("proj2"))?;
+        // let n_window_infer = config.n_window_infer;
+        let conv_chunksize = config.conv_chunksize;
+        Ok(Self {
+            n_window,
+            positional_embedding,
+            layers,
+            ln_post,
+            conv2d1,
+            conv2d2,
+            conv2d3,
+            conv_out,
+            proj1,
+            act,
+            proj2,
+            // n_window_infer,
+            conv_chunksize,
+        })
+    }
+
+    pub fn forward(&self, xs: &Tensor) -> Result<Tensor> {
+        // xs: (feature_dim, feature_len)
+        let feature_lens = xs.dim(1)?;
+        // let aftercnn_lens = get_feat_extract_output_lengths(feature_lens);
+        let chunk_num = feature_lens / (self.n_window * 2);
+        let mut chunk_lengths = vec![self.n_window * 2; chunk_num];
+        let last = feature_lens % (self.n_window * 2);
+        if last > 0 {
+            chunk_lengths.push(last);
+        }
+        let mut chunk_list = split_tensor(&xs.t()?, &chunk_lengths, 0)?;
+        if last > 0 {
+            let chunk_last = chunk_list
+                .pop()
+                .ok_or(anyhow::anyhow!(format!("chunk_list is empty")))?;
+            let pad_size = self.n_window * 2 - last;
+            let chunk_last = chunk_last.pad_with_zeros(0, 0, pad_size)?;
+            chunk_list.push(chunk_last);
+        }
+        let padded_feature = Tensor::stack(&chunk_list, 0)?.transpose(1, 2)?;
+        let feature_lens_after_cnn: Vec<usize> = chunk_lengths
+            .iter()
+            .map(|&i| get_feat_extract_output_lengths(i))
+            .collect();
+        let feature_len_after_cnn = feature_lens_after_cnn.iter().sum();
+        let padded_feature = padded_feature.unsqueeze(1)?;
+        let mut padded_embeds = vec![];
+        let feature_splits = split_tensor_with_size(&padded_feature, self.conv_chunksize, 0)?;
+        for chunk in feature_splits.iter() {
+            let padded_embed = self.conv2d1.forward(chunk)?.gelu()?;
+            let padded_embed = self.conv2d2.forward(&padded_embed)?.gelu()?;
+            let padded_embed = self.conv2d3.forward(&padded_embed)?.gelu()?;
+            padded_embeds.push(padded_embed);
+        }
+        let padded_embed = Tensor::cat(&padded_embeds, 0)?;
+        let (b, c, f, t) = padded_embed.dims4()?;
+        let padded_embed =
+            padded_embed
+                .permute((0, 3, 1, 2))?
+                .contiguous()?
+                .reshape((b, t, c * f))?;
+        let padded_embed = self.conv_out.forward(&padded_embed)?;
+        let padded_embed = self.positional_embedding.forward(&padded_embed, 0)?;
+        let padded_embed = padded_embed.flatten(0, 1)?;
+        let mut hidden_states = padded_embed
+            .narrow(0, 0, feature_len_after_cnn)?
+            .unsqueeze(0)?;
+        for layer in &self.layers {
+            hidden_states = layer.forward(&hidden_states, None)?;
+        }
+        let hidden_states = hidden_states.squeeze(0)?;
+        let hidden_states = self.ln_post.forward(&hidden_states)?;
+        let hidden_states = self.proj1.forward(&hidden_states)?.apply(&self.act)?;
+        let hidden_states = self.proj2.forward(&hidden_states)?;
+        Ok(hidden_states)
+    }
+}
+
+pub struct Qwen3ASRThinkerTextModel {
+    embed_tokens: Embedding,
+    layers: Vec<Qwen3DecoderLayer>,
+    norm: RmsNorm,
+    rotary_emb: Qwen3VLTextRotaryEmbedding,
+    mrope_section: Vec<usize>,
+}
+
+impl Qwen3ASRThinkerTextModel {
+    pub fn new(vb: VarBuilder, cfg: &Qwen3ASRTextConfig) -> Result<Self> {
+        let embed_tokens = embedding(cfg.vocab_size, cfg.hidden_size, vb.pp("embed_tokens"))?;
+        let mut layers = vec![];
+        let vb_layers = vb.pp("layers");
+        let qwen3cfg = qwen3asr_text_config2qwen3_config(cfg);
+        for i in 0..cfg.num_hidden_layers {
+            let layer = Qwen3DecoderLayer::new(&qwen3cfg, vb_layers.pp(i))?;
+            layers.push(layer);
+        }
+        let norm = rms_norm(cfg.hidden_size, cfg.rms_norm_eps, vb.pp("norm"))?;
+        let rotary_emb = Qwen3VLTextRotaryEmbedding::new(cfg.head_dim, cfg.rope_theta);
+        Ok(Self {
+            embed_tokens,
+            layers,
+            norm,
+            rotary_emb,
+            mrope_section: cfg.rope_scaling.mrope_section.clone(),
+        })
+    }
+
+    pub fn forward(
+        &mut self,
+        input_embeds: &Tensor,
+        seqlen_offset: usize,
+        position_ids: Option<&Tensor>,
+    ) -> Result<Tensor> {
+        let (b_size, seq_len, _) = input_embeds.dims3()?;
+        let position_ids = match position_ids {
+            Some(ids) => ids.clone(),
+            None => Tensor::arange(
+                seqlen_offset as u32,
+                (seq_len + seqlen_offset) as u32,
+                input_embeds.device(),
+            )?
+            .unsqueeze(0)?
+            .unsqueeze(0)?
+            .broadcast_as((3, b_size, seq_len))?,
+        };
+        let (cos, sin) = self.rotary_emb.forward_asr(
+            &position_ids,
+            input_embeds.dtype(),
+            self.mrope_section.clone(),
+        )?;
+        let mut xs = input_embeds.clone();
+        let attention_mask: Option<Tensor> = {
+            if seq_len <= 1 {
+                None
+            } else {
+                Some(prepare_causal_attention_mask(
+                    b_size,
+                    seq_len,
+                    0,
+                    input_embeds.device(),
+                )?)
+            }
+        };
+        for layer in self.layers.iter_mut() {
+            xs = layer.forward(&xs, &cos, &sin, attention_mask.as_ref())?;
+        }
+        let xs = self.norm.forward(&xs)?;
+        Ok(xs)
+    }
+
+    pub fn clear_kv_cache(&mut self) {
+        for layer in self.layers.iter_mut() {
+            layer.clear_kv_cache()
+        }
+    }
+}
+
+pub struct Qwen3ASRThinker {
+    audio_tower: Qwen3ASRAudioEncoder,
+    model: Qwen3ASRThinkerTextModel,
+    audio_token_id: u32,
+    lm_head: Linear,
+}
+
+impl Qwen3ASRThinker {
+    pub fn new(vb: VarBuilder, config: &ThinkerConfig) -> Result<Self> {
+        let audio_tower = Qwen3ASRAudioEncoder::new(vb.pp("audio_tower"), &config.audio_config)?;
+        let model = Qwen3ASRThinkerTextModel::new(vb.pp("model"), &config.text_config)?;
+        let lm_head = if config.text_config.tie_word_embeddings {
+            Linear::new(model.embed_tokens.embeddings().clone(), None)
+        } else {
+            linear_no_bias(
+                config.text_config.hidden_size,
+                config.text_config.vocab_size,
+                vb.pp("lm_head"),
+            )?
+        };
+        Ok(Self {
+            audio_tower,
+            model,
+            audio_token_id: config.audio_token_id,
+            lm_head,
+        })
+    }
+
+    pub fn forward(
+        &mut self,
+        input_ids: &Tensor,
+        seqlen_offset: usize,
+        input_features: Option<&Tensor>,
+    ) -> Result<Tensor> {
+        let mut input_embeds = self.model.embed_tokens.forward(input_ids)?;
+        if let Some(input_features) = input_features {
+            let audio_feature = self.audio_tower.forward(&input_features)?;
+            // println!("audio_feature: {}", audio_feature);
+            let audio_mask = get_equal_mask(input_ids, self.audio_token_id)?;
+            let n_audio_tokens = audio_mask.sum_all()?.to_scalar::<u32>()?;
+            if n_audio_tokens as usize != audio_feature.dim(0)? {
+                return Err(anyhow::anyhow!(format!(
+                    "n_audio_tokens num: {} not equal to audio_feature len: {}",
+                    n_audio_tokens,
+                    audio_feature.dim(0)?
+                )));
+            }
+            input_embeds = masked_scatter_dim0(&input_embeds, &audio_feature, &audio_mask)?;
+        }
+        let outputs = self.model.forward(&input_embeds, seqlen_offset, None)?;
+        let seq_len = outputs.dim(1)?;
+        let hidden_state = outputs.narrow(1, seq_len - 1, 1)?;
+        let logits = self.lm_head.forward(&hidden_state)?;
+        Ok(logits)
+    }
+    pub fn clear_kv_cache(&mut self) {
+        self.model.clear_kv_cache();
+    }
+}
+
+pub struct Qwen3ASRModel {
+    thinker: Qwen3ASRThinker,
+}
+
+impl Qwen3ASRModel {
+    pub fn new(vb: VarBuilder, config: &Qwen3ASRConfig) -> Result<Self> {
+        let thinker = Qwen3ASRThinker::new(vb.pp("thinker"), &config.thinker_config)?;
+        Ok(Self { thinker })
+    }
+
+    pub fn forward(
+        &mut self,
+        input_ids: &Tensor,
+        seqlen_offset: usize,
+        input_features: Option<&Tensor>,
+    ) -> Result<Tensor> {
+        let logits = self
+            .thinker
+            .forward(input_ids, seqlen_offset, input_features)?;
+        Ok(logits)
+    }
+    pub fn clear_kv_cache(&mut self) {
+        self.thinker.clear_kv_cache();
+    }
+}

src/models/qwen3_asr/processor.rs

@@ -0,0 +1,172 @@
+use aha_openai_dive::v1::resources::chat::ChatCompletionParameters;
+use anyhow::Result;
+use candle_core::{Device, Tensor};
+
+use crate::{
+    models::feature_extractor::{
+        config::FeatureExtractor, feature_extraction_whisper::WhisperFeatureExtractor,
+    },
+    tokenizer::TokenizerModel,
+    utils::{
+        audio_utils::{extract_audios, split_audio_into_chunks},
+        capitalize_first_letter,
+        tensor_utils::float_range_normalize,
+    },
+};
+
+pub struct Qwen3AsrProcessor {
+    device: Device,
+    sample_rate: usize,
+    support_language: Vec<String>,
+    max_asr_input_seconds: f32,
+    whisper_feature_extracor: WhisperFeatureExtractor,
+    audio_token: String,
+}
+
+impl Qwen3AsrProcessor {
+    pub fn new(device: &Device, config: &FeatureExtractor) -> Result<Self> {
+        let support_language: Vec<String> = vec![
+            "Chinese",
+            "English",
+            "Cantonese",
+            "Arabic",
+            "German",
+            "French",
+            "Spanish",
+            "Portuguese",
+            "Indonesian",
+            "Italian",
+            "Korean",
+            "Russian",
+            "Thai",
+            "Vietnamese",
+            "Japanese",
+            "Turkish",
+            "Hindi",
+            "Malay",
+            "Dutch",
+            "Swedish",
+            "Danish",
+            "Finnish",
+            "Polish",
+            "Czech",
+            "Filipino",
+            "Persian",
+            "Greek",
+            "Romanian",
+            "Hungarian",
+            "Macedonian",
+        ]
+        .iter()
+        .map(|s| s.to_string())
+        .collect();
+        let whisper_feature_extracor = WhisperFeatureExtractor::new(
+            config.feature_size,
+            config.hop_length,
+            // config.chunk_length,
+            config.n_fft,
+            config.dither,
+            // config.padding_value,
+            config.sampling_rate,
+            device,
+        )?;
+        Ok(Self {
+            device: device.clone(),
+            sample_rate: 16000,
+            support_language,
+            max_asr_input_seconds: 1200.0,
+            whisper_feature_extracor,
+            audio_token: "<|audio_pad|>".to_string(),
+        })
+    }
+
+    pub fn process_audio(&self, mes: &ChatCompletionParameters) -> Result<Vec<Tensor>> {
+        let audio_tensors = extract_audios(mes, &self.device, Some(self.sample_rate))?;
+        audio_tensors
+            .iter()
+            .map(|audio| float_range_normalize(&audio))
+            .collect()
+    }
+
+    pub fn validate_language(&self, lang: &String) -> bool {
+        self.support_language.contains(lang)
+    }
+
+    fn replace_special_tokens(&self, text: &str, token_len: usize) -> String {
+        let replace = "<|audio_placeholder|>".repeat(token_len as usize);
+        let text = text.replacen(&self.audio_token, &replace, 1);
+        let text = text.replace("<|audio_placeholder|>", &self.audio_token);
+        text
+    }
+
+    pub fn process_info(
+        &self,
+        mes: &ChatCompletionParameters,
+        render: &str,
+        tokenizer: &TokenizerModel,
+    ) -> Result<Vec<AudioData>> {
+        let audio_count = render
+            .matches("<|audio_start|><|audio_pad|><|audio_end|>")
+            .count();
+        let mut render = if audio_count > 1 {
+            render.replace(
+                &"<|audio_start|><|audio_pad|><|audio_end|>".repeat(audio_count),
+                "<|audio_start|><|audio_pad|><|audio_end|>",
+            )
+        } else {
+            render.to_string()
+        };
+        if let Some(map) = &mes.metadata
+            && map.contains_key("language")
+        {
+            let lang = map.get("language").unwrap();
+            let lang = capitalize_first_letter(lang);
+            if self.validate_language(&lang) {
+                render = format!("{}language {}'<asr_text>'", render, lang);
+            }
+        }
+        let audio_tensors = self.process_audio(mes)?;
+        let audio_len = audio_tensors.len();
+        if audio_len != audio_count {
+            return Err(anyhow::anyhow!("audio_pad num != audio num"));
+        }
+        let mut split_wavs = vec![];
+        for wav in audio_tensors.iter() {
+            let wavs = split_audio_into_chunks(wav, self.sample_rate, self.max_asr_input_seconds)?;
+            split_wavs.extend_from_slice(&wavs);
+        }
+        let mut audio_datas = vec![];
+        for wav in split_wavs.iter() {
+            let (input_features, _) =
+                self.whisper_feature_extracor
+                    .call(wav, self.sample_rate, false)?;
+            let audio_len = input_features.dim(2)?;
+            let output_len = get_feat_extract_output_lengths(audio_len);
+            let text = self.replace_special_tokens(&render, output_len);
+            let input_ids = tokenizer.text_encode(text, &self.device)?;
+            let input_features = input_features.squeeze(0)?;
+            let audio = AudioData {
+                input_features,
+                input_ids,
+            };
+            audio_datas.push(audio);
+        }
+        Ok(audio_datas)
+    }
+}
+
+pub struct AudioData {
+    pub input_features: Tensor,
+    pub input_ids: Tensor,
+}
+
+pub fn get_feat_extract_output_lengths(audio_len: usize) -> usize {
+    let input_len_leave = audio_len % 100;
+    let output_len = if input_len_leave > 0 {
+        let feat_lengths = (input_len_leave - 1) / 2 + 1;
+        ((feat_lengths - 1) / 2 + 1 - 1) / 2 + 1 + (audio_len / 100) * 13
+    } else {
+        (audio_len / 100) * 13
+    };
+    output_len
+}

src/models/rmbg2_0/generate.rs

@@ -49,7 +49,7 @@ impl RMBG2_0Model {
             img_std,
             device,
             dtype,
-            model_name: "RMBG2.0".to_string(),
+            model_name: "rmbg2.0".to_string(),
         })
     }
 

src/models/rmbg2_0/model.rs

@@ -7,7 +7,8 @@ use candle_nn::{
 
 use crate::{
     models::common::{
-        TwoLinearMLP, deform_conv2d_kernel, get_batch_norm, get_conv2d, get_layer_norm,
+        Conv2dWithBN, TwoLinearMLP, deform_conv2d_kernel, get_batch_norm, get_conv2d,
+        get_layer_norm,
     },
     utils::tensor_utils::{
         get_equal_mask, index_select_2d, interpolate_bilinear, split_tensor_with_size,
@@ -891,7 +892,7 @@ impl _ASPPModuleDeformable {
             padding,
             false,
         )?;
-        let bn = get_batch_norm(vb.pp("bn"), 1e-5, out_c)?;
+        let bn = get_batch_norm(vb.pp("bn"), 1e-5, out_c, true)?;
         Ok(Self { atrous_conv, bn })
     }
 
@@ -957,7 +958,8 @@ impl ASPPDeformable {
             1,
             false,
         )?;
-        let global_avg_pool_2 = get_batch_norm(vb.pp("global_avg_pool.2"), 1e-5, in_channelster)?;
+        let global_avg_pool_2 =
+            get_batch_norm(vb.pp("global_avg_pool.2"), 1e-5, in_channelster, true)?;
         let conv1 = get_conv2d(
             vb.pp("conv1"),
             in_channelster * (2 + parallel_block_sizes.len()),
@@ -969,7 +971,7 @@ impl ASPPDeformable {
             1,
             false,
         )?;
-        let bn1 = get_batch_norm(vb.pp("bn1"), 1e-5, out_c)?;
+        let bn1 = get_batch_norm(vb.pp("bn1"), 1e-5, out_c, true)?;
         Ok(Self {
             aspp1,
             aspp_deforms_0,
@@ -1031,8 +1033,8 @@ impl BasicDecBlk {
             1,
             true,
         )?;
-        let bn_in = get_batch_norm(vb.pp("bn_in"), 1e-5, inter_channels)?;
+        let bn_in = get_batch_norm(vb.pp("bn_in"), 1e-5, inter_channels, true)?;
-        let bn_out = get_batch_norm(vb.pp("bn_out"), 1e-5, out_c)?;
+        let bn_out = get_batch_norm(vb.pp("bn_out"), 1e-5, out_c, true)?;
         Ok(Self {
             conv_in,
             dec_att,
@@ -1072,32 +1074,6 @@ impl SimpleConvs {
     }
 }
 
-struct Conv2dWithBN {
-    conv_0: Conv2d,
-    bn_1: BatchNorm,
-}
-
-impl Conv2dWithBN {
-    pub fn new(
-        vb: VarBuilder,
-        in_c: usize,
-        out_c: usize,
-        ks: usize,
-        padding: usize,
-        stride: usize,
-    ) -> Result<Self> {
-        let conv_0 = get_conv2d(vb.pp("0"), in_c, out_c, ks, padding, stride, 1, 1, true)?;
-        let bn_1 = get_batch_norm(vb.pp("1"), 1e-5, out_c)?;
-        Ok(Self { conv_0, bn_1 })
-    }
-
-    pub fn forward(&self, x: &Tensor) -> Result<Tensor> {
-        let x = self.conv_0.forward(x)?;
-        let x = self.bn_1.forward_t(&x, false)?.relu()?;
-        Ok(x)
-    }
-}
-
 struct Decoder {
     ipt_blk5: SimpleConvs,
     ipt_blk4: SimpleConvs,
@@ -1207,9 +1183,12 @@ impl Decoder {
         // let conv_ms_spvn_2 =
         //     get_conv2d(vb.pp("conv_ms_spvn_2"), channels[3], 1, 1, 0, 1, 1, 1, true)?;
         let n = 16usize;
-        let gdt_convs_4 = Conv2dWithBN::new(vb.pp("gdt_convs_4"), channels[1], n, 3, 1, 1)?;
+        let gdt_convs_4 =
-        let gdt_convs_3 = Conv2dWithBN::new(vb.pp("gdt_convs_3"), channels[2], n, 3, 1, 1)?;
+            Conv2dWithBN::new(vb.pp("gdt_convs_4"), channels[1], n, 3, 1, 1, true, true)?;
-        let gdt_convs_2 = Conv2dWithBN::new(vb.pp("gdt_convs_2"), channels[3], n, 3, 1, 1)?;
+        let gdt_convs_3 =
+            Conv2dWithBN::new(vb.pp("gdt_convs_3"), channels[2], n, 3, 1, 1, true, true)?;
+        let gdt_convs_2 =
+            Conv2dWithBN::new(vb.pp("gdt_convs_2"), channels[3], n, 3, 1, 1, true, true)?;
 
         let gdt_convs_attn_4 = get_conv2d(vb.pp("gdt_convs_attn_4.0"), n, 1, 1, 0, 1, 1, 1, true)?;
         let gdt_convs_attn_3 = get_conv2d(vb.pp("gdt_convs_attn_3.0"), n, 1, 1, 0, 1, 1, 1, true)?;

src/models/voxcpm/minicpm4.rs

@@ -130,6 +130,9 @@ impl MiniCPMDecoderLayer {
             cfg.intermediate_size,
             candle_nn::Activation::Silu,
             false,
+            None,
+            None,
+            None,
         )?;
         let input_layernorm =
             rms_norm(cfg.hidden_size, cfg.rms_norm_eps, vb.pp("input_layernorm"))?;

src/models/w2v_bert_2_0/config.rs

@@ -0,0 +1,61 @@
+use candle_nn::Activation;
+
+#[derive(Debug, Clone, PartialEq, serde::Deserialize)]
+pub struct W2VBert2_0Config {
+    pub activation_dropout: f32,
+    pub adapter_act: String,
+    pub adapter_kernel_size: usize,
+    pub adapter_stride: usize,
+    pub add_adapter: bool,
+    pub apply_spec_augment: bool,
+    pub attention_dropout: f32,
+    pub bos_token_id: usize,
+    pub classifier_proj_size: usize,
+    pub codevector_dim: usize,
+    pub conformer_conv_dropout: f32,
+    pub contrastive_logits_temperature: f32,
+    pub conv_depthwise_kernel_size: usize,
+    pub ctc_loss_reduction: String,
+    pub ctc_zero_infinity: bool,
+    pub diversity_loss_weight: f32,
+    pub eos_token_id: usize,
+    pub feat_proj_dropout: f32,
+    pub feat_quantizer_dropout: f32,
+    pub feature_projection_input_dim: usize,
+    pub final_dropout: f32,
+    pub hidden_act: Activation,
+    pub hidden_dropout: f32,
+    pub hidden_size: usize,
+    pub initializer_range: f32,
+    pub intermediate_size: usize,
+    pub layer_norm_eps: f64,
+    pub layerdrop: f32,
+    pub left_max_position_embeddings: usize,
+    pub mask_feature_length: usize,
+    pub mask_feature_min_masks: usize,
+    pub mask_feature_prob: f32,
+    pub mask_time_length: usize,
+    pub mask_time_min_masks: usize,
+    pub mask_time_prob: f32,
+    pub max_source_positions: usize,
+    pub num_adapter_layers: usize,
+    pub num_attention_heads: usize,
+    pub num_codevector_groups: usize,
+    pub num_codevectors_per_group: usize,
+    pub num_hidden_layers: usize,
+    pub num_negatives: usize,
+    pub output_hidden_size: usize,
+    pub pad_token_id: usize,
+    pub position_embeddings_type: String,
+    pub proj_codevector_dim: usize,
+    pub right_max_position_embeddings: usize,
+    pub rotary_embedding_base: usize,
+    pub tdnn_dilation: Vec<usize>,
+    pub tdnn_dim: Vec<usize>,
+    pub tdnn_kernel: Vec<usize>,
+    pub torch_dtype: String,
+    pub use_intermediate_ffn_before_adapter: bool,
+    pub use_weighted_layer_sum: bool,
+    pub vocab_size: Option<usize>,
+    pub xvector_output_dim: usize,
+}

src/models/w2v_bert_2_0/mod.rs

@@ -0,0 +1,2 @@
+pub mod config;
+pub mod model;

src/models/w2v_bert_2_0/model.rs

@@ -0,0 +1,571 @@
+use anyhow::Result;
+use candle_core::{D, DType, Device, Tensor};
+use candle_nn::{
+    Activation, Conv1d, Embedding, Init, LayerNorm, Linear, Module, VarBuilder, embedding, linear,
+    linear_b,
+};
+
+use crate::{
+    models::{
+        common::{
+            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},
+    utils::{find_type_files, tensor_utils::masked_fill_zeros},
+};
+
+pub struct Wav2Vec2BertFeatureProjection {
+    layer_norm: LayerNorm,
+    projection: Linear,
+}
+
+impl Wav2Vec2BertFeatureProjection {
+    pub fn new(vb: VarBuilder, config: &W2VBert2_0Config) -> Result<Self> {
+        let layer_norm = get_layer_norm(
+            vb.pp("layer_norm"),
+            config.layer_norm_eps,
+            config.feature_projection_input_dim,
+        )?;
+        let projection = linear(
+            config.feature_projection_input_dim,
+            config.hidden_size,
+            vb.pp("projection"),
+        )?;
+        Ok(Self {
+            layer_norm,
+            projection,
+        })
+    }
+
+    pub fn forward(&self, xs: &Tensor) -> Result<(Tensor, Tensor)> {
+        let norm_xs = self.layer_norm.forward(xs)?;
+        let xs = self.projection.forward(&norm_xs)?;
+        Ok((xs, norm_xs))
+    }
+}
+
+pub struct Wav2Vec2BertSelfAttention {
+    q_proj: Linear,
+    k_proj: Linear,
+    v_proj: Linear,
+    o_proj: Linear,
+    head_dim: usize,
+    num_heads: usize,
+    position_embeddings_type: Option<String>,
+    linear_pos: Option<Linear>,
+    pos_bias_u: Option<Tensor>,
+    pos_bias_v: Option<Tensor>,
+    left_max_position_embeddings: usize,
+    right_max_position_embeddings: usize,
+    distance_embedding: Option<Embedding>,
+}
+
+impl Wav2Vec2BertSelfAttention {
+    pub fn new(
+        vb: VarBuilder,
+        config: &W2VBert2_0Config,
+        is_adapter_attention: bool,
+    ) -> Result<Self> {
+        let hidden_size = if is_adapter_attention {
+            config.hidden_size
+        } else {
+            config.output_hidden_size
+        };
+        let head_dim = hidden_size / config.num_attention_heads;
+        let num_heads = config.num_attention_heads;
+        let left_max_position_embeddings = config.left_max_position_embeddings;
+        let right_max_position_embeddings = config.right_max_position_embeddings;
+        let position_embeddings_type = if !is_adapter_attention {
+            Some(config.position_embeddings_type.clone())
+        } else {
+            None
+        };
+        let (linear_pos, pos_bias_u, pos_bias_v, distance_embedding) =
+            if let Some(pos_type) = &position_embeddings_type {
+                if pos_type.eq("relative") {
+                    let linear_pos = Some(linear_b(
+                        hidden_size,
+                        hidden_size,
+                        false,
+                        vb.pp("linear_pos"),
+                    )?);
+                    let pos_bias_u = Some(vb.get_with_hints(
+                        (config.num_attention_heads, head_dim),
+                        "pos_bias_u",
+                        Init::Const(0.),
+                    )?);
+                    let pos_bias_v = Some(vb.get_with_hints(
+                        (config.num_attention_heads, head_dim),
+                        "pos_bias_v",
+                        Init::Const(0.),
+                    )?);
+                    (linear_pos, pos_bias_u, pos_bias_v, None)
+                } else if pos_type.eq("relative_key") {
+                    let num_positions =
+                        left_max_position_embeddings + right_max_position_embeddings + 1;
+                    let distance_embedding = Some(embedding(
+                        num_positions,
+                        head_dim,
+                        vb.pp("distance_embedding"),
+                    )?);
+                    (None, None, None, distance_embedding)
+                } else {
+                    (None, None, None, None)
+                }
+            } else {
+                (None, None, None, None)
+            };
+        let q_proj = linear_b(hidden_size, hidden_size, true, vb.pp("linear_q"))?;
+        let k_proj = linear_b(hidden_size, hidden_size, true, vb.pp("linear_k"))?;
+        let v_proj = linear_b(hidden_size, hidden_size, true, vb.pp("linear_v"))?;
+        let o_proj = linear_b(hidden_size, hidden_size, true, vb.pp("linear_out"))?;
+        Ok(Self {
+            q_proj,
+            k_proj,
+            v_proj,
+            o_proj,
+            head_dim,
+            num_heads,
+            position_embeddings_type,
+            linear_pos,
+            pos_bias_u,
+            pos_bias_v,
+            left_max_position_embeddings,
+            right_max_position_embeddings,
+            distance_embedding,
+        })
+    }
+
+    pub fn forward(
+        &self,
+        xs: &Tensor,
+        cos: Option<&Tensor>,
+        sin: Option<&Tensor>,
+        attention_mask: Option<&Tensor>,
+    ) -> Result<Tensor> {
+        if let Some(pos_type) = &self.position_embeddings_type
+            && pos_type.eq("rotary")
+            && (cos.is_none() || sin.is_none())
+        {
+            return Err(anyhow::anyhow!(
+                "rotary type position cos and sin can not be none"
+            ));
+        }
+
+        let (b_sz, q_len, _) = xs.dims3()?;
+        let query_states = self.q_proj.forward(xs)?;
+        let key_states = self.k_proj.forward(xs)?;
+        let value_states = self.v_proj.forward(xs)?;
+        let query_states = query_states
+            .reshape((b_sz, q_len, self.num_heads, self.head_dim))?
+            .transpose(1, 2)?;
+        let key_states = key_states
+            .reshape((b_sz, q_len, self.num_heads, self.head_dim))?
+            .transpose(1, 2)?;
+        let value_states = value_states
+            .reshape((b_sz, q_len, self.num_heads, self.head_dim))?
+            .transpose(1, 2)?;
+        let (query_states, key_states) = if let Some(cos) = cos
+            && let Some(sin) = sin
+        {
+            apply_rotary_pos_emb(&query_states, &key_states, cos, sin, false)?
+        } else {
+            (query_states, key_states)
+        };
+        let scale = 1f64 / f64::sqrt(self.head_dim as f64);
+        let attention_mask = if let Some(pos_type) = &self.position_embeddings_type
+            && pos_type.eq("relative_key")
+            && let Some(embed) = &self.distance_embedding
+        {
+            let query_length = query_states.dim(2)?;
+            let key_length = key_states.dim(2)?;
+            let position_ids_l =
+                Tensor::arange(0i64, query_length as i64, xs.device())?.unsqueeze(D::Minus1)?;
+            let position_ids_r =
+                Tensor::arange(0i64, key_length as i64, xs.device())?.unsqueeze(0)?;
+            let distance = position_ids_r.broadcast_sub(&position_ids_l)?;
+            let distance = distance.clamp(
+                -(self.left_max_position_embeddings as i64),
+                self.right_max_position_embeddings as i64,
+            )?;
+            let distance = distance
+                .affine(1.0, self.left_max_position_embeddings as f64)?
+                .to_dtype(candle_core::DType::U32)?;
+            let pos_emb = embed.forward(&distance)?.to_dtype(query_states.dtype())?; // (seq_q, seq_k, dim)
+            let query_ = query_states.unsqueeze(D::Minus2)?; // (b, n_head, seq_q, 1, dim)
+            let pos_emb = pos_emb.unsqueeze(0)?.unsqueeze(0)?; // (1, 1, se_q, seq_k, dim)
+            // torch.einsum("bhld,lrd->bhlr", query, positional_embedding)
+            // (bs, n_head, seq_len, seq_len)
+            let relative_position_attn_weights = query_
+                .broadcast_mul(&pos_emb)?
+                .sum(D::Minus1)?
+                .affine(scale, 0.0)?;
+            if let Some(mask) = attention_mask {
+                // let mask = mask.unsqueeze(1)?.unsqueeze(D::Minus1)?;
+                Some(relative_position_attn_weights.broadcast_add(&mask)?)
+            } else {
+                Some(relative_position_attn_weights)
+            }
+        } else {
+            if let Some(mask) = attention_mask {
+                // let mask = mask.unsqueeze(1)?.unsqueeze(D::Minus1)?;
+                Some(mask.clone())
+            } else {
+                None
+            }
+        };
+
+        let attn_output = eager_attention_forward(
+            &query_states,
+            &key_states,
+            &value_states,
+            None,
+            attention_mask.as_ref(),
+            scale,
+        )?;
+        let attn_output = attn_output.reshape((b_sz, q_len, self.num_heads * self.head_dim))?;
+        let attn_output = attn_output.apply(&self.o_proj)?;
+        Ok(attn_output)
+    }
+}
+
+pub struct Wav2Vec2BertConvolutionModule {
+    layer_norm: LayerNorm,
+    pointwise_conv1: Conv1d,
+    glu: GLU,
+    conv_depthwise_kernel_size: usize,
+    depthwise_conv: Conv1d,
+    depthwise_layer_norm: LayerNorm,
+    act: Activation,
+    pointwise_conv2: Conv1d,
+}
+
+impl Wav2Vec2BertConvolutionModule {
+    pub fn new(vb: VarBuilder, config: &W2VBert2_0Config) -> Result<Self> {
+        let layer_norm = get_layer_norm(
+            vb.pp("layer_norm"),
+            config.layer_norm_eps,
+            config.hidden_size,
+        )?;
+        let pointwise_conv1 = get_conv1d(
+            vb.pp("pointwise_conv1"),
+            config.hidden_size,
+            2 * config.hidden_size,
+            1,
+            0,
+            1,
+            1,
+            1,
+            false,
+        )?;
+        let glu = GLU::new(1)?;
+        let conv_depthwise_kernel_size = config.conv_depthwise_kernel_size;
+        let depthwise_conv = get_conv1d(
+            vb.pp("depthwise_conv"),
+            config.hidden_size,
+            config.hidden_size,
+            conv_depthwise_kernel_size,
+            0,
+            1,
+            1,
+            config.hidden_size,
+            false,
+        )?;
+        let depthwise_layer_norm = get_layer_norm(
+            vb.pp("depthwise_layer_norm"),
+            config.layer_norm_eps,
+            config.hidden_size,
+        )?;
+        let pointwise_conv2 = get_conv1d(
+            vb.pp("pointwise_conv2"),
+            config.hidden_size,
+            config.hidden_size,
+            1,
+            0,
+            1,
+            1,
+            1,
+            false,
+        )?;
+        Ok(Self {
+            layer_norm,
+            pointwise_conv1,
+            glu,
+            conv_depthwise_kernel_size,
+            depthwise_conv,
+            depthwise_layer_norm,
+            act: config.hidden_act,
+            pointwise_conv2,
+        })
+    }
+
+    pub fn forward(&self, xs: &Tensor, mask: Option<&Tensor>) -> Result<Tensor> {
+        let mut xs = self.layer_norm.forward(xs)?;
+        if let Some(mask) = mask {
+            xs = masked_fill_zeros(&xs, mask)?;
+        }
+        let xs = xs.transpose(1, 2)?;
+        // (batch, 2*channel, dim)
+        let xs = self.pointwise_conv1.forward(&xs)?;
+        // (batch, channel, dim)
+        let xs = self.glu.forward(&xs)?;
+        let xs = xs.pad_with_zeros(D::Minus1, self.conv_depthwise_kernel_size - 1, 0)?;
+        let xs = self.depthwise_conv.forward(&xs)?;
+        let xs = self
+            .depthwise_layer_norm
+            .forward(&xs.transpose(1, 2)?)?
+            .transpose(1, 2)?;
+        let xs = xs.apply(&self.act)?;
+        let xs = self.pointwise_conv2.forward(&xs)?;
+        let xs = xs.transpose(1, 2)?;
+        Ok(xs)
+    }
+}
+
+pub struct Wav2Vec2BertEncoderLayer {
+    ffn1_layer_norm: LayerNorm,
+    ffn1: TwoLinearMLP,
+    self_attn_layer_norm: LayerNorm,
+    self_attn: Wav2Vec2BertSelfAttention,
+    conv_module: Wav2Vec2BertConvolutionModule,
+    ffn2_layer_norm: LayerNorm,
+    ffn2: TwoLinearMLP,
+    final_layer_norm: LayerNorm,
+}
+
+impl Wav2Vec2BertEncoderLayer {
+    pub fn new(vb: VarBuilder, config: &W2VBert2_0Config) -> Result<Self> {
+        let ffn1_layer_norm = get_layer_norm(
+            vb.pp("ffn1_layer_norm"),
+            config.layer_norm_eps,
+            config.hidden_size,
+        )?;
+        let ffn1 = TwoLinearMLP::new(
+            vb.pp("ffn1"),
+            config.hidden_size,
+            config.intermediate_size,
+            config.hidden_size,
+            config.hidden_act,
+            true,
+            "intermediate_dense",
+            "output_dense",
+        )?;
+        let self_attn_layer_norm = get_layer_norm(
+            vb.pp("self_attn_layer_norm"),
+            config.layer_norm_eps,
+            config.hidden_size,
+        )?;
+        let self_attn = Wav2Vec2BertSelfAttention::new(vb.pp("self_attn"), config, false)?;
+        let conv_module = Wav2Vec2BertConvolutionModule::new(vb.pp("conv_module"), config)?;
+        let ffn2_layer_norm = get_layer_norm(
+            vb.pp("ffn2_layer_norm"),
+            config.layer_norm_eps,
+            config.hidden_size,
+        )?;
+        let ffn2 = TwoLinearMLP::new(
+            vb.pp("ffn2"),
+            config.hidden_size,
+            config.intermediate_size,
+            config.hidden_size,
+            config.hidden_act,
+            true,
+            "intermediate_dense",
+            "output_dense",
+        )?;
+        let final_layer_norm = get_layer_norm(
+            vb.pp("final_layer_norm"),
+            config.layer_norm_eps,
+            config.hidden_size,
+        )?;
+        Ok(Self {
+            ffn1_layer_norm,
+            ffn1,
+            self_attn_layer_norm,
+            self_attn,
+            conv_module,
+            ffn2_layer_norm,
+            ffn2,
+            final_layer_norm,
+        })
+    }
+
+    pub fn forward(
+        &self,
+        xs: &Tensor,
+        cos: Option<&Tensor>,
+        sin: Option<&Tensor>,
+        attention_mask: Option<&Tensor>,
+        conv_attention_mask: Option<&Tensor>,
+    ) -> Result<Tensor> {
+        let residual = xs.clone();
+        let xs = self.ffn1_layer_norm.forward(xs)?;
+        let xs = self.ffn1.forward(&xs)?;
+        let residual = xs.affine(0.5, 0.0)?.add(&residual)?;
+        let xs = self.self_attn_layer_norm.forward(&residual)?;
+        let xs = self.self_attn.forward(&xs, cos, sin, attention_mask)?;
+        let residual = xs.add(&residual)?;
+        let xs = self.conv_module.forward(&residual, conv_attention_mask)?;
+        let residual = xs.add(&residual)?;
+        let xs = self.ffn2_layer_norm.forward(&residual)?;
+        let xs = self.ffn2.forward(&xs)?;
+        let xs = xs.affine(0.5, 0.0)?.add(&residual)?;
+        let xs = self.final_layer_norm.forward(&xs)?;
+        Ok(xs)
+    }
+}
+
+pub struct ModelOutput {
+    pub last_hidden_state: Tensor,
+    pub specify_layer_id_hidden_state: Option<Tensor>,
+    pub hidden_states: Option<Vec<Tensor>>,
+}
+
+pub struct Wav2Vec2BertEncoder {
+    embed_positions: Option<RoPE>,
+    layers: Vec<Wav2Vec2BertEncoderLayer>,
+}
+
+impl Wav2Vec2BertEncoder {
+    pub fn new(vb: VarBuilder, config: &W2VBert2_0Config) -> Result<Self> {
+        let embed_positions = if config.position_embeddings_type.eq("rotary") {
+            let dim = config.hidden_size / config.num_attention_heads;
+            let embed_positions = RoPE::new(dim, 10000.0, vb.device())?;
+            Some(embed_positions)
+        } else {
+            None
+        };
+        let vb_layers = vb.pp("layers");
+        let mut layers = vec![];
+        for i in 0..config.num_hidden_layers {
+            let layer = Wav2Vec2BertEncoderLayer::new(vb_layers.pp(i), config)?;
+            layers.push(layer);
+        }
+        Ok(Self {
+            embed_positions,
+            layers,
+        })
+    }
+
+    pub fn forward(
+        &self,
+        xs: &Tensor,
+        attention_mask: Option<&Tensor>,
+        layer_id: Option<usize>,
+        output_hidden_states: bool,
+    ) -> Result<ModelOutput> {
+        // xs: (bs, seq_len ,dim)
+        // attention_mask: Some: (bs, seq_len)
+        let (_, seq_len, _) = xs.dims3()?;
+        let conv_attention_mask = attention_mask;
+        let (mut xs, attention_mask) = if let Some(mask) = attention_mask {
+            let xs = masked_fill_zeros(xs, mask)?;
+            // (bs, 1, 1, seq_len)
+            let attention_mask = mask.unsqueeze(1)?.unsqueeze(1)?;
+            let neg_inf_t = attention_mask
+                .zeros_like()?
+                .to_dtype(xs.dtype())?
+                .affine(1.0, f64::NEG_INFINITY)?;
+            let attention_mask_f = attention_mask.to_dtype(xs.dtype())?;
+            let attention_mask = attention_mask
+                .where_cond(&attention_mask_f, &neg_inf_t)?
+                .to_dtype(xs.dtype())?
+                .affine(1.0, -1.0)?;
+            (xs, Some(attention_mask))
+        } else {
+            (xs.clone(), None)
+        };
+        let (cos, sin) = if let Some(embed_posi) = &self.embed_positions {
+            let (cos, sin) = embed_posi.forward(0, seq_len, xs.device())?;
+            (Some(cos), Some(sin))
+        } else {
+            (None, None)
+        };
+        let mut hidden_states: Vec<Tensor> = vec![];
+        let mut specify_layer_id_hidden_state = None;
+
+        for (i, layer) in (&self.layers).iter().enumerate() {
+            if output_hidden_states {
+                hidden_states.push(xs.clone());
+            } 
+            if let Some(id) = layer_id
+                && id == i
+            {
+                specify_layer_id_hidden_state = Some(xs.clone());
+            }
+            xs = layer.forward(
+                &xs,
+                cos.as_ref(),
+                sin.as_ref(),
+                attention_mask.as_ref(),
+                conv_attention_mask,
+            )?;                      
+        }
+        let hidden_states = if hidden_states.len() > 0 {
+            Some(hidden_states)
+        } else {
+            None
+        };
+        Ok(ModelOutput {
+            last_hidden_state: xs,
+            specify_layer_id_hidden_state,
+            hidden_states,
+        })
+    }
+}
+
+pub struct W2VBert2_0Model {
+    config: W2VBert2_0Config,
+    feature_projection: Wav2Vec2BertFeatureProjection,
+    masked_spec_embed: Option<Tensor>,
+    encoder: Wav2Vec2BertEncoder,
+    // config.add_adapter is false, adapter is None, Wav2Vec2BertAdapter not complish
+    // adapter: Option<Wav2Vec2BertAdapter>,
+    // config.use_intermediate_ffn_before_adapter is false, intermediate_ffn is None
+    // intermediate_ffn: Option<Wav2Vec2BertFeedForward>,
+}
+
+impl W2VBert2_0Model {
+    pub fn init(path: &str, device: &Device, dtype: DType) -> Result<Self> {
+        let config_path = path.to_string() + "/config.json";
+        let config: W2VBert2_0Config = serde_json::from_slice(&std::fs::read(config_path)?)?;
+        let model_list = find_type_files(path, "safetensors")?;
+        let vb = unsafe { VarBuilder::from_mmaped_safetensors(&model_list, dtype, device)? };
+        W2VBert2_0Model::new(vb, &config)
+    }
+
+    pub fn new(vb: VarBuilder, config: &W2VBert2_0Config) -> Result<Self> {
+        let feature_projection =
+            Wav2Vec2BertFeatureProjection::new(vb.pp("feature_projection"), config)?;
+        let masked_spec_embed = if config.mask_time_prob > 0.0 || config.mask_time_prob > 0.0 {
+            Some(
+                vb.get_with_hints(config.hidden_size, "masked_spec_embed", Init::Uniform {
+                    lo: 0.0,
+                    up: 1.0,
+                })?,
+            )
+        } else {
+            None
+        };
+        let encoder = Wav2Vec2BertEncoder::new(vb.pp("encoder"), config)?;
+        Ok(Self {
+            config: config.clone(),
+            feature_projection,
+            masked_spec_embed,
+            encoder,
+        })
+    }
+
+    pub fn forward(
+        &self,
+        xs: &Tensor,
+        attention_mask: Option<&Tensor>,
+        layer_id: Option<usize>,
+        output_hidden_states: bool,
+    ) -> Result<ModelOutput> {
+        let (xs, _) = self.feature_projection.forward(xs)?;
+        self.encoder
+            .forward(&xs, attention_mask, layer_id, output_hidden_states)
+    }
+}

src/position_embed/rope.rs

@@ -1,4 +1,4 @@
-use anyhow::Result;
+use anyhow::{Result, anyhow};
 use candle_core::{D, DType, Device, IndexOp, Tensor};
 use candle_transformers::models::deepseek2::SplitOp;
 
@@ -158,6 +158,68 @@ pub fn glm_asr_apply_rotary_pos_emb(
     Ok((q_embed, k_embed))
 }
 
+pub fn roformer_rotate(x: &Tensor) -> Result<Tensor> {
+    let dims = x.dims();
+    let last_dim = dims
+        .last()
+        .ok_or(anyhow!("Input tensor must have at least one dimension"))?;
+    if last_dim % 2 != 0 {
+        return Err(anyhow!(
+            "Last dimension size must be even, got {}",
+            last_dim
+        ));
+    }
+    let new_dims: Vec<usize> = dims[..dims.len() - 1]
+        .iter()
+        .copied()
+        .chain([last_dim / 2, 2])
+        .collect();
+    let x_reshape = x.reshape(new_dims)?;
+    let x_chunks = x_reshape.chunk(2, D::Minus1)?;
+    let x1 = &x_chunks[0];
+    let x2 = &x_chunks[1];
+    // let x1 = x_reshape.narrow(D::Minus1, 0, 1)?;
+    // let x2 = x_reshape.narrow(D::Minus1, 1, 1)?;
+    let x2_neg = x2.affine(-1.0, 0.0)?;
+    let rotate_x = Tensor::cat(&[&x2_neg, x1], D::Minus1)?;
+    Ok(rotate_x.flatten(D::Minus2, D::Minus1)?)
+}
+
+pub fn apply_rotary_pos_emb_roformer(
+    q: &Tensor,
+    k: &Tensor,
+    cos: &Tensor,
+    sin: &Tensor,
+    tof32: bool,
+) -> Result<(Tensor, Tensor)> {
+    let mut cos = cos.clone();
+    let mut sin = sin.clone();
+    if cos.rank() == 2 {
+        // (seq_len, head_dim) -> (1, 1, seq_len, head_dim)
+        cos = cos.unsqueeze(0)?.unsqueeze(0)?;
+        sin = sin.unsqueeze(0)?.unsqueeze(0)?;
+    }
+    if cos.rank() == 3 {
+        // (bs, seq_len, head_dim) -> (bs, 1, seq_len, head_dim)
+        cos = cos.unsqueeze(1)?;
+        sin = sin.unsqueeze(1)?;
+    }
+    let orig_dtype = q.dtype();
+    let q = if tof32 { &q.to_dtype(DType::F32)? } else { q };
+    let k = if tof32 { &k.to_dtype(DType::F32)? } else { k };
+    let cos = cos.to_dtype(q.dtype())?;
+    let sin = sin.to_dtype(q.dtype())?;
+    let q_embed = q
+        .broadcast_mul(&cos)?
+        .add(&roformer_rotate(q)?.broadcast_mul(&sin)?)?
+        .to_dtype(orig_dtype)?;
+    let k_embed = k
+        .broadcast_mul(&cos)?
+        .add(&roformer_rotate(k)?.broadcast_mul(&sin)?)?
+        .to_dtype(orig_dtype)?;
+    Ok((q_embed, k_embed))
+}
+
 #[derive(Debug, Clone)]
 pub struct Qwen2_5VLTextRotaryEmbedding {
     inv_freq: Vec<f32>,
@@ -278,6 +340,70 @@ impl Qwen3VLTextRotaryEmbedding {
         }
         Ok(freqs_t)
     }
+
+    pub fn apply_interleaved_mrope_asr(
+        &self,
+        freqs: &Tensor,
+        mrope_section: Vec<usize>,
+    ) -> Result<Tensor> {
+        let mut freqs_t = freqs.i(0)?.contiguous()?; //(3, bs, seq_len, head_dim //2) -> (bs, seq_len, head_dim //2)
+
+        // for dim in 1..3 {
+        for (dim, offset) in (1..3).enumerate() {
+            let dim = dim +1;
+            let length = mrope_section[dim];
+            let idx = Tensor::arange_step(offset as u32, length as u32, 3, freqs.device())?;
+            let src = freqs.i(dim)?.contiguous()?; // (bs, seq_len, head_dim //2)
+            let src = src.index_select(&idx, D::Minus1)?.contiguous()?;
+            let idx = idx
+                .unsqueeze(0)?
+                .unsqueeze(0)?
+                .broadcast_as(src.shape())?
+                .contiguous()?;
+            freqs_t = freqs_t.scatter(&idx, &src, D::Minus1)?;
+        }
+        Ok(freqs_t)
+    }
+
+    pub fn forward_asr(
+        &self,
+        position_ids: &Tensor,
+        dtype: DType,
+        mrope_section: Vec<usize>,
+    ) -> Result<(Tensor, Tensor)> {
+        // position_ids shape: (3, bs, position) -> (3, bs, 1, position)
+        let position_ids = if position_ids.rank() == 2 {
+            let (bs, len) = position_ids.dims2()?;
+            position_ids.unsqueeze(0)?.expand((3, bs, len))?
+        } else {
+            position_ids.clone()
+        };
+        let position_ids_expanded = position_ids
+            .unsqueeze(D::Minus2)?
+            .to_dtype(DType::F32)?
+            .contiguous()?;
+        // inv_freq Vec<f32> -> Tensor(1, 1, head_dim / 2, 1) -> (3, bs, head_dim / 2, 1)
+        let inv_freq_expanded = Tensor::from_vec(
+            self.inv_freq.clone(),
+            (1, 1, self.inv_freq.len(), 1),
+            position_ids.device(),
+        )?
+        .broadcast_as((3, position_ids.dim(1)?, self.inv_freq.len(), 1))?
+        .to_dtype(DType::F32)?
+        .contiguous()?;
+
+        // (3, bs, head_dim / 2, 1) matmul (3, bs, 1, position)
+        //    -> (3, bs, head_dim / 2, seq_len) -> (3, bs, seq_len, head_dim / 2)
+        let freqs = inv_freq_expanded
+            .matmul(&position_ids_expanded)?
+            .transpose(2, 3)?;
+        let freqs = self.apply_interleaved_mrope_asr(&freqs, mrope_section)?;
+        let emb = Tensor::cat(&[&freqs, &freqs], D::Minus1)?.contiguous()?;
+        let cos = emb.cos()?;
+        let sin = emb.sin()?;
+        Ok((cos.to_dtype(dtype)?, sin.to_dtype(dtype)?))
+    }
+
     pub fn forward(
         &self,
         position_ids: &Tensor,

src/tokenizer/mod.rs

@@ -1,6 +1,10 @@
 use anyhow::{Ok, Result, anyhow};
 use candle_core::{Device, Tensor};
-use tokenizers::Tokenizer;
+use serde_json::Value;
+use tokenizers::{
+    AddedToken, Tokenizer, decoders::byte_level::ByteLevel as ByteLevelDecoder, models::bpe::BPE,
+    pre_tokenizers::byte_level::ByteLevel,
+};
 
 pub struct TokenizerModel {
     pub tokenizer: Tokenizer,
@@ -14,12 +18,70 @@ impl TokenizerModel {
             "model path file not exists"
         );
         let tokenizer_file = path.clone() + "/tokenizer.json";
-        assert!(
+        let tokenizer = if std::path::Path::new(&tokenizer_file).exists() {
-            std::path::Path::new(&tokenizer_file).exists(),
+            Tokenizer::from_file(tokenizer_file)
-            "tokenizer.json not exists in model path"
+                .map_err(|e| anyhow!(format!("tokenizer from file error{}", e)))?
-        );
+        } else {
-        let tokenizer = Tokenizer::from_file(tokenizer_file)
+            // 如果不存在 tokenizer.json，尝试使用 vocab.json 和 merges.txt
-            .map_err(|e| anyhow!(format!("tokenizer from file error{}", e)))?;
+            let vocab_file = path.clone() + "/vocab.json";
+            let merges_file = path.clone() + "/merges.txt";
+            let config_file = path.clone() + "/tokenizer_config.json";
+
+            if !std::path::Path::new(&vocab_file).exists() {
+                return Err(anyhow!(
+                    "Neither tokenizer.json nor vocab.json found in model path"
+                ));
+            }
+
+            if !std::path::Path::new(&merges_file).exists() {
+                return Err(anyhow!(
+                    "Neither tokenizer.json nor merges.txt found in model path"
+                ));
+            }
+            // 创建 BPE 模型
+            let bpe = BPE::from_file(&vocab_file, &merges_file)
+                .build()
+                .map_err(|e| anyhow!(format!("failed to build BPE tokenizer: {}", e)))?;
+
+            // 创建分词器
+            let mut tokenizer = Tokenizer::new(bpe);
+            // 添加字节级预分词器，这会处理换行符等特殊字符
+            let byte_level_pre_tokenizer = ByteLevel::new(false, true, false);
+            tokenizer.with_pre_tokenizer(Some(byte_level_pre_tokenizer));
+            tokenizer.with_decoder(Some(ByteLevelDecoder::default()));
+            if std::path::Path::new(&config_file).exists() {
+                let config_content = std::fs::read_to_string(&config_file)?;
+                let config: Value = serde_json::from_str(&config_content)?;
+                if let Some(added_tokens_decoder) = config.get("added_tokens_decoder") {
+                    let mut special_tokens = Vec::new();
+
+                    if let Value::Object(tokens_map) = added_tokens_decoder {
+                        for (_, token_info) in tokens_map {
+                            if let Value::Object(token_obj) = token_info {
+                                if let Some(content_val) = token_obj.get("content") {
+                                    if let Some(content) = content_val.as_str() {
+                                        let special = token_obj
+                                            .get("special")
+                                            .and_then(|v| v.as_bool())
+                                            .unwrap_or(false);
+
+                                        let added_token =
+                                            AddedToken::from(content.to_string(), special);
+                                        special_tokens.push(added_token);
+                                    }
+                                }
+                            }
+                        }
+                    }
+
+                    // 添加所有特殊标记
+                    if !special_tokens.is_empty() {
+                        tokenizer.add_special_tokens(&special_tokens);
+                    }
+                }
+            }
+            tokenizer
+        };
         Ok(Self { tokenizer })
     }
 

src/utils/audio_utils.rs

@@ -32,7 +32,7 @@ use symphonia::core::meta::MetadataOptions;
 use symphonia::core::probe::Hint;
 
 use crate::utils::get_default_save_dir;
-use crate::utils::tensor_utils::{linspace, pad_replicate_last_dim};
+use crate::utils::tensor_utils::{linspace, log10, pad_reflect_last_dim, pad_replicate_last_dim, split_tensor};
 
 // 重采样方法枚举
 #[derive(Debug, Clone, Copy)]
@@ -569,6 +569,11 @@ pub fn load_audio_use_symphonia(audio_vec: Vec<u8>, device: &Device) -> Result<(
     Ok((audio_tensor, sample_rate as usize))
 }
 
+pub fn load_audio(path: &str, device: &Device) -> Result<(Tensor, usize)> {
+    let audio_vec = get_audio_bytes_vec(path)?;
+    load_audio_use_symphonia(audio_vec, device)
+}
+
 pub fn load_audio_with_resample(
     path: &str,
     device: &Device,
@@ -945,18 +950,44 @@ pub fn load_and_resample_audio_ffmpeg(
 //     Ok(audio)
 // }
 
+// pub fn create_hann_window(window_size: usize, dtype: DType, device: &Device) -> Result<Tensor> {
+//     let n = window_size as f64;
+//     let window: Vec<f32> = (0..window_size)
+//         .map(|i| {
+//             let i_f64 = i as f64;
+//             let val = 0.5 * (1.0 - (2.0 * PI * i_f64 / n).cos());
+//             val as f32
+//         })
+//         .collect();
+//     Ok(Tensor::from_vec(window, window_size, device)?.to_dtype(dtype)?)
+// }
+
 pub fn create_hann_window(window_size: usize, dtype: DType, device: &Device) -> Result<Tensor> {
-    let n = window_size as f64;
+    if window_size < 1 {
-    let window: Vec<f32> = (0..window_size)
+        return Err(anyhow::anyhow!("window_size must bigger than 0"));
+    }
+    if window_size == 1 {
+        return Ok(Tensor::new(1.0f32, device)?.to_dtype(dtype)?);
+    }
+    let n = window_size as f64 - 1.0;
+    let start = 1_i64 - window_size as i64;
+    let end = window_size as i64;
+    let window: Vec<f32> = (start..end)
+        .step_by(2)
         .map(|i| {
             let i_f64 = i as f64;
-            let val = 0.5 * (1.0 - (2.0 * PI * i_f64 / n).cos());
+            let val = 0.5 + 0.5 * (PI * i_f64 / n).cos();
             val as f32
         })
         .collect();
     Ok(Tensor::from_vec(window, window_size, device)?.to_dtype(dtype)?)
 }
 
+pub fn create_povey_window(window_size: usize, dtype: DType, device: &Device) -> Result<Tensor> {
+    let window = create_hann_window(window_size, dtype, device)?;
+    Ok(window.powf(0.85)?)
+}
+
 pub fn crate_hamming_window(
     window_size: usize,
     periodic: bool,
@@ -1166,6 +1197,22 @@ pub fn apply_stft(waveform: &Tensor) -> Result<Tensor> {
     Ok(magnitudes)
 }
 
+pub fn torch_stft(
+    waveform: &Tensor,
+    n_fft: usize,
+    hop_length: usize,
+    window: &Tensor,
+) -> Result<Tensor> {
+    // waveform: already padding
+    // (bs, n_frames, n_fft)
+    let frames = extract_frames(&waveform, n_fft, hop_length)?;
+    // 应用汉明窗口
+    let result = frames.broadcast_mul(window)?;
+    // 傅立叶变换
+    let magnitudes = apply_stft(&result)?;
+    Ok(magnitudes)
+}
+
 pub fn kaldi_fbank(
     waveform: &Tensor,
     mel_energies: &Tensor,
@@ -1489,3 +1536,92 @@ pub fn kaldi_get_mel_banks(
 
     Ok((bins_tensor, center_freqs))
 }
+
+pub fn spectrogram(
+    waveform: &Tensor,
+    window: &Tensor,
+    frame_length: usize,
+    hop_length: usize,
+    fft_length: usize,
+    power: Option<f32>,
+    center: bool,
+    preemphasis: f64,
+    mel_filters: Option<&Tensor>,
+    log_mel: Option<&str>,
+    mel_floor: f32,
+    remove_dc_offset: bool,
+) -> Result<Tensor> {
+    let waveform = if center {
+        let pad = frame_length / 2;
+        pad_reflect_last_dim(waveform, (pad, pad))?
+    } else {
+        waveform.clone()
+    };
+    let mut frames = extract_frames(&waveform, frame_length, hop_length)?;
+    if remove_dc_offset {
+        let row_means = frames.mean_keepdim(D::Minus1)?;
+        frames = frames.broadcast_sub(&row_means)?;
+    }
+
+    if preemphasis != 0.0 {
+        let buffer_0 = frames
+            .i((.., .., 0))?
+            .affine(1.0 - preemphasis, 0.0)?
+            .unsqueeze(D::Minus1)?;
+        let buffer_ = frames.i((.., .., 1..))?.sub(
+            &frames
+                .i((.., .., 0..frame_length - 1))?
+                .affine(preemphasis, 0.0)?,
+        )?;
+        frames = Tensor::cat(&[buffer_0, buffer_], D::Minus1)?;
+    }
+    let mut frames = frames.broadcast_mul(&window)?;
+    let pad_len = fft_length - frame_length;
+    if pad_len > 0 {
+        // (bs, nframes, frame_length) -> (bs, nframes, fft_length)
+        frames = frames.pad_with_zeros(D::Minus1, 0, pad_len)?;
+    }
+    let mut spectrogram = apply_stft(&frames)?; // stft已经做了pow(2.0)
+    spectrogram = spectrogram.transpose(D::Minus1, D::Minus2)?;
+    if let Some(mel_filters) = mel_filters {
+        let spect = mel_filters.t()?.broadcast_matmul(&spectrogram)?;
+        spectrogram = spect.maximum(
+            &Tensor::new(mel_floor, spect.device())?
+                .to_dtype(spect.dtype())?
+                .broadcast_as(spect.shape())?,
+        )?;
+    }
+    if let Some(_) = power
+        && let Some(log_mel) = log_mel
+    {
+        if log_mel == "log" {
+            spectrogram = spectrogram.log()?;
+        } else if log_mel == "log10" {
+            spectrogram = log10(&spectrogram)?;
+        } else {
+            return Err(anyhow!(
+                "dB not completed or Unknown log_mel option ".to_string()
+            ));
+        }
+    }
+    Ok(spectrogram)
+}
+
+pub fn split_audio_into_chunks(wav: &Tensor, sr: usize, max_chunk_sec: f32) -> Result<Vec<Tensor>> {
+    // wav: (1, len)
+    let total_len = wav.dim(1)?;
+    let total_sec = total_len as f32 / sr as f32;
+    let mut wavs = vec![];
+    if total_sec <= max_chunk_sec {
+        wavs.push(wav.clone());
+    } else {
+        let max_len = (max_chunk_sec * sr as f32).round() as usize;
+        let split_len = total_len / max_len;
+        let mut splits = vec![max_len; split_len];
+        let remain_len = total_len % max_len;
+        splits.push(remain_len);
+        let split_wav = split_tensor(wav, &splits, 1)?;
+        wavs.extend_from_slice(&split_wav);
+    }
+    Ok(wavs)
+}

src/utils/mod.rs

@@ -3,7 +3,8 @@ pub mod img_utils;
 pub mod tensor_utils;
 pub mod video_utils;
 
-use std::{fs, path::PathBuf, process::Command};
+use std::io::Read;
+use std::{collections::HashMap, fs, path::PathBuf, process::Command, time::Duration};
 
 use aha_openai_dive::v1::resources::{
     chat::{
@@ -14,10 +15,18 @@ use aha_openai_dive::v1::resources::{
     },
     shared::{FinishReason, Usage},
 };
-use anyhow::Result;
+use anyhow::{Result, anyhow};
-use candle_core::{DType, Device};
+use byteorder::{LittleEndian, ReadBytesExt};
+use candle_core::{
+    Context, DType, Device, Shape, Tensor,
+    pickle::{Object, Stack, TensorInfo, read_all_with_key},
+};
+use candle_nn::VarBuilder;
 use candle_transformers::generation::{LogitsProcessor, Sampling};
 use dirs::home_dir;
+use half::{bf16, f16, slice::HalfFloatSliceExt};
+use modelscope::ModelScope;
+use tokio::time::sleep;
 
 pub fn get_device(device: Option<&Device>) -> Device {
     match device {
@@ -128,6 +137,228 @@ pub fn find_type_files(path: &str, extension_type: &str) -> Result<Vec<String>>
     Ok(files)
 }
 
+pub fn get_vb_model_path(
+    model_path: String,
+    dtype: DType,
+    device: Device,
+    key: Option<&'_ str>,
+) -> Result<VarBuilder<'_>> {
+    let mut dict_to_hashmap = HashMap::new();
+    let dict = read_all_with_key(&model_path, key)?;
+    for (k, v) in dict {
+        dict_to_hashmap.insert(k, v);
+    }
+    let vb = VarBuilder::from_tensors(dict_to_hashmap, dtype, &device);
+    Ok(vb)
+}
+
+pub fn get_vb_extension(
+    path: String,
+    extension_type: String,
+    dtype: DType,
+    device: Device,
+    key: Option<&'_ str>,
+) -> Result<VarBuilder<'_>> {
+    let model_list = find_type_files(&path, &extension_type)?;
+    let mut dict_to_hashmap = HashMap::new();
+    for m in model_list {
+        let dict = read_all_with_key(m, key)?;
+        for (k, v) in dict {
+            dict_to_hashmap.insert(k, v);
+        }
+    }
+    let vb = VarBuilder::from_tensors(dict_to_hashmap, dtype, &device);
+    Ok(vb)
+}
+
+pub fn crate_tensor_from_reader<R: std::io::Read>(
+    shape: Shape,
+    dtype: DType,
+    reader: &mut R,
+) -> Result<Tensor> {
+    let elem_count = shape.elem_count();
+    match dtype {
+        DType::BF16 => {
+            let mut data_t = vec![bf16::ZERO; elem_count];
+            reader.read_u16_into::<LittleEndian>(data_t.reinterpret_cast_mut())?;
+            Ok(Tensor::from_vec(data_t, shape, &Device::Cpu)?)
+        }
+        DType::F16 => {
+            let mut data_t = vec![f16::ZERO; elem_count];
+            reader.read_u16_into::<LittleEndian>(data_t.reinterpret_cast_mut())?;
+            Ok(Tensor::from_vec(data_t, shape, &Device::Cpu)?)
+        }
+        DType::F32 => {
+            let mut data_t = vec![0f32; elem_count];
+            reader.read_f32_into::<LittleEndian>(&mut data_t)?;
+            Ok(Tensor::from_vec(data_t, shape, &Device::Cpu)?)
+        }
+        DType::F64 => {
+            let mut data_t = vec![0f64; elem_count];
+            reader.read_f64_into::<LittleEndian>(&mut data_t)?;
+            Ok(Tensor::from_vec(data_t, shape, &Device::Cpu)?)
+        }
+        DType::U8 => {
+            let mut data_t = vec![0u8; elem_count];
+            reader.read_exact(&mut data_t)?;
+            Ok(Tensor::from_vec(data_t, shape, &Device::Cpu)?)
+        }
+        DType::U32 => {
+            let mut data_t = vec![0u32; elem_count];
+            reader.read_u32_into::<LittleEndian>(&mut data_t)?;
+            Ok(Tensor::from_vec(data_t, shape, &Device::Cpu)?)
+        }
+        DType::I64 => {
+            let mut data_t = vec![0i64; elem_count];
+            reader.read_i64_into::<LittleEndian>(&mut data_t)?;
+            Ok(Tensor::from_vec(data_t, shape, &Device::Cpu)?)
+        }
+    }
+}
+
+pub fn read_pth_tensor_info_cycle<P: AsRef<std::path::Path>>(
+    path: P,
+    key: Option<&str>,
+) -> Result<Vec<(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)?;
+    let zip_file_names = zip
+        .file_names()
+        .map(|f| f.to_string())
+        .collect::<Vec<String>>();
+
+    let mut tensor_infos = vec![];
+    for file_name in zip_file_names.iter() {
+        if !file_name.ends_with("data.pkl") {
+            continue;
+        }
+        let dir_name = std::path::PathBuf::from(file_name.strip_suffix(".pkl").context("no .pkl")?);
+        let reader = zip.by_name(file_name)?;
+        let mut reader = std::io::BufReader::new(reader);
+        let mut stack = Stack::empty();
+        stack.read_loop(&mut reader)?;
+        let obj = stack.finalize()?;
+
+        let obj = match obj {
+            Object::Build { callable, args } => match *callable {
+                Object::Reduce { callable, args: _ } => match *callable {
+                    Object::Class {
+                        module_name,
+                        class_name,
+                    } if module_name == "__torch__" && class_name == "Module" => *args,
+                    _ => continue,
+                },
+                _ => continue,
+            },
+            obj => obj,
+        };
+
+        // If key is provided, then we need to extract the state_dict from the object.
+        let obj = if let Some(key) = key {
+            let multi_key: Vec<&str> = key.split(".").collect();
+            if multi_key.len() > 1 {
+                let mut current_obj = obj;
+                for k in multi_key.iter() {
+                    if let Object::Dict(key_values) = current_obj {
+                        current_obj = key_values
+                            .into_iter()
+                            .find(|(key_obj, _)| *key_obj == Object::Unicode(k.to_string()))
+                            .map(|(_, v)| v)
+                            .ok_or_else(|| anyhow!(format!("key '{}' not found", k)))?;
+                    } else {
+                        return Err(anyhow!(format!(
+                            "Expected dictionary at key '{}', but found other type",
+                            k
+                        )));
+                    }
+                }
+                current_obj
+            } else {
+                if let Object::Dict(key_values) = obj {
+                    key_values
+                        .into_iter()
+                        .find(|(k, _)| *k == Object::Unicode(key.to_owned()))
+                        .map(|(_, v)| v)
+                        .ok_or_else(|| anyhow!(format!("key {key} not found")))?
+                } else {
+                    obj
+                }
+            }
+        } else {
+            obj
+        };
+
+        // If the object is a dict, then we can extract the tensor info from it.
+        // NOTE: We are assuming that the `obj` is state_dict by this stage.
+        if let Object::Dict(key_values) = obj {
+            for (name, value) in key_values.into_iter() {
+                match value.into_tensor_info(name, &dir_name) {
+                    Ok(Some(tensor_info)) => tensor_infos.push(tensor_info),
+                    Ok(None) => {}
+                    Err(err) => eprintln!("skipping: {err:?}"),
+                }
+            }
+        }
+    }
+    let tensor_infos: HashMap<String, TensorInfo> = tensor_infos
+        .into_iter()
+        .map(|ti| (ti.name.to_string(), ti))
+        .collect();
+
+    let tensor_names = tensor_infos.keys();
+    let mut tensors = Vec::with_capacity(tensor_names.len());
+    for name in tensor_names {
+        let _ = match tensor_infos.get(name) {
+            None => {}
+            Some(tensor_info) => {
+                let zip_reader = std::io::BufReader::new(std::fs::File::open(&path)?);
+                let mut zip = zip::ZipArchive::new(zip_reader)?;
+                let mut reader = zip.by_name(&tensor_info.path)?;
+                let is_fortran_contiguous = tensor_info.layout.is_fortran_contiguous();
+                let rank = tensor_info.layout.shape().rank();
+
+                // Reading the data is a bit tricky as it can be strided, for now only support the basic
+                // case and when the tensor is fortran contiguous.
+                if !tensor_info.layout.is_contiguous() && !is_fortran_contiguous {
+                    return Err(anyhow!(format!(
+                        "cannot retrieve non-contiguous tensors {:?}",
+                        tensor_info.layout
+                    )));
+                }
+                let start_offset = tensor_info.layout.start_offset();
+                if start_offset > 0 {
+                    std::io::copy(
+                        &mut reader.by_ref().take(start_offset as u64),
+                        &mut std::io::sink(),
+                    )?;
+                }
+                let tensor = crate_tensor_from_reader(
+                    tensor_info.layout.shape().clone(),
+                    tensor_info.dtype,
+                    &mut reader,
+                )?;
+
+                if rank > 1 && is_fortran_contiguous {
+                    // Reverse the shape, e.g. Shape(2, 3, 4) -> Shape(4, 3, 2)
+                    let shape_reversed: Vec<_> =
+                        tensor_info.layout.dims().iter().rev().cloned().collect();
+                    let tensor = tensor.reshape(shape_reversed)?;
+
+                    // Permute (transpose) the dimensions, e.g. Shape(4, 3, 2) -> Shape(2, 3, 4)
+                    let dim_indeces_reversed: Vec<_> = (0..rank).rev().collect();
+                    let tensor = tensor.permute(dim_indeces_reversed)?;
+                    // Ok(Some(tensor))
+                    tensors.push((name.clone(), tensor));
+                } else {
+                    tensors.push((name.clone(), tensor));
+                }
+            }
+        };
+    }
+    Ok(tensors)
+}
+
 pub fn round_by_factor(num: u32, factor: u32) -> u32 {
     let round = (num as f32 / factor as f32).round() as u32;
     round * factor
@@ -481,6 +712,25 @@ pub fn extract_user_text(mes: &ChatCompletionParameters) -> Result<String> {
     Ok(ret)
 }
 
+pub fn extract_user_text_vec(mes: &ChatCompletionParameters) -> Result<Vec<String>> {
+    let mut ret = vec![];
+    for chat_mes in mes.messages.clone() {
+        if let ChatMessage::User { content, .. } = chat_mes.clone()
+            && let ChatMessageContent::ContentPart(part_vec) = content
+        {
+            for part in part_vec {
+                if let ChatMessageContentPart::Text(text_part) = part {
+                    let text = text_part.text;
+                    if text.chars().count() > 0 {
+                        ret.push(text);
+                    }
+                }
+            }
+        }
+    }
+    Ok(ret)
+}
+
 pub fn get_default_save_dir() -> Option<String> {
     home_dir().map(|mut path| {
         path.push(".aha");
@@ -491,6 +741,43 @@ pub fn get_default_save_dir() -> Option<String> {
     })
 }
 
+pub async fn download_model(
+    model_id: &str,
+    save_dir: &str,
+    max_retries: u32,
+) -> anyhow::Result<()> {
+    let mut attempts = 0u32;
+    loop {
+        attempts += 1;
+        println!(
+            "Attempting to download model (attempt {}/{})",
+            attempts, max_retries
+        );
+
+        match ModelScope::download(model_id, save_dir).await {
+            Ok(()) => {
+                println!("Model downloaded successfully");
+                return Ok(());
+            }
+            Err(e) => {
+                if attempts >= max_retries {
+                    return Err(anyhow::anyhow!(
+                        "Failed to download model after {} attempts. Last error: {}",
+                        max_retries,
+                        e
+                    ));
+                }
+
+                println!(
+                    "Download failed (attempt {}): {}. Retrying in 2 seconds...",
+                    attempts, e
+                );
+                sleep(Duration::from_secs(2)).await;
+            }
+        }
+    }
+}
+
 pub fn get_file_path(file: &str) -> Result<PathBuf> {
     let path = url::Url::parse(file)?;
     let path = path.to_file_path();
@@ -504,3 +791,14 @@ pub fn get_file_path(file: &str) -> Result<PathBuf> {
     };
     Ok(path)
 }
+
+pub fn capitalize_first_letter(input: &str) -> String {
+    if input.is_empty() {
+        return input.to_string();
+    }
+
+    let mut chars = input.chars();
+    let first_char = chars.next().unwrap().to_uppercase().collect::<String>();
+    let remaining = chars.as_str().to_lowercase();
+    format!("{}{}", first_char, remaining)
+}

src/utils/tensor_utils.rs

@@ -2,7 +2,23 @@ use anyhow::{Result, anyhow};
 use candle_core::{D, DType, Device, IndexOp, Tensor, shape::Dim};
 use candle_nn::ops::sigmoid;
 
-pub fn mask_filled(on_true: &Tensor, mask: &Tensor, on_false: f32) -> Result<Tensor> {
+pub enum PaddingSide {
+    Left,
+    Right,
+}
+
+pub fn masked_fill_zeros(hidden_states: &Tensor, mask: &Tensor) -> Result<Tensor> {
+    // hidden_states: (bs, seq_len, hidden_dim)
+    // mask: (bs, seq_len)
+    let on_false = hidden_states.zeros_like()?;
+    let mask = mask
+        .unsqueeze(D::Minus1)?
+        .broadcast_as(hidden_states.shape())?;
+    let hidden_states = mask.where_cond(&hidden_states, &on_false)?;
+    Ok(hidden_states)
+}
+
+pub fn attn_masked_fill(on_true: &Tensor, mask: &Tensor, on_false: f32) -> Result<Tensor> {
     let (mask_seq_len, _) = mask.dims2()?;
     let (_, _, seq_len, _) = on_true.dims4()?;
     assert!(
@@ -89,13 +105,14 @@ pub fn split_tensor_with_size<D: Dim>(
     let dim = dim.to_index(t.shape(), "split")?;
     let mut split_res = Vec::new();
     let dim_size = t.dim(dim)?;
-    assert_eq!(
+    // assert_eq!(
-        dim_size % splits_size,
+    //     dim_size % splits_size,
-        0,
+    //     0,
-        "input tensor dim size % splits_size must be equal to 0"
+    //     "input tensor dim size % splits_size must be equal to 0"
-    );
+    // );
-    for split in (0..dim_size).step_by(splits_size) {
+    for (i, split) in (0..dim_size).step_by(splits_size).enumerate() {
-        split_res.push(t.narrow(dim, split, splits_size)?);
+        let size = splits_size.min(dim_size - i*splits_size);
+        split_res.push(t.narrow(dim, split, size)?);
     }
     Ok(split_res)
 }
@@ -476,6 +493,44 @@ pub fn interpolate_linear_1d(
     Ok(output)
 }
 
+pub fn interpolate_nearest_1d(t: &Tensor, target_size: usize) -> Result<Tensor> {
+    // t: [b, channels, features]
+    if t.rank() != 3 {
+        return Err(anyhow::anyhow!(
+            "Input rank must have equal to 3 dimensions"
+        ));
+    }
+    let shape = t.dims();
+    let orig_size = shape[shape.len() - 1];
+    if orig_size == target_size {
+        return Ok(t.clone());
+    }
+
+    let (bs, channels, _) = t.dims3()?;
+    let mut output = Tensor::zeros((bs, channels, target_size), t.dtype(), t.device())?;
+    let coords = compute_1d_coords(orig_size, target_size, None)?;
+
+    for b in 0..bs {
+        for c in 0..channels {
+            let input_slice = t.i((b, c))?;
+            let mut out_i = Vec::new();
+
+            for &coord in coords.iter().take(target_size) {
+                // Nearest neighbor: round to nearest integer coordinate
+                let nearest_idx = coord.floor() as usize;
+                let clamped_idx = nearest_idx.min(orig_size - 1);
+
+                let value = input_slice.get(clamped_idx)?;
+                out_i.push(value);
+            }
+            let out_i = Tensor::stack(&out_i, 0)?.unsqueeze(0)?.unsqueeze(0)?;
+            output = output.slice_assign(&[(b..b + 1), (c..c + 1), (0..target_size)], &out_i)?;
+        }
+    }
+    output = output.contiguous()?;
+    Ok(output)
+}
+
 pub fn interpolate_bilinear(
     input: &Tensor,
     target_size: (usize, usize),
@@ -909,3 +964,99 @@ pub fn pad_replicate_last_dim(t: &Tensor, pad: (usize, usize)) -> Result<Tensor>
     }
     Ok(pad_tensor)
 }
+
+pub fn log10(t: &Tensor) -> Result<Tensor> {
+    Ok(t.log()?.affine(1.0 / 10.0_f64.ln(), 0.0)?)
+}
+
+pub fn z_score_normalize(t: &Tensor, dim: usize) -> Result<Tensor> {
+    let rank = t.rank();
+    if dim >= rank {
+        return Err(anyhow!(format!("input dim {} must < rank {}", dim, rank)));
+    }
+    Ok(t.broadcast_sub(&t.mean_keepdim(dim)?)?
+        .broadcast_div(&t.var_keepdim(dim)?.sqrt()?)?)
+}
+
+pub fn l2_normalize(t: &Tensor, dim: usize) -> Result<Tensor> {
+    let rank = t.rank();
+    if dim >= rank {
+        return Err(anyhow!(format!("input dim {} must < rank {}", dim, rank)));
+    }
+    let l2_norm = t.sqr()?.sum_keepdim(dim)?.sqrt()?;
+    Ok(t.broadcast_div(&l2_norm)?)
+}
+
+pub fn l1_normalize(t: &Tensor, dim: usize) -> Result<Tensor> {
+    let rank = t.rank();
+    if dim >= rank {
+        return Err(anyhow!(format!("input dim {} must < rank {}", dim, rank)));
+    }
+    let l1_norm = t.abs()?.sum_keepdim(dim)?;
+    Ok(t.broadcast_div(&l1_norm)?)
+}
+
+pub fn pool1d(xs: &Tensor, pool_size: usize, ceil_mode: bool, stype: &str) -> Result<Tensor> {
+    // xs: (bs, c, dim)
+    // ceil_mode: 是否保留不完整窗口，为true时通过pad实现
+    if pool_size == 0 {
+        return Err(anyhow!("pool_size must be greater than 0"));
+    }
+    let (bs, c, dim) = xs.dims3()?;
+    let xs_reshape = if ceil_mode {
+        let remain = dim % pool_size;
+        if remain > 0 {
+            let pad = pool_size - remain;
+            let xs_pad = pad_replicate_last_dim(xs, (0, pad))?;
+            xs_pad.reshape((bs, c, (), pool_size))?
+        } else {
+            xs.reshape((bs, c, (), pool_size))?
+        }
+    } else {
+        let remain = dim % pool_size;
+        if remain > 0 {
+            let xs_del = xs.narrow(D::Minus1, 0, dim - remain)?;
+            xs_del.reshape((bs, c, (), pool_size))?
+        } else {
+            xs.reshape((bs, c, (), pool_size))?
+        }
+    };
+    let xs_pool = match stype {
+        "avg" => xs_reshape.mean(D::Minus1)?,
+        "max" => xs_reshape.max(D::Minus1)?,
+        "min" => xs_reshape.min(D::Minus1)?,
+        _ => {
+            return Err(anyhow!(
+                "unsupported pool type: {}, supported types are: avg, max, min",
+                stype
+            ));
+        }
+    };
+    Ok(xs_pool)
+}
+
+pub fn statistics_pooling(xs: &Tensor, dim: D, keepdim: bool) -> Result<Tensor> {
+    let mean = xs.mean(dim)?;
+    let std = xs.var(dim)?.sqrt()?;
+    let mut stats = Tensor::cat(&[mean, std], D::Minus1)?;
+    if keepdim {
+        stats = stats.unsqueeze(dim)?;
+    }
+    Ok(stats)
+}
+pub fn float_range_normalize(t: &Tensor) -> Result<Tensor> {
+    let peak = t
+        .to_dtype(DType::F32)?
+        .abs()?
+        .max_all()?
+        .to_scalar::<f32>()?;
+    if peak == 0.0 {
+        return Ok(t.clone());
+    }
+    let mut t = t.clone();
+    if peak > 1.0 {
+        t = t.affine(1.0 / peak as f64, 0.0)?;
+    }
+    t = t.clamp(-1.0, 1.0)?;
+    Ok(t)
+}

tests/messy_test.rs

@@ -1,22 +1,31 @@
 // use std::io::Cursor;
 
-use aha::utils::audio_utils::create_hann_window;
+use std::time::Instant;
+
+use aha::utils::{tensor_utils::interpolate_nearest_1d};
 use anyhow::Result;
-use candle_core::DType;
+use candle_core::{Tensor};
 // use symphonia::core::io::MediaSourceStream;
 
 #[test]
 fn messy_test() -> Result<()> {
     // RUST_BACKTRACE=1 cargo test -F cuda,ffmpeg messy_test -r -- --nocapture
     let device = &candle_core::Device::Cpu;
+    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()?;
     // let vec_u8 = response.bytes()?.to_vec();
     // let mut content = Cursor::new(vec_u8);
     // let mss = MediaSourceStream::new(Box::new(content), Default::default());
-    let window = create_hann_window(400, DType::F32, device)?;
+    // let window = create_hann_window(400, DType::F32, device)?;
-    println!("window: {}", window);
+    // println!("window: {}", window);
     // let audio_path = "file:///home/jhq/Videos/voice_01.wav";
     // let audio_path = "/home/jhq/Videos/zh.mp3";
     // let audio_path = "/home/jhq/Videos/zh.mp3";

tests/test_index_tts2.rs

@@ -0,0 +1,48 @@
+use std::time::Instant;
+use anyhow::Result;
+
+use aha::models::index_tts2::{generate::IndexTTS2Generate, utils::download_index_tts2_need_model};
+use aha_openai_dive::v1::resources::chat::ChatCompletionParameters;
+
+#[tokio::test]
+async fn index_tts2_generate() -> Result<()> {
+    // RUST_BACKTRACE=1 cargo test -F cuda index_tts2_generate -r -- --nocapture
+    let save_dir =
+        aha::utils::get_default_save_dir().ok_or(anyhow::anyhow!("Failed to get save dir"))?;
+    let _ = download_index_tts2_need_model(Some(&save_dir)).await?;
+    let model_path = format!("{}/IndexTeam/IndexTTS-2", save_dir);
+    let message = r#"
+    {
+        "model": "index-tts2",
+        "messages": [
+            {
+                "role": "user",
+                "content": [
+                    {
+                        "type": "audio",
+                        "audio_url": 
+                        {
+                            "url": "file:///home/jhq/Videos/voice_01.wav"
+                        }
+                    },              
+                    {
+                        "type": "text", 
+                        "text": "你好啊"
+                    }
+                ]
+            }
+        ]
+    }
+    "#;
+    let mes: ChatCompletionParameters = serde_json::from_str(message)?;
+    let i_start = Instant::now();
+    let mut voxcpm_generate = IndexTTS2Generate::init(&model_path, None, None)?;
+    let i_duration = i_start.elapsed();
+    println!("Time elapsed in load model is: {:?}", i_duration);
+
+    let i_start = Instant::now();
+    let generate = voxcpm_generate.generate(mes)?;
+    let i_duration = i_start.elapsed();
+    println!("Time elapsed in generate is: {:?}", i_duration);
+    Ok(())
+}

tests/test_qwen3_asr.rs

@@ -0,0 +1,90 @@
+use std::{pin::pin, time::Instant};
+
+use aha::models::{GenerateModel, qwen3_asr::generate::Qwen3AsrGenerateModel};
+use aha_openai_dive::v1::resources::chat::ChatCompletionParameters;
+use anyhow::Result;
+use rocket::futures::StreamExt;
+#[test]
+fn qwen3_asr_generate() -> Result<()> {
+    // RUST_BACKTRACE=1 cargo test -F cuda qwen3_asr_generate -r -- --nocapture
+    let save_dir =
+        aha::utils::get_default_save_dir().ok_or(anyhow::anyhow!("Failed to get save dir"))?;
+    let model_path = format!("{}/Qwen/Qwen3-ASR-0.6B/", save_dir);  //Qwen/Qwen3-ASR-1.7B
+    let message = r#"
+    {
+        "model": "qwen3-asr",
+        "messages": [
+            {
+                "role": "user",
+                "content": [
+                    {
+                        "type": "audio",
+                        "audio_url": 
+                        {
+                            "url": "https://package-release.coderbox.cn/aiway/test/other/%E5%93%AA%E5%90%92.wav"
+                        }
+                    }
+                ]
+            }
+        ]
+    }
+    "#;
+    // "metadata": {"language": "Chinese"}
+    let mes: ChatCompletionParameters = serde_json::from_str(message)?;
+    let i_start = Instant::now();
+    let mut model = Qwen3AsrGenerateModel::init(&model_path, None, None)?;
+    let i_duration = i_start.elapsed();
+    println!("Time elapsed in load model is: {:?}", i_duration);
+    let i_start = Instant::now();
+    let res = model.generate(mes)?;
+    let i_duration = i_start.elapsed();
+    println!("generate: \n {:?}", res);
+    if res.usage.is_some() {
+        let num_token = res.usage.as_ref().unwrap().total_tokens;
+        let duration_secs = i_duration.as_secs_f64();
+        let tps = num_token as f64 / duration_secs;
+        println!("Tokens per second (TPS): {:.2}", tps);
+    }
+    println!("Time elapsed in generate is: {:?}", i_duration);
+    Ok(())
+}
+
+#[tokio::test]
+async fn qwen3_asr_stream() -> Result<()> {
+    // RUST_BACKTRACE=1 cargo test -F cuda qwen3_asr_stream -r -- --nocapture
+    let save_dir =
+        aha::utils::get_default_save_dir().ok_or(anyhow::anyhow!("Failed to get save dir"))?;
+    let model_path = format!("{}/Qwen/Qwen3-ASR-0.6B/", save_dir);
+    let message = r#"
+    {
+        "model": "qwen3-asr",
+        "messages": [
+            {
+                "role": "user",
+                "content": [
+                    {
+                        "type": "audio",
+                        "audio_url": 
+                        {
+                            "url": "https://package-release.coderbox.cn/aiway/test/other/%E5%93%AA%E5%90%92.wav"
+                        }
+                    }
+                ]
+            }
+        ]
+    }
+    "#;
+    let mes: ChatCompletionParameters = serde_json::from_str(message)?;
+    let i_start = Instant::now();
+    let mut fun_asr_model = Qwen3AsrGenerateModel::init(&model_path, None, None)?;
+    let i_duration = i_start.elapsed();
+    println!("Time elapsed in load model is: {:?}", i_duration);
+    let i_start = Instant::now();
+    let mut stream = pin!(fun_asr_model.generate_stream(mes)?);
+    while let Some(item) = stream.next().await {
+        println!("generate: \n {:?}", item);
+    }
+    let i_duration = i_start.elapsed();
+    println!("Time elapsed in generate is: {:?}", i_duration);
+    Ok(())
+}

tests/weight_test.rs

@@ -1,6 +1,6 @@
 use std::collections::HashMap;
 
-use aha::utils::{find_type_files, get_device};
+use aha::utils::{find_type_files, get_device, read_pth_tensor_info_cycle};
 use anyhow::Result;
 use candle_core::{Device, pickle::read_all_with_key, safetensors};
 use candle_nn::VarBuilder;
@@ -197,3 +197,41 @@ fn qwen3_weight() -> Result<()> {
     println!("model_list: {:?}", model_list);
     Ok(())
 }
+
+#[test]
+fn index_tts2_weight() -> Result<()> {
+    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); 
+    let s2mel_path = model_path+ "/s2mel.pth";
+    // let wac2vec2_path = model_path+ "/wav2vec2bert_stats.pt";
+    // let model_path = format!("{}/iic/speech_campplus_sv_zh-cn_16k-common/", save_dir);
+    // let campplus_path = model_path+ "/campplus_cn_common.bin";
+    // let model_list = find_type_files(&model_path, "safetensors")?;
+    let model_list = vec![s2mel_path];
+    // let mut dict_to_hashmap = HashMap::new();
+    // let mut dtype = candle_core::DType::F32;
+    for m in model_list {
+        // let dict = read_all_with_key(m, Some("state_dict"))?;
+        // let dict = read_all_with_key(m, Some("net"))?;
+        let dict = read_pth_tensor_info_cycle(m, Some("net.cfm"))?;
+        // dtype = dict[0].1.dtype();
+        for (k, v) in dict {
+            // if k.contains("model") {
+            //     println!("key: {}, tensor shape: {:?}", k, v);
+            // }
+            // dict_to_hashmap.insert(k, v);
+            println!("key: {}, tensor shape: {:?}", k, v);
+        }
+    }
+    // let device = Device::Cpu;
+    // let semantic_codec_path = save_dir.to_string() + "/amphion/MaskGCT/semantic_codec/model.safetensors" ;
+    // let model_list = vec![semantic_codec_path];
+    // for m in model_list {
+    //     let weights = safetensors::load(m, &device)?;
+    //     for (key, tensor) in weights.iter() {            
+    //         println!("=== {} === {:?}", key, tensor.shape());
+    //     }
+    // }
+    Ok(())
+}