From 963924a1f11f32c0b5ba013098875ce11489a2fb Mon Sep 17 00:00:00 2001 From: jhqxxx Date: Thu, 11 Dec 2025 18:33:35 +0800 Subject: [PATCH 1/2] stash save --- src/models/voxcpm/config.rs | 11 ++++++ src/models/voxcpm/generate.rs | 31 +++++++++++------ tests/config_tests.rs | 14 ++++++-- tests/test_voxcpm1_5.rs | 65 +++++++++++++++++++++++++++++++++++ tests/weight_test.rs | 20 +++++++++++ 5 files changed, 129 insertions(+), 12 deletions(-) create mode 100644 tests/test_voxcpm1_5.rs diff --git a/src/models/voxcpm/config.rs b/src/models/voxcpm/config.rs index f04195c..17466f4 100644 --- a/src/models/voxcpm/config.rs +++ b/src/models/voxcpm/config.rs @@ -51,6 +51,16 @@ pub struct VoxCPMDitConfig { pub cfm_config: CfmConfig, } +#[derive(Debug, Clone, PartialEq, serde::Deserialize)] +pub struct AudioVaeConfig { + pub encoder_dim: usize, + pub encoder_rates: Vec, + pub latent_dim: usize, + pub decoder_dim: usize, + pub decoder_rates: Vec, + pub sample_rate: usize, +} + #[derive(Debug, Clone, PartialEq, serde::Deserialize)] pub struct VoxCPMConfig { pub lm_config: VoxMiniCPM4Config, @@ -61,6 +71,7 @@ pub struct VoxCPMConfig { pub residual_lm_num_layers: usize, pub encoder_config: VoxCPMEncoderConfig, pub dit_config: VoxCPMDitConfig, + pub audio_vae_config: Option, pub max_length: usize, pub dtype: String, } diff --git a/src/models/voxcpm/generate.rs b/src/models/voxcpm/generate.rs index 2416821..3a12bca 100644 --- a/src/models/voxcpm/generate.rs +++ b/src/models/voxcpm/generate.rs @@ -6,7 +6,7 @@ use candle_nn::VarBuilder; use crate::{ models::voxcpm::{ - audio_vae::AudioVAE, config::VoxCPMConfig, model::VoxCPMModel, + audio_vae::AudioVAE, config::{AudioVaeConfig, VoxCPMConfig}, model::VoxCPMModel, tokenizer::SingleChineseTokenizer, }, utils::{find_type_files, get_device, get_dtype}, @@ -20,7 +20,8 @@ pub struct VoxCPMGenerate { impl VoxCPMGenerate { pub fn init(path: &str, device: Option<&Device>, dtype: Option) -> Result { let device = &get_device(device); - + let config_path = path.to_string() + "/config.json"; + let config: VoxCPMConfig = serde_json::from_slice(&std::fs::read(config_path)?)?; let model_list = find_type_files(path, "pth")?; // println!(" pth model_list: {:?}", model_list); let mut dict_to_hashmap = HashMap::new(); @@ -34,21 +35,31 @@ impl VoxCPMGenerate { } } let vb_vae = VarBuilder::from_tensors(dict_to_hashmap, vae_dtype, device); + let audio_config = match config.audio_vae_config.clone() { + Some(config) => config, + None => AudioVaeConfig { + encoder_dim: 128, + encoder_rates: vec![2, 5, 8, 8], + latent_dim: 64, + decoder_dim: 1536, + decoder_rates: vec![8, 8, 5, 2], + sample_rate: 16000 + } + }; let audio_vae = AudioVAE::new( vb_vae, - 128, - vec![2, 5, 8, 8], - Some(64), - 1536, - vec![8, 8, 5, 2], - 16000, + audio_config.encoder_dim, + audio_config.encoder_rates.clone(), + Some(audio_config.latent_dim), + audio_config.decoder_dim, + audio_config.decoder_rates.clone(), + audio_config.sample_rate, )?; let model_list = find_type_files(path, "bin")?; // println!(" bin model_list: {:?}", model_list); dict_to_hashmap = HashMap::new(); - let config_path = path.to_string() + "/config.json"; - let config: VoxCPMConfig = serde_json::from_slice(&std::fs::read(config_path)?)?; + let cfg_dtype = config.dtype.as_str(); let m_dtype = get_dtype(dtype, cfg_dtype); for m in model_list { diff --git a/tests/config_tests.rs b/tests/config_tests.rs index 238b486..7b1dfcb 100644 --- a/tests/config_tests.rs +++ b/tests/config_tests.rs @@ -28,8 +28,8 @@ fn minicpm4_config() -> Result<()> { #[test] fn voxcpm_config() -> Result<()> { - // cargo test -F cuda,flash-attn minicpm4_config -r -- --nocapture - // cargo test -F cuda minicpm4_config -- --nocapture + // cargo test -F cuda,flash-attn voxcpm_config -r -- --nocapture + // cargo test -F cuda voxcpm_config -r -- --nocapture let model_path = "/home/jhq/huggingface_model/openbmb/VoxCPM-0.5B/"; let config_path = model_path.to_string() + "/config.json"; let config: VoxCPMConfig = serde_json::from_slice(&std::fs::read(config_path)?)?; @@ -37,6 +37,16 @@ fn voxcpm_config() -> Result<()> { Ok(()) } +#[test] +fn voxcpm1_5_config() -> Result<()> { + // cargo test -F cuda voxcpm1_5_config -r -- --nocapture + let model_path = "/home/jhq/huggingface_model/OpenBMB/VoxCPM1.5/"; + let config_path = model_path.to_string() + "/config.json"; + let config: VoxCPMConfig = serde_json::from_slice(&std::fs::read(config_path)?)?; + println!("{:?}", config); + Ok(()) +} + #[test] fn qwen3vl_config() -> Result<()> { // cargo test -F cuda qwen3vl_config -r -- --nocapture diff --git a/tests/test_voxcpm1_5.rs b/tests/test_voxcpm1_5.rs new file mode 100644 index 0000000..c3044da --- /dev/null +++ b/tests/test_voxcpm1_5.rs @@ -0,0 +1,65 @@ +use std::time::Instant; + +use aha::{ + models::voxcpm::{generate::VoxCPMGenerate, tokenizer::SingleChineseTokenizer}, + utils::audio_utils::save_wav, +}; +use anyhow::{Ok, Result}; + +#[test] +fn voxcpm1_5_generate() -> Result<()> { + // RUST_BACKTRACE=1 cargo test -F cuda voxcpm1_5_generate -r -- --nocapture + let model_path = "/home/jhq/huggingface_model/OpenBMB/VoxCPM1.5/"; + + let i_start = Instant::now(); + let mut voxcpm_generate = VoxCPMGenerate::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_simple("太阳当空照，花儿对我笑，小鸟说早早早".to_string())?; + let generate = voxcpm_generate.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, + 6.0, + )?; + + // 创建prompt_cache + // let _ = voxcpm_generate.build_prompt_cache( + // "啥子小师叔，打狗还要看主人，你再要继续，我，就是你的对手".to_string(), + // "./assets/audio/voice_01.wav".to_string(), + // )?; + // // 使用prompt_cache生成语音 + // let generate = voxcpm_generate.generate_use_prompt_cache( + // "太阳当空照，花儿对我笑，小鸟说早早早".to_string(), + // 2, + // 100, + // 10, + // 2.0, + // false, + // 6.0, + // )?; + + let i_duration = i_start.elapsed(); + println!("Time elapsed in generate is: {:?}", i_duration); + save_wav(&generate, "voxcpm.wav")?; + Ok(()) +} + +#[test] +fn voxcpm1_5_tokenizer() -> Result<()> { + // RUST_BACKTRACE=1 cargo test -F cuda voxcpm1_5_tokenizer -r -- --nocapture + let model_path = "/home/jhq/huggingface_model/OpenBMB/VoxCPM1.5/"; + let tokenizer = SingleChineseTokenizer::new(model_path)?; + let ids = tokenizer.encode("你好啊，你吃饭了吗".to_string())?; + println!("ids: {:?}", ids); + Ok(()) +} diff --git a/tests/weight_test.rs b/tests/weight_test.rs index 32e3736..71dc0d5 100644 --- a/tests/weight_test.rs +++ b/tests/weight_test.rs @@ -46,6 +46,26 @@ fn voxcpm_weight() -> Result<()> { Ok(()) } +#[test] +fn voxcpm1_5_weight() -> Result<()> { + let model_path = "/home/jhq/huggingface_model/OpenBMB/VoxCPM1.5/"; + let model_list = find_type_files(model_path, "pth")?; + println!("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"))?; + dtype = dict[0].1.dtype(); + for (k, v) in dict { + println!("key: {}, tensor shape: {:?}", k, v); + dict_to_hashmap.insert(k, v); + } + } + + Ok(()) +} + #[test] fn qwen3vl_weight() -> Result<()> { let model_path = "/home/jhq/huggingface_model/Qwen/Qwen3-VL-4B-Instruct/"; From b9947a65169d0eb75e810565adb71f16d287a140 Mon Sep 17 00:00:00 2001 From: jhqxxx Date: Thu, 11 Dec 2025 23:14:43 +0800 Subject: [PATCH 2/2] update voxcpm code to support voxcpm1.5 model --- README.md | 1 + src/models/voxcpm/audio_vae.rs | 41 ++++++++++++++++++---------------- src/models/voxcpm/generate.rs | 30 +++++++++++++++---------- src/models/voxcpm/model.rs | 23 +++++++++++-------- src/utils/audio_utils.rs | 4 ++-- tests/test_voxcpm.rs | 10 ++++----- tests/test_voxcpm1_5.rs | 12 +++++----- 7 files changed, 68 insertions(+), 53 deletions(-) diff --git a/README.md b/README.md index aee2d83..ee16e79 100644 --- a/README.md +++ b/README.md @@ -19,6 +19,7 @@ * DeepSeek-OCR - 深度求索光学文字识别模型 * Hunyuan-OCR - 腾讯混元光学文字识别模型 * PaddleOCR-VL - 百度飞桨光学文字识别模型 +* VoxCPM1.5 - 面壁智能语音生成模型1.5版本 ## 计划支持 我们持续扩展支持的模型列表，欢迎贡献！ diff --git a/src/models/voxcpm/audio_vae.rs b/src/models/voxcpm/audio_vae.rs index 800a8cd..c7b87e6 100644 --- a/src/models/voxcpm/audio_vae.rs +++ b/src/models/voxcpm/audio_vae.rs @@ -281,7 +281,7 @@ impl CausalEncoderBlock { pub struct CausalEncoder { block0: WNCausalConv1d, - block1_4: Vec, + blocks: Vec, fc_mu: WNCausalConv1d, fc_logvar: WNCausalConv1d, } @@ -298,19 +298,19 @@ impl CausalEncoder { let mut groups; let block0 = WNCausalConv1d::new(vb.pp("block.0"), 1, d_model, 7, 1, 3, 1, 1)?; let vb_block = vb.pp("block"); - let mut block1_4 = Vec::new(); + let mut blocks = Vec::new(); for (i, stride) in strides.iter().enumerate() { d_model *= 2; groups = if depthwise { d_model / 2 } else { 1 }; let block_i = CausalEncoderBlock::new(vb_block.pp(i + 1), None, d_model, *stride, groups)?; - block1_4.push(block_i); + blocks.push(block_i); } let fc_mu = WNCausalConv1d::new(vb.pp("fc_mu"), d_model, laten_dim, 3, 1, 1, 1, 1)?; let fc_logvar = WNCausalConv1d::new(vb.pp("fc_logvar"), d_model, laten_dim, 3, 1, 1, 1, 1)?; Ok(Self { block0, - block1_4, + blocks, fc_mu, fc_logvar, }) @@ -318,7 +318,7 @@ impl CausalEncoder { pub fn forward(&self, x: &Tensor) -> Result<(Tensor, Tensor, Tensor)> { let mut hidden_state = self.block0.forward(x)?; - for block_i in &self.block1_4 { + for block_i in &self.blocks { hidden_state = block_i.forward(&hidden_state)?; } let mu = self.fc_mu.forward(&hidden_state)?; @@ -401,9 +401,9 @@ impl CausalDecoderBlock { pub struct CausalDecoder { model0: WNCausalConv1d, model1: WNCausalConv1d, - model2_5: Vec, - model6: Snake1d, - model7: WNCausalConv1d, + models: Vec, + model_minus_2: Snake1d, + model_minus_1: WNCausalConv1d, } impl CausalDecoder { @@ -413,6 +413,7 @@ impl CausalDecoder { channels: usize, rates: Vec, d_out: usize, + depthwise: bool, ) -> Result { let model0 = WNCausalConv1d::new( vb.pp("model.0"), @@ -427,11 +428,11 @@ impl CausalDecoder { 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(); + let mut models = Vec::new(); for (i, stride) in rates.iter().enumerate() { let input_dim = channels / 2_usize.pow(i as u32); output_dim = channels / 2_usize.pow((i + 1) as u32); - let groups = output_dim; + let groups = if depthwise { output_dim } else { 1 }; let model_i = CausalDecoderBlock::new( vb_model.pp(i + 2), input_dim, @@ -439,27 +440,28 @@ impl CausalDecoder { *stride, groups, )?; - model2_5.push(model_i); + models.push(model_i); } - let model6 = Snake1d::new(vb.pp("model.6"), output_dim)?; - let model7 = WNCausalConv1d::new(vb.pp("model.7"), output_dim, d_out, 7, 1, 3, 1, 1)?; + let idx = rates.len() + 2; + let model_minus_2 = Snake1d::new(vb_model.pp(idx), output_dim)?; + let model_minus_1 = WNCausalConv1d::new(vb_model.pp(idx+1), output_dim, d_out, 7, 1, 3, 1, 1)?; Ok(Self { model0, model1, - model2_5, - model6, - model7, + models, + model_minus_2, + model_minus_1, }) } 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 { + for model_i in &self.models { x = model_i.forward(&x)?; } - let x = self.model6.forward(&x)?; - let x = self.model7.forward(&x)?; + let x = self.model_minus_2.forward(&x)?; + let x = self.model_minus_1.forward(&x)?; let x = x.tanh()?; Ok(x) } @@ -506,6 +508,7 @@ impl AudioVAE { decoder_dim, decoder_rates.clone(), 1, + true, )?; let chunk_size = hop_length; Ok(Self { diff --git a/src/models/voxcpm/generate.rs b/src/models/voxcpm/generate.rs index 3a12bca..9dfee8a 100644 --- a/src/models/voxcpm/generate.rs +++ b/src/models/voxcpm/generate.rs @@ -56,21 +56,27 @@ impl VoxCPMGenerate { audio_config.sample_rate, )?; - let model_list = find_type_files(path, "bin")?; - // println!(" bin model_list: {:?}", model_list); - dict_to_hashmap = HashMap::new(); - let cfg_dtype = config.dtype.as_str(); let m_dtype = get_dtype(dtype, cfg_dtype); - for m in model_list { - let dict = read_all_with_key(m, Some("state_dict"))?; - for (k, v) in dict { - // println!("key: {}, tensor shape: {:?}", k, v); - dict_to_hashmap.insert(k, v); + + let model_list = find_type_files(path, "bin")?; + // voxcpm0.5B模型文件是.bin类型， voxcpm1.5模型文件是.safetensors类型 + let vb_voxcpm = if model_list.is_empty() { + let model_list = find_type_files(path, "safetensors")?; + unsafe { VarBuilder::from_mmaped_safetensors(&model_list, m_dtype, &device)? } + } else { + dict_to_hashmap = HashMap::new(); + let cfg_dtype = config.dtype.as_str(); + let m_dtype = get_dtype(dtype, cfg_dtype); + for m in model_list { + let dict = read_all_with_key(m, Some("state_dict"))?; + for (k, v) in dict { + // println!("key: {}, tensor shape: {:?}", k, v); + dict_to_hashmap.insert(k, v); + } } - } - // println!("model dtype: {:?}", m_dtype); - let vb_voxcpm = VarBuilder::from_tensors(dict_to_hashmap, m_dtype, device); + VarBuilder::from_tensors(dict_to_hashmap, m_dtype, device) + }; let tokenizer = SingleChineseTokenizer::new(path)?; let voxcpm = VoxCPMModel::new(vb_voxcpm, config, tokenizer, audio_vae)?; diff --git a/src/models/voxcpm/model.rs b/src/models/voxcpm/model.rs index 3e27c64..8f4e474 100644 --- a/src/models/voxcpm/model.rs +++ b/src/models/voxcpm/model.rs @@ -274,7 +274,8 @@ impl UnifiedCFM { let mut x = x.clone(); for step in 1..t_span_len { if use_cfg_zero_star && step <= zero_init_steps { - dphi_dt = Tensor::zeros(1, t_span.dtype(), t_span.device())?; + // dphi_dt = Tensor::zeros(1, t_span.dtype(), t_span.device())?; + dphi_dt = x.zeros_like()?; } else { let b = x.dim(0)?; // let x_in = Tensor::zeros((2*b, self.in_channels, x.dim(2)?), x.dtype(), x.device())?; @@ -517,16 +518,18 @@ impl VoxCPMModel { if audio.dim(1)? % patch_len != 0 { audio = audio.pad_with_zeros( D::Minus1, - 0, + // 0, + // patch_len - audio.dim(1)? % patch_len, patch_len - audio.dim(1)? % patch_len, + 0, )?; } let audio_feat = self.audio_vae.encode(&audio, Some(self.sample_rate))?; let audio_feat = audio_feat .reshape((self.audio_vae.latent_dim, (), self.patch_size))? .permute((1, 2, 0))?; - let dim0 = audio_feat.dim(0)? - 1; - let audio_feat = audio_feat.i(..dim0)?; + // let dim0 = audio_feat.dim(0)? - 1; + // let audio_feat = audio_feat.i(..dim0)?; 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)?; @@ -554,11 +557,13 @@ impl VoxCPMModel { } }; let target_text_length = self.tokenizer.encode(target_text)?.len(); - let max_len = if retry_badcase { - (target_text_length as f64 * retry_badcase_ratio_threshold + 10.0) as usize - } else { - max_len - }; + // let max_len = if retry_badcase { + // (target_text_length as f64 * retry_badcase_ratio_threshold + 10.0) as usize + // } else { + // max_len + // }; + let max_len = max_len + .min((target_text_length as f64 * retry_badcase_ratio_threshold + 10.0) as usize); let decode_audio = self._generate( &text_token, &text_mask, diff --git a/src/utils/audio_utils.rs b/src/utils/audio_utils.rs index adffb0b..a0ac671 100644 --- a/src/utils/audio_utils.rs +++ b/src/utils/audio_utils.rs @@ -278,10 +278,10 @@ pub fn load_audio_with_resample>( Ok(audio) } -pub fn save_wav(audio: &Tensor, save_path: &str) -> Result<()> { +pub fn save_wav(audio: &Tensor, save_path: &str, sample_rate: u32) -> Result<()> { let spec = hound::WavSpec { channels: 1, - sample_rate: 16000, + sample_rate, bits_per_sample: 16, sample_format: hound::SampleFormat::Int, }; diff --git a/tests/test_voxcpm.rs b/tests/test_voxcpm.rs index a3fd2a7..ab37ae3 100644 --- a/tests/test_voxcpm.rs +++ b/tests/test_voxcpm.rs @@ -20,10 +20,10 @@ fn voxcpm_generate() -> Result<()> { // let generate = voxcpm_generate.generate_simple("太阳当空照，花儿对我笑，小鸟说早早早".to_string())?; let generate = voxcpm_generate.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()), + 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, @@ -50,7 +50,7 @@ fn voxcpm_generate() -> Result<()> { let i_duration = i_start.elapsed(); println!("Time elapsed in generate is: {:?}", i_duration); - save_wav(&generate, "voxcpm.wav")?; + save_wav(&generate, "voxcpm.wav", 16000)?; Ok(()) } diff --git a/tests/test_voxcpm1_5.rs b/tests/test_voxcpm1_5.rs index c3044da..fd0a6f9 100644 --- a/tests/test_voxcpm1_5.rs +++ b/tests/test_voxcpm1_5.rs @@ -20,12 +20,12 @@ fn voxcpm1_5_generate() -> Result<()> { // let generate = voxcpm_generate.generate_simple("太阳当空照，花儿对我笑，小鸟说早早早".to_string())?; let generate = voxcpm_generate.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()), + Some("啥子小师叔，打狗还要看主人，你再要继续，我就是你的对手".to_string()), + Some("./assets/audio/voice_01.wav".to_string()), + // Some("一定被灰太狼给吃了，我已经为他准备好了花圈了".to_string()), + // Some("./assets/audio/voice_05.wav".to_string()), 2, - 100, + 4096, 10, 2.0, false, @@ -50,7 +50,7 @@ fn voxcpm1_5_generate() -> Result<()> { let i_duration = i_start.elapsed(); println!("Time elapsed in generate is: {:?}", i_duration); - save_wav(&generate, "voxcpm.wav")?; + save_wav(&generate, "voxcpm1_5.wav", 44100)?; Ok(()) }