From ed54347e8167f18cd78054d158ba36b7166bb553 Mon Sep 17 00:00:00 2001 From: AirHua <78540954+AirHua-byte@users.noreply.github.com> Date: Mon, 27 Jul 2026 14:31:23 +0800 Subject: [PATCH] Fix runaway MLX Qwen audio chunks (#964) * fix runaway MLX Qwen audio chunks * test: tighten runaway retry coverage --------- Co-authored-by: huanghua01 --- backend/backends/__init__.py | 15 +++ backend/routes/generations.py | 14 ++- backend/services/generation.py | 22 +++- backend/tests/test_qwen_runaway_retry.py | 117 ++++++++++++++++++ backend/utils/audio.py | 37 ++++++ backend/utils/chunked_tts.py | 82 +++++++++--- .../content/docs/developer/tts-generation.mdx | 4 +- 7 files changed, 268 insertions(+), 23 deletions(-) create mode 100644 backend/tests/test_qwen_runaway_retry.py diff --git a/backend/backends/__init__.py b/backend/backends/__init__.py index 2437a87b..142c430c 100644 --- a/backend/backends/__init__.py +++ b/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] diff --git a/backend/routes/generations.py b/backend/routes/generations.py index 215c96cb..fbbeece6 100644 --- a/backend/routes/generations.py +++ b/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 diff --git a/backend/services/generation.py b/backend/services/generation.py index ce8fe93c..a4b2e8a3 100644 --- a/backend/services/generation.py +++ b/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 diff --git a/backend/tests/test_qwen_runaway_retry.py b/backend/tests/test_qwen_runaway_retry.py new file mode 100644 index 00000000..650385f8 --- /dev/null +++ b/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] diff --git a/backend/utils/audio.py b/backend/utils/audio.py index 7e0fd6fd..bccab2c8 100644 --- a/backend/utils/audio.py +++ b/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, diff --git a/backend/utils/chunked_tts.py b/backend/utils/chunked_tts.py index 1f43379e..7dbefdb7 100644 --- a/backend/utils/chunked_tts.py +++ b/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 diff --git a/docs/content/docs/developer/tts-generation.mdx b/docs/content/docs/developer/tts-generation.mdx index 5b902755..0f685622 100644 --- a/docs/content/docs/developer/tts-generation.mdx +++ b/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.