Fix vibevoice TTS batched audio index by ebezzam · Pull Request #48902 · huggingface/transformers

ebezzam · 2026-09-17T14:06:50Z

In short, while benchmarking VibeVoice TTS for the Open TTS Leaderboard, I noticed some issues with batched inference. Namely cross-talk between individual samples when one audio finishes generating earlier than another audio.

The PR includes the fix needed in generation_vibevoice.py, and Claude helped put together a test (that would fail in the original case).

Long AI description 👇

_decode_audio_latent scatters the compacted diffusing rows back into a full-batch tensor before decoding, because the acoustic tokenizer's streaming padding_cache is indexed by batch row. Its output is therefore indexed by batch position, but the loop collecting the chunks read it by position among the diffusing rows.

The two agree only while every row emits an audio token on the same step. From the first step where they diverge -- in practice as soon as the shortest utterance in the batch emits EOS -- each remaining sequence starts collecting the previous row's audio, and the first one collects the zero-filled slot of a sequence that has already finished. The result is a clip that begins with its own speech and ends in silence or a neighbour's words, at the correct total duration: the model's own state uses the right indexing throughout, so only the saved waveform is misfiled and generation is otherwise unaffected.

Measured on 8 Seed-TTS English prompts with VibeVoice-1.5B: 41.3% WER at batch size 4 against 0.0% at batch size 1, and 0.0% at batch size 4 with this fix.

The added test generates a batch and the same sequences one at a time and compares them. Getting it to fail on the bug needs three things: rows that stop at different steps, an acoustic tokenizer whose output is actually distinguishable between rows (the tester defaults decode everything to ~1e-7, below the comparison tolerance), and fixed inputs -- ids_tensor's module-level RNG is not seeded by set_seed.

`_decode_audio_latent` scatters the compacted diffusing rows back into a full-batch tensor before decoding, because the acoustic tokenizer's streaming `padding_cache` is indexed by batch row. Its output is therefore indexed by batch position, but the loop collecting the chunks read it by position among the diffusing rows. The two agree only while every row emits an audio token on the same step. From the first step where they diverge -- in practice as soon as the shortest utterance in the batch emits EOS -- each remaining sequence starts collecting the previous row's audio, and the first one collects the zero-filled slot of a sequence that has already finished. The result is a clip that begins with its own speech and ends in silence or a neighbour's words, at the correct total duration: the model's own state uses the right indexing throughout, so only the saved waveform is misfiled and generation is otherwise unaffected. Measured on 8 Seed-TTS English prompts with VibeVoice-1.5B: 41.3% WER at batch size 4 against 0.0% at batch size 1, and 0.0% at batch size 4 with this fix. The added test generates a batch and the same sequences one at a time and compares them. Getting it to fail on the bug needs three things, all noted in its docstring: rows that stop at different steps, an acoustic tokenizer whose output is actually distinguishable between rows (the tester defaults decode everything to ~1e-7, below the comparison tolerance), and fixed inputs -- `ids_tensor`'s module-level RNG is not seeded by `set_seed`. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

ebezzam · 2026-09-17T14:09:22Z

run-slow: vibevoice

vasqu

Overall should be good just the test design should reuse our helpers and creators tbh

vasqu · 2026-09-17T14:11:45Z

@@ -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):


oh wow that is easy to miss tho tbf

vasqu · 2026-09-17T14:12:02Z

@@ -302,6 +296,73 @@ 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_generate_batched_matches_single(self):
+        # different default to trigger error for incorrect indexing of audio chunks


Can you reference the issue / PR

vasqu · 2026-09-17T14:12:50Z

+        model_tester = VibeVoiceModelTester(
+            self,
+            batch_size=4,
+            seq_length=4,
+            audio_config={
+                "model_type": "vibevoice_acoustic_tokenizer",
+                "hidden_size": 16,
+                "kernel_size": 3,
+                "num_filters": 4,
+                "downsampling_ratios": [2],
+                "depths": [1, 1],
+                "layer_scale_init_value": 0.1,
+                "initializer_range": 0.5,
+                "weight_init_value": 0.5,
+            },
+        )
+        config = model_tester.get_config()
+        seed = 7
+        input_ids = ids_tensor(
+            [model_tester.batch_size, model_tester.seq_length], model_tester.vocab_size, rng=random.Random(seed)
+        )
+        attention_mask = torch.ones_like(input_ids)


Imo, we should just be able to use our usual create input and model helpers no? No need to construct a whole new tester etc

vasqu · 2026-09-17T14:13:59Z

+        def zeros_instead_of_randn(*args, **kwargs):
+            return torch.zeros(*args, **kwargs)
+
+        with torch.no_grad(), patch("torch.randn", zeros_instead_of_randn):


no grad should already be in generate no?

vasqu · 2026-09-17T14:14:48Z

