import math import mlx.core as mx import mlx.nn as nn from mlx_audio.codec.models.mimi.modules.conv import ( Conv1d, ConvDownsample1d, ConvTrUpsample1d, get_extra_padding_for_conv1d, ) from mlx_audio.codec.models.mimi.modules.seanet import ( SeanetConfig, SeanetDecoder, SeanetEncoder, ) from mlx_audio.codec.models.mimi.modules.transformer import ( ProjectedTransformer, TransformerConfig, ) from .config import MimiConfig def _reset_kv_cache(cache) -> None: cache.keys = None cache.values = None cache.offset = 0 if hasattr(cache, "_idx"): cache._idx = 0 def pad_for_conv1d(x: mx.array, kernel_size: int, stride: int, padding_total: int = 0): extra_padding = get_extra_padding_for_conv1d( x, ksize=kernel_size, stride=stride, padding_total=padding_total ) if extra_padding <= 0: return x return mx.pad(x, pad_width=((0, 0), (0, 0), (0, extra_padding))) class DummyQuantizer(nn.Module): def __init__(self, dimension: int, output_dimension: int): super().__init__() self.output_proj = Conv1d(dimension, output_dimension, 1, bias=False) def __call__(self, x: mx.array) -> mx.array: return self.output_proj(x) class MimiAdapter(nn.Module): def __init__( self, encoder: SeanetEncoder, decoder: SeanetDecoder, quantizer: DummyQuantizer, frame_rate: float, encoder_frame_rate: float, sample_rate: int, channels: int, encoder_transformer: ProjectedTransformer, decoder_transformer: ProjectedTransformer, ): super().__init__() self.encoder = encoder self.decoder = decoder self.encoder_transformer = encoder_transformer self.decoder_transformer = decoder_transformer self.quantizer = quantizer self.frame_rate = frame_rate self.sample_rate = sample_rate self.channels = channels self.encoder_frame_rate = encoder_frame_rate self.dimension = encoder_transformer.transformer.cfg.d_model if encoder_frame_rate != frame_rate: if encoder_frame_rate <= frame_rate: raise ValueError("Cannot upsample with conv.") downsample_stride = encoder_frame_rate / frame_rate if downsample_stride != int(downsample_stride): raise ValueError( f"Only integer strides are supported, got {downsample_stride}" ) downsample_stride = int(downsample_stride) self.downsample = ConvDownsample1d( downsample_stride, dim=self.dimension, causal=True ) self.upsample = ConvTrUpsample1d( downsample_stride, dim=self.dimension, causal=True ) else: self.downsample = None self.upsample = None self.encoder_cache = self.encoder_transformer.make_cache() self.decoder_cache = self.decoder_transformer.make_cache() @property def frame_size(self) -> int: return int(self.sample_rate / self.frame_rate) def reset_state(self): self.encoder.reset_state() self.decoder.reset_state() if self.downsample is not None: self.downsample.reset_state() if self.upsample is not None: self.upsample.reset_state() for cache in self.encoder_cache: _reset_kv_cache(cache) for cache in self.decoder_cache: _reset_kv_cache(cache) def _to_framerate(self, x: mx.array): if self.encoder_frame_rate == self.frame_rate: return x if self.downsample is None: raise ValueError("Downsample module is missing.") return self.downsample(x) def _to_encoder_framerate(self, x: mx.array): if self.encoder_frame_rate == self.frame_rate: return x if self.upsample is None: raise ValueError("Upsample module is missing.") return self.upsample(x) def _to_encoder_framerate_step(self, x: mx.array): if self.encoder_frame_rate == self.frame_rate: return x if self.upsample is None: raise ValueError("Upsample module is missing.") return self.upsample.step(x) def encode_to_latent(self, x: mx.array) -> mx.array: if x.ndim != 3: raise ValueError( "MimiAdapter.encode_to_latent expects audio of shape [B, C, T]." ) self.encoder.reset_state() for cache in self.encoder_cache: _reset_kv_cache(cache) if self.downsample is not None: self.downsample.reset_state() frame_size = self.frame_size x = pad_for_conv1d(x, frame_size, frame_size) emb = self.encoder(x) emb = self.encoder_transformer(emb, cache=self.encoder_cache)[0] return self._to_framerate(emb) def decode_from_latent(self, latent: mx.array) -> mx.array: self.decoder.reset_state() for cache in self.decoder_cache: _reset_kv_cache(cache) if self.upsample is not None: self.upsample.reset_state() emb = self._to_encoder_framerate(latent) emb = self.decoder_transformer(emb, cache=self.decoder_cache)[0] return self.decoder(emb) def decode_step(self, latent: mx.array) -> mx.array: emb = self._to_encoder_framerate_step(latent) emb = self.decoder_transformer(emb, cache=self.decoder_cache)[0] return self.decoder.step(emb) @classmethod def from_config(cls, config: MimiConfig) -> "MimiAdapter": seanet_cfg = SeanetConfig( dimension=config.seanet.dimension, channels=config.seanet.channels, causal=True, nfilters=config.seanet.n_filters, nresidual_layers=config.seanet.n_residual_layers, ratios=config.seanet.ratios, ksize=config.seanet.kernel_size, residual_ksize=config.seanet.residual_kernel_size, last_ksize=config.seanet.last_kernel_size, dilation_base=config.seanet.dilation_base, pad_mode=config.seanet.pad_mode, true_skip=True, compress=config.seanet.compress, ) encoder = SeanetEncoder(seanet_cfg) decoder = SeanetDecoder(seanet_cfg) transformer_cfg = TransformerConfig( d_model=config.transformer.d_model, num_heads=config.transformer.num_heads, num_layers=config.transformer.num_layers, causal=True, norm_first=True, bias_ff=False, bias_attn=False, layer_scale=config.transformer.layer_scale, positional_embedding="rope", use_conv_block=False, cross_attention=False, conv_kernel_size=3, use_conv_bias=False, gating=False, norm="layer_norm", context=config.transformer.context, max_period=config.transformer.max_period, max_seq_len=8192, kv_repeat=1, dim_feedforward=config.transformer.dim_feedforward, conv_layout=True, ) output_dims = list(config.transformer.output_dimensions) encoder_transformer = ProjectedTransformer( transformer_cfg, input_dim=config.transformer.input_dimension, output_dims=output_dims, ) decoder_transformer = ProjectedTransformer( transformer_cfg, input_dim=config.transformer.input_dimension, output_dims=output_dims, ) quantizer = DummyQuantizer( dimension=config.quantizer.dimension, output_dimension=config.quantizer.output_dimension, ) encoder_frame_rate = config.sample_rate / math.prod(config.seanet.ratios) return cls( encoder=encoder, decoder=decoder, quantizer=quantizer, frame_rate=config.frame_rate, encoder_frame_rate=encoder_frame_rate, sample_rate=config.sample_rate, channels=config.channels, encoder_transformer=encoder_transformer, decoder_transformer=decoder_transformer, )