From 6de8add5ba099e79a84561b9f02f0d766b8e6a4b Mon Sep 17 00:00:00 2001 From: hemanth1999k Date: Thu, 4 Jun 2026 17:34:55 -0500 Subject: [PATCH] Guard audio tokens against Mimi's decodable range The codebook heads emit logits over audio_vocab_size (2051), so top-k sampling can produce IDs 0..2050. Mimi only decodes raw codec IDs in [0, cardinality-1] (0..2047); IDs 2048..2050 are non-codec special tokens. generate() fed raw samples straight into mimi.decode() with no range check, so a stray special token indexes outside Mimi's codebooks and crashes decode or yields garbage audio. Clamp the stacked frames into Mimi's decodable range (using the tokenizer's own cardinality) before decoding. This is a no-op for healthy generations and only engages on out-of-range tokens. Also return empty audio when the model emits EOS on the first frame, instead of crashing on torch.stack([]). --- generator.py | 19 ++++++++++++++++++- 1 file changed, 18 insertions(+), 1 deletion(-) diff --git a/generator.py b/generator.py index 8441cee..39b9e6d 100644 --- a/generator.py +++ b/generator.py @@ -165,7 +165,24 @@ def generate( ).unsqueeze(1) curr_pos = curr_pos[:, -1:] + 1 - audio = self._audio_tokenizer.decode(torch.stack(samples).permute(1, 2, 0)).squeeze(0).squeeze(0) + if not samples: + # Model emitted EOS on the very first frame: nothing to decode. + return torch.zeros(0, device=self.device) + + # (batch=1, num_codebooks, num_frames) + frames = torch.stack(samples).permute(1, 2, 0) + + # The codebook heads produce logits over `audio_vocab_size` (2051) entries, + # but Mimi can only decode raw codec IDs in [0, cardinality - 1] (2048 -> + # 0..2047). The extra IDs are non-codec special tokens; if one is ever + # sampled (top-k sampling makes this possible) it would index outside + # Mimi's codebooks and crash decode or emit garbage. Clamp into the + # decodable range as a defensive guard. + max_codec_id = self._audio_tokenizer.cardinality - 1 + if (frames > max_codec_id).any(): + frames = frames.clamp(0, max_codec_id) + + audio = self._audio_tokenizer.decode(frames).squeeze(0).squeeze(0) # This applies an imperceptible watermark to identify audio as AI-generated. # If using Miso TTS in another application, use your own private key and keep it secret.