+                    batched.audio[i],
+                    single.audio[0],
+                    msg=lambda m, i=i: f"Sequence {i} differs between batched and single-sample generation:\n{m}",
+                )


lets also make the 2 split test like in older models pls (e.g. see mamba2)

HuggingFaceDocBuilderDev · 2026-09-17T14:22:21Z

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

…m/transformers into fix-vibevoice-batched-audio-index

ebezzam · 2026-09-18T19:14:01Z

+    @pytest.mark.generate
+    def test_batched_equivalence_with_cache(self):
+        """Verifies that batched generation matches individual generation."""
+        self._check_batched_matches_single(use_cache=True)
+
+    @pytest.mark.generate
+    def test_batched_equivalence_without_cache(self):
+        """Verifies that batched generation matches individual generation, without cache."""
+        self._check_batched_matches_single(use_cache=False)


@vasqu is this what you meant by 2 split test like mamba2?

More like this

transformers/tests/models/mamba2/test_modeling_mamba2.py

Line 319 in 3713bd8

def test_kwargs_reach_mamba2_mixer(self):

But the design is good

ebezzam · 2026-09-18T19:15:53Z

+        batch_size, seq_length = 4, 4
+
+        config = self.model_tester.get_config()
+        # Change config so the decoded audio is distinguishable between rows.
+        config.audio_config.layer_scale_init_value = 0.1
+        config.audio_config.initializer_range = 0.5
+        config.audio_config.weight_init_value = 0.5
+
+        input_ids = ids_tensor([batch_size, seq_length], self.model_tester.vocab_size, rng=random.Random(seed))
+        attention_mask = torch.ones_like(input_ids)


re your comment, I'm not creating a whole new test this time, but are we go to overwrite some config parameters? so that the test would fail in the original case

Ah I think I confused you a bit / I worded it wrong: I just wanted the pattern as in https://github.com/ebezzam/transformers/blob/7b633379a4254ce92d599ae46e3e057810967837/tests/models/vibevoice/test_modeling_vibevoice.py#L231

vasqu

Last nits but overall good to go. Sorry i meant a different kind of split (i linked a bit more but on my phone so if its not clear just ping me again)

vasqu · 2026-09-18T19:49:19Z

+        batch_size, seq_length = 4, 4
+
+        config = self.model_tester.get_config()
+        # Change config so the decoded audio is distinguishable between rows.
+        config.audio_config.layer_scale_init_value = 0.1
+        config.audio_config.initializer_range = 0.5
+        config.audio_config.weight_init_value = 0.5
+
+        input_ids = ids_tensor([batch_size, seq_length], self.model_tester.vocab_size, rng=random.Random(seed))
+        attention_mask = torch.ones_like(input_ids)


Ah I think I confused you a bit / I worded it wrong: I just wanted the pattern as in https://github.com/ebezzam/transformers/blob/7b633379a4254ce92d599ae46e3e057810967837/tests/models/vibevoice/test_modeling_vibevoice.py#L231

vasqu · 2026-09-18T19:50:19Z

+        config.audio_config.initializer_range = 0.5
+        config.audio_config.weight_init_value = 0.5
+
+        input_ids = ids_tensor([batch_size, seq_length], self.model_tester.vocab_size, rng=random.Random(seed))


Then you use create config and input as base passed to the test (in tester)

vasqu · 2026-09-18T19:51:31Z

+    @pytest.mark.generate
+    def test_batched_equivalence_with_cache(self):
+        """Verifies that batched generation matches individual generation."""
+        self._check_batched_matches_single(use_cache=True)
+
+    @pytest.mark.generate
+    def test_batched_equivalence_without_cache(self):
+        """Verifies that batched generation matches individual generation, without cache."""
+        self._check_batched_matches_single(use_cache=False)


More like this

transformers/tests/models/mamba2/test_modeling_mamba2.py

Line 319 in 3713bd8

def test_kwargs_reach_mamba2_mixer(self):

But the design is good

ebezzam

@vasqu double checking if this is fine. from what I understood, you were looking to have a two-line pattern within the tests? 1. prepare inputs and 2. pass to tests

main question is if it's fine to add these optional args to the methods for preparing inputs? so we can keep the same small sizes for the other tests

ebezzam · 2026-09-21T13:24:18Z

@@ -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):


ok to allow optional input args? such that they default to current behavior, but for the new test allow creating a bigger batch and different settings to trigger the error (for the original code)

I like the idea! I'd only remove the rng, is it really needed?

nice! and no, not really needed. The original would sometimes pass (if it got lucky with the input), but we don't need it for the fix (from what I've observed over multiple runs). Removing it, launching slow tests, and merging 🙂

hmm actually seeding might be needed (got an error due to merge queue checks), as random input ids can lead to a flaky test:

