# Copyright (c) Kyutai, all rights reserved. # This source code is licensed under the license found in the # LICENSE file in the root directory of this source tree. import math from dataclasses import dataclass import mlx.core as mx import mlx.nn as nn from huggingface_hub import hf_hub_download from .modules import ( ConvDownsample1d, ConvTranspose1d, ConvTrUpsample1d, EuclideanCodebook, ProjectedTransformer, SeanetConfig, SeanetDecoder, SeanetEncoder, SplitResidualVectorQuantizer, TransformerConfig, ) def _reset_kv_cache(cache) -> None: cache.keys = None cache.values = None cache.offset = 0 if hasattr(cache, "_idx"): cache._idx = 0 @dataclass class MimiConfig: channels: int sample_rate: float frame_rate: float renormalize: bool seanet: SeanetConfig transformer: TransformerConfig quantizer_nq: int quantizer_bins: int quantizer_dim: int def mimi_202407(num_codebooks: int) -> MimiConfig: seanet = SeanetConfig( dimension=512, channels=1, causal=True, nfilters=64, nresidual_layers=1, ratios=[8, 6, 5, 4], ksize=7, residual_ksize=3, last_ksize=3, dilation_base=2, pad_mode="constant", true_skip=True, compress=2, ) transformer = TransformerConfig( d_model=seanet.dimension, num_heads=8, num_layers=8, causal=True, norm_first=True, bias_ff=False, bias_attn=False, layer_scale=0.01, positional_embedding="rope", use_conv_bias=True, gating=False, norm="layer_norm", context=250, max_period=10000, max_seq_len=8192, kv_repeat=1, dim_feedforward=2048, conv_layout=True, use_conv_block=False, cross_attention=False, conv_kernel_size=3, ) return MimiConfig( channels=1, sample_rate=24000, frame_rate=12.5, renormalize=True, seanet=seanet, transformer=transformer, quantizer_nq=num_codebooks, quantizer_bins=2048, quantizer_dim=256, ) class Mimi(nn.Module): def __init__(self, cfg: MimiConfig): super().__init__() dim = cfg.seanet.dimension self.cfg = cfg encoder_frame_rate = cfg.sample_rate / math.prod(cfg.seanet.ratios) downsample_stride = int(encoder_frame_rate / cfg.frame_rate) self.encoder = SeanetEncoder(cfg.seanet) self.decoder = SeanetDecoder(cfg.seanet) self.quantizer = SplitResidualVectorQuantizer( dim=cfg.quantizer_dim, input_dim=dim, output_dim=dim, nq=cfg.quantizer_nq, bins=cfg.quantizer_bins, ) self.encoder_transformer = ProjectedTransformer( cfg.transformer, input_dim=dim, output_dims=[dim], ) self.decoder_transformer = ProjectedTransformer( cfg.transformer, input_dim=dim, output_dims=[dim], ) self.downsample = ConvDownsample1d( stride=downsample_stride, dim=dim, causal=True, ) self.upsample = ConvTrUpsample1d( stride=downsample_stride, dim=dim, causal=True, ) self.encoder_cache = self.encoder_transformer.make_cache() self.decoder_cache = self.decoder_transformer.make_cache() def reset_state(self): self.encoder.reset_state() self.decoder.reset_state() for c in self.decoder_cache: _reset_kv_cache(c) for c in self.encoder_cache: _reset_kv_cache(c) def encode(self, xs: mx.array) -> mx.array: self.encoder.reset_state() for c in self.encoder_cache: _reset_kv_cache(c) xs = self.encoder(xs) xs = self.encoder_transformer(xs, cache=self.encoder_cache)[0] xs = self.downsample(xs) return self.quantizer.encode(xs) def decode(self, xs: mx.array) -> mx.array: self.decoder.reset_state() for c in self.decoder_cache: _reset_kv_cache(c) xs = self.quantizer.decode(xs) xs = self.upsample(xs) xs = self.decoder_transformer(xs, cache=self.decoder_cache)[0] return self.decoder(xs) def encode_step(self, xs: mx.array) -> mx.array: xs = self.encoder.step(xs) xs = self.encoder_transformer(xs, cache=self.encoder_cache)[0] xs = self.downsample.step(xs) xs = self.quantizer.encode(xs) return xs def decode_step(self, xs: mx.array) -> mx.array: xs = self.quantizer.decode(xs) xs = self.upsample.step(xs) xs = self.decoder_transformer(xs, cache=self.decoder_cache)[0] xs = self.decoder.step(xs) return xs def warmup(self): pcm = mx.zeros((1, 1, 1920 * 4)) codes = self.encode(pcm) pcm_out = self.decode(codes) mx.eval(pcm_out) @property def frame_rate(self) -> float: return self.cfg.frame_rate @property def sample_rate(self) -> float: return self.cfg.sample_rate def load_pytorch_weights( self, file: str, strict: bool = True, ) -> nn.Module: weights = [] for k, v in mx.load(file).items(): v: mx.array = v k: str = ".".join([s.removeprefix("_") for s in k.split(".")]) if k.startswith("encoder.model."): k = k.replace("encoder.model.", "encoder.") if k.startswith("decoder.model."): k = k.replace("decoder.model.", "decoder.") if k.endswith(".in_proj_weight"): k = k.replace(".in_proj_weight", ".in_proj.weight") if k.endswith(".linear1.weight"): k = k.replace(".linear1.weight", ".gating.linear1.weight") if k.endswith(".linear2.weight"): k = k.replace(".linear2.weight", ".gating.linear2.weight") # Awfully hardcoded matching between the pytorch layers and their mlx equivalent :( for layerIdx, decoderIdx in enumerate([2, 5, 8, 11]): k = k.replace( f"decoder.{decoderIdx}.", f"decoder.layers.{layerIdx}.upsample." ) k = k.replace( f"decoder.{decoderIdx + 1}.", f"decoder.layers.{layerIdx}.residuals.0.", ) for layerIdx, encoderIdx in enumerate([1, 4, 7, 10]): k = k.replace( f"encoder.{encoderIdx}.", f"encoder.layers.{layerIdx}.residuals.0." ) k = k.replace( f"encoder.{encoderIdx + 2}.", f"encoder.layers.{layerIdx}.downsample.", ) k = k.replace("decoder.0.", "decoder.init_conv1d.") k = k.replace("decoder.14.", "decoder.final_conv1d.") k = k.replace("encoder.0.", "encoder.init_conv1d.") k = k.replace("encoder.14.", "encoder.final_conv1d.") k = k.replace(".block.1.", ".block.0.") k = k.replace(".block.3.", ".block.1.") # PyTorch layout for conv weights is outC, inC, kSize, for MLX it's outC, kSize, inC if ( k.endswith(".conv.weight") or k.endswith(".output_proj.weight") or k.endswith(".input_proj.weight") ): v = v.swapaxes(-1, -2) # PyTorch layout for conv-transposed weights is (inC, outC/groups, kSize). # MLX expects (outC, kSize, inC/groups). Depthwise convtr needs a # different transpose because outC/groups == 1. if k.endswith(".convtr.weight"): if v.ndim == 3 and v.shape[1] == 1: v = v.transpose(0, 2, 1) else: v = v.transpose(1, 2, 0) weights.append((k, v)) m = self.load_weights(weights, strict=strict) def _filter_fn(module, name, _): if isinstance(module, EuclideanCodebook) and name == "initialized": module.update_in_place() if isinstance(module, ConvTranspose1d) and name == "weight": module.update_in_place() return True m.filter_and_map(_filter_fn) return m @classmethod def from_pretrained( cls, repo_id: str, filename: str = "tokenizer-e351c8d8-checkpoint125.safetensors", ) -> nn.Module: cfg = mimi_202407(32) model = cls(cfg) model_file = hf_hub_download(repo_id, filename) model.load_pytorch_weights(model_file, strict=True) return model class MimiStreamingDecoder: """Incremental decoder wrapper for the Mimi codec. This helper keeps the internal state of the Mimi model across calls and decodes audio tokens frame by frame using ``decode_step``. """ def __init__(self, mimi: "Mimi") -> None: # noqa: F821 - Mimi defined below self._mimi = mimi self.reset() def reset(self) -> None: """Reset the underlying codec state.""" self._mimi.decoder.reset_state() self._mimi.upsample.reset_state() for c in self._mimi.decoder_cache: _reset_kv_cache(c) def decode_frames(self, tokens: mx.array) -> mx.array: """Decode a sequence of audio tokens incrementally. Parameters ---------- tokens: Array of shape ``(B, C, T)`` or ``(C, T)`` containing the audio tokens to decode. ``B`` is the batch dimension, ``C`` is the number of codebooks and ``T`` the number of frames. Returns ------- mx.array The decoded waveform for the provided frames. """ if tokens.ndim == 2: tokens = mx.expand_dims(tokens, 0) pcm = [] for t in range(tokens.shape[-1]): step_tokens = tokens[:, :, t : t + 1] pcm.append(self._mimi.decode_step(step_tokens)) return mx.concat(pcm, axis=-1)