nexaai 1.0.29__cp310-cp310-macosx_14_0_universal2.whl

nexaai/mlx_backend/mlx_audio/codec/models/snac/layers.py

@@ -0,0 +1,306 @@
+import math
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .attention import LocalMHA
+
+
+def normalize_weight(x, except_dim=0):
+    if x.ndim != 3:
+        raise ValueError("Input tensor must have 3 dimensions")
+
+    axes = tuple(i for i in range(x.ndim) if i != except_dim)
+    return mx.sqrt(mx.sum(mx.power(x, 2), axis=axes, keepdims=True))
+
+
+class WNConv1d(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: int,
+        stride: int = 1,
+        padding: int = 0,
+        dilation: int = 1,
+        groups: int = 1,
+        bias: bool = True,
+    ):
+        super().__init__()
+
+        if bias:
+            self.bias = mx.zeros((out_channels,))
+
+        self.kernel_size = kernel_size
+        self.padding = padding
+        self.dilation = dilation
+        self.stride = stride
+        self.groups = groups
+
+        scale = math.sqrt(1 / (in_channels * kernel_size))
+        weight_init = mx.random.uniform(
+            low=-scale,
+            high=scale,
+            shape=(out_channels, kernel_size, in_channels // groups),
+        )
+        self.weight_g = normalize_weight(weight_init)
+        self.weight_v = weight_init / (self.weight_g + 1e-12)
+
+    def _extra_repr(self):
+        return (
+            f"kernel_size={self.kernel_size}, stride={self.stride}, "
+            f"padding={self.padding}, dilation={self.dilation}, "
+            f"groups={self.groups}, bias={'bias' in self}"
+        )
+
+    def __call__(self, x):
+        weight = self.weight_g * self.weight_v / normalize_weight(self.weight_v)
+        y = mx.conv1d(x, weight, self.stride, self.padding, self.dilation, self.groups)
+        if hasattr(self, "bias"):
+            y = y + self.bias
+        return y
+
+
+class WNConvTranspose1d(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: int,
+        stride: int = 1,
+        padding: int = 0,
+        dilation: int = 1,
+        output_padding: int = 0,
+        groups: int = 1,
+        bias: bool = True,
+    ):
+        super().__init__()
+
+        self.bias = mx.zeros((out_channels,)) if bias else None
+
+        self.kernel_size = kernel_size
+        self.padding = padding
+        self.dilation = dilation
+        self.stride = stride
+        self.output_padding = output_padding
+        self.groups = groups
+
+        scale = math.sqrt(1 / (in_channels * kernel_size))
+        weight_init = mx.random.uniform(
+            low=-scale,
+            high=scale,
+            shape=(in_channels, kernel_size, out_channels // groups),
+        )
+        self.weight_g = normalize_weight(weight_init, except_dim=0)
+        self.weight_v = weight_init / (self.weight_g + 1e-12)
+
+    def _extra_repr(self):
+        return (
+            f"kernel_size={self.kernel_size}, stride={self.stride}, "
+            f"padding={self.padding}, dilation={self.dilation}, "
+            f"output_padding={self.output_padding}, "
+            f"groups={self.groups}, bias={self.bias is not None}"
+        )
+
+    def __call__(self, x):
+        weight = (
+            self.weight_g
+            * self.weight_v
+            / normalize_weight(self.weight_v, except_dim=0)
+        )
+        y = mx.conv_transpose1d(
+            x,
+            weight.swapaxes(0, 2),  # mlx uses (out_channels, kernel_size, in_channels)
+            self.stride,
+            self.padding,
+            self.dilation,
+            self.groups,
+        )
+        if self.bias is not None:
+            y = y + self.bias
+        return y
+
+
+def snake(x, alpha):
+    shape = x.shape
+    x = x.reshape(shape[0], shape[1], -1)
+    recip = mx.reciprocal(alpha + 1e-9)
+    x = x + recip * mx.power(mx.sin(alpha * x), 2)
+    x = x.reshape(shape)
+    return x
+
+
+class Encoder(nn.Module):
+    def __init__(
+        self,
+        d_model=64,
+        strides=[3, 3, 7, 7],
+        depthwise=False,
+        attn_window_size=32,
+    ):
+        super().__init__()
+        layers = [WNConv1d(1, d_model, kernel_size=7, padding=3)]
+        for stride in strides:
+            d_model *= 2
+            groups = d_model // 2 if depthwise else 1
+            layers += [EncoderBlock(output_dim=d_model, stride=stride, groups=groups)]
+        if attn_window_size is not None:
+            layers += [LocalMHA(dim=d_model, window_size=attn_window_size)]
+        groups = d_model if depthwise else 1
+        layers += [
+            WNConv1d(d_model, d_model, kernel_size=7, padding=3, groups=groups),
+        ]
+        self.block = nn.Sequential(*layers)
+
+    def __call__(self, x):
+        return self.block(x).moveaxis(1, 2)
+
+
+class Decoder(nn.Module):
+    def __init__(
+        self,
+        input_channel,
+        channels,
+        rates,
+        noise=False,
+        depthwise=False,
+        attn_window_size=32,
+        d_out=1,
+    ):
+        super().__init__()
+        if depthwise:
+            layers = [
+                WNConv1d(
+                    input_channel,
+                    input_channel,
+                    kernel_size=7,
+                    padding=3,
+                    groups=input_channel,
+                ),
+                WNConv1d(input_channel, channels, kernel_size=1),
+            ]
+        else:
+            layers = [WNConv1d(input_channel, channels, kernel_size=7, padding=3)]
+
+        if attn_window_size is not None:
+            layers += [LocalMHA(dim=channels, window_size=attn_window_size)]
+
+        for i, stride in enumerate(rates):
+            input_dim = channels // (2**i)
+            output_dim = channels // (2 ** (i + 1))
+            groups = output_dim if depthwise else 1
+            layers.append(
+                DecoderBlock(input_dim, output_dim, stride, noise, groups=groups)
+            )
+
+        layers += [
+            Snake1d(output_dim),
+            WNConv1d(output_dim, d_out, kernel_size=7, padding=3),
+            nn.Tanh(),
+        ]
+        self.model = nn.Sequential(*layers)
+
+    def __call__(self, x):
+        x = self.model(x)
+        return x
+
+
+class ResidualUnit(nn.Module):
+    def __init__(self, dim=16, dilation=1, kernel=7, groups=1):
+        super().__init__()
+        pad = ((kernel - 1) * dilation) // 2
+        self.block = nn.Sequential(
+            Snake1d(dim),
+            WNConv1d(
+                dim,
+                dim,
+                kernel_size=kernel,
+                dilation=dilation,
+                padding=pad,
+                groups=groups,
+            ),
+            Snake1d(dim),
+            WNConv1d(dim, dim, kernel_size=1),
+        )
+
+    def __call__(self, x):
+        y = self.block(x)
+        pad = (x.shape[-1] - y.shape[-1]) // 2
+        if pad > 0:
+            x = x[..., pad:-pad]
+        return x + y
+
+
+class EncoderBlock(nn.Module):
+    def __init__(self, output_dim=16, input_dim=None, stride=1, groups=1):
+        super().__init__()
+        input_dim = input_dim or output_dim // 2
+        self.block = nn.Sequential(
+            ResidualUnit(input_dim, dilation=1, groups=groups),
+            ResidualUnit(input_dim, dilation=3, groups=groups),
+            ResidualUnit(input_dim, dilation=9, groups=groups),
+            Snake1d(input_dim),
+            WNConv1d(
+                input_dim,
+                output_dim,
+                kernel_size=2 * stride,
+                stride=stride,
+                padding=math.ceil(stride / 2),
+            ),
+        )
+
+    def __call__(self, x):
+        return self.block(x)
+
+
+class NoiseBlock(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+        self.linear = WNConv1d(dim, dim, kernel_size=1, bias=False)
+
+    def __call__(self, x):
+        B, C, T = x.shape
+        noise = mx.random.normal((B, 1, T))
+        h = self.linear(x)
+        n = noise * h
+        x = x + n
+        return x
+
+
+class DecoderBlock(nn.Module):
+    def __init__(self, input_dim=16, output_dim=8, stride=1, noise=False, groups=1):
+        super().__init__()
+
+        layers = [
+            Snake1d(input_dim),
+            WNConvTranspose1d(
+                input_dim,
+                output_dim,
+                kernel_size=2 * stride,
+                stride=stride,
+                padding=math.ceil(stride / 2),
+                output_padding=stride % 2,
+            ),
+        ]
+        if noise:
+            layers.append(NoiseBlock(output_dim))
+        layers.extend(
+            [
+                ResidualUnit(output_dim, dilation=1, groups=groups),
+                ResidualUnit(output_dim, dilation=3, groups=groups),
+                ResidualUnit(output_dim, dilation=9, groups=groups),
+            ]
+        )
+        self.block = nn.Sequential(*layers)
+
+    def __call__(self, x):
+        return self.block(x)
+
+
+class Snake1d(nn.Module):
+    def __init__(self, channels):
+        super().__init__()
+        self.alpha = mx.ones((1, channels, 1))
+
+    def __call__(self, x):
+        return snake(x, self.alpha.swapaxes(1, 2))

nexaai/mlx_backend/mlx_audio/codec/models/snac/snac.py

@@ -0,0 +1,154 @@
+import json
+import math
+from pathlib import Path
+from typing import List, Tuple
+
+import mlx.core as mx
+import mlx.nn as nn
+import numpy as np
+from huggingface_hub import snapshot_download
+
+from .layers import Decoder, Encoder
+from .vq import ResidualVectorQuantize
+
+
+class SNAC(nn.Module):
+    def __init__(
+        self,
+        sampling_rate=44100,
+        encoder_dim=64,
+        encoder_rates=[3, 3, 7, 7],
+        latent_dim=None,
+        decoder_dim=1536,
+        decoder_rates=[7, 7, 3, 3],
+        attn_window_size=32,
+        codebook_size=4096,
+        codebook_dim=8,
+        vq_strides=[8, 4, 2, 1],
+        noise=True,
+        depthwise=True,
+    ):
+        super().__init__()
+        self.sampling_rate = sampling_rate
+        self.encoder_dim = encoder_dim
+        self.encoder_rates = encoder_rates
+        self.decoder_dim = decoder_dim
+        self.decoder_rates = decoder_rates
+        if latent_dim is None:
+            latent_dim = encoder_dim * (2 ** len(encoder_rates))
+        self.latent_dim = latent_dim
+        self.hop_length = np.prod(encoder_rates)
+        self.encoder = Encoder(
+            encoder_dim,
+            encoder_rates,
+            depthwise=depthwise,
+            attn_window_size=attn_window_size,
+        )
+        self.n_codebooks = len(vq_strides)
+        self.codebook_size = codebook_size
+        self.codebook_dim = codebook_dim
+        self.vq_strides = vq_strides
+        self.attn_window_size = attn_window_size
+        self.quantizer = ResidualVectorQuantize(
+            input_dim=latent_dim,
+            codebook_size=codebook_size,
+            codebook_dim=codebook_dim,
+            vq_strides=vq_strides,
+        )
+        self.decoder = Decoder(
+            latent_dim,
+            decoder_dim,
+            decoder_rates,
+            noise,
+            depthwise=depthwise,
+            attn_window_size=attn_window_size,
+        )
+
+    def preprocess(self, audio_data):
+        length = audio_data.shape[-1]
+
+        def lcm(a, b):
+            return abs(a * b) // math.gcd(a, b)
+
+        lcm_value = self.vq_strides[0]
+        for i in range(1, len(self.vq_strides)):
+            lcm_value = lcm(lcm_value, self.vq_strides[i])
+
+        if self.attn_window_size:
+            lcm_value = lcm(lcm_value, self.attn_window_size)
+
+        pad_to = self.hop_length * lcm_value
+        right_pad = math.ceil(length / pad_to) * pad_to - length
+
+        # Pad the audio data
+        audio_data = mx.pad(audio_data, [(0, 0), (0, 0), (0, right_pad)])
+        return audio_data
+
+    def __call__(self, audio_data: mx.array) -> Tuple[mx.array, List[mx.array]]:
+        length = audio_data.shape[-1]
+        audio_data = self.preprocess(audio_data)
+        z = self.encoder(audio_data.moveaxis(1, 2))
+        z_q, codes = self.quantizer(z)
+        audio_hat = self.decoder(z_q)
+        return audio_hat[..., :length], codes
+
+    def encode(self, audio_data: mx.array) -> List[mx.array]:
+        audio_data = self.preprocess(audio_data)
+        z = self.encoder(audio_data.moveaxis(1, 2))
+        _, codes = self.quantizer(z)
+        return codes
+
+    def decode(self, codes: List[mx.array]) -> mx.array:
+        z_q = self.quantizer.from_codes(codes)
+        audio_hat = self.decoder(z_q.moveaxis(1, 2))
+        return audio_hat
+
+    def _extra_repr(self):
+        return (
+            f"sampling_rate={self.sampling_rate}, "
+            f"encoder_dim={self.encoder_dim}, "
+            f"encoder_rates={self.encoder_rates}, "
+            f"latent_dim={self.latent_dim}, "
+            f"decoder_dim={self.decoder_dim}, "
+            f"decoder_rates={self.decoder_rates}, "
+            f"codebook_size={self.codebook_size}, "
+            f"codebook_dim={self.codebook_dim}, "
+            f"vq_strides={self.vq_strides}"
+        )
+
+    @classmethod
+    def from_config(cls, config_path):
+        with open(config_path, "r") as f:
+            config = json.load(f)
+        model = cls(**config)
+        return model
+
+    @classmethod
+    def from_pretrained(cls, repo_id, **kwargs):
+        path = fetch_from_hub(repo_id)
+
+        if path is None:
+            raise ValueError(f"Could not find model {path}")
+
+        model_path = path / "model.safetensors"
+        config_path = path / "config.json"
+        snac = cls.from_config(config_path)
+
+        weights = mx.load(model_path.as_posix(), format="safetensors")
+        snac.load_weights(list(weights.items()))
+        mx.eval(snac.parameters())
+
+        return snac
+
+
+# fetch model from hub
+
+
+def fetch_from_hub(hf_repo: str) -> Path:
+    model_path = Path(
+        snapshot_download(
+            repo_id=hf_repo,
+            allow_patterns=["*.safetensors", "*.json"],
+        )
+    )
+    return model_path

nexaai/mlx_backend/mlx_audio/codec/models/snac/vq.py

@@ -0,0 +1,135 @@
+from typing import List
+
+import mlx.core as mx
+import mlx.nn as nn
+from einops.array_api import rearrange
+
+from .layers import WNConv1d
+
+
+class VectorQuantize(nn.Module):
+    def __init__(
+        self, input_dim: int, codebook_size: int, codebook_dim: int, stride: int = 1
+    ):
+        super().__init__()
+        self.codebook_size = codebook_size
+        self.codebook_dim = codebook_dim
+        self.stride = stride
+
+        self.in_proj = WNConv1d(input_dim, codebook_dim, kernel_size=1)
+        self.out_proj = WNConv1d(codebook_dim, input_dim, kernel_size=1)
+        self.codebook = nn.Embedding(codebook_size, codebook_dim)
+
+    def __call__(self, z):
+        z = z.moveaxis(1, 2)
+
+        if self.stride > 1:
+            kernel_size = self.stride
+            stride = self.stride
+            kernel = mx.ones((z.shape[2], kernel_size, 1)) / kernel_size
+            z = mx.conv1d(z, kernel, stride=stride, padding=0, groups=z.shape[2])
+
+        # Factorized codes - Project input into low-dimensional space
+        z_e = self.in_proj(z).moveaxis(1, 2)  # z_e : (B x D x T)
+        z_q, indices = self.decode_latents(z_e)
+
+        # Straight-through estimator: z_e + stop_gradient(z_q - z_e)
+        z_q = z_e + (z_q - z_e)
+
+        z_q = self.out_proj(z_q.moveaxis(1, 2)).moveaxis(1, 2)
+
+        if self.stride > 1:
+            # Implement repeat_interleave
+            shape = list(z_q.shape)
+            shape[-1] *= self.stride
+            # Create a new tensor with the expanded shape
+            expanded = mx.zeros(shape)
+
+            # Fill the expanded tensor with repeated values
+            for i in range(self.stride):
+                expanded[..., i :: self.stride] = z_q
+
+            z_q = expanded
+
+        return z_q, indices
+
+    def embed_code(self, embed_id):
+        return self.codebook.weight[embed_id]
+
+    def decode_code(self, embed_id):
+        return self.embed_code(embed_id).moveaxis(1, 2)
+
+    def decode_latents(self, latents):
+        encodings = rearrange(latents, "b d t -> (b t) d")
+        codebook = self.codebook.weight  # codebook: (N x D)
+
+        encodings = normalize(encodings)
+        codebook = normalize(codebook)
+
+        dist = (
+            mx.power(encodings, 2).sum(1, keepdims=True)
+            - 2 * encodings @ codebook.T
+            + mx.power(codebook, 2).sum(1, keepdims=True).T
+        )
+        min_dist = (-dist).argmax(1)
+        indices = rearrange(min_dist, "(b t) -> b t", b=latents.shape[0])
+        z_q = self.decode_code(indices)
+        return z_q, indices
+
+
+class ResidualVectorQuantize(nn.Module):
+    def __init__(
+        self,
+        input_dim: int = 512,
+        codebook_size: int = 1024,
+        codebook_dim: int = 8,
+        vq_strides: List[int] = [1, 1, 1, 1],
+    ):
+        super().__init__()
+        self.n_codebooks = len(vq_strides)
+        self.codebook_dim = codebook_dim
+        self.codebook_size = codebook_size
+        self.quantizers = [
+            VectorQuantize(input_dim, codebook_size, codebook_dim, stride)
+            for stride in vq_strides
+        ]
+
+    def __call__(self, z):
+        z_q = 0
+        residual = z
+        codes = []
+        for i, quantizer in enumerate(self.quantizers):
+            z_q_i, indices_i = quantizer(residual)
+            z_q = z_q + z_q_i
+            residual = residual - z_q_i
+            codes.append(indices_i)
+
+        return z_q, codes
+
+    def from_codes(self, codes: List[mx.array]) -> mx.array:
+        z_q = 0.0
+        for i in range(self.n_codebooks):
+            z_p_i = self.quantizers[i].decode_code(codes[i])
+            z_q_i = self.quantizers[i].out_proj(z_p_i.moveaxis(1, 2)).moveaxis(1, 2)
+
+            # Handle repeat_interleave for stride > 1
+            if self.quantizers[i].stride > 1:
+                stride = self.quantizers[i].stride
+                shape = list(z_q_i.shape)
+                shape[-1] *= stride
+                expanded = mx.zeros(shape)
+
+                for j in range(stride):
+                    expanded[..., j::stride] = z_q_i
+
+                z_q_i = expanded
+
+            z_q += z_q_i
+
+        return z_q
+
+
+def normalize(x, p=2.0, dim=1, eps=1e-12):
+    """L2 normalization function"""
+    norm = mx.power(mx.sum(mx.power(mx.abs(x), p), axis=dim, keepdims=True), 1 / p)
+    return x / mx.maximum(norm, eps)

nexaai/mlx_backend/mlx_audio/codec/models/vocos/__init__.py

@@ -0,0 +1 @@
+from .vocos import Vocos

nexaai/mlx_backend/mlx_audio/codec/models/vocos/mel.py

@@ -0,0 +1,33 @@
+from __future__ import annotations
+
+import mlx.core as mx
+
+from mlx_audio.utils import hanning, mel_filters, stft
+
+
+def log_mel_spectrogram(
+    audio: mx.array,
+    sample_rate: int = 24_000,
+    n_mels: int = 100,
+    n_fft: int = 1024,
+    hop_length: int = 256,
+    padding: int = 0,
+):
+    if not isinstance(audio, mx.array):
+        audio = mx.array(audio)
+
+    if padding > 0:
+        audio = mx.pad(audio, (0, padding))
+
+    freqs = stft(audio, window=hanning(n_fft), n_fft=n_fft, win_length=hop_length)
+    magnitudes = freqs[:-1, :].abs()
+    filters = mel_filters(
+        sample_rate=sample_rate,
+        n_fft=n_fft,
+        n_mels=n_mels,
+        norm=None,
+        mel_scale="htk",
+    )
+    mel_spec = magnitudes @ filters.T
+    log_spec = mx.maximum(mel_spec, 1e-5).log()
+    return mx.expand_dims(log_spec, axis=0)