| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """This example shows how to use vLLM for running offline inference |
| 4 | with the correct prompt format on audio language models. |
| 5 | |
| 6 | For most models, the prompt format should follow corresponding examples |
| 7 | on HuggingFace model repository. |
| 8 | """ |
| 9 | |
| 10 | import os |
| 11 | from typing import Any, NamedTuple |
| 12 | |
| 13 | from huggingface_hub import snapshot_download |
| 14 | from transformers import AutoTokenizer |
| 15 | |
| 16 | from vllm import LLM, EngineArgs, SamplingParams |
| 17 | from vllm.assets.audio import AudioAsset |
| 18 | from vllm.lora.request import LoRARequest |
| 19 | from vllm.utils.argparse_utils import FlexibleArgumentParser |
| 20 | |
| 21 | audio_assets = [AudioAsset("mary_had_lamb"), AudioAsset("winning_call")] |
| 22 | question_per_audio_count = { |
| 23 | 0: "What is 1+1?", |
| 24 | 1: "What is recited in the audio?", |
| 25 | 2: "What sport and what nursery rhyme are referenced?", |
| 26 | } |
| 27 | |
| 28 | |
| 29 | class ModelRequestData(NamedTuple): |
| 30 | engine_args: EngineArgs |
| 31 | prompt: str | None = None |
| 32 | prompt_token_ids: dict[str, list[int]] | None = None |
| 33 | multi_modal_data: dict[str, Any] | None = None |
| 34 | stop_token_ids: list[int] | None = None |
| 35 | lora_requests: list[LoRARequest] | None = None |
| 36 | |
| 37 | |
| 38 | # NOTE: The default `max_num_seqs` and `max_model_len` may result in OOM on |
| 39 | # lower-end GPUs. |
| 40 | # Unless specified, these settings have been tested to work on a single L4. |
| 41 | |
| 42 | |
| 43 | # AudioFlamingo3 |
| 44 | def run_audioflamingo3(question: str, audio_count: int) -> ModelRequestData: |
| 45 | model_name = "nvidia/audio-flamingo-3-hf" |
| 46 | engine_args = EngineArgs( |
| 47 | model=model_name, |
| 48 | max_model_len=4096, |
| 49 | max_num_seqs=2, |
| 50 | limit_mm_per_prompt={"audio": audio_count}, |
| 51 | enforce_eager=True, |
| 52 | ) |
| 53 | |
| 54 | # AudioFlamingo3 uses <sound> token for audio |
| 55 | audio_placeholder = "<sound>" * audio_count |
| 56 | |
| 57 | prompt = ( |
| 58 | "<|im_start|>system\n" |
| 59 | "You are a helpful assistant.<|im_end|>\n" |
| 60 | "<|im_start|>user\n" |
| 61 | f"{audio_placeholder}{question}<|im_end|>\n" |
| 62 | "<|im_start|>assistant\n" |
| 63 | ) |
| 64 | |
| 65 | return ModelRequestData( |
| 66 | engine_args=engine_args, |
| 67 | prompt=prompt, |
| 68 | ) |
| 69 | |
| 70 | |
| 71 | # CohereASR |
| 72 | def run_cohere_asr(question: str, audio_count: int) -> ModelRequestData: |
| 73 | assert audio_count == 1, "CohereASR only support single audio input per prompt" |
| 74 | model_name = "CohereLabs/cohere-transcribe-03-2026" |
| 75 | |
| 76 | prompt = ( |
| 77 | "<|startofcontext|><|startoftranscript|>" |
| 78 | "<|emo:undefined|><|en|><|en|><|pnc|><|noitn|>" |
| 79 | "<|notimestamp|><|nodiarize|>" |
| 80 | ) |
| 81 | engine_args = EngineArgs( |
| 82 | model=model_name, |
| 83 | limit_mm_per_prompt={"audio": audio_count}, |
| 84 | trust_remote_code=True, |
| 85 | ) |
| 86 | |
| 87 | return ModelRequestData( |
| 88 | engine_args=engine_args, |
| 89 | prompt=prompt, |
| 90 | ) |
| 91 | |
| 92 | |
| 93 | # Gemma3N |
| 94 | def run_gemma3n(question: str, audio_count: int) -> ModelRequestData: |
| 95 | model_name = "google/gemma-3n-E2B-it" |
| 96 | engine_args = EngineArgs( |
| 97 | model=model_name, |
| 98 | max_model_len=2048, |
| 99 | max_num_batched_tokens=2048, |
| 100 | max_num_seqs=2, |
| 101 | limit_mm_per_prompt={"audio": audio_count}, |
| 102 | enforce_eager=True, |
| 103 | ) |
| 104 | prompt = f"<start_of_turn>user\n<audio_soft_token>{question}" |
| 105 | "<end_of_turn>\n<start_of_turn>model\n" |
| 106 | return ModelRequestData( |
| 107 | engine_args=engine_args, |
| 108 | prompt=prompt, |
| 109 | ) |
| 110 | |
| 111 | |
| 112 | # GLM-ASR |
| 113 | def run_glmasr(question: str, audio_count: int) -> ModelRequestData: |
| 114 | model_name = "zai-org/GLM-ASR-Nano-2512" |
| 115 | |
| 116 | tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True) |
| 117 | |
| 118 | # GLM-ASR uses <|pad|> token for audio |
| 119 | audio_placeholder = "<|pad|>" * audio_count |
| 120 | |
| 121 | messages = [{"role": "user", "content": f"{audio_placeholder}{question}"}] |
| 122 | prompt = tokenizer.apply_chat_template( |
| 123 | messages, tokenize=False, add_generation_prompt=True |
| 124 | ) |
| 125 | |
| 126 | engine_args = EngineArgs( |
| 127 | model=model_name, |
| 128 | trust_remote_code=True, |
| 129 | max_model_len=4096, |
| 130 | max_num_seqs=2, |
| 131 | limit_mm_per_prompt={"audio": audio_count}, |
| 132 | ) |
| 133 | |
| 134 | return ModelRequestData( |
| 135 | engine_args=engine_args, |
| 136 | prompt=prompt, |
| 137 | ) |
| 138 | |
| 139 | |
| 140 | # FunAudioChat |
| 141 | def run_funaudiochat(question: str, audio_count: int) -> ModelRequestData: |
| 142 | # NOTE: FunAudioChat is not available on the HuggingFace Hub at the time of |
| 143 | # writing. Pass a local model path via `--model`. |
| 144 | model_name = "funaudiochat" |
| 145 | |
| 146 | engine_args = EngineArgs( |
| 147 | model=model_name, |
| 148 | max_model_len=4096, |
| 149 | max_num_seqs=2, |
| 150 | limit_mm_per_prompt={"audio": audio_count}, |
| 151 | enforce_eager=True, |
| 152 | ) |
| 153 | |
| 154 | audio_in_prompt = "".join( |
| 155 | ["<|audio_bos|><|AUDIO|><|audio_eos|>\n" for _ in range(audio_count)] |
| 156 | ) |
| 157 | prompt = f"{audio_in_prompt}{question}" |
| 158 | |
| 159 | return ModelRequestData( |
| 160 | engine_args=engine_args, |
| 161 | prompt=prompt, |
| 162 | ) |
| 163 | |
| 164 | |
| 165 | # Granite Speech |
| 166 | def run_granite_speech(question: str, audio_count: int) -> ModelRequestData: |
| 167 | # NOTE - the setting in this example are somewhat different from what is |
| 168 | # optimal for granite speech, and it is generally recommended to use beam |
| 169 | # search. Check the model README for suggested settings. |
| 170 | # https://huggingface.co/ibm-granite/granite-speech-3.3-8b |
| 171 | model_name = "ibm-granite/granite-speech-3.3-8b" |
| 172 | |
| 173 | engine_args = EngineArgs( |
| 174 | model=model_name, |
| 175 | trust_remote_code=True, |
| 176 | max_model_len=2048, |
| 177 | max_num_seqs=2, |
| 178 | enable_lora=True, |
| 179 | max_lora_rank=64, |
| 180 | limit_mm_per_prompt={"audio": audio_count}, |
| 181 | ) |
| 182 | |
| 183 | # The model has an audio-specific lora directly in its model dir; |
| 184 | # it should be enabled whenever you pass audio inputs to the model. |
| 185 | speech_lora_path = model_name |
| 186 | audio_placeholder = "<|audio|>" * audio_count |
| 187 | prompts = f"<|start_of_role|>system<|end_of_role|>Knowledge Cutoff Date: April 2024.\nToday's Date: December 19, 2024.\nYou are Granite, developed by IBM. You are a helpful AI assistant<|end_of_text|>\n<|start_of_role|>user<|end_of_role|>{audio_placeholder}{question}<|end_of_text|>\n<|start_of_role|>assistant<|end_of_role|>" # noqa: E501 |
| 188 | |
| 189 | return ModelRequestData( |
| 190 | engine_args=engine_args, |
| 191 | prompt=prompts, |
| 192 | lora_requests=[LoRARequest("speech", 1, speech_lora_path)], |
| 193 | ) |
| 194 | |
| 195 | |
| 196 | # Kimi-Audio-7B-Instruct |
| 197 | def run_kimi_audio(question: str, audio_count: int) -> ModelRequestData: |
| 198 | """Kimi-Audio-7B-Instruct for audio transcription and understanding.""" |
| 199 | model_name = "moonshotai/Kimi-Audio-7B-Instruct" |
| 200 | |
| 201 | engine_args = EngineArgs( |
| 202 | model=model_name, |
| 203 | trust_remote_code=True, |
| 204 | max_model_len=4096, |
| 205 | max_num_seqs=2, |
| 206 | limit_mm_per_prompt={"audio": audio_count}, |
| 207 | ) |
| 208 | |
| 209 | # Kimi-Audio uses <|im_kimia_text_blank|> as placeholder for audio features |
| 210 | audio_placeholder = "<|im_kimia_text_blank|>" * audio_count |
| 211 | # Default prompt for transcription |
| 212 | if not question: |
| 213 | question = "Please transcribe the audio" |
| 214 | prompt = f"{audio_placeholder}{question}" |
| 215 | |
| 216 | # Stop at EOS token (151644) to prevent repetition |
| 217 | return ModelRequestData( |
| 218 | engine_args=engine_args, |
| 219 | prompt=prompt, |
| 220 | stop_token_ids=[151644], |
| 221 | ) |
| 222 | |
| 223 | |
| 224 | # MiDashengLM |
| 225 | def run_midashenglm(question: str, audio_count: int): |
| 226 | model_name = "mispeech/midashenglm-7b" |
| 227 | |
| 228 | engine_args = EngineArgs( |
| 229 | model=model_name, |
| 230 | trust_remote_code=True, |
| 231 | max_model_len=4096, |
| 232 | max_num_seqs=5, |
| 233 | limit_mm_per_prompt={"audio": audio_count}, |
| 234 | ) |
| 235 | |
| 236 | audio_in_prompt = "".join( |
| 237 | ["<|audio_bos|><|AUDIO|><|audio_eos|>" for idx in range(audio_count)] |
| 238 | ) |
| 239 | |
| 240 | default_system = "You are a helpful language and speech assistant." |
| 241 | |
| 242 | prompt = ( |
| 243 | f"<|im_start|>system\n{default_system}<|im_end|>\n" |
| 244 | "<|im_start|>user\n" |
| 245 | f"{audio_in_prompt}{question}<|im_end|>\n" |
| 246 | "<|im_start|>assistant\n" |
| 247 | ) |
| 248 | return ModelRequestData( |
| 249 | engine_args=engine_args, |
| 250 | prompt=prompt, |
| 251 | ) |
| 252 | |
| 253 | |
| 254 | # MiniCPM-O |
| 255 | def run_minicpmo(question: str, audio_count: int) -> ModelRequestData: |
| 256 | model_name = "openbmb/MiniCPM-o-2_6" |
| 257 | tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True) |
| 258 | engine_args = EngineArgs( |
| 259 | model=model_name, |
| 260 | trust_remote_code=True, |
| 261 | max_model_len=4096, |
| 262 | max_num_seqs=2, |
| 263 | limit_mm_per_prompt={"audio": audio_count}, |
| 264 | ) |
| 265 | |
| 266 | stop_tokens = ["<|im_end|>", "<|endoftext|>"] |
| 267 | stop_token_ids = [tokenizer.convert_tokens_to_ids(i) for i in stop_tokens] |
| 268 | |
| 269 | audio_placeholder = "(<audio>./</audio>)" * audio_count |
| 270 | audio_chat_template = "{% for message in messages %}{{'<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>' + '\n'}}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant\n<|spk_bos|><|spk|><|spk_eos|><|tts_bos|>' }}{% endif %}" # noqa: E501 |
| 271 | messages = [{"role": "user", "content": f"{audio_placeholder}\n{question}"}] |
| 272 | prompt = tokenizer.apply_chat_template( |
| 273 | messages, |
| 274 | tokenize=False, |
| 275 | add_generation_prompt=True, |
| 276 | chat_template=audio_chat_template, |
| 277 | ) |
| 278 | |
| 279 | return ModelRequestData( |
| 280 | engine_args=engine_args, |
| 281 | prompt=prompt, |
| 282 | stop_token_ids=stop_token_ids, |
| 283 | ) |
| 284 | |
| 285 | |
| 286 | # Phi-4-multimodal-instruct |
| 287 | def run_phi4mm(question: str, audio_count: int) -> ModelRequestData: |
| 288 | """Phi-4-multimodal-instruct supports both image and audio inputs. Here, we |
| 289 | show how to process audio inputs. |
| 290 | """ |
| 291 | model_path = snapshot_download("microsoft/Phi-4-multimodal-instruct") |
| 292 | # Since the vision-lora and speech-lora co-exist with the base model, |
| 293 | # we have to manually specify the path of the lora weights. |
| 294 | speech_lora_path = os.path.join(model_path, "speech-lora") |
| 295 | placeholders = "".join([f"<|audio_{i + 1}|>" for i in range(audio_count)]) |
| 296 | |
| 297 | prompts = f"<|user|>{placeholders}{question}<|end|><|assistant|>" |
| 298 | |
| 299 | engine_args = EngineArgs( |
| 300 | model=model_path, |
| 301 | trust_remote_code=True, |
| 302 | max_model_len=12800, |
| 303 | max_num_seqs=2, |
| 304 | enable_lora=True, |
| 305 | max_lora_rank=320, |
| 306 | limit_mm_per_prompt={"audio": audio_count}, |
| 307 | ) |
| 308 | |
| 309 | return ModelRequestData( |
| 310 | engine_args=engine_args, |
| 311 | prompt=prompts, |
| 312 | lora_requests=[LoRARequest("speech", 1, speech_lora_path)], |
| 313 | ) |
| 314 | |
| 315 | |
| 316 | # Qwen2-Audio |
| 317 | def run_qwen2_audio(question: str, audio_count: int) -> ModelRequestData: |
| 318 | model_name = "Qwen/Qwen2-Audio-7B-Instruct" |
| 319 | |
| 320 | engine_args = EngineArgs( |
| 321 | model=model_name, |
| 322 | max_model_len=4096, |
| 323 | max_num_seqs=5, |
| 324 | limit_mm_per_prompt={"audio": audio_count}, |
| 325 | ) |
| 326 | |
| 327 | audio_in_prompt = "".join( |
| 328 | [ |
| 329 | f"Audio {idx + 1}: <|audio_bos|><|AUDIO|><|audio_eos|>\n" |
| 330 | for idx in range(audio_count) |
| 331 | ] |
| 332 | ) |
| 333 | |
| 334 | prompt = ( |
| 335 | "<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n" |
| 336 | "<|im_start|>user\n" |
| 337 | f"{audio_in_prompt}{question}<|im_end|>\n" |
| 338 | "<|im_start|>assistant\n" |
| 339 | ) |
| 340 | |
| 341 | return ModelRequestData( |
| 342 | engine_args=engine_args, |
| 343 | prompt=prompt, |
| 344 | ) |
| 345 | |
| 346 | |
| 347 | # Qwen2.5-Omni |
| 348 | def run_qwen2_5_omni(question: str, audio_count: int): |
| 349 | model_name = "Qwen/Qwen2.5-Omni-7B" |
| 350 | |
| 351 | engine_args = EngineArgs( |
| 352 | model=model_name, |
| 353 | max_model_len=4096, |
| 354 | max_num_seqs=5, |
| 355 | limit_mm_per_prompt={"audio": audio_count}, |
| 356 | ) |
| 357 | |
| 358 | audio_in_prompt = "".join( |
| 359 | ["<|audio_bos|><|AUDIO|><|audio_eos|>\n" for idx in range(audio_count)] |
| 360 | ) |
| 361 | |
| 362 | default_system = ( |
| 363 | "You are Qwen, a virtual human developed by the Qwen Team, Alibaba " |
| 364 | "Group, capable of perceiving auditory and visual inputs, as well as " |
| 365 | "generating text and speech." |
| 366 | ) |
| 367 | |
| 368 | prompt = ( |
| 369 | f"<|im_start|>system\n{default_system}<|im_end|>\n" |
| 370 | "<|im_start|>user\n" |
| 371 | f"{audio_in_prompt}{question}<|im_end|>\n" |
| 372 | "<|im_start|>assistant\n" |
| 373 | ) |
| 374 | return ModelRequestData( |
| 375 | engine_args=engine_args, |
| 376 | prompt=prompt, |
| 377 | ) |
| 378 | |
| 379 | |
| 380 | def run_qwen3_asr(question: str, audio_count: int) -> ModelRequestData: |
| 381 | model_name = "Qwen/Qwen3-Asr-1.7B" |
| 382 | |
| 383 | audio_in_prompt = "<|audio_start|><|audio_pad|><|audio_end|>\n" * audio_count |
| 384 | prompt = f"<|im_start|>user\n{audio_in_prompt}<|im_end|>\n<|im_start|>assistant\n" |
| 385 | |
| 386 | engine_args = EngineArgs( |
| 387 | model=model_name, |
| 388 | max_model_len=4096, |
| 389 | max_num_seqs=5, |
| 390 | limit_mm_per_prompt={"audio": audio_count}, |
| 391 | ) |
| 392 | |
| 393 | return ModelRequestData( |
| 394 | engine_args=engine_args, |
| 395 | prompt=prompt, |
| 396 | ) |
| 397 | |
| 398 | |
| 399 | # Ultravox 0.5-1B |
| 400 | def run_ultravox(question: str, audio_count: int) -> ModelRequestData: |
| 401 | model_name = "fixie-ai/ultravox-v0_5-llama-3_2-1b" |
| 402 | |
| 403 | tokenizer = AutoTokenizer.from_pretrained(model_name) |
| 404 | messages = [{"role": "user", "content": "<|audio|>\n" * audio_count + question}] |
| 405 | prompt = tokenizer.apply_chat_template( |
| 406 | messages, tokenize=False, add_generation_prompt=True |
| 407 | ) |
| 408 | |
| 409 | engine_args = EngineArgs( |
| 410 | model=model_name, |
| 411 | max_model_len=4096, |
| 412 | max_num_seqs=5, |
| 413 | trust_remote_code=True, |
| 414 | limit_mm_per_prompt={"audio": audio_count}, |
| 415 | ) |
| 416 | |
| 417 | return ModelRequestData( |
| 418 | engine_args=engine_args, |
| 419 | prompt=prompt, |
| 420 | ) |
| 421 | |
| 422 | |
| 423 | # Voxtral |
| 424 | # Make sure to install mistral-common[audio]. |
| 425 | def run_voxtral(question: str, audio_count: int) -> ModelRequestData: |
| 426 | from mistral_common.protocol.instruct.chunk import ( |
| 427 | AudioChunk, |
| 428 | TextChunk, |
| 429 | ) |
| 430 | from mistral_common.protocol.instruct.messages import ( |
| 431 | UserMessage, |
| 432 | ) |
| 433 | from mistral_common.protocol.instruct.request import ChatCompletionRequest |
| 434 | from mistral_common.tokens.tokenizers.audio import Audio |
| 435 | from mistral_common.tokens.tokenizers.mistral import MistralTokenizer |
| 436 | |
| 437 | model_name = "mistralai/Voxtral-Mini-3B-2507" |
| 438 | tokenizer = MistralTokenizer.from_hf_hub(model_name) |
| 439 | |
| 440 | engine_args = EngineArgs( |
| 441 | model=model_name, |
| 442 | max_model_len=8192, |
| 443 | max_num_seqs=2, |
| 444 | limit_mm_per_prompt={"audio": audio_count}, |
| 445 | config_format="mistral", |
| 446 | load_format="mistral", |
| 447 | tokenizer_mode="mistral", |
| 448 | enforce_eager=True, |
| 449 | enable_chunked_prefill=False, |
| 450 | ) |
| 451 | |
| 452 | text_chunk = TextChunk(text=question) |
| 453 | audios = [ |
| 454 | Audio.from_file(str(audio_assets[i].get_local_path()), strict=False) |
| 455 | for i in range(audio_count) |
| 456 | ] |
| 457 | audio_chunks = [AudioChunk.from_audio(audio) for audio in audios] |
| 458 | |
| 459 | messages = [UserMessage(content=[*audio_chunks, text_chunk])] |
| 460 | |
| 461 | req = ChatCompletionRequest(messages=messages, model=model_name) |
| 462 | |
| 463 | tokens = tokenizer.encode_chat_completion(req) |
| 464 | prompt_ids, audios = tokens.tokens, tokens.audios |
| 465 | |
| 466 | audios_and_sr = [(au.audio_array, au.sampling_rate) for au in audios] |
| 467 | |
| 468 | multi_modal_data = {"audio": audios_and_sr} |
| 469 | |
| 470 | return ModelRequestData( |
| 471 | engine_args=engine_args, |
| 472 | prompt_token_ids=prompt_ids, |
| 473 | multi_modal_data=multi_modal_data, |
| 474 | ) |
| 475 | |
| 476 | |
| 477 | # Whisper |
| 478 | def run_whisper(question: str, audio_count: int) -> ModelRequestData: |
| 479 | assert audio_count == 1, "Whisper only support single audio input per prompt" |
| 480 | model_name = "openai/whisper-large-v3-turbo" |
| 481 | |
| 482 | prompt = "<|startoftranscript|>" |
| 483 | |
| 484 | engine_args = EngineArgs( |
| 485 | model=model_name, |
| 486 | max_model_len=448, |
| 487 | max_num_seqs=5, |
| 488 | limit_mm_per_prompt={"audio": audio_count}, |
| 489 | ) |
| 490 | |
| 491 | return ModelRequestData( |
| 492 | engine_args=engine_args, |
| 493 | prompt=prompt, |
| 494 | ) |
| 495 | |
| 496 | |
| 497 | model_example_map = { |
| 498 | "audioflamingo3": run_audioflamingo3, |
| 499 | "cohere_asr": run_cohere_asr, |
| 500 | "funaudiochat": run_funaudiochat, |
| 501 | "gemma3n": run_gemma3n, |
| 502 | "glmasr": run_glmasr, |
| 503 | "granite_speech": run_granite_speech, |
| 504 | "kimi_audio": run_kimi_audio, |
| 505 | "midashenglm": run_midashenglm, |
| 506 | "minicpmo": run_minicpmo, |
| 507 | "phi4_mm": run_phi4mm, |
| 508 | "qwen2_audio": run_qwen2_audio, |
| 509 | "qwen2_5_omni": run_qwen2_5_omni, |
| 510 | "qwen3_asr": run_qwen3_asr, |
| 511 | "ultravox": run_ultravox, |
| 512 | "voxtral": run_voxtral, |
| 513 | "whisper": run_whisper, |
| 514 | } |
| 515 | |
| 516 | |
| 517 | def parse_args(): |
| 518 | parser = FlexibleArgumentParser( |
| 519 | description="Demo on using vLLM for offline inference with " |
| 520 | "audio language models" |
| 521 | ) |
| 522 | parser.add_argument( |
| 523 | "--model-type", |
| 524 | "-m", |
| 525 | type=str, |
| 526 | default="ultravox", |
| 527 | choices=model_example_map.keys(), |
| 528 | help='Huggingface "model_type".', |
| 529 | ) |
| 530 | parser.add_argument( |
| 531 | "--model", |
| 532 | type=str, |
| 533 | default=None, |
| 534 | help="Model ID or local path override. Required for funaudiochat.", |
| 535 | ) |
| 536 | parser.add_argument( |
| 537 | "--num-prompts", type=int, default=1, help="Number of prompts to run." |
| 538 | ) |
| 539 | parser.add_argument( |
| 540 | "--num-audios", |
| 541 | type=int, |
| 542 | default=1, |
| 543 | choices=[0, 1, 2], |
| 544 | help="Number of audio items per prompt.", |
| 545 | ) |
| 546 | parser.add_argument( |
| 547 | "--seed", |
| 548 | type=int, |
| 549 | default=0, |
| 550 | help="Set the seed when initializing `vllm.LLM`.", |
| 551 | ) |
| 552 | parser.add_argument( |
| 553 | "--tensor-parallel-size", |
| 554 | "-tp", |
| 555 | type=int, |
| 556 | default=None, |
| 557 | help="Tensor parallel size to override the model's default setting. ", |
| 558 | ) |
| 559 | |
| 560 | return parser.parse_args() |
| 561 | |
| 562 | |
| 563 | def main(args): |
| 564 | model = args.model_type |
| 565 | if model not in model_example_map: |
| 566 | raise ValueError(f"Model type {model} is not supported.") |
| 567 | |
| 568 | if model == "funaudiochat" and not args.model: |
| 569 | raise ValueError("--model is required when --model-type=funaudiochat") |
| 570 | |
| 571 | if args.tensor_parallel_size is not None and args.tensor_parallel_size < 1: |
| 572 | raise ValueError( |
| 573 | f"tensor_parallel_size must be a positive integer, " |
| 574 | f"got {args.tensor_parallel_size}" |
| 575 | ) |
| 576 | |
| 577 | audio_count = args.num_audios |
| 578 | req_data = model_example_map[model]( |
| 579 | question_per_audio_count[audio_count], audio_count |
| 580 | ) |
| 581 | if model == "funaudiochat": |
| 582 | req_data.engine_args.model = args.model |
| 583 | |
| 584 | # Disable other modalities to save memory |
| 585 | default_limits = {"image": 0, "video": 0, "audio": 0} |
| 586 | req_data.engine_args.limit_mm_per_prompt = default_limits | dict( |
| 587 | req_data.engine_args.limit_mm_per_prompt or {} |
| 588 | ) |
| 589 | |
| 590 | engine_args = vars(req_data.engine_args) | {"seed": args.seed} |
| 591 | if args.tensor_parallel_size is not None: |
| 592 | engine_args["tensor_parallel_size"] = args.tensor_parallel_size |
| 593 | llm = LLM(**engine_args) |
| 594 | |
| 595 | # We set temperature to 0.2 so that outputs can be different |
| 596 | # even when all prompts are identical when running batch inference. |
| 597 | sampling_params = SamplingParams( |
| 598 | temperature=0.2, max_tokens=64, stop_token_ids=req_data.stop_token_ids |
| 599 | ) |
| 600 | |
| 601 | def get_input(start, end): |
| 602 | mm_data = req_data.multi_modal_data |
| 603 | if not mm_data: |
| 604 | mm_data = {} |
| 605 | if end - start > 0: |
| 606 | mm_data = { |
| 607 | "audio": [ |
| 608 | asset.audio_and_sample_rate for asset in audio_assets[start:end] |
| 609 | ] |
| 610 | } |
| 611 | |
| 612 | inputs = {"multi_modal_data": mm_data} |
| 613 | |
| 614 | if req_data.prompt: |
| 615 | inputs["prompt"] = req_data.prompt |
| 616 | else: |
| 617 | inputs["prompt_token_ids"] = req_data.prompt_token_ids |
| 618 | |
| 619 | return inputs |
| 620 | |
| 621 | # Batch inference |
| 622 | assert args.num_prompts > 0 |
| 623 | if audio_count != 1: |
| 624 | inputs = get_input(0, audio_count) |
| 625 | inputs = [inputs] * args.num_prompts |
| 626 | else: |
| 627 | # For single audio input, we need to vary the audio input |
| 628 | # to avoid deduplication in vLLM engine. |
| 629 | inputs = [] |
| 630 | for i in range(args.num_prompts): |
| 631 | start = i % len(audio_assets) |
| 632 | inp = get_input(start, start + 1) |
| 633 | inputs.append(inp) |
| 634 | |
| 635 | # Add LoRA request if applicable |
| 636 | lora_request = ( |
| 637 | req_data.lora_requests * args.num_prompts if req_data.lora_requests else None |
| 638 | ) |
| 639 | |
| 640 | outputs = llm.generate( |
| 641 | inputs, |
| 642 | sampling_params=sampling_params, |
| 643 | lora_request=lora_request, |
| 644 | ) |
| 645 | |
| 646 | for o in outputs: |
| 647 | generated_text = o.outputs[0].text |
| 648 | print(generated_text) |
| 649 | |
| 650 | |
| 651 | if __name__ == "__main__": |
| 652 | args = parse_args() |
| 653 | main(args) |