@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 3a20450d31
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 11, 12:30 UTC

Add AudioFlamingo3 model support (#30539)

Signed-off-by: Lasha <26011196+lashahub@users.noreply.github.com> Signed-off-by: Lasha Koroshinadze <26011196+lashahub@users.noreply.github.com> Co-authored-by: Isotr0py <2037008807@qq.com> Co-authored-by: Jee Jee Li <pandaleefree@gmail.com> Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>

Lasha Koroshinadze committed Dec 14, 2025 at 05:14 UTC 3a20450d313e7bffc78f1a0d3628a0866b486883
9 files changed +989 -44
docs/models/supported_models.md
+1
@@ -659,6 +659,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
659 | Architecture | Models | Inputs | Example HF Models | [LoRA](../features/lora.md) | [PP](../serving/parallelism_scaling.md) |
660 |--------------|--------|--------|-------------------|----------------------|---------------------------|
661 | `AriaForConditionalGeneration` | Aria | T + I<sup>+</sup> | `rhymes-ai/Aria` | | |
662 +| `AudioFlamingo3ForConditionalGeneration` | AudioFlamingo3 | T + A<sup>+</sup> | `nvidia/audio-flamingo-3-hf`, `nvidia/music-flamingo-hf` | ✅︎ | ✅︎ |
663 | `AyaVisionForConditionalGeneration` | Aya Vision | T + I<sup>+</sup> | `CohereLabs/aya-vision-8b`, `CohereLabs/aya-vision-32b`, etc. | | ✅︎ |
664 | `BeeForConditionalGeneration` | Bee-8B | T + I<sup>E+</sup> | `Open-Bee/Bee-8B-RL`, `Open-Bee/Bee-8B-SFT` | | ✅︎ |
665 | `Blip2ForConditionalGeneration` | BLIP-2 | T + I<sup>E</sup> | `Salesforce/blip2-opt-2.7b`, `Salesforce/blip2-opt-6.7b`, etc. | | ✅︎ |
examples/offline_inference/audio_language.py
+73 -44
@@ -42,60 +42,31 @@ class ModelRequestData(NamedTuple):
42 # Unless specified, these settings have been tested to work on a single L4.
43
44
45 -# Voxtral
46 -# Make sure to install mistral-common[audio].
47 -def run_voxtral(question: str, audio_count: int) -> ModelRequestData:
48 - from mistral_common.audio import Audio
49 - from mistral_common.protocol.instruct.chunk import (
50 - AudioChunk,
51 - RawAudio,
52 - TextChunk,
53 - )
54 - from mistral_common.protocol.instruct.messages import (
55 - UserMessage,
56 - )
57 - from mistral_common.protocol.instruct.request import ChatCompletionRequest
58 - from mistral_common.tokens.tokenizers.mistral import MistralTokenizer
59 -
60 - model_name = "mistralai/Voxtral-Mini-3B-2507"
61 - tokenizer = MistralTokenizer.from_hf_hub(model_name)
62 -
45 +# AudioFlamingo3
46 +def run_audioflamingo3(question: str, audio_count: int) -> ModelRequestData:
47 + model_name = "nvidia/audio-flamingo-3-hf"
48 engine_args = EngineArgs(
49 model=model_name,
65 - max_model_len=8192,
50 + max_model_len=4096,
51 max_num_seqs=2,
52 limit_mm_per_prompt={"audio": audio_count},
68 - config_format="mistral",
69 - load_format="mistral",
70 - tokenizer_mode="mistral",
53 enforce_eager=True,
72 - enable_chunked_prefill=False,
54 )
55
75 - text_chunk = TextChunk(text=question)
76 - audios = [
77 - Audio.from_file(str(audio_assets[i].get_local_path()), strict=False)
78 - for i in range(audio_count)
79 - ]
80 - audio_chunks = [
81 - AudioChunk(input_audio=RawAudio.from_audio(audio)) for audio in audios
82 - ]
83 -
84 - messages = [UserMessage(content=[*audio_chunks, text_chunk])]
85 -
86 - req = ChatCompletionRequest(messages=messages, model=model_name)
87 -
88 - tokens = tokenizer.encode_chat_completion(req)
89 - prompt_ids, audios = tokens.tokens, tokens.audios
90 -
91 - audios_and_sr = [(au.audio_array, au.sampling_rate) for au in audios]
56 + # AudioFlamingo3 uses <sound> token for audio
57 + audio_placeholder = "<sound>" * audio_count
58
93 - multi_modal_data = {"audio": audios_and_sr}
59 + prompt = (
60 + "<|im_start|>system\n"
61 + "You are a helpful assistant.<|im_end|>\n"
62 + "<|im_start|>user\n"
63 + f"{audio_placeholder}{question}<|im_end|>\n"
64 + "<|im_start|>assistant\n"
65 + )
66
67 return ModelRequestData(
68 engine_args=engine_args,
97 - prompt_token_ids=prompt_ids,
98 - multi_modal_data=multi_modal_data,
69 + prompt=prompt,
70 )
71
72
@@ -361,6 +332,63 @@ def run_ultravox(question: str, audio_count: int) -> ModelRequestData:
332 )
333
334
335 +# Voxtral
336 +# Make sure to install mistral-common[audio].
337 +def run_voxtral(question: str, audio_count: int) -> ModelRequestData:
338 + from mistral_common.audio import Audio
339 + from mistral_common.protocol.instruct.chunk import (
340 + AudioChunk,
341 + RawAudio,
342 + TextChunk,
343 + )
344 + from mistral_common.protocol.instruct.messages import (
345 + UserMessage,
346 + )
347 + from mistral_common.protocol.instruct.request import ChatCompletionRequest
348 + from mistral_common.tokens.tokenizers.mistral import MistralTokenizer
349 +
350 + model_name = "mistralai/Voxtral-Mini-3B-2507"
351 + tokenizer = MistralTokenizer.from_hf_hub(model_name)
352 +
353 + engine_args = EngineArgs(
354 + model=model_name,
355 + max_model_len=8192,
356 + max_num_seqs=2,
357 + limit_mm_per_prompt={"audio": audio_count},
358 + config_format="mistral",
359 + load_format="mistral",
360 + tokenizer_mode="mistral",
361 + enforce_eager=True,
362 + enable_chunked_prefill=False,
363 + )
364 +
365 + text_chunk = TextChunk(text=question)
366 + audios = [
367 + Audio.from_file(str(audio_assets[i].get_local_path()), strict=False)
368 + for i in range(audio_count)
369 + ]
370 + audio_chunks = [
371 + AudioChunk(input_audio=RawAudio.from_audio(audio)) for audio in audios
372 + ]
373 +
374 + messages = [UserMessage(content=[*audio_chunks, text_chunk])]
375 +
376 + req = ChatCompletionRequest(messages=messages, model=model_name)
377 +
378 + tokens = tokenizer.encode_chat_completion(req)
379 + prompt_ids, audios = tokens.tokens, tokens.audios
380 +
381 + audios_and_sr = [(au.audio_array, au.sampling_rate) for au in audios]
382 +
383 + multi_modal_data = {"audio": audios_and_sr}
384 +
385 + return ModelRequestData(
386 + engine_args=engine_args,
387 + prompt_token_ids=prompt_ids,
388 + multi_modal_data=multi_modal_data,
389 + )
390 +
391 +
392 # Whisper
393 def run_whisper(question: str, audio_count: int) -> ModelRequestData:
394 assert audio_count == 1, "Whisper only support single audio input per prompt"
@@ -382,7 +410,7 @@ def run_whisper(question: str, audio_count: int) -> ModelRequestData:
410
411
412 model_example_map = {
385 - "voxtral": run_voxtral,
413 + "audioflamingo3": run_audioflamingo3,
414 "gemma3n": run_gemma3n,
415 "granite_speech": run_granite_speech,
416 "midashenglm": run_midashenglm,
@@ -392,6 +420,7 @@ model_example_map = {
420 "qwen2_audio": run_qwen2_audio,
421 "qwen2_5_omni": run_qwen2_5_omni,
422 "ultravox": run_ultravox,
423 + "voxtral": run_voxtral,
424 "whisper": run_whisper,
425 }
426
tests/models/fixtures/audioflamingo3/expected_results_batched.json new
+1
@@ -0,0 +1 @@
1 +{"transcriptions": ["There is no clear relationship between the barking and the music, as they seem to be independent of each other.", "(B) To indicate that language cannot express clearly, satirizing the inversion of black and white in the world"], "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], [5349, 8, 2014, 13216, 429, 4128, 4157, 3158, 9355, 11, 7578, 404, 4849, 279, 46488, 315, 3691, 323, 4158, 304, 279, 1879, 151645, 151671]]}
\ No newline at end of file
tests/models/fixtures/audioflamingo3/expected_results_single.json new
+1
@@ -0,0 +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
tests/models/multimodal/generation/test_audioflamingo3.py new
+142
@@ -0,0 +1,142 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +# Copyright 2025 The vLLM team.
5 +# Copyright 2025 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 +import json
21 +import os
22 +
23 +import pytest
24 +
25 +from tests.models.registry import HF_EXAMPLE_MODELS
26 +from vllm import LLM, SamplingParams
27 +
28 +MODEL_NAME = "nvidia/audio-flamingo-3-hf"
29 +
30 +
31 +def get_fixture_path(filename):
32 + return os.path.join(
33 + os.path.dirname(__file__), "../../fixtures/audioflamingo3", filename
34 + )
35 +
36 +
37 +@pytest.fixture(scope="module")
38 +def llm():
39 + # Check if the model is supported by the current transformers version
40 + model_info = HF_EXAMPLE_MODELS.get_hf_info("AudioFlamingo3ForConditionalGeneration")
41 + model_info.check_transformers_version(on_fail="skip")
42 +
43 + try:
44 + llm = LLM(
45 + model=MODEL_NAME,
46 + trust_remote_code=True,
47 + dtype="bfloat16",
48 + enforce_eager=True,
49 + limit_mm_per_prompt={"audio": 1},
50 + )
51 + return llm
52 + except Exception as e:
53 + pytest.skip(f"Failed to load model {MODEL_NAME}: {e}")
54 +
55 +
56 +def test_single_generation(llm):
57 + fixture_path = get_fixture_path("expected_results_single.json")
58 + if not os.path.exists(fixture_path):
59 + pytest.skip(f"Fixture not found: {fixture_path}")
60 +
61 + with open(fixture_path) as f:
62 + expected = json.load(f)
63 +
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 +
76 + sampling_params = SamplingParams(temperature=0.0, max_tokens=128)
77 +
78 + outputs = llm.chat(
79 + messages=messages,
80 + sampling_params=sampling_params,
81 + )
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
87 +
88 +
89 +def test_batched_generation(llm):
90 + fixture_path = get_fixture_path("expected_results_batched.json")
91 + if not os.path.exists(fixture_path):
92 + pytest.skip(f"Fixture not found: {fixture_path}")
93 +
94 + with open(fixture_path) as f:
95 + expected = json.load(f)
96 +
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 +
131 + sampling_params = SamplingParams(temperature=0.0, max_tokens=128)
132 +
133 + outputs = llm.chat(
134 + messages=conversations,
135 + sampling_params=sampling_params,
136 + )
137 +
138 + for i, output in enumerate(outputs):
139 + generated_text = output.outputs[0].text.strip()
140 + expected_text = expected["transcriptions"][i]
141 +
142 + assert expected_text in generated_text or generated_text in expected_text
tests/models/multimodal/processing/test_audioflamingo3.py new
+125
@@ -0,0 +1,125 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +# Copyright 2025 The vLLM team.
5 +# Copyright 2025 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 MockAudioFlamingo3Config(PretrainedConfig):
31 + model_type = "audioflamingo3"
32 +
33 + def __init__(self, **kwargs):
34 + super().__init__(**kwargs)
35 + self.audio_config = PretrainedConfig()
36 + self.text_config = PretrainedConfig()
37 +
38 +
39 +class MockAudioFlamingo3Processor:
40 + def __init__(self):
41 + self.audio_token = "<sound>"
42 + self.audio_token_id = 12345
43 + self.feature_extractor = MockFeatureExtractor()
44 +
45 + def __call__(self, text=None, audios=None, **kwargs):
46 + return {"input_ids": [1, 2, 3], "input_features": [np.zeros((3000, 80))]}
47 +
48 +
49 +class MockFeatureExtractor:
50 + def __init__(self):
51 + self.sampling_rate = 16000
52 + self.chunk_length = 30
53 +
54 +
55 +@pytest.fixture
56 +def mock_ctx():
57 + config = MockAudioFlamingo3Config()
58 +
59 + ctx = MagicMock()
60 + ctx.get_hf_config.return_value = config
61 + ctx.get_hf_processor.return_value = MockAudioFlamingo3Processor()
62 + ctx.model_config.hf_config = config
63 + return ctx
64 +
65 +
66 +@pytest.fixture(autouse=True)
67 +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 +
72 +
73 +def test_audio_chunk_counting(mock_ctx):
74 + from vllm.model_executor.models.audioflamingo3 import (
75 + AudioFlamingo3DummyInputsBuilder,
76 + AudioFlamingo3MultiModalProcessor,
77 + AudioFlamingo3ProcessingInfo,
78 + )
79 +
80 + info = AudioFlamingo3ProcessingInfo(mock_ctx)
81 + processor = AudioFlamingo3MultiModalProcessor(
82 + info, AudioFlamingo3DummyInputsBuilder(info)
83 + )
84 +
85 + sr = 16000
86 + audio_1 = np.zeros(30 * sr)
87 + audio_2 = np.zeros(45 * sr)
88 +
89 + mm_data = {"audio": [audio_1, audio_2]}
90 + prompt = "<|user|>Listen.<|end|>"
91 +
92 + from vllm.multimodal.processing import BaseMultiModalProcessor
93 +
94 + def mock_base_call(self, prompt, mm_data, mm_kwargs, tok_kwargs):
95 + return {"input_ids": [1, 2, 3], "input_features": torch.randn(1, 80, 3000)}
96 +
97 + with pytest.MonkeyPatch.context() as mp:
98 + mp.setattr(BaseMultiModalProcessor, "_call_hf_processor", mock_base_call)
99 +
100 + processed = processor._call_hf_processor(prompt, mm_data, {}, {})
101 +
102 + chunk_counts = processed["chunk_counts"]
103 +
104 + assert chunk_counts[0].item() == 1
105 + assert chunk_counts[1].item() == 2
106 + assert len(chunk_counts) == 2
107 +
108 +
109 +def test_dummy_data_generation(mock_ctx):
110 + from vllm.model_executor.models.audioflamingo3 import (
111 + AudioFlamingo3DummyInputsBuilder,
112 + AudioFlamingo3ProcessingInfo,
113 + )
114 +
115 + info = AudioFlamingo3ProcessingInfo(mock_ctx)
116 + builder = AudioFlamingo3DummyInputsBuilder(info)
117 +
118 + mm_counts = {"audio": 2}
119 + dummy_data = builder.get_dummy_mm_data(100, mm_counts, None)
120 +
121 + assert "audio" in dummy_data
122 + assert len(dummy_data["audio"]) == 2
123 +
124 + expected_len = 600 * 16000
125 + assert len(dummy_data["audio"][0]) == expected_len
tests/models/registry.py
+3
@@ -578,6 +578,9 @@ _AUTOMATIC_CONVERTED_MODELS = {
578 _MULTIMODAL_EXAMPLE_MODELS = {
579 # [Decoder-only]
580 "AriaForConditionalGeneration": _HfExamplesInfo("rhymes-ai/Aria"),
581 + "AudioFlamingo3ForConditionalGeneration": _HfExamplesInfo(
582 + "nvidia/audio-flamingo-3-hf", min_transformers_version="5.0.0.dev"
583 + ),
584 "AyaVisionForConditionalGeneration": _HfExamplesInfo("CohereLabs/aya-vision-8b"),
585 "BeeForConditionalGeneration": _HfExamplesInfo(
586 "Open-Bee/Bee-8B-RL",
vllm/model_executor/models/audioflamingo3.py new
+639
@@ -0,0 +1,639 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +# Copyright 2025 The vLLM team.
5 +# Copyright 2025 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 collections.abc import Iterable, Mapping, Sequence
21 +from typing import Annotated, Any, Literal, TypeAlias
22 +
23 +import torch
24 +import torch.nn as nn
25 +from transformers import BatchFeature, PretrainedConfig
26 +from transformers.models.audioflamingo3 import (
27 + AudioFlamingo3Config,
28 + AudioFlamingo3Processor,
29 +)
30 +from transformers.models.qwen2_audio import Qwen2AudioEncoder
31 +
32 +from vllm.config import VllmConfig
33 +from vllm.config.multimodal import BaseDummyOptions
34 +from vllm.model_executor.layers.activation import get_act_fn
35 +from vllm.model_executor.models.module_mapping import MultiModelKeys
36 +from vllm.multimodal import MULTIMODAL_REGISTRY
37 +from vllm.multimodal.inputs import (
38 + MultiModalDataDict,
39 + MultiModalFieldConfig,
40 + MultiModalKwargsItems,
41 +)
42 +from vllm.multimodal.parse import (
43 + DictEmbeddingItems,
44 + ModalityData,
45 + ModalityDataItems,
46 + MultiModalDataItems,
47 + MultiModalDataParser,
48 +)
49 +from vllm.multimodal.processing import (
50 + BaseMultiModalProcessor,
51 + BaseProcessingInfo,
52 + PromptReplacement,
53 + PromptUpdate,
54 + PromptUpdateDetails,
55 +)
56 +from vllm.multimodal.profiling import BaseDummyInputsBuilder
57 +from vllm.sequence import IntermediateTensors
58 +from vllm.utils.tensor_schema import TensorSchema, TensorShape
59 +
60 +from .interfaces import (
61 + MultiModalEmbeddings,
62 + SupportsLoRA,
63 + SupportsMultiModal,
64 + SupportsPP,
65 +)
66 +from .utils import (
67 + AutoWeightsLoader,
68 + init_vllm_registered_model,
69 + maybe_prefix,
70 +)
71 +
72 +MAX_AUDIO_LEN = 10 * 60
73 +
74 +
75 +# === Audio Inputs === #
76 +class AudioFlamingo3FeatureInputs(TensorSchema):
77 + """
78 + Dimensions:
79 + - num_chunks: Number of audio chunks (flattened)
80 + - nmb: Number of mel bins
81 + - num_audios: Number of original audio files
82 + """
83 +
84 + type: Literal["audio_features"]
85 + input_features: Annotated[
86 + torch.Tensor | list[torch.Tensor],
87 + TensorShape("num_chunks", "nmb", 3000),
88 + ]
89 +
90 + feature_attention_mask: Annotated[
91 + torch.Tensor,
92 + TensorShape("num_chunks", 3000),
93 + ]
94 +
95 + chunk_counts: Annotated[
96 + torch.Tensor,
97 + TensorShape("num_audios"),
98 + ]
99 +
100 +
101 +class AudioFlamingo3EmbeddingInputs(TensorSchema):
102 + """
103 + Dimensions:
104 + - bn: Batch size
105 + - naf: Number of audio features
106 + - hs: Hidden size (must match the hidden size of language model
107 + backbone)
108 + """
109 +
110 + type: Literal["audio_embeds"] = "audio_embeds"
111 +
112 + audio_embeds: Annotated[
113 + list[torch.Tensor],
114 + TensorShape("bn", "naf", "hs"),
115 + ]
116 +
117 +
118 +AudioFlamingo3Inputs: TypeAlias = (
119 + AudioFlamingo3FeatureInputs | AudioFlamingo3EmbeddingInputs
120 +)
121 +
122 +
123 +class AudioFlamingo3Encoder(Qwen2AudioEncoder):
124 + def __init__(
125 + self,
126 + config: PretrainedConfig,
127 + ):
128 + super().__init__(config)
129 + self.avg_pooler = nn.AvgPool1d(kernel_size=2, stride=2)
130 + # self.layer_norm is already initialized in super().__init__
131 +
132 + def forward(
133 + self,
134 + input_features: torch.Tensor | list[torch.Tensor],
135 + attention_mask: torch.Tensor = None,
136 + ):
137 + # input_features: (batch, num_mel_bins, seq_len)
138 + if isinstance(input_features, list):
139 + input_features = torch.stack(input_features)
140 +
141 + hidden_states = nn.functional.gelu(self.conv1(input_features))
142 + hidden_states = nn.functional.gelu(self.conv2(hidden_states))
143 + hidden_states = hidden_states.transpose(-1, -2)
144 + hidden_states = (
145 + hidden_states + self.embed_positions.weight[: hidden_states.size(-2), :]
146 + ).to(hidden_states.dtype)
147 +
148 + for layer in self.layers:
149 + layer_outputs = layer(hidden_states, attention_mask)
150 + hidden_states = layer_outputs[0]
151 +
152 + # AvgPool (time/2) + LayerNorm
153 + # hidden_states: (batch, seq_len, hidden_size)
154 + hidden_states = hidden_states.permute(0, 2, 1) # (batch, hidden_size, seq_len)
155 + hidden_states = self.avg_pooler(hidden_states)
156 + hidden_states = hidden_states.permute(
157 + 0, 2, 1
158 + ) # (batch, seq_len/2, hidden_size)
159 + hidden_states = self.layer_norm(hidden_states)
160 +
161 + return hidden_states
162 +
163 + def _get_feat_extract_output_lengths(self, input_lengths: torch.Tensor):
164 + """
165 + Computes the output length of the convolutional layers and the output length
166 + of the audio encoder
167 + """
168 + input_lengths = (input_lengths - 1) // 2 + 1
169 + output_lengths = (input_lengths - 2) // 2 + 1
170 + return input_lengths, output_lengths
171 +
172 +
173 +class AudioFlamingo3MultiModalProjector(nn.Module):
174 + def __init__(self, config: PretrainedConfig):
175 + super().__init__()
176 + self.linear_1 = nn.Linear(
177 + config.audio_config.hidden_size,
178 + config.text_config.hidden_size,
179 + bias=config.projector_bias,
180 + )
181 + self.act = get_act_fn(config.projector_hidden_act)
182 + self.linear_2 = nn.Linear(
183 + config.text_config.hidden_size,
184 + config.text_config.hidden_size,
185 + bias=config.projector_bias,
186 + )
187 +
188 + def forward(self, audio_features):
189 + hidden_states = self.linear_1(audio_features)
190 + hidden_states = self.act(hidden_states)
191 + hidden_states = self.linear_2(hidden_states)
192 + return hidden_states
193 +
194 +
195 +class AudioFlamingo3ProcessingInfo(BaseProcessingInfo):
196 + def get_hf_config(self):
197 + return self.ctx.get_hf_config(AudioFlamingo3Config)
198 +
199 + def get_hf_processor(self, **kwargs: object):
200 + return self.ctx.get_hf_processor(AudioFlamingo3Processor, **kwargs)
201 +
202 + def get_feature_extractor(self, **kwargs: object):
203 + hf_processor = self.get_hf_processor(**kwargs)
204 + feature_extractor = hf_processor.feature_extractor
205 + return feature_extractor
206 +
207 + def get_supported_mm_limits(self) -> Mapping[str, int | None]:
208 + return {"audio": None}
209 +
210 +
211 +class AudioFlamingo3DummyInputsBuilder(
212 + BaseDummyInputsBuilder[AudioFlamingo3ProcessingInfo]
213 +):
214 + def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
215 + num_audios = mm_counts.get("audio", 0)
216 + hf_processor = self.info.get_hf_processor()
217 + audio_token = hf_processor.audio_token
218 + return audio_token * num_audios
219 +
220 + def get_dummy_mm_data(
221 + self,
222 + seq_len: int,
223 + mm_counts: Mapping[str, int],
224 + mm_options: Mapping[str, BaseDummyOptions] | None = None,
225 + ) -> MultiModalDataDict:
226 + feature_extractor = self.info.get_feature_extractor()
227 + sampling_rate = feature_extractor.sampling_rate
228 + audio_len = MAX_AUDIO_LEN * sampling_rate
229 + num_audios = mm_counts.get("audio", 0)
230 + audio_overrides = mm_options.get("audio") if mm_options else None
231 +
232 + return {
233 + "audio": self._get_dummy_audios(
234 + length=audio_len,
235 + num_audios=num_audios,
236 + overrides=audio_overrides,
237 + )
238 + }
239 +
240 +
241 +def _audioflamingo3_field_config(hf_inputs: Mapping[str, torch.Tensor]):
242 + chunk_counts = hf_inputs.get("chunk_counts")
243 + if chunk_counts is not None:
244 + return dict(
245 + audio_embeds=MultiModalFieldConfig.batched("audio"),
246 + input_features=MultiModalFieldConfig.flat_from_sizes(
247 + "audio", chunk_counts, dim=0
248 + ),
249 + feature_attention_mask=MultiModalFieldConfig.flat_from_sizes(
250 + "audio", chunk_counts, dim=0
251 + ),
252 + chunk_counts=MultiModalFieldConfig.batched("audio"),
253 + )
254 + return dict(
255 + audio_embeds=MultiModalFieldConfig.batched("audio"),
256 + input_features=MultiModalFieldConfig.batched("audio"),
257 + feature_attention_mask=MultiModalFieldConfig.batched("audio"),
258 + chunk_counts=MultiModalFieldConfig.batched("audio"),
259 + )
260 +
261 +
262 +class AudioFlamingo3MultiModalDataParser(MultiModalDataParser):
263 + def _parse_audio_data(
264 + self,
265 + data: dict[str, torch.Tensor] | ModalityData[Any],
266 + ) -> ModalityDataItems[Any, Any] | None:
267 + if isinstance(data, dict):
268 + return DictEmbeddingItems(
269 + data,
270 + modality="audio",
271 + required_fields={"audio_embeds"},
272 + fields_factory=_audioflamingo3_field_config,
273 + )
274 + return super()._parse_audio_data(data)
275 +
276 +
277 +class AudioFlamingo3MultiModalProcessor(
278 + BaseMultiModalProcessor[AudioFlamingo3ProcessingInfo]
279 +):
280 + def _get_data_parser(self) -> MultiModalDataParser:
281 + feature_extractor = self.info.get_feature_extractor()
282 + return AudioFlamingo3MultiModalDataParser(
283 + target_sr=feature_extractor.sampling_rate
284 + )
285 +
286 + def _call_hf_processor(
287 + self,
288 + prompt: str,
289 + mm_data: dict[str, object],
290 + mm_kwargs: Mapping[str, Any],
291 + tok_kwargs: Mapping[str, object],
292 + ) -> BatchFeature:
293 + audios = mm_data.pop("audios", [])
294 + if audios:
295 + mm_data["audio"] = audios
296 +
297 + if not mm_data.get("audio", []):
298 + prompt_ids = self.info.get_tokenizer().encode(prompt)
299 + prompt_ids = self._apply_hf_processor_tokens_only(prompt_ids)
300 + return BatchFeature(dict(input_ids=[prompt_ids]), tensor_type="pt")
301 +
302 + feature_extractor = self.info.get_feature_extractor(**mm_kwargs)
303 + mm_kwargs = dict(
304 + **mm_kwargs,
305 + sampling_rate=feature_extractor.sampling_rate,
306 + )
307 +
308 + # Calculate chunk counts
309 + audio_list = mm_data.get("audio")
310 + if not isinstance(audio_list, list):
311 + audio_list = [audio_list]
312 +
313 + chunk_counts = []
314 + sampling_rate = feature_extractor.sampling_rate
315 + chunk_length = feature_extractor.chunk_length
316 + window_size = int(sampling_rate * chunk_length)
317 + # MAX_AUDIO_LEN is 10 * 60 in HF processor.
318 + max_windows = int(MAX_AUDIO_LEN // chunk_length)
319 +
320 + for audio in audio_list:
321 + # audio is numpy array or list
322 + n_samples = len(audio) if isinstance(audio, list) else audio.shape[0]
323 +
324 + n_win = max(1, (n_samples + window_size - 1) // window_size)
325 + if n_win > max_windows:
326 + n_win = max_windows
327 + chunk_counts.append(n_win)
328 +
329 + outputs = super()._call_hf_processor(
330 + prompt=prompt,
331 + mm_data=mm_data,
332 + mm_kwargs=mm_kwargs,
333 + tok_kwargs=tok_kwargs,
334 + )
335 +
336 + if "input_features_mask" in outputs:
337 + outputs["feature_attention_mask"] = outputs.pop("input_features_mask")
338 +
339 + outputs["chunk_counts"] = torch.tensor(chunk_counts, dtype=torch.long)
340 +
341 + return outputs
342 +
343 + def _get_mm_fields_config(
344 + self,
345 + hf_inputs: BatchFeature,
346 + hf_processor_mm_kwargs: Mapping[str, object],
347 + ) -> Mapping[str, MultiModalFieldConfig]:
348 + return _audioflamingo3_field_config(hf_inputs)
349 +
350 + def _get_prompt_updates(
351 + self,
352 + mm_items: MultiModalDataItems,
353 + hf_processor_mm_kwargs: Mapping[str, object],
354 + out_mm_kwargs: MultiModalKwargsItems,
355 + ) -> Sequence[PromptUpdate]:
356 + processor = self.info.get_hf_processor(**hf_processor_mm_kwargs)
357 + tokenizer = self.info.get_tokenizer()
358 + vocab = tokenizer.get_vocab()
359 +
360 + audio_token = getattr(processor, "audio_token", "<sound>")
361 + audio_token_id = vocab.get(audio_token)
362 + if audio_token_id is None:
363 + # Fallback if not found, though it should be there
364 + audio_token_id = processor.audio_token_id
365 +
366 + out_mm_data = out_mm_kwargs.get_data()
367 + feature_attention_mask = out_mm_data.get("feature_attention_mask")
368 + chunk_counts = out_mm_data.get("chunk_counts")
369 +
370 + def get_replacement_audioflamingo3(item_idx: int):
371 + if feature_attention_mask is not None:
372 + if chunk_counts is not None:
373 + counts = (
374 + chunk_counts.tolist()
375 + if isinstance(chunk_counts, torch.Tensor)
376 + else chunk_counts
377 + )
378 + start_idx = sum(counts[:item_idx])
379 + count = counts[item_idx]
380 + end_idx = start_idx + count
381 +
382 + if isinstance(feature_attention_mask, list):
383 + mask_list = feature_attention_mask[start_idx:end_idx]
384 + if len(mask_list) > 0 and isinstance(
385 + mask_list[0], torch.Tensor
386 + ):
387 + mask = torch.stack(mask_list)
388 + else:
389 + mask = torch.tensor(mask_list)
390 + else:
391 + mask = feature_attention_mask[start_idx:end_idx]
392 + else:
393 + # feature_attention_mask is list[Tensor] or Tensor
394 + if isinstance(feature_attention_mask, list):
395 + mask = feature_attention_mask[item_idx]
396 + else:
397 + mask = feature_attention_mask[item_idx].unsqueeze(0)
398 +
399 + # mask shape: (num_chunks, 3000)
400 + input_lengths = mask.sum(-1)
401 + conv_lengths = (input_lengths - 1) // 2 + 1
402 + audio_output_lengths = (conv_lengths - 2) // 2 + 1
403 + num_features = audio_output_lengths.sum().item()
404 + else:
405 + audio_embeds = out_mm_data["audio_embeds"][item_idx]
406 + num_features = audio_embeds.shape[0]
407 +
408 + if num_features == 0:
409 + raise ValueError("Audio is too short")
410 +
411 + audio_tokens = [audio_token_id] * int(num_features)
412 + return PromptUpdateDetails.select_token_id(
413 + audio_tokens,
414 + embed_token_id=audio_token_id,
415 + )
416 +
417 + return [
418 + PromptReplacement(
419 + modality="audio",
420 + target=audio_token,
421 + replacement=get_replacement_audioflamingo3,
422 + )
423 + ]
424 +
425 +
426 +@MULTIMODAL_REGISTRY.register_processor(
427 + AudioFlamingo3MultiModalProcessor,
428 + info=AudioFlamingo3ProcessingInfo,
429 + dummy_inputs=AudioFlamingo3DummyInputsBuilder,
430 +)
431 +class AudioFlamingo3ForConditionalGeneration(
432 + nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA
433 +):
434 + """
435 + AudioFlamingo3 model for conditional generation.
436 +
437 + This model integrates a Whisper-based audio encoder with a Qwen2 language model.
438 + It supports multi-chunk audio processing.
439 + """
440 +
441 + packed_modules_mapping = {
442 + "qkv_proj": ["q_proj", "k_proj", "v_proj"],
443 + "gate_up_proj": ["gate_proj", "up_proj"],
444 + }
445 +
446 + def get_mm_mapping(self) -> MultiModelKeys:
447 + """
448 + Get the module prefix in multimodal models
449 + """
450 + return MultiModelKeys.from_string_field(
451 + language_model="language_model.",
452 + connector="multi_modal_projector.",
453 + tower_model="audio_tower.",
454 + )
455 +
456 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
457 + super().__init__()
458 + config = vllm_config.model_config.hf_config
459 + quant_config = vllm_config.quant_config
460 + multimodal_config = vllm_config.model_config.multimodal_config
461 + self.config = config
462 + self.multimodal_config = multimodal_config
463 +
464 + self.audio_tower = AudioFlamingo3Encoder(
465 + config.audio_config,
466 + )
467 + self.multi_modal_projector = AudioFlamingo3MultiModalProjector(config)
468 +
469 + self.quant_config = quant_config
470 +
471 + self.language_model = init_vllm_registered_model(
472 + vllm_config=vllm_config,
473 + hf_config=config.text_config,
474 + prefix=maybe_prefix(prefix, "language_model"),
475 + architectures=["Qwen2ForCausalLM"],
476 + )
477 +
478 + self.make_empty_intermediate_tensors = (
479 + self.language_model.make_empty_intermediate_tensors
480 + )
481 +
482 + def _parse_and_validate_audio_input(
483 + self, **kwargs: object
484 + ) -> AudioFlamingo3Inputs | None:
485 + input_features = kwargs.pop("input_features", None)
486 + audio_embeds = kwargs.pop("audio_embeds", None)
487 + feature_attention_mask = kwargs.pop("feature_attention_mask", None)
488 + chunk_counts = kwargs.pop("chunk_counts", None)
489 +
490 + if input_features is None and audio_embeds is None:
491 + return None
492 +
493 + if audio_embeds is not None:
494 + return AudioFlamingo3EmbeddingInputs(
495 + type="audio_embeds", audio_embeds=audio_embeds
496 + )
497 +
498 + if input_features is not None:
499 + return AudioFlamingo3FeatureInputs(
500 + type="audio_features",
501 + input_features=input_features,
502 + feature_attention_mask=feature_attention_mask,
503 + chunk_counts=chunk_counts,
504 + )
505 +
506 + raise AssertionError("This line should be unreachable.")
507 +
508 + def _process_audio_input(
509 + self, audio_input: AudioFlamingo3Inputs
510 + ) -> torch.Tensor | tuple[torch.Tensor, ...]:
511 + if audio_input["type"] == "audio_embeds":
512 + audio_embeds = audio_input["audio_embeds"]
513 + return tuple(audio_embeds)
514 +
515 + input_features = audio_input["input_features"]
516 + feature_attention_mask = audio_input["feature_attention_mask"]
517 + chunk_counts = audio_input.get("chunk_counts")
518 +
519 + if isinstance(input_features, list):
520 + input_features = torch.cat(input_features, dim=0)
521 + feature_attention_mask = torch.cat(feature_attention_mask, dim=0)
522 +
523 + if chunk_counts is None:
524 + chunk_counts = [1] * input_features.shape[0]
525 + elif isinstance(chunk_counts, torch.Tensor):
526 + chunk_counts = chunk_counts.tolist()
527 + elif (
528 + isinstance(chunk_counts, list)
529 + and chunk_counts
530 + and isinstance(chunk_counts[0], torch.Tensor)
531 + ):
532 + chunk_counts = [c.item() for c in chunk_counts]
533 +
534 + # Calculate output lengths
535 + input_lengths = feature_attention_mask.sum(-1)
536 + # Conv downsampling
537 + conv_lengths = (input_lengths - 1) // 2 + 1
538 + # AvgPool downsampling
539 + audio_output_lengths = (conv_lengths - 2) // 2 + 1
540 +
541 + batch_size, _, max_mel_seq_len = input_features.shape
542 +
543 + # Calculate max_seq_len after convs (before pooling) for attention mask
544 + max_seq_len = (max_mel_seq_len - 1) // 2 + 1
545 +
546 + # Create a sequence tensor of shape (batch_size, max_seq_len)
547 + seq_range = (
548 + torch.arange(
549 + 0,
550 + max_seq_len,
551 + dtype=conv_lengths.dtype,
552 + device=conv_lengths.device,
553 + )
554 + .unsqueeze(0)
555 + .expand(batch_size, max_seq_len)
556 + )
557 + lengths_expand = conv_lengths.unsqueeze(-1).expand(batch_size, max_seq_len)
558 + # Create mask
559 + padding_mask = seq_range >= lengths_expand
560 +
561 + audio_attention_mask_ = padding_mask.view(batch_size, 1, 1, max_seq_len).expand(
562 + batch_size, 1, max_seq_len, max_seq_len
563 + )
564 + audio_attention_mask = audio_attention_mask_.to(
565 + dtype=self.audio_tower.conv1.weight.dtype,
566 + device=self.audio_tower.conv1.weight.device,
567 + )
568 + audio_attention_mask[audio_attention_mask_] = float("-inf")
569 +
570 + # Forward pass
571 + audio_features = self.audio_tower(
572 + input_features, attention_mask=audio_attention_mask
573 + )
574 +
575 + # Project
576 + audio_features = self.multi_modal_projector(audio_features)
577 +
578 + # Masking after pooling
579 + num_audios, max_audio_tokens, embed_dim = audio_features.shape
580 + audio_output_lengths = audio_output_lengths.unsqueeze(1)
581 + audio_features_mask = (
582 + torch.arange(max_audio_tokens)
583 + .expand(num_audios, max_audio_tokens)
584 + .to(audio_output_lengths.device)
585 + < audio_output_lengths
586 + )
587 + masked_audio_features = audio_features[audio_features_mask].view(-1, embed_dim)
588 +
589 + # Split to tuple of embeddings for individual audio input.
590 + chunk_embeddings = torch.split(
591 + masked_audio_features, audio_output_lengths.flatten().tolist()
592 + )
593 +
594 + grouped_embeddings = []
595 + current_idx = 0
596 + for count in chunk_counts:
597 + audio_chunks = chunk_embeddings[current_idx : current_idx + count]
598 + grouped_embeddings.append(torch.cat(audio_chunks, dim=0))
599 + current_idx += count
600 + return tuple(grouped_embeddings)
601 +
602 + def get_language_model(self) -> torch.nn.Module:
603 + return self.language_model
604 +
605 + def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings:
606 + audio_input = self._parse_and_validate_audio_input(**kwargs)
607 + if audio_input is None:
608 + return []
609 + masked_audio_features = self._process_audio_input(audio_input)
610 + return masked_audio_features
611 +
612 + def forward(
613 + self,
614 + input_ids: torch.Tensor,
615 + positions: torch.Tensor,
616 + intermediate_tensors: IntermediateTensors | None = None,
617 + inputs_embeds: torch.Tensor | None = None,
618 + **kwargs: object,
619 + ) -> torch.Tensor | IntermediateTensors:
620 + if intermediate_tensors is not None:
621 + inputs_embeds = None
622 +
623 + hidden_states = self.language_model.model(
624 + input_ids,
625 + positions,
626 + intermediate_tensors,
627 + inputs_embeds=inputs_embeds,
628 + )
629 + return hidden_states
630 +
631 + def compute_logits(
632 + self,
633 + hidden_states: torch.Tensor,
634 + ) -> torch.Tensor | None:
635 + return self.language_model.compute_logits(hidden_states)
636 +
637 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
638 + loader = AutoWeightsLoader(self)
639 + return loader.load_weights(weights)
vllm/model_executor/models/registry.py
+4
@@ -264,6 +264,10 @@ _CROSS_ENCODER_MODELS = {
264 _MULTIMODAL_MODELS = {
265 # [Decoder-only]
266 "AriaForConditionalGeneration": ("aria", "AriaForConditionalGeneration"),
267 + "AudioFlamingo3ForConditionalGeneration": (
268 + "audioflamingo3",
269 + "AudioFlamingo3ForConditionalGeneration",
270 + ),
271 "AyaVisionForConditionalGeneration": (
272 "aya_vision",
273 "AyaVisionForConditionalGeneration",