"""MiMo audio: tokenizer, encoding utilities, and audio encoder.""" # Audio tokenizer adapted from https://github.com/XiaomiMiMo/MiMo-Audio-Tokenizer.git import logging import math import os import typing as tp from dataclasses import dataclass from functools import wraps from typing import Optional, Tuple import torch import torch.nn as nn import torch.nn.functional as F from einops import rearrange from transformers.activations import ACT2FN from transformers.configuration_utils import PretrainedConfig from transformers.modeling_utils import PreTrainedModel from transformers.models.qwen2.configuration_qwen2 import Qwen2Config from transformers.models.qwen2.modeling_qwen2 import Qwen2Model from sglang.srt.layers.quantization.base_config import QuantizationConfig from sglang.srt.runtime_context import get_server_args from sglang.srt.utils import is_cuda if is_cuda(): from sgl_kernel.flash_attn import flash_attn_varlen_func else: def flash_attn_varlen_func(*args, **kwargs): raise RuntimeError("MiMoAudioTokenizer requires CUDA to run.") logger = logging.getLogger(__name__) def _compute_default_rope_parameters( config=None, device=None, seq_len=None, **rope_kwargs ): if config is not None and len(rope_kwargs) > 0: raise ValueError( "Unexpected arguments: `**rope_kwargs` and `config` are mutually exclusive" ) if len(rope_kwargs) > 0: base = rope_kwargs["base"] dim = rope_kwargs["dim"] elif config is not None: base = config.rope_theta partial_rotary_factor = ( config.partial_rotary_factor if hasattr(config, "partial_rotary_factor") else 1.0 ) head_dim = getattr(config, "head_dim", None) if head_dim is None: head_dim = config.hidden_size // config.num_attention_heads logger.info( "audio.head_dim not set; defaulting to hidden_size/num_heads = %d", head_dim, ) dim = int(head_dim * partial_rotary_factor) attention_factor = 1.0 inv_freq = 1.0 / ( base ** ( torch.arange(0, dim, 2, dtype=torch.int64).to( device=device, dtype=torch.float ) / dim ) ) return inv_freq, attention_factor _ROPE_INIT_FUNCTIONS = { "default": _compute_default_rope_parameters, } def _dynamic_rope_update(rope_forward): def longrope_frequency_update(self, position_ids, device): seq_len = torch.max(position_ids) + 1 if hasattr(self.config, "original_max_position_embeddings"): original_max_position_embeddings = ( self.config.original_max_position_embeddings ) else: original_max_position_embeddings = self.config.max_position_embeddings if seq_len > original_max_position_embeddings: if not hasattr(self, "long_inv_freq"): self.long_inv_freq, _ = self.rope_init_fn( self.config, device, seq_len=original_max_position_embeddings + 1 ) self.register_buffer("inv_freq", self.long_inv_freq, persistent=False) else: self.original_inv_freq = self.original_inv_freq.to(device) self.register_buffer("inv_freq", self.original_inv_freq, persistent=False) def dynamic_frequency_update(self, position_ids, device): seq_len = torch.max(position_ids) + 1 if seq_len > self.max_seq_len_cached: # growth inv_freq, self.attention_scaling = self.rope_init_fn( self.config, device, seq_len=seq_len ) self.register_buffer("inv_freq", inv_freq, persistent=False) self.max_seq_len_cached = seq_len if ( seq_len < self.original_max_seq_len and self.max_seq_len_cached > self.original_max_seq_len ): self.original_inv_freq = self.original_inv_freq.to(device) self.register_buffer("inv_freq", self.original_inv_freq, persistent=False) self.max_seq_len_cached = self.original_max_seq_len @wraps(rope_forward) def wrapper(self, x, position_ids): if "dynamic" in self.rope_type: dynamic_frequency_update(self, position_ids, device=x.device) elif self.rope_type == "longrope": longrope_frequency_update(self, position_ids, device=x.device) return rope_forward(self, x, position_ids) return wrapper class AudioRotaryEmbedding(nn.Module): def __init__(self, base, dim, max_seq_len, rope_type="default", device=None): super().__init__() self.max_seq_len = max_seq_len self.rope_type = rope_type self.rope_init_fn = _ROPE_INIT_FUNCTIONS[self.rope_type] inv_freq, self.attention_scaling = self.rope_init_fn( device=device, base=base, dim=dim ) self.register_buffer("inv_freq", inv_freq, persistent=False) self.original_inv_freq = self.inv_freq @torch.no_grad() @_dynamic_rope_update def forward(self, x, position_ids): inv_freq_expanded = self.inv_freq[:, None].float().expand(-1, 1).to(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" ) with torch.autocast(device_type=device_type, enabled=False): # Force float32 freqs = ( inv_freq_expanded.float() @ position_ids_expanded.float() ).transpose(0, 1) 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 EuclideanCodebook(nn.Module): """Codebook with Euclidean distance (inference-only).""" def __init__( self, dim: int, codebook_size: int, kmeans_init: bool = False, **kwargs ): super().__init__() init_fn = self._uniform_init if not kmeans_init else torch.zeros embed = init_fn(codebook_size, dim) self.codebook_size = codebook_size self.register_buffer("inited", torch.Tensor([not kmeans_init])) self.register_buffer("cluster_size", torch.zeros(codebook_size)) self.register_buffer("embed", embed) self.register_buffer("embed_avg", embed.clone()) def preprocess(self, x): x = rearrange(x, "... d -> (...) d") return x def quantize(self, x): embed = self.embed.t() dist_val = -( x.pow(2).sum(1, keepdim=True) - 2 * x @ embed + embed.pow(2).sum(0, keepdim=True) ) embed_ind = dist_val.max(dim=-1).indices return embed_ind def postprocess_emb(self, embed_ind, shape): return embed_ind.view(*shape[:-1]) def dequantize(self, embed_ind): quantize = F.embedding(embed_ind, self.embed) return quantize def encode(self, x): shape = x.shape x = self.preprocess(x) embed_ind = self.quantize(x) embed_ind = self.postprocess_emb(embed_ind, shape) return embed_ind def decode(self, embed_ind): quantize = self.dequantize(embed_ind) return quantize @staticmethod def _uniform_init(*shape: int): t = torch.empty(shape) nn.init.kaiming_uniform_(t) return t class VectorQuantization(nn.Module): """Vector quantization with euclidean distance (inference-only).""" def __init__( self, dim: int, codebook_size: int, codebook_dim: tp.Optional[int] = None, kmeans_init: bool = True, **kwargs, ): super().__init__() _codebook_dim: int = codebook_dim if codebook_dim is not None else dim requires_projection = _codebook_dim != dim self.project_in = ( nn.Linear(dim, _codebook_dim) if requires_projection else nn.Identity() ) self.project_out = ( nn.Linear(_codebook_dim, dim) if requires_projection else nn.Identity() ) self._codebook = EuclideanCodebook( dim=_codebook_dim, codebook_size=codebook_size, kmeans_init=kmeans_init, ) self.codebook_size = codebook_size @property def codebook(self): return self._codebook.embed def encode(self, x): x = self.project_in(x) embed_in = self._codebook.encode(x) return embed_in def decode(self, embed_ind): quantize = self._codebook.decode(embed_ind) quantize = self.project_out(quantize) return quantize class ResidualVectorQuantization(nn.Module): """Residual vector quantization implementation. Follows Algorithm 1. in https://arxiv.org/pdf/2107.03312.pdf """ def __init__(self, *, num_quantizers, codebook_size, **kwargs): super().__init__() if isinstance(codebook_size, int): codebook_size = [codebook_size] * num_quantizers elif len(codebook_size) < num_quantizers: codebook_size += [codebook_size[-1]] * (num_quantizers - len(codebook_size)) self.layers = nn.ModuleList( [ VectorQuantization(codebook_size=codebook_size[i], **kwargs) for i in range(num_quantizers) ] ) def encode( self, x: torch.Tensor, n_q: tp.Optional[int] = None, st: tp.Optional[int] = None ) -> torch.Tensor: residual = x all_indices = [] n_q = len(self.layers) if n_q is None else n_q st = 0 if st is None else st for layer in self.layers[st:n_q]: indices = layer.encode(residual) quantized = layer.decode(indices) residual = residual - quantized all_indices.append(indices) out_indices = torch.stack(all_indices) return out_indices def decode(self, q_indices: torch.Tensor, st: int = 0) -> torch.Tensor: quantized_out = self.layers[st].decode(q_indices[0]) for i in range(1, len(q_indices)): layer = self.layers[st + i] quantized = layer.decode(q_indices[i]) quantized_out = quantized_out + quantized return quantized_out class ResidualVectorQuantizer(nn.Module): """Residual Vector Quantizer (inference-only).""" def __init__( self, dimension: int = 256, n_q: int = 8, bins: int | list = 1024, kmeans_init: bool = True, **kwargs, ): super().__init__() self.n_q = n_q self.vq = ResidualVectorQuantization( dim=dimension, codebook_size=bins, num_quantizers=n_q, kmeans_init=kmeans_init, ) def encode( self, x: torch.Tensor, n_q: tp.Optional[int] = None, st: tp.Optional[int] = None ) -> torch.Tensor: n_q = n_q if n_q else self.n_q st = st or 0 codes = self.vq.encode(x, n_q=n_q, st=st) return codes def decode(self, codes: torch.Tensor, st: int = 0) -> torch.Tensor: quantized = self.vq.decode(codes, st=st) return quantized class MiMoAudioTokenizerConfig(PretrainedConfig): model_type = "mimo_audio_tokenizer" def __init__( self, max_audio_seconds: int = 1800, stride_size: int = 2, avg_pooler: int = 1, d_model: int = 768, scale_embedding: bool = True, kernel_size: int = 3, activation_function: str = "gelu", encoder_layers: int = 8, encoder_skip_layer_id: int = None, encoder_attention_heads: int = 12, encoder_ffn_dim: int = 3072, encoder_causal: bool = False, encoder_attn_window_size: list = None, decoder_layers: int = 8, decoder_attention_heads: int = 12, decoder_ffn_dim: int = 3072, decoder_kernel_size: int = 3, decoder_stride_size: int = 2, decoder_causal: bool = True, decoder_attn_window_size: list = None, nfft: int = 1024, vocoder_dim: int = 512, vocoder_intermediate_dim: int = 4096, vocoder_num_layers: int = 30, n_mels: int = 80, sampling_rate: int = 24000, hop_length: int = 240, window_size: int = 1024, vocoder_padding: str = "same", fmin: int = 0, fmax: int = None, num_quantizers: int = 12, codebook_size: list = None, threshold_ema_dead_code: int = 10, position_embedding_type: str = "rope", rope_theta: int = 10000, rope_type: str = "default", ln_type: str = "LayerNorm", vocoder_attention_heads: int = 4, vocoder_attn_window_size: list = None, use_istft_only: bool = False, hybrid_attention: bool = False, hybrid_block_size: int = 8, swa_per_block: int = 2, **kwargs, ): super().__init__(**kwargs) self.max_audio_seconds = max_audio_seconds self.stride_size = stride_size self.avg_pooler = avg_pooler self.d_model = d_model self.scale_embedding = scale_embedding self.kernel_size = kernel_size self.activation_function = activation_function self.encoder_layers = encoder_layers self.encoder_skip_layer_id = encoder_skip_layer_id self.encoder_attention_heads = encoder_attention_heads self.encoder_ffn_dim = encoder_ffn_dim self.encoder_causal = encoder_causal self.encoder_attn_window_size = ( encoder_attn_window_size if encoder_attn_window_size is not None else [-1, -1] ) self.decoder_layers = decoder_layers self.decoder_attention_heads = decoder_attention_heads self.decoder_ffn_dim = decoder_ffn_dim self.decoder_kernel_size = decoder_kernel_size self.decoder_stride_size = decoder_stride_size self.decoder_causal = decoder_causal self.decoder_attn_window_size = ( decoder_attn_window_size if decoder_attn_window_size is not None else [-1, -1] ) self.nfft = nfft self.vocoder_dim = vocoder_dim self.vocoder_intermediate_dim = vocoder_intermediate_dim self.vocoder_num_layers = vocoder_num_layers self.n_mels = n_mels self.sampling_rate = sampling_rate self.hop_length = hop_length self.window_size = window_size self.vocoder_padding = vocoder_padding self.fmin = fmin self.fmax = fmax self.num_quantizers = num_quantizers self.codebook_size = codebook_size if codebook_size is not None else [1024] self.threshold_ema_dead_code = threshold_ema_dead_code self.position_embedding_type = position_embedding_type self.rope_theta = rope_theta self.rope_type = rope_type self.ln_type = ln_type self.vocoder_attention_heads = vocoder_attention_heads self.vocoder_attn_window_size = ( vocoder_attn_window_size if vocoder_attn_window_size is not None else [40, 10] ) self.use_istft_only = use_istft_only self.hybrid_attention = hybrid_attention self.hybrid_block_size = hybrid_block_size self.swa_per_block = swa_per_block 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).to(inputs.device) sequence_mask = torch.lt(sequence_mask, inputs_length.reshape(bsz, 1)).view( bsz, tgt_len, 1 ) unpacking_index = torch.cumsum(sequence_mask.to(torch.int64).view(-1), dim=0) - 1 return sequence_mask, unpacking_index def unpack_hidden_states( hidden_states, lengths, sequence_mask=None, unpacking_index=None ): bsz = lengths.shape[0] if sequence_mask is None or unpacking_index is None: sequence_mask, unpacking_index = get_sequence_mask(hidden_states, lengths) hidden_states = torch.index_select(hidden_states, 0, unpacking_index).view( bsz, torch.max(lengths), hidden_states.shape[-1] ) return torch.where(sequence_mask, hidden_states, 0) def get_position_ids(lengths): total_len = lengths.sum() offset = torch.cat([torch.zeros(1).to(lengths), lengths[:-1].cumsum(dim=0)]) offset = torch.repeat_interleave(offset, lengths) return torch.arange(0, total_len).to(offset) - offset LAYER_NORM = {"LayerNorm": nn.LayerNorm} class AudioEncoderAttention(nn.Module): def __init__( self, embed_dim: int, num_heads: int, window_size: Tuple[int, int] = (-1, -1), causal: bool = False, ): super().__init__() self.embed_dim = embed_dim self.num_heads = num_heads self.head_dim = embed_dim // num_heads self.window_size = window_size self.causal = causal 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 forward( self, hidden_states: torch.Tensor, cu_seqlens: torch.Tensor, max_seqlen: int, rope_position_embeddings=None, ): bsz, _ = hidden_states.size() query_states = self.q_proj(hidden_states).view( bsz, self.num_heads, self.head_dim ) key_states = self.k_proj(hidden_states).view(bsz, self.num_heads, self.head_dim) value_states = self.v_proj(hidden_states).view( bsz, self.num_heads, self.head_dim ) if rope_position_embeddings is not None: cos, sin = rope_position_embeddings query_states, key_states = self.apply_rotary_pos_emb( query_states, key_states, cos, sin ) attn_output = flash_attn_varlen_func( query_states, key_states, value_states, cu_seqlens, cu_seqlens, max_seqlen, max_seqlen, causal=self.causal, window_size=self.window_size, ) attn_output = attn_output.reshape(bsz, self.embed_dim) attn_output = self.out_proj(attn_output) return attn_output @staticmethod def _rotate_half(x): x1 = x[..., : x.shape[-1] // 2] x2 = x[..., x.shape[-1] // 2 :] return torch.cat((-x2, x1), dim=-1) @classmethod def apply_rotary_pos_emb(cls, q, k, cos, sin, unsqueeze_dim=1): cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) q_embed = (q * cos) + (cls._rotate_half(q) * sin) k_embed = (k * cos) + (cls._rotate_half(k) * sin) return q_embed, k_embed class AudioEncoderTransformerLayer(nn.Module): def __init__( self, config: MiMoAudioTokenizerConfig, causal: bool, attn_window_size: Tuple[int, int] = (-1, -1), ): super().__init__() self.embed_dim = config.d_model self.self_attn = AudioEncoderAttention( embed_dim=self.embed_dim, num_heads=config.encoder_attention_heads, window_size=attn_window_size, causal=causal, ) self.self_attn_layer_norm = LAYER_NORM[config.ln_type](self.embed_dim) self.activation_fn = ACT2FN[config.activation_function] self.fc1 = nn.Linear(self.embed_dim, config.encoder_ffn_dim) self.fc2 = nn.Linear(config.encoder_ffn_dim, self.embed_dim) self.final_layer_norm = LAYER_NORM[config.ln_type](self.embed_dim) def forward( self, hidden_states: torch.Tensor, cu_seqlens: torch.Tensor, max_seqlen: int, rope_position_embeddings: Tuple[torch.Tensor, torch.Tensor], ) -> torch.Tensor: residual = hidden_states hidden_states = self.self_attn_layer_norm(hidden_states) hidden_states = self.self_attn( hidden_states, cu_seqlens, max_seqlen, rope_position_embeddings=rope_position_embeddings, ) hidden_states = residual + hidden_states residual = hidden_states hidden_states = self.final_layer_norm(hidden_states) hidden_states = self.activation_fn(self.fc1(hidden_states)) hidden_states = self.fc2(hidden_states) hidden_states = residual + hidden_states return hidden_states class AudioEncoder(nn.Module): def __init__( self, config: MiMoAudioTokenizerConfig, ): super().__init__() self.config = config self.max_source_positions = ( config.max_audio_seconds * config.sampling_rate // config.hop_length ) // config.stride_size self.embed_scale = math.sqrt(config.d_model) if config.scale_embedding else 1.0 self.skip_layer_idx = config.encoder_skip_layer_id self.conv1 = nn.Conv1d( config.n_mels, config.d_model, kernel_size=config.kernel_size, padding=1, ) self.conv2 = nn.Conv1d( config.d_model, config.d_model, kernel_size=config.kernel_size, stride=config.stride_size, padding=1, ) self.position_embedding = AudioRotaryEmbedding( config.rope_theta, config.d_model // config.encoder_attention_heads, self.max_source_positions, config.rope_type, ) attn_window_sizes = [] if config.hybrid_attention: for i in range(config.encoder_layers): if i % config.swa_per_block < config.swa_per_block - 1: attn_window_sizes.append(tuple(config.encoder_attn_window_size)) else: attn_window_sizes.append((-1, -1)) else: attn_window_sizes = [ tuple(config.encoder_attn_window_size) ] * config.encoder_layers self.layers = nn.ModuleList( [ AudioEncoderTransformerLayer( config=config, causal=config.encoder_causal, attn_window_size=attn_window_sizes[i], ) for i in range(config.encoder_layers) ] ) self.layer_norm = LAYER_NORM[config.ln_type](config.d_model) if config.avg_pooler != 1: self.down_sample_layer = nn.Sequential( nn.Conv1d( config.d_model, config.d_model, config.avg_pooler, config.avg_pooler, bias=False, ), nn.GELU(), ) self.down_sample_norm = LAYER_NORM[config.ln_type](config.d_model) else: self.down_sample_layer = None if config.num_quantizers != 0: self.quantizer = ResidualVectorQuantizer( dimension=config.d_model, n_q=config.num_quantizers, bins=config.codebook_size, threshold_ema_dead_code=config.threshold_ema_dead_code, ) else: self.quantizer = None def get_features(self, input_features, output_length): input_features = input_features.to(self.conv1.weight) inputs_embeds = nn.functional.gelu(self.conv1(input_features)) inputs_embeds = nn.functional.gelu(self.conv2(inputs_embeds)) inputs_embeds = inputs_embeds.permute(0, 2, 1) bsz, tgt_len, _ = inputs_embeds.size() hidden_states = inputs_embeds position_ids = get_position_ids(output_length).long().to(input_features.device) rope_position_embeddings = self.position_embedding(input_features, position_ids) attention_mask, unpacking_index = get_sequence_mask( hidden_states, output_length ) hidden_states = torch.masked_select(hidden_states, attention_mask).view( torch.sum(output_length), self.config.d_model ) cu_seqlens = F.pad( torch.cumsum(output_length, dim=0), (1, 0), "constant", 0 ).to(device=hidden_states.device, dtype=torch.int32) max_seqlen = torch.max(output_length).to(torch.int32).item() skip_connect_hidden_states = 0.0 for idx, encoder_layer in enumerate(self.layers): hidden_states = encoder_layer( hidden_states, cu_seqlens, max_seqlen, rope_position_embeddings=rope_position_embeddings, ) if (self.skip_layer_idx is not None) and idx == self.skip_layer_idx - 1: skip_connect_hidden_states = hidden_states.clone() hidden_states += skip_connect_hidden_states hidden_states = self.layer_norm(hidden_states) if self.down_sample_layer is not None: hidden_states = torch.index_select(hidden_states, 0, unpacking_index).view( bsz, tgt_len, self.config.d_model ) if hidden_states.size(1) % self.config.avg_pooler: pad_len = ( self.config.avg_pooler - hidden_states.size(1) % self.config.avg_pooler ) hidden_states = torch.nn.functional.pad( hidden_states, (0, 0, 0, pad_len), mode="constant", value=0.0 ) tgt_len += pad_len tgt_len = tgt_len // self.config.avg_pooler hidden_states = self.down_sample_layer(hidden_states.transpose(1, 2)) output_length = ( output_length // self.config.avg_pooler + (output_length % self.config.avg_pooler != 0).int() ) hidden_states = hidden_states.transpose(1, 2) attention_mask, unpacking_index = get_sequence_mask( hidden_states, output_length ) hidden_states = torch.masked_select(hidden_states, attention_mask).view( torch.sum(output_length), self.config.d_model ) hidden_states = self.down_sample_norm(hidden_states) return ( hidden_states, output_length, attention_mask, unpacking_index, tgt_len, bsz, ) def get_output_length(self, mel_len): tgt_len = mel_len + 3 - self.config.kernel_size return (tgt_len + 2 - self.config.kernel_size) // self.config.stride_size + 1 @torch.no_grad() def encode( self, input_features, input_lens=None, output_length=None, return_codes_only=False, n_q=None, use_quantizer=True, ): if output_length is None: output_length = self.get_output_length(input_lens) input_features = unpack_hidden_states(input_features, input_lens) hidden_states, output_length, attention_mask, unpacking_index, tgt_len, bsz = ( self.get_features( input_features=input_features.transpose(1, 2), output_length=output_length, ) ) dtype = hidden_states.dtype if use_quantizer and self.quantizer is not None: self.quantizer.float() codes = self.quantizer.encode(hidden_states.float(), n_q=n_q) if return_codes_only: return codes, output_length hidden_states = self.quantizer.decode(codes) hidden_states = hidden_states.to(dtype) else: codes = None hidden_states_packed = hidden_states.clone() hidden_states = torch.index_select(hidden_states, 0, unpacking_index).view( bsz, tgt_len, self.config.d_model ) hidden_states = torch.where(attention_mask, hidden_states, 0) return hidden_states, hidden_states_packed, output_length, codes @torch.no_grad() def decode_vq(self, codes): self.quantizer.float() return self.quantizer.decode(codes) class MiMoAudioTokenizer(PreTrainedModel): config_class = MiMoAudioTokenizerConfig def __init__(self, config: MiMoAudioTokenizerConfig): super().__init__(config) self.config = config self.sampling_rate = config.sampling_rate self.encoder = AudioEncoder(config=config) self.downsample_rate = int(config.hop_length * 2 * config.avg_pooler) def get_output_length(self, mel_len): tgt_len = mel_len + 3 - self.config.kernel_size return (tgt_len + 2 - self.config.kernel_size) // self.config.stride_size + 1 @torch.no_grad() def encode(self, mels, input_lens, use_quantizer=True): input_features = mels encoder_output_length = self.get_output_length(input_lens) hidden_states, hidden_states_packed, encoder_output_length, codes = ( self.encoder.encode( input_features, input_lens=input_lens, use_quantizer=use_quantizer ) ) return hidden_states, hidden_states_packed, encoder_output_length, codes def group_by_length(features: torch.Tensor, lengths: torch.Tensor, max_length: int): if features.size(0) != lengths.sum().item(): raise ValueError( f"Feature size mismatch: {features.size(0)} vs {lengths.sum().item()}" ) split_points = [] current_sum = 0 for i, seq_len in enumerate(lengths): if current_sum + seq_len > max_length and current_sum > 0: split_points.append(i) current_sum = seq_len.item() else: current_sum += seq_len.item() # Convert split points to group sizes group_sizes = [] prev = 0 for point in split_points: group_sizes.append(point - prev) prev = point if prev < len(lengths): group_sizes.append(len(lengths) - prev) len_groups = torch.split(lengths, group_sizes) feature_sizes = [group.sum().item() for group in len_groups] feature_groups = torch.split(features, feature_sizes) return feature_groups, len_groups @torch.no_grad() def encode_batch( audio_tokenizer_encoder, input_features: torch.Tensor, input_lens: torch.Tensor, max_length: int = 256000, ): feature_groups, len_groups = group_by_length(input_features, input_lens, max_length) encoded_parts = [] for features, lengths in zip(feature_groups, len_groups): codes, _ = audio_tokenizer_encoder.encode( # codes are also packed input_features=features, input_lens=lengths, return_codes_only=True ) encoded_parts.append(codes) return torch.cat(encoded_parts, dim=-1) def _segment_lengths_for_mel(mel: torch.Tensor, segment_size: int): """Split mel into segments of segment_size with a possible shorter remainder.""" input_len = mel.size(0) segs = [segment_size] * (input_len // segment_size) if input_len % segment_size > 0: segs.append(input_len % segment_size) return segs @torch.no_grad() def tokenize_audio_batch(mels, audio_tokenizer_encoder, segment_size=6000, device=None): """ Tokenize multiple mels in one encode_batch call. Returns list of code tensors, each [T_i, C] for that mel. """ if not mels: return [] if device is None: device = next(audio_tokenizer_encoder.parameters()).device # Build segment lengths per mel input_len_seg_per_mel = [_segment_lengths_for_mel(m, segment_size) for m in mels] input_lens_flat = [s for segs in input_len_seg_per_mel for s in segs] input_features = torch.cat([m.to(device) for m in mels], dim=0) input_lens_t = torch.tensor(input_lens_flat, dtype=torch.long, device=device) codes_packed = encode_batch( audio_tokenizer_encoder, input_features=input_features, input_lens=input_lens_t, ) codes = codes_packed.transpose(0, 1).detach() # [total_code_T, C] # Code length per mel: must match encoder's actual output (get_output_length + optional avg_pooler downsampling) code_lengths = [] for segs in input_len_seg_per_mel: out_len = audio_tokenizer_encoder.get_output_length( torch.tensor(segs, dtype=torch.long, device=device) ) if getattr(audio_tokenizer_encoder, "down_sample_layer", None) is not None: avg = audio_tokenizer_encoder.config.avg_pooler out_len = out_len // avg + (out_len % avg != 0).long() code_lengths.append(out_len.sum().item()) code_list = torch.split(codes, code_lengths) return list(code_list) @dataclass class MiMoAudioEncoderConfig: tokenizer_version: str = "v1" speech_vocab_size: str = "1025-1025-129-129-129-129-129-129" speech_zeroemb_idx: str = "1024-1024-128-128-128-128-128-128" group_size: int = 4 audio_channels: int = 8 input_local_layers: int = 6 input_local_dim: int = 1024 input_full_attention: bool = True input_local_attn_heads: int = 64 input_local_head_dim: int = 16 input_local_intermediate_size: int = 4096 input_local_hidden_dropout: float = 0.0 out_hidden_size: int = 4096 # mimo vl hidden dim rope_theta: float = 640000.0 partial_rotary_factor: float = 0.334 projection_layers: int = 1 add_post_norm: bool = False audio_segment_size: int = 6000 class AudioProjection(nn.Module): def __init__( self, input_size, hidden_size, output_size, quant_config: Optional[QuantizationConfig] = None, prefix: str = "", ) -> None: super().__init__() self.input_size = input_size self.hidden_size = hidden_size self.output_size = output_size self.mlp = nn.Sequential( nn.Linear(self.input_size, self.hidden_size, bias=False), nn.GELU(), nn.Linear(self.hidden_size, self.output_size, bias=False), ) def forward(self, x: torch.Tensor) -> torch.Tensor: return self.mlp(x) class MiMoV2AudioConfig: def __init__( self, speech_vocab_size: str | int = "1280", speech_lm_head_sizes: str | int | None = None, speech_zeroemb_idx: str | int = "1280", delay_pattern: str = "0-1-2-3-4-5-6-7-7-7-7-7-7-7-7-7-7-7-7-7", group_size: int = 4, audio_channels: int = 20, input_local_dim: int = 1024, input_local_layers: int = 6, input_local_attn_heads: int = 16, input_local_intermediate_size: int = 4096, input_local_rope_theta: float = 640000.0, input_local_partial_rotary_factor: float = 1.0, output_local_dim: int = 1024, output_local_layers: int = 6, output_local_attn_heads: int = 16, output_local_intermediate_size: int = 4096, output_local_rope_theta: float = 640000.0, output_local_partial_rotary_factor: float = 1.0, input_projection_layers: int = 2, output_projection_layers: int = 2, add_encoder_post_norm: bool = True, audio_config: dict = None, **kwargs, ): for key, value in kwargs.items(): setattr(self, key, value) if audio_config is not None: self._load_from_audio_config(audio_config) else: self.speech_vocab_size = speech_vocab_size self.speech_lm_head_sizes = ( speech_lm_head_sizes if speech_lm_head_sizes is not None else speech_vocab_size ) self.speech_zeroemb_idx = speech_zeroemb_idx self.delay_pattern = delay_pattern self.group_size = group_size self.audio_channels = audio_channels self.input_local_dim = input_local_dim self.input_local_layers = input_local_layers self.input_local_attn_heads = input_local_attn_heads self.input_local_intermediate_size = input_local_intermediate_size self.input_local_rope_theta = input_local_rope_theta self.input_local_partial_rotary_factor = input_local_partial_rotary_factor self.output_local_dim = output_local_dim self.output_local_layers = output_local_layers self.output_local_attn_heads = output_local_attn_heads self.output_local_intermediate_size = output_local_intermediate_size self.output_local_rope_theta = output_local_rope_theta self.output_local_partial_rotary_factor = output_local_partial_rotary_factor self.input_projection_layers = input_projection_layers self.output_projection_layers = output_projection_layers self.add_encoder_post_norm = add_encoder_post_norm self._attn_implementation_internal = "sdpa" def _load_from_audio_config(self, audio_config: dict): """Load audio parameters from audio_config dict in checkpoint. Uses naming that matches megatron2hf conversion output to minimize manual mapping. """ self.group_size = audio_config.get("group_size", 4) self.audio_channels = audio_config.get("audio_channels", 20) self.speech_vocab_size = audio_config.get("speech_vocab_size", "1280") self.speech_lm_head_sizes = audio_config.get( "speech_lm_head_sizes", self.speech_vocab_size ) self.speech_zeroemb_idx = audio_config.get("speech_zeroemb_idx", "1280") # Per-channel decode delays; len must equal audio_channels. self.delay_pattern = audio_config.get( "audio_output_delay_pattern", "0-1-2-3-4-5-6-7-7-7-7-7-7-7-7-7-7-7-7-7" ) self.input_local_dim = audio_config.get("input_local_dim", 1024) self.input_local_layers = audio_config.get("input_local_layers", 6) self.input_local_attn_heads = audio_config.get("input_local_attn_heads", 16) self.input_local_intermediate_size = audio_config.get( "input_local_intermediate_size", 4096 ) self.input_local_rope_theta = audio_config.get( "input_local_rope_theta", 640000.0 ) self.input_local_partial_rotary_factor = audio_config.get( "input_local_partial_rotary_factor", 1.0 ) self.output_local_dim = audio_config.get("output_local_dim", 1024) self.output_local_layers = audio_config.get("output_local_layers", 6) self.output_local_attn_heads = audio_config.get("output_local_attn_heads", 16) self.output_local_intermediate_size = audio_config.get( "output_local_intermediate_size", 4096 ) self.output_local_rope_theta = audio_config.get( "output_local_rope_theta", 640000.0 ) self.output_local_partial_rotary_factor = audio_config.get( "output_local_partial_rotary_factor", 1.0 ) self.input_projection_layers = audio_config.get("input_projection_layers", 2) self.output_projection_layers = audio_config.get("output_projection_layers", 2) self.add_encoder_post_norm = audio_config.get("add_encoder_post_norm", True) def _parse_maybe_list(self, value: str | int, length: int) -> list[int]: if isinstance(value, str) and "-" in value: return [int(s) for s in value.split("-")] return [int(value)] * length def parsed_speech_empty_ids(self): return self._parse_maybe_list(self.speech_zeroemb_idx, self.audio_channels) def parsed_speech_vocab_sizes(self): return self._parse_maybe_list(self.speech_vocab_size, self.audio_channels) def parsed_speech_lm_head_sizes(self): return self._parse_maybe_list(self.speech_lm_head_sizes, self.audio_channels) def parsed_delay_pattern(self): return self._parse_maybe_list(self.delay_pattern, self.audio_channels) def input_local_config(self): """Create config for input local transformer.""" config = Qwen2Config() for attr in dir(self): if not attr.startswith("_") and hasattr(config, attr): setattr(config, attr, getattr(self, attr)) config.hidden_size = self.input_local_dim config.num_hidden_layers = self.input_local_layers config.num_attention_heads = self.input_local_attn_heads config.num_key_value_heads = self.input_local_attn_heads config.head_dim = getattr( self, "input_local_head_dim", self.input_local_dim // self.input_local_attn_heads, ) config.intermediate_size = self.input_local_intermediate_size config.rope_theta = self.input_local_rope_theta config.partial_rotary_factor = self.input_local_partial_rotary_factor config._attn_implementation_internal = "sdpa" return config def output_local_config(self): """Create config for output local transformer.""" config = Qwen2Config() for attr in dir(self): if not attr.startswith("_") and hasattr(config, attr): setattr(config, attr, getattr(self, attr)) config.hidden_size = self.output_local_dim config.num_hidden_layers = self.output_local_layers config.num_attention_heads = self.output_local_attn_heads config.num_key_value_heads = self.output_local_attn_heads config.head_dim = self.output_local_dim // self.output_local_attn_heads config.intermediate_size = self.output_local_intermediate_size config.rope_theta = self.output_local_rope_theta config.partial_rotary_factor = self.output_local_partial_rotary_factor config._attn_implementation_internal = "sdpa" return config class AudioEncoderMixin: """LM model mixin that adds MiMo audio encoder components. Components are attached as top-level attributes (no ``audio_encoder.`` prefix), matching the checkpoint state_dict layout. Inner naming variations are normalized via ``AUDIO_WEIGHT_REMAP``. Hot config fields are cached as direct ``self.audio_*`` attributes at build time so helper methods can stay short and uniform (no ``self.audio_config.foo`` indirection inside hot paths). Subclasses call ``self.build_audio_encoder(audio_config)`` from their ``__init__`` after the language model is constructed; the mixin's ``get_audio_feature`` then handles audio item batching end-to-end. """ AUDIO_WEIGHT_REMAP: tuple[tuple[str, str], ...] = ( ("audio_projection", "projection"), ("speech_group_downcast", "projection"), ("audio_input_local_transformer", "input_local_transformer"), ) def build_audio_encoder(self, config) -> None: if not isinstance(config, MiMoV2AudioConfig): cfg_dict = vars(config) if hasattr(config, "__dict__") else config.__dict__ config = MiMoV2AudioConfig(**cfg_dict) self.audio_channels = config.audio_channels self.audio_group_size = config.group_size self.audio_segment_size = config.audio_segment_size self.audio_input_local_dim = config.input_local_dim self.audio_input_full_attention = config.input_full_attention self.audio_out_hidden_size = config.out_hidden_size speech_vocab_size = config._parse_maybe_list( config.speech_vocab_size, self.audio_channels ) speech_empty_ids = config._parse_maybe_list( config.speech_zeroemb_idx, self.audio_channels ) input_local_config = Qwen2Config( hidden_size=self.audio_input_local_dim, num_hidden_layers=config.input_local_layers, num_attention_heads=config.input_local_attn_heads, num_key_value_heads=config.input_local_attn_heads, intermediate_size=config.input_local_intermediate_size, attention_dropout=config.input_local_hidden_dropout, rope_theta=config.rope_theta, partial_rotary_factor=config.partial_rotary_factor, ) input_local_config.head_dim = config.input_local_head_dim self.input_local_transformer = Qwen2Model(input_local_config) if not config.add_post_norm: self.input_local_transformer.norm = nn.Identity() self.speech_embeddings = nn.ModuleList( [ nn.Embedding( speech_vocab_size[i], self.audio_input_local_dim, padding_idx=speech_empty_ids[i], ) for i in range(self.audio_channels) ] ) if config.projection_layers == 1: self.projection = nn.Linear( self.audio_input_local_dim * self.audio_group_size, self.audio_out_hidden_size, bias=False, ) elif config.projection_layers == 2: self.projection = AudioProjection( self.audio_input_local_dim * self.audio_group_size, self.audio_input_local_dim * self.audio_group_size * 4, self.audio_out_hidden_size, ) else: raise ValueError(f"Invalid projection layers: {config.projection_layers}") model_path = get_server_args().model_path if not os.path.isdir(model_path): from huggingface_hub import snapshot_download model_path = snapshot_download( model_path, allow_patterns=["audio_tokenizer/*"], ) audio_tokenizer_path = os.path.join(model_path, "audio_tokenizer") dev = torch.device(f"cuda:{torch.cuda.current_device()}") self.audio_tokenizer = self._load_mimo_audio_tokenizer( audio_tokenizer_path, dev ) @staticmethod def _load_mimo_audio_tokenizer( path: str, device: torch.device ) -> MiMoAudioTokenizer: """Load MiMoAudioTokenizer manually to avoid new-transformers compat issues.""" import json from safetensors.torch import load_file config_path = os.path.join(path, "config.json") with open(config_path) as f: config_dict = json.load(f) config = MiMoAudioTokenizer.config_class(**config_dict) model = MiMoAudioTokenizer(config) safetensors_path = os.path.join(path, "model.safetensors") bin_path = os.path.join(path, "pytorch_model.bin") if os.path.exists(safetensors_path): state_dict = load_file(safetensors_path, device="cpu") elif os.path.exists(bin_path): state_dict = torch.load(bin_path, map_location="cpu", weights_only=True) else: raise FileNotFoundError( f"No model weights found in {path} " "(expected model.safetensors or pytorch_model.bin)" ) # strict=False: upstream ckpt also carries decoder/vocoder weights # that this encoder-only MiMoAudioTokenizer doesn't materialize. model.load_state_dict(state_dict, strict=False) model = model.to(device=device, dtype=torch.bfloat16) model.eval() model.requires_grad_(False) return model def apply_input_local_transformer( self, speech_embeddings: torch.Tensor ) -> torch.Tensor: return self.input_local_transformer( inputs_embeds=speech_embeddings, return_dict=True, is_causal=not self.audio_input_full_attention, # for SDPA ).last_hidden_state # [T//group_size, group_size, input_local_dim] def apply_speech_embeddings(self, audio_codes: torch.Tensor) -> torch.Tensor: embeds = torch.zeros( (audio_codes.shape[0], self.audio_group_size, self.audio_input_local_dim), dtype=next(self.speech_embeddings[0].parameters()).dtype, device=audio_codes.device, ) for i in range(self.audio_channels): embeds.add_(self.speech_embeddings[i](audio_codes[:, :, i])) return embeds def pad_audio_codes(self, audio: torch.Tensor) -> torch.Tensor: T = audio.shape[0] audio = audio[:, : self.audio_channels] padded_T = ( (T + self.audio_group_size - 1) // self.audio_group_size * self.audio_group_size ) padded_audio = torch.cat( [ audio, torch.zeros( padded_T - T, self.audio_channels, dtype=torch.int32, device=audio.device, ) + audio[-1, :], ], dim=0, ) return padded_audio.reshape( padded_T // self.audio_group_size, self.audio_group_size, self.audio_channels, ) def get_audio_feature(self, items) -> torch.Tensor: """Compute audio features for a list of audio MultimodalDataItem. Each item.feature is either a mel tensor or a list of mel tensors (long audio split into chunks). """ all_mels = [] for item in items: f = item.feature if isinstance(f, (list, tuple)): all_mels.extend(f) else: all_mels.append(f) if not all_mels: device = next(self.projection.parameters()).device dtype = next(self.projection.parameters()).dtype return torch.empty( 0, self.audio_out_hidden_size, device=device, dtype=dtype ) device = next(self.audio_tokenizer.encoder.parameters()).device code_list = tokenize_audio_batch( all_mels, self.audio_tokenizer.encoder, segment_size=self.audio_segment_size, device=device, ) audio_codes = torch.cat([self.pad_audio_codes(c) for c in code_list], dim=0) embeds = self.apply_input_local_transformer( self.apply_speech_embeddings(audio_codes) ) return self.projection(embeds.reshape(embeds.shape[0], -1)) @classmethod def remap_audio_weight_name(cls, name: str) -> str: """Normalize inner audio weight name variations to canonical form.""" for src, dst in cls.AUDIO_WEIGHT_REMAP: if src in name: return name.replace(src, dst) return name