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

add voxcpm with some bug

Browse Source
This commit is contained in:
jhqxxx
2025-10-03 22:25:58 +08:00
parent ae36194a4b
commit f33eaaee0d
30 changed files with 2411 additions and 165 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/models/base_modules/mod.rs
+4 -2
View File
@@ -116,6 +116,7 @@ impl AttentionNobias {
cos: &Tensor,
sin: &Tensor,
attention_mask: Option<&Tensor>,
tof32: bool,
) -> Result<Tensor> {
let (b_sz, q_len, _) = xs.dims3()?;
let query_states = self.q_proj.forward(xs)?;
@@ -131,7 +132,7 @@ impl AttentionNobias {
.reshape((b_sz, q_len, self.num_kv_heads, self.head_dim))?
.transpose(1, 2)?;
let (query_states, key_states) =
apply_rotary_pos_emb(&query_states, &key_states, cos, sin)?;
apply_rotary_pos_emb(&query_states, &key_states, cos, sin, tof32)?;
let key_states = repeat_kv(key_states, self.num_kv_groups)?.contiguous()?;
let value_states = repeat_kv(value_states, self.num_kv_groups)?.contiguous()?;
@@ -184,6 +185,7 @@ impl AttentionNobias {
cos: &Tensor,
sin: &Tensor,
attention_mask: Option<&Tensor>,
tof32: bool,
) -> Result<Tensor> {
let (b_sz, q_len, _) = xs.dims3()?;
let query_states = self.q_proj.forward(xs)?;
@@ -199,7 +201,7 @@ impl AttentionNobias {
.reshape((b_sz, q_len, self.num_kv_heads, self.head_dim))?
.transpose(1, 2)?;
let (query_states, key_states) =
apply_rotary_pos_emb(&query_states, &key_states, cos, sin)?;
apply_rotary_pos_emb(&query_states, &key_states, cos, sin, tof32)?;
let (key_states, value_states) = match &self.kv_cache {
None => (key_states, value_states),
Some((prev_k, prev_v)) => {
src/models/minicpm4/config.rs
+1 -4
View File
@@ -23,10 +23,7 @@ pub struct MiniCPM4Config {
pub rope_scaling: RopeScalingConfig,
pub torch_dtype: String,
pub vocab_size: usize,
// pub use_mup: bool,
pub scale_emb:f32,
pub scale_emb: f64,
pub dim_model_base: usize,
pub scale_depth: f32,
// pub rope_theta: f32,
// pub kv_channels: i32,
}
src/models/minicpm4/generate.rs
+5 -6
View File
@@ -2,15 +2,14 @@ use crate::models::minicpm4::config::MiniCPM4Config;
use crate::models::minicpm4::model::MiniCPMModel;
// use crate::models::GenerateStream;
use crate::utils::utils::{
build_completion_chunk_response, build_completion_response, find_safetensors_files, get_device,
get_dtype, get_logit_processor,
build_completion_chunk_response, build_completion_response, find_type_files, get_device, get_dtype, get_logit_processor
};
use crate::{
chat_template::chat_template::ChatTemplate, models::GenerateModel,
tokenizer::tokenizer::TokenizerModel,
};
use anyhow::{Result, anyhow};
use candle_core::{D, DType, Device, IndexOp, Tensor};
use candle_core::{DType, Device, Tensor};
use candle_nn::VarBuilder;
use openai_dive::v1::resources::chat::{
ChatCompletionChunkResponse, ChatCompletionParameters, ChatCompletionResponse,
@@ -27,7 +26,7 @@ pub struct MiniCPMGenerateModel<'a> {
im_end_id: u32,
}
impl<'a> MiniCPMGenerateModel<'a> {
impl <'a> MiniCPMGenerateModel<'a> {
pub fn init(path: &str, device: Option<&Device>, dtype: Option<DType>) -> Result<Self> {
let chat_template = ChatTemplate::init(path)?;
let tokenizer = TokenizerModel::init(path)?;
@@ -38,7 +37,7 @@ impl<'a> MiniCPMGenerateModel<'a> {
let dtype = get_dtype(dtype, cfg_dtype);
let endoftext_id = cfg.eos_token_id[0];
let im_end_id = cfg.eos_token_id[1];
let model_list = find_safetensors_files(&path)?;
let model_list = find_type_files(&path, "safetensors")?;
let vb = unsafe { VarBuilder::from_mmaped_safetensors(&model_list, dtype, device)? };
let minicpm = MiniCPMModel::new(vb, cfg)?;
@@ -64,7 +63,7 @@ impl<'a> GenerateModel for MiniCPMGenerateModel<'a> {
let mut generate = Vec::new();
let sample_len = match mes.max_tokens {
Some(max) => max,
None => 512,
None => 2048,
};
for _ in 0..sample_len {
let logits = self.minicpm.forward_step(&input_ids, seqlen_offset)?;
src/models/minicpm4/model.rs
+79 -40
View File
@@ -7,38 +7,41 @@ use crate::{
utils::tensor_utils::prepare_causal_attention_mask,
};
use anyhow::{Ok, Result};
use candle_core::{D, DType, Device, Tensor, Var};
use candle_nn::{embedding, rms_norm, Embedding, Linear, Module, RmsNorm, VarBuilder};
use candle_core::{D, Device, Tensor};
use candle_nn::{Embedding, Linear, Module, RmsNorm, VarBuilder, embedding, rms_norm};
pub struct MiniCPMLongRoPE {
head_dim: usize,
rope_theta: f32,
max_position_embeddings: usize,
short_factor: Vec<f32>,
long_factor: Vec<f32>,
original_max_position_embeddings: usize,
max_seq_len_cached: usize,
scaling_factor: f64,
inv_freq: Tensor,
cos_cached: Tensor,
sin_cached: Tensor,
device: Device,
}
impl MiniCPMLongRoPE {
pub fn new(cfg: &MiniCPM4Config, device: &Device) -> Result<Self> {
let head_dim = cfg.hidden_size / cfg.num_attention_heads;
let rope_theta = 10000.0;
let max_position_embeddings = cfg.max_position_embeddings;
let short_factor = cfg.rope_scaling.short_factor.clone();
let long_factor = cfg.rope_scaling.short_factor.clone();
let original_max_position_embeddings = cfg.rope_scaling.original_max_position_embeddings;
let scale = max_position_embeddings / original_max_position_embeddings;
let max_position_embeddings = cfg.max_position_embeddings;
let scale = max_position_embeddings as f64 / original_max_position_embeddings as f64;
let scaling_factor =
(1.0 + (scale as f64).ln() + (original_max_position_embeddings as f64).ln()).sqrt();
(1.0 + scale.ln() / (original_max_position_embeddings as f64).ln()).sqrt();
let inv_freq = compute_default_rope_parameters(head_dim, rope_theta);
let inv_freq = Tensor::from_slice(&inv_freq, (1, inv_freq.len()), device)?;
let inv_freq =
Tensor::from_slice(&inv_freq, (1, inv_freq.len()), device)?;
let max_seq_len_cached = max_position_embeddings;
let t = Tensor::arange(0.0_f32, max_position_embeddings as f32, device)?
.reshape((max_position_embeddings, 1))?;
// short_factor.len() = 32
// head_dim = 1024 / 16 = 64, inv_freq.len() = 32
let ext_factors = Tensor::from_slice(&short_factor, (1, short_factor.len()), device)?;
let ext_factors =
Tensor::from_slice(&short_factor, (1, short_factor.len()), device)?;
let ext_factors = Tensor::ones_like(&ext_factors)?.div(&ext_factors)?;
// (seq_len, 1) matmul (1, 32) -> (seq_len, 32) * (1, 32)-> (seq_len, 32)
let freqs = t.matmul(&ext_factors)?.broadcast_mul(&inv_freq)?;
@@ -47,41 +50,46 @@ impl MiniCPMLongRoPE {
let cos_cached = emb.cos()?.affine(scaling_factor, 0.0)?;
let sin_cached = emb.sin()?.affine(scaling_factor, 0.0)?;
Ok(Self {
head_dim,
rope_theta,
max_position_embeddings,
short_factor,
long_factor,
original_max_position_embeddings,
max_seq_len_cached,
scaling_factor,
inv_freq,
cos_cached,
sin_cached,
device: device.clone(),
})
}
pub fn update_cos_sin_cache(&mut self, seqlen: usize, device: &Device) -> Result<()> {
let t = Tensor::arange(0.0_f32, seqlen as f32, device)?.reshape((seqlen, 1))?;
let mut ext_factors =
Tensor::from_slice(&self.short_factor, (1, self.short_factor.len()), device)?;
pub fn update_cos_sin_cache(&mut self, seqlen: usize) -> Result<()> {
self.max_seq_len_cached = seqlen;
let t = Tensor::arange(0.0_f32, seqlen as f32, &self.device)?
.reshape((seqlen, 1))?;
let mut ext_factors = Tensor::from_slice(
&self.short_factor,
(1, self.short_factor.len()),
&self.device,
)?;
if seqlen > self.original_max_position_embeddings {
ext_factors =
Tensor::from_slice(&self.long_factor, (1, self.long_factor.len()), device)?;
Tensor::from_slice(&self.long_factor, (1, self.long_factor.len()), &self.device)?;
}
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 scale = seqlen / self.original_max_position_embeddings;
let scaling_factor =
(1.0 + (scale as f64).ln() + (self.original_max_position_embeddings as f64).ln())
.sqrt();
let cos_cached = emb.cos()?.affine(scaling_factor, 0.0)?;
let sin_cached = emb.sin()?.affine(scaling_factor, 0.0)?;
let cos_cached = emb.cos()?.affine(self.scaling_factor, 0.0)?;
let sin_cached = emb.sin()?.affine(self.scaling_factor, 0.0)?;
self.cos_cached = cos_cached;
self.sin_cached = sin_cached;
Ok(())
}
pub fn forward(&self, pos_offset: usize, seqlen: usize) -> Result<(Tensor, Tensor)> {
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)?;
}
let cos = self.cos_cached.narrow(0, pos_offset, seqlen)?;
let sin = self.sin_cached.narrow(0, pos_offset, seqlen)?;
Ok((cos, sin))
}
}
@@ -133,13 +141,21 @@ impl MiniCPMDecoderLayer {
sin: &Tensor,
attention_mask: Option<&Tensor>,
) -> Result<Tensor> {
let residual = xs;
let residual = xs.clone();
let xs = self.input_layernorm.forward(xs)?;
let xs = self.self_attn.forward(&xs, cos, sin, attention_mask)?;
let xs = (xs + residual)?;
let xs = self.self_attn.forward(&xs, cos, sin, attention_mask, true)?;
let xs = (residual
+ xs.affine(
self.scale_depth as f64 / (self.num_hidden_layers as f64).sqrt(),
0.0,
))?;
let residual = &xs;
let xs = xs.apply(&self.post_attention_layernorm)?.apply(&self.mlp)?;
let xs = (residual + xs)?;
let xs = (residual
+ xs.affine(
self.scale_depth as f64 / (self.num_hidden_layers as f64).sqrt(),
0.0,
))?;
Ok(xs)
}
@@ -150,13 +166,21 @@ impl MiniCPMDecoderLayer {
sin: &Tensor,
attention_mask: Option<&Tensor>,
) -> Result<Tensor> {
let residual = xs;
let residual = xs.clone();
let xs = self.input_layernorm.forward(xs)?;
let xs = self.self_attn.forward_step(&xs, cos, sin, attention_mask)?;
let xs = (xs + residual)?;
let xs = self.self_attn.forward_step(&xs, cos, sin, attention_mask, true)?;
let xs = (residual
+ xs.affine(
self.scale_depth as f64 / (self.num_hidden_layers as f64).sqrt(),
0.0,
))?;
let residual = &xs;
let xs = xs.apply(&self.post_attention_layernorm)?.apply(&self.mlp)?;
let xs = (residual + xs)?;
let xs = (residual
+ xs.affine(
self.scale_depth as f64 / (self.num_hidden_layers as f64).sqrt(),
0.0,
))?;
Ok(xs)
}
pub fn clear_kv_cache(&mut self) {
@@ -175,6 +199,7 @@ pub struct MiniCPMModel {
impl MiniCPMModel {
pub fn new(vb: VarBuilder, cfg: MiniCPM4Config) -> Result<Self> {
let vb = vb.pp("model");
let embed_tokens = embedding(cfg.vocab_size, cfg.hidden_size, vb.pp("embed_tokens"))?;
let mut layers = Vec::with_capacity(cfg.num_hidden_layers);
let vb_layers = vb.pp("layers");
@@ -191,13 +216,16 @@ impl MiniCPMModel {
layers,
norm,
rope_emb,
lm_head
lm_head,
})
}
pub fn forward(&self, input_ids: &Tensor, position_id: usize) -> Result<Tensor> {
pub fn forward(&mut self, input_ids: &Tensor, position_id: usize) -> Result<Tensor> {
let (bs, seq_len) = input_ids.dims2()?;
let input_embeds = self.embed_tokens.forward(&input_ids)?;
let input_embeds = self
.embed_tokens
.forward(&input_ids)?
.affine(self.cfg.scale_emb, 0.0)?;
let attention_mask: Option<&Tensor> = {
if seq_len <= 1 {
None
@@ -210,7 +238,7 @@ impl MiniCPMModel {
)?)
}
};
let (cos, sin) = self.rope_emb.forward(position_id, seq_len)?;
let mut hidden_states = input_embeds;
for decode_layer in &self.layers {
@@ -218,13 +246,20 @@ impl MiniCPMModel {
}
hidden_states = self.norm.forward(&hidden_states)?;
let hidden_state = hidden_states.narrow(1, seq_len - 1, 1)?;
let hidden_state = hidden_state.affine(
1.0 / (self.cfg.hidden_size / self.cfg.dim_model_base) as f64,
0.0,
)?;
let logits = self.lm_head.forward(&hidden_state)?;
Ok(logits)
}
pub fn forward_step(&mut self, input_ids: &Tensor, position_id: usize) -> Result<Tensor> {
let (bs, seq_len) = input_ids.dims2()?;
let input_embeds = self.embed_tokens.forward(&input_ids)?;
let input_embeds = self
.embed_tokens
.forward(&input_ids)?
.affine(self.cfg.scale_emb, 0.0)?;
let attention_mask: Option<&Tensor> = {
if seq_len <= 1 {
None
@@ -237,14 +272,18 @@ impl MiniCPMModel {
)?)
}
};
let (cos, sin) = self.rope_emb.forward(position_id, seq_len)?;
let mut hidden_states = input_embeds;
for decode_layer in &mut self.layers {
hidden_states = decode_layer.forward_step(&hidden_states, &cos, &sin, attention_mask)?;
hidden_states =
decode_layer.forward_step(&hidden_states, &cos, &sin, attention_mask)?;
}
hidden_states = self.norm.forward(&hidden_states)?;
let hidden_state = hidden_states.narrow(1, seq_len - 1, 1)?;
let hidden_state = hidden_state.affine(
1.0 / (self.cfg.hidden_size / self.cfg.dim_model_base) as f64,
0.0,
)?;
let logits = self.lm_head.forward(&hidden_state)?;
Ok(logits)
}
src/models/mod.rs
+1 -4
View File
@@ -1,18 +1,15 @@
pub mod base_modules;
pub mod minicpm4;
pub mod qwen2_5vl;
pub mod voxcpm;
use anyhow::Result;
use candle_core::{DType, Device};
use openai_dive::v1::resources::chat::{
ChatCompletionChunkResponse, ChatCompletionParameters, ChatCompletionResponse,
};
use rocket::futures::Stream;
pub trait GenerateModel {
// fn init(path: &str, device: Option<&Device>, dtype: Option<DType>) -> Result<Self>
// where
// Self: Sized;
fn generate(&mut self, mes: ChatCompletionParameters) -> Result<ChatCompletionResponse>;
fn generate_stream(
&mut self,
src/models/qwen2_5vl/generate.rs
+4 -4
View File
@@ -1,8 +1,7 @@
// use crate::models::GenerateStream;
use crate::models::qwen2_5vl::config::Qwen2_5VLConfig;
use crate::utils::utils::{
build_completion_chunk_response, build_completion_response, find_safetensors_files, get_device,
get_dtype, get_logit_processor,
build_completion_chunk_response, build_completion_response, find_type_files, get_device, get_dtype, get_logit_processor
};
use crate::{
chat_template::chat_template::ChatTemplate,
@@ -43,7 +42,8 @@ impl<'a> Qwen2_5VLGenerateModel<'a> {
let pre_processor = Qwen2_5VLProcessor::new(device, dtype)?;
let endoftext_id = cfg.bos_token_id;
let im_end_id = cfg.eos_token_id;
let model_list = find_safetensors_files(&path)?;
// let model_list = find_safetensors_files(&path)?;
let model_list = find_type_files(&path, "safetensors")?;
let vb = unsafe { VarBuilder::from_mmaped_safetensors(&model_list, dtype, device)? };
let qwen2_5_vl = Qwen2_5VLModel::new(cfg, vb)?;
@@ -86,7 +86,7 @@ impl<'a> GenerateModel for Qwen2_5VLGenerateModel<'a> {
let mut generate = Vec::new();
let sample_len = match mes.max_tokens {
Some(max) => max,
None => 512,
None => 1024,
};
for _ in 0..sample_len {
let logits = self.qwen2_5_vl.forward(
src/models/qwen2_5vl/model.rs
+1 -1
View File
@@ -631,7 +631,7 @@ impl Qwen2_5VLTextAttention {
.reshape((b_sz, q_len, self.num_kv_heads, self.head_dim))?
.transpose(1, 2)?;
let (query_states, key_states) =
apply_rotary_pos_emb(&query_states, &key_states, cos, sin)?;
apply_rotary_pos_emb(&query_states, &key_states, cos, sin, false)?;
let (key_states, value_states) = match &self.kv_cache {
None => (key_states, value_states),
Some((prev_k, prev_v)) => {
src/models/voxcpm/audio_vae.rs
+586
View File
@@ -0,0 +1,586 @@
use anyhow::{Error, Ok, Result};
use candle_core::{D, IndexOp, Tensor};
use candle_nn::{Conv1d, Conv1dConfig, ConvTranspose1d, ConvTranspose1dConfig, Module, VarBuilder};
use std::result::Result::Ok as StdOk;
pub struct CausalConv1d {
conv1d: Conv1d,
padding: usize,
}
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,
stride: usize,
) -> Result<Self> {
let config = Conv1dConfig {
padding: 0,
stride,
dilation,
groups,
cudnn_fwd_algo: None,
};
let conv1d = Conv1d::new(weight, bias, config);
Ok(Self { conv1d, padding })
}
pub fn forward(&self, x: &Tensor) -> Result<Tensor> {
let x_pad = x.pad_with_zeros(D::Minus1, self.padding * 2, 0)?;
let x = self.conv1d.forward(&x_pad)?;
Ok(x)
}
}
pub struct CausalConvTranspose1d {
conv_transpose1d: ConvTranspose1d,
padding: usize,
output_padding: usize,
config: ConvTranspose1dConfig,
}
impl CausalConvTranspose1d {
pub fn new(
weight: Tensor,
bias: Option<Tensor>,
padding: usize,
dilation: usize,
output_padding: usize,
groups: usize,
stride: usize,
) -> Result<Self> {
let config = ConvTranspose1dConfig {
padding: 0,
output_padding,
stride,
dilation,
groups,
};
let conv_transpose1d = ConvTranspose1d::new(weight, bias, config.clone());
Ok(Self {
conv_transpose1d,
padding,
output_padding,
config
})
}
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)
}
}
pub struct WNCausalConv1d {
conv: CausalConv1d,
}
impl WNCausalConv1d {
pub fn new(
vb: VarBuilder,
in_c: usize,
out_c: usize,
kernel_size: usize,
dilation: usize,
padding: usize,
groups: usize,
stride: usize,
) -> Result<Self> {
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 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)
}
}
pub struct WNCausalConvTranspose1d {
conv_transpose: CausalConvTranspose1d,
}
impl WNCausalConvTranspose1d {
pub fn new(
vb: VarBuilder,
in_c: usize,
out_c: usize,
dilation: usize,
kernel_size: usize,
padding: usize,
output_padding: usize,
groups: usize,
stride: usize,
) -> Result<Self> {
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 weight_norm = weight_v.sqr()?.sum_keepdim(D::Minus1)?.sqrt()?;
let normalized_weight = weight_v.broadcast_div(&weight_norm)?;
let scaled_weight = normalized_weight.broadcast_mul(&weight_g)?;
let conv_transpose = CausalConvTranspose1d::new(
scaled_weight,
bias,
padding,
dilation,
output_padding,
groups,
stride,
)?;
Ok(Self { conv_transpose })
}
pub fn forward(&self, x: &Tensor) -> Result<Tensor> {
let x = self.conv_transpose.forward(x)?;
Ok(x)
}
}
pub struct Snake1d {
alpha: Tensor,
}
impl Snake1d {
pub fn new(vb: VarBuilder, channels: usize) -> Result<Self> {
let alpha = vb.get((1, channels, 1), "alpha")?;
Ok(Self { alpha })
}
pub fn forward(&self, x: &Tensor) -> Result<Tensor> {
let dims = x.dims();
let x = x.reshape((dims[0], dims[1], ()))?;
let alpha_ = self.alpha.affine(1.0, 1e-9)?.recip()?;
let alpha_ = x
.broadcast_mul(&self.alpha)?
.sin()?
.powf(2.0)?
.broadcast_mul(&alpha_)?;
let x = x.add(&alpha_)?;
let x = x.reshape(dims)?;
Ok(x)
}
}
pub struct CausalResidualUnit {
// pad: usize,
block0: Snake1d,
block1: WNCausalConv1d,
block2: Snake1d,
block3: WNCausalConv1d,
}
impl CausalResidualUnit {
pub fn new(
vb: VarBuilder,
dim: usize,
dilation: usize,
kernel: usize,
groups: usize,
) -> Result<Self> {
let pad = ((7 - 1) * dilation) / 2;
let block0 = Snake1d::new(vb.pp("block.0"), dim)?;
let block1 =
WNCausalConv1d::new(vb.pp("block.1"), dim, dim, kernel, dilation, pad, groups, 1)?;
let block2 = Snake1d::new(vb.pp("block.2"), dim)?;
let block3 = WNCausalConv1d::new(vb.pp("block.3"), dim, dim, 1, 1, 0, 1, 1)?;
Ok(Self {
// pad,
block0,
block1,
block2,
block3,
})
}
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();
let y = self.block0.forward(x)?;
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)?;
}
let x = y.add(&res_x)?;
Ok(x)
}
}
pub struct CausalEncoderBlock {
block0: CausalResidualUnit,
block1: CausalResidualUnit,
block2: CausalResidualUnit,
block3: Snake1d,
block4: WNCausalConv1d,
}
impl CausalEncoderBlock {
pub fn new(
vb: VarBuilder,
in_dim: Option<usize>,
out_dim: usize,
stride: usize,
groups: usize,
) -> Result<Self> {
let in_dim = match in_dim {
Some(d) => d,
None => out_dim / 2,
};
let block0 = CausalResidualUnit::new(vb.pp("block.0"), in_dim, 1, 7, groups)?;
let block1 = CausalResidualUnit::new(vb.pp("block.1"), in_dim, 3, 7, groups)?;
let block2 = CausalResidualUnit::new(vb.pp("block.2"), in_dim, 9, 7, groups)?;
let block3 = Snake1d::new(vb.pp("block.3"), in_dim)?;
let padding = (stride as f32 / 2.0).ceil() as usize;
let block4 = WNCausalConv1d::new(
vb.pp("block.4"),
in_dim,
out_dim,
2 * stride,
1,
padding,
1,
stride,
)?;
Ok(Self {
block0,
block1,
block2,
block3,
block4,
})
}
pub fn forward(&self, x: &Tensor) -> Result<Tensor> {
let x = self.block0.forward(x)?;
let x = self.block1.forward(&x)?;
let x = self.block2.forward(&x)?;
let x = self.block3.forward(&x)?;
let x = self.block4.forward(&x)?;
Ok(x)
}
}
pub struct CausalEncoder {
block0: WNCausalConv1d,
block1_4: Vec<CausalEncoderBlock>,
fc_mu: WNCausalConv1d,
fc_logvar: WNCausalConv1d,
}
impl CausalEncoder {
pub fn new(
vb: VarBuilder,
d_model: usize,
laten_dim: usize,
strides: Vec<usize>,
depthwise: bool,
) -> Result<Self> {
let mut d_model = d_model;
let mut groups = 1;
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)?;
block1_4.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,
fc_mu,
fc_logvar,
})
}
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 {
hidden_state = block_i.forward(&hidden_state)?;
}
let mu = self.fc_mu.forward(&hidden_state)?;
let logvar = self.fc_logvar.forward(&hidden_state)?;
Ok((hidden_state, mu, logvar))
}
}
pub struct NoiseBlock {
linear: WNCausalConv1d,
}
impl NoiseBlock {
pub fn new(vb: VarBuilder, dim: usize) -> Result<Self> {
let linear = WNCausalConv1d::new(vb.pp("linear"), dim, dim, 1, 1, 0, 1, 1)?;
Ok(Self { linear })
}
pub fn forward(&self, x: &Tensor) -> Result<Tensor> {
let (bs, _, t) = x.dims3()?;
let noise = Tensor::randn(0.0_f32, 1.0, (bs, 1, t), x.device())?.to_dtype(x.dtype())?;
let h = self.linear.forward(x)?;
let n = h.broadcast_mul(&noise)?;
let x = x.add(&n)?;
Ok(x)
}
}
pub struct CausalDecoderBlock {
block0: Snake1d,
block1: WNCausalConvTranspose1d,
block2: CausalResidualUnit,
block3: CausalResidualUnit,
block4: CausalResidualUnit,
}
impl CausalDecoderBlock {
pub fn new(
vb: VarBuilder,
input_dim: usize,
output_dim: usize,
stride: usize,
groups: usize,
) -> Result<Self> {
let block0 = Snake1d::new(vb.pp("block.0"), input_dim)?;
let padding = (stride as f32 / 2.0).ceil() as usize;
let block1 = WNCausalConvTranspose1d::new(
vb.pp("block.1"),
input_dim,
output_dim,
1,
2 * stride,
padding,
stride % 2,
1,
stride,
)?;
let block2 = CausalResidualUnit::new(vb.pp("block.2"), output_dim, 1, 7, groups)?;
let block3 = CausalResidualUnit::new(vb.pp("block.3"), output_dim, 3, 7, groups)?;
let block4 = CausalResidualUnit::new(vb.pp("block.4"), output_dim, 9, 7, groups)?;
Ok(Self {
block0,
block1,
block2,
block3,
block4,
})
}
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)
}
}
pub struct CausalDecoder {
model0: WNCausalConv1d,
model1: WNCausalConv1d,
model2_5: Vec<CausalDecoderBlock>,
model6: Snake1d,
model7: WNCausalConv1d,
}
impl CausalDecoder {
pub fn new(
vb: VarBuilder,
input_channel: usize,
channels: usize,
rates: Vec<usize>,
d_out: usize,
) -> Result<Self> {
let model0 = WNCausalConv1d::new(
vb.pp("model.0"),
input_channel,
input_channel,
7,
1,
3,
input_channel,
1,
)?;
let model1 = WNCausalConv1d::new(vb.pp("model.1"), input_channel, channels, 1, 1, 1, 1, 1)?;
let vb_model = vb.pp("model");
let mut output_dim = channels;
let mut model2_5 = 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 model_i = CausalDecoderBlock::new(
vb_model.pp(i + 2),
input_dim,
output_dim,
*stride,
groups,
)?;
model2_5.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)?;
Ok(Self {
model0,
model1,
model2_5,
model6,
model7,
})
}
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);
let mut x = self.model1.forward(&x)?;
for model_i in &self.model2_5 {
x = model_i.forward(&x)?;
}
let x = self.model6.forward(&x)?;
let x = self.model7.forward(&x)?;
let x = x.tanh()?;
Ok(x)
}
}
pub struct AudioVAE {
encoder_dim: usize,
encoder_rates: Vec<usize>,
decoder_dim: usize,
decoder_rates: Vec<usize>,
pub latent_dim: usize,
hop_length: usize,
encoder: CausalEncoder,
decoder: CausalDecoder,
pub sample_rate: usize,
pub chunk_size: usize,
}
impl AudioVAE {
pub fn new(
vb: VarBuilder,
encoder_dim: usize,
encoder_rates: Vec<usize>,
laten_dim: Option<usize>,
decoder_dim: usize,
decoder_rates: Vec<usize>,
sample_rate: usize,
) -> Result<Self> {
let latent_dim = match laten_dim {
Some(d) => d,
None => encoder_dim * (2_usize.pow(encoder_rates.len() as u32)),
};
let hop_length = encoder_rates.iter().product();
let encoder = CausalEncoder::new(
vb.pp("encoder"),
encoder_dim,
latent_dim,
encoder_rates.clone(),
true,
)?;
let decoder = CausalDecoder::new(
vb.pp("decoder"),
latent_dim,
decoder_dim,
decoder_rates.clone(),
1,
)?;
let chunk_size = hop_length;
Ok(Self {
encoder_dim,
encoder_rates,
decoder_dim,
decoder_rates,
latent_dim,
hop_length,
encoder,
decoder,
sample_rate,
chunk_size,
})
}
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
};
assert_eq!(sample_rate, self.sample_rate);
let pad_to = self.hop_length;
let length = audio_data.dim(D::Minus1)?;
let right_pad = (length as f32 / pad_to as f32).ceil() as usize * pad_to - length;
let audio_data = audio_data.pad_with_zeros(D::Minus1, 0, right_pad)?;
Ok(audio_data)
}
pub fn decode(&self, z: &Tensor) -> Result<Tensor> {
let x = self.decoder.forward(z)?;
Ok(x)
}
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()
};
let audio_data = self.preprocess(&audio_data, sample_rate)?;
let (_, mu, _) = self.encoder.forward(&audio_data)?;
Ok(mu)
}
}
src/models/voxcpm/config.rs
+33 -22
View File
@@ -1,15 +1,14 @@
use candle_nn::Activation;
#[derive(Debug, Clone, PartialEq, serde::Deserialize)]
pub struct RopeScalingConfig {
pub rope_type: String,
pub struct VoxRopeScalingConfig {
pub r#type: String,
pub long_factor: Vec<f32>,
pub short_factor: Vec<f32>,
pub original_max_position_embeddings: usize,
}
#[derive(Debug, Clone, PartialEq, serde::Deserialize)]
pub struct MiniCPM4Config {
pub struct VoxMiniCPM4Config {
pub bos_token_id: u32,
pub eos_token_id: u32,
pub hidden_size: usize,
@@ -19,38 +18,50 @@ pub struct MiniCPM4Config {
pub num_hidden_layers: usize,
pub num_key_value_heads: usize,
pub rms_norm_eps: f64,
pub rope_scaling: RopeScalingConfig,
pub torch_dtype: String,
pub rope_theta: f32,
pub rope_scaling: VoxRopeScalingConfig,
pub vocab_size: usize,
// pub use_mup: bool,
pub scale_emb:f32,
pub dim_model_base: usize,
pub scale_depth: f32,
// pub rope_theta: f32,
// pub kv_channels: i32,
pub use_mup: bool,
}
#[derive(Debug, Clone, PartialEq, serde::Deserialize)]
pub struct VoxCPMEncoderConfig {
hidden_dim: usize,
ffn_dim: usize,
num_heads: usize,
num_layers: usize,
pub hidden_dim: usize,
pub ffn_dim: usize,
pub num_heads: usize,
pub num_layers: usize,
}
#[derive(Debug, Clone, PartialEq, serde::Deserialize)]
pub struct CfmConfig {
sigma_min: f32,
solver: String,
t_scheduler: String,
inference_cfg_rate: f32,
pub sigma_min: f32,
pub solver: String,
pub t_scheduler: String,
pub inference_cfg_rate: f32,
}
#[derive(Debug, Clone, PartialEq, serde::Deserialize)]
pub struct VoxCPMDitConfig {
hidden_dim: usize,
ffn_dim: usize,
num_heads: usize,
num_layers: usize,
cfm_config: CfmConfig,
pub hidden_dim: usize,
pub ffn_dim: usize,
pub num_heads: usize,
pub num_layers: usize,
pub cfm_config: CfmConfig,
}
#[derive(Debug, Clone, PartialEq, serde::Deserialize)]
pub struct VoxCPMConfig {
pub lm_config: VoxMiniCPM4Config,
pub patch_size: usize,
pub feat_dim: usize,
pub scalar_quantization_latent_dim: usize,
pub scalar_quantization_scale: usize,
pub residual_lm_num_layers: usize,
pub encoder_config: VoxCPMEncoderConfig,
pub dit_config: VoxCPMDitConfig,
pub max_length: usize,
pub dtype: String,
}
src/models/voxcpm/minicpm4.rs
+330
View File
@@ -0,0 +1,330 @@
use std::{thread, time};
use crate::{
models::{
base_modules::{AttentionNobias, MLPNoBias},
voxcpm::config::VoxMiniCPM4Config,
},
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, Linear, Module, RmsNorm, VarBuilder, embedding, rms_norm};
pub struct MiniCPMLongRoPE {
short_factor: Vec<f32>,
long_factor: Vec<f32>,
original_max_position_embeddings: usize,
max_seq_len_cached: usize,
scaling_factor: f64,
inv_freq: Tensor,
cos_cached: Tensor,
sin_cached: Tensor,
device: Device,
dtype: DType,
}
impl MiniCPMLongRoPE {
pub fn new(cfg: &VoxMiniCPM4Config, device: &Device, dtype: DType) -> Result<Self> {
let head_dim = cfg.hidden_size / cfg.num_attention_heads;
let rope_theta = cfg.rope_theta;
let short_factor = cfg.rope_scaling.short_factor.clone();
let long_factor = cfg.rope_scaling.short_factor.clone();
let original_max_position_embeddings = cfg.rope_scaling.original_max_position_embeddings;
let max_position_embeddings = cfg.max_position_embeddings;
let scale = max_position_embeddings as f64 / original_max_position_embeddings as f64;
let scaling_factor =
(1.0 + scale.ln() / (original_max_position_embeddings as f64).ln()).sqrt();
let inv_freq = compute_default_rope_parameters(head_dim, rope_theta);
let inv_freq = Tensor::from_slice(&inv_freq, (1, inv_freq.len()), device)?;
let max_seq_len_cached = max_position_embeddings;
let t = Tensor::arange(0.0_f32, max_position_embeddings as f32, device)?
.reshape((max_position_embeddings, 1))?;
// short_factor.len() = 32
// head_dim = 1024 / 16 = 64, inv_freq.len() = 32
let ext_factors = Tensor::from_slice(&short_factor, (1, short_factor.len()), device)?;
let ext_factors = Tensor::ones_like(&ext_factors)?.div(&ext_factors)?;
// (seq_len, 1) matmul (1, 32) -> (seq_len, 32) * (1, 32)-> (seq_len, 32)
let freqs = t.matmul(&ext_factors)?.broadcast_mul(&inv_freq)?;
let emb = Tensor::cat(&[&freqs, &freqs], D::Minus1)?;
let cos_cached = emb.cos()?.affine(scaling_factor, 0.0)?.to_dtype(dtype)?;
let sin_cached = emb.sin()?.affine(scaling_factor, 0.0)?.to_dtype(dtype)?;
Ok(Self {
short_factor,
long_factor,
original_max_position_embeddings,
max_seq_len_cached,
scaling_factor,
inv_freq,
cos_cached,
sin_cached,
device: device.clone(),
dtype,
})
}
pub fn update_cos_sin_cache(&mut self, seqlen: usize) -> Result<()> {
self.max_seq_len_cached = seqlen;
let t = Tensor::arange(0.0_f32, seqlen as f32, &self.device)?.reshape((seqlen, 1))?;
let mut ext_factors = Tensor::from_slice(
&self.short_factor,
(1, self.short_factor.len()),
&self.device,
)?;
if seqlen > self.original_max_position_embeddings {
ext_factors =
Tensor::from_slice(&self.long_factor, (1, self.long_factor.len()), &self.device)?;
}
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)?;
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)?;
}
let cos = self.cos_cached.narrow(0, pos_offset, seqlen)?;
let sin = self.sin_cached.narrow(0, pos_offset, seqlen)?;
Ok((cos, sin))
}
}
pub struct MiniCPMDecoderLayer {
self_attn: AttentionNobias,
mlp: MLPNoBias,
input_layernorm: RmsNorm,
post_attention_layernorm: RmsNorm,
scale_depth: f32,
num_hidden_layers: usize,
use_mup: bool,
}
impl MiniCPMDecoderLayer {
pub fn new(vb: VarBuilder, cfg: &VoxMiniCPM4Config) -> Result<Self> {
let self_attn = AttentionNobias::new(
vb.pp("self_attn"),
cfg.hidden_size,
cfg.num_attention_heads,
cfg.num_key_value_heads,
)?;
let mlp = MLPNoBias::new(
vb.pp("mlp"),
cfg.hidden_size,
cfg.intermediate_size,
candle_nn::Activation::Silu,
)?;
let input_layernorm =
rms_norm(cfg.hidden_size, cfg.rms_norm_eps, vb.pp("input_layernorm"))?;
let post_attention_layernorm = rms_norm(
cfg.hidden_size,
cfg.rms_norm_eps,
vb.pp("post_attention_layernorm"),
)?;
Ok(Self {
self_attn,
mlp,
input_layernorm,
post_attention_layernorm,
scale_depth: cfg.scale_depth,
num_hidden_layers: cfg.num_hidden_layers,
use_mup: cfg.use_mup,
})
}
pub fn forward(
&self,
xs: &Tensor,
cos: &Tensor,
sin: &Tensor,
attention_mask: Option<&Tensor>,
) -> Result<Tensor> {
let residual = xs.clone();
let xs = self.input_layernorm.forward(xs)?;
let xs = self
.self_attn
.forward(&xs, cos, sin, attention_mask, true)?;
let xs = if self.use_mup {
let res_add = (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
};
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
+ 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
};
Ok(xs)
}
pub fn forward_step(
&mut self,
xs: &Tensor,
cos: &Tensor,
sin: &Tensor,
attention_mask: Option<&Tensor>,
) -> Result<Tensor> {
let residual = xs.clone();
let xs = self.input_layernorm.forward(xs)?;
let xs = self
.self_attn
.forward_step(&xs, cos, sin, attention_mask, true)?;
let xs = if self.use_mup {
let res_add = (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
};
let residual = &xs;
let xs = xs.apply(&self.post_attention_layernorm)?.apply(&self.mlp)?;
let xs = if self.use_mup {
let res_add = (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
};
Ok(xs)
}
pub fn clear_kv_cache(&mut self) {
self.self_attn.clear_kv_cache();
}
}
pub struct MiniCPMModel {
cfg: VoxMiniCPM4Config,
pub embed_tokens: Option<Embedding>,
layers: Vec<MiniCPMDecoderLayer>,
norm: RmsNorm,
rope_emb: MiniCPMLongRoPE,
// lm_head: Linear,
}
impl MiniCPMModel {
pub fn new(vb: VarBuilder, cfg: VoxMiniCPM4Config) -> Result<Self> {
// let vb = vb.pp("model");
let embed_tokens = if cfg.vocab_size > 0 {
Some(embedding(
cfg.vocab_size,
cfg.hidden_size,
vb.pp("embed_tokens"),
)?)
} else {
None
};
let mut layers = Vec::with_capacity(cfg.num_hidden_layers);
let vb_layers = vb.pp("layers");
for i in 0..cfg.num_hidden_layers {
let layer = MiniCPMDecoderLayer::new(vb_layers.pp(i), &cfg)?;
layers.push(layer);
}
let norm = rms_norm(cfg.hidden_size, cfg.rms_norm_eps, vb.pp("norm"))?;
let rope_emb = MiniCPMLongRoPE::new(&cfg, vb.device(), vb.dtype())?;
// let lm_head = Linear::new(embed_tokens.embeddings().clone(), None);
Ok(Self {
cfg,
embed_tokens,
layers,
norm,
rope_emb,
// lm_head,
})
}
pub fn forward(&mut self, input_embeds: &Tensor, position_id: usize, is_causal: bool) -> Result<Tensor> {
let (bs, seq_len, _) = input_embeds.dims3()?;
// let input_embeds = self
// .embed_tokens
// .forward(&input_ids)?
// .affine(self.cfg.scale_emb, 0.0)?;
let attention_mask: Option<&Tensor> = {
if !is_causal || seq_len <= 1 {
None
} else {
Some(&prepare_causal_attention_mask(
bs,
seq_len,
position_id,
input_embeds.device(),
)?)
}
};
let (cos, sin) = self.rope_emb.forward(position_id, seq_len)?;
let mut hidden_states = input_embeds.clone();
for decode_layer in &self.layers {
hidden_states = decode_layer.forward(&hidden_states, &cos, &sin, attention_mask)?;
}
hidden_states = self.norm.forward(&hidden_states)?;
Ok(hidden_states)
}
pub fn forward_step(&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"))
};
let (bs, seq_len, _) = input_embeds.dims3()?;
// let input_embeds = self
// .embed_tokens
// .forward(&input_ids)?
// .affine(self.cfg.scale_emb, 0.0)?;
let attention_mask: Option<&Tensor> = {
if seq_len <= 1 {
None
} else {
Some(&prepare_causal_attention_mask(
bs,
seq_len,
position_id,
input_embeds.device(),
)?)
}
};
let (cos, sin) = self.rope_emb.forward(position_id, seq_len)?;
let mut hidden_states = input_embeds.clone();
for decode_layer in &mut self.layers {
hidden_states =
decode_layer.forward_step(&hidden_states, &cos, &sin, attention_mask)?;
}
hidden_states = self.norm.forward(&hidden_states)?;
Ok(hidden_states)
}
pub fn clear_kv_cache(&mut self) {
for layer in self.layers.iter_mut() {
layer.clear_kv_cache()
}
}
}
src/models/voxcpm/mod.rs
+5
View File
@@ -0,0 +1,5 @@
pub mod config;
pub mod audio_vae;
pub mod minicpm4;
pub mod tokenizer;
pub mod model;
src/models/voxcpm/model.rs
+730
View File
@@ -0,0 +1,730 @@
use std::{cmp::max, f64, thread, time};
use anyhow::{Ok, Result};
use candle_core::{D, DType, Device, IndexOp, Tensor};
use candle_nn::{Linear, Module, VarBuilder, linear, linear_no_bias};
use candle_transformers::models::deepseek2::SplitOp;
use crate::{
models::voxcpm::{
audio_vae::{self, AudioVAE},
config::{CfmConfig, VoxCPMConfig, VoxMiniCPM4Config},
minicpm4::MiniCPMModel,
tokenizer::SingleChineseTokenizer,
},
utils::{audio_utils::load_audio_with_resample, tensor_utils::linspace},
};
pub struct ScalarQuantizationLayer {
scale: usize,
in_proj: Linear,
out_proj: Linear,
}
impl ScalarQuantizationLayer {
pub fn new(
vb: VarBuilder,
in_dim: usize,
out_dim: usize,
laten_dim: usize,
scale: usize,
) -> Result<Self> {
let in_proj = linear(in_dim, laten_dim, vb.pp("in_proj"))?;
let out_proj = linear(laten_dim, out_dim, vb.pp("out_proj"))?;
Ok(Self {
scale,
in_proj,
out_proj,
})
}
pub fn forward(&self, xs: &Tensor) -> Result<Tensor> {
let xs = self.in_proj.forward(xs)?;
let xs = xs.tanh()?;
let xs = xs
.affine(self.scale as f64, 0.0)?
.round()?
.affine(1.0 / self.scale as f64, 0.0)?;
let xs = self.out_proj.forward(&xs)?;
Ok(xs)
}
}
pub struct SinusoidalPosEmb {
dim: usize,
}
impl SinusoidalPosEmb {
pub fn new(dim: usize) -> Result<Self> {
assert_eq!(dim % 2, 0, "SinusoidalPosEmb requires dim to be even");
Ok(Self { dim })
}
pub fn forward(&self, x: &Tensor, scale: usize) -> Result<Tensor> {
let x = if x.rank() < 1 {
x.unsqueeze(0)?
} else {
x.clone()
};
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)?
.exp()?
.to_dtype(x.dtype())?;
let emb = x
.unsqueeze(D::Minus1)?
.contiguous()?
.matmul(&emb.unsqueeze(0)?.contiguous()?)?
.affine(scale as f64, 0.0)?;
let emb = Tensor::cat(&[emb.sin()?, emb.cos()?], D::Minus1)?;
Ok(emb)
}
}
pub struct TimestepEmbedding {
linear_1: Linear,
linear_2: Linear,
}
impl TimestepEmbedding {
pub fn new(
vb: VarBuilder,
in_channels: usize,
time_embed_dim: usize,
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()
} else {
time_embed_dim
};
let linear_2 = linear(time_embed_dim, time_embed_dim_out, vb.pp("linear_2"))?;
Ok(Self { linear_1, linear_2 })
}
pub fn forward(&self, sample: &Tensor) -> Result<Tensor> {
let sample = self.linear_1.forward(&sample)?.silu()?;
let sample = self.linear_2.forward(&sample)?;
Ok(sample)
}
}
pub struct VoxCPMLocDiT {
in_proj: Linear,
cond_proj: Linear,
out_proj: Linear,
time_embeddings: SinusoidalPosEmb,
time_mlp: TimestepEmbedding,
delta_time_mlp: TimestepEmbedding,
decoder: MiniCPMModel,
config: VoxMiniCPM4Config,
in_channels: usize,
}
impl VoxCPMLocDiT {
pub fn new(vb: VarBuilder, config: VoxMiniCPM4Config, in_channels: usize) -> Result<Self> {
let in_proj = linear(in_channels, config.hidden_size, vb.pp("in_proj"))?;
let cond_proj = linear(in_channels, config.hidden_size, vb.pp("cond_proj"))?;
let out_proj = linear(config.hidden_size, in_channels, vb.pp("out_proj"))?;
let time_embeddings = SinusoidalPosEmb::new(config.hidden_size)?;
let time_mlp = TimestepEmbedding::new(
vb.pp("time_mlp"),
config.hidden_size,
config.hidden_size,
None,
)?;
let delta_time_mlp = TimestepEmbedding::new(
vb.pp("delta_time_mlp"),
config.hidden_size,
config.hidden_size,
None,
)?;
assert_eq!(config.vocab_size, 0, "vocab_size must be 0 for local DiT");
let decoder = MiniCPMModel::new(vb.pp("decoder"), config.clone())?;
Ok(Self {
in_proj,
cond_proj,
out_proj,
time_embeddings,
time_mlp,
delta_time_mlp,
decoder,
config,
in_channels,
})
}
pub fn forward(
&mut self,
x: &Tensor,
mu: &Tensor,
t: &Tensor,
cond: &Tensor,
dt: &Tensor,
) -> Result<Tensor> {
let x = self.in_proj.forward(&x.transpose(1, 2)?.contiguous()?)?;
let cond = self
.cond_proj
.forward(&cond.transpose(1, 2)?.contiguous()?)?;
let prefix = cond.dims()[1];
let t = self.time_embeddings.forward(t, 1000)?.to_dtype(x.dtype())?;
let t = self.time_mlp.forward(&t)?;
let dt = self
.time_embeddings
.forward(dt, 1000)?
.to_dtype(x.dtype())?;
let dt = self.delta_time_mlp.forward(&dt)?;
let t = t.add(&dt)?;
let x = Tensor::cat(&[mu.add(&t)?.unsqueeze(1)?, cond, x], 1)?;
let hidden = self.decoder.forward(&x, 0, false)?;
let select_len = hidden.dims()[1] - (prefix + 1);
let hidden = hidden.narrow(1, prefix + 1, select_len)?;
let hidden = self.out_proj.forward(&hidden)?;
let hidden = hidden.transpose(1, 2)?.contiguous()?;
Ok(hidden)
}
}
pub struct UnifiedCFM {
solver: String,
sigma_min: f32,
t_scheduler: String,
in_channels: usize,
mean_mode: bool,
estimator: VoxCPMLocDiT,
}
impl UnifiedCFM {
pub fn new(
in_channels: usize,
cfm_params: CfmConfig,
estimator: VoxCPMLocDiT,
mean_mode: bool,
) -> Result<Self> {
let solver = cfm_params.solver;
let sigma_min = cfm_params.sigma_min;
let t_scheduler = cfm_params.t_scheduler;
Ok(Self {
solver,
sigma_min,
t_scheduler,
in_channels,
mean_mode,
estimator,
})
}
pub fn forward(
&mut self,
mu: &Tensor,
n_timesteps: usize,
patch_size: usize,
cond: &Tensor,
temperature: f64,
cfg_value: f64,
sway_sampling_coef: f64,
use_cfg_zero_star: bool,
) -> Result<Tensor> {
let (b, c) = mu.dims2()?;
let t = patch_size;
let dtype = mu.dtype();
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)?
.cos()?
.affine(1.0, -1.0)?
.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)
}
pub fn optimized_scale(
&self,
positive_flat: &Tensor,
negative_flat: &Tensor,
) -> Result<Tensor> {
let dot_product = positive_flat.mul(negative_flat)?.sum_keepdim(1)?;
let squared_norm = negative_flat.powf(2.0)?.sum_keepdim(1)?.affine(1.0, 1e-8)?;
let st_star = dot_product.div(&squared_norm)?;
Ok(st_star)
}
pub fn solve_euler(
&mut self,
x: &Tensor,
t_span: &Tensor,
mu: &Tensor,
cond: &Tensor,
cfg_value: f64,
use_cfg_zero_star: bool,
) -> Result<Tensor> {
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 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();
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())?;
} else {
let b = x.dim(0)?;
// let x_in = Tensor::zeros((2*b, self.in_channels, x.dim(2)?), x.dtype(), x.device())?;
let x_in = Tensor::cat(&[x.clone(), x.clone()], 0)?;
let mu_in = Tensor::zeros((b, mu.dim(1)?), x.dtype(), x.device())?;
let mu_in = Tensor::cat(&[mu.clone(), mu_in], 0)?;
let t_in = t.broadcast_as(2 * b)?;
let dt_in = if self.mean_mode {
dt.broadcast_as(2 * b)?
} else {
Tensor::zeros(2 * b, x.dtype(), x.device())?
};
let cond_in = Tensor::cat(&[cond, cond], 0)?;
dphi_dt = self
.estimator
.forward(&x_in, &mu_in, &t_in, &cond_in, &dt_in)?;
let split = dphi_dt.split(&[b, b], 0)?;
dphi_dt = split[0].clone();
let cfg_dphi_dt = split[1].clone();
let mut st_star = Tensor::ones(1, x.dtype(), x.device())?;
if use_cfg_zero_star {
let positive_flat = dphi_dt.reshape((b, ()))?;
let negative_flat = cfg_dphi_dt.reshape((b, ()))?;
st_star = self.optimized_scale(&positive_flat, &negative_flat)?;
let mut vec_shape = vec![b];
let vec_shape1 = vec![1; dphi_dt.rank() - 1];
vec_shape.extend_from_slice(&vec_shape1);
st_star = st_star.reshape(vec_shape)?;
}
let cfg = cfg_dphi_dt.broadcast_mul(&st_star)?;
dphi_dt = cfg.add(&dphi_dt.sub(&cfg)?.affine(cfg_value, 0.0)?)?;
}
x = x.broadcast_sub(&dphi_dt.broadcast_mul(&dt)?)?;
t = t.sub(&dt)?;
sol.push(x.clone());
if step < t_span_len - 1 {
dt = t.sub(&t_span.i(step + 1)?)?;
}
}
Ok(sol[sol.len() - 1].clone())
}
}
pub struct VoxCPMLocEnc {
special_token: Tensor,
in_proj: Linear,
encoder: MiniCPMModel,
hidden_size: usize,
}
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!(
config.vocab_size, 0,
"vocab_size must be 0 for local encoder"
);
let hidden_size = config.hidden_size;
let encoder = MiniCPMModel::new(vb.pp("encoder"), config)?;
Ok(Self {
special_token,
in_proj,
encoder,
hidden_size,
})
}
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)
}
}
pub struct VoxCPMModel {
config: VoxCPMConfig,
patch_size: usize,
audio_start_token: usize,
audio_end_token: usize,
chunk_size: usize,
sample_rate: usize,
tokenizer: SingleChineseTokenizer,
audio_vae: AudioVAE,
base_lm: MiniCPMModel,
residual_lm: MiniCPMModel,
feat_encoder: VoxCPMLocEnc,
feat_decoder: UnifiedCFM,
fsq_layer: ScalarQuantizationLayer,
enc_to_lm_proj: Linear,
lm_to_dit_proj: Linear,
res_to_dit_proj: Linear,
stop_proj: Linear,
stop_head: Linear,
device: Device,
dtype: DType,
}
impl VoxCPMModel {
pub fn new(
vb: VarBuilder,
config: VoxCPMConfig,
tokenizer: SingleChineseTokenizer,
audio_vae: AudioVAE,
) -> Result<Self> {
let base_lm = MiniCPMModel::new(vb.pp("base_lm"), config.lm_config.clone())?;
let audio_start_token = 101usize;
let audio_end_token = 102usize;
let mut residual_lm_config = config.lm_config.clone();
residual_lm_config.num_hidden_layers = config.residual_lm_num_layers;
residual_lm_config.vocab_size = 0;
let residual_lm = MiniCPMModel::new(vb.pp("residual_lm"), residual_lm_config)?;
let mut encoder_config = config.lm_config.clone();
encoder_config.hidden_size = config.encoder_config.hidden_dim;
encoder_config.intermediate_size = config.encoder_config.ffn_dim;
encoder_config.num_attention_heads = config.encoder_config.num_heads;
encoder_config.num_hidden_layers = config.encoder_config.num_layers;
encoder_config.vocab_size = 0;
let feat_encoder =
VoxCPMLocEnc::new(vb.pp("feat_encoder"), encoder_config, config.feat_dim)?;
let mut decoder_config = config.lm_config.clone();
decoder_config.hidden_size = config.dit_config.hidden_dim;
decoder_config.intermediate_size = config.dit_config.ffn_dim;
decoder_config.num_attention_heads = config.dit_config.num_heads;
decoder_config.num_hidden_layers = config.dit_config.num_layers;
decoder_config.vocab_size = 0;
let estimator = VoxCPMLocDiT::new(
vb.pp("feat_decoder.estimator"),
decoder_config,
config.feat_dim,
)?;
let feat_decoder = UnifiedCFM::new(
config.feat_dim,
config.dit_config.cfm_config.clone(),
estimator,
false,
)?;
let fsq_layer = ScalarQuantizationLayer::new(
vb.pp("fsq_layer"),
config.lm_config.hidden_size,
config.lm_config.hidden_size,
config.scalar_quantization_latent_dim,
config.scalar_quantization_scale,
)?;
let enc_to_lm_proj = linear(
config.encoder_config.hidden_dim,
config.lm_config.hidden_size,
vb.pp("enc_to_lm_proj"),
)?;
let lm_to_dit_proj = linear(
config.lm_config.hidden_size,
config.dit_config.hidden_dim,
vb.pp("lm_to_dit_proj"),
)?;
let res_to_dit_proj = linear(
config.lm_config.hidden_size,
config.dit_config.hidden_dim,
vb.pp("res_to_dit_proj"),
)?;
let stop_proj = linear(
config.lm_config.hidden_size,
config.lm_config.hidden_size,
vb.pp("stop_proj"),
)?;
let stop_head = linear_no_bias(config.lm_config.hidden_size, 2, vb.pp("stop_head"))?;
let patch_size = config.patch_size;
Ok(Self {
config,
patch_size,
audio_start_token,
audio_end_token,
chunk_size: audio_vae.chunk_size,
sample_rate: audio_vae.sample_rate,
tokenizer,
audio_vae,
base_lm,
residual_lm,
feat_encoder,
feat_decoder,
fsq_layer,
enc_to_lm_proj,
lm_to_dit_proj,
res_to_dit_proj,
stop_proj,
stop_head,
device: vb.device().clone(),
dtype: vb.dtype(),
})
}
pub fn generate(
&mut self,
target_text: String,
prompt_text: Option<String>,
prompt_wav_path: Option<String>,
min_len: usize,
max_len: usize,
inference_timesteps: usize,
cfg_value: f64,
retry_badcase: bool,
retry_badcase_max_times: usize,
retry_badcase_ratio_threshold: f64,
) -> Result<Tensor> {
let (text_token, text_mask, audio_feat, audio_mask) = match prompt_wav_path {
None => {
let text_token = self.tokenizer.encode(target_text.clone())?;
let text_token = Tensor::from_slice(&text_token, text_token.len(), &self.device)?;
let audio_start = Tensor::new(vec![self.audio_start_token as u32], &self.device)?;
let text_token = Tensor::cat(&[text_token, audio_start], D::Minus1)?;
let text_length = text_token.dim(0)?;
let audio_feat = Tensor::zeros(
(text_length, self.patch_size, self.audio_vae.latent_dim),
DType::F32,
&self.device,
)?;
let text_mask = Tensor::ones(text_length, self.dtype, &self.device)?;
let audio_mask = Tensor::zeros(text_length, self.dtype, &self.device)?;
(text_token, text_mask, audio_feat, audio_mask)
}
Some(path) => {
let text = prompt_text.unwrap_or("".to_string()) + &target_text;
let text_token = self.tokenizer.encode(text)?;
let text_token = Tensor::from_slice(&text_token, text_token.len(), &self.device)?;
let audio_start = Tensor::new(vec![self.audio_start_token as u32], &self.device)?;
let text_token = Tensor::cat(&[text_token, audio_start], D::Minus1)?;
let text_length = text_token.dim(0)?;
let mut audio =
load_audio_with_resample(path, self.device.clone(), Some(self.sample_rate))?;
let patch_len = self.patch_size * self.chunk_size;
if audio.dim(1)? % patch_len != 0 {
audio = audio.pad_with_zeros(
D::Minus1,
0,
patch_len - audio.dim(1)? % patch_len,
)?;
}
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)?;
println!("audio_feat: {:?}", audio_feat);
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)?;
let audio_pad_feat = Tensor::zeros(
(text_length, self.patch_size, self.audio_vae.latent_dim),
audio_feat.dtype(),
&self.device,
)?;
let audio_feat = Tensor::cat(&[audio_pad_feat, audio_feat], 0)?;
let text_mask = Tensor::cat(
&[
Tensor::ones(text_length, self.dtype, &self.device)?,
Tensor::zeros(audio_length, self.dtype, &self.device)?,
],
D::Minus1,
)?;
let audio_mask = Tensor::cat(
&[
Tensor::zeros(text_length, self.dtype, &self.device)?,
Tensor::ones(audio_length, self.dtype, &self.device)?,
],
D::Minus1,
)?;
(text_token, text_mask, audio_feat, audio_mask)
}
};
let text_token = text_token.unsqueeze(0)?;
let text_mask = text_mask.unsqueeze(0)?;
let audio_feat = audio_feat.unsqueeze(0)?.to_dtype(DType::BF16)?;
let audio_mask = audio_mask.unsqueeze(0)?;
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 latent_pred = self.inference(
&text_token,
&text_mask,
&audio_feat,
&audio_mask,
min_len,
max_len,
inference_timesteps,
cfg_value,
)?;
let decode_audio = self
.audio_vae
.decode(&latent_pred.to_dtype(DType::F32)?)?
.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)
}
pub fn inference(
&mut self,
text: &Tensor,
text_mask: &Tensor,
feat: &Tensor,
feat_mask: &Tensor,
min_len: usize,
max_len: usize,
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
.as_ref()
.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
.feat_decoder
.forward(
&dit_hidden,
inference_timesteps,
self.patch_size,
&cond,
1.0,
cfg_value,
1.0,
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;
}
position_id += seq_len;
seq_len = 1;
lm_hidden = self
.base_lm
.forward_step(&curr_embed.i((.., 0, ..))?, position_id)?
.squeeze(1)?;
lm_hidden = self.fsq_layer.forward(&lm_hidden)?;
residual_hidden = self
.residual_lm
.forward_step(&lm_hidden.add(&curr_embed.i((.., 0, ..))?)?, position_id)?
.squeeze(1)?;
}
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
+66
View File
@@ -0,0 +1,66 @@
use anyhow::{Ok, Result, anyhow};
use candle_core::Tensor;
use tokenizers::Tokenizer;
pub struct SingleChineseTokenizer {
tokenizer: Tokenizer,
multichar_tokens: Vec<String>,
}
impl SingleChineseTokenizer {
pub fn new(path: &str) -> Result<Self> {
let path = path.to_string();
assert!(
std::path::Path::new(&path).exists(),
"model path file not exists"
);
let tokenizer_file = path.clone() + "/tokenizer.json";
assert!(
std::path::Path::new(&tokenizer_file).exists(),
"tokenizer.json not exists in model path"
);
let tokenizer = Tokenizer::from_file(tokenizer_file)
.map_err(|e| anyhow!(format!("tokenizer from file error{}", e)))?;
let mut multichar_tokens = Vec::new();
for (token, _) in tokenizer.get_vocab(false) {
let len = token.chars().count();
if len >= 2 {
let is_chinese = token.chars().all(|c| {
let c_ = c as u32;
0x4E00 <= c_ && c_ <= 0x9FFF
});
if is_chinese {
multichar_tokens.push(token);
}
}
}
Ok(Self {
tokenizer,
multichar_tokens,
})
}
pub fn encode(&self, text: String) -> Result<Vec<u32>> {
let encode = self
.tokenizer
.encode(text, false)
.map_err(|e| anyhow!(format!("tokenizer encode error: {}", e)))?;
let tokens = encode.get_tokens();
println!("tokens: {:?}", tokens);
let mut split_character = Vec::new();
for token in tokens {
let clean_token = token.replace("▁", "to");
if self.multichar_tokens.contains(&clean_token) {
let chars: Vec<String> = clean_token.chars().map(|c| c.to_string()).collect();
split_character.extend(chars);
} else {
split_character.push(token.clone());
}
}
println!("split_character: {:?}", split_character);
let ids: Vec<u32> = split_character
.iter()
.filter_map(|c| self.tokenizer.token_to_id(c))
.collect();
Ok(ids)
}
}