diff --git a/docs/serving/speech_api.md b/docs/serving/speech_api.md index f5700a318f9..d7c738ce75f 100644 --- a/docs/serving/speech_api.md +++ b/docs/serving/speech_api.md @@ -978,7 +978,22 @@ See the [Gepard section of the online TTS hub](../user_guide/examples/online_ser | Model | Description | | ------- | ------------- | -| `k2-fsa/OmniVoice` | Pure-diffusion TTS. Supports voice cloning via `ref_audio` (with optional `ref_text`); no built-in voice presets. | +| `k2-fsa/OmniVoice` | OmniVoice generates mono speech at 24 kHz. It supports text-to-speech without a reference clip, voice cloning with `ref_audio` and optional `ref_text`, and voice design. It has no built-in voice presets. | + +For voice cloning, send `ref_audio`. If `ref_text` is missing or blank, the default deploy config leaves Whisper unloaded until the first such request, then loads `openai/whisper-large-v3-turbo` on the worker device and transcribes the clip. Set `additional_config.omnivoice_asr.load_asr_on_startup` to `true` to load it during worker startup. Set `additional_config.omnivoice_asr.asr_device` to `cuda:0` or `cpu` to override the worker device. Sending `ref_text` avoids ASR and uses the supplied transcript. + +OmniVoice prepares the reference clip before it builds the voice cloning input. It performs these steps: + +- It converts the clip to mono, resamples it to 24 kHz, and measures its original root mean square (RMS) level. +- When the original RMS is greater than zero and less than `0.1`, it raises the waveform level to `0.1` and keeps the original RMS for output volume control. +- When `ref_text` is missing, it trims a clip longer than 20 seconds at a detected silence. Supplied text skips this special long-audio trim. +- It removes middle silence and trims silence at both edges whether the transcript is automatic or supplied by the caller. +- It aligns the sample count down to the audio tokenizer hop size. +- When `ref_text` is missing, it sends the prepared waveform to Whisper. +- It adds a final period for English text or a Chinese full stop for Chinese text when the reference text has no ending punctuation. +- It sends the same prepared waveform to the audio tokenizer. + +OmniVoice processes voice-cloning audio after decoding. It removes middle gaps of at least 500 milliseconds. It trims edge silence while keeping 100 milliseconds at each edge. It scales cloned output using the original reference RMS when that value is below `0.1`, and leaves the level unchanged when the value is at least `0.1`. It fades the first and last 100 milliseconds, then adds 100 milliseconds of silence at each edge. Text-only output returns the decoder tensor without this voice-cloning postprocessing. OmniVoice uses packed variable-length attention for batched generator execution. The attention operator accepts FP16 and BF16 inputs, so its query, key, and value tensors are evaluated in BF16 even when the stage is configured with `dtype: float32`; the attention output is converted back to the model's diff --git a/docs/user_guide/examples/offline_inference/text_to_speech.md b/docs/user_guide/examples/offline_inference/text_to_speech.md index cee7c4a321e..940a6b05ca4 100644 --- a/docs/user_guide/examples/offline_inference/text_to_speech.md +++ b/docs/user_guide/examples/offline_inference/text_to_speech.md @@ -36,7 +36,7 @@ python examples/offline_inference/text_to_speech//end2end.py \ --ref-text "Transcript of the reference audio." ``` -`--ref-audio` and `--ref-text` are optional (text-only synthesis works without them) and must be provided together for voice cloning. The exotic scripts — Qwen3-TTS, Voxtral TTS, CosyVoice3 — accept additional model-specific flags documented in their per-model section below. +For most models, `--ref-audio` and `--ref-text` are optional for text-only synthesis but must be provided together for voice cloning. OmniVoice also accepts `--ref-audio` without `--ref-text` and transcribes the reference lazily with Whisper. Qwen3-TTS, Voxtral TTS, and CosyVoice3 accept additional model-specific flags documented in their sections below. --- @@ -252,7 +252,7 @@ Text → [Stage 0: AR] → Speech Tokens → [Stage 1: DiT + HiFT] → Audio (24 ## OmniVoice -Zero-shot multilingual TTS supporting 600+ languages, with three modes (auto / clone / design). +OmniVoice creates speech in more than 600 languages. It supports text-to-speech without a reference clip, voice cloning, and voice design. ### Prerequisites @@ -260,7 +260,7 @@ Zero-shot multilingual TTS supporting 600+ languages, with three modes (auto / c huggingface-cli download k2-fsa/OmniVoice ``` -Voice cloning requires `transformers>=5.3.0`. Auto and design modes work with `transformers>=4.57.0`. +Voice cloning requires `transformers>=5.3.0`. Automatic voice and voice design use `transformers>=4.57.0`. ### Quick start (auto voice) @@ -272,6 +272,8 @@ python examples/offline_inference/text_to_speech/omnivoice/end2end.py \ ### Voice cloning +With an explicit transcript, voice cloning does not load Whisper: + ```bash python examples/offline_inference/text_to_speech/omnivoice/end2end.py \ --model k2-fsa/OmniVoice \ @@ -280,6 +282,30 @@ python examples/offline_inference/text_to_speech/omnivoice/end2end.py \ --ref-text "This is the reference transcription." ``` +If `--ref-text` is omitted, OmniVoice loads `openai/whisper-large-v3-turbo` on the worker device and transcribes the reference audio after it prepares the waveform. The first request downloads and loads Whisper, so it takes longer and uses more device memory. Set `additional_config.omnivoice_asr.load_asr_on_startup` to `true` to load Whisper during worker startup, or set `additional_config.omnivoice_asr.asr_device` to use another GPU or the CPU. + +```bash +python examples/offline_inference/text_to_speech/omnivoice/end2end.py \ + --model k2-fsa/OmniVoice \ + --text "hello" \ + --ref-audio trump_ref.wav +``` + +### Reference audio processing + +OmniVoice prepares the reference clip before it transcribes the clip or builds the voice cloning input. It performs these steps: + +- It converts the clip to mono and resamples it to 24 kHz. +- It measures the original root mean square (RMS) level. When the RMS is greater than zero and below `0.1`, it raises the waveform level to `0.1` and keeps the original RMS for output volume control. +- When `--ref-text` is omitted, it trims a clip longer than 20 seconds at a detected silence. +- It removes middle silence and trims silence at both edges whether the transcript is automatic or supplied by the caller. +- It aligns the sample count down to the audio tokenizer hop size. +- When `--ref-text` is omitted, it sends the prepared waveform to Whisper. +- It adds a final period for English text or a Chinese full stop for Chinese text when the reference text has no ending punctuation. +- It sends the same prepared waveform to the audio tokenizer and uses the resulting reference tokens to estimate the target duration. + +OmniVoice processes voice cloning audio after decoding. It removes middle gaps of at least 500 milliseconds and trims edge silence while keeping 100 milliseconds at each edge. It uses the original reference RMS when the RMS is below `0.1`, and it leaves the level unchanged when the RMS is at least `0.1`. It fades the first and last 100 milliseconds, then adds 100 milliseconds of silence at both edges. Text only output returns the decoder tensor without this processing. + ### Voice design ```bash @@ -302,6 +328,11 @@ python examples/offline_inference/text_to_speech/omnivoice/end2end.py \ - Stage 0 (Generator): Qwen3-0.6B with 32-step iterative unmasking. - Stage 1 (Decoder): HiggsAudioV2 RVQ + DAC at 24 kHz. +- By default, the first request with reference audio and no `--ref-text` downloads + and loads `openai/whisper-large-v3-turbo`, so it has extra latency and GPU + memory use. Set `additional_config.omnivoice_asr.load_asr_on_startup` to `true` + to move that load to worker startup. +- Generated audio remains mono at 24 kHz. --- diff --git a/docs/user_guide/examples/online_serving/text_to_speech.md b/docs/user_guide/examples/online_serving/text_to_speech.md index 65268b9512a..c42dafbea2d 100644 --- a/docs/user_guide/examples/online_serving/text_to_speech.md +++ b/docs/user_guide/examples/online_serving/text_to_speech.md @@ -449,7 +449,7 @@ bash examples/online_serving/text_to_speech/glm_tts/run_gradio_demo.sh ## OmniVoice -Zero-shot multilingual TTS (600+ languages). Online serving currently exposes **auto voice** only; voice cloning and voice design are available offline. +OmniVoice creates speech in more than 600 languages. It supports text-to-speech without a reference clip, voice cloning, and voice design. ### Prerequisites @@ -457,7 +457,7 @@ Zero-shot multilingual TTS (600+ languages). Online serving currently exposes ** huggingface-cli download k2-fsa/OmniVoice ``` -Voice cloning (offline) needs `transformers>=5.3.0`; auto voice works with `transformers>=4.57.0`. +Voice cloning needs `transformers>=5.3.0`. When `ref_text` is omitted, OmniVoice loads `openai/whisper-large-v3-turbo` on the worker device to transcribe the reference audio. Set `additional_config.omnivoice_asr.load_asr_on_startup` to `true` to load Whisper during worker startup, or set `additional_config.omnivoice_asr.asr_device` to use another GPU or the CPU. ### Launch @@ -473,13 +473,35 @@ vllm serve k2-fsa/OmniVoice --omni --port 8091 --trust-remote-code cd examples/online_serving/text_to_speech/omnivoice python speech_client.py --text "Hello, how are you?" python speech_client.py --text "Bonjour, comment allez-vous?" --language French +python speech_client.py --text "hello" --ref-audio /path/to/ref_audio.wav ``` -The client supports `--api-base`, `--model`, `--text`, `--response-format`, `--language`, `--output`. +The client supports `--api-base`, `--model`, `--text`, `--response-format`, +`--language`, `--voice`, `--ref-audio`, `--ref-text`, `--instructions`, +`--seed`, and `--output`. ### Notes -- Voice cloning and voice design require offline inference; see the [offline OmniVoice section](https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/text_to_speech/README.md#omnivoice). +- The explicit `--ref-text` path is faster because it does not load Whisper. +- By default, the first request with reference audio and no reference text loads + `openai/whisper-large-v3-turbo` on the worker device, which adds latency and + device memory use. Set `additional_config.omnivoice_asr.load_asr_on_startup` to + `true` to move that load to worker startup. Generated audio is mono at 24 kHz. + +### Reference audio and output processing + +When a request includes `ref_audio`, OmniVoice prepares the reference clip before it transcribes the clip or builds the voice cloning input. It performs these steps: + +- It converts the clip to mono and resamples it to 24 kHz. +- It measures the original root mean square (RMS) level. When the RMS is greater than zero and below `0.1`, it raises the waveform level to `0.1` and keeps the original RMS for output volume control. +- When `ref_text` is omitted, it trims a clip longer than 20 seconds at a detected silence. +- It removes middle silence and trims silence at both edges whether the transcript is automatic or supplied by the caller. +- It aligns the sample count down to the audio tokenizer hop size. +- When `ref_text` is omitted, it sends the prepared waveform to Whisper. +- It adds a final period for English text or a Chinese full stop for Chinese text when the reference text has no ending punctuation. +- It sends the same prepared waveform to the audio tokenizer. + +After decoding, OmniVoice processes voice cloning output by removing middle gaps of at least 500 milliseconds and trimming edge silence while keeping 100 milliseconds at each edge. It uses the original reference RMS when the RMS is below `0.1`, and it leaves the level unchanged when the RMS is at least `0.1`. It fades the first and last 100 milliseconds, then adds 100 milliseconds of silence at both edges. Text only output returns the decoder tensor without this processing. --- diff --git a/examples/offline_inference/text_to_speech/README.md b/examples/offline_inference/text_to_speech/README.md index 1495a40068e..235f23c7f4c 100644 --- a/examples/offline_inference/text_to_speech/README.md +++ b/examples/offline_inference/text_to_speech/README.md @@ -44,7 +44,7 @@ python examples/offline_inference/text_to_speech//end2end.py \ --ref-text "Transcript of the reference audio." ``` -`--ref-audio` and `--ref-text` are optional (text-only synthesis works without them) and must be provided together for voice cloning. The exotic scripts — Qwen3-TTS, Voxtral TTS, CosyVoice3 — accept additional model-specific flags documented in their per-model section below. Qwen3-TTS in particular uses its own argparse surface (`--query-type`, `--audio-path`, etc.) and does not follow the common shape; see its section. +For most models, `--ref-audio` and `--ref-text` are optional for text-only synthesis but must be provided together for voice cloning. OmniVoice also accepts `--ref-audio` without `--ref-text` and transcribes the reference lazily with Whisper. Qwen3-TTS, Voxtral TTS, and CosyVoice3 accept additional model-specific flags documented in their sections below. Qwen3-TTS uses its own argparse surface with options such as `--query-type` and `--audio-path`, so it does not follow the common command shape. --- @@ -407,7 +407,7 @@ python examples/offline_inference/text_to_speech/moss_tts_nano/end2end.py \ ## OmniVoice -Zero-shot multilingual TTS supporting 600+ languages, with three modes (auto / clone / design). +OmniVoice creates speech in more than 600 languages. It supports text-to-speech without a reference clip, voice cloning, and voice design. ### Prerequisites @@ -415,7 +415,7 @@ Zero-shot multilingual TTS supporting 600+ languages, with three modes (auto / c huggingface-cli download k2-fsa/OmniVoice ``` -Voice cloning requires `transformers>=5.3.0`. Auto and design modes work with `transformers>=4.57.0`. +Voice cloning requires `transformers>=5.3.0`. Automatic voice and voice design use `transformers>=4.57.0`. ### Quick start (auto voice) @@ -427,6 +427,8 @@ python examples/offline_inference/text_to_speech/omnivoice/end2end.py \ ### Voice cloning +With an explicit transcript, voice cloning does not load Whisper: + ```bash python examples/offline_inference/text_to_speech/omnivoice/end2end.py \ --model k2-fsa/OmniVoice \ @@ -435,6 +437,30 @@ python examples/offline_inference/text_to_speech/omnivoice/end2end.py \ --ref-text "This is the reference transcription." ``` +Omit `--ref-text` to use automatic transcription. The default OmniVoice deploy config loads Whisper on the worker device when the first request needs it. Set `load_asr_on_startup: true` in `additional_config.omnivoice_asr` to load it during worker startup. Set `asr_device: "cuda:0"` or `asr_device: "cpu"` in the same block to override the worker device: + +```bash +python examples/offline_inference/text_to_speech/omnivoice/end2end.py \ + --model k2-fsa/OmniVoice \ + --text "hello" \ + --ref-audio trump_ref.wav +``` + +### Reference audio processing + +OmniVoice prepares the reference clip before it transcribes the clip or builds the voice cloning input. It performs these steps: + +- It converts the clip to mono and resamples it to 24 kHz. +- It measures the original root mean square (RMS) level. When the RMS is greater than zero and below `0.1`, it raises the waveform level to `0.1` and keeps the original RMS for output volume control. +- When `--ref-text` is omitted, it trims a clip longer than 20 seconds at a detected silence. +- It removes middle silence and trims silence at both edges whether the transcript is automatic or supplied by the caller. +- It aligns the sample count down to the audio tokenizer hop size. +- When `--ref-text` is omitted, it sends the prepared waveform to Whisper. +- It adds a final period for English text or a Chinese full stop for Chinese text when the reference text has no ending punctuation. +- It sends the same prepared waveform to the audio tokenizer. + +After decoding, OmniVoice applies silence removal, RMS scaling, fades, and padding only to voice-cloning output. Text-only output returns the decoder tensor unchanged. Output audio is mono at 24 kHz. + ### Voice design ```bash @@ -466,6 +492,11 @@ python examples/offline_inference/text_to_speech/omnivoice/end2end.py \ - Stage 0 (Generator): Qwen3-0.6B with 32-step iterative unmasking. - Stage 1 (Decoder): HiggsAudioV2 RVQ + DAC at 24 kHz. +- The default deploy config uses `openai/whisper-large-v3-turbo` for lazy ASR on + the worker device. Set `load_asr_on_startup: true` in `additional_config.omnivoice_asr` + to load it during worker startup. Set `asr_device` in the same block to use a + different GPU or the CPU. Supplying `--ref-text` avoids ASR. +- Generated audio remains mono at 24 kHz. --- diff --git a/examples/offline_inference/text_to_speech/omnivoice/end2end.py b/examples/offline_inference/text_to_speech/omnivoice/end2end.py index a54ae863d4f..bf040d14c96 100644 --- a/examples/offline_inference/text_to_speech/omnivoice/end2end.py +++ b/examples/offline_inference/text_to_speech/omnivoice/end2end.py @@ -1,10 +1,11 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """End-to-end OmniVoice TTS inference via vLLM-Omni. Supports: - Auto voice mode: text only → generated speech - Voice cloning mode: text + reference audio → cloned voice speech +- Automatic reference transcription: reference audio without reference text Usage: # Auto voice @@ -13,6 +14,11 @@ # Voice cloning python end2end.py --model k2-fsa/OmniVoice --text "Hello" \ --ref-audio ref.wav --ref-text "reference transcription" + + # Automatic reference transcription (lazy by default; set + # load_asr_on_startup: true to preload Whisper at worker startup) + python end2end.py --model k2-fsa/OmniVoice --text "hello" \ + --ref-audio trump_ref.wav """ import argparse @@ -113,7 +119,8 @@ def run_e2e(): audio_signal, sr = load_audio(args.ref_audio, sr=None) multi_modal_data["audio"] = (audio_signal.astype(np.float32), sr) - mm_processor_kwargs["ref_text"] = args.ref_text or "" + if args.ref_text is not None: + mm_processor_kwargs["ref_text"] = args.ref_text mm_processor_kwargs["sample_rate"] = sr if args.lang: diff --git a/examples/online_serving/text_to_speech/README.md b/examples/online_serving/text_to_speech/README.md index 9b07dd4d88a..c1f754c6b80 100644 --- a/examples/online_serving/text_to_speech/README.md +++ b/examples/online_serving/text_to_speech/README.md @@ -727,7 +727,7 @@ Then open in your browser. ## OmniVoice -Zero-shot multilingual TTS (600+ languages). Online serving currently exposes **auto voice** only; voice cloning and voice design are available offline. +OmniVoice creates speech in more than 600 languages. It supports text-to-speech without a reference clip, voice cloning, and voice design. ### Prerequisites @@ -735,7 +735,7 @@ Zero-shot multilingual TTS (600+ languages). Online serving currently exposes ** huggingface-cli download k2-fsa/OmniVoice ``` -Voice cloning (offline) needs `transformers>=5.3.0`; auto voice works with `transformers>=4.57.0`. +Voice cloning needs `transformers>=5.3.0`. The default OmniVoice deploy config loads `openai/whisper-large-v3-turbo` lazily on the worker device when `ref_text` is missing. Set `load_asr_on_startup: true` in `additional_config.omnivoice_asr` to load it during worker startup. Set `asr_device: "cuda:0"` or `asr_device: "cpu"` in the same block to override the worker device. ### Launch @@ -760,6 +760,11 @@ python speech_client.py \ --ref-audio /path/to/ref_audio.wav \ --ref-text "Bonjour, comment allez-vous?" +# Voice cloning with automatic reference transcription +python speech_client.py \ +--text "hello" \ +--ref-audio /path/to/ref_audio.wav + # Style instruction (voice design-style control) python speech_client.py \ --text "Bonjour, comment allez-vous?" \ @@ -772,6 +777,23 @@ python speech_client.py --text "Hello, how are you?" --seed 42 The client supports `--api-base`, `--model`, `--text`, `--response-format`, `--language`, `--voice`, `--ref-audio`, `--ref-text`, `--instructions`, `--seed`, and `--output`. +### Reference audio and output processing + +The explicit `--ref-text` path avoids Whisper. With the default config, missing `ref_text` loads `openai/whisper-large-v3-turbo` on the worker device during the first request that needs transcription. Set `load_asr_on_startup: true` in `additional_config.omnivoice_asr` to move that load to worker startup. Set `asr_device` in the same block to choose another GPU or the CPU. + +OmniVoice prepares the reference clip before it transcribes the clip or builds the voice cloning input. It performs these steps: + +- It converts the clip to mono and resamples it to 24 kHz. +- It measures the original root mean square (RMS) level. When the RMS is greater than zero and below `0.1`, it raises the waveform level to `0.1` and keeps the original RMS for output volume control. +- When `--ref-text` is omitted, it trims a clip longer than 20 seconds at a detected silence. Supplied text skips this special long-audio trim. +- It removes middle silence and trims silence at both edges whether the transcript is automatic or supplied by the caller. +- It aligns the sample count down to the audio tokenizer hop size. +- When `--ref-text` is omitted, it sends the prepared waveform to Whisper. +- It adds a final period for English text or a Chinese full stop for Chinese text when the reference text has no ending punctuation. +- It sends the same prepared waveform to the audio tokenizer. + +After decoding, OmniVoice applies silence removal, RMS scaling, fades, and padding only to voice-cloning output. Text-only output returns the decoder tensor unchanged. Output audio is mono at 24 kHz. + ## Qwen3-TTS Three model variants exposed via separate checkpoints: diff --git a/examples/online_serving/text_to_speech/omnivoice/speech_client.py b/examples/online_serving/text_to_speech/omnivoice/speech_client.py index 9c9b8128ea6..ee311762691 100644 --- a/examples/online_serving/text_to_speech/omnivoice/speech_client.py +++ b/examples/online_serving/text_to_speech/omnivoice/speech_client.py @@ -9,6 +9,9 @@ # Use a specific uploaded/supported voice python speech_client.py --text "Hello" --voice my_uploaded_voice + + # Clone a reference voice and let the server transcribe the reference audio + python speech_client.py --text "hello" --ref-audio /path/to/ref_audio.wav """ import argparse diff --git a/tests/diffusion/models/omnivoice/test_prompt_tokenization.py b/tests/diffusion/models/omnivoice/test_prompt_tokenization.py index b8fae42414a..46d0d08bef8 100644 --- a/tests/diffusion/models/omnivoice/test_prompt_tokenization.py +++ b/tests/diffusion/models/omnivoice/test_prompt_tokenization.py @@ -28,23 +28,26 @@ def test_prepare_request_tokenizes_text_once_with_reference_audio(audio_mode): pipeline.tokenizer = SimpleNamespace( encode=Mock(side_effect=[SimpleNamespace(ids=[1, 2]), SimpleNamespace(ids=[5, 6, 7])]) ) - pipeline.audio_tokenizer = Mock() + pipeline.audio_tokenizer = SimpleNamespace(config=SimpleNamespace(sample_rate=24000, hop_length=240)) ref_tokens = torch.full((8, 4), 8, dtype=torch.long) pipeline._encode_ref_audio = Mock(return_value=ref_tokens) pipeline._speaker_cache = Mock() - pipeline._speaker_cache.get.return_value = {"ref_audio_tokens": ref_tokens} if audio_mode == "cached" else None + pipeline._speaker_cache.get.return_value = ( + {"ref_audio_tokens": ref_tokens, "reference_rms": 0.2} if audio_mode == "cached" else None + ) prompt = { "text": "hello world", "mm_processor_kwargs": {"ref_text": "reference"}, } if audio_mode != "none": - prompt["multi_modal_data"] = {"audio": [(np.zeros(2400, dtype=np.float32), 24000)]} + prompt["multi_modal_data"] = {"audio": [(np.full(2400, 0.2, dtype=np.float32), 24000)]} prompt["voice_name"] = "speaker" prepared = pipeline._prepare_request_input(prompt, {"seed": 42}) assert pipeline.tokenizer.encode.call_count == 2 # One style and one text encoding. - assert pipeline.tokenizer.encode.call_args.args == ("<|text_start|>reference hello world<|text_end|>",) + expected_ref_text = "reference" if audio_mode == "none" else "reference." + assert pipeline.tokenizer.encode.call_args.args == (f"<|text_start|>{expected_ref_text} hello world<|text_end|>",) torch.testing.assert_close(prepared.input_ids[:5], torch.tensor([1, 2, 5, 6, 7]).unsqueeze(1).repeat(1, 8)) assert prepared.target_len == 2 assert prepared.seed == 42 @@ -54,7 +57,7 @@ def test_prepare_request_tokenizes_text_once_with_reference_audio(audio_mode): if audio_mode == "uncached": pipeline._encode_ref_audio.assert_called_once() audio, sr = pipeline._encode_ref_audio.call_args.args - assert tuple(audio.shape) == (2400,) + assert tuple(audio.shape) == (1, 2400) assert sr == 24000 pipeline._speaker_cache.put.assert_called_once() else: diff --git a/tests/e2e/offline_inference/test_omnivoice_expansion.py b/tests/e2e/offline_inference/test_omnivoice_expansion.py index dd65cbdccc6..d507d354d76 100644 --- a/tests/e2e/offline_inference/test_omnivoice_expansion.py +++ b/tests/e2e/offline_inference/test_omnivoice_expansion.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ E2E offline tests for OmniVoice TTS model with text input and audio output. @@ -12,13 +12,18 @@ import numpy as np import pytest +import soundfile as sf +from vllm.multimodal.media.audio import load_audio from tests.helpers.mark import hardware_test +from tests.helpers.media import convert_audio_file_to_text, get_asset_path from tests.helpers.runtime import OmniRunner from tests.helpers.stage_config import get_deploy_config_path +from vllm_omni.inputs.data import OmniDiffusionSamplingParams MODEL = "k2-fsa/OmniVoice" DEPLOY_CONFIG = get_deploy_config_path("omnivoice.yaml") +_OMNIVOICE_REF_AUDIO_SEED = 102 # OmniRunner tuple: model, deploy config path, extra Omni kwargs. _OMNI_RUNNER_PARAM = ( @@ -45,8 +50,6 @@ def test_omnivoice_text_to_audio(omni_runner: OmniRunner) -> None: Input Modal: text Output Modal: audio """ - from vllm_omni.inputs.data import OmniDiffusionSamplingParams - prompts = {"prompt": "Hello, this is a test for text to audio."} sampling_params_list = [OmniDiffusionSamplingParams()] @@ -56,12 +59,7 @@ def test_omnivoice_text_to_audio(omni_runner: OmniRunner) -> None: # Check final output has audio final_output = outputs[-1] - ro = final_output - assert ro is not None, "No request_output" - - mm = getattr(ro, "multimodal_output", None) - if not mm and ro.outputs: - mm = getattr(ro.outputs[0], "multimodal_output", None) + mm = final_output.multimodal_output assert mm is not None, "No multimodal_output" assert "audio" in mm, f"No 'audio' key in multimodal_output: {mm.keys()}" @@ -77,3 +75,48 @@ def test_omnivoice_text_to_audio(omni_runner: OmniRunner) -> None: assert rms > 0.01, f"Audio RMS too low ({rms:.4f}), likely silence" print(f"Generated audio: {len(audio_np) / 24000:.2f}s, rms={rms:.4f}") + + +@hardware_test(res={"cuda": "L4"}, num_cards=1) +def test_omnivoice_ref_audio_without_ref_text(omni_runner: OmniRunner, tmp_path) -> None: + """Reference audio without a transcript is transcribed before cloning.""" + ref_audio_path = get_asset_path("qwen3_tts/clone_2.wav") + audio_signal, sample_rate = load_audio(str(ref_audio_path), sr=None) + prompts = { + "prompt": "hello", + "multi_modal_data": { + "audio": (audio_signal.astype(np.float32), sample_rate), + }, + "mm_processor_kwargs": {"sample_rate": sample_rate}, + } + + def generate_audio(): + sampling_params_list = [OmniDiffusionSamplingParams(extra_args={"seed": _OMNIVOICE_REF_AUDIO_SEED})] + outputs = list(omni_runner.omni.generate(prompts, sampling_params_list=sampling_params_list)) + assert outputs, "No outputs generated" + final_output = outputs[-1] + mm = final_output.multimodal_output + assert mm is not None, "No multimodal_output" + + audio = mm.get("audio") + assert audio is not None, "No audio output" + audio_np = audio if isinstance(audio, np.ndarray) else audio.cpu().numpy().squeeze() + return np.asarray(audio_np).squeeze(), int(mm.get("sr", 24000)) + + audio_np, output_sample_rate = generate_audio() + assert audio_np.size > 1, "Audio output is empty" + assert np.isfinite(audio_np).all(), "Audio contains non-finite samples" + assert np.unique(audio_np).size > 1, "Audio has no sample variation" + rms = np.sqrt(np.mean(audio_np**2)) + assert rms > 0.01, f"Audio RMS too low ({rms:.4f}), likely white noise or silence" + assert output_sample_rate == 24000 + + repeated_audio, repeated_sample_rate = generate_audio() + assert repeated_sample_rate == output_sample_rate + assert repeated_audio.shape == audio_np.shape + np.testing.assert_allclose(repeated_audio, audio_np, rtol=1e-5, atol=1e-4) + + output_path = tmp_path / "omnivoice_ref_audio_only.wav" + sf.write(str(output_path), audio_np, output_sample_rate) + transcript = convert_audio_file_to_text(str(output_path), language="en").lower() + assert "hello" in transcript, f"Generated audio transcript does not contain 'hello': {transcript!r}" diff --git a/tests/e2e/online_serving/test_omnivoice_expansion.py b/tests/e2e/online_serving/test_omnivoice_expansion.py index c66750b7dce..b4b2c0eb899 100644 --- a/tests/e2e/online_serving/test_omnivoice_expansion.py +++ b/tests/e2e/online_serving/test_omnivoice_expansion.py @@ -14,7 +14,9 @@ os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" +import numpy as np import pytest +import soundfile as sf from openai import BadRequestError as OpenAIBadRequestError from tests.helpers.mark import hardware_test @@ -60,6 +62,7 @@ ] # Lower this in ``request_config`` via ``min_audio_bytes`` if a run produces legitimately short WAVs. _DEFAULT_MIN_AUDIO_BYTES = 5000 +_OMNIVOICE_REF_AUDIO_SEED = 102 REF_AUDIO_URL = get_asset_path("qwen3_tts/clone_2.wav", as_data_url=True) @@ -184,6 +187,8 @@ def test_speech_auto_voice_seed_deterministic(self, omni_server, online_client) r1 = online_client.send_audio_speech_request(cfg)[0] r2 = online_client.send_audio_speech_request(cfg)[0] + assert r1.audio_bytes is not None + assert r2.audio_bytes is not None assert r1.audio_bytes == r2.audio_bytes @hardware_test(res={"cuda": "L4"}, num_cards=1) @@ -210,16 +215,31 @@ class TestOmniVoiceVoiceCloning: @hardware_test(res={"cuda": "L4"}, num_cards=1) def test_voice_clone_ref_audio_only(self, omni_server, online_client) -> None: - """Test voice cloning with ref_audio only (x_vector mode).""" + """Test automatic reference transcription with ref_audio only.""" request_config = { "model": omni_server.model, - "input": get_prompt("text"), + "input": "hello", "ref_audio": REF_AUDIO_URL, "response_format": "wav", + "seed": _OMNIVOICE_REF_AUDIO_SEED, "timeout": 180.0, "min_audio_bytes": _DEFAULT_MIN_AUDIO_BYTES, + "transcript_language": "en", } - online_client.send_audio_speech_request(request_config) + response = online_client.send_audio_speech_request(request_config)[0] + assert response.audio_bytes is not None + audio, sample_rate = sf.read(io.BytesIO(response.audio_bytes), dtype="float32") + assert sample_rate == 24000 + assert np.isfinite(audio).all() + assert np.unique(audio).size > 1 + assert np.sqrt(np.mean(audio**2)) > 0.01 + + repeated_response = online_client.send_audio_speech_request(request_config)[0] + assert repeated_response.audio_bytes is not None + repeated_audio, repeated_sample_rate = sf.read(io.BytesIO(repeated_response.audio_bytes), dtype="float32") + assert repeated_sample_rate == sample_rate + assert repeated_audio.shape == audio.shape + np.testing.assert_allclose(repeated_audio, audio, rtol=1e-5, atol=1e-4) @hardware_test(res={"cuda": "L4"}, num_cards=1) def test_voice_clone_ref_audio_and_text(self, omni_server, online_client) -> None: diff --git a/tests/engine/test_omni_request_output.py b/tests/engine/test_omni_request_output.py index 0179a221bdd..e3c3c696d1f 100644 --- a/tests/engine/test_omni_request_output.py +++ b/tests/engine/test_omni_request_output.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Tests for OmniRequestOutput class.""" import numpy as np @@ -8,10 +8,32 @@ from vllm.outputs import CompletionOutput, RequestOutput from vllm_omni.outputs import OmniRequestOutput +from vllm_omni.outputs.mm_outputs import MultimodalCompletionOutput pytestmark = [pytest.mark.core_model, pytest.mark.cpu] +@pytest.mark.parametrize("payload_location", ["request", "completion", "missing"]) +def test_multimodal_output_audio_accessor(payload_location): + audio = {"audio": [0.1, -0.1], "sr": 24000} + output = OmniRequestOutput(request_id="audio-accessor") + if payload_location == "request": + output = OmniRequestOutput.from_diffusion(request_id="audio-accessor", images=[], multimodal_output=audio) + elif payload_location == "completion": + output.outputs = [ + MultimodalCompletionOutput( + index=0, + text="", + token_ids=[], + cumulative_logprob=None, + logprobs=None, + multimodal_output=audio, + ) + ] + + assert output.multimodal_output == ({} if payload_location == "missing" else audio) + + # ── helpers ──────────────────────────────────────────────────────────────── diff --git a/tests/model_executor/models/omnivoice/test_asr.py b/tests/model_executor/models/omnivoice/test_asr.py new file mode 100644 index 00000000000..c61e218730e --- /dev/null +++ b/tests/model_executor/models/omnivoice/test_asr.py @@ -0,0 +1,339 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from __future__ import annotations + +import numpy as np +import pytest +import torch + +from vllm_omni.diffusion.models.omnivoice import pipeline_omnivoice + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class _FakeASRModel: + def __init__(self): + self.to_calls = [] + + def to(self, device): + self.to_calls.append(device) + return self + + +class _FakeASRPipeline: + def __init__(self, result): + self.result = result + self.inputs = [] + self.model = _FakeASRModel() + self.device = torch.device("cpu") + + def __call__(self, audio_input): + self.inputs.append(audio_input) + return self.result + + +@pytest.fixture +def pipeline() -> pipeline_omnivoice.OmniVoicePipeline: + model = pipeline_omnivoice.OmniVoicePipeline.__new__(pipeline_omnivoice.OmniVoicePipeline) + torch.nn.Module.__init__(model) + model.device = torch.device("cpu") + model._asr_model_name = pipeline_omnivoice._ASR_MODEL_NAME + model._asr_device = "cpu" + model._asr_pipeline = None + return model + + +def _install_fake_asr(monkeypatch, result=None): + asr = _FakeASRPipeline(result={"text": " hello there "} if result is None else result) + load_calls = [] + + def load_pipeline(*args, **kwargs): + load_calls.append((args, kwargs)) + return asr + + monkeypatch.setattr(pipeline_omnivoice, "hf_pipeline", load_pipeline) + return asr, load_calls + + +def test_missing_reference_text_loads_asr_and_strips_transcript(pipeline, monkeypatch): + asr, load_calls = _install_fake_asr(monkeypatch) + + transcript = pipeline._resolve_ref_text((np.zeros((1, 4), dtype=np.float32), 16000), None) + + assert transcript == "hello there" + assert len(load_calls) == 1 + assert asr.inputs[0]["sampling_rate"] == 16000 + assert asr.inputs[0]["array"].shape == (4,) + + +def test_second_missing_reference_text_request_reuses_asr(pipeline, monkeypatch): + asr, load_calls = _install_fake_asr(monkeypatch) + ref_audio = (np.zeros(4, dtype=np.float32), 16000) + + assert pipeline._resolve_ref_text(ref_audio, None) == "hello there" + assert pipeline._resolve_ref_text(ref_audio, " ") == "hello there" + + assert len(load_calls) == 1 + assert len(asr.inputs) == 2 + + +def test_explicit_reference_text_bypasses_asr(pipeline, monkeypatch): + _install_fake_asr(monkeypatch) + ref_text = " caller supplied transcript " + + assert pipeline._resolve_ref_text((np.zeros(4), 16000), ref_text) == ref_text + assert pipeline._asr_pipeline is None + + +def test_text_only_input_bypasses_asr(pipeline, monkeypatch): + _install_fake_asr(monkeypatch) + + assert pipeline._resolve_ref_text(None, None) is None + assert pipeline._asr_pipeline is None + + +@pytest.mark.parametrize("ref_text", [None, "", " ", "\t\n"]) +def test_empty_reference_text_uses_asr(pipeline, monkeypatch, ref_text): + _install_fake_asr(monkeypatch) + + assert pipeline._resolve_ref_text((np.zeros(4), 16000), ref_text) == "hello there" + + +def test_empty_asr_result_is_rejected(pipeline, monkeypatch): + _install_fake_asr(monkeypatch, result={"text": " \t"}) + + with pytest.raises(ValueError, match="empty reference transcription"): + pipeline._resolve_ref_text((np.zeros(4), 16000), None) + + +def test_malformed_asr_result_is_rejected(pipeline, monkeypatch): + _install_fake_asr(monkeypatch, result={"chunks": []}) + + with pytest.raises(RuntimeError, match="without a 'text' field"): + pipeline._resolve_ref_text((np.zeros(4), 16000), None) + + +def test_asr_load_failure_includes_checkpoint_and_device(pipeline, monkeypatch): + def load_pipeline(*args, **kwargs): + raise OSError("download failed") + + monkeypatch.setattr(pipeline_omnivoice, "hf_pipeline", load_pipeline) + + with pytest.raises(RuntimeError) as exc_info: + pipeline._load_asr_pipeline() + + message = str(exc_info.value) + assert pipeline_omnivoice._ASR_MODEL_NAME in message + assert "cpu" in message + assert "download failed" in message + + +def test_asr_device_placement_failure_is_reported_and_can_be_retried(pipeline, monkeypatch): + asr, load_calls = _install_fake_asr(monkeypatch) + pipeline._asr_device = "cuda:1" + + def move_model(device): + raise RuntimeError(f"cannot move to {device}") + + monkeypatch.setattr(asr.model, "to", move_model) + + with pytest.raises(RuntimeError) as exc_info: + pipeline._load_asr_pipeline() + + message = str(exc_info.value) + assert pipeline_omnivoice._ASR_MODEL_NAME in message + assert "cuda:1" in message + assert "cannot move" in message + assert pipeline._asr_pipeline is None + + monkeypatch.setattr(asr.model, "to", lambda device: asr.model) + assert pipeline._load_asr_pipeline() is asr + assert len(load_calls) == 2 + assert asr.device == torch.device("cuda:1") + + +@pytest.mark.parametrize( + ("additional_config", "expected"), + [ + ({}, (False, pipeline_omnivoice._ASR_MODEL_NAME, None)), + ( + {"omnivoice_asr": {"load_asr_on_startup": False, "asr_device": None}}, + (False, pipeline_omnivoice._ASR_MODEL_NAME, None), + ), + ( + { + "omnivoice_asr": { + "load_asr_on_startup": True, + "asr_model_name": "local/whisper", + "asr_device": "cuda:1", + } + }, + (True, "local/whisper", "cuda:1"), + ), + ], +) +def test_asr_config_defaults_and_values(additional_config, expected): + assert pipeline_omnivoice._parse_asr_config(additional_config) == expected + + +@pytest.mark.parametrize( + "additional_config", + [ + {"omnivoice_asr": {"load_asr_on_startup": "true"}}, + {"omnivoice_asr": {"asr_model_name": " "}}, + {"omnivoice_asr": {"asr_device": ""}}, + {"omnivoice_asr": []}, + {"omnivoice_asr": None}, + ], +) +def test_asr_config_rejects_invalid_values(additional_config): + with pytest.raises((TypeError, ValueError)): + pipeline_omnivoice._parse_asr_config(additional_config) + + +def test_asr_loader_uses_configured_model_and_device(pipeline, monkeypatch): + asr, load_calls = _install_fake_asr(monkeypatch) + pipeline._asr_model_name = "local/whisper" + pipeline._asr_device = "cpu" + + pipeline._load_asr_pipeline() + + assert asr is pipeline._asr_pipeline + assert asr.model.to_calls == [torch.device("cpu")] + assert asr.device == torch.device("cpu") + assert load_calls == [ + ( + ("automatic-speech-recognition",), + {"model": "local/whisper", "dtype": torch.float32, "device": "cpu"}, + ) + ] + + +def test_eager_asr_is_loaded_during_initialization(monkeypatch): + model = pipeline_omnivoice.OmniVoicePipeline.__new__(pipeline_omnivoice.OmniVoicePipeline) + torch.nn.Module.__init__(model) + load_calls = [] + + def load_asr(self): + load_calls.append((self._asr_model_name, self._asr_device)) + + monkeypatch.setattr(pipeline_omnivoice.OmniVoicePipeline, "_load_asr_pipeline", load_asr) + + model._initialize_asr( + { + "omnivoice_asr": { + "load_asr_on_startup": True, + "asr_model_name": "local/whisper", + "asr_device": "cpu", + } + } + ) + + assert load_calls == [("local/whisper", "cpu")] + + +def test_lazy_asr_is_not_loaded_during_initialization(monkeypatch): + model = pipeline_omnivoice.OmniVoicePipeline.__new__(pipeline_omnivoice.OmniVoicePipeline) + torch.nn.Module.__init__(model) + load_calls = [] + + def load_asr(self): + load_calls.append(True) + + monkeypatch.setattr(pipeline_omnivoice.OmniVoicePipeline, "_load_asr_pipeline", load_asr) + + model._initialize_asr({"omnivoice_asr": {"load_asr_on_startup": False}}) + + assert load_calls == [] + + +def test_lazy_asr_uses_worker_device_on_first_request(pipeline, monkeypatch): + asr, load_calls = _install_fake_asr(monkeypatch) + pipeline._asr_device = None + + assert pipeline._resolve_ref_text((np.zeros(4, dtype=np.float32), 16000), None) == "hello there" + assert asr is pipeline._asr_pipeline + assert load_calls[0][1]["device"] == torch.device("cpu") + + +@pytest.mark.parametrize( + "waveform", + [ + np.arange(4, dtype=np.float32)[None, :], + torch.arange(4, dtype=torch.float32)[None, :], + ], +) +def test_asr_input_accepts_numpy_and_torch_waveforms_without_mutating(waveform, pipeline, monkeypatch): + asr, _ = _install_fake_asr(monkeypatch) + original = waveform.clone() if isinstance(waveform, torch.Tensor) else waveform.copy() + + pipeline._transcribe_ref_audio((waveform, 22050)) + + asr_input = asr.inputs[0] + assert asr_input["array"].shape == (4,) + np.testing.assert_array_equal(asr_input["array"], np.arange(4, dtype=np.float32)) + asr_input["array"][0] = 100 + if isinstance(waveform, torch.Tensor): + assert torch.equal(waveform, original) + else: + np.testing.assert_array_equal(waveform, original) + + +@pytest.mark.parametrize( + ("device", "dtype"), + [ + (torch.device("cpu"), torch.float32), + (torch.device("cuda:0"), torch.float16), + ("xpu:0", torch.float16), + ], +) +def test_loader_receives_checkpoint_device_and_dtype(pipeline, monkeypatch, device, dtype): + asr, load_calls = _install_fake_asr(monkeypatch) + pipeline._asr_device = str(device) + + pipeline._load_asr_pipeline() + + assert load_calls[0][1]["model"] == "openai/whisper-large-v3-turbo" + assert load_calls[0][1]["device"] == str(device) + assert load_calls[0][1]["dtype"] is dtype + assert asr.model.to_calls == [torch.device(device)] + assert asr.device == torch.device(device) + + +def test_transcript_is_combined_before_target_text(pipeline, monkeypatch): + _install_fake_asr(monkeypatch) + ref_audio = (np.zeros(4, dtype=np.float32), 16000) + + ref_text = pipeline._resolve_ref_text(ref_audio, None) + + assert pipeline_omnivoice._combine_text("hello", ref_text) == "hello there hello" + + +def test_voice_clone_duration_uses_resolved_text_and_reference_token_count(pipeline): + calls = [] + + class DurationEstimator: + def estimate_duration(self, text, ref_text, ref_duration): + calls.append((text, ref_text, ref_duration)) + return 17.8 + + pipeline.duration_estimator = DurationEstimator() + ref_audio_tokens = torch.zeros(8, 73, dtype=torch.long) + + assert pipeline._estimate_target_len("target", "transcribed reference", ref_audio_tokens) == 17 + assert calls == [("target", "transcribed reference", 73)] + + +def test_text_only_duration_keeps_fallback_inputs(pipeline): + calls = [] + + class DurationEstimator: + def estimate_duration(self, text, ref_text, ref_duration): + calls.append((text, ref_text, ref_duration)) + return 4 + + pipeline.duration_estimator = DurationEstimator() + + assert pipeline._estimate_target_len("target", None, None) == 4 + assert calls == [("target", "Nice to meet you.", 25)] diff --git a/tests/model_executor/models/omnivoice/test_audio_processing.py b/tests/model_executor/models/omnivoice/test_audio_processing.py new file mode 100644 index 00000000000..39909940255 --- /dev/null +++ b/tests/model_executor/models/omnivoice/test_audio_processing.py @@ -0,0 +1,609 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from __future__ import annotations + +from collections import OrderedDict +from pathlib import Path +from types import SimpleNamespace + +import numpy as np +import pytest +import soundfile as sf +import torch + +from vllm_omni.diffusion.models.omnivoice import audio as audio_utils +from vllm_omni.diffusion.models.omnivoice import pipeline_omnivoice +from vllm_omni.diffusion.models.omnivoice.audio import ( + PreparedReferenceAudio, + add_reference_punctuation, + postprocess_generated_audio, + prepare_reference_audio, + remove_silence, +) +from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch +from vllm_omni.diffusion.worker.utils import StepRequestState +from vllm_omni.inputs.data import OmniDiffusionSamplingParams +from vllm_omni.utils.speaker_cache import SpeakerEmbeddingCache + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +_SAMPLE_RATE = 24000 +_HOP_LENGTH = 960 + + +def _prepare(waveform, sample_rate=_SAMPLE_RATE): + return prepare_reference_audio( + waveform, + sample_rate, + target_sample_rate=_SAMPLE_RATE, + hop_length=_HOP_LENGTH, + trim_long=False, + ) + + +def test_reference_audio_normalization_and_hop_alignment(): + waveform = np.full(1010, 0.05, dtype=np.float32) + + prepared = _prepare(waveform) + + assert prepared.waveform.shape == (1, _HOP_LENGTH) + assert prepared.original_rms == pytest.approx(0.05, abs=1e-6) + assert np.sqrt(np.mean(prepared.waveform**2)) == pytest.approx(0.1, abs=1e-4) + + +def test_reference_audio_normalization_clips_peaky_quiet_waveform(): + waveform = np.zeros(7056, dtype=np.float32) + waveform[waveform.size // 2] = 1.0 + + prepared = _prepare(waveform) + + assert prepared.original_rms == pytest.approx(0.01190476, abs=1e-8) + assert np.max(np.abs(prepared.waveform)) == 1.0 + + +def test_reference_audio_at_or_above_target_rms_is_not_rescaled(): + prepared = _prepare(np.full(960, 0.2, dtype=np.float32)) + + assert prepared.original_rms == pytest.approx(0.2, abs=1e-6) + assert np.sqrt(np.mean(prepared.waveform**2)) == pytest.approx(0.2, abs=1e-4) + + +def test_reference_audio_resamples_and_downmixes(): + waveform = np.full((2, 800), 0.1, dtype=np.float32) + + prepared = _prepare(waveform, sample_rate=16000) + + assert prepared.waveform.shape == (1, 960) + assert prepared.sample_rate == _SAMPLE_RATE + + +@pytest.mark.parametrize("waveform", [np.full(960, 0.1, dtype=np.float32), torch.full((960,), 0.1)]) +def test_reference_audio_does_not_mutate_input(waveform): + original = waveform.clone() if isinstance(waveform, torch.Tensor) else waveform.copy() + + _prepare(waveform) + + if isinstance(waveform, torch.Tensor): + assert torch.equal(waveform, original) + else: + np.testing.assert_array_equal(waveform, original) + + +def test_reference_audio_rejects_empty_after_silence_removal(): + with pytest.raises(ValueError, match="empty after silence removal"): + _prepare(np.zeros(960, dtype=np.float32)) + + +def test_reference_audio_rejects_channels_that_cancel_when_downmixed(): + waveform = np.stack([np.full(960, 0.1), np.full(960, -0.1)]) + + with pytest.raises(ValueError, match="empty after silence removal"): + _prepare(waveform) + + +def test_reference_audio_matches_official_asset_preparation(): + asset_path = Path(__file__).resolve().parents[4] / "tests/assets/qwen3_tts/clone_2.wav" + waveform, sample_rate = sf.read(asset_path, always_2d=False) + + prepared = prepare_reference_audio( + waveform, + sample_rate, + target_sample_rate=_SAMPLE_RATE, + hop_length=_HOP_LENGTH, + trim_long=True, + ) + + assert prepared.waveform.shape == (1, 179520) + assert prepared.original_rms == pytest.approx(0.07468347, abs=1e-5) + + +@pytest.mark.parametrize( + ("text", "expected"), + [ + ("hello", "hello."), + ("你好", "你好。"), + ("hello!", "hello!"), + ("hello,", "hello,"), + ("hello]", "hello]"), + ("", ""), + ], +) +def test_reference_punctuation(text, expected): + assert add_reference_punctuation(text) == expected + + +def test_remove_silence_uses_reference_thresholds(): + waveform = np.concatenate( + [ + np.zeros(2400, dtype=np.float32), + np.full(12000, 0.2, dtype=np.float32), + np.zeros(9600, dtype=np.float32), + np.full(12000, 0.2, dtype=np.float32), + np.zeros(7200, dtype=np.float32), + ] + )[np.newaxis, :] + + processed = remove_silence( + waveform, + _SAMPLE_RATE, + middle_silence_ms=200, + leading_silence_ms=100, + trailing_silence_ms=200, + ) + + assert processed.shape[-1] < waveform.shape[-1] + assert processed.shape[-1] > 24000 + + +@pytest.mark.parametrize( + ("gap_ms", "expected_removed_ms"), + [(199, 0), (200, 0), (500, 100)], +) +def test_remove_silence_preserves_middle_gap_boundary(gap_ms, expected_removed_ms): + active = np.full(_SAMPLE_RATE, 0.2, dtype=np.float32) + gap = np.zeros(round(gap_ms * _SAMPLE_RATE / 1000), dtype=np.float32) + waveform = np.concatenate([active, gap, active])[np.newaxis, :] + + processed = remove_silence( + waveform, + _SAMPLE_RATE, + middle_silence_ms=200, + leading_silence_ms=0, + trailing_silence_ms=0, + ) + + removed_ms = round(1000 * (waveform.shape[-1] - processed.shape[-1]) / _SAMPLE_RATE) + assert removed_ms == expected_removed_ms + + +@pytest.mark.parametrize( + ("gap_level", "is_silent"), + [ + (10 ** (-50 / 20) * 0.99, True), + (10 ** (-50 / 20), True), + (10 ** (-50 / 20) * 1.01, False), + ], +) +def test_remove_silence_uses_reference_threshold(gap_level, is_silent): + active = np.full(_SAMPLE_RATE, 0.2, dtype=np.float32) + gap = np.full(round(250 * _SAMPLE_RATE / 1000), gap_level, dtype=np.float32) + waveform = np.concatenate([active, gap, active])[np.newaxis, :] + + silent_ranges = audio_utils._detect_silence_ranges( + waveform, + _SAMPLE_RATE, + min_silence_ms=200, + silence_threshold_db=-50, + seek_step_ms=10, + ) + + assert bool(silent_ranges) is is_silent + + +@pytest.mark.parametrize("duration_s", [19.9, 20.0, 20.1]) +def test_trim_long_audio_only_runs_above_twenty_seconds(duration_s): + waveform = np.full((1, round(duration_s * _SAMPLE_RATE)), 0.2, dtype=np.float32) + + trimmed = audio_utils.trim_long_audio(waveform, _SAMPLE_RATE) + + if duration_s <= 20.0: + assert trimmed.shape == waveform.shape + else: + assert trimmed.shape[-1] == 15 * _SAMPLE_RATE + + +def test_generated_audio_scales_with_reference_rms_and_pads(): + audio = np.full((1, 10000), 0.2, dtype=np.float32) + + processed = postprocess_generated_audio( + audio, + sample_rate=_SAMPLE_RATE, + reference_rms=0.05, + ) + + # NumPy keeps the source sample count and uses sample-based padding. + assert processed.shape[-1] == 14800 + assert processed[0, 5000] == pytest.approx(0.1, abs=1e-5) + assert processed[0, 0] == 0.0 + assert processed[0, -1] == 0.0 + + +def test_generated_audio_keeps_reference_rms_at_or_above_target(): + processed = postprocess_generated_audio( + np.full((1, 10000), 0.2, dtype=np.float32), + sample_rate=_SAMPLE_RATE, + reference_rms=0.2, + ) + + assert processed[0, 5000] == pytest.approx(0.2, abs=3e-5) + + +def test_generated_audio_uses_output_silence_threshold(monkeypatch): + calls = [] + + def fake_remove_silence(audio, sample_rate, **kwargs): + calls.append((sample_rate, kwargs)) + return audio + + monkeypatch.setattr(audio_utils, "remove_silence", fake_remove_silence) + audio_utils.postprocess_generated_audio( + np.ones((1, 1000), dtype=np.float32), + sample_rate=_SAMPLE_RATE, + reference_rms=0.05, + ) + + assert calls == [ + ( + _SAMPLE_RATE, + { + "middle_silence_ms": 500, + "leading_silence_ms": 100, + "trailing_silence_ms": 100, + }, + ) + ] + + +def test_generated_zero_audio_remains_finite(): + processed = postprocess_generated_audio( + np.zeros((1, 1000), dtype=np.float32), + sample_rate=_SAMPLE_RATE, + reference_rms=0.05, + ) + + assert np.isfinite(processed).all() + + +class _FakeASR: + def __init__(self, text: str): + self.text = text + self.inputs: list[dict[str, object]] = [] + + def __call__(self, audio_input): + self.inputs.append(audio_input) + return {"text": self.text} + + +class _FakeGenerator: + def __init__(self, num_codebooks: int): + self.num_codebooks = num_codebooks + + def __call__(self, **kwargs): + target_len = sum(kwargs["target_lens"]) + return torch.zeros((1, self.num_codebooks, target_len), dtype=torch.long) + + +class _FakeDecoder: + def __call__(self, tokens): + return torch.full((1, 1, 10000), 0.2, dtype=torch.float32) + + +class _FakeDurationEstimator: + def __init__(self): + self.calls = [] + + def estimate_duration(self, text, ref_text, ref_audio_tokens): + self.calls.append((text, ref_text, ref_audio_tokens)) + return 4 + + +def _build_fake_pipeline(monkeypatch, prepared_waveform, *, asr_text="reference transcript"): + model = pipeline_omnivoice.OmniVoicePipeline.__new__(pipeline_omnivoice.OmniVoicePipeline) + torch.nn.Module.__init__(model) + model.device = torch.device("cpu") + model.config = SimpleNamespace(num_audio_codebook=2, audio_mask_id=-1) + model.audio_tokenizer = SimpleNamespace(config=SimpleNamespace(sample_rate=_SAMPLE_RATE, hop_length=_HOP_LENGTH)) + model.tokenizer = SimpleNamespace(encode=lambda text: SimpleNamespace(ids=[1, 2])) + model.generator = _FakeGenerator(num_codebooks=2) + model.decoder = _FakeDecoder() + model.duration_estimator = _FakeDurationEstimator() + model._asr_pipeline = _FakeASR(asr_text) + model.num_step = 1 + model.guidance_scale = 1.0 + model.t_shift = 1.0 + model.layer_penalty_factor = 0.0 + model.position_temperature = 1.0 + model.class_temperature = 1.0 + model.sample_rate = _SAMPLE_RATE + model._inline_reference_cache = OrderedDict() + + prepare_calls = [] + encoded_calls = [] + + def fake_prepare(waveform, sample_rate, **kwargs): + prepare_calls.append((waveform, sample_rate, kwargs)) + return PreparedReferenceAudio(prepared_waveform, _SAMPLE_RATE, 0.07) + + def fake_encode(self, audio_signal, sample_rate): + encoded_calls.append((audio_signal.detach().clone(), sample_rate)) + return torch.zeros((2, 3), dtype=torch.long) + + monkeypatch.setattr(pipeline_omnivoice, "prepare_reference_audio", fake_prepare) + monkeypatch.setattr(pipeline_omnivoice.OmniVoicePipeline, "_encode_ref_audio", fake_encode) + return model, prepare_calls, encoded_calls + + +def _request(*prompts): + return DiffusionRequestBatch( + requests=[ + OmniDiffusionRequest( + prompt=prompt, + sampling_params=OmniDiffusionSamplingParams(), + request_id=str(index), + ) + for index, prompt in enumerate(prompts) + ], + ) + + +def test_pipeline_uses_one_prepared_waveform_for_asr_and_tokenizer(monkeypatch): + prepared = np.arange(_HOP_LENGTH, dtype=np.float32)[np.newaxis, :] + model, prepare_calls, encoded_calls = _build_fake_pipeline(monkeypatch, prepared) + + result = model.forward( + _request( + { + "prompt": "hello", + "multi_modal_data": {"audio": (np.ones(1000, dtype=np.float32), 16000)}, + } + ) + ) + + assert result[0].error is None + assert len(prepare_calls) == 1 + assert prepare_calls[0][1] == 16000 + assert prepare_calls[0][2]["trim_long"] is True + assert prepare_calls[0][2]["target_sample_rate"] == _SAMPLE_RATE + assert prepare_calls[0][2]["hop_length"] == _HOP_LENGTH + assert len(model._asr_pipeline.inputs) == 1 + np.testing.assert_array_equal(model._asr_pipeline.inputs[0]["array"], prepared[0]) + assert model._asr_pipeline.inputs[0]["sampling_rate"] == _SAMPLE_RATE + assert len(encoded_calls) == 1 + torch.testing.assert_close(encoded_calls[0][0], torch.from_numpy(prepared)) + assert encoded_calls[0][1] == _SAMPLE_RATE + assert model.duration_estimator.calls == [("hello", "reference transcript.", 3)] + + +def test_pipeline_explicit_reference_text_prepares_audio_without_asr(monkeypatch): + prepared = np.ones((1, _HOP_LENGTH), dtype=np.float32) + model, prepare_calls, encoded_calls = _build_fake_pipeline(monkeypatch, prepared) + + result = model.forward( + _request( + { + "prompt": "hello", + "ref_audio": (np.ones(1000, dtype=np.float32), 16000), + "ref_text": "caller supplied transcript", + } + ) + ) + + assert result[0].error is None + assert model._asr_pipeline.inputs == [] + assert prepare_calls[0][2]["trim_long"] is False + assert len(encoded_calls) == 1 + assert model.duration_estimator.calls == [("hello", "caller supplied transcript.", 3)] + + +def test_pipeline_text_only_skips_reference_audio_processing(monkeypatch): + model, prepare_calls, encoded_calls = _build_fake_pipeline( + monkeypatch, + np.ones((1, _HOP_LENGTH), dtype=np.float32), + ) + + result = model.forward(_request("hello")) + + assert result[0].error is None + torch.testing.assert_close(result[0].output, torch.full((1, 1, 10000), 0.2)) + assert prepare_calls == [] + assert encoded_calls == [] + assert model._asr_pipeline.inputs == [] + assert model.duration_estimator.calls == [("hello", "Nice to meet you.", 25)] + + +def test_pipeline_text_only_bypasses_audio_postprocessing(monkeypatch): + model, _, _ = _build_fake_pipeline( + monkeypatch, + np.ones((1, _HOP_LENGTH), dtype=np.float32), + ) + + def fail_postprocess(*args, **kwargs): + raise AssertionError("text-only output must bypass voice-cloning postprocessing") + + monkeypatch.setattr(pipeline_omnivoice, "postprocess_generated_audio", fail_postprocess) + + result = model.forward(_request("hello")) + + assert result[0].error is None + + +def test_pipeline_reuses_inline_reference_cache(monkeypatch): + prepared = np.arange(_HOP_LENGTH, dtype=np.float32)[np.newaxis, :] + model, prepare_calls, encoded_calls = _build_fake_pipeline(monkeypatch, prepared) + request = _request( + { + "prompt": "hello", + "ref_audio": (np.ones(1000, dtype=np.float32), 16000), + } + ) + + first = model.forward(request) + second = model.forward(request) + + assert first[0].error is None + assert second[0].error is None + assert len(prepare_calls) == 2 + assert len(model._asr_pipeline.inputs) == 1 + assert len(encoded_calls) == 1 + + +def test_pipeline_keeps_explicit_and_asr_cache_entries_separate(monkeypatch): + prepared = np.arange(_HOP_LENGTH, dtype=np.float32)[np.newaxis, :] + model, prepare_calls, encoded_calls = _build_fake_pipeline(monkeypatch, prepared) + explicit_request = _request( + { + "prompt": "hello", + "ref_audio": (np.ones(1000, dtype=np.float32), 16000), + "ref_text": "caller supplied transcript", + } + ) + missing_request = _request( + { + "prompt": "hello", + "ref_audio": (np.ones(1000, dtype=np.float32), 16000), + } + ) + + model.forward(explicit_request) + model.forward(missing_request) + model.forward(missing_request) + + assert len(prepare_calls) == 3 + assert len(encoded_calls) == 2 + assert len(model._asr_pipeline.inputs) == 1 + assert {key[0] for key in model._inline_reference_cache} == {"explicit", "asr"} + + +def test_named_cache_separates_explicit_and_asr_preparation(monkeypatch): + prepared = np.arange(_HOP_LENGTH, dtype=np.float32)[np.newaxis, :] + model, prepare_calls, encoded_calls = _build_fake_pipeline(monkeypatch, prepared) + model._speaker_cache = SpeakerEmbeddingCache(max_bytes=1024 * 1024) + explicit_request = _request( + { + "prompt": "hello", + "ref_audio": (np.ones(1000, dtype=np.float32), 16000), + "ref_text": "caller supplied transcript", + "voice_name": "alice", + "voice_created_at": 1, + } + ) + missing_request = _request( + { + "prompt": "hello", + "ref_audio": (np.ones(1000, dtype=np.float32), 16000), + "voice_name": "alice", + "voice_created_at": 1, + } + ) + + model.forward(explicit_request) + model.forward(missing_request) + model.forward(missing_request) + + assert model._speaker_cache.stats()["entries"] == 2 + assert len(prepare_calls) == 2 + assert len(encoded_calls) == 2 + assert len(model._asr_pipeline.inputs) == 1 + + +def test_inline_reference_cache_eviction_is_bounded(monkeypatch): + model, _, _ = _build_fake_pipeline( + monkeypatch, + np.ones((1, _HOP_LENGTH), dtype=np.float32), + ) + + for index in range(pipeline_omnivoice._INLINE_CACHE_MAX_ENTRIES + 1): + model._put_inline_cache(("asr", index), {"index": index}) + + assert len(model._inline_reference_cache) == pipeline_omnivoice._INLINE_CACHE_MAX_ENTRIES + assert ("asr", 0) not in model._inline_reference_cache + assert ("asr", pipeline_omnivoice._INLINE_CACHE_MAX_ENTRIES) in model._inline_reference_cache + + +def test_inline_reference_cache_key_preserves_original_rms(): + waveform = np.ones((1, _HOP_LENGTH), dtype=np.float32) + quiet = PreparedReferenceAudio(waveform, _SAMPLE_RATE, 0.05) + loud = PreparedReferenceAudio(waveform, _SAMPLE_RATE, 0.08) + + assert pipeline_omnivoice.OmniVoicePipeline._inline_cache_key( + quiet, "asr" + ) != pipeline_omnivoice.OmniVoicePipeline._inline_cache_key(loud, "asr") + + +@pytest.mark.parametrize("cache_kind", ["inline", "named"]) +def test_batch_keeps_reference_rms_per_request_and_on_cache_hits(monkeypatch, cache_kind): + model, _, encoded_calls = _build_fake_pipeline(monkeypatch, np.ones((1, _HOP_LENGTH), dtype=np.float32)) + model._speaker_cache = SpeakerEmbeddingCache(max_bytes=1024 * 1024) + + def prepare(waveform, sample_rate, **kwargs): + return PreparedReferenceAudio(waveform[np.newaxis, :], sample_rate, float(waveform[0])) + + monkeypatch.setattr(pipeline_omnivoice, "prepare_reference_audio", prepare) + prompts = ["hello", {"prompt": ""}] + for index, rms in enumerate((0.03, 0.08)): + prompt = { + "prompt": "hello", + "ref_audio": (np.full(_HOP_LENGTH, rms, dtype=np.float32), _SAMPLE_RATE), + } + if cache_kind == "named": + prompt["voice_name"] = f"speaker-{index}" + prompts.append(prompt) + + for _ in range(2): + outputs = model.forward(_request(*prompts)) + assert len(outputs) == 4 + torch.testing.assert_close(outputs[0].output, torch.full((1, 1, 10000), 0.2)) + assert outputs[1].error == "Empty text prompt" + for output, rms in zip(outputs[2:], (0.03, 0.08)): + assert output.error is None + assert output.output.shape == (1, 1, 14800) + assert output.output[0, 0, 5000].item() == pytest.approx(0.2 * rms / 0.1, abs=1e-5) + assert len(encoded_calls) == 2 + assert len(model._asr_pipeline.inputs) == 2 + + +@pytest.mark.parametrize("clone", [False, True]) +def test_step_execution_preserves_reference_postprocessing(monkeypatch, clone): + model, prepare_calls, _ = _build_fake_pipeline(monkeypatch, np.ones((1, _HOP_LENGTH), dtype=np.float32)) + prompt: str | dict = "hello" + if clone: + prompt = { + "prompt": "hello", + "ref_audio": (np.ones(_HOP_LENGTH, dtype=np.float32), _SAMPLE_RATE), + } + else: + + def fail_postprocess(*args, **kwargs): + raise AssertionError("text-only output must bypass voice-cloning postprocessing") + + monkeypatch.setattr(pipeline_omnivoice, "postprocess_generated_audio", fail_postprocess) + + state = StepRequestState( + request_id="step-request", + prompt=prompt, + sampling=OmniDiffusionSamplingParams(num_inference_steps=1), + ) + model.prepare_encode(state) + output = model.post_decode(state) + + assert output.error is None + if clone: + assert state.extra["reference_rms"] == 0.07 + assert output.output.shape == (1, 1, 14800) + assert output.output[0, 0, 5000].item() == pytest.approx(0.14, abs=1e-5) + else: + assert state.extra["reference_rms"] is None + assert prepare_calls == [] + torch.testing.assert_close(output.output, torch.full((1, 1, 10000), 0.2)) diff --git a/vllm_omni/deploy/omnivoice.yaml b/vllm_omni/deploy/omnivoice.yaml index df9e6a5daf0..0028529a42d 100644 --- a/vllm_omni/deploy/omnivoice.yaml +++ b/vllm_omni/deploy/omnivoice.yaml @@ -11,5 +11,12 @@ stages: trust_remote_code: true distributed_executor_backend: "mp" dtype: "float32" + additional_config: + omnivoice_asr: + # false defers Whisper until ref_audio is sent without ref_text; + # it does not disable automatic transcription. Set true to load at startup. + load_asr_on_startup: false + asr_model_name: "openai/whisper-large-v3-turbo" + asr_device: null max_num_seqs: 8 request_batch_max_wait_ms: 10 diff --git a/vllm_omni/diffusion/models/omnivoice/audio.py b/vllm_omni/diffusion/models/omnivoice/audio.py new file mode 100644 index 00000000000..b20a3f93c6a --- /dev/null +++ b/vllm_omni/diffusion/models/omnivoice/audio.py @@ -0,0 +1,405 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Audio preparation and output processing for OmniVoice.""" + +from __future__ import annotations + +from dataclasses import dataclass + +import numpy as np +import torch +import torchaudio + +_END_PUNCTUATION = { + ";", + ":", + ",", + ".", + "!", + "?", + "…", + ")", + "]", + "}", + '"', + "'", + "“", + "”", + "‘", + "’", + "；", + "：", + "，", + "。", + "！", + "？", + "、", + "）", + "】", +} + + +@dataclass(frozen=True) +class PreparedReferenceAudio: + """Reference waveform shared by ASR and the audio tokenizer.""" + + waveform: np.ndarray + sample_rate: int + original_rms: float + + +def _segment_length_ms(num_samples: int, sample_rate: int) -> int: + """Return the rounded millisecond length used by the reference algorithm.""" + return round(1000 * num_samples / sample_rate) + + +def _frame_index(milliseconds: int, sample_rate: int, num_samples: int) -> int: + """Map a millisecond position to a bounded sample index.""" + return min(num_samples, max(0, int(milliseconds * sample_rate / 1000.0))) + + +def _slice_milliseconds(audio: np.ndarray, start_ms: int, end_ms: int | None, sample_rate: int) -> np.ndarray: + num_samples = audio.shape[-1] + start = _frame_index(start_ms, sample_rate, num_samples) + end = num_samples if end_ms is None else _frame_index(end_ms, sample_rate, num_samples) + return audio[:, start:end] + + +def _rms(audio: np.ndarray) -> float: + if audio.size == 0: + return 0.0 + squared = np.square(np.asarray(audio, dtype=np.float32), dtype=np.float32) + return float(np.sqrt(np.mean(squared, dtype=np.float32))) + + +def _dbfs(audio: np.ndarray) -> float: + rms = _rms(audio) + if rms == 0.0: + return -float("inf") + return 20.0 * float(np.log10(rms)) + + +def _window_rms(audio: np.ndarray, starts_ms: list[int], window_ms: int, sample_rate: int) -> np.ndarray: + starts = np.fromiter( + (_frame_index(start, sample_rate, audio.shape[-1]) for start in starts_ms), + dtype=np.int64, + ) + ends = np.fromiter( + (_frame_index(start + window_ms, sample_rate, audio.shape[-1]) for start in starts_ms), + dtype=np.int64, + ) + squared = np.square(audio, dtype=np.float32) + cumulative = np.concatenate( + [np.zeros((audio.shape[0], 1), dtype=np.float64), np.cumsum(squared, axis=-1, dtype=np.float64)], + axis=-1, + ) + sums = cumulative[:, ends] - cumulative[:, starts] + counts = ends - starts + rms = np.zeros(len(starts_ms), dtype=np.float32) + valid = counts > 0 + if np.any(valid): + means = sums[:, valid].sum(axis=0) / (counts[valid] * audio.shape[0]) + rms[valid] = np.sqrt(means).astype(np.float32) + return rms + + +def _detect_silence_ranges( + audio: np.ndarray, + sample_rate: int, + *, + min_silence_ms: int, + silence_threshold_db: float, + seek_step_ms: int, +) -> list[list[int]]: + """Find silent ranges with the same window and range rules as pydub.""" + segment_length_ms = _segment_length_ms(audio.shape[-1], sample_rate) + if segment_length_ms < min_silence_ms: + return [] + + threshold = 10.0 ** (silence_threshold_db / 20.0) + last_start = segment_length_ms - min_silence_ms + silence_starts = list(range(0, last_start + 1, seek_step_ms)) + if last_start % seek_step_ms: + silence_starts.append(last_start) + + window_rms = _window_rms(audio, silence_starts, min_silence_ms, sample_rate) + silent_starts = [start for start, rms in zip(silence_starts, window_rms) if rms <= threshold] + if not silent_starts: + return [] + + silent_ranges: list[list[int]] = [] + previous = silent_starts[0] + current_start = previous + for start in silent_starts[1:]: + continuous = start == previous + seek_step_ms + has_gap = start > previous + min_silence_ms + if not continuous and has_gap: + silent_ranges.append([current_start, previous + min_silence_ms]) + current_start = start + previous = start + silent_ranges.append([current_start, previous + min_silence_ms]) + return silent_ranges + + +def _detect_nonsilent_ranges( + audio: np.ndarray, + sample_rate: int, + *, + min_silence_ms: int, + silence_threshold_db: float, + seek_step_ms: int, +) -> list[list[int]]: + segment_length_ms = _segment_length_ms(audio.shape[-1], sample_rate) + silent_ranges = _detect_silence_ranges( + audio, + sample_rate, + min_silence_ms=min_silence_ms, + silence_threshold_db=silence_threshold_db, + seek_step_ms=seek_step_ms, + ) + if not silent_ranges: + return [[0, segment_length_ms]] + if silent_ranges[0] == [0, segment_length_ms]: + return [] + + nonsilent_ranges: list[list[int]] = [] + previous_end = 0 + for start, end in silent_ranges: + nonsilent_ranges.append([previous_end, start]) + previous_end = end + if previous_end != segment_length_ms: + nonsilent_ranges.append([previous_end, segment_length_ms]) + if nonsilent_ranges and nonsilent_ranges[0] == [0, 0]: + nonsilent_ranges.pop(0) + return nonsilent_ranges + + +def _split_on_silence( + audio: np.ndarray, + sample_rate: int, + *, + min_silence_ms: int, + silence_threshold_db: float, + keep_silence_ms: int, + seek_step_ms: int, +) -> np.ndarray: + ranges = _detect_nonsilent_ranges( + audio, + sample_rate, + min_silence_ms=min_silence_ms, + silence_threshold_db=silence_threshold_db, + seek_step_ms=seek_step_ms, + ) + output_ranges = [[start - keep_silence_ms, end + keep_silence_ms] for start, end in ranges] + for previous, current in zip(output_ranges, output_ranges[1:]): + if current[0] < previous[1]: + split = (previous[1] + current[0]) // 2 + previous[1] = split + current[0] = split + + segment_length_ms = _segment_length_ms(audio.shape[-1], sample_rate) + segments = [ + _slice_milliseconds(audio, max(start, 0), min(end, segment_length_ms), sample_rate) + for start, end in output_ranges + ] + if not segments: + return np.empty((audio.shape[0], 0), dtype=np.float32) + return np.concatenate(segments, axis=-1) + + +def _detect_leading_silence_ms(audio: np.ndarray, sample_rate: int, *, threshold_db: float, chunk_ms: int = 10) -> int: + segment_length_ms = _segment_length_ms(audio.shape[-1], sample_rate) + trim_ms = 0 + while trim_ms < segment_length_ms: + chunk = _slice_milliseconds(audio, trim_ms, trim_ms + chunk_ms, sample_rate) + if _dbfs(chunk) >= threshold_db: + break + trim_ms += chunk_ms + return min(trim_ms, segment_length_ms) + + +def remove_silence( + audio: np.ndarray, + sample_rate: int, + *, + middle_silence_ms: int, + leading_silence_ms: int, + trailing_silence_ms: int, +) -> np.ndarray: + """Remove long gaps and trim edge silence using official thresholds.""" + wave = np.asarray(audio, dtype=np.float32) + if wave.ndim == 1: + wave = wave[np.newaxis, :] + if wave.ndim != 2: + raise ValueError(f"OmniVoice audio must be 1D or 2D, got {wave.ndim} dimensions.") + if middle_silence_ms > 0: + wave = _split_on_silence( + wave, + sample_rate, + min_silence_ms=middle_silence_ms, + silence_threshold_db=-50, + keep_silence_ms=middle_silence_ms, + seek_step_ms=10, + ) + + start_idx = _detect_leading_silence_ms(wave, sample_rate, threshold_db=-50) + start_idx = max(0, start_idx - leading_silence_ms) + wave = _slice_milliseconds(wave, start_idx, None, sample_rate) + + wave = wave[:, ::-1] + start_idx = _detect_leading_silence_ms(wave, sample_rate, threshold_db=-50) + start_idx = max(0, start_idx - trailing_silence_ms) + wave = _slice_milliseconds(wave, start_idx, None, sample_rate) + wave = wave[:, ::-1] + + return np.ascontiguousarray(wave, dtype=np.float32) + + +def trim_long_audio( + audio: np.ndarray, + sample_rate: int, + *, + max_duration_s: float = 15.0, + min_duration_s: float = 3.0, + trim_threshold_s: float = 20.0, +) -> np.ndarray: + """Trim long reference audio at a suitable silence boundary.""" + if audio.shape[-1] / sample_rate <= trim_threshold_s: + return audio + + non_silent_ranges = _detect_nonsilent_ranges( + np.asarray(audio, dtype=np.float32), + sample_rate=sample_rate, + min_silence_ms=100, + silence_threshold_db=-40, + seek_step_ms=10, + ) + if not non_silent_ranges: + return audio + + max_ms = int(max_duration_s * 1000) + min_ms = int(min_duration_s * 1000) + best_split = 0 + for start, end in non_silent_ranges: + if start > best_split and start <= max_ms: + best_split = start + if end > max_ms: + break + + if best_split < min_ms: + best_split = min(max_ms, _segment_length_ms(audio.shape[-1], sample_rate)) + return _slice_milliseconds(np.asarray(audio, dtype=np.float32), 0, best_split, sample_rate) + + +def prepare_reference_audio( + waveform: np.ndarray | torch.Tensor, + sample_rate: int, + *, + target_sample_rate: int, + hop_length: int, + trim_long: bool, +) -> PreparedReferenceAudio: + """Prepare reference audio before ASR and audio-tokenizer encoding.""" + if isinstance(waveform, torch.Tensor): + waveform = waveform.detach().cpu().numpy() + waveform = np.array(waveform, dtype=np.float32, copy=True) + if waveform.ndim == 1: + waveform = waveform[np.newaxis, :] + elif waveform.ndim != 2: + raise ValueError(f"OmniVoice reference audio must be 1D or 2D, got {waveform.ndim} dimensions.") + if waveform.shape[0] > 1: + waveform = waveform.mean(axis=0, keepdims=True) + if not waveform.size or not np.any(waveform): + raise ValueError("Reference audio is empty after silence removal.") + + if sample_rate != target_sample_rate: + waveform = torchaudio.functional.resample( + torch.from_numpy(waveform), + orig_freq=sample_rate, + new_freq=target_sample_rate, + ).numpy() + + original_rms = float(np.sqrt(np.mean(waveform**2))) if waveform.size else 0.0 + if 0 < original_rms < 0.1: + waveform = waveform * (0.1 / original_rms) + waveform = np.clip(waveform, -1.0, 1.0) + + if trim_long: + waveform = trim_long_audio(waveform, target_sample_rate) + waveform = remove_silence( + waveform, + target_sample_rate, + middle_silence_ms=200, + leading_silence_ms=100, + trailing_silence_ms=200, + ) + if waveform.shape[-1] == 0: + raise ValueError("Reference audio is empty after silence removal.") + + remainder = waveform.shape[-1] % hop_length + if remainder: + waveform = waveform[:, :-remainder] + if waveform.shape[-1] == 0: + raise ValueError("Reference audio is shorter than one audio-tokenizer hop after preprocessing.") + + return PreparedReferenceAudio( + waveform=np.ascontiguousarray(waveform, dtype=np.float32), + sample_rate=target_sample_rate, + original_rms=original_rms, + ) + + +def add_reference_punctuation(text: str) -> str: + """Add an English or Chinese sentence terminator when one is missing.""" + text = text.strip() + if not text: + return text + if text[-1] in _END_PUNCTUATION: + return text + if any("\u4e00" <= char <= "\u9fff" for char in text): + return f"{text}。" + return f"{text}." + + +def postprocess_generated_audio( + audio: np.ndarray, + *, + sample_rate: int, + reference_rms: float, +) -> np.ndarray: + """Apply official OmniVoice output silence, volume, fade, and padding rules.""" + audio = np.array(audio, dtype=np.float32, copy=True) + if audio.ndim == 1: + audio = audio[np.newaxis, :] + if audio.ndim != 2: + raise ValueError(f"OmniVoice generated audio must be 1D or 2D, got {audio.ndim} dimensions.") + + audio = remove_silence( + audio, + sample_rate, + middle_silence_ms=500, + leading_silence_ms=100, + trailing_silence_ms=100, + ) + + if audio.shape[-1] == 0: + return audio + + if reference_rms < 0.1: + audio = audio * (reference_rms / 0.1) + + fade_samples = int(0.1 * sample_rate) + if fade_samples > 0: + fade_length = min(fade_samples, audio.shape[-1] // 2) + if fade_length > 0: + fade_in = np.linspace(0, 1, fade_length, dtype=np.float32)[np.newaxis, :] + fade_out = np.linspace(1, 0, fade_length, dtype=np.float32)[np.newaxis, :] + audio[:, :fade_length] *= fade_in + audio[:, -fade_length:] *= fade_out + + pad_samples = int(0.1 * sample_rate) + if pad_samples > 0: + padding = np.zeros((audio.shape[0], pad_samples), dtype=np.float32) + audio = np.concatenate([padding, audio, padding], axis=-1) + return np.ascontiguousarray(audio, dtype=np.float32) diff --git a/vllm_omni/diffusion/models/omnivoice/pipeline_omnivoice.py b/vllm_omni/diffusion/models/omnivoice/pipeline_omnivoice.py index bad1c6154e7..d94da8b070d 100644 --- a/vllm_omni/diffusion/models/omnivoice/pipeline_omnivoice.py +++ b/vllm_omni/diffusion/models/omnivoice/pipeline_omnivoice.py @@ -11,17 +11,20 @@ from __future__ import annotations +import hashlib import json import math import os import random import re -from collections.abc import Iterable, Sequence +from collections import OrderedDict +from collections.abc import Iterable, Mapping, Sequence from dataclasses import dataclass from typing import Any, ClassVar import numpy as np import torch +import torchaudio from tokenizers import Tokenizer as HFTokenizer from torch import nn from vllm.logger import init_logger @@ -29,6 +32,11 @@ from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig from vllm_omni.diffusion.distributed.utils import get_local_device from vllm_omni.diffusion.models.interface import SupportAudioOutput +from vllm_omni.diffusion.models.omnivoice.audio import ( + add_reference_punctuation, + postprocess_generated_audio, + prepare_reference_audio, +) from vllm_omni.diffusion.worker.input_batch import InputBatch from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch from vllm_omni.diffusion.worker.utils import StepRequestState @@ -48,10 +56,42 @@ except ImportError: HiggsAudioV2TokenizerModel = None -import torchaudio +try: + from transformers import pipeline as hf_pipeline +except ImportError: + hf_pipeline = None logger = init_logger(__name__) +_ASR_MODEL_NAME = "openai/whisper-large-v3-turbo" +_INLINE_CACHE_MAX_ENTRIES = 8 + + +def _parse_asr_config(additional_config: Mapping[str, object] | None) -> tuple[bool, str, str | None]: + """Return validated OmniVoice ASR settings from the model config.""" + if additional_config is None: + additional_config = {} + if not isinstance(additional_config, Mapping): + raise TypeError(f"additional_config must be a mapping or None, got {type(additional_config)!r}") + + raw_config = additional_config.get("omnivoice_asr", {}) + if not isinstance(raw_config, Mapping): + raise TypeError(f"additional_config['omnivoice_asr'] must be a mapping, got {type(raw_config)!r}") + + load_asr_on_startup = raw_config.get("load_asr_on_startup", False) + if not isinstance(load_asr_on_startup, bool): + raise TypeError("additional_config['omnivoice_asr']['load_asr_on_startup'] must be a bool") + + model_name = raw_config.get("asr_model_name", _ASR_MODEL_NAME) + if not isinstance(model_name, str) or not model_name.strip(): + raise ValueError("additional_config['omnivoice_asr']['asr_model_name'] must be a non-empty string") + + asr_device = raw_config.get("asr_device") + if asr_device is not None and (not isinstance(asr_device, str) or not asr_device.strip()): + raise ValueError("additional_config['omnivoice_asr']['asr_device'] must be a non-empty string") + + return load_asr_on_startup, model_name.strip(), asr_device.strip() if asr_device is not None else None + @dataclass class _PreparedOmniVoiceRequest: @@ -60,6 +100,7 @@ class _PreparedOmniVoiceRequest: cond_len: int target_len: int seed: int | None + reference_rms: float | None = None def get_omnivoice_post_process_func(od_config: OmniDiffusionConfig): @@ -159,6 +200,8 @@ def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""): self.od_config = od_config self.device = get_local_device() self.model_path = od_config.model + self._initialize_asr(getattr(od_config, "additional_config", None)) + self._inline_reference_cache: OrderedDict[tuple[object, ...], dict[str, object]] = OrderedDict() # Resolve model path (HF hub ID → local cache) if not os.path.isdir(self.model_path): @@ -206,6 +249,101 @@ def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""): self.class_temperature = self.config.class_temperature self.sample_rate = self.config.sample_rate + def _initialize_asr(self, additional_config: Mapping[str, object] | None) -> None: + load_asr_on_startup, self._asr_model_name, self._asr_device = _parse_asr_config(additional_config) + self._asr_pipeline = None + if load_asr_on_startup: + self._load_asr_pipeline() + + def _load_asr_pipeline(self): + """Load the configured reference-audio ASR pipeline.""" + if self._asr_pipeline is not None: + return self._asr_pipeline + asr_device = self._asr_device if self._asr_device is not None else self.device + if hf_pipeline is None: + raise RuntimeError( + "OmniVoice automatic transcription requires the Hugging Face " + f"ASR pipeline ({self._asr_model_name!r}) on device {asr_device}." + ) + + asr_dtype = torch.float16 if str(asr_device).lower().startswith(("cuda", "xpu")) else torch.float32 + logger.info( + "Loading OmniVoice ASR model %s on %s", + self._asr_model_name, + asr_device, + ) + try: + asr_pipeline = hf_pipeline( + "automatic-speech-recognition", + model=self._asr_model_name, + dtype=asr_dtype, + device=asr_device, + ) + # Transformers keeps a newly loaded pipeline model on CPU when a + # torch.distributed process group is already initialized. vLLM + # workers always have one, so explicitly honor the configured + # device after construction, as other auxiliary models do here. + target_device = torch.device(asr_device) + asr_pipeline.model = asr_pipeline.model.to(target_device) + asr_pipeline.device = target_device + except Exception as exc: + raise RuntimeError( + f"Failed to load OmniVoice ASR model {self._asr_model_name!r} on device {asr_device}: {exc}" + ) from exc + self._asr_pipeline = asr_pipeline + logger.info("OmniVoice ASR model loaded on %s", asr_device) + return self._asr_pipeline + + @torch.inference_mode() + def _transcribe_ref_audio(self, ref_audio) -> str: + """Transcribe a reference waveform with the lazily loaded ASR model.""" + waveform, sr = ref_audio + if isinstance(waveform, torch.Tensor): + waveform = waveform.detach().cpu().numpy() + waveform = np.squeeze(np.array(waveform, copy=True)) + result = self._load_asr_pipeline()( + { + "array": waveform, + "sampling_rate": int(sr), + } + ) + if not isinstance(result, dict) or "text" not in result: + raise RuntimeError("OmniVoice ASR returned a malformed result without a 'text' field.") + transcript = result["text"] + if not isinstance(transcript, str): + raise RuntimeError("OmniVoice ASR returned a malformed result: 'text' must be a string.") + transcript = transcript.strip() + if not transcript: + raise ValueError("OmniVoice ASR returned an empty reference transcription.") + return transcript + + def _resolve_ref_text(self, ref_audio, ref_text: str | None) -> str | None: + """Resolve missing reference text only when reference audio is present.""" + if ref_audio is not None and (ref_text is None or not ref_text.strip()): + logger.debug("Automatically transcribing OmniVoice reference audio") + return self._transcribe_ref_audio(ref_audio) + return ref_text + + def _estimate_target_len( + self, + text: str, + ref_text: str | None, + ref_audio_tokens: torch.Tensor | None, + ) -> int: + """Estimate target audio tokens using resolved voice-clone conditioning.""" + if ref_audio_tokens is None or not ref_text: + ref_text = "Nice to meet you." + num_ref_audio_tokens = 25 + else: + num_ref_audio_tokens = ref_audio_tokens.size(-1) + + target_len = self.duration_estimator.estimate_duration( + text, + ref_text, + num_ref_audio_tokens, + ) + return max(1, int(target_len)) + def _encode_ref_audio(self, audio_signal: torch.Tensor, sr: int) -> torch.Tensor: """Encode reference audio to 8-codebook tokens for voice cloning.""" if self.audio_tokenizer is None: @@ -226,6 +364,30 @@ def _encode_ref_audio(self, audio_signal: torch.Tensor, sr: int) -> torch.Tensor tokens = tokens.squeeze(0) # [8, T_ref] return tokens + @staticmethod + def _inline_cache_key(prepared_audio, preparation_mode: str) -> tuple[object, ...]: + waveform = np.ascontiguousarray(prepared_audio.waveform, dtype=np.float32) + digest = hashlib.sha256(waveform.tobytes()).digest() + return ( + preparation_mode, + digest, + tuple(waveform.shape), + int(prepared_audio.sample_rate), + float(prepared_audio.original_rms), + ) + + def _get_inline_cache(self, key: tuple[object, ...]) -> dict[str, object] | None: + cached = self._inline_reference_cache.pop(key, None) + if cached is not None: + self._inline_reference_cache[key] = cached + return cached + + def _put_inline_cache(self, key: tuple[object, ...], artifacts: dict[str, object]) -> None: + self._inline_reference_cache[key] = artifacts + self._inline_reference_cache.move_to_end(key) + while len(self._inline_reference_cache) > _INLINE_CACHE_MAX_ENTRIES: + self._inline_reference_cache.popitem(last=False) + def _prepare_request_input( self, prompt: Any, @@ -271,6 +433,10 @@ def _prepare_request_input( lang = mm_kwargs.get("lang") if instruct is None: instruct = mm_kwargs.get("instruct") + + if isinstance(ref_text, str): + ref_text = ref_text.strip() or None + if not text: return DiffusionOutput(error="Empty text prompt") lang = lang or "None" @@ -280,49 +446,119 @@ def _prepare_request_input( if not text: return DiffusionOutput(error="Empty text prompt") - target_len = self.duration_estimator.estimate_duration(text, "Nice to meet you.", 25) - target_len = max(1, int(target_len)) - - style_text = f"<|denoise|><|lang_start|>{lang}<|lang_end|><|instruct_start|>{instruct}<|instruct_end|>" - full_text = _combine_text(ref_text=ref_text, text=text) - wrapped_text = f"<|text_start|>{full_text}<|text_end|>" - style_tokens = self.tokenizer.encode(style_text).ids - text_tokens = _tokenize_with_nonverbal_tags(wrapped_text, self.tokenizer) - encoding_ids = style_tokens + text_tokens - text_tokens_tensor = torch.tensor(encoding_ids, dtype=torch.long, device=self.device) - text_len = text_tokens_tensor.shape[0] + device = self.device + num_cb = self.config.num_audio_codebook + mask_id = self.config.audio_mask_id ref_audio_tokens = None + reference_rms = None if ref_audio is not None: + needs_asr = ref_text is None if self.audio_tokenizer is None: raise RuntimeError( "Voice cloning requires transformers>=5.3.0. Try: uv pip install 'transformers>=5.3.0'" ) - cache_key = None + preparation_mode = "asr" if needs_asr else "explicit" + _cache_key = None if voice_name: - cache_key = self._speaker_cache.make_cache_key( + _cache_key = self._speaker_cache.make_cache_key( voice_name, - model_type="omnivoice", + model_type=f"omnivoice_{preparation_mode}", created_at=int(prompt.get("voice_created_at") or 0), ) - cached = self._speaker_cache.get(cache_key) - if cached is not None: - ref_audio_tokens = cached["ref_audio_tokens"].to(self.device) - cache_key = None + cached = self._speaker_cache.get(_cache_key) + if cached is not None and (not needs_asr or cached.get("ref_text")): + ref_audio_tokens = cached["ref_audio_tokens"].to(device) + reference_rms = cached["reference_rms"] + if needs_asr: + ref_text = cached["ref_text"] + _cache_key = None # hit → don't store again logger.debug("Speaker cache HIT for OmniVoice speaker '%s'", voice_name) if ref_audio_tokens is None: - audio_signal, sr = ref_audio - if isinstance(audio_signal, np.ndarray): - audio_signal = torch.from_numpy(audio_signal).float() - ref_audio_tokens = self._encode_ref_audio(audio_signal, int(sr)).to(self.device) - if cache_key is not None: - self._speaker_cache.put(cache_key, {"ref_audio_tokens": ref_audio_tokens.cpu()}) + audio_signal, sample_rate = ref_audio + try: + prepared_audio = prepare_reference_audio( + audio_signal, + int(sample_rate), + target_sample_rate=self.audio_tokenizer.config.sample_rate, + hop_length=self.audio_tokenizer.config.hop_length, + trim_long=needs_asr, + ) + except (RuntimeError, ValueError) as exc: + return DiffusionOutput(error=str(exc)) + reference_rms = prepared_audio.original_rms + reference_duration = prepared_audio.waveform.shape[-1] / prepared_audio.sample_rate + if reference_duration > 20.0: + logger.warning( + "OmniVoice reference audio is %.1fs long (>20s); this may increase memory use " + "and reduce cloning quality.", + reference_duration, + ) + + inline_cache_key = None + if not voice_name: + inline_cache_key = self._inline_cache_key(prepared_audio, preparation_mode) + cached = self._get_inline_cache(inline_cache_key) + if cached is not None: + ref_audio_tokens = cached["ref_audio_tokens"].to(device) + reference_rms = cached["reference_rms"] + if needs_asr: + ref_text = cached["ref_text"] + logger.debug("Inline OmniVoice reference cache HIT") + + if ref_audio_tokens is None and needs_asr: + try: + ref_text = self._resolve_ref_text( + (prepared_audio.waveform, prepared_audio.sample_rate), + ref_text, + ) + except (RuntimeError, ValueError) as exc: + return DiffusionOutput(error=str(exc)) + if ref_audio_tokens is None: + ref_audio_tokens = self._encode_ref_audio( + torch.from_numpy(prepared_audio.waveform), + prepared_audio.sample_rate, + ).to(device) + + # Store named and inline entries for the matching preparation mode. + if _cache_key is not None: + self._speaker_cache.put( + _cache_key, + { + "ref_audio_tokens": ref_audio_tokens.cpu(), + "ref_text": ref_text if needs_asr else None, + "reference_rms": reference_rms, + }, + ) logger.debug("Speaker cache STORE for OmniVoice speaker '%s'", voice_name) + elif not voice_name: + self._put_inline_cache( + inline_cache_key, + { + "ref_audio_tokens": ref_audio_tokens.cpu(), + "ref_text": ref_text if needs_asr else None, + "reference_rms": reference_rms, + }, + ) + logger.debug("Inline OmniVoice reference cache STORE") + + if ref_text: + ref_text = add_reference_punctuation(ref_text) + + target_len = self._estimate_target_len(text, ref_text, ref_audio_tokens) + + # Build text prompt with control tokens + style_text = f"<|denoise|><|lang_start|>{lang}<|lang_end|><|instruct_start|>{instruct}<|instruct_end|>" + full_text = _combine_text(ref_text=ref_text, text=text) + wrapped_text = f"<|text_start|>{full_text}<|text_end|>" + style_tokens = self.tokenizer.encode(style_text).ids + text_tokens = _tokenize_with_nonverbal_tags(wrapped_text, self.tokenizer) + encoding_ids = style_tokens + text_tokens + text_tokens = torch.tensor(encoding_ids, dtype=torch.long, device=device) + text_len = text_tokens.shape[0] - num_cb = self.config.num_audio_codebook - mask_id = self.config.audio_mask_id - text_ids = text_tokens_tensor.unsqueeze(0).repeat(num_cb, 1) + text_ids = text_tokens.unsqueeze(0).repeat(num_cb, 1) target_ids = torch.full((num_cb, target_len), mask_id, dtype=torch.long, device=self.device) cond_ids = ( torch.cat([text_ids, ref_audio_tokens, target_ids], dim=1) @@ -343,6 +579,7 @@ def _prepare_request_input( cond_len=cond_len, target_len=target_len, seed=seed, + reference_rms=reference_rms, ) def _collate_request_inputs( @@ -431,6 +668,7 @@ def prepare_encode(self, state: StepRequestState) -> StepRequestState: state.extra["target_len"] = target_len state.extra["audio_mask"] = audio_mask state.extra["tokens"] = tokens + state.extra["reference_rms"] = prepared_request.reference_rms return state def denoise_step(self, input_batch: InputBatch, *, states: Sequence[StepRequestState] | None = None, **kwargs: Any): @@ -540,11 +778,23 @@ def step_scheduler(self, state: StepRequestState, noise_pred: torch.Tensor, **kw state.latents = noise_pred state.step_index += 1 + def _postprocess_audio(self, audio: torch.Tensor, reference_rms: float) -> torch.Tensor: + audio_np = audio.squeeze(0).detach().cpu().float().numpy() + audio_np = postprocess_generated_audio( + audio_np, + sample_rate=self.sample_rate, + reference_rms=reference_rms, + ) + return torch.from_numpy(audio_np).unsqueeze(0) + def post_decode(self, state: StepRequestState, **kwargs: Any): tokens = state.extra["tokens"] if tokens.dim() == 2: tokens = tokens.unsqueeze(0) audio = self.decoder(tokens) + reference_rms = state.extra["reference_rms"] + if reference_rms is not None: + audio = self._postprocess_audio(audio, reference_rms) return DiffusionOutput(output=audio) @torch.inference_mode() @@ -596,6 +846,9 @@ def forward(self, req: DiffusionRequestBatch) -> list[DiffusionOutput]: for i, target_len in enumerate(batch_target_len): request_tokens = tokens[:, :, target_offset : target_offset + target_len] audio = self.decoder(request_tokens) + reference_rms = prepared_requests[i].reference_rms + if reference_rms is not None: + audio = self._postprocess_audio(audio, reference_rms) outputs[prepared_indices[i]] = DiffusionOutput(output=audio) target_offset += target_len return outputs diff --git a/vllm_omni/entrypoints/openai/serving_speech.py b/vllm_omni/entrypoints/openai/serving_speech.py index cf684aaaae4..97271148349 100644 --- a/vllm_omni/entrypoints/openai/serving_speech.py +++ b/vllm_omni/entrypoints/openai/serving_speech.py @@ -2586,8 +2586,7 @@ async def _create_diffusion_speech( # to allow tests to catch invalid voices generally request.voice = self._get_normalized_voice(request.voice) - # Assume that this will follow the same adapter pattern - # once all RFC is implemented + assert self._adapter is not None validation_error = validate_tts_request(self._adapter, request) if validation_error is not None: raise ValueError(validation_error) diff --git a/vllm_omni/entrypoints/openai/tts_adapters/omnivoice.py b/vllm_omni/entrypoints/openai/tts_adapters/omnivoice.py index b2baa72898e..8766204bd96 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/omnivoice.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/omnivoice.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """OmniVoice serving adapter (AR engine_client path). OmniVoice can also be served through the pure-diffusion engine via @@ -41,7 +42,7 @@ async def build( if request.ref_audio: wav, sr, _ = await server._resolve_ref_audio(request.ref_audio) prompt["ref_audio"] = (np.asarray(wav, dtype=np.float32), sr) - if request.ref_text: + if request.ref_text and request.ref_text.strip(): prompt["ref_text"] = request.ref_text if request.voice: voice_lower = request.voice.lower()