Merge branch 'voxcpm1.5'

README.md

@@ -19,6 +19,7 @@
 * DeepSeek-OCR - 深度求索光学文字识别模型
 * Hunyuan-OCR - 腾讯混元光学文字识别模型
 * PaddleOCR-VL - 百度飞桨光学文字识别模型
+* VoxCPM1.5 - 面壁智能语音生成模型1.5版本
 
 ## 计划支持
 我们持续扩展支持的模型列表，欢迎贡献！

src/models/voxcpm/audio_vae.rs

@@ -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>,
+    models: Vec<CausalDecoderBlock>,
-    model6: Snake1d,
+    model_minus_2: Snake1d,
-    model7: WNCausalConv1d,
+    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 idx = rates.len() + 2;
-        let model7 = WNCausalConv1d::new(vb.pp("model.7"), output_dim, d_out, 7, 1, 3, 1, 1)?;
+        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,
+            models,
-            model6,
+            model_minus_2,
-            model7,
+            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.model_minus_2.forward(&x)?;
-        let x = self.model7.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

@@ -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

@@ -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,
+            audio_config.encoder_dim,
-            vec![2, 5, 8, 8],
+            audio_config.encoder_rates.clone(),
-            Some(64),
+            Some(audio_config.latent_dim),
-            1536,
+            audio_config.decoder_dim,
-            vec![8, 8, 5, 2],
+            audio_config.decoder_rates.clone(),
-            16000,
+            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"))?;
+        let model_list = find_type_files(path, "bin")?;
-            for (k, v) in dict {
+        // voxcpm0.5B模型文件是.bin类型， voxcpm1.5模型文件是.safetensors类型
-                // println!("key: {}, tensor shape: {:?}", k, v);
+        let vb_voxcpm = if model_list.is_empty() {
-                dict_to_hashmap.insert(k, v);
+            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);
+                }
             }
-        }
+            VarBuilder::from_tensors(dict_to_hashmap, m_dtype, device)
-        // println!("model dtype: {:?}", m_dtype);
+        };
-        let vb_voxcpm = 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

@@ -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 dim0 = audio_feat.dim(0)? - 1;
-                let audio_feat = audio_feat.i(..dim0)?;
+                // 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 {
+        // let max_len = if retry_badcase {
-            (target_text_length as f64 * retry_badcase_ratio_threshold + 10.0) as usize
+        //     (target_text_length as f64 * retry_badcase_ratio_threshold + 10.0) as usize
-        } else {
+        // } else {
-            max_len
+        //     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

@@ -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

@@ -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,flash-attn voxcpm_config -r -- --nocapture
-    // cargo test -F cuda minicpm4_config -- --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

@@ -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("啥子小师叔，打狗还要看主人，你再要继续，我，就是你的对手".to_string()),
-        // Some("./assets/audio/voice_01.wav".to_string()),
+        Some("./assets/audio/voice_01.wav".to_string()),
-        Some("一定被灰太狼给吃了，我已经为他准备好了花圈了".to_string()),
+        // Some("一定被灰太狼给吃了，我已经为他准备好了花圈了".to_string()),
-        Some("./assets/audio/voice_05.wav".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

@@ -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

@@ -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/";