Guard audio tokens against Mimi's decodable range (relates to #6)

[hemanth1999k]

Background

While looking into the gibberish/garbled-output reports (#6), I found a concrete robustness gap on the decode path. This does not claim to fix the model-quality side of #6 (that's a training/sampling-stability matter), but it removes one real failure mode and is a no-op for healthy generations.

The bug

The codebook heads are sized to audio_vocab_size = 2051:

self.codebook0_head = nn.Linear(backbone_dim, config.audio_vocab_size, bias=False)
self.audio_head = nn.Parameter(torch.empty(config.audio_num_codebooks - 1, decoder_dim, config.audio_vocab_size))

so top-k sampling can produce IDs 0..2050. But Mimi only decodes raw codec IDs in [0, cardinality - 1] — cardinality == 2048, i.e. 0..2047 (see MimiModel.cardinality in moshi 0.2.2). IDs 2048..2050 are non-codec special tokens.

generate() feeds the raw samples straight into the decoder with no range check:

audio = self._audio_tokenizer.decode(torch.stack(samples).permute(1, 2, 0))...

The only stop condition is torch.all(sample == 0), which catches a fully zero frame but not a frame that merely contains a stray out-of-range token. When one is sampled, it indexes outside Mimi's codebooks and either crashes decode or produces garbage audio. (This matches the token-range mismatch @robbiemu documented in #1's comments.)

The fix

• Clamp the stacked frames into Mimi's decodable range using the tokenizer's own cardinality, right before decode. No-op for any generation that never emits an out-of-range token — it only engages on the 2048–2050 case, turning a crash/garbage into valid output. Zero regression risk for the common path.
• Return empty audio when the model emits EOS on the very first frame, instead of crashing on torch.stack([]).

One file, +18/−1.

Open question for maintainers

I went with clamp-at-decode because it's the most conservative choice and never alters a healthy generation. But the intended semantics of the 2048–2050 special tokens depend on how the model was trained, which only you know. If those IDs are meant to signal stop/EOS, the better fix would be to break generation on encountering one (like the all-zero EOS), or to mask them out of the sampling logits so they can never be drawn. Happy to switch to whichever matches the training setup — just let me know.

Relates to #6; addresses the token-range issue raised in #1.