Consistent output shape from `get_image_features` by zucchini-nlp · Pull Request #46405 · huggingface/transformers

zucchini-nlp · 2026-06-04T09:19:20Z

What does this PR do?

as per title, branches off from #45783

The pooled image (OR video, NOT audio yet) feature output will now always satisfy three cond, where image features can be a list or a 3D tensor:

len(image_features) = len(input_images) # NOTE: or num_videos, not number of total video frames
image_features[0].ndim == 2
image_features[0].shape == {{actual seq length of this image, LM hidden size}}

Could make it complete BC and return the "correct" output when a certain flag is passed, but for now decided to update directly smaller utility fn in a breaking way

HuggingFaceDocBuilderDev · 2026-06-04T09:31:25Z

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

zucchini-nlp · 2026-06-04T10:00:15Z

run-slow: aya_vision, cohere2_vision, deepseek_ocr2, gemma4, paddleocr_vl, qwen2_5_omni, qwen3_omni_moe

github-actions · 2026-06-04T10:01:37Z

Workflow Run ⚙️

This comment contains run-slow, running the specified jobs:

models: ["models/aya_vision", "models/cohere2_vision", "models/deepseek_ocr2", "models/gemma4", "models/paddleocr_vl", "models/qwen2_5_omni", "models/qwen3_omni_moe"]
quantizations: []

github-actions · 2026-06-04T11:40:20Z

CI Results

Workflow Run ⚙️

Commit Info

Context	Commit	Description
RUN	5e5f7bbe	workflow commit (merge commit)
PR	52788433	branch commit (from PR)
main	b07d99be	base commit (on `main`)

✅ No failing test specific to this PR 🎉 👏 !

vasqu

Some initial comments. I think it's overall fine to break but I would want @tomaarsen's opinion as well since he worked quite a bit on this and uses it downstream so just wanna make sure we coordinate

Other than that some design questions

Do we not cat in get image features, why?
We sometimes have to calculate the expected number of tokens produced by the vision tower; would it make sense to align across processor/modeling by having some global function?

vasqu · 2026-06-17T17:27:36Z

+        image_outputs.pooler_output = pooler_output.reshape(
+            selected_image_feature.shape[0], -1, self.config.text_config.hidden_size
+        )



Does this reshape not affect the downstream forward?

nope, if used in core VLM it won't affect. I reshaped inside VLMs forward when needed

vasqu · 2026-06-17T17:28:40Z

                pixel_values, pixel_values_local, num_local_patches, return_dict=True
            ).pooler_output
-            image_features = image_features.to(inputs_embeds.device, inputs_embeds.dtype)
+            image_features = torch.cat(image_features, dim=0).to(inputs_embeds.device, inputs_embeds.dtype)


any reason the cat is not directly in get img features?

then we have a vague 2D tensor of (total_seq_length, dim) while I am moving towardsa unified output. So the output will show clearly how many images there are and make it easy to get an embeddings of image_i

Ah true, yea that makes sense

vasqu · 2026-06-17T17:29:56Z

