nexaai 1.0.29__cp310-cp310-macosx_14_0_universal2.whl

nexaai/mlx_backend/mlx_audio/codec/models/mimi/modules/kv_cache.py

@@ -0,0 +1,199 @@
+# Most of the code below comes from:
+# https://github.com/ml-explore/mlx-examples/blob/6c2369e4b97f49fb5906ec46033497b39931b25d/llms/mlx_lm/models/base.py#L1
+# Copyright © 2023-2024 Apple Inc.
+
+from __future__ import annotations
+
+import inspect
+from dataclasses import dataclass
+from typing import Any
+
+import mlx.core as mx
+
+
+class KVCache:
+
+    def __init__(self, head_dim, n_kv_heads):
+        self.n_kv_heads = n_kv_heads
+        if isinstance(head_dim, int):
+            self.k_head_dim = self.v_head_dim = head_dim
+        elif isinstance(head_dim, tuple) and len(head_dim) == 2:
+            self.k_head_dim, self.v_head_dim = head_dim
+        else:
+            raise ValueError("head_dim must be an int or a tuple of two ints")
+        self.keys = None
+        self.values = None
+        self.offset = 0
+        self.step = 256
+
+    def update_and_fetch(self, keys, values) -> tuple[mx.array, mx.array]:
+        prev = self.offset
+        if self.keys is None or (prev + keys.shape[2]) > self.keys.shape[2]:
+            B = keys.shape[0]
+            n_steps = (self.step + keys.shape[2] - 1) // self.step
+            k_shape = (B, self.n_kv_heads, n_steps * self.step, self.k_head_dim)
+            v_shape = (B, self.n_kv_heads, n_steps * self.step, self.v_head_dim)
+            new_k = mx.zeros(k_shape, keys.dtype)
+            new_v = mx.zeros(v_shape, values.dtype)
+            if self.keys is not None:
+                assert self.values is not None
+                if prev % self.step != 0:
+                    self.keys = self.keys[..., :prev, :]
+                    self.values = self.values[..., :prev, :]
+                self.keys = mx.concatenate([self.keys, new_k], axis=2)
+                self.values = mx.concatenate([self.values, new_v], axis=2)
+            else:
+                self.keys, self.values = new_k, new_v
+
+        self.offset += keys.shape[2]
+        self.keys[..., prev : self.offset, :] = keys
+        assert self.values is not None
+        self.values[..., prev : self.offset, :] = values
+        return self.keys[..., : self.offset, :], self.values[..., : self.offset, :]
+
+    def reset(self):
+        self.offset = 0
+        self.keys = None
+        self.values = None
+
+    @property
+    def state(self):
+        return self.keys, self.values
+
+
+class RotatingKVCache:
+
+    def __init__(self, head_dim, n_kv_heads, max_size, keep=0, step=256):
+        self.n_kv_heads = n_kv_heads
+        if isinstance(head_dim, int):
+            self.k_head_dim = self.v_head_dim = head_dim
+        elif isinstance(head_dim, tuple) and len(head_dim) == 2:
+            self.k_head_dim, self.v_head_dim = head_dim
+        else:
+            raise ValueError("head_dim must be an int or a tuple of two ints")
+        self.keep = keep
+        self.keys = None
+        self.values = None
+        self.offset = 0
+        self.max_size = max_size
+        self.step = step
+        self._idx = 0
+
+    def _trim(self, trim_size, v, append=None):
+        to_cat = []
+        if trim_size > 0:
+            to_cat = [v[..., : self.keep, :], v[..., trim_size + self.keep :, :]]
+        else:
+            to_cat = [v]
+        if append is not None:
+            to_cat.append(append)
+        return mx.concatenate(to_cat, axis=2)
+
+    def update_and_fetch(self, keys, values) -> tuple[mx.array, mx.array]:
+        prev = self.offset
+        B, _, S = keys.shape[:3]
+
+        # Prefill mode
+        if S > 1:
+            if self.keys is None:
+                self.keys = keys
+                self.values = values
+            else:
+                # The largest size is self.max_size + S - 1 to ensure
+                # every token gets at least self.max_size context
+                trim_size = self.keys.shape[2] - self.max_size + 1
+                self.keys = self._trim(trim_size, self.keys, keys)
+                self.values = self._trim(trim_size, self.values, values)
+            self.offset += S
+            self._idx = self.keys.shape[2]
+            return self.keys, self.values
+
+        # Generation mode
+        # May not have hit the max size yet, so potentially
+        # keep growing the cache
+        if self.keys is None or (
+            prev >= self.keys.shape[2] and self.keys.shape[2] < self.max_size
+        ):
+            new_size = min(self.step, self.max_size - prev)
+            k_shape = (B, self.n_kv_heads, new_size, self.k_head_dim)
+            v_shape = (B, self.n_kv_heads, new_size, self.v_head_dim)
+            new_k = mx.zeros(k_shape, keys.dtype)
+            new_v = mx.zeros(v_shape, values.dtype)
+            if self.keys is not None:
+                assert self.values is not None
+                self.keys = mx.concatenate([self.keys, new_k], axis=2)
+                self.values = mx.concatenate([self.values, new_v], axis=2)
+            else:
+                self.keys, self.values = new_k, new_v
+            self._idx = prev
+
+        # Trim if needed
+        trim_size = self.keys.shape[2] - self.max_size
+        if trim_size > 0:
+            self.keys = self._trim(trim_size, self.keys)
+            self.values = self._trim(trim_size, self.values)
+            self._idx = self.max_size
+
+        # Rotate
+        if self._idx == self.max_size:
+            self._idx = self.keep
+
+        # Assign
+        self.keys[..., self._idx : self._idx + 1, :] = keys
+        assert self.values is not None
+        self.values[..., self._idx : self._idx + 1, :] = values
+        self.offset += 1
+        self._idx += 1
+
+        # If the buffer is not full, slice off the end
+        if self.offset < self.max_size:
+            return self.keys[..., : self.offset, :], self.values[..., : self.offset, :]
+        return self.keys, self.values
+
+    def reset(self):
+        self.offset = 0
+        self._idx = 0
+        self.keys = None
+        self.values = None
+
+    @property
+    def state(self):
+        return self.keys, self.values
+
+
+@dataclass
+class BaseModelArgs:
+    @classmethod
+    def from_dict(cls, params):
+        return cls(
+            **{
+                k: v
+                for k, v in params.items()
+                if k in inspect.signature(cls).parameters
+            }
+        )
+
+
+def create_additive_causal_mask(N: int, offset: int = 0):
+    rinds = mx.arange(offset + N)
+    linds = mx.arange(offset, offset + N) if offset else rinds
+    mask = linds[:, None] < rinds[None]
+    return mask * -1e9
+
+
+def create_attention_mask(h: mx.array, cache: Any | None = None):
+    T = h.shape[1]
+    if T > 1:
+        if cache is not None and cache[0] is not None:
+            c = cache[0]
+            if isinstance(c, RotatingKVCache):
+                offset = min(c.max_size - 1, c.offset)
+            else:
+                offset = c.offset
+        else:
+            offset = 0
+        mask = create_additive_causal_mask(T, offset)
+        mask = mask.astype(h.dtype)
+    else:
+        mask = None
+    return mask

