Fix rotary/kvcache OOB with speaker_audio + CUDA/CPU device mix

zonos/model.py

@@ -199,7 +199,8 @@ class Zonos(nn.Module):
         audio_seq_len = prefix_audio_len + max_new_tokens
         seq_len = prefix_conditioning.shape[1] + audio_seq_len
 
-        inference_params = self.setup_cache(batch_size=batch_size * 2, max_seqlen=seq_len)
+        # +9 headroom for the delay-pattern tail (delayed_codes is 9 longer than audio_seq_len)
+        inference_params = self.setup_cache(batch_size=batch_size * 2, max_seqlen=seq_len + 9)
 
         codes = torch.full((batch_size, 9, audio_seq_len), unknown_token, device="cuda")
         if audio_prefix_codes is not None:
@@ -237,7 +238,7 @@ class Zonos(nn.Module):
             next_token = sample_from_logits(logits, generated_tokens=delayed_codes[..., :offset], **sampling_params)
             eos_in_cb0 = next_token[:, 0] == self.eos_token_id
 
-            remaining_steps[eos_in_cb0[:, 0]] = torch.minimum(remaining_steps[eos_in_cb0[:, 0]], torch.tensor(9))
+            remaining_steps[eos_in_cb0[:, 0]] = remaining_steps[eos_in_cb0[:, 0]].clamp(max=9)
             stopping |= eos_in_cb0[:, 0]
 
             eos_codebook_idx = 9 - remaining_steps