# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 # This file was automatically generated from src/transformers/models/neucodec/modular_neucodec.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_neucodec.py file directly. One of our CI enforces this. # 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 # Copyright 2026 The HuggingFace Inc. team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import math from collections.abc import Callable from dataclasses import dataclass import numpy as np import torch.nn as nn from torch.nn import Parameter from ... import initialization as init from ...activations import ACT2FN from ...cache_utils import Cache from ...integrations import use_kernel_forward_from_hub, use_kernelized_func from ...masking_utils import create_bidirectional_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel from ...processing_utils import Unpack from ...utils import ModelOutput, TransformersKwargs, auto_docstring, can_return_tuple, is_torch_available from ...utils.deprecation import deprecate_kwarg from ...utils.generic import maybe_autocast from ..auto import AutoModel from .configuration_neucodec import NeuCodecConfig if is_torch_available(): import torch import torch.nn.functional as F @auto_docstring @dataclass class NeuCodecOutput(ModelOutput): r""" audio_values (`torch.FloatTensor` of shape `(batch_size, 1, sequence_length)`, *optional*): Decoded audio waveform values in the time domain, obtained using the decoder part of NeuCodec. These represent the reconstructed audio signal. audio_codes (`torch.LongTensor` of shape `(batch_size, 1, codes_length)`, *optional*): Discrete code embeddings computed using `model.encode`. These are the quantized representations of the input audio used for further processing or generation. latents (`torch.Tensor` of shape `(batch_size, dimension, time_steps)`): Quantized continuous representation of input's embedding. audio_codes_mask (`torch.int32` of shape `(batch_size, 1, codes_length)`, *optional*): Downsampled `padding_mask` for indicating valid audio codes in `audio_codes`. """ audio_values: torch.FloatTensor | None = None audio_codes: torch.LongTensor | None = None latents: torch.Tensor | None = None audio_codes_mask: torch.Tensor | None = None @auto_docstring @dataclass class NeuCodecEncoderOutput(ModelOutput): r""" audio_codes (`torch.LongTensor` of shape `(batch_size, 1, codes_length)`, *optional*): Discrete code embeddings computed using `model.encode`. These represent the compressed, quantized form of the input audio signal that can be used for storage, transmission, or generation. latents (`torch.Tensor` of shape `(batch_size, dimension, time_steps)`): Quantized continuous representation of input's embedding. audio_codes_mask (`torch.int32` of shape `(batch_size, 1, codes_length)`, *optional*): Downsampled `padding_mask` for indicating valid audio codes in `audio_codes`. """ audio_codes: torch.LongTensor | None = None latents: torch.Tensor | None = None audio_codes_mask: torch.Tensor | None = None @auto_docstring @dataclass class NeuCodecDecoderOutput(ModelOutput): r""" audio_values (`torch.FloatTensor` of shape `(batch_size, 1, segment_length)`, *optional*): Decoded audio waveform values in the time domain, obtained by converting the discrete codes back into continuous audio signals. This represents the reconstructed audio that can be played back. """ audio_values: torch.FloatTensor | None = None class NeuCodecMLP(nn.Module): def __init__(self, config: NeuCodecConfig): super().__init__() self.config = config self.activation_fn = ACT2FN[config.hidden_act] self.fc1 = nn.Linear(config.hidden_size, config.intermediate_size, bias=False) self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size, bias=False) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: hidden_states = self.fc1(hidden_states) hidden_states = self.activation_fn(hidden_states) hidden_states = self.fc2(hidden_states) return hidden_states def rotate_half(x): """Rotates half the hidden dims of the input.""" x1 = x[..., : x.shape[-1] // 2] x2 = x[..., x.shape[-1] // 2 :] return torch.cat((-x2, x1), dim=-1) @use_kernel_forward_from_hub("rotary_pos_emb") def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): """Applies Rotary Position Embedding to the query and key tensors. Args: q (`torch.Tensor`): The query tensor. k (`torch.Tensor`): The key tensor. cos (`torch.Tensor`): The cosine part of the rotary embedding. sin (`torch.Tensor`): The sine part of the rotary embedding. unsqueeze_dim (`int`, *optional*, defaults to 1): The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2. Returns: `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding. """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: """ This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) """ batch, num_key_value_heads, slen, head_dim = hidden_states.shape if n_rep == 1: return hidden_states hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim) return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) def eager_attention_forward( module: nn.Module, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, attention_mask: torch.Tensor | None, scaling: float, dropout: float = 0.0, **kwargs: Unpack[TransformersKwargs], ): key_states = repeat_kv(key, module.num_key_value_groups) value_states = repeat_kv(value, module.num_key_value_groups) attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling if attention_mask is not None: attn_weights = attn_weights + attention_mask attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype) attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training) attn_output = torch.matmul(attn_weights, value_states) attn_output = attn_output.transpose(1, 2).contiguous() return attn_output, attn_weights @use_kernelized_func(apply_rotary_pos_emb) class NeuCodecAttention(nn.Module): """Multi-headed attention from 'Attention Is All You Need' paper""" def __init__(self, config: NeuCodecConfig, layer_idx: int): super().__init__() self.config = config self.layer_idx = layer_idx self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads) self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads self.scaling = self.head_dim**-0.5 self.attention_dropout = config.attention_dropout self.is_causal = False self.q_proj = nn.Linear( config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias ) self.k_proj = nn.Linear( config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias ) self.v_proj = nn.Linear( config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias ) self.o_proj = nn.Linear( config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias ) def forward( self, hidden_states: torch.Tensor, position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, attention_mask: torch.Tensor | None = None, past_key_values: Cache | None = None, **kwargs: Unpack[TransformersKwargs], ) -> tuple[torch.Tensor, torch.Tensor]: input_shape = hidden_states.shape[:-1] hidden_shape = (*input_shape, -1, self.head_dim) query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2) key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2) value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2) cos, sin = position_embeddings # NeuCodec uses position_ids of shape (1, num_attention_heads) so cos/sin have shape # (batch, num_attention_heads, head_dim). unsqueeze_dim=2 broadcasts correctly against # q/k of shape (batch, num_heads, seq_len, head_dim), unlike Llama's default of 1. query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, unsqueeze_dim=2) if past_key_values is not None: key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx) attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface( self.config._attn_implementation, eager_attention_forward ) attn_output, attn_weights = attention_interface( self, query_states, key_states, value_states, attention_mask, dropout=0.0 if not self.training else self.attention_dropout, scaling=self.scaling, **kwargs, ) attn_output = attn_output.reshape(*input_shape, -1).contiguous() attn_output = self.o_proj(attn_output) return attn_output, attn_weights @use_kernel_forward_from_hub("RMSNorm") class NeuCodecRMSNorm(nn.Module): def __init__(self, hidden_size, eps: float = 1e-6) -> None: """ NeuCodecRMSNorm is equivalent to T5LayerNorm """ super().__init__() self.weight = nn.Parameter(torch.ones(hidden_size)) self.variance_epsilon = eps def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: input_dtype = hidden_states.dtype hidden_states = hidden_states.to(torch.float32) variance = hidden_states.pow(2).mean(-1, keepdim=True) hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) return self.weight * hidden_states.to(input_dtype) def extra_repr(self): return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}" class NeuCodecDecoderLayer(GradientCheckpointingLayer): def __init__(self, config: NeuCodecConfig, layer_idx: int): super().__init__() self.hidden_size = config.hidden_size self.self_attn = NeuCodecAttention(config=config, layer_idx=layer_idx) self.mlp = NeuCodecMLP(config) self.input_layernorm = NeuCodecRMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.post_attention_layernorm = NeuCodecRMSNorm(config.hidden_size, eps=config.rms_norm_eps) def forward( self, hidden_states: torch.Tensor, attention_mask: torch.Tensor | None = None, position_ids: torch.LongTensor | None = None, past_key_values: Cache | None = None, use_cache: bool | None = False, position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, **kwargs: Unpack[TransformersKwargs], ) -> torch.Tensor: residual = hidden_states hidden_states = self.input_layernorm(hidden_states) # Self Attention hidden_states, _ = self.self_attn( hidden_states=hidden_states, attention_mask=attention_mask, position_ids=position_ids, past_key_values=past_key_values, use_cache=use_cache, position_embeddings=position_embeddings, **kwargs, ) hidden_states = residual + hidden_states # Fully Connected residual = hidden_states hidden_states = self.post_attention_layernorm(hidden_states) hidden_states = self.mlp(hidden_states) hidden_states = residual + hidden_states return hidden_states class NeuCodecSnakeBeta(nn.Module): """ A modified Snake function which uses separate parameters for the magnitude of the periodic components Shape: - Input: (B, C, T) - Output: (B, C, T), same shape as the input Parameters: - alpha - trainable parameter that controls frequency - beta - trainable parameter that controls magnitude References: - This activation function is a modified version based on this paper by Liu Ziyin, Tilman Hartwig, Masahito Ueda: https://huggingface.co/papers/2006.08195 """ def __init__(self, in_features, alpha=1.0): super().__init__() self.in_features = in_features # initialize alpha self.alpha = Parameter(torch.zeros(in_features) * alpha) self.beta = Parameter(torch.zeros(in_features) * alpha) self.no_div_by_zero = 0.000000001 def forward(self, hidden_states): """ Forward pass of the function. Applies the function to the input elementwise. SnakeBeta ∶= x + 1/b * sin^2 (xa) """ alpha = self.alpha.unsqueeze(0).unsqueeze(-1) # line up with x to [B, C, T] beta = self.beta.unsqueeze(0).unsqueeze(-1) alpha = torch.exp(alpha) beta = torch.exp(beta) hidden_states = hidden_states + (1.0 / (beta + self.no_div_by_zero)) * torch.pow( torch.sin(hidden_states * alpha), 2 ) return hidden_states def kaiser_sinc_filter1d(cutoff, half_width, kernel_size): """Generates a 1D Kaiser-windowed sinc filter. Args: cutoff (float): Normalized cutoff frequency (0 to 0.5). half_width (float): Transition bandwidth. kernel_size (int): Number of filter taps. Returns: torch.Tensor: A tensor of shape (1, 1, kernel_size) representing the filter. """ is_even = kernel_size % 2 == 0 half_size = kernel_size // 2 # Compute Kaiser window parameters delta_f = 4 * half_width attenuation = 2.285 * (half_size - 1) * math.pi * delta_f + 7.95 if attenuation > 50.0: beta = 0.1102 * (attenuation - 8.7) elif attenuation >= 21.0: beta = 0.5842 * (attenuation - 21) ** 0.4 + 0.07886 * (attenuation - 21.0) else: beta = 0.0 kaiser_window = torch.kaiser_window(kernel_size, beta=beta, periodic=False, dtype=torch.float32) # Compute time indices if is_even: time_indices = torch.arange(-half_size, half_size) + 0.5 else: time_indices = torch.arange(kernel_size) - half_size # Compute sinc filter if cutoff == 0: return torch.zeros((1, 1, kernel_size), dtype=torch.float32) # Ensures correct shape sinc_filter = torch.sinc(2 * cutoff * time_indices) normalized_filter = 2 * cutoff * kaiser_window * sinc_filter # Normalize to ensure sum = 1 (avoid leakage of constant component) normalized_filter /= normalized_filter.sum() return normalized_filter.view(1, 1, kernel_size) class NeuCodecDownSample1d(nn.Module): def __init__(self, ratio=2, kernel_size=None): super().__init__() cutoff = 0.5 / ratio half_width = 0.6 / ratio self.cutoff = cutoff self.half_width = half_width self.kernel_size = kernel_size if cutoff < 0.0: raise ValueError("Minimum cutoff must be larger than zero.") if cutoff > 0.5: raise ValueError("A cutoff above 0.5 does not make sense.") self.even = kernel_size % 2 == 0 self.pad_left = kernel_size // 2 - int(self.even) self.pad_right = kernel_size // 2 self.stride = ratio filter = kaiser_sinc_filter1d(cutoff, half_width, kernel_size) self.filter = nn.Buffer(filter, persistent=False) def forward(self, hidden_states): channels = hidden_states.shape[1] hidden_states = F.pad(hidden_states, (self.pad_left, self.pad_right), mode="replicate") out = F.conv1d( hidden_states, # add casting to avoid dtype mismatch for SDPA self.filter.to(hidden_states.dtype).expand(channels, -1, -1), stride=self.stride, groups=channels, ) return out class NeuCodecUpSample1d(nn.Module): def __init__(self, ratio=2, kernel_size=None): super().__init__() self.ratio = ratio self.kernel_size = int(6 * ratio // 2) * 2 if kernel_size is None else kernel_size self.stride = ratio self.pad = self.kernel_size // ratio - 1 self.pad_left = self.pad * self.stride + (self.kernel_size - self.stride) // 2 self.pad_right = self.pad * self.stride + (self.kernel_size - self.stride + 1) // 2 filter = kaiser_sinc_filter1d(cutoff=0.5 / ratio, half_width=0.6 / ratio, kernel_size=self.kernel_size) self.filter = nn.Buffer(filter, persistent=False) def forward(self, hidden_states): channels = hidden_states.shape[1] hidden_states = F.pad(hidden_states, (self.pad, self.pad), mode="replicate") hidden_states = self.ratio * F.conv_transpose1d( hidden_states, # add casting to avoid dtype mismatch for SDPA self.filter.to(hidden_states.dtype).expand(channels, -1, -1), stride=self.stride, groups=channels, ) hidden_states = hidden_states[..., self.pad_left : -self.pad_right] return hidden_states class NeuCodecFiniteScalarQuantization(nn.Module): """ Finite Scalar Quantization (FSQ) module that quantizes continuous latent representations into discrete codes. Original code: https://github.com/lucidrains/vector-quantize-pytorch/blob/353d46027888dfb140c3c65a67a7356f1492d71d/vector_quantize_pytorch/finite_scalar_quantization.py#L64 Original modeling uses `ResidualFSQ` with a single quantizer: https://huggingface.co/HKUSTAudio/neucodec/blob/main/vq/codec_decoder_vocos.py#L389 But we can directly use FSQ since a main feature of NeuCodec is that it uses a single codebook. """ def __init__(self, config: NeuCodecConfig): super().__init__() self.quantization_levels = list(config.quantization_levels) levels, basis, codebook = self._compute_buffers() self.levels = nn.Buffer(levels, persistent=False) self.basis = nn.Buffer(basis, persistent=False) self.codebook = nn.Buffer(codebook, persistent=False) def _compute_buffers(self, device=None): """Compute the levels, basis, and codebook buffers for the FSQ quantizer.""" levels = torch.tensor(self.quantization_levels, dtype=torch.int32, device=device) basis = torch.cumprod( torch.tensor([1] + self.quantization_levels[:-1], device=device), dim=0, dtype=torch.int32 ) indices = torch.arange(int(np.prod(self.quantization_levels)), device=device).unsqueeze(-1) level_indices = (indices // basis) % levels half_width = levels // 2 codebook = (level_indices - half_width) / half_width return levels, basis, codebook def _indices_to_codes(self, indices: torch.Tensor) -> torch.Tensor: """ Convert integer codebook indices to normalized per-dimension codes in [-1, 1]. """ indices = indices.unsqueeze(-1) level_indices = (indices // self.basis) % self.levels half_width = self.levels // 2 codes = (level_indices - half_width) / half_width return codes def bound(self, hidden_states: torch.Tensor, eps: float = 1e-3) -> torch.Tensor: """ Constrain `hidden_states` to the valid quantization range for each dimension. Uses a scaled tanh to soft-clip values into the interval $[-(L-1)/2, (L-1)/2]$ (offset by 0.5 for even-level dimensions), where $L$ is the number of quantization levels. The small `eps` margin prevents values from saturating exactly at the boundary, which would zero out gradients. Args: hidden_states (`torch.Tensor`): Continuous input to be bounded. eps (`float`, *optional*, defaults to `1e-3`): Small margin added to the level range to avoid gradient saturation at boundaries. Returns: `torch.Tensor`: Bounded values in the valid quantization range. """ half_range = (self.levels - 1) * (1 + eps) / 2 offset = torch.where(self.levels % 2 == 0, 0.5, 0.0) shift = (offset / half_range).atanh() return (hidden_states + shift).tanh() * half_range - offset def forward(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: # NOTE: could rewrite to pass tensor to a decorator such that device type is handled internally original_dtype = hidden_states.dtype device_type = ( hidden_states.device.type if isinstance(hidden_states.device.type, str) and hidden_states.device.type != "mps" else "cpu" ) with maybe_autocast(device_type=device_type, enabled=False): # Force float32 hidden_states = hidden_states.float() half_width = self.levels // 2 # Quantize: bound and round with straight-through gradient hidden_states = self.bound(hidden_states) rounded = hidden_states.round() codes = hidden_states + (rounded - hidden_states).detach() codes = codes / half_width # Code to indices code_scaled = (codes * half_width) + half_width indices = (code_scaled * self.basis).sum(dim=-1).to(torch.int32) return codes.to(original_dtype), indices class NeuCodecISTFTHead(nn.Module): """ Head for converting decoder outputs to waveform via STFT projection and ISTFT. Uses custom "same" padding ISTFT from Vocos: https://github.com/gemelo-ai/vocos/blob/c859e3b7b534f3776a357983029d34170ddd6fc3/vocos/spectral_ops.py#L47 """ def __init__(self, config: NeuCodecConfig): super().__init__() self.linear = nn.Linear(config.hidden_size, config.n_fft + 2) self.n_fft = config.n_fft self.hop_length = config.hop_length self.padding = (self.n_fft - self.hop_length) // 2 window = torch.hann_window(config.n_fft) self.window = nn.Buffer(window, persistent=False) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: stft_pred = self.linear(hidden_states) # Keep the frequency axis last for `torch.polar` compatibility with torch.compile magnitude, phase = stft_pred.chunk(2, dim=-1) # Cast to float32: complex exponential and irfft are not supported for fp16 (ComplexHalf) magnitude = magnitude.float() phase = phase.float() # Clamp like original: https://huggingface.co/HKUSTAudio/neucodec/blob/main/vq/codec_decoder_vocos.py#L138 magnitude = torch.exp(magnitude).clamp(max=1e2) spectrogram_complex = torch.polar(magnitude, phase) # Back to audio (ISTFT with manual "same" padding: torch.istft lacks a native same-padding mode, # so we use irfft + fold with explicit pre-computed padding to replicate it) time_frames = torch.fft.irfft(spectrogram_complex, self.n_fft, dim=-1, norm="backward").transpose(1, 2) time_frames = time_frames * self.window[None, :, None] num_frames = spectrogram_complex.shape[1] output_size = (num_frames - 1) * self.hop_length + self.n_fft audio = F.fold( time_frames, output_size=(1, output_size), kernel_size=(1, self.n_fft), stride=(1, self.hop_length), )[:, 0, 0, self.padding : -self.padding] # Normalize window_envelope = F.fold( self.window.square().expand(1, num_frames, -1).transpose(1, 2), output_size=(1, output_size), kernel_size=(1, self.n_fft), stride=(1, self.hop_length), ).squeeze()[self.padding : -self.padding] # Clamp as expected by original: https://huggingface.co/HKUSTAudio/neucodec/blob/main/vq/codec_decoder_vocos.py#L82 window_envelope = window_envelope.clamp(min=1e-11) audio = audio / window_envelope return audio.unsqueeze(1) @auto_docstring class NeuCodecPreTrainedModel(PreTrainedModel): config: NeuCodecConfig base_model_prefix = "neucodec" input_modalities = ("audio",) supports_gradient_checkpointing = True _no_split_modules = None _skip_keys_device_placement = ["past_key_values"] _supports_flash_attn = True _supports_sdpa = True _supports_flex_attn = True _supports_cache_class = True _supports_attention_backend = True _can_compile_fullgraph = True main_input_name = "input_values" _can_record_outputs = { "hidden_states": NeuCodecDecoderLayer, "attentions": NeuCodecDecoderLayer, } def _init_weights(self, module): super()._init_weights(module) if isinstance(module, NeuCodecSnakeBeta): init.zeros_(module.alpha) init.zeros_(module.beta) elif isinstance(module, NeuCodecISTFTHead): window = torch.hann_window(module.n_fft) init.copy_(module.window, window) elif isinstance(module, NeuCodecFiniteScalarQuantization): levels, basis, codebook = module._compute_buffers(device=module.levels.device) init.copy_(module.levels, levels) init.copy_(module.basis, basis) init.copy_(module.codebook, codebook) elif isinstance(module, NeuCodecUpSample1d): filter_tensor = kaiser_sinc_filter1d(0.5 / module.ratio, 0.6 / module.ratio, module.kernel_size) init.copy_(module.filter, filter_tensor) elif isinstance(module, NeuCodecDownSample1d): filter_tensor = kaiser_sinc_filter1d(module.cutoff, module.half_width, module.kernel_size) init.copy_(module.filter, filter_tensor) class NeuCodecRotaryEmbedding(nn.Module): @deprecate_kwarg("device", version="5.18") def __init__(self, config: NeuCodecConfig, device=None): super().__init__() self.max_seq_len_cached = config.max_position_embeddings self.original_max_seq_len = config.max_position_embeddings self.config = config self.rope_type = self.config.rope_parameters["rope_type"] rope_init_fn: Callable = self.compute_default_rope_parameters if self.rope_type != "default": rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type] inv_freq, self.attention_scaling = rope_init_fn(self.config, device) self.inv_freq = nn.Buffer(inv_freq, persistent=False) self.original_inv_freq = nn.Buffer(inv_freq.clone(), persistent=False) @staticmethod @deprecate_kwarg("device", version="5.18") def compute_default_rope_parameters(config: NeuCodecConfig, device=None, **kwargs) -> tuple[torch.Tensor, float]: """ Computes the inverse frequencies according to the original RoPE implementation Args: config ([`~transformers.PreTrainedConfig`]): The model configuration. Returns: Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the post-processing scaling factor applied to the computed cos/sin (unused in this type of RoPE). """ base = config.rope_parameters["rope_theta"] dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads attention_factor = 1.0 # Unused in this type of RoPE # Compute the inverse frequencies inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float) / dim)) return inv_freq.to(device), attention_factor @torch.no_grad() @dynamic_rope_update # power user: used with advanced RoPE types (e.g. dynamic rope) def forward(self, x, position_ids): inv_freq_expanded = ( self.inv_freq[None, :, None].expand(position_ids.shape[0], -1, 1).to(dtype=torch.float, device=x.device) ) position_ids_expanded = position_ids[:, None, :].float() device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu" # Disable any outside autocast context if any, to really force fp32 with maybe_autocast(device_type=device_type, enabled=False): freqs = (inv_freq_expanded @ position_ids_expanded).transpose(1, 2) emb = torch.cat((freqs, freqs), dim=-1) cos = emb.cos() * self.attention_scaling sin = emb.sin() * self.attention_scaling return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) class NeuCodecAntiAliasedActivation1d(nn.Module): def __init__( self, activation, up_ratio: int = 2, down_ratio: int = 2, up_kernel_size: int = 12, down_kernel_size: int = 12, ): super().__init__() if not callable(activation): raise TypeError("Activation function must be callable") self.act = activation self.upsample = NeuCodecUpSample1d(up_ratio, up_kernel_size) self.downsample = NeuCodecDownSample1d(down_ratio, down_kernel_size) def forward(self, hidden_states): hidden_states = self.upsample(hidden_states) hidden_states = self.act(hidden_states) hidden_states = self.downsample(hidden_states) return hidden_states class NeuCodecResidualUnit(nn.Module): """ A residual unit composed of Snake1d and weight-normalized Conv1d layers with dilations. """ def __init__(self, dimension, dilation): super().__init__() pad = ((7 - 1) * dilation) // 2 self.snake1 = NeuCodecAntiAliasedActivation1d(activation=NeuCodecSnakeBeta(dimension)) self.conv1 = nn.Conv1d(dimension, dimension, kernel_size=7, dilation=dilation, padding=pad) self.snake2 = NeuCodecAntiAliasedActivation1d(activation=NeuCodecSnakeBeta(dimension)) self.conv2 = nn.Conv1d(dimension, dimension, kernel_size=1) def forward(self, hidden_state): """ Forward pass through the residual unit. 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 NeuCodecEncoderBlock(nn.Module): """Encoder block used in NEUCODEC encoder.""" def __init__(self, config: NeuCodecConfig, stride: int = 1, stride_index: int = 1): super().__init__() dimension = config.encoder_hidden_size * 2**stride_index self.res_unit1 = NeuCodecResidualUnit(dimension // 2, dilation=1) self.res_unit2 = NeuCodecResidualUnit(dimension // 2, dilation=3) self.res_unit3 = NeuCodecResidualUnit(dimension // 2, dilation=9) self.snake1 = NeuCodecAntiAliasedActivation1d(activation=NeuCodecSnakeBeta(dimension // 2)) self.conv1 = nn.Conv1d( dimension // 2, dimension, kernel_size=2 * stride, stride=stride, padding=math.ceil(stride / 2) ) 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 NeuCodecEncoder(nn.Module): """NEUCODEC Encoder""" def __init__(self, config: NeuCodecConfig): super().__init__() # 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(config.downsampling_ratios): stride_index = stride_index + 1 self.block += [NeuCodecEncoderBlock(config, stride=stride, stride_index=stride_index)] self.block = nn.ModuleList(self.block) d_model = config.encoder_hidden_size * 2 ** len(config.downsampling_ratios) self.snake1 = NeuCodecAntiAliasedActivation1d(activation=NeuCodecSnakeBeta(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) return hidden_state class NeuCodecResNetBlock(nn.Module): def __init__(self, config: NeuCodecConfig): super().__init__() self.norm1 = nn.GroupNorm(num_groups=32, num_channels=config.hidden_size, eps=1e-6, affine=True) self.activation1 = nn.SiLU() self.conv1 = nn.Conv1d(config.hidden_size, config.hidden_size, kernel_size=3, stride=1, padding=1) self.norm2 = nn.GroupNorm(num_groups=32, num_channels=config.hidden_size, eps=1e-6, affine=True) self.activation2 = nn.SiLU() self.activation_dropout = config.activation_dropout self.conv2 = nn.Conv1d(config.hidden_size, config.hidden_size, kernel_size=3, stride=1, padding=1) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: hidden_states = hidden_states.transpose(1, 2) residual = hidden_states hidden_states = self.norm1(hidden_states) hidden_states = self.activation1(hidden_states) hidden_states = self.conv1(hidden_states) hidden_states = self.norm2(hidden_states) hidden_states = self.activation2(hidden_states) hidden_states = nn.functional.dropout(hidden_states, p=self.activation_dropout, training=self.training) hidden_states = self.conv2(hidden_states) return (hidden_states + residual).transpose(1, 2) class NeuCodecQuantizer(nn.Module): def __init__(self, config: NeuCodecConfig): super().__init__() self.quantizer = NeuCodecFiniteScalarQuantization(config) self.project_in = nn.Linear(config.quantization_dim, len(config.quantization_levels)) self.project_out = nn.Linear(len(config.quantization_levels), config.quantization_dim) def from_codes(self, indices: torch.Tensor) -> torch.Tensor: indices = indices.squeeze(-1) # Remove channel dimension codes = self.quantizer.codebook[indices] return self.project_out(codes) def forward(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: hidden_states = self.project_in(hidden_states) original_dtype = hidden_states.dtype hidden_states = self.quantizer.bound(hidden_states) # For consistency with original checkpoint quantized_out, indices = self.quantizer(hidden_states) quantized_out = self.project_out(quantized_out.to(original_dtype)) indices = indices.unsqueeze(-1) # Add channel dimension for single codebook return quantized_out, indices class NeuCodecDecoder(nn.Module): """Vocos-based decoder with ResNet, Transformer, and ISTFT head for audio reconstruction.""" def __init__(self, config: NeuCodecConfig): super().__init__() self.fc = nn.Linear(config.hidden_size + config.semantic_model_config.hidden_size, config.hidden_size) self.embed = nn.Conv1d(config.hidden_size, config.hidden_size, kernel_size=7, padding=3) self.prior_net = nn.ModuleList([NeuCodecResNetBlock(config), NeuCodecResNetBlock(config)]) self.num_attention_heads = config.num_attention_heads self.rotary_emb = NeuCodecRotaryEmbedding(config=config) self.layers = nn.ModuleList( [NeuCodecDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] ) self.post_net = nn.ModuleList([NeuCodecResNetBlock(config), NeuCodecResNetBlock(config)]) self.norm = nn.LayerNorm(config.hidden_size, eps=1e-6) self.head = NeuCodecISTFTHead(config) def forward(self, hidden_states: torch.Tensor, **kwargs) -> torch.Tensor: hidden_states = self.fc(hidden_states) hidden_states = hidden_states.transpose(1, 2) hidden_states = self.embed(hidden_states) hidden_states = hidden_states.transpose(1, 2) # Conv ResNet for layer in self.prior_net: hidden_states = layer(hidden_states) # Transformer: (batch, time, hidden) # position_ids uses num_attention_heads so that RoPE produces cos/sin of shape (batch, num_heads, head_dim), # which broadcasts correctly against q/k of shape (batch, num_heads, seq_len, head_dim) via unsqueeze_dim=2 # in `apply_rotary_pos_emb`. NOTE: this is non-standard and could be unsafe under tensor parallelism # (TP shards see only a subset of heads), but TP is not used for this model in practice. position_ids = torch.arange(self.num_attention_heads, device=hidden_states.device).unsqueeze(0) position_embeddings = self.rotary_emb(hidden_states, position_ids) for layer in self.layers: hidden_states = layer(hidden_states, position_embeddings=position_embeddings, **kwargs) # Conv ResNet for layer in self.post_net: hidden_states = layer(hidden_states) return self.head(self.norm(hidden_states)) class NeuCodecSemanticAdapter(nn.Module): def __init__(self, config: NeuCodecConfig): super().__init__() self.conv1 = nn.Conv1d( in_channels=config.semantic_model_config.hidden_size, out_channels=config.semantic_model_config.hidden_size, kernel_size=3, padding=1, bias=False, ) self.act1 = nn.ReLU() self.conv2 = nn.Conv1d( config.semantic_model_config.hidden_size, config.semantic_model_config.hidden_size, kernel_size=3, padding=1, bias=True, ) self.act2 = nn.ReLU() self.conv3 = nn.Conv1d( config.semantic_model_config.hidden_size, config.semantic_model_config.hidden_size, kernel_size=3, padding=1, bias=True, ) self.conv4 = nn.Conv1d( in_channels=config.semantic_model_config.hidden_size, out_channels=config.semantic_model_config.hidden_size, kernel_size=3, padding=1, bias=False, ) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: hidden_states = self.conv1(hidden_states) hidden_states = self.act1(hidden_states) residual = hidden_states hidden_states = self.conv2(hidden_states) hidden_states = self.act2(hidden_states) hidden_states = self.conv3(hidden_states) hidden_states = hidden_states + residual hidden_states = self.conv4(hidden_states) return hidden_states @auto_docstring(custom_intro="NeuCodec neural audio codec model.") class NeuCodecModel(NeuCodecPreTrainedModel): config_class = NeuCodecConfig def __init__(self, config: NeuCodecConfig): super().__init__(config) # `Xcodec2Model.hop_length` mirrors `config.hop_length`, which for NeuCodec is expressed in the decoder's # (24kHz) domain. The mask arithmetic in `encode()` operates on `input_values` in the encoder's (16kHz) # domain, so it must use the un-rescaled hop length instead. self.hop_length = config.encoder_hop_length self.semantic_encoder = AutoModel.from_config(config.semantic_model_config) self.semantic_adapter = NeuCodecSemanticAdapter(config) self.acoustic_encoder = NeuCodecEncoder(config) self.fc_encoder = nn.Linear( config.hidden_size + config.semantic_model_config.hidden_size, config.hidden_size + config.semantic_model_config.hidden_size, ) self.quantizer = NeuCodecQuantizer(config) self.acoustic_decoder = NeuCodecDecoder(config) self.sample_rate_conversion_factor = config.output_sampling_rate / config.input_sampling_rate self.post_init() @auto_docstring @can_return_tuple def encode( self, input_values: torch.Tensor, input_features: torch.Tensor, padding_mask: torch.Tensor | None = None, input_features_mask: torch.Tensor | None = None, output_latents: bool = False, **kwargs: Unpack[TransformersKwargs], ) -> tuple | NeuCodecEncoderOutput: r""" input_values (`torch.Tensor` of shape `(batch_size, 1, sequence_length)`): Input audio waveform. input_features (`torch.Tensor` of shape `(batch_size, mel_bins, time_steps)`): Input audio mel spectrogram for semantic encoding. padding_mask (`torch.Tensor` of shape `(batch_size, 1, sequence_length)`): Padding mask used to pad `input_values`. input_features_mask (`torch.Tensor` of shape `(batch_size, time_steps)`, *optional*): Attention mask for the spectrogram input to the semantic encoder. `1` for valid frames, `0` for padding. output_latents (`bool`, *optional*, defaults to `False`): Whether to return the continuous latent representation from the quantizer. """ # Semantic embedding with torch.no_grad(): semantic_output = self.semantic_encoder(input_features, attention_mask=input_features_mask) semantic_hidden_states = semantic_output.last_hidden_state.transpose(1, 2) semantic_hidden_states = self.semantic_adapter(semantic_hidden_states) # Acoustic embedding acoustic_hidden_states = self.acoustic_encoder(input_values) # The two branches downsample independently and can differ by a frame; trim to the shorter one, matching # the reference: https://github.com/neuphonic/neucodec/blob/ed3e6cd1bdc374ce14a21355e5eee66a777149ce/neucodec/model.py#L173 min_length = min(acoustic_hidden_states.shape[-1], semantic_hidden_states.shape[-1]) acoustic_hidden_states = acoustic_hidden_states[..., :min_length] semantic_hidden_states = semantic_hidden_states[..., :min_length] hidden_states = torch.cat([semantic_hidden_states, acoustic_hidden_states], dim=1) hidden_states = self.fc_encoder(hidden_states.transpose(1, 2)) # Quantize latents, audio_codes = self.quantizer(hidden_states) latents = latents.transpose(1, 2) audio_codes = audio_codes.transpose(1, 2) # If provided, compute corresponding padding mask for audio codes audio_codes_mask = None if padding_mask is not None: audio_length = padding_mask.sum(dim=-1, keepdim=True) audio_length = audio_length * self.sample_rate_conversion_factor token_length = audio_length // self.hop_length idx = torch.arange(audio_codes.shape[-1], device=padding_mask.device).view(1, -1) audio_codes_mask = (idx < token_length).to(padding_mask.dtype) return NeuCodecEncoderOutput( audio_codes=audio_codes, latents=latents if output_latents else None, audio_codes_mask=audio_codes_mask, ) @auto_docstring @can_return_tuple def decode( self, audio_codes: torch.Tensor | None = None, latents: torch.Tensor | None = None, audio_codes_mask: torch.Tensor | None = None, **kwargs: Unpack[TransformersKwargs], ) -> tuple | NeuCodecDecoderOutput: r""" audio_codes (`torch.LongTensor` of shape `(batch_size, 1, codes_length)`): Discrete code indices computed using `model.encode`. latents (torch.Tensor of shape `(batch_size, dimension, time_steps)`, *optional*): Quantized continuous representation of input. audio_codes_mask (`torch.Tensor` of shape `(batch_size, codes_length)`, *optional*): Mask marking valid (non-padding) positions, as returned by `model.encode`. Needed for correct batched decoding of variable-length audio. """ if latents is None and audio_codes is None: raise ValueError("Either `latents` or `audio_codes` must be provided.") if audio_codes is not None: latents = self.quantizer.from_codes(audio_codes.transpose(1, 2)) else: latents = latents.transpose(1, 2) # Difference to xcodec2 is the additional mask attention_mask = create_bidirectional_mask( config=self.config, inputs_embeds=latents, attention_mask=audio_codes_mask, ) recon_audio = self.acoustic_decoder(latents, attention_mask=attention_mask, **kwargs) return NeuCodecDecoderOutput(audio_values=recon_audio) @auto_docstring @can_return_tuple def forward( self, input_values: torch.Tensor, input_features: torch.Tensor, padding_mask: torch.Tensor | None = None, input_features_mask: torch.Tensor | None = None, output_latents: bool = False, **kwargs: Unpack[TransformersKwargs], ) -> tuple | NeuCodecOutput: r""" input_values (`torch.Tensor` of shape `(batch_size, 1, sequence_length)`): Input audio waveform, sampled at `config.input_sampling_rate`. input_features (`torch.Tensor` of shape `(batch_size, mel_bins, time_steps)`): Input audio mel spectrogram for semantic encoding. padding_mask (`torch.Tensor` of shape `(batch_size, 1, sequence_length)`): Padding mask used to pad `input_values`. input_features_mask (`torch.Tensor` of shape `(batch_size, time_steps)`, *optional*): Attention mask for the spectrogram input to the semantic encoder. `1` for valid frames, `0` for padding. output_latents (`bool`, *optional*, defaults to `False`): Whether to return the continuous latent representation from the quantizer. Examples: ```python >>> from datasets import load_dataset >>> from transformers import AutoFeatureExtractor, NeuCodecModel >>> dataset = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation") >>> audio = dataset["train"]["audio"][0]["array"] >>> model_id = "neuphonic/neucodec" >>> model = NeuCodecModel.from_pretrained(model_id) >>> feature_extractor = AutoFeatureExtractor.from_pretrained(model_id) >>> inputs = feature_extractor(audio=audio, sampling_rate=feature_extractor.sampling_rate, return_tensors="pt") >>> outputs = model(**inputs) >>> audio_codes = outputs.audio_codes >>> audio_values = outputs.audio_values # sampled at 24kHz ```""" # NeuCodec's decoder outputs audio at `output_sampling_rate`, which differs from the `input_sampling_rate` of input_length = input_values.shape[-1] output_length = int(input_length * self.sample_rate_conversion_factor) encoder_outputs = self.encode( input_values, input_features=input_features, padding_mask=padding_mask, input_features_mask=input_features_mask, output_latents=True, return_dict=True, ) audio_values = self.decode( latents=encoder_outputs.latents, audio_codes_mask=encoder_outputs.audio_codes_mask, return_dict=True, **kwargs, )[0][..., :output_length] return NeuCodecOutput( audio_values=audio_values, audio_codes=encoder_outputs.audio_codes, latents=encoder_outputs.latents if output_latents else None, audio_codes_mask=encoder_outputs.audio_codes_mask, ) __all__ = ["NeuCodecModel", "NeuCodecPreTrainedModel"]