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

Merge branch 'pr/ariesdevil/1' into fmt

Browse Source
This commit is contained in:
jhqxxx
2025-10-15 23:24:51 +08:00
parent b9dcaec3e6 9b9a8f2c73
commit a5e463ec61
40 changed files with 873 additions and 836 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/models/voxcpm/audio_vae.rs
+12 -18
View File
@@ -1,7 +1,6 @@
use anyhow::{Ok, Result};
use candle_core::{D, Tensor};
use candle_nn::{Conv1d, Conv1dConfig, ConvTranspose1d, ConvTranspose1dConfig, Module, VarBuilder};
use std::{result::Result::Ok as StdOk};
pub struct CausalConv1d {
conv1d: Conv1d,
@@ -65,7 +64,7 @@ impl CausalConvTranspose1d {
groups,
};
let conv_transpose1d = ConvTranspose1d::new(weight, bias, config.clone());
let conv_transpose1d = ConvTranspose1d::new(weight, bias, config);
Ok(Self {
conv_transpose1d,
padding,
@@ -95,13 +94,10 @@ impl WNCausalConv1d {
groups: usize,
stride: usize,
) -> Result<Self> {
let in_c = in_c / groups;
let in_c = in_c / groups;
let weight_g = vb.get((out_c, 1, 1), "weight_g")?;
let weight_v = vb.get((out_c, in_c, kernel_size), "weight_v")?;
let bias = match vb.get(out_c, "bias") {
StdOk(b) => Some(b),
Err(_) => None,
};
let bias = vb.get(out_c, "bias").ok();
let weight_norm = weight_v.sqr()?.sum_keepdim(1)?.sum_keepdim(2)?.sqrt()?;
let normalized_weight = weight_v.broadcast_div(&weight_norm)?;
let scaled_weight = normalized_weight.broadcast_mul(&weight_g)?;
@@ -133,10 +129,7 @@ impl WNCausalConvTranspose1d {
let in_c = in_c / groups;
let weight_g = vb.get((in_c, 1, 1), "weight_g")?;
let weight_v = vb.get((in_c, out_c, kernel_size), "weight_v")?;
let bias = match vb.get(out_c, "bias") {
StdOk(b) => Some(b),
Err(_) => None,
};
let bias = vb.get(out_c, "bias").ok();
let weight_norm = weight_v.sqr()?.sum_keepdim(1)?.sum_keepdim(2)?.sqrt()?;
let normalized_weight = weight_v.broadcast_div(&weight_norm)?;
let scaled_weight = normalized_weight.broadcast_mul(&weight_g)?;
@@ -302,14 +295,15 @@ impl CausalEncoder {
depthwise: bool,
) -> Result<Self> {
let mut d_model = d_model;
let mut groups = 1;
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();
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)?;
let block_i =
CausalEncoderBlock::new(vb_block.pp(i + 1), None, d_model, *stride, groups)?;
block1_4.push(block_i);
}
let fc_mu = WNCausalConv1d::new(vb.pp("fc_mu"), d_model, laten_dim, 3, 1, 1, 1, 1)?;
@@ -458,8 +452,8 @@ impl CausalDecoder {
})
}
pub fn forward(&self, x: &Tensor) -> Result<Tensor> {
let x = self.model0.forward(x)?;
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 {
x = model_i.forward(&x)?;
@@ -528,10 +522,10 @@ impl AudioVAE {
})
}
pub fn preprocess(&self, audio_data: &Tensor, sample_rate: Option<usize>) -> Result<Tensor>{
pub fn preprocess(&self, audio_data: &Tensor, sample_rate: Option<usize>) -> Result<Tensor> {
let sample_rate = match sample_rate {
Some(r) => r,
None => self.sample_rate
None => self.sample_rate,
};
assert_eq!(sample_rate, self.sample_rate);
let pad_to = self.hop_length;
@@ -549,7 +543,7 @@ impl AudioVAE {
pub fn encode(&self, audio_data: &Tensor, sample_rate: Option<usize>) -> Result<Tensor> {
let audio_data = match audio_data.rank() {
2 => audio_data.unsqueeze(1)?,
_ => audio_data.clone()
_ => audio_data.clone(),
};
let audio_data = self.preprocess(&audio_data, sample_rate)?;
let (_, mu, _) = self.encoder.forward(&audio_data)?;
src/models/voxcpm/config.rs
+2 -3
View File
@@ -1,4 +1,3 @@
#[derive(Debug, Clone, PartialEq, serde::Deserialize)]
pub struct VoxRopeScalingConfig {
pub r#type: String,
@@ -21,7 +20,7 @@ pub struct VoxMiniCPM4Config {
pub rope_theta: f32,
pub rope_scaling: VoxRopeScalingConfig,
pub vocab_size: usize,
pub scale_emb:f32,
pub scale_emb: f32,
pub dim_model_base: usize,
pub scale_depth: f32,
pub use_mup: bool,
@@ -64,4 +63,4 @@ pub struct VoxCPMConfig {
pub dit_config: VoxCPMDitConfig,
pub max_length: usize,
pub dtype: String,
}
}
src/models/voxcpm/generate.rs
+8 -7
View File
@@ -1,15 +1,16 @@
use std::collections::HashMap;
use anyhow::{Ok, Result};
use candle_core::{DType, Device, Tensor, pickle::read_all_with_key};
use candle_nn::VarBuilder;
use crate::{
models::voxcpm::{
audio_vae::AudioVAE, config::VoxCPMConfig, model::VoxCPMModel,
tokenizer::SingleChineseTokenizer,
},
utils::utils::{find_type_files, get_device, get_dtype},
utils::{find_type_files, get_device, get_dtype},
};
use anyhow::{Ok, Result};
use candle_core::{DType, Device, Tensor, pickle::read_all_with_key};
use candle_nn::VarBuilder;
pub struct VoxCPMGenerate {
voxcpm: VoxCPMModel,
@@ -19,7 +20,7 @@ pub struct VoxCPMGenerate {
impl VoxCPMGenerate {
pub fn init(path: &str, device: Option<&Device>, dtype: Option<DType>) -> Result<Self> {
let device = &get_device(device);
let model_list = find_type_files(path, "pth")?;
// println!(" pth model_list: {:?}", model_list);
let mut dict_to_hashmap = HashMap::new();
@@ -32,7 +33,7 @@ impl VoxCPMGenerate {
dict_to_hashmap.insert(k, v);
}
}
let vb_vae = VarBuilder::from_tensors(dict_to_hashmap, vae_dtype, &device);
let vb_vae = VarBuilder::from_tensors(dict_to_hashmap, vae_dtype, device);
let audio_vae = AudioVAE::new(
vb_vae,
128,
@@ -49,7 +50,7 @@ impl VoxCPMGenerate {
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 mut m_dtype = get_dtype(dtype, cfg_dtype);
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 {
src/models/voxcpm/minicpm4.rs
+36 -29
View File
@@ -1,3 +1,6 @@
use anyhow::{Ok, Result, anyhow};
use candle_core::{D, DType, Device, Tensor};
use candle_nn::{Embedding, Module, RmsNorm, VarBuilder, embedding, rms_norm};
use crate::{
models::{
@@ -7,9 +10,6 @@ use crate::{
position_embed::rope::compute_default_rope_parameters,
utils::tensor_utils::prepare_causal_attention_mask,
};
use anyhow::{anyhow, Ok, Result};
use candle_core::{DType, Device, Tensor, D};
use candle_nn::{Embedding, Module, RmsNorm, VarBuilder, embedding, rms_norm};
pub struct MiniCPMLongRoPE {
short_factor: Vec<f32>,
@@ -77,15 +77,21 @@ impl MiniCPMLongRoPE {
let ext_factors = Tensor::ones_like(&ext_factors)?.div(&ext_factors)?;
let freqs = t.matmul(&ext_factors)?.broadcast_mul(&self.inv_freq)?;
let emb = Tensor::cat(&[&freqs, &freqs], D::Minus1)?;
let cos_cached = emb.cos()?.affine(self.scaling_factor, 0.0)?.to_dtype(self.dtype)?;
let sin_cached = emb.sin()?.affine(self.scaling_factor, 0.0)?.to_dtype(self.dtype)?;
let cos_cached = emb
.cos()?
.affine(self.scaling_factor, 0.0)?
.to_dtype(self.dtype)?;
let sin_cached = emb
.sin()?
.affine(self.scaling_factor, 0.0)?
.to_dtype(self.dtype)?;
self.cos_cached = cos_cached;
self.sin_cached = sin_cached;
Ok(())
}
pub fn forward(&mut self, pos_offset: usize, seqlen: usize) -> Result<(Tensor, Tensor)> {
if pos_offset + seqlen > self.max_seq_len_cached {
let _ = self.update_cos_sin_cache(pos_offset + seqlen)?;
self.update_cos_sin_cache(pos_offset + seqlen)?;
}
let cos = self.cos_cached.narrow(0, pos_offset, seqlen)?;
let sin = self.sin_cached.narrow(0, pos_offset, seqlen)?;
@@ -149,29 +155,25 @@ impl MiniCPMDecoderLayer {
.self_attn
.forward(&xs, cos, sin, attention_mask, true)?;
let xs = if self.use_mup {
let res_add = (residual
(residual
+ xs.affine(
self.scale_depth as f64 / (self.num_hidden_layers as f64).sqrt(),
0.0,
))?;
res_add
))?
} else {
let res_add = (residual + xs)?;
res_add
(residual + xs)?
};
let residual = xs.clone();
let xs = xs.apply(&self.post_attention_layernorm)?;
let xs = xs.apply(&self.mlp)?;
let xs = if self.use_mup {
let res_add = (residual
(residual
+ xs.affine(
self.scale_depth as f64 / (self.num_hidden_layers as f64).sqrt(),
0.0,
))?;
res_add
))?
} else {
let res_add = (residual + xs)?;
res_add
(residual + xs)?
};
Ok(xs)
}
@@ -189,28 +191,24 @@ impl MiniCPMDecoderLayer {
.self_attn
.forward_with_cache(&xs, cos, sin, attention_mask, true)?;
let xs = if self.use_mup {
let res_add = (residual
(residual
+ xs.affine(
self.scale_depth as f64 / (self.num_hidden_layers as f64).sqrt(),
0.0,
)?)?;
res_add
)?)?
} else {
let res_add = (residual + xs)?;
res_add
(residual + xs)?
};
let residual = &xs;
let xs = xs.apply(&self.post_attention_layernorm)?.apply(&self.mlp)?;
let xs = if self.use_mup {
let res_add = (residual
(residual
+ xs.affine(
self.scale_depth as f64 / (self.num_hidden_layers as f64).sqrt(),
0.0,
)?)?;
res_add
)?)?
} else {
let res_add = (residual + xs)?;
res_add
(residual + xs)?
};
Ok(xs)
}
@@ -257,7 +255,12 @@ impl MiniCPMModel {
})
}
pub fn forward(&mut self, input_embeds: &Tensor, position_id: usize, is_causal: bool) -> Result<Tensor> {
pub fn forward(
&mut self,
input_embeds: &Tensor,
position_id: usize,
is_causal: bool,
) -> Result<Tensor> {
let (bs, seq_len, _) = input_embeds.dims3()?;
let attention_mask: Option<&Tensor> = {
if !is_causal || seq_len <= 1 {
@@ -280,11 +283,15 @@ impl MiniCPMModel {
Ok(hidden_states)
}
pub fn forward_with_cache(&mut self, input_embeds: &Tensor, position_id: usize) -> Result<Tensor> {
pub fn forward_with_cache(
&mut self,
input_embeds: &Tensor,
position_id: usize,
) -> Result<Tensor> {
let input_embeds = match input_embeds.rank() {
2 => input_embeds.unsqueeze(1)?,
3 => input_embeds.clone(),
_ => return Err(anyhow!("MiniCPMModelinput_embeds illigal"))
_ => return Err(anyhow!("MiniCPMModelinput_embeds illigal")),
};
let (bs, seq_len, _) = input_embeds.dims3()?;
let attention_mask: Option<&Tensor> = {
src/models/voxcpm/mod.rs
+3 -3
View File
@@ -1,6 +1,6 @@
pub mod config;
pub mod audio_vae;
pub mod config;
pub mod generate;
pub mod minicpm4;
pub mod tokenizer;
pub mod model;
pub mod generate;
pub mod tokenizer;
src/models/voxcpm/model.rs
+14 -10
View File
@@ -7,7 +7,7 @@ use candle_transformers::models::deepseek2::SplitOp;
use crate::{
models::voxcpm::{
audio_vae::{AudioVAE},
audio_vae::AudioVAE,
config::{CfmConfig, VoxCPMConfig, VoxMiniCPM4Config},
minicpm4::MiniCPMModel,
tokenizer::SingleChineseTokenizer,
@@ -67,7 +67,7 @@ impl SinusoidalPosEmb {
let half_dim = self.dim / 2;
let dif = 10000.0_f64.ln() / (half_dim - 1) as f64;
let emb = Tensor::arange(0.0, half_dim as f32, x.device())?
.affine(-1.0 * dif, 0.0)?
.affine(-dif, 0.0)?
.exp()?
.to_dtype(x.dtype())?;
@@ -94,8 +94,8 @@ impl TimestepEmbedding {
out_dim: Option<usize>,
) -> Result<Self> {
let linear_1 = linear(in_channels, time_embed_dim, vb.pp("linear_1"))?;
let time_embed_dim_out = if out_dim.is_some() {
out_dim.unwrap()
let time_embed_dim_out = if let Some(out_dim) = out_dim {
out_dim
} else {
time_embed_dim
};
@@ -104,7 +104,7 @@ impl TimestepEmbedding {
}
pub fn forward(&self, sample: &Tensor) -> Result<Tensor> {
let sample = self.linear_1.forward(&sample)?.silu()?;
let sample = self.linear_1.forward(sample)?.silu()?;
let sample = self.linear_2.forward(&sample)?;
Ok(sample)
}
@@ -199,7 +199,7 @@ pub struct UnifiedCFM {
impl UnifiedCFM {
pub fn new(
in_channels: usize,
cfm_params: CfmConfig,
_cfm_params: CfmConfig,
estimator: VoxCPMLocDiT,
mean_mode: bool,
) -> Result<Self> {
@@ -270,7 +270,7 @@ impl UnifiedCFM {
let mut sol = Vec::new();
let t_span_len = t_span.dim(0)?;
let zero_init_steps = max(1, (t_span_len as f32 * 0.04) as usize);
let mut dphi_dt = Tensor::zeros(1, t_span.dtype(), t_span.device())?;
let mut dphi_dt;
let mut x = x.clone();
for step in 1..t_span_len {
if use_cfg_zero_star && step <= zero_init_steps {
@@ -307,7 +307,7 @@ impl UnifiedCFM {
let cfg = cfg_dphi_dt.broadcast_mul(&st_star)?;
dphi_dt = cfg.add(&dphi_dt.sub(&cfg)?.affine(cfg_value, 0.0)?)?; // step步的预测噪声
}
x = x.broadcast_sub(&dphi_dt.broadcast_mul(&dt)?)?; // 逐步去噪
x = x.broadcast_sub(&dphi_dt.broadcast_mul(&dt)?)?; // 逐步去噪
t = t.sub(&dt)?;
sol.push(x.clone());
if step < t_span_len - 1 {
@@ -642,7 +642,9 @@ impl VoxCPMModel {
let mut pred_feat_seq = Vec::new();
let mut position_id = 0;
let mut seq_len = t;
let enc_outputs = self.base_lm.forward_with_cache(&combined_embed, position_id)?;
let enc_outputs = self
.base_lm
.forward_with_cache(&combined_embed, position_id)?;
let enc_outputs = self
.fsq_layer
.forward(&enc_outputs)?
@@ -653,7 +655,9 @@ impl VoxCPMModel {
let input_embeds =
enc_outputs.add(&feat_mask.unsqueeze(D::Minus1)?.broadcast_mul(&feat_embed)?)?;
let residual_enc_outputs = self.residual_lm.forward_with_cache(&input_embeds, position_id)?;
let residual_enc_outputs = self
.residual_lm
.forward_with_cache(&input_embeds, position_id)?;
let mut residual_hidden = residual_enc_outputs.i((.., t - 1, ..))?;
for i in 0..max_len {
src/models/voxcpm/tokenizer.rs
+1 -1
View File
@@ -26,7 +26,7 @@ impl SingleChineseTokenizer {
if len >= 2 {
let is_chinese = token.chars().all(|c| {
let c_ = c as u32;
0x4E00 <= c_ && c_ <= 0x9FFF
(0x4E00..=0x9FFF).contains(&c_)
});
if is_chinese {
multichar_tokens.push(token);