AlignmentStreamAnalyzer leaks forward hooks per generate() call, causing state corruption across runs

Summary

Repeated calls to ChatterboxMultilingualTTS.generate() on the same engine instance leak forward hooks on the transformer's attention layers. Each call registers new hooks via AlignmentStreamAnalyzer.__init__ but never removes the old ones. The stale hooks corrupt alignment state used to decide when to force EOS, degrading output quality and — on some platforms — collapsing generated audio to ~0.4 s of garbage after the first call.

The English-only ChatterboxTTS path is not affected because self.hp.is_multilingual is False and the analyzer is never constructed.

Observed behaviour

macOS (CPU/MPS) — severe:

Call 1: duration= 18.42s  samples=883200  total_stale_hooks=1
Call 2: duration=  0.40s  samples=19200   total_stale_hooks=2
Call 3: duration=  0.40s  samples=19200   total_stale_hooks=3

Windows 11, RTX 3080 Ti, CUDA 13.1, torch 2.6.0+cu124 — moderate:

Call 1: duration=  4.64s  samples=111360  total_stale_hooks=33
Call 2: duration=  4.84s  samples=116160  total_stale_hooks=36
Call 3: duration=  4.52s  samples=108480  total_stale_hooks=39
Call 4: duration=  5.12s  samples=122880  total_stale_hooks=42
Call 5: duration=  4.36s  samples=104640  total_stale_hooks=45

On Windows/CUDA the output doesn't collapse to 0.4 s, but the hook count grows monotonically (+3 per call) confirming the leak. All calls trigger forced EOS via long_tail — with the workaround applied (clearing hooks before each call), the hook count stays at 3 and behaviour becomes consistent.

Expected: every call returns full-length synthesis and total_stale_hooks stays constant.

Minimal reproducer

import torch
from chatterbox.mtl_tts import ChatterboxMultilingualTTS

engine = ChatterboxMultilingualTTS.from_pretrained(device="cpu")
text = "This is a short English sentence used to reproduce the hook leak bug."

for i in range(1, 6):
    with torch.no_grad():
        wav = engine.generate(
            text,
            language_id="en",
            audio_prompt_path="sample_voice.wav",  # any short reference clip
        )
    hooks = sum(
        len(layer.self_attn._forward_hooks) for layer in engine.t3.tfmr.layers
    )
    print(f"Call {i}: {wav.shape[-1] / engine.sr:5.2f}s  hooks={hooks}")

Root cause

1. Hook handle discarded — alignment_stream_analyzer.py, line 84:

target_layer = tfmr.layers[layer_idx].self_attn
target_layer.register_forward_hook(attention_forward_hook)

The RemovableHandle returned by register_forward_hook is discarded, so the hook can never be removed. _add_attention_spy is called once per entry in LLAMA_ALIGNED_HEADS (three entries), so each analyzer registers three hooks.

2. New analyzer every call — t3.py, lines 273-287:

self.compiled = False is force-set on line 273 immediately before the if not self.compiled guard, so a new AlignmentStreamAnalyzer is built on every generate() call. After N calls, the targeted attention layers carry N stale hooks, all closing over earlier analyzer instances whose last_aligned_attns buffers still update on every forward pass.

3. Config values overwritten — alignment_stream_analyzer.py, lines 85-90:

tfmr.config.output_attentions and tfmr.config._attn_implementation are mutated on construction. Each new analyzer saves the already-mutated values as "original", so the true originals (sdpa, output_attentions=False) are lost after the first call.

End-user workaround

for layer in engine.t3.tfmr.layers:
    layer.self_attn._forward_hooks.clear()
wav = engine.generate(text, language_id="en", audio_prompt_path="voice.wav")

This relies on the private _forward_hooks attribute and does not restore the leaked config values, so it's only a stopgap.

Proposed fix

The fix touches two files:

alignment_stream_analyzer.py:

• Store every RemovableHandle on self._hook_handles in _add_attention_spy
• Save original tfmr.config.output_attentions and _attn_implementation exactly once (guarded by self._config_patched)
• Add close() method that removes all hooks and restores config. Also implements __enter__/__exit__ for context-manager use

t3.py:

• Wrap the generation body in try: ... finally: alignment_stream_analyzer.close() so hooks are always removed, even on exception

Happy to open a PR with the patch if this direction looks right.

Environment

• chatterbox-tts 0.1.7
• torch 2.6.0
• transformers 5.2.0
• Python 3.11
• Reproduced on: macOS 14 (CPU/MPS), Linux CUDA, Windows 11 CUDA (RTX 3080 Ti)