tckm.yan/aha
0
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge branch 'voxcpm1.5'

Browse Source
This commit is contained in:
jhqxxx
2025-12-11 23:15:05 +08:00
parent 6430338e95 b9947a6516
commit d282a0cf40
10 changed files with 190 additions and 58 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+1
View File
@@ -19,6 +19,7 @@
* DeepSeek-OCR - 深度求索光学文字识别模型
* Hunyuan-OCR - 腾讯混元光学文字识别模型
* PaddleOCR-VL - 百度飞桨光学文字识别模型
* VoxCPM1.5 - 面壁智能语音生成模型1.5版本
## 计划支持
我们持续扩展支持的模型列表，欢迎贡献！
src/models/voxcpm/audio_vae.rs
+22 -19
View File
@@ -281,7 +281,7 @@ impl CausalEncoderBlock {
pub struct CausalEncoder {
block0: WNCausalConv1d,
block1_4: Vec<CausalEncoderBlock>,
blocks: Vec<CausalEncoderBlock>,
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<CausalDecoderBlock>,
model6: Snake1d,
model7: WNCausalConv1d,
models: Vec<CausalDecoderBlock>,
model_minus_2: Snake1d,
model_minus_1: WNCausalConv1d,
}
impl CausalDecoder {
@@ -413,6 +413,7 @@ impl CausalDecoder {
channels: usize,
rates: Vec<usize>,
d_out: usize,
depthwise: bool,
) -> Result<Self> {
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<Tensor> {
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 {
src/models/voxcpm/config.rs
+11
View File
@@ -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<usize>,
pub latent_dim: usize,
pub decoder_dim: usize,
pub decoder_rates: Vec<usize>,
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<AudioVaeConfig>,
pub max_length: usize,
pub dtype: String,
}
src/models/voxcpm/generate.rs
+38 -21
View File
@@ -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<DType>) -> Result<Self> {
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,32 +35,48 @@ 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 {
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)?;
src/models/voxcpm/model.rs
+14 -9
View File
@@ -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,
src/utils/audio_utils.rs
+2 -2
View File
@@ -278,10 +278,10 @@ pub fn load_audio_with_resample<P: AsRef<Path>>(
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,
};
tests/config_tests.rs
+12 -2
View File
@@ -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
tests/test_voxcpm.rs
+5 -5
View File
@@ -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(())
}
tests/test_voxcpm1_5.rs
+65
View File
@@ -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,
4096,
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, "voxcpm1_5.wav", 44100)?;
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(())
}
tests/weight_test.rs
+20
View File
@@ -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/";