Add X-Codec model

docs/source/en/_toctree.yml

@@ -677,6 +677,8 @@
         title: UL2
       - local: model_doc/umt5
         title: UMT5
+      - local: model_doc/xcodec
+        title: X-CODEC
       - local: model_doc/xmod
         title: X-MOD
       - local: model_doc/xglm

docs/source/en/model_doc/xcodec.md

@@ -0,0 +1,93 @@
+<!--Copyright 2023 The HuggingFace Team. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
+the License. You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
+an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
+specific language governing permissions and limitations under the License.
+
+⚠️ Note that this file is in Markdown but contain specific syntax for our doc-builder (similar to MDX) that may not be
+rendered properly in your Markdown viewer.
+
+-->
+
+# X-Codec
+
+<div class="flex flex-wrap space-x-1">
+<img alt="PyTorch" src="https://img.shields.io/badge/PyTorch-DE3412?style=flat&logo=pytorch&logoColor=white">
+</div>
+
+## Overview
+
+The X-Codec model was proposed in [Codec Does Matter: Exploring the Semantic Shortcoming of Codec for Audio Language Model](https://arxiv.org/abs/2408.17175) by Zhen Ye, Peiwen Sun, Jiahe Lei, Hongzhan Lin, Xu Tan, Zheqi Dai, Qiuqiang Kong, Jianyi Chen, Jiahao Pan, Qifeng Liu, Yike Guo, Wei Xue
+
+The X-Codec model is a neural audio codec that integrates semantic information from self-supervised models (e.g., HuBERT) alongside traditional acoustic information. This enables :
+
+- **Music continuation** : Better modeling of musical semantics yields more coherent continuations.
+- **Text-to-Sound Synthesis** : X-Codec captures semantic alignment between text prompts and generated audio.
+- **Semantic aware audio tokenization**: X-Codec is used as an audio tokenizer in the YuE lyrics to song generation model.
+
+The abstract of the paper states the following:
+
+*Recent advancements in audio generation have been significantly propelled by the capabilities of Large Language Models (LLMs). The existing research on audio LLM has primarily focused on enhancing the architecture and scale of audio language models, as well as leveraging larger datasets, and generally, acoustic codecs, such as EnCodec, are used for audio tokenization. However, these codecs were originally designed for audio compression, which may lead to suboptimal performance in the context of audio LLM. Our research aims to address the shortcomings of current audio LLM codecs, particularly their challenges in maintaining semantic integrity in generated audio. For instance, existing methods like VALL-E, which condition acoustic token generation on text transcriptions, often suffer from content inaccuracies and elevated word error rates (WER) due to semantic misinterpretations of acoustic tokens, resulting in word skipping and errors. To overcome these issues, we propose a straightforward yet effective approach called X-Codec. X-Codec incorporates semantic features from a pre-trained semantic encoder before the Residual Vector Quantization (RVQ) stage and introduces a semantic reconstruction loss after RVQ. By enhancing the semantic ability of the codec, X-Codec significantly reduces WER in speech synthesis tasks and extends these benefits to non-speech applications, including music and sound generation. Our experiments in text-to-speech, music continuation, and text-to-sound tasks demonstrate that integrating semantic information substantially improves the overall performance of language models in audio generation.* 
+
+Demos can be found in this [post](https://x-codec-audio.github.io/).
+
+
+This model was contributed by [Manal El Aidouni](https://huggingface.co/Manel). The original code can be found [here](https://github.com/zhenye234/xcodec) and original checkpoint [here](https://huggingface.co/ZhenYe234/xcodec/blob/main/xcodec_speech_hubert_librispeech.pth).
+
+
+
+## Usage example 
+
+Here is a quick example of how to encode and decode an audio using this model:
+
+```python 
+from datasets import load_dataset, Audio
+from transformers import XcodecModel, AutoFeatureExtractor
+dummy_dataset = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation")
+
+# load model and feature extractor
+model = XcodecModel.from_pretrained("Manel/X-Codec")
+feature_extractor = AutoFeatureExtractor.from_pretrained("Manel/X-Codec")
+# load audio sample
+dummy_dataset = dummy_dataset.cast_column("audio", Audio(sampling_rate=feature_extractor.sampling_rate))
+audio_sample = dummy_dataset[-1]["audio"]["array"]
+inputs = feature_extractor(raw_audio=audio_sample, sampling_rate=feature_extractor.sampling_rate, return_tensors="pt")
+
+encoder_outputs = model.encode(inputs["input_values"])
+decoder_outputs = model.decode(encoder_outputs.audio_codes)
+audio_values = decoder_outputs.audio_values
+
+# or the equivalent with a forward pass
+audio_values = model(inputs["input_values"]).audio_values
+
+```
+To listen to the original and reconstructed audio, run the snippet below and then open the generated `original.wav` and `reconstruction.wav` files in your music player to compare.
+
+```python
+import soundfile as sf
+
+original = audio_sample
+reconstruction = audio_values[0].cpu().detach().numpy()
+sampling_rate = feature_extractor.sampling_rate
+
+sf.write("original.wav", original, sampling_rate)
+sf.write("reconstruction.wav", reconstruction.T, sampling_rate)
+```
+
+
+## XcodecConfig
+
+[[autodoc]] XcodecConfig
+
+
+## XcodecModel
+
+[[autodoc]] XcodecModel
+    - decode
+    - encode
+    - forward

src/transformers/models/__init__.py

@@ -344,6 +344,7 @@
     from .wavlm import *
     from .whisper import *
     from .x_clip import *
+    from .xcodec import *
     from .xglm import *
     from .xlm import *
     from .xlm_roberta import *

src/transformers/models/auto/configuration_auto.py

@@ -400,6 +400,7 @@
         ("wavlm", "WavLMConfig"),
         ("whisper", "WhisperConfig"),
         ("xclip", "XCLIPConfig"),
+        ("xcodec", "XcodecConfig"),
         ("xglm", "XGLMConfig"),
         ("xlm", "XLMConfig"),
         ("xlm-prophetnet", "XLMProphetNetConfig"),
@@ -815,6 +816,7 @@
         ("wavlm", "WavLM"),
         ("whisper", "Whisper"),
         ("xclip", "X-CLIP"),
+        ("xcodec", "X-CODEC"),
         ("xglm", "XGLM"),
         ("xlm", "XLM"),
         ("xlm-prophetnet", "XLM-ProphetNet"),

src/transformers/models/auto/feature_extraction_auto.py

@@ -115,6 +115,7 @@
         ("wavlm", "Wav2Vec2FeatureExtractor"),
         ("whisper", "WhisperFeatureExtractor"),
         ("xclip", "CLIPFeatureExtractor"),
+        ("xcodec", "EncodecFeatureExtractor"),
         ("yolos", "YolosFeatureExtractor"),
     ]
 )

src/transformers/models/auto/modeling_auto.py

@@ -370,6 +370,7 @@
         ("wavlm", "WavLMModel"),
         ("whisper", "WhisperModel"),
         ("xclip", "XCLIPModel"),
+        ("xcodec", "XcodecModel"),
         ("xglm", "XGLMModel"),
         ("xlm", "XLMModel"),
         ("xlm-prophetnet", "XLMProphetNetModel"),

src/transformers/models/xcodec/__init__.py

@@ -0,0 +1,27 @@
+# Copyright 2024 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from typing import TYPE_CHECKING
+
+from ...utils import _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+    from .configuration_xcodec import *
+    from .modeling_xcodec import *
+else:
+    import sys
+
+    _file = globals()["__file__"]
+    sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)

