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

update voxcpm

Browse Source
This commit is contained in:
jhqxxx
2025-10-10 14:17:42 +08:00
parent f33eaaee0d
commit fb746842d7
13 changed files with 135 additions and 200 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/models/voxcpm/audio_vae.rs
+10 -42
View File
@@ -1,7 +1,7 @@
use anyhow::{Error, Ok, Result};
use candle_core::{D, IndexOp, Tensor};
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;
use std::{result::Result::Ok as StdOk, thread, time};
pub struct CausalConv1d {
conv1d: Conv1d,
@@ -9,13 +9,9 @@ pub struct CausalConv1d {
}
impl CausalConv1d {
// CausalConv1d::new(scaled_weight, bias, padding, dilation, stride)?;
pub fn new(
weight: Tensor,
bias: Option<Tensor>,
// in_c: usize,
// out_c: usize,
// kernel_size: usize,
padding: usize,
dilation: usize,
groups: usize,
@@ -59,7 +55,7 @@ impl CausalConvTranspose1d {
) -> Result<Self> {
let config = ConvTranspose1dConfig {
padding: 0,
output_padding,
output_padding: 0,
stride,
dilation,
groups,
@@ -74,23 +70,10 @@ impl CausalConvTranspose1d {
})
}
pub fn forward(&self, x: &Tensor) -> Result<Tensor> {
println!("transpose conv input x: {:?}", x);
println!("transpose conv config stride: {:?}", self.config.stride);
println!("transpose conv config padding: {:?}", self.config.padding);
println!("transpose conv config output_padding: {:?}", self.config.output_padding);
println!("transpose conv config groups: {:?}", self.config.groups);
println!("transpose conv config dilation: {:?}", self.config.dilation);
println!("transpose conv config weight: {:?}", self.conv_transpose1d.weight());
let x = self.conv_transpose1d.forward(x)?;
println!("transpose conv after x: {:?}", x);
println!("transpose conv after self.padding: {:?}", self.padding);
println!("transpose conv after self.output_padding: {:?}", self.output_padding);
let last_dim = x.dim(D::Minus1)?;
let select_num = last_dim - (self.padding * 2 - self.output_padding);
println!("transpose conv after select_num: {:?}", select_num);
let x = x.narrow(D::Minus1, 0, select_num)?;
println!("transpose conv after x: {:?}", x);
Ok(x)
}
}
@@ -109,24 +92,21 @@ 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 weight_norm = weight_v.sqr()?.sum_keepdim(D::Minus1)?.sqrt()?;
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)?;
let conv = CausalConv1d::new(scaled_weight, bias, padding, dilation, groups, stride)?;
Ok(Self { conv })
}
pub fn forward(&self, x: &Tensor) -> Result<Tensor> {
println!("conv1d: x: {:?}", x);
println!("conv weight: : {:?}", self.conv.conv1d.weight());
let x = self.conv.forward(x)?;
println!("conv1d: WN causal x: {:?}", x);
Ok(x)
}
}
@@ -154,7 +134,7 @@ impl WNCausalConvTranspose1d {
StdOk(b) => Some(b),
Err(_) => None,
};
let weight_norm = weight_v.sqr()?.sum_keepdim(D::Minus1)?.sqrt()?;
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)?;
let conv_transpose = CausalConvTranspose1d::new(
@@ -230,7 +210,6 @@ impl CausalResidualUnit {
}
pub fn forward(&self, x: &Tensor) -> Result<Tensor> {
println!("causal residual unit x: {:?}", x);
// let orig_dim = x.dims();
let last_dim_x = x.dim(D::Minus1)?;
let mut res_x = x.clone();
@@ -238,11 +217,8 @@ impl CausalResidualUnit {
let y = self.block1.forward(&y)?;
let y = self.block2.forward(&y)?;
let y = self.block3.forward(&y)?;
println!("causal residual unit y: {:?}", y);
// let dim = y.dims();
let last_dim_y = y.dim(D::Minus1)?;
println!("last_dim_x: {:?}", last_dim_x);
println!("last_dim_y: {:?}", last_dim_y);
let pad = (last_dim_x - last_dim_y) / 2;
if pad > 0 {
res_x = res_x.narrow(D::Minus1, pad, last_dim_y)?;
@@ -415,17 +391,11 @@ impl CausalDecoderBlock {
}
pub fn forward(&self, x: &Tensor) -> Result<Tensor> {
println!("decoder block x : {:?}", x);
let x = self.block0.forward(x)?;
println!("decoder block0 x : {:?}", x);
let x = self.block1.forward(&x)?;
println!("decoder block1 x : {:?}", x);
let x = self.block2.forward(&x)?;
println!("decoder block2 x : {:?}", x);
let x = self.block3.forward(&x)?;
println!("decoder block3 x : {:?}", x);
let x = self.block4.forward(&x)?;
println!("decoder block4 x : {:?}", x);
Ok(x)
}
}
@@ -456,7 +426,7 @@ impl CausalDecoder {
input_channel,
1,
)?;
let model1 = WNCausalConv1d::new(vb.pp("model.1"), input_channel, channels, 1, 1, 1, 1, 1)?;
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();
@@ -484,10 +454,8 @@ impl CausalDecoder {
})
}
pub fn forward(&self, x: &Tensor) -> Result<Tensor> {
print!("audio_vae decoder input x shape: {:?}", x);
let x = self.model0.forward(x)?;
print!("audio_vae decoder model0 x shape: {:?}", 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)?;
src/models/voxcpm/model.rs
+11 -49
View File
@@ -72,10 +72,10 @@ impl SinusoidalPosEmb {
.to_dtype(x.dtype())?;
let emb = x
.unsqueeze(D::Minus1)?
.unsqueeze(1)?
.contiguous()?
.matmul(&emb.unsqueeze(0)?.contiguous()?)?
.affine(scale as f64, 0.0)?;
.affine(scale as f64, 0.0)?
.matmul(&emb.unsqueeze(0)?.contiguous()?)?;
let emb = Tensor::cat(&[emb.sin()?, emb.cos()?], D::Minus1)?;
Ok(emb)
}
@@ -167,7 +167,7 @@ impl VoxCPMLocDiT {
let cond = self
.cond_proj
.forward(&cond.transpose(1, 2)?.contiguous()?)?;
let prefix = cond.dims()[1];
let prefix = cond.dim(1)?;
let t = self.time_embeddings.forward(t, 1000)?.to_dtype(x.dtype())?;
let t = self.time_mlp.forward(&t)?;
let dt = self
@@ -233,7 +233,6 @@ impl UnifiedCFM {
let z = Tensor::randn(0.0f32, 1.0, (b, self.in_channels, t), mu.device())?
.to_dtype(dtype)?
.affine(temperature, 0.0)?;
println!("z: {}", z);
let t_span = linspace(1.0, 0.0, n_timesteps + 1, mu.device())?.to_dtype(dtype)?;
let t_span = t_span
.affine(f64::consts::PI / 2.0, 0.0)?
@@ -242,11 +241,6 @@ impl UnifiedCFM {
.add(&t_span)?
.affine(sway_sampling_coef, 0.0)?
.add(&t_span)?;
println!("t_span: {}", t_span);
println!("mu: {}", mu);
println!("cond: {}", cond);
println!("cfg_value: {}", cfg_value);
println!("use_cfg_zero_star: {}", use_cfg_zero_star);
let x = self.solve_euler(&z, &t_span, mu, cond, cfg_value, use_cfg_zero_star)?;
Ok(x)
}
@@ -274,7 +268,7 @@ impl UnifiedCFM {
let mut t = t_span.i(0)?;
let mut dt = t.sub(&t_span.i(1)?)?;
let mut sol = Vec::new();
let t_span_len = t_span.dims1()?;
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 x = x.clone();
@@ -320,7 +314,8 @@ impl UnifiedCFM {
dt = t.sub(&t_span.i(step + 1)?)?;
}
}
Ok(sol[sol.len() - 1].clone())
let ret = sol[sol.len() - 1].clone();
Ok(ret)
}
}
@@ -333,8 +328,6 @@ pub struct VoxCPMLocEnc {
impl VoxCPMLocEnc {
pub fn new(vb: VarBuilder, config: VoxMiniCPM4Config, input_dim: usize) -> Result<Self> {
// let special_token = Tensor::randn(0.0f32, 1.0, (1, 1, 1, config.hidden_size), vb.device())?
// .to_dtype(vb.dtype())?;
let special_token = vb.get((1, 1, 1, config.hidden_size), "special_token")?;
let in_proj = linear(input_dim, config.hidden_size, vb.pp("in_proj"))?;
assert_eq!(
@@ -354,16 +347,12 @@ impl VoxCPMLocEnc {
pub fn forward(&mut self, x: &Tensor) -> Result<Tensor> {
let (b, t, p, d) = x.dims4()?;
let x = self.in_proj.forward(x)?;
println!("VoxCPMLocEnc: in_proj: {}", x);
let special_tokens = self.special_token.expand((b, t, 1, self.hidden_size))?;
let x = Tensor::cat(&[special_tokens, x], 2)?;
println!("VoxCPMLocEnc: cat: {}", x);
let (b, t, p, c) = x.dims4()?;
let x = x.reshape((b * t, p, c))?;
let outputs = self.encoder.forward(&x, 0, false)?;
println!("VoxCPMLocEnc: encoder: {}", outputs);
let cls_output = outputs.i((.., 0, ..))?;
println!("VoxCPMLocEnc: cls_output: {}", cls_output);
let cls_output = cls_output.reshape((b, t, c))?;
Ok(cls_output)
}
@@ -537,10 +526,8 @@ impl VoxCPMModel {
let audio_feat = audio_feat
.reshape((self.audio_vae.latent_dim, (), self.patch_size))?
.permute((1, 2, 0))?;
let dim0 = audio_feat.dim(0)?;
println!("audio_feat: {:?}", audio_feat);
let dim0 = audio_feat.dim(0)? - 1;
let audio_feat = audio_feat.i(..dim0)?;
println!("audio_feat --: {:?}", audio_feat);
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)?;
@@ -594,7 +581,6 @@ impl VoxCPMModel {
.squeeze(1)?;
let decode_audio_len = decode_audio.dim(D::Minus1)? - 640 - 640;
let decode_audio = decode_audio.narrow(D::Minus1, 640, decode_audio_len)?;
println!("decode_audio: {}", decode_audio);
Ok(decode_audio)
}
@@ -609,21 +595,15 @@ impl VoxCPMModel {
inference_timesteps: usize,
cfg_value: f64,
) -> Result<Tensor> {
println!("text: {}", text);
println!("text_mask: {}", text_mask);
println!("feat: {}", feat);
println!("feat_mask: {}", feat_mask);
let (b, t, p, d) = feat.dims4()?;
let feat_embed = self.feat_encoder.forward(feat)?; // [b, t, h_feat]
println!("feat_embed: {}", feat_embed);
let feat_embed = self.enc_to_lm_proj.forward(&feat_embed)?;
println!("feat_embed: {}", feat_embed);
let scale_emb = if self.config.lm_config.use_mup {
self.config.lm_config.scale_emb
} else {
1.0
};
let text_embed = self
.base_lm
.embed_tokens
@@ -631,41 +611,32 @@ impl VoxCPMModel {
.unwrap()
.forward(text)?
.affine(scale_emb as f64, 0.0)?;
println!("text_embed: {}", text_embed);
let combined_embed = text_mask
.unsqueeze(D::Minus1)?
.broadcast_mul(&text_embed)?
.add(&feat_mask.unsqueeze(D::Minus1)?.broadcast_mul(&feat_embed)?)?;
println!("combined_embed: {}", combined_embed);
let mut prefix_feat_cond = feat.i((.., t - 1, ..))?;
let mut pred_feat_seq = Vec::new();
let mut position_id = 0;
let mut seq_len = t;
let enc_outputs = self.base_lm.forward_step(&combined_embed, position_id)?;
println!("base_lm enc_outputs: {}", enc_outputs);
let enc_outputs = self
.fsq_layer
.forward(&enc_outputs)?
.broadcast_mul(&feat_mask.unsqueeze(D::Minus1)?)?
.add(&enc_outputs.broadcast_mul(&text_mask.unsqueeze(D::Minus1)?)?)?;
println!("fsq_layer enc_outputs: {}", enc_outputs);
let mut lm_hidden = enc_outputs.i((.., t - 1, ..))?;
println!("lm_hidden shape: {:?}", lm_hidden);
let input_embeds =
enc_outputs.add(&feat_mask.unsqueeze(D::Minus1)?.broadcast_mul(&feat_embed)?)?;
let residual_enc_outputs = self.residual_lm.forward_step(&input_embeds, position_id)?;
println!("residual_lm residual_enc_outputs: {}", residual_enc_outputs);
let mut residual_hidden = residual_enc_outputs.i((.., t - 1, ..))?;
for i in 0..max_len {
let dit_hidden_1 = self.lm_to_dit_proj.forward(&lm_hidden)?; // [b, h_dit]
println!("dit_hidden_1: {}", dit_hidden_1);
let dit_hidden_2 = self.res_to_dit_proj.forward(&residual_hidden)?; // [b, h_dit]
println!("dit_hidden_2: {}", dit_hidden_2);
let dit_hidden = dit_hidden_1.add(&dit_hidden_2)?;
println!("dit_hidden: {}", dit_hidden);
let cond = prefix_feat_cond.transpose(1, 2)?.contiguous()?;
let pred_feat = self
@@ -681,24 +652,18 @@ impl VoxCPMModel {
true,
)?
.transpose(1, 2)?; // [b, p, d]
println!("pred_feat: {}", pred_feat);
let curr_embed = self.feat_encoder.forward(&pred_feat.unsqueeze(1)?)?; // [b, 1, c]
let curr_embed = self.enc_to_lm_proj.forward(&curr_embed)?;
println!("curr_embed: {}", curr_embed);
pred_feat_seq.push(pred_feat.unsqueeze(1)?);
prefix_feat_cond = pred_feat;
println!("lm_hidden: {}", lm_hidden);
let stop_flag = self.stop_proj.forward(&lm_hidden)?.silu()?;
println!("stop_flag: {}", stop_flag);
let stop_flag = self
.stop_head
.forward(&stop_flag)?
.argmax(D::Minus1)?
.i(0)?
.to_scalar::<u32>()?;
println!("i: {}, stop_flag: {}", i, stop_flag);
if i > min_len && stop_flag == 1 {
break;
}
@@ -716,15 +681,12 @@ impl VoxCPMModel {
}
let pred_seq = Tensor::cat(&pred_feat_seq, 1)?; // (b, t, p, d)
let (b, t, p, d) = pred_seq.dims4()?;
println!("pred_seq: {:?}", pred_seq);
let feat_pred = pred_seq
.permute((0, 3, 1, 2))?
.reshape((b, d, ()))?
.contiguous()?;
println!("feat_pred: {:?}", feat_pred);
self.base_lm.clear_kv_cache();
self.residual_lm.clear_kv_cache();
Ok(feat_pred)
}
}
src/models/voxcpm/tokenizer.rs
+2 -2
View File
@@ -45,7 +45,7 @@ impl SingleChineseTokenizer {
.encode(text, false)
.map_err(|e| anyhow!(format!("tokenizer encode error: {}", e)))?;
let tokens = encode.get_tokens();
println!("tokens: {:?}", tokens);
// println!("tokens: {:?}", tokens);
let mut split_character = Vec::new();
for token in tokens {
let clean_token = token.replace("▁", "to");
@@ -56,7 +56,7 @@ impl SingleChineseTokenizer {
split_character.push(token.clone());
}
}
println!("split_character: {:?}", split_character);
// println!("split_character: {:?}", split_character);
let ids: Vec<u32> = split_character
.iter()
.filter_map(|c| self.tokenizer.token_to_id(c))