Fix vibevoice TTS batched audio index

src/transformers/models/vibevoice/generation_vibevoice.py

@@ -688,8 +688,8 @@ def _sample(
                 )
                 audio_output = self._decode_audio_latent(audio_latent, diffusion_mask, batch_size, acoustic_cache)
                 acoustic_cache = audio_output.padding_cache
-                for i, sample_idx in enumerate(diffusion_mask.nonzero(as_tuple=False).view(-1)):
-                    audio_chunks[sample_idx.item()].append(audio_output.audio[i])
+                for sample_idx in diffusion_mask.nonzero(as_tuple=False).view(-1):
+                    audio_chunks[sample_idx.item()].append(audio_output.audio[sample_idx])
 
                 # prepare inputs for next LM step
                 semantic_outputs = self.model.semantic_tokenizer_encoder(

tests/models/vibevoice/test_modeling_vibevoice.py

@@ -14,8 +14,10 @@
 
 import copy
 import json
+import random
 import unittest
 from pathlib import Path
+from unittest.mock import patch
 
 import pytest
 
@@ -46,14 +48,7 @@
 
 
 class DummyNoiseScheduler:
-    """
-    A simple dummy noise scheduler for testing purposes.
-
-    Contrary to real schedulers, `step` returns a *deterministic* output that does not depend on the (randomly
-    sampled) input latent. The denoised latent is fed back into the language model as the next-step embedding, so a
-    random latent would make generated sequences differ between two `generate` calls (the global RNG state advances),
-    breaking tests that compare two runs (e.g. dynamic vs static cache, eager vs compiled).
-    """
+    """A simple dummy noise scheduler for testing purposes."""
 
     def __init__(self):
         self.num_inference_steps = None
@@ -65,9 +60,8 @@ class StepOutput:
             def __init__(self, prev_sample):
                 self.prev_sample = prev_sample
 
-        # Deterministic output: ignore the random input latent and noise estimate (see class docstring)
-        prev_sample = torch.zeros_like(sample) + 0.1 * timestep.to(sample.dtype) / 1000
-        return StepOutput(prev_sample)
+        # Simple update
+        return StepOutput(sample - 0.1 * eps)
 
     def set_timesteps(self, num_inference_steps):
         self.num_inference_steps = num_inference_steps
@@ -141,10 +135,10 @@ def __init__(
         self.num_hidden_layers = text_config["num_hidden_layers"]
         self.pad_token_id = text_config["pad_token_id"]
 
-    def get_config(self):
+    def get_config(self, audio_config=None):
         return VibeVoiceConfig(
             text_config=self.text_config,
-            audio_config=self.audio_config,
+            audio_config=audio_config if audio_config is not None else self.audio_config,
             semantic_model_config=self.semantic_model_config,
             diffusion_head_config=self.diffusion_head_config,
             use_cache=self.use_cache,
@@ -155,10 +149,12 @@ def get_config(self):
             audio_token_id=5,  # Instead of default 151654
         )
 
-    def prepare_config_and_inputs(self):
-        config = self.get_config()
-        input_ids = ids_tensor([self.batch_size, self.seq_length], self.vocab_size)
-        attention_mask = torch.ones([self.batch_size, self.seq_length], dtype=torch.long, device=torch_device)
+    def prepare_config_and_inputs(self, batch_size=None, seq_length=None, rng=None, audio_config=None):
+        batch_size = batch_size if batch_size is not None else self.batch_size
+        seq_length = seq_length if seq_length is not None else self.seq_length
+        config = self.get_config(audio_config=audio_config)
+        input_ids = ids_tensor([batch_size, seq_length], self.vocab_size, rng=rng)
+        attention_mask = torch.ones([batch_size, seq_length], dtype=torch.long, device=torch_device)
         return config, input_ids, attention_mask
 
     def prepare_config_and_inputs_for_common(self):
@@ -178,6 +174,50 @@ def create_and_check_model(self, config, input_ids, attention_mask):
         self.parent.assertIsNotNone(result.logits)
         self.parent.assertEqual(result.logits.shape, (self.batch_size, self.seq_length, self.vocab_size))
 
+    def create_and_check_batched_matches_single(self, config, input_ids, attention_mask, use_cache=True):
+        # Fixed weights, so that the decoded audio is reproducible across runs.
+        set_seed(7)
+        model = VibeVoiceForConditionalGeneration(config=config).to(torch_device)
+
+        # No `min_new_tokens`: the rows have to be free to stop at different steps.
+        generate_kwargs = {
+            "noise_scheduler": DummyNoiseScheduler(),
+            "max_new_tokens": 20,
+            "do_sample": False,
+            "return_dict_in_generate": True,
+            "guidance_scale": 1.3,
+            "num_diffusion_steps": 10,
+            "use_cache": use_cache,
+        }
+
+        # Initialize diffusion with same input for comparable outputs
+        def zeros_instead_of_randn(*args, **kwargs):
+            return torch.zeros(*args, **kwargs)
+
+        with patch("torch.randn", zeros_instead_of_randn):
+            batched = model.generate(input_ids=input_ids, attention_mask=attention_mask, **generate_kwargs)
+            per_sample = [
+                model.generate(
+                    input_ids=input_ids[i : i + 1],
+                    attention_mask=attention_mask[i : i + 1],
+                    **generate_kwargs,
+                )
+                for i in range(input_ids.shape[0])
+            ]
+
+        for i, single in enumerate(per_sample):
+            self.parent.assertEqual(
+                batched.audio[i] is None,
+                single.audio[0] is None,
+                msg=f"Sequence {i}: batched and single-sample generation disagree on whether audio was produced",
+            )
+            if batched.audio[i] is not None:
+                torch.testing.assert_close(
+                    batched.audio[i],
+                    single.audio[0],
+                    msg=lambda m, i=i: f"Sequence {i} differs between batched and single-sample generation:\n{m}",
+                )
+
 
 class VibeVoiceForConditionalGenerationTest(ModelTesterMixin, GenerationTesterMixin, unittest.TestCase):
     all_model_classes = (VibeVoiceForConditionalGeneration,) if is_torch_available() else ()
@@ -302,6 +342,44 @@ def test_vibevoice_generate_max_new_tokens(self):
         self.assertIsNotNone(output.audio)
         self.assertEqual(len(output.audio), self.model_tester.batch_size)
 
+    @pytest.mark.generate
+    def test_batched_equivalence_with_cache(self):
+        """
+        Each decoded audio chunk must be attributed to the sequence that produced it, see
+        https://github.com/huggingface/transformers/pull/48902.
+        """
+        # Use different input settings to trigger different stopping times for each row, so that a wrong row/audio attribution is observable.
+        config_and_inputs = self.model_tester.prepare_config_and_inputs(
+            batch_size=4,
+            seq_length=4,
+            rng=random.Random(7),
+            audio_config={
+                **self.model_tester.audio_config,
+                "layer_scale_init_value": 0.1,
+                "initializer_range": 0.5,
+            },
+        )
+        self.model_tester.create_and_check_batched_matches_single(*config_and_inputs, use_cache=True)
+
+    @pytest.mark.generate
+    def test_batched_equivalence_without_cache(self):
+        """
+        Each decoded audio chunk must be attributed to the sequence that produced it, see
+        https://github.com/huggingface/transformers/pull/48902.
+        """
+        # Use different input settings to trigger different stopping times for each row, so that a wrong row/audio attribution is observable.
+        config_and_inputs = self.model_tester.prepare_config_and_inputs(
+            batch_size=4,
+            seq_length=4,
+            rng=random.Random(7),
+            audio_config={
+                **self.model_tester.audio_config,
+                "layer_scale_init_value": 0.1,
+                "initializer_range": 0.5,
+            },
+        )
+        self.model_tester.create_and_check_batched_matches_single(*config_and_inputs, use_cache=False)
+
     @unittest.skip(reason="Vibevoice has a special cache format so skipping for now")
     def test_cached_decode_matches_cacheless(self):
         pass