update voxcpm
This commit is contained in:
+34
-51
@@ -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(())
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user