jhqxxx
584 lines
19 KiB
Rust
584 lines
19 KiB
Rust
use anyhow::{Result, anyhow};
|
|
use candle_core::{D, IndexOp, Tensor};
|
|
use candle_nn::{
|
|
Conv2d, Embedding, Init, Linear, Module, RmsNorm, VarBuilder, embedding, linear, linear_b,
|
|
rms_norm,
|
|
};
|
|
|
|
use crate::{
|
|
models::{
|
|
common::{GateUpDownMLP, NaiveAttnTwoLinearMLPBlock, eager_attention_forward, get_conv2d},
|
|
hunyuan_ocr::config::{HunYuanVLConfig, HunYuanVLVisionConfig},
|
|
},
|
|
position_embed::rope::{RoPE, apply_rotary_pos_emb, get_xd_cos_sin},
|
|
utils::tensor_utils::{
|
|
interpolate_bilinear, masked_scatter_dim0, prepare_causal_attention_mask, split_tensor,
|
|
},
|
|
};
|
|
|
|
pub struct HunYuanVisionPatchEmbed {
|
|
patch_embedding: Conv2d,
|
|
// position_embedding: Embedding,
|
|
num_channels: usize,
|
|
patch_size: usize,
|
|
// num_positions: usize,
|
|
// position_edge: usize,
|
|
embed_dim: usize,
|
|
patch_pos_embed: Tensor,
|
|
}
|
|
|
|
impl HunYuanVisionPatchEmbed {
|
|
pub fn new(vb: VarBuilder, config: &HunYuanVLVisionConfig) -> Result<Self> {
|
|
let patch_embedding = get_conv2d(
|
|
vb.pp("patch_embedding"),
|
|
config.num_channels,
|
|
config.hidden_size,
|
|
config.patch_size,
|
|
0,
|
|
config.patch_size,
|
|
1,
|
|
1,
|
|
true,
|
|
)?;
|
|
let num_channels = config.num_channels;
|
|
let patch_size = config.patch_size;
|
|
let position_edge = config.max_image_size / patch_size;
|
|
let num_positions = (position_edge).pow(2) + 1;
|
|
let embed_dim = config.hidden_size;
|
|
let position_embedding = embedding(num_positions, embed_dim, vb.pp("position_embedding"))?;
|
|
let patch_pos_embed = position_embedding
|
|
.embeddings()
|
|
.i(1..)?
|
|
.reshape((1, position_edge, position_edge, embed_dim))?
|
|
.permute((0, 3, 1, 2))?;
|
|
Ok(Self {
|
|
patch_embedding,
|
|
// position_embedding,
|
|
num_channels,
|
|
patch_size,
|
|
// num_positions,
|
|
// position_edge,
|
|
embed_dim,
|
|
patch_pos_embed,
|
|
})
|
|
}
|
|
|
|
pub fn forward(&self, pixel_values: &Tensor, grid_thw: &Tensor) -> Result<Tensor> {
|
|
let (num_patches, _) = pixel_values.dims2()?;
|
|
let pixel_values = pixel_values.reshape((
|
|
num_patches,
|
|
self.num_channels,
|
|
self.patch_size,
|
|
self.patch_size,
|
|
))?;
|
|
let patch_embeds = self.patch_embedding.forward(&pixel_values)?;
|
|
let patch_embeds = patch_embeds
|
|
.squeeze(D::Minus1)?
|
|
.squeeze(D::Minus1)?
|
|
.unsqueeze(0)?;
|
|
let mut patch_pos_embed_list = vec![];
|
|
let img_num = grid_thw.dim(0)?;
|
|
for i in 0..img_num {
|
|
let grid_i = grid_thw.i(i)?;
|
|
let grid_h = grid_i.i(1)?.to_scalar::<u32>()? as usize;
|
|
let grid_w = grid_i.i(2)?.to_scalar::<u32>()? as usize;
|
|
let patch_pos_embed_ =
|
|
interpolate_bilinear(&self.patch_pos_embed, (grid_h, grid_w), Some(false))?;
|
|
let patch_pos_embed_ = patch_pos_embed_
|
|
.reshape((self.embed_dim, ()))?
|
|
.transpose(0, 1)?
|
|
.unsqueeze(0)?;
|
|
patch_pos_embed_list.push(patch_pos_embed_);
|
|
}
|
|
let patch_pos_embed = Tensor::cat(&patch_pos_embed_list, 1)?;
|
|
let embedding = patch_embeds.add(&patch_pos_embed)?;
|
|
Ok(embedding)
|
|
}
|
|
}
|
|
|
|
pub struct HunYuanVisionPatchMerger {
|
|
proj_0: Conv2d,
|
|
proj_2: Conv2d,
|
|
mlp: Linear,
|
|
image_newline: Tensor,
|
|
image_begin: Tensor,
|
|
image_end: Tensor,
|
|
// image_sep: Tensor,
|
|
before_rms: RmsNorm,
|
|
after_rms: RmsNorm,
|
|
}
|
|
|
|
impl HunYuanVisionPatchMerger {
|
|
pub fn new(vb: VarBuilder, config: &HunYuanVLVisionConfig) -> Result<Self> {
|
|
let proj_0 = get_conv2d(
|
|
vb.pp("proj.0"),
|
|
config.hidden_size,
|
|
config.hidden_size * 2,
|
|
config.spatial_merge_size,
|
|
0,
|
|
config.spatial_merge_size,
|
|
1,
|
|
1,
|
|
true,
|
|
)?;
|
|
let proj_2 = get_conv2d(
|
|
vb.pp("proj.2"),
|
|
config.hidden_size * 2,
|
|
config.hidden_size * 4,
|
|
1,
|
|
0,
|
|
1,
|
|
1,
|
|
1,
|
|
true,
|
|
)?;
|
|
let mlp = linear(config.hidden_size * 4, config.out_hidden_size, vb.pp("mlp"))?;
|
|
let image_newline =
|
|
vb.get_with_hints(config.hidden_size * 4, "image_newline", Init::Const(0.))?;
|
|
let image_begin =
|
|
vb.get_with_hints(config.out_hidden_size, "image_begin", Init::Const(0.))?;
|
|
let image_end = vb.get_with_hints(config.out_hidden_size, "image_end", Init::Const(0.))?;
|
|
// let image_sep = vb.get_with_hints(config.out_hidden_size, "image_sep", Init::Const(0.))?;
|
|
let before_rms = rms_norm(config.hidden_size, config.rms_norm_eps, vb.pp("before_rms"))?;
|
|
let after_rms = rms_norm(
|
|
config.out_hidden_size,
|
|
config.rms_norm_eps,
|
|
vb.pp("after_rms"),
|
|
)?;
|
|
Ok(Self {
|
|
proj_0,
|
|
proj_2,
|
|
mlp,
|
|
image_newline,
|
|
image_begin,
|
|
image_end,
|
|
// image_sep,
|
|
before_rms,
|
|
after_rms,
|
|
})
|
|
}
|
|
pub fn forward(&self, xs: &Tensor, size: (usize, usize)) -> Result<Tensor> {
|
|
let xs = self.before_rms.forward(xs)?;
|
|
let (h, w) = size;
|
|
let xs = xs.permute((0, 2, 1))?.reshape((xs.dim(0)?, (), h, w))?;
|
|
let xs = self.proj_0.forward(&xs)?.gelu()?;
|
|
let xs = self.proj_2.forward(&xs)?;
|
|
let (b, c, h, _) = xs.dims4()?;
|
|
let image_newline = self
|
|
.image_newline
|
|
.reshape((1, c, 1, 1))?
|
|
.broadcast_as((b, c, h, 1))?
|
|
.to_dtype(xs.dtype())?;
|
|
let xs = Tensor::cat(&[xs, image_newline], D::Minus1)?;
|
|
let xs = xs.reshape((b, c, ()))?.permute((0, 2, 1))?;
|
|
let xs = self.mlp.forward(&xs)?;
|
|
let begin = self
|
|
.image_begin
|
|
.reshape((1, 1, ()))?
|
|
.broadcast_as((b, 1, xs.dim(D::Minus1)?))?
|
|
.to_dtype(xs.dtype())?;
|
|
let end = self
|
|
.image_end
|
|
.reshape((1, 1, ()))?
|
|
.broadcast_as((b, 1, xs.dim(D::Minus1)?))?
|
|
.to_dtype(xs.dtype())?;
|
|
let xs = Tensor::cat(&[begin, xs, end], 1)?;
|
|
let xs = self.after_rms.forward(&xs)?;
|
|
Ok(xs)
|
|
}
|
|
}
|
|
|
|
pub struct HunYuanVisionTransformer {
|
|
embeddings: HunYuanVisionPatchEmbed,
|
|
layers: Vec<NaiveAttnTwoLinearMLPBlock>,
|
|
perceive: HunYuanVisionPatchMerger,
|
|
}
|
|
|
|
impl HunYuanVisionTransformer {
|
|
pub fn new(vb: VarBuilder, config: &HunYuanVLVisionConfig) -> Result<Self> {
|
|
let embeddings = HunYuanVisionPatchEmbed::new(vb.pp("embeddings"), config)?;
|
|
let mut layers = vec![];
|
|
let vb_layers = vb.pp("layers");
|
|
for i in 0..config.num_hidden_layers {
|
|
let layer_i = NaiveAttnTwoLinearMLPBlock::new(
|
|
vb_layers.pp(i),
|
|
config.hidden_size,
|
|
config.num_attention_heads,
|
|
None,
|
|
None,
|
|
true,
|
|
"self_attn",
|
|
None,
|
|
config.intermediate_size,
|
|
config.hidden_act,
|
|
true,
|
|
"mlp",
|
|
"dense_h_to_4h",
|
|
"dense_4h_to_h",
|
|
config.rms_norm_eps,
|
|
"input_layernorm",
|
|
"post_attention_layernorm",
|
|
)?;
|
|
layers.push(layer_i);
|
|
}
|
|
let perceive = HunYuanVisionPatchMerger::new(vb.pp("perceive"), config)?;
|
|
Ok(Self {
|
|
embeddings,
|
|
layers,
|
|
perceive,
|
|
})
|
|
}
|
|
|
|
pub fn forward(&self, xs: &Tensor, grid_thw: &Tensor) -> Result<Tensor> {
|
|
let mut hidden_states = self.embeddings.forward(xs, grid_thw)?;
|
|
for layer in &self.layers {
|
|
hidden_states = layer.forward(&hidden_states, None, None, None, false)?;
|
|
}
|
|
let mut cu_seqlens = vec![];
|
|
for i in 0..grid_thw.dim(0)? {
|
|
let [_, h, w] = grid_thw.i(i)?.to_vec1::<u32>()?[..] else {
|
|
return Err(anyhow!(format!("grid_thw Expected exactly 3 elements")));
|
|
};
|
|
cu_seqlens.push((h * w) as usize);
|
|
}
|
|
let split_items = split_tensor(&hidden_states, &cu_seqlens, 1)?;
|
|
let mut processed_item = vec![];
|
|
for i in 0..grid_thw.dim(0)? {
|
|
let [_, h, w] = grid_thw.i(i)?.to_vec1::<u32>()?[..] else {
|
|
return Err(anyhow!(format!("grid_thw Expected exactly 3 elements")));
|
|
};
|
|
let processed = self
|
|
.perceive
|
|
.forward(&split_items[i], (h as usize, w as usize))?;
|
|
processed_item.push(processed);
|
|
}
|
|
let xs = Tensor::cat(&processed_item, 1)?;
|
|
Ok(xs)
|
|
}
|
|
}
|
|
|
|
pub struct HunYuanVLAttention {
|
|
q_proj: Linear,
|
|
k_proj: Linear,
|
|
v_proj: Linear,
|
|
o_proj: Linear,
|
|
query_layernorm: RmsNorm,
|
|
key_layernorm: RmsNorm,
|
|
num_attention_heads: usize,
|
|
num_key_value_heads: usize,
|
|
num_kv_groups: usize,
|
|
head_dim: usize,
|
|
scaling: f64,
|
|
kv_cache: Option<(Tensor, Tensor)>,
|
|
}
|
|
|
|
impl HunYuanVLAttention {
|
|
pub fn new(
|
|
vb: VarBuilder,
|
|
hidden_size: usize,
|
|
head_dim: usize,
|
|
num_attention_heads: usize,
|
|
num_key_value_heads: usize,
|
|
attention_bias: bool,
|
|
rms_norm_eps: f64,
|
|
) -> Result<Self> {
|
|
let num_kv_groups = num_attention_heads / num_key_value_heads;
|
|
let scaling = 1f64 / f64::sqrt(head_dim as f64);
|
|
let q_proj = linear_b(
|
|
hidden_size,
|
|
num_attention_heads * head_dim,
|
|
attention_bias,
|
|
vb.pp("q_proj"),
|
|
)?;
|
|
let k_proj = linear_b(
|
|
hidden_size,
|
|
num_key_value_heads * head_dim,
|
|
attention_bias,
|
|
vb.pp("k_proj"),
|
|
)?;
|
|
let v_proj = linear_b(
|
|
hidden_size,
|
|
num_key_value_heads * head_dim,
|
|
attention_bias,
|
|
vb.pp("v_proj"),
|
|
)?;
|
|
let o_proj = linear_b(
|
|
num_attention_heads * head_dim,
|
|
hidden_size,
|
|
attention_bias,
|
|
vb.pp("o_proj"),
|
|
)?;
|
|
|
|
let query_layernorm = rms_norm(head_dim, rms_norm_eps, vb.pp("query_layernorm"))?;
|
|
let key_layernorm = rms_norm(head_dim, rms_norm_eps, vb.pp("key_layernorm"))?;
|
|
Ok(Self {
|
|
q_proj,
|
|
k_proj,
|
|
v_proj,
|
|
o_proj,
|
|
query_layernorm,
|
|
key_layernorm,
|
|
num_attention_heads,
|
|
num_key_value_heads,
|
|
num_kv_groups,
|
|
head_dim,
|
|
scaling,
|
|
kv_cache: None,
|
|
})
|
|
}
|
|
|
|
pub fn forward(
|
|
&mut self,
|
|
xs: &Tensor,
|
|
cos: &Tensor,
|
|
sin: &Tensor,
|
|
attention_mask: Option<&Tensor>,
|
|
) -> Result<Tensor> {
|
|
let (b_sz, q_len, _) = xs.dims3()?;
|
|
let query_states = self
|
|
.q_proj
|
|
.forward(xs)?
|
|
.reshape((b_sz, q_len, self.num_attention_heads, self.head_dim))?
|
|
.transpose(1, 2)?;
|
|
|
|
let key_states = self
|
|
.k_proj
|
|
.forward(xs)?
|
|
.reshape((b_sz, q_len, self.num_key_value_heads, self.head_dim))?
|
|
.transpose(1, 2)?;
|
|
let value_states = self.v_proj.forward(xs)?;
|
|
let value_states = value_states
|
|
.reshape((b_sz, q_len, self.num_key_value_heads, self.head_dim))?
|
|
.transpose(1, 2)?;
|
|
let (query_states, key_states) =
|
|
apply_rotary_pos_emb(&query_states, &key_states, cos, sin, false)?;
|
|
let query_states = self.query_layernorm.forward(&query_states)?;
|
|
let key_states = self.key_layernorm.forward(&key_states)?;
|
|
let (key_states, value_states) = match &self.kv_cache {
|
|
None => (key_states, value_states),
|
|
Some((prev_k, prev_v)) => {
|
|
let key_states = Tensor::cat(&[prev_k, &key_states], 2)?;
|
|
let value_states = Tensor::cat(&[prev_v, &value_states], 2)?;
|
|
(key_states, value_states)
|
|
}
|
|
};
|
|
self.kv_cache = Some((key_states.clone(), value_states.clone()));
|
|
let attn_output = eager_attention_forward(
|
|
&query_states,
|
|
&key_states,
|
|
&value_states,
|
|
Some(self.num_kv_groups),
|
|
attention_mask,
|
|
self.scaling,
|
|
)?;
|
|
let attn_output =
|
|
attn_output.reshape((b_sz, q_len, self.num_attention_heads * self.head_dim))?;
|
|
let attn_output = attn_output.apply(&self.o_proj)?;
|
|
Ok(attn_output)
|
|
}
|
|
|
|
pub fn clear_kv_cache(&mut self) {
|
|
self.kv_cache = None
|
|
}
|
|
}
|
|
|
|
pub struct HunYuanVLDecoderLayer {
|
|
self_attn: HunYuanVLAttention,
|
|
mlp: GateUpDownMLP,
|
|
input_layernorm: RmsNorm,
|
|
post_attention_layernorm: RmsNorm,
|
|
}
|
|
|
|
impl HunYuanVLDecoderLayer {
|
|
pub fn new(config: &HunYuanVLConfig, vb: VarBuilder) -> Result<Self> {
|
|
let self_attn = HunYuanVLAttention::new(
|
|
vb.pp("self_attn"),
|
|
config.hidden_size,
|
|
config.head_dim,
|
|
config.num_attention_heads,
|
|
config.num_key_value_heads,
|
|
config.attention_bias,
|
|
config.rms_norm_eps,
|
|
)?;
|
|
let mlp = GateUpDownMLP::new(
|
|
vb.pp("mlp"),
|
|
config.hidden_size,
|
|
config.intermediate_size,
|
|
config.hidden_act,
|
|
false,
|
|
None,
|
|
None,
|
|
None,
|
|
)?;
|
|
let input_layernorm = rms_norm(
|
|
config.hidden_size,
|
|
config.rms_norm_eps,
|
|
vb.pp("input_layernorm"),
|
|
)?;
|
|
let post_attention_layernorm = rms_norm(
|
|
config.hidden_size,
|
|
config.rms_norm_eps,
|
|
vb.pp("post_attention_layernorm"),
|
|
)?;
|
|
Ok(Self {
|
|
self_attn,
|
|
mlp,
|
|
input_layernorm,
|
|
post_attention_layernorm,
|
|
})
|
|
}
|
|
|
|
pub fn forward(
|
|
&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(&xs, cos, sin, attention_mask)?;
|
|
let xs = residual.add(&xs)?;
|
|
let residual = xs.clone();
|
|
let xs = self.post_attention_layernorm.forward(&xs)?;
|
|
let xs = self.mlp.forward(&xs)?;
|
|
let xs = residual.add(&xs)?;
|
|
Ok(xs)
|
|
}
|
|
|
|
pub fn clear_kv_cache(&mut self) {
|
|
self.self_attn.clear_kv_cache();
|
|
}
|
|
}
|
|
|
|
pub struct HunYuanVLTextModel {
|
|
embed_tokens: Embedding,
|
|
layers: Vec<HunYuanVLDecoderLayer>,
|
|
norm: RmsNorm,
|
|
rope: RoPE,
|
|
xdrope_section: Vec<usize>,
|
|
}
|
|
|
|
impl HunYuanVLTextModel {
|
|
pub fn new(vb: VarBuilder, config: &HunYuanVLConfig) -> Result<Self> {
|
|
let embed_tokens = embedding(config.vocab_size, config.hidden_size, vb.pp("embed_tokens"))?;
|
|
let mut layers = vec![];
|
|
let vb_layers = vb.pp("layers");
|
|
for i in 0..config.num_hidden_layers {
|
|
let layer = HunYuanVLDecoderLayer::new(config, vb_layers.pp(i))?;
|
|
layers.push(layer);
|
|
}
|
|
let norm = rms_norm(config.hidden_size, config.rms_norm_eps, vb.pp("norm"))?;
|
|
let base = config.rope_theta
|
|
* config
|
|
.rope_scaling
|
|
.alpha
|
|
.powf(config.head_dim as f64 / (config.head_dim - 2) as f64);
|
|
let rope = RoPE::new(config.head_dim, base as f32, vb.device())?;
|
|
let xdrope_section = config.rope_scaling.xdrope_section.clone();
|
|
Ok(Self {
|
|
embed_tokens,
|
|
layers,
|
|
norm,
|
|
rope,
|
|
xdrope_section,
|
|
})
|
|
}
|
|
|
|
pub fn forward(
|
|
&mut self,
|
|
inputs_embeds: &Tensor,
|
|
position_ids: Option<&Tensor>,
|
|
seqlen_offset: usize,
|
|
) -> Result<Tensor> {
|
|
let (b_size, seq_len, _) = inputs_embeds.dims3()?;
|
|
|
|
let attention_mask: Option<Tensor> = {
|
|
if seq_len <= 1 {
|
|
None
|
|
} else {
|
|
Some(prepare_causal_attention_mask(
|
|
b_size,
|
|
seq_len,
|
|
0,
|
|
inputs_embeds.device(),
|
|
)?)
|
|
}
|
|
};
|
|
|
|
let (cos, sin) = self
|
|
.rope
|
|
.forward(seqlen_offset, seq_len, inputs_embeds.device())?;
|
|
let mut xs = inputs_embeds.clone();
|
|
for (i, layer) in self.layers.iter_mut().enumerate() {
|
|
if i == 0
|
|
&& let Some(position_ids) = position_ids
|
|
{
|
|
let (cos, sin) =
|
|
get_xd_cos_sin(&cos, &sin, position_ids, self.xdrope_section.clone())?;
|
|
xs = layer.forward(&xs, &cos, &sin, attention_mask.as_ref())?;
|
|
} else {
|
|
xs = layer.forward(&xs, &cos, &sin, attention_mask.as_ref())?;
|
|
}
|
|
}
|
|
let xs = self.norm.forward(&xs)?;
|
|
Ok(xs)
|
|
}
|
|
|
|
pub fn clear_kv_cache(&mut self) {
|
|
for layer in self.layers.iter_mut() {
|
|
layer.clear_kv_cache()
|
|
}
|
|
}
|
|
}
|
|
|
|
pub struct HunyuanVLModel {
|
|
// config: HunYuanVLConfig,
|
|
vit: HunYuanVisionTransformer,
|
|
model: HunYuanVLTextModel,
|
|
lm_head: Linear,
|
|
}
|
|
|
|
impl HunyuanVLModel {
|
|
pub fn new(vb: VarBuilder, config: HunYuanVLConfig) -> Result<Self> {
|
|
let vit = HunYuanVisionTransformer::new(vb.pp("vit"), &config.vision_config)?;
|
|
let model = HunYuanVLTextModel::new(vb.pp("model"), &config)?;
|
|
let lm_head = Linear::new(model.embed_tokens.embeddings().clone(), None);
|
|
Ok(Self {
|
|
// config,
|
|
vit,
|
|
model,
|
|
lm_head,
|
|
})
|
|
}
|
|
pub fn forward(
|
|
&mut self,
|
|
input_ids: &Tensor,
|
|
pixel_values: Option<&Tensor>,
|
|
image_grid_thw: Option<&Tensor>,
|
|
image_mask: Option<&Tensor>,
|
|
position_ids: Option<&Tensor>,
|
|
seqlen_offset: usize,
|
|
) -> Result<Tensor> {
|
|
let mut inputs_embeds = self.model.embed_tokens.forward(input_ids)?;
|
|
if let Some(pixel_values) = pixel_values
|
|
&& let Some(grid_thw) = image_grid_thw
|
|
&& let Some(image_mask) = image_mask
|
|
{
|
|
let image_embeds = self.vit.forward(pixel_values, grid_thw)?.squeeze(0)?;
|
|
inputs_embeds = masked_scatter_dim0(&inputs_embeds, &image_embeds, image_mask)?;
|
|
}
|
|
let outputs = self
|
|
.model
|
|
.forward(&inputs_embeds, position_ids, seqlen_offset)?;
|
|
let seq_len = outputs.dim(1)?;
|
|
let hidden_state = outputs.narrow(1, seq_len - 1, 1)?;
|
|
let logits = self.lm_head.forward(&hidden_state)?;
|
|
Ok(logits)
|
|
}
|
|
|
|
pub fn clear_kv_cache(&mut self) {
|
|
self.model.clear_kv_cache();
|
|
}
|
|
}
|