yiyexy commited on
Commit
5a75eaf
·
verified ·
1 Parent(s): 6885831

fix(codec): merge per-canvas image_grid_thw into [[N,H,W]] per video

Browse files

The codec branch in __call__ was emitting per-canvas image_grid_thw rows
(N, 3) = [[1, h, w]]*N, which caused the vision encoder's _build_cu_seqlens
(fixed_t=4) to treat each canvas as an isolated 1-frame sample instead of
grouping canvases into 4-frame attention windows. This silently degraded
video benchmark scores by ~1pt on videoeval_pro (45.07 vs reference 46.4).

Merge same-size canvases per video into a single [[N, H, W]] row so the
encoder restores the cross-canvas self-attention windows that the training
pipeline assumes.

Files changed (1) hide show
  1. processing_llava_onevision2.py +24 -1
processing_llava_onevision2.py CHANGED
@@ -308,7 +308,30 @@ class LlavaOnevision2Processor:
308
  all_patch_positions.append(patch_positions)
309
 
310
  out["pixel_values"] = torch.cat(all_pixel_values, dim=0)
311
- out["image_grid_thw"] = torch.cat(all_grid_thw, dim=0)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
312
  out["patch_positions"] = torch.cat(all_patch_positions, dim=0)
313
  text = rewritten_texts
314
  # Codec branch handled the video. Suppress the frame-sampling block below.
 
308
  all_patch_positions.append(patch_positions)
309
 
310
  out["pixel_values"] = torch.cat(all_pixel_values, dim=0)
311
+ # Merge per-canvas grid_thw rows into a single [N, H, W] row per
312
+ # video so the vision encoder's _build_cu_seqlens (fixed_t=4)
313
+ # groups canvases into 4-frame attention windows, matching the
314
+ # training pipeline. Without this merge the encoder treats each
315
+ # canvas as an isolated 1-frame sample (no cross-canvas
316
+ # self-attention), which silently degrades video benchmark
317
+ # scores by ~1pt on videoeval_pro.
318
+ merged_grid_thw_rows = []
319
+ for grid in all_grid_thw:
320
+ if (
321
+ grid.shape[0] > 1
322
+ and bool(torch.all(grid[:, 1] == grid[0, 1]).item())
323
+ and bool(torch.all(grid[:, 2] == grid[0, 2]).item())
324
+ ):
325
+ merged_grid_thw_rows.append(
326
+ torch.tensor(
327
+ [[int(grid.shape[0]), int(grid[0, 1]), int(grid[0, 2])]],
328
+ dtype=grid.dtype,
329
+ device=grid.device,
330
+ )
331
+ )
332
+ else:
333
+ merged_grid_thw_rows.append(grid)
334
+ out["image_grid_thw"] = torch.cat(merged_grid_thw_rows, dim=0)
335
  out["patch_positions"] = torch.cat(all_patch_positions, dim=0)
336
  text = rewritten_texts
337
  # Codec branch handled the video. Suppress the frame-sampling block below.