@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / e7767eccae
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 6, 05:45 UTC

Fix AudioFlamingo3/MusicFlamingo HF parity and RoTE handling (#37643)

Signed-off-by: Lasha <26011196+lashahub@users.noreply.github.com>

Lasha Koroshinadze committed Mar 22, 2026 at 22:29 UTC e7767eccae3c6f6a8fdf659407ea70a2ef0953b1
12 files changed +1159 -245
docs/models/supported_models.md
+2 -1
@@ -535,7 +535,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
535 | Architecture | Models | Inputs | Example HF Models | [LoRA](../features/lora.md) | [PP](../serving/parallelism_scaling.md) |
536 | ------------ | ------ | ------ | ----------------- | -------------------- | ------------------------- |
537 | `AriaForConditionalGeneration` | Aria | T + I<sup>+</sup> | `rhymes-ai/Aria` | | |
538 -| `AudioFlamingo3ForConditionalGeneration` | AudioFlamingo3 | T + A | `nvidia/audio-flamingo-3-hf`, `nvidia/music-flamingo-2601-hf` | ✅︎ | ✅︎ |
538 +| `AudioFlamingo3ForConditionalGeneration` | AudioFlamingo3 | T + A | `nvidia/audio-flamingo-3-hf`, `nvidia/music-flamingo-hf` | ✅︎ | ✅︎ |
539 | `AyaVisionForConditionalGeneration` | Aya Vision | T + I<sup>+</sup> | `CohereLabs/aya-vision-8b`, `CohereLabs/aya-vision-32b`, etc. | | ✅︎ |
540 | `BagelForConditionalGeneration` | BAGEL | T + I<sup>+</sup> | `ByteDance-Seed/BAGEL-7B-MoT` | ✅︎ | ✅︎ |
541 | `BeeForConditionalGeneration` | Bee-8B | T + I<sup>E+</sup> | `Open-Bee/Bee-8B-RL`, `Open-Bee/Bee-8B-SFT` | | ✅︎ |
@@ -586,6 +586,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
586 | `Mistral3ForConditionalGeneration` | Mistral3 (HF Transformers) | T + I<sup>+</sup> | `mistralai/Mistral-Small-3.1-24B-Instruct-2503`, etc. | ✅︎ | ✅︎ |
587 | `MolmoForCausalLM` | Molmo | T + I<sup>+</sup> | `allenai/Molmo-7B-D-0924`, `allenai/Molmo-7B-O-0924`, etc. | ✅︎ | ✅︎ |
588 | `Molmo2ForConditionalGeneration` | Molmo2 | T + I<sup>+</sup> / V | `allenai/Molmo2-4B`, `allenai/Molmo2-8B`, `allenai/Molmo2-O-7B` | ✅︎ | ✅︎ |
589 +| `MusicFlamingoForConditionalGeneration` | MusicFlamingo | T + A | `nvidia/music-flamingo-2601-hf`, `nvidia/music-flamingo-think-2601-hf` | ✅︎ | ✅︎ |
590 | `NVLM_D_Model` | NVLM-D 1.0 | T + I<sup>+</sup> | `nvidia/NVLM-D-72B`, etc. | | ✅︎ |
591 | `OpenCUAForConditionalGeneration` | OpenCUA-7B | T + I<sup>E+</sup> | `xlangai/OpenCUA-7B` | ✅︎ | ✅︎ |
592 | `OpenPanguVLForConditionalGeneration` | openpangu-VL | T + I<sup>E+</sup> + V<sup>E+</sup> | `FreedomIntelligence/openPangu-VL-7B` | ✅︎ | ✅︎ |
examples/offline_inference/audio_language.py
+12 -2
@@ -104,12 +104,22 @@ def run_musicflamingo(question: str, audio_count: int) -> ModelRequestData:
104 enforce_eager=True,
105 )
106
107 - # MusicFlamingo uses <sound> token for audio
107 + # MusicFlamingo prompt placeholders use <sound>; vLLM's MusicFlamingo
108 + # multimodal processor expands each one into <|sound_bos|> + audio tokens +
109 + # <|sound_eos|> based on extracted audio feature lengths.
110 audio_placeholder = "<sound>" * audio_count
111 + system_prompt = (
112 + "You are Music Flamingo, a multimodal assistant for language and music. "
113 + "On each turn you receive an audio clip which contains music and optional "
114 + "text, you will receive at least one or both; use your world knowledge and "
115 + "reasoning to help the user with any task. Interpret the entirety of the "
116 + "content any input music--regardlenss of whether the user calls it audio, "
117 + "music, or sound."
118 + )
119
120 prompt = (
121 "<|im_start|>system\n"
112 - "You are a helpful assistant.<|im_end|>\n"
122 + f"{system_prompt}<|im_end|>\n"
123 "<|im_start|>user\n"
124 f"{audio_placeholder}{question}<|im_end|>\n"
125 "<|im_start|>assistant\n"
tests/models/fixtures/audioflamingo3/expected_results_single.json
+1 -1
@@ -1 +1 @@
1 -{"transcriptions": ["The content of the input audio is 'you can ask why over and over and over again forever even if one day we explain every physical interaction and scientific law and hope and dream and regret with a single elegant equation'."], "token_ids": [[785, 2213, 315, 279, 1946, 7699, 374, 364, 9330, 646, 2548, 3170, 916, 323, 916, 323, 916, 1549, 15683, 1496, 421, 825, 1899, 582, 10339, 1449, 6961, 16230, 323, 12344, 2329, 323, 3900, 323, 7904, 323, 22231, 448, 264, 3175, 25777, 23606, 4427, 151645]]}
\ No newline at end of file
1 +{"transcriptions": ["There is no clear relationship between the barking and the music, as they seem to be independent of each other."], "token_ids": [[3862, 374, 902, 2797, 5025, 1948, 279, 293, 33452, 323, 279, 4627, 11, 438, 807, 2803, 311, 387, 9489, 315, 1817, 1008, 13, 151645]]}
tests/models/fixtures/musicflamingo/expected_results_batched.json new
+1
@@ -0,0 +1 @@
1 +{"transcriptions": ["This track is an energetic Eurodance / Dance‑Pop anthem that blends the bright, melodic sensibilities of mainstream pop with the driving, club‑ready pulse of classic Eurodance. The duration of the piece is ", "**Verse 1**\nMidnight cravings in bloom, lights flicker in the room, pepperoni dreams arise, pizza party on your skies\n\n**Verse 2**\nCheese melts on the crust, in flavor we trust, boxes stacked to the"], "token_ids": [[1986, 3754, 374, 458, 44855, 19461, 98875, 378, 107, 14, 378, 107, 35, 681, 55964, 11598, 55564, 429, 57843, 279, 9906, 11, 10581, 52760, 6097, 13450, 315, 20729, 2420, 448, 279, 9842, 11, 6335, 55964, 2307, 27235, 315, 11416, 19461, 98875, 13, 220, 576, 8090, 315, 279, 6573, 374, 220], [334, 68043, 220, 16, 1019, 33648, 9287, 88828, 304, 51454, 11, 12711, 28347, 261, 304, 279, 3054, 11, 24353, 20783, 18707, 30789, 11, 22502, 4614, 389, 697, 49293, 271, 334, 68043, 220, 17, 1019, 26843, 2367, 98091, 389, 279, 39612, 11, 304, 17172, 582, 6950, 11, 14697, 41315, 311, 279]]}
tests/models/fixtures/musicflamingo/expected_results_single.json new
+1
@@ -0,0 +1 @@
1 +{"transcriptions": ["This track is an energetic Eurodance / Dance‑Pop anthem that blends the bright, melodic sensibilities of mainstream pop with the driving, club‑ready pulse of classic Eurodance. The duration of the piece is "], "token_ids": [[1986, 3754, 374, 458, 44855, 19461, 98875, 378, 107, 14, 378, 107, 35, 681, 55964, 11598, 55564, 429, 57843, 279, 9906, 11, 10581, 52760, 6097, 13450, 315, 20729, 2420, 448, 279, 9842, 11, 6335, 55964, 2307, 27235, 315, 11416, 19461, 98875, 13, 220, 576, 8090, 315, 279, 6573, 374, 220]]}
tests/models/multimodal/generation/test_audioflamingo3.py
+89 -60
@@ -26,6 +26,54 @@ from tests.models.registry import HF_EXAMPLE_MODELS
26 from vllm import LLM, SamplingParams
27
28 MODEL_NAME = "nvidia/audio-flamingo-3-hf"
29 +SINGLE_CONVERSATION = [
30 + {
31 + "role": "user",
32 + "content": [
33 + {
34 + "type": "text",
35 + "text": "What is surprising about the relationship between "
36 + "the barking and the music?",
37 + },
38 + {
39 + "type": "audio_url",
40 + "audio_url": {
41 + "url": "https://huggingface.co/datasets/nvidia/AudioSkills/"
42 + "resolve/main/assets/"
43 + "dogs_barking_in_sync_with_the_music.wav",
44 + },
45 + },
46 + ],
47 + }
48 +]
49 +BATCHED_CONVERSATIONS = [
50 + SINGLE_CONVERSATION,
51 + [
52 + {
53 + "role": "user",
54 + "content": [
55 + {
56 + "type": "text",
57 + "text": "Why is the philosopher's name mentioned in the "
58 + "lyrics? (A) To express a sense of nostalgia "
59 + "(B) To indicate that language cannot express clearly, "
60 + "satirizing the inversion of black and white in the world "
61 + "(C) To add depth and complexity to the lyrics "
62 + "(D) To showcase the wisdom and influence of the "
63 + "philosopher",
64 + },
65 + {
66 + "type": "audio_url",
67 + "audio_url": {
68 + "url": "https://huggingface.co/datasets/nvidia/"
69 + "AudioSkills/resolve/main/assets/"
70 + "Ch6Ae9DT6Ko_00-04-03_00-04-31.wav",
71 + },
72 + },
73 + ],
74 + }
75 + ],
76 +]
77
78
79 def get_fixture_path(filename):
@@ -34,21 +82,29 @@ def get_fixture_path(filename):
82 )
83
84
85 +def assert_output_matches(output, expected_text, expected_token_ids):
86 + generated = output.outputs[0]
87 + assert generated.text.strip() == expected_text
88 + actual_token_ids = list(generated.token_ids)
89 + assert (
90 + actual_token_ids == expected_token_ids
91 + or actual_token_ids == expected_token_ids[:-1]
92 + or actual_token_ids[:-1] == expected_token_ids
93 + )
94 +
95 +
96 @pytest.fixture(scope="module")
97 def llm():
39 - # Check if the model is supported by the current transformers version
98 model_info = HF_EXAMPLE_MODELS.get_hf_info("AudioFlamingo3ForConditionalGeneration")
99 model_info.check_transformers_version(on_fail="skip")
100
101 try:
44 - llm = LLM(
102 + return LLM(
103 model=MODEL_NAME,
46 - trust_remote_code=True,
104 dtype="bfloat16",
105 enforce_eager=True,
106 limit_mm_per_prompt={"audio": 1},
107 )
51 - return llm
108 except Exception as e:
109 pytest.skip(f"Failed to load model {MODEL_NAME}: {e}")
110
@@ -61,29 +117,17 @@ def test_single_generation(llm):
117 with open(fixture_path) as f:
118 expected = json.load(f)
119
64 - audio_url = "https://huggingface.co/datasets/nvidia/AudioSkills/resolve/main/assets/Why_do_we_ask_questions_converted.wav"
65 -
66 - messages = [
67 - {
68 - "role": "user",
69 - "content": [
70 - {"type": "audio_url", "audio_url": {"url": audio_url}},
71 - {"type": "text", "text": "Transcribe the input speech."},
72 - ],
73 - }
74 - ]
75 -
120 sampling_params = SamplingParams(temperature=0.0, max_tokens=128)
121
122 outputs = llm.chat(
79 - messages=messages,
123 + messages=SINGLE_CONVERSATION,
124 sampling_params=sampling_params,
125 )
82 - generated_text = outputs[0].outputs[0].text.strip()
83 -
84 - expected_text = expected["transcriptions"][0]
85 -
86 - assert expected_text in generated_text or generated_text in expected_text
126 + assert_output_matches(
127 + outputs[0],
128 + expected["transcriptions"][0],
129 + expected["token_ids"][0],
130 + )
131
132
133 def test_batched_generation(llm):
@@ -94,49 +138,34 @@ def test_batched_generation(llm):
138 with open(fixture_path) as f:
139 expected = json.load(f)
140
97 - items = [
98 - {
99 - "audio_url": "https://huggingface.co/datasets/nvidia/AudioSkills/resolve/main/assets/dogs_barking_in_sync_with_the_music.wav",
100 - "question": "What is surprising about the relationship "
101 - "between the barking and the music?",
102 - "expected_idx": 0,
103 - },
104 - {
105 - "audio_url": "https://huggingface.co/datasets/nvidia/AudioSkills/resolve/main/assets/Ch6Ae9DT6Ko_00-04-03_00-04-31.wav",
106 - "question": (
107 - "Why is the philosopher's name mentioned in the lyrics? "
108 - "(A) To express a sense of nostalgia "
109 - "(B) To indicate that language cannot express clearly, "
110 - "satirizing the inversion of black and white in the world "
111 - "(C) To add depth and complexity to the lyrics "
112 - "(D) To showcase the wisdom and influence of the philosopher"
113 - ),
114 - "expected_idx": 1,
115 - },
116 - ]
117 -
118 - conversations = []
119 - for item in items:
120 - messages = [
121 - {
122 - "role": "user",
123 - "content": [
124 - {"type": "audio_url", "audio_url": {"url": item["audio_url"]}},
125 - {"type": "text", "text": item["question"]},
126 - ],
127 - }
128 - ]
129 - conversations.append(messages)
130 -
141 sampling_params = SamplingParams(temperature=0.0, max_tokens=128)
142
143 outputs = llm.chat(
134 - messages=conversations,
144 + messages=BATCHED_CONVERSATIONS,
145 sampling_params=sampling_params,
146 )
147
148 for i, output in enumerate(outputs):
139 - generated_text = output.outputs[0].text.strip()
140 - expected_text = expected["transcriptions"][i]
149 + assert_output_matches(
150 + output,
151 + expected["transcriptions"][i],
152 + expected["token_ids"][i],
153 + )
154 +
155 +
156 +def test_single_and_batched_generation_match(llm):
157 + sampling_params = SamplingParams(temperature=0.0, max_tokens=128)
158
142 - assert expected_text in generated_text or generated_text in expected_text
159 + single_output = llm.chat(
160 + messages=SINGLE_CONVERSATION,
161 + sampling_params=sampling_params,
162 + )[0]
163 + batched_output = llm.chat(
164 + messages=BATCHED_CONVERSATIONS,
165 + sampling_params=sampling_params,
166 + )[0]
167 +
168 + assert single_output.outputs[0].text == batched_output.outputs[0].text
169 + assert list(single_output.outputs[0].token_ids) == list(
170 + batched_output.outputs[0].token_ids
171 + )
tests/models/multimodal/generation/test_musicflamingo.py new
+146
@@ -0,0 +1,146 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import json
5 +import os
6 +
7 +import pytest
8 +
9 +from tests.models.registry import HF_EXAMPLE_MODELS
10 +from vllm import LLM, SamplingParams
11 +
12 +MODEL_NAME = "nvidia/music-flamingo-2601-hf"
13 +SINGLE_CONVERSATION = [
14 + {
15 + "role": "user",
16 + "content": [
17 + {
18 + "type": "text",
19 + "text": "Describe this track in full detail - tell me the "
20 + "genre, tempo, and key, then dive into the instruments, "
21 + "production style, and overall mood it creates.",
22 + },
23 + {
24 + "type": "audio_url",
25 + "audio_url": {
26 + "url": "https://huggingface.co/datasets/nvidia/AudioSkills/"
27 + "resolve/main/assets/song_1.mp3",
28 + },
29 + },
30 + ],
31 + }
32 +]
33 +BATCHED_CONVERSATIONS = [
34 + SINGLE_CONVERSATION,
35 + [
36 + {
37 + "role": "user",
38 + "content": [
39 + {
40 + "type": "text",
41 + "text": "Generate a structured lyric sheet from the input music.",
42 + },
43 + {
44 + "type": "audio_url",
45 + "audio_url": {
46 + "url": "https://huggingface.co/datasets/nvidia/"
47 + "AudioSkills/resolve/main/assets/song_2.mp3",
48 + },
49 + },
50 + ],
51 + }
52 + ],
53 +]
54 +
55 +
56 +def get_fixture_path(filename):
57 + return os.path.join(
58 + os.path.dirname(__file__), "../../fixtures/musicflamingo", filename
59 + )
60 +
61 +
62 +def assert_output_matches(output, expected_text, expected_token_ids):
63 + generated = output.outputs[0]
64 + assert generated.text == expected_text
65 + actual_token_ids = list(generated.token_ids)
66 + assert (
67 + actual_token_ids == expected_token_ids
68 + or actual_token_ids == expected_token_ids[:-1]
69 + or actual_token_ids[:-1] == expected_token_ids
70 + )
71 +
72 +
73 +@pytest.fixture(scope="module")
74 +def llm():
75 + model_info = HF_EXAMPLE_MODELS.get_hf_info("MusicFlamingoForConditionalGeneration")
76 + model_info.check_transformers_version(on_fail="skip")
77 +
78 + try:
79 + return LLM(
80 + model=MODEL_NAME,
81 + dtype="bfloat16",
82 + enforce_eager=True,
83 + max_model_len=8192,
84 + limit_mm_per_prompt={"audio": 1},
85 + )
86 + except Exception as e:
87 + pytest.skip(f"Failed to load model {MODEL_NAME}: {e}")
88 +
89 +
90 +def test_single_generation(llm):
91 + fixture_path = get_fixture_path("expected_results_single.json")
92 + if not os.path.exists(fixture_path):
93 + pytest.skip(f"Fixture not found: {fixture_path}")
94 +
95 + with open(fixture_path) as f:
96 + expected = json.load(f)
97 +
98 + outputs = llm.chat(
99 + messages=SINGLE_CONVERSATION,
100 + sampling_params=SamplingParams(temperature=0.0, max_tokens=50),
101 + )
102 +
103 + assert_output_matches(
104 + outputs[0],
105 + expected["transcriptions"][0],
106 + expected["token_ids"][0],
107 + )
108 +
109 +
110 +def test_batched_generation(llm):
111 + fixture_path = get_fixture_path("expected_results_batched.json")
112 + if not os.path.exists(fixture_path):
113 + pytest.skip(f"Fixture not found: {fixture_path}")
114 +
115 + with open(fixture_path) as f:
116 + expected = json.load(f)
117 +
118 + outputs = llm.chat(
119 + messages=BATCHED_CONVERSATIONS,
120 + sampling_params=SamplingParams(temperature=0.0, max_tokens=50),
121 + )
122 +
123 + for i, output in enumerate(outputs):
124 + assert_output_matches(
125 + output,
126 + expected["transcriptions"][i],
127 + expected["token_ids"][i],
128 + )
129 +
130 +
131 +def test_single_and_batched_generation_match(llm):
132 + sampling_params = SamplingParams(temperature=0.0, max_tokens=50)
133 +
134 + single_output = llm.chat(
135 + messages=SINGLE_CONVERSATION,
136 + sampling_params=sampling_params,
137 + )[0]
138 + batched_output = llm.chat(
139 + messages=BATCHED_CONVERSATIONS,
140 + sampling_params=sampling_params,
141 + )[0]
142 +
143 + assert single_output.outputs[0].text == batched_output.outputs[0].text
144 + assert list(single_output.outputs[0].token_ids) == list(
145 + batched_output.outputs[0].token_ids
146 + )
tests/models/multimodal/processing/test_audioflamingo3.py
+105 -3
@@ -40,6 +40,7 @@ class MockAudioFlamingo3Processor:
40 def __init__(self):
41 self.audio_token = "<sound>"
42 self.audio_token_id = 12345
43 + self.max_audio_len = 60
44 self.feature_extractor = MockFeatureExtractor()
45
46 def __call__(self, text=None, audios=None, **kwargs):
@@ -65,7 +66,6 @@ def mock_ctx():
66
67 @pytest.fixture(autouse=True)
68 def check_transformers_version():
68 - # Check if the model is supported by the current transformers version
69 model_info = HF_EXAMPLE_MODELS.get_hf_info("AudioFlamingo3ForConditionalGeneration")
70 model_info.check_transformers_version(on_fail="skip")
71
@@ -84,7 +84,7 @@ def test_audio_chunk_counting(mock_ctx):
84
85 sr = 16000
86 audio_1 = np.zeros(30 * sr)
87 - audio_2 = np.zeros(45 * sr)
87 + audio_2 = np.zeros(75 * sr)
88
89 mm_data = {"audio": [audio_1, audio_2]}
90 prompt = "<|user|>Listen.<|end|>"
@@ -121,5 +121,107 @@ def test_dummy_data_generation(mock_ctx):
121 assert "audio" in dummy_data
122 assert len(dummy_data["audio"]) == 2
123
124 - expected_len = 600 * 16000
124 + expected_len = 60 * 16000
125 assert len(dummy_data["audio"][0]) == expected_len
126 +
127 +
128 +def test_audio_token_count_matches_hf_processor_math():
129 + from vllm.model_executor.models.audioflamingo3 import (
130 + _count_audio_tokens_from_mask,
131 + )
132 +
133 + feature_attention_mask = torch.zeros((3, 3000), dtype=torch.long)
134 + feature_attention_mask[0, :2999] = 1
135 + feature_attention_mask[1, :2999] = 1
136 + feature_attention_mask[2, :1500] = 1
137 + chunk_counts = torch.tensor([2, 1], dtype=torch.long)
138 +
139 + assert (
140 + _count_audio_tokens_from_mask(feature_attention_mask, chunk_counts, 0) == 1499
141 + )
142 + assert _count_audio_tokens_from_mask(feature_attention_mask, chunk_counts, 1) == 375
143 +
144 +
145 +def test_audio_feature_pipeline_matches_hf_small_config():
146 + from transformers.models.audioflamingo3 import (
147 + modeling_audioflamingo3 as hf_audioflamingo3_modeling,
148 + )
149 + from transformers.models.audioflamingo3.configuration_audioflamingo3 import (
150 + AudioFlamingo3Config,
151 + )
152 +
153 + from vllm.model_executor.models.audioflamingo3 import (
154 + AudioFlamingo3Encoder,
155 + AudioFlamingo3MultiModalProjector,
156 + _build_audio_encoder_attention_mask,
157 + _flatten_valid_audio_embeddings,
158 + )
159 +
160 + text_config = {
161 + "model_type": "qwen2",
162 + "intermediate_size": 64,
163 + "initializer_range": 0.02,
164 + "hidden_size": 32,
165 + "max_position_embeddings": 1024,
166 + "num_hidden_layers": 2,
167 + "num_attention_heads": 4,
168 + "num_key_value_heads": 2,
169 + "vocab_size": 128,
170 + "pad_token_id": 1,
171 + "use_mrope": False,
172 + }
173 + audio_config = {
174 + "hidden_size": 16,
175 + "num_attention_heads": 4,
176 + "intermediate_size": 32,
177 + "num_hidden_layers": 2,
178 + "num_mel_bins": 80,
179 + "max_source_positions": 1500,
180 + "dropout": 0.0,
181 + "attention_dropout": 0.0,
182 + "activation_dropout": 0.0,
183 + "encoder_layerdrop": 0.0,
184 + }
185 +
186 + torch.manual_seed(0)
187 + config = AudioFlamingo3Config(
188 + text_config=text_config,
189 + audio_config=audio_config,
190 + audio_token_id=0,
191 + )
192 + hf_model = hf_audioflamingo3_modeling.AudioFlamingo3ForConditionalGeneration(
193 + config
194 + ).eval()
195 +
196 + vllm_encoder = AudioFlamingo3Encoder(config.audio_config).eval()
197 + vllm_encoder.load_state_dict(hf_model.audio_tower.state_dict())
198 +
199 + vllm_projector = AudioFlamingo3MultiModalProjector(config).eval()
200 + vllm_projector.load_state_dict(hf_model.multi_modal_projector.state_dict())
201 +
202 + input_features = torch.randn(3, 80, 3000)
203 + feature_attention_mask = torch.zeros(3, 3000, dtype=torch.bool)
204 + feature_attention_mask[0, :3000] = True
205 + feature_attention_mask[1, :2500] = True
206 + feature_attention_mask[2, :1500] = True
207 +
208 + hf_output = hf_model.get_audio_features(
209 + input_features,
210 + feature_attention_mask,
211 + return_dict=True,
212 + ).pooler_output
213 + vllm_attention_mask = _build_audio_encoder_attention_mask(
214 + feature_attention_mask,
215 + dtype=vllm_encoder.conv1.weight.dtype,
216 + device=vllm_encoder.conv1.weight.device,
217 + )
218 + vllm_hidden_states = vllm_encoder(
219 + input_features,
220 + attention_mask=vllm_attention_mask,
221 + )
222 + vllm_output, _ = _flatten_valid_audio_embeddings(
223 + vllm_projector(vllm_hidden_states),
224 + feature_attention_mask,
225 + )
226 +
227 + torch.testing.assert_close(vllm_output, hf_output)
tests/models/multimodal/processing/test_musicflamingo.py new
+222
@@ -0,0 +1,222 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +# Copyright 2026 The vLLM team.
5 +# Copyright 2026 NVIDIA CORPORATION and the HuggingFace Inc. team. All rights
6 +# reserved.
7 +#
8 +# Licensed under the Apache License, Version 2.0 (the "License");
9 +# you may not use this file except in compliance with the License.
10 +# You may obtain a copy of the License at
11 +#
12 +# http://www.apache.org/licenses/LICENSE-2.0
13 +#
14 +# Unless required by applicable law or agreed to in writing, software
15 +# distributed under the License is distributed on an "AS IS" BASIS,
16 +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
17 +# See the License for the specific language governing permissions and
18 +# limitations under the License.
19 +
20 +from unittest.mock import MagicMock
21 +
22 +import numpy as np
23 +import pytest
24 +import torch
25 +from transformers import PretrainedConfig
26 +
27 +from tests.models.registry import HF_EXAMPLE_MODELS
28 +
29 +
30 +class MockMusicFlamingoConfig(PretrainedConfig):
31 + model_type = "musicflamingo"
32 +
33 + def __init__(self, **kwargs):
34 + super().__init__(**kwargs)
35 + self.audio_config = PretrainedConfig()
36 + self.text_config = PretrainedConfig()
37 +
38 +
39 +class MockMusicFlamingoProcessor:
40 + def __init__(self):
41 + self.audio_token = "<sound>"
42 + self.audio_token_id = 12345
43 + self.audio_bos_token = "<|sound_bos|>"
44 + self.audio_bos_token_id = 12346
45 + self.audio_eos_token = "<|sound_eos|>"
46 + self.audio_eos_token_id = 12347
47 + self.max_audio_len = 1200
48 + self.feature_extractor = MockFeatureExtractor()
49 +
50 +
51 +class MockFeatureExtractor:
52 + def __init__(self):
53 + self.sampling_rate = 16000
54 + self.chunk_length = 30
55 +
56 +
57 +@pytest.fixture
58 +def mock_ctx():
59 + config = MockMusicFlamingoConfig()
60 +
61 + ctx = MagicMock()
62 + ctx.get_hf_config.return_value = config
63 + ctx.get_hf_processor.return_value = MockMusicFlamingoProcessor()
64 + ctx.model_config.hf_config = config
65 + return ctx
66 +
67 +
68 +@pytest.fixture(autouse=True)
69 +def check_transformers_version():
70 + model_info = HF_EXAMPLE_MODELS.get_hf_info("MusicFlamingoForConditionalGeneration")
71 + model_info.check_transformers_version(on_fail="skip")
72 +
73 +
74 +def test_musicflamingo_chunk_counting_uses_rote_timestamps(mock_ctx, monkeypatch):
75 + from vllm.model_executor.models.musicflamingo import (
76 + MusicFlamingoDummyInputsBuilder,
77 + MusicFlamingoMultiModalProcessor,
78 + MusicFlamingoProcessingInfo,
79 + )
80 +
81 + info = MusicFlamingoProcessingInfo(mock_ctx)
82 + processor = MusicFlamingoMultiModalProcessor(
83 + info, MusicFlamingoDummyInputsBuilder(info)
84 + )
85 +
86 + sr = 16000
87 + audio_1 = np.zeros(30 * sr)
88 + audio_2 = np.zeros(45 * sr)
89 +
90 + mm_data = {"audio": [audio_1, audio_2]}
91 + prompt = "<|user|>Listen.<|end|>"
92 +
93 + from vllm.multimodal.processing import BaseMultiModalProcessor
94 +
95 + def mock_base_call(self, prompt, mm_data, mm_kwargs, tok_kwargs):
96 + del self, prompt, mm_data, mm_kwargs, tok_kwargs
97 + return {
98 + "input_ids": [1, 2, 3],
99 + "input_features": torch.randn(3, 80, 3000),
100 + "rote_timestamps": torch.randn(3, 750),
101 + }
102 +
103 + monkeypatch.setattr(BaseMultiModalProcessor, "_call_hf_processor", mock_base_call)
104 +
105 + processed = processor._call_hf_processor(prompt, mm_data, {}, {})
106 +
107 + chunk_counts = processed["chunk_counts"]
108 +
109 + assert chunk_counts.tolist() == [1, 2]
110 + assert "rote_timestamps" in processed
111 +
112 +
113 +def test_musicflamingo_dummy_text_uses_plain_audio_tokens(mock_ctx):
114 + from vllm.model_executor.models.musicflamingo import (
115 + MusicFlamingoDummyInputsBuilder,
116 + MusicFlamingoProcessingInfo,
117 + )
118 +
119 + info = MusicFlamingoProcessingInfo(mock_ctx)
120 + builder = MusicFlamingoDummyInputsBuilder(info)
121 +
122 + assert builder.get_dummy_text({"audio": 2}) == "<sound><sound>"
123 +
124 +
125 +def test_musicflamingo_audio_feature_pipeline_matches_hf_small_config():
126 + from transformers.models.musicflamingo import (
127 + modeling_musicflamingo as hf_musicflamingo_modeling,
128 + )
129 + from transformers.models.musicflamingo.configuration_musicflamingo import (
130 + MusicFlamingoConfig,
131 + )
132 +
133 + from vllm.model_executor.models.audioflamingo3 import (
134 + _build_audio_encoder_attention_mask,
135 + _flatten_valid_audio_embeddings,
136 + )
137 + from vllm.model_executor.models.musicflamingo import (
138 + MusicFlamingoEncoder,
139 + MusicFlamingoMultiModalProjector,
140 + MusicFlamingoRotaryEmbedding,
141 + apply_rotary_time_emb,
142 + )
143 +
144 + text_config = {
145 + "model_type": "qwen2",
146 + "intermediate_size": 64,
147 + "initializer_range": 0.02,
148 + "hidden_size": 32,
149 + "max_position_embeddings": 1024,
150 + "num_hidden_layers": 2,
151 + "num_attention_heads": 4,
152 + "num_key_value_heads": 2,
153 + "vocab_size": 128,
154 + "pad_token_id": 1,
155 + "use_mrope": False,
156 + }
157 + audio_config = {
158 + "hidden_size": 16,
159 + "num_attention_heads": 4,
160 + "intermediate_size": 32,
161 + "num_hidden_layers": 2,
162 + "num_mel_bins": 80,
163 + "max_source_positions": 1500,
164 + "dropout": 0.0,
165 + "attention_dropout": 0.0,
166 + "activation_dropout": 0.0,
167 + "encoder_layerdrop": 0.0,
168 + }
169 +
170 + torch.manual_seed(0)
171 + config = MusicFlamingoConfig(
172 + text_config=text_config,
173 + audio_config=audio_config,
174 + audio_token_id=0,
175 + head_dim=8,
176 + rope_parameters={"rope_type": "default", "rope_theta": 2048},
177 + )
178 + hf_model = hf_musicflamingo_modeling.MusicFlamingoForConditionalGeneration(
179 + config
180 + ).eval()
181 +
182 + vllm_encoder = MusicFlamingoEncoder(config.audio_config).eval()
183 + vllm_encoder.load_state_dict(hf_model.audio_tower.state_dict())
184 +
185 + vllm_projector = MusicFlamingoMultiModalProjector(config).eval()
186 + vllm_projector.load_state_dict(hf_model.multi_modal_projector.state_dict())
187 +
188 + vllm_rope = MusicFlamingoRotaryEmbedding(config).eval()
189 + vllm_rope.load_state_dict(hf_model.pos_emb.state_dict(), strict=False)
190 +
191 + input_features = torch.randn(3, 80, 3000)
192 + feature_attention_mask = torch.zeros(3, 3000, dtype=torch.bool)
193 + feature_attention_mask[0, :3000] = True
194 + feature_attention_mask[1, :2500] = True
195 + feature_attention_mask[2, :1500] = True
196 + rote_timestamps = (
197 + torch.arange(750, dtype=torch.float32).unsqueeze(0).repeat(3, 1) * 0.04
198 + )
199 +
200 + hf_output = hf_model.get_audio_features(
201 + input_features,
202 + feature_attention_mask,
203 + rote_timestamps=rote_timestamps,
204 + return_dict=True,
205 + ).pooler_output
206 + vllm_attention_mask = _build_audio_encoder_attention_mask(
207 + feature_attention_mask,
208 + dtype=vllm_encoder.conv1.weight.dtype,
209 + device=vllm_encoder.conv1.weight.device,
210 + )
211 + vllm_hidden_states = vllm_encoder(
212 + input_features,
213 + attention_mask=vllm_attention_mask,
214 + )
215 + cos, sin = vllm_rope(rote_timestamps, seq_len=vllm_hidden_states.shape[-2])
216 + vllm_hidden_states = apply_rotary_time_emb(vllm_hidden_states, cos, sin)
217 + vllm_output, _ = _flatten_valid_audio_embeddings(
218 + vllm_projector(vllm_hidden_states),
219 + feature_attention_mask,
220 + )
221 +
222 + torch.testing.assert_close(vllm_output, hf_output)
tests/models/registry.py
+1 -1
@@ -752,7 +752,7 @@ _MULTIMODAL_EXAMPLE_MODELS = {
752 "nvidia/audio-flamingo-3-hf", min_transformers_version="5.0.0"
753 ),
754 "MusicFlamingoForConditionalGeneration": _HfExamplesInfo(
755 - "nvidia/music-flamingo-2601-hf", min_transformers_version="5.0.0.dev"
755 + "nvidia/music-flamingo-2601-hf", min_transformers_version="5.3.0"
756 ),
757 "AyaVisionForConditionalGeneration": _HfExamplesInfo("CohereLabs/aya-vision-8b"),
758 "BagelForConditionalGeneration": _HfExamplesInfo("ByteDance-Seed/BAGEL-7B-MoT"),
vllm/model_executor/models/audioflamingo3.py
+173 -134
@@ -69,10 +69,7 @@ from .utils import (
69 maybe_prefix,
70 )
71
72 -MAX_AUDIO_LEN = 10 * 60
72
74 -
75 -# === Audio Inputs === #
73 class AudioFlamingo3FeatureInputs(TensorSchema):
74 """
75 Dimensions:
@@ -127,14 +124,12 @@ class AudioFlamingo3Encoder(Qwen2AudioEncoder):
124 ):
125 super().__init__(config)
126 self.avg_pooler = nn.AvgPool1d(kernel_size=2, stride=2)
130 - # self.layer_norm is already initialized in super().__init__
127
128 def forward(
129 self,
130 input_features: torch.Tensor | list[torch.Tensor],
131 attention_mask: torch.Tensor = None,
132 ):
137 - # input_features: (batch, num_mel_bins, seq_len)
133 if isinstance(input_features, list):
134 input_features = torch.stack(input_features)
135
@@ -146,17 +141,14 @@ class AudioFlamingo3Encoder(Qwen2AudioEncoder):
141 ).to(hidden_states.dtype)
142
143 for layer in self.layers:
149 - # Qwen2AudioEncoderLayer expects layer_head_mask as third arg.
150 - layer_outputs = layer(hidden_states, attention_mask, None)
151 - hidden_states = layer_outputs[0]
144 + layer_outputs = layer(hidden_states, attention_mask)
145 + hidden_states = (
146 + layer_outputs[0] if isinstance(layer_outputs, tuple) else layer_outputs
147 + )
148
153 - # AvgPool (time/2) + LayerNorm
154 - # hidden_states: (batch, seq_len, hidden_size)
155 - hidden_states = hidden_states.permute(0, 2, 1) # (batch, hidden_size, seq_len)
149 + hidden_states = hidden_states.permute(0, 2, 1)
150 hidden_states = self.avg_pooler(hidden_states)
157 - hidden_states = hidden_states.permute(
158 - 0, 2, 1
159 - ) # (batch, seq_len/2, hidden_size)
151 + hidden_states = hidden_states.permute(0, 2, 1)
152 hidden_states = self.layer_norm(hidden_states)
153
154 return hidden_states
@@ -193,22 +185,6 @@ class AudioFlamingo3MultiModalProjector(nn.Module):
185 return hidden_states
186
187
196 -class AudioFlamingo3MultiModalDataParser(MultiModalDataParser):
197 - def _parse_audio_data(
198 - self,
199 - data: dict[str, torch.Tensor] | ModalityData[Any],
200 - ) -> ModalityDataItems[Any, Any] | None:
201 - if isinstance(data, dict):
202 - return DictEmbeddingItems(
203 - data,
204 - modality="audio",
205 - required_fields={"audio_embeds"},
206 - fields_factory=_audioflamingo3_field_config,
207 - )
208 -
209 - return super()._parse_audio_data(data)
210 -
211 -
188 class AudioFlamingo3ProcessingInfo(BaseProcessingInfo):
189 def get_hf_config(self):
190 return self.ctx.get_hf_config(AudioFlamingo3Config)
@@ -217,20 +193,17 @@ class AudioFlamingo3ProcessingInfo(BaseProcessingInfo):
193 return self.ctx.get_hf_processor(AudioFlamingo3Processor, **kwargs)
194
195 def get_feature_extractor(self, **kwargs: object):
220 - hf_processor = self.get_hf_processor(**kwargs)
221 - feature_extractor = hf_processor.feature_extractor
222 - return feature_extractor
196 + return self.get_hf_processor(**kwargs).feature_extractor
197
224 - def get_data_parser(self):
198 + def get_data_parser(self) -> MultiModalDataParser:
199 feature_extractor = self.get_feature_extractor()
226 -
200 return AudioFlamingo3MultiModalDataParser(
201 target_sr=feature_extractor.sampling_rate,
202 expected_hidden_size=self._get_expected_hidden_size(),
203 )
204
205 def get_supported_mm_limits(self) -> Mapping[str, int | None]:
233 - return {"audio": 1}
206 + return {"audio": None}
207
208
209 class AudioFlamingo3DummyInputsBuilder(
@@ -248,9 +221,10 @@ class AudioFlamingo3DummyInputsBuilder(
221 mm_counts: Mapping[str, int],
222 mm_options: Mapping[str, BaseDummyOptions],
223 ) -> MultiModalDataDict:
224 + hf_processor = self.info.get_hf_processor()
225 feature_extractor = self.info.get_feature_extractor()
226 sampling_rate = feature_extractor.sampling_rate
253 - audio_len = MAX_AUDIO_LEN * sampling_rate
227 + audio_len = int(hf_processor.max_audio_len * sampling_rate)
228 num_audios = mm_counts.get("audio", 0)
229 audio_overrides = mm_options.get("audio")
230
@@ -284,6 +258,118 @@ def _audioflamingo3_field_config(hf_inputs: Mapping[str, torch.Tensor]):
258 )
259
260
261 +def _get_audio_post_pool_output_lengths(input_lengths: torch.Tensor) -> torch.Tensor:
262 + conv_lengths = (input_lengths - 1) // 2 + 1
263 + return (conv_lengths - 2) // 2 + 1
264 +
265 +
266 +def _build_audio_encoder_attention_mask(
267 + feature_attention_mask: torch.Tensor,
268 + *,
269 + dtype: torch.dtype,
270 + device: torch.device,
271 +) -> torch.Tensor:
272 + input_lengths = feature_attention_mask.sum(-1).to(torch.long)
273 + conv_lengths = (input_lengths - 1) // 2 + 1
274 +
275 + batch_size, max_mel_seq_len = feature_attention_mask.shape
276 + max_seq_len = (max_mel_seq_len - 1) // 2 + 1
277 +
278 + seq_range = (
279 + torch.arange(
280 + max_seq_len,
281 + dtype=conv_lengths.dtype,
282 + device=conv_lengths.device,
283 + )
284 + .unsqueeze(0)
285 + .expand(batch_size, max_seq_len)
286 + )
287 + padding_mask = seq_range >= conv_lengths[:, None]
288 +
289 + attention_mask = padding_mask.view(batch_size, 1, 1, max_seq_len).expand(
290 + batch_size, 1, max_seq_len, max_seq_len
291 + )
292 + attention_mask = attention_mask.to(dtype=dtype, device=device)
293 + attention_mask.masked_fill_(padding_mask[:, None, None, :], float("-inf"))
294 +
295 + return attention_mask
296 +
297 +
298 +def _flatten_valid_audio_embeddings(
299 + audio_embeddings: torch.Tensor,
300 + feature_attention_mask: torch.Tensor,
301 +) -> tuple[torch.Tensor, torch.Tensor]:
302 + input_lengths = feature_attention_mask.sum(-1).to(torch.long)
303 + output_lengths = _get_audio_post_pool_output_lengths(input_lengths)
304 + valid_mask = (
305 + torch.arange(audio_embeddings.shape[1], device=output_lengths.device)[None, :]
306 + < output_lengths[:, None]
307 + )
308 +
309 + return audio_embeddings[valid_mask], output_lengths
310 +
311 +
312 +def _count_audio_tokens_from_mask(
313 + feature_attention_mask: torch.Tensor | list[torch.Tensor],
314 + chunk_counts: torch.Tensor | list[torch.Tensor] | list[int] | None,
315 + item_idx: int,
316 +) -> int:
317 + if chunk_counts is not None:
318 + if isinstance(chunk_counts, torch.Tensor):
319 + counts = chunk_counts.tolist()
320 + elif chunk_counts and isinstance(chunk_counts[0], torch.Tensor):
321 + counts = [count.item() for count in chunk_counts]
322 + else:
323 + counts = chunk_counts
324 +
325 + start_idx = sum(counts[:item_idx])
326 + count = counts[item_idx]
327 + end_idx = start_idx + count
328 +
329 + if isinstance(feature_attention_mask, list):
330 + sample_mask = feature_attention_mask[start_idx:end_idx]
331 + if len(sample_mask) == 0:
332 + raise ValueError("Expected non-empty audio mask slice.")
333 + if isinstance(sample_mask[0], torch.Tensor):
334 + sample_mask = torch.stack(sample_mask)
335 + else:
336 + sample_mask = torch.tensor(sample_mask)
337 + else:
338 + sample_mask = feature_attention_mask[start_idx:end_idx]
339 + else:
340 + if isinstance(feature_attention_mask, list):
341 + sample_mask = feature_attention_mask[item_idx]
342 + else:
343 + sample_mask = feature_attention_mask[item_idx]
344 +
345 + if sample_mask.ndim == 1:
346 + sample_input_lengths = sample_mask.sum().unsqueeze(0)
347 + else:
348 + # Match the HF processor, which derives placeholder lengths from the
349 + # total pre-encoder feature length for each original audio sample.
350 + sample_input_lengths = sample_mask.sum().reshape(1)
351 +
352 + post_lengths = _get_audio_post_pool_output_lengths(
353 + sample_input_lengths.to(torch.long)
354 + )
355 + return int(post_lengths[0].item())
356 +
357 +
358 +class AudioFlamingo3MultiModalDataParser(MultiModalDataParser):
359 + def _parse_audio_data(
360 + self,
361 + data: dict[str, torch.Tensor] | ModalityData[Any],
362 + ) -> ModalityDataItems[Any, Any] | None:
363 + if isinstance(data, dict):
364 + return DictEmbeddingItems(
365 + data,
366 + modality="audio",
367 + required_fields={"audio_embeds"},
368 + fields_factory=_audioflamingo3_field_config,
369 + )
370 + return super()._parse_audio_data(data)
371 +
372 +
373 class AudioFlamingo3MultiModalProcessor(
374 BaseMultiModalProcessor[AudioFlamingo3ProcessingInfo]
375 ):
@@ -303,13 +389,13 @@ class AudioFlamingo3MultiModalProcessor(
389 prompt_ids = self._apply_hf_processor_tokens_only(prompt_ids)
390 return BatchFeature(dict(input_ids=[prompt_ids]), tensor_type="pt")
391
306 - feature_extractor = self.info.get_feature_extractor(**mm_kwargs)
392 + processor = self.info.get_hf_processor(**mm_kwargs)
393 + feature_extractor = processor.feature_extractor
394 mm_kwargs = dict(
395 **mm_kwargs,
396 sampling_rate=feature_extractor.sampling_rate,
397 )
398
312 - # Calculate chunk counts
399 audio_list = mm_data.get("audio")
400 if not isinstance(audio_list, list):
401 audio_list = [audio_list]
@@ -318,8 +404,7 @@ class AudioFlamingo3MultiModalProcessor(
404 sampling_rate = feature_extractor.sampling_rate
405 chunk_length = feature_extractor.chunk_length
406 window_size = int(sampling_rate * chunk_length)
321 - # MAX_AUDIO_LEN is 10 * 60 in HF processor.
322 - max_windows = int(MAX_AUDIO_LEN // chunk_length)
407 + max_windows = int(processor.max_audio_len // chunk_length)
408
409 for audio in audio_list:
410 # audio is numpy array or list
@@ -364,7 +449,6 @@ class AudioFlamingo3MultiModalProcessor(
449 audio_token = getattr(processor, "audio_token", "<sound>")
450 audio_token_id = vocab.get(audio_token)
451 if audio_token_id is None:
367 - # Fallback if not found, though it should be there
452 audio_token_id = processor.audio_token_id
453
454 out_mm_data = out_mm_kwargs.get_data()
@@ -373,38 +457,11 @@ class AudioFlamingo3MultiModalProcessor(
457
458 def get_replacement_audioflamingo3(item_idx: int):
459 if feature_attention_mask is not None:
376 - if chunk_counts is not None:
377 - counts = (
378 - chunk_counts.tolist()
379 - if isinstance(chunk_counts, torch.Tensor)
380 - else chunk_counts
381 - )
382 - start_idx = sum(counts[:item_idx])
383 - count = counts[item_idx]
384 - end_idx = start_idx + count
385 -
386 - if isinstance(feature_attention_mask, list):
387 - mask_list = feature_attention_mask[start_idx:end_idx]
388 - if len(mask_list) > 0 and isinstance(
389 - mask_list[0], torch.Tensor
390 - ):
391 - mask = torch.stack(mask_list)
392 - else:
393 - mask = torch.tensor(mask_list)
394 - else:
395 - mask = feature_attention_mask[start_idx:end_idx]
396 - else:
397 - # feature_attention_mask is list[Tensor] or Tensor
398 - if isinstance(feature_attention_mask, list):
399 - mask = feature_attention_mask[item_idx]
400 - else:
401 - mask = feature_attention_mask[item_idx].unsqueeze(0)
402 -
403 - # mask shape: (num_chunks, 3000)
404 - input_lengths = mask.sum(-1)
405 - conv_lengths = (input_lengths - 1) // 2 + 1
406 - audio_output_lengths = (conv_lengths - 2) // 2 + 1
407 - num_features = audio_output_lengths.sum().item()
460 + num_features = _count_audio_tokens_from_mask(
461 + feature_attention_mask,
462 + chunk_counts,
463 + item_idx,
464 + )
465 else:
466 audio_embeds = out_mm_data["audio_embeds"][item_idx]
467 num_features = audio_embeds.shape[0]
@@ -435,13 +492,6 @@ class AudioFlamingo3MultiModalProcessor(
492 class AudioFlamingo3ForConditionalGeneration(
493 nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA
494 ):
438 - """
439 - AudioFlamingo3 model for conditional generation.
440 -
441 - This model integrates a Whisper-based audio encoder with a Qwen2 language model.
442 - It supports multi-chunk audio processing.
443 - """
444 -
495 packed_modules_mapping = {
496 "qkv_proj": ["q_proj", "k_proj", "v_proj"],
497 "gate_up_proj": ["gate_proj", "up_proj"],
@@ -517,6 +567,25 @@ class AudioFlamingo3ForConditionalGeneration(
567 audio_embeds = audio_input["audio_embeds"]
568 return tuple(audio_embeds)
569
570 + (
571 + input_features,
572 + feature_attention_mask,
573 + chunk_counts,
574 + ) = self._normalize_audio_feature_inputs(audio_input)
575 + audio_hidden_states = self._encode_audio_features(
576 + input_features,
577 + feature_attention_mask,
578 + )
579 + audio_features = self.multi_modal_projector(audio_hidden_states)
580 + return self._group_audio_embeddings(
581 + audio_features,
582 + feature_attention_mask,
583 + chunk_counts,
584 + )
585 +
586 + def _normalize_audio_feature_inputs(
587 + self, audio_input: AudioFlamingo3FeatureInputs
588 + ) -> tuple[torch.Tensor, torch.Tensor, list[int]]:
589 input_features = audio_input["input_features"]
590 feature_attention_mask = audio_input["feature_attention_mask"]
591 chunk_counts = audio_input.get("chunk_counts")
@@ -534,66 +603,36 @@ class AudioFlamingo3ForConditionalGeneration(
603 and chunk_counts
604 and isinstance(chunk_counts[0], torch.Tensor)
605 ):
537 - chunk_counts = [c.item() for c in chunk_counts]
538 -
539 - # Calculate output lengths
540 - input_lengths = feature_attention_mask.sum(-1)
541 - # Conv downsampling
542 - conv_lengths = (input_lengths - 1) // 2 + 1
543 - # AvgPool downsampling
544 - audio_output_lengths = (conv_lengths - 2) // 2 + 1
545 -
546 - batch_size, _, max_mel_seq_len = input_features.shape
547 -
548 - # Calculate max_seq_len after convs (before pooling) for attention mask
549 - max_seq_len = (max_mel_seq_len - 1) // 2 + 1
550 -
551 - # Create a sequence tensor of shape (batch_size, max_seq_len)
552 - seq_range = (
553 - torch.arange(
554 - 0,
555 - max_seq_len,
556 - dtype=conv_lengths.dtype,
557 - device=conv_lengths.device,
558 - )
559 - .unsqueeze(0)
560 - .expand(batch_size, max_seq_len)
561 - )
562 - lengths_expand = conv_lengths.unsqueeze(-1).expand(batch_size, max_seq_len)
563 - # Create mask
564 - padding_mask = seq_range >= lengths_expand
606 + chunk_counts = [count.item() for count in chunk_counts]
607
566 - audio_attention_mask_ = padding_mask.view(batch_size, 1, 1, max_seq_len).expand(
567 - batch_size, 1, max_seq_len, max_seq_len
568 - )
569 - audio_attention_mask = audio_attention_mask_.to(
608 + return input_features, feature_attention_mask, chunk_counts
609 +
610 + def _encode_audio_features(
611 + self,
612 + input_features: torch.Tensor,
613 + feature_attention_mask: torch.Tensor,
614 + ) -> torch.Tensor:
615 + audio_attention_mask = _build_audio_encoder_attention_mask(
616 + feature_attention_mask,
617 dtype=self.audio_tower.conv1.weight.dtype,
618 device=self.audio_tower.conv1.weight.device,
619 )
573 - audio_attention_mask[audio_attention_mask_] = float("-inf")
620
575 - # Forward pass
576 - audio_features = self.audio_tower(
577 - input_features, attention_mask=audio_attention_mask
578 - )
621 + return self.audio_tower(input_features, attention_mask=audio_attention_mask)
622
580 - # Project
581 - audio_features = self.multi_modal_projector(audio_features)
582 -
583 - # Masking after pooling
584 - num_audios, max_audio_tokens, embed_dim = audio_features.shape
585 - audio_output_lengths = audio_output_lengths.unsqueeze(1)
586 - audio_features_mask = (
587 - torch.arange(max_audio_tokens)
588 - .expand(num_audios, max_audio_tokens)
589 - .to(audio_output_lengths.device)
590 - < audio_output_lengths
623 + def _group_audio_embeddings(
624 + self,
625 + audio_features: torch.Tensor,
626 + feature_attention_mask: torch.Tensor,
627 + chunk_counts: list[int],
628 + ) -> tuple[torch.Tensor, ...]:
629 + masked_audio_features, audio_output_lengths = _flatten_valid_audio_embeddings(
630 + audio_features,
631 + feature_attention_mask,
632 )
592 - masked_audio_features = audio_features[audio_features_mask].view(-1, embed_dim)
593 -
594 - # Split to tuple of embeddings for individual audio input.
633 chunk_embeddings = torch.split(
596 - masked_audio_features, audio_output_lengths.flatten().tolist()
634 + masked_audio_features,
635 + audio_output_lengths.tolist(),
636 )
637
638 grouped_embeddings = []
@@ -613,7 +652,7 @@ class AudioFlamingo3ForConditionalGeneration(
652
653 def forward(
654 self,
616 - input_ids: torch.Tensor | None,
655 + input_ids: torch.Tensor,
656 positions: torch.Tensor,
657 intermediate_tensors: IntermediateTensors | None = None,
658 inputs_embeds: torch.Tensor | None = None,
vllm/model_executor/models/musicflamingo.py
+406 -43
@@ -1,63 +1,209 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 -"""MusicFlamingo model adapter.
4 +# Copyright 2026 The vLLM team.
5 +# Copyright 2026 NVIDIA CORPORATION and the HuggingFace Inc. team. All rights
6 +# reserved.
7 +#
8 +# Licensed under the Apache License, Version 2.0 (the "License");
9 +# you may not use this file except in compliance with the License.
10 +# You may obtain a copy of the License at
11 +#
12 +# http://www.apache.org/licenses/LICENSE-2.0
13 +#
14 +# Unless required by applicable law or agreed to in writing, software
15 +# distributed under the License is distributed on an "AS IS" BASIS,
16 +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
17 +# See the License for the specific language governing permissions and
18 +# limitations under the License.
19
6 -MusicFlamingo shares the AudioFlamingo3 architecture, so we reuse the same
7 -implementation and multimodal processor, while accepting MusicFlamingo config
8 -and processor classes when available.
9 -"""
20 +from collections.abc import Callable, Mapping, Sequence
21 +from math import pi
22 +from typing import Annotated, Any, Optional, TypeAlias
23
11 -from collections.abc import Mapping
12 -
13 -from transformers.models.audioflamingo3 import (
14 - AudioFlamingo3Config,
15 - AudioFlamingo3Processor,
24 +import torch
25 +from torch import Tensor, broadcast_tensors, nn
26 +from transformers import BatchFeature
27 +from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS
28 +from transformers.models.musicflamingo import (
29 + MusicFlamingoConfig,
30 + MusicFlamingoProcessor,
31 )
32
33 +from vllm.config import VllmConfig
34 +from vllm.config.multimodal import BaseDummyOptions
35 from vllm.multimodal import MULTIMODAL_REGISTRY
19 -from vllm.multimodal.processing import BaseProcessingInfo
36 +from vllm.multimodal.inputs import (
37 + MultiModalDataDict,
38 + MultiModalFieldConfig,
39 + MultiModalKwargsItems,
40 +)
41 +from vllm.multimodal.parse import (
42 + DictEmbeddingItems,
43 + ModalityData,
44 + ModalityDataItems,
45 + MultiModalDataItems,
46 + MultiModalDataParser,
47 +)
48 +from vllm.multimodal.processing import (
49 + PromptReplacement,
50 + PromptUpdate,
51 + PromptUpdateDetails,
52 +)
53 +from vllm.utils.tensor_schema import TensorShape
54
55 from .audioflamingo3 import (
56 AudioFlamingo3DummyInputsBuilder,
57 + AudioFlamingo3EmbeddingInputs,
58 + AudioFlamingo3Encoder,
59 + AudioFlamingo3FeatureInputs,
60 AudioFlamingo3ForConditionalGeneration,
61 AudioFlamingo3MultiModalDataParser,
62 AudioFlamingo3MultiModalProcessor,
63 + AudioFlamingo3MultiModalProjector,
64 + AudioFlamingo3ProcessingInfo,
65 + _audioflamingo3_field_config,
66 + _count_audio_tokens_from_mask,
67 )
68
28 -try:
29 - # Optional dependency: use MusicFlamingo classes when transformers provides them.
30 - from transformers.models.musicflamingo import (
31 - MusicFlamingoConfig,
32 - MusicFlamingoProcessor,
33 - )
34 -except Exception: # pragma: no cover - optional dependency
35 - MusicFlamingoConfig = None
36 - MusicFlamingoProcessor = None
37 -
38 -
39 -class MusicFlamingoProcessingInfo(BaseProcessingInfo):
40 - def get_hf_config(self):
41 - if MusicFlamingoConfig is None:
42 - return self.ctx.get_hf_config(AudioFlamingo3Config)
43 - return self.ctx.get_hf_config((MusicFlamingoConfig, AudioFlamingo3Config))
44 -
45 - def get_hf_processor(self, **kwargs: object):
46 - if MusicFlamingoProcessor is None:
47 - return self.ctx.get_hf_processor(AudioFlamingo3Processor, **kwargs)
48 - # Tuple triggers AutoProcessor path and accepts either processor class.
49 - return self.ctx.get_hf_processor(
50 - (MusicFlamingoProcessor, AudioFlamingo3Processor), **kwargs
69 +
70 +def rotate_half(x):
71 + x = x.reshape(*x.shape[:-1], -1, 2)
72 + x1, x2 = x.unbind(dim=-1)
73 + x = torch.stack((-x2, x1), dim=-1)
74 + return x.flatten(-2)
75 +
76 +
77 +def apply_rotary_time_emb(hidden_states, cos, sin):
78 + original_dtype = hidden_states.dtype
79 + hidden_states = hidden_states.to(torch.float64)
80 + cos = cos.to(hidden_states)
81 + sin = sin.to(hidden_states)
82 + rot_dim = cos.shape[-1]
83 + if rot_dim > hidden_states.shape[-1]:
84 + raise ValueError(
85 + f"feature dimension {hidden_states.shape[-1]} is not of "
86 + f"sufficient size to rotate in all the positions {rot_dim}"
87 )
88
53 - def get_feature_extractor(self, **kwargs: object):
54 - hf_processor = self.get_hf_processor(**kwargs)
55 - return hf_processor.feature_extractor
89 + rotated = hidden_states[..., :rot_dim]
90 + passthrough = hidden_states[..., rot_dim:]
91 + rotated = (rotated * cos) + (rotate_half(rotated) * sin)
92 + return torch.cat((rotated, passthrough), dim=-1).to(original_dtype)
93
57 - def get_data_parser(self):
58 - feature_extractor = self.get_feature_extractor()
94
60 - return AudioFlamingo3MultiModalDataParser(
95 +class MusicFlamingoRotaryEmbedding(nn.Module):
96 + inv_freq: torch.Tensor
97 +
98 + def __init__(self, config: MusicFlamingoConfig, device=None):
99 + super().__init__()
100 + self.max_seq_len_cached = config.max_position_embeddings
101 + self.original_max_seq_len = config.max_position_embeddings
102 +
103 + self.config = config
104 + self.rope_type = self.config.rope_parameters["rope_type"]
105 + rope_init_fn: Callable = self.compute_default_rope_parameters
106 + if self.rope_type != "default":
107 + rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
108 + inv_freq, self.attention_scaling = rope_init_fn(self.config, device)
109 +
110 + self.register_buffer("inv_freq", inv_freq, persistent=False)
111 + self.register_buffer("original_inv_freq", inv_freq.clone(), persistent=False)
112 + position_angles = self._compute_position_angles(self.inv_freq)
113 + self.register_buffer("position_angles", position_angles, persistent=False)
114 +
115 + @staticmethod
116 + def compute_default_rope_parameters(
117 + config: MusicFlamingoConfig | None = None,
118 + device: Optional["torch.device"] = None,
119 + seq_len: int | None = None,
120 + ) -> tuple["torch.Tensor", float]:
121 + del seq_len
122 + base = config.rope_parameters["rope_theta"]
123 + dim = getattr(config, "head_dim", None) or (
124 + config.hidden_size // config.num_attention_heads
125 + )
126 + attention_factor = 1.0
127 +
128 + inv_freq = 1.0 / (
129 + base
130 + ** (
131 + torch.arange(0, dim, 2, dtype=torch.int64).to(
132 + device=device,
133 + dtype=torch.float,
134 + )
135 + / dim
136 + )
137 + )
138 + return inv_freq, attention_factor
139 +
140 + def _compute_position_angles(self, inv_freq):
141 + positions = torch.arange(
142 + int(self.max_seq_len_cached),
143 + device=inv_freq.device,
144 + dtype=inv_freq.dtype,
145 + )
146 + positions = positions / self.max_seq_len_cached * (2 * pi)
147 + position_angles = positions.unsqueeze(-1) * inv_freq
148 + position_angles = torch.repeat_interleave(position_angles, 2, dim=-1)
149 + return position_angles.to(dtype=inv_freq.dtype)
150 +
151 + @torch.no_grad()
152 + def forward(self, timestamps: Tensor, seq_len: int) -> tuple[Tensor, Tensor]:
153 + batch_positions = torch.arange(
154 + timestamps.shape[0],
155 + device=self.inv_freq.device,
156 + dtype=self.inv_freq.dtype,
157 + )
158 + batch_positions = batch_positions / self.max_seq_len_cached
159 + batch_freqs = batch_positions.unsqueeze(-1) * self.inv_freq
160 + batch_freqs = torch.repeat_interleave(batch_freqs, 2, dim=-1)
161 +
162 + batch_freqs = batch_freqs[:, None, :]
163 + time_freqs = self.position_angles[:seq_len][None, :, :]
164 + batch_freqs, time_freqs = broadcast_tensors(batch_freqs, time_freqs)
165 + freqs = torch.cat((batch_freqs, time_freqs), dim=-1)
166 + angle = (-timestamps * 2 * pi).to(freqs)
167 + freqs = freqs * angle.unsqueeze(-1)
168 + return freqs.cos(), freqs.sin()
169 +
170 +
171 +class MusicFlamingoFeatureInputs(AudioFlamingo3FeatureInputs):
172 + rote_timestamps: Annotated[
173 + torch.Tensor,
174 + TensorShape(
175 + "num_chunks",
176 + "num_audio_time_steps",
177 + dynamic_dims={"num_audio_time_steps"},
178 + ),
179 + ]
180 +
181 +
182 +MusicFlamingoEmbeddingInputs = AudioFlamingo3EmbeddingInputs
183 +
184 +MusicFlamingoInputs: TypeAlias = (
185 + MusicFlamingoFeatureInputs | MusicFlamingoEmbeddingInputs
186 +)
187 +
188 +
189 +class MusicFlamingoEncoder(AudioFlamingo3Encoder):
190 + pass
191 +
192 +
193 +class MusicFlamingoMultiModalProjector(AudioFlamingo3MultiModalProjector):
194 + pass
195 +
196 +
197 +class MusicFlamingoProcessingInfo(AudioFlamingo3ProcessingInfo):
198 + def get_hf_config(self) -> MusicFlamingoConfig:
199 + return self.ctx.get_hf_config(MusicFlamingoConfig)
200 +
201 + def get_hf_processor(self, **kwargs: object) -> MusicFlamingoProcessor:
202 + return self.ctx.get_hf_processor(MusicFlamingoProcessor, **kwargs)
203 +
204 + def get_data_parser(self) -> MultiModalDataParser:
205 + feature_extractor = self.get_feature_extractor()
206 + return MusicFlamingoMultiModalDataParser(
207 target_sr=feature_extractor.sampling_rate,
208 expected_hidden_size=self._get_expected_hidden_size(),
209 )
@@ -67,13 +213,230 @@ class MusicFlamingoProcessingInfo(BaseProcessingInfo):
213
214
215 class MusicFlamingoDummyInputsBuilder(AudioFlamingo3DummyInputsBuilder):
70 - pass
216 + def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
217 + num_audios = mm_counts.get("audio", 0)
218 + hf_processor = self.info.get_hf_processor()
219 + return hf_processor.audio_token * num_audios
220 +
221 + def get_dummy_mm_data(
222 + self,
223 + seq_len: int,
224 + mm_counts: Mapping[str, int],
225 + mm_options: Mapping[str, BaseDummyOptions],
226 + ) -> MultiModalDataDict:
227 + hf_processor = self.info.get_hf_processor()
228 + feature_extractor = self.info.get_feature_extractor()
229 + sampling_rate = feature_extractor.sampling_rate
230 + audio_len = int(hf_processor.max_audio_len * sampling_rate)
231 + num_audios = mm_counts.get("audio", 0)
232 + audio_overrides = mm_options.get("audio")
233 +
234 + return {
235 + "audio": self._get_dummy_audios(
236 + length=audio_len,
237 + num_audios=num_audios,
238 + overrides=audio_overrides,
239 + )
240 + }
241 +
242 +
243 +def _musicflamingo_field_config(hf_inputs: Mapping[str, torch.Tensor]):
244 + fields = dict(_audioflamingo3_field_config(hf_inputs))
245 + chunk_counts = hf_inputs.get("chunk_counts")
246 + if chunk_counts is not None:
247 + fields["rote_timestamps"] = MultiModalFieldConfig.flat_from_sizes(
248 + "audio", chunk_counts, dim=0
249 + )
250 + else:
251 + fields["rote_timestamps"] = MultiModalFieldConfig.batched("audio")
252 + return fields
253 +
254 +
255 +class MusicFlamingoMultiModalDataParser(AudioFlamingo3MultiModalDataParser):
256 + def _parse_audio_data(
257 + self,
258 + data: dict[str, torch.Tensor] | ModalityData[Any],
259 + ) -> ModalityDataItems[Any, Any] | None:
260 + if isinstance(data, dict):
261 + return DictEmbeddingItems(
262 + data,
263 + modality="audio",
264 + required_fields={"audio_embeds"},
265 + fields_factory=_musicflamingo_field_config,
266 + )
267 + return super()._parse_audio_data(data)
268 +
269 +
270 +class MusicFlamingoMultiModalProcessor(AudioFlamingo3MultiModalProcessor):
271 + def _call_hf_processor(
272 + self,
273 + prompt: str,
274 + mm_data: dict[str, object],
275 + mm_kwargs: Mapping[str, Any],
276 + tok_kwargs: Mapping[str, object],
277 + ) -> BatchFeature:
278 + outputs = super()._call_hf_processor(
279 + prompt=prompt,
280 + mm_data=mm_data,
281 + mm_kwargs=mm_kwargs,
282 + tok_kwargs=tok_kwargs,
283 + )
284 +
285 + audio_data = mm_data.get("audio")
286 + if audio_data is None:
287 + return outputs
288 +
289 + audio_list = audio_data if isinstance(audio_data, list) else [audio_data]
290 + if len(audio_list) == 0:
291 + return outputs
292 +
293 + processor = self.info.get_hf_processor(**mm_kwargs)
294 + feature_extractor = processor.feature_extractor
295 + sampling_rate = feature_extractor.sampling_rate
296 + chunk_length = feature_extractor.chunk_length
297 + window_size = int(sampling_rate * chunk_length)
298 + max_windows = int(processor.max_audio_len // chunk_length)
299 +
300 + chunk_counts = []
301 + for audio in audio_list:
302 + n_samples = len(audio) if isinstance(audio, list) else audio.shape[0]
303 + n_win = max(1, (n_samples + window_size - 1) // window_size)
304 + chunk_counts.append(min(n_win, max_windows))
305 + outputs["chunk_counts"] = torch.tensor(chunk_counts, dtype=torch.long)
306 +
307 + if "rote_timestamps" not in outputs:
308 + raise KeyError(
309 + "MusicFlamingoProcessor output must include `rote_timestamps`."
310 + )
311 +
312 + return outputs
313 +
314 + def _get_mm_fields_config(
315 + self,
316 + hf_inputs: BatchFeature,
317 + hf_processor_mm_kwargs: Mapping[str, object],
318 + ) -> Mapping[str, MultiModalFieldConfig]:
319 + return _musicflamingo_field_config(hf_inputs)
320 +
321 + def _get_prompt_updates(
322 + self,
323 + mm_items: MultiModalDataItems,
324 + hf_processor_mm_kwargs: Mapping[str, object],
325 + out_mm_kwargs: MultiModalKwargsItems,
326 + ) -> Sequence[PromptUpdate]:
327 + processor = self.info.get_hf_processor(**hf_processor_mm_kwargs)
328 + tokenizer = self.info.get_tokenizer()
329 + vocab = tokenizer.get_vocab()
330 +
331 + audio_token = processor.audio_token
332 + audio_token_id = vocab.get(audio_token, processor.audio_token_id)
333 +
334 + audio_bos_token = processor.audio_bos_token
335 + audio_bos_token_id = vocab.get(audio_bos_token, processor.audio_bos_token_id)
336 +
337 + audio_eos_token = processor.audio_eos_token
338 + audio_eos_token_id = vocab.get(audio_eos_token, processor.audio_eos_token_id)
339 +
340 + out_mm_data = out_mm_kwargs.get_data()
341 + feature_attention_mask = out_mm_data.get("feature_attention_mask")
342 + chunk_counts = out_mm_data.get("chunk_counts")
343 +
344 + def get_replacement_musicflamingo(item_idx: int):
345 + if feature_attention_mask is not None:
346 + num_features = _count_audio_tokens_from_mask(
347 + feature_attention_mask,
348 + chunk_counts,
349 + item_idx,
350 + )
351 + else:
352 + audio_embeds = out_mm_data["audio_embeds"][item_idx]
353 + num_features = audio_embeds.shape[0]
354 +
355 + if num_features == 0:
356 + raise ValueError("Audio is too short")
357 +
358 + full_tokens = [
359 + audio_bos_token_id,
360 + *([audio_token_id] * int(num_features)),
361 + audio_eos_token_id,
362 + ]
363 +
364 + return PromptUpdateDetails.select_token_id(
365 + full_tokens,
366 + embed_token_id=audio_token_id,
367 + )
368 +
369 + return [
370 + PromptReplacement(
371 + modality="audio",
372 + target=audio_token,
373 + replacement=get_replacement_musicflamingo,
374 + )
375 + ]
376
377
378 @MULTIMODAL_REGISTRY.register_processor(
74 - AudioFlamingo3MultiModalProcessor,
379 + MusicFlamingoMultiModalProcessor,
380 info=MusicFlamingoProcessingInfo,
381 dummy_inputs=MusicFlamingoDummyInputsBuilder,
382 )
383 class MusicFlamingoForConditionalGeneration(AudioFlamingo3ForConditionalGeneration):
79 - """MusicFlamingo model for conditional generation."""
384 + """vLLM MusicFlamingo model aligned with HF modular_musicflamingo."""
385 +
386 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
387 + super().__init__(vllm_config=vllm_config, prefix=prefix)
388 + self.audio_tower = MusicFlamingoEncoder(self.config.audio_config)
389 + self.multi_modal_projector = MusicFlamingoMultiModalProjector(self.config)
390 + self.pos_emb = MusicFlamingoRotaryEmbedding(self.config)
391 +
392 + def _parse_and_validate_audio_input(
393 + self, **kwargs: object
394 + ) -> MusicFlamingoInputs | None:
395 + rote_timestamps = kwargs.pop("rote_timestamps", None)
396 + audio_input = super()._parse_and_validate_audio_input(**kwargs)
397 + if audio_input is None or audio_input["type"] == "audio_embeds":
398 + return audio_input
399 +
400 + return MusicFlamingoFeatureInputs(
401 + type="audio_features",
402 + input_features=audio_input["input_features"],
403 + feature_attention_mask=audio_input["feature_attention_mask"],
404 + chunk_counts=audio_input["chunk_counts"],
405 + rote_timestamps=rote_timestamps,
406 + )
407 +
408 + def _process_audio_input(
409 + self, audio_input: MusicFlamingoInputs
410 + ) -> torch.Tensor | tuple[torch.Tensor, ...]:
411 + if audio_input["type"] == "audio_embeds":
412 + return super()._process_audio_input(audio_input)
413 +
414 + rote_timestamps = audio_input["rote_timestamps"]
415 + if rote_timestamps is None:
416 + raise ValueError(
417 + "MusicFlamingo audio feature inputs must include `rote_timestamps`."
418 + )
419 + if isinstance(rote_timestamps, list):
420 + rote_timestamps = torch.cat(rote_timestamps, dim=0)
421 +
422 + (
423 + input_features,
424 + feature_attention_mask,
425 + chunk_counts,
426 + ) = self._normalize_audio_feature_inputs(audio_input)
427 + hidden_states = self._encode_audio_features(
428 + input_features,
429 + feature_attention_mask,
430 + )
431 + cos, sin = self.pos_emb(
432 + rote_timestamps.to(hidden_states.device),
433 + seq_len=hidden_states.shape[-2],
434 + )
435 + hidden_states = apply_rotary_time_emb(hidden_states, cos, sin)
436 + audio_features = self.multi_modal_projector(hidden_states)
437 +
438 + return self._group_audio_embeddings(
439 + audio_features,
440 + feature_attention_mask,
441 + chunk_counts,
442 + )