diff --git a/MIGRATION_GUIDE_V5.md b/MIGRATION_GUIDE_V5.md index 7c755fa0c3b2..f42a7bb2c039 100644 --- a/MIGRATION_GUIDE_V5.md +++ b/MIGRATION_GUIDE_V5.md @@ -453,6 +453,37 @@ We dropped support for two torch APIs: Those APIs were deprecated by the PyTorch team, and we're instead focusing on the supported APIs `dynamo` and `export`. +### Feature extraction helpers: `get_*_features` + +Many multi-modal models expose convenience methods such as `get_text_features`, `get_image_features`, `get_audio_features`, and `get_video_features` to run inference on a single modality without calling `model(**inputs)` directly. + +Starting with v5, these 4 helper methods now return a `BaseModelOutputWithPooling` (or a subclass) instead of only a pooled embedding tensor: + +- `last_hidden_state`: unpooled token/patch/frame embeddings for the requested modality. +- `pooler_output`: pooled representation (what most models previously returned from `get_*_features`). +- `hidden_states`: full hidden states for all layers when `output_hidden_states=True` is passed. +- `attentions`: attention maps when `output_attentions=True` is passed. + +> [!IMPORTANT] +> There is **no single universal shape** for `last_hidden_state` or `pooler_output`. It's recommended to inspect a small forward pass before making assumptions about shapes or semantics. + +If your code previously did something like this: + +```python +text_embeddings = model.get_text_features(**inputs) +``` + +and you used `text_embeddings` as a tensor, you should now explicitly use `return_dict=True` take the `pooler_output` field from the returned `BaseModelOutputWithPooling`: + +```python +outputs = model.get_text_features(**inputs, return_dict=True) +text_embeddings = outputs.pooler_output +``` + +This will match the previous behavior in the large majority of cases. If your model-specific implementation returned a tuple of results before, those values should now be accessible as fields on the corresponding `BaseModelOutputWithPooling` subclass. + +Linked PR: https://github.com/huggingface/transformers/pull/42564 + ## Quantization changes We clean up the quantization API in transformers, and significantly refactor the weight loading as highlighted @@ -558,7 +589,7 @@ Linked PRs: - `use_mps_device` -> mps will be used by default if detected - `fp16_backend` and `half_precision_backend` -> we will only rely on torch.amp as everything has been upstream to torch - `no_cuda` -> `use_cpu` -- ` include_tokens_per_second` -> `include_num_input_tokens_seen` +- `include_tokens_per_second` -> `include_num_input_tokens_seen` - `use_legacy_prediction_loop` -> we only use `evaluation_loop` function from now on ### Removing deprecated arguments in `Trainer` @@ -574,7 +605,7 @@ Linked PRs: ### New defaults for `Trainer` -- `use_cache` in the model config will be set to `False`. You can still change the cache value through `TrainingArguments` `usel_cache` argument if needed. +- `use_cache` in the model config will be set to `False`. You can still change the cache value through `TrainingArguments` `use_cache` argument if needed. ## Pipelines diff --git a/docs/source/en/model_doc/aimv2.md b/docs/source/en/model_doc/aimv2.md index acf9c4de12fe..eb0fdd3b47cd 100644 --- a/docs/source/en/model_doc/aimv2.md +++ b/docs/source/en/model_doc/aimv2.md @@ -89,6 +89,8 @@ probs = outputs.logits_per_image.softmax(dim=-1) [[autodoc]] Aimv2Model - forward + - get_text_features + - get_image_features ## Aimv2VisionModel diff --git a/docs/source/en/model_doc/aria.md b/docs/source/en/model_doc/aria.md index 1d33981b334d..44875b48315d 100644 --- a/docs/source/en/model_doc/aria.md +++ b/docs/source/en/model_doc/aria.md @@ -175,3 +175,4 @@ print(response) [[autodoc]] AriaForConditionalGeneration - forward + - get_image_features diff --git a/docs/source/en/model_doc/audioflamingo3.md b/docs/source/en/model_doc/audioflamingo3.md index c9e670274063..6c585e8adb51 100644 --- a/docs/source/en/model_doc/audioflamingo3.md +++ b/docs/source/en/model_doc/audioflamingo3.md @@ -401,3 +401,4 @@ are forwarded, so you can tweak padding or tensor formats just like when calling [[autodoc]] AudioFlamingo3ForConditionalGeneration - forward + - get_audio_features diff --git a/docs/source/en/model_doc/aya_vision.md b/docs/source/en/model_doc/aya_vision.md index 4f3a77007d6a..99ed5d28cf4d 100644 --- a/docs/source/en/model_doc/aya_vision.md +++ b/docs/source/en/model_doc/aya_vision.md @@ -274,3 +274,4 @@ print(processor.tokenizer.decode(generated[0], skip_special_tokens=True)) [[autodoc]] AyaVisionForConditionalGeneration - forward + - get_image_features diff --git a/docs/source/en/model_doc/blip-2.md b/docs/source/en/model_doc/blip-2.md index e2a260f8def7..43abbb6d31a2 100644 --- a/docs/source/en/model_doc/blip-2.md +++ b/docs/source/en/model_doc/blip-2.md @@ -97,6 +97,7 @@ If you're interested in submitting a resource to be included here, please feel f [[autodoc]] Blip2ForConditionalGeneration - forward - generate + - get_image_features ## Blip2ForImageTextRetrieval diff --git a/docs/source/en/model_doc/chameleon.md b/docs/source/en/model_doc/chameleon.md index 3a9f1a707a46..9faeade71bde 100644 --- a/docs/source/en/model_doc/chameleon.md +++ b/docs/source/en/model_doc/chameleon.md @@ -203,8 +203,10 @@ model = ChameleonForConditionalGeneration.from_pretrained( [[autodoc]] ChameleonModel - forward + - get_image_features ## ChameleonForConditionalGeneration [[autodoc]] ChameleonForConditionalGeneration - forward + - get_image_features diff --git a/docs/source/en/model_doc/cohere2_vision.md b/docs/source/en/model_doc/cohere2_vision.md index 49771d1feca4..b8a1ceae0e55 100644 --- a/docs/source/en/model_doc/cohere2_vision.md +++ b/docs/source/en/model_doc/cohere2_vision.md @@ -125,11 +125,13 @@ print(outputs) [[autodoc]] Cohere2VisionForConditionalGeneration - forward + - get_image_features ## Cohere2VisionModel [[autodoc]] Cohere2VisionModel - forward + - get_image_features ## Cohere2VisionImageProcessorFast diff --git a/docs/source/en/model_doc/deepseek_vl.md b/docs/source/en/model_doc/deepseek_vl.md index 7c5502849b7e..e88ad717561e 100644 --- a/docs/source/en/model_doc/deepseek_vl.md +++ b/docs/source/en/model_doc/deepseek_vl.md @@ -223,6 +223,7 @@ model = DeepseekVLForConditionalGeneration.from_pretrained( [[autodoc]] DeepseekVLModel - forward + - get_image_features ## DeepseekVLForConditionalGeneration diff --git a/docs/source/en/model_doc/deepseek_vl_hybrid.md b/docs/source/en/model_doc/deepseek_vl_hybrid.md index 35cf380f95ba..879f3afa51fc 100644 --- a/docs/source/en/model_doc/deepseek_vl_hybrid.md +++ b/docs/source/en/model_doc/deepseek_vl_hybrid.md @@ -222,6 +222,7 @@ model = DeepseekVLHybridForConditionalGeneration.from_pretrained( [[autodoc]] DeepseekVLHybridModel - forward + - get_image_features ## DeepseekVLHybridForConditionalGeneration diff --git a/docs/source/en/model_doc/edgetam.md b/docs/source/en/model_doc/edgetam.md index 7e52f74da982..173b89533c83 100644 --- a/docs/source/en/model_doc/edgetam.md +++ b/docs/source/en/model_doc/edgetam.md @@ -330,3 +330,4 @@ EdgeTAM can use masks from previous predictions as input to refine segmentation: [[autodoc]] EdgeTamModel - forward + - get_image_features diff --git a/docs/source/en/model_doc/edgetam_video.md b/docs/source/en/model_doc/edgetam_video.md index 733591d6565b..2274d9eb9ae2 100644 --- a/docs/source/en/model_doc/edgetam_video.md +++ b/docs/source/en/model_doc/edgetam_video.md @@ -294,3 +294,4 @@ Tracked 2 objects through 200 frames [[autodoc]] EdgeTamVideoModel - forward + - get_image_features diff --git a/docs/source/en/model_doc/ernie4_5_vl_moe.md b/docs/source/en/model_doc/ernie4_5_vl_moe.md index 71ffc1ba0a97..715bc55184bc 100644 --- a/docs/source/en/model_doc/ernie4_5_vl_moe.md +++ b/docs/source/en/model_doc/ernie4_5_vl_moe.md @@ -222,8 +222,12 @@ print(output_text) [[autodoc]] Ernie4_5_VL_MoeModel - forward + - get_video_features + - get_image_features ## Ernie4_5_VL_MoeForConditionalGeneration [[autodoc]] Ernie4_5_VL_MoeForConditionalGeneration - forward + - get_video_features + - get_image_features diff --git a/docs/source/en/model_doc/fast_vlm.md b/docs/source/en/model_doc/fast_vlm.md index acb1aeb7c0c6..fcd28e1130fc 100644 --- a/docs/source/en/model_doc/fast_vlm.md +++ b/docs/source/en/model_doc/fast_vlm.md @@ -171,3 +171,4 @@ Flash Attention 2 is an even faster, optimized version of the previous optimizat [[autodoc]] FastVlmForConditionalGeneration - forward + - get_image_features diff --git a/docs/source/en/model_doc/florence2.md b/docs/source/en/model_doc/florence2.md index bc7dd1368e3d..2ef085502824 100644 --- a/docs/source/en/model_doc/florence2.md +++ b/docs/source/en/model_doc/florence2.md @@ -177,11 +177,13 @@ print(parsed_answer) [[autodoc]] Florence2Model - forward + - get_image_features ## Florence2ForConditionalGeneration [[autodoc]] Florence2ForConditionalGeneration - forward + - get_image_features ## Florence2VisionBackbone diff --git a/docs/source/en/model_doc/gemma3.md b/docs/source/en/model_doc/gemma3.md index 381a6a4ff0f6..948b0ed829c2 100644 --- a/docs/source/en/model_doc/gemma3.md +++ b/docs/source/en/model_doc/gemma3.md @@ -271,6 +271,7 @@ visualizer("What is shown in this image?") [[autodoc]] Gemma3ForConditionalGeneration - forward + - get_image_features ## Gemma3ForSequenceClassification diff --git a/docs/source/en/model_doc/gemma3n.md b/docs/source/en/model_doc/gemma3n.md index edfde4041395..a3531262c36b 100644 --- a/docs/source/en/model_doc/gemma3n.md +++ b/docs/source/en/model_doc/gemma3n.md @@ -188,6 +188,8 @@ echo -e "Plants create energy through a process known as" | transformers run --t [[autodoc]] Gemma3nModel - forward + - get_image_features + - get_audio_features ## Gemma3nForCausalLM @@ -198,6 +200,7 @@ echo -e "Plants create energy through a process known as" | transformers run --t [[autodoc]] Gemma3nForConditionalGeneration - forward + - get_image_features [altup]: https://proceedings.neurips.cc/paper_files/paper/2023/hash/f2059277ac6ce66e7e5543001afa8bb5-Abstract-Conference.html [attention-mask-viz]: https://github.com/huggingface/transformers/blob/beb9b5b02246b9b7ee81ddf938f93f44cfeaad19/src/transformers/utils/attention_visualizer.py#L139 diff --git a/docs/source/en/model_doc/glm46v.md b/docs/source/en/model_doc/glm46v.md index bc5cbdc4ee43..c0877e66bf8b 100644 --- a/docs/source/en/model_doc/glm46v.md +++ b/docs/source/en/model_doc/glm46v.md @@ -78,8 +78,12 @@ This model was contributed by [Raushan Turganbay](https://huggingface.co/Raushan [[autodoc]] Glm46VModel - forward + - get_video_features + - get_image_features ## Glm46VForConditionalGeneration [[autodoc]] Glm46VForConditionalGeneration - forward + - get_video_features + - get_image_features diff --git a/docs/source/en/model_doc/glm4v.md b/docs/source/en/model_doc/glm4v.md index 206287f9d576..d199d5b03e2d 100644 --- a/docs/source/en/model_doc/glm4v.md +++ b/docs/source/en/model_doc/glm4v.md @@ -215,19 +215,23 @@ print(output_text) ## Glm4vVisionModel [[autodoc]] Glm4vVisionModel -- forward + - forward ## Glm4vTextModel [[autodoc]] Glm4vTextModel -- forward + - forward ## Glm4vModel [[autodoc]] Glm4vModel -- forward + - forward + - get_video_features + - get_image_features ## Glm4vForConditionalGeneration [[autodoc]] Glm4vForConditionalGeneration -- forward + - forward + - get_video_features + - get_image_features diff --git a/docs/source/en/model_doc/glm4v_moe.md b/docs/source/en/model_doc/glm4v_moe.md index a67906ea001e..9ba76d624288 100644 --- a/docs/source/en/model_doc/glm4v_moe.md +++ b/docs/source/en/model_doc/glm4v_moe.md @@ -76,8 +76,12 @@ This model was contributed by [Raushan Turganbay](https://huggingface.co/Raushan [[autodoc]] Glm4vMoeModel - forward + - get_video_features + - get_image_features ## Glm4vMoeForConditionalGeneration [[autodoc]] Glm4vMoeForConditionalGeneration - forward + - get_video_features + - get_image_features diff --git a/docs/source/en/model_doc/glm_image.md b/docs/source/en/model_doc/glm_image.md index 4b4ff609d2dd..2d08b6b7db1b 100644 --- a/docs/source/en/model_doc/glm_image.md +++ b/docs/source/en/model_doc/glm_image.md @@ -16,7 +16,7 @@ limitations under the License. ⚠️ Note that this file is in Markdown but contain specific syntax for our doc-builder (similar to MDX) that may not be rendered properly in your Markdown viewer. --> -*This model was released on 2026-01-10 and added to Hugging Face Transformers on 2026-01-10.* +*This model was released on 2026-01-10 and added to Hugging Face Transformers on 2026-01-13.* # GlmImage @@ -199,8 +199,10 @@ print(f"Output tokens: {output_tokens}") [[autodoc]] GlmImageModel - forward + - get_image_features ## GlmImageForConditionalGeneration [[autodoc]] GlmImageForConditionalGeneration - forward + - get_image_features diff --git a/docs/source/en/model_doc/got_ocr2.md b/docs/source/en/model_doc/got_ocr2.md index 8deb85fb144c..3611917dcdf6 100644 --- a/docs/source/en/model_doc/got_ocr2.md +++ b/docs/source/en/model_doc/got_ocr2.md @@ -291,3 +291,4 @@ alt="drawing" width="600"/> [[autodoc]] GotOcr2ForConditionalGeneration - forward + - get_image_features diff --git a/docs/source/en/model_doc/granite_speech.md b/docs/source/en/model_doc/granite_speech.md index 286d62bd9341..72699bf442cb 100644 --- a/docs/source/en/model_doc/granite_speech.md +++ b/docs/source/en/model_doc/granite_speech.md @@ -170,3 +170,4 @@ for i, transcription in enumerate(transcriptions): [[autodoc]] GraniteSpeechForConditionalGeneration - forward + - get_audio_features diff --git a/docs/source/en/model_doc/idefics2.md b/docs/source/en/model_doc/idefics2.md index 858f5de9b5b8..38ba46e7a0f6 100644 --- a/docs/source/en/model_doc/idefics2.md +++ b/docs/source/en/model_doc/idefics2.md @@ -208,11 +208,13 @@ A list of official Hugging Face and community (indicated by 🌎) resources to h [[autodoc]] Idefics2Model - forward + - get_image_features ## Idefics2ForConditionalGeneration [[autodoc]] Idefics2ForConditionalGeneration - forward + - get_image_features ## Idefics2ImageProcessor diff --git a/docs/source/en/model_doc/idefics3.md b/docs/source/en/model_doc/idefics3.md index 2a4d9eb242f8..c8b59da5f492 100644 --- a/docs/source/en/model_doc/idefics3.md +++ b/docs/source/en/model_doc/idefics3.md @@ -70,11 +70,13 @@ This model was contributed by [amyeroberts](https://huggingface.co/amyeroberts) [[autodoc]] Idefics3Model - forward + - get_image_features ## Idefics3ForConditionalGeneration [[autodoc]] Idefics3ForConditionalGeneration - forward + - get_image_features ## Idefics3ImageProcessor diff --git a/docs/source/en/model_doc/instructblip.md b/docs/source/en/model_doc/instructblip.md index e2fef6c1e35c..ef2d96d97213 100644 --- a/docs/source/en/model_doc/instructblip.md +++ b/docs/source/en/model_doc/instructblip.md @@ -78,3 +78,4 @@ The attributes can be obtained from model config, as `model.config.num_query_tok [[autodoc]] InstructBlipForConditionalGeneration - forward - generate + - get_image_features diff --git a/docs/source/en/model_doc/instructblipvideo.md b/docs/source/en/model_doc/instructblipvideo.md index 3a36d5366b66..27852b2f2039 100644 --- a/docs/source/en/model_doc/instructblipvideo.md +++ b/docs/source/en/model_doc/instructblipvideo.md @@ -83,3 +83,4 @@ The attributes can be obtained from model config, as `model.config.num_query_tok [[autodoc]] InstructBlipVideoForConditionalGeneration - forward - generate + - get_video_features diff --git a/docs/source/en/model_doc/internvl.md b/docs/source/en/model_doc/internvl.md index e2a712697d1d..a838d9f14531 100644 --- a/docs/source/en/model_doc/internvl.md +++ b/docs/source/en/model_doc/internvl.md @@ -339,11 +339,13 @@ This example showcases how to handle a batch of chat conversations with interlea [[autodoc]] InternVLModel - forward + - get_image_features ## InternVLForConditionalGeneration [[autodoc]] InternVLForConditionalGeneration - forward + - get_image_features ## InternVLProcessor diff --git a/docs/source/en/model_doc/janus.md b/docs/source/en/model_doc/janus.md index 916592489b5f..56b9ef5d0ff5 100644 --- a/docs/source/en/model_doc/janus.md +++ b/docs/source/en/model_doc/janus.md @@ -229,6 +229,7 @@ for i, image in enumerate(images['pixel_values']): [[autodoc]] JanusModel - forward + - get_image_features ## JanusForConditionalGeneration diff --git a/docs/source/en/model_doc/kosmos-2.md b/docs/source/en/model_doc/kosmos-2.md index c449dfd71ad1..c51498418c16 100644 --- a/docs/source/en/model_doc/kosmos-2.md +++ b/docs/source/en/model_doc/kosmos-2.md @@ -96,6 +96,7 @@ This model was contributed by [Yih-Dar SHIEH](https://huggingface.co/ydshieh). T [[autodoc]] Kosmos2Model - forward + - get_image_features ## Kosmos2ForConditionalGeneration diff --git a/docs/source/en/model_doc/lfm2_vl.md b/docs/source/en/model_doc/lfm2_vl.md index de77b984369d..d9c101439e34 100644 --- a/docs/source/en/model_doc/lfm2_vl.md +++ b/docs/source/en/model_doc/lfm2_vl.md @@ -92,8 +92,10 @@ processor.batch_decode(outputs, skip_special_tokens=True)[0] [[autodoc]] Lfm2VlModel - forward + - get_image_features ## Lfm2VlForConditionalGeneration [[autodoc]] Lfm2VlForConditionalGeneration - forward + - get_image_features diff --git a/docs/source/en/model_doc/lighton_ocr.md b/docs/source/en/model_doc/lighton_ocr.md index 274a2c473848..77671206c2ae 100644 --- a/docs/source/en/model_doc/lighton_ocr.md +++ b/docs/source/en/model_doc/lighton_ocr.md @@ -73,8 +73,10 @@ print(output_text) [[autodoc]] LightOnOcrModel - forward + - get_image_features ## LightOnOcrForConditionalGeneration [[autodoc]] LightOnOcrForConditionalGeneration - forward + - get_image_features diff --git a/docs/source/en/model_doc/llama4.md b/docs/source/en/model_doc/llama4.md index 6ccb3190e485..f430366f49b1 100644 --- a/docs/source/en/model_doc/llama4.md +++ b/docs/source/en/model_doc/llama4.md @@ -426,6 +426,7 @@ model = Llama4ForConditionalGeneration.from_pretrained( [[autodoc]] Llama4ForConditionalGeneration - forward + - get_image_features ## Llama4ForCausalLM diff --git a/docs/source/en/model_doc/llava.md b/docs/source/en/model_doc/llava.md index 2adaa15db8f3..bb24a909d42b 100644 --- a/docs/source/en/model_doc/llava.md +++ b/docs/source/en/model_doc/llava.md @@ -260,3 +260,4 @@ A list of official Hugging Face and community (indicated by 🌎) resources to h [[autodoc]] LlavaForConditionalGeneration - forward + - get_image_features diff --git a/docs/source/en/model_doc/llava_next.md b/docs/source/en/model_doc/llava_next.md index 11a7a8dfe76c..7e78219bc633 100644 --- a/docs/source/en/model_doc/llava_next.md +++ b/docs/source/en/model_doc/llava_next.md @@ -216,3 +216,4 @@ print(processor.decode(output[0], skip_special_tokens=True)) [[autodoc]] LlavaNextForConditionalGeneration - forward + - get_image_features diff --git a/docs/source/en/model_doc/llava_next_video.md b/docs/source/en/model_doc/llava_next_video.md index d7c6b82104d5..06c0686f5cfd 100644 --- a/docs/source/en/model_doc/llava_next_video.md +++ b/docs/source/en/model_doc/llava_next_video.md @@ -259,3 +259,5 @@ model = LlavaNextVideoForConditionalGeneration.from_pretrained( [[autodoc]] LlavaNextVideoForConditionalGeneration - forward + - get_image_features + - get_video_features diff --git a/docs/source/en/model_doc/llava_onevision.md b/docs/source/en/model_doc/llava_onevision.md index c271421862fe..0a292bfd8b10 100644 --- a/docs/source/en/model_doc/llava_onevision.md +++ b/docs/source/en/model_doc/llava_onevision.md @@ -322,3 +322,5 @@ model = LlavaOnevisionForConditionalGeneration.from_pretrained( [[autodoc]] LlavaOnevisionForConditionalGeneration - forward + - get_image_features + - get_video_features diff --git a/docs/source/en/model_doc/lw_detr.md b/docs/source/en/model_doc/lw_detr.md index fb09718fb56a..8204ada503af 100644 --- a/docs/source/en/model_doc/lw_detr.md +++ b/docs/source/en/model_doc/lw_detr.md @@ -13,7 +13,7 @@ specific language governing permissions and limitations under the License. rendered properly in your Markdown viewer. --> -*This model was released on 2024-04-05 and added to Hugging Face Transformers on 2026-01-10.* +*This model was released on 2024-04-05 and added to Hugging Face Transformers on 2026-01-12.*
diff --git a/docs/source/en/model_doc/mistral3.md b/docs/source/en/model_doc/mistral3.md index 7911200ad744..60338fb99682 100644 --- a/docs/source/en/model_doc/mistral3.md +++ b/docs/source/en/model_doc/mistral3.md @@ -254,3 +254,4 @@ messages = [ [[autodoc]] Mistral3ForConditionalGeneration - forward + - get_image_features diff --git a/docs/source/en/model_doc/ovis2.md b/docs/source/en/model_doc/ovis2.md index eacce1b30b38..0476c19339c4 100644 --- a/docs/source/en/model_doc/ovis2.md +++ b/docs/source/en/model_doc/ovis2.md @@ -95,6 +95,7 @@ with torch.inference_mode(): [[autodoc]] Ovis2ForConditionalGeneration - forward + - get_image_features ## Ovis2ImageProcessor diff --git a/docs/source/en/model_doc/paddleocr_vl.md b/docs/source/en/model_doc/paddleocr_vl.md index cc4df0774316..07e68e488cb5 100644 --- a/docs/source/en/model_doc/paddleocr_vl.md +++ b/docs/source/en/model_doc/paddleocr_vl.md @@ -206,6 +206,7 @@ model = AutoModelForImageTextToText.from_pretrained("PaddlePaddle/PaddleOCR-VL", [[autodoc]] PaddleOCRVLForConditionalGeneration - forward + - get_image_features ## PaddleOCRVLConfig diff --git a/docs/source/en/model_doc/paligemma.md b/docs/source/en/model_doc/paligemma.md index 638d5f47ebc2..0d984a5e220e 100644 --- a/docs/source/en/model_doc/paligemma.md +++ b/docs/source/en/model_doc/paligemma.md @@ -185,3 +185,4 @@ visualizer(" What is in this image?") [[autodoc]] PaliGemmaForConditionalGeneration - forward + - get_image_features diff --git a/docs/source/en/model_doc/qwen2_5_vl.md b/docs/source/en/model_doc/qwen2_5_vl.md index 693172157317..f579dc85f1b0 100644 --- a/docs/source/en/model_doc/qwen2_5_vl.md +++ b/docs/source/en/model_doc/qwen2_5_vl.md @@ -257,8 +257,12 @@ model = Qwen2_5_VLForConditionalGeneration.from_pretrained( [[autodoc]] Qwen2_5_VLModel - forward + - get_video_features + - get_image_features ## Qwen2_5_VLForConditionalGeneration [[autodoc]] Qwen2_5_VLForConditionalGeneration - forward + - get_video_features + - get_image_features diff --git a/docs/source/en/model_doc/qwen2_vl.md b/docs/source/en/model_doc/qwen2_vl.md index 6960f7f07b03..84017f0028ee 100644 --- a/docs/source/en/model_doc/qwen2_vl.md +++ b/docs/source/en/model_doc/qwen2_vl.md @@ -313,8 +313,12 @@ model = Qwen2VLForConditionalGeneration.from_pretrained( [[autodoc]] Qwen2VLModel - forward + - get_video_features + - get_image_features ## Qwen2VLForConditionalGeneration [[autodoc]] Qwen2VLForConditionalGeneration - forward + - get_video_features + - get_image_features diff --git a/docs/source/en/model_doc/qwen3_vl.md b/docs/source/en/model_doc/qwen3_vl.md index 856c92cf9897..f5542da032e6 100644 --- a/docs/source/en/model_doc/qwen3_vl.md +++ b/docs/source/en/model_doc/qwen3_vl.md @@ -112,8 +112,12 @@ print(output_text) [[autodoc]] Qwen3VLModel - forward + - get_video_features + - get_image_features ## Qwen3VLForConditionalGeneration [[autodoc]] Qwen3VLForConditionalGeneration - forward + - get_video_features + - get_image_features diff --git a/docs/source/en/model_doc/qwen3_vl_moe.md b/docs/source/en/model_doc/qwen3_vl_moe.md index 771f6d411cf2..00923c7a245e 100644 --- a/docs/source/en/model_doc/qwen3_vl_moe.md +++ b/docs/source/en/model_doc/qwen3_vl_moe.md @@ -103,8 +103,12 @@ print(output_text) [[autodoc]] Qwen3VLMoeModel - forward + - get_video_features + - get_image_features ## Qwen3VLMoeForConditionalGeneration [[autodoc]] Qwen3VLMoeForConditionalGeneration - forward + - get_video_features + - get_image_features diff --git a/docs/source/en/model_doc/sam2.md b/docs/source/en/model_doc/sam2.md index f0363616867a..3d0514de57cb 100644 --- a/docs/source/en/model_doc/sam2.md +++ b/docs/source/en/model_doc/sam2.md @@ -351,3 +351,4 @@ SAM2 can use masks from previous predictions as input to refine segmentation: [[autodoc]] Sam2Model - forward + - get_image_features diff --git a/docs/source/en/model_doc/sam2_video.md b/docs/source/en/model_doc/sam2_video.md index b55daf141250..efc2b4c43207 100644 --- a/docs/source/en/model_doc/sam2_video.md +++ b/docs/source/en/model_doc/sam2_video.md @@ -316,3 +316,4 @@ A list of official Hugging Face and community (indicated by 🌎) resources to h [[autodoc]] Sam2VideoModel - forward - propagate_in_video_iterator + - get_image_features diff --git a/docs/source/en/model_doc/sam3.md b/docs/source/en/model_doc/sam3.md index 3ec8b7d7b69f..bae0fe97bde3 100644 --- a/docs/source/en/model_doc/sam3.md +++ b/docs/source/en/model_doc/sam3.md @@ -437,3 +437,4 @@ SAM3 uses the following label conventions: [[autodoc]] Sam3Model - forward + - get_text_features diff --git a/docs/source/en/model_doc/sam3_tracker.md b/docs/source/en/model_doc/sam3_tracker.md index 8aa2342e4e7a..c64c8b711c45 100644 --- a/docs/source/en/model_doc/sam3_tracker.md +++ b/docs/source/en/model_doc/sam3_tracker.md @@ -319,6 +319,7 @@ Sam3Tracker can use masks from previous predictions as input to refine segmentat [[autodoc]] Sam3TrackerModel - forward + - get_image_features ## Sam3TrackerPreTrainedModel diff --git a/docs/source/en/model_doc/sam3_tracker_video.md b/docs/source/en/model_doc/sam3_tracker_video.md index 97c13d05a8f1..5c7a7679600e 100644 --- a/docs/source/en/model_doc/sam3_tracker_video.md +++ b/docs/source/en/model_doc/sam3_tracker_video.md @@ -303,3 +303,4 @@ Tracked 2 objects through 180 frames [[autodoc]] Sam3TrackerVideoModel - forward - propagate_in_video_iterator + - get_image_features diff --git a/docs/source/en/model_doc/smolvlm.md b/docs/source/en/model_doc/smolvlm.md index 63fd34731038..762ff70408c0 100644 --- a/docs/source/en/model_doc/smolvlm.md +++ b/docs/source/en/model_doc/smolvlm.md @@ -187,11 +187,13 @@ print(generated_texts[0]) [[autodoc]] SmolVLMModel - forward + - get_image_features ## SmolVLMForConditionalGeneration [[autodoc]] SmolVLMForConditionalGeneration - forward + - get_image_features ## SmolVLMImageProcessor diff --git a/docs/source/en/model_doc/t5gemma2.md b/docs/source/en/model_doc/t5gemma2.md index bcadd6022b45..346f7ac0ae6b 100644 --- a/docs/source/en/model_doc/t5gemma2.md +++ b/docs/source/en/model_doc/t5gemma2.md @@ -107,6 +107,7 @@ print(processor.decode(generation[0])) [[autodoc]] T5Gemma2ForConditionalGeneration - forward + - get_image_features ## T5Gemma2ForSequenceClassification diff --git a/docs/source/en/model_doc/video_llama_3.md b/docs/source/en/model_doc/video_llama_3.md index dc5c8f1c2a22..7cb0aea5a2d2 100644 --- a/docs/source/en/model_doc/video_llama_3.md +++ b/docs/source/en/model_doc/video_llama_3.md @@ -218,6 +218,8 @@ model = VideoLlama3ForConditionalGeneration.from_pretrained( [[autodoc]] VideoLlama3Model - forward + - get_video_features + - get_image_features ## VideoLlama3VisionModel @@ -228,3 +230,5 @@ model = VideoLlama3ForConditionalGeneration.from_pretrained( [[autodoc]] VideoLlama3ForConditionalGeneration - forward + - get_video_features + - get_image_features diff --git a/docs/source/en/model_doc/video_llava.md b/docs/source/en/model_doc/video_llava.md index 24437684716f..ccc77575d451 100644 --- a/docs/source/en/model_doc/video_llava.md +++ b/docs/source/en/model_doc/video_llava.md @@ -223,3 +223,4 @@ model = VideoLlavaForConditionalGeneration.from_pretrained( [[autodoc]] VideoLlavaForConditionalGeneration - forward + - get_image_features diff --git a/docs/source/en/model_doc/vipllava.md b/docs/source/en/model_doc/vipllava.md index 336d0c100a52..3a687cc271ca 100644 --- a/docs/source/en/model_doc/vipllava.md +++ b/docs/source/en/model_doc/vipllava.md @@ -109,3 +109,4 @@ A chat between a curious human and an artificial intelligence assistant. The ass [[autodoc]] VipLlavaForConditionalGeneration - forward + - get_image_features diff --git a/docs/source/en/model_doc/vision-text-dual-encoder.md b/docs/source/en/model_doc/vision-text-dual-encoder.md index d4ba9878bdbd..39ba554a0455 100644 --- a/docs/source/en/model_doc/vision-text-dual-encoder.md +++ b/docs/source/en/model_doc/vision-text-dual-encoder.md @@ -48,3 +48,5 @@ new zero-shot vision tasks such as image classification or retrieval. [[autodoc]] VisionTextDualEncoderModel - forward + - get_text_features + - get_image_features diff --git a/docs/source/en/model_doc/voxtral.md b/docs/source/en/model_doc/voxtral.md index b22bb18efab1..605d04a5f15f 100644 --- a/docs/source/en/model_doc/voxtral.md +++ b/docs/source/en/model_doc/voxtral.md @@ -370,3 +370,4 @@ This model was contributed by [Eustache Le Bihan](https://huggingface.co/eustlb) [[autodoc]] VoxtralForConditionalGeneration - forward + - get_audio_features diff --git a/docs/source/ko/model_doc/llama4.md b/docs/source/ko/model_doc/llama4.md index 49642a135b74..0272de22341a 100644 --- a/docs/source/ko/model_doc/llama4.md +++ b/docs/source/ko/model_doc/llama4.md @@ -420,24 +420,24 @@ model = Llama4ForConditionalGeneration.from_pretrained( ## Llama4ForConditionalGeneration [[autodoc]] Llama4ForConditionalGeneration -- forward + - forward ## Llama4ForCausalLM [[autodoc]] Llama4ForCausalLM -- forward + - forward ## Llama4TextModel [[autodoc]] Llama4TextModel -- forward + - forward ## Llama4ForCausalLM [[autodoc]] Llama4ForCausalLM -- forward + - forward ## Llama4VisionModel [[autodoc]] Llama4VisionModel -- forward \ No newline at end of file + - forward \ No newline at end of file diff --git a/src/transformers/models/afmoe/modeling_afmoe.py b/src/transformers/models/afmoe/modeling_afmoe.py index 1a311011d4a8..15e88fc1f00b 100644 --- a/src/transformers/models/afmoe/modeling_afmoe.py +++ b/src/transformers/models/afmoe/modeling_afmoe.py @@ -571,7 +571,7 @@ def forward( cache_position: torch.LongTensor | None = None, use_cache: bool | None = None, **kwargs: Unpack[TransformersKwargs], - ) -> MoeModelOutputWithPast: + ) -> tuple | MoeModelOutputWithPast: if (input_ids is None) ^ (inputs_embeds is not None): raise ValueError("You must specify exactly one of input_ids or inputs_embeds") diff --git a/src/transformers/models/afmoe/modular_afmoe.py b/src/transformers/models/afmoe/modular_afmoe.py index d3e9b5ffcce1..d81a659e905b 100644 --- a/src/transformers/models/afmoe/modular_afmoe.py +++ b/src/transformers/models/afmoe/modular_afmoe.py @@ -392,7 +392,7 @@ def forward( cache_position: torch.LongTensor | None = None, use_cache: bool | None = None, **kwargs: Unpack[TransformersKwargs], - ) -> MoeModelOutputWithPast: + ) -> tuple | MoeModelOutputWithPast: if (input_ids is None) ^ (inputs_embeds is not None): raise ValueError("You must specify exactly one of input_ids or inputs_embeds") diff --git a/src/transformers/models/aimv2/modeling_aimv2.py b/src/transformers/models/aimv2/modeling_aimv2.py index c9e3c5f1b50f..1000765e53a8 100644 --- a/src/transformers/models/aimv2/modeling_aimv2.py +++ b/src/transformers/models/aimv2/modeling_aimv2.py @@ -36,7 +36,7 @@ from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel from ...processing_utils import Unpack -from ...utils import ModelOutput, TransformersKwargs, auto_docstring, can_return_tuple, filter_out_non_signature_kwargs +from ...utils import ModelOutput, TransformersKwargs, auto_docstring, can_return_tuple from ...utils.generic import check_model_inputs from .configuration_aimv2 import Aimv2Config, Aimv2TextConfig, Aimv2VisionConfig @@ -603,19 +603,16 @@ def __init__(self, config: Aimv2Config): self.post_init() - @filter_out_non_signature_kwargs() + @can_return_tuple @auto_docstring def get_text_features( self, input_ids: torch.Tensor, attention_mask: torch.Tensor | None = None, position_ids: torch.Tensor | None = None, - ) -> torch.FloatTensor: + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: r""" - Returns: - text_features (`torch.FloatTensor` of shape `(batch_size, output_dim`): The text embeddings obtained by - applying the projection layer to the pooled output of [`Aimv2TextModel`]. - Examples: ```python @@ -634,24 +631,23 @@ def get_text_features( input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids, + return_dict=True, + **kwargs, ) pooled_output = text_outputs.pooler_output - text_features = self.text_projection(pooled_output) + text_outputs.pooler_output = self.text_projection(pooled_output) - return text_features + return text_outputs - @filter_out_non_signature_kwargs() + @can_return_tuple @auto_docstring def get_image_features( self, pixel_values: torch.FloatTensor, interpolate_pos_encoding: bool = False, - ) -> torch.FloatTensor: + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: r""" - Returns: - image_features (`torch.FloatTensor` of shape `(batch_size, output_dim`): The image embeddings obtained by - applying the projection layer to the pooled output of [`Aimv2VisionModel`]. - Examples: ```python @@ -673,11 +669,13 @@ def get_image_features( vision_outputs: BaseModelOutputWithPooling = self.vision_model( pixel_values=pixel_values, interpolate_pos_encoding=interpolate_pos_encoding, + return_dict=True, + **kwargs, ) pooled_output = vision_outputs.pooler_output - image_features = self.visual_projection(pooled_output) + vision_outputs.pooler_output = self.visual_projection(pooled_output) - return image_features + return vision_outputs @auto_docstring @can_return_tuple diff --git a/src/transformers/models/align/modeling_align.py b/src/transformers/models/align/modeling_align.py index 43dfe428ffdb..98d14fbb307c 100644 --- a/src/transformers/models/align/modeling_align.py +++ b/src/transformers/models/align/modeling_align.py @@ -31,8 +31,9 @@ BaseModelOutputWithPoolingAndNoAttention, ) from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel +from ...processing_utils import Unpack from ...pytorch_utils import apply_chunking_to_forward -from ...utils import ModelOutput, auto_docstring, can_return_tuple, filter_out_non_signature_kwargs, logging +from ...utils import ModelOutput, TransformersKwargs, auto_docstring, can_return_tuple, logging from .configuration_align import AlignConfig, AlignTextConfig, AlignVisionConfig @@ -1092,7 +1093,7 @@ def __init__(self, config: AlignConfig): # Initialize weights and apply final processing self.post_init() - @filter_out_non_signature_kwargs() + @can_return_tuple @auto_docstring def get_text_features( self, @@ -1101,12 +1102,9 @@ def get_text_features( token_type_ids: torch.Tensor | None = None, position_ids: torch.Tensor | None = None, inputs_embeds: torch.Tensor | None = None, - ) -> torch.FloatTensor: + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: r""" - Returns: - text_features (`torch.FloatTensor` of shape `(batch_size, output_dim`): The text embeddings obtained by - applying the projection layer to the pooled output of [`AlignTextModel`]. - Examples: ```python @@ -1120,26 +1118,26 @@ def get_text_features( >>> with torch.inference_mode(): ... text_features = model.get_text_features(**inputs) ```""" - text_outputs = self.text_model( + text_outputs: BaseModelOutputWithPooling = self.text_model( input_ids=input_ids, attention_mask=attention_mask, token_type_ids=token_type_ids, position_ids=position_ids, inputs_embeds=inputs_embeds, + return_dict=True, + **kwargs, ) last_hidden_state = text_outputs[0][:, 0, :] - text_features = self.text_projection(last_hidden_state) + text_outputs.pooler_output = self.text_projection(last_hidden_state) - return text_features + return text_outputs - @filter_out_non_signature_kwargs() + @can_return_tuple @auto_docstring - def get_image_features(self, pixel_values: torch.FloatTensor) -> torch.FloatTensor: + def get_image_features( + self, pixel_values: torch.FloatTensor, **kwargs: Unpack[TransformersKwargs] + ) -> tuple | BaseModelOutputWithPooling: r""" - Returns: - image_features (`torch.FloatTensor` of shape `(batch_size, output_dim`): The image embeddings obtained by - applying the projection layer to the pooled output of [`AlignVisionModel`]. - Examples: ```python @@ -1157,9 +1155,7 @@ def get_image_features(self, pixel_values: torch.FloatTensor) -> torch.FloatTens >>> with torch.inference_mode(): ... image_features = model.get_image_features(**inputs) ```""" - vision_outputs = self.vision_model(pixel_values=pixel_values) - image_features = vision_outputs.pooler_output - return image_features + return self.vision_model(pixel_values=pixel_values, **kwargs) @can_return_tuple @auto_docstring diff --git a/src/transformers/models/altclip/modeling_altclip.py b/src/transformers/models/altclip/modeling_altclip.py index 669308f332eb..13c94fa88362 100755 --- a/src/transformers/models/altclip/modeling_altclip.py +++ b/src/transformers/models/altclip/modeling_altclip.py @@ -31,8 +31,9 @@ BaseModelOutputWithPoolingAndProjection, ) from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel +from ...processing_utils import Unpack from ...pytorch_utils import apply_chunking_to_forward -from ...utils import ModelOutput, auto_docstring, can_return_tuple, filter_out_non_signature_kwargs, logging, torch_int +from ...utils import ModelOutput, TransformersKwargs, auto_docstring, can_return_tuple, logging, torch_int from ...utils.generic import is_flash_attention_requested from .configuration_altclip import AltCLIPConfig, AltCLIPTextConfig, AltCLIPVisionConfig @@ -1154,7 +1155,7 @@ def __init__(self, config: AltCLIPConfig): # Initialize weights and apply final processing self.post_init() - @filter_out_non_signature_kwargs() + @can_return_tuple @auto_docstring def get_text_features( self, @@ -1162,12 +1163,9 @@ def get_text_features( attention_mask: torch.Tensor | None = None, position_ids: torch.Tensor | None = None, token_type_ids: torch.Tensor | None = None, - ) -> torch.FloatTensor: + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: r""" - Returns: - text_features (`torch.FloatTensor` of shape `(batch_size, output_dim`): The text embeddings obtained by - applying the projection layer to the pooled output of [`AltCLIPTextModel`]. - Examples: ```python @@ -1181,29 +1179,28 @@ def get_text_features( >>> with torch.inference_mode(): ... text_features = model.get_text_features(**inputs) ```""" - text_outputs = self.text_model( + text_outputs: BaseModelOutputWithPoolingAndProjection = self.text_model( input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids, token_type_ids=token_type_ids, + return_dict=True, + **kwargs, ) pooled_output = text_outputs.pooler_output - text_features = self.text_projection(pooled_output) + text_outputs.pooler_output = self.text_projection(pooled_output) - return text_features + return text_outputs - @filter_out_non_signature_kwargs() + @can_return_tuple @auto_docstring def get_image_features( self, pixel_values: torch.FloatTensor, interpolate_pos_encoding: bool = False, - ) -> torch.FloatTensor: + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: r""" - Returns: - image_features (`torch.FloatTensor` of shape `(batch_size, output_dim`): The image embeddings obtained by - applying the projection layer to the pooled output of [`AltCLIPVisionModel`]. - Examples: ```python @@ -1224,11 +1221,13 @@ def get_image_features( vision_outputs = self.vision_model( pixel_values=pixel_values, interpolate_pos_encoding=interpolate_pos_encoding, + return_dict=True, + **kwargs, ) pooled_output = vision_outputs.pooler_output - image_features = self.visual_projection(pooled_output) + vision_outputs.pooler_output = self.visual_projection(pooled_output) - return image_features + return vision_outputs @auto_docstring def forward( diff --git a/src/transformers/models/aria/modeling_aria.py b/src/transformers/models/aria/modeling_aria.py index c3965ff88930..91aecac948b7 100644 --- a/src/transformers/models/aria/modeling_aria.py +++ b/src/transformers/models/aria/modeling_aria.py @@ -32,7 +32,12 @@ from ...masking_utils import create_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer -from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast, ModelOutput +from ...modeling_outputs import ( + BaseModelOutputWithPast, + BaseModelOutputWithPooling, + CausalLMOutputWithPast, + ModelOutput, +) from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel from ...processing_utils import Unpack @@ -910,33 +915,25 @@ def get_input_embeddings(self): def set_input_embeddings(self, value): self.language_model.set_input_embeddings(value) + @check_model_inputs(tie_last_hidden_states=False) + @auto_docstring( + custom_intro="Obtains image last hidden states from the vision tower and apply multimodal projection." + ) def get_image_features( self, pixel_values: torch.FloatTensor, pixel_mask: torch.FloatTensor | None = None, vision_feature_layer: int = -1, - ): - """ - Obtains image last hidden states from the vision tower and apply multimodal projection. - - Args: - pixel_values (`torch.FloatTensor]` of shape `(batch_size, channels, height, width)`): - The tensors corresponding to the input images. - pixel_mask (`torch.FloatTensor]`, *optional*): - The tensors corresponding to the input image mask. - vision_feature_layer (`Union[int, list[int]]`, *optional*): - The index of the layer to select the vision feature. If multiple indices are provided, - the vision feature of the corresponding indices will be concatenated to form the - vision features. - Returns: - image_features (`torch.Tensor`): Image feature tensor of shape `(num_images, image_length, embed_dim)`). - """ - vision_feature_layer = ( - vision_feature_layer if vision_feature_layer is not None else self.config.vision_feature_layer - ) + output_hidden_states: bool | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: patch_attention_mask = self._create_patch_attention_mask(pixel_mask) image_outputs = self.vision_tower( - pixel_values, patch_attention_mask=patch_attention_mask, output_hidden_states=True + pixel_values, + patch_attention_mask=patch_attention_mask, + output_hidden_states=True, # Ignore arg on purpose + return_dict=True, + **kwargs, ) image_attn_mask = None if patch_attention_mask is not None: @@ -944,8 +941,9 @@ def get_image_features( image_attn_mask = torch.logical_not(flattened_mask) selected_image_feature = image_outputs.hidden_states[vision_feature_layer] - image_features = self.multi_modal_projector(selected_image_feature, attn_mask=image_attn_mask) - return image_features + image_outputs.pooler_output = self.multi_modal_projector(selected_image_feature, attn_mask=image_attn_mask) + + return image_outputs def get_placeholder_mask( self, input_ids: torch.LongTensor, inputs_embeds: torch.FloatTensor, image_features: torch.FloatTensor @@ -995,7 +993,8 @@ def forward( pixel_values=pixel_values, pixel_mask=pixel_mask, vision_feature_layer=self.config.vision_feature_layer, - ) + return_dict=True, + ).pooler_output image_features = image_features.to(inputs_embeds.device, inputs_embeds.dtype) special_image_mask = self.get_placeholder_mask( input_ids, inputs_embeds=inputs_embeds, image_features=image_features @@ -1069,16 +1068,19 @@ def set_input_embeddings(self, value): def get_output_embeddings(self) -> nn.Module: return self.lm_head + @auto_docstring def get_image_features( self, pixel_values: torch.FloatTensor, pixel_mask: torch.FloatTensor | None = None, vision_feature_layer: int = -1, - ): + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: return self.model.get_image_features( pixel_values=pixel_values, pixel_mask=pixel_mask, vision_feature_layer=vision_feature_layer, + **kwargs, ) @can_return_tuple diff --git a/src/transformers/models/aria/modular_aria.py b/src/transformers/models/aria/modular_aria.py index 032349de6d88..cd84941a9c34 100644 --- a/src/transformers/models/aria/modular_aria.py +++ b/src/transformers/models/aria/modular_aria.py @@ -36,10 +36,12 @@ validate_preprocess_arguments, ) from ...modeling_flash_attention_utils import FlashAttentionKwargs +from ...modeling_outputs import BaseModelOutputWithPooling from ...modeling_utils import PreTrainedModel from ...processing_utils import ImagesKwargs, MultiModalData, ProcessingKwargs, ProcessorMixin, Unpack from ...tokenization_python import PreTokenizedInput, TextInput from ...utils import TensorType, TransformersKwargs, auto_docstring, can_return_tuple, logging +from ...utils.generic import check_model_inputs from ..auto import CONFIG_MAPPING, AutoConfig, AutoTokenizer from ..llama.configuration_llama import LlamaConfig from ..llama.modeling_llama import ( @@ -1258,33 +1260,25 @@ def _create_patch_attention_mask(self, pixel_mask): ) return (patches_subgrid.sum(dim=(-1, -2)) > 0).bool() + @check_model_inputs(tie_last_hidden_states=False) + @auto_docstring( + custom_intro="Obtains image last hidden states from the vision tower and apply multimodal projection." + ) def get_image_features( self, pixel_values: torch.FloatTensor, pixel_mask: torch.FloatTensor | None = None, vision_feature_layer: int = -1, - ): - """ - Obtains image last hidden states from the vision tower and apply multimodal projection. - - Args: - pixel_values (`torch.FloatTensor]` of shape `(batch_size, channels, height, width)`): - The tensors corresponding to the input images. - pixel_mask (`torch.FloatTensor]`, *optional*): - The tensors corresponding to the input image mask. - vision_feature_layer (`Union[int, list[int]]`, *optional*): - The index of the layer to select the vision feature. If multiple indices are provided, - the vision feature of the corresponding indices will be concatenated to form the - vision features. - Returns: - image_features (`torch.Tensor`): Image feature tensor of shape `(num_images, image_length, embed_dim)`). - """ - vision_feature_layer = ( - vision_feature_layer if vision_feature_layer is not None else self.config.vision_feature_layer - ) + output_hidden_states: bool | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: patch_attention_mask = self._create_patch_attention_mask(pixel_mask) image_outputs = self.vision_tower( - pixel_values, patch_attention_mask=patch_attention_mask, output_hidden_states=True + pixel_values, + patch_attention_mask=patch_attention_mask, + output_hidden_states=True, # Ignore arg on purpose + return_dict=True, + **kwargs, ) image_attn_mask = None if patch_attention_mask is not None: @@ -1292,9 +1286,12 @@ def get_image_features( image_attn_mask = torch.logical_not(flattened_mask) selected_image_feature = image_outputs.hidden_states[vision_feature_layer] - image_features = self.multi_modal_projector(selected_image_feature, attn_mask=image_attn_mask) - return image_features + image_outputs.pooler_output = self.multi_modal_projector(selected_image_feature, attn_mask=image_attn_mask) + return image_outputs + + @can_return_tuple + @auto_docstring def forward( self, input_ids: torch.LongTensor | None = None, @@ -1317,7 +1314,8 @@ def forward( pixel_values=pixel_values, pixel_mask=pixel_mask, vision_feature_layer=self.config.vision_feature_layer, - ) + return_dict=True, + ).pooler_output image_features = image_features.to(inputs_embeds.device, inputs_embeds.dtype) special_image_mask = self.get_placeholder_mask( input_ids, inputs_embeds=inputs_embeds, image_features=image_features @@ -1354,16 +1352,19 @@ def forward( class AriaForConditionalGeneration(LlavaForConditionalGeneration): _tied_weights_keys = {"lm_head.weight": "model.language_model.embed_tokens.weight"} + @auto_docstring def get_image_features( self, pixel_values: torch.FloatTensor, pixel_mask: torch.FloatTensor | None = None, vision_feature_layer: int = -1, - ): + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: return self.model.get_image_features( pixel_values=pixel_values, pixel_mask=pixel_mask, vision_feature_layer=vision_feature_layer, + **kwargs, ) @can_return_tuple diff --git a/src/transformers/models/audioflamingo3/modeling_audioflamingo3.py b/src/transformers/models/audioflamingo3/modeling_audioflamingo3.py index a8dfe8c0d575..f88a19796f34 100644 --- a/src/transformers/models/audioflamingo3/modeling_audioflamingo3.py +++ b/src/transformers/models/audioflamingo3/modeling_audioflamingo3.py @@ -19,7 +19,6 @@ # See the License for the specific language governing permissions and # limitations under the License. - import math from collections.abc import Callable @@ -32,10 +31,11 @@ from ...masking_utils import create_bidirectional_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer -from ...modeling_outputs import BaseModelOutput, CausalLMOutputWithPast +from ...modeling_outputs import BaseModelOutputWithPooling, CausalLMOutputWithPast from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel from ...processing_utils import Unpack from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, logging +from ...utils.generic import check_model_inputs from ..auto import AutoModel, AutoModelForCausalLM from .configuration_audioflamingo3 import AudioFlamingo3Config, AudioFlamingo3EncoderConfig @@ -280,6 +280,11 @@ class AudioFlamingo3Encoder(AudioFlamingo3PreTrainedModel): input_modalities = "audio" _no_split_modules = ["AudioFlamingo3EncoderLayer"] + _can_record_outputs = { + "hidden_states": AudioFlamingo3EncoderLayer, + "attentions": AudioFlamingo3Attention, + } + def __init__(self, config: AudioFlamingo3EncoderConfig): super().__init__(config) self.dropout = config.dropout @@ -316,13 +321,13 @@ def get_input_embeddings(self) -> nn.Module: def set_input_embeddings(self, value: nn.Module): self.conv1 = value - @can_return_tuple + @check_model_inputs def forward( self, input_features: torch.Tensor, input_features_mask: torch.Tensor | None = None, **kwargs, - ): + ) -> tuple | BaseModelOutputWithPooling: r""" Args: input_features (`torch.FloatTensor` of shape `(batch_size, feature_size, sequence_length)`): @@ -366,7 +371,7 @@ def forward( hidden_states = self.avg_pooler(hidden_states).permute(0, 2, 1) hidden_states = self.layer_norm(hidden_states) - return BaseModelOutput( + return BaseModelOutputWithPooling( last_hidden_state=hidden_states, ) @@ -441,35 +446,40 @@ def set_decoder(self, decoder): def get_decoder(self): return self.language_model.get_decoder() + @can_return_tuple + @auto_docstring( + custom_intro="This method is used to get the audio embeddings from input features (a log mel spectrogram), meaning inferring the audio encoder and the multi-modal projector." + ) def get_audio_features( - self, input_features: torch.FloatTensor, input_features_mask: torch.Tensor - ) -> torch.FloatTensor: - """ - This method is used to get the audio embeddings from input features (a log mel spectrogram), meaning inferring the audio encoder and the multi-modal projector. - Args: - input_features (`torch.FloatTensor`): - Float values of mel features extracted from the raw speech waveform. Raw speech waveform can be - obtained by loading a `.flac` or `.wav` audio file into an array of type `list[float]` or a - `numpy.ndarray`, *e.g.* via the soundfile library (`pip install soundfile`). To prepare the array into - `input_features`, the [`AutoFeatureExtractor`] should be used for extracting the mel features, padding - and conversion into a tensor of type `torch.FloatTensor`. See [`~WhisperFeatureExtractor.__call__`] - input_features_mask (`torch.Tensor` of shape `(batch_size, feature_sequence_length)`): - Mask to avoid performing attention on padded feature indices. - - Returns: - `torch.FloatTensor`: - The audio embeddings. + self, + input_features: torch.FloatTensor, + input_features_mask: torch.Tensor, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + input_features (`torch.FloatTensor`): + Float values of mel features extracted from the raw speech waveform. Raw speech waveform can be + obtained by loading a `.flac` or `.wav` audio file into an array of type `list[float]` or a + `numpy.ndarray`, *e.g.* via the soundfile library (`pip install soundfile`). To prepare the array into + `input_features`, the [`AutoFeatureExtractor`] should be used for extracting the mel features, padding + and conversion into a tensor of type `torch.FloatTensor`. See [`~WhisperFeatureExtractor.__call__`] + input_features_mask (`torch.Tensor` of shape `(batch_size, feature_sequence_length)`): + Mask to avoid performing attention on padded feature indices. """ # Encode audio - encoder_output = self.audio_tower(input_features, input_features_mask=input_features_mask) - audio_embeds = self.multi_modal_projector(encoder_output.last_hidden_state) + audio_output = self.audio_tower( + input_features, input_features_mask=input_features_mask, return_dict=True, **kwargs + ) + audio_embeds = self.multi_modal_projector(audio_output.last_hidden_state) # Mask according to avg pooling (which is after attention blocks) post_lengths = (input_features_mask.sum(-1) - 2) // 2 + 1 valid_mask = torch.arange(audio_embeds.shape[1], device=post_lengths.device)[None, :] < post_lengths[:, None] audio_embeds = audio_embeds[valid_mask.to(audio_embeds.device)] - return audio_embeds + audio_output.pooler_output = audio_embeds + + return audio_output @can_return_tuple @auto_docstring @@ -555,7 +565,7 @@ def forward( inputs_embeds = self.get_input_embeddings()(input_ids) if input_features is not None and input_ids is not None: - audio_embeds = self.get_audio_features(input_features, input_features_mask) + audio_embeds = self.get_audio_features(input_features, input_features_mask, return_dict=True).pooler_output # replace text-audio token placeholders with audio embeddings audio_token_mask = (input_ids == self.config.audio_token_id).unsqueeze(-1) diff --git a/src/transformers/models/audioflamingo3/modular_audioflamingo3.py b/src/transformers/models/audioflamingo3/modular_audioflamingo3.py index 0eaa735ef85c..b846957940cc 100644 --- a/src/transformers/models/audioflamingo3/modular_audioflamingo3.py +++ b/src/transformers/models/audioflamingo3/modular_audioflamingo3.py @@ -13,28 +13,32 @@ # See the License for the specific language governing permissions and # limitations under the License. - import torch from torch import nn from ...activations import ACT2FN from ...cache_utils import Cache from ...masking_utils import create_bidirectional_mask -from ...modeling_outputs import BaseModelOutput, CausalLMOutputWithPast +from ...modeling_outputs import BaseModelOutputWithPooling, CausalLMOutputWithPast from ...processing_utils import Unpack from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, logging +from ...utils.generic import check_model_inputs from ..qwen2_audio.modeling_qwen2_audio import ( Qwen2AudioEncoder, Qwen2AudioPreTrainedModel, ) from ..voxtral.modeling_voxtral import VoxtralForConditionalGeneration, VoxtralMultiModalProjector -from ..whisper.modeling_whisper import WhisperEncoderLayer +from ..whisper.modeling_whisper import WhisperAttention, WhisperEncoderLayer from .configuration_audioflamingo3 import AudioFlamingo3Config logger = logging.get_logger(__name__) +class AudioFlamingo3Attention(WhisperAttention): + pass + + class AudioFlamingo3EncoderLayer(WhisperEncoderLayer): pass @@ -53,13 +57,18 @@ class AudioFlamingo3Encoder(Qwen2AudioEncoder): AudioFlamingo3 encoder: Whisper encoder, average pool (time/2), then LayerNorm. """ - @can_return_tuple + _can_record_outputs = { + "hidden_states": AudioFlamingo3EncoderLayer, + "attentions": AudioFlamingo3Attention, + } + + @check_model_inputs def forward( self, input_features: torch.Tensor, input_features_mask: torch.Tensor | None = None, **kwargs, - ): + ) -> tuple | BaseModelOutputWithPooling: r""" Args: input_features (`torch.FloatTensor` of shape `(batch_size, feature_size, sequence_length)`): @@ -103,7 +112,7 @@ def forward( hidden_states = self.avg_pooler(hidden_states).permute(0, 2, 1) hidden_states = self.layer_norm(hidden_states) - return BaseModelOutput( + return BaseModelOutputWithPooling( last_hidden_state=hidden_states, ) @@ -138,35 +147,40 @@ class AudioFlamingo3ForConditionalGeneration(VoxtralForConditionalGeneration): def __init__(self, config): super().__init__(config) + @can_return_tuple + @auto_docstring( + custom_intro="This method is used to get the audio embeddings from input features (a log mel spectrogram), meaning inferring the audio encoder and the multi-modal projector." + ) def get_audio_features( - self, input_features: torch.FloatTensor, input_features_mask: torch.Tensor - ) -> torch.FloatTensor: - """ - This method is used to get the audio embeddings from input features (a log mel spectrogram), meaning inferring the audio encoder and the multi-modal projector. - Args: - input_features (`torch.FloatTensor`): - Float values of mel features extracted from the raw speech waveform. Raw speech waveform can be - obtained by loading a `.flac` or `.wav` audio file into an array of type `list[float]` or a - `numpy.ndarray`, *e.g.* via the soundfile library (`pip install soundfile`). To prepare the array into - `input_features`, the [`AutoFeatureExtractor`] should be used for extracting the mel features, padding - and conversion into a tensor of type `torch.FloatTensor`. See [`~WhisperFeatureExtractor.__call__`] - input_features_mask (`torch.Tensor` of shape `(batch_size, feature_sequence_length)`): - Mask to avoid performing attention on padded feature indices. - - Returns: - `torch.FloatTensor`: - The audio embeddings. + self, + input_features: torch.FloatTensor, + input_features_mask: torch.Tensor, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + input_features (`torch.FloatTensor`): + Float values of mel features extracted from the raw speech waveform. Raw speech waveform can be + obtained by loading a `.flac` or `.wav` audio file into an array of type `list[float]` or a + `numpy.ndarray`, *e.g.* via the soundfile library (`pip install soundfile`). To prepare the array into + `input_features`, the [`AutoFeatureExtractor`] should be used for extracting the mel features, padding + and conversion into a tensor of type `torch.FloatTensor`. See [`~WhisperFeatureExtractor.__call__`] + input_features_mask (`torch.Tensor` of shape `(batch_size, feature_sequence_length)`): + Mask to avoid performing attention on padded feature indices. """ # Encode audio - encoder_output = self.audio_tower(input_features, input_features_mask=input_features_mask) - audio_embeds = self.multi_modal_projector(encoder_output.last_hidden_state) + audio_output = self.audio_tower( + input_features, input_features_mask=input_features_mask, return_dict=True, **kwargs + ) + audio_embeds = self.multi_modal_projector(audio_output.last_hidden_state) # Mask according to avg pooling (which is after attention blocks) post_lengths = (input_features_mask.sum(-1) - 2) // 2 + 1 valid_mask = torch.arange(audio_embeds.shape[1], device=post_lengths.device)[None, :] < post_lengths[:, None] audio_embeds = audio_embeds[valid_mask.to(audio_embeds.device)] - return audio_embeds + audio_output.pooler_output = audio_embeds + + return audio_output @can_return_tuple @auto_docstring @@ -252,7 +266,7 @@ def forward( inputs_embeds = self.get_input_embeddings()(input_ids) if input_features is not None and input_ids is not None: - audio_embeds = self.get_audio_features(input_features, input_features_mask) + audio_embeds = self.get_audio_features(input_features, input_features_mask, return_dict=True).pooler_output # replace text-audio token placeholders with audio embeddings audio_token_mask = (input_ids == self.config.audio_token_id).unsqueeze(-1) diff --git a/src/transformers/models/aya_vision/modeling_aya_vision.py b/src/transformers/models/aya_vision/modeling_aya_vision.py index fee54aadfce1..91517071cb17 100644 --- a/src/transformers/models/aya_vision/modeling_aya_vision.py +++ b/src/transformers/models/aya_vision/modeling_aya_vision.py @@ -26,10 +26,10 @@ from ...activations import ACT2FN from ...cache_utils import Cache from ...generation import GenerationMixin -from ...modeling_outputs import BaseModelOutputWithPast, ModelOutput +from ...modeling_outputs import BaseModelOutputWithPast, BaseModelOutputWithPooling, ModelOutput from ...modeling_utils import PreTrainedModel from ...processing_utils import Unpack -from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, torch_compilable_check +from ...utils import TransformersKwargs, auto_docstring, torch_compilable_check from ...utils.generic import check_model_inputs from ..auto import AutoModel from .configuration_aya_vision import AyaVisionConfig @@ -179,44 +179,26 @@ def get_input_embeddings(self): def set_input_embeddings(self, value): self.language_model.set_input_embeddings(value) + @check_model_inputs(tie_last_hidden_states=False) + @auto_docstring( + custom_intro="Obtains image last hidden states from the vision tower and apply multimodal projection." + ) def get_image_features( self, pixel_values: torch.FloatTensor, vision_feature_layer: int | list[int] | None = None, vision_feature_select_strategy: str | None = None, - **kwargs, - ): - """ - Obtains image last hidden states from the vision tower and apply multimodal projection. - - Args: - pixel_values (`torch.FloatTensor]` of shape `(batch_size, channels, height, width)`): - The tensors corresponding to the input images. - vision_feature_layer (`Union[int, list[int]]`, *optional*): - The index of the layer to select the vision feature. If multiple indices are provided, - the vision feature of the corresponding indices will be concatenated to form the - vision features. - vision_feature_select_strategy (`str`, *optional*): - The feature selection strategy used to select the vision feature from the vision backbone. - Can be one of `"default"` or `"full"` - Returns: - image_features (`torch.Tensor`): Image feature tensor of shape `(num_images, image_length, embed_dim)`). - """ - vision_feature_layer = ( - vision_feature_layer if vision_feature_layer is not None else self.config.vision_feature_layer - ) - vision_feature_select_strategy = ( - vision_feature_select_strategy - if vision_feature_select_strategy is not None - else self.config.vision_feature_select_strategy - ) - - if vision_feature_select_strategy not in ["default", "full"]: - raise ValueError(f"Unexpected select feature strategy: {self.config.vision_feature_select_strategy}") - + output_hidden_states: bool | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: kwargs = {k: v for k, v in kwargs.items() if v is not None} # this is not memory efficient at all (output_hidden_states=True) will save all the hidden states. - image_outputs = self.vision_tower(pixel_values, output_hidden_states=True, **kwargs) + image_outputs = self.vision_tower( + pixel_values, + output_hidden_states=True, # Ignore arg on purpose + return_dict=True, + **kwargs, + ) # If we have one vision feature layer, return the corresponding hidden states, # otherwise, select the hidden states of each feature layer and concatenate them @@ -231,8 +213,9 @@ def get_image_features( hs_pool = [hs[:, 1:] for hs in hs_pool] selected_image_feature = torch.cat(hs_pool, dim=-1) - image_features = self.multi_modal_projector(selected_image_feature) - return image_features + image_outputs.pooler_output = self.multi_modal_projector(selected_image_feature) + + return image_outputs def get_placeholder_mask( self, input_ids: torch.LongTensor, inputs_embeds: torch.FloatTensor, image_features: torch.FloatTensor @@ -274,15 +257,6 @@ def forward( cache_position: torch.LongTensor | None = None, **kwargs: Unpack[TransformersKwargs], ) -> tuple | AyaVisionModelOutputWithPast: - vision_feature_layer = ( - vision_feature_layer if vision_feature_layer is not None else self.config.vision_feature_layer - ) - vision_feature_select_strategy = ( - vision_feature_select_strategy - if vision_feature_select_strategy is not None - else self.config.vision_feature_select_strategy - ) - if (input_ids is None) ^ (inputs_embeds is not None): raise ValueError("You must specify exactly one of input_ids or inputs_embeds") @@ -294,7 +268,8 @@ def forward( pixel_values=pixel_values, vision_feature_layer=vision_feature_layer, vision_feature_select_strategy=vision_feature_select_strategy, - ) + return_dict=True, + ).pooler_output image_features = image_features.to(inputs_embeds.device, inputs_embeds.dtype) special_image_mask = self.get_placeholder_mask( input_ids, inputs_embeds=inputs_embeds, image_features=image_features @@ -349,13 +324,14 @@ def set_input_embeddings(self, value): def get_output_embeddings(self) -> nn.Module: return self.lm_head + @auto_docstring def get_image_features( self, pixel_values: torch.FloatTensor, vision_feature_layer: int | list[int] | None = None, vision_feature_select_strategy: str | None = None, - **kwargs, - ): + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: return self.model.get_image_features( pixel_values=pixel_values, vision_feature_layer=vision_feature_layer, @@ -363,7 +339,7 @@ def get_image_features( **kwargs, ) - @can_return_tuple + @check_model_inputs(tie_last_hidden_states=False) @auto_docstring def forward( self, @@ -417,15 +393,6 @@ def forward( >>> gen_tokens = model.generate(**inputs, max_new_tokens=300, do_sample=True, temperature=0.3) >>> processor.tokenizer.decode(gen_tokens[0][inputs.input_ids.shape[1]:], skip_special_tokens=True) ```""" - vision_feature_layer = ( - vision_feature_layer if vision_feature_layer is not None else self.config.vision_feature_layer - ) - vision_feature_select_strategy = ( - vision_feature_select_strategy - if vision_feature_select_strategy is not None - else self.config.vision_feature_select_strategy - ) - outputs = self.model( input_ids=input_ids, pixel_values=pixel_values, diff --git a/src/transformers/models/aya_vision/modular_aya_vision.py b/src/transformers/models/aya_vision/modular_aya_vision.py index e0c4b23c4edb..378a826e1f2e 100644 --- a/src/transformers/models/aya_vision/modular_aya_vision.py +++ b/src/transformers/models/aya_vision/modular_aya_vision.py @@ -22,13 +22,13 @@ LlavaModel, LlavaModelOutputWithPast, LlavaPreTrainedModel, - TransformersKwargs, ) from ...activations import ACT2FN from ...cache_utils import Cache +from ...modeling_outputs import BaseModelOutputWithPooling from ...processing_utils import Unpack -from ...utils import auto_docstring, logging +from ...utils import TransformersKwargs, auto_docstring, logging from ...utils.generic import check_model_inputs from .configuration_aya_vision import AyaVisionConfig @@ -104,44 +104,26 @@ class AyaVisionModelOutputWithPast(LlavaModelOutputWithPast): class AyaVisionModel(LlavaModel): # Unlike LLaVA, the model doesn't have to deal with Pixtral-style image states + @check_model_inputs(tie_last_hidden_states=False) + @auto_docstring( + custom_intro="Obtains image last hidden states from the vision tower and apply multimodal projection." + ) def get_image_features( self, pixel_values: torch.FloatTensor, vision_feature_layer: int | list[int] | None = None, vision_feature_select_strategy: str | None = None, - **kwargs, - ): - """ - Obtains image last hidden states from the vision tower and apply multimodal projection. - - Args: - pixel_values (`torch.FloatTensor]` of shape `(batch_size, channels, height, width)`): - The tensors corresponding to the input images. - vision_feature_layer (`Union[int, list[int]]`, *optional*): - The index of the layer to select the vision feature. If multiple indices are provided, - the vision feature of the corresponding indices will be concatenated to form the - vision features. - vision_feature_select_strategy (`str`, *optional*): - The feature selection strategy used to select the vision feature from the vision backbone. - Can be one of `"default"` or `"full"` - Returns: - image_features (`torch.Tensor`): Image feature tensor of shape `(num_images, image_length, embed_dim)`). - """ - vision_feature_layer = ( - vision_feature_layer if vision_feature_layer is not None else self.config.vision_feature_layer - ) - vision_feature_select_strategy = ( - vision_feature_select_strategy - if vision_feature_select_strategy is not None - else self.config.vision_feature_select_strategy - ) - - if vision_feature_select_strategy not in ["default", "full"]: - raise ValueError(f"Unexpected select feature strategy: {self.config.vision_feature_select_strategy}") - + output_hidden_states: bool | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: kwargs = {k: v for k, v in kwargs.items() if v is not None} # this is not memory efficient at all (output_hidden_states=True) will save all the hidden states. - image_outputs = self.vision_tower(pixel_values, output_hidden_states=True, **kwargs) + image_outputs = self.vision_tower( + pixel_values, + output_hidden_states=True, # Ignore arg on purpose + return_dict=True, + **kwargs, + ) # If we have one vision feature layer, return the corresponding hidden states, # otherwise, select the hidden states of each feature layer and concatenate them @@ -156,8 +138,9 @@ def get_image_features( hs_pool = [hs[:, 1:] for hs in hs_pool] selected_image_feature = torch.cat(hs_pool, dim=-1) - image_features = self.multi_modal_projector(selected_image_feature) - return image_features + image_outputs.pooler_output = self.multi_modal_projector(selected_image_feature) + + return image_outputs @check_model_inputs @auto_docstring @@ -175,15 +158,6 @@ def forward( cache_position: torch.LongTensor | None = None, **kwargs: Unpack[TransformersKwargs], ) -> tuple | AyaVisionModelOutputWithPast: - vision_feature_layer = ( - vision_feature_layer if vision_feature_layer is not None else self.config.vision_feature_layer - ) - vision_feature_select_strategy = ( - vision_feature_select_strategy - if vision_feature_select_strategy is not None - else self.config.vision_feature_select_strategy - ) - if (input_ids is None) ^ (inputs_embeds is not None): raise ValueError("You must specify exactly one of input_ids or inputs_embeds") @@ -195,7 +169,8 @@ def forward( pixel_values=pixel_values, vision_feature_layer=vision_feature_layer, vision_feature_select_strategy=vision_feature_select_strategy, - ) + return_dict=True, + ).pooler_output image_features = image_features.to(inputs_embeds.device, inputs_embeds.dtype) special_image_mask = self.get_placeholder_mask( input_ids, inputs_embeds=inputs_embeds, image_features=image_features diff --git a/src/transformers/models/blip/modeling_blip.py b/src/transformers/models/blip/modeling_blip.py index 23079f977145..e79ccdfa7d5e 100644 --- a/src/transformers/models/blip/modeling_blip.py +++ b/src/transformers/models/blip/modeling_blip.py @@ -24,7 +24,11 @@ from ...activations import ACT2FN from ...generation import GenerationMixin from ...modeling_layers import GradientCheckpointingLayer -from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling +from ...modeling_outputs import ( + BaseModelOutput, + BaseModelOutputWithPooling, + BaseModelOutputWithPoolingAndCrossAttentions, +) from ...modeling_utils import PreTrainedModel from ...processing_utils import Unpack from ...utils import ModelOutput, TransformersKwargs, auto_docstring, can_return_tuple, logging, torch_int @@ -558,18 +562,16 @@ def get_input_embeddings(self): def set_input_embeddings(self, value): self.text_model.set_input_embeddings(value) + @can_return_tuple @auto_docstring def get_text_features( self, input_ids: torch.Tensor | None = None, attention_mask: torch.Tensor | None = None, position_ids: torch.Tensor | None = None, - ) -> torch.FloatTensor: + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: r""" - Returns: - text_features (`torch.FloatTensor` of shape `(batch_size, output_dim`): The text embeddings obtained by - applying the projection layer to the pooled output of [`BlipTextModel`]. - Examples: ```python @@ -581,28 +583,27 @@ def get_text_features( >>> inputs = processor(text=["a photo of a cat", "a photo of a dog"], padding=True, return_tensors="pt") >>> text_features = model.get_text_features(**inputs) ```""" - text_outputs = self.text_model( + text_outputs: BaseModelOutputWithPoolingAndCrossAttentions = self.text_model( input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids, + return_dict=True, + **kwargs, ) + pooled_output = text_outputs.pooler_output + text_outputs.pooler_output = self.text_projection(pooled_output) - pooled_output = text_outputs[1] - text_features = self.text_projection(pooled_output) - - return text_features + return text_outputs + @can_return_tuple @auto_docstring def get_image_features( self, pixel_values: torch.FloatTensor | None = None, interpolate_pos_encoding: bool = False, - ) -> torch.FloatTensor: + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: r""" - Returns: - image_features (`torch.FloatTensor` of shape `(batch_size, output_dim`): The image embeddings obtained by - applying the projection layer to the pooled output of [`BlipVisionModel`]. - Examples: ```python @@ -621,15 +622,16 @@ def get_image_features( >>> image_features = model.get_image_features(**inputs) ```""" - vision_outputs = self.vision_model( + vision_outputs: BaseModelOutputWithPooling = self.vision_model( pixel_values=pixel_values, interpolate_pos_encoding=interpolate_pos_encoding, + return_dict=True, + **kwargs, ) + pooled_output = vision_outputs.pooler_output + vision_outputs.pooler_output = self.visual_projection(pooled_output) - pooled_output = vision_outputs[1] # pooled_output - image_features = self.visual_projection(pooled_output) - - return image_features + return vision_outputs @auto_docstring def get_multimodal_features( diff --git a/src/transformers/models/blip_2/modeling_blip_2.py b/src/transformers/models/blip_2/modeling_blip_2.py index e23c96d8dffb..0ec86489224e 100644 --- a/src/transformers/models/blip_2/modeling_blip_2.py +++ b/src/transformers/models/blip_2/modeling_blip_2.py @@ -28,6 +28,7 @@ from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import ( BaseModelOutput, + BaseModelOutputWithPast, BaseModelOutputWithPastAndCrossAttentions, BaseModelOutputWithPooling, BaseModelOutputWithPoolingAndCrossAttentions, @@ -54,6 +55,20 @@ logger = logging.get_logger(__name__) +@dataclass +@auto_docstring +class BaseModelOutputWithVisionQformerOutputs(BaseModelOutputWithPooling): + r""" + vision_outputs (`BaseModelOutputWithPooling`): + Outputs of the vision encoder. + qformer_outputs (`BaseModelOutputWithPoolingAndCrossAttentions`): + Outputs of the Q-Former (Querying Transformer). + """ + + vision_outputs: BaseModelOutputWithPooling | None = None + qformer_outputs: BaseModelOutputWithPoolingAndCrossAttentions | None = None + + @dataclass @auto_docstring( custom_intro=""" @@ -76,9 +91,9 @@ class Blip2ForConditionalGenerationModelOutput(ModelOutput): loss: tuple[torch.FloatTensor] | None = None logits: tuple[torch.FloatTensor] | None = None - vision_outputs: torch.FloatTensor | None = None - qformer_outputs: tuple[torch.FloatTensor] | None = None - language_model_outputs: tuple[torch.FloatTensor] | None = None + vision_outputs: BaseModelOutputWithPooling | None = None + qformer_outputs: BaseModelOutputWithPoolingAndCrossAttentions | None = None + language_model_outputs: CausalLMOutputWithPast | Seq2SeqLMOutput | None = None def to_tuple(self) -> tuple[Any]: return tuple( @@ -465,7 +480,6 @@ def forward( @auto_docstring -# Copied from transformers.models.blip.modeling_blip.BlipVisionModel with Blip->Blip2, BLIP->BLIP_2 class Blip2VisionModel(Blip2PreTrainedModel): main_input_name = "pixel_values" input_modalities = ("image",) @@ -1067,7 +1081,7 @@ def get_encoder(self, modality=None): else: return super().get_encoder(modality=modality) - @filter_out_non_signature_kwargs() + @can_return_tuple @auto_docstring def get_text_features( self, @@ -1076,7 +1090,8 @@ def get_text_features( decoder_input_ids: torch.Tensor | None = None, decoder_attention_mask: torch.Tensor | None = None, labels: torch.Tensor | None = None, - ) -> torch.FloatTensor | CausalLMOutputWithPast: + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: r""" decoder_input_ids (`torch.LongTensor` of shape `(batch_size, target_sequence_length)`, *optional*): Indices of decoder input sequence tokens in the vocabulary. @@ -1095,10 +1110,6 @@ def get_text_features( Default behavior: generate a tensor that ignores pad tokens in `decoder_input_ids`. Causal mask will also be used by default. - Returns: - text_outputs (``torch.FloatTensor`): - The language model's last hidden states. - Examples: ```python >>> import torch @@ -1113,36 +1124,34 @@ def get_text_features( ```""" if self.config.use_decoder_only_language_model: - text_outputs: CausalLMOutputWithPast = self.language_model( + text_outputs: BaseModelOutputWithPast = self.language_model.base_model( input_ids=input_ids, attention_mask=attention_mask, return_dict=True, + **kwargs, ) else: - inputs_embeds = self.language_model.get_input_embeddings()(input_ids) - text_outputs: Seq2SeqLMOutput = self.language_model( - inputs_embeds=inputs_embeds, + text_outputs: BaseModelOutputWithPastAndCrossAttentions = self.language_model.get_encoder()( + input_ids=input_ids, attention_mask=attention_mask, - decoder_input_ids=decoder_input_ids, - decoder_attention_mask=decoder_attention_mask, - labels=labels, return_dict=True, + **kwargs, ) + return BaseModelOutputWithPooling( + last_hidden_state=text_outputs.last_hidden_state, + hidden_states=text_outputs.hidden_states, + attentions=text_outputs.attentions, + ) - return text_outputs.logits - - @filter_out_non_signature_kwargs() + @can_return_tuple @auto_docstring def get_image_features( self, pixel_values: torch.FloatTensor, interpolate_pos_encoding: bool = False, - ) -> torch.FloatTensor | CausalLMOutputWithPast: + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: r""" - Returns: - vision_outputs (`torch.FloatTensor`): - The vision model's last layer pooled logits. - Examples: ```python >>> import torch @@ -1159,14 +1168,12 @@ def get_image_features( >>> with torch.inference_mode(): ... image_outputs = model.get_image_features(**inputs) ```""" - vision_outputs = self.vision_model( + return self.vision_model( pixel_values=pixel_values, interpolate_pos_encoding=interpolate_pos_encoding, - return_dict=True, + **kwargs, ) - return vision_outputs.pooler_output - @filter_out_non_signature_kwargs() @auto_docstring def get_qformer_features( @@ -1608,25 +1615,33 @@ def _preprocess_accelerate(self): if hasattr(self.language_model, "_hf_hook"): self.language_model._hf_hook.io_same_device = True # For `generate` compatibility + @can_return_tuple + @auto_docstring def get_image_features( self, pixel_values: torch.FloatTensor, interpolate_pos_encoding: bool | None = False, - return_dict: bool | None = False, - ): - """ - Encodes images into continuous embeddings that can be forwarded to the language model. - - Args: - pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): - The tensors corresponding to the input images. + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithVisionQformerOutputs: + r""" + pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): + The tensors corresponding to the input images. """ # step 1: forward the images through the vision encoder, # to get image embeddings of shape (batch_size, seq_len, hidden_size) - vision_outputs = self.vision_model( + vision_outputs: BaseModelOutputWithPooling = self.vision_model( pixel_values=pixel_values, interpolate_pos_encoding=interpolate_pos_encoding, return_dict=True, + **kwargs, + ) + vision_outputs = BaseModelOutputWithVisionQformerOutputs( + last_hidden_state=vision_outputs.last_hidden_state, + pooler_output=vision_outputs.pooler_output, + hidden_states=vision_outputs.hidden_states, + attentions=vision_outputs.attentions, + vision_outputs=vision_outputs, + qformer_outputs=None, ) image_embeds = vision_outputs[0] @@ -1634,23 +1649,24 @@ def get_image_features( image_attention_mask = torch.ones(image_embeds.size()[:-1], dtype=torch.long, device=image_embeds.device) query_tokens = self.query_tokens.expand(image_embeds.shape[0], -1, -1) - query_outputs = self.qformer( + qformer_outputs = self.qformer( query_embeds=query_tokens, encoder_hidden_states=image_embeds, encoder_attention_mask=image_attention_mask, return_dict=True, ) - query_output = query_outputs[0] + vision_outputs.qformer_outputs = qformer_outputs + query_output = qformer_outputs[0] # Qformer is kept in fp32, we downcast the output back if needed if query_output.dtype != image_embeds.dtype: query_output = query_output.to(image_embeds.dtype) # step 3: use the language model, conditioned on the query outputs and the prompt - language_model_inputs = self.language_projection(query_output) - if return_dict: - return language_model_inputs, vision_outputs, query_outputs - return language_model_inputs + image_features = self.language_projection(query_output) + vision_outputs.pooler_output = image_features + + return vision_outputs def get_placeholder_mask(self, input_ids: torch.LongTensor, inputs_embeds: torch.FloatTensor): """ @@ -1755,9 +1771,13 @@ def forward( two ```""" - language_model_inputs, vision_outputs, query_outputs = self.get_image_features( + image_features: BaseModelOutputWithVisionQformerOutputs = self.get_image_features( pixel_values, interpolate_pos_encoding=interpolate_pos_encoding, return_dict=True ) + language_model_inputs = image_features.pooler_output + qformer_outputs = image_features.qformer_outputs + vision_outputs = image_features.vision_outputs + if inputs_embeds is None: inputs_embeds = self.get_input_embeddings()(input_ids) @@ -1807,7 +1827,7 @@ def forward( loss=loss, logits=logits, vision_outputs=vision_outputs, - qformer_outputs=query_outputs, + qformer_outputs=qformer_outputs, language_model_outputs=outputs, ) diff --git a/src/transformers/models/blt/modeling_blt.py b/src/transformers/models/blt/modeling_blt.py index d1812a825538..0dea1ef44a67 100644 --- a/src/transformers/models/blt/modeling_blt.py +++ b/src/transformers/models/blt/modeling_blt.py @@ -1226,7 +1226,7 @@ def forward( use_cache: bool | None = None, cache_position: torch.LongTensor | None = None, **kwargs: Unpack[TransformersKwargs], - ) -> BaseModelOutputWithPast: + ) -> tuple | BaseModelOutputWithPast: if (input_ids is None) ^ (inputs_embeds is not None): raise ValueError("You must specify exactly one of input_ids or inputs_embeds") diff --git a/src/transformers/models/blt/modular_blt.py b/src/transformers/models/blt/modular_blt.py index 421c2e286c40..e27a24ee59d7 100644 --- a/src/transformers/models/blt/modular_blt.py +++ b/src/transformers/models/blt/modular_blt.py @@ -950,7 +950,7 @@ def forward( use_cache: bool | None = None, cache_position: torch.LongTensor | None = None, **kwargs: Unpack[TransformersKwargs], - ) -> BaseModelOutputWithPast: + ) -> tuple | BaseModelOutputWithPast: if (input_ids is None) ^ (inputs_embeds is not None): raise ValueError("You must specify exactly one of input_ids or inputs_embeds") diff --git a/src/transformers/models/chameleon/modeling_chameleon.py b/src/transformers/models/chameleon/modeling_chameleon.py index cf40ed0612f0..d4ec10c50bdc 100644 --- a/src/transformers/models/chameleon/modeling_chameleon.py +++ b/src/transformers/models/chameleon/modeling_chameleon.py @@ -14,6 +14,7 @@ """PyTorch Chameleon model.""" from collections.abc import Callable +from dataclasses import dataclass from functools import cached_property from typing import Optional @@ -27,7 +28,7 @@ from ...masking_utils import create_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer -from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast +from ...modeling_outputs import BaseModelOutputWithPast, BaseModelOutputWithPooling, CausalLMOutputWithPast from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel from ...processing_utils import Unpack @@ -38,13 +39,30 @@ logging, torch_compilable_check, ) -from ...utils.generic import maybe_autocast +from ...utils.generic import check_model_inputs, maybe_autocast from .configuration_chameleon import ChameleonConfig, ChameleonVQVAEConfig logger = logging.get_logger(__name__) +@dataclass +@auto_docstring +class ChameleonVQVAEModelOutput(BaseModelOutputWithPooling): + r""" + quantized_last_hidden_state (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): + Quantized last hidden state from the VQ-VAE model. + image_tokens (`torch.FloatTensor` of shape `(batch_size, config.vocab_size`): + Indices of the image tokens predicted by the VQ-VAE model. + embedding_loss (`torch.FloatTensor`): + The embedding loss computed during quantization. + """ + + quantized_last_hidden_state: torch.FloatTensor | None = None + image_tokens: torch.FloatTensor | None = None + embedding_loss: torch.FloatTensor | None = None + + # Copied from transformers.models.llama.modeling_llama.LlamaRMSNorm with Llama->Chameleon class ChameleonRMSNorm(nn.Module): def __init__(self, hidden_size, eps=1e-6): @@ -781,6 +799,10 @@ class ChameleonPreTrainedModel(PreTrainedModel): _can_compile_fullgraph = True _supports_flex_attn = True _supports_attention_backend = True + _can_record_outputs = { + "hidden_states": [ChameleonDecoderLayer, ChameleonSwinDecoderLayer], + "attentions": ChameleonAttention, + } @auto_docstring( @@ -798,6 +820,10 @@ class ChameleonVQVAE(ChameleonPreTrainedModel): "ChameleonVQVAEEncoderAttnBlock", "ChameleonVQVAEEncoderResnetBlock", ] + _can_record_outputs = { + "hidden_states": ChameleonVQVAEEncoderResnetBlock, + "attentions": ChameleonVQVAEEncoderAttnBlock, + } def __init__(self, config: ChameleonVQVAEConfig): super().__init__(config) @@ -809,11 +835,19 @@ def __init__(self, config: ChameleonVQVAEConfig): self.eval() # Chameleon's VQ model is frozen self.post_init() - def encode(self, pixel_values: torch.LongTensor): + @check_model_inputs + def encode( + self, pixel_values: torch.LongTensor, **kwargs: Unpack[TransformersKwargs] + ) -> ChameleonVQVAEModelOutput: hidden_states = self.encoder(pixel_values) - hidden_states = self.quant_conv(hidden_states) - quant, emb_loss, indices = self.quantize(hidden_states) - return quant, emb_loss, indices + conv_hidden_states = self.quant_conv(hidden_states) + quantized_last_hidden_state, emb_loss, indices = self.quantize(conv_hidden_states) + return ChameleonVQVAEModelOutput( + last_hidden_state=hidden_states, + quantized_last_hidden_state=quantized_last_hidden_state, + image_tokens=indices, + embedding_loss=emb_loss, + ) @auto_docstring @@ -848,23 +882,25 @@ def get_image_tokens(self, pixel_values: torch.FloatTensor): The tensors corresponding to the input images. """ batch_size = pixel_values.shape[0] - _, _, image_toks = self.vqmodel.encode(pixel_values) - bpe_toks = self.vocabulary_mapping.convert_img2bpe(image_toks) + vqmodel_outputs: ChameleonVQVAEModelOutput = self.vqmodel.encode(pixel_values, return_dict=True) + bpe_toks = self.vocabulary_mapping.convert_img2bpe(vqmodel_outputs.image_tokens) bpe_toks = bpe_toks.view(batch_size, -1) return bpe_toks - def get_image_features(self, pixel_values: torch.FloatTensor): - """ - Tokenizes images into discrete tokens with VQGAN module and embeds - them with text embeddings layer - - Args: - pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)): - The tensors corresponding to the input images. + @can_return_tuple + @auto_docstring( + custom_intro="Tokenizes images into discrete tokens with VQGAN module and embeds them with text embeddings layer." + ) + def get_image_features( + self, pixel_values: torch.FloatTensor, **kwargs: Unpack[TransformersKwargs] + ) -> tuple | BaseModelOutputWithPooling: + r""" + pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): + The tensors corresponding to the input images. """ - image_tokens = self.get_image_tokens(pixel_values) - vision_embeddings = self.get_input_embeddings()(image_tokens) - return vision_embeddings + vqmodel_outputs: ChameleonVQVAEModelOutput = self.vqmodel.encode(pixel_values, return_dict=True, **kwargs) + vqmodel_outputs.pooler_output = self.get_input_embeddings()(vqmodel_outputs.image_tokens) + return vqmodel_outputs def get_placeholder_mask( self, input_ids: torch.LongTensor, inputs_embeds: torch.FloatTensor, image_features: torch.FloatTensor @@ -926,11 +962,11 @@ def forward( inputs_embeds = self.embed_tokens(input_ids) if pixel_values is not None: - image_embeds = self.get_image_features(pixel_values) + image_features = self.get_image_features(pixel_values, return_dict=True).pooler_output special_image_mask = self.get_placeholder_mask( - input_ids, inputs_embeds=inputs_embeds, image_features=image_embeds + input_ids, inputs_embeds=inputs_embeds, image_features=image_features ) - inputs_embeds = inputs_embeds.masked_scatter(special_image_mask, image_embeds) + inputs_embeds = inputs_embeds.masked_scatter(special_image_mask, image_features) # torch.jit.trace() doesn't support cache objects in the output if use_cache and past_key_values is None and not torch.jit.is_tracing(): @@ -1022,8 +1058,11 @@ def __init__(self, config): def get_image_tokens(self, pixel_values): return self.model.get_image_tokens(pixel_values) - def get_image_features(self, pixel_values): - return self.model.get_image_features(pixel_values) + @auto_docstring + def get_image_features( + self, pixel_values: torch.FloatTensor, **kwargs: Unpack[TransformersKwargs] + ) -> tuple | BaseModelOutputWithPooling: + return self.model.get_image_features(pixel_values, **kwargs) @can_return_tuple @auto_docstring diff --git a/src/transformers/models/chinese_clip/modeling_chinese_clip.py b/src/transformers/models/chinese_clip/modeling_chinese_clip.py index b14942c2a482..e3e45c96d5d5 100644 --- a/src/transformers/models/chinese_clip/modeling_chinese_clip.py +++ b/src/transformers/models/chinese_clip/modeling_chinese_clip.py @@ -30,8 +30,9 @@ BaseModelOutputWithPoolingAndCrossAttentions, ) from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel +from ...processing_utils import Unpack from ...pytorch_utils import apply_chunking_to_forward -from ...utils import ModelOutput, auto_docstring, can_return_tuple, filter_out_non_signature_kwargs, logging, torch_int +from ...utils import ModelOutput, TransformersKwargs, auto_docstring, can_return_tuple, logging, torch_int from .configuration_chinese_clip import ChineseCLIPConfig, ChineseCLIPTextConfig, ChineseCLIPVisionConfig @@ -1000,7 +1001,7 @@ def __init__(self, config: ChineseCLIPConfig): # Initialize weights and apply final processing self.post_init() - @filter_out_non_signature_kwargs() + @can_return_tuple @auto_docstring def get_text_features( self, @@ -1008,12 +1009,9 @@ def get_text_features( attention_mask: torch.Tensor | None = None, token_type_ids: torch.Tensor | None = None, position_ids: torch.Tensor | None = None, - ) -> torch.FloatTensor: + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: r""" - Returns: - text_features (`torch.FloatTensor` of shape `(batch_size, output_dim`): The text embeddings obtained by - applying the projection layer to the final [CLS] hidden state of Text-Transformer. - Examples: ```python @@ -1028,30 +1026,28 @@ def get_text_features( ... text_features = model.get_text_features(**inputs) >>> text_features = text_features / text_features.norm(p=2, dim=-1, keepdim=True) ```""" - text_outputs = self.text_model( + text_outputs: BaseModelOutputWithPooling = self.text_model( input_ids=input_ids, attention_mask=attention_mask, token_type_ids=token_type_ids, position_ids=position_ids, + return_dict=True, + **kwargs, ) + pooled_output = text_outputs.last_hidden_state[:, 0, :] + text_outputs.pooler_output = self.text_projection(pooled_output) - pooled_output = text_outputs[0][:, 0, :] - text_features = self.text_projection(pooled_output) - - return text_features + return text_outputs - @filter_out_non_signature_kwargs() + @can_return_tuple @auto_docstring def get_image_features( self, pixel_values: torch.FloatTensor, interpolate_pos_encoding: bool = False, - ) -> torch.FloatTensor: + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: r""" - Returns: - image_features (`torch.FloatTensor` of shape `(batch_size, output_dim`): The image embeddings obtained by - applying the projection layer to the final [CLS] hidden state of Vision-Transformer. - Examples: ```python @@ -1074,12 +1070,13 @@ def get_image_features( vision_outputs: BaseModelOutputWithPooling = self.vision_model( pixel_values=pixel_values, interpolate_pos_encoding=interpolate_pos_encoding, + return_dict=True, + **kwargs, ) - pooled_output = vision_outputs.pooler_output - image_features = self.visual_projection(pooled_output) + vision_outputs.pooler_output = self.visual_projection(pooled_output) - return image_features + return vision_outputs @can_return_tuple @auto_docstring diff --git a/src/transformers/models/clap/modeling_clap.py b/src/transformers/models/clap/modeling_clap.py index 1fbade060c91..64223c23e8c1 100644 --- a/src/transformers/models/clap/modeling_clap.py +++ b/src/transformers/models/clap/modeling_clap.py @@ -32,8 +32,9 @@ BaseModelOutputWithPoolingAndCrossAttentions, ) from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel +from ...processing_utils import Unpack from ...pytorch_utils import apply_chunking_to_forward, meshgrid -from ...utils import ModelOutput, auto_docstring, can_return_tuple, filter_out_non_signature_kwargs, logging, torch_int +from ...utils import ModelOutput, TransformersKwargs, auto_docstring, can_return_tuple, logging, torch_int from .configuration_clap import ClapAudioConfig, ClapConfig, ClapTextConfig @@ -1552,19 +1553,16 @@ def __init__(self, config: ClapConfig): # Initialize weights and apply final processing self.post_init() - @filter_out_non_signature_kwargs() + @can_return_tuple @auto_docstring def get_text_features( self, input_ids: torch.Tensor, attention_mask: torch.Tensor | None = None, position_ids: torch.Tensor | None = None, - ) -> torch.FloatTensor: + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: r""" - Returns: - text_features (`torch.FloatTensor` of shape `(batch_size, output_dim`): The text embeddings obtained by - applying the projection layer to the pooled output of [`ClapTextModel`]. - Examples: ```python @@ -1579,30 +1577,31 @@ def get_text_features( ... text_features = model.get_text_features(**inputs) ```""" text_outputs: BaseModelOutputWithPooling = self.text_model( - input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + return_dict=True, + **kwargs, ) text_features = self.text_projection(text_outputs.pooler_output) - text_features = F.normalize(text_features, dim=-1) + text_outputs.pooler_output = F.normalize(text_features, dim=-1) - return text_features + return text_outputs - @filter_out_non_signature_kwargs() + @can_return_tuple @auto_docstring def get_audio_features( self, input_features: torch.Tensor, is_longer: torch.Tensor | None = None, attention_mask: torch.Tensor | None = None, - ) -> torch.FloatTensor: + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: r""" is_longer (`torch.FloatTensor`, of shape `(batch_size, 1)`, *optional*): Whether the audio clip is longer than `max_length`. If `True`, a feature fusion will be enabled to enhance the features. - Returns: - audio_features (`torch.FloatTensor` of shape `(batch_size, output_dim`): The audio embeddings obtained by - applying the projection layer to the pooled output of [`ClapAudioModel`]. - Examples: ```python @@ -1618,12 +1617,12 @@ def get_audio_features( ... audio_features = model.get_audio_features(**inputs) ```""" audio_outputs: BaseModelOutputWithPooling = self.audio_model( - input_features=input_features, is_longer=is_longer + input_features=input_features, is_longer=is_longer, return_dict=True, **kwargs ) audio_features = self.audio_projection(audio_outputs.pooler_output) - audio_features = F.normalize(audio_features, dim=-1) + audio_outputs.pooler_output = F.normalize(audio_features, dim=-1) - return audio_features + return audio_outputs @can_return_tuple @auto_docstring diff --git a/src/transformers/models/clip/modeling_clip.py b/src/transformers/models/clip/modeling_clip.py index 1154a3b47500..54e02843c2de 100644 --- a/src/transformers/models/clip/modeling_clip.py +++ b/src/transformers/models/clip/modeling_clip.py @@ -32,7 +32,6 @@ TransformersKwargs, auto_docstring, can_return_tuple, - filter_out_non_signature_kwargs, logging, torch_int, ) @@ -782,19 +781,16 @@ def __init__(self, config: CLIPConfig): # Initialize weights and apply final processing self.post_init() - @filter_out_non_signature_kwargs() + @can_return_tuple @auto_docstring def get_text_features( self, input_ids: torch.Tensor, attention_mask: torch.Tensor | None = None, position_ids: torch.Tensor | None = None, - ) -> torch.FloatTensor: + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: r""" - Returns: - text_features (`torch.FloatTensor` of shape `(batch_size, output_dim`): The text embeddings obtained by - applying the projection layer to the pooled output of [`CLIPTextModel`]. - Examples: ```python @@ -813,24 +809,23 @@ def get_text_features( input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids, + return_dict=True, + **kwargs, ) pooled_output = text_outputs.pooler_output - text_features = self.text_projection(pooled_output) + text_outputs.pooler_output = self.text_projection(pooled_output) - return text_features + return text_outputs - @filter_out_non_signature_kwargs() + @can_return_tuple @auto_docstring def get_image_features( self, pixel_values: torch.FloatTensor, interpolate_pos_encoding: bool = False, - ) -> torch.FloatTensor: + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: r""" - Returns: - image_features (`torch.FloatTensor` of shape `(batch_size, output_dim`): The image embeddings obtained by - applying the projection layer to the pooled output of [`CLIPVisionModel`]. - Examples: ```python @@ -852,11 +847,13 @@ def get_image_features( vision_outputs: BaseModelOutputWithPooling = self.vision_model( pixel_values=pixel_values, interpolate_pos_encoding=interpolate_pos_encoding, + return_dict=True, + **kwargs, ) pooled_output = vision_outputs.pooler_output - image_features = self.visual_projection(pooled_output) + vision_outputs.pooler_output = self.visual_projection(pooled_output) - return image_features + return vision_outputs @can_return_tuple @auto_docstring diff --git a/src/transformers/models/clipseg/modeling_clipseg.py b/src/transformers/models/clipseg/modeling_clipseg.py index 6d2af9299fdc..818ebab02f88 100644 --- a/src/transformers/models/clipseg/modeling_clipseg.py +++ b/src/transformers/models/clipseg/modeling_clipseg.py @@ -28,7 +28,8 @@ from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel -from ...utils import ModelOutput, auto_docstring, can_return_tuple, filter_out_non_signature_kwargs, logging, torch_int +from ...processing_utils import Unpack +from ...utils import ModelOutput, TransformersKwargs, auto_docstring, can_return_tuple, logging, torch_int from ...utils.generic import is_flash_attention_requested from .configuration_clipseg import CLIPSegConfig, CLIPSegTextConfig, CLIPSegVisionConfig @@ -850,19 +851,16 @@ def __init__(self, config: CLIPSegConfig): # Initialize weights and apply final processing self.post_init() - @filter_out_non_signature_kwargs() + @can_return_tuple @auto_docstring def get_text_features( self, input_ids: torch.Tensor, attention_mask: torch.Tensor | None = None, position_ids: torch.Tensor | None = None, - ) -> torch.FloatTensor: + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: r""" - Returns: - text_features (`torch.FloatTensor` of shape `(batch_size, output_dim`): The text embeddings obtained by - applying the projection layer to the pooled output of [`CLIPSegTextModel`]. - Examples: ```python @@ -880,24 +878,23 @@ def get_text_features( input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids, + return_dict=True, + **kwargs, ) pooled_output = text_outputs.pooler_output - text_features = self.text_projection(pooled_output) + text_outputs.pooler_output = self.text_projection(pooled_output) - return text_features + return text_outputs - @filter_out_non_signature_kwargs() + @can_return_tuple @auto_docstring def get_image_features( self, pixel_values: torch.FloatTensor, interpolate_pos_encoding: bool = True, - ) -> torch.FloatTensor: + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: r""" - Returns: - image_features (`torch.FloatTensor` of shape `(batch_size, output_dim`): The image embeddings obtained by - applying the projection layer to the pooled output of [`CLIPSegVisionModel`]. - Examples: ```python @@ -919,11 +916,13 @@ def get_image_features( vision_outputs: BaseModelOutputWithPooling = self.vision_model( pixel_values=pixel_values, interpolate_pos_encoding=interpolate_pos_encoding, + return_dict=True, + **kwargs, ) pooled_output = vision_outputs.pooler_output - image_features = self.visual_projection(pooled_output) + vision_outputs.pooler_output = self.visual_projection(pooled_output) - return image_features + return vision_outputs @auto_docstring def forward( @@ -1219,13 +1218,13 @@ def get_conditional_embeddings( with torch.no_grad(): conditional_embeddings = self.clip.get_text_features( input_ids, attention_mask=attention_mask, position_ids=position_ids - ) + ).pooler_output elif conditional_pixel_values is not None: # compute conditional embeddings from images if len(conditional_pixel_values) != batch_size: raise ValueError("Make sure to pass as many prompt images as there are query images") with torch.no_grad(): - conditional_embeddings = self.clip.get_image_features(conditional_pixel_values) + conditional_embeddings = self.clip.get_image_features(conditional_pixel_values).pooler_output else: raise ValueError( "Invalid conditional, should be either provided as `input_ids` or `conditional_pixel_values`" diff --git a/src/transformers/models/clvp/modeling_clvp.py b/src/transformers/models/clvp/modeling_clvp.py index ac637c92f163..1a9fcccea86d 100644 --- a/src/transformers/models/clvp/modeling_clvp.py +++ b/src/transformers/models/clvp/modeling_clvp.py @@ -35,10 +35,13 @@ CausalLMOutputWithCrossAttentions, ) from ...modeling_utils import PreTrainedModel +from ...processing_utils import Unpack from ...pytorch_utils import Conv1D, isin_mps_friendly from ...utils import ( ModelOutput, + TransformersKwargs, auto_docstring, + can_return_tuple, logging, ) from .configuration_clvp import ( @@ -1489,36 +1492,23 @@ def fix_speech_decoder_output(self, speech_ids: torch.LongTensor) -> torch.LongT return speech_ids + @can_return_tuple + @auto_docstring( + custom_intro=""" + This method can be used to extract text_embeds from a text. The text embeddings obtained by applying the + projection layer to the pooled output of the CLVP text encoder model. + """ + ) def get_text_features( self, input_ids: torch.LongTensor | None = None, text_encoder_inputs_embeds: torch.FloatTensor | None = None, attention_mask: torch.LongTensor | None = None, - ) -> torch.FloatTensor: + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | ClvpEncoderOutput: r""" - This method can be used to extract text_embeds from a text. The text embeddings obtained by applying the - projection layer to the pooled output of the CLVP text encoder model. - - Args: - input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): - Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you - provide it. - - [What are input IDs?](../glossary#input-ids) - text_encoder_inputs_embeds (`torch.FloatTensor`, *optional*): - inputs_embeds for the text encoder model passed in place of `input_ids`. - attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): - Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`: - - - 1 for tokens that are **not masked**, - - 0 for tokens that are **masked**. - - [What are attention masks?](../glossary#attention-mask) - - Returns: - `torch.FloatTensor` of shape `(batch_size, output_dim)`: - The text embeddings obtained by applying the projection layer to the pooled output of the CLVP Text - Model. + text_encoder_inputs_embeds (`torch.FloatTensor`, *optional*): + inputs_embeds for the text encoder model passed in place of `input_ids`. Examples: @@ -1537,15 +1527,14 @@ def get_text_features( >>> text_embeds = model.get_text_features(input_ids=processor_output["input_ids"]) ``` """ - - outputs = self.text_encoder_model( + return self.text_encoder_model( input_ids=input_ids, inputs_embeds=text_encoder_inputs_embeds, attention_mask=attention_mask, + return_dict=True, + **kwargs, ) - return outputs[0] - def get_speech_features( self, speech_ids: torch.LongTensor | None = None, diff --git a/src/transformers/models/cohere2_vision/modeling_cohere2_vision.py b/src/transformers/models/cohere2_vision/modeling_cohere2_vision.py index 98b41e94bc1c..0123d7064a41 100644 --- a/src/transformers/models/cohere2_vision/modeling_cohere2_vision.py +++ b/src/transformers/models/cohere2_vision/modeling_cohere2_vision.py @@ -25,10 +25,10 @@ from ...cache_utils import Cache from ...generation import GenerationMixin from ...modeling_flash_attention_utils import FlashAttentionKwargs -from ...modeling_outputs import BaseModelOutputWithPast, ModelOutput +from ...modeling_outputs import BaseModelOutputWithPast, BaseModelOutputWithPooling, ModelOutput from ...modeling_utils import PreTrainedModel from ...processing_utils import Unpack -from ...utils import TransformersKwargs, auto_docstring, torch_compilable_check +from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, torch_compilable_check from ...utils.generic import check_model_inputs from ..auto import AutoModel from .configuration_cohere2_vision import Cohere2VisionConfig @@ -165,22 +165,18 @@ def get_input_embeddings(self): def set_input_embeddings(self, value): self.language_model.set_input_embeddings(value) - def get_image_features(self, pixel_values: torch.FloatTensor): - """ - Obtains image last hidden states from the vision tower and apply multimodal projection. - - Args: - pixel_values (`torch.FloatTensor]` of shape `(batch_size, num_patches, channels, height, width)`) - The tensors corresponding to the input images. - Returns: - image_features (List[`torch.Tensor`]): List of image feature tensor, each contains all the visual feature of all patches - and are of shape `(num_patches, image_length, embed_dim)`). - """ + @can_return_tuple + @auto_docstring( + custom_intro="Obtains image last hidden states from the vision tower and apply multimodal projection." + ) + def get_image_features( + self, pixel_values: torch.FloatTensor, **kwargs: Unpack[TransformersKwargs] + ) -> tuple | BaseModelOutputWithPooling: + image_outputs = self.vision_tower(pixel_values, return_dict=True, **kwargs) + selected_image_feature = image_outputs.last_hidden_state + image_outputs.pooler_output = self.multi_modal_projector(selected_image_feature) - image_features = self.vision_tower(pixel_values, output_hidden_states=True) - selected_image_feature = image_features.last_hidden_state - image_features = self.multi_modal_projector(selected_image_feature) - return image_features + return image_outputs def get_placeholder_mask( self, input_ids: torch.LongTensor, inputs_embeds: torch.FloatTensor, image_features: torch.FloatTensor @@ -227,7 +223,7 @@ def forward( inputs_embeds = self.get_input_embeddings()(input_ids) if pixel_values is not None: - image_features = self.get_image_features(pixel_values) + image_features = self.get_image_features(pixel_values, return_dict=True).pooler_output image_features = image_features.to(inputs_embeds.device, inputs_embeds.dtype) special_image_mask = self.get_placeholder_mask( input_ids, inputs_embeds=inputs_embeds, image_features=image_features @@ -277,8 +273,11 @@ def set_input_embeddings(self, value): def get_output_embeddings(self) -> nn.Module: return self.lm_head - def get_image_features(self, pixel_values: torch.FloatTensor): - return self.model.get_image_features(pixel_values=pixel_values) + @auto_docstring + def get_image_features( + self, pixel_values: torch.FloatTensor, **kwargs: Unpack[TransformersKwargs] + ) -> tuple | BaseModelOutputWithPooling: + return self.model.get_image_features(pixel_values=pixel_values, **kwargs) @check_model_inputs @auto_docstring diff --git a/src/transformers/models/cohere2_vision/modular_cohere2_vision.py b/src/transformers/models/cohere2_vision/modular_cohere2_vision.py index ba1ca307286a..ffe43efe3b80 100644 --- a/src/transformers/models/cohere2_vision/modular_cohere2_vision.py +++ b/src/transformers/models/cohere2_vision/modular_cohere2_vision.py @@ -32,8 +32,9 @@ from ...image_processing_utils import BatchFeature from ...image_utils import ImageInput from ...modeling_flash_attention_utils import FlashAttentionKwargs +from ...modeling_outputs import BaseModelOutputWithPooling from ...processing_utils import ImagesKwargs, Unpack -from ...utils import TransformersKwargs, auto_docstring, logging +from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, logging from ...utils.generic import check_model_inputs from .configuration_cohere2_vision import Cohere2VisionConfig @@ -95,22 +96,18 @@ class Cohere2VisionPreTrainedModel(AyaVisionPreTrainedModel): class Cohere2VisionModel(AyaVisionModel): _checkpoint_conversion_mapping = {} - def get_image_features(self, pixel_values: torch.FloatTensor): - """ - Obtains image last hidden states from the vision tower and apply multimodal projection. - - Args: - pixel_values (`torch.FloatTensor]` of shape `(batch_size, num_patches, channels, height, width)`) - The tensors corresponding to the input images. - Returns: - image_features (List[`torch.Tensor`]): List of image feature tensor, each contains all the visual feature of all patches - and are of shape `(num_patches, image_length, embed_dim)`). - """ - - image_features = self.vision_tower(pixel_values, output_hidden_states=True) - selected_image_feature = image_features.last_hidden_state - image_features = self.multi_modal_projector(selected_image_feature) - return image_features + @can_return_tuple + @auto_docstring( + custom_intro="Obtains image last hidden states from the vision tower and apply multimodal projection." + ) + def get_image_features( + self, pixel_values: torch.FloatTensor, **kwargs: Unpack[TransformersKwargs] + ) -> tuple | BaseModelOutputWithPooling: + image_outputs = self.vision_tower(pixel_values, return_dict=True, **kwargs) + selected_image_feature = image_outputs.last_hidden_state + image_outputs.pooler_output = self.multi_modal_projector(selected_image_feature) + + return image_outputs @check_model_inputs @auto_docstring @@ -133,7 +130,7 @@ def forward( inputs_embeds = self.get_input_embeddings()(input_ids) if pixel_values is not None: - image_features = self.get_image_features(pixel_values) + image_features = self.get_image_features(pixel_values, return_dict=True).pooler_output image_features = image_features.to(inputs_embeds.device, inputs_embeds.dtype) special_image_mask = self.get_placeholder_mask( input_ids, inputs_embeds=inputs_embeds, image_features=image_features @@ -162,8 +159,11 @@ def forward( class Cohere2VisionForConditionalGeneration(AyaVisionForConditionalGeneration): _checkpoint_conversion_mapping = {} - def get_image_features(self, pixel_values: torch.FloatTensor): - return self.model.get_image_features(pixel_values=pixel_values) + @auto_docstring + def get_image_features( + self, pixel_values: torch.FloatTensor, **kwargs: Unpack[TransformersKwargs] + ) -> tuple | BaseModelOutputWithPooling: + return self.model.get_image_features(pixel_values=pixel_values, **kwargs) @check_model_inputs @auto_docstring diff --git a/src/transformers/models/colqwen2/modeling_colqwen2.py b/src/transformers/models/colqwen2/modeling_colqwen2.py index 5f8194da4cdc..931d6d5b2582 100644 --- a/src/transformers/models/colqwen2/modeling_colqwen2.py +++ b/src/transformers/models/colqwen2/modeling_colqwen2.py @@ -172,7 +172,9 @@ def forward( inputs_embeds = self.vlm.get_input_embeddings()(input_ids) if pixel_values is not None: - image_embeds = self.vlm.model.visual(pixel_values, grid_thw=image_grid_thw) + image_embeds = self.vlm.model.visual( + pixel_values, grid_thw=image_grid_thw, return_dict=True + ).pooler_output image_mask = ( (input_ids == self.config.vlm_config.image_token_id).unsqueeze(-1).expand_as(inputs_embeds) ) diff --git a/src/transformers/models/colqwen2/modular_colqwen2.py b/src/transformers/models/colqwen2/modular_colqwen2.py index f66658e1e82b..4d135a1046e9 100644 --- a/src/transformers/models/colqwen2/modular_colqwen2.py +++ b/src/transformers/models/colqwen2/modular_colqwen2.py @@ -316,7 +316,9 @@ def forward( inputs_embeds = self.vlm.get_input_embeddings()(input_ids) if pixel_values is not None: - image_embeds = self.vlm.model.visual(pixel_values, grid_thw=image_grid_thw) + image_embeds = self.vlm.model.visual( + pixel_values, grid_thw=image_grid_thw, return_dict=True + ).pooler_output image_mask = ( (input_ids == self.config.vlm_config.image_token_id).unsqueeze(-1).expand_as(inputs_embeds) ) diff --git a/src/transformers/models/deepseek_vl/modeling_deepseek_vl.py b/src/transformers/models/deepseek_vl/modeling_deepseek_vl.py index c8d398a0f1fc..0113c1e5235c 100644 --- a/src/transformers/models/deepseek_vl/modeling_deepseek_vl.py +++ b/src/transformers/models/deepseek_vl/modeling_deepseek_vl.py @@ -26,7 +26,7 @@ from ...cache_utils import Cache from ...generation import GenerationMixin -from ...modeling_outputs import ModelOutput +from ...modeling_outputs import BaseModelOutputWithPooling, ModelOutput from ...modeling_utils import PreTrainedModel from ...processing_utils import Unpack from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, torch_compilable_check @@ -153,10 +153,15 @@ def get_input_embeddings(self): def set_input_embeddings(self, value): self.language_model.set_input_embeddings(value) - def get_image_features(self, pixel_values): - image_embeds = self.vision_model(pixel_values) - image_embeds = self.aligner(image_embeds.last_hidden_state) - return image_embeds + @can_return_tuple + @auto_docstring + def get_image_features( + self, pixel_values: torch.FloatTensor, **kwargs: Unpack[TransformersKwargs] + ) -> tuple | BaseModelOutputWithPooling: + vision_outputs = self.vision_model(pixel_values, return_dict=True, **kwargs) + vision_outputs.pooler_output = self.aligner(vision_outputs.last_hidden_state) + + return vision_outputs def get_placeholder_mask( self, input_ids: torch.LongTensor, inputs_embeds: torch.FloatTensor, image_features: torch.FloatTensor @@ -205,7 +210,7 @@ def forward( inputs_embeds = self.get_input_embeddings()(input_ids) if pixel_values is not None: - image_embeds = self.get_image_features(pixel_values) + image_embeds = self.get_image_features(pixel_values, return_dict=True).pooler_output image_features = image_embeds.reshape(-1, inputs_embeds.shape[-1]) image_features = image_features.to(inputs_embeds.device, inputs_embeds.dtype) image_attention_mask = self.get_placeholder_mask( diff --git a/src/transformers/models/deepseek_vl_hybrid/modeling_deepseek_vl_hybrid.py b/src/transformers/models/deepseek_vl_hybrid/modeling_deepseek_vl_hybrid.py index 56c7cf288d7c..15f653d7d5c0 100644 --- a/src/transformers/models/deepseek_vl_hybrid/modeling_deepseek_vl_hybrid.py +++ b/src/transformers/models/deepseek_vl_hybrid/modeling_deepseek_vl_hybrid.py @@ -18,7 +18,6 @@ # See the License for the specific language governing permissions and # limitations under the License. - from dataclasses import dataclass import torch @@ -27,7 +26,7 @@ from ... import initialization as init from ...cache_utils import Cache from ...generation import GenerationMixin -from ...modeling_outputs import ModelOutput +from ...modeling_outputs import BaseModelOutputWithPooling, ModelOutput from ...modeling_utils import PreTrainedModel from ...processing_utils import Unpack from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, torch_compilable_check @@ -35,6 +34,30 @@ from .configuration_deepseek_vl_hybrid import DeepseekVLHybridConfig +@dataclass +@auto_docstring +class BaseModelOutputWithHighResVisionEncodings(BaseModelOutputWithPooling): + r""" + high_res_vision_last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`): + Sequence of hidden-states at the output of the last layer of the high resolution vision model. + high_res_vision_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the high resolution vision model has an embedding layer, + + one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the high resolution vision model at the output of each layer plus the optional initial embedding outputs. + high_res_vision_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)` from the high resolution vision model. + + Attentions weights after the attention softmax, used to compute the weighted average in the self-attention + heads. + """ + + high_res_vision_last_hidden_state: torch.FloatTensor | None = None + high_res_vision_hidden_states: tuple[torch.FloatTensor] | None = None + high_res_vision_attentions: tuple[torch.FloatTensor] | None = None + + @dataclass @auto_docstring( custom_intro=""" @@ -269,11 +292,27 @@ def get_input_embeddings(self): def set_input_embeddings(self, value): self.language_model.set_input_embeddings(value) - def get_image_features(self, pixel_values, high_res_pixel_values): - vision_encodings = self.get_low_res_image_features(pixel_values) - high_res_vision_encodings = self.get_high_res_image_features(high_res_pixel_values) - images_embeds = self.aligner(vision_encodings, high_res_vision_encodings) - return images_embeds + @can_return_tuple + @auto_docstring(custom_args=DEEPSEEK_VL_COMMON_CUSTOM_ARGS) + def get_image_features( + self, + pixel_values: torch.FloatTensor, + high_res_pixel_values: torch.FloatTensor, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithHighResVisionEncodings: + low_res_outputs = self.get_low_res_image_features(pixel_values, **kwargs) + high_res_outputs = self.get_high_res_image_features(high_res_pixel_values, **kwargs) + image_features = self.aligner(low_res_outputs.last_hidden_state, high_res_outputs.last_hidden_state) + + return BaseModelOutputWithHighResVisionEncodings( + last_hidden_state=low_res_outputs.last_hidden_state, + pooler_output=image_features, + hidden_states=low_res_outputs.hidden_states, + attentions=low_res_outputs.attentions, + high_res_vision_last_hidden_state=high_res_outputs.last_hidden_state, + high_res_vision_hidden_states=high_res_outputs.hidden_states, + high_res_vision_attentions=high_res_outputs.attentions, + ) def get_placeholder_mask( self, input_ids: torch.LongTensor, inputs_embeds: torch.FloatTensor, image_features: torch.FloatTensor @@ -336,7 +375,7 @@ def forward( image_attention_mask = input_ids == self.config.image_token_id image_attention_mask = image_attention_mask.unsqueeze(-1).expand_as(inputs_embeds).to(inputs_embeds.device) - image_embeds = self.get_image_features(pixel_values, high_res_pixel_values) + image_embeds = self.get_image_features(pixel_values, high_res_pixel_values, return_dict=True).pooler_output image_features = image_embeds.reshape(-1, inputs_embeds.shape[-1]) image_features = image_features.to(inputs_embeds.device, inputs_embeds.dtype) inputs_embeds = inputs_embeds.masked_scatter(image_attention_mask, image_features) @@ -360,21 +399,25 @@ def forward( image_hidden_states=image_embeds if pixel_values is not None else None, ) - def get_low_res_image_features(self, pixel_values): - output = self.vision_model(pixel_values) - output = output[0] - return output + def get_low_res_image_features(self, pixel_values: torch.FloatTensor, **kwargs: Unpack[TransformersKwargs]): + return self.vision_model(pixel_values, return_dict=True, **kwargs) - def get_high_res_image_features(self, pixel_values): - output = self.high_res_vision_model( + def get_high_res_image_features( + self, + pixel_values: torch.FloatTensor, + output_hidden_states: bool | None = None, + **kwargs: Unpack[TransformersKwargs], + ): + high_res_outputs = self.high_res_vision_model( pixel_values=pixel_values, - output_hidden_states=True, + output_hidden_states=True, # Ignore arg on purpose return_dict=True, + **kwargs, ) - last_hidden_state = output.last_hidden_state + last_hidden_state = high_res_outputs.last_hidden_state last_hidden_state = self.high_res_vision_proj(last_hidden_state) - hidden_states = output.hidden_states + hidden_states = high_res_outputs.hidden_states global_hidden_state = hidden_states[self.global_attn_index + 1] # +1 for embedding layer global_hidden_state = self.high_res_vision_neck(global_hidden_state) global_hidden_state = self.high_res_vision_proj(global_hidden_state) @@ -384,8 +427,9 @@ def get_high_res_image_features(self, pixel_values): # batch_size, hidden_size, height, width -> batch_size, seq_len, hidden_size output = output.permute(0, 2, 3, 1) output = output.reshape(output.shape[0], -1, output.shape[-1]) + high_res_outputs.last_hidden_state = output - return output + return high_res_outputs class DeepseekVLHybridForConditionalGeneration(DeepseekVLHybridPreTrainedModel, GenerationMixin): diff --git a/src/transformers/models/deepseek_vl_hybrid/modular_deepseek_vl_hybrid.py b/src/transformers/models/deepseek_vl_hybrid/modular_deepseek_vl_hybrid.py index 13a02b6cec3a..8488833e6aaf 100644 --- a/src/transformers/models/deepseek_vl_hybrid/modular_deepseek_vl_hybrid.py +++ b/src/transformers/models/deepseek_vl_hybrid/modular_deepseek_vl_hybrid.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +from dataclasses import dataclass from typing import Optional, Union import torch @@ -43,6 +44,7 @@ valid_images, validate_preprocess_arguments, ) +from ...modeling_outputs import BaseModelOutputWithPooling from ...processing_utils import ImagesKwargs, Unpack from ...tokenization_utils_base import ( PreTokenizedInput, @@ -144,6 +146,30 @@ def __init__( ) +@dataclass +@auto_docstring +class BaseModelOutputWithHighResVisionEncodings(BaseModelOutputWithPooling): + r""" + high_res_vision_last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`): + Sequence of hidden-states at the output of the last layer of the high resolution vision model. + high_res_vision_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): + Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the high resolution vision model has an embedding layer, + + one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the high resolution vision model at the output of each layer plus the optional initial embedding outputs. + high_res_vision_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)` from the high resolution vision model. + + Attentions weights after the attention softmax, used to compute the weighted average in the self-attention + heads. + """ + + high_res_vision_last_hidden_state: torch.FloatTensor | None = None + high_res_vision_hidden_states: tuple[torch.FloatTensor] | None = None + high_res_vision_attentions: tuple[torch.FloatTensor] | None = None + + class DeepseekVLHybridBaseModelOutputWithPast(IdeficsBaseModelOutputWithPast): pass @@ -249,21 +275,25 @@ def __init__(self, config): super().__init__(config) - def get_low_res_image_features(self, pixel_values): - output = self.vision_model(pixel_values) - output = output[0] - return output + def get_low_res_image_features(self, pixel_values: torch.FloatTensor, **kwargs: Unpack[TransformersKwargs]): + return self.vision_model(pixel_values, return_dict=True, **kwargs) - def get_high_res_image_features(self, pixel_values): - output = self.high_res_vision_model( + def get_high_res_image_features( + self, + pixel_values: torch.FloatTensor, + output_hidden_states: bool | None = None, + **kwargs: Unpack[TransformersKwargs], + ): + high_res_outputs = self.high_res_vision_model( pixel_values=pixel_values, - output_hidden_states=True, + output_hidden_states=True, # Ignore arg on purpose return_dict=True, + **kwargs, ) - last_hidden_state = output.last_hidden_state + last_hidden_state = high_res_outputs.last_hidden_state last_hidden_state = self.high_res_vision_proj(last_hidden_state) - hidden_states = output.hidden_states + hidden_states = high_res_outputs.hidden_states global_hidden_state = hidden_states[self.global_attn_index + 1] # +1 for embedding layer global_hidden_state = self.high_res_vision_neck(global_hidden_state) global_hidden_state = self.high_res_vision_proj(global_hidden_state) @@ -273,14 +303,31 @@ def get_high_res_image_features(self, pixel_values): # batch_size, hidden_size, height, width -> batch_size, seq_len, hidden_size output = output.permute(0, 2, 3, 1) output = output.reshape(output.shape[0], -1, output.shape[-1]) + high_res_outputs.last_hidden_state = output - return output + return high_res_outputs - def get_image_features(self, pixel_values, high_res_pixel_values): - vision_encodings = self.get_low_res_image_features(pixel_values) - high_res_vision_encodings = self.get_high_res_image_features(high_res_pixel_values) - images_embeds = self.aligner(vision_encodings, high_res_vision_encodings) - return images_embeds + @can_return_tuple + @auto_docstring(custom_args=DEEPSEEK_VL_COMMON_CUSTOM_ARGS) + def get_image_features( + self, + pixel_values: torch.FloatTensor, + high_res_pixel_values: torch.FloatTensor, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithHighResVisionEncodings: + low_res_outputs = self.get_low_res_image_features(pixel_values, **kwargs) + high_res_outputs = self.get_high_res_image_features(high_res_pixel_values, **kwargs) + image_features = self.aligner(low_res_outputs.last_hidden_state, high_res_outputs.last_hidden_state) + + return BaseModelOutputWithHighResVisionEncodings( + last_hidden_state=low_res_outputs.last_hidden_state, + pooler_output=image_features, + hidden_states=low_res_outputs.hidden_states, + attentions=low_res_outputs.attentions, + high_res_vision_last_hidden_state=high_res_outputs.last_hidden_state, + high_res_vision_hidden_states=high_res_outputs.hidden_states, + high_res_vision_attentions=high_res_outputs.attentions, + ) @can_return_tuple @auto_docstring(custom_args=DEEPSEEK_VL_COMMON_CUSTOM_ARGS) @@ -319,7 +366,7 @@ def forward( image_attention_mask = input_ids == self.config.image_token_id image_attention_mask = image_attention_mask.unsqueeze(-1).expand_as(inputs_embeds).to(inputs_embeds.device) - image_embeds = self.get_image_features(pixel_values, high_res_pixel_values) + image_embeds = self.get_image_features(pixel_values, high_res_pixel_values, return_dict=True).pooler_output image_features = image_embeds.reshape(-1, inputs_embeds.shape[-1]) image_features = image_features.to(inputs_embeds.device, inputs_embeds.dtype) inputs_embeds = inputs_embeds.masked_scatter(image_attention_mask, image_features) diff --git a/src/transformers/models/dpt/modeling_dpt.py b/src/transformers/models/dpt/modeling_dpt.py index 48e5253413ea..737986492dd5 100755 --- a/src/transformers/models/dpt/modeling_dpt.py +++ b/src/transformers/models/dpt/modeling_dpt.py @@ -990,11 +990,12 @@ def forward( if labels is not None: raise NotImplementedError("Training is not implemented yet") + kwargs["output_hidden_states"] = True if self.backbone is not None: - outputs = self.backbone.forward_with_filtered_kwargs(pixel_values, output_hidden_states=True, **kwargs) + outputs = self.backbone.forward_with_filtered_kwargs(pixel_values, **kwargs) hidden_states = outputs.feature_maps else: - outputs = self.dpt(pixel_values, output_hidden_states=True, **kwargs) + outputs = self.dpt(pixel_values, **kwargs) hidden_states = outputs.hidden_states # only keep certain features based on config.backbone_out_indices # note that the hidden_states also include the initial embeddings @@ -1123,9 +1124,8 @@ def forward( if labels is not None and self.config.num_labels == 1: raise ValueError("The number of labels should be greater than one") - outputs: BaseModelOutputWithPoolingAndIntermediateActivations = self.dpt( - pixel_values, output_hidden_states=True, **kwargs - ) + kwargs["output_hidden_states"] = True + outputs: BaseModelOutputWithPoolingAndIntermediateActivations = self.dpt(pixel_values, **kwargs) hidden_states = outputs.hidden_states # only keep certain features based on config.backbone_out_indices diff --git a/src/transformers/models/edgetam/modeling_edgetam.py b/src/transformers/models/edgetam/modeling_edgetam.py index 0f278913bb45..34c59b50a8d4 100644 --- a/src/transformers/models/edgetam/modeling_edgetam.py +++ b/src/transformers/models/edgetam/modeling_edgetam.py @@ -32,11 +32,11 @@ from ... import initialization as init from ...activations import ACT2FN -from ...modeling_outputs import BaseModelOutput +from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel from ...processing_utils import Unpack from ...pytorch_utils import compile_compatible_method_lru_cache -from ...utils import ModelOutput, auto_docstring, logging +from ...utils import ModelOutput, auto_docstring, can_return_tuple, logging from ...utils.generic import TransformersKwargs, check_model_inputs, is_flash_attention_requested from ..auto import AutoModel from .configuration_edgetam import ( @@ -47,11 +47,6 @@ ) -# fix this in modular -if True: - from ..timm_wrapper.modeling_timm_wrapper import TimmWrapperModel - - logger = logging.get_logger(__name__) @@ -83,16 +78,10 @@ def forward(self, features: torch.Tensor) -> torch.Tensor: @dataclass @auto_docstring(custom_intro="Base class for the vision encoder's outputs.") -class EdgeTamVisionEncoderOutput(ModelOutput): +class EdgeTamVisionEncoderOutput(BaseModelOutputWithPooling): r""" last_hidden_state (`torch.FloatTensor` of shape `(batch_size, height, width, hidden_size)`): Sequence of hidden-states at the output of the last layer of the model. - fpn_hidden_states (`tuple(torch.FloatTensor)`): - Tuple of `torch.FloatTensor` (one for each feature level, from high to low resolution) of shape - `(batch_size, hidden_size, height, width)`. Feature maps from the Feature Pyramid Network neck. - fpn_position_encoding (`tuple(torch.FloatTensor)`): - Tuple of `torch.FloatTensor` (one for each feature level, from high to low resolution) of shape - `(batch_size, hidden_size, height, width)`. Positional encodings corresponding to the `fpn_hidden_states`. hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, + one for the output of each stage) of shape `(batch_size, height, width, hidden_size)`. Hidden-states of the @@ -101,13 +90,16 @@ class EdgeTamVisionEncoderOutput(ModelOutput): Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, sequence_length)`. Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads. + fpn_hidden_states (`tuple(torch.FloatTensor)`): + Tuple of `torch.FloatTensor` (one for each feature level, from high to low resolution) of shape + `(batch_size, hidden_size, height, width)`. Feature maps from the Feature Pyramid Network neck. + fpn_position_encoding (`tuple(torch.FloatTensor)`): + Tuple of `torch.FloatTensor` (one for each feature level, from high to low resolution) of shape + `(batch_size, hidden_size, height, width)`. Positional encodings corresponding to the `fpn_hidden_states`. """ - last_hidden_state: torch.FloatTensor | None = None fpn_hidden_states: torch.FloatTensor | None = None fpn_position_encoding: torch.FloatTensor | None = None - hidden_states: tuple[torch.FloatTensor, ...] | None = None - attentions: tuple[torch.FloatTensor, ...] | None = None def eager_attention_forward( @@ -437,7 +429,9 @@ def forward(self, hidden_states: torch.Tensor) -> tuple[tuple[torch.Tensor, ...] class EdgeTamVisionModel(EdgeTamPreTrainedModel): config_class = EdgeTamVisionConfig main_input_name = "pixel_values" - _can_record_outputs = {"hidden_states": TimmWrapperModel, "attentions": TimmWrapperModel} + # TODO: TimmWrapper models aren't compatible with _can_record_outputs yet. We specifically set this to + # an empty dict to avoid the _can_record_outputs from Sam2VisionModel being inherited here. + _can_record_outputs = {} def __init__(self, config: EdgeTamVisionConfig): super().__init__(config) @@ -460,7 +454,7 @@ def forward( raise ValueError("You have to specify pixel_values") # Forward through backbone - backbone_output = self.backbone(pixel_values) + backbone_output = self.backbone(pixel_values, **kwargs) intermediate_hidden_states = backbone_output.last_hidden_state intermediate_hidden_states = [hidden_state.permute(0, 2, 3, 1) for hidden_state in intermediate_hidden_states] @@ -473,6 +467,7 @@ def forward( last_hidden_state=intermediate_hidden_states[-1], fpn_hidden_states=fpn_hidden_states, fpn_position_encoding=fpn_position_encoding, + hidden_states=backbone_output.hidden_states, ) @@ -981,7 +976,8 @@ def get_image_embeddings( Input pixel values """ batch_size = pixel_values.shape[0] - feature_maps, _, _, _ = self.get_image_features(pixel_values, **kwargs) + image_outputs = self.get_image_features(pixel_values, return_dict=True, **kwargs) + feature_maps = image_outputs.fpn_hidden_states # add no memory embedding to the last feature map feature_maps[-1] = feature_maps[-1] + self.no_memory_embedding @@ -1137,10 +1133,12 @@ def forward( vision_hidden_states = None if pixel_values is not None: - feature_maps, _, vision_hidden_states, vision_attentions = self.get_image_features( - pixel_values, - **kwargs, + image_outputs: EdgeTamVisionEncoderOutput = self.get_image_features( + pixel_values, return_dict=True, **kwargs ) + feature_maps = image_outputs.fpn_hidden_states + vision_hidden_states = image_outputs.hidden_states + vision_attentions = image_outputs.attentions # add no memory embedding to the last feature map feature_maps[-1] = feature_maps[-1] + self.no_memory_embedding @@ -1200,34 +1198,18 @@ def forward( vision_attentions=vision_attentions, ) + @can_return_tuple + @auto_docstring def get_image_features( self, pixel_values: torch.FloatTensor, **kwargs: Unpack[TransformersKwargs], - ) -> tuple[ - list[torch.Tensor], - list[torch.Tensor], - tuple[torch.FloatTensor, ...] | None, - tuple[torch.FloatTensor, ...] | None, - ]: + ) -> tuple | EdgeTamVisionEncoderOutput: r""" - Extract and preprocess image features using the vision encoder. - - Args: - pixel_values (`torch.FloatTensor`): - Input pixel values of shape `(batch_size, num_channels, height, width)`. - - Returns: - `tuple`: A tuple containing: - - feature_maps (`list[torch.Tensor]`): List of feature maps from different levels. - - feature_maps_position_embeddings (`list[torch.Tensor]`): List of positional embeddings for each feature level. - - vision_hidden_states (`tuple[torch.FloatTensor]`, *optional*): Hidden states from the vision encoder. - - vision_attentions (`tuple[torch.FloatTensor]`, *optional*): Attention weights from the vision encoder. + pixel_values (`torch.FloatTensor`): + Input pixel values of shape `(batch_size, num_channels, height, width)`. """ - vision_outputs: EdgeTamVisionEncoderOutput = self.vision_encoder( - pixel_values, - **kwargs, - ) + vision_outputs: EdgeTamVisionEncoderOutput = self.vision_encoder(pixel_values, return_dict=True, **kwargs) feature_maps = vision_outputs.fpn_hidden_states feature_maps_position_embeddings = vision_outputs.fpn_position_encoding @@ -1244,8 +1226,10 @@ def get_image_features( feature_map_position_embedding.flatten(2).permute(2, 0, 1) for feature_map_position_embedding in feature_maps_position_embeddings ] + vision_outputs.fpn_hidden_states = feature_maps + vision_outputs.fpn_position_encoding = feature_maps_position_embeddings - return feature_maps, feature_maps_position_embeddings, vision_outputs.hidden_states, vision_outputs.attentions + return vision_outputs __all__ = ["EdgeTamModel", "EdgeTamVisionModel", "EdgeTamPreTrainedModel"] diff --git a/src/transformers/models/edgetam/modular_edgetam.py b/src/transformers/models/edgetam/modular_edgetam.py index 9aa74c70d0a5..7113d1a8dadc 100644 --- a/src/transformers/models/edgetam/modular_edgetam.py +++ b/src/transformers/models/edgetam/modular_edgetam.py @@ -37,11 +37,6 @@ ) -# fix this in modular -if True: - from ..timm_wrapper.modeling_timm_wrapper import TimmWrapperModel - - class EdgeTamVisionConfig(PreTrainedConfig): r""" This is the configuration class to store the configuration of a [`EdgeTamVisionModel`]. It is used to instantiate a SAM @@ -188,7 +183,9 @@ def _init_weights(self, module): class EdgeTamVisionModel(Sam2VisionModel): config_class = EdgeTamVisionConfig main_input_name = "pixel_values" - _can_record_outputs = {"hidden_states": TimmWrapperModel, "attentions": TimmWrapperModel} + # TODO: TimmWrapper models aren't compatible with _can_record_outputs yet. We specifically set this to + # an empty dict to avoid the _can_record_outputs from Sam2VisionModel being inherited here. + _can_record_outputs = {} def get_input_embeddings(self): raise NotImplementedError("Can't get input embeddings from timm wrapper model") @@ -203,7 +200,7 @@ def forward( raise ValueError("You have to specify pixel_values") # Forward through backbone - backbone_output = self.backbone(pixel_values) + backbone_output = self.backbone(pixel_values, **kwargs) intermediate_hidden_states = backbone_output.last_hidden_state intermediate_hidden_states = [hidden_state.permute(0, 2, 3, 1) for hidden_state in intermediate_hidden_states] @@ -216,6 +213,7 @@ def forward( last_hidden_state=intermediate_hidden_states[-1], fpn_hidden_states=fpn_hidden_states, fpn_position_encoding=fpn_position_encoding, + hidden_states=backbone_output.hidden_states, ) diff --git a/src/transformers/models/edgetam_video/modeling_edgetam_video.py b/src/transformers/models/edgetam_video/modeling_edgetam_video.py index 9ac2cee80398..ef8623eb3bfe 100644 --- a/src/transformers/models/edgetam_video/modeling_edgetam_video.py +++ b/src/transformers/models/edgetam_video/modeling_edgetam_video.py @@ -37,11 +37,11 @@ from ...activations import ACT2FN from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer -from ...modeling_outputs import BaseModelOutput +from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel from ...processing_utils import Unpack from ...pytorch_utils import compile_compatible_method_lru_cache -from ...utils import ModelOutput, auto_docstring, logging +from ...utils import ModelOutput, auto_docstring, can_return_tuple, logging from ...utils.generic import TransformersKwargs, is_flash_attention_requested from ..auto import AutoModel from .configuration_edgetam_video import ( @@ -119,16 +119,10 @@ def forward(self, hidden_states): @dataclass @auto_docstring(custom_intro="Base class for the vision encoder's outputs.") -class EdgeTamVideoVisionEncoderOutput(ModelOutput): +class EdgeTamVideoVisionEncoderOutput(BaseModelOutputWithPooling): r""" last_hidden_state (`torch.FloatTensor` of shape `(batch_size, height, width, hidden_size)`): Sequence of hidden-states at the output of the last layer of the model. - fpn_hidden_states (`tuple(torch.FloatTensor)`): - Tuple of `torch.FloatTensor` (one for each feature level, from high to low resolution) of shape - `(batch_size, hidden_size, height, width)`. Feature maps from the Feature Pyramid Network neck. - fpn_position_encoding (`tuple(torch.FloatTensor)`): - Tuple of `torch.FloatTensor` (one for each feature level, from high to low resolution) of shape - `(batch_size, hidden_size, height, width)`. Positional encodings corresponding to the `fpn_hidden_states`. hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, + one for the output of each stage) of shape `(batch_size, height, width, hidden_size)`. Hidden-states of the @@ -137,13 +131,16 @@ class EdgeTamVideoVisionEncoderOutput(ModelOutput): Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, sequence_length)`. Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads. + fpn_hidden_states (`tuple(torch.FloatTensor)`): + Tuple of `torch.FloatTensor` (one for each feature level, from high to low resolution) of shape + `(batch_size, hidden_size, height, width)`. Feature maps from the Feature Pyramid Network neck. + fpn_position_encoding (`tuple(torch.FloatTensor)`): + Tuple of `torch.FloatTensor` (one for each feature level, from high to low resolution) of shape + `(batch_size, hidden_size, height, width)`. Positional encodings corresponding to the `fpn_hidden_states`. """ - last_hidden_state: torch.FloatTensor | None = None fpn_hidden_states: torch.FloatTensor | None = None fpn_position_encoding: torch.FloatTensor | None = None - hidden_states: tuple[torch.FloatTensor, ...] | None = None - attentions: tuple[torch.FloatTensor, ...] | None = None class EdgeTamVideoVisionRotaryEmbedding(nn.Module): @@ -2086,7 +2083,8 @@ def get_image_embeddings( Input pixel values """ batch_size = pixel_values.shape[0] - feature_maps, _, _, _ = self.get_image_features(pixel_values, **kwargs) + image_outputs = self.get_image_features(pixel_values, return_dict=True, **kwargs) + feature_maps = image_outputs.fpn_hidden_states # add no memory embedding to the last feature map feature_maps[-1] = feature_maps[-1] + self.no_memory_embedding @@ -2231,34 +2229,18 @@ def forward( frame_idx=frame_idx, ) + @can_return_tuple + @auto_docstring def get_image_features( self, pixel_values: torch.FloatTensor, **kwargs: Unpack[TransformersKwargs], - ) -> tuple[ - list[torch.Tensor], - list[torch.Tensor], - tuple[torch.FloatTensor, ...] | None, - tuple[torch.FloatTensor, ...] | None, - ]: + ) -> tuple | EdgeTamVideoVisionEncoderOutput: r""" - Extract and preprocess image features using the vision encoder. - - Args: - pixel_values (`torch.FloatTensor`): - Input pixel values of shape `(batch_size, num_channels, height, width)`. - - Returns: - `tuple`: A tuple containing: - - feature_maps (`list[torch.Tensor]`): List of feature maps from different levels. - - feature_maps_position_embeddings (`list[torch.Tensor]`): List of positional embeddings for each feature level. - - vision_hidden_states (`tuple[torch.FloatTensor]`, *optional*): Hidden states from the vision encoder. - - vision_attentions (`tuple[torch.FloatTensor]`, *optional*): Attention weights from the vision encoder. + pixel_values (`torch.FloatTensor`): + Input pixel values of shape `(batch_size, num_channels, height, width)`. """ - vision_outputs: EdgeTamVideoVisionEncoderOutput = self.vision_encoder( - pixel_values, - **kwargs, - ) + vision_outputs: EdgeTamVideoVisionEncoderOutput = self.vision_encoder(pixel_values, return_dict=True, **kwargs) feature_maps = vision_outputs.fpn_hidden_states feature_maps_position_embeddings = vision_outputs.fpn_position_encoding @@ -2275,8 +2257,10 @@ def get_image_features( feature_map_position_embedding.flatten(2).permute(2, 0, 1) for feature_map_position_embedding in feature_maps_position_embeddings ] + vision_outputs.fpn_hidden_states = feature_maps + vision_outputs.fpn_position_encoding = feature_maps_position_embeddings - return feature_maps, feature_maps_position_embeddings, vision_outputs.hidden_states, vision_outputs.attentions + return vision_outputs def _prepare_vision_features( self, @@ -2293,7 +2277,9 @@ def _prepare_vision_features( else: # Compute features using image encoder image_batch = inference_session.get_frame(frame_idx).unsqueeze(0) # Add batch dimension - vision_feats, vision_pos_embeds, _, _ = self.get_image_features(image_batch) + image_outputs = self.get_image_features(image_batch, return_dict=True) + vision_feats = image_outputs.fpn_hidden_states + vision_pos_embeds = image_outputs.fpn_position_embeddings # Cache features inference_session.cache.cache_vision_features( frame_idx, {"vision_feats": vision_feats, "vision_pos_embeds": vision_pos_embeds} @@ -2398,10 +2384,10 @@ def _single_frame_forward( vision_hidden_states = None if pixel_values is not None: - feature_maps, _, vision_hidden_states, vision_attentions = self.get_image_features( - pixel_values, - **kwargs, - ) + image_outputs = self.get_image_features(pixel_values, return_dict=True, **kwargs) + feature_maps = image_outputs.fpn_hidden_states + vision_hidden_states = image_outputs.hidden_states + vision_attentions = image_outputs.attentions # add no memory embedding to the last feature map feature_maps[-1] = feature_maps[-1] + self.no_memory_embedding diff --git a/src/transformers/models/emu3/modeling_emu3.py b/src/transformers/models/emu3/modeling_emu3.py index b793a024b9d0..2d41f4aead43 100644 --- a/src/transformers/models/emu3/modeling_emu3.py +++ b/src/transformers/models/emu3/modeling_emu3.py @@ -21,6 +21,7 @@ import math from collections.abc import Callable +from dataclasses import dataclass from functools import cached_property from typing import Optional @@ -35,7 +36,7 @@ from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub, use_kernelized_func from ...masking_utils import create_causal_mask from ...modeling_layers import GradientCheckpointingLayer -from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast +from ...modeling_outputs import BaseModelOutputWithPast, BaseModelOutputWithPooling, CausalLMOutputWithPast from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel from ...processing_utils import Unpack @@ -44,6 +45,17 @@ from .configuration_emu3 import Emu3Config, Emu3TextConfig, Emu3VQVAEConfig +@dataclass +@auto_docstring +class Emu3VQVAEModelOutput(BaseModelOutputWithPooling): + r""" + image_tokens (`torch.LongTensor` of shape `(batch_size, config.vocab_size`): + Indices of the image tokens predicted by the VQ-VAE model. + """ + + image_tokens: torch.LongTensor | None = None + + def rotate_half(x): """Rotates half the hidden dims of the input.""" x1 = x[..., : x.shape[-1] // 2] @@ -937,6 +949,10 @@ class Emu3VQVAE(PreTrainedModel): "Emu3VQVAEResnetBlock", "Emu3VQVAEVectorQuantizer", ] + _can_record_outputs = { + "hidden_states": [Emu3VQVAEResnetBlock, Emu3VQVAETemporalResnetBlock], + "attentions": Emu3VQVAEAttentionBlock, + } @torch.no_grad() def _init_weights(self, module): @@ -986,7 +1002,10 @@ def __init__(self, config: Emu3VQVAEConfig): self.post_init() - def encode(self, pixel_values: torch.Tensor, image_sizes: torch.Tensor): + @check_model_inputs + def encode( + self, pixel_values: torch.Tensor, image_sizes: torch.Tensor, **kwargs: Unpack[TransformersKwargs] + ) -> Emu3VQVAEModelOutput: is_image = pixel_values.ndim == 4 if is_image: temporal = self.config.temporal_downsample_factor @@ -998,12 +1017,12 @@ def encode(self, pixel_values: torch.Tensor, image_sizes: torch.Tensor): hidden_states = self.encoder(pixel_values) # b t c h w -> b c t h w - hidden_states = hidden_states.permute(0, 2, 1, 3, 4) - hidden_states = self.quant_conv(hidden_states) + conv_hidden_states = hidden_states.permute(0, 2, 1, 3, 4) + conv_hidden_states = self.quant_conv(conv_hidden_states) # b c t h w -> b t c h w - hidden_states = hidden_states.permute(0, 2, 1, 3, 4) - codes = self.quantize(hidden_states) + conv_hidden_states = conv_hidden_states.permute(0, 2, 1, 3, 4) + codes = self.quantize(conv_hidden_states) image_tokens = codes.squeeze(1) if is_image else codes @@ -1012,7 +1031,10 @@ def encode(self, pixel_values: torch.Tensor, image_sizes: torch.Tensor): for single_image, size in zip(image_tokens, image_sizes) ] - return image_tokens + return Emu3VQVAEModelOutput( + last_hidden_state=hidden_states, + image_tokens=image_tokens, + ) def decode(self, hidden_states: torch.Tensor): is_image = hidden_states.ndim == 3 @@ -1110,6 +1132,10 @@ class Emu3PreTrainedModel(PreTrainedModel): _can_compile_fullgraph = True _supports_flex_attn = True _supports_attention_backend = True + _can_record_outputs = { + "hidden_states": Emu3DecoderLayer, + "attentions": Emu3Attention, + } class Emu3RotaryEmbedding(nn.Module): @@ -1179,11 +1205,6 @@ def forward(self, x, position_ids): @auto_docstring class Emu3TextModel(Emu3PreTrainedModel): - _can_record_outputs = { - "hidden_states": Emu3DecoderLayer, - "attentions": Emu3Attention, - } - def __init__(self, config: Emu3Config): super().__init__(config) self.padding_idx = config.pad_token_id @@ -1357,7 +1378,7 @@ def get_input_embeddings(self): def set_input_embeddings(self, value): self.text_model.set_input_embeddings(value) - def get_image_tokens(self, pixel_values: torch.FloatTensor, image_sizes: torch.LongTensor): + def get_image_tokens(self, pixel_values: torch.FloatTensor, image_sizes: torch.LongTensor) -> torch.LongTensor: """ Tokenizes images into discrete tokens with VQGAN module. Converts obtained image tokens into BPE tokens and wraps with "boi" and "eoi" @@ -1369,28 +1390,40 @@ def get_image_tokens(self, pixel_values: torch.FloatTensor, image_sizes: torch.L image_sizes (`torch.LongTensor` of shape `(batch_size, 2)`): The sizes of the images in the batch, being (height, width) for each image. """ - image_tokens_list = self.vqmodel.encode(pixel_values, image_sizes) - bpe_tokens_list = [self.vocabulary_mapping.convert_img2bpe(tokens).flatten() for tokens in image_tokens_list] + vqmodel_outputs: Emu3VQVAEModelOutput = self.vqmodel.encode(pixel_values, image_sizes, return_dict=True) + bpe_tokens_list = [ + self.vocabulary_mapping.convert_img2bpe(tokens).flatten() for tokens in vqmodel_outputs.image_tokens + ] bpe_tokens = torch.cat(bpe_tokens_list) return bpe_tokens - def get_image_features(self, pixel_values: torch.FloatTensor, image_sizes: torch.LongTensor): - """ - Tokenizes images into discrete tokens with VQGAN module and embeds - them with text embeddings layer - - Args: - pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)): - The tensors corresponding to the input images. + @can_return_tuple + @auto_docstring( + custom_intro="Tokenizes images into discrete tokens with VQGAN module and embeds them with text embeddings layer" + ) + def get_image_features( + self, pixel_values: torch.FloatTensor, image_sizes: torch.LongTensor, **kwargs: Unpack[TransformersKwargs] + ) -> tuple | Emu3VQVAEModelOutput: + r""" + pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)): + The tensors corresponding to the input images. """ - image_tokens = self.get_image_tokens(pixel_values, image_sizes) + vqmodel_outputs: Emu3VQVAEModelOutput = self.vqmodel.encode( + pixel_values, image_sizes, return_dict=True, **kwargs + ) split_sizes = [ (height // self.vqmodel.vision_spatial_factor) * (width // self.vqmodel.vision_spatial_factor + 1) for height, width in image_sizes ] - image_features = self.get_input_embeddings()(image_tokens) - image_features = torch.split(image_features, split_sizes) - return image_features + bpe_tokens_list = [ + self.vocabulary_mapping.convert_img2bpe(tokens).flatten() for tokens in vqmodel_outputs.image_tokens + ] + bpe_tokens = torch.cat(bpe_tokens_list) + image_embeddings = self.get_input_embeddings()(bpe_tokens) + image_features = torch.split(image_embeddings, split_sizes) + vqmodel_outputs.pooler_output = image_features + + return vqmodel_outputs @torch.no_grad() def decode_image_tokens(self, image_tokens: torch.LongTensor, height: int, width: int): @@ -1465,12 +1498,12 @@ def forward( inputs_embeds = self.get_input_embeddings()(input_ids) if pixel_values is not None: - image_embeds = self.get_image_features(pixel_values, image_sizes) - image_embeds = torch.cat(image_embeds, dim=0) + image_features = self.get_image_features(pixel_values, image_sizes).pooler_output + image_features = torch.cat(image_features, dim=0) special_image_mask = self.get_placeholder_mask( - input_ids, inputs_embeds=inputs_embeds, image_features=image_embeds + input_ids, inputs_embeds=inputs_embeds, image_features=image_features ) - inputs_embeds = inputs_embeds.masked_scatter(special_image_mask, image_embeds) + inputs_embeds = inputs_embeds.masked_scatter(special_image_mask, image_features) # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) outputs = self.text_model( diff --git a/src/transformers/models/emu3/modular_emu3.py b/src/transformers/models/emu3/modular_emu3.py index 167ad64b14b7..ac4c79ae3cf2 100644 --- a/src/transformers/models/emu3/modular_emu3.py +++ b/src/transformers/models/emu3/modular_emu3.py @@ -14,6 +14,7 @@ # limitations under the License. import math +from dataclasses import dataclass from functools import cached_property import torch @@ -23,10 +24,11 @@ from ... import initialization as init from ...cache_utils import Cache from ...generation import GenerationMixin -from ...modeling_outputs import CausalLMOutputWithPast +from ...modeling_outputs import BaseModelOutputWithPooling, CausalLMOutputWithPast from ...modeling_utils import PreTrainedModel from ...processing_utils import Unpack from ...utils import auto_docstring, can_return_tuple, logging, torch_compilable_check +from ...utils.generic import check_model_inputs from ..chameleon.modeling_chameleon import ( ChameleonPreTrainedModel, ChameleonVQVAEEncoderConvDownsample, @@ -39,6 +41,17 @@ logger = logging.get_logger(__name__) +@dataclass +@auto_docstring +class Emu3VQVAEModelOutput(BaseModelOutputWithPooling): + r""" + image_tokens (`torch.LongTensor` of shape `(batch_size, config.vocab_size`): + Indices of the image tokens predicted by the VQ-VAE model. + """ + + image_tokens: torch.LongTensor | None = None + + class Emu3Attention(LlamaAttention): pass @@ -686,6 +699,10 @@ class Emu3VQVAE(PreTrainedModel): "Emu3VQVAEResnetBlock", "Emu3VQVAEVectorQuantizer", ] + _can_record_outputs = { + "hidden_states": [Emu3VQVAEResnetBlock, Emu3VQVAETemporalResnetBlock], + "attentions": Emu3VQVAEAttentionBlock, + } @torch.no_grad() def _init_weights(self, module): @@ -735,7 +752,10 @@ def __init__(self, config: Emu3VQVAEConfig): self.post_init() - def encode(self, pixel_values: torch.Tensor, image_sizes: torch.Tensor): + @check_model_inputs + def encode( + self, pixel_values: torch.Tensor, image_sizes: torch.Tensor, **kwargs: Unpack[TransformersKwargs] + ) -> Emu3VQVAEModelOutput: is_image = pixel_values.ndim == 4 if is_image: temporal = self.config.temporal_downsample_factor @@ -747,12 +767,12 @@ def encode(self, pixel_values: torch.Tensor, image_sizes: torch.Tensor): hidden_states = self.encoder(pixel_values) # b t c h w -> b c t h w - hidden_states = hidden_states.permute(0, 2, 1, 3, 4) - hidden_states = self.quant_conv(hidden_states) + conv_hidden_states = hidden_states.permute(0, 2, 1, 3, 4) + conv_hidden_states = self.quant_conv(conv_hidden_states) # b c t h w -> b t c h w - hidden_states = hidden_states.permute(0, 2, 1, 3, 4) - codes = self.quantize(hidden_states) + conv_hidden_states = conv_hidden_states.permute(0, 2, 1, 3, 4) + codes = self.quantize(conv_hidden_states) image_tokens = codes.squeeze(1) if is_image else codes @@ -761,7 +781,10 @@ def encode(self, pixel_values: torch.Tensor, image_sizes: torch.Tensor): for single_image, size in zip(image_tokens, image_sizes) ] - return image_tokens + return Emu3VQVAEModelOutput( + last_hidden_state=hidden_states, + image_tokens=image_tokens, + ) def decode(self, hidden_states: torch.Tensor): is_image = hidden_states.ndim == 3 @@ -849,14 +872,13 @@ class Emu3PreTrainedModel(ChameleonPreTrainedModel, Emu3VQVAE): ] _supports_flex_attn = True _supports_attention_backend = True - - -class Emu3TextModel(LlamaModel, Emu3PreTrainedModel): _can_record_outputs = { "hidden_states": Emu3DecoderLayer, "attentions": Emu3Attention, } + +class Emu3TextModel(LlamaModel, Emu3PreTrainedModel): def __init__(self, config: Emu3Config): super().__init__(config) self.layers = nn.ModuleList( @@ -910,7 +932,7 @@ def get_input_embeddings(self): def set_input_embeddings(self, value): self.text_model.set_input_embeddings(value) - def get_image_tokens(self, pixel_values: torch.FloatTensor, image_sizes: torch.LongTensor): + def get_image_tokens(self, pixel_values: torch.FloatTensor, image_sizes: torch.LongTensor) -> torch.LongTensor: """ Tokenizes images into discrete tokens with VQGAN module. Converts obtained image tokens into BPE tokens and wraps with "boi" and "eoi" @@ -922,28 +944,40 @@ def get_image_tokens(self, pixel_values: torch.FloatTensor, image_sizes: torch.L image_sizes (`torch.LongTensor` of shape `(batch_size, 2)`): The sizes of the images in the batch, being (height, width) for each image. """ - image_tokens_list = self.vqmodel.encode(pixel_values, image_sizes) - bpe_tokens_list = [self.vocabulary_mapping.convert_img2bpe(tokens).flatten() for tokens in image_tokens_list] + vqmodel_outputs: Emu3VQVAEModelOutput = self.vqmodel.encode(pixel_values, image_sizes, return_dict=True) + bpe_tokens_list = [ + self.vocabulary_mapping.convert_img2bpe(tokens).flatten() for tokens in vqmodel_outputs.image_tokens + ] bpe_tokens = torch.cat(bpe_tokens_list) return bpe_tokens - def get_image_features(self, pixel_values: torch.FloatTensor, image_sizes: torch.LongTensor): - """ - Tokenizes images into discrete tokens with VQGAN module and embeds - them with text embeddings layer - - Args: - pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)): - The tensors corresponding to the input images. + @can_return_tuple + @auto_docstring( + custom_intro="Tokenizes images into discrete tokens with VQGAN module and embeds them with text embeddings layer" + ) + def get_image_features( + self, pixel_values: torch.FloatTensor, image_sizes: torch.LongTensor, **kwargs: Unpack[TransformersKwargs] + ) -> tuple | Emu3VQVAEModelOutput: + r""" + pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)): + The tensors corresponding to the input images. """ - image_tokens = self.get_image_tokens(pixel_values, image_sizes) + vqmodel_outputs: Emu3VQVAEModelOutput = self.vqmodel.encode( + pixel_values, image_sizes, return_dict=True, **kwargs + ) split_sizes = [ (height // self.vqmodel.vision_spatial_factor) * (width // self.vqmodel.vision_spatial_factor + 1) for height, width in image_sizes ] - image_features = self.get_input_embeddings()(image_tokens) - image_features = torch.split(image_features, split_sizes) - return image_features + bpe_tokens_list = [ + self.vocabulary_mapping.convert_img2bpe(tokens).flatten() for tokens in vqmodel_outputs.image_tokens + ] + bpe_tokens = torch.cat(bpe_tokens_list) + image_embeddings = self.get_input_embeddings()(bpe_tokens) + image_features = torch.split(image_embeddings, split_sizes) + vqmodel_outputs.pooler_output = image_features + + return vqmodel_outputs @torch.no_grad() def decode_image_tokens(self, image_tokens: torch.LongTensor, height: int, width: int): @@ -1018,12 +1052,12 @@ def forward( inputs_embeds = self.get_input_embeddings()(input_ids) if pixel_values is not None: - image_embeds = self.get_image_features(pixel_values, image_sizes) - image_embeds = torch.cat(image_embeds, dim=0) + image_features = self.get_image_features(pixel_values, image_sizes).pooler_output + image_features = torch.cat(image_features, dim=0) special_image_mask = self.get_placeholder_mask( - input_ids, inputs_embeds=inputs_embeds, image_features=image_embeds + input_ids, inputs_embeds=inputs_embeds, image_features=image_features ) - inputs_embeds = inputs_embeds.masked_scatter(special_image_mask, image_embeds) + inputs_embeds = inputs_embeds.masked_scatter(special_image_mask, image_features) # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) outputs = self.text_model( diff --git a/src/transformers/models/ernie4_5_vl_moe/modeling_ernie4_5_vl_moe.py b/src/transformers/models/ernie4_5_vl_moe/modeling_ernie4_5_vl_moe.py index a5d38649cf26..3b02f84c8d84 100644 --- a/src/transformers/models/ernie4_5_vl_moe/modeling_ernie4_5_vl_moe.py +++ b/src/transformers/models/ernie4_5_vl_moe/modeling_ernie4_5_vl_moe.py @@ -34,7 +34,7 @@ from ...masking_utils import create_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer -from ...modeling_outputs import MoeCausalLMOutputWithPast, MoeModelOutputWithPast +from ...modeling_outputs import BaseModelOutputWithPooling, MoeCausalLMOutputWithPast, MoeModelOutputWithPast from ...modeling_rope_utils import dynamic_rope_update from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel from ...processing_utils import Unpack @@ -552,6 +552,142 @@ def forward( return hidden_states +def rotate_half(x): + """Rotates half the hidden dims of the input.""" + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + +def apply_rotary_pos_emb_vision( + q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor +) -> tuple[torch.Tensor, torch.Tensor]: + orig_q_dtype = q.dtype + orig_k_dtype = k.dtype + q, k = q.float(), k.float() + cos, sin = cos.unsqueeze(-2).float(), sin.unsqueeze(-2).float() + q_embed = (q * cos) + (rotate_half(q) * sin) + k_embed = (k * cos) + (rotate_half(k) * sin) + q_embed = q_embed.to(orig_q_dtype) + k_embed = k_embed.to(orig_k_dtype) + return q_embed, k_embed + + +class Ernie4_5_VL_MoeVisionAttention(nn.Module): + def __init__(self, config: Ernie4_5_VL_MoeVisionConfig) -> None: + super().__init__() + self.dim = config.hidden_size + self.num_heads = config.num_heads + self.head_dim = self.dim // self.num_heads + self.num_key_value_groups = 1 # needed for eager attention + self.qkv = nn.Linear(self.dim, self.dim * 3, bias=True) + self.proj = nn.Linear(self.dim, self.dim) + self.scaling = self.head_dim**-0.5 + self.config = config + self.attention_dropout = 0.0 + self.is_causal = False + + def forward( + self, + hidden_states: torch.Tensor, + cu_seqlens: torch.Tensor, + rotary_pos_emb: torch.Tensor | None = None, + position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, + **kwargs, + ) -> torch.Tensor: + seq_length = hidden_states.shape[0] + query_states, key_states, value_states = ( + self.qkv(hidden_states).reshape(seq_length, 3, self.num_heads, -1).permute(1, 0, 2, 3).unbind(0) + ) + cos, sin = position_embeddings + query_states, key_states = apply_rotary_pos_emb_vision(query_states, key_states, cos, sin) + + query_states = query_states.transpose(0, 1).unsqueeze(0) + key_states = key_states.transpose(0, 1).unsqueeze(0) + value_states = value_states.transpose(0, 1).unsqueeze(0) + + attention_interface: Callable = eager_attention_forward + if self.config._attn_implementation != "eager": + attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation] + + if is_flash_attention_requested(self.config): + # Flash Attention: Use cu_seqlens for variable length attention + max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max() + attn_output, _ = attention_interface( + self, + query_states, + key_states, + value_states, + attention_mask=None, + scaling=self.scaling, + dropout=0.0 if not self.training else self.attention_dropout, + cu_seq_lens_q=cu_seqlens, + cu_seq_lens_k=cu_seqlens, + max_length_q=max_seqlen, + max_length_k=max_seqlen, + is_causal=False, + **kwargs, + ) + else: + # Other implementations: Process each chunk separately + lengths = cu_seqlens[1:] - cu_seqlens[:-1] + splits = [ + torch.split(tensor, lengths.tolist(), dim=2) for tensor in (query_states, key_states, value_states) + ] + + attn_outputs = [ + attention_interface( + self, + q, + k, + v, + attention_mask=None, + scaling=self.scaling, + dropout=0.0 if not self.training else self.attention_dropout, + is_causal=False, + **kwargs, + )[0] + for q, k, v in zip(*splits) + ] + attn_output = torch.cat(attn_outputs, dim=1) + + attn_output = attn_output.reshape(seq_length, -1).contiguous() + attn_output = self.proj(attn_output) + return attn_output + + +class Ernie4_5_VL_MoeVisionBlock(GradientCheckpointingLayer): + def __init__(self, config) -> None: + super().__init__() + + self.norm1 = nn.LayerNorm(config.hidden_size, config.rms_norm_eps) + self.norm2 = nn.LayerNorm(config.hidden_size, config.rms_norm_eps) + self.attn = Ernie4_5_VL_MoeVisionAttention(config=config) + self.mlp = Ernie4_5VLVisionMLP( + dim=config.hidden_size, + hidden_dim=config.intermediate_size, + hidden_act=config.hidden_act, + ) + + def forward( + self, + hidden_states: torch.Tensor, + cu_seqlens: torch.Tensor, + rotary_pos_emb: torch.Tensor | None = None, + position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, + **kwargs, + ) -> torch.Tensor: + hidden_states = hidden_states + self.attn( + self.norm1(hidden_states), + cu_seqlens=cu_seqlens, + rotary_pos_emb=rotary_pos_emb, + position_embeddings=position_embeddings, + **kwargs, + ) + hidden_states = hidden_states + self.mlp(self.norm2(hidden_states)) + return hidden_states + + @auto_docstring class Ernie4_5_VL_MoePreTrainedModel(PreTrainedModel): config: Ernie4_5_VL_MoeConfig @@ -736,148 +872,17 @@ def forward(self, seqlen: int) -> torch.Tensor: return freqs -def rotate_half(x): - """Rotates half the hidden dims of the input.""" - x1 = x[..., : x.shape[-1] // 2] - x2 = x[..., x.shape[-1] // 2 :] - return torch.cat((-x2, x1), dim=-1) - - -def apply_rotary_pos_emb_vision( - q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor -) -> tuple[torch.Tensor, torch.Tensor]: - orig_q_dtype = q.dtype - orig_k_dtype = k.dtype - q, k = q.float(), k.float() - cos, sin = cos.unsqueeze(-2).float(), sin.unsqueeze(-2).float() - q_embed = (q * cos) + (rotate_half(q) * sin) - k_embed = (k * cos) + (rotate_half(k) * sin) - q_embed = q_embed.to(orig_q_dtype) - k_embed = k_embed.to(orig_k_dtype) - return q_embed, k_embed - - -class Ernie4_5_VL_MoeVisionAttention(nn.Module): - def __init__(self, config: Ernie4_5_VL_MoeVisionConfig) -> None: - super().__init__() - self.dim = config.hidden_size - self.num_heads = config.num_heads - self.head_dim = self.dim // self.num_heads - self.num_key_value_groups = 1 # needed for eager attention - self.qkv = nn.Linear(self.dim, self.dim * 3, bias=True) - self.proj = nn.Linear(self.dim, self.dim) - self.scaling = self.head_dim**-0.5 - self.config = config - self.attention_dropout = 0.0 - self.is_causal = False - - def forward( - self, - hidden_states: torch.Tensor, - cu_seqlens: torch.Tensor, - rotary_pos_emb: torch.Tensor | None = None, - position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, - **kwargs, - ) -> torch.Tensor: - seq_length = hidden_states.shape[0] - query_states, key_states, value_states = ( - self.qkv(hidden_states).reshape(seq_length, 3, self.num_heads, -1).permute(1, 0, 2, 3).unbind(0) - ) - cos, sin = position_embeddings - query_states, key_states = apply_rotary_pos_emb_vision(query_states, key_states, cos, sin) - - query_states = query_states.transpose(0, 1).unsqueeze(0) - key_states = key_states.transpose(0, 1).unsqueeze(0) - value_states = value_states.transpose(0, 1).unsqueeze(0) - - attention_interface: Callable = eager_attention_forward - if self.config._attn_implementation != "eager": - attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation] - - if is_flash_attention_requested(self.config): - # Flash Attention: Use cu_seqlens for variable length attention - max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max() - attn_output, _ = attention_interface( - self, - query_states, - key_states, - value_states, - attention_mask=None, - scaling=self.scaling, - dropout=0.0 if not self.training else self.attention_dropout, - cu_seq_lens_q=cu_seqlens, - cu_seq_lens_k=cu_seqlens, - max_length_q=max_seqlen, - max_length_k=max_seqlen, - is_causal=False, - **kwargs, - ) - else: - # Other implementations: Process each chunk separately - lengths = cu_seqlens[1:] - cu_seqlens[:-1] - splits = [ - torch.split(tensor, lengths.tolist(), dim=2) for tensor in (query_states, key_states, value_states) - ] - - attn_outputs = [ - attention_interface( - self, - q, - k, - v, - attention_mask=None, - scaling=self.scaling, - dropout=0.0 if not self.training else self.attention_dropout, - is_causal=False, - **kwargs, - )[0] - for q, k, v in zip(*splits) - ] - attn_output = torch.cat(attn_outputs, dim=1) - - attn_output = attn_output.reshape(seq_length, -1).contiguous() - attn_output = self.proj(attn_output) - return attn_output - - -class Ernie4_5_VL_MoeVisionBlock(GradientCheckpointingLayer): - def __init__(self, config) -> None: - super().__init__() - - self.norm1 = nn.LayerNorm(config.hidden_size, config.rms_norm_eps) - self.norm2 = nn.LayerNorm(config.hidden_size, config.rms_norm_eps) - self.attn = Ernie4_5_VL_MoeVisionAttention(config=config) - self.mlp = Ernie4_5VLVisionMLP( - dim=config.hidden_size, - hidden_dim=config.intermediate_size, - hidden_act=config.hidden_act, - ) - - def forward( - self, - hidden_states: torch.Tensor, - cu_seqlens: torch.Tensor, - rotary_pos_emb: torch.Tensor | None = None, - position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, - **kwargs, - ) -> torch.Tensor: - hidden_states = hidden_states + self.attn( - self.norm1(hidden_states), - cu_seqlens=cu_seqlens, - rotary_pos_emb=rotary_pos_emb, - position_embeddings=position_embeddings, - **kwargs, - ) - hidden_states = hidden_states + self.mlp(self.norm2(hidden_states)) - return hidden_states - - @auto_docstring class Ernie4_5_VL_MoeVisionTransformerPretrainedModel(Ernie4_5_VL_MoePreTrainedModel): config: Ernie4_5_VL_MoeVisionConfig input_modalities = ("image", "video") _no_split_modules = ["Ernie4_5_VL_MoeVisionBlock"] _input_embed_layer = "patch_embed" + _can_record_outputs = { + "router_logits": OutputRecorder(Ernie4_5_VL_MoeMoeBlock, index=1), + "hidden_states": Ernie4_5_VL_MoeVisionBlock, + "attentions": Ernie4_5_VL_MoeVisionAttention, + } def __init__(self, config) -> None: super().__init__(config) @@ -928,13 +933,10 @@ def rot_pos_emb(self, grid_thw): rotary_pos_emb = rotary_pos_emb_full[pos_ids].flatten(1) return rotary_pos_emb - @auto_docstring + @check_model_inputs def forward( - self, - hidden_states: torch.Tensor, - grid_thw: torch.Tensor, - **kwargs, - ) -> torch.Tensor: + self, hidden_states: torch.Tensor, grid_thw: torch.Tensor, **kwargs: Unpack[TransformersKwargs] + ) -> tuple | BaseModelOutputWithPooling: r""" grid_thw (`torch.LongTensor` of shape `(num_images, 3)`): The temporal, height and width dimensions of feature shape for each image. Each row contains [t, h, w] values. @@ -962,7 +964,7 @@ def forward( **kwargs, ) hidden_states = self.ln(hidden_states) - return hidden_states + return BaseModelOutputWithPooling(last_hidden_state=hidden_states) class Ernie4_5_VL_MoeVisionMLP(nn.Module): @@ -1264,43 +1266,51 @@ def get_rope_index( return position_ids, mrope_position_deltas + @can_return_tuple + @auto_docstring def get_video_features( - self, pixel_values_videos: torch.FloatTensor, video_grid_thw: torch.LongTensor | None = None - ): - """ - Encodes videos into continuous embeddings that can be forwarded to the language model. - - Args: - pixel_values_videos (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): - The tensors corresponding to the input videos. - video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*): - The temporal, height and width of feature shape of each video in LLM. + self, + pixel_values_videos: torch.FloatTensor, + video_grid_thw: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + pixel_values_videos (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): + The tensors corresponding to the input videos. + video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*): + The temporal, height and width of feature shape of each video in LLM. """ - video_embeds = self.vision_tower(pixel_values_videos, video_grid_thw) - video_embeds = self.resampler_model(video_embeds, video_grid_thw) + video_outputs = self.vision_tower(pixel_values_videos, video_grid_thw, return_dict=True, **kwargs) + video_embeds = self.resampler_model(video_outputs.last_hidden_state, video_grid_thw) split_sizes = ( video_grid_thw.prod(-1) // self.vision_tower.spatial_merge_size**2 // self.resampler_model.temporal_merge_size ).tolist() video_embeds = torch.split(video_embeds, split_sizes) - return video_embeds - - def get_image_features(self, pixel_values: torch.FloatTensor, image_grid_thw: torch.LongTensor | None = None): - """ - Encodes images into continuous embeddings that can be forwarded to the language model. + video_outputs.pooler_output = video_embeds + return video_outputs - Args: - pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): - The tensors corresponding to the input images. - image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): - The temporal, height and width of feature shape of each image in LLM. + @can_return_tuple + @auto_docstring + def get_image_features( + self, + pixel_values: torch.FloatTensor, + image_grid_thw: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): + The tensors corresponding to the input images. + image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): + The temporal, height and width of feature shape of each image in LLM. """ - image_embeds = self.vision_tower(pixel_values, image_grid_thw) - image_embeds = self.resampler_model(image_embeds, image_grid_thw) + image_outputs = self.vision_tower(pixel_values, image_grid_thw, return_dict=True, **kwargs) + image_embeds = self.resampler_model(image_outputs.last_hidden_state, image_grid_thw) split_sizes = (image_grid_thw.prod(-1) // self.vision_tower.spatial_merge_size**2).tolist() image_embeds = torch.split(image_embeds, split_sizes) - return image_embeds + image_outputs.pooler_output = image_embeds + return image_outputs def get_placeholder_mask( self, @@ -1379,7 +1389,7 @@ def forward( inputs_embeds = self.get_input_embeddings()(input_ids) if pixel_values is not None: - image_embeds = self.get_image_features(pixel_values, image_grid_thw) + image_embeds = self.get_image_features(pixel_values, image_grid_thw, return_dict=True).pooler_output image_embeds = torch.cat(image_embeds, dim=0).to(inputs_embeds.device, inputs_embeds.dtype) image_mask, _ = self.get_placeholder_mask( input_ids, inputs_embeds=inputs_embeds, image_features=image_embeds @@ -1387,7 +1397,7 @@ def forward( inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds) if pixel_values_videos is not None: - video_embeds = self.get_video_features(pixel_values_videos, video_grid_thw) + video_embeds = self.get_video_features(pixel_values_videos, video_grid_thw, return_dict=True).pooler_output video_embeds = torch.cat(video_embeds, dim=0).to(inputs_embeds.device, inputs_embeds.dtype) _, video_mask = self.get_placeholder_mask( input_ids, inputs_embeds=inputs_embeds, video_features=video_embeds @@ -1599,13 +1609,37 @@ def get_input_embeddings(self): def set_input_embeddings(self, value): self.model.set_input_embeddings(value) + @auto_docstring def get_video_features( - self, pixel_values_videos: torch.FloatTensor, video_grid_thw: torch.LongTensor | None = None - ): - return self.model.get_video_features(pixel_values_videos, video_grid_thw) + self, + pixel_values_videos: torch.FloatTensor, + video_grid_thw: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + pixel_values_videos (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): + The tensors corresponding to the input videos. + video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*): + The temporal, height and width of feature shape of each video in LLM. + """ + return self.model.get_video_features( + pixel_values_videos=pixel_values_videos, video_grid_thw=video_grid_thw, **kwargs + ) - def get_image_features(self, pixel_values: torch.FloatTensor, image_grid_thw: torch.LongTensor | None = None): - return self.model.get_image_features(pixel_values, image_grid_thw) + @auto_docstring + def get_image_features( + self, + pixel_values: torch.FloatTensor, + image_grid_thw: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): + The tensors corresponding to the input images. + image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): + The temporal, height and width of feature shape of each image in LLM. + """ + return self.model.get_image_features(pixel_values=pixel_values, image_grid_thw=image_grid_thw, **kwargs) @auto_docstring @can_return_tuple diff --git a/src/transformers/models/ernie4_5_vl_moe/modular_ernie4_5_vl_moe.py b/src/transformers/models/ernie4_5_vl_moe/modular_ernie4_5_vl_moe.py index 09f67b989fcc..1093d0d9edc9 100644 --- a/src/transformers/models/ernie4_5_vl_moe/modular_ernie4_5_vl_moe.py +++ b/src/transformers/models/ernie4_5_vl_moe/modular_ernie4_5_vl_moe.py @@ -48,7 +48,7 @@ from ...masking_utils import create_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer -from ...modeling_outputs import MoeCausalLMOutputWithPast, MoeModelOutputWithPast +from ...modeling_outputs import BaseModelOutputWithPooling, MoeCausalLMOutputWithPast, MoeModelOutputWithPast from ...modeling_rope_utils import dynamic_rope_update from ...modeling_utils import PreTrainedModel from ...processing_utils import Unpack @@ -80,6 +80,7 @@ Qwen2_5_VisionRotaryEmbedding, Qwen2_5_VLModel, Qwen2_5_VLPreTrainedModel, + Qwen2_5_VLVisionAttention, Qwen2_5_VLVisionBlock, ) from ..qwen2_vl.configuration_qwen2_vl import Qwen2VLVisionConfig @@ -699,6 +700,23 @@ def forward( return hidden_states +class Ernie4_5_VL_MoeVisionAttention(Qwen2_5_VLVisionAttention): + pass + + +class Ernie4_5_VL_MoeVisionBlock(Qwen2_5_VLVisionBlock): + def __init__(self, config) -> None: + super().__init__(config, None) + + self.norm1 = nn.LayerNorm(config.hidden_size, config.rms_norm_eps) + self.norm2 = nn.LayerNorm(config.hidden_size, config.rms_norm_eps) + self.mlp = Ernie4_5VLVisionMLP( + dim=config.hidden_size, + hidden_dim=config.intermediate_size, + hidden_act=config.hidden_act, + ) + + class Ernie4_5_VL_MoePreTrainedModel(Qwen2_5_VLPreTrainedModel): _can_compile_fullgraph = False @@ -841,20 +859,13 @@ class Ernie4_5_VL_MoeVisionRotaryEmbedding(Qwen2_5_VisionRotaryEmbedding): pass -class Ernie4_5_VL_MoeVisionBlock(Qwen2_5_VLVisionBlock): - def __init__(self, config) -> None: - super().__init__(config, None) - - self.norm1 = nn.LayerNorm(config.hidden_size, config.rms_norm_eps) - self.norm2 = nn.LayerNorm(config.hidden_size, config.rms_norm_eps) - self.mlp = Ernie4_5VLVisionMLP( - dim=config.hidden_size, - hidden_dim=config.intermediate_size, - hidden_act=config.hidden_act, - ) - - class Ernie4_5_VL_MoeVisionTransformerPretrainedModel(Qwen2VisionTransformerPretrainedModel): + _can_record_outputs = { + "router_logits": OutputRecorder(Ernie4_5_VL_MoeMoeBlock, index=1), + "hidden_states": Ernie4_5_VL_MoeVisionBlock, + "attentions": Ernie4_5_VL_MoeVisionAttention, + } + def __init__(self, config) -> None: super().__init__(config) @@ -877,12 +888,10 @@ def get_dtype(self): def get_device(self): raise AttributeError("Ernie 4.5 VL Moe does not need this!") + @check_model_inputs def forward( - self, - hidden_states: torch.Tensor, - grid_thw: torch.Tensor, - **kwargs, - ) -> torch.Tensor: + self, hidden_states: torch.Tensor, grid_thw: torch.Tensor, **kwargs: Unpack[TransformersKwargs] + ) -> tuple | BaseModelOutputWithPooling: hidden_states = self.patch_embed(hidden_states) rotary_pos_emb = self.rot_pos_emb(grid_thw) emb = torch.cat((rotary_pos_emb, rotary_pos_emb), dim=-1) @@ -906,7 +915,7 @@ def forward( **kwargs, ) hidden_states = self.ln(hidden_states) - return hidden_states + return BaseModelOutputWithPooling(last_hidden_state=hidden_states) class Ernie4_5_VL_MoeVisionMLP(nn.Module): @@ -1259,43 +1268,39 @@ def get_rope_index( return position_ids, mrope_position_deltas + @can_return_tuple + @auto_docstring def get_video_features( - self, pixel_values_videos: torch.FloatTensor, video_grid_thw: torch.LongTensor | None = None - ): - """ - Encodes videos into continuous embeddings that can be forwarded to the language model. - - Args: - pixel_values_videos (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): - The tensors corresponding to the input videos. - video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*): - The temporal, height and width of feature shape of each video in LLM. - """ - video_embeds = self.vision_tower(pixel_values_videos, video_grid_thw) - video_embeds = self.resampler_model(video_embeds, video_grid_thw) + self, + pixel_values_videos: torch.FloatTensor, + video_grid_thw: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + video_outputs = self.vision_tower(pixel_values_videos, video_grid_thw, return_dict=True, **kwargs) + video_embeds = self.resampler_model(video_outputs.last_hidden_state, video_grid_thw) split_sizes = ( video_grid_thw.prod(-1) // self.vision_tower.spatial_merge_size**2 // self.resampler_model.temporal_merge_size ).tolist() video_embeds = torch.split(video_embeds, split_sizes) - return video_embeds - - def get_image_features(self, pixel_values: torch.FloatTensor, image_grid_thw: torch.LongTensor | None = None): - """ - Encodes images into continuous embeddings that can be forwarded to the language model. + video_outputs.pooler_output = video_embeds + return video_outputs - Args: - pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): - The tensors corresponding to the input images. - image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): - The temporal, height and width of feature shape of each image in LLM. - """ - image_embeds = self.vision_tower(pixel_values, image_grid_thw) - image_embeds = self.resampler_model(image_embeds, image_grid_thw) + @can_return_tuple + @auto_docstring + def get_image_features( + self, + pixel_values: torch.FloatTensor, + image_grid_thw: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + image_outputs = self.vision_tower(pixel_values, image_grid_thw, return_dict=True, **kwargs) + image_embeds = self.resampler_model(image_outputs.last_hidden_state, image_grid_thw) split_sizes = (image_grid_thw.prod(-1) // self.vision_tower.spatial_merge_size**2).tolist() image_embeds = torch.split(image_embeds, split_sizes) - return image_embeds + image_outputs.pooler_output = image_embeds + return image_outputs @auto_docstring @can_return_tuple @@ -1333,7 +1338,7 @@ def forward( inputs_embeds = self.get_input_embeddings()(input_ids) if pixel_values is not None: - image_embeds = self.get_image_features(pixel_values, image_grid_thw) + image_embeds = self.get_image_features(pixel_values, image_grid_thw, return_dict=True).pooler_output image_embeds = torch.cat(image_embeds, dim=0).to(inputs_embeds.device, inputs_embeds.dtype) image_mask, _ = self.get_placeholder_mask( input_ids, inputs_embeds=inputs_embeds, image_features=image_embeds @@ -1341,7 +1346,7 @@ def forward( inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds) if pixel_values_videos is not None: - video_embeds = self.get_video_features(pixel_values_videos, video_grid_thw) + video_embeds = self.get_video_features(pixel_values_videos, video_grid_thw, return_dict=True).pooler_output video_embeds = torch.cat(video_embeds, dim=0).to(inputs_embeds.device, inputs_embeds.dtype) _, video_mask = self.get_placeholder_mask( input_ids, inputs_embeds=inputs_embeds, video_features=video_embeds @@ -1392,6 +1397,14 @@ def __init__(self, config): self.num_experts = config.text_config.moe_num_experts self.num_experts_per_tok = config.text_config.moe_k + @auto_docstring + def get_video_features(self, **super_kwargs): + return super().get_video_features(**super_kwargs) + + @auto_docstring + def get_image_features(self, **super_kwargs): + return super().get_image_features(**super_kwargs) + def prepare_inputs_for_generation( self, input_ids, diff --git a/src/transformers/models/fast_vlm/configuration_fast_vlm.py b/src/transformers/models/fast_vlm/configuration_fast_vlm.py index c3ed8ae7f1d1..46e5a6ccbf76 100644 --- a/src/transformers/models/fast_vlm/configuration_fast_vlm.py +++ b/src/transformers/models/fast_vlm/configuration_fast_vlm.py @@ -18,7 +18,6 @@ # See the License for the specific language governing permissions and # limitations under the License. - from ...configuration_utils import PreTrainedConfig from ..auto import CONFIG_MAPPING, AutoConfig diff --git a/src/transformers/models/fast_vlm/modeling_fast_vlm.py b/src/transformers/models/fast_vlm/modeling_fast_vlm.py index 19b1587d277b..e2e2f2bb90b7 100644 --- a/src/transformers/models/fast_vlm/modeling_fast_vlm.py +++ b/src/transformers/models/fast_vlm/modeling_fast_vlm.py @@ -18,7 +18,6 @@ # See the License for the specific language governing permissions and # limitations under the License. - from dataclasses import dataclass import torch @@ -27,10 +26,11 @@ from ...activations import ACT2FN from ...cache_utils import Cache from ...generation import GenerationMixin -from ...modeling_outputs import BaseModelOutputWithPast, ModelOutput +from ...modeling_outputs import BaseModelOutputWithPast, BaseModelOutputWithPooling, ModelOutput from ...modeling_utils import PreTrainedModel from ...processing_utils import Unpack -from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, torch_compilable_check +from ...utils import TransformersKwargs, auto_docstring, torch_compilable_check +from ...utils.generic import check_model_inputs from ..auto import AutoModel from .configuration_fast_vlm import FastVlmConfig @@ -114,45 +114,36 @@ def get_input_embeddings(self): def set_input_embeddings(self, value): self.language_model.set_input_embeddings(value) + @check_model_inputs(tie_last_hidden_states=False) + @auto_docstring( + custom_intro="Obtains image last hidden states from the vision tower and apply multimodal projection." + ) def get_image_features( self, pixel_values: torch.FloatTensor, vision_feature_layer: int | list[int] | None = None, vision_feature_select_strategy: str | None = None, - **kwargs, - ): - """ - Obtains image last hidden states from the vision tower and apply multimodal projection. - - Args: - pixel_values (`torch.FloatTensor]` of shape `(batch_size, channels, height, width)`): - The tensors corresponding to the input images. - vision_feature_layer (`Union[int, list[int]]`, *optional*): - The index/indices of the layer to select the vision feature. Only -1 supported. - vision_feature_select_strategy (`str`, *optional*): - The feature selection strategy used to select the vision feature from the vision backbone. - Only "full" supported. - Returns: - image_features (`torch.Tensor`): Image feature tensor of shape `(num_images, image_length, embed_dim)`). + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + pixel_values (`torch.FloatTensor]` of shape `(batch_size, channels, height, width)`): + The tensors corresponding to the input images. + vision_feature_layer (`Union[int, list[int]]`, *optional*): + The index/indices of the layer to select the vision feature. Only -1 supported. + vision_feature_select_strategy (`str`, *optional*): + The feature selection strategy used to select the vision feature from the vision backbone. + Only "full" supported. """ - vision_feature_layer = ( - vision_feature_layer if vision_feature_layer is not None else self.config.vision_feature_layer - ) - vision_feature_select_strategy = ( - vision_feature_select_strategy - if vision_feature_select_strategy is not None - else self.config.vision_feature_select_strategy - ) - kwargs = {k: v for k, v in kwargs.items() if v is not None} - image_outputs = self.vision_tower(pixel_values, **kwargs) + image_outputs = self.vision_tower(pixel_values, return_dict=True, **kwargs) # since the vision tower is hybrid in FastVLM, its output needs to be handled differently from Llava selected_image_feature = image_outputs.last_hidden_state selected_image_feature = selected_image_feature.flatten(2).permute(0, 2, 1) image_features = self.multi_modal_projector(selected_image_feature) - image_features = list(image_features) - return image_features + image_outputs.pooler_output = list(image_features) + + return image_outputs def get_placeholder_mask( self, input_ids: torch.LongTensor, inputs_embeds: torch.FloatTensor, image_features: torch.FloatTensor @@ -178,7 +169,7 @@ def get_placeholder_mask( ) return special_image_mask - @can_return_tuple + @check_model_inputs(tie_last_hidden_states=False) @auto_docstring def forward( self, @@ -191,7 +182,6 @@ def forward( vision_feature_layer: int | list[int] | None = None, vision_feature_select_strategy: str | None = None, cache_position: torch.LongTensor | None = None, - image_sizes: torch.Tensor | None = None, **kwargs: Unpack[TransformersKwargs], ) -> tuple | FastVlmModelOutputWithPast: r""" @@ -201,15 +191,6 @@ def forward( vision_feature_select_strategy (`str`, *optional*): The feature selection strategy used to select the vision feature from the vision backbone. Only "full" supported. """ - vision_feature_layer = ( - vision_feature_layer if vision_feature_layer is not None else self.config.vision_feature_layer - ) - vision_feature_select_strategy = ( - vision_feature_select_strategy - if vision_feature_select_strategy is not None - else self.config.vision_feature_select_strategy - ) - if (input_ids is None) ^ (inputs_embeds is not None): raise ValueError("You must specify exactly one of input_ids or inputs_embeds") @@ -221,8 +202,8 @@ def forward( pixel_values=pixel_values, vision_feature_layer=vision_feature_layer, vision_feature_select_strategy=vision_feature_select_strategy, - image_sizes=image_sizes, - ) + return_dict=True, + ).pooler_output image_features = torch.cat(image_features, dim=0).to(inputs_embeds.device, inputs_embeds.dtype) special_image_mask = self.get_placeholder_mask( input_ids, inputs_embeds=inputs_embeds, image_features=image_features @@ -301,13 +282,14 @@ def set_input_embeddings(self, value): def get_output_embeddings(self) -> nn.Module: return self.lm_head + @auto_docstring def get_image_features( self, pixel_values: torch.FloatTensor, vision_feature_layer: int | list[int] | None = None, vision_feature_select_strategy: str | None = None, - **kwargs, - ): + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: return self.model.get_image_features( pixel_values=pixel_values, vision_feature_layer=vision_feature_layer, @@ -315,7 +297,7 @@ def get_image_features( **kwargs, ) - @can_return_tuple + @check_model_inputs(tie_last_hidden_states=False) @auto_docstring def forward( self, @@ -330,7 +312,6 @@ def forward( labels: torch.LongTensor | None = None, cache_position: torch.LongTensor | None = None, logits_to_keep: int | torch.Tensor = 0, - image_sizes: torch.Tensor | None = None, **kwargs: Unpack[TransformersKwargs], ) -> tuple | FastVlmCausalLMOutputWithPast: r""" @@ -378,15 +359,6 @@ def forward( >>> print(processor.batch_decode(generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]) system\n You are a helpful assistant.\n user\n What are these?\n assistant\n The image depicts a traditional Chinese street... ```""" - vision_feature_layer = ( - vision_feature_layer if vision_feature_layer is not None else self.config.vision_feature_layer - ) - vision_feature_select_strategy = ( - vision_feature_select_strategy - if vision_feature_select_strategy is not None - else self.config.vision_feature_select_strategy - ) - outputs = self.model( input_ids=input_ids, pixel_values=pixel_values, @@ -397,7 +369,6 @@ def forward( vision_feature_layer=vision_feature_layer, vision_feature_select_strategy=vision_feature_select_strategy, cache_position=cache_position, - image_sizes=image_sizes, **kwargs, ) diff --git a/src/transformers/models/fast_vlm/modular_fast_vlm.py b/src/transformers/models/fast_vlm/modular_fast_vlm.py index 2b68619bc726..fbe891b334fd 100644 --- a/src/transformers/models/fast_vlm/modular_fast_vlm.py +++ b/src/transformers/models/fast_vlm/modular_fast_vlm.py @@ -12,18 +12,23 @@ # See the License for the specific language governing permissions and # limitations under the License. - import torch from torch import nn from ...activations import ACT2FN +from ...cache_utils import Cache from ...configuration_utils import PreTrainedConfig -from ...utils import auto_docstring +from ...modeling_outputs import BaseModelOutputWithPooling +from ...processing_utils import Unpack +from ...utils import TransformersKwargs, auto_docstring +from ...utils.generic import check_model_inputs from ..auto import CONFIG_MAPPING from ..llava.configuration_llava import LlavaConfig from ..llava.modeling_llava import ( + LlavaCausalLMOutputWithPast, LlavaForConditionalGeneration, LlavaModel, + LlavaModelOutputWithPast, LlavaMultiModalProjector, LlavaPreTrainedModel, ) @@ -155,53 +160,62 @@ class FastVlmPreTrainedModel(LlavaPreTrainedModel): pass +class FastVlmModelOutputWithPast(LlavaModelOutputWithPast): + pass + + class FastVlmModel(LlavaModel): _checkpoint_conversion_mapping = {} def __init__(self, config: FastVlmConfig): super().__init__(config) + @check_model_inputs(tie_last_hidden_states=False) + @auto_docstring( + custom_intro="Obtains image last hidden states from the vision tower and apply multimodal projection." + ) def get_image_features( self, pixel_values: torch.FloatTensor, vision_feature_layer: int | list[int] | None = None, vision_feature_select_strategy: str | None = None, - **kwargs, - ): - """ - Obtains image last hidden states from the vision tower and apply multimodal projection. - - Args: - pixel_values (`torch.FloatTensor]` of shape `(batch_size, channels, height, width)`): - The tensors corresponding to the input images. - vision_feature_layer (`Union[int, list[int]]`, *optional*): - The index/indices of the layer to select the vision feature. Only -1 supported. - vision_feature_select_strategy (`str`, *optional*): - The feature selection strategy used to select the vision feature from the vision backbone. - Only "full" supported. - Returns: - image_features (`torch.Tensor`): Image feature tensor of shape `(num_images, image_length, embed_dim)`). + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + pixel_values (`torch.FloatTensor]` of shape `(batch_size, channels, height, width)`): + The tensors corresponding to the input images. + vision_feature_layer (`Union[int, list[int]]`, *optional*): + The index/indices of the layer to select the vision feature. Only -1 supported. + vision_feature_select_strategy (`str`, *optional*): + The feature selection strategy used to select the vision feature from the vision backbone. + Only "full" supported. """ - vision_feature_layer = ( - vision_feature_layer if vision_feature_layer is not None else self.config.vision_feature_layer - ) - vision_feature_select_strategy = ( - vision_feature_select_strategy - if vision_feature_select_strategy is not None - else self.config.vision_feature_select_strategy - ) - kwargs = {k: v for k, v in kwargs.items() if v is not None} - image_outputs = self.vision_tower(pixel_values, **kwargs) + image_outputs = self.vision_tower(pixel_values, return_dict=True, **kwargs) # since the vision tower is hybrid in FastVLM, its output needs to be handled differently from Llava selected_image_feature = image_outputs.last_hidden_state selected_image_feature = selected_image_feature.flatten(2).permute(0, 2, 1) image_features = self.multi_modal_projector(selected_image_feature) - image_features = list(image_features) - return image_features + image_outputs.pooler_output = list(image_features) + + return image_outputs - def forward(self, **super_kwargs): + @check_model_inputs(tie_last_hidden_states=False) + @auto_docstring + def forward( + self, + input_ids: torch.LongTensor | None = None, + pixel_values: torch.FloatTensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + inputs_embeds: torch.FloatTensor | None = None, + vision_feature_layer: int | list[int] | None = None, + vision_feature_select_strategy: str | None = None, + cache_position: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | FastVlmModelOutputWithPast: r""" vision_feature_layer (`Union[int, list[int], NoneType]`, *optional*): The index of the layer to select the vision feature. If multiple indices are provided, the vision feature of the @@ -209,7 +223,45 @@ def forward(self, **super_kwargs): vision_feature_select_strategy (`str`, *optional*): The feature selection strategy used to select the vision feature from the vision backbone. Only "full" supported. """ - super().forward(**super_kwargs) + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") + + if inputs_embeds is None: + inputs_embeds = self.get_input_embeddings()(input_ids) + + if pixel_values is not None: + image_features = self.get_image_features( + pixel_values=pixel_values, + vision_feature_layer=vision_feature_layer, + vision_feature_select_strategy=vision_feature_select_strategy, + return_dict=True, + ).pooler_output + image_features = torch.cat(image_features, dim=0).to(inputs_embeds.device, inputs_embeds.dtype) + special_image_mask = self.get_placeholder_mask( + input_ids, inputs_embeds=inputs_embeds, image_features=image_features + ) + inputs_embeds = inputs_embeds.masked_scatter(special_image_mask, image_features) + + outputs = self.language_model( + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + cache_position=cache_position, + **kwargs, + ) + + return FastVlmModelOutputWithPast( + last_hidden_state=outputs.last_hidden_state, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + image_hidden_states=image_features if pixel_values is not None else None, + ) + + +class FastVlmCausalLMOutputWithPast(LlavaCausalLMOutputWithPast): + pass @auto_docstring( @@ -220,7 +272,23 @@ def forward(self, **super_kwargs): class FastVlmForConditionalGeneration(LlavaForConditionalGeneration): _checkpoint_conversion_mapping = {} - def forward(self, **super_kwargs): + @check_model_inputs(tie_last_hidden_states=False) + @auto_docstring + def forward( + self, + input_ids: torch.LongTensor | None = None, + pixel_values: torch.FloatTensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + inputs_embeds: torch.FloatTensor | None = None, + vision_feature_layer: int | list[int] | None = None, + vision_feature_select_strategy: str | None = None, + labels: torch.LongTensor | None = None, + cache_position: torch.LongTensor | None = None, + logits_to_keep: int | torch.Tensor = 0, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | FastVlmCausalLMOutputWithPast: r""" vision_feature_layer (`Union[int, list[int], NoneType]`, *optional*): The index of the layer to select the vision feature. If multiple indices are provided, the vision feature of the @@ -266,7 +334,38 @@ def forward(self, **super_kwargs): >>> print(processor.batch_decode(generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]) system\n You are a helpful assistant.\n user\n What are these?\n assistant\n The image depicts a traditional Chinese street... ```""" - super().forward(**super_kwargs) + outputs = self.model( + input_ids=input_ids, + pixel_values=pixel_values, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + vision_feature_layer=vision_feature_layer, + vision_feature_select_strategy=vision_feature_select_strategy, + cache_position=cache_position, + **kwargs, + ) + + hidden_states = outputs[0] + # Only compute necessary logits, and do not upcast them to float if we are not computing the loss + slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep + logits = self.lm_head(hidden_states[:, slice_indices, :]) + + loss = None + if labels is not None: + loss = self.loss_function( + logits=logits, labels=labels, vocab_size=self.config.text_config.vocab_size, **kwargs + ) + + return FastVlmCausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + image_hidden_states=outputs.image_hidden_states, + ) __all__ = ["FastVlmForConditionalGeneration", "FastVlmModel", "FastVlmPreTrainedModel", "FastVlmConfig"] diff --git a/src/transformers/models/flava/modeling_flava.py b/src/transformers/models/flava/modeling_flava.py index e24b37abdd6b..87f7766e6631 100644 --- a/src/transformers/models/flava/modeling_flava.py +++ b/src/transformers/models/flava/modeling_flava.py @@ -27,7 +27,8 @@ from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling from ...modeling_utils import PreTrainedModel -from ...utils import ModelOutput, auto_docstring, filter_out_non_signature_kwargs, logging, torch_int +from ...processing_utils import Unpack +from ...utils import ModelOutput, TransformersKwargs, auto_docstring, can_return_tuple, logging, torch_int from .configuration_flava import ( FlavaConfig, FlavaImageCodebookConfig, @@ -997,7 +998,7 @@ def __init__(self, config: FlavaConfig): # Initialize weights and apply final processing self.post_init() - @filter_out_non_signature_kwargs() + @can_return_tuple @auto_docstring def get_text_features( self, @@ -1005,7 +1006,8 @@ def get_text_features( attention_mask: torch.Tensor | None = None, token_type_ids: torch.Tensor | None = None, position_ids: torch.Tensor | None = None, - ) -> torch.FloatTensor: + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: r""" input_ids (`torch.LongTensor` of shape `(batch_size, text_seq_length)`): Indices of input sequence tokens in the vocabulary. Indices can be obtained using [`AutoTokenizer`]. See @@ -1018,10 +1020,6 @@ def get_text_features( - 1 corresponds to a *sentence B* token. [What are token type IDs?](../glossary#token-type-ids) - Returns: - text_features (`torch.FloatTensor` of shape `(batch_size, output_dim`): The text embeddings obtained by - applying the projection layer to the pooled output of [`FlavaTextModel`]. - Examples: ```python @@ -1043,13 +1041,15 @@ def get_text_features( attention_mask=attention_mask, token_type_ids=token_type_ids, position_ids=position_ids, + return_dict=True, + **kwargs, ) - pooled_output = text_outputs.last_hidden_state - text_features = self.text_projection(pooled_output) + last_hidden_state = text_outputs.last_hidden_state + text_outputs.pooler_output = self.text_projection(last_hidden_state) - return text_features + return text_outputs - @filter_out_non_signature_kwargs() + @can_return_tuple @auto_docstring def get_image_features( self, @@ -1057,15 +1057,12 @@ def get_image_features( bool_masked_pos: torch.BoolTensor | None = None, interpolate_pos_encoding: bool | None = None, attention_mask: torch.Tensor | None = None, - ) -> torch.FloatTensor: + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: r""" bool_masked_pos (`torch.BoolTensor` of shape `(batch_size, image_num_patches)`): Boolean masked positions. Indicates which patches are masked (1) and which aren't (0). - Returns: - image_features (`torch.FloatTensor` of shape `(batch_size, output_dim`): The image embeddings obtained by - applying the projection layer to the pooled output of [`FlavaImageModel`]. - Examples: ```python @@ -1090,11 +1087,13 @@ def get_image_features( bool_masked_pos=bool_masked_pos, attention_mask=attention_mask, interpolate_pos_encoding=interpolate_pos_encoding, + return_dict=True, + **kwargs, ) - pooled_output = image_outputs.last_hidden_state - image_features = self.image_projection(pooled_output) + last_hidden_state = image_outputs.last_hidden_state + image_outputs.pooler_output = self.image_projection(last_hidden_state) - return image_features + return image_outputs @auto_docstring def forward( diff --git a/src/transformers/models/florence2/modeling_florence2.py b/src/transformers/models/florence2/modeling_florence2.py index 21ffca5bacda..f4fe735ce4f5 100644 --- a/src/transformers/models/florence2/modeling_florence2.py +++ b/src/transformers/models/florence2/modeling_florence2.py @@ -29,7 +29,7 @@ from ...activations import ACT2FN from ...cache_utils import Cache from ...generation import GenerationMixin -from ...modeling_outputs import Seq2SeqLMOutput, Seq2SeqModelOutput +from ...modeling_outputs import BaseModelOutputWithPooling, Seq2SeqLMOutput, Seq2SeqModelOutput from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel from ...processing_utils import Unpack from ...utils import ( @@ -40,6 +40,7 @@ logging, torch_compilable_check, ) +from ...utils.generic import check_model_inputs from ..auto import AutoModel from .configuration_florence2 import Florence2Config, Florence2VisionConfig @@ -498,6 +499,10 @@ class Florence2VisionPreTrainedModel(PreTrainedModel): _supports_flex_attn = True _can_compile_fullgraph = True + _can_record_outputs = { + "hidden_states": Florence2VisionBlock, + "attentions": [Florence2VisionChannelAttention, Florence2VisionWindowAttention], + } @auto_docstring @@ -548,12 +553,18 @@ def __init__(self, config: Florence2VisionConfig): # Initialize weights and apply final processing self.post_init() - def forward(self, hidden_states: torch.Tensor, **kwargs): + @check_model_inputs + def forward( + self, hidden_states: torch.Tensor, **kwargs: Unpack[TransformersKwargs] + ) -> tuple | BaseModelOutputWithPooling: for conv, block in zip(self.convs, self.blocks): hidden_states = conv(hidden_states) for layer in block: hidden_states = layer(hidden_states) - return hidden_states + + return BaseModelOutputWithPooling( + last_hidden_state=hidden_states, + ) class Florence2MultiModalProjector(nn.Module): @@ -671,19 +682,21 @@ def get_input_embeddings(self): def set_input_embeddings(self, value): self.language_model.set_input_embeddings(value) - def get_image_features(self, pixel_values: torch.Tensor, **kwargs): + @can_return_tuple + @auto_docstring( + custom_intro="Obtains image last hidden states from the vision tower and apply multimodal projection." + ) + def get_image_features( + self, pixel_values: torch.Tensor, **kwargs: Unpack[TransformersKwargs] + ) -> tuple | BaseModelOutputWithPooling: + r""" + pixel_values (`torch.FloatTensor]` of shape `(batch_size, channels, height, width)`): + The tensors corresponding to the input images. """ - Obtains image last hidden states from the vision tower and apply multimodal projection. + image_outputs = self.vision_tower(pixel_values, return_dict=True, **kwargs) + image_outputs.pooler_output = self.multi_modal_projector(image_outputs.last_hidden_state) - Args: - pixel_values (`torch.FloatTensor]` of shape `(batch_size, channels, height, width)`): - The tensors corresponding to the input images. - Returns: - image_features (`torch.Tensor`): Image feature tensor of shape `(num_images, image_length, embed_dim)`). - """ - image_features = self.vision_tower(pixel_values, **kwargs) - image_embeds = self.multi_modal_projector(image_features) - return image_embeds + return image_outputs def get_placeholder_mask( self, input_ids: torch.LongTensor, inputs_embeds: torch.FloatTensor, image_features: torch.FloatTensor @@ -743,7 +756,7 @@ def forward( inputs_embeds = self.get_input_embeddings()(input_ids) if pixel_values is not None: - image_features = self.get_image_features(pixel_values) + image_features = self.get_image_features(pixel_values, return_dict=True).pooler_output image_features = image_features.to(inputs_embeds.device, inputs_embeds.dtype) special_image_mask = self.get_placeholder_mask( input_ids, inputs_embeds=inputs_embeds, image_features=image_features @@ -838,7 +851,10 @@ def set_input_embeddings(self, value): def get_output_embeddings(self) -> nn.Module: return self.lm_head - def get_image_features(self, pixel_values: torch.Tensor, **kwargs): + @auto_docstring + def get_image_features( + self, pixel_values: torch.Tensor, **kwargs: Unpack[TransformersKwargs] + ) -> tuple | BaseModelOutputWithPooling: return self.model.get_image_features(pixel_values=pixel_values, **kwargs) @can_return_tuple @@ -920,7 +936,7 @@ def forward( output_hidden_states=output_hidden_states, return_dict=True, cache_position=cache_position, - # **kwargs, ## TODO: add back when Bart attention is refactored and takes kwargs + **kwargs, ) hidden_states = outputs[0] @@ -1003,7 +1019,7 @@ def _prepare_encoder_decoder_kwargs_for_generation( inputs_embeds = self.get_input_embeddings()(inputs_tensor) if pixel_values is not None: - image_features = self.get_image_features(pixel_values) + image_features = self.get_image_features(pixel_values, return_dict=True).pooler_output image_features = image_features.to(inputs_embeds.device, inputs_embeds.dtype) special_image_mask = self.get_placeholder_mask( inputs_tensor, inputs_embeds=inputs_embeds, image_features=image_features diff --git a/src/transformers/models/florence2/modular_florence2.py b/src/transformers/models/florence2/modular_florence2.py index dc6fa1755712..f94974769f7a 100644 --- a/src/transformers/models/florence2/modular_florence2.py +++ b/src/transformers/models/florence2/modular_florence2.py @@ -27,11 +27,12 @@ from ...configuration_utils import PreTrainedConfig from ...feature_extraction_utils import BatchFeature from ...image_utils import ImageInput -from ...modeling_outputs import Seq2SeqLMOutput, Seq2SeqModelOutput +from ...modeling_outputs import BaseModelOutputWithPooling, Seq2SeqLMOutput, Seq2SeqModelOutput from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel from ...processing_utils import MultiModalData, ProcessorMixin, Unpack from ...tokenization_utils_base import PreTokenizedInput, TextInput from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, is_torch_available, logging +from ...utils.generic import check_model_inputs from ..auto import CONFIG_MAPPING, AutoConfig from ..bart.modeling_bart import eager_attention_forward, shift_tokens_right from ..beit.modeling_beit import BeitDropPath @@ -1344,6 +1345,10 @@ class Florence2VisionPreTrainedModel(PreTrainedModel): _supports_flex_attn = True _can_compile_fullgraph = True + _can_record_outputs = { + "hidden_states": Florence2VisionBlock, + "attentions": [Florence2VisionChannelAttention, Florence2VisionWindowAttention], + } @auto_docstring @@ -1394,12 +1399,18 @@ def __init__(self, config: Florence2VisionConfig): # Initialize weights and apply final processing self.post_init() - def forward(self, hidden_states: torch.Tensor, **kwargs): + @check_model_inputs + def forward( + self, hidden_states: torch.Tensor, **kwargs: Unpack[TransformersKwargs] + ) -> tuple | BaseModelOutputWithPooling: for conv, block in zip(self.convs, self.blocks): hidden_states = conv(hidden_states) for layer in block: hidden_states = layer(hidden_states) - return hidden_states + + return BaseModelOutputWithPooling( + last_hidden_state=hidden_states, + ) class Florence2MultiModalProjector(nn.Module): @@ -1503,19 +1514,21 @@ def get_encoder(self, modality=None): else: return super().get_encoder(modality=modality) - def get_image_features(self, pixel_values: torch.Tensor, **kwargs): + @can_return_tuple + @auto_docstring( + custom_intro="Obtains image last hidden states from the vision tower and apply multimodal projection." + ) + def get_image_features( + self, pixel_values: torch.Tensor, **kwargs: Unpack[TransformersKwargs] + ) -> tuple | BaseModelOutputWithPooling: + r""" + pixel_values (`torch.FloatTensor]` of shape `(batch_size, channels, height, width)`): + The tensors corresponding to the input images. """ - Obtains image last hidden states from the vision tower and apply multimodal projection. + image_outputs = self.vision_tower(pixel_values, return_dict=True, **kwargs) + image_outputs.pooler_output = self.multi_modal_projector(image_outputs.last_hidden_state) - Args: - pixel_values (`torch.FloatTensor]` of shape `(batch_size, channels, height, width)`): - The tensors corresponding to the input images. - Returns: - image_features (`torch.Tensor`): Image feature tensor of shape `(num_images, image_length, embed_dim)`). - """ - image_features = self.vision_tower(pixel_values, **kwargs) - image_embeds = self.multi_modal_projector(image_features) - return image_embeds + return image_outputs @can_return_tuple @auto_docstring @@ -1551,7 +1564,7 @@ def forward( inputs_embeds = self.get_input_embeddings()(input_ids) if pixel_values is not None: - image_features = self.get_image_features(pixel_values) + image_features = self.get_image_features(pixel_values, return_dict=True).pooler_output image_features = image_features.to(inputs_embeds.device, inputs_embeds.dtype) special_image_mask = self.get_placeholder_mask( input_ids, inputs_embeds=inputs_embeds, image_features=image_features @@ -1609,7 +1622,10 @@ class Florence2ForConditionalGeneration(LlavaForConditionalGeneration): "lm_head.weight": "model.language_model.shared.weight", } - def get_image_features(self, pixel_values: torch.Tensor, **kwargs): + @auto_docstring + def get_image_features( + self, pixel_values: torch.Tensor, **kwargs: Unpack[TransformersKwargs] + ) -> tuple | BaseModelOutputWithPooling: return self.model.get_image_features(pixel_values=pixel_values, **kwargs) @can_return_tuple @@ -1691,7 +1707,7 @@ def forward( output_hidden_states=output_hidden_states, return_dict=True, cache_position=cache_position, - # **kwargs, ## TODO: add back when Bart attention is refactored and takes kwargs + **kwargs, ) hidden_states = outputs[0] @@ -1740,7 +1756,7 @@ def _prepare_encoder_decoder_kwargs_for_generation( inputs_embeds = self.get_input_embeddings()(inputs_tensor) if pixel_values is not None: - image_features = self.get_image_features(pixel_values) + image_features = self.get_image_features(pixel_values, return_dict=True).pooler_output image_features = image_features.to(inputs_embeds.device, inputs_embeds.dtype) special_image_mask = self.get_placeholder_mask( inputs_tensor, inputs_embeds=inputs_embeds, image_features=image_features diff --git a/src/transformers/models/fuyu/modeling_fuyu.py b/src/transformers/models/fuyu/modeling_fuyu.py index 2fae3597a54b..63f7be1f27b7 100644 --- a/src/transformers/models/fuyu/modeling_fuyu.py +++ b/src/transformers/models/fuyu/modeling_fuyu.py @@ -18,10 +18,11 @@ from ...cache_utils import Cache from ...generation import GenerationMixin -from ...modeling_outputs import CausalLMOutputWithPast +from ...modeling_outputs import BaseModelOutputWithPooling, CausalLMOutputWithPast from ...modeling_utils import PreTrainedModel from ...models.auto.modeling_auto import AutoModel -from ...utils import auto_docstring, can_return_tuple, logging, torch_compilable_check +from ...processing_utils import Unpack +from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, logging, torch_compilable_check from .configuration_fuyu import FuyuConfig @@ -113,15 +114,17 @@ def gather_continuous_embeddings( ) return output_embeddings - def get_image_features(self, pixel_values: torch.FloatTensor, **kwargs): - """ - Encodes images into continuous embeddings that can be forwarded to the language model. - - Args: - pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): - The tensors corresponding to the input images. + @can_return_tuple + @auto_docstring + def get_image_features( + self, pixel_values: torch.FloatTensor, **kwargs: Unpack[TransformersKwargs] + ) -> tuple | BaseModelOutputWithPooling: + r""" + pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): + The tensors corresponding to the input images. """ - return self.vision_embed_tokens(pixel_values) + patch_embeddings = self.vision_embed_tokens(pixel_values) + return BaseModelOutputWithPooling(last_hidden_state=patch_embeddings) def get_placeholder_mask( self, input_ids: torch.LongTensor, inputs_embeds: torch.FloatTensor, image_features: torch.FloatTensor @@ -200,7 +203,8 @@ def forward( inputs_embeds = self.language_model.get_input_embeddings()(input_ids) if image_patches is not None: - patch_embeddings = self.get_image_features(image_patches).to(inputs_embeds.device, inputs_embeds.dtype) + patch_embeddings = self.get_image_features(image_patches, return_dict=True).last_hidden_state + patch_embeddings = patch_embeddings.to(inputs_embeds.device, inputs_embeds.dtype) special_image_mask = self.get_placeholder_mask( input_ids, inputs_embeds=inputs_embeds, image_features=patch_embeddings ) diff --git a/src/transformers/models/gemma3/modeling_gemma3.py b/src/transformers/models/gemma3/modeling_gemma3.py index 2fc26eca6c6d..9893d7a0f9a8 100644 --- a/src/transformers/models/gemma3/modeling_gemma3.py +++ b/src/transformers/models/gemma3/modeling_gemma3.py @@ -33,7 +33,12 @@ from ...integrations import use_kernel_func_from_hub, use_kernelized_func from ...masking_utils import create_causal_mask, create_masks_for_generate, create_sliding_window_causal_mask from ...modeling_layers import GenericForSequenceClassification, GradientCheckpointingLayer -from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast, SequenceClassifierOutputWithPast +from ...modeling_outputs import ( + BaseModelOutputWithPast, + BaseModelOutputWithPooling, + CausalLMOutputWithPast, + SequenceClassifierOutputWithPast, +) from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel from ...processing_utils import Unpack @@ -834,19 +839,16 @@ def get_input_embeddings(self): def set_input_embeddings(self, value): self.language_model.set_input_embeddings(value) - def get_image_features(self, pixel_values: torch.Tensor) -> torch.Tensor: - """ - Projects the last hidden state from the vision model into language model space. + @can_return_tuple + @auto_docstring(custom_intro="Projects the last hidden state from the vision model into language model space.") + def get_image_features( + self, pixel_values: torch.FloatTensor, **kwargs: Unpack[TransformersKwargs] + ) -> tuple | BaseModelOutputWithPooling: + vision_outputs = self.vision_tower(pixel_values=pixel_values, return_dict=True, **kwargs) + last_hidden_state = vision_outputs.last_hidden_state + vision_outputs.pooler_output = self.multi_modal_projector(last_hidden_state) - Args: - pixel_values (`torch.FloatTensor]` of shape `(batch_size, channels, height, width)`) - The tensors corresponding to the input images. - Returns: - image_features (`torch.Tensor`): Image feature tensor of shape `(num_images, image_length, embed_dim)`). - """ - vision_outputs = self.vision_tower(pixel_values=pixel_values).last_hidden_state - image_features = self.multi_modal_projector(vision_outputs) - return image_features + return vision_outputs def get_placeholder_mask( self, input_ids: torch.LongTensor, inputs_embeds: torch.FloatTensor, image_features: torch.FloatTensor @@ -937,7 +939,7 @@ def forward( # Merge text and images if pixel_values is not None: - image_features = self.get_image_features(pixel_values) + image_features = self.get_image_features(pixel_values, return_dict=True).pooler_output image_features = image_features.to(inputs_embeds.device, inputs_embeds.dtype) special_image_mask = self.get_placeholder_mask( input_ids, inputs_embeds=inputs_embeds, image_features=image_features @@ -1007,8 +1009,9 @@ def get_input_embeddings(self): def set_input_embeddings(self, value): self.model.set_input_embeddings(value) - def get_image_features(self, pixel_values): - return self.model.get_image_features(pixel_values) + @auto_docstring + def get_image_features(self, pixel_values: torch.FloatTensor, **kwargs: Unpack[TransformersKwargs]): + return self.model.get_image_features(pixel_values, **kwargs) @can_return_tuple @auto_docstring diff --git a/src/transformers/models/gemma3/modular_gemma3.py b/src/transformers/models/gemma3/modular_gemma3.py index 39054ac5bc67..31424e010f76 100644 --- a/src/transformers/models/gemma3/modular_gemma3.py +++ b/src/transformers/models/gemma3/modular_gemma3.py @@ -23,7 +23,7 @@ from ...configuration_utils import PreTrainedConfig, layer_type_validation from ...masking_utils import create_causal_mask, create_masks_for_generate, create_sliding_window_causal_mask from ...modeling_layers import GenericForSequenceClassification, GradientCheckpointingLayer -from ...modeling_outputs import BaseModelOutputWithPast, SequenceClassifierOutputWithPast +from ...modeling_outputs import BaseModelOutputWithPast, BaseModelOutputWithPooling, SequenceClassifierOutputWithPast from ...modeling_rope_utils import ( ROPE_INIT_FUNCTIONS, RopeParameters, @@ -799,19 +799,16 @@ def __init__(self, config: Gemma3Config): super().__init__(config) del self.text_config_dtype - def get_image_features(self, pixel_values: torch.Tensor) -> torch.Tensor: - """ - Projects the last hidden state from the vision model into language model space. + @can_return_tuple + @auto_docstring(custom_intro="Projects the last hidden state from the vision model into language model space.") + def get_image_features( + self, pixel_values: torch.FloatTensor, **kwargs: Unpack[TransformersKwargs] + ) -> tuple | BaseModelOutputWithPooling: + vision_outputs = self.vision_tower(pixel_values=pixel_values, return_dict=True, **kwargs) + last_hidden_state = vision_outputs.last_hidden_state + vision_outputs.pooler_output = self.multi_modal_projector(last_hidden_state) - Args: - pixel_values (`torch.FloatTensor]` of shape `(batch_size, channels, height, width)`) - The tensors corresponding to the input images. - Returns: - image_features (`torch.Tensor`): Image feature tensor of shape `(num_images, image_length, embed_dim)`). - """ - vision_outputs = self.vision_tower(pixel_values=pixel_values).last_hidden_state - image_features = self.multi_modal_projector(vision_outputs) - return image_features + return vision_outputs @can_return_tuple @auto_docstring @@ -851,7 +848,7 @@ def forward( # Merge text and images if pixel_values is not None: - image_features = self.get_image_features(pixel_values) + image_features = self.get_image_features(pixel_values, return_dict=True).pooler_output image_features = image_features.to(inputs_embeds.device, inputs_embeds.dtype) special_image_mask = self.get_placeholder_mask( input_ids, inputs_embeds=inputs_embeds, image_features=image_features diff --git a/src/transformers/models/gemma3n/modeling_gemma3n.py b/src/transformers/models/gemma3n/modeling_gemma3n.py index 1815e49d09ae..500e770c3df1 100644 --- a/src/transformers/models/gemma3n/modeling_gemma3n.py +++ b/src/transformers/models/gemma3n/modeling_gemma3n.py @@ -34,7 +34,7 @@ from ...integrations import use_kernelized_func from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask from ...modeling_layers import GradientCheckpointingLayer -from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast +from ...modeling_outputs import BaseModelOutputWithPast, BaseModelOutputWithPooling, CausalLMOutputWithPast from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel from ...processing_utils import Unpack @@ -50,6 +50,17 @@ from .configuration_gemma3n import Gemma3nAudioConfig, Gemma3nConfig, Gemma3nTextConfig, Gemma3nVisionConfig +@dataclass +@auto_docstring +class Gemma3nAudioEncoderModelOutput(BaseModelOutputWithPooling): + """ + audio_mel_mask (`torch.FloatTensor`, *optional*): + A torch.BoolTensor of shape `(batch_size, num_frames)` + """ + + audio_mel_mask: torch.BoolTensor | None = None + + @dataclass @auto_docstring( custom_intro=""" @@ -930,9 +941,10 @@ def __init__(self, config: Gemma3nAudioConfig): ) self.post_init() + @check_model_inputs def forward( - self, audio_mel: torch.Tensor, audio_mel_mask: torch.BoolTensor, **kwargs - ) -> tuple[torch.Tensor, torch.BoolTensor]: + self, audio_mel: torch.Tensor, audio_mel_mask: torch.BoolTensor, **kwargs: Unpack[TransformersKwargs] + ) -> tuple | Gemma3nAudioEncoderModelOutput: """Encodes a batch of MELs. Args: @@ -983,7 +995,10 @@ def forward( current_mask = current_mask[:, :: self.config.conf_reduction_factor] audio_encodings = audio_encodings.masked_fill(current_mask.unsqueeze(-1), 0.0) - return audio_encodings, current_mask + return Gemma3nAudioEncoderModelOutput( + last_hidden_state=audio_encodings, + audio_mel_mask=current_mask, + ) class Gemma3nTextScaledWordEmbedding(nn.Embedding): @@ -1931,30 +1946,27 @@ def get_input_embeddings(self): def set_input_embeddings(self, value): self.language_model.set_input_embeddings(value) - def get_image_features(self, pixel_values: torch.Tensor) -> torch.Tensor: - """ - Projects the last hidden state from the vision model into language model space. - - Args: - pixel_values (`torch.FloatTensor]` of shape `(batch_size, channels, height, width)`) - The tensors corresponding to the input images. - - Returns: - image_features (`torch.Tensor`): Image feature tensor of shape `(num_images, image_length, embed_dim)`). - """ - vision_outputs = self.vision_tower( - pixel_values=pixel_values, do_pooling=False, return_dict=True - ).last_hidden_state + @can_return_tuple + @auto_docstring(custom_intro="Projects the last hidden state from the vision model into language model space.") + def get_image_features( + self, + pixel_values: torch.FloatTensor, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + vision_outputs = self.vision_tower(pixel_values=pixel_values, do_pooling=False, return_dict=True, **kwargs) + last_hidden_state = vision_outputs.last_hidden_state # Convert from (batch, channels, height, width) to (batch, height * width, channels) where: # height == width and height * width == Gemma3nConfig.vision_soft_tokens_per_image. - vision_outputs = vision_outputs.reshape( - vision_outputs.shape[0], + last_hidden_state = last_hidden_state.reshape( + last_hidden_state.shape[0], self.config.vision_config.hidden_size, self.config.vision_soft_tokens_per_image, ).permute(0, 2, 1) # Normalize and embed the soft tokens into language model space. - vision_outputs *= self.config.vision_config.hidden_size**0.5 - return self.embed_vision(inputs_embeds=vision_outputs) + last_hidden_state *= self.config.vision_config.hidden_size**0.5 + vision_outputs.pooler_output = self.embed_vision(inputs_embeds=last_hidden_state) + + return vision_outputs def get_placeholder_mask( self, @@ -2087,7 +2099,7 @@ def forward( # Merge text and images if pixel_values is not None: - image_features = self.get_image_features(pixel_values) + image_features = self.get_image_features(pixel_values, return_dict=True).pooler_output image_features = image_features.to(inputs_embeds.device, inputs_embeds.dtype) special_image_mask, _ = self.get_placeholder_mask( input_ids, inputs_embeds=inputs_embeds, image_features=image_features @@ -2096,7 +2108,9 @@ def forward( # Merge text and audio if input_features is not None and input_features_mask is not None: - audio_features, audio_mask = self.get_audio_features(input_features, ~input_features_mask) + audio_features = self.get_audio_features(input_features, ~input_features_mask, return_dict=True) + audio_features = audio_features.pooler_output + audio_mask = audio_features.audio_mel_mask # The Gemma3nProcessor expects all audio will be 30s in length and inserts 188 audio soft tokens into the # text to account for this. However, the audio preprocessing and encoder do not gurarantee they will @@ -2142,23 +2156,27 @@ def forward( audio_hidden_states=audio_features if input_features is not None else None, ) + @can_return_tuple + @auto_docstring(custom_intro="Projects the last hidden state from the audio encoder into language model space.") def get_audio_features( - self, input_features: torch.Tensor, input_features_mask: torch.Tensor - ) -> tuple[torch.Tensor, torch.Tensor]: + self, + input_features: torch.Tensor, + input_features_mask: torch.Tensor, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | Gemma3nAudioEncoderModelOutput: + r""" + input_features (`torch.FloatTensor]` of shape `(num_images, seq_length, num_features)`): + The tensors corresponding to the input audio. + input_features_mask (`torch.FloatTensor]` of shape `(num_images, seq_length)`): + The attention mask for the input audio. """ - Projects the last hidden state from the audio encoder into language model space. + audio_outputs: Gemma3nAudioEncoderModelOutput = self.audio_tower( + input_features, input_features_mask, return_dict=True, **kwargs + ) + audio_embeds = self.embed_audio(inputs_embeds=audio_outputs.last_hidden_state) + audio_outputs.pooler_output = audio_embeds - Args: - input_features (`torch.FloatTensor]` of shape `(num_images, seq_length, num_features)`): - The tensors corresponding to the input audio. - input_features_mask (`torch.FloatTensor]` of shape `(num_images, seq_length)`): - The attention mask for the input audio. - - Returns: - audio_features (`torch.Tensor`): Audio feature tensor of shape `(num_images, audio_length, embed_dim)`). - """ - audio_outputs, audio_mask = self.audio_tower(input_features, input_features_mask) - return self.embed_audio(inputs_embeds=audio_outputs), audio_mask + return audio_outputs @auto_docstring( @@ -2183,8 +2201,9 @@ def get_input_embeddings(self): def set_input_embeddings(self, value): self.model.set_input_embeddings(value) - def get_image_features(self, pixel_values): - return self.model.get_image_features(pixel_values) + @auto_docstring + def get_image_features(self, pixel_values: torch.FloatTensor, **kwargs: Unpack[TransformersKwargs]): + return self.model.get_image_features(pixel_values, **kwargs) @can_return_tuple @auto_docstring diff --git a/src/transformers/models/gemma3n/modular_gemma3n.py b/src/transformers/models/gemma3n/modular_gemma3n.py index d7265a44f92a..fd95a8b37366 100644 --- a/src/transformers/models/gemma3n/modular_gemma3n.py +++ b/src/transformers/models/gemma3n/modular_gemma3n.py @@ -14,6 +14,7 @@ # limitations under the License. import math from collections.abc import Callable, Sequence +from dataclasses import dataclass from typing import Any import torch @@ -25,7 +26,7 @@ from ...cache_utils import Cache, DynamicCache from ...configuration_utils import PreTrainedConfig, layer_type_validation from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask -from ...modeling_outputs import BaseModelOutputWithPast +from ...modeling_outputs import BaseModelOutputWithPast, BaseModelOutputWithPooling from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, RopeParameters from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel from ...processing_utils import Unpack @@ -630,6 +631,17 @@ def __init__( self.tie_word_embeddings = tie_word_embeddings +@dataclass +@auto_docstring +class Gemma3nAudioEncoderModelOutput(BaseModelOutputWithPooling): + """ + audio_mel_mask (`torch.FloatTensor`, *optional*): + A torch.BoolTensor of shape `(batch_size, num_frames)` + """ + + audio_mel_mask: torch.BoolTensor | None = None + + class Gemma3nModelOutputWithPast(PaligemmaModelOutputWithPast): r""" past_key_values (`Cache`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`): @@ -1486,9 +1498,10 @@ def __init__(self, config: Gemma3nAudioConfig): ) self.post_init() + @check_model_inputs def forward( - self, audio_mel: torch.Tensor, audio_mel_mask: torch.BoolTensor, **kwargs - ) -> tuple[torch.Tensor, torch.BoolTensor]: + self, audio_mel: torch.Tensor, audio_mel_mask: torch.BoolTensor, **kwargs: Unpack[TransformersKwargs] + ) -> tuple | Gemma3nAudioEncoderModelOutput: """Encodes a batch of MELs. Args: @@ -1539,7 +1552,10 @@ def forward( current_mask = current_mask[:, :: self.config.conf_reduction_factor] audio_encodings = audio_encodings.masked_fill(current_mask.unsqueeze(-1), 0.0) - return audio_encodings, current_mask + return Gemma3nAudioEncoderModelOutput( + last_hidden_state=audio_encodings, + audio_mel_mask=current_mask, + ) # ==== Language Model ==== @@ -2197,30 +2213,27 @@ def __init__(self, config: Gemma3nConfig): self.embed_vision = Gemma3nMultimodalEmbedder(config.vision_config, config.text_config) self.embed_audio = Gemma3nMultimodalEmbedder(config.audio_config, config.text_config) - def get_image_features(self, pixel_values: torch.Tensor) -> torch.Tensor: - """ - Projects the last hidden state from the vision model into language model space. - - Args: - pixel_values (`torch.FloatTensor]` of shape `(batch_size, channels, height, width)`) - The tensors corresponding to the input images. - - Returns: - image_features (`torch.Tensor`): Image feature tensor of shape `(num_images, image_length, embed_dim)`). - """ - vision_outputs = self.vision_tower( - pixel_values=pixel_values, do_pooling=False, return_dict=True - ).last_hidden_state + @can_return_tuple + @auto_docstring(custom_intro="Projects the last hidden state from the vision model into language model space.") + def get_image_features( + self, + pixel_values: torch.FloatTensor, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + vision_outputs = self.vision_tower(pixel_values=pixel_values, do_pooling=False, return_dict=True, **kwargs) + last_hidden_state = vision_outputs.last_hidden_state # Convert from (batch, channels, height, width) to (batch, height * width, channels) where: # height == width and height * width == Gemma3nConfig.vision_soft_tokens_per_image. - vision_outputs = vision_outputs.reshape( - vision_outputs.shape[0], + last_hidden_state = last_hidden_state.reshape( + last_hidden_state.shape[0], self.config.vision_config.hidden_size, self.config.vision_soft_tokens_per_image, ).permute(0, 2, 1) # Normalize and embed the soft tokens into language model space. - vision_outputs *= self.config.vision_config.hidden_size**0.5 - return self.embed_vision(inputs_embeds=vision_outputs) + last_hidden_state *= self.config.vision_config.hidden_size**0.5 + vision_outputs.pooler_output = self.embed_vision(inputs_embeds=last_hidden_state) + + return vision_outputs def get_placeholder_mask( self, @@ -2353,7 +2366,7 @@ def forward( # Merge text and images if pixel_values is not None: - image_features = self.get_image_features(pixel_values) + image_features = self.get_image_features(pixel_values, return_dict=True).pooler_output image_features = image_features.to(inputs_embeds.device, inputs_embeds.dtype) special_image_mask, _ = self.get_placeholder_mask( input_ids, inputs_embeds=inputs_embeds, image_features=image_features @@ -2362,7 +2375,9 @@ def forward( # Merge text and audio if input_features is not None and input_features_mask is not None: - audio_features, audio_mask = self.get_audio_features(input_features, ~input_features_mask) + audio_features = self.get_audio_features(input_features, ~input_features_mask, return_dict=True) + audio_features = audio_features.pooler_output + audio_mask = audio_features.audio_mel_mask # The Gemma3nProcessor expects all audio will be 30s in length and inserts 188 audio soft tokens into the # text to account for this. However, the audio preprocessing and encoder do not gurarantee they will @@ -2408,23 +2423,27 @@ def forward( audio_hidden_states=audio_features if input_features is not None else None, ) + @can_return_tuple + @auto_docstring(custom_intro="Projects the last hidden state from the audio encoder into language model space.") def get_audio_features( - self, input_features: torch.Tensor, input_features_mask: torch.Tensor - ) -> tuple[torch.Tensor, torch.Tensor]: + self, + input_features: torch.Tensor, + input_features_mask: torch.Tensor, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | Gemma3nAudioEncoderModelOutput: + r""" + input_features (`torch.FloatTensor]` of shape `(num_images, seq_length, num_features)`): + The tensors corresponding to the input audio. + input_features_mask (`torch.FloatTensor]` of shape `(num_images, seq_length)`): + The attention mask for the input audio. """ - Projects the last hidden state from the audio encoder into language model space. - - Args: - input_features (`torch.FloatTensor]` of shape `(num_images, seq_length, num_features)`): - The tensors corresponding to the input audio. - input_features_mask (`torch.FloatTensor]` of shape `(num_images, seq_length)`): - The attention mask for the input audio. + audio_outputs: Gemma3nAudioEncoderModelOutput = self.audio_tower( + input_features, input_features_mask, return_dict=True, **kwargs + ) + audio_embeds = self.embed_audio(inputs_embeds=audio_outputs.last_hidden_state) + audio_outputs.pooler_output = audio_embeds - Returns: - audio_features (`torch.Tensor`): Audio feature tensor of shape `(num_images, audio_length, embed_dim)`). - """ - audio_outputs, audio_mask = self.audio_tower(input_features, input_features_mask) - return self.embed_audio(inputs_embeds=audio_outputs), audio_mask + return audio_outputs @auto_docstring( diff --git a/src/transformers/models/glm46v/modeling_glm46v.py b/src/transformers/models/glm46v/modeling_glm46v.py index bc70cf8a26f1..de7b79923826 100644 --- a/src/transformers/models/glm46v/modeling_glm46v.py +++ b/src/transformers/models/glm46v/modeling_glm46v.py @@ -28,7 +28,7 @@ from ...cache_utils import Cache from ...generation import GenerationMixin -from ...modeling_outputs import ModelOutput +from ...modeling_outputs import BaseModelOutputWithPooling, ModelOutput from ...modeling_utils import PreTrainedModel from ...processing_utils import Unpack from ...utils import ( @@ -296,17 +296,19 @@ def get_rope_index( return position_ids, mrope_position_deltas + @can_return_tuple + @auto_docstring def get_video_features( - self, pixel_values_videos: torch.FloatTensor, video_grid_thw: torch.LongTensor | None = None - ): - """ - Encodes videos into continuous embeddings that can be forwarded to the language model. - - Args: - pixel_values_videos (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): - The tensors corresponding to the input videos. - video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*): - The temporal, height and width of feature shape of each video in LLM. + self, + pixel_values_videos: torch.FloatTensor, + video_grid_thw: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + pixel_values_videos (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): + The tensors corresponding to the input videos. + video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*): + The temporal, height and width of feature shape of each video in LLM. """ pixel_values_videos = pixel_values_videos.type(self.visual.dtype) # reshape video_grid_thw -> [b, 3] -> [1, h, w] * frames @@ -315,26 +317,36 @@ def get_video_features( repeated_row = torch.tensor([1, h.item(), w.item()]).unsqueeze(0).repeat(t, 1) temp_frames_hw.append(repeated_row) flattened_video_grid_thw = torch.cat(temp_frames_hw, dim=0) - video_embeds = self.visual(pixel_values_videos, grid_thw=flattened_video_grid_thw) + vision_outputs = self.visual( + pixel_values_videos, grid_thw=flattened_video_grid_thw, return_dict=True, **kwargs + ) split_sizes = (video_grid_thw.prod(-1) // self.visual.spatial_merge_size**2).tolist() - video_embeds = torch.split(video_embeds, split_sizes) - return video_embeds + video_embeds = torch.split(vision_outputs.pooler_output, split_sizes) + vision_outputs.pooler_output = video_embeds - def get_image_features(self, pixel_values: torch.FloatTensor, image_grid_thw: torch.LongTensor | None = None): - """ - Encodes images into continuous embeddings that can be forwarded to the language model. + return vision_outputs - Args: - pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): - The tensors corresponding to the input images. - image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): - The temporal, height and width of feature shape of each image in LLM. + @can_return_tuple + @auto_docstring + def get_image_features( + self, + pixel_values: torch.FloatTensor, + image_grid_thw: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): + The tensors corresponding to the input images. + image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): + The temporal, height and width of feature shape of each image in LLM. """ pixel_values = pixel_values.type(self.visual.dtype) - image_embeds = self.visual(pixel_values, grid_thw=image_grid_thw) + vision_outputs = self.visual(pixel_values, grid_thw=image_grid_thw, return_dict=True, **kwargs) split_sizes = (image_grid_thw.prod(-1) // self.visual.spatial_merge_size**2).tolist() - image_embeds = torch.split(image_embeds, split_sizes) - return image_embeds + image_embeds = torch.split(vision_outputs.pooler_output, split_sizes) + vision_outputs.pooler_output = image_embeds + + return vision_outputs def get_placeholder_mask( self, @@ -410,13 +422,13 @@ def forward( inputs_embeds = self.get_input_embeddings()(input_ids) if pixel_values is not None: - image_embeds = self.get_image_features(pixel_values, image_grid_thw) + image_embeds = self.get_image_features(pixel_values, image_grid_thw, return_dict=True).pooler_output image_embeds = torch.cat(image_embeds, dim=0).to(inputs_embeds.device, inputs_embeds.dtype) image_mask, _ = self.get_placeholder_mask(input_ids, inputs_embeds, image_features=image_embeds) inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds) if pixel_values_videos is not None: - video_embeds = self.get_video_features(pixel_values_videos, video_grid_thw) + video_embeds = self.get_video_features(pixel_values_videos, video_grid_thw, return_dict=True).pooler_output video_embeds = torch.cat(video_embeds, dim=0).to(inputs_embeds.device, inputs_embeds.dtype) _, video_mask = self.get_placeholder_mask(input_ids, inputs_embeds, video_features=video_embeds) inputs_embeds = inputs_embeds.masked_scatter(video_mask, video_embeds) @@ -534,13 +546,37 @@ def get_input_embeddings(self): def set_input_embeddings(self, value): self.model.set_input_embeddings(value) + @auto_docstring def get_video_features( - self, pixel_values_videos: torch.FloatTensor, video_grid_thw: torch.LongTensor | None = None - ): - return self.model.get_video_features(pixel_values_videos, video_grid_thw) + self, + pixel_values_videos: torch.FloatTensor, + video_grid_thw: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + pixel_values_videos (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): + The tensors corresponding to the input videos. + video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*): + The temporal, height and width of feature shape of each video in LLM. + """ + return self.model.get_video_features( + pixel_values_videos=pixel_values_videos, video_grid_thw=video_grid_thw, **kwargs + ) - def get_image_features(self, pixel_values: torch.FloatTensor, image_grid_thw: torch.LongTensor | None = None): - return self.model.get_image_features(pixel_values, image_grid_thw) + @auto_docstring + def get_image_features( + self, + pixel_values: torch.FloatTensor, + image_grid_thw: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): + The tensors corresponding to the input images. + image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): + The temporal, height and width of feature shape of each image in LLM. + """ + return self.model.get_image_features(pixel_values=pixel_values, image_grid_thw=image_grid_thw, **kwargs) @can_return_tuple @auto_docstring diff --git a/src/transformers/models/glm4v/modeling_glm4v.py b/src/transformers/models/glm4v/modeling_glm4v.py index a2c0dff04f99..fb6de9c68dc6 100644 --- a/src/transformers/models/glm4v/modeling_glm4v.py +++ b/src/transformers/models/glm4v/modeling_glm4v.py @@ -35,7 +35,7 @@ from ...masking_utils import create_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer -from ...modeling_outputs import BaseModelOutputWithPast, ModelOutput +from ...modeling_outputs import BaseModelOutputWithPast, BaseModelOutputWithPooling, ModelOutput from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel from ...processing_utils import Unpack @@ -701,6 +701,10 @@ class Glm4vVisionModel(Glm4vPreTrainedModel): config: Glm4vVisionConfig input_modalities = ("image", "video") _no_split_modules = ["Glm4vVisionBlock"] + _can_record_outputs = { + "hidden_states": Glm4vVisionBlock, + "attentions": Glm4vVisionAttention, + } def __init__(self, config) -> None: super().__init__(config) @@ -759,13 +763,16 @@ def rot_pos_emb(self, grid_thw): rotary_pos_emb = rotary_pos_emb_full[pos_ids].flatten(1) return rotary_pos_emb, pos_ids - def forward(self, hidden_states: torch.Tensor, grid_thw: torch.Tensor, **kwargs) -> torch.Tensor: - """ - Args: - hidden_states (`torch.Tensor` of shape `(seq_len, hidden_size)`): - The final hidden states of the model. - grid_thw (`torch.Tensor` of shape `(num_images_or_videos, 3)`): - The temporal, height and width of feature shape of each image in LLM. + @check_model_inputs + @auto_docstring + def forward( + self, hidden_states: torch.Tensor, grid_thw: torch.Tensor, **kwargs: Unpack[TransformersKwargs] + ) -> tuple | BaseModelOutputWithPooling: + r""" + hidden_states (`torch.Tensor` of shape `(seq_len, hidden_size)`): + The final hidden states of the model. + grid_thw (`torch.Tensor` of shape `(num_images_or_videos, 3)`): + The temporal, height and width of feature shape of each image in LLM. Returns: `torch.Tensor`: hidden_states. @@ -799,6 +806,7 @@ def forward(self, hidden_states: torch.Tensor, grid_thw: torch.Tensor, **kwargs) hidden_states, cu_seqlens=cu_seqlens, position_embeddings=position_embeddings, + **kwargs, ) hidden_states = self.post_layernorm(hidden_states) @@ -809,8 +817,12 @@ def forward(self, hidden_states: torch.Tensor, grid_thw: torch.Tensor, **kwargs) hidden_states = hidden_states.permute(0, 3, 1, 2) hidden_states = self.downsample(hidden_states).view(-1, self.config.out_hidden_size) - hidden_states = self.merger(hidden_states) - return hidden_states + merged_hidden_states = self.merger(hidden_states) + + return BaseModelOutputWithPooling( + last_hidden_state=hidden_states, + pooler_output=merged_hidden_states, + ) @auto_docstring @@ -1134,17 +1146,19 @@ def get_rope_index( return position_ids, mrope_position_deltas + @can_return_tuple + @auto_docstring def get_video_features( - self, pixel_values_videos: torch.FloatTensor, video_grid_thw: torch.LongTensor | None = None - ): - """ - Encodes videos into continuous embeddings that can be forwarded to the language model. - - Args: - pixel_values_videos (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): - The tensors corresponding to the input videos. - video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*): - The temporal, height and width of feature shape of each video in LLM. + self, + pixel_values_videos: torch.FloatTensor, + video_grid_thw: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + pixel_values_videos (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): + The tensors corresponding to the input videos. + video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*): + The temporal, height and width of feature shape of each video in LLM. """ pixel_values_videos = pixel_values_videos.type(self.visual.dtype) # reshape video_grid_thw -> [b, 3] -> [1, h, w] * frames @@ -1153,26 +1167,36 @@ def get_video_features( repeated_row = torch.tensor([1, h.item(), w.item()]).unsqueeze(0).repeat(t, 1) temp_frames_hw.append(repeated_row) flattened_video_grid_thw = torch.cat(temp_frames_hw, dim=0) - video_embeds = self.visual(pixel_values_videos, grid_thw=flattened_video_grid_thw) + vision_outputs = self.visual( + pixel_values_videos, grid_thw=flattened_video_grid_thw, return_dict=True, **kwargs + ) split_sizes = (video_grid_thw.prod(-1) // self.visual.spatial_merge_size**2).tolist() - video_embeds = torch.split(video_embeds, split_sizes) - return video_embeds + video_embeds = torch.split(vision_outputs.pooler_output, split_sizes) + vision_outputs.pooler_output = video_embeds - def get_image_features(self, pixel_values: torch.FloatTensor, image_grid_thw: torch.LongTensor | None = None): - """ - Encodes images into continuous embeddings that can be forwarded to the language model. + return vision_outputs - Args: - pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): - The tensors corresponding to the input images. - image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): - The temporal, height and width of feature shape of each image in LLM. + @can_return_tuple + @auto_docstring + def get_image_features( + self, + pixel_values: torch.FloatTensor, + image_grid_thw: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): + The tensors corresponding to the input images. + image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): + The temporal, height and width of feature shape of each image in LLM. """ pixel_values = pixel_values.type(self.visual.dtype) - image_embeds = self.visual(pixel_values, grid_thw=image_grid_thw) + vision_outputs = self.visual(pixel_values, grid_thw=image_grid_thw, return_dict=True, **kwargs) split_sizes = (image_grid_thw.prod(-1) // self.visual.spatial_merge_size**2).tolist() - image_embeds = torch.split(image_embeds, split_sizes) - return image_embeds + image_embeds = torch.split(vision_outputs.pooler_output, split_sizes) + vision_outputs.pooler_output = image_embeds + + return vision_outputs def get_placeholder_mask( self, @@ -1248,13 +1272,13 @@ def forward( inputs_embeds = self.get_input_embeddings()(input_ids) if pixel_values is not None: - image_embeds = self.get_image_features(pixel_values, image_grid_thw) + image_embeds = self.get_image_features(pixel_values, image_grid_thw, return_dict=True).pooler_output image_embeds = torch.cat(image_embeds, dim=0).to(inputs_embeds.device, inputs_embeds.dtype) image_mask, _ = self.get_placeholder_mask(input_ids, inputs_embeds, image_features=image_embeds) inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds) if pixel_values_videos is not None: - video_embeds = self.get_video_features(pixel_values_videos, video_grid_thw) + video_embeds = self.get_video_features(pixel_values_videos, video_grid_thw, return_dict=True).pooler_output video_embeds = torch.cat(video_embeds, dim=0).to(inputs_embeds.device, inputs_embeds.dtype) _, video_mask = self.get_placeholder_mask(input_ids, inputs_embeds, video_features=video_embeds) inputs_embeds = inputs_embeds.masked_scatter(video_mask, video_embeds) @@ -1372,13 +1396,37 @@ def get_input_embeddings(self): def set_input_embeddings(self, value): self.model.set_input_embeddings(value) + @auto_docstring def get_video_features( - self, pixel_values_videos: torch.FloatTensor, video_grid_thw: torch.LongTensor | None = None - ): - return self.model.get_video_features(pixel_values_videos, video_grid_thw) + self, + pixel_values_videos: torch.FloatTensor, + video_grid_thw: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + pixel_values_videos (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): + The tensors corresponding to the input videos. + video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*): + The temporal, height and width of feature shape of each video in LLM. + """ + return self.model.get_video_features( + pixel_values_videos=pixel_values_videos, video_grid_thw=video_grid_thw, **kwargs + ) - def get_image_features(self, pixel_values: torch.FloatTensor, image_grid_thw: torch.LongTensor | None = None): - return self.model.get_image_features(pixel_values, image_grid_thw) + @auto_docstring + def get_image_features( + self, + pixel_values: torch.FloatTensor, + image_grid_thw: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): + The tensors corresponding to the input images. + image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): + The temporal, height and width of feature shape of each image in LLM. + """ + return self.model.get_image_features(pixel_values=pixel_values, image_grid_thw=image_grid_thw, **kwargs) @can_return_tuple @auto_docstring diff --git a/src/transformers/models/glm4v/modular_glm4v.py b/src/transformers/models/glm4v/modular_glm4v.py index e7941ccb79c8..0f1d57404d50 100644 --- a/src/transformers/models/glm4v/modular_glm4v.py +++ b/src/transformers/models/glm4v/modular_glm4v.py @@ -29,7 +29,7 @@ from ...masking_utils import create_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer -from ...modeling_outputs import BaseModelOutputWithPast +from ...modeling_outputs import BaseModelOutputWithPast, BaseModelOutputWithPooling from ...modeling_rope_utils import RopeParameters from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel from ...processing_utils import Unpack @@ -720,6 +720,10 @@ class Glm4vVisionModel(Glm4vPreTrainedModel): config: Glm4vVisionConfig input_modalities = ("image", "video") _no_split_modules = ["Glm4vVisionBlock"] + _can_record_outputs = { + "hidden_states": Glm4vVisionBlock, + "attentions": Glm4vVisionAttention, + } def __init__(self, config) -> None: super().__init__(config) @@ -778,13 +782,16 @@ def rot_pos_emb(self, grid_thw): rotary_pos_emb = rotary_pos_emb_full[pos_ids].flatten(1) return rotary_pos_emb, pos_ids - def forward(self, hidden_states: torch.Tensor, grid_thw: torch.Tensor, **kwargs) -> torch.Tensor: - """ - Args: - hidden_states (`torch.Tensor` of shape `(seq_len, hidden_size)`): - The final hidden states of the model. - grid_thw (`torch.Tensor` of shape `(num_images_or_videos, 3)`): - The temporal, height and width of feature shape of each image in LLM. + @check_model_inputs + @auto_docstring + def forward( + self, hidden_states: torch.Tensor, grid_thw: torch.Tensor, **kwargs: Unpack[TransformersKwargs] + ) -> tuple | BaseModelOutputWithPooling: + r""" + hidden_states (`torch.Tensor` of shape `(seq_len, hidden_size)`): + The final hidden states of the model. + grid_thw (`torch.Tensor` of shape `(num_images_or_videos, 3)`): + The temporal, height and width of feature shape of each image in LLM. Returns: `torch.Tensor`: hidden_states. @@ -818,6 +825,7 @@ def forward(self, hidden_states: torch.Tensor, grid_thw: torch.Tensor, **kwargs) hidden_states, cu_seqlens=cu_seqlens, position_embeddings=position_embeddings, + **kwargs, ) hidden_states = self.post_layernorm(hidden_states) @@ -828,8 +836,12 @@ def forward(self, hidden_states: torch.Tensor, grid_thw: torch.Tensor, **kwargs) hidden_states = hidden_states.permute(0, 3, 1, 2) hidden_states = self.downsample(hidden_states).view(-1, self.config.out_hidden_size) - hidden_states = self.merger(hidden_states) - return hidden_states + merged_hidden_states = self.merger(hidden_states) + + return BaseModelOutputWithPooling( + last_hidden_state=hidden_states, + pooler_output=merged_hidden_states, + ) class Glm4vTextModel(Qwen2_5_VLTextModel): @@ -1127,17 +1139,19 @@ def get_rope_index( return position_ids, mrope_position_deltas + @can_return_tuple + @auto_docstring def get_video_features( - self, pixel_values_videos: torch.FloatTensor, video_grid_thw: torch.LongTensor | None = None - ): - """ - Encodes videos into continuous embeddings that can be forwarded to the language model. - - Args: - pixel_values_videos (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): - The tensors corresponding to the input videos. - video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*): - The temporal, height and width of feature shape of each video in LLM. + self, + pixel_values_videos: torch.FloatTensor, + video_grid_thw: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + pixel_values_videos (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): + The tensors corresponding to the input videos. + video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*): + The temporal, height and width of feature shape of each video in LLM. """ pixel_values_videos = pixel_values_videos.type(self.visual.dtype) # reshape video_grid_thw -> [b, 3] -> [1, h, w] * frames @@ -1146,10 +1160,14 @@ def get_video_features( repeated_row = torch.tensor([1, h.item(), w.item()]).unsqueeze(0).repeat(t, 1) temp_frames_hw.append(repeated_row) flattened_video_grid_thw = torch.cat(temp_frames_hw, dim=0) - video_embeds = self.visual(pixel_values_videos, grid_thw=flattened_video_grid_thw) + vision_outputs = self.visual( + pixel_values_videos, grid_thw=flattened_video_grid_thw, return_dict=True, **kwargs + ) split_sizes = (video_grid_thw.prod(-1) // self.visual.spatial_merge_size**2).tolist() - video_embeds = torch.split(video_embeds, split_sizes) - return video_embeds + video_embeds = torch.split(vision_outputs.pooler_output, split_sizes) + vision_outputs.pooler_output = video_embeds + + return vision_outputs def get_placeholder_mask( self, @@ -1225,13 +1243,13 @@ def forward( inputs_embeds = self.get_input_embeddings()(input_ids) if pixel_values is not None: - image_embeds = self.get_image_features(pixel_values, image_grid_thw) + image_embeds = self.get_image_features(pixel_values, image_grid_thw, return_dict=True).pooler_output image_embeds = torch.cat(image_embeds, dim=0).to(inputs_embeds.device, inputs_embeds.dtype) image_mask, _ = self.get_placeholder_mask(input_ids, inputs_embeds, image_features=image_embeds) inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds) if pixel_values_videos is not None: - video_embeds = self.get_video_features(pixel_values_videos, video_grid_thw) + video_embeds = self.get_video_features(pixel_values_videos, video_grid_thw, return_dict=True).pooler_output video_embeds = torch.cat(video_embeds, dim=0).to(inputs_embeds.device, inputs_embeds.dtype) _, video_mask = self.get_placeholder_mask(input_ids, inputs_embeds, video_features=video_embeds) inputs_embeds = inputs_embeds.masked_scatter(video_mask, video_embeds) diff --git a/src/transformers/models/glm4v_moe/modeling_glm4v_moe.py b/src/transformers/models/glm4v_moe/modeling_glm4v_moe.py index 1f0ac5e14380..bc84e5b83e00 100644 --- a/src/transformers/models/glm4v_moe/modeling_glm4v_moe.py +++ b/src/transformers/models/glm4v_moe/modeling_glm4v_moe.py @@ -35,7 +35,7 @@ from ...masking_utils import create_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer -from ...modeling_outputs import ModelOutput, MoeModelOutputWithPast +from ...modeling_outputs import BaseModelOutputWithPooling, ModelOutput, MoeModelOutputWithPast from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel from ...processing_utils import Unpack @@ -767,6 +767,10 @@ class Glm4vMoeVisionModel(Glm4vMoePreTrainedModel): config: Glm4vMoeVisionConfig input_modalities = ("image", "video") _no_split_modules = ["Glm4vMoeVisionBlock"] + _can_record_outputs = { + "hidden_states": Glm4vMoeVisionBlock, + "attentions": Glm4vMoeVisionAttention, + } def __init__(self, config) -> None: super().__init__(config) @@ -825,13 +829,16 @@ def rot_pos_emb(self, grid_thw): rotary_pos_emb = rotary_pos_emb_full[pos_ids].flatten(1) return rotary_pos_emb, pos_ids - def forward(self, hidden_states: torch.Tensor, grid_thw: torch.Tensor, **kwargs) -> torch.Tensor: - """ - Args: - hidden_states (`torch.Tensor` of shape `(seq_len, hidden_size)`): - The final hidden states of the model. - grid_thw (`torch.Tensor` of shape `(num_images_or_videos, 3)`): - The temporal, height and width of feature shape of each image in LLM. + @check_model_inputs + @auto_docstring + def forward( + self, hidden_states: torch.Tensor, grid_thw: torch.Tensor, **kwargs: Unpack[TransformersKwargs] + ) -> tuple | BaseModelOutputWithPooling: + r""" + hidden_states (`torch.Tensor` of shape `(seq_len, hidden_size)`): + The final hidden states of the model. + grid_thw (`torch.Tensor` of shape `(num_images_or_videos, 3)`): + The temporal, height and width of feature shape of each image in LLM. Returns: `torch.Tensor`: hidden_states. @@ -865,6 +872,7 @@ def forward(self, hidden_states: torch.Tensor, grid_thw: torch.Tensor, **kwargs) hidden_states, cu_seqlens=cu_seqlens, position_embeddings=position_embeddings, + **kwargs, ) hidden_states = self.post_layernorm(hidden_states) @@ -875,8 +883,12 @@ def forward(self, hidden_states: torch.Tensor, grid_thw: torch.Tensor, **kwargs) hidden_states = hidden_states.permute(0, 3, 1, 2) hidden_states = self.downsample(hidden_states).view(-1, self.config.out_hidden_size) - hidden_states = self.merger(hidden_states) - return hidden_states + merged_hidden_states = self.merger(hidden_states) + + return BaseModelOutputWithPooling( + last_hidden_state=hidden_states, + pooler_output=merged_hidden_states, + ) class Glm4vMoeTextRotaryEmbedding(nn.Module): @@ -989,7 +1001,7 @@ def forward( use_cache: bool | None = None, cache_position: torch.LongTensor | None = None, **kwargs: Unpack[FlashAttentionKwargs], - ) -> MoeModelOutputWithPast: + ) -> tuple | MoeModelOutputWithPast: if (input_ids is None) ^ (inputs_embeds is not None): raise ValueError("You must specify exactly one of input_ids or inputs_embeds") @@ -1303,17 +1315,19 @@ def get_rope_index( return position_ids, mrope_position_deltas + @can_return_tuple + @auto_docstring def get_video_features( - self, pixel_values_videos: torch.FloatTensor, video_grid_thw: torch.LongTensor | None = None - ): - """ - Encodes videos into continuous embeddings that can be forwarded to the language model. - - Args: - pixel_values_videos (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): - The tensors corresponding to the input videos. - video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*): - The temporal, height and width of feature shape of each video in LLM. + self, + pixel_values_videos: torch.FloatTensor, + video_grid_thw: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + pixel_values_videos (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): + The tensors corresponding to the input videos. + video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*): + The temporal, height and width of feature shape of each video in LLM. """ pixel_values_videos = pixel_values_videos.type(self.visual.dtype) # reshape video_grid_thw -> [b, 3] -> [1, h, w] * frames @@ -1322,26 +1336,36 @@ def get_video_features( repeated_row = torch.tensor([1, h.item(), w.item()]).unsqueeze(0).repeat(t, 1) temp_frames_hw.append(repeated_row) flattened_video_grid_thw = torch.cat(temp_frames_hw, dim=0) - video_embeds = self.visual(pixel_values_videos, grid_thw=flattened_video_grid_thw) + vision_outputs = self.visual( + pixel_values_videos, grid_thw=flattened_video_grid_thw, return_dict=True, **kwargs + ) split_sizes = (video_grid_thw.prod(-1) // self.visual.spatial_merge_size**2).tolist() - video_embeds = torch.split(video_embeds, split_sizes) - return video_embeds + video_embeds = torch.split(vision_outputs.pooler_output, split_sizes) + vision_outputs.pooler_output = video_embeds - def get_image_features(self, pixel_values: torch.FloatTensor, image_grid_thw: torch.LongTensor | None = None): - """ - Encodes images into continuous embeddings that can be forwarded to the language model. + return vision_outputs - Args: - pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): - The tensors corresponding to the input images. - image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): - The temporal, height and width of feature shape of each image in LLM. + @can_return_tuple + @auto_docstring + def get_image_features( + self, + pixel_values: torch.FloatTensor, + image_grid_thw: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): + The tensors corresponding to the input images. + image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): + The temporal, height and width of feature shape of each image in LLM. """ pixel_values = pixel_values.type(self.visual.dtype) - image_embeds = self.visual(pixel_values, grid_thw=image_grid_thw) + vision_outputs = self.visual(pixel_values, grid_thw=image_grid_thw, return_dict=True, **kwargs) split_sizes = (image_grid_thw.prod(-1) // self.visual.spatial_merge_size**2).tolist() - image_embeds = torch.split(image_embeds, split_sizes) - return image_embeds + image_embeds = torch.split(vision_outputs.pooler_output, split_sizes) + vision_outputs.pooler_output = image_embeds + + return vision_outputs def get_placeholder_mask( self, @@ -1417,13 +1441,13 @@ def forward( inputs_embeds = self.get_input_embeddings()(input_ids) if pixel_values is not None: - image_embeds = self.get_image_features(pixel_values, image_grid_thw) + image_embeds = self.get_image_features(pixel_values, image_grid_thw, return_dict=True).pooler_output image_embeds = torch.cat(image_embeds, dim=0).to(inputs_embeds.device, inputs_embeds.dtype) image_mask, _ = self.get_placeholder_mask(input_ids, inputs_embeds, image_features=image_embeds) inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds) if pixel_values_videos is not None: - video_embeds = self.get_video_features(pixel_values_videos, video_grid_thw) + video_embeds = self.get_video_features(pixel_values_videos, video_grid_thw, return_dict=True).pooler_output video_embeds = torch.cat(video_embeds, dim=0).to(inputs_embeds.device, inputs_embeds.dtype) _, video_mask = self.get_placeholder_mask(input_ids, inputs_embeds, video_features=video_embeds) inputs_embeds = inputs_embeds.masked_scatter(video_mask, video_embeds) @@ -1594,13 +1618,37 @@ def get_input_embeddings(self): def set_input_embeddings(self, value): self.model.set_input_embeddings(value) + @auto_docstring def get_video_features( - self, pixel_values_videos: torch.FloatTensor, video_grid_thw: torch.LongTensor | None = None - ): - return self.model.get_video_features(pixel_values_videos, video_grid_thw) + self, + pixel_values_videos: torch.FloatTensor, + video_grid_thw: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + pixel_values_videos (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): + The tensors corresponding to the input videos. + video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*): + The temporal, height and width of feature shape of each video in LLM. + """ + return self.model.get_video_features( + pixel_values_videos=pixel_values_videos, video_grid_thw=video_grid_thw, **kwargs + ) - def get_image_features(self, pixel_values: torch.FloatTensor, image_grid_thw: torch.LongTensor | None = None): - return self.model.get_image_features(pixel_values, image_grid_thw) + @auto_docstring + def get_image_features( + self, + pixel_values: torch.FloatTensor, + image_grid_thw: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): + The tensors corresponding to the input images. + image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): + The temporal, height and width of feature shape of each image in LLM. + """ + return self.model.get_image_features(pixel_values=pixel_values, image_grid_thw=image_grid_thw, **kwargs) @auto_docstring @check_model_inputs diff --git a/src/transformers/models/glm4v_moe/modular_glm4v_moe.py b/src/transformers/models/glm4v_moe/modular_glm4v_moe.py index 5392289f5572..dd59ae412916 100644 --- a/src/transformers/models/glm4v_moe/modular_glm4v_moe.py +++ b/src/transformers/models/glm4v_moe/modular_glm4v_moe.py @@ -415,7 +415,7 @@ def forward( use_cache: bool | None = None, cache_position: torch.LongTensor | None = None, **kwargs: Unpack[FlashAttentionKwargs], - ) -> MoeModelOutputWithPast: + ) -> tuple | MoeModelOutputWithPast: if (input_ids is None) ^ (inputs_embeds is not None): raise ValueError("You must specify exactly one of input_ids or inputs_embeds") diff --git a/src/transformers/models/glm_image/modeling_glm_image.py b/src/transformers/models/glm_image/modeling_glm_image.py index fb6fb23db4d7..b694814e82df 100644 --- a/src/transformers/models/glm_image/modeling_glm_image.py +++ b/src/transformers/models/glm_image/modeling_glm_image.py @@ -32,7 +32,7 @@ from ...masking_utils import create_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer -from ...modeling_outputs import BaseModelOutputWithPast, ModelOutput +from ...modeling_outputs import BaseModelOutputWithPast, BaseModelOutputWithPooling, ModelOutput from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel from ...processing_utils import Unpack @@ -612,6 +612,23 @@ def forward(self, hidden_state: torch.Tensor): return hidden_state_quant, loss, min_encoding_indices +@dataclass +@auto_docstring +class GlmImageVQVAEModelOutput(BaseModelOutputWithPooling): + r""" + quantized_last_hidden_state (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): + Quantized last hidden state from the VQ-VAE model. + image_tokens (`torch.FloatTensor` of shape `(batch_size, config.vocab_size`): + Indices of the image tokens predicted by the VQ-VAE model. + embedding_loss (`torch.FloatTensor`): + The embedding loss computed during quantization. + """ + + quantized_last_hidden_state: torch.FloatTensor | None = None + image_tokens: torch.FloatTensor | None = None + embedding_loss: torch.FloatTensor | None = None + + @auto_docstring( custom_intro=""" The VQ-VAE model used in GlmImage for encoding/decoding images into discrete tokens. @@ -625,6 +642,7 @@ class GlmImageVQVAE(GlmImagePreTrainedModel): _no_split_modules = [ "GlmImageVQVAEVectorQuantizer", ] + _can_record_outputs = {} def __init__(self, config: GlmImageVQVAEConfig): super().__init__(config) @@ -634,16 +652,26 @@ def __init__(self, config: GlmImageVQVAEConfig): self.eval() # GlmImage's VQ model is frozen self.post_init() - def encode(self, hidden_states): - hidden_states = self.quant_conv(hidden_states) - quant, emb_loss, indices = self.quantize(hidden_states) - return quant, emb_loss, indices + @check_model_inputs + def encode(self, hidden_states) -> GlmImageVQVAEModelOutput: + conv_hidden_states = self.quant_conv(hidden_states) + quantized_last_hidden_state, emb_loss, indices = self.quantize(conv_hidden_states) + return GlmImageVQVAEModelOutput( + last_hidden_state=hidden_states, + quantized_last_hidden_state=quantized_last_hidden_state, + image_tokens=indices, + embedding_loss=emb_loss, + ) class GlmImageVisionModel(GlmImagePreTrainedModel): config: GlmImageVisionConfig input_modalities = ("image",) _no_split_modules = ["GlmImageVisionBlock"] + _can_record_outputs = { + "hidden_states": GlmImageVisionBlock, + "attentions": GlmImageVisionAttention, + } main_input_name = "pixel_values" def __init__(self, config: GlmImageVisionConfig) -> None: @@ -688,13 +716,16 @@ def rot_pos_emb(self, grid_thw): pos_ids = torch.cat(pos_ids, dim=0) return pos_ids - def forward(self, pixel_values: torch.Tensor, grid_thw: torch.Tensor, **kwargs) -> torch.Tensor: - """ - Args: - pixel_values (`torch.Tensor` of shape `(total_patches, num_channels * patch_size * patch_size)`): - Packed pixel values. - grid_thw (`torch.Tensor` of shape `(num_images, 3)`): - The temporal, height and width of feature shape of each image. + @check_model_inputs + @auto_docstring + def forward( + self, pixel_values: torch.Tensor, grid_thw: torch.Tensor, **kwargs: Unpack[TransformersKwargs] + ) -> tuple | BaseModelOutputWithPooling: + r""" + pixel_values (`torch.Tensor` of shape `(total_patches, num_channels * patch_size * patch_size)`): + Packed pixel values. + grid_thw (`torch.Tensor` of shape `(num_images, 3)`): + The temporal, height and width of feature shape of each image. Returns: `torch.Tensor` of shape `(total_patches, hidden_size)`: Hidden states. @@ -723,7 +754,8 @@ def forward(self, pixel_values: torch.Tensor, grid_thw: torch.Tensor, **kwargs) hidden_states, cu_seqlens=cu_seqlens, ) - return hidden_states + + return BaseModelOutputWithPooling(last_hidden_state=hidden_states) class GlmImageTextRotaryEmbedding(nn.Module): @@ -1161,21 +1193,27 @@ def get_rope_index( return position_ids, mrope_position_deltas - def get_image_features(self, pixel_values: torch.FloatTensor, image_grid_thw: torch.LongTensor | None = None): - """ - Encodes images into continuous embeddings that can be forwarded to the language model. - - Args: - pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): - The tensors corresponding to the input images. - image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): - The temporal, height and width of feature shape of each image in LLM. + @can_return_tuple + @auto_docstring + def get_image_features( + self, + pixel_values: torch.FloatTensor, + image_grid_thw: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): + The tensors corresponding to the input images. + image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): + The temporal, height and width of feature shape of each image in LLM. """ pixel_values = pixel_values.type(self.visual.dtype) - image_embeds = self.visual(pixel_values, grid_thw=image_grid_thw) + vision_outputs = self.visual(pixel_values, grid_thw=image_grid_thw, return_dict=True, **kwargs) split_sizes = (image_grid_thw.prod(-1) // self.visual.spatial_merge_size**2).tolist() - image_embeds = torch.split(image_embeds, split_sizes) - return image_embeds + image_embeds = torch.split(vision_outputs.last_hidden_state, split_sizes) + vision_outputs.pooler_output = image_embeds + + return vision_outputs def get_placeholder_mask( self, @@ -1233,7 +1271,7 @@ def forward( raise ValueError("You must specify exactly one of input_ids or inputs_embeds") if pixel_values is not None: - image_embeds = self.get_image_features(pixel_values, image_grid_thw[:-1]) + image_embeds = self.get_image_features(pixel_values, image_grid_thw[:-1], return_dict=True).pooler_output image_embeds = torch.cat(image_embeds, dim=0) image_ids = self.get_image_tokens(image_embeds, image_grid_thw[:-1]) image_ids = image_ids.view(-1).to(input_ids.device) @@ -1319,8 +1357,8 @@ def get_image_tokens( grid_t, grid_h, grid_w = image_grid_thw[i].tolist() hs = hs.view(grid_t, grid_h, grid_w, hidden_size) hs = hs.permute(0, 3, 1, 2).contiguous() - _, _, image_toks = self.vqmodel.encode(hs) - all_image_toks.append(image_toks) + vqmodel_outputs: GlmImageVQVAEModelOutput = self.vqmodel.encode(hs) + all_image_toks.append(vqmodel_outputs.image_tokens) return torch.cat(all_image_toks, dim=0) @@ -1369,8 +1407,20 @@ def __init__(self, config): # Initialize weights and apply final processing self.post_init() - def get_image_features(self, pixel_values: torch.FloatTensor, image_grid_thw: torch.LongTensor | None = None): - return self.model.get_image_features(pixel_values, image_grid_thw) + @auto_docstring + def get_image_features( + self, + pixel_values: torch.FloatTensor, + image_grid_thw: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): + The tensors corresponding to the input images. + image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): + The temporal, height and width of feature shape of each image in LLM. + """ + return self.model.get_image_features(pixel_values, image_grid_thw, **kwargs) def get_image_tokens(self, hidden_states: torch.FloatTensor, image_grid_thw: torch.LongTensor | None = None): return self.model.get_image_tokens(hidden_states, image_grid_thw) diff --git a/src/transformers/models/glm_image/modular_glm_image.py b/src/transformers/models/glm_image/modular_glm_image.py index 38ac344475c2..d2f4604a4e72 100644 --- a/src/transformers/models/glm_image/modular_glm_image.py +++ b/src/transformers/models/glm_image/modular_glm_image.py @@ -25,11 +25,13 @@ from ...feature_extraction_utils import BatchFeature from ...generation import GenerationMixin from ...image_utils import ImageInput +from ...modeling_outputs import BaseModelOutputWithPooling from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel from ...processing_utils import ImagesKwargs, ProcessorMixin, Unpack from ...tokenization_utils_base import PreTokenizedInput, TextInput -from ...utils import TransformersKwargs, is_torch_available, logging -from ..chameleon.modeling_chameleon import ChameleonVQVAE, ChameleonVQVAEVectorQuantizer +from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, is_torch_available, logging +from ...utils.generic import check_model_inputs +from ..chameleon.modeling_chameleon import ChameleonVQVAE, ChameleonVQVAEModelOutput, ChameleonVQVAEVectorQuantizer from ..glm4v.configuration_glm4v import Glm4vTextConfig, Glm4vVisionConfig from ..glm4v.modeling_glm4v import ( Glm4vCausalLMOutputWithPast, @@ -515,19 +517,29 @@ def forward(self, hidden_state: torch.Tensor): return hidden_state_quant, loss, min_encoding_indices +class GlmImageVQVAEModelOutput(ChameleonVQVAEModelOutput): + pass + + class GlmImageVQVAE(ChameleonVQVAE): _no_split_modules = [ "GlmImageVQVAEVectorQuantizer", ] + _can_record_outputs = {} def __init__(self, config: GlmImageVQVAEConfig): super().__init__(config) del self.encoder def encode(self, hidden_states): - hidden_states = self.quant_conv(hidden_states) - quant, emb_loss, indices = self.quantize(hidden_states) - return quant, emb_loss, indices + conv_hidden_states = self.quant_conv(hidden_states) + quantized_last_hidden_state, emb_loss, indices = self.quantize(conv_hidden_states) + return GlmImageVQVAEModelOutput( + last_hidden_state=hidden_states, + quantized_last_hidden_state=quantized_last_hidden_state, + image_tokens=indices, + embedding_loss=emb_loss, + ) class GlmImageVisionModel(Glm4vVisionModel): @@ -573,13 +585,16 @@ def rot_pos_emb(self, grid_thw): pos_ids = torch.cat(pos_ids, dim=0) return pos_ids - def forward(self, pixel_values: torch.Tensor, grid_thw: torch.Tensor, **kwargs) -> torch.Tensor: - """ - Args: - pixel_values (`torch.Tensor` of shape `(total_patches, num_channels * patch_size * patch_size)`): - Packed pixel values. - grid_thw (`torch.Tensor` of shape `(num_images, 3)`): - The temporal, height and width of feature shape of each image. + @check_model_inputs + @auto_docstring + def forward( + self, pixel_values: torch.Tensor, grid_thw: torch.Tensor, **kwargs: Unpack[TransformersKwargs] + ) -> tuple | BaseModelOutputWithPooling: + r""" + pixel_values (`torch.Tensor` of shape `(total_patches, num_channels * patch_size * patch_size)`): + Packed pixel values. + grid_thw (`torch.Tensor` of shape `(num_images, 3)`): + The temporal, height and width of feature shape of each image. Returns: `torch.Tensor` of shape `(total_patches, hidden_size)`: Hidden states. @@ -608,7 +623,8 @@ def forward(self, pixel_values: torch.Tensor, grid_thw: torch.Tensor, **kwargs) hidden_states, cu_seqlens=cu_seqlens, ) - return hidden_states + + return BaseModelOutputWithPooling(last_hidden_state=hidden_states) class GlmImageTextModel(Glm4vTextModel): @@ -879,13 +895,35 @@ def get_image_tokens( grid_t, grid_h, grid_w = image_grid_thw[i].tolist() hs = hs.view(grid_t, grid_h, grid_w, hidden_size) hs = hs.permute(0, 3, 1, 2).contiguous() - _, _, image_toks = self.vqmodel.encode(hs) - all_image_toks.append(image_toks) + vqmodel_outputs: GlmImageVQVAEModelOutput = self.vqmodel.encode(hs) + all_image_toks.append(vqmodel_outputs.image_tokens) return torch.cat(all_image_toks, dim=0) def get_video_features(self): raise AttributeError("Not needed for GlmImage") + @can_return_tuple + @auto_docstring + def get_image_features( + self, + pixel_values: torch.FloatTensor, + image_grid_thw: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): + The tensors corresponding to the input images. + image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): + The temporal, height and width of feature shape of each image in LLM. + """ + pixel_values = pixel_values.type(self.visual.dtype) + vision_outputs = self.visual(pixel_values, grid_thw=image_grid_thw, return_dict=True, **kwargs) + split_sizes = (image_grid_thw.prod(-1) // self.visual.spatial_merge_size**2).tolist() + image_embeds = torch.split(vision_outputs.last_hidden_state, split_sizes) + vision_outputs.pooler_output = image_embeds + + return vision_outputs + def get_placeholder_mask( self, input_ids: torch.LongTensor, @@ -940,7 +978,7 @@ def forward( raise ValueError("You must specify exactly one of input_ids or inputs_embeds") if pixel_values is not None: - image_embeds = self.get_image_features(pixel_values, image_grid_thw[:-1]) + image_embeds = self.get_image_features(pixel_values, image_grid_thw[:-1], return_dict=True).pooler_output image_embeds = torch.cat(image_embeds, dim=0) image_ids = self.get_image_tokens(image_embeds, image_grid_thw[:-1]) image_ids = image_ids.view(-1).to(input_ids.device) @@ -1020,8 +1058,20 @@ def __init__(self, config): # Initialize weights and apply final processing self.post_init() - def get_image_features(self, pixel_values: torch.FloatTensor, image_grid_thw: torch.LongTensor | None = None): - return self.model.get_image_features(pixel_values, image_grid_thw) + @auto_docstring + def get_image_features( + self, + pixel_values: torch.FloatTensor, + image_grid_thw: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): + The tensors corresponding to the input images. + image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): + The temporal, height and width of feature shape of each image in LLM. + """ + return self.model.get_image_features(pixel_values, image_grid_thw, **kwargs) def get_image_tokens(self, hidden_states: torch.FloatTensor, image_grid_thw: torch.LongTensor | None = None): return self.model.get_image_tokens(hidden_states, image_grid_thw) diff --git a/src/transformers/models/glmasr/modeling_glmasr.py b/src/transformers/models/glmasr/modeling_glmasr.py index 555ce5ac7472..9a08c0ec1adb 100644 --- a/src/transformers/models/glmasr/modeling_glmasr.py +++ b/src/transformers/models/glmasr/modeling_glmasr.py @@ -26,12 +26,12 @@ from ...generation import GenerationMixin from ...integrations import use_kernelized_func from ...modeling_layers import GradientCheckpointingLayer -from ...modeling_outputs import BaseModelOutput, CausalLMOutputWithPast +from ...modeling_outputs import BaseModelOutputWithPooling, CausalLMOutputWithPast from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel from ...processing_utils import Unpack -from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, is_torch_available -from ...utils.generic import check_model_inputs, maybe_autocast +from ...utils import TransformersKwargs, auto_docstring, is_torch_available +from ...utils.generic import can_return_tuple, check_model_inputs, maybe_autocast from ..auto import AutoModel, AutoModelForCausalLM from .configuration_glmasr import GlmAsrConfig, GlmAsrEncoderConfig @@ -292,6 +292,10 @@ class GlmAsrEncoder(GlmAsrPreTrainedModel): main_input_name = "input_features" input_modalities = "audio" _no_split_modules = ["GlmAsrEncoderLayer"] + _can_record_outputs = { + "hidden_states": GlmAsrEncoderLayer, + "attentions": GlmAsrAttention, + } def __init__(self, config: GlmAsrEncoderConfig): super().__init__(config) @@ -322,7 +326,7 @@ def forward(self, input_features, **kwargs: Unpack[TransformersKwargs]): hidden_states = encoder_layer(hidden_states, position_embeddings=position_embeddings, **kwargs) hidden_states = self.norm(hidden_states) - return BaseModelOutput(last_hidden_state=hidden_states) + return BaseModelOutputWithPooling(last_hidden_state=hidden_states) class GlmAsrMultiModalProjector(nn.Module): @@ -382,26 +386,27 @@ def set_decoder(self, decoder): def get_decoder(self): return self.language_model.get_decoder() + @can_return_tuple + @auto_docstring( + custom_intro="Compute audio embeddings from log-mel input features using the audio encoder and multi-modal projector." + ) def get_audio_features( - self, input_features: torch.FloatTensor, input_features_mask: torch.Tensor - ) -> torch.FloatTensor: - """ - This method is used to get the audio embeddings from input features (a log mel spectrogram), meaning inferring the audio encoder and the multi-modal projector. - Args: - input_features (`torch.FloatTensor`): - Float values of mel features extracted from the raw speech waveform. Raw speech waveform can be - obtained by loading a `.flac` or `.wav` audio file into an array of type `list[float]` or a - `numpy.ndarray`, *e.g.* via the soundfile library (`pip install soundfile`). To prepare the array into - `input_features`, the [`AutoFeatureExtractor`] should be used for extracting the mel features, padding - and conversion into a tensor of type `torch.FloatTensor`. See [`~WhisperFeatureExtractor.__call__`] - input_features_mask (`torch.Tensor` of shape `(batch_size, feature_sequence_length)`): - Mask to avoid performing attention on padded feature indices. - - Returns: - `torch.FloatTensor`: - The audio embeddings. + self, + input_features: torch.FloatTensor, + input_features_mask: torch.Tensor, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + input_features (`torch.FloatTensor`): + Float values of mel features extracted from the raw speech waveform. Raw speech waveform can be + obtained by loading a `.flac` or `.wav` audio file into an array of type `list[float]` or a + `numpy.ndarray`, *e.g.* via the soundfile library (`pip install soundfile`). To prepare the array into + `input_features`, the [`AutoFeatureExtractor`] should be used for extracting the mel features, padding + and conversion into a tensor of type `torch.FloatTensor`. See [`~WhisperFeatureExtractor.__call__`] + input_features_mask (`torch.Tensor` of shape `(batch_size, feature_sequence_length)`): + Mask to avoid performing attention on padded feature indices. """ - audio_outputs = self.audio_tower(input_features) + audio_outputs = self.audio_tower(input_features, return_dict=True, **kwargs) audio_hidden_states = audio_outputs.last_hidden_state audio_hidden_states = audio_hidden_states.reshape( input_features.shape[0], -1, self.config.audio_config.intermediate_size @@ -415,8 +420,9 @@ def get_audio_features( post_lengths = (audio_lengths - merge_factor) // merge_factor + 1 valid_mask = torch.arange(audio_embeds.shape[1], device=post_lengths.device)[None, :] < post_lengths[:, None] - audio_embeds = audio_embeds[valid_mask.to(audio_embeds.device)] - return audio_embeds + audio_outputs.pooler_output = audio_embeds[valid_mask.to(audio_embeds.device)] + + return audio_outputs @can_return_tuple @auto_docstring @@ -468,7 +474,7 @@ def forward( inputs_embeds = self.get_input_embeddings()(input_ids) if input_features is not None and input_ids is not None: - audio_embeds = self.get_audio_features(input_features, input_features_mask) + audio_embeds = self.get_audio_features(input_features, input_features_mask, return_dict=True).pooler_output # replace text-audio token placeholders with audio embeddings audio_token_mask = (input_ids == self.config.audio_token_id).unsqueeze(-1) diff --git a/src/transformers/models/glmasr/modular_glmasr.py b/src/transformers/models/glmasr/modular_glmasr.py index dc6128d23de1..a81a6ed7748d 100644 --- a/src/transformers/models/glmasr/modular_glmasr.py +++ b/src/transformers/models/glmasr/modular_glmasr.py @@ -21,11 +21,11 @@ from ...cache_utils import Cache from ...feature_extraction_utils import BatchFeature from ...modeling_layers import GradientCheckpointingLayer -from ...modeling_outputs import BaseModelOutput, CausalLMOutputWithPast +from ...modeling_outputs import BaseModelOutputWithPooling, CausalLMOutputWithPast from ...modeling_utils import ALL_ATTENTION_FUNCTIONS from ...processing_utils import Unpack from ...utils import TransformersKwargs, auto_docstring, is_torch_available, logging -from ...utils.generic import check_model_inputs +from ...utils.generic import can_return_tuple, check_model_inputs from ..audioflamingo3.modeling_audioflamingo3 import ( AudioFlamingo3ForConditionalGeneration, AudioFlamingo3MultiModalProjector, @@ -305,6 +305,10 @@ class GlmAsrEncoder(GlmAsrPreTrainedModel): main_input_name = "input_features" input_modalities = "audio" _no_split_modules = ["GlmAsrEncoderLayer"] + _can_record_outputs = { + "hidden_states": GlmAsrEncoderLayer, + "attentions": GlmAsrAttention, + } def __init__(self, config: GlmAsrEncoderConfig): super().__init__(config) @@ -335,7 +339,7 @@ def forward(self, input_features, **kwargs: Unpack[TransformersKwargs]): hidden_states = encoder_layer(hidden_states, position_embeddings=position_embeddings, **kwargs) hidden_states = self.norm(hidden_states) - return BaseModelOutput(last_hidden_state=hidden_states) + return BaseModelOutputWithPooling(last_hidden_state=hidden_states) class GlmAsrMultiModalProjector(AudioFlamingo3MultiModalProjector): @@ -351,10 +355,17 @@ def __init__(self, config: GlmAsrConfig): """ ) class GlmAsrForConditionalGeneration(AudioFlamingo3ForConditionalGeneration): + @can_return_tuple + @auto_docstring( + custom_intro="Compute audio embeddings from log-mel input features using the audio encoder and multi-modal projector." + ) def get_audio_features( - self, input_features: torch.FloatTensor, input_features_mask: torch.Tensor - ) -> torch.FloatTensor: - audio_outputs = self.audio_tower(input_features) + self, + input_features: torch.FloatTensor, + input_features_mask: torch.Tensor, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + audio_outputs = self.audio_tower(input_features, return_dict=True, **kwargs) audio_hidden_states = audio_outputs.last_hidden_state audio_hidden_states = audio_hidden_states.reshape( input_features.shape[0], -1, self.config.audio_config.intermediate_size @@ -368,8 +379,9 @@ def get_audio_features( post_lengths = (audio_lengths - merge_factor) // merge_factor + 1 valid_mask = torch.arange(audio_embeds.shape[1], device=post_lengths.device)[None, :] < post_lengths[:, None] - audio_embeds = audio_embeds[valid_mask.to(audio_embeds.device)] - return audio_embeds + audio_outputs.pooler_output = audio_embeds[valid_mask.to(audio_embeds.device)] + + return audio_outputs def forward( self, diff --git a/src/transformers/models/got_ocr2/configuration_got_ocr2.py b/src/transformers/models/got_ocr2/configuration_got_ocr2.py index 253cdd5612d6..4048ddaa4328 100644 --- a/src/transformers/models/got_ocr2/configuration_got_ocr2.py +++ b/src/transformers/models/got_ocr2/configuration_got_ocr2.py @@ -18,7 +18,6 @@ # See the License for the specific language governing permissions and # limitations under the License. - from ...configuration_utils import PreTrainedConfig from ..auto import CONFIG_MAPPING, AutoConfig diff --git a/src/transformers/models/got_ocr2/modeling_got_ocr2.py b/src/transformers/models/got_ocr2/modeling_got_ocr2.py index 34d3bca3acea..e5f2d747cf1c 100644 --- a/src/transformers/models/got_ocr2/modeling_got_ocr2.py +++ b/src/transformers/models/got_ocr2/modeling_got_ocr2.py @@ -18,7 +18,6 @@ # See the License for the specific language governing permissions and # limitations under the License. - import collections from dataclasses import dataclass @@ -26,17 +25,16 @@ import torch.nn as nn import torch.nn.functional as F -from transformers.utils.generic import check_model_inputs - from ... import initialization as init from ...activations import ACT2FN from ...cache_utils import Cache from ...generation import GenerationMixin from ...modeling_layers import GradientCheckpointingLayer -from ...modeling_outputs import BaseModelOutputWithPast, ModelOutput +from ...modeling_outputs import BaseModelOutputWithPast, BaseModelOutputWithPooling, ModelOutput from ...modeling_utils import PreTrainedModel from ...processing_utils import Unpack from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, torch_compilable_check +from ...utils.generic import check_model_inputs from ..auto import AutoModel from .configuration_got_ocr2 import GotOcr2Config, GotOcr2VisionConfig @@ -439,7 +437,7 @@ def get_input_embeddings(self): @check_model_inputs(tie_last_hidden_states=False) def forward( self, pixel_values: torch.FloatTensor | None = None, **kwargs: Unpack[TransformersKwargs] - ) -> GotOcr2VisionEncoderOutput: + ) -> tuple | GotOcr2VisionEncoderOutput: if pixel_values is None: raise ValueError("You have to specify pixel_values") @@ -550,20 +548,20 @@ def get_input_embeddings(self): def set_input_embeddings(self, value): self.language_model.set_input_embeddings(value) + @can_return_tuple + @auto_docstring( + custom_intro="Obtains image last hidden states from the vision tower and apply multimodal projection." + ) def get_image_features( self, pixel_values: torch.FloatTensor, - ): - """ - Obtains image last hidden states from the vision tower and apply multimodal projection. + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + image_outputs = self.vision_tower(pixel_values, return_dict=True, **kwargs) + last_hidden_state = image_outputs.last_hidden_state + image_outputs.pooler_output = self.multi_modal_projector(last_hidden_state) - Args: - pixel_values (`torch.FloatTensor]` of shape `(batch_size, channels, height, width)`) - Returns: - image_features (`torch.Tensor`): Image feature tensor of shape `(num_images, image_length, embed_dim)`). - """ - image_outputs = self.vision_tower(pixel_values).last_hidden_state - return self.multi_modal_projector(image_outputs) + return image_outputs def get_placeholder_mask( self, input_ids: torch.LongTensor, inputs_embeds: torch.FloatTensor, image_features: torch.FloatTensor @@ -589,7 +587,7 @@ def get_placeholder_mask( ) return special_image_mask - @can_return_tuple + @check_model_inputs(tie_last_hidden_states=False) @auto_docstring def forward( self, @@ -619,7 +617,9 @@ def forward( inputs_embeds = self.get_input_embeddings()(input_ids) if pixel_values is not None: - image_features = self.get_image_features(pixel_values=pixel_values.to(inputs_embeds.dtype)) + image_features = self.get_image_features( + pixel_values=pixel_values.to(inputs_embeds.dtype), return_dict=True + ).pooler_output image_features = image_features.to(inputs_embeds.device, inputs_embeds.dtype) special_image_mask = self.get_placeholder_mask( input_ids, inputs_embeds=inputs_embeds, image_features=image_features @@ -677,19 +677,11 @@ def set_input_embeddings(self, value): def get_output_embeddings(self) -> nn.Module: return self.lm_head + @auto_docstring def get_image_features( - self, - pixel_values: torch.FloatTensor, - vision_feature_layer: int | list[int] | None = None, - vision_feature_select_strategy: str | None = None, - **kwargs, - ): - return self.model.get_image_features( - pixel_values=pixel_values, - vision_feature_layer=vision_feature_layer, - vision_feature_select_strategy=vision_feature_select_strategy, - **kwargs, - ) + self, pixel_values: torch.FloatTensor, **kwargs: Unpack[TransformersKwargs] + ) -> tuple | BaseModelOutputWithPooling: + return self.model.get_image_features(pixel_values=pixel_values, **kwargs) @can_return_tuple @auto_docstring diff --git a/src/transformers/models/got_ocr2/modular_got_ocr2.py b/src/transformers/models/got_ocr2/modular_got_ocr2.py index 28484a502f5d..cf7a3e584549 100644 --- a/src/transformers/models/got_ocr2/modular_got_ocr2.py +++ b/src/transformers/models/got_ocr2/modular_got_ocr2.py @@ -12,13 +12,13 @@ # See the License for the specific language governing permissions and # limitations under the License. - import torch import torch.nn as nn from ... import initialization as init from ...cache_utils import Cache from ...configuration_utils import PreTrainedConfig +from ...modeling_outputs import BaseModelOutputWithPooling from ...modeling_utils import PreTrainedModel from ...processing_utils import Unpack from ...utils import auto_docstring, can_return_tuple, logging @@ -300,20 +300,20 @@ def __init__(self, config: GotOcr2Config): super().__init__(config) self.vision_tower = GotOcr2VisionEncoder(config.vision_config) + @can_return_tuple + @auto_docstring( + custom_intro="Obtains image last hidden states from the vision tower and apply multimodal projection." + ) def get_image_features( self, pixel_values: torch.FloatTensor, - ): - """ - Obtains image last hidden states from the vision tower and apply multimodal projection. + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + image_outputs = self.vision_tower(pixel_values, return_dict=True, **kwargs) + last_hidden_state = image_outputs.last_hidden_state + image_outputs.pooler_output = self.multi_modal_projector(last_hidden_state) - Args: - pixel_values (`torch.FloatTensor]` of shape `(batch_size, channels, height, width)`) - Returns: - image_features (`torch.Tensor`): Image feature tensor of shape `(num_images, image_length, embed_dim)`). - """ - image_outputs = self.vision_tower(pixel_values).last_hidden_state - return self.multi_modal_projector(image_outputs) + return image_outputs def forward( self, @@ -343,7 +343,9 @@ def forward( inputs_embeds = self.get_input_embeddings()(input_ids) if pixel_values is not None: - image_features = self.get_image_features(pixel_values=pixel_values.to(inputs_embeds.dtype)) + image_features = self.get_image_features( + pixel_values=pixel_values.to(inputs_embeds.dtype), return_dict=True + ).pooler_output image_features = image_features.to(inputs_embeds.device, inputs_embeds.dtype) special_image_mask = self.get_placeholder_mask( input_ids, inputs_embeds=inputs_embeds, image_features=image_features @@ -468,6 +470,12 @@ def forward( image_hidden_states=outputs.image_hidden_states, ) + @auto_docstring + def get_image_features( + self, pixel_values: torch.FloatTensor, **kwargs: Unpack[TransformersKwargs] + ) -> tuple | BaseModelOutputWithPooling: + return self.model.get_image_features(pixel_values=pixel_values, **kwargs) + __all__ = [ "GotOcr2VisionConfig", diff --git a/src/transformers/models/granite_speech/modeling_granite_speech.py b/src/transformers/models/granite_speech/modeling_granite_speech.py index b34500c69b83..5f6bf6461141 100644 --- a/src/transformers/models/granite_speech/modeling_granite_speech.py +++ b/src/transformers/models/granite_speech/modeling_granite_speech.py @@ -22,9 +22,11 @@ from ... import initialization as init from ...cache_utils import Cache from ...generation import GenerationMixin -from ...modeling_outputs import ModelOutput +from ...modeling_outputs import BaseModelOutputWithPooling, ModelOutput from ...modeling_utils import PreTrainedModel -from ...utils import auto_docstring, is_peft_available, logging +from ...processing_utils import Unpack +from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, is_peft_available, logging +from ...utils.generic import check_model_inputs from ..auto import AutoModel, AutoModelForCausalLM from .configuration_granite_speech import GraniteSpeechConfig, GraniteSpeechEncoderConfig @@ -248,10 +250,38 @@ def forward(self, hidden_states: torch.Tensor, attention_dists: torch.Tensor) -> return hidden_states -class GraniteSpeechCTCEncoder(nn.Module): +@auto_docstring +class GraniteSpeechPreTrainedModel(PreTrainedModel): + config: GraniteSpeechConfig + input_modalities = ("audio", "text") + + _supports_flash_attn = False # `blip_2_qformer` dependency does not allow for this + _supports_sdpa = True + + @torch.no_grad() + def _init_weights(self, module: nn.Module): + """Initialize the weights.""" + super()._init_weights(module) + if isinstance(module, GraniteSpeechEncoderProjector): + init.normal_(module.query) + elif isinstance(module, GraniteSpeechCTCEncoder): + context_size = module.config.context_size + seq = torch.arange(context_size) + relpos_dist = seq.view(-1, 1) - seq.view(1, -1) + attention_dists = torch.clamp(relpos_dist, -context_size, context_size) + module.config.max_pos_emb + init.copy_(module.attention_dists, attention_dists) + + +class GraniteSpeechCTCEncoder(GraniteSpeechPreTrainedModel): + config: GraniteSpeechEncoderConfig + input_modalities = "audio" + _can_record_outputs = { + "hidden_states": GraniteSpeechConformerBlock, + "attentions": GraniteSpeechConformerAttention, + } + def __init__(self, config: GraniteSpeechEncoderConfig): - super().__init__() - self.config = config + super().__init__(config) # Precompute clamped relative positional encoding distances seq = torch.arange(config.context_size) @@ -264,8 +294,12 @@ def __init__(self, config: GraniteSpeechEncoderConfig): self.out = nn.Linear(config.hidden_dim, config.output_dim, bias=True) self.out_mid = nn.Linear(config.output_dim, config.hidden_dim, bias=True) self.num_layers = config.num_layers + self.post_init() - def forward(self, hidden_states: torch.Tensor): + @check_model_inputs + def forward( + self, hidden_states: torch.Tensor, **kwargs: Unpack[TransformersKwargs] + ) -> tuple | BaseModelOutputWithPooling: hidden_states = self.input_linear(hidden_states) for idx, layer in enumerate(self.layers, start=1): hidden_states = layer(hidden_states, attention_dists=self.attention_dists) @@ -274,29 +308,8 @@ def forward(self, hidden_states: torch.Tensor): hidden_states_mid = hidden_states.clone() hidden_states_mid = self.out(hidden_states_mid) hidden_states += self.out_mid(nn.Softmax(dim=-1)(hidden_states_mid)) - return hidden_states - -@auto_docstring -class GraniteSpeechPreTrainedModel(PreTrainedModel): - config: GraniteSpeechConfig - input_modalities = ("audio", "text") - - _supports_flash_attn = False # `blip_2_qformer` dependency does not allow for this - _supports_sdpa = True - - @torch.no_grad() - def _init_weights(self, module: nn.Module): - """Initialize the weights.""" - super()._init_weights(module) - if isinstance(module, GraniteSpeechEncoderProjector): - init.normal_(module.query) - elif isinstance(module, GraniteSpeechCTCEncoder): - context_size = module.config.context_size - seq = torch.arange(context_size) - relpos_dist = seq.view(-1, 1) - seq.view(1, -1) - attention_dists = torch.clamp(relpos_dist, -context_size, context_size) + module.config.max_pos_emb - init.copy_(module.attention_dists, attention_dists) + return BaseModelOutputWithPooling(last_hidden_state=hidden_states) @auto_docstring( @@ -344,11 +357,16 @@ def get_input_embeddings(self): def get_output_embeddings(self): return self.language_model.get_output_embeddings() - def get_audio_features(self, input_features: torch.Tensor) -> torch.Tensor: - """Get the audio features to merged into the multimodal embeddings.""" - encoder_embeds = self.encoder(input_features) - projected_embeds = self.projector(encoder_embeds) - return projected_embeds + @can_return_tuple + @auto_docstring + def get_audio_features( + self, input_features: torch.Tensor, **kwargs: Unpack[TransformersKwargs] + ) -> tuple | BaseModelOutputWithPooling: + audio_outputs = self.encoder(input_features, return_dict=True, **kwargs) + projected_embeds = self.projector(audio_outputs.last_hidden_state) + audio_outputs.pooler_output = projected_embeds + + return audio_outputs @auto_docstring def forward( @@ -405,7 +423,7 @@ def forward( if input_features.dtype != self.dtype: input_features = input_features.to(self.dtype) # Get the audio features from the encoder / projector - audio_embeds = self.get_audio_features(input_features) + audio_embeds = self.get_audio_features(input_features, return_dict=True).pooler_output # Merge the audio features into the LLM embeddings inputs_embeds = self.get_merged_audio_embeddings( diff --git a/src/transformers/models/groupvit/modeling_groupvit.py b/src/transformers/models/groupvit/modeling_groupvit.py index f724435cf8fd..c7b9870cd606 100644 --- a/src/transformers/models/groupvit/modeling_groupvit.py +++ b/src/transformers/models/groupvit/modeling_groupvit.py @@ -27,7 +27,8 @@ from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling from ...modeling_utils import PreTrainedModel -from ...utils import ModelOutput, auto_docstring, filter_out_non_signature_kwargs, logging, torch_int +from ...processing_utils import Unpack +from ...utils import ModelOutput, TransformersKwargs, auto_docstring, can_return_tuple, logging, torch_int from .configuration_groupvit import GroupViTConfig, GroupViTTextConfig, GroupViTVisionConfig @@ -1227,19 +1228,16 @@ def __init__(self, config: GroupViTConfig): # Initialize weights and apply final processing self.post_init() - @filter_out_non_signature_kwargs() + @can_return_tuple @auto_docstring def get_text_features( self, input_ids: torch.Tensor, attention_mask: torch.Tensor | None = None, position_ids: torch.Tensor | None = None, - ) -> torch.FloatTensor: + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: r""" - Returns: - text_features (`torch.FloatTensor` of shape `(batch_size, output_dim`): The text embeddings obtained by - applying the projection layer to the pooled output of [`GroupViTTextModel`]. - Examples: ```python @@ -1257,18 +1255,22 @@ def get_text_features( input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids, + return_dict=True, + **kwargs, ) - text_features = self.text_projection(text_outputs.pooler_output) - return text_features + pooled_output = text_outputs.pooler_output + text_outputs.pooler_output = self.text_projection(pooled_output) - @filter_out_non_signature_kwargs() + return text_outputs + + @can_return_tuple @auto_docstring - def get_image_features(self, pixel_values: torch.Tensor) -> torch.FloatTensor: + def get_image_features( + self, + pixel_values: torch.Tensor, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: r""" - Returns: - image_features (`torch.FloatTensor` of shape `(batch_size, output_dim`): The image embeddings obtained by - applying the projection layer to the pooled output of [`GroupViTVisionModel`]. - Examples: ```python @@ -1287,9 +1289,10 @@ def get_image_features(self, pixel_values: torch.Tensor) -> torch.FloatTensor: >>> with torch.inference_mode(): ... image_features = model.get_image_features(**inputs) ```""" - vision_outputs: BaseModelOutputWithPooling = self.vision_model(pixel_values) - image_features = self.visual_projection(vision_outputs.pooler_output) - return image_features + vision_outputs: BaseModelOutputWithPooling = self.vision_model(pixel_values, return_dict=True, **kwargs) + vision_outputs.pooler_output = self.visual_projection(vision_outputs.pooler_output) + + return vision_outputs @auto_docstring def forward( diff --git a/src/transformers/models/idefics2/modeling_idefics2.py b/src/transformers/models/idefics2/modeling_idefics2.py index 196e256db317..e1ef275d519a 100644 --- a/src/transformers/models/idefics2/modeling_idefics2.py +++ b/src/transformers/models/idefics2/modeling_idefics2.py @@ -26,7 +26,7 @@ from ...modeling_attn_mask_utils import _prepare_4d_attention_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer -from ...modeling_outputs import BaseModelOutput, ModelOutput +from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling, ModelOutput from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel from ...processing_utils import Unpack from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, is_torchdynamo_compiling, logging @@ -827,17 +827,19 @@ def inputs_merger( inputs_embeds = inputs_embeds.masked_scatter(special_image_mask, image_hidden_states) return inputs_embeds + @can_return_tuple + @auto_docstring def get_image_features( - self, pixel_values: torch.FloatTensor, pixel_attention_mask: torch.LongTensor | None = None - ): - """ - Encodes images into continuous embeddings that can be forwarded to the language model. - - Args: - pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): - The tensors corresponding to the input images. - pixel_attention_mask (`torch.LongTensor`, *optional*): - The attention mask indicating padded regions in the image. + self, + pixel_values: torch.FloatTensor, + pixel_attention_mask: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): + The tensors corresponding to the input images. + pixel_attention_mask (`torch.LongTensor`, *optional*): + The attention mask indicating padded regions in the image. """ batch_size, num_images, num_channels, height, width = pixel_values.shape pixel_values = pixel_values.to(dtype=self.dtype) # fp16 compatibility @@ -865,15 +867,18 @@ def get_image_features( patches_subgrid = patches_subgrid.unfold(dimension=2, size=patch_size, step=patch_size) patch_attention_mask = (patches_subgrid.sum(dim=(-1, -2)) == patch_size * patch_size).bool() # Get sequence from the vision encoder - image_hidden_states = self.vision_model(pixel_values=pixel_values, patch_attention_mask=patch_attention_mask) - image_hidden_states = image_hidden_states.last_hidden_state + image_outputs = self.vision_model( + pixel_values=pixel_values, patch_attention_mask=patch_attention_mask, return_dict=True, **kwargs + ) + image_hidden_states = image_outputs.last_hidden_state # Modality projection & resampling - image_hidden_states = self.connector( + image_features = self.connector( image_hidden_states, attention_mask=patch_attention_mask.view(pixel_values.size(0), -1) ) - image_hidden_states = image_hidden_states.view(-1, image_hidden_states.shape[-1]) - return image_hidden_states + image_outputs.pooler_output = image_features.view(-1, image_features.shape[-1]) + + return image_outputs @can_return_tuple @auto_docstring( @@ -933,7 +938,9 @@ def forward( if pixel_values is not None and image_hidden_states is not None: raise ValueError("You cannot specify both pixel_values and image_hidden_states at the same time") elif pixel_values is not None: - image_hidden_states = self.get_image_features(pixel_values, pixel_attention_mask) + image_hidden_states = self.get_image_features( + pixel_values, pixel_attention_mask, return_dict=True + ).pooler_output elif image_hidden_states is not None: image_hidden_states = image_hidden_states.to(dtype=self.dtype, device=input_ids.device) @@ -991,10 +998,22 @@ def get_input_embeddings(self): def set_input_embeddings(self, value): self.model.text_model.set_input_embeddings(value) + @auto_docstring def get_image_features( - self, pixel_values: torch.FloatTensor, pixel_attention_mask: torch.LongTensor | None = None - ): - return self.model.get_image_features(pixel_values=pixel_values, pixel_attention_mask=pixel_attention_mask) + self, + pixel_values: torch.FloatTensor, + pixel_attention_mask: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): + The tensors corresponding to the input images. + pixel_attention_mask (`torch.LongTensor`, *optional*): + The attention mask indicating padded regions in the image. + """ + return self.model.get_image_features( + pixel_values=pixel_values, pixel_attention_mask=pixel_attention_mask, **kwargs + ) @can_return_tuple @auto_docstring diff --git a/src/transformers/models/idefics3/modeling_idefics3.py b/src/transformers/models/idefics3/modeling_idefics3.py index 705b3e9d6061..d1ac39b42115 100644 --- a/src/transformers/models/idefics3/modeling_idefics3.py +++ b/src/transformers/models/idefics3/modeling_idefics3.py @@ -25,7 +25,7 @@ from ...masking_utils import create_bidirectional_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer -from ...modeling_outputs import BaseModelOutput, ModelOutput +from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling, ModelOutput from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel from ...processing_utils import Unpack from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, logging @@ -570,17 +570,19 @@ def inputs_merger( inputs_embeds = inputs_embeds.masked_scatter(special_image_mask, image_hidden_states) return inputs_embeds + @can_return_tuple + @auto_docstring def get_image_features( - self, pixel_values: torch.FloatTensor, pixel_attention_mask: torch.LongTensor | None = None - ): - """ - Encodes images into continuous embeddings that can be forwarded to the language model. - - Args: - pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): - The tensors corresponding to the input images. - pixel_attention_mask (`torch.LongTensor`, *optional*): - The attention mask indicating padded regions in the image. + self, + pixel_values: torch.FloatTensor, + pixel_attention_mask: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): + The tensors corresponding to the input images. + pixel_attention_mask (`torch.LongTensor`, *optional*): + The attention mask indicating padded regions in the image. """ batch_size, num_images, num_channels, height, width = pixel_values.shape pixel_values = pixel_values.to(dtype=self.dtype) # fp16 compatibility @@ -609,12 +611,16 @@ def get_image_features( patch_attention_mask = (patches_subgrid.sum(dim=(-1, -2)) > 0).bool() # Get sequence from the vision encoder - image_hidden_states = self.vision_model(pixel_values=pixel_values, patch_attention_mask=patch_attention_mask) - image_hidden_states.last_hidden_state + image_outputs = self.vision_model( + pixel_values=pixel_values, patch_attention_mask=patch_attention_mask, return_dict=True, **kwargs + ) + image_hidden_states = image_outputs.last_hidden_state # Modality projection & resampling - image_hidden_states = self.connector(image_hidden_states.last_hidden_state) - return image_hidden_states + image_features = self.connector(image_hidden_states) + image_outputs.pooler_output = image_features + + return image_outputs @can_return_tuple @auto_docstring( @@ -682,7 +688,9 @@ def forward( if pixel_values is not None and image_hidden_states is not None: raise ValueError("You cannot specify both pixel_values and image_hidden_states at the same time") elif pixel_values is not None: - image_hidden_states = self.get_image_features(pixel_values, pixel_attention_mask) + image_hidden_states = self.get_image_features( + pixel_values, pixel_attention_mask, return_dict=True + ).pooler_output elif image_hidden_states is not None: image_hidden_states = image_hidden_states.to(dtype=self.dtype, device=input_ids.device) @@ -745,10 +753,22 @@ def get_input_embeddings(self): def set_input_embeddings(self, value): self.model.text_model.set_input_embeddings(value) + @auto_docstring def get_image_features( - self, pixel_values: torch.FloatTensor, pixel_attention_mask: torch.LongTensor | None = None - ): - return self.model.get_image_features(pixel_values=pixel_values, pixel_attention_mask=pixel_attention_mask) + self, + pixel_values: torch.FloatTensor, + pixel_attention_mask: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): + The tensors corresponding to the input images. + pixel_attention_mask (`torch.LongTensor`, *optional*): + The attention mask indicating padded regions in the image. + """ + return self.model.get_image_features( + pixel_values=pixel_values, pixel_attention_mask=pixel_attention_mask, **kwargs + ) @can_return_tuple @auto_docstring diff --git a/src/transformers/models/instructblip/modeling_instructblip.py b/src/transformers/models/instructblip/modeling_instructblip.py index 79f1565099c1..78ab67d2c09a 100644 --- a/src/transformers/models/instructblip/modeling_instructblip.py +++ b/src/transformers/models/instructblip/modeling_instructblip.py @@ -31,6 +31,8 @@ BaseModelOutputWithPastAndCrossAttentions, BaseModelOutputWithPooling, BaseModelOutputWithPoolingAndCrossAttentions, + CausalLMOutputWithPast, + Seq2SeqLMOutput, ) from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel from ...processing_utils import Unpack @@ -44,6 +46,20 @@ logger = logging.get_logger(__name__) +@dataclass +@auto_docstring +class BaseModelOutputWithVisionQformerOutputs(BaseModelOutputWithPooling): + r""" + vision_outputs (`BaseModelOutputWithPooling`): + Outputs of the vision encoder. + qformer_outputs (`BaseModelOutputWithPoolingAndCrossAttentions`): + Outputs of the Q-Former (Querying Transformer). + """ + + vision_outputs: BaseModelOutputWithPooling | None = None + qformer_outputs: BaseModelOutputWithPoolingAndCrossAttentions | None = None + + @dataclass @auto_docstring( custom_intro=""" @@ -67,9 +83,9 @@ class InstructBlipForConditionalGenerationModelOutput(ModelOutput): loss: tuple[torch.FloatTensor] | None = None logits: tuple[torch.FloatTensor] | None = None - vision_outputs: torch.FloatTensor | None = None - qformer_outputs: tuple[torch.FloatTensor] | None = None - language_model_outputs: tuple[torch.FloatTensor] | None = None + vision_outputs: BaseModelOutputWithPooling | None = None + qformer_outputs: BaseModelOutputWithPoolingAndCrossAttentions | None = None + language_model_outputs: CausalLMOutputWithPast | Seq2SeqLMOutput | None = None def to_tuple(self) -> tuple[Any]: return tuple( @@ -371,7 +387,6 @@ def forward( return BaseModelOutput(last_hidden_state=hidden_states) -# Copied from transformers.models.blip.modeling_blip.BlipVisionModel with Blip->InstructBlip, BLIP->INSTRUCTBLIP class InstructBlipVisionModel(InstructBlipPreTrainedModel): main_input_name = "pixel_values" input_modalities = ("image",) @@ -1173,28 +1188,44 @@ def _preprocess_accelerate(self): if hasattr(self.language_model, "_hf_hook"): self.language_model._hf_hook.io_same_device = True # For `generate` compatibility + @can_return_tuple + @auto_docstring def get_image_features( self, pixel_values: torch.FloatTensor, qformer_input_ids: torch.LongTensor, qformer_attention_mask: torch.LongTensor | None = None, interpolate_pos_encoding: bool | None = False, - return_dict: bool | None = False, - ): - """ - Encodes images into continuous embeddings that can be forwarded to the language model. + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithVisionQformerOutputs: + r""" + pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): + The tensors corresponding to the input images. + qformer_input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Indices of input sequence tokens in the vocabulary of the Q-Former. Input tokens can optionally be provided + to serve as text prompt, which the Q-Former model will encode. - Args: - pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): - The tensors corresponding to the input images. + Indices can be obtained using [`InstructBlipProcessor`]. See [`InstructBlipProcessor.__call__`] for + details. + + [What are input IDs?](../glossary#input-ids) + qformer_attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): + Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`: + + - 1 for tokens that are **not masked**, + - 0 for tokens that are **masked**. + + [What are attention masks?](../glossary#attention-mask) """ # step 1: forward the images through the vision encoder, # to get image embeddings of shape (batch_size, seq_len, hidden_size) - vision_outputs = self.vision_model( + vision_outputs: BaseModelOutputWithPooling = self.vision_model( pixel_values=pixel_values, interpolate_pos_encoding=interpolate_pos_encoding, return_dict=True, + **kwargs, ) + vision_outputs = BaseModelOutputWithVisionQformerOutputs(**vision_outputs, vision_outputs=vision_outputs) image_embeds = vision_outputs[0] # step 2: forward the query tokens through the QFormer, using the image embeddings for cross-attention @@ -1206,21 +1237,23 @@ def get_image_features( if qformer_attention_mask is None: qformer_attention_mask = torch.ones_like(qformer_input_ids) qformer_attention_mask = torch.cat([query_attention_mask, qformer_attention_mask], dim=1) - query_outputs = self.qformer( + qformer_outputs = self.qformer( input_ids=qformer_input_ids, attention_mask=qformer_attention_mask, query_embeds=query_tokens, encoder_hidden_states=image_embeds, encoder_attention_mask=image_attention_mask, return_dict=True, + **kwargs, ) - query_output = query_outputs[0][:, : query_tokens.size(1), :] + vision_outputs.qformer_outputs = qformer_outputs + query_output = qformer_outputs[0][:, : query_tokens.size(1), :] # step 3: use the language model, conditioned on the query outputs and the prompt - language_model_inputs = self.language_projection(query_output) - if return_dict: - return language_model_inputs, vision_outputs, query_outputs - return language_model_inputs + image_features = self.language_projection(query_output) + vision_outputs.pooler_output = image_features + + return vision_outputs def get_placeholder_mask(self, input_ids: torch.LongTensor, inputs_embeds: torch.FloatTensor): """ @@ -1314,13 +1347,16 @@ def forward( The unusual aspect of this image is that a man is ironing clothes on the back of a yellow SUV, which is parked in the middle of a busy city street. This is an unconventional approach to ironing clothes, as it requires the man to balance himself and his ironing equipment on top of the vehicle while navigating through traffic. Additionally, the presence of taxis and other vehicles in the scene further emphasizes the unusual nature of this situation. ```""" - language_model_inputs, vision_outputs, query_outputs = self.get_image_features( + image_features: BaseModelOutputWithVisionQformerOutputs = self.get_image_features( pixel_values, qformer_input_ids=qformer_input_ids, qformer_attention_mask=qformer_attention_mask, interpolate_pos_encoding=interpolate_pos_encoding, return_dict=True, ) + language_model_inputs = image_features.pooler_output + qformer_outputs = image_features.qformer_outputs + vision_outputs = image_features.vision_outputs if inputs_embeds is None: inputs_embeds = self.get_input_embeddings()(input_ids) @@ -1362,7 +1398,7 @@ def forward( loss=loss, logits=logits, vision_outputs=vision_outputs, - qformer_outputs=query_outputs, + qformer_outputs=qformer_outputs, language_model_outputs=outputs, ) @@ -1405,13 +1441,14 @@ def generate( self._preprocess_accelerate() batch_size = pixel_values.shape[0] - language_model_inputs, vision_outputs, query_outputs = self.get_image_features( + image_features: BaseModelOutputWithVisionQformerOutputs = self.get_image_features( pixel_values, qformer_input_ids=qformer_input_ids, qformer_attention_mask=qformer_attention_mask, interpolate_pos_encoding=interpolate_pos_encoding, return_dict=True, ) + language_model_inputs = image_features.pooler_output if inputs_embeds is None: if input_ids is None: diff --git a/src/transformers/models/instructblipvideo/modeling_instructblipvideo.py b/src/transformers/models/instructblipvideo/modeling_instructblipvideo.py index b21d0c4739d1..121db617af5c 100644 --- a/src/transformers/models/instructblipvideo/modeling_instructblipvideo.py +++ b/src/transformers/models/instructblipvideo/modeling_instructblipvideo.py @@ -37,6 +37,8 @@ BaseModelOutputWithPastAndCrossAttentions, BaseModelOutputWithPooling, BaseModelOutputWithPoolingAndCrossAttentions, + CausalLMOutputWithPast, + Seq2SeqLMOutput, ) from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel from ...processing_utils import Unpack @@ -906,9 +908,9 @@ class InstructBlipVideoForConditionalGenerationModelOutput(ModelOutput): loss: tuple[torch.FloatTensor] | None = None logits: tuple[torch.FloatTensor] | None = None - vision_outputs: torch.FloatTensor | None = None - qformer_outputs: tuple[torch.FloatTensor] | None = None - language_model_outputs: tuple[torch.FloatTensor] | None = None + vision_outputs: BaseModelOutputWithPooling | None = None + qformer_outputs: BaseModelOutputWithPoolingAndCrossAttentions | None = None + language_model_outputs: CausalLMOutputWithPast | Seq2SeqLMOutput | None = None def to_tuple(self) -> tuple[Any]: return tuple( @@ -1120,6 +1122,20 @@ def forward( ) +@dataclass +@auto_docstring +class BaseModelOutputWithVisionQformerOutputs(BaseModelOutputWithPooling): + r""" + vision_outputs (`BaseModelOutputWithPooling`): + Outputs of the vision encoder. + qformer_outputs (`BaseModelOutputWithPoolingAndCrossAttentions`): + Outputs of the Q-Former (Querying Transformer). + """ + + vision_outputs: BaseModelOutputWithPooling | None = None + qformer_outputs: BaseModelOutputWithPoolingAndCrossAttentions | None = None + + @auto_docstring( custom_intro=""" InstructBlipVideo Model for generating text given an image and an optional text prompt. The model consists of a vision @@ -1203,23 +1219,6 @@ def _preprocess_accelerate(self): if hasattr(self.language_model, "_hf_hook"): self.language_model._hf_hook.io_same_device = True # For `generate` compatibility - def get_image_features( - self, - pixel_values: torch.FloatTensor, - qformer_input_ids: torch.LongTensor, - qformer_attention_mask: torch.LongTensor | None = None, - interpolate_pos_encoding: bool | None = False, - return_dict: bool | None = False, - ): - """ - Encodes images into continuous embeddings that can be forwarded to the language model. - - Args: - pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): - The tensors corresponding to the input images. - """ - pass - def get_placeholder_mask(self, input_ids: torch.LongTensor, inputs_embeds: torch.FloatTensor): """ Obtains multimodal placeholder mask from `input_ids` or `inputs_embeds`. @@ -1320,15 +1319,17 @@ def forward( ```""" return_dict = return_dict if return_dict is not None else self.config.use_return_dict - language_model_inputs, vision_outputs, query_outputs = self.get_video_features( + video_features: BaseModelOutputWithVisionQformerOutputs = self.get_video_features( pixel_values, qformer_input_ids=qformer_input_ids, qformer_attention_mask=qformer_attention_mask, interpolate_pos_encoding=interpolate_pos_encoding, return_dict=True, + **kwargs, ) - vision_outputs = vision_outputs.to_tuple() if not return_dict else vision_outputs - query_outputs = query_outputs.to_tuple() if not return_dict else query_outputs + language_model_inputs = video_features.pooler_output + qformer_outputs = video_features.qformer_outputs + vision_outputs = video_features.vision_outputs if inputs_embeds is None: inputs_embeds = self.get_input_embeddings()(input_ids) @@ -1377,7 +1378,7 @@ def forward( loss=loss, logits=logits, vision_outputs=vision_outputs, - qformer_outputs=query_outputs, + qformer_outputs=qformer_outputs, language_model_outputs=outputs, ) @@ -1420,13 +1421,14 @@ def generate( self._preprocess_accelerate() batch_size = pixel_values.shape[0] - language_model_inputs, vision_outputs, query_outputs = self.get_video_features( + video_features: BaseModelOutputWithVisionQformerOutputs = self.get_video_features( pixel_values, qformer_input_ids=qformer_input_ids, qformer_attention_mask=qformer_attention_mask, interpolate_pos_encoding=interpolate_pos_encoding, return_dict=True, ) + language_model_inputs = video_features.pooler_output if inputs_embeds is None: if input_ids is None: @@ -1451,30 +1453,42 @@ def generate( return outputs + @can_return_tuple + @auto_docstring def get_video_features( self, pixel_values: torch.FloatTensor, qformer_input_ids: torch.LongTensor, qformer_attention_mask: torch.LongTensor | None = None, interpolate_pos_encoding: bool | None = False, - return_dict: bool | None = False, - ): - """ - Encodes images into continuous embeddings that can be forwarded to the language model. - - Args: - pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): - The tensors corresponding to the input images. + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithVisionQformerOutputs: + r""" + pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): + The tensors corresponding to the input images. + qformer_input_ids (`torch.LongTensor` of shape (batch_size, sequence_length)): + The sequence used as a prompt to be fed to the Q-Former module. + qformer_attention_mask (`torch.LongTensor` of shape (batch_size, sequence_length), *optional*): + Mask to avoid performing attention on padding token indices. """ # step 1: forward the images through the vision encoder, # we process in a batched way, later unbatch it back (video has frames=4 always) batch_size, frames, channel, height, width = pixel_values.shape pixel_values = pixel_values.reshape(batch_size * frames, channel, height, width) - vision_outputs = self.vision_model( + vision_outputs: BaseModelOutputWithPooling = self.vision_model( pixel_values=pixel_values, interpolate_pos_encoding=interpolate_pos_encoding, return_dict=True, + **kwargs, + ) + vision_outputs = BaseModelOutputWithVisionQformerOutputs( + last_hidden_state=vision_outputs.last_hidden_state, + pooler_output=vision_outputs.pooler_output, + hidden_states=vision_outputs.hidden_states, + attentions=vision_outputs.attentions, + vision_outputs=vision_outputs, + qformer_outputs=None, ) image_embeds = vision_outputs[0] @@ -1491,24 +1505,26 @@ def get_video_features( qformer_input_ids = qformer_input_ids.repeat_interleave(frames, dim=0) qformer_attention_mask = qformer_attention_mask.repeat_interleave(frames, dim=0) qformer_attention_mask = torch.cat([query_attention_mask, qformer_attention_mask], dim=1) - query_outputs = self.qformer( + qformer_outputs = self.qformer( input_ids=qformer_input_ids, attention_mask=qformer_attention_mask, query_embeds=query_tokens, encoder_hidden_states=image_embeds, encoder_attention_mask=image_attention_mask, return_dict=True, + **kwargs, ) - query_output = query_outputs[0][:, : query_tokens.size(1), :] + vision_outputs.qformer_outputs = qformer_outputs + query_output = qformer_outputs[0][:, : query_tokens.size(1), :] # step 3: use the language model, conditioned on the query outputs and the prompt - language_model_inputs = self.language_projection(query_output) + video_features = self.language_projection(query_output) # unbatch inputs back, each video-frame gets `num_query_tokens` seq length - language_model_inputs = language_model_inputs.reshape(batch_size, self.config.num_query_tokens * frames, -1) - if return_dict: - return language_model_inputs, vision_outputs, query_outputs - return language_model_inputs + video_features = video_features.reshape(batch_size, self.config.num_query_tokens * frames, -1) + vision_outputs.pooler_output = video_features + + return vision_outputs __all__ = [ diff --git a/src/transformers/models/instructblipvideo/modular_instructblipvideo.py b/src/transformers/models/instructblipvideo/modular_instructblipvideo.py index 29e0e24479cf..8b76bfc6fba1 100644 --- a/src/transformers/models/instructblipvideo/modular_instructblipvideo.py +++ b/src/transformers/models/instructblipvideo/modular_instructblipvideo.py @@ -20,6 +20,7 @@ InstructBlipVisionConfig, ) from transformers.models.instructblip.modeling_instructblip import ( + BaseModelOutputWithVisionQformerOutputs, InstructBlipForConditionalGeneration, InstructBlipForConditionalGenerationModelOutput, InstructBlipModel, @@ -31,9 +32,10 @@ from ...configuration_utils import PreTrainedConfig from ...modeling_flash_attention_utils import FlashAttentionKwargs +from ...modeling_outputs import BaseModelOutputWithPooling from ...models.auto.modeling_auto import MODEL_FOR_CAUSAL_LM_MAPPING_NAMES from ...processing_utils import Unpack -from ...utils import logging +from ...utils import auto_docstring, can_return_tuple, logging from ..auto import CONFIG_MAPPING, AutoConfig @@ -283,30 +285,42 @@ def forward( class InstructBlipVideoForConditionalGeneration(InstructBlipForConditionalGeneration): + @can_return_tuple + @auto_docstring def get_video_features( self, pixel_values: torch.FloatTensor, qformer_input_ids: torch.LongTensor, qformer_attention_mask: torch.LongTensor | None = None, interpolate_pos_encoding: bool | None = False, - return_dict: bool | None = False, - ): - """ - Encodes images into continuous embeddings that can be forwarded to the language model. - - Args: - pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): - The tensors corresponding to the input images. + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithVisionQformerOutputs: + r""" + pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): + The tensors corresponding to the input images. + qformer_input_ids (`torch.LongTensor` of shape (batch_size, sequence_length)): + The sequence used as a prompt to be fed to the Q-Former module. + qformer_attention_mask (`torch.LongTensor` of shape (batch_size, sequence_length), *optional*): + Mask to avoid performing attention on padding token indices. """ # step 1: forward the images through the vision encoder, # we process in a batched way, later unbatch it back (video has frames=4 always) batch_size, frames, channel, height, width = pixel_values.shape pixel_values = pixel_values.reshape(batch_size * frames, channel, height, width) - vision_outputs = self.vision_model( + vision_outputs: BaseModelOutputWithPooling = self.vision_model( pixel_values=pixel_values, interpolate_pos_encoding=interpolate_pos_encoding, return_dict=True, + **kwargs, + ) + vision_outputs = BaseModelOutputWithVisionQformerOutputs( + last_hidden_state=vision_outputs.last_hidden_state, + pooler_output=vision_outputs.pooler_output, + hidden_states=vision_outputs.hidden_states, + attentions=vision_outputs.attentions, + vision_outputs=vision_outputs, + qformer_outputs=None, ) image_embeds = vision_outputs[0] @@ -323,35 +337,29 @@ def get_video_features( qformer_input_ids = qformer_input_ids.repeat_interleave(frames, dim=0) qformer_attention_mask = qformer_attention_mask.repeat_interleave(frames, dim=0) qformer_attention_mask = torch.cat([query_attention_mask, qformer_attention_mask], dim=1) - query_outputs = self.qformer( + qformer_outputs = self.qformer( input_ids=qformer_input_ids, attention_mask=qformer_attention_mask, query_embeds=query_tokens, encoder_hidden_states=image_embeds, encoder_attention_mask=image_attention_mask, return_dict=True, + **kwargs, ) - query_output = query_outputs[0][:, : query_tokens.size(1), :] + vision_outputs.qformer_outputs = qformer_outputs + query_output = qformer_outputs[0][:, : query_tokens.size(1), :] # step 3: use the language model, conditioned on the query outputs and the prompt - language_model_inputs = self.language_projection(query_output) + video_features = self.language_projection(query_output) # unbatch inputs back, each video-frame gets `num_query_tokens` seq length - language_model_inputs = language_model_inputs.reshape(batch_size, self.config.num_query_tokens * frames, -1) - if return_dict: - return language_model_inputs, vision_outputs, query_outputs - return language_model_inputs + video_features = video_features.reshape(batch_size, self.config.num_query_tokens * frames, -1) + vision_outputs.pooler_output = video_features - # Model supports only videos - def get_image_features( - self, - pixel_values: torch.FloatTensor, - qformer_input_ids: torch.LongTensor, - qformer_attention_mask: torch.LongTensor | None = None, - interpolate_pos_encoding: bool | None = False, - return_dict: bool | None = False, - ): - pass + return vision_outputs + + def get_image_features(**super_kwargs): + raise AttributeError("No need to inherit as this architecture only supports videos.") def get_placeholder_mask(self, input_ids: torch.LongTensor, inputs_embeds: torch.FloatTensor): """ @@ -451,15 +459,17 @@ def forward( ```""" return_dict = return_dict if return_dict is not None else self.config.use_return_dict - language_model_inputs, vision_outputs, query_outputs = self.get_video_features( + video_features: BaseModelOutputWithVisionQformerOutputs = self.get_video_features( pixel_values, qformer_input_ids=qformer_input_ids, qformer_attention_mask=qformer_attention_mask, interpolate_pos_encoding=interpolate_pos_encoding, return_dict=True, + **kwargs, ) - vision_outputs = vision_outputs.to_tuple() if not return_dict else vision_outputs - query_outputs = query_outputs.to_tuple() if not return_dict else query_outputs + language_model_inputs = video_features.pooler_output + qformer_outputs = video_features.qformer_outputs + vision_outputs = video_features.vision_outputs if inputs_embeds is None: inputs_embeds = self.get_input_embeddings()(input_ids) @@ -508,7 +518,7 @@ def forward( loss=loss, logits=logits, vision_outputs=vision_outputs, - qformer_outputs=query_outputs, + qformer_outputs=qformer_outputs, language_model_outputs=outputs, ) @@ -551,13 +561,14 @@ def generate( self._preprocess_accelerate() batch_size = pixel_values.shape[0] - language_model_inputs, vision_outputs, query_outputs = self.get_video_features( + video_features: BaseModelOutputWithVisionQformerOutputs = self.get_video_features( pixel_values, qformer_input_ids=qformer_input_ids, qformer_attention_mask=qformer_attention_mask, interpolate_pos_encoding=interpolate_pos_encoding, return_dict=True, ) + language_model_inputs = video_features.pooler_output if inputs_embeds is None: if input_ids is None: diff --git a/src/transformers/models/internvl/modeling_internvl.py b/src/transformers/models/internvl/modeling_internvl.py index 5176a5776630..498caae1044e 100644 --- a/src/transformers/models/internvl/modeling_internvl.py +++ b/src/transformers/models/internvl/modeling_internvl.py @@ -35,14 +35,7 @@ from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPast, BaseModelOutputWithPooling from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel from ...processing_utils import Unpack -from ...utils import ( - ModelOutput, - TransformersKwargs, - auto_docstring, - can_return_tuple, - torch_compilable_check, - torch_int, -) +from ...utils import ModelOutput, TransformersKwargs, auto_docstring, torch_compilable_check, torch_int from ...utils.generic import check_model_inputs from ..auto import AutoModel from .configuration_internvl import InternVLConfig, InternVLVisionConfig @@ -551,39 +544,33 @@ def get_input_embeddings(self): def set_input_embeddings(self, value): self.language_model.set_input_embeddings(value) + @check_model_inputs(tie_last_hidden_states=False) + @auto_docstring( + custom_intro="Obtains image last hidden states from the vision tower and apply multimodal projection." + ) def get_image_features( self, pixel_values: torch.FloatTensor, vision_feature_layer: int | list[int] | None = None, vision_feature_select_strategy: str | None = None, - **kwargs, - ): - """ - Obtains image last hidden states from the vision tower and apply multimodal projection. - - Args: - pixel_values (`torch.FloatTensor]` of shape `(batch_size, channels, height, width)`) - The tensors corresponding to the input images. - vision_feature_layer (`int` or `list[int]`): - Layer index or list of layer indices to extract features from. - Returns: - vision_features (`torch.Tensor`): Image feature tensor of shape `(num_images, image_length, embed_dim)`. + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + pixel_values (`torch.FloatTensor]` of shape `(batch_size, channels, height, width)`) + The tensors corresponding to the input images. + vision_feature_layer (`int` or `list[int]`): + Layer index or list of layer indices to extract features from. """ - vision_feature_layer = ( - vision_feature_layer if vision_feature_layer is not None else self.config.vision_feature_layer - ) - vision_feature_select_strategy = ( - vision_feature_select_strategy - if vision_feature_select_strategy is not None - else self.config.vision_feature_select_strategy - ) pixel_values = pixel_values.to(dtype=self.dtype) # fp16 compatibility downsample_ratio = self.config.downsample_ratio + if vision_feature_layer != -1: + kwargs["output_hidden_states"] = True + vision_outputs = self.vision_tower(pixel_values=pixel_values, return_dict=True, **kwargs) if vision_feature_layer == -1: - vision_features = self.vision_tower(pixel_values=pixel_values).last_hidden_state + vision_features = vision_outputs.last_hidden_state else: - vision_features = self.vision_model(pixel_values=pixel_values).hidden_states[vision_feature_layer] + vision_features = vision_outputs.hidden_states[vision_feature_layer] if vision_feature_select_strategy == "default": vision_features = vision_features[:, 1:, :] @@ -603,7 +590,9 @@ def get_image_features( # Project features through multi-modal projector vision_features = self.multi_modal_projector(vision_features) - return vision_features + vision_outputs.pooler_output = vision_features + + return vision_outputs def get_placeholder_mask( self, input_ids: torch.LongTensor, inputs_embeds: torch.FloatTensor, image_features: torch.FloatTensor @@ -629,7 +618,7 @@ def get_placeholder_mask( ) return special_image_mask - @can_return_tuple + @check_model_inputs(tie_last_hidden_states=False) @auto_docstring def forward( self, @@ -644,15 +633,6 @@ def forward( cache_position: torch.LongTensor | None = None, **kwargs: Unpack[TransformersKwargs], ) -> tuple | InternVLModelOutputWithPast: - vision_feature_layer = ( - vision_feature_layer if vision_feature_layer is not None else self.config.vision_feature_layer - ) - vision_feature_select_strategy = ( - vision_feature_select_strategy - if vision_feature_select_strategy is not None - else self.config.vision_feature_select_strategy - ) - if (input_ids is None) ^ (inputs_embeds is not None): raise ValueError("You must specify exactly one of input_ids or inputs_embeds") @@ -664,7 +644,8 @@ def forward( pixel_values=pixel_values, vision_feature_layer=vision_feature_layer, vision_feature_select_strategy=vision_feature_select_strategy, - ) + return_dict=True, + ).pooler_output image_features = image_features.to(inputs_embeds.device, inputs_embeds.dtype) special_image_mask = self.get_placeholder_mask( input_ids, inputs_embeds=inputs_embeds, image_features=image_features @@ -783,13 +764,14 @@ def set_input_embeddings(self, value): def get_output_embeddings(self) -> nn.Module: return self.lm_head + @auto_docstring def get_image_features( self, pixel_values: torch.FloatTensor, vision_feature_layer: int | list[int] | None = None, vision_feature_select_strategy: str | None = None, - **kwargs, - ): + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: return self.model.get_image_features( pixel_values=pixel_values, vision_feature_layer=vision_feature_layer, @@ -797,7 +779,7 @@ def get_image_features( **kwargs, ) - @can_return_tuple + @check_model_inputs(tie_last_hidden_states=False) @auto_docstring def forward( self, @@ -850,15 +832,6 @@ def forward( >>> print(processor.decode(generate_ids[0, inputs["input_ids"].shape[1] :], skip_special_tokens=True)) The images depict the Statue of Liberty and the Golden Gate Bridge. ```""" - vision_feature_layer = ( - vision_feature_layer if vision_feature_layer is not None else self.config.vision_feature_layer - ) - vision_feature_select_strategy = ( - vision_feature_select_strategy - if vision_feature_select_strategy is not None - else self.config.vision_feature_select_strategy - ) - outputs = self.model( input_ids=input_ids, pixel_values=pixel_values, diff --git a/src/transformers/models/internvl/modular_internvl.py b/src/transformers/models/internvl/modular_internvl.py index 25f031577ca2..d5ec73e498da 100644 --- a/src/transformers/models/internvl/modular_internvl.py +++ b/src/transformers/models/internvl/modular_internvl.py @@ -27,7 +27,7 @@ from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel from ...processing_utils import Unpack -from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, torch_int +from ...utils import TransformersKwargs, auto_docstring, torch_int from ...utils.generic import check_model_inputs from ..clip.modeling_clip import CLIPMLP from ..janus.modeling_janus import JanusVisionAttention @@ -484,39 +484,33 @@ def pixel_shuffle(self, vision_features: torch.Tensor, scale_factor: float = 0.5 return vision_features + @check_model_inputs(tie_last_hidden_states=False) + @auto_docstring( + custom_intro="Obtains image last hidden states from the vision tower and apply multimodal projection." + ) def get_image_features( self, pixel_values: torch.FloatTensor, vision_feature_layer: int | list[int] | None = None, vision_feature_select_strategy: str | None = None, - **kwargs, - ): - """ - Obtains image last hidden states from the vision tower and apply multimodal projection. - - Args: - pixel_values (`torch.FloatTensor]` of shape `(batch_size, channels, height, width)`) - The tensors corresponding to the input images. - vision_feature_layer (`int` or `list[int]`): - Layer index or list of layer indices to extract features from. - Returns: - vision_features (`torch.Tensor`): Image feature tensor of shape `(num_images, image_length, embed_dim)`. + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + pixel_values (`torch.FloatTensor]` of shape `(batch_size, channels, height, width)`) + The tensors corresponding to the input images. + vision_feature_layer (`int` or `list[int]`): + Layer index or list of layer indices to extract features from. """ - vision_feature_layer = ( - vision_feature_layer if vision_feature_layer is not None else self.config.vision_feature_layer - ) - vision_feature_select_strategy = ( - vision_feature_select_strategy - if vision_feature_select_strategy is not None - else self.config.vision_feature_select_strategy - ) pixel_values = pixel_values.to(dtype=self.dtype) # fp16 compatibility downsample_ratio = self.config.downsample_ratio + if vision_feature_layer != -1: + kwargs["output_hidden_states"] = True + vision_outputs = self.vision_tower(pixel_values=pixel_values, return_dict=True, **kwargs) if vision_feature_layer == -1: - vision_features = self.vision_tower(pixel_values=pixel_values).last_hidden_state + vision_features = vision_outputs.last_hidden_state else: - vision_features = self.vision_model(pixel_values=pixel_values).hidden_states[vision_feature_layer] + vision_features = vision_outputs.hidden_states[vision_feature_layer] if vision_feature_select_strategy == "default": vision_features = vision_features[:, 1:, :] @@ -536,9 +530,11 @@ def get_image_features( # Project features through multi-modal projector vision_features = self.multi_modal_projector(vision_features) - return vision_features + vision_outputs.pooler_output = vision_features - @can_return_tuple + return vision_outputs + + @check_model_inputs(tie_last_hidden_states=False) @auto_docstring def forward( self, @@ -553,15 +549,6 @@ def forward( cache_position: torch.LongTensor | None = None, **kwargs: Unpack[TransformersKwargs], ) -> tuple | InternVLModelOutputWithPast: - vision_feature_layer = ( - vision_feature_layer if vision_feature_layer is not None else self.config.vision_feature_layer - ) - vision_feature_select_strategy = ( - vision_feature_select_strategy - if vision_feature_select_strategy is not None - else self.config.vision_feature_select_strategy - ) - if (input_ids is None) ^ (inputs_embeds is not None): raise ValueError("You must specify exactly one of input_ids or inputs_embeds") @@ -573,7 +560,8 @@ def forward( pixel_values=pixel_values, vision_feature_layer=vision_feature_layer, vision_feature_select_strategy=vision_feature_select_strategy, - ) + return_dict=True, + ).pooler_output image_features = image_features.to(inputs_embeds.device, inputs_embeds.dtype) special_image_mask = self.get_placeholder_mask( input_ids, inputs_embeds=inputs_embeds, image_features=image_features diff --git a/src/transformers/models/janus/modeling_janus.py b/src/transformers/models/janus/modeling_janus.py index 5a912358700d..7ef79803672f 100644 --- a/src/transformers/models/janus/modeling_janus.py +++ b/src/transformers/models/janus/modeling_janus.py @@ -431,136 +431,14 @@ def forward( return BaseModelOutput(last_hidden_state=hidden_states) -class JanusAttention(nn.Module): - """Multi-headed attention from 'Attention Is All You Need' paper""" - - def __init__(self, config): - super().__init__() - self.config = config - self.embed_dim = config.hidden_size - self.num_heads = config.num_attention_heads - self.head_dim = self.embed_dim // self.num_heads - if self.head_dim * self.num_heads != self.embed_dim: - raise ValueError( - f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim} and `num_heads`:" - f" {self.num_heads})." - ) - self.scale = self.head_dim**-0.5 - self.is_causal = False - self.attention_dropout = config.attention_dropout - - # small tweak here compared to CLIP, no bias here - self.qkv = nn.Linear(self.embed_dim, 3 * self.embed_dim, bias=False) - - if config.qkv_bias: - q_bias = nn.Parameter(torch.zeros(self.embed_dim)) - v_bias = nn.Parameter(torch.zeros(self.embed_dim)) - else: - q_bias = None - v_bias = None - - if q_bias is not None: - qkv_bias = torch.cat((q_bias, torch.zeros_like(v_bias, requires_grad=False), v_bias)) - self.qkv.bias = nn.Parameter(qkv_bias) - - self.projection = nn.Linear(self.embed_dim, self.embed_dim) - - def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int): - return tensor.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous() - - def forward( - self, - hidden_states: torch.Tensor, - **kwargs, - ) -> tuple[torch.Tensor, torch.Tensor | None, tuple[torch.Tensor] | None]: - """Input shape: Batch x Time x Channel""" - - bsz, tgt_len, embed_dim = hidden_states.size() - - mixed_qkv = self.qkv(hidden_states) - - mixed_qkv = mixed_qkv.reshape(bsz, tgt_len, 3, self.num_heads, embed_dim // self.num_heads).permute( - 2, 0, 3, 1, 4 - ) - query_states, key_states, value_states = mixed_qkv[0], mixed_qkv[1], mixed_qkv[2] - - attention_interface: Callable = eager_attention_forward - - if self.config._attn_implementation != "eager": - attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation] - - attn_output, attn_weights = attention_interface( - self, - query_states, - key_states, - value_states, - attention_mask=None, - dropout=0.0 if not self.training else self.attention_dropout, - scaling=self.scale, - **kwargs, - ) - - attn_output = attn_output.reshape(bsz, tgt_len, -1).contiguous() - attn_output = self.projection(attn_output) - - return attn_output, attn_weights - - -class JanusMLP(nn.Module): - def __init__(self, config): - super().__init__() - self.config = config - self.activation_fn = ACT2FN[config.hidden_act] - self.fc1 = nn.Linear(config.hidden_size, config.intermediate_size) - self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size) - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - hidden_states = self.fc1(hidden_states) - hidden_states = self.activation_fn(hidden_states) - hidden_states = self.fc2(hidden_states) - return hidden_states - - -class JanusEncoderLayer(GradientCheckpointingLayer): - def __init__(self, config: JanusConfig): - super().__init__() - self.embed_dim = config.hidden_size - self.self_attn = JanusAttention(config) - self.layer_norm1 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps) - self.mlp = JanusMLP(config) - self.layer_norm2 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps) - - @auto_docstring - def forward( - self, - hidden_states: torch.Tensor, - **kwargs: Unpack[TransformersKwargs], - ) -> torch.FloatTensor: - residual = hidden_states - - hidden_states = self.layer_norm1(hidden_states) - hidden_states, _ = self.self_attn( - hidden_states=hidden_states, - **kwargs, - ) - hidden_states = hidden_states + residual - residual = hidden_states - hidden_states = self.layer_norm2(hidden_states) - hidden_states = self.mlp(hidden_states) - - hidden_states = hidden_states + residual - - return hidden_states - - @auto_docstring class JanusVisionModel(JanusPreTrainedModel): main_input_name = "pixel_values" input_modalities = ("image",) config: JanusVisionConfig _can_record_outputs = { - "hidden_states": JanusEncoderLayer, - "attentions": JanusAttention, + "hidden_states": JanusVisionEncoderLayer, + "attentions": JanusVisionAttention, } def __init__(self, config: JanusVisionConfig): @@ -957,6 +835,23 @@ def forward(self, hidden_state: torch.FloatTensor) -> torch.FloatTensor: return hidden_state +@dataclass +@auto_docstring +class JanusVQVAEModelOutput(BaseModelOutputWithPooling): + r""" + quantized_last_hidden_state (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): + Quantized last hidden state from the VQ-VAE model. + image_tokens (`torch.FloatTensor` of shape `(batch_size, config.vocab_size`): + Indices of the image tokens predicted by the VQ-VAE model. + embedding_loss (`torch.FloatTensor`): + The embedding loss computed during quantization. + """ + + quantized_last_hidden_state: torch.FloatTensor | None = None + image_tokens: torch.FloatTensor | None = None + embedding_loss: torch.FloatTensor | None = None + + @auto_docstring( custom_intro=""" The VQ-VAE model used in Janus for encoding/decoding images into discrete tokens. @@ -972,6 +867,10 @@ class JanusVQVAE(JanusPreTrainedModel): "JanusVQVAEResnetBlock", "JanusVQVAEVectorQuantizer", ] + _can_record_outputs = { + "hidden_states": JanusVQVAEResnetBlock, + "attentions": JanusVQVAEAttnBlock, + } main_input_name = "pixel_values" def __init__(self, config: JanusVQVAEConfig): @@ -986,11 +885,17 @@ def __init__(self, config: JanusVQVAEConfig): self.gradient_checkpointing = False self.post_init() - def encode(self, pixel_values: torch.LongTensor): + @check_model_inputs + def encode(self, pixel_values: torch.LongTensor, **kwargs: Unpack[TransformersKwargs]) -> JanusVQVAEModelOutput: hidden_states = self.encoder(pixel_values) - hidden_states = self.quant_conv(hidden_states) - quant, emb_loss, indices = self.quantize(hidden_states) - return quant, emb_loss, indices + conv_hidden_states = self.quant_conv(hidden_states) + quantized_last_hidden_state, emb_loss, indices = self.quantize(conv_hidden_states) + return JanusVQVAEModelOutput( + last_hidden_state=hidden_states, + quantized_last_hidden_state=quantized_last_hidden_state, + image_tokens=indices, + embedding_loss=emb_loss, + ) def decode(self, image_tokens: torch.LongTensor) -> torch.FloatTensor: """ @@ -1019,10 +924,10 @@ def forward( **kwargs, ) -> tuple[torch.FloatTensor, torch.FloatTensor]: batch_size = pixel_values.shape[0] - quant, embedding_loss, indices = self.encode(pixel_values) - decoded_pixel_values = self.decode(indices.view(batch_size, -1)) + encode_outputs = self.encode(pixel_values, return_dict=True, **kwargs) + decoded_pixel_values = self.decode(encode_outputs.image_tokens.view(batch_size, -1)) - return JanusVQVAEOutput(decoded_pixel_values, embedding_loss) + return JanusVQVAEOutput(decoded_pixel_values, encode_outputs.embedding_loss) class JanusVQVAEAlignerMLP(nn.Module): @@ -1092,10 +997,15 @@ def get_input_embeddings(self): def set_input_embeddings(self, value): self.language_model.set_input_embeddings(value) - def get_image_features(self, pixel_values): - image_embeds = self.vision_model(pixel_values) - image_embeds = self.aligner(image_embeds.last_hidden_state) - return image_embeds + @can_return_tuple + @auto_docstring + def get_image_features( + self, pixel_values: torch.FloatTensor, **kwargs: Unpack[TransformersKwargs] + ) -> tuple | BaseModelOutputWithPooling: + vision_outputs = self.vision_model(pixel_values, return_dict=True, **kwargs) + vision_outputs.pooler_output = self.aligner(vision_outputs.last_hidden_state) + + return vision_outputs def get_placeholder_mask( self, input_ids: torch.LongTensor, inputs_embeds: torch.FloatTensor, image_features: torch.FloatTensor @@ -1144,7 +1054,7 @@ def forward( inputs_embeds = self.get_input_embeddings()(input_ids) if pixel_values is not None: - image_embeds = self.get_image_features(pixel_values) + image_embeds = self.get_image_features(pixel_values, return_dict=True).pooler_output image_features = image_embeds.reshape(-1, inputs_embeds.shape[-1]) image_features = image_features.to(inputs_embeds.device, inputs_embeds.dtype) image_attention_mask = self.get_placeholder_mask( diff --git a/src/transformers/models/janus/modular_janus.py b/src/transformers/models/janus/modular_janus.py index d8431bb20333..c636b69b47a4 100644 --- a/src/transformers/models/janus/modular_janus.py +++ b/src/transformers/models/janus/modular_janus.py @@ -41,7 +41,7 @@ valid_images, validate_preprocess_arguments, ) -from ...modeling_outputs import ModelOutput +from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling, ModelOutput from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel from ...processing_utils import ImagesKwargs, Unpack from ...utils import ( @@ -552,10 +552,42 @@ def __init__(self, config: JanusVisionConfig): class JanusVisionModel(Blip2VisionModel): + _can_record_outputs = { + "hidden_states": JanusVisionEncoderLayer, + "attentions": JanusVisionAttention, + } + def __init__(self, config: JanusVisionConfig): super().__init__(config) self.encoder = JanusVisionEncoder(config) + def forward( + self, + pixel_values: torch.FloatTensor | None = None, + interpolate_pos_encoding: bool = False, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + if pixel_values is None: + raise ValueError("You have to specify pixel_values") + + hidden_states = self.embeddings(pixel_values, interpolate_pos_encoding=interpolate_pos_encoding) + + encoder_outputs: BaseModelOutput = self.encoder( + inputs_embeds=hidden_states, + **kwargs, + ) + + last_hidden_state = encoder_outputs.last_hidden_state + last_hidden_state = self.post_layernorm(last_hidden_state) + + pooled_output = last_hidden_state[:, 0, :] + pooled_output = self.post_layernorm(pooled_output) + + return BaseModelOutputWithPooling( + last_hidden_state=last_hidden_state, + pooler_output=pooled_output, + ) + class JanusVisionAlignerMLP(nn.Module): def __init__(self, config: JanusVisionConfig): @@ -792,6 +824,10 @@ class JanusVQVAE(ChameleonVQVAE): "JanusVQVAEResnetBlock", "JanusVQVAEVectorQuantizer", ] + _can_record_outputs = { + "hidden_states": JanusVQVAEResnetBlock, + "attentions": JanusVQVAEAttnBlock, + } main_input_name = "pixel_values" def __init__(self, config: JanusVQVAEConfig): @@ -829,10 +865,10 @@ def forward( **kwargs, ) -> tuple[torch.FloatTensor, torch.FloatTensor]: batch_size = pixel_values.shape[0] - quant, embedding_loss, indices = self.encode(pixel_values) - decoded_pixel_values = self.decode(indices.view(batch_size, -1)) + encode_outputs = self.encode(pixel_values, return_dict=True, **kwargs) + decoded_pixel_values = self.decode(encode_outputs.image_tokens.view(batch_size, -1)) - return JanusVQVAEOutput(decoded_pixel_values, embedding_loss) + return JanusVQVAEOutput(decoded_pixel_values, encode_outputs.embedding_loss) class JanusVQVAEAlignerMLP(nn.Module): @@ -902,10 +938,15 @@ def get_input_embeddings(self): def set_input_embeddings(self, value): self.language_model.set_input_embeddings(value) - def get_image_features(self, pixel_values): - image_embeds = self.vision_model(pixel_values) - image_embeds = self.aligner(image_embeds.last_hidden_state) - return image_embeds + @can_return_tuple + @auto_docstring + def get_image_features( + self, pixel_values: torch.FloatTensor, **kwargs: Unpack[TransformersKwargs] + ) -> tuple | BaseModelOutputWithPooling: + vision_outputs = self.vision_model(pixel_values, return_dict=True, **kwargs) + vision_outputs.pooler_output = self.aligner(vision_outputs.last_hidden_state) + + return vision_outputs def get_placeholder_mask( self, input_ids: torch.LongTensor, inputs_embeds: torch.FloatTensor, image_features: torch.FloatTensor @@ -954,7 +995,7 @@ def forward( inputs_embeds = self.get_input_embeddings()(input_ids) if pixel_values is not None: - image_embeds = self.get_image_features(pixel_values) + image_embeds = self.get_image_features(pixel_values, return_dict=True).pooler_output image_features = image_embeds.reshape(-1, inputs_embeds.shape[-1]) image_features = image_features.to(inputs_embeds.device, inputs_embeds.dtype) image_attention_mask = self.get_placeholder_mask( diff --git a/src/transformers/models/kosmos2/modeling_kosmos2.py b/src/transformers/models/kosmos2/modeling_kosmos2.py index 92480cda06f5..bcfff60f0b02 100644 --- a/src/transformers/models/kosmos2/modeling_kosmos2.py +++ b/src/transformers/models/kosmos2/modeling_kosmos2.py @@ -14,6 +14,7 @@ """PyTorch KOSMOS-2 model.""" import math +import warnings from collections.abc import Callable from dataclasses import dataclass from typing import Any @@ -74,6 +75,21 @@ def _make_causal_mask( return mask[None, None, :, :].expand(bsz, 1, tgt_len, tgt_len + past_key_values_length) +@dataclass +@auto_docstring +class BaseModelOutputWithProjectionAttentions(BaseModelOutputWithPooling): + r""" + projection_attentions (`tuple(torch.FloatTensor)`): + Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, + sequence_length)`. + + Attentions weights given by `Kosmos2ImageToTextProjection`, after the attention softmax, used to compute + the weighted average in the self-attention heads. + """ + + projection_attentions: tuple[torch.FloatTensor] | None = None + + @dataclass @auto_docstring( custom_intro=""" @@ -408,7 +424,6 @@ def forward( return outputs -# Copied from transformers.models.altclip.modeling_altclip.AltCLIPEncoder with AltCLIP->Kosmos2Vision class Kosmos2VisionEncoder(nn.Module): """ Transformer encoder consisting of `config.num_hidden_layers` self attention layers. Each layer is a @@ -491,7 +506,7 @@ def forward( if output_hidden_states: encoder_states = encoder_states + (hidden_states,) - return BaseModelOutput( + return BaseModelOutputWithProjectionAttentions( last_hidden_state=hidden_states, hidden_states=encoder_states, attentions=all_attentions ) @@ -1207,7 +1222,7 @@ def forward( interpolate_pos_encoding: bool = False, return_dict: bool | None = None, **kwargs, - ) -> tuple | BaseModelOutputWithPooling: + ) -> tuple | BaseModelOutputWithProjectionAttentions: return self.model( pixel_values=pixel_values, output_attentions=output_attentions, @@ -1492,36 +1507,37 @@ def get_input_embeddings(self) -> nn.Module: def set_input_embeddings(self, value): self.text_model.model.embed_tokens = value + @can_return_tuple + @auto_docstring def get_image_features( self, pixel_values: torch.FloatTensor, - return_attentions: bool | None = False, interpolate_pos_encoding: bool | None = False, - ): - """ - Encodes images into continuous embeddings that can be forwarded to the language model. + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithProjectionAttentions: + if "return_attentions" in kwargs: + warnings.warn( + "`return_attentions` is deprecated and will be removed in a future version. Please use `return_dict`" + " and access `projection_attentions` from the returned `ModelOutput` instead.", + FutureWarning, + ) + kwargs.pop("return_attentions", None) - Args: - pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): - The tensors corresponding to the input images. - return_attentions (`bool`, *optional*, defaults to `False`): - Whether to return `projection_attentions` or not. - interpolate_pos_encoding (`bool`, *optional*, defaults to `False`): - Whether to interpolate positional embeddings or not. - """ - vision_model_output = self.vision_model( + vision_output: BaseModelOutputWithProjectionAttentions = self.vision_model( pixel_values=pixel_values, interpolate_pos_encoding=interpolate_pos_encoding, + return_dict=True, + **kwargs, ) # The whole `last_hidden_state` through `post_layernorm` instead of just `pooled_output`. - image_embeds = self.vision_model.model.post_layernorm(vision_model_output[0]) + image_embeds = self.vision_model.model.post_layernorm(vision_output[0]) # normalized features image_embeds = nn.functional.normalize(image_embeds, dim=-1) image_embeds, projection_attentions = self.image_to_text_projection(image_embeds) + vision_output.pooler_output = image_embeds + vision_output.projection_attentions = projection_attentions - if return_attentions: - return image_embeds, projection_attentions - return image_embeds + return vision_output @can_return_tuple @auto_docstring @@ -1593,9 +1609,11 @@ def forward( if image_embeds is None: if pixel_values is None: raise ValueError("You have to specify either `pixel_values` or `image_embeds`.") - image_embeds, projection_attentions = self.get_image_features( - pixel_values, return_attentions=True, interpolate_pos_encoding=interpolate_pos_encoding + image_features = self.get_image_features( + pixel_values, interpolate_pos_encoding=interpolate_pos_encoding, return_dict=True ) + image_embeds = image_features.pooler_output + projection_attentions = image_features.projection_attentions outputs = self.text_model( input_ids=input_ids, diff --git a/src/transformers/models/lfm2_vl/modeling_lfm2_vl.py b/src/transformers/models/lfm2_vl/modeling_lfm2_vl.py index 8362c9c2f30a..5ff43fadaa10 100755 --- a/src/transformers/models/lfm2_vl/modeling_lfm2_vl.py +++ b/src/transformers/models/lfm2_vl/modeling_lfm2_vl.py @@ -26,7 +26,7 @@ from ...activations import ACT2FN from ...cache_utils import Cache from ...generation import GenerationMixin -from ...modeling_outputs import BaseModelOutputWithPast, ModelOutput +from ...modeling_outputs import BaseModelOutputWithPast, BaseModelOutputWithPooling, ModelOutput from ...modeling_utils import PreTrainedModel from ...processing_utils import Unpack from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, torch_compilable_check @@ -161,37 +161,39 @@ def get_input_embeddings(self): def set_input_embeddings(self, value): self.language_model.set_input_embeddings(value) + @can_return_tuple + @auto_docstring( + custom_intro="Obtains image last hidden states from the vision tower and apply multimodal projection." + ) def get_image_features( self, pixel_values: torch.FloatTensor, spatial_shapes: torch.Tensor, pixel_attention_mask: torch.Tensor, - **kwargs, - ) -> list[torch.Tensor]: - """ - Obtains image last hidden states from the vision tower and apply multimodal projection. - - Args: - pixel_values (`torch.FloatTensor]` of shape `(batch_size, channels, height, width)`): - The tensors corresponding to the input images. - spatial_shapes (`torch.Tensor` of shape `(batch_size, 2)`): - The spatial shapes of the input images. - pixel_attention_mask (`torch.Tensor` of shape `(batch_size, height, width)`): - The pixel attention mask of the input images. - Returns: - image_features (`list[torch.Tensor]`): Image feature tensor of shape `(num_images, image_length, embed_dim)`). + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + pixel_values (`torch.FloatTensor]` of shape `(batch_size, channels, height, width)`): + The tensors corresponding to the input images. + spatial_shapes (`torch.Tensor` of shape `(batch_size, 2)`): + The spatial shapes of the input images. + pixel_attention_mask (`torch.Tensor` of shape `(batch_size, height, width)`): + The pixel attention mask of the input images. """ image_outputs = self.vision_tower( pixel_values=pixel_values, spatial_shapes=spatial_shapes, pixel_attention_mask=pixel_attention_mask, - ).last_hidden_state + return_dict=True, + **kwargs, + ) + last_hidden_state = image_outputs.last_hidden_state img_feature_lengths = pixel_attention_mask.sum(dim=1) image_features = [] - for img_idx in range(image_outputs.size(0)): - feature = image_outputs[img_idx] + for img_idx in range(last_hidden_state.size(0)): + feature = last_hidden_state[img_idx] # unpad the image representation feature = feature[: img_feature_lengths[img_idx], :].unsqueeze(0) @@ -206,7 +208,8 @@ def get_image_features( img_embedding = img_embedding.reshape(-1, img_embedding.size(-1)) image_features.append(img_embedding) - return image_features + image_outputs.pooler_output = image_features + return image_outputs def get_placeholder_mask( self, input_ids: torch.LongTensor, inputs_embeds: torch.FloatTensor, image_features: torch.FloatTensor @@ -266,7 +269,8 @@ def forward( pixel_values=pixel_values, spatial_shapes=spatial_shapes, pixel_attention_mask=pixel_attention_mask, - ) + return_dict=True, + ).pooler_output image_features = torch.cat(image_features, dim=0).to(inputs_embeds.device, inputs_embeds.dtype) special_image_mask = self.get_placeholder_mask( input_ids=input_ids, @@ -318,13 +322,22 @@ def set_input_embeddings(self, value): def get_output_embeddings(self) -> nn.Module: return self.lm_head + @auto_docstring def get_image_features( self, pixel_values: torch.FloatTensor, spatial_shapes: torch.Tensor, pixel_attention_mask: torch.Tensor, - **kwargs, - ): + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + pixel_values (`torch.FloatTensor]` of shape `(batch_size, channels, height, width)`): + The tensors corresponding to the input images. + spatial_shapes (`torch.Tensor` of shape `(batch_size, 2)`): + The spatial shapes of the input images. + pixel_attention_mask (`torch.Tensor` of shape `(batch_size, height, width)`): + The pixel attention mask of the input images. + """ return self.model.get_image_features( pixel_values=pixel_values, spatial_shapes=spatial_shapes, diff --git a/src/transformers/models/lfm2_vl/modular_lfm2_vl.py b/src/transformers/models/lfm2_vl/modular_lfm2_vl.py index d5eed3493fc4..2ff5a055fff9 100644 --- a/src/transformers/models/lfm2_vl/modular_lfm2_vl.py +++ b/src/transformers/models/lfm2_vl/modular_lfm2_vl.py @@ -18,6 +18,7 @@ from ...activations import ACT2FN from ...cache_utils import Cache +from ...modeling_outputs import BaseModelOutputWithPooling from ...processing_utils import Unpack from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, logging, torch_compilable_check from ..llava.modeling_llava import ( @@ -91,37 +92,39 @@ class Lfm2VlModel(LlavaModel): def __init__(self, config: Lfm2VlConfig): super().__init__(config) + @can_return_tuple + @auto_docstring( + custom_intro="Obtains image last hidden states from the vision tower and apply multimodal projection." + ) def get_image_features( self, pixel_values: torch.FloatTensor, spatial_shapes: torch.Tensor, pixel_attention_mask: torch.Tensor, - **kwargs, - ) -> list[torch.Tensor]: - """ - Obtains image last hidden states from the vision tower and apply multimodal projection. - - Args: - pixel_values (`torch.FloatTensor]` of shape `(batch_size, channels, height, width)`): - The tensors corresponding to the input images. - spatial_shapes (`torch.Tensor` of shape `(batch_size, 2)`): - The spatial shapes of the input images. - pixel_attention_mask (`torch.Tensor` of shape `(batch_size, height, width)`): - The pixel attention mask of the input images. - Returns: - image_features (`list[torch.Tensor]`): Image feature tensor of shape `(num_images, image_length, embed_dim)`). + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + pixel_values (`torch.FloatTensor]` of shape `(batch_size, channels, height, width)`): + The tensors corresponding to the input images. + spatial_shapes (`torch.Tensor` of shape `(batch_size, 2)`): + The spatial shapes of the input images. + pixel_attention_mask (`torch.Tensor` of shape `(batch_size, height, width)`): + The pixel attention mask of the input images. """ image_outputs = self.vision_tower( pixel_values=pixel_values, spatial_shapes=spatial_shapes, pixel_attention_mask=pixel_attention_mask, - ).last_hidden_state + return_dict=True, + **kwargs, + ) + last_hidden_state = image_outputs.last_hidden_state img_feature_lengths = pixel_attention_mask.sum(dim=1) image_features = [] - for img_idx in range(image_outputs.size(0)): - feature = image_outputs[img_idx] + for img_idx in range(last_hidden_state.size(0)): + feature = last_hidden_state[img_idx] # unpad the image representation feature = feature[: img_feature_lengths[img_idx], :].unsqueeze(0) @@ -136,7 +139,8 @@ def get_image_features( img_embedding = img_embedding.reshape(-1, img_embedding.size(-1)) image_features.append(img_embedding) - return image_features + image_outputs.pooler_output = image_features + return image_outputs def get_placeholder_mask( self, input_ids: torch.LongTensor, inputs_embeds: torch.FloatTensor, image_features: torch.FloatTensor @@ -196,7 +200,8 @@ def forward( pixel_values=pixel_values, spatial_shapes=spatial_shapes, pixel_attention_mask=pixel_attention_mask, - ) + return_dict=True, + ).pooler_output image_features = torch.cat(image_features, dim=0).to(inputs_embeds.device, inputs_embeds.dtype) special_image_mask = self.get_placeholder_mask( input_ids=input_ids, @@ -227,13 +232,22 @@ def forward( class Lfm2VlForConditionalGeneration(LlavaForConditionalGeneration): _checkpoint_conversion_mapping = {} + @auto_docstring def get_image_features( self, pixel_values: torch.FloatTensor, spatial_shapes: torch.Tensor, pixel_attention_mask: torch.Tensor, - **kwargs, - ): + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + pixel_values (`torch.FloatTensor]` of shape `(batch_size, channels, height, width)`): + The tensors corresponding to the input images. + spatial_shapes (`torch.Tensor` of shape `(batch_size, 2)`): + The spatial shapes of the input images. + pixel_attention_mask (`torch.Tensor` of shape `(batch_size, height, width)`): + The pixel attention mask of the input images. + """ return self.model.get_image_features( pixel_values=pixel_values, spatial_shapes=spatial_shapes, diff --git a/src/transformers/models/lighton_ocr/modeling_lighton_ocr.py b/src/transformers/models/lighton_ocr/modeling_lighton_ocr.py index 12ae454a9571..c95633254ae9 100644 --- a/src/transformers/models/lighton_ocr/modeling_lighton_ocr.py +++ b/src/transformers/models/lighton_ocr/modeling_lighton_ocr.py @@ -25,10 +25,11 @@ from ...cache_utils import Cache from ...generation import GenerationMixin from ...integrations import use_kernel_forward_from_hub -from ...modeling_outputs import BaseModelOutputWithPast, ModelOutput +from ...modeling_outputs import BaseModelOutputWithPast, BaseModelOutputWithPooling, ModelOutput from ...modeling_utils import PreTrainedModel from ...processing_utils import Unpack from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, torch_compilable_check +from ...utils.generic import check_model_inputs from ..auto import AutoModel from .configuration_lighton_ocr import LightOnOcrConfig @@ -170,27 +171,22 @@ def get_input_embeddings(self): def set_input_embeddings(self, value): self.language_model.set_input_embeddings(value) - def get_image_features(self, pixel_values: torch.Tensor, image_sizes: torch.Tensor | list): - """ - Obtains image features from the vision encoder and projection. - - Args: - pixel_values: Image tensors - image_sizes: Tensor or list of (height, width) pairs for each image - - Returns: - List of image feature tensors, one per image - """ - visual_features = self.vision_encoder(pixel_values, image_sizes=image_sizes).last_hidden_state - - image_features = self.vision_projection(visual_features.squeeze(0), image_sizes) + @can_return_tuple + @auto_docstring + def get_image_features( + self, pixel_values: torch.Tensor, image_sizes: torch.Tensor | list, **kwargs: Unpack[TransformersKwargs] + ) -> tuple | BaseModelOutputWithPooling: + image_outputs = self.vision_encoder(pixel_values, image_sizes=image_sizes, return_dict=True, **kwargs) + image_features = image_outputs.last_hidden_state + image_features = self.vision_projection(image_features.squeeze(0), image_sizes) # Split features per image based on the effective patch size downsample_ratio = self.config.vision_config.patch_size * self.config.spatial_merge_size split_sizes = [(height // downsample_ratio) * (width // downsample_ratio) for height, width in image_sizes] image_features = torch.split(image_features, split_sizes) + image_outputs.pooler_output = image_features - return image_features + return image_outputs def get_placeholder_mask( self, input_ids: torch.LongTensor, inputs_embeds: torch.FloatTensor, image_features: torch.FloatTensor @@ -247,7 +243,9 @@ def forward( inputs_embeds = self.get_input_embeddings()(input_ids) if pixel_values is not None: - image_features = self.get_image_features(pixel_values=pixel_values, image_sizes=image_sizes) + image_features = self.get_image_features( + pixel_values=pixel_values, image_sizes=image_sizes, return_dict=True + ).pooler_output image_features = torch.cat(image_features, dim=0).to(inputs_embeds.device, inputs_embeds.dtype) special_image_mask = self.get_placeholder_mask( input_ids, inputs_embeds=inputs_embeds, image_features=image_features @@ -330,10 +328,13 @@ def set_input_embeddings(self, value): def get_output_embeddings(self) -> nn.Module: return self.lm_head - def get_image_features(self, pixel_values: torch.FloatTensor, image_sizes: torch.Tensor, **kwargs): - return self.model.get_image_features(pixel_values=pixel_values, image_sizes=image_sizes) + @auto_docstring + def get_image_features( + self, pixel_values: torch.FloatTensor, image_sizes: torch.Tensor, **kwargs: Unpack[TransformersKwargs] + ) -> tuple | BaseModelOutputWithPooling: + return self.model.get_image_features(pixel_values=pixel_values, image_sizes=image_sizes, **kwargs) - @can_return_tuple + @check_model_inputs(tie_last_hidden_states=False) @auto_docstring def forward( self, diff --git a/src/transformers/models/lighton_ocr/modular_lighton_ocr.py b/src/transformers/models/lighton_ocr/modular_lighton_ocr.py index d7527f29d752..310118ac1a4b 100644 --- a/src/transformers/models/lighton_ocr/modular_lighton_ocr.py +++ b/src/transformers/models/lighton_ocr/modular_lighton_ocr.py @@ -21,6 +21,7 @@ from ...configuration_utils import PretrainedConfig from ...feature_extraction_utils import BatchFeature from ...image_utils import ImageInput +from ...modeling_outputs import BaseModelOutputWithPooling from ...modeling_utils import PreTrainedModel from ...processing_utils import ( MultiModalData, @@ -305,27 +306,22 @@ def __init__(self, config: LightOnOcrConfig): self.language_model = AutoModel.from_config(config.text_config) self.post_init() - def get_image_features(self, pixel_values: torch.Tensor, image_sizes: torch.Tensor | list): - """ - Obtains image features from the vision encoder and projection. - - Args: - pixel_values: Image tensors - image_sizes: Tensor or list of (height, width) pairs for each image - - Returns: - List of image feature tensors, one per image - """ - visual_features = self.vision_encoder(pixel_values, image_sizes=image_sizes).last_hidden_state - - image_features = self.vision_projection(visual_features.squeeze(0), image_sizes) + @can_return_tuple + @auto_docstring + def get_image_features( + self, pixel_values: torch.Tensor, image_sizes: torch.Tensor | list, **kwargs: Unpack[TransformersKwargs] + ) -> tuple | BaseModelOutputWithPooling: + image_outputs = self.vision_encoder(pixel_values, image_sizes=image_sizes, return_dict=True, **kwargs) + image_features = image_outputs.last_hidden_state + image_features = self.vision_projection(image_features.squeeze(0), image_sizes) # Split features per image based on the effective patch size downsample_ratio = self.config.vision_config.patch_size * self.config.spatial_merge_size split_sizes = [(height // downsample_ratio) * (width // downsample_ratio) for height, width in image_sizes] image_features = torch.split(image_features, split_sizes) + image_outputs.pooler_output = image_features - return image_features + return image_outputs @can_return_tuple @auto_docstring @@ -358,7 +354,9 @@ def forward( inputs_embeds = self.get_input_embeddings()(input_ids) if pixel_values is not None: - image_features = self.get_image_features(pixel_values=pixel_values, image_sizes=image_sizes) + image_features = self.get_image_features( + pixel_values=pixel_values, image_sizes=image_sizes, return_dict=True + ).pooler_output image_features = torch.cat(image_features, dim=0).to(inputs_embeds.device, inputs_embeds.dtype) special_image_mask = self.get_placeholder_mask( input_ids, inputs_embeds=inputs_embeds, image_features=image_features @@ -390,8 +388,11 @@ def forward( class LightOnOcrForConditionalGeneration(Mistral3ForConditionalGeneration): _checkpoint_conversion_mapping = {} - def get_image_features(self, pixel_values: torch.FloatTensor, image_sizes: torch.Tensor, **kwargs): - return self.model.get_image_features(pixel_values=pixel_values, image_sizes=image_sizes) + @auto_docstring + def get_image_features( + self, pixel_values: torch.FloatTensor, image_sizes: torch.Tensor, **kwargs: Unpack[TransformersKwargs] + ) -> tuple | BaseModelOutputWithPooling: + return self.model.get_image_features(pixel_values=pixel_values, image_sizes=image_sizes, **kwargs) __all__ = [ diff --git a/src/transformers/models/llama4/modeling_llama4.py b/src/transformers/models/llama4/modeling_llama4.py index 801fafdcdcd5..cbc8b42e9114 100644 --- a/src/transformers/models/llama4/modeling_llama4.py +++ b/src/transformers/models/llama4/modeling_llama4.py @@ -31,7 +31,13 @@ from ...masking_utils import create_causal_mask, create_chunked_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer -from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPast, CausalLMOutputWithPast, ModelOutput +from ...modeling_outputs import ( + BaseModelOutput, + BaseModelOutputWithPast, + BaseModelOutputWithPooling, + CausalLMOutputWithPast, + ModelOutput, +) from ...modeling_rope_utils import ( ROPE_INIT_FUNCTIONS, dynamic_rope_update, @@ -1201,30 +1207,23 @@ def set_decoder(self, decoder): def get_decoder(self): return self.language_model.get_decoder() + @check_model_inputs(tie_last_hidden_states=False) + @auto_docstring(custom_intro="Obtains image last hidden states from the vision tower and apply al projection.") def get_image_features( self, pixel_values: torch.FloatTensor, vision_feature_select_strategy: str, - **kwargs, - ): - """ - Obtains image last hidden states from the vision tower and apply al projection. - - Args: - pixel_values (`torch.FloatTensor]` of shape `(batch_size, channels, height, width)`) - The tensors corresponding to the input images. - vision_feature_select_strategy (`str`): - The feature selection strategy used to select the vision feature from the vision backbone. - Can be one of `"default"` or `"full"` - Returns: - image_features (`torch.Tensor`): Image feature tensor of shape `(num_images, image_length, embed_dim)`). + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + pixel_values (`torch.FloatTensor]` of shape `(batch_size, channels, height, width)`) + The tensors corresponding to the input images. + vision_feature_select_strategy (`str`): + The feature selection strategy used to select the vision feature from the vision backbone. + Can be one of `"default"` or `"full"` """ - if vision_feature_select_strategy not in ["default", "full"]: - raise ValueError(f"Unexpected select feature strategy: {self.vision_feature_select_strategy}") kwargs = {k: v for k, v in kwargs.items() if v is not None} - image_outputs = self.vision_model(pixel_values, output_hidden_states=False, **kwargs) - hidden_state = image_outputs.last_hidden_state - return hidden_state + return self.vision_model(pixel_values, **kwargs) def get_placeholder_mask( self, input_ids: torch.LongTensor, inputs_embeds: torch.FloatTensor, image_features: torch.FloatTensor @@ -1249,6 +1248,7 @@ def get_placeholder_mask( ) return special_image_mask + @check_model_inputs(tie_last_hidden_states=False) @auto_docstring def forward( self, @@ -1301,11 +1301,6 @@ def forward( output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states ) return_dict = return_dict if return_dict is not None else self.config.use_return_dict - vision_feature_select_strategy = ( - vision_feature_select_strategy - if vision_feature_select_strategy is not None - else self.config.vision_config.vision_feature_select_strategy - ) if (input_ids is None) ^ (inputs_embeds is not None): raise ValueError("You must specify exactly one of input_ids or inputs_embeds") @@ -1322,7 +1317,8 @@ def forward( image_features = self.get_image_features( pixel_values=pixel_values, vision_feature_select_strategy=vision_feature_select_strategy, - ) + return_dict=True, + ).pooler_output vision_flat = image_features.view(-1, image_features.size(-1)) projected_vision_flat = self.multi_modal_projector(vision_flat).to( diff --git a/src/transformers/models/llava/modeling_llava.py b/src/transformers/models/llava/modeling_llava.py index 004051e3a5d3..9fbd58c786df 100644 --- a/src/transformers/models/llava/modeling_llava.py +++ b/src/transformers/models/llava/modeling_llava.py @@ -21,10 +21,11 @@ from ...activations import ACT2FN from ...cache_utils import Cache from ...generation import GenerationMixin -from ...modeling_outputs import BaseModelOutputWithPast, ModelOutput +from ...modeling_outputs import BaseModelOutputWithPast, BaseModelOutputWithPooling, ModelOutput from ...modeling_utils import PreTrainedModel from ...processing_utils import Unpack -from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, logging, torch_compilable_check +from ...utils import TransformersKwargs, auto_docstring, logging, torch_compilable_check +from ...utils.generic import check_model_inputs from ..auto import AutoModel from .configuration_llava import LlavaConfig @@ -145,44 +146,26 @@ def get_input_embeddings(self): def set_input_embeddings(self, value): self.language_model.set_input_embeddings(value) + @check_model_inputs(tie_last_hidden_states=False) + @auto_docstring( + custom_intro="Obtains image last hidden states from the vision tower and apply multimodal projection." + ) def get_image_features( self, pixel_values: torch.FloatTensor, vision_feature_layer: int | list[int] | None = None, vision_feature_select_strategy: str | None = None, - **kwargs, - ): - """ - Obtains image last hidden states from the vision tower and apply multimodal projection. - - Args: - pixel_values (`torch.FloatTensor]` of shape `(batch_size, channels, height, width)`): - The tensors corresponding to the input images. - vision_feature_layer (`Union[int, list[int]]`, *optional*): - The index of the layer to select the vision feature. If multiple indices are provided, - the vision feature of the corresponding indices will be concatenated to form the - vision features. - vision_feature_select_strategy (`str`, *optional*): - The feature selection strategy used to select the vision feature from the vision backbone. - Can be one of `"default"` or `"full"` - Returns: - image_features (`torch.Tensor`): Image feature tensor of shape `(num_images, image_length, embed_dim)`). - """ - vision_feature_layer = ( - vision_feature_layer if vision_feature_layer is not None else self.config.vision_feature_layer - ) - vision_feature_select_strategy = ( - vision_feature_select_strategy - if vision_feature_select_strategy is not None - else self.config.vision_feature_select_strategy - ) - - if vision_feature_select_strategy not in ["default", "full"]: - raise ValueError(f"Unexpected select feature strategy: {self.config.vision_feature_select_strategy}") - + output_hidden_states: bool | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: kwargs = {k: v for k, v in kwargs.items() if v is not None} # this is not memory efficient at all (output_hidden_states=True) will save all the hidden states. - image_outputs = self.vision_tower(pixel_values, output_hidden_states=True, **kwargs) + image_outputs = self.vision_tower( + pixel_values, + output_hidden_states=True, # Ignore arg on purpose + return_dict=True, + **kwargs, + ) # If we have one vision feature layer, return the corresponding hidden states, # otherwise, select the hidden states of each feature layer and concatenate them @@ -199,7 +182,9 @@ def get_image_features( image_features = self.multi_modal_projector(selected_image_feature) - if "image_sizes" in kwargs: + # If image_sizes is provided, we need to split the image features accordingly, + # but only if the image_sizes is not None (the default in this and related architectures) + if kwargs.get("image_sizes") is not None: split_sizes = ( (torch.as_tensor(kwargs["image_sizes"], device=image_features.device) // self.vision_tower.patch_size) .prod(dim=-1) @@ -208,7 +193,9 @@ def get_image_features( image_features = torch.split(image_features.squeeze(0), split_sizes) else: image_features = list(image_features) - return image_features + image_outputs.pooler_output = image_features + + return image_outputs def get_placeholder_mask( self, input_ids: torch.LongTensor, inputs_embeds: torch.FloatTensor, image_features: torch.FloatTensor @@ -234,7 +221,7 @@ def get_placeholder_mask( ) return special_image_mask - @can_return_tuple + @check_model_inputs(tie_last_hidden_states=False) @auto_docstring def forward( self, @@ -250,15 +237,6 @@ def forward( image_sizes: torch.Tensor | None = None, **kwargs: Unpack[TransformersKwargs], ) -> tuple | LlavaModelOutputWithPast: - vision_feature_layer = ( - vision_feature_layer if vision_feature_layer is not None else self.config.vision_feature_layer - ) - vision_feature_select_strategy = ( - vision_feature_select_strategy - if vision_feature_select_strategy is not None - else self.config.vision_feature_select_strategy - ) - if (input_ids is None) ^ (inputs_embeds is not None): raise ValueError("You must specify exactly one of input_ids or inputs_embeds") @@ -271,7 +249,8 @@ def forward( vision_feature_layer=vision_feature_layer, vision_feature_select_strategy=vision_feature_select_strategy, image_sizes=image_sizes, - ) + return_dict=True, + ).pooler_output image_features = torch.cat(image_features, dim=0).to(inputs_embeds.device, inputs_embeds.dtype) special_image_mask = self.get_placeholder_mask( input_ids, inputs_embeds=inputs_embeds, image_features=image_features @@ -325,13 +304,14 @@ def set_input_embeddings(self, value): def get_output_embeddings(self) -> nn.Module: return self.lm_head + @auto_docstring def get_image_features( self, pixel_values: torch.FloatTensor, vision_feature_layer: int | list[int] | None = None, vision_feature_select_strategy: str | None = None, - **kwargs, - ): + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: return self.model.get_image_features( pixel_values=pixel_values, vision_feature_layer=vision_feature_layer, @@ -339,7 +319,7 @@ def get_image_features( **kwargs, ) - @can_return_tuple + @check_model_inputs(tie_last_hidden_states=False) @auto_docstring def forward( self, @@ -384,15 +364,6 @@ def forward( >>> processor.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] "USER: \nWhat's the content of the image? ASSISTANT: The image features a busy city street with a stop sign prominently displayed" ```""" - vision_feature_layer = ( - vision_feature_layer if vision_feature_layer is not None else self.config.vision_feature_layer - ) - vision_feature_select_strategy = ( - vision_feature_select_strategy - if vision_feature_select_strategy is not None - else self.config.vision_feature_select_strategy - ) - outputs = self.model( input_ids=input_ids, pixel_values=pixel_values, diff --git a/src/transformers/models/llava_next/modeling_llava_next.py b/src/transformers/models/llava_next/modeling_llava_next.py index 4ad555ce9f2f..311f958005b5 100644 --- a/src/transformers/models/llava_next/modeling_llava_next.py +++ b/src/transformers/models/llava_next/modeling_llava_next.py @@ -26,10 +26,11 @@ from ...generation import GenerationMixin from ...image_processing_utils import select_best_resolution from ...modeling_flash_attention_utils import FlashAttentionKwargs -from ...modeling_outputs import BaseModelOutputWithPast, ModelOutput +from ...modeling_outputs import BaseModelOutputWithPast, BaseModelOutputWithPooling, ModelOutput from ...modeling_utils import PreTrainedModel from ...processing_utils import Unpack from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, logging, torch_compilable_check +from ...utils.generic import check_model_inputs from ..auto import AutoModel from .configuration_llava_next import LlavaNextConfig @@ -343,41 +344,32 @@ def pack_image_features(self, image_features, image_sizes, vision_feature_select feature_lens = torch.tensor(feature_lens, dtype=torch.long, device=image_features[0].device) return new_image_features, feature_lens + @check_model_inputs(tie_last_hidden_states=False) + @auto_docstring( + custom_intro="Obtains image last hidden states from the vision tower and apply multimodal projection." + ) def get_image_features( self, pixel_values: torch.FloatTensor, image_sizes: torch.Tensor, vision_feature_layer: int | list[int] | None = None, vision_feature_select_strategy: str | None = None, - ): - """ - Obtains image last hidden states from the vision tower and apply multimodal projection. - - Args: - pixel_values (`torch.FloatTensor]` of shape `(batch_size, num_patches, channels, height, width)`) - The tensors corresponding to the input images. - image_sizes (`torch.Tensor` of shape `(num_images, 2)`) - Actual image size of each images (H, W). - vision_feature_layer (`Union[int, list[int]]`, *optional*): - The index of the layer to select the vision feature. If multiple indices are provided, - the vision feature of the corresponding indices will be concatenated to form the - vision features. - vision_feature_select_strategy (`str`, *optional*): - The feature selection strategy used to select the vision feature from the vision backbone. - Can be one of `"default"` or `"full"` - Returns: - image_features (list[`torch.Tensor`]): List of image feature tensor, each contains all the visual feature of all patches - and are of shape `(num_patches, image_length, embed_dim)`). + output_hidden_states: bool | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + pixel_values (`torch.FloatTensor]` of shape `(batch_size, num_patches, channels, height, width)`) + The tensors corresponding to the input images. + image_sizes (`torch.Tensor` of shape `(num_images, 2)`) + Actual image size of each images (H, W). + vision_feature_layer (`Union[int, list[int]]`, *optional*): + The index of the layer to select the vision feature. If multiple indices are provided, + the vision feature of the corresponding indices will be concatenated to form the + vision features. + vision_feature_select_strategy (`str`, *optional*): + The feature selection strategy used to select the vision feature from the vision backbone. + Can be one of `"default"` or `"full"` """ - vision_feature_layer = ( - vision_feature_layer if vision_feature_layer is not None else self.config.vision_feature_layer - ) - vision_feature_select_strategy = ( - vision_feature_select_strategy - if vision_feature_select_strategy is not None - else self.config.vision_feature_select_strategy - ) - # ! infer image_num_patches from image_sizes image_num_patches = [ image_size_to_num_patches( @@ -395,13 +387,18 @@ def get_image_features( # otherwise has to be stacked from list of (num_patches, num_channels, height, width) raise ValueError(f"pixel_values of shape {pixel_values.shape}, expect to be of 4 or 5 dimensions") - image_features = self.vision_tower(pixel_values, output_hidden_states=True) + image_outputs = self.vision_tower( + pixel_values, + output_hidden_states=True, # Ignore arg on purpose + return_dict=True, + **kwargs, + ) # If we have one vision feature layer, return the corresponding hidden states, # otherwise, select the hidden states of each feature layer and concatenate them if isinstance(vision_feature_layer, int): - selected_image_feature = image_features.hidden_states[vision_feature_layer] + selected_image_feature = image_outputs.hidden_states[vision_feature_layer] else: - hs_pool = [image_features.hidden_states[layer_idx] for layer_idx in vision_feature_layer] + hs_pool = [image_outputs.hidden_states[layer_idx] for layer_idx in vision_feature_layer] selected_image_feature = torch.cat(hs_pool, dim=-1) if vision_feature_select_strategy == "default": @@ -417,7 +414,9 @@ def get_image_features( vision_feature_select_strategy=vision_feature_select_strategy, image_newline=self.image_newline, ) - return image_features + image_outputs.pooler_output = image_features + + return image_outputs def get_placeholder_mask( self, input_ids: torch.LongTensor, inputs_embeds: torch.FloatTensor, image_features: torch.FloatTensor @@ -473,14 +472,6 @@ def forward( output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states ) return_dict = return_dict if return_dict is not None else self.config.use_return_dict - vision_feature_layer = ( - vision_feature_layer if vision_feature_layer is not None else self.config.vision_feature_layer - ) - vision_feature_select_strategy = ( - vision_feature_select_strategy - if vision_feature_select_strategy is not None - else self.config.vision_feature_select_strategy - ) if (input_ids is None) ^ (inputs_embeds is not None): raise ValueError("You must specify exactly one of input_ids or inputs_embeds") @@ -494,7 +485,8 @@ def forward( image_sizes, vision_feature_layer=vision_feature_layer, vision_feature_select_strategy=vision_feature_select_strategy, - ) + return_dict=True, + ).pooler_output image_features = torch.cat(image_features, dim=0).to(inputs_embeds.device, inputs_embeds.dtype) special_image_mask = self.get_placeholder_mask( input_ids, inputs_embeds=inputs_embeds, image_features=image_features @@ -561,21 +553,37 @@ def pack_image_features(self, image_features, image_sizes, vision_feature_select image_newline=image_newline, ) + @auto_docstring def get_image_features( self, pixel_values: torch.FloatTensor, image_sizes: torch.Tensor, vision_feature_layer: int | list[int] | None = None, vision_feature_select_strategy: str | None = None, - ): + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + pixel_values (`torch.FloatTensor]` of shape `(batch_size, num_patches, channels, height, width)`) + The tensors corresponding to the input images. + image_sizes (`torch.Tensor` of shape `(num_images, 2)`) + Actual image size of each images (H, W). + vision_feature_layer (`Union[int, list[int]]`, *optional*): + The index of the layer to select the vision feature. If multiple indices are provided, + the vision feature of the corresponding indices will be concatenated to form the + vision features. + vision_feature_select_strategy (`str`, *optional*): + The feature selection strategy used to select the vision feature from the vision backbone. + Can be one of `"default"` or `"full"` + """ return self.model.get_image_features( pixel_values=pixel_values, image_sizes=image_sizes, vision_feature_layer=vision_feature_layer, vision_feature_select_strategy=vision_feature_select_strategy, + **kwargs, ) - @can_return_tuple + @check_model_inputs(tie_last_hidden_states=False) @auto_docstring def forward( self, @@ -632,14 +640,6 @@ def forward( output_hidden_states = ( output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states ) - vision_feature_layer = ( - vision_feature_layer if vision_feature_layer is not None else self.config.vision_feature_layer - ) - vision_feature_select_strategy = ( - vision_feature_select_strategy - if vision_feature_select_strategy is not None - else self.config.vision_feature_select_strategy - ) outputs = self.model( input_ids, diff --git a/src/transformers/models/llava_next_video/modeling_llava_next_video.py b/src/transformers/models/llava_next_video/modeling_llava_next_video.py index 6ff06be1af70..687899c8f5cd 100644 --- a/src/transformers/models/llava_next_video/modeling_llava_next_video.py +++ b/src/transformers/models/llava_next_video/modeling_llava_next_video.py @@ -31,10 +31,11 @@ from ...generation import GenerationMixin from ...image_processing_utils import select_best_resolution from ...modeling_flash_attention_utils import FlashAttentionKwargs -from ...modeling_outputs import BaseModelOutputWithPast, ModelOutput +from ...modeling_outputs import BaseModelOutputWithPast, BaseModelOutputWithPooling, ModelOutput from ...modeling_utils import PreTrainedModel from ...processing_utils import Unpack -from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, logging, torch_compilable_check +from ...utils import TransformersKwargs, auto_docstring, logging, torch_compilable_check +from ...utils.generic import check_model_inputs from ..auto import AutoModel from .configuration_llava_next_video import LlavaNextVideoConfig @@ -395,41 +396,32 @@ def pack_image_features(self, image_features, image_sizes, vision_feature_select feature_lens = torch.tensor(feature_lens, dtype=torch.long, device=image_features[0].device) return new_image_features, feature_lens + @check_model_inputs(tie_last_hidden_states=False) + @auto_docstring( + custom_intro="Obtains image last hidden states from the vision tower and apply multimodal projection." + ) def get_image_features( self, pixel_values: torch.FloatTensor, image_sizes: torch.Tensor, vision_feature_layer: int | list[int] | None = None, vision_feature_select_strategy: str | None = None, - ): - """ - Obtains image last hidden states from the vision tower and apply multimodal projection. - - Args: - pixel_values (`torch.FloatTensor]` of shape `(batch_size, num_patches, channels, height, width)`) - The tensors corresponding to the input images. - image_sizes (`torch.Tensor` of shape `(num_images, 2)`) - Actual image size of each images (H, W). - vision_feature_layer (`Union[int, list[int]]`, *optional*): - The index of the layer to select the vision feature. If multiple indices are provided, - the vision feature of the corresponding indices will be concatenated to form the - vision features. - vision_feature_select_strategy (`str`, *optional*): - The feature selection strategy used to select the vision feature from the vision backbone. - Can be one of `"default"` or `"full"` - Returns: - image_features (list[`torch.Tensor`]): List of image feature tensor, each contains all the visual feature of all patches - and are of shape `(num_patches, image_length, embed_dim)`). + output_hidden_states: bool | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + pixel_values (`torch.FloatTensor]` of shape `(batch_size, num_patches, channels, height, width)`) + The tensors corresponding to the input images. + image_sizes (`torch.Tensor` of shape `(num_images, 2)`) + Actual image size of each images (H, W). + vision_feature_layer (`Union[int, list[int]]`, *optional*): + The index of the layer to select the vision feature. If multiple indices are provided, + the vision feature of the corresponding indices will be concatenated to form the + vision features. + vision_feature_select_strategy (`str`, *optional*): + The feature selection strategy used to select the vision feature from the vision backbone. + Can be one of `"default"` or `"full"` """ - vision_feature_layer = ( - vision_feature_layer if vision_feature_layer is not None else self.config.vision_feature_layer - ) - vision_feature_select_strategy = ( - vision_feature_select_strategy - if vision_feature_select_strategy is not None - else self.config.vision_feature_select_strategy - ) - # ! infer image_num_patches from image_sizes image_num_patches = [ image_size_to_num_patches( @@ -447,13 +439,18 @@ def get_image_features( # otherwise has to be stacked from list of (num_patches, num_channels, height, width) raise ValueError(f"pixel_values of shape {pixel_values.shape}, expect to be of 4 or 5 dimensions") - image_features = self.vision_tower(pixel_values, output_hidden_states=True) + image_outputs = self.vision_tower( + pixel_values, + output_hidden_states=True, # Ignore arg on purpose + return_dict=True, + **kwargs, + ) # If we have one vision feature layer, return the corresponding hidden states, # otherwise, select the hidden states of each feature layer and concatenate them if isinstance(vision_feature_layer, int): - selected_image_feature = image_features.hidden_states[vision_feature_layer] + selected_image_feature = image_outputs.hidden_states[vision_feature_layer] else: - hs_pool = [image_features.hidden_states[layer_idx] for layer_idx in vision_feature_layer] + hs_pool = [image_outputs.hidden_states[layer_idx] for layer_idx in vision_feature_layer] selected_image_feature = torch.cat(hs_pool, dim=-1) if vision_feature_select_strategy == "default": @@ -467,7 +464,9 @@ def get_image_features( vision_feature_select_strategy, image_newline=self.image_newline, ) - return image_features + image_outputs.pooler_output = image_features + + return image_outputs def get_placeholder_mask( self, @@ -510,7 +509,7 @@ def get_placeholder_mask( ) return special_image_mask, special_video_mask - @can_return_tuple + @check_model_inputs(tie_last_hidden_states=False) @auto_docstring def forward( self, @@ -542,14 +541,6 @@ def forward( output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states ) return_dict = return_dict if return_dict is not None else self.config.use_return_dict - vision_feature_layer = ( - vision_feature_layer if vision_feature_layer is not None else self.config.vision_feature_layer - ) - vision_feature_select_strategy = ( - vision_feature_select_strategy - if vision_feature_select_strategy is not None - else self.config.vision_feature_select_strategy - ) if (input_ids is None) ^ (inputs_embeds is not None): raise ValueError("You must specify exactly one of input_ids or inputs_embeds") @@ -563,7 +554,8 @@ def forward( image_sizes, vision_feature_layer=vision_feature_layer, vision_feature_select_strategy=vision_feature_select_strategy, - ) + return_dict=True, + ).pooler_output image_features = torch.cat(image_features, dim=0).to(inputs_embeds.device, inputs_embeds.dtype) special_image_mask, _ = self.get_placeholder_mask( input_ids, inputs_embeds=inputs_embeds, image_features=image_features @@ -575,7 +567,8 @@ def forward( pixel_values_videos, vision_feature_layer=vision_feature_layer, vision_feature_select_strategy=vision_feature_select_strategy, - ) + return_dict=True, + ).pooler_output video_features = [feature.flatten(0, 1) for feature in video_features] video_feature_lens = [feature.size(0) for feature in video_features] video_features = torch.cat(video_features, dim=0) @@ -609,48 +602,44 @@ def forward( video_hidden_states=video_features if pixel_values_videos is not None else None, ) + @check_model_inputs(tie_last_hidden_states=False) + @auto_docstring( + custom_intro="Obtains video last hidden states from the vision tower and apply multimodal projection." + ) def get_video_features( self, pixel_values: torch.FloatTensor, vision_feature_layer: int | list[int] | None = None, vision_feature_select_strategy: str | None = None, - ): - """ - Obtains video last hidden states from the vision tower and apply multimodal projection. - - Args: - pixel_values (`torch.FloatTensor]` of shape `(batch_size, num_frames, channels, height, width)`) - The tensors corresponding to the input video. - vision_feature_layer (`Union[int, list[int]]`, *optional;*): - The index of the layer to select the vision feature. If multiple indices are provided, - the vision feature of the corresponding indices will be concatenated to form the - vision features. - vision_feature_select_strategy (`str`, *optional*): - The feature selection strategy used to select the vision feature from the vision backbone. - Can be one of `"default"` or `"full"` - Returns: - video_features (list[`torch.Tensor`]): List of video feature tensor, each contains all the visual feature of all patches - and are of shape `(num_videos, video_length, embed_dim)`). + output_hidden_states: bool | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + pixel_values (`torch.FloatTensor]` of shape `(batch_size, num_frames, channels, height, width)`) + The tensors corresponding to the input video. + vision_feature_layer (`Union[int, list[int]]`, *optional;*): + The index of the layer to select the vision feature. If multiple indices are provided, + the vision feature of the corresponding indices will be concatenated to form the + vision features. + vision_feature_select_strategy (`str`, *optional*): + The feature selection strategy used to select the vision feature from the vision backbone. + Can be one of `"default"` or `"full"` """ - vision_feature_layer = ( - vision_feature_layer if vision_feature_layer is not None else self.config.vision_feature_layer - ) - vision_feature_select_strategy = ( - vision_feature_select_strategy - if vision_feature_select_strategy is not None - else self.config.vision_feature_select_strategy - ) - batch_size, frames, channels, height, width = pixel_values.shape pixel_values = pixel_values.reshape(batch_size * frames, channels, height, width) - video_features = self.vision_tower(pixel_values, output_hidden_states=True) + video_outputs = self.vision_tower( + pixel_values, + output_hidden_states=True, # Ignore arg on purpose + return_dict=True, + **kwargs, + ) # If we have one vision feature layer, return the corresponding hidden states, # otherwise, select the hidden states of each feature layer and concatenate them if isinstance(vision_feature_layer, int): - selected_video_features = video_features.hidden_states[vision_feature_layer] + selected_video_features = video_outputs.hidden_states[vision_feature_layer] else: - hs_pool = [video_features.hidden_states[layer_idx] for layer_idx in vision_feature_layer] + hs_pool = [video_outputs.hidden_states[layer_idx] for layer_idx in vision_feature_layer] selected_video_features = torch.cat(hs_pool, dim=-1) if vision_feature_select_strategy == "default": @@ -660,7 +649,9 @@ def get_video_features( video_features = self.vision_resampler(selected_video_features) video_features = self.multi_modal_projector(video_features) video_features = torch.split(video_features, frames, dim=0) - return video_features + video_outputs.pooler_output = video_features + + return video_outputs @auto_docstring( @@ -701,21 +692,37 @@ def pack_image_features(self, image_features, image_sizes, vision_feature_select image_newline=image_newline, ) + @auto_docstring def get_image_features( self, pixel_values: torch.FloatTensor, image_sizes: torch.Tensor, vision_feature_layer: int | list[int] | None = None, vision_feature_select_strategy: str | None = None, - ): + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + pixel_values (`torch.FloatTensor]` of shape `(batch_size, num_patches, channels, height, width)`) + The tensors corresponding to the input images. + image_sizes (`torch.Tensor` of shape `(num_images, 2)`) + Actual image size of each images (H, W). + vision_feature_layer (`Union[int, list[int]]`, *optional*): + The index of the layer to select the vision feature. If multiple indices are provided, + the vision feature of the corresponding indices will be concatenated to form the + vision features. + vision_feature_select_strategy (`str`, *optional*): + The feature selection strategy used to select the vision feature from the vision backbone. + Can be one of `"default"` or `"full"` + """ return self.model.get_image_features( pixel_values=pixel_values, image_sizes=image_sizes, vision_feature_layer=vision_feature_layer, vision_feature_select_strategy=vision_feature_select_strategy, + **kwargs, ) - @can_return_tuple + @check_model_inputs(tie_last_hidden_states=False) @auto_docstring def forward( self, @@ -806,14 +813,6 @@ def forward( output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states ) return_dict = return_dict if return_dict is not None else self.config.use_return_dict - vision_feature_layer = ( - vision_feature_layer if vision_feature_layer is not None else self.config.vision_feature_layer - ) - vision_feature_select_strategy = ( - vision_feature_select_strategy - if vision_feature_select_strategy is not None - else self.config.vision_feature_select_strategy - ) outputs = self.model( input_ids=input_ids, @@ -948,16 +947,30 @@ def _prepare_4d_causal_attention_mask_with_cache_position( return causal_mask + @auto_docstring def get_video_features( self, pixel_values: torch.FloatTensor, vision_feature_layer: int | list[int] | None = None, vision_feature_select_strategy: str | None = None, - ): + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + pixel_values (`torch.FloatTensor]` of shape `(batch_size, num_frames, channels, height, width)`) + The tensors corresponding to the input video. + vision_feature_layer (`Union[int, list[int]]`, *optional;*): + The index of the layer to select the vision feature. If multiple indices are provided, + the vision feature of the corresponding indices will be concatenated to form the + vision features. + vision_feature_select_strategy (`str`, *optional*): + The feature selection strategy used to select the vision feature from the vision backbone. + Can be one of `"default"` or `"full"` + """ return self.model.get_video_features( pixel_values=pixel_values, vision_feature_layer=vision_feature_layer, vision_feature_select_strategy=vision_feature_select_strategy, + **kwargs, ) diff --git a/src/transformers/models/llava_next_video/modular_llava_next_video.py b/src/transformers/models/llava_next_video/modular_llava_next_video.py index b8c37dfe8b0b..345798810cb3 100644 --- a/src/transformers/models/llava_next_video/modular_llava_next_video.py +++ b/src/transformers/models/llava_next_video/modular_llava_next_video.py @@ -31,8 +31,10 @@ from ...cache_utils import Cache from ...configuration_utils import PreTrainedConfig from ...modeling_flash_attention_utils import FlashAttentionKwargs +from ...modeling_outputs import BaseModelOutputWithPooling from ...processing_utils import Unpack -from ...utils import logging, torch_compilable_check +from ...utils import auto_docstring, logging, torch_compilable_check +from ...utils.generic import check_model_inputs from ..auto import CONFIG_MAPPING, AutoConfig @@ -268,41 +270,32 @@ def __init__(self, config: LlavaNextVideoConfig, **super_kwargs): self.vision_resampler = LlavaNextVideoPooler(config) self.post_init() + @check_model_inputs(tie_last_hidden_states=False) + @auto_docstring( + custom_intro="Obtains image last hidden states from the vision tower and apply multimodal projection." + ) def get_image_features( self, pixel_values: torch.FloatTensor, image_sizes: torch.Tensor, vision_feature_layer: int | list[int] | None = None, vision_feature_select_strategy: str | None = None, - ): - """ - Obtains image last hidden states from the vision tower and apply multimodal projection. - - Args: - pixel_values (`torch.FloatTensor]` of shape `(batch_size, num_patches, channels, height, width)`) - The tensors corresponding to the input images. - image_sizes (`torch.Tensor` of shape `(num_images, 2)`) - Actual image size of each images (H, W). - vision_feature_layer (`Union[int, list[int]]`, *optional*): - The index of the layer to select the vision feature. If multiple indices are provided, - the vision feature of the corresponding indices will be concatenated to form the - vision features. - vision_feature_select_strategy (`str`, *optional*): - The feature selection strategy used to select the vision feature from the vision backbone. - Can be one of `"default"` or `"full"` - Returns: - image_features (list[`torch.Tensor`]): List of image feature tensor, each contains all the visual feature of all patches - and are of shape `(num_patches, image_length, embed_dim)`). + output_hidden_states: bool | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + pixel_values (`torch.FloatTensor]` of shape `(batch_size, num_patches, channels, height, width)`) + The tensors corresponding to the input images. + image_sizes (`torch.Tensor` of shape `(num_images, 2)`) + Actual image size of each images (H, W). + vision_feature_layer (`Union[int, list[int]]`, *optional*): + The index of the layer to select the vision feature. If multiple indices are provided, + the vision feature of the corresponding indices will be concatenated to form the + vision features. + vision_feature_select_strategy (`str`, *optional*): + The feature selection strategy used to select the vision feature from the vision backbone. + Can be one of `"default"` or `"full"` """ - vision_feature_layer = ( - vision_feature_layer if vision_feature_layer is not None else self.config.vision_feature_layer - ) - vision_feature_select_strategy = ( - vision_feature_select_strategy - if vision_feature_select_strategy is not None - else self.config.vision_feature_select_strategy - ) - # ! infer image_num_patches from image_sizes image_num_patches = [ image_size_to_num_patches( @@ -320,13 +313,18 @@ def get_image_features( # otherwise has to be stacked from list of (num_patches, num_channels, height, width) raise ValueError(f"pixel_values of shape {pixel_values.shape}, expect to be of 4 or 5 dimensions") - image_features = self.vision_tower(pixel_values, output_hidden_states=True) + image_outputs = self.vision_tower( + pixel_values, + output_hidden_states=True, # Ignore arg on purpose + return_dict=True, + **kwargs, + ) # If we have one vision feature layer, return the corresponding hidden states, # otherwise, select the hidden states of each feature layer and concatenate them if isinstance(vision_feature_layer, int): - selected_image_feature = image_features.hidden_states[vision_feature_layer] + selected_image_feature = image_outputs.hidden_states[vision_feature_layer] else: - hs_pool = [image_features.hidden_states[layer_idx] for layer_idx in vision_feature_layer] + hs_pool = [image_outputs.hidden_states[layer_idx] for layer_idx in vision_feature_layer] selected_image_feature = torch.cat(hs_pool, dim=-1) if vision_feature_select_strategy == "default": @@ -340,50 +338,48 @@ def get_image_features( vision_feature_select_strategy, image_newline=self.image_newline, ) - return image_features + image_outputs.pooler_output = image_features + return image_outputs + + @check_model_inputs(tie_last_hidden_states=False) + @auto_docstring( + custom_intro="Obtains video last hidden states from the vision tower and apply multimodal projection." + ) def get_video_features( self, pixel_values: torch.FloatTensor, vision_feature_layer: int | list[int] | None = None, vision_feature_select_strategy: str | None = None, - ): - """ - Obtains video last hidden states from the vision tower and apply multimodal projection. - - Args: - pixel_values (`torch.FloatTensor]` of shape `(batch_size, num_frames, channels, height, width)`) - The tensors corresponding to the input video. - vision_feature_layer (`Union[int, list[int]]`, *optional;*): - The index of the layer to select the vision feature. If multiple indices are provided, - the vision feature of the corresponding indices will be concatenated to form the - vision features. - vision_feature_select_strategy (`str`, *optional*): - The feature selection strategy used to select the vision feature from the vision backbone. - Can be one of `"default"` or `"full"` - Returns: - video_features (list[`torch.Tensor`]): List of video feature tensor, each contains all the visual feature of all patches - and are of shape `(num_videos, video_length, embed_dim)`). + output_hidden_states: bool | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + pixel_values (`torch.FloatTensor]` of shape `(batch_size, num_frames, channels, height, width)`) + The tensors corresponding to the input video. + vision_feature_layer (`Union[int, list[int]]`, *optional;*): + The index of the layer to select the vision feature. If multiple indices are provided, + the vision feature of the corresponding indices will be concatenated to form the + vision features. + vision_feature_select_strategy (`str`, *optional*): + The feature selection strategy used to select the vision feature from the vision backbone. + Can be one of `"default"` or `"full"` """ - vision_feature_layer = ( - vision_feature_layer if vision_feature_layer is not None else self.config.vision_feature_layer - ) - vision_feature_select_strategy = ( - vision_feature_select_strategy - if vision_feature_select_strategy is not None - else self.config.vision_feature_select_strategy - ) - batch_size, frames, channels, height, width = pixel_values.shape pixel_values = pixel_values.reshape(batch_size * frames, channels, height, width) - video_features = self.vision_tower(pixel_values, output_hidden_states=True) + video_outputs = self.vision_tower( + pixel_values, + output_hidden_states=True, # Ignore arg on purpose + return_dict=True, + **kwargs, + ) # If we have one vision feature layer, return the corresponding hidden states, # otherwise, select the hidden states of each feature layer and concatenate them if isinstance(vision_feature_layer, int): - selected_video_features = video_features.hidden_states[vision_feature_layer] + selected_video_features = video_outputs.hidden_states[vision_feature_layer] else: - hs_pool = [video_features.hidden_states[layer_idx] for layer_idx in vision_feature_layer] + hs_pool = [video_outputs.hidden_states[layer_idx] for layer_idx in vision_feature_layer] selected_video_features = torch.cat(hs_pool, dim=-1) if vision_feature_select_strategy == "default": @@ -393,7 +389,9 @@ def get_video_features( video_features = self.vision_resampler(selected_video_features) video_features = self.multi_modal_projector(video_features) video_features = torch.split(video_features, frames, dim=0) - return video_features + video_outputs.pooler_output = video_features + + return video_outputs def get_placeholder_mask( self, @@ -436,6 +434,8 @@ def get_placeholder_mask( ) return special_image_mask, special_video_mask + @check_model_inputs(tie_last_hidden_states=False) + @auto_docstring def forward( self, input_ids: torch.LongTensor | None = None, @@ -460,14 +460,6 @@ def forward( output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states ) return_dict = return_dict if return_dict is not None else self.config.use_return_dict - vision_feature_layer = ( - vision_feature_layer if vision_feature_layer is not None else self.config.vision_feature_layer - ) - vision_feature_select_strategy = ( - vision_feature_select_strategy - if vision_feature_select_strategy is not None - else self.config.vision_feature_select_strategy - ) if (input_ids is None) ^ (inputs_embeds is not None): raise ValueError("You must specify exactly one of input_ids or inputs_embeds") @@ -481,7 +473,8 @@ def forward( image_sizes, vision_feature_layer=vision_feature_layer, vision_feature_select_strategy=vision_feature_select_strategy, - ) + return_dict=True, + ).pooler_output image_features = torch.cat(image_features, dim=0).to(inputs_embeds.device, inputs_embeds.dtype) special_image_mask, _ = self.get_placeholder_mask( input_ids, inputs_embeds=inputs_embeds, image_features=image_features @@ -493,7 +486,8 @@ def forward( pixel_values_videos, vision_feature_layer=vision_feature_layer, vision_feature_select_strategy=vision_feature_select_strategy, - ) + return_dict=True, + ).pooler_output video_features = [feature.flatten(0, 1) for feature in video_features] video_feature_lens = [feature.size(0) for feature in video_features] video_features = torch.cat(video_features, dim=0) @@ -529,18 +523,34 @@ def forward( class LlavaNextVideoForConditionalGeneration(LlavaNextForConditionalGeneration): + @auto_docstring def get_video_features( self, pixel_values: torch.FloatTensor, vision_feature_layer: int | list[int] | None = None, vision_feature_select_strategy: str | None = None, - ): + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + pixel_values (`torch.FloatTensor]` of shape `(batch_size, num_frames, channels, height, width)`) + The tensors corresponding to the input video. + vision_feature_layer (`Union[int, list[int]]`, *optional;*): + The index of the layer to select the vision feature. If multiple indices are provided, + the vision feature of the corresponding indices will be concatenated to form the + vision features. + vision_feature_select_strategy (`str`, *optional*): + The feature selection strategy used to select the vision feature from the vision backbone. + Can be one of `"default"` or `"full"` + """ return self.model.get_video_features( pixel_values=pixel_values, vision_feature_layer=vision_feature_layer, vision_feature_select_strategy=vision_feature_select_strategy, + **kwargs, ) + @check_model_inputs(tie_last_hidden_states=False) + @auto_docstring def forward( self, input_ids: torch.LongTensor | None = None, @@ -630,14 +640,6 @@ def forward( output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states ) return_dict = return_dict if return_dict is not None else self.config.use_return_dict - vision_feature_layer = ( - vision_feature_layer if vision_feature_layer is not None else self.config.vision_feature_layer - ) - vision_feature_select_strategy = ( - vision_feature_select_strategy - if vision_feature_select_strategy is not None - else self.config.vision_feature_select_strategy - ) outputs = self.model( input_ids=input_ids, diff --git a/src/transformers/models/llava_onevision/modeling_llava_onevision.py b/src/transformers/models/llava_onevision/modeling_llava_onevision.py index 9f920125ce6f..3c7931e2e4ce 100644 --- a/src/transformers/models/llava_onevision/modeling_llava_onevision.py +++ b/src/transformers/models/llava_onevision/modeling_llava_onevision.py @@ -31,10 +31,11 @@ from ...generation import GenerationMixin from ...image_processing_utils import select_best_resolution from ...modeling_flash_attention_utils import FlashAttentionKwargs -from ...modeling_outputs import BaseModelOutputWithPast, ModelOutput +from ...modeling_outputs import BaseModelOutputWithPast, BaseModelOutputWithPooling, ModelOutput from ...modeling_utils import PreTrainedModel from ...processing_utils import Unpack -from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, torch_compilable_check +from ...utils import TransformersKwargs, auto_docstring, torch_compilable_check +from ...utils.generic import check_model_inputs from ..auto import AutoModel from .configuration_llava_onevision import LlavaOnevisionConfig @@ -299,7 +300,7 @@ def pack_image_features(self, image_features, image_sizes, image_newline=None, v image_newline (`torch.Tensor` of shape `(embed_dim)`) New line embedding vector. vision_aspect_ratio (`str`, *optional*, "anyres_max_9"): - Aspect ratio used when processong image features. The default value is "anyres_max_9". + Aspect ratio used when processing image features. The default value is "anyres_max_9". Returns: image_features (`torch.Tensor` of shape `(all_feat_len, embed_dim)`) feature_lens (`list[int]`) @@ -352,6 +353,10 @@ def pack_image_features(self, image_features, image_sizes, image_newline=None, v feature_lens = torch.tensor(feature_lens, dtype=torch.long, device=image_features[0].device) return new_image_features, feature_lens + @check_model_inputs(tie_last_hidden_states=False) + @auto_docstring( + custom_intro="Obtains image last hidden states from the vision tower and apply multimodal projection." + ) def get_image_features( self, pixel_values: torch.FloatTensor, @@ -360,40 +365,17 @@ def get_image_features( vision_feature_select_strategy: str | None = None, vision_aspect_ratio: str | None = None, batch_num_images: torch.LongTensor | None = None, - ): - """ - Obtains image last hidden states from the vision tower and apply multimodal projection. - - Args: - pixel_values (`torch.FloatTensor]` of shape `(batch_size, num_patches, channels, height, width)`) - The tensors corresponding to the input images. - image_sizes (`torch.Tensor` of shape `(num_images, 2)`) - Actual image size of each images (H, W). - vision_feature_layer (`Union[int, list[int]]`): - The index of the layer to select the vision feature. If multiple indices are provided, - the vision feature of the corresponding indices will be concatenated to form the - vision features. - vision_feature_select_strategy (`str`): - The feature selection strategy used to select the vision feature from the vision backbone. - Can be one of `"default"` or `"full"` - batch_num_images (`torch.LongTensor`, *optional*): - Number of images in each sample. - Returns: - image_features (list[`torch.Tensor`]): List of image feature tensor, each contains all the visual feature of all patches - and are of shape `(num_patches, image_length, embed_dim)`). + output_hidden_states: bool | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + image_sizes (`torch.Tensor` of shape `(num_images, 2)`): + Actual image size of each images (H, W). + vision_aspect_ratio (`str`, *optional*, defaults to `"anyres_max_9"`): + Aspect ratio used when processing image features. The default value is "anyres_max_9". + batch_num_images (`torch.LongTensor`, *optional*): + Number of images in each sample. """ - vision_feature_layer = ( - vision_feature_layer if vision_feature_layer is not None else self.config.vision_feature_layer - ) - vision_feature_select_strategy = ( - vision_feature_select_strategy - if vision_feature_select_strategy is not None - else self.config.vision_feature_select_strategy - ) - vision_aspect_ratio = ( - vision_aspect_ratio if vision_aspect_ratio is not None else self.config.vision_aspect_ratio - ) - # ! infer image_num_patches from image_sizes if batch_num_images is None: # treat this as a single-image case for backward compatibility @@ -418,13 +400,18 @@ def get_image_features( # otherwise has to be stacked from list of (num_patches, num_channels, height, width) raise ValueError(f"pixel_values of shape {pixel_values.shape}, expect to be of 4 or 5 dimensions") - image_features = self.vision_tower(pixel_values, output_hidden_states=True) + image_outputs = self.vision_tower( + pixel_values, + output_hidden_states=True, # Ignore arg on purpose + return_dict=True, + **kwargs, + ) # If we have one vision feature layer, return the corresponding hidden states, # otherwise, select the hidden states of each feature layer and concatenate them if isinstance(vision_feature_layer, int): - selected_image_feature = image_features.hidden_states[vision_feature_layer] + selected_image_feature = image_outputs.hidden_states[vision_feature_layer] else: - hs_pool = [image_features.hidden_states[layer_idx] for layer_idx in vision_feature_layer] + hs_pool = [image_outputs.hidden_states[layer_idx] for layer_idx in vision_feature_layer] selected_image_feature = torch.cat(hs_pool, dim=-1) if vision_feature_select_strategy == "default": @@ -438,7 +425,9 @@ def get_image_features( image_newline=self.image_newline, vision_aspect_ratio=vision_aspect_ratio, ) - return image_features + image_outputs.pooler_output = image_features + + return image_outputs def get_placeholder_mask( self, @@ -481,7 +470,7 @@ def get_placeholder_mask( ) return special_image_mask, special_video_mask - @can_return_tuple + @check_model_inputs(tie_last_hidden_states=False) @auto_docstring def forward( self, @@ -509,7 +498,7 @@ def forward( image_sizes_videos (`torch.LongTensor` of shape `(batch_size, frames, 2)`, *optional*): The sizes of the videos in the batch, being (height, width) for each frame in the video. vision_aspect_ratio (`str`, *optional*, defaults to `"anyres_max_9"`): - Aspect ratio used when processong image features. The default value is "anyres_max_9". + Aspect ratio used when processing image features. The default value is "anyres_max_9". batch_num_images (`torch.LongTensor`, *optional*): Number of images in each sample. """ @@ -519,17 +508,6 @@ def forward( output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states ) return_dict = return_dict if return_dict is not None else self.config.use_return_dict - vision_feature_layer = ( - vision_feature_layer if vision_feature_layer is not None else self.config.vision_feature_layer - ) - vision_feature_select_strategy = ( - vision_feature_select_strategy - if vision_feature_select_strategy is not None - else self.config.vision_feature_select_strategy - ) - vision_aspect_ratio = ( - vision_aspect_ratio if vision_aspect_ratio is not None else self.config.vision_aspect_ratio - ) if (input_ids is None) ^ (inputs_embeds is not None): raise ValueError("You must specify exactly one of input_ids or inputs_embeds") @@ -545,7 +523,8 @@ def forward( vision_feature_layer=vision_feature_layer, vision_feature_select_strategy=vision_feature_select_strategy, batch_num_images=batch_num_images, - ) + return_dict=True, + ).pooler_output image_features = torch.cat(image_features, dim=0) image_features = image_features.to(inputs_embeds.device, inputs_embeds.dtype) special_image_mask, _ = self.get_placeholder_mask( @@ -559,7 +538,8 @@ def forward( pixel_values_videos, vision_feature_layer=vision_feature_layer, vision_feature_select_strategy=vision_feature_select_strategy, - ) + return_dict=True, + ).pooler_output image_newline = ( self.image_newline[None, None, :].repeat(video_features.shape[0], 1, 1).to(video_features.device) ) @@ -592,39 +572,44 @@ def forward( video_hidden_states=video_features if pixel_values_videos is not None else None, ) + @check_model_inputs(tie_last_hidden_states=False) + @auto_docstring( + custom_intro="Obtains video last hidden states from the vision tower, apply multimodal projection and pooling." + ) def get_video_features( self, pixel_values: torch.FloatTensor, - vision_feature_layer: int | list[int], - vision_feature_select_strategy: str, - ): - """ - Obtains video last hidden states from the vision tower, apply multimodal projection and pooling. - - Args: - pixel_values (`torch.FloatTensor]` of shape `(batch_size, num_frames, channels, height, width)`) - The tensors corresponding to the input video. - vision_feature_layer (`Union[int, list[int]], *optional*, defaults to -2`): - The index of the layer to select the vision feature. If multiple indices are provided, - the vision feature of the corresponding indices will be concatenated to form the - vision features. - vision_feature_select_strategy (`str`): - The feature selection strategy used to select the vision feature from the vision backbone. - Can be one of `"default"` or `"full"` - Returns: - video_features (list[`torch.Tensor`]): List of video feature tensor, each contains all the visual feature of all patches - and are of shape `(num_videos, video_length, embed_dim)`). + vision_feature_layer: int | list[int] | None = None, + vision_feature_select_strategy: str | None = None, + output_hidden_states: bool | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + pixel_values (`torch.FloatTensor]` of shape `(batch_size, num_frames, channels, height, width)`) + The tensors corresponding to the input video. + vision_feature_layer (`Union[int, list[int]], *optional*, defaults to -2`): + The index of the layer to select the vision feature. If multiple indices are provided, + the vision feature of the corresponding indices will be concatenated to form the + vision features. + vision_feature_select_strategy (`str`): + The feature selection strategy used to select the vision feature from the vision backbone. + Can be one of `"default"` or `"full"` """ batch_size, frames, channels, height, width = pixel_values.shape pixel_values = pixel_values.view(batch_size * frames, channels, height, width) - video_features = self.vision_tower(pixel_values, output_hidden_states=True) + vision_outputs = self.vision_tower( + pixel_values, + output_hidden_states=True, # Ignore arg on purpose + return_dict=True, + **kwargs, + ) # If we have one vision feature layer, return the corresponding hidden states, # otherwise, select the hidden states of each feature layer and concatenate them if isinstance(vision_feature_layer, int): - selected_video_feature = video_features.hidden_states[vision_feature_layer] + selected_video_feature = vision_outputs.hidden_states[vision_feature_layer] else: - hs_pool = [video_features.hidden_states[layer_idx] for layer_idx in vision_feature_layer] + hs_pool = [vision_outputs.hidden_states[layer_idx] for layer_idx in vision_feature_layer] selected_video_feature = torch.cat(hs_pool, dim=-1) if vision_feature_select_strategy == "default": @@ -633,8 +618,9 @@ def get_video_features( video_features = self.apply_pooling(video_features) video_features = video_features.reshape(batch_size, frames * video_features.shape[1], -1) + vision_outputs.pooler_output = video_features - return video_features + return vision_outputs def apply_pooling(self, image_features): height = width = self.config.vision_config.image_size // self.config.vision_config.patch_size @@ -689,21 +675,36 @@ def pack_image_features(self, image_features, image_sizes, vision_feature_select image_newline=image_newline, ) + @auto_docstring def get_image_features( self, pixel_values: torch.FloatTensor, image_sizes: torch.Tensor, vision_feature_layer: int | list[int] | None = None, vision_feature_select_strategy: str | None = None, - ): + vision_aspect_ratio: str | None = None, + batch_num_images: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + image_sizes (`torch.Tensor` of shape `(num_images, 2)`): + Actual image size of each images (H, W). + vision_aspect_ratio (`str`, *optional*, defaults to `"anyres_max_9"`): + Aspect ratio used when processing image features. The default value is "anyres_max_9". + batch_num_images (`torch.LongTensor`, *optional*): + Number of images in each sample. + """ return self.model.get_image_features( pixel_values=pixel_values, image_sizes=image_sizes, vision_feature_layer=vision_feature_layer, vision_feature_select_strategy=vision_feature_select_strategy, + vision_aspect_ratio=vision_aspect_ratio, + batch_num_images=batch_num_images, + **kwargs, ) - @can_return_tuple + @check_model_inputs(tie_last_hidden_states=False) @auto_docstring def forward( self, @@ -733,7 +734,7 @@ def forward( image_sizes_videos (`torch.LongTensor` of shape `(batch_size, frames, 2)`, *optional*): The sizes of the videos in the batch, being (height, width) for each frame in the video. vision_aspect_ratio (`str`, *optional*, defaults to `"anyres_max_9"`): - Aspect ratio used when processong image features. The default value is "anyres_max_9". + Aspect ratio used when processing image features. The default value is "anyres_max_9". batch_num_images (`torch.LongTensor`, *optional*): Number of images in each sample. labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): @@ -776,17 +777,6 @@ def forward( output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states ) return_dict = return_dict if return_dict is not None else self.config.use_return_dict - vision_feature_layer = ( - vision_feature_layer if vision_feature_layer is not None else self.config.vision_feature_layer - ) - vision_feature_select_strategy = ( - vision_feature_select_strategy - if vision_feature_select_strategy is not None - else self.config.vision_feature_select_strategy - ) - vision_aspect_ratio = ( - vision_aspect_ratio if vision_aspect_ratio is not None else self.config.vision_aspect_ratio - ) outputs = self.model( input_ids=input_ids, @@ -927,16 +917,30 @@ def _prepare_4d_causal_attention_mask_with_cache_position( return causal_mask + @auto_docstring def get_video_features( self, pixel_values: torch.FloatTensor, vision_feature_layer: int | list[int] | None = None, vision_feature_select_strategy: str | None = None, - ): + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + pixel_values (`torch.FloatTensor]` of shape `(batch_size, num_frames, channels, height, width)`) + The tensors corresponding to the input video. + vision_feature_layer (`Union[int, list[int]]`, *optional;*): + The index of the layer to select the vision feature. If multiple indices are provided, + the vision feature of the corresponding indices will be concatenated to form the + vision features. + vision_feature_select_strategy (`str`, *optional*): + The feature selection strategy used to select the vision feature from the vision backbone. + Can be one of `"default"` or `"full"` + """ return self.model.get_video_features( pixel_values=pixel_values, vision_feature_layer=vision_feature_layer, vision_feature_select_strategy=vision_feature_select_strategy, + **kwargs, ) diff --git a/src/transformers/models/llava_onevision/modular_llava_onevision.py b/src/transformers/models/llava_onevision/modular_llava_onevision.py index 1df8da4380c3..6c49d02662f2 100644 --- a/src/transformers/models/llava_onevision/modular_llava_onevision.py +++ b/src/transformers/models/llava_onevision/modular_llava_onevision.py @@ -45,13 +45,10 @@ get_image_size, ) from ...modeling_flash_attention_utils import FlashAttentionKwargs +from ...modeling_outputs import BaseModelOutputWithPooling from ...processing_utils import Unpack -from ...utils import ( - TensorType, - auto_docstring, - can_return_tuple, - logging, -) +from ...utils import TensorType, auto_docstring, logging +from ...utils.generic import check_model_inputs from .image_processing_llava_onevision import LlavaOnevisionImageProcessorKwargs @@ -246,7 +243,7 @@ def pack_image_features(self, image_features, image_sizes, image_newline=None, v image_newline (`torch.Tensor` of shape `(embed_dim)`) New line embedding vector. vision_aspect_ratio (`str`, *optional*, "anyres_max_9"): - Aspect ratio used when processong image features. The default value is "anyres_max_9". + Aspect ratio used when processing image features. The default value is "anyres_max_9". Returns: image_features (`torch.Tensor` of shape `(all_feat_len, embed_dim)`) feature_lens (`list[int]`) @@ -313,6 +310,10 @@ def apply_pooling(self, image_features): image_features = image_features.view(batch_frames, -1, dim) return image_features + @check_model_inputs(tie_last_hidden_states=False) + @auto_docstring( + custom_intro="Obtains image last hidden states from the vision tower and apply multimodal projection." + ) def get_image_features( self, pixel_values: torch.FloatTensor, @@ -321,40 +322,17 @@ def get_image_features( vision_feature_select_strategy: str | None = None, vision_aspect_ratio: str | None = None, batch_num_images: torch.LongTensor | None = None, - ): - """ - Obtains image last hidden states from the vision tower and apply multimodal projection. - - Args: - pixel_values (`torch.FloatTensor]` of shape `(batch_size, num_patches, channels, height, width)`) - The tensors corresponding to the input images. - image_sizes (`torch.Tensor` of shape `(num_images, 2)`) - Actual image size of each images (H, W). - vision_feature_layer (`Union[int, list[int]]`): - The index of the layer to select the vision feature. If multiple indices are provided, - the vision feature of the corresponding indices will be concatenated to form the - vision features. - vision_feature_select_strategy (`str`): - The feature selection strategy used to select the vision feature from the vision backbone. - Can be one of `"default"` or `"full"` - batch_num_images (`torch.LongTensor`, *optional*): - Number of images in each sample. - Returns: - image_features (list[`torch.Tensor`]): List of image feature tensor, each contains all the visual feature of all patches - and are of shape `(num_patches, image_length, embed_dim)`). + output_hidden_states: bool | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + image_sizes (`torch.Tensor` of shape `(num_images, 2)`): + Actual image size of each images (H, W). + vision_aspect_ratio (`str`, *optional*, defaults to `"anyres_max_9"`): + Aspect ratio used when processing image features. The default value is "anyres_max_9". + batch_num_images (`torch.LongTensor`, *optional*): + Number of images in each sample. """ - vision_feature_layer = ( - vision_feature_layer if vision_feature_layer is not None else self.config.vision_feature_layer - ) - vision_feature_select_strategy = ( - vision_feature_select_strategy - if vision_feature_select_strategy is not None - else self.config.vision_feature_select_strategy - ) - vision_aspect_ratio = ( - vision_aspect_ratio if vision_aspect_ratio is not None else self.config.vision_aspect_ratio - ) - # ! infer image_num_patches from image_sizes if batch_num_images is None: # treat this as a single-image case for backward compatibility @@ -379,13 +357,18 @@ def get_image_features( # otherwise has to be stacked from list of (num_patches, num_channels, height, width) raise ValueError(f"pixel_values of shape {pixel_values.shape}, expect to be of 4 or 5 dimensions") - image_features = self.vision_tower(pixel_values, output_hidden_states=True) + image_outputs = self.vision_tower( + pixel_values, + output_hidden_states=True, # Ignore arg on purpose + return_dict=True, + **kwargs, + ) # If we have one vision feature layer, return the corresponding hidden states, # otherwise, select the hidden states of each feature layer and concatenate them if isinstance(vision_feature_layer, int): - selected_image_feature = image_features.hidden_states[vision_feature_layer] + selected_image_feature = image_outputs.hidden_states[vision_feature_layer] else: - hs_pool = [image_features.hidden_states[layer_idx] for layer_idx in vision_feature_layer] + hs_pool = [image_outputs.hidden_states[layer_idx] for layer_idx in vision_feature_layer] selected_image_feature = torch.cat(hs_pool, dim=-1) if vision_feature_select_strategy == "default": @@ -399,41 +382,48 @@ def get_image_features( image_newline=self.image_newline, vision_aspect_ratio=vision_aspect_ratio, ) - return image_features + image_outputs.pooler_output = image_features + + return image_outputs + @check_model_inputs(tie_last_hidden_states=False) + @auto_docstring( + custom_intro="Obtains video last hidden states from the vision tower, apply multimodal projection and pooling." + ) def get_video_features( self, pixel_values: torch.FloatTensor, - vision_feature_layer: int | list[int], - vision_feature_select_strategy: str, - ): - """ - Obtains video last hidden states from the vision tower, apply multimodal projection and pooling. - - Args: - pixel_values (`torch.FloatTensor]` of shape `(batch_size, num_frames, channels, height, width)`) - The tensors corresponding to the input video. - vision_feature_layer (`Union[int, list[int]], *optional*, defaults to -2`): - The index of the layer to select the vision feature. If multiple indices are provided, - the vision feature of the corresponding indices will be concatenated to form the - vision features. - vision_feature_select_strategy (`str`): - The feature selection strategy used to select the vision feature from the vision backbone. - Can be one of `"default"` or `"full"` - Returns: - video_features (list[`torch.Tensor`]): List of video feature tensor, each contains all the visual feature of all patches - and are of shape `(num_videos, video_length, embed_dim)`). + vision_feature_layer: int | list[int] | None = None, + vision_feature_select_strategy: str | None = None, + output_hidden_states: bool | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + pixel_values (`torch.FloatTensor]` of shape `(batch_size, num_frames, channels, height, width)`) + The tensors corresponding to the input video. + vision_feature_layer (`Union[int, list[int]], *optional*, defaults to -2`): + The index of the layer to select the vision feature. If multiple indices are provided, + the vision feature of the corresponding indices will be concatenated to form the + vision features. + vision_feature_select_strategy (`str`): + The feature selection strategy used to select the vision feature from the vision backbone. + Can be one of `"default"` or `"full"` """ batch_size, frames, channels, height, width = pixel_values.shape pixel_values = pixel_values.view(batch_size * frames, channels, height, width) - video_features = self.vision_tower(pixel_values, output_hidden_states=True) + vision_outputs = self.vision_tower( + pixel_values, + output_hidden_states=True, # Ignore arg on purpose + return_dict=True, + **kwargs, + ) # If we have one vision feature layer, return the corresponding hidden states, # otherwise, select the hidden states of each feature layer and concatenate them if isinstance(vision_feature_layer, int): - selected_video_feature = video_features.hidden_states[vision_feature_layer] + selected_video_feature = vision_outputs.hidden_states[vision_feature_layer] else: - hs_pool = [video_features.hidden_states[layer_idx] for layer_idx in vision_feature_layer] + hs_pool = [vision_outputs.hidden_states[layer_idx] for layer_idx in vision_feature_layer] selected_video_feature = torch.cat(hs_pool, dim=-1) if vision_feature_select_strategy == "default": @@ -442,9 +432,12 @@ def get_video_features( video_features = self.apply_pooling(video_features) video_features = video_features.reshape(batch_size, frames * video_features.shape[1], -1) + vision_outputs.pooler_output = video_features - return video_features + return vision_outputs + @check_model_inputs(tie_last_hidden_states=False) + @auto_docstring def forward( self, input_ids: torch.LongTensor | None = None, @@ -471,7 +464,7 @@ def forward( image_sizes_videos (`torch.LongTensor` of shape `(batch_size, frames, 2)`, *optional*): The sizes of the videos in the batch, being (height, width) for each frame in the video. vision_aspect_ratio (`str`, *optional*, defaults to `"anyres_max_9"`): - Aspect ratio used when processong image features. The default value is "anyres_max_9". + Aspect ratio used when processing image features. The default value is "anyres_max_9". batch_num_images (`torch.LongTensor`, *optional*): Number of images in each sample. """ @@ -481,17 +474,6 @@ def forward( output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states ) return_dict = return_dict if return_dict is not None else self.config.use_return_dict - vision_feature_layer = ( - vision_feature_layer if vision_feature_layer is not None else self.config.vision_feature_layer - ) - vision_feature_select_strategy = ( - vision_feature_select_strategy - if vision_feature_select_strategy is not None - else self.config.vision_feature_select_strategy - ) - vision_aspect_ratio = ( - vision_aspect_ratio if vision_aspect_ratio is not None else self.config.vision_aspect_ratio - ) if (input_ids is None) ^ (inputs_embeds is not None): raise ValueError("You must specify exactly one of input_ids or inputs_embeds") @@ -507,7 +489,8 @@ def forward( vision_feature_layer=vision_feature_layer, vision_feature_select_strategy=vision_feature_select_strategy, batch_num_images=batch_num_images, - ) + return_dict=True, + ).pooler_output image_features = torch.cat(image_features, dim=0) image_features = image_features.to(inputs_embeds.device, inputs_embeds.dtype) special_image_mask, _ = self.get_placeholder_mask( @@ -521,7 +504,8 @@ def forward( pixel_values_videos, vision_feature_layer=vision_feature_layer, vision_feature_select_strategy=vision_feature_select_strategy, - ) + return_dict=True, + ).pooler_output image_newline = ( self.image_newline[None, None, :].repeat(video_features.shape[0], 1, 1).to(video_features.device) ) @@ -556,7 +540,7 @@ def forward( class LlavaOnevisionForConditionalGeneration(LlavaNextVideoForConditionalGeneration): - @can_return_tuple + @check_model_inputs(tie_last_hidden_states=False) @auto_docstring def forward( self, @@ -586,7 +570,7 @@ def forward( image_sizes_videos (`torch.LongTensor` of shape `(batch_size, frames, 2)`, *optional*): The sizes of the videos in the batch, being (height, width) for each frame in the video. vision_aspect_ratio (`str`, *optional*, defaults to `"anyres_max_9"`): - Aspect ratio used when processong image features. The default value is "anyres_max_9". + Aspect ratio used when processing image features. The default value is "anyres_max_9". batch_num_images (`torch.LongTensor`, *optional*): Number of images in each sample. labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): @@ -629,17 +613,6 @@ def forward( output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states ) return_dict = return_dict if return_dict is not None else self.config.use_return_dict - vision_feature_layer = ( - vision_feature_layer if vision_feature_layer is not None else self.config.vision_feature_layer - ) - vision_feature_select_strategy = ( - vision_feature_select_strategy - if vision_feature_select_strategy is not None - else self.config.vision_feature_select_strategy - ) - vision_aspect_ratio = ( - vision_aspect_ratio if vision_aspect_ratio is not None else self.config.vision_aspect_ratio - ) outputs = self.model( input_ids=input_ids, @@ -725,6 +698,35 @@ def prepare_inputs_for_generation( return model_inputs + @auto_docstring + def get_image_features( + self, + pixel_values: torch.FloatTensor, + image_sizes: torch.Tensor, + vision_feature_layer: int | list[int] | None = None, + vision_feature_select_strategy: str | None = None, + vision_aspect_ratio: str | None = None, + batch_num_images: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + image_sizes (`torch.Tensor` of shape `(num_images, 2)`): + Actual image size of each images (H, W). + vision_aspect_ratio (`str`, *optional*, defaults to `"anyres_max_9"`): + Aspect ratio used when processing image features. The default value is "anyres_max_9". + batch_num_images (`torch.LongTensor`, *optional*): + Number of images in each sample. + """ + return self.model.get_image_features( + pixel_values=pixel_values, + image_sizes=image_sizes, + vision_feature_layer=vision_feature_layer, + vision_feature_select_strategy=vision_feature_select_strategy, + vision_aspect_ratio=vision_aspect_ratio, + batch_num_images=batch_num_images, + **kwargs, + ) + __all__ = [ "LlavaOnevisionImageProcessorFast", diff --git a/src/transformers/models/metaclip_2/modeling_metaclip_2.py b/src/transformers/models/metaclip_2/modeling_metaclip_2.py index 727893ebe306..2440e5fcb58f 100644 --- a/src/transformers/models/metaclip_2/modeling_metaclip_2.py +++ b/src/transformers/models/metaclip_2/modeling_metaclip_2.py @@ -18,14 +18,7 @@ from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling, ImageClassifierOutput from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel from ...processing_utils import Unpack -from ...utils import ( - ModelOutput, - TransformersKwargs, - auto_docstring, - can_return_tuple, - filter_out_non_signature_kwargs, - torch_int, -) +from ...utils import ModelOutput, TransformersKwargs, auto_docstring, can_return_tuple, torch_int from ...utils.generic import check_model_inputs from .configuration_metaclip_2 import MetaClip2Config, MetaClip2TextConfig, MetaClip2VisionConfig @@ -798,19 +791,16 @@ def __init__(self, config: MetaClip2Config): # Initialize weights and apply final processing self.post_init() - @filter_out_non_signature_kwargs() + @can_return_tuple @auto_docstring def get_text_features( self, - input_ids: torch.Tensor | None = None, + input_ids: torch.Tensor, attention_mask: torch.Tensor | None = None, position_ids: torch.Tensor | None = None, - ) -> torch.FloatTensor: + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: r""" - Returns: - text_features (`torch.FloatTensor` of shape `(batch_size, output_dim`): The text embeddings obtained by - applying the projection layer to the pooled output of [`MetaClip2TextModel`]. - Examples: ```python @@ -826,24 +816,23 @@ def get_text_features( input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids, + return_dict=True, + **kwargs, ) pooled_output = text_outputs.pooler_output - text_features = self.text_projection(pooled_output) + text_outputs.pooler_output = self.text_projection(pooled_output) - return text_features + return text_outputs - @filter_out_non_signature_kwargs() + @can_return_tuple @auto_docstring def get_image_features( self, pixel_values: torch.FloatTensor | None = None, interpolate_pos_encoding: bool = False, - ) -> torch.FloatTensor: + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: r""" - Returns: - image_features (`torch.FloatTensor` of shape `(batch_size, output_dim`): The image embeddings obtained by - applying the projection layer to the pooled output of [`MetaClip2VisionModel`]. - Examples: ```python @@ -864,11 +853,13 @@ def get_image_features( vision_outputs: BaseModelOutputWithPooling = self.vision_model( pixel_values=pixel_values, interpolate_pos_encoding=interpolate_pos_encoding, + return_dict=True, + **kwargs, ) pooled_output = vision_outputs.pooler_output - image_features = self.visual_projection(pooled_output) + vision_outputs.pooler_output = self.visual_projection(pooled_output) - return image_features + return vision_outputs @can_return_tuple @auto_docstring diff --git a/src/transformers/models/metaclip_2/modular_metaclip_2.py b/src/transformers/models/metaclip_2/modular_metaclip_2.py index 0f4dfd732089..4fc1aa6e3d6a 100644 --- a/src/transformers/models/metaclip_2/modular_metaclip_2.py +++ b/src/transformers/models/metaclip_2/modular_metaclip_2.py @@ -555,17 +555,16 @@ def forward( **kwargs, ) + @can_return_tuple + @auto_docstring def get_text_features( self, - input_ids: torch.Tensor | None = None, + input_ids: torch.Tensor, attention_mask: torch.Tensor | None = None, position_ids: torch.Tensor | None = None, - ): + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: r""" - Returns: - text_features (`torch.FloatTensor` of shape `(batch_size, output_dim`): The text embeddings obtained by - applying the projection layer to the pooled output of [`MetaClip2TextModel`]. - Examples: ```python @@ -581,18 +580,19 @@ def get_text_features( input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids, + return_dict=True, + **kwargs, ) + @can_return_tuple + @auto_docstring def get_image_features( self, pixel_values: torch.FloatTensor | None = None, interpolate_pos_encoding: bool = False, - ): + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: r""" - Returns: - image_features (`torch.FloatTensor` of shape `(batch_size, output_dim`): The image embeddings obtained by - applying the projection layer to the pooled output of [`MetaClip2VisionModel`]. - Examples: ```python @@ -613,6 +613,8 @@ def get_image_features( return super().get_image_features( pixel_values=pixel_values, interpolate_pos_encoding=interpolate_pos_encoding, + return_dict=True, + **kwargs, ) diff --git a/src/transformers/models/minimax/modeling_minimax.py b/src/transformers/models/minimax/modeling_minimax.py index 841c81694c8c..a2dff7e9401b 100644 --- a/src/transformers/models/minimax/modeling_minimax.py +++ b/src/transformers/models/minimax/modeling_minimax.py @@ -654,7 +654,7 @@ def forward( use_cache: bool | None = None, cache_position: torch.LongTensor | None = None, **kwargs: Unpack[TransformersKwargs], - ) -> MoeModelOutputWithPast: + ) -> tuple | MoeModelOutputWithPast: if (input_ids is None) ^ (inputs_embeds is not None): raise ValueError("You must specify exactly one of input_ids or inputs_embeds") diff --git a/src/transformers/models/minimax/modular_minimax.py b/src/transformers/models/minimax/modular_minimax.py index 5e2f67bc8150..c05e7db364f5 100644 --- a/src/transformers/models/minimax/modular_minimax.py +++ b/src/transformers/models/minimax/modular_minimax.py @@ -546,7 +546,7 @@ def forward( use_cache: bool | None = None, cache_position: torch.LongTensor | None = None, **kwargs: Unpack[TransformersKwargs], - ) -> MoeModelOutputWithPast: + ) -> tuple | MoeModelOutputWithPast: if (input_ids is None) ^ (inputs_embeds is not None): raise ValueError("You must specify exactly one of input_ids or inputs_embeds") diff --git a/src/transformers/models/mistral3/modeling_mistral3.py b/src/transformers/models/mistral3/modeling_mistral3.py index 91d009bdb632..03112e6fb77d 100644 --- a/src/transformers/models/mistral3/modeling_mistral3.py +++ b/src/transformers/models/mistral3/modeling_mistral3.py @@ -28,10 +28,11 @@ from ...cache_utils import Cache from ...generation import GenerationMixin from ...integrations import use_kernel_forward_from_hub -from ...modeling_outputs import BaseModelOutputWithPast, ModelOutput +from ...modeling_outputs import BaseModelOutputWithPast, BaseModelOutputWithPooling, ModelOutput from ...modeling_utils import PreTrainedModel from ...processing_utils import Unpack -from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, torch_compilable_check +from ...utils import TransformersKwargs, auto_docstring, torch_compilable_check +from ...utils.generic import check_model_inputs from ..auto import AutoModel from .configuration_mistral3 import Mistral3Config @@ -214,35 +215,27 @@ def get_input_embeddings(self): def set_input_embeddings(self, value): self.language_model.set_input_embeddings(value) + @check_model_inputs(tie_last_hidden_states=False) + @auto_docstring( + custom_intro="Obtains image last hidden states from the vision tower and apply multimodal projection." + ) def get_image_features( self, pixel_values: torch.FloatTensor, image_sizes: torch.Tensor, vision_feature_layer: int | list[int] | None = None, - **kwargs, - ): - """ - Obtains image last hidden states from the vision tower and apply multimodal projection. - - Args: - pixel_values (`torch.FloatTensor]` of shape `(batch_size, channels, height, width)`): - The tensors corresponding to the input images. - vision_feature_layer (`Union[int, list[int]]`, *optional*): - The index of the layer to select the vision feature. If multiple indices are provided, - the vision feature of the corresponding indices will be concatenated to form the - vision features. - image_sizes (`torch.Tensor`, *optional*): - Tensor containing the image sizes as returned by the processor. - Returns: - image_features (`torch.Tensor`): Image feature tensor of shape `(num_images, image_length, embed_dim)`). - """ - vision_feature_layer = ( - vision_feature_layer if vision_feature_layer is not None else self.config.vision_feature_layer - ) - + output_hidden_states: bool | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: kwargs = {k: v for k, v in kwargs.items() if v is not None} # this is not memory efficient at all (output_hidden_states=True) will save all the hidden states. - image_outputs = self.vision_tower(pixel_values, image_sizes=image_sizes, output_hidden_states=True, **kwargs) + image_outputs = self.vision_tower( + pixel_values, + image_sizes=image_sizes, + output_hidden_states=True, # Ignore arg on purpose + return_dict=True, + **kwargs, + ) # If we have one vision feature layer, return the corresponding hidden states, # otherwise, select the hidden states of each feature layer and concatenate them if isinstance(vision_feature_layer, int): @@ -257,7 +250,9 @@ def get_image_features( (torch.as_tensor(image_sizes, device=image_features.device) // downsample_ratio).prod(dim=-1).tolist() ) image_features = torch.split(image_features.squeeze(0), split_sizes) - return image_features + image_outputs.pooler_output = image_features + + return image_outputs def get_placeholder_mask( self, input_ids: torch.LongTensor, inputs_embeds: torch.FloatTensor, image_features: torch.FloatTensor @@ -283,7 +278,7 @@ def get_placeholder_mask( ) return special_image_mask - @can_return_tuple + @check_model_inputs(tie_last_hidden_states=False) @auto_docstring def forward( self, @@ -307,9 +302,6 @@ def forward( output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states ) return_dict = return_dict if return_dict is not None else self.config.use_return_dict - vision_feature_layer = ( - vision_feature_layer if vision_feature_layer is not None else self.config.vision_feature_layer - ) if (input_ids is None) ^ (inputs_embeds is not None): raise ValueError("You must specify exactly one of input_ids or inputs_embeds") @@ -322,7 +314,8 @@ def forward( pixel_values=pixel_values, vision_feature_layer=vision_feature_layer, image_sizes=image_sizes, - ) + return_dict=True, + ).pooler_output image_features = torch.cat(image_features, dim=0).to(inputs_embeds.device, inputs_embeds.dtype) special_image_mask = self.get_placeholder_mask( input_ids, inputs_embeds=inputs_embeds, image_features=image_features @@ -380,13 +373,14 @@ def set_input_embeddings(self, value): def get_output_embeddings(self) -> nn.Module: return self.lm_head + @auto_docstring def get_image_features( self, pixel_values: torch.FloatTensor, image_sizes: torch.Tensor, vision_feature_layer: int | list[int] | None = None, - **kwargs, - ): + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: return self.model.get_image_features( pixel_values=pixel_values, image_sizes=image_sizes, @@ -394,7 +388,7 @@ def get_image_features( **kwargs, ) - @can_return_tuple + @check_model_inputs(tie_last_hidden_states=False) @auto_docstring def forward( self, diff --git a/src/transformers/models/mistral3/modular_mistral3.py b/src/transformers/models/mistral3/modular_mistral3.py index 525b7b014f77..7ebf180b3e3e 100644 --- a/src/transformers/models/mistral3/modular_mistral3.py +++ b/src/transformers/models/mistral3/modular_mistral3.py @@ -18,8 +18,10 @@ from ...activations import ACT2FN from ...cache_utils import Cache +from ...modeling_outputs import BaseModelOutputWithPooling from ...processing_utils import Unpack -from ...utils import logging +from ...utils import auto_docstring, logging +from ...utils.generic import check_model_inputs from ..llava.modeling_llava import ( LlavaCausalLMOutputWithPast, LlavaForConditionalGeneration, @@ -118,35 +120,27 @@ class Mistral3PreTrainedModel(LlavaPreTrainedModel): class Mistral3Model(LlavaModel): + @check_model_inputs(tie_last_hidden_states=False) + @auto_docstring( + custom_intro="Obtains image last hidden states from the vision tower and apply multimodal projection." + ) def get_image_features( self, pixel_values: torch.FloatTensor, image_sizes: torch.Tensor, vision_feature_layer: int | list[int] | None = None, - **kwargs, - ): - """ - Obtains image last hidden states from the vision tower and apply multimodal projection. - - Args: - pixel_values (`torch.FloatTensor]` of shape `(batch_size, channels, height, width)`): - The tensors corresponding to the input images. - vision_feature_layer (`Union[int, list[int]]`, *optional*): - The index of the layer to select the vision feature. If multiple indices are provided, - the vision feature of the corresponding indices will be concatenated to form the - vision features. - image_sizes (`torch.Tensor`, *optional*): - Tensor containing the image sizes as returned by the processor. - Returns: - image_features (`torch.Tensor`): Image feature tensor of shape `(num_images, image_length, embed_dim)`). - """ - vision_feature_layer = ( - vision_feature_layer if vision_feature_layer is not None else self.config.vision_feature_layer - ) - + output_hidden_states: bool | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: kwargs = {k: v for k, v in kwargs.items() if v is not None} # this is not memory efficient at all (output_hidden_states=True) will save all the hidden states. - image_outputs = self.vision_tower(pixel_values, image_sizes=image_sizes, output_hidden_states=True, **kwargs) + image_outputs = self.vision_tower( + pixel_values, + image_sizes=image_sizes, + output_hidden_states=True, # Ignore arg on purpose + return_dict=True, + **kwargs, + ) # If we have one vision feature layer, return the corresponding hidden states, # otherwise, select the hidden states of each feature layer and concatenate them if isinstance(vision_feature_layer, int): @@ -161,8 +155,12 @@ def get_image_features( (torch.as_tensor(image_sizes, device=image_features.device) // downsample_ratio).prod(dim=-1).tolist() ) image_features = torch.split(image_features.squeeze(0), split_sizes) - return image_features + image_outputs.pooler_output = image_features + + return image_outputs + @check_model_inputs(tie_last_hidden_states=False) + @auto_docstring def forward( self, input_ids: torch.LongTensor | None = None, @@ -185,9 +183,6 @@ def forward( output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states ) return_dict = return_dict if return_dict is not None else self.config.use_return_dict - vision_feature_layer = ( - vision_feature_layer if vision_feature_layer is not None else self.config.vision_feature_layer - ) if (input_ids is None) ^ (inputs_embeds is not None): raise ValueError("You must specify exactly one of input_ids or inputs_embeds") @@ -200,7 +195,8 @@ def forward( pixel_values=pixel_values, vision_feature_layer=vision_feature_layer, image_sizes=image_sizes, - ) + return_dict=True, + ).pooler_output image_features = torch.cat(image_features, dim=0).to(inputs_embeds.device, inputs_embeds.dtype) special_image_mask = self.get_placeholder_mask( input_ids, inputs_embeds=inputs_embeds, image_features=image_features @@ -230,13 +226,14 @@ def forward( class Mistral3ForConditionalGeneration(LlavaForConditionalGeneration): + @auto_docstring def get_image_features( self, pixel_values: torch.FloatTensor, image_sizes: torch.Tensor, vision_feature_layer: int | list[int] | None = None, - **kwargs, - ): + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: return self.model.get_image_features( pixel_values=pixel_values, image_sizes=image_sizes, diff --git a/src/transformers/models/moonshine/modeling_moonshine.py b/src/transformers/models/moonshine/modeling_moonshine.py index 3a2c45d4002c..7d8731e0ddc0 100644 --- a/src/transformers/models/moonshine/modeling_moonshine.py +++ b/src/transformers/models/moonshine/modeling_moonshine.py @@ -556,7 +556,7 @@ def forward( input_values: torch.FloatTensor, attention_mask: torch.Tensor | None = None, **kwargs: Unpack[TransformersKwargs], - ) -> BaseModelOutputWithPast: + ) -> tuple | BaseModelOutputWithPast: r""" Args: input_values (`torch.FloatTensor` of shape `(batch_size, audio_length)`): diff --git a/src/transformers/models/moonshine/modular_moonshine.py b/src/transformers/models/moonshine/modular_moonshine.py index b38bcf051fe6..6517494ffe7f 100644 --- a/src/transformers/models/moonshine/modular_moonshine.py +++ b/src/transformers/models/moonshine/modular_moonshine.py @@ -519,7 +519,7 @@ def forward( input_values: torch.FloatTensor, attention_mask: torch.Tensor | None = None, **kwargs: Unpack[TransformersKwargs], - ) -> BaseModelOutputWithPast: + ) -> tuple | BaseModelOutputWithPast: r""" Args: input_values (`torch.FloatTensor` of shape `(batch_size, audio_length)`): diff --git a/src/transformers/models/ovis2/modeling_ovis2.py b/src/transformers/models/ovis2/modeling_ovis2.py index b1f8662706ed..57e5bcd305cb 100644 --- a/src/transformers/models/ovis2/modeling_ovis2.py +++ b/src/transformers/models/ovis2/modeling_ovis2.py @@ -31,14 +31,26 @@ from ...generation import GenerationMixin from ...integrations import use_kernel_forward_from_hub from ...modeling_layers import GradientCheckpointingLayer -from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPast +from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPast, BaseModelOutputWithPooling from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel from ...processing_utils import Unpack from ...utils import ModelOutput, TransformersKwargs, auto_docstring, can_return_tuple, torch_compilable_check +from ...utils.generic import check_model_inputs from ..auto import AutoModel from .configuration_ovis2 import Ovis2Config, Ovis2VisionConfig +@dataclass +@auto_docstring +class BaseModelOutputWithVisualIndicatorFeatures(BaseModelOutputWithPooling): + r""" + visual_indicator_features (`torch.FloatTensor` of shape `(batch_size, visual_indicator_size)`): + Visual indicator features extracted from the model, which can be used for auxiliary tasks or further processing. + """ + + visual_indicator_features: torch.FloatTensor | None = None + + @dataclass @auto_docstring( custom_intro=""" @@ -260,71 +272,10 @@ def forward(self, x): return down_proj -class Ovis2Attention(nn.Module): - """Multi-headed attention from 'Attention Is All You Need' paper""" - - def __init__(self, config): - super().__init__() - self.config = config - self.embed_dim = config.hidden_size - self.num_heads = config.num_attention_heads - self.head_dim = self.embed_dim // self.num_heads - if self.head_dim * self.num_heads != self.embed_dim: - raise ValueError( - f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim} and `num_heads`:" - f" {self.num_heads})." - ) - self.scale = self.head_dim**-0.5 - self.dropout = config.attention_dropout - self.is_causal = False - self.k_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=config.qkv_bias) - self.v_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=config.qkv_bias) - self.q_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=config.qkv_bias) - self.out_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=config.qkv_bias) - - def forward( - self, - hidden_states: torch.Tensor, - attention_mask: torch.Tensor | None = None, - **kwargs, - ) -> tuple[torch.Tensor, torch.Tensor | None]: - """Input shape: Batch x Time x Channel""" - - batch_size, seq_length, embed_dim = hidden_states.shape - - queries = self.q_proj(hidden_states) - keys = self.k_proj(hidden_states) - values = self.v_proj(hidden_states) - - queries = queries.view(batch_size, seq_length, self.num_heads, self.head_dim).transpose(1, 2) - keys = keys.view(batch_size, seq_length, self.num_heads, self.head_dim).transpose(1, 2) - values = values.view(batch_size, seq_length, self.num_heads, self.head_dim).transpose(1, 2) - - attention_interface: Callable = eager_attention_forward - if self.config._attn_implementation != "eager": - attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation] - - attn_output, attn_weights = attention_interface( - self, - queries, - keys, - values, - attention_mask, - is_causal=self.is_causal, - scaling=self.scale, - dropout=0.0 if not self.training else self.dropout, - ) - - attn_output = attn_output.reshape(batch_size, seq_length, embed_dim).contiguous() - attn_output = self.out_proj(attn_output) - - return attn_output, attn_weights - - class Ovis2VisionEncoderLayer(GradientCheckpointingLayer): def __init__(self, config: Ovis2VisionConfig): super().__init__() - self.attention = Ovis2Attention(config) + self.attention = Ovis2VisionAttention(config) self.ffn = Ovis2MLP(config) self.rms_norm1 = Ovis2RMSNorm(config.hidden_size, config.rms_norm_eps) self.rms_norm2 = Ovis2RMSNorm(config.hidden_size, config.rms_norm_eps) @@ -447,6 +398,10 @@ def hard_softmax(logits: torch.Tensor, dim: int): class Ovis2VisionModel(Ovis2PreTrainedModel): config: Ovis2VisionConfig + _can_record_outputs = { + "hidden_states": Ovis2VisionEncoderLayer, + "attentions": Ovis2VisionAttention, + } def __init__(self, config: Ovis2VisionConfig): super().__init__(config) @@ -463,7 +418,10 @@ def __init__(self, config: Ovis2VisionConfig): self.post_init() - def forward(self, pixel_values: torch.FloatTensor, **kwargs) -> tuple[torch.Tensor, torch.Tensor]: + @check_model_inputs + def forward( + self, pixel_values: torch.FloatTensor, **kwargs: Unpack[TransformersKwargs] + ) -> tuple | BaseModelOutputWithVisualIndicatorFeatures: outputs = self.transformer(pixel_values, **kwargs) last_hidden_state = outputs[0] if self.config.hidden_stride > 1: @@ -496,7 +454,10 @@ def forward(self, pixel_values: torch.FloatTensor, **kwargs) -> tuple[torch.Tens elif self.config.tokenize_function == "softmax": prob_token = nn.functional.softmax(logits, dim=-1) - return prob_token + return BaseModelOutputWithVisualIndicatorFeatures( + last_hidden_state=last_hidden_state, + pooler_output=prob_token, + ) @auto_docstring( @@ -524,27 +485,17 @@ def get_input_embeddings(self): def set_input_embeddings(self, value): self.language_model.set_input_embeddings(value) + @can_return_tuple + @auto_docstring( + custom_intro="Obtains image last hidden states from the vision tower and apply multimodal projection." + ) def get_image_features( self, pixel_values: torch.FloatTensor, - ) -> torch.FloatTensor: - """ - Obtains image last hidden states from the vision tower and apply multimodal projection. - - Args: - pixel_values (`torch.FloatTensor]` of shape `(batch_size, channels, height, width)`): - The tensors corresponding to the input images. - vision_feature_layer (`Union[int, list[int]]`, *optional*): - The index of the layer to select the vision feature. If multiple indices are provided, - the vision feature of the corresponding indices will be concatenated to form the - vision features. - vision_feature_select_strategy (`str`, *optional*): - The feature selection strategy used to select the vision feature from the vision backbone. - Can be one of `"default"` or `"full"` - Returns: - image_features (`torch.Tensor`): Image feature tensor of shape `(num_images, image_length, embed_dim)`). - """ - image_features = self.vision_tower(pixel_values) + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithVisualIndicatorFeatures: + image_outputs = self.vision_tower(pixel_values, return_dict=True, **kwargs) + image_features = image_outputs.pooler_output batch_size, img_seq_len, _ = image_features.shape padding_tensor = torch.zeros( (batch_size, img_seq_len, self.vision_tower.num_visual_indicator_tokens), @@ -561,9 +512,10 @@ def get_image_features( self.visual_vocab_size, dtype=torch.long, ).to(image_features.device) - visual_indicator_features = self.visual_embeddings_table(visual_indicator) + image_outputs.pooler_output = image_features + image_outputs.visual_indicator_features = self.visual_embeddings_table(visual_indicator) - return image_features, visual_indicator_features + return image_outputs def get_placeholder_mask( self, input_ids: torch.LongTensor, inputs_embeds: torch.FloatTensor, image_features: torch.FloatTensor @@ -620,7 +572,9 @@ def forward( inputs_embeds = self.get_input_embeddings()(input_ids) if pixel_values is not None: - image_features, visual_indicator_features = self.get_image_features(pixel_values=pixel_values) + image_outputs = self.get_image_features(pixel_values=pixel_values, return_dict=True) + image_features = image_outputs.pooler_output + visual_indicator_features = image_outputs.visual_indicator_features special_image_mask = self.get_placeholder_mask( input_ids, @@ -688,8 +642,11 @@ def set_input_embeddings(self, value): def get_output_embeddings(self) -> nn.Module: return self.lm_head - def get_image_features(self, pixel_values: torch.FloatTensor): - return self.model.get_image_features(pixel_values=pixel_values) + @auto_docstring + def get_image_features( + self, pixel_values: torch.FloatTensor, **kwargs: Unpack[TransformersKwargs] + ) -> tuple | BaseModelOutputWithVisualIndicatorFeatures: + return self.model.get_image_features(pixel_values=pixel_values, **kwargs) @can_return_tuple @auto_docstring diff --git a/src/transformers/models/ovis2/modular_ovis2.py b/src/transformers/models/ovis2/modular_ovis2.py index 409f9e37177a..aa659655fb49 100644 --- a/src/transformers/models/ovis2/modular_ovis2.py +++ b/src/transformers/models/ovis2/modular_ovis2.py @@ -13,6 +13,7 @@ # limitations under the License. import math +from dataclasses import dataclass import torch from torch import nn @@ -20,10 +21,11 @@ from ... import initialization as init from ...cache_utils import Cache from ...generation import GenerationMixin -from ...modeling_outputs import BaseModelOutput +from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling from ...modeling_utils import PreTrainedModel from ...processing_utils import Unpack from ...utils import TransformersKwargs, auto_docstring, can_return_tuple +from ...utils.generic import check_model_inputs from ..aimv2.modeling_aimv2 import Aimv2Attention, Aimv2EncoderLayer from ..auto import AutoModel from ..llama.modeling_llama import LlamaMLP, LlamaRMSNorm @@ -43,6 +45,17 @@ def hard_softmax(logits: torch.Tensor, dim: int): return ret +@dataclass +@auto_docstring +class BaseModelOutputWithVisualIndicatorFeatures(BaseModelOutputWithPooling): + r""" + visual_indicator_features (`torch.FloatTensor` of shape `(batch_size, visual_indicator_size)`): + Visual indicator features extracted from the model, which can be used for auxiliary tasks or further processing. + """ + + visual_indicator_features: torch.FloatTensor | None = None + + class Ovis2ModelOutputWithPast(LlavaNextModelOutputWithPast): pass @@ -83,7 +96,9 @@ class Ovis2VisionAttention(Aimv2Attention): class Ovis2VisionEncoderLayer(Aimv2EncoderLayer): - pass + def __init__(self, config: Ovis2VisionConfig): + super().__init__() + self.attention = Ovis2VisionAttention(config) class Ovis2VisionEncoder(SiglipEncoder): @@ -166,6 +181,10 @@ def _init_weights(self, module): class Ovis2VisionModel(Ovis2PreTrainedModel): config: Ovis2VisionConfig + _can_record_outputs = { + "hidden_states": Ovis2VisionEncoderLayer, + "attentions": Ovis2VisionAttention, + } def __init__(self, config: Ovis2VisionConfig): super().__init__(config) @@ -182,7 +201,10 @@ def __init__(self, config: Ovis2VisionConfig): self.post_init() - def forward(self, pixel_values: torch.FloatTensor, **kwargs) -> tuple[torch.Tensor, torch.Tensor]: + @check_model_inputs + def forward( + self, pixel_values: torch.FloatTensor, **kwargs: Unpack[TransformersKwargs] + ) -> tuple | BaseModelOutputWithVisualIndicatorFeatures: outputs = self.transformer(pixel_values, **kwargs) last_hidden_state = outputs[0] if self.config.hidden_stride > 1: @@ -215,7 +237,10 @@ def forward(self, pixel_values: torch.FloatTensor, **kwargs) -> tuple[torch.Tens elif self.config.tokenize_function == "softmax": prob_token = nn.functional.softmax(logits, dim=-1) - return prob_token + return BaseModelOutputWithVisualIndicatorFeatures( + last_hidden_state=last_hidden_state, + pooler_output=prob_token, + ) class Ovis2Model(LlavaModel): @@ -232,11 +257,17 @@ def __init__(self, config: Ovis2Config): self.language_model = AutoModel.from_config(config.text_config) del self.multi_modal_projector + @can_return_tuple + @auto_docstring( + custom_intro="Obtains image last hidden states from the vision tower and apply multimodal projection." + ) def get_image_features( self, pixel_values: torch.FloatTensor, - ) -> torch.FloatTensor: - image_features = self.vision_tower(pixel_values) + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithVisualIndicatorFeatures: + image_outputs = self.vision_tower(pixel_values, return_dict=True, **kwargs) + image_features = image_outputs.pooler_output batch_size, img_seq_len, _ = image_features.shape padding_tensor = torch.zeros( (batch_size, img_seq_len, self.vision_tower.num_visual_indicator_tokens), @@ -253,9 +284,10 @@ def get_image_features( self.visual_vocab_size, dtype=torch.long, ).to(image_features.device) - visual_indicator_features = self.visual_embeddings_table(visual_indicator) + image_outputs.pooler_output = image_features + image_outputs.visual_indicator_features = self.visual_embeddings_table(visual_indicator) - return image_features, visual_indicator_features + return image_outputs @can_return_tuple @auto_docstring @@ -288,7 +320,9 @@ def forward( inputs_embeds = self.get_input_embeddings()(input_ids) if pixel_values is not None: - image_features, visual_indicator_features = self.get_image_features(pixel_values=pixel_values) + image_outputs = self.get_image_features(pixel_values=pixel_values, return_dict=True) + image_features = image_outputs.pooler_output + visual_indicator_features = image_outputs.visual_indicator_features special_image_mask = self.get_placeholder_mask( input_ids, @@ -344,8 +378,11 @@ def __init__(self, config: Ovis2Config): super().__init__(config) self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) - def get_image_features(self, pixel_values: torch.FloatTensor): - return self.model.get_image_features(pixel_values=pixel_values) + @auto_docstring + def get_image_features( + self, pixel_values: torch.FloatTensor, **kwargs: Unpack[TransformersKwargs] + ) -> tuple | BaseModelOutputWithVisualIndicatorFeatures: + return self.model.get_image_features(pixel_values=pixel_values, **kwargs) @can_return_tuple @auto_docstring diff --git a/src/transformers/models/owlv2/modeling_owlv2.py b/src/transformers/models/owlv2/modeling_owlv2.py index 5cf69e944dd3..ded895ebb814 100644 --- a/src/transformers/models/owlv2/modeling_owlv2.py +++ b/src/transformers/models/owlv2/modeling_owlv2.py @@ -25,10 +25,12 @@ from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling from ...modeling_utils import PreTrainedModel +from ...processing_utils import Unpack from ...utils import ( ModelOutput, + TransformersKwargs, auto_docstring, - filter_out_non_signature_kwargs, + can_return_tuple, is_vision_available, logging, torch_int, @@ -972,23 +974,20 @@ def __init__(self, config: Owlv2Config): # Initialize weights and apply final processing self.post_init() - @filter_out_non_signature_kwargs() + @can_return_tuple @auto_docstring def get_text_features( self, input_ids: torch.Tensor, attention_mask: torch.Tensor | None = None, - ) -> torch.FloatTensor: + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: r""" input_ids (`torch.LongTensor` of shape `(batch_size * num_max_text_queries, sequence_length)`): Indices of input sequence tokens in the vocabulary. Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and [`PreTrainedTokenizer.__call__`] for details. [What are input IDs?](../glossary#input-ids) - Returns: - text_features (`torch.FloatTensor` of shape `(batch_size, output_dim`): The text embeddings obtained by - applying the projection layer to the pooled output of [`Owlv2TextModel`]. - Examples: ```python >>> import torch @@ -1003,23 +1002,26 @@ def get_text_features( ... text_features = model.get_text_features(**inputs) ```""" # Get embeddings for all text queries in all batch samples - text_outputs: BaseModelOutputWithPooling = self.text_model(input_ids=input_ids, attention_mask=attention_mask) - text_features = self.text_projection(text_outputs.pooler_output) + text_outputs: BaseModelOutputWithPooling = self.text_model( + input_ids=input_ids, + attention_mask=attention_mask, + return_dict=True, + **kwargs, + ) + pooled_output = text_outputs.pooler_output + text_outputs.pooler_output = self.text_projection(pooled_output) - return text_features + return text_outputs - @filter_out_non_signature_kwargs() + @can_return_tuple @auto_docstring def get_image_features( self, pixel_values: torch.Tensor, interpolate_pos_encoding: bool = False, - ) -> torch.FloatTensor: + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: r""" - Returns: - image_features (`torch.FloatTensor` of shape `(batch_size, output_dim`): The image embeddings obtained by - applying the projection layer to the pooled output of [`Owlv2VisionModel`]. - Examples: ```python >>> import torch @@ -1039,10 +1041,12 @@ def get_image_features( vision_outputs: BaseModelOutputWithPooling = self.vision_model( pixel_values=pixel_values, interpolate_pos_encoding=interpolate_pos_encoding, + return_dict=True, + **kwargs, ) - image_features = self.visual_projection(vision_outputs.pooler_output) + vision_outputs.pooler_output = self.visual_projection(vision_outputs.pooler_output) - return image_features + return vision_outputs @auto_docstring def forward( diff --git a/src/transformers/models/owlvit/modeling_owlvit.py b/src/transformers/models/owlvit/modeling_owlvit.py index ed8069996c4c..bbda1109fd3e 100644 --- a/src/transformers/models/owlvit/modeling_owlvit.py +++ b/src/transformers/models/owlvit/modeling_owlvit.py @@ -25,10 +25,12 @@ from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling from ...modeling_utils import PreTrainedModel +from ...processing_utils import Unpack from ...utils import ( ModelOutput, + TransformersKwargs, auto_docstring, - filter_out_non_signature_kwargs, + can_return_tuple, is_vision_available, logging, torch_int, @@ -953,23 +955,20 @@ def __init__(self, config: OwlViTConfig): # Initialize weights and apply final processing self.post_init() - @filter_out_non_signature_kwargs() + @can_return_tuple @auto_docstring def get_text_features( self, input_ids: torch.Tensor, attention_mask: torch.Tensor | None = None, - ) -> torch.FloatTensor: + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: r""" input_ids (`torch.LongTensor` of shape `(batch_size * num_max_text_queries, sequence_length)`): Indices of input sequence tokens in the vocabulary. Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and [`PreTrainedTokenizer.__call__`] for details. [What are input IDs?](../glossary#input-ids) - Returns: - text_features (`torch.FloatTensor` of shape `(batch_size, output_dim`): The text embeddings obtained by - applying the projection layer to the pooled output of [`OwlViTTextModel`]. - Examples: ```python >>> import torch @@ -984,23 +983,26 @@ def get_text_features( ... text_features = model.get_text_features(**inputs) ```""" # Get embeddings for all text queries in all batch samples - text_outputs: BaseModelOutputWithPooling = self.text_model(input_ids=input_ids, attention_mask=attention_mask) - text_features = self.text_projection(text_outputs.pooler_output) + text_outputs: BaseModelOutputWithPooling = self.text_model( + input_ids=input_ids, + attention_mask=attention_mask, + return_dict=True, + **kwargs, + ) + pooled_output = text_outputs.pooler_output + text_outputs.pooler_output = self.text_projection(pooled_output) - return text_features + return text_outputs - @filter_out_non_signature_kwargs() + @can_return_tuple @auto_docstring def get_image_features( self, pixel_values: torch.Tensor, interpolate_pos_encoding: bool = False, - ) -> torch.FloatTensor: + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: r""" - Returns: - image_features (`torch.FloatTensor` of shape `(batch_size, output_dim`): The image embeddings obtained by - applying the projection layer to the pooled output of [`OwlViTVisionModel`]. - Examples: ```python >>> import torch @@ -1020,10 +1022,12 @@ def get_image_features( vision_outputs: BaseModelOutputWithPooling = self.vision_model( pixel_values=pixel_values, interpolate_pos_encoding=interpolate_pos_encoding, + return_dict=True, + **kwargs, ) - image_features = self.visual_projection(vision_outputs.pooler_output) + vision_outputs.pooler_output = self.visual_projection(vision_outputs.pooler_output) - return image_features + return vision_outputs @auto_docstring def forward( diff --git a/src/transformers/models/paddleocr_vl/modeling_paddleocr_vl.py b/src/transformers/models/paddleocr_vl/modeling_paddleocr_vl.py index 0a3878422458..f813ac0c10ea 100644 --- a/src/transformers/models/paddleocr_vl/modeling_paddleocr_vl.py +++ b/src/transformers/models/paddleocr_vl/modeling_paddleocr_vl.py @@ -544,42 +544,6 @@ def forward( ) -class PaddleOCRVisionModel(PaddleOCRVLPreTrainedModel): - config: PaddleOCRVisionConfig - main_input_name = "pixel_values" - input_modalities = "image" - - def __init__(self, config: PaddleOCRVisionConfig): - super().__init__(config) - - self.vision_model = PaddleOCRVisionTransformer(config) - - # Initialize weights and apply final processing - self.post_init() - - def forward( - self, - pixel_values: torch.FloatTensor, - cu_seqlens: torch.Tensor, - image_grid_thw: list[tuple[int, int, int] | list[tuple[int, int, int]]] | None = None, - **kwargs, - ) -> BaseModelOutputWithPooling: - """ - Args: - pixel_values (`torch.FloatTensor` of shape `(batch_size, sequence_length, image_channels, patch_size, patch_size)`): - The tensors corresponding to the input images. - cu_seqlens (`torch.Tensor` of shape `(num_images + 1,)`): - The cumulative sequence lengths of each image or video feature. - image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): - The temporal, height and width of feature shape of each image in LLM. - """ - return self.vision_model( - pixel_values=pixel_values, - cu_seqlens=cu_seqlens, - image_grid_thw=image_grid_thw, - ) - - class PaddleOCRVisionEmbeddings(nn.Module): def __init__(self, config: PaddleOCRVisionConfig): super().__init__() @@ -970,6 +934,47 @@ def forward( ) +class PaddleOCRVisionModel(PaddleOCRVLPreTrainedModel): + config: PaddleOCRVisionConfig + main_input_name = "pixel_values" + input_modalities = "image" + _can_record_outputs = { + "hidden_states": PaddleOCRVisionEncoderLayer, + "attentions": PaddleOCRVisionAttention, + } + + def __init__(self, config: PaddleOCRVisionConfig): + super().__init__(config) + + self.vision_model = PaddleOCRVisionTransformer(config) + + # Initialize weights and apply final processing + self.post_init() + + @check_model_inputs(tie_last_hidden_states=False) + def forward( + self, + pixel_values: torch.FloatTensor, + cu_seqlens: torch.Tensor, + image_grid_thw: list[tuple[int, int, int] | list[tuple[int, int, int]]] | None = None, + **kwargs, + ) -> tuple | BaseModelOutputWithPooling: + """ + Args: + pixel_values (`torch.FloatTensor` of shape `(batch_size, sequence_length, image_channels, patch_size, patch_size)`): + The tensors corresponding to the input images. + cu_seqlens (`torch.Tensor` of shape `(num_images + 1,)`): + The cumulative sequence lengths of each image or video feature. + image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): + The temporal, height and width of feature shape of each image in LLM. + """ + return self.vision_model( + pixel_values=pixel_values, + cu_seqlens=cu_seqlens, + image_grid_thw=image_grid_thw, + ) + + @dataclass @auto_docstring( custom_intro=""" @@ -1196,33 +1201,19 @@ def get_rope_index( return position_ids, mrope_position_deltas - def get_video_features( - self, pixel_values_videos: torch.FloatTensor, video_grid_thw: torch.LongTensor | None = None - ): - """ - Encodes videos into continuous embeddings that can be forwarded to the language model. - - Args: - pixel_values_videos (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): - The tensors corresponding to the input videos. - video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*): - The temporal, height and width of feature shape of each video in LLM. - """ - pixel_values_videos = pixel_values_videos.type(self.visual.dtype) - video_embeds = self.visual(pixel_values_videos, grid_thw=video_grid_thw) - split_sizes = (video_grid_thw.prod(-1) // self.visual.spatial_merge_size**2).tolist() - video_embeds = torch.split(video_embeds, split_sizes) - return video_embeds - - def get_image_features(self, pixel_values: torch.FloatTensor, image_grid_thw: torch.LongTensor | None = None): - """ - Encodes images into continuous embeddings that can be forwarded to the language model. - - Args: - pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): - The tensors corresponding to the input images. - image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): - The temporal, height and width of feature shape of each image in LLM. + @can_return_tuple + @auto_docstring + def get_image_features( + self, + pixel_values: torch.FloatTensor, + image_grid_thw: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): + The tensors corresponding to the input images. + image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): + The temporal, height and width of feature shape of each image in LLM. """ pixel_values = pixel_values.type(self.visual.dtype).unsqueeze(0) cu_seqlens = torch.repeat_interleave(image_grid_thw[:, 1] * image_grid_thw[:, 2], image_grid_thw[:, 0]).cumsum( @@ -1238,10 +1229,14 @@ def get_image_features(self, pixel_values: torch.FloatTensor, image_grid_thw: to pixel_values=pixel_values, image_grid_thw=image_grid_thw, cu_seqlens=cu_seqlens, + return_dict=True, + **kwargs, ) image_embeds = vision_outputs.last_hidden_state image_embeds = self.projector(image_embeds, image_grid_thw) - return image_embeds + vision_outputs.pooler_output = image_embeds + + return vision_outputs def get_placeholder_mask( self, input_ids: torch.LongTensor, inputs_embeds: torch.FloatTensor, image_features: torch.FloatTensor @@ -1292,9 +1287,8 @@ def forward( inputs_embeds = self.language_model.embed_tokens(input_ids) if pixel_values is not None: - image_embeds = self.get_image_features(pixel_values, image_grid_thw).to( - inputs_embeds.device, inputs_embeds.dtype - ) + image_embeds = self.get_image_features(pixel_values, image_grid_thw, return_dict=True).pooler_output + image_embeds = image_embeds.to(inputs_embeds.device, inputs_embeds.dtype) image_mask = self.get_placeholder_mask(input_ids, inputs_embeds=inputs_embeds, image_features=image_embeds) inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds) @@ -1360,13 +1354,20 @@ def get_input_embeddings(self): def set_input_embeddings(self, value): self.model.set_input_embeddings(value) - def get_video_features( - self, pixel_values_videos: torch.FloatTensor, video_grid_thw: torch.LongTensor | None = None - ): - return self.model.get_video_features(pixel_values_videos, video_grid_thw) - - def get_image_features(self, pixel_values: torch.FloatTensor, image_grid_thw: torch.LongTensor | None = None): - return self.model.get_image_features(pixel_values, image_grid_thw) + @auto_docstring + def get_image_features( + self, + pixel_values: torch.FloatTensor, + image_grid_thw: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): + The tensors corresponding to the input images. + image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): + The temporal, height and width of feature shape of each image in LLM. + """ + return self.model.get_image_features(pixel_values=pixel_values, image_grid_thw=image_grid_thw, **kwargs) @can_return_tuple @auto_docstring diff --git a/src/transformers/models/paddleocr_vl/modular_paddleocr_vl.py b/src/transformers/models/paddleocr_vl/modular_paddleocr_vl.py index 23e5cc9e556d..6a320b66440d 100644 --- a/src/transformers/models/paddleocr_vl/modular_paddleocr_vl.py +++ b/src/transformers/models/paddleocr_vl/modular_paddleocr_vl.py @@ -863,42 +863,6 @@ def forward( ) -class PaddleOCRVisionModel(PaddleOCRVLPreTrainedModel): - config: PaddleOCRVisionConfig - main_input_name = "pixel_values" - input_modalities = "image" - - def __init__(self, config: PaddleOCRVisionConfig): - super().__init__(config) - - self.vision_model = PaddleOCRVisionTransformer(config) - - # Initialize weights and apply final processing - self.post_init() - - def forward( - self, - pixel_values: torch.FloatTensor, - cu_seqlens: torch.Tensor, - image_grid_thw: list[tuple[int, int, int] | list[tuple[int, int, int]]] | None = None, - **kwargs, - ) -> BaseModelOutputWithPooling: - """ - Args: - pixel_values (`torch.FloatTensor` of shape `(batch_size, sequence_length, image_channels, patch_size, patch_size)`): - The tensors corresponding to the input images. - cu_seqlens (`torch.Tensor` of shape `(num_images + 1,)`): - The cumulative sequence lengths of each image or video feature. - image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): - The temporal, height and width of feature shape of each image in LLM. - """ - return self.vision_model( - pixel_values=pixel_values, - cu_seqlens=cu_seqlens, - image_grid_thw=image_grid_thw, - ) - - class PaddleOCRVisionEmbeddings(SiglipVisionEmbeddings): def __init__(self, config: PaddleOCRVisionConfig): super().__init__() @@ -1089,6 +1053,47 @@ def forward( ) +class PaddleOCRVisionModel(PaddleOCRVLPreTrainedModel): + config: PaddleOCRVisionConfig + main_input_name = "pixel_values" + input_modalities = "image" + _can_record_outputs = { + "hidden_states": PaddleOCRVisionEncoderLayer, + "attentions": PaddleOCRVisionAttention, + } + + def __init__(self, config: PaddleOCRVisionConfig): + super().__init__(config) + + self.vision_model = PaddleOCRVisionTransformer(config) + + # Initialize weights and apply final processing + self.post_init() + + @check_model_inputs(tie_last_hidden_states=False) + def forward( + self, + pixel_values: torch.FloatTensor, + cu_seqlens: torch.Tensor, + image_grid_thw: list[tuple[int, int, int] | list[tuple[int, int, int]]] | None = None, + **kwargs, + ) -> tuple | BaseModelOutputWithPooling: + """ + Args: + pixel_values (`torch.FloatTensor` of shape `(batch_size, sequence_length, image_channels, patch_size, patch_size)`): + The tensors corresponding to the input images. + cu_seqlens (`torch.Tensor` of shape `(num_images + 1,)`): + The cumulative sequence lengths of each image or video feature. + image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): + The temporal, height and width of feature shape of each image in LLM. + """ + return self.vision_model( + pixel_values=pixel_values, + cu_seqlens=cu_seqlens, + image_grid_thw=image_grid_thw, + ) + + class PaddleOCRVLModelOutputWithPast(Qwen2VLModelOutputWithPast): pass @@ -1116,15 +1121,22 @@ def get_input_embeddings(self): def set_input_embeddings(self, value): self.language_model.embed_tokens = value - def get_image_features(self, pixel_values: torch.FloatTensor, image_grid_thw: torch.LongTensor | None = None): - """ - Encodes images into continuous embeddings that can be forwarded to the language model. + def get_video_features(self): + raise AttributeError("PaddleOCRVLModel does not support video.") - Args: - pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): - The tensors corresponding to the input images. - image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): - The temporal, height and width of feature shape of each image in LLM. + @can_return_tuple + @auto_docstring + def get_image_features( + self, + pixel_values: torch.FloatTensor, + image_grid_thw: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): + The tensors corresponding to the input images. + image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): + The temporal, height and width of feature shape of each image in LLM. """ pixel_values = pixel_values.type(self.visual.dtype).unsqueeze(0) cu_seqlens = torch.repeat_interleave(image_grid_thw[:, 1] * image_grid_thw[:, 2], image_grid_thw[:, 0]).cumsum( @@ -1140,10 +1152,14 @@ def get_image_features(self, pixel_values: torch.FloatTensor, image_grid_thw: to pixel_values=pixel_values, image_grid_thw=image_grid_thw, cu_seqlens=cu_seqlens, + return_dict=True, + **kwargs, ) image_embeds = vision_outputs.last_hidden_state image_embeds = self.projector(image_embeds, image_grid_thw) - return image_embeds + vision_outputs.pooler_output = image_embeds + + return vision_outputs def get_placeholder_mask( self, input_ids: torch.LongTensor, inputs_embeds: torch.FloatTensor, image_features: torch.FloatTensor @@ -1194,9 +1210,8 @@ def forward( inputs_embeds = self.language_model.embed_tokens(input_ids) if pixel_values is not None: - image_embeds = self.get_image_features(pixel_values, image_grid_thw).to( - inputs_embeds.device, inputs_embeds.dtype - ) + image_embeds = self.get_image_features(pixel_values, image_grid_thw, return_dict=True).pooler_output + image_embeds = image_embeds.to(inputs_embeds.device, inputs_embeds.dtype) image_mask = self.get_placeholder_mask(input_ids, inputs_embeds=inputs_embeds, image_features=image_embeds) inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds) @@ -1248,6 +1263,9 @@ class PaddleOCRVLForConditionalGeneration(Qwen2VLForConditionalGeneration): } _keys_to_ignore_on_load_unexpected = ["packing_position_embedding", "vision_model.head"] + def get_video_features(self): + raise AttributeError("PaddleOCRVLForConditionalGeneration does not support video.") + @can_return_tuple @auto_docstring def forward( diff --git a/src/transformers/models/paligemma/modeling_paligemma.py b/src/transformers/models/paligemma/modeling_paligemma.py index 269a53065903..f5998f0607b6 100644 --- a/src/transformers/models/paligemma/modeling_paligemma.py +++ b/src/transformers/models/paligemma/modeling_paligemma.py @@ -24,7 +24,7 @@ from ...generation import GenerationMixin from ...masking_utils import create_masks_for_generate from ...modeling_flash_attention_utils import FlashAttentionKwargs -from ...modeling_outputs import BaseModelOutputWithPast +from ...modeling_outputs import BaseModelOutputWithPast, BaseModelOutputWithPooling from ...modeling_utils import PreTrainedModel from ...processing_utils import Unpack from ...utils import ( @@ -258,21 +258,20 @@ def get_input_embeddings(self): def set_input_embeddings(self, value): self.language_model.set_input_embeddings(value) - def get_image_features(self, pixel_values: torch.FloatTensor): - """ - Obtains image last hidden states from the vision tower and apply multimodal projection. - - Args: - pixel_values (`torch.FloatTensor]` of shape `(batch_size, channels, height, width)`) - The tensors corresponding to the input images. - Returns: - image_features (`torch.Tensor`): Image feature tensor of shape `(num_images, image_length, embed_dim)`). - """ - image_outputs = self.vision_tower(pixel_values) + @can_return_tuple + @auto_docstring( + custom_intro="Obtains image last hidden states from the vision tower and apply multimodal projection." + ) + def get_image_features( + self, pixel_values: torch.FloatTensor, **kwargs: Unpack[TransformersKwargs] + ) -> tuple | BaseModelOutputWithPooling: + image_outputs = self.vision_tower(pixel_values, return_dict=True, **kwargs) selected_image_feature = image_outputs.last_hidden_state image_features = self.multi_modal_projector(selected_image_feature) image_features = image_features / (self.config.text_config.hidden_size**0.5) - return image_features + image_outputs.pooler_output = image_features + + return image_outputs def get_placeholder_mask( self, input_ids: torch.LongTensor, inputs_embeds: torch.FloatTensor, image_features: torch.FloatTensor @@ -376,7 +375,7 @@ def forward( # Merge text and images if pixel_values is not None: - image_features = self.get_image_features(pixel_values) + image_features = self.get_image_features(pixel_values, return_dict=True).pooler_output image_features = image_features.to(inputs_embeds.device, inputs_embeds.dtype) special_image_mask = self.get_placeholder_mask( input_ids, inputs_embeds=inputs_embeds, image_features=image_features @@ -445,8 +444,9 @@ def get_input_embeddings(self): def set_input_embeddings(self, value): self.model.set_input_embeddings(value) - def get_image_features(self, pixel_values): - return self.model.get_image_features(pixel_values) + @auto_docstring + def get_image_features(self, pixel_values: torch.FloatTensor, **kwargs: Unpack[TransformersKwargs]): + return self.model.get_image_features(pixel_values, **kwargs) @can_return_tuple @auto_docstring diff --git a/src/transformers/models/pe_audio/modeling_pe_audio.py b/src/transformers/models/pe_audio/modeling_pe_audio.py index 0daf26ea1649..948cd6e1fd16 100644 --- a/src/transformers/models/pe_audio/modeling_pe_audio.py +++ b/src/transformers/models/pe_audio/modeling_pe_audio.py @@ -642,7 +642,7 @@ def forward( input_values: torch.Tensor, padding_mask: torch.Tensor | None = None, **kwargs, - ) -> BaseModelOutputWithPooling: + ) -> tuple | BaseModelOutputWithPooling: inputs_embeds, padding_mask = self.embedder(input_values, padding_mask=padding_mask) inputs_embeds, attention_mask = self.patch_embedder(inputs_embeds, padding_mask=padding_mask) diff --git a/src/transformers/models/pe_audio/modular_pe_audio.py b/src/transformers/models/pe_audio/modular_pe_audio.py index b5e594393344..84a1ad80df85 100644 --- a/src/transformers/models/pe_audio/modular_pe_audio.py +++ b/src/transformers/models/pe_audio/modular_pe_audio.py @@ -117,7 +117,7 @@ def forward( input_values: torch.Tensor, padding_mask: torch.Tensor | None = None, **kwargs, - ) -> BaseModelOutputWithPooling: + ) -> tuple | BaseModelOutputWithPooling: inputs_embeds, padding_mask = self.embedder(input_values, padding_mask=padding_mask) inputs_embeds, attention_mask = self.patch_embedder(inputs_embeds, padding_mask=padding_mask) diff --git a/src/transformers/models/pe_audio_video/modeling_pe_audio_video.py b/src/transformers/models/pe_audio_video/modeling_pe_audio_video.py index bf47a4ea519d..0fb693d67941 100644 --- a/src/transformers/models/pe_audio_video/modeling_pe_audio_video.py +++ b/src/transformers/models/pe_audio_video/modeling_pe_audio_video.py @@ -589,7 +589,7 @@ def forward( padding_mask: torch.Tensor | None = None, padding_mask_videos: torch.Tensor | None = None, **kwargs, - ) -> PeAudioVideoEncoderOutput: + ) -> tuple | PeAudioVideoEncoderOutput: inputs_embeds, padding_mask, audio_output, video_output = self.embedder( input_values, pixel_values_videos, diff --git a/src/transformers/models/pe_audio_video/modular_pe_audio_video.py b/src/transformers/models/pe_audio_video/modular_pe_audio_video.py index 7127de028c6a..78bd0a044259 100644 --- a/src/transformers/models/pe_audio_video/modular_pe_audio_video.py +++ b/src/transformers/models/pe_audio_video/modular_pe_audio_video.py @@ -378,7 +378,7 @@ def forward( padding_mask: torch.Tensor | None = None, padding_mask_videos: torch.Tensor | None = None, **kwargs, - ) -> PeAudioVideoEncoderOutput: + ) -> tuple | PeAudioVideoEncoderOutput: inputs_embeds, padding_mask, audio_output, video_output = self.embedder( input_values, pixel_values_videos, diff --git a/src/transformers/models/pe_video/modeling_pe_video.py b/src/transformers/models/pe_video/modeling_pe_video.py index 6eb3392ded5f..a94e53b77dc4 100644 --- a/src/transformers/models/pe_video/modeling_pe_video.py +++ b/src/transformers/models/pe_video/modeling_pe_video.py @@ -526,7 +526,7 @@ def forward( pixel_values_videos: torch.Tensor, padding_mask_videos: torch.Tensor | None = None, **kwargs, - ) -> BaseModelOutputWithPooling: + ) -> tuple | BaseModelOutputWithPooling: inputs_embeds, padding_mask = self.embedder(pixel_values_videos, padding_mask=padding_mask_videos) inputs_embeds, attention_mask = self.patch_embedder(inputs_embeds, padding_mask=padding_mask) @@ -570,27 +570,39 @@ def __init__(self, config: PeVideoConfig): self.post_init() - def get_text_features(self, input_ids, attention_mask=None): - # TODO: should it be named feature or embeds - text_outputs: MaskedLMOutput = self.text_model( - input_ids=input_ids, - attention_mask=attention_mask, - return_dict=True, - ) - - text_features = text_outputs.last_hidden_state - text_features = self.text_video_head(text_features) - return text_features + @can_return_tuple + @auto_docstring + def get_text_features( + self, + input_ids: torch.Tensor, + attention_mask: torch.Tensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + text_outputs: BaseModelOutputWithPooling = self.text_model( + input_ids=input_ids, + attention_mask=attention_mask, + return_dict=True, + **kwargs, + ) + text_outputs.pooler_output = self.text_video_head(text_outputs.last_hidden_state) + return text_outputs - def get_video_features(self, pixel_values_videos, padding_mask_videos=None): - # TODO: should it be named feature or embeds - video_outputs: BaseModelOutputWithPooling = self.video_encoder( - pixel_values_videos=pixel_values_videos, - padding_mask_videos=padding_mask_videos, - return_dict=True, - ) - video_features = self.video_head(video_outputs.pooler_output) - return video_features + @can_return_tuple + @auto_docstring + def get_video_features( + self, + pixel_values_videos: torch.Tensor, + padding_mask_videos: torch.Tensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + video_outputs: BaseModelOutputWithPooling = self.video_encoder( + pixel_values_videos=pixel_values_videos, + padding_mask_videos=padding_mask_videos, + return_dict=True, + **kwargs, + ) + video_outputs.pooler_output = self.video_head(video_outputs.pooler_output) + return video_outputs @can_return_tuple def forward( diff --git a/src/transformers/models/pe_video/modular_pe_video.py b/src/transformers/models/pe_video/modular_pe_video.py index 687701a0664a..38576598cf44 100644 --- a/src/transformers/models/pe_video/modular_pe_video.py +++ b/src/transformers/models/pe_video/modular_pe_video.py @@ -20,7 +20,8 @@ from ...modeling_attn_mask_utils import _prepare_4d_attention_mask from ...modeling_outputs import BaseModelOutputWithPooling, MaskedLMOutput -from ...utils import ModelOutput, auto_docstring, can_return_tuple +from ...processing_utils import Unpack +from ...utils import ModelOutput, TransformersKwargs, auto_docstring, can_return_tuple from ...utils.generic import check_model_inputs from ..auto import AutoModel, AutoModelForImageClassification from ..pe_audio_video.modeling_pe_audio_video import ( @@ -106,7 +107,7 @@ def forward( pixel_values_videos: torch.Tensor, padding_mask_videos: torch.Tensor | None = None, **kwargs, - ) -> BaseModelOutputWithPooling: + ) -> tuple | BaseModelOutputWithPooling: inputs_embeds, padding_mask = self.embedder(pixel_values_videos, padding_mask=padding_mask_videos) inputs_embeds, attention_mask = self.patch_embedder(inputs_embeds, padding_mask=padding_mask) @@ -150,27 +151,39 @@ def __init__(self, config: PeVideoConfig): self.post_init() - def get_text_features(self, input_ids, attention_mask=None): - # TODO: should it be named feature or embeds - text_outputs: MaskedLMOutput = self.text_model( - input_ids=input_ids, - attention_mask=attention_mask, - return_dict=True, - ) - - text_features = text_outputs.last_hidden_state - text_features = self.text_video_head(text_features) - return text_features - - def get_video_features(self, pixel_values_videos, padding_mask_videos=None): - # TODO: should it be named feature or embeds - video_outputs: BaseModelOutputWithPooling = self.video_encoder( - pixel_values_videos=pixel_values_videos, - padding_mask_videos=padding_mask_videos, - return_dict=True, - ) - video_features = self.video_head(video_outputs.pooler_output) - return video_features + @can_return_tuple + @auto_docstring + def get_text_features( + self, + input_ids: torch.Tensor, + attention_mask: torch.Tensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + text_outputs: BaseModelOutputWithPooling = self.text_model( + input_ids=input_ids, + attention_mask=attention_mask, + return_dict=True, + **kwargs, + ) + text_outputs.pooler_output = self.text_video_head(text_outputs.last_hidden_state) + return text_outputs + + @can_return_tuple + @auto_docstring + def get_video_features( + self, + pixel_values_videos: torch.Tensor, + padding_mask_videos: torch.Tensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + video_outputs: BaseModelOutputWithPooling = self.video_encoder( + pixel_values_videos=pixel_values_videos, + padding_mask_videos=padding_mask_videos, + return_dict=True, + **kwargs, + ) + video_outputs.pooler_output = self.video_head(video_outputs.pooler_output) + return video_outputs @can_return_tuple def forward( diff --git a/src/transformers/models/perception_lm/modeling_perception_lm.py b/src/transformers/models/perception_lm/modeling_perception_lm.py index 0f2d4bd9d99b..7a202554a6eb 100644 --- a/src/transformers/models/perception_lm/modeling_perception_lm.py +++ b/src/transformers/models/perception_lm/modeling_perception_lm.py @@ -26,9 +26,10 @@ from ...cache_utils import Cache from ...generation import GenerationMixin -from ...modeling_outputs import BaseModelOutputWithPast, ModelOutput +from ...modeling_outputs import BaseModelOutputWithPast, BaseModelOutputWithPooling, ModelOutput from ...modeling_utils import PreTrainedModel -from ...utils import auto_docstring, can_return_tuple, torch_compilable_check +from ...processing_utils import Unpack +from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, torch_compilable_check from ..auto import AutoModel from .configuration_perception_lm import PerceptionLMConfig @@ -178,26 +179,23 @@ def get_input_embeddings(self): def set_input_embeddings(self, value): self.language_model.set_input_embeddings(value) + @can_return_tuple + @auto_docstring( + custom_intro="Obtains image last hidden states from the vision tower and apply multimodal projection." + ) def get_image_features( self, pixel_values: torch.FloatTensor, - **kwargs, - ): - """ - Obtains image last hidden states from the vision tower and apply multimodal projection. - - Args: - pixel_values (`torch.FloatTensor]` of shape `(batch_size, num_tiles, channels, height, width)`) - The tensors corresponding to the input images. - Returns: - image_features (`torch.Tensor`): Image feature tensor of shape `(num_tiles, num_patches, embed_dim)`). - """ - image_outputs = self.vision_tower(pixel_values.flatten(0, 1)) - image_outputs = image_outputs.last_hidden_state + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + image_outputs = self.vision_tower(pixel_values.flatten(0, 1), return_dict=True, **kwargs) + last_hidden_state = image_outputs.last_hidden_state if self.config.vision_use_cls_token: - image_outputs = image_outputs[:, 1:, :] - image_features = self.multi_modal_projector(image_outputs) - return image_features + last_hidden_state = last_hidden_state[:, 1:, :] + image_features = self.multi_modal_projector(last_hidden_state) + image_outputs.pooler_output = image_features + + return image_outputs def get_placeholder_mask( self, @@ -274,7 +272,7 @@ def forward( image_features = None if pixel_values is not None: - image_features = self.get_image_features(pixel_values=pixel_values) + image_features = self.get_image_features(pixel_values=pixel_values, return_dict=True).pooler_output image_features = image_features.to(inputs_embeds.device, dtype=inputs_embeds.dtype) special_image_mask, _ = self.get_placeholder_mask( input_ids, inputs_embeds=inputs_embeds, image_features=image_features @@ -283,7 +281,7 @@ def forward( video_features = None if pixel_values_videos is not None: - video_features = self.get_image_features(pixel_values=pixel_values_videos) + video_features = self.get_image_features(pixel_values=pixel_values_videos, return_dict=True).pooler_output video_features = video_features.to(inputs_embeds.device, dtype=inputs_embeds.dtype) _, special_video_mask = self.get_placeholder_mask( input_ids, inputs_embeds=inputs_embeds, video_features=video_features diff --git a/src/transformers/models/perception_lm/modular_perception_lm.py b/src/transformers/models/perception_lm/modular_perception_lm.py index 7a3d383e2711..ccd335a84ae6 100644 --- a/src/transformers/models/perception_lm/modular_perception_lm.py +++ b/src/transformers/models/perception_lm/modular_perception_lm.py @@ -19,7 +19,15 @@ from torch import nn from ...cache_utils import Cache -from ...utils import auto_docstring, can_return_tuple, logging, torch_compilable_check +from ...modeling_outputs import BaseModelOutputWithPooling +from ...processing_utils import Unpack +from ...utils import ( + TransformersKwargs, + auto_docstring, + can_return_tuple, + logging, + torch_compilable_check, +) from ..auto import AutoModel from ..llava.modeling_llava import ( LlavaCausalLMOutputWithPast, @@ -139,26 +147,23 @@ def __init__(self, config: PerceptionLMConfig): self.multi_modal_projector = PerceptionLMMultiModalProjector(config) self.language_model = AutoModel.from_config(config.text_config) + @can_return_tuple + @auto_docstring( + custom_intro="Obtains image last hidden states from the vision tower and apply multimodal projection." + ) def get_image_features( self, pixel_values: torch.FloatTensor, - **kwargs, - ): - """ - Obtains image last hidden states from the vision tower and apply multimodal projection. - - Args: - pixel_values (`torch.FloatTensor]` of shape `(batch_size, num_tiles, channels, height, width)`) - The tensors corresponding to the input images. - Returns: - image_features (`torch.Tensor`): Image feature tensor of shape `(num_tiles, num_patches, embed_dim)`). - """ - image_outputs = self.vision_tower(pixel_values.flatten(0, 1)) - image_outputs = image_outputs.last_hidden_state + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + image_outputs = self.vision_tower(pixel_values.flatten(0, 1), return_dict=True, **kwargs) + last_hidden_state = image_outputs.last_hidden_state if self.config.vision_use_cls_token: - image_outputs = image_outputs[:, 1:, :] - image_features = self.multi_modal_projector(image_outputs) - return image_features + last_hidden_state = last_hidden_state[:, 1:, :] + image_features = self.multi_modal_projector(last_hidden_state) + image_outputs.pooler_output = image_features + + return image_outputs def get_placeholder_mask( self, @@ -235,7 +240,7 @@ def forward( image_features = None if pixel_values is not None: - image_features = self.get_image_features(pixel_values=pixel_values) + image_features = self.get_image_features(pixel_values=pixel_values, return_dict=True).pooler_output image_features = image_features.to(inputs_embeds.device, dtype=inputs_embeds.dtype) special_image_mask, _ = self.get_placeholder_mask( input_ids, inputs_embeds=inputs_embeds, image_features=image_features @@ -244,7 +249,7 @@ def forward( video_features = None if pixel_values_videos is not None: - video_features = self.get_image_features(pixel_values=pixel_values_videos) + video_features = self.get_image_features(pixel_values=pixel_values_videos, return_dict=True).pooler_output video_features = video_features.to(inputs_embeds.device, dtype=inputs_embeds.dtype) _, special_video_mask = self.get_placeholder_mask( input_ids, inputs_embeds=inputs_embeds, video_features=video_features diff --git a/src/transformers/models/phi4_multimodal/modeling_phi4_multimodal.py b/src/transformers/models/phi4_multimodal/modeling_phi4_multimodal.py index 5e4aa50bb47d..38995c5f167d 100644 --- a/src/transformers/models/phi4_multimodal/modeling_phi4_multimodal.py +++ b/src/transformers/models/phi4_multimodal/modeling_phi4_multimodal.py @@ -447,7 +447,7 @@ def forward( pixel_values, patch_attention_mask: torch.BoolTensor | None = None, **kwargs: Unpack[TransformersKwargs], - ) -> BaseModelOutputWithPooling: + ) -> tuple | BaseModelOutputWithPooling: batch_size = pixel_values.size(0) if patch_attention_mask is None: patch_attention_mask = torch.ones( @@ -1568,7 +1568,7 @@ def forward( output_hidden_states: bool | None = None, cache_position: torch.LongTensor | None = None, **kwargs, - ) -> BaseModelOutputWithPast: + ) -> tuple | BaseModelOutputWithPast: r""" image_pixel_values (`torch.FloatTensor`, *optional*): If the input contains images, these correspond to the pixel values after transformations (as returned by diff --git a/src/transformers/models/phi4_multimodal/modular_phi4_multimodal.py b/src/transformers/models/phi4_multimodal/modular_phi4_multimodal.py index ccd202c62785..2043cd0077e9 100644 --- a/src/transformers/models/phi4_multimodal/modular_phi4_multimodal.py +++ b/src/transformers/models/phi4_multimodal/modular_phi4_multimodal.py @@ -688,7 +688,7 @@ def forward( pixel_values, patch_attention_mask: torch.BoolTensor | None = None, **kwargs: Unpack[TransformersKwargs], - ) -> BaseModelOutputWithPooling: + ) -> tuple | BaseModelOutputWithPooling: batch_size = pixel_values.size(0) if patch_attention_mask is None: patch_attention_mask = torch.ones( @@ -1505,7 +1505,7 @@ def forward( output_hidden_states: bool | None = None, cache_position: torch.LongTensor | None = None, **kwargs, - ) -> BaseModelOutputWithPast: + ) -> tuple | BaseModelOutputWithPast: r""" image_pixel_values (`torch.FloatTensor`, *optional*): If the input contains images, these correspond to the pixel values after transformations (as returned by diff --git a/src/transformers/models/qwen2_5_omni/modeling_qwen2_5_omni.py b/src/transformers/models/qwen2_5_omni/modeling_qwen2_5_omni.py index d9c064952d65..24633e9ae0bc 100644 --- a/src/transformers/models/qwen2_5_omni/modeling_qwen2_5_omni.py +++ b/src/transformers/models/qwen2_5_omni/modeling_qwen2_5_omni.py @@ -37,13 +37,20 @@ from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer -from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPast, ModelOutput +from ...modeling_outputs import BaseModelOutputWithPast, BaseModelOutputWithPooling, ModelOutput from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel from ...processing_utils import Unpack -from ...utils import TransformersKwargs, auto_docstring, check_torch_load_is_safe, logging, torch_compilable_check +from ...utils import ( + TransformersKwargs, + auto_docstring, + can_return_tuple, + check_torch_load_is_safe, + logging, + torch_compilable_check, +) from ...utils.deprecation import deprecate_kwarg -from ...utils.generic import is_flash_attention_requested, maybe_autocast +from ...utils.generic import check_model_inputs, is_flash_attention_requested, maybe_autocast from ...utils.hub import cached_file from ..qwen2.modeling_qwen2 import Qwen2RMSNorm from .configuration_qwen2_5_omni import ( @@ -779,6 +786,10 @@ class Qwen2_5OmniAudioEncoder(Qwen2_5OmniPreTrainedModel): input_modalities = "audio" _no_split_modules = ["Qwen2_5OmniAudioEncoderLayer"] _supports_sdpa = True + _can_record_outputs = { + "hidden_states": Qwen2_5OmniAudioEncoderLayer, + "attentions": Qwen2_5OmniAudioAttention, + } def __init__(self, config: Qwen2_5OmniAudioEncoderConfig): super().__init__(config) @@ -831,14 +842,9 @@ def _prepare_attention_mask(self, inputs_tensor: torch.Tensor, cu_seqlens: torch attention_mask[..., cu_seqlens[i - 1] : cu_seqlens[i], cu_seqlens[i - 1] : cu_seqlens[i]] = 0 return attention_mask + @check_model_inputs(tie_last_hidden_states=False) @auto_docstring - def forward( - self, - input_features, - feature_lens=None, - aftercnn_lens=None, - **kwargs, - ): + def forward(self, input_features, feature_lens=None, aftercnn_lens=None, **kwargs: Unpack[TransformersKwargs]): r""" feature_lens (`torch.LongTensor` of shape `(batch_size,)`): mel length @@ -888,7 +894,7 @@ def forward( each_audio_states = self.proj(each_audio_states) token_audio_list.append(each_audio_states) token_audio = torch.cat(token_audio_list, dim=0) - return BaseModelOutput(last_hidden_state=token_audio) + return BaseModelOutputWithPooling(last_hidden_state=token_audio) def padded_and_mask_function(self, tensor_list, tensor_len, padding_value=0, padding_side="right"): """ @@ -1142,6 +1148,10 @@ class Qwen2_5OmniVisionEncoder(Qwen2_5OmniPreTrainedModel): config: Qwen2_5OmniVisionEncoderConfig _no_split_modules = ["Qwen2_5OmniVisionBlock"] _input_embed_layer = "patch_embed" + _can_record_outputs = { + "hidden_states": Qwen2_5OmniVisionBlock, + "attentions": Qwen2_5OmniVisionAttention, + } input_modalities = ("image", "video") def __init__(self, config: Qwen2_5OmniVisionEncoderConfig, *inputs, **kwargs) -> None: @@ -1241,7 +1251,10 @@ def get_window_index(self, grid_thw): return window_index, cu_window_seqlens - def forward(self, hidden_states: torch.Tensor, grid_thw: torch.Tensor, **kwargs) -> torch.Tensor: + @check_model_inputs + def forward( + self, hidden_states: torch.Tensor, grid_thw: torch.Tensor, **kwargs: Unpack[TransformersKwargs] + ) -> tuple | BaseModelOutputWithPooling: """ Args: hidden_states (`torch.Tensor` of shape `(seq_len, hidden_size)`): @@ -1294,11 +1307,15 @@ def forward(self, hidden_states: torch.Tensor, grid_thw: torch.Tensor, **kwargs) rotary_pos_emb=rotary_pos_emb, **kwargs, ) - hidden_states = self.merger(hidden_states) + + merged_hidden_states = self.merger(hidden_states) reverse_indices = torch.argsort(window_index) - hidden_states = hidden_states[reverse_indices, :] + merged_hidden_states = merged_hidden_states[reverse_indices, :] - return hidden_states + return BaseModelOutputWithPooling( + last_hidden_state=hidden_states, + pooler_output=merged_hidden_states, + ) class Qwen2_5OmniRotaryEmbedding(nn.Module): @@ -1782,52 +1799,56 @@ def get_input_embeddings(self): def set_input_embeddings(self, value): self.model.set_input_embeddings(value) + @can_return_tuple + @auto_docstring def get_video_features( - self, pixel_values_videos: torch.FloatTensor, video_grid_thw: torch.LongTensor | None = None - ): - """ - Encodes videos into continuous embeddings that can be forwarded to the language model. - - Args: - pixel_values_videos (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): - The tensors corresponding to the input videos. - video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*): - The temporal, height and width of feature shape of each video in LLM. + self, + pixel_values_videos: torch.FloatTensor, + video_grid_thw: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + pixel_values_videos (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): + The tensors corresponding to the input videos. + video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*): + The temporal, height and width of feature shape of each video in LLM. """ pixel_values_videos = pixel_values_videos.type(self.visual.dtype) - video_embeds = self.visual(pixel_values_videos, grid_thw=video_grid_thw) - return video_embeds - - def get_image_features(self, pixel_values: torch.FloatTensor, image_grid_thw: torch.LongTensor | None = None): - """ - Encodes images into continuous embeddings that can be forwarded to the language model. + return self.visual(pixel_values_videos, grid_thw=video_grid_thw, **kwargs) - Args: - pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): - The tensors corresponding to the input images. - image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): - The temporal, height and width of feature shape of each image in LLM. + @can_return_tuple + @auto_docstring + def get_image_features( + self, + pixel_values: torch.FloatTensor, + image_grid_thw: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): + The tensors corresponding to the input images. + image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): + The temporal, height and width of feature shape of each image in LLM. """ pixel_values = pixel_values.type(self.visual.dtype) - image_embeds = self.visual(pixel_values, grid_thw=image_grid_thw) - return image_embeds + return self.visual(pixel_values, grid_thw=image_grid_thw, **kwargs) + @can_return_tuple + @auto_docstring def get_audio_features( self, input_features: torch.FloatTensor, feature_attention_mask: torch.LongTensor | None = None, audio_feature_lengths: torch.LongTensor | None = None, - ): - """ - Encodes audios into continuous embeddings that can be forwarded to the language model. - - Args: - input_features (`torch.FloatTensor`): - The tensors corresponding to the input audios. - feature_attention_mask (`torch.LongTensor`, *optional*): - Mask to avoid performing attention on padding feature indices. Mask values selected in `[0, 1]`: - audio_feature_lengths (`torch.LongTensor` of shape `(num_audios)`, *optional*): - The length of feature shape of each audio in LLM. + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + input_features (`torch.FloatTensor`): + The tensors corresponding to the input audios. + feature_attention_mask (`torch.LongTensor`, *optional*): + Mask to avoid performing attention on padding feature indices. Mask values selected in `[0, 1]`: + audio_feature_lengths (`torch.LongTensor` of shape `(num_audios)`, *optional*): + The length of feature shape of each audio in LLM. """ if feature_attention_mask is not None: audio_feature_lengths = torch.sum(feature_attention_mask, dim=1) @@ -1843,13 +1864,13 @@ def get_audio_features( input_features, feature_lens=feature_lens, aftercnn_lens=audio_feat_lengths, + return_dict=True, + **kwargs, ) - audio_features = audio_outputs.last_hidden_state - - if audio_features.shape[0] != sum(audio_output_lengths.tolist()): + if audio_outputs.last_hidden_state.shape[0] != sum(audio_output_lengths.tolist()): raise ValueError("length of audio_features should match audio_output_lengths") - return audio_features + return audio_outputs def get_placeholder_mask( self, @@ -1999,13 +2020,14 @@ def forward( input_features, feature_attention_mask=feature_attention_mask, audio_feature_lengths=audio_feature_lengths, - ) + return_dict=True, + ).last_hidden_state audio_features = audio_features.to(inputs_embeds.device, inputs_embeds.dtype) _, _, audio_mask = self.get_placeholder_mask(input_ids, inputs_embeds=inputs_embeds) inputs_embeds = inputs_embeds.masked_scatter(audio_mask, audio_features) if pixel_values is not None: - image_embeds = self.get_image_features(pixel_values, image_grid_thw) + image_embeds = self.get_image_features(pixel_values, image_grid_thw, return_dict=True).pooler_output image_embeds = image_embeds.to(inputs_embeds.device, inputs_embeds.dtype) image_mask, _, _ = self.get_placeholder_mask( input_ids, inputs_embeds=inputs_embeds, image_features=image_embeds @@ -2013,7 +2035,7 @@ def forward( inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds) if pixel_values_videos is not None: - video_embeds = self.get_video_features(pixel_values_videos, video_grid_thw) + video_embeds = self.get_video_features(pixel_values_videos, video_grid_thw, return_dict=True).pooler_output video_embeds = video_embeds.to(inputs_embeds.device, inputs_embeds.dtype) _, video_mask, _ = self.get_placeholder_mask( input_ids, inputs_embeds=inputs_embeds, video_features=video_embeds diff --git a/src/transformers/models/qwen2_5_omni/modular_qwen2_5_omni.py b/src/transformers/models/qwen2_5_omni/modular_qwen2_5_omni.py index 6cb77771653b..6ec82e156cda 100644 --- a/src/transformers/models/qwen2_5_omni/modular_qwen2_5_omni.py +++ b/src/transformers/models/qwen2_5_omni/modular_qwen2_5_omni.py @@ -29,13 +29,20 @@ from ...cache_utils import Cache from ...configuration_utils import PreTrainedConfig, layer_type_validation from ...generation import GenerationMixin -from ...modeling_outputs import BaseModelOutput, ModelOutput +from ...modeling_outputs import BaseModelOutputWithPooling, ModelOutput from ...modeling_rope_utils import RopeParameters from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel from ...processing_utils import Unpack -from ...utils import TransformersKwargs, auto_docstring, check_torch_load_is_safe, logging, torch_compilable_check +from ...utils import ( + TransformersKwargs, + auto_docstring, + can_return_tuple, + check_torch_load_is_safe, + logging, + torch_compilable_check, +) from ...utils.deprecation import deprecate_kwarg -from ...utils.generic import is_flash_attention_requested +from ...utils.generic import check_model_inputs, is_flash_attention_requested from ...utils.hub import cached_file from ..llama.modeling_llama import LlamaRotaryEmbedding, rotate_half from ..qwen2_5_vl.configuration_qwen2_5_vl import Qwen2_5_VLVisionConfig @@ -1660,6 +1667,10 @@ class Qwen2_5OmniAudioEncoder(Qwen2_5OmniPreTrainedModel): input_modalities = "audio" _no_split_modules = ["Qwen2_5OmniAudioEncoderLayer"] _supports_sdpa = True + _can_record_outputs = { + "hidden_states": Qwen2_5OmniAudioEncoderLayer, + "attentions": Qwen2_5OmniAudioAttention, + } def __init__(self, config: Qwen2_5OmniAudioEncoderConfig): super().__init__(config) @@ -1712,14 +1723,9 @@ def _prepare_attention_mask(self, inputs_tensor: torch.Tensor, cu_seqlens: torch attention_mask[..., cu_seqlens[i - 1] : cu_seqlens[i], cu_seqlens[i - 1] : cu_seqlens[i]] = 0 return attention_mask + @check_model_inputs(tie_last_hidden_states=False) @auto_docstring - def forward( - self, - input_features, - feature_lens=None, - aftercnn_lens=None, - **kwargs, - ): + def forward(self, input_features, feature_lens=None, aftercnn_lens=None, **kwargs: Unpack[TransformersKwargs]): r""" feature_lens (`torch.LongTensor` of shape `(batch_size,)`): mel length @@ -1769,7 +1775,7 @@ def forward( each_audio_states = self.proj(each_audio_states) token_audio_list.append(each_audio_states) token_audio = torch.cat(token_audio_list, dim=0) - return BaseModelOutput(last_hidden_state=token_audio) + return BaseModelOutputWithPooling(last_hidden_state=token_audio) def padded_and_mask_function(self, tensor_list, tensor_len, padding_value=0, padding_side="right"): """ @@ -1946,12 +1952,19 @@ class Qwen2_5OmniVisionEncoder(Qwen2_5_VisionTransformerPretrainedModel): input_modalities = ("image", "video") _no_split_modules = ["Qwen2_5OmniVisionBlock"] _input_embed_layer = "patch_embed" + _can_record_outputs = { + "hidden_states": Qwen2_5OmniVisionBlock, + "attentions": Qwen2_5OmniVisionAttention, + } def __init__(self, config: Qwen2_5OmniVisionEncoderConfig, *inputs, **kwargs) -> None: super().__init__(config, *inputs, **kwargs) self.blocks = nn.ModuleList([Qwen2_5OmniVisionBlock(config) for _ in range(config.depth)]) - def forward(self, hidden_states: torch.Tensor, grid_thw: torch.Tensor, **kwargs) -> torch.Tensor: + @check_model_inputs + def forward( + self, hidden_states: torch.Tensor, grid_thw: torch.Tensor, **kwargs: Unpack[TransformersKwargs] + ) -> tuple | BaseModelOutputWithPooling: """ Args: hidden_states (`torch.Tensor` of shape `(seq_len, hidden_size)`): @@ -2004,11 +2017,15 @@ def forward(self, hidden_states: torch.Tensor, grid_thw: torch.Tensor, **kwargs) rotary_pos_emb=rotary_pos_emb, **kwargs, ) - hidden_states = self.merger(hidden_states) + + merged_hidden_states = self.merger(hidden_states) reverse_indices = torch.argsort(window_index) - hidden_states = hidden_states[reverse_indices, :] + merged_hidden_states = merged_hidden_states[reverse_indices, :] - return hidden_states + return BaseModelOutputWithPooling( + last_hidden_state=hidden_states, + pooler_output=merged_hidden_states, + ) class Qwen2_5OmniRotaryEmbedding(Qwen2VLRotaryEmbedding): @@ -2088,52 +2105,56 @@ def get_input_embeddings(self): def set_input_embeddings(self, value): self.model.set_input_embeddings(value) + @can_return_tuple + @auto_docstring def get_video_features( - self, pixel_values_videos: torch.FloatTensor, video_grid_thw: torch.LongTensor | None = None - ): - """ - Encodes videos into continuous embeddings that can be forwarded to the language model. - - Args: - pixel_values_videos (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): - The tensors corresponding to the input videos. - video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*): - The temporal, height and width of feature shape of each video in LLM. + self, + pixel_values_videos: torch.FloatTensor, + video_grid_thw: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + pixel_values_videos (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): + The tensors corresponding to the input videos. + video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*): + The temporal, height and width of feature shape of each video in LLM. """ pixel_values_videos = pixel_values_videos.type(self.visual.dtype) - video_embeds = self.visual(pixel_values_videos, grid_thw=video_grid_thw) - return video_embeds - - def get_image_features(self, pixel_values: torch.FloatTensor, image_grid_thw: torch.LongTensor | None = None): - """ - Encodes images into continuous embeddings that can be forwarded to the language model. + return self.visual(pixel_values_videos, grid_thw=video_grid_thw, **kwargs) - Args: - pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): - The tensors corresponding to the input images. - image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): - The temporal, height and width of feature shape of each image in LLM. + @can_return_tuple + @auto_docstring + def get_image_features( + self, + pixel_values: torch.FloatTensor, + image_grid_thw: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): + The tensors corresponding to the input images. + image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): + The temporal, height and width of feature shape of each image in LLM. """ pixel_values = pixel_values.type(self.visual.dtype) - image_embeds = self.visual(pixel_values, grid_thw=image_grid_thw) - return image_embeds + return self.visual(pixel_values, grid_thw=image_grid_thw, **kwargs) + @can_return_tuple + @auto_docstring def get_audio_features( self, input_features: torch.FloatTensor, feature_attention_mask: torch.LongTensor | None = None, audio_feature_lengths: torch.LongTensor | None = None, - ): - """ - Encodes audios into continuous embeddings that can be forwarded to the language model. - - Args: - input_features (`torch.FloatTensor`): - The tensors corresponding to the input audios. - feature_attention_mask (`torch.LongTensor`, *optional*): - Mask to avoid performing attention on padding feature indices. Mask values selected in `[0, 1]`: - audio_feature_lengths (`torch.LongTensor` of shape `(num_audios)`, *optional*): - The length of feature shape of each audio in LLM. + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + input_features (`torch.FloatTensor`): + The tensors corresponding to the input audios. + feature_attention_mask (`torch.LongTensor`, *optional*): + Mask to avoid performing attention on padding feature indices. Mask values selected in `[0, 1]`: + audio_feature_lengths (`torch.LongTensor` of shape `(num_audios)`, *optional*): + The length of feature shape of each audio in LLM. """ if feature_attention_mask is not None: audio_feature_lengths = torch.sum(feature_attention_mask, dim=1) @@ -2149,13 +2170,13 @@ def get_audio_features( input_features, feature_lens=feature_lens, aftercnn_lens=audio_feat_lengths, + return_dict=True, + **kwargs, ) - audio_features = audio_outputs.last_hidden_state - - if audio_features.shape[0] != sum(audio_output_lengths.tolist()): + if audio_outputs.last_hidden_state.shape[0] != sum(audio_output_lengths.tolist()): raise ValueError("length of audio_features should match audio_output_lengths") - return audio_features + return audio_outputs def get_placeholder_mask( self, @@ -2305,13 +2326,14 @@ def forward( input_features, feature_attention_mask=feature_attention_mask, audio_feature_lengths=audio_feature_lengths, - ) + return_dict=True, + ).last_hidden_state audio_features = audio_features.to(inputs_embeds.device, inputs_embeds.dtype) _, _, audio_mask = self.get_placeholder_mask(input_ids, inputs_embeds=inputs_embeds) inputs_embeds = inputs_embeds.masked_scatter(audio_mask, audio_features) if pixel_values is not None: - image_embeds = self.get_image_features(pixel_values, image_grid_thw) + image_embeds = self.get_image_features(pixel_values, image_grid_thw, return_dict=True).pooler_output image_embeds = image_embeds.to(inputs_embeds.device, inputs_embeds.dtype) image_mask, _, _ = self.get_placeholder_mask( input_ids, inputs_embeds=inputs_embeds, image_features=image_embeds @@ -2319,7 +2341,7 @@ def forward( inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds) if pixel_values_videos is not None: - video_embeds = self.get_video_features(pixel_values_videos, video_grid_thw) + video_embeds = self.get_video_features(pixel_values_videos, video_grid_thw, return_dict=True).pooler_output video_embeds = video_embeds.to(inputs_embeds.device, inputs_embeds.dtype) _, video_mask, _ = self.get_placeholder_mask( input_ids, inputs_embeds=inputs_embeds, video_features=video_embeds diff --git a/src/transformers/models/qwen2_5_vl/modeling_qwen2_5_vl.py b/src/transformers/models/qwen2_5_vl/modeling_qwen2_5_vl.py index b77bd6e79602..29f66d6cd204 100644 --- a/src/transformers/models/qwen2_5_vl/modeling_qwen2_5_vl.py +++ b/src/transformers/models/qwen2_5_vl/modeling_qwen2_5_vl.py @@ -37,12 +37,12 @@ from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer -from ...modeling_outputs import BaseModelOutputWithPast, ModelOutput +from ...modeling_outputs import BaseModelOutputWithPast, BaseModelOutputWithPooling, ModelOutput from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel from ...processing_utils import Unpack from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, logging, torch_compilable_check -from ...utils.generic import is_flash_attention_requested, maybe_autocast +from ...utils.generic import check_model_inputs, is_flash_attention_requested, maybe_autocast from ..qwen2.modeling_qwen2 import Qwen2RMSNorm from .configuration_qwen2_5_vl import Qwen2_5_VLConfig, Qwen2_5_VLTextConfig, Qwen2_5_VLVisionConfig @@ -316,6 +316,10 @@ class Qwen2_5_VisionTransformerPretrainedModel(Qwen2_5_VLPreTrainedModel): config: Qwen2_5_VLVisionConfig _no_split_modules = ["Qwen2_5_VLVisionBlock"] _input_embed_layer = "patch_embed" + _can_record_outputs = { + "hidden_states": Qwen2_5_VLVisionBlock, + "attentions": Qwen2_5_VLVisionAttention, + } def __init__(self, config, *inputs, **kwargs) -> None: super().__init__(config, *inputs, **kwargs) @@ -415,7 +419,10 @@ def get_window_index(self, grid_thw): return window_index, cu_window_seqlens - def forward(self, hidden_states: torch.Tensor, grid_thw: torch.Tensor, **kwargs) -> torch.Tensor: + @check_model_inputs + def forward( + self, hidden_states: torch.Tensor, grid_thw: torch.Tensor, **kwargs: Unpack[TransformersKwargs] + ) -> tuple | BaseModelOutputWithPooling: """ Args: hidden_states (`torch.Tensor` of shape `(seq_len, hidden_size)`): @@ -469,11 +476,14 @@ def forward(self, hidden_states: torch.Tensor, grid_thw: torch.Tensor, **kwargs) **kwargs, ) - hidden_states = self.merger(hidden_states) + merged_hidden_states = self.merger(hidden_states) reverse_indices = torch.argsort(window_index) - hidden_states = hidden_states[reverse_indices, :] + merged_hidden_states = merged_hidden_states[reverse_indices, :] - return hidden_states + return BaseModelOutputWithPooling( + last_hidden_state=hidden_states, + pooler_output=merged_hidden_states, + ) @dataclass @@ -1168,39 +1178,49 @@ def get_rope_index( return position_ids, mrope_position_deltas + @can_return_tuple + @auto_docstring def get_video_features( - self, pixel_values_videos: torch.FloatTensor, video_grid_thw: torch.LongTensor | None = None - ): - """ - Encodes videos into continuous embeddings that can be forwarded to the language model. - - Args: - pixel_values_videos (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): - The tensors corresponding to the input videos. - video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*): - The temporal, height and width of feature shape of each video in LLM. + self, + pixel_values_videos: torch.FloatTensor, + video_grid_thw: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + pixel_values_videos (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): + The tensors corresponding to the input videos. + video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*): + The temporal, height and width of feature shape of each video in LLM. """ pixel_values_videos = pixel_values_videos.type(self.visual.dtype) - video_embeds = self.visual(pixel_values_videos, grid_thw=video_grid_thw) + vision_outputs = self.visual(pixel_values_videos, grid_thw=video_grid_thw, return_dict=True, **kwargs) split_sizes = (video_grid_thw.prod(-1) // self.visual.spatial_merge_size**2).tolist() - video_embeds = torch.split(video_embeds, split_sizes) - return video_embeds + video_embeds = torch.split(vision_outputs.pooler_output, split_sizes) + vision_outputs.pooler_output = video_embeds - def get_image_features(self, pixel_values: torch.FloatTensor, image_grid_thw: torch.LongTensor | None = None): - """ - Encodes images into continuous embeddings that can be forwarded to the language model. + return vision_outputs - Args: - pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): - The tensors corresponding to the input images. - image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): - The temporal, height and width of feature shape of each image in LLM. + @can_return_tuple + @auto_docstring + def get_image_features( + self, + pixel_values: torch.FloatTensor, + image_grid_thw: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): + The tensors corresponding to the input images. + image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): + The temporal, height and width of feature shape of each image in LLM. """ pixel_values = pixel_values.type(self.visual.dtype) - image_embeds = self.visual(pixel_values, grid_thw=image_grid_thw) + vision_outputs = self.visual(pixel_values, grid_thw=image_grid_thw, return_dict=True, **kwargs) split_sizes = (image_grid_thw.prod(-1) // self.visual.spatial_merge_size**2).tolist() - image_embeds = torch.split(image_embeds, split_sizes) - return image_embeds + image_embeds = torch.split(vision_outputs.pooler_output, split_sizes) + vision_outputs.pooler_output = image_embeds + + return vision_outputs def get_placeholder_mask( self, @@ -1285,7 +1305,7 @@ def forward( inputs_embeds = self.get_input_embeddings()(input_ids) if pixel_values is not None: - image_embeds = self.get_image_features(pixel_values, image_grid_thw) + image_embeds = self.get_image_features(pixel_values, image_grid_thw, return_dict=True).pooler_output image_embeds = torch.cat(image_embeds, dim=0).to(inputs_embeds.device, inputs_embeds.dtype) image_mask, _ = self.get_placeholder_mask( input_ids, inputs_embeds=inputs_embeds, image_features=image_embeds @@ -1293,7 +1313,7 @@ def forward( inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds) if pixel_values_videos is not None: - video_embeds = self.get_video_features(pixel_values_videos, video_grid_thw) + video_embeds = self.get_video_features(pixel_values_videos, video_grid_thw, return_dict=True).pooler_output video_embeds = torch.cat(video_embeds, dim=0).to(inputs_embeds.device, inputs_embeds.dtype) _, video_mask = self.get_placeholder_mask( input_ids, inputs_embeds=inputs_embeds, video_features=video_embeds @@ -1394,13 +1414,37 @@ def get_input_embeddings(self): def set_input_embeddings(self, value): self.model.set_input_embeddings(value) + @auto_docstring def get_video_features( - self, pixel_values_videos: torch.FloatTensor, video_grid_thw: torch.LongTensor | None = None - ): - return self.model.get_video_features(pixel_values_videos, video_grid_thw) + self, + pixel_values_videos: torch.FloatTensor, + video_grid_thw: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + pixel_values_videos (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): + The tensors corresponding to the input videos. + video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*): + The temporal, height and width of feature shape of each video in LLM. + """ + return self.model.get_video_features( + pixel_values_videos=pixel_values_videos, video_grid_thw=video_grid_thw, **kwargs + ) - def get_image_features(self, pixel_values: torch.FloatTensor, image_grid_thw: torch.LongTensor | None = None): - return self.model.get_image_features(pixel_values, image_grid_thw) + @auto_docstring + def get_image_features( + self, + pixel_values: torch.FloatTensor, + image_grid_thw: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): + The tensors corresponding to the input images. + image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): + The temporal, height and width of feature shape of each image in LLM. + """ + return self.model.get_image_features(pixel_values=pixel_values, image_grid_thw=image_grid_thw, **kwargs) @can_return_tuple @auto_docstring diff --git a/src/transformers/models/qwen2_5_vl/modular_qwen2_5_vl.py b/src/transformers/models/qwen2_5_vl/modular_qwen2_5_vl.py index bd0935db283f..6e266e062f16 100644 --- a/src/transformers/models/qwen2_5_vl/modular_qwen2_5_vl.py +++ b/src/transformers/models/qwen2_5_vl/modular_qwen2_5_vl.py @@ -30,10 +30,12 @@ from ...feature_extraction_utils import BatchFeature from ...image_utils import ImageInput from ...modeling_layers import GradientCheckpointingLayer +from ...modeling_outputs import BaseModelOutputWithPooling from ...modeling_utils import PreTrainedModel from ...processing_utils import MultiModalData, ProcessingKwargs, Unpack from ...tokenization_utils_base import PreTokenizedInput, TextInput from ...utils import logging +from ...utils.generic import check_model_inputs from ...video_utils import VideoInput from ..qwen2_vl.configuration_qwen2_vl import Qwen2VLConfig, Qwen2VLTextConfig from ..qwen2_vl.modeling_qwen2_vl import ( @@ -177,6 +179,10 @@ class Qwen2_5_VisionTransformerPretrainedModel(Qwen2_5_VLPreTrainedModel): config: Qwen2_5_VLVisionConfig _no_split_modules = ["Qwen2_5_VLVisionBlock"] _input_embed_layer = "patch_embed" + _can_record_outputs = { + "hidden_states": Qwen2_5_VLVisionBlock, + "attentions": Qwen2_5_VLVisionAttention, + } def __init__(self, config, *inputs, **kwargs) -> None: super().__init__(config, *inputs, **kwargs) @@ -276,7 +282,10 @@ def get_window_index(self, grid_thw): return window_index, cu_window_seqlens - def forward(self, hidden_states: torch.Tensor, grid_thw: torch.Tensor, **kwargs) -> torch.Tensor: + @check_model_inputs + def forward( + self, hidden_states: torch.Tensor, grid_thw: torch.Tensor, **kwargs: Unpack[TransformersKwargs] + ) -> tuple | BaseModelOutputWithPooling: """ Args: hidden_states (`torch.Tensor` of shape `(seq_len, hidden_size)`): @@ -330,11 +339,14 @@ def forward(self, hidden_states: torch.Tensor, grid_thw: torch.Tensor, **kwargs) **kwargs, ) - hidden_states = self.merger(hidden_states) + merged_hidden_states = self.merger(hidden_states) reverse_indices = torch.argsort(window_index) - hidden_states = hidden_states[reverse_indices, :] + merged_hidden_states = merged_hidden_states[reverse_indices, :] - return hidden_states + return BaseModelOutputWithPooling( + last_hidden_state=hidden_states, + pooler_output=merged_hidden_states, + ) class Qwen2_5_VLModelOutputWithPast(Qwen2VLModelOutputWithPast): @@ -578,7 +590,7 @@ def forward( inputs_embeds = self.get_input_embeddings()(input_ids) if pixel_values is not None: - image_embeds = self.get_image_features(pixel_values, image_grid_thw) + image_embeds = self.get_image_features(pixel_values, image_grid_thw, return_dict=True).pooler_output image_embeds = torch.cat(image_embeds, dim=0).to(inputs_embeds.device, inputs_embeds.dtype) image_mask, _ = self.get_placeholder_mask( input_ids, inputs_embeds=inputs_embeds, image_features=image_embeds @@ -586,7 +598,7 @@ def forward( inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds) if pixel_values_videos is not None: - video_embeds = self.get_video_features(pixel_values_videos, video_grid_thw) + video_embeds = self.get_video_features(pixel_values_videos, video_grid_thw, return_dict=True).pooler_output video_embeds = torch.cat(video_embeds, dim=0).to(inputs_embeds.device, inputs_embeds.dtype) _, video_mask = self.get_placeholder_mask( input_ids, inputs_embeds=inputs_embeds, video_features=video_embeds diff --git a/src/transformers/models/qwen2_vl/modeling_qwen2_vl.py b/src/transformers/models/qwen2_vl/modeling_qwen2_vl.py index 39c19a45280f..7eb1829d17c4 100644 --- a/src/transformers/models/qwen2_vl/modeling_qwen2_vl.py +++ b/src/transformers/models/qwen2_vl/modeling_qwen2_vl.py @@ -34,7 +34,7 @@ from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer -from ...modeling_outputs import BaseModelOutputWithPast, ModelOutput +from ...modeling_outputs import BaseModelOutputWithPast, BaseModelOutputWithPooling, ModelOutput from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel from ...processing_utils import Unpack @@ -45,10 +45,8 @@ logging, torch_compilable_check, ) -from ...utils.generic import is_flash_attention_requested, maybe_autocast -from ..qwen2.modeling_qwen2 import ( - Qwen2RMSNorm, -) +from ...utils.generic import check_model_inputs, is_flash_attention_requested, maybe_autocast +from ..qwen2.modeling_qwen2 import Qwen2RMSNorm from .configuration_qwen2_vl import Qwen2VLConfig, Qwen2VLTextConfig, Qwen2VLVisionConfig @@ -681,6 +679,10 @@ class Qwen2VisionTransformerPretrainedModel(Qwen2VLPreTrainedModel): input_modalities = ("image", "video") _no_split_modules = ["Qwen2VLVisionBlock"] _input_embed_layer = "patch_embed" + _can_record_outputs = { + "hidden_states": Qwen2VLVisionBlock, + "attentions": VisionAttention, + } def __init__(self, config) -> None: super().__init__(config) @@ -739,12 +741,13 @@ def rot_pos_emb(self, grid_thw): rotary_pos_emb = rotary_pos_emb_full[pos_ids].flatten(1) return rotary_pos_emb + @check_model_inputs @auto_docstring def forward( self, hidden_states: torch.Tensor, grid_thw: torch.Tensor, - **kwargs, + **kwargs: Unpack[TransformersKwargs], ) -> torch.Tensor: r""" grid_thw (`torch.LongTensor` of shape `(num_images, 3)`): @@ -773,7 +776,12 @@ def forward( **kwargs, ) - return self.merger(hidden_states) + merged_hidden_states = self.merger(hidden_states) + + return BaseModelOutputWithPooling( + last_hidden_state=hidden_states, + pooler_output=merged_hidden_states, + ) @auto_docstring @@ -1104,39 +1112,49 @@ def get_rope_index( return position_ids, mrope_position_deltas + @can_return_tuple + @auto_docstring def get_video_features( - self, pixel_values_videos: torch.FloatTensor, video_grid_thw: torch.LongTensor | None = None - ): - """ - Encodes videos into continuous embeddings that can be forwarded to the language model. - - Args: - pixel_values_videos (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): - The tensors corresponding to the input videos. - video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*): - The temporal, height and width of feature shape of each video in LLM. + self, + pixel_values_videos: torch.FloatTensor, + video_grid_thw: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + pixel_values_videos (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): + The tensors corresponding to the input videos. + video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*): + The temporal, height and width of feature shape of each video in LLM. """ pixel_values_videos = pixel_values_videos.type(self.visual.dtype) - video_embeds = self.visual(pixel_values_videos, grid_thw=video_grid_thw) + vision_outputs = self.visual(pixel_values_videos, grid_thw=video_grid_thw, return_dict=True, **kwargs) split_sizes = (video_grid_thw.prod(-1) // self.visual.spatial_merge_size**2).tolist() - video_embeds = torch.split(video_embeds, split_sizes) - return video_embeds + video_embeds = torch.split(vision_outputs.pooler_output, split_sizes) + vision_outputs.pooler_output = video_embeds - def get_image_features(self, pixel_values: torch.FloatTensor, image_grid_thw: torch.LongTensor | None = None): - """ - Encodes images into continuous embeddings that can be forwarded to the language model. + return vision_outputs - Args: - pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): - The tensors corresponding to the input images. - image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): - The temporal, height and width of feature shape of each image in LLM. + @can_return_tuple + @auto_docstring + def get_image_features( + self, + pixel_values: torch.FloatTensor, + image_grid_thw: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): + The tensors corresponding to the input images. + image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): + The temporal, height and width of feature shape of each image in LLM. """ pixel_values = pixel_values.type(self.visual.dtype) - image_embeds = self.visual(pixel_values, grid_thw=image_grid_thw) + vision_outputs = self.visual(pixel_values, grid_thw=image_grid_thw, return_dict=True, **kwargs) split_sizes = (image_grid_thw.prod(-1) // self.visual.spatial_merge_size**2).tolist() - image_embeds = torch.split(image_embeds, split_sizes) - return image_embeds + image_embeds = torch.split(vision_outputs.pooler_output, split_sizes) + vision_outputs.pooler_output = image_embeds + + return vision_outputs def get_placeholder_mask( self, @@ -1218,7 +1236,7 @@ def forward( inputs_embeds = self.get_input_embeddings()(input_ids) if pixel_values is not None: - image_embeds = self.get_image_features(pixel_values, image_grid_thw) + image_embeds = self.get_image_features(pixel_values, image_grid_thw, return_dict=True).pooler_output image_embeds = torch.cat(image_embeds, dim=0).to(inputs_embeds.device, inputs_embeds.dtype) image_mask, _ = self.get_placeholder_mask( input_ids, inputs_embeds=inputs_embeds, image_features=image_embeds @@ -1226,7 +1244,7 @@ def forward( inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds) if pixel_values_videos is not None: - video_embeds = self.get_video_features(pixel_values_videos, video_grid_thw) + video_embeds = self.get_video_features(pixel_values_videos, video_grid_thw, return_dict=True).pooler_output video_embeds = torch.cat(video_embeds, dim=0).to(inputs_embeds.device, inputs_embeds.dtype) _, video_mask = self.get_placeholder_mask( input_ids, inputs_embeds=inputs_embeds, video_features=video_embeds @@ -1293,13 +1311,37 @@ def get_input_embeddings(self): def set_input_embeddings(self, value): self.model.set_input_embeddings(value) + @auto_docstring def get_video_features( - self, pixel_values_videos: torch.FloatTensor, video_grid_thw: torch.LongTensor | None = None - ): - return self.model.get_video_features(pixel_values_videos, video_grid_thw) + self, + pixel_values_videos: torch.FloatTensor, + video_grid_thw: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + pixel_values_videos (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): + The tensors corresponding to the input videos. + video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*): + The temporal, height and width of feature shape of each video in LLM. + """ + return self.model.get_video_features( + pixel_values_videos=pixel_values_videos, video_grid_thw=video_grid_thw, **kwargs + ) - def get_image_features(self, pixel_values: torch.FloatTensor, image_grid_thw: torch.LongTensor | None = None): - return self.model.get_image_features(pixel_values, image_grid_thw) + @auto_docstring + def get_image_features( + self, + pixel_values: torch.FloatTensor, + image_grid_thw: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): + The tensors corresponding to the input images. + image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): + The temporal, height and width of feature shape of each image in LLM. + """ + return self.model.get_image_features(pixel_values=pixel_values, image_grid_thw=image_grid_thw, **kwargs) @can_return_tuple @auto_docstring diff --git a/src/transformers/models/qwen3_omni_moe/modeling_qwen3_omni_moe.py b/src/transformers/models/qwen3_omni_moe/modeling_qwen3_omni_moe.py index 3fb0fb1b0323..cfcfcec4e2c7 100644 --- a/src/transformers/models/qwen3_omni_moe/modeling_qwen3_omni_moe.py +++ b/src/transformers/models/qwen3_omni_moe/modeling_qwen3_omni_moe.py @@ -44,8 +44,8 @@ from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import ( - BaseModelOutput, BaseModelOutputWithPast, + BaseModelOutputWithPooling, CausalLMOutputWithPast, MoeCausalLMOutputWithPast, MoeModelOutputWithPast, @@ -74,6 +74,17 @@ ) +@dataclass +@auto_docstring +class BaseModelOutputWithDeepstackFeatures(BaseModelOutputWithPooling): + r""" + deepstack_features (`List[torch.FloatTensor]`, *optional*): + List of hidden-states (feature maps) from deepstack layers. + """ + + deepstack_features: list[torch.FloatTensor] | None = None + + class SinusoidsPositionEmbedding(nn.Module): def __init__(self, length, channels, max_timescale=10000): super().__init__() @@ -676,6 +687,10 @@ class Qwen3OmniMoeAudioEncoder(Qwen3OmniMoePreTrainedModel): input_modalities = "audio" _no_split_modules = ["Qwen3OmniMoeAudioEncoderLayer"] _supports_sdpa = True + _can_record_outputs = { + "hidden_states": Qwen3OmniMoeAudioEncoderLayer, + "attentions": Qwen3OmniMoeAudioAttention, + } def __init__(self, config: Qwen3OmniMoeAudioEncoderConfig): super().__init__(config) @@ -736,6 +751,7 @@ def _prepare_attention_mask(self, inputs_tensor: torch.Tensor, cu_seqlens: torch attention_mask[..., cu_seqlens[i - 1] : cu_seqlens[i], cu_seqlens[i - 1] : cu_seqlens[i]] = 0 return attention_mask + @check_model_inputs(tie_last_hidden_states=False) @auto_docstring def forward( self, @@ -805,7 +821,7 @@ def forward( hidden_states = self.proj1(hidden_states) hidden_states = self.act(hidden_states) hidden_states = self.proj2(hidden_states) - return BaseModelOutput(last_hidden_state=hidden_states) + return BaseModelOutputWithPooling(last_hidden_state=hidden_states) def padded_and_mask_function(self, tensor_list, tensor_len, padding_value=0, padding_side="right"): """ @@ -998,6 +1014,25 @@ def forward(self, seqlen: int) -> torch.Tensor: return freqs +class Qwen3OmniMoeTextTopKRouter(nn.Module): + def __init__(self, config): + super().__init__() + self.top_k = config.num_experts_per_tok + self.num_experts = config.num_experts + self.hidden_dim = config.hidden_size + self.weight = nn.Parameter(torch.zeros(self.num_experts, self.hidden_dim)) + + def forward(self, hidden_states): + hidden_states = hidden_states.reshape(-1, self.hidden_dim) + router_logits = F.linear(hidden_states, self.weight) # (seq_len, num_experts) + router_logits = torch.nn.functional.softmax(router_logits, dtype=torch.float, dim=-1) + router_top_value, router_indices = torch.topk(router_logits, self.top_k, dim=-1) # (seq_len, top_k) + router_top_value /= router_top_value.sum(dim=-1, keepdim=True) + router_top_value = router_top_value.to(router_logits.dtype) + router_scores = router_top_value + return router_logits, router_scores, router_indices + + class Qwen3OmniMoeVisionMLP(nn.Module): def __init__(self, config): super().__init__() @@ -1011,26 +1046,6 @@ def forward(self, hidden_state): return self.linear_fc2(self.act_fn(self.linear_fc1(hidden_state))) -class Qwen3OmniMoeVisionPatchEmbed(nn.Module): - def __init__(self, config) -> None: - super().__init__() - self.patch_size = config.patch_size - self.temporal_patch_size = config.temporal_patch_size - self.in_channels = config.in_channels - self.embed_dim = config.hidden_size - - kernel_size = [self.temporal_patch_size, self.patch_size, self.patch_size] - self.proj = nn.Conv3d(self.in_channels, self.embed_dim, kernel_size=kernel_size, stride=kernel_size, bias=True) - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - target_dtype = self.proj.weight.dtype - hidden_states = hidden_states.view( - -1, self.in_channels, self.temporal_patch_size, self.patch_size, self.patch_size - ) - hidden_states = self.proj(hidden_states.to(dtype=target_dtype)).view(-1, self.embed_dim) - return hidden_states - - class Qwen3OmniMoeVisionBlock(GradientCheckpointingLayer): def __init__(self, config, attn_implementation: str = "sdpa") -> None: super().__init__() @@ -1058,9 +1073,34 @@ def forward( return hidden_states +class Qwen3OmniMoeVisionPatchEmbed(nn.Module): + def __init__(self, config) -> None: + super().__init__() + self.patch_size = config.patch_size + self.temporal_patch_size = config.temporal_patch_size + self.in_channels = config.in_channels + self.embed_dim = config.hidden_size + + kernel_size = [self.temporal_patch_size, self.patch_size, self.patch_size] + self.proj = nn.Conv3d(self.in_channels, self.embed_dim, kernel_size=kernel_size, stride=kernel_size, bias=True) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + target_dtype = self.proj.weight.dtype + hidden_states = hidden_states.view( + -1, self.in_channels, self.temporal_patch_size, self.patch_size, self.patch_size + ) + hidden_states = self.proj(hidden_states.to(dtype=target_dtype)).view(-1, self.embed_dim) + return hidden_states + + class Qwen3OmniMoeVisionEncoder(Qwen3OmniMoePreTrainedModel): config: Qwen3OmniMoeVisionEncoderConfig _no_split_modules = ["Qwen3OmniMoeVisionBlock"] + _can_record_outputs = { + "router_logits": OutputRecorder(Qwen3OmniMoeTextTopKRouter, layer_name="mlp.gate", index=0), + "hidden_states": Qwen3OmniMoeVisionBlock, + "attentions": Qwen3OmniMoeVisionAttention, + } def __init__(self, config, *inputs, **kwargs) -> None: super().__init__(config, *inputs, **kwargs) @@ -1198,7 +1238,10 @@ def fast_pos_embed_interpolate(self, grid_thw): patch_pos_embeds = torch.cat(patch_pos_embeds_permute) return patch_pos_embeds - def forward(self, hidden_states: torch.Tensor, grid_thw: torch.Tensor, **kwargs) -> torch.Tensor: + @check_model_inputs + def forward( + self, hidden_states: torch.Tensor, grid_thw: torch.Tensor, **kwargs: Unpack[TransformersKwargs] + ) -> tuple | BaseModelOutputWithDeepstackFeatures: """ Args: hidden_states (`torch.Tensor` of shape `(seq_len, hidden_size)`): @@ -1246,9 +1289,13 @@ def forward(self, hidden_states: torch.Tensor, grid_thw: torch.Tensor, **kwargs) ) deepstack_feature_lists.append(deepstack_feature) - hidden_states = self.merger(hidden_states) + merged_hidden_states = self.merger(hidden_states) - return hidden_states, deepstack_feature_lists + return BaseModelOutputWithDeepstackFeatures( + last_hidden_state=hidden_states, + pooler_output=merged_hidden_states, + deepstack_features=deepstack_feature_lists, + ) @property def deepstack_merger_list(self): @@ -1935,52 +1982,56 @@ def get_input_embeddings(self): def set_input_embeddings(self, value): self.model.set_input_embeddings(value) + @can_return_tuple + @auto_docstring def get_video_features( - self, pixel_values_videos: torch.FloatTensor, video_grid_thw: torch.LongTensor | None = None - ): - """ - Encodes videos into continuous embeddings that can be forwarded to the language model. - - Args: - pixel_values_videos (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): - The tensors corresponding to the input videos. - video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*): - The temporal, height and width of feature shape of each video in LLM. + self, + pixel_values_videos: torch.FloatTensor, + video_grid_thw: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + pixel_values_videos (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): + The tensors corresponding to the input videos. + video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*): + The temporal, height and width of feature shape of each video in LLM. """ pixel_values_videos = pixel_values_videos.type(self.visual.dtype) - video_embeds = self.visual(pixel_values_videos, grid_thw=video_grid_thw) - return video_embeds + return self.visual(pixel_values_videos, grid_thw=video_grid_thw, **kwargs) - def get_image_features(self, pixel_values: torch.FloatTensor, image_grid_thw: torch.LongTensor | None = None): - """ - Encodes images into continuous embeddings that can be forwarded to the language model. - - Args: - pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): - The tensors corresponding to the input images. - image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): - The temporal, height and width of feature shape of each image in LLM. + @can_return_tuple + @auto_docstring + def get_image_features( + self, + pixel_values: torch.FloatTensor, + image_grid_thw: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): + The tensors corresponding to the input images. + image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): + The temporal, height and width of feature shape of each image in LLM. """ pixel_values = pixel_values.type(self.visual.dtype) - image_embeds = self.visual(pixel_values, grid_thw=image_grid_thw) - return image_embeds + return self.visual(pixel_values, grid_thw=image_grid_thw, **kwargs) + @can_return_tuple + @auto_docstring def get_audio_features( self, input_features: torch.FloatTensor, feature_attention_mask: torch.LongTensor | None = None, audio_feature_lengths: torch.LongTensor | None = None, - ): - """ - Encodes audios into continuous embeddings that can be forwarded to the language model. - - Args: - input_features (`torch.FloatTensor`): - The tensors corresponding to the input audios. - feature_attention_mask (`torch.LongTensor`, *optional*): - Mask to avoid performing attention on padding feature indices. Mask values selected in `[0, 1]`: - audio_feature_lengths (`torch.LongTensor` of shape `(num_audios)`, *optional*): - The length of feature shape of each audio in LLM. + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + input_features (`torch.FloatTensor`): + The tensors corresponding to the input audios. + feature_attention_mask (`torch.LongTensor`, *optional*): + Mask to avoid performing attention on padding feature indices. Mask values selected in `[0, 1]`: + audio_feature_lengths (`torch.LongTensor` of shape `(num_audios)`, *optional*): + The length of feature shape of each audio in LLM. """ if feature_attention_mask is not None: audio_feature_lengths = torch.sum(feature_attention_mask, dim=1) @@ -1992,10 +2043,11 @@ def get_audio_features( audio_outputs = self.audio_tower( input_features, feature_lens=feature_lens, + return_dict=True, + **kwargs, ) - audio_features = audio_outputs.last_hidden_state - return audio_features + return audio_outputs def get_placeholder_mask( self, @@ -2144,13 +2196,18 @@ def forward( input_features, feature_attention_mask=feature_attention_mask, audio_feature_lengths=audio_feature_lengths, - ) + return_dict=True, + ).last_hidden_state audio_features = audio_features.to(inputs_embeds.device, inputs_embeds.dtype) _, _, audio_mask = self.get_placeholder_mask(input_ids, inputs_embeds=inputs_embeds) inputs_embeds = inputs_embeds.masked_scatter(audio_mask, audio_features) if pixel_values is not None: - image_embeds, image_embeds_multiscale = self.get_image_features(pixel_values, image_grid_thw) + image_outputs: BaseModelOutputWithDeepstackFeatures = self.get_image_features( + pixel_values, image_grid_thw, return_dict=True + ) + image_embeds = image_outputs.pooler_output + image_embeds_multiscale = image_outputs.deepstack_features image_embeds = image_embeds.to(inputs_embeds.device, inputs_embeds.dtype) image_mask, _, _ = self.get_placeholder_mask( input_ids, inputs_embeds=inputs_embeds, image_features=image_embeds @@ -2158,7 +2215,9 @@ def forward( inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds) if pixel_values_videos is not None: - video_embeds, video_embeds_multiscale = self.get_video_features(pixel_values_videos, video_grid_thw) + video_embeds, video_embeds_multiscale = self.get_video_features( + pixel_values_videos, video_grid_thw, return_dict=True + ).pooler_output video_embeds = video_embeds.to(inputs_embeds.device, inputs_embeds.dtype) _, video_mask, _ = self.get_placeholder_mask( diff --git a/src/transformers/models/qwen3_omni_moe/modular_qwen3_omni_moe.py b/src/transformers/models/qwen3_omni_moe/modular_qwen3_omni_moe.py index 5d293e53783b..9bb913b232b3 100644 --- a/src/transformers/models/qwen3_omni_moe/modular_qwen3_omni_moe.py +++ b/src/transformers/models/qwen3_omni_moe/modular_qwen3_omni_moe.py @@ -34,8 +34,8 @@ from ...masking_utils import create_causal_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import ( - BaseModelOutput, BaseModelOutputWithPast, + BaseModelOutputWithPooling, CausalLMOutputWithPast, MoeCausalLMOutputWithPast, MoeModelOutputWithPast, @@ -100,6 +100,17 @@ logger = logging.get_logger(__name__) +@dataclass +@auto_docstring +class BaseModelOutputWithDeepstackFeatures(BaseModelOutputWithPooling): + r""" + deepstack_features (`List[torch.FloatTensor]`, *optional*): + List of hidden-states (feature maps) from deepstack layers. + """ + + deepstack_features: list[torch.FloatTensor] | None = None + + def _get_feat_extract_output_lengths(input_lengths): """ Computes the output length of the convolutional layers and the output length of the audio encoder @@ -1305,7 +1316,7 @@ def forward( hidden_states = self.proj1(hidden_states) hidden_states = self.act(hidden_states) hidden_states = self.proj2(hidden_states) - return BaseModelOutput(last_hidden_state=hidden_states) + return BaseModelOutputWithPooling(last_hidden_state=hidden_states) class Qwen3OmniMoeVisionAttention(Qwen3VLMoeVisionAttention): @@ -1441,22 +1452,22 @@ def __init__(self, config): self.num_experts_per_tok = config.text_config.num_experts_per_tok self.router_aux_loss_coef = config.text_config.router_aux_loss_coef + @can_return_tuple + @auto_docstring def get_audio_features( self, input_features: torch.FloatTensor, feature_attention_mask: torch.LongTensor | None = None, audio_feature_lengths: torch.LongTensor | None = None, - ): - """ - Encodes audios into continuous embeddings that can be forwarded to the language model. - - Args: - input_features (`torch.FloatTensor`): - The tensors corresponding to the input audios. - feature_attention_mask (`torch.LongTensor`, *optional*): - Mask to avoid performing attention on padding feature indices. Mask values selected in `[0, 1]`: - audio_feature_lengths (`torch.LongTensor` of shape `(num_audios)`, *optional*): - The length of feature shape of each audio in LLM. + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + input_features (`torch.FloatTensor`): + The tensors corresponding to the input audios. + feature_attention_mask (`torch.LongTensor`, *optional*): + Mask to avoid performing attention on padding feature indices. Mask values selected in `[0, 1]`: + audio_feature_lengths (`torch.LongTensor` of shape `(num_audios)`, *optional*): + The length of feature shape of each audio in LLM. """ if feature_attention_mask is not None: audio_feature_lengths = torch.sum(feature_attention_mask, dim=1) @@ -1468,10 +1479,11 @@ def get_audio_features( audio_outputs = self.audio_tower( input_features, feature_lens=feature_lens, + return_dict=True, + **kwargs, ) - audio_features = audio_outputs.last_hidden_state - return audio_features + return audio_outputs @can_return_tuple @auto_docstring @@ -1515,13 +1527,18 @@ def forward( input_features, feature_attention_mask=feature_attention_mask, audio_feature_lengths=audio_feature_lengths, - ) + return_dict=True, + ).last_hidden_state audio_features = audio_features.to(inputs_embeds.device, inputs_embeds.dtype) _, _, audio_mask = self.get_placeholder_mask(input_ids, inputs_embeds=inputs_embeds) inputs_embeds = inputs_embeds.masked_scatter(audio_mask, audio_features) if pixel_values is not None: - image_embeds, image_embeds_multiscale = self.get_image_features(pixel_values, image_grid_thw) + image_outputs: BaseModelOutputWithDeepstackFeatures = self.get_image_features( + pixel_values, image_grid_thw, return_dict=True + ) + image_embeds = image_outputs.pooler_output + image_embeds_multiscale = image_outputs.deepstack_features image_embeds = image_embeds.to(inputs_embeds.device, inputs_embeds.dtype) image_mask, _, _ = self.get_placeholder_mask( input_ids, inputs_embeds=inputs_embeds, image_features=image_embeds @@ -1529,7 +1546,9 @@ def forward( inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds) if pixel_values_videos is not None: - video_embeds, video_embeds_multiscale = self.get_video_features(pixel_values_videos, video_grid_thw) + video_embeds, video_embeds_multiscale = self.get_video_features( + pixel_values_videos, video_grid_thw, return_dict=True + ).pooler_output video_embeds = video_embeds.to(inputs_embeds.device, inputs_embeds.dtype) _, video_mask, _ = self.get_placeholder_mask( diff --git a/src/transformers/models/qwen3_vl/modeling_qwen3_vl.py b/src/transformers/models/qwen3_vl/modeling_qwen3_vl.py index f8c4ce45d5ee..a06ab3e2ca7c 100644 --- a/src/transformers/models/qwen3_vl/modeling_qwen3_vl.py +++ b/src/transformers/models/qwen3_vl/modeling_qwen3_vl.py @@ -34,7 +34,7 @@ from ...masking_utils import create_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer -from ...modeling_outputs import BaseModelOutputWithPast, ModelOutput +from ...modeling_outputs import BaseModelOutputWithPast, BaseModelOutputWithPooling, ModelOutput from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel from ...processing_utils import Unpack @@ -43,6 +43,17 @@ from .configuration_qwen3_vl import Qwen3VLConfig, Qwen3VLTextConfig, Qwen3VLVisionConfig +@dataclass +@auto_docstring +class BaseModelOutputWithDeepstackFeatures(BaseModelOutputWithPooling): + r""" + deepstack_features (`List[torch.FloatTensor]`, *optional*): + List of hidden-states (feature maps) from deepstack layers. + """ + + deepstack_features: list[torch.FloatTensor] | None = None + + class Qwen3VLVisionMLP(nn.Module): def __init__(self, config): super().__init__() @@ -602,6 +613,10 @@ def _init_weights(self, module): class Qwen3VLVisionModel(Qwen3VLPreTrainedModel): config: Qwen3VLVisionConfig _no_split_modules = ["Qwen3VLVisionBlock"] + _can_record_outputs = { + "hidden_states": Qwen3VLVisionBlock, + "attentions": Qwen3VLVisionAttention, + } def __init__(self, config, *inputs, **kwargs) -> None: super().__init__(config, *inputs, **kwargs) @@ -739,7 +754,10 @@ def fast_pos_embed_interpolate(self, grid_thw): patch_pos_embeds = torch.cat(patch_pos_embeds_permute) return patch_pos_embeds - def forward(self, hidden_states: torch.Tensor, grid_thw: torch.Tensor, **kwargs) -> torch.Tensor: + @check_model_inputs + def forward( + self, hidden_states: torch.Tensor, grid_thw: torch.Tensor, **kwargs: Unpack[TransformersKwargs] + ) -> tuple | BaseModelOutputWithDeepstackFeatures: """ Args: hidden_states (`torch.Tensor` of shape `(seq_len, hidden_size)`): @@ -787,9 +805,13 @@ def forward(self, hidden_states: torch.Tensor, grid_thw: torch.Tensor, **kwargs) ) deepstack_feature_lists.append(deepstack_feature) - hidden_states = self.merger(hidden_states) + merged_hidden_states = self.merger(hidden_states) - return hidden_states, deepstack_feature_lists + return BaseModelOutputWithDeepstackFeatures( + last_hidden_state=hidden_states, + pooler_output=merged_hidden_states, + deepstack_features=deepstack_feature_lists, + ) @auto_docstring( @@ -1066,36 +1088,47 @@ def get_rope_index( return position_ids, mrope_position_deltas + @can_return_tuple + @auto_docstring def get_video_features( - self, pixel_values_videos: torch.FloatTensor, video_grid_thw: torch.LongTensor | None = None - ): - """ - Encodes videos into continuous embeddings that can be forwarded to the language model. The deepstack visual features are also returned. - - Args: - pixel_values_videos (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): - The tensors corresponding to the input videos. - video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*): - The temporal, height and width of feature shape of each video in LLM. + self, + pixel_values_videos: torch.FloatTensor, + video_grid_thw: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithDeepstackFeatures: + r""" + pixel_values_videos (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): + The tensors corresponding to the input videos. + video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*): + The temporal, height and width of feature shape of each video in LLM. """ # Same implementation as for images - return self.get_image_features(pixel_values_videos, video_grid_thw) - - def get_image_features(self, pixel_values: torch.FloatTensor, image_grid_thw: torch.LongTensor | None = None): - """ - Encodes images into continuous embeddings that can be forwarded to the language model. The deepstack visual features are also returned. + return self.get_image_features(pixel_values_videos, video_grid_thw, **kwargs) - Args: - pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): - The tensors corresponding to the input images. - image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): - The temporal, height and width of feature shape of each image in LLM. + @can_return_tuple + @auto_docstring + def get_image_features( + self, + pixel_values: torch.FloatTensor, + image_grid_thw: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithDeepstackFeatures: + r""" + pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): + The tensors corresponding to the input images. + image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): + The temporal, height and width of feature shape of each image in LLM. """ pixel_values = pixel_values.type(self.visual.dtype) - image_embeds, deepstack_image_embeds = self.visual(pixel_values, grid_thw=image_grid_thw) + vision_output: BaseModelOutputWithDeepstackFeatures = self.visual( + pixel_values, grid_thw=image_grid_thw, return_dict=True, **kwargs + ) + image_embeds = vision_output.pooler_output split_sizes = (image_grid_thw.prod(-1) // self.visual.spatial_merge_size**2).tolist() image_embeds = torch.split(image_embeds, split_sizes) - return image_embeds, deepstack_image_embeds + vision_output.pooler_output = image_embeds + + return vision_output def get_placeholder_mask( self, @@ -1170,7 +1203,11 @@ def forward( video_mask = None if pixel_values is not None: - image_embeds, deepstack_image_embeds = self.get_image_features(pixel_values, image_grid_thw) + image_outputs: BaseModelOutputWithDeepstackFeatures = self.get_image_features( + pixel_values, image_grid_thw, return_dict=True + ) + image_embeds = image_outputs.pooler_output + deepstack_image_embeds = image_outputs.deepstack_features image_embeds = torch.cat(image_embeds, dim=0).to(inputs_embeds.device, inputs_embeds.dtype) image_mask, _ = self.get_placeholder_mask( input_ids, inputs_embeds=inputs_embeds, image_features=image_embeds @@ -1178,7 +1215,11 @@ def forward( inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds) if pixel_values_videos is not None: - video_embeds, deepstack_video_embeds = self.get_video_features(pixel_values_videos, video_grid_thw) + video_outputs: BaseModelOutputWithDeepstackFeatures = self.get_video_features( + pixel_values_videos, video_grid_thw, return_dict=True + ) + video_embeds = video_outputs.pooler_output + deepstack_video_embeds = video_outputs.deepstack_features video_embeds = torch.cat(video_embeds, dim=0).to(inputs_embeds.device, inputs_embeds.dtype) _, video_mask = self.get_placeholder_mask( input_ids, inputs_embeds=inputs_embeds, video_features=video_embeds @@ -1298,13 +1339,37 @@ def get_input_embeddings(self): def set_input_embeddings(self, value): self.model.set_input_embeddings(value) + @auto_docstring def get_video_features( - self, pixel_values_videos: torch.FloatTensor, video_grid_thw: torch.LongTensor | None = None - ): - return self.model.get_video_features(pixel_values_videos, video_grid_thw) + self, + pixel_values_videos: torch.FloatTensor, + video_grid_thw: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithDeepstackFeatures: + r""" + pixel_values_videos (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): + The tensors corresponding to the input videos. + video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*): + The temporal, height and width of feature shape of each video in LLM. + """ + return self.model.get_video_features( + pixel_values_videos=pixel_values_videos, video_grid_thw=video_grid_thw, **kwargs + ) - def get_image_features(self, pixel_values: torch.FloatTensor, image_grid_thw: torch.LongTensor | None = None): - return self.model.get_image_features(pixel_values, image_grid_thw) + @auto_docstring + def get_image_features( + self, + pixel_values: torch.FloatTensor, + image_grid_thw: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithDeepstackFeatures: + r""" + pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): + The tensors corresponding to the input images. + image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): + The temporal, height and width of feature shape of each image in LLM. + """ + return self.model.get_image_features(pixel_values=pixel_values, image_grid_thw=image_grid_thw, **kwargs) @can_return_tuple def forward( diff --git a/src/transformers/models/qwen3_vl/modular_qwen3_vl.py b/src/transformers/models/qwen3_vl/modular_qwen3_vl.py index b6813585698c..8c38b4391c57 100644 --- a/src/transformers/models/qwen3_vl/modular_qwen3_vl.py +++ b/src/transformers/models/qwen3_vl/modular_qwen3_vl.py @@ -14,6 +14,7 @@ """PyTorch Qwen3-VL model.""" from collections.abc import Callable +from dataclasses import dataclass from typing import Any import numpy as np @@ -29,7 +30,7 @@ from ...image_utils import ImageInput from ...masking_utils import create_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs -from ...modeling_outputs import BaseModelOutputWithPast +from ...modeling_outputs import BaseModelOutputWithPast, BaseModelOutputWithPooling from ...modeling_rope_utils import RopeParameters, dynamic_rope_update from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel from ...processing_utils import ProcessingKwargs, Unpack @@ -65,6 +66,17 @@ logger = logging.get_logger(__name__) +@dataclass +@auto_docstring +class BaseModelOutputWithDeepstackFeatures(BaseModelOutputWithPooling): + r""" + deepstack_features (`List[torch.FloatTensor]`, *optional*): + List of hidden-states (feature maps) from deepstack layers. + """ + + deepstack_features: list[torch.FloatTensor] | None = None + + class Qwen3VLVisionConfig(PreTrainedConfig): model_type = "qwen3_vl" base_config_key = "vision_config" @@ -499,6 +511,10 @@ def _init_weights(self, module): class Qwen3VLVisionModel(Qwen3VLPreTrainedModel): config: Qwen3VLVisionConfig _no_split_modules = ["Qwen3VLVisionBlock"] + _can_record_outputs = { + "hidden_states": Qwen3VLVisionBlock, + "attentions": Qwen3VLVisionAttention, + } def __init__(self, config, *inputs, **kwargs) -> None: super().__init__(config, *inputs, **kwargs) @@ -636,7 +652,10 @@ def fast_pos_embed_interpolate(self, grid_thw): patch_pos_embeds = torch.cat(patch_pos_embeds_permute) return patch_pos_embeds - def forward(self, hidden_states: torch.Tensor, grid_thw: torch.Tensor, **kwargs) -> torch.Tensor: + @check_model_inputs + def forward( + self, hidden_states: torch.Tensor, grid_thw: torch.Tensor, **kwargs: Unpack[TransformersKwargs] + ) -> tuple | BaseModelOutputWithDeepstackFeatures: """ Args: hidden_states (`torch.Tensor` of shape `(seq_len, hidden_size)`): @@ -684,9 +703,13 @@ def forward(self, hidden_states: torch.Tensor, grid_thw: torch.Tensor, **kwargs) ) deepstack_feature_lists.append(deepstack_feature) - hidden_states = self.merger(hidden_states) + merged_hidden_states = self.merger(hidden_states) - return hidden_states, deepstack_feature_lists + return BaseModelOutputWithDeepstackFeatures( + last_hidden_state=hidden_states, + pooler_output=merged_hidden_states, + deepstack_features=deepstack_feature_lists, + ) @auto_docstring( @@ -939,36 +962,47 @@ def get_rope_index( return position_ids, mrope_position_deltas - def get_image_features(self, pixel_values: torch.FloatTensor, image_grid_thw: torch.LongTensor | None = None): - """ - Encodes images into continuous embeddings that can be forwarded to the language model. The deepstack visual features are also returned. - - Args: - pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): - The tensors corresponding to the input images. - image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): - The temporal, height and width of feature shape of each image in LLM. + @can_return_tuple + @auto_docstring + def get_image_features( + self, + pixel_values: torch.FloatTensor, + image_grid_thw: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithDeepstackFeatures: + r""" + pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): + The tensors corresponding to the input images. + image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): + The temporal, height and width of feature shape of each image in LLM. """ pixel_values = pixel_values.type(self.visual.dtype) - image_embeds, deepstack_image_embeds = self.visual(pixel_values, grid_thw=image_grid_thw) + vision_output: BaseModelOutputWithDeepstackFeatures = self.visual( + pixel_values, grid_thw=image_grid_thw, return_dict=True, **kwargs + ) + image_embeds = vision_output.pooler_output split_sizes = (image_grid_thw.prod(-1) // self.visual.spatial_merge_size**2).tolist() image_embeds = torch.split(image_embeds, split_sizes) - return image_embeds, deepstack_image_embeds + vision_output.pooler_output = image_embeds - def get_video_features( - self, pixel_values_videos: torch.FloatTensor, video_grid_thw: torch.LongTensor | None = None - ): - """ - Encodes videos into continuous embeddings that can be forwarded to the language model. The deepstack visual features are also returned. + return vision_output - Args: - pixel_values_videos (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): - The tensors corresponding to the input videos. - video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*): - The temporal, height and width of feature shape of each video in LLM. + @can_return_tuple + @auto_docstring + def get_video_features( + self, + pixel_values_videos: torch.FloatTensor, + video_grid_thw: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithDeepstackFeatures: + r""" + pixel_values_videos (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): + The tensors corresponding to the input videos. + video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*): + The temporal, height and width of feature shape of each video in LLM. """ # Same implementation as for images - return self.get_image_features(pixel_values_videos, video_grid_thw) + return self.get_image_features(pixel_values_videos, video_grid_thw, **kwargs) @auto_docstring @check_model_inputs @@ -1002,7 +1036,11 @@ def forward( video_mask = None if pixel_values is not None: - image_embeds, deepstack_image_embeds = self.get_image_features(pixel_values, image_grid_thw) + image_outputs: BaseModelOutputWithDeepstackFeatures = self.get_image_features( + pixel_values, image_grid_thw, return_dict=True + ) + image_embeds = image_outputs.pooler_output + deepstack_image_embeds = image_outputs.deepstack_features image_embeds = torch.cat(image_embeds, dim=0).to(inputs_embeds.device, inputs_embeds.dtype) image_mask, _ = self.get_placeholder_mask( input_ids, inputs_embeds=inputs_embeds, image_features=image_embeds @@ -1010,7 +1048,11 @@ def forward( inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds) if pixel_values_videos is not None: - video_embeds, deepstack_video_embeds = self.get_video_features(pixel_values_videos, video_grid_thw) + video_outputs: BaseModelOutputWithDeepstackFeatures = self.get_video_features( + pixel_values_videos, video_grid_thw, return_dict=True + ) + video_embeds = video_outputs.pooler_output + deepstack_video_embeds = video_outputs.deepstack_features video_embeds = torch.cat(video_embeds, dim=0).to(inputs_embeds.device, inputs_embeds.dtype) _, video_mask = self.get_placeholder_mask( input_ids, inputs_embeds=inputs_embeds, video_features=video_embeds @@ -1089,6 +1131,14 @@ class Qwen3VLForConditionalGeneration(Qwen2_5_VLForConditionalGeneration): config: Qwen3VLConfig _checkpoint_conversion_mapping = {} + @auto_docstring + def get_image_features(self, **super_kwargs) -> tuple | BaseModelOutputWithDeepstackFeatures: + return super().get_image_features(**super_kwargs) + + @auto_docstring + def get_video_features(self, **super_kwargs) -> tuple | BaseModelOutputWithDeepstackFeatures: + return super().get_video_features(**super_kwargs) + @can_return_tuple def forward( self, diff --git a/src/transformers/models/qwen3_vl_moe/modeling_qwen3_vl_moe.py b/src/transformers/models/qwen3_vl_moe/modeling_qwen3_vl_moe.py index fc44a924bfb7..d697286d0eca 100644 --- a/src/transformers/models/qwen3_vl_moe/modeling_qwen3_vl_moe.py +++ b/src/transformers/models/qwen3_vl_moe/modeling_qwen3_vl_moe.py @@ -39,7 +39,7 @@ from ...masking_utils import create_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer -from ...modeling_outputs import BaseModelOutputWithPast, ModelOutput +from ...modeling_outputs import BaseModelOutputWithPast, BaseModelOutputWithPooling, ModelOutput from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel from ...processing_utils import Unpack @@ -415,55 +415,6 @@ def forward(self, seqlen: int) -> torch.Tensor: return freqs -class Qwen3VLMoeVisionMLP(nn.Module): - def __init__(self, config): - super().__init__() - self.hidden_size = config.hidden_size - self.intermediate_size = config.intermediate_size - self.linear_fc1 = nn.Linear(self.hidden_size, self.intermediate_size, bias=True) - self.linear_fc2 = nn.Linear(self.intermediate_size, self.hidden_size, bias=True) - self.act_fn = ACT2FN[config.hidden_act] - - def forward(self, hidden_state): - return self.linear_fc2(self.act_fn(self.linear_fc1(hidden_state))) - - -class Qwen3VLMoeVisionPatchEmbed(nn.Module): - def __init__(self, config) -> None: - super().__init__() - self.patch_size = config.patch_size - self.temporal_patch_size = config.temporal_patch_size - self.in_channels = config.in_channels - self.embed_dim = config.hidden_size - - kernel_size = [self.temporal_patch_size, self.patch_size, self.patch_size] - self.proj = nn.Conv3d(self.in_channels, self.embed_dim, kernel_size=kernel_size, stride=kernel_size, bias=True) - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - target_dtype = self.proj.weight.dtype - hidden_states = hidden_states.view( - -1, self.in_channels, self.temporal_patch_size, self.patch_size, self.patch_size - ) - hidden_states = self.proj(hidden_states.to(dtype=target_dtype)).view(-1, self.embed_dim) - return hidden_states - - -class Qwen3VLMoeVisionPatchMerger(nn.Module): - def __init__(self, config: Qwen3VLMoeVisionConfig, use_postshuffle_norm=False) -> None: - super().__init__() - self.hidden_size = config.hidden_size * (config.spatial_merge_size**2) - self.use_postshuffle_norm = use_postshuffle_norm - self.norm = nn.LayerNorm(self.hidden_size if use_postshuffle_norm else config.hidden_size, eps=1e-6) - self.linear_fc1 = nn.Linear(self.hidden_size, self.hidden_size) - self.act_fn = nn.GELU() - self.linear_fc2 = nn.Linear(self.hidden_size, config.out_hidden_size) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - x = self.norm(x.view(-1, self.hidden_size) if self.use_postshuffle_norm else x).view(-1, self.hidden_size) - x = self.linear_fc2(self.act_fn(self.linear_fc1(x))) - return x - - def apply_rotary_pos_emb_vision( q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor ) -> tuple[torch.Tensor, torch.Tensor]: @@ -561,6 +512,19 @@ def forward( return attn_output +class Qwen3VLMoeVisionMLP(nn.Module): + def __init__(self, config): + super().__init__() + self.hidden_size = config.hidden_size + self.intermediate_size = config.intermediate_size + self.linear_fc1 = nn.Linear(self.hidden_size, self.intermediate_size, bias=True) + self.linear_fc2 = nn.Linear(self.intermediate_size, self.hidden_size, bias=True) + self.act_fn = ACT2FN[config.hidden_act] + + def forward(self, hidden_state): + return self.linear_fc2(self.act_fn(self.linear_fc1(hidden_state))) + + class Qwen3VLMoeVisionBlock(GradientCheckpointingLayer): def __init__(self, config, attn_implementation: str = "sdpa") -> None: super().__init__() @@ -588,9 +552,61 @@ def forward( return hidden_states +@dataclass +@auto_docstring +class BaseModelOutputWithDeepstackFeatures(BaseModelOutputWithPooling): + r""" + deepstack_features (`List[torch.FloatTensor]`, *optional*): + List of hidden-states (feature maps) from deepstack layers. + """ + + deepstack_features: list[torch.FloatTensor] | None = None + + +class Qwen3VLMoeVisionPatchEmbed(nn.Module): + def __init__(self, config) -> None: + super().__init__() + self.patch_size = config.patch_size + self.temporal_patch_size = config.temporal_patch_size + self.in_channels = config.in_channels + self.embed_dim = config.hidden_size + + kernel_size = [self.temporal_patch_size, self.patch_size, self.patch_size] + self.proj = nn.Conv3d(self.in_channels, self.embed_dim, kernel_size=kernel_size, stride=kernel_size, bias=True) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + target_dtype = self.proj.weight.dtype + hidden_states = hidden_states.view( + -1, self.in_channels, self.temporal_patch_size, self.patch_size, self.patch_size + ) + hidden_states = self.proj(hidden_states.to(dtype=target_dtype)).view(-1, self.embed_dim) + return hidden_states + + +class Qwen3VLMoeVisionPatchMerger(nn.Module): + def __init__(self, config: Qwen3VLMoeVisionConfig, use_postshuffle_norm=False) -> None: + super().__init__() + self.hidden_size = config.hidden_size * (config.spatial_merge_size**2) + self.use_postshuffle_norm = use_postshuffle_norm + self.norm = nn.LayerNorm(self.hidden_size if use_postshuffle_norm else config.hidden_size, eps=1e-6) + self.linear_fc1 = nn.Linear(self.hidden_size, self.hidden_size) + self.act_fn = nn.GELU() + self.linear_fc2 = nn.Linear(self.hidden_size, config.out_hidden_size) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.norm(x.view(-1, self.hidden_size) if self.use_postshuffle_norm else x).view(-1, self.hidden_size) + x = self.linear_fc2(self.act_fn(self.linear_fc1(x))) + return x + + class Qwen3VLMoeVisionModel(Qwen3VLMoePreTrainedModel): config: Qwen3VLMoeVisionConfig _no_split_modules = ["Qwen3VLMoeVisionBlock"] + _can_record_outputs = { + "router_logits": OutputRecorder(Qwen3VLMoeTextTopKRouter, layer_name="mlp.gate", index=0), + "hidden_states": Qwen3VLMoeVisionBlock, + "attentions": Qwen3VLMoeVisionAttention, + } def __init__(self, config, *inputs, **kwargs) -> None: super().__init__(config, *inputs, **kwargs) @@ -728,7 +744,10 @@ def fast_pos_embed_interpolate(self, grid_thw): patch_pos_embeds = torch.cat(patch_pos_embeds_permute) return patch_pos_embeds - def forward(self, hidden_states: torch.Tensor, grid_thw: torch.Tensor, **kwargs) -> torch.Tensor: + @check_model_inputs + def forward( + self, hidden_states: torch.Tensor, grid_thw: torch.Tensor, **kwargs: Unpack[TransformersKwargs] + ) -> tuple | BaseModelOutputWithDeepstackFeatures: """ Args: hidden_states (`torch.Tensor` of shape `(seq_len, hidden_size)`): @@ -776,9 +795,13 @@ def forward(self, hidden_states: torch.Tensor, grid_thw: torch.Tensor, **kwargs) ) deepstack_feature_lists.append(deepstack_feature) - hidden_states = self.merger(hidden_states) + merged_hidden_states = self.merger(hidden_states) - return hidden_states, deepstack_feature_lists + return BaseModelOutputWithDeepstackFeatures( + last_hidden_state=hidden_states, + pooler_output=merged_hidden_states, + deepstack_features=deepstack_feature_lists, + ) class Qwen3VLMoeTextRotaryEmbedding(nn.Module): @@ -1198,36 +1221,47 @@ def get_rope_index( return position_ids, mrope_position_deltas + @can_return_tuple + @auto_docstring def get_video_features( - self, pixel_values_videos: torch.FloatTensor, video_grid_thw: torch.LongTensor | None = None - ): - """ - Encodes videos into continuous embeddings that can be forwarded to the language model. The deepstack visual features are also returned. - - Args: - pixel_values_videos (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): - The tensors corresponding to the input videos. - video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*): - The temporal, height and width of feature shape of each video in LLM. + self, + pixel_values_videos: torch.FloatTensor, + video_grid_thw: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithDeepstackFeatures: + r""" + pixel_values_videos (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): + The tensors corresponding to the input videos. + video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*): + The temporal, height and width of feature shape of each video in LLM. """ # Same implementation as for images - return self.get_image_features(pixel_values_videos, video_grid_thw) + return self.get_image_features(pixel_values_videos, video_grid_thw, **kwargs) - def get_image_features(self, pixel_values: torch.FloatTensor, image_grid_thw: torch.LongTensor | None = None): - """ - Encodes images into continuous embeddings that can be forwarded to the language model. The deepstack visual features are also returned. - - Args: - pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): - The tensors corresponding to the input images. - image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): - The temporal, height and width of feature shape of each image in LLM. + @can_return_tuple + @auto_docstring + def get_image_features( + self, + pixel_values: torch.FloatTensor, + image_grid_thw: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithDeepstackFeatures: + r""" + pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): + The tensors corresponding to the input images. + image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): + The temporal, height and width of feature shape of each image in LLM. """ pixel_values = pixel_values.type(self.visual.dtype) - image_embeds, deepstack_image_embeds = self.visual(pixel_values, grid_thw=image_grid_thw) + vision_output: BaseModelOutputWithDeepstackFeatures = self.visual( + pixel_values, grid_thw=image_grid_thw, return_dict=True, **kwargs + ) + image_embeds = vision_output.pooler_output split_sizes = (image_grid_thw.prod(-1) // self.visual.spatial_merge_size**2).tolist() image_embeds = torch.split(image_embeds, split_sizes) - return image_embeds, deepstack_image_embeds + vision_output.pooler_output = image_embeds + + return vision_output def get_placeholder_mask( self, @@ -1302,7 +1336,11 @@ def forward( video_mask = None if pixel_values is not None: - image_embeds, deepstack_image_embeds = self.get_image_features(pixel_values, image_grid_thw) + image_outputs: BaseModelOutputWithDeepstackFeatures = self.get_image_features( + pixel_values, image_grid_thw, return_dict=True + ) + image_embeds = image_outputs.pooler_output + deepstack_image_embeds = image_outputs.deepstack_features image_embeds = torch.cat(image_embeds, dim=0).to(inputs_embeds.device, inputs_embeds.dtype) image_mask, _ = self.get_placeholder_mask( input_ids, inputs_embeds=inputs_embeds, image_features=image_embeds @@ -1310,7 +1348,11 @@ def forward( inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds) if pixel_values_videos is not None: - video_embeds, deepstack_video_embeds = self.get_video_features(pixel_values_videos, video_grid_thw) + video_outputs: BaseModelOutputWithDeepstackFeatures = self.get_video_features( + pixel_values_videos, video_grid_thw, return_dict=True + ) + video_embeds = video_outputs.pooler_output + deepstack_video_embeds = video_outputs.deepstack_features video_embeds = torch.cat(video_embeds, dim=0).to(inputs_embeds.device, inputs_embeds.dtype) _, video_mask = self.get_placeholder_mask( input_ids, inputs_embeds=inputs_embeds, video_features=video_embeds @@ -1483,13 +1525,37 @@ def get_input_embeddings(self): def set_input_embeddings(self, value): self.model.set_input_embeddings(value) + @auto_docstring def get_video_features( - self, pixel_values_videos: torch.FloatTensor, video_grid_thw: torch.LongTensor | None = None - ): - return self.model.get_video_features(pixel_values_videos, video_grid_thw) + self, + pixel_values_videos: torch.FloatTensor, + video_grid_thw: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithDeepstackFeatures: + r""" + pixel_values_videos (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): + The tensors corresponding to the input videos. + video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*): + The temporal, height and width of feature shape of each video in LLM. + """ + return self.model.get_video_features( + pixel_values_videos=pixel_values_videos, video_grid_thw=video_grid_thw, **kwargs + ) - def get_image_features(self, pixel_values: torch.FloatTensor, image_grid_thw: torch.LongTensor | None = None): - return self.model.get_image_features(pixel_values, image_grid_thw) + @auto_docstring + def get_image_features( + self, + pixel_values: torch.FloatTensor, + image_grid_thw: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithDeepstackFeatures: + r""" + pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): + The tensors corresponding to the input images. + image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): + The temporal, height and width of feature shape of each image in LLM. + """ + return self.model.get_image_features(pixel_values=pixel_values, image_grid_thw=image_grid_thw, **kwargs) @can_return_tuple def forward( diff --git a/src/transformers/models/qwen3_vl_moe/modular_qwen3_vl_moe.py b/src/transformers/models/qwen3_vl_moe/modular_qwen3_vl_moe.py index a00569d1d7ab..946e972d0fd9 100644 --- a/src/transformers/models/qwen3_vl_moe/modular_qwen3_vl_moe.py +++ b/src/transformers/models/qwen3_vl_moe/modular_qwen3_vl_moe.py @@ -24,6 +24,7 @@ from ...modeling_utils import PreTrainedModel from ...processing_utils import Unpack from ...utils import TransformersKwargs, can_return_tuple, logging +from ...utils.generic import OutputRecorder from ..qwen3_moe.modeling_qwen3_moe import ( Qwen3MoeDecoderLayer, Qwen3MoeExperts, @@ -39,6 +40,8 @@ Qwen3VLModel, Qwen3VLTextAttention, Qwen3VLTextModel, + Qwen3VLVisionAttention, + Qwen3VLVisionBlock, Qwen3VLVisionModel, Qwen3VLVisionRotaryEmbedding, ) @@ -314,10 +317,22 @@ class Qwen3VLMoeVisionRotaryEmbedding(Qwen3VLVisionRotaryEmbedding): pass -class Qwen3VLMoeVisionModel(Qwen3VLVisionModel): +class Qwen3VLMoeVisionAttention(Qwen3VLVisionAttention): + pass + + +class Qwen3VLMoeVisionBlock(Qwen3VLVisionBlock): pass +class Qwen3VLMoeVisionModel(Qwen3VLVisionModel): + _can_record_outputs = { + "router_logits": OutputRecorder(Qwen3VLMoeTextTopKRouter, layer_name="mlp.gate", index=0), + "hidden_states": Qwen3VLMoeVisionBlock, + "attentions": Qwen3VLMoeVisionAttention, + } + + class Qwen3VLMoeTextModel(Qwen3VLTextModel): pass diff --git a/src/transformers/models/sam/modeling_sam.py b/src/transformers/models/sam/modeling_sam.py index 18b5a86eb8c9..ef1181ba641a 100644 --- a/src/transformers/models/sam/modeling_sam.py +++ b/src/transformers/models/sam/modeling_sam.py @@ -1056,7 +1056,7 @@ def get_input_embeddings(self): @check_model_inputs(tie_last_hidden_states=False) def forward( self, pixel_values: torch.FloatTensor | None = None, **kwargs: Unpack[TransformersKwargs] - ) -> SamVisionEncoderOutput: + ) -> tuple | SamVisionEncoderOutput: if pixel_values is None: raise ValueError("You have to specify pixel_values") diff --git a/src/transformers/models/sam2/modeling_sam2.py b/src/transformers/models/sam2/modeling_sam2.py index 0baa5aaa3461..7d93142aff65 100644 --- a/src/transformers/models/sam2/modeling_sam2.py +++ b/src/transformers/models/sam2/modeling_sam2.py @@ -33,11 +33,16 @@ from ... import initialization as init from ...activations import ACT2FN from ...modeling_layers import GradientCheckpointingLayer -from ...modeling_outputs import BaseModelOutput +from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel from ...processing_utils import Unpack from ...pytorch_utils import compile_compatible_method_lru_cache -from ...utils import ModelOutput, auto_docstring, logging +from ...utils import ( + ModelOutput, + auto_docstring, + can_return_tuple, + logging, +) from ...utils.generic import TransformersKwargs, check_model_inputs, is_flash_attention_requested from ..auto import AutoModel from .configuration_sam2 import ( @@ -54,16 +59,10 @@ @dataclass @auto_docstring(custom_intro="Base class for the vision encoder's outputs.") -class Sam2VisionEncoderOutput(ModelOutput): +class Sam2VisionEncoderOutput(BaseModelOutputWithPooling): r""" last_hidden_state (`torch.FloatTensor` of shape `(batch_size, height, width, hidden_size)`): Sequence of hidden-states at the output of the last layer of the model. - fpn_hidden_states (`tuple(torch.FloatTensor)`): - Tuple of `torch.FloatTensor` (one for each feature level, from high to low resolution) of shape - `(batch_size, hidden_size, height, width)`. Feature maps from the Feature Pyramid Network neck. - fpn_position_encoding (`tuple(torch.FloatTensor)`): - Tuple of `torch.FloatTensor` (one for each feature level, from high to low resolution) of shape - `(batch_size, hidden_size, height, width)`. Positional encodings corresponding to the `fpn_hidden_states`. hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, + one for the output of each stage) of shape `(batch_size, height, width, hidden_size)`. Hidden-states of the @@ -72,13 +71,16 @@ class Sam2VisionEncoderOutput(ModelOutput): Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, sequence_length)`. Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads. + fpn_hidden_states (`tuple(torch.FloatTensor)`): + Tuple of `torch.FloatTensor` (one for each feature level, from high to low resolution) of shape + `(batch_size, hidden_size, height, width)`. Feature maps from the Feature Pyramid Network neck. + fpn_position_encoding (`tuple(torch.FloatTensor)`): + Tuple of `torch.FloatTensor` (one for each feature level, from high to low resolution) of shape + `(batch_size, hidden_size, height, width)`. Positional encodings corresponding to the `fpn_hidden_states`. """ - last_hidden_state: torch.FloatTensor | None = None fpn_hidden_states: torch.FloatTensor | None = None fpn_position_encoding: torch.FloatTensor | None = None - hidden_states: tuple[torch.FloatTensor, ...] | None = None - attentions: tuple[torch.FloatTensor, ...] | None = None @dataclass @@ -1341,7 +1343,8 @@ def get_image_embeddings( Input pixel values """ batch_size = pixel_values.shape[0] - feature_maps, _, _, _ = self.get_image_features(pixel_values, **kwargs) + image_outputs = self.get_image_features(pixel_values, return_dict=True, **kwargs) + feature_maps = image_outputs.fpn_hidden_states # add no memory embedding to the last feature map feature_maps[-1] = feature_maps[-1] + self.no_memory_embedding @@ -1497,10 +1500,10 @@ def forward( vision_hidden_states = None if pixel_values is not None: - feature_maps, _, vision_hidden_states, vision_attentions = self.get_image_features( - pixel_values, - **kwargs, - ) + image_outputs: Sam2VisionEncoderOutput = self.get_image_features(pixel_values, return_dict=True, **kwargs) + feature_maps = image_outputs.fpn_hidden_states + vision_hidden_states = image_outputs.hidden_states + vision_attentions = image_outputs.attentions # add no memory embedding to the last feature map feature_maps[-1] = feature_maps[-1] + self.no_memory_embedding @@ -1560,34 +1563,18 @@ def forward( vision_attentions=vision_attentions, ) + @can_return_tuple + @auto_docstring def get_image_features( self, pixel_values: torch.FloatTensor, **kwargs: Unpack[TransformersKwargs], - ) -> tuple[ - list[torch.Tensor], - list[torch.Tensor], - tuple[torch.FloatTensor, ...] | None, - tuple[torch.FloatTensor, ...] | None, - ]: + ) -> tuple | Sam2VisionEncoderOutput: r""" - Extract and preprocess image features using the vision encoder. - - Args: - pixel_values (`torch.FloatTensor`): - Input pixel values of shape `(batch_size, num_channels, height, width)`. - - Returns: - `tuple`: A tuple containing: - - feature_maps (`list[torch.Tensor]`): List of feature maps from different levels. - - feature_maps_position_embeddings (`list[torch.Tensor]`): List of positional embeddings for each feature level. - - vision_hidden_states (`tuple[torch.FloatTensor]`, *optional*): Hidden states from the vision encoder. - - vision_attentions (`tuple[torch.FloatTensor]`, *optional*): Attention weights from the vision encoder. + pixel_values (`torch.FloatTensor`): + Input pixel values of shape `(batch_size, num_channels, height, width)`. """ - vision_outputs: Sam2VisionEncoderOutput = self.vision_encoder( - pixel_values, - **kwargs, - ) + vision_outputs: Sam2VisionEncoderOutput = self.vision_encoder(pixel_values, return_dict=True, **kwargs) feature_maps = vision_outputs.fpn_hidden_states feature_maps_position_embeddings = vision_outputs.fpn_position_encoding @@ -1604,8 +1591,10 @@ def get_image_features( feature_map_position_embedding.flatten(2).permute(2, 0, 1) for feature_map_position_embedding in feature_maps_position_embeddings ] + vision_outputs.fpn_hidden_states = feature_maps + vision_outputs.fpn_position_encoding = feature_maps_position_embeddings - return feature_maps, feature_maps_position_embeddings, vision_outputs.hidden_states, vision_outputs.attentions + return vision_outputs __all__ = ["Sam2Model", "Sam2VisionModel", "Sam2PreTrainedModel", "Sam2HieraDetModel"] diff --git a/src/transformers/models/sam2/modular_sam2.py b/src/transformers/models/sam2/modular_sam2.py index 5e8a0c65b40a..42f30cafd05a 100644 --- a/src/transformers/models/sam2/modular_sam2.py +++ b/src/transformers/models/sam2/modular_sam2.py @@ -36,14 +36,10 @@ pil_torch_interpolation_mapping, ) from ...modeling_layers import GradientCheckpointingLayer +from ...modeling_outputs import BaseModelOutputWithPooling from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel from ...processing_utils import ImagesKwargs, Unpack -from ...utils import ( - ModelOutput, - TensorType, - auto_docstring, - logging, -) +from ...utils import ModelOutput, TensorType, auto_docstring, can_return_tuple, logging from ...utils.generic import TransformersKwargs, check_model_inputs, is_flash_attention_requested from ..auto import AutoModel from ..maskformer.modeling_maskformer import MaskFormerSinePositionEmbedding @@ -296,16 +292,10 @@ def resize(self): @dataclass @auto_docstring(custom_intro="Base class for the vision encoder's outputs.") -class Sam2VisionEncoderOutput(ModelOutput): +class Sam2VisionEncoderOutput(BaseModelOutputWithPooling): r""" last_hidden_state (`torch.FloatTensor` of shape `(batch_size, height, width, hidden_size)`): Sequence of hidden-states at the output of the last layer of the model. - fpn_hidden_states (`tuple(torch.FloatTensor)`): - Tuple of `torch.FloatTensor` (one for each feature level, from high to low resolution) of shape - `(batch_size, hidden_size, height, width)`. Feature maps from the Feature Pyramid Network neck. - fpn_position_encoding (`tuple(torch.FloatTensor)`): - Tuple of `torch.FloatTensor` (one for each feature level, from high to low resolution) of shape - `(batch_size, hidden_size, height, width)`. Positional encodings corresponding to the `fpn_hidden_states`. hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, + one for the output of each stage) of shape `(batch_size, height, width, hidden_size)`. Hidden-states of the @@ -314,13 +304,16 @@ class Sam2VisionEncoderOutput(ModelOutput): Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, sequence_length)`. Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads. + fpn_hidden_states (`tuple(torch.FloatTensor)`): + Tuple of `torch.FloatTensor` (one for each feature level, from high to low resolution) of shape + `(batch_size, hidden_size, height, width)`. Feature maps from the Feature Pyramid Network neck. + fpn_position_encoding (`tuple(torch.FloatTensor)`): + Tuple of `torch.FloatTensor` (one for each feature level, from high to low resolution) of shape + `(batch_size, hidden_size, height, width)`. Positional encodings corresponding to the `fpn_hidden_states`. """ - last_hidden_state: torch.FloatTensor | None = None fpn_hidden_states: torch.FloatTensor | None = None fpn_position_encoding: torch.FloatTensor | None = None - hidden_states: tuple[torch.FloatTensor, ...] | None = None - attentions: tuple[torch.FloatTensor, ...] | None = None @dataclass @@ -1233,7 +1226,8 @@ def get_image_embeddings( Input pixel values """ batch_size = pixel_values.shape[0] - feature_maps, _, _, _ = self.get_image_features(pixel_values, **kwargs) + image_outputs = self.get_image_features(pixel_values, return_dict=True, **kwargs) + feature_maps = image_outputs.fpn_hidden_states # add no memory embedding to the last feature map feature_maps[-1] = feature_maps[-1] + self.no_memory_embedding @@ -1246,34 +1240,18 @@ def get_image_embeddings( return image_embeddings + @can_return_tuple + @auto_docstring def get_image_features( self, pixel_values: torch.FloatTensor, **kwargs: Unpack[TransformersKwargs], - ) -> tuple[ - list[torch.Tensor], - list[torch.Tensor], - tuple[torch.FloatTensor, ...] | None, - tuple[torch.FloatTensor, ...] | None, - ]: + ) -> tuple | Sam2VisionEncoderOutput: r""" - Extract and preprocess image features using the vision encoder. - - Args: - pixel_values (`torch.FloatTensor`): - Input pixel values of shape `(batch_size, num_channels, height, width)`. - - Returns: - `tuple`: A tuple containing: - - feature_maps (`list[torch.Tensor]`): List of feature maps from different levels. - - feature_maps_position_embeddings (`list[torch.Tensor]`): List of positional embeddings for each feature level. - - vision_hidden_states (`tuple[torch.FloatTensor]`, *optional*): Hidden states from the vision encoder. - - vision_attentions (`tuple[torch.FloatTensor]`, *optional*): Attention weights from the vision encoder. + pixel_values (`torch.FloatTensor`): + Input pixel values of shape `(batch_size, num_channels, height, width)`. """ - vision_outputs: Sam2VisionEncoderOutput = self.vision_encoder( - pixel_values, - **kwargs, - ) + vision_outputs: Sam2VisionEncoderOutput = self.vision_encoder(pixel_values, return_dict=True, **kwargs) feature_maps = vision_outputs.fpn_hidden_states feature_maps_position_embeddings = vision_outputs.fpn_position_encoding @@ -1290,8 +1268,10 @@ def get_image_features( feature_map_position_embedding.flatten(2).permute(2, 0, 1) for feature_map_position_embedding in feature_maps_position_embeddings ] + vision_outputs.fpn_hidden_states = feature_maps + vision_outputs.fpn_position_encoding = feature_maps_position_embeddings - return feature_maps, feature_maps_position_embeddings, vision_outputs.hidden_states, vision_outputs.attentions + return vision_outputs @check_model_inputs @auto_docstring @@ -1403,10 +1383,10 @@ def forward( vision_hidden_states = None if pixel_values is not None: - feature_maps, _, vision_hidden_states, vision_attentions = self.get_image_features( - pixel_values, - **kwargs, - ) + image_outputs: Sam2VisionEncoderOutput = self.get_image_features(pixel_values, return_dict=True, **kwargs) + feature_maps = image_outputs.fpn_hidden_states + vision_hidden_states = image_outputs.hidden_states + vision_attentions = image_outputs.attentions # add no memory embedding to the last feature map feature_maps[-1] = feature_maps[-1] + self.no_memory_embedding diff --git a/src/transformers/models/sam2_video/modeling_sam2_video.py b/src/transformers/models/sam2_video/modeling_sam2_video.py index 16f330a5a88a..12a9dfc4d1be 100644 --- a/src/transformers/models/sam2_video/modeling_sam2_video.py +++ b/src/transformers/models/sam2_video/modeling_sam2_video.py @@ -35,11 +35,11 @@ from ...activations import ACT2FN from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer -from ...modeling_outputs import BaseModelOutput +from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel from ...processing_utils import Unpack from ...pytorch_utils import compile_compatible_method_lru_cache -from ...utils import ModelOutput, auto_docstring, logging +from ...utils import ModelOutput, auto_docstring, can_return_tuple, logging from ...utils.generic import OutputRecorder, TransformersKwargs, is_flash_attention_requested from ..auto import AutoModel from .configuration_sam2_video import Sam2VideoConfig, Sam2VideoMaskDecoderConfig, Sam2VideoPromptEncoderConfig @@ -1152,16 +1152,10 @@ def forward(self, input_coords, input_shape=None): @dataclass @auto_docstring(custom_intro="Base class for the vision encoder's outputs.") -class Sam2VideoVisionEncoderOutput(ModelOutput): +class Sam2VideoVisionEncoderOutput(BaseModelOutputWithPooling): r""" last_hidden_state (`torch.FloatTensor` of shape `(batch_size, height, width, hidden_size)`): Sequence of hidden-states at the output of the last layer of the model. - fpn_hidden_states (`tuple(torch.FloatTensor)`): - Tuple of `torch.FloatTensor` (one for each feature level, from high to low resolution) of shape - `(batch_size, hidden_size, height, width)`. Feature maps from the Feature Pyramid Network neck. - fpn_position_encoding (`tuple(torch.FloatTensor)`): - Tuple of `torch.FloatTensor` (one for each feature level, from high to low resolution) of shape - `(batch_size, hidden_size, height, width)`. Positional encodings corresponding to the `fpn_hidden_states`. hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, + one for the output of each stage) of shape `(batch_size, height, width, hidden_size)`. Hidden-states of the @@ -1170,13 +1164,16 @@ class Sam2VideoVisionEncoderOutput(ModelOutput): Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, sequence_length)`. Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads. + fpn_hidden_states (`tuple(torch.FloatTensor)`): + Tuple of `torch.FloatTensor` (one for each feature level, from high to low resolution) of shape + `(batch_size, hidden_size, height, width)`. Feature maps from the Feature Pyramid Network neck. + fpn_position_encoding (`tuple(torch.FloatTensor)`): + Tuple of `torch.FloatTensor` (one for each feature level, from high to low resolution) of shape + `(batch_size, hidden_size, height, width)`. Positional encodings corresponding to the `fpn_hidden_states`. """ - last_hidden_state: torch.FloatTensor | None = None fpn_hidden_states: torch.FloatTensor | None = None fpn_position_encoding: torch.FloatTensor | None = None - hidden_states: tuple[torch.FloatTensor, ...] | None = None - attentions: tuple[torch.FloatTensor, ...] | None = None class Sam2VideoMaskEmbedding(nn.Module): @@ -1664,7 +1661,8 @@ def get_image_embeddings( Input pixel values """ batch_size = pixel_values.shape[0] - feature_maps, _, _, _ = self.get_image_features(pixel_values, **kwargs) + image_outputs = self.get_image_features(pixel_values, return_dict=True, **kwargs) + feature_maps = image_outputs.fpn_hidden_states # add no memory embedding to the last feature map feature_maps[-1] = feature_maps[-1] + self.no_memory_embedding @@ -1835,34 +1833,18 @@ def forward( frame_idx=frame_idx, ) + @can_return_tuple + @auto_docstring def get_image_features( self, pixel_values: torch.FloatTensor, **kwargs: Unpack[TransformersKwargs], - ) -> tuple[ - list[torch.Tensor], - list[torch.Tensor], - tuple[torch.FloatTensor, ...] | None, - tuple[torch.FloatTensor, ...] | None, - ]: + ) -> tuple | Sam2VideoVisionEncoderOutput: r""" - Extract and preprocess image features using the vision encoder. - - Args: - pixel_values (`torch.FloatTensor`): - Input pixel values of shape `(batch_size, num_channels, height, width)`. - - Returns: - `tuple`: A tuple containing: - - feature_maps (`list[torch.Tensor]`): List of feature maps from different levels. - - feature_maps_position_embeddings (`list[torch.Tensor]`): List of positional embeddings for each feature level. - - vision_hidden_states (`tuple[torch.FloatTensor]`, *optional*): Hidden states from the vision encoder. - - vision_attentions (`tuple[torch.FloatTensor]`, *optional*): Attention weights from the vision encoder. + pixel_values (`torch.FloatTensor`): + Input pixel values of shape `(batch_size, num_channels, height, width)`. """ - vision_outputs: Sam2VideoVisionEncoderOutput = self.vision_encoder( - pixel_values, - **kwargs, - ) + vision_outputs: Sam2VideoVisionEncoderOutput = self.vision_encoder(pixel_values, return_dict=True, **kwargs) feature_maps = vision_outputs.fpn_hidden_states feature_maps_position_embeddings = vision_outputs.fpn_position_encoding @@ -1879,8 +1861,10 @@ def get_image_features( feature_map_position_embedding.flatten(2).permute(2, 0, 1) for feature_map_position_embedding in feature_maps_position_embeddings ] + vision_outputs.fpn_hidden_states = feature_maps + vision_outputs.fpn_position_encoding = feature_maps_position_embeddings - return feature_maps, feature_maps_position_embeddings, vision_outputs.hidden_states, vision_outputs.attentions + return vision_outputs def _prepare_vision_features( self, @@ -1897,7 +1881,9 @@ def _prepare_vision_features( else: # Compute features using image encoder image_batch = inference_session.get_frame(frame_idx).unsqueeze(0) # Add batch dimension - vision_feats, vision_pos_embeds, _, _ = self.get_image_features(image_batch) + image_outputs = self.get_image_features(image_batch, return_dict=True) + vision_feats = image_outputs.fpn_hidden_states + vision_pos_embeds = image_outputs.fpn_position_embeddings # Cache features inference_session.cache.cache_vision_features( frame_idx, {"vision_feats": vision_feats, "vision_pos_embeds": vision_pos_embeds} @@ -2002,10 +1988,10 @@ def _single_frame_forward( vision_hidden_states = None if pixel_values is not None: - feature_maps, _, vision_hidden_states, vision_attentions = self.get_image_features( - pixel_values, - **kwargs, - ) + image_outputs = self.get_image_features(pixel_values, return_dict=True, **kwargs) + feature_maps = image_outputs.fpn_hidden_states + vision_hidden_states = image_outputs.hidden_states + vision_attentions = image_outputs.attentions # add no memory embedding to the last feature map feature_maps[-1] = feature_maps[-1] + self.no_memory_embedding diff --git a/src/transformers/models/sam2_video/modular_sam2_video.py b/src/transformers/models/sam2_video/modular_sam2_video.py index c980ca1478fb..c4ba568beb48 100644 --- a/src/transformers/models/sam2_video/modular_sam2_video.py +++ b/src/transformers/models/sam2_video/modular_sam2_video.py @@ -1534,7 +1534,9 @@ def _prepare_vision_features( else: # Compute features using image encoder image_batch = inference_session.get_frame(frame_idx).unsqueeze(0) # Add batch dimension - vision_feats, vision_pos_embeds, _, _ = self.get_image_features(image_batch) + image_outputs = self.get_image_features(image_batch, return_dict=True) + vision_feats = image_outputs.fpn_hidden_states + vision_pos_embeds = image_outputs.fpn_position_embeddings # Cache features inference_session.cache.cache_vision_features( frame_idx, {"vision_feats": vision_feats, "vision_pos_embeds": vision_pos_embeds} @@ -1639,10 +1641,10 @@ def _single_frame_forward( vision_hidden_states = None if pixel_values is not None: - feature_maps, _, vision_hidden_states, vision_attentions = self.get_image_features( - pixel_values, - **kwargs, - ) + image_outputs = self.get_image_features(pixel_values, return_dict=True, **kwargs) + feature_maps = image_outputs.fpn_hidden_states + vision_hidden_states = image_outputs.hidden_states + vision_attentions = image_outputs.attentions # add no memory embedding to the last feature map feature_maps[-1] = feature_maps[-1] + self.no_memory_embedding diff --git a/src/transformers/models/sam3/modeling_sam3.py b/src/transformers/models/sam3/modeling_sam3.py index b0e233c67952..f068c9f61a95 100644 --- a/src/transformers/models/sam3/modeling_sam3.py +++ b/src/transformers/models/sam3/modeling_sam3.py @@ -12,7 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. - import math from collections.abc import Callable, Iterable from dataclasses import dataclass @@ -32,12 +31,13 @@ from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import ( BaseModelOutput, + BaseModelOutputWithPooling, ModelOutput, ) from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel from ...processing_utils import Unpack from ...pytorch_utils import compile_compatible_method_lru_cache -from ...utils import auto_docstring, logging +from ...utils import auto_docstring, can_return_tuple, logging from ...utils.generic import TransformersKwargs, check_model_inputs, is_flash_attention_requested from ..auto import AutoModel from .configuration_sam3 import ( @@ -56,23 +56,16 @@ @dataclass @auto_docstring -class Sam3VisionEncoderOutput(ModelOutput): +class Sam3VisionEncoderOutput(BaseModelOutputWithPooling): r""" fpn_hidden_states (`tuple[torch.FloatTensor]`): Tuple of multi-level FPN feature maps. fpn_position_encoding (`tuple[torch.FloatTensor]`): Tuple of position encodings for each FPN level. - hidden_states (`tuple[torch.FloatTensor]`, *optional*): - Tuple of hidden states from all ViT layers. - attentions (`tuple[torch.FloatTensor]`, *optional*): - Tuple of attention weights from all ViT layers. """ - last_hidden_state: torch.FloatTensor = None fpn_hidden_states: tuple[torch.FloatTensor, ...] = None fpn_position_encoding: tuple[torch.FloatTensor, ...] = None - hidden_states: tuple[torch.FloatTensor] | None = None - attentions: tuple[torch.FloatTensor] | None = None @dataclass @@ -1406,7 +1399,7 @@ def forward( text_mask: torch.Tensor | None = None, spatial_sizes: list[tuple[int, int]] | None = None, **kwargs: Unpack[TransformersKwargs], - ): + ) -> tuple | Sam3DETREncoderOutput: """ Forward pass for the DETR encoder. @@ -1701,7 +1694,7 @@ def forward( text_mask: torch.Tensor | None = None, spatial_shapes: torch.Tensor | None = None, **kwargs: Unpack[TransformersKwargs], - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None, torch.Tensor | None]: + ) -> tuple | Sam3DETRDecoderOutput: """ Forward pass for the DETR decoder. @@ -2017,7 +2010,7 @@ def forward( prompt_features: torch.Tensor | None = None, prompt_mask: torch.Tensor | None = None, **kwargs: Unpack[TransformersKwargs], - ) -> dict[str, torch.Tensor]: + ) -> tuple | Sam3MaskDecoderOutput: """ Args: decoder_queries: Decoder output queries [batch_size, num_queries, hidden_size] @@ -2147,18 +2140,15 @@ def __init__(self, config: Sam3Config): self.post_init() + @can_return_tuple @auto_docstring def get_text_features( self, input_ids: torch.LongTensor, attention_mask: torch.Tensor | None = None, **kwargs: Unpack[TransformersKwargs], - ) -> torch.FloatTensor: + ) -> tuple | BaseModelOutputWithPooling: r""" - Returns: - text_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`): - Text embeddings that can be passed as `text_embeds` to the forward method. - Example: ```python @@ -2171,7 +2161,7 @@ def get_text_features( >>> # Pre-compute text embeddings >>> text_inputs = processor(text="cat", return_tensors="pt") - >>> text_embeds = model.get_text_features(**text_inputs) + >>> text_embeds = model.get_text_features(**text_inputs).pooler_output >>> # Reuse text embeddings for multiple images >>> img_url = "http://images.cocodataset.org/val2017/000000077595.jpg" @@ -2180,11 +2170,13 @@ def get_text_features( >>> outputs = model(pixel_values=img_inputs.pixel_values, text_embeds=text_embeds) ``` """ - text_features = self.text_encoder( - input_ids=input_ids, attention_mask=attention_mask, **kwargs - ).last_hidden_state - text_features = self.text_projection(text_features) - return text_features + text_outputs = self.text_encoder( + input_ids=input_ids, attention_mask=attention_mask, return_dict=True, **kwargs + ) + last_hidden_state = text_outputs.last_hidden_state + text_outputs.pooler_output = self.text_projection(last_hidden_state) + + return text_outputs @auto_docstring def get_vision_features( @@ -2193,10 +2185,6 @@ def get_vision_features( **kwargs: Unpack[TransformersKwargs], ) -> Sam3VisionEncoderOutput: r""" - Returns: - vision_embeds (`Sam3VisionEncoderOutput`): - Vision embeddings that can be passed as `vision_embeds` to the forward method. - Example: ```python @@ -2289,7 +2277,9 @@ def forward( fpn_position_encoding = vision_outputs.fpn_position_encoding[:-1] if text_embeds is None: - text_features = self.get_text_features(input_ids=input_ids, attention_mask=attention_mask, **kwargs) + text_features = self.get_text_features( + input_ids=input_ids, attention_mask=attention_mask, return_dict=True + ).pooler_output else: text_features = text_embeds diff --git a/src/transformers/models/sam3_tracker/configuration_sam3_tracker.py b/src/transformers/models/sam3_tracker/configuration_sam3_tracker.py index 28c4bce87eca..b4e75edde397 100644 --- a/src/transformers/models/sam3_tracker/configuration_sam3_tracker.py +++ b/src/transformers/models/sam3_tracker/configuration_sam3_tracker.py @@ -18,7 +18,6 @@ # See the License for the specific language governing permissions and # limitations under the License. - from ...configuration_utils import PreTrainedConfig from ..auto import CONFIG_MAPPING, AutoConfig diff --git a/src/transformers/models/sam3_tracker/modeling_sam3_tracker.py b/src/transformers/models/sam3_tracker/modeling_sam3_tracker.py index 291d0a5f3ddb..78674ca80427 100644 --- a/src/transformers/models/sam3_tracker/modeling_sam3_tracker.py +++ b/src/transformers/models/sam3_tracker/modeling_sam3_tracker.py @@ -18,7 +18,6 @@ # See the License for the specific language governing permissions and # limitations under the License. - from collections.abc import Callable from dataclasses import dataclass @@ -32,10 +31,10 @@ from ... import initialization as init from ...activations import ACT2FN -from ...modeling_outputs import BaseModelOutput +from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel from ...processing_utils import Unpack -from ...utils import ModelOutput, auto_docstring, logging +from ...utils import ModelOutput, auto_docstring, can_return_tuple, logging from ...utils.generic import TransformersKwargs, check_model_inputs, is_flash_attention_requested from ..auto import AutoModel from .configuration_sam3_tracker import Sam3TrackerConfig, Sam3TrackerMaskDecoderConfig, Sam3TrackerPromptEncoderConfig @@ -738,16 +737,10 @@ def _dynamic_multimask_via_stability(self, all_mask_logits, all_iou_scores): @dataclass @auto_docstring(custom_intro="Base class for the vision encoder's outputs.") -class Sam3TrackerVisionEncoderOutput(ModelOutput): +class Sam3TrackerVisionEncoderOutput(BaseModelOutputWithPooling): r""" last_hidden_state (`torch.FloatTensor` of shape `(batch_size, height, width, hidden_size)`): Sequence of hidden-states at the output of the last layer of the model. - fpn_hidden_states (`tuple(torch.FloatTensor)`): - Tuple of `torch.FloatTensor` (one for each feature level, from high to low resolution) of shape - `(batch_size, hidden_size, height, width)`. Feature maps from the Feature Pyramid Network neck. - fpn_position_encoding (`tuple(torch.FloatTensor)`): - Tuple of `torch.FloatTensor` (one for each feature level, from high to low resolution) of shape - `(batch_size, hidden_size, height, width)`. Positional encodings corresponding to the `fpn_hidden_states`. hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, + one for the output of each stage) of shape `(batch_size, height, width, hidden_size)`. Hidden-states of the @@ -756,13 +749,16 @@ class Sam3TrackerVisionEncoderOutput(ModelOutput): Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, sequence_length)`. Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads. + fpn_hidden_states (`tuple(torch.FloatTensor)`): + Tuple of `torch.FloatTensor` (one for each feature level, from high to low resolution) of shape + `(batch_size, hidden_size, height, width)`. Feature maps from the Feature Pyramid Network neck. + fpn_position_encoding (`tuple(torch.FloatTensor)`): + Tuple of `torch.FloatTensor` (one for each feature level, from high to low resolution) of shape + `(batch_size, hidden_size, height, width)`. Positional encodings corresponding to the `fpn_hidden_states`. """ - last_hidden_state: torch.FloatTensor | None = None fpn_hidden_states: torch.FloatTensor | None = None fpn_position_encoding: torch.FloatTensor | None = None - hidden_states: tuple[torch.FloatTensor, ...] | None = None - attentions: tuple[torch.FloatTensor, ...] | None = None @auto_docstring( @@ -841,7 +837,8 @@ def get_image_embeddings( Input pixel values """ batch_size = pixel_values.shape[0] - feature_maps, _, _, _ = self.get_image_features(pixel_values, **kwargs) + image_outputs = self.get_image_features(pixel_values, return_dict=True, **kwargs) + feature_maps = image_outputs.fpn_hidden_states # add no memory embedding to the last feature map feature_maps[-1] = feature_maps[-1] + self.no_memory_embedding @@ -997,10 +994,12 @@ def forward( vision_hidden_states = None if pixel_values is not None: - feature_maps, _, vision_hidden_states, vision_attentions = self.get_image_features( - pixel_values, - **kwargs, + image_outputs: Sam3TrackerVisionEncoderOutput = self.get_image_features( + pixel_values, return_dict=True, **kwargs ) + feature_maps = image_outputs.fpn_hidden_states + vision_hidden_states = image_outputs.hidden_states + vision_attentions = image_outputs.attentions # add no memory embedding to the last feature map feature_maps[-1] = feature_maps[-1] + self.no_memory_embedding @@ -1060,34 +1059,18 @@ def forward( vision_attentions=vision_attentions, ) + @can_return_tuple + @auto_docstring def get_image_features( self, pixel_values: torch.FloatTensor, **kwargs: Unpack[TransformersKwargs], - ) -> tuple[ - list[torch.Tensor], - list[torch.Tensor], - tuple[torch.FloatTensor, ...] | None, - tuple[torch.FloatTensor, ...] | None, - ]: + ) -> tuple | Sam3TrackerVisionEncoderOutput: r""" - Extract and preprocess image features using the vision encoder. - - Args: - pixel_values (`torch.FloatTensor`): - Input pixel values of shape `(batch_size, num_channels, height, width)`. - - Returns: - `tuple`: A tuple containing: - - feature_maps (`list[torch.Tensor]`): List of feature maps from different levels. - - feature_maps_position_embeddings (`list[torch.Tensor]`): List of positional embeddings for each feature level. - - vision_hidden_states (`tuple[torch.FloatTensor]`, *optional*): Hidden states from the vision encoder. - - vision_attentions (`tuple[torch.FloatTensor]`, *optional*): Attention weights from the vision encoder. + pixel_values (`torch.FloatTensor`): + Input pixel values of shape `(batch_size, num_channels, height, width)`. """ - vision_outputs: Sam3TrackerVisionEncoderOutput = self.vision_encoder( - pixel_values, - **kwargs, - ) + vision_outputs: Sam3TrackerVisionEncoderOutput = self.vision_encoder(pixel_values, return_dict=True, **kwargs) feature_maps = vision_outputs.fpn_hidden_states feature_maps_position_embeddings = vision_outputs.fpn_position_encoding @@ -1104,8 +1087,10 @@ def get_image_features( feature_map_position_embedding.flatten(2).permute(2, 0, 1) for feature_map_position_embedding in feature_maps_position_embeddings ] + vision_outputs.fpn_hidden_states = feature_maps + vision_outputs.fpn_position_encoding = feature_maps_position_embeddings - return feature_maps, feature_maps_position_embeddings, vision_outputs.hidden_states, vision_outputs.attentions + return vision_outputs __all__ = ["Sam3TrackerModel", "Sam3TrackerPreTrainedModel"] diff --git a/src/transformers/models/sam3_tracker/modular_sam3_tracker.py b/src/transformers/models/sam3_tracker/modular_sam3_tracker.py index 716f1e655237..acc59c9e5842 100644 --- a/src/transformers/models/sam3_tracker/modular_sam3_tracker.py +++ b/src/transformers/models/sam3_tracker/modular_sam3_tracker.py @@ -12,7 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. - import torch from ... import initialization as init diff --git a/src/transformers/models/sam3_tracker/processing_sam3_tracker.py b/src/transformers/models/sam3_tracker/processing_sam3_tracker.py index d9beb6ee65da..73296fd6c1d4 100644 --- a/src/transformers/models/sam3_tracker/processing_sam3_tracker.py +++ b/src/transformers/models/sam3_tracker/processing_sam3_tracker.py @@ -18,7 +18,6 @@ # See the License for the specific language governing permissions and # limitations under the License. - from copy import deepcopy import numpy as np diff --git a/src/transformers/models/sam3_tracker_video/modeling_sam3_tracker_video.py b/src/transformers/models/sam3_tracker_video/modeling_sam3_tracker_video.py index 1b781d25ec3c..6a07c21dbd18 100644 --- a/src/transformers/models/sam3_tracker_video/modeling_sam3_tracker_video.py +++ b/src/transformers/models/sam3_tracker_video/modeling_sam3_tracker_video.py @@ -36,12 +36,12 @@ from ...activations import ACT2FN from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer -from ...modeling_outputs import BaseModelOutput +from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel from ...processing_utils import Unpack from ...pytorch_utils import compile_compatible_method_lru_cache -from ...utils import ModelOutput, auto_docstring, logging -from ...utils.generic import OutputRecorder, TransformersKwargs, is_flash_attention_requested +from ...utils import ModelOutput, TransformersKwargs, auto_docstring, can_return_tuple, logging +from ...utils.generic import OutputRecorder, is_flash_attention_requested from ..auto import AutoModel from .configuration_sam3_tracker_video import ( Sam3TrackerVideoConfig, @@ -1136,16 +1136,10 @@ def forward( @dataclass @auto_docstring(custom_intro="Base class for the vision encoder's outputs.") -class Sam3TrackerVideoVisionEncoderOutput(ModelOutput): +class Sam3TrackerVideoVisionEncoderOutput(BaseModelOutputWithPooling): r""" last_hidden_state (`torch.FloatTensor` of shape `(batch_size, height, width, hidden_size)`): Sequence of hidden-states at the output of the last layer of the model. - fpn_hidden_states (`tuple(torch.FloatTensor)`): - Tuple of `torch.FloatTensor` (one for each feature level, from high to low resolution) of shape - `(batch_size, hidden_size, height, width)`. Feature maps from the Feature Pyramid Network neck. - fpn_position_encoding (`tuple(torch.FloatTensor)`): - Tuple of `torch.FloatTensor` (one for each feature level, from high to low resolution) of shape - `(batch_size, hidden_size, height, width)`. Positional encodings corresponding to the `fpn_hidden_states`. hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, + one for the output of each stage) of shape `(batch_size, height, width, hidden_size)`. Hidden-states of the @@ -1154,13 +1148,16 @@ class Sam3TrackerVideoVisionEncoderOutput(ModelOutput): Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, sequence_length)`. Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads. + fpn_hidden_states (`tuple(torch.FloatTensor)`): + Tuple of `torch.FloatTensor` (one for each feature level, from high to low resolution) of shape + `(batch_size, hidden_size, height, width)`. Feature maps from the Feature Pyramid Network neck. + fpn_position_encoding (`tuple(torch.FloatTensor)`): + Tuple of `torch.FloatTensor` (one for each feature level, from high to low resolution) of shape + `(batch_size, hidden_size, height, width)`. Positional encodings corresponding to the `fpn_hidden_states`. """ - last_hidden_state: torch.FloatTensor | None = None fpn_hidden_states: torch.FloatTensor | None = None fpn_position_encoding: torch.FloatTensor | None = None - hidden_states: tuple[torch.FloatTensor, ...] | None = None - attentions: tuple[torch.FloatTensor, ...] | None = None class Sam3TrackerVideoPositionalEmbedding(nn.Module): @@ -1687,7 +1684,8 @@ def get_image_embeddings( Input pixel values """ batch_size = pixel_values.shape[0] - feature_maps, _, _, _ = self.get_image_features(pixel_values, **kwargs) + image_outputs = self.get_image_features(pixel_values, return_dict=True, **kwargs) + feature_maps = image_outputs.fpn_hidden_states # add no memory embedding to the last feature map feature_maps[-1] = feature_maps[-1] + self.no_memory_embedding @@ -1858,33 +1856,19 @@ def forward( frame_idx=frame_idx, ) + @can_return_tuple + @auto_docstring def get_image_features( self, pixel_values: torch.FloatTensor, **kwargs: Unpack[TransformersKwargs], - ) -> tuple[ - list[torch.Tensor], - list[torch.Tensor], - tuple[torch.FloatTensor, ...] | None, - tuple[torch.FloatTensor, ...] | None, - ]: + ) -> tuple | Sam3TrackerVideoVisionEncoderOutput: r""" - Extract and preprocess image features using the vision encoder. - - Args: - pixel_values (`torch.FloatTensor`): - Input pixel values of shape `(batch_size, num_channels, height, width)`. - - Returns: - `tuple`: A tuple containing: - - feature_maps (`list[torch.Tensor]`): List of feature maps from different levels. - - feature_maps_position_embeddings (`list[torch.Tensor]`): List of positional embeddings for each feature level. - - vision_hidden_states (`tuple[torch.FloatTensor]`, *optional*): Hidden states from the vision encoder. - - vision_attentions (`tuple[torch.FloatTensor]`, *optional*): Attention weights from the vision encoder. + pixel_values (`torch.FloatTensor`): + Input pixel values of shape `(batch_size, num_channels, height, width)`. """ vision_outputs: Sam3TrackerVideoVisionEncoderOutput = self.vision_encoder( - pixel_values, - **kwargs, + pixel_values, return_dict=True, **kwargs ) feature_maps = vision_outputs.fpn_hidden_states @@ -1902,8 +1886,10 @@ def get_image_features( feature_map_position_embedding.flatten(2).permute(2, 0, 1) for feature_map_position_embedding in feature_maps_position_embeddings[:-1] ] + vision_outputs.fpn_hidden_states = feature_maps + vision_outputs.fpn_position_encoding = feature_maps_position_embeddings - return feature_maps, feature_maps_position_embeddings, vision_outputs.hidden_states, vision_outputs.attentions + return vision_outputs def _prepare_vision_features( self, @@ -1920,7 +1906,9 @@ def _prepare_vision_features( else: # Compute features using image encoder image_batch = inference_session.get_frame(frame_idx).unsqueeze(0) # Add batch dimension - vision_feats, vision_pos_embeds, _, _ = self.get_image_features(image_batch) + image_outputs = self.get_image_features(image_batch, return_dict=True) + vision_feats = image_outputs.fpn_hidden_states + vision_pos_embeds = image_outputs.fpn_position_embeddings # Cache features inference_session.cache.cache_vision_features( frame_idx, {"vision_feats": vision_feats, "vision_pos_embeds": vision_pos_embeds} @@ -2025,10 +2013,10 @@ def _single_frame_forward( vision_hidden_states = None if pixel_values is not None: - feature_maps, _, vision_hidden_states, vision_attentions = self.get_image_features( - pixel_values, - **kwargs, - ) + image_outputs = self.get_image_features(pixel_values, return_dict=True, **kwargs) + feature_maps = image_outputs.fpn_hidden_states + vision_hidden_states = image_outputs.hidden_states + vision_attentions = image_outputs.attentions # add no memory embedding to the last feature map feature_maps[-1] = feature_maps[-1] + self.no_memory_embedding diff --git a/src/transformers/models/sam3_tracker_video/modular_sam3_tracker_video.py b/src/transformers/models/sam3_tracker_video/modular_sam3_tracker_video.py index 17b126f6c167..1d357fe923b0 100644 --- a/src/transformers/models/sam3_tracker_video/modular_sam3_tracker_video.py +++ b/src/transformers/models/sam3_tracker_video/modular_sam3_tracker_video.py @@ -17,7 +17,7 @@ from ...configuration_utils import PreTrainedConfig from ...processing_utils import Unpack -from ...utils.generic import TransformersKwargs +from ...utils import TransformersKwargs, auto_docstring, can_return_tuple from ..auto import CONFIG_MAPPING, AutoConfig, AutoModel from ..sam2_video.configuration_sam2_video import Sam2VideoMaskDecoderConfig, Sam2VideoPromptEncoderConfig from ..sam2_video.modeling_sam2_video import ( @@ -544,33 +544,19 @@ def __init__(self, config: Sam3TrackerVideoConfig, remove_vision_encoder: bool = self.post_init() + @can_return_tuple + @auto_docstring def get_image_features( self, pixel_values: torch.FloatTensor, **kwargs: Unpack[TransformersKwargs], - ) -> tuple[ - list[torch.Tensor], - list[torch.Tensor], - tuple[torch.FloatTensor, ...] | None, - tuple[torch.FloatTensor, ...] | None, - ]: + ) -> tuple | Sam3TrackerVideoVisionEncoderOutput: r""" - Extract and preprocess image features using the vision encoder. - - Args: - pixel_values (`torch.FloatTensor`): - Input pixel values of shape `(batch_size, num_channels, height, width)`. - - Returns: - `tuple`: A tuple containing: - - feature_maps (`list[torch.Tensor]`): List of feature maps from different levels. - - feature_maps_position_embeddings (`list[torch.Tensor]`): List of positional embeddings for each feature level. - - vision_hidden_states (`tuple[torch.FloatTensor]`, *optional*): Hidden states from the vision encoder. - - vision_attentions (`tuple[torch.FloatTensor]`, *optional*): Attention weights from the vision encoder. + pixel_values (`torch.FloatTensor`): + Input pixel values of shape `(batch_size, num_channels, height, width)`. """ vision_outputs: Sam3TrackerVideoVisionEncoderOutput = self.vision_encoder( - pixel_values, - **kwargs, + pixel_values, return_dict=True, **kwargs ) feature_maps = vision_outputs.fpn_hidden_states @@ -588,8 +574,10 @@ def get_image_features( feature_map_position_embedding.flatten(2).permute(2, 0, 1) for feature_map_position_embedding in feature_maps_position_embeddings[:-1] ] + vision_outputs.fpn_hidden_states = feature_maps + vision_outputs.fpn_position_encoding = feature_maps_position_embeddings - return feature_maps, feature_maps_position_embeddings, vision_outputs.hidden_states, vision_outputs.attentions + return vision_outputs __all__ = [ diff --git a/src/transformers/models/siglip/modeling_siglip.py b/src/transformers/models/siglip/modeling_siglip.py index 2001f861b424..657d5e3e4554 100644 --- a/src/transformers/models/siglip/modeling_siglip.py +++ b/src/transformers/models/siglip/modeling_siglip.py @@ -35,7 +35,6 @@ TransformersKwargs, auto_docstring, can_return_tuple, - filter_out_non_signature_kwargs, torch_int, ) from ...utils.generic import check_model_inputs, is_flash_attention_requested @@ -621,10 +620,6 @@ def forward( class SiglipVisionTransformer(SiglipPreTrainedModel): _input_embed_layer = "patch_embedding" - _can_record_outputs = { - "hidden_states": SiglipEncoderLayer, - "attentions": SiglipAttention, - } def __init__(self, config: SiglipVisionConfig): super().__init__(config) @@ -640,7 +635,6 @@ def __init__(self, config: SiglipVisionConfig): self.post_init() - @check_model_inputs(tie_last_hidden_states=False) @auto_docstring def forward( self, @@ -789,19 +783,16 @@ def get_input_embeddings(self) -> nn.Module: def set_input_embeddings(self, value: nn.Module): self.text_model.embeddings.token_embedding = value - @filter_out_non_signature_kwargs() + @can_return_tuple @auto_docstring def get_text_features( self, input_ids: torch.Tensor, attention_mask: torch.Tensor | None = None, position_ids: torch.Tensor | None = None, - ) -> torch.FloatTensor: + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: r""" - Returns: - text_features (`torch.FloatTensor` of shape `(batch_size, output_dim`): The text embeddings obtained by - applying the projection layer to the pooled output of [`SiglipTextModel`]. - Examples: ```python @@ -816,28 +807,22 @@ def get_text_features( >>> with torch.no_grad(): ... text_features = model.get_text_features(**inputs) ```""" - text_outputs: BaseModelOutputWithPooling = self.text_model( + return self.text_model( input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids, + **kwargs, ) - pooled_output = text_outputs.pooler_output - - return pooled_output - @filter_out_non_signature_kwargs() + @can_return_tuple @auto_docstring def get_image_features( self, pixel_values: torch.FloatTensor, interpolate_pos_encoding: bool = False, **kwargs: Unpack[TransformersKwargs], - ) -> torch.FloatTensor: + ) -> tuple | BaseModelOutputWithPooling: r""" - Returns: - image_features (`torch.FloatTensor` of shape `(batch_size, output_dim`): The image embeddings obtained by - applying the projection layer to the pooled output of [`SiglipVisionModel`]. - Examples: ```python @@ -856,14 +841,11 @@ def get_image_features( >>> with torch.no_grad(): ... image_features = model.get_image_features(**inputs) ```""" - vision_outputs: BaseModelOutputWithPooling = self.vision_model( + return self.vision_model( pixel_values=pixel_values, interpolate_pos_encoding=interpolate_pos_encoding, **kwargs, ) - pooled_output = vision_outputs.pooler_output - - return pooled_output # NOTE: SiglipModel uses Pretrained backbones, so we don't need to add `check_model_inputs` here @can_return_tuple diff --git a/src/transformers/models/siglip2/modeling_siglip2.py b/src/transformers/models/siglip2/modeling_siglip2.py index 7db28ba2a577..3182fac5ca8c 100644 --- a/src/transformers/models/siglip2/modeling_siglip2.py +++ b/src/transformers/models/siglip2/modeling_siglip2.py @@ -36,7 +36,7 @@ from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling, ImageClassifierOutput from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel from ...processing_utils import Unpack -from ...utils import ModelOutput, TransformersKwargs, auto_docstring, can_return_tuple, filter_out_non_signature_kwargs +from ...utils import ModelOutput, TransformersKwargs, auto_docstring, can_return_tuple from ...utils.generic import check_model_inputs, is_flash_attention_requested from .configuration_siglip2 import Siglip2Config, Siglip2TextConfig, Siglip2VisionConfig @@ -529,10 +529,6 @@ def forward( class Siglip2VisionTransformer(Siglip2PreTrainedModel): _input_embed_layer = "patch_embedding" - _can_record_outputs = { - "hidden_states": Siglip2EncoderLayer, - "attentions": Siglip2Attention, - } def __init__(self, config: Siglip2VisionConfig): super().__init__(config) @@ -548,7 +544,6 @@ def __init__(self, config: Siglip2VisionConfig): self.post_init() - @check_model_inputs(tie_last_hidden_states=False) @auto_docstring def forward( self, @@ -769,9 +764,7 @@ def forward( pixel_values: torch.FloatTensor, pixel_attention_mask: torch.Tensor, spatial_shapes: torch.LongTensor, - output_attentions: bool | None = None, - output_hidden_states: bool | None = None, - **kwargs, + **kwargs: Unpack[TransformersKwargs], ) -> BaseModelOutputWithPooling: r""" pixel_attention_mask (`torch.Tensor` of shape `(batch_size, image_size, image_size)`, *optional*): @@ -802,8 +795,7 @@ def forward( pixel_values=pixel_values, attention_mask=pixel_attention_mask, spatial_shapes=spatial_shapes, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, + **kwargs, ) @@ -849,19 +841,16 @@ def get_input_embeddings(self) -> nn.Module: def set_input_embeddings(self, value: nn.Module): self.text_model.embeddings.token_embedding = value - @filter_out_non_signature_kwargs() + @can_return_tuple @auto_docstring def get_text_features( self, input_ids: torch.Tensor, attention_mask: torch.Tensor | None = None, position_ids: torch.Tensor | None = None, - ) -> torch.FloatTensor: + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: r""" - Returns: - text_features (`torch.FloatTensor` of shape `(batch_size, output_dim`): The text embeddings obtained by - applying the projection layer to the pooled output of [`Siglip2TextModel`]. - Examples: ```python @@ -876,33 +865,28 @@ def get_text_features( >>> with torch.no_grad(): ... text_features = model.get_text_features(**inputs) ```""" - text_outputs: BaseModelOutputWithPooling = self.text_model( + return self.text_model( input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids, + **kwargs, ) - pooled_output = text_outputs.pooler_output - - return pooled_output - @filter_out_non_signature_kwargs() + @can_return_tuple @auto_docstring def get_image_features( self, pixel_values: torch.FloatTensor | None = None, pixel_attention_mask: torch.Tensor | None = None, spatial_shapes: torch.LongTensor | None = None, - ) -> torch.FloatTensor: + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: r""" pixel_attention_mask (`torch.Tensor` of shape `(batch_size, image_size, image_size)`, *optional*): Mask to avoid performing attention on padding pixel indices. spatial_shapes (`torch.LongTensor` of shape `(batch_size, 2)`): Tensor containing the spatial dimensions (height, width) of the input images. - Returns: - image_features (`torch.FloatTensor` of shape `(batch_size, output_dim`): The image embeddings obtained by - applying the projection layer to the pooled output of [`Siglip2VisionModel`]. - Examples: ```python @@ -922,14 +906,12 @@ def get_image_features( ... image_features = model.get_image_features(**inputs) ``` """ - vision_outputs: BaseModelOutputWithPooling = self.vision_model( + return self.vision_model( pixel_values=pixel_values, attention_mask=pixel_attention_mask, spatial_shapes=spatial_shapes, + **kwargs, ) - pooled_output = vision_outputs.pooler_output - - return pooled_output # NOTE: Siglip2Model uses Pretrained backbones, so we don't need to add `check_model_inputs` here @can_return_tuple diff --git a/src/transformers/models/siglip2/modular_siglip2.py b/src/transformers/models/siglip2/modular_siglip2.py index 830378b5344a..9c022de838e9 100644 --- a/src/transformers/models/siglip2/modular_siglip2.py +++ b/src/transformers/models/siglip2/modular_siglip2.py @@ -34,7 +34,8 @@ ) from ...modeling_attn_mask_utils import _prepare_4d_attention_mask -from ...utils import auto_docstring, filter_out_non_signature_kwargs +from ...processing_utils import Unpack +from ...utils import TransformersKwargs, auto_docstring, can_return_tuple from ...utils.generic import check_model_inputs, is_flash_attention_requested @@ -321,9 +322,7 @@ def forward( pixel_values: torch.FloatTensor, pixel_attention_mask: torch.Tensor, spatial_shapes: torch.LongTensor, - output_attentions: bool | None = None, - output_hidden_states: bool | None = None, - **kwargs, + **kwargs: Unpack[TransformersKwargs], ) -> BaseModelOutputWithPooling: r""" pixel_attention_mask (`torch.Tensor` of shape `(batch_size, image_size, image_size)`, *optional*): @@ -354,31 +353,27 @@ def forward( pixel_values=pixel_values, attention_mask=pixel_attention_mask, spatial_shapes=spatial_shapes, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, + **kwargs, ) class Siglip2Model(SiglipModel): # Update: add `spatial_shapes` and `pixel_attention_mask` - @filter_out_non_signature_kwargs() + @can_return_tuple @auto_docstring def get_image_features( self, pixel_values: torch.FloatTensor | None = None, pixel_attention_mask: torch.Tensor | None = None, spatial_shapes: torch.LongTensor | None = None, - ) -> torch.FloatTensor: + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: r""" pixel_attention_mask (`torch.Tensor` of shape `(batch_size, image_size, image_size)`, *optional*): Mask to avoid performing attention on padding pixel indices. spatial_shapes (`torch.LongTensor` of shape `(batch_size, 2)`): Tensor containing the spatial dimensions (height, width) of the input images. - Returns: - image_features (`torch.FloatTensor` of shape `(batch_size, output_dim`): The image embeddings obtained by - applying the projection layer to the pooled output of [`Siglip2VisionModel`]. - Examples: ```python @@ -398,14 +393,12 @@ def get_image_features( ... image_features = model.get_image_features(**inputs) ``` """ - vision_outputs: BaseModelOutputWithPooling = self.vision_model( + return self.vision_model( pixel_values=pixel_values, attention_mask=pixel_attention_mask, spatial_shapes=spatial_shapes, + **kwargs, ) - pooled_output = vision_outputs.pooler_output - - return pooled_output # Update: add `spatial_shapes` and `pixel_attention_mask` def forward( diff --git a/src/transformers/models/smolvlm/modeling_smolvlm.py b/src/transformers/models/smolvlm/modeling_smolvlm.py index 67665a79af7a..745817caa8b6 100644 --- a/src/transformers/models/smolvlm/modeling_smolvlm.py +++ b/src/transformers/models/smolvlm/modeling_smolvlm.py @@ -31,7 +31,7 @@ from ...masking_utils import create_bidirectional_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer -from ...modeling_outputs import BaseModelOutput, ModelOutput +from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling, ModelOutput from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel from ...processing_utils import Unpack from ...utils import ( @@ -524,17 +524,21 @@ def inputs_merger( merged_embeds = torch.where(image_mask.unsqueeze(-1), image_embeds, inputs_embeds) return merged_embeds + @can_return_tuple + @auto_docstring( + custom_intro="Encodes images into continuous embeddings that can be forwarded to the language model." + ) def get_image_features( - self, pixel_values: torch.FloatTensor, pixel_attention_mask: torch.LongTensor | None = None - ): - """ - Encodes images into continuous embeddings that can be forwarded to the language model. - - Args: - pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): - The tensors corresponding to the input images. - pixel_attention_mask (`torch.LongTensor`, *optional*): - The attention mask indicating padded regions in the image. + self, + pixel_values: torch.FloatTensor, + pixel_attention_mask: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): + The tensors corresponding to the input images. + pixel_attention_mask (`torch.LongTensor`, *optional*): + The attention mask indicating padded regions in the image. """ batch_size, num_images, num_channels, height, width = pixel_values.shape pixel_values = pixel_values.to(dtype=self.dtype) # fp16 compatibility @@ -565,12 +569,16 @@ def get_image_features( patch_attention_mask = (patches_subgrid.sum(dim=(-1, -2)) > 0).bool() # Get sequence from the vision encoder - image_hidden_states = self.vision_model(pixel_values=pixel_values, patch_attention_mask=patch_attention_mask) - image_hidden_states = image_hidden_states.last_hidden_state + image_outputs = self.vision_model( + pixel_values=pixel_values, patch_attention_mask=patch_attention_mask, return_dict=True, **kwargs + ) + image_hidden_states = image_outputs.last_hidden_state # Modality projection & resampling - image_hidden_states = self.connector(image_hidden_states) - return image_hidden_states + image_features = self.connector(image_hidden_states) + image_outputs.pooler_output = image_features + + return image_outputs @can_return_tuple @auto_docstring( @@ -639,7 +647,10 @@ def forward( raise ValueError("You cannot specify both pixel_values and image_hidden_states at the same time") if pixel_values is not None: - image_hidden_states = self.get_image_features(pixel_values, pixel_attention_mask).to(inputs_embeds.device) + image_hidden_states = self.get_image_features( + pixel_values, pixel_attention_mask, return_dict=True + ).pooler_output + image_hidden_states = image_hidden_states.to(inputs_embeds.device) elif image_hidden_states is not None: image_hidden_states = image_hidden_states.to(dtype=self.dtype, device=inputs_embeds.device) @@ -730,10 +741,22 @@ def get_input_embeddings(self): def set_input_embeddings(self, value): self.model.text_model.set_input_embeddings(value) + @auto_docstring def get_image_features( - self, pixel_values: torch.FloatTensor, pixel_attention_mask: torch.LongTensor | None = None - ): - return self.model.get_image_features(pixel_values=pixel_values, pixel_attention_mask=pixel_attention_mask) + self, + pixel_values: torch.FloatTensor, + pixel_attention_mask: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): + The tensors corresponding to the input images. + pixel_attention_mask (`torch.LongTensor`, *optional*): + The attention mask indicating padded regions in the image. + """ + return self.model.get_image_features( + pixel_values=pixel_values, pixel_attention_mask=pixel_attention_mask, **kwargs + ) @can_return_tuple @auto_docstring diff --git a/src/transformers/models/smolvlm/modular_smolvlm.py b/src/transformers/models/smolvlm/modular_smolvlm.py index 0ed5c4fda655..2298499b670c 100644 --- a/src/transformers/models/smolvlm/modular_smolvlm.py +++ b/src/transformers/models/smolvlm/modular_smolvlm.py @@ -19,8 +19,9 @@ from ...cache_utils import Cache, DynamicCache from ...generation import GenerationConfig from ...modeling_flash_attention_utils import FlashAttentionKwargs +from ...modeling_outputs import BaseModelOutputWithPooling from ...processing_utils import Unpack -from ...utils import auto_docstring, can_return_tuple, logging, torch_compilable_check +from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, logging, torch_compilable_check from ..idefics3.configuration_idefics3 import Idefics3Config, Idefics3VisionConfig from ..idefics3.image_processing_idefics3 import Idefics3ImageProcessor from ..idefics3.image_processing_idefics3_fast import Idefics3ImageProcessorFast @@ -191,17 +192,21 @@ def inputs_merger( merged_embeds = torch.where(image_mask.unsqueeze(-1), image_embeds, inputs_embeds) return merged_embeds + @can_return_tuple + @auto_docstring( + custom_intro="Encodes images into continuous embeddings that can be forwarded to the language model." + ) def get_image_features( - self, pixel_values: torch.FloatTensor, pixel_attention_mask: torch.LongTensor | None = None - ): - """ - Encodes images into continuous embeddings that can be forwarded to the language model. - - Args: - pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): - The tensors corresponding to the input images. - pixel_attention_mask (`torch.LongTensor`, *optional*): - The attention mask indicating padded regions in the image. + self, + pixel_values: torch.FloatTensor, + pixel_attention_mask: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): + The tensors corresponding to the input images. + pixel_attention_mask (`torch.LongTensor`, *optional*): + The attention mask indicating padded regions in the image. """ batch_size, num_images, num_channels, height, width = pixel_values.shape pixel_values = pixel_values.to(dtype=self.dtype) # fp16 compatibility @@ -232,12 +237,16 @@ def get_image_features( patch_attention_mask = (patches_subgrid.sum(dim=(-1, -2)) > 0).bool() # Get sequence from the vision encoder - image_hidden_states = self.vision_model(pixel_values=pixel_values, patch_attention_mask=patch_attention_mask) - image_hidden_states = image_hidden_states.last_hidden_state + image_outputs = self.vision_model( + pixel_values=pixel_values, patch_attention_mask=patch_attention_mask, return_dict=True, **kwargs + ) + image_hidden_states = image_outputs.last_hidden_state # Modality projection & resampling - image_hidden_states = self.connector(image_hidden_states) - return image_hidden_states + image_features = self.connector(image_hidden_states) + image_outputs.pooler_output = image_features + + return image_outputs @can_return_tuple @auto_docstring( @@ -300,7 +309,10 @@ def forward( raise ValueError("You cannot specify both pixel_values and image_hidden_states at the same time") if pixel_values is not None: - image_hidden_states = self.get_image_features(pixel_values, pixel_attention_mask).to(inputs_embeds.device) + image_hidden_states = self.get_image_features( + pixel_values, pixel_attention_mask, return_dict=True + ).pooler_output + image_hidden_states = image_hidden_states.to(inputs_embeds.device) elif image_hidden_states is not None: image_hidden_states = image_hidden_states.to(dtype=self.dtype, device=inputs_embeds.device) diff --git a/src/transformers/models/starcoder2/modeling_starcoder2.py b/src/transformers/models/starcoder2/modeling_starcoder2.py index 100da24a3c1d..8613fef45695 100644 --- a/src/transformers/models/starcoder2/modeling_starcoder2.py +++ b/src/transformers/models/starcoder2/modeling_starcoder2.py @@ -364,7 +364,7 @@ def forward( use_cache: bool | None = None, cache_position: torch.LongTensor | None = None, **kwargs: Unpack[TransformersKwargs], - ) -> BaseModelOutputWithPast: + ) -> tuple | BaseModelOutputWithPast: if (input_ids is None) ^ (inputs_embeds is not None): raise ValueError("You must specify exactly one of input_ids or inputs_embeds") diff --git a/src/transformers/models/starcoder2/modular_starcoder2.py b/src/transformers/models/starcoder2/modular_starcoder2.py index 478593ecb60d..f6fd0841b217 100644 --- a/src/transformers/models/starcoder2/modular_starcoder2.py +++ b/src/transformers/models/starcoder2/modular_starcoder2.py @@ -153,7 +153,7 @@ def forward( use_cache: bool | None = None, cache_position: torch.LongTensor | None = None, **kwargs: Unpack[TransformersKwargs], - ) -> BaseModelOutputWithPast: + ) -> tuple | BaseModelOutputWithPast: if (input_ids is None) ^ (inputs_embeds is not None): raise ValueError("You must specify exactly one of input_ids or inputs_embeds") diff --git a/src/transformers/models/switch_transformers/modeling_switch_transformers.py b/src/transformers/models/switch_transformers/modeling_switch_transformers.py index 39ad3ffe95d2..cb30a42ac764 100644 --- a/src/transformers/models/switch_transformers/modeling_switch_transformers.py +++ b/src/transformers/models/switch_transformers/modeling_switch_transformers.py @@ -691,7 +691,7 @@ def forward( use_cache=None, cache_position=None, **kwargs: Unpack[TransformersKwargs], - ): + ) -> tuple | MoEModelOutputWithPastAndCrossAttentions: if (input_ids is None) ^ (inputs_embeds is not None): raise ValueError("You must specify exactly one of input_ids or inputs_embeds") diff --git a/src/transformers/models/switch_transformers/modular_switch_transformers.py b/src/transformers/models/switch_transformers/modular_switch_transformers.py index 5ff5a23826c5..64f91cc20c39 100644 --- a/src/transformers/models/switch_transformers/modular_switch_transformers.py +++ b/src/transformers/models/switch_transformers/modular_switch_transformers.py @@ -447,7 +447,7 @@ def forward( use_cache=None, cache_position=None, **kwargs: Unpack[TransformersKwargs], - ): + ) -> tuple | MoEModelOutputWithPastAndCrossAttentions: if (input_ids is None) ^ (inputs_embeds is not None): raise ValueError("You must specify exactly one of input_ids or inputs_embeds") diff --git a/src/transformers/models/t5gemma/modeling_t5gemma.py b/src/transformers/models/t5gemma/modeling_t5gemma.py index 2fe975e71ba9..f6bc2ca10eae 100644 --- a/src/transformers/models/t5gemma/modeling_t5gemma.py +++ b/src/transformers/models/t5gemma/modeling_t5gemma.py @@ -691,7 +691,7 @@ def forward( position_ids: torch.LongTensor | None = None, inputs_embeds: torch.FloatTensor | None = None, **kwargs: Unpack[TransformersKwargs], - ) -> BaseModelOutput: + ) -> tuple | BaseModelOutput: if (input_ids is None) ^ (inputs_embeds is not None): raise ValueError("You must specify exactly one of input_ids or inputs_embeds") @@ -790,7 +790,7 @@ def forward( encoder_hidden_states: torch.Tensor | None = None, encoder_attention_mask: torch.Tensor | None = None, **kwargs: Unpack[TransformersKwargs], - ) -> BaseModelOutputWithPastAndCrossAttentions: + ) -> tuple | BaseModelOutputWithPastAndCrossAttentions: if (input_ids is None) ^ (inputs_embeds is not None): raise ValueError("You must specify exactly one of input_ids or inputs_embeds") if encoder_hidden_states is None: diff --git a/src/transformers/models/t5gemma/modular_t5gemma.py b/src/transformers/models/t5gemma/modular_t5gemma.py index 5c48354f57e5..c55d20ba7b66 100644 --- a/src/transformers/models/t5gemma/modular_t5gemma.py +++ b/src/transformers/models/t5gemma/modular_t5gemma.py @@ -701,7 +701,7 @@ def forward( position_ids: torch.LongTensor | None = None, inputs_embeds: torch.FloatTensor | None = None, **kwargs: Unpack[TransformersKwargs], - ) -> BaseModelOutput: + ) -> tuple | BaseModelOutput: if (input_ids is None) ^ (inputs_embeds is not None): raise ValueError("You must specify exactly one of input_ids or inputs_embeds") @@ -800,7 +800,7 @@ def forward( encoder_hidden_states: torch.Tensor | None = None, encoder_attention_mask: torch.Tensor | None = None, **kwargs: Unpack[TransformersKwargs], - ) -> BaseModelOutputWithPastAndCrossAttentions: + ) -> tuple | BaseModelOutputWithPastAndCrossAttentions: if (input_ids is None) ^ (inputs_embeds is not None): raise ValueError("You must specify exactly one of input_ids or inputs_embeds") if encoder_hidden_states is None: diff --git a/src/transformers/models/t5gemma2/configuration_t5gemma2.py b/src/transformers/models/t5gemma2/configuration_t5gemma2.py index 87a83649346b..88e6599ab4a3 100644 --- a/src/transformers/models/t5gemma2/configuration_t5gemma2.py +++ b/src/transformers/models/t5gemma2/configuration_t5gemma2.py @@ -609,6 +609,7 @@ def __setattr__(self, key, value): "attention_dropout", "vocab_size", "dtype", + "return_dict", ] if key in shared_attr_with_submodules: diff --git a/src/transformers/models/t5gemma2/modeling_t5gemma2.py b/src/transformers/models/t5gemma2/modeling_t5gemma2.py index f350ebe0aeaf..da8bdc4905bc 100644 --- a/src/transformers/models/t5gemma2/modeling_t5gemma2.py +++ b/src/transformers/models/t5gemma2/modeling_t5gemma2.py @@ -36,6 +36,7 @@ from ...modeling_outputs import ( BaseModelOutput, BaseModelOutputWithPastAndCrossAttentions, + BaseModelOutputWithPooling, Seq2SeqLMOutput, Seq2SeqModelOutput, SequenceClassifierOutput, @@ -806,13 +807,19 @@ def __init__( # Initialize weights and apply final processing self.post_init() - def get_image_features(self, pixel_values: torch.Tensor) -> torch.Tensor: - """Convert pixel image to image features via the encoder and projector.""" + @can_return_tuple + @auto_docstring + def get_image_features( + self, pixel_values: torch.Tensor, **kwargs: Unpack[TransformersKwargs] + ) -> tuple | BaseModelOutputWithPooling: # pixel_values: (batch_size, channels, height, width) # image_features: Image feature tensor of shape (num_images, image_length, embed_dim). - vision_outputs = self.vision_tower(pixel_values=pixel_values).last_hidden_state - image_features = self.multi_modal_projector(vision_outputs) - return image_features + vision_outputs = self.vision_tower(pixel_values=pixel_values, return_dict=True, **kwargs) + last_hidden_state = vision_outputs.last_hidden_state + image_features = self.multi_modal_projector(last_hidden_state) + vision_outputs.pooler_output = image_features + + return vision_outputs def get_image_placeholder_mask( self, @@ -851,7 +858,7 @@ def preprocess_image_features( inputs_embeds: torch.FloatTensor | None = None, ): """Convert pixel images to image features and merge into input embeds.""" - image_features = self.get_image_features(pixel_values) + image_features = self.get_image_features(pixel_values, return_dict=True).pooler_output image_features = image_features.to(inputs_embeds.device, inputs_embeds.dtype) image_mask = self.get_image_placeholder_mask( @@ -1216,8 +1223,11 @@ def get_encoder(self): def get_decoder(self): return self.model.get_decoder() - def get_image_features(self, pixel_values): - return self.get_encoder().get_image_features(pixel_values) + @auto_docstring + def get_image_features( + self, pixel_values: torch.Tensor, **kwargs: Unpack[TransformersKwargs] + ) -> tuple | BaseModelOutputWithPooling: + return self.get_encoder().get_image_features(pixel_values, **kwargs) @property def vision_tower(self): diff --git a/src/transformers/models/t5gemma2/modular_t5gemma2.py b/src/transformers/models/t5gemma2/modular_t5gemma2.py index 47932d547e08..3245a041d20b 100644 --- a/src/transformers/models/t5gemma2/modular_t5gemma2.py +++ b/src/transformers/models/t5gemma2/modular_t5gemma2.py @@ -28,6 +28,7 @@ from ...modeling_outputs import ( BaseModelOutput, BaseModelOutputWithPastAndCrossAttentions, + BaseModelOutputWithPooling, Seq2SeqLMOutput, Seq2SeqModelOutput, SequenceClassifierOutput, @@ -482,6 +483,7 @@ def __setattr__(self, key, value): "attention_dropout", "vocab_size", "dtype", + "return_dict", ] if key in shared_attr_with_submodules: @@ -857,13 +859,19 @@ def __init__( # Initialize weights and apply final processing self.post_init() - def get_image_features(self, pixel_values: torch.Tensor) -> torch.Tensor: - """Convert pixel image to image features via the encoder and projector.""" + @can_return_tuple + @auto_docstring + def get_image_features( + self, pixel_values: torch.Tensor, **kwargs: Unpack[TransformersKwargs] + ) -> tuple | BaseModelOutputWithPooling: # pixel_values: (batch_size, channels, height, width) # image_features: Image feature tensor of shape (num_images, image_length, embed_dim). - vision_outputs = self.vision_tower(pixel_values=pixel_values).last_hidden_state - image_features = self.multi_modal_projector(vision_outputs) - return image_features + vision_outputs = self.vision_tower(pixel_values=pixel_values, return_dict=True, **kwargs) + last_hidden_state = vision_outputs.last_hidden_state + image_features = self.multi_modal_projector(last_hidden_state) + vision_outputs.pooler_output = image_features + + return vision_outputs def get_image_placeholder_mask( self, @@ -902,7 +910,7 @@ def preprocess_image_features( inputs_embeds: torch.FloatTensor | None = None, ): """Convert pixel images to image features and merge into input embeds.""" - image_features = self.get_image_features(pixel_values) + image_features = self.get_image_features(pixel_values, return_dict=True).pooler_output image_features = image_features.to(inputs_embeds.device, inputs_embeds.dtype) image_mask = self.get_image_placeholder_mask( @@ -1254,8 +1262,11 @@ def get_encoder(self): def get_decoder(self): return self.model.get_decoder() - def get_image_features(self, pixel_values): - return self.get_encoder().get_image_features(pixel_values) + @auto_docstring + def get_image_features( + self, pixel_values: torch.Tensor, **kwargs: Unpack[TransformersKwargs] + ) -> tuple | BaseModelOutputWithPooling: + return self.get_encoder().get_image_features(pixel_values, **kwargs) @property def vision_tower(self): diff --git a/src/transformers/models/timm_wrapper/modeling_timm_wrapper.py b/src/transformers/models/timm_wrapper/modeling_timm_wrapper.py index 461ce2ab30ba..0d944fc02421 100644 --- a/src/transformers/models/timm_wrapper/modeling_timm_wrapper.py +++ b/src/transformers/models/timm_wrapper/modeling_timm_wrapper.py @@ -168,6 +168,7 @@ def forward( output_hidden_states: bool | list[int] | None = None, return_dict: bool | None = None, do_pooling: bool | None = None, + use_cache: bool | None = None, **kwargs, ) -> TimmWrapperModelOutput | tuple[Tensor, ...]: r""" diff --git a/src/transformers/models/video_llama_3/modeling_video_llama_3.py b/src/transformers/models/video_llama_3/modeling_video_llama_3.py index 2a35bdf2b68f..d19162af17a3 100644 --- a/src/transformers/models/video_llama_3/modeling_video_llama_3.py +++ b/src/transformers/models/video_llama_3/modeling_video_llama_3.py @@ -30,7 +30,7 @@ from ...cache_utils import Cache from ...generation import GenerationMixin from ...modeling_layers import GradientCheckpointingLayer -from ...modeling_outputs import BaseModelOutput, ModelOutput +from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling, ModelOutput from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel from ...processing_utils import Unpack from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, torch_compilable_check @@ -546,54 +546,62 @@ def get_input_embeddings(self): def set_input_embeddings(self, value): self.language_model.set_input_embeddings(value) + @can_return_tuple + @auto_docstring def get_video_features( self, pixel_values_videos: torch.FloatTensor, video_grid_thw: torch.LongTensor, video_merge_sizes: torch.LongTensor, - ): - """ - Encodes videos into continuous embeddings that can be forwarded to the language model. - - Args: - pixel_values_videos (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): - The tensors corresponding to the input videos. - video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*): - The temporal, height and width of feature shape of each video in LLM. - video_merge_sizes (`torch.Tensor` of shape `(num_videos,)`): - The spatial downsampling ratio of each video feature. + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + pixel_values_videos (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): + The tensors corresponding to the input videos. + video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*): + The temporal, height and width of feature shape of each video in LLM. + video_merge_sizes (`torch.Tensor` of shape `(num_videos,)`): + The spatial downsampling ratio of each video feature. """ - return self.get_image_features(pixel_values_videos, video_grid_thw, video_merge_sizes) + return self.get_image_features( + pixel_values=pixel_values_videos, + image_grid_thw=video_grid_thw, + image_merge_sizes=video_merge_sizes, + **kwargs, + ) + @can_return_tuple + @auto_docstring def get_image_features( self, pixel_values: torch.FloatTensor, image_grid_thw: torch.LongTensor, image_merge_sizes: torch.LongTensor, - ): - """ - Encodes images into continuous embeddings that can be forwarded to the language model. - - Args: - pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): - The tensors corresponding to the input images. - image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): - The temporal, height and width of feature shape of each image in LLM. - image_merge_sizes (`torch.Tensor` of shape `(num_images,)`): - The spatial downsampling ratio of each image feature. + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): + The tensors corresponding to the input images. + image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): + The temporal, height and width of feature shape of each image in LLM. + image_merge_sizes (`torch.Tensor` of shape `(num_images,)`): + The spatial downsampling ratio of each image feature. """ - image_embeds = self.vision_model( + vision_outputs = self.vision_model( pixel_values=pixel_values, grid_thw=image_grid_thw, merge_sizes=image_merge_sizes, return_dict=True, - ).last_hidden_state - image_embeds = self.projector(image_embeds) + **kwargs, + ) + last_hidden_state = vision_outputs.last_hidden_state + image_embeds = self.projector(last_hidden_state) split_sizes = image_grid_thw.prod(dim=1) // (image_merge_sizes**2) image_embeds = torch.split(image_embeds, split_sizes.tolist()) + vision_outputs.pooler_output = image_embeds - return image_embeds + return vision_outputs def get_placeholder_mask( self, @@ -673,7 +681,9 @@ def forward( image_embeds = None if pixel_values is not None: - image_embeds = self.get_image_features(pixel_values, image_grid_thw, image_merge_sizes) + image_embeds = self.get_image_features( + pixel_values, image_grid_thw, image_merge_sizes, return_dict=True + ).pooler_output image_embeds = torch.cat(image_embeds, dim=0).to(inputs_embeds.device, inputs_embeds.dtype) image_mask, _ = self.get_placeholder_mask( input_ids, inputs_embeds=inputs_embeds, image_features=image_embeds @@ -682,7 +692,9 @@ def forward( video_embeds = None if pixel_values_videos is not None: - video_embeds = self.get_video_features(pixel_values_videos, video_grid_thw, video_merge_sizes) + video_embeds = self.get_video_features( + pixel_values_videos, video_grid_thw, video_merge_sizes, return_dict=True + ).pooler_output video_embeds = torch.cat(video_embeds, dim=0).to(inputs_embeds.device, inputs_embeds.dtype) if video_compression_mask is not None: video_embeds = video_embeds[video_compression_mask.to(video_embeds.device)] @@ -765,13 +777,37 @@ def get_input_embeddings(self): def set_input_embeddings(self, value): self.model.set_input_embeddings(value) + @auto_docstring def get_video_features( - self, pixel_values_videos: torch.FloatTensor, video_grid_thw: torch.LongTensor | None = None - ): - return self.model.get_video_features(pixel_values_videos, video_grid_thw) + self, + pixel_values_videos: torch.FloatTensor, + video_grid_thw: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + pixel_values_videos (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): + The tensors corresponding to the input videos. + video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*): + The temporal, height and width of feature shape of each video in LLM. + """ + return self.model.get_video_features( + pixel_values_videos=pixel_values_videos, video_grid_thw=video_grid_thw, **kwargs + ) - def get_image_features(self, pixel_values: torch.FloatTensor, image_grid_thw: torch.LongTensor | None = None): - return self.model.get_image_features(pixel_values, image_grid_thw) + @auto_docstring + def get_image_features( + self, + pixel_values: torch.FloatTensor, + image_grid_thw: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): + The tensors corresponding to the input images. + image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): + The temporal, height and width of feature shape of each image in LLM. + """ + return self.model.get_image_features(pixel_values=pixel_values, image_grid_thw=image_grid_thw, **kwargs) @can_return_tuple @auto_docstring diff --git a/src/transformers/models/video_llama_3/modular_video_llama_3.py b/src/transformers/models/video_llama_3/modular_video_llama_3.py index a9629f90b53d..caeb8483afdd 100644 --- a/src/transformers/models/video_llama_3/modular_video_llama_3.py +++ b/src/transformers/models/video_llama_3/modular_video_llama_3.py @@ -37,7 +37,7 @@ valid_images, validate_preprocess_arguments, ) -from ...modeling_outputs import BaseModelOutput, ModelOutput +from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling, ModelOutput from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel from ...processing_utils import Unpack from ...tokenization_utils_base import PreTokenizedInput, TextInput @@ -594,54 +594,62 @@ def __init__(self, config: VideoLlama3Config): def get_rope_index(self): raise AttributeError("Not needed for VideoLLaMA3") + @can_return_tuple + @auto_docstring def get_video_features( self, pixel_values_videos: torch.FloatTensor, video_grid_thw: torch.LongTensor, video_merge_sizes: torch.LongTensor, - ): - """ - Encodes videos into continuous embeddings that can be forwarded to the language model. - - Args: - pixel_values_videos (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): - The tensors corresponding to the input videos. - video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*): - The temporal, height and width of feature shape of each video in LLM. - video_merge_sizes (`torch.Tensor` of shape `(num_videos,)`): - The spatial downsampling ratio of each video feature. + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + pixel_values_videos (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): + The tensors corresponding to the input videos. + video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*): + The temporal, height and width of feature shape of each video in LLM. + video_merge_sizes (`torch.Tensor` of shape `(num_videos,)`): + The spatial downsampling ratio of each video feature. """ - return self.get_image_features(pixel_values_videos, video_grid_thw, video_merge_sizes) + return self.get_image_features( + pixel_values=pixel_values_videos, + image_grid_thw=video_grid_thw, + image_merge_sizes=video_merge_sizes, + **kwargs, + ) + @can_return_tuple + @auto_docstring def get_image_features( self, pixel_values: torch.FloatTensor, image_grid_thw: torch.LongTensor, image_merge_sizes: torch.LongTensor, - ): - """ - Encodes images into continuous embeddings that can be forwarded to the language model. - - Args: - pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): - The tensors corresponding to the input images. - image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): - The temporal, height and width of feature shape of each image in LLM. - image_merge_sizes (`torch.Tensor` of shape `(num_images,)`): - The spatial downsampling ratio of each image feature. + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): + The tensors corresponding to the input images. + image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): + The temporal, height and width of feature shape of each image in LLM. + image_merge_sizes (`torch.Tensor` of shape `(num_images,)`): + The spatial downsampling ratio of each image feature. """ - image_embeds = self.vision_model( + vision_outputs = self.vision_model( pixel_values=pixel_values, grid_thw=image_grid_thw, merge_sizes=image_merge_sizes, return_dict=True, - ).last_hidden_state - image_embeds = self.projector(image_embeds) + **kwargs, + ) + last_hidden_state = vision_outputs.last_hidden_state + image_embeds = self.projector(last_hidden_state) split_sizes = image_grid_thw.prod(dim=1) // (image_merge_sizes**2) image_embeds = torch.split(image_embeds, split_sizes.tolist()) + vision_outputs.pooler_output = image_embeds - return image_embeds + return vision_outputs @can_return_tuple def forward( @@ -680,7 +688,9 @@ def forward( image_embeds = None if pixel_values is not None: - image_embeds = self.get_image_features(pixel_values, image_grid_thw, image_merge_sizes) + image_embeds = self.get_image_features( + pixel_values, image_grid_thw, image_merge_sizes, return_dict=True + ).pooler_output image_embeds = torch.cat(image_embeds, dim=0).to(inputs_embeds.device, inputs_embeds.dtype) image_mask, _ = self.get_placeholder_mask( input_ids, inputs_embeds=inputs_embeds, image_features=image_embeds @@ -689,7 +699,9 @@ def forward( video_embeds = None if pixel_values_videos is not None: - video_embeds = self.get_video_features(pixel_values_videos, video_grid_thw, video_merge_sizes) + video_embeds = self.get_video_features( + pixel_values_videos, video_grid_thw, video_merge_sizes, return_dict=True + ).pooler_output video_embeds = torch.cat(video_embeds, dim=0).to(inputs_embeds.device, inputs_embeds.dtype) if video_compression_mask is not None: video_embeds = video_embeds[video_compression_mask.to(video_embeds.device)] diff --git a/src/transformers/models/video_llava/modeling_video_llava.py b/src/transformers/models/video_llava/modeling_video_llava.py index 0742205829f3..fe86cbb2512f 100644 --- a/src/transformers/models/video_llava/modeling_video_llava.py +++ b/src/transformers/models/video_llava/modeling_video_llava.py @@ -23,10 +23,11 @@ from ...cache_utils import Cache from ...generation import GenerationMixin from ...modeling_flash_attention_utils import FlashAttentionKwargs -from ...modeling_outputs import ModelOutput +from ...modeling_outputs import BaseModelOutputWithPooling, ModelOutput from ...modeling_utils import PreTrainedModel from ...processing_utils import Unpack -from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, logging, torch_compilable_check +from ...utils import TransformersKwargs, auto_docstring, logging, torch_compilable_check +from ...utils.generic import check_model_inputs from ..auto import AutoModel from .configuration_video_llava import VideoLlavaConfig @@ -169,86 +170,82 @@ def get_input_embeddings(self): def set_input_embeddings(self, value): self.language_model.set_input_embeddings(value) + @check_model_inputs(tie_last_hidden_states=False) + @auto_docstring( + custom_intro="Obtains image last hidden states from the vision tower and apply multimodal projection." + ) def get_image_features( self, pixel_values_images: torch.FloatTensor, vision_feature_layer: int | list[int] | None = None, vision_feature_select_strategy: str | None = None, - ): - """ - Obtains image last hidden states from the vision tower and apply multimodal projection. - - Args: - pixel_values_images (`torch.FloatTensor]` of shape `(batch_size, channels, height, width)`) - The tensors corresponding to the input images. - vision_feature_layer (`Union[int, list[int]]`, *optional*): - The index of the layer to select the vision feature. If multiple indices are provided, - the vision feature of the corresponding indices will be concatenated to form the - vision features. - vision_feature_select_strategy (`str`, *optional*): - The feature selection strategy used to select the vision feature from the vision backbone. - Can be one of `"default"` or `"full"` - Returns: - image_features (`torch.Tensor`): Image feature tensor of shape `(num_images, image_length, embed_dim)`). + output_hidden_states: bool | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + pixel_values_images (`torch.FloatTensor]` of shape `(batch_size, channels, height, width)`) + The tensors corresponding to the input images. + vision_feature_layer (`Union[int, list[int]]`, *optional*): + The index of the layer to select the vision feature. If multiple indices are provided, + the vision feature of the corresponding indices will be concatenated to form the + vision features. + vision_feature_select_strategy (`str`, *optional*): + The feature selection strategy used to select the vision feature from the vision backbone. + Can be one of `"default"` or `"full"` """ - vision_feature_layer = ( - vision_feature_layer if vision_feature_layer is not None else self.config.vision_feature_layer - ) - vision_feature_select_strategy = ( - vision_feature_select_strategy - if vision_feature_select_strategy is not None - else self.config.vision_feature_select_strategy + image_outputs = self.image_tower( + pixel_values_images, + output_hidden_states=True, # Ignore arg on purpose + return_dict=True, + **kwargs, ) - if vision_feature_select_strategy not in ["default", "full"]: - raise ValueError(f"Unexpected select feature strategy: {self.config.vision_feature_select_strategy}") - - image_outputs = self.image_tower(pixel_values_images, output_hidden_states=True) - # If we have one vision feature layer, return the corresponding hidden states, # otherwise, select the hidden states of each feature layer and concatenate them if isinstance(vision_feature_layer, int): - image_outputs = image_outputs.hidden_states[vision_feature_layer] + selected_hidden_state = image_outputs.hidden_states[vision_feature_layer] if vision_feature_select_strategy == "default": - image_outputs = image_outputs[:, 1:] + selected_hidden_state = selected_hidden_state[:, 1:] else: hs_pool = [image_outputs.hidden_states[layer_idx] for layer_idx in vision_feature_layer] # For default; crop CLS from each hidden state in the hidden state pool if vision_feature_select_strategy == "default": hs_pool = [hs[:, 1:] for hs in hs_pool] - image_outputs = torch.cat(hs_pool, dim=-1) + selected_hidden_state = torch.cat(hs_pool, dim=-1) - image_features = self.multi_modal_projector(image_outputs) + image_features = self.multi_modal_projector(selected_hidden_state) + image_outputs.pooler_output = image_features - return image_features + return image_outputs + @check_model_inputs(tie_last_hidden_states=False) + @auto_docstring( + custom_intro="Obtains video last hidden states from the vision tower and apply multimodal projection." + ) def get_video_features( self, pixel_values_videos: torch.FloatTensor, vision_feature_layer: int | list[int] | None = None, - ): - """ - Obtains video last hidden states from the vision tower and apply multimodal projection. - - Args: - pixel_values_videos (`torch.FloatTensor]` of shape `(batch_size, num_frames, channels, height, width)`) - The tensors corresponding to the input videos. - vision_feature_layer (`Union[int, list[int]]`, *optional*): - The index of the layer to select the vision feature. If multiple indices are provided, - the vision feature of the corresponding indices will be concatenated to form the - vision features. - Returns: - video_features (`torch.Tensor`): Video feature tensor of shape `(num_videos * num_frames, image_length, embed_dim)`). - frames (`int`): Number of frames the videos have. + output_hidden_states: bool | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + pixel_values_videos (`torch.FloatTensor]` of shape `(batch_size, num_frames, channels, height, width)`) + The tensors corresponding to the input videos. + vision_feature_layer (`Union[int, list[int]]`, *optional*): + The index of the layer to select the vision feature. If multiple indices are provided, + the vision feature of the corresponding indices will be concatenated to form the + vision features. """ - vision_feature_layer = ( - vision_feature_layer if vision_feature_layer is not None else self.config.vision_feature_layer - ) - batch_size_vid, num_frames, channels, height, width = pixel_values_videos.shape pixel_values = pixel_values_videos.reshape(batch_size_vid * num_frames, channels, height, width) - video_outputs = self.video_tower(pixel_values, output_hidden_states=True) + video_outputs = self.video_tower( + pixel_values, + output_hidden_states=True, # Ignore arg on purpose + return_dict=True, + **kwargs, + ) # If we have one vision feature layer, return the corresponding hidden states, # otherwise, select the hidden states of each feature layer and concatenate them @@ -259,8 +256,9 @@ def get_video_features( video_features = torch.cat(hs_pool, dim=-1) video_features = self.multi_modal_projector(video_features) + video_outputs.pooler_output = video_features - return video_features, num_frames + return video_outputs def get_placeholder_mask( self, @@ -303,7 +301,7 @@ def get_placeholder_mask( ) return special_image_mask, special_video_mask - @can_return_tuple + @check_model_inputs(tie_last_hidden_states=False) @auto_docstring def forward( self, @@ -334,14 +332,6 @@ def forward( output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states ) return_dict = return_dict if return_dict is not None else self.config.use_return_dict - vision_feature_layer = ( - vision_feature_layer if vision_feature_layer is not None else self.config.vision_feature_layer - ) - vision_feature_select_strategy = ( - vision_feature_select_strategy - if vision_feature_select_strategy is not None - else self.config.vision_feature_select_strategy - ) if (input_ids is None) ^ (inputs_embeds is not None): raise ValueError("You must specify exactly one of input_ids or inputs_embeds") @@ -354,7 +344,8 @@ def forward( pixel_values_images, vision_feature_layer=vision_feature_layer, vision_feature_select_strategy=vision_feature_select_strategy, - ) + return_dict=True, + ).pooler_output image_features = image_features.to(inputs_embeds.device, inputs_embeds.dtype) special_image_mask, _ = self.get_placeholder_mask( input_ids, inputs_embeds=inputs_embeds, image_features=image_features @@ -362,9 +353,9 @@ def forward( inputs_embeds = inputs_embeds.masked_scatter(special_image_mask, image_features) if pixel_values_videos is not None: - video_features, num_frames = self.get_video_features( - pixel_values_videos=pixel_values_videos, vision_feature_layer=vision_feature_layer - ) + video_features = self.get_video_features( + pixel_values_videos=pixel_values_videos, vision_feature_layer=vision_feature_layer, return_dict=True + ).pooler_output video_features = video_features.to(inputs_embeds.device, inputs_embeds.dtype) _, special_video_mask = self.get_placeholder_mask( input_ids, inputs_embeds=inputs_embeds, video_features=video_features @@ -424,19 +415,33 @@ def set_input_embeddings(self, value): def get_output_embeddings(self) -> nn.Module: return self.lm_head + @auto_docstring def get_image_features( self, pixel_values_images: torch.FloatTensor, vision_feature_layer: int | list[int] | None = None, vision_feature_select_strategy: str | None = None, - ): + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + pixel_values_images (`torch.FloatTensor]` of shape `(batch_size, channels, height, width)`) + The tensors corresponding to the input images. + vision_feature_layer (`Union[int, list[int]]`, *optional*): + The index of the layer to select the vision feature. If multiple indices are provided, + the vision feature of the corresponding indices will be concatenated to form the + vision features. + vision_feature_select_strategy (`str`, *optional*): + The feature selection strategy used to select the vision feature from the vision backbone. + Can be one of `"default"` or `"full"` + """ return self.model.get_image_features( pixel_values_images=pixel_values_images, vision_feature_layer=vision_feature_layer, vision_feature_select_strategy=vision_feature_select_strategy, + **kwargs, ) - @can_return_tuple + @check_model_inputs(tie_last_hidden_states=False) @auto_docstring def forward( self, @@ -478,7 +483,6 @@ def forward( >>> from huggingface_hub import hf_hub_download >>> from transformers import VideoLlavaProcessor, VideoLlavaForConditionalGeneration - >>> def read_video_pyav(container, indices): ... ''' ... Decode the video with PyAV decoder. @@ -539,14 +543,6 @@ def forward( output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states ) return_dict = return_dict if return_dict is not None else self.config.use_return_dict - vision_feature_layer = ( - vision_feature_layer if vision_feature_layer is not None else self.config.vision_feature_layer - ) - vision_feature_select_strategy = ( - vision_feature_select_strategy - if vision_feature_select_strategy is not None - else self.config.vision_feature_select_strategy - ) outputs = self.model( input_ids=input_ids, diff --git a/src/transformers/models/vipllava/modeling_vipllava.py b/src/transformers/models/vipllava/modeling_vipllava.py index e0dc823af7b7..a5d287c94b31 100644 --- a/src/transformers/models/vipllava/modeling_vipllava.py +++ b/src/transformers/models/vipllava/modeling_vipllava.py @@ -18,7 +18,6 @@ # See the License for the specific language governing permissions and # limitations under the License. - from dataclasses import dataclass import torch @@ -27,9 +26,11 @@ from ...activations import ACT2FN from ...cache_utils import Cache from ...generation import GenerationMixin -from ...modeling_outputs import BaseModelOutputWithPast, ModelOutput +from ...modeling_outputs import BaseModelOutputWithPast, BaseModelOutputWithPooling, ModelOutput from ...modeling_utils import PreTrainedModel -from ...utils import auto_docstring, can_return_tuple, torch_compilable_check +from ...processing_utils import Unpack +from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, torch_compilable_check +from ...utils.generic import check_model_inputs from ..auto import AutoModel from .configuration_vipllava import VipLlavaConfig @@ -149,25 +150,33 @@ def get_input_embeddings(self): def set_input_embeddings(self, value): self.language_model.set_input_embeddings(value) + @can_return_tuple + @auto_docstring( + custom_intro="Obtains image last hidden states from the vision tower and apply multimodal projection." + ) def get_image_features( - self, pixel_values: torch.FloatTensor, vision_feature_layers: int | list[int] | None = None - ): - """ - Obtains image last hidden states from the vision tower and apply multimodal projection. - - Args: - pixel_values (`torch.FloatTensor]` of shape `(batch_size, channels, height, width)`) - The tensors corresponding to the input images. - vision_feature_layers (`Union[int, list[int]]`): - The vision feature layer, or the list of indexes of the layers to select - the vision feature. - Returns: - image_features (`torch.Tensor`): Image feature tensor of shape `(num_images, image_length, embed_dim)`). + self, + pixel_values: torch.FloatTensor, + vision_feature_layers: int | list[int] | None = None, + output_hidden_states: bool | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + pixel_values (`torch.FloatTensor]` of shape `(batch_size, channels, height, width)`): + The tensors corresponding to the input images. + vision_feature_layers (`Union[int, list[int]]`, *optional*): + The vision feature layer, or the list of indexes of the layers to select + the vision feature. """ vision_feature_layers = ( vision_feature_layers if vision_feature_layers is not None else self.config.vision_feature_layers ) - image_outputs = self.vision_tower(pixel_values, output_hidden_states=True) + image_outputs = self.vision_tower( + pixel_values, + output_hidden_states=True, # Ignore arg on purpose + return_dict=True, + **kwargs, + ) # If multiple feature layers are provided (which is usually the case) # then the image features are concatenated after the CLS is removed. @@ -178,7 +187,9 @@ def get_image_features( image_features = [image_outputs.hidden_states[index][:, 1:] for index in vision_feature_layers] image_features = torch.cat(image_features, dim=-1) image_features = self.multi_modal_projector(image_features) - return image_features + image_outputs.pooler_output = image_features + + return image_outputs def get_placeholder_mask( self, input_ids: torch.LongTensor, inputs_embeds: torch.FloatTensor, image_features: torch.FloatTensor @@ -243,8 +254,8 @@ def forward( if pixel_values is not None: image_features = self.get_image_features( - pixel_values=pixel_values, vision_feature_layers=vision_feature_layers - ) + pixel_values=pixel_values, vision_feature_layers=vision_feature_layers, return_dict=True + ).pooler_output image_features = image_features.to(inputs_embeds.device, inputs_embeds.dtype) special_image_mask = self.get_placeholder_mask( input_ids, inputs_embeds=inputs_embeds, image_features=image_features @@ -303,12 +314,25 @@ def set_input_embeddings(self, value): def get_output_embeddings(self) -> nn.Module: return self.lm_head + @auto_docstring def get_image_features( - self, pixel_values: torch.FloatTensor, vision_feature_layers: int | list[int] | None = None - ): - return self.model.get_image_features(pixel_values=pixel_values, vision_feature_layers=vision_feature_layers) + self, + pixel_values: torch.FloatTensor, + vision_feature_layers: int | list[int] | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + pixel_values (`torch.FloatTensor]` of shape `(batch_size, channels, height, width)`): + The tensors corresponding to the input images. + vision_feature_layers (`Union[int, list[int]]`, *optional*): + The vision feature layer, or the list of indexes of the layers to select + the vision feature. + """ + return self.model.get_image_features( + pixel_values=pixel_values, vision_feature_layers=vision_feature_layers, **kwargs + ) - @can_return_tuple + @check_model_inputs(tie_last_hidden_states=False) @auto_docstring def forward( self, diff --git a/src/transformers/models/vipllava/modular_vipllava.py b/src/transformers/models/vipllava/modular_vipllava.py index b8b033b83d35..3b9a2dfdf284 100644 --- a/src/transformers/models/vipllava/modular_vipllava.py +++ b/src/transformers/models/vipllava/modular_vipllava.py @@ -12,7 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. - import torch from torch import nn @@ -26,7 +25,9 @@ from ...activations import ACT2FN from ...cache_utils import Cache -from ...utils import auto_docstring, logging +from ...modeling_outputs import BaseModelOutputWithPooling +from ...processing_utils import Unpack +from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, logging from .configuration_vipllava import VipLlavaConfig @@ -70,25 +71,33 @@ class VipLlavaPreTrainedModel(LlavaPreTrainedModel): class VipLlavaModel(LlavaModel): + @can_return_tuple + @auto_docstring( + custom_intro="Obtains image last hidden states from the vision tower and apply multimodal projection." + ) def get_image_features( - self, pixel_values: torch.FloatTensor, vision_feature_layers: int | list[int] | None = None - ): - """ - Obtains image last hidden states from the vision tower and apply multimodal projection. - - Args: - pixel_values (`torch.FloatTensor]` of shape `(batch_size, channels, height, width)`) - The tensors corresponding to the input images. - vision_feature_layers (`Union[int, list[int]]`): - The vision feature layer, or the list of indexes of the layers to select - the vision feature. - Returns: - image_features (`torch.Tensor`): Image feature tensor of shape `(num_images, image_length, embed_dim)`). + self, + pixel_values: torch.FloatTensor, + vision_feature_layers: int | list[int] | None = None, + output_hidden_states: bool | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + pixel_values (`torch.FloatTensor]` of shape `(batch_size, channels, height, width)`): + The tensors corresponding to the input images. + vision_feature_layers (`Union[int, list[int]]`, *optional*): + The vision feature layer, or the list of indexes of the layers to select + the vision feature. """ vision_feature_layers = ( vision_feature_layers if vision_feature_layers is not None else self.config.vision_feature_layers ) - image_outputs = self.vision_tower(pixel_values, output_hidden_states=True) + image_outputs = self.vision_tower( + pixel_values, + output_hidden_states=True, # Ignore arg on purpose + return_dict=True, + **kwargs, + ) # If multiple feature layers are provided (which is usually the case) # then the image features are concatenated after the CLS is removed. @@ -99,7 +108,9 @@ def get_image_features( image_features = [image_outputs.hidden_states[index][:, 1:] for index in vision_feature_layers] image_features = torch.cat(image_features, dim=-1) image_features = self.multi_modal_projector(image_features) - return image_features + image_outputs.pooler_output = image_features + + return image_outputs @auto_docstring def forward( @@ -140,8 +151,8 @@ def forward( if pixel_values is not None: image_features = self.get_image_features( - pixel_values=pixel_values, vision_feature_layers=vision_feature_layers - ) + pixel_values=pixel_values, vision_feature_layers=vision_feature_layers, return_dict=True + ).pooler_output image_features = image_features.to(inputs_embeds.device, inputs_embeds.dtype) special_image_mask = self.get_placeholder_mask( input_ids, inputs_embeds=inputs_embeds, image_features=image_features @@ -172,10 +183,23 @@ def forward( class VipLlavaForConditionalGeneration(LlavaForConditionalGeneration): + @auto_docstring def get_image_features( - self, pixel_values: torch.FloatTensor, vision_feature_layers: int | list[int] | None = None - ): - return self.model.get_image_features(pixel_values=pixel_values, vision_feature_layers=vision_feature_layers) + self, + pixel_values: torch.FloatTensor, + vision_feature_layers: int | list[int] | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + pixel_values (`torch.FloatTensor]` of shape `(batch_size, channels, height, width)`): + The tensors corresponding to the input images. + vision_feature_layers (`Union[int, list[int]]`, *optional*): + The vision feature layer, or the list of indexes of the layers to select + the vision feature. + """ + return self.model.get_image_features( + pixel_values=pixel_values, vision_feature_layers=vision_feature_layers, **kwargs + ) def forward( self, diff --git a/src/transformers/models/vision_text_dual_encoder/modeling_vision_text_dual_encoder.py b/src/transformers/models/vision_text_dual_encoder/modeling_vision_text_dual_encoder.py index ba2cc6e551c2..77aa6feb384a 100755 --- a/src/transformers/models/vision_text_dual_encoder/modeling_vision_text_dual_encoder.py +++ b/src/transformers/models/vision_text_dual_encoder/modeling_vision_text_dual_encoder.py @@ -18,7 +18,8 @@ from ...modeling_outputs import BaseModelOutputWithPooling from ...modeling_utils import PreTrainedModel -from ...utils import auto_docstring, filter_out_non_signature_kwargs, logging +from ...processing_utils import Unpack +from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, logging from ..auto.configuration_auto import AutoConfig from ..auto.modeling_auto import AutoModel from ..clip.modeling_clip import CLIPOutput, CLIPVisionConfig, CLIPVisionModel @@ -101,7 +102,7 @@ def __init__( self.post_init() - @filter_out_non_signature_kwargs() + @can_return_tuple @auto_docstring def get_text_features( self, @@ -109,12 +110,9 @@ def get_text_features( attention_mask: torch.Tensor | None = None, position_ids: torch.Tensor | None = None, token_type_ids: torch.Tensor | None = None, - ) -> torch.FloatTensor: + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: r""" - Returns: - text_features (`torch.FloatTensor` of shape `(batch_size, output_dim`): The text embeddings obtained by - applying the projection layer to the pooled output of [`CLIPTextModel`]. - Examples: ```python @@ -133,19 +131,20 @@ def get_text_features( attention_mask=attention_mask, position_ids=position_ids, token_type_ids=token_type_ids, + return_dict=True, + **kwargs, ) - text_features = self.text_projection(text_outputs.pooler_output) + pooled_output = text_outputs.pooler_output + text_outputs.pooler_output = self.text_projection(pooled_output) - return text_features + return text_outputs - @filter_out_non_signature_kwargs() + @can_return_tuple @auto_docstring - def get_image_features(self, pixel_values: torch.Tensor) -> torch.FloatTensor: + def get_image_features( + self, pixel_values: torch.Tensor, **kwargs: Unpack[TransformersKwargs] + ) -> tuple | BaseModelOutputWithPooling: r""" - Returns: - image_features (`torch.FloatTensor` of shape `(batch_size, output_dim`): The image embeddings obtained by - applying the projection layer to the pooled output of [`CLIPVisionModel`]. - Examples: ```python @@ -164,10 +163,10 @@ def get_image_features(self, pixel_values: torch.Tensor) -> torch.FloatTensor: >>> with torch.inference_mode(): ... image_features = model.get_image_features(**inputs) ```""" - vision_outputs = self.vision_model(pixel_values=pixel_values) - image_features = self.visual_projection(vision_outputs.pooler_output) + vision_outputs = self.vision_model(pixel_values=pixel_values, return_dict=True, **kwargs) + vision_outputs.pooler_output = self.visual_projection(vision_outputs.pooler_output) - return image_features + return vision_outputs @auto_docstring def forward( diff --git a/src/transformers/models/voxtral/modeling_voxtral.py b/src/transformers/models/voxtral/modeling_voxtral.py index 43b5d8bffaab..594ca09efe55 100644 --- a/src/transformers/models/voxtral/modeling_voxtral.py +++ b/src/transformers/models/voxtral/modeling_voxtral.py @@ -29,7 +29,7 @@ from ...cache_utils import Cache from ...generation import GenerationMixin from ...modeling_layers import GradientCheckpointingLayer -from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPast, CausalLMOutputWithPast +from ...modeling_outputs import BaseModelOutputWithPast, BaseModelOutputWithPooling, CausalLMOutputWithPast from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel from ...processing_utils import Unpack from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, logging @@ -296,7 +296,7 @@ def forward( input_features, attention_mask=None, **kwargs: Unpack[TransformersKwargs], - ): + ) -> tuple | BaseModelOutputWithPooling: r""" Args: input_features (`torch.LongTensor` of shape `(batch_size, feature_size, sequence_length)`): @@ -333,7 +333,7 @@ def forward( hidden_states = self.layer_norm(hidden_states) - return BaseModelOutput( + return BaseModelOutputWithPooling( last_hidden_state=hidden_states, ) @@ -397,26 +397,28 @@ def set_decoder(self, decoder): def get_decoder(self): return self.language_model.get_decoder() - def get_audio_features(self, input_features: torch.FloatTensor): - """ - This method is used to get the audio embeddings from input features (a log mel spectrogram), meaning inferring the audio encoder and the multi-modal projector. - Args: - input_features (`torch.FloatTensor`): - Float values of mel features extracted from the raw speech waveform. Raw speech waveform can be - obtained by loading a `.flac` or `.wav` audio file into an array of type `list[float]` or a - `numpy.ndarray`, *e.g.* via the soundfile library (`pip install soundfile`). To prepare the array into - `input_features`, the [`AutoFeatureExtractor`] should be used for extracting the mel features, padding - and conversion into a tensor of type `torch.FloatTensor`. See [`~WhisperFeatureExtractor.__call__`] - - Returns: - `torch.FloatTensor`: - The audio embeddings. + @can_return_tuple + @auto_docstring( + custom_intro="This method is used to get the audio embeddings from input features (a log mel spectrogram), meaning inferring the audio encoder and the multi-modal projector." + ) + def get_audio_features( + self, input_features: torch.FloatTensor, **kwargs: Unpack[TransformersKwargs] + ) -> tuple | BaseModelOutputWithPooling: + r""" + input_features (`torch.FloatTensor`): + Float values of mel features extracted from the raw speech waveform. Raw speech waveform can be + obtained by loading a `.flac` or `.wav` audio file into an array of type `list[float]` or a + `numpy.ndarray`, *e.g.* via the soundfile library (`pip install soundfile`). To prepare the array into + `input_features`, the [`AutoFeatureExtractor`] should be used for extracting the mel features, padding + and conversion into a tensor of type `torch.FloatTensor`. See [`~WhisperFeatureExtractor.__call__`] """ - audio_outputs = self.audio_tower(input_features) + audio_outputs = self.audio_tower(input_features, return_dict=True, **kwargs) audio_hidden_states = audio_outputs.last_hidden_state audio_hidden_states = audio_hidden_states.reshape(-1, self.config.audio_config.intermediate_size) audio_embeds = self.multi_modal_projector(audio_hidden_states) - return audio_embeds + audio_outputs.pooler_output = audio_embeds + + return audio_outputs @can_return_tuple @auto_docstring @@ -471,7 +473,7 @@ def forward( inputs_embeds = self.get_input_embeddings()(input_ids) if input_features is not None and input_ids is not None: - audio_embeds = self.get_audio_features(input_features) + audio_embeds = self.get_audio_features(input_features, return_dict=True).pooler_output # replace text-audio token placeholders with audio embeddings audio_token_mask = (input_ids == self.config.audio_token_id).unsqueeze(-1) diff --git a/src/transformers/models/voxtral/modular_voxtral.py b/src/transformers/models/voxtral/modular_voxtral.py index fab6bb44d65e..7fddbcd29648 100644 --- a/src/transformers/models/voxtral/modular_voxtral.py +++ b/src/transformers/models/voxtral/modular_voxtral.py @@ -19,7 +19,11 @@ from ...activations import ACT2FN from ...cache_utils import Cache from ...generation import GenerationMixin -from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPast, CausalLMOutputWithPast +from ...modeling_outputs import ( + BaseModelOutputWithPast, + BaseModelOutputWithPooling, + CausalLMOutputWithPast, +) from ...processing_utils import Unpack from ...utils import TransformersKwargs, auto_docstring, can_return_tuple from ...utils.generic import check_model_inputs @@ -67,7 +71,7 @@ def forward( input_features, attention_mask=None, **kwargs: Unpack[TransformersKwargs], - ): + ) -> tuple | BaseModelOutputWithPooling: r""" Args: input_features (`torch.LongTensor` of shape `(batch_size, feature_size, sequence_length)`): @@ -104,7 +108,7 @@ def forward( hidden_states = self.layer_norm(hidden_states) - return BaseModelOutput( + return BaseModelOutputWithPooling( last_hidden_state=hidden_states, ) @@ -159,26 +163,28 @@ def set_decoder(self, decoder): def get_decoder(self): return self.language_model.get_decoder() - def get_audio_features(self, input_features: torch.FloatTensor): - """ - This method is used to get the audio embeddings from input features (a log mel spectrogram), meaning inferring the audio encoder and the multi-modal projector. - Args: - input_features (`torch.FloatTensor`): - Float values of mel features extracted from the raw speech waveform. Raw speech waveform can be - obtained by loading a `.flac` or `.wav` audio file into an array of type `list[float]` or a - `numpy.ndarray`, *e.g.* via the soundfile library (`pip install soundfile`). To prepare the array into - `input_features`, the [`AutoFeatureExtractor`] should be used for extracting the mel features, padding - and conversion into a tensor of type `torch.FloatTensor`. See [`~WhisperFeatureExtractor.__call__`] - - Returns: - `torch.FloatTensor`: - The audio embeddings. + @can_return_tuple + @auto_docstring( + custom_intro="This method is used to get the audio embeddings from input features (a log mel spectrogram), meaning inferring the audio encoder and the multi-modal projector." + ) + def get_audio_features( + self, input_features: torch.FloatTensor, **kwargs: Unpack[TransformersKwargs] + ) -> tuple | BaseModelOutputWithPooling: + r""" + input_features (`torch.FloatTensor`): + Float values of mel features extracted from the raw speech waveform. Raw speech waveform can be + obtained by loading a `.flac` or `.wav` audio file into an array of type `list[float]` or a + `numpy.ndarray`, *e.g.* via the soundfile library (`pip install soundfile`). To prepare the array into + `input_features`, the [`AutoFeatureExtractor`] should be used for extracting the mel features, padding + and conversion into a tensor of type `torch.FloatTensor`. See [`~WhisperFeatureExtractor.__call__`] """ - audio_outputs = self.audio_tower(input_features) + audio_outputs = self.audio_tower(input_features, return_dict=True, **kwargs) audio_hidden_states = audio_outputs.last_hidden_state audio_hidden_states = audio_hidden_states.reshape(-1, self.config.audio_config.intermediate_size) audio_embeds = self.multi_modal_projector(audio_hidden_states) - return audio_embeds + audio_outputs.pooler_output = audio_embeds + + return audio_outputs @can_return_tuple @auto_docstring @@ -233,7 +239,7 @@ def forward( inputs_embeds = self.get_input_embeddings()(input_ids) if input_features is not None and input_ids is not None: - audio_embeds = self.get_audio_features(input_features) + audio_embeds = self.get_audio_features(input_features, return_dict=True).pooler_output # replace text-audio token placeholders with audio embeddings audio_token_mask = (input_ids == self.config.audio_token_id).unsqueeze(-1) diff --git a/src/transformers/models/x_clip/modeling_x_clip.py b/src/transformers/models/x_clip/modeling_x_clip.py index aa23e75a304d..f08c4a695e2b 100644 --- a/src/transformers/models/x_clip/modeling_x_clip.py +++ b/src/transformers/models/x_clip/modeling_x_clip.py @@ -27,7 +27,15 @@ from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel -from ...utils import ModelOutput, auto_docstring, can_return_tuple, filter_out_non_signature_kwargs, logging, torch_int +from ...processing_utils import Unpack +from ...utils import ( + ModelOutput, + TransformersKwargs, + auto_docstring, + can_return_tuple, + logging, + torch_int, +) from ...utils.generic import is_flash_attention_requested from .configuration_x_clip import XCLIPConfig, XCLIPTextConfig, XCLIPVisionConfig @@ -1195,19 +1203,16 @@ def __init__(self, config: XCLIPConfig): # Initialize weights and apply final processing self.post_init() - @filter_out_non_signature_kwargs() + @can_return_tuple @auto_docstring def get_text_features( self, input_ids: torch.Tensor | None = None, attention_mask: torch.Tensor | None = None, position_ids: torch.Tensor | None = None, - ) -> torch.FloatTensor: + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: r""" - Returns: - text_features (`torch.FloatTensor` of shape `(batch_size, output_dim`): The text embeddings obtained by - applying the projection layer to the pooled output of [`XCLIPTextModel`]. - Examples: ```python @@ -1225,22 +1230,22 @@ def get_text_features( input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids, + return_dict=True, + **kwargs, ) - text_features = self.text_projection(text_outputs.pooler_output) - return text_features + pooled_output = text_outputs.pooler_output + text_outputs.pooler_output = self.text_projection(pooled_output) - @filter_out_non_signature_kwargs() + return text_outputs + + @can_return_tuple @auto_docstring def get_video_features( self, pixel_values: torch.Tensor, - ) -> torch.FloatTensor: + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: r""" - Returns: - video_features (`torch.FloatTensor` of shape `(batch_size, output_dim`): The video embeddings obtained by - applying the projection layer to the pooled output of [`XCLIPVisionModel`] and - [`XCLIPMultiframeIntegrationTransformer`]. - Examples: ```python @@ -1313,17 +1318,17 @@ def get_video_features( batch_size, num_frames, num_channels, height, width = pixel_values.shape pixel_values = pixel_values.reshape(-1, num_channels, height, width) - vision_outputs: BaseModelOutputWithPooling = self.vision_model(pixel_values=pixel_values) - - video_embeds = vision_outputs.pooler_output + video_outputs: BaseModelOutputWithPooling = self.vision_model( + pixel_values=pixel_values, return_dict=True, **kwargs + ) + video_embeds = video_outputs.pooler_output video_embeds = self.visual_projection(video_embeds) cls_features = video_embeds.view(batch_size, num_frames, -1) + mit_outputs: BaseModelOutputWithPooling = self.mit(cls_features, return_dict=True, **kwargs) + video_outputs.pooler_output = mit_outputs.pooler_output - mit_outputs: BaseModelOutputWithPooling = self.mit(cls_features) - video_embeds = mit_outputs.pooler_output - - return video_embeds + return video_outputs @auto_docstring def forward( diff --git a/src/transformers/utils/generic.py b/src/transformers/utils/generic.py index 53aee7450972..e530b5f9acc8 100644 --- a/src/transformers/utils/generic.py +++ b/src/transformers/utils/generic.py @@ -875,34 +875,46 @@ def check_model_inputs(func=None, *, tie_last_hidden_states=True): def wrapped_fn(func): @wraps(func) def wrapper(self, *args, **kwargs): - use_cache_arg_index = None - if "use_cache" in func.__code__.co_varnames: - use_cache_arg_index = func.__code__.co_varnames.index("use_cache") - 1 # -1 for self - - if ( - use_cache_arg_index is not None - and len(args) > use_cache_arg_index - and args[use_cache_arg_index] is not None - ): - use_cache = args[use_cache_arg_index] - elif kwargs.get("use_cache") is not None: - use_cache = kwargs["use_cache"] - else: - use_cache = getattr(self.config, "use_cache", None) - - if use_cache is not None: - if getattr(self, "gradient_checkpointing", False) and self.training and use_cache: - logger.warning_once( - "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`." - ) - use_cache = False - - if use_cache_arg_index is not None and len(args) > use_cache_arg_index: - args = list(args) - args[use_cache_arg_index] = use_cache - args = tuple(args) + args_with_config_defaults = [ + "use_cache", + "vision_feature_layer", + "vision_feature_select_strategy", + "vision_aspect_ratio", + ] + for arg_name in args_with_config_defaults: + arg_index = None + if arg_name in func.__code__.co_varnames: + arg_index = func.__code__.co_varnames.index(arg_name) - 1 # -1 for self + + if arg_index is not None and len(args) > arg_index and args[arg_index] is not None: + arg_value = args[arg_index] + elif kwargs.get(arg_name) is not None: + arg_value = kwargs[arg_name] else: - kwargs["use_cache"] = use_cache + arg_value = getattr(self.config, arg_name, None) + + if arg_value is not None: + # Arg-specific handling + if arg_name == "use_cache": + if getattr(self, "gradient_checkpointing", False) and self.training and arg_value: + logger.warning_once( + "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`." + ) + arg_value = False + elif arg_name == "vision_feature_select_strategy": + valid_strategies = ["default", "full"] + if arg_value not in valid_strategies: + raise ValueError( + f"`Unexpected select feature strategy: {arg_value}. " + f"Please select from {valid_strategies}." + ) + + if arg_index is not None and len(args) > arg_index: + args = list(args) + args[arg_index] = arg_value + args = tuple(args) + else: + kwargs[arg_name] = arg_value return_dict = kwargs.pop("return_dict", None) if return_dict is None: diff --git a/tests/models/align/test_modeling_align.py b/tests/models/align/test_modeling_align.py index 863623d2f3e0..16429a31b029 100644 --- a/tests/models/align/test_modeling_align.py +++ b/tests/models/align/test_modeling_align.py @@ -14,6 +14,7 @@ """Testing suite for the PyTorch ALIGN model.""" import inspect +import math import tempfile import unittest @@ -59,6 +60,7 @@ def __init__( batch_size=12, image_size=32, num_channels=3, + depth_coefficient=3.1, kernel_sizes=[3, 3, 5], in_channels=[32, 16, 24], out_channels=[16, 24, 30], @@ -73,6 +75,7 @@ def __init__( self.batch_size = batch_size self.image_size = image_size self.num_channels = num_channels + self.depth_coefficient = depth_coefficient self.kernel_sizes = kernel_sizes self.in_channels = in_channels self.out_channels = out_channels @@ -92,6 +95,7 @@ def prepare_config_and_inputs(self): def get_config(self): return AlignVisionConfig( num_channels=self.num_channels, + depth_coefficient=self.depth_coefficient, kernel_sizes=self.kernel_sizes, in_channels=self.in_channels, out_channels=self.out_channels, @@ -434,6 +438,8 @@ class AlignModelTest(ModelTesterMixin, PipelineTesterMixin, unittest.TestCase): test_resize_embeddings = False test_attention_outputs = False + has_attentions = False + skip_test_image_features_output_shape = True # Align uses index -3 for hidden_size instead of -1 def setUp(self): self.model_tester = AlignModelTester(self) @@ -493,6 +499,12 @@ def test_load_vision_text_config(self): text_config = AlignTextConfig.from_pretrained(tmp_dir_name) self.assertDictEqual(config.text_config.to_dict(), text_config.to_dict()) + def _image_features_get_expected_num_attentions(self, model_tester=None): + return sum( + math.ceil(self.model_tester.vision_model_tester.depth_coefficient * repeat) + for repeat in self.model_tester.vision_model_tester.num_block_repeats + ) + @slow def test_model_from_pretrained(self): model_name = "kakaobrain/align-base" diff --git a/tests/models/blip/test_modeling_blip.py b/tests/models/blip/test_modeling_blip.py index acfe8f405a53..20ac995ca703 100644 --- a/tests/models/blip/test_modeling_blip.py +++ b/tests/models/blip/test_modeling_blip.py @@ -463,7 +463,7 @@ def test_get_image_features(self): model.eval() image_features = model.get_image_features(**inputs_dict) self.assertEqual( - image_features.shape, + image_features.pooler_output.shape, ( self.model_tester.batch_size, model.projection_dim, @@ -482,7 +482,7 @@ def test_get_text_features(self): model.eval() text_features = model.get_text_features(**inputs_dict) self.assertEqual( - text_features.shape, + text_features.pooler_output.shape, ( self.model_tester.batch_size, model.projection_dim, diff --git a/tests/models/blip_2/test_modeling_blip_2.py b/tests/models/blip_2/test_modeling_blip_2.py index aa1878e473c9..999eafbf94d1 100644 --- a/tests/models/blip_2/test_modeling_blip_2.py +++ b/tests/models/blip_2/test_modeling_blip_2.py @@ -688,6 +688,25 @@ def get_config(self): is_encoder_decoder=True, ) + def prepare_config_and_inputs_for_common(self): + config_and_inputs = self.prepare_config_and_inputs() + ( + config, + input_ids, + decoder_input_ids, + attention_mask, + decoder_attention_mask, + labels, + ) = config_and_inputs + inputs_dict = { + "input_ids": input_ids, + "attention_mask": attention_mask, + "decoder_input_ids": decoder_input_ids, + "decoder_attention_mask": decoder_attention_mask, + "labels": labels, + } + return config, inputs_dict + # this model tester uses an encoder-decoder language model (T5) class Blip2ModelTester: @@ -952,13 +971,12 @@ def test_get_text_features(self): inputs_dict = { "input_ids": torch.LongTensor([[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]]).to(torch_device), "attention_mask": torch.LongTensor([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]).to(torch_device), - "decoder_input_ids": torch.LongTensor([[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]]).to(torch_device), } model = Blip2Model(config).to(torch_device) model.eval() text_features = model.get_text_features(**inputs_dict) - self.assertEqual(text_features[0].shape, (10, config.text_config.vocab_size)) + self.assertEqual(text_features[0].shape, (1, 10, config.text_config.hidden_size)) def test_get_image_features(self): config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common() @@ -972,7 +990,7 @@ def test_get_image_features(self): model.eval() image_features = model.get_image_features(**inputs_dict) self.assertEqual( - image_features[0].shape, + image_features.pooler_output[0].shape, (config.vision_config.hidden_size,), ) diff --git a/tests/models/chameleon/test_modeling_chameleon.py b/tests/models/chameleon/test_modeling_chameleon.py index b2521fba555b..78dcbafef6ca 100644 --- a/tests/models/chameleon/test_modeling_chameleon.py +++ b/tests/models/chameleon/test_modeling_chameleon.py @@ -17,6 +17,7 @@ import unittest import requests +from parameterized import parameterized from transformers import BitsAndBytesConfig, ChameleonConfig, is_torch_available, is_vision_available from transformers.testing_utils import ( @@ -74,7 +75,11 @@ def __init__( pad_token_id=0, vq_num_embeds=5, vq_embed_dim=5, + vq_resolution=512, vq_channel_multiplier=[1, 2], + vq_num_res_blocks=2, + vq_attn_resolutions=None, + vq_attn_type="vanilla", vq_img_token_start_id=10, # has to be less than vocab size when added with vq_num_embeds scope=None, ): @@ -104,7 +109,11 @@ def __init__( self.scope = scope self.vq_num_embeds = vq_num_embeds self.vq_embed_dim = vq_embed_dim + self.vq_resolution = vq_resolution self.vq_channel_multiplier = vq_channel_multiplier + self.vq_num_res_blocks = vq_num_res_blocks + self.vq_attn_resolutions = vq_attn_resolutions + self.vq_attn_type = vq_attn_type self.vq_img_token_start_id = vq_img_token_start_id def prepare_config_and_inputs(self): @@ -165,9 +174,13 @@ def get_vq_config(self): "embed_dim": self.vq_embed_dim, "num_embeddings": self.vq_num_embeds, "latent_channels": self.vq_embed_dim, + "resolution": self.vq_resolution, "in_channels": 3, "base_channels": 32, # we have a GroupNorm of 32 groups, so can't do less "channel_multiplier": self.vq_channel_multiplier, + "num_res_blocks": self.vq_num_res_blocks, + "attn_resolutions": self.vq_attn_resolutions, + "attn_type": self.vq_attn_type, } def create_and_check_model(self, config, input_ids, input_mask, sequence_labels, token_labels, choice_labels): @@ -219,6 +232,19 @@ def test_model(self): def test_batching_equivalence(self): pass + @unittest.skip("Skip get_image_features tests as those are tested via ChameleonVision2SeqModelTest instead") + @parameterized.expand([True, False, None]) + def test_get_image_features_output(self, return_dict: bool | None): + pass + + @unittest.skip("Skip get_image_features tests as those are tested via ChameleonVision2SeqModelTest instead") + def test_get_image_features_hidden_states(self): + pass + + @unittest.skip("Skip get_image_features tests as those are tested via ChameleonVision2SeqModelTest instead") + def test_get_image_features_attentions(self): + pass + class ChameleonVision2SeqModelTester(ChameleonModelTester): def __init__(self, parent, image_size=10, **kwargs): @@ -322,6 +348,24 @@ def test_mismatching_num_image_tokens(self): pixel_values = torch.cat([pixel_values, pixel_values], dim=0) _ = model(input_ids=input_ids, pixel_values=pixel_values) + def _image_features_get_expected_num_hidden_states(self, model_tester=None): + if model_tester is None: + model_tester = self.model_tester + # The number of ChameleonVQVAEEncoderResnetBlock instances, plus 1 for before the block + return len(model_tester.vq_channel_multiplier) * model_tester.vq_num_res_blocks + 3 + + def _image_features_get_expected_num_attentions(self, model_tester=None): + if model_tester is None: + model_tester = self.model_tester + # The number of ChameleonVQVAEEncoderAttnBlock instances + if ( + model_tester.vq_attn_resolutions + and model_tester.vq_resolution in model_tester.vq_attn_resolutions + and model_tester.vq_attn_type == "vanilla" + ): + return len(model_tester.vq_channel_multiplier) * model_tester.vq_num_res_blocks + 1 + return 1 + @require_torch class ChameleonIntegrationTest(unittest.TestCase): diff --git a/tests/models/clip/test_modeling_clip.py b/tests/models/clip/test_modeling_clip.py index ceda8e9cf3fa..7c6d4bab40f5 100644 --- a/tests/models/clip/test_modeling_clip.py +++ b/tests/models/clip/test_modeling_clip.py @@ -587,6 +587,26 @@ def test_sdpa_can_dispatch_on_flash(self): def test_sdpa_can_compile_dynamic(self): self.skipTest(reason="CLIP model can't be compiled dynamic, error in clip_loss`") + @unittest.skip(reason="The CLIP family currently does not work with output_attentions.") + def test_get_text_features_attentions(self): + # This test should no longer be skipped once this architecture is refactored to work with output_attentions. + pass + + @unittest.skip(reason="The CLIP family currently does not work with output_hidden_states.") + def test_get_text_features_hidden_states(self): + # This test should no longer be skipped once this architecture is refactored to work with output_hidden_states. + pass + + @unittest.skip(reason="The CLIP family currently does not work with output_attentions.") + def test_get_image_features_attentions(self): + # This test should no longer be skipped once this architecture is refactored to work with output_attentions. + pass + + @unittest.skip(reason="The CLIP family currently does not work with output_hidden_states.") + def test_get_image_features_hidden_states(self): + # This test should no longer be skipped once this architecture is refactored to work with output_hidden_states. + pass + class CLIPForImageClassificationModelTester(CLIPModelTester): def __init__(self, parent): diff --git a/tests/models/edgetam/test_modeling_edgetam.py b/tests/models/edgetam/test_modeling_edgetam.py index 508ab62d416a..36d0f3ac21fd 100644 --- a/tests/models/edgetam/test_modeling_edgetam.py +++ b/tests/models/edgetam/test_modeling_edgetam.py @@ -13,10 +13,12 @@ # limitations under the License. """Testing suite for the PyTorch EDGETAM model.""" +import copy import gc import unittest import requests +from parameterized import parameterized from transformers import ( EdgeTamConfig, @@ -36,7 +38,7 @@ from transformers.video_utils import load_video from ...test_configuration_common import ConfigTester -from ...test_modeling_common import ModelTesterMixin, floats_tensor +from ...test_modeling_common import TEST_EAGER_MATCHES_SDPA_INFERENCE_PARAMETERIZATION, ModelTesterMixin, floats_tensor from ...test_pipeline_mixin import PipelineTesterMixin @@ -327,6 +329,13 @@ def test_flash_attn_kernels_mps_inference_equivalence(self): def test_eager_matches_sdpa_generate(self): pass + @parameterized.expand(TEST_EAGER_MATCHES_SDPA_INFERENCE_PARAMETERIZATION) + @unittest.skip("Test requires hidden_states in outputs, which is not available in EdgeTamModel.forward's output") + def test_eager_matches_sdpa_inference( + self, name, dtype, padding_side, use_attention_mask, output_attentions, enable_kernels + ): + pass + @unittest.skip("Flash attn test is not configured correctly as we need to configure vision/timm model to 'eager'.") def test_flash_attn_2_inference_equivalence(self): pass @@ -356,6 +365,67 @@ def test_model_from_pretrained(self): def test_sdpa_can_compile_dynamic(self): self.skipTest(reason="EDGETAM model can't be compiled dynamic yet") + def test_model_outputs_equivalence(self): + # Modified from upstream to remove output_hidden_states as the timm model doesn't support it + config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common() + + def set_nan_tensor_to_zero(t): + t[t != t] = 0 + return t + + def check_equivalence(model, tuple_inputs, dict_inputs, additional_kwargs={}): + with torch.no_grad(): + tuple_output = model(**tuple_inputs, return_dict=False, **additional_kwargs) + dict_output = model(**dict_inputs, return_dict=True, **additional_kwargs).to_tuple() + + def recursive_check(tuple_object, dict_object): + if isinstance(tuple_object, (list, tuple)): + for tuple_iterable_value, dict_iterable_value in zip(tuple_object, dict_object): + recursive_check(tuple_iterable_value, dict_iterable_value) + elif isinstance(tuple_object, dict): + for tuple_iterable_value, dict_iterable_value in zip( + tuple_object.values(), dict_object.values() + ): + recursive_check(tuple_iterable_value, dict_iterable_value) + elif tuple_object is None: + return + # model might return non-tensors objects (e.g. Cache class) + elif isinstance(tuple_object, torch.Tensor): + self.assertTrue( + torch.allclose( + set_nan_tensor_to_zero(tuple_object), set_nan_tensor_to_zero(dict_object), atol=1e-5 + ), + msg=( + "Tuple and dict output are not equal. Difference:" + f" {torch.max(torch.abs(tuple_object - dict_object))}. Tuple has `nan`:" + f" {torch.isnan(tuple_object).any()} and `inf`: {torch.isinf(tuple_object)}. Dict has" + f" `nan`: {torch.isnan(dict_object).any()} and `inf`: {torch.isinf(dict_object)}." + ), + ) + + recursive_check(tuple_output, dict_output) + + for model_class in self.all_model_classes: + model = model_class(copy.deepcopy(config)) + model.to(torch_device) + model.eval() + + tuple_inputs = self._prepare_for_class(inputs_dict, model_class) + dict_inputs = self._prepare_for_class(inputs_dict, model_class) + check_equivalence(model, tuple_inputs, dict_inputs) + + tuple_inputs = self._prepare_for_class(inputs_dict, model_class, return_labels=True) + dict_inputs = self._prepare_for_class(inputs_dict, model_class, return_labels=True) + check_equivalence(model, tuple_inputs, dict_inputs) + + @unittest.skip("Cannot set `output_attentions` for timm models.") + def test_get_image_features_attentions(self): + pass + + @unittest.skip("Cannot set `output_hidden_states` for this timm model.") + def test_get_image_features_hidden_states(self): + pass + def prepare_image(): img_url = "https://huggingface.co/datasets/hf-internal-testing/sam2-fixtures/resolve/main/truck.jpg" diff --git a/tests/models/emu3/test_modeling_emu3.py b/tests/models/emu3/test_modeling_emu3.py index 5803b0be37b6..91edf5ecf0cf 100644 --- a/tests/models/emu3/test_modeling_emu3.py +++ b/tests/models/emu3/test_modeling_emu3.py @@ -167,6 +167,7 @@ def __init__( temporal_downsample_factor=1, base_channels=32, vq_channel_multiplier=[1, 2, 1], + vq_num_res_blocks=2, image_seq_length=12, vq_img_token_start_id=3, ): @@ -189,6 +190,7 @@ def __init__( self.codebook_size = codebook_size self.temporal_downsample_factor = temporal_downsample_factor self.vq_channel_multiplier = vq_channel_multiplier + self.vq_num_res_blocks = vq_num_res_blocks self.vq_img_token_start_id = vq_img_token_start_id self.base_channels = base_channels self.seq_length = seq_length + image_seq_length @@ -289,6 +291,7 @@ class Emu3Vision2TextModelTest(ModelTesterMixin, GenerationTesterMixin, Pipeline if is_torch_available() else {} ) + skip_test_image_features_output_shape = True # Emu3 uses index -3 for hidden_size instead of -1 def setUp(self): self.model_tester = Emu3Vision2TextModelTester(self) @@ -320,6 +323,21 @@ def test_cpu_offload(self): def test_generate_with_static_cache(self): pass + def _image_features_get_expected_num_attentions(self, model_tester=None): + if model_tester is None: + model_tester = self.model_tester + # The number of Emu3VQVAEAttentionBlock instances in the encoder, assumes that attn_resolutions is empty (default) + # 0 via down due to attn_resolutions being empty, 1 via middle block, 0 via up due to attn_resolutions being empty + return 1 + + def _image_features_get_expected_num_hidden_states(self, model_tester=None): + if model_tester is None: + model_tester = self.model_tester + # The number of Emu3VQVAEResnetBlock and Emu3VQVAETemporalResnetBlock instances in the encoder, plus 1 for before the block + # up_down_blocks for down, 2 for middle, vq_num_res_blocks for Emu3VQVAETemporalResnetBlock + up_down_blocks = len(model_tester.vq_channel_multiplier) * model_tester.vq_num_res_blocks + return up_down_blocks + 2 + model_tester.vq_num_res_blocks + 1 + @require_torch class Emu3IntegrationTest(unittest.TestCase): diff --git a/tests/models/ernie4_5_vl_moe/test_modeling_ernie4_5_vl_moe.py b/tests/models/ernie4_5_vl_moe/test_modeling_ernie4_5_vl_moe.py index dbc41f22064b..d734e73be849 100644 --- a/tests/models/ernie4_5_vl_moe/test_modeling_ernie4_5_vl_moe.py +++ b/tests/models/ernie4_5_vl_moe/test_modeling_ernie4_5_vl_moe.py @@ -266,6 +266,34 @@ def test_inputs_embeds_matches_input_ids(self): def test_multi_gpu_data_parallel_forward(self): pass + def _video_features_prepare_config_and_inputs(self): + """ + Helper method to extract only video-related inputs from the full set of inputs, for testing `get_video_features`. + + The superclass method simply calls the model_tester.prepare_config_and_inputs_for_common(), + but that method only prepared image inputs, i.e. where the temporal dimension in grid_thw is 1. + This override prepares proper video inputs with 12 frames. + """ + config = self.model_tester.get_config() + patch_size = config.vision_config.patch_size + batch_size = self.model_tester.batch_size + image_size = self.model_tester.image_size + num_channels = self.model_tester.num_channels + num_frames = 12 + pixel_values_videos = floats_tensor( + [num_frames * batch_size * (image_size**2) // (patch_size**2), num_channels * (patch_size**2)] + ) + + patches_per_side = image_size // patch_size + video_grid_thw = torch.tensor( + [[num_frames, patches_per_side, patches_per_side]] * batch_size, device=torch_device + ) + inputs_dict = { + "pixel_values_videos": pixel_values_videos, + "video_grid_thw": video_grid_thw, + } + return config, inputs_dict + @slow @require_torch_large_accelerator(memory=64) # Tested on A100 diff --git a/tests/models/fast_vlm/test_modeling_fast_vlm.py b/tests/models/fast_vlm/test_modeling_fast_vlm.py index eb2a5f5ce425..73645dbfeab0 100644 --- a/tests/models/fast_vlm/test_modeling_fast_vlm.py +++ b/tests/models/fast_vlm/test_modeling_fast_vlm.py @@ -176,6 +176,7 @@ class FastVlmForConditionalGenerationModelTest(ModelTesterMixin, GenerationTeste if is_torch_available() else {} ) + skip_test_image_features_output_shape = True # FastVLM uses index -3 for hidden_size instead of -1 _is_composite = True @@ -227,6 +228,15 @@ def test_mismatching_num_image_tokens(self): def test_can_be_initialized_on_meta(self): pass + @unittest.skip("Cannot set output_attentions on timm models.") + def test_get_image_features_attentions(self): + pass + + def _image_features_get_expected_num_hidden_states(self, model_tester=None): + # For models that rely on timm for their vision backend, it's hard to infer how many layers the model has + # from the timm config alone. So, we're just hardcoding the expected number of hidden states here. + return 2 + @require_torch @slow diff --git a/tests/models/flava/test_modeling_flava.py b/tests/models/flava/test_modeling_flava.py index 58e2785938c6..19a9771edc67 100644 --- a/tests/models/flava/test_modeling_flava.py +++ b/tests/models/flava/test_modeling_flava.py @@ -762,7 +762,7 @@ def __init__( image_codebook_kwargs = {} self.parent = parent - self.image_model_tester = FlavaImageModelTester(parent, **image_kwargs) + self.vision_model_tester = FlavaImageModelTester(parent, **image_kwargs) self.text_model_tester = FlavaTextModelTester(parent, **text_kwargs) self.multimodal_model_tester = FlavaMultimodalModelTester(parent, **multimodal_kwargs) self.image_codebook_tester = FlavaImageCodebookTester(parent, **image_codebook_kwargs) @@ -778,7 +778,7 @@ def test_config(self): self.config_tester.run_common_tests() def prepare_config_and_inputs_for_common(self): - _, pixel_values, bool_masked_pos = self.image_model_tester.prepare_config_and_inputs() + _, pixel_values, bool_masked_pos = self.vision_model_tester.prepare_config_and_inputs() _, input_ids, token_type_ids, attention_mask = self.text_model_tester.prepare_config_and_inputs() config = self.get_config() @@ -793,7 +793,7 @@ def prepare_config_and_inputs_for_common(self): def get_config(self): return FlavaConfig( - image_config=self.image_model_tester.get_config(), + image_config=self.vision_model_tester.get_config(), text_config=self.text_model_tester.get_config(), multimodal_config=self.multimodal_model_tester.get_config(), image_codebook_config=self.image_codebook_tester.get_config(), @@ -818,14 +818,14 @@ def _test_model(self, config, inputs, test_image=False, test_text=False): pixel_values=inputs["pixel_values"] if test_image else None, bool_masked_pos=inputs["bool_masked_pos"] if test_image else None, ) - image_size = (self.image_model_tester.image_size, self.image_model_tester.image_size) - patch_size = (self.image_model_tester.patch_size, self.image_model_tester.patch_size) + image_size = (self.vision_model_tester.image_size, self.vision_model_tester.image_size) + patch_size = (self.vision_model_tester.patch_size, self.vision_model_tester.patch_size) num_patches = (image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0]) if test_image: self.parent.assertEqual( result.image_embeddings.shape, - (self.image_model_tester.batch_size, num_patches + 1, self.image_model_tester.hidden_size), + (self.vision_model_tester.batch_size, num_patches + 1, self.vision_model_tester.hidden_size), ) else: self.parent.assertIsNone(result.image_embeddings) @@ -927,7 +927,7 @@ class FlavaForPreTrainingTester(FlavaModelTester): model_class = FlavaForPreTraining def prepare_config_and_inputs_for_common(self): - _, pixel_values, bool_masked_pos = self.image_model_tester.prepare_config_and_inputs() + _, pixel_values, bool_masked_pos = self.vision_model_tester.prepare_config_and_inputs() _, input_ids, token_type_ids, attention_mask = self.text_model_tester.prepare_config_and_inputs() config = self.get_config() @@ -937,7 +937,7 @@ def prepare_config_and_inputs_for_common(self): mlm_labels[:, :] = config.ce_ignore_index mlm_labels[:, 1:3] = input_ids[:, 1:3] mim_labels = torch.randint( - 0, self.image_model_tester.vocab_size, bool_masked_pos.size(), device=bool_masked_pos.device + 0, self.vision_model_tester.vocab_size, bool_masked_pos.size(), device=bool_masked_pos.device ).long() mim_labels[bool_masked_pos.ne(True)] = config.ce_ignore_index itm_labels = torch.ones(mlm_labels.size(0), device=bool_masked_pos.device).long() @@ -970,14 +970,14 @@ def _test_model(self, config, inputs, test_image=False, test_text=False): itm_labels=inputs["itm_labels"], return_loss=inputs["return_loss"], ) - image_size = (self.image_model_tester.image_size, self.image_model_tester.image_size) - patch_size = (self.image_model_tester.patch_size, self.image_model_tester.patch_size) + image_size = (self.vision_model_tester.image_size, self.vision_model_tester.image_size) + patch_size = (self.vision_model_tester.patch_size, self.vision_model_tester.patch_size) num_patches = (image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0]) if test_image: self.parent.assertEqual( result.image_embeddings.shape, - (self.image_model_tester.batch_size, num_patches + 1, self.image_model_tester.hidden_size), + (self.vision_model_tester.batch_size, num_patches + 1, self.vision_model_tester.hidden_size), ) if not test_text: self.parent.assertEqual( @@ -986,7 +986,7 @@ def _test_model(self, config, inputs, test_image=False, test_text=False): ) self.parent.assertEqual( result.mim_logits.shape, - (inputs["bool_masked_pos"].sum().item(), self.image_model_tester.vocab_size), + (inputs["bool_masked_pos"].sum().item(), self.vision_model_tester.vocab_size), ) else: @@ -1035,15 +1035,15 @@ def _test_model(self, config, inputs, test_image=False, test_text=False): ) self.parent.assertEqual( result.mmm_image_logits.shape, - (inputs["bool_masked_pos"].sum().item(), self.image_model_tester.vocab_size), + (inputs["bool_masked_pos"].sum().item(), self.vision_model_tester.vocab_size), ) self.parent.assertEqual( result.contrastive_logits_per_image.shape, - (self.image_model_tester.batch_size, self.text_model_tester.batch_size), + (self.vision_model_tester.batch_size, self.text_model_tester.batch_size), ) self.parent.assertEqual( result.contrastive_logits_per_text.shape, - (self.text_model_tester.batch_size, self.image_model_tester.batch_size), + (self.text_model_tester.batch_size, self.vision_model_tester.batch_size), ) for item in [ diff --git a/tests/models/florence2/test_modeling_florence2.py b/tests/models/florence2/test_modeling_florence2.py index c2dc6fc436a7..fcdce182fda2 100644 --- a/tests/models/florence2/test_modeling_florence2.py +++ b/tests/models/florence2/test_modeling_florence2.py @@ -238,8 +238,9 @@ class Florence2ForConditionalGenerationModelTest( if is_torch_available() else {} ) + skip_test_image_features_output_shape = True # Florence2 uses index -3 for hidden_size instead of -1 - test_attention_outputs = False + has_attentions = False _is_composite = True def setUp(self): diff --git a/tests/models/fuyu/test_modeling_fuyu.py b/tests/models/fuyu/test_modeling_fuyu.py index 3ffd29deaa56..6f06c8d5f68d 100644 --- a/tests/models/fuyu/test_modeling_fuyu.py +++ b/tests/models/fuyu/test_modeling_fuyu.py @@ -246,6 +246,24 @@ def test_sdpa_padding_matches_padding_free_with_position_ids(self): def test_model_base_model_prefix(self): pass + def _image_features_prepare_config_and_inputs(self): + """ + Helper method to extract only image-related inputs from the full set of inputs, for testing `get_image_features`. + + The Fuyu model uses image_patches, except for get_image_features, where they're called pixel_values. + """ + config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common() + inputs_dict = {"pixel_values": inputs_dict["image_patches"]} + return config, inputs_dict + + @unittest.skip("Skip get_image_features tests as Fuyu's image features originate from a simple Linear") + def test_get_image_features_hidden_states(self): + pass + + @unittest.skip("Skip get_image_features tests as Fuyu's image features originate from a simple Linear") + def test_get_image_features_attentions(self): + pass + @slow @require_torch_accelerator diff --git a/tests/models/glm_image/test_modeling_glm_image.py b/tests/models/glm_image/test_modeling_glm_image.py index 85e72a12efc7..6f06d1417734 100644 --- a/tests/models/glm_image/test_modeling_glm_image.py +++ b/tests/models/glm_image/test_modeling_glm_image.py @@ -91,7 +91,6 @@ def __init__( "depth": 2, "hidden_act": "gelu", "hidden_size": 32, - "out_hidden_size": 16, "intermediate_size": 22, "patch_size": 16, "spatial_merge_size": 1, @@ -397,6 +396,19 @@ def test_flash_attn_2_fp32_ln(self): def test_flash_attn_2_from_config(self): pass + def _image_features_prepare_config_and_inputs(self): + """ + Helper method to extract only image-related inputs from the full set of inputs, for testing `get_image_features`. + + GlmImage internally preprocesses the image_grid_thw input by slicing off the last entry, + so we need to prepare inputs accordingly for testing get_image_features. We also discard text-related inputs. + """ + config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common() + inputs_dict["image_grid_thw"] = inputs_dict["image_grid_thw"][:-1] + del inputs_dict["input_ids"] + del inputs_dict["attention_mask"] + return config, inputs_dict + @require_torch @slow diff --git a/tests/models/groupvit/test_modeling_groupvit.py b/tests/models/groupvit/test_modeling_groupvit.py index 6565630e0fbf..89d1befd0c3d 100644 --- a/tests/models/groupvit/test_modeling_groupvit.py +++ b/tests/models/groupvit/test_modeling_groupvit.py @@ -574,6 +574,17 @@ def test_model_from_pretrained(self): model = GroupViTModel.from_pretrained(model_name) self.assertIsNotNone(model) + def _image_features_get_expected_num_attentions(self, model_tester=None): + if model_tester is None: + model_tester = self.model_tester.vision_model_tester + # GroupViT returns attention grouping of each stage + return sum(g > 0 for g in self.model_tester.vision_model_tester.num_group_tokens) + + def _image_features_get_expected_num_hidden_states(self, model_tester=None): + if model_tester is None: + model_tester = self.model_tester.vision_model_tester + return model_tester.expected_num_hidden_layers + # We will verify our results on an image of cute cats def prepare_img(): diff --git a/tests/models/idefics2/test_modeling_idefics2.py b/tests/models/idefics2/test_modeling_idefics2.py index 8e681b91718e..18db4501c185 100644 --- a/tests/models/idefics2/test_modeling_idefics2.py +++ b/tests/models/idefics2/test_modeling_idefics2.py @@ -177,6 +177,8 @@ class Idefics2ModelTest(ModelTesterMixin, unittest.TestCase): """ all_model_classes = (Idefics2Model,) if is_torch_available() else () + # Idefics2 merges batch_size and num_frames in the first output dimension + skip_test_image_features_output_shape = True test_resize_embeddings = True _is_composite = True @@ -367,6 +369,9 @@ class Idefics2ForConditionalGenerationModelTest(GenerationTesterMixin, ModelTest all_model_classes = (Idefics2ForConditionalGeneration,) if is_torch_available() else () pipeline_model_mapping = {"image-text-to-text": Idefics2ForConditionalGeneration} if is_torch_available() else () + skip_test_image_features_output_shape = ( + True # Idefics2 merges batch_size and num_frames in the first output dimension + ) test_resize_embeddings = True diff --git a/tests/models/idefics3/test_modeling_idefics3.py b/tests/models/idefics3/test_modeling_idefics3.py index 4642444be237..3a9b402e96c6 100644 --- a/tests/models/idefics3/test_modeling_idefics3.py +++ b/tests/models/idefics3/test_modeling_idefics3.py @@ -167,6 +167,8 @@ class Idefics3ModelTest(ModelTesterMixin, unittest.TestCase): """ all_model_classes = (Idefics3Model,) if is_torch_available() else () + # Idefics3 merges batch_size and num_frames in the first output dimension + skip_test_image_features_output_shape = True test_resize_embeddings = True @@ -332,6 +334,9 @@ class Idefics3ForConditionalGenerationModelTest(GenerationTesterMixin, ModelTest all_model_classes = (Idefics3ForConditionalGeneration,) if is_torch_available() else () pipeline_model_mapping = {"image-text-to-text": Idefics3ForConditionalGeneration} if is_torch_available() else () + skip_test_image_features_output_shape = ( + True # Idefics3 merges batch_size and num_frames in the first output dimension + ) test_resize_embeddings = True diff --git a/tests/models/instructblip/test_modeling_instructblip.py b/tests/models/instructblip/test_modeling_instructblip.py index 831b2bb5f800..d2ecfe3eb93f 100644 --- a/tests/models/instructblip/test_modeling_instructblip.py +++ b/tests/models/instructblip/test_modeling_instructblip.py @@ -615,6 +615,18 @@ def test_sdpa_can_dispatch_composite_models(self): ): raise ValueError("The eager model should not have SDPA attention layers") + def _image_features_prepare_config_and_inputs(self): + """ + Helper method to extract only image-related inputs from the full set of inputs, for testing `get_image_features`. + + InstructBlip's `get_image_features` uses `qformer_input_ids` and `qformer_attention_mask` along with `pixel_values`, + so we override this method to keep those, and only discard `input_ids` and `attention_mask`. + """ + config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common() + del inputs_dict["input_ids"] + del inputs_dict["attention_mask"] + return config, inputs_dict + # We will verify our results on an image of cute cats def prepare_img(): diff --git a/tests/models/instructblipvideo/test_modeling_instructblipvideo.py b/tests/models/instructblipvideo/test_modeling_instructblipvideo.py index 7426204bb01b..3abf88e79239 100644 --- a/tests/models/instructblipvideo/test_modeling_instructblipvideo.py +++ b/tests/models/instructblipvideo/test_modeling_instructblipvideo.py @@ -482,6 +482,8 @@ class InstructBlipVideoForConditionalGenerationDecoderOnlyTest( (InstructBlipVideoForConditionalGeneration, InstructBlipVideoModel) if is_torch_available() else () ) additional_model_inputs = ["qformer_input_ids", "input_ids"] + # InstructBlipVideo merges batch_size and num_frames in the first output dimension + skip_test_video_features_output_shape = True test_resize_embeddings = True test_attention_outputs = False @@ -623,6 +625,18 @@ def test_sdpa_can_dispatch_composite_models(self): ): raise ValueError("The eager model should not have SDPA attention layers") + def _video_features_prepare_config_and_inputs(self): + """ + Helper method to extract only video-related inputs from the full set of inputs, for testing `get_video_features`. + + InstructBlip's `get_video_features` uses `qformer_input_ids` and `qformer_attention_mask` along with `pixel_values`, + so we override this method to keep those, and only discard `input_ids` and `attention_mask`. + """ + config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common() + del inputs_dict["input_ids"] + del inputs_dict["attention_mask"] + return config, inputs_dict + # We will verify our results on an image of cute cats def prepare_video(): diff --git a/tests/models/lighton_ocr/test_modeling_lighton_ocr.py b/tests/models/lighton_ocr/test_modeling_lighton_ocr.py index 6dd0fc8be378..528e9fdd9530 100644 --- a/tests/models/lighton_ocr/test_modeling_lighton_ocr.py +++ b/tests/models/lighton_ocr/test_modeling_lighton_ocr.py @@ -225,6 +225,8 @@ class LightOnOcrForConditionalGenerationModelTest(ModelTesterMixin, GenerationTe else () ) pipeline_model_mapping = {"image-text-to-text": LightOnOcrForConditionalGeneration} if is_torch_available() else {} + # LightOnOcr uses a PixtralVisionModel, which merges batch_size and num_patches in index 1, with index 0 hardcoded to 1 + skip_test_image_features_output_shape = True _is_composite = True @@ -421,7 +423,7 @@ def test_get_image_features(self): image_features_list = model.get_image_features( pixel_values=input_dict["pixel_values"].to(torch_device), image_sizes=input_dict["image_sizes"], - ) + ).pooler_output # Check that features are returned as a list self.assertIsNotNone(image_features_list) diff --git a/tests/models/llava_next/test_modeling_llava_next.py b/tests/models/llava_next/test_modeling_llava_next.py index c93ed8a18e70..a663d4d0d157 100644 --- a/tests/models/llava_next/test_modeling_llava_next.py +++ b/tests/models/llava_next/test_modeling_llava_next.py @@ -200,6 +200,8 @@ class LlavaNextForConditionalGenerationModelTest(ModelTesterMixin, GenerationTes if is_torch_available() else {} ) + # Llava-NeXT merges batch_size and num_patches in the first output dimension + skip_test_image_features_output_shape = True _is_composite = True def setUp(self): diff --git a/tests/models/llava_next_video/test_modeling_llava_next_video.py b/tests/models/llava_next_video/test_modeling_llava_next_video.py index 394be32cbba5..862544a856d6 100644 --- a/tests/models/llava_next_video/test_modeling_llava_next_video.py +++ b/tests/models/llava_next_video/test_modeling_llava_next_video.py @@ -206,6 +206,10 @@ class LlavaNextVideoForConditionalGenerationModelTest(ModelTesterMixin, Generati if is_torch_available() else () ) + # LlavaNextVideo merges batch_size and num_patches in the first output dimension + skip_test_image_features_output_shape = True + # LlavaNextVideo merges batch_size and num_frames in the first output dimension + skip_test_video_features_output_shape = True _is_composite = True @@ -324,6 +328,17 @@ def test_flash_attn_2_fp32_ln(self): def test_flash_attention_2_padding_matches_padding_free_with_position_ids(self): pass + def _video_features_prepare_config_and_inputs(self): + """ + Helper method to extract only image-related inputs from the full set of inputs, for testing `get_image_features`. + + Despite using `pixel_values_videos` in forward, LlavaNextVideo's `get_video_features` method + instead uses `pixel_values` as input, so we need to override the inputs accordingly. + """ + config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common() + inputs_dict = {"pixel_values": inputs_dict["pixel_values_videos"]} + return config, inputs_dict + @require_torch class LlavaNextVideoForConditionalGenerationIntegrationTest(unittest.TestCase): diff --git a/tests/models/llava_onevision/test_modeling_llava_onevision.py b/tests/models/llava_onevision/test_modeling_llava_onevision.py index a4000b5ccef3..99a33b8fc1dd 100644 --- a/tests/models/llava_onevision/test_modeling_llava_onevision.py +++ b/tests/models/llava_onevision/test_modeling_llava_onevision.py @@ -201,6 +201,10 @@ class LlavaOnevisionForConditionalGenerationModelTest(ModelTesterMixin, Generati if is_torch_available() else {} ) + # LlavaOnevision merges batch_size and num_patches in the first output dimension + skip_test_image_features_output_shape = True + # LlavaOnevision merges batch_size and num_frames in the first output dimension + skip_test_video_features_output_shape = True # MP works but offload doesn't work when the MultiheadAttention is offloaded # TODO: One potential solution would be to add to set preload_module_classes = ["Siglip2MultiheadAttentionPoolingHead"] @@ -287,6 +291,26 @@ def test_training_gradient_checkpointing_use_reentrant_true(self): def test_flash_attention_2_padding_matches_padding_free_with_position_ids(self): pass + def _video_features_prepare_config_and_inputs(self): + """ + Helper method to extract only video-related inputs from the full set of inputs, for testing `get_video_features`. + + The superclass method will rename "pixel_values" to "pixel_values_videos" automatically, but LlavaOnevision's + `get_video_features` uses "pixel_values" as input, so we need to override the inputs accordingly. + """ + pixel_values_videos = floats_tensor( + [ + self.model_tester.batch_size, + 8, + self.model_tester.vision_config["num_channels"], + self.model_tester.vision_config["image_size"], + self.model_tester.vision_config["image_size"], + ] + ) + config = self.model_tester.get_config() + inputs_dict = {"pixel_values": pixel_values_videos} + return config, inputs_dict + @require_torch class LlavaOnevisionForConditionalGenerationIntegrationTest(unittest.TestCase): diff --git a/tests/models/metaclip_2/test_modeling_metaclip_2.py b/tests/models/metaclip_2/test_modeling_metaclip_2.py index 825326ac228c..211e6bae096f 100644 --- a/tests/models/metaclip_2/test_modeling_metaclip_2.py +++ b/tests/models/metaclip_2/test_modeling_metaclip_2.py @@ -599,6 +599,26 @@ def test_sdpa_can_dispatch_on_flash(self): def test_sdpa_can_compile_dynamic(self): self.skipTest(reason="MetaClip2 model can't be compiled dynamic, error in metaclip_2_loss`") + @unittest.skip(reason="The MetaCLIP2 family currently does not work with output_attentions.") + def test_get_text_features_attentions(self): + # This test should no longer be skipped once this architecture is refactored to work with output_attentions. + pass + + @unittest.skip(reason="The MetaCLIP2 family currently does not work with output_hidden_states.") + def test_get_text_features_hidden_states(self): + # This test should no longer be skipped once this architecture is refactored to work with output_hidden_states. + pass + + @unittest.skip(reason="The MetaCLIP2 family currently does not work with output_attentions.") + def test_get_image_features_attentions(self): + # This test should no longer be skipped once this architecture is refactored to work with output_attentions. + pass + + @unittest.skip(reason="The MetaCLIP2 family currently does not work with output_hidden_states.") + def test_get_image_features_hidden_states(self): + # This test should no longer be skipped once this architecture is refactored to work with output_hidden_states. + pass + class MetaClip2ForImageClassificationModelTester(MetaClip2ModelTester): def __init__(self, parent): diff --git a/tests/models/mistral3/test_modeling_mistral3.py b/tests/models/mistral3/test_modeling_mistral3.py index db652198c9ce..3159ee012c29 100644 --- a/tests/models/mistral3/test_modeling_mistral3.py +++ b/tests/models/mistral3/test_modeling_mistral3.py @@ -174,6 +174,8 @@ class Mistral3ModelTest(ModelTesterMixin, GenerationTesterMixin, PipelineTesterM if is_torch_available() else {} ) + # Mistral3 merges batch_size and num_patches in index 1, with index 0 hardcoded to 1 + skip_test_image_features_output_shape = True _is_composite = True def setUp(self): diff --git a/tests/models/ovis2/test_modeling_ovis2.py b/tests/models/ovis2/test_modeling_ovis2.py index 9b4717947fc9..a72ce04d0cfd 100644 --- a/tests/models/ovis2/test_modeling_ovis2.py +++ b/tests/models/ovis2/test_modeling_ovis2.py @@ -178,6 +178,8 @@ class Ovis2ModelTest(ModelTesterMixin, GenerationTesterMixin, PipelineTesterMixi if is_torch_available() else {} ) + # Ovis2 post-processes the last_hidden_state to hidden_size * hidden_stride**2 + skip_test_image_features_output_shape = True _is_composite = True def setUp(self): diff --git a/tests/models/perception_lm/test_modeling_perception_lm.py b/tests/models/perception_lm/test_modeling_perception_lm.py index a1aba5565482..7c800f909676 100644 --- a/tests/models/perception_lm/test_modeling_perception_lm.py +++ b/tests/models/perception_lm/test_modeling_perception_lm.py @@ -366,6 +366,15 @@ def test_retain_grad_hidden_states_attentions(self): def test_generate_compilation_all_outputs(self): pass + @unittest.skip("Cannot set output_attentions on timm models.") + def test_get_image_features_attentions(self): + pass + + def _image_features_get_expected_num_hidden_states(self, model_tester=None): + # For models that rely on timm for their vision backend, it's hard to infer how many layers the model has + # from the timm config alone. So, we're just hardcoding the expected number of hidden states here. + return 2 + TEST_MODEL_PATH = "facebook/Perception-LM-1B" diff --git a/tests/models/qwen2_5_omni/test_modeling_qwen2_5_omni.py b/tests/models/qwen2_5_omni/test_modeling_qwen2_5_omni.py index 60a0aab38559..7de1544d6893 100644 --- a/tests/models/qwen2_5_omni/test_modeling_qwen2_5_omni.py +++ b/tests/models/qwen2_5_omni/test_modeling_qwen2_5_omni.py @@ -268,6 +268,7 @@ class Qwen2_5OmniThinkerForConditionalGenerationModelTest( # ) # FIXME @raushan Omni tests take ages because the model is big. Try to make it even smaller pipeline_model_mapping = {} + skip_test_audio_features_output_shape = True # Qwen2_5Omni merges batch_size and audio_output_lengths in index 0 _is_composite = True model_split_percents = [0.5, 0.9] diff --git a/tests/models/qwen3_omni_moe/test_modeling_qwen3_omni_moe.py b/tests/models/qwen3_omni_moe/test_modeling_qwen3_omni_moe.py index 56d9b02177be..34a2e7f5bd7b 100644 --- a/tests/models/qwen3_omni_moe/test_modeling_qwen3_omni_moe.py +++ b/tests/models/qwen3_omni_moe/test_modeling_qwen3_omni_moe.py @@ -262,7 +262,7 @@ class Qwen3OmniMoeThinkerForConditionalGenerationModelTest(ModelTesterMixin, Gen all_model_classes = (Qwen3OmniMoeThinkerForConditionalGeneration,) if is_torch_available() else () all_generative_model_classes = (Qwen3OmniMoeThinkerForConditionalGeneration,) if is_torch_available() else () - + skip_test_audio_features_output_shape = True # Qwen3OmniMoe merges batch_size and audio_output_lengths in index 0 _is_composite = True model_split_percents = [0.5, 0.9] @@ -618,6 +618,36 @@ def test_get_rope_index_video_with_audio(self): self.assertTrue(torch.equal(position_ids, expected_position_ids)) + def _image_features_get_expected_num_attentions(self, model_tester=None): + if model_tester is None: + model_tester = self.model_tester + return model_tester.vision_config["depth"] + + def _image_features_get_expected_num_hidden_states(self, model_tester=None): + if model_tester is None: + model_tester = self.model_tester + return model_tester.vision_config["depth"] + 1 + + def _audio_features_get_expected_num_attentions(self, model_tester=None): + if model_tester is None: + model_tester = self.model_tester + return model_tester.audio_config["encoder_layers"] + + def _audio_features_get_expected_num_hidden_states(self, model_tester=None): + if model_tester is None: + model_tester = self.model_tester + return model_tester.audio_config["encoder_layers"] + 1 + + def _video_features_get_expected_num_attentions(self, model_tester=None): + if model_tester is None: + model_tester = self.model_tester + return model_tester.vision_config["depth"] + + def _video_features_get_expected_num_hidden_states(self, model_tester=None): + if model_tester is None: + model_tester = self.model_tester + return model_tester.vision_config["depth"] + 1 + @require_torch class Qwen3OmniModelIntegrationTest(unittest.TestCase): diff --git a/tests/models/sam2/test_modeling_sam2.py b/tests/models/sam2/test_modeling_sam2.py index a1c895c6c13e..300f872ea082 100644 --- a/tests/models/sam2/test_modeling_sam2.py +++ b/tests/models/sam2/test_modeling_sam2.py @@ -715,6 +715,16 @@ def test_model_from_pretrained(self): def test_sdpa_can_compile_dynamic(self): self.skipTest(reason="SAM2 model can't be compiled dynamic yet") + def _image_features_get_expected_num_attentions(self, model_tester=None): + if model_tester is None: + model_tester = self.model_tester + return sum(model_tester.blocks_per_stage) + + def _image_features_get_expected_num_hidden_states(self, model_tester=None): + if model_tester is None: + model_tester = self.model_tester + return sum(model_tester.blocks_per_stage) + 1 + def prepare_image(): img_url = "https://huggingface.co/datasets/hf-internal-testing/sam2-fixtures/resolve/main/truck.jpg" diff --git a/tests/models/sam3/test_modeling_sam3.py b/tests/models/sam3/test_modeling_sam3.py index 809a2beeb0d8..ae20e522d260 100644 --- a/tests/models/sam3/test_modeling_sam3.py +++ b/tests/models/sam3/test_modeling_sam3.py @@ -339,7 +339,7 @@ def get_config(self): "hidden_size": 32, "intermediate_size": 64, "projection_dim": 32, - "num_hidden_layers": 1, + "num_hidden_layers": self.num_hidden_layers, "num_attention_heads": 4, "max_position_embeddings": 32, # Keep at 32 for stability "hidden_act": "gelu", @@ -750,8 +750,8 @@ def test_forward_with_text_embeds(self): # First get text embeddings with torch.no_grad(): text_embeds = model.get_text_features( - input_ids=inputs_dict["input_ids"], attention_mask=inputs_dict["attention_mask"] - ) + input_ids=inputs_dict["input_ids"], attention_mask=inputs_dict["attention_mask"], return_dict=True + ).pooler_output # Forward with text_embeds (remove input_ids) inputs_with_embeds = { diff --git a/tests/models/seamless_m4t/test_modeling_seamless_m4t.py b/tests/models/seamless_m4t/test_modeling_seamless_m4t.py index 6ddf803879b9..338a8b396ec8 100644 --- a/tests/models/seamless_m4t/test_modeling_seamless_m4t.py +++ b/tests/models/seamless_m4t/test_modeling_seamless_m4t.py @@ -614,7 +614,7 @@ class SeamlessM4TGenerationTest(unittest.TestCase): # test generation of: SeamlessM4TModel, SeamlessM4TForSpeechToSpeech, SeamlessM4TForSpeechToText, SeamlessM4TForTextToSpeech def setUp(self): - self.speech_model_tester = SeamlessM4TModelTester(self, input_modality="speech") + self.audio_model_tester = SeamlessM4TModelTester(self, input_modality="speech") self.text_model_tester = SeamlessM4TModelTester(self, input_modality="text") self.tmpdirname = tempfile.mkdtemp() @@ -648,7 +648,7 @@ def prepare_text_input(self): return config, input_dict def prepare_speech_input(self): - config, inputs, decoder_input_ids, input_mask, lm_labels = self.speech_model_tester.prepare_config_and_inputs() + config, inputs, decoder_input_ids, input_mask, lm_labels = self.audio_model_tester.prepare_config_and_inputs() input_dict = { "input_features": inputs, @@ -661,7 +661,7 @@ def prepare_speech_input(self): return config, input_dict def prepare_speech_and_text_input(self): - config, inputs, decoder_input_ids, input_mask, lm_labels = self.speech_model_tester.prepare_config_and_inputs() + config, inputs, decoder_input_ids, input_mask, lm_labels = self.audio_model_tester.prepare_config_and_inputs() input_speech = { "input_features": inputs, diff --git a/tests/models/seamless_m4t_v2/test_modeling_seamless_m4t_v2.py b/tests/models/seamless_m4t_v2/test_modeling_seamless_m4t_v2.py index bd4852ae3efd..32aa09f615e8 100644 --- a/tests/models/seamless_m4t_v2/test_modeling_seamless_m4t_v2.py +++ b/tests/models/seamless_m4t_v2/test_modeling_seamless_m4t_v2.py @@ -621,7 +621,7 @@ class SeamlessM4Tv2GenerationTest(unittest.TestCase): # test generation of: SeamlessM4Tv2Model, SeamlessM4Tv2ForSpeechToSpeech, SeamlessM4Tv2ForSpeechToText, SeamlessM4Tv2ForTextToSpeech def setUp(self): - self.speech_model_tester = SeamlessM4Tv2ModelTester(self, input_modality="speech") + self.audio_model_tester = SeamlessM4Tv2ModelTester(self, input_modality="speech") self.text_model_tester = SeamlessM4Tv2ModelTester(self, input_modality="text") self.tmpdirname = tempfile.mkdtemp() @@ -672,7 +672,7 @@ def prepare_text_input(self, tgt_lang): return config, input_dict def prepare_speech_input(self): - config, inputs, decoder_input_ids, input_mask, lm_labels = self.speech_model_tester.prepare_config_and_inputs() + config, inputs, decoder_input_ids, input_mask, lm_labels = self.audio_model_tester.prepare_config_and_inputs() input_dict = { "input_features": inputs, @@ -685,7 +685,7 @@ def prepare_speech_input(self): return config, input_dict def prepare_speech_and_text_input(self): - config, inputs, decoder_input_ids, input_mask, lm_labels = self.speech_model_tester.prepare_config_and_inputs() + config, inputs, decoder_input_ids, input_mask, lm_labels = self.audio_model_tester.prepare_config_and_inputs() input_speech = { "input_features": inputs, diff --git a/tests/models/siglip/test_modeling_siglip.py b/tests/models/siglip/test_modeling_siglip.py index 6dda8e6fb689..ef5d69b95dd5 100644 --- a/tests/models/siglip/test_modeling_siglip.py +++ b/tests/models/siglip/test_modeling_siglip.py @@ -504,6 +504,26 @@ def test_load_vision_text_config(self): text_config = SiglipTextConfig.from_pretrained(tmp_dir_name) self.assertDictEqual(config.text_config.to_dict(), text_config.to_dict()) + @unittest.skip(reason="The SigLIP family currently does not work with output_attentions.") + def test_get_text_features_attentions(self): + # This test should no longer be skipped once this architecture is refactored to work with output_attentions. + pass + + @unittest.skip(reason="The SigLIP family currently does not work with output_hidden_states.") + def test_get_text_features_hidden_states(self): + # This test should no longer be skipped once this architecture is refactored to work with output_hidden_states. + pass + + @unittest.skip(reason="The SigLIP family currently does not work with output_attentions.") + def test_get_image_features_attentions(self): + # This test should no longer be skipped once this architecture is refactored to work with output_attentions. + pass + + @unittest.skip(reason="The SigLIP family currently does not work with output_hidden_states.") + def test_get_image_features_hidden_states(self): + # This test should no longer be skipped once this architecture is refactored to work with output_hidden_states. + pass + @slow def test_model_from_pretrained(self): model_name = "google/siglip-base-patch16-224" diff --git a/tests/models/siglip2/test_modeling_siglip2.py b/tests/models/siglip2/test_modeling_siglip2.py index 2c2108407c3d..e2049a89fbbf 100644 --- a/tests/models/siglip2/test_modeling_siglip2.py +++ b/tests/models/siglip2/test_modeling_siglip2.py @@ -577,6 +577,26 @@ def test_load_vision_text_config(self): text_config = Siglip2TextConfig.from_pretrained(tmp_dir_name) self.assertDictEqual(config.text_config.to_dict(), text_config.to_dict()) + @unittest.skip(reason="The SigLIP2 family currently does not work with output_attentions.") + def test_get_text_features_attentions(self): + # This test should no longer be skipped once this architecture is refactored to work with output_attentions. + pass + + @unittest.skip(reason="The SigLIP2 family currently does not work with output_hidden_states.") + def test_get_text_features_hidden_states(self): + # This test should no longer be skipped once this architecture is refactored to work with output_hidden_states. + pass + + @unittest.skip(reason="The SigLIP2 family currently does not work with output_attentions.") + def test_get_image_features_attentions(self): + # This test should no longer be skipped once this architecture is refactored to work with output_attentions. + pass + + @unittest.skip(reason="The SigLIP2 family currently does not work with output_hidden_states.") + def test_get_image_features_hidden_states(self): + # This test should no longer be skipped once this architecture is refactored to work with output_hidden_states. + pass + @slow def test_model_from_pretrained(self): model_name = "google/siglip2-base-patch16-naflex" diff --git a/tests/models/smolvlm/test_modeling_smolvlm.py b/tests/models/smolvlm/test_modeling_smolvlm.py index e61db1f9022a..5b40e5c993d3 100644 --- a/tests/models/smolvlm/test_modeling_smolvlm.py +++ b/tests/models/smolvlm/test_modeling_smolvlm.py @@ -168,7 +168,7 @@ class SmolVLMModelTest(ModelTesterMixin, unittest.TestCase): """ all_model_classes = (SmolVLMModel,) if is_torch_available() else () - + skip_test_image_features_output_shape = True # SmolVLM merges batch_size with num_images in index 0 test_resize_embeddings = True def setUp(self): @@ -339,6 +339,7 @@ class SmolVLMForConditionalGenerationModelTest( if is_torch_available() else () ) + skip_test_image_features_output_shape = True # SmolVLM merges batch_size with num_images in index 0 test_resize_embeddings = True def setUp(self): diff --git a/tests/models/video_llava/test_modeling_video_llava.py b/tests/models/video_llava/test_modeling_video_llava.py index 4d4db3f3a482..a24a6602e910 100644 --- a/tests/models/video_llava/test_modeling_video_llava.py +++ b/tests/models/video_llava/test_modeling_video_llava.py @@ -202,7 +202,8 @@ class VideoLlavaForConditionalGenerationModelTest(ModelTesterMixin, GenerationTe if is_torch_available() else () ) - + # VideoLlava merges batch_size and num_frames in the first output dimension + skip_test_video_features_output_shape = True test_resize_embeddings = True _is_composite = True diff --git a/tests/models/x_clip/test_modeling_x_clip.py b/tests/models/x_clip/test_modeling_x_clip.py index dbd09deaa4c2..5b4aac347e79 100644 --- a/tests/models/x_clip/test_modeling_x_clip.py +++ b/tests/models/x_clip/test_modeling_x_clip.py @@ -507,7 +507,8 @@ def prepare_config_and_inputs_for_common(self): class XCLIPModelTest(ModelTesterMixin, PipelineTesterMixin, unittest.TestCase): all_model_classes = (XCLIPModel,) if is_torch_available() else () pipeline_model_mapping = {"feature-extraction": XCLIPModel} if is_torch_available() else {} - + # XCLIP merges batch_size and num_frames in the first output dimension + skip_test_video_features_output_shape = True test_resize_embeddings = False test_attention_outputs = False maxdiff = None @@ -567,6 +568,20 @@ def test_model_from_pretrained(self): model = XCLIPModel.from_pretrained(model_name) self.assertIsNotNone(model) + def _video_features_prepare_config_and_inputs(self): + """ + Helper method to extract only video-related inputs from the full set of inputs, for testing `get_video_features`. + + The model_tester.vision_model_tester.prepare_config_and_inputs() method prepares image inputs + where the batch size * time dimension is flattened. So, instead we use the model_tester.prepare_config_and_inputs() + which prepares video inputs with shape (batch_size, num_frames, num_channels, height, width) instead. + """ + config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common() + del inputs_dict["input_ids"] + del inputs_dict["attention_mask"] + del inputs_dict["return_loss"] + return config, inputs_dict + # We will verify our results on a spaghetti video def prepare_video(): diff --git a/tests/test_modeling_common.py b/tests/test_modeling_common.py index 0a5023bd3462..0b1177fe7f92 100755 --- a/tests/test_modeling_common.py +++ b/tests/test_modeling_common.py @@ -109,6 +109,7 @@ CONFIG_NAME, GENERATION_CONFIG_NAME, SAFE_WEIGHTS_NAME, + ModelOutput, is_torch_bf16_available_on_device, is_torch_fp16_available_on_device, ) @@ -4634,6 +4635,715 @@ def test_can_load_from_already_mapped_keys(self): # Make sure both saved state_dict are identical self.assertTrue(compare_state_dicts(model.state_dict(), model_reloaded.state_dict())) + def _text_features_prepare_config_and_inputs(self): + """ + Helper method to extract only text-related inputs from the full set of inputs, for testing `get_text_features`. + + Specifically, it tests both the model_tester and its text_model_tester (if any), + and filters for "input_ids", "token_type_ids", and "attention_mask" keys. + """ + config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common() + if hasattr(self.model_tester, "text_model_tester"): + _, inputs_dict = self.model_tester.text_model_tester.prepare_config_and_inputs_for_common() + else: + inputs_dict = { + key: value + for key, value in inputs_dict.items() + if key in ["input_ids", "token_type_ids", "attention_mask"] + } + return config, inputs_dict + + def _image_features_prepare_config_and_inputs(self): + """ + Helper method to extract only image-related inputs from the full set of inputs, for testing `get_image_features`. + + Specifically, it tests both the model_tester and its vision_model_tester (if any), + and filters for keys related to images. It excludes video-related keys, but allows + "spatial_shapes" and "qformer_input_ids" keys as required by some architectures. + """ + config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common() + if hasattr(self.model_tester, "vision_model_tester"): + _, inputs_dict = self.model_tester.vision_model_tester.prepare_config_and_inputs_for_common() + else: + inputs_dict = { + key: value + for key, value in inputs_dict.items() + if ("pixel" in key or "image" in key) + and "video" not in key + or key in ["spatial_shapes", "qformer_input_ids"] + } + return config, inputs_dict + + def _audio_features_prepare_config_and_inputs(self): + """ + Helper method to extract only audio-related inputs from the full set of inputs, for testing `get_audio_features`. + + Specifically, it tests both the model_tester and its audio_model_tester (if any), + and filters for keys related to audio. + """ + config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common() + if hasattr(self.model_tester, "audio_model_tester"): + _, inputs_dict = self.model_tester.audio_model_tester.prepare_config_and_inputs_for_common() + else: + inputs_dict = { + key: value + for key, value in inputs_dict.items() + if "audio" in key + or "input_values" in key + or "input_features" in key + or key in ["padding_mask", "is_longer", "feature_attention_mask"] + } + return config, inputs_dict + + def _video_features_prepare_config_and_inputs(self): + """ + Helper method to extract only video-related inputs from the full set of inputs, for testing `get_video_features`. + + Specifically, it tests both the model_tester and its video_model_tester (if any), + and filters for keys related to videos. It also handles key renaming for video inputs + if there is no dedicated video_model_tester. + """ + config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common() + if hasattr(self.model_tester, "video_model_tester"): + _, inputs_dict = self.model_tester.video_model_tester.prepare_config_and_inputs_for_common() + else: + key_mappings = { + "pixel_values": "pixel_values_videos", + "image_grid_thw": "video_grid_thw", + "image_merge_sizes": "video_merge_sizes", + } + + for src_key, dst_key in key_mappings.items(): + if src_key in inputs_dict and dst_key not in inputs_dict: + inputs_dict[dst_key] = inputs_dict.pop(src_key) + + allowed_non_video_keys = {"vision_feature_layer", "vision_feature_select_strategy", "cu_seqlens"} + inputs_dict = { + key: value for key, value in inputs_dict.items() if "video" in key or key in allowed_non_video_keys + } + return config, inputs_dict + + def _text_features_get_expected_num_attentions(self, model_tester=None): + if model_tester is None: + model_tester = self.model_tester + + if hasattr(model_tester, "text_model_tester"): + return self._text_features_get_expected_num_attentions(model_tester.text_model_tester) + if hasattr(model_tester, "expected_num_hidden_layers"): + return model_tester.expected_num_hidden_layers - 1 + if hasattr(model_tester, "num_hidden_layers"): + return model_tester.num_hidden_layers + raise ValueError("Cannot determine the expected number of layers for text features") + + def _text_features_get_expected_num_hidden_states(self, model_tester=None): + return self._text_features_get_expected_num_attentions(model_tester) + 1 + + def _image_features_get_expected_num_attentions(self, model_tester=None): + if model_tester is None: + model_tester = self.model_tester + if hasattr(model_tester, "vision_model_tester"): + return self._image_features_get_expected_num_attentions(model_tester.vision_model_tester) + elif ( + hasattr(model_tester, "vision_config") + and isinstance(model_tester.vision_config, dict) + and "num_hidden_layers" in model_tester.vision_config + ): + return model_tester.vision_config["num_hidden_layers"] + + if hasattr(model_tester, "expected_num_hidden_layers"): + return model_tester.expected_num_hidden_layers - 1 + elif hasattr(model_tester, "num_hidden_layers"): + return model_tester.num_hidden_layers + raise ValueError("Cannot determine the expected number of layers for image features") + + def _image_features_get_expected_num_hidden_states(self, model_tester=None): + return self._image_features_get_expected_num_attentions(model_tester) + 1 + + def _audio_features_get_expected_num_attentions(self, model_tester=None): + if model_tester is None: + model_tester = self.model_tester + + if hasattr(model_tester, "audio_model_tester"): + return self._audio_features_get_expected_num_attentions(model_tester.audio_model_tester) + elif ( + hasattr(model_tester, "audio_config") + and isinstance(model_tester.audio_config, dict) + and "num_hidden_layers" in model_tester.audio_config + ): + return model_tester.audio_config["num_hidden_layers"] + + if hasattr(model_tester, "expected_num_hidden_layers"): + return model_tester.expected_num_hidden_layers - 1 + elif hasattr(model_tester, "num_hidden_layers"): + return model_tester.num_hidden_layers + raise ValueError("Cannot determine the expected number of layers for audio features") + + def _audio_features_get_expected_num_hidden_states(self, model_tester=None): + return self._audio_features_get_expected_num_attentions(model_tester) + 1 + + def _video_features_get_expected_num_attentions(self, model_tester=None): + if model_tester is None: + model_tester = self.model_tester + + if hasattr(model_tester, "video_model_tester"): + return self._video_features_get_expected_num_attentions(model_tester.video_model_tester) + if hasattr(model_tester, "vision_model_tester"): + return self._video_features_get_expected_num_attentions(model_tester.vision_model_tester) + elif ( + hasattr(model_tester, "video_config") + and isinstance(model_tester.video_config, dict) + and "num_hidden_layers" in model_tester.video_config + ): + return model_tester.video_config["num_hidden_layers"] + + if hasattr(model_tester, "expected_num_hidden_layers"): + return model_tester.expected_num_hidden_layers - 1 + elif hasattr(model_tester, "num_hidden_layers"): + return model_tester.num_hidden_layers + raise ValueError("Cannot determine the expected number of layers for video features") + + def _video_features_get_expected_num_hidden_states(self, model_tester=None): + return self._video_features_get_expected_num_attentions(model_tester) + 1 + + @parameterized.expand([True, False, None]) + def test_get_text_features_output(self, return_dict: bool | None): + for model_class in self.all_model_classes: + if not hasattr(model_class, "get_text_features"): + continue + + config, inputs_dict = self._text_features_prepare_config_and_inputs() + if return_dict is not None: + config.return_dict = return_dict + + model = model_class(config).eval() + model = model.to(torch_device) + + torch.manual_seed(0) + with torch.no_grad(): + outputs = model.get_text_features(**inputs_dict) + + if return_dict in (True, None): + self.assertTrue(isinstance(outputs, ModelOutput), "get_text_features() must return a BaseModelOutput") + self.assertTrue( + hasattr(outputs, "last_hidden_state"), + "get_text_features() must return a BaseModelOutput with last_hidden_state", + ) + self.assertTrue( + hasattr(outputs, "pooler_output"), + "get_text_features() must return a BaseModelOutput with pooler_output", + ) + self.assertTrue( + hasattr(outputs, "hidden_states"), + "get_text_features() must return a BaseModelOutput with hidden_states", + ) + if self.has_attentions: + self.assertTrue( + hasattr(outputs, "attentions"), + "get_text_features() must return a BaseModelOutput with attentions", + ) + + # Test against (batch_size, seq_len, hidden_size) + last_hidden_state = outputs.last_hidden_state + expected_hidden_size = config.text_config.hidden_size + expected_shape = ( + inputs_dict["input_ids"].shape[0], + inputs_dict["input_ids"].shape[1], + expected_hidden_size, + ) + self.assertEqual(last_hidden_state.shape, expected_shape, "last_hidden_state shape mismatch") + + else: + self.assertIsInstance(outputs, tuple, "get_text_features() must return a tuple if return_dict=False") + + def test_get_text_features_hidden_states(self): + def check_hidden_states_output(inputs_dict, config, model_class): + model = model_class(copy.deepcopy(config)) + model.to(torch_device) + model.eval() + + with torch.no_grad(): + outputs = model.get_text_features(**inputs_dict) + # hidden_states = outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states + hidden_states = outputs.hidden_states + expected_num_hidden_states = self._text_features_get_expected_num_hidden_states() + self.assertIsNotNone(hidden_states, "hidden_states should not be None") + self.assertEqual(len(hidden_states), expected_num_hidden_states, "Number of hidden states layers mismatch") + + for model_class in self.all_model_classes: + if not hasattr(model_class, "get_text_features"): + continue + + config, inputs_dict = self._text_features_prepare_config_and_inputs() + + inputs_dict["output_hidden_states"] = True + check_hidden_states_output(inputs_dict, config, model_class) + + # check that output_hidden_states also work using config + del inputs_dict["output_hidden_states"] + config.output_hidden_states = True + for k in config.sub_configs: + if getattr(config, k) is not None: + getattr(config, k).output_hidden_states = True + + check_hidden_states_output(inputs_dict, config, model_class) + + def test_get_text_features_attentions(self): + def check_attentions_output(inputs_dict, config, model_class): + model = model_class(copy.deepcopy(config)) + model.set_attn_implementation("eager") + model.to(torch_device) + model.eval() + + with torch.no_grad(): + outputs = model.get_text_features(**inputs_dict) + attentions = outputs.attentions + # model.text_model(**inputs_dict) also no attentions for aimv2 + expected_num_attentions = self._text_features_get_expected_num_attentions() + self.assertIsNotNone(attentions, "attentions should not be None") + self.assertEqual(len(attentions), expected_num_attentions, "Number of attention layers mismatch") + + if not self.has_attentions: + return + + for model_class in self.all_model_classes: + if not hasattr(model_class, "get_text_features"): + continue + + config, inputs_dict = self._text_features_prepare_config_and_inputs() + inputs_dict["output_hidden_states"] = False + inputs_dict["output_attentions"] = True + check_attentions_output(inputs_dict, config, model_class) + + # check that output_attentions also work using config + del inputs_dict["output_attentions"] + config.output_attentions = True + for k in config.sub_configs: + if getattr(config, k) is not None: + getattr(config, k).output_attentions = True + + check_attentions_output(inputs_dict, config, model_class) + + @parameterized.expand([True, False, None]) + def test_get_image_features_output(self, return_dict: bool | None): + for model_class in self.all_model_classes: + if not hasattr(model_class, "get_image_features"): + continue + + config, inputs_dict = self._image_features_prepare_config_and_inputs() + if return_dict is not None: + config.return_dict = return_dict + + model = model_class(config).eval() + model = model.to(torch_device) + + torch.manual_seed(0) + with torch.no_grad(): + outputs = model.get_image_features(**inputs_dict) + + if return_dict in (True, None): + self.assertTrue(isinstance(outputs, ModelOutput), "get_image_features() must return a BaseModelOutput") + self.assertTrue( + hasattr(outputs, "last_hidden_state"), + "get_image_features() must return a BaseModelOutput with last_hidden_state", + ) + self.assertTrue( + hasattr(outputs, "pooler_output"), + "get_image_features() must return a BaseModelOutput with pooler_output", + ) + self.assertTrue( + hasattr(outputs, "hidden_states"), + "get_image_features() must return a BaseModelOutput with hidden_states", + ) + if self.has_attentions: + self.assertTrue( + hasattr(outputs, "attentions"), + "get_image_features() must return a BaseModelOutput with attentions", + ) + + if getattr(self, "skip_test_image_features_output_shape", False): + return + + last_hidden_state_shape = outputs.last_hidden_state.shape + batch_size = ( + inputs_dict["pixel_values"].shape[0] + if "pixel_values" in inputs_dict + else inputs_dict["pixel_values_images"].shape[0] + ) + self.assertEqual( + last_hidden_state_shape[0], + batch_size, + f"batch_size mismatch, full shape: {last_hidden_state_shape}", + ) + + vision_config = config.vision_config if hasattr(config, "vision_config") else config + vision_config = ( + vision_config.backbone_config if hasattr(vision_config, "backbone_config") else vision_config + ) + vision_config = vision_config.vq_config if hasattr(vision_config, "vq_config") else vision_config + vision_config = vision_config.model_args if hasattr(vision_config, "model_args") else vision_config + attribute_candidates = [ + "embed_dim_per_stage", + "embed_dim", + "embed_dims", + "out_hidden_size", + "hidden_size", + "hidden_dim", + ] + hidden_size = None + for attr in attribute_candidates: + if hasattr(vision_config, attr): + hidden_size = getattr(vision_config, attr) + break + elif isinstance(vision_config, dict) and attr in vision_config: + hidden_size = vision_config[attr] + break + else: + raise ValueError("Cannot find the hidden size attribute in vision_config") + if isinstance(hidden_size, (list, tuple)): + hidden_size = hidden_size[-1] + self.assertEqual( + last_hidden_state_shape[-1], + hidden_size, + f"hidden_size mismatch, full shape: {last_hidden_state_shape}", + ) + + else: + self.assertIsInstance(outputs, tuple, "get_image_features() must return a tuple if return_dict=False") + + def test_get_image_features_hidden_states(self): + def check_hidden_states_output(inputs_dict, config, model_class): + model = model_class(copy.deepcopy(config)) + model.to(torch_device) + model.eval() + + with torch.no_grad(): + outputs = model.get_image_features(**inputs_dict) + # hidden_states = outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states + hidden_states = outputs.hidden_states + expected_num_hidden_states = self._image_features_get_expected_num_hidden_states() + self.assertIsNotNone(hidden_states, "hidden_states should not be None") + self.assertEqual(len(hidden_states), expected_num_hidden_states, "Number of hidden states layers mismatch") + + for model_class in self.all_model_classes: + if not hasattr(model_class, "get_image_features"): + continue + + config, inputs_dict = self._image_features_prepare_config_and_inputs() + + inputs_dict["output_hidden_states"] = True + check_hidden_states_output(inputs_dict, config, model_class) + + # check that output_hidden_states also work using config + del inputs_dict["output_hidden_states"] + config.output_hidden_states = True + for k in config.sub_configs: + if getattr(config, k) is not None: + getattr(config, k).output_hidden_states = True + + check_hidden_states_output(inputs_dict, config, model_class) + + def test_get_image_features_attentions(self): + def check_attentions_output(inputs_dict, config, model_class): + model = model_class(copy.deepcopy(config)) + model.set_attn_implementation("eager") + model.to(torch_device) + model.eval() + + with torch.no_grad(): + outputs = model.get_image_features(**inputs_dict) + attentions = outputs.attentions + # model.text_model(**inputs_dict) also no attentions for aimv2 + expected_num_attentions = self._image_features_get_expected_num_attentions() + self.assertIsNotNone(attentions, "attentions should not be None") + self.assertEqual(len(attentions), expected_num_attentions, "Number of attention layers mismatch") + + if not self.has_attentions: + return + + for model_class in self.all_model_classes: + if not hasattr(model_class, "get_image_features"): + continue + + config, inputs_dict = self._image_features_prepare_config_and_inputs() + inputs_dict["output_hidden_states"] = False + inputs_dict["output_attentions"] = True + check_attentions_output(inputs_dict, config, model_class) + + # check that output_attentions also work using config + del inputs_dict["output_attentions"] + config.output_attentions = True + for k in config.sub_configs: + if getattr(config, k) is not None: + getattr(config, k).output_attentions = True + + check_attentions_output(inputs_dict, config, model_class) + + @parameterized.expand([True, False, None]) + def test_get_audio_features_output(self, return_dict: bool | None): + for model_class in self.all_model_classes: + if not hasattr(model_class, "get_audio_features"): + continue + + config, inputs_dict = self._audio_features_prepare_config_and_inputs() + if return_dict is not None: + config.return_dict = return_dict + + model = model_class(config).eval() + model = model.to(torch_device) + + torch.manual_seed(0) + with torch.no_grad(): + outputs = model.get_audio_features(**inputs_dict) + + if return_dict in (True, None): + self.assertTrue(isinstance(outputs, ModelOutput), "get_audio_features() must return a BaseModelOutput") + self.assertTrue( + hasattr(outputs, "last_hidden_state"), + "get_audio_features() must return a BaseModelOutput with last_hidden_state", + ) + self.assertTrue( + hasattr(outputs, "pooler_output"), + "get_audio_features() must return a BaseModelOutput with pooler_output", + ) + self.assertTrue( + hasattr(outputs, "hidden_states"), + "get_audio_features() must return a BaseModelOutput with hidden_states", + ) + if self.has_attentions: + self.assertTrue( + hasattr(outputs, "attentions"), + "get_audio_features() must return a BaseModelOutput with attentions", + ) + + if getattr(self, "skip_test_audio_features_output_shape", False): + return + + last_hidden_state_shape = outputs.last_hidden_state.shape + batch_size = inputs_dict["input_features"].shape[0] + self.assertEqual( + last_hidden_state_shape[0], + batch_size, + f"batch_size mismatch, full shape: {last_hidden_state_shape}", + ) + + audio_config = config.audio_config if hasattr(config, "audio_config") else config + if hasattr(audio_config, "projection_dim"): + hidden_size = audio_config.projection_dim + elif hasattr(audio_config, "hidden_size"): + hidden_size = audio_config.hidden_size + elif hasattr(audio_config, "encoder_config"): + hidden_size = audio_config.encoder_config.hidden_dim + elif hasattr(audio_config, "encoder_ffn_dim"): + hidden_size = audio_config.encoder_ffn_dim + self.assertEqual( + last_hidden_state_shape[-1], + hidden_size, + f"hidden_size mismatch, full shape: {last_hidden_state_shape}", + ) + + else: + self.assertIsInstance(outputs, tuple, "get_audio_features() must return a tuple if return_dict=False") + + def test_get_audio_features_hidden_states(self): + def check_hidden_states_output(inputs_dict, config, model_class): + model = model_class(copy.deepcopy(config)) + model.to(torch_device) + model.eval() + + with torch.no_grad(): + outputs = model.get_audio_features(**inputs_dict) + hidden_states = outputs.hidden_states + expected_num_hidden_states = self._audio_features_get_expected_num_hidden_states() + self.assertIsNotNone(hidden_states, "hidden_states should not be None") + self.assertEqual(len(hidden_states), expected_num_hidden_states, "Number of hidden states layers mismatch") + + for model_class in self.all_model_classes: + if not hasattr(model_class, "get_audio_features"): + continue + + config, inputs_dict = self._audio_features_prepare_config_and_inputs() + + inputs_dict["output_hidden_states"] = True + check_hidden_states_output(inputs_dict, config, model_class) + + # check that output_hidden_states also work using config + del inputs_dict["output_hidden_states"] + config.output_hidden_states = True + for k in config.sub_configs: + if getattr(config, k) is not None: + getattr(config, k).output_hidden_states = True + + check_hidden_states_output(inputs_dict, config, model_class) + + def test_get_audio_features_attentions(self): + def check_attentions_output(inputs_dict, config, model_class): + model = model_class(copy.deepcopy(config)) + model.set_attn_implementation("eager") + model.to(torch_device) + model.eval() + + with torch.no_grad(): + outputs = model.get_audio_features(**inputs_dict) + attentions = outputs.attentions + expected_num_attentions = self._audio_features_get_expected_num_attentions() + self.assertIsNotNone(attentions, "attentions should not be None") + self.assertEqual(len(attentions), expected_num_attentions, "Number of attention layers mismatch") + + if not self.has_attentions: + return + + for model_class in self.all_model_classes: + if not hasattr(model_class, "get_audio_features"): + continue + + config, inputs_dict = self._audio_features_prepare_config_and_inputs() + inputs_dict["output_hidden_states"] = False + inputs_dict["output_attentions"] = True + check_attentions_output(inputs_dict, config, model_class) + + # check that output_attentions also work using config + del inputs_dict["output_attentions"] + config.output_attentions = True + for k in config.sub_configs: + if getattr(config, k) is not None: + getattr(config, k).output_attentions = True + + check_attentions_output(inputs_dict, config, model_class) + + @parameterized.expand([True, False, None]) + def test_get_video_features_output(self, return_dict: bool | None): + for model_class in self.all_model_classes: + if not hasattr(model_class, "get_video_features"): + continue + + config, inputs_dict = self._video_features_prepare_config_and_inputs() + if return_dict is not None: + config.return_dict = return_dict + + model = model_class(config).eval() + model = model.to(torch_device) + + torch.manual_seed(0) + with torch.no_grad(): + outputs = model.get_video_features(**inputs_dict) + + if return_dict in (True, None): + self.assertTrue(isinstance(outputs, ModelOutput), "get_video_features() must return a BaseModelOutput") + self.assertTrue( + hasattr(outputs, "last_hidden_state"), + "get_video_features() must return a BaseModelOutput with last_hidden_state", + ) + self.assertTrue( + hasattr(outputs, "pooler_output"), + "get_video_features() must return a BaseModelOutput with pooler_output", + ) + self.assertTrue( + hasattr(outputs, "hidden_states"), + "get_video_features() must return a BaseModelOutput with hidden_states", + ) + if self.has_attentions: + self.assertTrue( + hasattr(outputs, "attentions"), + "get_video_features() must return a BaseModelOutput with attentions", + ) + + if getattr(self, "skip_test_video_features_output_shape", False): + return + + last_hidden_state_shape = outputs.last_hidden_state.shape + if "pixel_values_videos" in inputs_dict: + batch_size = inputs_dict["pixel_values_videos"].shape[0] + elif "pixel_values" in inputs_dict: + batch_size = inputs_dict["pixel_values"].shape[0] + self.assertEqual( + last_hidden_state_shape[0], + batch_size, + f"batch_size mismatch, full shape: {last_hidden_state_shape}", + ) + video_config = config + if hasattr(config, "video_config"): + video_config = config.video_config + elif hasattr(config, "vision_config"): + video_config = config.vision_config + if hasattr(video_config, "out_hidden_size"): + hidden_size = video_config.out_hidden_size + else: + hidden_size = video_config.hidden_size + self.assertEqual( + last_hidden_state_shape[-1], + hidden_size, + f"hidden_size mismatch, full shape: {last_hidden_state_shape}", + ) + + else: + self.assertIsInstance(outputs, tuple, "get_video_features() must return a tuple if return_dict=False") + + def test_get_video_features_hidden_states(self): + def check_hidden_states_output(inputs_dict, config, model_class): + model = model_class(copy.deepcopy(config)) + model.to(torch_device) + model.eval() + + with torch.no_grad(): + outputs = model.get_video_features(**inputs_dict) + hidden_states = outputs.hidden_states + expected_num_hidden_states = self._video_features_get_expected_num_hidden_states() + self.assertIsNotNone(hidden_states, "hidden_states should not be None") + self.assertEqual(len(hidden_states), expected_num_hidden_states, "Number of hidden states layers mismatch") + + for model_class in self.all_model_classes: + if not hasattr(model_class, "get_video_features"): + continue + + config, inputs_dict = self._video_features_prepare_config_and_inputs() + + inputs_dict["output_hidden_states"] = True + check_hidden_states_output(inputs_dict, config, model_class) + + # check that output_hidden_states also work using config + del inputs_dict["output_hidden_states"] + config.output_hidden_states = True + for k in config.sub_configs: + if getattr(config, k) is not None: + getattr(config, k).output_hidden_states = True + + check_hidden_states_output(inputs_dict, config, model_class) + + def test_get_video_features_attentions(self): + def check_attentions_output(inputs_dict, config, model_class): + model = model_class(copy.deepcopy(config)) + model.set_attn_implementation("eager") + model.to(torch_device) + model.eval() + + with torch.no_grad(): + outputs = model.get_video_features(**inputs_dict) + attentions = outputs.attentions + expected_num_attentions = self._video_features_get_expected_num_attentions() + self.assertIsNotNone(attentions, "attentions should not be None") + self.assertEqual(len(attentions), expected_num_attentions, "Number of attention layers mismatch") + + if not self.has_attentions: + return + + for model_class in self.all_model_classes: + if not hasattr(model_class, "get_video_features"): + continue + + config, inputs_dict = self._video_features_prepare_config_and_inputs() + inputs_dict["output_hidden_states"] = False + inputs_dict["output_attentions"] = True + check_attentions_output(inputs_dict, config, model_class) + + # check that output_attentions also work using config + del inputs_dict["output_attentions"] + config.output_attentions = True + for k in config.sub_configs: + if getattr(config, k) is not None: + getattr(config, k).output_attentions = True + + check_attentions_output(inputs_dict, config, model_class) + global_rng = random.Random() diff --git a/utils/check_config_attributes.py b/utils/check_config_attributes.py index 996e14c5e7f5..d2e0d83e1719 100644 --- a/utils/check_config_attributes.py +++ b/utils/check_config_attributes.py @@ -164,6 +164,10 @@ "pretraining_tp", "use_sliding_window", "max_window_layers", + # vision attributes that may be used indirectly via check_model_inputs + "vision_feature_layer", + "vision_feature_select_strategy", + "vision_aspect_ratio", )