Fix runaway MLX Qwen audio chunks (#964)

* fix runaway MLX Qwen audio chunks

* test: tighten runaway retry coverage

---------

Co-authored-by: huanghua01 <[email protected]>

backend/backends/__init__.py

@@ -56,6 +56,7 @@ class ModelConfig:
     model_size: str = "default"
     size_mb: int = 0
     needs_trim: bool = False
+    retries_runaway: bool = False
     supports_instruct: bool = False
     languages: list[str] = field(default_factory=lambda: ["en"])
 
@@ -232,6 +233,10 @@ def _get_qwen_model_configs() -> list[ModelConfig]:
         repo_1_7b = "Qwen/Qwen3-TTS-12Hz-1.7B-Base"
         repo_0_6b = "Qwen/Qwen3-TTS-12Hz-0.6B-Base"
 
+    # mlx-audio can continue after an EOS miss with silence followed by
+    # codec noise. Retry only the affected text as smaller chunks.
+    retries_runaway = backend_type == "mlx"
+
     return [
         ModelConfig(
             model_name="qwen-tts-1.7B",
@@ -240,6 +245,7 @@ def _get_qwen_model_configs() -> list[ModelConfig]:
             hf_repo_id=repo_1_7b,
             model_size="1.7B",
             size_mb=3500,
+            retries_runaway=retries_runaway,
             supports_instruct=False,  # Base model drops instruct silently
             languages=["zh", "en", "ja", "ko", "de", "fr", "ru", "pt", "es", "it"],
         ),
@@ -250,6 +256,7 @@ def _get_qwen_model_configs() -> list[ModelConfig]:
             hf_repo_id=repo_0_6b,
             model_size="0.6B",
             size_mb=1200,
+            retries_runaway=retries_runaway,
             supports_instruct=False,
             languages=["zh", "en", "ja", "ko", "de", "fr", "ru", "pt", "es", "it"],
         ),
@@ -504,6 +511,14 @@ def engine_needs_trim(engine: str) -> bool:
     return False
 
 
+def engine_retries_runaway(engine: str) -> bool:
+    """Whether unstable output should be retried in smaller chunks."""
+    for cfg in get_tts_model_configs():
+        if cfg.engine == engine:
+            return cfg.retries_runaway
+    return False
+
+
 def engine_has_model_sizes(engine: str) -> bool:
     """Whether this engine supports multiple model sizes (only Qwen currently)."""
     configs = [c for c in get_tts_model_configs() if c.engine == engine]

backend/routes/generations.py

@@ -321,7 +321,13 @@ async def stream_speech(
     db: Session = Depends(get_db),
 ):
     """Generate speech and stream the WAV audio directly without saving to disk."""
-    from ..backends import get_tts_backend_for_engine, ensure_model_cached_or_raise, load_engine_model, engine_needs_trim
+    from ..backends import (
+        engine_needs_trim,
+        engine_retries_runaway,
+        ensure_model_cached_or_raise,
+        get_tts_backend_for_engine,
+        load_engine_model,
+    )
 
     profile = await profiles.get_profile(data.profile_id, db)
     if not profile:
@@ -347,10 +353,15 @@ async def stream_speech(
     from ..utils.chunked_tts import generate_chunked
 
     trim_fn = None
+    runaway_detector = None
     if engine_needs_trim(engine):
         from ..utils.audio import trim_tts_output
 
         trim_fn = trim_tts_output
+    if engine_retries_runaway(engine):
+        from ..utils.audio import has_tts_runaway
+
+        runaway_detector = has_tts_runaway
 
     audio, sample_rate = await generate_chunked(
         tts_model,
@@ -362,6 +373,7 @@ async def stream_speech(
         max_chunk_chars=data.max_chunk_chars,
         crossfade_ms=data.crossfade_ms,
         trim_fn=trim_fn,
+        runaway_detector=runaway_detector,
     )
 
     effects_chain_config = None

backend/services/generation.py

@@ -48,9 +48,14 @@ async def run_generation(
     This is the single entry point for all background generation work.
     It is designed to be enqueued via ``services.task_queue.enqueue_generation``.
     """
-    from ..backends import load_engine_model, get_tts_backend_for_engine, engine_needs_trim
+    from ..backends import (
+        engine_needs_trim,
+        engine_retries_runaway,
+        get_tts_backend_for_engine,
+        load_engine_model,
+    )
     from ..utils.chunked_tts import generate_chunked
-    from ..utils.audio import normalize_audio, save_audio, trim_tts_output
+    from ..utils.audio import has_tts_runaway, normalize_audio, save_audio, trim_tts_output
 
     task_manager = get_task_manager()
     bg_db = next(get_db())
@@ -72,12 +77,14 @@ async def run_generation(
 
         await history.update_generation_status(generation_id, "generating", bg_db)
         trim_fn = trim_tts_output if engine_needs_trim(engine) else None
+        runaway_detector = has_tts_runaway if engine_retries_runaway(engine) else None
 
         gen_kwargs: dict = dict(
             language=language,
             seed=seed if mode != "regenerate" else None,
             instruct=instruct,
             trim_fn=trim_fn,
+            runaway_detector=runaway_detector,
         )
         if max_chunk_chars is not None:
             gen_kwargs["max_chunk_chars"] = max_chunk_chars
@@ -267,9 +274,14 @@ async def generate_audio_sync(
     normalize, then encodes in-memory via :func:`tts.audio_to_wav_bytes`
     (same helper ``/generate/stream`` uses).
     """
-    from ..backends import load_engine_model, get_tts_backend_for_engine, engine_needs_trim
+    from ..backends import (
+        engine_needs_trim,
+        engine_retries_runaway,
+        get_tts_backend_for_engine,
+        load_engine_model,
+    )
     from ..utils.chunked_tts import generate_chunked
-    from ..utils.audio import normalize_audio, trim_tts_output
+    from ..utils.audio import has_tts_runaway, normalize_audio, trim_tts_output
     from . import tts
 
     bg_db = next(get_db())
@@ -287,12 +299,14 @@ async def generate_audio_sync(
         bg_db.close()
 
     trim_fn = trim_tts_output if engine_needs_trim(engine) else None
+    runaway_detector = has_tts_runaway if engine_retries_runaway(engine) else None
 
     gen_kwargs: dict = dict(
         language=language,
         seed=seed,
         instruct=instruct,
         trim_fn=trim_fn,
+        runaway_detector=runaway_detector,
     )
     if max_chunk_chars is not None:
         gen_kwargs["max_chunk_chars"] = max_chunk_chars

backend/tests/test_qwen_runaway_retry.py

@@ -0,0 +1,117 @@
+"""Regression coverage for runaway MLX Qwen TTS output."""
+
+from unittest.mock import patch
+
+import numpy as np
+import pytest
+
+from backend.backends import engine_needs_trim, engine_retries_runaway
+from backend.utils.audio import has_tts_runaway
+from backend.utils.chunked_tts import generate_chunked
+
+SAMPLE_RATE = 1000
+
+
+def test_mlx_qwen_enables_runaway_retry_without_aggressive_trim():
+    with patch("backend.backends.get_backend_type", return_value="mlx"):
+        assert engine_needs_trim("qwen") is False
+        assert engine_retries_runaway("qwen") is True
+
+
+def test_pytorch_qwen_keeps_runaway_retry_disabled():
+    with patch("backend.backends.get_backend_type", return_value="pytorch"):
+        assert engine_needs_trim("qwen") is False
+        assert engine_retries_runaway("qwen") is False
+
+
+def test_detector_flags_long_internal_silence():
+    speech = np.full(2 * SAMPLE_RATE, 0.2, dtype=np.float32)
+    runaway_gap = np.zeros(2500, dtype=np.float32)
+    hallucinated_noise = np.full(2 * SAMPLE_RATE, 0.8, dtype=np.float32)
+    audio = np.concatenate([speech, runaway_gap, hallucinated_noise])
+
+    assert has_tts_runaway(audio, SAMPLE_RATE) is True
+
+
+def test_detector_ignores_normal_internal_pause():
+    speech = np.full(SAMPLE_RATE, 0.2, dtype=np.float32)
+    normal_pause = np.zeros(1200, dtype=np.float32)
+    audio = np.concatenate([speech, normal_pause, speech])
+
+    assert has_tts_runaway(audio, SAMPLE_RATE) is False
+
+
+def test_trailing_silence_is_not_a_runaway():
+    speech = np.full(SAMPLE_RATE, 0.2, dtype=np.float32)
+    trailing_silence = np.zeros(2 * SAMPLE_RATE, dtype=np.float32)
+
+    assert (
+        has_tts_runaway(
+            np.concatenate([speech, trailing_silence]),
+            SAMPLE_RATE,
+        )
+        is False
+    )
+
+
+@pytest.mark.asyncio
+async def test_runaway_chunk_is_retried_as_smaller_chunks():
+    class FakeBackend:
+        def __init__(self):
+            self.calls = []
+
+        async def generate(self, text, *_args):
+            self.calls.append(text)
+            if len(text) > 200:
+                speech = np.full(SAMPLE_RATE, 0.2, dtype=np.float32)
+                silence = np.zeros(2500, dtype=np.float32)
+                noise = np.full(SAMPLE_RATE, 0.8, dtype=np.float32)
+                return np.concatenate([speech, silence, noise]), SAMPLE_RATE
+            return np.full(SAMPLE_RATE, 0.2, dtype=np.float32), SAMPLE_RATE
+
+    backend = FakeBackend()
+    text = f"{'A' * 119}. {'B' * 119}."
+
+    audio, sample_rate = await generate_chunked(
+        backend,
+        text,
+        {},
+        max_chunk_chars=800,
+        crossfade_ms=50,
+        runaway_detector=has_tts_runaway,
+    )
+
+    assert sample_rate == SAMPLE_RATE
+    assert backend.calls == [text, f"{'A' * 119}.", f"{'B' * 119}."]
+    assert len(audio) == 1950
+
+
+@pytest.mark.asyncio
+async def test_persistent_runaway_fails_instead_of_returning_corrupt_audio():
+    class AlwaysRunawayBackend:
+        def __init__(self):
+            self.calls = []
+
+        async def generate(self, text, *_args):
+            self.calls.append(text)
+            speech = np.full(SAMPLE_RATE, 0.2, dtype=np.float32)
+            silence = np.zeros(2500, dtype=np.float32)
+            noise = np.full(SAMPLE_RATE, 0.8, dtype=np.float32)
+            return np.concatenate([speech, silence, noise]), SAMPLE_RATE
+
+    backend = AlwaysRunawayBackend()
+    text = f"{'A' * 119}. {'B' * 119}."
+
+    with pytest.raises(
+        RuntimeError,
+        match="remained unstable after retrying smaller text chunks",
+    ):
+        await generate_chunked(
+            backend,
+            text,
+            {},
+            max_chunk_chars=800,
+            runaway_detector=has_tts_runaway,
+        )
+
+    assert [len(call) for call in backend.calls] == [241, 120, 100]

backend/utils/audio.py

@@ -110,6 +110,43 @@ def save_audio(
         raise OSError(f"Failed to save audio to {path}: {e}") from e
 
 
+def has_tts_runaway(
+    audio: np.ndarray,
+    sample_rate: int = 24000,
+    frame_ms: int = 20,
+    silence_threshold_db: float = -40.0,
+    max_internal_silence_ms: int = 2000,
+) -> bool:
+    """Detect speech followed by a long silence and then more output.
+
+    This shape is a reliable signal that a TTS model missed EOS and resumed
+    with hallucinated speech or codec noise. Leading and trailing silence do
+    not count because they are not bounded by non-silent audio.
+    """
+    frame_len = int(sample_rate * frame_ms / 1000)
+    if frame_len == 0 or len(audio) < frame_len:
+        return False
+
+    n_frames = len(audio) // frame_len
+    threshold_linear = 10 ** (silence_threshold_db / 20)
+    max_silence_frames = int(max_internal_silence_ms / frame_ms)
+    seen_speech = False
+    consecutive_silence = 0
+
+    for i in range(n_frames):
+        frame = audio[i * frame_len : (i + 1) * frame_len]
+        is_speech = np.sqrt(np.mean(frame**2)) >= threshold_linear
+        if is_speech:
+            if seen_speech and consecutive_silence >= max_silence_frames:
+                return True
+            seen_speech = True
+            consecutive_silence = 0
+        elif seen_speech:
+            consecutive_silence += 1
+
+    return False
+
+
 def trim_tts_output(
     audio: np.ndarray,
     sample_rate: int = 24000,

backend/utils/chunked_tts.py

@@ -20,6 +20,8 @@ logger = logging.getLogger("voicebox.chunked-tts")
 # Default chunk size in characters.  Can be overridden per-request via
 # the ``max_chunk_chars`` field on GenerationRequest.
 DEFAULT_MAX_CHUNK_CHARS = 800
+MAX_RUNAWAY_RETRIES = 2
+MIN_RUNAWAY_RETRY_CHARS = 100
 
 # Common abbreviations that should NOT be treated as sentence endings.
 # Lowercase for case-insensitive matching.
@@ -211,6 +213,7 @@ async def generate_chunked(
     max_chunk_chars: int = DEFAULT_MAX_CHUNK_CHARS,
     crossfade_ms: int = 50,
     trim_fn=None,
+    runaway_detector=None,
 ) -> Tuple[np.ndarray, int]:
     """Generate audio with automatic chunking for long text.
 
@@ -239,25 +242,75 @@ async def generate_chunked(
         Optional ``(audio, sample_rate) -> audio`` post-processing
         function applied to each chunk before concatenation (e.g.
         ``trim_tts_output`` for Chatterbox engines).
+    runaway_detector : callable | None
+        Optional ``(audio, sample_rate) -> bool`` detector. When it flags
+        unstable output, the affected text is split in half and retried.
 
     Returns
     -------
     (audio, sample_rate) : Tuple[np.ndarray, int]
     """
+    async def generate_one(
+        chunk_text: str,
+        chunk_seed: int | None,
+        retry_depth: int = 0,
+    ) -> tuple[np.ndarray, int]:
+        chunk_audio, chunk_sr = await backend.generate(
+            chunk_text,
+            voice_prompt,
+            language,
+            chunk_seed,
+            instruct,
+        )
+
+        if runaway_detector is not None and runaway_detector(chunk_audio, chunk_sr):
+            if retry_depth >= MAX_RUNAWAY_RETRIES or len(chunk_text) <= MIN_RUNAWAY_RETRY_CHARS:
+                raise RuntimeError(
+                    "TTS output remained unstable after retrying smaller text chunks"
+                )
+
+            retry_max_chars = max(MIN_RUNAWAY_RETRY_CHARS, len(chunk_text) // 2)
+            retry_chunks = split_text_into_chunks(chunk_text, retry_max_chars)
+            if len(retry_chunks) <= 1:
+                raise RuntimeError("Unable to split unstable TTS output for retry")
+
+            logger.warning(
+                "Detected unstable TTS output for %d chars; retrying as %d smaller chunks",
+                len(chunk_text),
+                len(retry_chunks),
+            )
+            retry_audio: list[np.ndarray] = []
+            for i, retry_text in enumerate(retry_chunks):
+                retry_seed = (
+                    chunk_seed + ((retry_depth + 1) * 1000) + i
+                    if chunk_seed is not None
+                    else None
+                )
+                audio, sample_rate = await generate_one(
+                    retry_text,
+                    retry_seed,
+                    retry_depth + 1,
+                )
+                retry_audio.append(np.asarray(audio, dtype=np.float32))
+
+            return (
+                concatenate_audio_chunks(
+                    retry_audio,
+                    sample_rate,
+                    crossfade_ms=crossfade_ms,
+                ),
+                sample_rate,
+            )
+
+        if trim_fn is not None:
+            chunk_audio = trim_fn(chunk_audio, chunk_sr)
+        return np.asarray(chunk_audio, dtype=np.float32), chunk_sr
+
     chunks = split_text_into_chunks(text, max_chunk_chars)
 
     if len(chunks) <= 1:
         # Short text — single-shot fast path
-        audio, sample_rate = await backend.generate(
-            text,
-            voice_prompt,
-            language,
-            seed,
-            instruct,
-        )
-        if trim_fn is not None:
-            audio = trim_fn(audio, sample_rate)
-        return audio, sample_rate
+        return await generate_one(text, seed)
 
     # Long text — chunked generation
     logger.info(
@@ -281,17 +334,12 @@ async def generate_chunked(
         # always produces the same output.
         chunk_seed = (seed + i) if seed is not None else None
 
-        chunk_audio, chunk_sr = await backend.generate(
+        chunk_audio, chunk_sr = await generate_one(
             chunk_text,
-            voice_prompt,
-            language,
             chunk_seed,
-            instruct,
         )
-        if trim_fn is not None:
-            chunk_audio = trim_fn(chunk_audio, chunk_sr)
 
-        audio_chunks.append(np.asarray(chunk_audio, dtype=np.float32))
+        audio_chunks.append(chunk_audio)
         if sample_rate is None:
             sample_rate = chunk_sr
 

docs/content/docs/developer/tts-generation.mdx

@@ -49,6 +49,7 @@ class ModelConfig:
     model_size: str = "default"
     size_mb: int = 0
     needs_trim: bool = False
+    retries_runaway: bool = False
     supports_instruct: bool = False
     languages: list[str] = field(default_factory=lambda: ["en"])
 ```
@@ -59,6 +60,7 @@ Registry helpers in `backends/__init__.py` replace what used to be per-engine `i
 - `get_tts_model_configs()` — only TTS variants
 - `get_model_config(model_name)` — lookup by name
 - `engine_needs_trim(engine)` — whether output should run through `trim_tts_output()`
+- `engine_retries_runaway(engine)` — whether unstable output should be retried as smaller chunks
 - `load_engine_model(engine, model_size)` — downloads + loads, handles engines with multiple sizes
 - `get_tts_backend_for_engine(engine)` — thread-safe backend factory with double-checked locking
 
@@ -152,7 +154,7 @@ The request path from frontend to audio file:
 
 6. **Inference** — the engine's `generate()` returns `(audio_array, sample_rate)`.
 
-7. **Post-process** — if `engine_needs_trim(engine)` is True, `trim_tts_output()` strips trailing silence. Effects chains (if any) are applied per generation version, not the clean version.
+7. **Validate and post-process** — engines with `retries_runaway=True` retry unstable output as smaller chunks. If `engine_needs_trim(engine)` is True, `trim_tts_output()` strips trailing silence. Effects chains (if any) are applied per generation version, not the clean version.
 
 8. **Persist** — audio is written to the generations directory, a row is inserted into the `generations` table, and the response includes the generation metadata.