From 963924a1f11f32c0b5ba013098875ce11489a2fb Mon Sep 17 00:00:00 2001 From: jhqxxx Date: Thu, 11 Dec 2025 18:33:35 +0800 Subject: [PATCH] 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/";