# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """ This example shows how to use vLLM-omni for running offline inference with the correct prompt format on Qwen2.5-Omni """ import os from typing import NamedTuple import soundfile as sf from vllm.assets.audio import AudioAsset from vllm.assets.image import ImageAsset from vllm.assets.video import VideoAsset from vllm.multimodal.image import convert_image_mode from vllm.sampling_params import SamplingParams from vllm.utils import FlexibleArgumentParser from vllm_omni import Omni SEED = 42 class QueryResult(NamedTuple): inputs: dict limit_mm_per_prompt: dict[str, int] # NOTE: The default `max_num_seqs` and `max_model_len` may result in OOM on # lower-end GPUs. # Unless specified, these settings have been tested to work on a single L4. default_system = ( "You are Qwen, a virtual human developed by the Qwen Team, Alibaba " "Group, capable of perceiving auditory and visual inputs, as well as " "generating text and speech." ) def get_text_query(question: str = None) -> QueryResult: if question is None: question = "Explain the system architecture for a scalable audio generation pipeline. Answer in 15 words." prompt = ( f"<|im_start|>system\n{default_system}<|im_end|>\n" "<|im_start|>user\n" f"{question}<|im_end|>\n" f"<|im_start|>assistant\n" ) return QueryResult( inputs={ "prompt": prompt, }, limit_mm_per_prompt={}, ) def get_mixed_modalities_query() -> QueryResult: question = "What is recited in the audio? What is the content of this image? Why is this video funny?" prompt = ( f"<|im_start|>system\n{default_system}<|im_end|>\n" "<|im_start|>user\n<|audio_bos|><|AUDIO|><|audio_eos|>" "<|vision_bos|><|IMAGE|><|vision_eos|>" "<|vision_bos|><|VIDEO|><|vision_eos|>" f"{question}<|im_end|>\n" f"<|im_start|>assistant\n" ) return QueryResult( inputs={ "prompt": prompt, "multi_modal_data": { "audio": AudioAsset("mary_had_lamb").audio_and_sample_rate, "image": convert_image_mode(ImageAsset("cherry_blossom").pil_image, "RGB"), "video": VideoAsset(name="baby_reading", num_frames=16).np_ndarrays, }, }, limit_mm_per_prompt={"audio": 1, "image": 1, "video": 1}, ) def get_use_audio_in_video_query() -> QueryResult: question = "Describe the content of the video, then convert what the baby say into text." prompt = ( f"<|im_start|>system\n{default_system}<|im_end|>\n" "<|im_start|>user\n<|vision_bos|><|VIDEO|><|vision_eos|>" f"{question}<|im_end|>\n" f"<|im_start|>assistant\n" ) asset = VideoAsset(name="baby_reading", num_frames=16) audio = asset.get_audio(sampling_rate=16000) return QueryResult( inputs={ "prompt": prompt, "multi_modal_data": { "video": asset.np_ndarrays, "audio": audio, }, "mm_processor_kwargs": { "use_audio_in_video": True, }, }, limit_mm_per_prompt={"audio": 1, "video": 1}, ) def get_multi_audios_query() -> QueryResult: question = "Are these two audio clips the same?" prompt = ( f"<|im_start|>system\n{default_system}<|im_end|>\n" "<|im_start|>user\n<|audio_bos|><|AUDIO|><|audio_eos|>" "<|audio_bos|><|AUDIO|><|audio_eos|>" f"{question}<|im_end|>\n" f"<|im_start|>assistant\n" ) return QueryResult( inputs={ "prompt": prompt, "multi_modal_data": { "audio": [ AudioAsset("winning_call").audio_and_sample_rate, AudioAsset("mary_had_lamb").audio_and_sample_rate, ], }, }, limit_mm_per_prompt={ "audio": 2, }, ) query_map = { "mixed_modalities": get_mixed_modalities_query, "use_audio_in_video": get_use_audio_in_video_query, "multi_audios": get_multi_audios_query, "text": get_text_query, } def main(args): model_name = "Qwen/Qwen2.5-Omni-7B" query_result = query_map[args.query_type]() omni_llm = Omni( model=model_name, log_stats=args.enable_stats, log_file=("omni_llm_pipeline.log" if args.enable_stats else None), init_sleep_seconds=args.init_sleep_seconds, batch_timeout=args.batch_timeout, init_timeout=args.init_timeout, shm_threshold_bytes=args.shm_threshold_bytes, ) thinker_sampling_params = SamplingParams( temperature=0.0, # Deterministic - no randomness top_p=1.0, # Disable nucleus sampling top_k=-1, # Disable top-k sampling max_tokens=2048, seed=SEED, # Fixed seed for sampling detokenize=True, repetition_penalty=1.1, ) talker_sampling_params = SamplingParams( temperature=0.9, top_p=0.8, top_k=40, max_tokens=2048, seed=SEED, # Fixed seed for sampling detokenize=True, repetition_penalty=1.05, stop_token_ids=[8294], ) code2wav_sampling_params = SamplingParams( temperature=0.0, # Deterministic - no randomness top_p=1.0, # Disable nucleus sampling top_k=-1, # Disable top-k sampling max_tokens=2048, seed=SEED, # Fixed seed for sampling detokenize=True, repetition_penalty=1.1, ) sampling_params_list = [ thinker_sampling_params, talker_sampling_params, code2wav_sampling_params, ] if args.txt_prompts is None: prompts = [query_result.inputs for _ in range(args.num_prompts)] else: assert args.query_type == "text", "txt-prompts is only supported for text query type" with open(args.txt_prompts, encoding="utf-8") as f: lines = [ln.strip() for ln in f.readlines()] prompts = [get_text_query(ln).inputs for ln in lines if ln != ""] print(f"[Info] Loaded {len(prompts)} prompts from {args.txt_prompts}") omni_outputs = omni_llm.generate(prompts, sampling_params_list) # Determine output directory: prefer --output-dir; fallback to --output-wav output_dir = args.output_dir if getattr(args, "output_dir", None) else args.output_wav os.makedirs(output_dir, exist_ok=True) for stage_outputs in omni_outputs: if stage_outputs.final_output_type == "text": for output in stage_outputs.request_output: request_id = int(output.request_id) text_output = output.outputs[0].text # Save aligned text file per request prompt_text = prompts[request_id]["prompt"] out_txt = os.path.join(output_dir, f"{request_id:05d}.txt") lines = [] lines.append("Prompt:\n") lines.append(str(prompt_text) + "\n") lines.append("vllm_text_output:\n") lines.append(str(text_output).strip() + "\n") try: with open(out_txt, "w", encoding="utf-8") as f: f.writelines(lines) except Exception as e: print(f"[Warn] Failed writing text file {out_txt}: {e}") print(f"Request ID: {request_id}, Text saved to {out_txt}") elif stage_outputs.final_output_type == "audio": for output in stage_outputs.request_output: request_id = int(output.request_id) audio_tensor = output.multimodal_output["audio"] output_wav = os.path.join(output_dir, f"output_{output.request_id}.wav") sf.write(output_wav, audio_tensor.detach().cpu().numpy(), samplerate=24000) print(f"Request ID: {request_id}, Saved audio to {output_wav}") def parse_args(): parser = FlexibleArgumentParser(description="Demo on using vLLM for offline inference with audio language models") parser.add_argument( "--query-type", "-q", type=str, default="mixed_modalities", choices=query_map.keys(), help="Query type.", ) parser.add_argument( "--enable-stats", action="store_true", default=False, help="Enable writing detailed statistics (default: disabled)", ) parser.add_argument( "--init-sleep-seconds", type=int, default=20, help="Sleep seconds after starting each stage process to allow initialization (default: 20)", ) parser.add_argument( "--batch-timeout", type=int, default=5, help="Timeout for batching in seconds (default: 5)", ) parser.add_argument( "--init-timeout", type=int, default=300, help="Timeout for initializing stages in seconds (default: 300)", ) parser.add_argument( "--shm-threshold-bytes", type=int, default=65536, help="Threshold for using shared memory in bytes (default: 65536)", ) parser.add_argument( "--output-wav", default="output_audio", help="[Deprecated] Output wav directory (use --output-dir).", ) parser.add_argument( "--num-prompts", type=int, default=1, help="Number of prompts to generate.", ) parser.add_argument( "--txt-prompts", type=str, default=None, help="Path to a .txt file with one prompt per line (preferred).", ) return parser.parse_args() if __name__ == "__main__": args = parse_args() main(args)