Repository: huggingface/transformers
Pull request: Fix Bark failing tests

# Goal of this PR

There are three failing tests for Bark, see [here](https://github.com/huggingface/transformers/actions/runs/16282903241/job/45976157467) for modeling ones.

---

## 1) `tests/models/bark/test_modeling_bark.py::BarkModelIntegrationTests::test_generate_end_to_end_with_args`

```bash
FAILED tests/models/bark/test_modeling_bark.py::BarkModelIntegrationTests::test_generate_end_to_end_with_args - RuntimeError: shape '[1, 518400]' is invalid for input of size 40192
```
Fails because the vocab size is wrongly configured. By default it takes [this value](https://github.com/huggingface/transformers/blob/565dd0bad74a46d85c41e2d870f803d9e7a1a94e/src/transformers/generation/utils.py#L3992), but this is the _input vocab size_ for Bark (instead of output vocab size), which causes failure when reshaping [here](https://github.com/huggingface/transformers/blob/565dd0bad74a46d85c41e2d870f803d9e7a1a94e/src/transformers/generation/utils.py#L4128).

So a new condition should be added to extract the correct value.

@gante I see a [TODO](https://github.com/huggingface/transformers/blob/2b819ba4e383dd5b82d6d251039e30c3e94761b1/src/transformers/generation/utils.py#L3986) about standardizing the special cases. Is proposed solution for now?


## 2) RESOLVED ~~`tests/models/bark/test_modeling_bark.py::BarkModelIntegrationTests::test_generate_batching`~~

```bash
FAILED tests/models/bark/test_modeling_bark.py::BarkModelIntegrationTests::test_generate_batching - RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu! (when checking argument for argument index in method wrapper_CUDA__index_select)
```

Related to https://github.com/huggingface/transformers/issues/34634 and fixed by https://github.com/huggingface/transformers/pull/38985


## 3) `tests/models/bark/test_processor_bark.py::BarkProcessorTest::test_save_load_pretrained_additional_features`

This one is very subtle. Fails on first run with no cached models:
```bash
src/transformers/models/bark/processing_bark.py:208: ValueError
---------------------------------- Captured stderr call -------------------------------
HTTP Error 429 thrown while requesting HEAD https://huggingface.co/ylacombe/bark-small/resolve/main/speaker_embeddings/pl_speaker_3_semantic_prompt.npy
Retrying in 1s [Retry 1/5].
HTTP Error 429 thrown while requesting HEAD https://huggingface.co/ylacombe/bark-small/resolve/main/speaker_embeddings/pl_speaker_3_semantic_prompt.npy
Retrying in 2s [Retry 2/5].
HTTP Error 429 thrown while requesting HEAD https://huggingface.co/ylacombe/bark-small/resolve/main/speaker_embeddings/pl_speaker_3_semantic_prompt.npy
Retrying in 4s [Retry 3/5].
HTTP Error 429 thrown while requesting HEAD https://huggingface.co/ylacombe/bark-small/resolve/main/speaker_embeddings/pl_speaker_3_semantic_prompt.npy
Retrying in 8s [Retry 4/5].
HTTP Error 429 thrown while requesting HEAD https://huggingface.co/ylacombe/bark-small/resolve/main/speaker_embeddings/pl_speaker_3_semantic_prompt.npy
Retrying in 8s [Retry 5/5].
HTTP Error 429 thrown while requesting HEAD https://huggingface.co/ylacombe/bark-small/resolve/main/speaker_embeddings/pl_speaker_3_semantic_prompt.npy
----------------------------------- Captured log call -----------------------------------
WARNING  huggingface_hub.utils._http:_http.py:315 HTTP Error 429 thrown while requesting HEAD https://huggingface.co/ylacombe/bark-small/resolve/main/speaker_embeddings/pl_speaker_3_semantic_prompt.npy
WARNING  huggingface_hub.utils._http:_http.py:332 Retrying in 1s [Retry 1/5].
WARNING  huggingface_hub.utils._http:_http.py:315 HTTP Error 429 thrown while requesting HEAD https://huggingface.co/ylacombe/bark-small/resolve/main/speaker_embeddings/pl_speaker_3_semantic_prompt.npy
WARNING  huggingface_hub.utils._http:_http.py:332 Retrying in 2s [Retry 2/5].
WARNING  huggingface_hub.utils._http:_http.py:315 HTTP Error 429 thrown while requesting HEAD https://huggingface.co/ylacombe/bark-small/resolve/main/speaker_embeddings/pl_speaker_3_semantic_prompt.npy
WARNING  huggingface_hub.utils._http:_http.py:332 Retrying in 4s [Retry 3/5].
WARNING  huggingface_hub.utils._http:_http.py:315 HTTP Error 429 thrown while requesting HEAD https://huggingface.co/ylacombe/bark-small/resolve/main/speaker_embeddings/pl_speaker_3_semantic_prompt.npy
WARNING  huggingface_hub.utils._http:_http.py:332 Retrying in 8s [Retry 4/5].
WARNING  huggingface_hub.utils._http:_http.py:315 HTTP Error 429 thrown while requesting HEAD https://huggingface.co/ylacombe/bark-small/resolve/main/speaker_embeddings/pl_speaker_3_semantic_prompt.npy
WARNING  huggingface_hub.utils._http:_http.py:332 Retrying in 8s [Retry 5/5].
WARNING  huggingface_hub.utils._http:_http.py:315 HTTP Error 429 thrown while requesting HEAD https://huggingface.co/ylacombe/bark-small/resolve/main/speaker_embeddings/pl_speaker_3_semantic_prompt.npy
=============================================================================== short test summary info ===============================================================================
FAILED tests/models/bark/test_processor_bark.py::BarkProcessorTest::test_save_load_pretrained_additional_features - ValueError: `ylacombe/bark-small/speaker_embeddings/pl_speaker_3_semantic_prompt.npy` does not exists
```
But may pass after multiple tries (maybe because missing files are downloaded).

It can explain issues like OP of https://github.com/huggingface/transformers/issues/34634 where a voice preset may not be (first) available.

Also this error helped me realize the official checkpoints are still pointing to `ylacombe`'s checkpoints. I opened PRs on the hub to fix this.

Changed files:
filename: src/transformers/generation/utils.py
patch: @@ -3158,6 +3158,8 @@ def _beam_search(
             vocab_size = self.config.audio_vocab_size
         elif self.__class__.__name__ == "ImageGPTForCausalImageModeling":
             vocab_size = self.get_output_embeddings().out_features
+        elif self.__class__.__name__ == "BarkSemanticModel":
+            vocab_size = self.config.output_vocab_size
         else:
             vocab_size = self.config.get_text_config().vocab_size
         decoder_prompt_len = cur_len

filename: src/transformers/models/bark/processing_bark.py
patch: @@ -53,9 +53,9 @@ class BarkProcessor(ProcessorMixin):
     attributes = ["tokenizer"]
 
     preset_shape = {
-        "semantic_prompt": 1,
-        "coarse_prompt": 2,
-        "fine_prompt": 2,
+        "semantic_prompt": 1,  # 1D array of shape (X,)
+        "coarse_prompt": 2,  # 2D array of shape (2,X)
+        "fine_prompt": 2,  # 2D array of shape (8,X)
     }
 
     def __init__(self, tokenizer, speaker_embeddings=None):
@@ -115,6 +115,9 @@ def from_pretrained(
         else:
             speaker_embeddings = None
 
+        if speaker_embeddings is not None:
+            if "repo_or_path" in speaker_embeddings:
+                speaker_embeddings["repo_or_path"] = pretrained_processor_name_or_path
         tokenizer = AutoTokenizer.from_pretrained(pretrained_processor_name_or_path, **kwargs)
 
         return cls(tokenizer=tokenizer, speaker_embeddings=speaker_embeddings)
@@ -154,22 +157,21 @@ def save_pretrained(
 
             embeddings_dict["repo_or_path"] = save_directory
 
-            for prompt_key in self.speaker_embeddings:
-                if prompt_key != "repo_or_path":
-                    voice_preset = self._load_voice_preset(prompt_key)
+            for prompt_key in self.available_voice_presets:
+                voice_preset = self._load_voice_preset(prompt_key)
 
-                    tmp_dict = {}
-                    for key in self.speaker_embeddings[prompt_key]:
-                        np.save(
-                            os.path.join(
-                                embeddings_dict["repo_or_path"], speaker_embeddings_directory, f"{prompt_key}_{key}"
-                            ),
-                            voice_preset[key],
-                            allow_pickle=False,
-                        )
-                        tmp_dict[key] = os.path.join(speaker_embeddings_directory, f"{prompt_key}_{key}.npy")
+                tmp_dict = {}
+                for key in self.speaker_embeddings[prompt_key]:
+                    np.save(
+                        os.path.join(
+                            embeddings_dict["repo_or_path"], speaker_embeddings_directory, f"{prompt_key}_{key}"
+                        ),
+                        voice_preset[key],
+                        allow_pickle=False,
+                    )
+                    tmp_dict[key] = os.path.join(speaker_embeddings_directory, f"{prompt_key}_{key}.npy")
 
-                    embeddings_dict[prompt_key] = tmp_dict
+                embeddings_dict[prompt_key] = tmp_dict
 
             with open(os.path.join(save_directory, speaker_embeddings_dict_path), "w") as fp:
                 json.dump(embeddings_dict, fp)
@@ -223,6 +225,45 @@ def _validate_voice_preset_dict(self, voice_preset: Optional[dict] = None):
             if len(voice_preset[key].shape) != self.preset_shape[key]:
                 raise ValueError(f"{key} voice preset must be a {str(self.preset_shape[key])}D ndarray.")
 
+    @property
+    def available_voice_presets(self) -> list:
+        """
+        Returns a list of available voice presets.
+
+        Returns:
+            `list[str]`: A list of voice preset names.
+        """
+        if self.speaker_embeddings is None:
+            return []
+
+        voice_presets = list(self.speaker_embeddings.keys())
+        if "repo_or_path" in voice_presets:
+            voice_presets.remove("repo_or_path")
+        return voice_presets
+
+    def _verify_speaker_embeddings(self, remove_unavailable: bool = True):
+        # check which actually downloaded properly / are available
+        unavailable_keys = []
+        if self.speaker_embeddings is not None:
+            for voice_preset in self.available_voice_presets:
+                try:
+                    voice_preset_dict = self._load_voice_preset(voice_preset)
+                except ValueError:
+                    # error from `_load_voice_preset` of path not existing
+                    unavailable_keys.append(voice_preset)
+                    continue
+                self._validate_voice_preset_dict(voice_preset_dict)
+
+            if unavailable_keys:
+                logger.warning(
+                    f"The following {len(unavailable_keys)} speaker embeddings are not available: {unavailable_keys} "
+                    "If you would like to use them, please check the paths or try downloading them again."
+                )
+
+            if remove_unavailable:
+                for voice_preset in unavailable_keys:
+                    del self.speaker_embeddings[voice_preset]
+
     def __call__(
         self,
         text=None,
@@ -248,7 +289,8 @@ def __call__(
             voice_preset (`str`, `dict[np.ndarray]`):
                 The voice preset, i.e the speaker embeddings. It can either be a valid voice_preset name, e.g
                 `"en_speaker_1"`, or directly a dictionary of `np.ndarray` embeddings for each submodel of `Bark`. Or
-                it can be a valid file name of a local `.npz` single voice preset.
+                it can be a valid file name of a local `.npz` single voice preset containing the keys
+                `"semantic_prompt"`, `"coarse_prompt"` and `"fine_prompt"`.
             return_tensors (`str` or [`~utils.TensorType`], *optional*):
                 If set, will return tensors of a particular framework. Acceptable values are:
 

filename: tests/models/bark/test_processing_bark.py
patch: @@ -55,6 +55,13 @@ def test_save_load_pretrained_additional_features(self):
             pretrained_processor_name_or_path=self.checkpoint,
             speaker_embeddings_dict_path=self.speaker_embeddings_dict_path,
         )
+
+        # TODO (ebezzam) not all speaker embedding are properly downloaded.
+        # My hypothesis: there are many files (~700 speaker embeddings) and some fail to download (not the same at different first runs)
+        # https://github.com/huggingface/transformers/blob/967045082faaaaf3d653bfe665080fd746b2bb60/src/transformers/models/bark/processing_bark.py#L89
+        # https://github.com/huggingface/transformers/blob/967045082faaaaf3d653bfe665080fd746b2bb60/src/transformers/models/bark/processing_bark.py#L188
+        # So for testing purposes, we will remove the unavailable speaker embeddings before saving.
+        processor._verify_speaker_embeddings(remove_unavailable=True)
         processor.save_pretrained(
             self.tmpdirname,
             speaker_embeddings_dict_path=self.speaker_embeddings_dict_path,