+        k_squared = int((image_position_ids.shape[1] // output_length) ** 0.5) ** 2
+        non_pad_mask = (image_position_ids != -1).all(dim=-1)
+        split_sizes = (non_pad_mask.sum(dim=-1) // k_squared).tolist()
+        vision_outputs.pooler_output = torch.split(pooler_output, split_sizes)


Ig this is similar to how the processor calculates the number of images/videos? Would be nice if we could have a cross ref

nah, it is taken from Gemma4Pooling layer but yeah, will add a ref

vasqu · 2026-06-17T17:30:42Z

        image_embeds = self.projector(image_embeds, image_grid_thw)
-        vision_outputs.pooler_output = image_embeds
+        split_sizes = (image_grid_thw.prod(-1) // self.config.vision_config.spatial_merge_size**2).tolist()
+        vision_outputs.pooler_output = torch.split(image_embeds, split_sizes)


Hmm ok this will happen a few times; adjusting my previous comment a bit. Would it make sense to have a global function for this that is used across processor and model as here?

haha, this happens for all qwen-based models which is indeed a lot atm. Wdym by a global fn since I'm not sure it's worth for a one-liner torch.split

I thought maybe a static function that can be both consumed by the model / processor to get the expected number here. But yea indeed it's not too worthwhile for now

vasqu · 2026-06-17T17:34:56Z

We also need to add 🚨 for sure

zucchini-nlp · 2026-06-18T07:42:09Z

Do we not cat in get image features, why?
Answered in comments, to get clearly split encodings per image so we don't have to guess. It'll make to easy to get embeddings of image-i and expand/crop/do whatever you need

We sometimes have to calculate the expected number of tokens produced by the vision tower; would it make sense to align across processor/modeling by having some global function?

You mean gemma4 (which I also replied is based on modeling, not processor)? Otherwise we always pass the number of grids (as in qwen) or just split based on input pixel shape

zucchini-nlp · 2026-06-18T07:43:07Z

Oh yeah, Tom requested this in the past for last_hidden_states which I personally didn't need for generation 😅 Leaving it open, and working only on pooler_output for now

vasqu · 2026-06-18T07:48:34Z

Yea I think we are mostly fine then, still want to wait on @tomaarsen if he has any opinions on this and maybe @remi-or using it for CB(?)

zucchini-nlp · 2026-06-18T07:52:03Z

yeah, let's wait for ppl affected. I also commented on CB to let Remi know :)

remi-or · 2026-06-18T11:43:55Z

In general, anything that makes shapes more static or standardize inputs is good for CB, so this is great!

The only problem that this could create is if torch.split is called with a number of splits that cannot be determined beforehand, which would not be CUDA graphable. Eventually, we will want CUDA graphs around the encoder.
But, since multimodal for CB has not landed yet, and the first iteration of that won't have CUDA graphs around the encoder, maybe we solve that problem when we get to it. Then, maybe an option to return the inputs before the split would do the trick.

tomaarsen · 2026-06-18T11:54:25Z

I updated my ST test suite a bit here, and then ran it in conjunction with this branch: no issues. The code also looks solid from my perspective, i.e. more standardized is more better 😄

@remi-or 's concern is fair though.

Tom Aarsen

zucchini-nlp · 2026-06-18T12:16:32Z

if torch.split is called with a number of splits that cannot be determined beforehand

indeed, the number if splits would be dynamic each call depending on input image. Even though it is part of inputs, the split will cause graphbreak. And now that I think about it, we're interfering with Ilyas' efforts to make VLMs exportable 😿

@vasqu I'll tag you after checking in with exporters and finding a better way to make everyone happy

vasqu · 2026-06-18T12:54:45Z

I think we can revisit this after #41992? I'm not sure but I don't think it would interfere with Ilyas work as we forward the inputs once and capture the necessary metadata based on that - so yea it should be fine on that side

zucchini-nlp · 2026-07-17T11:52:46Z

@vasqu , let's merge? I rebased main and fixed new models

vasqu

One small comment, but

Let's add 🚨 as this is technically still breaking (in a good way)
Get everything on board cc @tomaarsen @IlyasMoutawwakil if you could have a final check as well so we don't break exporters and ST
Fix fast CI 😬

But yes, I think it's mergable. Just want everyone to be aware that it happens

tomaarsen

I did some tests for Sentence Transformers, and found that this PR shouldn't affect it. ST only really uses the get_..._features methods when the forward does not contain a last_hidden_state (e.g. CLIP, SigLIP, etc.), as then the get_..._features methods must be used to get per-modality outputs.

For all of these modern architectures, the forward accepts any combination of modalities already, so I don't have to rely on the get_..._features methods.

I did write two small comments though, based on my and Fable's review of the work.

zucchini-nlp · 2026-07-20T09:03:08Z

Thanks for review, fixing that and hopefully merging later in the day

github-actions · 2026-07-22T08:14:27Z

[For maintainers] Suggested jobs to run (before merge)

run-slow: aya_vision, cohere2_vision, deepseek_ocr2, diffusion_gemma, gemma4, inkling, kimi_k25, paddleocr_vl, qwen2_5_omni, qwen3_omni_moe

github-actions · 2026-07-22T08:40:36Z

CI recap

Dashboard: View test results in Grafana
Latest run: 29903802249:1
Result: success | Jobs: 15 | Tests: 172,426 | Failures: 0 | Duration: 15h 17m

* upstream/main: (39 commits) Remove deprecated training args and `is_fast` property (huggingface#46917) Consistent output shape from `get_image_features` (huggingface#46405) Fix multi-device mxfp4 dequantization race in `_convert_moe_packed_tensors` (huggingface#47423) fix failed test cases for qwen3_omni_moe model (huggingface#47449) Fix Hunyuan-VL PIL image resize parity with reference preprocessing (huggingface#47233) Move `value` padding into the attention interfaces that need it (huggingface#47451) Simplify function dispatch for linear attention (huggingface#47450) [cache] Allow sliding window layers to be roll-backed for speculative decoding (huggingface#47447) Fix double-shifted training loss in GitForCausalLM (huggingface#47395) Fix CohereASR training-loss double-shift (same as Moonshine fix huggingface#46784) (huggingface#46895) Warn when `group_by_length` is silently ignored for iterable datasets (huggingface#47379) Update bug report list (huggingface#46607) Fix shape mismatch in KyutaiSpeechToText `generate()` last window (huggingface#46952) Optimize flash attention max seqlen computation in vision attention (huggingface#47170) fix: remove unreachable return in special token builder (huggingface#47420) Add Harry to slow CI (huggingface#47454) BLT: vectorize patch length processing (huggingface#47385) Fix `TrackioCallback` fails to log evaluation metrics after training ends (huggingface#46935) [Kimi] add integration tests (huggingface#47383) Fix typo in `MusicgenForCausalLM.generate()` (huggingface#46974) ...

done

6130e27

fix repo

5278843

and videos

80f4bba

zucchini-nlp requested a review from vasqu June 4, 2026 12:42

Merge branch 'main' into image-output-shapes

c5fb2dc

vasqu reviewed Jun 17, 2026

View reviewed changes

zucchini-nlp mentioned this pull request Jun 18, 2026

[CB] [Major] Add multimodal support to continous batching #46376

Open

zucchini-nlp added 3 commits July 17, 2026 12:06

Merge remote-tracking branch 'upstream/main' into image-output-shapes

3454874

new models

a8fc78f

kimi

69ac8e8

vasqu approved these changes Jul 17, 2026

View reviewed changes

Comment thread tests/models/gemma4/test_modeling_gemma4.py

tomaarsen reviewed Jul 20, 2026

View reviewed changes

Comment thread src/transformers/models/kimi_k25/modeling_kimi_k25.py Outdated

Comment thread src/transformers/models/gemma4/modular_gemma4.py Outdated

zucchini-nlp added 3 commits July 20, 2026 11:30

comments

230d7b7

fix diffusion gemma

3a70eaa

Merge branch 'main' into image-output-shapes

a9e3701

zucchini-nlp enabled auto-merge July 22, 2026 08:13

style

c4ffcb7

zucchini-nlp added this pull request to the merge queue Jul 22, 2026

Merged via the queue into huggingface:main with commit 378d2f6 Jul 22, 2026
103 checks passed

zucchini-nlp deleted the image-output-shapes branch July 22, 2026 08:52

Uh oh!

Conversation

zucchini-nlp commented Jun 4, 2026 • edited by github-actions Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do?

Uh oh!

HuggingFaceDocBuilderDev commented Jun 4, 2026

Uh oh!

zucchini-nlp commented Jun 4, 2026

Uh oh!

github-actions Bot commented Jun 4, 2026

Uh oh!

github-actions Bot commented Jun 4, 2026

CI Results

Commit Info

Uh oh!

vasqu left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

vasqu commented Jun 17, 2026

Uh oh!

zucchini-nlp commented Jun 18, 2026

Uh oh!

zucchini-nlp commented Jun 18, 2026

Uh oh!

vasqu commented Jun 18, 2026

Uh oh!

zucchini-nlp commented Jun 18, 2026

Uh oh!

remi-or commented Jun 18, 2026

Uh oh!

tomaarsen commented Jun 18, 2026

Uh oh!

zucchini-nlp commented Jun 18, 2026

Uh oh!

vasqu commented Jun 18, 2026

Uh oh!

zucchini-nlp commented Jul 17, 2026

Uh oh!

vasqu left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

tomaarsen left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

zucchini-nlp commented Jul 20, 2026

Uh oh!

github-actions Bot commented Jul 22, 2026

Uh oh!

github-actions Bot commented Jul 22, 2026

CI recap

Uh oh!

Uh oh!

Reviewers

Assignees

zucchini-nlp commented Jun 4, 2026 •

edited by github-actions Bot

Loading

vasqu left a comment •

edited

Loading