293 lines
9.1 KiB
Python
293 lines
9.1 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""
|
|
This example shows how to use vLLM for running offline inference
|
|
with the correct prompt format on Qwen3-Omni (thinker only).
|
|
"""
|
|
|
|
import os
|
|
from typing import NamedTuple
|
|
|
|
import soundfile as sf
|
|
from vllm import SamplingParams
|
|
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.utils import FlexibleArgumentParser
|
|
|
|
from vllm_omni.entrypoints.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_video_query(question: str = None) -> QueryResult:
|
|
if question is None:
|
|
question = "Why is this video funny?"
|
|
prompt = (
|
|
f"<|im_start|>system\n{default_system}<|im_end|>\n"
|
|
"<|im_start|>user\n<|vision_start|><|video_pad|><|vision_end|>"
|
|
f"{question}<|im_end|>\n"
|
|
f"<|im_start|>assistant\n"
|
|
)
|
|
|
|
return QueryResult(
|
|
inputs={
|
|
"prompt": prompt,
|
|
"multi_modal_data": {
|
|
"video": VideoAsset(name="baby_reading", num_frames=16).np_ndarrays,
|
|
},
|
|
},
|
|
limit_mm_per_prompt={"video": 1},
|
|
)
|
|
|
|
|
|
def get_image_query(question: str = None) -> QueryResult:
|
|
if question is None:
|
|
question = "What is the content of this image?"
|
|
prompt = (
|
|
f"<|im_start|>system\n{default_system}<|im_end|>\n"
|
|
"<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>"
|
|
f"{question}<|im_end|>\n"
|
|
f"<|im_start|>assistant\n"
|
|
)
|
|
return QueryResult(
|
|
inputs={
|
|
"prompt": prompt,
|
|
"multi_modal_data": {
|
|
"image": convert_image_mode(ImageAsset("cherry_blossom").pil_image, "RGB"),
|
|
},
|
|
},
|
|
limit_mm_per_prompt={"image": 1},
|
|
)
|
|
|
|
|
|
def get_audio_query(question: str = None) -> QueryResult:
|
|
if question is None:
|
|
question = "What is the content of this audio?"
|
|
prompt = (
|
|
f"<|im_start|>system\n{default_system}<|im_end|>\n"
|
|
"<|im_start|>user\n<|audio_start|><|audio_pad|><|audio_end|>"
|
|
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,
|
|
},
|
|
},
|
|
limit_mm_per_prompt={"audio": 1},
|
|
)
|
|
|
|
|
|
query_map = {
|
|
"text": get_text_query,
|
|
"use_audio": get_audio_query,
|
|
"use_image": get_image_query,
|
|
"use_video": get_video_query,
|
|
}
|
|
|
|
|
|
def main(args):
|
|
model_name = "Qwen/Qwen3-Omni-30B-A3B-Instruct"
|
|
query_result = query_map[args.query_type]()
|
|
|
|
omni_llm = Omni(
|
|
model=model_name,
|
|
stage_configs_path=args.stage_configs_path,
|
|
)
|
|
|
|
thinker_sampling_params = SamplingParams(
|
|
temperature=0.4,
|
|
top_p=0.9,
|
|
top_k=-1,
|
|
max_tokens=1200,
|
|
repetition_penalty=1.05,
|
|
logit_bias={},
|
|
seed=SEED,
|
|
)
|
|
|
|
talker_sampling_params = SamplingParams(
|
|
temperature=0.9,
|
|
top_k=50,
|
|
max_tokens=4096,
|
|
seed=SEED,
|
|
detokenize=False,
|
|
repetition_penalty=1.05,
|
|
stop_token_ids=[2150], # TALKER_CODEC_EOS_TOKEN_ID
|
|
)
|
|
|
|
# Sampling parameters for Code2Wav stage (audio generation)
|
|
code2wav_sampling_params = SamplingParams(
|
|
temperature=0.0,
|
|
top_p=1.0,
|
|
top_k=-1,
|
|
max_tokens=4096 * 16,
|
|
seed=SEED,
|
|
detokenize=True,
|
|
repetition_penalty=1.1,
|
|
)
|
|
|
|
sampling_params_list = [
|
|
thinker_sampling_params,
|
|
talker_sampling_params, # code predictor is integrated into talker for Qwen3 Omni
|
|
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")
|
|
|
|
# Convert to numpy array and ensure correct format
|
|
audio_numpy = audio_tensor.float().detach().cpu().numpy()
|
|
|
|
# Ensure audio is 1D (flatten if needed)
|
|
if audio_numpy.ndim > 1:
|
|
audio_numpy = audio_numpy.flatten()
|
|
|
|
# Save audio file with explicit WAV format
|
|
sf.write(output_wav, audio_numpy, samplerate=24000, format="WAV")
|
|
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).",
|
|
)
|
|
parser.add_argument(
|
|
"--stage-configs-path",
|
|
type=str,
|
|
default=None,
|
|
help="Path to a stage configs file.",
|
|
)
|
|
|
|
return parser.parse_args()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
args = parse_args()
|
|
main(args)
|