update cu_seqlens_repeat
This commit is contained in:
@@ -516,19 +516,11 @@ impl Qwen3VLVisionModel {
|
||||
let sin = emb.sin()?;
|
||||
let cu_seqlens = grid_thw.i((.., 1))?.mul(&grid_thw.i((.., 2))?)?;
|
||||
let grid_t = grid_thw.i((.., 0))?.to_vec1::<u32>()?;
|
||||
let cu_seqlens_full = match cu_seqlens.rank() {
|
||||
1 => cu_seqlens.repeat(grid_t[0] as usize)?,
|
||||
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 as usize)?);
|
||||
}
|
||||
Tensor::cat(&cu_seqlens_repeat, 0)?.flatten_all()?
|
||||
}
|
||||
_ => {
|
||||
return Err(anyhow!(format!("create cu_seqlens error")));
|
||||
}
|
||||
};
|
||||
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 as usize)?);
|
||||
}
|
||||
let cu_seqlens_full = Tensor::cat(&cu_seqlens_repeat, 0)?.flatten_all()?;
|
||||
let cu_seqlens = cu_seqlens_full
|
||||
.to_dtype(DType::F64)?
|
||||
.cumsum(0)?
|
||||
@@ -895,6 +887,8 @@ impl Qwen3VLModel {
|
||||
let mut v_thw_vec = Vec::new();
|
||||
for (index, t) in grid_t.iter().enumerate() {
|
||||
let mut thw_i = thw.i(index)?.to_vec1::<u32>()?;
|
||||
// [12, 30, 50]
|
||||
// [1, 30, 50]*t
|
||||
thw_i[0] = 1;
|
||||
v_thw_vec.push(
|
||||
Tensor::new(thw_i, thw.device())?
|
||||
|
||||
Reference in New Issue
Block a user