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

mstar/api_server/openai/adapters.py

@@ -185,6 +185,16 @@ 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
+
     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

@@ -3,47 +3,125 @@
 Non-streaming returns the full audio as a container blob (WAV by default).
 Streaming returns a single open-ended WAV response (header + PCM16 frames) as
 the audio is produced.
+
+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.
+
+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.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))
+    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
 
 
 async def create_speech(api, model_name, adapter, req, raw_request=None):  # noqa: ARG001
     args = adapter.speech_to_request(req, api.upload_dir)
     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)))
     if req.stream:
+        pending = [submit(i) for i in range(min(lookahead, len(chunks)))]
+        # 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])
+        first = await anext(first_iter, None)
+        _raise_if_error(first)
         return StreamingResponse(
-            _stream_wav(api, request_id, sample_rate),
+            _stream_wav(api, submit, len(chunks), pending, first_iter, first, sample_rate),
             media_type="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] = []
+    pending: list[str] = [submit(i) for i in range(min(lookahead, len(chunks)))]
+    for index in range(len(chunks)):
+        if len(pending) < len(chunks):
+            pending.append(submit(len(pending)))
+        results = await api.collect_results(pending[index], raw_request)
+        pcm_parts.append(b"".join(c.data for c in results if c.modality == "audio"))
+    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_wav(api, request_id, sample_rate):
+def _raise_if_error(chunk) -> None:
+    """A data-worker failure arrives as an ``error`` chunk; turn it into the HTTP error it carries."""
+    if chunk is not None and chunk.modality == "error":
+        raise 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 _stream_wav(api, submit: Callable[[int], str], num_chunks: int, pending: list[str],
+                      first_iter, first, sample_rate: int):
     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
+    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
+        async for c in iterator:
+            # Mid-stream the status is already sent; closing the stream is the
+            # only honest signal left, so raise rather than end quietly.
+            _raise_if_error(c)
+            if c.modality == "audio" and c.data:
+                yield c.data

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