Repository: huggingface/transformers
Pull request: [higgs-audio-v2] fix sampling

# What does this PR do?

As per title

Changed files:
filename: src/transformers/models/higgs_audio_v2/generation_higgs_audio_v2.py
patch: @@ -26,7 +26,13 @@
     LogitsProcessorList,
     StoppingCriteriaList,
 )
-from ...generation.logits_process import LogitsProcessor
+from ...generation.logits_process import (
+    InfNanRemoveLogitsProcessor,
+    LogitsProcessor,
+    TemperatureLogitsWarper,
+    TopKLogitsWarper,
+    TopPLogitsWarper,
+)
 from ...generation.streamers import BaseStreamer
 from ...generation.utils import GenerateNonBeamOutput
 from ...utils import add_start_docstrings, logging
@@ -192,23 +198,47 @@ class HiggsAudioV2GenerationOutput(GenerateDecoderOnlyOutput):
 
 
 class HiggsAudioV2GenerationMixin(GenerationMixin):
+    # Logits processors that only operate on scores and are safe to apply per-codebook.
+    # Other processors (e.g. RepetitionPenaltyLogitsProcessor) use input_ids to index into
+    # scores and are incompatible with audio codebook logits.
+    _supported_logits_processor_types = (
+        TemperatureLogitsWarper,
+        TopKLogitsWarper,
+        TopPLogitsWarper,
+        InfNanRemoveLogitsProcessor,
+    )
+
     def _get_logits_processor(self, *args, **kwargs):
-        if kwargs.get("logits_processor") is None:
-            logits_processor = LogitsProcessorList()
-        else:
-            logits_processor = kwargs.get("logits_processor")
-
-        logits_processor.append(
-            HiggsAudioV2DelayPatternLogitsProcessor(
-                delay_pattern=[el + 1 for el in range(self.config.num_codebooks)],
-                audio_bos_token_id=self.config.audio_bos_token_id,
-                audio_eos_token_id=self.config.audio_delay_token_id,
-                audio_stream_bos_id=self.config.audio_stream_bos_id,
-                audio_stream_eos_id=self.config.audio_stream_eos_id,
-                num_codebooks=self.config.num_codebooks,
-                codebook_size=self.config.codebook_size,
+        parent_processors = super()._get_logits_processor(*args, **kwargs)
+
+        unsupported = [p for p in parent_processors if not isinstance(p, self._supported_logits_processor_types)]
+        if unsupported:
+            unsupported_names = [type(p).__name__ for p in unsupported]
+            raise ValueError(
+                f"HiggsAudioV2 generates audio codebook logits, not text logits. "
+                f"The following logits processors are not compatible: {unsupported_names}. "
+                f"Only the following processors are supported: "
+                f"{[t.__name__ for t in self._supported_logits_processor_types]}."
             )
+
+        delay_pattern_processor = HiggsAudioV2DelayPatternLogitsProcessor(
+            delay_pattern=[el + 1 for el in range(self.config.num_codebooks)],
+            audio_bos_token_id=self.config.audio_bos_token_id,
+            audio_eos_token_id=self.config.audio_delay_token_id,
+            audio_stream_bos_id=self.config.audio_stream_bos_id,
+            audio_stream_eos_id=self.config.audio_stream_eos_id,
+            num_codebooks=self.config.num_codebooks,
+            codebook_size=self.config.codebook_size,
         )
+
+        # The delay pattern processor must run first: it reshapes scores from flat
+        # (batch_size, num_codebooks * codebook_size) to per-codebook (batch_size * num_codebooks, codebook_size).
+        # The sampling warpers (temperature, top_k, top_p) then correctly apply per-codebook.
+        # Without this ordering, top_k/top_p would filter across all codebooks combined,
+        # zeroing out entire codebooks and producing NaN after softmax.
+        logits_processor = LogitsProcessorList()
+        logits_processor.append(delay_pattern_processor)
+        logits_processor.extend(parent_processors)
         return logits_processor
 
     def _prepare_generation_config(
@@ -294,7 +324,7 @@ def _sample(
             # (the clone itself is always small)
             next_token_logits = outputs.logits[:, -1, :].to(copy=True, dtype=torch.float32, device=input_ids.device)
 
-            # pre-process distribution
+            # pre-process distribution (delay pattern reshapes to per-codebook, then warpers apply per-codebook)
             next_token_scores = logits_processor(input_ids, next_token_logits)
 
             # ===========================

filename: tests/models/higgs_audio_v2/test_modeling_higgs_audio_v2.py
patch: @@ -191,6 +191,22 @@ def setUp(self):
         self.model_tester = HiggsAudioV2ModelTester(self)
         self.config_tester = ConfigTester(self, config_class=HiggsAudioV2Config)
 
+    def _get_logits_processor_kwargs(self, do_sample=False, config=None):
+        # The base `GenerationTesterMixin` adds logits processors (e.g. repetition_penalty, bad_words_ids)
+        # that index into scores using `input_ids`. HiggsAudioV2 outputs audio codebook logits with shape
+        # (batch, num_codebooks, codebook_size), so those processors cause index-out-of-bounds errors.
+        # Override to return an empty dict (or only sampling params) to skip incompatible processors.
+        logits_processor_kwargs = {}
+        if do_sample:
+            logits_processor_kwargs.update(
+                {
+                    "top_k": 10,
+                    "top_p": 0.7,
+                    "temperature": 0.7,
+                }
+            )
+        return logits_processor_kwargs
+
     def test_config(self):
         self.config_tester.run_common_tests()