Aluode/PerceptionLabPortable

1# coding=utf-82# Copyright 2025 The HuggingFace Inc. team. All rights reserved.3#4# Licensed under the Apache License, Version 2.0 (the "License");5# you may not use this file except in compliance with the License.6# You may obtain a copy of the License at7#8#     http://www.apache.org/licenses/LICENSE-2.09#10# Unless required by applicable law or agreed to in writing, software11# distributed under the License is distributed on an "AS IS" BASIS,12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.13# See the License for the specific language governing permissions and14# limitations under the License.15"""Transformers Xcodec model."""16 17import math18from dataclasses import dataclass19from typing import Optional, Union20 21import torch22import torch.nn as nn23import torch.nn.functional as F24 25from ...modeling_utils import PreTrainedAudioTokenizerBase26from ...utils import ModelOutput, auto_docstring27from ..auto import AutoModel28from .configuration_xcodec import XcodecConfig29 30 31@dataclass32class XcodecOutput(ModelOutput):33    """34    Args:35        audio_codes (`torch.LongTensor`  of shape `(batch_size, num_quantizers, codes_length)`, *optional*):36            Discrete code indices computed using `model.encode`.37        audio_values (`torch.FloatTensor` of shape `(batch_size, channels, num_samples)`, *optional*)38            Decoded audio values obtained using the decoder part of Xcodec.39    """40 41    audio_codes: Optional[torch.LongTensor] = None42    audio_values: Optional[torch.FloatTensor] = None43 44 45@dataclass46class XcodecEncoderOutput(ModelOutput):47    """48    Args:49        audio_codes (`torch.LongTensor`  of shape `(batch_size, num_quantizers, codes_length)`, *optional*):50            Discrete code indices computed using `model.encode`.51    """52 53    audio_codes: Optional[torch.LongTensor] = None54 55 56@dataclass57class XcodecDecoderOutput(ModelOutput):58    """59    Args:60        audio_values (`torch.FloatTensor`  of shape `(batch_size, channels, num_samples)`, *optional*):61            Decoded audio values obtained using the decoder part of Xcodec.62    """63 64    audio_values: Optional[torch.FloatTensor] = None65 66 67class ResidualUnit(nn.Module):68    """Residual block for SemanticEncoder and SemanticDecoder used in Xcodec."""69 70    def __init__(self, config: XcodecConfig, in_channels: int, out_channels: int, dilation: int):71        super().__init__()72        self.activation = nn.ELU()73        padding = ((config.unit_kernel_size - 1) // 2) * dilation74        self.conv1 = nn.Conv1d(75            in_channels,76            out_channels,77            config.unit_kernel_size,78            stride=1,79            padding=padding,80            dilation=dilation,81            groups=1,82            bias=False,83        )84        self.conv2 = nn.Conv1d(in_channels=out_channels, out_channels=out_channels, kernel_size=1, bias=False)85 86    def forward(self, hidden_state: torch.Tensor) -> torch.Tensor:87        output_tensor = self.activation(hidden_state)88        output_tensor = self.conv1(output_tensor)89        output_tensor = self.activation(output_tensor)90        output_tensor = self.conv2(output_tensor)91        return hidden_state + output_tensor92 93 94class SemanticEncoderBlock(nn.Module):95    def __init__(self, config: XcodecConfig, in_channels: int, out_channels: int, stride: int):96        super().__init__()97        self.res_units = nn.ModuleList(98            [ResidualUnit(config, in_channels, in_channels, dilation) for dilation in config.block_dilations]99        )100 101        # special case: stride=1, do not use kernel=2102        kernel = 3 if stride == 1 else (2 * stride)103        padding = (kernel - 1) // 2104        self.conv = nn.Conv1d(in_channels, out_channels, kernel_size=kernel, stride=stride, padding=padding, bias=True)105 106    def forward(self, hidden_state: torch.Tensor) -> torch.Tensor:107        for unit in self.res_units:108            hidden_state = unit(hidden_state)109        hidden_state = self.conv(hidden_state)110        return hidden_state111 112 113class SemanticEncoder(nn.Module):114    def __init__(self, config):115        super().__init__()116        if len(config.strides) != len(config.channel_ratios):117            raise ValueError("Number of strides must match the number of channel_ratios.")118        self.conv = nn.Conv1d(119            config.semantic_hidden_size,120            config.semantic_hidden_size,121            config.kernel_size,122            1,123            config.kernel_size // 2,124            bias=False,125        )126 127        in_channels = config.semantic_hidden_size128        conv_blocks = []129        for i, stride in enumerate(config.strides):130            out_channels = int(config.semantic_hidden_size * config.channel_ratios[i])131            conv_blocks += [SemanticEncoderBlock(config, in_channels, out_channels, stride)]132            in_channels = out_channels133 134        self.conv_blocks = nn.ModuleList(conv_blocks)135 136    def forward(self, hidden_state: torch.Tensor) -> torch.Tensor:137        hidden_state = self.conv(hidden_state)138        for block in self.conv_blocks:139            hidden_state = block(hidden_state)140        return hidden_state141 142 143class SemanticDecoderBlock(nn.Module):144    def __init__(self, config: XcodecConfig, in_channels: int, out_channels: int, stride: int):145        super().__init__()146        if stride == 1:147            self.conv = nn.Conv1d(148                in_channels,149                out_channels,150                kernel_size=3,151                stride=1,152                padding=1,153                bias=True,154            )155        else:156            kernel_size = 2 * stride157            padding = (stride + 1) // 2158            output_padding = 1 if stride % 2 == 1 else 0159            self.conv = nn.ConvTranspose1d(160                in_channels, out_channels, kernel_size, stride, padding, output_padding, bias=False161            )162 163        self.res_units = nn.ModuleList(164            [ResidualUnit(config, out_channels, out_channels, dilation) for dilation in config.block_dilations]165        )166 167    def forward(self, hidden_state: torch.Tensor) -> torch.Tensor:168        hidden_state = self.conv(hidden_state)169        for unit in self.res_units:170            hidden_state = unit(hidden_state)171        return hidden_state172 173 174class SemanticDecoder(nn.Module):175    def __init__(self, config):176        super().__init__()177        self.conv1 = nn.Conv1d(178            in_channels=config.semantic_hidden_size,179            out_channels=int(config.semantic_hidden_size * config.channel_ratios[0]),180            kernel_size=config.kernel_size,181            stride=1,182            padding=config.kernel_size // 2,183            bias=False,184        )185        conv_blocks = []186        for i, stride in enumerate(config.strides):187            in_channels = int(config.semantic_hidden_size * config.channel_ratios[i])188 189            if i < (len(config.channel_ratios) - 1):190                out_channels = int(config.semantic_hidden_size * config.channel_ratios[i + 1])191            else:192                out_channels = config.semantic_hidden_size193 194            conv_blocks += [SemanticDecoderBlock(config, in_channels, out_channels, stride)]195 196        self.conv_blocks = nn.ModuleList(conv_blocks)197        self.conv2 = nn.Conv1d(198            config.semantic_hidden_size,199            config.semantic_hidden_size,200            config.kernel_size,201            stride=1,202            padding=config.kernel_size // 2,203            bias=False,204        )205 206    def forward(self, hidden_state: torch.Tensor) -> torch.Tensor:207        hidden_state = self.conv1(hidden_state)208        for block in self.conv_blocks:209            hidden_state = block(hidden_state)210        hidden_state = self.conv2(hidden_state)211        return hidden_state212 213 214class XcodecEuclideanCodebook(nn.Module):215    """Codebook with Euclidean distance."""216 217    def __init__(self, config):218        super().__init__()219        embed = torch.zeros(config.codebook_size, config.codebook_dim)220        self.codebook_size = config.codebook_size221        self.register_buffer("inited", torch.Tensor([True]))222        self.register_buffer("cluster_size", torch.zeros(config.codebook_size))223        self.register_buffer("embed", embed)224        self.register_buffer("embed_avg", embed.clone())225 226    # Copied from transformers.models.encodec.modeling_encodec.EncodecEuclideanCodebook.quantize227    def quantize(self, hidden_states):228        embed = self.embed.t()229        scaled_states = hidden_states.pow(2).sum(1, keepdim=True)230        dist = -(scaled_states - 2 * hidden_states @ embed + embed.pow(2).sum(0, keepdim=True))231        embed_ind = dist.max(dim=-1).indices232        return embed_ind233 234    def encode(self, hidden_states):235        shape = hidden_states.shape236        hidden_states = hidden_states.reshape((-1, shape[-1]))237        embed_ind = self.quantize(hidden_states)238        embed_ind = embed_ind.view(*shape[:-1])239        return embed_ind240 241    def decode(self, embed_ind):242        quantized = F.embedding(embed_ind, self.embed)243        return quantized244 245 246class XcodecVectorQuantization(nn.Module):247    """248    Vector quantization implementation. Currently supports only euclidean distance.249    """250 251    def __init__(self, config: XcodecConfig):252        super().__init__()253        self.codebook = XcodecEuclideanCodebook(config)254 255    # Copied from transformers.models.encodec.modeling_encodec.EncodecVectorQuantization.encode256    def encode(self, hidden_states):257        hidden_states = hidden_states.permute(0, 2, 1)258        embed_in = self.codebook.encode(hidden_states)259        return embed_in260 261    # Copied from transformers.models.encodec.modeling_encodec.EncodecVectorQuantization.decode262    def decode(self, embed_ind):263        quantize = self.codebook.decode(embed_ind)264        quantize = quantize.permute(0, 2, 1)265        return quantize266 267 268class XcodecResidualVectorQuantization(nn.Module):269    """270    Residual vector quantization implementation. Follows Algorithm 1 in https://huggingface.co/papers/2107.03312271    """272 273    def __init__(self, config: XcodecConfig):274        super().__init__()275        self.quantizers = nn.ModuleList([XcodecVectorQuantization(config) for _ in range(config.num_quantizers)])276        self.frame_rate = config.frame_rate277        self.codebook_size = config.codebook_size278        self.num_quantizers = config.num_quantizers279 280    def get_bandwidth_per_quantizer(self):281        """Return bandwidth per quantizer."""282        return math.log2(self.codebook_size) * self.frame_rate / 1000283 284    def get_num_quantizers_for_bandwidth(self, bandwidth=None) -> int:285        """Return num_quantizers based on specified target bandwidth."""286        bw_per_q = self.get_bandwidth_per_quantizer()287        num_quantizers = self.num_quantizers288        if bandwidth is not None and bandwidth > 0.0:289            num_quantizers = int(max(1, math.floor(bandwidth / bw_per_q)))290        return num_quantizers291 292    def encode(self, embeddings: torch.Tensor, bandwidth=None) -> torch.Tensor:293        """294        Encode the input tensor into discrete indices using RVQ, with the number of quantizers selected based on the given bandwidth.295        Each quantizer /codebook residually quantizes the input and returns the nearest indices in terms of Euclidian distance.296        """297        num_quantizers = self.get_num_quantizers_for_bandwidth(bandwidth)298        residual = embeddings299        all_indices = []300        for quantizer in self.quantizers[:num_quantizers]:301            indices = quantizer.encode(residual)302            quantized = quantizer.decode(indices)303            residual = residual - quantized304            all_indices.append(indices)305        out_indices = torch.stack(all_indices)306        return out_indices307 308    def decode(self, codes: torch.Tensor) -> torch.Tensor:309        """Decode the given codes to their quantized representation."""310        quantized_out = torch.tensor(0.0, device=codes.device)311        for i, indices in enumerate(codes):312            quantizer = self.quantizers[i]313            quantized = quantizer.decode(indices)314            quantized_out = quantized_out + quantized315        return quantized_out316 317 318@auto_docstring319class XcodecPreTrainedModel(PreTrainedAudioTokenizerBase):320    """321    An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained322    models.323    """324 325    config_class = XcodecConfig326    base_model_prefix = "xcodec"327    main_input_name = "input_values"328 329    def _init_weights(self, module):330        """Initialize the weights"""331        if isinstance(module, nn.Linear):332            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)333            if module.bias is not None:334                module.bias.data.zero_()335        elif isinstance(module, (nn.LayerNorm, nn.GroupNorm)):336            module.bias.data.zero_()337            module.weight.data.fill_(1.0)338        elif isinstance(module, nn.Conv1d):339            nn.init.kaiming_normal_(module.weight)340            if module.bias is not None:341                k = math.sqrt(module.groups / (module.in_channels * module.kernel_size[0]))342                nn.init.uniform_(module.bias, a=-k, b=k)343        elif module.__class__.__name__ == "Snake1d":344            module.alpha.data.fill_(1.0)345        elif isinstance(module, nn.ConvTranspose1d):346            module.reset_parameters()347        elif isinstance(module, nn.Embedding):348            module.weight.data.normal_(mean=0.0, std=0.02)349        elif isinstance(module, XcodecModel):350            # The conv1d are not handled correctly, as `self.acoustic_encoder/decoder` are initialized from a PreTrainedModel,351            # but then only the submodules are used (which are not PreTrainedModels...) -> here we reinit them as in DacModel352            for submodule in module.acoustic_encoder.modules():353                if isinstance(submodule, nn.Conv1d):354                    nn.init.trunc_normal_(submodule.weight, std=0.02)355                    nn.init.constant_(submodule.bias, 0)356            for submodule in module.acoustic_decoder.modules():357                if isinstance(submodule, nn.Conv1d):358                    nn.init.trunc_normal_(submodule.weight, std=0.02)359                    nn.init.constant_(submodule.bias, 0)360 361    def apply_weight_norm(self):362        """Apply weight norm in the acoustic encoder and decoder because the original checkpoint has weight norm applied."""363        weight_norm = torch.nn.utils.weight_norm364        if hasattr(torch.nn.utils.parametrizations, "weight_norm"):365            weight_norm = torch.nn.utils.parametrizations.weight_norm366 367        weight_norm(self.acoustic_encoder.conv1)368        weight_norm(self.acoustic_encoder.conv2)369 370        for block in self.acoustic_encoder.block:371            weight_norm(block.conv1)372            for res_unit in (block.res_unit1, block.res_unit2, block.res_unit3):373                weight_norm(res_unit.conv1)374                weight_norm(res_unit.conv2)375 376        weight_norm(self.acoustic_decoder.conv1, name="weight")377        weight_norm(self.acoustic_decoder.conv2, name="weight")378 379        for block in self.acoustic_decoder.block:380            weight_norm(block.conv_t1, name="weight")381            for res_unit in (block.res_unit1, block.res_unit2, block.res_unit3):382                weight_norm(res_unit.conv1, name="weight")383                weight_norm(res_unit.conv2, name="weight")384 385    def remove_weight_norm(self):386        """Remove the weight norm from the acoustic encoder and decoder."""387        for module in (self.acoustic_encoder, self.acoustic_decoder):388            for m in module.modules():389                try:390                    torch.nn.utils.remove_weight_norm(m, name="weight")391                except (ValueError, AttributeError):392                    pass393                if hasattr(m, "parametrizations") and "weight" in m.parametrizations:394                    torch.nn.utils.parametrize.remove_parametrizations(m, "weight", leave_parametrized=True)395 396 397@auto_docstring(custom_intro="""The Xcodec neural audio codec model.""")398class XcodecModel(XcodecPreTrainedModel):399    def __init__(self, config):400        super().__init__(config)401        self.config = config402        self.pad = config.hop_length // 2403        acoustic_model = AutoModel.from_config(config.acoustic_model_config)404        self.acoustic_encoder = acoustic_model.encoder405        self.acoustic_decoder = acoustic_model.decoder406        self._adjust_dac_decoder(self.acoustic_decoder)407        self.encoder_semantic = SemanticEncoder(config)408        self.decoder_semantic = SemanticDecoder(config)409        self.semantic_model = AutoModel.from_config(config.semantic_model_config).eval()410        self.fc = nn.Linear(config.hidden_size, config.hidden_size)411        self.fc1 = nn.Linear(config.hidden_size, config.semantic_model_config.hidden_size)412        self.fc2 = nn.Linear(config.hidden_size, config.acoustic_model_config.hidden_size)413        self.quantizer = XcodecResidualVectorQuantization(config)414 415        # Initialize weights and apply final processing416        self.post_init()417 418    @staticmethod419    def _adjust_dac_decoder(decoder: nn.Module):420        r"""421        DAC implemented in Xcodec is slightly different from the HF version.422        DAC in Xcodec adjusts the output padding in every ConvTranspose1d in the decoder and removes423        the final `nn.Tanh` activation function.424        """425        for module in decoder.modules():426            if isinstance(module, nn.ConvTranspose1d):427                stride = module.stride[0] if isinstance(module.stride, tuple) else module.stride428                module.output_padding = (stride % 2,)429        if hasattr(decoder, "tanh") and isinstance(decoder.tanh, nn.Tanh):430            decoder.tanh = nn.Identity()431 432    def _extract_semantic_features(self, input_values: torch.FloatTensor) -> torch.FloatTensor:433        input_values = input_values[:, 0, :]434        input_values = F.pad(input_values, (self.pad, self.pad))435        with torch.no_grad():436            outputs = self.semantic_model(input_values, output_hidden_states=True)437            hidden_states = outputs.hidden_states438 439        stacked = torch.stack(hidden_states, dim=1)440        return stacked.mean(dim=1)441 442    @auto_docstring443    def encode(444        self,445        input_values: torch.Tensor,446        bandwidth: Optional[float] = None,447        return_dict: Optional[bool] = None,448    ) -> Union[torch.Tensor, XcodecEncoderOutput]:449        r"""450        input_values (`torch.FloatTensor` of shape `(batch_size, channels, num_samples)`):451            Float values of the input audio waveform.452        bandwidth (`float`, *optional*):453            The target bandwidth in (kbps) supports only values in `config.target_bandwidths`.454            Defaults to the highest available bandwidth `4.0` kbps.455        return_dict (`bool`, *optional*):456            Whether or not to return a [`~utils.ModelOutput`].457 458        Returns:459            `torch.LongTensor` of shape `(batch_size, num_quantizers, codes_length)` containing the discrete encoded audio codes.460        """461        return_dict = return_dict if return_dict is not None else self.config.return_dict462 463        channels = input_values.shape[1]464        if channels != 1:465            raise ValueError(f"Audio must be mono, but got {channels}")466 467        if bandwidth is None:468            bandwidth = self.config.target_bandwidths[-1]469        elif bandwidth not in self.config.target_bandwidths:470            raise ValueError(471                f"This model doesn't support the bandwidth {bandwidth}. Select one of {self.config.target_bandwidths}."472            )473 474        e_semantic_input = self._extract_semantic_features(input_values).detach()475        e_semantic = self.encoder_semantic(e_semantic_input.transpose(1, 2))476        e_acoustic = self.acoustic_encoder(input_values)477 478        if e_acoustic.shape[2] != e_semantic.shape[2]:479            # make sure they line up if frames don't match480            e_acoustic = self.acoustic_encoder(F.pad(input_values[:, 0, :], (self.pad, self.pad)).unsqueeze(1))481 482        embeddings = torch.cat([e_acoustic, e_semantic], dim=1)483        embeddings = self.fc(embeddings.transpose(1, 2)).transpose(1, 2)484        audio_codes = self.quantizer.encode(embeddings, bandwidth)485        audio_codes = audio_codes.transpose(0, 1)486 487        if not return_dict:488            return audio_codes489 490        return XcodecEncoderOutput(audio_codes)491 492    @auto_docstring493    def decode(494        self,495        audio_codes: torch.Tensor,496        return_dict: Optional[bool] = None,497    ) -> Union[torch.Tensor, XcodecDecoderOutput]:498        r"""499        audio_codes (`torch.LongTensor`  of shape `(batch_size, num_quantizers, codes_length)`):500            Discrete code indices computed using `model.encode`.501        return_dict (`bool`, *optional*):502            Whether or not to return a [`~utils.ModelOutput`]503 504        Returns:505            Decoded audio values of shape `(batch_size, channels, num_samples)` obtained using the decoder part of506            Xcodec.507        """508        return_dict = return_dict if return_dict is not None else self.config.return_dict509 510        audio_codes = audio_codes.transpose(0, 1)511        quantized = self.quantizer.decode(audio_codes)512        quantized_acoustic = self.fc2(quantized.transpose(1, 2)).transpose(1, 2)513        audio_values = self.acoustic_decoder(quantized_acoustic)514 515        if not return_dict:516            return audio_values517 518        return XcodecDecoderOutput(audio_values)519 520    @auto_docstring521    def forward(522        self,523        input_values: torch.Tensor,524        audio_codes: Optional[torch.Tensor] = None,525        bandwidth: Optional[float] = None,526        return_dict: Optional[bool] = None,527    ) -> Union[tuple[torch.Tensor, torch.Tensor], XcodecOutput]:528        r"""529        input_values (`torch.FloatTensor` of shape `(batch_size, channels, num_samples)`):530            The raw float values of the input audio waveform.531        audio_codes (`torch.LongTensor`  of shape `(batch_size, num_quantizers, codes_length)`:532            Discrete code indices computed using `model.encode`.533        bandwidth (`float`, *optional*):534            Target bandwidth in kbps. Must be one of `config.target_bandwidths`. Defaults to the highest available bandwidth.535        bandwidth (`float`, *optional*):536            Target bandwidth in kbps. Must be one of `config.target_bandwidths`. Defaults to the highest available bandwidth.537        return_dict (`bool`, *optional*):538            Whether to return a [`XcodecOutput`] instead of a plain tuple.539 540        Returns:541            `XcodecOutput` or tuple `(audio_codes, audio_values)`:542            - `audio_codes` of shape `(batch_size, num_quantizers, codes_length)`: the quantized discrete codes.543            - `audio_values` of shape `(batch_size, channels, num_samples)`: the reconstructed audio waveform given the codes.544 545        Example:546 547        ```python548        >>> from datasets import load_dataset549        >>> from transformers import AutoFeatureExtractor, XcodecModel550 551        >>> model_id = "hf-audio/xcodec-hubert-librispeech"552        >>> model = XcodecModel.from_pretrained(model_id)553        >>> feature_extractor = AutoFeatureExtractor.from_pretrained(model_id)554 555        >>> dataset = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation")556        >>> dataset = dataset.cast_column("audio", Audio(sampling_rate=feature_extractor.sampling_rate))557        >>> audio_sample = dataset[0]['audio']['array']558 559        >>> inputs = feature_extractor(raw_audio=audio_sample, return_tensors="pt")560 561        >>> outputs = model(**inputs)562        >>> audio_codes = outputs.audio_codes563        >>> audio_values = outputs.audio_values564        ```565        """566        return_dict = return_dict if return_dict is not None else self.config.return_dict567        length = input_values.shape[-1]568 569        if audio_codes is None:570            audio_codes = self.encode(input_values, bandwidth, return_dict=False)571 572        audio_values = self.decode(audio_codes, return_dict=return_dict)[0][..., :length]573 574        if not return_dict:575            return (audio_codes, audio_values)576 577        return XcodecOutput(audio_codes=audio_codes, audio_values=audio_values)578 579 580__all__ = ["XcodecModel", "XcodecPreTrainedModel"]581