fix(backend): prevent unbounded memory accumulation over consecutive TTS generations (#923) - #1032
Open
devangkantharia wants to merge 1 commit into
Open
Conversation
Contributor
📝 WalkthroughWalkthroughThe change runs TTS inference under ChangesGeneration memory cleanup
Estimated code review effort: 3 (Moderate) | ~20 minutes Sequence Diagram(s)sequenceDiagram
participant GenerationService
participant TTSBackend
participant DeviceCache
GenerationService->>TTSBackend: Generate audio in inference mode
TTSBackend-->>GenerationService: Return waveform
GenerationService->>DeviceCache: Clear device cache during finalization
DeviceCache-->>GenerationService: Suppress cleanup errors
Possibly related issues
Possibly related PRs
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 4
🧹 Nitpick comments (2)
backend/backends/base.py (2)
186-187: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winVerify the XPU availability guard.
hasattr(torch, "xpu")only confirms that the namespace exists. PyTorch exposestorch.xpu.is_available()as the runtime check and documents XPU as lazily importable. (docs.pytorch.org) If"xpu"reaches this helper while no XPU is usable, the code still callstorch.xpu.empty_cache(). Add an availability check and anempty_cachemethod check, or enforce that callers pass only an available device.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@backend/backends/base.py` around lines 186 - 187, Update the XPU branch in the device-cache helper to verify torch.xpu is runtime-available via is_available() before clearing the cache, and also confirm empty_cache is callable before invoking it; do not rely on hasattr(torch, "xpu") alone.
182-190: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winVerify that cleanup does not block the async event loop.
gc.collect()and the device-cache calls execute synchronously.run_generationandgenerate_audio_synccall this helper fromasync deffunctions. A full collection can delay queued jobs and the synchronous/profiles/{id}/speakresponse. Offload cleanup to a worker thread or dedicated executor. Keeptask_manager.complete_generationafter cleanup if it releases queue capacity.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@backend/backends/base.py` around lines 182 - 190, Update the cleanup helper containing gc.collect and device cache calls so all synchronous cleanup work runs in a worker thread or dedicated executor when invoked by async run_generation and generate_audio_sync flows. Preserve the existing device-specific cleanup behavior, and ensure task_manager.complete_generation remains after cleanup so queue capacity is released only once cleanup finishes.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@backend/backends/chatterbox_backend.py`:
- Around line 206-221: Add a public device property to the Chatterbox backend
class in backend/backends/chatterbox_backend.py at lines 206-221, returning
_device when set and "cpu" otherwise. Add the same property to the corresponding
Chatterbox Turbo backend class in backend/backends/chatterbox_turbo_backend.py
at lines 187-201 so tts_model.device exposes the selected device for cache
cleanup.
In `@backend/services/generation.py`:
- Around line 159-164: Update both cleanup finalizers around empty_device_cache
to record cleanup failures instead of silently passing, including the engine and
resolved device in a warning or metric. Preserve the original generation
exception and successful generation result while ensuring cleanup observability
does not replace or re-raise the primary error.
- Around line 159-164: Add a public device property to both ChatterboxTTSBackend
and ChatterboxTurboTTSBackend that returns their configured _device value, so
generation.py finalizer cleanup receives the selected device instead of
defaulting to CPU.
- Around line 322-337: The generate_audio_sync flow currently starts its
finalizer after backend/model/prompt setup, so setup failures bypass
empty_device_cache. Move get_tts_backend_for_engine, load_engine_model, and
create_voice_prompt_for_profile into the same outer try/finally as
generate_chunked, while preserving bg_db.close() in its existing nested cleanup.
---
Nitpick comments:
In `@backend/backends/base.py`:
- Around line 186-187: Update the XPU branch in the device-cache helper to
verify torch.xpu is runtime-available via is_available() before clearing the
cache, and also confirm empty_cache is callable before invoking it; do not rely
on hasattr(torch, "xpu") alone.
- Around line 182-190: Update the cleanup helper containing gc.collect and
device cache calls so all synchronous cleanup work runs in a worker thread or
dedicated executor when invoked by async run_generation and generate_audio_sync
flows. Preserve the existing device-specific cleanup behavior, and ensure
task_manager.complete_generation remains after cleanup so queue capacity is
released only once cleanup finishes.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: d5514121-3c18-4848-8245-37d7c47d2dfc
📒 Files selected for processing (8)
backend/backends/base.pybackend/backends/chatterbox_backend.pybackend/backends/chatterbox_turbo_backend.pybackend/backends/kokoro_backend.pybackend/backends/luxtts_backend.pybackend/backends/pytorch_backend.pybackend/backends/qwen_custom_voice_backend.pybackend/services/generation.py
Comment on lines
+206
to
+221
| with torch.inference_mode(): | ||
| wav = self.model.generate( | ||
| text, | ||
| language_id=language, | ||
| audio_prompt_path=ref_audio, | ||
| exaggeration=lang_defaults["exaggeration"], | ||
| cfg_weight=lang_defaults["cfg_weight"], | ||
| temperature=lang_defaults["temperature"], | ||
| repetition_penalty=lang_defaults["repetition_penalty"], | ||
| ) | ||
|
|
||
| # Convert tensor -> numpy | ||
| if isinstance(wav, torch.Tensor): | ||
| audio = wav.squeeze().cpu().numpy().astype(np.float32) | ||
| else: | ||
| audio = np.asarray(wav, dtype=np.float32) |
Contributor
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Expose the selected device for both Chatterbox backends.
The generation finalizer reads tts_model.device, but both backends store the selected device only in _device. On CUDA or XPU, empty_device_cache receives "cpu" and does not clear the active device cache.
backend/backends/chatterbox_backend.py#L206-L221: add a publicdeviceproperty that returns_deviceor"cpu".backend/backends/chatterbox_turbo_backend.py#L187-L201: add the same publicdeviceproperty.
Proposed fix
class ChatterboxTTSBackend:
+ `@property`
+ def device(self) -> str:
+ return self._device or "cpu"
class ChatterboxTurboTTSBackend:
+ `@property`
+ def device(self) -> str:
+ return self._device or "cpu"📍 Affects 2 files
backend/backends/chatterbox_backend.py#L206-L221(this comment)backend/backends/chatterbox_turbo_backend.py#L187-L201
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@backend/backends/chatterbox_backend.py` around lines 206 - 221, Add a public
device property to the Chatterbox backend class in
backend/backends/chatterbox_backend.py at lines 206-221, returning _device when
set and "cpu" otherwise. Add the same property to the corresponding Chatterbox
Turbo backend class in backend/backends/chatterbox_turbo_backend.py at lines
187-201 so tts_model.device exposes the selected device for cache cleanup.
Comment on lines
+159
to
+164
| try: | ||
| from ..backends.base import empty_device_cache | ||
| device = getattr(tts_model, "device", "cpu") if "tts_model" in locals() else "cpu" | ||
| empty_device_cache(device) | ||
| except Exception: | ||
| pass |
Contributor
There was a problem hiding this comment.
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
Make cleanup failures observable.
Both finalizers catch Exception and execute pass. A successful generation can therefore be reported as complete even when cache cleanup failed. Log a warning or metric with the engine and device while preserving the original generation error.
Also applies to: 332-337
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@backend/services/generation.py` around lines 159 - 164, Update both cleanup
finalizers around empty_device_cache to record cleanup failures instead of
silently passing, including the engine and resolved device in a warning or
metric. Preserve the original generation exception and successful generation
result while ensuring cleanup observability does not replace or re-raise the
primary error.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
rg -n -C 4 'def device|self\._device|device = getattr' \
backend/backends/chatterbox_backend.py \
backend/backends/chatterbox_turbo_backend.py \
backend/backends/luxtts_backend.py \
backend/backends/kokoro_backend.py \
backend/services/generation.pyRepository: jamiepine/voicebox
Length of output: 10865
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== generation.py outlines =="
ast-grep outline backend/services/generation.py --view compact || true
echo
echo "== generation.py relevant sections =="
sed -n '110,175p' backend/services/generation.py
sed -n '300,345p' backend/services/generation.py
echo
echo "== backend device attributes =="
python3 - <<'PY'
import ast
from pathlib import Path
files = [
Path("backend/backends/chatterbox_backend.py"),
Path("backend/backends/chatterbox_turbo_backend.py"),
Path("backend/backends/luxtts_backend.py"),
Path("backend/backends/kokoro_backend.py"),
]
for path in files:
tree = ast.parse(path.read_text())
classes = [n for n in ast.walk(tree) if isinstance(n, ast.ClassDef)]
print(f"--- {path} classes: {[c.name for c in classes]}")
for cls in classes:
names = []
assigns = []
for node in ast.walk(cls):
if isinstance(node, ast.FunctionDef) and node.name == "device":
names.append(f"def {node.name}")
if isinstance(node, ast.Attribute) and isinstance(node.value, ast.Name) and node.value.id == "self":
names.append(node.attr)
if isinstance(node, ast.Assign):
for target in node.targets:
if isinstance(target, ast.Attribute) and isinstance(target.value, ast.Name) and target.value.id == "self":
assigns.append(target.attr)
print(" ", path.name, cls.name, sorted(set(names + assigns)))
print("== getattr None behavior ==")
class C:
pass
print(genattr(C()), getattr(C(), "device", "cpu"))
print(genattr(C), getattr(object(), "device", "cpu"))
PYRepository: jamiepine/voicebox
Length of output: 4989
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== generation.py outlines =="
ast-grep outline backend/services/generation.py --view compact || true
echo
echo "== generation.py relevant sections =="
sed -n '110,175p' backend/services/generation.py
sed -n '300,345p' backend/services/generation.py
echo
echo "== backend device attributes =="
python3 - <<'PY'
import ast
from pathlib import Path
files = [
Path("backend/backends/chatterbox_backend.py"),
Path("backend/backends/chatterbox_turbo_backend.py"),
Path("backend/backends/luxtts_backend.py"),
Path("backend/backends/kokoro_backend.py"),
]
for path in files:
tree = ast.parse(path.read_text())
classes = [n for n in ast.walk(tree) if isinstance(n, ast.ClassDef)]
print(f"--- {path} classes: {[c.name for c in classes]}")
for cls in classes:
names = []
assigns = []
for node in ast.walk(cls):
if isinstance(node, ast.FunctionDef) and node.name == "device":
names.append(f"def {node.name}")
if isinstance(node, ast.Attribute) and isinstance(node.value, ast.Name) and node.value.id == "self":
names.append(node.attr)
if isinstance(node, ast.Assign):
for target in node.targets:
if isinstance(target, ast.Attribute) and isinstance(target.value, ast.Name) and target.value.id == "self":
assigns.append(target.attr)
print(" ", path.name, cls.name, sorted(set(names + assigns)))
print("== getattr None behavior ==")
class C:
pass
print(genattr(C()), getattr(C(), "device", "cpu"))
print(genattr(C), getattr(C), getattr(object(), "device", "cpu"))
PYRepository: jamiepine/voicebox
Length of output: 4989
Add a public device property to the Chatterbox backends.
backend/services/generation.py reads tts_model.device in its finalizers; LuxTTS and Kokoro expose it, but ChatterboxTTSBackend and ChatterboxTurboTTSBackend only set _device. That defaults Chatterbox generation cleanup to CPU and can skip the selected device cache.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@backend/services/generation.py` around lines 159 - 164, Add a public device
property to both ChatterboxTTSBackend and ChatterboxTurboTTSBackend that returns
their configured _device value, so generation.py finalizer cleanup receives the
selected device instead of defaulting to CPU.
Comment on lines
+322
to
+337
| try: | ||
| audio, sample_rate = await generate_chunked( | ||
| tts_model, text, voice_prompt, **gen_kwargs | ||
| ) | ||
|
|
||
| if normalize: | ||
| audio = normalize_audio(audio) | ||
| if normalize: | ||
| audio = normalize_audio(audio) | ||
|
|
||
| return tts.audio_to_wav_bytes(audio, sample_rate) | ||
| return tts.audio_to_wav_bytes(audio, sample_rate) | ||
| finally: | ||
| try: | ||
| from ..backends.base import empty_device_cache | ||
| device = getattr(tts_model, "device", "cpu") if "tts_model" in locals() else "cpu" | ||
| empty_device_cache(device) | ||
| except Exception: | ||
| pass |
Contributor
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== locate generation.py =="
fd -a 'generation\.py$' . || true
echo "== file stats =="
wc -l backend/services/generation.py 2>/dev/null || true
echo "== relevant generation.py lines 260-360 =="
sed -n '260,360p' backend/services/generation.py 2>/dev/null | nl -ba -v260
echo "== search for empty_device_cache definitions/usages =="
rg -n "empty_device_cache|load_engine_model|get_tts_backend_for_engine|create_voice_prompt_for_profile|run_generation" backend/services/generation.py backend -S || true
echo "== outline generation.py =="
ast-grep outline backend/services/generation.py --view compact 2>/dev/null || trueRepository: jamiepine/voicebox
Length of output: 434
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== relevant generation.py lines 260-360 =="
sed -n '260,360p' backend/services/generation.py
echo "== search for empty_device_cache definitions/usages and setup =="
rg -n "empty_device_cache|load_engine_model|get_tts_backend_for_engine|create_voice_prompt_for_profile|run_generation|class .*Generation|def .*generation" backend/services/generation.py backend -S || true
echo "== control-flow summary around finalizers =="
python3 - <<'PY'
from pathlib import Path
p = Path("backend/services/generation.py")
text = p.read_text().splitlines()
for i, line in enumerate(text, 1):
if "def " in line or "try:" in line or "finally:" in line or "empty_device_cache" in line or "load_engine_model" in line or "get_tts_backend_for_engine" in line or "create_voice_prompt_for_profile" in line or "return tts.audio_to_wav_bytes" in line or "bg_db.close()" in line:
print(f"{i:04d}: {line}")
PYRepository: jamiepine/voicebox
Length of output: 17283
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== run_generation try/finalizer context =="
sed -n '54,165p' backend/services/generation.py
echo "== load engine model and backend setup snippets =="
sed -n '520,675p' backend/backends/__init__.py
echo "== empty_device_cache context =="
sed -n '150,190p' backend/backends/base.pyRepository: jamiepine/voicebox
Length of output: 11845
Extend the finalizer to cover model setup.
generate_audio_sync starts the finalizer only after get_tts_backend_for_engine, load_engine_model, and create_voice_prompt_for_profile. If setup raises after allocating tensors, this synchronous path skips empty_device_cache, while run_generation already includes setup and cleanup in one try/finally. Wrap setup and generation together, and keep bg_db.close() in the nested cleanup.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@backend/services/generation.py` around lines 322 - 337, The
generate_audio_sync flow currently starts its finalizer after
backend/model/prompt setup, so setup failures bypass empty_device_cache. Move
get_tts_backend_for_engine, load_engine_model, and
create_voice_prompt_for_profile into the same outer try/finally as
generate_chunked, while preserving bg_db.close() in its existing nested cleanup.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary & Context
Fixes #923 where
voicebox-servermemory accumulated continuously over consecutive speech synthesis runs on CPU (as well as GPU/MPS) backends. On long texts or consecutive generation requests, memory would grow step-wise (+120 MB to +250 MB per run), eventually causing RAM inflation and stalling subsequent queue items ingeneratingstate.Root Cause Analysis
torch.inference_mode()ortorch.no_grad()contexts across TTS engine backends. PyTorch retained computation graphs and intermediate activation tensors in RAM across consecutive synthesis passes.empty_device_cache()lacked explicit Python garbage collection (gc.collect()) for CPU memory reclamation and did not flush Apple Silicon MPS pools.Proposed Changes
1. PyTorch Inference Mode Guards
Wrapped all forward synthesis calls in
with torch.inference_mode():across engine drivers:backend/backends/kokoro_backend.pybackend/backends/qwen_custom_voice_backend.pybackend/backends/pytorch_backend.pybackend/backends/luxtts_backend.pybackend/backends/chatterbox_backend.pybackend/backends/chatterbox_turbo_backend.py2. Mandatory Post-Generation Reclaim Hooks
Added
empty_device_cache(device)invocations in thefinallyblocks ofrun_generation()andgenerate_audio_sync()withinbackend/services/generation.py.3. Comprehensive Device & Host Cache Clearing
Updated
empty_device_cache(device)inbackend/backends/base.pyto triggergc.collect()across all platforms and addedtorch.mps.empty_cache()support for Apple Silicon.Empirical Benchmark & Verification Results
Test Environment & Cross-Platform Roadmap
CUDA_VISIBLE_DEVICES="-1"), Kokoro 82M engine (Heartfemale preset voice).Benchmark 1: 10-Run Baseline Test (97 Words / 685 Characters)
Benchmark 2: Exhaustive 7-Run Stress Test (195 Words / 1,533 Characters Passage)
Key Observations:
PowerShell Reproduction & Testing Snippet
I used the following PowerShell script to verify this fix against a running dev server (
http://localhost:17493):