Qwen3-TTS PR 3: sentence chunking and error status for /v1/audio/speech

mstar/api_server/entrypoint.py

@@ -766,6 +766,17 @@ def abort_request(self, request_id: str) -> None:
         logger.info("Client cancelled request %s; releasing resources", request_id)
         self.preprocess_worker.abort_request(request_id)
 
+    def release_request(self, request_id: str) -> None:
+        """Drop a request whose results will never be read: abort it if it is
+        still running, else just forget it (an ended request needs no abort)."""
+        with self.request_lock:
+            req = self.pending_requests.get(request_id)
+            finished = req is not None and req.event.is_set()
+            if finished:
+                self.pending_requests.pop(request_id)
+        if not finished:
+            self.abort_request(request_id)
+
     # ----------------------------------------------------------
     # Cleanup
     # ----------------------------------------------------------

mstar/api_server/openai/adapters.py

@@ -185,6 +185,18 @@ class OpenAIAdapter:
     supports_videos: bool = False   # POST /v1/videos/generations
     supports_realtime: bool = False  # /v1/realtime (bidirectional speech WebSocket)
 
+    # ``/v1/audio/speech`` sentence chunking (``serving_speech``): inputs of at
+    # least ``speech_chunk_min_chars`` characters are split into sentence
+    # groups of about ``speech_chunk_max_chars`` and synthesized as ordered
+    # sub-requests. ``None`` keeps the whole text in one request unless the
+    # client sends ``sentence_chunking: true``.
+    speech_chunk_min_chars: int | None = None
+    speech_chunk_max_chars: int = 400
+    # sub-requests kept in flight ahead of the one being streamed
+    speech_chunk_lookahead: int = 2
+    # most sub-requests one input may split into; longer inputs get a 400
+    speech_chunk_max_pieces: int = 32
+
     def chat_to_request(self, req: ChatCompletionRequest, upload_dir: Path) -> SubmitArgs:  # noqa: ARG002
         # Output modalities vary by model: e.g. Qwen3-Omni speech output also
         # emits text, whereas BAGEL chat is text-only.

mstar/api_server/openai/serving_speech.py

@@ -5,15 +5,61 @@
 a WAV header followed by PCM16 frames, or, with ``response_format="pcm"``,
 the bare PCM16 frames that the OpenAI TTS clients (LiveKit, Pipecat) expect.
 ``/v1/audio/voices`` lists the ``voice`` values the served model accepts.
