update voxcpm

Cargo.lock

@@ -36,7 +36,6 @@ dependencies = [
  "openai_dive",
  "reqwest",
  "rocket",
- "rubato",
  "serde",
  "serde_json",
  "tokenizers",
@@ -2626,15 +2625,6 @@ dependencies = [
  "zerocopy",
 ]
 
-[[package]]
-name = "primal-check"
-version = "0.3.4"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "dc0d895b311e3af9902528fbb8f928688abbd95872819320517cc24ca6b2bd08"
-dependencies = [
- "num-integer",
-]
-
 [[package]]
 name = "proc-macro-crate"
 version = "3.4.0"
@@ -2918,15 +2908,6 @@ dependencies = [
  "crossbeam-utils",
 ]
 
-[[package]]
-name = "realfft"
-version = "3.5.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "f821338fddb99d089116342c46e9f1fbf3828dba077674613e734e01d6ea8677"
-dependencies = [
- "rustfft",
-]
-
 [[package]]
 name = "reborrow"
 version = "0.5.5"
@@ -3153,18 +3134,6 @@ dependencies = [
  "uncased",
 ]
 
-[[package]]
-name = "rubato"
-version = "0.16.2"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "5258099699851cfd0082aeb645feb9c084d9a5e1f1b8d5372086b989fc5e56a1"
-dependencies = [
- "num-complex",
- "num-integer",
- "num-traits",
- "realfft",
-]
-
 [[package]]
 name = "rustc-demangle"
 version = "0.1.26"
@@ -3177,20 +3146,6 @@ version = "2.1.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "357703d41365b4b27c590e3ed91eabb1b663f07c4c084095e60cbed4362dff0d"
 
-[[package]]
-name = "rustfft"
-version = "6.4.1"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "21db5f9893e91f41798c88680037dba611ca6674703c1a18601b01a72c8adb89"
-dependencies = [
- "num-complex",
- "num-integer",
- "num-traits",
- "primal-check",
- "strength_reduce",
- "transpose",
-]
-
 [[package]]
 name = "rustix"
 version = "1.1.2"
@@ -3523,12 +3478,6 @@ version = "1.1.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "a2eb9349b6444b326872e140eb1cf5e7c522154d69e7a0ffb0fb81c06b37543f"
 
-[[package]]
-name = "strength_reduce"
-version = "0.2.4"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "fe895eb47f22e2ddd4dabc02bce419d2e643c8e3b585c78158b349195bc24d82"
-
 [[package]]
 name = "strsim"
 version = "0.11.1"
@@ -4042,16 +3991,6 @@ dependencies = [
  "tracing-log",
 ]
 
-[[package]]
-name = "transpose"
-version = "0.2.3"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "1ad61aed86bc3faea4300c7aee358b4c6d0c8d6ccc36524c96e4c92ccf26e77e"
-dependencies = [
- "num-integer",
- "strength_reduce",
-]
-
 [[package]]
 name = "try-lock"
 version = "0.2.5"

Cargo.toml

@@ -27,7 +27,6 @@ chrono = "0.4.42"
 rocket = "0.5.1"
 tokio = "1.47.1"
 hound = "3.5.1"
-rubato = "0.16.2"
 
 [features]
 flash-attn=["candle-flash-attn"]

src/models/voxcpm/audio_vae.rs

@@ -1,7 +1,7 @@
-use anyhow::{Error, Ok, Result};
+use anyhow::{Ok, Result};
-use candle_core::{D, IndexOp, Tensor};
+use candle_core::{D, Tensor};
 use candle_nn::{Conv1d, Conv1dConfig, ConvTranspose1d, ConvTranspose1dConfig, Module, VarBuilder};
-use std::result::Result::Ok as StdOk;
+use std::{result::Result::Ok as StdOk, thread, time};
 
 pub struct CausalConv1d {
     conv1d: Conv1d,
@@ -9,13 +9,9 @@ pub struct CausalConv1d {
 }
 
 impl CausalConv1d {
-    // CausalConv1d::new(scaled_weight, bias, padding, dilation, stride)?;
     pub fn new(
         weight: Tensor,
         bias: Option<Tensor>,
-        // in_c: usize,
-        // out_c: usize,
-        // kernel_size: usize,
         padding: usize,
         dilation: usize,
         groups: usize,
@@ -59,7 +55,7 @@ impl CausalConvTranspose1d {
     ) -> Result<Self> {
         let config = ConvTranspose1dConfig {
             padding: 0,
-            output_padding,
+            output_padding: 0,
             stride,
             dilation,
             groups,
@@ -74,23 +70,10 @@ impl CausalConvTranspose1d {
         })
     }
     pub fn forward(&self, x: &Tensor) -> Result<Tensor> {
-        println!("transpose conv input x: {:?}", x);
-        println!("transpose conv config stride: {:?}", self.config.stride);
-        println!("transpose conv config padding: {:?}", self.config.padding);
-        println!("transpose conv config output_padding: {:?}", self.config.output_padding);
-        println!("transpose conv config groups: {:?}", self.config.groups);
-        println!("transpose conv config dilation: {:?}", self.config.dilation);
-        println!("transpose conv config weight: {:?}", self.conv_transpose1d.weight());
-
         let x = self.conv_transpose1d.forward(x)?;
-        println!("transpose conv after x: {:?}", x);
-        println!("transpose conv after self.padding: {:?}", self.padding);
-        println!("transpose conv after self.output_padding: {:?}", self.output_padding);
         let last_dim = x.dim(D::Minus1)?;
         let select_num = last_dim - (self.padding * 2 - self.output_padding);
-        println!("transpose conv after select_num: {:?}", select_num);
         let x = x.narrow(D::Minus1, 0, select_num)?;
-        println!("transpose conv after x: {:?}", x);
         Ok(x)
     }
 }
@@ -109,24 +92,21 @@ impl WNCausalConv1d {
         groups: usize,
         stride: usize,
     ) -> Result<Self> {
-        let in_c = in_c / groups;
+        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 = match vb.get(out_c, "bias") {
             StdOk(b) => Some(b),
             Err(_) => None,
         };
-        let weight_norm = weight_v.sqr()?.sum_keepdim(D::Minus1)?.sqrt()?;
+        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 conv = CausalConv1d::new(scaled_weight, bias, padding, dilation, groups, stride)?;
         Ok(Self { conv })
     }
     pub fn forward(&self, x: &Tensor) -> Result<Tensor> {
-        println!("conv1d: x: {:?}", x);
-        println!("conv weight: : {:?}", self.conv.conv1d.weight());
         let x = self.conv.forward(x)?;
-        println!("conv1d: WN causal x: {:?}", x);
         Ok(x)
     }
 }
@@ -154,7 +134,7 @@ impl WNCausalConvTranspose1d {
             StdOk(b) => Some(b),
             Err(_) => None,
         };
-        let weight_norm = weight_v.sqr()?.sum_keepdim(D::Minus1)?.sqrt()?;
+        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 conv_transpose = CausalConvTranspose1d::new(
@@ -230,7 +210,6 @@ impl CausalResidualUnit {
     }
 
     pub fn forward(&self, x: &Tensor) -> Result<Tensor> {
-        println!("causal residual unit x: {:?}", x);
         // let orig_dim = x.dims();
         let last_dim_x = x.dim(D::Minus1)?;
         let mut res_x = x.clone();
@@ -238,11 +217,8 @@ impl CausalResidualUnit {
         let y = self.block1.forward(&y)?;
         let y = self.block2.forward(&y)?;
         let y = self.block3.forward(&y)?;
-        println!("causal residual unit y: {:?}", y);
         // let dim = y.dims();
         let last_dim_y = y.dim(D::Minus1)?;
-        println!("last_dim_x: {:?}", last_dim_x);
-        println!("last_dim_y: {:?}", last_dim_y);
         let pad = (last_dim_x - last_dim_y) / 2;
         if pad > 0 {
             res_x = res_x.narrow(D::Minus1, pad, last_dim_y)?;
@@ -415,17 +391,11 @@ impl CausalDecoderBlock {
     }
 
     pub fn forward(&self, x: &Tensor) -> Result<Tensor> {
-        println!("decoder block x : {:?}", x);
         let x = self.block0.forward(x)?;
-        println!("decoder block0 x : {:?}", x);
         let x = self.block1.forward(&x)?;
-        println!("decoder block1 x : {:?}", x);
         let x = self.block2.forward(&x)?;
-        println!("decoder block2 x : {:?}", x);
         let x = self.block3.forward(&x)?;
-        println!("decoder block3 x : {:?}", x);
         let x = self.block4.forward(&x)?;
-        println!("decoder block4 x : {:?}", x);
         Ok(x)
     }
 }
@@ -456,7 +426,7 @@ impl CausalDecoder {
             input_channel,
             1,
         )?;
-        let model1 = WNCausalConv1d::new(vb.pp("model.1"), input_channel, channels, 1, 1, 1, 1, 1)?;
+        let model1 = WNCausalConv1d::new(vb.pp("model.1"), input_channel, channels, 1, 1, 0, 1, 1)?;
         let vb_model = vb.pp("model");
         let mut output_dim = channels;
         let mut model2_5 = Vec::new();
@@ -484,10 +454,8 @@ impl CausalDecoder {
         })
     }
 
-    pub fn forward(&self, x: &Tensor) -> Result<Tensor> {
+    pub fn forward(&self, x: &Tensor) -> Result<Tensor> {   
-        print!("audio_vae decoder input x shape: {:?}", x);
+        let x = self.model0.forward(x)?;    
-        let x = self.model0.forward(x)?;
-        print!("audio_vae decoder model0 x shape: {:?}", x);
         let mut x = self.model1.forward(&x)?;
         for model_i in &self.model2_5 {
             x = model_i.forward(&x)?;

src/models/voxcpm/model.rs

@@ -72,10 +72,10 @@ impl SinusoidalPosEmb {
             .to_dtype(x.dtype())?;
 
         let emb = x
-            .unsqueeze(D::Minus1)?
+            .unsqueeze(1)?
             .contiguous()?
-            .matmul(&emb.unsqueeze(0)?.contiguous()?)?
+            .affine(scale as f64, 0.0)?
-            .affine(scale as f64, 0.0)?;
+            .matmul(&emb.unsqueeze(0)?.contiguous()?)?;
         let emb = Tensor::cat(&[emb.sin()?, emb.cos()?], D::Minus1)?;
         Ok(emb)
     }
@@ -167,7 +167,7 @@ impl VoxCPMLocDiT {
         let cond = self
             .cond_proj
             .forward(&cond.transpose(1, 2)?.contiguous()?)?;
-        let prefix = cond.dims()[1];
+        let prefix = cond.dim(1)?;
         let t = self.time_embeddings.forward(t, 1000)?.to_dtype(x.dtype())?;
         let t = self.time_mlp.forward(&t)?;
         let dt = self
@@ -233,7 +233,6 @@ impl UnifiedCFM {
         let z = Tensor::randn(0.0f32, 1.0, (b, self.in_channels, t), mu.device())?
             .to_dtype(dtype)?
             .affine(temperature, 0.0)?;
-        println!("z: {}", z);
         let t_span = linspace(1.0, 0.0, n_timesteps + 1, mu.device())?.to_dtype(dtype)?;
         let t_span = t_span
             .affine(f64::consts::PI / 2.0, 0.0)?
@@ -242,11 +241,6 @@ impl UnifiedCFM {
             .add(&t_span)?
             .affine(sway_sampling_coef, 0.0)?
             .add(&t_span)?;
-        println!("t_span: {}", t_span);
-        println!("mu: {}", mu);
-        println!("cond: {}", cond);
-        println!("cfg_value: {}", cfg_value);
-        println!("use_cfg_zero_star: {}", use_cfg_zero_star);
         let x = self.solve_euler(&z, &t_span, mu, cond, cfg_value, use_cfg_zero_star)?;
         Ok(x)
     }
@@ -274,7 +268,7 @@ impl UnifiedCFM {
         let mut t = t_span.i(0)?;
         let mut dt = t.sub(&t_span.i(1)?)?;
         let mut sol = Vec::new();
-        let t_span_len = t_span.dims1()?;
+        let t_span_len = t_span.dim(0)?;
         let zero_init_steps = max(1, (t_span_len as f32 * 0.04) as usize);
         let mut dphi_dt = Tensor::zeros(1, t_span.dtype(), t_span.device())?;
         let mut x = x.clone();
@@ -320,7 +314,8 @@ impl UnifiedCFM {
                 dt = t.sub(&t_span.i(step + 1)?)?;
             }
         }
-        Ok(sol[sol.len() - 1].clone())
+        let ret = sol[sol.len() - 1].clone();
+        Ok(ret)
     }
 }
 
@@ -333,8 +328,6 @@ pub struct VoxCPMLocEnc {
 
 impl VoxCPMLocEnc {
     pub fn new(vb: VarBuilder, config: VoxMiniCPM4Config, input_dim: usize) -> Result<Self> {
-        // let special_token = Tensor::randn(0.0f32, 1.0, (1, 1, 1, config.hidden_size), vb.device())?
-        //     .to_dtype(vb.dtype())?;
         let special_token = vb.get((1, 1, 1, config.hidden_size), "special_token")?;
         let in_proj = linear(input_dim, config.hidden_size, vb.pp("in_proj"))?;
         assert_eq!(
@@ -354,16 +347,12 @@ impl VoxCPMLocEnc {
     pub fn forward(&mut self, x: &Tensor) -> Result<Tensor> {
         let (b, t, p, d) = x.dims4()?;
         let x = self.in_proj.forward(x)?;
-        println!("VoxCPMLocEnc: in_proj: {}", x);
         let special_tokens = self.special_token.expand((b, t, 1, self.hidden_size))?;
         let x = Tensor::cat(&[special_tokens, x], 2)?;
-        println!("VoxCPMLocEnc: cat: {}", x);
         let (b, t, p, c) = x.dims4()?;
         let x = x.reshape((b * t, p, c))?;
         let outputs = self.encoder.forward(&x, 0, false)?;
-        println!("VoxCPMLocEnc: encoder: {}", outputs);
         let cls_output = outputs.i((.., 0, ..))?;
-        println!("VoxCPMLocEnc: cls_output: {}", cls_output);
         let cls_output = cls_output.reshape((b, t, c))?;
         Ok(cls_output)
     }
@@ -537,10 +526,8 @@ impl VoxCPMModel {
                 let audio_feat = audio_feat
                     .reshape((self.audio_vae.latent_dim, (), self.patch_size))?
                     .permute((1, 2, 0))?;
-                let dim0 = audio_feat.dim(0)?;
+                let dim0 = audio_feat.dim(0)? - 1;
-                println!("audio_feat: {:?}", audio_feat);
                 let audio_feat = audio_feat.i(..dim0)?;
-                println!("audio_feat --: {:?}", audio_feat);
                 let audio_length = audio_feat.dim(0)?;
                 let text_pad_token = Tensor::zeros(audio_length, DType::U32, &self.device)?;
                 let text_token = Tensor::cat(&[text_token, text_pad_token], D::Minus1)?;
@@ -594,7 +581,6 @@ impl VoxCPMModel {
             .squeeze(1)?;
         let decode_audio_len = decode_audio.dim(D::Minus1)? - 640 - 640;
         let decode_audio = decode_audio.narrow(D::Minus1, 640, decode_audio_len)?;
-        println!("decode_audio: {}", decode_audio);
         Ok(decode_audio)
     }
 
@@ -609,21 +595,15 @@ impl VoxCPMModel {
         inference_timesteps: usize,
         cfg_value: f64,
     ) -> Result<Tensor> {
-        println!("text: {}", text);
-        println!("text_mask: {}", text_mask);
-        println!("feat: {}", feat);
-        println!("feat_mask: {}", feat_mask);
         let (b, t, p, d) = feat.dims4()?;
         let feat_embed = self.feat_encoder.forward(feat)?; // [b, t, h_feat]
-        println!("feat_embed: {}", feat_embed);
         let feat_embed = self.enc_to_lm_proj.forward(&feat_embed)?;
-        println!("feat_embed: {}", feat_embed);
-        
         let scale_emb = if self.config.lm_config.use_mup {
             self.config.lm_config.scale_emb
         } else {
             1.0
         };
+
         let text_embed = self
             .base_lm
             .embed_tokens
@@ -631,41 +611,32 @@ impl VoxCPMModel {
             .unwrap()
             .forward(text)?
             .affine(scale_emb as f64, 0.0)?;
-        println!("text_embed: {}", text_embed);
         let combined_embed = text_mask
             .unsqueeze(D::Minus1)?
             .broadcast_mul(&text_embed)?
             .add(&feat_mask.unsqueeze(D::Minus1)?.broadcast_mul(&feat_embed)?)?;
-        println!("combined_embed: {}", combined_embed);
-
         let mut prefix_feat_cond = feat.i((.., t - 1, ..))?;
         let mut pred_feat_seq = Vec::new();
         let mut position_id = 0;
         let mut seq_len = t;
         let enc_outputs = self.base_lm.forward_step(&combined_embed, position_id)?;
-        println!("base_lm enc_outputs: {}", enc_outputs);
         let enc_outputs = self
             .fsq_layer
             .forward(&enc_outputs)?
             .broadcast_mul(&feat_mask.unsqueeze(D::Minus1)?)?
             .add(&enc_outputs.broadcast_mul(&text_mask.unsqueeze(D::Minus1)?)?)?;
-        println!("fsq_layer enc_outputs: {}", enc_outputs);
+
         let mut lm_hidden = enc_outputs.i((.., t - 1, ..))?;
-        println!("lm_hidden shape: {:?}", lm_hidden);
+
-        
         let input_embeds =
             enc_outputs.add(&feat_mask.unsqueeze(D::Minus1)?.broadcast_mul(&feat_embed)?)?;
         let residual_enc_outputs = self.residual_lm.forward_step(&input_embeds, position_id)?;
-        println!("residual_lm residual_enc_outputs: {}", residual_enc_outputs);
         let mut residual_hidden = residual_enc_outputs.i((.., t - 1, ..))?;
 
         for i in 0..max_len {
             let dit_hidden_1 = self.lm_to_dit_proj.forward(&lm_hidden)?; // [b, h_dit]
-            println!("dit_hidden_1: {}", dit_hidden_1);
             let dit_hidden_2 = self.res_to_dit_proj.forward(&residual_hidden)?; // [b, h_dit]
-            println!("dit_hidden_2: {}", dit_hidden_2);
             let dit_hidden = dit_hidden_1.add(&dit_hidden_2)?;
-            println!("dit_hidden: {}", dit_hidden);
             let cond = prefix_feat_cond.transpose(1, 2)?.contiguous()?;
 
             let pred_feat = self
@@ -681,24 +652,18 @@ impl VoxCPMModel {
                     true,
                 )?
                 .transpose(1, 2)?; // [b, p, d]
-            println!("pred_feat: {}", pred_feat);
-
             let curr_embed = self.feat_encoder.forward(&pred_feat.unsqueeze(1)?)?; // [b, 1, c]
             let curr_embed = self.enc_to_lm_proj.forward(&curr_embed)?;
-            println!("curr_embed: {}", curr_embed);
             pred_feat_seq.push(pred_feat.unsqueeze(1)?);
 
             prefix_feat_cond = pred_feat;
-            println!("lm_hidden: {}", lm_hidden);
             let stop_flag = self.stop_proj.forward(&lm_hidden)?.silu()?;
-            println!("stop_flag: {}", stop_flag);
             let stop_flag = self
                 .stop_head
                 .forward(&stop_flag)?
                 .argmax(D::Minus1)?
                 .i(0)?
                 .to_scalar::<u32>()?;
-            println!("i: {}, stop_flag: {}", i, stop_flag);
             if i > min_len && stop_flag == 1 {
                 break;
             }
@@ -716,15 +681,12 @@ impl VoxCPMModel {
         }
         let pred_seq = Tensor::cat(&pred_feat_seq, 1)?; // (b, t, p, d)
         let (b, t, p, d) = pred_seq.dims4()?;
-        println!("pred_seq: {:?}", pred_seq);
         let feat_pred = pred_seq
             .permute((0, 3, 1, 2))?
             .reshape((b, d, ()))?
             .contiguous()?;
-        println!("feat_pred: {:?}", feat_pred);
         self.base_lm.clear_kv_cache();
         self.residual_lm.clear_kv_cache();
-
         Ok(feat_pred)
     }
 }

src/models/voxcpm/tokenizer.rs

@@ -45,7 +45,7 @@ impl SingleChineseTokenizer {
             .encode(text, false)
             .map_err(|e| anyhow!(format!("tokenizer encode error: {}", e)))?;
         let tokens = encode.get_tokens();
-        println!("tokens: {:?}", tokens);
+        // println!("tokens: {:?}", tokens);
         let mut split_character = Vec::new();
         for token in tokens {
             let clean_token = token.replace("▁", "to");
@@ -56,7 +56,7 @@ impl SingleChineseTokenizer {
                 split_character.push(token.clone());
             }
         }
-        println!("split_character: {:?}", split_character);
+        // println!("split_character: {:?}", split_character);
         let ids: Vec<u32> = split_character
             .iter()
             .filter_map(|c| self.tokenizer.token_to_id(c))

src/utils/audio_utils.rs

@@ -1,11 +1,9 @@
 use anyhow::{Result, anyhow};
 use candle_core::{D, DType, Device, Tensor};
-use candle_nn::{conv1d_no_bias, Conv1d, Conv1dConfig, Module};
+use candle_nn::{Conv1d, Conv1dConfig, Module, conv1d_no_bias};
 use hound::{SampleFormat, WavReader};
 use rocket::futures::future::ok;
-use rubato::{
+
-    Resampler, SincFixedIn, SincInterpolationParameters, SincInterpolationType, WindowFunction,
-};
 use std::f64::consts::PI;
 use std::path::Path;
 
@@ -89,7 +87,7 @@ pub fn get_sinc_resample_kernel(
         ResamplingMethod::SincInterpKaiser => {
             let beta_val = beta.unwrap_or(14.769656459379492);
             let i0_beta = i0(beta_val);
-            
+
             let normalized_t = t.affine(1.0 / lowpass_filter_width as f64, 0.0)?;
             let arg = (1.0 - normalized_t.sqr()?)?;
             // 处理arg为负数的情况
@@ -97,7 +95,10 @@ pub fn get_sinc_resample_kernel(
             let sqrt_dims = sqrt_arg.dims();
             let sqrt_arg_vec = sqrt_arg.flatten_all()?.to_vec1::<f32>()?;
 
-            let window_val:Vec<f32> = sqrt_arg_vec.iter().map(|x| i0(beta_val * x) / i0_beta).collect();
+            let window_val: Vec<f32> = sqrt_arg_vec
+                .iter()
+                .map(|x| i0(beta_val * x) / i0_beta)
+                .collect();
             let window = Tensor::new(window_val, device)?.reshape(sqrt_dims)?;
             window
         }
@@ -130,20 +131,21 @@ pub fn apply_sinc_resample_kernel(
 
     // 获取波形形状
     let dims = waveform.dims();
-    let waveform_flat = waveform.reshape(((), dims[dims.len()-1]))?;
+    let waveform_flat = waveform.reshape(((), dims[dims.len() - 1]))?;
 
     let (num_wavs, length) = waveform_flat.dims2()?;
-    let padded_waveform = waveform.pad_with_zeros(D::Minus1, width as usize, (width+orig_freq) as usize)?;
+    let padded_waveform =
+        waveform.pad_with_zeros(D::Minus1, width as usize, (width + orig_freq) as usize)?;
 
     // 添加通道维度 [batch_size, 1, padded_length]
     let waveform_3d = padded_waveform.unsqueeze(1)?;
     let config = Conv1dConfig {
-            padding: 0,
+        padding: 0,
-            stride: orig_freq as usize,
+        stride: orig_freq as usize,
-            dilation: 1,
+        dilation: 1,
-            groups: 1,
+        groups: 1,
-            cudnn_fwd_algo: None,
+        cudnn_fwd_algo: None,
-        };
+    };
 
     let conv1d = Conv1d::new(kernel.clone(), None, config);
     // 执行卷积
@@ -153,16 +155,15 @@ pub fn apply_sinc_resample_kernel(
 
     // 转置并重塑 [batch_size, output_length * new_freq_reduced]
     let conv_transposed = conv_output.transpose(1, 2)?.reshape((num_wavs, ()))?;
-    
+
     // 计算目标长度
-    let target_length =
+    let target_length = ((new_freq as f64 * length as f64) / orig_freq as f64).ceil() as usize;
-        ((new_freq as f64 * length as f64) / orig_freq as f64).ceil() as usize;
 
     // 截取目标长度
     let resampled_flat =
         conv_transposed.narrow(1, 0, target_length.min(conv_transposed.dim(1)?))?;
     let mut new_dims = dims.to_vec();
-    let last_dim = new_dims.len()-1;
+    let last_dim = new_dims.len() - 1;
     new_dims[last_dim] = resampled_flat.dim(1)?;
     // 恢复原始批次形状
 
@@ -182,9 +183,7 @@ pub fn resample(
     beta: Option<f32>,
 ) -> Result<Tensor> {
     if orig_freq <= 0 || new_freq <= 0 {
-        return Err(anyhow!(
+        return Err(anyhow!("Frequencies must be positive".to_string(),));
-            "Frequencies must be positive".to_string(),
-        ));
     }
 
     if orig_freq == new_freq {
@@ -226,11 +225,27 @@ pub fn load_audio<P: AsRef<Path>>(path: P, device: Device) -> Result<(Tensor, us
     let spec = reader.spec();
     let samples: Vec<f32> = match spec.sample_format {
         SampleFormat::Int => {
-            // 将整数样本转换为浮点数 [-1.0, 1.0]
+            // 将整数样本转换为浮点数 [-1.0, 1.0]            
-            let max_value = match spec.bits_per_sample {
+            println!("spec.bits_per_sample: {}", spec.bits_per_sample);
-                8 => i8::MAX as f32,
+            let samples = match spec.bits_per_sample {
-                16 => i16::MAX as f32,
+                8 => {
-                24 => 8388607.0,
+                    reader
+                    .samples::<i8>()
+                    .map(|s| s.map(|sample| sample as f32 / i8::MAX as f32))
+                    .collect::<Result<Vec<_>, _>>()?
+                },
+                16 => {
+                    reader
+                    .samples::<i16>()
+                    .map(|s| s.map(|sample| sample as f32 / i16::MAX as f32))
+                    .collect::<Result<Vec<_>, _>>()?
+                },
+                24 => {
+                    reader
+                    .samples::<i32>()
+                    .map(|s| s.map(|sample| sample as f32 / 8388607.0))
+                    .collect::<Result<Vec<_>, _>>()?
+                },
                 _ => {
                     return Err(anyhow::anyhow!(
                         "Unsupported bit depth: {}",
@@ -238,10 +253,7 @@ pub fn load_audio<P: AsRef<Path>>(path: P, device: Device) -> Result<(Tensor, us
                     ));
                 }
             };
-            reader
+            samples
-                .samples::<i16>()
-                .map(|s| s.map(|sample| sample as f32 / max_value))
-                .collect::<Result<Vec<_>, _>>()?
         }
         SampleFormat::Float => {
             // 直接读取浮点数样本
@@ -258,6 +270,7 @@ pub fn load_audio<P: AsRef<Path>>(path: P, device: Device) -> Result<(Tensor, us
         &device,
     )?
     .t()?;
+    // println!("audio channels: {}", spec.channels);
     if spec.channels > 1 {
         // 对channel通道求平均， channel维度变为1
         audio_tensor = audio_tensor.mean_keepdim(0)?;
@@ -277,3 +290,25 @@ pub fn load_audio_with_resample<P: AsRef<Path>>(
     }
     Ok(audio)
 }
+
+pub fn save_wav(audio: &Tensor, save_path: &str) -> Result<()> {
+    let spec = hound::WavSpec {
+        channels: 1,
+        sample_rate: 16000,
+        bits_per_sample: 16,
+        sample_format: hound::SampleFormat::Int,
+    };
+    assert_eq!(audio.dim(0)?, 1, "audio channel must be 1");
+    let max = audio.abs()?.max_all()?;
+    let max = max.to_scalar::<f32>()?;
+    let ratio = if max > 1.0 { 32767.0 / max } else { 32767.0 };
+    let audio = audio.squeeze(0)?;
+    let audio_vec = audio.to_vec1::<f32>()?;
+    let mut writer = hound::WavWriter::create(save_path, spec).unwrap();
+    for i in audio_vec {
+        let sample_i16 = (i * ratio).round() as i16;
+        writer.write_sample(sample_i16).unwrap();
+    }
+    writer.finalize().unwrap();
+    Ok(())
+}

tests/test_voxcpm.rs

@@ -1,49 +1,81 @@
-use std::collections::HashMap;
 use anyhow::{Ok, Result};
+use std::collections::HashMap;
 
-use aha::{models::voxcpm::{audio_vae::AudioVAE, config::VoxCPMConfig, model::VoxCPMModel, tokenizer::SingleChineseTokenizer}, utils::utils::{find_type_files, get_device}};
+use aha::{
+    models::voxcpm::{
+        audio_vae::AudioVAE, config::VoxCPMConfig, model::VoxCPMModel,
+        tokenizer::SingleChineseTokenizer,
+    },
+    utils::{
+        audio_utils::save_wav,
+        utils::{find_type_files, get_device},
+    },
+};
 use candle_core::pickle::read_all_with_key;
 use candle_nn::VarBuilder;
 
-
 #[test]
 fn voxcpm_generate() -> Result<()> {
+    // RUST_BACKTRACE=1 cargo test -F cuda,flash-attn voxcpm_generate -- --nocapture
     let model_path = "/home/jhq/huggingface_model/openbmb/VoxCPM-0.5B/";
     let model_list = find_type_files(&model_path, "pth")?;
     println!(" pth model_list: {:?}", model_list);
     let dev = get_device(None);
     let mut dict_to_hashmap = HashMap::new();
     let mut dtype = candle_core::DType::F32;
-    for m in model_list {  
+    for m in model_list {
-        let dict = read_all_with_key(m, Some("state_dict"))?;        
+        let dict = read_all_with_key(m, Some("state_dict"))?;
         dtype = dict[0].1.dtype();
         for (k, v) in dict {
             // println!("key: {}, tensor shape: {:?}", k, v);
+            // if k.contains("decoder.model.2.block.1") {
+            //     println!("val: {}", v);
+            // }
             dict_to_hashmap.insert(k, v);
-        }        
+        }
     }
     let vb = VarBuilder::from_tensors(dict_to_hashmap, dtype, &dev);
-    let audio_vae = AudioVAE::new(vb, 128, vec![2, 5, 8, 8], Some(64), 1536, vec![8, 8, 5, 2], 16000)?;
+    let audio_vae = AudioVAE::new(
+        vb,
+        128,
+        vec![2, 5, 8, 8],
+        Some(64),
+        1536,
+        vec![8, 8, 5, 2],
+        16000,
+    )?;
     println!("audio vae load down");
     let model_list = find_type_files(&model_path, "bin")?;
     println!(" bin model_list: {:?}", model_list);
     dict_to_hashmap = HashMap::new();
-    for m in model_list {  
+    for m in model_list {
-        let dict = read_all_with_key(m, Some("state_dict"))?;        
+        let dict = read_all_with_key(m, Some("state_dict"))?;
         dtype = dict[0].1.dtype();
         for (k, v) in dict {
             // println!("key: {}, tensor shape: {:?}", k, v);
             dict_to_hashmap.insert(k, v);
-        }        
+        }
     }
     let vb_vox = VarBuilder::from_tensors(dict_to_hashmap, dtype, &dev);
     let config_path = model_path.to_string() + "/config.json";
     let config: VoxCPMConfig = serde_json::from_slice(&std::fs::read(config_path)?)?;
     let tokenizer = SingleChineseTokenizer::new(model_path)?;
     let mut voxcpm = VoxCPMModel::new(vb_vox, config, tokenizer, audio_vae)?;
-    let generate = voxcpm.generate("你好啊，这是初始测试语句".to_string(), None, None, 2, 30, 10, 2.0, false, 3, 6.0)?;
+    let generate = voxcpm.generate(
-    // let audio_path = "./assets/audio/example.wav";
+        "太阳当空照，花儿对我笑，小鸟说早早早".to_string(),
-    
+        Some("啥子小师叔，打狗还要看主人，你再要继续，我，就是你的对手".to_string()),
+        Some("./assets/audio/voice_01.wav".to_string()),
+        // Some("一定被灰太狼给吃了，我已经为他准备好了花圈了".to_string()),
+        // Some("./assets/audio/voice_05.wav".to_string()),
+        2,
+        100,
+        10,
+        2.0,
+        false,
+        3,
+        6.0,
+    )?;
+    let _ = save_wav(&generate, "voxcpm_init.wav")?;
     Ok(())
 }
 
@@ -54,4 +86,4 @@ fn voxcpm_tokenizer() -> Result<()> {
     let ids = tokenizer.encode("你好啊，你吃饭了吗".to_string())?;
     println!("ids: {:?}", ids);
     Ok(())
-}
+}

tests/weight_test.rs

@@ -41,4 +41,4 @@ fn voxcpm_weight() -> Result<()> {
     let contain_key = vb.contains_tensor("encoder.block.4.block.2.block.3.weight_g");
     println!("contain encoder.block.4.block.2.block.3.weight_g: {}", contain_key);
     Ok(())
-}
+}