some of the audio in the batch don't generate audio

the argmax for deciding to emit audio is near-tie, which can cause a diff between batched/single due to float noise

(FYI: obsevations from profiling with Claude over 70 seeds)

or is it better to use the is_flaky decorator?

hmm yea no worries then, it was just a if it works situation :D

ebezzam · 2026-09-21T13:25:50Z

+        config_and_inputs = self.model_tester.prepare_config_and_inputs(
+            batch_size=4,
+            seq_length=4,
+            rng=random.Random(0),
+            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)


and so that we can do the two line pattern that I think you intended?

ebezzam · 2026-09-21T14:36:53Z

run-slow: vibevoice

github-actions · 2026-09-21T14:38:04Z

AMD CI

Workflow Run ⚙️

This comment contains run-slow, running the specified jobs on AMD:

models: ["models/vibevoice"]

github-actions · 2026-09-21T14:38:21Z

Nvidia CI

Workflow Run ⚙️

This comment contains run-slow, running the specified jobs on Nvidia:

models: ["models/vibevoice"]
quantizations: []

github-actions · 2026-09-21T14:42:33Z

CI Results (Nvidia)

Workflow Run ⚙️

Commit Info

Context	Commit	Description
RUN	4d1dba9b	workflow commit (merge commit)
PR	44e912e3	branch commit (from PR)
main	d53e5876	base commit (on `main`)

✅ No failing test specific to this PR 🎉 👏 !

github-actions · 2026-09-21T14:42:47Z

CI Results (AMD)

Workflow Run ⚙️

Commit Info

Context	Commit	Description
RUN	4d1dba9b	workflow commit (merge commit)
PR	44e912e3	branch commit (from PR)
main	d53e5876	base commit (on `main`)

✅ No failing test specific to this PR 🎉 👏 !

github-actions · 2026-09-21T15:48:51Z

[For maintainers] Suggested jobs to run (before merge)

run-slow: vibevoice

github-actions · 2026-09-21T15:52:15Z

CI recap

Dashboard: View test results in Grafana
Latest run: 35616095056:1
Result: success | Jobs: 4 | Tests: 255 | Failures: 1 | Duration: 3m 30s

ebezzam and others added 4 commits September 16, 2026 19:23

CLean up.

0688025

Nit.

b585887

nit.

f7c0e61

ebezzam commented Sep 17, 2026

View reviewed changes

Comment thread tests/models/vibevoice/test_modeling_vibevoice.py Outdated

Remove guard

1dd327f

ebezzam requested a review from vasqu September 17, 2026 14:09

vasqu reviewed Sep 17, 2026

View reviewed changes

ebezzam added 3 commits September 17, 2026 17:28

address some comments

cef4d5c

Merge branch 'fix-vibevoice-batched-audio-index' of github.com:ebezza…

3a04c37

…m/transformers into fix-vibevoice-batched-audio-index

split test

7b63337

ebezzam commented Sep 18, 2026

View reviewed changes

ebezzam requested a review from vasqu September 18, 2026 19:16

vasqu approved these changes Sep 18, 2026

View reviewed changes

Follow two line pattern of other tests

608bc0b

ebezzam commented Sep 21, 2026

View reviewed changes

remove rng

44e912e

ebezzam added this pull request to the merge queue Sep 21, 2026

github-merge-queue Bot removed this pull request from the merge queue due to failed status checks Sep 21, 2026

add seed to avoid flakiness

1400d2f

ebezzam added this pull request to the merge queue Sep 21, 2026

Merged via the queue into huggingface:main with commit a10365f Sep 21, 2026
41 checks passed

ebezzam deleted the fix-vibevoice-batched-audio-index branch September 21, 2026 16:48

Conversation

ebezzam commented Sep 17, 2026 • edited by github-actions Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Long AI description 👇

Uh oh!

Uh oh!

ebezzam commented Sep 17, 2026

Uh oh!

vasqu left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

HuggingFaceDocBuilderDev commented Sep 17, 2026

Uh oh!

ebezzam Sep 18, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

vasqu left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ebezzam left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ebezzam commented Sep 21, 2026

Uh oh!

github-actions Bot commented Sep 21, 2026

AMD CI

Uh oh!

github-actions Bot commented Sep 21, 2026

Nvidia CI

Uh oh!

github-actions Bot commented Sep 21, 2026

CI Results (Nvidia)

Commit Info

Uh oh!

github-actions Bot commented Sep 21, 2026

CI Results (AMD)

Commit Info

Uh oh!

Uh oh!

github-actions Bot commented Sep 21, 2026

Uh oh!

github-actions Bot commented Sep 21, 2026

CI recap

Uh oh!

Uh oh!

ebezzam commented Sep 17, 2026 •

edited by github-actions Bot

Loading

ebezzam Sep 18, 2026 •

edited

Loading