[Feature] Flush the speech WebSocket on input.done without closing it (#5517)
Signed-off-by: iancarrasco-b10 <ian.carrasco@baseten.co>
This commit is contained in:
1 parent
0feb92bdfa
commit
23bb4395bb
8 files changed
+688
-302
No files matched your search
@@ -276,20 +276,42 @@ Client -> Server:
|
||||
|
||||
| Message | Description |
|
||||
|---------|-------------|
|
||||
| `{"type": "session.config", ...}` | Session configuration (sent once, first message) |
|
||||
| `{"type": "session.config", ...}` | Session configuration (first message; may be resent between utterances to change it) |
|
||||
| `{"type": "input.text", "text": "..."}` | Text chunk |
|
||||
| `{"type": "input.done"}` | End of input, flushes remaining buffer |
|
||||
| `{"type": "input.done"}` | End of utterance: flushes the buffer and keeps the connection open |
|
||||
| `{"type": "session.close"}` | End of connection |
|
||||
|
||||
Server -> Client:
|
||||
|
||||
| Message | Description |
|
||||
|---------|-------------|
|
||||
| `{"type": "audio.start", "sentence_index": 0, "sentence_text": "...", "format": "pcm", "sample_rate": 24000}` | Audio generation starting for the buffered input |
|
||||
| `{"type": "audio.start", "utterance_index": 0, "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 the buffered input |
|
||||
| `{"type": "session.done", "total_sentences": N}` | Session complete |
|
||||
| `{"type": "audio.done", "utterance_index": 0, "sentence_index": 0, "total_bytes": 96000, "error": false}` | Audio complete for the buffered input |
|
||||
| `{"type": "session.done", "utterance_index": 0, "total_sentences": N}` | Flushed utterance complete |
|
||||
| `{"type": "error", "message": "..."}` | Non-fatal error |
|
||||
|
||||
### Flushing vs. Closing
|
||||
|
||||
`input.done` is a flush, not a disconnect. The server synthesizes the buffered
|
||||
text, emits `session.done`, and then waits on the same connection for the next
|
||||
utterance, so a client that speaks repeatedly (for example one driven by an
|
||||
upstream LLM) pays the WebSocket handshake once instead of once per utterance.
|
||||
|
||||
* The session config is sticky. Send `input.text` again straight after
|
||||
`session.done` to reuse it, or send another `session.config` first to change
|
||||
voice, format, or reference audio. A `session.config` sent while text is
|
||||
still buffered is rejected so no pending input is silently dropped.
|
||||
* An utterance is the flush unit, not a linguistic one: it is whatever text was
|
||||
buffered when `input.done` arrived, of any length, synthesized as one request.
|
||||
`utterance_index` counts those flushes across the connection, so it tells you
|
||||
which `input.done` a frame belongs to. `sentence_index` counts within one
|
||||
flush and so pairs with `total_sentences`, which means every utterance reports
|
||||
`sentence_index: 0` of `total_sentences: 1` (or `0` for an empty buffer).
|
||||
* End the connection with `session.close`, or by closing the socket. An idle
|
||||
connection is still closed after the server's idle timeout, which now also
|
||||
applies to the gap between utterances.
|
||||
|
||||
### Session Config Parameters
|
||||
|
||||
All REST API parameters are supported, plus:
|
||||
|
||||
@@ -422,12 +422,19 @@ curl -X POST http://localhost:8091/v1/audio/speech \
|
||||
Raw PCM streaming requires `stream_format="audio"`, `response_format="pcm"`, and `async_chunk: true` on the stage config (default in `qwen3_tts.yaml`). `speed` is not supported when streaming.
|
||||
|
||||
### Streaming WebSocket
|
||||
The `/v1/audio/speech/stream` endpoint accepts text incrementally, splits it at sentence boundaries, and emits one PCM stream per sentence:
|
||||
The `/v1/audio/speech/stream` endpoint accepts text incrementally and synthesizes the buffered text as one continuous request on `input.done`:
|
||||
```bash
|
||||
python qwen3_tts/streaming_speech_client.py --text "Hello world. How are you? I am fine."
|
||||
python qwen3_tts/streaming_speech_client.py --text "..." --simulate-stt --stt-delay 0.1
|
||||
```
|
||||
|
||||
`input.done` flushes without closing, so repeating `--text` synthesizes several utterances over one connection:
|
||||
```bash
|
||||
python qwen3_tts/streaming_speech_client.py \
|
||||
--text "First utterance." \
|
||||
--text "Second utterance, same connection."
|
||||
```
|
||||
|
||||
### Gradio demos
|
||||
```bash
|
||||
./qwen3_tts/run_gradio_demo.sh # CustomVoice (default)
|
||||
|
||||
@@ -537,12 +537,19 @@ curl -X POST http://localhost:8091/v1/audio/speech \
|
||||
Raw PCM streaming requires `stream_format="audio"`, `response_format="pcm"`, and `async_chunk: true` on the stage config (default in `qwen3_tts.yaml`). `speed` is not supported when streaming.
|
||||
|
||||
### Streaming WebSocket
|
||||
The `/v1/audio/speech/stream` endpoint accepts text incrementally, splits it at sentence boundaries, and emits one PCM stream per sentence:
|
||||
The `/v1/audio/speech/stream` endpoint accepts text incrementally and synthesizes the buffered text as one continuous request on `input.done`:
|
||||
```bash
|
||||
python qwen3_tts/streaming_speech_client.py --text "Hello world. How are you? I am fine."
|
||||
python qwen3_tts/streaming_speech_client.py --text "..." --simulate-stt --stt-delay 0.1
|
||||
```
|
||||
|
||||
`input.done` flushes without closing, so repeating `--text` synthesizes several utterances over one connection:
|
||||
```bash
|
||||
python qwen3_tts/streaming_speech_client.py \
|
||||
--text "First utterance." \
|
||||
--text "Second utterance, same connection."
|
||||
```
|
||||
|
||||
To receive word-level timestamps, launch the server with a forced aligner:
|
||||
```bash
|
||||
vllm-omni serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \
|
||||
|
||||
@@ -7,6 +7,12 @@ Usage:
|
||||
# Send full text at once
|
||||
python streaming_speech_client.py --text "Hello world. How are you? I am fine."
|
||||
|
||||
# Several utterances over one connection: input.done flushes each of them
|
||||
# and the connection stays open, so only the first pays the handshake
|
||||
python streaming_speech_client.py \
|
||||
--text "First utterance." \
|
||||
--text "Second utterance, same connection."
|
||||
|
||||
# Pick a built-in speaker for CustomVoice models
|
||||
python streaming_speech_client.py \
|
||||
--text "打开电子枪,关闭扫描,切换到低倍模式。" \
|
||||
@@ -64,18 +70,27 @@ except ImportError:
|
||||
raise SystemExit(1)
|
||||
|
||||
|
||||
def save_audio_file(output_dir: str, sentence_index: int, response_format: str, chunks: list[bytes]) -> str:
|
||||
def frame_basename(msg: dict) -> str:
|
||||
"""Name a file after the utterance and sentence a frame belongs to.
|
||||
|
||||
Every utterance is sentence 0, so the utterance index is what keeps the
|
||||
files of one connection from overwriting each other.
|
||||
"""
|
||||
return f"utterance_{msg['utterance_index']:03d}_sentence_{msg['sentence_index']:03d}"
|
||||
|
||||
|
||||
def save_audio_file(output_dir: str, basename: str, response_format: str, chunks: list[bytes]) -> str:
|
||||
audio_bytes = b"".join(chunks)
|
||||
if response_format == "pcm":
|
||||
raw_path = os.path.join(output_dir, f"sentence_{sentence_index:03d}.pcm")
|
||||
raw_path = os.path.join(output_dir, f"{basename}.pcm")
|
||||
with open(raw_path, "wb") as f:
|
||||
f.write(audio_bytes)
|
||||
|
||||
wav_path = os.path.join(output_dir, f"sentence_{sentence_index:03d}.wav")
|
||||
wav_path = os.path.join(output_dir, f"{basename}.wav")
|
||||
save_pcm_wav(wav_path, audio_bytes)
|
||||
return wav_path
|
||||
|
||||
filename = os.path.join(output_dir, f"sentence_{sentence_index:03d}.{response_format}")
|
||||
filename = os.path.join(output_dir, f"{basename}.{response_format}")
|
||||
with open(filename, "wb") as f:
|
||||
f.write(audio_bytes)
|
||||
return filename
|
||||
@@ -95,183 +110,207 @@ def timestamp_words_text(timestamps: list[dict] | None) -> str:
|
||||
return " ".join(str(item.get("word", "")) for item in timestamps if item.get("word"))
|
||||
|
||||
|
||||
async def send_utterance(ws, text: str, simulate_stt: bool, stt_delay: float) -> None:
|
||||
"""Send one utterance's text, then input.done to flush it."""
|
||||
if simulate_stt:
|
||||
words = text.split(" ")
|
||||
for i, word in enumerate(words):
|
||||
chunk = word + (" " if i < len(words) - 1 else "")
|
||||
await ws.send(
|
||||
json.dumps(
|
||||
{
|
||||
"type": "input.text",
|
||||
"text": chunk,
|
||||
}
|
||||
)
|
||||
)
|
||||
print(f" Sent: {chunk!r}")
|
||||
await asyncio.sleep(stt_delay)
|
||||
else:
|
||||
await ws.send(
|
||||
json.dumps(
|
||||
{
|
||||
"type": "input.text",
|
||||
"text": text,
|
||||
}
|
||||
)
|
||||
)
|
||||
print(f"Sent full text: {text!r}")
|
||||
|
||||
await ws.send(json.dumps({"type": "input.done"}))
|
||||
print("Sent input.done")
|
||||
|
||||
|
||||
async def receive_utterance(
|
||||
ws,
|
||||
output_dir: str,
|
||||
response_format: str,
|
||||
word_timestamps: bool,
|
||||
save_chunks: bool,
|
||||
) -> None:
|
||||
"""Consume frames until session.done marks the flushed utterance complete."""
|
||||
current_label = "utterance 0"
|
||||
current_sentence_text = ""
|
||||
current_chunks: list[bytes] = []
|
||||
current_timestamps: list[dict] = []
|
||||
# Distinguish `null` (aligner failed) from `[]` (silence / no tokens):
|
||||
# the sentence is a failure only if no frame ever carried timestamps.
|
||||
alignment_frame_seen = False
|
||||
|
||||
while True:
|
||||
message = await ws.recv()
|
||||
|
||||
if isinstance(message, bytes):
|
||||
current_chunks.append(message)
|
||||
print(f" Received audio chunk for {current_label}: {len(message)} bytes")
|
||||
else:
|
||||
# JSON frame
|
||||
msg = json.loads(message)
|
||||
msg_type = msg.get("type")
|
||||
|
||||
if msg_type == "audio.start":
|
||||
current_label = f"utterance {msg['utterance_index']}"
|
||||
current_sentence_text = msg.get("sentence_text", "")
|
||||
current_chunks = []
|
||||
current_timestamps = []
|
||||
alignment_frame_seen = False
|
||||
print(f" [{current_label}] Generating: {current_sentence_text!r}")
|
||||
elif msg_type == "audio.chunk":
|
||||
audio = base64.b64decode(msg["audio_b64"])
|
||||
current_chunks.append(audio)
|
||||
chunk_timestamps = msg.get("timestamps")
|
||||
# Sentence-level alignment sends a trailing timestamp-only
|
||||
# frame with empty audio; don't write an empty chunk file.
|
||||
if save_chunks and audio:
|
||||
chunk_base = os.path.join(
|
||||
output_dir,
|
||||
f"{frame_basename(msg)}_chunk_{msg['chunk_id']:03d}",
|
||||
)
|
||||
chunk_pcm = f"{chunk_base}.pcm"
|
||||
chunk_wav = f"{chunk_base}.wav"
|
||||
chunk_json = f"{chunk_base}_timestamps.json"
|
||||
with open(chunk_pcm, "wb") as f:
|
||||
f.write(audio)
|
||||
save_pcm_wav(chunk_wav, audio, int(msg.get("sample_rate", 24000)))
|
||||
with open(chunk_json, "w", encoding="utf-8") as f:
|
||||
json.dump(
|
||||
{
|
||||
"utterance_index": msg["utterance_index"],
|
||||
"sentence_index": msg["sentence_index"],
|
||||
"chunk_id": msg["chunk_id"],
|
||||
"chunk_start_ms": msg.get("chunk_start_ms"),
|
||||
"chunk_end_ms": msg.get("chunk_end_ms"),
|
||||
"sentence_text": current_sentence_text,
|
||||
"covered_text": timestamp_words_text(chunk_timestamps),
|
||||
"timestamps": chunk_timestamps,
|
||||
},
|
||||
f,
|
||||
ensure_ascii=False,
|
||||
indent=2,
|
||||
)
|
||||
if chunk_timestamps is None:
|
||||
print(f" [{current_label}] chunk {msg['chunk_id']}: {len(audio)} bytes, timestamps unavailable")
|
||||
else:
|
||||
alignment_frame_seen = True
|
||||
current_timestamps.extend(chunk_timestamps)
|
||||
print(
|
||||
f" [{current_label}] chunk {msg['chunk_id']}: "
|
||||
f"{len(audio)} bytes, {len(chunk_timestamps)} timestamp(s)"
|
||||
)
|
||||
if save_chunks and audio:
|
||||
print(f" saved chunk -> {chunk_wav}")
|
||||
elif msg_type == "audio.done":
|
||||
filename = save_audio_file(
|
||||
output_dir,
|
||||
frame_basename(msg),
|
||||
response_format,
|
||||
current_chunks,
|
||||
)
|
||||
if word_timestamps:
|
||||
# `null` on failure, else the accumulated list (`[]` for silence).
|
||||
sentence_timestamps = current_timestamps if alignment_frame_seen else None
|
||||
ts_filename = os.path.join(
|
||||
output_dir,
|
||||
f"{frame_basename(msg)}_timestamps.json",
|
||||
)
|
||||
with open(ts_filename, "w", encoding="utf-8") as f:
|
||||
json.dump(
|
||||
{
|
||||
"utterance_index": msg["utterance_index"],
|
||||
"sentence_index": msg["sentence_index"],
|
||||
"sentence_text": current_sentence_text,
|
||||
"covered_text": timestamp_words_text(sentence_timestamps),
|
||||
"timestamps": sentence_timestamps,
|
||||
},
|
||||
f,
|
||||
ensure_ascii=False,
|
||||
indent=2,
|
||||
)
|
||||
print(
|
||||
f" [{current_label}] Done"
|
||||
f" bytes={msg.get('total_bytes', len(b''.join(current_chunks)))}"
|
||||
f" error={msg.get('error', False)}"
|
||||
f" -> {filename}"
|
||||
)
|
||||
if word_timestamps:
|
||||
print(f" [{current_label}] Timestamps -> {ts_filename}")
|
||||
current_chunks = []
|
||||
current_timestamps = []
|
||||
alignment_frame_seen = False
|
||||
elif msg_type == "session.done":
|
||||
print(f"\nUtterance {msg['utterance_index']} complete: {msg['total_sentences']} sentence(s) generated")
|
||||
return
|
||||
elif msg_type == "error":
|
||||
print(f" ERROR: {msg['message']}")
|
||||
else:
|
||||
print(f" Unknown message: {msg}")
|
||||
|
||||
|
||||
async def stream_tts(
|
||||
url: str,
|
||||
text: str,
|
||||
texts: list[str],
|
||||
config: dict,
|
||||
output_dir: str,
|
||||
simulate_stt: bool = False,
|
||||
stt_delay: float = 0.1,
|
||||
save_chunks: bool = False,
|
||||
) -> None:
|
||||
"""Connect to the streaming TTS endpoint and process audio responses."""
|
||||
"""Connect to the streaming TTS endpoint and process audio responses.
|
||||
|
||||
Every text is synthesized over the same connection: input.done flushes an
|
||||
utterance without closing, so only the first one pays the handshake.
|
||||
"""
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
response_format = config.get("response_format", "wav")
|
||||
word_timestamps = bool(config.get("word_timestamps", False))
|
||||
|
||||
async with websockets.connect(url) as ws:
|
||||
# 1. Send session config
|
||||
# 1. Send session config; it stays in effect for every utterance below.
|
||||
config_msg = {"type": "session.config", **config}
|
||||
await ws.send(json.dumps(config_msg))
|
||||
print(f"Sent session config: {config}")
|
||||
|
||||
# 2. Send text (either all at once or word-by-word)
|
||||
async def send_text():
|
||||
if simulate_stt:
|
||||
words = text.split(" ")
|
||||
for i, word in enumerate(words):
|
||||
chunk = word + (" " if i < len(words) - 1 else "")
|
||||
await ws.send(
|
||||
json.dumps(
|
||||
{
|
||||
"type": "input.text",
|
||||
"text": chunk,
|
||||
}
|
||||
)
|
||||
)
|
||||
print(f" Sent: {chunk!r}")
|
||||
await asyncio.sleep(stt_delay)
|
||||
else:
|
||||
await ws.send(
|
||||
json.dumps(
|
||||
{
|
||||
"type": "input.text",
|
||||
"text": text,
|
||||
}
|
||||
)
|
||||
)
|
||||
print(f"Sent full text: {text!r}")
|
||||
|
||||
# 3. Signal end of input
|
||||
await ws.send(json.dumps({"type": "input.done"}))
|
||||
print("Sent input.done")
|
||||
|
||||
# Run sender and receiver concurrently
|
||||
sender_task = asyncio.create_task(send_text())
|
||||
|
||||
response_format = config.get("response_format", "wav")
|
||||
word_timestamps = bool(config.get("word_timestamps", False))
|
||||
current_sentence_index = 0
|
||||
current_sentence_text = ""
|
||||
current_chunks: list[bytes] = []
|
||||
current_timestamps: list[dict] = []
|
||||
# Distinguish `null` (aligner failed) from `[]` (silence / no tokens):
|
||||
# the sentence is a failure only if no frame ever carried timestamps.
|
||||
alignment_frame_seen = False
|
||||
|
||||
try:
|
||||
while True:
|
||||
message = await ws.recv()
|
||||
|
||||
if isinstance(message, bytes):
|
||||
current_chunks.append(message)
|
||||
print(f" Received audio chunk for sentence {current_sentence_index}: {len(message)} bytes")
|
||||
else:
|
||||
# JSON frame
|
||||
msg = json.loads(message)
|
||||
msg_type = msg.get("type")
|
||||
|
||||
if msg_type == "audio.start":
|
||||
current_sentence_index = msg["sentence_index"]
|
||||
current_sentence_text = msg.get("sentence_text", "")
|
||||
current_chunks = []
|
||||
current_timestamps = []
|
||||
alignment_frame_seen = False
|
||||
print(f" [sentence {msg['sentence_index']}] Generating: {current_sentence_text!r}")
|
||||
elif msg_type == "audio.chunk":
|
||||
audio = base64.b64decode(msg["audio_b64"])
|
||||
current_chunks.append(audio)
|
||||
chunk_timestamps = msg.get("timestamps")
|
||||
# Sentence-level alignment sends a trailing timestamp-only
|
||||
# frame with empty audio; don't write an empty chunk file.
|
||||
if save_chunks and audio:
|
||||
chunk_base = os.path.join(
|
||||
output_dir,
|
||||
f"sentence_{msg['sentence_index']:03d}_chunk_{msg['chunk_id']:03d}",
|
||||
)
|
||||
chunk_pcm = f"{chunk_base}.pcm"
|
||||
chunk_wav = f"{chunk_base}.wav"
|
||||
chunk_json = f"{chunk_base}_timestamps.json"
|
||||
with open(chunk_pcm, "wb") as f:
|
||||
f.write(audio)
|
||||
save_pcm_wav(chunk_wav, audio, int(msg.get("sample_rate", 24000)))
|
||||
with open(chunk_json, "w", encoding="utf-8") as f:
|
||||
json.dump(
|
||||
{
|
||||
"sentence_index": msg["sentence_index"],
|
||||
"chunk_id": msg["chunk_id"],
|
||||
"chunk_start_ms": msg.get("chunk_start_ms"),
|
||||
"chunk_end_ms": msg.get("chunk_end_ms"),
|
||||
"sentence_text": current_sentence_text,
|
||||
"covered_text": timestamp_words_text(chunk_timestamps),
|
||||
"timestamps": chunk_timestamps,
|
||||
},
|
||||
f,
|
||||
ensure_ascii=False,
|
||||
indent=2,
|
||||
)
|
||||
if chunk_timestamps is None:
|
||||
print(
|
||||
f" [sentence {msg['sentence_index']}] chunk {msg['chunk_id']}: "
|
||||
f"{len(audio)} bytes, timestamps unavailable"
|
||||
)
|
||||
else:
|
||||
alignment_frame_seen = True
|
||||
current_timestamps.extend(chunk_timestamps)
|
||||
print(
|
||||
f" [sentence {msg['sentence_index']}] chunk {msg['chunk_id']}: "
|
||||
f"{len(audio)} bytes, {len(chunk_timestamps)} timestamp(s)"
|
||||
)
|
||||
if save_chunks and audio:
|
||||
print(f" saved chunk -> {chunk_wav}")
|
||||
elif msg_type == "audio.done":
|
||||
filename = save_audio_file(
|
||||
output_dir,
|
||||
msg["sentence_index"],
|
||||
response_format,
|
||||
current_chunks,
|
||||
)
|
||||
if word_timestamps:
|
||||
# `null` on failure, else the accumulated list (`[]` for silence).
|
||||
sentence_timestamps = current_timestamps if alignment_frame_seen else None
|
||||
ts_filename = os.path.join(
|
||||
output_dir,
|
||||
f"sentence_{msg['sentence_index']:03d}_timestamps.json",
|
||||
)
|
||||
with open(ts_filename, "w", encoding="utf-8") as f:
|
||||
json.dump(
|
||||
{
|
||||
"sentence_index": msg["sentence_index"],
|
||||
"sentence_text": current_sentence_text,
|
||||
"covered_text": timestamp_words_text(sentence_timestamps),
|
||||
"timestamps": sentence_timestamps,
|
||||
},
|
||||
f,
|
||||
ensure_ascii=False,
|
||||
indent=2,
|
||||
)
|
||||
print(
|
||||
f" [sentence {msg['sentence_index']}] Done"
|
||||
f" bytes={msg.get('total_bytes', len(b''.join(current_chunks)))}"
|
||||
f" error={msg.get('error', False)}"
|
||||
f" -> {filename}"
|
||||
)
|
||||
if word_timestamps:
|
||||
print(f" [sentence {msg['sentence_index']}] Timestamps -> {ts_filename}")
|
||||
current_chunks = []
|
||||
current_timestamps = []
|
||||
alignment_frame_seen = False
|
||||
elif msg_type == "session.done":
|
||||
print(f"\nSession complete: {msg['total_sentences']} sentence(s) generated")
|
||||
break
|
||||
elif msg_type == "error":
|
||||
print(f" ERROR: {msg['message']}")
|
||||
else:
|
||||
print(f" Unknown message: {msg}")
|
||||
finally:
|
||||
sender_task.cancel()
|
||||
for text in texts:
|
||||
# 2. Send text (either all at once or word-by-word), then flush.
|
||||
sender_task = asyncio.create_task(send_utterance(ws, text, simulate_stt, stt_delay))
|
||||
try:
|
||||
await sender_task
|
||||
except asyncio.CancelledError:
|
||||
pass # Task cancellation is expected during shutdown
|
||||
await receive_utterance(
|
||||
ws,
|
||||
output_dir=output_dir,
|
||||
response_format=response_format,
|
||||
word_timestamps=word_timestamps,
|
||||
save_chunks=save_chunks,
|
||||
)
|
||||
finally:
|
||||
sender_task.cancel()
|
||||
try:
|
||||
await sender_task
|
||||
except asyncio.CancelledError:
|
||||
pass # Task cancellation is expected during shutdown
|
||||
|
||||
# 3. Release the connection once there is nothing left to say.
|
||||
await ws.send(json.dumps({"type": "session.close"}))
|
||||
print("Sent session.close")
|
||||
|
||||
print(f"\nAudio files saved to: {output_dir}/")
|
||||
|
||||
@@ -286,7 +325,8 @@ def main():
|
||||
parser.add_argument(
|
||||
"--text",
|
||||
required=True,
|
||||
help="Text to synthesize",
|
||||
action="append",
|
||||
help="Text to synthesize; repeat to send several utterances over one connection",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output-dir",
|
||||
@@ -385,7 +425,7 @@ def main():
|
||||
asyncio.run(
|
||||
stream_tts(
|
||||
url=args.url,
|
||||
text=args.text,
|
||||
texts=args.text,
|
||||
config=config,
|
||||
output_dir=args.output_dir,
|
||||
simulate_stt=args.simulate_stt,
|
||||
|
||||
@@ -33,26 +33,63 @@ tts_ws_server_params = [
|
||||
]
|
||||
|
||||
|
||||
async def _run_ws_session(host: str, port: int, model: str) -> dict:
|
||||
uri = f"ws://{host}:{port}/v1/audio/speech/stream"
|
||||
def _session_config(model: str) -> str:
|
||||
return json.dumps(
|
||||
{
|
||||
"type": "session.config",
|
||||
"model": model,
|
||||
"voice": "vivian",
|
||||
"language": "English",
|
||||
"response_format": "pcm",
|
||||
"stream_audio": True,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
async def _collect_utterance(ws) -> dict:
|
||||
"""Read frames until session.done closes out one flushed utterance."""
|
||||
starts: list[dict] = []
|
||||
dones: list[dict] = []
|
||||
chunk_lengths: dict[int, list[int]] = {}
|
||||
session_done: dict | None = None
|
||||
|
||||
while True:
|
||||
message = await asyncio.wait_for(ws.recv(), timeout=180)
|
||||
if isinstance(message, bytes):
|
||||
if not starts:
|
||||
raise AssertionError("Received audio bytes before audio.start")
|
||||
sentence_index = starts[-1]["sentence_index"]
|
||||
chunk_lengths.setdefault(sentence_index, []).append(len(message))
|
||||
continue
|
||||
|
||||
payload = json.loads(message)
|
||||
msg_type = payload.get("type")
|
||||
if msg_type == "audio.start":
|
||||
starts.append(payload)
|
||||
chunk_lengths.setdefault(payload["sentence_index"], [])
|
||||
elif msg_type == "audio.done":
|
||||
dones.append(payload)
|
||||
elif msg_type == "session.done":
|
||||
session_done = payload
|
||||
break
|
||||
elif msg_type == "error":
|
||||
raise AssertionError(f"WebSocket error: {payload['message']}")
|
||||
else:
|
||||
raise AssertionError(f"Unexpected WebSocket message: {payload}")
|
||||
|
||||
return {
|
||||
"starts": starts,
|
||||
"dones": dones,
|
||||
"chunk_lengths": chunk_lengths,
|
||||
"session_done": session_done,
|
||||
}
|
||||
|
||||
|
||||
async def _run_ws_session(host: str, port: int, model: str) -> dict:
|
||||
uri = f"ws://{host}:{port}/v1/audio/speech/stream"
|
||||
|
||||
async with websockets.connect(uri, max_size=None) as ws:
|
||||
await ws.send(
|
||||
json.dumps(
|
||||
{
|
||||
"type": "session.config",
|
||||
"model": model,
|
||||
"voice": "vivian",
|
||||
"language": "English",
|
||||
"response_format": "pcm",
|
||||
"stream_audio": True,
|
||||
}
|
||||
)
|
||||
)
|
||||
await ws.send(_session_config(model))
|
||||
await ws.send(
|
||||
json.dumps(
|
||||
{
|
||||
@@ -67,36 +104,23 @@ async def _run_ws_session(host: str, port: int, model: str) -> dict:
|
||||
)
|
||||
await ws.send(json.dumps({"type": "input.done"}))
|
||||
|
||||
while True:
|
||||
message = await asyncio.wait_for(ws.recv(), timeout=180)
|
||||
if isinstance(message, bytes):
|
||||
if not starts:
|
||||
raise AssertionError("Received audio bytes before audio.start")
|
||||
sentence_index = starts[-1]["sentence_index"]
|
||||
chunk_lengths.setdefault(sentence_index, []).append(len(message))
|
||||
continue
|
||||
return await _collect_utterance(ws)
|
||||
|
||||
payload = json.loads(message)
|
||||
msg_type = payload.get("type")
|
||||
if msg_type == "audio.start":
|
||||
starts.append(payload)
|
||||
chunk_lengths.setdefault(payload["sentence_index"], [])
|
||||
elif msg_type == "audio.done":
|
||||
dones.append(payload)
|
||||
elif msg_type == "session.done":
|
||||
session_done = payload
|
||||
break
|
||||
elif msg_type == "error":
|
||||
raise AssertionError(f"WebSocket error: {payload['message']}")
|
||||
else:
|
||||
raise AssertionError(f"Unexpected WebSocket message: {payload}")
|
||||
|
||||
return {
|
||||
"starts": starts,
|
||||
"dones": dones,
|
||||
"chunk_lengths": chunk_lengths,
|
||||
"session_done": session_done,
|
||||
}
|
||||
async def _run_ws_reused_connection(host: str, port: int, model: str, texts: list[str]) -> list[dict]:
|
||||
"""Synthesize several utterances over a single connection."""
|
||||
uri = f"ws://{host}:{port}/v1/audio/speech/stream"
|
||||
results: list[dict] = []
|
||||
|
||||
async with websockets.connect(uri, max_size=None) as ws:
|
||||
await ws.send(_session_config(model))
|
||||
for text in texts:
|
||||
await ws.send(json.dumps({"type": "input.text", "text": text}))
|
||||
await ws.send(json.dumps({"type": "input.done"}))
|
||||
results.append(await _collect_utterance(ws))
|
||||
await ws.send(json.dumps({"type": "session.close"}))
|
||||
|
||||
return results
|
||||
|
||||
|
||||
class TestQwen3TTSWebSocket:
|
||||
@@ -113,6 +137,7 @@ class TestQwen3TTSWebSocket:
|
||||
session_done = result["session_done"]
|
||||
|
||||
assert session_done is not None
|
||||
assert session_done["utterance_index"] == 0
|
||||
assert session_done["total_sentences"] == 1
|
||||
assert len(starts) == 1
|
||||
assert len(dones) == 1
|
||||
@@ -131,3 +156,32 @@ class TestQwen3TTSWebSocket:
|
||||
assert total_bytes > 0
|
||||
assert chunk_lengths[sentence_index], f"Expected binary PCM frames for sentence {sentence_index}"
|
||||
assert sum(chunk_lengths[sentence_index]) == total_bytes
|
||||
|
||||
@pytest.mark.advanced_model
|
||||
@pytest.mark.tts
|
||||
@hardware_test(res={"cuda": "L4"}, num_cards=1)
|
||||
@pytest.mark.parametrize("omni_server", tts_ws_server_params, indirect=True)
|
||||
def test_input_done_flushes_without_closing_connection(self, omni_server) -> None:
|
||||
# input.done must flush the buffer and leave the connection usable:
|
||||
# the second utterance is synthesized without a second handshake.
|
||||
texts = [
|
||||
"This is the first utterance sent over the websocket connection.",
|
||||
"This is the second utterance, reusing the very same connection.",
|
||||
]
|
||||
results = asyncio.run(_run_ws_reused_connection(omni_server.host, omni_server.port, omni_server.model, texts))
|
||||
|
||||
assert len(results) == len(texts)
|
||||
for expected_index, result in enumerate(results):
|
||||
assert result["session_done"] == {
|
||||
"type": "session.done",
|
||||
"utterance_index": expected_index,
|
||||
"total_sentences": 1,
|
||||
}
|
||||
assert len(result["starts"]) == 1
|
||||
assert len(result["dones"]) == 1
|
||||
# utterance_index counts the flushes of the connection, while
|
||||
# sentence_index stays within the flush it belongs to.
|
||||
assert result["starts"][0]["utterance_index"] == expected_index
|
||||
assert result["starts"][0]["sentence_index"] == 0
|
||||
assert result["dones"][0]["error"] is False
|
||||
assert result["dones"][0]["total_bytes"] > 0
|
||||
@@ -72,13 +72,162 @@ class TestStreamingSpeechWebSocket:
|
||||
assert audio.startswith(b"RIFF")
|
||||
|
||||
done = ws.receive_json()
|
||||
assert done == {"type": "audio.done", "sentence_index": 0, "total_bytes": len(audio), "error": False}
|
||||
assert done == {
|
||||
"type": "audio.done",
|
||||
"utterance_index": 0,
|
||||
"sentence_index": 0,
|
||||
"total_bytes": len(audio),
|
||||
"error": False,
|
||||
}
|
||||
|
||||
session_done = ws.receive_json()
|
||||
assert session_done == {"type": "session.done", "total_sentences": 1}
|
||||
assert session_done == {"type": "session.done", "utterance_index": 0, "total_sentences": 1}
|
||||
|
||||
assert speech_service._generate_audio_bytes.await_count == 1
|
||||
|
||||
def test_input_done_flushes_and_keeps_connection_open(self, mocker: MockerFixture):
|
||||
app, speech_service = _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"})
|
||||
|
||||
# The same connection serves several utterances; the config sent
|
||||
# once at the top keeps applying and utterance_index rises.
|
||||
for expected_index, text in enumerate(("First utterance. ", "Second utterance. ")):
|
||||
ws.send_json({"type": "input.text", "text": text})
|
||||
ws.send_json({"type": "input.done"})
|
||||
|
||||
start = ws.receive_json()
|
||||
assert start["type"] == "audio.start"
|
||||
assert start["utterance_index"] == expected_index
|
||||
assert start["sentence_text"] == text.strip()
|
||||
assert ws.receive_bytes().startswith(b"RIFF")
|
||||
assert ws.receive_json()["type"] == "audio.done"
|
||||
assert ws.receive_json() == {
|
||||
"type": "session.done",
|
||||
"utterance_index": expected_index,
|
||||
"total_sentences": 1,
|
||||
}
|
||||
|
||||
assert speech_service._generate_audio_bytes.await_count == 2
|
||||
assert [call.args[0].voice for call in speech_service._generate_audio_bytes.await_args_list] == [
|
||||
"Vivian",
|
||||
"Vivian",
|
||||
]
|
||||
|
||||
def test_sentence_index_stays_within_the_flushed_utterance(self, mocker: MockerFixture):
|
||||
# sentence_index counts within one flush and utterance_index counts the
|
||||
# flushes, so a late utterance never reports "sentence 2 of 1".
|
||||
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"})
|
||||
|
||||
for expected_index in range(3):
|
||||
ws.send_json({"type": "input.text", "text": "Hello world. "})
|
||||
ws.send_json({"type": "input.done"})
|
||||
|
||||
start = ws.receive_json()
|
||||
assert start["utterance_index"] == expected_index
|
||||
assert start["sentence_index"] == 0
|
||||
ws.receive_bytes()
|
||||
|
||||
done = ws.receive_json()
|
||||
assert done["utterance_index"] == expected_index
|
||||
assert done["sentence_index"] == 0
|
||||
|
||||
session_done = ws.receive_json()
|
||||
assert session_done["utterance_index"] == expected_index
|
||||
assert session_done["total_sentences"] == 1
|
||||
assert start["sentence_index"] < session_done["total_sentences"]
|
||||
|
||||
def test_session_config_between_utterances_replaces_config(self, mocker: MockerFixture):
|
||||
app, speech_service = _build_test_app(mocker=mocker)
|
||||
|
||||
with TestClient(app) as client:
|
||||
with client.websocket_connect("/v1/audio/speech/stream") as ws:
|
||||
for expected_index, voice in enumerate(("Vivian", "Serena")):
|
||||
ws.send_json({"type": "session.config", "voice": voice})
|
||||
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()["type"] == "audio.done"
|
||||
assert ws.receive_json() == {
|
||||
"type": "session.done",
|
||||
"utterance_index": expected_index,
|
||||
"total_sentences": 1,
|
||||
}
|
||||
|
||||
assert [call.args[0].voice for call in speech_service._generate_audio_bytes.await_args_list] == [
|
||||
"Vivian",
|
||||
"Serena",
|
||||
]
|
||||
|
||||
def test_session_config_rejected_while_input_is_buffered(self, mocker: MockerFixture):
|
||||
app, speech_service = _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": "Hello world. "})
|
||||
ws.send_json({"type": "session.config", "voice": "Serena"})
|
||||
|
||||
error = ws.receive_json()
|
||||
assert error["type"] == "error"
|
||||
assert "while input is buffered" in error["message"]
|
||||
|
||||
# The buffered text survives the rejected reconfiguration.
|
||||
ws.send_json({"type": "input.done"})
|
||||
assert ws.receive_json()["sentence_text"] == "Hello world."
|
||||
ws.receive_bytes()
|
||||
assert ws.receive_json()["type"] == "audio.done"
|
||||
assert ws.receive_json() == {"type": "session.done", "utterance_index": 0, "total_sentences": 1}
|
||||
|
||||
assert speech_service._generate_audio_bytes.await_args_list[0].args[0].voice == "Vivian"
|
||||
|
||||
def test_session_close_ends_connection(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": "Hello world. "})
|
||||
ws.send_json({"type": "input.done"})
|
||||
|
||||
assert ws.receive_json()["type"] == "audio.start"
|
||||
ws.receive_bytes()
|
||||
assert ws.receive_json()["type"] == "audio.done"
|
||||
assert ws.receive_json() == {"type": "session.done", "utterance_index": 0, "total_sentences": 1}
|
||||
|
||||
ws.send_json({"type": "session.close"})
|
||||
with pytest.raises(WebSocketDisconnect):
|
||||
ws.receive_json()
|
||||
|
||||
def test_idle_timeout_closes_reused_connection(self, mocker: MockerFixture):
|
||||
app, _ = _build_test_app(idle_timeout=0.05, 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": "Hello world. "})
|
||||
ws.send_json({"type": "input.done"})
|
||||
|
||||
assert ws.receive_json()["type"] == "audio.start"
|
||||
ws.receive_bytes()
|
||||
assert ws.receive_json()["type"] == "audio.done"
|
||||
assert ws.receive_json() == {"type": "session.done", "utterance_index": 0, "total_sentences": 1}
|
||||
|
||||
# Holding a connection open is not free: an idle client still
|
||||
# gets timed out after the flush.
|
||||
assert ws.receive_json() == {
|
||||
"type": "error",
|
||||
"message": "Idle timeout: no message received",
|
||||
}
|
||||
|
||||
def test_streaming_multiple_binary_frames(self, mocker: MockerFixture):
|
||||
captured_requests = []
|
||||
|
||||
@@ -125,9 +274,15 @@ class TestStreamingSpeechWebSocket:
|
||||
assert ws.receive_bytes() == b"\x06"
|
||||
|
||||
done = ws.receive_json()
|
||||
assert done == {"type": "audio.done", "sentence_index": 0, "total_bytes": 6, "error": False}
|
||||
assert done == {
|
||||
"type": "audio.done",
|
||||
"utterance_index": 0,
|
||||
"sentence_index": 0,
|
||||
"total_bytes": 6,
|
||||
"error": False,
|
||||
}
|
||||
|
||||
assert ws.receive_json() == {"type": "session.done", "total_sentences": 1}
|
||||
assert ws.receive_json() == {"type": "session.done", "utterance_index": 0, "total_sentences": 1}
|
||||
|
||||
assert len(captured_requests) == 1
|
||||
assert captured_requests[0].stream is True
|
||||
@@ -210,6 +365,8 @@ class TestStreamingSpeechWebSocket:
|
||||
# is fully aligned.
|
||||
chunk = ws.receive_json()
|
||||
assert chunk["type"] == "audio.chunk"
|
||||
assert chunk["utterance_index"] == 0
|
||||
assert chunk["sentence_index"] == 0
|
||||
assert chunk["chunk_id"] == 0
|
||||
assert chunk["chunk_start_ms"] == 0
|
||||
assert chunk["chunk_end_ms"] == 500
|
||||
@@ -240,7 +397,13 @@ class TestStreamingSpeechWebSocket:
|
||||
]
|
||||
|
||||
done = ws.receive_json()
|
||||
assert done == {"type": "audio.done", "sentence_index": 0, "total_bytes": 2000, "error": False}
|
||||
assert done == {
|
||||
"type": "audio.done",
|
||||
"utterance_index": 0,
|
||||
"sentence_index": 0,
|
||||
"total_bytes": 2000,
|
||||
"error": False,
|
||||
}
|
||||
|
||||
assert captured_requests[0].word_timestamps is True
|
||||
assert mock_align.await_count == 1
|
||||
@@ -309,11 +472,12 @@ class TestStreamingSpeechWebSocket:
|
||||
assert ws.receive_bytes()
|
||||
assert ws.receive_json() == {
|
||||
"type": "audio.done",
|
||||
"utterance_index": 0,
|
||||
"sentence_index": 0,
|
||||
"total_bytes": 36,
|
||||
"error": False,
|
||||
}
|
||||
assert ws.receive_json() == {"type": "session.done", "total_sentences": 1}
|
||||
assert ws.receive_json() == {"type": "session.done", "utterance_index": 0, "total_sentences": 1}
|
||||
|
||||
def test_invalid_streaming_config(self, mocker: MockerFixture):
|
||||
app, _ = _build_test_app(mocker=mocker)
|
||||
@@ -341,7 +505,7 @@ class TestStreamingSpeechWebSocket:
|
||||
ws.send_json({"type": "input.text", "text": ""})
|
||||
ws.send_json({"type": "input.done"})
|
||||
|
||||
assert ws.receive_json() == {"type": "session.done", "total_sentences": 0}
|
||||
assert ws.receive_json() == {"type": "session.done", "utterance_index": 0, "total_sentences": 0}
|
||||
|
||||
assert speech_service._generate_audio_bytes.await_count == 0
|
||||
|
||||
@@ -361,11 +525,12 @@ class TestStreamingSpeechWebSocket:
|
||||
ws.receive_bytes()
|
||||
assert ws.receive_json() == {
|
||||
"type": "audio.done",
|
||||
"utterance_index": 0,
|
||||
"sentence_index": 0,
|
||||
"total_bytes": 36,
|
||||
"error": False,
|
||||
}
|
||||
assert ws.receive_json() == {"type": "session.done", "total_sentences": 1}
|
||||
assert ws.receive_json() == {"type": "session.done", "utterance_index": 0, "total_sentences": 1}
|
||||
|
||||
def test_unknown_message_type_keeps_session_open(self, mocker: MockerFixture):
|
||||
app, _ = _build_test_app(mocker=mocker)
|
||||
@@ -384,12 +549,13 @@ class TestStreamingSpeechWebSocket:
|
||||
ws.receive_bytes()
|
||||
assert ws.receive_json() == {
|
||||
"type": "audio.done",
|
||||
"utterance_index": 0,
|
||||
"sentence_index": 0,
|
||||
"total_bytes": 36,
|
||||
"error": False,
|
||||
}
|
||||
|
||||
assert ws.receive_json() == {"type": "session.done", "total_sentences": 1}
|
||||
assert ws.receive_json() == {"type": "session.done", "utterance_index": 0, "total_sentences": 1}
|
||||
|
||||
def test_config_timeout_closes_session(self, mocker: MockerFixture):
|
||||
app, _ = _build_test_app(config_timeout=0.01, mocker=mocker)
|
||||
@@ -416,10 +582,19 @@ class TestStreamingSpeechWebSocket:
|
||||
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}
|
||||
assert ws.receive_json() == {
|
||||
"type": "error",
|
||||
"message": "Generation failed for utterance 0, sentence 0: boom",
|
||||
}
|
||||
assert ws.receive_json() == {
|
||||
"type": "audio.done",
|
||||
"utterance_index": 0,
|
||||
"sentence_index": 0,
|
||||
"total_bytes": 0,
|
||||
"error": True,
|
||||
}
|
||||
|
||||
assert ws.receive_json() == {"type": "session.done", "total_sentences": 1}
|
||||
assert ws.receive_json() == {"type": "session.done", "utterance_index": 0, "total_sentences": 1}
|
||||
|
||||
def test_streaming_generation_error_marks_audio_done(self, mocker: MockerFixture):
|
||||
speech_service = mocker.MagicMock(spec=OmniOpenAIServingSpeech)
|
||||
@@ -453,16 +628,17 @@ class TestStreamingSpeechWebSocket:
|
||||
assert ws.receive_bytes() == b"\x01\x02"
|
||||
assert ws.receive_json() == {
|
||||
"type": "error",
|
||||
"message": "Generation failed for sentence 0: stream boom",
|
||||
"message": "Generation failed for utterance 0, sentence 0: stream boom",
|
||||
}
|
||||
assert ws.receive_json() == {
|
||||
"type": "audio.done",
|
||||
"utterance_index": 0,
|
||||
"sentence_index": 0,
|
||||
"total_bytes": 2,
|
||||
"error": True,
|
||||
}
|
||||
|
||||
assert ws.receive_json() == {"type": "session.done", "total_sentences": 1}
|
||||
assert ws.receive_json() == {"type": "session.done", "utterance_index": 0, "total_sentences": 1}
|
||||
|
||||
def test_invalid_input_text_type_returns_validation_error(self, mocker: MockerFixture):
|
||||
app, speech_service = _build_test_app(mocker=mocker)
|
||||
@@ -478,7 +654,7 @@ class TestStreamingSpeechWebSocket:
|
||||
}
|
||||
|
||||
ws.send_json({"type": "input.done"})
|
||||
assert ws.receive_json() == {"type": "session.done", "total_sentences": 0}
|
||||
assert ws.receive_json() == {"type": "session.done", "utterance_index": 0, "total_sentences": 0}
|
||||
|
||||
assert speech_service._generate_audio_bytes.await_count == 0
|
||||
|
||||
@@ -497,7 +673,7 @@ class TestStreamingSpeechWebSocket:
|
||||
}
|
||||
|
||||
ws.send_json({"type": "input.done"})
|
||||
assert ws.receive_json() == {"type": "session.done", "total_sentences": 0}
|
||||
assert ws.receive_json() == {"type": "session.done", "utterance_index": 0, "total_sentences": 0}
|
||||
|
||||
assert speech_service._generate_audio_bytes.await_count == 0
|
||||
|
||||
@@ -550,7 +726,15 @@ class TestStreamingSpeechWebSocket:
|
||||
config.word_timestamps = False
|
||||
|
||||
with pytest.raises(WebSocketDisconnect):
|
||||
asyncio.run(handler._generate_and_send(websocket, config, "Hello world.", 0))
|
||||
asyncio.run(
|
||||
handler._generate_and_send(
|
||||
websocket,
|
||||
config,
|
||||
"Hello world.",
|
||||
utterance_index=0,
|
||||
sentence_index=0,
|
||||
)
|
||||
)
|
||||
|
||||
speech_service.engine_client.abort.assert_awaited_once_with("req-abort")
|
||||
assert websocket.send_json.await_count == 2
|
||||
|
||||
@@ -1572,8 +1572,9 @@ async def delete_voice(name: str, raw_request: Request):
|
||||
async def streaming_speech(websocket: WebSocket):
|
||||
"""WebSocket endpoint for streaming text input TTS.
|
||||
|
||||
Accepts text incrementally, splits at sentence boundaries, and
|
||||
returns audio per sentence. See serving_speech_stream.py for protocol.
|
||||
Accepts text incrementally and returns audio for the buffered text on
|
||||
input.done, which flushes without closing the connection. See
|
||||
serving_speech_stream.py for protocol.
|
||||
"""
|
||||
handler = getattr(websocket.app.state, "openai_streaming_speech", None)
|
||||
if handler is None:
|
||||
|
||||
@@ -3,26 +3,46 @@
|
||||
Accepts text incrementally via WebSocket, buffers it until input.done, and
|
||||
generates audio once for the buffered input using the existing TTS pipeline.
|
||||
|
||||
input.done is a flush, not a close: it ends the current utterance and the
|
||||
connection stays open, so the next utterance reuses the same connection
|
||||
instead of paying another WebSocket handshake. A connection ends on
|
||||
session.close, on the idle timeout, or when the client closes the socket.
|
||||
The session config is sticky across flushes and can be replaced by sending
|
||||
another session.config between utterances.
|
||||
|
||||
"Utterance" here names the flush unit rather than any linguistic unit: it is
|
||||
whatever text the client had buffered when it sent input.done, from a word to
|
||||
several paragraphs, synthesized as a single request.
|
||||
|
||||
Protocol:
|
||||
Client -> Server:
|
||||
{"type": "session.config", ...} # Session config (sent once first)
|
||||
{"type": "session.config", ...} # Session config (first message; repeatable)
|
||||
{"type": "input.text", "text": "..."} # Text chunks
|
||||
{"type": "input.done"} # End of input
|
||||
{"type": "input.done"} # End of utterance, flush and keep connection open
|
||||
{"type": "session.close"} # End of connection
|
||||
|
||||
Server -> Client (default, word_timestamps=false):
|
||||
{"type": "audio.start", "sentence_index": 0, "sentence_text": "...", "format": "wav"}
|
||||
{"type": "audio.start", "utterance_index": 0, "sentence_index": 0,
|
||||
"sentence_text": "...", "format": "wav"}
|
||||
<binary frame: audio bytes>
|
||||
...
|
||||
{"type": "audio.done", "sentence_index": 0}
|
||||
{"type": "session.done", "total_sentences": N}
|
||||
{"type": "audio.done", "utterance_index": 0, "sentence_index": 0}
|
||||
{"type": "session.done", "utterance_index": 0, "total_sentences": N}
|
||||
{"type": "error", "message": "..."}
|
||||
# session.done ends the flushed utterance, not the connection. An
|
||||
# utterance is just the flush unit: whatever text was buffered when
|
||||
# input.done arrived, of any length. utterance_index counts those
|
||||
# flushes across the connection, while sentence_index counts within
|
||||
# one of them and so pairs with total_sentences.
|
||||
|
||||
Server -> Client (when word_timestamps=true):
|
||||
{"type": "audio.start", "sentence_index": 0, "sentence_text": "...", "format": "pcm"}
|
||||
{"type": "audio.chunk", "sentence_index": 0, "chunk_id": 0, "audio_b64": "<base64 PCM>", "timestamps": null}
|
||||
{"type": "audio.start", "utterance_index": 0, "sentence_index": 0,
|
||||
"sentence_text": "...", "format": "pcm"}
|
||||
{"type": "audio.chunk", "utterance_index": 0, "sentence_index": 0, "chunk_id": 0,
|
||||
"audio_b64": "<base64 PCM>", "timestamps": null}
|
||||
...
|
||||
{"type": "audio.chunk", "audio_b64": "", "timestamps": [{"word", "start_ms", "end_ms"}, ...]}
|
||||
{"type": "audio.done", "sentence_index": 0}
|
||||
{"type": "audio.done", "utterance_index": 0, "sentence_index": 0}
|
||||
# Audio is JSON base64 PCM (not binary). A trailing empty-audio chunk carries the
|
||||
# full sentence-relative alignment. timestamps: list = aligned, [] = silence, null = failed.
|
||||
"""
|
||||
@@ -57,9 +77,11 @@ _MAX_INPUT_TEXT_MESSAGE_SIZE = 128 * 1024
|
||||
class OmniStreamingSpeechHandler:
|
||||
"""Handles WebSocket sessions for streaming text-input TTS.
|
||||
|
||||
Each WebSocket connection is an independent session. Text arrives
|
||||
incrementally, is buffered until input.done, and audio is generated once
|
||||
for the buffered input using the existing OmniOpenAIServingSpeech pipeline.
|
||||
A connection carries one or more utterances. Text arrives incrementally,
|
||||
is buffered until input.done, and audio is generated once for the
|
||||
buffered input using the existing OmniOpenAIServingSpeech pipeline. The
|
||||
connection outlives each utterance so a client can keep synthesizing on
|
||||
it without reconnecting.
|
||||
|
||||
Args:
|
||||
speech_service: The existing TTS serving instance (reused for
|
||||
@@ -79,54 +101,64 @@ class OmniStreamingSpeechHandler:
|
||||
self._config_timeout = config_timeout
|
||||
|
||||
async def handle_session(self, websocket: WebSocket) -> None:
|
||||
"""Main session loop for a single WebSocket connection."""
|
||||
"""Main loop for a single WebSocket connection.
|
||||
|
||||
Serves any number of utterances: text is buffered until input.done,
|
||||
which flushes it as one TTS request and then leaves the connection
|
||||
open for the next one. Rejecting a message is only fatal before the
|
||||
first valid session.config; afterwards the error is reported and the
|
||||
connection survives.
|
||||
"""
|
||||
await websocket.accept()
|
||||
|
||||
config: StreamingSpeechSessionConfig | None = None
|
||||
text_parts: list[str] = []
|
||||
utterance_index = 0
|
||||
|
||||
try:
|
||||
# 1. Wait for session.config
|
||||
config = await self._receive_config(websocket)
|
||||
if config is None:
|
||||
return # Error already sent, connection closing
|
||||
|
||||
# Validate model if specified
|
||||
if config.model and hasattr(self._speech_service, "_check_model"):
|
||||
error = await self._speech_service._check_model(
|
||||
OpenAICreateSpeechRequest(input="ping", model=config.model)
|
||||
)
|
||||
if error is not None:
|
||||
await self._send_error(websocket, str(error))
|
||||
return
|
||||
|
||||
text_parts: list[str] = []
|
||||
|
||||
# 2. Receive text chunks until input.done
|
||||
while True:
|
||||
try:
|
||||
raw = await asyncio.wait_for(
|
||||
websocket.receive_text(),
|
||||
timeout=self._idle_timeout,
|
||||
timeout=self._config_timeout if config is None else self._idle_timeout,
|
||||
)
|
||||
except asyncio.TimeoutError:
|
||||
await self._send_error(websocket, "Idle timeout: no message received")
|
||||
if config is None:
|
||||
await self._send_error(websocket, "Timeout waiting for session.config")
|
||||
else:
|
||||
await self._send_error(websocket, "Idle timeout: no message received")
|
||||
return
|
||||
|
||||
if len(raw) > _MAX_INPUT_TEXT_MESSAGE_SIZE:
|
||||
await self._send_error(websocket, "input.text message too large")
|
||||
continue
|
||||
|
||||
try:
|
||||
msg = json.loads(raw)
|
||||
except json.JSONDecodeError:
|
||||
await self._send_error(websocket, "Invalid JSON message")
|
||||
continue
|
||||
|
||||
if not isinstance(msg, dict):
|
||||
await self._send_error(websocket, "WebSocket messages must be JSON objects")
|
||||
msg = await self._parse_message(websocket, raw)
|
||||
if msg is None:
|
||||
if config is None:
|
||||
return # Malformed handshake, connection closing
|
||||
continue
|
||||
|
||||
msg_type = msg.get("type")
|
||||
|
||||
if msg_type == "input.text":
|
||||
if msg_type == "session.config":
|
||||
if text_parts:
|
||||
await self._send_error(
|
||||
websocket,
|
||||
"session.config cannot be applied while input is buffered; send input.done first",
|
||||
)
|
||||
continue
|
||||
new_config = await self._build_config(websocket, msg)
|
||||
if new_config is None:
|
||||
if config is None:
|
||||
return # Error already sent, connection closing
|
||||
continue # Keep serving with the previous config
|
||||
config = new_config
|
||||
|
||||
elif config is None:
|
||||
await self._send_error(
|
||||
websocket,
|
||||
f"Expected session.config, got: {msg_type}",
|
||||
)
|
||||
return
|
||||
|
||||
elif msg_type == "input.text":
|
||||
text = msg.get("text", "")
|
||||
if not isinstance(text, str):
|
||||
await self._send_error(websocket, "input.text requires a string value")
|
||||
@@ -135,17 +167,31 @@ class OmniStreamingSpeechHandler:
|
||||
|
||||
elif msg_type == "input.done":
|
||||
full_text = "".join(text_parts).strip()
|
||||
text_parts.clear()
|
||||
total_sentences = 0
|
||||
if full_text:
|
||||
await self._generate_and_send(websocket, config, full_text, 0)
|
||||
# However long the buffered text is, the pipeline takes
|
||||
# it as one request, so every flush is sentence 0 of 1.
|
||||
await self._generate_and_send(
|
||||
websocket,
|
||||
config,
|
||||
full_text,
|
||||
utterance_index=utterance_index,
|
||||
sentence_index=0,
|
||||
)
|
||||
total_sentences = 1
|
||||
|
||||
await websocket.send_json(
|
||||
{
|
||||
"type": "session.done",
|
||||
"utterance_index": utterance_index,
|
||||
"total_sentences": total_sentences,
|
||||
}
|
||||
)
|
||||
utterance_index += 1
|
||||
|
||||
elif msg_type == "session.close":
|
||||
await websocket.close()
|
||||
return
|
||||
|
||||
else:
|
||||
@@ -163,44 +209,50 @@ class OmniStreamingSpeechHandler:
|
||||
except Exception:
|
||||
logger.debug("Failed to send error to streaming speech client", exc_info=True)
|
||||
|
||||
async def _receive_config(self, websocket: WebSocket) -> StreamingSpeechSessionConfig | None:
|
||||
"""Wait for and validate the session.config message."""
|
||||
try:
|
||||
raw = await asyncio.wait_for(
|
||||
websocket.receive_text(),
|
||||
timeout=self._config_timeout,
|
||||
)
|
||||
except asyncio.TimeoutError:
|
||||
await self._send_error(websocket, "Timeout waiting for session.config")
|
||||
return None
|
||||
async def _parse_message(self, websocket: WebSocket, raw: str) -> dict | None:
|
||||
"""Decode one client message, or report why it was rejected.
|
||||
|
||||
if len(raw) > _MAX_CONFIG_MESSAGE_SIZE:
|
||||
await self._send_error(websocket, "session.config message too large")
|
||||
Size limits are per message type: session.config carries ref_audio
|
||||
payloads and gets the larger budget.
|
||||
"""
|
||||
if len(raw) > max(_MAX_CONFIG_MESSAGE_SIZE, _MAX_INPUT_TEXT_MESSAGE_SIZE):
|
||||
await self._send_error(websocket, "WebSocket message too large")
|
||||
return None
|
||||
|
||||
try:
|
||||
msg = json.loads(raw)
|
||||
except json.JSONDecodeError:
|
||||
await self._send_error(websocket, "Invalid JSON in session.config")
|
||||
await self._send_error(websocket, "Invalid JSON message")
|
||||
return None
|
||||
|
||||
if not isinstance(msg, dict):
|
||||
await self._send_error(websocket, "session.config must be a JSON object")
|
||||
await self._send_error(websocket, "WebSocket messages must be JSON objects")
|
||||
return None
|
||||
|
||||
if msg.get("type") != "session.config":
|
||||
await self._send_error(
|
||||
websocket,
|
||||
f"Expected session.config, got: {msg.get('type')}",
|
||||
)
|
||||
if msg.get("type") == "session.config":
|
||||
if len(raw) > _MAX_CONFIG_MESSAGE_SIZE:
|
||||
await self._send_error(websocket, "session.config message too large")
|
||||
return None
|
||||
elif len(raw) > _MAX_INPUT_TEXT_MESSAGE_SIZE:
|
||||
await self._send_error(websocket, "input.text message too large")
|
||||
return None
|
||||
|
||||
return msg
|
||||
|
||||
async def _build_config(self, websocket: WebSocket, msg: dict) -> StreamingSpeechSessionConfig | None:
|
||||
"""Validate a session.config message and its model."""
|
||||
try:
|
||||
config = StreamingSpeechSessionConfig(**{k: v for k, v in msg.items() if k != "type"})
|
||||
except ValidationError as e:
|
||||
await self._send_error(websocket, f"Invalid session config: {e}")
|
||||
return None
|
||||
|
||||
if config.model and hasattr(self._speech_service, "_check_model"):
|
||||
error = await self._speech_service._check_model(OpenAICreateSpeechRequest(input="ping", model=config.model))
|
||||
if error is not None:
|
||||
await self._send_error(websocket, str(error))
|
||||
return None
|
||||
|
||||
return config
|
||||
|
||||
async def _generate_and_send(
|
||||
@@ -208,9 +260,15 @@ class OmniStreamingSpeechHandler:
|
||||
websocket: WebSocket,
|
||||
config: StreamingSpeechSessionConfig,
|
||||
sentence_text: str,
|
||||
*,
|
||||
utterance_index: int,
|
||||
sentence_index: int,
|
||||
) -> None:
|
||||
"""Generate audio for a single sentence and send it over WebSocket."""
|
||||
"""Generate audio for a single sentence and send it over WebSocket.
|
||||
|
||||
``utterance_index`` identifies the flush this sentence belongs to and
|
||||
``sentence_index`` its position inside that flush.
|
||||
"""
|
||||
response_format = config.response_format or "wav"
|
||||
|
||||
# Reject unmet word-timestamps preconditions early with a clear reason.
|
||||
@@ -253,6 +311,7 @@ class OmniStreamingSpeechHandler:
|
||||
|
||||
start_payload = {
|
||||
"type": "audio.start",
|
||||
"utterance_index": utterance_index,
|
||||
"sentence_index": sentence_index,
|
||||
"sentence_text": sentence_text,
|
||||
"format": response_format,
|
||||
@@ -277,6 +336,7 @@ class OmniStreamingSpeechHandler:
|
||||
request_id=request_id,
|
||||
generator=generator,
|
||||
sentence_text=sentence_text,
|
||||
utterance_index=utterance_index,
|
||||
sentence_index=sentence_index,
|
||||
language=config.language,
|
||||
)
|
||||
@@ -298,13 +358,22 @@ class OmniStreamingSpeechHandler:
|
||||
raise
|
||||
except Exception as e:
|
||||
generation_failed = True
|
||||
logger.error("Generation failed for sentence %d: %s", sentence_index, e)
|
||||
await self._send_error(websocket, f"Generation failed for sentence {sentence_index}: {e}")
|
||||
logger.error(
|
||||
"Generation failed for utterance %d, sentence %d: %s",
|
||||
utterance_index,
|
||||
sentence_index,
|
||||
e,
|
||||
)
|
||||
await self._send_error(
|
||||
websocket,
|
||||
f"Generation failed for utterance {utterance_index}, sentence {sentence_index}: {e}",
|
||||
)
|
||||
finally:
|
||||
try:
|
||||
await websocket.send_json(
|
||||
{
|
||||
"type": "audio.done",
|
||||
"utterance_index": utterance_index,
|
||||
"sentence_index": sentence_index,
|
||||
"total_bytes": total_bytes,
|
||||
"error": generation_failed,
|
||||
@@ -320,6 +389,7 @@ class OmniStreamingSpeechHandler:
|
||||
request_id: str,
|
||||
generator,
|
||||
sentence_text: str,
|
||||
utterance_index: int,
|
||||
sentence_index: int,
|
||||
language: str | None = None,
|
||||
) -> int:
|
||||
@@ -350,6 +420,7 @@ class OmniStreamingSpeechHandler:
|
||||
await websocket.send_json(
|
||||
{
|
||||
"type": "audio.chunk",
|
||||
"utterance_index": utterance_index,
|
||||
"sentence_index": sentence_index,
|
||||
"chunk_id": chunk_id,
|
||||
"chunk_start_ms": chunk_start_ms,
|
||||
|
||||
Reference in new issue
Block a user