Merge branch 'pr/ariesdevil/1' into fmt
This commit is contained in:
@@ -1,5 +1,5 @@
|
||||
use anyhow::Result;
|
||||
use candle_core::{Tensor, D};
|
||||
use candle_core::{D, Tensor};
|
||||
use candle_nn::{Activation, Linear, Module, VarBuilder, linear, linear_no_bias};
|
||||
|
||||
use crate::{position_embed::rope::apply_rotary_pos_emb, utils::tensor_utils::repeat_kv};
|
||||
@@ -89,7 +89,12 @@ pub struct AttentionNobias {
|
||||
}
|
||||
|
||||
impl AttentionNobias {
|
||||
pub fn new(vb: VarBuilder, hidden_size: usize, num_attention_heads: usize, num_key_value_heads: usize) -> Result<Self> {
|
||||
pub fn new(
|
||||
vb: VarBuilder,
|
||||
hidden_size: usize,
|
||||
num_attention_heads: usize,
|
||||
num_key_value_heads: usize,
|
||||
) -> Result<Self> {
|
||||
let num_kv_groups = num_attention_heads / num_key_value_heads;
|
||||
let head_dim = hidden_size / num_attention_heads;
|
||||
let q_proj = linear_no_bias(hidden_size, num_attention_heads * head_dim, vb.pp("q_proj"))?;
|
||||
@@ -146,11 +151,12 @@ impl AttentionNobias {
|
||||
let attn_weights = (attn_weights * scale)?;
|
||||
let attn_weights = match attention_mask {
|
||||
None => attn_weights,
|
||||
Some(mask) => attn_weights.broadcast_add(&mask.to_dtype(attn_weights.dtype())?)?,
|
||||
Some(mask) => {
|
||||
attn_weights.broadcast_add(&mask.to_dtype(attn_weights.dtype())?)?
|
||||
}
|
||||
};
|
||||
let attn_weights = candle_nn::ops::softmax_last_dim(&attn_weights)?;
|
||||
let attn_weights = attn_weights.matmul(&value_states)?;
|
||||
attn_weights
|
||||
attn_weights.matmul(&value_states)?
|
||||
}
|
||||
#[cfg(feature = "flash-attn")]
|
||||
{
|
||||
@@ -225,11 +231,12 @@ impl AttentionNobias {
|
||||
let attn_weights = (attn_weights * scale)?;
|
||||
let attn_weights = match attention_mask {
|
||||
None => attn_weights,
|
||||
Some(mask) => attn_weights.broadcast_add(&mask.to_dtype(attn_weights.dtype())?)?,
|
||||
Some(mask) => {
|
||||
attn_weights.broadcast_add(&mask.to_dtype(attn_weights.dtype())?)?
|
||||
}
|
||||
};
|
||||
let attn_weights = candle_nn::ops::softmax_last_dim(&attn_weights)?;
|
||||
let attn_weights = attn_weights.matmul(&value_states)?;
|
||||
attn_weights
|
||||
attn_weights.matmul(&value_states)?
|
||||
}
|
||||
#[cfg(feature = "flash-attn")]
|
||||
{
|
||||
|
||||
@@ -26,4 +26,4 @@ pub struct MiniCPM4Config {
|
||||
pub scale_emb: f64,
|
||||
pub dim_model_base: usize,
|
||||
pub scale_depth: f32,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,13 +1,3 @@
|
||||
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_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::{DType, Device, Tensor};
|
||||
use candle_nn::VarBuilder;
|
||||
@@ -17,6 +7,15 @@ use openai_dive::v1::resources::chat::{
|
||||
use rocket::async_stream::stream;
|
||||
use rocket::futures::Stream;
|
||||
|
||||
use crate::models::minicpm4::config::MiniCPM4Config;
|
||||
use crate::models::minicpm4::model::MiniCPMModel;
|
||||
// use crate::models::GenerateStream;
|
||||
use crate::utils::{
|
||||
build_completion_chunk_response, build_completion_response, find_type_files, get_device,
|
||||
get_dtype, get_logit_processor,
|
||||
};
|
||||
use crate::{chat_template::ChatTemplate, models::GenerateModel, tokenizer::TokenizerModel};
|
||||
|
||||
pub struct MiniCPMGenerateModel<'a> {
|
||||
chat_template: ChatTemplate<'a>,
|
||||
tokenizer: TokenizerModel,
|
||||
@@ -26,7 +25,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)?;
|
||||
@@ -37,7 +36,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_type_files(&path, "safetensors")?;
|
||||
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)?;
|
||||
|
||||
@@ -53,7 +52,6 @@ impl <'a> MiniCPMGenerateModel<'a> {
|
||||
}
|
||||
|
||||
impl<'a> GenerateModel for MiniCPMGenerateModel<'a> {
|
||||
|
||||
fn generate(&mut self, mes: ChatCompletionParameters) -> Result<ChatCompletionResponse> {
|
||||
let mut logit_processor = get_logit_processor(mes.temperature, mes.top_p);
|
||||
let mes_render = self.chat_template.apply_chat_template(&mes)?;
|
||||
@@ -61,10 +59,7 @@ impl<'a> GenerateModel for MiniCPMGenerateModel<'a> {
|
||||
let mut seq_len = input_ids.dim(1)?;
|
||||
let mut seqlen_offset = 0;
|
||||
let mut generate = Vec::new();
|
||||
let sample_len = match mes.max_tokens {
|
||||
Some(max) => max,
|
||||
None => 2048,
|
||||
};
|
||||
let sample_len = mes.max_tokens.unwrap_or(2048);
|
||||
for _ in 0..sample_len {
|
||||
let logits = self.minicpm.forward_with_cache(&input_ids, seqlen_offset)?;
|
||||
let logits = logits.squeeze(0)?.squeeze(0)?.to_dtype(DType::F32)?;
|
||||
@@ -91,10 +86,7 @@ impl<'a> GenerateModel for MiniCPMGenerateModel<'a> {
|
||||
let mut input_ids = self.tokenizer.text_encode(mes_render, &self.device)?;
|
||||
let mut seq_len = input_ids.dim(1)?;
|
||||
let mut seqlen_offset = 0;
|
||||
let sample_len = match mes.max_tokens {
|
||||
Some(max) => max,
|
||||
None => 512,
|
||||
};
|
||||
let sample_len = mes.max_tokens.unwrap_or(512);
|
||||
let stream = stream! {
|
||||
let mut error_tokens = Vec::new();
|
||||
for _ in 0..sample_len {
|
||||
@@ -105,7 +97,7 @@ impl<'a> GenerateModel for MiniCPMGenerateModel<'a> {
|
||||
let logits = logits.squeeze(0)?.squeeze(0)?.to_dtype(DType::F32)?;
|
||||
let next_token = logit_processor.sample(&logits)?;
|
||||
let mut decode_ids = Vec::new();
|
||||
if error_tokens.len() > 0 {
|
||||
if !error_tokens.is_empty(){
|
||||
decode_ids.extend_from_slice(&error_tokens);
|
||||
}
|
||||
decode_ids.push(next_token);
|
||||
|
||||
@@ -1,3 +1,3 @@
|
||||
pub mod config;
|
||||
pub mod generate;
|
||||
pub mod model;
|
||||
pub mod generate;
|
||||
@@ -1,3 +1,7 @@
|
||||
use anyhow::{Ok, Result};
|
||||
use candle_core::{D, Device, Tensor};
|
||||
use candle_nn::{Embedding, Linear, Module, RmsNorm, VarBuilder, embedding, rms_norm};
|
||||
|
||||
use crate::{
|
||||
models::{
|
||||
common::{AttentionNobias, MLPNoBias},
|
||||
@@ -6,9 +10,6 @@ use crate::{
|
||||
position_embed::rope::compute_default_rope_parameters,
|
||||
utils::tensor_utils::prepare_causal_attention_mask,
|
||||
};
|
||||
use anyhow::{Ok, Result};
|
||||
use candle_core::{D, Device, Tensor};
|
||||
use candle_nn::{Embedding, Linear, Module, RmsNorm, VarBuilder, embedding, rms_norm};
|
||||
|
||||
pub struct MiniCPMLongRoPE {
|
||||
short_factor: Vec<f32>,
|
||||
@@ -33,15 +34,13 @@ impl MiniCPMLongRoPE {
|
||||
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 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)?;
|
||||
@@ -63,8 +62,7 @@ impl MiniCPMLongRoPE {
|
||||
}
|
||||
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 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()),
|
||||
@@ -85,7 +83,7 @@ impl MiniCPMLongRoPE {
|
||||
}
|
||||
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)?;
|
||||
@@ -143,7 +141,9 @@ impl MiniCPMDecoderLayer {
|
||||
) -> 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 = 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(),
|
||||
@@ -168,7 +168,9 @@ impl MiniCPMDecoderLayer {
|
||||
) -> Result<Tensor> {
|
||||
let residual = xs.clone();
|
||||
let xs = self.input_layernorm.forward(xs)?;
|
||||
let xs = self.self_attn.forward_with_cache(&xs, cos, sin, attention_mask, true)?;
|
||||
let xs = self
|
||||
.self_attn
|
||||
.forward_with_cache(&xs, cos, sin, attention_mask, true)?;
|
||||
let xs = (residual
|
||||
+ xs.affine(
|
||||
self.scale_depth as f64 / (self.num_hidden_layers as f64).sqrt(),
|
||||
@@ -220,11 +222,11 @@ impl MiniCPMModel {
|
||||
})
|
||||
}
|
||||
|
||||
pub fn forward(&mut 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)?
|
||||
.forward(input_ids)?
|
||||
.affine(self.cfg.scale_emb, 0.0)?;
|
||||
let attention_mask: Option<&Tensor> = {
|
||||
if seq_len <= 1 {
|
||||
@@ -238,7 +240,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 {
|
||||
@@ -258,7 +260,7 @@ impl MiniCPMModel {
|
||||
let (bs, seq_len) = input_ids.dims2()?;
|
||||
let input_embeds = self
|
||||
.embed_tokens
|
||||
.forward(&input_ids)?
|
||||
.forward(input_ids)?
|
||||
.affine(self.cfg.scale_emb, 0.0)?;
|
||||
let attention_mask: Option<&Tensor> = {
|
||||
if seq_len <= 1 {
|
||||
|
||||
@@ -72,8 +72,8 @@ pub struct VisionSetting {
|
||||
pub image_std: Vec<f32>,
|
||||
}
|
||||
|
||||
impl VisionSetting {
|
||||
pub fn default() -> Self {
|
||||
impl Default for VisionSetting {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
image_factor: 28,
|
||||
min_pixels: 4 * 28 * 28,
|
||||
@@ -90,8 +90,8 @@ impl VisionSetting {
|
||||
fps: 2.0,
|
||||
fps_min_frames: 4,
|
||||
fps_max_frames: 768,
|
||||
image_mean: vec![0.48145466_f32, 0.4578275, 0.40821073],
|
||||
image_std: vec![0.26862954, 0.26130258, 0.27577711],
|
||||
image_mean: vec![0.48145466_f32, 0.4578275f32, 0.40821073f32],
|
||||
image_std: vec![0.26862954f32, 0.2613026f32, 0.2757771f32],
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,16 +1,4 @@
|
||||
// use crate::models::GenerateStream;
|
||||
use crate::models::qwen2_5vl::config::Qwen2_5VLConfig;
|
||||
use crate::utils::utils::{
|
||||
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,
|
||||
qwen2_5vl::{model::Qwen2_5VLModel, processor::Qwen2_5VLProcessor},
|
||||
},
|
||||
tokenizer::tokenizer::TokenizerModel,
|
||||
};
|
||||
use anyhow::{Result, anyhow};
|
||||
use candle_core::{D, DType, Device, IndexOp, Tensor};
|
||||
use candle_nn::VarBuilder;
|
||||
@@ -20,6 +8,20 @@ use openai_dive::v1::resources::chat::{
|
||||
use rocket::async_stream::stream;
|
||||
use rocket::futures::Stream;
|
||||
|
||||
use crate::models::qwen2_5vl::config::Qwen2_5VLConfig;
|
||||
use crate::utils::{
|
||||
build_completion_chunk_response, build_completion_response, find_type_files, get_device,
|
||||
get_dtype, get_logit_processor,
|
||||
};
|
||||
use crate::{
|
||||
chat_template::ChatTemplate,
|
||||
models::{
|
||||
GenerateModel,
|
||||
qwen2_5vl::{model::Qwen2_5VLModel, processor::Qwen2_5VLProcessor},
|
||||
},
|
||||
tokenizer::TokenizerModel,
|
||||
};
|
||||
|
||||
pub struct Qwen2_5VLGenerateModel<'a> {
|
||||
chat_template: ChatTemplate<'a>,
|
||||
tokenizer: TokenizerModel,
|
||||
@@ -43,7 +45,7 @@ impl<'a> Qwen2_5VLGenerateModel<'a> {
|
||||
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_type_files(&path, "safetensors")?;
|
||||
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)?;
|
||||
|
||||
@@ -60,7 +62,6 @@ impl<'a> Qwen2_5VLGenerateModel<'a> {
|
||||
}
|
||||
|
||||
impl<'a> GenerateModel for Qwen2_5VLGenerateModel<'a> {
|
||||
|
||||
fn generate(&mut self, mes: ChatCompletionParameters) -> Result<ChatCompletionResponse> {
|
||||
let mut logit_processor = get_logit_processor(mes.temperature, mes.top_p);
|
||||
let mes_render = self.chat_template.apply_chat_template(&mes)?;
|
||||
@@ -84,10 +85,7 @@ impl<'a> GenerateModel for Qwen2_5VLGenerateModel<'a> {
|
||||
.broadcast_sub(&Tensor::new(vec![1_u32], input_ids.device())?)?;
|
||||
|
||||
let mut generate = Vec::new();
|
||||
let sample_len = match mes.max_tokens {
|
||||
Some(max) => max,
|
||||
None => 1024,
|
||||
};
|
||||
let sample_len = mes.max_tokens.unwrap_or(1024);
|
||||
for _ in 0..sample_len {
|
||||
let logits = self.qwen2_5_vl.forward(
|
||||
&input_ids,
|
||||
@@ -145,10 +143,7 @@ impl<'a> GenerateModel for Qwen2_5VLGenerateModel<'a> {
|
||||
.to_dtype(candle_core::DType::U32)?
|
||||
.broadcast_sub(&Tensor::new(vec![1_u32], input_ids.device())?)?;
|
||||
|
||||
let sample_len = match mes.max_tokens {
|
||||
Some(max) => max,
|
||||
None => 512,
|
||||
};
|
||||
let sample_len = mes.max_tokens.unwrap_or(512);
|
||||
let stream = stream! {
|
||||
let mut error_tokens = Vec::new();
|
||||
let mut pixel_values = pixel_values.as_ref();
|
||||
@@ -170,7 +165,7 @@ impl<'a> GenerateModel for Qwen2_5VLGenerateModel<'a> {
|
||||
let logits = logits.squeeze(0)?.squeeze(0)?.to_dtype(DType::F32)?;
|
||||
let next_token = logit_processor.sample(&logits)?;
|
||||
let mut decode_ids = Vec::new();
|
||||
if error_tokens.len() > 0 {
|
||||
if !error_tokens.is_empty() {
|
||||
decode_ids.extend_from_slice(&error_tokens);
|
||||
}
|
||||
decode_ids.push(next_token);
|
||||
|
||||
@@ -1,18 +1,21 @@
|
||||
use crate::{
|
||||
models::qwen2_5vl::config::{Qwen2_5VLConfig, RopeScaling},
|
||||
position_embed::rope::{
|
||||
apply_rotary_pos_emb, apply_rotary_pos_emb_vision, Qwen2_5VLTextRotaryEmbedding, Qwen2_5VisionRotaryEmbedding
|
||||
},
|
||||
utils::tensor_utils::{
|
||||
get_equal_mask, get_vision_next_indices, masked_scatter_dim0, nonzero_index, repeat_kv, safe_arg_sort_last_dim, zero_index
|
||||
},
|
||||
};
|
||||
use anyhow::{Result, anyhow};
|
||||
use candle_core::{D, DType, Device, IndexOp, Tensor};
|
||||
use candle_nn::{
|
||||
Activation, Init, Linear, Module, RmsNorm, VarBuilder, linear, linear_no_bias, rms_norm,
|
||||
};
|
||||
|
||||
use crate::{
|
||||
models::qwen2_5vl::config::{Qwen2_5VLConfig, RopeScaling},
|
||||
position_embed::rope::{
|
||||
Qwen2_5VLTextRotaryEmbedding, Qwen2_5VisionRotaryEmbedding, apply_rotary_pos_emb,
|
||||
apply_rotary_pos_emb_vision,
|
||||
},
|
||||
utils::tensor_utils::{
|
||||
get_equal_mask, get_vision_next_indices, masked_scatter_dim0, nonzero_index, repeat_kv,
|
||||
safe_arg_sort_last_dim, zero_index,
|
||||
},
|
||||
};
|
||||
|
||||
pub struct Qwen2_5VisionPatchEmbed {
|
||||
conv3d_weight: Tensor,
|
||||
}
|
||||
@@ -175,7 +178,7 @@ impl Qwen2_5VLVisionAttention {
|
||||
let attn_weights = query_states
|
||||
.matmul(&key_states.transpose(D::Minus2, D::Minus1)?)?
|
||||
.broadcast_mul(&self.scale)?;
|
||||
let attn_weights = attn_weights.broadcast_add(&attention_mask)?;
|
||||
let attn_weights = attn_weights.broadcast_add(attention_mask)?;
|
||||
let attn_weights = candle_nn::ops::softmax_last_dim(&attn_weights)?;
|
||||
attn_weights.matmul(&value_states)?
|
||||
};
|
||||
@@ -495,7 +498,7 @@ impl Qwen2_5VLVisionModel {
|
||||
2 => {
|
||||
let mut cu_seqlens_repeat = Vec::new();
|
||||
for (index, t) in grid_t.iter().enumerate() {
|
||||
cu_seqlens_repeat.push(cu_seqlens.i(index)?.repeat(t.clone() as usize)?);
|
||||
cu_seqlens_repeat.push(cu_seqlens.i(index)?.repeat(*t as usize)?);
|
||||
}
|
||||
Tensor::cat(&cu_seqlens_repeat, 0)?.flatten_all()?
|
||||
}
|
||||
@@ -521,7 +524,7 @@ impl Qwen2_5VLVisionModel {
|
||||
hidden_states.device(),
|
||||
hidden_states.dtype(),
|
||||
)?;
|
||||
let mut attention_mask = attention_mask_window.clone();
|
||||
let mut attention_mask;
|
||||
for (layer_num, block) in self.blocks.iter().enumerate() {
|
||||
if self.fullatt_block_indexes.contains(&layer_num) {
|
||||
attention_mask = attention_mask_full.clone();
|
||||
@@ -537,7 +540,6 @@ impl Qwen2_5VLVisionModel {
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
struct Qwen2_5VLTextMLP {
|
||||
gate_proj: Linear,
|
||||
@@ -658,8 +660,7 @@ impl Qwen2_5VLTextAttention {
|
||||
Some(mask) => attn_weights.broadcast_add(mask)?,
|
||||
};
|
||||
let attn_weights = candle_nn::ops::softmax_last_dim(&attn_weights)?;
|
||||
let attn_weights = attn_weights.matmul(&value_states)?;
|
||||
attn_weights
|
||||
attn_weights.matmul(&value_states)?
|
||||
}
|
||||
#[cfg(feature = "flash-attn")]
|
||||
{
|
||||
@@ -897,12 +898,9 @@ impl Qwen2_5VLModel {
|
||||
let mut mrope_position_deltas: Vec<i64> = Vec::new();
|
||||
if image_grid_thw.is_some() || video_grid_thw.is_some() {
|
||||
let total_input_ids = input_ids.clone();
|
||||
let mut mask_;
|
||||
if mask.is_none() {
|
||||
mask_ = Tensor::ones_like(&total_input_ids)?;
|
||||
} else {
|
||||
mask_ = mask.unwrap().clone();
|
||||
}
|
||||
let mask_ = mask
|
||||
.cloned()
|
||||
.unwrap_or(Tensor::ones_like(&total_input_ids)?);
|
||||
let mut position_ids = Tensor::ones(
|
||||
(3, input_ids.dim(0)?, input_ids.dim(1)?),
|
||||
input_ids.dtype(),
|
||||
@@ -950,7 +948,7 @@ impl Qwen2_5VLModel {
|
||||
let llm_grid_h = thw[1] / spatial_merge_size as u32;
|
||||
let llm_grid_w = thw[2] / spatial_merge_size as u32;
|
||||
let text_len = text_end - text_start;
|
||||
let start_idx = if llm_pos_ids_list.len() > 0 {
|
||||
let start_idx = if !llm_pos_ids_list.is_empty() {
|
||||
llm_pos_ids_list[llm_pos_ids_list.len() - 1]
|
||||
.max_all()?
|
||||
.to_scalar::<u32>()?
|
||||
@@ -1024,7 +1022,7 @@ impl Qwen2_5VLModel {
|
||||
};
|
||||
|
||||
if text_start < input_ids_i.dim(0)? as u32 {
|
||||
let start_idx = if llm_pos_ids_list.len() > 0 {
|
||||
let start_idx = if !llm_pos_ids_list.is_empty() {
|
||||
llm_pos_ids_list[llm_pos_ids_list.len() - 1]
|
||||
.max_all()?
|
||||
.to_scalar::<u32>()?
|
||||
@@ -1051,66 +1049,61 @@ impl Qwen2_5VLModel {
|
||||
if mrope_position_deltas.rank() == 1 {
|
||||
mrope_position_deltas = mrope_position_deltas.unsqueeze(0)?;
|
||||
}
|
||||
return Ok((position_ids.contiguous()?, mrope_position_deltas));
|
||||
} else {
|
||||
if mask.is_some() {
|
||||
let mut position_ids = mask
|
||||
.unwrap()
|
||||
.to_dtype(candle_core::DType::F64)?
|
||||
.cumsum(D::Minus1)?
|
||||
.to_dtype(candle_core::DType::U32)?
|
||||
.broadcast_sub(&Tensor::new(vec![1_u32], input_ids.device())?)?;
|
||||
for i in 0..position_ids.dim(0)? {
|
||||
let mut position_ids_i = position_ids.i(i)?;
|
||||
let mask_i = mask.unwrap().i(i)?;
|
||||
// 如果有pad, 将填充位置置为1
|
||||
// 当bs>1, 可能存在不同序列长度,需要添加pad使seq_len长度一致
|
||||
if mask_i.sum_all()?.to_scalar::<u32>()? != mask_i.dim(0)? as u32 {
|
||||
let zero_indices = zero_index(&mask_i)?;
|
||||
let replace_1 = Tensor::ones(
|
||||
zero_indices.dim(0)?,
|
||||
candle_core::DType::U32,
|
||||
input_ids.device(),
|
||||
)?;
|
||||
position_ids_i = position_ids_i
|
||||
.scatter(&zero_indices, &replace_1, 0)?
|
||||
.unsqueeze(0)?;
|
||||
position_ids = position_ids.slice_assign(
|
||||
&[(i..i + 1), (0..position_ids.dim(1)?)],
|
||||
&position_ids_i,
|
||||
)?;
|
||||
}
|
||||
}
|
||||
position_ids = position_ids
|
||||
.unsqueeze(0)?
|
||||
.broadcast_as((3, input_ids.dim(0)?, input_ids.dim(1)?))?
|
||||
.contiguous()?;
|
||||
let mut mrope_position_deltas = position_ids
|
||||
.max(0)?
|
||||
.max(D::Minus1)?
|
||||
.broadcast_sub(&Tensor::new(
|
||||
vec![mask.unwrap().dim(D::Minus1)? as u32 - 1],
|
||||
Ok((position_ids.contiguous()?, mrope_position_deltas))
|
||||
} else if let Some(mask) = mask {
|
||||
let mut position_ids = mask
|
||||
.to_dtype(candle_core::DType::F64)?
|
||||
.cumsum(D::Minus1)?
|
||||
.to_dtype(candle_core::DType::U32)?
|
||||
.broadcast_sub(&Tensor::new(vec![1_u32], input_ids.device())?)?;
|
||||
for i in 0..position_ids.dim(0)? {
|
||||
let mut position_ids_i = position_ids.i(i)?;
|
||||
let mask_i = mask.i(i)?;
|
||||
// 如果有pad, 将填充位置置为1
|
||||
// 当bs>1, 可能存在不同序列长度,需要添加pad使seq_len长度一致
|
||||
if mask_i.sum_all()?.to_scalar::<u32>()? != mask_i.dim(0)? as u32 {
|
||||
let zero_indices = zero_index(&mask_i)?;
|
||||
let replace_1 = Tensor::ones(
|
||||
zero_indices.dim(0)?,
|
||||
candle_core::DType::U32,
|
||||
input_ids.device(),
|
||||
)?)?
|
||||
.contiguous()?;
|
||||
if mrope_position_deltas.rank() == 1 {
|
||||
mrope_position_deltas = mrope_position_deltas.unsqueeze(0)?;
|
||||
)?;
|
||||
position_ids_i = position_ids_i
|
||||
.scatter(&zero_indices, &replace_1, 0)?
|
||||
.unsqueeze(0)?;
|
||||
position_ids = position_ids
|
||||
.slice_assign(&[(i..i + 1), (0..position_ids.dim(1)?)], &position_ids_i)?;
|
||||
}
|
||||
return Ok((position_ids, mrope_position_deltas));
|
||||
} else {
|
||||
let position_ids =
|
||||
Tensor::arange(0_u32, input_ids.dim(D::Minus1)? as u32, input_ids.device())?
|
||||
.unsqueeze(0)?
|
||||
.unsqueeze(0)?
|
||||
.broadcast_as((3, input_ids.dim(0)?, input_ids.dim(D::Minus1)?))?
|
||||
.contiguous()?;
|
||||
let mrope_position_deltas = Tensor::zeros(
|
||||
(input_ids.dim(0)?, 1),
|
||||
input_ids.dtype(),
|
||||
input_ids.device(),
|
||||
)?;
|
||||
Ok((position_ids, mrope_position_deltas))
|
||||
}
|
||||
position_ids = position_ids
|
||||
.unsqueeze(0)?
|
||||
.broadcast_as((3, input_ids.dim(0)?, input_ids.dim(1)?))?
|
||||
.contiguous()?;
|
||||
let mut mrope_position_deltas = position_ids
|
||||
.max(0)?
|
||||
.max(D::Minus1)?
|
||||
.broadcast_sub(&Tensor::new(
|
||||
vec![mask.dim(D::Minus1)? as u32 - 1],
|
||||
input_ids.device(),
|
||||
)?)?
|
||||
.contiguous()?;
|
||||
if mrope_position_deltas.rank() == 1 {
|
||||
mrope_position_deltas = mrope_position_deltas.unsqueeze(0)?;
|
||||
}
|
||||
Ok((position_ids, mrope_position_deltas))
|
||||
} else {
|
||||
let position_ids =
|
||||
Tensor::arange(0_u32, input_ids.dim(D::Minus1)? as u32, input_ids.device())?
|
||||
.unsqueeze(0)?
|
||||
.unsqueeze(0)?
|
||||
.broadcast_as((3, input_ids.dim(0)?, input_ids.dim(D::Minus1)?))?
|
||||
.contiguous()?;
|
||||
let mrope_position_deltas = Tensor::zeros(
|
||||
(input_ids.dim(0)?, 1),
|
||||
input_ids.dtype(),
|
||||
input_ids.device(),
|
||||
)?;
|
||||
Ok((position_ids, mrope_position_deltas))
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1127,14 +1120,14 @@ impl Qwen2_5VLModel {
|
||||
second_per_grid_ts: Option<Vec<f32>>,
|
||||
) -> Result<Tensor> {
|
||||
// input_ids shape: (bs, seq_len)
|
||||
let mut inputs_embeds = self.model.embed_tokens.forward(&input_ids)?;
|
||||
let mut inputs_embeds = self.model.embed_tokens.forward(input_ids)?;
|
||||
// inputs_embeds shape: (bs, seq_len, hidden_dim)
|
||||
if pixel_values.is_some() && image_grid_thw.is_some() {
|
||||
if let Some(pixel_values) = pixel_values
|
||||
&& let Some(image_grid_thw) = image_grid_thw
|
||||
{
|
||||
// image_embed shape: (seq_len, hidden_dim)
|
||||
let image_embed = self
|
||||
.visual
|
||||
.forward(pixel_values.unwrap(), image_grid_thw.unwrap())?;
|
||||
let vision_mask = get_equal_mask(&input_ids, self.cfg.image_token_id as u32)?;
|
||||
let image_embed = self.visual.forward(pixel_values, image_grid_thw)?;
|
||||
let vision_mask = get_equal_mask(input_ids, self.cfg.image_token_id as u32)?;
|
||||
|
||||
let n_image_tokens = vision_mask.sum_all()?.to_scalar::<u32>()?;
|
||||
if n_image_tokens as usize != image_embed.dim(0)? {
|
||||
@@ -1146,12 +1139,12 @@ impl Qwen2_5VLModel {
|
||||
}
|
||||
inputs_embeds = masked_scatter_dim0(&inputs_embeds, &image_embed, &vision_mask)?;
|
||||
}
|
||||
if pixel_values_video.is_some() && video_grid_thw.is_some() {
|
||||
let video_embed = self
|
||||
.visual
|
||||
.forward(pixel_values_video.unwrap(), video_grid_thw.unwrap())?;
|
||||
if let Some(pixel_values_video) = pixel_values_video
|
||||
&& let Some(video_grid_thw) = video_grid_thw
|
||||
{
|
||||
let video_embed = self.visual.forward(pixel_values_video, video_grid_thw)?;
|
||||
|
||||
let vision_mask = get_equal_mask(&input_ids, self.cfg.video_token_id as u32)?;
|
||||
let vision_mask = get_equal_mask(input_ids, self.cfg.video_token_id as u32)?;
|
||||
let n_video_tokens = vision_mask.sum_all()?.to_scalar::<u32>()?;
|
||||
if n_video_tokens as usize != video_embed.dim(0)? {
|
||||
return Err(anyhow!(format!(
|
||||
@@ -1162,8 +1155,8 @@ impl Qwen2_5VLModel {
|
||||
}
|
||||
inputs_embeds = masked_scatter_dim0(&inputs_embeds, &video_embed, &vision_mask)?;
|
||||
}
|
||||
let mut position_ids;
|
||||
let mut rope_deltas;
|
||||
let position_ids;
|
||||
let rope_deltas;
|
||||
if (cache_position.is_some() && cache_position.unwrap().i(0)?.to_scalar::<u32>()? == 0)
|
||||
|| self.rope_deltas.is_none()
|
||||
{
|
||||
@@ -1177,12 +1170,11 @@ impl Qwen2_5VLModel {
|
||||
self.rope_deltas = Some(rope_deltas);
|
||||
} else {
|
||||
let (bs, seq_len, _) = inputs_embeds.dims3()?;
|
||||
let delta = if cache_position.is_some() {
|
||||
let delta = if let Some(cache_position) = cache_position {
|
||||
cache_position
|
||||
.unwrap()
|
||||
.i(0)?
|
||||
.to_dtype(self.rope_deltas.as_ref().unwrap().dtype())?
|
||||
.broadcast_add(&self.rope_deltas.as_ref().unwrap())?
|
||||
.broadcast_add(self.rope_deltas.as_ref().unwrap())?
|
||||
.contiguous()?
|
||||
.to_dtype(candle_core::DType::U32)?
|
||||
} else {
|
||||
|
||||
@@ -13,7 +13,7 @@ use crate::{
|
||||
models::qwen2_5vl::config::VisionSetting,
|
||||
utils::{
|
||||
img_utils::get_image,
|
||||
utils::{ceil_by_factor, floor_by_factor, round_by_factor},
|
||||
{ceil_by_factor, floor_by_factor, round_by_factor},
|
||||
},
|
||||
};
|
||||
|
||||
@@ -63,22 +63,15 @@ impl Qwen2_5VLProcessor {
|
||||
vision_map.insert("image".to_string(), Vec::new());
|
||||
vision_map.insert("video".to_string(), Vec::new());
|
||||
for chat_mes in mes.messages.clone() {
|
||||
match chat_mes {
|
||||
ChatMessage::User { content, name } => match content {
|
||||
ChatMessageContent::ContentPart(part_vec) => {
|
||||
for part in part_vec {
|
||||
match part {
|
||||
ChatMessageContentPart::Image(img_part) => {
|
||||
let img_url = img_part.image_url;
|
||||
vision_map.get_mut("image").unwrap().push(img_url.url);
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
if let ChatMessage::User { content, .. } = chat_mes
|
||||
&& let ChatMessageContent::ContentPart(part_vec) = content
|
||||
{
|
||||
for part in part_vec {
|
||||
if let ChatMessageContentPart::Image(img_part) = part {
|
||||
let img_url = img_part.image_url;
|
||||
vision_map.get_mut("image").unwrap().push(img_url.url);
|
||||
}
|
||||
_ => {}
|
||||
},
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(vision_map)
|
||||
@@ -107,9 +100,7 @@ impl Qwen2_5VLProcessor {
|
||||
// 0-255 rescale to 0-1
|
||||
let img_tensor = img_tensor.affine(1.0 / 255.0, 0.)?;
|
||||
// normalize
|
||||
let img_tensor = img_tensor
|
||||
.broadcast_sub(&img_mean)?
|
||||
.broadcast_div(&img_std)?;
|
||||
let img_tensor = img_tensor.broadcast_sub(img_mean)?.broadcast_div(img_std)?;
|
||||
// (c, h, w) => (1, c, h, w)
|
||||
let img_tensor = img_tensor.unsqueeze(0)?;
|
||||
Ok(img_tensor)
|
||||
@@ -169,7 +160,7 @@ impl Qwen2_5VLProcessor {
|
||||
let mut vision_grid_thws_vec = Vec::new();
|
||||
|
||||
for img in imgs {
|
||||
let img_tensor = self.process_img(&img, &img_mean, &img_std)?;
|
||||
let img_tensor = self.process_img(&img, img_mean, img_std)?;
|
||||
let img_tensor = Tensor::cat(&[&img_tensor, &img_tensor], 0)?.contiguous()?;
|
||||
let (img_tensor, grid_thw) = self.process_vision_tensor(&img_tensor)?;
|
||||
pixel_values_vec.push(img_tensor);
|
||||
@@ -196,8 +187,8 @@ impl Qwen2_5VLProcessor {
|
||||
let video_tensor = single_video.to_dtype(self.dtype)?.affine(1.0 / 255.0, 0.)?;
|
||||
// normalize
|
||||
let video_tensor = video_tensor
|
||||
.broadcast_sub(&img_mean)?
|
||||
.broadcast_div(&img_std)?
|
||||
.broadcast_sub(img_mean)?
|
||||
.broadcast_div(img_std)?
|
||||
.contiguous()?;
|
||||
let (video_tensor, video_grid_thw) = self.process_vision_tensor(&video_tensor)?;
|
||||
pixel_values_vec.push(video_tensor);
|
||||
@@ -238,7 +229,7 @@ impl Qwen2_5VLProcessor {
|
||||
Err(e) => println!("get_image err: {:?}", e),
|
||||
};
|
||||
}
|
||||
if file_vec.len() > 0 {
|
||||
if !file_vec.is_empty() {
|
||||
let vision_input = self.process_images(file_vec, &img_mean, &img_std);
|
||||
match vision_input {
|
||||
Ok(img_input) => {
|
||||
@@ -258,7 +249,7 @@ impl Qwen2_5VLProcessor {
|
||||
Err(e) => println!("get_video_data err: {:?}", e),
|
||||
};
|
||||
}
|
||||
if file_vec.len() > 0 {
|
||||
if !file_vec.is_empty() {
|
||||
let vision_input = self.process_videos(file_vec, &img_mean, &img_std);
|
||||
match vision_input {
|
||||
Ok(video_input) => {
|
||||
@@ -280,10 +271,10 @@ impl Qwen2_5VLProcessor {
|
||||
}
|
||||
let merge_length = self.vision_setting.merge_size.pow(2);
|
||||
let mut text = text.to_string();
|
||||
if image_grid_thw.is_some() {
|
||||
if let Some(ref image_grid_thw) = image_grid_thw {
|
||||
let mut index = 0;
|
||||
while text.contains(&self.image_token) {
|
||||
let grid_i = image_grid_thw.as_ref().unwrap().i(index)?;
|
||||
let grid_i = image_grid_thw.i(index)?;
|
||||
let repeat_num =
|
||||
grid_i.to_vec1::<u32>()?.iter().product::<u32>() as usize / merge_length;
|
||||
let replace = "<|placeholder|>".repeat(repeat_num);
|
||||
@@ -292,10 +283,10 @@ impl Qwen2_5VLProcessor {
|
||||
}
|
||||
text = text.replace("<|placeholder|>", &self.image_token);
|
||||
}
|
||||
if video_grid_thw.is_some() {
|
||||
if let Some(ref video_grid_thw) = video_grid_thw {
|
||||
let mut index = 0;
|
||||
while text.contains(&self.video_token) {
|
||||
let grid_i = video_grid_thw.as_ref().unwrap().i(index)?;
|
||||
let grid_i = video_grid_thw.i(index)?;
|
||||
let repeat_num =
|
||||
grid_i.to_vec1::<u32>()?.iter().product::<u32>() as usize / merge_length;
|
||||
let replace = "<|placeholder|>".repeat(repeat_num);
|
||||
@@ -336,15 +327,15 @@ pub fn smart_resize(
|
||||
}
|
||||
let mut h_bar = std::cmp::max(image_factor, round_by_factor(img_h, image_factor));
|
||||
let mut w_bar = std::cmp::max(image_factor, round_by_factor(img_w, image_factor));
|
||||
let mut max_pixels = 0u32;
|
||||
let mut min_pixels = 0u32;
|
||||
if is_img {
|
||||
min_pixels = vision_setting.min_pixels;
|
||||
max_pixels = vision_setting.max_pixels;
|
||||
|
||||
let (min_pixels, max_pixels) = if is_img {
|
||||
(vision_setting.min_pixels, vision_setting.max_pixels)
|
||||
} else {
|
||||
min_pixels = vision_setting.video_min_pixels;
|
||||
max_pixels = vision_setting.video_max_pixels;
|
||||
}
|
||||
(
|
||||
vision_setting.video_min_pixels,
|
||||
vision_setting.video_max_pixels,
|
||||
)
|
||||
};
|
||||
if h_bar * w_bar > max_pixels {
|
||||
let beta = ((img_h * img_w) as f32 / max_pixels as f32).sqrt();
|
||||
h_bar = floor_by_factor(img_h as f32 / beta, image_factor);
|
||||
@@ -419,7 +410,7 @@ pub fn get_video_data(
|
||||
|decoder: &mut ffmpeg::decoder::Video| -> Result<()> {
|
||||
let mut decoded = ffmpeg::frame::Video::empty();
|
||||
while decoder.receive_frame(&mut decoded).is_ok() {
|
||||
if frame_id % sample_interval == 0 {
|
||||
if frame_id.is_multiple_of(sample_interval) {
|
||||
let mut rgb_frame = ffmpeg::frame::Video::empty();
|
||||
scaler
|
||||
.run(&decoded, &mut rgb_frame)
|
||||
|
||||
@@ -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)?;
|
||||
|
||||
@@ -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,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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> = {
|
||||
|
||||
@@ -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;
|
||||
|
||||
+14
-10
@@ -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 {
|
||||
|
||||
@@ -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);
|
||||
|
||||
Reference in New Issue
Block a user