Repository: huggingface/transformers
Pull request: [Mimi] Calibrate to ensure encoder streaming performs correctly

# What does this PR do?

1. According to the paper, this model is designed to reference 250 contexts (10 seconds), but the current implementation uses DynamicCache without employing create_sliding_window_causal_mask, causing it to reference too much context for long inputs.
2. Without using create_sliding_window_causal_mask, streaming results differ from running the entire process at once.
3. When `use_streaming=True`, failing to enable `use_cache` in the transformer layer to reference the previous cache similarly produces different results.


## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [x] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#create-a-pull-request),
      Pull Request section?
- [ ] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [ ] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [ ] Did you write any new necessary tests?


## Who can review?

@eustlb @ebezzam @vasqu @zucchini-nlp


Changed files:
filename: src/transformers/models/mimi/modeling_mimi.py
patch: @@ -24,7 +24,7 @@
 from ... import initialization as init
 from ...activations import ACT2FN
 from ...cache_utils import Cache, DynamicCache, StaticCache
-from ...masking_utils import create_causal_mask
+from ...masking_utils import create_sliding_window_causal_mask
 from ...modeling_flash_attention_utils import flash_attn_supports_top_left_mask, is_flash_attn_available
 from ...modeling_layers import GradientCheckpointingLayer
 from ...modeling_outputs import BaseModelOutputWithPast
@@ -1108,7 +1108,7 @@ def forward(
         if position_ids is None:
             position_ids = cache_position.unsqueeze(0)
 
-        causal_mask = create_causal_mask(
+        causal_mask = create_sliding_window_causal_mask(
             config=self.config,
             inputs_embeds=hidden_states,
             attention_mask=attention_mask,
@@ -1476,6 +1476,7 @@ def _encode_frame(
         padding_mask: int,
         past_key_values: Cache | None = None,
         padding_cache: MimiConv1dPaddingCache | None = None,
+        use_streaming: bool | None = None,
         return_dict: bool | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor | None]:
         """
@@ -1487,7 +1488,10 @@ def _encode_frame(
 
         # TODO: @eustlb, convert the padding mask to attention mask.
         encoder_outputs = self.encoder_transformer(
-            embeddings.transpose(1, 2), past_key_values=past_key_values, return_dict=return_dict
+            embeddings.transpose(1, 2),
+            past_key_values=past_key_values,
+            use_cache=use_streaming,
+            return_dict=return_dict,
         )
         if return_dict:
             past_key_values = encoder_outputs.get("past_key_values")
@@ -1610,6 +1614,7 @@ def encode(
             padding_mask.bool(),
             past_key_values=encoder_past_key_values,
             padding_cache=padding_cache,
+            use_streaming=use_streaming,
             return_dict=return_dict,
         )
 

filename: tests/models/mimi/test_modeling_mimi.py
patch: @@ -23,6 +23,7 @@
 from pytest import mark
 
 from transformers import AutoFeatureExtractor, MimiConfig, set_seed
+from transformers.audio_utils import load_audio
 from transformers.testing_utils import (
     is_flaky,
     is_torch_available,
@@ -506,3 +507,59 @@ def test_integration(self):
                 # the RMSE of two random gaussian noise vectors with ~N(0, 1) is around 1.0
                 rmse = compute_rmse(arr, arr_enc_dec)
                 self.assertTrue(np.abs(rmse - expected_rmse) < 1e-5)
+
+    def test_integration_longform(self):
+        """
+        Test Mimi on a longer audio (~45s) that exceeds the sliding window context (250 frames = 20s).
+        reproducer: https://gist.github.com/eustlb/34f79f34d423ccf8983c2c6c8dab2bcc
+        """
+
+        expected_rmses = {
+            "8": 0.00067151,
+            "32": 0.00049521,
+        }
+        expected_codesums = {
+            "8": 4621433,
+            "32": 18446927,
+        }
+
+        model_id = "kyutai/mimi"
+
+        processor = AutoFeatureExtractor.from_pretrained(model_id)
+        audio_sample = load_audio(
+            "https://huggingface.co/datasets/hf-internal-testing/dummy-audio-samples/resolve/main/obama_first_45_secs.mp3",
+            processor.sampling_rate,
+        )
+
+        inputs = processor(
+            raw_audio=audio_sample,
+            sampling_rate=processor.sampling_rate,
+            return_tensors="pt",
+        ).to(torch_device)
+
+        for use_cache in [False, True]:
+            model = MimiModel.from_pretrained(model_id, use_cache=use_cache).to(torch_device)
+            for num_codebooks, expected_rmse in expected_rmses.items():
+                with torch.no_grad():
+                    encoder_outputs = model.encode(inputs["input_values"], num_quantizers=int(num_codebooks))
+
+                    audio_code_sums = encoder_outputs[0].sum().item()
+
+                    self.assertTrue(
+                        np.abs(audio_code_sums - expected_codesums[num_codebooks]) <= (3e-3 * audio_code_sums)
+                    )
+
+                    input_values_dec = model.decode(encoder_outputs[0], padding_mask=inputs["padding_mask"])[0]
+                    input_values_enc_dec = model(
+                        inputs["input_values"], inputs["padding_mask"], num_quantizers=int(num_codebooks)
+                    )[1]
+
+                torch.testing.assert_close(input_values_dec, input_values_enc_dec)
+
+                self.assertTrue(inputs["input_values"].shape == input_values_enc_dec.shape)
+
+                arr = inputs["input_values"][0].cpu().numpy()
+                arr_enc_dec = input_values_enc_dec[0].cpu().numpy()
+
+                rmse = compute_rmse(arr, arr_enc_dec)
+                self.assertTrue(np.abs(rmse - expected_rmse) < 1e-5)