src/transformers/models/xcodec/configuration_xcodec.py

@@ -0,0 +1,186 @@
+# coding=utf-8
+# Copyright 2024 Descript and The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Xcodec model configuration"""
+
+import math
+from typing import Optional, Union
+
+import numpy as np
+
+from transformers import DacConfig, HubertConfig
+
+from ...configuration_utils import PretrainedConfig
+from ...utils import logging
+
+
+logger = logging.get_logger(__name__)
+
+
+class XcodecConfig(PretrainedConfig):
+    r"""
+    This is the configuration class to store the configuration of an [`XcodecModel`]. It is used to instantiate a
+    Xcodec model according to the specified arguments, defining the model architecture. Instantiating a configuration
+    with the defaults will yield a similar configuration to that of the
+    [Manel/X-Codec](https://huggingface.co/Manel/X-Codec) architecture.
+
+    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+    documentation from [`PretrainedConfig`] for more information.
+
+    Args:
+        target_bandwidths (`List[float]`, *optional*, defaults to `[0.5, 1, 1.5, 2, 4]`):
+            The range of different bandwidths (in kbps) the model can encode audio with.
+        audio_channels (`int`, *optional*, defaults to 1):
+            Number of channels in the audio data. Either 1 for mono or 2 for stereo.
+        sample_rate (`int`, *optional*, defaults to 16000):
+            The sampling rate at which the audio waveform should be digitalized, in hertz (Hz).
+        input_channels (`int`, *optional*, defaults to 768):
+            Number of channels of the input to the first convolution in the semantic encoder.
+        encoder_channels (`int`, *optional*, defaults to 768):
+            Number of hidden channels in each semantic encoder block.
+        kernel_size (`int`, *optional*, defaults to 3):
+            Kernel size for the initial semantic convolution.
+        channel_ratios (`List[float]`, *optional*, defaults to `[1, 1]`):
+            Expansion factors for the number of output channels in each semantic block.
+        strides (`List[int]`, *optional*, defaults to `[1, 1]`):
+            Strides for each semantic encoder block.
+        block_dilations (`List[int]`, *optional*, defaults to `[1, 1]`):
+            Dilation factors for the residual units in semantic blocks.
+        unit_kernel_size (`int`, *optional*, defaults to 3):
+            Kernel size inside each ResidualUnit in semantic blocks.
+        decoder_channels (`int`, *optional*, defaults to 768):
+            Number of hidden channels in each semantic decoder block.
+        output_channels (`int`, *optional*, defaults to 768):
+            Number of output channels in the semantic decoder.
+        codebook_size (`int`, *optional*, defaults to 1024):
+            Number of entries in each residual quantizer’s codebook.
+        num_quantizers (`int`, *optional*, defaults to 8):
+            Number of sequential quantizers (codebooks) in the RVQ stack.
+        codebook_dim (`int`, *optional*, defaults to 1024):
+            Dimensionality of each codebook vector.
+        initializer_range (`float`, *optional*, defaults to 0.02):
+            Standard deviation of the truncated normal initializer for all weight matrices.
+        hidden_dim (`int`, *optional*, defaults to 1024):
+            Dimensionality of the joint acoustic+semantic FC layer.
+        intermediate_dim (`int`, *optional*, defaults to 768):
+            Dimensionality of the next FC layer in the decoder path.
+        output_dim (`int`, *optional*, defaults to 256):
+            Dimensionality of the final FC layer before feeding into the acoustic decoder.
+        acoustic_model_config (`Union[Dict, DacConfig]`, *optional*):
+            An instance of the configuration for the acoustic (DAC) model.
+        semantic_model_config (`Union[Dict, HubertConfig]`, *optional*):
+            An instance of the configuration object for the semantic (HuBERT) model.
+
+    Example:
+
+    ```python
+    >>> from transformers import XcodecModel, XcodecConfig
+
+    >>> # Initializing a " " style configuration
+    >>> configuration = XcodecConfig()
+
+    >>> # Initializing a model (with random weights) from the " " style configuration
+    >>> model = XcodecModel(configuration)
+
+    >>> # Accessing the model configuration
+    >>> configuration = model.config
+    ```"""
+
+    model_type = "xcodec"
+
+    sub_configs = {
+        "acoustic_model_config": DacConfig,
+        "semantic_model_config": HubertConfig,
+    }
+
+    def __init__(
+        self,
+        target_bandwidths: Optional[list[float]] = None,
+        audio_channels: int = 1,
+        sample_rate: int = 16000,
+        input_channels: int = 768,
+        encoder_channels: int = 768,
+        kernel_size: int = 3,
+        channel_ratios: list[float] = [1, 1],
+        strides: list[int] = [1, 1],
+        block_dilations: list[int] = [1, 1],
+        unit_kernel_size: int = 3,
+        decoder_channels: int = 768,
+        output_channels: int = 768,
+        codebook_size: int = 1024,
+        num_quantizers: int = 8,
+        codebook_dim: int = 1024,
+        initializer_range: float = 0.02,
+        hidden_dim: int = 1024,
+        intermediate_dim: int = 768,
+        output_dim: int = 256,
+        acoustic_model_config: Union[dict, DacConfig] = None,
+        semantic_model_config: Union[dict, HubertConfig] = None,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        if acoustic_model_config is None:
+            self.acoustic_model_config = DacConfig(
+                encoder_hidden_size=64,
+                downsampling_ratios=[8, 5, 4, 2],
+                decoder_hidden_size=1024,
+                upsampling_ratios=[8, 5, 4, 2],
+                hidden_size=256,
+            )
+        elif isinstance(acoustic_model_config, dict):
+            self.acoustic_model_config = DacConfig(**acoustic_model_config)
+        elif isinstance(acoustic_model_config, DacConfig):
+            self.acoustic_model_config = acoustic_model_config
+
+        if semantic_model_config is None:
+            self.semantic_model_config = HubertConfig()
+        elif isinstance(semantic_model_config, dict):
+            self.semantic_model_config = HubertConfig(**semantic_model_config)
+        elif isinstance(semantic_model_config, HubertConfig):
+            self.semantic_model_config = semantic_model_config
+
+        if target_bandwidths is None:
+            target_bandwidths = [0.5, 1, 1.5, 2, 4]
+
+        self.target_bandwidths = target_bandwidths
+        self.audio_channels = audio_channels
+        self.sample_rate = sample_rate
+        self.input_channels = input_channels
+        self.encoder_channels = encoder_channels
+        self.kernel_size = kernel_size
+        self.channel_ratios = channel_ratios
+        self.strides = strides
+        self.block_dilations = block_dilations
+        self.unit_kernel_size = unit_kernel_size
+        self.decoder_channels = decoder_channels
+        self.output_channels = output_channels
+        self.codebook_size = codebook_size
+        self.num_quantizers = num_quantizers
+        self.codebook_dim = codebook_dim
+        self.initializer_range = initializer_range
+        self.hidden_dim = hidden_dim
+        self.intermediate_dim = intermediate_dim
+        self.output_dim = output_dim
+
+    @property
+    def frame_rate(self) -> int:
+        return math.ceil(self.sample_rate / np.prod(self.acoustic_model_config.upsampling_ratios))
+
+    @property
+    def hop_length(self) -> int:
+        return int(np.prod(self.acoustic_model_config.downsampling_ratios))
+
+
+__all__ = ["XcodecConfig"]