diff --git a/Cargo.lock b/Cargo.lock index 87ba2e6..5f23147 100644 --- a/Cargo.lock +++ b/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" diff --git a/Cargo.toml b/Cargo.toml index 05727ab..990f0c1 100644 --- a/Cargo.toml +++ b/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"] diff --git a/assets/audio/example.wav b/assets/audio/example.wav deleted file mode 100644 index dcc2b46..0000000 Binary files a/assets/audio/example.wav and /dev/null differ diff --git a/assets/audio/voice_01.wav b/assets/audio/voice_01.wav new file mode 100644 index 0000000..224a4e9 Binary files /dev/null and b/assets/audio/voice_01.wav differ diff --git a/assets/audio/voice_05.wav b/assets/audio/voice_05.wav new file mode 100644 index 0000000..2dcc49d Binary files /dev/null and b/assets/audio/voice_05.wav differ diff --git a/assets/audio/voice_06.wav b/assets/audio/voice_06.wav deleted file mode 100644 index 7444036..0000000 Binary files a/assets/audio/voice_06.wav and /dev/null differ diff --git a/src/models/voxcpm/audio_vae.rs b/src/models/voxcpm/audio_vae.rs index 1c6c918..0c478a4 100644 --- a/src/models/voxcpm/audio_vae.rs +++ b/src/models/voxcpm/audio_vae.rs @@ -1,7 +1,7 @@ -use anyhow::{Error, Ok, Result}; -use candle_core::{D, IndexOp, Tensor}; +use anyhow::{Ok, Result}; +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, - // in_c: usize, - // out_c: usize, - // kernel_size: usize, padding: usize, dilation: usize, groups: usize, @@ -59,7 +55,7 @@ impl CausalConvTranspose1d { ) -> Result { 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 { - 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 { - 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 { - 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 { - 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 { - 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 { - print!("audio_vae decoder input x shape: {:?}", x); - let x = self.model0.forward(x)?; - print!("audio_vae decoder model0 x shape: {:?}", x); + pub fn forward(&self, x: &Tensor) -> Result { + let x = self.model0.forward(x)?; let mut x = self.model1.forward(&x)?; for model_i in &self.model2_5 { x = model_i.forward(&x)?; diff --git a/src/models/voxcpm/model.rs b/src/models/voxcpm/model.rs index 08a48ab..cfd695e 100644 --- a/src/models/voxcpm/model.rs +++ b/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 { - // 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 { 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)?; - println!("audio_feat: {:?}", audio_feat); + let dim0 = audio_feat.dim(0)? - 1; 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 { - 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::()?; - 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) } } diff --git a/src/models/voxcpm/tokenizer.rs b/src/models/voxcpm/tokenizer.rs index ca0ee8c..0c913d3 100644 --- a/src/models/voxcpm/tokenizer.rs +++ b/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 = split_character .iter() .filter_map(|c| self.tokenizer.token_to_id(c)) diff --git a/src/utils/audio_utils.rs b/src/utils/audio_utils.rs index 40c5cd3..8f16e26 100644 --- a/src/utils/audio_utils.rs +++ b/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::()?; - let window_val:Vec = sqrt_arg_vec.iter().map(|x| i0(beta_val * x) / i0_beta).collect(); + let window_val: Vec = 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, - stride: orig_freq as usize, - dilation: 1, - groups: 1, - cudnn_fwd_algo: None, - }; + padding: 0, + stride: orig_freq as usize, + dilation: 1, + groups: 1, + 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 = - ((new_freq as f64 * length as f64) / orig_freq as f64).ceil() as usize; + let target_length = ((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, ) -> Result { if orig_freq <= 0 || new_freq <= 0 { - return Err(anyhow!( - "Frequencies must be positive".to_string(), - )); + return Err(anyhow!("Frequencies must be positive".to_string(),)); } if orig_freq == new_freq { @@ -226,11 +225,27 @@ pub fn load_audio>(path: P, device: Device) -> Result<(Tensor, us let spec = reader.spec(); let samples: Vec = match spec.sample_format { SampleFormat::Int => { - // 将整数样本转换为浮点数 [-1.0, 1.0] - let max_value = match spec.bits_per_sample { - 8 => i8::MAX as f32, - 16 => i16::MAX as f32, - 24 => 8388607.0, + // 将整数样本转换为浮点数 [-1.0, 1.0] + println!("spec.bits_per_sample: {}", spec.bits_per_sample); + let samples = match spec.bits_per_sample { + 8 => { + reader + .samples::() + .map(|s| s.map(|sample| sample as f32 / i8::MAX as f32)) + .collect::, _>>()? + }, + 16 => { + reader + .samples::() + .map(|s| s.map(|sample| sample as f32 / i16::MAX as f32)) + .collect::, _>>()? + }, + 24 => { + reader + .samples::() + .map(|s| s.map(|sample| sample as f32 / 8388607.0)) + .collect::, _>>()? + }, _ => { return Err(anyhow::anyhow!( "Unsupported bit depth: {}", @@ -238,10 +253,7 @@ pub fn load_audio>(path: P, device: Device) -> Result<(Tensor, us )); } }; - reader - .samples::() - .map(|s| s.map(|sample| sample as f32 / max_value)) - .collect::, _>>()? + samples } SampleFormat::Float => { // 直接读取浮点数样本 @@ -258,6 +270,7 @@ pub fn load_audio>(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>( } 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::()?; + let ratio = if max > 1.0 { 32767.0 / max } else { 32767.0 }; + let audio = audio.squeeze(0)?; + let audio_vec = audio.to_vec1::()?; + 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(()) +} diff --git a/tests/test_voxcpm.rs b/tests/test_voxcpm.rs index af8f617..93bd724 100644 --- a/tests/test_voxcpm.rs +++ b/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 { - let dict = read_all_with_key(m, Some("state_dict"))?; + for m in model_list { + 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 { - let dict = read_all_with_key(m, Some("state_dict"))?; + for m in model_list { + 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 audio_path = "./assets/audio/example.wav"; - + let generate = voxcpm.generate( + "太阳当空照，花儿对我笑，小鸟说早早早".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(()) -} \ No newline at end of file +} diff --git a/tests/weight_test.rs b/tests/weight_test.rs index 117d2fa..462fbdc 100644 --- a/tests/weight_test.rs +++ b/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(()) -} \ No newline at end of file +} diff --git a/voxcpm_init.wav b/voxcpm_init.wav new file mode 100644 index 0000000..3669a44 Binary files /dev/null and b/voxcpm_init.wav differ