nexaai/mlx_backend/mlx_audio/codec/models/mimi/modules/quantization.py

@@ -0,0 +1,179 @@
+# 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.
+
+from __future__ import annotations
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .conv import Conv1d
+
+
+class EuclideanCodebook(nn.Module):
+    def __init__(self, dim: int, codebook_size: int):
+        super().__init__()
+        self._epsilon = 1e-5
+        self._dim = dim
+        self.initialized = mx.zeros([1], dtype=mx.float32)
+        self.embedding_sum = mx.zeros([codebook_size, dim], dtype=mx.float32)
+        self.cluster_usage = mx.zeros([codebook_size], dtype=mx.float32)
+        cluster_usage = mx.maximum(self.cluster_usage, self._epsilon)[:, None]
+        self._embedding = self.embedding_sum / cluster_usage
+        self._c2 = self._embedding.square().sum(axis=-1) / 2
+
+    def update(self, parameters: dict) -> nn.Module:
+        super().update(parameters)
+        cluster_usage = mx.maximum(self.cluster_usage, self._epsilon)[:, None]
+        self._embedding = self.embedding_sum / cluster_usage
+        self._c2 = self._embedding.square().sum(axis=-1) / 2
+        return self
+
+    def encode(self, xs: mx.array) -> mx.array:
+        target_shape = xs.shape[:-1]
+        xs = xs.flatten(end_axis=-2)
+        dot_prod = xs @ self._embedding.swapaxes(-1, -2)
+        return (self._c2 - dot_prod).argmin(axis=-1).reshape(target_shape)
+
+    def decode(self, xs: mx.array) -> mx.array:
+        target_shape = list(xs.shape) + [self._dim]
+        return mx.take(self._embedding, xs.flatten(), axis=0).reshape(target_shape)
+
+
+class VectorQuantization(nn.Module):
+    def __init__(self, dim: int, codebook_size: int, codebook_dim: int | None):
+        super().__init__()
+        codebook_dim = dim if codebook_dim is None else codebook_dim
+        if dim == codebook_dim:
+            self.project_in = None
+            self.project_out = None
+        else:
+            self.project_in = nn.Linear(dim, codebook_dim)
+            self.project_out = nn.Linear(codebook_dim, dim)
+        self.codebook = EuclideanCodebook(dim=codebook_dim, codebook_size=codebook_size)
+
+    def encode(self, xs: mx.array) -> mx.array:
+        xs = xs.swapaxes(-1, -2)
+        if self.project_in is not None:
+            xs = self.project_in(xs)
+        return self.codebook.encode(xs)
+
+    def decode(self, xs: mx.array) -> mx.array:
+        xs = self.codebook.decode(xs)
+        if self.project_out is not None:
+            xs = self.project_out(xs)
+        return xs.swapaxes(-1, -2)
+
+
+class ResidualVectorQuantization(nn.Module):
+    def __init__(self, nq: int, dim: int, codebook_size: int, codebook_dim: int | None):
+        super().__init__()
+        layers = []
+        for _ in range(nq):
+            vq = VectorQuantization(
+                dim=dim,
+                codebook_size=codebook_size,
+                codebook_dim=codebook_dim,
+            )
+            layers.append(vq)
+        self.layers = layers
+
+    def encode(self, xs: mx.array) -> mx.array:
+        codes = []
+        residual = xs
+        for layer in self.layers:
+            indices = layer.encode(residual)
+            quantized = layer.decode(indices)
+            residual = residual - quantized
+            codes.append(indices)
+        return mx.stack(codes, axis=0)
+
+    def decode(self, xs: mx.array) -> mx.array:
+        seq_len = xs.shape[0]
+        quantized = self.layers[0].decode(xs[0])
+        for i in range(1, seq_len):
+            quantized = quantized + self.layers[i].decode(xs[i])
+        return quantized
+
+
+class ResidualVectorQuantizer(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        input_dim: int | None,
+        output_dim: int | None,
+        nq: int,
+        bins: int,
+        force_projection: bool,
+    ):
+        super().__init__()
+        input_dim = dim if input_dim is None else input_dim
+        output_dim = dim if output_dim is None else output_dim
+        if input_dim == dim and not force_projection:
+            self.input_proj = None
+        else:
+            self.input_proj = Conv1d(input_dim, dim, 1, bias=False)
+        if output_dim == dim and not force_projection:
+            self.output_proj = None
+        else:
+            self.output_proj = Conv1d(dim, output_dim, 1, bias=False)
+        self.vq = ResidualVectorQuantization(
+            nq=nq,
+            dim=dim,
+            codebook_size=bins,
+            codebook_dim=None,
+        )
+
+    def encode(self, xs: mx.array) -> mx.array:
+        if self.input_proj is not None:
+            xs = self.input_proj(xs)
+        return self.vq.encode(xs).swapaxes(0, 1)
+
+    def decode(self, xs: mx.array) -> mx.array:
+        xs = xs.swapaxes(0, 1)
+        quantized = self.vq.decode(xs)
+        if self.output_proj is not None:
+            quantized = self.output_proj(quantized)
+        return quantized
+
+
+class SplitResidualVectorQuantizer(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        input_dim: int | None,
+        output_dim: int | None,
+        nq: int,
+        bins: int,
+    ):
+        super().__init__()
+        self._nq = nq
+        self.rvq_first = ResidualVectorQuantizer(
+            dim=dim,
+            input_dim=input_dim,
+            output_dim=output_dim,
+            nq=1,
+            bins=bins,
+            force_projection=True,
+        )
+        self.rvq_rest = ResidualVectorQuantizer(
+            dim=dim,
+            input_dim=input_dim,
+            output_dim=output_dim,
+            nq=nq - 1,
+            bins=bins,
+            force_projection=True,
+        )
+
+    def encode(self, xs: mx.array) -> mx.array:
+        codes = self.rvq_first.encode(xs)
+        if self._nq > 1:
+            rest_codes = self.rvq_rest.encode(xs)
+            codes = mx.concat([codes, rest_codes], axis=1)
+        return codes
+
+    def decode(self, xs: mx.array) -> mx.array:
+        quantized = self.rvq_first.decode(xs[:, :1])
+        if self._nq > 1:
+            quantized = quantized + self.rvq_rest.decode(xs[:, 1:])
+        return quantized

