From 26dfc5ab38753411eb4878214b74f315d08e9d83 Mon Sep 17 00:00:00 2001 From: Mustafa Alammar Date: Mon, 10 Aug 2026 20:31:13 -0700 Subject: [PATCH] fix: wrap generate() in inference_mode() to stop a host-RAM leak from watermarking perth's apply_watermark() (called internally by ChatterboxTTS.generate()/ ChatterboxTurboTTS.generate()) runs its own conv-based encoder forward pass with no no_grad()/inference_mode() guard of its own. The call site in engine.py only wrapped generate() in torch.autocast(), which doesn't disable gradient tracking -- so autograd built and retained a full backward graph for that forward pass on every single call, none of which is ever backpropped through. Under sustained load this leaks real host RAM that never gets freed. Fix: add torch.inference_mode() around the same call site. No vendored code needs to change -- inference_mode() as a context manager covers everything called inside it, including the internal watermark step. --- engine.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/engine.py b/engine.py index a2be1e9..fd60490 100644 --- a/engine.py +++ b/engine.py @@ -487,7 +487,16 @@ def synthesize( # Call the core model's generate method. # autocast promotes float32 inputs to bfloat16 to match T3/S3Gen weights, # keeping numerically sensitive ops (softmax, norms) in float32 automatically. - with torch.autocast("cuda", dtype=torch.bfloat16, enabled=BF16_ENABLED): + # inference_mode matters here beyond the usual speed/memory win: the + # watermarking step inside generate() (perth's apply_watermark(), called + # internally by ChatterboxTTS.generate()/ChatterboxTurboTTS.generate()) + # runs its own conv-based encoder forward pass with no grad-mode guard of + # its own. Without an outer no_grad()/inference_mode() context, autograd + # builds and retains a full backward graph for that forward pass on every + # single call, none of which is ever backpropped through -- a real, + # measured host-RAM leak under sustained load. See the PR description for + # profiling details. + with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16, enabled=BF16_ENABLED): if loaded_model_type == "multilingual": wav_tensor = chatterbox_model.generate( text=text,