From ec443b8db6a4b7a592fbd6151b08a80d3f3de544 Mon Sep 17 00:00:00 2001 From: Mustafa Alammar Date: Fri, 24 Jul 2026 23:34:31 -0700 Subject: [PATCH 1/4] fix: bound the voice conditioning cache to stop it exhausting VRAM MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit _conds_cache is a plain dict keyed by (path, mtime, exaggeration) and cleared only by unload_model(), so every distinct reference voice pins GPU tensors for the lifetime of the process. Each entry costs ~175 MB of VRAM. Measured on a P106-100 (5.93 GiB) with chatterbox-turbo, three 10-minute soaks cycling 28 reference voices at ~45 req/min: cache size requests failed headroom unbounded 453 406 0.11 GiB 2 89 1 0.10-0.49 GiB 0 88 0 0.28-0.40 GiB Live tensors climbed 4.09 -> 5.46 GiB and stayed there; the card then failed 2 MiB allocations while holding ~260 MB reserved-but-free. The same load with a single voice ran clean, which isolates the cache rather than allocator fragmentation. Larger cards are not immune, only slower to get there: /upload_reference stores each uploaded file under its own path, so every distinct voice a deployment ever serves adds a permanent entry. Adds an LRU bound via CONDS_CACHE_MAX (default 4, ~700 MB). A cache hit now refreshes recency, so a rotating workload evicts least-recently-used rather than by insertion order. Set 0 to disable caching entirely — a miss only re-encodes the reference audio, measured at +0.14s per request, a fixed cost that is ~13% on an 84-character request and within noise by 240 characters. Co-Authored-By: Claude Opus 5 --- engine.py | 44 ++++++++++++++++++++++++++++++++++++++++++-- 1 file changed, 42 insertions(+), 2 deletions(-) diff --git a/engine.py b/engine.py index a2be1e9..f6b5104 100644 --- a/engine.py +++ b/engine.py @@ -5,6 +5,7 @@ import logging import os import random +from collections import OrderedDict import numpy as np import torch from typing import Optional, Tuple @@ -114,7 +115,43 @@ def _resolve_bf16_setting() -> bool: # Voice conditioning cache: avoids re-encoding the same voice file on every request. # Key: (resolved_path, file_mtime, exaggeration) — mtime invalidates if file changes. -_conds_cache: dict = {} +# +# Bounded, because each entry pins GPU tensors for the process lifetime and this +# was previously an unbounded dict cleared only by unload_model(). Measured on a +# P106-100 (5.93 GiB) with chatterbox-turbo: cycling 28 reference voices grew +# live tensors from 4.09 to 5.46 GiB and left 0.11 GiB of headroom, after which +# 406 of 453 requests failed. The same load with one voice ran clean. +# +# This is not only a small-card concern. /upload_reference stores each uploaded +# file under its own path, so every distinct voice a deployment ever serves mints +# a permanent key — a larger card buys a higher ceiling, not a different outcome. +# +# CONDS_CACHE_MAX bounds the entry count. Measured cost is ~175 MB of VRAM per +# entry (live tensors grew 1.37 GiB across exactly 8 cached voices), so this is a +# far more expensive cache than it looks — budget it against free VRAM, not +# against an entry count that sounds small. +# +# The default of 4 costs ~700 MB, which leaves a 6 GB card ~1.1 GiB for +# generation after the 4.09 GiB model. On a 12 GB card 16+ is comfortable. Set 0 +# to disable caching and re-encode the voice on every request. +_CONDS_CACHE_MAX: int = max(0, int(os.environ.get("CONDS_CACHE_MAX", "4"))) +_conds_cache: "OrderedDict[tuple, object]" = OrderedDict() + + +def _conds_cache_store(key: tuple, conds) -> None: + """Insert under an LRU bound, evicting the least recently used entries. + + Eviction drops the last reference to that entry's tensors; the caching + allocator reuses the freed blocks, so VRAM is recovered without an explicit + empty_cache() on the request path. + """ + if _CONDS_CACHE_MAX == 0: + return + _conds_cache[key] = conds + _conds_cache.move_to_end(key) + while len(_conds_cache) > _CONDS_CACHE_MAX: + evicted_key, _ = _conds_cache.popitem(last=False) + logger.debug(f"Voice cache evicted (LRU, max={_CONDS_CACHE_MAX}): {evicted_key[0]}") def _conds_cache_key(path: str, exaggeration: float) -> tuple: @@ -481,6 +518,9 @@ def synthesize( conds_key = _conds_cache_key(audio_prompt_path, ex_for_key) if conds_key in _conds_cache: chatterbox_model.conds = _conds_cache[conds_key] + # Refresh recency, or the LRU bound would evict by insertion order + # and throw away the hottest voice under a rotating workload. + _conds_cache.move_to_end(conds_key) effective_prompt = None # conds already set, skip prepare_conditionals logger.debug(f"Voice cache hit: {audio_prompt_path}") @@ -509,7 +549,7 @@ def synthesize( # Store conds in cache after first compute for this voice. if conds_key is not None and effective_prompt is not None: if chatterbox_model.conds is not None: - _conds_cache[conds_key] = chatterbox_model.conds + _conds_cache_store(conds_key, chatterbox_model.conds) logger.debug(f"Cached voice conditionals for: {audio_prompt_path}") # The ChatterboxTTS.generate method already returns a CPU tensor. From a7a6eb00cf004b2d80069eaae73376f4f271f34c Mon Sep 17 00:00:00 2001 From: Mustafa Alammar Date: Wed, 29 Jul 2026 21:50:41 -0700 Subject: [PATCH 2/4] fix: make the voice cache LRU thread-safe MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit synthesize() runs concurrently — server.py dispatches it through loop.run_in_executor() — so several threads touch _conds_cache at once. The cache read was two non-atomic steps (membership check, then index) and the write three (insert, refresh, evict), so a thread could evict a key between another thread's check and its use and raise KeyError. The unbounded version could not hit this: nothing was ever removed, so a key that existed stayed. Adding eviction introduced the failure mode, which makes this a regression in the LRU commit rather than a pre-existing bug. Serializes cache access behind a lock and replaces the check-then-index read with a single locked lookup that also refreshes recency. Verified with 8 threads x 4000 operations over 12 voices at cap 4 (constant eviction): 0 errors, bound held. Also stops a bad CONDS_CACHE_MAX from breaking module import — int() on garbage raised ValueError at import time, which would have failed the server at startup on a typo. Now warns and falls back to the default. Co-Authored-By: Claude Opus 5 --- engine.py | 61 ++++++++++++++++++++++++++++++++++++++++++++----------- 1 file changed, 49 insertions(+), 12 deletions(-) diff --git a/engine.py b/engine.py index f6b5104..fc48086 100644 --- a/engine.py +++ b/engine.py @@ -5,6 +5,7 @@ import logging import os import random +import threading from collections import OrderedDict import numpy as np import torch @@ -134,9 +135,41 @@ def _resolve_bf16_setting() -> bool: # The default of 4 costs ~700 MB, which leaves a 6 GB card ~1.1 GiB for # generation after the 4.09 GiB model. On a 12 GB card 16+ is comfortable. Set 0 # to disable caching and re-encode the voice on every request. -_CONDS_CACHE_MAX: int = max(0, int(os.environ.get("CONDS_CACHE_MAX", "4"))) +def _resolve_conds_cache_max() -> int: + """Read CONDS_CACHE_MAX, tolerating garbage rather than failing to import.""" + raw = os.environ.get("CONDS_CACHE_MAX", "4").strip() + try: + return max(0, int(raw)) + except ValueError: + logger.warning( + f"CONDS_CACHE_MAX={raw!r} is not an integer; using the default of 4." + ) + return 4 + + +_CONDS_CACHE_MAX: int = _resolve_conds_cache_max() _conds_cache: "OrderedDict[tuple, object]" = OrderedDict() +# synthesize() runs concurrently: server.py dispatches it through +# loop.run_in_executor(), so several threads can touch this cache at once. +# Reads are two steps (look up, then refresh recency) and writes are three +# (insert, refresh, evict), none of which are atomic — without this lock a +# thread can evict a key between another thread's lookup and its use, raising +# KeyError. The unbounded version could not hit this because nothing was ever +# removed. +_conds_cache_lock = threading.Lock() + + +def _conds_cache_get(key: tuple): + """Return cached conds for `key` and mark it most-recently-used, else None.""" + if _CONDS_CACHE_MAX == 0: + return None + with _conds_cache_lock: + conds = _conds_cache.get(key) + if conds is not None: + _conds_cache.move_to_end(key) + return conds + def _conds_cache_store(key: tuple, conds) -> None: """Insert under an LRU bound, evicting the least recently used entries. @@ -147,11 +180,14 @@ def _conds_cache_store(key: tuple, conds) -> None: """ if _CONDS_CACHE_MAX == 0: return - _conds_cache[key] = conds - _conds_cache.move_to_end(key) - while len(_conds_cache) > _CONDS_CACHE_MAX: - evicted_key, _ = _conds_cache.popitem(last=False) - logger.debug(f"Voice cache evicted (LRU, max={_CONDS_CACHE_MAX}): {evicted_key[0]}") + with _conds_cache_lock: + _conds_cache[key] = conds + _conds_cache.move_to_end(key) + while len(_conds_cache) > _CONDS_CACHE_MAX: + evicted_key, _ = _conds_cache.popitem(last=False) + logger.debug( + f"Voice cache evicted (LRU, max={_CONDS_CACHE_MAX}): {evicted_key[0]}" + ) def _conds_cache_key(path: str, exaggeration: float) -> tuple: @@ -516,11 +552,11 @@ def synthesize( if audio_prompt_path and hasattr(chatterbox_model, "conds"): ex_for_key = 0.0 if loaded_model_type == "turbo" else exaggeration conds_key = _conds_cache_key(audio_prompt_path, ex_for_key) - if conds_key in _conds_cache: - chatterbox_model.conds = _conds_cache[conds_key] - # Refresh recency, or the LRU bound would evict by insertion order - # and throw away the hottest voice under a rotating workload. - _conds_cache.move_to_end(conds_key) + # Single locked lookup that also refreshes recency — checking + # membership and then indexing would race with eviction. + cached_conds = _conds_cache_get(conds_key) + if cached_conds is not None: + chatterbox_model.conds = cached_conds effective_prompt = None # conds already set, skip prepare_conditionals logger.debug(f"Voice cache hit: {audio_prompt_path}") @@ -630,7 +666,8 @@ def reload_model() -> bool: MODEL_LOADED = False loaded_model_type = None loaded_model_class_name = None - _conds_cache.clear() + with _conds_cache_lock: + _conds_cache.clear() logger.info("Voice conditioning cache cleared.") # 3. Force Python Garbage Collection From 7194fb885dce3ba9069865088f5ab2a47a952ce5 Mon Sep 17 00:00:00 2001 From: Mustafa Alammar Date: Wed, 29 Jul 2026 22:45:04 -0700 Subject: [PATCH 3/4] feat: batch a request's chunks through turbo in one forward pass MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit /tts already splits long text into chunks that all share one voice, then synthesized them one at a time. Turbo's decode loop (T3.inference_turbo) is written with batch dimensions throughout and uses HuggingFace logits processors, which are batch-correct, so those chunks can go through a single forward pass. Measured on a P106-100 with chatterbox-turbo, four chunks, both paths warm: sequential 5.12s -> batched 2.62s = 1.96x T3 decode alone scales further, and costs almost no extra VRAM because the weights dominate and only the KV cache grows: batch T3 time s/item speedup peak VRAM 1 0.98s 0.976 1.00x 4.50 GiB 2 1.15s 0.575 1.70x 4.62 GiB 4 1.75s 0.438 2.23x 4.84 GiB 8 2.76s 0.346 2.82x 5.29 GiB TTS_BATCH_SIZE caps the batch (default 4; batch 8 leaves only ~0.6 GiB on a 6 GB card). synthesize_batch() returns (None, None) rather than erroring whenever batching would change behaviour, and /tts falls back to the per-chunk path: - model is not turbo (only inference_turbo accepts a batch) - a seed was requested: sequential chunks each re-seed from the same value while a batch draws all rows from one generator, so seeded output would stop being reproducible - a single chunk, or batching disabled Two details that are easy to get wrong: inference_turbo only breaks once EVERY row emits stop on the same step, so rows that finish early keep sampling. Those trailing tokens fall below the OOV limit and would survive turbo's own `tokens < 6561` filter as audible gibberish, so each row is truncated at its first stop token before filtering. The first batched call on a cold process costs ~10s while CUDA loads kernels for an unseen batch dimension (13.5s first call vs 2.8s steady state). The cost is one-time rather than per-shape — after warming at batch 2, fresh batch sizes of 3 and 4 ran at full speed — so warmup_batch() pays it during startup instead of letting the first real request absorb it. s3gen still runs per row: it accepted a batch in testing but only with equal-length token sequences, and real chunks differ in length. T3 decode dominates the cost anyway. Co-Authored-By: Claude Opus 5 --- engine.py | 197 +++++++++++++++++++++++++++++++++++++++++++++++++++++- server.py | 97 ++++++++++++++++++--------- 2 files changed, 262 insertions(+), 32 deletions(-) diff --git a/engine.py b/engine.py index fc48086..d352f22 100644 --- a/engine.py +++ b/engine.py @@ -6,10 +6,11 @@ import os import random import threading +import time from collections import OrderedDict import numpy as np import torch -from typing import Optional, Tuple +from typing import List, Optional, Tuple from pathlib import Path from chatterbox.tts import ChatterboxTTS # Main TTS engine class @@ -497,6 +498,200 @@ def load_model() -> bool: return False +# --- Chunk batching ----------------------------------------------------------- +# Turbo's decode loop (T3.inference_turbo) is written with batch dimensions +# throughout and uses HuggingFace logits processors, which are batch-correct. +# Measured on a P106-100, batching identical-voice chunks scales well and costs +# almost no VRAM, because the weights dominate and only the KV cache grows: +# +# batch T3 time s/item speedup peak VRAM +# 1 0.98s 0.976 1.00x 4.50 GiB +# 2 1.15s 0.575 1.70x 4.62 GiB +# 4 1.75s 0.438 2.23x 4.84 GiB +# 8 2.76s 0.346 2.82x 5.29 GiB +# +# /tts already splits long text into chunks that all share one voice, so those +# chunks can go through a single forward pass with no cross-request queueing and +# no added latency. TTS_BATCH_SIZE caps it; 4 is the defensible ceiling on a 6 GB +# card (batch 8 leaves only ~0.6 GiB, and real chunks are longer than the short +# test strings above). +def _resolve_batch_size() -> int: + raw = os.environ.get("TTS_BATCH_SIZE", "4").strip() + try: + return max(1, int(raw)) + except ValueError: + logger.warning(f"TTS_BATCH_SIZE={raw!r} is not an integer; using 4.") + return 4 + + +TTS_BATCH_SIZE: int = _resolve_batch_size() + +# Token ids at or above this are out-of-vocabulary for s3gen; turbo's own +# generate() filters on the same constant. +_S3GEN_TOKEN_LIMIT = 6561 + + +def batching_available(text_count: int, seed: int) -> bool: + """Whether synthesize_batch() can serve this request. + + Declines rather than silently changing behaviour when: + - the model is not turbo (only inference_turbo takes a batch), + - there is nothing to gain (one chunk, or batching disabled), + - a seed was requested. Sequential chunks each re-seed from the same value, + while a batch draws all rows from one generator, so seeded output would + stop being reproducible. Reproducibility wins over speed here. + """ + return ( + MODEL_LOADED + and chatterbox_model is not None + and loaded_model_type == "turbo" + and TTS_BATCH_SIZE > 1 + and text_count > 1 + and seed == 0 + ) + + +def _resolve_conds(audio_prompt_path: Optional[str], exaggeration: float): + """Set chatterbox_model.conds for this voice, using the cache when possible.""" + if not audio_prompt_path or not hasattr(chatterbox_model, "conds"): + return + ex_for_key = 0.0 if loaded_model_type == "turbo" else exaggeration + conds_key = _conds_cache_key(audio_prompt_path, ex_for_key) + cached = _conds_cache_get(conds_key) + if cached is not None: + chatterbox_model.conds = cached + logger.debug(f"Voice cache hit (batch): {audio_prompt_path}") + return + chatterbox_model.prepare_conditionals(audio_prompt_path, exaggeration=exaggeration) + if chatterbox_model.conds is not None: + _conds_cache_store(conds_key, chatterbox_model.conds) + logger.debug(f"Cached voice conditionals (batch) for: {audio_prompt_path}") + + +def warmup_batch(audio_prompt_path: Optional[str]) -> bool: + """Pay the one-time batched-kernel setup cost at startup instead of on a request. + + The first batched forward pass on a cold process costs ~10s on a P106-100 + while CUDA selects and loads kernels for a batch dimension the model has not + seen. Measured: 13.5s first call, then 2.8s steady state. The cost is one-time + and NOT per-shape — after warming at batch 2, fresh batch sizes of 3 and 4 ran + at full speed — so a single small warmup covers every later batch. + + Without this the first real multi-chunk request eats the whole penalty, which + would read as a latency regression rather than a speedup. + """ + if not batching_available(text_count=2, seed=0) or not audio_prompt_path: + return False + try: + started = time.time() + wavs, _ = synthesize_batch( + texts=["Warming up the batched path.", "Second row for the batch."], + audio_prompt_path=audio_prompt_path, + ) + if wavs is None: + logger.warning("Batch warmup did not run; first batched request will be slow.") + return False + logger.info(f"Batched path warmed in {time.time() - started:.1f}s.") + return True + except Exception as e: + logger.warning(f"Batch warmup failed ({e}); first batched request will be slow.") + return False + + +def synthesize_batch( + texts: List[str], + audio_prompt_path: Optional[str] = None, + temperature: float = 0.8, + exaggeration: float = 0.5, + cfg_weight: float = 0.5, + seed: int = 0, + language: str = "en", +) -> Tuple[Optional[List[torch.Tensor]], Optional[int]]: + """Synthesize several texts sharing one voice in batched forward passes. + + Returns (list of wav tensors aligned with `texts`, sample_rate), or + (None, None) if batching does not apply or fails — callers must fall back to + per-item synthesize() on None rather than surfacing an error. + """ + if not batching_available(len(texts), seed): + return None, None + + try: + # Imported here so a non-turbo install never needs these symbols. + from chatterbox.tts_turbo import punc_norm, S3GEN_SIL + + model = chatterbox_model + _resolve_conds(audio_prompt_path, exaggeration) + if getattr(model, "conds", None) is None: + logger.warning("Batch synthesis has no voice conditionals; falling back.") + return None, None + + stop_token = model.t3.hp.stop_speech_token + conds = model.conds + wavs: List[torch.Tensor] = [] + + for start in range(0, len(texts), TTS_BATCH_SIZE): + group = texts[start : start + TTS_BATCH_SIZE] + encoded = model.tokenizer( + [punc_norm(t) for t in group], + return_tensors="pt", + padding=True, + truncation=True, + ) + text_tokens = encoded.input_ids.to(model.device) + + with torch.autocast("cuda", dtype=torch.bfloat16, enabled=BF16_ENABLED): + all_tokens = model.t3.inference_turbo( + t3_cond=conds.t3, + text_tokens=text_tokens, + temperature=temperature, + ) + + for row_idx in range(all_tokens.size(0)): + row = all_tokens[row_idx] + # inference_turbo only breaks once EVERY row has emitted stop, so + # rows that finished early keep sampling. Those trailing tokens are + # below the OOV limit and would survive the filter as audible + # gibberish, so truncate at the first stop before filtering. + hit = (row == stop_token).nonzero() + if hit.numel(): + row = row[: hit[0, 0]] + row = row[row < _S3GEN_TOKEN_LIMIT] + if row.numel() == 0: + logger.warning( + f"Batch row {start + row_idx} produced no usable tokens; falling back." + ) + return None, None + + silence = torch.tensor( + [S3GEN_SIL] * 3, dtype=torch.long, device=model.device + ) + row = torch.cat([row, silence]) + + # s3gen runs per row: it accepted a batch in testing, but only with + # equal-length token sequences, and real chunks differ in length. + # Padding plus masking would be needed to batch this stage safely, + # and T3 decode dominates the cost anyway. + with torch.autocast("cuda", dtype=torch.bfloat16, enabled=BF16_ENABLED): + wav, _ = model.s3gen.inference( + speech_tokens=row, ref_dict=conds.gen, n_cfm_timesteps=2 + ) + wav_np = wav.squeeze(0).detach().cpu().numpy() + wav_np = model.watermarker.apply_watermark(wav_np, sample_rate=model.sr) + wavs.append(torch.from_numpy(wav_np).unsqueeze(0)) + + logger.info( + f"Batched {len(texts)} chunk(s) in " + f"{(len(texts) + TTS_BATCH_SIZE - 1) // TTS_BATCH_SIZE} pass(es) " + f"(batch size {TTS_BATCH_SIZE})." + ) + return wavs, model.sr + + except Exception as e: + logger.error(f"Batched synthesis failed, falling back: {e}", exc_info=True) + return None, None + + def synthesize( text: str, audio_prompt_path: Optional[str] = None, diff --git a/server.py b/server.py index e435185..d068ad4 100644 --- a/server.py +++ b/server.py @@ -159,6 +159,19 @@ async def lifespan(app: FastAPI): ) else: logger.info("TTS Model loaded successfully via engine.") + + # Pay the one-time batched-kernel setup cost now rather than on the + # first multi-chunk request, which would otherwise take ~10s longer + # than an unbatched one and look like a regression. No-op unless the + # loaded model supports batching. + try: + voices_dir = get_predefined_voices_path(ensure_absolute=True) + warm_voice = next(iter(sorted(voices_dir.glob("*.wav"))), None) + if warm_voice is not None: + engine.warmup_batch(str(warm_voice)) + except Exception as warm_exc: + logger.warning(f"Skipped batch warmup: {warm_exc}") + host_address = get_host() server_port = get_port() browser_thread = threading.Thread( @@ -1068,40 +1081,62 @@ async def _stream_generator(): ) # --- End streaming fork --- + # Resolve generation parameters once — the batched and per-chunk paths below + # both use these values. + _prompt_for_engine = ( + str(audio_prompt_path_for_engine) if audio_prompt_path_for_engine else None + ) + _temperature = ( + request.temperature + if request.temperature is not None + else get_gen_default_temperature() + ) + _exaggeration = ( + request.exaggeration + if request.exaggeration is not None + else get_gen_default_exaggeration() + ) + _cfg_weight = ( + request.cfg_weight + if request.cfg_weight is not None + else get_gen_default_cfg_weight() + ) + _seed = request.seed if request.seed is not None else get_gen_default_seed() + _language = ( + request.language if request.language is not None else get_gen_default_language() + ) + + # Chunk batching: every chunk of this request shares one voice, so turbo can + # generate them in a single forward pass (~2.2x at batch 4, measured on a + # P106-100). Returns None when batching does not apply — non-turbo model, a + # seeded request, or a single chunk — leaving the per-chunk path unchanged. + batched_wavs, batched_sr = engine.synthesize_batch( + texts=text_chunks, + audio_prompt_path=_prompt_for_engine, + temperature=_temperature, + exaggeration=_exaggeration, + cfg_weight=_cfg_weight, + seed=_seed, + language=_language, + ) + if batched_wavs is not None: + perf_monitor.record(f"Engine batch-synthesized {len(text_chunks)} chunk(s)") + for i, chunk in enumerate(text_chunks): logger.info(f"Synthesizing chunk {i+1}/{len(text_chunks)}...") try: - chunk_audio_tensor, chunk_sr_from_engine = engine.synthesize( - text=chunk, - audio_prompt_path=( - str(audio_prompt_path_for_engine) - if audio_prompt_path_for_engine - else None - ), - temperature=( - request.temperature - if request.temperature is not None - else get_gen_default_temperature() - ), - exaggeration=( - request.exaggeration - if request.exaggeration is not None - else get_gen_default_exaggeration() - ), - cfg_weight=( - request.cfg_weight - if request.cfg_weight is not None - else get_gen_default_cfg_weight() - ), - seed=( - request.seed if request.seed is not None else get_gen_default_seed() - ), - language=( - request.language - if request.language is not None - else get_gen_default_language() - ), - ) + if batched_wavs is not None: + chunk_audio_tensor, chunk_sr_from_engine = batched_wavs[i], batched_sr + else: + chunk_audio_tensor, chunk_sr_from_engine = engine.synthesize( + text=chunk, + audio_prompt_path=_prompt_for_engine, + temperature=_temperature, + exaggeration=_exaggeration, + cfg_weight=_cfg_weight, + seed=_seed, + language=_language, + ) perf_monitor.record(f"Engine synthesized chunk {i+1}") if chunk_audio_tensor is None or chunk_sr_from_engine is None: From d8925ddf52b28e5af259e0cf0780d34a8dd20d3f Mon Sep 17 00:00:00 2001 From: Mustafa Alammar Date: Wed, 29 Jul 2026 23:08:52 -0700 Subject: [PATCH 4/4] fix: close a conds race and an unguarded index in chunk batching MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Review of PR #161 found two defects before merge: 1. _resolve_conds() set chatterbox_model.conds and let the caller read it back via a second statement. That attribute is shared process-wide, so a concurrent request could reassign it between the two statements and the batch would generate in the wrong voice. Now returns the Conditionals object directly instead of round-tripping through the shared attribute. 2. synthesize_batch() could hand back a wavs list shorter than the input list (defensive case, not observed) and server.py indexed it unconditionally — an IndexError/500 instead of the intended fallback to per-chunk synthesis. Now refuses (returns None, None) on a length mismatch, and server.py's index is bounds-checked as a second line of defense. Re-verified on chattybox after the fix: 4-chunk batch still correct and 1.88x faster than sequential with both paths warm (4.99s -> 2.65s), all distinctness and duration-sanity checks pass. Co-Authored-By: Claude Sonnet 5 --- engine.py | 34 +++++++++++++++++++++++++--------- server.py | 2 +- 2 files changed, 26 insertions(+), 10 deletions(-) diff --git a/engine.py b/engine.py index d352f22..8170f4f 100644 --- a/engine.py +++ b/engine.py @@ -552,20 +552,29 @@ def batching_available(text_count: int, seed: int) -> bool: def _resolve_conds(audio_prompt_path: Optional[str], exaggeration: float): - """Set chatterbox_model.conds for this voice, using the cache when possible.""" + """Return the Conditionals for this voice, using the cache when possible. + + Returns the object rather than leaving the caller to read + chatterbox_model.conds back: that attribute is shared process-wide, so a + concurrent request can reassign it between the two statements and the batch + would generate in the wrong voice. Holding a local reference sidesteps that + entirely — the batch uses the conds it resolved, whatever else touches the + model meanwhile. + """ if not audio_prompt_path or not hasattr(chatterbox_model, "conds"): - return + return None ex_for_key = 0.0 if loaded_model_type == "turbo" else exaggeration conds_key = _conds_cache_key(audio_prompt_path, ex_for_key) cached = _conds_cache_get(conds_key) if cached is not None: - chatterbox_model.conds = cached logger.debug(f"Voice cache hit (batch): {audio_prompt_path}") - return + return cached chatterbox_model.prepare_conditionals(audio_prompt_path, exaggeration=exaggeration) - if chatterbox_model.conds is not None: - _conds_cache_store(conds_key, chatterbox_model.conds) + conds = chatterbox_model.conds + if conds is not None: + _conds_cache_store(conds_key, conds) logger.debug(f"Cached voice conditionals (batch) for: {audio_prompt_path}") + return conds def warmup_batch(audio_prompt_path: Optional[str]) -> bool: @@ -621,13 +630,12 @@ def synthesize_batch( from chatterbox.tts_turbo import punc_norm, S3GEN_SIL model = chatterbox_model - _resolve_conds(audio_prompt_path, exaggeration) - if getattr(model, "conds", None) is None: + conds = _resolve_conds(audio_prompt_path, exaggeration) + if conds is None: logger.warning("Batch synthesis has no voice conditionals; falling back.") return None, None stop_token = model.t3.hp.stop_speech_token - conds = model.conds wavs: List[torch.Tensor] = [] for start in range(0, len(texts), TTS_BATCH_SIZE): @@ -680,6 +688,14 @@ def synthesize_batch( wav_np = model.watermarker.apply_watermark(wav_np, sample_rate=model.sr) wavs.append(torch.from_numpy(wav_np).unsqueeze(0)) + # Callers index this list per chunk, so a short list would be an + # IndexError rather than a graceful fallback. Refuse instead. + if len(wavs) != len(texts): + logger.error( + f"Batch produced {len(wavs)} wav(s) for {len(texts)} chunk(s); falling back." + ) + return None, None + logger.info( f"Batched {len(texts)} chunk(s) in " f"{(len(texts) + TTS_BATCH_SIZE - 1) // TTS_BATCH_SIZE} pass(es) " diff --git a/server.py b/server.py index d068ad4..e194ee9 100644 --- a/server.py +++ b/server.py @@ -1125,7 +1125,7 @@ async def _stream_generator(): for i, chunk in enumerate(text_chunks): logger.info(f"Synthesizing chunk {i+1}/{len(text_chunks)}...") try: - if batched_wavs is not None: + if batched_wavs is not None and i < len(batched_wavs): chunk_audio_tensor, chunk_sr_from_engine = batched_wavs[i], batched_sr else: chunk_audio_tensor, chunk_sr_from_engine = engine.synthesize(