From 7ae5a60f239c8996c21bbe0da0a1bad54129db6d Mon Sep 17 00:00:00 2001 From: Eustache Le Bihan Date: Wed, 5 Nov 2025 15:38:02 +0100 Subject: [PATCH 1/5] add can_return_tuple to dac --- src/transformers/models/dac/modeling_dac.py | 24 +++------------------ 1 file changed, 3 insertions(+), 21 deletions(-) diff --git a/src/transformers/models/dac/modeling_dac.py b/src/transformers/models/dac/modeling_dac.py index 81cfcbb931d4..a37ae7e425cd 100644 --- a/src/transformers/models/dac/modeling_dac.py +++ b/src/transformers/models/dac/modeling_dac.py @@ -24,7 +24,7 @@ import torch.nn.functional as F from ...modeling_utils import PreTrainedAudioTokenizerBase -from ...utils import ModelOutput, auto_docstring +from ...utils import ModelOutput, auto_docstring, can_return_tuple from .configuration_dac import DacConfig @@ -575,12 +575,12 @@ def __init__(self, config: DacConfig): # Initialize weights and apply final processing self.post_init() + @can_return_tuple @auto_docstring def encode( self, input_values: torch.Tensor, n_quantizers: Optional[int] = None, - return_dict: Optional[bool] = None, ): r""" input_values (`torch.Tensor of shape `(batch_size, 1, time_steps)`): @@ -588,8 +588,6 @@ def encode( n_quantizers (int, *optional*): Number of quantizers to use. If None, all quantizers are used. Default is None. """ - return_dict = return_dict if return_dict is not None else self.config.return_dict - quantized_representation = self.encoder(input_values) quantized_representation, audio_codes, projected_latents, commitment_loss, codebook_loss = self.quantizer( quantized_representation, n_quantizers @@ -597,17 +595,14 @@ def encode( loss = self.config.commitment_loss_weight * commitment_loss + self.config.codebook_loss_weight * codebook_loss - if not return_dict: - return (loss, quantized_representation, audio_codes, projected_latents) - return DacEncoderOutput(loss, quantized_representation, audio_codes, projected_latents) + @can_return_tuple @auto_docstring def decode( self, quantized_representation: Optional[torch.Tensor] = None, audio_codes: Optional[torch.Tensor] = None, - return_dict: Optional[bool] = None, ): r""" quantized_representation (torch.Tensor of shape `(batch_size, dimension, time_steps)`, *optional*): @@ -616,23 +611,16 @@ 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: raise ValueError("Either `quantized_representation` or `audio_codes` must be provided.") - return_dict = return_dict if return_dict is not None else self.config.return_dict - if audio_codes is not None: quantized_representation = self.quantizer.from_codes(audio_codes)[0] audio_values = self.decoder(quantized_representation).squeeze(1) - if not return_dict: - return (audio_values,) - return DacDecoderOutput(audio_values) @auto_docstring @@ -640,7 +628,6 @@ def forward( self, input_values: torch.Tensor, n_quantizers: Optional[int] = None, - return_dict: Optional[bool] = None, ): r""" input_values (`torch.Tensor` of shape `(batch_size, 1, time_steps)`): @@ -669,8 +656,6 @@ def forward( >>> # or the equivalent with a forward pass >>> audio_values = model(inputs["input_values"]).audio_values ```""" - - 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( @@ -678,9 +663,6 @@ def forward( ) audio_values = self.decode(quantized_representation, return_dict=False)[0][..., :length] - if not return_dict: - return (loss, audio_values, quantized_representation, audio_codes, projected_latents) - return DacOutput(loss, audio_values, quantized_representation, audio_codes, projected_latents) From 85048f9f2f73f4f7e9d1847c85f5d1cd78d60b6f Mon Sep 17 00:00:00 2001 From: Eustache Le Bihan Date: Wed, 5 Nov 2025 15:43:44 +0100 Subject: [PATCH 2/5] xcodec cleaning --- .../models/xcodec/modeling_xcodec.py | 1010 ++++++++++------- .../models/xcodec/modular_xcodec.py | 136 +++ 2 files changed, 767 insertions(+), 379 deletions(-) create mode 100644 src/transformers/models/xcodec/modular_xcodec.py diff --git a/src/transformers/models/xcodec/modeling_xcodec.py b/src/transformers/models/xcodec/modeling_xcodec.py index 774f9c74b8de..6a8f5016f707 100644 --- a/src/transformers/models/xcodec/modeling_xcodec.py +++ b/src/transformers/models/xcodec/modeling_xcodec.py @@ -1,5 +1,11 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/xcodec/modular_xcodec.py. +# Do NOT edit this file manually as any edits will be overwritten by the generation of +# the file from the modular. If any change should be done, please apply the change to the +# modular_xcodec.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 # coding=utf-8 -# Copyright 2025 The HuggingFace Inc. team. All rights reserved. +# Copyright 2025 Boson AI and 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. @@ -12,218 +18,277 @@ # 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. -"""Transformers Xcodec model.""" import math from dataclasses import dataclass -from typing import Optional, Union +from typing import Optional +import numpy as np import torch import torch.nn as nn import torch.nn.functional as F +import torchaudio from ...modeling_utils import PreTrainedAudioTokenizerBase -from ...utils import ModelOutput, auto_docstring -from ..auto import AutoModel +from ...utils import ModelOutput, auto_docstring, can_return_tuple +from ..auto.modeling_auto import AutoModel from .configuration_xcodec import XcodecConfig -@dataclass -class XcodecOutput(ModelOutput): +class Snake1d(nn.Module): """ - Args: - audio_codes (`torch.LongTensor` of shape `(batch_size, num_quantizers, codes_length)`, *optional*): - Discrete code indices computed using `model.encode`. - audio_values (`torch.FloatTensor` of shape `(batch_size, channels, num_samples)`, *optional*) - Decoded audio values obtained using the decoder part of Xcodec. + A 1-dimensional Snake activation function module. """ - audio_codes: Optional[torch.LongTensor] = None - audio_values: Optional[torch.FloatTensor] = None + def __init__(self, hidden_dim): + super().__init__() + self.alpha = nn.Parameter(torch.ones(1, hidden_dim, 1)) + def forward(self, hidden_states): + shape = hidden_states.shape + hidden_states = hidden_states.reshape(shape[0], shape[1], -1) + hidden_states = hidden_states + (self.alpha + 1e-9).reciprocal() * torch.sin(self.alpha * hidden_states).pow(2) + hidden_states = hidden_states.reshape(shape) + return hidden_states -@dataclass -class XcodecEncoderOutput(ModelOutput): + +class XcodecResidualUnit(nn.Module): """ - Args: - audio_codes (`torch.LongTensor` of shape `(batch_size, num_quantizers, codes_length)`, *optional*): - Discrete code indices computed using `model.encode`. + A residual unit composed of Snake1d and weight-normalized Conv1d layers with dilations. """ - audio_codes: Optional[torch.LongTensor] = None + def __init__(self, dimension: int = 16, dilation: int = 1): + super().__init__() + pad = ((7 - 1) * dilation) // 2 + self.snake1 = Snake1d(dimension) + self.conv1 = nn.Conv1d(dimension, dimension, kernel_size=7, dilation=dilation, padding=pad) + self.snake2 = Snake1d(dimension) + self.conv2 = nn.Conv1d(dimension, dimension, kernel_size=1) -@dataclass -class XcodecDecoderOutput(ModelOutput): - """ - Args: - audio_values (`torch.FloatTensor` of shape `(batch_size, channels, num_samples)`, *optional*): - Decoded audio values obtained using the decoder part of Xcodec. - """ + def forward(self, hidden_state): + """ + Forward pass through the residual unit. - audio_values: Optional[torch.FloatTensor] = None + Args: + hidden_state (`torch.Tensor` of shape `(batch_size, channels, time_steps)`): + Input tensor . + + Returns: + output_tensor (`torch.Tensor` of shape `(batch_size, channels, time_steps)`): + Input tensor after passing through the residual unit. + """ + output_tensor = hidden_state + output_tensor = self.conv1(self.snake1(output_tensor)) + output_tensor = self.conv2(self.snake2(output_tensor)) + padding = (hidden_state.shape[-1] - output_tensor.shape[-1]) // 2 + if padding > 0: + hidden_state = hidden_state[..., padding:-padding] + output_tensor = hidden_state + output_tensor + return output_tensor -class ResidualUnit(nn.Module): - """Residual block for SemanticEncoder and SemanticDecoder used in Xcodec.""" - def __init__(self, config: XcodecConfig, in_channels: int, out_channels: int, dilation: int): +class XcodecEncoderBlock(nn.Module): + """Encoder block used in XCODEC encoder.""" + + def __init__(self, config: XcodecConfig, stride: int = 1, stride_index: int = 1): super().__init__() - self.activation = nn.ELU() - padding = ((config.unit_kernel_size - 1) // 2) * dilation + + dimension = config.encoder_hidden_size * 2**stride_index + self.res_unit1 = XcodecResidualUnit(dimension // 2, dilation=1) + self.res_unit2 = XcodecResidualUnit(dimension // 2, dilation=3) + self.res_unit3 = XcodecResidualUnit(dimension // 2, dilation=9) + self.snake1 = Snake1d(dimension // 2) self.conv1 = nn.Conv1d( - in_channels, - out_channels, - config.unit_kernel_size, - stride=1, - padding=padding, - dilation=dilation, - groups=1, - bias=False, + dimension // 2, dimension, kernel_size=2 * stride, stride=stride, padding=math.ceil(stride / 2) ) - self.conv2 = nn.Conv1d(in_channels=out_channels, out_channels=out_channels, kernel_size=1, bias=False) - def forward(self, hidden_state: torch.Tensor) -> torch.Tensor: - output_tensor = self.activation(hidden_state) - output_tensor = self.conv1(output_tensor) - output_tensor = self.activation(output_tensor) - output_tensor = self.conv2(output_tensor) - return hidden_state + output_tensor + def forward(self, hidden_state): + hidden_state = self.res_unit1(hidden_state) + hidden_state = self.res_unit2(hidden_state) + hidden_state = self.snake1(self.res_unit3(hidden_state)) + hidden_state = self.conv1(hidden_state) + + return hidden_state + +class XcodecAcousticEncoder(nn.Module): + """XCODEC_ACOUSTIC Encoder""" -class SemanticEncoderBlock(nn.Module): - def __init__(self, config: XcodecConfig, in_channels: int, out_channels: int, stride: int): + def __init__(self, config: XcodecConfig): super().__init__() - self.res_units = nn.ModuleList( - [ResidualUnit(config, in_channels, in_channels, dilation) for dilation in config.block_dilations] - ) - # special case: stride=1, do not use kernel=2 - kernel = 3 if stride == 1 else (2 * stride) - padding = (kernel - 1) // 2 - self.conv = nn.Conv1d(in_channels, out_channels, kernel_size=kernel, stride=stride, padding=padding, bias=True) + strides = config.downsampling_ratios + # Create first convolution + self.conv1 = nn.Conv1d(1, config.encoder_hidden_size, kernel_size=7, padding=3) + + self.block = [] + # Create EncoderBlocks that double channels as they downsample by `stride` + for stride_index, stride in enumerate(strides): + stride_index = stride_index + 1 + self.block += [XcodecEncoderBlock(config, stride=stride, stride_index=stride_index)] + + self.block = nn.ModuleList(self.block) + d_model = config.encoder_hidden_size * 2**stride_index + self.snake1 = Snake1d(d_model) + self.conv2 = nn.Conv1d(d_model, config.hidden_size, kernel_size=3, padding=1) + + def forward(self, hidden_state): + hidden_state = self.conv1(hidden_state) + + for module in self.block: + hidden_state = module(hidden_state) + + hidden_state = self.snake1(hidden_state) + hidden_state = self.conv2(hidden_state) - def forward(self, hidden_state: torch.Tensor) -> torch.Tensor: - for unit in self.res_units: - hidden_state = unit(hidden_state) - hidden_state = self.conv(hidden_state) return hidden_state -class SemanticEncoder(nn.Module): +class XcodecSemanticEncoderLayer(nn.Module): def __init__(self, config): super().__init__() - if len(config.strides) != len(config.channel_ratios): - raise ValueError("Number of strides must match the number of channel_ratios.") - self.conv = nn.Conv1d( - config.semantic_hidden_size, - config.semantic_hidden_size, - config.kernel_size, - 1, - config.kernel_size // 2, - bias=False, - ) + self.act_fn = nn.ELU() + self.conv1 = nn.Conv1d(config.semantic_hidden_size, config.semantic_hidden_size, kernel_size=3, bias=False) + self.conv2 = nn.Conv1d(config.semantic_hidden_size, config.semantic_hidden_size, kernel_size=1, bias=False) + self.conv3 = nn.Conv1d(config.semantic_hidden_size, config.semantic_hidden_size, kernel_size=3, bias=False) + + def forward(self, hidden_state): + residual = hidden_state + hidden_state = self.conv1(self.act_fn(hidden_state)) + hidden_state = self.conv2(self.act_fn(hidden_state)) + hidden_state = self.conv3(residual + hidden_state) + return hidden_state - in_channels = config.semantic_hidden_size - conv_blocks = [] - for i, stride in enumerate(config.strides): - out_channels = int(config.semantic_hidden_size * config.channel_ratios[i]) - conv_blocks += [SemanticEncoderBlock(config, in_channels, out_channels, stride)] - in_channels = out_channels - self.conv_blocks = nn.ModuleList(conv_blocks) +class XcodecSemanticEncoder(nn.Module): + def __init__(self, config): + super().__init__() + self.input_conv = nn.Conv1d( + config.semantic_hidden_size, config.semantic_hidden_size, kernel_size=3, padding=1, bias=False + ) + self.layers = nn.ModuleList([XcodecSemanticEncoderLayer(config) for _ in range(config.num_layers)]) - def forward(self, hidden_state: torch.Tensor) -> torch.Tensor: - hidden_state = self.conv(hidden_state) - for block in self.conv_blocks: - hidden_state = block(hidden_state) + def forward(self, hidden_state): + hidden_state = self.input_conv(hidden_state) + for layer in self.layers: + hidden_state = layer(hidden_state) return hidden_state -class SemanticDecoderBlock(nn.Module): - def __init__(self, config: XcodecConfig, in_channels: int, out_channels: int, stride: int): +class XcodecEncoder(nn.Module): + def __init__(self, config: XcodecConfig): super().__init__() - if stride == 1: - self.conv = nn.Conv1d( - in_channels, - out_channels, - kernel_size=3, - stride=1, - padding=1, - bias=True, - ) - else: - kernel_size = 2 * stride - padding = (stride + 1) // 2 - output_padding = 1 if stride % 2 == 1 else 0 - self.conv = nn.ConvTranspose1d( - in_channels, out_channels, kernel_size, stride, padding, output_padding, bias=False + self.acoustic_encoder = XcodecAcousticEncoder(config) + self.semantic_encoder = XcodecSemanticEncoder(config) + self.semantic_model = AutoModel.from_config(config.semantic_config) + + self.sample_rate = config.sample_rate + self.semantic_sample_rate = config.semantic_sample_rate + + def _extract_semantic_features(self, input_values): + with torch.no_grad(): + input_values = torchaudio.functional.resample( + input_values, self.config.sampling_rate, self.config.semantic_sample_rate ) + input_values = input_values[:, 0, :] + input_values = F.pad(input_values, (self.config.pad, self.config.pad)) + outputs = self.semantic_model(input_values, output_hidden_states=True) + hidden_states = outputs.hidden_states + stacked = torch.stack(hidden_states, dim=1) + semantic_features = stacked.mean(dim=1) + semantic_features = semantic_features[:, :: self.config.semantic_downsample_factor, :] + return semantic_features - self.res_units = nn.ModuleList( - [ResidualUnit(config, out_channels, out_channels, dilation) for dilation in config.block_dilations] - ) + def forward(self, input_values: torch.Tensor): + acoustic_embeds = self.acoustic_encoder(input_values) - def forward(self, hidden_state: torch.Tensor) -> torch.Tensor: - hidden_state = self.conv(hidden_state) - for unit in self.res_units: - hidden_state = unit(hidden_state) - return hidden_state + semantic_features = self._extract_semantic_features(input_values) + semantic_embeds = self.semantic_encoder(semantic_features.transpose(1, 2)) + input_embeds = torch.cat([acoustic_embeds, semantic_embeds], dim=1) + return input_embeds -class SemanticDecoder(nn.Module): - def __init__(self, config): + +class XcodecDecoderBlock(nn.Module): + """Decoder block used in XCODEC decoder.""" + + def __init__(self, config, stride: int = 1, stride_index: int = 1): super().__init__() - self.conv1 = nn.Conv1d( - in_channels=config.semantic_hidden_size, - out_channels=int(config.semantic_hidden_size * config.channel_ratios[0]), - kernel_size=config.kernel_size, - stride=1, - padding=config.kernel_size // 2, - bias=False, - ) - conv_blocks = [] - for i, stride in enumerate(config.strides): - in_channels = int(config.semantic_hidden_size * config.channel_ratios[i]) - - if i < (len(config.channel_ratios) - 1): - out_channels = int(config.semantic_hidden_size * config.channel_ratios[i + 1]) - else: - out_channels = config.semantic_hidden_size - - conv_blocks += [SemanticDecoderBlock(config, in_channels, out_channels, stride)] - - self.conv_blocks = nn.ModuleList(conv_blocks) - self.conv2 = nn.Conv1d( - config.semantic_hidden_size, - config.semantic_hidden_size, - config.kernel_size, - stride=1, - padding=config.kernel_size // 2, - bias=False, + input_dim = config.decoder_hidden_size // 2**stride_index + output_dim = config.decoder_hidden_size // 2 ** (stride_index + 1) + self.snake1 = Snake1d(input_dim) + self.conv_t1 = nn.ConvTranspose1d( + input_dim, + output_dim, + kernel_size=2 * stride, + stride=stride, + padding=math.ceil(stride / 2), + output_padding=(stride % 2,), ) - def forward(self, hidden_state: torch.Tensor) -> torch.Tensor: - hidden_state = self.conv1(hidden_state) - for block in self.conv_blocks: - hidden_state = block(hidden_state) - hidden_state = self.conv2(hidden_state) + self.res_unit1 = XcodecResidualUnit(output_dim, dilation=1) + self.res_unit2 = XcodecResidualUnit(output_dim, dilation=3) + self.res_unit3 = XcodecResidualUnit(output_dim, dilation=9) + + def forward(self, hidden_state): + hidden_state = self.snake1(hidden_state) + hidden_state = self.conv_t1(hidden_state) + hidden_state = self.res_unit1(hidden_state) + hidden_state = self.res_unit2(hidden_state) + hidden_state = self.res_unit3(hidden_state) + return hidden_state +class XcodecDecoder(nn.Module): + """XCODEC Decoder""" + + def __init__(self, config: XcodecConfig): + super().__init__() + input_channel = config.acoustic_hidden_size + channels = config.decoder_hidden_size + strides = config.upsampling_ratios + + # Add first conv layer + self.conv1 = nn.Conv1d(input_channel, channels, kernel_size=7, padding=3) + + # Add upsampling + MRF blocks + block = [] + for stride_index, stride in enumerate(strides): + block += [XcodecDecoderBlock(config, stride, stride_index)] + + self.block = nn.ModuleList(block) + output_dim = config.decoder_hidden_size // 2 ** (stride_index + 1) + self.snake1 = Snake1d(output_dim) + self.conv2 = nn.Conv1d(output_dim, 1, kernel_size=7, padding=3) + + def forward(self, hidden_states): + hidden_states = self.conv1(hidden_states) + for layer in self.block: + hidden_states = layer(hidden_states) + + hidden_states = self.snake1(hidden_states) + hidden_states = self.conv2(hidden_states) + + return hidden_states + + class XcodecEuclideanCodebook(nn.Module): """Codebook with Euclidean distance.""" - def __init__(self, config): + def __init__(self, config: XcodecConfig): super().__init__() embed = torch.zeros(config.codebook_size, config.codebook_dim) + self.codebook_size = config.codebook_size + self.register_buffer("inited", torch.Tensor([True])) self.register_buffer("cluster_size", torch.zeros(config.codebook_size)) self.register_buffer("embed", embed) self.register_buffer("embed_avg", embed.clone()) - # Copied from transformers.models.encodec.modeling_encodec.EncodecEuclideanCodebook.quantize def quantize(self, hidden_states): embed = self.embed.t() scaled_states = hidden_states.pow(2).sum(1, keepdim=True) @@ -233,14 +298,17 @@ def quantize(self, hidden_states): def encode(self, hidden_states): shape = hidden_states.shape + # pre-process hidden_states = hidden_states.reshape((-1, shape[-1])) + # quantize embed_ind = self.quantize(hidden_states) + # post-process embed_ind = embed_ind.view(*shape[:-1]) return embed_ind def decode(self, embed_ind): - quantized = F.embedding(embed_ind, self.embed) - return quantized + quantize = nn.functional.embedding(embed_ind, self.embed) + return quantize class XcodecVectorQuantization(nn.Module): @@ -252,13 +320,11 @@ def __init__(self, config: XcodecConfig): super().__init__() self.codebook = XcodecEuclideanCodebook(config) - # Copied from transformers.models.encodec.modeling_encodec.EncodecVectorQuantization.encode def encode(self, hidden_states): hidden_states = hidden_states.permute(0, 2, 1) embed_in = self.codebook.encode(hidden_states) return embed_in - # Copied from transformers.models.encodec.modeling_encodec.EncodecVectorQuantization.decode def decode(self, embed_ind): quantize = self.codebook.decode(embed_ind) quantize = quantize.permute(0, 2, 1) @@ -266,255 +332,461 @@ def decode(self, embed_ind): class XcodecResidualVectorQuantization(nn.Module): - """ - Residual vector quantization implementation. Follows Algorithm 1 in https://huggingface.co/papers/2107.03312 - """ + """Residual Vector Quantizer.""" def __init__(self, config: XcodecConfig): super().__init__() - self.quantizers = nn.ModuleList([XcodecVectorQuantization(config) for _ in range(config.num_quantizers)]) - self.frame_rate = config.frame_rate self.codebook_size = config.codebook_size + self.frame_rate = config.frame_rate self.num_quantizers = config.num_quantizers + self.layers = nn.ModuleList([XcodecVectorQuantization(config) for _ in range(config.num_quantizers)]) - def get_bandwidth_per_quantizer(self): - """Return bandwidth per quantizer.""" - return math.log2(self.codebook_size) * self.frame_rate / 1000 - - def get_num_quantizers_for_bandwidth(self, bandwidth=None) -> int: + def get_num_quantizers_for_bandwidth(self, bandwidth: Optional[float] = None) -> int: """Return num_quantizers based on specified target bandwidth.""" - bw_per_q = self.get_bandwidth_per_quantizer() + bw_per_q = math.log2(self.codebook_size) * self.frame_rate num_quantizers = self.num_quantizers if bandwidth is not None and bandwidth > 0.0: - num_quantizers = int(max(1, math.floor(bandwidth / bw_per_q))) + num_quantizers = int(max(1, math.floor(bandwidth * 1000 / bw_per_q))) return num_quantizers - def encode(self, embeddings: torch.Tensor, bandwidth=None) -> torch.Tensor: + def encode(self, embeddings: torch.Tensor, bandwidth: Optional[float] = None) -> torch.Tensor: """ - Encode the input tensor into discrete indices using RVQ, with the number of quantizers selected based on the given bandwidth. - Each quantizer /codebook residually quantizes the input and returns the nearest indices in terms of Euclidian distance. + Encode a given input tensor with the specified frame rate at the given bandwidth. The RVQ encode method sets + the appropriate number of quantizers to use and returns indices for each quantizer. """ num_quantizers = self.get_num_quantizers_for_bandwidth(bandwidth) residual = embeddings all_indices = [] - for quantizer in self.quantizers[:num_quantizers]: - indices = quantizer.encode(residual) - quantized = quantizer.decode(indices) + for layer in self.layers[:num_quantizers]: + indices = layer.encode(residual) + quantized = layer.decode(indices) residual = residual - quantized all_indices.append(indices) out_indices = torch.stack(all_indices) return out_indices def decode(self, codes: torch.Tensor) -> torch.Tensor: - """Decode the given codes to their quantized representation.""" + """Decode the given codes to the quantized representation.""" quantized_out = torch.tensor(0.0, device=codes.device) for i, indices in enumerate(codes): - quantizer = self.quantizers[i] - quantized = quantizer.decode(indices) + layer = self.layers[i] + quantized = layer.decode(indices) quantized_out = quantized_out + quantized return quantized_out +@dataclass @auto_docstring -class XcodecPreTrainedModel(PreTrainedAudioTokenizerBase): +class XcodecOutput(ModelOutput): + r""" + loss (`torch.Tensor`): + Loss from the encoder model, comprising the weighted combination of the commitment and codebook losses. + audio_values (`torch.Tensor` of shape `(batch_size, input_length)`): + Reconstructed audio data. + quantized_representation (`torch.Tensor` of shape `(batch_size, dimension, time_steps)`): + Quantized continuous representation of input. + audio_codes (`torch.LongTensor` of shape `(batch_size, num_codebooks, time_steps)`): + Codebook indices for each codebook (quantized discrete representation of input). + projected_latents (`torch.Tensor` of shape `(batch_size, num_codebooks * dimension, time_steps)`): + Projected latents (continuous representation of input before quantization). + """ + + loss: Optional[torch.FloatTensor] = None + audio_values: Optional[torch.FloatTensor] = None + quantized_representation: Optional[torch.FloatTensor] = None + audio_codes: Optional[torch.LongTensor] = None + projected_latents: Optional[torch.FloatTensor] = None + + +@dataclass +@auto_docstring +class XcodecEncoderOutput(ModelOutput): + r""" + loss (`torch.Tensor`): + Loss from the encoder model, comprising the weighted combination of the commitment and codebook losses. + quantized_representation (`torch.Tensor` of shape `(batch_size, dimension, time_steps)`, *optional*): + Quantized continuous representation of input. + audio_codes (`torch.Tensor` of shape `(batch_size, num_codebooks, time_steps)`, *optional*): + Codebook indices for each codebook (quantized discrete representation of input). + projected_latents (`torch.Tensor` of shape `(batch_size, num_codebooks * dimension, time_steps)`, *optional*): + Projected latents (continuous representation of input before quantization). """ - An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained - models. + + loss: Optional[torch.FloatTensor] = None + quantized_representation: Optional[torch.FloatTensor] = None + audio_codes: Optional[torch.FloatTensor] = None + projected_latents: Optional[torch.FloatTensor] = None + + +@dataclass +@auto_docstring +class XcodecDecoderOutput(ModelOutput): + r""" + audio_values (`torch.FloatTensor` of shape `(batch_size, input_length)`, *optional*): + Decoded audio values, obtained using the decoder part of Xcodec. """ - config_class = XcodecConfig + audio_values: Optional[torch.FloatTensor] = None + + +class XcodecVectorQuantize(nn.Module): + """ + Implementation of VQ similar to Karpathy's repo (https://github.com/karpathy/deep-vector-quantization) + + Additionally uses following tricks from improved VQGAN + (https://huggingface.co/papers/2110.04627): + 1. Factorized codes: Perform nearest neighbor lookup in low-dimensional space + for improved codebook usage + 2. l2-normalized codes: Converts euclidean distance to cosine similarity which + improves training stability + """ + + def __init__(self, config: XcodecConfig): + super().__init__() + + self.codebook_dim = config.codebook_dim + self.in_proj = nn.Conv1d(config.hidden_size, config.codebook_dim, kernel_size=1) + self.out_proj = nn.Conv1d(config.codebook_dim, config.hidden_size, kernel_size=1) + self.codebook = nn.Embedding(config.codebook_size, config.codebook_dim) + + def forward(self, hidden_state): + """ + Quantizes the input tensor using a fixed codebook and returns the corresponding codebook vectors. + + Args: + hidden_state (`torch.FloatTensor` of shape `(batch_size, dimension, time_steps)`): + Input tensor. + + Returns: + quantized_representation (`torch.Tensor`of shape `(batch_size, dimension, time_steps)`): + Quantized continuous representation of input. + commitment_loss (`torch.FloatTensor`of shape `(1)`): + Commitment loss to train encoder to predict vectors closer to codebook entries. + codebook_loss (`torch.FloatTensor`of shape `(1)`): + Codebook loss to update the codebook. + audio_codes (`torch.LongTensor` of shape `(batch_size, time_steps)`): + Codebook indices for each codebook, quantized discrete representation of input. + projected_latents (torch.FloatTensor of shape `(batch_size, num_codebooks * dimension, time_steps)`): + Projected latents (continuous representation of input before quantization). + """ + + projected_latents = self.in_proj(hidden_state) + quantized_representation, audio_codes = self.decode_latents(projected_latents) + + commitment_loss = F.mse_loss(projected_latents, quantized_representation.detach(), reduction="mean") + codebook_loss = F.mse_loss(quantized_representation, projected_latents.detach(), reduction="mean") + # noop in forward pass, straight-through gradient estimator in backward pass + quantized_representation = projected_latents + (quantized_representation - projected_latents).detach() + quantized_representation = self.out_proj(quantized_representation) + + return quantized_representation, commitment_loss, codebook_loss, audio_codes, projected_latents + + def decode_latents(self, hidden_states): + batch_size, hidden_dim, sequence_length = hidden_states.shape + encodings = hidden_states.permute(0, 2, 1).reshape(batch_size * sequence_length, hidden_dim) + codebook = self.codebook.weight # codebook: (N x D) + + # L2 normalize encodings and codebook (ViT-VQGAN) + encodings = F.normalize(encodings) + codebook = F.normalize(codebook) + + # Compute euclidean distance with codebook + l2_norm = encodings.pow(2).sum(1, keepdim=True) + dist = -(l2_norm - 2 * encodings @ codebook.t()) + codebook.pow(2).sum(1, keepdim=True).t() + + indices = dist.max(1)[1] + indices = indices.reshape(hidden_states.size(0), -1) + quantized_representation = self.codebook(indices).transpose(1, 2) + return quantized_representation, indices + + +class XcodecResidualVectorQuantize(nn.Module): + """ + ResidualVectorQuantize block - Introduced in SoundStream: An end2end neural audio codec (https://huggingface.co/papers/2107.03312) + """ + + def __init__(self, config: XcodecConfig): + super().__init__() + + n_codebooks = config.n_codebooks + quantizer_dropout = config.quantizer_dropout + + self.n_codebooks = n_codebooks + + self.quantizers = nn.ModuleList([XcodecVectorQuantize(config) for i in range(config.n_codebooks)]) + self.quantizer_dropout = quantizer_dropout + + def forward(self, hidden_state, n_quantizers: Optional[int] = None): + """ + Quantizes the input tensor using a fixed set of codebooks and returns corresponding codebook vectors. + Args: + hidden_state (`torch.Tensor` of shape `(batch_size, dimension, time_steps)`): + Input tensor to be quantized. + n_quantizers (`int`, *optional*): + Number of quantizers to use. If specified and `self.quantizer_dropout` is True, + this argument is ignored during training, and a random number of quantizers is used. + + Returns: + quantized_representation (`torch.Tensor` of shape `(batch_size, dimension, time_steps)`): + Quantized continuous representation of input. + audio_codes (`torch.Tensor` of shape `(batch_size, num_codebooks, time_steps)`): + Codebook indices for each codebook (quantized discrete representation of input). + projected_latents (`torch.Tensor` of shape `(batch_size, num_codebooks * dimension, time_steps)`): + Projected latents (continuous representation of input before quantization). + commitment_loss (`torch.Tensor` of shape `(1)`): + Commitment loss to train the encoder to predict vectors closer to codebook entries. + codebook_loss (`torch.Tensor` of shape `(1)`): + Codebook loss to update the codebook. + """ + + quantized_representation = 0 + residual = hidden_state + commitment_loss = 0 + codebook_loss = 0 + + audio_codes = [] + projected_latents = [] + + n_quantizers = n_quantizers if n_quantizers is not None else self.n_codebooks + if self.training: + n_quantizers = torch.ones((hidden_state.shape[0],)) * self.n_codebooks + 1 + dropout = torch.randint(1, self.n_codebooks + 1, (hidden_state.shape[0],)) + n_dropout = int(hidden_state.shape[0] * self.quantizer_dropout) + n_quantizers[:n_dropout] = dropout[:n_dropout] + n_quantizers = n_quantizers.to(hidden_state.device) + + for i, quantizer in enumerate(self.quantizers): + if self.training is False and i >= n_quantizers: + break + + quantized_representation_i, commitment_loss_i, codebook_loss_i, indices_i, projected_latents_i = quantizer( + residual + ) + + # Create mask to apply quantizer dropout + mask = torch.full((hidden_state.shape[0],), fill_value=i, device=hidden_state.device) < n_quantizers + quantized_representation = quantized_representation + quantized_representation_i * mask[:, None, None] + residual = residual - quantized_representation_i + + # Sum losses + commitment_loss += commitment_loss_i * mask + codebook_loss += codebook_loss_i * mask + + audio_codes.append(indices_i) + projected_latents.append(projected_latents_i) + + audio_codes = torch.stack(audio_codes, dim=1) + projected_latents = torch.cat(projected_latents, dim=1) + + return quantized_representation, audio_codes, projected_latents, commitment_loss, codebook_loss + + def from_codes(self, audio_codes: torch.Tensor): + """ + Reconstructs the continuous representation from quantized codes. + + Args: + audio_codes (`torch.Tensor` of shape `(batch_size, num_codebooks, time_steps)`): + Quantized discrete representation of input. + + Returns: + quantized_representation (`torch.Tensor`): + Quantized continuous representation of input. + projected_latents (`torch.Tensor`): + List of projected latents (continuous representations of input before quantization) + for each codebook. + audio_codes (`torch.Tensor`): + Codebook indices for each codebook. + """ + quantized_representation = 0.0 + projected_latents = [] + n_codebooks = audio_codes.shape[1] + for i in range(n_codebooks): + projected_latents_i = self.quantizers[i].codebook(audio_codes[:, i, :]).transpose(1, 2) + projected_latents.append(projected_latents_i) + quantized_representation += self.quantizers[i].out_proj(projected_latents_i) + return quantized_representation, torch.cat(projected_latents, dim=1), audio_codes + + def from_latents(self, latents: torch.Tensor): + """Reconstructs the quantized representation from unquantized latents. + + Args: + latents (`torch.Tensor` of shape `(batch_size, total_latent_dimension, time_steps)`): + Continuous representation of input after projection. + + Returns: + quantized_representation (`torch.Tensor` of shape `(batch_size, dimension, time_steps)`): + Quantized representation of the full-projected space. + quantized_latents (`torch.Tensor` of shape `(batch_size, dimension, time_steps)`): + Quantized representation of the latent space (continuous representation before quantization). + """ + quantized_representation = 0 + quantized_latents = [] + codes = [] + codebook_dims_tensor = torch.tensor([0] + [q.codebook_dim for q in self.quantizers]) + dims = torch.cumsum(codebook_dims_tensor, dim=0) + + n_codebooks = np.where(dims <= latents.shape[1])[0].max(axis=0, keepdims=True)[0] + for i in range(n_codebooks): + hidden_dim_j, hidden_dim_k = dims[i], dims[i + 1] + quantized_latents_i, codes_i = self.quantizers[i].decode_latents(latents[:, hidden_dim_j:hidden_dim_k, :]) + quantized_latents.append(quantized_latents_i) + codes.append(codes_i) + + quantized_representation_i = self.quantizers[i].out_proj(quantized_latents_i) + quantized_representation = quantized_representation + quantized_representation_i + + return quantized_representation, torch.cat(quantized_latents, dim=1) + + +@auto_docstring +class XcodecPreTrainedModel(PreTrainedAudioTokenizerBase): + config: XcodecConfig base_model_prefix = "xcodec" main_input_name = "input_values" - input_modalities = "audio" def _init_weights(self, module): - """Initialize the weights""" - if isinstance(module, nn.Linear): - module.weight.data.normal_(mean=0.0, std=self.config.initializer_range) - if module.bias is not None: - module.bias.data.zero_() - elif isinstance(module, (nn.LayerNorm, nn.GroupNorm)): - module.bias.data.zero_() - module.weight.data.fill_(1.0) - elif isinstance(module, nn.Conv1d): - nn.init.kaiming_normal_(module.weight) - if module.bias is not None: - k = math.sqrt(module.groups / (module.in_channels * module.kernel_size[0])) - nn.init.uniform_(module.bias, a=-k, b=k) - elif module.__class__.__name__ == "Snake1d": + if isinstance(module, nn.Conv1d): + nn.init.trunc_normal_(module.weight, std=0.02) + nn.init.constant_(module.bias, 0) + elif isinstance(module, Snake1d): module.alpha.data.fill_(1.0) elif isinstance(module, nn.ConvTranspose1d): module.reset_parameters() elif isinstance(module, nn.Embedding): module.weight.data.normal_(mean=0.0, std=0.02) - elif isinstance(module, XcodecModel): - # The conv1d are not handled correctly, as `self.acoustic_encoder/decoder` are initialized from a PreTrainedModel, - # but then only the submodules are used (which are not PreTrainedModels...) -> here we reinit them as in DacModel - for submodule in module.acoustic_encoder.modules(): - if isinstance(submodule, nn.Conv1d): - nn.init.trunc_normal_(submodule.weight, std=0.02) - nn.init.constant_(submodule.bias, 0) - for submodule in module.acoustic_decoder.modules(): - if isinstance(submodule, nn.Conv1d): - nn.init.trunc_normal_(submodule.weight, std=0.02) - nn.init.constant_(submodule.bias, 0) def apply_weight_norm(self): - """Apply weight norm in the acoustic encoder and decoder because the original checkpoint has weight norm applied.""" - weight_norm = torch.nn.utils.weight_norm - if hasattr(torch.nn.utils.parametrizations, "weight_norm"): - weight_norm = torch.nn.utils.parametrizations.weight_norm - - weight_norm(self.acoustic_encoder.conv1) - weight_norm(self.acoustic_encoder.conv2) - - for block in self.acoustic_encoder.block: - weight_norm(block.conv1) - for res_unit in (block.res_unit1, block.res_unit2, block.res_unit3): - weight_norm(res_unit.conv1) - weight_norm(res_unit.conv2) - - weight_norm(self.acoustic_decoder.conv1, name="weight") - weight_norm(self.acoustic_decoder.conv2, name="weight") - - for block in self.acoustic_decoder.block: - weight_norm(block.conv_t1, name="weight") - for res_unit in (block.res_unit1, block.res_unit2, block.res_unit3): - weight_norm(res_unit.conv1, name="weight") - weight_norm(res_unit.conv2, name="weight") + weight_norm = nn.utils.weight_norm + if hasattr(nn.utils.parametrizations, "weight_norm"): + weight_norm = nn.utils.parametrizations.weight_norm + + for layer in self.quantizer.quantizers: + weight_norm(layer.in_proj) + weight_norm(layer.out_proj) + + weight_norm(self.encoder.conv1) + weight_norm(self.encoder.conv2) + + for layer in self.encoder.block: + weight_norm(layer.conv1) + weight_norm(layer.res_unit1.conv1) + weight_norm(layer.res_unit1.conv2) + weight_norm(layer.res_unit2.conv1) + weight_norm(layer.res_unit2.conv2) + weight_norm(layer.res_unit3.conv1) + weight_norm(layer.res_unit3.conv2) + + weight_norm(self.decoder.conv1) + weight_norm(self.decoder.conv2) + + for layer in self.decoder.block: + weight_norm(layer.conv_t1) + weight_norm(layer.res_unit1.conv1) + weight_norm(layer.res_unit1.conv2) + weight_norm(layer.res_unit2.conv1) + weight_norm(layer.res_unit2.conv2) + weight_norm(layer.res_unit3.conv1) + weight_norm(layer.res_unit3.conv2) def remove_weight_norm(self): - """Remove the weight norm from the acoustic encoder and decoder.""" - for module in (self.acoustic_encoder, self.acoustic_decoder): - for m in module.modules(): - try: - torch.nn.utils.remove_weight_norm(m, name="weight") - except (ValueError, AttributeError): - pass - if hasattr(m, "parametrizations") and "weight" in m.parametrizations: - torch.nn.utils.parametrize.remove_parametrizations(m, "weight", leave_parametrized=True) - - -@auto_docstring(custom_intro="""The Xcodec neural audio codec model.""") + for layer in self.quantizer.quantizers: + nn.utils.remove_weight_norm(layer.in_proj) + nn.utils.remove_weight_norm(layer.out_proj) + + nn.utils.remove_weight_norm(self.encoder.conv1) + nn.utils.remove_weight_norm(self.encoder.conv2) + + for layer in self.encoder.block: + nn.utils.remove_weight_norm(layer.conv1) + nn.utils.remove_weight_norm(layer.res_unit1.conv1) + nn.utils.remove_weight_norm(layer.res_unit1.conv2) + nn.utils.remove_weight_norm(layer.res_unit2.conv1) + nn.utils.remove_weight_norm(layer.res_unit2.conv2) + nn.utils.remove_weight_norm(layer.res_unit3.conv1) + nn.utils.remove_weight_norm(layer.res_unit3.conv2) + + nn.utils.remove_weight_norm(self.decoder.conv1) + nn.utils.remove_weight_norm(self.decoder.conv2) + + for layer in self.decoder.block: + nn.utils.remove_weight_norm(layer.conv_t1) + nn.utils.remove_weight_norm(layer.res_unit1.conv1) + nn.utils.remove_weight_norm(layer.res_unit1.conv2) + nn.utils.remove_weight_norm(layer.res_unit2.conv1) + nn.utils.remove_weight_norm(layer.res_unit2.conv2) + nn.utils.remove_weight_norm(layer.res_unit3.conv1) + nn.utils.remove_weight_norm(layer.res_unit3.conv2) + + +@auto_docstring( + custom_intro=""" + The XCODEC (Descript Audio Codec) model. + """ +) class XcodecModel(XcodecPreTrainedModel): - def __init__(self, config): + input_modalities = "audio" + + def __init__(self, config: XcodecConfig): super().__init__(config) self.config = config - self.pad = config.hop_length // 2 - acoustic_model = AutoModel.from_config(config.acoustic_model_config) - self.acoustic_encoder = acoustic_model.encoder - self.acoustic_decoder = acoustic_model.decoder - self._adjust_dac_decoder(self.acoustic_decoder) - self.encoder_semantic = SemanticEncoder(config) - self.decoder_semantic = SemanticDecoder(config) - self.semantic_model = AutoModel.from_config(config.semantic_model_config).eval() - self.fc = nn.Linear(config.hidden_size, config.hidden_size) - self.fc1 = nn.Linear(config.hidden_size, config.semantic_model_config.hidden_size) - self.fc2 = nn.Linear(config.hidden_size, config.acoustic_model_config.hidden_size) - self.quantizer = XcodecResidualVectorQuantization(config) - # Initialize weights and apply final processing - self.post_init() + self.encoder = XcodecEncoder(config) + self.decoder = XcodecDecoder(config) - @staticmethod - def _adjust_dac_decoder(decoder: nn.Module): - r""" - DAC implemented in Xcodec is slightly different from the HF version. - DAC in Xcodec adjusts the output padding in every ConvTranspose1d in the decoder and removes - the final `nn.Tanh` activation function. - """ - for module in decoder.modules(): - if isinstance(module, nn.ConvTranspose1d): - stride = module.stride[0] if isinstance(module.stride, tuple) else module.stride - module.output_padding = (stride % 2,) - if hasattr(decoder, "tanh") and isinstance(decoder.tanh, nn.Tanh): - decoder.tanh = nn.Identity() - - def _extract_semantic_features(self, input_values: torch.FloatTensor) -> torch.FloatTensor: - input_values = input_values[:, 0, :] - input_values = F.pad(input_values, (self.pad, self.pad)) - with torch.no_grad(): - outputs = self.semantic_model(input_values, output_hidden_states=True) - hidden_states = outputs.hidden_states + self.quantizer = XcodecResidualVectorQuantize(config) + + self.bits_per_codebook = int(math.log2(self.config.codebook_size)) + if 2**self.bits_per_codebook != self.config.codebook_size: + raise ValueError("The codebook_size must be a power of 2.") - stacked = torch.stack(hidden_states, dim=1) - return stacked.mean(dim=1) + # Initialize weights and apply final processing + self.post_init() + @can_return_tuple @auto_docstring def encode( self, input_values: torch.Tensor, - bandwidth: Optional[float] = None, - return_dict: Optional[bool] = None, - ) -> Union[torch.Tensor, XcodecEncoderOutput]: + n_quantizers: Optional[int] = None, + ): 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. + input_values (`torch.Tensor of shape `(batch_size, 1, time_steps)`): + Input audio data to encode, + n_quantizers (int, *optional*): + Number of quantizers to use. If None, all quantizers are used. Default is None. """ - return_dict = return_dict if return_dict is not None else self.config.return_dict - - 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] - elif bandwidth not in self.config.target_bandwidths: - raise ValueError( - f"This model doesn't support the bandwidth {bandwidth}. Select one of {self.config.target_bandwidths}." - ) - - e_semantic_input = self._extract_semantic_features(input_values).detach() - e_semantic = self.encoder_semantic(e_semantic_input.transpose(1, 2)) - e_acoustic = self.acoustic_encoder(input_values) - - if e_acoustic.shape[2] != e_semantic.shape[2]: - # make sure they line up if frames don't match - e_acoustic = self.acoustic_encoder(F.pad(input_values[:, 0, :], (self.pad, self.pad)).unsqueeze(1)) - - embeddings = torch.cat([e_acoustic, e_semantic], dim=1) - embeddings = self.fc(embeddings.transpose(1, 2)).transpose(1, 2) - audio_codes = self.quantizer.encode(embeddings, bandwidth) - audio_codes = audio_codes.transpose(0, 1) + quantized_representation = self.encoder(input_values) + quantized_representation, audio_codes, projected_latents, commitment_loss, codebook_loss = self.quantizer( + quantized_representation, n_quantizers + ) - if not return_dict: - return audio_codes + loss = self.config.commitment_loss_weight * commitment_loss + self.config.codebook_loss_weight * codebook_loss - return XcodecEncoderOutput(audio_codes) + return XcodecEncoderOutput(loss, quantized_representation, audio_codes, projected_latents) + @can_return_tuple @auto_docstring def decode( self, - audio_codes: torch.Tensor, - return_dict: Optional[bool] = None, - ) -> Union[torch.Tensor, XcodecDecoderOutput]: + quantized_representation: Optional[torch.Tensor] = None, + audio_codes: Optional[torch.Tensor] = None, + ): 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. + quantized_representation (torch.Tensor of shape `(batch_size, dimension, time_steps)`, *optional*): + Quantized continuous representation of input. + audio_codes (`torch.Tensor` of shape `(batch_size, num_codebooks, time_steps)`, *optional*): + 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 = return_dict if return_dict is not None else self.config.return_dict - audio_codes = audio_codes.transpose(0, 1) - quantized = self.quantizer.decode(audio_codes) - quantized_acoustic = self.fc2(quantized.transpose(1, 2)).transpose(1, 2) - audio_values = self.acoustic_decoder(quantized_acoustic) + if quantized_representation is None and audio_codes is None: + raise ValueError("Either `quantized_representation` or `audio_codes` must be provided.") + + if audio_codes is not None: + quantized_representation = self.quantizer.from_codes(audio_codes)[0] - if not return_dict: - return audio_values + audio_values = self.decoder(quantized_representation).squeeze(1) return XcodecDecoderOutput(audio_values) @@ -522,60 +794,40 @@ def decode( def forward( self, input_values: torch.Tensor, - audio_codes: Optional[torch.Tensor] = None, - bandwidth: Optional[float] = None, - return_dict: Optional[bool] = None, - ) -> Union[tuple[torch.Tensor, torch.Tensor], XcodecOutput]: + n_quantizers: Optional[int] = None, + ): 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. + input_values (`torch.Tensor` of shape `(batch_size, 1, time_steps)`): + Audio data to encode. + n_quantizers (`int`, *optional*): + Number of quantizers to use. If `None`, all quantizers are used. Default is `None`. - Returns: - `XcodecOutput` or tuple `(audio_codes, audio_values)`: - - `audio_codes` of shape `(batch_size, num_quantizers, codes_length)`: the quantized discrete codes. - - `audio_values` of shape `(batch_size, channels, num_samples)`: the reconstructed audio waveform given the codes. - - Example: + Examples: ```python - >>> from datasets import load_dataset - >>> from transformers import AutoFeatureExtractor, XcodecModel - - >>> model_id = "hf-audio/xcodec-hubert-librispeech" - >>> model = XcodecModel.from_pretrained(model_id) - >>> feature_extractor = AutoFeatureExtractor.from_pretrained(model_id) - - >>> dataset = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation") - >>> dataset = dataset.cast_column("audio", Audio(sampling_rate=feature_extractor.sampling_rate)) - >>> audio_sample = dataset[0]['audio']['array'] - - >>> inputs = feature_extractor(raw_audio=audio_sample, return_tensors="pt") - - >>> outputs = model(**inputs) - >>> audio_codes = outputs.audio_codes - >>> audio_values = outputs.audio_values - ``` - """ - return_dict = return_dict if return_dict is not None else self.config.return_dict + >>> from datasets import load_dataset, Audio + >>> from transformers import XcodecModel, AutoProcessor + >>> librispeech_dummy = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation") + + >>> model = XcodecModel.from_pretrained("descript/xcodec_16khz") + >>> processor = AutoProcessor.from_pretrained("descript/xcodec_16khz") + >>> librispeech_dummy = librispeech_dummy.cast_column("audio", Audio(sampling_rate=processor.sampling_rate)) + >>> audio_sample = librispeech_dummy[-1]["audio"]["array"] + >>> inputs = processor(raw_audio=audio_sample, sampling_rate=processor.sampling_rate, return_tensors="pt") + + >>> encoder_outputs = model.encode(inputs["input_values"]) + >>> # Get the intermediate audio codes + >>> audio_codes = encoder_outputs.audio_codes + >>> # Reconstruct the audio from its quantized representation + >>> audio_values = model.decode(encoder_outputs.quantized_representation) + >>> # or the equivalent with a forward pass + >>> audio_values = model(inputs["input_values"]).audio_values + ```""" 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][..., :length] - - if not return_dict: - return (audio_codes, audio_values) - - return XcodecOutput(audio_codes=audio_codes, audio_values=audio_values) - + loss, quantized_representation, audio_codes, projected_latents = self.encode( + input_values, n_quantizers, return_dict=False + ) + audio_values = self.decode(quantized_representation, return_dict=False)[0][..., :length] -__all__ = ["XcodecModel", "XcodecPreTrainedModel"] + return XcodecOutput(loss, audio_values, quantized_representation, audio_codes, projected_latents) diff --git a/src/transformers/models/xcodec/modular_xcodec.py b/src/transformers/models/xcodec/modular_xcodec.py new file mode 100644 index 000000000000..baca1fe47569 --- /dev/null +++ b/src/transformers/models/xcodec/modular_xcodec.py @@ -0,0 +1,136 @@ +# coding=utf-8 +# Copyright 2025 Boson AI and 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. +"""XcodecModel model.""" + +import math +from typing import Optional, Union + +import torch +import torch.nn as nn +import torch.nn.functional as F +import torchaudio + +from .configuration_xcodec import XcodecConfig +from ...utils import auto_docstring, can_return_tuple +from ..auto.modeling_auto import AutoModel +from ..dac.modeling_dac import DacDecoder, DacDecoderBlock, DacEncoder, DacModel +from ..encodec.modeling_encodec import EncodecResidualVectorQuantizer + + +class XcodecAcousticEncoder(DacEncoder):... + + +class XcodecSemanticEncoderLayer(nn.Module): + def __init__(self, config): + super().__init__() + self.act_fn = nn.ELU() + self.conv1 = nn.Conv1d(config.semantic_hidden_size, config.semantic_hidden_size, kernel_size=3, bias=False) + self.conv2 = nn.Conv1d(config.semantic_hidden_size, config.semantic_hidden_size, kernel_size=1, bias=False) + self.conv3 = nn.Conv1d(config.semantic_hidden_size, config.semantic_hidden_size, kernel_size=3, bias=False) + + def forward(self, hidden_state): + residual = hidden_state + hidden_state = self.conv1(self.act_fn(hidden_state)) + hidden_state = self.conv2(self.act_fn(hidden_state)) + hidden_state = self.conv3(residual + hidden_state) + return hidden_state + + +class XcodecSemanticEncoder(nn.Module): + def __init__(self, config): + super().__init__() + self.input_conv = nn.Conv1d(config.semantic_hidden_size, config.semantic_hidden_size, kernel_size=3, padding=1, bias=False) + self.layers = nn.ModuleList( + [XcodecSemanticEncoderLayer(config) for _ in range(config.num_layers)] + ) + + def forward(self, hidden_state): + hidden_state = self.input_conv(hidden_state) + for layer in self.layers: + hidden_state = layer(hidden_state) + return hidden_state + + +class XcodecEncoder(nn.Module): + def __init__(self, config: XcodecConfig): + super().__init__() + self.acoustic_encoder = XcodecAcousticEncoder(config) + self.semantic_encoder = XcodecSemanticEncoder(config) + self.semantic_model = AutoModel.from_config(config.semantic_config) + + self.sample_rate = config.sample_rate + self.semantic_sample_rate = config.semantic_sample_rate + + def _extract_semantic_features(self, input_values): + with torch.no_grad(): + input_values = torchaudio.functional.resample( + input_values, self.config.sampling_rate, self.config.semantic_sample_rate + ) + input_values = input_values[:, 0, :] + input_values = F.pad(input_values, (self.config.pad, self.config.pad)) + outputs = self.semantic_model(input_values, output_hidden_states=True) + hidden_states = outputs.hidden_states + stacked = torch.stack(hidden_states, dim=1) + semantic_features = stacked.mean(dim=1) + semantic_features = semantic_features[:, :: self.config.semantic_downsample_factor, :] + return semantic_features + + def forward(self, input_values: torch.Tensor): + acoustic_embeds = self.acoustic_encoder(input_values) + + semantic_features = self._extract_semantic_features(input_values) + semantic_embeds = self.semantic_encoder(semantic_features.transpose(1, 2)) + + input_embeds = torch.cat([acoustic_embeds, semantic_embeds], dim=1) + return input_embeds + + +class XcodecDecoderBlock(DacDecoderBlock): + def __init__(self, config, stride: int = 1, stride_index: int = 1): + super().__init__(config) + input_dim = config.decoder_hidden_size // 2**stride_index + output_dim = config.decoder_hidden_size // 2 ** (stride_index + 1) + self.conv_t1 = nn.ConvTranspose1d( + input_dim, + output_dim, + kernel_size=2 * stride, + stride=stride, + padding=math.ceil(stride / 2), + output_padding=(stride % 2,), + ) + + +class XcodecDecoder(DacDecoder): + def __init__(self, config: XcodecConfig): + super().__init__(config) + input_channel = config.acoustic_hidden_size + del self.tanh + + def forward(self, hidden_states): + hidden_states = self.conv1(hidden_states) + for layer in self.block: + hidden_states = layer(hidden_states) + + hidden_states = self.snake1(hidden_states) + hidden_states = self.conv2(hidden_states) + + return hidden_states + + +class XcodecResidualVectorQuantization(EncodecResidualVectorQuantizer): ... + + +# TODO: @eustlb, consider adding bandwidth handling in dac in order to standardize codec usage +class XcodecModel(DacModel): ... \ No newline at end of file From d69d0ef741f14d9a9336a1d4800127293316be13 Mon Sep 17 00:00:00 2001 From: Eustache Le Bihan Date: Wed, 5 Nov 2025 15:59:36 +0100 Subject: [PATCH 3/5] clean config --- .../models/xcodec/configuration_xcodec.py | 182 ++++++++---------- .../models/xcodec/modeling_xcodec.py | 2 +- .../models/xcodec/modular_xcodec.py | 3 +- 3 files changed, 86 insertions(+), 101 deletions(-) diff --git a/src/transformers/models/xcodec/configuration_xcodec.py b/src/transformers/models/xcodec/configuration_xcodec.py index 7281e1719033..44ff37e3c68f 100644 --- a/src/transformers/models/xcodec/configuration_xcodec.py +++ b/src/transformers/models/xcodec/configuration_xcodec.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 The HuggingFace Inc. team. All rights reserved. +# Copyright 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. @@ -12,37 +12,34 @@ # 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 AutoConfig, DacConfig, HubertConfig, WavLMConfig +from ...configuration_utils import PretrainedConfig +from ..auto import CONFIG_MAPPING, AutoConfig -from ...configuration_utils import PreTrainedConfig -from ...utils import logging - -logger = logging.get_logger(__name__) - - -class XcodecConfig(PreTrainedConfig): +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. + XcodecModel according to the specified arguments, defining the model architecture. - Configuration objects inherit from [`PreTrainedConfig`] and can be used to control the model outputs. Read the - documentation from [`PreTrainedConfig`] for more information. + Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the + documentation from [`PretrainedConfig`] for more information. + + Instantiating a configuration with the defaults will yield a similar configuration to that of the + [bosonai/higgs-audio-v2-tokenizer](https://huggingface.co/bosonai/higgs-audio-v2-tokenizer) architecture. 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. - sample_rate (`int`, *optional*, defaults to 16000): + sample_rate (`int`, *optional*, defaults to 24000): The sampling rate at which the audio waveform should be digitalized, in hertz (Hz). + semantic_sample_rate (`int`, *optional*, defaults to 16000): + The sampling rate at which the semantic model was trained with, in hertz (Hz). kernel_size (`int`, *optional*, defaults to 3): Kernel size for the initial semantic convolution. channel_ratios (`List[float]`, *optional*, defaults to `[1, 1]`): @@ -55,132 +52,121 @@ class XcodecConfig(PreTrainedConfig): Kernel size inside each ResidualUnit in semantic blocks. codebook_size (`int`, *optional*, defaults to 1024): Number of entries in each residual quantizer's codebook. - codebook_dim (`int`, *optional*): - Dimensionality of each codebook vector. Defaults to sum of hidden size of acoustic and semantic models. + num_quantizers (`int`, *optional*, defaults to 8): + Number of sequential quantizers (codebooks) in the RVQ stack. + codebook_dim (`int`, *optional*, defaults to 64): + Dimensionality of each codebook vector. initializer_range (`float`, *optional*, defaults to 0.02): Standard deviation of the truncated normal initializer for all weight matrices. acoustic_model_config (`Union[Dict, DacConfig]`, *optional*): An instance of the configuration for the acoustic (DAC) model. - semantic_model_config (`Union[Dict, HubertConfig, WavLMConfig]`, *optional*): + semantic_config (`Union[Dict, HubertConfig]`, *optional*): An instance of the configuration object for the semantic (HuBERT) model. + downsample_mode (`str`, *optional*, defaults to `"step_down"`): + The downsample mode for the semantic features. + pad (`int`, *optional*, defaults to 160): + Padding size for the input to the semantic model. + downsample_factor (`int`, *optional*, defaults to 320): + The downsample factor used to compute actual downsample factor for the semantic features. Example: ```python >>> from transformers import XcodecModel, XcodecConfig - >>> # Initializing configuration + >>> # Initializing a " " style configuration >>> configuration = XcodecConfig() - >>> # Initializing a model (with random weights) from the configuration + >>> # 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": AutoConfig, + model_type = "higgs_audio_tokenizer" + sub_configs = {"semantic_config": AutoConfig} + attribute_map = { + "sampling_rate": "sample_rate", } - def __init__( - self, - target_bandwidths: Optional[list[float]] = None, - sample_rate: int = 16000, - 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, - codebook_size: int = 1024, - codebook_dim: Optional[int] = None, - initializer_range: float = 0.02, - acoustic_model_config: Optional[Union[dict, DacConfig]] = None, - semantic_model_config: Optional[Union[dict, HubertConfig]] = None, + self, + encoder_hidden_size=64, + acoustic_hidden_size=256, + downsampling_ratios=[8, 5, 4, 2, 3], + decoder_hidden_size=1024, + upsampling_ratios=[8, 5, 4, 2, 3], + n_codebooks=9, + strides=[1, 1], + channel_ratios=[1, 1], + kernel_size=3, + block_dilations=[1, 1], + semantic_config=None, + unit_kernel_size=3, + target_bandwidths=[0.5, 1, 1.5, 2, 4], + sample_rate=24000, + semantic_sample_rate=16000, + codebook_size=1024, + num_quantizers=8, + codebook_dim=64, + initializer_range=0.02, + pad=160, **kwargs, ): - if acoustic_model_config is None: - self.acoustic_model_config = DacConfig( - encoder_hidden_size=64, - # NOTE: original DAC uses [2, 4, 8, 8] `downsampling ratios`, namely reverse of `upsampling_ratios` - # (not sure if intentional by Xcodec but we keep it) - 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 - else: - raise ValueError( - f"acoustic_model_config must be a dict or DacConfig instance, but got {type(acoustic_model_config)}" - ) - - if semantic_model_config is None: - self.semantic_model_config = HubertConfig() - elif isinstance(semantic_model_config, dict): - if "_name_or_path" in semantic_model_config: - # If the config is a path, load it using AutoConfig - self.semantic_model_config = AutoConfig.from_pretrained(semantic_model_config["_name_or_path"]) - else: - # assume HubertConfig as probably created from scratch - logger.warning( - "Could not determine semantic model type from config architecture. Defaulting to `HubertConfig`." - ) - self.semantic_model_config = HubertConfig(**semantic_model_config) - elif isinstance(semantic_model_config, WavLMConfig) or isinstance(semantic_model_config, HubertConfig): - self.semantic_model_config = semantic_model_config - else: - raise ValueError( - f"semantic_model_config must be a dict, HubertConfig, or WavLMConfig instance, but got {type(semantic_model_config)}" - ) - - if target_bandwidths is None: - target_bandwidths = [0.5, 1, 1.5, 2, 4] + # Handle semantic model config + if isinstance(semantic_config, dict): + semantic_config["model_type"] = semantic_config.get("model_type", "hubert") + semantic_config = CONFIG_MAPPING[semantic_config["model_type"]](**semantic_config) + elif semantic_config is None: + semantic_config = CONFIG_MAPPING["hubert"]() + + self.semantic_config = semantic_config + + self.acoustic_hidden_size = acoustic_hidden_size + self.encoder_hidden_size = encoder_hidden_size + self.downsampling_ratios = downsampling_ratios + self.n_codebooks = n_codebooks - self.target_bandwidths = target_bandwidths - self.sample_rate = sample_rate - self.kernel_size = kernel_size - self.channel_ratios = channel_ratios self.strides = strides + self.channel_ratios = channel_ratios + self.kernel_size = kernel_size self.block_dilations = block_dilations self.unit_kernel_size = unit_kernel_size + + self.decoder_hidden_size = decoder_hidden_size + self.target_bandwidths = target_bandwidths + self.sample_rate = sample_rate + self.semantic_sample_rate = semantic_sample_rate self.codebook_size = codebook_size - self.initializer_range = initializer_range - if codebook_dim is None: - codebook_dim = self.acoustic_model_config.hidden_size + self.semantic_model_config.hidden_size + self.num_quantizers = num_quantizers self.codebook_dim = codebook_dim + self.initializer_range = initializer_range + self.pad = pad + self.upsampling_ratios = upsampling_ratios super().__init__(**kwargs) @property def frame_rate(self) -> int: - return math.ceil(self.sample_rate / self.hop_length) - - @property - def semantic_hidden_size(self) -> int: - return self.semantic_model_config.hidden_size + return math.ceil(self.sample_rate / np.prod(self.upsampling_ratios)) @property def hop_length(self) -> int: - return int(np.prod(self.acoustic_model_config.downsampling_ratios)) + return int(np.prod(self.downsampling_ratios)) @property - def codebook_nbits(self) -> int: - return math.ceil(math.log2(self.codebook_size)) + def semantic_hidden_size(self) -> int: + return self.semantic_config.hidden_size @property def hidden_size(self) -> int: - return self.acoustic_model_config.hidden_size + self.semantic_model_config.hidden_size + return self.acoustic_hidden_size + self.semantic_hidden_size @property - def num_quantizers(self) -> int: - return int(1000 * self.target_bandwidths[-1] // (self.frame_rate * self.codebook_nbits)) + def semantic_downsample_factor(self) -> int: + return int( + self.hop_length / (self.sample_rate / self.semantic_sample_rate) / self.downsample_factor + ) __all__ = ["XcodecConfig"] diff --git a/src/transformers/models/xcodec/modeling_xcodec.py b/src/transformers/models/xcodec/modeling_xcodec.py index 6a8f5016f707..495972d59f77 100644 --- a/src/transformers/models/xcodec/modeling_xcodec.py +++ b/src/transformers/models/xcodec/modeling_xcodec.py @@ -5,7 +5,7 @@ # modular_xcodec.py file directly. One of our CI enforces this. # 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 # coding=utf-8 -# Copyright 2025 Boson AI and The HuggingFace Team. All rights reserved. +# Copyright 2025 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. diff --git a/src/transformers/models/xcodec/modular_xcodec.py b/src/transformers/models/xcodec/modular_xcodec.py index baca1fe47569..b00614d6d21d 100644 --- a/src/transformers/models/xcodec/modular_xcodec.py +++ b/src/transformers/models/xcodec/modular_xcodec.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2025 Boson AI and The HuggingFace Team. All rights reserved. +# Copyright 2025 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. @@ -12,7 +12,6 @@ # 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. -"""XcodecModel model.""" import math from typing import Optional, Union From 51b4f4c3af423be47376b75507fb48bc104c16a3 Mon Sep 17 00:00:00 2001 From: Eustache Le Bihan Date: Wed, 5 Nov 2025 16:02:25 +0100 Subject: [PATCH 4/5] make --- src/transformers/models/xcodec/configuration_xcodec.py | 3 +-- src/transformers/models/xcodec/modular_xcodec.py | 6 ++---- 2 files changed, 3 insertions(+), 6 deletions(-) diff --git a/src/transformers/models/xcodec/configuration_xcodec.py b/src/transformers/models/xcodec/configuration_xcodec.py index 44ff37e3c68f..523158ee36e4 100644 --- a/src/transformers/models/xcodec/configuration_xcodec.py +++ b/src/transformers/models/xcodec/configuration_xcodec.py @@ -14,7 +14,6 @@ # limitations under the License. import math -from typing import Optional, Union import numpy as np @@ -90,7 +89,7 @@ class XcodecConfig(PretrainedConfig): "sampling_rate": "sample_rate", } def __init__( - self, + self, encoder_hidden_size=64, acoustic_hidden_size=256, downsampling_ratios=[8, 5, 4, 2, 3], diff --git a/src/transformers/models/xcodec/modular_xcodec.py b/src/transformers/models/xcodec/modular_xcodec.py index b00614d6d21d..ccaa0dede3a1 100644 --- a/src/transformers/models/xcodec/modular_xcodec.py +++ b/src/transformers/models/xcodec/modular_xcodec.py @@ -14,18 +14,16 @@ # limitations under the License. import math -from typing import Optional, Union import torch import torch.nn as nn import torch.nn.functional as F import torchaudio -from .configuration_xcodec import XcodecConfig -from ...utils import auto_docstring, can_return_tuple from ..auto.modeling_auto import AutoModel from ..dac.modeling_dac import DacDecoder, DacDecoderBlock, DacEncoder, DacModel from ..encodec.modeling_encodec import EncodecResidualVectorQuantizer +from .configuration_xcodec import XcodecConfig class XcodecAcousticEncoder(DacEncoder):... @@ -132,4 +130,4 @@ class XcodecResidualVectorQuantization(EncodecResidualVectorQuantizer): ... # TODO: @eustlb, consider adding bandwidth handling in dac in order to standardize codec usage -class XcodecModel(DacModel): ... \ No newline at end of file +class XcodecModel(DacModel): ... From cc6e0a49d435504eda1a8078a9c53ca78e23621c Mon Sep 17 00:00:00 2001 From: Eustache Le Bihan Date: Thu, 6 Nov 2025 13:42:43 +0100 Subject: [PATCH 5/5] updates --- .../models/xcodec/configuration_xcodec.py | 16 +- .../models/xcodec/modeling_xcodec.py | 330 ++++++------------ .../models/xcodec/modular_xcodec.py | 59 +++- 3 files changed, 170 insertions(+), 235 deletions(-) diff --git a/src/transformers/models/xcodec/configuration_xcodec.py b/src/transformers/models/xcodec/configuration_xcodec.py index 523158ee36e4..27c963bfc89d 100644 --- a/src/transformers/models/xcodec/configuration_xcodec.py +++ b/src/transformers/models/xcodec/configuration_xcodec.py @@ -89,12 +89,12 @@ class XcodecConfig(PretrainedConfig): "sampling_rate": "sample_rate", } def __init__( - self, + self, encoder_hidden_size=64, acoustic_hidden_size=256, - downsampling_ratios=[8, 5, 4, 2, 3], + downsampling_ratios=[8, 5, 4, 2], decoder_hidden_size=1024, - upsampling_ratios=[8, 5, 4, 2, 3], + upsampling_ratios=[8, 5, 4, 2], n_codebooks=9, strides=[1, 1], channel_ratios=[1, 1], @@ -143,6 +143,10 @@ def __init__( self.pad = pad self.upsampling_ratios = upsampling_ratios + self.num_layers = 2 + self.num_residual_layers = 2 + self.hidden_size = 256 + super().__init__(**kwargs) @property @@ -157,9 +161,9 @@ def hop_length(self) -> int: def semantic_hidden_size(self) -> int: return self.semantic_config.hidden_size - @property - def hidden_size(self) -> int: - return self.acoustic_hidden_size + self.semantic_hidden_size + # @property + # def hidden_size(self) -> int: + # return self.acoustic_hidden_size + self.semantic_hidden_size @property def semantic_downsample_factor(self) -> int: diff --git a/src/transformers/models/xcodec/modeling_xcodec.py b/src/transformers/models/xcodec/modeling_xcodec.py index 495972d59f77..4ddc39b1c3f5 100644 --- a/src/transformers/models/xcodec/modeling_xcodec.py +++ b/src/transformers/models/xcodec/modeling_xcodec.py @@ -30,7 +30,7 @@ import torchaudio from ...modeling_utils import PreTrainedAudioTokenizerBase -from ...utils import ModelOutput, auto_docstring, can_return_tuple +from ...utils import ModelOutput, auto_docstring from ..auto.modeling_auto import AutoModel from .configuration_xcodec import XcodecConfig @@ -112,6 +112,8 @@ def forward(self, hidden_state): return hidden_state +class XcodecAcousticEncoder(nn.Module): + """XCODEC_ACOUSTIC Encoder""" class XcodecAcousticEncoder(nn.Module): """XCODEC_ACOUSTIC Encoder""" @@ -146,13 +148,14 @@ def forward(self, hidden_state): return hidden_state -class XcodecSemanticEncoderLayer(nn.Module): +class XcodecSemanticEncoderResidualLayer(nn.Module): def __init__(self, config): super().__init__() self.act_fn = nn.ELU() - self.conv1 = nn.Conv1d(config.semantic_hidden_size, config.semantic_hidden_size, kernel_size=3, bias=False) + self.conv1 = nn.Conv1d( + config.semantic_hidden_size, config.semantic_hidden_size, kernel_size=3, padding=1, bias=False + ) self.conv2 = nn.Conv1d(config.semantic_hidden_size, config.semantic_hidden_size, kernel_size=1, bias=False) - self.conv3 = nn.Conv1d(config.semantic_hidden_size, config.semantic_hidden_size, kernel_size=3, bias=False) def forward(self, hidden_state): residual = hidden_state @@ -162,6 +165,21 @@ def forward(self, hidden_state): return hidden_state +class XcodecSemanticEncoderLayer(nn.Module): + def __init__(self, config): + super().__init__() + self.residual_layers = nn.ModuleList( + [XcodecSemanticEncoderResidualLayer(config) for _ in range(config.num_residual_layers)] + ) + self.conv = nn.Conv1d(config.semantic_hidden_size, config.semantic_hidden_size, kernel_size=3, padding=1) + + def forward(self, hidden_state): + for layer in self.residual_layers: + hidden_state = layer(hidden_state) + hidden_state = self.conv(hidden_state) + return hidden_state + + class XcodecSemanticEncoder(nn.Module): def __init__(self, config): super().__init__() @@ -183,6 +201,10 @@ def __init__(self, config: XcodecConfig): self.acoustic_encoder = XcodecAcousticEncoder(config) self.semantic_encoder = XcodecSemanticEncoder(config) self.semantic_model = AutoModel.from_config(config.semantic_config) + self.linear = nn.Linear( + config.semantic_hidden_size + config.acoustic_hidden_size, + config.semantic_hidden_size + config.acoustic_hidden_size, + ) self.sample_rate = config.sample_rate self.semantic_sample_rate = config.semantic_sample_rate @@ -211,8 +233,8 @@ def forward(self, input_values: torch.Tensor): return input_embeds -class XcodecDecoderBlock(nn.Module): - """Decoder block used in XCODEC decoder.""" +class XcodecAcousticDecoderBlock(nn.Module): + """Decoder block used in XCODEC_ACOUSTIC decoder.""" def __init__(self, config, stride: int = 1, stride_index: int = 1): super().__init__() @@ -241,9 +263,42 @@ def forward(self, hidden_state): return hidden_state + input_embeds = torch.cat([acoustic_embeds, semantic_embeds], dim=1) + return input_embeds -class XcodecDecoder(nn.Module): - """XCODEC Decoder""" +class XcodecDecoderBlock(nn.Module): + """Decoder block used in XCODEC decoder.""" + + def __init__(self, config: XcodecConfig, stride: int = 1, stride_index: int = 1): + super().__init__() + + input_dim = config.decoder_hidden_size // 2**stride_index + output_dim = config.decoder_hidden_size // 2 ** (stride_index + 1) + self.snake1 = Snake1d(input_dim) + self.conv_t1 = nn.ConvTranspose1d( + input_dim, + output_dim, + kernel_size=2 * stride, + stride=stride, + padding=math.ceil(stride / 2), + ) + + self.res_unit1 = XcodecResidualUnit(output_dim, dilation=1) + self.res_unit2 = XcodecResidualUnit(output_dim, dilation=3) + self.res_unit3 = XcodecResidualUnit(output_dim, dilation=9) + + def forward(self, hidden_state): + hidden_state = self.snake1(hidden_state) + hidden_state = self.conv_t1(hidden_state) + hidden_state = self.res_unit1(hidden_state) + hidden_state = self.res_unit2(hidden_state) + hidden_state = self.res_unit3(hidden_state) + + return hidden_state + + +class XcodecAcousticDecoder(nn.Module): + """XCODEC_ACOUSTIC Decoder""" def __init__(self, config: XcodecConfig): super().__init__() @@ -275,6 +330,18 @@ def forward(self, hidden_states): return hidden_states +class XcodecSemanticDecoder(nn.Module): + def __init__(self, config): + super().__init__() + self.input_conv = nn.Conv1d( + config.semantic_hidden_size, config.semantic_hidden_size, kernel_size=3, padding=1, bias=False + ) + self.layers = nn.ModuleList([XcodecSemanticEncoderLayer(config) for _ in range(config.num_layers)]) + self.output_conv = nn.Conv1d( + config.semantic_hidden_size, config.semantic_hidden_size, kernel_size=3, padding=1, bias=False + ) + + class XcodecEuclideanCodebook(nn.Module): """Codebook with Euclidean distance.""" @@ -331,7 +398,7 @@ def decode(self, embed_ind): return quantize -class XcodecResidualVectorQuantization(nn.Module): +class XcodecResidualVectorQuantizer(nn.Module): """Residual Vector Quantizer.""" def __init__(self, config: XcodecConfig): @@ -429,215 +496,6 @@ class XcodecDecoderOutput(ModelOutput): audio_values: Optional[torch.FloatTensor] = None -class XcodecVectorQuantize(nn.Module): - """ - Implementation of VQ similar to Karpathy's repo (https://github.com/karpathy/deep-vector-quantization) - - Additionally uses following tricks from improved VQGAN - (https://huggingface.co/papers/2110.04627): - 1. Factorized codes: Perform nearest neighbor lookup in low-dimensional space - for improved codebook usage - 2. l2-normalized codes: Converts euclidean distance to cosine similarity which - improves training stability - """ - - def __init__(self, config: XcodecConfig): - super().__init__() - - self.codebook_dim = config.codebook_dim - self.in_proj = nn.Conv1d(config.hidden_size, config.codebook_dim, kernel_size=1) - self.out_proj = nn.Conv1d(config.codebook_dim, config.hidden_size, kernel_size=1) - self.codebook = nn.Embedding(config.codebook_size, config.codebook_dim) - - def forward(self, hidden_state): - """ - Quantizes the input tensor using a fixed codebook and returns the corresponding codebook vectors. - - Args: - hidden_state (`torch.FloatTensor` of shape `(batch_size, dimension, time_steps)`): - Input tensor. - - Returns: - quantized_representation (`torch.Tensor`of shape `(batch_size, dimension, time_steps)`): - Quantized continuous representation of input. - commitment_loss (`torch.FloatTensor`of shape `(1)`): - Commitment loss to train encoder to predict vectors closer to codebook entries. - codebook_loss (`torch.FloatTensor`of shape `(1)`): - Codebook loss to update the codebook. - audio_codes (`torch.LongTensor` of shape `(batch_size, time_steps)`): - Codebook indices for each codebook, quantized discrete representation of input. - projected_latents (torch.FloatTensor of shape `(batch_size, num_codebooks * dimension, time_steps)`): - Projected latents (continuous representation of input before quantization). - """ - - projected_latents = self.in_proj(hidden_state) - quantized_representation, audio_codes = self.decode_latents(projected_latents) - - commitment_loss = F.mse_loss(projected_latents, quantized_representation.detach(), reduction="mean") - codebook_loss = F.mse_loss(quantized_representation, projected_latents.detach(), reduction="mean") - # noop in forward pass, straight-through gradient estimator in backward pass - quantized_representation = projected_latents + (quantized_representation - projected_latents).detach() - quantized_representation = self.out_proj(quantized_representation) - - return quantized_representation, commitment_loss, codebook_loss, audio_codes, projected_latents - - def decode_latents(self, hidden_states): - batch_size, hidden_dim, sequence_length = hidden_states.shape - encodings = hidden_states.permute(0, 2, 1).reshape(batch_size * sequence_length, hidden_dim) - codebook = self.codebook.weight # codebook: (N x D) - - # L2 normalize encodings and codebook (ViT-VQGAN) - encodings = F.normalize(encodings) - codebook = F.normalize(codebook) - - # Compute euclidean distance with codebook - l2_norm = encodings.pow(2).sum(1, keepdim=True) - dist = -(l2_norm - 2 * encodings @ codebook.t()) + codebook.pow(2).sum(1, keepdim=True).t() - - indices = dist.max(1)[1] - indices = indices.reshape(hidden_states.size(0), -1) - quantized_representation = self.codebook(indices).transpose(1, 2) - return quantized_representation, indices - - -class XcodecResidualVectorQuantize(nn.Module): - """ - ResidualVectorQuantize block - Introduced in SoundStream: An end2end neural audio codec (https://huggingface.co/papers/2107.03312) - """ - - def __init__(self, config: XcodecConfig): - super().__init__() - - n_codebooks = config.n_codebooks - quantizer_dropout = config.quantizer_dropout - - self.n_codebooks = n_codebooks - - self.quantizers = nn.ModuleList([XcodecVectorQuantize(config) for i in range(config.n_codebooks)]) - self.quantizer_dropout = quantizer_dropout - - def forward(self, hidden_state, n_quantizers: Optional[int] = None): - """ - Quantizes the input tensor using a fixed set of codebooks and returns corresponding codebook vectors. - Args: - hidden_state (`torch.Tensor` of shape `(batch_size, dimension, time_steps)`): - Input tensor to be quantized. - n_quantizers (`int`, *optional*): - Number of quantizers to use. If specified and `self.quantizer_dropout` is True, - this argument is ignored during training, and a random number of quantizers is used. - - Returns: - quantized_representation (`torch.Tensor` of shape `(batch_size, dimension, time_steps)`): - Quantized continuous representation of input. - audio_codes (`torch.Tensor` of shape `(batch_size, num_codebooks, time_steps)`): - Codebook indices for each codebook (quantized discrete representation of input). - projected_latents (`torch.Tensor` of shape `(batch_size, num_codebooks * dimension, time_steps)`): - Projected latents (continuous representation of input before quantization). - commitment_loss (`torch.Tensor` of shape `(1)`): - Commitment loss to train the encoder to predict vectors closer to codebook entries. - codebook_loss (`torch.Tensor` of shape `(1)`): - Codebook loss to update the codebook. - """ - - quantized_representation = 0 - residual = hidden_state - commitment_loss = 0 - codebook_loss = 0 - - audio_codes = [] - projected_latents = [] - - n_quantizers = n_quantizers if n_quantizers is not None else self.n_codebooks - if self.training: - n_quantizers = torch.ones((hidden_state.shape[0],)) * self.n_codebooks + 1 - dropout = torch.randint(1, self.n_codebooks + 1, (hidden_state.shape[0],)) - n_dropout = int(hidden_state.shape[0] * self.quantizer_dropout) - n_quantizers[:n_dropout] = dropout[:n_dropout] - n_quantizers = n_quantizers.to(hidden_state.device) - - for i, quantizer in enumerate(self.quantizers): - if self.training is False and i >= n_quantizers: - break - - quantized_representation_i, commitment_loss_i, codebook_loss_i, indices_i, projected_latents_i = quantizer( - residual - ) - - # Create mask to apply quantizer dropout - mask = torch.full((hidden_state.shape[0],), fill_value=i, device=hidden_state.device) < n_quantizers - quantized_representation = quantized_representation + quantized_representation_i * mask[:, None, None] - residual = residual - quantized_representation_i - - # Sum losses - commitment_loss += commitment_loss_i * mask - codebook_loss += codebook_loss_i * mask - - audio_codes.append(indices_i) - projected_latents.append(projected_latents_i) - - audio_codes = torch.stack(audio_codes, dim=1) - projected_latents = torch.cat(projected_latents, dim=1) - - return quantized_representation, audio_codes, projected_latents, commitment_loss, codebook_loss - - def from_codes(self, audio_codes: torch.Tensor): - """ - Reconstructs the continuous representation from quantized codes. - - Args: - audio_codes (`torch.Tensor` of shape `(batch_size, num_codebooks, time_steps)`): - Quantized discrete representation of input. - - Returns: - quantized_representation (`torch.Tensor`): - Quantized continuous representation of input. - projected_latents (`torch.Tensor`): - List of projected latents (continuous representations of input before quantization) - for each codebook. - audio_codes (`torch.Tensor`): - Codebook indices for each codebook. - """ - quantized_representation = 0.0 - projected_latents = [] - n_codebooks = audio_codes.shape[1] - for i in range(n_codebooks): - projected_latents_i = self.quantizers[i].codebook(audio_codes[:, i, :]).transpose(1, 2) - projected_latents.append(projected_latents_i) - quantized_representation += self.quantizers[i].out_proj(projected_latents_i) - return quantized_representation, torch.cat(projected_latents, dim=1), audio_codes - - def from_latents(self, latents: torch.Tensor): - """Reconstructs the quantized representation from unquantized latents. - - Args: - latents (`torch.Tensor` of shape `(batch_size, total_latent_dimension, time_steps)`): - Continuous representation of input after projection. - - Returns: - quantized_representation (`torch.Tensor` of shape `(batch_size, dimension, time_steps)`): - Quantized representation of the full-projected space. - quantized_latents (`torch.Tensor` of shape `(batch_size, dimension, time_steps)`): - Quantized representation of the latent space (continuous representation before quantization). - """ - quantized_representation = 0 - quantized_latents = [] - codes = [] - codebook_dims_tensor = torch.tensor([0] + [q.codebook_dim for q in self.quantizers]) - dims = torch.cumsum(codebook_dims_tensor, dim=0) - - n_codebooks = np.where(dims <= latents.shape[1])[0].max(axis=0, keepdims=True)[0] - for i in range(n_codebooks): - hidden_dim_j, hidden_dim_k = dims[i], dims[i + 1] - quantized_latents_i, codes_i = self.quantizers[i].decode_latents(latents[:, hidden_dim_j:hidden_dim_k, :]) - quantized_latents.append(quantized_latents_i) - codes.append(codes_i) - - quantized_representation_i = self.quantizers[i].out_proj(quantized_latents_i) - quantized_representation = quantized_representation + quantized_representation_i - - return quantized_representation, torch.cat(quantized_latents, dim=1) - - @auto_docstring class XcodecPreTrainedModel(PreTrainedAudioTokenizerBase): config: XcodecConfig @@ -647,7 +505,7 @@ class XcodecPreTrainedModel(PreTrainedAudioTokenizerBase): def _init_weights(self, module): if isinstance(module, nn.Conv1d): nn.init.trunc_normal_(module.weight, std=0.02) - nn.init.constant_(module.bias, 0) + # nn.init.constant_(module.bias, 0) elif isinstance(module, Snake1d): module.alpha.data.fill_(1.0) elif isinstance(module, nn.ConvTranspose1d): @@ -731,13 +589,20 @@ def __init__(self, config: XcodecConfig): self.config = config self.encoder = XcodecEncoder(config) - self.decoder = XcodecDecoder(config) + self.decoder = XcodecAcousticDecoder(config) - self.quantizer = XcodecResidualVectorQuantize(config) + self.quantizer = XcodecResidualVectorQuantizer(config) self.bits_per_codebook = int(math.log2(self.config.codebook_size)) if 2**self.bits_per_codebook != self.config.codebook_size: raise ValueError("The codebook_size must be a power of 2.") + self.semantic_decoder = XcodecSemanticDecoder(config) + self.semantic_proj = nn.Linear( + config.semantic_hidden_size + config.acoustic_hidden_size, config.semantic_hidden_size + ) + self.acoustic_proj = nn.Linear( + config.semantic_hidden_size + config.acoustic_hidden_size, config.acoustic_hidden_size + ) # Initialize weights and apply final processing self.post_init() @@ -748,6 +613,7 @@ def encode( self, input_values: torch.Tensor, n_quantizers: Optional[int] = None, + return_dict: Optional[bool] = None, ): r""" input_values (`torch.Tensor of shape `(batch_size, 1, time_steps)`): @@ -755,6 +621,8 @@ def encode( n_quantizers (int, *optional*): Number of quantizers to use. If None, all quantizers are used. Default is None. """ + return_dict = return_dict if return_dict is not None else self.config.return_dict + quantized_representation = self.encoder(input_values) quantized_representation, audio_codes, projected_latents, commitment_loss, codebook_loss = self.quantizer( quantized_representation, n_quantizers @@ -762,6 +630,9 @@ def encode( loss = self.config.commitment_loss_weight * commitment_loss + self.config.codebook_loss_weight * codebook_loss + if not return_dict: + return (loss, quantized_representation, audio_codes, projected_latents) + return XcodecEncoderOutput(loss, quantized_representation, audio_codes, projected_latents) @can_return_tuple @@ -770,6 +641,7 @@ def decode( self, quantized_representation: Optional[torch.Tensor] = None, audio_codes: Optional[torch.Tensor] = None, + return_dict: Optional[bool] = None, ): r""" quantized_representation (torch.Tensor of shape `(batch_size, dimension, time_steps)`, *optional*): @@ -778,16 +650,23 @@ 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 [`XcodecDecoderOutput`] instead of a plain tuple. """ if quantized_representation is None and audio_codes is None: raise ValueError("Either `quantized_representation` or `audio_codes` must be provided.") + return_dict = return_dict if return_dict is not None else self.config.return_dict + if audio_codes is not None: quantized_representation = self.quantizer.from_codes(audio_codes)[0] audio_values = self.decoder(quantized_representation).squeeze(1) + if not return_dict: + return (audio_values,) + return XcodecDecoderOutput(audio_values) @auto_docstring @@ -795,6 +674,7 @@ def forward( self, input_values: torch.Tensor, n_quantizers: Optional[int] = None, + return_dict: Optional[bool] = None, ): r""" input_values (`torch.Tensor` of shape `(batch_size, 1, time_steps)`): @@ -823,6 +703,8 @@ def forward( >>> # or the equivalent with a forward pass >>> audio_values = model(inputs["input_values"]).audio_values ```""" + + 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( @@ -830,4 +712,10 @@ def forward( ) audio_values = self.decode(quantized_representation, return_dict=False)[0][..., :length] + if not return_dict: + return (loss, audio_values, quantized_representation, audio_codes, projected_latents) + return XcodecOutput(loss, audio_values, quantized_representation, audio_codes, projected_latents) + + +__all__ = ["XcodecModel"] diff --git a/src/transformers/models/xcodec/modular_xcodec.py b/src/transformers/models/xcodec/modular_xcodec.py index ccaa0dede3a1..38fcd6680180 100644 --- a/src/transformers/models/xcodec/modular_xcodec.py +++ b/src/transformers/models/xcodec/modular_xcodec.py @@ -26,16 +26,15 @@ from .configuration_xcodec import XcodecConfig -class XcodecAcousticEncoder(DacEncoder):... +class XcodecAcousticEncoder(DacEncoder): ... -class XcodecSemanticEncoderLayer(nn.Module): +class XcodecSemanticEncoderResidualLayer(nn.Module): def __init__(self, config): super().__init__() self.act_fn = nn.ELU() - self.conv1 = nn.Conv1d(config.semantic_hidden_size, config.semantic_hidden_size, kernel_size=3, bias=False) + self.conv1 = nn.Conv1d(config.semantic_hidden_size, config.semantic_hidden_size, kernel_size=3, padding=1, bias=False) self.conv2 = nn.Conv1d(config.semantic_hidden_size, config.semantic_hidden_size, kernel_size=1, bias=False) - self.conv3 = nn.Conv1d(config.semantic_hidden_size, config.semantic_hidden_size, kernel_size=3, bias=False) def forward(self, hidden_state): residual = hidden_state @@ -45,6 +44,21 @@ def forward(self, hidden_state): return hidden_state +class XcodecSemanticEncoderLayer(nn.Module): + def __init__(self, config): + super().__init__() + self.residual_layers = nn.ModuleList( + [XcodecSemanticEncoderResidualLayer(config) for _ in range(config.num_residual_layers)] + ) + self.conv = nn.Conv1d(config.semantic_hidden_size, config.semantic_hidden_size, kernel_size=3, padding=1) + + def forward(self, hidden_state): + for layer in self.residual_layers: + hidden_state = layer(hidden_state) + hidden_state = self.conv(hidden_state) + return hidden_state + + class XcodecSemanticEncoder(nn.Module): def __init__(self, config): super().__init__() @@ -66,6 +80,10 @@ def __init__(self, config: XcodecConfig): self.acoustic_encoder = XcodecAcousticEncoder(config) self.semantic_encoder = XcodecSemanticEncoder(config) self.semantic_model = AutoModel.from_config(config.semantic_config) + self.linear = nn.Linear( + config.semantic_hidden_size + config.acoustic_hidden_size, + config.semantic_hidden_size + config.acoustic_hidden_size + ) self.sample_rate = config.sample_rate self.semantic_sample_rate = config.semantic_sample_rate @@ -94,7 +112,7 @@ def forward(self, input_values: torch.Tensor): return input_embeds -class XcodecDecoderBlock(DacDecoderBlock): +class XcodecAcousticDecoderBlock(DacDecoderBlock): def __init__(self, config, stride: int = 1, stride_index: int = 1): super().__init__(config) input_dim = config.decoder_hidden_size // 2**stride_index @@ -109,7 +127,7 @@ def __init__(self, config, stride: int = 1, stride_index: int = 1): ) -class XcodecDecoder(DacDecoder): +class XcodecAcousticDecoder(DacDecoder): def __init__(self, config: XcodecConfig): super().__init__(config) input_channel = config.acoustic_hidden_size @@ -126,8 +144,33 @@ def forward(self, hidden_states): return hidden_states -class XcodecResidualVectorQuantization(EncodecResidualVectorQuantizer): ... +class XcodecSemanticDecoder(nn.Module): + def __init__(self, config): + super().__init__() + self.input_conv = nn.Conv1d(config.semantic_hidden_size, config.semantic_hidden_size, kernel_size=3, padding=1, bias=False) + self.layers = nn.ModuleList( + [XcodecSemanticEncoderLayer(config) for _ in range(config.num_layers)] + ) + self.output_conv = nn.Conv1d(config.semantic_hidden_size, config.semantic_hidden_size, kernel_size=3, padding=1, bias=False) + + +class XcodecResidualVectorQuantizer(EncodecResidualVectorQuantizer): ... # TODO: @eustlb, consider adding bandwidth handling in dac in order to standardize codec usage -class XcodecModel(DacModel): ... +class XcodecModel(DacModel): + def __init__(self, config: XcodecConfig): + super().__init__(config) + self.decoder = XcodecAcousticDecoder(config) + self.semantic_decoder = XcodecSemanticDecoder(config) + self.semantic_proj = nn.Linear( + config.semantic_hidden_size + config.acoustic_hidden_size, + config.semantic_hidden_size + ) + self.acoustic_proj = nn.Linear( + config.semantic_hidden_size + config.acoustic_hidden_size, + config.acoustic_hidden_size + ) + + +__all__ = ["XcodecModel"]