# coding=utf-8 # Copyright 2024 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. """Transformers XYTokenizer model.""" import math from collections import defaultdict from dataclasses import asdict, dataclass from typing import Optional, Tuple, Union, List import numpy as np import torch import torch.distributed as dist import torch.nn as nn import torch.nn.functional as F from einops import rearrange from torch.nn.utils.parametrizations import weight_norm from transformers.activations import ACT2FN from transformers.modeling_utils import PreTrainedAudioTokenizerBase from transformers.utils import ModelOutput, logging from transformers.feature_extraction_utils import BatchFeature from .configuration_xy_tokenizer import XYTokenizerConfig from .feature_extraction_xy_tokenizer import ExtractorIterator logger = logging.get_logger(__name__) # ----------------------------------------------- # # Model Output Dataclasses # # ----------------------------------------------- # @dataclass class XYTokenizerEncodeOutput(ModelOutput): """ Output type of [`XYTokenizerModel.encode`]. Args: quantized_representation (`torch.FloatTensor` of shape `(batch_size, hidden_dim, sequence_length)`): The quantized continuous representation of the input audio. This is the output of the quantizer. audio_codes (`torch.LongTensor` of shape `(num_codebooks, batch_size, sequence_length)`): The discrete codes from the quantizer for each codebook. codes_lengths (`torch.LongTensor` of shape `(batch_size,)`): The valid length of each sequence in `audio_codes`. commit_loss (`torch.FloatTensor`, *optional*): The commitment loss from the vector quantizer. overlap_seconds (`int`, *optional*): The duration of the overlap in seconds between adjacent audio chunks. """ quantized_representation: torch.FloatTensor = None audio_codes: torch.LongTensor = None codes_lengths: torch.LongTensor = None commit_loss: Optional[torch.FloatTensor] = None overlap_seconds: Optional[int] = None @dataclass class XYTokenizerDecodeOutput(ModelOutput): """ Output type of [`XYTokenizerModel.decode`]. Args: audio_values (`torch.FloatTensor` of shape `(batch_size, 1, sequence_length)`): The reconstructed audio waveform. output_length (`torch.LongTensor` of shape `(batch_size,)`): The valid length of each sequence in `audio_values`. """ audio_values: torch.FloatTensor = None output_length: Optional[torch.LongTensor] = None @dataclass class XYTokenizerModelOutput(ModelOutput): """ Output type of [`XYTokenizerModel`]'s forward pass. Args: audio_values (`torch.FloatTensor` of shape `(batch_size, 1, sequence_length)`): The reconstructed audio waveform. output_length (`torch.LongTensor` of shape `(batch_size,)`): The valid length of each sequence in `audio_values`. quantized_representation (`torch.FloatTensor` of shape `(batch_size, hidden_dim, sequence_length)`): The quantized continuous representation of the input audio. This is the output of the quantizer. audio_codes (`torch.LongTensor` of shape `(num_codebooks, batch_size, sequence_length)`): The discrete codes from the quantizer for each codebook. codes_lengths (`torch.LongTensor` of shape `(batch_size,)`): The valid length of each sequence in `audio_codes`. commit_loss (`torch.FloatTensor`, *optional*): The commitment loss from the vector quantizer. """ audio_values: torch.FloatTensor = None output_length: torch.LongTensor = None quantized_representation: torch.FloatTensor = None audio_codes: torch.LongTensor = None codes_lengths: torch.LongTensor = None commit_loss: Optional[torch.FloatTensor] = None @dataclass class VectorQuantizerConfig: """Configuration for the VectorQuantize module.""" commitment: float = 1.0 decay: float = 0.99 epsilon: float = 1e-5 threshold_ema_dead: int = 2 kmeans_init: bool = True kmeans_iters: int = 10 # ----------------------------------------------- # # All Helper Modules (Copied from source) # # ----------------------------------------------- # def sinusoids(length, channels, max_timescale=10000, device=None): assert channels % 2 == 0 log_timescale_increment = np.log(max_timescale) / (channels // 2 - 1) inv_timescales = torch.exp(-log_timescale_increment * torch.arange(channels // 2)) scaled_time = torch.arange(length, device=device)[:, np.newaxis] * inv_timescales[np.newaxis, :] return torch.cat([torch.sin(scaled_time), torch.cos(scaled_time)], dim=1) def get_sequence_mask(inputs, inputs_length): if inputs.dim() == 3: bsz, tgt_len, _ = inputs.size() else: bsz, tgt_len = inputs_length.shape[0], torch.max(inputs_length) sequence_mask = torch.arange(0, tgt_len, device=inputs.device) sequence_mask = torch.lt(sequence_mask, inputs_length.reshape(bsz, 1)).view(bsz, tgt_len, 1) return sequence_mask class RMSNorm(nn.Module): def __init__(self, hidden_size, eps=1e-6): super().__init__() self.weight = nn.Parameter(torch.ones(hidden_size)) self.variance_epsilon = eps def forward(self, hidden_states): variance = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=True) hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) if self.weight.dtype in [torch.float16, torch.bfloat16]: hidden_states = hidden_states.to(self.weight.dtype) return self.weight * hidden_states class VarLenAttention(nn.Module): def __init__(self, embed_dim, num_heads, causal=False, dropout=0.0): super().__init__() self.embed_dim = embed_dim self.num_heads = num_heads self.head_dim = embed_dim // num_heads assert embed_dim % num_heads == 0, "embed_dim must be divisible by num_heads" self.causal = causal self.dropout = nn.Dropout(dropout) self.scaling = self.head_dim ** -0.5 self.k_proj = nn.Linear(embed_dim, embed_dim, bias=False) self.v_proj = nn.Linear(embed_dim, embed_dim, bias=True) self.q_proj = nn.Linear(embed_dim, embed_dim, bias=True) self.out_proj = nn.Linear(embed_dim, embed_dim, bias=True) def _create_attention_mask(self, seq_len, max_len, device, dtype): bsz = seq_len.size(0) mask = torch.ones(bsz, 1, max_len, max_len, device=device, dtype=dtype) seq_indices = torch.arange(max_len, device=device).unsqueeze(0) seq_len_expanded = seq_len.unsqueeze(1) valid_mask = seq_indices < seq_len_expanded.unsqueeze(-1) mask = mask * (valid_mask.unsqueeze(2) & valid_mask.unsqueeze(3)).to(dtype) if self.causal: causal_mask = torch.triu(torch.ones(max_len, max_len, device=device, dtype=torch.bool), diagonal=1) mask = mask * (~causal_mask.unsqueeze(0).unsqueeze(1)).to(dtype) mask = mask + (1.0 - mask) * torch.finfo(dtype).min return mask def forward(self, hidden_states: torch.Tensor, seq_len: torch.Tensor) -> torch.Tensor: bsz, max_len, _ = hidden_states.size() query = self.q_proj(hidden_states) * self.scaling key = self.k_proj(hidden_states) value = self.v_proj(hidden_states) query = query.view(bsz, max_len, self.num_heads, self.head_dim).transpose(1, 2) key = key.view(bsz, max_len, self.num_heads, self.head_dim).transpose(1, 2) value = value.view(bsz, max_len, self.num_heads, self.head_dim).transpose(1, 2) attn_scores = torch.matmul(query, key.transpose(-1, -2)) attn_mask = self._create_attention_mask(seq_len, max_len, hidden_states.device, attn_scores.dtype) attn_scores = attn_scores + attn_mask attn_weights = F.softmax(attn_scores, dim=-1) attn_weights = self.dropout(attn_weights) attn_output = torch.matmul(attn_weights, value) attn_output = attn_output.transpose(1, 2).contiguous().view(bsz, max_len, self.embed_dim) attn_output = self.out_proj(attn_output) return attn_output class OmniWhisperMLP(nn.Module): def __init__(self, activation_function="gelu", d_model=1280, ffn_dim=5120): super().__init__() self.activation_fn = ACT2FN[activation_function] self.fc1 = nn.Linear(d_model, ffn_dim) self.fc2 = nn.Linear(ffn_dim, d_model) def forward(self, hidden_states): hidden_states = self.activation_fn(self.fc1(hidden_states)) return self.fc2(hidden_states) class OmniWhisperTransformerLayer(nn.Module): def __init__(self, activation_function="gelu", d_model=1280, attention_heads=20, ffn_dim=5120, causal=False, ln_type="LayerNorm", attn_type="varlen"): super().__init__() self.embed_dim = d_model if attn_type != "varlen": raise ValueError(f"Unknown attn_type: {attn_type}. Only 'varlen' is supported.") self.self_attn = VarLenAttention(self.embed_dim, attention_heads, causal) if ln_type == "LayerNorm": self.self_attn_layer_norm = nn.LayerNorm(self.embed_dim) elif ln_type == "RMSNorm": self.self_attn_layer_norm = RMSNorm(self.embed_dim) else: raise ValueError(f"Unknown ln_type: {ln_type}") self.mlp = OmniWhisperMLP(activation_function, d_model, ffn_dim) if ln_type == "LayerNorm": self.final_layer_norm = nn.LayerNorm(self.embed_dim) elif ln_type == "RMSNorm": self.final_layer_norm = RMSNorm(self.embed_dim) else: raise ValueError(f"Unknown ln_type: {ln_type}") def forward(self, hidden_states: torch.Tensor, seq_len: torch.Tensor) -> torch.Tensor: residual = hidden_states hidden_states = self.self_attn_layer_norm(hidden_states) hidden_states = self.self_attn(hidden_states, seq_len) hidden_states = residual + hidden_states residual = hidden_states hidden_states = self.final_layer_norm(hidden_states) hidden_states = self.mlp(hidden_states) hidden_states = residual + hidden_states if (hidden_states.dtype == torch.float16 or hidden_states.dtype == torch.bfloat16) and \ (torch.isinf(hidden_states).any() or torch.isnan(hidden_states).any()): clamp_value = torch.finfo(hidden_states.dtype).max - 1000 hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value) return hidden_states class OmniAudioEncoder(nn.Module): def __init__( self, num_mel_bins=128, sampling_rate=16000, hop_length=160, stride_size=2, kernel_size=3, d_model=1280, scale_embedding=True, max_audio_seconds=30, encoder_layers=32, encoder_attention_heads=20, encoder_ffn_dim=5120, activation_function="gelu", attn_type="varlen" ): super().__init__() self.max_source_positions = (max_audio_seconds * sampling_rate // hop_length) // stride_size self.embed_scale = math.sqrt(d_model) if scale_embedding else 1.0 self.num_mel_bins, self.d_model, self.stride_size = num_mel_bins, d_model, stride_size self.conv1 = nn.Conv1d(num_mel_bins, d_model, kernel_size=kernel_size, padding=1) self.conv2 = nn.Conv1d(d_model, d_model, kernel_size=kernel_size, stride=stride_size, padding=1) self.register_buffer("positional_embedding", sinusoids(self.max_source_positions, d_model)) self.layers = nn.ModuleList([ OmniWhisperTransformerLayer(activation_function, d_model, encoder_attention_heads, encoder_ffn_dim, False, attn_type=attn_type) for _ in range(encoder_layers) ]) self.layer_norm = nn.LayerNorm(d_model) def forward(self, input_features, input_length, output_hidden_states=False): input_features = input_features.to(self.conv1.weight.dtype) inputs_embeds = F.gelu(self.conv1(input_features)) inputs_embeds = F.gelu(self.conv2(inputs_embeds)) output_length = (input_length // self.stride_size).long() hidden_states = inputs_embeds.permute(0, 2, 1) bsz, tgt_len, _ = hidden_states.size() pos_embed = self.positional_embedding[:tgt_len] if tgt_len < self.positional_embedding.shape[0] else self.positional_embedding hidden_states = (hidden_states.to(torch.float32) + pos_embed).to(hidden_states.dtype) attention_mask = get_sequence_mask(hidden_states, output_length) all_hidden = () if output_hidden_states else None for layer in self.layers: if output_hidden_states: all_hidden += (hidden_states,) hidden_states = layer(hidden_states, output_length) hidden_states = self.layer_norm(hidden_states) if output_hidden_states: all_hidden += (hidden_states,) hidden_states = torch.where(attention_mask, hidden_states, 0).transpose(1, 2) if not output_hidden_states: return hidden_states, output_length return hidden_states, output_length, all_hidden class OmniAudioDecoder(nn.Module): def __init__( self, num_mel_bins=128, sampling_rate=16000, hop_length=160, stride_size=2, kernel_size=3, d_model=1280, scale_embedding=True, max_audio_seconds=30, decoder_layers=32, decoder_attention_heads=20, decoder_ffn_dim=5120, activation_function="gelu", attn_type="varlen" ): super().__init__() self.max_source_positions = (max_audio_seconds * sampling_rate // hop_length) // stride_size self.embed_scale = math.sqrt(d_model) if scale_embedding else 1.0 self.num_mel_bins, self.d_model, self.stride_size = num_mel_bins, d_model, stride_size self.deconv1 = nn.ConvTranspose1d(d_model, d_model, kernel_size, stride_size, padding=0, output_padding=0) self.deconv2 = nn.ConvTranspose1d(d_model, num_mel_bins, kernel_size, stride=1, padding=0) self.register_buffer("positional_embedding", sinusoids(self.max_source_positions, d_model)) self.layers = nn.ModuleList([ OmniWhisperTransformerLayer(activation_function, d_model, decoder_attention_heads, decoder_ffn_dim, False, attn_type=attn_type) for _ in range(decoder_layers) ]) self.layer_norm = nn.LayerNorm(d_model) def forward(self, hidden_states, input_length): hidden_states = hidden_states.transpose(1, 2) bsz, tgt_len, _ = hidden_states.size() pos_embed = self.positional_embedding[:tgt_len] if tgt_len < self.positional_embedding.shape[0] else self.positional_embedding hidden_states = (hidden_states.to(torch.float32) + pos_embed).to(hidden_states.dtype) attention_mask = get_sequence_mask(hidden_states, input_length) for layer in self.layers: hidden_states = layer(hidden_states, input_length) hidden_states = self.layer_norm(hidden_states) hidden_states = torch.where(attention_mask, hidden_states, 0).permute(0, 2, 1) output_features = F.gelu(self.deconv1(hidden_states)) output_features = F.gelu(self.deconv2(output_features)) expected_length = tgt_len * self.stride_size if output_features.size(2) > expected_length: output_features = output_features[:, :, :expected_length] output_length = input_length * self.stride_size return output_features, output_length class ResidualDownConv(nn.Module): def __init__(self, d_model=1280, avg_pooler=4): super().__init__() self.d_model, self.avg_pooler = d_model, avg_pooler self.intermediate_dim = d_model * avg_pooler self.gate_proj = nn.Conv1d(d_model, self.intermediate_dim, avg_pooler, avg_pooler, bias=False) self.up_proj = nn.Conv1d(d_model, self.intermediate_dim, avg_pooler, avg_pooler, bias=False) self.down_proj = nn.Linear(self.intermediate_dim, self.intermediate_dim, bias=False) self.act_fn = ACT2FN['silu'] self.layer_norm = nn.LayerNorm(self.intermediate_dim) def forward(self, x, input_length): output_length = input_length // self.avg_pooler x = x.transpose(1, 2) batch_size, seq_len, _ = x.shape if seq_len % self.avg_pooler != 0: pad_size = self.avg_pooler - seq_len % self.avg_pooler x = F.pad(x, (0, 0, 0, pad_size), "constant", 0) # Pad sequence dim xt = x.permute(0, 2, 1) g, u = self.gate_proj(xt).permute(0, 2, 1), self.up_proj(xt).permute(0, 2, 1) x = x.reshape(batch_size, -1, self.intermediate_dim) c = self.down_proj(self.act_fn(g) * u) res = self.layer_norm(c + x).transpose(1, 2) return res, output_length class UpConv(nn.Module): def __init__(self, d_model=1280, stride=4): super().__init__() self.d_model, self.stride = d_model, stride self.up_conv = nn.ConvTranspose1d(self.stride * d_model, d_model, stride, stride, bias=False) def forward(self, x, input_length): res = self.up_conv(x) output_length = input_length * self.stride return res, output_length class Transformer(nn.Module): def __init__( self, input_dim=1280, d_model=1280, output_dim=1280, max_source_positions=1500, encoder_layers=32, encoder_attention_heads=20, encoder_ffn_dim=5120, activation_function="gelu", attn_type="varlen" ): super().__init__() self.input_dim, self.d_model, self.output_dim, self.max_source_positions = input_dim, d_model, output_dim, max_source_positions self.proj = nn.Linear(input_dim, d_model, bias=True) if input_dim != d_model else None self.register_buffer("positional_embedding", sinusoids(self.max_source_positions, d_model)) self.layers = nn.ModuleList([ OmniWhisperTransformerLayer(activation_function, d_model, encoder_attention_heads, encoder_ffn_dim, False, attn_type=attn_type) for _ in range(encoder_layers) ]) self.layer_norm = nn.LayerNorm(d_model) self.out_proj = nn.Linear(d_model, output_dim, bias=True) if output_dim != d_model else None def forward(self, input_features, input_length, output_hidden_states=False): output_length = input_length.long() hidden_states = self.proj(input_features.permute(0, 2, 1)).permute(0, 2, 1) if self.proj else input_features hidden_states = hidden_states.permute(0, 2, 1) bsz, tgt_len, _ = hidden_states.size() pos_embed = self.positional_embedding[:tgt_len] if tgt_len < self.positional_embedding.shape[0] else self.positional_embedding hidden_states = (hidden_states.to(torch.float32) + pos_embed).to(hidden_states.dtype) attention_mask = get_sequence_mask(hidden_states, output_length) all_hidden = () if output_hidden_states else None for layer in self.layers: if output_hidden_states: all_hidden += (hidden_states,) hidden_states = layer(hidden_states, output_length) hidden_states = self.layer_norm(hidden_states) if output_hidden_states: all_hidden += (hidden_states,) hidden_states = torch.where(attention_mask, hidden_states, 0).transpose(1, 2) if self.out_proj: hidden_states = self.out_proj(hidden_states.permute(0, 2, 1)).permute(0, 2, 1) if not output_hidden_states: return hidden_states, output_length return hidden_states, output_length, all_hidden # Note: The other helper classes like STFT, ISTFT, Vocos, VectorQuantize, etc., # would be placed here. For brevity, they are omitted but are required dependencies. # Assuming they are defined in the same way as the user provided code. # The code below will assume these classes are defined in the current scope. # ... [Paste all other helper classes here] ... class ISTFT(nn.Module): def __init__(self, n_fft: int, hop_length: int, win_length: int, padding: str = "same"): super().__init__() if padding not in ["center", "same"]: raise ValueError("Padding must be 'center' or 'same'.") self.padding, self.n_fft, self.hop_length, self.win_length = padding, n_fft, hop_length, win_length self.register_buffer("window", torch.hann_window(win_length)) def forward(self, spec: torch.Tensor) -> torch.Tensor: if self.padding == "center": return torch.istft(spec, self.n_fft, self.hop_length, self.win_length, self.window, center=True) elif self.padding == "same": pad = (self.win_length - self.hop_length) // 2 else: raise ValueError("Padding must be 'center' or 'same'.") B, N, T = spec.shape ifft = torch.fft.irfft(spec, self.n_fft, dim=1, norm="backward") * self.window[None, :, None] output_size = (T - 1) * self.hop_length + self.win_length y = F.fold(ifft, (1, output_size), (1, self.win_length), stride=(1, self.hop_length))[:, 0, 0, pad:-pad] window_sq = self.window.square().expand(1, T, -1).transpose(1, 2) window_envelope = torch.nn.functional.fold( window_sq, output_size=(1, output_size), kernel_size=(1, self.win_length), stride=(1, self.hop_length), ).squeeze()[pad:-pad] assert (window_envelope > 1e-11).all() return y / window_envelope class FourierHead(nn.Module): def forward(self, x: torch.Tensor) -> torch.Tensor: raise NotImplementedError("Subclasses must implement the forward method.") class ISTFTHead(FourierHead): def __init__(self, dim: int, n_fft: int, hop_length: int, padding: str = "same"): super().__init__() self.out = nn.Linear(dim, n_fft + 2) self.istft = ISTFT(n_fft, hop_length, n_fft, padding) def forward(self, x: torch.Tensor) -> torch.Tensor: x = self.out(x).transpose(1, 2) mag, p = x.chunk(2, dim=1) mag = torch.exp(mag).clip(max=1e2) s = mag.float() * (torch.cos(p).float() + 1j * torch.sin(p).float()) return self.istft(s).to(x.dtype) class AdaLayerNorm(nn.Module): def __init__(self, num_embeddings: int, embedding_dim: int, eps: float = 1e-6): super().__init__() self.eps, self.dim = eps, embedding_dim self.scale = nn.Embedding(num_embeddings, embedding_dim) self.shift = nn.Embedding(num_embeddings, embedding_dim) torch.nn.init.ones_(self.scale.weight) torch.nn.init.zeros_(self.shift.weight) def forward(self, x: torch.Tensor, cond_embedding_id: torch.Tensor) -> torch.Tensor: scale, shift = self.scale(cond_embedding_id), self.shift(cond_embedding_id) x = F.layer_norm(x, (self.dim,), eps=self.eps) return x * scale + shift class ConvNeXtBlock(nn.Module): def __init__(self, dim, intermediate_dim, layer_scale_init_value, adanorm_num_embeddings=None): super().__init__() self.dwconv = nn.Conv1d(dim, dim, 7, 1, 3, groups=dim) self.adanorm = adanorm_num_embeddings is not None self.norm = AdaLayerNorm(adanorm_num_embeddings, dim) if self.adanorm else nn.LayerNorm(dim, eps=1e-6) self.pwconv1 = nn.Linear(dim, intermediate_dim) self.act = nn.GELU() self.pwconv2 = nn.Linear(intermediate_dim, dim) self.gamma = nn.Parameter(layer_scale_init_value * torch.ones(dim), requires_grad=True) if layer_scale_init_value > 0 else None def forward(self, x, cond_embedding_id=None): res = x x = self.dwconv(x).transpose(1, 2) x = self.norm(x, cond_embedding_id) if self.adanorm else self.norm(x) x = self.pwconv2(self.act(self.pwconv1(x))) if self.gamma is not None: x = self.gamma * x x = res + x.transpose(1, 2) return x class Backbone(nn.Module): def forward(self, x: torch.Tensor, **kwargs) -> torch.Tensor: raise NotImplementedError("Subclasses must implement the forward method.") class VocosBackbone(Backbone): def __init__(self, input_channels, dim, intermediate_dim, num_layers, layer_scale_init_value=None, adanorm_num_embeddings=None): super().__init__() self.input_channels, self.embed = input_channels, nn.Conv1d(input_channels, dim, 7, 1, 3) self.adanorm = adanorm_num_embeddings is not None self.norm = AdaLayerNorm(adanorm_num_embeddings, dim) if self.adanorm else nn.LayerNorm(dim, eps=1e-6) self.convnext = nn.ModuleList([ConvNeXtBlock(dim, intermediate_dim, layer_scale_init_value or 1/num_layers, adanorm_num_embeddings) for _ in range(num_layers)]) self.final_layer_norm = nn.LayerNorm(dim, eps=1e-6) self.apply(self._init_weights) def _init_weights(self, m): if isinstance(m, (nn.Conv1d, nn.Linear)): nn.init.trunc_normal_(m.weight, std=0.02) if m.bias is not None: nn.init.constant_(m.bias, 0) def forward(self, x: torch.Tensor, **kwargs) -> torch.Tensor: x = self.embed(x).transpose(1, 2) x = self.norm(x, kwargs.get("bandwidth_id")) if self.adanorm else self.norm(x) x = x.transpose(1, 2) for block in self.convnext: x = block(x, kwargs.get("bandwidth_id")) return self.final_layer_norm(x.transpose(1, 2)) class Vocos(nn.Module): def __init__(self, input_channels=128, dim=512, intermediate_dim=4096, num_layers=30, n_fft=640, hop_size=160, padding="same", adanorm_num_embeddings=None): super().__init__() self.backbone = VocosBackbone(input_channels, dim, intermediate_dim, num_layers, adanorm_num_embeddings=adanorm_num_embeddings) self.head = ISTFTHead(dim, n_fft, hop_size, padding) self.hop_size = hop_size def forward(self, x, input_length): x = self.backbone(x) x = self.head(x) return x[:, None, :], input_length * self.hop_size def WNConv1d(*args, **kwargs): return weight_norm(nn.Conv1d(*args, **kwargs)) def ema_inplace(moving_avg, new, decay): moving_avg.data.mul_(decay).add_(new.float(), alpha=(1 - decay)) def sample_vectors(samples, num): num_samples, device = samples.shape[0], samples.device indices = torch.randperm(num_samples, device=device)[:num] if num_samples >= num else torch.randint(0, num_samples, (num,), device=device) return samples[indices].float() def kmeans(samples, num_clusters, num_iters=10): dim, means = samples.shape[-1], sample_vectors(samples, num_clusters).float() for _ in range(num_iters): dists = -(samples.float().pow(2).sum(1, keepdim=True) - 2 * samples.float() @ means.t() + means.t().float().pow(2).sum(0, keepdim=True)) buckets = dists.max(dim=-1).indices bins = torch.bincount(buckets, minlength=num_clusters) zero_mask = bins == 0 bins_min_clamped = bins.masked_fill(zero_mask, 1) new_means = buckets.new_zeros(num_clusters, dim, dtype=torch.float32).scatter_add_(0, buckets.unsqueeze(1).expand(-1, dim), samples.float()) / bins_min_clamped[..., None] means = torch.where(zero_mask[..., None], means, new_means) dists = -(samples.float().pow(2).sum(1, keepdim=True) - 2 * samples.float() @ means.t() + means.t().float().pow(2).sum(0, keepdim=True)) return means, torch.bincount(dists.max(dim=-1).indices, minlength=num_clusters).float() class VectorQuantize(nn.Module): def __init__(self, input_dim, codebook_size, codebook_dim, commitment=1.0, decay=0.99, epsilon=1e-5, threshold_ema_dead=2, kmeans_init=True, kmeans_iters=10): super().__init__() self.input_dim, self.codebook_size, self.codebook_dim = input_dim, codebook_size, codebook_dim self.commitment, self.decay, self.epsilon, self.threshold_ema_dead = commitment, decay, epsilon, threshold_ema_dead self.kmeans_init, self.kmeans_iters = kmeans_init, kmeans_iters self.in_project = WNConv1d(input_dim, codebook_dim, 1) if input_dim != codebook_dim else nn.Identity() self.out_project = WNConv1d(codebook_dim, input_dim, 1) if codebook_dim != input_dim else nn.Identity() self.register_buffer("codebook", torch.zeros(codebook_size, codebook_dim) if kmeans_init else torch.randn(codebook_size, codebook_dim)) self.register_buffer("inited", torch.tensor(not kmeans_init, dtype=torch.bool)) self.register_buffer("cluster_size", torch.zeros(codebook_size)) self.register_buffer("embed_avg", self.codebook.clone()) def ema_update(self, encodings, embed_onehot): encodings, embed_onehot = encodings.float(), embed_onehot.float() cluster_size_new, embed_sum = embed_onehot.sum(0), encodings.t() @ embed_onehot if dist.is_initialized(): dist.all_reduce(cluster_size_new) dist.all_reduce(embed_sum) ema_inplace(self.cluster_size, cluster_size_new, self.decay) ema_inplace(self.embed_avg, embed_sum.t(), self.decay) cluster_size = (self.cluster_size + self.epsilon) / (self.cluster_size.sum() + self.codebook_size * self.epsilon) * self.cluster_size.sum() self.codebook.copy_(self.embed_avg / cluster_size.unsqueeze(1)) def replace_dead_codes(self, encodings): if self.threshold_ema_dead == 0: return dead_mask = self.cluster_size < self.threshold_ema_dead if dead_mask.any(): samples = sample_vectors(encodings.float(), self.codebook_size) if not dist.is_initialized() or dist.get_rank() == 0 else torch.zeros_like(self.codebook) if dist.is_initialized(): dist.broadcast(samples, src=0) self.codebook[dead_mask] = samples[:dead_mask.sum()].to(self.codebook.dtype) def init_codebook(self, encodings): if self.inited.item(): return if not dist.is_initialized() or dist.get_rank() == 0: embed, cluster_sizes = kmeans(encodings.float(), self.codebook_size, self.kmeans_iters) else: embed, cluster_sizes = torch.zeros(self.codebook_size, self.codebook_dim, device=encodings.device), torch.zeros(self.codebook_size, device=encodings.device) if dist.is_initialized(): dist.broadcast(embed, src=0) dist.broadcast(cluster_sizes, src=0) self.codebook.copy_(embed) self.embed_avg.copy_(embed.clone()) self.cluster_size.copy_(cluster_sizes) self.inited.fill_(True) def forward(self, z): z_e = self.in_project(z.float()) encodings = rearrange(z_e, "b d t -> (b t) d") if self.kmeans_init and not self.inited.item(): self.init_codebook(encodings) dist = encodings.pow(2).sum(1, keepdim=True) - 2 * encodings @ self.codebook.float().t() + self.codebook.float().pow(2).sum(1, keepdim=True).t() indices = rearrange((-dist).max(1)[1], "(b t) -> b t", b=z.size(0)) z_q = self.decode_code(indices) commit_loss = F.mse_loss(z_e, z_q.detach(), reduction="none").mean([1, 2]) * self.commitment if self.training and torch.is_grad_enabled(): self.ema_update(encodings, F.one_hot(indices.view(-1), self.codebook_size)) self.replace_dead_codes(encodings) z_q = self.out_project(z_e + (z_q - z_e).detach()) return z_q, commit_loss, torch.tensor(0.0, device=z.device), indices, z_e def decode_code(self, embed_id): return F.embedding(embed_id, self.codebook.float()).transpose(1, 2) class ResidualVQ(nn.Module): def __init__( self, input_dim: int = 1280, rvq_dim: int = None, output_dim: int = None, num_quantizers: int = 32, codebook_size: int = 1024, codebook_dim: int = 8, quantizer_dropout: float = 0.5, skip_rvq_ratio: float = 0.0, vq_config: VectorQuantizerConfig = None, **kwargs ): super().__init__() self.input_dim, self.rvq_dim, self.output_dim = input_dim, rvq_dim, output_dim or input_dim self.num_quantizers, self.codebook_size, self.codebook_dim = num_quantizers, codebook_size, codebook_dim self.quantizer_dropout, self.skip_rvq_ratio = quantizer_dropout, skip_rvq_ratio self.input_proj = WNConv1d(input_dim, rvq_dim, 1) if input_dim != rvq_dim else nn.Identity() self.output_proj = WNConv1d(rvq_dim, self.output_dim, 1) if rvq_dim != self.output_dim else nn.Identity() if vq_config is None: vq_config = VectorQuantizerConfig() quantizer_kwargs = asdict(vq_config) self.quantizers = nn.ModuleList([VectorQuantize(rvq_dim, codebook_size, codebook_dim, **quantizer_kwargs, **kwargs) for _ in range(num_quantizers)]) def forward(self, z, input_length, n_quantizers: int = None): z = self.input_proj(z) with torch.autocast('cuda', enabled=False): batch_size, _, max_time = z.shape device = z.device mask = torch.arange(max_time, device=device).expand(batch_size, max_time) < input_length.unsqueeze(1) quantized_out = torch.zeros_like(z) residual = z.clone().float() all_commit_losses = [] all_indices = [] all_quantized = [] # --- Complexity Reduction Start --- # 1. Extracted logic for determining quantizer numbers and skip mask n_q_tensor = self._get_n_quantizers_tensor(batch_size, device, n_quantizers) skip_mask = self._get_skip_mask(batch_size, device) # --- Complexity Reduction End --- max_q_to_run = self.num_quantizers if self.training else (n_quantizers or self.num_quantizers) for i, quantizer in enumerate(self.quantizers[:max_q_to_run]): # Create a mask for which batch items are active in this iteration active_in_iteration_mask = (i < n_q_tensor) # Skip quantization for items that are not active if not active_in_iteration_mask.any(): # If no items are active, we can add placeholders and continue # This branch is less common but handles the case where all items have dropped out all_commit_losses.append(torch.tensor(0.0, device=device)) all_indices.append(torch.zeros(batch_size, max_time, dtype=torch.long, device=device)) all_quantized.append(torch.zeros_like(z)) continue masked_residual = residual * mask.unsqueeze(1) # --- Complexity Reduction Start --- # 2. Extracted quantization step logic z_q_i, commit_loss_i, indices_i = self._quantize_step(quantizer, masked_residual, skip_mask) # --- Complexity Reduction End --- # Create a mask for updating tensors (batch items active in this iteration AND within valid length) update_mask = (active_in_iteration_mask.view(-1, 1, 1) & mask.unsqueeze(1)) quantized_out += z_q_i * update_mask residual -= z_q_i * update_mask # Calculate average commitment loss only for active items commit_loss_i = commit_loss_i[active_in_iteration_mask].mean() if active_in_iteration_mask.any() else torch.tensor(0.0, device=device) all_commit_losses.append(commit_loss_i) all_indices.append(indices_i) all_quantized.append(z_q_i) # Pad the outputs if the loop was exited early (e.g., in eval mode with n_quantizers) num_loops_done = len(all_commit_losses) if num_loops_done < self.num_quantizers: remaining = self.num_quantizers - num_loops_done all_commit_losses.extend([torch.tensor(0.0, device=device)] * remaining) all_indices.extend([torch.zeros(batch_size, max_time, dtype=torch.long, device=device)] * remaining) all_quantized.extend([torch.zeros_like(z)] * remaining) quantized_out = self.output_proj(quantized_out) all_indices_tensor = torch.stack(all_indices) all_commit_losses_tensor = torch.stack(all_commit_losses) all_quantized_tensor = torch.stack(all_quantized) return ( quantized_out, all_indices_tensor, all_commit_losses_tensor, all_quantized_tensor, input_length, ) def decode_codes(self, codes): nq, B, T = codes.shape emb = torch.zeros(B, self.rvq_dim, T, device=codes.device, dtype=torch.float32) for i, quantizer in enumerate(self.quantizers[:nq]): emb += quantizer.decode_code(codes[i]) return self.output_proj(emb) def _get_n_quantizers_tensor(self, batch_size: int, device: torch.device, n_quantizers_override: Optional[int] = None) -> torch.Tensor: """ Determines the number of quantizers to use for each item in the batch, applying dropout during training. """ # If not training or dropout is disabled, use the override or default number of quantizers is_training = self.training and torch.is_grad_enabled() if not is_training or self.quantizer_dropout == 0: num_q = n_quantizers_override or self.num_quantizers return torch.full((batch_size,), num_q, dtype=torch.long, device=device) # During training, apply quantizer dropout n_q_tensor = torch.full((batch_size,), self.num_quantizers, device=device) n_dropout = int(batch_size * self.quantizer_dropout) if n_dropout > 0: dropout_indices = torch.randperm(batch_size, device=device)[:n_dropout] dropout_values = torch.randint(1, self.num_quantizers + 1, (n_dropout,), device=device) n_q_tensor[dropout_indices] = dropout_values return n_q_tensor def _get_skip_mask(self, batch_size: int, device: torch.device) -> Optional[torch.Tensor]: """Generates a mask for skipping RVQ during training if skip_rvq_ratio > 0.""" is_training = self.training and torch.is_grad_enabled() if not is_training or self.skip_rvq_ratio <= 0: return None skip_mask = torch.rand(batch_size, device=device) < self.skip_rvq_ratio # Ensure at least one sample is not skipped to avoid errors in modules like DDP if skip_mask.all(): skip_mask[0] = False return skip_mask def _quantize_step(self, quantizer, residual, skip_mask): """Helper to perform one step of quantization, handling the skip logic.""" # The main logic is for non-skipped samples z_q_i, commit_loss_i, _, indices_i, z_e_i = quantizer(residual.float()) # If skipping is active, overwrite the results for the masked samples if skip_mask is not None: # For skipped samples, the "quantized" output is the residual itself # and the loss is zero. skip_mask_expanded = skip_mask.view(-1, 1, 1) z_q_i = torch.where(skip_mask_expanded, residual, z_q_i) commit_loss_i = torch.where(skip_mask, torch.zeros_like(commit_loss_i), commit_loss_i) return z_q_i, commit_loss_i, indices_i # ----------------------------------------------- # # PreTrainedModel Base Class # # ----------------------------------------------- # class XYTokenizerPreTrainedModel(PreTrainedAudioTokenizerBase): """ An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained models. """ config_class = XYTokenizerConfig base_model_prefix = "xy_tokenizer" main_input_name = "input_values" _supports_grad_checkpointing = True def _init_weights(self, module): """Initialize the weights.""" if isinstance(module, (nn.Linear, nn.Conv1d, nn.ConvTranspose1d)): module.weight.data.normal_(mean=0.0, std=0.02) if module.bias is not None: module.bias.data.zero_() elif isinstance(module, nn.Embedding): module.weight.data.normal_(mean=0.0, std=0.02) if module.padding_idx is not None: module.weight.data[module.padding_idx].zero_() def _set_gradient_checkpointing(self, module, value=False): if isinstance(module, (OmniAudioEncoder, OmniAudioDecoder, Transformer)): module.gradient_checkpointing = value # ----------------------------------------------- # # Main Model Class # # ----------------------------------------------- # class XYTokenizerModel(XYTokenizerPreTrainedModel): def __init__(self, config: XYTokenizerConfig): super().__init__(config) # Reconstruct the nested parameter dictionaries from the flat config # This is a bit of a boilerplate but necessary to reuse the original module code. # A more integrated approach would refactor the sub-modules to accept the flat config directly. self.config = config params = config.params self.semantic_encoder = OmniAudioEncoder(**params['semantic_encoder_kwargs']) self.semantic_encoder_adapter = Transformer(**params['semantic_encoder_adapter_kwargs']) self.acoustic_encoder = OmniAudioEncoder(**params['acoustic_encoder_kwargs']) self.pre_rvq_adapter = Transformer(**params['pre_rvq_adapter_kwargs']) self.downsample = ResidualDownConv(**params['downsample_kwargs']) self.quantizer = ResidualVQ(**params['quantizer_kwargs']) self.post_rvq_adapter = Transformer(**params['post_rvq_adapter_kwargs']) self.upsample = UpConv(**params['upsample_kwargs']) self.acoustic_decoder = OmniAudioDecoder(**params['acoustic_decoder_kwargs']) self.enhanced_vocos = Vocos(**params['vocos_kwargs']) self.feature_extractor = params['feature_extractor_kwargs'] # Store some config values for easier access self.encoder_downsample_rate = config.encoder_downsample_rate self.nq = params['quantizer_kwargs']['num_quantizers'] # Initialize weights and apply final processing self.post_init() def _get_feat_extract_output_lengths(self, input_lengths: Optional[torch.Tensor]): """ Computes the output lengths of the feature extractor. """ def _get_out_len(in_len): return (in_len - self.feature_extractor["n_fft"]) // self.feature_extractor["hop_length"] + 1 if input_lengths is None: return None return torch.tensor([_get_out_len(l) for l in input_lengths], device=self.device) def scale_window_size(self, boundaries, scaling_factor): scaling_range = [] scaling_boundaries = [] for left_boundary, right_boundary in boundaries: scaling_left_boundary = left_boundary// scaling_factor scaling_right_boundary = right_boundary // scaling_factor scaling_range.append(scaling_right_boundary-scaling_left_boundary) scaling_boundaries.append(slice(scaling_left_boundary, scaling_right_boundary)) return scaling_range, scaling_boundaries @torch.inference_mode def encode( self, features: Union[BatchFeature, ExtractorIterator], n_quantizers: Optional[int] = None, return_dict: Optional[bool] = True, ) -> Union[XYTokenizerEncodeOutput, Tuple]: r""" Encodes the input audio waveform into discrete codes. Args: features (`BatchFeature` or `ExtractorIterator`): A single batch of features or an iterator that yields batches of chunks for long audio files. The iterator is expected to yield `BatchFeature` dicts which must contain a `sequence_ids` tensor of shape `(batch_size,)` mapping each item in the chunk to its original sequence. n_quantizers (`int`, *optional*): The number of quantizers to use. If not specified, all quantizers are used. return_dict (`bool`, *optional*): Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. Returns: [`XYTokenizerEncodeOutput`] or `tuple(torch.FloatTensor)` """ assert isinstance(features, (BatchFeature, ExtractorIterator)) # Handle single batch case if isinstance(features, BatchFeature): return self._encode(features, n_quantizers, return_dict) # Handle streaming/chunked case else: # Use a dictionary to group chunks by their original sequence ID encodings = defaultdict(lambda: {"zq": [], "codes": [], "length": 0}) commit_losses = [] total_frames = 0 # 1. Iterate through chunks and store intermediate results for chunk_features in features: # Always use return_dict=True for easier access to named outputs chunk_output = self._encode(chunk_features, n_quantizers, return_dict=True) valid_code_lengths, valid_code_ranges = self.scale_window_size(chunk_features["input_lengths"], self.encoder_downsample_rate) # Accumulate weighted commit loss chunk_length = chunk_output.codes_lengths.sum().item() valid_chunk_length = sum(valid_code_lengths) if chunk_output.commit_loss is not None and valid_chunk_length > 0: commit_loss = chunk_output.commit_loss / chunk_length * valid_chunk_length commit_losses.append((commit_loss.cpu(), valid_chunk_length)) total_frames += valid_chunk_length # Group results by original sequence ID for i, seq_id in enumerate(chunk_features["chunk_seq_no"].tolist()): valid_code_range = valid_code_ranges[i] if valid_code_range.stop > 0: encodings[seq_id]["zq"].append(chunk_output.quantized_representation[i:i+1, :, valid_code_range]) encodings[seq_id]["codes"].append(chunk_output.audio_codes[:, i:i+1, valid_code_range]) # Add the valid length of this chunk to the total for this sequence encodings[seq_id]["length"] += valid_code_lengths[i] final_outputs = [] for seq_id, seq_data in encodings.items(): final_outputs.append({ "zq": torch.cat(seq_data["zq"], dim=2), "codes": torch.cat(seq_data["codes"], dim=2), "length": seq_data["length"] }) # 3. Pad all sequences to the same length and stack into a batch max_len = max(seq["zq"].shape[2] for seq in final_outputs) batch_zq = [] batch_codes = [] batch_lengths = [] for seq in final_outputs: pad_amount = max_len - seq["zq"].shape[2] # Pad on the right side of the last dimension (time) padded_zq = F.pad(seq["zq"], (0, pad_amount)) padded_codes = F.pad(seq["codes"], (0, pad_amount)) batch_zq.append(padded_zq) batch_codes.append(padded_codes) batch_lengths.append(seq["length"]) # Stack the list of tensors into a single batch tensor quantized_representation = torch.cat(batch_zq, dim=0) audio_codes = torch.cat(batch_codes, dim=0) codes_lengths = torch.tensor(batch_lengths, dtype=torch.long, device=self.device) # 4. Calculate final commit loss if total_frames > 0: # Weighted average of commit losses commit_loss = sum(loss * length for loss, length in commit_losses) / total_frames commit_loss = commit_loss.to(self.device) else: commit_loss = torch.tensor(0.0, device=self.device) if not return_dict: return (quantized_representation, audio_codes, codes_lengths, commit_loss) return XYTokenizerEncodeOutput( quantized_representation=quantized_representation, audio_codes=audio_codes, codes_lengths=codes_lengths, commit_loss=commit_loss, overlap_seconds=features.overlap_seconds, ) def _encode( self, features: BatchFeature, n_quantizers: Optional[int] = None, return_dict: Optional[bool] = True, ) -> Union[XYTokenizerEncodeOutput, Tuple]: input_mel = features['input_features'].to(self.device, dtype=self.dtype) mel_attention_mask = features['attention_mask'].to(self.device) mel_output_length = mel_attention_mask.sum(dim=-1).long() # --- Encoder Path --- semantic_encoder_output, semantic_encoder_output_length = self.semantic_encoder(input_mel, mel_output_length) semantic_adapter_output, _ = self.semantic_encoder_adapter(semantic_encoder_output, semantic_encoder_output_length) acoustic_encoder_output, acoustic_encoder_output_length = self.acoustic_encoder(input_mel, mel_output_length) concated_channel = torch.cat([semantic_adapter_output, acoustic_encoder_output], dim=1) pre_rvq_adapter_output, pre_rvq_adapter_output_length = self.pre_rvq_adapter(concated_channel, acoustic_encoder_output_length) downsample_output, downsample_output_length = self.downsample(pre_rvq_adapter_output, pre_rvq_adapter_output_length) n_quantizers = n_quantizers or self.quantizer.num_quantizers zq, codes, vq_loss, _, quantizer_output_length = self.quantizer(downsample_output, downsample_output_length, n_quantizers=n_quantizers) if not return_dict: return (zq, codes, quantizer_output_length, vq_loss) return XYTokenizerEncodeOutput( quantized_representation=zq, audio_codes=codes, codes_lengths=quantizer_output_length, commit_loss=vq_loss.mean() ) @torch.inference_mode def decode( self, audio_codes: Union[torch.Tensor, XYTokenizerEncodeOutput], overlap_seconds: int = 10, return_dict: Optional[bool] = True, ) -> Union[XYTokenizerDecodeOutput, Tuple]: r""" Decodes discrete codes back into an audio waveform. Args: audio_codes (`torch.LongTensor` of shape `(num_codebooks, batch_size, sequence_length)`): The discrete codes from the quantizer for each codebook. codes_lengths (`torch.LongTensor` of shape `(batch_size,)`, *optional*): The valid length of each sequence in `audio_codes`. If not provided, it's assumed to be the full length. return_dict (`bool`, *optional*): Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. Returns: [`XYTokenizerDecodeOutput`] or `tuple(torch.FloatTensor)` """ assert not isinstance(audio_codes, tuple), "try to set param `return_dict=True` for `codec.encode()` function" assert isinstance(audio_codes, (torch.Tensor, XYTokenizerEncodeOutput)), \ "only accept `torch.Tensor` or `XYTokenizerEncodeOutput` for `codec.decode()` function" if isinstance(audio_codes, XYTokenizerEncodeOutput): audio_codes = audio_codes.audio_codes if hasattr(audio_codes, "overlap_seconds"): overlap_seconds = audio_codes.overlap_seconds if overlap_seconds is None: overlap_seconds = 0 chunk_length = self.feature_extractor["chunk_length"] duration_seconds = chunk_length - overlap_seconds chunk_code_length = int(chunk_length * self.feature_extractor["sampling_rate"] // self.config.encoder_downsample_rate) # Maximum code length per chunk duration_code_length = int(duration_seconds * self.feature_extractor["sampling_rate"] // self.config.encoder_downsample_rate) # Valid code length per chunk duration_wav_length = duration_code_length * self.config.decoder_upsample_rate # Valid waveform length per chunk # Get maximum code length batch_size = audio_codes.shape[1] codes_list = [audio_codes[:, i, :] for i in range(batch_size)] max_code_length = max(codes.shape[-1] for codes in codes_list) batch_size = len(codes_list) codes_tensor = torch.zeros(self.nq, batch_size, max_code_length, device=self.device, dtype=torch.long) code_lengths = torch.zeros(batch_size, dtype=torch.long, device=self.device) for i, codes in enumerate(codes_list): codes_tensor[:, i, :codes.shape[-1]] = codes.to(self.device) code_lengths[i] = codes.shape[-1] # (B,) # Calculate number of chunks needed max_chunks = (max_code_length + duration_code_length - 1) // duration_code_length wav_list = [] # Process the entire batch in chunks for chunk_idx in range(max_chunks): start = chunk_idx * duration_code_length end = min(start + chunk_code_length, max_code_length) chunk_codes = codes_tensor[:, :, start:end] # (nq, B, T') chunk_code_lengths = torch.clamp(code_lengths - start, 0, end - start) # (B,) # Skip empty chunks if chunk_code_lengths.max() == 0: continue # Decode result = self._decode(chunk_codes, chunk_code_lengths) # {"y": (B, 1, T'), "output_length": (B,)} chunk_wav = result["audio_values"] # (B, 1, T') chunk_wav_lengths = result["output_length"] # (B,) # Extract valid portion valid_wav_lengths = torch.clamp(chunk_wav_lengths, 0, duration_wav_length) # (B,) valid_chunk_wav = torch.zeros(batch_size, 1, duration_wav_length, device=self.device) for b in range(batch_size): if valid_wav_lengths[b] > 0: valid_chunk_wav[b, :, :valid_wav_lengths[b]] = chunk_wav[b, :, :valid_wav_lengths[b]] # (B, 1, valid_wav_length) wav_list.append(valid_chunk_wav) # (B, 1, valid_wav_length) # Concatenate all chunks if wav_list: wav_tensor = torch.cat(wav_list, dim=-1) # (B, 1, T_total) syn_wav_list = [wav_tensor[i, :, :code_lengths[i] * self.config.decoder_upsample_rate] for i in range(batch_size)] # B * (1, T,) else: syn_wav_list = [torch.zeros(1, 0, device=self.device) for _ in range(batch_size)] # B * (1, 0,) if not return_dict: return (syn_wav_list,) return XYTokenizerDecodeOutput( audio_values=syn_wav_list ) def _decode( self, audio_codes: torch.Tensor, codes_lengths: Optional[torch.Tensor] = None, return_dict: Optional[bool] = True, ) -> Union[XYTokenizerDecodeOutput, Tuple]: r""" Decodes discrete codes back into an audio waveform. Args: audio_codes (`torch.LongTensor` of shape `(num_codebooks, batch_size, sequence_length)`): The discrete codes from the quantizer for each codebook. codes_lengths (`torch.LongTensor` of shape `(batch_size,)`, *optional*): The valid length of each sequence in `audio_codes`. If not provided, it's assumed to be the full length. return_dict (`bool`, *optional*): Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. Returns: [`XYTokenizerDecodeOutput`] or `tuple(torch.FloatTensor)` """ return_dict = return_dict if return_dict is not None else self.config.use_return_dict if codes_lengths is None: codes_lengths = torch.full((audio_codes.shape[1],), audio_codes.shape[2], device=self.device) # --- Decoder Path --- zq = self.quantizer.decode_codes(audio_codes) post_rvq_adapter_output, post_rvq_adapter_output_length = self.post_rvq_adapter(zq, codes_lengths) upsample_output, upsample_output_length = self.upsample(post_rvq_adapter_output, post_rvq_adapter_output_length) acoustic_decoder_output, acoustic_decoder_output_length = self.acoustic_decoder(upsample_output, upsample_output_length) y, vocos_output_length = self.enhanced_vocos(acoustic_decoder_output, acoustic_decoder_output_length) if not return_dict: return (y, vocos_output_length) return XYTokenizerDecodeOutput( audio_values=y, output_length=vocos_output_length ) def forward( self, input_values: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, n_quantizers: Optional[int] = None, return_dict: Optional[bool] = True, ) -> Union[XYTokenizerModelOutput, Tuple]: r""" The forward method that handles the full encoding and decoding process. Args: input_values (`torch.FloatTensor` of shape `(batch_size, sequence_length)`): Float values of the input audio waveform. attention_mask (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): Mask to avoid performing attention on padding token indices. n_quantizers (`int`, *optional*): The number of quantizers to use for encoding. If not specified, all quantizers are used. return_dict (`bool`, *optional*): Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. Examples: ```python >>> from transformers import AutoModel, AutoFeatureExtractor >>> from datasets import load_dataset, Audio >>> import torch >>> # This is a placeholder model name, replace with the actual one on the Hub >>> model_id = "your-namespace/xy-tokenizer-model" >>> model = AutoModel.from_pretrained(model_id) >>> # The feature extractor config is part of the model config, so it can be loaded this way >>> feature_extractor = AutoFeatureExtractor.from_pretrained(model_id) >>> # Load a dummy audio dataset >>> ds = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation") >>> audio_sample = ds[0]["audio"]["array"] >>> sampling_rate = ds[0]["audio"]["sampling_rate"] >>> # Process audio >>> inputs = feature_extractor(audio_sample, sampling_rate=feature_extractor.sampling_rate, return_tensors="pt") >>> # Encode to get codes >>> with torch.no_grad(): ... encoder_output = model.encode(inputs["input_values"], attention_mask=inputs["attention_mask"]) ... audio_codes = encoder_output.audio_codes >>> # Decode from codes >>> with torch.no_grad(): ... decoder_output = model.decode(audio_codes) ... reconstructed_audio = decoder_output.audio_values >>> # Full forward pass >>> with torch.no_grad(): ... model_output = model(**inputs) ... reconstructed_audio_fwd = model_output.audio_values >>> print(reconstructed_audio.shape) torch.Size([1, 1, 147200]) >>> print(torch.allclose(reconstructed_audio, reconstructed_audio_fwd)) True ``` Returns: [`XYTokenizerModelOutput`] or `tuple(torch.FloatTensor)` """ return_dict = return_dict if return_dict is not None else self.config.use_return_dict encoder_outputs = self.encode( input_values=input_values, attention_mask=attention_mask, n_quantizers=n_quantizers, return_dict=True ) decoder_outputs = self.decode( audio_codes=encoder_outputs, return_dict=True ) if not return_dict: return ( decoder_outputs.audio_values, decoder_outputs.output_length, encoder_outputs.quantized_representation, encoder_outputs.audio_codes, encoder_outputs.codes_lengths, encoder_outputs.commit_loss ) return XYTokenizerModelOutput( audio_values=decoder_outputs.audio_values, output_length=decoder_outputs.output_length, quantized_representation=encoder_outputs.quantized_representation, audio_codes=encoder_outputs.audio_codes, codes_lengths=encoder_outputs.codes_lengths, commit_loss=encoder_outputs.commit_loss )