Clean up XCodec and other codecs (#40348)

* Clean up xcodec addition.

* Clean up config.

* Switch to fixtures test.

* Small stuff.

* Polish XCodec and standardize across codecs.

* Update src/transformers/models/xcodec/modeling_xcodec.py

Co-authored-by: Anton Vlasjuk <[email protected]>

* Format and fix test.

* Update tol.

---------

Co-authored-by: Anton Vlasjuk <[email protected]>

Author: ebezzam

src/transformers/models/dac/feature_extraction_dac.py

@@ -150,10 +150,11 @@ def __call__(
             max_length=max_length,
             truncation=truncation,
             padding=padding,
-            return_attention_mask=False,
+            return_attention_mask=padding,
             pad_to_multiple_of=self.hop_length,
         )
-
+        if padding:
+            padded_inputs["padding_mask"] = padded_inputs.pop("attention_mask")
         if padding:
             padded_inputs.input_values = padded_inputs.input_values[:, np.newaxis, :]
 

src/transformers/models/dac/modeling_dac.py

@@ -613,6 +613,8 @@ def decode(
             The codebook indices for each codebook, representing the quantized discrete
             representation of the input. This parameter should be provided if you want
             to decode directly from the audio codes (it will overwrite quantized_representation).
+        return_dict (`bool`, *optional*, defaults to `True`):
+            Whether to return a [`DacDecoderOutput`] instead of a plain tuple.
         """
 
         if quantized_representation is None and audio_codes is None:
@@ -667,6 +669,7 @@ def forward(
 
         return_dict = return_dict if return_dict is not None else self.config.return_dict
         length = input_values.shape[-1]
+
         loss, quantized_representation, audio_codes, projected_latents = self.encode(
             input_values, n_quantizers, return_dict=False
         )

src/transformers/models/encodec/modeling_encodec.py

@@ -21,7 +21,7 @@
 import torch
 from torch import nn
 
-from ...modeling_utils import PreTrainedModel
+from ...modeling_utils import PreTrainedAudioTokenizerBase
 from ...utils import (
     ModelOutput,
     auto_docstring,
@@ -449,7 +449,7 @@ def decode(self, codes: torch.Tensor) -> torch.Tensor:
 
 
 @auto_docstring
-class EncodecPreTrainedModel(PreTrainedModel):
+class EncodecPreTrainedModel(PreTrainedAudioTokenizerBase):
     config: EncodecConfig
     base_model_prefix = "encodec"
     main_input_name = "input_values"

src/transformers/models/xcodec/modeling_xcodec.py

@@ -22,7 +22,7 @@
 import torch.nn as nn
 import torch.nn.functional as F
 
-from ...modeling_utils import PreTrainedModel
+from ...modeling_utils import PreTrainedAudioTokenizerBase
 from ...utils import ModelOutput, auto_docstring
 from ..auto import AutoModel
 from .configuration_xcodec import XcodecConfig
@@ -316,7 +316,7 @@ def decode(self, codes: torch.Tensor) -> torch.Tensor:
 
 
 @auto_docstring
-class XcodecPreTrainedModel(PreTrainedModel):
+class XcodecPreTrainedModel(PreTrainedAudioTokenizerBase):
     """
     An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
     models.
@@ -325,7 +325,6 @@ class XcodecPreTrainedModel(PreTrainedModel):
     config_class = XcodecConfig
     base_model_prefix = "xcodec"
     main_input_name = "input_values"
-    supports_gradient_checkpointing = False
 
     def _init_weights(self, module):
         """Initialize the weights"""
@@ -427,34 +426,24 @@ def encode(
         input_values: torch.Tensor,
         bandwidth: Optional[float] = None,
         return_dict: Optional[bool] = None,
-        **kwargs,
     ) -> Union[torch.Tensor, XcodecEncoderOutput]:
-        """
-        Encodes the input audio waveform into discrete audio codes.
-
-        Args:
-            input_values (`torch.FloatTensor` of shape `(batch_size, channels, num_samples)`):
-                Float values of the input audio waveform.
-            bandwidth (`float`, *optional*):
-                The target bandwidth in (kbps) supports only values in `config.target_bandwidths`.
-                Defaults to the highest available bandwidth `4.0` kbps.
-            return_dict (`bool`, *optional*):
-                Whether or not to return a [`~utils.ModelOutput`].
+        r"""
+        input_values (`torch.FloatTensor` of shape `(batch_size, channels, num_samples)`):
+            Float values of the input audio waveform.
+        bandwidth (`float`, *optional*):
+            The target bandwidth in (kbps) supports only values in `config.target_bandwidths`.
+            Defaults to the highest available bandwidth `4.0` kbps.
+        return_dict (`bool`, *optional*):
+            Whether or not to return a [`~utils.ModelOutput`].
 
         Returns:
             `torch.LongTensor` of shape `(batch_size, num_quantizers, codes_length)` containing the discrete encoded audio codes.
         """
         return_dict = return_dict if return_dict is not None else self.config.return_dict
 
-        if input_values.ndim != 3:
-            raise ValueError(
-                f"Expected input shape (batch_size, channels, num_samples), but got shape {input_values.shape}"
-            )
-
-        _, channels, self._input_length = input_values.shape
-
-        if channels not in (1, 2):
-            raise ValueError(f"Number of audio channels must be 1 or 2, but got {channels}")
+        channels = input_values.shape[1]
+        if channels != 1:
+            raise ValueError(f"Audio must be mono, but got {channels}")
 
         if bandwidth is None:
             bandwidth = self.config.target_bandwidths[-1]
@@ -483,22 +472,19 @@ def encode(
 
     @auto_docstring
     def decode(
-        self, audio_codes: torch.Tensor, return_dict: Optional[bool] = None, **kwargs
+        self,
+        audio_codes: torch.Tensor,
+        return_dict: Optional[bool] = None,
     ) -> Union[torch.Tensor, XcodecDecoderOutput]:
-        """
-        Decode the given discrete codes into an output audio waveform.
-
-        The produced audio waveform is longer than the audio input, so it's automatically trimmed to match the original input.
-
-        Args:
-            audio_codes (`torch.LongTensor`  of shape `(batch_size, num_quantizers, codes_length)`):
-                Discrete code indices computed using `model.encode`.
-
-            return_dict (`bool`, *optional*):
-                Whether or not to return a [`~utils.ModelOutput`]
+        r"""
+        audio_codes (`torch.LongTensor`  of shape `(batch_size, num_quantizers, codes_length)`):
+            Discrete code indices computed using `model.encode`.
+        return_dict (`bool`, *optional*):
+            Whether or not to return a [`~utils.ModelOutput`]
 
         Returns:
-            Decoded audio values of shape `(batch_size, channels, num_samples)` obtained using the decoder part of Xcodec.
+            Decoded audio values of shape `(batch_size, channels, num_samples)` obtained using the decoder part of
+            Xcodec.
         """
         return_dict = return_dict if return_dict is not None else self.config.return_dict
 
@@ -507,13 +493,6 @@ def decode(
         quantized_acoustic = self.fc2(quantized.transpose(1, 2)).transpose(1, 2)
         audio_values = self.acoustic_decoder(quantized_acoustic)
 
-        if getattr(self, "_input_length", None) is not None:
-            output_length = audio_values.shape[-1]
-            if self._input_length != output_length:
-                extra = output_length - self._input_length
-                start = extra // 2
-                audio_values = audio_values[..., start : start + self._input_length]
-
         if not return_dict:
             return audio_values
 
@@ -526,20 +505,18 @@ def forward(
         audio_codes: Optional[torch.Tensor] = None,
         bandwidth: Optional[float] = None,
         return_dict: Optional[bool] = None,
-        **kwargs,
     ) -> Union[tuple[torch.Tensor, torch.Tensor], XcodecOutput]:
-        """
-        Encodes and quantizes the input audio into discrete codes, then decodes those codes back into an audio waveform.
-
-        Args:
-            input_values (`torch.FloatTensor` of shape `(batch_size, channels, num_samples)`):
-                The raw float values of the input audio waveform.
-            audio_codes (`torch.LongTensor`  of shape `(batch_size, num_quantizers, codes_length)`:
-                Discrete code indices computed using `model.encode`.
-            bandwidth (`float`, *optional*):
-                Target bandwidth in kbps. Must be one of `config.target_bandwidths`. Defaults to the highest available bandwidth.
-            return_dict (`bool`, *optional*):
-                Whether to return a [`XcodecOutput`] instead of a plain tuple.
+        r"""
+        input_values (`torch.FloatTensor` of shape `(batch_size, channels, num_samples)`):
+            The raw float values of the input audio waveform.
+        audio_codes (`torch.LongTensor`  of shape `(batch_size, num_quantizers, codes_length)`:
+            Discrete code indices computed using `model.encode`.
+        bandwidth (`float`, *optional*):
+            Target bandwidth in kbps. Must be one of `config.target_bandwidths`. Defaults to the highest available bandwidth.
+        bandwidth (`float`, *optional*):
+            Target bandwidth in kbps. Must be one of `config.target_bandwidths`. Defaults to the highest available bandwidth.
+        return_dict (`bool`, *optional*):
+            Whether to return a [`XcodecOutput`] instead of a plain tuple.
 
         Returns:
             `XcodecOutput` or tuple `(audio_codes, audio_values)`:
@@ -568,11 +545,12 @@ def forward(
         ```
         """
         return_dict = return_dict if return_dict is not None else self.config.return_dict
+        length = input_values.shape[-1]
 
         if audio_codes is None:
             audio_codes = self.encode(input_values, bandwidth, return_dict=False)
 
-        audio_values = self.decode(audio_codes, return_dict=return_dict)[0]
+        audio_values = self.decode(audio_codes, return_dict=return_dict)[0][..., :length]
 
         if not return_dict:
             return (audio_codes, audio_values)

tests/models/dac/test_modeling_dac.py

@@ -527,7 +527,7 @@ def compute_rmse(arr1, arr2):
 }
 EXPECTED_QUANT_CODEBOOK_LOSS = {
     "dac_16khz": 20.7299,
-    "dac_24khz": 22.6652,
+    "dac_24khz": 22.6602,
     "dac_44khz": 16.2168,
 }
 EXPECTED_CODEC_ERROR = {
@@ -793,7 +793,7 @@ def test_integration(self, model_name):
                 atol=1e-6,
             )
             torch.testing.assert_close(
-                quantizer_outputs[4].squeeze().item(), EXPECTED_QUANT_CODEBOOK_LOSS[model_name], rtol=1e-6, atol=1e-6
+                quantizer_outputs[4].squeeze().item(), EXPECTED_QUANT_CODEBOOK_LOSS[model_name], rtol=1e-4, atol=1e-4
             )
 
             # compare decoder outputs

tests/models/xcodec/test_modeling_xcodec.py

@@ -67,9 +67,10 @@ def __init__(
         self.num_samples = num_samples
 
     def prepare_config_and_inputs(self):
-        input_values = floats_tensor([self.batch_size, self.num_channels, self.num_samples], scale=1.0)
         config = self.get_config()
-        inputs_dict = {"input_values": input_values}
+        inputs_dict = {
+            "input_values": floats_tensor([self.batch_size, self.num_channels, self.num_samples], scale=1.0)
+        }
         return config, inputs_dict
 
     def prepare_config_and_inputs_for_common(self):
@@ -82,7 +83,6 @@ def prepare_config_and_inputs_for_model_class(self, model_class):
         inputs_dict["audio_codes"] = ids_tensor(
             [self.batch_size, config.num_quantizers, codes_length], config.codebook_size
         )
-
         return config, inputs_dict
 
     def get_config(self):
@@ -94,8 +94,7 @@ def get_config(self):
 
     def create_and_check_model_forward(self, config, inputs_dict):
         model = XcodecModel(config=config).to(torch_device).eval()
-        input_values = inputs_dict["input_values"]
-        result = model(input_values)
+        result = model(input_values=inputs_dict["input_values"])
         self.parent.assertEqual(result.audio_values.shape, (self.batch_size, self.num_channels, self.num_samples))