index tts stash save

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

@@ -118,7 +118,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

@@ -25,7 +25,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"
@@ -77,7 +77,7 @@ case $MODEL_ALIAS in
     "paddleocr-vl")
         MODEL_ID="PaddlePaddle/PaddleOCR-VL"
         ;;
-    "RMBG2.0")
+    "rmbg2.0")
         MODEL_ID="briaai/RMBG-2.0"
         ;;
     "voxcpm")

src/main.rs

@@ -1,6 +1,6 @@
 use std::{net::IpAddr, str::FromStr, time::Duration};
 
-use aha::{models::WhichModel, utils::get_default_save_dir};
+use aha::{models::WhichModel, utils::{download_model, get_default_save_dir}};
 use clap::Parser;
 use modelscope::ModelScope;
 use rocket::{
@@ -34,38 +34,38 @@ struct Args {
     #[arg(long)]
     download_retries: Option<u32>,
 }
-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
-        );
+// 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
-                    ));
-                }
+//         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;
-            }
-        }
-    }
-}
+//                 println!(
+//                     "Download failed (attempt {}): {}. Retrying in 2 seconds...",
+//                     attempts, e
+//                 );
+//                 sleep(Duration::from_secs(2)).await;
+//             }
+//         }
+//     }
+// }
 
 #[tokio::main]
 async fn main() -> anyhow::Result<()> {

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,
-    LayerNorm, LayerNormConfig, Linear, Module, RmsNorm, VarBuilder, batch_norm, conv1d,
-    conv1d_no_bias, conv2d, conv2d_no_bias, embedding, layer_norm, linear_b, linear_no_bias,
-    rms_norm,
+    Activation, BatchNorm, BatchNormConfig, Conv1d, Conv1dConfig, Conv2d, Conv2dConfig,
+    ConvTranspose1d, ConvTranspose1dConfig, Embedding, LayerNorm, LayerNormConfig, Linear, Module,
+    ModuleT, RmsNorm, VarBuilder, batch_norm, conv1d, conv1d_no_bias, conv2d, conv2d_no_bias,
+    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 up_proj = linear_b(hidden_size, intermediate_size, bias, vb.pp("up_proj"))?;
-        let down_proj = linear_b(intermediate_size, hidden_size, bias, vb.pp("down_proj"))?;
+        let gate_pp_name = gate_pp_name.unwrap_or("gate_proj");
+        let up_pp_name = up_pp_name.unwrap_or("up_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/mod.rs

@@ -0,0 +1 @@
+pub mod seamless_m4t_feature_extractor;

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/processor.rs

@@ -9,8 +9,9 @@ use crate::{
     utils::{
         audio_utils::{
             apply_stft, create_hann_window, extract_audios, extract_frames, mel_filter_bank,
+            torch_stft,
         },
-        tensor_utils::{pad_reflect_last_dim, split_tensor},
+        tensor_utils::{log10, pad_reflect_last_dim, split_tensor},
     },
 };
 
@@ -87,19 +88,21 @@ impl GlmAsrNanoProcessor {
         let waveform = pad_reflect_last_dim(waveform, (pad, pad))?;
         let (_, samples) = waveform.dims2()?;
 
-        // 计算输出维度
+        // // (bs, n_frames, n_fft)
+        // let frames = extract_frames(&waveform, self.n_fft, self.hop_length)?;
+        // // 应用汉明窗口
+        // 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)?
+            .transpose(D::Minus1, D::Minus2)?;
         let n_frames = (samples - self.n_fft) / self.hop_length + 1;
-        // (bs, n_frames, n_fft)
-        let frames = extract_frames(&waveform, self.n_fft, self.hop_length)?;
-        // 应用汉明窗口
-        let result = frames.broadcast_mul(&self.window)?;
-        // 傅立叶变换
-        let magnitudes = apply_stft(&result)?.transpose(D::Minus1, D::Minus2)?;
         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 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)?;

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, spectrogram,
+            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, IndexOp, 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,8 +1,12 @@
 pub mod common;
+pub mod campplus;
 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;
@@ -10,6 +14,7 @@ pub mod qwen3;
 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,
@@ -51,7 +56,7 @@ pub enum WhichModel {
     HunyuanOCR,
     #[value(name = "paddleocr-vl")]
     PaddleOCRVL,
-    #[value(name = "RMBG2.0")]
+    #[value(name = "rmbg2.0")]
     RMBG2_0,
     #[value(name = "voxcpm")]
     VoxCPM,

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/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_out = get_batch_norm(vb.pp("bn_out"), 1e-5, out_c)?;
+        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, 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_3 = Conv2dWithBN::new(vb.pp("gdt_convs_3"), channels[2], n, 3, 1, 1)?;
-        let gdt_convs_2 = Conv2dWithBN::new(vb.pp("gdt_convs_2"), channels[3], n, 3, 1, 1)?;
+        let gdt_convs_4 =
+            Conv2dWithBN::new(vb.pp("gdt_convs_4"), channels[1], n, 3, 1, 1, true, true)?;
+        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, NaiveAttention, 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>,

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};
 
 // 重采样方法枚举
 #[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;
-    let window: Vec<f32> = (0..window_size)
+    if window_size < 1 {
+        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,73 @@ 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)
+}

src/utils/mod.rs

@@ -3,7 +3,8 @@ pub mod img_utils;
 pub mod tensor_utils;
 pub mod video_utils;
 
-use std::{fs, process::Command};
+use std::io::Read;
+use std::{collections::HashMap, fs, 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 candle_core::{DType, Device};
+use anyhow::{Result, anyhow};
+use byteorder::{LittleEndian, ReadBytesExt};
+use candle_core::{
+    Context, DType, Device, Shape, Tensor,
+    pickle::{Object, PthTensors, 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
@@ -490,3 +721,40 @@ pub fn get_default_save_dir() -> Option<String> {
         path.to_string_lossy().to_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;
+            }
+        }
+    }
+}

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!(
@@ -476,6 +492,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 +963,83 @@ 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)
+}

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::{audio_utils::create_hann_window, tensor_utils::interpolate_nearest_1d};
 use anyhow::Result;
-use candle_core::DType;
+use candle_core::{DType, 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)?;
-    println!("window: {}", window);
+    // let window = create_hann_window(400, DType::F32, device)?;
+    // 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/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(())
+}