+
+Long inputs can be synthesized as ordered sentence chunks (one engine request
+per chunk, ``speech_chunking.split_sentences``): the adapter's
+``speech_chunk_min_chars`` turns this on for long texts, and a client can force
+or suppress it per request with ``sentence_chunking: true|false``. The next
+chunks are submitted while the current one streams, so the engine batches them
+and playback never waits for a prefill. An input needing more than the
+adapter's ``speech_chunk_max_pieces`` chunks is rejected.
+
+A streaming request only commits to HTTP 200 once its first result chunk has
+arrived and is not an error; an error chunk before that becomes the HTTP error
+it carries (the non-streaming path gets the same from ``collect_results``).
 """
 
 from __future__ import annotations
 
+from collections.abc import Callable
+
+from fastapi import HTTPException
 from fastapi.responses import Response, StreamingResponse
 
 from mstar.api_server import media_io
 from mstar.api_server.openai._util import rid
 from mstar.api_server.openai.protocol import VoiceCard, VoiceList
+from mstar.api_server.openai.speech_chunking import split_sentences
+
+
+def _plan_chunks(req, adapter, text: str) -> list[str]:
+    """The texts to synthesize, in playback order (``[text]`` when unchunked)."""
+    # Extra request fields live in pydantic's ``model_extra``; plain objects
+    # (tests, other frontends) may carry the attribute directly.
+    requested = (getattr(req, "model_extra", None) or {}).get("sentence_chunking")
+    if requested is None:
+        requested = getattr(req, "sentence_chunking", None)
+    min_chars = getattr(adapter, "speech_chunk_min_chars", None)
+    if requested is False or (requested is None and (min_chars is None or len(text) < min_chars)):
+        return [text]
+    chunks = split_sentences(text, max_chars=getattr(adapter, "speech_chunk_max_chars", 400))
+    max_pieces = getattr(adapter, "speech_chunk_max_pieces", 32)
+    if len(chunks) > max_pieces:
+        raise HTTPException(
+            status_code=400,
+            detail=f"input splits into {len(chunks)} sentence chunks; at most {max_pieces} are allowed",
+        )
+    return chunks or [text]
+
+
+def _chunk_kwargs(model_kwargs: dict, index: int) -> dict:
+    """Per-chunk model kwargs: identical, except a client seed advances per chunk."""
+    kwargs = dict(model_kwargs)
+    kwargs.pop("sentence_chunking", None)
+    seed = kwargs.get("seed")
+    if isinstance(seed, int) and not isinstance(seed, bool):
+        kwargs["seed"] = seed + index
+    return kwargs
 
 
 def list_voices(api) -> VoiceList | None:
@@ -33,34 +79,117 @@ async def create_speech(api, model_name, adapter, req, raw_request=None):  # noq
     request_id = rid("speech")
     sample_rate = api.model.get_output_sample_rate("audio") if api.model is not None else 24000
     fmt = (req.response_format or "wav").lower()
+    chunks = _plan_chunks(req, adapter, args.text or "")
 
-    api.submit_request(
-        text=args.text,
-        file_paths=args.file_paths,
-        input_modalities=args.input_modalities,
-        output_modalities=args.output_modalities,
-        model_kwargs=args.model_kwargs,
-        streaming=bool(req.stream),
-        request_id=request_id,
-    )
+    def submit(index: int) -> str:
+        chunk_id = request_id if len(chunks) == 1 else f"{request_id}-{index}"
+        return api.submit_request(
+            text=chunks[index],
+            file_paths=args.file_paths,
+            input_modalities=args.input_modalities,
+            output_modalities=args.output_modalities,
+            model_kwargs=_chunk_kwargs(args.model_kwargs, index),
+            streaming=bool(req.stream),
+            request_id=chunk_id,
+        )
 
+    lookahead = max(1, int(getattr(adapter, "speech_chunk_lookahead", 2)))
+    pending = [submit(i) for i in range(min(lookahead, len(chunks)))]
     if req.stream:
         raw_pcm = fmt == "pcm"
+        # Look at the first result before committing to a 200: a request the
+        # engine rejects (bad voice, dead worker, ...) must surface as an HTTP
+        # error, not as an empty WAV.
+        first_iter = api.iter_result_chunks(pending[0])
+        try:
+            first = await anext(first_iter, None)
+            if _is_error(first):
+                # Read it to the end so the finished request isn't aborted.
+                await _drain(first_iter)
+                raise _http_error(first)
+        except BaseException:
+            # The first piece released itself; release the ones submitted ahead.
+            _release_unread(api, pending, 1)
+            raise
         return StreamingResponse(
-            _stream_pcm(api, request_id, sample_rate, with_wav_header=not raw_pcm),
+            _stream_pcm(api, submit, len(chunks), pending, first_iter, first, sample_rate,
+                        with_wav_header=not raw_pcm),
             media_type="audio/pcm" if raw_pcm else "audio/wav",
             headers={"Cache-Control": "no-cache"},
         )
 
-    chunks = await api.collect_results(request_id, raw_request)
-    pcm = b"".join(c.data for c in chunks if c.modality == "audio")
-    audio_bytes, mime = media_io.pcm16_to_container(pcm, sample_rate, fmt)
+    pcm_parts: list[bytes] = []
+    read = 0  # pieces collect_results has released; the rest are released on early exit
+    try:
+        for index in range(len(chunks)):
+            if len(pending) < len(chunks):
+                pending.append(submit(len(pending)))
+            try:
+                results = await api.collect_results(pending[index], raw_request)
+            except HTTPException:
+                read = index + 1
+                raise
+            read = index + 1
+            if not results and raw_request is not None and await raw_request.is_disconnected():
+                break
+            pcm_parts.append(b"".join(c.data for c in results if c.modality == "audio"))
+    finally:
+        _release_unread(api, pending, read)
+    audio_bytes, mime = media_io.pcm16_to_container(b"".join(pcm_parts), sample_rate, fmt)
     return Response(content=audio_bytes, media_type=mime)
 
 
-async def _stream_pcm(api, request_id, sample_rate, with_wav_header=True):
-    if with_wav_header:
-        yield media_io.wav_stream_header(sample_rate)
-    async for c in api.iter_result_chunks(request_id):
-        if c.modality == "audio" and c.data:
-            yield c.data
+def _is_error(chunk) -> bool:
+    return chunk is not None and chunk.modality == "error"
+
+
+def _http_error(chunk) -> HTTPException:
+    """A data-worker failure arrives as an ``error`` chunk; the HTTP error it carries."""
+    return HTTPException(
+        status_code=int((chunk.metadata or {}).get("status", 500)),
+        detail=chunk.data.decode("utf-8", "replace") if isinstance(chunk.data, bytes) else str(chunk.data),
+    )
+
+
+async def _drain(iterator) -> None:
+    async for _ in iterator:
+        pass
+
+
+def _release_unread(api, pending: list[str], read: int) -> None:
+    """Release every submitted piece from ``pending[read]`` on (not yet fully read)."""
+    for chunk_id in pending[read:]:
+        api.release_request(chunk_id)
+
+
+async def _stream_pcm(api, submit: Callable[[int], str], num_chunks: int, pending: list[str],
+                      first_iter, first, sample_rate: int, with_wav_header: bool = True):
+    read = 0  # pieces read to the end; the rest are released if the stream stops early
+    try:
+        if with_wav_header:
+            yield media_io.wav_stream_header(sample_rate)
+        for index in range(num_chunks):
+            if len(pending) < num_chunks:
+                # Keep the next chunk generating while this one plays.
+                pending.append(submit(len(pending)))
+            if index == 0:
+                iterator, head = first_iter, first
+            else:
+                iterator, head = api.iter_result_chunks(pending[index]), None
+            if head is not None and head.modality == "audio" and head.data:
+                yield head.data
+            error = None
+            async for c in iterator:
+                # Read past an error to the end: leaving early would abort a
+                # request that has already finished.
+                if error is None and _is_error(c):
+                    error = c
+                elif error is None and c.modality == "audio" and c.data:
+                    yield c.data
+            read = index + 1
+            if error is not None:
+                # Mid-stream the status is already sent; closing the stream is
+                # the only honest signal left, so raise rather than end quietly.
+                raise _http_error(error)
+    finally:
+        _release_unread(api, pending, read)

mstar/api_server/openai/speech_chunking.py

@@ -0,0 +1,100 @@
+"""Sentence chunking for ``/v1/audio/speech``.
+
+Long inputs are split into sentence groups that are synthesized as separate
+requests and played back in order. This keeps every autoregressive
+text-to-speech request inside the length its model was trained for, lets the
+engine batch the pieces of one long input like independent requests, and
+starts the second piece while the first is still streaming (see
+``serving_speech``). Splitting is purely textual and model-agnostic.
+"""
+
+from __future__ import annotations
+
+import re
+
+# One sentence: a lazy body up to a Latin terminator (with any closing quotes
+# or brackets) that precedes whitespace or the end, a CJK terminator, a
+# paragraph break, or the end of the text.
+_SENTENCE = re.compile(
+    r""".+?(?:
+        [.!?;]["'”’)\]]*(?=\s|\Z)
+      | [。！？；]
+      | (?=\n[ \t]*\n)
+      | \Z
+    )""",
+    re.VERBOSE | re.DOTALL,
+)
+# Soft break points inside an over-long sentence, most preferred first.
+_SOFT_BREAKS = (
+    re.compile(r"(?<=[,;:，；：])\s*"),
+    re.compile(r"\s+"),
+)
+
+
+def _hard_wrap(sentence: str, max_chars: int) -> list[str]:
+    """Split one over-long sentence at clause boundaries, then at spaces."""
+    pieces = [sentence]
+    for pattern in _SOFT_BREAKS:
+        wrapped: list[str] = []
+        for piece in pieces:
+            if len(piece) <= max_chars:
+                wrapped.append(piece)
+                continue
+            current = ""
+            for part in (p.strip() for p in pattern.split(piece)):
+                if not part:
+                    continue
+                candidate = f"{current} {part}" if current else part
+                if current and len(candidate) > max_chars:
+                    wrapped.append(current)
+                    current = part
+                else:
+                    current = candidate
+            if current:
+                wrapped.append(current)
+        pieces = wrapped
+    return pieces
+
+
+def split_sentences(text: str, max_chars: int = 400, min_chars: int = 24) -> list[str]:
+    """Group ``text`` into sentence chunks of at most ``max_chars`` characters.
+
+    Sentences are never cut unless one alone exceeds ``max_chars`` (then it is
+    wrapped at clause boundaries, or spaces as a last resort). A trailing
+    fragment shorter than ``min_chars`` is merged into its predecessor so the
+    model is not asked to voice a lone "Okay." Returns ``[text]`` when nothing
+    needs splitting and ``[]`` for blank input.
+    """
+    if max_chars <= 0:
+        raise ValueError("max_chars must be positive")
+    text = text.strip()
+    if not text:
+        return []
+    if len(text) <= max_chars:
+        return [text]
+
+    sentences: list[str] = []
+    for match in _SENTENCE.finditer(text):
+        piece = " ".join(match.group(0).split())
+        if not piece:
+            continue
+        sentences.extend(_hard_wrap(piece, max_chars) if len(piece) > max_chars else [piece])
+
+    chunks: list[str] = []
+    current = ""
+    for sentence in sentences:
+        if current and len(current) + 1 + len(sentence) > max_chars:
+            chunks.append(current)
+            current = sentence
+        else:
+            current = f"{current} {sentence}".strip()
+    if current:
+        chunks.append(current)
+
+    if (
+        len(chunks) > 1
+        and len(chunks[-1]) < min_chars
+        and len(chunks[-2]) + 1 + len(chunks[-1]) <= max_chars * 5 // 4
+    ):
+        chunks[-2:] = [f"{chunks[-2]} {chunks[-1]}"]
+    return chunks