# Copyright (c) 2025, Prince Canuma and contributors (https://github.com/Blaizzy/mlx-audio) import math import re from typing import Dict, List, Optional, Tuple import mlx.core as mx import mlx.nn as nn import numpy as np from mlx_lm.models.cache import KVCache from mlx_audio.codec.models.mimi.mimi import _reset_kv_cache from mlx_audio.codec.models.mimi.modules import ( ConvDownsample1d, ProjectedTransformer, SeanetConfig, SeanetEncoder, ) from mlx_audio.codec.models.mimi.modules import ( SplitResidualVectorQuantizer as MimiSplitRVQ, ) from mlx_audio.codec.models.mimi.modules import TransformerConfig from .config import ( Qwen3TTSTokenizerConfig, Qwen3TTSTokenizerDecoderConfig, Qwen3TTSTokenizerEncoderConfig, ) class CausalConv1d(nn.Module): """Causal 1D convolution with proper padding. Supports grouped convolutions via nn.Conv1d(groups=...). All data flows in NLC format [batch, time, channels]. """ def __init__( self, in_channels: int, out_channels: int, kernel_size: int, stride: int = 1, dilation: int = 1, groups: int = 1, ): super().__init__() self.stride = stride effective_kernel = (kernel_size - 1) * dilation + 1 self.padding = effective_kernel - stride self._buffer = None # Streaming state buffer self.conv = nn.Conv1d( in_channels, out_channels, kernel_size, stride=stride, padding=0, dilation=dilation, groups=groups, ) def __call__(self, x: mx.array) -> mx.array: # x: [batch, time, channels] (NLC format) if self.padding > 0: x = mx.pad(x, [(0, 0), (self.padding, 0), (0, 0)]) return self.conv(x) def step(self, x: mx.array) -> mx.array: """Incremental streaming step using internal buffer.""" if self.padding > 0: if self._buffer is not None: x = mx.concatenate([self._buffer, x], axis=1) else: x = mx.pad(x, [(0, 0), (self.padding, 0), (0, 0)]) self._buffer = x[:, -self.padding :, :] return self.conv(x) def reset_state(self): """Reset streaming buffer.""" self._buffer = None class CausalTransposeConv1d(nn.Module): """Causal transposed 1D convolution for upsampling. NLC format.""" def __init__( self, in_channels: int, out_channels: int, kernel_size: int, stride: int = 1, ): super().__init__() self.conv = nn.ConvTranspose1d( in_channels, out_channels, kernel_size, stride=stride, padding=0 ) self.trim_right = kernel_size - stride def __call__(self, x: mx.array) -> mx.array: # x: [batch, time, channels] (NLC format) x = self.conv(x) if self.trim_right > 0: x = x[:, : -self.trim_right, :] return x class SnakeBeta(nn.Module): """Snake activation with learnable alpha and beta parameters. SnakeBeta(x) = x + (1/beta) * sin^2(x * alpha) """ def __init__(self, channels: int, alpha: float = 1.0): super().__init__() self.channels = channels self.alpha = mx.zeros((channels,)) self.beta = mx.zeros((channels,)) self.eps = 1e-9 def __call__(self, x: mx.array) -> mx.array: alpha = mx.exp(self.alpha) beta = mx.exp(self.beta) return x + (1.0 / (beta + self.eps)) * mx.power(mx.sin(x * alpha), 2) class ConvNeXtBlock(nn.Module): """ConvNeXt block for feature processing.""" def __init__(self, dim: int): super().__init__() self.dwconv = CausalConv1d(dim, dim, kernel_size=7, groups=dim) self.norm = nn.LayerNorm(dim, eps=1e-6) self.pwconv1 = nn.Linear(dim, 4 * dim) self.pwconv2 = nn.Linear(4 * dim, dim) self.gamma = mx.ones((dim,)) * 1e-6 def __call__(self, x: mx.array) -> mx.array: # x: [batch, time, channels] (NLC format) residual = x x = self.dwconv(x) x = self.norm(x) x = self.pwconv1(x) x = nn.gelu(x) x = self.pwconv2(x) x = self.gamma * x return residual + x def step(self, x: mx.array) -> mx.array: residual = x x = self.dwconv.step(x) x = self.norm(x) x = self.pwconv1(x) x = nn.gelu(x) x = self.pwconv2(x) x = self.gamma * x return residual + x class DecoderRMSNorm(nn.Module): """RMS normalization for decoder.""" def __init__(self, hidden_size: int, eps: float = 1e-6): super().__init__() self.weight = mx.ones((hidden_size,)) self.eps = eps def __call__(self, x: mx.array) -> mx.array: x_float = x.astype(mx.float32) variance = mx.mean(x_float**2, axis=-1, keepdims=True) x_normed = x_float * mx.rsqrt(variance + self.eps) return (self.weight * x_normed).astype(x.dtype) class LayerScale(nn.Module): """Layer scale for residual connections.""" def __init__(self, channels: int, initial_scale: float = 0.01): super().__init__() self.scale = mx.ones((channels,)) * initial_scale def __call__(self, x: mx.array) -> mx.array: return self.scale * x class DecoderRotaryEmbedding(nn.Module): """Rotary position embedding for decoder transformer.""" def __init__( self, dim: int, max_position_embeddings: int = 8000, base: float = 10000.0 ): super().__init__() self.dim = dim self.max_position_embeddings = max_position_embeddings self.base = base inv_freq = 1.0 / (base ** (mx.arange(0, dim, 2, dtype=mx.float32) / dim)) self._inv_freq = inv_freq def __call__( self, x: mx.array, position_ids: mx.array ) -> Tuple[mx.array, mx.array]: inv_freq = self._inv_freq[None, :, None].astype(mx.float32) pos = position_ids[:, None, :].astype(mx.float32) freqs = (inv_freq * pos).transpose(0, 2, 1) emb = mx.concatenate([freqs, freqs], axis=-1) cos = mx.cos(emb).astype(x.dtype) sin = mx.sin(emb).astype(x.dtype) return cos, sin def rotate_half(x: mx.array) -> mx.array: x1 = x[..., : x.shape[-1] // 2] x2 = x[..., x.shape[-1] // 2 :] return mx.concatenate([-x2, x1], axis=-1) def apply_rotary_pos_emb( q: mx.array, k: mx.array, cos: mx.array, sin: mx.array ) -> Tuple[mx.array, mx.array]: cos = mx.expand_dims(cos, axis=1) sin = mx.expand_dims(sin, axis=1) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed class DecoderAttention(nn.Module): """Multi-head attention for decoder transformer.""" def __init__(self, config: Qwen3TTSTokenizerDecoderConfig, layer_idx: int): super().__init__() self.config = config self.layer_idx = layer_idx self.head_dim = config.head_dim self.num_heads = config.num_attention_heads self.num_kv_heads = config.num_key_value_heads self.num_kv_groups = self.num_heads // self.num_kv_heads self.scale = self.head_dim**-0.5 self.sliding_window = config.sliding_window self.q_proj = nn.Linear( config.hidden_size, self.num_heads * self.head_dim, bias=config.attention_bias, ) self.k_proj = nn.Linear( config.hidden_size, self.num_kv_heads * self.head_dim, bias=config.attention_bias, ) self.v_proj = nn.Linear( config.hidden_size, self.num_kv_heads * self.head_dim, bias=config.attention_bias, ) self.o_proj = nn.Linear( self.num_heads * self.head_dim, config.hidden_size, bias=config.attention_bias, ) def __call__( self, x: mx.array, position_embeddings: Tuple[mx.array, mx.array], mask: Optional[mx.array] = None, cache: Optional[KVCache] = None, ) -> mx.array: batch, seq_len, _ = x.shape q = ( self.q_proj(x) .reshape(batch, seq_len, self.num_heads, self.head_dim) .transpose(0, 2, 1, 3) ) k = ( self.k_proj(x) .reshape(batch, seq_len, self.num_kv_heads, self.head_dim) .transpose(0, 2, 1, 3) ) v = ( self.v_proj(x) .reshape(batch, seq_len, self.num_kv_heads, self.head_dim) .transpose(0, 2, 1, 3) ) cos, sin = position_embeddings q, k = apply_rotary_pos_emb(q, k, cos, sin) # Handle KV cache using mlx_lm's KVCache if cache is not None: k, v = cache.update_and_fetch(k, v) # Use fast scaled dot product attention with GQA support output = mx.fast.scaled_dot_product_attention( q, k, v, scale=self.scale, mask=mask ) output = output.transpose(0, 2, 1, 3).reshape(batch, seq_len, -1) return self.o_proj(output) class DecoderMLP(nn.Module): """MLP for decoder transformer.""" def __init__(self, config: Qwen3TTSTokenizerDecoderConfig): super().__init__() self.gate_proj = nn.Linear( config.hidden_size, config.intermediate_size, bias=False ) self.up_proj = nn.Linear( config.hidden_size, config.intermediate_size, bias=False ) self.down_proj = nn.Linear( config.intermediate_size, config.hidden_size, bias=False ) def __call__(self, x: mx.array) -> mx.array: return self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x)) class DecoderTransformerLayer(nn.Module): """Transformer layer for decoder.""" def __init__(self, config: Qwen3TTSTokenizerDecoderConfig, layer_idx: int): super().__init__() self.self_attn = DecoderAttention(config, layer_idx) self.mlp = DecoderMLP(config) self.input_layernorm = DecoderRMSNorm(config.hidden_size, config.rms_norm_eps) self.post_attention_layernorm = DecoderRMSNorm( config.hidden_size, config.rms_norm_eps ) self.self_attn_layer_scale = LayerScale( config.hidden_size, config.layer_scale_initial_scale ) self.mlp_layer_scale = LayerScale( config.hidden_size, config.layer_scale_initial_scale ) def __call__( self, x: mx.array, position_embeddings: Tuple[mx.array, mx.array], mask: Optional[mx.array] = None, cache: Optional[KVCache] = None, ) -> mx.array: residual = x x = self.input_layernorm(x) x = self.self_attn(x, position_embeddings, mask, cache) x = residual + self.self_attn_layer_scale(x) residual = x x = self.post_attention_layernorm(x) x = self.mlp(x) x = residual + self.mlp_layer_scale(x) return x class DecoderTransformer(nn.Module): """Transformer model for decoder.""" def __init__(self, config: Qwen3TTSTokenizerDecoderConfig): super().__init__() self.config = config self.layers = [ DecoderTransformerLayer(config, i) for i in range(config.num_hidden_layers) ] self.norm = DecoderRMSNorm(config.hidden_size, config.rms_norm_eps) self.rotary_emb = DecoderRotaryEmbedding( config.head_dim, config.max_position_embeddings, config.rope_theta ) self.input_proj = nn.Linear(config.latent_dim, config.hidden_size) self.output_proj = nn.Linear(config.hidden_size, config.latent_dim) def make_cache(self) -> List[KVCache]: """Create KV cache for all layers.""" return [KVCache() for _ in self.layers] def __call__( self, inputs_embeds: mx.array, mask: Optional[mx.array] = None, cache: Optional[List[KVCache]] = None, ) -> mx.array: batch, seq_len, _ = inputs_embeds.shape x = self.input_proj(inputs_embeds) # Position ids - use KVCache.offset for position tracking offset = cache[0].offset if cache is not None else 0 position_ids = mx.arange(offset, offset + seq_len)[None, :] position_ids = mx.broadcast_to(position_ids, (batch, seq_len)) position_embeddings = self.rotary_emb(x, position_ids) if mask is None and seq_len > 1: total_len = offset + seq_len mask = nn.MultiHeadAttention.create_additive_causal_mask(total_len) mask = mask[-seq_len:] # Only query positions need masking rows mask = mask.astype(x.dtype) for i, layer in enumerate(self.layers): layer_cache = cache[i] if cache is not None else None x = layer(x, position_embeddings, mask, layer_cache) x = self.norm(x) x = self.output_proj(x) return x class EuclideanCodebook(nn.Module): """Euclidean codebook for vector quantization. Uses nn.Embedding for the codebook to ensure proper weight loading. The embedding is computed from cluster_usage and embedding_sum during sanitization. """ def __init__(self, dim: int, codebook_size: int, eps: float = 1e-5): super().__init__() self.dim = dim self.codebook_size = codebook_size self.eps = eps # Use Embedding layer for proper weight loading self.embed = nn.Embedding(codebook_size, dim) def decode(self, codes: mx.array) -> mx.array: # codes: [batch, time] return self.embed(codes) class VectorQuantization(nn.Module): """Vector quantization layer.""" def __init__( self, dim: int, codebook_size: int, codebook_dim: Optional[int] = None, eps: float = 1e-5, ): super().__init__() codebook_dim = codebook_dim or dim requires_projection = codebook_dim != dim if requires_projection: self.project_out = nn.Linear(codebook_dim, dim) else: self.project_out = None self.codebook = EuclideanCodebook(codebook_dim, codebook_size, eps) self.codebook_size = codebook_size def decode(self, codes: mx.array) -> mx.array: # codes: [batch, time] quantized = self.codebook.decode(codes) # [batch, time, codebook_dim] if self.project_out is not None: quantized = self.project_out(quantized) quantized = mx.transpose(quantized, (0, 2, 1)) # [batch, dim, time] return quantized class ResidualVectorQuantization(nn.Module): """Residual vector quantization with multiple codebooks.""" def __init__( self, num_quantizers: int, dim: int, codebook_size: int, codebook_dim: Optional[int] = None, ): super().__init__() self.layers = [ VectorQuantization(dim, codebook_size, codebook_dim) for _ in range(num_quantizers) ] def decode(self, codes: mx.array) -> mx.array: # codes: [num_quantizers, batch, time] quantized = mx.zeros( (codes.shape[1], self.layers[0].codebook.dim, codes.shape[2]) ) for idx, layer_codes in enumerate(codes): quantized = quantized + self.layers[idx].decode(layer_codes) return quantized class ResidualVectorQuantizer(nn.Module): """Residual vector quantizer with input/output projections.""" def __init__( self, dimension: int = 128, input_dimension: Optional[int] = None, output_dimension: Optional[int] = None, n_q: int = 8, bins: int = 1024, force_projection: bool = False, ): super().__init__() self.n_q = n_q self.dimension = dimension self.input_dimension = input_dimension or dimension self.output_dimension = output_dimension or dimension self.bins = bins if self.input_dimension == self.dimension and not force_projection: self.input_proj = None else: self.input_proj = nn.Conv1d( self.input_dimension, self.dimension, 1, bias=False ) if self.output_dimension == self.dimension and not force_projection: self.output_proj = None else: self.output_proj = nn.Conv1d( self.dimension, self.output_dimension, 1, bias=False ) self.vq = ResidualVectorQuantization( num_quantizers=n_q, dim=dimension, codebook_size=bins, ) def decode(self, codes: mx.array) -> mx.array: # codes: [batch, num_quantizers, time] codes = mx.transpose(codes, (1, 0, 2)) # [num_quantizers, batch, time] quantized = self.vq.decode(codes) # [batch, dim, time] if self.output_proj is not None: # Conv1d expects [batch, time, channels], transpose for MLX quantized = mx.transpose(quantized, (0, 2, 1)) # [batch, time, dim] quantized = self.output_proj(quantized) # [batch, time, output_dim] quantized = mx.transpose(quantized, (0, 2, 1)) # [batch, output_dim, time] return quantized class SplitResidualVectorQuantizer(nn.Module): """Split RVQ with separate quantizers for semantic and acoustic.""" def __init__( self, n_q: int = 8, n_q_semantic: int = 1, dimension: int = 128, input_dimension: Optional[int] = None, output_dimension: Optional[int] = None, bins: int = 1024, ): super().__init__() self.n_q_semantic = n_q_semantic self.n_q_acoustic = n_q - n_q_semantic self.rvq_first = ResidualVectorQuantizer( dimension=dimension, input_dimension=input_dimension, output_dimension=output_dimension, n_q=n_q_semantic, bins=bins, force_projection=True, ) self.rvq_rest = ResidualVectorQuantizer( dimension=dimension, input_dimension=input_dimension, output_dimension=output_dimension, n_q=n_q - n_q_semantic, bins=bins, force_projection=True, ) def decode(self, codes: mx.array) -> mx.array: # codes: [batch, num_quantizers, time] quantized = self.rvq_first.decode(codes[:, : self.n_q_semantic]) if codes.shape[1] > self.n_q_semantic: quantized = quantized + self.rvq_rest.decode(codes[:, self.n_q_semantic :]) return quantized class DecoderResidualUnit(nn.Module): """Residual unit for decoder. PyTorch weight keys: - act1.alpha, act1.beta (SnakeBeta) - conv1.conv.weight, conv1.conv.bias (CausalConv1d) - act2.alpha, act2.beta (SnakeBeta) - conv2.conv.weight, conv2.conv.bias (CausalConv1d) """ def __init__(self, dim: int, dilation: int = 1): super().__init__() self.act1 = SnakeBeta(dim) self.conv1 = CausalConv1d(dim, dim, kernel_size=7, dilation=dilation) self.act2 = SnakeBeta(dim) self.conv2 = CausalConv1d(dim, dim, kernel_size=1) def __call__(self, x: mx.array) -> mx.array: residual = x x = self.act1(x) x = self.conv1(x) x = self.act2(x) x = self.conv2(x) return x + residual def step(self, x: mx.array) -> mx.array: """Streaming step — uses conv buffers.""" residual = x x = self.act1(x) x = self.conv1.step(x) x = self.act2(x) x = self.conv2.step(x) return x + residual class DecoderBlockUpsample(nn.Module): """Upsample layer wrapper for decoder block. PyTorch weight keys: conv.weight, conv.bias Note: Implements causal transpose conv inline to match PyTorch key structure. """ def __init__(self, in_dim: int, out_dim: int, upsample_rate: int): super().__init__() kernel_size = 2 * upsample_rate self.conv = nn.ConvTranspose1d( in_dim, out_dim, kernel_size, stride=upsample_rate, padding=0 ) # Trim from the right for causal behavior (matches Encodec/DAC/Mimi) self.trim_right = kernel_size - upsample_rate self._overflow = None # Overlap-add buffer for streaming def __call__(self, x: mx.array) -> mx.array: # x: [batch, time, channels] (NLC format) x = self.conv(x) if self.trim_right > 0: x = x[:, : -self.trim_right, :] return x def step(self, x: mx.array) -> mx.array: """Streaming step with overlap-add for ConvTranspose1d.""" y = self.conv(x) if self._overflow is not None: ov_len = self._overflow.shape[1] y = mx.concatenate( [y[:, :ov_len, :] + self._overflow, y[:, ov_len:, :]], axis=1 ) if self.trim_right > 0: self._overflow = y[:, -self.trim_right :, :] y = y[:, : -self.trim_right, :] return y def reset_state(self): """Reset overlap-add buffer.""" self._overflow = None class DecoderBlock(nn.Module): """Decoder block with upsampling. PyTorch structure - self.block is a ModuleList: - block[0]: SnakeBeta (alpha, beta) - block[1]: ConvWrapper (conv.weight, conv.bias) - block[2:]: DecoderResidualUnits (act1, conv1, act2, conv2) """ def __init__(self, config: Qwen3TTSTokenizerDecoderConfig, layer_idx: int): super().__init__() in_dim = config.decoder_dim // (2**layer_idx) out_dim = config.decoder_dim // (2 ** (layer_idx + 1)) upsample_rate = config.upsample_rates[layer_idx] # PyTorch uses self.block as a ModuleList self.block = [ SnakeBeta(in_dim), # block[0]: snake DecoderBlockUpsample(in_dim, out_dim, upsample_rate), # block[1]: upsample DecoderResidualUnit(out_dim, dilation=1), # block[2]: residual unit DecoderResidualUnit(out_dim, dilation=3), # block[3]: residual unit DecoderResidualUnit(out_dim, dilation=9), # block[4]: residual unit ] def __call__(self, x: mx.array) -> mx.array: for layer in self.block: x = layer(x) return x def step(self, x: mx.array) -> mx.array: """Streaming step — snake is stateless, upsample and residual units use buffers.""" x = self.block[0](x) # SnakeBeta (stateless) x = self.block[1].step(x) # DecoderBlockUpsample (overlap-add) for unit in self.block[2:]: # DecoderResidualUnits x = unit.step(x) return x class DecoderInitialConv(nn.Module): """Wrapper for initial decoder conv to match PyTorch weight structure. PyTorch key: decoder.decoder.0.conv.{weight,bias} NLC format throughout. """ def __init__(self, latent_dim: int, decoder_dim: int, kernel_size: int = 7): super().__init__() self.conv = nn.Conv1d(latent_dim, decoder_dim, kernel_size, padding=0) self.kernel_size = kernel_size self._buffer = None def __call__(self, x: mx.array) -> mx.array: # x: [batch, time, channels] (NLC format) x = mx.pad(x, [(0, 0), (self.kernel_size - 1, 0), (0, 0)]) return self.conv(x) def step(self, x: mx.array) -> mx.array: """Streaming step using internal buffer.""" padding = self.kernel_size - 1 if padding > 0: if self._buffer is not None: x = mx.concatenate([self._buffer, x], axis=1) else: x = mx.pad(x, [(0, 0), (padding, 0), (0, 0)]) self._buffer = x[:, -padding:, :] return self.conv(x) def reset_state(self): self._buffer = None class DecoderOutputSnake(nn.Module): """Output SnakeBeta layer - decoder.decoder[5]. PyTorch key: decoder.decoder.5.{alpha, beta} NLC format: alpha/beta broadcast on last dim. """ def __init__(self, channels: int): super().__init__() self.alpha = mx.zeros((channels,)) self.beta = mx.zeros((channels,)) self.eps = 1e-9 def __call__(self, x: mx.array) -> mx.array: alpha = mx.exp(self.alpha) beta = mx.exp(self.beta) return x + (1.0 / (beta + self.eps)) * mx.power(mx.sin(x * alpha), 2) class DecoderOutputConv(nn.Module): """Output conv layer - decoder.decoder[6]. PyTorch key: decoder.decoder.6.conv.{weight, bias} NLC format throughout. """ def __init__(self, channels: int, kernel_size: int = 7): super().__init__() self.conv = nn.Conv1d(channels, 1, kernel_size, padding=0) self.kernel_size = kernel_size self._buffer = None def __call__(self, x: mx.array) -> mx.array: # x: [batch, time, channels] (NLC format) x = mx.pad(x, [(0, 0), (self.kernel_size - 1, 0), (0, 0)]) return self.conv(x) def step(self, x: mx.array) -> mx.array: """Streaming step using internal buffer.""" padding = self.kernel_size - 1 if padding > 0: if self._buffer is not None: x = mx.concatenate([self._buffer, x], axis=1) else: x = mx.pad(x, [(0, 0), (padding, 0), (0, 0)]) self._buffer = x[:, -padding:, :] return self.conv(x) def reset_state(self): self._buffer = None class Qwen3TTSSpeechTokenizerDecoder(nn.Module): """Full decoder for speech tokenizer.""" def __init__(self, config: Qwen3TTSTokenizerDecoderConfig): super().__init__() self.config = config self.total_upsample = int( np.prod(config.upsample_rates + config.upsampling_ratios) ) # Streaming state self._transformer_cache = None # Transformer self.pre_transformer = DecoderTransformer(config) # Quantizer self.quantizer = SplitResidualVectorQuantizer( dimension=config.codebook_dim // 2, n_q=config.num_quantizers, n_q_semantic=config.num_semantic_quantizers, bins=config.codebook_size, input_dimension=config.codebook_dim, output_dimension=config.codebook_dim, ) # Pre-conv self.pre_conv = CausalConv1d( config.codebook_dim, config.latent_dim, kernel_size=3 ) # Upsampling blocks (matches PyTorch's self.upsample ModuleList) # Each upsample is a list: [CausalTransposeConv, ConvNeXtBlock] # PyTorch key: decoder.upsample.{idx}.{0|1}.{attr} self.upsample = [ [ CausalTransposeConv1d( config.latent_dim, config.latent_dim, factor, factor ), ConvNeXtBlock(config.latent_dim), ] for factor in config.upsampling_ratios ] # Main decoder (matches PyTorch's self.decoder ModuleList) # [0]: Initial conv, [1:5]: 4 DecoderBlocks, [5]: Output snake, [6]: Output conv # PyTorch key: decoder.decoder.{idx}.{attr} output_dim = config.decoder_dim // (2 ** len(config.upsample_rates)) self.decoder = [ DecoderInitialConv(config.latent_dim, config.decoder_dim, 7), # [0] *[ DecoderBlock(config, i) for i in range(len(config.upsample_rates)) ], # [1-4] DecoderOutputSnake(output_dim), # [5] DecoderOutputConv(output_dim, 7), # [6] ] def __call__(self, codes: mx.array) -> mx.array: """ Args: codes: [batch, num_quantizers, time] Returns: audio: [batch, 1, samples] """ if codes.shape[1] != self.config.num_quantizers: raise ValueError( f"Expected {self.config.num_quantizers} layers of codes, got {codes.shape[1]}" ) # Dequantize: [batch, codebook_dim, time] (NCL) hidden = self.quantizer.decode(codes) # Convert to NLC for the entire decoder pipeline hidden = mx.transpose(hidden, (0, 2, 1)) # [batch, time, codebook_dim] # Pre-conv (NLC throughout) hidden = self.pre_conv(hidden) # [batch, time, latent_dim] # Transformer (already expects NLC) hidden = self.pre_transformer(hidden) # Upsampling (NLC throughout) for upsample_layers in self.upsample: for layer in upsample_layers: hidden = layer(hidden) # Decoder (NLC throughout) wav = hidden for decoder_layer in self.decoder: wav = decoder_layer(wav) # Convert back to NCL for output wav = mx.transpose(wav, (0, 2, 1)) # [batch, 1, samples] return mx.clip(wav, -1.0, 1.0) def reset_streaming_state(self): """Reset all conv buffers, overflow state, and transformer KV cache.""" self._transformer_cache = None for _, m in self.named_modules(): if hasattr(m, "reset_state"): m.reset_state() def streaming_step(self, codes: mx.array) -> mx.array: """Incrementally decode new codes using conv buffers and transformer KV cache. Processes only the new tokens, maintaining state from previous calls. Must call reset_streaming_state() before starting a new stream. Args: codes: [batch, num_quantizers, new_time] — only the NEW codes Returns: audio: [batch, 1, new_samples] """ # Initialize transformer KV cache on first call if self._transformer_cache is None: self._transformer_cache = self.pre_transformer.make_cache() # Dequantize: [batch, codebook_dim, time] (NCL) hidden = self.quantizer.decode(codes) # Convert to NLC hidden = mx.transpose(hidden, (0, 2, 1)) # Pre-conv with buffer hidden = self.pre_conv.step(hidden) # Transformer with KV cache hidden = self.pre_transformer(hidden, cache=self._transformer_cache) # Upsampling: transpose conv (stateless) + ConvNeXtBlock (buffered) for upsample_layers in self.upsample: hidden = upsample_layers[0](hidden) # CausalTransposeConv1d hidden = upsample_layers[1].step(hidden) # ConvNeXtBlock # Decoder pipeline wav = self.decoder[0].step(hidden) # DecoderInitialConv for block in self.decoder[1:-2]: # DecoderBlocks wav = block.step(wav) wav = self.decoder[-2](wav) # DecoderOutputSnake (stateless) wav = self.decoder[-1].step(wav) # DecoderOutputConv # Convert back to NCL for output wav = mx.transpose(wav, (0, 2, 1)) return mx.clip(wav, -1.0, 1.0) def chunked_decode( self, codes: mx.array, chunk_size: int = 300, left_context_size: int = 25, ) -> mx.array: """Decode in chunks to handle long sequences.""" wavs = [] start_index = 0 while start_index < codes.shape[-1]: end_index = min(start_index + chunk_size, codes.shape[-1]) context_size = ( left_context_size if start_index - left_context_size > 0 else start_index ) codes_chunk = codes[..., start_index - context_size : end_index] wav_chunk = self(codes_chunk) wavs.append(wav_chunk[..., context_size * self.total_upsample :]) start_index = end_index return mx.concatenate(wavs, axis=-1) class Qwen3TTSSpeechTokenizerEncoder(nn.Module): """Encoder for the speech tokenizer using Mimi components. Architecture: SeanetEncoder -> ProjectedTransformer -> ConvDownsample1d -> SplitRVQ """ def __init__(self, config: Qwen3TTSTokenizerEncoderConfig): super().__init__() self.config = config self.valid_num_quantizers = 16 # Only first 16 quantizers are used for ICL # Build SeanetConfig from encoder config seanet_cfg = SeanetConfig( dimension=config.hidden_size, channels=config.audio_channels, causal=config.use_causal_conv, nfilters=config.num_filters, nresidual_layers=config.num_residual_layers, ratios=config.upsampling_ratios, ksize=config.kernel_size, residual_ksize=config.residual_kernel_size, last_ksize=config.last_kernel_size, dilation_base=config.dilation_growth_rate, pad_mode="constant", true_skip=not config.use_conv_shortcut, compress=config.compress, ) self.encoder = SeanetEncoder(seanet_cfg) # Build TransformerConfig transformer_cfg = TransformerConfig( d_model=config.hidden_size, num_heads=config.num_attention_heads, num_layers=config.num_hidden_layers, causal=config.use_causal_conv, norm_first=True, bias_ff=False, bias_attn=False, layer_scale=config.layer_scale_initial_scale, positional_embedding="rope", use_conv_block=False, cross_attention=False, conv_kernel_size=3, use_conv_bias=True, gating=False, norm="layer_norm", context=config.sliding_window, max_period=int(config.rope_theta), max_seq_len=config.max_position_embeddings, kv_repeat=config.num_attention_heads // config.num_key_value_heads, dim_feedforward=config.intermediate_size, conv_layout=True, rope_traditional=False, ) self.encoder_transformer = ProjectedTransformer( transformer_cfg, input_dim=config.hidden_size, output_dims=[config.hidden_size], ) # ConvDownsample1d: stride = encoder_frame_rate / frame_rate encoder_frame_rate = config.sampling_rate / math.prod(config.upsampling_ratios) downsample_stride = int(encoder_frame_rate / config.frame_rate) self.downsample = ConvDownsample1d( stride=downsample_stride, dim=config.hidden_size, causal=config.use_causal_conv, ) # SplitResidualVectorQuantizer self.quantizer = MimiSplitRVQ( dim=config.codebook_dim, input_dim=config.hidden_size, output_dim=config.hidden_size, nq=config.num_quantizers, bins=config.codebook_size, ) self.encoder_cache = self.encoder_transformer.make_cache() def encode(self, audio: mx.array) -> mx.array: """Encode audio waveform to codes. Args: audio: [batch, 1, samples] audio waveform Returns: codes: [batch, num_quantizers, time] """ self.encoder.reset_state() for c in self.encoder_cache: _reset_kv_cache(c) xs = self.encoder(audio) # Create causal attention mask (the model was trained with causal attention) seq_len = xs.shape[-1] # NCL format, time is last dim mask = mx.full((seq_len, seq_len), -mx.inf, dtype=xs.dtype) mask = mx.triu(mask, k=1) # Upper triangle = -inf, lower triangle + diag = 0 mask = mask[None, None, :, :] # [1, 1, seq_len, seq_len] xs = self.encoder_transformer(xs, cache=self.encoder_cache, mask=mask)[0] xs = self.downsample(xs) codes = self.quantizer.encode(xs) return codes[:, : self.valid_num_quantizers, :] class Qwen3TTSSpeechTokenizer(nn.Module): """Full speech tokenizer model.""" def __init__(self, config: Qwen3TTSTokenizerConfig): super().__init__() self.config = config self.encoder_valid_num_quantizers = config.encoder_valid_num_quantizers self.input_sample_rate = config.input_sample_rate self.output_sample_rate = config.output_sample_rate self.decode_upsample_rate = config.decode_upsample_rate self.encode_downsample_rate = config.encode_downsample_rate # Decoder self.decoder = Qwen3TTSSpeechTokenizerDecoder(config.decoder_config) # Encoder (for ICL voice cloning) if config.encoder_config is not None: self.encoder_model = Qwen3TTSSpeechTokenizerEncoder(config.encoder_config) else: self.encoder_model = None def encode(self, audio: mx.array) -> mx.array: """Encode audio waveform to codes. Args: audio: [batch, 1, samples] audio waveform Returns: codes: [batch, num_quantizers, time] """ if self.encoder_model is None: raise ValueError("Encoder not available for this speech tokenizer") return self.encoder_model.encode(audio) @property def has_encoder(self) -> bool: return self.encoder_model is not None def decode(self, audio_codes: mx.array) -> mx.array: """ Decode audio codes to waveform. Args: audio_codes: [batch, time, num_quantizers] Returns: audio: [batch, samples] """ # Transpose to [batch, num_quantizers, time] codes = mx.transpose(audio_codes, (0, 2, 1)) wav = self.decoder.chunked_decode(codes).squeeze(1) # Trim based on valid lengths audio_lengths = (audio_codes[..., 0] > 0).sum( axis=1 ) * self.decode_upsample_rate return wav, audio_lengths def batch_decode( self, codes_list: list[mx.array] ) -> tuple[list[mx.array], list[int]]: """ Decode a list of variable-length code sequences in a single batched pass. Pads codes to the longest sequence, decodes the full batch through the vocoder, then trims each output to its valid length. Args: codes_list: List of code arrays, each [1, time_i, num_quantizers] (or [time_i, num_quantizers]). Sequences may have different lengths. Returns: audios: List of [samples_i] arrays, one per input sequence. lengths: List of valid sample counts per sequence. """ if not codes_list: return [], [] # Normalize to 3-D: [1, time, Q] normed = [] for c in codes_list: if c.ndim == 2: c = c[None] # [time, Q] -> [1, time, Q] normed.append(c) seq_lens = [c.shape[1] for c in normed] max_len = max(seq_lens) num_q = normed[0].shape[2] # Pad with 0 (silence code) to equalize lengths padded = [] for c in normed: pad_len = max_len - c.shape[1] if pad_len > 0: padding = mx.zeros((1, pad_len, num_q), dtype=c.dtype) c = mx.concatenate([c, padding], axis=1) padded.append(c) batch_codes = mx.concatenate(padded, axis=0) # [batch, max_len, Q] # Transpose to [batch, Q, max_len] and decode codes_t = mx.transpose(batch_codes, (0, 2, 1)) wav_batch = self.decoder.chunked_decode(codes_t).squeeze(1) # [batch, samples] mx.eval(wav_batch) # Compute per-sequence valid lengths and trim audio_lengths = [int(sl) * self.decode_upsample_rate for sl in seq_lens] audios = [] for b, valid_samples in enumerate(audio_lengths): audio = wav_batch[b] if valid_samples > 0 and valid_samples < audio.shape[0]: audio = audio[:valid_samples] audios.append(audio) del wav_batch, batch_codes, codes_t mx.clear_cache() return audios, audio_lengths def streaming_decode(self, audio_codes: mx.array, chunk_tokens: int = 100): """ Decode audio codes to waveform in a streaming fashion. Yields audio chunks as they're decoded, reducing peak memory usage. Args: audio_codes: [batch, time, num_quantizers] chunk_tokens: Number of tokens to decode per chunk Yields: audio_chunk: [batch, samples] decoded audio for each chunk """ # Transpose to [batch, num_quantizers, time] codes = mx.transpose(audio_codes, (0, 2, 1)) total_tokens = codes.shape[-1] left_context_size = 25 start_index = 0 while start_index < total_tokens: end_index = min(start_index + chunk_tokens, total_tokens) context_size = ( left_context_size if start_index - left_context_size > 0 else start_index ) codes_chunk = codes[..., start_index - context_size : end_index] wav_chunk = self.decoder(codes_chunk) wav_chunk = wav_chunk[..., context_size * self.decode_upsample_rate :] wav_chunk = wav_chunk.squeeze(1) # Evaluate immediately to free computation graph mx.eval(wav_chunk) yield wav_chunk # Clear cache after each chunk mx.clear_cache() start_index = end_index @staticmethod def sanitize(weights: Dict[str, mx.array]) -> Dict[str, mx.array]: """Sanitize weights from PyTorch to MLX format.""" from .qwen3_tts import check_array_shape_qwen3 sanitized = {} # Collect decoder codebook weights to compute embeddings codebook_data = {} # Collect encoder transformer q/k/v for in_proj concatenation encoder_transformer_qkv: Dict[int, Dict[str, mx.array]] = {} # Collect encoder codebook data (cluster_usage + embedding_sum) encoder_codebook_data: Dict[str, Dict[str, mx.array]] = {} # SeanetEncoder layer mapping: # N=0: init_conv, N=3,6,9,12: downsample, N=14: final_conv # N=1,4,7,10: residual blocks (with .block.1 and .block.3 sub-keys) seanet_conv_map = { 0: "encoder_model.encoder.init_conv1d", 3: "encoder_model.encoder.layers.0.downsample", 6: "encoder_model.encoder.layers.1.downsample", 9: "encoder_model.encoder.layers.2.downsample", 12: "encoder_model.encoder.layers.3.downsample", 14: "encoder_model.encoder.final_conv1d", } # Residual block layers: N -> encoder layer index seanet_residual_map = {1: 0, 4: 1, 7: 2, 10: 3} # Block index -> residual conv index seanet_block_map = {1: 0, 3: 1} for k, v in weights.items(): if k.startswith("encoder."): # --- Handle encoder weights --- # SeanetEncoder convolutions if k.startswith("encoder.encoder.layers."): parts = k.split(".") n = int(parts[3]) # encoder.encoder.layers.{N}... if "block" in k: # Residual block: encoder.encoder.layers.{N}.block.{B}.conv.{w/b} if n not in seanet_residual_map: continue layer_idx = seanet_residual_map[n] block_idx = int(parts[5]) # .block.{B} if block_idx not in seanet_block_map: continue conv_idx = seanet_block_map[block_idx] base_path = f"encoder_model.encoder.layers.{layer_idx}.residuals.0.block.{conv_idx}" suffix = ".".join(parts[6:]) # conv.weight or conv.bias else: # Direct conv: encoder.encoder.layers.{N}.conv.{w/b} if n not in seanet_conv_map: continue base_path = seanet_conv_map[n] suffix = ".".join(parts[4:]) # conv.weight or conv.bias new_key = f"{base_path}.conv.{suffix}" # Transpose conv weights: PyTorch [out, in, kernel] -> MLX [out, kernel, in] if "weight" in suffix and len(v.shape) == 3: v = v.swapaxes(-1, -2) sanitized[new_key] = v # Encoder transformer layers elif k.startswith("encoder.encoder_transformer.layers."): # encoder.encoder_transformer.layers.{i}.self_attn.q_proj.weight parts = k.split(".") layer_idx = int( parts[3] ) # encoder.encoder_transformer.layers.{i}... rest = ".".join(parts[4:]) # Collect q/k/v for in_proj concatenation if "self_attn.q_proj.weight" in rest: if layer_idx not in encoder_transformer_qkv: encoder_transformer_qkv[layer_idx] = {} encoder_transformer_qkv[layer_idx]["q"] = v elif "self_attn.k_proj.weight" in rest: if layer_idx not in encoder_transformer_qkv: encoder_transformer_qkv[layer_idx] = {} encoder_transformer_qkv[layer_idx]["k"] = v elif "self_attn.v_proj.weight" in rest: if layer_idx not in encoder_transformer_qkv: encoder_transformer_qkv[layer_idx] = {} encoder_transformer_qkv[layer_idx]["v"] = v elif "self_attn.o_proj.weight" in rest: new_key = f"encoder_model.encoder_transformer.transformer.layers.{layer_idx}.self_attn.out_proj.weight" sanitized[new_key] = v elif "mlp.fc1.weight" in rest: new_key = f"encoder_model.encoder_transformer.transformer.layers.{layer_idx}.gating.linear1.weight" sanitized[new_key] = v elif "mlp.fc2.weight" in rest: new_key = f"encoder_model.encoder_transformer.transformer.layers.{layer_idx}.gating.linear2.weight" sanitized[new_key] = v elif "input_layernorm.weight" in rest: new_key = f"encoder_model.encoder_transformer.transformer.layers.{layer_idx}.norm1.weight" sanitized[new_key] = v elif "input_layernorm.bias" in rest: new_key = f"encoder_model.encoder_transformer.transformer.layers.{layer_idx}.norm1.bias" sanitized[new_key] = v elif "post_attention_layernorm.weight" in rest: new_key = f"encoder_model.encoder_transformer.transformer.layers.{layer_idx}.norm2.weight" sanitized[new_key] = v elif "post_attention_layernorm.bias" in rest: new_key = f"encoder_model.encoder_transformer.transformer.layers.{layer_idx}.norm2.bias" sanitized[new_key] = v elif "self_attn_layer_scale.scale" in rest: new_key = f"encoder_model.encoder_transformer.transformer.layers.{layer_idx}.layer_scale_1.scale" sanitized[new_key] = v elif "mlp_layer_scale.scale" in rest: new_key = f"encoder_model.encoder_transformer.transformer.layers.{layer_idx}.layer_scale_2.scale" sanitized[new_key] = v # Encoder downsample conv elif k.startswith("encoder.downsample."): suffix = k.replace("encoder.downsample.", "") new_key = f"encoder_model.downsample.conv.conv.{suffix}" if "weight" in suffix and len(v.shape) == 3: v = v.swapaxes(-1, -2) sanitized[new_key] = v # Encoder quantizer elif k.startswith("encoder.quantizer."): rest = k.replace("encoder.quantizer.", "") # Codebook data (cluster_usage / embed_sum) # HF format: .layers.{i}.codebook.{cluster_usage,embed_sum,initialized} if ( ".codebook.cluster_usage" in rest or ".codebook.embed_sum" in rest ): # Extract base path for grouping base = rest.rsplit(".codebook.", 1)[0] if base not in encoder_codebook_data: encoder_codebook_data[base] = {} if "cluster_usage" in rest: encoder_codebook_data[base]["cluster_usage"] = v elif "embed_sum" in rest: encoder_codebook_data[base]["embedding_sum"] = v elif ".codebook.initialized" in rest: pass # Skip initialized flag # Input/output projections elif "input_proj.weight" in rest or "output_proj.weight" in rest: if "semantic_residual_vector_quantizer" in rest: proj_type = ( "input_proj" if "input_proj" in rest else "output_proj" ) new_key = ( f"encoder_model.quantizer.rvq_first.{proj_type}.weight" ) else: proj_type = ( "input_proj" if "input_proj" in rest else "output_proj" ) new_key = ( f"encoder_model.quantizer.rvq_rest.{proj_type}.weight" ) # Conv1d weight: PyTorch [out, in, kernel] -> MLX [out, kernel, in] if len(v.shape) == 3: v = v.swapaxes(-1, -2) sanitized[new_key] = v else: # --- Handle decoder weights (existing logic) --- # Collect codebook cluster_usage and embedding_sum for later processing if "_codebook.cluster_usage" in k or "_codebook.embedding_sum" in k: base_path = k.rsplit("._codebook.", 1)[0] if base_path not in codebook_data: codebook_data[base_path] = {} if "cluster_usage" in k: codebook_data[base_path]["cluster_usage"] = v else: codebook_data[base_path]["embedding_sum"] = v continue new_key = k # Identify ConvTranspose1d weights (for upsampling layers) is_transpose_conv = ("upsample" in k and ".0.conv.weight" in k) or ( "decoder.decoder" in k and "block.1.conv.weight" in k ) if is_transpose_conv and len(v.shape) == 3: v = v if check_array_shape_qwen3(v) else mx.transpose(v, (1, 2, 0)) elif "conv.weight" in k and len(v.shape) == 3: v = v if check_array_shape_qwen3(v) else mx.transpose(v, (0, 2, 1)) elif "_proj.weight" in k and len(v.shape) == 3: v = v if check_array_shape_qwen3(v) else mx.transpose(v, (0, 2, 1)) sanitized[new_key] = v # Process encoder transformer q/k/v into combined in_proj weights for layer_idx, qkv in encoder_transformer_qkv.items(): if "q" in qkv and "k" in qkv and "v" in qkv: in_proj_weight = mx.concatenate([qkv["q"], qkv["k"], qkv["v"]], axis=0) new_key = f"encoder_model.encoder_transformer.transformer.layers.{layer_idx}.self_attn.in_proj.weight" sanitized[new_key] = in_proj_weight # Process encoder codebook data into embedding_sum and cluster_usage for base_path, data in encoder_codebook_data.items(): if "cluster_usage" in data and "embedding_sum" in data: # Map HF quantizer paths to Mimi paths if "semantic_residual_vector_quantizer" in base_path: m = re.search(r"layers\.(\d+)", base_path) if m: layer_idx = int(m.group(1)) prefix = f"encoder_model.quantizer.rvq_first.vq.layers.{layer_idx}.codebook" sanitized[f"{prefix}.embedding_sum"] = data["embedding_sum"] sanitized[f"{prefix}.cluster_usage"] = data["cluster_usage"] elif "acoustic_residual_vector_quantizer" in base_path: m = re.search(r"layers\.(\d+)", base_path) if m: layer_idx = int(m.group(1)) prefix = f"encoder_model.quantizer.rvq_rest.vq.layers.{layer_idx}.codebook" sanitized[f"{prefix}.embedding_sum"] = data["embedding_sum"] sanitized[f"{prefix}.cluster_usage"] = data["cluster_usage"] # Compute decoder embeddings from cluster_usage and embedding_sum eps = 1e-5 for base_path, data in codebook_data.items(): if "cluster_usage" in data and "embedding_sum" in data: cluster_usage = data["cluster_usage"] embedding_sum = data["embedding_sum"] embedding = embedding_sum / mx.clip(cluster_usage[:, None], eps, None) new_key = f"{base_path}.codebook.embed.weight" sanitized[new_key] = embedding return sanitized