293 lines
9.6 KiB
Python
293 lines
9.6 KiB
Python
# 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)
|