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

update voxcpm code to support voxcpm1.5 model

Browse Source
This commit is contained in:
jhqxxx
2025-12-11 23:14:43 +08:00
parent 963924a1f1
commit b9947a6516
7 changed files with 68 additions and 53 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/generate.rs
+18 -12
View File
@@ -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)?;
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/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
+6 -6
View File
@@ -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(())
}