Skip to content

Commit 7c19bb0

Browse files
committed
Thread image_grid_thw through Qwen3-VL, add smart-resize + multi-image
Per-image grid dimensions are now a real model input rather than being inferred from sqrt(num_patches). This fixes correctness for non-square images and lets one forward pass handle multiple images of different sizes in the same prompt (concatenated patch sequence + per-image grid_thw). Featurizer changes: - smart_resize mirrors qwen-vl-utils: preserves aspect ratio, rounds to a multiple of patch_size * merge_size, clamps total pixels between min_pixels and max_pixels - :quality preset (:low / :medium / :high) over explicit pixel caps - accepts a list of images of different sizes; concatenates patches, returns image_grid_thw of shape {num_images, 3} - patches are emitted in windowed order so the merger reshape stays shape-agnostic (4 consecutive patches = one 2x2 merge block) Vision encoder: - image_grid_thw declared as a model input - per-patch (row, col, grid_h, grid_w, image_id) derived from grid_thw via Nx ops; drives bilinear pos-embed interpolation and 2D rotary - block-diagonal attention mask (image_id == image_id) so patches from one image cannot attend to patches from another - patch merger and deepstack merger reshapes no longer assume a square single-image grid Multimodal: - image_grid_thw plumbed as an optional model input and forwarded into the vision sub-model Validation: - 13 new tests covering smart_resize aspect ratio, quality presets, multi-image concat, windowed-layout invariant, single+multi-image end-to-end with the tiny model - Full fast suite: 285 passed, 0 regressions - Real Qwen3-VL-2B-Instruct on a 640x480 COCO image: 8/10 top-token agreement vs HuggingFace transformers reference; top-3 identical and in same order Refs #442.
1 parent 8548ee3 commit 7c19bb0

6 files changed

Lines changed: 797 additions & 441 deletions

File tree

lib/bumblebee/multimodal/qwen3_vl.ex

Lines changed: 13 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -33,9 +33,15 @@ defmodule Bumblebee.Multimodal.Qwen3VL do
3333
3434
* `"pixel_values"` - `{num_patches, flattened_patch_size}`
3535
36-
Pre-extracted image/video patches from the featurizer. The shape is
37-
`{num_patches, channels * temporal_patch_size * patch_size * patch_size}`.
38-
For a 384x384 image with default settings, this is `{576, 1536}`.
36+
Concatenated, pre-extracted image/video patches from the featurizer.
37+
Shape is `{num_patches, channels * temporal_patch_size * patch_size * patch_size}`.
38+
39+
* `"image_grid_thw"` - `{num_images, 3}`
40+
41+
Per-image grid dimensions `[temporal, height, width]` in patch
42+
units. Threaded into the vision encoder so it can compute correct
43+
per-patch positions for variable image sizes and multiple images
44+
per prompt.
3945
4046
* `"input_ids"` - `{batch_size, sequence_length}`
4147
@@ -92,6 +98,7 @@ defmodule Bumblebee.Multimodal.Qwen3VL do
9298

9399
%{
94100
"pixel_values" => Nx.template({num_patches, flattened_patch_size}, :f32),
101+
"image_grid_thw" => Nx.template({1, 3}, :s64),
95102
"input_ids" => Nx.template({1, 1}, :u32)
96103
}
97104
end
@@ -114,7 +121,8 @@ defmodule Bumblebee.Multimodal.Qwen3VL do
114121
Bumblebee.build_model(spec.vision_spec)
115122
|> Bumblebee.Utils.Axon.prefix_names("vision_model.")
116123
|> Bumblebee.Utils.Axon.plug_inputs(%{
117-
"pixel_values" => inputs["pixel_values"]
124+
"pixel_values" => inputs["pixel_values"],
125+
"image_grid_thw" => inputs["image_grid_thw"]
118126
})
119127

120128
# Get vision embeddings using correct Axon.nx pattern
@@ -194,6 +202,7 @@ defmodule Bumblebee.Multimodal.Qwen3VL do
194202

195203
Bumblebee.Utils.Model.inputs_to_map([
196204
Axon.input("pixel_values", optional: true, shape: vision_shape),
205+
Axon.input("image_grid_thw", optional: true, shape: {nil, 3}),
197206
Axon.input("input_ids", shape: text_shape),
198207
Axon.input("attention_mask", optional: true, shape: text_shape),
199208
Axon.input("position_ids", optional: true, shape: text_shape),

0 commit comments

Comments
 (0)