@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 195 lines 5.75 KB
Raw
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 Qwen2.5-Omni (thinker only).
5 """
6
7 from typing import NamedTuple
8
9 from vllm import LLM, SamplingParams
10 from vllm.assets.audio import AudioAsset
11 from vllm.assets.image import ImageAsset
12 from vllm.assets.video import VideoAsset
13 from vllm.multimodal.image import convert_image_mode
14 from vllm.utils.argparse_utils import FlexibleArgumentParser
15
16
17 class QueryResult(NamedTuple):
18 inputs: dict
19 limit_mm_per_prompt: dict[str, int]
20
21
22 # NOTE: The default `max_num_seqs` and `max_model_len` may result in OOM on
23 # lower-end GPUs.
24 # Unless specified, these settings have been tested to work on a single L4.
25
26 default_system = (
27 "You are Qwen, a virtual human developed by the Qwen Team, Alibaba "
28 "Group, capable of perceiving auditory and visual inputs, as well as "
29 "generating text and speech."
30 )
31
32
33 def get_mixed_modalities_query() -> QueryResult:
34 question = (
35 "What is recited in the audio? "
36 "What is the content of this image? Why is this video funny?"
37 )
38 prompt = (
39 f"<|im_start|>system\n{default_system}<|im_end|>\n"
40 "<|im_start|>user\n<|audio_bos|><|AUDIO|><|audio_eos|>"
41 "<|vision_bos|><|IMAGE|><|vision_eos|>"
42 "<|vision_bos|><|VIDEO|><|vision_eos|>"
43 f"{question}<|im_end|>\n"
44 f"<|im_start|>assistant\n"
45 )
46 return QueryResult(
47 inputs={
48 "prompt": prompt,
49 "multi_modal_data": {
50 "audio": AudioAsset("mary_had_lamb").audio_and_sample_rate,
51 "image": convert_image_mode(
52 ImageAsset("cherry_blossom").pil_image, "RGB"
53 ),
54 "video": VideoAsset(name="baby_reading", num_frames=16).np_ndarrays,
55 },
56 },
57 limit_mm_per_prompt={"audio": 1, "image": 1, "video": 1},
58 )
59
60
61 def get_use_audio_in_video_query() -> QueryResult:
62 question = (
63 "Describe the content of the video, then convert what the baby say into text."
64 )
65 prompt = (
66 f"<|im_start|>system\n{default_system}<|im_end|>\n"
67 "<|im_start|>user\n<|vision_bos|><|VIDEO|><|vision_eos|>"
68 f"{question}<|im_end|>\n"
69 f"<|im_start|>assistant\n"
70 )
71 asset = VideoAsset(name="baby_reading", num_frames=16)
72 audio = asset.get_audio(sampling_rate=16000)
73
74 return QueryResult(
75 inputs={
76 "prompt": prompt,
77 "multi_modal_data": {
78 "video": asset.np_ndarrays,
79 "audio": audio,
80 },
81 "mm_processor_kwargs": {
82 "use_audio_in_video": True,
83 },
84 },
85 limit_mm_per_prompt={"audio": 1, "video": 1},
86 )
87
88
89 def get_multi_audios_query() -> QueryResult:
90 question = "Are these two audio clips the same?"
91 prompt = (
92 f"<|im_start|>system\n{default_system}<|im_end|>\n"
93 "<|im_start|>user\n<|audio_bos|><|AUDIO|><|audio_eos|>"
94 "<|audio_bos|><|AUDIO|><|audio_eos|>"
95 f"{question}<|im_end|>\n"
96 f"<|im_start|>assistant\n"
97 )
98 return QueryResult(
99 inputs={
100 "prompt": prompt,
101 "multi_modal_data": {
102 "audio": [
103 AudioAsset("winning_call").audio_and_sample_rate,
104 AudioAsset("mary_had_lamb").audio_and_sample_rate,
105 ],
106 },
107 },
108 limit_mm_per_prompt={
109 "audio": 2,
110 },
111 )
112
113
114 def get_multi_images_query() -> QueryResult:
115 question = "What are the differences between these two images?"
116 prompt = (
117 f"<|im_start|>system\n{default_system}<|im_end|>\n"
118 "<|im_start|>user\n<|vision_bos|><|IMAGE|><|vision_eos|>"
119 "<|vision_bos|><|IMAGE|><|vision_eos|>"
120 f"{question}<|im_end|>\n"
121 f"<|im_start|>assistant\n"
122 )
123 return QueryResult(
124 inputs={
125 "prompt": prompt,
126 "multi_modal_data": {
127 "image": [
128 convert_image_mode(ImageAsset("cherry_blossom").pil_image, "RGB"),
129 convert_image_mode(ImageAsset("stop_sign").pil_image, "RGB"),
130 ],
131 },
132 },
133 limit_mm_per_prompt={
134 "image": 2,
135 },
136 )
137
138
139 query_map = {
140 "mixed_modalities": get_mixed_modalities_query,
141 "use_audio_in_video": get_use_audio_in_video_query,
142 "multi_audios": get_multi_audios_query,
143 "multi_images": get_multi_images_query,
144 }
145
146
147 def main(args):
148 model_name = "Qwen/Qwen2.5-Omni-7B"
149 query_result = query_map[args.query_type]()
150
151 llm = LLM(
152 model=model_name,
153 max_model_len=5632,
154 max_num_seqs=5,
155 limit_mm_per_prompt=query_result.limit_mm_per_prompt,
156 seed=args.seed,
157 )
158
159 # We set temperature to 0.2 so that outputs can be different
160 # even when all prompts are identical when running batch inference.
161 sampling_params = SamplingParams(temperature=0.2, max_tokens=64)
162
163 outputs = llm.generate(query_result.inputs, sampling_params=sampling_params)
164
165 for o in outputs:
166 generated_text = o.outputs[0].text
167 print(generated_text)
168
169
170 def parse_args():
171 parser = FlexibleArgumentParser(
172 description="Demo on using vLLM for offline inference with "
173 "audio language models"
174 )
175 parser.add_argument(
176 "--query-type",
177 "-q",
178 type=str,
179 default="mixed_modalities",
180 choices=query_map.keys(),
181 help="Query type.",
182 )
183 parser.add_argument(
184 "--seed",
185 type=int,
186 default=0,
187 help="Set the seed when initializing `vllm.LLM`.",
188 )
189
190 return parser.parse_args()
191
192
193 if __name__ == "__main__":
194 args = parse_args()
195 main(args)