[Bugfix] Stream Qwen3-TTS WebSocket input as one request (#4731)

Signed-off-by: Sy03 <1370724210@qq.com>
This commit is contained in:
Sy03 authored and GitHub committed 2026-06-26 07:48:30 +00:00
1 parent 580a0efe40
commit c39ac9fffd
7 files changed
+34 -432

No files matched your search

+3 -4
View File
@@ -233,9 +233,9 @@ Server -> Client:
| Message | Description |
|---------|-------------|
| `{"type": "audio.start", "sentence_index": 0, "sentence_text": "...", "format": "pcm", "sample_rate": 24000}` | Audio generation starting for a sentence |
| `{"type": "audio.start", "sentence_index": 0, "sentence_text": "...", "format": "pcm", "sample_rate": 24000}` | Audio generation starting for the buffered input |
| Binary frame | Raw audio bytes (one or more PCM chunks when `stream_audio=true`) |
| `{"type": "audio.done", "sentence_index": 0, "total_bytes": 96000, "error": false}` | Audio complete for a sentence |
| `{"type": "audio.done", "sentence_index": 0, "total_bytes": 96000, "error": false}` | Audio complete for the buffered input |
| `{"type": "session.done", "total_sentences": N}` | Session complete |
| `{"type": "error", "message": "..."}` | Non-fatal error |
@@ -245,8 +245,7 @@ All REST API parameters are supported, plus:
| Parameter | Type | Default | Description |
|-----------|------|---------|-------------|
| `stream_audio` | bool | false | Stream one or more PCM chunks per sentence over WebSocket |
| `split_granularity` | string | "sentence" | Text splitting granularity |
| `stream_audio` | bool | false | Stream one or more PCM chunks for the buffered input over WebSocket |
```bash
@@ -113,9 +113,9 @@ class TestQwen3TTSWebSocket:
session_done = result["session_done"]
assert session_done is not None
assert session_done["total_sentences"] == 2
assert len(starts) == 2
assert len(dones) == 2
assert session_done["total_sentences"] == 1
assert len(starts) == 1
assert len(dones) == 1
for idx, start in enumerate(starts):
assert start["type"] == "audio.start"
@@ -60,6 +60,7 @@ class TestStreamingSpeechWebSocket:
with client.websocket_connect("/v1/audio/speech/stream") as ws:
ws.send_json({"type": "session.config", "voice": "Vivian"})
ws.send_json({"type": "input.text", "text": "Hello world. "})
ws.send_json({"type": "input.done"})
start = ws.receive_json()
assert start["type"] == "audio.start"
@@ -73,7 +74,6 @@ class TestStreamingSpeechWebSocket:
done = ws.receive_json()
assert done == {"type": "audio.done", "sentence_index": 0, "total_bytes": len(audio), "error": False}
ws.send_json({"type": "input.done"})
session_done = ws.receive_json()
assert session_done == {"type": "session.done", "total_sentences": 1}
@@ -113,6 +113,7 @@ class TestStreamingSpeechWebSocket:
}
)
ws.send_json({"type": "input.text", "text": "Hello world. "})
ws.send_json({"type": "input.done"})
start = ws.receive_json()
assert start["type"] == "audio.start"
@@ -126,7 +127,6 @@ class TestStreamingSpeechWebSocket:
done = ws.receive_json()
assert done == {"type": "audio.done", "sentence_index": 0, "total_bytes": 6, "error": False}
ws.send_json({"type": "input.done"})
assert ws.receive_json() == {"type": "session.done", "total_sentences": 1}
assert len(captured_requests) == 1
@@ -150,6 +150,7 @@ class TestStreamingSpeechWebSocket:
}
)
ws.send_json({"type": "input.text", "text": "Hello world. "})
ws.send_json({"type": "input.done"})
error = ws.receive_json()
assert error["type"] == "error"
@@ -199,6 +200,7 @@ class TestStreamingSpeechWebSocket:
}
)
ws.send_json({"type": "input.text", "text": "Hello world. "})
ws.send_json({"type": "input.done"})
start = ws.receive_json()
assert start["type"] == "audio.start"
@@ -279,6 +281,7 @@ class TestStreamingSpeechWebSocket:
}
)
ws.send_json({"type": "input.text", "text": "Hello world. "})
ws.send_json({"type": "input.done"})
assert ws.receive_json()["type"] == "audio.start"
# Real-time audio chunk (timestamps null), then the timestamp frame.
@@ -342,16 +345,19 @@ class TestStreamingSpeechWebSocket:
assert speech_service._generate_audio_bytes.await_count == 0
def test_multiple_sentences_increment_indices(self, mocker: MockerFixture):
def test_multiple_sentences_are_buffered_into_one_request(self, mocker: MockerFixture):
app, _ = _build_test_app(mocker=mocker)
with TestClient(app) as client:
with client.websocket_connect("/v1/audio/speech/stream") as ws:
ws.send_json({"type": "session.config", "voice": "Vivian"})
ws.send_json({"type": "input.text", "text": "First sentence. Second sentence. "})
ws.send_json({"type": "input.text", "text": "First sentence. "})
ws.send_json({"type": "input.text", "text": "Second sentence. "})
ws.send_json({"type": "input.done"})
first_start = ws.receive_json()
assert first_start["sentence_index"] == 0
start = ws.receive_json()
assert start["sentence_index"] == 0
assert start["sentence_text"] == "First sentence. Second sentence."
ws.receive_bytes()
assert ws.receive_json() == {
"type": "audio.done",
@@ -359,19 +365,7 @@ class TestStreamingSpeechWebSocket:
"total_bytes": 36,
"error": False,
}
second_start = ws.receive_json()
assert second_start["sentence_index"] == 1
ws.receive_bytes()
assert ws.receive_json() == {
"type": "audio.done",
"sentence_index": 1,
"total_bytes": 36,
"error": False,
}
ws.send_json({"type": "input.done"})
assert ws.receive_json() == {"type": "session.done", "total_sentences": 2}
assert ws.receive_json() == {"type": "session.done", "total_sentences": 1}
def test_unknown_message_type_keeps_session_open(self, mocker: MockerFixture):
app, _ = _build_test_app(mocker=mocker)
@@ -385,6 +379,7 @@ class TestStreamingSpeechWebSocket:
assert error == {"type": "error", "message": "Unknown message type: unknown"}
ws.send_json({"type": "input.text", "text": "Hello world. "})
ws.send_json({"type": "input.done"})
assert ws.receive_json()["type"] == "audio.start"
ws.receive_bytes()
assert ws.receive_json() == {
@@ -394,7 +389,6 @@ class TestStreamingSpeechWebSocket:
"error": False,
}
ws.send_json({"type": "input.done"})
assert ws.receive_json() == {"type": "session.done", "total_sentences": 1}
def test_config_timeout_closes_session(self, mocker: MockerFixture):
@@ -419,12 +413,12 @@ class TestStreamingSpeechWebSocket:
with client.websocket_connect("/v1/audio/speech/stream") as ws:
ws.send_json({"type": "session.config", "voice": "Vivian"})
ws.send_json({"type": "input.text", "text": "Hello world. "})
ws.send_json({"type": "input.done"})
assert ws.receive_json()["type"] == "audio.start"
assert ws.receive_json() == {"type": "error", "message": "Generation failed for sentence 0: boom"}
assert ws.receive_json() == {"type": "audio.done", "sentence_index": 0, "total_bytes": 0, "error": True}
ws.send_json({"type": "input.done"})
assert ws.receive_json() == {"type": "session.done", "total_sentences": 1}
def test_streaming_generation_error_marks_audio_done(self, mocker: MockerFixture):
@@ -453,6 +447,7 @@ class TestStreamingSpeechWebSocket:
}
)
ws.send_json({"type": "input.text", "text": "Hello world. "})
ws.send_json({"type": "input.done"})
assert ws.receive_json()["type"] == "audio.start"
assert ws.receive_bytes() == b"\x01\x02"
@@ -467,7 +462,6 @@ class TestStreamingSpeechWebSocket:
"error": True,
}
ws.send_json({"type": "input.done"})
assert ws.receive_json() == {"type": "session.done", "total_sentences": 1}
def test_invalid_input_text_type_returns_validation_error(self, mocker: MockerFixture):
@@ -1,253 +0,0 @@
"""Tests for SentenceSplitter used in streaming TTS input."""
import pytest
from vllm_omni.entrypoints.openai.text_splitter import SentenceSplitter
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
class TestSentenceSplitterEnglish:
"""Tests for English sentence splitting."""
def test_single_sentence_no_boundary(self):
splitter = SentenceSplitter()
result = splitter.add_text("Hello world")
assert result == []
assert splitter.buffer == "Hello world"
def test_single_sentence_with_boundary(self):
splitter = SentenceSplitter()
result = splitter.add_text("Hello world. How are you?")
assert len(result) == 1
assert result[0] == "Hello world."
def test_multiple_sentences(self):
splitter = SentenceSplitter()
result = splitter.add_text("Hello. How are you? I am fine! ")
assert len(result) == 3
assert result[0] == "Hello."
assert result[1] == "How are you?"
assert result[2] == "I am fine!"
def test_exclamation_mark(self):
splitter = SentenceSplitter()
result = splitter.add_text("Wow, that is great! Tell me more.")
assert len(result) == 1
assert result[0] == "Wow, that is great!"
def test_question_mark(self):
splitter = SentenceSplitter()
result = splitter.add_text("Can you hear me? I hope so.")
assert len(result) == 1
assert result[0] == "Can you hear me?"
class TestSentenceSplitterChinese:
"""Tests for CJK sentence splitting."""
def test_chinese_period(self):
splitter = SentenceSplitter()
result = splitter.add_text("你好世界。你好吗")
assert len(result) == 1
assert result[0] == "你好世界。"
def test_chinese_exclamation(self):
splitter = SentenceSplitter()
result = splitter.add_text("太好了!谢谢你")
assert len(result) == 1
assert result[0] == "太好了!"
def test_chinese_question(self):
splitter = SentenceSplitter()
result = splitter.add_text("你是谁?我是小明")
assert len(result) == 1
assert result[0] == "你是谁?"
def test_chinese_comma_no_split(self):
"""Chinese commas are clause-level and should not trigger a split."""
splitter = SentenceSplitter()
result = splitter.add_text("你好,世界")
assert result == []
assert splitter.buffer == "你好,世界"
def test_chinese_semicolon_no_split(self):
"""Chinese semicolons are clause-level and should not trigger a split."""
splitter = SentenceSplitter()
result = splitter.add_text("第一点;第二点")
assert result == []
assert splitter.buffer == "第一点;第二点"
def test_chinese_multiple(self):
splitter = SentenceSplitter()
result = splitter.add_text("你好!你好吗?我很好。")
assert len(result) == 3
assert result[0] == "你好!"
assert result[1] == "你好吗?"
assert result[2] == "我很好。"
class TestSentenceSplitterMixed:
"""Tests for mixed-language sentence splitting."""
def test_mixed_english_chinese(self):
splitter = SentenceSplitter()
result = splitter.add_text("Hello世界。How are you? ")
assert len(result) == 2
assert result[0] == "Hello世界。"
assert result[1] == "How are you?"
class TestSentenceSplitterIncremental:
"""Tests for incremental (multi-chunk) text input."""
def test_accumulation_across_chunks(self):
splitter = SentenceSplitter()
# First chunk: no boundary
result1 = splitter.add_text("Hello ")
assert result1 == []
# Second chunk: completes a sentence
result2 = splitter.add_text("world. How")
assert len(result2) == 1
assert result2[0] == "Hello world."
assert splitter.buffer == "How"
def test_word_by_word(self):
splitter = SentenceSplitter()
words = ["Hello, ", "how ", "are ", "you? ", "I ", "am ", "fine."]
all_sentences = []
for word in words:
all_sentences.extend(splitter.add_text(word))
assert len(all_sentences) == 1
assert all_sentences[0] == "Hello, how are you?"
# "I am fine." stays in buffer (no trailing whitespace after period)
def test_three_chunks(self):
splitter = SentenceSplitter()
splitter.add_text("The quick brown ")
splitter.add_text("fox jumps. ")
result = splitter.add_text("Over the lazy dog. ")
# "The quick brown fox jumps." should have been returned on second chunk
# "Over the lazy dog." on third chunk
assert len(result) == 1
assert result[0] == "Over the lazy dog."
class TestSentenceSplitterFlush:
"""Tests for flush behavior."""
def test_flush_returns_remaining(self):
splitter = SentenceSplitter()
splitter.add_text("Hello world")
result = splitter.flush()
assert result == "Hello world"
assert splitter.buffer == ""
def test_flush_empty_buffer(self):
splitter = SentenceSplitter()
result = splitter.flush()
assert result is None
def test_flush_after_sentence(self):
splitter = SentenceSplitter()
splitter.add_text("Hello world. Remaining text")
result = splitter.flush()
assert result == "Remaining text"
def test_flush_whitespace_only(self):
splitter = SentenceSplitter()
splitter.add_text("Hello. ")
# "Hello." extracted, buffer is " "
result = splitter.flush()
# Whitespace-only should return None
assert result is None
def test_flush_clears_buffer(self):
splitter = SentenceSplitter()
splitter.add_text("some text")
splitter.flush()
assert splitter.buffer == ""
# Second flush should return None
assert splitter.flush() is None
class TestSentenceSplitterEdgeCases:
"""Edge case tests."""
def test_empty_input(self):
splitter = SentenceSplitter()
result = splitter.add_text("")
assert result == []
assert splitter.buffer == ""
def test_none_like_empty(self):
"""Empty string should not affect buffer."""
splitter = SentenceSplitter()
splitter.add_text("Hello")
splitter.add_text("")
assert splitter.buffer == "Hello"
def test_only_punctuation(self):
splitter = SentenceSplitter()
result = splitter.add_text(". ")
# "." is 1 char, below default min_sentence_length of 2
# It will be carried forward
assert result == []
def test_min_sentence_length(self):
splitter = SentenceSplitter(min_sentence_length=10)
result = splitter.add_text("Hi. Hello world. ")
# "Hi." is 3 chars (< 10), so it gets carried to "Hello world."
assert len(result) == 1
assert "Hi." in result[0]
assert "Hello world." in result[0]
def test_short_segments_are_carried_until_long_enough(self):
splitter = SentenceSplitter(min_sentence_length=10)
result = splitter.add_text("Hi. Ok. Hello there. ")
assert result == ["Hi.Ok.Hello there."]
assert splitter.buffer == ""
def test_min_sentence_length_zero(self):
splitter = SentenceSplitter(min_sentence_length=0)
result = splitter.add_text("A. B. ")
assert len(result) == 2
def test_no_boundary_then_flush(self):
splitter = SentenceSplitter()
result = splitter.add_text("Hello world how are you")
assert result == []
flushed = splitter.flush()
assert flushed == "Hello world how are you"
def test_consecutive_punctuation(self):
splitter = SentenceSplitter()
result = splitter.add_text("Really?! Yes, really. ")
assert len(result) >= 1
def test_reuse_after_flush(self):
"""Splitter can be reused after flush."""
splitter = SentenceSplitter()
splitter.add_text("First session.")
splitter.flush()
result = splitter.add_text("Second session. More text")
assert len(result) == 1
assert result[0] == "Second session."
assert splitter.buffer == "More text"
class TestSentenceSplitterBufferLimit:
"""Tests for buffer overflow protection."""
def test_buffer_overflow_raises(self):
from vllm_omni.entrypoints.openai.text_splitter import _MAX_BUFFER_SIZE
splitter = SentenceSplitter()
# Fill buffer just under the limit
splitter.add_text("x" * (_MAX_BUFFER_SIZE - 1))
# One more char should exceed the limit
with pytest.raises(ValueError, match="exceeded maximum size"):
splitter.add_text("xx")
@@ -476,13 +476,6 @@ class StreamingSpeechSessionConfig(BaseModel):
"Requires response_format='pcm'. Speed adjustment is not supported when streaming."
),
)
split_granularity: Literal["sentence", "clause"] = Field(
default="sentence",
description=(
"Text splitting granularity: 'sentence' splits on .!?。!?, "
"'clause' also splits on CJK commas , and semicolons ;."
),
)
word_timestamps: bool = Field(
default=False,
description=(
@@ -1,7 +1,7 @@
"""WebSocket handler for streaming text input TTS.
Accepts text incrementally via WebSocket, buffers and splits at sentence
boundaries, and generates audio per sentence using the existing TTS pipeline.
Accepts text incrementally via WebSocket, buffers it until input.done, and
generates audio once for the buffered input using the existing TTS pipeline.
Protocol:
Client -> Server:
@@ -41,11 +41,6 @@ from vllm_omni.entrypoints.openai.protocol.audio import (
StreamingSpeechSessionConfig,
)
from vllm_omni.entrypoints.openai.serving_speech import OmniOpenAIServingSpeech
from vllm_omni.entrypoints.openai.text_splitter import (
SPLIT_CLAUSE,
SPLIT_SENTENCE,
SentenceSplitter,
)
from vllm_omni.utils.forced_aligner import ForcedAlignerLoadError
from vllm_omni.utils.forced_aligner import align as forced_align
@@ -63,8 +58,8 @@ class OmniStreamingSpeechHandler:
"""Handles WebSocket sessions for streaming text-input TTS.
Each WebSocket connection is an independent session. Text arrives
incrementally, is split at sentence boundaries, and audio is generated
per sentence using the existing OmniOpenAIServingSpeech pipeline.
incrementally, is buffered until input.done, and audio is generated once
for the buffered input using the existing OmniOpenAIServingSpeech pipeline.
Args:
speech_service: The existing TTS serving instance (reused for
@@ -102,9 +97,7 @@ class OmniStreamingSpeechHandler:
await self._send_error(websocket, str(error))
return
boundary_re = SPLIT_CLAUSE if config.split_granularity == "clause" else SPLIT_SENTENCE
splitter = SentenceSplitter(boundary_re=boundary_re)
sentence_index = 0
text_parts: list[str] = []
# 2. Receive text chunks until input.done
while True:
@@ -138,23 +131,19 @@ class OmniStreamingSpeechHandler:
if not isinstance(text, str):
await self._send_error(websocket, "input.text requires a string value")
continue
sentences = splitter.add_text(text)
for sentence in sentences:
await self._generate_and_send(websocket, config, sentence, sentence_index)
sentence_index += 1
text_parts.append(text)
elif msg_type == "input.done":
# Flush remaining buffer
remaining = splitter.flush()
if remaining:
await self._generate_and_send(websocket, config, remaining, sentence_index)
sentence_index += 1
full_text = "".join(text_parts).strip()
total_sentences = 0
if full_text:
await self._generate_and_send(websocket, config, full_text, 0)
total_sentences = 1
# Send session.done
await websocket.send_json(
{
"type": "session.done",
"total_sentences": sentence_index,
"total_sentences": total_sentences,
}
)
return
@@ -1,120 +0,0 @@
"""Multi-language sentence boundary detector for streaming TTS input.
Buffers incoming text and splits at sentence boundaries (English and CJK),
yielding complete sentences for audio generation.
"""
import re
from re import Pattern
# Maximum buffer size (in characters) to prevent unbounded memory growth.
_MAX_BUFFER_SIZE = 100_000 # ~100 KB of text
# Sentence-level: .!? + CJK sentence-ending 。!?
# NOTE: English requires trailing whitespace to confirm a boundary —
# end-of-string is NOT treated as a boundary (that is what flush() is for).
SPLIT_SENTENCE = re.compile(
r"(?<=[.!?])\s+"
r"|(?<=[。!?])"
)
# Clause-level: adds CJK commas , and semicolons ;
SPLIT_CLAUSE = re.compile(
r"(?<=[.!?])\s+"
r"|(?<=[。!?,;])"
)
# Default alias
_SENTENCE_BOUNDARY_RE = SPLIT_SENTENCE
class SentenceSplitter:
"""Incremental sentence splitter for streaming text input.
Buffers text and yields complete sentences when boundaries are detected.
Designed for TTS pipelines where text arrives incrementally (e.g., from STT).
Args:
min_sentence_length: Minimum character length for a sentence.
Sentences shorter than this are kept in the buffer to avoid
splitting on abbreviations like "Dr." or "U.S.".
boundary_re: Custom compiled regex for sentence boundaries.
Use ``SPLIT_SENTENCE`` (default) for sentence-level splitting,
``SPLIT_CLAUSE`` for finer-grained clause-level splitting,
or pass your own ``re.Pattern``.
"""
def __init__(
self,
min_sentence_length: int = 2,
boundary_re: Pattern[str] | None = None,
) -> None:
self._buffer: str = ""
self._min_sentence_length = min_sentence_length
self._boundary_re = boundary_re or _SENTENCE_BOUNDARY_RE
@property
def buffer(self) -> str:
"""Current buffered text."""
return self._buffer
def add_text(self, text: str) -> list[str]:
"""Add text to the buffer and return any complete sentences.
Args:
text: Incoming text chunk.
Returns:
List of complete sentences extracted from the buffer.
May be empty if no sentence boundary was found.
Raises:
ValueError: If the buffer exceeds the maximum size.
"""
if not text:
return []
self._buffer += text
if len(self._buffer) > _MAX_BUFFER_SIZE:
raise ValueError(
f"Text buffer exceeded maximum size ({_MAX_BUFFER_SIZE} chars). "
"Consider adding sentence-ending punctuation to your input."
)
return self._extract_sentences()
def flush(self) -> str | None:
"""Flush remaining buffered text as a final sentence.
Returns:
The remaining buffered text (stripped), or None if buffer is empty.
"""
remaining = self._buffer.strip()
self._buffer = ""
return remaining if remaining else None
def _extract_sentences(self) -> list[str]:
"""Split buffer at sentence boundaries, keeping incomplete text buffered."""
parts = self._boundary_re.split(self._buffer)
if len(parts) <= 1:
# No boundary found — keep everything in buffer
return []
sentences: list[str] = []
carry = ""
# All parts except the last are complete sentences
for i in range(len(parts) - 1):
text = carry + parts[i]
carry = ""
stripped = text.strip()
if len(stripped) >= self._min_sentence_length:
sentences.append(stripped)
elif stripped:
# Too short (e.g. "Dr.") — carry forward to next part
carry = text
# else: empty, skip
# Last part stays in buffer (may be incomplete)
self._buffer = carry + parts[-1]
return sentences