valenti/spark-pocket-tts
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Fix memory leak by using no_grad in all threads (#16)

Browse Source 
`torch.no_grad` is thread local, I now run the autoregressive generation
in a separate thread, in addition of the thread running mimi, and forgot
to add `torch.no_grad` to the function inside this thread.
This commit is contained in:
Gabriel de Marmiesse
2026-01-13 18:07:02 +01:00
committed by GitHub
parent b2397fcbcf
commit 222714c495
1 changed files with 2 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pocket_tts/models/tts_model.py
+2
View File
@@ -261,6 +261,7 @@ class TTSModel(nn.Module):
conditioning = F.linear(latents, self.flow_lm.speaker_proj_weight) conditioning = F.linear(latents, self.flow_lm.speaker_proj_weight)
return conditioning return conditioning
@torch.no_grad
def _decode_audio_worker(self, latents_queue: queue.Queue, result_queue: queue.Queue): def _decode_audio_worker(self, latents_queue: queue.Queue, result_queue: queue.Queue):
"""Worker thread function for decoding audio latents from queue with immediate streaming.""" """Worker thread function for decoding audio latents from queue with immediate streaming."""
try: try:
@@ -509,6 +510,7 @@ class TTSModel(nn.Module):
generation_thread = threading.Thread(target=run_generation, daemon=True) generation_thread = threading.Thread(target=run_generation, daemon=True)
generation_thread.start() generation_thread.start()
@torch.no_grad
def _autoregressive_generation( def _autoregressive_generation(
self, model_state: dict, max_gen_len: int, frames_after_eos: int, latents_queue: queue.Queue self, model_state: dict, max_gen_len: int, frames_after_eos: int, latents_queue: queue.Queue
): ):