|
| 1 | +# SPDX-License-Identifier: Apache-2.0 |
| 2 | +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | +""" |
| 4 | +Pytest configuration and fixtures for vllm-omni tests. |
| 5 | +""" |
| 6 | + |
| 7 | +from typing import Any |
| 8 | + |
| 9 | +import pytest |
| 10 | +from vllm.distributed.parallel_state import cleanup_dist_env_and_memory |
| 11 | +from vllm.sampling_params import SamplingParams |
| 12 | + |
| 13 | +from vllm_omni.entrypoints.omni import Omni |
| 14 | + |
| 15 | +PromptAudioInput = list[tuple[Any, int]] | tuple[Any, int] | None |
| 16 | +PromptImageInput = list[Any] | Any | None |
| 17 | +PromptVideoInput = list[Any] | Any | None |
| 18 | + |
| 19 | + |
| 20 | +class OmniRunner: |
| 21 | + """ |
| 22 | + Test runner for Omni models. |
| 23 | + """ |
| 24 | + |
| 25 | + def __init__( |
| 26 | + self, |
| 27 | + model_name: str, |
| 28 | + seed: int = 42, |
| 29 | + init_sleep_seconds: int = 20, |
| 30 | + batch_timeout: int = 10, |
| 31 | + init_timeout: int = 300, |
| 32 | + shm_threshold_bytes: int = 65536, |
| 33 | + log_stats: bool = False, |
| 34 | + stage_configs_path: str | None = None, |
| 35 | + **kwargs, |
| 36 | + ) -> None: |
| 37 | + """ |
| 38 | + Initialize an OmniRunner for testing. |
| 39 | +
|
| 40 | + Args: |
| 41 | + model_name: The model name or path |
| 42 | + seed: Random seed for reproducibility |
| 43 | + init_sleep_seconds: Sleep time after starting each stage |
| 44 | + batch_timeout: Timeout for batching in seconds |
| 45 | + init_timeout: Timeout for initializing stages in seconds |
| 46 | + shm_threshold_bytes: Threshold for using shared memory |
| 47 | + log_stats: Enable detailed statistics logging |
| 48 | + stage_configs_path: Optional path to YAML stage config file |
| 49 | + **kwargs: Additional arguments passed to Omni |
| 50 | + """ |
| 51 | + self.model_name = model_name |
| 52 | + self.seed = seed |
| 53 | + |
| 54 | + self.omni = Omni( |
| 55 | + model=model_name, |
| 56 | + log_stats=log_stats, |
| 57 | + init_sleep_seconds=init_sleep_seconds, |
| 58 | + batch_timeout=batch_timeout, |
| 59 | + init_timeout=init_timeout, |
| 60 | + shm_threshold_bytes=shm_threshold_bytes, |
| 61 | + stage_configs_path=stage_configs_path, |
| 62 | + **kwargs, |
| 63 | + ) |
| 64 | + |
| 65 | + def get_default_sampling_params_list(self) -> list[SamplingParams]: |
| 66 | + """ |
| 67 | + Get a list of default sampling parameters for all stages. |
| 68 | +
|
| 69 | + Returns: |
| 70 | + List of SamplingParams with default decoding for each stage |
| 71 | + """ |
| 72 | + return [st.default_sampling_params for st in self.omni.instance.stage_list] |
| 73 | + |
| 74 | + def get_omni_inputs( |
| 75 | + self, |
| 76 | + prompts: list[str] | str, |
| 77 | + system_prompt: str | None = None, |
| 78 | + audios: PromptAudioInput = None, |
| 79 | + images: PromptImageInput = None, |
| 80 | + videos: PromptVideoInput = None, |
| 81 | + mm_processor_kwargs: dict[str, Any] | None = None, |
| 82 | + ) -> list[dict[str, Any]]: |
| 83 | + """ |
| 84 | + Construct Omni input format from prompts and multimodal data. |
| 85 | +
|
| 86 | + Args: |
| 87 | + prompts: Text prompt(s) - either a single string or list of strings |
| 88 | + system_prompt: Optional system prompt (defaults to Qwen system prompt) |
| 89 | + audios: Audio input(s) - tuple of (audio_array, sample_rate) or list of tuples |
| 90 | + images: Image input(s) - PIL Image or list of PIL Images |
| 91 | + videos: Video input(s) - numpy array or list of numpy arrays |
| 92 | + mm_processor_kwargs: Optional processor kwargs (e.g., use_audio_in_video) |
| 93 | +
|
| 94 | + Returns: |
| 95 | + List of prompt dictionaries suitable for Omni.generate() |
| 96 | + """ |
| 97 | + if system_prompt is None: |
| 98 | + system_prompt = ( |
| 99 | + "You are Qwen, a virtual human developed by the Qwen Team, Alibaba " |
| 100 | + "Group, capable of perceiving auditory and visual inputs, as well as " |
| 101 | + "generating text and speech." |
| 102 | + ) |
| 103 | + |
| 104 | + if isinstance(prompts, str): |
| 105 | + prompts = [prompts] |
| 106 | + |
| 107 | + def _normalize_mm_input(mm_input, num_prompts): |
| 108 | + if mm_input is None: |
| 109 | + return [None] * num_prompts |
| 110 | + if isinstance(mm_input, list): |
| 111 | + if len(mm_input) != num_prompts: |
| 112 | + raise ValueError( |
| 113 | + f"Multimodal input list length ({len(mm_input)}) must match prompts length ({num_prompts})" |
| 114 | + ) |
| 115 | + return mm_input |
| 116 | + return [mm_input] * num_prompts |
| 117 | + |
| 118 | + num_prompts = len(prompts) |
| 119 | + audios_list = _normalize_mm_input(audios, num_prompts) |
| 120 | + images_list = _normalize_mm_input(images, num_prompts) |
| 121 | + videos_list = _normalize_mm_input(videos, num_prompts) |
| 122 | + |
| 123 | + omni_inputs = [] |
| 124 | + for i, prompt_text in enumerate(prompts): |
| 125 | + user_content = "" |
| 126 | + multi_modal_data = {} |
| 127 | + |
| 128 | + audio = audios_list[i] |
| 129 | + if audio is not None: |
| 130 | + if isinstance(audio, list): |
| 131 | + for _ in audio: |
| 132 | + user_content += "<|audio_bos|><|AUDIO|><|audio_eos|>" |
| 133 | + multi_modal_data["audio"] = audio |
| 134 | + else: |
| 135 | + user_content += "<|audio_bos|><|AUDIO|><|audio_eos|>" |
| 136 | + multi_modal_data["audio"] = audio |
| 137 | + |
| 138 | + image = images_list[i] |
| 139 | + if image is not None: |
| 140 | + if isinstance(image, list): |
| 141 | + for _ in image: |
| 142 | + user_content += "<|vision_bos|><|IMAGE|><|vision_eos|>" |
| 143 | + multi_modal_data["image"] = image |
| 144 | + else: |
| 145 | + user_content += "<|vision_bos|><|IMAGE|><|vision_eos|>" |
| 146 | + multi_modal_data["image"] = image |
| 147 | + |
| 148 | + video = videos_list[i] |
| 149 | + if video is not None: |
| 150 | + if isinstance(video, list): |
| 151 | + for _ in video: |
| 152 | + user_content += "<|vision_bos|><|VIDEO|><|vision_eos|>" |
| 153 | + multi_modal_data["video"] = video |
| 154 | + else: |
| 155 | + user_content += "<|vision_bos|><|VIDEO|><|vision_eos|>" |
| 156 | + multi_modal_data["video"] = video |
| 157 | + |
| 158 | + user_content += prompt_text |
| 159 | + |
| 160 | + full_prompt = ( |
| 161 | + f"<|im_start|>system\n{system_prompt}<|im_end|>\n" |
| 162 | + f"<|im_start|>user\n{user_content}<|im_end|>\n" |
| 163 | + f"<|im_start|>assistant\n" |
| 164 | + ) |
| 165 | + |
| 166 | + input_dict: dict[str, Any] = {"prompt": full_prompt} |
| 167 | + if multi_modal_data: |
| 168 | + input_dict["multi_modal_data"] = multi_modal_data |
| 169 | + if mm_processor_kwargs: |
| 170 | + input_dict["mm_processor_kwargs"] = mm_processor_kwargs |
| 171 | + |
| 172 | + omni_inputs.append(input_dict) |
| 173 | + |
| 174 | + return omni_inputs |
| 175 | + |
| 176 | + def generate( |
| 177 | + self, |
| 178 | + prompts: list[dict[str, Any]], |
| 179 | + sampling_params_list: list[SamplingParams] | None = None, |
| 180 | + ) -> list[Any]: |
| 181 | + """ |
| 182 | + Generate outputs for the given prompts. |
| 183 | +
|
| 184 | + Args: |
| 185 | + prompts: List of prompt dictionaries with 'prompt' and optionally |
| 186 | + 'multi_modal_data' keys |
| 187 | + sampling_params_list: List of sampling parameters for each stage. |
| 188 | + If None, uses default parameters. |
| 189 | +
|
| 190 | + Returns: |
| 191 | + List of OmniRequestOutput objects from stages with final_output=True |
| 192 | + """ |
| 193 | + if sampling_params_list is None: |
| 194 | + sampling_params_list = self.get_default_sampling_params_list() |
| 195 | + |
| 196 | + return self.omni.generate(prompts, sampling_params_list) |
| 197 | + |
| 198 | + def generate_multimodal( |
| 199 | + self, |
| 200 | + prompts: list[str] | str, |
| 201 | + sampling_params_list: list[SamplingParams] | None = None, |
| 202 | + system_prompt: str | None = None, |
| 203 | + audios: PromptAudioInput = None, |
| 204 | + images: PromptImageInput = None, |
| 205 | + videos: PromptVideoInput = None, |
| 206 | + mm_processor_kwargs: dict[str, Any] | None = None, |
| 207 | + ) -> list[Any]: |
| 208 | + """ |
| 209 | + Convenience method to generate with multimodal inputs. |
| 210 | +
|
| 211 | + Args: |
| 212 | + prompts: Text prompt(s) |
| 213 | + sampling_params_list: List of sampling parameters for each stage |
| 214 | + system_prompt: Optional system prompt |
| 215 | + audios: Audio input(s) |
| 216 | + images: Image input(s) |
| 217 | + videos: Video input(s) |
| 218 | + mm_processor_kwargs: Optional processor kwargs |
| 219 | +
|
| 220 | + Returns: |
| 221 | + List of OmniRequestOutput objects from stages with final_output=True |
| 222 | + """ |
| 223 | + omni_inputs = self.get_omni_inputs( |
| 224 | + prompts=prompts, |
| 225 | + system_prompt=system_prompt, |
| 226 | + audios=audios, |
| 227 | + images=images, |
| 228 | + videos=videos, |
| 229 | + mm_processor_kwargs=mm_processor_kwargs, |
| 230 | + ) |
| 231 | + return self.generate(omni_inputs, sampling_params_list) |
| 232 | + |
| 233 | + def __enter__(self): |
| 234 | + """Context manager entry.""" |
| 235 | + return self |
| 236 | + |
| 237 | + def __exit__(self, exc_type, exc_val, exc_tb): |
| 238 | + """Context manager exit - cleanup resources.""" |
| 239 | + self.close() |
| 240 | + del self.omni |
| 241 | + cleanup_dist_env_and_memory() |
| 242 | + |
| 243 | + def close(self): |
| 244 | + """Close and cleanup the Omni instance.""" |
| 245 | + if hasattr(self.omni.instance, "close"): |
| 246 | + self.omni.instance.close() |
| 247 | + |
| 248 | + |
| 249 | +@pytest.fixture(scope="session") |
| 250 | +def omni_runner(): |
| 251 | + return OmniRunner |
0 commit comments