Skip to content

fix(backend): prevent unbounded memory accumulation over consecutive TTS generations (#923) #1032

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
devangkantharia wants to merge 1 commit into
base: main
Choose a base branch
from
Open

fix(backend): prevent unbounded memory accumulation over consecutive TTS generations (#923) #1032

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 9 additions & 3 deletions backend/backends/base.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -171,17 +171,23 @@ def check_cuda_compatibility() -> tuple[bool, str | None]:

def empty_device_cache(device: str) -> None:
"""
Free cached memory on the given device (CUDA or XPU).
Free cached memory and unreferenced tensors on the given device (CUDA, XPU, MPS, CPU).

Backends should call this after unloading models so VRAM is returned
to the OS.
Backends call this after model unloading and post-generation cleanup to return
memory to the OS and prevent process heap accumulation.
"""
import gc
import torch

gc.collect()

if device == "cuda" and torch.cuda.is_available():
torch.cuda.empty_cache()
elif device == "xpu" and hasattr(torch, "xpu"):
torch.xpu.empty_cache()
elif device == "mps" and hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
if hasattr(torch.mps, "empty_cache"):
torch.mps.empty_cache()


def manual_seed(seed: int, device: str) -> None:
Expand Down
31 changes: 16 additions & 15 deletions backend/backends/chatterbox_backend.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -203,21 +203,22 @@ def _generate_sync():

logger.info(f"[Chatterbox] Generating: lang={language}")

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)
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)
Comment on lines +206 to +221

@coderabbitai coderabbitai Bot Aug 9, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🩺 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 public device property that returns _device or "cpu".
  • backend/backends/chatterbox_turbo_backend.py#L187-L201: add the same public device property.
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.


sample_rate = getattr(self.model, "sr", None) or getattr(self.model, "sample_rate", 24000)

Expand Down
29 changes: 15 additions & 14 deletions backend/backends/chatterbox_turbo_backend.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -184,20 +184,21 @@ def _generate_sync():

logger.info("[Chatterbox Turbo] Generating (English)")

wav = self.model.generate(
text,
audio_prompt_path=ref_audio,
temperature=0.8,
top_k=1000,
top_p=0.95,
repetition_penalty=1.2,
)

# Convert tensor -> numpy
if isinstance(wav, torch.Tensor):
audio = wav.squeeze().cpu().numpy().astype(np.float32)
else:
audio = np.asarray(wav, dtype=np.float32)
with torch.inference_mode():
wav = self.model.generate(
text,
audio_prompt_path=ref_audio,
temperature=0.8,
top_k=1000,
top_p=0.95,
repetition_penalty=1.2,
)

# Convert tensor -> numpy
if isinstance(wav, torch.Tensor):
audio = wav.squeeze().cpu().numpy().astype(np.float32)
else:
audio = np.asarray(wav, dtype=np.float32)

sample_rate = getattr(self.model, "sr", None) or getattr(self.model, "sample_rate", 24000)

Expand Down
13 changes: 7 additions & 6 deletions backend/backends/kokoro_backend.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -276,12 +276,13 @@ def _generate_sync():

# Generate all chunks and concatenate
audio_chunks = []
for result in pipeline(text, voice=voice_name, speed=1.0):
if result.audio is not None:
chunk = result.audio
if isinstance(chunk, torch.Tensor):
chunk = chunk.detach().cpu().numpy()
audio_chunks.append(chunk.squeeze())
with torch.inference_mode():
for result in pipeline(text, voice=voice_name, speed=1.0):
if result.audio is not None:
chunk = result.audio
if isinstance(chunk, torch.Tensor):
chunk = chunk.detach().cpu().numpy()
audio_chunks.append(chunk.squeeze())

if not audio_chunks:
# Return 1 second of silence as fallback
Expand Down
25 changes: 14 additions & 11 deletions backend/backends/luxtts_backend.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -167,18 +167,21 @@ def _generate_sync():
if seed is not None:
manual_seed(seed, self.device)

wav = self.model.generate_speech(
text=text,
encode_dict=voice_prompt,
num_steps=4,
guidance_scale=3.0,
t_shift=0.5,
speed=1.0,
return_smooth=False, # 48kHz output
)
import torch

with torch.inference_mode():
wav = self.model.generate_speech(
text=text,
encode_dict=voice_prompt,
num_steps=4,
guidance_scale=3.0,
t_shift=0.5,
speed=1.0,
return_smooth=False, # 48kHz output
)

# LuxTTS returns a tensor (may be on GPU/MPS), move to CPU first
audio = wav.detach().cpu().numpy().squeeze()
# LuxTTS returns a tensor (may be on GPU/MPS), move to CPU first
audio = wav.detach().cpu().numpy().squeeze()
return audio, 48000

return await asyncio.to_thread(_generate_sync)
13 changes: 7 additions & 6 deletions backend/backends/pytorch_backend.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -231,12 +231,13 @@ def _generate_sync():

# See _create_prompt_sync comment — inference runs with the
# process's default HF_HUB_OFFLINE state (issue #462).
wavs, sample_rate = self.model.generate_voice_clone(
text=text,
voice_clone_prompt=voice_prompt,
language=LANGUAGE_CODE_TO_NAME.get(language, "auto"),
instruct=instruct,
)
with torch.inference_mode():
wavs, sample_rate = self.model.generate_voice_clone(
text=text,
voice_clone_prompt=voice_prompt,
language=LANGUAGE_CODE_TO_NAME.get(language, "auto"),
instruct=instruct,
)
return wavs[0], sample_rate

# Run blocking inference in thread pool to avoid blocking event loop
Expand Down
3 changes: 2 additions & 1 deletion backend/backends/qwen_custom_voice_backend.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -207,7 +207,8 @@ def _generate_sync():
# state. Forcing offline here (issue #462) regressed online
# users whose libraries issue legitimate metadata lookups
# during generation.
wavs, sample_rate = self.model.generate_custom_voice(**kwargs)
with torch.inference_mode():
wavs, sample_rate = self.model.generate_custom_voice(**kwargs)
return wavs[0], sample_rate

audio, sample_rate = await asyncio.to_thread(_generate_sync)
Expand Down
26 changes: 20 additions & 6 deletions backend/services/generation.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -156,6 +156,12 @@ async def run_generation(
finally:
task_manager.complete_generation(generation_id)
bg_db.close()
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
Comment on lines +159 to +164

@coderabbitai coderabbitai Bot Aug 9, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🩺 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.py

Repository: 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"))
PY

Repository: 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"))
PY

Repository: 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.



def _notify_speak_end(generation_id: str, *, status: str) -> None:
Expand Down Expand Up @@ -313,14 +319,22 @@ async def generate_audio_sync(
if crossfade_ms is not None:
gen_kwargs["crossfade_ms"] = crossfade_ms

audio, sample_rate = await generate_chunked(
tts_model, text, voice_prompt, **gen_kwargs
)
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
Comment on lines +322 to +337

@coderabbitai coderabbitai Bot Aug 9, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🩺 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 || true

Repository: 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}")
PY

Repository: 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.py

Repository: 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.



def _save_regenerate(
Expand Down