From 7a34298743049cfa06a280865ef5fe0dc8b2e949 Mon Sep 17 00:00:00 2001 From: Divin Honnappa Date: Tue, 21 Apr 2026 19:05:33 -0500 Subject: [PATCH 01/45] Remove dead code (#2998) Signed-off-by: Divin Honnappa Co-authored-by: SYLAR <125541396+lishunyang12@users.noreply.github.com> --- .../scripts/hardware_ci/run-amd-test.sh | 21 ------------------- 1 file changed, 21 deletions(-) diff --git a/.buildkite/scripts/hardware_ci/run-amd-test.sh b/.buildkite/scripts/hardware_ci/run-amd-test.sh index 2707dfefe..96c139c8f 100755 --- a/.buildkite/scripts/hardware_ci/run-amd-test.sh +++ b/.buildkite/scripts/hardware_ci/run-amd-test.sh @@ -11,15 +11,6 @@ set -o pipefail export PYTHONPATH=".." # Print ROCm version -echo "--- Confirming Clean Initial State" -while true; do - sleep 3 - if grep -q clean /opt/amdgpu/etc/gpu_state; then - echo "GPUs state is \"clean\"" - break - fi -done - echo "--- ROCm info" rocminfo @@ -51,18 +42,6 @@ cleanup_docker() { # Call the cleanup docker function cleanup_docker -echo "--- Resetting GPUs" - -echo "reset" > /opt/amdgpu/etc/gpu_state - -while true; do - sleep 3 - if grep -q clean /opt/amdgpu/etc/gpu_state; then - echo "GPUs state is \"clean\"" - break - fi -done - echo "--- Pulling container" ## Temporary change to use AMD Docker Hub to store the vllm-omni image # to bypass the rate limit issue with ECR Public Gallery. From e6cf26efb707a9d296afff51843d3182c7f33421 Mon Sep 17 00:00:00 2001 From: ChipMates <156656427+ChipMates@users.noreply.github.com> Date: Wed, 22 Apr 2026 05:04:54 +0200 Subject: [PATCH 02/45] [Qwen3TTS][Bugfix] Guard inner CUDA graph replay during outer capture (#2910) Signed-off-by: ChipMates --- .../models/qwen3_tts/cuda_graph_decoder_wrapper.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/vllm_omni/model_executor/models/qwen3_tts/cuda_graph_decoder_wrapper.py b/vllm_omni/model_executor/models/qwen3_tts/cuda_graph_decoder_wrapper.py index 8f7eeb542..26d69fe22 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/cuda_graph_decoder_wrapper.py +++ b/vllm_omni/model_executor/models/qwen3_tts/cuda_graph_decoder_wrapper.py @@ -141,6 +141,15 @@ class CUDAGraphDecoderWrapper: if not self.enabled or not self._warmed_up or codes.shape[0] != 1: return self.decoder(codes) + # Inner CUDA graph replay is illegal while an outer stream capture is + # active (e.g. vLLM's cudagraph_mode=FULL warmup on Stage 1). Fall back + # to eager in that case so the outer capture can complete. The guard is + # a no-op at runtime: is_current_stream_capturing() returns False + # outside the startup capture window, so normal inference still hits + # the graph fast path. + if torch.cuda.is_current_stream_capturing(): + return self.decoder(codes) + actual_size = codes.shape[-1] padded_size = self._get_padded_size(actual_size) From 2903f528d2b73387eb6a065d2159b6ff950a6c73 Mon Sep 17 00:00:00 2001 From: Greenhand Tsang <74168957+greenhandzpx@users.noreply.github.com> Date: Wed, 22 Apr 2026 11:35:39 +0800 Subject: [PATCH 03/45] [Bugfix][Model] Fix Qwen3-TTS Code2Wav max_model_len validation (#2508) Signed-off-by: greenhandzpx Co-authored-by: greenhandzpx --- tests/engine/test_arg_utils.py | 35 ++++++++++++++++++++++++++++++++++ vllm_omni/engine/arg_utils.py | 7 +++++++ 2 files changed, 42 insertions(+) diff --git a/tests/engine/test_arg_utils.py b/tests/engine/test_arg_utils.py index 4d69f24c5..0d61f6a67 100644 --- a/tests/engine/test_arg_utils.py +++ b/tests/engine/test_arg_utils.py @@ -139,6 +139,41 @@ def test_from_cli_args_picks_up_stage_configs_path(): assert args.custom_pipeline_args is None +def test_qwen3_tts_code2wav_injects_max_position_embeddings(monkeypatch): + """Ensure Code2Wav mirrors stage max_model_len into nested HF overrides. + + Qwen3-TTS Code2Wav is a pure decoder stage whose runtime max_model_len can + legitimately exceed the base checkpoint's default text max length. Recent + vLLM validates these values during ModelConfig creation, so we inject + ``talker_config.max_position_embeddings`` before delegating to vLLM. + """ + captured: dict[str, object] = {} + baseline_config = Mock() + + def fake_create_model_config(self): + captured["hf_overrides"] = self.hf_overrides + return baseline_config + + monkeypatch.setattr(EngineArgs, "create_model_config", fake_create_model_config) + monkeypatch.setattr( + OmniModelConfig, + "from_vllm_model_config", + classmethod(lambda cls, model_config, **omni_kwargs: model_config), + ) + + OmniEngineArgs( + model_arch="Qwen3TTSCode2Wav", + max_model_len=65536, + ).create_model_config() + + assert captured["hf_overrides"] == { + "architectures": ["Qwen3TTSCode2Wav"], + "talker_config": { + "max_position_embeddings": 65536, + }, + } + + def test_stage_specific_text_config_override(): """Ensure dependent attributes are updated when using stage-specific config.""" vllm_config = EngineArgs().create_model_config() diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index c93bd32c2..a2fa3295e 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -246,6 +246,13 @@ class OmniEngineArgs(EngineArgs): if model_type is not None: self.hf_overrides.setdefault("model_type", model_type) + # Stage wrappers (e.g. Code2Wav) may need max_model_len larger + # than the base checkpoint's text max_position_embeddings. + if self.model_arch == "Qwen3TTSCode2Wav" and self.max_model_len is not None: + self.hf_overrides.setdefault("talker_config", {}).setdefault( + "max_position_embeddings", int(self.max_model_len) + ) + # For models whose HF config.json is empty or lacks model_type # (e.g. CosyVoice3), AutoConfig.from_pretrained fails because it # cannot determine which config class to use from the empty dict. From 8e1bbc9ab58816fd80e7722e29189b9f294cee66 Mon Sep 17 00:00:00 2001 From: Chen-Yo Sun Date: Tue, 21 Apr 2026 20:49:50 -0700 Subject: [PATCH 04/45] [Feature][Voxtral TTS] Support per-model extra sampling params & add cfg_alpha for Voxtral TTS (#2338) Signed-off-by: Chen-Yo Sun Signed-off-by: Yueqian Lin Co-authored-by: Yueqian Lin --- .../offline_inference/voxtral_tts/end2end.py | 11 ++++++ .../online_serving/voxtral_tts/gradio_demo.py | 27 +++++++++++-- .../test_cuda_graph_acoustic_transformer.py | 32 ++++++++++------ vllm_omni/config/model.py | 1 + vllm_omni/deploy/voxtral_tts.yaml | 2 + vllm_omni/engine/arg_utils.py | 2 + vllm_omni/engine/stage_init_utils.py | 4 ++ .../entrypoints/openai/protocol/audio.py | 6 ++- .../entrypoints/openai/serving_speech.py | 18 ++++++++- ...cuda_graph_acoustic_transformer_wrapper.py | 38 ++++++++++++------- .../models/voxtral_tts/voxtral_tts.py | 33 +++++++++++++++- .../voxtral_tts_audio_generation.py | 14 +++++-- vllm_omni/worker/gpu_model_runner.py | 9 +++++ 13 files changed, 160 insertions(+), 37 deletions(-) diff --git a/examples/offline_inference/voxtral_tts/end2end.py b/examples/offline_inference/voxtral_tts/end2end.py index cf28f917e..0a6f88715 100644 --- a/examples/offline_inference/voxtral_tts/end2end.py +++ b/examples/offline_inference/voxtral_tts/end2end.py @@ -298,6 +298,12 @@ def parse_args() -> Namespace: default=None, help="Voice to use instead of audio file.", ) + parser.add_argument( + "--cfg-alpha", + type=float, + default=None, + help="CFG alpha for flow-matching guidance (default: use value from stage config, typically 1.2).", + ) return parser.parse_args() @@ -349,8 +355,13 @@ def main(args: Any) -> None: inputs = compose_request(model_name, text_chunk, audio_prompt_file, args) + extra_args = {} + if args.cfg_alpha is not None: + extra_args["cfg_alpha"] = args.cfg_alpha + sampling_params = SamplingParams( max_tokens=max_num_tokens, + extra_args=extra_args if extra_args else None, ) sampling_params_list = [ sampling_params, diff --git a/examples/online_serving/voxtral_tts/gradio_demo.py b/examples/online_serving/voxtral_tts/gradio_demo.py index 35d6b590c..7905c6261 100644 --- a/examples/online_serving/voxtral_tts/gradio_demo.py +++ b/examples/online_serving/voxtral_tts/gradio_demo.py @@ -216,6 +216,7 @@ def make_update_voice_dropdown(language_voices: dict[str, list[str]]): def run_inference( voice_name: str, text_prompt: str, + cfg_alpha: float, base_url: str, model: str, ) -> tuple[int, np.ndarray]: @@ -233,6 +234,7 @@ def run_inference( "model": model, "response_format": "wav", "voice": voice_name, + "extra_params": {"cfg_alpha": cfg_alpha}, } response = httpx.post( @@ -377,6 +379,14 @@ def main( placeholder="Enter the text you want to synthesize...", lines=4, ) + cfg_alpha_slider = gr.Slider( + minimum=1.0, + maximum=2.0, + step=0.1, + value=1.2, + label="CFG Alpha", + info="Flow-matching guidance strength (default: 1.2)", + ) with gr.Row(): reset_btn = gr.Button("Clear") submit_btn = gr.Button("Generate audio", interactive=False) @@ -415,9 +425,9 @@ def main( ) # --- Wiring inference + persistence to the button --- - def _on_submit(voice: str, text: str): + def _on_submit(voice: str, text: str, cfg_alpha: float): assert text.strip() != "" - sr, audio_array = run_inference(voice, text, base_url, model) + sr, audio_array = run_inference(voice, text, cfg_alpha, base_url, model) if outputs_dir is not None: share_id, saved_audio_path = _save_example( outputs_dir, @@ -432,7 +442,7 @@ def main( submit_btn.click( fn=_on_submit, - inputs=[voice_name, text_prompt], + inputs=[voice_name, text_prompt, cfg_alpha_slider], outputs=[output_audio, share_link_box], ) @@ -446,6 +456,7 @@ def main( language, # language_dropdown voice, # voice_name "", # text_prompt + 1.2, # cfg_alpha_slider None, # output_audio gr.update(interactive=False), # submit_btn "", # share_link_box @@ -456,7 +467,15 @@ def main( reset_btn.click( fn=make_on_reset(languages, language_voices), inputs=[], - outputs=[language_dropdown, voice_name, text_prompt, output_audio, submit_btn, share_link_box], + outputs=[ + language_dropdown, + voice_name, + text_prompt, + cfg_alpha_slider, + output_audio, + submit_btn, + share_link_box, + ], ) def make_load_from_share(outputs_dir: Path | None, languages: list[str], language_voices: dict[str, list[str]]): diff --git a/tests/model_executor/models/voxtral_tts/test_cuda_graph_acoustic_transformer.py b/tests/model_executor/models/voxtral_tts/test_cuda_graph_acoustic_transformer.py index 847adae06..c7b023361 100644 --- a/tests/model_executor/models/voxtral_tts/test_cuda_graph_acoustic_transformer.py +++ b/tests/model_executor/models/voxtral_tts/test_cuda_graph_acoustic_transformer.py @@ -137,7 +137,7 @@ class SyntheticModel(nn.Module): def compute_mm_logits( self, hidden_states: torch.Tensor, - mm_sampling_tensors=None, + cfg_alpha: torch.Tensor, ): """Eager fallback path: replicate what the wrapper does.""" at = self.acoustic_transformer @@ -216,6 +216,10 @@ def _random_hidden(batch_size, device=DEVICE, dtype=torch.bfloat16): return torch.randn(batch_size, HIDDEN_DIM, device=device, dtype=dtype) +def _cfg_alpha(batch_size, value=1.2, device=DEVICE): + return torch.full((batch_size,), value, device=device, dtype=torch.float32) + + def _unpack_audio_codes(result): """Unpack (fake_eos, {"audio": [list of tensors]}) into (fake_eos, audio_codes).""" fake_eos, mm_tokens = result @@ -235,7 +239,7 @@ def test_exact_size_output_format(model, wrapper, batch_size): """Graph path returns correctly shaped and bounded outputs.""" hidden = _random_hidden(batch_size) with torch.no_grad(): - graph_eos, graph_codes = _unpack_audio_codes(wrapper(hidden)) + graph_eos, graph_codes = _unpack_audio_codes(wrapper(hidden, cfg_alpha=_cfg_alpha(batch_size))) assert graph_eos.shape == (batch_size,) assert graph_codes.shape == (batch_size, 1 + N_ACOUSTIC_CODEBOOK) # fake_eos should be 0.0 or 1.0 @@ -248,11 +252,12 @@ def test_exact_size_output_format(model, wrapper, batch_size): def test_exact_size_deterministic(model, wrapper, batch_size): """Same input + same RNG state produces identical CUDA graph output.""" hidden = _random_hidden(batch_size) + cfg_alpha = _cfg_alpha(batch_size) with torch.no_grad(): torch.manual_seed(42) - eos1, codes1 = _unpack_audio_codes(wrapper(hidden)) + eos1, codes1 = _unpack_audio_codes(wrapper(hidden, cfg_alpha=cfg_alpha)) torch.manual_seed(42) - eos2, codes2 = _unpack_audio_codes(wrapper(hidden)) + eos2, codes2 = _unpack_audio_codes(wrapper(hidden, cfg_alpha=cfg_alpha)) torch.testing.assert_close(eos1, eos2, atol=0, rtol=0) torch.testing.assert_close(codes1, codes2, atol=0, rtol=0) @@ -267,7 +272,7 @@ def test_padded_output_shape(model, wrapper, batch_size): """Padded decode must return output trimmed to actual batch size.""" hidden = _random_hidden(batch_size) with torch.no_grad(): - graph_eos, graph_codes = _unpack_audio_codes(wrapper(hidden)) + graph_eos, graph_codes = _unpack_audio_codes(wrapper(hidden, cfg_alpha=_cfg_alpha(batch_size))) assert graph_eos.shape == (batch_size,) assert graph_codes.shape == (batch_size, 1 + N_ACOUSTIC_CODEBOOK) @@ -277,7 +282,7 @@ def test_padded_output_bounded(model, wrapper, batch_size): """Padded output audio codes should be non-negative integers.""" hidden = _random_hidden(batch_size) with torch.no_grad(): - graph_eos, graph_codes = _unpack_audio_codes(wrapper(hidden)) + graph_eos, graph_codes = _unpack_audio_codes(wrapper(hidden, cfg_alpha=_cfg_alpha(batch_size))) # fake_eos should be 0.0 or 1.0 assert torch.all((graph_eos == 0.0) | (graph_eos == 1.0)) # Audio codes should be non-negative @@ -293,11 +298,12 @@ def test_padded_output_bounded(model, wrapper, batch_size): def test_fallback_eager_exact_match(model, wrapper, batch_size): """Cudagraph fallback to eager. Two eager runs should produce identical results.""" hidden = _random_hidden(batch_size) + alpha = _cfg_alpha(batch_size) with torch.no_grad(): torch.manual_seed(100) - eager_eos, eager_codes = _unpack_audio_codes(model.compute_mm_logits(hidden)) + eager_eos, eager_codes = _unpack_audio_codes(model.compute_mm_logits(hidden, cfg_alpha=alpha)) torch.manual_seed(100) - graph_eos, graph_codes = _unpack_audio_codes(wrapper(hidden)) + graph_eos, graph_codes = _unpack_audio_codes(wrapper(hidden, cfg_alpha=alpha)) torch.testing.assert_close(graph_eos, eager_eos, atol=0, rtol=0) torch.testing.assert_close(graph_codes, eager_codes, atol=0, rtol=0) @@ -310,12 +316,13 @@ def test_fallback_eager_exact_match(model, wrapper, batch_size): def test_disabled_wrapper_matches_eager(model, wrapper): """Cudagraph fallback to eager. Two eager runs should produce identical results.""" hidden = _random_hidden(4) + alpha = _cfg_alpha(4) wrapper.enabled = False with torch.no_grad(): torch.manual_seed(200) - eager_eos, eager_codes = _unpack_audio_codes(model.compute_mm_logits(hidden)) + eager_eos, eager_codes = _unpack_audio_codes(model.compute_mm_logits(hidden, cfg_alpha=alpha)) torch.manual_seed(200) - graph_eos, graph_codes = _unpack_audio_codes(wrapper(hidden)) + graph_eos, graph_codes = _unpack_audio_codes(wrapper(hidden, cfg_alpha=alpha)) wrapper.enabled = True torch.testing.assert_close(graph_eos, eager_eos, atol=0, rtol=0) torch.testing.assert_close(graph_codes, eager_codes, atol=0, rtol=0) @@ -329,10 +336,11 @@ def test_disabled_wrapper_matches_eager(model, wrapper): def test_deterministic_across_calls(model, wrapper): """Same input + same RNG state. Two cudagraph runs should produce identical results.""" hidden = _random_hidden(4) + alpha = _cfg_alpha(4) with torch.no_grad(): torch.manual_seed(300) - eos1, codes1 = _unpack_audio_codes(wrapper(hidden)) + eos1, codes1 = _unpack_audio_codes(wrapper(hidden, cfg_alpha=alpha)) torch.manual_seed(300) - eos2, codes2 = _unpack_audio_codes(wrapper(hidden)) + eos2, codes2 = _unpack_audio_codes(wrapper(hidden, cfg_alpha=alpha)) torch.testing.assert_close(eos1, eos2, atol=0, rtol=0) torch.testing.assert_close(codes1, codes2, atol=0, rtol=0) diff --git a/vllm_omni/config/model.py b/vllm_omni/config/model.py index 96a34a8d7..af89184a6 100644 --- a/vllm_omni/config/model.py +++ b/vllm_omni/config/model.py @@ -114,6 +114,7 @@ class OmniModelConfig(ModelConfig): codec_frame_rate_hz: float | None = None task_type: str | None = None enable_sleep_mode: bool = False + has_sampling_extra_args: bool = False @property def registry(self): diff --git a/vllm_omni/deploy/voxtral_tts.yaml b/vllm_omni/deploy/voxtral_tts.yaml index b0899a399..87d999c67 100644 --- a/vllm_omni/deploy/voxtral_tts.yaml +++ b/vllm_omni/deploy/voxtral_tts.yaml @@ -36,6 +36,8 @@ stages: max_tokens: 2048 seed: 42 repetition_penalty: 1.1 + extra_args: + cfg_alpha: 1.2 tokenizer_mode: mistral config_format: mistral load_format: mistral diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index a2fa3295e..89139bf1b 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -170,6 +170,7 @@ class OmniEngineArgs(EngineArgs): output_modalities: list[str] | None = None log_stats: bool = False custom_pipeline_args: dict[str, Any] | None = None + has_sampling_extra_args: bool = False def __post_init__(self) -> None: if self.worker_cls is None: @@ -326,6 +327,7 @@ class OmniEngineArgs(EngineArgs): subtalker_sampling_params=self.subtalker_sampling_params, omni_kv_config=self.omni_kv_config, task_type=self.task_type, + has_sampling_extra_args=self.has_sampling_extra_args, ) return omni_config diff --git a/vllm_omni/engine/stage_init_utils.py b/vllm_omni/engine/stage_init_utils.py index cd6cd6a69..94b9faa80 100644 --- a/vllm_omni/engine/stage_init_utils.py +++ b/vllm_omni/engine/stage_init_utils.py @@ -411,6 +411,10 @@ def build_engine_args_dict( if stage_type != "diffusion": resolve_worker_cls(engine_args_dict) + # Check whether the stage's default_sampling_params defines extra_args. + default_sp = _to_dict(getattr(stage_config, "default_sampling_params", {})) + engine_args_dict["has_sampling_extra_args"] = bool(default_sp.get("extra_args")) + return engine_args_dict diff --git a/vllm_omni/entrypoints/openai/protocol/audio.py b/vllm_omni/entrypoints/openai/protocol/audio.py index 8468efd86..59b5777a8 100644 --- a/vllm_omni/entrypoints/openai/protocol/audio.py +++ b/vllm_omni/entrypoints/openai/protocol/audio.py @@ -1,5 +1,5 @@ import math -from typing import Literal +from typing import Any, Literal import numpy as np from pydantic import AliasChoices, BaseModel, Field, field_validator, model_validator @@ -74,6 +74,10 @@ class OpenAICreateSpeechRequest(BaseModel): ge=0, description="Per-request initial chunk size override. If null, computed dynamically based on server load.", ) + extra_params: dict[str, Any] | None = Field( + default=None, + description=("Optional model-specific parameters passed directly to the model's extra_args."), + ) @field_validator("stream_format") @classmethod diff --git a/vllm_omni/entrypoints/openai/serving_speech.py b/vllm_omni/entrypoints/openai/serving_speech.py index c275c7795..5f686c666 100644 --- a/vllm_omni/entrypoints/openai/serving_speech.py +++ b/vllm_omni/entrypoints/openai/serving_speech.py @@ -8,13 +8,14 @@ import re import struct import time from concurrent.futures import ThreadPoolExecutor +from http import HTTPStatus from pathlib import Path from typing import Any import numpy as np import soundfile as sf import torch -from fastapi import Request, UploadFile +from fastapi import HTTPException, Request, UploadFile from fastapi.responses import Response, StreamingResponse from transformers.utils.hub import cached_file from vllm.entrypoints.launcher import terminate_if_errored @@ -1664,6 +1665,21 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): max_tokens, ) + # Apply model-specific extra parameters + if request.extra_params is not None and sampling_params_list: + if not isinstance(request.extra_params, dict): + raise HTTPException( + status_code=HTTPStatus.BAD_REQUEST.value, + detail="extra_params must be a JSON object/dict.", + ) + import copy + + sampling_params_list = copy.deepcopy(sampling_params_list) + if sampling_params_list[0].extra_args is None: + sampling_params_list[0].extra_args = {} + sampling_params_list[0].extra_args.update(request.extra_params) + logger.info("Applied extra_params: %s", request.extra_params) + # Fish defaults come from stage_configs YAML. Only override when the caller # explicitly requests a different generation length. if self._is_fish_speech and request.max_new_tokens is not None and sampling_params_list: diff --git a/vllm_omni/model_executor/models/voxtral_tts/cuda_graph_acoustic_transformer_wrapper.py b/vllm_omni/model_executor/models/voxtral_tts/cuda_graph_acoustic_transformer_wrapper.py index ff053342d..d7407afe5 100644 --- a/vllm_omni/model_executor/models/voxtral_tts/cuda_graph_acoustic_transformer_wrapper.py +++ b/vllm_omni/model_executor/models/voxtral_tts/cuda_graph_acoustic_transformer_wrapper.py @@ -48,13 +48,13 @@ class CUDAGraphAcousticTransformerWrapper: self.n_acoustic_codebook = self.acoustic_transformer.model_args.n_acoustic_codebook self.acoustic_embeddings_levels = self.acoustic_transformer.acoustic_embeddings_levels - self.cfg_alpha = 1.2 self.n_steps = self.acoustic_transformer.acoustic_transformer_args.n_decoding_steps # Graph storage self.graphs: dict[int, CUDAGraph] = {} self.static_inputs: dict[int, torch.Tensor] = {} self.static_noise: dict[int, torch.Tensor] = {} + self.static_cfg_alpha: dict[int, torch.Tensor] = {} self.static_fake_eos: dict[int, torch.Tensor] = {} self.static_audio_codes: dict[int, torch.Tensor] = {} @@ -80,8 +80,10 @@ class CUDAGraphAcousticTransformerWrapper: # Phase 1: Eager warmup for ALL capture sizes for size in self.capture_sizes: dummy = torch.zeros(size, hidden_dim, device=device, dtype=dtype) + dummy_cfg_alpha = torch.full((size, 1), 1.2, device=device, dtype=dtype) + dummy_noise = torch.randn(size, self.n_acoustic_codebook, device=device, dtype=dtype) with torch.no_grad(): - self._forward_cudagraph_compatible(dummy) + self._forward_cudagraph_compatible(dummy, cfg_alpha=dummy_cfg_alpha, noise=dummy_noise) torch.cuda.synchronize(device) @@ -105,7 +107,12 @@ class CUDAGraphAcousticTransformerWrapper: len(self.capture_sizes), ) - def _forward_cudagraph_compatible(self, hidden_states: torch.Tensor, noise: torch.Tensor | None = None): + def _forward_cudagraph_compatible( + self, + hidden_states: torch.Tensor, + cfg_alpha: torch.Tensor, + noise: torch.Tensor, + ): """ The actual computation captured by the CUDA graph. @@ -117,6 +124,7 @@ class CUDAGraphAcousticTransformerWrapper: - Calls _predict_velocity directly - Uses a pre-allocated noise buffer to avoid baking random state into the CUDA graph + - Uses a pre-allocated cfg_alpha buffer for per-request CFG strength """ at = self.acoustic_transformer B = hidden_states.shape[0] @@ -132,10 +140,7 @@ class CUDAGraphAcousticTransformerWrapper: # --- Flow matching: Euler ODE --- should_decode = semantic_code.squeeze(1) != self.end_audio_token_id - if noise is not None: - x = noise - else: - x = torch.randn(B, self.n_acoustic_codebook, device=hidden_states.device, dtype=hidden_states.dtype) + x = noise # Pre-compute zero hidden states for unconditional CFG branch hidden_states_zero = torch.zeros_like(hidden_states) @@ -154,8 +159,8 @@ class CUDAGraphAcousticTransformerWrapper: v_all = at._predict_velocity(x_t=x_batched, llm_output=llm_batched, t_emb=t_emb_batched) v_t, uncond_v_t = v_all[:B], v_all[B:] - # CFG combination - v_t = self.cfg_alpha * v_t + (1 - self.cfg_alpha) * uncond_v_t + # CFG combination (cfg_alpha is (B, 1), v_t is (B, C)) + v_t = cfg_alpha * v_t + (1 - cfg_alpha) * uncond_v_t x = x + v_t * dt @@ -188,10 +193,11 @@ class CUDAGraphAcousticTransformerWrapper: """Capture a CUDA graph for a specific batch size.""" static_input = torch.zeros(size, hidden_dim, device=device, dtype=dtype) static_noise = torch.randn(size, self.n_acoustic_codebook, device=device, dtype=dtype) + static_cfg_alpha = torch.full((size, 1), 1.2, device=device, dtype=dtype) # Stabilizing eager run with torch.no_grad(): - _ = self._forward_cudagraph_compatible(static_input, noise=static_noise) + _ = self._forward_cudagraph_compatible(static_input, cfg_alpha=static_cfg_alpha, noise=static_noise) torch.cuda.synchronize(device) @@ -199,12 +205,13 @@ class CUDAGraphAcousticTransformerWrapper: with torch.no_grad(): with torch.cuda.graph(graph, pool=current_platform.get_global_graph_pool()): static_fake_eos, static_audio_codes = self._forward_cudagraph_compatible( - static_input, noise=static_noise + static_input, cfg_alpha=static_cfg_alpha, noise=static_noise ) self.graphs[size] = graph self.static_inputs[size] = static_input self.static_noise[size] = static_noise + self.static_cfg_alpha[size] = static_cfg_alpha self.static_fake_eos[size] = static_fake_eos self.static_audio_codes[size] = static_audio_codes @@ -218,6 +225,7 @@ class CUDAGraphAcousticTransformerWrapper: def __call__( self, hidden_states: torch.Tensor, + cfg_alpha: torch.Tensor, ) -> tuple[torch.Tensor, dict[str, list[torch.Tensor]] | None]: """ Drop-in replacement for model.compute_mm_logits(). @@ -229,16 +237,20 @@ class CUDAGraphAcousticTransformerWrapper: actual_size = hidden_states.shape[0] if not self.enabled or not self._warmed_up: - return self.model.compute_mm_logits(hidden_states) + return self.model.compute_mm_logits(hidden_states, cfg_alpha=cfg_alpha) padded_size = self._get_padded_size(actual_size) if padded_size is None or padded_size not in self.graphs: - return self.model.compute_mm_logits(hidden_states) + return self.model.compute_mm_logits(hidden_states, cfg_alpha=cfg_alpha) # Zero static input, then copy actual data self.static_inputs[padded_size].zero_() self.static_inputs[padded_size][:actual_size] = hidden_states + # Copy per-request cfg_alpha into static buffer (pad with 1.2 default) + self.static_cfg_alpha[padded_size].fill_(1.2) + self.static_cfg_alpha[padded_size][:actual_size, 0] = cfg_alpha + # Fill noise buffer with fresh random values before replay so the # flow-matching ODE starts from different initial noise each time. self.static_noise[padded_size].normal_() diff --git a/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts.py b/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts.py index 127171067..c78089150 100644 --- a/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts.py +++ b/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts.py @@ -283,6 +283,30 @@ class VoxtralTTSForConditionalGeneration( multimodal_outputs={"audio": batch_audio_arrays}, ) + _DEFAULT_CFG_ALPHA = 1.2 + + def _extract_cfg_alpha(self, input_hidden_states: torch.Tensor, **kwargs) -> torch.Tensor: + """Extract per-request cfg_alpha from sampling_extra_args. + + Returns a 1-D tensor of shape (B,) with per-request cfg_alpha values. + Falls back to default if sampling_extra_args is missing or incomplete. + """ + B = input_hidden_states.shape[0] + sampling_extra_args = kwargs.get("sampling_extra_args") + if sampling_extra_args is None: + return torch.full( + (B,), + self._DEFAULT_CFG_ALPHA, + device=input_hidden_states.device, + dtype=input_hidden_states.dtype, + ) + cfg_alpha_values = [ea.get("cfg_alpha", self._DEFAULT_CFG_ALPHA) for ea in sampling_extra_args] + return torch.tensor( + cfg_alpha_values, + device=input_hidden_states.device, + dtype=input_hidden_states.dtype, + ) + def make_omni_output( self, model_outputs: torch.Tensor | OmniOutput | tuple, logits_index: int | None = None, **kwargs ) -> OmniOutput: @@ -291,10 +315,15 @@ class VoxtralTTSForConditionalGeneration( hidden_states = model_outputs assert logits_index is not None input_hidden_states = hidden_states[logits_index] + cfg_alpha = self._extract_cfg_alpha(input_hidden_states, **kwargs) if self._cudagraph_acoustic_transformer is not None: - fake_eos, multimodal_outputs = self._cudagraph_acoustic_transformer(input_hidden_states) + fake_eos, multimodal_outputs = self._cudagraph_acoustic_transformer( + input_hidden_states, cfg_alpha=cfg_alpha + ) else: - fake_eos, multimodal_outputs = self.model.compute_mm_logits(input_hidden_states) + fake_eos, multimodal_outputs = self.model.compute_mm_logits( + input_hidden_states, cfg_alpha=cfg_alpha + ) hidden_states[logits_index, 0] = fake_eos return OmniOutput( text_hidden_states=hidden_states, diff --git a/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts_audio_generation.py b/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts_audio_generation.py index cd67e4f07..8b7dd7d13 100644 --- a/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts_audio_generation.py +++ b/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts_audio_generation.py @@ -438,8 +438,6 @@ class FlowMatchingAudioTransformer(nn.Module): # Flow matching constants self._n_steps = args.n_decoding_steps - # TODO(chenyo): hardcoded, need to fix - self._cfg_alpha = 1.2 self._noise_scale = 1.0 self.register_buffer( "_timesteps", @@ -512,6 +510,7 @@ class FlowMatchingAudioTransformer(nn.Module): self, semantic_code: torch.Tensor, llm_hidden: torch.Tensor, + cfg_alpha: torch.Tensor, ) -> torch.Tensor: B = semantic_code.shape[0] @@ -525,6 +524,10 @@ class FlowMatchingAudioTransformer(nn.Module): timesteps = self._timesteps.to(dtype=llm_hidden.dtype) llm_hidden_zero = torch.zeros_like(llm_hidden) + # Reshape cfg_alpha for broadcasting: (B,) -> (B, 1) + cfg_alpha = cfg_alpha.to(dtype=llm_hidden.dtype, device=llm_hidden.device) + cfg_alpha = cfg_alpha.unsqueeze(1) # (B, 1) for broadcasting with (B, C) + # Euler integration with batched conditional + unconditional velocity sampled = x_0 for i in range(len(timesteps) - 1): @@ -544,7 +547,7 @@ class FlowMatchingAudioTransformer(nn.Module): t_emb=t_emb_batched, ) v_t, uncond_v_t = v_all[:B], v_all[B:] - v_t = self._cfg_alpha * v_t + (1 - self._cfg_alpha) * uncond_v_t + v_t = cfg_alpha * v_t + (1 - cfg_alpha) * uncond_v_t sampled = sampled + v_t * dt @@ -585,6 +588,7 @@ class FlowMatchingAudioTransformer(nn.Module): def forward( self, llm_hidden: torch.Tensor, + cfg_alpha: torch.Tensor, ) -> torch.Tensor: # llm_hidden: BxD semantic_logit = self.semantic_codebook_output(llm_hidden).float() @@ -594,10 +598,10 @@ class FlowMatchingAudioTransformer(nn.Module): # semantic_logit: Bx1 semantic_code = semantic_logit.argmax(dim=-1, keepdim=True) - # acoustic codes, TODO(@chenyo): config sampling acoustic_codes = self.decode_one_frame( semantic_code.squeeze(1), llm_hidden, + cfg_alpha=cfg_alpha, ) audio_codes = torch.concatenate( @@ -1035,11 +1039,13 @@ class VoxtralTTSAudioGenerationForConditionalGeneration(nn.Module, SupportsMulti def compute_mm_logits( self, hidden_states: torch.Tensor, + cfg_alpha: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor | None]: audio_codes = None mm_tokens = None audio_codes = self.acoustic_transformer( llm_hidden=hidden_states, + cfg_alpha=cfg_alpha, ) fake_eos = torch.where( audio_codes[:, 0] == AudioSpecialTokens.id(AudioSpecialTokens.end_audio), diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py index d3ccbaaf3..7a6f3b453 100644 --- a/vllm_omni/worker/gpu_model_runner.py +++ b/vllm_omni/worker/gpu_model_runner.py @@ -1038,6 +1038,15 @@ class OmniGPUModelRunner(GPUModelRunner): import traceback traceback.print_exc() + + if getattr(self.model_config, "has_sampling_extra_args", False): + extra_args_list: list[dict] = [] + for req_id in self.input_batch.req_ids: + req = self.requests[req_id] + sp = req.sampling_params if req else None + extra_args_list.append(sp.extra_args if sp and sp.extra_args else {}) + model_kwargs_extra["sampling_extra_args"] = extra_args_list + return model_kwargs_extra def _process_additional_information_updates( From e82c0a5e1d08ffbbe131b0e785cfa9e42ca86074 Mon Sep 17 00:00:00 2001 From: ChenWenjing <54166744+Shirley125@users.noreply.github.com> Date: Wed, 22 Apr 2026 12:08:56 +0800 Subject: [PATCH 05/45] [ci]add merge ci for streaming input (#2965) Signed-off-by: CHEN <116010019@link.cuhk.edu.cn> --- .buildkite/test-merge.yml | 40 +++++++++++++++++++ .../test_qwen3_omni_realtime_websocket.py | 29 +++++++------- tests/helpers/media.py | 13 ++++-- 3 files changed, 64 insertions(+), 18 deletions(-) rename tests/{e2e/online_serving => entrypoints/openai_api}/test_qwen3_omni_realtime_websocket.py (91%) diff --git a/.buildkite/test-merge.yml b/.buildkite/test-merge.yml index 785cc58fa..691f3f876 100644 --- a/.buildkite/test-merge.yml +++ b/.buildkite/test-merge.yml @@ -275,6 +275,46 @@ steps: path: /mnt/hf-cache type: DirectoryOrCreate + - label: "Audio Streaming Input Test with H100" + timeout_in_minutes: 30 + depends_on: upload-merge-pipeline + commands: + - export VLLM_TEST_CLEAN_GPU_MEMORY="1" + - pytest -s -v tests/entrypoints/openai_api/test_qwen3_omni_realtime_websocket.py -m "advanced_model" --run-level "advanced_model" + agents: + queue: "mithril-h100-pool" + plugins: + - kubernetes: + podSpec: + containers: + - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT + resources: + limits: + nvidia.com/gpu: 2 + volumeMounts: + - name: devshm + mountPath: /dev/shm + - name: hf-cache + mountPath: /root/.cache/huggingface + env: + - name: HF_HOME + value: /root/.cache/huggingface + - name: HF_TOKEN + valueFrom: + secretKeyRef: + name: hf-token-secret + key: token + nodeSelector: + node.kubernetes.io/instance-type: gpu-h100-sxm + volumes: + - name: devshm + emptyDir: + medium: Memory + - name: hf-cache + hostPath: + path: /mnt/hf-cache + type: DirectoryOrCreate + - label: "Diffusion Image Edit Test with H100 (1 GPU)" timeout_in_minutes: 20 depends_on: upload-merge-pipeline diff --git a/tests/e2e/online_serving/test_qwen3_omni_realtime_websocket.py b/tests/entrypoints/openai_api/test_qwen3_omni_realtime_websocket.py similarity index 91% rename from tests/e2e/online_serving/test_qwen3_omni_realtime_websocket.py rename to tests/entrypoints/openai_api/test_qwen3_omni_realtime_websocket.py index f3b261081..90f8897c5 100644 --- a/tests/e2e/online_serving/test_qwen3_omni_realtime_websocket.py +++ b/tests/entrypoints/openai_api/test_qwen3_omni_realtime_websocket.py @@ -23,33 +23,28 @@ from tests.helpers.media import ( generate_synthetic_audio, ) from tests.helpers.runtime import OmniServerParams -from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config +from tests.helpers.stage_config import get_deploy_config_path os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" MODEL = "Qwen/Qwen3-Omni-30B-A3B-Instruct" +# Synthetic input for realtime E2E (``generate_synthetic_audio``); distinct cache file per phrase. +REALTIME_SYNTH_PHRASE_TEXT = "Translate into Chinese: Beijing is the Capital of China" + # The new-schema CI overlay bakes in async_chunk: False and covers CUDA/ROCm/XPU # via its ``platforms:`` section, so one path serves all three. default_stage_config = get_deploy_config_path("ci/qwen3_omni_moe.yaml") - -def _realtime_stage_config_path() -> str: - """CI omni layout without async_chunk; stage 0 thinker max_tokens=10.""" - return modify_stage_config( - default_stage_config, - updates={"stages": {0: {"default_sampling_params.max_tokens": 10}}}, - ) - - realtime_server_params = [ pytest.param( OmniServerParams( model=MODEL, - stage_config_path=_realtime_stage_config_path(), + stage_config_path=default_stage_config, use_stage_cli=True, + server_args=["--no-async-chunk"], ), - id="thinker_max_tokens_10", + id="default", ), ] @@ -169,9 +164,15 @@ class TestQwen3OmniRealtimeWebSocket: """ Short streamed 16 kHz mono PCM16 input; expect streamed PCM16 audio deltas and transcription. Verify Whisper(output audio) aligns with model text (same idea - as multimodal omni e2e). + as multimodal omni e2e). Input speech is synthesized from + ``REALTIME_SYNTH_PHRASE_TEXT``. """ - syn = generate_synthetic_audio(10, 1, sample_rate=16000) + syn = generate_synthetic_audio( + 10, + 1, + sample_rate=16000, + phrase_text=REALTIME_SYNTH_PHRASE_TEXT, + ) wav_bytes = base64.b64decode(syn["base64"]) pcm16 = _pcm16_mono_16k_from_wav_bytes(wav_bytes) diff --git a/tests/helpers/media.py b/tests/helpers/media.py index 4463acbbb..2e99788b1 100644 --- a/tests/helpers/media.py +++ b/tests/helpers/media.py @@ -3,6 +3,7 @@ import base64 import concurrent.futures import gc +import hashlib import io import logging import math @@ -66,6 +67,7 @@ def generate_synthetic_audio( num_channels: int, sample_rate: int = 48000, *, + phrase_text: str = "test", force_regenerate: bool = False, cache_dir: Path | str | None = None, ) -> dict[str, Any]: @@ -74,12 +76,16 @@ def generate_synthetic_audio( Caches the WAV under ``cache_dir`` when given, else under the default temp subdirectory. Reuses the file when the same - ``duration`` / ``num_channels`` / ``sample_rate`` are requested unless - ``force_regenerate`` is true. + ``duration`` / ``num_channels`` / ``sample_rate`` / ``phrase_text`` are + requested unless ``force_regenerate`` is true. + + The cache filename includes a SHA-256 digest of ``phrase_text`` so different + phrases never share a WAV cache entry. """ root = _resolve_synthetic_media_cache_dir(cache_dir) root.mkdir(parents=True, exist_ok=True) - cache_path = root / f"synth_audio_d{duration}_ch{num_channels}_sr{sample_rate}.wav" + phrase_key = hashlib.sha256(phrase_text.encode("utf-8")).hexdigest() + cache_path = root / f"synth_audio_d{duration}_ch{num_channels}_sr{sample_rate}_pt{phrase_key}.wav" if not force_regenerate and cache_path.is_file(): data, _sr = sf.read(str(cache_path), dtype="float32", always_2d=True) @@ -204,7 +210,6 @@ def generate_synthetic_audio( enhanced = enhanced / peak * 0.95 return enhanced.astype(np.float32) - phrase_text = "test" num_samples = int(sample_rate * max(1, duration)) audio_data = np.zeros((num_samples, num_channels), dtype=np.float32) From e1cdd0c23bb3cb589479a6ea47d1c920d8006059 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=B1=AA=E5=BF=97=E9=B9=8F?= Date: Wed, 22 Apr 2026 12:36:20 +0800 Subject: [PATCH 06/45] [BugFix][CI]: Fix test_omni_sleep_mode ci bug (#3010) Signed-off-by: princepride --- tests/e2e/offline_inference/test_omni_sleep_mode.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/e2e/offline_inference/test_omni_sleep_mode.py b/tests/e2e/offline_inference/test_omni_sleep_mode.py index 2ad0b53b0..fc4001ff5 100644 --- a/tests/e2e/offline_inference/test_omni_sleep_mode.py +++ b/tests/e2e/offline_inference/test_omni_sleep_mode.py @@ -123,7 +123,7 @@ async def test_multistage_sleep_h100(tp_size): pass acks = await engine.sleep(stage_ids=[0, 1], level=2) - assert len(acks) == 2 * tp_size + assert len(acks) == 2 await engine.wake_up(stage_ids=[0, 1]) async for _ in engine.generate("verify", sampling_params=[SamplingParams(), sp]): From 554233208f7619094fc1304bee42670eda9325a1 Mon Sep 17 00:00:00 2001 From: Jared Wen Date: Wed, 22 Apr 2026 15:09:28 +0800 Subject: [PATCH 07/45] [ConfigRefactor] GLM-Image (#2977) Signed-off-by: JaredforReal Signed-off-by: Jared Wen Co-authored-by: SYLAR <125541396+lishunyang12@users.noreply.github.com> --- .../offline_inference/glm_image/README.md | 145 ----- .../offline_inference/glm_image/end2end.py | 525 ------------------ .../offline_inference/glm_image/run_i2i.sh | 93 ---- .../offline_inference/glm_image/run_t2i.sh | 87 --- examples/online_serving/glm_image/README.md | 204 ------- .../glm_image/openai_chat_client.py | 172 ------ .../glm_image/run_curl_image_edit.sh | 61 -- .../glm_image/run_curl_text_to_image.sh | 30 - .../online_serving/glm_image/run_server.sh | 12 - vllm_omni/config/pipeline_registry.py | 4 + vllm_omni/config/stage_config.py | 39 ++ vllm_omni/deploy/glm_image.yaml | 38 ++ .../models/glm_image/__init__.py | 15 +- .../models/glm_image/pipeline.py | 47 ++ .../stage_configs/glm_image.yaml | 74 --- .../glm_image_muilticonnector.yaml | 87 --- 16 files changed, 142 insertions(+), 1491 deletions(-) delete mode 100644 examples/offline_inference/glm_image/README.md delete mode 100644 examples/offline_inference/glm_image/end2end.py delete mode 100755 examples/offline_inference/glm_image/run_i2i.sh delete mode 100755 examples/offline_inference/glm_image/run_t2i.sh delete mode 100644 examples/online_serving/glm_image/README.md delete mode 100644 examples/online_serving/glm_image/openai_chat_client.py delete mode 100755 examples/online_serving/glm_image/run_curl_image_edit.sh delete mode 100755 examples/online_serving/glm_image/run_curl_text_to_image.sh delete mode 100755 examples/online_serving/glm_image/run_server.sh create mode 100644 vllm_omni/deploy/glm_image.yaml create mode 100644 vllm_omni/model_executor/models/glm_image/pipeline.py delete mode 100644 vllm_omni/model_executor/stage_configs/glm_image.yaml delete mode 100644 vllm_omni/model_executor/stage_configs/glm_image_muilticonnector.yaml diff --git a/examples/offline_inference/glm_image/README.md b/examples/offline_inference/glm_image/README.md deleted file mode 100644 index c3c7c2916..000000000 --- a/examples/offline_inference/glm_image/README.md +++ /dev/null @@ -1,145 +0,0 @@ -# GLM-Image Multistage End-to-End Inference - -This example demonstrates how to run GLM-Image with the vLLM-Omni multistage architecture. - -## Architecture - -GLM-Image uses a 2-stage pipeline: - -``` -┌─────────────────────────────────────────────────────────────┐ -│ GLM-Image Pipeline │ -├─────────────────────────────────────────────────────────────┤ -│ │ -│ Stage 0 (AR Model) Stage 1 (Diffusion) │ -│ ┌─────────────────┐ ┌─────────────────────┐ │ -│ │ vLLM-optimized │ │ GlmImagePipeline │ │ -│ │ GlmImageFor │ prior │ ┌───────────────┐ │ │ -│ │ Conditional │──tokens───►│ │ DiT Denoiser │ │ │ -│ │ Generation │ │ └───────────────┘ │ │ -│ │ (9B AR model) │ │ │ │ │ -│ └─────────────────┘ │ ▼ │ │ -│ ▲ │ ┌───────────────┐ │ │ -│ │ │ │ VAE Decode │──┼──► Image -│ Text/Image │ └───────────────┘ │ │ -│ Input └─────────────────────┘ │ -│ │ -└─────────────────────────────────────────────────────────────┘ -``` - -## Features - -- **vLLM-optimized AR**: Uses PagedAttention and tensor parallelism for faster prior token generation -- **Flexible deployment**: AR and Diffusion stages can run on different GPUs -- **Text-to-Image**: Generate images from text descriptions -- **Image-to-Image**: Edit existing images with text prompts - -## Usage - -### Text-to-Image - -```bash -python end2end.py \ - --config-path ../../../vllm_omni/model_executor/stage_configs/glm_image.yaml \ - --prompt "A beautiful sunset over the ocean with sailing boats" \ - --height 1024 \ - --width 1024 \ - --output output_t2i.png -``` - -### Image-to-Image (Image Editing) - -```bash -python end2end.py \ - --config-path ../../../vllm_omni/model_executor/stage_configs/glm_image.yaml \ - --prompt "Transform this scene into a winter wonderland" \ - --image input.png \ - --output output_i2i.png -``` - -### With Custom Parameters - -```bash -python end2end.py \ - --model-path /path/to/glm-image \ - --config-path ../../../vllm_omni/model_executor/stage_configs/glm_image.yaml \ - --prompt "A photorealistic cat sitting on a window sill" \ - --height 1024 \ - --width 1024 \ - --num-inference-steps 50 \ - --guidance-scale 1.5 \ - --seed 42 \ - --output output.png -``` - -## Shell Scripts - -### Run Text-to-Image - -```bash -./run_t2i.sh -``` - -### Run Image-to-Image - -```bash -./run_i2i.sh --image /path/to/input.png -``` - -## Stage Configuration - -The stage config (`glm_image.yaml`) defines: - -- **Stage 0 (AR)**: Uses `GPUARWorker` with vLLM engine - - - Model: `GlmImageForConditionalGeneration` - - Output: `token_ids` (prior tokens) - -- **Stage 1 (Diffusion)**: Uses diffusion engine - - Model: `GlmImagePipeline` - - Output: Generated image - -See `vllm_omni/model_executor/stage_configs/glm_image.yaml` for full configuration. - -## Comparison with Single-Stage - -| Aspect | Single-Stage (transformers) | Multistage (vLLM) | -| ----------- | --------------------------- | ------------------- | -| AR Model | transformers native | vLLM PagedAttention | -| Memory | Higher (no KV cache opt) | Lower (optimized) | -| Throughput | Lower | Higher | -| Flexibility | Single GPU | Multi-GPU support | - -## Troubleshooting - -### OOM Error - -Try reducing memory usage: - -```bash -# In glm_image.yaml, adjust: -gpu_memory_utilization: 0.5 # Reduce from 0.6 -``` - -### Slow Initialization - -The first run loads model weights. Subsequent runs are faster: - -```bash ---stage-init-timeout 900 # Increase timeout for slow storage -``` - -### `Transformers does not recognize this architecture` Error - -Your have to upgrade `transformers` package to `5.3.0` or above: - -``` -pip install --upgrade transformers -``` - -## Requirements - -- vLLM-Omni with GLM-Image support -- CUDA-capable GPU (recommended: H100/A100 with 80GB) -- GLM-Image model weights -- `transformers` v5.3.0 or above diff --git a/examples/offline_inference/glm_image/end2end.py b/examples/offline_inference/glm_image/end2end.py deleted file mode 100644 index 7ae9478ba..000000000 --- a/examples/offline_inference/glm_image/end2end.py +++ /dev/null @@ -1,525 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -End-to-end offline inference example for GLM-Image with multistage architecture. - -This script tests the multistage pipeline where: -- Stage 0 (AR): vLLM-optimized GlmImageForConditionalGeneration generates prior_token_ids -- Stage 1 (Diffusion): GlmImagePipeline performs DiT denoising + VAE decode - -Usage (text-to-image): - python end2end.py \ - --model-path /path/to/glm-image \ - --config-path /path/to/glm_image.yaml \ - --prompt "A beautiful sunset over the ocean" \ - --output output_t2i.png - -Usage (image-to-image / image edit): - python end2end.py \ - --model-path /path/to/glm-image \ - --config-path /path/to/glm_image.yaml \ - --prompt "Make it look like winter" \ - --image input.png \ - --output output_i2i.png - -Usage (with custom parameters): - python end2end.py \ - --model-path /path/to/glm-image \ - --config-path /path/to/glm_image.yaml \ - --prompt "A cat sitting on a window sill" \ - --height 1024 \ - --width 1024 \ - --num-inference-steps 50 \ - --guidance-scale 1.5 \ - --seed 42 - -For more options, run: - python end2end.py --help -""" - -import argparse -import os -import time -from pathlib import Path - -from PIL import Image - -from vllm_omni.entrypoints.omni import Omni -from vllm_omni.inputs.data import OmniDiffusionSamplingParams - -# Default stage config path (relative to vllm_omni package) -DEFAULT_CONFIG_PATH = "vllm_omni/model_executor/stage_configs/glm_image.yaml" - -SEED = 42 - -# GLM-Image special tokens -GLM_IMAGE_EOS_TOKEN_ID = 16385 # eos_token_id from generation_config.json -GLM_IMAGE_VISION_VOCAB_SIZE = 16512 # top_k should be vision_vocab_size - - -def compute_max_tokens(height: int, width: int, factor: int = 32, is_i2i: bool = False) -> int: - """ - Compute max_new_tokens for GLM-Image AR generation. - - GLM-Image generation differs by mode: - - text-to-image (t2i): small preview + large target + EOS - - image-to-image (i2i): large target + EOS - - Args: - height: Target image height in pixels - width: Target image width in pixels - factor: Downsampling factor (32 for GLM-Image AR output) - is_i2i: Whether the request is image-to-image mode - - Returns: - Total number of tokens to generate for the specified mode - """ - # Large image tokens (target resolution) - token_h = height // factor - token_w = width // factor - large_tokens = token_h * token_w - - # Small preview tokens (half resolution in each dimension) - import math - - ratio = token_h / token_w if token_w > 0 else 1.0 - small_token_h = max(1, int(math.sqrt(ratio) * (factor // 2))) - small_token_w = max(1, int(math.sqrt(1 / ratio) * (factor // 2))) - small_tokens = small_token_h * small_token_w - - if is_i2i: - return large_tokens + 1 - return small_tokens + large_tokens + 1 - - -def load_image(image_path: str) -> Image.Image: - """Load an image from file path.""" - if not os.path.exists(image_path): - raise FileNotFoundError(f"Image file not found: {image_path}") - return Image.open(image_path).convert("RGB") - - -def save_image(image: Image.Image, output_path: str) -> None: - """Save an image to file path.""" - output_dir = os.path.dirname(output_path) - if output_dir: - os.makedirs(output_dir, exist_ok=True) - image.save(output_path) - print(f"Image saved to: {output_path}") - - -def build_prompt_for_t2i( - prompt: str, - height: int = 1024, - width: int = 1024, -) -> dict: - """ - Build prompt dict for text-to-image generation. - - Args: - prompt: Text description for image generation - height: Target image height - width: Target image width - - Returns: - Dict containing prompt and generation parameters - """ - return { - "prompt": prompt, - "height": height, - "width": width, - # Pass target dimensions to AR processor for proper grid token generation - "mm_processor_kwargs": { - "target_h": height, - "target_w": width, - }, - } - - -def build_prompt_for_i2i( - prompt: str, - image: Image.Image, - height: int | None = None, - width: int | None = None, -) -> dict: - """ - Build prompt dict for image-to-image generation. - - Args: - prompt: Text description for image editing - image: Source image for editing - height: Target image height (default: use source image size) - width: Target image width (default: use source image size) - - Returns: - Dict containing prompt, image, and generation parameters - """ - # Use source image dimensions if not specified - if height is None: - height = image.height - if width is None: - width = image.width - - return { - "prompt": prompt, - "multi_modal_data": { - "image": image, - }, - "height": height, - "width": width, - # Pass target dimensions to AR processor for proper grid token generation - "mm_processor_kwargs": { - "target_h": height, - "target_w": width, - }, - } - - -def main(args: argparse.Namespace) -> None: - """Main entry point for GLM-Image end-to-end inference.""" - print("=" * 60) - print("GLM-Image Multistage End-to-End Inference") - print("=" * 60) - - # Validate arguments - if not args.prompt: - raise ValueError("--prompt is required") - - # Determine config path - config_path = args.config_path - if config_path is None: - # Try to find default config - if os.path.exists(DEFAULT_CONFIG_PATH): - config_path = DEFAULT_CONFIG_PATH - else: - # Try relative to script location - script_dir = Path(__file__).parent.parent.parent.parent - config_path = script_dir / "vllm_omni/model_executor/stage_configs/glm_image.yaml" - if not config_path.exists(): - raise FileNotFoundError( - f"Stage config not found. Please specify --config-path. Tried: {DEFAULT_CONFIG_PATH}" - ) - config_path = str(config_path) - - print(f"Model path: {args.model_path}") - print(f"Config path: {config_path}") - print(f"Prompt: {args.prompt}") - - # Load source image for image-to-image mode - source_image = None - if args.image: - print(f"Source image: {args.image}") - source_image = load_image(args.image) - print(f" Image size: {source_image.width}x{source_image.height}") - - # Build prompt based on mode - if source_image is not None: - # Image-to-image mode - prompt_dict = build_prompt_for_i2i( - prompt=args.prompt, - image=source_image, - height=args.height, - width=args.width, - ) - mode = "image-to-image" - else: - # Text-to-image mode - prompt_dict = build_prompt_for_t2i( - prompt=args.prompt, - height=args.height or 1024, - width=args.width or 1024, - ) - mode = "text-to-image" - - print(f"Mode: {mode}") - print(f"Target size: {prompt_dict.get('height', 1024)}x{prompt_dict.get('width', 1024)}") - - # Add generation parameters to prompt - prompt_dict["seed"] = args.seed - prompt_dict["num_inference_steps"] = args.num_inference_steps - prompt_dict["guidance_scale"] = args.guidance_scale - - if args.negative_prompt: - prompt_dict["negative_prompt"] = args.negative_prompt - - # Build cache-dit config if requested - cache_config = None - if args.cache_backend == "cache_dit": - cache_config = { - "Fn_compute_blocks": 1, - "Bn_compute_blocks": 0, - "max_warmup_steps": 4, - "residual_diff_threshold": 0.24, - "max_continuous_cached_steps": 3, - "enable_taylorseer": False, - "taylorseer_order": 1, - "scm_steps_mask_policy": None, - "scm_steps_policy": "dynamic", - } - - # Initialize Omni with multistage config - print("\nInitializing Omni with multistage pipeline...") - print(f" Cache backend: {args.cache_backend or 'None (no acceleration)'}") - start_time = time.time() - - omni = Omni( - model=args.model_path, - stage_configs_path=config_path, - log_stats=args.enable_stats, - stage_init_timeout=args.stage_init_timeout, - cache_backend=args.cache_backend, - cache_config=cache_config, - enable_cache_dit_summary=getattr(args, "enable_cache_dit_summary", False), - enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler, - ) - - init_time = time.time() - start_time - print(f"Initialization completed in {init_time:.2f}s") - - # Prepare prompts (support batch generation) - prompts = [prompt_dict for _ in range(args.num_prompts)] - - # No explicit sampling_params for diffusion - parameters are in prompt_dict - # For multistage, the AR stage may need sampling params - from vllm import SamplingParams - - # Compute max_tokens dynamically based on target image size - target_height = prompt_dict.get("height", 1024) - target_width = prompt_dict.get("width", 1024) - is_i2i = source_image is not None - calculated_max_tokens = compute_max_tokens(target_height, target_width, is_i2i=is_i2i) - - # Use calculated value unless user explicitly specified a different value - # Default args.max_tokens is 16384 (very large), so prefer calculated value - effective_max_tokens = calculated_max_tokens if args.max_tokens == 16384 else args.max_tokens - - if args.verbose: - print( - f"AR max_tokens: {effective_max_tokens} " - f"(calculated: {calculated_max_tokens}, arg: {args.max_tokens}, mode: {'i2i' if is_i2i else 't2i'})" - ) - - # IMPORTANT: GLM-Image AR model requires these exact sampling parameters - # from generation_config.json for proper image token generation. - # - temperature=0.9, top_p=0.75, top_k=16512 (vision_vocab_size) - # - stop_token_ids=[16385] (eos_token_id) is CRITICAL to stop generation - ar_sampling_params = SamplingParams( - temperature=0.9, # From generation_config.json - top_p=0.75, # From generation_config.json - top_k=GLM_IMAGE_VISION_VOCAB_SIZE, # 16512, vision vocabulary size - max_tokens=effective_max_tokens, - stop_token_ids=[GLM_IMAGE_EOS_TOKEN_ID], # 16385, CRITICAL for stopping - seed=args.seed, - detokenize=False, - # Keep target size available in runner/model for deterministic M-RoPE - # decode grids in t2i (no mm_features available in this path). - extra_args={ - "target_h": int(target_height), - "target_w": int(target_width), - }, - ) - - # For diffusion stage, sampling_params contains diffusion-specific parameters - # These are passed as kwargs to the diffusion engine - diffusion_sampling_params = OmniDiffusionSamplingParams( - num_inference_steps=args.num_inference_steps, - guidance_scale=args.guidance_scale, - height=prompt_dict.get("height", 1024), - width=prompt_dict.get("width", 1024), - seed=args.seed, - ) - - # For multistage, we need sampling_params for each stage - # Stage 0 (AR): SamplingParams for vLLM - # Stage 1 (Diffusion): dict with diffusion kwargs - sampling_params_list = [ar_sampling_params, diffusion_sampling_params] - - # Run generation - print(f"\nGenerating {args.num_prompts} image(s)...") - gen_start_time = time.time() - - output_dir = os.path.dirname(args.output) if args.output else "outputs" - if output_dir: - os.makedirs(output_dir, exist_ok=True) - - output_count = 0 - for stage_outputs in omni.generate(prompts, sampling_params_list, py_generator=True): - output = stage_outputs.request_output - if stage_outputs.final_output_type == "image": - request_id = output.request_id - - # Get generated images - images = output.images if hasattr(output, "images") else [] - if not images and hasattr(output, "multimodal_output"): - images = output.multimodal_output.get("images", []) - - # Save each generated image - for idx, img in enumerate(images): - if args.num_prompts == 1 and len(images) == 1: - output_path = args.output - else: - base, ext = os.path.splitext(args.output) - output_path = f"{base}_{request_id}_{idx}{ext}" - - if isinstance(img, Image.Image): - save_image(img, output_path) - else: - print(f"Warning: Unexpected image type for request {request_id}: {type(img)}") - - output_count += 1 - - elif stage_outputs.final_output_type == "text": - # AR stage output (intermediate, for debugging) - if args.verbose: - print(f"AR output for request {output.request_id}:") - print(f" Token count: {len(output.outputs[0].token_ids)}") - - gen_time = time.time() - gen_start_time - print(f"\nGeneration completed in {gen_time:.2f}s") - print(f"Generated {output_count} image(s)") - - # Cleanup - omni.close() - print("\nDone!") - - -def parse_args() -> argparse.Namespace: - """Parse command line arguments.""" - parser = argparse.ArgumentParser( - description="GLM-Image Multistage End-to-End Inference", - formatter_class=argparse.RawDescriptionHelpFormatter, - epilog=__doc__, - ) - - # Required arguments - parser.add_argument( - "--model-path", - type=str, - default="zai-org/GLM-Image", - help="Path to GLM-Image model directory or HuggingFace model ID", - ) - parser.add_argument( - "--prompt", - type=str, - required=True, - help="Text prompt for image generation", - ) - - # Optional arguments - parser.add_argument( - "--config-path", - type=str, - default=None, - help="Path to stage config YAML file (default: auto-detect)", - ) - parser.add_argument( - "--image", - type=str, - default=None, - help="Path to source image for image-to-image mode", - ) - parser.add_argument( - "--output", - type=str, - default="output_glm_image.png", - help="Output image path (default: output_glm_image.png)", - ) - parser.add_argument( - "--negative-prompt", - type=str, - default=None, - help="Negative prompt for classifier-free guidance", - ) - - # Generation parameters - parser.add_argument( - "--height", - type=int, - default=None, - help="Output image height (default: 1024 for t2i, source size for i2i)", - ) - parser.add_argument( - "--width", - type=int, - default=None, - help="Output image width (default: 1024 for t2i, source size for i2i)", - ) - parser.add_argument( - "--num-inference-steps", - type=int, - default=50, - help="Number of diffusion denoising steps (default: 50)", - ) - parser.add_argument( - "--guidance-scale", - type=float, - default=1.5, - help="Classifier-free guidance scale (default: 1.5)", - ) - parser.add_argument( - "--seed", - type=int, - default=SEED, - help=f"Random seed for reproducibility (default: {SEED})", - ) - parser.add_argument( - "--max-tokens", - type=int, - default=16384, - help="Maximum tokens for AR generation (default: 16384)", - ) - - # Batch processing - parser.add_argument( - "--num-prompts", - type=int, - default=1, - help="Number of images to generate (default: 1)", - ) - - # Cache acceleration - parser.add_argument( - "--cache-backend", - type=str, - default=None, - choices=["cache_dit"], - help="Cache backend for DiT acceleration. Default: None (no cache).", - ) - parser.add_argument( - "--enable-cache-dit-summary", - action="store_true", - help="Enable cache-dit summary logging after diffusion forward passes.", - ) - - # Runtime options - parser.add_argument( - "--enable-stats", - action="store_true", - default=False, - help="Enable statistics logging", - ) - parser.add_argument( - "--stage-init-timeout", - type=int, - default=600, - help="Timeout for stage initialization in seconds (default: 600)", - ) - parser.add_argument( - "--verbose", - "-v", - action="store_true", - default=False, - help="Enable verbose output", - ) - parser.add_argument( - "--enable-diffusion-pipeline-profiler", - action="store_true", - help="Enable diffusion pipeline profiler to display stage durations.", - ) - - return parser.parse_args() - - -if __name__ == "__main__": - args = parse_args() - main(args) diff --git a/examples/offline_inference/glm_image/run_i2i.sh b/examples/offline_inference/glm_image/run_i2i.sh deleted file mode 100755 index f81b157b0..000000000 --- a/examples/offline_inference/glm_image/run_i2i.sh +++ /dev/null @@ -1,93 +0,0 @@ -#!/bin/bash -# SPDX-License-Identifier: Apache-2.0 -# Run GLM-Image image-to-image (editing) with multistage pipeline - -set -e - -# Default values -MODEL_PATH="${MODEL_PATH:-/path/to/glm-image}" -CONFIG_PATH="${CONFIG_PATH:-vllm_omni/model_executor/stage_configs/glm_image.yaml}" -PROMPT="${PROMPT:-Transform this image into an oil painting style}" -INPUT_IMAGE="" -OUTPUT="${OUTPUT:-output_i2i.png}" -NUM_STEPS="${NUM_STEPS:-50}" -GUIDANCE="${GUIDANCE:-1.5}" -SEED="${SEED:-42}" - -# Parse command line arguments -while [[ $# -gt 0 ]]; do - case $1 in - --model-path) - MODEL_PATH="$2" - shift 2 - ;; - --config-path) - CONFIG_PATH="$2" - shift 2 - ;; - --prompt) - PROMPT="$2" - shift 2 - ;; - --image) - INPUT_IMAGE="$2" - shift 2 - ;; - --output) - OUTPUT="$2" - shift 2 - ;; - --num-steps) - NUM_STEPS="$2" - shift 2 - ;; - --guidance) - GUIDANCE="$2" - shift 2 - ;; - --seed) - SEED="$2" - shift 2 - ;; - *) - echo "Unknown option: $1" - exit 1 - ;; - esac -done - -# Check if input image is provided -if [ -z "${INPUT_IMAGE}" ]; then - echo "Error: --image is required for image-to-image mode" - echo "Usage: ./run_i2i.sh --image /path/to/input.png [--prompt \"edit instruction\"]" - exit 1 -fi - -if [ ! -f "${INPUT_IMAGE}" ]; then - echo "Error: Input image not found: ${INPUT_IMAGE}" - exit 1 -fi - -echo "==============================================" -echo "GLM-Image Image-to-Image Generation" -echo "==============================================" -echo "Model: ${MODEL_PATH}" -echo "Config: ${CONFIG_PATH}" -echo "Input: ${INPUT_IMAGE}" -echo "Prompt: ${PROMPT}" -echo "Output: ${OUTPUT}" -echo "Steps: ${NUM_STEPS}" -echo "Guidance: ${GUIDANCE}" -echo "Seed: ${SEED}" -echo "==============================================" - -python end2end.py \ - --model-path "${MODEL_PATH}" \ - --config-path "${CONFIG_PATH}" \ - --prompt "${PROMPT}" \ - --image "${INPUT_IMAGE}" \ - --output "${OUTPUT}" \ - --num-inference-steps "${NUM_STEPS}" \ - --guidance-scale "${GUIDANCE}" \ - --seed "${SEED}" \ - --verbose diff --git a/examples/offline_inference/glm_image/run_t2i.sh b/examples/offline_inference/glm_image/run_t2i.sh deleted file mode 100755 index 5d249960b..000000000 --- a/examples/offline_inference/glm_image/run_t2i.sh +++ /dev/null @@ -1,87 +0,0 @@ -#!/bin/bash -# SPDX-License-Identifier: Apache-2.0 -# Run GLM-Image text-to-image generation with multistage pipeline - -set -e - -# Default values -MODEL_PATH="${MODEL_PATH:-/path/to/glm-image}" -CONFIG_PATH="${CONFIG_PATH:-vllm_omni/model_executor/stage_configs/glm_image.yaml}" -PROMPT="${PROMPT:-A beautiful mountain landscape with snow-capped peaks and a clear blue lake}" -OUTPUT="${OUTPUT:-output_t2i.png}" -HEIGHT="${HEIGHT:-1024}" -WIDTH="${WIDTH:-1024}" -NUM_STEPS="${NUM_STEPS:-50}" -GUIDANCE="${GUIDANCE:-1.5}" -SEED="${SEED:-42}" - -# Parse command line arguments -while [[ $# -gt 0 ]]; do - case $1 in - --model-path) - MODEL_PATH="$2" - shift 2 - ;; - --config-path) - CONFIG_PATH="$2" - shift 2 - ;; - --prompt) - PROMPT="$2" - shift 2 - ;; - --output) - OUTPUT="$2" - shift 2 - ;; - --height) - HEIGHT="$2" - shift 2 - ;; - --width) - WIDTH="$2" - shift 2 - ;; - --num-steps) - NUM_STEPS="$2" - shift 2 - ;; - --guidance) - GUIDANCE="$2" - shift 2 - ;; - --seed) - SEED="$2" - shift 2 - ;; - *) - echo "Unknown option: $1" - exit 1 - ;; - esac -done - -echo "==============================================" -echo "GLM-Image Text-to-Image Generation" -echo "==============================================" -echo "Model: ${MODEL_PATH}" -echo "Config: ${CONFIG_PATH}" -echo "Prompt: ${PROMPT}" -echo "Output: ${OUTPUT}" -echo "Size: ${WIDTH}x${HEIGHT}" -echo "Steps: ${NUM_STEPS}" -echo "Guidance: ${GUIDANCE}" -echo "Seed: ${SEED}" -echo "==============================================" - -python end2end.py \ - --model-path "${MODEL_PATH}" \ - --config-path "${CONFIG_PATH}" \ - --prompt "${PROMPT}" \ - --output "${OUTPUT}" \ - --height "${HEIGHT}" \ - --width "${WIDTH}" \ - --num-inference-steps "${NUM_STEPS}" \ - --guidance-scale "${GUIDANCE}" \ - --seed "${SEED}" \ - --verbose diff --git a/examples/online_serving/glm_image/README.md b/examples/online_serving/glm_image/README.md deleted file mode 100644 index 7ba4b501c..000000000 --- a/examples/online_serving/glm_image/README.md +++ /dev/null @@ -1,204 +0,0 @@ -# GLM-Image Online Serving - -This example demonstrates how to deploy GLM-Image for online image generation using vLLM-Omni. - -## 🛠️ Installation - -Please refer to [README.md](../../../README.md) - -## Run examples (GLM-Image) - -**Note**: These examples work with the default configuration on **2× NVIDIA A100 (80GB)** or equivalent. Stage 0 (AR) and Stage 1 (Diffusion) each use one GPU by default. For single-GPU setups, modify the stage configuration to share the same device. - -### Launch the Server - -```bash -# Use default configuration -vllm serve zai-org/GLM-Image --omni --port 8091 -``` - -Or use the convenience script: - -```bash -cd examples/online_serving/glm_image -bash run_server.sh -``` - -If you have a custom stage configs file: - -```bash -vllm serve zai-org/GLM-Image --omni --port 8091 --stage-configs-path /path/to/glm_image.yaml -``` - -### Send Requests - -Get into the glm_image folder: - -```bash -cd examples/online_serving/glm_image -``` - -Send request via Python: - -```bash -python openai_chat_client.py --prompt "A cute cat sitting on a window sill" -``` - -The Python client supports the following command-line arguments: - -- `--prompt` (or `-p`): Text prompt for generation (default: `A beautiful sunset over the ocean with sailing boats`) -- `--output` (or `-o`): Output file path (default: `glm_image_output.png`) -- `--server` (or `-s`): Server URL (default: `http://localhost:8091`) -- `--image` (or `-i`): Input image path (for image-to-image editing) -- `--height`: Image height in pixels (default: 1024) -- `--width`: Image width in pixels (default: 1024) -- `--steps`: Number of inference steps (default: 50) -- `--guidance-scale`: Classifier-free guidance scale (default: 1.5) -- `--seed`: Random seed (default: 42) -- `--negative`: Negative prompt - -## Modality Control - -GLM-Image supports **text-to-image** and **image-to-image** modes. - -The default yaml configuration deploys AR on GPU 0 and DiT on GPU 1. You can use the default configuration file: [`glm_image.yaml`](../../../vllm_omni/model_executor/stage_configs/glm_image.yaml) - -| Mode | Input | Output | Description | -| -------------- | ------------ | ------ | ---------------------------------- | -| Text-to-Image | Text | Image | Generate images from text prompts | -| Image-to-Image | Image + Text | Image | Edit images with text instructions | - -### Text-to-Image - -```bash -python openai_chat_client.py \ - --prompt "A photorealistic mountain landscape at sunset" \ - --height 1024 \ - --width 1024 \ - --output landscape.png - -# Or use the curl script: -bash run_curl_text_to_image.sh "A futuristic city skyline at night" -``` - -### Image-to-Image (Image Editing) - -```bash -python openai_chat_client.py \ - --prompt "Convert this image to watercolor style" \ - --image input.png \ - --output watercolor.png - -# Or use the curl script: -bash run_curl_image_edit.sh input.png "Convert to watercolor style" -``` - -For general-purpose request methods (curl, OpenAI SDK, Python `requests`), see -the [Text-to-Image](../text_to_image/README.md) and -[Image-to-Image](../image_to_image/README.md) guides. - -## Generation Parameters - -When using `/v1/chat/completions`, pass these inside `extra_body` in the curl -JSON, or via the `extra_body` keyword argument in the OpenAI Python SDK. -When using the dedicated `/v1/images/generations` or `/v1/images/edits` -endpoints, pass the supported generation controls as top-level fields directly. -For image dimensions and count, use `size` and `n` rather than `height` or -`width`. - -| Parameter | Type | Default | Description | -| --------------------- | ----- | ------- | ----------------------------------- | -| `height` | int | 1024 | Image height in pixels | -| `width` | int | 1024 | Image width in pixels | -| `num_inference_steps` | int | 50 | Number of diffusion denoising steps | -| `guidance_scale` | float | 1.5 | Classifier-free guidance scale | -| `seed` | int | None | Optional random seed; `/v1/images/*` generates one server-side if omitted | -| `negative_prompt` | str | None | Negative prompt | - -## Response Format - -```json -{ - "id": "chatcmpl-xxx", - "created": 1234567890, - "model": "zai-org/GLM-Image", - "choices": [ - { - "index": 0, - "message": { - "role": "assistant", - "content": [ - { - "type": "image_url", - "image_url": { - "url": "data:image/png;base64,..." - } - } - ] - }, - "finish_reason": "stop" - } - ], - "usage": {} -} -``` - -## Extract Image - -```bash -# From a saved JSON response -cat response.json | jq -r '.choices[0].message.content[0].image_url.url' | cut -d',' -f2- | base64 -d > output.png -``` - -## Architecture - -GLM-Image uses a 2-stage multistage pipeline: - -``` -Stage 0 (AR Model) Stage 1 (Diffusion) -┌───────────────────┐ ┌─────────────────────┐ -│ vLLM-optimized │ prior │ GlmImagePipeline │ -│ GlmImageFor │──tokens──►│ ┌───────────────┐ │ -│ Conditional │ │ │ DiT Denoiser │ │ -│ Generation │ │ └───────┬───────┘ │ -│ (9B AR model) │ │ ▼ │ -└───────────────────┘ │ ┌───────────────┐ │ - ▲ │ │ VAE Decode │──┼──► Image - │ │ └───────────────┘ │ - Text / Image └─────────────────────┘ - Input -``` - -## VRAM Requirements - -| Stage | VRAM | -| :---------------- | :--------------------- | -| Stage-0 (AR) | **~18 GiB + KV Cache** | -| Stage-1 (DiT+VAE) | **~20 GiB** | -| Total | **~38 GiB + KV Cache** | - -## File Description - -| File | Description | -| --------------------------- | ------------------------------------- | -| `run_server.sh` | Server startup script | -| `run_curl_text_to_image.sh` | Text-to-image curl example | -| `run_curl_image_edit.sh` | Image-to-image (editing) curl example | -| `openai_chat_client.py` | Python client (t2i + i2i) | - -## FAQ - -- If you encounter OOM errors, adjust `gpu_memory_utilization` in the stage config: - -```yaml -# In glm_image.yaml, reduce from default 0.6: -gpu_memory_utilization: 0.5 -``` - -- The first request may be slow due to model warmup. Subsequent requests will be faster. - -- If you encounter `Transformers does not recognize this architecture` error, your have to upgrade `transformers` package to `5.3.0` or above: - -``` -pip install --upgrade transformers -``` diff --git a/examples/online_serving/glm_image/openai_chat_client.py b/examples/online_serving/glm_image/openai_chat_client.py deleted file mode 100644 index e142b0719..000000000 --- a/examples/online_serving/glm_image/openai_chat_client.py +++ /dev/null @@ -1,172 +0,0 @@ -#!/usr/bin/env python3 -""" -GLM-Image OpenAI-compatible chat client for text-to-image and image-to-image. - -Usage: - # Text-to-image - python openai_chat_client.py --prompt "A cute cat" --output output.png - - # Image-to-image (image editing) - python openai_chat_client.py --prompt "Convert to watercolor style" --image input.png --output output.png -""" - -import argparse -import base64 -from pathlib import Path - -import requests - - -def generate_image( - prompt: str, - server_url: str = "http://localhost:8091", - image_path: str | None = None, - height: int = 1024, - width: int = 1024, - steps: int = 50, - guidance_scale: float = 1.5, - seed: int | None = None, - negative_prompt: str | None = None, -) -> bytes | None: - """Generate or edit an image using the chat completions API. - - Args: - prompt: Text description or editing instruction - server_url: Server URL - image_path: Path to input image (for image-to-image editing) - height: Image height in pixels - width: Image width in pixels - steps: Number of inference steps - guidance_scale: Classifier-free guidance scale - seed: Random seed for reproducibility - negative_prompt: Negative prompt - - Returns: - Image bytes or None if failed - """ - # Build message content - content: list[dict] = [{"type": "text", "text": prompt}] - - if image_path: - img_path = Path(image_path) - if not img_path.exists(): - print(f"Error: Image file not found: {image_path}") - return None - b64_data = base64.b64encode(img_path.read_bytes()).decode("utf-8") - suffix = img_path.suffix.lstrip(".").lower() - mime = {"jpg": "jpeg", "jpeg": "jpeg", "png": "png", "webp": "webp"}.get(suffix, "png") - content.append( - { - "type": "image_url", - "image_url": {"url": f"data:image/{mime};base64,{b64_data}"}, - } - ) - - messages = [{"role": "user", "content": content}] - - # Build request payload - extra_body: dict = { - "height": height, - "width": width, - "num_inference_steps": steps, - "guidance_scale": guidance_scale, - } - if seed is not None: - extra_body["seed"] = seed - if negative_prompt: - extra_body["negative_prompt"] = negative_prompt - - payload = {"messages": messages, "extra_body": extra_body} - - # Send request - try: - mode = "image-to-image" if image_path else "text-to-image" - print(f"Sending {mode} request to {server_url}...") - response = requests.post( - f"{server_url}/v1/chat/completions", - headers={"Content-Type": "application/json"}, - json=payload, - timeout=600, - ) - response.raise_for_status() - data = response.json() - - # Extract image from response - choices = data.get("choices", []) - for choice in choices: - choice_content = choice.get("message", {}).get("content") - if isinstance(choice_content, list): - for item in choice_content: - if isinstance(item, dict) and "image_url" in item: - img_url = item["image_url"].get("url", "") - if img_url.startswith("data:image"): - _, b64 = img_url.split(",", 1) - return base64.b64decode(b64) - - print(f"Unexpected response format: {data}") - return None - - except Exception as e: - print(f"Error: {e}") - return None - - -def main(): - parser = argparse.ArgumentParser(description="GLM-Image chat client") - parser.add_argument( - "--prompt", - "-p", - default="A beautiful sunset over the ocean with sailing boats", - help="Text prompt", - ) - parser.add_argument("--output", "-o", default="glm_image_output.png", help="Output file") - parser.add_argument("--server", "-s", default="http://localhost:8091", help="Server URL") - - # Image-to-image - parser.add_argument( - "--image", - "-i", - type=str, - help="Input image path (for image-to-image editing)", - ) - - # Generation parameters - parser.add_argument("--height", type=int, default=1024, help="Image height") - parser.add_argument("--width", type=int, default=1024, help="Image width") - parser.add_argument("--steps", type=int, default=50, help="Inference steps") - parser.add_argument("--guidance-scale", type=float, default=1.5, help="CFG guidance scale") - parser.add_argument("--seed", type=int, default=42, help="Random seed") - parser.add_argument("--negative", help="Negative prompt") - - args = parser.parse_args() - - mode = "image-to-image" if args.image else "text-to-image" - print(f"Mode: {mode}") - print(f"Prompt: {args.prompt}") - if args.image: - print(f"Input image: {args.image}") - - image_bytes = generate_image( - prompt=args.prompt, - server_url=args.server, - image_path=args.image, - height=args.height, - width=args.width, - steps=args.steps, - guidance_scale=args.guidance_scale, - seed=args.seed, - negative_prompt=args.negative, - ) - - if image_bytes: - output_path = Path(args.output) - output_path.write_bytes(image_bytes) - print(f"Image saved to: {output_path}") - print(f"Size: {len(image_bytes) / 1024:.1f} KB") - else: - print("Failed to generate image") - exit(1) - - -if __name__ == "__main__": - main() diff --git a/examples/online_serving/glm_image/run_curl_image_edit.sh b/examples/online_serving/glm_image/run_curl_image_edit.sh deleted file mode 100755 index bb1e851ba..000000000 --- a/examples/online_serving/glm_image/run_curl_image_edit.sh +++ /dev/null @@ -1,61 +0,0 @@ -#!/bin/bash -# GLM-Image image-edit (image-to-image) curl example - -set -euo pipefail - -if [[ $# -lt 2 ]]; then - echo "Usage: $0 \"\" [output_file]" >&2 - exit 1 -fi - -INPUT_IMG=$1 -PROMPT=$2 -SERVER="${SERVER:-http://localhost:8091}" -CURRENT_TIME=$(date +%Y%m%d%H%M%S) -OUTPUT="${3:-glm_image_i2i_${CURRENT_TIME}.png}" - -if [[ ! -f "$INPUT_IMG" ]]; then - echo "Input image not found: $INPUT_IMG" >&2 - exit 1 -fi - -# base64 encode (macOS uses -i, Linux uses -w0) -if [[ "$(uname)" == "Darwin" ]]; then - IMG_B64=$(base64 < "$INPUT_IMG" | tr -d '\n') -else - IMG_B64=$(base64 -w0 "$INPUT_IMG") -fi - -REQUEST_JSON=$( - jq -n --arg prompt "$PROMPT" --arg img "$IMG_B64" '{ - messages: [{ - role: "user", - content: [ - {"type": "text", "text": $prompt}, - {"type": "image_url", "image_url": {"url": ("data:image/png;base64," + $img)}} - ] - }], - extra_body: { - height: 1024, - width: 1024, - num_inference_steps: 50, - guidance_scale: 1.5, - seed: 42 - } - }' -) - -echo "Generating edited image..." -echo "Server: $SERVER" -echo "Prompt: $PROMPT" -echo "Input : $INPUT_IMG" -echo "Output: $OUTPUT" - -curl -s "$SERVER/v1/chat/completions" \ - -H "Content-Type: application/json" \ - -d "$REQUEST_JSON" \ - | jq -r '.choices[0].message.content[0].image_url.url' \ - | cut -d',' -f2- \ - | base64 -d > "$OUTPUT" - -echo "Image saved to: $OUTPUT" diff --git a/examples/online_serving/glm_image/run_curl_text_to_image.sh b/examples/online_serving/glm_image/run_curl_text_to_image.sh deleted file mode 100755 index aecb6953c..000000000 --- a/examples/online_serving/glm_image/run_curl_text_to_image.sh +++ /dev/null @@ -1,30 +0,0 @@ -#!/bin/bash -# GLM-Image text-to-image curl example - -set -euo pipefail - -PROMPT="${1:-A beautiful sunset over the ocean with sailing boats}" -SERVER="${SERVER:-http://localhost:8091}" -OUTPUT="${OUTPUT:-glm_image_t2i_output.png}" - -echo "Generating image..." -echo "Server: $SERVER" -echo "Prompt: $PROMPT" -echo "Output: $OUTPUT" - -curl -s "$SERVER/v1/chat/completions" \ - -H "Content-Type: application/json" \ - -d "{ - \"messages\": [ - {\"role\": \"user\", \"content\": \"$PROMPT\"} - ], - \"extra_body\": { - \"height\": 1024, - \"width\": 1024, - \"num_inference_steps\": 50, - \"guidance_scale\": 1.5, - \"seed\": 42 - } - }" | jq -r '.choices[0].message.content[0].image_url.url' | cut -d',' -f2- | base64 -d > "$OUTPUT" - -echo "Image saved to: $OUTPUT" diff --git a/examples/online_serving/glm_image/run_server.sh b/examples/online_serving/glm_image/run_server.sh deleted file mode 100755 index b47d9f885..000000000 --- a/examples/online_serving/glm_image/run_server.sh +++ /dev/null @@ -1,12 +0,0 @@ -#!/bin/bash -# GLM-Image online serving startup script - -MODEL="${MODEL:-zai-org/GLM-Image}" -PORT="${PORT:-8091}" - -echo "Starting GLM-Image server..." -echo "Model: $MODEL" -echo "Port: $PORT" - -vllm serve "$MODEL" --omni \ - --port "$PORT" diff --git a/vllm_omni/config/pipeline_registry.py b/vllm_omni/config/pipeline_registry.py index 69a004015..98aa13286 100644 --- a/vllm_omni/config/pipeline_registry.py +++ b/vllm_omni/config/pipeline_registry.py @@ -43,6 +43,10 @@ _OMNI_PIPELINES: dict[str, tuple[str, str]] = { "vllm_omni.model_executor.models.qwen3_tts.pipeline", "QWEN3_TTS_PIPELINE", ), + "glm_image": ( + "vllm_omni.model_executor.models.glm_image.pipeline", + "GLM_IMAGE_PIPELINE", + ), "voxcpm2": ( "vllm_omni.model_executor.models.voxcpm2.pipeline", "VOXCPM2_PIPELINE", diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index 9839a5389..6bd2faf7e 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -169,6 +169,12 @@ class StagePipelineConfig: prompt_expand_func: str | None = None cfg_kv_collect_func: str | None = None omni_kv_config: dict[str, Any] | None = None + # Model subdirectory indirections: for multi-component HF repos where the + # stage's config/tokenizer lives in a subdirectory (e.g. GLM-Image's AR + # config is in ``vision_language_encoder/``). Consumed at stage-init time + # by ``stage_init_utils._resolve_model_tokenizer_paths``. + model_subdir: str | None = None + tokenizer_subdir: str | None = None extras: dict[str, Any] = field(default_factory=dict) @@ -185,6 +191,11 @@ class PipelineConfig: # matches ``hf_config.architectures[*]`` against this tuple to route # to the correct pipeline. Leave empty for models with unique model_type. hf_architectures: tuple[str, ...] = () + # Diffusers pipeline class name: for models that ship a ``model_index.json`` + # (no root ``config.json``), the ``_class_name`` field is matched against + # this value to auto-detect the pipeline. Only needed for diffusers-style + # multi-component repos (e.g. GLM-Image). ``None`` = not a diffusers model. + diffusers_class_name: str | None = None def get_stage(self, stage_id: int) -> StagePipelineConfig | None: """Look up a stage by its ID.""" @@ -702,6 +713,13 @@ def _build_engine_args( engine_args["engine_output_type"] = ps.engine_output_type if next_stage_proc: engine_args["custom_process_next_stage_input_func"] = next_stage_proc + # Subdirectory indirections from StagePipelineConfig (structural, not + # deployment knobs). Deploy YAML ``engine_extras`` can still override + # these per-stage if needed. + if ps.model_subdir: + engine_args["model_subdir"] = ps.model_subdir + if ps.tokenizer_subdir: + engine_args["tokenizer_subdir"] = ps.tokenizer_subdir # Pipeline-wide top-level DeployConfig settings, applied to every stage. for name in _PIPELINE_WIDE_ENGINE_FIELDS: @@ -1378,6 +1396,27 @@ class StageConfigFactory: except Exception as e: logger.debug(f"Failed to auto-detect model type for {model}: {e}") + # Fallback for diffusers-style models: check model_index.json. + # Some models (e.g. GLM-Image) have no root config.json but ship a + # model_index.json with _class_name that maps to a pipeline key via + # PipelineConfig.diffusers_class_name. + try: + from vllm.transformers_utils.config import get_hf_file_to_dict + + model_index = get_hf_file_to_dict("model_index.json", model, revision=None) + if model_index and "_class_name" in model_index: + class_name = model_index["_class_name"] + for pipeline_cfg in _PIPELINE_REGISTRY.values(): + if pipeline_cfg.diffusers_class_name == class_name: + logger.info( + "Detected pipeline %r from model_index.json (_class_name=%r)", + pipeline_cfg.model_type, + class_name, + ) + return pipeline_cfg.model_type, None + except Exception: + pass + # Final fallback: some models (e.g. CosyVoice3) ship an empty # config.json and rely on naming conventions. Match the model path # basename against registered pipeline keys — longest match wins diff --git a/vllm_omni/deploy/glm_image.yaml b/vllm_omni/deploy/glm_image.yaml new file mode 100644 index 000000000..47cd02aa8 --- /dev/null +++ b/vllm_omni/deploy/glm_image.yaml @@ -0,0 +1,38 @@ +# GLM-Image deploy: AR (stage 0) + Diffusion (stage 1). +# Topology declared in vllm_omni/model_executor/models/glm_image/pipeline.py. +# +# Fields omitted from a stage fall back to StageDeployConfig dataclass +# defaults (see vllm_omni/config/stage_config.py). Pipeline-wide settings +# (trust_remote_code, distributed_executor_backend, etc.) use DeployConfig +# defaults unless overridden here. +async_chunk: false + +stages: + # Stage 0: AR Model (GlmImageForConditionalGeneration) + # Generates prior_token_ids for conditioning the diffusion process. + - stage_id: 0 + max_num_seqs: 1 + gpu_memory_utilization: 0.6 + enforce_eager: false + max_num_batched_tokens: 32768 + devices: "0" + default_sampling_params: + temperature: 0.9 + top_p: 0.75 + top_k: 16512 + stop_token_ids: [16385] + seed: 42 + detokenize: false + + # Stage 1: Diffusion (DiT + VAE) + # Receives prior_token_ids from AR, performs denoising + VAE decode. + - stage_id: 1 + max_num_seqs: 1 + enforce_eager: true + devices: "1" + default_sampling_params: + seed: 42 + num_inference_steps: 50 + guidance_scale: 1.5 + height: 1024 + width: 1024 diff --git a/vllm_omni/model_executor/models/glm_image/__init__.py b/vllm_omni/model_executor/models/glm_image/__init__.py index d37044c09..6153f07c0 100644 --- a/vllm_omni/model_executor/models/glm_image/__init__.py +++ b/vllm_omni/model_executor/models/glm_image/__init__.py @@ -1,3 +1,16 @@ -from .glm_image_ar import GlmImageForConditionalGeneration +def __getattr__(name: str): + """Lazy import to avoid importing transformers.models.glm_image at module init. + + The AR model depends on ``transformers.models.glm_image`` which is only + available when the model weights are present. Importing it eagerly in + ``__init__.py`` breaks the pipeline registry lookup for environments that + don't have this custom transformers extension installed. + """ + if name == "GlmImageForConditionalGeneration": + from .glm_image_ar import GlmImageForConditionalGeneration + + return GlmImageForConditionalGeneration + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + __all__ = ["GlmImageForConditionalGeneration"] diff --git a/vllm_omni/model_executor/models/glm_image/pipeline.py b/vllm_omni/model_executor/models/glm_image/pipeline.py new file mode 100644 index 000000000..bd041292c --- /dev/null +++ b/vllm_omni/model_executor/models/glm_image/pipeline.py @@ -0,0 +1,47 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""GLM-Image pipeline topologies (frozen). +Two-stage (default): + Stage 0: AR — multimodal understanding + token_ids generation + Stage 1: DiT — diffusion image generation +""" + +from vllm_omni.config.stage_config import ( + PipelineConfig, + StageExecutionType, + StagePipelineConfig, +) + +GLM_IMAGE_PIPELINE = PipelineConfig( + model_type="glm_image", + model_arch="GlmImageForConditionalGeneration", + hf_architectures=("GlmImageForConditionalGeneration",), + diffusers_class_name="GlmImagePipeline", + stages=( + StagePipelineConfig( + stage_id=0, + model_stage="ar", + execution_type=StageExecutionType.LLM_AR, + requires_multimodal_data=True, + input_sources=(), + final_output=False, + owns_tokenizer=True, + model_arch="GlmImageForConditionalGeneration", + engine_output_type="token_ids", + model_subdir="vision_language_encoder", + tokenizer_subdir="processor", + ), + StagePipelineConfig( + stage_id=1, + model_stage="dit", + execution_type=StageExecutionType.DIFFUSION, + input_sources=(0,), + requires_multimodal_data=True, + final_output=True, + final_output_type="image", + model_arch="GlmImagePipeline", + custom_process_input_func="vllm_omni.model_executor.stage_input_processors.glm_image.ar2diffusion", + omni_kv_config={"need_recv_cache": False}, + ), + ), +) diff --git a/vllm_omni/model_executor/stage_configs/glm_image.yaml b/vllm_omni/model_executor/stage_configs/glm_image.yaml deleted file mode 100644 index f3ed6c721..000000000 --- a/vllm_omni/model_executor/stage_configs/glm_image.yaml +++ /dev/null @@ -1,74 +0,0 @@ -# Stage config for running GLM-Image with 2-stage architecture -# Stage 0: AR Model (vLLM implementation) - generates prior_token_ids -# Stage 1: Diffusion (DiT + VAE) - denoising and image decoding - -stage_args: - # Stage 0: AR Model (GlmImageForConditionalGeneration) - # This stage uses the vLLM-optimized AR model to generate prior tokens - # for conditioning the diffusion process. - - stage_id: 0 - stage_type: llm - runtime: - process: true - devices: "0" - requires_multimodal_data: true # Required for i2i mode to receive source images - engine_args: - model_stage: ar - max_num_seqs: 1 - model_arch: GlmImageForConditionalGeneration - model_subdir: vision_language_encoder # AR model config.json is in this subdirectory - tokenizer_subdir: processor # Use processor's tokenizer (not ByT5 from tokenizer/) - worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker - scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler - gpu_memory_utilization: 0.6 - enforce_eager: false - trust_remote_code: true - engine_output_type: token_ids # Output prior_token_ids for diffusion stage - distributed_executor_backend: "mp" - enable_prefix_caching: false - max_num_batched_tokens: 32768 - final_output: false # AR is not the final output - is_comprehension: true - default_sampling_params: - temperature: 0.9 # From model's generation_config.json - top_p: 0.75 # From model's generation_config.json - top_k: 16512 # vision_vocab_size from generation_config.json - stop_token_ids: [16385] # eos_token_id from generation_config.json - seed: 42 - detokenize: false - - # Stage 1: Diffusion (DiT + VAE) - # This stage receives prior_token_ids from AR and performs denoising + VAE decode - - stage_id: 1 - stage_type: diffusion - runtime: - process: true - devices: "1" # Can use different GPU, or same GPU if memory allows - requires_multimodal_data: true # Required for i2i mode to pass condition images - engine_args: - model_stage: dit - max_num_seqs: 1 - model_arch: GlmImagePipeline # Required for diffusion model class resolution - # Diffusion-specific parameters - num_gpus: 1 - enforce_eager: true - trust_remote_code: true - distributed_executor_backend: "mp" - engine_input_source: [0] # Input from AR stage - custom_process_input_func: vllm_omni.model_executor.stage_input_processors.glm_image.ar2diffusion - final_output: true - final_output_type: image - default_sampling_params: - # Diffusion-specific parameters only (no LLM params like temperature/top_p/top_k) - seed: 42 - num_inference_steps: 50 - guidance_scale: 1.5 - height: 1024 - width: 1024 - -# Top-level runtime config -runtime: - enabled: true - edges: - - from: 0 # AR → Diffusion: trigger after AR completes - to: 1 diff --git a/vllm_omni/model_executor/stage_configs/glm_image_muilticonnector.yaml b/vllm_omni/model_executor/stage_configs/glm_image_muilticonnector.yaml deleted file mode 100644 index 2a85a6dad..000000000 --- a/vllm_omni/model_executor/stage_configs/glm_image_muilticonnector.yaml +++ /dev/null @@ -1,87 +0,0 @@ -# Stage config for running GLM-Image with 2-stage architecture (MultiConnector version) -# Stage 0: AR Model (vLLM implementation) - generates prior_token_ids -# Stage 1: Diffusion (DiT + VAE) - denoising and image decoding -# -# This config uses OmniConnectors for inter-stage communication, -# enabling efficient tensor transfer between stages on different processes/nodes. - -stage_args: - # Stage 0: AR Model (GlmImageForConditionalGeneration) - # This stage uses the vLLM-optimized AR model to generate prior tokens - # for conditioning the diffusion process. - - stage_id: 0 - stage_type: llm - runtime: - process: true - devices: "0" - requires_multimodal_data: true # Required for i2i mode to receive source images - engine_args: - model_stage: ar - max_num_seqs: 1 - model_arch: GlmImageForConditionalGeneration - worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker - scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler - gpu_memory_utilization: 0.6 - enforce_eager: false - trust_remote_code: true - engine_output_type: token_ids # Output prior_token_ids for diffusion stage - distributed_executor_backend: "mp" - enable_prefix_caching: false - max_num_batched_tokens: 32768 - hf_config_name: vision_language_encoder # Subfolder in model path - final_output: false # AR is not the final output - is_comprehension: true - default_sampling_params: - temperature: 0.9 # From model's generation_config.json - top_p: 0.75 # From model's generation_config.json - top_k: 16512 # vision_vocab_size from generation_config.json - stop_token_ids: [16385] # eos_token_id from generation_config.json - seed: 42 - detokenize: false - - # Stage 1: Diffusion (DiT + VAE) - # This stage receives prior_token_ids from AR and performs denoising + VAE decode - - stage_id: 1 - stage_type: diffusion - runtime: - process: true - devices: "1" # Use separate GPU for diffusion - requires_multimodal_data: true # Required for i2i mode to pass condition images - engine_args: - model_stage: dit - max_num_seqs: 1 - # Diffusion-specific parameters - num_gpus: 1 - enforce_eager: true - trust_remote_code: true - distributed_executor_backend: "mp" - engine_input_source: [0] # Input from AR stage - custom_process_input_func: vllm_omni.model_executor.stage_input_processors.glm_image.ar2diffusion - final_output: true - final_output_type: image - default_sampling_params: - seed: 42 - num_inference_steps: 50 - guidance_scale: 1.5 - height: 1024 - width: 1024 - -# Top-level runtime config with MultiConnector support -runtime: - enabled: true - edges: - - from: 0 # AR → Diffusion - to: 1 - -# OmniConnector configuration for efficient inter-stage tensor transfer -connectors: - - type: tensor_transfer - source_stage: 0 - target_stage: 1 - # Transfer prior_token_ids efficiently between stages - fields: - - name: prior_token_ids - dtype: int64 - - name: prior_token_image_ids - dtype: int64 - optional: true From e18cb89d5a8b6c306c84f753fca71802faac4b57 Mon Sep 17 00:00:00 2001 From: wangyu <53896905+yenuo26@users.noreply.github.com> Date: Wed, 22 Apr 2026 15:14:20 +0800 Subject: [PATCH 08/45] [CI] Update test markers and configurations to use 'full_model' for L4 nightly tests (#2641) Signed-off-by: wangyu <410167048@qq.com> --- .buildkite/test-nightly.yml | 24 +++--- .buildkite/test-ready.yml | 17 ---- benchmarks/accuracy/README.md | 4 +- benchmarks/accuracy/image_to_image/README.md | 4 +- docs/contributing/ci/CI_5levels.md | 10 ++- .../l4_functionality_tests.inc.md | 2 +- docs/contributing/ci/test_guide.md | 33 +++++--- docs/contributing/ci/tests_markers.md | 3 +- pyproject.toml | 3 +- tests/benchmarks/test_accuracy_bench_utils.py | 5 +- tests/dfx/perf/scripts/run_benchmark.py | 4 + .../perf/scripts/run_diffusion_benchmark.py | 5 +- .../models/glm_image/test_glm_image_sp.py | 5 +- .../test_omni_coord_client_for_stage.py | 2 - tests/e2e/accuracy/test_gebench_h100_smoke.py | 4 +- .../accuracy/test_gedit_bench_h100_smoke.py | 4 +- tests/e2e/accuracy/test_qwen_image.py | 5 +- tests/e2e/accuracy/test_qwen_image_edit.py | 7 +- tests/e2e/accuracy/test_qwen_image_layered.py | 5 +- .../test_wan22_i2v_video_similarity.py | 8 +- .../test_quantization_fp8.py | 5 +- .../test_stable_audio_expansion.py | 4 +- .../online_serving/test_bagel_expansion.py | 4 +- .../test_dynin_omni_expansion.py | 79 ++----------------- .../online_serving/test_flux2_expansion.py | 4 +- .../test_flux2_klein_inpaint_expansion.py | 8 +- .../test_flux_2_dev_expansion.py | 4 +- .../test_flux_kontext_expansion.py | 12 +-- .../test_hunyuan_video_15_expansion.py | 4 +- .../test_longcat_image_edit_expansion.py | 4 +- .../test_longcat_image_expansion.py | 4 +- .../online_serving/test_nextstep_expansion.py | 4 +- .../test_qwen3_omni_expansion.py | 38 +-------- .../test_qwen3_tts_base_expansion.py | 14 ++-- .../test_qwen3_tts_customvoice_expansion.py | 16 ++-- .../test_qwen_image_edit_expansion.py | 6 +- .../test_qwen_image_expansion.py | 6 +- .../test_qwen_image_layered_expansion.py | 8 +- .../e2e/online_serving/test_sd3_expansion.py | 4 +- .../online_serving/test_wan22_expansion.py | 4 +- .../test_wan_2_1_vace_expansion.py | 4 +- .../online_serving/test_zimage_expansion.py | 4 +- .../offline_inference/test_text_to_image.py | 3 +- .../online_serving/test_qwen2_5_omni.py | 21 +---- .../online_serving/test_qwen3_omni.py | 21 +---- .../online_serving/test_text_to_image.py | 2 +- tests/helpers/assertions.py | 38 ++++++--- tests/helpers/fixtures/run_args.py | 4 +- tests/helpers/fixtures/runtime.py | 2 +- tests/helpers/runtime.py | 10 --- 50 files changed, 173 insertions(+), 322 deletions(-) diff --git a/.buildkite/test-nightly.yml b/.buildkite/test-nightly.yml index c33d7b4d1..e6912bd49 100644 --- a/.buildkite/test-nightly.yml +++ b/.buildkite/test-nightly.yml @@ -13,7 +13,7 @@ steps: - label: ":full_moon: Omni · Function Test" timeout_in_minutes: 90 commands: - - pytest -s -v tests/e2e/online_serving/test_*_expansion.py -m "advanced_model and H100 and omni" --run-level "advanced_model" + - pytest -s -v tests/e2e/ -m "full_model and H100 and omni" --run-level "full_model" --ignore=tests/e2e/accuracy agents: queue: "mithril-h100-pool" plugins: @@ -52,7 +52,7 @@ steps: timeout_in_minutes: 90 commands: - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1" - - pytest -s -v tests/examples/ -m "advanced_model and omni and L4" --run-level "advanced_model" + - pytest -s -v tests/examples/ -m "full_model and omni and L4" --run-level "full_model" agents: queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU plugins: @@ -70,7 +70,7 @@ steps: - label: ":full_moon: Omni · Doc Test with H100" timeout_in_minutes: 90 commands: - - pytest -s -v tests/examples/ -m "advanced_model and omni and H100" --run-level "advanced_model" + - pytest -s -v tests/examples/ -m "full_model and omni and H100" --run-level "full_model" agents: queue: "mithril-h100-pool" plugins: @@ -160,7 +160,7 @@ steps: timeout_in_minutes: 90 commands: - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1" - - pytest -s -v tests/e2e/online_serving/test_*_expansion.py -m "advanced_model and L4 and omni" --run-level "advanced_model" + - pytest -s -v tests/e2e/ -m "full_model and L4 and omni" --run-level "full_model" --ignore=tests/e2e/accuracy agents: queue: "gpu_1_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU plugins: @@ -233,7 +233,7 @@ steps: - label: ":full_moon: Diffusion X2I(&A&T) · Function Test with H100" timeout_in_minutes: 120 commands: - - pytest -s -v tests/e2e/online_serving/test_*_expansion.py -k "not test_wan22_expansion and not test_wan_2_1_vace_expansion and not hunyuan" -m "advanced_model and diffusion and H100" --run-level "advanced_model" + - pytest -sv tests/e2e/ -k "not test_wan and not hunyuan" -m "full_model and diffusion and H100" --run-level "full_model" --ignore=tests/e2e/accuracy agents: queue: "mithril-h100-pool" plugins: @@ -271,7 +271,7 @@ steps: - label: ":full_moon: Diffusion X2I(&A&T) · Function Test with L4" timeout_in_minutes: 60 commands: - - pytest -s -v tests/e2e/online_serving/test_*_expansion.py -k "not test_wan22_expansion and not test_wan_2_1_vace_expansion and not hunyuan" -m "advanced_model and diffusion and L4" --run-level "advanced_model" + - pytest -sv tests/e2e/ -k "not test_wan and not hunyuan" -m "full_model and diffusion and L4" --run-level "full_model" --ignore=tests/e2e/accuracy agents: queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU plugins: @@ -290,7 +290,7 @@ steps: timeout_in_minutes: 60 commands: - export VLLM_TEST_CLEAN_GPU_MEMORY="1" - - pytest -s -v tests/examples/*/test_text_to_image.py -m "advanced_model and example and H100" --run-level "advanced_model" + - pytest -s -v tests/examples/*/test_text_to_image.py -m "full_model and example and H100" --run-level "full_model" agents: queue: "mithril-h100-pool" plugins: @@ -328,7 +328,7 @@ steps: - label: ":full_moon: Diffusion X2I(&A&T) · GEBench Accuracy Test" timeout_in_minutes: 60 commands: - - pytest -s -v tests/e2e/accuracy/test_gebench_h100_smoke.py --run-level advanced_model --gebench-model Qwen/Qwen-Image-2512 --accuracy-judge-model QuantTrio/Qwen3-VL-30B-A3B-Instruct-AWQ --accuracy-gpu 0 --gebench-port 8093 --accuracy-workers 1 + - pytest -s -v tests/e2e/accuracy/test_gebench_h100_smoke.py --run-level full_model --gebench-model Qwen/Qwen-Image-2512 --accuracy-judge-model QuantTrio/Qwen3-VL-30B-A3B-Instruct-AWQ --accuracy-gpu 0 --gebench-port 8093 --accuracy-workers 1 - buildkite-agent artifact upload "tests/e2e/accuracy/artifacts/gebench_qwen-image-2512/summary*.json" agents: queue: "mithril-h100-pool" @@ -367,7 +367,7 @@ steps: - label: ":full_moon: Diffusion X2I(&A&T) · GEdit-Bench Accuracy Test" timeout_in_minutes: 60 commands: - - pytest -s -v tests/e2e/accuracy/test_gedit_bench_h100_smoke.py --run-level advanced_model --gedit-model Qwen/Qwen-Image-Edit --accuracy-judge-model QuantTrio/Qwen3-VL-30B-A3B-Instruct-AWQ --accuracy-gpu 0 --gedit-port 8093 --gedit-samples-per-group 20 --accuracy-workers 1 + - pytest -s -v tests/e2e/accuracy/test_gedit_bench_h100_smoke.py --run-level full_model --gedit-model Qwen/Qwen-Image-Edit --accuracy-judge-model QuantTrio/Qwen3-VL-30B-A3B-Instruct-AWQ --accuracy-gpu 0 --gedit-port 8093 --gedit-samples-per-group 20 --accuracy-workers 1 - buildkite-agent artifact upload "tests/e2e/accuracy/artifacts/gedit_scores_qwen-image-edit/qwen-image-edit_all_all_vie_score_*.csv" - buildkite-agent artifact upload "tests/e2e/accuracy/artifacts/gedit_scores_qwen-image-edit/qwen-image-edit_all_all_summary_*.json" agents: @@ -409,7 +409,7 @@ steps: - label: ":full_moon: Diffusion X2I(&A&T) · Accuracy Test" timeout_in_minutes: 180 commands: - - pytest -s -v tests/e2e/accuracy/test_qwen_image*.py --run-level advanced_model + - pytest -s -v tests/e2e/accuracy/test_qwen_image*.py --run-level full_model agents: queue: "mithril-h100-pool" plugins: @@ -514,7 +514,7 @@ steps: - label: ":full_moon: Diffusion X2V · Function Test" timeout_in_minutes: 90 commands: - - pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py tests/e2e/online_serving/test_wan_2_1_vace_expansion.py tests/e2e/online_serving/test_hunyuan_video_15_expansion.py -m "advanced_model" --run-level "advanced_model" + - pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py tests/e2e/online_serving/test_wan_2_1_vace_expansion.py tests/e2e/online_serving/test_hunyuan_video_15_expansion.py -m "full_model" --run-level "full_model" agents: queue: "mithril-h100-pool" plugins: @@ -552,7 +552,7 @@ steps: - label: ":full_moon: Diffusion X2V · Accuracy Test" timeout_in_minutes: 180 commands: - - pytest -s -v tests/e2e/accuracy/wan22_i2v/test_wan22_i2v_video_similarity.py -m advanced_model --run-level advanced_model + - pytest -s -v tests/e2e/accuracy/wan22_i2v/test_wan22_i2v_video_similarity.py -m full_model --run-level full_model agents: queue: "mithril-h100-pool" plugins: diff --git a/.buildkite/test-ready.yml b/.buildkite/test-ready.yml index 3ca1747fe..080f18885 100644 --- a/.buildkite/test-ready.yml +++ b/.buildkite/test-ready.yml @@ -120,23 +120,6 @@ steps: volumes: - "/fsx/hf_cache:/fsx/hf_cache" - - label: "Audio Generation Model Test" - depends_on: upload-ready-pipeline - commands: - - timeout 20m pytest -s -v tests/e2e/offline_inference/test_stable_audio_expansion.py -m "advanced_model and diffusion and L4" --run-level advanced_model - agents: - queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU - plugins: - - docker#v5.2.0: - image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT - always-pull: true - propagate-environment: true - environment: - - "HF_HOME=/fsx/hf_cache" - - "HF_TOKEN" - volumes: - - "/fsx/hf_cache:/fsx/hf_cache" - - label: "Diffusion Cache Backend Test" depends_on: upload-ready-pipeline commands: diff --git a/benchmarks/accuracy/README.md b/benchmarks/accuracy/README.md index 0d73215b6..dbe20916a 100644 --- a/benchmarks/accuracy/README.md +++ b/benchmarks/accuracy/README.md @@ -23,5 +23,5 @@ Test guidance: - Local static/self-checks live in `tests/benchmarks/test_accuracy_bench_utils.py`. - End-to-end generation/evaluation should be validated in a remote GPU environment. In the current repo marker system there is `L4` but no `L5` - marker, so benchmark smoke tests should be wired as `advanced_model + - benchmark + L4` when GPU capacity is available. + marker, so benchmark smoke tests should be wired as `full_model + + benchmark + L4` for nightly when GPU capacity is available. diff --git a/benchmarks/accuracy/image_to_image/README.md b/benchmarks/accuracy/image_to_image/README.md index ee1d58f10..86e7b0cf3 100644 --- a/benchmarks/accuracy/image_to_image/README.md +++ b/benchmarks/accuracy/image_to_image/README.md @@ -99,5 +99,5 @@ Notes: - This flow requires the optional Hugging Face `datasets` package. - `generate` writes `generation_manifest.json` with local output coverage. - The current repo marker set exposes `L4` but not `L5`, so if you promote an - end-to-end smoke test into CI, use the existing `advanced_model`, `benchmark`, - and `L4` markers or introduce a new repo-wide marker explicitly first. + end-to-end smoke test into CI, use the `full_model`, `benchmark`, + and `L4` markers for nightly (or `advanced_model` for merge) or introduce a new repo-wide marker explicitly first. diff --git a/docs/contributing/ci/CI_5levels.md b/docs/contributing/ci/CI_5levels.md index 2452ef5d4..7bad897e7 100644 --- a/docs/contributing/ci/CI_5levels.md +++ b/docs/contributing/ci/CI_5levels.md @@ -418,13 +418,13 @@ L3 level testing executes after code is merged into the main branch. Its core pu **Explanation**: - @pytest.mark.advanced_model: Marks the test as L3 or L4 level, indicating that this test case performs deep validation, using real models for performance, integration, and accuracy testing. This forms a "basic-advanced" correspondence with the core_model mark at the L2 level. + @pytest.mark.advanced_model: Marks the test as L3 merge level, indicating deep validation with real models. @pytest.mark.full_model: Marks L4 nightly-only suites (e.g. `test_*_expansion.py`, doc examples). @pytest.mark.core_model: Marks the test as L1 or L2 level, indicating that this test case validates the basic functionality of the core model. It uses mock weights and only checks if the relevant interface functions correctly. @pytest.mark.parametrize: A parameterization decorator that allows abstracting test data into parameters, enabling reuse of the same test logic across different data configurations. indirect=True indicates that parameters will be passed to the fixture for processing. - **Notes**: If you believe the test case only needs to execute basic run logic at the PR-level CI, you can mark it only with @pytest.mark.core_model. If you believe it only needs to execute deep validation run logic at the merge or nightly level, you can mark it only with @pytest.mark.advanced_model. If you believe the test case needs to accommodate both basic run and deep validation test logic, you should mark it with both @pytest.mark.core_model and @pytest.mark.advanced_model. + **Notes**: If you believe the test case only needs to execute basic run logic at the PR-level CI, you can mark it only with @pytest.mark.core_model. If you believe it only needs to execute deep validation at merge (L3), use @pytest.mark.advanced_model. For L4 nightly-only expansion and doc-example tests, use @pytest.mark.full_model with `--run-level full_model`. If the test case needs both basic run and deep validation, mark with @pytest.mark.core_model and the appropriate L3/L4 marker (`advanced_model` and/or `full_model`). **2.4.2 Test Function Definition and Documentation** @@ -516,9 +516,11 @@ L3 level testing executes after code is merged into the main branch. Its core pu **Single Request**: The comment clearly states this is a single-request completion test. For concurrent testing, it can be extended to multiple requests using request_num = n. - **Implicit Validation**: The `send_omni_request` and `send_diffusion_request` methods internally includes validation logic dynamically selected based on the --run-level parameter: core_model performs basic validation, while advanced_model performs deep validation. + **Implicit Validation**: The `send_omni_request` and `send_diffusion_request` methods internally includes validation logic dynamically selected based on the --run-level parameter: core_model performs basic validation, while advanced_model and full_model perform deep validation. -- ***Run Command***: `pytest -s -v /tests/e2e/online_serving/test_{model_name}.py -m advanced_model --run-level=advanced_model` +- ***Run Command (L3 merge)***: `pytest -s -v /tests/e2e/online_serving/test_{model_name}.py -m advanced_model --run-level=advanced_model` + +- ***Run Command (L4 nightly expansion)***: `pytest -s -v /tests/e2e/online_serving/test_{model_name}_expansion.py -m full_model --run-level=full_model` ## Chapter 3: L4 Level Testing - Full Functionality, Performance, and Documentation Testing diff --git a/docs/contributing/ci/test_examples/l4_functionality_tests.inc.md b/docs/contributing/ci/test_examples/l4_functionality_tests.inc.md index ab4deecd6..e1309b1ad 100644 --- a/docs/contributing/ci/test_examples/l4_functionality_tests.inc.md +++ b/docs/contributing/ci/test_examples/l4_functionality_tests.inc.md @@ -37,7 +37,7 @@ Currently all the features are available in online serving mode. Hence, only nee **Code Style** - Validation: test that the multimodal output files of your model have the correct shapes. `OpenAIClientHandler.send_diffusion_request` should have taken care of this. -- Test marks: always add `advanced_model` and `diffusion`. Add GPU-related marks if needed. Ref: [Markers for Tests](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/ci/tests_markers/). +- Test marks: always add `full_model` and `diffusion` for L4 nightly `test_*_expansion.py` cases. Add GPU-related marks if needed. Ref: [Markers for Tests](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/ci/tests_markers/). - To maximize code reuse, you may refer to - `tests/conftest.py` for `omni_server` (running server in subprocess) and `openai_client` fixtures (sending requests and validating output), `generate_synthetic_image` and `assert_XXX_valid` helper. - `tests/helpers/mark.py` for `@hardware_test(...)` and `hardware_marks`. diff --git a/docs/contributing/ci/test_guide.md b/docs/contributing/ci/test_guide.md index 08b2e3b4e..0c9ec2fb3 100644 --- a/docs/contributing/ci/test_guide.md +++ b/docs/contributing/ci/test_guide.md @@ -42,31 +42,46 @@ Our test scripts use the pytest framework. First, please use `git clone https:// ``` The latest test commands for various test suites can be found in the [pipeline](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-ready.yml). -=== "L3 level & L4 level" +=== "L3 level" ```bash pytest -s -v -m "advanced_model" --run-level=advanced_model ``` - If you only want to run L3 test case, you can use: - ```bash - pytest -s -v e2e/ --ignore-glob='*expansion.py' -m "advanced_model" --run-level=advanced_model - ``` If you only want to run a specific test case, you can use: ```bash pytest -s -v test_xxxx.py --run-level=advanced_model ``` If you only want to run specific test cases on a particular platform, you can use: ```bash - pytest -s -v -m "core_model and distributed_cuda and L4" --run-level=core_model + pytest -s -v -m "advanced_model and distributed_cuda and L4" --run-level=advanced_model + ``` + The latest L3 test commands for various test suites can be found in the [pipeline](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-merge.yml). + + +=== "L4 level" + + ```bash + cd tests + pytest -s -v -m "full_model" --run-level=full_model + ``` + If you only want to run a specific test case, you can use: + ```bash + pytest -s -v test_xxxx.py --run-level=full_model + ``` + If you only want to run specific test cases on a particular platform, you can use: + ```bash + pytest -s -v -m "full_model and distributed_cuda and L4" --run-level=full_model ``` Note: To run performance tests (defaults to ``test_qwen_omni.json``; use ``--test-config-file tests/dfx/perf/tests/test_tts.json`` for TTS): ```bash pytest -s -v tests/dfx/perf/scripts/run_benchmark.py ``` + The latest L4 (nightly) test commands use the `full_model` marker and `--run-level full_model` (see [test-nightly.yml](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-nightly.yml) and [test-nightly-diffusion.yml](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-nightly-diffusion.yml)). Example: - The latest L3 test commands for various test suites can be found in the [pipeline](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-merge.yml). - - The latest L4 test commands for various test suites can be found in the [pipeline](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-nightly.yml). + ```bash + cd tests + pytest -s -v -m "full_model and omni and H100" --run-level=full_model + ``` You can find more information about markers in the documentation: [marker doc](./tests_markers.md) diff --git a/docs/contributing/ci/tests_markers.md b/docs/contributing/ci/tests_markers.md index 7628db284..6130541a6 100644 --- a/docs/contributing/ci/tests_markers.md +++ b/docs/contributing/ci/tests_markers.md @@ -8,7 +8,8 @@ Defined in `pyproject.toml`: | Marker | Description | | ------------------ | --------------------------------------------------------- | | `core_model` | L1&L2 tests (run in each PR) | -| `advanced_model` | L3&L4 level tests (run in each merge or nightly) | +| `advanced_model` | L3 tests (run on each merge to main) | +| `full_model` | L4 tests (run nightly) | | `diffusion` | Diffusion model tests | | `omni` | Omni model tests | | `cache` | Cache backend tests | diff --git a/pyproject.toml b/pyproject.toml index b6f4092fd..0958b9395 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -179,7 +179,8 @@ addopts = [ markers = [ # ci/cd required "core_model: L1&L2 tests (run in each PR)", - "advanced_model: L3&L4 level tests (run in each merge or nightly)", + "advanced_model: L3 level tests (run on each merge)", + "full_model: L4 level tests (run nightly)", # function module markers "diffusion: Diffusion model tests", "omni: Omni model tests", diff --git a/tests/benchmarks/test_accuracy_bench_utils.py b/tests/benchmarks/test_accuracy_bench_utils.py index a0479fb1b..8425804ae 100644 --- a/tests/benchmarks/test_accuracy_bench_utils.py +++ b/tests/benchmarks/test_accuracy_bench_utils.py @@ -6,6 +6,9 @@ from pathlib import Path import pytest from PIL import Image +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + + REPO_ROOT = Path(__file__).resolve().parents[2] if str(REPO_ROOT) not in sys.path: sys.path.insert(0, str(REPO_ROOT)) @@ -37,8 +40,6 @@ from benchmarks.accuracy.text_to_image.gbench import ( summarize_gebench_results, ) -pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] - def test_summarize_gebench_generated_records_groups_by_type(): records = [ diff --git a/tests/dfx/perf/scripts/run_benchmark.py b/tests/dfx/perf/scripts/run_benchmark.py index 13011b4bd..f14638c15 100644 --- a/tests/dfx/perf/scripts/run_benchmark.py +++ b/tests/dfx/perf/scripts/run_benchmark.py @@ -17,6 +17,9 @@ from tests.dfx.conftest import ( ) from tests.helpers.runtime import OmniServer +pytestmark = [pytest.mark.full_model, pytest.mark.omni] + + os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0" @@ -293,6 +296,7 @@ def assert_result( print(f"ERROR: Test results exceeded baseline: {metric_name}: {current_value} < {baseline_value}") +@pytest.mark.benchmark @pytest.mark.parametrize("omni_server", test_params, indirect=True) @pytest.mark.parametrize("benchmark_params", benchmark_indices, indirect=True) def test_performance_benchmark(omni_server, benchmark_params): diff --git a/tests/dfx/perf/scripts/run_diffusion_benchmark.py b/tests/dfx/perf/scripts/run_diffusion_benchmark.py index 23efa8bb0..d41f80dc7 100644 --- a/tests/dfx/perf/scripts/run_diffusion_benchmark.py +++ b/tests/dfx/perf/scripts/run_diffusion_benchmark.py @@ -32,6 +32,8 @@ from typing import Any, cast import psutil import pytest +pytestmark = [pytest.mark.diffusion, pytest.mark.full_model] + os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0" os.environ.setdefault("DIFFUSION_ATTENTION_BACKEND", "FLASH_ATTN") @@ -662,8 +664,7 @@ def assert_result(result: dict[str, Any], params: dict[str, Any]) -> None: # --------------------------------------------------------------------------- # Test entry point # --------------------------------------------------------------------------- - - +@pytest.mark.benchmark @pytest.mark.parametrize( "diffusion_server", server_params, diff --git a/tests/diffusion/models/glm_image/test_glm_image_sp.py b/tests/diffusion/models/glm_image/test_glm_image_sp.py index 40d1c8730..06a1a116d 100644 --- a/tests/diffusion/models/glm_image/test_glm_image_sp.py +++ b/tests/diffusion/models/glm_image/test_glm_image_sp.py @@ -6,6 +6,8 @@ import pytest from vllm_omni.diffusion.data import DiffusionParallelConfig +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + @pytest.fixture(scope="function", autouse=True) def setup_sp_groups(mocker): @@ -24,9 +26,6 @@ def setup_sp_groups(mocker): yield -pytestmark = [pytest.mark.core_model, pytest.mark.cpu] - - def test_glm_image_sp_plan_defined(): """Test that _sp_plan is properly defined on GlmImageTransformer2DModel.""" from vllm_omni.diffusion.models.glm_image.glm_image_transformer import ( diff --git a/tests/distributed/omni_coordinator/test_omni_coord_client_for_stage.py b/tests/distributed/omni_coordinator/test_omni_coord_client_for_stage.py index f095dfd49..0ba19c7ff 100644 --- a/tests/distributed/omni_coordinator/test_omni_coord_client_for_stage.py +++ b/tests/distributed/omni_coordinator/test_omni_coord_client_for_stage.py @@ -17,8 +17,6 @@ from vllm_omni.distributed.omni_coordinator import ( pytestmark = [pytest.mark.core_model, pytest.mark.cpu] -pytestmark = [pytest.mark.core_model, pytest.mark.cpu] - def _bind_router() -> tuple[zmq.Context, zmq.Socket, str]: ctx = zmq.Context.instance() diff --git a/tests/e2e/accuracy/test_gebench_h100_smoke.py b/tests/e2e/accuracy/test_gebench_h100_smoke.py index 748919269..2702710e4 100644 --- a/tests/e2e/accuracy/test_gebench_h100_smoke.py +++ b/tests/e2e/accuracy/test_gebench_h100_smoke.py @@ -9,10 +9,10 @@ from benchmarks.accuracy.text_to_image.gbench import main as gbench_main from tests.e2e.accuracy.helpers import infer_model_label, reset_artifact_dir from tests.helpers.mark import hardware_test +pytestmark = [pytest.mark.diffusion, pytest.mark.full_model] + -@pytest.mark.advanced_model @pytest.mark.benchmark -@pytest.mark.diffusion @hardware_test(res={"cuda": "H100"}, num_cards=1) def test_gebench_h100_smoke( gebench_accuracy_servers, diff --git a/tests/e2e/accuracy/test_gedit_bench_h100_smoke.py b/tests/e2e/accuracy/test_gedit_bench_h100_smoke.py index 6227d6368..789f7ec93 100644 --- a/tests/e2e/accuracy/test_gedit_bench_h100_smoke.py +++ b/tests/e2e/accuracy/test_gedit_bench_h100_smoke.py @@ -10,10 +10,10 @@ from benchmarks.accuracy.image_to_image.gedit_bench import main as gedit_main from tests.e2e.accuracy.helpers import infer_model_label, reset_artifact_dir from tests.helpers.mark import hardware_test +pytestmark = [pytest.mark.diffusion, pytest.mark.full_model] + -@pytest.mark.advanced_model @pytest.mark.benchmark -@pytest.mark.diffusion @hardware_test(res={"cuda": "H100"}, num_cards=1) def test_gedit_bench_h100_smoke( gedit_accuracy_servers, diff --git a/tests/e2e/accuracy/test_qwen_image.py b/tests/e2e/accuracy/test_qwen_image.py index 8922d9d10..4b8215d54 100644 --- a/tests/e2e/accuracy/test_qwen_image.py +++ b/tests/e2e/accuracy/test_qwen_image.py @@ -17,6 +17,9 @@ from tests.helpers.env import run_post_test_cleanup, run_pre_test_cleanup from tests.helpers.mark import hardware_test from tests.helpers.runtime import OmniServer +pytestmark = [pytest.mark.full_model, pytest.mark.diffusion] + + MODEL_ID = "Qwen/Qwen-Image" MODEL_ENV_VAR = "QWEN_IMAGE_MODEL" PROMPT = "A photo of a cat sitting on a laptop keyboard, digital art style." @@ -99,9 +102,7 @@ def _run_diffusers_qwen_image(*, model: str, output_path: Path) -> Image.Image: run_post_test_cleanup(enable_force=True) -@pytest.mark.advanced_model @pytest.mark.benchmark -@pytest.mark.diffusion @hardware_test(res={"cuda": "H100"}, num_cards=1) def test_qwen_image_matches_diffusers(accuracy_artifact_root: Path) -> None: model = _model_name() diff --git a/tests/e2e/accuracy/test_qwen_image_edit.py b/tests/e2e/accuracy/test_qwen_image_edit.py index e17aca6e9..07deecca9 100644 --- a/tests/e2e/accuracy/test_qwen_image_edit.py +++ b/tests/e2e/accuracy/test_qwen_image_edit.py @@ -15,6 +15,9 @@ from tests.helpers.env import run_post_test_cleanup, run_pre_test_cleanup from tests.helpers.mark import hardware_test from tests.helpers.runtime import OmniServer +pytestmark = [pytest.mark.full_model, pytest.mark.diffusion] + + SINGLE_MODEL = "Qwen/Qwen-Image-Edit" MULTIPLE_MODEL = "Qwen/Qwen-Image-Edit-2509" WIDTH = 512 @@ -169,9 +172,7 @@ def _diffusers_output_multiple_image( ) -@pytest.mark.advanced_model @pytest.mark.benchmark -@pytest.mark.diffusion @hardware_test(res={"cuda": "H100"}, num_cards=1) def test_qwen_image_edit_single_matches_diffusers( accuracy_artifact_root: Path, @@ -196,9 +197,7 @@ def test_qwen_image_edit_single_matches_diffusers( ) -@pytest.mark.advanced_model @pytest.mark.benchmark -@pytest.mark.diffusion @hardware_test(res={"cuda": "H100"}, num_cards=1) @pytest.mark.skip( reason="Skipping as the second image seems to be ignored by the API. Will come back to this later after #2772 is merged." diff --git a/tests/e2e/accuracy/test_qwen_image_layered.py b/tests/e2e/accuracy/test_qwen_image_layered.py index 0ab9cb323..30ad2966f 100644 --- a/tests/e2e/accuracy/test_qwen_image_layered.py +++ b/tests/e2e/accuracy/test_qwen_image_layered.py @@ -17,6 +17,9 @@ from tests.helpers.env import run_post_test_cleanup, run_pre_test_cleanup from tests.helpers.mark import hardware_test from tests.helpers.runtime import OmniServer +pytestmark = [pytest.mark.full_model, pytest.mark.diffusion] + + MODEL_ID = "Qwen/Qwen-Image-Layered" MODEL_ENV_VAR = "QWEN_IMAGE_LAYERED_MODEL" PROMPT = "decompose into layers" @@ -126,9 +129,7 @@ def _run_diffusers_qwen_image_layered(*, model: str, input_image: Image.Image, o run_post_test_cleanup(enable_force=True) -@pytest.mark.advanced_model @pytest.mark.benchmark -@pytest.mark.diffusion @hardware_test(res={"cuda": "H100"}, num_cards=1) def test_qwen_image_layered_matches_diffusers(accuracy_artifact_root: Path, qwen_bear_image: Image.Image) -> None: model = _model_name() diff --git a/tests/e2e/accuracy/wan22_i2v/test_wan22_i2v_video_similarity.py b/tests/e2e/accuracy/wan22_i2v/test_wan22_i2v_video_similarity.py index 3aa5da85c..1caef3bff 100644 --- a/tests/e2e/accuracy/wan22_i2v/test_wan22_i2v_video_similarity.py +++ b/tests/e2e/accuracy/wan22_i2v/test_wan22_i2v_video_similarity.py @@ -50,6 +50,8 @@ from tests.e2e.accuracy.wan22_i2v.wan22_i2v_video_similarity_common import ( from tests.helpers.mark import hardware_test from tests.helpers.runtime import OmniServerParams +pytestmark = [pytest.mark.diffusion, pytest.mark.full_model] + def test_parse_video_metadata_extracts_dimensions_and_fps() -> None: payload = { @@ -537,9 +539,7 @@ def _generate_offline_video(*, image_source: str) -> tuple[Path, Path]: return offline_path, offline_metadata_path -@pytest.mark.advanced_model @pytest.mark.benchmark -@pytest.mark.diffusion @hardware_test(res={"cuda": "H100"}, num_cards=1) def test_wan22_i2v_diffusers_offline_generates_video( wan22_i2v_image_source: str | None, @@ -563,9 +563,7 @@ def test_wan22_i2v_diffusers_offline_generates_video( assert offline_metadata["frame_count"] == NUM_FRAMES -@pytest.mark.advanced_model @pytest.mark.benchmark -@pytest.mark.diffusion @hardware_test(res={"cuda": "H100"}, num_cards=2) @pytest.mark.parametrize("omni_server", SERVER_CASES, indirect=True) def test_wan22_i2v_online_serving_generates_video( @@ -594,9 +592,7 @@ def test_wan22_i2v_online_serving_generates_video( assert online_metadata["frame_count"] == NUM_FRAMES -@pytest.mark.advanced_model @pytest.mark.benchmark -@pytest.mark.diffusion @hardware_test(res={"cuda": "H100"}, num_cards=2) def test_wan22_i2v_serving_matches_diffusers_video_similarity( wan22_i2v_image_source: str | None, diff --git a/tests/e2e/offline_inference/test_quantization_fp8.py b/tests/e2e/offline_inference/test_quantization_fp8.py index 92cf351e3..fce7bfa27 100644 --- a/tests/e2e/offline_inference/test_quantization_fp8.py +++ b/tests/e2e/offline_inference/test_quantization_fp8.py @@ -27,9 +27,6 @@ Usage: """ import os - -os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" - from pathlib import Path from typing import Any @@ -44,6 +41,8 @@ from vllm_omni.platforms import current_omni_platform pytestmark = [pytest.mark.core_model, pytest.mark.diffusion] +os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" + # ─── helpers ────────────────────────────────────────────────────────────────── diff --git a/tests/e2e/offline_inference/test_stable_audio_expansion.py b/tests/e2e/offline_inference/test_stable_audio_expansion.py index 2fc07eb24..a5d3e6d22 100644 --- a/tests/e2e/offline_inference/test_stable_audio_expansion.py +++ b/tests/e2e/offline_inference/test_stable_audio_expansion.py @@ -22,6 +22,8 @@ from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.outputs import OmniRequestOutput from vllm_omni.platforms import current_omni_platform +pytestmark = [pytest.mark.full_model, pytest.mark.diffusion] + _SAMPLE_RATE = 44100 _CLIP_DURATION_S = 2.0 @@ -72,8 +74,6 @@ def generate_stable_audio_short_clip( return audio -@pytest.mark.advanced_model -@pytest.mark.diffusion @pytest.mark.cache @hardware_test(res={"cuda": "L4", "xpu": "B60"}) def test_stable_audio_quantization_and_teacache() -> None: diff --git a/tests/e2e/online_serving/test_bagel_expansion.py b/tests/e2e/online_serving/test_bagel_expansion.py index aa289c6d9..21142439b 100644 --- a/tests/e2e/online_serving/test_bagel_expansion.py +++ b/tests/e2e/online_serving/test_bagel_expansion.py @@ -19,6 +19,8 @@ import pytest from tests.helpers.mark import hardware_marks from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler, dummy_messages_from_mix_data +pytestmark = [pytest.mark.diffusion, pytest.mark.full_model] + PROMPT = "A futuristic city skyline at twilight, cyberpunk style, ultra-detailed, high resolution." NEGATIVE_PROMPT = "low quality, blurry, distorted, deformed, watermark" @@ -112,8 +114,6 @@ def _get_diffusion_feature_cases(model: str): ] -@pytest.mark.advanced_model -@pytest.mark.diffusion @pytest.mark.parametrize( "omni_server", _get_diffusion_feature_cases("ByteDance-Seed/BAGEL-7B-MoT"), diff --git a/tests/e2e/online_serving/test_dynin_omni_expansion.py b/tests/e2e/online_serving/test_dynin_omni_expansion.py index b63fcbd75..da179dbc8 100644 --- a/tests/e2e/online_serving/test_dynin_omni_expansion.py +++ b/tests/e2e/online_serving/test_dynin_omni_expansion.py @@ -5,29 +5,26 @@ Example online tests for Dynin-Omni model. """ import base64 -import gc import os from io import BytesIO from pathlib import Path -import numpy as np import pytest -import soundfile as sf from vllm.assets.image import ImageAsset from tests.helpers.mark import hardware_test -from tests.helpers.media import convert_audio_bytes_to_text from tests.helpers.runtime import OmniServerParams +pytestmark = [pytest.mark.full_model, pytest.mark.omni] + os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0" MODEL = "snu-aidas/Dynin-Omni" STAGE_CONFIG = str(Path(__file__).parent.parent / "stage_configs" / "dynin_omni_ci.yaml") -_WHISPER_SAMPLE_RATE_HZ = 16_000 T2I_PROMPT = "A high quality detailed living room interior photo." -T2S_PROMPT = "Please read this sentence naturally: Hello from Dynin-Omni online serving." +T2S_PROMPT = "Please read this sentence naturally: Hello from online serving." I2I_PROMPT = "Transform this outdoor nature boardwalk scene into a painting style with vivid colors." TEST_PARAMS = [OmniServerParams(model=MODEL, stage_config_path=STAGE_CONFIG, stage_init_timeout=600)] @@ -35,65 +32,6 @@ _STAGE_COUNT = 3 _I2I_STAGE_SAMPLING = {"max_tokens": 1, "temperature": 0.0, "top_p": 1.0, "detokenize": False} -def _prepare_audio_waveform_for_whisper(audio_data: np.ndarray, samplerate: int) -> np.ndarray: - """Normalize decoded audio into a mono 16 kHz float32 waveform for Whisper.""" - if samplerate <= 0: - raise ValueError(f"Invalid audio sample rate: {samplerate}") - - waveform = np.asarray(audio_data, dtype=np.float32) - if waveform.ndim == 0: - raise ValueError("Audio waveform must have at least one dimension") - if waveform.ndim > 1: - waveform = np.mean(waveform, axis=1) - if waveform.size == 0: - raise ValueError("Empty audio waveform") - - if samplerate != _WHISPER_SAMPLE_RATE_HZ: - target_num_samples = max(int(round(waveform.shape[0] * _WHISPER_SAMPLE_RATE_HZ / samplerate)), 1) - source_positions = np.arange(waveform.shape[0], dtype=np.float64) - target_positions = np.linspace( - 0.0, - max(waveform.shape[0] - 1, 0), - num=target_num_samples, - dtype=np.float64, - ) - waveform = np.interp(target_positions, source_positions, waveform).astype(np.float32) - - return np.ascontiguousarray(np.clip(waveform, -1.0, 1.0), dtype=np.float32) - - -def _convert_audio_bytes_to_text_without_ffmpeg(raw_bytes: bytes) -> str: - """Dynin t2s keeps Whisper transcription local to this test module and avoids ffmpeg.""" - import whisper - - data, samplerate = sf.read(BytesIO(raw_bytes), dtype="float32", always_2d=True) - audio_waveform = _prepare_audio_waveform_for_whisper(data, samplerate) - - model = whisper.load_model("small", device="cpu") - try: - transcript = model.transcribe( - audio_waveform, - temperature=0.0, - word_timestamps=True, - condition_on_previous_text=False, - )["text"] - finally: - del model - gc.collect() - - return transcript or "" - - -@pytest.fixture -def dynin_t2s_openai_client(openai_client, monkeypatch): - monkeypatch.setattr( - convert_audio_bytes_to_text, - "convert_audio_bytes_to_text", - _convert_audio_bytes_to_text_without_ffmpeg, - ) - return openai_client - - def _build_t2i_messages(prompt: str) -> list[dict]: return [{"role": "user", "content": [{"type": "text", "text": f"<|t2i|> {prompt}"}]}] @@ -118,8 +56,6 @@ def _build_i2i_messages(prompt: str) -> list[dict]: ] -@pytest.mark.advanced_model -@pytest.mark.omni @hardware_test(res={"cuda": "H100", "rocm": "MI325"}) @pytest.mark.parametrize("omni_server", TEST_PARAMS, indirect=True) def test_send_i2i_request_001(omni_server, openai_client) -> None: @@ -134,8 +70,6 @@ def test_send_i2i_request_001(omni_server, openai_client) -> None: openai_client.send_diffusion_request(request_config) -@pytest.mark.advanced_model -@pytest.mark.omni @hardware_test(res={"cuda": "H100", "rocm": "MI325"}) @pytest.mark.parametrize("omni_server", TEST_PARAMS, indirect=True) def test_send_t2i_request_001(omni_server, openai_client) -> None: @@ -147,14 +81,13 @@ def test_send_t2i_request_001(omni_server, openai_client) -> None: openai_client.send_diffusion_request(request_config) -@pytest.mark.core_model -@pytest.mark.omni @hardware_test(res={"cuda": "H100", "rocm": "MI325"}) @pytest.mark.parametrize("omni_server", TEST_PARAMS, indirect=True) -def test_send_t2s_request_001(omni_server, dynin_t2s_openai_client) -> None: +def test_send_t2s_request_001(omni_server, openai_client) -> None: request_config = { "model": omni_server.model, "messages": _build_t2s_messages(T2S_PROMPT), "modalities": ["audio"], + "audio_ref_text": T2S_PROMPT, } - dynin_t2s_openai_client.send_omni_request(request_config) + openai_client.send_omni_request(request_config) diff --git a/tests/e2e/online_serving/test_flux2_expansion.py b/tests/e2e/online_serving/test_flux2_expansion.py index ce06ad564..9a2b164b3 100644 --- a/tests/e2e/online_serving/test_flux2_expansion.py +++ b/tests/e2e/online_serving/test_flux2_expansion.py @@ -14,6 +14,8 @@ import pytest from tests.helpers.mark import hardware_marks from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler +pytestmark = [pytest.mark.diffusion, pytest.mark.full_model] + FOUR_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "L4"}, num_cards=4) POSITIVE_PROMPT = "A cat sitting on a windowsill" NEGATIVE_PROMPT = "blurry, low quality" @@ -86,8 +88,6 @@ def _get_diffusion_feature_cases(model: str): ] -@pytest.mark.advanced_model -@pytest.mark.diffusion @pytest.mark.parametrize( "omni_server", _get_diffusion_feature_cases( diff --git a/tests/e2e/online_serving/test_flux2_klein_inpaint_expansion.py b/tests/e2e/online_serving/test_flux2_klein_inpaint_expansion.py index 0c45bb33f..a3d4f004e 100644 --- a/tests/e2e/online_serving/test_flux2_klein_inpaint_expansion.py +++ b/tests/e2e/online_serving/test_flux2_klein_inpaint_expansion.py @@ -16,6 +16,8 @@ from PIL import Image, ImageDraw from tests.helpers.runtime import OmniServer, OmniServerParams +pytestmark = [pytest.mark.full_model, pytest.mark.diffusion] + MODEL = "black-forest-labs/FLUX.2-klein-4B" _HEIGHT = 512 @@ -66,8 +68,6 @@ def _send_edit_request(host: str, port: int, model: str, image_b64: str, mask_b6 return response.json() -@pytest.mark.advanced_model -@pytest.mark.diffusion @pytest.mark.parametrize("omni_server", _get_diffusion_feature_cases(MODEL), indirect=True) def test_flux2_klein_inpaint_basic(omni_server: OmniServer): input_image_b64 = _image_to_base64_jpeg(Image.new("RGB", (_WIDTH, _HEIGHT), (128, 128, 128))) @@ -92,8 +92,6 @@ def test_flux2_klein_inpaint_basic(omni_server: OmniServer): assert img.size == (_WIDTH, _HEIGHT) -@pytest.mark.advanced_model -@pytest.mark.diffusion @pytest.mark.parametrize("omni_server", _get_diffusion_feature_cases(MODEL), indirect=True) def test_flux2_klein_inpaint_deterministic(omni_server: OmniServer): input_image_b64 = _image_to_base64_jpeg(Image.new("RGB", (_WIDTH, _HEIGHT), (128, 128, 128))) @@ -137,8 +135,6 @@ def test_flux2_klein_inpaint_deterministic(omni_server: OmniServer): ) -@pytest.mark.advanced_model -@pytest.mark.diffusion @pytest.mark.parametrize("omni_server", _get_diffusion_feature_cases(MODEL), indirect=True) def test_flux2_klein_inpaint_multiple_outputs(omni_server: OmniServer): input_image_b64 = _image_to_base64_jpeg(Image.new("RGB", (_WIDTH, _HEIGHT), (128, 128, 128))) diff --git a/tests/e2e/online_serving/test_flux_2_dev_expansion.py b/tests/e2e/online_serving/test_flux_2_dev_expansion.py index addf6f002..953cb448a 100644 --- a/tests/e2e/online_serving/test_flux_2_dev_expansion.py +++ b/tests/e2e/online_serving/test_flux_2_dev_expansion.py @@ -17,6 +17,8 @@ import pytest from tests.helpers.mark import hardware_marks from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler, dummy_messages_from_mix_data +pytestmark = [pytest.mark.diffusion, pytest.mark.full_model] + MODEL = "black-forest-labs/FLUX.2-dev" PROMPT = "A cinematic mountain landscape at sunrise, dramatic clouds, ultra-detailed, realistic photography." NEGATIVE_PROMPT = "low quality, blurry, distorted, deformed, watermark" @@ -54,8 +56,6 @@ def _get_flux_2_dev_feature_cases(model: str): ] -@pytest.mark.advanced_model -@pytest.mark.diffusion @pytest.mark.parametrize( "omni_server", _get_flux_2_dev_feature_cases(MODEL), diff --git a/tests/e2e/online_serving/test_flux_kontext_expansion.py b/tests/e2e/online_serving/test_flux_kontext_expansion.py index 574bb3db4..fd7d7d1b4 100644 --- a/tests/e2e/online_serving/test_flux_kontext_expansion.py +++ b/tests/e2e/online_serving/test_flux_kontext_expansion.py @@ -8,6 +8,8 @@ import pytest from tests.helpers.media import generate_synthetic_image from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler, dummy_messages_from_mix_data +pytestmark = [pytest.mark.diffusion, pytest.mark.full_model] + EDIT_PROMPT = "Transform this modern, geometrist image into a Vincent van Gogh style impressionist painting." NEGATIVE_PROMPT = "blurry, low quality, modern, geometrist" MODEL = "black-forest-labs/FLUX.1-Kontext-dev" @@ -29,8 +31,6 @@ def _get_diffusion_feature_cases(model: str): ] -@pytest.mark.advanced_model -@pytest.mark.diffusion @pytest.mark.parametrize( "omni_server", _get_diffusion_feature_cases(MODEL), @@ -54,8 +54,6 @@ def test_flux_kontext_text_to_image(omni_server: OmniServer, openai_client: Open openai_client.send_diffusion_request(request_config) -@pytest.mark.advanced_model -@pytest.mark.diffusion @pytest.mark.parametrize( "omni_server", _get_diffusion_feature_cases(MODEL), @@ -83,8 +81,6 @@ def test_flux_kontext_image_edit(omni_server: OmniServer, openai_client: OpenAIC openai_client.send_diffusion_request(request_config) -@pytest.mark.advanced_model -@pytest.mark.diffusion @pytest.mark.parametrize( "omni_server", _get_diffusion_feature_cases(MODEL), @@ -110,8 +106,6 @@ def test_flux_kontext_image_edit_no_negative(omni_server: OmniServer, openai_cli openai_client.send_diffusion_request(request_config) -@pytest.mark.advanced_model -@pytest.mark.diffusion @pytest.mark.parametrize( "omni_server", _get_diffusion_feature_cases(MODEL), @@ -135,8 +129,6 @@ def test_flux_kontext_high_resolution(omni_server: OmniServer, openai_client: Op openai_client.send_diffusion_request(request_config) -@pytest.mark.advanced_model -@pytest.mark.diffusion @pytest.mark.parametrize( "omni_server", _get_diffusion_feature_cases(MODEL), diff --git a/tests/e2e/online_serving/test_hunyuan_video_15_expansion.py b/tests/e2e/online_serving/test_hunyuan_video_15_expansion.py index 7af4c4344..681d8ff23 100644 --- a/tests/e2e/online_serving/test_hunyuan_video_15_expansion.py +++ b/tests/e2e/online_serving/test_hunyuan_video_15_expansion.py @@ -14,6 +14,8 @@ import pytest from tests.helpers.mark import hardware_marks from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler +pytestmark = [pytest.mark.diffusion, pytest.mark.full_model] + PROMPT = "A cat walking across a sunlit garden, cinematic lighting, slow motion." NEGATIVE_PROMPT = "low quality, blurry, distorted" @@ -62,8 +64,6 @@ def _get_diffusion_feature_cases(model: str): ] -@pytest.mark.advanced_model -@pytest.mark.diffusion @pytest.mark.parametrize( "omni_server", _get_diffusion_feature_cases(MODEL), diff --git a/tests/e2e/online_serving/test_longcat_image_edit_expansion.py b/tests/e2e/online_serving/test_longcat_image_edit_expansion.py index 9a96280f2..4be5fe42d 100644 --- a/tests/e2e/online_serving/test_longcat_image_edit_expansion.py +++ b/tests/e2e/online_serving/test_longcat_image_edit_expansion.py @@ -17,6 +17,8 @@ from tests.helpers.mark import hardware_marks from tests.helpers.media import generate_synthetic_image from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler, dummy_messages_from_mix_data +pytestmark = [pytest.mark.diffusion, pytest.mark.full_model] + EDIT_PROMPT = "Transform this modern image into a cinematic animation style with vibrant colors and soft lighting." NEGATIVE_PROMPT = "blurry, low quality, distorted, oversaturated" SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"}) @@ -50,8 +52,6 @@ def _get_diffusion_feature_cases(model: str): ] -@pytest.mark.advanced_model -@pytest.mark.diffusion @pytest.mark.parametrize( "omni_server", _get_diffusion_feature_cases("meituan-longcat/LongCat-Image-Edit"), diff --git a/tests/e2e/online_serving/test_longcat_image_expansion.py b/tests/e2e/online_serving/test_longcat_image_expansion.py index 7d12db6a2..b9fd858b0 100644 --- a/tests/e2e/online_serving/test_longcat_image_expansion.py +++ b/tests/e2e/online_serving/test_longcat_image_expansion.py @@ -16,6 +16,8 @@ import pytest from tests.helpers.mark import hardware_marks from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler, dummy_messages_from_mix_data +pytestmark = [pytest.mark.diffusion, pytest.mark.full_model] + TEXT_TO_IMAGE_PROMPT = ( "A cinematic illustration of a cat typing on a silver laptop, soft window light, highly detailed." ) @@ -51,8 +53,6 @@ def _get_diffusion_feature_cases(model: str): ] -@pytest.mark.advanced_model -@pytest.mark.diffusion @pytest.mark.parametrize( "omni_server", _get_diffusion_feature_cases("meituan-longcat/LongCat-Image"), diff --git a/tests/e2e/online_serving/test_nextstep_expansion.py b/tests/e2e/online_serving/test_nextstep_expansion.py index 004a0967b..75290db5d 100644 --- a/tests/e2e/online_serving/test_nextstep_expansion.py +++ b/tests/e2e/online_serving/test_nextstep_expansion.py @@ -14,6 +14,8 @@ from tests.helpers.runtime import ( dummy_messages_from_mix_data, ) +pytestmark = [pytest.mark.full_model, pytest.mark.diffusion] + # L4: 4 GPUs + TP=4; XPU B60: 2 cards (use num_cards={"cuda": 4, "xpu": 4} if needed) FOUR_CARD_MARKS = hardware_marks( res={"cuda": "L4", "xpu": "B60"}, @@ -45,8 +47,6 @@ def _get_diffusion_feature_cases(model: str): ] -@pytest.mark.advanced_model -@pytest.mark.diffusion @pytest.mark.parametrize( "omni_server", _get_diffusion_feature_cases(model=os.environ.get("VLLM_TEST_NEXTSTEP_MODEL", _DEFAULT_MODEL)), diff --git a/tests/e2e/online_serving/test_qwen3_omni_expansion.py b/tests/e2e/online_serving/test_qwen3_omni_expansion.py index 6be2ebb2a..1388f7543 100644 --- a/tests/e2e/online_serving/test_qwen3_omni_expansion.py +++ b/tests/e2e/online_serving/test_qwen3_omni_expansion.py @@ -6,8 +6,6 @@ E2E Online tests for Qwen3-Omni model. import os -os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" - import pytest from tests.helpers.mark import hardware_test @@ -15,6 +13,10 @@ from tests.helpers.media import generate_synthetic_audio, generate_synthetic_ima from tests.helpers.runtime import OmniServerParams, dummy_messages_from_mix_data from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config +pytestmark = [pytest.mark.full_model, pytest.mark.omni] + +os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" + model = "Qwen/Qwen3-Omni-30B-A3B-Instruct" AUDIO_KEY = ["test"] @@ -130,8 +132,6 @@ def get_max_batch_size(size_type="few"): return batch_sizes.get(size_type, 5) -@pytest.mark.advanced_model -@pytest.mark.omni @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_text_to_audio_001(omni_server, openai_client) -> None: @@ -154,8 +154,6 @@ def test_text_to_audio_001(omni_server, openai_client) -> None: openai_client.send_omni_request(request_config) -@pytest.mark.advanced_model -@pytest.mark.omni @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) @pytest.mark.parametrize("omni_server", test_params + test_token_params, indirect=True) def test_text_to_text_audio_001(omni_server, openai_client) -> None: @@ -176,8 +174,6 @@ def test_text_to_text_audio_001(omni_server, openai_client) -> None: openai_client.send_omni_request(request_config, request_num=get_max_batch_size()) -@pytest.mark.advanced_model -@pytest.mark.omni @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_text_video_to_text_001(omni_server, openai_client) -> None: @@ -202,8 +198,6 @@ def test_text_video_to_text_001(omni_server, openai_client) -> None: openai_client.send_omni_request(request_config, request_num=get_max_batch_size()) -@pytest.mark.advanced_model -@pytest.mark.omni @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) @pytest.mark.parametrize("omni_server", test_params + test_token_params, indirect=True) def test_text_audio_to_text_audio_001(omni_server, openai_client) -> None: @@ -227,8 +221,6 @@ def test_text_audio_to_text_audio_001(omni_server, openai_client) -> None: openai_client.send_omni_request(request_config) -@pytest.mark.advanced_model -@pytest.mark.omni @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) @pytest.mark.parametrize("omni_server", test_params + test_token_params, indirect=True) def test_text_audio_to_text_audio_002(omni_server, openai_client) -> None: @@ -254,8 +246,6 @@ def test_text_audio_to_text_audio_002(omni_server, openai_client) -> None: openai_client.send_omni_request(request_config, request_num=get_max_batch_size()) -@pytest.mark.advanced_model -@pytest.mark.omni @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) @pytest.mark.parametrize("omni_server", test_params + test_token_params, indirect=True) def test_text_image_to_text_audio_001(omni_server, openai_client) -> None: @@ -280,8 +270,6 @@ def test_text_image_to_text_audio_001(omni_server, openai_client) -> None: openai_client.send_omni_request(request_config) -@pytest.mark.advanced_model -@pytest.mark.omni @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) @pytest.mark.parametrize("omni_server", test_params + test_token_params, indirect=True) def test_large_image_to_text_audio_001(omni_server, openai_client) -> None: @@ -310,8 +298,6 @@ def test_large_image_to_text_audio_001(omni_server, openai_client) -> None: openai_client.send_omni_request(request_config, request_num=get_max_batch_size()) -@pytest.mark.advanced_model -@pytest.mark.omni @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) @pytest.mark.parametrize("omni_server", test_params + test_token_params, indirect=True) def test_text_video_to_text_audio_001(omni_server, openai_client) -> None: @@ -338,8 +324,6 @@ def test_text_video_to_text_audio_001(omni_server, openai_client) -> None: @pytest.mark.skip(reason="There is a known issue with shape mismatch error.") -@pytest.mark.advanced_model -@pytest.mark.omni @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) @pytest.mark.parametrize("omni_server", test_params + test_token_params, indirect=True) def test_mix_to_text_audio_001(omni_server, openai_client) -> None: @@ -369,8 +353,6 @@ def test_mix_to_text_audio_001(omni_server, openai_client) -> None: openai_client.send_omni_request(request_config, request_num=get_max_batch_size()) -@pytest.mark.advanced_model -@pytest.mark.omni @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) @pytest.mark.parametrize("omni_server", test_params, indirect=True) @pytest.mark.skip(reason="issue: #2827") @@ -398,8 +380,6 @@ def test_audio_in_video_001(omni_server, openai_client) -> None: openai_client.send_omni_request(request_config) -@pytest.mark.advanced_model -@pytest.mark.omni @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_audio_in_video_002(omni_server, openai_client) -> None: @@ -427,8 +407,6 @@ def test_audio_in_video_002(omni_server, openai_client) -> None: openai_client.send_omni_request(request_config, request_num=get_max_batch_size()) -@pytest.mark.advanced_model -@pytest.mark.omni @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_one_word_prompt_001(omni_server, openai_client) -> None: @@ -463,8 +441,6 @@ def test_one_word_prompt_001(omni_server, openai_client) -> None: print(f"Similarity assertion failed, retrying {attempt + 2}/{_max_retries}: {e!r}") -@pytest.mark.advanced_model -@pytest.mark.omni @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_speaker_001(omni_server, openai_client) -> None: @@ -490,8 +466,6 @@ def test_speaker_001(omni_server, openai_client) -> None: openai_client.send_omni_request(request_config) -@pytest.mark.advanced_model -@pytest.mark.omni @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) @pytest.mark.parametrize("omni_server", test_params, indirect=True) @pytest.mark.skip(reason="Known issue: occasional inaccuracy in voice recognition.") @@ -528,8 +502,6 @@ def test_speaker_002(omni_server, openai_client) -> None: print(f"Gender assertion failed, retrying {attempt + 2}/{_max_retries}: {e!r}") -@pytest.mark.advanced_model -@pytest.mark.omni @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_speaker_003(omni_server, openai_client) -> None: @@ -555,8 +527,6 @@ def test_speaker_003(omni_server, openai_client) -> None: openai_client.send_omni_request(request_config) -@pytest.mark.advanced_model -@pytest.mark.omni @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_language_001(omni_server, openai_client) -> None: diff --git a/tests/e2e/online_serving/test_qwen3_tts_base_expansion.py b/tests/e2e/online_serving/test_qwen3_tts_base_expansion.py index 527eeaef4..d86f96af0 100644 --- a/tests/e2e/online_serving/test_qwen3_tts_base_expansion.py +++ b/tests/e2e/online_serving/test_qwen3_tts_base_expansion.py @@ -9,15 +9,17 @@ actual model inference, not mocks. import os -os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" -os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0" - import pytest from tests.helpers.mark import hardware_test from tests.helpers.runtime import OmniServerParams from tests.helpers.stage_config import get_deploy_config_path +pytestmark = [pytest.mark.full_model, pytest.mark.omni] + +os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" +os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0" + MODEL = "Qwen/Qwen3-TTS-12Hz-0.6B-Base" REF_AUDIO_URL = "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen3-TTS-Repo/clone_2.wav" @@ -61,9 +63,6 @@ tts_server_params = [ ] -@pytest.mark.advanced_model -@pytest.mark.core_model -@pytest.mark.omni @hardware_test(res={"cuda": "L4"}, num_cards=1) @pytest.mark.parametrize("omni_server", tts_server_params, indirect=True) def test_voice_clone_streaming_001(omni_server, openai_client) -> None: @@ -89,9 +88,6 @@ def test_voice_clone_streaming_001(omni_server, openai_client) -> None: openai_client.send_audio_speech_request(request_config, request_num=get_max_batch_size("few")) -@pytest.mark.advanced_model -@pytest.mark.core_model -@pytest.mark.omni @hardware_test(res={"cuda": "L4"}, num_cards=1) @pytest.mark.parametrize("omni_server", tts_server_params, indirect=True) def test_response_format_001(omni_server, openai_client) -> None: diff --git a/tests/e2e/online_serving/test_qwen3_tts_customvoice_expansion.py b/tests/e2e/online_serving/test_qwen3_tts_customvoice_expansion.py index b32303842..d97e70e41 100644 --- a/tests/e2e/online_serving/test_qwen3_tts_customvoice_expansion.py +++ b/tests/e2e/online_serving/test_qwen3_tts_customvoice_expansion.py @@ -9,15 +9,17 @@ actual model inference, not mocks. import os -os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" -os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0" - import pytest from tests.helpers.mark import hardware_test from tests.helpers.runtime import OmniServerParams from tests.helpers.stage_config import get_deploy_config_path +pytestmark = [pytest.mark.full_model, pytest.mark.omni] + +os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" +os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0" + MODEL = "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice" @@ -59,8 +61,6 @@ tts_server_params = [ ] -@pytest.mark.advanced_model -@pytest.mark.omni @hardware_test(res={"cuda": "L4"}, num_cards=1) @pytest.mark.parametrize("omni_server", tts_server_params, indirect=True) def test_voice_001(omni_server, openai_client) -> None: @@ -92,8 +92,6 @@ def test_voice_001(omni_server, openai_client) -> None: raise -@pytest.mark.advanced_model -@pytest.mark.omni @hardware_test(res={"cuda": "L4"}, num_cards=1) @pytest.mark.parametrize("omni_server", tts_server_params, indirect=True) def test_voice_002(omni_server, openai_client) -> None: @@ -117,8 +115,6 @@ def test_voice_002(omni_server, openai_client) -> None: openai_client.send_audio_speech_request(request_config) -@pytest.mark.advanced_model -@pytest.mark.omni @hardware_test(res={"cuda": "L4"}, num_cards=1) @pytest.mark.parametrize("omni_server", tts_server_params, indirect=True) def test_voice_003(omni_server, openai_client) -> None: @@ -142,8 +138,6 @@ def test_voice_003(omni_server, openai_client) -> None: openai_client.send_audio_speech_request(request_config) -@pytest.mark.advanced_model -@pytest.mark.omni @hardware_test(res={"cuda": "L4"}, num_cards=1) @pytest.mark.parametrize("omni_server", tts_server_params, indirect=True) def test_language_001(omni_server, openai_client) -> None: diff --git a/tests/e2e/online_serving/test_qwen_image_edit_expansion.py b/tests/e2e/online_serving/test_qwen_image_edit_expansion.py index c2461977d..bb8bf7dca 100644 --- a/tests/e2e/online_serving/test_qwen_image_edit_expansion.py +++ b/tests/e2e/online_serving/test_qwen_image_edit_expansion.py @@ -11,6 +11,8 @@ from tests.helpers.mark import hardware_marks from tests.helpers.media import generate_synthetic_image from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler, dummy_messages_from_mix_data +pytestmark = [pytest.mark.diffusion, pytest.mark.full_model] + EDIT_PROMPT = "Transform this modern, geometrist image into a Vincent van Gogh style impressionist painting." MULTI_EDIT_PROMPT = ( "Transform the first image into a Dadaism collage art. " @@ -108,8 +110,6 @@ def _get_diffusion_feature_cases(model: str): ] -@pytest.mark.advanced_model -@pytest.mark.diffusion @pytest.mark.parametrize( "omni_server", _get_diffusion_feature_cases("Qwen/Qwen-Image-Edit"), @@ -138,8 +138,6 @@ def test_qwen_image_edit(omni_server: OmniServer, openai_client: OpenAIClientHan openai_client.send_diffusion_request(request_config) -@pytest.mark.advanced_model -@pytest.mark.diffusion @pytest.mark.parametrize( "omni_server", _get_diffusion_feature_cases("Qwen/Qwen-Image-Edit-2509"), diff --git a/tests/e2e/online_serving/test_qwen_image_expansion.py b/tests/e2e/online_serving/test_qwen_image_expansion.py index b6f91c13d..0830e4944 100644 --- a/tests/e2e/online_serving/test_qwen_image_expansion.py +++ b/tests/e2e/online_serving/test_qwen_image_expansion.py @@ -15,6 +15,8 @@ import pytest from tests.helpers.mark import hardware_marks from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler, dummy_messages_from_mix_data +pytestmark = [pytest.mark.diffusion, pytest.mark.full_model] + T2I_PROMPT = "A photo of a cat sitting on a laptop keyboard, digital art style." NEGATIVE_PROMPT = "blurry, low quality" SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"}) @@ -117,8 +119,6 @@ def _get_diffusion_feature_cases(model: str): ] -@pytest.mark.advanced_model -@pytest.mark.diffusion @pytest.mark.parametrize( "omni_server", _get_diffusion_feature_cases("Qwen/Qwen-Image"), @@ -142,8 +142,6 @@ def test_qwen_image(omni_server: OmniServer, openai_client: OpenAIClientHandler) openai_client.send_diffusion_request(request_config) -@pytest.mark.advanced_model -@pytest.mark.diffusion @pytest.mark.parametrize( "omni_server", _get_diffusion_feature_cases("Qwen/Qwen-Image-2512"), diff --git a/tests/e2e/online_serving/test_qwen_image_layered_expansion.py b/tests/e2e/online_serving/test_qwen_image_layered_expansion.py index b958cfc05..4e79beab7 100644 --- a/tests/e2e/online_serving/test_qwen_image_layered_expansion.py +++ b/tests/e2e/online_serving/test_qwen_image_layered_expansion.py @@ -18,6 +18,8 @@ from tests.helpers.mark import hardware_marks from tests.helpers.media import decode_b64_image, generate_synthetic_image from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler, dummy_messages_from_mix_data +pytestmark = [pytest.mark.diffusion, pytest.mark.full_model] + MODEL = "Qwen/Qwen-Image-Layered" EDIT_PROMPT = "Decompose this image into layers." EMPTY_EDIT_PROMPT = "" @@ -71,8 +73,6 @@ FEATURE_CASES = [ ] -@pytest.mark.advanced_model -@pytest.mark.diffusion @pytest.mark.parametrize("omni_server", FEATURE_CASES, indirect=True) def test_feature(omni_server: OmniServer, openai_client: OpenAIClientHandler): """Test feature combinations with Qwen-Image-Layered.""" @@ -149,8 +149,6 @@ def _collect_image_url_items(openai_client: OpenAIClientHandler, request_config: return image_items -@pytest.mark.advanced_model -@pytest.mark.diffusion @pytest.mark.parametrize( "omni_server, expected_layers", LAYERS_GUARD_CASES, @@ -224,8 +222,6 @@ PROMPT_CASES = [ ] -@pytest.mark.advanced_model -@pytest.mark.diffusion @pytest.mark.parametrize("omni_server", PROMPT_CASES, indirect=True) def test_empty_prompt(omni_server: OmniServer, openai_client: OpenAIClientHandler): """Test feature combinations with Qwen-Image-Layered.""" diff --git a/tests/e2e/online_serving/test_sd3_expansion.py b/tests/e2e/online_serving/test_sd3_expansion.py index 37e590ba3..09b50d2e5 100644 --- a/tests/e2e/online_serving/test_sd3_expansion.py +++ b/tests/e2e/online_serving/test_sd3_expansion.py @@ -7,6 +7,8 @@ import pytest from tests.helpers.mark import hardware_marks from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler +pytestmark = [pytest.mark.diffusion, pytest.mark.full_model] + FOUR_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "L4"}, num_cards=4) POSITIVE_PROMPT = "A serene mountain landscape at sunset" NEGATIVE_PROMPT = "blurry, low quality, distorted" @@ -35,8 +37,6 @@ def _get_diffusion_feature_cases(model: str): ] -@pytest.mark.advanced_model -@pytest.mark.diffusion @pytest.mark.parametrize( "omni_server", _get_diffusion_feature_cases( diff --git a/tests/e2e/online_serving/test_wan22_expansion.py b/tests/e2e/online_serving/test_wan22_expansion.py index 7e5bc9121..36035f649 100644 --- a/tests/e2e/online_serving/test_wan22_expansion.py +++ b/tests/e2e/online_serving/test_wan22_expansion.py @@ -23,6 +23,8 @@ from tests.helpers.mark import hardware_marks from tests.helpers.media import generate_synthetic_image from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler +pytestmark = [pytest.mark.diffusion, pytest.mark.full_model] + PROMPT = "Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage." NEGATIVE_PROMPT = "low quality, blurry, distorted face, extra limbs, bad anatomy, watermark, logo, text, ugly, deformed, mutated, jpeg artifacts" SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"}) @@ -79,8 +81,6 @@ def _get_wan22_feature_cases(): return cases -@pytest.mark.advanced_model -@pytest.mark.diffusion @pytest.mark.parametrize( "omni_server", _get_wan22_feature_cases(), diff --git a/tests/e2e/online_serving/test_wan_2_1_vace_expansion.py b/tests/e2e/online_serving/test_wan_2_1_vace_expansion.py index 1f7a8c072..4f0e96447 100644 --- a/tests/e2e/online_serving/test_wan_2_1_vace_expansion.py +++ b/tests/e2e/online_serving/test_wan_2_1_vace_expansion.py @@ -26,6 +26,8 @@ import pytest from tests.helpers.mark import hardware_marks from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler +pytestmark = [pytest.mark.diffusion, pytest.mark.full_model] + MODEL = "Wan-AI/Wan2.1-VACE-1.3B-diffusers" PROMPT = "A cat walking slowly across a sunlit garden path" @@ -131,8 +133,6 @@ def _get_vace_feature_cases(): ] -@pytest.mark.advanced_model -@pytest.mark.diffusion @pytest.mark.parametrize( "omni_server", _get_vace_feature_cases(), diff --git a/tests/e2e/online_serving/test_zimage_expansion.py b/tests/e2e/online_serving/test_zimage_expansion.py index 679233c82..d730db0ca 100644 --- a/tests/e2e/online_serving/test_zimage_expansion.py +++ b/tests/e2e/online_serving/test_zimage_expansion.py @@ -16,6 +16,8 @@ import pytest from tests.helpers.mark import hardware_marks from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler +pytestmark = [pytest.mark.diffusion, pytest.mark.full_model] + MODEL = "Tongyi-MAI/Z-Image-Turbo" PROMPT = "A high-detail studio photo of an orange tabby cat sitting on a laptop keyboard." @@ -101,8 +103,6 @@ def _get_diffusion_feature_cases(): ] -@pytest.mark.advanced_model -@pytest.mark.diffusion @pytest.mark.parametrize( "omni_server", _get_diffusion_feature_cases(), diff --git a/tests/examples/offline_inference/test_text_to_image.py b/tests/examples/offline_inference/test_text_to_image.py index f24506587..041c32dc4 100644 --- a/tests/examples/offline_inference/test_text_to_image.py +++ b/tests/examples/offline_inference/test_text_to_image.py @@ -11,7 +11,8 @@ from tests.examples.helpers import EXAMPLES, ExampleRunner, ReadmeSnippet from tests.helpers.assertions import assert_image_valid from tests.helpers.mark import hardware_marks -pytestmark = [pytest.mark.advanced_model, pytest.mark.example, *hardware_marks(res={"cuda": "H100"})] +pytestmark = [pytest.mark.full_model, pytest.mark.example, *hardware_marks(res={"cuda": "H100"})] + T2I_SCRIPT = EXAMPLES / "offline_inference" / "text_to_image" / "text_to_image.py" README_PATH = T2I_SCRIPT.with_name("README.md") diff --git a/tests/examples/online_serving/test_qwen2_5_omni.py b/tests/examples/online_serving/test_qwen2_5_omni.py index b8acfef84..b3e49b8d9 100644 --- a/tests/examples/online_serving/test_qwen2_5_omni.py +++ b/tests/examples/online_serving/test_qwen2_5_omni.py @@ -4,9 +4,6 @@ See examples/online_serving/qwen2_5_omni/README.md """ import os - -os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" - from pathlib import Path import pytest @@ -21,7 +18,9 @@ from tests.helpers.media import convert_audio_file_to_text, cosine_similarity_te from tests.helpers.runtime import OmniServerParams from tests.helpers.stage_config import get_deploy_config_path -pytestmark = [pytest.mark.advanced_model, pytest.mark.example] +pytestmark = [pytest.mark.full_model, pytest.mark.example, pytest.mark.omni] + +os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" models = ["Qwen/Qwen2.5-Omni-7B"] @@ -40,8 +39,6 @@ test_params = [ common_args = ["python", os.path.join(example_dir, "openai_chat_completion_client_for_multimodal_generation.py")] -@pytest.mark.advanced_model -@pytest.mark.omni @hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards={"cuda": 4, "rocm": 2}) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_send_multimodal_request_001(omni_server) -> None: @@ -77,8 +74,6 @@ def test_send_multimodal_request_001(omni_server) -> None: # TODO: Verify the E2E latency after confirmation baseline. -@pytest.mark.advanced_model -@pytest.mark.omni @hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards={"cuda": 4, "rocm": 2}) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_send_multimodal_request_002(omni_server) -> None: @@ -114,8 +109,6 @@ def test_send_multimodal_request_002(omni_server) -> None: # TODO: Verify the E2E latency after confirmation baseline. -@pytest.mark.advanced_model -@pytest.mark.omni @hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards={"cuda": 4, "rocm": 2}) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_send_multimodal_request_003(omni_server) -> None: @@ -141,8 +134,6 @@ def test_send_multimodal_request_003(omni_server) -> None: # TODO: Verify the E2E latency after confirmation baseline. -@pytest.mark.advanced_model -@pytest.mark.omni @hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards={"cuda": 4, "rocm": 2}) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_modality_control_001(omni_server) -> None: @@ -171,8 +162,6 @@ def test_modality_control_001(omni_server) -> None: # TODO: Verify the E2E latency after confirmation baseline. -@pytest.mark.advanced_model -@pytest.mark.omni @hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards={"cuda": 4, "rocm": 2}) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_modality_control_002(omni_server) -> None: @@ -200,8 +189,6 @@ def test_modality_control_002(omni_server) -> None: # TODO: Verify the E2E latency after confirmation baseline. -@pytest.mark.advanced_model -@pytest.mark.omni @hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards={"cuda": 4, "rocm": 2}) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_modality_control_003(omni_server) -> None: @@ -238,8 +225,6 @@ def test_modality_control_003(omni_server) -> None: # TODO: Verify the E2E latency after confirmation baseline. -@pytest.mark.advanced_model -@pytest.mark.omni @hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards={"cuda": 4, "rocm": 2}) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_stream_001(omni_server) -> None: diff --git a/tests/examples/online_serving/test_qwen3_omni.py b/tests/examples/online_serving/test_qwen3_omni.py index e2c3b3384..e52a2bf5a 100644 --- a/tests/examples/online_serving/test_qwen3_omni.py +++ b/tests/examples/online_serving/test_qwen3_omni.py @@ -4,9 +4,6 @@ See examples/online_serving/qwen3_omni/README.md """ import os - -os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" - from pathlib import Path import pytest @@ -21,7 +18,9 @@ from tests.helpers.media import convert_audio_file_to_text, cosine_similarity_te from tests.helpers.runtime import OmniServerParams from tests.helpers.stage_config import get_deploy_config_path -pytestmark = [pytest.mark.advanced_model, pytest.mark.example] +pytestmark = [pytest.mark.full_model, pytest.mark.example, pytest.mark.omni] + +os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" models = ["Qwen/Qwen3-Omni-30B-A3B-Instruct"] @@ -39,8 +38,6 @@ test_params = [ common_args = ["python", os.path.join(example_dir, "openai_chat_completion_client_for_multimodal_generation.py")] -@pytest.mark.advanced_model -@pytest.mark.omni @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_send_multimodal_request_001(omni_server) -> None: @@ -69,8 +66,6 @@ def test_send_multimodal_request_001(omni_server) -> None: # TODO: Verify the E2E latency after confirmation baseline. -@pytest.mark.advanced_model -@pytest.mark.omni @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_send_multimodal_request_002(omni_server) -> None: @@ -102,8 +97,6 @@ def test_send_multimodal_request_002(omni_server) -> None: # TODO: Verify the E2E latency after confirmation baseline. -@pytest.mark.advanced_model -@pytest.mark.omni @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_send_multimodal_request_003(omni_server) -> None: @@ -119,8 +112,6 @@ def test_send_multimodal_request_003(omni_server) -> None: # TODO: Verify the E2E latency after confirmation baseline. -@pytest.mark.advanced_model -@pytest.mark.omni @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_modality_control_001(omni_server) -> None: @@ -143,8 +134,6 @@ def test_modality_control_001(omni_server) -> None: # TODO: Verify the E2E latency after confirmation baseline. -@pytest.mark.advanced_model -@pytest.mark.omni @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_modality_control_002(omni_server) -> None: @@ -167,8 +156,6 @@ def test_modality_control_002(omni_server) -> None: # TODO: Verify the E2E latency after confirmation baseline. -@pytest.mark.advanced_model -@pytest.mark.omni @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_modality_control_003(omni_server) -> None: @@ -199,8 +186,6 @@ def test_modality_control_003(omni_server) -> None: # TODO: Verify the E2E latency after confirmation baseline. -@pytest.mark.advanced_model -@pytest.mark.omni @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_stream_001(omni_server) -> None: diff --git a/tests/examples/online_serving/test_text_to_image.py b/tests/examples/online_serving/test_text_to_image.py index 6f89cb549..ee0a1fedb 100644 --- a/tests/examples/online_serving/test_text_to_image.py +++ b/tests/examples/online_serving/test_text_to_image.py @@ -18,7 +18,7 @@ from tests.helpers.assertions import assert_image_valid from tests.helpers.mark import hardware_marks from tests.helpers.runtime import OmniServer, OmniServerParams -pytestmark = [pytest.mark.advanced_model, pytest.mark.example, *hardware_marks(res={"cuda": "H100"})] +pytestmark = [pytest.mark.full_model, pytest.mark.example, *hardware_marks(res={"cuda": "H100"})] T2I_ONLINE_CLIENT = EXAMPLES / "online_serving" / "text_to_image" / "openai_chat_client.py" EXAMPLE_OUTPUT_SUBFOLDER = "example_online_t2i" diff --git a/tests/helpers/assertions.py b/tests/helpers/assertions.py index 533806477..604b76b62 100644 --- a/tests/helpers/assertions.py +++ b/tests/helpers/assertions.py @@ -11,7 +11,9 @@ import numpy as np import soundfile as sf from PIL import Image -from tests.helpers.media import cosine_similarity_text +from tests.helpers.media import ( + cosine_similarity_text, +) _GENDER_PIPELINE = None _GENDER_PIPELINE_LOCK = threading.Lock() @@ -56,7 +58,7 @@ def assert_image_diffusion_response( f"Expected {num_outputs_per_prompt} images, got {len(response.images)}" ) - if run_level == "advanced_model": + if run_level in {"advanced_model", "full_model"}: width = extra_body.get("width") height = extra_body.get("height") @@ -407,7 +409,8 @@ def assert_omni_response(response: Any, request_config: dict[str, Any], run_leve modalities = request_config.get("modalities", ["text", "audio"]) - if run_level == "advanced_model": + if run_level in {"advanced_model", "full_model"}: + # Verify output success if "audio" in modalities: assert response.audio_content is not None, "No audio output is generated" print(f"audio content is: {response.audio_content}") @@ -417,12 +420,11 @@ def assert_omni_response(response: Any, request_config: dict[str, Any], run_leve response.audio_bytes, speaker, ) - if "text" in modalities: assert response.text_content is not None, "No text output is generated" print(f"text content is: {response.text_content}") - # Verify image description + # Verify keywords in output word_types = ["text", "image", "audio", "video"] keywords_dict = request_config.get("key_words", {}) for word_type in word_types: @@ -441,11 +443,25 @@ def assert_omni_response(response: Any, request_config: dict[str, Any], run_leve ) # Verify similarity (Whisper transcript vs streamed/detokenized text) - if "text" in modalities and "audio" in modalities: - assert response.similarity is not None and response.similarity > 0.9, ( - "The audio content is not same as the text" - ) - print(f"similarity is: {response.similarity}") + if "audio" in modalities: + audio_ref_text = request_config.get("audio_ref_text") + if "text" in modalities: + transcript = (response.audio_content or "").strip() + text_output = (response.text_content or "").strip() + similarity = cosine_similarity_text( + transcript.lower(), + text_output.lower(), + ) + assert similarity > 0.9, "The audio content is not same as the text" + print(f"similarity is: {similarity}") + if audio_ref_text: + audio_similarity = cosine_similarity_text( + response.audio_content.lower(), + str(audio_ref_text).lower(), + ) + assert audio_similarity > 0.9, ( + f"The audio content does not match reference text: similarity={audio_similarity:.3f}" + ) def assert_audio_speech_response(response: Any, request_config: dict[str, Any], run_level: str) -> None: @@ -464,7 +480,7 @@ def assert_audio_speech_response(response: Any, request_config: dict[str, Any], elif req_fmt == "wav" and response.audio_format: assert req_fmt in response.audio_format - if run_level == "advanced_model" and req_fmt != "pcm": + if run_level in {"advanced_model", "full_model"} and req_fmt != "pcm": expected_text = request_config.get("input") if expected_text: transcript = (response.audio_content or "").strip() diff --git a/tests/helpers/fixtures/run_args.py b/tests/helpers/fixtures/run_args.py index b18a64b98..975584d20 100644 --- a/tests/helpers/fixtures/run_args.py +++ b/tests/helpers/fixtures/run_args.py @@ -6,8 +6,8 @@ def pytest_addoption(parser): "--run-level", action="store", default="core_model", - choices=["core_model", "advanced_model"], - help="Test level to run: L2, L3", + choices=["core_model", "advanced_model", "full_model"], + help="Test level to run: L2, L3, L4", ) diff --git a/tests/helpers/fixtures/runtime.py b/tests/helpers/fixtures/runtime.py index 77a468f2f..4cae13cd6 100644 --- a/tests/helpers/fixtures/runtime.py +++ b/tests/helpers/fixtures/runtime.py @@ -36,7 +36,7 @@ def omni_server(request: pytest.FixtureRequest, run_level: str, model_prefix: st model = model_prefix + params.model port = params.port stage_config_path = params.stage_config_path - if run_level == "advanced_model" and stage_config_path is not None: + if run_level in {"advanced_model", "full_model"} and stage_config_path is not None: with open(stage_config_path, encoding="utf-8") as f: cfg = yaml.safe_load(f) or {} # Strip ``load_format: dummy`` (CI overlay default) so advanced_model diff --git a/tests/helpers/runtime.py b/tests/helpers/runtime.py index 9520aff0c..190cb66b4 100644 --- a/tests/helpers/runtime.py +++ b/tests/helpers/runtime.py @@ -35,7 +35,6 @@ from tests.helpers.env import run_forced_gpu_cleanup_round from tests.helpers.media import ( _merge_base64_audio_to_segment, convert_audio_bytes_to_text, - cosine_similarity_text, decode_b64_image, ) from vllm_omni.config.stage_config import resolve_deploy_yaml @@ -500,7 +499,6 @@ class OmniResponse: audio_content: str | None = None audio_format: str | None = None audio_bytes: bytes | None = None - similarity: float | None = None e2e_latency: float | None = None success: bool = False error_message: str | None = None @@ -542,19 +540,15 @@ class OpenAIClientHandler: text_content += content result.e2e_latency = time.perf_counter() - start_time audio_content = None - similarity = None if audio_data: merged_seg = _merge_base64_audio_to_segment(audio_data) wav_buf = BytesIO() merged_seg.export(wav_buf, format="wav") result.audio_bytes = wav_buf.getvalue() audio_content = convert_audio_bytes_to_text(result.audio_bytes) - if audio_content and text_content: - similarity = cosine_similarity_text(audio_content.lower(), text_content.lower()) result.text_content = text_content result.audio_data = audio_data result.audio_content = audio_content - result.similarity = similarity result.success = True except Exception as e: result.error_message = f"Stream processing error: {str(e)}" @@ -578,15 +572,11 @@ class OpenAIClientHandler: result.cached_tokens = details.cached_tokens result.e2e_latency = time.perf_counter() - start_time audio_content = None - similarity = None if audio_data: result.audio_bytes = base64.b64decode(audio_data) audio_content = convert_audio_bytes_to_text(result.audio_bytes) - if audio_content and text_content: - similarity = cosine_similarity_text(audio_content.lower(), text_content.lower()) result.text_content = text_content result.audio_content = audio_content - result.similarity = similarity result.success = True except Exception as e: result.error_message = f"Non-stream processing error: {str(e)}" From 9d1392dcfc34da09d3aca6f77c2777a6727a3ba8 Mon Sep 17 00:00:00 2001 From: GXIN <37653830+gxxx-hum@users.noreply.github.com> Date: Wed, 22 Apr 2026 18:15:08 +0800 Subject: [PATCH 09/45] [NPU] Support code predictor NPU graph (#2695) Signed-off-by: XIN GAO <1037396230@qq.com> Signed-off-by: gcanlin Co-authored-by: gcanlin --- pyproject.toml | 1 + vllm_omni/deploy/qwen3_omni_moe.yaml | 4 +- vllm_omni/deploy/qwen3_tts.yaml | 13 -- .../models/common/qwen3_code_predictor.py | 137 +++++++++++++++--- .../qwen3_omni_moe_code_predictor_mtp.py | 3 +- 5 files changed, 123 insertions(+), 35 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 0958b9395..c7b2be20e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -242,4 +242,5 @@ ue = "ue" semantics = "semantics" fullset = "fullset" Vai = "Vai" +tockens = "tockens" CANN = "CANN" diff --git a/vllm_omni/deploy/qwen3_omni_moe.yaml b/vllm_omni/deploy/qwen3_omni_moe.yaml index fb8b61621..39baed6bd 100644 --- a/vllm_omni/deploy/qwen3_omni_moe.yaml +++ b/vllm_omni/deploy/qwen3_omni_moe.yaml @@ -46,7 +46,6 @@ stages: - stage_id: 2 gpu_memory_utilization: 0.1 - max_num_seqs: 1 enforce_eager: true async_scheduling: false max_num_batched_tokens: 51200 @@ -67,10 +66,11 @@ platforms: - stage_id: 0 gpu_memory_utilization: 0.6 tensor_parallel_size: 2 + max_num_batched_tokens: 8192 devices: "0,1" - stage_id: 1 gpu_memory_utilization: 0.6 - enforce_eager: true + max_num_batched_tokens: 8192 devices: "2" - stage_id: 2 gpu_memory_utilization: 0.3 diff --git a/vllm_omni/deploy/qwen3_tts.yaml b/vllm_omni/deploy/qwen3_tts.yaml index 5948007a9..522ea7c58 100644 --- a/vllm_omni/deploy/qwen3_tts.yaml +++ b/vllm_omni/deploy/qwen3_tts.yaml @@ -71,16 +71,3 @@ stages: max_tokens: 65536 seed: 42 repetition_penalty: 1.0 - -platforms: - npu: - stages: - # NPU does not yet support async-scheduling for TTS, and the - # talker fits at max_num_seqs=1 only. - - stage_id: 0 - max_num_seqs: 1 - enforce_eager: true - async_scheduling: false - - stage_id: 1 - gpu_memory_utilization: 0.2 - async_scheduling: false diff --git a/vllm_omni/model_executor/models/common/qwen3_code_predictor.py b/vllm_omni/model_executor/models/common/qwen3_code_predictor.py index 3a904442f..ecfeec381 100644 --- a/vllm_omni/model_executor/models/common/qwen3_code_predictor.py +++ b/vllm_omni/model_executor/models/common/qwen3_code_predictor.py @@ -122,6 +122,9 @@ class CodePredictorAttention(nn.Module): super().__init__() self.num_heads = config.num_attention_heads self.num_kv_heads = config.num_key_value_heads + assert self.num_heads % self.num_kv_heads == 0 + self.is_gqa = self.num_kv_heads != self.num_heads + self.num_queries_per_kv = self.num_heads // self.num_kv_heads self.head_dim = getattr( config, "head_dim", @@ -129,7 +132,7 @@ class CodePredictorAttention(nn.Module): ) self.hidden_size = config.hidden_size self.scaling = self.head_dim**-0.5 - self._use_gqa = self.num_kv_heads != self.num_heads + self.max_seq = int(config.num_code_groups) + 1 # Separate q/k/v projections matching HF (no fused packing) bias = getattr(config, "attention_bias", False) @@ -140,6 +143,73 @@ class CodePredictorAttention(nn.Module): self.q_norm = _RMSNorm(self.head_dim, eps=config.rms_norm_eps) self.k_norm = _RMSNorm(self.head_dim, eps=config.rms_norm_eps) + if current_omni_platform.is_npu(): + if self.max_seq > 2048: + raise ValueError( + "Qwen3-TTS code predictor NPU fusion attention uses a fixed 2048x2048 " + f"causal mask, but max_seq={self.max_seq} exceeds the mask size." + ) + # Ascend SDPA is_causal migration example uses a fixed 2048x2048 + # compressed causal mask with sparse_mode=2. + fusion_mask = torch.triu( + torch.ones(2048, 2048, dtype=torch.bool), + diagonal=1, + ) + self.register_buffer("_fusion_causal_mask", fusion_mask, persistent=False) + + def _forward_npu_attention( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + bsz: int, + seq_len: int, + ) -> torch.Tensor: + import torch_npu + + q_f, k_f, v_f = q, k, v + if self.is_gqa: + k_f = ( + k[:, :, None, :, :] + .expand(bsz, self.num_kv_heads, self.num_queries_per_kv, seq_len, self.head_dim) + .reshape(bsz, self.num_heads, seq_len, self.head_dim) + ) + v_f = ( + v[:, :, None, :, :] + .expand(bsz, self.num_kv_heads, self.num_queries_per_kv, seq_len, self.head_dim) + .reshape(bsz, self.num_heads, seq_len, self.head_dim) + ) + + mask = self._fusion_causal_mask + mask = mask.contiguous() + q_f = q_f.contiguous() + k_f = k_f.contiguous() + v_f = v_f.contiguous() + return torch_npu.npu_fusion_attention( + q_f, + k_f, + v_f, + self.num_heads, + "BNSD", + pse=None, + padding_mask=None, + atten_mask=mask, + scale=float(self.scaling), + keep_prob=1.0, + # Keep torch_npu's API spelling. + pre_tockens=2147483647, + next_tockens=2147483647, + inner_precise=0, + prefix=None, + actual_seq_qlen=None, + actual_seq_kvlen=None, + # Ascend SDPA is_causal migration example uses sparse_mode=2. + sparse_mode=2, + gen_mask_parallel=True, + # Keep sync=True for the NPU fused attention path. + sync=True, + )[0] + def forward( self, hidden_states: torch.Tensor, @@ -160,14 +230,17 @@ class CodePredictorAttention(nn.Module): q = (q * cos) + (_rotate_half(q) * sin) k = (k * cos) + (_rotate_half(k) * sin) - attn_out = F.scaled_dot_product_attention( - q, - k, - v, - scale=self.scaling, - is_causal=True, - enable_gqa=self._use_gqa, - ) + if not current_omni_platform.is_npu(): + attn_out = F.scaled_dot_product_attention( + q, + k, + v, + scale=self.scaling, + is_causal=True, + enable_gqa=self.is_gqa, + ) + else: + attn_out = self._forward_npu_attention(q, k, v, bsz, seq_len) attn_out = attn_out.transpose(1, 2).reshape(bsz, seq_len, -1) return self.o_proj(attn_out) @@ -384,7 +457,7 @@ class CodePredictorWrapper(nn.Module): self._bucket_pos_ids: dict[int, torch.Tensor] = {} self._lm_heads_list: list[nn.Module] | None = None self._codec_embeds_list: list[nn.Module] | None = None - self._cuda_graphs: dict[int, tuple[torch.cuda.CUDAGraph, torch.Tensor]] = {} + self._device_graphs: dict[int, tuple] = {} # (graph, static_output) per bucket def get_input_embeddings(self) -> nn.ModuleList: return self.model.get_input_embeddings() @@ -412,7 +485,7 @@ class CodePredictorWrapper(nn.Module): self._proj_buf = torch.zeros(bsz, max_seq, self._cp_hidden, dtype=dtype, device=device) def _setup_compile(self) -> None: - """Lazily set up torch.compile with optional CUDA graph capture.""" + """Lazily set up torch.compile with optional device graph capture.""" if self._compiled_model_fwd is not None: return @@ -424,8 +497,16 @@ class CodePredictorWrapper(nn.Module): self._codec_embeds_list = list(self.model.codec_embedding) if not current_omni_platform.supports_torch_inductor(): - logger.warning_once("code_predictor: torch.compile disabled") + # NPU or other platforms without Inductor support self._compiled_model_fwd = self.model.forward + + if current_omni_platform.is_npu() and self._wrapper_config.use_cuda_graphs: + # For NPU, use eager + NPU graphs (no torch.compile) + self._warmup_buckets() + self._capture_npu_graphs() + logger.info("code_predictor: eager mode + NPU graphs") + else: + logger.warning_once("code_predictor: torch.compile disabled") return # torch.compile fuses RMSNorm/RoPE in ways that lose float32 @@ -491,10 +572,28 @@ class CodePredictorWrapper(nn.Module): with torch.cuda.graph(g, pool=pool): static_output = self._compiled_model_fwd(static_input, pos_ids) - self._cuda_graphs[bsz] = (g, static_output) + self._device_graphs[bsz] = (g, static_output) logger.info("code_predictor: captured CUDA graphs for buckets %s", self._bucket_sizes) + def _capture_npu_graphs(self) -> None: + """Capture an NPU graph per bucket using torch_npu's NPUGraph.""" + max_seq = self._num_groups + 1 + proj_buf = self._proj_buf + pool = torch.npu.graph_pool_handle() + + for bsz in self._bucket_sizes: + static_input = proj_buf[:bsz, :max_seq, :] + pos_ids = self._bucket_pos_ids[bsz] + + g = torch.npu.NPUGraph() + with torch.npu.graph(g, pool=pool): + static_output = self._compiled_model_fwd(static_input, pos_ids) + + self._device_graphs[bsz] = (g, static_output) + + logger.info("code_predictor: captured NPU graphs for buckets %s", self._bucket_sizes) + # ------------------------------------------------------------------ # Forward -- re-prefill + inline sampling # ------------------------------------------------------------------ @@ -544,8 +643,8 @@ class CodePredictorWrapper(nn.Module): torch.arange(max_seq, device=device, dtype=torch.long).unsqueeze(0).expand(padded_bsz, -1).contiguous() ) - # Use captured CUDA graph if available, otherwise call compiled fn. - cuda_graph_entry = self._cuda_graphs.get(padded_bsz) + # Use captured device graph if available, otherwise call compiled fn. + device_graph_entry = self._device_graphs.get(padded_bsz) # Prepare sampling parameters stored_mode = self._wrapper_config.sampling_mode == "stored" @@ -570,10 +669,10 @@ class CodePredictorWrapper(nn.Module): # Autoregressive loop: predict layers 1..G-1 for step in range(1, num_groups): - # Run transformer (CUDA graph replay or compiled forward) - if cuda_graph_entry is not None: - cuda_graph_entry[0].replay() - hidden_out = cuda_graph_entry[1] + # Run transformer (device graph replay or compiled forward) + if device_graph_entry is not None: + device_graph_entry[0].replay() + hidden_out = device_graph_entry[1] else: hidden_out = model_fwd(proj_buf[:padded_bsz, :max_seq, :], full_pos_ids) diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_code_predictor_mtp.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_code_predictor_mtp.py index 819e22e18..73c7e41d2 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_code_predictor_mtp.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_code_predictor_mtp.py @@ -6,6 +6,7 @@ from vllm_omni.model_executor.models.common.qwen3_code_predictor import ( CodePredictorWrapper, CodePredictorWrapperConfig, ) +from vllm_omni.platforms import current_omni_platform class Qwen3OmniMoeTalkerCodePredictor(CodePredictorWrapper): @@ -17,7 +18,7 @@ class Qwen3OmniMoeTalkerCodePredictor(CodePredictorWrapper): vllm_config=vllm_config, cp_config=cp_config, wrapper_config=CodePredictorWrapperConfig( - use_cuda_graphs=False, + use_cuda_graphs=current_omni_platform.is_npu(), use_parallel_embedding=True, use_projection=False, return_proj_buf=True, From 9337bec9adf3ad62a6efa969606c4d10986ebc32 Mon Sep 17 00:00:00 2001 From: amy-why-3459 Date: Wed, 22 Apr 2026 19:08:10 +0800 Subject: [PATCH 10/45] [Tests] Modify test cases (#2991) Signed-off-by: amy-why-3459 --- tests/dfx/perf/tests/test_qwen_omni.json | 64 +++++++++---------- .../test_qwen3_omni_expansion.py | 1 + 2 files changed, 33 insertions(+), 32 deletions(-) diff --git a/tests/dfx/perf/tests/test_qwen_omni.json b/tests/dfx/perf/tests/test_qwen_omni.json index ca3eb5557..eda9720c4 100644 --- a/tests/dfx/perf/tests/test_qwen_omni.json +++ b/tests/dfx/perf/tests/test_qwen_omni.json @@ -10,16 +10,16 @@ "dataset_name": "random", "backend": "openai-chat-omni", "endpoint": "/v1/chat/completions", - "num_prompts": [4, 16, 40], - "max_concurrency": [1, 4, 10], + "num_prompts": [4, 16, 32, 64, 128], + "max_concurrency": [1, 4, 8, 16, 32], "random_input_len": 2500, "random_output_len": 900, "ignore_eos": true, "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { - "mean_ttft_ms": [1000, 3000, 5000], - "mean_audio_ttfp_ms": [30000, 60000, 90000], - "mean_audio_rtf": [0.35, 0.45, 0.55] + "mean_ttft_ms": [1000, 3000, 5000, 7000, 9000], + "mean_audio_ttfp_ms": [30000, 60000, 90000, 120000, 150000], + "mean_audio_rtf": [0.35, 0.45, 0.55, 0.65, 0.75] } }, { @@ -52,7 +52,7 @@ "backend": "openai-chat-omni", "endpoint": "/v1/chat/completions", "num_prompts": [40], - "request_rate": [0.3], + "request_rate": [0.5], "random_input_len": 100, "random_output_len": 100, "random_range_ratio": 0.0, @@ -69,9 +69,9 @@ }, "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { - "mean_ttft_ms": [4000], - "mean_audio_ttfp_ms": [13000], - "mean_audio_rtf": [0.35] + "mean_ttft_ms": [6000], + "mean_audio_ttfp_ms": [15000], + "mean_audio_rtf": [0.45] } }, { @@ -79,7 +79,7 @@ "backend": "openai-chat-omni", "endpoint": "/v1/chat/completions", "num_prompts": [100], - "request_rate": [0.5], + "request_rate": [1.0], "random_input_len": 100, "random_output_len": 100, "random_range_ratio": 0.0, @@ -98,9 +98,9 @@ }, "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { - "mean_ttft_ms": [6000], - "mean_audio_ttfp_ms": [15000], - "mean_audio_rtf": [0.45] + "mean_ttft_ms": [12000], + "mean_audio_ttfp_ms": [18000], + "mean_audio_rtf": [0.9] } } ] @@ -116,16 +116,16 @@ "dataset_name": "random", "backend": "openai-chat-omni", "endpoint": "/v1/chat/completions", - "num_prompts": [4, 16, 40], - "max_concurrency": [1, 4, 10], + "num_prompts": [4, 16, 32, 64], + "max_concurrency": [1, 4, 8, 16], "random_input_len": 2500, "random_output_len": 900, "ignore_eos": true, "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { - "mean_ttft_ms": [1000, 3000, 5000], - "mean_audio_ttfp_ms": [1000, 3000, 5000], - "mean_audio_rtf": [0.2, 0.35, 0.6] + "mean_ttft_ms": [1000, 3000, 5000, 7000], + "mean_audio_ttfp_ms": [1000, 3000, 5000, 7000], + "mean_audio_rtf": [0.2, 0.35, 0.6, 0.85] } }, { @@ -158,7 +158,7 @@ "backend": "openai-chat-omni", "endpoint": "/v1/chat/completions", "num_prompts": [40], - "request_rate": [0.3], + "request_rate": [0.5], "random_input_len": 100, "random_output_len": 100, "random_range_ratio": 0.0, @@ -175,9 +175,9 @@ }, "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { - "mean_ttft_ms": [4000], - "mean_audio_ttfp_ms": [4000], - "mean_audio_rtf": [0.4] + "mean_ttft_ms": [6000], + "mean_audio_ttfp_ms": [6000], + "mean_audio_rtf": [0.7] } }, { @@ -185,7 +185,7 @@ "backend": "openai-chat-omni", "endpoint": "/v1/chat/completions", "num_prompts": [100], - "request_rate": [0.5], + "request_rate": [1.0], "random_input_len": 100, "random_output_len": 100, "random_range_ratio": 0.0, @@ -204,17 +204,17 @@ }, "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { - "mean_ttft_ms": [6000], - "mean_audio_ttfp_ms": [6000], - "mean_audio_rtf": [0.7] + "mean_ttft_ms": [12000], + "mean_audio_ttfp_ms": [12000], + "mean_audio_rtf": [1.0] } }, { "dataset_name": "random", "backend": "openai-chat-omni", "endpoint": "/v1/chat/completions", - "num_prompts": [4, 16, 40], - "max_concurrency": [1, 4, 10], + "num_prompts": [4, 16, 32, 64, 128], + "max_concurrency": [1, 4, 8, 16, 32], "random_input_len": 2500, "random_output_len": 900, "ignore_eos": true, @@ -223,7 +223,7 @@ }, "percentile-metrics": "ttft,tpot,itl,e2el", "baseline": { - "mean_ttft_ms": [1000, 3000, 5000] + "mean_ttft_ms": [1000, 3000, 5000, 7000, 9000] } }, { @@ -257,7 +257,7 @@ "backend": "openai-chat-omni", "endpoint": "/v1/chat/completions", "num_prompts": [40], - "request_rate": [0.3], + "request_rate": [0.5], "random_input_len": 100, "random_output_len": 100, "random_range_ratio": 0.0, @@ -277,7 +277,7 @@ }, "percentile-metrics": "ttft,tpot,itl,e2el", "baseline": { - "mean_ttft_ms": [4000] + "mean_ttft_ms": [6000] } }, { @@ -285,7 +285,7 @@ "backend": "openai-chat-omni", "endpoint": "/v1/chat/completions", "num_prompts": [100], - "request_rate": [0.5], + "request_rate": [1.0], "random_input_len": 100, "random_output_len": 100, "random_range_ratio": 0.0, diff --git a/tests/e2e/online_serving/test_qwen3_omni_expansion.py b/tests/e2e/online_serving/test_qwen3_omni_expansion.py index 1388f7543..2ebf5c7e3 100644 --- a/tests/e2e/online_serving/test_qwen3_omni_expansion.py +++ b/tests/e2e/online_serving/test_qwen3_omni_expansion.py @@ -84,6 +84,7 @@ test_params = [ model=model, stage_config_path=get_async_chunk_config(default_path), use_stage_cli=True, + server_args=["--async-chunk"], ), id="async_chunk", ), From ee15f39ea785b31492027248ba62f9f70ba5c58a Mon Sep 17 00:00:00 2001 From: Liang Lv Date: Wed, 22 Apr 2026 19:54:45 +0800 Subject: [PATCH 11/45] [AutoRound] Support Qwen Omni W4A16 quantization model (#2670) Signed-off-by: lvliang-intel Co-authored-by: SYLAR <125541396+lishunyang12@users.noreply.github.com> --- .../diffusion/quantization/autoround.md | 2 + .../offline_inference/qwen2_5_omni/end2end.py | 19 +- .../quantization/test_component_routing.py | 408 ++++++++++++++++++ .../test_qwen2_5_omni_autoround_w4a16.py | 170 ++++++++ .../test_qwen3_omni_autoround_w4a16.py | 205 +++++++++ vllm_omni/config/model.py | 19 + vllm_omni/engine/stage_init_utils.py | 8 +- .../qwen3_omni/qwen3_omni_moe_talker.py | 18 +- vllm_omni/quantization/__init__.py | 2 + vllm_omni/quantization/factory.py | 6 +- vllm_omni/quantization/inc_config.py | 140 ++++++ 11 files changed, 989 insertions(+), 8 deletions(-) create mode 100644 tests/diffusion/quantization/test_component_routing.py create mode 100644 tests/e2e/offline_inference/test_qwen2_5_omni_autoround_w4a16.py create mode 100644 tests/e2e/offline_inference/test_qwen3_omni_autoround_w4a16.py create mode 100644 vllm_omni/quantization/inc_config.py diff --git a/docs/user_guide/diffusion/quantization/autoround.md b/docs/user_guide/diffusion/quantization/autoround.md index 48df176b0..d06627d40 100644 --- a/docs/user_guide/diffusion/quantization/autoround.md +++ b/docs/user_guide/diffusion/quantization/autoround.md @@ -72,6 +72,8 @@ At load time: | Model | HF Checkpoint | Scheme | Group Size | Backend | |-------|--------------|--------|------------|---------| | FLUX.1-dev | `vllm-project-org/FLUX.1-dev-AutoRound-w4a16` | W4A16 | 128 | GPTQ-Marlin | +| Qwen2.5-Omni-7B | `Intel/Qwen2.5-Omni-7B-int4-AutoRound` | W4A16 | 128 | GPTQ-Marlin | +| Qwen3-Omni-30B-A3B-Instruct | `Intel/Qwen3-Omni-30B-A3B-Instruct-int4-AutoRound` | W4A16 | 128 | GPTQ-Marlin | ## Creating a Quantized Checkpoint diff --git a/examples/offline_inference/qwen2_5_omni/end2end.py b/examples/offline_inference/qwen2_5_omni/end2end.py index dfe124700..a65c554a9 100644 --- a/examples/offline_inference/qwen2_5_omni/end2end.py +++ b/examples/offline_inference/qwen2_5_omni/end2end.py @@ -5,6 +5,7 @@ This example shows how to use vLLM-Omni for running offline inference with the correct prompt format on Qwen2.5-Omni """ +import json import os import time from typing import NamedTuple @@ -289,7 +290,10 @@ query_map = { def main(args): - model_name = "Qwen/Qwen2.5-Omni-7B" + model_name = args.model + quantization_config = None + if args.quantization_config is not None: + quantization_config = json.loads(args.quantization_config) # Get paths from args video_path = getattr(args, "video_path", None) @@ -320,6 +324,7 @@ def main(args): query_result = query_func(audio_path=audio_path, sampling_rate=sampling_rate) else: query_result = query_func() + args.quantization_config = quantization_config omni = Omni.from_cli_args(args, model=model_name) thinker_sampling_params = SamplingParams( temperature=0.0, # Deterministic - no randomness @@ -424,6 +429,18 @@ def main(args): def parse_args(): parser = FlexibleArgumentParser(description="Demo on using vLLM for offline inference with audio language models") + parser.add_argument( + "--model", + type=str, + default="Qwen/Qwen2.5-Omni-7B", + help="Model name or local path.", + ) + parser.add_argument( + "--quantization-config", + type=str, + default=None, + help="Optional JSON string forwarded to Omni(quantization_config=...).", + ) parser.add_argument( "--query-type", "-q", diff --git a/tests/diffusion/quantization/test_component_routing.py b/tests/diffusion/quantization/test_component_routing.py new file mode 100644 index 000000000..c8b3837e2 --- /dev/null +++ b/tests/diffusion/quantization/test_component_routing.py @@ -0,0 +1,408 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for component routing for quantization.""" + +from unittest.mock import MagicMock + +import pytest +import torch +from vllm.model_executor.layers.quantization.base_config import ( + QuantizationConfig, +) +from vllm.model_executor.models.utils import WeightsMapper + +from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker import ( + PRE_QUANTIZED_METHODS, +) +from vllm_omni.quantization.component_config import ( + ComponentQuantizationConfig, +) +from vllm_omni.quantization.inc_config import OmniINCConfig + +pytestmark = [pytest.mark.core_model] + + +# --------------------------------------------------------------------------- +# Helpers: lightweight mock quant configs +# --------------------------------------------------------------------------- + + +class _MockQuantConfig(QuantizationConfig): + """Minimal mock that only implements get_name().""" + + def __init__(self, name: str, **attrs): + self._name = name + for k, v in attrs.items(): + setattr(self, k, v) + + def get_name(self) -> str: + return self._name + + def get_quant_method(self, layer, prefix): + return MagicMock() + + @classmethod + def get_supported_act_dtypes(cls): + return [torch.bfloat16, torch.float16] + + def get_min_capability(self): + return 0 + + @classmethod + def from_config(cls, config): + raise NotImplementedError + + def get_config_filenames(self): + return [] + + +def _make_inc_config(block_names="thinker.model.layers,talker.model.layers", extra_config=None): + """Create a real OmniINCConfig with block_name_to_quantize.""" + return OmniINCConfig( + weight_bits=4, + group_size=128, + sym=True, + block_name_to_quantize=block_names, + extra_config=extra_config or {}, + ) + + +THINKER_MAPPER = WeightsMapper( + orig_to_new_prefix={ + "thinker.lm_head.": "language_model.lm_head.", + "thinker.model.": "language_model.model.", + "thinker.": "", + } +) + +TALKER_MAPPER = WeightsMapper( + orig_to_new_prefix={ + "talker.codec_head.": "language_model.lm_head.", + "talker.model.": "language_model.model.", + "talker.thinker_to_talker_proj.": "thinker_to_talker_proj.", + "talker.": "", + } +) + + +# =================================================================== +# 1. OmniINCConfig.apply_vllm_mapper +# =================================================================== + + +class TestApplyVllmMapper: + def test_inc_csv_string_normalized_to_list(self): + """CSV string block_name_to_quantize is split into a list.""" + cfg = _make_inc_config("thinker.model.layers,talker.model.layers") + cfg.apply_vllm_mapper(THINKER_MAPPER) + assert isinstance(cfg.block_name_to_quantize, list) + + def test_thinker_blocks_remapped(self): + """thinker.model.layers -> language_model.model.layers after apply_vllm_mapper.""" + cfg = _make_inc_config("thinker.model.layers,talker.model.layers") + cfg.apply_vllm_mapper(THINKER_MAPPER) + assert any("language_model.model.layers" in b for b in cfg.block_name_to_quantize) + + def test_cross_stage_blocks_kept_unchanged(self): + """Blocks not matching any mapper prefix are kept unchanged (harmless).""" + cfg = _make_inc_config("thinker.model.layers,talker.model.layers") + cfg.apply_vllm_mapper(THINKER_MAPPER) + # talker.model.layers doesn't match any thinker mapper prefix → stays as-is + assert "talker.model.layers" in cfg.block_name_to_quantize + + def test_talker_remap(self): + """talker.model.layers -> language_model.model.layers with talker mapper.""" + cfg = _make_inc_config("thinker.model.layers,talker.model.layers") + cfg.apply_vllm_mapper(TALKER_MAPPER) + assert any("language_model.model.layers" in b for b in cfg.block_name_to_quantize) + # thinker.model.layers doesn't match talker mapper → stays as-is + assert "thinker.model.layers" in cfg.block_name_to_quantize + + def test_extra_config_keys_remapped(self): + """Regex keys in extra_config get their escaped-dot prefixes remapped.""" + extra = { + r".*thinker\.model\.layers\.0\.mlp\.gate.*": {"bits": 16, "data_type": "float"}, + } + cfg = _make_inc_config("thinker.model.layers", extra_config=extra) + cfg.apply_vllm_mapper(THINKER_MAPPER) + # The key should now reference the vLLM runtime path + assert any("language_model" in k for k in cfg.extra_config) + # Original thinker\.model prefix should be replaced + assert not any(r"thinker\.model" in k for k in cfg.extra_config) + + def test_single_block_name(self): + """Only one block name (not CSV) still works.""" + cfg = _make_inc_config("thinker.model.layers") + cfg.apply_vllm_mapper(THINKER_MAPPER) + assert any("language_model.model.layers" in b for b in cfg.block_name_to_quantize) + + def test_already_list_block_names(self): + """block_name_to_quantize already a list (not CSV string) works.""" + cfg = _make_inc_config(["thinker.model.layers", "talker.model.layers"]) + cfg.apply_vllm_mapper(THINKER_MAPPER) + assert isinstance(cfg.block_name_to_quantize, list) + assert any("language_model.model.layers" in b for b in cfg.block_name_to_quantize) + + def test_mutates_in_place(self): + """apply_vllm_mapper mutates the config in place (same as upstream INCConfig).""" + cfg = _make_inc_config("thinker.model.layers") + original_id = id(cfg) + cfg.apply_vllm_mapper(THINKER_MAPPER) + assert id(cfg) == original_id + + # -- Stage prefix tests (runtime prefix = container + internal name) -- + + def test_thinker_block_has_stage_prefix(self): + """Mapped block name must start with 'thinker.' so runtime startswith() works.""" + cfg = _make_inc_config("thinker.model.layers,talker.model.layers") + cfg.apply_vllm_mapper(THINKER_MAPPER) + assert "thinker.language_model.model.layers" in cfg.block_name_to_quantize + + def test_talker_block_has_stage_prefix(self): + """Mapped block name must start with 'talker.' so runtime startswith() works.""" + cfg = _make_inc_config("thinker.model.layers,talker.model.layers") + cfg.apply_vllm_mapper(TALKER_MAPPER) + assert "talker.language_model.model.layers" in cfg.block_name_to_quantize + + def test_thinker_block_matches_runtime_prefix(self): + """Simulates get_layer_config's startswith() check for FusedMoE layers.""" + cfg = _make_inc_config("thinker.model.layers,talker.model.layers") + cfg.apply_vllm_mapper(THINKER_MAPPER) + runtime_prefix = "thinker.language_model.model.layers.0.mlp.experts" + assert any(runtime_prefix.startswith(b) for b in cfg.block_name_to_quantize) + + def test_talker_block_matches_runtime_prefix(self): + """Simulates get_layer_config's startswith() check for talker FusedMoE.""" + cfg = _make_inc_config("thinker.model.layers,talker.model.layers") + cfg.apply_vllm_mapper(TALKER_MAPPER) + runtime_prefix = "talker.language_model.model.layers.0.mlp.experts" + assert any(runtime_prefix.startswith(b) for b in cfg.block_name_to_quantize) + + def test_extra_config_plain_key_has_stage_prefix(self): + """Plain extra_config keys are remapped with stage prefix.""" + extra = { + "talker.model.layers.0.mlp.shared_expert_gate": {"bits": 16}, + } + cfg = _make_inc_config("talker.model.layers", extra_config=extra) + cfg.apply_vllm_mapper(TALKER_MAPPER) + assert "talker.language_model.model.layers.0.mlp.shared_expert_gate" in cfg.extra_config + + def test_extra_config_regex_key_still_works(self): + """Regex extra_config keys use re.search so no stage prefix needed.""" + import re + + extra = { + r".*thinker\.model\.layers\.0\.mlp\.gate.*": {"bits": 16}, + } + cfg = _make_inc_config("thinker.model.layers", extra_config=extra) + cfg.apply_vllm_mapper(THINKER_MAPPER) + runtime_name = "thinker.language_model.model.layers.0.mlp.gate" + matched = any(re.search(k, runtime_name) for k in cfg.extra_config) + assert matched + + +# =================================================================== +# 2. OmniINCConfig upgrade helpers +# =================================================================== + + +class TestOmniINCConfigUpgrade: + def test_maybe_upgrade_none(self): + assert OmniINCConfig.maybe_upgrade(None) is None + + def test_maybe_upgrade_non_inc(self): + """Non-INC configs are passed through unchanged.""" + cfg = _MockQuantConfig("fp8") + assert OmniINCConfig.maybe_upgrade(cfg) is cfg + + def test_maybe_upgrade_already_omni(self): + """Already OmniINCConfig is returned as-is.""" + cfg = _make_inc_config() + assert OmniINCConfig.maybe_upgrade(cfg) is cfg + + def test_maybe_upgrade_vanilla_inc(self): + """Vanilla INCConfig is promoted to OmniINCConfig.""" + from vllm.model_executor.layers.quantization.inc import INCConfig + + vanilla = INCConfig(weight_bits=4, group_size=128, sym=True) + upgraded = OmniINCConfig.maybe_upgrade(vanilla) + assert isinstance(upgraded, OmniINCConfig) + assert upgraded.weight_bits == 4 + assert upgraded.group_size == 128 + + +# =================================================================== +# 2. Three-branch thinker routing (simulated) +# =================================================================== + + +def _simulate_thinker_routing(quant_config): + """Simulate the three-branch routing in thinker __init__. + + Returns (visual_quant_config, language_quant_config, wrapped_vllm_quant). + """ + if isinstance(quant_config, ComponentQuantizationConfig): + visual_quant_config = quant_config.resolve("visual") + language_quant_config = quant_config.resolve("language_model") + return visual_quant_config, language_quant_config, quant_config + elif quant_config is not None: + if quant_config.get_name() in PRE_QUANTIZED_METHODS: + return quant_config, quant_config, quant_config + else: + language_quant_config = quant_config + wrapped = ComponentQuantizationConfig( + component_configs={"language_model": quant_config}, + default_config=None, + ) + return None, language_quant_config, wrapped + else: + return None, None, None + + +class TestThinkerRouting: + def test_none(self): + vis, lang, wrapped = _simulate_thinker_routing(None) + assert vis is None + assert lang is None + assert wrapped is None + + @pytest.mark.parametrize("method", ["modelopt", "modelopt_fp4", "modelopt_mxfp8"]) + def test_pre_quantized_all_components(self, method): + """Pre-quantized methods pass config to all components.""" + cfg = _MockQuantConfig(method) + vis, lang, wrapped = _simulate_thinker_routing(cfg) + assert vis is cfg + assert lang is cfg + assert wrapped is cfg + + def test_fp8_dynamic_language_only(self): + """fp8 dynamic: visual=None, language gets original config.""" + cfg = _MockQuantConfig("fp8") + vis, lang, wrapped = _simulate_thinker_routing(cfg) + assert vis is None + assert lang is cfg + assert isinstance(wrapped, ComponentQuantizationConfig) + assert wrapped.resolve("language_model") is cfg + assert wrapped.resolve("visual") is None + + def test_inc_autoround_language_only(self): + """INC/AutoRound: not in _PRE_QUANTIZED_METHODS -> wrapped like fp8.""" + cfg = _MockQuantConfig("inc") + vis, lang, wrapped = _simulate_thinker_routing(cfg) + assert vis is None + assert lang is cfg + assert isinstance(wrapped, ComponentQuantizationConfig) + + def test_component_config_passthrough(self): + """Explicit ComponentQuantizationConfig is used directly.""" + inner_fp8 = _MockQuantConfig("fp8") + inner_modelopt = _MockQuantConfig("modelopt") + cqc = ComponentQuantizationConfig( + component_configs={ + "visual": inner_modelopt, + "language_model": inner_fp8, + } + ) + vis, lang, wrapped = _simulate_thinker_routing(cqc) + assert vis is inner_modelopt + assert lang is inner_fp8 + assert wrapped is cqc + + +# =================================================================== +# 3. Talker visual routing (init_multi_modal guard) +# =================================================================== + + +def _simulate_talker_visual_routing(quant_config): + """Simulate the talker init_multi_modal visual routing.""" + if quant_config is not None and quant_config.get_name() in PRE_QUANTIZED_METHODS: + return quant_config + return None + + +class TestTalkerVisualRouting: + def test_none(self): + assert _simulate_talker_visual_routing(None) is None + + @pytest.mark.parametrize("method", ["modelopt", "modelopt_fp4", "modelopt_mxfp8"]) + def test_pre_quantized_passes_through(self, method): + """Pre-quantized methods pass quant config to visual.""" + cfg = _MockQuantConfig(method) + assert _simulate_talker_visual_routing(cfg) is cfg + + def test_fp8_blocked(self): + """fp8 dynamic must NOT be passed to visual.""" + cfg = _MockQuantConfig("fp8") + assert _simulate_talker_visual_routing(cfg) is None + + def test_inc_blocked(self): + """INC/AutoRound must NOT be passed to visual (not in _PRE_QUANTIZED_METHODS).""" + cfg = _MockQuantConfig("inc") + assert _simulate_talker_visual_routing(cfg) is None + + +# =================================================================== +# 4. ComponentQuantizationConfig.resolve +# =================================================================== + + +class TestComponentResolve: + def test_longest_prefix_match(self): + a = _MockQuantConfig("a") + b = _MockQuantConfig("b") + cqc = ComponentQuantizationConfig(component_configs={"language_model": a, "language_model.model": b}) + assert cqc.resolve("language_model.model.layers.0") is b + assert cqc.resolve("language_model.lm_head") is a + + def test_no_match_returns_default(self): + a = _MockQuantConfig("a") + default = _MockQuantConfig("default") + cqc = ComponentQuantizationConfig( + component_configs={"language_model": a}, + default_config=default, + ) + assert cqc.resolve("visual") is default + + def test_no_match_no_default_returns_none(self): + a = _MockQuantConfig("a") + cqc = ComponentQuantizationConfig( + component_configs={"language_model": a}, + ) + assert cqc.resolve("visual") is None + + def test_get_name(self): + cqc = ComponentQuantizationConfig(component_configs={}) + assert cqc.get_name() == "component" + + def test_get_quant_method_delegates(self): + """get_quant_method dispatches to the resolved config.""" + inner = _MockQuantConfig("fp8") + cqc = ComponentQuantizationConfig( + component_configs={"language_model": inner}, + ) + layer = MagicMock() + result = cqc.get_quant_method(layer, "language_model.model.layers.0.mlp") + assert result is not None # delegates to inner.get_quant_method + + def test_get_quant_method_returns_none_for_unmatched(self): + """get_quant_method returns None when no config matches.""" + inner = _MockQuantConfig("fp8") + cqc = ComponentQuantizationConfig( + component_configs={"language_model": inner}, + ) + layer = MagicMock() + result = cqc.get_quant_method(layer, "visual.blocks.0.mlp") + assert result is None + + def test_min_capability(self): + a = _MockQuantConfig("a") + a.get_min_capability = lambda: 80 + b = _MockQuantConfig("b") + b.get_min_capability = lambda: 70 + cqc = ComponentQuantizationConfig(component_configs={"x": a, "y": b}) + assert cqc.get_min_capability() == 70 + + def test_min_capability_empty(self): + cqc = ComponentQuantizationConfig(component_configs={}) + assert cqc.get_min_capability() == 0 diff --git a/tests/e2e/offline_inference/test_qwen2_5_omni_autoround_w4a16.py b/tests/e2e/offline_inference/test_qwen2_5_omni_autoround_w4a16.py new file mode 100644 index 000000000..e59a8c7d7 --- /dev/null +++ b/tests/e2e/offline_inference/test_qwen2_5_omni_autoround_w4a16.py @@ -0,0 +1,170 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""E2E tests for Qwen2.5-Omni AutoRound W4A16 quantized inference. + +These tests cover text, audio, image, video, and mixed-modality inputs +to verify multimodal understanding with quantized weights. + +Requirements: + - CUDA GPUs (4x L4 / 24 GB or equivalent) + - The quantized model checkpoint (Intel/Qwen2.5-Omni-7B-int4-AutoRound) +""" + +import os + +import pytest + +from tests.helpers.mark import hardware_test +from tests.helpers.media import ( + generate_synthetic_audio, + generate_synthetic_image, + generate_synthetic_video, +) +from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config + +QUANTIZED_MODEL = "Intel/Qwen2.5-Omni-7B-int4-AutoRound" +BASELINE_MODEL = "Qwen/Qwen2.5-Omni-7B" + +# Allow overriding via environment for local testing +QUANTIZED_MODEL = os.environ.get("QWEN2_5_OMNI_AUTOROUND_MODEL", QUANTIZED_MODEL) +BASELINE_MODEL = os.environ.get("QWEN2_5_OMNI_BASELINE_MODEL", BASELINE_MODEL) + +_CI_DEPLOY = get_deploy_config_path("ci/qwen2_5_omni.yaml") + + +def _get_stage_config(): + """Build a CI-friendly stage config with eager mode.""" + return modify_stage_config( + _CI_DEPLOY, + updates={ + "stages": { + 0: {"enforce_eager": True}, + 1: {"enforce_eager": True}, + }, + }, + ) + + +stage_config = _get_stage_config() + +# Parametrise: (model, stage_config) +quant_params = [(QUANTIZED_MODEL, stage_config)] + + +# ------------------------------------------------------------------ +# Test: text-only input → text output +# ------------------------------------------------------------------ + + +@pytest.mark.advanced_model +@pytest.mark.omni +@hardware_test(res={"cuda": "L4"}, num_cards=4) +@pytest.mark.parametrize("omni_runner", quant_params, indirect=True) +def test_text_to_text(omni_runner, omni_runner_handler): + """Text input → text output with W4A16 quantized Qwen2.5-Omni.""" + request_config = { + "prompts": "What is the capital of China?", + "modalities": ["text"], + } + response = omni_runner_handler.send_request(request_config) + assert response.success, f"Request failed: {response.error_message}" + assert response.text_content and len(response.text_content.strip()) > 0 + + +# ------------------------------------------------------------------ +# Test: audio input → text output +# ------------------------------------------------------------------ + + +@pytest.mark.advanced_model +@pytest.mark.omni +@hardware_test(res={"cuda": "L4"}, num_cards=4) +@pytest.mark.parametrize("omni_runner", quant_params, indirect=True) +def test_audio_to_text(omni_runner, omni_runner_handler): + """Audio input → text output with W4A16 quantized Qwen2.5-Omni.""" + audio = generate_synthetic_audio(1, 1, 16000)["np_array"] + if len(audio.shape) == 2: + audio = audio.squeeze() + + request_config = { + "prompts": "What is the content of this audio?", + "audios": (audio, 16000), + "modalities": ["text"], + } + response = omni_runner_handler.send_request(request_config) + assert response.success, f"Request failed: {response.error_message}" + assert response.text_content and len(response.text_content.strip()) > 0 + + +# ------------------------------------------------------------------ +# Test: image input → text output +# ------------------------------------------------------------------ + + +@pytest.mark.advanced_model +@pytest.mark.omni +@hardware_test(res={"cuda": "L4"}, num_cards=4) +@pytest.mark.parametrize("omni_runner", quant_params, indirect=True) +def test_image_to_text(omni_runner, omni_runner_handler): + """Image input → text output with W4A16 quantized Qwen2.5-Omni.""" + image = generate_synthetic_image(16, 16)["np_array"] + + request_config = { + "prompts": "Describe what you see in this image.", + "images": image, + "modalities": ["text"], + } + response = omni_runner_handler.send_request(request_config) + assert response.success, f"Request failed: {response.error_message}" + assert response.text_content and len(response.text_content.strip()) > 0 + + +# ------------------------------------------------------------------ +# Test: video input → text output +# ------------------------------------------------------------------ + + +@pytest.mark.advanced_model +@pytest.mark.omni +@hardware_test(res={"cuda": "L4"}, num_cards=4) +@pytest.mark.parametrize("omni_runner", quant_params, indirect=True) +def test_video_to_text(omni_runner, omni_runner_handler): + """Video input → text output with W4A16 quantized Qwen2.5-Omni.""" + video = generate_synthetic_video(16, 16, 30)["np_array"] + + request_config = { + "prompts": "Describe the video briefly.", + "videos": video, + "modalities": ["text"], + } + response = omni_runner_handler.send_request(request_config) + assert response.success, f"Request failed: {response.error_message}" + assert response.text_content and len(response.text_content.strip()) > 0 + + +# ------------------------------------------------------------------ +# Test: mixed modality (audio + image + video) → audio output +# ------------------------------------------------------------------ + + +@pytest.mark.advanced_model +@pytest.mark.omni +@hardware_test(res={"cuda": "L4"}, num_cards=4) +@pytest.mark.parametrize("omni_runner", quant_params, indirect=True) +def test_mix_to_audio(omni_runner, omni_runner_handler): + """Mixed-modality input → audio output with W4A16 quantized Qwen2.5-Omni.""" + video = generate_synthetic_video(16, 16, 30)["np_array"] + image = generate_synthetic_image(16, 16)["np_array"] + audio = generate_synthetic_audio(1, 1, 16000)["np_array"] + if len(audio.shape) == 2: + audio = audio.squeeze() + + request_config = { + "prompts": "What is recited in the audio? What is in this image? Describe the video briefly.", + "videos": video, + "images": image, + "audios": (audio, 16000), + "modalities": ["audio"], + } + response = omni_runner_handler.send_request(request_config) + assert response.success, f"Request failed: {response.error_message}" diff --git a/tests/e2e/offline_inference/test_qwen3_omni_autoround_w4a16.py b/tests/e2e/offline_inference/test_qwen3_omni_autoround_w4a16.py new file mode 100644 index 000000000..3a3c874b6 --- /dev/null +++ b/tests/e2e/offline_inference/test_qwen3_omni_autoround_w4a16.py @@ -0,0 +1,205 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""E2E tests for Qwen3-Omni AutoRound W4A16 quantized inference. + +These tests cover text, audio, image, video, and mixed-modality inputs +to verify multimodal understanding with quantized weights. + +Requirements: + - CUDA GPUs (2x H100-80G or equivalent) + - The quantized model checkpoint (Intel/Qwen3-Omni-30B-A3B-Instruct-int4-AutoRound) +""" + +import os + +import pytest + +from tests.helpers.mark import hardware_test +from tests.helpers.media import ( + generate_synthetic_audio, + generate_synthetic_image, + generate_synthetic_video, +) +from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config + +QUANTIZED_MODEL = "Intel/Qwen3-Omni-30B-A3B-Instruct-int4-AutoRound" +BASELINE_MODEL = "Qwen/Qwen3-Omni-30B-A3B-Instruct" + +# Allow overriding via environment for local testing +QUANTIZED_MODEL = os.environ.get("QWEN3_OMNI_AUTOROUND_MODEL", QUANTIZED_MODEL) +BASELINE_MODEL = os.environ.get("QWEN3_OMNI_BASELINE_MODEL", BASELINE_MODEL) + +_CI_DEPLOY = get_deploy_config_path("ci/qwen3_omni_moe.yaml") + + +@pytest.fixture(scope="module", autouse=True) +def _qwen3_omni_env(): + """Set env vars required by multi-stage worker spawning. + + Must run before CUDA context init. Reverted after every test module + so that values do not leak into unrelated test files. + """ + with pytest.MonkeyPatch.context() as mp: + mp.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn") + mp.setenv("VLLM_TEST_CLEAN_GPU_MEMORY", "0") + yield + + +def _get_stage_config(): + """Build a CI-friendly stage config with eager mode.""" + return modify_stage_config( + _CI_DEPLOY, + updates={ + "stages": { + 0: {"enforce_eager": True}, + 1: {"enforce_eager": True}, + }, + }, + ) + + +stage_config = _get_stage_config() + +# Parametrise: (model, stage_config) +quant_params = [(QUANTIZED_MODEL, stage_config)] + + +# ------------------------------------------------------------------ +# Test: text-only input → text output +# ------------------------------------------------------------------ + + +@pytest.mark.advanced_model +@pytest.mark.omni +@hardware_test(res={"cuda": "H100"}, num_cards=2) +@pytest.mark.parametrize("omni_runner", quant_params, indirect=True) +def test_text_to_text(omni_runner, omni_runner_handler): + """Text input → text output with W4A16 quantized Qwen3-Omni.""" + request_config = { + "prompts": "What is the capital of France?", + "modalities": ["text"], + } + response = omni_runner_handler.send_request(request_config) + assert response.success, f"Request failed: {response.error_message}" + assert response.text_content and len(response.text_content.strip()) > 0 + + +# ------------------------------------------------------------------ +# Test: audio input → text output +# ------------------------------------------------------------------ + + +@pytest.mark.advanced_model +@pytest.mark.omni +@hardware_test(res={"cuda": "H100"}, num_cards=2) +@pytest.mark.parametrize("omni_runner", quant_params, indirect=True) +def test_audio_to_text(omni_runner, omni_runner_handler): + """Audio input → text output with W4A16 quantized Qwen3-Omni.""" + audio = generate_synthetic_audio(1, 1, 16000)["np_array"] + if len(audio.shape) == 2: + audio = audio.squeeze() + + request_config = { + "prompts": "What is the content of this audio?", + "audios": (audio, 16000), + "modalities": ["text"], + } + response = omni_runner_handler.send_request(request_config) + assert response.success, f"Request failed: {response.error_message}" + assert response.text_content and len(response.text_content.strip()) > 0 + + +# ------------------------------------------------------------------ +# Test: image input → text output +# ------------------------------------------------------------------ + + +@pytest.mark.advanced_model +@pytest.mark.omni +@hardware_test(res={"cuda": "H100"}, num_cards=2) +@pytest.mark.parametrize("omni_runner", quant_params, indirect=True) +def test_image_to_text(omni_runner, omni_runner_handler): + """Image input → text output with W4A16 quantized Qwen3-Omni.""" + image = generate_synthetic_image(16, 16)["np_array"] + + request_config = { + "prompts": "Describe what you see in this image.", + "images": image, + "modalities": ["text"], + } + response = omni_runner_handler.send_request(request_config) + assert response.success, f"Request failed: {response.error_message}" + assert response.text_content and len(response.text_content.strip()) > 0 + + +# ------------------------------------------------------------------ +# Test: video input → text output +# ------------------------------------------------------------------ + + +@pytest.mark.advanced_model +@pytest.mark.omni +@hardware_test(res={"cuda": "H100"}, num_cards=2) +@pytest.mark.parametrize("omni_runner", quant_params, indirect=True) +def test_video_to_text(omni_runner, omni_runner_handler): + """Video input → text output with W4A16 quantized Qwen3-Omni.""" + video = generate_synthetic_video(224, 224, 300)["np_array"] + + request_config = { + "prompts": "Describe the video briefly.", + "videos": video, + "modalities": ["text"], + } + response = omni_runner_handler.send_request(request_config) + assert response.success, f"Request failed: {response.error_message}" + assert response.text_content and len(response.text_content.strip()) > 0 + + +# ------------------------------------------------------------------ +# Test: video input → audio output +# ------------------------------------------------------------------ + + +@pytest.mark.advanced_model +@pytest.mark.omni +@hardware_test(res={"cuda": "H100"}, num_cards=2) +@pytest.mark.parametrize("omni_runner", quant_params, indirect=True) +def test_video_to_audio(omni_runner, omni_runner_handler): + """Video input → audio output with W4A16 quantized Qwen3-Omni.""" + video = generate_synthetic_video(224, 224, 300)["np_array"] + + request_config = { + "prompts": "Describe the video briefly.", + "videos": video, + "modalities": ["audio"], + } + response = omni_runner_handler.send_request(request_config) + assert response.success, f"Request failed: {response.error_message}" + + +# ------------------------------------------------------------------ +# Test: mixed modality (audio + image + video) → audio output +# ------------------------------------------------------------------ + + +@pytest.mark.advanced_model +@pytest.mark.omni +@hardware_test(res={"cuda": "H100"}, num_cards=2) +@pytest.mark.parametrize("omni_runner", quant_params, indirect=True) +def test_mix_to_audio(omni_runner, omni_runner_handler): + """Mixed-modality input → audio output with W4A16 quantized Qwen3-Omni.""" + video = generate_synthetic_video(224, 224, 300)["np_array"] + image = generate_synthetic_image(16, 16)["np_array"] + audio = generate_synthetic_audio(1, 1, 16000)["np_array"] + if len(audio.shape) == 2: + audio = audio.squeeze() + + request_config = { + "prompts": "What is recited in the audio? What is in this image? Describe the video briefly.", + "videos": video, + "images": image, + "audios": (audio, 16000), + "modalities": ["audio"], + } + response = omni_runner_handler.send_request(request_config) + assert response.success, f"Request failed: {response.error_message}" diff --git a/vllm_omni/config/model.py b/vllm_omni/config/model.py index af89184a6..217a600b4 100644 --- a/vllm_omni/config/model.py +++ b/vllm_omni/config/model.py @@ -49,6 +49,25 @@ class OmniModelArchConfigConvertor(ModelArchConfigConvertorBase): if quant_cfg is not None: return quant_cfg + # Fall back to top-level quantization_config + top_quant = super().get_quantization_config() + if top_quant is not None: + block_names = top_quant.get("block_name_to_quantize") + if block_names is not None: + # NOTE: This assumes stage_config_name follows the HF + # ``_config`` convention (e.g. thinker_config → + # prefix "thinker."). removesuffix is a no-op when + # the suffix doesn't match, so a non-standard name + # would just use itself as prefix — safe but worth + # verifying if new stage names are introduced. + hf_prefix = self.stage_config_name.removesuffix("_config") + "." + if isinstance(block_names, str): + block_names = [b.strip() for b in block_names.split(",")] + if isinstance(block_names, list) and not any(b.startswith(hf_prefix) for b in block_names): + # This stage is not listed → no quantization. + return None + return top_quant + # For non-thinker stages (talker, code2wav) whose text_config # has no quantization_config, return None so quantization is # not applied to stages that were not quantized. diff --git a/vllm_omni/engine/stage_init_utils.py b/vllm_omni/engine/stage_init_utils.py index 94b9faa80..cc7676ba5 100644 --- a/vllm_omni/engine/stage_init_utils.py +++ b/vllm_omni/engine/stage_init_utils.py @@ -14,7 +14,7 @@ import multiprocessing as mp import os import time from collections.abc import Callable, Sequence -from dataclasses import dataclass +from dataclasses import dataclass, replace from typing import Any, Literal from vllm.logger import init_logger @@ -29,6 +29,7 @@ from vllm_omni.entrypoints.stage_utils import _to_dict, set_stage_devices from vllm_omni.entrypoints.utils import filter_dataclass_kwargs, resolve_model_config_path from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniSamplingParams from vllm_omni.platforms import current_omni_platform +from vllm_omni.quantization.inc_config import OmniINCConfig logger = init_logger(__name__) @@ -459,6 +460,11 @@ def build_vllm_config( ) executor_class = Executor.get_class(vllm_config) + # Upgrade vanilla INCConfig to OmniINCConfig for multi-stage models. + upgraded = OmniINCConfig.maybe_upgrade(vllm_config.quant_config) + if upgraded is not vllm_config.quant_config: + vllm_config = replace(vllm_config, quant_config=upgraded) + return vllm_config, executor_class diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py index bb491d01b..28f49918f 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py @@ -100,11 +100,23 @@ class Qwen3OmniMoeTalkerForConditionalGeneration( def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() talker_config: Qwen3OmniMoeTalkerConfig = vllm_config.model_config.hf_config - rope_params = talker_config.text_config.rope_scaling + rope_params = getattr(talker_config.text_config, "rope_scaling", None) if rope_params is None: - # Newer transformers use rope_parameters instead of rope_scaling rope_params = getattr(talker_config.text_config, "rope_parameters", None) or {} - rope_params["rope_theta"] = talker_config.text_config.rope_theta + rope_params = dict(rope_params) + # In transformers <5.0.0, rope_theta is a top-level config attribute + # (e.g. config.text_config.rope_theta = 1000000.0). + # In transformers >=5.0.0 (PR #39847), rope_theta moved inside the + # rope_parameters dict (e.g. config.text_config.rope_parameters = + # {"rope_theta": 1000000.0, "rope_type": "default"}). + # Use setdefault so we never overwrite a value already present. + # Precedence: rope_params["rope_theta"] (already set) + # > text_config.rope_theta (transformers <5.0.0 top-level attr) + # > 1000000 (Qwen3 Omni default) + rope_params.setdefault( + "rope_theta", + getattr(talker_config.text_config, "rope_theta", 1000000), + ) talker_config.text_config.rope_parameters = rope_params quant_config = vllm_config.quant_config if isinstance(quant_config, ComponentQuantizationConfig): diff --git a/vllm_omni/quantization/__init__.py b/vllm_omni/quantization/__init__.py index b50fe8607..a709dc74d 100644 --- a/vllm_omni/quantization/__init__.py +++ b/vllm_omni/quantization/__init__.py @@ -13,6 +13,7 @@ Adds per-component quantization for multi-stage models. from .component_config import ComponentQuantizationConfig from .factory import SUPPORTED_QUANTIZATION_METHODS, build_quant_config +from .inc_config import OmniINCConfig # DiffusionGGUFConfig is NOT imported here to avoid pulling in # GGUF -> fused_moe -> pynvml at module load time. @@ -20,5 +21,6 @@ from .factory import SUPPORTED_QUANTIZATION_METHODS, build_quant_config __all__ = [ "build_quant_config", "ComponentQuantizationConfig", + "OmniINCConfig", "SUPPORTED_QUANTIZATION_METHODS", ] diff --git a/vllm_omni/quantization/factory.py b/vllm_omni/quantization/factory.py index f85589d69..5af47a1f7 100644 --- a/vllm_omni/quantization/factory.py +++ b/vllm_omni/quantization/factory.py @@ -43,16 +43,16 @@ def _build_int8(**kw: Any) -> QuantizationConfig: def _build_inc(**kw: Any) -> QuantizationConfig: """Lazy import for INC/AutoRound config with checkpoint kwarg normalization.""" - from vllm.model_executor.layers.quantization.inc import INCConfig + from .inc_config import OmniINCConfig # Map checkpoint key 'bits' to INCConfig's 'weight_bits' if "bits" in kw and "weight_bits" not in kw: kw["weight_bits"] = kw.pop("bits") # Filter to only valid INCConfig params - valid = set(inspect.signature(INCConfig.__init__).parameters) - {"self"} + valid = set(inspect.signature(OmniINCConfig.__init__).parameters) - {"self"} filtered = {k: v for k, v in kw.items() if k in valid} - return INCConfig(**filtered) + return OmniINCConfig(**filtered) _OVERRIDES: dict[str, Callable[..., QuantizationConfig]] = { diff --git a/vllm_omni/quantization/inc_config.py b/vllm_omni/quantization/inc_config.py new file mode 100644 index 000000000..fb4851cca --- /dev/null +++ b/vllm_omni/quantization/inc_config.py @@ -0,0 +1,140 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Extended INC/AutoRound config for multi-stage omni models.""" + +from __future__ import annotations + +from os.path import commonprefix +from typing import TYPE_CHECKING, Any + +from vllm.model_executor.layers.quantization.inc import INCConfig +from vllm.model_executor.models.utils import WeightsMapper + +if TYPE_CHECKING: + from vllm.model_executor.layers.quantization.base_config import ( + QuantizationConfig, + ) + +_REGEX_SPECIAL_CHARS = frozenset(r"*+?^$()[]{}|\\") + + +def _stage_prefix(prefix_map: dict[str, str | None]) -> str: + """Derive the container/stage prefix from mapper source keys.""" + cp = commonprefix(list(prefix_map.keys())) + dot = cp.rfind(".") + return cp[: dot + 1] if dot >= 0 else "" + + +def _map_with_stage_prefix( + items: list[str], + prefix_map: dict[str, str | None], + stage: str, +) -> list[str]: + """Apply *prefix_map* to each item and prepend *stage* to mapped items.""" + sorted_keys = sorted(prefix_map, key=len, reverse=True) + result: list[str] = [] + for item in items: + new_item = item + for orig in sorted_keys: + if item.startswith(orig): + new_val = prefix_map[orig] or "" + new_item = stage + new_val + item[len(orig) :] + break + result.append(new_item) + return result + + +class OmniINCConfig(INCConfig): + """INCConfig extended with multi-stage prefix remapping.""" + + # ------------------------------------------------------------------ + # Core integration: called by vLLM's configure_quant_config() + # ------------------------------------------------------------------ + + def apply_vllm_mapper(self, hf_to_vllm_mapper: WeightsMapper) -> None: + """Remap HF checkpoint names to vLLM runtime prefixes.""" + prefix_map = getattr(hf_to_vllm_mapper, "orig_to_new_prefix", None) or {} + stage = _stage_prefix(prefix_map) if prefix_map else "" + + # -- Normalize CSV string ----------------------------------------- + if isinstance(self.block_name_to_quantize, str): + self.block_name_to_quantize = [b.strip() for b in self.block_name_to_quantize.split(",") if b.strip()] + + # -- block_name_to_quantize ---------------------------------------- + if self.block_name_to_quantize is not None: + if prefix_map and stage: + self.block_name_to_quantize = _map_with_stage_prefix( + self.block_name_to_quantize, + prefix_map, + stage, + ) + else: + self.block_name_to_quantize = hf_to_vllm_mapper.apply_list(self.block_name_to_quantize) + + # -- extra_config -------------------------------------------------- + if self.extra_config is not None and prefix_map: + new_extra: dict[str, Any] = {} + sorted_keys = sorted(prefix_map, key=len, reverse=True) + + # Build escaped-dot map for regex pattern keys + escaped_map: dict[str, str] = {} + for orig, new in prefix_map.items(): + escaped_map[orig.replace(".", r"\.")] = (new or "").replace(".", r"\.") + escaped_sorted = sorted(escaped_map, key=len, reverse=True) + + for key, val in self.extra_config.items(): + is_regex = any(c in _REGEX_SPECIAL_CHARS for c in key) + if is_regex: + # Regex keys: escaped-dot substring replacement. + # re.search matches anywhere so no stage prefix needed. + new_key = key + for esc_orig in escaped_sorted: + if esc_orig in new_key: + new_key = new_key.replace( + esc_orig, + escaped_map[esc_orig], + 1, + ) + break + else: + # Plain keys: prefix replacement + stage prefix. + new_key = key + for orig in sorted_keys: + if key.startswith(orig): + new_val = prefix_map[orig] or "" + new_key = stage + new_val + key[len(orig) :] + break + new_extra[new_key] = val + self.extra_config = new_extra + elif self.extra_config is not None: + self.extra_config = hf_to_vllm_mapper.apply_dict(self.extra_config) + + # ------------------------------------------------------------------ + # Upgrading a vanilla INCConfig created by vLLM + # ------------------------------------------------------------------ + + @classmethod + def from_inc_config(cls, inc: INCConfig) -> OmniINCConfig: + """Promote a vanilla :class:`INCConfig` to :class:`OmniINCConfig`. + + Copies all attributes so that the new instance is a drop-in + replacement. + """ + omni = object.__new__(cls) + omni.__dict__.update(inc.__dict__) + return omni + + @classmethod + def maybe_upgrade(cls, quant_config: QuantizationConfig | None) -> QuantizationConfig | None: + """Upgrade *quant_config* to :class:`OmniINCConfig` if applicable. + + Returns the original config unchanged when it is not an INC + config or is already an :class:`OmniINCConfig`. + """ + if quant_config is None: + return None + if isinstance(quant_config, cls): + return quant_config + if isinstance(quant_config, INCConfig): + return cls.from_inc_config(quant_config) + return quant_config From a664e407168713e8752b8535873c773111b0352d Mon Sep 17 00:00:00 2001 From: Ricardo Noriega Date: Wed, 22 Apr 2026 15:08:19 +0200 Subject: [PATCH 12/45] Fix links in glm_image.md (#3028) Signed-off-by: Ricardo Noriega --- docs/user_guide/examples/online_serving/glm_image.md | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/docs/user_guide/examples/online_serving/glm_image.md b/docs/user_guide/examples/online_serving/glm_image.md index f7027b906..b67eed1e0 100644 --- a/docs/user_guide/examples/online_serving/glm_image.md +++ b/docs/user_guide/examples/online_serving/glm_image.md @@ -104,9 +104,9 @@ guides. When using `/v1/chat/completions`, pass these inside `extra_body` in the curl JSON, or via the `extra_body` keyword argument in the OpenAI Python SDK (see the -[Diffusion Chat API guide](../../../../serving/diffusion_chat_api.md)). -When using the dedicated [`/v1/images/generations`](../../../../serving/image_generation_api.md) -or [`/v1/images/edits`](../../../../serving/image_edit_api.md) endpoints, pass +[Diffusion Chat API guide](../../../serving/diffusion_chat_api.md)). +When using the dedicated [`/v1/images/generations`](../../../serving/image_generation_api.md) +or [`/v1/images/edits`](../../../serving/image_edit_api.md) endpoints, pass the supported generation controls as top-level fields directly. For image dimensions and count, use `size` and `n` rather than `height` or `width`. From d8cc7a04096aa58df0d4a149aebc2f72ea37970e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Zeyu=20Huang=20=7C=20=E9=BB=83=E6=BE=A4=E5=AE=87?= <11222265+fhfuih@users.noreply.github.com> Date: Thu, 23 Apr 2026 00:07:18 +0800 Subject: [PATCH 13/45] [feat]: General diffusers adapter backend to run diffusion models (#2724) Signed-off-by: Huang, Zeyu <11222265+fhfuih@users.noreply.github.com> --- docs/.nav.yml | 1 + .../diffusers_pipeline_adapter.md | 93 +++++ .../diffusers_pipeline_adapter/README.md | 83 ++++ .../stage_config.yaml | 31 ++ tests/diffusion/test_diffusers_adapter.py | 186 +++++++++ .../online_serving/test_diffusers_adapter.py | 56 +++ vllm_omni/diffusion/data.py | 77 ++-- .../model_loader/diffusers_loader.py | 18 +- .../models/diffusers_adapter/__init__.py | 12 + .../pipeline_diffusers_adapter.py | 358 ++++++++++++++++++ vllm_omni/diffusion/registry.py | 5 + vllm_omni/engine/async_omni_engine.py | 6 + vllm_omni/entrypoints/cli/serve.py | 34 ++ 13 files changed, 936 insertions(+), 24 deletions(-) create mode 100644 docs/user_guide/examples/online_serving/diffusers_pipeline_adapter.md create mode 100644 examples/online_serving/diffusers_pipeline_adapter/README.md create mode 100644 examples/online_serving/diffusers_pipeline_adapter/stage_config.yaml create mode 100644 tests/diffusion/test_diffusers_adapter.py create mode 100644 tests/e2e/online_serving/test_diffusers_adapter.py create mode 100644 vllm_omni/diffusion/models/diffusers_adapter/__init__.py create mode 100644 vllm_omni/diffusion/models/diffusers_adapter/pipeline_diffusers_adapter.py diff --git a/docs/.nav.yml b/docs/.nav.yml index b0022965a..aebd560ca 100644 --- a/docs/.nav.yml +++ b/docs/.nav.yml @@ -34,6 +34,7 @@ nav: - Online Serving: - BAGEL-7B-MoT: user_guide/examples/online_serving/bagel.md - vLLM-Omni Helm Chart: user_guide/examples/online_serving/chart-helm.md + - Diffusers Backend Adapter Example: user_guide/examples/online_serving/diffusers_pipeline_adapter.md - Fish Speech S2 Pro: user_guide/examples/online_serving/fish_speech.md - GLM-Image Online Serving: user_guide/examples/online_serving/glm_image.md - Image-To-Image: user_guide/examples/online_serving/image_to_image.md diff --git a/docs/user_guide/examples/online_serving/diffusers_pipeline_adapter.md b/docs/user_guide/examples/online_serving/diffusers_pipeline_adapter.md new file mode 100644 index 000000000..ac88071d5 --- /dev/null +++ b/docs/user_guide/examples/online_serving/diffusers_pipeline_adapter.md @@ -0,0 +1,93 @@ +# Diffusers Backend Adapter Example + +Source . + + +This example demonstrates how to serve any 🤗 Diffusers pipeline through vLLM-Omni +using the `diffusers` load format. + +## Supported Models + +Any model loadable via `DiffusionPipeline.from_pretrained()` should be supported, including text-to-image, image-to-image, text-to-video, image-to-video, and text-to-audio. + +## Limitations + +The diffusers backend is a black-box adapter. The following features are NOT yet supported. +It is not guaranteed whether they will be supported in the future. + +- CFG parallel execution +- Sequence parallel execution +- TeaCache / Cache-DiT acceleration +- Step-wise execution (continuous batching) + +For these features, it is recommended to use natively supported pipelines instead. + +## Usage + +### Option 1: CLI arguments + +```bash +vllm serve "stable-diffusion-v1-5/stable-diffusion-v1-5" \ + --omni \ + --diffusion-load-format diffusers \ + --diffusers-load-kwargs '{"use_safetensors": true}' \ + --diffusers-call-kwargs '{"num_inference_steps": 30, "guidance_scale": 7.5}' +``` + +`--diffusers-load-kwargs` and `--diffusers-call-kwargs` are only valid together with `--diffusion-load-format diffusers`. + +### Option 2: Stage config YAML + +```bash +vllm serve stable-diffusion-v1-5/stable-diffusion-v1-5 --stage-configs-path examples/online_serving/diffusers_pipeline_adapter/stage_config.yaml --omni +``` + +The particular fields of interest are `model`, `diffusion_load_format`, `diffusers_load_kwargs`, and `diffusers_call_kwargs` under `engine_args`. They are the same as the CLI arguments. + +## Send a Request + +```bash +curl http://localhost:8000/v1/images/generations \ + -H "Content-Type: application/json" \ + -d '{ + "model": "stable-diffusion-v1-5/stable-diffusion-v1-5", + "prompt": "a photo of an astronaut riding a horse on mars", + "n": 1, + "size": "512x512" + }' +``` + +Or refer to other documentation pages on how to request a particular input/output modality, such as `examples/online_serving/text_to_image/openai_chat_client.py`. + +## Configuration Reference + +For the diffusers adapter, set options under **`engine_args`**: + +### `diffusion_load_format: "diffusers"` + +This field selects the Hugging Face diffusers adapter path (see `DiffusersPipelineLoader`). + +### `diffusers_load_kwargs` + +Passed to `DiffusionPipeline.from_pretrained()`. + +This is suitable for model-specific configurations not available through the vLLM-Omni interface (such as `Omni.__init__()`, `vllm serve` CLI arguments, and stage config YAML fields outside `diffusers_load_kwargs`). + +When a parameter is available in the vLLM-Omni interface, it will be adapted here. +But if that parameter is simultaneously set in both the vLLM-Omni interface and `diffusers_load_kwargs`, the **latter** will take precedence. + +### `diffusers_call_kwargs` + +Passed to `pipeline.__call__()`. + +This is suitable for sampling parameters not available through the vLLM-Omni interface (such as `Omni.generate()` and online serving payloads). + +When a parameter is available in the vLLM-Omni interface, it will be adapted here. +But if that parameter is simultaneously set in both the vLLM-Omni interface and `diffusers_call_kwargs`, the **former** will take precedence (because it is set at request time). + +## Example materials + +??? abstract "stage_config.yaml" + ``````yaml + --8<-- "examples/online_serving/diffusers_pipeline_adapter/stage_config.yaml" + `````` diff --git a/examples/online_serving/diffusers_pipeline_adapter/README.md b/examples/online_serving/diffusers_pipeline_adapter/README.md new file mode 100644 index 000000000..8dbf9369a --- /dev/null +++ b/examples/online_serving/diffusers_pipeline_adapter/README.md @@ -0,0 +1,83 @@ +# Diffusers Backend Adapter Example + +This example demonstrates how to serve any 🤗 Diffusers pipeline through vLLM-Omni +using the `diffusers` load format. + +## Supported Models + +Any model loadable via `DiffusionPipeline.from_pretrained()` should be supported, including text-to-image, image-to-image, text-to-video, image-to-video, and text-to-audio. + +## Limitations + +The diffusers backend is a black-box adapter. The following features are NOT yet supported. +It is not guaranteed whether they will be supported in the future. + +- CFG parallel execution +- Sequence parallel execution +- TeaCache / Cache-DiT acceleration +- Step-wise execution (continuous batching) + +For these features, it is recommended to use natively supported pipelines instead. + +## Usage + +### Option 1: CLI arguments + +```bash +vllm serve "stable-diffusion-v1-5/stable-diffusion-v1-5" \ + --omni \ + --diffusion-load-format diffusers \ + --diffusers-load-kwargs '{"use_safetensors": true}' \ + --diffusers-call-kwargs '{"num_inference_steps": 30, "guidance_scale": 7.5}' +``` + +`--diffusers-load-kwargs` and `--diffusers-call-kwargs` are only valid together with `--diffusion-load-format diffusers`. + +### Option 2: Stage config YAML + +```bash +vllm serve stable-diffusion-v1-5/stable-diffusion-v1-5 --stage-configs-path examples/online_serving/diffusers_pipeline_adapter/stage_config.yaml --omni +``` + +The particular fields of interest are `model`, `diffusion_load_format`, `diffusers_load_kwargs`, and `diffusers_call_kwargs` under `engine_args`. They are the same as the CLI arguments. + +## Send a Request + +```bash +curl http://localhost:8000/v1/images/generations \ + -H "Content-Type: application/json" \ + -d '{ + "model": "stable-diffusion-v1-5/stable-diffusion-v1-5", + "prompt": "a photo of an astronaut riding a horse on mars", + "n": 1, + "size": "512x512" + }' +``` + +Or refer to other documentation pages on how to request a particular input/output modality, such as `examples/online_serving/text_to_image/openai_chat_client.py`. + +## Configuration Reference + +For the diffusers adapter, set options under **`engine_args`**: + +### `diffusion_load_format: "diffusers"` + +This field selects the Hugging Face diffusers adapter path (see `DiffusersPipelineLoader`). + +### `diffusers_load_kwargs` + +Passed to `DiffusionPipeline.from_pretrained()`. + +This is suitable for model-specific configurations not available through the vLLM-Omni interface (such as `Omni.__init__()`, `vllm serve` CLI arguments, and stage config YAML fields outside `diffusers_load_kwargs`). + +When a parameter is available in the vLLM-Omni interface, it will be adapted here. +But if that parameter is simultaneously set in both the vLLM-Omni interface and `diffusers_load_kwargs`, the **latter** will take precedence. + +### `diffusers_call_kwargs` + +Passed to `pipeline.__call__()`. + +This is suitable for sampling parameters not available through the vLLM-Omni interface (such as `Omni.generate()` and online serving payloads). + +When a parameter is available in the vLLM-Omni interface, it will be adapted here. +But if that parameter is simultaneously set in both the vLLM-Omni interface and `diffusers_call_kwargs`, the **former** will take precedence (because it is set at request time). diff --git a/examples/online_serving/diffusers_pipeline_adapter/stage_config.yaml b/examples/online_serving/diffusers_pipeline_adapter/stage_config.yaml new file mode 100644 index 000000000..7c96eb6c1 --- /dev/null +++ b/examples/online_serving/diffusers_pipeline_adapter/stage_config.yaml @@ -0,0 +1,31 @@ +# Example stage config for diffusers backend +# This config demonstrates serving Stable Diffusion 1.5 via the diffusers adapter. +# Users should copy and modify this for their own models. + +model_type: diffusion + +stage_args: + - stage_id: 0 + stage_type: diffusion + engine_args: + model_stage: diffusion + model: "stable-diffusion-v1-5/stable-diffusion-v1-5" + distributed_executor_backend: "mp" + # gpu_memory_utilization: 0.9 + engine_output_type: image + # Select the HF diffusers adapter + diffusion_load_format: "diffusers" + # model_class_name: "DiffusersAdapterPipeline" # default when diffusion_load_format is diffusers + diffusers_load_kwargs: + # Passed to DiffusionPipeline.from_pretrained(). + # Good for model-specific loading parameters not covered by OmniDiffusionConfig. + # During model load time, parameters here override their counterparts in the vLLM-Omni interface. + use_safetensors: true + diffusers_call_kwargs: + # Passed to pipeline.__call__(). + # Good for model-specific sampling parameters not covered by OmniDiffusionSamplingParams. + # During request time, parameters here are overridden by the counterparts in OmniDiffusionSamplingParams. + num_inference_steps: 30 + guidance_scale: 7.5 + final_output: true + final_output_type: image diff --git a/tests/diffusion/test_diffusers_adapter.py b/tests/diffusion/test_diffusers_adapter.py new file mode 100644 index 000000000..ac2ec2e3f --- /dev/null +++ b/tests/diffusion/test_diffusers_adapter.py @@ -0,0 +1,186 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from collections import namedtuple +from types import SimpleNamespace + +import pytest +import torch +from diffusers import DiffusionPipeline +from PIL import Image + +from vllm_omni.diffusion.data import ( + DiffusionOutput, + DiffusionParallelConfig, + OmniDiffusionConfig, +) +from vllm_omni.diffusion.models.diffusers_adapter import DiffusersAdapterPipeline +from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.inputs.data import OmniDiffusionSamplingParams + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + + +def _make_od_config(**overrides) -> OmniDiffusionConfig: + od_config = OmniDiffusionConfig( + model="test/model", + model_class_name="DiffusersAdapterPipeline", + dtype=torch.float16, + diffusion_load_format="diffusers", + diffusers_load_kwargs={}, + diffusers_call_kwargs={}, + output_type="pil", + parallel_config=DiffusionParallelConfig(cfg_parallel_size=1, sequence_parallel_size=1), + cache_backend="none", + ) + for key, value in overrides.items(): + setattr(od_config, key, value) + return od_config + + +def _make_request(**overrides) -> OmniDiffusionRequest: + prompt = overrides.pop("prompt", "a test prompt") + negative_prompt = overrides.pop("negative_prompt", None) + prompt_obj: dict[str, str] = {"prompt": prompt} + if negative_prompt is not None: + prompt_obj["negative_prompt"] = negative_prompt + + defaults = { + "prompts": [prompt_obj], + "sampling_params": OmniDiffusionSamplingParams( + num_inference_steps=20, + guidance_scale=7.5, + height=16, + width=16, + num_frames=1, + num_outputs_per_prompt=1, + seed=42, + output_type="pil", + generator_device="cpu", + ), + } + defaults.update(overrides) + return OmniDiffusionRequest(**defaults) + + +class TestDiffusersAdapterPipeline: + def test_adapter_forward_returns_output(self, mocker): + od_config = _make_od_config() + request = _make_request() + stub_image = Image.new("RGB", (request.sampling_params.width, request.sampling_params.height)) # pyright: ignore[reportArgumentType] + + adapter = DiffusersAdapterPipeline(od_config=od_config) + MockPipelineOutput = namedtuple("MockPipelineOutput", ["image"]) + MockPipeline = type("MockPipeline", (DiffusionPipeline,), {}) + adapter._pipeline = MockPipeline() + + mocker.patch.object( + MockPipeline, + "__call__", + return_value=MockPipelineOutput(image=stub_image), + ) + output = adapter.forward(request) + + assert isinstance(output, DiffusionOutput) + assert isinstance(output.output, MockPipelineOutput) + assert output.output.image is stub_image + + @pytest.mark.parametrize( + "feature_id", + ["cfg_parallel", "ulysses", "ring", "teacache", "cache_dit", "enforce_eager", "quantization"], + ) + def test_adapter_guard_unsupported_feature(self, feature_id): + if feature_id == "cfg_parallel": + od_config = _make_od_config( + parallel_config=DiffusionParallelConfig(cfg_parallel_size=2, sequence_parallel_size=1), + cache_backend="none", + ) + elif feature_id == "ulysses": + od_config = _make_od_config( + parallel_config=DiffusionParallelConfig(cfg_parallel_size=1, ulysses_degree=2), + cache_backend="none", + ) + elif feature_id == "ring": + od_config = _make_od_config( + parallel_config=DiffusionParallelConfig(cfg_parallel_size=1, ring_degree=2), + cache_backend="none", + ) + elif feature_id == "teacache": + od_config = _make_od_config( + parallel_config=DiffusionParallelConfig(cfg_parallel_size=1, sequence_parallel_size=1), + cache_backend="tea_cache", + ) + elif feature_id == "cache_dit": + od_config = _make_od_config( + parallel_config=DiffusionParallelConfig(cfg_parallel_size=1, sequence_parallel_size=1), + cache_backend="cache_dit", + ) + elif feature_id == "enforce_eager": + od_config = _make_od_config(enforce_eager=True) + elif feature_id == "quantization": + od_config = _make_od_config(quantization_config=SimpleNamespace(quant_method="fp8")) + else: + raise ValueError(f"Unknown feature ID: {feature_id}") + + with pytest.raises(NotImplementedError): + DiffusersAdapterPipeline(od_config=od_config) + + def test_adapter_guard_unknown_output_type(self, mocker): + """Test that the adapter wraps an unknown output type as-is. + This is useful when `return_dict=True` and the diffusers pipeline returns an OrderedDict subclass.""" + + adapter = DiffusersAdapterPipeline(od_config=_make_od_config()) + raw_output = {"unexpected": "dict-output"} + + MockPipeline = type("MockPipeline", (DiffusionPipeline,), {}) + adapter._pipeline = MockPipeline() + + mocker.patch.object( + MockPipeline, + "__call__", + return_value=raw_output, + ) + output = adapter.forward(_make_request()) + + assert isinstance(output, DiffusionOutput) + assert output.output == raw_output + + def test_adapter_build_call_kwargs(self): + adapter = DiffusersAdapterPipeline( + od_config=_make_od_config( + diffusers_call_kwargs={ + "guidance_scale": 1.25, + "eta": 0.3, + "output_type": "np", + } + ) + ) + req = _make_request( + prompt="a cat on mars", + negative_prompt="low quality", + sampling_params=OmniDiffusionSamplingParams( + num_inference_steps=9, + guidance_scale=8.0, + height=320, + width=640, + num_frames=8, + num_outputs_per_prompt=2, + seed=123, + output_type="pil", + ), + ) + + kwargs = adapter._build_call_kwargs(req) + + assert kwargs["prompt"] == "a cat on mars" + assert kwargs["negative_prompt"] == "low quality" + assert kwargs["num_inference_steps"] == 9 + assert kwargs["guidance_scale"] == 8.0 + assert kwargs["height"] == 320 + assert kwargs["width"] == 640 + assert kwargs["num_frames"] == 8 + assert kwargs["num_images_per_prompt"] == 2 + assert kwargs["output_type"] == "pil" + assert isinstance(kwargs["generator"], torch.Generator) + assert kwargs["generator"].device.type == "cpu" + assert kwargs["generator"].initial_seed() == 123 diff --git a/tests/e2e/online_serving/test_diffusers_adapter.py b/tests/e2e/online_serving/test_diffusers_adapter.py new file mode 100644 index 000000000..8b41db13a --- /dev/null +++ b/tests/e2e/online_serving/test_diffusers_adapter.py @@ -0,0 +1,56 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""End-to-end tests for DiffusersAdapterPipeline. + +It tests the full user flow of launching a diffusers-backed model and running inference. +""" + +import pytest +from PIL import Image + +from tests.helpers.mark import hardware_test +from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler, dummy_messages_from_mix_data + +pytestmark = [pytest.mark.diffusion, pytest.mark.core_model] + + +@pytest.mark.parametrize( + "omni_server", + [ + OmniServerParams( + model="tiny-random/Qwen-Image", + server_args=[ + "--diffusion-load-format", + "diffusers", + "--diffusers-call-kwargs", + '{"height": 512, "width": 0}', # deliberately weird width to be overridden + ], + ), + ], + indirect=True, +) +@hardware_test(res={"cuda": "L4"}, num_cards=1) +def test_t2i_with_diffusers_adapter( + omni_server: OmniServer, + openai_client: OpenAIClientHandler, +): + messages = dummy_messages_from_mix_data(content_text="a photo of an astronaut riding a horse on mars") + + request_config = { + "model": omni_server.model, + "messages": messages, + "extra_body": { + "width": 512, + "num_inference_steps": 2, + "negative_prompt": "blurry", + "true_cfg_scale": 4.0, + "seed": 42, + }, + } + + response = openai_client.send_diffusion_request(request_config) + image: Image.Image = response[0].images[0] # pyright: ignore[reportOptionalSubscript] + + # Request config has incomplete width/height, so internal assertion in `send_diffusion_request` is incomplete. + assert image.size == (512, 512) diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index 012c0130c..b9307657f 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -452,7 +452,14 @@ class OmniDiffusionConfig: custom_pipeline_args: dict[str, Any] | None = None # Diffusion model loading format - diffusion_load_format: str = "default" # "default", "custom_pipeline", "dummy" + # "default", "custom_pipeline", "dummy", "diffusers" (HF diffusers adapter) + diffusion_load_format: str = "default" + + # Diffusers adapter kwargs + # kwargs forwarded to DiffusionPipeline.from_pretrained() + diffusers_load_kwargs: dict[str, Any] = field(default_factory=dict) + # kwargs forwarded to pipeline.__call__() + diffusers_call_kwargs: dict[str, Any] = field(default_factory=dict) # http server endpoint config, would be ignored in local mode host: str | None = None @@ -648,6 +655,12 @@ class OmniDiffusionConfig: elif self.max_cpu_loras < 1: raise ValueError("max_cpu_loras must be >= 1 for diffusion LoRA") + if self.diffusion_load_format != "diffusers" and (self.diffusers_load_kwargs or self.diffusers_call_kwargs): + raise ValueError( + "diffusers_load_kwargs and diffusers_call_kwargs are only " + "valid together with diffusion_load_format=diffusers" + ) + def set_tf_model_config(self, tf_config: "TransformerConfig") -> None: """Assign `tf_model_config` and propagate quantization if detected. @@ -686,6 +699,10 @@ class OmniDiffusionConfig: """ from vllm.transformers_utils.config import get_hf_file_to_dict + # Default model_class_name for diffusers adapter + if self.model_class_name is None and self.diffusion_load_format == "diffusers": + self.model_class_name = "DiffusersAdapterPipeline" + try: config_dict = get_hf_file_to_dict("model_index.json", self.model) if config_dict is not None: @@ -693,32 +710,42 @@ class OmniDiffusionConfig: self.model_class_name = config_dict.get("_class_name", None) self.update_multimodal_support() - tf_config_dict = get_hf_file_to_dict("transformer/config.json", self.model) - self.tf_model_config = TransformerConfig.from_dict(tf_config_dict) + # Skip transformer config loading for diffusers adapter + # (non-DiT models don't have a separate transformer folder/config) + if self.diffusion_load_format == "diffusers": + self.tf_model_config = TransformerConfig() + else: + tf_config_dict = get_hf_file_to_dict("transformer/config.json", self.model) + self.tf_model_config = TransformerConfig.from_dict(tf_config_dict) else: raise FileNotFoundError("model_index.json not found") except (AttributeError, OSError, ValueError, FileNotFoundError): - cfg = get_hf_file_to_dict("config.json", self.model) - if cfg is None: - raise ValueError(f"Could not find config.json or model_index.json for model {self.model}") - - self.tf_model_config = TransformerConfig.from_dict(cfg) - model_type = cfg.get("model_type") - architectures = cfg.get("architectures") or [] - - if model_type == "bagel" or "BagelForConditionalGeneration" in architectures: - self.model_class_name = "BagelPipeline" + # Skip transformer config loading for diffusers adapter + # (non-DiT models don't have a separate transformer folder/config) + if self.diffusion_load_format == "diffusers": self.tf_model_config = TransformerConfig() - self.update_multimodal_support() - elif model_type == "nextstep": - if self.model_class_name is None: - self.model_class_name = "NextStep11Pipeline" - self.tf_model_config = TransformerConfig() - self.update_multimodal_support() - elif architectures and len(architectures) == 1: - self.model_class_name = architectures[0] else: - raise + cfg = get_hf_file_to_dict("config.json", self.model) + if cfg is None: + raise ValueError(f"Could not find config.json or model_index.json for model {self.model}") + + self.tf_model_config = TransformerConfig.from_dict(cfg) + model_type = cfg.get("model_type") + architectures = cfg.get("architectures") or [] + + if model_type == "bagel" or "BagelForConditionalGeneration" in architectures: + self.model_class_name = "BagelPipeline" + self.tf_model_config = TransformerConfig() + self.update_multimodal_support() + elif model_type == "nextstep": + if self.model_class_name is None: + self.model_class_name = "NextStep11Pipeline" + self.tf_model_config = TransformerConfig() + self.update_multimodal_support() + elif architectures and len(architectures) == 1: + self.model_class_name = architectures[0] + else: + raise @classmethod def from_kwargs(cls, **kwargs: Any) -> "OmniDiffusionConfig": @@ -743,6 +770,12 @@ class OmniDiffusionConfig: cache_backend = os.environ.get("DIFFUSION_CACHE_BACKEND") or os.environ.get("DIFFUSION_CACHE_ADAPTER") kwargs["cache_backend"] = cache_backend.lower() if cache_backend else "none" + # Falsy-value check for not-None fields (convert potential None values in YAML config to empty containers) + if "diffusers_load_kwargs" in kwargs and kwargs["diffusers_load_kwargs"] is None: + kwargs["diffusers_load_kwargs"] = {} + if "diffusers_call_kwargs" in kwargs and kwargs["diffusers_call_kwargs"] is None: + kwargs["diffusers_call_kwargs"] = {} + # Filter kwargs to only include valid fields valid_fields = {f.name for f in fields(cls)} filtered_kwargs = {k: v for k, v in kwargs.items() if k in valid_fields} diff --git a/vllm_omni/diffusion/model_loader/diffusers_loader.py b/vllm_omni/diffusion/model_loader/diffusers_loader.py index 146afb26f..91f3574b1 100644 --- a/vllm_omni/diffusion/model_loader/diffusers_loader.py +++ b/vllm_omni/diffusion/model_loader/diffusers_loader.py @@ -32,6 +32,7 @@ from vllm.utils.torch_utils import set_default_torch_dtype from vllm_omni.diffusion.data import OmniDiffusionConfig from vllm_omni.diffusion.distributed.hsdp import HSDPInferenceConfig from vllm_omni.diffusion.model_loader.gguf_adapters import get_gguf_adapter +from vllm_omni.diffusion.models.diffusers_adapter.pipeline_diffusers_adapter import DiffusersAdapterPipeline from vllm_omni.diffusion.registry import initialize_model if TYPE_CHECKING: @@ -257,11 +258,14 @@ class DiffusersPipelineLoader: self, od_config: OmniDiffusionConfig, load_device: str, - load_format: str = "default", + load_format: str | None = "default", custom_pipeline_name: str | None = None, device: torch.device | None = None, ) -> nn.Module: """Load a model with the given configurations.""" + if load_format is None: + load_format = "default" + # CPU offload + FP8: load weights on device for FP8 quantization if load_device == "cpu" and od_config.quantization_config is not None: load_device = device.type @@ -277,11 +281,21 @@ class DiffusersPipelineLoader: with target_device: if load_format == "default": model = initialize_model(od_config) + elif load_format == "diffusers": + model = DiffusersAdapterPipeline(od_config=od_config, device=target_device) elif load_format == "custom_pipeline": model_cls = resolve_obj_by_qualname(custom_pipeline_name) model = model_cls(od_config=od_config) + else: + # 'dummy' format should not call this function at all + raise ValueError(f"Unknown load_format: {load_format}") logger.debug("Loading weights on %s ...", load_device) - if self._is_gguf_quantization(od_config): + if load_format == "diffusers": + # DiffusersAdapterPipeline.load_weights() calls + # DiffusionPipeline.from_pretrained() internally — it does + # NOT use our native (customized) pipeline classes. + cast(DiffusersAdapterPipeline, model).load_weights() + elif self._is_gguf_quantization(od_config): self._load_weights_with_gguf(model, od_config) else: # Quantization does not happen in `load_weights` but after it diff --git a/vllm_omni/diffusion/models/diffusers_adapter/__init__.py b/vllm_omni/diffusion/models/diffusers_adapter/__init__.py new file mode 100644 index 000000000..c8dd51c8e --- /dev/null +++ b/vllm_omni/diffusion/models/diffusers_adapter/__init__.py @@ -0,0 +1,12 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Diffusers backend adapter for vLLM-Omni.""" + +from vllm_omni.diffusion.models.diffusers_adapter.pipeline_diffusers_adapter import ( + DiffusersAdapterPipeline, +) + +__all__ = [ + "DiffusersAdapterPipeline", +] diff --git a/vllm_omni/diffusion/models/diffusers_adapter/pipeline_diffusers_adapter.py b/vllm_omni/diffusion/models/diffusers_adapter/pipeline_diffusers_adapter.py new file mode 100644 index 000000000..8a1fdfc08 --- /dev/null +++ b/vllm_omni/diffusion/models/diffusers_adapter/pipeline_diffusers_adapter.py @@ -0,0 +1,358 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Diffusers backend adapter for vLLM-Omni. + +Provides a black-box wrapper around any 🤗 Diffusers pipeline, enabling +vLLM-Omni to directly serve Diffusers models with near-zero per-model code. + +The adapter delegates full pipeline execution to diffusers' ``__call__()``. +It does NOT support: +- CFG parallel (diffusers handles CFG via guidance_scale internally) +- Sequence parallel (requires model-specific attention surgery) +- TeaCache / Cache-DiT (requires hooking into transformer blocks) +- Step-wise execution (continuous batching) +""" + +import logging +import os +from typing import Any + +import torch +from diffusers.pipelines.pipeline_utils import DiffusionPipeline +from torch import nn + +from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig +from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin +from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.inputs.data import OmniPromptType +from vllm_omni.platforms import current_omni_platform + +logger = logging.getLogger(__name__) + + +class DiffusersAdapterPipeline(nn.Module, DiffusionPipelineProfilerMixin): + """Black-box adapter that delegates full pipeline execution to a diffusers pipeline. + + Usage:: + + adapter = DiffusersAdapterPipeline(od_config=od_config) + adapter.load_weights() # calls DiffusionPipeline.from_pretrained() + output = adapter.forward(req) + + Step-wise execution is explicitly rejected — diffusers encapsulates the + full denoising loop internally. Use native pipelines for continuous + batching mode. + """ + + supports_step_execution: bool = False + + def __init__(self, *, od_config: OmniDiffusionConfig, device: torch.device | None = None): + super().__init__() + self._pipeline: DiffusionPipeline + self.od_config = od_config + self.device = device + self._capabilities: dict[str, Any] = {} + self._raise_unsupported_features() + + self.setup_diffusion_pipeline_profiler( + enable_diffusion_pipeline_profiler=od_config.enable_diffusion_pipeline_profiler, + profiler_targets=["forward"], + ) + if od_config.enable_diffusion_pipeline_profiler: + logger.info("Profiling enabled for DiffusersAdapterPipeline. Only 'forward' is supported.") + + # ------------------------------------------------------------------ + # Weight loading + # ------------------------------------------------------------------ + + def load_weights(self) -> None: + """Load the diffusers pipeline via ``DiffusionPipeline.from_pretrained()``.""" + + model_id = self.od_config.model + dtype = self.od_config.dtype + + load_kwargs = { + "torch_dtype": dtype, + **self.od_config.diffusers_load_kwargs, + } + logger.debug(f"Loading diffusers pipeline with kwargs: {load_kwargs}") + + self._pipeline = DiffusionPipeline.from_pretrained( + model_id, + **load_kwargs, + ).to(self.device) + + # CPU offloading + if self.od_config.enable_layerwise_offload: + self._pipeline.enable_sequential_cpu_offload() + elif self.od_config.enable_cpu_offload: + self._pipeline.enable_model_cpu_offload() + + # VAE slicing and tiling: try-catch because not all models have VAE + if self.od_config.vae_use_slicing: + try: + self._pipeline.enable_vae_slicing() + except Exception as e: + logger.warning( + f"Failed to enable VAE slicing for diffusers pipeline {self._pipeline.__class__.__name__}: {e}" + ) + if self.od_config.vae_use_tiling: + try: + self._pipeline.enable_vae_tiling() + except Exception as e: + logger.warning( + f"Failed to enable VAE tiling for diffusers pipeline {self._pipeline.__class__.__name__}: {e}" + ) + + # Attention backend + self._set_attention_backend() + + # ------------------------------------------------------------------ + # Step-wise execution — explicitly rejected + # ------------------------------------------------------------------ + + def prepare_encode(self, **_: Any) -> Any: + raise NotImplementedError( + "Step-wise execution is not yet supported with the diffusers backend. " + "Use a native pipeline for continuous batching mode." + ) + + def denoise_step(self, **_: Any) -> torch.Tensor | None: + raise NotImplementedError( + "Step-wise execution is not yet supported with the diffusers backend. " + "Use a native pipeline for continuous batching mode." + ) + + def step_scheduler(self, **_: Any) -> None: + raise NotImplementedError( + "Step-wise execution is not yet supported with the diffusers backend. " + "Use a native pipeline for continuous batching mode." + ) + + def post_decode(self, **_: Any) -> Any: + raise NotImplementedError( + "Step-wise execution is not yet supported with the diffusers backend. " + "Use a native pipeline for continuous batching mode." + ) + + # ------------------------------------------------------------------ + # Forward pass + # ------------------------------------------------------------------ + + def forward(self, req: OmniDiffusionRequest) -> DiffusionOutput: + """Full delegation to diffusers ``pipeline.__call__()``.""" + + kwargs = self._build_call_kwargs(req) + logger.debug(f"Calling diffusers pipeline with kwargs: {kwargs}") + + with torch.inference_mode(): + output = self._pipeline(**kwargs) # pyright: ignore[reportCallIssue] + + return self._wrap_output(output) + + # ------------------------------------------------------------------ + # Validation guards + # ------------------------------------------------------------------ + + def _raise_unsupported_features(self) -> None: + """Raise an error for incompatible feature switches.""" + pc = self.od_config.parallel_config + if pc.cfg_parallel_size > 1: + raise NotImplementedError( + "CFG parallel is not supported with the diffusers backend. " + "Diffusers handles CFG internally via guidance_scale." + ) + if pc.sequence_parallel_size is not None and pc.sequence_parallel_size > 1: + raise NotImplementedError( + "Sequence parallel is not supported with the diffusers backend. " + "It requires model-specific attention surgery." + ) + if self.od_config.cache_backend not in ("none", None): + raise NotImplementedError( + f"Cache backend '{self.od_config.cache_backend}' is not supported " + "with the diffusers backend. TeaCache/Cache-DiT require hooking " + "into individual transformer blocks." + ) + if self.od_config.enforce_eager: + raise NotImplementedError( + "Eager execution is not supported with the diffusers backend. " + "Use a native pipeline for continuous batching mode." + ) + if self.od_config.quantization_config is not None: + raise NotImplementedError( + "Quantization is not supported with the diffusers backend. Use a native pipeline for quantization." + ) + + # ------------------------------------------------------------------ + # Wrap settings, inputs, and outputs + # ------------------------------------------------------------------ + + def _set_attention_backend(self) -> None: + """Set the attention backend. + + Roughly follow the logic in vllm_omni/diffusion/attention/backends/utils/fa.py, + But also consider the available attention backends in diffusers. + (See: https://huggingface.co/docs/diffusers/optimization/attention_backends) + """ + if not hasattr(self._pipeline, "transformer"): + logging.info("No transformer found in diffusers pipeline. Skipping attention backend setting.") + return + + attention_backend_config = self.od_config.attention_backend or os.environ.get("DIFFUSION_ATTENTION_BACKEND") + attention_backend_attempts: list[str] = [] + match attention_backend_config: + case "FLASH_ATTN" | None: + if current_omni_platform.is_rocm(): + attention_backend_attempts.append("aiter") + elif current_omni_platform.is_xpu(): + attention_backend_attempts.append("_native_xla") + elif current_omni_platform.is_musa(): + logger.warning( + "Unknown diffusers attention backend option for MUSA platform. Falling back to SDPA." + ) + attention_backend_attempts.append("native") + else: + attention_backend_attempts.extend( + [ + "_flash_3_hub", + "_flash_3_varlen_hub", + "_flash_3", + "_flash_varlen_3", + "flash_hub", + "flash_varlen_hub", + "flash", + "flash_varlen", + "_native_flash", + ] + ) + case "SAGE_ATTN": + attention_backend_attempts.extend(["sage_hub", "sage", "sage", "sage_varlen"]) + case "ASCEND": + attention_backend_attempts.append("_native_npu") + case "TORCH_SDPA": + attention_backend_attempts.append("native") + case _: + logger.warning(f"Invalid attention backend: {attention_backend_config}. Falling back to SDPA.") + attention_backend_attempts.append("native") + + attempt_errors: list[str] = [] + set_backend: str | None = None + for backend in attention_backend_attempts: + try: + self._pipeline.transformer.set_attention_backend(backend) + set_backend = backend + break + except Exception as e: + attempt_errors.append(str(e)) + + # If all attempts fail, fallback to SDPA and warn the user about the failures + if len(attempt_errors) == len(attention_backend_attempts): + self._pipeline.transformer.set_attention_backend("native") + logger.warning( + f"Failed to set attention backend '{attention_backend_config}' for " + f"diffusers pipeline {self._pipeline.__class__.__name__}. " + "Falling back to SDPA. " + f"The following attempts were made: {dict(zip(attention_backend_attempts, attempt_errors))}" + ) + return + + # If some attempts fail, only warn the user about the failures + logger.info( + f"Set diffusers attention backend to '{set_backend}', adapted from " + f"user config value '{attention_backend_config}'." + ) + if len(attempt_errors) > 0: + logger.warning( + f"The following failed attempts were made before choosing this diffusers backend: " + f"{dict(zip(attention_backend_attempts, attempt_errors))}" + ) + + def _build_call_kwargs(self, req: OmniDiffusionRequest) -> dict[str, Any]: + """Translate ``OmniDiffusionRequest`` into diffusers ``__call__`` kwargs.""" + sampling = req.sampling_params + prompt, neg_prompt = self._extract_prompt(req.prompts) + + # Merge user-provided call kwargs from stage/CLI defaults. + # Request-time parameters take precedence over stage-config defaults + call_kwargs = self.od_config.diffusers_call_kwargs + kwargs: dict[str, Any] = { + **call_kwargs, + "prompt": prompt, + "num_inference_steps": sampling.num_inference_steps, + "guidance_scale": sampling.guidance_scale, + "output_type": sampling.output_type or self.od_config.output_type, + } + + if sampling.height is not None: + kwargs["height"] = sampling.height + if sampling.width is not None: + kwargs["width"] = sampling.width + if sampling.num_frames is not None and sampling.num_frames > 1: + kwargs["num_frames"] = sampling.num_frames + if sampling.num_outputs_per_prompt is not None and sampling.num_outputs_per_prompt > 1: + kwargs["num_images_per_prompt"] = sampling.num_outputs_per_prompt + + if neg_prompt is not None: + kwargs["negative_prompt"] = neg_prompt + + if sampling.generator is not None: + kwargs["generator"] = sampling.generator + elif sampling.seed is not None: + kwargs["generator"] = torch.Generator(device=sampling.generator_device).manual_seed(sampling.seed) + else: + kwargs["generator"] = torch.Generator(device=sampling.generator_device) + + if sampling.latents is not None: + kwargs["latents"] = sampling.latents + + return kwargs + + @staticmethod + def _extract_prompt(prompt_obj: list[OmniPromptType]) -> tuple[str | list[str], str | list[str] | None]: + """Extract the text prompts and negative prompts from a list of prompt objects.""" + if len(prompt_obj) == 1: + if isinstance(prompt_obj[0], str): + return prompt_obj[0], None + else: + return prompt_obj[0].get("prompt", ""), prompt_obj[0].get("negative_prompt", None) + + prompts = [] + negative_prompts: list[str] | None = [] + for prompt in prompt_obj: + if isinstance(prompt, str): + prompts.append(prompt) + else: + prompts.append(prompt.get("prompt", "")) + negative_prompts.append(prompt.get("negative_prompt", "")) + if all(not np for np in negative_prompts): + negative_prompts = None + return prompts, negative_prompts + + @staticmethod + def _extract_negative_prompt(prompt_obj: Any) -> str | None: + """Extract the negative prompt from a prompt object, if present.""" + if isinstance(prompt_obj, dict): + return prompt_obj.get("negative_prompt") + return getattr(prompt_obj, "negative_prompt", None) + + def _wrap_output(self, output: Any) -> DiffusionOutput: + """Convert diffusers pipeline output to ``DiffusionOutput``. + + Diffusers output types: + - ``ImagePipelineOutput(images=...)`` — text2img, img2img + - ``VideoPipelineOutput(frames=...)`` — text2vid, img2vid + """ + from vllm_omni.diffusion.data import DiffusionOutput + + if hasattr(output, "images"): + # Preserve diffusers image format (`output_type`) + return DiffusionOutput(output=output.images) + + if hasattr(output, "frames"): + # Preserve diffusers video format (`output_type`) + return DiffusionOutput(output=output.frames) + + if hasattr(output, "audios"): + return DiffusionOutput(output=output.audios) + + return DiffusionOutput(output=output) diff --git a/vllm_omni/diffusion/registry.py b/vllm_omni/diffusion/registry.py index 4001109cc..2da2a68f4 100644 --- a/vllm_omni/diffusion/registry.py +++ b/vllm_omni/diffusion/registry.py @@ -208,6 +208,11 @@ _DIFFUSION_MODELS = { "pipeline_omnivoice", "OmniVoicePipeline", ), + "DiffusersAdapterPipeline": ( + "diffusers_adapter", + "pipeline_diffusers_adapter", + "DiffusersAdapterPipeline", + ), } diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index a26102d05..197c99fb2 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -1319,6 +1319,12 @@ class AsyncOmniEngine: if "dtype" in normalized_kwargs: stage_engine_args["dtype"] = normalized_kwargs["dtype"] + # New split fields for diffusers adapter kwargs. + if kwargs.get("diffusers_load_kwargs") is not None: + stage_engine_args["diffusers_load_kwargs"] = kwargs["diffusers_load_kwargs"] + if kwargs.get("diffusers_call_kwargs") is not None: + stage_engine_args["diffusers_call_kwargs"] = kwargs["diffusers_call_kwargs"] + default_stage_cfg = [ { "stage_id": 0, diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index e355ee679..841858778 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -267,6 +267,40 @@ class OmniServeCommand(CLISubcommand): default=None, help="Override the diffusion pipeline class name (e.g. LTX2ImageToVideoPipeline).", ) + omni_config_group.add_argument( + "--diffusion-load-format", + dest="diffusion_load_format", + type=str, + default=None, + choices=["default", "custom_pipeline", "dummy", "diffusers"], + help=( + "How to load the diffusion pipeline: native/registry (default), " + "custom_pipeline, dummy, or diffusers for the HF diffusers adapter." + ), + ) + omni_config_group.add_argument( + "--diffusers-load-kwargs", + dest="diffusers_load_kwargs", + type=json.loads, + default="{}", + help=( + "JSON object passed to DiffusionPipeline.from_pretrained()." + "It overrides corresponding parameters in the standard vLLM-Omni interface." + '(e.g. \'{"use_safetensors": true, "variant": "fp16"}\').' + ), + ) + omni_config_group.add_argument( + "--diffusers-call-kwargs", + dest="diffusers_call_kwargs", + type=json.loads, + default="{}", + help=( + "JSON object passed to pipeline.__call__(). " + "Useful for model-specific sampling parameters not covered by the vLLM-Omni interface." + "During request time, it is overridden by corresponding parameters in the vLLM-Omni interface." + '(e.g. \'{"num_inference_steps": 30, "guidance_scale": 7.5}\').' + ), + ) omni_config_group.add_argument( "--usp", "--ulysses-degree", From ead87aa52140e7a83ff3b2bf6bea770a6686dc44 Mon Sep 17 00:00:00 2001 From: NumberWan Date: Thu, 23 Apr 2026 00:57:01 +0800 Subject: [PATCH 14/45] [BugFix] Surface diffusion metrics in chat completions; sanitize Bagel img2img mm kwargs (#2932) Signed-off-by: NumberWan --- benchmarks/diffusion/backends.py | 12 ++ .../test_diffusion_backends_metrics.py | 107 ++++++++++++++++++ tests/e2e/online_serving/test_bagel_online.py | 4 +- .../openai_api/test_serving_chat_metrics.py | 35 ++++++ vllm_omni/entrypoints/openai/serving_chat.py | 8 +- .../model_executor/models/bagel/bagel.py | 11 +- 6 files changed, 173 insertions(+), 4 deletions(-) create mode 100644 tests/benchmarks/test_diffusion_backends_metrics.py diff --git a/benchmarks/diffusion/backends.py b/benchmarks/diffusion/backends.py index 13ce7c830..d33160f13 100644 --- a/benchmarks/diffusion/backends.py +++ b/benchmarks/diffusion/backends.py @@ -122,6 +122,18 @@ async def async_request_chat_completions( output.peak_memory_mb = first_item.get("peak_memory_mb", 0.0) except (IndexError, TypeError, AttributeError): pass + + if (not output.stage_durations or output.peak_memory_mb == 0.0) and isinstance( + resp_json.get("metrics"), dict + ): + m = resp_json["metrics"] + if not output.stage_durations and isinstance(m.get("stage_durations"), dict): + output.stage_durations = m.get("stage_durations") or {} + if output.peak_memory_mb == 0.0 and m.get("peak_memory_mb") is not None: + try: + output.peak_memory_mb = float(m.get("peak_memory_mb") or 0.0) + except (TypeError, ValueError): + pass else: output.error = f"HTTP {response.status}: {await response.text()}" output.success = False diff --git a/tests/benchmarks/test_diffusion_backends_metrics.py b/tests/benchmarks/test_diffusion_backends_metrics.py new file mode 100644 index 000000000..2d51d0f1d --- /dev/null +++ b/tests/benchmarks/test_diffusion_backends_metrics.py @@ -0,0 +1,107 @@ +import pytest + +from benchmarks.diffusion.backends import RequestFuncInput, async_request_chat_completions + + +class _MockResponse: + def __init__(self, payload: dict, status: int = 200): + self._payload = payload + self.status = status + + async def __aenter__(self): + return self + + async def __aexit__(self, exc_type, exc, tb): + return False + + async def json(self): + return self._payload + + async def text(self): + return str(self._payload) + + +class _MockSession: + def __init__(self, payload: dict): + self._payload = payload + + def post(self, *args, **kwargs): + return _MockResponse(self._payload) + + +@pytest.mark.core_model +@pytest.mark.benchmark +@pytest.mark.cpu +@pytest.mark.asyncio +async def test_chat_completions_metrics_fallback_to_top_level(): + payload = { + "choices": [ + { + "message": { + "content": [ + { + "type": "image_url", + "image_url": {"url": "data:image/png;base64,abc"}, + } + ] + } + } + ], + "metrics": { + "stage_durations": {"diffusion": 1.25}, + "peak_memory_mb": 4096.0, + }, + } + + output = await async_request_chat_completions( + RequestFuncInput( + prompt="draw a cat", + api_url="http://test.local/v1/chat/completions", + model="ByteDance-Seed/BAGEL-7B-MoT", + ), + session=_MockSession(payload), + ) + + assert output.success is True + assert output.stage_durations == {"diffusion": 1.25} + assert output.peak_memory_mb == 4096.0 + + +@pytest.mark.core_model +@pytest.mark.benchmark +@pytest.mark.cpu +@pytest.mark.asyncio +async def test_chat_completions_metrics_message_level_takes_precedence(): + payload = { + "choices": [ + { + "message": { + "content": [ + { + "type": "image_url", + "image_url": {"url": "data:image/png;base64,abc"}, + "stage_durations": {"message_stage": 0.7}, + "peak_memory_mb": 1234.0, + } + ] + } + } + ], + "metrics": { + "stage_durations": {"top_level_stage": 9.9}, + "peak_memory_mb": 9999.0, + }, + } + + output = await async_request_chat_completions( + RequestFuncInput( + prompt="draw a dog", + api_url="http://test.local/v1/chat/completions", + model="ByteDance-Seed/BAGEL-7B-MoT", + ), + session=_MockSession(payload), + ) + + assert output.success is True + assert output.stage_durations == {"message_stage": 0.7} + assert output.peak_memory_mb == 1234.0 diff --git a/tests/e2e/online_serving/test_bagel_online.py b/tests/e2e/online_serving/test_bagel_online.py index 745e9a116..1878fcd6e 100644 --- a/tests/e2e/online_serving/test_bagel_online.py +++ b/tests/e2e/online_serving/test_bagel_online.py @@ -107,7 +107,7 @@ def test_bagel_text2img_online(omni_server, openai_client) -> None: @hardware_test(res={"cuda": "H100", "rocm": "MI325"}) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_bagel_img2img_online(omni_server, openai_client) -> None: - """Test Bagel img2img via OpenAI-compatible chat completions API.""" + """Test Bagel img2img with explicit height/width in chat completions API.""" input_image = ImageAsset("2560px-Gfp-wisconsin-madison-the-nature-boardwalk").pil_image.convert("RGB") buffer = BytesIO() input_image.save(buffer, format="JPEG") @@ -118,6 +118,8 @@ def test_bagel_img2img_online(omni_server, openai_client) -> None: "messages": _build_img2img_messages(IMG2IMG_PROMPT, image_b64), "modalities": ["image"], "extra_body": { + "height": 512, + "width": 512, "num_inference_steps": 2, "guidance_scale": 0.0, "seed": 42, diff --git a/tests/entrypoints/openai_api/test_serving_chat_metrics.py b/tests/entrypoints/openai_api/test_serving_chat_metrics.py index d25af6c38..0647c40a3 100644 --- a/tests/entrypoints/openai_api/test_serving_chat_metrics.py +++ b/tests/entrypoints/openai_api/test_serving_chat_metrics.py @@ -1,6 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 """Unit tests for OmniChatCompletionResponse/StreamResponse metrics field.""" +from types import SimpleNamespace + import pytest pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -50,3 +52,36 @@ def test_omni_chat_completion_stream_response_metrics(): ) assert response.modality == "audio" assert response.metrics == {"stage_latency": 0.5} + + +def test_create_image_choice_exposes_diffusion_metrics(): + """Ensure image chat content exposes profiler metrics for clients.""" + from PIL import Image + + from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat + + stage_durations = {"prefill": 0.12, "diffusion": 1.23} + peak_memory_mb = 3210.5 + omni_outputs = SimpleNamespace( + request_output=None, + stage_durations=stage_durations, + peak_memory_mb=peak_memory_mb, + images=[Image.new("RGB", (2, 2), color=(255, 0, 0))], + ) + + choices = OmniOpenAIServingChat._create_image_choice( # type: ignore[misc] + None, + omni_outputs=omni_outputs, + role="assistant", + request=SimpleNamespace(return_token_ids=False), + ) + + assert len(choices) == 1 + content = choices[0].message.content + assert isinstance(content, list) + assert len(content) == 1 + first_item = content[0] + assert first_item["type"] == "image_url" + assert first_item["image_url"]["url"].startswith("data:image/png;base64,") + assert first_item["stage_durations"] == stage_durations + assert first_item["peak_memory_mb"] == peak_memory_mb diff --git a/vllm_omni/entrypoints/openai/serving_chat.py b/vllm_omni/entrypoints/openai/serving_chat.py index 06b739e3b..faf3ea320 100644 --- a/vllm_omni/entrypoints/openai/serving_chat.py +++ b/vllm_omni/entrypoints/openai/serving_chat.py @@ -1636,7 +1636,13 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): logger.warning(f"Unsupported final output type: {omni_outputs.final_output_type}") continue if omni_outputs.metrics: - response_metrics = omni_outputs.metrics + response_metrics = dict(omni_outputs.metrics) + if omni_outputs.final_output_type == "image": + # Expose diffusion profiler metrics on the top-level response for benchmarks / clients. + if response_metrics is None: + response_metrics = {} + response_metrics.setdefault("stage_durations", omni_outputs.stage_durations or {}) + response_metrics.setdefault("peak_memory_mb", float(omni_outputs.peak_memory_mb or 0.0)) choices.extend(choices_data) response = OmniChatCompletionResponse( diff --git a/vllm_omni/model_executor/models/bagel/bagel.py b/vllm_omni/model_executor/models/bagel/bagel.py index cbb775680..9f07ef6a0 100644 --- a/vllm_omni/model_executor/models/bagel/bagel.py +++ b/vllm_omni/model_executor/models/bagel/bagel.py @@ -203,6 +203,13 @@ class OmniBagelDataParser(MultiModalDataParser): class OmniBagelMultiModalProcessor(BaseMultiModalProcessor[OmniBagelProcessingInfo]): IMG2IMG_PLACEHOLDER = "<|fim_middle|>" + @staticmethod + def _mm_kwargs_for_bagel_img2img_hf(mm_kwargs: Mapping[str, object]) -> dict[str, object]: + # OpenAI / GLM-style serving may pass target_h/target_w for output grid sizing. + # BagelProcessor does not accept these in img2img mode; strip here so callers + # (e.g. serving_chat) can stay model-agnostic. + return {k: v for k, v in mm_kwargs.items() if k not in ("target_h", "target_w")} + def _cached_apply_hf_processor(self, inputs, timing_ctx): # img2img: prompt text must be modified based on mm data presence, # so text and mm data cannot be tokenized separately — bypass cache. @@ -248,7 +255,7 @@ class OmniBagelMultiModalProcessor(BaseMultiModalProcessor[OmniBagelProcessingIn if "images" in img2img_data: del img2img_data["images"] img2img_data["images"] = img2img_data.pop("pixel_values_img2img") - kwargs_img2img = dict(mm_kwargs) + kwargs_img2img = self._mm_kwargs_for_bagel_img2img_hf(mm_kwargs) kwargs_img2img["is_img2img"] = True out_img2img = super()._call_hf_processor(prompt, img2img_data, kwargs_img2img, tok_kwargs) if "pixel_values" in out_img2img: @@ -262,7 +269,7 @@ class OmniBagelMultiModalProcessor(BaseMultiModalProcessor[OmniBagelProcessingIn elif has_img2img: mm_data = dict(mm_data) mm_data["images"] = mm_data.pop("pixel_values_img2img") - mm_kwargs = dict(mm_kwargs) + mm_kwargs = self._mm_kwargs_for_bagel_img2img_hf(mm_kwargs) mm_kwargs["is_img2img"] = True outputs = super()._call_hf_processor(prompt, mm_data, mm_kwargs, tok_kwargs) if "pixel_values" in outputs: From e3b0afbb97a026e8fa348c11a06408bab33d55d1 Mon Sep 17 00:00:00 2001 From: Nick Cao Date: Wed, 22 Apr 2026 16:23:00 -0400 Subject: [PATCH 15/45] [Refactor] Let diffusion pipelines declare offloadable modules via SupportsModuleOffload (#2427) Signed-off-by: Nick Cao Co-authored-by: Claude --- .claude/skills/add-diffusion-model/SKILL.md | 40 ++- .../diffusion/cpu_offload_diffusion.md | 55 +++- .../offloader/test_module_collector.py | 240 ++++++++++++++++++ .../diffusion/models/bagel/pipeline_bagel.py | 17 +- vllm_omni/diffusion/models/interface.py | 21 ++ .../diffusion/models/ltx2/pipeline_ltx2.py | 9 +- .../models/ltx2/pipeline_ltx2_image2video.py | 9 +- .../models/omnigen2/pipeline_omnigen2.py | 9 +- .../diffusion/offloader/layerwise_backend.py | 13 +- .../diffusion/offloader/module_collector.py | 134 +++++++--- .../diffusion/offloader/sequential_backend.py | 6 +- 11 files changed, 491 insertions(+), 62 deletions(-) create mode 100644 tests/diffusion/offloader/test_module_collector.py diff --git a/.claude/skills/add-diffusion-model/SKILL.md b/.claude/skills/add-diffusion-model/SKILL.md index a7e0bbf9a..0b979e1a9 100644 --- a/.claude/skills/add-diffusion-model/SKILL.md +++ b/.claude/skills/add-diffusion-model/SKILL.md @@ -282,10 +282,11 @@ For Omni or custom models, create: Required updates: 1. `docs/user_guide/diffusion/parallelism_acceleration.md` — parallelism support table -2. `docs/user_guide/diffusion/teacache.md` — if TeaCache supported -3. `docs/user_guide/diffusion/cache_dit_acceleration.md` — if Cache-DiT supported -4. `examples/offline_inference/xxx/README.md` — offline example docs -5. `examples/online_serve/xxx/README.md` — online serve docs +2. `docs/user_guide/diffusion/cpu_offload_diffusion.md` — if CPU offload supported (add to supported models table) +3. `docs/user_guide/diffusion/teacache.md` — if TeaCache supported +4. `docs/user_guide/diffusion/cache_dit_acceleration.md` — if Cache-DiT supported +5. `examples/offline_inference/xxx/README.md` — offline example docs +6. `examples/online_serve/xxx/README.md` — online serve docs ### Step 8: Add E2E Tests (Recommended) @@ -512,6 +513,37 @@ After adding parallelism support, update: 1. `docs/user_guide/diffusion/parallelism_acceleration.md` — add your model to the support table 2. Record which parallelism methods are supported (USP, Ring, CFG, TP, HSDP, VAE-Patch) +### Step 11: Add CPU Offload Support + +Implement `SupportsModuleOffload` on your pipeline class to enable +`--enable-cpu-offload` and `--enable-layerwise-offload`. The protocol +declares which submodules the offloader should manage: + +```python +from typing import ClassVar +from vllm_omni.diffusion.models.interface import SupportsModuleOffload + +class YourPipeline(nn.Module, SupportsModuleOffload): + _dit_modules: ClassVar[list[str]] = ["transformer"] + _encoder_modules: ClassVar[list[str]] = ["text_encoder"] + _vae_modules: ClassVar[list[str]] = ["vae"] + _resident_modules: ClassVar[list[str]] = [] # optional +``` + +- `_dit_modules`: denoising submodules (kept on GPU during diffusion loop) +- `_encoder_modules`: encoder/vision submodules (offloaded to CPU during diffusion loop) +- `_vae_modules`: VAE(s) (handled by both sequential and layerwise backends) +- `_resident_modules`: additional modules to pin on GPU during layerwise + offloading (e.g. embedders, connectors). Only used by the layerwise + backend. Optional — defaults to `[]`. + +All attribute names support dotted paths for nested submodules +(e.g. `"pipe.transformer"`, `"bagel.time_embedder"`). + +Pipelines without `SupportsModuleOffload` fall back to scanning +well-known attribute names (`transformer`, `text_encoder`, `vae`, +etc.), which fails for non-standard names. + --- ## Iterative Development Tips diff --git a/docs/user_guide/diffusion/cpu_offload_diffusion.md b/docs/user_guide/diffusion/cpu_offload_diffusion.md index f80005ccb..b85c7e408 100644 --- a/docs/user_guide/diffusion/cpu_offload_diffusion.md +++ b/docs/user_guide/diffusion/cpu_offload_diffusion.md @@ -36,6 +36,45 @@ m = Omni(model="Wan-AI/Wan2.2-T2V-A14B-Diffusers", enable_cpu_offload=True) vllm-omni serve diffusion Wan-AI/Wan2.2-T2V-A14B-Diffusers --enable-cpu-offload ``` +### To Support a Model + +Implement the `SupportsModuleOffload` protocol to declare which +submodules participate in offloading: + +```python +from typing import ClassVar +from vllm_omni.diffusion.models.interface import SupportsModuleOffload + +class MyPipeline(nn.Module, SupportsModuleOffload): + _dit_modules: ClassVar[list[str]] = ["transformer"] + _encoder_modules: ClassVar[list[str]] = ["text_encoder", "vision_model"] + _vae_modules: ClassVar[list[str]] = ["vae"] + _resident_modules: ClassVar[list[str]] = [] # optional + + def __init__(self): + super().__init__() + self.transformer = ... # DiT — stays on GPU during denoising + self.text_encoder = ... # Encoder — offloaded to CPU during denoising + self.vision_model = ... # Encoder — offloaded to CPU during denoising + self.vae = ... # VAE — always on GPU +``` + +- `_dit_modules`: attribute names of denoising submodules (kept on GPU + during the diffusion loop). +- `_encoder_modules`: attribute names of encoder/vision submodules + (offloaded to CPU during the diffusion loop). +- `_vae_modules`: attribute names of VAE(s) (always kept on GPU, not + part of the mutual exclusion hooks). +- `_resident_modules`: attribute names of small submodules that must + stay on GPU during layerwise offloading (e.g. embedders, connectors). + Optional — defaults to `[]`. + +All attribute names support dotted paths for nested submodules +(e.g. `"pipe.transformer"`, `"bagel.time_embedder"`). + +Both DiT and encoder lists are needed because the offload hooks use +mutual exclusion: when one group runs, the other moves to CPU. + ### Limitations - Cold start latency increases - Adds overhead from CPU-GPU transfers between encoder and denoising phases @@ -116,11 +155,19 @@ class Flux2Transformer2DModel(nn.Module): **Module Discovery** -The offloader automatically discovers pipeline components: +The offloader discovers pipeline components in two ways: -- **DiT modules**: `transformer`, `transformer_2`, `dit` -- **Encoders**: `text_encoder`, `text_encoder_2`, `text_encoder_3`, `image_encoder` -- **VAE**: `vae` +1. **Protocol-based** (preferred): If the pipeline implements + `SupportsModuleOffload`, its `_dit_modules`, `_encoder_modules`, + `_vae_modules`, and `_resident_modules` class variables are used + directly. All attribute names support dotted paths (e.g. + `"pipe.transformer"`, `"bagel.time_embedder"`) for nested submodules. + +2. **Fallback attribute scan**: Otherwise, the offloader scans for + well-known attribute names: + - **DiT modules**: `transformer`, `transformer_2`, `dit`, `sr_dit`, `language_model`, `transformer_blocks`, `model` + - **Encoders**: `text_encoder`, `text_encoder_2`, `text_encoder_3`, `image_encoder` + - **VAE**: `vae`, `audio_vae` **Hook System** diff --git a/tests/diffusion/offloader/test_module_collector.py b/tests/diffusion/offloader/test_module_collector.py new file mode 100644 index 000000000..ab15ad8df --- /dev/null +++ b/tests/diffusion/offloader/test_module_collector.py @@ -0,0 +1,240 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Unit tests for ModuleDiscovery and SupportsModuleOffload.""" + +from typing import ClassVar + +import pytest +from torch import nn + +from vllm_omni.diffusion.models.interface import SupportsModuleOffload +from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery + +pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model] + +# NOTE: tests for skipped/warned attributes verify the *behavioral* +# outcome (attribute excluded from results) but do not assert on log +# output. vllm's logger sets propagate=False, preventing caplog from +# capturing records. See https://github.com/pytest-dev/pytest/issues/3697 + + +# --------------------------------------------------------------------------- +# Test pipelines +# --------------------------------------------------------------------------- + + +class FallbackPipeline(nn.Module): + """Pipeline with standard attribute names (no protocol).""" + + def __init__(self): + super().__init__() + self.transformer = nn.Linear(10, 10) + self.text_encoder = nn.Linear(10, 10) + self.text_encoder_2 = nn.Linear(10, 10) + self.vae = nn.Linear(10, 10) + + +class NonModuleAttrPipeline(nn.Module): + """Pipeline where an attribute is not an nn.Module (fallback path).""" + + def __init__(self): + super().__init__() + self.transformer = nn.Linear(10, 10) + self.text_encoder = "not_a_module" + self.vae = nn.Linear(10, 10) + + +class DuplicateAttrPipeline(nn.Module): + """Pipeline where two encoder attrs point to the same module.""" + + def __init__(self): + super().__init__() + self.transformer = nn.Linear(10, 10) + encoder = nn.Linear(10, 10) + self.text_encoder = encoder + self.text_encoder_2 = encoder + self.vae = nn.Linear(10, 10) + + +class ProtocolPipeline(nn.Module, SupportsModuleOffload): + """Pipeline with non-standard names, using the protocol.""" + + _dit_modules: ClassVar[list[str]] = ["gen_transformer"] + _encoder_modules: ClassVar[list[str]] = ["mllm", "vision_model"] + _vae_modules: ClassVar[list[str]] = ["gen_vae"] + + def __init__(self): + super().__init__() + self.gen_transformer = nn.Linear(10, 10) + self.mllm = nn.Linear(10, 10) + self.vision_model = nn.Linear(10, 10) + self.gen_vae = nn.Linear(10, 10) + # Standard name present but NOT declared — should be ignored + self.transformer = nn.Linear(10, 10) + + +class MissingAttrPipeline(nn.Module, SupportsModuleOffload): + """Pipeline that declares a non-existent attribute.""" + + _dit_modules: ClassVar[list[str]] = ["transformer"] + _encoder_modules: ClassVar[list[str]] = ["nonexistent_encoder"] + _vae_modules: ClassVar[list[str]] = ["vae"] + + def __init__(self): + super().__init__() + self.transformer = nn.Linear(10, 10) + self.vae = nn.Linear(10, 10) + + +class MissingIntermediatePipeline(nn.Module, SupportsModuleOffload): + """Pipeline with dotted path referencing non-existent intermediate.""" + + _dit_modules: ClassVar[list[str]] = ["nonexistent.transformer"] + _encoder_modules: ClassVar[list[str]] = [] + _vae_modules: ClassVar[list[str]] = [] + + def __init__(self): + super().__init__() + + +class NestedPipeline(nn.Module, SupportsModuleOffload): + """Pipeline with nested modules accessed via dotted paths.""" + + _dit_modules: ClassVar[list[str]] = ["pipe.transformer"] + _encoder_modules: ClassVar[list[str]] = ["pipe.text_encoder"] + _vae_modules: ClassVar[list[str]] = ["vae"] + + def __init__(self): + super().__init__() + self.pipe = nn.Module() + self.pipe.transformer = nn.Linear(10, 10) + self.pipe.text_encoder = nn.Linear(10, 10) + self.vae = nn.Linear(10, 10) + + +class ResidentPipeline(nn.Module, SupportsModuleOffload): + """Pipeline with resident modules that must stay on GPU.""" + + _dit_modules: ClassVar[list[str]] = ["language_model.model"] + _encoder_modules: ClassVar[list[str]] = [] + _vae_modules: ClassVar[list[str]] = ["vae"] + _resident_modules: ClassVar[list[str]] = [ + "bagel.time_embedder", + "bagel.vae2llm", + ] + + def __init__(self): + super().__init__() + self.language_model = nn.Module() + self.language_model.model = nn.Linear(10, 10) + self.bagel = nn.Module() + self.bagel.time_embedder = nn.Linear(10, 10) + self.bagel.vae2llm = nn.Linear(10, 10) + self.vae = nn.Linear(10, 10) + + +class MultiVaePipeline(nn.Module, SupportsModuleOffload): + """Pipeline with multiple VAEs.""" + + _dit_modules: ClassVar[list[str]] = ["transformer"] + _encoder_modules: ClassVar[list[str]] = ["text_encoder"] + _vae_modules: ClassVar[list[str]] = ["vae", "audio_vae"] + + def __init__(self): + super().__init__() + self.transformer = nn.Linear(10, 10) + self.text_encoder = nn.Linear(10, 10) + self.vae = nn.Linear(10, 10) + self.audio_vae = nn.Linear(10, 10) + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- + + +class TestFallbackDiscovery: + """Test the fallback attribute scan (no SupportsModuleOffload).""" + + def test_discovers_standard_attrs(self): + pipeline = FallbackPipeline() + result = ModuleDiscovery.discover(pipeline) + + assert not isinstance(pipeline, SupportsModuleOffload) + assert result.dit_names == ["transformer"] + assert result.dits[0] is pipeline.transformer + assert result.encoder_names == ["text_encoder", "text_encoder_2"] + assert result.vaes[0] is pipeline.vae + assert result.resident_modules == [] + + def test_deduplicates_encoders(self): + pipeline = DuplicateAttrPipeline() + result = ModuleDiscovery.discover(pipeline) + + assert len(result.encoders) == 1 + assert result.encoder_names == ["text_encoder"] + + def test_skips_non_module_attr(self): + pipeline = NonModuleAttrPipeline() + result = ModuleDiscovery.discover(pipeline) + + assert len(result.encoders) == 0 + + +class TestProtocolDiscovery: + """Test discovery via SupportsModuleOffload protocol.""" + + def test_discovers_declared_attrs_and_ignores_undeclared(self): + pipeline = ProtocolPipeline() + result = ModuleDiscovery.discover(pipeline) + + assert isinstance(pipeline, SupportsModuleOffload) + assert result.dit_names == ["gen_transformer"] + assert result.encoder_names == ["mllm", "vision_model"] + assert len(result.vaes) == 1 + # self.transformer exists but is NOT in _dit_modules + assert "transformer" not in result.dit_names + # No _resident_modules declared — defaults to empty + assert result.resident_modules == [] + + def test_skips_missing_attr(self): + pipeline = MissingAttrPipeline() + result = ModuleDiscovery.discover(pipeline) + + assert len(result.encoders) == 0 + + def test_skips_missing_intermediate(self): + result = ModuleDiscovery.discover(MissingIntermediatePipeline()) + + assert len(result.dits) == 0 + + def test_dotted_path_resolves_nested_modules(self): + pipeline = NestedPipeline() + result = ModuleDiscovery.discover(pipeline) + + assert result.dit_names == ["pipe.transformer"] + assert result.dits[0] is pipeline.pipe.transformer + assert result.encoder_names == ["pipe.text_encoder"] + assert result.encoders[0] is pipeline.pipe.text_encoder + assert result.vaes[0] is pipeline.vae + + def test_resident_modules(self): + pipeline = ResidentPipeline() + result = ModuleDiscovery.discover(pipeline) + + assert result.resident_names == [ + "bagel.time_embedder", + "bagel.vae2llm", + ] + assert result.resident_modules[0] is pipeline.bagel.time_embedder + assert result.resident_modules[1] is pipeline.bagel.vae2llm + assert result.dits[0] is pipeline.language_model.model + + def test_multiple_vaes(self): + pipeline = MultiVaePipeline() + result = ModuleDiscovery.discover(pipeline) + + assert len(result.vaes) == 2 + assert result.vaes[0] is pipeline.vae + assert result.vaes[1] is pipeline.audio_vae diff --git a/vllm_omni/diffusion/models/bagel/pipeline_bagel.py b/vllm_omni/diffusion/models/bagel/pipeline_bagel.py index 90baf5f67..0290e9502 100644 --- a/vllm_omni/diffusion/models/bagel/pipeline_bagel.py +++ b/vllm_omni/diffusion/models/bagel/pipeline_bagel.py @@ -12,6 +12,7 @@ from collections.abc import Iterable from copy import deepcopy from dataclasses import dataclass from math import isqrt +from typing import ClassVar import numpy as np import torch @@ -26,6 +27,7 @@ from vllm.transformers_utils.configs.bagel import BagelConfig from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig from vllm_omni.diffusion.distributed.utils import get_local_device from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader +from vllm_omni.diffusion.models.interface import SupportsModuleOffload from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific @@ -148,12 +150,25 @@ class SiglipNaViTWrapper(nn.Module): return outputs.last_hidden_state.squeeze(0) -class BagelPipeline(nn.Module, DiffusionPipelineProfilerMixin): +class BagelPipeline(nn.Module, SupportsModuleOffload, DiffusionPipelineProfilerMixin): """Bagel generation pipeline (MoT) packaged for vllm-omni diffusion engine. This pipeline is self-contained and uses the ported Bagel core files. """ + _dit_modules: ClassVar[list[str]] = ["language_model.model"] + _encoder_modules: ClassVar[list[str]] = [] + _vae_modules: ClassVar[list[str]] = ["vae"] + _resident_modules: ClassVar[list[str]] = [ + "bagel.time_embedder", + "bagel.vae2llm", + "bagel.llm2vae", + "bagel.latent_pos_embed", + "bagel.vit_model", + "bagel.connector", + "bagel.vit_pos_embed", + ] + def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""): super().__init__() self.od_config = od_config diff --git a/vllm_omni/diffusion/models/interface.py b/vllm_omni/diffusion/models/interface.py index ef906472b..00d54420d 100644 --- a/vllm_omni/diffusion/models/interface.py +++ b/vllm_omni/diffusion/models/interface.py @@ -58,6 +58,27 @@ class SupportsStepExecution(Protocol): """Decode output after denoise loop.""" +@runtime_checkable +class SupportsModuleOffload(Protocol): + """Declares which submodules participate in CPU offload. + + All attribute names support dotted paths for nested submodules + (e.g. ``"pipe.transformer"``). + + Attributes: + _dit_modules: Denoising submodules (on GPU during diffusion). + _encoder_modules: Encoder submodules (offloaded during diffusion). + _vae_modules: VAE(s) (always on GPU). + _resident_modules: Extra modules pinned on GPU during layerwise + offloading. Optional, defaults to ``[]``. + """ + + _dit_modules: ClassVar[list[str]] + _encoder_modules: ClassVar[list[str]] + _vae_modules: ClassVar[list[str]] + _resident_modules: ClassVar[list[str]] = [] + + def supports_step_execution(pipeline: object) -> bool: """Return whether `pipeline` implements :class:`SupportsStepExecution`.""" diff --git a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2.py b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2.py index f06ffab16..4f62d72c9 100644 --- a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2.py +++ b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2.py @@ -9,7 +9,7 @@ import json import os from collections.abc import Iterable from contextlib import nullcontext -from typing import Any +from typing import Any, ClassVar import numpy as np import torch @@ -34,6 +34,7 @@ from vllm_omni.diffusion.distributed.utils import get_local_device from vllm_omni.diffusion.lora.manager import DiffusionLoRAManager from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader from vllm_omni.diffusion.models.dmd2 import DMD2PipelineMixin +from vllm_omni.diffusion.models.interface import SupportsModuleOffload from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.lora.request import LoRARequest @@ -1152,9 +1153,13 @@ class LTX2Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin): return loader.load_weights(weights) -class LTX2TwoStagesPipeline(nn.Module): +class LTX2TwoStagesPipeline(nn.Module, SupportsModuleOffload): """LTX2TwoStagesPipeline is for two stages image to video generation""" + _dit_modules: ClassVar[list[str]] = ["pipe.transformer"] + _encoder_modules: ClassVar[list[str]] = ["pipe.text_encoder"] + _vae_modules: ClassVar[list[str]] = ["pipe.vae", "pipe.audio_vae"] + def __init__( self, *, diff --git a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_image2video.py b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_image2video.py index 50a71a54b..4cc65f749 100644 --- a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_image2video.py +++ b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_image2video.py @@ -6,7 +6,7 @@ from __future__ import annotations import copy import os from collections.abc import Iterable -from typing import Any +from typing import Any, ClassVar import numpy as np import PIL.Image @@ -26,6 +26,7 @@ from vllm_omni.diffusion.distributed.utils import get_local_device from vllm_omni.diffusion.lora.manager import DiffusionLoRAManager from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader from vllm_omni.diffusion.models.dmd2 import DMD2PipelineMixin +from vllm_omni.diffusion.models.interface import SupportsModuleOffload from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.lora.request import LoRARequest @@ -733,11 +734,15 @@ class LTX2ImageToVideoPipeline(LTX2Pipeline): return DiffusionOutput(output=(video, audio)) -class LTX2ImageToVideoTwoStagesPipeline(nn.Module): +class LTX2ImageToVideoTwoStagesPipeline(nn.Module, SupportsModuleOffload): """LTXImageToVideoTwoStagesPipeline is for two stages image to video generation""" support_image_input = True + _dit_modules: ClassVar[list[str]] = ["pipe.transformer"] + _encoder_modules: ClassVar[list[str]] = ["pipe.text_encoder"] + _vae_modules: ClassVar[list[str]] = ["pipe.vae", "pipe.audio_vae"] + def __init__( self, *, diff --git a/vllm_omni/diffusion/models/omnigen2/pipeline_omnigen2.py b/vllm_omni/diffusion/models/omnigen2/pipeline_omnigen2.py index 04720c932..46d634bfd 100644 --- a/vllm_omni/diffusion/models/omnigen2/pipeline_omnigen2.py +++ b/vllm_omni/diffusion/models/omnigen2/pipeline_omnigen2.py @@ -8,7 +8,7 @@ import os import warnings from collections.abc import Iterable from dataclasses import dataclass -from typing import Any +from typing import Any, ClassVar import numpy as np import PIL.Image @@ -32,6 +32,7 @@ from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin from vllm_omni.diffusion.distributed.utils import get_local_device from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader +from vllm_omni.diffusion.models.interface import SupportsModuleOffload from vllm_omni.diffusion.models.omnigen2.omnigen2_transformer import ( OmniGen2RotaryPosEmbed, OmniGen2Transformer2DModel, @@ -620,7 +621,7 @@ def retrieve_timesteps( return timesteps, num_inference_steps -class OmniGen2Pipeline(CFGParallelMixin, nn.Module): +class OmniGen2Pipeline(CFGParallelMixin, nn.Module, SupportsModuleOffload): """ Pipeline for text-to-image generation using OmniGen2. @@ -634,6 +635,10 @@ class OmniGen2Pipeline(CFGParallelMixin, nn.Module): od_config (OmniDiffusionConfig): The OmniDiffusion configuration. """ + _dit_modules: ClassVar[list[str]] = ["transformer"] + _encoder_modules: ClassVar[list[str]] = ["mllm"] + _vae_modules: ClassVar[list[str]] = ["vae"] + def __init__( self, *, diff --git a/vllm_omni/diffusion/offloader/layerwise_backend.py b/vllm_omni/diffusion/offloader/layerwise_backend.py index d1216e2f2..9979d01b1 100644 --- a/vllm_omni/diffusion/offloader/layerwise_backend.py +++ b/vllm_omni/diffusion/offloader/layerwise_backend.py @@ -298,13 +298,20 @@ class LayerWiseOffloadBackend(OffloadBackend): for enc in modules.encoders: enc.to(self.device) - # Move VAE to GPU if available - if modules.vae is not None: + # Move VAE(s) to GPU if available + for vae in modules.vaes: try: - modules.vae.to(self.device, non_blocking=True) + vae.to(self.device, non_blocking=True) except Exception as exc: logger.debug("Failed to move VAE to GPU: %s", exc) + # Move resident modules to GPU (small modules needed every forward) + for name, module in zip(modules.resident_names, modules.resident_modules): + try: + module.to(self.device) + except Exception as exc: + logger.debug("Failed to move resident module %s to GPU: %s", name, exc) + logger.info("Applying layer-wise offloading on %s", modules.dit_names) # Apply block-wise offloading hook for each of the blocks in DiT model(s) diff --git a/vllm_omni/diffusion/offloader/module_collector.py b/vllm_omni/diffusion/offloader/module_collector.py index a09a33700..dfd81e98b 100644 --- a/vllm_omni/diffusion/offloader/module_collector.py +++ b/vllm_omni/diffusion/offloader/module_collector.py @@ -1,11 +1,14 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from dataclasses import dataclass +from dataclasses import dataclass, field +from operator import attrgetter from torch import nn from vllm.logger import init_logger +from vllm_omni.diffusion.models.interface import SupportsModuleOffload + logger = init_logger(__name__) @@ -15,15 +18,81 @@ class PipelineModules: dit_names: list[str] encoders: list[nn.Module] encoder_names: list[str] - vae: nn.Module | None = None + vaes: list[nn.Module] + resident_modules: list[nn.Module] = field(default_factory=list) + resident_names: list[str] = field(default_factory=list) class ModuleDiscovery: - """Discovers pipeline components for offloading""" + """Discovers pipeline components for offloading. - DIT_ATTRS = ["transformer", "transformer_2", "dit", "sr_dit", "language_model", "transformer_blocks", "model"] - ENCODER_ATTRS = ["text_encoder", "text_encoder_2", "text_encoder_3", "image_encoder"] - VAE_ATTRS = ["vae", "audio_vae"] + If the pipeline implements :class:`SupportsModuleOffload`, + its ``_dit_modules``, ``_encoder_modules``, and ``_vae_modules`` + class variables are used directly. Otherwise, falls back to + scanning well-known attribute names. + """ + + # Fallback attribute names for pipelines that do not implement + # SupportsModuleOffload. + _FALLBACK_DIT_ATTRS = [ + "transformer", + "transformer_2", + "dit", + "sr_dit", + "language_model", + "transformer_blocks", + "model", + ] + _FALLBACK_ENCODER_ATTRS = [ + "text_encoder", + "text_encoder_2", + "text_encoder_3", + "image_encoder", + ] + _FALLBACK_VAE_ATTRS = [ + "vae", + "audio_vae", + ] + + @staticmethod + def _collect_modules( + pipeline: nn.Module, + attr_names: list[str], + *, + warn_missing: bool = False, + ) -> tuple[list[nn.Module], list[str]]: + """Resolve attribute names to (module, name) pairs, skipping missing. + + Supports dotted paths via :func:`operator.attrgetter`. + Warns on missing attributes when *warn_missing* is True. + """ + modules: list[nn.Module] = [] + names: list[str] = [] + seen: set[int] = set() + for attr in attr_names: + try: + module = attrgetter(attr)(pipeline) + except AttributeError: + module = None + if module is None: + if warn_missing: + logger.warning( + "Pipeline declares '%s' as offloadable but the attribute does not exist or is None", + attr, + ) + continue + if not isinstance(module, nn.Module): + logger.warning( + "Expected '%s' to be nn.Module, got %r", + attr, + type(module), + ) + continue + if id(module) not in seen: + seen.add(id(module)) + modules.append(module) + names.append(attr) + return modules, names @staticmethod def discover(pipeline: nn.Module) -> PipelineModules: @@ -35,46 +104,29 @@ class ModuleDiscovery: Returns: PipelineModules with lists of discovered modules and names """ - # Collect DiT/transformer modules - dit_modules: list[nn.Module] = [] - dit_names: list[str] = [] - for attr in ModuleDiscovery.DIT_ATTRS: - if not hasattr(pipeline, attr): - continue - module_obj = getattr(pipeline, attr) - if module_obj is None: - continue + declared = isinstance(pipeline, SupportsModuleOffload) + if declared: + dit_attrs = pipeline._dit_modules + enc_attrs = pipeline._encoder_modules + vae_attrs = pipeline._vae_modules + res_attrs = pipeline._resident_modules + else: + dit_attrs = ModuleDiscovery._FALLBACK_DIT_ATTRS + enc_attrs = ModuleDiscovery._FALLBACK_ENCODER_ATTRS + vae_attrs = ModuleDiscovery._FALLBACK_VAE_ATTRS + res_attrs = [] - if not isinstance(module_obj, nn.Module): - logger.warning(f"Expected {attr} to be nn.Module, got {type(module_obj)!r}") - continue - - if module_obj in dit_modules: - continue - - dit_modules.append(module_obj) - dit_names.append(attr) - - # Collect all encoders - encoders: list[nn.Module] = [] - encoder_names: list[str] = [] - for attr in ModuleDiscovery.ENCODER_ATTRS: - if hasattr(pipeline, attr) and getattr(pipeline, attr) is not None: - encoders.append(getattr(pipeline, attr)) - encoder_names.append(attr) - - # Collect VAE - vae = None - for attr in ModuleDiscovery.VAE_ATTRS: - module = getattr(pipeline, attr, None) - if module is not None: - vae = module - break + dit_modules, dit_names = ModuleDiscovery._collect_modules(pipeline, dit_attrs, warn_missing=declared) + encoders, encoder_names = ModuleDiscovery._collect_modules(pipeline, enc_attrs, warn_missing=declared) + vaes, _ = ModuleDiscovery._collect_modules(pipeline, vae_attrs, warn_missing=declared) + residents, resident_names = ModuleDiscovery._collect_modules(pipeline, res_attrs, warn_missing=declared) return PipelineModules( dits=dit_modules, dit_names=dit_names, encoders=encoders, encoder_names=encoder_names, - vae=vae, + vaes=vaes, + resident_modules=residents, + resident_names=resident_names, ) diff --git a/vllm_omni/diffusion/offloader/sequential_backend.py b/vllm_omni/diffusion/offloader/sequential_backend.py index 46f48e99c..06454ad5c 100644 --- a/vllm_omni/diffusion/offloader/sequential_backend.py +++ b/vllm_omni/diffusion/offloader/sequential_backend.py @@ -210,10 +210,10 @@ class ModelLevelOffloadBackend(OffloadBackend): for enc in modules.encoders: enc.to(self.device) - # Move VAE to GPU if available - if modules.vae is not None: + # Move VAE(s) to GPU if available + for vae in modules.vaes: try: - modules.vae.to(self.device, non_blocking=True) + vae.to(self.device, non_blocking=True) except Exception as exc: logger.debug("Failed to move VAE to GPU: %s", exc) From fd9a0ee27f7d70da5f6b344a93ac1175f556f286 Mon Sep 17 00:00:00 2001 From: wuhang Date: Thu, 23 Apr 2026 04:32:43 +0800 Subject: [PATCH 16/45] [Feature] Failure message shows more details (#2961) Signed-off-by: wuhang --- .../openai_api/test_serving_speech.py | 72 +++++++++++++++ tests/entrypoints/test_omni_entrypoints.py | 23 ++++- vllm_omni/engine/orchestrator.py | 5 + vllm_omni/entrypoints/async_omni.py | 24 ++++- vllm_omni/entrypoints/omni_base.py | 28 +++++- vllm_omni/entrypoints/openai/api_server.py | 91 +++++++++++++------ 6 files changed, 205 insertions(+), 38 deletions(-) diff --git a/tests/entrypoints/openai_api/test_serving_speech.py b/tests/entrypoints/openai_api/test_serving_speech.py index b388b1860..b78d62d9e 100644 --- a/tests/entrypoints/openai_api/test_serving_speech.py +++ b/tests/entrypoints/openai_api/test_serving_speech.py @@ -18,6 +18,7 @@ from pydantic import ValidationError from pytest_mock import MockerFixture from vllm.entrypoints.openai.engine.protocol import ErrorInfo, ErrorResponse +from vllm_omni.entrypoints.omni_base import OmniEngineDeadError from vllm_omni.entrypoints.openai import api_server as api_server_module from vllm_omni.entrypoints.openai.audio_utils_mixin import AudioMixin from vllm_omni.entrypoints.openai.protocol.audio import ( @@ -1707,6 +1708,77 @@ def test_api_server_create_speech_wraps_error_response_status(mocker: MockerFixt assert response.status_code == 400 +def test_api_server_create_speech_engine_error_response_includes_request_and_stage_id(mocker: MockerFixture): + handler = mocker.MagicMock() + handler.create_speech = mocker.AsyncMock( + side_effect=OmniEngineDeadError( + "engine dead", + error_stage_id=1, + ) + ) + + terminate_mock = mocker.patch.object(api_server_module, "terminate_if_errored") + + app = FastAPI() + app.state.args = SimpleNamespace(log_error_stack=False) + app.state.openai_serving_speech = handler + app.state.engine_client = SimpleNamespace( + engine=SimpleNamespace(is_alive=lambda: False), + errored=True, + ) + app.state.server = SimpleNamespace() + scope = { + "type": "http", + "app": app, + "method": "POST", + "path": "/v1/audio/speech", + "headers": [], + "query_string": b"", + "client": ("127.0.0.1", 12345), + "server": ("testserver", 80), + "scheme": "http", + } + raw_request = Request(scope) + raw_request.state.request_metadata = SimpleNamespace(request_id="speech-req-1") + request = OpenAICreateSpeechRequest(input="Hello") + + response = asyncio.run(api_server_module.create_speech(request, raw_request)) + + assert isinstance(response, JSONResponse) + assert response.status_code == 500 + assert response.body.decode("utf-8") == ( + '{"error":{"message":"engine dead","type":"InternalServerError","param":null,' + '"code":500,"request_id":"speech-req-1","error_stage_id":1}}' + ) + terminate_mock.assert_called_once() + + +def test_omni_engine_error_handler_includes_request_and_stage_id(mocker: MockerFixture): + app = FastAPI() + app.state.args = SimpleNamespace(log_error_stack=False) + app.state.engine_client = SimpleNamespace( + engine=SimpleNamespace(is_alive=lambda: False), + errored=True, + ) + app.state.server = SimpleNamespace() + + terminate_mock = mocker.patch.object(api_server_module, "terminate_if_errored") + api_server_module._register_omni_exception_handlers(app) + + @app.get("/boom") + async def boom(request: Request): + request.state.request_metadata = SimpleNamespace(request_id="speech-req-1") + exc = OmniEngineDeadError("engine dead", error_stage_id=1) + raise exc + + response = TestClient(app).get("/boom") + + assert response.status_code == 500 + assert response.json()["error"]["request_id"] == "speech-req-1" + assert response.json()["error"]["error_stage_id"] == 1 + terminate_mock.assert_called_once() + + class TestWAVHeaderGeneration: """Unit tests for WAV header generation with placeholder values.""" diff --git a/tests/entrypoints/test_omni_entrypoints.py b/tests/entrypoints/test_omni_entrypoints.py index 289cf2673..341bf6bc0 100644 --- a/tests/entrypoints/test_omni_entrypoints.py +++ b/tests/entrypoints/test_omni_entrypoints.py @@ -14,6 +14,7 @@ from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError from vllm_omni.entrypoints.async_omni import AsyncOmni from vllm_omni.entrypoints.omni import Omni +from vllm_omni.entrypoints.omni_base import OmniEngineDeadError pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -335,6 +336,18 @@ def _enqueue_error_message(engine: FakeAsyncOmniEngine, msg: dict[str, Any]) -> ) +def _enqueue_fatal_error_message(engine: FakeAsyncOmniEngine, msg: dict[str, Any]) -> None: + engine.output_q.put_nowait( + { + "type": "error", + "fatal": True, + "request_id": msg["request_id"], + "stage_id": 2, + "error": "engine dead", + } + ) + + @pytest.mark.asyncio async def test_get_supported_tasks_returns_engine_supported_tasks(): omni = object.__new__(AsyncOmni) @@ -546,18 +559,22 @@ async def test_async_omni_abort_forwards_to_engine(monkeypatch: pytest.MonkeyPat @pytest.mark.asyncio -async def test_async_omni_propagates_engine_error(monkeypatch: pytest.MonkeyPatch): - engine = FakeAsyncOmniEngine(stage_metadata=THREE_STAGE_META, on_add_request=_enqueue_error_message) +async def test_async_omni_propagates_fatal_error_context(monkeypatch: pytest.MonkeyPatch): + engine = FakeAsyncOmniEngine(stage_metadata=THREE_STAGE_META, on_add_request=_enqueue_fatal_error_message) _patch_engine(monkeypatch, engine) app = AsyncOmni("dummy-model") try: - with pytest.raises(RuntimeError, match="engine boom"): + with pytest.raises(EngineDeadError, match="engine dead") as exc_info: async for _ in app.generate(prompt="hello", request_id="req-1"): pass finally: app.shutdown() + assert isinstance(exc_info.value, OmniEngineDeadError) + assert str(exc_info.value) == "engine dead" + assert getattr(exc_info.value, "error_stage_id") == 2 + def test_omni_generate_py_generator_yields_final_outputs_for_each_request(monkeypatch: pytest.MonkeyPatch): sampling_params = [SamplingParams(max_tokens=8) for _ in range(3)] diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py index ca60c7531..ff5117934 100644 --- a/vllm_omni/engine/orchestrator.py +++ b/vllm_omni/engine/orchestrator.py @@ -181,6 +181,7 @@ class Orchestrator: self._shutdown_event = asyncio.Event() self._stages_shutdown = False self._fatal_error: str | None = None + self._fatal_error_stage_id: int | None = None async def run(self) -> None: """Main entry point for the Orchestrator event loop.""" @@ -331,6 +332,7 @@ class Orchestrator: e, ) self._fatal_error = str(e) + self._fatal_error_stage_id = stage_id for req_id, req_state in list(self.request_states.items()): if stage_id in req_state.stage_submit_ts: await self.output_async_queue.put( @@ -339,6 +341,7 @@ class Orchestrator: "error": str(e), "fatal": True, "request_id": req_id, + "stage_id": stage_id, } ) self.request_states.pop(req_id, None) @@ -1135,6 +1138,7 @@ class Orchestrator: "error": self._fatal_error, "fatal": True, "request_id": req_id, + "stage_id": self._fatal_error_stage_id, } ) notified.add(req_id) @@ -1150,6 +1154,7 @@ class Orchestrator: "error": self._fatal_error, "fatal": True, "request_id": req_id, + "stage_id": self._fatal_error_stage_id, } ) self.request_states.pop(req_id, None) diff --git a/vllm_omni/entrypoints/async_omni.py b/vllm_omni/entrypoints/async_omni.py index 67425ea2e..119232ce7 100644 --- a/vllm_omni/entrypoints/async_omni.py +++ b/vllm_omni/entrypoints/async_omni.py @@ -27,7 +27,10 @@ from vllm.v1.engine.exceptions import EngineDeadError from vllm_omni.diffusion.data import OmniACK, OmniSleepTask, OmniWakeTask from vllm_omni.entrypoints.client_request_state import ClientRequestState -from vllm_omni.entrypoints.omni_base import OmniBase +from vllm_omni.entrypoints.omni_base import ( + OmniBase, + OmniEngineDeadError, +) from vllm_omni.metrics.stats import OrchestratorAggregator as OrchestratorMetrics from vllm_omni.outputs import OmniRequestOutput from vllm_omni.platforms import current_omni_platform @@ -494,7 +497,10 @@ class AsyncOmni(EngineClient, OmniBase): stage_id = result.get("stage_id", 0) if result.get("type") == "error" and result.get("fatal"): - raise EngineDeadError(result.get("error", "")) + raise OmniEngineDeadError( + result.get("error", ""), + error_stage_id=result.get("stage_id"), + ) # Check for errors if "error" in result: @@ -580,6 +586,18 @@ class AsyncOmni(EngineClient, OmniBase): except asyncio.CancelledError: raise + except OmniEngineDeadError as e: + logger.error("[AsyncOmni] Engine dead: %s", e) + for req_state in list(self.request_states.values()): + error_msg = { + "type": "error", + "error": str(e), + "fatal": True, + "request_id": req_state.request_id, + } + if e.error_stage_id is not None: + error_msg["stage_id"] = e.error_stage_id + await req_state.queue.put(error_msg) except EngineDeadError as e: logger.error("[AsyncOmni] Engine dead: %s", e) for req_state in list(self.request_states.values()): @@ -884,7 +902,7 @@ class AsyncOmni(EngineClient, OmniBase): @property def dead_error(self) -> BaseException: """EngineClient abstract property implementation.""" - return EngineDeadError() + return OmniEngineDeadError() # ==================== EngineClient Interface ==================== diff --git a/vllm_omni/entrypoints/omni_base.py b/vllm_omni/entrypoints/omni_base.py index 3180c9c80..d922ab072 100644 --- a/vllm_omni/entrypoints/omni_base.py +++ b/vllm_omni/entrypoints/omni_base.py @@ -28,6 +28,23 @@ if TYPE_CHECKING: logger = init_logger(__name__) +class OmniEngineDeadError(EngineDeadError): + _DEFAULT_MESSAGE = EngineDeadError().args[0] + error_stage_id: int | None + + def __init__( + self, + message: str | None = None, + *, + error_stage_id: int | None = None, + suppress_context: bool = False, + ) -> None: + resolved_message = message or self._DEFAULT_MESSAGE + Exception.__init__(self, resolved_message) + self.__suppress_context__ = suppress_context + self.error_stage_id = error_stage_id + + def _weak_shutdown_engine(engine: AsyncOmniEngine) -> None: """Best-effort engine cleanup for GC finalization.""" try: @@ -296,8 +313,12 @@ class OmniBase(PDDisaggregationMixin): if msg_type == "error": error_text = msg.get("error", "Orchestrator returned an error message") + stage_id = msg.get("stage_id") if msg.get("fatal"): - raise EngineDeadError(error_text) + raise OmniEngineDeadError( + error_text, + error_stage_id=stage_id, + ) raise RuntimeError(error_text) if msg_type != "output": @@ -352,7 +373,10 @@ class OmniBase(PDDisaggregationMixin): ) # NOTE: O(n_stages) check for every error. if self.errored: - raise EngineDeadError(error_text) + raise OmniEngineDeadError( + error_text, + error_stage_id=stage_id, + ) raise EngineGenerateError(error_text) def _process_single_result( diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index 41ccbcf1d..4454f5bda 100644 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/vllm_omni/entrypoints/openai/api_server.py @@ -220,33 +220,64 @@ def _register_omni_exception_handlers(app) -> None: req: Request, exc: EngineDeadError | EngineGenerateError, ): - request_id = req.state.request_metadata.request_id if hasattr(req.state, "request_metadata") else None + request_id = _get_request_id_from_request(req) if req.app.state.args.log_error_stack: logger.exception("Engine Exception caught. Request id: %s", request_id) - engine = req.app.state.engine_client - if isinstance(exc, EngineDeadError): - # Log Omni-specific diagnostic information for dead engines. - orchestrator_alive = engine.engine.is_alive() if hasattr(engine, "engine") else "N/A" - logger.error( - "EngineDeadError: orchestrator_alive=%s, errored=%s, request_id=%s", - orchestrator_alive, - engine.errored, - request_id, - ) - - terminate_if_errored( - server=req.app.state.server, - engine=engine, - ) - err = create_error_response(exc) - return JSONResponse(err.model_dump(), status_code=err.error.code) + return _create_engine_error_json_response(req, exc) app.exception_handler(EngineGenerateError)(omni_engine_error_handler) app.exception_handler(EngineDeadError)(omni_engine_error_handler) +def _get_request_id_from_request(req: Request) -> str | None: + return req.state.request_metadata.request_id if hasattr(req.state, "request_metadata") else None + + +def _build_engine_error_payload( + exc: EngineDeadError | EngineGenerateError, + *, + request_id: str | None, +) -> tuple[dict[str, Any], int]: + err = create_error_response(exc) + payload = err.model_dump() + error_body = payload.get("error", {}) + + error_body["request_id"] = request_id + error_body["error_stage_id"] = getattr(exc, "error_stage_id", None) + + return payload, err.error.code + + +def _create_engine_error_json_response( + req: Request, + exc: EngineDeadError | EngineGenerateError, +) -> JSONResponse: + request_id = _get_request_id_from_request(req) + error_stage_id = getattr(exc, "error_stage_id", None) + engine = req.app.state.engine_client + + if isinstance(exc, EngineDeadError): + # Log Omni-specific diagnostic information for dead engines. + orchestrator_alive = engine.engine.is_alive() if hasattr(engine, "engine") else "N/A" + logger.error( + "EngineDeadError: orchestrator_alive=%s, errored=%s, request_id=%s, error_stage_id=%s", + orchestrator_alive, + engine.errored, + request_id, + error_stage_id, + ) + + terminate_if_errored( + server=req.app.state.server, + engine=engine, + ) + + payload, status_code = _build_engine_error_payload(exc, request_id=request_id) + return JSONResponse(content=payload, status_code=status_code) + + class _DiffusionServingModels: """Minimal OpenAIServingModels implementation for diffusion-only servers. @@ -930,8 +961,8 @@ async def create_chat_completion(request: ChatCompletionRequest, raw_request: Re return base_server.create_error_response(message="The model does not support Chat Completions API") try: generator = await handler.create_chat_completion(request, raw_request) - except (EngineGenerateError, EngineDeadError): - raise # Propagate to the global Omni exception handler + except (EngineGenerateError, EngineDeadError) as exc: + return _create_engine_error_json_response(raw_request, exc) except Exception as e: logger.exception("Chat completion failed: %s", e) raise HTTPException(status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value, detail=str(e)) from e @@ -1027,8 +1058,8 @@ async def create_speech(request: OpenAICreateSpeechRequest, raw_request: Request status_code=result.error.code if result.error else 400, ) return result - except (EngineGenerateError, EngineDeadError): - raise # Propagate to the global Omni exception handler + except (EngineGenerateError, EngineDeadError) as exc: + return _create_engine_error_json_response(raw_request, exc) except Exception as e: raise HTTPException(status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value, detail=str(e)) from e @@ -1063,8 +1094,8 @@ async def create_speech_batch(request: BatchSpeechRequest, raw_request: Request) status_code=result.error.code if result.error else 400, ) return JSONResponse(content=result.model_dump()) - except (EngineGenerateError, EngineDeadError): - raise # Propagate to the global Omni exception handler + except (EngineGenerateError, EngineDeadError) as exc: + return _create_engine_error_json_response(raw_request, exc) except ValueError as e: raise HTTPException(status_code=HTTPStatus.BAD_REQUEST.value, detail=str(e)) from e except Exception as e: @@ -1522,8 +1553,8 @@ async def generate_images(request: ImageGenerationRequest, raw_request: Request) data=image_data, ) - except (EngineGenerateError, EngineDeadError): - raise # Propagate to the global Omni exception handler + except (EngineGenerateError, EngineDeadError) as exc: + return _create_engine_error_json_response(raw_request, exc) except HTTPException: raise except ValueError as e: @@ -1759,8 +1790,8 @@ async def edit_images( size=size_str, ) - except (EngineGenerateError, EngineDeadError): - raise # Propagate to the global Omni exception handler + except (EngineGenerateError, EngineDeadError) as exc: + return _create_engine_error_json_response(raw_request, exc) except HTTPException: raise except ValueError as e: @@ -2481,8 +2512,8 @@ async def create_video_sync( status_code=HTTPStatus.GATEWAY_TIMEOUT.value, detail=f"Video generation timed out after {VIDEO_SYNC_TIMEOUT_S}s.", ) - except (EngineGenerateError, EngineDeadError): - raise # Propagate to the global Omni exception handler + except (EngineGenerateError, EngineDeadError) as exc: + return _create_engine_error_json_response(raw_request, exc) except HTTPException: raise except Exception as exc: From 3aa84c47422887562203895c53b54acc86be6a12 Mon Sep 17 00:00:00 2001 From: amy-why-3459 Date: Thu, 23 Apr 2026 09:40:43 +0800 Subject: [PATCH 17/45] [Tests]Add accuracy benchmark L4 test cases (#2843) Signed-off-by: amy-why-3459 --- .buildkite/test-nightly.yml | 47 +- pyproject.toml | 14 +- tests/benchmarks/test_accuracy_bench_utils.py | 61 +++ tests/e2e/accuracy/qwen3_omni/__init__.py | 2 + .../qwen3_omni/qwen3_omni_acc_bench_core.py | 201 ++++++++ .../qwen3_omni/run_qwen_omni_acc_benchmark.py | 428 ++++++++++++++++++ .../accuracy/qwen3_omni/test_qwen3_omni.py | 137 ++++++ .../data_modules/daily_omni_dataset.py | 148 +++++- .../data_modules/daily_omni_eval.py | 13 +- .../data_modules/seed_tts_dataset.py | 13 +- .../benchmarks/data_modules/seed_tts_eval.py | 6 +- vllm_omni/entrypoints/cli/benchmark/serve.py | 4 +- 12 files changed, 1044 insertions(+), 30 deletions(-) create mode 100644 tests/e2e/accuracy/qwen3_omni/__init__.py create mode 100644 tests/e2e/accuracy/qwen3_omni/qwen3_omni_acc_bench_core.py create mode 100644 tests/e2e/accuracy/qwen3_omni/run_qwen_omni_acc_benchmark.py create mode 100644 tests/e2e/accuracy/qwen3_omni/test_qwen3_omni.py diff --git a/.buildkite/test-nightly.yml b/.buildkite/test-nightly.yml index e6912bd49..82d862661 100644 --- a/.buildkite/test-nightly.yml +++ b/.buildkite/test-nightly.yml @@ -105,6 +105,51 @@ steps: path: /mnt/hf-cache type: DirectoryOrCreate + - label: ":full_moon: Omni · Accuracy Test" + timeout_in_minutes: 180 + commands: + - export SEED_TTS_WER_EVAL=1 + - export SEED_TTS_EVAL_DEVICE=cuda:1 + - | + set +e + pytest -s -v tests/e2e/accuracy/qwen3_omni/test_qwen3_omni.py -m "full_model" --run-level full_model + EXIT=$$? + buildkite-agent artifact upload "tests/e2e/accuracy/qwen3_omni/results/qwen_omni_acc/*.json" + exit $$EXIT + agents: + queue: "mithril-h100-pool" + plugins: + - kubernetes: + podSpec: + containers: + - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT + resources: + limits: + nvidia.com/gpu: 2 + volumeMounts: + - name: devshm + mountPath: /dev/shm + - name: hf-cache + mountPath: /root/.cache/huggingface + env: + - name: HF_HOME + value: /root/.cache/huggingface + - name: HF_TOKEN + valueFrom: + secretKeyRef: + name: hf-token-secret + key: token + nodeSelector: + node.kubernetes.io/instance-type: gpu-h100-sxm + volumes: + - name: devshm + emptyDir: + medium: Memory + - name: hf-cache + hostPath: + path: /mnt/hf-cache + type: DirectoryOrCreate + - label: ":full_moon: Omni · Perf Test" key: nightly-omni-performance timeout_in_minutes: 180 @@ -514,7 +559,7 @@ steps: - label: ":full_moon: Diffusion X2V · Function Test" timeout_in_minutes: 90 commands: - - pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py tests/e2e/online_serving/test_wan_2_1_vace_expansion.py tests/e2e/online_serving/test_hunyuan_video_15_expansion.py -m "full_model" --run-level "full_model" + - pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py tests/e2e/online_serving/test_wan_2_1_vace_expansion.py tests/e2e/online_serving/test_hunyuan_video_15_expansion.py --run-level "full_model" agents: queue: "mithril-h100-pool" plugins: diff --git a/pyproject.toml b/pyproject.toml index c7b2be20e..8346693f1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -56,23 +56,17 @@ dev = [ "opencc>=1.2.0", "mistune>=3.2.0", # for example tests "torchmetrics>=1.4.0", # for accuracy similarity metrics -] - -demo = [ - "gradio>=6.7.0", -] - -# Seed-TTS serve benchmark WER (BytedanceSpeech/seed-tts-eval run_wer.py protocol). -seed-tts-eval = [ "jiwer>=3.0.0", "zhon>=2.0.0", "zhconv>=1.4.2", "scipy>=1.10.0", - "soundfile>=0.12.0", - "transformers>=4.36.0", "funasr>=1.0.0", ] +demo = [ + "gradio>=6.7.0", +] + docs = [ "mkdocs>=1.5.0", "mkdocs-api-autonav", diff --git a/tests/benchmarks/test_accuracy_bench_utils.py b/tests/benchmarks/test_accuracy_bench_utils.py index 8425804ae..6ceebb11b 100644 --- a/tests/benchmarks/test_accuracy_bench_utils.py +++ b/tests/benchmarks/test_accuracy_bench_utils.py @@ -1,6 +1,9 @@ # ruff: noqa: E402, I001 +import argparse import math +import os import sys +import types from pathlib import Path import pytest @@ -39,6 +42,64 @@ from benchmarks.accuracy.text_to_image.gbench import ( summarize_generated_records as summarize_gebench_generated_records, summarize_gebench_results, ) +from tests.e2e.accuracy.qwen3_omni.qwen3_omni_acc_bench_core import seed_tts_bench_argv +from tests.e2e.accuracy.qwen3_omni.run_qwen_omni_acc_benchmark import sync_dataset_env_from_ns +from vllm_omni.benchmarks.data_modules.seed_tts_dataset import resolve_seed_tts_root + + +def test_seed_tts_bench_argv_preserves_hf_repo_id_from_env(monkeypatch): + monkeypatch.setenv("VLLM_SEED_TTS_DATASET_PATH", "zhaochenyang20/seed-tts-eval") + monkeypatch.delenv("VLLM_SEED_TTS_REPO", raising=False) + + argv = seed_tts_bench_argv(locale="en") + + dataset_idx = argv.index("--dataset-path") + assert argv[dataset_idx + 1] == "zhaochenyang20/seed-tts-eval" + + +def test_sync_dataset_env_preserves_seed_tts_hf_repo_id(monkeypatch): + ns = argparse.Namespace( + daily_omni_repo=None, + daily_omni_qa_json=None, + daily_omni_video_dir=None, + seed_tts_dataset_path="zhaochenyang20/seed-tts-eval", + seed_tts_root=None, + ) + + monkeypatch.delenv("VLLM_SEED_TTS_DATASET_PATH", raising=False) + sync_dataset_env_from_ns(ns) + + assert os.environ["VLLM_SEED_TTS_DATASET_PATH"] == "zhaochenyang20/seed-tts-eval" + + +def test_resolve_seed_tts_root_downloads_only_requested_locale(monkeypatch, tmp_path: Path): + downloaded_root = tmp_path / "seed_tts_cache" + (downloaded_root / "zh" / "prompt-wavs").mkdir(parents=True) + (downloaded_root / "zh" / "meta.lst").write_text("", encoding="utf-8") + captured: dict[str, object] = {} + + def fake_snapshot_download(*, repo_id, repo_type, allow_patterns): + captured["repo_id"] = repo_id + captured["repo_type"] = repo_type + captured["allow_patterns"] = allow_patterns + return str(downloaded_root) + + monkeypatch.setitem( + sys.modules, + "huggingface_hub", + types.SimpleNamespace(snapshot_download=fake_snapshot_download), + ) + + resolved = resolve_seed_tts_root( + "zhaochenyang20/seed-tts-eval", + explicit_root=None, + locale="zh", + ) + + assert resolved == downloaded_root.resolve() + assert captured["repo_id"] == "zhaochenyang20/seed-tts-eval" + assert captured["repo_type"] == "dataset" + assert captured["allow_patterns"] == ["zh/**"] def test_summarize_gebench_generated_records_groups_by_type(): diff --git a/tests/e2e/accuracy/qwen3_omni/__init__.py b/tests/e2e/accuracy/qwen3_omni/__init__.py new file mode 100644 index 000000000..79a31c4f1 --- /dev/null +++ b/tests/e2e/accuracy/qwen3_omni/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Qwen3-Omni accuracy benchmarks (Daily-Omni / Seed-TTS ``vllm bench serve --omni``).""" diff --git a/tests/e2e/accuracy/qwen3_omni/qwen3_omni_acc_bench_core.py b/tests/e2e/accuracy/qwen3_omni/qwen3_omni_acc_bench_core.py new file mode 100644 index 000000000..2ce86d504 --- /dev/null +++ b/tests/e2e/accuracy/qwen3_omni/qwen3_omni_acc_bench_core.py @@ -0,0 +1,201 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Shared helpers for Qwen3-Omni Daily-Omni / Seed-TTS ``vllm bench serve --omni`` accuracy runs. + +Local dataset paths are **optional**. When ``VLLM_DAILY_OMNI_QA_JSON`` + ``VLLM_DAILY_OMNI_VIDEO_DIR`` +point to existing files, those are used with inline video. Otherwise the benchmark falls back to +the HuggingFace dataset id (``liarliar/Daily-Omni``); QA loads via ``datasets``, and the first +bench request that needs media downloads ``Videos.tar`` from the Hub when no video dir is set. + +Similarly for Seed-TTS: a local directory wins; otherwise ``--dataset-path`` uses the Hub id +and ``huggingface_hub.snapshot_download`` inside ``resolve_seed_tts_root`` pulls files on demand. + +Use :func:`build_acc_benchmark_cli_argv` to assemble ``argv`` for a live Omni server (host/port/model +and small bench defaults) before ``parse_args`` / ``run_acc_benchmark`` in the accuracy driver. +""" + +from __future__ import annotations + +import json +import os +import shutil +import subprocess +from pathlib import Path +from typing import Any, Protocol + +DEFAULT_DAILY_OMNI_HF_REPO = "liarliar/Daily-Omni" +DEFAULT_SEED_TTS_HF_REPO = "zhaochenyang20/seed-tts-eval" + + +class OmniBenchServerEndpoint(Protocol): + """Anything with ``host`` / ``port`` / ``model`` (e.g. :class:`tests.conftest.OmniServer`).""" + + host: str + port: int + model: str + + +def build_acc_benchmark_cli_argv( + server: OmniBenchServerEndpoint, + *, + skip_seed: bool, + skip_daily: bool, + num_prompts: int | None = None, + max_concurrency: int | None = None, +) -> list[str]: + """Prefix argv for :func:`run_qwen_omni_acc_benchmark.parse_acc_benchmark_args` + :func:`run_acc_benchmark`. + + Wires ``--host`` / ``--port`` / ``--model`` to a running Omni OpenAI server, sets small + ``--num-prompts`` / ``--max-concurrency`` defaults (overridable via ``ACC_BENCH_NUM_PROMPTS`` / + ``ACC_BENCH_MAX_CONCURRENCY``), and when Daily-Omni runs adds ``--daily-omni-repo`` so Hub QA + matches :func:`daily_omni_bench_argv` once ``run_acc_benchmark`` mirrors ``--daily-omni-repo`` into env. + """ + n_prompts = int(os.environ.get("ACC_BENCH_NUM_PROMPTS", "2000")) if num_prompts is None else int(num_prompts) + n_conc = int(os.environ.get("ACC_BENCH_MAX_CONCURRENCY", "10")) if max_concurrency is None else int(max_concurrency) + argv = [ + "--host", + server.host, + "--port", + str(server.port), + "--model", + server.model, + "--num-prompts", + str(n_prompts), + "--max-concurrency", + str(n_conc), + ] + if not skip_daily: + repo = os.environ.get("VLLM_DAILY_OMNI_REPO", DEFAULT_DAILY_OMNI_HF_REPO).strip() or DEFAULT_DAILY_OMNI_HF_REPO + argv.extend(["--daily-omni-repo", repo]) + if skip_seed: + argv.append("--skip-seed-tts") + if skip_daily: + argv.append("--skip-daily-omni") + return argv + + +def daily_omni_bench_argv() -> list[str]: + """CLI args for Daily-Omni (after ``vllm bench serve --omni``).""" + qa = os.environ.get("VLLM_DAILY_OMNI_QA_JSON", "").strip() + vd = os.environ.get("VLLM_DAILY_OMNI_VIDEO_DIR", "").strip() + if qa and vd: + qap = Path(qa).expanduser() + vdp = Path(vd).expanduser() + if qap.is_file() and vdp.is_dir(): + return [ + "--dataset-name", + "daily-omni", + "--daily-omni-qa-json", + str(qap.resolve()), + "--daily-omni-video-dir", + str(vdp.resolve()), + "--daily-omni-inline-local-video", + ] + repo = os.environ.get("VLLM_DAILY_OMNI_REPO", DEFAULT_DAILY_OMNI_HF_REPO).strip() or DEFAULT_DAILY_OMNI_HF_REPO + return [ + "--dataset-name", + "daily-omni", + "--dataset-path", + repo, + ] + + +def seed_tts_bench_argv(*, locale: str = "en") -> list[str]: + """CLI args for Seed-TTS (after ``vllm bench serve --omni``).""" + dp = os.environ.get("VLLM_SEED_TTS_DATASET_PATH", "").strip() + if dp: + p = Path(dp).expanduser() + # Preserve Hugging Face repo ids verbatim. Only canonicalize to an + # absolute path when the value actually exists as a local directory. + dataset_path = str(p.resolve()) if p.exists() and p.is_dir() else dp + else: + dataset_path = ( + os.environ.get("VLLM_SEED_TTS_REPO", DEFAULT_SEED_TTS_HF_REPO).strip() or DEFAULT_SEED_TTS_HF_REPO + ) + out = ["--dataset-name", "seed-tts", "--dataset-path", dataset_path] + root = os.environ.get("SEED_TTS_ROOT", "").strip() + if root: + out.extend(["--seed-tts-root", str(Path(root).expanduser().resolve())]) + out.extend(["--seed-tts-locale", locale]) + return out + + +def find_vllm_cli() -> str: + exe = shutil.which("vllm") + if not exe: + raise FileNotFoundError("Could not find `vllm` on PATH (install vLLM-Omni with CLI entrypoints).") + return exe + + +def run_vllm_bench_subprocess(vllm: str, argv: list[str], *, extra_env: dict[str, str] | None = None) -> None: + env = os.environ.copy() + if extra_env: + env.update(extra_env) + subprocess.run([vllm, *argv], env=env, check=True) + + +def load_benchmark_result(path: Path) -> dict[str, Any]: + with path.open(encoding="utf-8") as f: + return json.load(f) + + +def build_serve_common_argv( + *, + host: str, + port: int, + model: str, + num_prompts: int, + max_concurrency: int, + num_warmups: int, + percentile_metrics: str, + result_dir: Path, + result_filename: str, + ready_check_timeout_sec: int | None = None, +) -> list[str]: + out = [ + "bench", + "serve", + "--omni", + "--host", + host, + "--port", + str(port), + "--model", + model, + "--endpoint", + "/v1/chat/completions", + "--backend", + "openai-chat-omni", + "--request-rate", + "inf", + "--num-prompts", + str(num_prompts), + "--max-concurrency", + str(max_concurrency), + "--no-oversample", + "--num-warmups", + str(num_warmups), + "--percentile-metrics", + percentile_metrics, + "--save-result", + "--result-dir", + str(result_dir), + "--result-filename", + result_filename, + ] + if ready_check_timeout_sec is not None: + out.extend(["--ready-check-timeout-sec", str(int(ready_check_timeout_sec))]) + return out + + +def assert_daily_omni_scored(result: dict[str, Any]) -> None: + acc = result.get("daily_omni_accuracy") + assert acc is not None, "daily_omni_accuracy missing — wrong dataset or benchmark wiring" + assert int(result.get("daily_omni_evaluated_ok", 0) or 0) > 0, "no successful MCQ rows (daily_omni_evaluated_ok==0)" + + +def assert_seed_tts_scored(result: dict[str, Any]) -> None: + err = result.get("seed_tts_eval_setup_error") + assert not err, f"Seed-TTS eval deps/setup failed: {err}" + assert int(result.get("seed_tts_content_evaluated", 0) or 0) > 0, ( + "seed_tts_content_evaluated==0 — enable WER eval and check PCM capture / modalities" + ) diff --git a/tests/e2e/accuracy/qwen3_omni/run_qwen_omni_acc_benchmark.py b/tests/e2e/accuracy/qwen3_omni/run_qwen_omni_acc_benchmark.py new file mode 100644 index 000000000..d30457dcd --- /dev/null +++ b/tests/e2e/accuracy/qwen3_omni/run_qwen_omni_acc_benchmark.py @@ -0,0 +1,428 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +"""Accuracy (and light perf) checks for Qwen3-Omni via ``vllm bench serve --omni``. + +The standalone CLI uses small ``--num-prompts`` / ``--max-concurrency`` defaults suitable for +L4-style smoke runs against an already-running server. The pytest wrappers in +``tests/e2e/accuracy/qwen3_omni/test_qwen3_omni.py`` may still require larger GPUs (currently +H100 / MI325) because they launch the live Omni server inside the test. + +1. **Daily-Omni** — MCQ accuracy fields in the saved JSON (``daily_omni_accuracy``, …); by default the + run **fails** if accuracy is strictly below **0.69** (``--min-daily-omni-accuracy`` / ``ACC_BENCH_MIN_DAILY_OMNI_ACCURACY``). +2. **Seed-TTS** — ``seed-tts-eval``-style metrics when ``--seed-tts-wer-eval`` is used + (WER / SIM / UTMOS keys from :func:`compute_seed_tts_wer_metrics`). + +Prerequisites +------------- +* A running Omni OpenAI-compatible server (same machine or reachable host), e.g.:: + + vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8000 + + On L4 you may need a smaller checkpoint, quantization, or tighter engine flags; this script + only drives the **client** benchmark. + +* ``vllm`` CLI from **vLLM-Omni** (so ``bench serve`` registers ``daily-omni`` / ``seed-tts``). + +* **Daily-Omni** — if local ``qa.json`` + ``Videos/`` are not both provided (CLI or matching env), + the client passes ``--dataset-path`` with a Hub id (default ``liarliar/Daily-Omni``). The **child** + ``vllm bench serve`` process then loads QA via ``datasets.load_dataset`` (needs ``pip install datasets``, + network or HF cache). Without ``--daily-omni-video-dir``, the benchmark **lazily** downloads and + extracts ``Videos.tar`` from the Hub (``huggingface_hub``) on first multimodal request. Override + the dataset repo with ``--daily-omni-repo`` or ``VLLM_DAILY_OMNI_REPO``; override the tar repo + with ``VLLM_DAILY_OMNI_MEDIA_REPO`` if needed. + +* **Seed-TTS** optional extras for WER/SIM/UTMOS:: + + pip install 'vllm-omni[seed-tts-eval]' + +Examples +-------- +Pytest (same checks; needs a running server):: + + pytest -sv tests/e2e/accuracy/qwen3_omni/test_qwen3_omni.py + +Smoke on localhost (server already up):: + + python tests/e2e/accuracy/qwen3_omni/run_qwen_omni_acc_benchmark.py \\ + --model Qwen/Qwen3-Omni-30B-A3B-Instruct \\ + --daily-omni-qa-json ./qa.json \\ + --daily-omni-video-dir ./Videos \\ + --seed-tts-dataset-path ./seed-tts-eval + +Skip one suite, tighten gates:: + + python tests/e2e/accuracy/qwen3_omni/run_qwen_omni_acc_benchmark.py \\ + --skip-daily-omni \\ + --max-seed-tts-mean-wer 0.35 \\ + --min-seed-tts-mean-sim 0.75 +""" + +from __future__ import annotations + +import argparse +import contextlib +import json +import os +import sys +from datetime import datetime +from pathlib import Path +from typing import Any + +from tests.e2e.accuracy.qwen3_omni.qwen3_omni_acc_bench_core import ( + build_serve_common_argv, + daily_omni_bench_argv, + find_vllm_cli, + load_benchmark_result, + run_vllm_bench_subprocess, + seed_tts_bench_argv, +) + +_REPO_ROOT = Path(__file__).resolve().parents[4] + + +def _repo_root() -> Path: + return _REPO_ROOT + + +def _default_result_dir() -> Path: + return Path(__file__).resolve().parent / "results" / "qwen_omni_acc" + + +def _validate_daily_omni(result: dict[str, Any], *, min_accuracy: float | None) -> list[str]: + errs: list[str] = [] + acc = result.get("daily_omni_accuracy") + if acc is None: + errs.append("Missing daily_omni_accuracy (wrong dataset or no gold-evaluated rows).") + return errs + ev = int(result.get("daily_omni_evaluated_ok", 0) or 0) + if ev <= 0: + errs.append("daily_omni_evaluated_ok is 0; no successful MCQ rows to score.") + if min_accuracy is not None and float(acc) + 1e-12 < float(min_accuracy): + errs.append(f"daily_omni_accuracy={acc:.6f} < --min-daily-omni-accuracy={min_accuracy}") + return errs + + +def _validate_seed_tts( + result: dict[str, Any], + *, + max_mean_wer: float | None, + min_mean_sim: float | None, + min_mean_utmos: float | None, +) -> list[str]: + errs: list[str] = [] + setup = result.get("seed_tts_eval_setup_error") + if setup: + errs.append(f"Seed-TTS eval setup failed: {setup}") + return errs + n = int(result.get("seed_tts_content_evaluated", 0) or 0) + if n <= 0: + errs.append("seed_tts_content_evaluated is 0 (enable --seed-tts-wer-eval and check PCM capture).") + mean_wer = result.get("seed_tts_content_error_mean") + if mean_wer is not None and max_mean_wer is not None and float(mean_wer) > float(max_mean_wer) + 1e-12: + errs.append(f"seed_tts_content_error_mean (WER)={mean_wer:.6f} > --max-seed-tts-mean-wer={max_mean_wer}") + sim_m = result.get("seed_tts_sim_mean") + if sim_m is not None and min_mean_sim is not None and float(sim_m) + 1e-12 < float(min_mean_sim): + errs.append(f"seed_tts_sim_mean={sim_m:.6f} < --min-seed-tts-mean-sim={min_mean_sim}") + ut_m = result.get("seed_tts_utmos_mean") + if ut_m is not None and min_mean_utmos is not None and float(ut_m) + 1e-12 < float(min_mean_utmos): + errs.append(f"seed_tts_utmos_mean={ut_m:.6f} < --min-seed-tts-mean-utmos={min_mean_utmos}") + return errs + + +def sync_dataset_env_from_ns(ns: argparse.Namespace) -> None: + """Mirror CLI path flags into env vars read by ``daily_omni_bench_argv`` / ``seed_tts_bench_argv``.""" + repo = getattr(ns, "daily_omni_repo", None) + if repo is not None and str(repo).strip(): + os.environ["VLLM_DAILY_OMNI_REPO"] = str(repo).strip() + if ns.daily_omni_qa_json is not None: + os.environ["VLLM_DAILY_OMNI_QA_JSON"] = str(Path(ns.daily_omni_qa_json).expanduser().resolve()) + if ns.daily_omni_video_dir is not None: + os.environ["VLLM_DAILY_OMNI_VIDEO_DIR"] = str(Path(ns.daily_omni_video_dir).expanduser().resolve()) + if ns.seed_tts_dataset_path is not None: + # ``--seed-tts-dataset-path`` accepts either a local directory or a + # Hugging Face repo id. Only resolve to an absolute filesystem path + # when the value actually exists locally; otherwise preserve the repo + # string verbatim so downstream code can pass it to snapshot_download. + raw = str(ns.seed_tts_dataset_path).strip() + p = Path(raw).expanduser() + os.environ["VLLM_SEED_TTS_DATASET_PATH"] = str(p.resolve()) if p.exists() and p.is_dir() else raw + if ns.seed_tts_root is not None: + os.environ["SEED_TTS_ROOT"] = str(Path(ns.seed_tts_root).expanduser().resolve()) + + +@contextlib.contextmanager +def _preserve_benchmark_dataset_env() -> Any: + """Save/restore dataset-related env vars so benchmark tests don't leak state.""" + keys = ( + "VLLM_DAILY_OMNI_REPO", + "VLLM_DAILY_OMNI_QA_JSON", + "VLLM_DAILY_OMNI_VIDEO_DIR", + "VLLM_SEED_TTS_DATASET_PATH", + "SEED_TTS_ROOT", + ) + original = {k: os.environ.get(k) for k in keys} + try: + yield + finally: + for key, value in original.items(): + if value is None: + os.environ.pop(key, None) + else: + os.environ[key] = value + + +def _build_common_args(ns: argparse.Namespace, *, result_filename: str) -> list[str]: + return build_serve_common_argv( + host=ns.host, + port=ns.port, + model=ns.model, + num_prompts=ns.num_prompts, + max_concurrency=ns.max_concurrency, + num_warmups=ns.num_warmups, + percentile_metrics=ns.percentile_metrics, + result_dir=ns.result_dir, + result_filename=result_filename, + ready_check_timeout_sec=ns.ready_check_timeout_sec, + ) + + +def run_daily_omni(ns: argparse.Namespace, vllm: str) -> Path: + ns.result_dir.mkdir(parents=True, exist_ok=True) + tag = datetime.now().strftime("%Y%m%d-%H%M%S") + result_filename = f"qwen_omni_acc_daily_omni_{tag}.json" + extra = json.loads(ns.daily_extra_body_json) + argv = ( + _build_common_args(ns, result_filename=result_filename) + + daily_omni_bench_argv() + + [ + "--daily-omni-input-mode", + ns.daily_omni_input_mode, + "--extra-body", + json.dumps(extra, ensure_ascii=False, separators=(",", ":")), + ] + ) + if ns.daily_omni_save_eval_items: + argv.append("--daily-omni-save-eval-items") + print("\n$", vllm, *argv, "\n", flush=True) + run_vllm_bench_subprocess(vllm, argv) + out = Path(ns.result_dir) / result_filename + if not out.is_file(): + raise FileNotFoundError(f"Expected result JSON at {out}") + return out + + +def run_seed_tts(ns: argparse.Namespace, vllm: str) -> Path: + ns.result_dir.mkdir(parents=True, exist_ok=True) + tag = datetime.now().strftime("%Y%m%d-%H%M%S") + result_filename = f"qwen_omni_acc_seed_tts_{tag}.json" + extra = json.loads(ns.seed_extra_body_json) + argv = ( + _build_common_args(ns, result_filename=result_filename) + + seed_tts_bench_argv(locale=ns.seed_tts_locale) + + [ + "--seed-tts-wer-eval", + "--extra-body", + json.dumps(extra, ensure_ascii=False, separators=(",", ":")), + ] + ) + if ns.seed_tts_wer_save_items: + argv.append("--seed-tts-wer-save-items") + if ns.seed_tts_file_ref_audio: + argv.append("--seed-tts-file-ref-audio") + extra_env: dict[str, str] = {"SEED_TTS_WER_EVAL": "1"} + if ns.seed_tts_eval_device: + extra_env["SEED_TTS_EVAL_DEVICE"] = ns.seed_tts_eval_device + print("\n$", vllm, *argv, "\n", flush=True) + run_vllm_bench_subprocess(vllm, argv, extra_env=extra_env) + out = Path(ns.result_dir) / result_filename + if not out.is_file(): + raise FileNotFoundError(f"Expected result JSON at {out}") + return out + + +def build_arg_parser() -> argparse.ArgumentParser: + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("--host", default=os.environ.get("ACC_BENCH_HOST", "127.0.0.1")) + p.add_argument("--port", type=int, default=int(os.environ.get("ACC_BENCH_PORT", "8000"))) + p.add_argument( + "--model", + default=os.environ.get( + "ACC_BENCH_MODEL", + "Qwen/Qwen3-Omni-30B-A3B-Instruct", + ), + help="Model id passed to ``vllm bench serve`` (must match the running server).", + ) + p.add_argument("--num-prompts", type=int, default=int(os.environ.get("ACC_BENCH_NUM_PROMPTS", "2000"))) + p.add_argument("--max-concurrency", type=int, default=int(os.environ.get("ACC_BENCH_MAX_CONCURRENCY", "10"))) + p.add_argument("--num-warmups", type=int, default=int(os.environ.get("ACC_BENCH_NUM_WARMUPS", "0"))) + p.add_argument( + "--percentile-metrics", + default=os.environ.get("ACC_BENCH_PERCENTILE_METRICS", "ttft,tpot,itl,e2el,audio_ttfp,audio_rtf"), + ) + p.add_argument( + "--ready-check-timeout-sec", + type=int, + default=None, + help="If set, forwarded to ``vllm bench serve`` (probe first request until success). " + "Omit to use upstream default (typically skip).", + ) + p.add_argument( + "--result-dir", + type=Path, + default=Path(os.environ.get("ACC_BENCH_RESULT_DIR", str(_default_result_dir()))), + ) + + p.add_argument("--skip-daily-omni", action="store_true") + p.add_argument("--skip-seed-tts", action="store_true") + + p.add_argument( + "--daily-omni-repo", + type=str, + default=None, + help="Hugging Face dataset id for Daily-Omni Hub mode (sets VLLM_DAILY_OMNI_REPO). " + "Ignored when local qa.json + video dir are used.", + ) + p.add_argument( + "--daily-omni-qa-json", + type=Path, + default=None, + help="Optional local qa.json; if omitted with no env, uses Hub liarliar/Daily-Omni.", + ) + p.add_argument( + "--daily-omni-video-dir", + type=Path, + default=None, + help="Optional local Videos root; if omitted, media is fetched lazily from Hub Videos.tar.", + ) + p.add_argument("--daily-omni-input-mode", choices=("all", "visual", "audio"), default="all") + p.add_argument( + "--daily-extra-body-json", + default='{"modalities":["text"]}', + help="JSON merged into each chat request for Daily-Omni (default matches common L4 / text-output runs).", + ) + p.add_argument( + "--daily-omni-save-eval-items", + action="store_true", + help="Sets env via CLI flag so per-item rows are stored in the result JSON.", + ) + p.add_argument( + "--min-daily-omni-accuracy", + type=float, + default=float((os.environ.get("ACC_BENCH_MIN_DAILY_OMNI_ACCURACY") or "0.69").strip() or "0.69"), + help="Fail when daily_omni_accuracy is strictly below this threshold (0–1). " + "Default baseline 0.69; override with env ACC_BENCH_MIN_DAILY_OMNI_ACCURACY or pass 0 to disable the floor.", + ) + + p.add_argument( + "--seed-tts-dataset-path", + type=str, + default=None, + help="Optional local root or Hub id; if omitted, uses zhaochenyang20/seed-tts-eval.", + ) + p.add_argument("--seed-tts-root", type=Path, default=None, help="Optional override for Seed-TTS filesystem root.") + p.add_argument("--seed-tts-locale", choices=("en", "zh"), default="en") + p.add_argument( + "--seed-extra-body-json", + default='{"modalities":["text","audio"]}', + help="JSON for Seed-TTS chat requests (must include audio for synthesis + PCM capture).", + ) + p.add_argument("--seed-tts-wer-save-items", action="store_true") + p.add_argument( + "--seed-tts-file-ref-audio", + action="store_true", + help="Use file:// ref_audio; server must allow local media paths.", + ) + p.add_argument( + "--seed-tts-eval-device", + default=os.environ.get("SEED_TTS_EVAL_DEVICE"), + help="Sets SEED_TTS_EVAL_DEVICE for Whisper / WavLM / UTMOS (e.g. cuda:0).", + ) + p.add_argument( + "--max-seed-tts-mean-wer", + type=float, + default=0.02, + help="If set, fail when seed_tts_content_error_mean is strictly above this value.", + ) + p.add_argument( + "--min-seed-tts-mean-sim", + type=float, + default=None, + help="If set, fail when seed_tts_sim_mean is strictly below this value.", + ) + p.add_argument( + "--min-seed-tts-mean-utmos", + type=float, + default=None, + help="If set, fail when seed_tts_utmos_mean is strictly below this value.", + ) + return p + + +def parse_acc_benchmark_args(argv: list[str] | None = None) -> argparse.Namespace: + """Parse CLI args; when ``argv`` is ``None``, use ``sys.argv[1:]`` (standalone script).""" + if argv is None: + argv = sys.argv[1:] + return build_arg_parser().parse_args(argv) + + +def run_acc_benchmark(ns: argparse.Namespace) -> int: + """Run Daily-Omni and/or Seed-TTS client benches against a running server; return 0 on success.""" + failed: list[str] = [] + + with _preserve_benchmark_dataset_env(): + sync_dataset_env_from_ns(ns) + + vllm = find_vllm_cli() + print(f"Using vLLM CLI: {vllm}", flush=True) + print(f"Repo root (for cwd reference): {_repo_root()}", flush=True) + + if not ns.skip_daily_omni: + path = run_daily_omni(ns, vllm) + print(f"\n[Daily-Omni] result JSON: {path}", flush=True) + data = load_benchmark_result(path) + errs = _validate_daily_omni(data, min_accuracy=ns.min_daily_omni_accuracy) + if errs: + failed.extend([f"[Daily-Omni] {e}" for e in errs]) + else: + print( + f"[Daily-Omni] daily_omni_accuracy={data.get('daily_omni_accuracy')} " + f"evaluated_ok={data.get('daily_omni_evaluated_ok')}", + flush=True, + ) + + if not ns.skip_seed_tts: + path = run_seed_tts(ns, vllm) + print(f"\n[Seed-TTS] result JSON: {path}", flush=True) + data = load_benchmark_result(path) + errs = _validate_seed_tts( + data, + max_mean_wer=ns.max_seed_tts_mean_wer, + min_mean_sim=ns.min_seed_tts_mean_sim, + min_mean_utmos=ns.min_seed_tts_mean_utmos, + ) + if errs: + failed.extend([f"[Seed-TTS] {e}" for e in errs]) + else: + print( + f"[Seed-TTS] mean_wer={data.get('seed_tts_content_error_mean')} " + f"mean_sim={data.get('seed_tts_sim_mean')} mean_utmos={data.get('seed_tts_utmos_mean')} " + f"evaluated={data.get('seed_tts_content_evaluated')}", + flush=True, + ) + + if failed: + print("\nACCURACY CHECK FAILED:", file=sys.stderr) + for line in failed: + print(f" - {line}", file=sys.stderr) + return 1 + + print("\nAll configured accuracy checks passed.", flush=True) + return 0 + + +def main() -> int: + return run_acc_benchmark(parse_acc_benchmark_args()) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/e2e/accuracy/qwen3_omni/test_qwen3_omni.py b/tests/e2e/accuracy/qwen3_omni/test_qwen3_omni.py new file mode 100644 index 000000000..773f7c110 --- /dev/null +++ b/tests/e2e/accuracy/qwen3_omni/test_qwen3_omni.py @@ -0,0 +1,137 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Qwen3-Omni accuracy benchmarks (Daily-Omni MCQ + Seed-TTS WER) via ``vllm bench serve --omni``. + +Starts a **module-scoped** Omni OpenAI-compatible server (same pattern as ``tests/dfx/perf`` and +``tests/e2e/online_serving/test_qwen3_omni.py``), then runs the client benches against +``omni_server.host`` / ``omni_server.port`` / ``omni_server.model``. + +**Daily-Omni from Hugging Face:** unless ``VLLM_DAILY_OMNI_QA_JSON`` and ``VLLM_DAILY_OMNI_VIDEO_DIR`` +point at a full local tree, the bench uses ``--dataset-path`` (default ``liarliar/Daily-Omni`` via +``VLLM_DAILY_OMNI_REPO`` / ``--daily-omni-repo``). QA loads through ``datasets``; ``Videos.tar`` is +downloaded and extracted under ``HF_HOME`` on demand. The tests patch in +``--daily-omni-inline-local-video`` so multimodal payloads use data URLs (no +``--allowed-local-media-path`` on the server). Use small ``--num-prompts`` defaults suitable for CI +(override with ``ACC_BENCH_NUM_PROMPTS`` / ``ACC_BENCH_MAX_CONCURRENCY``; see +:func:`tests.e2e.accuracy.qwen3_omni.qwen3_omni_acc_bench_core.build_acc_benchmark_cli_argv`). + +This package lives under ``tests/e2e/accuracy/qwen3_omni/``, so pytest still loads +``tests/e2e/accuracy/conftest.py``, which imports ``tests.conftest`` (heavy deps: ``vllm``, ``torch``, …). +A broken or partial install can therefore **fail during collection** before these tests run. + +If ``vllm`` is not on ``PATH``, the tests **skip** instead of erroring. Without +``VLLM_SKIP_ACC_BENCH=1``, a failed bench still yields a **failed** run (non-zero subprocess exit). + +Run:: + + pytest -sv tests/e2e/accuracy/qwen3_omni/test_qwen3_omni.py + +Only the subprocess accuracy marker:: + + pytest -sv tests/e2e/accuracy/qwen3_omni/test_qwen3_omni.py -m qwen3_omni_acc + +Skip when you do not have GPUs, a server, or datasets (CI opt-out):: + + VLLM_SKIP_ACC_BENCH=1 pytest -sv tests/e2e/accuracy/qwen3_omni/test_qwen3_omni.py + +Standalone CLI (expects a server already up; uses ``ACC_BENCH_*`` env defaults):: + + python tests/e2e/accuracy/qwen3_omni/run_qwen_omni_acc_benchmark.py --help +""" + +from __future__ import annotations + +from pathlib import Path + +import pytest + +from tests.e2e.accuracy.qwen3_omni import run_qwen_omni_acc_benchmark as _acc_bench +from tests.e2e.accuracy.qwen3_omni.qwen3_omni_acc_bench_core import ( + build_acc_benchmark_cli_argv, + find_vllm_cli, +) +from tests.helpers.mark import hardware_test +from tests.helpers.runtime import OmniServerParams +from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config +from vllm_omni.platforms import current_omni_platform + +_E2E_ROOT = Path(__file__).resolve().parent.parent.parent + +models = ["Qwen/Qwen3-Omni-30B-A3B-Instruct"] + +pytestmark = [pytest.mark.full_model, pytest.mark.omni] + +_CI_DEPLOY = get_deploy_config_path("ci/qwen3_omni_moe.yaml") + + +def get_chunk_config(config_path: str | None = None): + """Load the qwen3_omni CI deploy yaml with async_chunk modifications for streaming mode.""" + if config_path is None: + config_path = _CI_DEPLOY + # TODO: remove this workaround once legacy `stage_args` path is deleted. + # The pipeline (qwen3_omni/pipeline.py) already wires + # thinker2talker_async_chunk / talker2code2wav_async_chunk on stage 0/1, + # so only async_chunk needs flipping. Writing nested `engine_args:` into + # the new-schema overlay trips _parse_stage_deploy's legacy branch and + # drops flat fields (load_format, max_num_seqs, ...). + return modify_stage_config(config_path, updates={"async_chunk": True}) + + +if current_omni_platform.is_xpu(): + stage_configs = [_CI_DEPLOY] +else: # CUDA + ROCm MI325 share the same deploy config + stage_configs = [get_chunk_config()] + +test_params = [ + OmniServerParams(model=model, stage_config_path=stage_config) for model in models for stage_config in stage_configs +] + + +def _require_vllm_cli() -> None: + try: + find_vllm_cli() + except FileNotFoundError as exc: + pytest.skip(str(exc)) + + +@pytest.fixture(autouse=True) +def _daily_omni_hub_inline_media(monkeypatch: pytest.MonkeyPatch) -> None: + """Hub / lazy-cache mode uses local files → default ``file://`` needs server allowlist. + + ``run_qwen_omni_acc_benchmark`` binds ``daily_omni_bench_argv`` at import time; patch that copy + so we append ``--daily-omni-inline-local-video`` whenever the core helper did not already set it + (local qa.json + video-dir mode already passes the flag). + """ + orig = _acc_bench.daily_omni_bench_argv + + def _wrapped() -> list[str]: + out = list(orig()) + if "--daily-omni-inline-local-video" not in out: + out.append("--daily-omni-inline-local-video") + return out + + monkeypatch.setattr(_acc_bench, "daily_omni_bench_argv", _wrapped) + monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn") + monkeypatch.setenv("VLLM_TEST_CLEAN_GPU_MEMORY", "0") + + +@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) +@pytest.mark.parametrize("omni_server", test_params, indirect=True) +def test_qwen3_omni_daily_omni_accuracy_bench(omni_server) -> None: + _require_vllm_cli() + pytest.importorskip("datasets") + pytest.importorskip("huggingface_hub") + ns = _acc_bench.parse_acc_benchmark_args( + build_acc_benchmark_cli_argv(omni_server, skip_seed=True, skip_daily=False) + ) + assert _acc_bench.run_acc_benchmark(ns) == 0 + + +@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) +@pytest.mark.parametrize("omni_server", test_params, indirect=True) +def test_qwen3_omni_seed_tts_wer_bench(omni_server) -> None: + _require_vllm_cli() + pytest.importorskip("huggingface_hub") + ns = _acc_bench.parse_acc_benchmark_args( + build_acc_benchmark_cli_argv(omni_server, skip_seed=False, skip_daily=True) + ) + assert _acc_bench.run_acc_benchmark(ns) == 0 diff --git a/vllm_omni/benchmarks/data_modules/daily_omni_dataset.py b/vllm_omni/benchmarks/data_modules/daily_omni_dataset.py index 01b86d0fd..65918414f 100644 --- a/vllm_omni/benchmarks/data_modules/daily_omni_dataset.py +++ b/vllm_omni/benchmarks/data_modules/daily_omni_dataset.py @@ -9,7 +9,9 @@ Supports loading QA metadata from: - Local JSON file (``qa_json_path``): recommended for offline/air-gapped environments - HuggingFace datasets (``dataset_path``): legacy online mode -The videos must be separately downloaded and extracted from Videos.tar. +Video/audio files normally come from extracted ``Videos.tar``. When ``--daily-omni-video-dir`` +is not set, the first request that needs on-disk media downloads that archive from the Hugging Face +dataset repo (``huggingface_hub``) and caches it under ``HF_HOME``. Why ``BenchmarkDataset`` instead of ``HuggingFaceDataset``? vLLM's ``HuggingFaceDataset`` is a thin wrapper whose ``__init__`` always ends by calling @@ -52,6 +54,9 @@ Usage: import base64 import json import logging +import os +import shutil +import tarfile from dataclasses import dataclass from pathlib import Path from typing import Any, Literal @@ -73,6 +78,103 @@ except ImportError: logger = logging.getLogger(__name__) +def _daily_omni_hf_cache_root() -> Path: + return Path(os.environ.get("HF_HOME", Path.home() / ".cache" / "huggingface")).expanduser().resolve() + + +def _daily_omni_tar_fingerprint(tar_path: Path) -> str: + st = tar_path.stat() + return f"v1:{st.st_size}:{int(st.st_mtime_ns)}" + + +def _daily_omni_find_videos_root_in_extract(tmp: Path) -> Path: + """Return directory whose children are ``video_id`` folders with ``*_video.mp4``.""" + videos = tmp / "Videos" + if videos.is_dir(): + return videos + for child in sorted(tmp.iterdir()): + if child.is_dir() and not child.name.startswith("."): + probe = child / f"{child.name}_video.mp4" + if probe.is_file(): + return tmp + raise RuntimeError( + f"Unrecognized layout after extracting Daily-Omni Videos.tar under {tmp} " + "(expected top-level 'Videos/' or per-video_id subdirs)." + ) + + +def ensure_daily_omni_hub_videos_dir(repo_id: str) -> Path: + """Download ``Videos.tar`` from the Hugging Face dataset repo and return the ``Videos`` root. + + The returned path matches ``--daily-omni-video-dir`` (directory containing ``{{video_id}}/``). + + Cached under ``HF_HOME`` / ``vllm_omni/daily_omni_media/``. Reuses extraction when the + tarball fingerprint matches. + + Raises: + ImportError: if ``huggingface_hub`` is not installed. + FileNotFoundError / RuntimeError: if the archive is missing or malformed. + """ + rid = (repo_id or "").strip() + if not rid: + raise ValueError("repo_id is required to download Daily-Omni Videos.tar") + + try: + from huggingface_hub import hf_hub_download + except ImportError as e: + raise ImportError( + "Daily-Omni Hub media download requires huggingface_hub. " + "Install it (e.g. with vLLM) or provide --daily-omni-video-dir with a local extract." + ) from e + + safe = rid.replace("/", "__").replace("\\", "_") + staging_root = _daily_omni_hf_cache_root() / "vllm_omni" / "daily_omni_media" / safe + videos_dir = staging_root / "Videos" + marker = staging_root / ".videos_extracted" + + tar_path: Path | None = None + for fname in ("Videos.tar", "videos.tar"): + try: + tar_path = Path(hf_hub_download(repo_id=rid, filename=fname, repo_type="dataset")) + break + except Exception: + continue + if tar_path is None or not tar_path.is_file(): + raise FileNotFoundError( + f"Could not download Videos.tar from Hugging Face dataset {rid!r} (tried Videos.tar / videos.tar)." + ) + + fp = _daily_omni_tar_fingerprint(tar_path) + if marker.is_file() and videos_dir.is_dir(): + try: + if marker.read_text(encoding="utf-8").strip() == fp: + next(videos_dir.iterdir()) + logger.info("Reusing cached Daily-Omni Videos extract at %s", videos_dir) + return videos_dir + except (OSError, StopIteration): + shutil.rmtree(videos_dir, ignore_errors=True) + marker.unlink(missing_ok=True) + + staging_root.mkdir(parents=True, exist_ok=True) + work = staging_root / "_extract_work" + shutil.rmtree(work, ignore_errors=True) + work.mkdir(parents=True) + try: + logger.info("Extracting Daily-Omni Videos.tar from %s (repo=%s)", tar_path, rid) + with tarfile.open(tar_path, "r:*") as tf: + tf.extractall(path=work, filter="data") + found = _daily_omni_find_videos_root_in_extract(work) + if videos_dir.exists(): + shutil.rmtree(videos_dir, ignore_errors=True) + shutil.move(str(found), str(videos_dir)) + finally: + shutil.rmtree(work, ignore_errors=True) + + marker.write_text(fp, encoding="utf-8") + logger.info("Daily-Omni Hub media ready at %s", videos_dir) + return videos_dir + + class _ListDatasetIterator: """Simple iterator wrapper around a list to mimic HuggingFace streaming dataset behavior.""" @@ -142,7 +244,10 @@ class DailyOmniDataset(BenchmarkDataset): - Local JSON file (``qa_json_path``): recommended for offline/air-gapped environments - HuggingFace datasets (``dataset_path``): legacy online mode - The videos must be separately downloaded and extracted from Videos.tar. + Video/audio files normally come from extracted ``Videos.tar``. When ``video_dir`` is not set, + the first sample that needs on-disk media downloads that archive from the Hugging Face dataset + repo (env ``VLLM_DAILY_OMNI_MEDIA_REPO`` overrides the repo id; else ``dataset_path`` or + :data:`DEFAULT_HF_DATASET_ID`). Args: qa_json_path: Path to local qa.json file (offline mode, preferred). When provided, @@ -151,7 +256,8 @@ class DailyOmniDataset(BenchmarkDataset): ``qa_json_path`` is not provided (legacy online mode). dataset_split: Dataset split to use (default: "train"). Used only in online mode. random_seed: Random seed for shuffling - video_dir: Directory containing extracted video files (default: None) + video_dir: Directory containing extracted video files (default: None; may be filled lazily + from Hub — see above). input_mode: Which modalities to send, matching upstream Daily-Omni ``--input_mode``: ``all`` — video + WAV (default; official audio-visual protocol); ``visual`` — video only; @@ -308,13 +414,14 @@ class DailyOmniDataset(BenchmarkDataset): "Install with: pip install datasets, or use local JSON mode instead." ) - ds = load_dataset( - self.dataset_path, - name=self.dataset_subset, - split=self.dataset_split, - streaming=self._hf_streaming, - trust_remote_code=self.trust_remote_code, - ) + load_kw: dict[str, Any] = { + "split": self.dataset_split, + "streaming": self._hf_streaming, + "trust_remote_code": self.trust_remote_code, + } + if self.dataset_subset is not None: + load_kw["name"] = self.dataset_subset + ds = load_dataset(self.dataset_path, **load_kw) if not getattr(self, "disable_shuffle", False): ds = ds.shuffle(seed=self.random_seed) self.data = ds @@ -569,6 +676,17 @@ class DailyOmniDataset(BenchmarkDataset): "audio_url": {"url": path.as_uri()}, } + def _lazy_ensure_hub_media_dir(self) -> None: + """If ``video_dir`` was not configured, download and extract ``Videos.tar`` once from HF.""" + if self.video_dir is not None: + return + repo = os.environ.get("VLLM_DAILY_OMNI_MEDIA_REPO", "").strip() + if not repo: + repo = (self.dataset_path or "").strip() + if not repo: + repo = self.DEFAULT_HF_DATASET_ID + self.video_dir = ensure_daily_omni_hub_videos_dir(repo) + def _get_video_content( self, video_id: str, @@ -589,6 +707,8 @@ class DailyOmniDataset(BenchmarkDataset): url = f"https://{url.lstrip('/')}" return {"type": "video_url", "video_url": {"url": url}} + self._lazy_ensure_hub_media_dir() + if self.video_dir and video_id: video_path = self._resolve_local_video_path(video_id) if video_path is not None: @@ -615,7 +735,13 @@ class DailyOmniDataset(BenchmarkDataset): return None def _get_audio_content(self, video_id: str) -> dict[str, Any] | None: - """Resolve extracted WAV for OpenAI-style ``audio_url`` (local files only).""" + """Resolve extracted WAV for OpenAI-style ``audio_url`` (local files under ``video_dir``). + + Uses the same tree as video (``{video_id}/{video_id}_audio.wav``), including after lazy + Hub ``Videos.tar`` extraction when ``video_dir`` was unset. + """ + self._lazy_ensure_hub_media_dir() + if not self.video_dir or not video_id: logger.warning( "Daily-Omni input_mode %r requires --daily-omni-video-dir with %s", diff --git a/vllm_omni/benchmarks/data_modules/daily_omni_eval.py b/vllm_omni/benchmarks/data_modules/daily_omni_eval.py index ecc9edc84..f191cf2fe 100644 --- a/vllm_omni/benchmarks/data_modules/daily_omni_eval.py +++ b/vllm_omni/benchmarks/data_modules/daily_omni_eval.py @@ -5,7 +5,8 @@ Compares model ``generated_text`` to dataset ``Answer`` (A/B/C/D). **Alignment with open-source** (`Lliar-liar/Daily-Omni` ``test_model/.../testmodel.py``): - Answer extraction defaults to the same rules as ``extract_choice_letter`` (strip after an - ``assistant`` marker, then leading ``A``–``D``, else first ``\\b[A-D]\\b``). Set env + ``assistant`` marker, then leading ``A``–``D``, else ``\\b[A-D]\\b``, else a CJK-safe + non-letter-boundary pass). Set env ``DAILY_OMNI_EXTRACT_MODE=relaxed`` to use the older vLLM-Omni heuristics (last ``answer:``, tail scan, etc.). - Overall accuracy comparable to the official script uses **successful HTTP responses only** as @@ -52,6 +53,16 @@ def extract_choice_letter_official(text: str | None) -> str | None: fallback = re.search(r"\b([A-D])\b", candidate.upper()) if fallback: return fallback.group(1) + # ``\b`` is ASCII/Latin-word-centric; CJK (e.g. "选B", "答案：B") has no boundary before B. + loose = list( + re.finditer( + r"(?:[^A-Za-z]|^)([A-D])(?:[^A-Za-z]|$)", + candidate, + flags=re.IGNORECASE, + ) + ) + if loose: + return loose[-1].group(1).upper() return None diff --git a/vllm_omni/benchmarks/data_modules/seed_tts_dataset.py b/vllm_omni/benchmarks/data_modules/seed_tts_dataset.py index ca6de4cb2..4f072e1a2 100644 --- a/vllm_omni/benchmarks/data_modules/seed_tts_dataset.py +++ b/vllm_omni/benchmarks/data_modules/seed_tts_dataset.py @@ -94,7 +94,7 @@ def _load_meta_rows(meta_file: Path) -> list[_SeedTTSRow]: return rows -def resolve_seed_tts_root(dataset_path: str | None, *, explicit_root: str | None) -> Path: +def resolve_seed_tts_root(dataset_path: str | None, *, explicit_root: str | None, locale: str = "en") -> Path: """Return directory containing ``{locale}/meta.lst`` and ``{locale}/prompt-wavs/``.""" if explicit_root: root = Path(explicit_root).expanduser().resolve() @@ -117,7 +117,14 @@ def resolve_seed_tts_root(dataset_path: str | None, *, explicit_root: str | None "Install huggingface_hub to download Seed-TTS from the Hub, or clone the dataset " "locally and pass --dataset-path / --seed-tts-root to that directory." ) from e - cache = snapshot_download(repo_id=repo_id, repo_type="dataset") + # Download only the requested locale subtree instead of the whole dataset + # repo. This avoids large, flaky nightly downloads when we only need e.g. + # ``en/meta.lst`` + ``en/prompt-wavs/**``. + cache = snapshot_download( + repo_id=repo_id, + repo_type="dataset", + allow_patterns=[f"{locale}/**"], + ) return Path(cache).resolve() @@ -172,7 +179,7 @@ class SeedTTSDataset(BenchmarkDataset): disable_shuffle=disable_shuffle, **kwargs, ) - self._root = resolve_seed_tts_root(self.dataset_path, explicit_root=self._explicit_root) + self._root = resolve_seed_tts_root(self.dataset_path, explicit_root=self._explicit_root, locale=self.locale) self._rows: list[_SeedTTSRow] = [] self.load_data() diff --git a/vllm_omni/benchmarks/data_modules/seed_tts_eval.py b/vllm_omni/benchmarks/data_modules/seed_tts_eval.py index d5f1b6470..d8c37af13 100644 --- a/vllm_omni/benchmarks/data_modules/seed_tts_eval.py +++ b/vllm_omni/benchmarks/data_modules/seed_tts_eval.py @@ -1,7 +1,7 @@ """Seed-TTS WER aligned with Bytedance ``seed-tts-eval`` / ``run_wer.py``. Matches the published protocol (see Hugging Face dataset card and -https://github.com/BytedanceSpeech/seed-tts-eval): +https://github.com/zhaochenyang20/seed-tts-eval): - **EN**: ``openai/whisper-large-v3`` via ``transformers``, audio resampled to **16 kHz** (same as ``run_wer.py``). @@ -489,7 +489,7 @@ def compute_seed_tts_wer_metrics( sim_failed = 0 sim_skipped_no_ref = 0 utmos_failed = 0 - utmos_on = _eval_submetric_enabled("SEED_TTS_UTMOS_EVAL", default=True) + utmos_on = _eval_submetric_enabled("SEED_TTS_UTMOS_EVAL", default=False) for req, out in zip(input_requests, outputs, strict=True): assert isinstance(req, SeedTTSSampleRequest) @@ -622,7 +622,7 @@ def compute_seed_tts_wer_metrics( errs.append(wer) sim_v: float | None = None - if _eval_submetric_enabled("SEED_TTS_SIM_EVAL", default=True): + if _eval_submetric_enabled("SEED_TTS_SIM_EVAL", default=False): ref_path = getattr(req, "seed_tts_ref_wav_path", "") or "" if ref_path and os.path.isfile(ref_path): try: diff --git a/vllm_omni/entrypoints/cli/benchmark/serve.py b/vllm_omni/entrypoints/cli/benchmark/serve.py index d281432e5..44d38303f 100644 --- a/vllm_omni/entrypoints/cli/benchmark/serve.py +++ b/vllm_omni/entrypoints/cli/benchmark/serve.py @@ -33,6 +33,8 @@ def add_daily_omni_cli_args(parser: argparse.ArgumentParser) -> None: default=None, help="Root directory of extracted Daily-Omni videos (contents of Videos.tar: " "each video_id in its own subdir with {video_id}_video.mp4). " + "If omitted, Videos.tar is downloaded from the Hugging Face dataset repo on first multimodal " + "request. " "When using file URLs, you MUST start the vLLM server with " "--allowed-local-media-path set to this same directory (or a parent), " "otherwise requests fail with 'Cannot load local files without " @@ -116,7 +118,7 @@ def add_seed_tts_cli_args(parser: argparse.ArgumentParser) -> None: default=False, help="Keep synthesized audio as 24 kHz mono PCM for WER (works with " "--backend openai-audio-speech or openai-chat-omni). Scoring follows " - "BytedanceSpeech/seed-tts-eval (Whisper-large-v3 / Paraformer-zh + jiwer). " + "zhaochenyang20/seed-tts-eval (Whisper-large-v3 / Paraformer-zh + jiwer). " "Sets SEED_TTS_WER_EVAL=1. Install: pip install 'vllm-omni[seed-tts-eval]'. " "Optional: SEED_TTS_EVAL_DEVICE, SEED_TTS_HF_WHISPER_MODEL.", ) From 8384dead5624d757f90bbaef2727cbfe67e15859 Mon Sep 17 00:00:00 2001 From: R0CKSTAR Date: Thu, 23 Apr 2026 11:07:02 +0800 Subject: [PATCH 18/45] [MUSA][Feat] Upgrade MATE to match CUDA/ROCm behavior on FA (#2766) Signed-off-by: Xiaodong Ye --- requirements/musa.txt | 4 +++- vllm_omni/diffusion/attention/backends/flash_attn.py | 12 +----------- vllm_omni/diffusion/attention/backends/utils/fa.py | 6 +++--- 3 files changed, 7 insertions(+), 15 deletions(-) diff --git a/requirements/musa.txt b/requirements/musa.txt index 112f32604..c100c70cf 100644 --- a/requirements/musa.txt +++ b/requirements/musa.txt @@ -1,4 +1,6 @@ -r common.txt # MUSA platform dependencies -torchada>=0.1.46 +torchada>=0.1.50 onnxruntime>=1.23.2 +mate>=0.2.0 +flash_attn_3>=0.1.4 diff --git a/vllm_omni/diffusion/attention/backends/flash_attn.py b/vllm_omni/diffusion/attention/backends/flash_attn.py index b6ab3a57a..d38ea4f6e 100644 --- a/vllm_omni/diffusion/attention/backends/flash_attn.py +++ b/vllm_omni/diffusion/attention/backends/flash_attn.py @@ -96,7 +96,7 @@ class FlashAttentionImpl(AttentionImpl): value: torch.Tensor, attn_metadata: AttentionMetadata = None, ) -> torch.Tensor: - """CUDA/ROCm flash attention implementation.""" + """CUDA/ROCm/MUSA flash attention implementation.""" from vllm_omni.diffusion.attention.backends.utils.fa import ( HAS_FLASH_ATTN, flash_attn_func, @@ -209,13 +209,3 @@ class FlashAttentionImpl(AttentionImpl): layout="BNSD", ) return output - - def forward_musa( - self, - query: torch.Tensor, - key: torch.Tensor, - value: torch.Tensor, - attn_metadata: AttentionMetadata = None, - ) -> torch.Tensor: - # XXX (MUSA): MUSA uses the same implementation as XPU (mate only provides flash_attn_varlen_func) - return self.forward_xpu(query, key, value, attn_metadata) diff --git a/vllm_omni/diffusion/attention/backends/utils/fa.py b/vllm_omni/diffusion/attention/backends/utils/fa.py index fe6051f8b..18886871a 100644 --- a/vllm_omni/diffusion/attention/backends/utils/fa.py +++ b/vllm_omni/diffusion/attention/backends/utils/fa.py @@ -41,7 +41,7 @@ elif current_omni_platform.is_xpu(): pass elif current_omni_platform.is_musa(): try: - from mate import flash_attn_varlen_func # noqa: F401 + from flash_attn_interface import flash_attn_func, flash_attn_varlen_func # noqa: F401 except (ImportError, ModuleNotFoundError): pass else: @@ -82,8 +82,8 @@ HAS_FLASH_ATTN = flash_attn_func is not None or flash_attn_varlen_func is not No @lru_cache(maxsize=1) def is_mate_available() -> bool: - """Check if MATE (MUSA Flash Attention) is available.""" - return current_omni_platform.is_musa() and flash_attn_varlen_func is not None + """Check if MATE (MUSA AI Tensor Engine) is available.""" + return current_omni_platform.is_musa() and flash_attn_func is not None or flash_attn_varlen_func is not None def _index_first_axis(tensor, indices): From d4cbdff9cdbe7536f6062a8f9daa07fa6552b899 Mon Sep 17 00:00:00 2001 From: WeiQing Chen <40507679+david6666666@users.noreply.github.com> Date: Thu, 23 Apr 2026 11:30:47 +0800 Subject: [PATCH 19/45] [Doc] Add diffusion attention backend docs (#3011) Signed-off-by: david6666666 <530634352@qq.com> --- docs/.nav.yml | 1 + .../diffusion/attention_backends.md | 120 ++++++++++++++++++ .../offline_inference/text_to_video.md | 2 + 3 files changed, 123 insertions(+) create mode 100644 docs/user_guide/diffusion/attention_backends.md diff --git a/docs/.nav.yml b/docs/.nav.yml index aebd560ca..8b0f9a023 100644 --- a/docs/.nav.yml +++ b/docs/.nav.yml @@ -65,6 +65,7 @@ nav: - FP8: user_guide/diffusion/quantization/fp8.md - Int8: user_guide/diffusion/quantization/int8.md - GGUF: user_guide/diffusion/quantization/gguf.md + - Attention Backends: user_guide/diffusion/attention_backends.md - Frame Interpolation: user_guide/diffusion/frame_interpolation.md - Parallelism: - Overview: user_guide/diffusion/parallelism/overview.md diff --git a/docs/user_guide/diffusion/attention_backends.md b/docs/user_guide/diffusion/attention_backends.md new file mode 100644 index 000000000..692bcc0f9 --- /dev/null +++ b/docs/user_guide/diffusion/attention_backends.md @@ -0,0 +1,120 @@ +# Diffusion Attention Backends + +This document describes the diffusion attention backends available in vLLM-Omni, how to select them, and how to use SageAttention. + +## Overview + +Diffusion attention backend selection is controlled by the `DIFFUSION_ATTENTION_BACKEND` environment variable and resolved in `vllm_omni.diffusion.attention.selector`. + +This backend is used by diffusion attention layers such as the DiT attention in video and image generation models. + +On CUDA, the practical choices today are: + +- `FLASH_ATTN`: FlashAttention backend. This is the default on supported CUDA systems when FlashAttention is installed. +- `TORCH_SDPA`: PyTorch `scaled_dot_product_attention`. +- `SAGE_ATTN`: SageAttention backend, if `sageattention` is installed. + +If `DIFFUSION_ATTENTION_BACKEND` is unset, vLLM-Omni asks the current platform to choose the default backend. On CUDA, that normally means `FLASH_ATTN` when available, otherwise `TORCH_SDPA`. + +## Backend Options + +| Value | Notes | +|---|---| +| `FLASH_ATTN` | Default on CUDA when FlashAttention is available. Good default for most diffusion workloads. | +| `TORCH_SDPA` | Most conservative fallback. Useful for debugging or compatibility. | +| `SAGE_ATTN` | Requires `sageattention`. Can improve performance on some workloads, but output quality must be validated model-by-model. | + +## Selection Priority + +Diffusion attention backend selection follows this order: + +1. `DIFFUSION_ATTENTION_BACKEND` +2. Platform default + +Example: + +```bash +export DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN +``` + +## SageAttention Installation + +vLLM-Omni expects SageAttention to be installed into the same Python environment as vLLM-Omni. + +Build from source: + +```bash +git clone https://github.com/thu-ml/SageAttention.git +cd SageAttention + +export EXT_PARALLEL=4 NVCC_APPEND_FLAGS="--threads 8" MAX_JOBS=32 +pip install . --no-build-isolation +``` + +Quick check: + +```bash +python -c "import sageattention; print(sageattention.__file__)" +``` + +## Usage + +### Enable SageAttention + +Example: HunyuanVideo-1.5 text-to-video + +```bash +DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN python examples/offline_inference/text_to_video/text_to_video.py \ + --model hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v \ + --prompt "A dog running across a field of golden wheat." \ + --height 480 --width 832 --num-frames 33 \ + --num-inference-steps 30 --seed 42 --guidance-scale 6.0 \ + --tensor-parallel-size 2 \ + --output ../tmp/hv15_modelopt_sage.mp4 +``` + +Example: Wan2.2 TI2V 5B + +```bash +DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN python examples/offline_inference/text_to_video/text_to_video.py \ + --model Wan-AI/Wan2.2-TI2V-5B-Diffusers \ + --prompt "A dog running across a field of golden wheat." \ + --height 704 --width 1280 --num-frames 49 \ + --num-inference-steps 30 --seed 42 --guidance-scale 5.0 \ + --tensor-parallel-size 2 \ + --output outputs/wan22_sage.mp4 +``` + +### Compare Against FlashAttention + +Unset the backend override, or explicitly use `FLASH_ATTN`: + +```bash +python examples/offline_inference/text_to_video/text_to_video.py \ + --model Wan-AI/Wan2.2-TI2V-5B-Diffusers \ + --prompt "A dog running across a field of golden wheat." \ + --height 704 --width 1280 --num-frames 49 \ + --num-inference-steps 30 --seed 42 --guidance-scale 5.0 \ + --tensor-parallel-size 2 \ + --output outputs/wan22_fa3.mp4 +``` + +## Validation Guidance + +Do not assume that a faster attention backend is numerically interchangeable with `FLASH_ATTN`. + +Always compare: + +- End-to-end runtime +- DiT / diffusion stage runtime +- Output quality against a known-good baseline + +At minimum, keep the same: + +- model +- prompt +- seed +- resolution +- frame count +- inference steps +- parallel config diff --git a/docs/user_guide/examples/offline_inference/text_to_video.md b/docs/user_guide/examples/offline_inference/text_to_video.md index 4288c089c..a09dbfc97 100644 --- a/docs/user_guide/examples/offline_inference/text_to_video.md +++ b/docs/user_guide/examples/offline_inference/text_to_video.md @@ -5,6 +5,8 @@ Source Date: Thu, 23 Apr 2026 11:31:13 +0800 Subject: [PATCH 20/45] [Bugfix] Fix default sampling params for /v1/videos on main (#3049) Signed-off-by: david6666666 <530634352@qq.com> --- .../openai_api/test_video_server.py | 116 ++++++++++++++++++ vllm_omni/entrypoints/openai/api_server.py | 6 +- vllm_omni/entrypoints/openai/serving_video.py | 41 +++++-- 3 files changed, 148 insertions(+), 15 deletions(-) diff --git a/tests/entrypoints/openai_api/test_video_server.py b/tests/entrypoints/openai_api/test_video_server.py index 1cc561665..a29f4493c 100644 --- a/tests/entrypoints/openai_api/test_video_server.py +++ b/tests/entrypoints/openai_api/test_video_server.py @@ -29,6 +29,7 @@ from vllm_omni.entrypoints.openai.protocol.videos import ( from vllm_omni.entrypoints.openai.serving_video import OmniOpenAIServingVideo from vllm_omni.entrypoints.openai.storage import LocalStorageManager from vllm_omni.entrypoints.openai.stores import AsyncDictStore, TaskRegistry +from vllm_omni.inputs.data import OmniDiffusionSamplingParams pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -60,6 +61,7 @@ class MockVideoResult: class FakeAsyncOmni: def __init__(self): self.stage_configs = [SimpleNamespace(stage_type="diffusion")] + self.default_sampling_params_list = [OmniDiffusionSamplingParams()] self.captured_prompt = None self.captured_sampling_params_list = None @@ -442,6 +444,87 @@ def test_frame_interpolation_params_pass_to_diffusion_sampling_params(test_clien assert captured.frame_interpolation_model_path == "local-rife" +def test_default_sampling_params_apply_to_video_requests(test_client, mocker: MockerFixture): + mocker.patch( + "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes", + return_value=b"fake-video", + ) + engine = test_client.app.state.openai_serving_video._engine_client + engine.default_sampling_params_list = [ + OmniDiffusionSamplingParams( + num_inference_steps=4, + guidance_scale=7.5, + generator_device="cpu", + enable_frame_interpolation=True, + frame_interpolation_exp=2, + frame_interpolation_scale=0.5, + frame_interpolation_model_path="default-rife", + ) + ] + + response = test_client.post( + "/v1/videos", + data={ + "prompt": "default param pass-through", + }, + ) + + assert response.status_code == 200 + video_id = response.json()["id"] + _wait_for_status(test_client, video_id, VideoGenerationStatus.COMPLETED.value) + + captured = engine.captured_sampling_params_list[0] + assert captured.num_inference_steps == 4 + assert captured.guidance_scale == 7.5 + assert captured.generator_device == "cpu" + assert captured.enable_frame_interpolation is True + assert captured.frame_interpolation_exp == 2 + assert captured.frame_interpolation_scale == 0.5 + assert captured.frame_interpolation_model_path == "default-rife" + + +def test_request_params_override_default_video_sampling_params(test_client, mocker: MockerFixture): + mocker.patch( + "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes", + return_value=b"fake-video", + ) + engine = test_client.app.state.openai_serving_video._engine_client + engine.default_sampling_params_list = [ + OmniDiffusionSamplingParams( + num_inference_steps=4, + guidance_scale=7.5, + enable_frame_interpolation=True, + frame_interpolation_exp=2, + frame_interpolation_scale=0.5, + frame_interpolation_model_path="default-rife", + ) + ] + + response = test_client.post( + "/v1/videos", + data={ + "prompt": "explicit override", + "num_inference_steps": "8", + "enable_frame_interpolation": "false", + "frame_interpolation_exp": "1", + "frame_interpolation_scale": "1.0", + "frame_interpolation_model_path": "custom-rife", + }, + ) + + assert response.status_code == 200 + video_id = response.json()["id"] + _wait_for_status(test_client, video_id, VideoGenerationStatus.COMPLETED.value) + + captured = engine.captured_sampling_params_list[0] + assert captured.num_inference_steps == 8 + assert captured.guidance_scale == 7.5 + assert captured.enable_frame_interpolation is False + assert captured.frame_interpolation_exp == 1 + assert captured.frame_interpolation_scale == 1.0 + assert captured.frame_interpolation_model_path == "custom-rife" + + def test_worker_fps_multiplier_is_applied_to_async_encoding(test_client, mocker: MockerFixture): fps_values = [] engine = test_client.app.state.openai_serving_video._engine_client @@ -1149,6 +1232,39 @@ def test_sync_frame_interpolation_params_pass_to_sampling_params(test_client, mo assert kwargs["fps"] == 8 +def test_sync_default_sampling_params_apply_to_video_requests(test_client, mocker: MockerFixture): + _mock_encode_video_bytes(mocker) + engine = test_client.app.state.openai_serving_video._engine_client + engine.default_sampling_params_list = [ + OmniDiffusionSamplingParams( + num_inference_steps=4, + guidance_scale=7.5, + enable_frame_interpolation=True, + frame_interpolation_exp=2, + frame_interpolation_scale=0.5, + frame_interpolation_model_path="default-rife", + ) + ] + + response = test_client.post( + "/v1/videos/sync", + data={ + "prompt": "sync default param pass-through", + "fps": "8", + }, + ) + + assert response.status_code == 200 + engine = test_client.app.state.openai_serving_video._engine_client + captured = engine.captured_sampling_params_list[0] + assert captured.num_inference_steps == 4 + assert captured.guidance_scale == 7.5 + assert captured.enable_frame_interpolation is True + assert captured.frame_interpolation_exp == 2 + assert captured.frame_interpolation_scale == 0.5 + assert captured.frame_interpolation_model_path == "default-rife" + + def test_worker_fps_multiplier_is_applied_to_sync_encoding(test_client, mocker: MockerFixture): engine = test_client.app.state.openai_serving_video._engine_client fps_values = [] diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index 4454f5bda..4d348a0d8 100644 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/vllm_omni/entrypoints/openai/api_server.py @@ -2361,9 +2361,9 @@ async def _parse_video_form( true_cfg_scale: float | None = Form(default=None), seed: int | None = Form(default=None), negative_prompt: str | None = Form(default=None), - enable_frame_interpolation: bool = Form(default=False), - frame_interpolation_exp: int = Form(default=1, ge=1), - frame_interpolation_scale: float = Form(default=1.0, gt=0.0), + enable_frame_interpolation: bool | None = Form(default=None), + frame_interpolation_exp: int | None = Form(default=None, ge=1), + frame_interpolation_scale: float | None = Form(default=None, gt=0.0), frame_interpolation_model_path: str | None = Form(default=None), lora: str | None = Form(default=None), extra_params: str | None = Form(default=None), diff --git a/vllm_omni/entrypoints/openai/serving_video.py b/vllm_omni/entrypoints/openai/serving_video.py index 741295c7c..a4be330eb 100644 --- a/vllm_omni/entrypoints/openai/serving_video.py +++ b/vllm_omni/entrypoints/openai/serving_video.py @@ -3,6 +3,7 @@ from __future__ import annotations +import copy import time from dataclasses import dataclass from http import HTTPStatus @@ -95,7 +96,7 @@ class OmniOpenAIServingVideo: if request.negative_prompt is not None: prompt["negative_prompt"] = request.negative_prompt - gen_params = OmniDiffusionSamplingParams() + gen_params = self._resolve_default_sampling_params() input_image = None if reference_image is None else reference_image.data vp = request.resolve_video_params() @@ -113,22 +114,27 @@ class OmniOpenAIServingVideo: if vp.fps is not None: gen_params.fps = vp.fps gen_params.frame_rate = float(vp.fps) - gen_params.enable_frame_interpolation = request.enable_frame_interpolation - gen_params.frame_interpolation_exp = request.frame_interpolation_exp - gen_params.frame_interpolation_scale = request.frame_interpolation_scale - gen_params.frame_interpolation_model_path = request.frame_interpolation_model_path + provided_fields = request.model_fields_set + if "enable_frame_interpolation" in provided_fields: + gen_params.enable_frame_interpolation = request.enable_frame_interpolation + if "frame_interpolation_exp" in provided_fields: + gen_params.frame_interpolation_exp = request.frame_interpolation_exp + if "frame_interpolation_scale" in provided_fields: + gen_params.frame_interpolation_scale = request.frame_interpolation_scale + if "frame_interpolation_model_path" in provided_fields: + gen_params.frame_interpolation_model_path = request.frame_interpolation_model_path - if request.num_inference_steps is not None: + if "num_inference_steps" in provided_fields and request.num_inference_steps is not None: gen_params.num_inference_steps = request.num_inference_steps - if request.guidance_scale is not None: + if "guidance_scale" in provided_fields and request.guidance_scale is not None: gen_params.guidance_scale = request.guidance_scale - if request.guidance_scale_2 is not None: + if "guidance_scale_2" in provided_fields and request.guidance_scale_2 is not None: gen_params.guidance_scale_2 = request.guidance_scale_2 - if request.true_cfg_scale is not None: + if "true_cfg_scale" in provided_fields and request.true_cfg_scale is not None: gen_params.true_cfg_scale = request.true_cfg_scale - if request.seed is not None: + if "seed" in provided_fields and request.seed is not None: gen_params.seed = request.seed - if request.boundary_ratio is not None: + if "boundary_ratio" in provided_fields and request.boundary_ratio is not None: gen_params.boundary_ratio = request.boundary_ratio logger.info( @@ -136,7 +142,7 @@ class OmniOpenAIServingVideo: request.boundary_ratio, gen_params.boundary_ratio, ) - if request.flow_shift is not None: + if "flow_shift" in provided_fields and request.flow_shift is not None: gen_params.extra_args["flow_shift"] = request.flow_shift # Apply model-specific extra parameters @@ -263,6 +269,17 @@ class OmniOpenAIServingVideo: return int(multiplier) return 1 + def _resolve_default_sampling_params(self) -> OmniDiffusionSamplingParams: + default_sampling_params_list = getattr(self._engine_client, "default_sampling_params_list", None) + if default_sampling_params_list: + for params in default_sampling_params_list: + if isinstance(params, OmniDiffusionSamplingParams): + # Requests mutate sampling params in-place, including + # nested dict fields like extra_args. Deep-copy the stage + # defaults so one request cannot leak state into another. + return copy.deepcopy(params) + return OmniDiffusionSamplingParams() + @staticmethod def _apply_lora(lora_body: Any, gen_params: OmniDiffusionSamplingParams) -> None: try: From 3a32cd6780f74049325bc0baabd13cd8b2b47bec Mon Sep 17 00:00:00 2001 From: Yuanheng Zhao <54058983+yuanheng-zhao@users.noreply.github.com> Date: Thu, 23 Apr 2026 12:04:52 +0800 Subject: [PATCH 21/45] [Model] Ming-flash-omni-2.0 Omni-Speech and TTS (#2890) Signed-off-by: Yuanheng Zhao Signed-off-by: yuanheng Signed-off-by: LHXuuu Co-authored-by: LHXuuu Co-authored-by: Hongsheng Liu --- .../ming_flash_omni/README.md | 98 +- .../ming_flash_omni/end2end.py | 26 +- .../ming_flash_omni_tts/README.md | 47 + .../ming_flash_omni_tts/end2end.py | 129 ++ .../online_serving/ming_flash_omni/README.md | 183 +-- .../ming_flash_omni_tts/README.md | 54 + .../ming_flash_omni_tts/run_server.sh | 25 + .../ming_flash_omni_tts/speech_client.py | 93 ++ recipes/README.md | 2 + recipes/inclusionAI/Ming-flash-omni-2.0.md | 210 +++ .../offline_inference/test_ming_flash_omni.py | 53 + .../bailingmm_moe_v2_lite_ci.yaml | 48 +- ...bailingmm_moe_v2_lite_thinker_only_ci.yaml | 35 + .../models/ming_flash_omni/test_talker_cfm.py | 146 +++ .../ming_flash_omni/test_talker_modules.py | 162 +++ vllm_omni/config/model.py | 17 +- vllm_omni/entrypoints/openai/serving_chat.py | 26 +- .../entrypoints/openai/serving_speech.py | 103 ++ .../models/ming_flash_omni/__init__.py | 2 + .../models/ming_flash_omni/audio_vae.py | 390 ++++++ .../models/ming_flash_omni/ming_flash_omni.py | 106 +- .../ming_flash_omni/ming_flash_omni_talker.py | 586 +++++++++ .../models/ming_flash_omni/prompt_utils.py | 50 + .../models/ming_flash_omni/spk_embedding.py | 44 + .../models/ming_flash_omni/talker_module.py | 1151 +++++++++++++++++ .../models/ming_flash_omni/text_processing.py | 535 ++++++++ .../models/ming_flash_omni/voice_presets.py | 289 +++++ vllm_omni/model_executor/models/registry.py | 5 + .../stage_configs/ming_flash_omni.yaml | 70 + .../stage_configs/ming_flash_omni_tts.yaml | 32 + .../stage_input_processors/ming_flash_omni.py | 104 ++ .../configs/ming_flash_omni.py | 60 +- 32 files changed, 4620 insertions(+), 261 deletions(-) create mode 100644 examples/offline_inference/ming_flash_omni_tts/README.md create mode 100644 examples/offline_inference/ming_flash_omni_tts/end2end.py create mode 100644 examples/online_serving/ming_flash_omni_tts/README.md create mode 100755 examples/online_serving/ming_flash_omni_tts/run_server.sh create mode 100644 examples/online_serving/ming_flash_omni_tts/speech_client.py create mode 100644 recipes/inclusionAI/Ming-flash-omni-2.0.md create mode 100644 tests/e2e/stage_configs/bailingmm_moe_v2_lite_thinker_only_ci.yaml create mode 100644 tests/model_executor/models/ming_flash_omni/test_talker_cfm.py create mode 100644 tests/model_executor/models/ming_flash_omni/test_talker_modules.py create mode 100644 vllm_omni/model_executor/models/ming_flash_omni/audio_vae.py create mode 100644 vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_talker.py create mode 100644 vllm_omni/model_executor/models/ming_flash_omni/prompt_utils.py create mode 100644 vllm_omni/model_executor/models/ming_flash_omni/spk_embedding.py create mode 100644 vllm_omni/model_executor/models/ming_flash_omni/talker_module.py create mode 100644 vllm_omni/model_executor/models/ming_flash_omni/text_processing.py create mode 100644 vllm_omni/model_executor/models/ming_flash_omni/voice_presets.py create mode 100644 vllm_omni/model_executor/stage_configs/ming_flash_omni.yaml create mode 100644 vllm_omni/model_executor/stage_configs/ming_flash_omni_tts.yaml create mode 100644 vllm_omni/model_executor/stage_input_processors/ming_flash_omni.py diff --git a/examples/offline_inference/ming_flash_omni/README.md b/examples/offline_inference/ming_flash_omni/README.md index 7414163fc..be90b408d 100644 --- a/examples/offline_inference/ming_flash_omni/README.md +++ b/examples/offline_inference/ming_flash_omni/README.md @@ -1,75 +1,91 @@ # Ming-flash-omni 2.0 -[Ming-flash-omni-2.0](https://github.com/inclusionAI/Ming) is an omni-modal model supporting text, image, video, and audio understanding, with outputs in text, image, and audio. For now, Ming-flash-omni-2.0 in vLLM-Omni is supported with thinker stage (multi-modal understanding). +[Ming-flash-omni-2.0](https://github.com/inclusionAI/Ming) is an omni-modal model supporting text, image, video, and audio understanding, with text and speech outputs. + +vLLM-Omni supports two deployment modes: + +| Mode | Stage config | Output | +|------|-------------|--------| +| Thinker only (multimodal understanding) | `ming_flash_omni_thinker.yaml` (default `--omni`) | Text | +| Thinker + Talker (omni-speech) | `ming_flash_omni.yaml` | Text + Audio | + +For standalone TTS (talker only), see [`examples/offline_inference/ming_flash_omni_tts/`](../ming_flash_omni_tts/). ## Setup Please refer to the [stage configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) to configure memory allocation appropriately for your hardware setup. +The default `--omni` flag runs thinker only. For omni-speech, pass the two-stage config explicitly: + +```bash +--stage-configs-path vllm_omni/model_executor/stage_configs/ming_flash_omni.yaml +``` + ## Run examples -### Text-only +The end-to-end script defaults to built-in assets; pass `--image-path`, +`--audio-path`, or `--video-path` to override. + ```bash +# Text-only python examples/offline_inference/ming_flash_omni/end2end.py --query-type text + +# Image / audio / video / mixed understanding +python examples/offline_inference/ming_flash_omni/end2end.py --query-type use_image +python examples/offline_inference/ming_flash_omni/end2end.py --query-type use_audio +python examples/offline_inference/ming_flash_omni/end2end.py --query-type use_video --num-frames 16 +python examples/offline_inference/ming_flash_omni/end2end.py --query-type use_mixed_modalities \ + --image-path /path/to/image.jpg --audio-path /path/to/audio.wav ``` #### Reasoning (Thinking Mode) -Reasoning (Thinking) mode is enabled via applying "detailed thinking on" when building the system prompt template (in `apply_chat_template`). - -In the end2end example, a default problem for thinking mode is provided, as referred to the example usage of Ming's cookbook; -To utilize it, you have to download the example figure from https://github.com/inclusionAI/Ming/blob/3954fcb880ff5e61ff128bcf7f1ec344d46a6fe3/figures/cases/3_0.png +Reasoning ("detailed thinking on") is applied by the script when +`--query-type reasoning` is set. The default prompt matches Ming's cookbook +and expects the reference figure from the upstream repo — see +`get_reasoning_query` in `end2end.py`. ```bash python examples/offline_inference/ming_flash_omni/end2end.py -q reasoning --image-path ./3_0.png ``` -### Image understanding -```bash -python examples/offline_inference/ming_flash_omni/end2end.py --query-type use_image +### Omni-speech (thinker + talker) -# With a local image -python examples/offline_inference/ming_flash_omni/end2end.py --query-type use_image --image-path /path/to/image.jpg +To enable spoken output, use the two-stage config and request `audio` (or `text,audio`) modalities. +The thinker processes your multimodal input, generates text, then the talker synthesises the response as speech. + +**Audio-only output** (speech response, no text): +```bash +python examples/offline_inference/ming_flash_omni/end2end.py \ + --query-type text \ + --stage-configs-path vllm_omni/model_executor/stage_configs/ming_flash_omni.yaml \ + --modalities audio \ + --output-dir output_ming_omni_speech ``` -### Audio understanding +**Both text and audio output**: ```bash -python examples/offline_inference/ming_flash_omni/end2end.py --query-type use_audio - -# With a local audio file -python examples/offline_inference/ming_flash_omni/end2end.py --query-type use_audio --audio-path /path/to/audio.wav +python examples/offline_inference/ming_flash_omni/end2end.py \ + --query-type use_audio \ + --stage-configs-path vllm_omni/model_executor/stage_configs/ming_flash_omni.yaml \ + --modalities text,audio \ + --output-dir output_ming_omni_speech ``` -### Video understanding -```bash -python examples/offline_inference/ming_flash_omni/end2end.py --query-type use_video +Generated `.wav` files are saved to `--output-dir` (default `output_ming`), one per request. -# With a local video and custom frame count -python examples/offline_inference/ming_flash_omni/end2end.py --query-type use_video --video-path /path/to/video.mp4 --num-frames 16 -``` - -### Mixed modalities (image + audio) -```bash -python examples/offline_inference/ming_flash_omni/end2end.py --query-type use_mixed_modalities \ - --image-path /path/to/image.jpg \ - --audio-path /path/to/audio.wav -``` - -If media file paths are not provided, the script uses built-in default assets. +The stage config allocates thinker on GPUs 0–3 and talker on GPU 3 by default. Adjust `devices` in the YAML to match your hardware. ### Modality control -To control output modalities (e.g. text-only output): -```bash -python examples/offline_inference/ming_flash_omni/end2end.py --query-type use_audio --modalities text -``` -*For now, only text output is supported* +| `--modalities` | Thinker output | Talker | Saved files | +|---------------|----------------|--------|-------------| +| `text` (default) | Text | Not run | `.txt` | +| `audio` | Text (internal) | Runs | `.wav` | +| `text,audio` | Text | Runs | `.txt` + `.wav` | -### Custom stage config -```bash -python examples/offline_inference/ming_flash_omni/end2end.py --query-type use_image \ - --stage-configs-path /path/to/your_config.yaml -``` +Pass `--stage-configs-path /path/to/your_config.yaml` to any of the commands +above to override the stage config. ## Online serving diff --git a/examples/offline_inference/ming_flash_omni/end2end.py b/examples/offline_inference/ming_flash_omni/end2end.py index 8f8730131..e00dcea7b 100644 --- a/examples/offline_inference/ming_flash_omni/end2end.py +++ b/examples/offline_inference/ming_flash_omni/end2end.py @@ -7,6 +7,7 @@ import time from typing import NamedTuple import numpy as np +import soundfile as sf import vllm from PIL import Image from transformers import AutoProcessor @@ -319,7 +320,16 @@ def main(args): seed=SEED, detokenize=True, ) - sampling_params_list = [thinker_sampling_params] + # Talker (ming_tts) uses a custom generation loop (CFM + AudioVAE); + # vLLM sampling is a no-op here — max_tokens=1 just satisfies the scheduler. + talker_sampling_params = SamplingParams( + temperature=0.0, + max_tokens=1, + ) + all_sampling_params = [thinker_sampling_params, talker_sampling_params] + # Match sampling params to the number of configured stages + # (thinker-only yaml → 1, thinker+talker yaml → 2). + sampling_params_list = all_sampling_params[: omni.num_stages] prompts = [query_result.inputs for _ in range(args.num_prompts)] @@ -362,7 +372,19 @@ def main(args): print(f"Failed to write output file {out_txt}: {e}") elif stage_outputs.final_output_type == "audio": - raise NotImplementedError("Add audio example after talker supported.") + request_id = output.request_id + mm = output.outputs[0].multimodal_output + if mm and "audio" in mm: + audio = mm["audio"] + sr_raw = mm.get("sr", 44100) + sample_rate = int(sr_raw.item() if hasattr(sr_raw, "item") else sr_raw) + audio_numpy = audio.float().squeeze().cpu().numpy() + output_wav = os.path.join(output_dir, f"{request_id}.wav") + sf.write(output_wav, audio_numpy, samplerate=sample_rate, format="WAV") + print( + f"Request ID: {request_id}, audio saved to {output_wav} " + f"({len(audio_numpy) / sample_rate:.2f}s, {sample_rate}Hz)" + ) processed_count += 1 if profiler_enabled and processed_count >= total_requests: diff --git a/examples/offline_inference/ming_flash_omni_tts/README.md b/examples/offline_inference/ming_flash_omni_tts/README.md new file mode 100644 index 000000000..15b84041d --- /dev/null +++ b/examples/offline_inference/ming_flash_omni_tts/README.md @@ -0,0 +1,47 @@ +# Ming-flash-omni Standalone TTS (Offline) + +This example runs **Ming-flash-omni-2.0 talker-only** offline inference with: + +- `model`: `Jonathan1909/Ming-flash-omni-2.0` +- `stage config`: `vllm_omni/model_executor/stage_configs/ming_flash_omni_tts.yaml` + +It follows the Ming cookbook parameter style: + +- `prompt`: `"Please generate speech based on the following description.\n"` +- `max_decode_steps`: `200` +- `cfg`: `2.0` +- `sigma`: `0.25` +- `temperature`: `0.0` + +## Quick Start + +```bash +python examples/offline_inference/ming_flash_omni_tts/end2end.py --case style +``` + +## Cases + +```bash +# Style +python examples/offline_inference/ming_flash_omni_tts/end2end.py --case style + +# IP +python examples/offline_inference/ming_flash_omni_tts/end2end.py --case ip + +# Basic (speed/pitch/volume control) +python examples/offline_inference/ming_flash_omni_tts/end2end.py --case basic +``` + +## Useful Arguments + +- `--text`: override default text in the selected case +- `--output`: custom output wav path +- `--model`: local model path or HF repo id +- `--stage-configs-path`: custom talker stage config path +- `--log-stats`: enable runtime stats logs + +## Notes + +- This directory is for **standalone talker deployment (TTS)**. +- For Ming thinker multimodal understanding examples, see: + `examples/offline_inference/ming_flash_omni/`. diff --git a/examples/offline_inference/ming_flash_omni_tts/end2end.py b/examples/offline_inference/ming_flash_omni_tts/end2end.py new file mode 100644 index 000000000..0a2d25646 --- /dev/null +++ b/examples/offline_inference/ming_flash_omni_tts/end2end.py @@ -0,0 +1,129 @@ +"""Offline e2e example for Ming-flash-omni-2.0 standalone talker (TTS).""" + +import os +from typing import Any + +import soundfile as sf +import torch +from vllm.utils.argparse_utils import FlexibleArgumentParser + +os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" + +from vllm_omni.entrypoints.omni import Omni +from vllm_omni.inputs.data import OmniTokensPrompt +from vllm_omni.model_executor.models.ming_flash_omni.prompt_utils import ( + DEFAULT_PROMPT, + create_instruction, +) + +MODEL_NAME = "Jonathan1909/Ming-flash-omni-2.0" +DEFAULT_STAGE_CONFIG = "vllm_omni/model_executor/stage_configs/ming_flash_omni_tts.yaml" + + +def get_messages(case: str, text_override: str | None) -> dict[str, Any]: + if case == "style": + text = text_override or "我会一直在这里陪着你，直到你慢慢、慢慢地沉入那个最温柔的梦里……好吗？" + instruction = create_instruction( + { + "风格": "这是一种ASMR耳语，属于一种旨在引发特殊感官体验的创意风格。这个女性使用轻柔的普通话进行耳语，声音气音成分重。音量极低，紧贴麦克风，语速极慢，旨在制造触发听者颅内快感的声学刺激。", + } + ) + return { + "prompt": DEFAULT_PROMPT, + "text": text, + "instruction": instruction, + "use_zero_spk_emb": True, + } + if case == "ip": + text = text_override or "这款产品的名字，叫变态坑爹牛肉丸。" + return { + "prompt": DEFAULT_PROMPT, + "text": text, + "instruction": create_instruction({"IP": "灵小甄"}), + "use_zero_spk_emb": True, + } + if case == "basic": + text = text_override or "我们当迎着阳光辛勤耕作，去摘取，去制作，去品尝，去馈赠。" + return { + "prompt": DEFAULT_PROMPT, + "text": text, + "instruction": create_instruction({"语速": "快速", "基频": "中", "音量": "中"}), + "use_zero_spk_emb": True, + } + raise ValueError(f"Unknown case: {case}") + + +def save_audio(mm: dict[str, Any], output_path: str) -> None: + if not mm or "audio" not in mm: + raise RuntimeError("No audio found in model output") + audio = mm["audio"] + sr_raw = mm.get("sr", 44100) + if isinstance(sr_raw, torch.Tensor): + sample_rate = int(sr_raw.item()) + else: + sample_rate = int(sr_raw) + waveform = audio.squeeze().float().cpu().numpy() + sf.write(output_path, waveform, sample_rate) + print(f"Saved {output_path} ({len(waveform) / sample_rate:.2f}s, {sample_rate}Hz)") + + +def parse_args(): + parser = FlexibleArgumentParser(description="Ming-flash-omni standalone talker offline e2e example") + parser.add_argument("--model", type=str, default=MODEL_NAME, help="Model name or local path.") + parser.add_argument( + "--stage-configs-path", + type=str, + default=DEFAULT_STAGE_CONFIG, + help="Path to stage configs yaml for standalone talker deployment.", + ) + parser.add_argument( + "--case", + type=str, + default="style", + choices=["style", "ip", "basic"], + help="Example case.", + ) + parser.add_argument("--text", type=str, default=None, help="Override default text for the selected case.") + parser.add_argument("--output", type=str, default=None, help="Output wav path.") + parser.add_argument("--log-stats", action="store_true", default=False, help="Enable stats logging.") + parser.add_argument("--init-timeout", type=int, default=600, help="Engine init timeout in seconds.") + parser.add_argument("--stage-init-timeout", type=int, default=300, help="Single stage init timeout in seconds.") + return parser.parse_args() + + +def main(): + args = parse_args() + + omni = Omni( + model=args.model, + stage_configs_path=args.stage_configs_path, + trust_remote_code=True, + log_stats=args.log_stats, + init_timeout=args.init_timeout, + stage_init_timeout=args.stage_init_timeout, + ) + + messages = get_messages(args.case, args.text) + decode_args = { + # Standalone TTS deployment + "ming_task": "instruct", + "max_decode_steps": 200, + "cfg": 2.0, + "sigma": 0.25, + "temperature": 0.0, + } + req = OmniTokensPrompt( + prompt_token_ids=[0], + additional_information={**messages, **decode_args}, + ) + + outputs = omni.generate(req) + mm = outputs[0].outputs[0].multimodal_output + + output_path = args.output or f"output_{args.case}.wav" + save_audio(mm, output_path) + omni.close() + + +if __name__ == "__main__": + main() diff --git a/examples/online_serving/ming_flash_omni/README.md b/examples/online_serving/ming_flash_omni/README.md index 502232725..8b7d03e21 100644 --- a/examples/online_serving/ming_flash_omni/README.md +++ b/examples/online_serving/ming_flash_omni/README.md @@ -4,38 +4,49 @@ Please refer to [README.md](../../../README.md) +## Deployment modes + +| Mode | Launch command | Output | +|------|---------------|--------| +| Thinker only (multimodal understanding) | `vllm serve ... --omni` | Text | +| Thinker + Talker (omni-speech) | `vllm serve ... --omni --stage-configs-path ming_flash_omni.yaml` | Text + Audio | + +For standalone TTS (talker only), see [`examples/online_serving/ming_flash_omni_tts/`](../ming_flash_omni_tts/). + ## Run examples (Ming-flash-omni 2.0) ### Launch the Server +**Thinker only (text output):** ```bash vllm serve Jonathan1909/Ming-flash-omni-2.0 --omni --port 8091 ``` -If you have custom stage configs file, launch the server with command below +**Thinker + Talker (omni-speech, text + audio output):** ```bash -vllm serve Jonathan1909/Ming-flash-omni-2.0 --omni --port 8091 --stage-configs-path /path/to/stage_configs_file +vllm serve Jonathan1909/Ming-flash-omni-2.0 --omni --port 8091 \ + --stage-configs-path vllm_omni/model_executor/stage_configs/ming_flash_omni.yaml ``` +Pass `--stage-configs-path /path/to/your_config.yaml` to use a custom stage +config. + ### Send Multi-modal Request -#### Send request via python +Shared Python client (supports `text | use_image | use_audio | use_video | +use_mixed_modalities`; pass `--image-path` / `--audio-path` / `--video-path` +for local files or URLs, `--modalities text` for output, `--help` for the +full flag list): ```bash -python examples/online_serving/openai_chat_completion_client_for_multimodal_generation.py --model Jonathan1909/Ming-flash-omni-2.0 --query-type use_mixed_modalities --port 8091 --host "localhost" --modalities text +python examples/online_serving/openai_chat_completion_client_for_multimodal_generation.py \ + --model Jonathan1909/Ming-flash-omni-2.0 \ + --query-type use_mixed_modalities \ + --port 8091 --host localhost \ + --modalities text ``` -The Python client supports the following command-line arguments: - -- `--query-type` (or `-q`): Query type. Options: `text`, `use_audio`, `use_image`, `use_video`, `use_mixed_modalities` -- `--video-path` (or `-v`): Path to local video file or URL. If not provided and query-type uses video, uses default video URL. Supports local file paths (automatically encoded to base64) or HTTP/HTTPS URLs. Example: `--video-path /path/to/video.mp4` or `--video-path https://example.com/video.mp4` -- `--image-path` (or `-i`): Path to local image file or URL. If not provided and query-type uses image, uses default image URL. Supports local file paths (automatically encoded to base64) or HTTP/HTTPS URLs and common image formats: JPEG, PNG, GIF, WebP. Example: `--image-path /path/to/image.jpg` or `--image-path https://example.com/image.png` -- `--audio-path` (or `-a`): Path to local audio file or URL. If not provided and query-type uses audio, uses default audio URL. Supports local file paths (automatically encoded to base64) or HTTP/HTTPS URLs and common audio formats: MP3, WAV, OGG, FLAC, M4A. Example: `--audio-path /path/to/audio.wav` or `--audio-path https://example.com/audio.mp3` -- `--prompt` (or `-p`): Custom text prompt/question. If not provided, uses default prompt for the selected query type. Example: `--prompt "What are the main activities shown in this video?"` -- `--modalities`: Output modalities. For now, only `text` is supported. Example: `--modalities text` - - -#### Send request via curl +Parameterized curl wrapper in this directory: ```bash bash run_curl_multimodal_generation.sh text @@ -47,83 +58,17 @@ bash run_curl_multimodal_generation.sh use_mixed_modalities ## Modality control -Ming-flash-omni 2.0 currently supports text output only (thinker stage). +| `modalities` | Server config | Output | +|-------------|--------------|--------| +| `["text"]` or omitted | Thinker only | Text | +| `["audio"]` | Thinker + Talker | Audio (speech) | +| `["text", "audio"]` | Thinker + Talker | Text + Audio | -| Modalities | Output | -|------------|--------| -| `["text"]` | Text only | -| Not specified | Text only (default) | +For ready-to-copy curl examples (text / audio / multimodal input, SSE +streaming, reasoning mode), see the recipe at +[`recipes/inclusionAI/Ming-flash-omni-2.0.md`](../../../recipes/inclusionAI/Ming-flash-omni-2.0.md). -### Using curl - -```bash -curl http://localhost:8091/v1/chat/completions \ - -H "Content-Type: application/json" \ - -d '{ - "model": "Jonathan1909/Ming-flash-omni-2.0", - "messages": [ - {"role": "system", "content": [{"type": "text", "text": "你是一个友好的AI助手。\n\ndetailed thinking off"}]}, - {"role": "user", "content": "请详细介绍鹦鹉的生活习性。"} - ], - "modalities": ["text"] - }' -``` - -### Using OpenAI Python SDK - -```python -from openai import OpenAI - -client = OpenAI(base_url="http://localhost:8091/v1", api_key="EMPTY") - -response = client.chat.completions.create( - model="Jonathan1909/Ming-flash-omni-2.0", - messages=[ - {"role": "system", "content": [{"type": "text", "text": "你是一个友好的AI助手。\n\ndetailed thinking off"}]}, - {"role": "user", "content": "请详细介绍鹦鹉的生活习性。"}, - ], - modalities=["text"], -) -print(response.choices[0].message.content) -``` - -### Multi-modal input with OpenAI Python SDK - -```python -from openai import OpenAI - -client = OpenAI(base_url="http://localhost:8091/v1", api_key="EMPTY") - -response = client.chat.completions.create( - model="Jonathan1909/Ming-flash-omni-2.0", - messages=[ - {"role": "system", "content": [{"type": "text", "text": "你是一个友好的AI助手。\n\ndetailed thinking off"}]}, - { - "role": "user", - "content": [ - {"type": "image_url", "image_url": {"url": "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/cherry_blossom.jpg"}}, - {"type": "text", "text": "Describe this image in detail."}, - ], - }, - ], - modalities=["text"], -) -print(response.choices[0].message.content) -``` - -## Streaming Output - -To enable streaming output: - -```bash -python examples/online_serving/openai_chat_completion_client_for_multimodal_generation.py \ - --query-type use_image \ - --model Jonathan1909/Ming-flash-omni-2.0 \ - --modalities text \ - --stream -``` - -Or with the OpenAI Python SDK: +## OpenAI Python SDK — streaming ```python from openai import OpenAI @@ -146,59 +91,5 @@ for chunk in response: print() ``` -Or using curl: - -```bash -curl http://localhost:8091/v1/chat/completions \ - -H "Content-Type: application/json" \ - -d '{ - "model": "Jonathan1909/Ming-flash-omni-2.0", - "messages": [ - {"role": "system", "content": [{"type": "text", "text": "你是一个友好的AI助手。\n\ndetailed thinking off"}]}, - {"role": "user", "content": "请详细介绍鹦鹉的生活习性。"} - ], - "modalities": ["text"], - "stream": true, - }' -``` - - -## Reasoning (Thinking Mode) - -To enable reasoning/thinking mode, change `detailed thinking off` to `detailed thinking on` in the system prompt: - -### Using curl - -```bash -curl http://localhost:8091/v1/chat/completions \ - -H "Content-Type: application/json" \ - -d '{ - "model": "Jonathan1909/Ming-flash-omni-2.0", - "messages": [ - {"role": "system", "content": [{"type": "text", "text": "你是一个友好的AI助手。\n\ndetailed thinking on"}]}, - {"role": "user", "content": [ - {"type": "image_url", "image_url": {"url": "https://example.com/math_problem.png"}}, - {"type": "text", "text": "Solve this math problem step by step."} - ]} - ], - "modalities": ["text"] - }' -``` - -### Using OpenAI Python SDK - -```python -from openai import OpenAI - -client = OpenAI(base_url="http://localhost:8091/v1", api_key="EMPTY") - -response = client.chat.completions.create( - model="Jonathan1909/Ming-flash-omni-2.0", - messages=[ - {"role": "system", "content": [{"type": "text", "text": "你是一个友好的AI助手。\n\ndetailed thinking on"}]}, - {"role": "user", "content": "If a train travels 120 km in 2 hours, what is its average speed?"}, - ], - modalities=["text"], -) -print(response.choices[0].message.content) -``` +The `--stream` flag on the Python client script above shows the same pattern +driven by the shared multimodal client. diff --git a/examples/online_serving/ming_flash_omni_tts/README.md b/examples/online_serving/ming_flash_omni_tts/README.md new file mode 100644 index 000000000..e56d9ea12 --- /dev/null +++ b/examples/online_serving/ming_flash_omni_tts/README.md @@ -0,0 +1,54 @@ +# Ming-flash-omni Standalone TTS (Online Serving) + +This directory contains online e2e examples for **Ming-flash-omni-2.0 standalone talker deployment**. + +Server uses: + +- `model`: `Jonathan1909/Ming-flash-omni-2.0` +- `stage config`: `vllm_omni/model_executor/stage_configs/ming_flash_omni_tts.yaml` + +## Launch the Server + +```bash +# from repo root +bash examples/online_serving/ming_flash_omni_tts/run_server.sh +``` + +Equivalent manual command: + +```bash +vllm serve Jonathan1909/Ming-flash-omni-2.0 \ + --stage-configs-path vllm_omni/model_executor/stage_configs/ming_flash_omni_tts.yaml \ + --host 0.0.0.0 \ + --port 8091 \ + --trust-remote-code \ + --omni +``` + +## Send TTS Request + +Python client: + +```bash +python examples/online_serving/ming_flash_omni_tts/speech_client.py \ + --text "我们当迎着阳光辛勤耕作，去摘取，去制作，去品尝，去馈赠。" \ + --output ming_online.wav +``` + +Long-form `instructions` (e.g. ASMR whisper style) via the client: + +```bash +python examples/online_serving/ming_flash_omni_tts/speech_client.py \ + --text "我会一直在这里陪着你，直到你慢慢、慢慢地沉入那个最温柔的梦里……好吗？" \ + --instructions "这是一种ASMR耳语，属于一种旨在引发特殊感官体验的创意风格。这个女性使用轻柔的普通话进行耳语，声音气音成分重。音量极低，紧贴麦克风，语速极慢，旨在制造触发听者颅内快感的声学刺激。" \ + --output ming_online_asmr.wav +``` + +## Notes + +- This is the **online serving** counterpart of `examples/offline_inference/ming_flash_omni_tts/`. +- The server uses `use_zero_spk_emb=True` and the default decode args + (`max_decode_steps=200`, `cfg=2.0`, `sigma=0.25`, `temperature=0.0`). + For other caption fields (`语速`, `基频`, `IP`, BGM, etc.) or overriding + decode args, use the offline e2e example where `additional_information` + is set explicitly. diff --git a/examples/online_serving/ming_flash_omni_tts/run_server.sh b/examples/online_serving/ming_flash_omni_tts/run_server.sh new file mode 100755 index 000000000..91be35c6c --- /dev/null +++ b/examples/online_serving/ming_flash_omni_tts/run_server.sh @@ -0,0 +1,25 @@ +#!/bin/bash +# Launch vLLM-Omni server for Ming-flash-omni-2.0 standalone talker (TTS). +# +# Usage: +# ./run_server.sh +# MODEL=/path/to/local/model ./run_server.sh +# PORT=8091 ./run_server.sh +# HOST=127.0.0.1 ./run_server.sh # bind only to loopback + +set -e + +MODEL="${MODEL:-Jonathan1909/Ming-flash-omni-2.0}" +HOST="${HOST:-0.0.0.0}" +PORT="${PORT:-8091}" +STAGE_CONFIG="${STAGE_CONFIG:-vllm_omni/model_executor/stage_configs/ming_flash_omni_tts.yaml}" + +echo "Starting Ming standalone TTS server with model: $MODEL" +echo "Stage config: $STAGE_CONFIG" + +vllm serve "$MODEL" \ + --stage-configs-path "$STAGE_CONFIG" \ + --host "$HOST" \ + --port "$PORT" \ + --trust-remote-code \ + --omni diff --git a/examples/online_serving/ming_flash_omni_tts/speech_client.py b/examples/online_serving/ming_flash_omni_tts/speech_client.py new file mode 100644 index 000000000..b3fceba25 --- /dev/null +++ b/examples/online_serving/ming_flash_omni_tts/speech_client.py @@ -0,0 +1,93 @@ +"""Client for Ming standalone TTS via /v1/audio/speech endpoint.""" + +import argparse +import json +import sys + +import httpx + +DEFAULT_API_BASE = "http://localhost:8091" +DEFAULT_API_KEY = "EMPTY" +DEFAULT_MODEL = "Jonathan1909/Ming-flash-omni-2.0" + + +def run_tts(args) -> None: + payload = { + "model": args.model, + "input": args.text, + "response_format": args.response_format, + } + + instructions = args.instructions + if args.instruction_json: + if instructions: + sys.exit("--instructions and --instruction-json are mutually exclusive") + + try: + parsed = json.loads(args.instruction_json) + except json.JSONDecodeError as exc: + sys.exit(f"--instruction-json must be valid JSON: {exc}") + if not isinstance(parsed, dict): + sys.exit("--instruction-json must decode to a JSON object") + # Re-encode with ensure_ascii=False so UTF-8 Chinese keys/values + # arrive at the server intact rather than as \\uXXXX escapes. + instructions = json.dumps(parsed, ensure_ascii=False) + if instructions: + payload["instructions"] = instructions + + print(f"Model: {args.model}") + print(f"Text: {args.text}") + print("Generating audio...") + + api_url = f"{args.api_base}/v1/audio/speech" + headers = { + "Content-Type": "application/json", + "Authorization": f"Bearer {args.api_key}", + } + + with httpx.Client(timeout=300.0) as client: + response = client.post(api_url, json=payload, headers=headers) + + if response.status_code != 200: + print(f"Error: {response.status_code}") + print(response.text) + return + + output_path = args.output or "ming_tts_output.wav" + with open(output_path, "wb") as f: + f.write(response.content) + print(f"Audio saved to: {output_path}") + + +def main(): + parser = argparse.ArgumentParser(description="Ming standalone TTS speech client") + parser.add_argument("--api-base", default=DEFAULT_API_BASE, help="API base URL") + parser.add_argument("--api-key", default=DEFAULT_API_KEY, help="API key") + parser.add_argument("--model", "-m", default=DEFAULT_MODEL, help="Model name or local path") + parser.add_argument("--text", required=True, help="Text to synthesize") + parser.add_argument( + "--response-format", + default="wav", + choices=["wav", "mp3", "flac", "pcm", "aac", "opus"], + help="Audio format (default: wav)", + ) + parser.add_argument("--output", "-o", default=None, help="Output file path") + parser.add_argument( + "--instructions", + default=None, + help="Free-form style description (mapped to caption 风格 on the server).", + ) + parser.add_argument( + "--instruction-json", + default=None, + help=( + "Structured caption JSON forwarded as `instructions`. Accepts Ming " + "caption keys: 方言, 风格, 语速, 基频, 音量, 情感, IP, 说话人, BGM. " + ), + ) + args = parser.parse_args() + run_tts(args) + + +if __name__ == "__main__": + main() diff --git a/recipes/README.md b/recipes/README.md index 01ecc41f1..69ce4d750 100644 --- a/recipes/README.md +++ b/recipes/README.md @@ -29,6 +29,8 @@ recipes/ multimodal chat on `1x A100 80GB` - [`Wan-AI/Wan2.2-I2V.md`](./Wan-AI/Wan2.2-I2V.md): image-to-video serving recipe for Wan2.2 14B on `8x Ascend NPU (A2/A3)` +- [`inclusionAI/Ming-flash-omni-2.0.md`](./inclusionAI/Ming-flash-omni-2.0.md): + online serving recipe for multimodal chat (`4x H100 80GB`) and standalone TTS (`1x H100 80GB`) Within a single recipe file, include different hardware support sections such as `GPU`, `ROCm`, and `NPU`, and add concrete tested configurations like diff --git a/recipes/inclusionAI/Ming-flash-omni-2.0.md b/recipes/inclusionAI/Ming-flash-omni-2.0.md new file mode 100644 index 000000000..873158c8a --- /dev/null +++ b/recipes/inclusionAI/Ming-flash-omni-2.0.md @@ -0,0 +1,210 @@ +# Ming-flash-omni 2.0 for omni-speech chat and standalone TTS + +## Summary + +- Vendor: inclusionAI +- Model: `Jonathan1909/Ming-flash-omni-2.0` +- Task: Multimodal chat with text, image, audio, or video input; standalone text-to-speech (TTS); +and image generation +- Mode: Online serving with the OpenAI-compatible API +- Maintainer: Community + +## When to use this recipe + +Use this recipe when you want a known-good starting point for serving +`Jonathan1909/Ming-flash-omni-2.0` with vLLM-Omni in one of three modes: + +- **Thinker only** — multimodal understanding with text output. +- **Thinker + Talker (omni-speech)** — multimodal understanding with text and spoken output. +- **Talker only (TTS)** — standalone text-to-speech via the OpenAI `/v1/audio/speech` endpoint. + +## References + +- Upstream model: + [`inclusionAI/Ming`](https://github.com/inclusionAI/Ming) +- For offline inference and additional client variants, see + `examples/offline_inference/ming_flash_omni{,_tts}/` and + `examples/online_serving/ming_flash_omni{,_tts}/`. + + +## Hardware Support + +This recipe documents reference GPU configurations for the two-stage +omni-speech deployment and the standalone TTS deployment. +Other hardware and configurations are welcome as community validation lands. + +## GPU + +### 4x H100 80GB — omni-speech/chat (thinker + talker) + +The bundled `ming_flash_omni.yaml` runs the thinker with tensor parallel size +4 on GPUs 0–3 and the talker on GPU 3. +Adjust `devices` in the YAML to match your hardware. + +#### Environment + +- OS: Linux +- Python: 3.10+ +- CUDA Driver Version: 590.48.01 +- CUDA 12.5 +- vLLM version: 0.19.0 +- vLLM-Omni version or commit: 0.19.0rc1 + +#### Command + +Thinker only (text output): + +```bash +vllm serve Jonathan1909/Ming-flash-omni-2.0 --omni --port 8091 +``` + +Thinker + talker (text and/or audio output): + +```bash +vllm serve Jonathan1909/Ming-flash-omni-2.0 \ + --omni \ + --port 8091 \ + --stage-configs-path vllm_omni/model_executor/stage_configs/ming_flash_omni.yaml \ + --log-stats +``` + +`--log-stats` is optional but recommended while validating the deployment. + +#### Verification + +Text output from a multimodal (image) input: + +```bash +curl http://localhost:8091/v1/chat/completions \ + -H "Content-Type: application/json" \ + -d '{ + "model": "Jonathan1909/Ming-flash-omni-2.0", + "messages": [ + {"role": "system", "content": [{"type": "text", "text": "你是一个友好的AI助手。\n\ndetailed thinking off"}]}, + {"role": "user", "content": [ + {"type": "image_url", "image_url": {"url": "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/cherry_blossom.jpg"}}, + {"type": "text", "text": "Describe this image in detail."} + ]} + ], + "modalities": ["text"] + }' +``` + +Spoken response from a text query (save the WAV bytes): + +```bash +curl http://localhost:8091/v1/chat/completions \ + -H "Content-Type: application/json" \ + -d '{ + "model": "Jonathan1909/Ming-flash-omni-2.0", + "messages": [ + {"role": "system", "content": [{"type": "text", "text": "你是一个友好的AI助手。\n\ndetailed thinking off"}]}, + {"role": "user", "content": "请详细介绍鹦鹉的生活习性。"} + ], + "modalities": ["audio"] + }' | jq -r '.choices[0].message.audio.data' | base64 -d > ming_omni_parrot.wav +``` + +Text + audio output from an audio input (swap `audio_url` for `video_url` +or `image_url` to exercise the other multimodal input paths): + +```bash +curl http://localhost:8091/v1/chat/completions \ + -H "Content-Type: application/json" \ + -d '{ + "model": "Jonathan1909/Ming-flash-omni-2.0", + "messages": [ + {"role": "system", "content": [{"type": "text", "text": "你是一个友好的AI助手。\n\ndetailed thinking off"}]}, + {"role": "user", "content": [ + {"type": "audio_url", "audio_url": {"url": "https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/mary_had_lamb.ogg"}}, + {"type": "text", "text": "Please recognize the language of this speech and transcribe it. Format: oral."} + ]} + ], + "modalities": ["text", "audio"] + }' | jq -r '.choices[0].message.content' +``` + +Streaming text output via SSE (set `"stream": true`): + +```bash +curl -N http://localhost:8091/v1/chat/completions \ + -H "Content-Type: application/json" \ + -d '{ + "model": "Jonathan1909/Ming-flash-omni-2.0", + "messages": [ + {"role": "system", "content": [{"type": "text", "text": "你是一个友好的AI助手。\n\ndetailed thinking off"}]}, + {"role": "user", "content": "请详细介绍鹦鹉的生活习性。"} + ], + "modalities": ["text"], + "stream": true + }' +``` + +Each SSE event carries a `data:` line with a chat-completion chunk; text +deltas appear at `choices[0].delta.content`. + +#### Notes + +- Output modality is selected by the request body: `"modalities": ["text"]`, + `["audio"]`, or `["text", "audio"]`. The two-stage omni-speech server must be launched + for any request containing `audio`. +- Reasoning mode: flip the system prompt suffix from `detailed thinking off` + to `detailed thinking on` in any request above. +- Memory usage: size depends on output modalities and multimodal input; leave + headroom for video frames and audio caches. + +### 1x H100 80GB — standalone TTS (talker only) + +The bundled `ming_flash_omni_tts.yaml` runs the talker on a single GPU and exposes the OpenAI `/v1/audio/speech` endpoint. + +#### Environment + +- OS: Linux +- Python: 3.10+ +- CUDA Driver Version: 590.48.01 +- CUDA 12.5 +- vLLM version: 0.19.0 +- vLLM-Omni version or commit: 0.19.0rc1 + +#### Command + +```bash +vllm serve Jonathan1909/Ming-flash-omni-2.0 \ + --omni \ + --stage-configs-path vllm_omni/model_executor/stage_configs/ming_flash_omni_tts.yaml \ + --port 8091 \ + --log-stats +``` + +`--log-stats` is optional but recommended while validating the deployment. + +#### Verification + +Basic curl: + +```bash +curl -X POST http://localhost:8091/v1/audio/speech \ + -H "Content-Type: application/json" \ + -d '{ + "model": "Jonathan1909/Ming-flash-omni-2.0", + "input": "我会一直在这里陪着你。", + "response_format": "wav" + }' --output ming_online.wav +``` + +Speaker selection (e.g. `lingguang`): + +```bash +curl -X POST http://localhost:8091/v1/audio/speech \ + -H "Content-Type: application/json" \ + -d '{ + "model": "Jonathan1909/Ming-flash-omni-2.0", + "input": "春天来了，万物复苏，大地一片生机盎然。田野里的油菜花开得金灿灿的，蜜蜂在花丛中忙碌地采蜜。远处的山坡上，桃花和杏花竞相绽放，粉的白的交织在一起，美不胜收。清晨的微风带着泥土的芬芳，轻轻拂过脸颊，让人感到无比惬意。孩子们在田间小路上追逐嬉戏，老人们坐在门前晒太阳，享受着这份宁静与美好。", + "speaker": "lingguang", + "response_format": "wav" + }' --output ming_online_lingguang.wav +``` + +#### Notes + +- The OpenAI `instructions` field is forwarded to the talker as the caption JSON — pass a raw string for `风格` (style) only, or a JSON-encoded object for multiple entries such as `方言` (dialect) and `情感` (emotion). diff --git a/tests/e2e/offline_inference/test_ming_flash_omni.py b/tests/e2e/offline_inference/test_ming_flash_omni.py index c591e910a..ca0b0fe0d 100644 --- a/tests/e2e/offline_inference/test_ming_flash_omni.py +++ b/tests/e2e/offline_inference/test_ming_flash_omni.py @@ -36,6 +36,21 @@ def build_prompt(user_text: str) -> str: def get_eager_config(): + path = modify_stage_config( + str(Path(__file__).parent.parent / "stage_configs" / "bailingmm_moe_v2_lite_thinker_only_ci.yaml"), + updates={ + "stage_args": { + 0: { + "engine_args.enforce_eager": "true", + }, + }, + }, + ) + return path + + +def get_eager_tts_config(): + """Thinker+talker CI config with enforce_eager on the thinker stage.""" path = modify_stage_config( str(Path(__file__).parent.parent / "stage_configs" / "bailingmm_moe_v2_lite_ci.yaml"), updates={ @@ -49,9 +64,14 @@ def get_eager_config(): return path +# Thinker-only config — used by text-output tests. stage_configs = [get_eager_config()] test_params = [(model, stage_config) for model in models for stage_config in stage_configs] +# Thinker+talker config — used by audio-output tests. +stage_configs_tts = [get_eager_tts_config()] +test_params_tts = [(model, stage_config) for model in models for stage_config in stage_configs_tts] + @pytest.mark.core_model @pytest.mark.omni @@ -140,3 +160,36 @@ def test_mixed_to_text(omni_runner, omni_runner_handler) -> None: request_config = {"prompts": prompt, "images": image, "audios": audio, "modalities": ["text"]} omni_runner_handler.send_request(request_config) + + +@pytest.mark.advanced_model +@pytest.mark.omni +@hardware_test(res={"cuda": "H100"}, num_cards=4) +@pytest.mark.parametrize("omni_runner", test_params_tts, indirect=True) +def test_text_to_audio(omni_runner, omni_runner_handler) -> None: + """ + Test text input with audio output via the thinker+talker pipeline. + Input Modal: text + Output Modal: audio + """ + prompt = build_prompt("请简单介绍一下北京。") + request_config = {"prompts": prompt, "modalities": ["audio"]} + + omni_runner_handler.send_request(request_config) + + +@pytest.mark.advanced_model +@pytest.mark.omni +@hardware_test(res={"cuda": "H100"}, num_cards=4) +@pytest.mark.parametrize("omni_runner", test_params_tts, indirect=True) +def test_image_to_audio(omni_runner, omni_runner_handler) -> None: + """ + Test image + text input with audio output via the thinker+talker pipeline. + Input Modal: image + text + Output Modal: audio + """ + image = generate_synthetic_image(224, 224)["np_array"] + prompt = build_prompt(f"{IMAGE_TOKEN}Describe this image briefly.") + request_config = {"prompts": prompt, "images": image, "modalities": ["audio"]} + + omni_runner_handler.send_request(request_config) diff --git a/tests/e2e/stage_configs/bailingmm_moe_v2_lite_ci.yaml b/tests/e2e/stage_configs/bailingmm_moe_v2_lite_ci.yaml index fb0c72cc5..53cc73ce0 100644 --- a/tests/e2e/stage_configs/bailingmm_moe_v2_lite_ci.yaml +++ b/tests/e2e/stage_configs/bailingmm_moe_v2_lite_ci.yaml @@ -1,4 +1,8 @@ -# Thinker stage only +# CI stage config for Ming-flash-omni-2.0 thinker+talker pipeline. +# Stage 0: Thinker (multimodal understanding -> text generation) +# Stage 1: Talker (text -> audio waveform via CFM + AudioVAE) +# +# The following config has been verified on 4x H100-80G GPUs stage_args: - stage_id: 0 stage_type: llm @@ -10,10 +14,12 @@ stage_args: model_arch: MingFlashOmniForConditionalGeneration worker_type: ar scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler - gpu_memory_utilization: 0.9 + gpu_memory_utilization: 0.74 enforce_eager: false trust_remote_code: true - engine_output_type: latent + # Ming Thinker -> talker bridge reads detokenised text from + # source_output.outputs[0].text (not hidden states). + engine_output_type: text distributed_executor_backend: "mp" enable_prefix_caching: false max_num_batched_tokens: 32768 @@ -22,14 +28,48 @@ stage_args: hf_config_name: llm_config load_format: dummy mm_processor_cache_gb: 0 + compilation_config: + pass_config: + # Disable fused all-reduce to avoid a vllm/flashinfer version mismatch. + fuse_allreduce_rms: false final_output: true final_output_type: text is_comprehension: true default_sampling_params: - temperature: 0.4 + temperature: 0.0 top_p: 0.9 max_tokens: 100 repetition_penalty: 1.05 seed: 42 detokenize: true ignore_eos: false + + - stage_id: 1 + stage_type: llm + runtime: + devices: "3" + max_batch_size: 1 + engine_args: + model_stage: ming_tts + model_arch: MingFlashOmniTalkerForConditionalGeneration + worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker + scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler + gpu_memory_utilization: 0.18 + enforce_eager: true + trust_remote_code: true + engine_output_type: audio + enable_prefix_caching: false + max_num_batched_tokens: 1000000 + tokenizer_subdir: talker/llm + # The HF repo ships BailingMM2Config (thinker-only) at root; + # OmniModelConfig treats that as "stage does not share outer mrope". + hf_config_name: talker_config + load_format: dummy + engine_input_source: [0] + custom_process_input_func: vllm_omni.model_executor.stage_input_processors.ming_flash_omni.thinker2talker + final_output: true + final_output_type: audio + default_sampling_params: + temperature: 0.0 + max_tokens: 1 + seed: 42 diff --git a/tests/e2e/stage_configs/bailingmm_moe_v2_lite_thinker_only_ci.yaml b/tests/e2e/stage_configs/bailingmm_moe_v2_lite_thinker_only_ci.yaml new file mode 100644 index 000000000..fb0c72cc5 --- /dev/null +++ b/tests/e2e/stage_configs/bailingmm_moe_v2_lite_thinker_only_ci.yaml @@ -0,0 +1,35 @@ +# Thinker stage only +stage_args: + - stage_id: 0 + stage_type: llm + runtime: + devices: "0,1,2,3" + max_batch_size: 1 + engine_args: + model_stage: thinker + model_arch: MingFlashOmniForConditionalGeneration + worker_type: ar + scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler + gpu_memory_utilization: 0.9 + enforce_eager: false + trust_remote_code: true + engine_output_type: latent + distributed_executor_backend: "mp" + enable_prefix_caching: false + max_num_batched_tokens: 32768 + max_model_len: 32768 + tensor_parallel_size: 4 + hf_config_name: llm_config + load_format: dummy + mm_processor_cache_gb: 0 + final_output: true + final_output_type: text + is_comprehension: true + default_sampling_params: + temperature: 0.4 + top_p: 0.9 + max_tokens: 100 + repetition_penalty: 1.05 + seed: 42 + detokenize: true + ignore_eos: false diff --git a/tests/model_executor/models/ming_flash_omni/test_talker_cfm.py b/tests/model_executor/models/ming_flash_omni/test_talker_cfm.py new file mode 100644 index 000000000..419ce00da --- /dev/null +++ b/tests/model_executor/models/ming_flash_omni/test_talker_cfm.py @@ -0,0 +1,146 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from __future__ import annotations + +from types import SimpleNamespace + +import pytest + +from vllm_omni.model_executor.models.ming_flash_omni.talker_module import ( + CFM, + Aggregator, + CFMGraphExecutor, + CFMGraphExecutorPool, + DiT, +) + +torch = pytest.importorskip("torch") +pytest.importorskip("x_transformers") + +pytestmark = [ + pytest.mark.core_model, + pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA for graph capture"), +] + +_LATENT_DIM = 8 +_PATCH_SIZE = 4 +_HIS_PATCH_SIZE = 8 +_LLM_HIDDEN = 16 +_DIT_HIDDEN = 32 +_AGG_HIDDEN = 32 +_NUM_HEADS = 4 +_DEPTH = 2 +_STEPS = 5 +_DTYPE = torch.float32 + + +def _warmup_pipeline(cfm: CFM, aggregator: Aggregator, stop_head: torch.nn.Linear, device: torch.device) -> None: + llm_cond = torch.randn(1, 1, _LLM_HIDDEN, device=device, dtype=_DTYPE) + lat_cond = torch.randn(1, _HIS_PATCH_SIZE, _LATENT_DIM, device=device, dtype=_DTYPE) + y0 = torch.randn(1, _PATCH_SIZE, _LATENT_DIM, device=device, dtype=_DTYPE) + t = torch.linspace(0.0, 1.0, _STEPS + 1, device=device, dtype=_DTYPE) + sde_args = torch.tensor([2.0, 0.25, 0.0], device=device, dtype=_DTYPE) + sde_rnd = torch.randn(_STEPS, 1, _PATCH_SIZE, _LATENT_DIM, device=device, dtype=_DTYPE) + + with torch.no_grad(): + gen_lat = cfm.sample(llm_cond, lat_cond, y0, t, sde_args, sde_rnd) + aggregator(gen_lat) + stop_head(llm_cond[:, -1, :]).softmax(dim=-1) + torch.cuda.synchronize(device) + + +def _build_pipeline(): + device = torch.device("cuda") + dit = ( + DiT( + in_channels=_LATENT_DIM, + hidden_size=_DIT_HIDDEN, + depth=_DEPTH, + num_heads=_NUM_HEADS, + mlp_ratio=2.0, + llm_cond_dim=_LLM_HIDDEN, + ) + .to(device=device, dtype=_DTYPE) + .eval() + ) + cfm = CFM(dit, steps=_STEPS, sway_sampling_coef=-1.0).to(device=device, dtype=_DTYPE).eval() + aggregator = ( + Aggregator( + in_channels=_LATENT_DIM, + hidden_size=_AGG_HIDDEN, + depth=_DEPTH, + num_heads=_NUM_HEADS, + mlp_ratio=2.0, + llm_input_dim=_LLM_HIDDEN, + ) + .to(device=device, dtype=_DTYPE) + .eval() + ) + stop_head = torch.nn.Linear(_LLM_HIDDEN, 2).to(device=device, dtype=_DTYPE).eval() + + config = SimpleNamespace(steps=_STEPS, patch_size=_PATCH_SIZE) + _warmup_pipeline(cfm, aggregator, stop_head, device) + return config, cfm, aggregator, stop_head, device + + +class TestCFMGraphExecutor: + """Capture once, replay twice: outputs must stay consistently-shaped.""" + + def test_execute_shapes_and_replay(self) -> None: + config, cfm, aggregator, stop_head, device = _build_pipeline() + executor = CFMGraphExecutor(config, cfm, aggregator, stop_head) + + bsz = 1 + input_tensor = torch.randn(bsz, 1, _LLM_HIDDEN, device=device, dtype=_DTYPE) + his_lat = torch.randn(bsz, _HIS_PATCH_SIZE, _LATENT_DIM, device=device, dtype=_DTYPE) + + gen_lat, inputs_embeds, stop_out = executor.execute(input_tensor, his_lat) + torch.cuda.synchronize() + + assert gen_lat.shape == (bsz, _PATCH_SIZE, _LATENT_DIM) + assert inputs_embeds.shape == (bsz, 1, _LLM_HIDDEN) + assert stop_out.shape == (bsz, 2) + assert torch.isfinite(gen_lat).all() + assert torch.isfinite(inputs_embeds).all() + # stop_head output is softmax-normalized across the last dim. + assert torch.allclose(stop_out.sum(dim=-1), torch.ones(bsz, device=device, dtype=_DTYPE), atol=1e-4) + + # Replay the captured graph with fresh inputs — shapes must match. + new_input = torch.randn_like(input_tensor) + new_his = torch.randn_like(his_lat) + gen_lat2, inputs_embeds2, stop_out2 = executor.execute(new_input, new_his) + torch.cuda.synchronize() + assert gen_lat2.shape == gen_lat.shape + assert inputs_embeds2.shape == inputs_embeds.shape + assert stop_out2.shape == stop_out.shape + assert executor.initialized is True + + def test_execute_is_noninplace_on_inputs(self) -> None: + config, cfm, aggregator, stop_head, device = _build_pipeline() + executor = CFMGraphExecutor(config, cfm, aggregator, stop_head) + + input_tensor = torch.randn(1, 1, _LLM_HIDDEN, device=device, dtype=_DTYPE) + his_lat = torch.randn(1, _HIS_PATCH_SIZE, _LATENT_DIM, device=device, dtype=_DTYPE) + snapshot_input = input_tensor.clone() + snapshot_his = his_lat.clone() + + executor.execute(input_tensor, his_lat) + torch.cuda.synchronize() + assert torch.equal(input_tensor, snapshot_input) + assert torch.equal(his_lat, snapshot_his) + + +class TestCFMGraphExecutorPool: + def test_pool_acquires_and_releases(self) -> None: + config, cfm, aggregator, stop_head, device = _build_pipeline() + pool = CFMGraphExecutorPool(config, cfm, aggregator, stop_head, pool_size=2) + + input_tensor = torch.randn(1, 1, _LLM_HIDDEN, device=device, dtype=_DTYPE) + his_lat = torch.randn(1, _HIS_PATCH_SIZE, _LATENT_DIM, device=device, dtype=_DTYPE) + + gen_lat, inputs_embeds, stop_out = pool.execute(input_tensor, his_lat) + torch.cuda.synchronize() + assert gen_lat.shape == (1, _PATCH_SIZE, _LATENT_DIM) + assert inputs_embeds.shape == (1, 1, _LLM_HIDDEN) + assert stop_out.shape == (1, 2) + assert pool.pool.qsize() == 2 diff --git a/tests/model_executor/models/ming_flash_omni/test_talker_modules.py b/tests/model_executor/models/ming_flash_omni/test_talker_modules.py new file mode 100644 index 000000000..4cbbc887a --- /dev/null +++ b/tests/model_executor/models/ming_flash_omni/test_talker_modules.py @@ -0,0 +1,162 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from __future__ import annotations + +import pytest + +from vllm_omni.model_executor.models.ming_flash_omni.talker_module import CFM, Aggregator, DiT + +torch = pytest.importorskip("torch") +pytest.importorskip("x_transformers") + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +_LATENT_DIM = 8 +_PATCH_SIZE = 4 +_HIS_PATCH_SIZE = 8 +_LLM_HIDDEN = 16 +_DIT_HIDDEN = 32 +_AGG_HIDDEN = 32 +_NUM_HEADS = 4 +_DEPTH = 2 +_STEPS = 5 + + +def _make_dit() -> DiT: + return DiT( + in_channels=_LATENT_DIM, + hidden_size=_DIT_HIDDEN, + depth=_DEPTH, + num_heads=_NUM_HEADS, + mlp_ratio=2.0, + llm_cond_dim=_LLM_HIDDEN, + ) + + +def _make_aggregator() -> Aggregator: + return Aggregator( + in_channels=_LATENT_DIM, + hidden_size=_AGG_HIDDEN, + depth=_DEPTH, + num_heads=_NUM_HEADS, + mlp_ratio=2.0, + llm_input_dim=_LLM_HIDDEN, + ) + + +class TestDiTDummyForward: + """DiT with dummy weights runs forward + CFG-doubled forward.""" + + def test_forward_shape(self) -> None: + dit = _make_dit().eval() + bsz = 2 + x = torch.randn(bsz, _PATCH_SIZE, _LATENT_DIM) + t = torch.zeros(bsz) + c = torch.randn(bsz, 1, _LLM_HIDDEN) + latent_history = torch.randn(bsz, _HIS_PATCH_SIZE, _LATENT_DIM) + + with torch.no_grad(): + out = dit(x, t, c, latent_history) + + # Output preserves the concatenated (history + time/cond prefix + x) + # token axis: history + 1 (time+cond) + patch. + assert out.shape == (bsz, _HIS_PATCH_SIZE + 1 + _PATCH_SIZE, _LATENT_DIM) + + def test_forward_with_cfg_trims_to_patch(self) -> None: + dit = _make_dit().eval() + bsz = 1 + x = torch.randn(bsz, _PATCH_SIZE, _LATENT_DIM) + t = torch.zeros(()) + c = torch.randn(bsz, 1, _LLM_HIDDEN) + latent_history = torch.randn(bsz, _HIS_PATCH_SIZE, _LATENT_DIM) + + with torch.no_grad(): + out = dit.forward_with_cfg(x, t, c, latent_history) + + # CFG doubles the batch and trims the output to the patch window. + assert out.shape == (2 * bsz, _PATCH_SIZE, _LATENT_DIM) + + +class TestAggregatorDummyForward: + """Aggregator with dummy weights maps latent patch -> LLM hidden.""" + + def test_forward_shape(self) -> None: + agg = _make_aggregator().eval() + bsz = 3 + gen_lat = torch.randn(bsz, _PATCH_SIZE, _LATENT_DIM) + + with torch.no_grad(): + out = agg(gen_lat) + + assert out.shape == (bsz, 1, _LLM_HIDDEN) + + def test_forward_is_finite(self) -> None: + agg = _make_aggregator().eval() + gen_lat = torch.randn(1, _PATCH_SIZE, _LATENT_DIM) + with torch.no_grad(): + out = agg(gen_lat) + assert torch.isfinite(out).all() + + +class TestCFMSampleDummy: + """CFM.sample drives DiT.forward_with_cfg through the integration loop.""" + + def test_sample_shape_and_finite(self) -> None: + cfm = CFM(_make_dit(), steps=_STEPS, sway_sampling_coef=-1.0).eval() + bsz = 1 + llm_cond = torch.randn(bsz, 1, _LLM_HIDDEN) + lat_cond = torch.randn(bsz, _HIS_PATCH_SIZE, _LATENT_DIM) + y0 = torch.randn(bsz, _PATCH_SIZE, _LATENT_DIM) + # Grid used by the talker; must span [0, 1] inclusive. + t = torch.linspace(0.0, 1.0, _STEPS + 1) + sde_args = torch.tensor([2.0, 0.0, 0.0]) # cfg=2.0, sigma=0, temp=0 + sde_rnd = torch.zeros(_STEPS, bsz, _PATCH_SIZE, _LATENT_DIM) + + with torch.no_grad(): + out = cfm.sample(llm_cond, lat_cond, y0, t, sde_args, sde_rnd) + + assert out.shape == y0.shape + assert torch.isfinite(out).all() + + def test_sample_zero_cfg_reduces_to_unguided(self) -> None: + """With cfg=0 the guidance term drops, but output shape is still valid.""" + cfm = CFM(_make_dit(), steps=_STEPS, sway_sampling_coef=None).eval() + bsz = 2 + llm_cond = torch.randn(bsz, 1, _LLM_HIDDEN) + lat_cond = torch.randn(bsz, _HIS_PATCH_SIZE, _LATENT_DIM) + y0 = torch.zeros(bsz, _PATCH_SIZE, _LATENT_DIM) + t = torch.linspace(0.0, 1.0, _STEPS + 1) + sde_args = torch.tensor([0.0, 0.0, 0.0]) + sde_rnd = torch.zeros(_STEPS, bsz, _PATCH_SIZE, _LATENT_DIM) + + with torch.no_grad(): + out = cfm.sample(llm_cond, lat_cond, y0, t, sde_args, sde_rnd) + + assert out.shape == (bsz, _PATCH_SIZE, _LATENT_DIM) + assert torch.isfinite(out).all() + + +class TestTalkerPipelineDummyWiring: + """End-to-end wiring of DiT -> CFM.sample -> Aggregator with dummy weights.""" + + def test_cfm_then_aggregator(self) -> None: + dit = _make_dit().eval() + cfm = CFM(dit, steps=_STEPS, sway_sampling_coef=-1.0).eval() + agg = _make_aggregator().eval() + + bsz = 1 + llm_cond = torch.randn(bsz, 1, _LLM_HIDDEN) + lat_cond = torch.randn(bsz, _HIS_PATCH_SIZE, _LATENT_DIM) + y0 = torch.randn(bsz, _PATCH_SIZE, _LATENT_DIM) + t = torch.linspace(0.0, 1.0, _STEPS + 1) + sde_args = torch.tensor([2.0, 0.0, 0.0]) + sde_rnd = torch.zeros(_STEPS, bsz, _PATCH_SIZE, _LATENT_DIM) + + with torch.no_grad(): + gen_lat = cfm.sample(llm_cond, lat_cond, y0, t, sde_args, sde_rnd) + agg_out = agg(gen_lat) + + assert gen_lat.shape == (bsz, _PATCH_SIZE, _LATENT_DIM) + assert agg_out.shape == (bsz, 1, _LLM_HIDDEN) + assert torch.isfinite(agg_out).all() diff --git a/vllm_omni/config/model.py b/vllm_omni/config/model.py index 217a600b4..72731eda8 100644 --- a/vllm_omni/config/model.py +++ b/vllm_omni/config/model.py @@ -5,7 +5,10 @@ from pydantic import ConfigDict, TypeAdapter from vllm.config import ModelConfig from vllm.config.utils import config from vllm.logger import init_logger -from vllm.transformers_utils.config import get_hf_text_config +from vllm.transformers_utils.config import ( + get_hf_text_config, + thinker_uses_mrope, +) from vllm.transformers_utils.model_arch_config_convertor import ( ModelArchConfigConvertorBase, ) @@ -145,6 +148,18 @@ class OmniModelConfig(ModelConfig): return [self.model_arch] return super().architectures + @property + def uses_mrope(self) -> bool: + if self.hf_config_name is not None: + # talker_config/thinker_config/etc + stage_config = getattr(self.hf_config, self.hf_config_name, None) + if stage_config is None: + # Check the named sub-config's text_config directly. + # Handles mrope resolution of stage-specific cls + # (e.g., talker runs as a standalone cls) + return thinker_uses_mrope(self.hf_config) + return super().uses_mrope + @property def embedding_size(self): if self.hf_config_name is not None: diff --git a/vllm_omni/entrypoints/openai/serving_chat.py b/vllm_omni/entrypoints/openai/serving_chat.py index faf3ea320..fb96b397e 100644 --- a/vllm_omni/entrypoints/openai/serving_chat.py +++ b/vllm_omni/entrypoints/openai/serving_chat.py @@ -559,6 +559,16 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): engine_prompt["additional_information"] = {} engine_prompt["additional_information"]["language"] = [language.strip()] + # Style instruction — used by Ming-flash-omni instruct TTS path + # (ming_task="instruct"). For the omni speech path the thinker2talker + # bridge drops this field to match upstream omni_audio_generation + # which hardcodes instruction=None. + instructions = getattr(request, "instructions", None) + if instructions is not None and isinstance(instructions, str) and instructions.strip(): + if "additional_information" not in engine_prompt or engine_prompt["additional_information"] is None: + engine_prompt["additional_information"] = {} + engine_prompt["additional_information"]["instruction"] = instructions.strip() + return conversation, [engine_prompt] async def _inject_audio_from_video_urls( @@ -1936,7 +1946,8 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): final_res = omni_outputs.request_output # OMNI: Access multimodal_output from CompletionOutput (outputs[0]), not from RequestOutput # Reference: examples/offline_inference/qwen3_omni/end2end.py line 421 - audio_data = final_res.outputs[0].multimodal_output.get("audio") + mm_output = final_res.outputs[0].multimodal_output + audio_data = mm_output.get("audio") if isinstance(audio_data, list): if stream: audio_tensor = audio_data[-1] @@ -1950,9 +1961,20 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): if audio_tensor.ndim > 1: audio_tensor = audio_tensor.flatten() + # Prefer the talker-reported sample rate when present. Qwen3-Omni + # omits "sr" and runs at 24kHz; Ming-flash-omni surfaces a 44.1kHz + # AudioVAE rate via multimodal_output["sr"]. + sr_raw = mm_output.get("sr") + if sr_raw is None: + sample_rate = 24000 + elif hasattr(sr_raw, "item"): + sample_rate = int(sr_raw.item()) + else: + sample_rate = int(sr_raw) + audio_obj = CreateAudio( audio_tensor=audio_tensor, - sample_rate=24000, + sample_rate=sample_rate, response_format="wav", speed=1.0, stream_format="audio", diff --git a/vllm_omni/entrypoints/openai/serving_speech.py b/vllm_omni/entrypoints/openai/serving_speech.py index 5f686c666..4946368e9 100644 --- a/vllm_omni/entrypoints/openai/serving_speech.py +++ b/vllm_omni/entrypoints/openai/serving_speech.py @@ -45,6 +45,12 @@ from vllm_omni.model_executor.models.fish_speech.prompt_utils import ( estimate_fish_voice_clone_prompt_len_from_normalized, normalize_fish_voice_clone_texts, ) +from vllm_omni.model_executor.models.ming_flash_omni.prompt_utils import ( + DEFAULT_PROMPT as MING_DEFAULT_PROMPT, +) +from vllm_omni.model_executor.models.ming_flash_omni.prompt_utils import ( + create_instruction as ming_create_instruction, +) from vllm_omni.outputs import OmniRequestOutput logger = init_logger(__name__) @@ -57,6 +63,7 @@ _COSYVOICE3_TTS_MODEL_STAGES = {"cosyvoice3_talker"} _OMNIVOICE_TTS_MODEL_STAGES = {"omnivoice_generator"} _VOXCPM_TTS_MODEL_STAGES = {"latent_generator", "vae"} _VOXCPM2_TTS_MODEL_STAGES = {"latent_generator"} +_MING_TTS_MODEL_STAGES = {"ming_tts"} _TTS_MODEL_STAGES: set[str] = ( _VOXTRAL_TTS_MODEL_STAGES | _QWEN3_TTS_MODEL_STAGES @@ -65,6 +72,7 @@ _TTS_MODEL_STAGES: set[str] = ( | _OMNIVOICE_TTS_MODEL_STAGES | _VOXCPM_TTS_MODEL_STAGES | _VOXCPM2_TTS_MODEL_STAGES + | _MING_TTS_MODEL_STAGES ) _TTS_LANGUAGES: set[str] = { "Auto", @@ -313,6 +321,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): for stage in self.engine_client.stage_configs ) return "voxcpm" if has_vae_stage or model_stage == "vae" else "voxcpm2" + if model_stage in _MING_TTS_MODEL_STAGES: + return "ming_flash_omni_tts" return None def _compute_max_instructions_length(self) -> int: @@ -336,6 +346,11 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): def _load_supported_speakers(self) -> set[str]: """Load supported speakers (case-insensitive) from the model configuration.""" + if self._tts_model_type == "ming_flash_omni_tts": + # Ming-flash-omni drives speaker selection via the caption JSON + # (audio_sequence[0]["说话人"]) rather than a spk_id table, so there + # is no static speaker list to surface here. + return set() try: if self._tts_model_type == "voxcpm": return set() @@ -837,6 +852,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): return self._validate_voxcpm_request(request) if self._tts_model_type == "voxcpm2": return None # VoxCPM2 accepts any text input + if self._tts_model_type == "ming_flash_omni_tts": + return self._validate_ming_tts_request(request) return self._validate_qwen_tts_request(request) def _voxcpm2_encode(self, text: str) -> list[int]: @@ -858,6 +875,41 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): ids = self._voxcpm2_tokenizer.encode(text, add_special_tokens=True) return split_multichar_chinese(ids, self._voxcpm2_split_map) + def _validate_ming_tts_request(self, request: OpenAICreateSpeechRequest) -> str | None: + """Validate Ming-flash-omni standalone-talker request parameters.""" + if not request.input or not request.input.strip(): + return "Input text cannot be empty" + if request.instructions is not None: + if not isinstance(request.instructions, str): + return "instructions must be a string" + if len(request.instructions) > self._max_instructions_length: + return f"instructions exceeds max length {self._max_instructions_length}" + + if request.task_type is not None: + return "'task_type' is not supported for Ming-flash-omni TTS" + if request.language is not None: + return "'language' is not supported for Ming-flash-omni TTS (language is inferred from input text)" + if request.x_vector_only_mode is not None: + return "'x_vector_only_mode' is not supported for Ming-flash-omni TTS" + if request.initial_codec_chunk_frames is not None: + return "'initial_codec_chunk_frames' is not supported for Ming-flash-omni TTS" + + # Per-request voice cloning from raw audio is not yet wired up: Ming + # extracts spk_emb / prompt_wav_lat / prompt_wav_emb model-side via + # register_prompt_wav() at engine init. For ad-hoc cloning, callers + # should pre-compute speaker_embedding and pass it directly. + if request.ref_audio is not None: + return ( + "'ref_audio' is not yet supported for Ming-flash-omni TTS; " + "use a preset 'voice' or 'speaker_embedding' instead" + ) + if request.ref_text is not None: + return "'ref_text' is not yet supported for Ming-flash-omni TTS" + + if request.max_new_tokens is not None and request.max_new_tokens <= 0: + return "'max_new_tokens' must be a positive integer" + return None + def _validate_ref_audio_format(self, ref_audio: str) -> str | None: """Validate ref_audio is a supported URI format. Returns error or None.""" if not ( @@ -1520,6 +1572,52 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): }, } + # ---- Ming-flash-omni standalone-talker (TTS) helpers ---- + + def _build_ming_prompt(self, request: OpenAICreateSpeechRequest) -> dict[str, Any]: + # request.instructions accepts two forms: + # 1. Plain text: mapped to the caption's 风格 (style) field + # 2. JSON object: parsed and splatted into the caption. Unlocks + # Unknown keys are dropped by `ming_create_instruction`. + caption_fields: dict[str, Any] = {} + if request.instructions: + stripped = request.instructions.strip() + if stripped.startswith("{"): + try: + parsed = json.loads(stripped) + except json.JSONDecodeError: + parsed = None + if isinstance(parsed, dict): + caption_fields.update(parsed) + else: + caption_fields["风格"] = request.instructions + else: + caption_fields["风格"] = request.instructions + + has_spk_emb = request.speaker_embedding is not None + + # TTS path applies ming task type `instruct`. + # voice_name enables talker-side voice preset resolution (e.g. "DB30"). + additional_information: dict[str, Any] = { + "ming_task": "instruct", + "prompt": MING_DEFAULT_PROMPT, + "text": request.input, + "instruction": ming_create_instruction(caption_fields), + "voice_name": request.voice or None, + "use_zero_spk_emb": not has_spk_emb, + "max_decode_steps": request.max_new_tokens or _TTS_MAX_NEW_TOKENS_MAX, + "cfg": 2.0, + "sigma": 0.25, + "temperature": 0.0, + } + if has_spk_emb: + # Passed as plain float list + additional_information["spk_emb"] = list(request.speaker_embedding) + return { + "prompt_token_ids": [0], + "additional_information": additional_information, + } + # ---- Common speech generation helpers ---- async def _prepare_speech_generation( @@ -1582,6 +1680,9 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): elif self._tts_model_type == "cosyvoice3": prompt = await self._build_cosyvoice3_prompt(request) tts_params = {} + elif self._tts_model_type == "ming_flash_omni_tts": + prompt = self._build_ming_prompt(request) + tts_params = {} else: tts_params = self._build_tts_params(request) # Resolve ref_audio (explicit or auto-set for uploaded voices) @@ -1629,6 +1730,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): model_type = "voxcpm" elif self._tts_model_type == "voxcpm2": model_type = "voxcpm2" + elif self._tts_model_type == "ming_flash_omni_tts": + model_type = "ming_flash_omni_tts" elif self._is_tts: model_type = tts_params.get("task_type", ["unknown"])[0] else: diff --git a/vllm_omni/model_executor/models/ming_flash_omni/__init__.py b/vllm_omni/model_executor/models/ming_flash_omni/__init__.py index d7fa44fd7..4cd086e64 100644 --- a/vllm_omni/model_executor/models/ming_flash_omni/__init__.py +++ b/vllm_omni/model_executor/models/ming_flash_omni/__init__.py @@ -2,6 +2,7 @@ # Copyright 2025 The vLLM-Omni team. from .ming_flash_omni import MingFlashOmniForConditionalGeneration +from .ming_flash_omni_talker import MingFlashOmniTalkerForConditionalGeneration from .ming_flash_omni_thinker import ( MingFlashOmniThinkerDummyInputsBuilder, MingFlashOmniThinkerForConditionalGeneration, @@ -11,6 +12,7 @@ from .ming_flash_omni_thinker import ( __all__ = [ "MingFlashOmniForConditionalGeneration", + "MingFlashOmniTalkerForConditionalGeneration", "MingFlashOmniThinkerForConditionalGeneration", "MingFlashOmniThinkerProcessingInfo", "MingFlashOmniThinkerMultiModalProcessor", diff --git a/vllm_omni/model_executor/models/ming_flash_omni/audio_vae.py b/vllm_omni/model_executor/models/ming_flash_omni/audio_vae.py new file mode 100644 index 000000000..9d5c266b4 --- /dev/null +++ b/vllm_omni/model_executor/models/ming_flash_omni/audio_vae.py @@ -0,0 +1,390 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright 2025 The vLLM-Omni team. +# Copyright (c) Ant Group. All rights reserved. +# Adapted from: +# https://github.com/inclusionAI/Ming/tree/e58533db227031990c5a6864dcf5f08fb53ed0d2/AudioVAE + +from __future__ import annotations + +import torch +import torch.nn as nn +import torch.nn.functional as F +from transformers import PretrainedConfig, PreTrainedModel, Qwen2Config, Qwen2Model +from vllm.logger import init_logger + +logger = init_logger(__name__) +try: + import flash_attn # noqa: F401 +except (ImportError, ModuleNotFoundError): + flash_attn = None + logger.warning( + "flash_attn is not available, the model may not yield the " + "exactly same result as the transformers implementation " + "in the audio tower part." + ) + + +class AudioVAEConfig(PretrainedConfig): + model_type = "audio_vae" + + def __init__( + self, + sample_rate: int = 44100, + enc_kwargs: dict | None = None, + dec_kwargs: dict | None = None, + init_method: str = "kaiming", + patch_size: int = 4, + **kwargs, + ): + self.sample_rate = sample_rate + self.enc_kwargs = enc_kwargs or {} + self.dec_kwargs = dec_kwargs or {} + self.init_method = init_method + self.patch_size = patch_size + super().__init__(**kwargs) + + +class ISTFT(nn.Module): + def __init__(self, n_fft: int, hop_length: int, win_length: int, padding: str = "same"): + super().__init__() + if padding not in ["center", "same"]: + raise ValueError("Padding must be 'center' or 'same'.") + self.padding = padding + self.n_fft = n_fft + self.hop_length = hop_length + self.win_length = win_length + window = torch.hann_window(win_length) + self.register_buffer("window", window) + self.buffer_len = self.win_length - self.hop_length + + def _buffer_process(self, x, buffer, pad, last_chunk=False, streaming=False): + if streaming: + if buffer is None: + x = x[:, pad:] + if buffer is not None: + x[:, : self.buffer_len] += buffer + buffer = x[:, -self.buffer_len :] + if not last_chunk: + x = x[:, : -self.buffer_len] + else: + x = x[:, :-pad] + else: + x = x[:, pad:-pad] + return x, buffer + + def forward(self, spec, audio_buffer=None, window_buffer=None, streaming=False, last_chunk=False): + if self.padding == "center": + return torch.istft(spec, self.n_fft, self.hop_length, self.win_length, self.window, center=True) + elif self.padding == "same": + pad = (self.win_length - self.hop_length) // 2 + else: + raise ValueError("Padding must be 'center' or 'same'.") + + B, N, T = spec.shape + ifft = torch.fft.irfft(spec, self.n_fft, dim=1, norm="backward") + ifft = ifft * self.window[None, :, None] + + output_size = (T - 1) * self.hop_length + self.win_length + y = torch.nn.functional.fold( + ifft, + output_size=(1, output_size), + kernel_size=(1, self.win_length), + stride=(1, self.hop_length), + )[:, 0, 0, :] + + y, audio_buffer = self._buffer_process(y, audio_buffer, pad, last_chunk=last_chunk, streaming=streaming) + + window_sq = self.window.square().expand(1, T, -1).transpose(1, 2) + window_envelope = ( + torch.nn.functional.fold( + window_sq, + output_size=(1, output_size), + kernel_size=(1, self.win_length), + stride=(1, self.hop_length), + ) + .squeeze(0) + .squeeze(0) + ) + + window_envelope, window_buffer = self._buffer_process( + window_envelope, window_buffer, pad, last_chunk=last_chunk, streaming=streaming + ) + window_envelope = window_envelope.squeeze() + + assert (window_envelope > 1e-11).all() + y = y / window_envelope + + return y, audio_buffer, window_buffer + + +class ISTFTHead(nn.Module): + def __init__(self, dim: int, n_fft: int, hop_length: int, padding: str = "same"): + super().__init__() + out_dim = n_fft + 2 + self.out = nn.Linear(dim, out_dim) + self.istft = ISTFT(n_fft=n_fft, hop_length=hop_length, win_length=n_fft, padding=padding) + + def forward(self, x, audio_buffer=None, window_buffer=None, streaming=False, last_chunk=False): + x_pred = self.out(x) + x_pred = x_pred.transpose(1, 2) + mag, p = x_pred.chunk(2, dim=1) + mag = torch.exp(mag) + mag = torch.clip(mag, max=1e2) + x = torch.cos(p) + y = torch.sin(p) + S = mag * (x + 1j * y) + audio, audio_buffer, window_buffer = self.istft( + S, audio_buffer=audio_buffer, window_buffer=window_buffer, streaming=streaming, last_chunk=last_chunk + ) + return audio.unsqueeze(1), x_pred, audio_buffer, window_buffer + + +class StreamingLinearUpsample(nn.Module): + def __init__(self, scale_factor=4): + super().__init__() + self.scale_factor = scale_factor + self.upsampler = nn.Upsample(scale_factor=scale_factor, mode="linear", align_corners=False) + + def forward(self, x, state=None, is_last=False): + if state is None: + state = {"prev_chunk": None, "history_last": None, "is_first": True} + + if x is None and not is_last: + return None, state + + if state["is_first"] and is_last: + out = self.upsampler(x.transpose(1, 2)).transpose(1, 2) + return out, None + + output_chunks = [] + + if state["is_first"]: + state["prev_chunk"] = x + state["is_first"] = False + if not is_last: + return None, state + + if state["prev_chunk"] is not None: + p = state["prev_chunk"].transpose(1, 2) + + if state["history_last"] is None: + lookahead = x[:, :1, :].transpose(1, 2) + inp = torch.cat([p, lookahead], dim=2) + up = self.upsampler(inp) + out_prev = up[:, :, : p.size(2) * self.scale_factor] + else: + lookahead = x[:, :1, :].transpose(1, 2) + inp = torch.cat([state["history_last"], p, lookahead], dim=2) + up = self.upsampler(inp) + start = self.scale_factor + end = start + p.size(2) * self.scale_factor + out_prev = up[:, :, start:end] + + output_chunks.append(out_prev.transpose(1, 2)) + state["history_last"] = p[:, :, -1:] + state["prev_chunk"] = x + + if is_last: + p = state["prev_chunk"].transpose(1, 2) + inp = torch.cat([state["history_last"], p], dim=2) + up = self.upsampler(inp) + out_last = up[:, :, self.scale_factor :] + output_chunks.append(out_last.transpose(1, 2)) + state = None + + final_out = torch.cat(output_chunks, dim=1) if output_chunks else None + return final_out, state + + +class Decoder(nn.Module): + def __init__(self, decoder_args, output_dim=320, latent_dim=64, patch_size=-1): + super().__init__() + config = Qwen2Config.from_dict(config_dict=decoder_args) + if flash_attn is None: + config._attn_implementation = "sdpa" + self.decoder = Qwen2Model(config) + self.output_dim = output_dim + self.latent_dim = latent_dim + self.fc1 = nn.Linear(latent_dim, config.hidden_size) + self.hop_length = output_dim + self.head = ISTFTHead( + dim=config.hidden_size, n_fft=self.hop_length * 4, hop_length=self.hop_length, padding="same" + ) + self.patch_size = patch_size + if self.patch_size != -1: + self.upsampling = StreamingLinearUpsample(scale_factor=patch_size) + + def low_level_reconstruct(self, x, past_key_values=None, use_cache=False, stream_state=None, last_chunk=False): + upsample_state, audio_buffer, window_buffer = stream_state + bsz, device, dtype = x.size(0), x.device, x.dtype + x = self.fc1(x) + if self.patch_size != -1: + if use_cache: + x, upsample_state = self.upsampling(x, state=upsample_state, is_last=last_chunk) + if x is None: + stream_state = (upsample_state, audio_buffer, window_buffer) + return torch.empty(bsz, 1, 0, device=device, dtype=dtype), stream_state, past_key_values + else: + x = self.upsampling.upsampler(x.transpose(1, 2)).transpose(1, 2) + + hidden_states_list = [] + + if use_cache and getattr(self.decoder.config, "sliding_window", None) is not None: + sw_size = self.decoder.config.sliding_window + target_len = sw_size - 1 + if past_key_values is None: + past_len = 0 + elif hasattr(past_key_values, "get_seq_length"): + past_len = past_key_values.get_seq_length() + elif isinstance(past_key_values, tuple) and len(past_key_values) > 0: + past_len = past_key_values[0][0].shape[-2] + else: + past_len = 0 + + curr_len = x.shape[1] + + if past_len < target_len and (past_len + curr_len) >= sw_size: + fill_len = target_len - past_len + x_fill = x[:, :fill_len, :] + outputs = self.decoder(inputs_embeds=x_fill, past_key_values=past_key_values, use_cache=use_cache) + hidden_states_list.append(outputs.last_hidden_state) + past_key_values = outputs.past_key_values + x = x[:, fill_len:, :] + + outputs = self.decoder(inputs_embeds=x, past_key_values=past_key_values, use_cache=use_cache) + hidden_states_list.append(outputs.last_hidden_state) + past_key_values = outputs.past_key_values + + if len(hidden_states_list) > 1: + full_hidden_state = torch.cat(hidden_states_list, dim=1) + else: + full_hidden_state = hidden_states_list[0] + + x_out, _, audio_buffer, window_buffer = self.head( + full_hidden_state, + streaming=use_cache, + audio_buffer=audio_buffer, + window_buffer=window_buffer, + last_chunk=last_chunk, + ) + + stream_state = (upsample_state, audio_buffer, window_buffer) + return x_out, stream_state, past_key_values + + +class Encoder(nn.Module): + def __init__(self, encoder_args, input_dim=320, hop_size=320, latent_dim=64, patch_size=-1): + super().__init__() + config = Qwen2Config.from_dict(config_dict=encoder_args) + if flash_attn is None: + config._attn_implementation = "sdpa" + self.encoder = Qwen2Model(config) + self.input_dim = input_dim + self.hop_size = hop_size + self.latent_dim = latent_dim + self.fc1 = nn.Linear(input_dim, config.hidden_size, bias=False) + self.fc2 = nn.Linear(config.hidden_size, config.hidden_size) + self.fc3 = nn.Linear(config.hidden_size, latent_dim * 2) + self.norm = nn.LayerNorm(config.hidden_size) + self.patch_size = patch_size + if patch_size != -1: + config.num_hidden_layers = 4 + self.aggregator = Qwen2Model(config) + self.cls_embed = nn.Parameter(torch.rand(1, 1, config.hidden_size)) + self.cls_embed.data.normal_(0, 0.02) + + def get_frames(self, x): + num_frames_total = (x.size(-1) + self.hop_size - 1) // self.hop_size + expected_len = (num_frames_total - 1) * self.hop_size + self.input_dim + padding_needed = expected_len - x.size(-1) + waveform = F.pad(x, (0, padding_needed), value=0.0) + frames = waveform.unfold(dimension=-1, size=self.input_dim, step=self.hop_size) + return frames + + def pad_patch_insert_cls(self, x): + bsz, _, dim = x.size() + num_frame = x.size(1) + r = num_frame % self.patch_size + pad_num = self.patch_size - r if r else 0 + x = F.pad(x, (0, 0, 0, pad_num), value=0.0) + x = x.reshape(-1, self.patch_size, dim) + x = torch.cat((x, self.cls_embed.expand(x.size(0), -1, -1)), dim=1) + x = x.reshape(bsz, -1, dim) + return x + + def forward(self, waveform): + x = self.get_frames(waveform) + x = self.fc1(x) + x = self.fc2(x) + x = self.encoder(inputs_embeds=x) + x = x.last_hidden_state + + if self.patch_size != -1: + x = self.pad_patch_insert_cls(x) + x = self.aggregator(inputs_embeds=x) + x = x.last_hidden_state + bsz, _, dim = x.size() + x = x.reshape(-1, self.patch_size + 1, dim) + x = x[:, -1:, :].reshape(bsz, -1, dim) + + x = self.fc3(x) + return x, waveform.unsqueeze(1) + + +class AudioVAE(PreTrainedModel): + config_class = AudioVAEConfig + + def __init__(self, config: AudioVAEConfig): + super().__init__(config) + self.encoder = Encoder( + encoder_args=config.enc_kwargs["backbone"], + input_dim=config.enc_kwargs["input_dim"], + hop_size=config.enc_kwargs.get("hop_size", 320), + latent_dim=config.enc_kwargs["latent_dim"], + patch_size=config.patch_size, + ) + self.decoder = Decoder( + decoder_args=config.dec_kwargs["backbone"], + output_dim=config.dec_kwargs["output_dim"], + latent_dim=config.dec_kwargs["latent_dim"], + patch_size=config.patch_size, + ) + self.post_init() + + def _init_weights(self, module): + std = 0.02 + if isinstance(module, nn.Linear): + if self.config.init_method == "kaiming": + nn.init.kaiming_normal_(module.weight, mode="fan_in", nonlinearity="relu") + else: + module.weight.data.normal_(mean=0.0, std=std) + if module.bias is not None: + module.bias.data.zero_() + elif isinstance(module, nn.Embedding): + module.weight.data.normal_(mean=0.0, std=std) + if module.padding_idx is not None: + module.weight.data[module.padding_idx].zero_() + + def encode_latent(self, waveform, waveform_length): + from diffusers.models.autoencoders.autoencoder_oobleck import OobleckDiagonalGaussianDistribution + + frame_num = torch.ceil(waveform_length / self.config.enc_kwargs["input_dim"]).to(torch.int32) + if self.config.patch_size != -1: + frame_num = torch.ceil(frame_num / self.config.patch_size) + h, y = self.encoder(waveform) + h = h.transpose(1, 2) + + posterior = OobleckDiagonalGaussianDistribution(h) + latent = posterior.sample() + latent = latent.transpose(1, 2) + return latent, frame_num + + def decode(self, latent, past_key_values=None, use_cache=False, stream_state=(None, None, None), last_chunk=False): + waveform, stream_state, past_key_values = self.decoder.low_level_reconstruct( + latent, + past_key_values=past_key_values, + use_cache=use_cache, + stream_state=stream_state, + last_chunk=last_chunk, + ) + return waveform, stream_state, past_key_values diff --git a/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni.py b/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni.py index 87728890b..462c20438 100644 --- a/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni.py +++ b/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni.py @@ -16,7 +16,23 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Ming-flash-omni-2.0 unified model (thinker + imagegen + talker).""" +"""Ming-flash-omni-2.0 thinker / image-gen wrapper. + +This class is the multimodal-registered entry point for Ming stages that +share the thinker's backbone: comprehension / text generation (`thinker`) +and diffusion conditioning for image generation (`imagegen`, not yet +implemented). + +The talker deliberately lives elsewhere. Upstream Ming hands text (not +hidden states) from the thinker to the talker, and the talker then +tokenises that string with its own Qwen2 tokenizer and runs an entirely +self-contained LLM + CFM + AudioVAE pipeline. Because it has no +multimodal inputs, it belongs in the non-MM-registered +`MingFlashOmniTalkerForConditionalGeneration` — routing it through +this wrapper would force it through vLLM's multimodal preprocess path +and trigger a hidden-size mismatch between the outer Ming config +(4096, thinker's LLM) and the talker's Qwen2 backbone (896). +""" from collections.abc import Iterable @@ -40,7 +56,10 @@ from vllm.sequence import IntermediateTensors from vllm_omni.model_executor.custom_process_mixin import CustomProcessMixin from vllm_omni.model_executor.models.output_templates import OmniOutput from vllm_omni.model_executor.models.utils import add_prefix_to_loaded_weights -from vllm_omni.transformers_utils.configs.ming_flash_omni import BailingMM2Config, MingFlashOmniConfig +from vllm_omni.transformers_utils.configs.ming_flash_omni import ( + BailingMM2Config, + MingFlashOmniConfig, +) from .ming_flash_omni_thinker import ( MingFlashOmniThinkerDummyInputsBuilder, @@ -63,7 +82,7 @@ class MingFlashOmniForConditionalGeneration( SupportsMRoPE, CustomProcessMixin, ): - """Unified Ming-flash-omni-2.0 model combining thinker, imagegen, and talker.""" + """Ming-flash-omni-2.0 thinker + image-gen wrapper.""" supports_multimodal = True requires_raw_input_tokens: bool = True @@ -75,19 +94,26 @@ class MingFlashOmniForConditionalGeneration( self.has_postprocess = False config = vllm_config.model_config.hf_config - - self.vllm_config = vllm_config self.config = config - - if isinstance(config, MingFlashOmniConfig): - thinker_config = config.thinker_config - else: - thinker_config = config - - self.thinker_config: BailingMM2Config = thinker_config self.model_stage = vllm_config.model_config.model_stage + if self.model_stage == "talker": + raise ValueError( + "MingFlashOmniForConditionalGeneration does not support " + "model_stage='talker'. Use " + "model_arch='MingFlashOmniTalkerForConditionalGeneration' " + "directly — the talker has a self-contained LLM that " + "tokenises text itself and does not need the multimodal " + "preprocess path. See stage_configs/ming_flash_omni.yaml " + "stage 1 and stage_configs/ming_flash_omni_tts.yaml." + ) + if self.model_stage == "thinker": + if isinstance(config, MingFlashOmniConfig): + thinker_config: BailingMM2Config = config.thinker_config + else: + thinker_config: BailingMM2Config = config + thinker_vllm_config = vllm_config.with_hf_config( thinker_config, architectures=["MingFlashOmniThinkerForConditionalGeneration"] ) @@ -97,8 +123,7 @@ class MingFlashOmniForConditionalGeneration( architectures=["MingFlashOmniThinkerForConditionalGeneration"], ) self.model = self.thinker - self.imagegen = None - self.talker = None + self.make_empty_intermediate_tensors = self.thinker.make_empty_intermediate_tensors elif self.model_stage == "imagegen": # TODO: Implement image generator stage @@ -106,22 +131,12 @@ class MingFlashOmniForConditionalGeneration( "Image generation stage is not yet implemented. Please use model_stage='thinker' for now." ) - elif self.model_stage == "talker": - # TODO: Implement talker (TTS) stage - raise NotImplementedError( - "Talker (TTS) stage is not yet implemented. Please use model_stage='thinker' for now." - ) - else: raise ValueError( - f"Invalid model_stage: {self.model_stage}. Must be one of: 'thinker', 'imagegen', 'talker'" + f"Invalid model_stage: {self.model_stage!r}. Must be one of: 'thinker', 'imagegen'. " + f"For the talker stage, use MingFlashOmniTalkerForConditionalGeneration directly." ) - # Set up intermediate tensors - self.make_empty_intermediate_tensors = ( - self.thinker.make_empty_intermediate_tensors if self.model_stage == "thinker" else lambda: None - ) - def forward( self, input_ids: torch.Tensor, @@ -154,7 +169,7 @@ class MingFlashOmniForConditionalGeneration( ): if hasattr(self.model, "sample"): return self.model.sample(logits, sampling_metadata) - raise NotImplementedError("sample method not available on current stage") + return None def get_mrope_input_positions(self, *args, **kwargs): if hasattr(self.model, "get_mrope_input_positions"): @@ -162,36 +177,9 @@ class MingFlashOmniForConditionalGeneration( raise NotImplementedError("get_mrope_input_positions not available on current stage") def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loaded_weights = set() - thinker_weights = [] - imagegen_weights = [] - talker_weights = [] - - for name, value in weights: - if name.startswith("thinker."): - thinker_weights.append((name, value)) - elif name.startswith("imagegen."): - imagegen_weights.append((name, value)) - elif name.startswith("talker."): - talker_weights.append((name, value)) - else: - # Weights without prefix go to thinker by default - thinker_weights.append((name, value)) - - if self.model_stage == "thinker" and thinker_weights: - # Remove "thinker." prefix before loading - thinker_weights_stripped = [ - (name.replace("thinker.", "", 1) if name.startswith("thinker.") else name, value) - for name, value in thinker_weights - ] - thinker_loaded = self.thinker.load_weights(thinker_weights_stripped) - thinker_loaded = add_prefix_to_loaded_weights(thinker_loaded, "thinker") - loaded_weights.update(thinker_loaded) - - # TODO: Load imagegen weights when implemented - # TODO: Load talker weights when implemented - - return loaded_weights + stripped = ((name.removeprefix("thinker."), value) for name, value in weights) + thinker_loaded = self.thinker.load_weights(stripped) + return add_prefix_to_loaded_weights(thinker_loaded, "thinker") def get_mm_mapping(self) -> MultiModelKeys: return MultiModelKeys.from_string_field( @@ -202,9 +190,7 @@ class MingFlashOmniForConditionalGeneration( @property def sampler(self): - if hasattr(self.model, "sampler"): - return self.model.sampler - return None + return getattr(self.model, "sampler", None) def embed_input_ids( self, diff --git a/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_talker.py b/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_talker.py new file mode 100644 index 000000000..08ed9e854 --- /dev/null +++ b/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_talker.py @@ -0,0 +1,586 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright 2025 The vLLM-Omni team. +# Copyright (c) Ant Group. All rights reserved. +# Adapted from: +# https://github.com/inclusionAI/Ming/blob/e58533db227031990c5a6864dcf5f08fb53ed0d2/modeling_bailing_talker.py +"""Ming-flash-omni-2.0 talker (TTS) stage model.""" + +from __future__ import annotations + +import glob as glob_module +import os +from collections.abc import Iterable +from dataclasses import dataclass +from functools import cached_property +from typing import Any + +import torch +import torch.nn as nn +from safetensors.torch import load_file +from transformers import AutoTokenizer, Qwen2Config, Qwen2Model +from transformers.utils.hub import cached_file +from vllm.config import VllmConfig +from vllm.logger import init_logger +from vllm.model_executor.models.utils import AutoWeightsLoader +from vllm.sequence import IntermediateTensors + +from vllm_omni.model_executor.custom_process_mixin import CustomProcessMixin +from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific +from vllm_omni.model_executor.models.output_templates import OmniOutput +from vllm_omni.transformers_utils.configs.ming_flash_omni import MingFlashOmniTalkerConfig + +from .audio_vae import AudioVAE, AudioVAEConfig +from .prompt_utils import DEFAULT_PROMPT as MING_DEFAULT_PROMPT +from .talker_module import CFM, Aggregator, DiT, MingAudioGenerator, build_tts_input +from .text_processing import segment_and_normalize +from .voice_presets import VoicePresetRegistry + +logger = init_logger(__name__) + + +@dataclass(slots=True) +class _GenerationParams: + """Resolved sampling / decoding parameters for one forward call.""" + + prompt: str + instruction: str | None + cfg: float + sigma: float + temperature: float + max_steps: int + use_zero_spk_emb: bool + max_text_length: int + use_static_cache: bool + stream_decode: bool + + +@dataclass(slots=True) +class _VoiceContext: + """Voice cloning inputs resolved from request info + presets.""" + + spk_emb: Any # list[Tensor] | Tensor | list[float] | None + prompt_text: str | None + prompt_wav_lat: torch.Tensor | None + prompt_wav_emb: torch.Tensor | None + already_projected: bool + + +class MingFlashOmniTalkerForConditionalGeneration(nn.Module, CustomProcessMixin): + """Ming-flash-omni-2.0 talker stage: text -> audio waveform. + + Uses Qwen2 LLM + CFM (Conditional Flow Matching with DiT) + Aggregator + in an autoregressive loop to produce continuous audio latents, then + AudioVAE decodes latents to waveforms. + """ + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + self.have_multimodal_outputs = True + self.has_preprocess = False + self.has_postprocess = False + + self.vllm_config = vllm_config + root_config = vllm_config.model_config.hf_config + + model_path = vllm_config.model_config.model + self._model_path = model_path + self.talker_dir = ( + os.path.join(model_path, "talker") if os.path.isdir(os.path.join(model_path, "talker")) else model_path + ) + + # When used standalone (model_arch=MingFlashOmniTalkerForConditionalGeneration), + # the root hf_config may be BailingMM2Config (thinker-only) due to model file structure + # Resolve talker config from talker/config.json in that case. + config = ( + root_config + if isinstance(root_config, MingFlashOmniTalkerConfig) + else self._resolve_talker_config(root_config, self.talker_dir, model_path) + ) + self.config = config + + self._standalone = prefix in ("", "talker") + if self._standalone: + self.allow_patterns_overrides = ["talker/model*.safetensors"] + self.fall_back_to_pt_during_load = False + + # LLM + llm_config = self._resolve_llm_config(config, self.talker_dir, model_path) + llm_config._attn_implementation = "sdpa" + self.llm_config = llm_config + self.hidden_size = llm_config.hidden_size + self.latent_dim = config.latent_dim + self.patch_size = config.patch_size + self.his_patch_size = config.history_patch_size + self.cfg_strength = config.cfg_strength + + self.model = Qwen2Model(llm_config) + self.cfm = CFM( + DiT(llm_input_dim=self.hidden_size, **config.flowmodel), + steps=config.steps, + ) + self.aggregator = Aggregator(llm_input_dim=self.hidden_size, **config.aggregator) + self.stop_head = nn.Linear(self.hidden_size, 2, bias=True) + # CAMPPlus 192-dim -> hidden + self.spk_head = nn.Linear(192, self.hidden_size, bias=True) + + # AudioVAE + self.audio_vae, self._vae_weight_source = self._init_audio_vae(config, self.talker_dir, model_path) + + self._use_cuda_graphs = not vllm_config.model_config.enforce_eager + + self.audio_generator = MingAudioGenerator( + config=self.config, + llm_config=self.llm_config, + model=self.model, + cfm=self.cfm, + aggregator=self.aggregator, + stop_head=self.stop_head, + audio_vae=self.audio_vae, + patch_size=self.patch_size, + his_patch_size=self.his_patch_size, + latent_dim=self.latent_dim, + cfg_strength=self.cfg_strength, + use_cuda_graphs=self._use_cuda_graphs, + ) + self.voice_presets = VoicePresetRegistry( + talker_dir=self.talker_dir, + model_path=self._model_path, + download_dir=vllm_config.load_config.download_dir, + audio_vae=self.audio_vae, + aggregator=self.aggregator, + spk_head=self.spk_head, + patch_size=self.patch_size, + ) + + @property + def device(self) -> torch.device: + return next(self.model.parameters()).device + + @property + def dtype(self) -> torch.dtype: + return next(self.model.parameters()).dtype + + @cached_property + def tokenizer(self): + # Lazy Qwen2 tokenizer resolution: + # 1. Try local dirs first (talker/llm, talker, and then model root). + # 2. HF repo-id fallback: talker/llm is the canonical tokenizer location. + candidates = (os.path.join(self.talker_dir, "llm"), self.talker_dir, self._model_path) + for path in candidates: + if os.path.isdir(path): + try: + logger.debug("Resolving talker tokenizer from local dir %s", path) + return AutoTokenizer.from_pretrained(path, trust_remote_code=True) + except Exception: + continue + for subfolder in ("talker/llm", "llm"): + try: + logger.debug("Resolving talker tokenizer from HF subfolder %s", subfolder) + return AutoTokenizer.from_pretrained(self._model_path, subfolder=subfolder, trust_remote_code=True) + except Exception: + continue + logger.debug("Falling back to raw model_path tokenizer resolution") + return AutoTokenizer.from_pretrained(self._model_path, trust_remote_code=True) + + @staticmethod + def _resolve_talker_config(config, talker_dir: str, model_path: str) -> MingFlashOmniTalkerConfig: + """Resolve MingFlashOmniTalkerConfig when the root config is not one. + + This happens in standalone TTS mode where hf_config is BailingMM2Config. + """ + # If the root config wraps a talker_config, use it + talker_config = getattr(config, "talker_config", None) + if isinstance(talker_config, MingFlashOmniTalkerConfig): + return talker_config + + # Try loading from talker/config.json + if os.path.isdir(talker_dir): + try: + resolved = MingFlashOmniTalkerConfig.from_pretrained(talker_dir) + logger.info("Resolved talker config from %s", talker_dir) + return resolved + except Exception: + pass + + try: + resolved = MingFlashOmniTalkerConfig.from_pretrained(model_path, subfolder="talker", trust_remote_code=True) + logger.info("Resolved talker config from %s/talker (HF hub)", model_path) + return resolved + except Exception as e: + raise ValueError( + f"Cannot resolve MingFlashOmniTalkerConfig. The root config " + f"is {type(config).__name__}, and talker/config.json was not " + f"found at {talker_dir} or via HF hub: {e}" + ) from e + + @staticmethod + def _resolve_llm_config(config: MingFlashOmniTalkerConfig, talker_dir: str, model_path: str) -> Qwen2Config: + """Resolve the Qwen2 LLM config for the talker backbone.""" + + if config.llm_config is not None: + return Qwen2Config(**config.llm_config) if isinstance(config.llm_config, dict) else config.llm_config + + # Try local talker/llm directory + llm_dir = os.path.join(talker_dir, "llm") + if os.path.isdir(llm_dir): + return Qwen2Config.from_pretrained(llm_dir) + + # HF hub fallback + for subfolder in ("talker/llm", "llm"): + try: + return Qwen2Config.from_pretrained(model_path, subfolder=subfolder, trust_remote_code=True) + except Exception: + continue + + raise ValueError( + f"Cannot find talker LLM config at {llm_dir}. " + "Either provide llm_config in MingFlashOmniTalkerConfig or " + "ensure the model path contains talker/llm/config.json." + ) + + @staticmethod + def _init_audio_vae( + config: MingFlashOmniTalkerConfig, talker_dir: str, model_path: str + ) -> tuple[AudioVAE | None, str | tuple[str, str] | None]: + """Initialize AudioVAE and return (vae, weight_source). + + weight_source is either a local directory path (str) or an + (repo_id, subfolder) tuple for HF hub downloads, or None. + """ + vae_path = config.audio_vae_path or os.path.join(talker_dir, "vae") + + # Try local directory first + if os.path.isdir(vae_path): + try: + vae_config = AudioVAEConfig.from_pretrained(vae_path) + vae = AudioVAE(vae_config) + logger.info("Initialized AudioVAE from %s (sr=%d)", vae_path, vae_config.sample_rate) + return vae, vae_path + except Exception as e: + logger.warning("Failed to initialize AudioVAE from %s: %s", vae_path, e) + return None, None + + # HF hub fallback + for subfolder in ("talker/vae", "vae"): + try: + vae_config = AudioVAEConfig.from_pretrained(model_path, subfolder=subfolder, trust_remote_code=True) + vae = AudioVAE(vae_config) + logger.info(f"Initialized AudioVAE from {model_path}/{subfolder}") + return vae, (model_path, subfolder) + except Exception: + continue + + logger.info("AudioVAE not found at %s; waveform decoding unavailable", vae_path) + return None, None + + def compute_logits(self, hidden_states: torch.Tensor, sampling_metadata=None) -> torch.Tensor | None: + return None + + def sample(self, logits: torch.Tensor, sampling_metadata): + return None + + def embed_input_ids( + self, + input_ids: torch.Tensor, + multimodal_embeddings=None, + is_multimodal=None, + ) -> torch.Tensor: + return self.model.get_input_embeddings()(input_ids) + + def make_empty_intermediate_tensors( + self, batch_size: int, dtype: torch.dtype, device: torch.device + ) -> IntermediateTensors | None: + return None + + def get_dummy_runtime_additional_information(self, num_reqs: int) -> list[dict[str, object]]: + info: dict[str, object] = {"text": "dummy", "use_zero_spk_emb": True, "max_steps": 1} + return [info for _ in range(num_reqs)] + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + runtime_additional_information: list[dict] | None = None, + **kwargs, + ) -> OmniOutput: + """Run TTS generation and return audio output. + + The full autoregressive generation loop is executed inside this method. + """ + additional_info = self._extract_additional_info(runtime_additional_information) + params = self._resolve_generation_params(additional_info) + voice = self._resolve_voice(additional_info) + + latents = self._generate_latents( + input_ids=input_ids, + inputs_embeds=inputs_embeds, + text=additional_info.get("text", ""), + params=params, + voice=voice, + ) + return self._decode_to_output(latents, stream_decode=params.stream_decode) + + @staticmethod + def _extract_additional_info( + runtime_additional_information: list[dict] | None, + ) -> dict[str, Any]: + if runtime_additional_information and len(runtime_additional_information) > 0: + return runtime_additional_information[0] or {} + return {} + + def _resolve_generation_params(self, additional_info: dict[str, Any]) -> _GenerationParams: + # "omni" : thinker -> talker hand-off with hardcoded defaults + # "instruct": standalone TTS with caller-supplied sampling knobs + ming_task = additional_info.get("ming_task", "instruct") + + if ming_task == "omni": + prompt = MING_DEFAULT_PROMPT + instruction = None + use_zero_spk_emb = additional_info.get("spk_emb") is None + cfg = 2.0 + sigma = 0.25 + temperature = 0.0 + max_steps = 200 + else: + prompt = additional_info.get("prompt", MING_DEFAULT_PROMPT) + instruction = additional_info.get("instruction", None) + use_zero_spk_emb = additional_info.get("use_zero_spk_emb", False) + cfg = additional_info.get("cfg", self.cfg_strength) + sigma = additional_info.get("sigma", 0.25) + temperature = additional_info.get("temperature", 0.0) + max_steps = int(additional_info.get("max_steps", additional_info.get("max_decode_steps", 200))) + + return _GenerationParams( + prompt=prompt, + instruction=instruction, + cfg=cfg, + sigma=sigma, + temperature=temperature, + max_steps=max_steps, + use_zero_spk_emb=use_zero_spk_emb, + max_text_length=int(additional_info.get("max_text_length", 50)), + use_static_cache=bool(additional_info.get("use_static_cache", True)), + stream_decode=bool(additional_info.get("stream_decode", True)), + ) + + def _resolve_voice(self, additional_info: dict[str, Any]) -> _VoiceContext: + spk_emb = additional_info.get("spk_emb", None) + prompt_text = additional_info.get("prompt_text", None) + prompt_wav_lat = additional_info.get("prompt_wav_lat", None) + prompt_wav_emb = additional_info.get("prompt_wav_emb", None) + already_projected = False + + voice_name = additional_info.get("voice_name", None) + if voice_name and spk_emb is None and voice_name in self.voice_presets: + preset = self.voice_presets.get(voice_name) or {} + prompt_wav_lat = preset.get("prompt_wav_lat") + prompt_wav_emb = preset.get("prompt_wav_emb") + spk_emb = preset.get("spk_emb") + already_projected = True + if prompt_text is None: + prompt_text = preset.get("prompt_text") + + return _VoiceContext( + spk_emb=spk_emb, + prompt_text=prompt_text, + prompt_wav_lat=prompt_wav_lat, + prompt_wav_emb=prompt_wav_emb, + already_projected=already_projected, + ) + + def _project_spk_emb( + self, spk_emb: Any, already_projected: bool, use_zero_spk_emb: bool + ) -> list[torch.Tensor] | None: + if spk_emb is None: + if use_zero_spk_emb: + return [torch.zeros(1, self.hidden_size, device=self.device, dtype=self.dtype)] + return None + + if already_projected: + return spk_emb if isinstance(spk_emb, list) else [spk_emb] + + if isinstance(spk_emb, torch.Tensor): + tensors = [spk_emb] + elif isinstance(spk_emb, list) and spk_emb and isinstance(spk_emb[0], (int, float)): + tensors = [torch.tensor(spk_emb, dtype=self.dtype).unsqueeze(0)] + elif isinstance(spk_emb, list): + tensors = spk_emb + else: + tensors = [spk_emb] + return [self.spk_head(t.to(device=self.device, dtype=self.dtype)) for t in tensors] + + def _generate_latents( + self, + *, + input_ids: torch.Tensor, + inputs_embeds: torch.Tensor | None, + text: str, + params: _GenerationParams, + voice: _VoiceContext, + ) -> list[torch.Tensor]: + generator = self.audio_generator + + if inputs_embeds is not None: + # Caller pre-built embeddings — run a single AR pass. + return generator.generate_latents( + inputs_embeds=inputs_embeds, + prompt_wav_lat=voice.prompt_wav_lat, + max_steps=params.max_steps, + cfg=params.cfg, + sigma=params.sigma, + temperature=params.temperature, + use_static_cache=params.use_static_cache, + ) + + spk_emb = self._project_spk_emb(voice.spk_emb, voice.already_projected, params.use_zero_spk_emb) + text_segments = segment_and_normalize(text, max_length=params.max_text_length) if text else [] + + if not text_segments: + # vLLM passes 1D input_ids; Qwen2Model expects (batch, seq). + inputs_embeds = self.model.get_input_embeddings()(input_ids.to(self.device)).unsqueeze(0) + return generator.generate_latents( + inputs_embeds=inputs_embeds, + prompt_wav_lat=voice.prompt_wav_lat, + max_steps=params.max_steps, + cfg=params.cfg, + sigma=params.sigma, + temperature=params.temperature, + use_static_cache=params.use_static_cache, + ) + + all_latents: list[torch.Tensor] = [] + for segment in text_segments: + seg_embeds, _ = build_tts_input( + tokenizer=self.tokenizer, + embed_tokens=self.model.get_input_embeddings(), + device=self.device, + dtype=torch.bfloat16, + text=segment, + prompt=params.prompt, + spk_emb=spk_emb, + instruction=params.instruction, + prompt_text=voice.prompt_text, + prompt_wav_emb=voice.prompt_wav_emb, + ) + effective_max_steps = generator.duration_capped_steps(len(segment), params.max_steps) + all_latents.extend( + generator.generate_latents( + inputs_embeds=seg_embeds, + prompt_wav_lat=voice.prompt_wav_lat, + max_steps=effective_max_steps, + cfg=params.cfg, + sigma=params.sigma, + temperature=params.temperature, + use_static_cache=params.use_static_cache, + ) + ) + return all_latents + + def _decode_to_output(self, latents: list[torch.Tensor], *, stream_decode: bool) -> OmniOutput: + multimodal_outputs: dict[str, Any] = {} + if latents and self.audio_vae is not None: + waveform = self.audio_generator.decode_to_waveform(latents, stream_decode=stream_decode) + if not stream_decode: + waveform = self.audio_generator.trim_trailing_silence(waveform) + multimodal_outputs["audio"] = waveform.detach().float().cpu() + multimodal_outputs["sr"] = torch.tensor(self.audio_vae.config.sample_rate) + elif latents: + all_lat = torch.cat(latents, dim=1) + multimodal_outputs["audio_latents"] = all_lat.detach().float().cpu() + + return OmniOutput(text_hidden_states=None, multimodal_outputs=multimodal_outputs) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + """Load weights for all talker components. + + The talker's HF checkpoint (talker/model.safetensors) stores + weights with prefixes matching this module's submodule names directly. + And AudioVAE weights live in a separate file under talker/vae/ + """ + # Standalone: bypass the default loader's iterator (torch.load on + # .safetensors crashes) and read talker/model*.safetensors directly. + if self._standalone: + weights = self._iter_talker_safetensors() + + loader = AutoWeightsLoader( + self, + skip_prefixes=["audio_vae."], # loaded separately + skip_substrs=["rotary_embed.inv_freq"], # non-persistent buffer + ) + loaded = loader.load_weights(weights) + logger.info("Loaded %d talker weights from checkpoint", len(loaded)) + + if self.audio_vae is not None and self._vae_weight_source is not None: + loaded.update(self._load_vae_weights()) + + # Register voice presets after all weights (incl. VAE) are loaded. + try: + self.voice_presets.load_presets_from_manifest(device=self.device, dtype=self.dtype) + except Exception as e: # pragma: no cover — best-effort + logger.warning("Voice preset loading failed (non-fatal): %s", e) + + return loaded + + def _iter_talker_safetensors(self) -> Iterable[tuple[str, torch.Tensor]]: + """Yield (name, tensor) pairs from talker/model*.safetensors.""" + model_path = self._model_path + # Try local path first + for candidate in (os.path.join(model_path, "talker"), model_path): + sf_files = sorted(glob_module.glob(os.path.join(candidate, "model*.safetensors"))) + if sf_files: + for sf_path in sf_files: + yield from load_file(sf_path, device="cpu").items() + return + + # HF hub fallback: download only the talker checkpoint files + model_root = download_weights_from_hf_specific( + model_path, + self.vllm_config.load_config.download_dir, + allow_patterns=["talker/model*.safetensors"], + ) + talker_dir = os.path.join(model_root, "talker") + sf_files = sorted(glob_module.glob(os.path.join(talker_dir, "model*.safetensors"))) + if not sf_files: + raise RuntimeError(f"No talker safetensors found under {model_root}. Expected talker/model*.safetensors.") + for sf_path in sf_files: + yield from load_file(sf_path, device="cpu").items() + + def _load_vae_weights(self) -> set[str]: + """Load AudioVAE weights from talker/vae/model.safetensors.""" + if self.audio_vae is None or self._vae_weight_source is None: + return set() + + # Resolve safetensors file paths from the weight source + safetensors_files: list[str] = [] + source = self._vae_weight_source + if isinstance(source, str): + # Local directory path + safetensors_files = sorted(glob_module.glob(os.path.join(source, "*.safetensors"))) + elif isinstance(source, tuple): + # (repo_id, subfolder) for HF hub + repo_id, subfolder = source + for filename in ("model.safetensors", "diffusion_pytorch_model.safetensors"): + try: + cached = cached_file(repo_id, filename, subfolder=subfolder) + except Exception: + cached = None + if cached is not None: + safetensors_files.append(cached) + break + + if not safetensors_files: + logger.warning("No AudioVAE safetensors files found for source=%s", source) + return set() + + vae_state_keys = set(self.audio_vae.state_dict().keys()) + vae_loader = AutoWeightsLoader(self.audio_vae) + loaded: set[str] = set() + for sf_path in safetensors_files: + file_weights = load_file(sf_path, device="cpu") + matched = ((name, tensor) for name, tensor in file_weights.items() if name in vae_state_keys) + loaded.update(f"audio_vae.{name}" for name in vae_loader.load_weights(matched)) + + logger.info("Loaded %d AudioVAE weights from %s", len(loaded), source) + return loaded diff --git a/vllm_omni/model_executor/models/ming_flash_omni/prompt_utils.py b/vllm_omni/model_executor/models/ming_flash_omni/prompt_utils.py new file mode 100644 index 000000000..4271114bc --- /dev/null +++ b/vllm_omni/model_executor/models/ming_flash_omni/prompt_utils.py @@ -0,0 +1,50 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright 2025 The vLLM-Omni team. +# Copyright (c) Ant Group. All rights reserved. +# Adapted from Ming repo's usage cookbook: +# https://github.com/inclusionAI/Ming/blob/3954fcb880ff5e61ff128bcf7f1ec344d46a6fe3/cookbook.ipynb +"""Shared prompt-building helpers for Ming-flash-omni standalone talker.""" + +import copy +import json +from typing import Any + +DEFAULT_PROMPT = "Please generate speech based on the following description.\n" + +BASE_CAPTION_TEMPLATE: dict[str, Any] = { + "audio_sequence": [ + { + "序号": 1, + "说话人": "speaker_1", + "方言": None, + "风格": None, + "语速": None, + "基频": None, + "音量": None, + "情感": None, + "BGM": { + "Genre": None, + "Mood": None, + "Instrument": None, + "Theme": None, + "ENV": None, + "SNR": None, + }, + "IP": None, + } + ] +} + + +def create_instruction(user_input: dict[str, Any]) -> str: + """Return a JSON caption string for ``audio_sequence[0]``. + + Only keys already present on the base template are merged in; unknown + keys are silently ignored to keep the output schema stable. + """ + caption = copy.deepcopy(BASE_CAPTION_TEMPLATE) + item = caption["audio_sequence"][0] + for key, value in user_input.items(): + if key in item: + item[key] = value + return json.dumps(caption, ensure_ascii=False) diff --git a/vllm_omni/model_executor/models/ming_flash_omni/spk_embedding.py b/vllm_omni/model_executor/models/ming_flash_omni/spk_embedding.py new file mode 100644 index 000000000..68dbfe650 --- /dev/null +++ b/vllm_omni/model_executor/models/ming_flash_omni/spk_embedding.py @@ -0,0 +1,44 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright 2025 The vLLM-Omni team. +# Copyright (c) Ant Group. All rights reserved. +# Ported from: +# https://github.com/inclusionAI/Ming/blob/e58533db227031990c5a6864dcf5f08fb53ed0d2/modeling_bailing_talker.py + +from __future__ import annotations + +import torch + + +class SpkembExtractor: + """CAMPPlus ONNX-based speaker embedding extractor (runs on CPU).""" + + def __init__(self, campplus_model: str, target_sr: int = 16000): + import onnxruntime + import torchaudio.compliance.kaldi as kaldi + + self.kaldi = kaldi + option = onnxruntime.SessionOptions() + option.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_ALL + option.intra_op_num_threads = 2 + self.campplus_session = onnxruntime.InferenceSession( + campplus_model, sess_options=option, providers=["CPUExecutionProvider"] + ) + self.target_sr = target_sr + + def _extract_spk_embedding(self, speech): + feat = self.kaldi.fbank(speech, num_mel_bins=80, dither=0, sample_frequency=16000) + feat = feat - feat.mean(dim=0, keepdim=True) + embedding = ( + self.campplus_session.run( + None, + {self.campplus_session.get_inputs()[0].name: feat.unsqueeze(dim=0).cpu().numpy()}, + )[0] + .flatten() + .tolist() + ) + embedding = torch.tensor([embedding]) + return embedding + + def __call__(self, waveform, **kwargs) -> torch.Tensor | None: + spk_emb = self._extract_spk_embedding(waveform) + return spk_emb diff --git a/vllm_omni/model_executor/models/ming_flash_omni/talker_module.py b/vllm_omni/model_executor/models/ming_flash_omni/talker_module.py new file mode 100644 index 000000000..80acbaad0 --- /dev/null +++ b/vllm_omni/model_executor/models/ming_flash_omni/talker_module.py @@ -0,0 +1,1151 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright 2025 The vLLM-Omni team. +# Copyright (c) Ant Group. All rights reserved. +# Ported from: +# https://github.com/inclusionAI/Ming/blob/e58533db227031990c5a6864dcf5f08fb53ed0d2/talker_module/dit.py +# +# Ported from: +# https://github.com/inclusionAI/Ming/blob/e58533db227031990c5a6864dcf5f08fb53ed0d2/talker_module/modules.py +# Ported from: +# https://github.com/inclusionAI/Ming/blob/e58533db227031990c5a6864dcf5f08fb53ed0d2/talker_module/cfm.py + +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# Partial of the following source code +# is licensed under the license found in the +# LICENSE file in the root directory of this source tree. +# -------------------------------------------------------- +# References: +# GLIDE: https://github.com/openai/glide-text2im +# MAE: https://github.com/facebookresearch/mae/blob/main/models_mae.py +# -------------------------------------------------------- +import logging +import math +from functools import cached_property +from queue import Queue +from threading import Lock + +import torch +import torch.nn as nn +import torch.nn.functional as F +from transformers import PreTrainedTokenizerBase, Qwen2Config, Qwen2Model, StaticCache +from vllm.logger import init_logger +from x_transformers.x_transformers import RotaryEmbedding, apply_rotary_pos_emb + +from .audio_vae import AudioVAE + +logger = init_logger(__name__) + + +######################################################################## +# DiT Modules +# Ported from: +# https://github.com/inclusionAI/Ming/blob/e58533db227031990c5a6864dcf5f08fb53ed0d2/talker_module/modules.py +# Ported from: +# https://github.com/inclusionAI/Ming/blob/e58533db227031990c5a6864dcf5f08fb53ed0d2/talker_module/dit.py +######################################################################## + + +class RMSNorm(nn.Module): + def __init__(self, dim: int, eps: float = 1e-6): + super().__init__() + self.eps = eps + self.weight = nn.Parameter(torch.ones(dim)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if self.weight.dtype in [torch.float16, torch.bfloat16]: + x = x.to(self.weight.dtype) + x = F.rms_norm(x, normalized_shape=(x.shape[-1],), weight=self.weight, eps=self.eps) + return x + + +class FeedForward(nn.Module): + def __init__( + self, dim: int, dim_out: int | None = None, mult: float = 4, dropout: float = 0.0, approximate: str = "none" + ): + super().__init__() + inner_dim = int(dim * mult) + dim_out = dim_out if dim_out is not None else dim + + activation = nn.GELU(approximate=approximate) + project_in = nn.Sequential(nn.Linear(dim, inner_dim), activation) + self.ff = nn.Sequential(project_in, nn.Dropout(dropout), nn.Linear(inner_dim, dim_out)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.ff(x) + + +class Attention(nn.Module): + def __init__( + self, + dim: int, + heads: int = 8, + dim_head: int = 64, + dropout: float = 0.0, + qk_norm: str | None = None, + pe_attn_head: int | None = None, + attn_mask_enabled: bool = True, + ): + super().__init__() + self.dim = dim + self.heads = heads + self.inner_dim = dim_head * heads + self.dropout = dropout + + self.to_q = nn.Linear(dim, self.inner_dim) + self.to_k = nn.Linear(dim, self.inner_dim) + self.to_v = nn.Linear(dim, self.inner_dim) + if qk_norm is None: + self.q_norm = None + self.k_norm = None + elif qk_norm == "rms_norm": + self.q_norm = RMSNorm(dim_head) + self.k_norm = RMSNorm(dim_head) + else: + raise ValueError(f"Unimplemented qk_norm: {qk_norm}") + + self.to_out = nn.ModuleList([]) + self.to_out.append(nn.Linear(self.inner_dim, dim)) + self.to_out.append(nn.Dropout(dropout)) + + self.pe_attn_head = pe_attn_head + self.attn_mask_enabled = attn_mask_enabled + + def forward( + self, + x: torch.Tensor, + mask: torch.Tensor | None = None, + rope: tuple[torch.Tensor, torch.Tensor | None] | None = None, + ) -> torch.Tensor: + batch_size = x.shape[0] + + query = self.to_q(x) + key = self.to_k(x) + value = self.to_v(x) + + inner_dim = key.shape[-1] + head_dim = inner_dim // self.heads + query = query.view(batch_size, -1, self.heads, head_dim).transpose(1, 2) + key = key.view(batch_size, -1, self.heads, head_dim).transpose(1, 2) + value = value.view(batch_size, -1, self.heads, head_dim).transpose(1, 2) + + if self.q_norm is not None: + query = self.q_norm(query) + if self.k_norm is not None: + key = self.k_norm(key) + + if rope is not None: + freqs, xpos_scale = rope + q_xpos_scale, k_xpos_scale = (xpos_scale, xpos_scale**-1.0) if xpos_scale is not None else (1.0, 1.0) + + if self.pe_attn_head is not None: + on = self.pe_attn_head + query[:, :on, :, :] = apply_rotary_pos_emb(query[:, :on, :, :], freqs, q_xpos_scale) + key[:, :on, :, :] = apply_rotary_pos_emb(key[:, :on, :, :], freqs, k_xpos_scale) + else: + query = apply_rotary_pos_emb(query, freqs, q_xpos_scale) + key = apply_rotary_pos_emb(key, freqs, k_xpos_scale) + + if self.attn_mask_enabled and mask is not None: + valid_sample_indices = mask.any(dim=1) + final_output = torch.zeros_like(query).to(query.device) + + attn_mask = mask[valid_sample_indices] + query = query[valid_sample_indices] + key = key[valid_sample_indices] + value = value[valid_sample_indices] + attn_mask = attn_mask.unsqueeze(1).unsqueeze(1) + attn_mask = attn_mask.expand(valid_sample_indices.sum().item(), self.heads, query.shape[-2], key.shape[-2]) + else: + attn_mask = None + + x = F.scaled_dot_product_attention(query, key, value, attn_mask=attn_mask, dropout_p=0.0, is_causal=False) + if self.attn_mask_enabled and mask is not None: + final_output[valid_sample_indices] = x + x = final_output + + x = x.transpose(1, 2).reshape(batch_size, -1, self.heads * head_dim) + x = x.to(query.dtype) + + x = self.to_out[0](x) + x = self.to_out[1](x) + + if mask is not None: + mask = mask.unsqueeze(-1) + x = x.masked_fill(~mask, 0.0) + + return x + + +class DiTBlock(nn.Module): + """A DiT block with pre-norm and residual connections.""" + + def __init__( + self, + hidden_size: int, + num_heads: int, + mlp_ratio: float = 4.0, + dropout: float = 0.1, + qk_norm: str | None = None, + pe_attn_head: int | None = None, + attn_mask_enabled: bool = True, + **kwargs, + ): + super().__init__() + self.norm1 = RMSNorm(hidden_size) + self.attn = Attention( + dim=hidden_size, + heads=num_heads, + dim_head=hidden_size // num_heads, + dropout=dropout, + qk_norm=qk_norm, + pe_attn_head=pe_attn_head, + attn_mask_enabled=attn_mask_enabled, + ) + self.norm2 = RMSNorm(hidden_size) + self.mlp = FeedForward(dim=hidden_size, mult=mlp_ratio, dropout=dropout, approximate="tanh") + + def forward( + self, + x: torch.Tensor, + mask: torch.Tensor | None, + rope: tuple[torch.Tensor, torch.Tensor | None] | None, + ) -> torch.Tensor: + x = x + self.attn(self.norm1(x), mask=mask, rope=rope) + x = x + self.mlp(self.norm2(x)) + return x + + +class FinalLayer(nn.Module): + """The final layer of DiT.""" + + def __init__(self, hidden_size: int, out_channels: int): + super().__init__() + self.norm_final = RMSNorm(hidden_size) + self.linear = nn.Linear(hidden_size, out_channels, bias=True) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.norm_final(x) + x = self.linear(x) + return x + + +class SinusPositionEmbedding(nn.Module): + def __init__(self, dim: int): + super().__init__() + self.dim = dim + + def forward(self, x: torch.Tensor, scale: float = 1000) -> torch.Tensor: + device = x.device + half_dim = self.dim // 2 + emb = math.log(10000) / (half_dim - 1) + emb = torch.exp(torch.arange(half_dim, device=device).float() * -emb) + emb = scale * x.unsqueeze(1) * emb.unsqueeze(0) + emb = torch.cat((emb.sin(), emb.cos()), dim=-1) + return emb + + +class TimestepEmbedder(nn.Module): + def __init__(self, dim: int, freq_embed_dim: int = 256): + super().__init__() + self.time_embed = SinusPositionEmbedding(freq_embed_dim) + self.time_mlp = nn.Sequential(nn.Linear(freq_embed_dim, dim), nn.SiLU(), nn.Linear(dim, dim)) + + def forward(self, timestep: torch.Tensor) -> torch.Tensor: + time_hidden = self.time_embed(timestep) + time_hidden = time_hidden.to(timestep.dtype) + time = self.time_mlp(time_hidden) + return time + + +class CondEmbedder(nn.Module): + """Embeds LLM hidden states with optional CFG dropout.""" + + def __init__(self, input_feature_size: int, hidden_size: int): + super().__init__() + self.cond_embedder = nn.Linear(input_feature_size, hidden_size) + + def forward(self, llm_cond: torch.Tensor) -> torch.Tensor: + return self.cond_embedder(llm_cond) + + +class DiT(nn.Module): + """Diffusion model with a Transformer backbone for audio latent generation.""" + + def __init__( + self, + in_channels: int = 64, + hidden_size: int = 1024, + depth: int = 28, + num_heads: int = 16, + mlp_ratio: float = 4.0, + llm_cond_dim: int = 896, + **kwargs, + ): + super().__init__() + self.in_channels = in_channels + self.out_channels = in_channels + self.num_heads = num_heads + + self.t_embedder = TimestepEmbedder(hidden_size) + self.x_embedder = nn.Linear(in_channels, hidden_size) + self.c_embedder = CondEmbedder(llm_cond_dim, hidden_size) + if "spk_dim" in kwargs: + self.spk_embedder = nn.Linear(kwargs["spk_dim"], hidden_size) + else: + self.spk_embedder = None + self.hidden_size = hidden_size + + self.rotary_embed = RotaryEmbedding(hidden_size // num_heads) + + self.blocks = nn.ModuleList( + [DiTBlock(hidden_size, num_heads, mlp_ratio=mlp_ratio, **kwargs) for _ in range(depth)] + ) + self.final_layer = FinalLayer(hidden_size, self.out_channels) + + def forward( + self, + x: torch.Tensor, + t: torch.Tensor, + c: torch.Tensor, + latent_history: torch.Tensor, + spk_emb: torch.Tensor | None = None, + ) -> torch.Tensor: + x = torch.cat([latent_history, x], dim=1) + x = self.x_embedder(x) + t = self.t_embedder(t).unsqueeze(1) + c = self.c_embedder(c) + y = t + c + if spk_emb is None: + assert self.spk_embedder is None + x = torch.cat([y, x], dim=1) + else: + x = torch.cat([self.spk_embedder(spk_emb), y, x], dim=1) + rope = self.rotary_embed.forward_from_seq_len(x.shape[1]) + + for block in self.blocks: + x = block(x, None, rope) + x = self.final_layer(x) + return x + + def forward_with_cfg( + self, + x: torch.Tensor, + t: torch.Tensor, + c: torch.Tensor, + latent_history: torch.Tensor, + spk_emb: torch.Tensor | None = None, + ) -> torch.Tensor: + """Forward with classifier-free guidance (doubles batch for CFG).""" + x = torch.cat([x, x], dim=0) + latent_history = torch.cat([latent_history, latent_history], dim=0) + fake_latent = torch.zeros_like(c) + c = torch.cat([c, fake_latent], dim=0) + if t.ndim == 0: + t = t.repeat(x.shape[0]) + if spk_emb is not None: + spk_emb = torch.cat([spk_emb, spk_emb], dim=0) + model_out = self.forward(x, t, c, latent_history, spk_emb) + return model_out[:, -x.shape[1] :, :] + + +######################################################################################### +# CFM +# Ported from: +# https://github.com/inclusionAI/Ming/blob/e58533db227031990c5a6864dcf5f08fb53ed0d2/talker_module/cfm.py +######################################################################################### + + +def get_epss_timesteps(n, device, dtype): + dt = 1 / 32 + predefined_timesteps = { + 5: [0, 2, 4, 8, 16, 32], + 6: [0, 2, 4, 6, 8, 16, 32], + 7: [0, 2, 4, 6, 8, 16, 24, 32], + 10: [0, 2, 4, 6, 8, 12, 16, 20, 24, 28, 32], + 12: [0, 2, 4, 6, 8, 10, 12, 14, 16, 20, 24, 28, 32], + 16: [0, 1, 2, 3, 4, 5, 6, 7, 8, 10, 12, 14, 16, 20, 24, 28, 32], + } + t = predefined_timesteps.get(n, []) + if not t: + return torch.linspace(0, 1, n + 1, device=device, dtype=dtype) + return dt * torch.tensor(t, device=device, dtype=dtype) + + +class CFM(nn.Module): + """Conditional Flow Matching module for audio latent generation.""" + + def __init__(self, model: nn.Module, steps: int = 10, sway_sampling_coef: float | None = -1.0): + """ + Args: + model: DiT used for the velocity prediction. + steps: number of integration steps per sample call. + sway_sampling_coef: coefficient used to skew the integration + grid towards low-noise timesteps. Defaults to -1.0 which + packs more steps near t=0, where prediction error is highest. + Set to `None` to use the linear grid as-is. + """ + super().__init__() + self.model = model + self.steps = steps + self.sway_sampling_coef = sway_sampling_coef + + @torch.no_grad() + def sample( + self, + llm_cond: torch.Tensor, + lat_cond: torch.Tensor, + y0: torch.Tensor, + t: torch.Tensor, + sde_args: torch.Tensor, + sde_rnd: torch.Tensor, + ): + """Sample audio latent via ODE/SDE integration with CFG. + + Args: + llm_cond: LLM hidden state (B, 1, hidden_size) + lat_cond: latent history (B, his_patch_size, latent_dim) + y0: initial noise (B, patch_size, latent_dim) + t: timesteps from get_epss_timesteps + sde_args: [cfg_strength, sigma, temperature] + sde_rnd: random noise for SDE steps (steps, B, patch_size, latent_dim) + """ + + def fn(fn_t, x): + pred_cfg = self.model.forward_with_cfg(x, fn_t, llm_cond, lat_cond, None) + pred, null_pred = torch.chunk(pred_cfg, 2, dim=0) + return pred + (pred - null_pred) * sde_args[0] + + if self.sway_sampling_coef is not None: + t = t + self.sway_sampling_coef * (torch.cos(torch.pi / 2 * t) - 1 + t) + + for step in range(self.steps): + dt = t[step + 1] - t[step] + y0 = y0 + fn(t[step], y0) * dt + y0 = y0 + sde_args[1] * (sde_args[2] ** 0.5) * (dt.abs() ** 0.5) * sde_rnd[step] + + return y0 + + +class CFMGraphExecutor: + """CUDA graph-accelerated executor for CFM + Aggregator + StopHead pipeline.""" + + def __init__(self, config, cfm, aggregator, stop_head: nn.Linear): + self.config = config + self.cfm = cfm + self.aggregator = aggregator + self.stop_head = stop_head + self.initialized = False + + self.last_hidden_state_placeholder = None + self.his_lat_placeholder = None + self.randn_like_placeholder = None + self.t_placeholder = None + self.sde_args_placeholder = None + self.sde_rnd_placeholder = None + self.gen_lat_placeholder = None + self.inputs_embeds_placeholder = None + self.stop_out_placeholder = None + self.graph = None + + def execute( + self, + input_tensor: torch.Tensor, + his_lat: torch.Tensor, + cfg_strength: float = 2.0, + sigma: float = 0.25, + temperature: float = 0.0, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + bat_size, his_patch_size, z_dim = his_lat.shape + randn_tensor = torch.randn( + (bat_size, self.config.patch_size, z_dim), device=input_tensor.device, dtype=input_tensor.dtype + ) + t = get_epss_timesteps(self.config.steps, device=input_tensor.device, dtype=input_tensor.dtype) + sde_rnd = torch.randn( + (self.config.steps, *randn_tensor.shape), device=input_tensor.device, dtype=input_tensor.dtype + ) + + if not self.initialized: + self._initialize_graph(input_tensor, his_lat, randn_tensor, sde_rnd) + + self.last_hidden_state_placeholder.copy_(input_tensor) + self.his_lat_placeholder.copy_(his_lat) + self.randn_like_placeholder.copy_(randn_tensor) + self.t_placeholder.copy_(t) + self.sde_args_placeholder[0] = cfg_strength + self.sde_args_placeholder[1] = sigma + self.sde_args_placeholder[2] = temperature + self.sde_rnd_placeholder.copy_(sde_rnd) + + self.graph.replay() + + gen_lat = torch.empty_like(self.gen_lat_placeholder) + gen_lat.copy_(self.gen_lat_placeholder) + + inputs_embeds = torch.empty_like(self.inputs_embeds_placeholder) + inputs_embeds.copy_(self.inputs_embeds_placeholder) + + stop_out = torch.empty_like(self.stop_out_placeholder) + stop_out.copy_(self.stop_out_placeholder) + + return gen_lat, inputs_embeds, stop_out + + def _initialize_graph( + self, + input_tensor: torch.Tensor, + his_lat: torch.Tensor, + randn_tensor: torch.Tensor, + sde_rnd: torch.Tensor, + ) -> None: + self.last_hidden_state_placeholder = torch.empty_like(input_tensor) + self.his_lat_placeholder = torch.empty_like(his_lat) + self.randn_like_placeholder = torch.empty_like(randn_tensor) + self.t_placeholder = get_epss_timesteps(self.config.steps, device=input_tensor.device, dtype=input_tensor.dtype) + self.sde_args_placeholder = torch.empty(3, device=input_tensor.device, dtype=input_tensor.dtype) + self.sde_rnd_placeholder = torch.empty_like(sde_rnd) + + self.graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(self.graph): + self.gen_lat_placeholder = self.cfm.sample( + self.last_hidden_state_placeholder, + self.his_lat_placeholder, + self.randn_like_placeholder, + self.t_placeholder, + self.sde_args_placeholder, + self.sde_rnd_placeholder, + ) + self.inputs_embeds_placeholder = self.aggregator(self.gen_lat_placeholder) + self.stop_out_placeholder = self.stop_head(self.last_hidden_state_placeholder[:, -1, :]).softmax(dim=-1) + + self.initialized = True + + +class CFMGraphExecutorPool: + """Thread-safe pool of CFMGraphExecutors for concurrent inference.""" + + def __init__(self, config, cfm, aggregator, stop_head: nn.Linear, pool_size: int = 1): + self.config = config + self.cfm = cfm + self.aggregator = aggregator + self.stop_head = stop_head + self.pool_size = pool_size + self.pool = Queue(maxsize=pool_size) + self.lock = Lock() + + for _ in range(pool_size): + executor = CFMGraphExecutor(config, cfm, aggregator, stop_head) + self.pool.put(executor) + + def acquire(self) -> CFMGraphExecutor: + return self.pool.get() + + def release(self, executor: CFMGraphExecutor) -> None: + self.pool.put(executor) + + def execute( + self, + input_tensor: torch.Tensor, + his_lat: torch.Tensor, + cfg_strength: float = 2.0, + sigma: float = 0.25, + temperature: float = 0.0, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + executor = self.acquire() + try: + return executor.execute( + input_tensor, his_lat, cfg_strength=cfg_strength, sigma=sigma, temperature=temperature + ) + finally: + self.release(executor) + + +######################################################################## +# Audio Postprocess +# Adapted from: +# https://github.com/inclusionAI/Ming/blob/e58533db227031990c5a6864dcf5f08fb53ed0d2/modeling_bailing_talker.py +######################################################################## + + +@torch.no_grad() +def resample(waveform: torch.Tensor, orig_sr: int, target_sr: int) -> torch.Tensor: + """Resample a waveform via linear interpolation (no torchaudio dep). + + Args: + waveform: Tensor shaped ``(..., num_samples)``. + orig_sr: Source sample rate (Hz); must be > 0. + target_sr: Target sample rate (Hz); must be > 0. + + Raises: + ValueError: If sample rates are non-positive, the waveform is empty, + or the resampled length would round to zero. + """ + if orig_sr <= 0: + raise ValueError(f"orig_sr must be positive, got {orig_sr}") + if target_sr <= 0: + raise ValueError(f"target_sr must be positive, got {target_sr}") + if waveform.numel() == 0 or waveform.shape[-1] == 0: + raise ValueError("waveform must contain at least one sample") + if orig_sr == target_sr: + return waveform + + ratio = target_sr / orig_sr + new_len = int(waveform.shape[-1] * ratio) + if new_len <= 0: + raise ValueError( + f"resampled waveform would be empty for input length {waveform.shape[-1]}, " + f"orig_sr={orig_sr}, target_sr={target_sr}" + ) + return torch.nn.functional.interpolate( + waveform.unsqueeze(0), + size=new_len, + mode="linear", + align_corners=False, + ).squeeze(0) + + +def trim_trailing_silence( + waveform: torch.Tensor, + sample_rate: int, + sil_th: float = 1e-3, + tail_silence_s: float = 0.3, +) -> torch.Tensor: + """Trim low-energy tail while keeping a short trailing silence. + + Works on 2-D ``(channels, samples)`` or 3-D ``(batch, channels, samples)`` + tensors. Any other shape is returned unchanged. + """ + if waveform.numel() == 0: + return waveform + + original_dim = waveform.dim() + if original_dim == 3: + speech = waveform[:, 0, :] + elif original_dim == 2: + speech = waveform + else: + return waveform + + frame_step = int(sample_rate * 0.1) + frame_size = int(sample_rate * 0.1) + if speech.shape[-1] < frame_size: + keep = min(speech.shape[-1], int(tail_silence_s * sample_rate)) + trimmed = speech[..., :keep] + else: + num_frame = (speech.shape[-1] - frame_size) // frame_step + 1 + cur_len = (num_frame - 1) * frame_step + frame_size + speech = speech[..., :cur_len] + spe_frames = speech.unfold(-1, frame_size, frame_step) + scores = spe_frames.abs().mean(dim=-1) + scores = scores.mean(dim=list(range(scores.dim() - 1))) + idx = scores.shape[0] - 1 + while idx >= 0 and scores[idx] <= sil_th: + idx -= 1 + if idx < 0: + keep = min(speech.shape[-1], int(tail_silence_s * sample_rate)) + trimmed = speech[..., :keep] + else: + non_sil_len = idx * frame_step + frame_size + int(tail_silence_s * sample_rate) + non_sil_len = min(non_sil_len, speech.shape[-1]) + trimmed = speech[..., :non_sil_len] + + if original_dim == 3: + return trimmed.unsqueeze(1) + return trimmed + + +def silence_holder( + speech: torch.Tensor, + sample_rate: int, + sil_cache: dict | None = None, + last_chunk: bool = True, + sil_th: float = 1e-3, + last_sil: float = 0.3, +) -> tuple[torch.Tensor, dict]: + """Ming-style streaming silence holder. + + Used during streaming VAE decode to defer emission of silent regions + until a non-silent frame arrives (or the stream ends). ``sil_cache`` + is carried across chunks and updated in place. + """ + if speech.numel() == 0: + return speech, sil_cache or {"holder": [], "buffer": []} + + frame_step = int(sample_rate * 0.1) + frame_size = int(sample_rate * 0.1) + if sil_cache is None: + sil_cache = {"holder": [], "buffer": []} + + if sil_cache["buffer"]: + speech = torch.cat([*sil_cache["buffer"], speech], dim=-1) + sil_cache["buffer"] = [] + + if speech.shape[-1] < frame_size: + sil_cache["buffer"].append(speech) + if last_chunk: + speech = torch.cat(sil_cache["holder"] + sil_cache["buffer"], dim=-1) + return speech[..., : int(last_sil * sample_rate)], sil_cache + return torch.zeros((*speech.shape[:-1], 0), device=speech.device, dtype=speech.dtype), sil_cache + + num_frame = (speech.shape[-1] - frame_size) // frame_step + 1 + cur_len = (num_frame - 1) * frame_step + frame_size + if speech.shape[-1] > cur_len: + sil_cache["buffer"].append(speech[..., cur_len:]) + speech = speech[..., :cur_len] + + spe_frames = speech.unfold(-1, frame_size, frame_step) + scores = spe_frames.abs().mean(dim=-1) + scores = scores.mean(dim=list(range(scores.dim() - 1))) + idx = scores.shape[0] - 1 + while idx >= 0 and scores[idx] <= sil_th: + idx -= 1 + + if idx < 0: + sil_cache["holder"].append(speech) + if last_chunk: + speech = torch.cat(sil_cache["holder"] + sil_cache["buffer"], dim=-1) + return speech[..., : int(last_sil * sample_rate)], sil_cache + return torch.zeros((*speech.shape[:-1], 0), device=speech.device, dtype=speech.dtype), sil_cache + + non_sil_len = idx * frame_step + frame_size + if last_chunk: + non_sil_len += int(last_sil * sample_rate) + non_sil_len = min(non_sil_len, speech.shape[-1]) + speech_out = torch.cat([*sil_cache["holder"], speech[..., :non_sil_len]], dim=-1) + sil_cache["holder"] = [] + if non_sil_len < speech.shape[-1]: + sil_cache["holder"].append(speech[..., non_sil_len:]) + return speech_out, sil_cache + + +######################################################################## +# Audio Postprocess +# Ported from: +# https://github.com/inclusionAI/Ming/blob/e58533db227031990c5a6864dcf5f08fb53ed0d2/talker_module/aggregator.py +######################################################################## + + +class Aggregator(nn.Module): + """Maps generated audio latent patches back to LLM embedding space.""" + + def __init__( + self, + in_channels: int = 64, + hidden_size: int = 1152, + depth: int = 28, + num_heads: int = 16, + mlp_ratio: float = 4.0, + llm_input_dim: int = 896, + **kwargs, + ): + super().__init__() + self.in_channels = in_channels + self.out_channels = in_channels + self.num_heads = num_heads + + self.word_embedder = nn.Embedding(1, hidden_size) + self.x_embedder = nn.Linear(in_channels, hidden_size) + self.hidden_size = hidden_size + + self.rotary_embed = RotaryEmbedding(hidden_size // num_heads) + + self.blocks = nn.ModuleList( + [DiTBlock(hidden_size, num_heads, mlp_ratio=mlp_ratio, **kwargs) for _ in range(depth)] + ) + self.final_layer = FinalLayer(hidden_size, llm_input_dim) + + def forward(self, x: torch.Tensor, mask: torch.Tensor | None = None) -> torch.Tensor: + x = self.x_embedder(x) + cls_embed = self.word_embedder(torch.zeros((x.shape[0], 1), dtype=torch.long, device=x.device)) + x = torch.cat([cls_embed, x], dim=1) + + rope = self.rotary_embed.forward_from_seq_len(x.shape[1]) + if mask is not None: + mask_pad = mask.clone().detach()[:, :1] + mask = torch.cat([mask_pad, mask], dim=-1) + for block in self.blocks: + x = block(x, mask, rope) + x = self.final_layer(x) + x = x[:, :1, :] + return x + + +######################################################################## +# Prompt Builder +# Adapted from: +# https://github.com/inclusionAI/Ming/blob/e58533db227031990c5a6864dcf5f08fb53ed0d2/modeling_bailing_talker.py +######################################################################## + +_MUSIC_TAGS = ("Genre: ", "Mood: ", "Instrument: ", "Theme: ", "Duration: ") + + +def _looks_like_music_prompt(text: str) -> bool: + return all(tag in text for tag in _MUSIC_TAGS) + + +def build_tts_input( + *, + tokenizer: PreTrainedTokenizerBase, + embed_tokens: torch.nn.Module, + device: torch.device, + dtype: torch.dtype, + text: str, + prompt: str, + spk_emb: list[torch.Tensor] | None = None, + instruction: str | None = None, + prompt_text: str | None = None, + prompt_wav_emb: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Build (inputs_embeds, input_ids) for one TTS segment. + + Args: + tokenizer: HF tokenizer + embed_tokens: The LLM's input-embedding module + device: Device to place the returned tensors on. + dtype: dtype for the returned `inputs_embeds`. + text: Text to synthesize. + prompt: System-level instruction prompt prepended to the user turn. + spk_emb: Optional list of speaker embeddings already projected into + LLM hidden dim; each is injected at a `<|vision_start|>` slot. + instruction: Optional free-form instruction + prompt_text: Reference text for zero-shot voice cloning. + prompt_wav_emb: Reference-wav embeddings to inject. + """ + spk_emb_prompt: list[int] = [] + if spk_emb is not None: + for i in range(len(spk_emb)): + spk_emb_prompt.extend( + tokenizer.encode(f" speaker_{i + 1}:") + + tokenizer.encode("<|vision_start|>") + + tokenizer.encode("<|vision_pad|>") + + tokenizer.encode("<|vision_end|>\n") + ) + + instruction_prompt: list[int] = [] + if instruction is not None: + instruction_prompt = tokenizer.encode(instruction) + tokenizer.encode("<|im_end|>") + + prompt_text_token: list[int] = [] + prompt_latent_token: list[int] = [] + if prompt_wav_emb is not None and prompt_text is not None: + prompt_text_token = tokenizer.encode(prompt_text) + prompt_latent_token = tokenizer.encode("") * prompt_wav_emb.size(1) + + prompt2 = [] if _looks_like_music_prompt(text) else tokenizer.encode(" Text input:\n") + + input_part = ( + tokenizer.encode("<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n") + + tokenizer.encode("<|im_start|>user\n") + + tokenizer.encode(prompt) + + spk_emb_prompt + + prompt2 + + prompt_text_token + + tokenizer.encode(text) + + tokenizer.encode("<|im_end|>\n") + + tokenizer.encode("<|im_start|>assistant\n") + + instruction_prompt + + tokenizer.encode("