nexaai/mlx_backend/mlx_audio/codec/models/mimi/modules/seanet.py

@@ -0,0 +1,314 @@
+# 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.
+
+from dataclasses import dataclass
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .conv import StreamableConv1d, StreamableConvTranspose1d
+
+
+@dataclass
+class SeanetConfig:
+    dimension: int
+    channels: int
+    causal: bool
+    nfilters: int
+    nresidual_layers: int
+    ratios: list[int]
+    ksize: int
+    residual_ksize: int
+    last_ksize: int
+    dilation_base: int
+    pad_mode: str
+    true_skip: bool
+    compress: int
+
+
+class StreamingAdd(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self._lhs = None
+        self._rhs = None
+
+    def step(self, lhs: mx.array, rhs: mx.array) -> mx.array:
+        if self._lhs is not None:
+            lhs = mx.concat([self._lhs, lhs], axis=-1)
+            self._lhs = None
+        if self._rhs is not None:
+            rhs = mx.concat([self._rhs, rhs], axis=-1)
+            self._rhs = None
+        lhs_l = lhs.shape[-1]
+        rhs_l = rhs.shape[-1]
+        if lhs_l == rhs_l:
+            return lhs + rhs
+        elif lhs_l < rhs_l:
+            self._rhs = rhs[..., lhs_l:]
+            return lhs + rhs[..., :lhs_l]
+        else:
+            self._lhs = lhs[..., rhs_l:]
+            return lhs[..., :rhs_l] + rhs
+
+
+class SeanetResnetBlock(nn.Module):
+    def __init__(self, cfg: SeanetConfig, dim: int, ksizes_and_dilations: list):
+        super().__init__()
+        block = []
+        hidden = dim // cfg.compress
+        for i, (ksize, dilation) in enumerate(ksizes_and_dilations):
+            in_channels = dim if i == 0 else hidden
+            out_channels = dim if i == len(ksizes_and_dilations) - 1 else hidden
+            c = StreamableConv1d(
+                in_channels=in_channels,
+                out_channels=out_channels,
+                ksize=ksize,
+                stride=1,
+                dilation=dilation,
+                groups=1,
+                bias=True,
+                causal=cfg.causal,
+                pad_mode=cfg.pad_mode,
+            )
+            block.append(c)
+        self.block = block
+        self.streaming_add = StreamingAdd()
+
+        if cfg.true_skip:
+            self.shortcut = None
+        else:
+            self.shortcut = StreamableConv1d(
+                in_channels=dim,
+                out_channels=dim,
+                ksize=1,
+                stride=1,
+                dilation=1,
+                groups=1,
+                bias=True,
+                causal=cfg.causal,
+                pad_mode=cfg.pad_mode,
+            )
+
+    def reset_state(self):
+        if self.shortcut is not None:
+            self.shortcut.reset_state()
+        for b in self.block:
+            b.reset_state()
+
+    def __call__(self, xs: mx.array) -> mx.array:
+        residual = xs
+        for b in self.block:
+            xs = b(nn.elu(xs, alpha=1.0))
+        if self.shortcut is None:
+            xs = xs + residual
+        else:
+            xs = xs + self.shortcut(residual)
+        return xs
+
+    def step(self, xs: mx.array) -> mx.array:
+        residual = xs
+        for b in self.block:
+            xs = b.step(nn.elu(xs, alpha=1.0))
+        if self.shortcut is None:
+            xs = self.streaming_add.step(xs, residual)
+        else:
+            xs = self.streaming_add.step(xs, self.shortcut.step(residual))
+        return xs
+
+
+class EncoderLayer(nn.Module):
+    def __init__(self, cfg: SeanetConfig, ratio: int, mult: int):
+        super().__init__()
+        residuals = []
+        dilation = 1
+        for _ in range(cfg.nresidual_layers):
+            b = SeanetResnetBlock(
+                cfg,
+                dim=mult * cfg.nfilters,
+                ksizes_and_dilations=[(cfg.residual_ksize, dilation), (1, 1)],
+            )
+            residuals.append(b)
+            dilation *= cfg.dilation_base
+        self.residuals = residuals
+        self.downsample = StreamableConv1d(
+            in_channels=mult * cfg.nfilters,
+            out_channels=mult * cfg.nfilters * 2,
+            ksize=ratio * 2,
+            stride=ratio,
+            dilation=1,
+            groups=1,
+            bias=True,
+            causal=True,
+            pad_mode=cfg.pad_mode,
+        )
+
+    def reset_state(self):
+        self.downsample.reset_state()
+        for r in self.residuals:
+            r.reset_state()
+
+    def __call__(self, xs: mx.array) -> mx.array:
+        for r in self.residuals:
+            xs = r(xs)
+        return self.downsample(nn.elu(xs, alpha=1.0))
+
+    def step(self, xs: mx.array) -> mx.array:
+        for r in self.residuals:
+            xs = r.step(xs)
+        return self.downsample.step(nn.elu(xs, alpha=1.0))
+
+
+class SeanetEncoder(nn.Module):
+    def __init__(self, cfg: SeanetConfig):
+        super().__init__()
+        mult = 1
+        self.init_conv1d = StreamableConv1d(
+            in_channels=cfg.channels,
+            out_channels=mult * cfg.nfilters,
+            ksize=cfg.ksize,
+            stride=1,
+            dilation=1,
+            groups=1,
+            bias=True,
+            causal=cfg.causal,
+            pad_mode=cfg.pad_mode,
+        )
+        layers = []
+        for ratio in reversed(cfg.ratios):
+            layers.append(EncoderLayer(cfg, ratio=ratio, mult=mult))
+            mult *= 2
+        self.layers = layers
+        self.final_conv1d = StreamableConv1d(
+            in_channels=mult * cfg.nfilters,
+            out_channels=cfg.dimension,
+            ksize=cfg.last_ksize,
+            stride=1,
+            dilation=1,
+            groups=1,
+            bias=True,
+            causal=cfg.causal,
+            pad_mode=cfg.pad_mode,
+        )
+
+    def reset_state(self):
+        self.init_conv1d.reset_state()
+        self.final_conv1d.reset_state()
+        for layer in self.layers:
+            layer.reset_state()
+
+    def __call__(self, xs: mx.array) -> mx.array:
+        xs = self.init_conv1d(xs)
+        for layer in self.layers:
+            xs = layer(xs)
+        xs = nn.elu(xs, alpha=1.0)
+        return self.final_conv1d(xs)
+
+    def step(self, xs: mx.array) -> mx.array:
+        xs = self.init_conv1d.step(xs)
+        for layer in self.layers:
+            xs = layer.step(xs)
+        xs = nn.elu(xs, alpha=1.0)
+        return self.final_conv1d.step(xs)
+
+
+class DecoderLayer(nn.Module):
+    def __init__(self, cfg: SeanetConfig, ratio: int, mult: int):
+        super().__init__()
+        self.upsample = StreamableConvTranspose1d(
+            in_channels=mult * cfg.nfilters,
+            out_channels=mult * cfg.nfilters // 2,
+            ksize=ratio * 2,
+            stride=ratio,
+            groups=1,
+            bias=True,
+            causal=cfg.causal,
+        )
+        residuals = []
+        dilation = 1
+        for _ in range(cfg.nresidual_layers):
+            r = SeanetResnetBlock(
+                cfg,
+                dim=mult * cfg.nfilters // 2,
+                ksizes_and_dilations=[(cfg.residual_ksize, dilation), (1, 1)],
+            )
+            residuals.append(r)
+            dilation *= cfg.dilation_base
+        self.residuals = residuals
+
+    def reset_state(self):
+        self.upsample.reset_state()
+        for r in self.residuals:
+            r.reset_state()
+
+    def __call__(self, xs: mx.array) -> mx.array:
+        xs = self.upsample(nn.elu(xs, alpha=1.0))
+        for r in self.residuals:
+            xs = r(xs)
+        return xs
+
+    def step(self, xs: mx.array) -> mx.array:
+        xs = self.upsample.step(nn.elu(xs, alpha=1.0))
+        for r in self.residuals:
+            xs = r.step(xs)
+        return xs
+
+
+class SeanetDecoder(nn.Module):
+    def __init__(self, cfg: SeanetConfig):
+        super().__init__()
+        mult = 1 << len(cfg.ratios)
+        self.init_conv1d = StreamableConv1d(
+            in_channels=cfg.dimension,
+            out_channels=mult * cfg.nfilters,
+            ksize=cfg.ksize,
+            stride=1,
+            dilation=1,
+            groups=1,
+            bias=True,
+            causal=cfg.causal,
+            pad_mode=cfg.pad_mode,
+        )
+        layers = []
+        for ratio in cfg.ratios:
+            layers.append(DecoderLayer(cfg, ratio=ratio, mult=mult))
+            mult //= 2
+        self.layers = layers
+        self.final_conv1d = StreamableConv1d(
+            in_channels=cfg.nfilters,
+            out_channels=cfg.channels,
+            ksize=cfg.last_ksize,
+            stride=1,
+            dilation=1,
+            groups=1,
+            bias=True,
+            causal=cfg.causal,
+            pad_mode=cfg.pad_mode,
+        )
+
+    def reset_state(self):
+        self.init_conv1d.reset_state()
+        self.final_conv1d.reset_state()
+        for layer in self.layers:
+            layer.reset_state()
+
+    def __call__(self, xs: mx.array) -> mx.array:
+        xs = self.init_conv1d(xs)
+        for layer in self.layers:
+            xs = layer(xs)
+        xs = nn.elu(xs, alpha=1.0)
+        return self.final_conv1d(xs)
+
+    def step(self, xs: mx.array) -> mx.array:
+        xs = self.init_conv1d.step(xs)
+        for layer in self.layers:
+            xs = layer.step(xs)
+        xs = nn.elu(xs, alpha=1.0)
+        return self.final_conv1d.step(xs)
+
+
+class Seanet(nn.Module):
+    def __init__(self, cfg: SeanetConfig):
+        super().__init__()
+        self.encoder = SeanetEncoder(cfg)
+        self.decoder = SeanetDecoder(cfg)