From 9b0df3909ae06207337cff817ce47e09cbbfe697 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E7=8E=8B=E9=87=91=E6=97=AD?= <105263726+wjinxu@users.noreply.github.com> Date: Wed, 12 Aug 2026 13:33:34 +0800 Subject: [PATCH] [Feature] timestamps for TTS via a forced-aligner pooling (#4795) Signed-off-by: wjinxu <1299461899@qq.com> Co-authored-by: xurui Co-authored-by: Claude Opus 4.8 (1M context) --- .../text_to_speech/README.md | 16 +- .../qwen3_tts/word_timestamps.py | 172 -------- .../online_serving/text_to_speech/README.md | 7 + tests/config/test_forced_aligner_injection.py | 71 ++++ .../sched/test_omni_ar_scheduler_logprobs.py | 71 ++++ .../openai_api/test_serving_speech_stream.py | 60 +-- .../test_forced_aligner_stage.py | 117 ++++++ tests/utils/test_forced_aligner.py | 26 +- .../utils/test_qwen3_force_align_processor.py | 17 - vllm_omni/config/config_factory.py | 4 + vllm_omni/config/model.py | 3 + vllm_omni/config/stage_config.py | 9 +- vllm_omni/core/sched/omni_ar_scheduler.py | 16 +- .../core/sched/omni_generation_scheduler.py | 22 +- vllm_omni/core/sched/omni_scheduler_mixin.py | 25 ++ .../deploy/qwen3_tts_forced_aligner.yaml | 14 +- vllm_omni/engine/arg_utils.py | 8 + vllm_omni/engine/orchestrator.py | 10 +- vllm_omni/engine/stage_init_utils.py | 20 +- vllm_omni/entrypoints/cli/serve.py | 12 +- vllm_omni/entrypoints/openai/api_server.py | 6 +- .../entrypoints/openai/protocol/audio.py | 10 +- .../entrypoints/openai/serving_speech.py | 88 +++- .../openai/serving_speech_stream.py | 83 ++-- .../stage_input_processors/forced_aligner.py | 173 ++++++++ vllm_omni/utils/forced_aligner.py | 375 ++++++------------ .../utils/qwen3_force_align_processor.py | 33 +- vllm_omni/worker/gpu_ar_model_runner.py | 9 +- 28 files changed, 850 insertions(+), 627 deletions(-) delete mode 100644 examples/offline_inference/text_to_speech/qwen3_tts/word_timestamps.py create mode 100644 tests/config/test_forced_aligner_injection.py create mode 100644 tests/model_executor/stage_input_processors/test_forced_aligner_stage.py create mode 100644 vllm_omni/model_executor/stage_input_processors/forced_aligner.py diff --git a/examples/offline_inference/text_to_speech/README.md b/examples/offline_inference/text_to_speech/README.md index 56135b55e..d246f9e45 100644 --- a/examples/offline_inference/text_to_speech/README.md +++ b/examples/offline_inference/text_to_speech/README.md @@ -372,18 +372,10 @@ python examples/offline_inference/text_to_speech/qwen3_tts/end2end.py \ Streaming requires `async_chunk: true` in the stage config. ### Word Timestamps -Generate a WAV offline and a JSON sidecar with word-level timestamps from -`Qwen/Qwen3-ForcedAligner-0.6B`: -```bash -python examples/offline_inference/text_to_speech/qwen3_tts/word_timestamps.py \ - --model Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \ - --forced-aligner Qwen/Qwen3-ForcedAligner-0.6B \ - --text "Hello world." \ - --output-dir /tmp/qwen3_tts_timestamps -``` -The script writes `qwen3_tts_word_timestamps.wav` and -`qwen3_tts_word_timestamps.json`. On machines without a local CUDA toolkit, -set `VLLM_USE_FLASHINFER_SAMPLER=0` to avoid FlashInfer sampler JIT. +Word-level timestamps are currently a serving-path feature: launch +`vllm-omni serve` with `--forced-aligner` and request `word_timestamps` +(see `examples/online_serving/text_to_speech/README.md`). An offline +example is not provided in this release. ### Batched decoding The Code2Wav stage supports batched decoding through the SpeechTokenizer. Pass multiple prompts via `--txt-prompts` and set `--batch-size` accordingly. To raise `max_num_seqs` on either stage, point `--stage-configs-path` at a stage configs YAML with the desired values (see `vllm_omni/model_executor/stage_configs/` for templates): diff --git a/examples/offline_inference/text_to_speech/qwen3_tts/word_timestamps.py b/examples/offline_inference/text_to_speech/qwen3_tts/word_timestamps.py deleted file mode 100644 index 581be2a7d..000000000 --- a/examples/offline_inference/text_to_speech/qwen3_tts/word_timestamps.py +++ /dev/null @@ -1,172 +0,0 @@ -"""Offline Qwen3-TTS example with word-level timestamps. - -This demo runs Qwen3-TTS offline, aligns the synthesized audio in memory with -the same shared forced aligner used by the streaming server path, then writes -the WAV and JSON sidecar: - - python examples/offline_inference/text_to_speech/qwen3_tts/word_timestamps.py \ - --model Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \ - --forced-aligner Qwen/Qwen3-ForcedAligner-0.6B - -On machines without a local CUDA toolkit, set ``VLLM_USE_FLASHINFER_SAMPLER=0`` -to avoid FlashInfer sampler JIT during warmup. -""" - -from __future__ import annotations - -import asyncio -import json -import os -from pathlib import Path -from typing import Any - -os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn") -# The in-process offline aligner consumes token_classify outputs. Keeping V1 -# multiprocessing off avoids msgspec serialization issues for token-wise logits. -os.environ.setdefault("VLLM_ENABLE_V1_MULTIPROCESSING", "0") - -import numpy as np -import soundfile as sf -import torch -from end2end import _estimate_prompt_len -from vllm.utils.argparse_utils import FlexibleArgumentParser - -from vllm_omni import Omni -from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults -from vllm_omni.utils.forced_aligner import align, build_forced_aligner_config - - -def _default_stage_config() -> str: - repo_root = Path(__file__).resolve().parents[4] - return str(repo_root / "vllm_omni" / "deploy" / "qwen3_tts.yaml") - - -def _build_custom_voice_input(args: Any) -> dict[str, Any]: - additional_information = { - "task_type": ["CustomVoice"], - "text": [args.text], - "language": [args.language], - "speaker": [args.speaker], - "instruct": [args.instructions], - "max_new_tokens": [args.max_new_tokens], - } - return { - "prompt_token_ids": [0] * _estimate_prompt_len(additional_information, args.model), - "additional_information": additional_information, - } - - -def _audio_tensor_and_sample_rate(mm: dict[str, Any]) -> tuple[torch.Tensor, int]: - audio_data = mm["audio"] - sr_raw = mm["sr"] - sr_val = sr_raw[-1] if isinstance(sr_raw, list) and sr_raw else sr_raw - sample_rate = int(sr_val.item()) if hasattr(sr_val, "item") else int(sr_val) - audio_tensor = torch.cat(audio_data, dim=-1) if isinstance(audio_data, list) else audio_data - return audio_tensor.float().detach().cpu().flatten(), sample_rate - - -def _float_audio_to_pcm16_bytes(audio: torch.Tensor) -> bytes: - samples = audio.numpy() - pcm = (np.clip(samples, -1.0, 1.0) * 32767.0).astype(" list[dict[str, Any]] | None: - aligner_config = build_forced_aligner_config(args) - if aligner_config is None: - raise ValueError("--forced-aligner or --forced-aligner-config is required") - timestamps = await align( - audio=_float_audio_to_pcm16_bytes(audio), - text=args.text, - sample_rate=sample_rate, - config=aligner_config, - language=args.language, - ) - if timestamps is None: - return None - return [ - { - "word": item.word, - "start_ms": item.start_ms, - "end_ms": item.end_ms, - } - for item in timestamps - ] - - -def main(args: Any) -> None: - output_dir = Path(args.output_dir) - output_dir.mkdir(parents=True, exist_ok=True) - - omni_kwargs = vars(args).copy() - for key in ( - "forced_aligner", - "text", - "language", - "speaker", - "instructions", - "max_new_tokens", - "output_dir", - ): - omni_kwargs.pop(key, None) - omni_kwargs["stage_configs_path"] = args.stage_configs_path or _default_stage_config() - omni_kwargs["log_stats"] = args.log_stats - omni = Omni(**omni_kwargs) - - prompt = _build_custom_voice_input(args) - final_output = None - for stage_outputs in omni.generate([prompt]): - final_output = stage_outputs.request_output - if final_output is None: - raise RuntimeError("Qwen3-TTS did not produce an output.") - - mm = final_output.outputs[0].multimodal_output - audio, sample_rate = _audio_tensor_and_sample_rate(mm) - - wav_path = output_dir / "qwen3_tts_word_timestamps.wav" - # Align the in-memory tensor, not the saved WAV. The WAV is only a demo - # artifact so users can listen to the same audio referenced by the sidecar. - timestamps = asyncio.run(_run_alignment(args, audio, sample_rate)) - sf.write(wav_path, audio.numpy(), samplerate=sample_rate, format="WAV") - - sidecar = { - "text": args.text, - "sample_rate": sample_rate, - "audio_path": str(wav_path), - "timestamps": timestamps, - } - json_path = output_dir / "qwen3_tts_word_timestamps.json" - json_path.write_text(json.dumps(sidecar, ensure_ascii=False, indent=2) + "\n", encoding="utf-8") - - print(f"Saved audio: {wav_path}") - print(f"Saved timestamps: {json_path}") - print(json.dumps(sidecar["timestamps"], ensure_ascii=False, indent=2)) - - -def parse_args() -> Any: - parser = FlexibleArgumentParser(description="Offline Qwen3-TTS word timestamp example") - parser.add_argument("--model", default="Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice", help="Qwen3-TTS model path/name") - parser.add_argument( - "--forced-aligner", - default=None, - help="Qwen3 forced aligner model path/name", - ) - parser.add_argument( - "--forced-aligner-config", - default=None, - help="Optional YAML file for forced aligner settings (incl. gpu_memory_utilization)", - ) - parser.add_argument("--stage-configs-path", default=None, help="Qwen3-TTS deploy YAML") - parser.add_argument("--text", default="Hello world.", help="Text to synthesize and align") - parser.add_argument("--language", default="English", help="Qwen3-TTS language field") - parser.add_argument("--speaker", default="Vivian", help="CustomVoice speaker name") - parser.add_argument("--instructions", default="", help="Optional speaking style instruction") - parser.add_argument("--max-new-tokens", type=int, default=2048, help="TTS max_new_tokens") - parser.add_argument("--output-dir", default="output_audio", help="Directory for WAV and JSON sidecar") - parser.add_argument("--log-stats", action="store_true", default=False, help="Enable vLLM-Omni stats logging") - nullify_stage_engine_defaults(parser) - return parser.parse_args() - - -if __name__ == "__main__": - main(parse_args()) diff --git a/examples/online_serving/text_to_speech/README.md b/examples/online_serving/text_to_speech/README.md index 2f9d84bf3..1ea191a37 100644 --- a/examples/online_serving/text_to_speech/README.md +++ b/examples/online_serving/text_to_speech/README.md @@ -605,6 +605,13 @@ python qwen3_tts/streaming_speech_client.py \ The client writes one PCM file per sentence and a matching `sentence_XXX_timestamps.json` sidecar. +Non-streaming requests can also ask for timestamps: pass +`"word_timestamps": true` to `POST /v1/audio/speech` and read the +`X-Word-Timestamps` response header (JSON list of `{word, start_ms, end_ms}`, +ASCII-escaped). Past 4 KB the header is replaced by +`X-Word-Timestamps-Omitted: oversize; bytes=; limit=4096` and the audio +still returns — use the WebSocket path for long transcripts. + To *see* the alignment instead of reading a JSON sidecar, run the word-timestamp Gradio demo (server must be launched with `--forced-aligner`): ```bash diff --git a/tests/config/test_forced_aligner_injection.py b/tests/config/test_forced_aligner_injection.py new file mode 100644 index 000000000..73a7bcda2 --- /dev/null +++ b/tests/config/test_forced_aligner_injection.py @@ -0,0 +1,71 @@ +import pytest + +import vllm_omni.config.pipeline_registry # noqa: F401 (populate registry) +from vllm_omni.config.pipeline_registry import OMNI_PIPELINES as _PIPELINE_REGISTRY +from vllm_omni.config.stage_config import ( + DeployConfig, + StageExecutionType, +) +from vllm_omni.model_executor.stage_input_processors.forced_aligner import ( + POOLING_OUTPUT_DECODER_PATH, + TIMESTAMPS_MODALITY, +) +from vllm_omni.utils.forced_aligner import _DEFAULT_CONFIG_PATH, inject_forced_aligner_stage + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_inject_appends_pooling_aligner_stage(): + pipeline = _PIPELINE_REGISTRY["qwen3_tts"] + deploy = DeployConfig() + n0 = len(pipeline.stages) + + ext_pipeline, ext_deploy = inject_forced_aligner_stage( + pipeline, deploy, {"forced_aligner": "/models/Qwen3-ForcedAligner-0.6B"} + ) + + # A new pooling stage is appended; the original pipeline is unchanged. + assert len(pipeline.stages) == n0 + assert len(ext_pipeline.stages) == n0 + 1 + # The original deploy config is not mutated either (double-injection safe). + assert len(deploy.stages) == 0 + assert len(ext_deploy.stages) == 1 + aligner = ext_pipeline.stages[-1] + # Pooling stage reuses LLM_AR (+ runner=pooling), no dedicated LLM_POOLING. + assert aligner.execution_type == StageExecutionType.LLM_AR + assert aligner.final_output is True + assert aligner.final_output_type == TIMESTAMPS_MODALITY + assert aligner.input_sources == (n0 - 1,) + assert aligner.requires_multimodal_data is True + assert aligner.owns_tokenizer is True + + # The deploy stage carries the separate aligner checkpoint + decoder hook. + ds = ext_deploy.stages[-1] + assert ds.engine_extras["model"] == "/models/Qwen3-ForcedAligner-0.6B" + assert ds.engine_extras["runner"] == "pooling" + assert ds.engine_extras["pooling_output_decoder"] == POOLING_OUTPUT_DECODER_PATH + + +def test_inject_from_config_only(): + # --forced-aligner-config alone (model in the YAML) must also inject. + pipeline = _PIPELINE_REGISTRY["qwen3_tts"] + deploy = DeployConfig() + n0 = len(pipeline.stages) + + ext_pipeline, ext_deploy = inject_forced_aligner_stage( + pipeline, deploy, {"forced_aligner": None, "forced_aligner_config": str(_DEFAULT_CONFIG_PATH)} + ) + + assert len(ext_pipeline.stages) == n0 + 1 + assert ext_pipeline.stages[-1].final_output_type == TIMESTAMPS_MODALITY + assert ext_deploy.stages[-1].engine_extras["model"] == "Qwen/Qwen3-ForcedAligner-0.6B" + + +def test_inject_noop_without_forced_aligner(): + pipeline = _PIPELINE_REGISTRY["qwen3_tts"] + deploy = DeployConfig() + + ext_pipeline, ext_deploy = inject_forced_aligner_stage(pipeline, deploy, {}) + + assert ext_pipeline is pipeline + assert len(ext_deploy.stages) == 0 diff --git a/tests/core/sched/test_omni_ar_scheduler_logprobs.py b/tests/core/sched/test_omni_ar_scheduler_logprobs.py index 2205ab32e..604904dd0 100644 --- a/tests/core/sched/test_omni_ar_scheduler_logprobs.py +++ b/tests/core/sched/test_omni_ar_scheduler_logprobs.py @@ -27,6 +27,7 @@ _MIXIN_UPDATE_HELPERS = ( "_attach_finished_request_sets", "_attach_scheduler_stats", "_capture_omni_connector_output", + "_maybe_decode_pooling_output", ) @@ -138,6 +139,7 @@ def _make_scheduler_stub(requests: list[_Request]) -> SimpleNamespace: pending_stop_after_extraction=set(), waiting_for_transfer_free=set(), _new_prompt_len_snapshot={}, + _pooling_output_decoder=None, finished_req_ids=set(), finished_req_ids_dict=defaultdict(set), kv_cache_manager=SimpleNamespace(take_events=lambda: None), @@ -267,3 +269,72 @@ def test_invalid_logprobs_finish_only_the_affected_scheduler_request() -> None: assert output_by_id["bad"].new_token_ids == [] assert output_by_id["good"].new_token_ids == [8] np.testing.assert_array_equal(output_by_id["good"].new_logprobs.logprob_token_ids[:, 0], [8]) + + +def _pooling_model_runner_output(pooler_tensor) -> SimpleNamespace: + return SimpleNamespace( + sampled_token_ids=[[]], + logprobs=None, + prompt_logprobs_dict={}, + pooler_output=[pooler_tensor], + num_nans_in_logits=None, + kv_connector_output=None, + cudagraph_stats=None, + req_id_to_index={"req": 0}, + routed_experts=None, + ) + + +def _make_pooling_stub(decoder) -> tuple[_Request, SimpleNamespace]: + request = _Request("req") + request.sampling_params = SimpleNamespace(num_logprobs=None) + request.pooling_params = SimpleNamespace() + scheduler = _make_scheduler_stub([request]) + scheduler._pooling_output_decoder = decoder + scheduler.vllm_config = SimpleNamespace(model_config=SimpleNamespace(hf_config=None)) + _bind_request_lifecycle(scheduler, update_request=lambda req, tokens: (tokens, False)) + return request, scheduler + + +def test_pooling_decode_success_emits_decoded_payload() -> None: + import torch + + decoded = torch.tensor([[1, 2]], dtype=torch.int32) + request, scheduler = _make_pooling_stub(lambda logits, req, hf: decoded) + + scheduler_output = SimpleNamespace( + num_scheduled_tokens={"req": 1}, + scheduled_spec_decode_tokens={}, + num_invalid_spec_tokens=0, + ) + outputs = OmniARScheduler.update_from_output( + scheduler, scheduler_output, _pooling_model_runner_output(torch.zeros(2, 4)) + ) + (output,) = outputs[0].outputs + + assert output.finish_reason is FinishReason.STOP + assert output.pooling_output is decoded + + +def test_pooling_decode_failure_finishes_request_with_error() -> None: + import torch + + def raising_decoder(logits, req, hf): + raise ValueError("boom") + + request, scheduler = _make_pooling_stub(raising_decoder) + + scheduler_output = SimpleNamespace( + num_scheduled_tokens={"req": 1}, + scheduled_spec_decode_tokens={}, + num_invalid_spec_tokens=0, + ) + outputs = OmniARScheduler.update_from_output( + scheduler, scheduler_output, _pooling_model_runner_output(torch.zeros(2, 4)) + ) + (output,) = outputs[0].outputs + + # Decoder failure is a request error (500), not an empty success. + assert output.finish_reason is FinishReason.ERROR + assert output.pooling_output is None + assert "pooling output decode failed" in str(request.stop_reason) diff --git a/tests/entrypoints/openai_api/test_serving_speech_stream.py b/tests/entrypoints/openai_api/test_serving_speech_stream.py index 44e84b346..37e3509f5 100644 --- a/tests/entrypoints/openai_api/test_serving_speech_stream.py +++ b/tests/entrypoints/openai_api/test_serving_speech_stream.py @@ -3,6 +3,7 @@ import base64 from types import SimpleNamespace import pytest +import torch from fastapi import FastAPI, WebSocket from pytest_mock import MockerFixture from starlette.testclient import TestClient @@ -11,11 +12,24 @@ from starlette.websockets import WebSocketDisconnect from vllm_omni.entrypoints.openai import serving_speech_stream as streaming_speech_module from vllm_omni.entrypoints.openai.serving_speech import OmniOpenAIServingSpeech from vllm_omni.entrypoints.openai.serving_speech_stream import OmniStreamingSpeechHandler -from vllm_omni.utils.forced_aligner import WordTimestamp +from vllm_omni.model_executor.stage_input_processors.forced_aligner import ALIGNER_WORDS_KEY pytestmark = [pytest.mark.core_model, pytest.mark.cpu] +def _fake_aligner_res(pairs, words): + """Build a stand-in for the forced-aligner stage's pooling output. + + Mirrors what rides the generator in production: ``res.outputs.data`` is a + ``[n_words, 2]`` int32 tensor of ``[start_ms, end_ms]`` and the word strings + travel in ``additional_information``. Decoded by ``extract_word_timestamps``. + """ + return SimpleNamespace( + outputs=SimpleNamespace(data=torch.tensor(pairs, dtype=torch.int32), additional_information=None), + additional_information={ALIGNER_WORDS_KEY: list(words)}, + ) + + def _build_test_app( speech_service=None, *, @@ -28,9 +42,9 @@ def _build_test_app( speech_service = mocker.MagicMock(spec=OmniOpenAIServingSpeech) speech_service._generate_audio_bytes = mocker.AsyncMock(return_value=(b"RIFF" + b"\x00" * 32, "audio/wav")) speech_service._prepare_speech_generation = mocker.AsyncMock(return_value=("req-1", object(), {})) - speech_service.forced_aligner_config = None + speech_service.forced_aligner_enabled = False - async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False): + async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False, collect=None): for chunk in (b"\x01\x02", b"\x03\x04\x05"): yield (chunk, 24000) if include_sample_rate else chunk @@ -235,7 +249,7 @@ class TestStreamingSpeechWebSocket: speech_service._generate_audio_bytes = mocker.AsyncMock(return_value=(b"", "audio/wav")) speech_service.engine_client = mocker.MagicMock() speech_service.engine_client.abort = mocker.AsyncMock() - speech_service.forced_aligner_config = None + speech_service.forced_aligner_enabled = False async def mock_prepare_speech_generation(request): captured_requests.append(request) @@ -311,13 +325,13 @@ class TestStreamingSpeechWebSocket: assert error["type"] == "error" assert "without --forced-aligner" in error["message"] - def test_word_timestamps_emit_sidecar_json_frame(self, mocker: MockerFixture): + def test_word_timestamps_emit_pipeline_json_frame(self, mocker: MockerFixture): captured_requests = [] speech_service = mocker.MagicMock(spec=OmniOpenAIServingSpeech) speech_service._generate_audio_bytes = mocker.AsyncMock(return_value=(b"", "audio/wav")) speech_service.engine_client = mocker.MagicMock() speech_service.engine_client.abort = mocker.AsyncMock() - speech_service.forced_aligner_config = SimpleNamespace(model="aligner") + speech_service.forced_aligner_enabled = True async def mock_prepare_speech_generation(request): captured_requests.append(request) @@ -328,19 +342,15 @@ class TestStreamingSpeechWebSocket: first_chunk = b"\x01" * 1000 second_chunk = b"\x02" * 1000 - async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False): + # The forced-aligner stage rides the same generator: its pooling output + # is surfaced via the ``collect`` channel once the audio has streamed. + async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False, collect=None): for chunk in (first_chunk, second_chunk): yield (chunk, 1000) if include_sample_rate else chunk + if collect is not None: + collect["aligner_res"] = _fake_aligner_res([[0, 200], [200, 900]], ["Hello", "world"]) speech_service._generate_pcm_chunks = mock_generate_pcm_chunks - # Sentence-level: aligner runs once over the whole sentence audio. - mock_align = mocker.AsyncMock( - return_value=[ - WordTimestamp("Hello", 0, 200), - WordTimestamp("world", 200, 900), - ] - ) - mocker.patch.object(streaming_speech_module, "forced_align", mock_align) app, _ = _build_test_app(speech_service) with TestClient(app) as client: @@ -406,7 +416,6 @@ class TestStreamingSpeechWebSocket: } assert captured_requests[0].word_timestamps is True - assert mock_align.await_count == 1 def test_word_timestamps_emit_word_dicts(self, mocker: MockerFixture): # The streaming layer forwards the aligner's (already monotonic, @@ -415,21 +424,16 @@ class TestStreamingSpeechWebSocket: speech_service._generate_audio_bytes = mocker.AsyncMock(return_value=(b"", "audio/wav")) speech_service.engine_client = mocker.MagicMock() speech_service.engine_client.abort = mocker.AsyncMock() - speech_service.forced_aligner_config = SimpleNamespace(model="aligner") + speech_service.forced_aligner_enabled = True speech_service._prepare_speech_generation = mocker.AsyncMock(return_value=("req", object(), {})) - async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False): + async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False, collect=None): chunk = b"\x01" * 1000 yield (chunk, 1000) if include_sample_rate else chunk + if collect is not None: + collect["aligner_res"] = _fake_aligner_res([[0, 1000], [1000, 1200]], ["Hello", "world"]) speech_service._generate_pcm_chunks = mock_generate_pcm_chunks - mock_align = mocker.AsyncMock( - return_value=[ - WordTimestamp("Hello", 0, 1000), - WordTimestamp("world", 1000, 1200), - ] - ) - mocker.patch.object(streaming_speech_module, "forced_align", mock_align) app, _ = _build_test_app(speech_service) with TestClient(app) as client: @@ -572,7 +576,7 @@ class TestStreamingSpeechWebSocket: speech_service._generate_pcm_chunks = mocker.AsyncMock() speech_service.engine_client = mocker.MagicMock() speech_service.engine_client.abort = mocker.AsyncMock() - speech_service.forced_aligner_config = None + speech_service.forced_aligner_enabled = False app, _ = _build_test_app(speech_service) with TestClient(app) as client: @@ -602,7 +606,7 @@ class TestStreamingSpeechWebSocket: speech_service._prepare_speech_generation = mocker.AsyncMock(return_value=("req-stream-err", object(), {})) speech_service.engine_client = mocker.MagicMock() speech_service.engine_client.abort = mocker.AsyncMock() - speech_service.forced_aligner_config = None + speech_service.forced_aligner_enabled = False async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False): yield b"\x01\x02" @@ -696,7 +700,7 @@ class TestStreamingSpeechWebSocket: speech_service._prepare_speech_generation = mocker.AsyncMock(return_value=("req-abort", object(), {})) speech_service.engine_client = mocker.MagicMock() speech_service.engine_client.abort = mocker.AsyncMock() - speech_service.forced_aligner_config = None + speech_service.forced_aligner_enabled = False async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False): yield b"\x01\x02" diff --git a/tests/model_executor/stage_input_processors/test_forced_aligner_stage.py b/tests/model_executor/stage_input_processors/test_forced_aligner_stage.py new file mode 100644 index 000000000..9b1d5ec2a --- /dev/null +++ b/tests/model_executor/stage_input_processors/test_forced_aligner_stage.py @@ -0,0 +1,117 @@ +from types import SimpleNamespace + +import numpy as np +import pytest +import torch + +from vllm_omni.model_executor.stage_input_processors.forced_aligner import ( + ALIGNER_WORDS_KEY, + _extract_text, + _extract_waveform, + code2wav2aligner, + decode_pooling_output, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class _FakeCompletion: + def __init__(self, multimodal_output): + self.multimodal_output = multimodal_output + + +class _FakeOutput: + def __init__(self, multimodal_output, finished=True): + self.finished = finished + self.outputs = [_FakeCompletion(multimodal_output)] + + +def _audio_mm(n=2400, sr=24000): + return {"audio": torch.zeros(n, dtype=torch.float32), "sr": torch.tensor(sr, dtype=torch.int32)} + + +def test_extract_waveform_from_audio_key(): + wav, sr = _extract_waveform(_audio_mm(n=100, sr=16000)) + assert isinstance(wav, np.ndarray) and wav.dtype == np.float32 + assert wav.shape == (100,) + assert sr == 16000 + + +def test_extract_waveform_falls_back_to_model_outputs_and_default_sr(): + wav, sr = _extract_waveform({"model_outputs": [torch.ones(50)]}) + assert wav.shape == (50,) + assert sr == 24000 # default when sr missing + + +def test_extract_text_unwraps_additional_information_list(): + prompt = {"additional_information": {"text": ["Hello world."]}} + assert _extract_text(prompt) == "Hello world." + + +def test_code2wav2aligner_builds_text_prompt_and_audio(): + prompt = {"additional_information": {"text": ["U.S.A test"], "language": ["English"]}} + out = code2wav2aligner([_FakeOutput(_audio_mm())], prompt) + + assert len(out) == 1 + item = out[0] + # text prompt: tokenization + audio features deferred to the stage preprocessor + assert "prompt_token_ids" not in item + # prompt built from official segmentation: "U.S.A" -> "USA" + assert "USAtest" in item["prompt"] + # audio carried as (np.float32, sr) tuple + audio, sr = item["multi_modal_data"]["audio"] + assert isinstance(audio, np.ndarray) and sr == 24000 + assert item["additional_information"]["aligner_words"] == [["USA", "test"]] + assert item["additional_information"]["language"] == ["English"] + + +def test_code2wav2aligner_skips_unfinished_and_empty(): + # unfinished -> skipped + assert code2wav2aligner([_FakeOutput(_audio_mm(), finished=False)], {}) == [] + # finished but empty waveform -> skipped + empty = _FakeOutput({"audio": torch.zeros(0), "sr": torch.tensor(24000)}) + assert code2wav2aligner([empty], {"additional_information": {"text": ["hi"]}}) == [] + + +def test_decode_pooling_output_produces_word_timestamps_ms(): + # 2 words -> 4 markers (id=99). classify_num=5, seg=200ms. + hf_config = SimpleNamespace( + timestamp_token_id=99, + timestamp_segment_time=200.0, + thinker_config=SimpleNamespace(classify_num=5), + ) + # prompt: marker positions at indices 2,3,5,6 + prompt_token_ids = [1, 1, 99, 99, 1, 99, 99] + request = SimpleNamespace( + prompt_token_ids=prompt_token_ids, + additional_information={ALIGNER_WORDS_KEY: [["hello", "world"]]}, + ) + logits = torch.zeros((7, 5), dtype=torch.float32) + logits[2, 0] = 1.0 # word0 start -> bin 0 + logits[3, 2] = 1.0 # word0 end -> bin 2 + logits[5, 2] = 1.0 # word1 start -> bin 2 + logits[6, 4] = 1.0 # word1 end -> bin 4 + + ms = decode_pooling_output(logits, request, hf_config) + + # decoder returns a bare [n_words, 2] int32 tensor (not a dict) + assert ms.tolist() == [[0, 400], [400, 800]] # bins x 200ms + + +def test_decode_pooling_output_handles_hf_config_without_thinker(): + # hf_config without thinker_config falls back to itself for classify_num. + hf_config = SimpleNamespace( + timestamp_token_id=99, + timestamp_segment_time=100.0, + classify_num=4, + ) + request = SimpleNamespace( + prompt_token_ids=[99, 99], + additional_information={ALIGNER_WORDS_KEY: [["hi"]]}, + ) + logits = torch.zeros((2, 4), dtype=torch.float32) + logits[0, 1] = 1.0 + logits[1, 3] = 1.0 + + ms = decode_pooling_output(logits, request, hf_config) + assert ms.tolist() == [[100, 300]] diff --git a/tests/utils/test_forced_aligner.py b/tests/utils/test_forced_aligner.py index 4964f3900..0e15f5b64 100644 --- a/tests/utils/test_forced_aligner.py +++ b/tests/utils/test_forced_aligner.py @@ -105,9 +105,7 @@ forced_aligner: """, encoding="utf-8", ) - args = type("Args", (), {"forced_aligner": None, "forced_aligner_config": str(cfg)})() - - out = forced_aligner.build_forced_aligner_config(args) + out = forced_aligner.build_forced_aligner_config(None, str(cfg)) assert out == forced_aligner.ForcedAlignerConfig( model="Qwen/Qwen3-ForcedAligner-0.6B", @@ -127,16 +125,7 @@ def test_build_config_cli_model_overrides_yaml(tmp_path): "forced_aligner:\n model: old\n gpu_memory_utilization: 0.2\n dtype: float16\n", encoding="utf-8", ) - args = type( - "Args", - (), - { - "forced_aligner": "new", - "forced_aligner_config": str(cfg), - }, - )() - - out = forced_aligner.build_forced_aligner_config(args) + out = forced_aligner.build_forced_aligner_config("new", str(cfg)) assert out is not None # --forced-aligner overrides the YAML model; gpu_memory_utilization/dtype @@ -148,16 +137,7 @@ def test_build_config_cli_model_overrides_yaml(tmp_path): def test_build_config_from_cli_model_uses_default_yaml(): - args = type( - "Args", - (), - { - "forced_aligner": "local-aligner", - "forced_aligner_config": None, - }, - )() - - out = forced_aligner.build_forced_aligner_config(args) + out = forced_aligner.build_forced_aligner_config("local-aligner", None) assert out is not None assert out.model == "local-aligner" diff --git a/tests/utils/test_qwen3_force_align_processor.py b/tests/utils/test_qwen3_force_align_processor.py index c9af7167e..4a928d0e5 100644 --- a/tests/utils/test_qwen3_force_align_processor.py +++ b/tests/utils/test_qwen3_force_align_processor.py @@ -46,20 +46,3 @@ def test_fix_timestamp_repairs_dip_onto_monotonic_sequence(): def test_fix_timestamp_passes_through_monotonic(): assert processor.fix_timestamp([0, 100, 100, 250]) == [0, 100, 100, 250] - - -def test_resolve_timestamp_token_id_defaults_to_marker_token(): - # Regression: the default must resolve the same marker that - # build_prompt inserts, not None (which would degrade every request). - seen = {} - - class FakeTokenizer: - def convert_tokens_to_ids(self, token): - seen["token"] = token - return 151705 - - tid = processor.resolve_timestamp_token_id(FakeTokenizer()) - - assert tid == 151705 - assert seen["token"] == processor.TIMESTAMP_TOKEN - assert processor.TIMESTAMP_TOKEN in processor.build_prompt(["hello", "world"]) diff --git a/vllm_omni/config/config_factory.py b/vllm_omni/config/config_factory.py index ed9b05f1e..a1778eda3 100644 --- a/vllm_omni/config/config_factory.py +++ b/vllm_omni/config/config_factory.py @@ -409,6 +409,10 @@ class StageConfigFactory: if cli_async_chunk is not None: deploy_cfg.async_chunk = bool(cli_async_chunk) + from vllm_omni.utils.forced_aligner import inject_forced_aligner_stage + + pipeline_cfg, deploy_cfg = inject_forced_aligner_stage(pipeline_cfg, deploy_cfg, cli_overrides) + stages = merge_pipeline_deploy(pipeline_cfg, deploy_cfg, cli_overrides) # Overlay declarative parallel strategies (opt-in) before CLI overrides. diff --git a/vllm_omni/config/model.py b/vllm_omni/config/model.py index f830fc37e..9164e1352 100644 --- a/vllm_omni/config/model.py +++ b/vllm_omni/config/model.py @@ -129,6 +129,9 @@ class OmniModelConfig(ModelConfig): model_arch: str | None = None worker_type: str | None = None engine_output_type: str | None = None + # Optional dotted path of a per-stage pooling-output decoder applied + # worker-side before IPC. Read by the AR scheduler. + pooling_output_decoder: str | None = None hf_config_name: str | None = None custom_process_next_stage_input_func: str | None = None stage_connector_config: dict[str, Any] = field( diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index 55c4a46cc..fa8e45295 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -340,6 +340,7 @@ class StageDeployConfig: output_connectors: dict[str, str] | None = None input_connectors: dict[str, str] | None = None default_sampling_params: dict[str, Any] | None = None + default_pooling_params: dict[str, Any] | None = None subtalker_sampling_params: dict[str, Any] | None = None # === Generic stage engine fields === @@ -509,6 +510,7 @@ _STAGE_RESERVED_KEYS = frozenset( "output_connectors", "input_connectors", "default_sampling_params", + "default_pooling_params", "engine_extras", "engine_args", "runtime", @@ -562,11 +564,14 @@ def _parse_stage_deploy(stage_data: dict[str, Any]) -> StageDeployConfig: kwargs["output_connectors"] = stage_data.get("output_connectors") kwargs["input_connectors"] = stage_data.get("input_connectors") kwargs["default_sampling_params"] = stage_data.get("default_sampling_params") + kwargs["default_pooling_params"] = stage_data.get("default_pooling_params") kwargs["engine_extras"] = _get_recursively_merged_dict(explicit_engine_extras, flat_args) return StageDeployConfig(**kwargs) -_DEEP_MERGE_KEYS = frozenset({"default_sampling_params", "subtalker_sampling_params", "engine_extras", "engine_args"}) +_DEEP_MERGE_KEYS = frozenset( + {"default_sampling_params", "default_pooling_params", "subtalker_sampling_params", "engine_extras", "engine_args"} +) def _deep_merge_stage(base: dict, overlay: dict) -> dict: @@ -886,6 +891,8 @@ def _build_extras( sampling.update(ps.sampling_constraints) if sampling: extras["default_sampling_params"] = sampling + if ds is not None and ds.default_pooling_params: + extras["default_pooling_params"] = dict(ds.default_pooling_params) if ds is not None and ds.output_connectors: extras["output_connectors"] = dict(ds.output_connectors) if ds is not None and ds.input_connectors: diff --git a/vllm_omni/core/sched/omni_ar_scheduler.py b/vllm_omni/core/sched/omni_ar_scheduler.py index 7cf3b5a6d..66da231a6 100644 --- a/vllm_omni/core/sched/omni_ar_scheduler.py +++ b/vllm_omni/core/sched/omni_ar_scheduler.py @@ -415,6 +415,17 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): finish_reason = None routed_experts = None + # Decode the pooling output before stop handling so a decoder + # failure finishes the request with FinishReason.ERROR (500). + try: + pooling_output_payload = self._maybe_decode_pooling_output(request, pooler_output) + except Exception as exc: + logger.exception("[pooling] decoder hook failed for request %s", req_id) + pooling_output_payload = None + request.status = RequestStatus.FINISHED_ERROR + request.stop_reason = f"pooling output decode failed: {exc}" + request.resumable = False + # Check for stop and update request status. if new_token_ids: num_sampled_tokens = len(new_token_ids) @@ -428,7 +439,8 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): new_logprobs = logprobs.slice_request(req_index, len(new_token_ids)) elif request.pooling_params and pooler_output is not None: # Pooling stops as soon as there is output. - request.status = RequestStatus.FINISHED_STOPPED + if request.status != RequestStatus.FINISHED_ERROR: + request.status = RequestStatus.FINISHED_STOPPED stopped = True # If criteria returns True, it means we must STOP the request. @@ -522,7 +534,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): finish_reason=finish_reason, new_logprobs=new_logprobs, new_prompt_logprobs_tensors=prompt_logprobs_tensors, - pooling_output=pooler_output, + pooling_output=pooling_output_payload, multimodal_output=mm_output, stop_reason=request.stop_reason, prefill_stats=request.take_prefill_stats(), diff --git a/vllm_omni/core/sched/omni_generation_scheduler.py b/vllm_omni/core/sched/omni_generation_scheduler.py index cf3678251..ae19a41a3 100644 --- a/vllm_omni/core/sched/omni_generation_scheduler.py +++ b/vllm_omni/core/sched/omni_generation_scheduler.py @@ -5,6 +5,7 @@ from collections import defaultdict from typing import Any from vllm.compilation.cuda_graph import CUDAGraphStat +from vllm.logger import init_logger from vllm.v1.core.kv_cache_manager import KVCacheBlocks from vllm.v1.core.sched.interface import PauseState from vllm.v1.core.sched.output import SchedulerOutput @@ -25,6 +26,8 @@ from vllm_omni.core.sched.output import OmniCachedRequestData, OmniNewRequestDat from vllm_omni.core.sched.utils import omni_routed_experts_for_request from vllm_omni.outputs import OmniModelRunnerOutput +logger = init_logger(__name__) + class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): def __init__(self, *args, **kwargs): @@ -442,10 +445,22 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): is_segment_finished = False routed_experts = None + # Decode the pooling output before stop handling so a decoder + # failure finishes the request with FinishReason.ERROR (500). + try: + pooling_output_payload = self._maybe_decode_pooling_output(request, pooler_output) + except Exception as exc: + logger.exception("[pooling] decoder hook failed for request %s", req_id) + pooling_output_payload = None + request.status = RequestStatus.FINISHED_ERROR + request.stop_reason = f"pooling output decode failed: {exc}" + request.resumable = False + # One-shot generation request: finish after its current input unit # has been fully processed. if ( - request.status == RequestStatus.FINISHED_STOPPED + request.status == RequestStatus.FINISHED_ERROR + or request.status == RequestStatus.FINISHED_STOPPED or (self.chunk_transfer_adapter is None and request.num_computed_tokens >= request.num_prompt_tokens) or ( self.chunk_transfer_adapter is not None @@ -453,7 +468,8 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): and request.num_computed_tokens >= len(request.prompt_token_ids) ) ): - request.status = RequestStatus.FINISHED_STOPPED + if request.status != RequestStatus.FINISHED_ERROR: + request.status = RequestStatus.FINISHED_STOPPED # Optional: set a stop_reason for front-end clarity # (does not affect protocol) stopped = True @@ -508,7 +524,7 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): finish_reason=finish_reason, new_logprobs=new_logprobs, new_prompt_logprobs_tensors=prompt_logprobs_tensors, - pooling_output=pooler_output, + pooling_output=pooling_output_payload, multimodal_output=mm_output, stop_reason=request.stop_reason, prefill_stats=request.take_prefill_stats(), diff --git a/vllm_omni/core/sched/omni_scheduler_mixin.py b/vllm_omni/core/sched/omni_scheduler_mixin.py index d0a4cefab..e0df774c0 100644 --- a/vllm_omni/core/sched/omni_scheduler_mixin.py +++ b/vllm_omni/core/sched/omni_scheduler_mixin.py @@ -5,10 +5,12 @@ import time from collections.abc import Iterable from typing import Any +import torch from vllm.compilation.cuda_graph import CUDAGraphStat from vllm.distributed.kv_events import KVEventBatch from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats from vllm.logger import init_logger +from vllm.utils.import_utils import resolve_obj_by_qualname from vllm.v1.core.sched.output import SchedulerOutput from vllm.v1.core.sched.utils import remove_all from vllm.v1.engine import ( @@ -80,6 +82,29 @@ class OmniSchedulerMixin: else None ) self._latest_omni_connector_output = None + # Optional per-stage pooling-output decoder hook (dotted path in + # model_config); applied worker-side before IPC. + self._pooling_output_decoder = None + _decoder_path = getattr(model_config, "pooling_output_decoder", None) + if _decoder_path: + self._pooling_output_decoder = resolve_obj_by_qualname(str(_decoder_path)) + + def _maybe_decode_pooling_output(self, request: Request, pooler_output: Any) -> Any: + """Apply the stage's pooling-output decoder hook to the pooler tensor + before IPC, or pass it through unchanged when none is configured. + Decoder exceptions propagate; callers fail the request with + FinishReason.ERROR rather than emitting an empty success.""" + if self._pooling_output_decoder is None: + return pooler_output + if pooler_output is None or getattr(request, "pooling_params", None) is None: + return pooler_output + if not isinstance(pooler_output, torch.Tensor): + return pooler_output + return self._pooling_output_decoder( + pooler_output, + request, + self.vllm_config.model_config.hf_config, + ) def _free_input_coordinator_request(self, request_id: str) -> None: """Prune full-payload coordinator state for a completed request.""" diff --git a/vllm_omni/deploy/qwen3_tts_forced_aligner.yaml b/vllm_omni/deploy/qwen3_tts_forced_aligner.yaml index 41145343d..2aecdbf9b 100644 --- a/vllm_omni/deploy/qwen3_tts_forced_aligner.yaml +++ b/vllm_omni/deploy/qwen3_tts_forced_aligner.yaml @@ -1,10 +1,12 @@ -# Reference deploy knobs for Qwen/Qwen3-ForcedAligner-0.6B, the shared forced -# aligner for streaming TTS word timestamps. Pass with --forced-aligner-config; -# --forced-aligner overrides the model path without editing this file. +# Default deploy knobs for Qwen/Qwen3-ForcedAligner-0.6B, the TTS word-timestamp +# aligner. These configure the pooling stage that --forced-aligner appends to +# the pipeline tail (see utils.forced_aligner.inject_forced_aligner_stage); they are +# NOT a standalone engine. Pass --forced-aligner-config to override this file; +# --forced-aligner sets just the model path. # -# Run as LLM(runner="pooling") with pooling_task="token_classify". Prompt tokens, -# chat template and word segmentation are Qwen-specific and live in -# qwen3_force_align_processor (not configurable here). +# The stage runs the model with runner="pooling" + pooling_task="token_classify". +# Prompt tokens, chat template and word segmentation are Qwen-specific and live +# in qwen3_force_align_processor (not configurable here). forced_aligner: model: Qwen/Qwen3-ForcedAligner-0.6B diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index a66f39ce3..1d9eb803b 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -191,6 +191,8 @@ class OmniEngineArgs(EngineArgs): quantization_config: Any | None = None force_cutlass_fp8: bool | None = None worker_type: str | None = None + # Dotted path of a per-stage pooling-output decoder applied worker-side. + pooling_output_decoder: str | None = None task_type: str | None = None worker_cls: str = None # type: ignore[assignment] # Upstream default is "auto"; omni resolves # in __post_init__ based on worker_type (ar/generation), so None is safe here. @@ -380,6 +382,7 @@ class OmniEngineArgs(EngineArgs): model_stage=self.model_stage, model_arch=self.model_arch, worker_type=self.worker_type, + pooling_output_decoder=self.pooling_output_decoder, engine_output_type=self.engine_output_type, hf_config_name=self.hf_config_name, custom_process_next_stage_input_func=self.custom_process_next_stage_input_func, @@ -475,6 +478,11 @@ class OrchestratorArgs: # === Mode Switches (orchestrator reads, DeployConfig redistributes) === async_chunk: bool | None = None + # === Forced aligner (orchestrator injects a pooling stage; never a per-stage knob) === + forced_aligner: str | None = None + forced_aligner_config: str | None = None + forced_aligner_device: str | None = None + # === Observability === log_stats: bool = False enable_orch_monitor: bool = False diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py index c539e9559..5646e82eb 100644 --- a/vllm_omni/engine/orchestrator.py +++ b/vllm_omni/engine/orchestrator.py @@ -1543,11 +1543,19 @@ class Orchestrator: return request processor = self._get_stage_input_processor(next_stage_id) + # A pooling stage is driven by PoolingParams; vLLM validates them against + # the stage's supported tasks, so advertise the model's pooling tasks (or + # the configured task) instead of generate. + if isinstance(params, PoolingParams): + model_cfg = next_pool.stage_vllm_config.model_config + supported_tasks = tuple(getattr(model_cfg, "supported_tasks", ()) or ()) or (params.task,) + else: + supported_tasks = ("generate",) request = processor.process_inputs( request_id=req_id, prompt=next_input, params=params, - supported_tasks=("generate",), + supported_tasks=supported_tasks, arrival_time=_time.time(), resumable=resumable, ) diff --git a/vllm_omni/engine/stage_init_utils.py b/vllm_omni/engine/stage_init_utils.py index b0a95d22c..c2e1ab5bc 100644 --- a/vllm_omni/engine/stage_init_utils.py +++ b/vllm_omni/engine/stage_init_utils.py @@ -20,6 +20,7 @@ from dataclasses import dataclass, fields, replace from typing import Any, Literal, cast from vllm.logger import init_logger +from vllm.pooling_params import PoolingParams from vllm.renderers import BaseRenderer from vllm.sampling_params import SamplingParams from vllm.tokenizers import cached_tokenizer_from_config @@ -395,8 +396,19 @@ def extract_legacy_stage_metadata(stage_config: Any) -> StageMetadata: final_output_type: str | None = stage_config.final_output_type default_sp = _to_dict(_get_attr_or_item(stage_config, "default_sampling_params", {})) - SPClass = SamplingParams if stage_type == "llm" else OmniDiffusionSamplingParams - default_sampling_params: OmniSamplingParams = SPClass(**default_sp) + # A pooling stage carries its task via default_pooling_params, set where the + # stage is declared. + default_pp = _to_dict(_get_attr_or_item(stage_config, "default_pooling_params", {})) + # A pooling stage is an LLM stage run with runner="pooling" (vLLM's + # is_pooling_model signal); pick params by that signal, not execution_type. + is_pooling = str(engine_args.get("runner", "")).lower() == "pooling" + default_params: OmniSamplingParams | PoolingParams + if stage_type == "diffusion": + default_params = OmniDiffusionSamplingParams(**default_sp) + elif is_pooling: + default_params = PoolingParams(**default_pp) + else: # generative llm: ar / generation + default_params = SamplingParams(**default_sp) custom_process_input_func: Callable | None = None _cpif_path = _get_attr_or_item(stage_config, "custom_process_input_func") @@ -428,7 +440,7 @@ def extract_legacy_stage_metadata(stage_config: Any) -> StageMetadata: engine_input_source=engine_input_source, final_output=final_output, final_output_type=final_output_type, - default_sampling_params=default_sampling_params, + default_sampling_params=default_params, custom_process_input_func=custom_process_input_func, model_stage=model_stage, runtime_cfg=runtime_cfg, @@ -448,7 +460,7 @@ def extract_legacy_stage_metadata(stage_config: Any) -> StageMetadata: engine_input_source=engine_input_source, final_output=final_output, final_output_type=final_output_type, - default_sampling_params=default_sampling_params, + default_sampling_params=default_params, custom_process_input_func=custom_process_input_func, model_stage=model_stage, runtime_cfg=runtime_cfg, diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index 815ce4d6e..dfa222cd1 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -237,9 +237,8 @@ class OmniServeCommand(CLISubcommand): "VoiceDesign, or Base, while diffusion models may use it to select " "task-specific weights. If omitted, the model default is used.", ) - # Forced aligner / word timestamps. --forced-aligner is the opt-in - # toggle; heavier knobs (gpu_memory_utilization, dtype, max_model_len) - # live in the deploy YAML passed via --forced-aligner-config. + # Forced aligner / word timestamps. Either flag opts in (--forced-aligner-config + # alone works when its YAML sets the model); heavier knobs live in that YAML. omni_config_group.add_argument( "--forced-aligner", type=str, @@ -260,6 +259,13 @@ class OmniServeCommand(CLISubcommand): "flag, when set, overrides the YAML model field." ), ) + omni_config_group.add_argument( + "--forced-aligner-device", + type=str, + default=None, + help="Device(s) for the forced-aligner stage (e.g. '2'). Defaults to " + "sharing an existing stage's GPU when unset.", + ) omni_config_group.add_argument( "--deploy-config", type=str, diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index 940f62f2e..0ea0f3a3f 100755 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/vllm_omni/entrypoints/openai/api_server.py @@ -1127,7 +1127,11 @@ async def omni_init_app_state( state.openai_serving_models, request_logger=request_logger, model_name=model_name, - forced_aligner_config=build_forced_aligner_config(args), + forced_aligner_enabled=build_forced_aligner_config( + getattr(args, "forced_aligner", None), + getattr(args, "forced_aligner_config", None), + ) + is not None, ) # Warm up speech pipeline (CUDA Graph capture, torch.compile) so the first diff --git a/vllm_omni/entrypoints/openai/protocol/audio.py b/vllm_omni/entrypoints/openai/protocol/audio.py index 1564c4523..b2b55bc52 100644 --- a/vllm_omni/entrypoints/openai/protocol/audio.py +++ b/vllm_omni/entrypoints/openai/protocol/audio.py @@ -161,10 +161,12 @@ class OpenAICreateSpeechRequest(BaseModel): word_timestamps: bool = Field( default=False, description=( - "When true, the server runs a shared forced aligner alongside the streamed " - "audio and emits per-chunk word timestamps. Requires the server to be " - "launched with --forced-aligner pointing at an aligner model. No effect " - "when streaming is off." + "When true, non-streaming responses carry per-word timestamps in the " + "X-Word-Timestamps header (JSON list of {word, start_ms, end_ms}, " + "ASCII-escaped; replaced by X-Word-Timestamps-Omitted past 4 KB). " + "Requires the server to be launched with --forced-aligner (400 otherwise). " + "Not supported with stream=true; for streaming use the WebSocket " + "/v1/audio/speech/stream path." ), ) diff --git a/vllm_omni/entrypoints/openai/serving_speech.py b/vllm_omni/entrypoints/openai/serving_speech.py index 6b7ca1318..88e7624f1 100644 --- a/vllm_omni/entrypoints/openai/serving_speech.py +++ b/vllm_omni/entrypoints/openai/serving_speech.py @@ -423,7 +423,9 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): def __init__(self, *args, **kwargs): self.model_name = kwargs.pop("model_name", None) - self.forced_aligner_config: Any | None = kwargs.pop("forced_aligner_config", None) + # True when the server was launched with --forced-aligner (a pooling + # aligner stage is appended to the pipeline). Gates word_timestamps. + self.forced_aligner_enabled: bool = bool(kwargs.pop("forced_aligner_enabled", False)) super().__init__(*args, **kwargs) self._init_speaker_storage() @@ -2186,6 +2188,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): request_start_s: float | None = None, include_sample_rate: bool = False, usage_acc: SpeechOutputTokenCounter | None = None, + collect: dict | None = None, ): """Generate audio chunks for streaming response. @@ -2218,6 +2221,9 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): usage_acc.observe(res) audio_output, audio_key = self._extract_audio_output(res) if audio_key is None: + # Stash the aligner's timestamps output for streaming callers. + if collect is not None and self._is_timestamps_output(res): + collect["aligner_res"] = res continue sr_raw = audio_output.get("sr") @@ -2406,6 +2412,13 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): data = json.dumps(payload, separators=(",", ":")) yield f"event: speech.audio.error\ndata: {data}\n\n" + @staticmethod + def _is_timestamps_output(res) -> bool: + """True when ``res`` is the forced-aligner stage's terminal timestamps output.""" + from vllm_omni.model_executor.stage_input_processors.forced_aligner import TIMESTAMPS_MODALITY + + return getattr(res, "final_output_type", None) == TIMESTAMPS_MODALITY + @staticmethod def _extract_audio_output(res) -> tuple[dict | None, str | None]: """Return (audio_output dict, audio key) or (None, None). @@ -3332,11 +3345,20 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): stage0_params.extra_args = {} stage0_params.extra_args.setdefault("tts_local_seed", int(default_seed)) + # When word_timestamps is requested, also ask for the aligner stage's + # output so the orchestrator drives the request through the forced-aligner + # stage (final_stage_id extends to it). Harmless if no aligner stage exists. + output_modalities = ["audio"] + if getattr(request, "word_timestamps", False): + from vllm_omni.model_executor.stage_input_processors.forced_aligner import TIMESTAMPS_MODALITY + + output_modalities.append(TIMESTAMPS_MODALITY) + generator = self.engine_client.generate( prompt=prompt, request_id=request_id, sampling_params_list=sampling_params_list, - output_modalities=["audio"], + output_modalities=output_modalities, ) self._track_ref_audio_artifact_warmup( request_id, @@ -3345,17 +3367,22 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): ) return request_id, generator, tts_params - async def _generate_pcm_chunks(self, generator, request_id: str, *, include_sample_rate: bool = False): + async def _generate_pcm_chunks( + self, generator, request_id: str, *, include_sample_rate: bool = False, collect: dict | None = None + ): """Yield raw PCM byte chunks from the engine generator. Delegates to ``_generate_audio_chunks`` with ``response_format="pcm"``. Used by the WebSocket streaming handler and ``_iter_pcm_audio_bytes``. + ``collect`` (when given) receives the forced-aligner stage's pooling + output under ``"aligner_res"`` for downstream word-timestamp extraction. """ async for chunk in self._generate_audio_chunks( generator, request_id, response_format="pcm", include_sample_rate=include_sample_rate, + collect=collect, ): yield chunk @@ -3375,6 +3402,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): request_id: str | None = None, usage_out: list[SpeechTokenUsage] | None = None, has_inline_ref_audio: bool | None = None, + collect: dict | None = None, ) -> tuple[bytes | str, str]: # ``usage_out`` is an opt-in output channel: when a list is passed, the # computed SpeechTokenUsage is appended to it. The return stays a @@ -3398,9 +3426,20 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): # Non-streaming is FINAL_ONLY, so the stage-0 output carries the full # token sequence; the counter records its length for output_tokens. usage_acc = SpeechOutputTokenCounter() + audio_res: OmniRequestOutput | None = None + aligner_res: OmniRequestOutput | None = None async for res in generator: final_output = res usage_acc.observe(res) + # The generator yields both the audio output (Code2Wav) and, with + # a forced-aligner stage, a timestamps output. Keep the audio res + # for the WAV and the aligner res for word timestamps. + if self._is_timestamps_output(res): + aligner_res = res + else: + _, audio_key = self._extract_audio_output(res) + if audio_key is not None: + audio_res = res if not is_moss: continue try: @@ -3422,10 +3461,26 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): if final_output is None: raise ValueError("No output generated from the model.") - audio_output, audio_key = self._extract_audio_output(final_output) + # Extract audio from the audio-bearing res (not necessarily the last + # yielded one, which may be the aligner's timestamps output). + audio_source = audio_res if audio_res is not None else final_output + audio_output, audio_key = self._extract_audio_output(audio_source) if audio_key is None: raise ValueError("TTS model did not produce audio output.") + # Surface forced-aligner word timestamps to the caller (set as a + # response header) when requested and an aligner stage produced them. + if collect is not None and getattr(request, "word_timestamps", False): + from vllm_omni.utils.forced_aligner import extract_word_timestamps + + ts = ( + extract_word_timestamps(aligner_res, request.input, getattr(request, "language", None)) + if aligner_res is not None + else None + ) + if ts is not None: + collect["word_timestamps"] = ts + audio_tensor = audio_output[audio_key] sr_raw = audio_output.get("sr", 24000) sr_val = sr_raw[-1] if isinstance(sr_raw, list) and sr_raw else sr_raw @@ -3695,6 +3750,11 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): "/v1/audio/speech/stream path. Use session.config with " "stream_audio=true and response_format='pcm'." ) + if request.word_timestamps and not self.forced_aligner_enabled: + # Fail loud instead of silently returning 200 with no timestamps header. + return self.create_error_response( + "word_timestamps=true requires the server to be launched with --forced-aligner." + ) if request.is_raw_audio_stream(): response_format, error = self._validate_speech_streaming_request( @@ -3739,7 +3799,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): media_type="text/event-stream", ) - audio_bytes, media_type = await self._generate_audio_bytes(request, request_id=request_id) + collect: dict = {} + audio_bytes, media_type = await self._generate_audio_bytes(request, request_id=request_id, collect=collect) total_ms = (time.perf_counter() - request_start_s) * 1000.0 logger.info( "[SpeechE2E] request_id=%s stream=false status=ok total_ms=%.2f response_bytes=%d", @@ -3747,7 +3808,22 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): total_ms, len(audio_bytes) if isinstance(audio_bytes, (bytes, bytearray)) else len(str(audio_bytes)), ) - return Response(content=audio_bytes, media_type=media_type) + headers = {} + if collect.get("word_timestamps") is not None: + # Default ensure_ascii keeps the header latin-1 encodable (non-ASCII words \uXXXX-escaped). + ts_json = json.dumps(collect["word_timestamps"]) + # Cap at 4 KB: oversized headers turn into opaque 502s at common reverse-proxy defaults. + if len(ts_json) <= 4096: + headers["X-Word-Timestamps"] = ts_json + else: + # Marker header so clients can tell an oversized alignment from no alignment. + headers["X-Word-Timestamps-Omitted"] = f"oversize; bytes={len(ts_json)}; limit=4096" + logger.warning( + "X-Word-Timestamps header omitted: %d bytes exceeds the 4 KB budget " + "(use the WebSocket streaming path for long transcripts)", + len(ts_json), + ) + return Response(content=audio_bytes, media_type=media_type, headers=headers) except asyncio.CancelledError: total_ms = (time.perf_counter() - request_start_s) * 1000.0 diff --git a/vllm_omni/entrypoints/openai/serving_speech_stream.py b/vllm_omni/entrypoints/openai/serving_speech_stream.py index 102cfda45..2b71ad6c7 100644 --- a/vllm_omni/entrypoints/openai/serving_speech_stream.py +++ b/vllm_omni/entrypoints/openai/serving_speech_stream.py @@ -61,8 +61,7 @@ from vllm_omni.entrypoints.openai.protocol.audio import ( StreamingSpeechSessionConfig, ) from vllm_omni.entrypoints.openai.serving_speech import OmniOpenAIServingSpeech -from vllm_omni.utils.forced_aligner import ForcedAlignerLoadError -from vllm_omni.utils.forced_aligner import align as forced_align +from vllm_omni.utils.forced_aligner import extract_word_timestamps logger = init_logger(__name__) @@ -273,7 +272,7 @@ class OmniStreamingSpeechHandler: # Reject unmet word-timestamps preconditions early with a clear reason. if config.word_timestamps: - if self._speech_service.forced_aligner_config is None: + if not self._speech_service.forced_aligner_enabled: await self._send_error( websocket, "word_timestamps=true but the server was launched without " @@ -285,7 +284,8 @@ class OmniStreamingSpeechHandler: await self._send_error( websocket, "word_timestamps=true requires stream_audio=true and " - "response_format='pcm' (the aligner consumes raw PCM).", + "response_format='pcm' (timestamps ride the per-sentence " + "PCM audio.chunk stream).", ) return @@ -395,19 +395,20 @@ class OmniStreamingSpeechHandler: ) -> int: """Stream PCM as JSON ``audio.chunk`` frames, aligned per sentence. - Forward each PCM chunk live (``timestamps: null``) while buffering the - sentence audio, then run the forced aligner once over the whole - sentence and emit a final empty-audio ``audio.chunk`` with the word - timestamps. On aligner failure timestamps is ``null``; for silence it - is ``[]`` (audio always flows regardless). + Forward each PCM chunk live (``timestamps: null``). The forced-aligner + pipeline stage (appended when the server is launched with + ``--forced-aligner``) consumes the synthesized audio internally and its + pooling output rides the same generator, so once the audio finishes we + pull the word timestamps straight off that aligner output and emit a + final empty-audio ``audio.chunk`` carrying them. Timestamps is ``null`` + when the aligner produced none; audio always flows regardless. """ - aligner_config = self._speech_service.forced_aligner_config - assert aligner_config is not None # gated by the precondition check - - sentence_audio = bytearray() + audio_bytes_seen = 0 total_bytes = 0 sample_rate = _PCM_SAMPLE_RATE chunk_id = 0 + # Receives the aligner stage's pooling output from the generator. + collect: dict = {} async def send_chunk( chunk: bytes, @@ -437,62 +438,30 @@ class OmniStreamingSpeechHandler: generator, request_id, include_sample_rate=True, + collect=collect, ) ) as stream: async for chunk, chunk_sample_rate in stream: sample_rate = chunk_sample_rate - chunk_start_ms = int(round((len(sentence_audio) / _BYTES_PER_SAMPLE / sample_rate) * 1000.0)) - sentence_audio.extend(chunk) - chunk_end_ms = int(round((len(sentence_audio) / _BYTES_PER_SAMPLE / sample_rate) * 1000.0)) + chunk_start_ms = int(round((audio_bytes_seen / _BYTES_PER_SAMPLE / sample_rate) * 1000.0)) + audio_bytes_seen += len(chunk) + chunk_end_ms = int(round((audio_bytes_seen / _BYTES_PER_SAMPLE / sample_rate) * 1000.0)) total_bytes += len(chunk) # Audio first, timestamps after the whole sentence is aligned. await send_chunk(chunk, chunk_sample_rate, None, chunk_start_ms, chunk_end_ms) - # Single alignment pass over the full sentence, then emit timestamps. - # A load/config failure is permanent, so surface the reason once; audio - # has already streamed, so the trailing frame still carries null. - try: - timestamps_payload = await self._align_sentence( - audio=bytes(sentence_audio), - text=sentence_text, - sample_rate=sample_rate, - config=aligner_config, - language=language, - ) - except ForcedAlignerLoadError as exc: - await self._send_error(websocket, f"forced aligner unavailable: {exc}") - timestamps_payload = None - sentence_end_ms = int(round((len(sentence_audio) / _BYTES_PER_SAMPLE / sample_rate) * 1000.0)) + # Pull word timestamps off the aligner stage's pooling output (it rode + # the same generator); extract_word_timestamps re-segments the sentence + # text for the word strings when the aligner output doesn't carry them. + aligner_res = collect.get("aligner_res") + timestamps_payload = ( + extract_word_timestamps(aligner_res, sentence_text, language) if aligner_res is not None else None + ) + sentence_end_ms = int(round((audio_bytes_seen / _BYTES_PER_SAMPLE / sample_rate) * 1000.0)) await send_chunk(b"", sample_rate, timestamps_payload, 0, sentence_end_ms) return total_bytes - @staticmethod - async def _align_sentence( - *, - audio: bytes, - text: str, - sample_rate: int, - config, - language: str | None = None, - ) -> list[dict] | None: - """Convert a sentence alignment into JSON word-timestamp dicts. - - Returns ``None`` on aligner failure, ``[]`` when it ran but produced no - tokens. Monotonic, non-overlapping bounds are guaranteed by the decoder. - ``language`` is forwarded to word segmentation. - """ - aligned = await forced_align( - audio=audio, - text=text, - sample_rate=sample_rate, - config=config, - language=language, - ) - if aligned is None: - return None - return [{"word": ts.word, "start_ms": ts.start_ms, "end_ms": ts.end_ms} for ts in aligned] - @staticmethod async def _send_error(websocket: WebSocket, message: str) -> None: """Send an error message to the client.""" diff --git a/vllm_omni/model_executor/stage_input_processors/forced_aligner.py b/vllm_omni/model_executor/stage_input_processors/forced_aligner.py new file mode 100644 index 000000000..7646c6f73 --- /dev/null +++ b/vllm_omni/model_executor/stage_input_processors/forced_aligner.py @@ -0,0 +1,173 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""stage input processor for a forced-aligner stage. + +Bridges a Code2Wav (audio) stage to a pooling forced-aligner stage: + + Talker (LLM_AR) -> Code2Wav (LLM_GENERATION) -> ForcedAligner (LLM_AR + runner=pooling) + +``code2wav2aligner`` is a synchronous ``custom_process_input_func`` (the same +shape as ``qwen3_tts.talker2code2wav``). It runs in the orchestrator process, +reads the synthesized waveform from the upstream Code2Wav output and the +ground-truth text from the original request, then builds the aligner prompt +(```` markers per word) + audio multimodal input. +""" + +from __future__ import annotations + +from typing import Any + +import numpy as np +import torch +from vllm.logger import init_logger + +from vllm_omni.model_executor.stage_input_processors.tts_utils import ( + extract_language_from_prompt, +) +from vllm_omni.utils.qwen3_force_align_processor import build_prompt, segment_words + +logger = init_logger(__name__) + +# Shared keys / names for the forced-aligner stage payloads, so the producer +# (this processor), the worker-side decode, and the serving consumer don't +# duplicate magic strings. +ALIGNER_WORDS_KEY = "aligner_words" +ALIGNER_DURATION_MS_KEY = "aligner_audio_duration_ms" +TIMESTAMPS_MODALITY = "timestamps" +# Dotted path of the pooling-output decoder hook (read by the AR scheduler for +# a stage whose engine_args set ``pooling_output_decoder``). +POOLING_OUTPUT_DECODER_PATH = "vllm_omni.model_executor.stage_input_processors.forced_aligner.decode_pooling_output" + + +def decode_pooling_output(logits: Any, request: Any, hf_config: Any) -> Any: + """Worker-side decoder for the forced-aligner pooling stage. + + Converts the aligner's ``[n_token, classify_num]`` logits into a bare + ``[n_words, 2]`` int32 tensor of per-word ``[start_ms, end_ms]`` (word + strings travel separately via ``additional_information``). + """ + from vllm_omni.engine.serialization import deserialize_additional_information + from vllm_omni.utils.forced_aligner import _decode_timestamps + + ts_id = int(getattr(hf_config, "timestamp_token_id")) + thinker_cfg = getattr(hf_config, "thinker_config", hf_config) + classify_num = int(getattr(thinker_cfg, "classify_num")) + seg_ms = float(getattr(hf_config, "timestamp_segment_time")) + + prompt_token_ids = list(getattr(request, "prompt_token_ids", []) or []) + positions = [i for i, t in enumerate(prompt_token_ids) if t == ts_id] + + info = deserialize_additional_information(getattr(request, "additional_information", None)) or {} + words_field = info.get(ALIGNER_WORDS_KEY) + words = words_field[0] if isinstance(words_field, list) and words_field else [] + + # Real waveform duration (stashed by code2wav2aligner) drives the decode clamp; fall back to the + # theoretical ceiling (clamp no-op) if absent. + duration_field = info.get(ALIGNER_DURATION_MS_KEY) + duration_ms = ( + float(duration_field[0]) + if isinstance(duration_field, list) and duration_field + else float(classify_num) * seg_ms + ) + + timestamps = _decode_timestamps( + logits=logits, + words=list(words), + timestamp_positions=positions, + classify_num=classify_num, + timestamp_segment_time_ms=seg_ms, + audio_duration_ms=duration_ms, + ) + # No failure sentinel: an empty [0, 2] tensor decodes downstream to timestamps: null. + ms = torch.tensor( + [[t.start_ms, t.end_ms] for t in timestamps], + dtype=torch.int32, + ).reshape(-1, 2) + return ms + + +def _extract_waveform(multimodal_output: dict[str, Any]) -> tuple[np.ndarray, int]: + """Pull a mono float32 waveform + sample rate from a Code2Wav output.""" + wav = multimodal_output.get("audio") + if wav is None: + wav = multimodal_output.get("model_outputs") + if isinstance(wav, list): + wav = wav[0] if wav else None + if isinstance(wav, torch.Tensor): + wav_np = wav.detach().cpu().to(torch.float32).flatten().numpy() + elif isinstance(wav, np.ndarray): + wav_np = wav.astype(np.float32).flatten() + else: + wav_np = np.zeros(0, dtype=np.float32) + + sr_raw = multimodal_output.get("sr") + if isinstance(sr_raw, list): + sr_raw = sr_raw[-1] if sr_raw else None + if isinstance(sr_raw, torch.Tensor): + sr = int(sr_raw.reshape(-1)[-1].item()) if sr_raw.numel() > 0 else 24000 + elif sr_raw is None: + sr = 24000 + else: + sr = int(sr_raw) + return wav_np, sr + + +def _extract_text(prompt: Any, index: int = 0) -> str: + """Read the ground-truth text the aligner needs from the original prompt.""" + if not isinstance(prompt, dict): + return "" + text = (prompt.get("additional_information") or {}).get("text") + if isinstance(text, list): + return str(text[index]) if index < len(text) else "" + return str(text) if text else "" + + +def code2wav2aligner( + source_outputs: list[Any], + prompt: Any = None, + _requires_multimodal_data: bool = False, + _streaming_context: Any | None = None, +) -> list[Any]: + """Build forced-aligner stage inputs from finished Code2Wav audio. + + Emits a **text prompt** (``prompt`` string + ``multi_modal_data`` audio). + Tokenization and audio feature extraction are deferred to the aligner + stage's input preprocessor (run by the orchestrator), so this function + needs neither the aligner tokenizer nor its mm processor and stays + unit-testable on CPU. + """ + from vllm_omni.inputs.data import OmniTextPrompt + + aligner_inputs: list[OmniTextPrompt] = [] + for i, src in enumerate(source_outputs): + if not getattr(src, "finished", False): + # Sentence-final alignment: only run once the audio is complete. + continue + mm = src.outputs[0].multimodal_output or {} + wav_np, sample_rate = _extract_waveform(mm) + if wav_np.size == 0: + logger.warning("code2wav2aligner: empty waveform for request; skipping alignment.") + continue + + text = _extract_text(prompt, index=i) + language_list = extract_language_from_prompt(prompt, index=i) + language = language_list[0] if language_list else None + words = segment_words(text, language) + prompt_str = build_prompt(words) + + additional_information: dict[str, Any] = { + "text": [text], + ALIGNER_WORDS_KEY: [words], + ALIGNER_DURATION_MS_KEY: [float(wav_np.size) / float(sample_rate) * 1000.0], + } + if language_list is not None: + additional_information["language"] = language_list + + aligner_inputs.append( + OmniTextPrompt( + prompt=prompt_str, + multi_modal_data={"audio": (wav_np, sample_rate)}, + additional_information=additional_information, + ) + ) + return aligner_inputs diff --git a/vllm_omni/utils/forced_aligner.py b/vllm_omni/utils/forced_aligner.py index 2a97df3d8..3b578ee51 100644 --- a/vllm_omni/utils/forced_aligner.py +++ b/vllm_omni/utils/forced_aligner.py @@ -1,30 +1,29 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Shared forced aligner for streaming TTS word timestamps. +"""Forced-aligner config + word-timestamp decoding for TTS. -Hosts one lazy-loaded, in-process pooling ``vllm.LLM`` shared by the TTS -frontend. ``llm.encode`` is sync/blocking, so :func:`align` runs it in -``asyncio.to_thread``, serialized on a lock since the offline LLM is not -thread-safe. See :func:`align` for the return-value contract. +The aligner runs as a pooling stage appended to the pipeline by +``--forced-aligner`` (see :func:`inject_forced_aligner_stage`). Qwen-specific +word segmentation / prompt building / marker repair live in +:mod:`vllm_omni.utils.qwen3_force_align_processor`. """ from __future__ import annotations -import asyncio import logging -import os -import threading -from dataclasses import dataclass, field +from dataclasses import dataclass, replace from pathlib import Path -from typing import Any +from typing import TYPE_CHECKING, Any import numpy as np import yaml from vllm_omni.utils import qwen3_force_align_processor as _processor -logger = logging.getLogger(__name__) +if TYPE_CHECKING: + from vllm_omni.config.stage_config import DeployConfig, PipelineConfig +logger = logging.getLogger(__name__) _DEFAULT_CONFIG_PATH = Path(__file__).resolve().parents[1] / "deploy" / "qwen3_tts_forced_aligner.yaml" @@ -42,7 +41,8 @@ class WordTimestamp: @dataclass(frozen=True, slots=True) class ForcedAlignerConfig: - """Plain-data config for constructing the aligner ``vllm.LLM`` (lazy-loaded).""" + """Plain-data config for the forced-aligner pipeline stage; fields are + consumed by :func:`inject_forced_aligner_stage`.""" model: str runner: str | None = None @@ -52,7 +52,6 @@ class ForcedAlignerConfig: dtype: str | None = None max_model_len: int | None = None trust_remote_code: bool | None = None - extra_llm_kwargs: dict[str, Any] = field(default_factory=dict) def _load_forced_aligner_yaml(path: str | Path) -> dict[str, Any]: @@ -61,8 +60,8 @@ def _load_forced_aligner_yaml(path: str | Path) -> dict[str, Any]: return dict(raw.get("forced_aligner", raw)) -def build_forced_aligner_config(args: Any) -> ForcedAlignerConfig | None: - """Build a config from CLI args, or ``None`` when the feature is off. +def build_forced_aligner_config(model: str | None, config_path: str | None = None) -> ForcedAlignerConfig | None: + """Build a config from the CLI values, or ``None`` when the feature is off. Precedence (lowest to highest): the packaged default YAML (:data:`_DEFAULT_CONFIG_PATH`, Qwen deploy defaults) -> a user YAML passed @@ -70,254 +69,94 @@ def build_forced_aligner_config(args: Any) -> ForcedAlignerConfig | None: feature is off (returns ``None``) unless a model resolves from this chain. Per-field overrides such as ``gpu_memory_utilization`` live in the YAML. """ - config_path = getattr(args, "forced_aligner_config", None) config_data: dict[str, Any] = {} - model = getattr(args, "forced_aligner", None) if config_path or model: config_data.update(_load_forced_aligner_yaml(_DEFAULT_CONFIG_PATH)) if config_path: config_data.update(_load_forced_aligner_yaml(config_path)) if model: config_data["model"] = str(model) - model = config_data.get("model") - if not model: + if not config_data.get("model"): return None allowed = set(ForcedAlignerConfig.__dataclass_fields__) return ForcedAlignerConfig(**{k: v for k, v in config_data.items() if k in allowed}) -class ForcedAlignerLoadError(RuntimeError): - """The aligner model/config could not be loaded (permanent until restart). - - Distinguished from a per-request alignment failure so the streaming layer - can report a clear reason once rather than degrading every request to - ``timestamps: null``. +def inject_forced_aligner_stage( + pipeline: PipelineConfig, + deploy: DeployConfig, + cli_overrides: dict[str, Any], +) -> tuple[PipelineConfig, DeployConfig]: + """Append a forced-aligner pooling stage to the pipeline tail when + ``--forced-aligner`` or ``--forced-aligner-config`` resolves an aligner + model; no-op otherwise. The stage runs the aligner model with + ``runner="pooling"``, consuming the previous stage's audio and emitting a + terminal word-timestamps side-output. """ + if not (cli_overrides.get("forced_aligner") or cli_overrides.get("forced_aligner_config")): + return pipeline, deploy + from vllm_omni.config.stage_config import ( + StageDeployConfig, + StageExecutionType, + StagePipelineConfig, + ) + from vllm_omni.model_executor.stage_input_processors.forced_aligner import ( + POOLING_OUTPUT_DECODER_PATH, + TIMESTAMPS_MODALITY, + ) -# --- Singleton state --- -# A single LLM serves the whole API server. ``_lock`` guards the lazy -# constructor; ``_encode_lock`` serializes ``encode`` because the offline -# ``vllm.LLM`` is a sync batch interface that is not safe to call from multiple -# threads at once (concurrent WebSocket sessions otherwise race the engine). -_lock = threading.Lock() -_encode_lock = asyncio.Lock() -_llm: Any = None -_classify_num: int | None = None -_timestamp_token_id: int | None = None -_timestamp_segment_time_ms: float | None = None -_loaded_config: ForcedAlignerConfig | None = None + fa = build_forced_aligner_config( + cli_overrides.get("forced_aligner"), + cli_overrides.get("forced_aligner_config"), + ) + if fa is None: + return pipeline, deploy + _FA = "vllm_omni.model_executor.stage_input_processors.forced_aligner" + new_id = len(pipeline.stages) + aligner_ps = StagePipelineConfig( + stage_id=new_id, + model_stage="forced_aligner", + # Pooling stage = LLM_AR stage run with runner="pooling" (vLLM derives + # is_pooling_model from runner_type). + execution_type=StageExecutionType.LLM_AR, + input_sources=(new_id - 1,), + final_output=True, + final_output_type=TIMESTAMPS_MODALITY, + owns_tokenizer=True, + requires_multimodal_data=True, + model_arch=fa.architecture, + custom_process_input_func=f"{_FA}.code2wav2aligner", + ) + extended = replace(pipeline, stages=pipeline.stages + (aligner_ps,)) -async def align( - *, - audio: bytes, - text: str, - sample_rate: int, - config: ForcedAlignerConfig, - language: str | None = None, -) -> list[WordTimestamp] | None: - """Run one forced-alignment pass. - - Args: - audio: Signed-int16 little-endian mono PCM bytes. - text: Ground-truth text whose words to align. - sample_rate: Sample rate of ``audio`` in Hz. - config: Aligner config (same instance for every call across the - server's lifetime; reload requires a server restart). - language: Optional language hint for word segmentation. ``None`` / - ``"auto"`` use the space + Chinese-mixed path; ``"japanese"`` / - ``"korean"`` (or their codes) need ``qwen_asr`` installed for - faithful tokenisation, else they degrade to whitespace splitting. - - Returns: - List of :class:`WordTimestamp` on success (possibly empty for - silence / no aligned tokens), ``None`` on a per-request alignment - failure (decode error). - - Raises: - ForcedAlignerLoadError: the model/config could not be loaded. This is - permanent until restart, so callers surface it once instead of - degrading every request to ``None``. - """ - async with _encode_lock: - try: - await asyncio.to_thread(_ensure_loaded, config) - except Exception as exc: # noqa: BLE001 - logger.exception("Forced aligner failed to load") - raise ForcedAlignerLoadError(str(exc)) from exc - try: - return await asyncio.to_thread(_align_sync, audio, text, sample_rate, config, language) - except Exception: # noqa: BLE001 - logger.exception("Forced aligner failed for text=%r", text) - return None - - -def _align_sync( - audio: bytes, - text: str, - sample_rate: int, - config: ForcedAlignerConfig, - language: str | None = None, -) -> list[WordTimestamp]: - # Precondition: align() has already run _ensure_loaded(config) under - # _encode_lock and owns the load-failure -> ForcedAlignerLoadError mapping. - # Don't reload here, so a load error can't be swallowed as a per-request None. - audio_arr = _pcm_bytes_to_float32(audio) - if audio_arr.size == 0: - return [] - audio_duration_ms = (audio_arr.size / sample_rate) * 1000.0 - - # Segment once and reuse for both the prompt and the decode: the word - # units MUST match between the two or the markers drift out of sync. - words = _processor.segment_words(text, language) - prompt = _processor.build_prompt(words) - request = { - "prompt": prompt, - "multi_modal_data": {"audio": (audio_arr, sample_rate)}, + engine_extras: dict[str, Any] = { + "model": fa.model, + "runner": fa.runner or "pooling", + # Worker-side decoder hook: pooler logits -> word-timestamps tensor. + "pooling_output_decoder": POOLING_OUTPUT_DECODER_PATH, } - - from vllm.pooling_params import PoolingParams - - outputs = _llm.encode( # type: ignore[union-attr] - [request], - pooling_params=PoolingParams(), - pooling_task=config.pooling_task or "token_classify", - use_tqdm=False, + if fa.architecture: + engine_extras["hf_overrides"] = {"architectures": [fa.architecture]} + if fa.trust_remote_code is not None: + engine_extras["trust_remote_code"] = fa.trust_remote_code + if fa.dtype is not None: + engine_extras["dtype"] = fa.dtype + aligner_ds = StageDeployConfig( + stage_id=new_id, + devices=cli_overrides.get("forced_aligner_device"), + gpu_memory_utilization=fa.gpu_memory_utilization, + max_model_len=fa.max_model_len, + engine_extras=engine_extras, + default_pooling_params={"task": fa.pooling_task or "token_classify"}, + # Pooling is one-shot -> sync AR scheduler (the async one assumes + # token-by-token generation). + async_scheduling=False, ) - if not outputs: - return [] - - result = outputs[0] - logits = result.outputs.data # [n_token, classify_num] - prompt_token_ids = list(result.prompt_token_ids) - timestamp_positions = [i for i, tid in enumerate(prompt_token_ids) if tid == _timestamp_token_id] - if not timestamp_positions: - logger.warning( - "No <|timestamp|> tokens found in prompt for text=%r; aligner returned %d rows.", - text, - logits.shape[0] if hasattr(logits, "shape") else len(logits), - ) - return [] - - return _decode_timestamps( - logits=logits, - words=words, - timestamp_positions=timestamp_positions, - classify_num=_classify_num, - timestamp_segment_time_ms=_timestamp_segment_time_ms, - audio_duration_ms=audio_duration_ms, - ) - - -def _ensure_loaded(config: ForcedAlignerConfig) -> None: - """Lazy-load the singleton ``vllm.LLM`` under lock; idempotent.""" - global _llm, _classify_num, _timestamp_token_id, _timestamp_segment_time_ms, _loaded_config - - if _llm is not None: - if _loaded_config is not None and _loaded_config.model != config.model: - # Multiple configs from different requests — refuse rather - # than swap models silently. A server restart is required. - raise RuntimeError( - f"Forced aligner already loaded with config={_loaded_config!r}; " - f"cannot serve a request that asks for {config!r}. " - "Restart the server to change the aligner config." - ) - return - - with _lock: - if _llm is not None: - return # raced; another caller did the load - - # Lazy import: vllm pulls torch + CUDA, which we want to avoid at - # module import time. - from vllm import LLM - - logger.info( - "Loading forced aligner %s (gpu_memory_utilization=%s)", - config.model, - config.gpu_memory_utilization if config.gpu_memory_utilization is not None else "default", - ) - llm_kwargs: dict[str, Any] = {"model": config.model} - if config.runner is not None: - llm_kwargs["runner"] = config.runner - if config.architecture is not None: - llm_kwargs["hf_overrides"] = {"architectures": [config.architecture]} - if config.gpu_memory_utilization is not None: - llm_kwargs["gpu_memory_utilization"] = config.gpu_memory_utilization - if config.trust_remote_code is not None: - llm_kwargs["trust_remote_code"] = config.trust_remote_code - if config.dtype is not None: - llm_kwargs["dtype"] = config.dtype - if config.max_model_len is not None: - llm_kwargs["max_model_len"] = config.max_model_len - llm_kwargs.update(config.extra_llm_kwargs) - - # Force the aligner engine in-process: in multiprocessing mode the - # token_classify ``pooling_output`` fails to round-trip over the - # engine-core IPC on current vLLM. Hard-set (not setdefault) so a - # globally exported ``=1`` can't push it into a subprocess; save/restore - # to avoid leaking the override to anything built later in this process. - prev_mp = os.environ.get("VLLM_ENABLE_V1_MULTIPROCESSING") - os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0" - try: - llm = LLM(**llm_kwargs) - finally: - if prev_mp is None: - os.environ.pop("VLLM_ENABLE_V1_MULTIPROCESSING", None) - else: - os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = prev_mp - - thinker_config = getattr(llm.llm_engine.model_config.hf_config, "thinker_config", None) - if thinker_config is None or not hasattr(thinker_config, "classify_num"): - raise RuntimeError( - "Loaded aligner has no thinker_config.classify_num; " - "expected a Qwen3ASRForcedAlignerForTokenClassification checkpoint." - ) - timestamp_segment_time_ms = getattr(llm.llm_engine.model_config.hf_config, "timestamp_segment_time", None) - if timestamp_segment_time_ms is None: - raise RuntimeError( - "Loaded aligner has no timestamp_segment_time; expected a Qwen3ASR forced aligner checkpoint." - ) - - tokenizer = llm.get_tokenizer() - timestamp_token_id = _processor.resolve_timestamp_token_id(tokenizer) - - # Publish in this order so a concurrent reader either sees - # _llm == None (will block on the lock) or sees a fully - # initialized aligner. - _classify_num = int(thinker_config.classify_num) - _timestamp_token_id = timestamp_token_id - _timestamp_segment_time_ms = float(timestamp_segment_time_ms) - _loaded_config = config - _llm = llm - - logger.info( - "Forced aligner ready: timestamp_token_id=%d, classify_num=%d, timestamp_segment_time_ms=%.1f", - timestamp_token_id, - _classify_num, - _timestamp_segment_time_ms, - ) - - -# --- pure helpers (testable without a GPU / vllm) --- -# -# Word segmentation, prompt building, timestamp repair and marker-token -# lookup are Qwen-specific and live in -# :mod:`vllm_omni.utils.qwen3_force_align_processor`. This module owns only -# the generic vLLM orchestration. - - -def _pcm_bytes_to_float32(audio: bytes) -> np.ndarray: - """Decode signed-int16 mono PCM bytes into a [-1, 1] float32 array.""" - if not audio: - return np.zeros(0, dtype=np.float32) - if len(audio) % 2 != 0: - # Drop a trailing odd byte rather than raise; keeps streaming - # robust against off-by-one chunk boundaries. - audio = audio[:-1] - pcm = np.frombuffer(audio, dtype=np.int16) - return (pcm.astype(np.float32) / 32768.0).copy() + deploy = replace(deploy, stages=list(deploy.stages) + [aligner_ds]) + logger.info("[stage_init] Injected forced-aligner stage %d (model=%s)", new_id, fa.model) + return extended, deploy def _decode_timestamps( @@ -382,18 +221,40 @@ def _decode_timestamps( return out -# Test hooks --------------------------------------------------------------- -# Tests need a way to reset module state without restarting Python. Not -# part of the public API; do not call in production code. +def extract_word_timestamps(res: Any, text: str, language: str | None = None) -> list[dict] | None: + """Read a forced-aligner stage's engine output into per-word timestamps. + Returns ``[{word, start_ms, end_ms}, ...]`` from the stage's pooling output + (``res.outputs.data`` is an int32 ``[n_words, 2]`` tensor). Words come from + re-segmenting ``text``; any mismatch or non-aligner input returns ``None``. + """ + if not text: + return None + try: + outs = getattr(res, "outputs", None) + if outs is None: + return None + # The aligner stage's res.outputs is a single PoolingOutput (not a list). + # Its .data is a [n_words, 2] int32 tensor of [start_ms, end_ms] per word. + out0 = outs[0] if isinstance(outs, (list, tuple)) else outs + ms = getattr(out0, "data", None) + if ms is None or not hasattr(ms, "tolist"): + return None + ms_list = ms.tolist() + if not ms_list or not isinstance(ms_list[0], (list, tuple)): + return None -def _reset_for_tests() -> None: - """Drop the cached aligner state so the next call reloads.""" - global _llm, _classify_num, _timestamp_token_id, _timestamp_segment_time_ms, _loaded_config - with _lock: - _llm = None - _classify_num = None - _timestamp_token_id = None - _timestamp_segment_time_ms = None - _loaded_config = None - _processor._reset_for_tests() + words = _processor.segment_words(text, language) + if len(words) != len(ms_list): + # Segmentation drifted from the engine side; pairing positionally would mislabel every word. + logger.warning( + "word/timestamp count mismatch (%d words vs %d intervals); returning no timestamps", + len(words), + len(ms_list), + ) + return None + + return [{"word": w, "start_ms": int(p[0]), "end_ms": int(p[1])} for w, p in zip(words, ms_list)] + except Exception: + logger.debug("failed to extract word timestamps from aligner output", exc_info=True) + return None diff --git a/vllm_omni/utils/qwen3_force_align_processor.py b/vllm_omni/utils/qwen3_force_align_processor.py index 66a084b44..e70d270db 100644 --- a/vllm_omni/utils/qwen3_force_align_processor.py +++ b/vllm_omni/utils/qwen3_force_align_processor.py @@ -12,15 +12,15 @@ """Qwen3 forced-aligner text/timestamp processor. This is the model-specific half of upstream's ``Qwen3ForceAlignProcessor``: -it turns text into the aligner's word units and prompt, repairs the predicted -timestamp bins, and resolves the marker token id. The generic "run a vLLM -pooling model" half lives in :mod:`vllm_omni.utils.forced_aligner`, which -calls into this module. +it turns text into the aligner's word units and prompt, and repairs the +predicted timestamp bins. It feeds the forced-aligner pipeline stage (input +glue in :mod:`vllm_omni.model_executor.stage_input_processors.forced_aligner`, +timestamp decoding in :mod:`vllm_omni.utils.forced_aligner`). Keeping the Qwen-specific pieces here marks the seam for the model-agnostic aligner the issue asks for: a different aligner family would supply its own processor exposing the same small surface — :func:`segment_words`, -:func:`build_prompt`, :func:`fix_timestamp`, :func:`resolve_timestamp_token_id`. +:func:`build_prompt`, :func:`fix_timestamp`. Word segmentation prefers qwen_asr's official ``Qwen3ForceAlignProcessor`` when installed (full multilingual fidelity, incl. Japanese/Korean) and @@ -260,26 +260,3 @@ def fix_timestamp(values: list[int]) -> list[int]: i = j return [int(r) for r in result] - - -def resolve_timestamp_token_id(tokenizer: Any, timestamp_token: str = TIMESTAMP_TOKEN) -> int: - """Look up the integer id of the timestamp special token.""" - convert = getattr(tokenizer, "convert_tokens_to_ids", None) - if not callable(convert): - raise RuntimeError("Aligner tokenizer has no convert_tokens_to_ids method.") - tid = convert(timestamp_token) - if isinstance(tid, list): - tid = tid[0] if tid else None - if tid is None or (isinstance(tid, int) and tid < 0): - raise RuntimeError( - f"Aligner tokenizer does not recognise {timestamp_token!r} (got id={tid}). " - "Check the model card; the marker token may use a different name." - ) - return int(tid) - - -def _reset_for_tests() -> None: - """Drop the cached official-processor probe so the next call re-checks.""" - global _official_processor, _official_processor_unavailable - _official_processor = None - _official_processor_unavailable = False diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py index b68c722ec..0f65489b7 100644 --- a/vllm_omni/worker/gpu_ar_model_runner.py +++ b/vllm_omni/worker/gpu_ar_model_runner.py @@ -1360,7 +1360,12 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): hidden_states = model_output aux_hidden_states = None - hidden_states, multimodal_outputs = self.extract_multimodal_outputs(model_output) + # A pooling stage skips mm extraction + the omni prefix cache (both assume an AR + # [n_tok, hidden] output) and exits through the shared _pool() branch below. + if self.is_pooling_model: + multimodal_outputs = None + else: + hidden_states, multimodal_outputs = self.extract_multimodal_outputs(model_output) hidden_states_cpu = None # Async-write pipeline (replaces the per-step blocking @@ -1369,7 +1374,7 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): # the actual CPU scatter into ``hidden_states_cache`` / # ``mm_outputs_cache`` happens in ``drain_ready_async_writes`` # at the top of subsequent execute_model() calls. - if self.omni_prefix_cache is not None and get_pp_group().is_last_rank: + if not self.is_pooling_model and self.omni_prefix_cache is not None and get_pp_group().is_last_rank: hs_for_cache = hidden_states if self._model_needs_full_prefix_hidden_states() else None # Some models (e.g. qwen3-tts-talker) opt out of full-hidden-state # prefix caching but the downstream pooler payload path still