diff --git a/README.md b/README.md index aee2d83..ee16e79 100644 --- a/README.md +++ b/README.md @@ -19,6 +19,7 @@ * DeepSeek-OCR - 深度求索光学文字识别模型 * Hunyuan-OCR - 腾讯混元光学文字识别模型 * PaddleOCR-VL - 百度飞桨光学文字识别模型 +* VoxCPM1.5 - 面壁智能语音生成模型1.5版本 ## 计划支持 我们持续扩展支持的模型列表，欢迎贡献！ diff --git a/src/models/voxcpm/audio_vae.rs b/src/models/voxcpm/audio_vae.rs index 800a8cd..c7b87e6 100644 --- a/src/models/voxcpm/audio_vae.rs +++ b/src/models/voxcpm/audio_vae.rs @@ -281,7 +281,7 @@ impl CausalEncoderBlock { pub struct CausalEncoder { block0: WNCausalConv1d, - block1_4: Vec, + blocks: Vec, 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, - model6: Snake1d, - model7: WNCausalConv1d, + models: Vec, + model_minus_2: Snake1d, + model_minus_1: WNCausalConv1d, } impl CausalDecoder { @@ -413,6 +413,7 @@ impl CausalDecoder { channels: usize, rates: Vec, d_out: usize, + depthwise: bool, ) -> Result { 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 { 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 { diff --git a/src/models/voxcpm/generate.rs b/src/models/voxcpm/generate.rs index 3a12bca..9dfee8a 100644 --- a/src/models/voxcpm/generate.rs +++ b/src/models/voxcpm/generate.rs @@ -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)?; diff --git a/src/models/voxcpm/model.rs b/src/models/voxcpm/model.rs index 3e27c64..8f4e474 100644 --- a/src/models/voxcpm/model.rs +++ b/src/models/voxcpm/model.rs @@ -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, diff --git a/src/utils/audio_utils.rs b/src/utils/audio_utils.rs index adffb0b..a0ac671 100644 --- a/src/utils/audio_utils.rs +++ b/src/utils/audio_utils.rs @@ -278,10 +278,10 @@ pub fn load_audio_with_resample>( 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, }; diff --git a/tests/test_voxcpm.rs b/tests/test_voxcpm.rs index a3fd2a7..ab37ae3 100644 --- a/tests/test_voxcpm.rs +++ b/tests/test_voxcpm.rs @@ -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(()) } diff --git a/tests/test_voxcpm1_5.rs b/tests/test_voxcpm1_5.rs index c3044da..fd0a6f9 100644 --- a/tests/test_voxcpm1_5.rs +++ b/tests/test_voxcpm1_5.rs @@ -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(()) }