nexaai 1.0.29__cp310-cp310-macosx_14_0_universal2.whl

nexaai/mlx_backend/mlx_audio/tts/models/spark/modules/speaker/pooling_layers.py

@@ -0,0 +1,297 @@
+# Copyright (c) 2021 Shuai Wang (wsstriving@gmail.com)
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+Pooling functions to aggregate frame-level deep features
+into segment-level speaker embeddings
+
+High-order statistics are surprisingly effective, TSDP acts similarly as TSTP,
+even though we remove the mean statistic, on Voxceleb.
+"""
+import mlx.core as mx
+import mlx.nn as nn
+
+
+class TAP(nn.Module):
+    """
+    Temporal average pooling, only first-order mean is considered
+    """
+
+    def __init__(self, in_dim=0, **kwargs):
+        super(TAP, self).__init__()
+        self.in_dim = in_dim
+
+    def __call__(self, x):
+        pooling_mean = mx.mean(x, axis=-1)
+        # To be compatable with 2D input
+        pooling_mean = pooling_mean.flatten(start_axis=1)
+        return pooling_mean
+
+    def get_out_dim(self):
+        self.out_dim = self.in_dim
+        return self.out_dim
+
+
+class TSDP(nn.Module):
+    """
+    Temporal standard deviation pooling, only second-order std is considered
+    """
+
+    def __init__(self, in_dim=0, **kwargs):
+        super(TSDP, self).__init__()
+        self.in_dim = in_dim
+
+    def __call__(self, x):
+        # The last dimension is the temporal axis
+        pooling_std = mx.sqrt(mx.var(x, axis=-1) + 1e-7)
+        pooling_std = pooling_std.flatten(start_axis=1)
+        return pooling_std
+
+    def get_out_dim(self):
+        self.out_dim = self.in_dim
+        return self.out_dim
+
+
+class TSTP(nn.Module):
+    """
+    Temporal statistics pooling, concatenate mean and std, which is used in
+    x-vector
+    Comment: simple concatenation can not make full use of both statistics
+    """
+
+    def __init__(self, in_dim=0, **kwargs):
+        super(TSTP, self).__init__()
+        self.in_dim = in_dim
+
+    def __call__(self, x):
+        # The last dimension is the temporal axis
+        pooling_mean = mx.mean(x, axis=-1)
+        pooling_std = mx.sqrt(mx.var(x, axis=-1) + 1e-7)
+        pooling_mean = pooling_mean.flatten(start_axis=1)
+        pooling_std = pooling_std.flatten(start_axis=1)
+        stats = mx.concatenate((pooling_mean, pooling_std), axis=1)
+        return stats
+
+    def get_out_dim(self):
+        self.out_dim = self.in_dim * 2
+        return self.out_dim
+
+
+class ASTP(nn.Module):
+    """Attentive statistics pooling: Channel- and context-dependent
+    statistics pooling, first used in ECAPA_TDNN.
+    """
+
+    def __init__(self, in_dim, bottleneck_dim=128, global_context_att=False, **kwargs):
+        super(ASTP, self).__init__()
+        self.in_dim = in_dim
+        self.global_context_att = global_context_att
+
+        # Use Conv1d with stride == 1 rather than Linear, then we don't
+        # need to transpose inputs.
+        if global_context_att:
+            self.linear1 = nn.Conv1d(
+                in_dim * 3, bottleneck_dim, kernel_size=1
+            )  # equals W and b in the paper
+        else:
+            self.linear1 = nn.Conv1d(
+                in_dim, bottleneck_dim, kernel_size=1
+            )  # equals W and b in the paper
+        self.linear2 = nn.Conv1d(
+            bottleneck_dim, in_dim, kernel_size=1
+        )  # equals V and k in the paper
+
+    def __call__(self, x):
+        """
+        x: a 3-dimensional tensor in tdnn-based architecture (B,F,T)
+            or a 4-dimensional tensor in resnet architecture (B,C,F,T)
+            0-dim: batch-dimension, last-dim: time-dimension (frame-dimension)
+        """
+        if len(x.shape) == 4:
+            x = x.reshape(x.shape[0], x.shape[1] * x.shape[2], x.shape[3])
+        assert len(x.shape) == 3
+
+        if self.global_context_att:
+            context_mean = mx.mean(x, axis=-1)[:, :, None]
+            context_mean = mx.broadcast_to(context_mean, x.shape)
+            context_std = mx.sqrt(mx.var(x, axis=-1) + 1e-7)[:, :, None]
+            context_std = mx.broadcast_to(context_std, x.shape)
+            x_in = mx.concatenate((x, context_mean, context_std), axis=1)
+        else:
+            x_in = x
+
+        # DON'T use ReLU here! ReLU may be hard to converge.
+        alpha = mx.tanh(
+            self.linear1(x_in.transpose(0, 2, 1)).transpose(0, 2, 1)
+        )  # alpha = F.relu(self.linear1(x_in))
+        alpha = mx.softmax(
+            self.linear2(alpha.transpose(0, 2, 1)).transpose(0, 2, 1), axis=2
+        )
+        mean = mx.sum(alpha * x, axis=2)
+        var = mx.sum(alpha * (x**2), axis=2) - mean**2
+        std = mx.sqrt(mx.clip(var, 1e-7, None))
+        return mx.concatenate([mean, std], axis=1)
+
+    def get_out_dim(self):
+        self.out_dim = 2 * self.in_dim
+        return self.out_dim
+
+
+class MHASTP(nn.Module):
+    """Multi head attentive statistics pooling
+    Reference:
+        Self Multi-Head Attention for Speaker Recognition
+        https://arxiv.org/pdf/1906.09890.pdf
+    """
+
+    def __init__(
+        self, in_dim, layer_num=2, head_num=2, d_s=1, bottleneck_dim=64, **kwargs
+    ):
+        super(MHASTP, self).__init__()
+        assert (
+            in_dim % head_num
+        ) == 0  # make sure that head num can be divided by input_dim
+        self.in_dim = in_dim
+        self.head_num = head_num
+        d_model = int(in_dim / head_num)
+        channel_dims = [bottleneck_dim for i in range(layer_num + 1)]
+        if d_s > 1:
+            d_s = d_model
+        else:
+            d_s = 1
+        self.d_s = d_s
+        channel_dims[0], channel_dims[-1] = d_model, d_s
+        self.heads_att_trans = []
+        for i in range(self.head_num):
+            layers = []
+            for j in range(layer_num - 1):
+                layers.extend(
+                    [
+                        nn.Conv1d(channel_dims[j], channel_dims[j + 1], 1, 1),
+                        nn.Tanh(),
+                    ]
+                )
+            layers.append(
+                nn.Conv1d(channel_dims[layer_num - 1], channel_dims[layer_num], 1, 1)
+            )
+            self.heads_att_trans.append(nn.Sequential(*layers))
+
+    def __call__(self, input):
+        """
+        input: a 3-dimensional tensor in xvector architecture
+            or a 4-dimensional tensor in resnet architecture
+            0-dim: batch-dimension, last-dim: time-dimension (frame-dimension)
+        """
+        if len(input.shape) == 4:  # B x F x T
+            input = input.reshape(
+                input.shape[0], input.shape[1] * input.shape[2], input.shape[3]
+            )
+        assert len(input.shape) == 3
+        bs, f_dim, t_dim = input.shape
+        chunks = mx.split(input, self.head_num, axis=1)
+        # split
+        chunks_out = []
+        for i, layer in enumerate(self.heads_att_trans):
+            att_score = layer(chunks[i].transpose(0, 2, 1)).transpose(0, 2, 1)
+            alpha = mx.softmax(att_score, axis=-1)
+            mean = mx.sum(alpha * chunks[i], axis=2)
+            var = mx.sum(alpha * chunks[i] ** 2, axis=2) - mean**2
+            std = mx.sqrt(mx.clip(var, 1e-7, None))
+            chunks_out.append(mx.concatenate((mean, std), axis=1))
+        out = mx.concatenate(chunks_out, axis=1)
+        return out
+
+    def get_out_dim(self):
+        self.out_dim = 2 * self.in_dim
+        return self.out_dim
+
+
+class MQMHASTP(nn.Module):
+    """An attentive pooling
+    Reference:
+        multi query multi head attentive statistics pooling
+        https://arxiv.org/pdf/2110.05042.pdf
+    Args:
+        in_dim: the feature dimension of input
+        layer_num: the number of layer in the pooling layer
+        query_num: the number of querys
+        head_num: the number of heads
+        bottleneck_dim: the bottleneck dimension
+
+    SA (H = 1, Q = 1, n = 2, d_s = 1) ref:
+        https://www.danielpovey.com/files/2018_interspeech_xvector_attention.pdf
+    MHA (H > 1, Q = 1, n = 1, d_s = 1) ref:
+        https://arxiv.org/pdf/1906.09890.pdf
+    AS (H = 1, Q > 1, n = 2, d_s = 1) ref:
+        https://arxiv.org/pdf/1803.10963.pdf
+    VSA (H = 1, Q > 1, n = 2, d_s = d_h) ref:
+        http://www.interspeech2020.org/uploadfile/pdf/Mon-2-10-5.pdf
+    """
+
+    def __init__(
+        self,
+        in_dim,
+        layer_num=2,
+        query_num=2,
+        head_num=8,
+        d_s=2,
+        bottleneck_dim=64,
+        **kwargs
+    ):
+        super(MQMHASTP, self).__init__()
+        self.n_query = [
+            MHASTP(
+                in_dim,
+                layer_num=layer_num,
+                head_num=head_num,
+                d_s=d_s,
+                bottleneck_dim=bottleneck_dim,
+            )
+            for i in range(query_num)
+        ]
+        self.query_num = query_num
+        self.in_dim = in_dim
+
+    def __call__(self, input):
+        """
+        input: a 3-dimensional tensor in xvector architecture
+            or a 4-dimensional tensor in resnet architecture
+            0-dim: batch-dimension, last-dim: time-dimension (frame-dimension)
+        """
+        if len(input.shape) == 4:  # B x F x T
+            input = input.reshape(
+                input.shape[0], input.shape[1] * input.shape[2], input.shape[3]
+            )
+        assert len(input.shape) == 3
+        res = []
+        for i, layer in enumerate(self.n_query):
+            res.append(layer(input))
+        out = mx.concatenate(res, axis=-1)
+        return out
+
+    def get_out_dim(self):
+        self.out_dim = self.in_dim * 2 * self.query_num
+        return self.out_dim
+
+
+if __name__ == "__main__":
+    data = mx.random.normal(shape=(16, 512, 10, 35))
+    # model = StatisticsPooling()
+    model = MQMHASTP(512 * 10)
+    model = MHASTP(512 * 10)
+    model = MQMHASTP(512 * 10, context=False)
+    print(model)
+
+    out = model(data)
+    print(out.shape)
+    print(model.get_out_dim())

nexaai/mlx_backend/mlx_audio/tts/models/spark/modules/speaker/speaker_encoder.py

@@ -0,0 +1,155 @@
+# Copyright (c) 2025 SparkAudio
+#               2025 Xinsheng Wang (w.xinshawn@gmail.com)
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import List, Tuple
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from mlx_audio.tts.models.spark.modules.residual_fsq import ResidualFSQ
+from mlx_audio.tts.models.spark.modules.speaker.ecapa_tdnn import ECAPA_TDNN_GLOB_c512
+from mlx_audio.tts.models.spark.modules.speaker.perceiver_encoder import (
+    PerceiverResampler,
+)
+
+# from mlx_audio.codec.models.descript.nn.quantize import ResidualVectorQuantize
+
+
+"""
+x-vector + d-vector
+"""
+
+
+class SpeakerEncoder(nn.Module):
+
+    def __init__(
+        self,
+        input_dim: int = 100,
+        out_dim: int = 512,
+        latent_dim: int = 128,
+        token_num: int = 32,
+        fsq_levels: List[int] = [4, 4, 4, 4, 4, 4],
+        fsq_num_quantizers: int = 1,
+    ):
+        super(SpeakerEncoder, self).__init__()
+
+        self.speaker_encoder = ECAPA_TDNN_GLOB_c512(
+            feat_dim=input_dim, embed_dim=out_dim
+        )
+        self.perceiver_sampler = PerceiverResampler(
+            dim=latent_dim, dim_context=512 * 3, num_latents=token_num
+        )
+        self.quantizer = ResidualFSQ(
+            dim=latent_dim,
+            num_quantizers=fsq_num_quantizers,
+            levels=fsq_levels,
+            is_channel_first=True,
+            quantize_dropout=False,
+        )
+
+        self.project = nn.Linear(latent_dim * token_num, out_dim)
+
+    def get_codes_from_indices(self, indices: mx.array) -> mx.array:
+        zq = self.quantizer.get_codes_from_indices(indices.transpose(1, 2))
+        return zq.transpose(0, 2, 1)
+
+    def get_indices(self, mels: mx.array) -> mx.array:
+        mels = mels.transpose(0, 2, 1)
+        x = self.perceiver_sampler(mels).transpose(0, 2, 1)
+        zq, indices = self.quantizer(x)
+        return indices
+
+    def __call__(self, mels: mx.array) -> Tuple[mx.array, mx.array]:
+        """
+        Args:
+            mels: (B, D_mel, T1)
+
+        Return:
+            x_vector: (B, out_dim)
+            d_vector: (B, out_dim)
+        """
+        # mels = mels.transpose(1,2)
+
+        x_vector, features = self.speaker_encoder(mels, True)
+        x = self.perceiver_sampler(features.transpose(0, 2, 1)).transpose(0, 2, 1)
+        z_q, indices = self.quantizer(x)  # zq: (B, latent_dim, T2, latent_dim)
+        x = z_q.reshape(z_q.shape[0], -1)
+        d_vector = self.project(x)
+
+        return x_vector, d_vector
+
+    def tokenize(self, mels: mx.array) -> mx.array:
+        """tokenize the input mel spectrogram"""
+        _, features = self.speaker_encoder(mels, True)
+        x = self.perceiver_sampler(features.transpose(0, 2, 1)).transpose(0, 2, 1)
+        z_q, indices = self.quantizer(x)
+        return indices
+
+    def detokenize(self, indices: mx.array) -> mx.array:
+        zq = self.quantizer.get_output_from_indices(indices.swapaxes(-1, -2)).swapaxes(
+            -1, -2
+        )
+        x = zq.reshape(zq.shape[0], -1)
+        d_vector = self.project(x)
+        return d_vector
+
+    def sanitize(self, weights):
+        sanitized_weights = {}
+        for k, v in weights.items():
+            if (
+                ".conv.weight" in k
+                or ("convs." in k and "weight" in k)
+                or ("speaker_encoder.pool.linear" in k and "weight" in k)
+            ):
+                if v.shape[1] > v.shape[-1]:
+                    sanitized_weights[k] = v.transpose(0, 2, 1)
+                else:
+                    sanitized_weights[k] = v
+            else:
+                sanitized_weights[k] = v
+        return sanitized_weights
+
+
+if __name__ == "__main__":
+    from mlx.utils import tree_flatten
+
+    model = SpeakerEncoder(
+        input_dim=100,
+        latent_dim=128,
+        token_num=32,
+        fsq_levels=[4, 4, 4, 4, 4, 4],
+        fsq_num_quantizers=1,
+    )
+    mel = mx.random.normal(shape=(8, 200, 100), scale=1.0)
+    x_vector, d_vector = model(mel)
+    print("x-vector shape", x_vector.shape)
+    print("d-vector shape", d_vector.shape)
+
+    indices = model.tokenize(mel)
+    print("indices shape", indices.shape)
+    d_vector_post = model.detokenize(indices)
+    print("d-vector shape", d_vector_post.shape)
+    if d_vector_post.all() == d_vector.all():
+        print("d-vector post and d-vector are the same")
+    else:
+        print("d-vector post and d-vector are different")
+
+    num_params = 0
+
+    weights = dict(tree_flatten(model.parameters()))
+
+    for k, v in weights.items():
+        num_params += v.size
+    print("{} M".format(num_params / 1e6))