update voxcpm

This commit is contained in:
jhqxxx
2025-10-10 20:36:52 +08:00
parent fb746842d7
commit 39536df06b
14 changed files with 381 additions and 122 deletions
+34 -51
View File
@@ -1,9 +1,9 @@
use anyhow::{Ok, Result};
use std::collections::HashMap;
use std::{collections::HashMap, time::Instant};
use aha::{
models::voxcpm::{
audio_vae::AudioVAE, config::VoxCPMConfig, model::VoxCPMModel,
audio_vae::AudioVAE, config::VoxCPMConfig, generate::VoxCPMGenerate, model::VoxCPMModel,
tokenizer::SingleChineseTokenizer,
},
utils::{
@@ -18,64 +18,47 @@ use candle_nn::VarBuilder;
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"))?;
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 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(),
)?;
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"))?;
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(
// 使用prompt_cache生成语音
let generate = voxcpm_generate.generate_use_prompt_cache(
"太阳当空照,花儿对我笑,小鸟说早早早".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")?;
let i_duration = i_start.elapsed();
println!("Time elapsed in generate is: {:?}", i_duration);
let _ = save_wav(&generate, "voxcpm.wav")?;
Ok(())
}