nexaai 1.0.29__cp310-cp310-macosx_14_0_universal2.whl

nexaai/binds/metal/py-lib/mlx_audio/tts/models/sesame/watermarking.py

@@ -0,0 +1,105 @@
+import argparse
+
+import mlx.core as mx
+import numpy as np
+import silentcipher
+import soundfile as sf
+from scipy import signal
+
+# This watermark key is public, it is not secure.
+# If using CSM 1B in another application, use a new private key and keep it secret.
+CSM_1B_GH_WATERMARK = [212, 211, 146, 56, 201]
+
+
+def cli_check_audio() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--audio_path", type=str, required=True)
+    args = parser.parse_args()
+    check_audio_from_file(args.audio_path)
+
+
+def load_watermarker() -> silentcipher.server.Model:
+    model = silentcipher.get_model(
+        model_type="44.1k",
+    )
+    return model
+
+
+def resample_audio(audio: np.ndarray, orig_sr: int, target_sr: int) -> np.ndarray:
+    gcd = np.gcd(orig_sr, target_sr)
+    up = target_sr // gcd
+    down = orig_sr // gcd
+    resampled = signal.resample_poly(audio, up, down, padtype="edge")
+    return resampled
+
+
+def watermark(
+    watermarker: silentcipher.server.Model,
+    audio_array: mx.array,
+    sample_rate: int,
+    watermark_key: list[int],
+) -> tuple[mx.array, int]:
+    audio_array = np.array(audio_array, dtype=np.float32)
+
+    if sample_rate != 44100:
+        audio_array_44khz = resample_audio(audio_array, sample_rate, 44100)
+    else:
+        audio_array_44khz = audio_array
+
+    encoded, *_ = watermarker.encode_wav(
+        audio_array_44khz, 44100, watermark_key, calc_sdr=False, message_sdr=36
+    )
+
+    if sample_rate != 44100:
+        encoded = resample_audio(encoded, 44100, sample_rate)
+
+    return encoded
+
+
+def verify(
+    watermarker: silentcipher.server.Model,
+    watermarked_audio: mx.array,
+    sample_rate: int,
+    watermark_key: list[int],
+) -> bool:
+    if sample_rate != 44100:
+        watermarked_audio_44khz = resample_audio(watermarked_audio, sample_rate, 44100)
+    else:
+        watermarked_audio_44khz = watermarked_audio
+
+    result = watermarker.decode_wav(
+        watermarked_audio_44khz, 44100, phase_shift_decoding=True
+    )
+
+    is_watermarked = result["status"]
+    if is_watermarked:
+        is_csm_watermarked = result["messages"][0] == watermark_key
+    else:
+        is_csm_watermarked = False
+
+    return is_watermarked and is_csm_watermarked
+
+
+def check_audio_from_file(audio_path: str) -> None:
+    watermarker = load_watermarker()
+    audio_array, sample_rate = load_audio(audio_path)
+    is_watermarked = verify(watermarker, audio_array, sample_rate, CSM_1B_GH_WATERMARK)
+    outcome = "Watermarked" if is_watermarked else "Not watermarked"
+    print(f"{outcome}: {audio_path}")
+
+
+def load_audio(audio_path: str) -> tuple[mx.array, int]:
+    audio_array_np, sample_rate = sf.read(audio_path, always_2d=True)
+
+    if audio_array_np.shape[1] > 1:
+        audio_array_np = audio_array_np.mean(axis=1)
+    else:
+        audio_array_np = audio_array_np.squeeze()
+
+    audio_array = mx.array(audio_array_np)
+
+    return audio_array, int(sample_rate)
+
+
+if __name__ == "__main__":
+    cli_check_audio()

nexaai/binds/metal/py-lib/mlx_audio/tts/models/spark/__init__.py

@@ -0,0 +1 @@
+from .spark import Model, ModelConfig

nexaai/binds/metal/py-lib/mlx_audio/tts/models/spark/audio_tokenizer.py

@@ -0,0 +1,138 @@
+from pathlib import Path
+from typing import Any, Dict, Tuple, Union
+
+import mlx.core as mx
+import numpy as np
+
+from mlx_audio.stt.models.wav2vec.feature_extractor import Wav2Vec2FeatureExtractor
+from mlx_audio.stt.models.wav2vec.wav2vec import Wav2Vec2Model
+
+from .bicodec import BiCodec
+from .utils.audio import load_audio
+from .utils.file import load_config
+
+
+class BiCodecTokenizer:
+    """BiCodec tokenizer for handling audio input and tokenization."""
+
+    def __init__(self, model_dir: Path, **kwargs):
+        super().__init__()
+        """
+        Args:
+            model_dir: Path to the model directory.
+            device: Device to run the model on (default is GPU if available).
+        """
+        self.model_dir = model_dir
+        self.config = load_config(f"{model_dir}/audio_tokenizer_config.yaml")
+        self._initialize_model()
+
+    def _initialize_model(self):
+        """Load and initialize the BiCodec model and Wav2Vec2 feature extractor."""
+        self.model = BiCodec.load_from_checkpoint(f"{self.model_dir}/BiCodec")
+        self.processor = Wav2Vec2FeatureExtractor.from_pretrained(
+            f"{self.model_dir}/wav2vec2-large-xlsr-53"
+        )
+        self.feature_extractor = Wav2Vec2Model.from_pretrained(
+            f"{self.model_dir}/wav2vec2-large-xlsr-53"
+        )
+        self.feature_extractor.config.output_hidden_states = True
+
+    def get_ref_clip(self, wav: np.ndarray) -> np.ndarray:
+        """Get reference audio clip for speaker embedding."""
+        ref_segment_length = (
+            int(self.config["sample_rate"] * self.config["ref_segment_duration"])
+            // self.config["latent_hop_length"]
+            * self.config["latent_hop_length"]
+        )
+        wav_length = len(wav)
+
+        if ref_segment_length > wav_length:
+            # Repeat and truncate to handle insufficient length
+            wav = np.tile(wav, ref_segment_length // wav_length + 1)
+
+        return wav[:ref_segment_length]
+
+    def process_audio(
+        self, wav_path: Union[Path, mx.array]
+    ) -> Tuple[np.ndarray, mx.array]:
+        """load auido and get reference audio from wav path"""
+        if isinstance(wav_path, Path) or isinstance(wav_path, str):
+            wav = load_audio(
+                wav_path,
+                sampling_rate=self.config["sample_rate"],
+                volume_normalize=self.config["volume_normalize"],
+            )
+        elif isinstance(wav_path, mx.array):
+            wav = wav_path
+        else:
+            raise ValueError(f"Invalid input type: {type(wav_path)}")
+
+        wav_ref = self.get_ref_clip(wav)
+
+        return wav, wav_ref[None, ...]
+
+    def extract_wav2vec2_features(self, wavs: mx.array) -> mx.array:
+        """extract wav2vec2 features"""
+        inputs = self.processor(
+            wavs,
+            sampling_rate=16000,
+            return_tensors="mx",
+            padding=True,
+            output_hidden_states=True,
+        )["input_values"]
+        feat = self.feature_extractor(inputs)
+        feats_mix = (
+            feat.hidden_states[11] + feat.hidden_states[14] + feat.hidden_states[16]
+        ) / 3
+
+        return feats_mix
+
+    def tokenize_batch(self, batch: Dict[str, Any]) -> Tuple[mx.array, mx.array]:
+        """tokenize the batch of audio
+
+        Args:
+            batch:
+                wavs (List[np.ndarray]): batch of audio
+                ref_wavs (mx.array): reference audio. shape: (batch_size, seq_len)
+
+        Returns:
+            semantic_tokens: semantic tokens. shape: (batch_size, seq_len, latent_dim)
+            global_tokens: global tokens. shape: (batch_size, seq_len, global_dim)
+        """
+        feats = self.extract_wav2vec2_features(batch["wav"])
+        batch["feat"] = feats
+
+        semantic_tokens, global_tokens = self.model.tokenize(batch)
+
+        return global_tokens, semantic_tokens
+
+    def tokenize(self, audio_path: str) -> Tuple[mx.array, mx.array]:
+        """tokenize the audio"""
+        wav, ref_wav = self.process_audio(audio_path)
+        feat = self.extract_wav2vec2_features(wav)
+        batch = {
+            "wav": wav[None, ...],
+            "ref_wav": ref_wav,
+            "feat": feat,
+        }
+        semantic_tokens, global_tokens = self.model.tokenize(batch)
+
+        return global_tokens, semantic_tokens
+
+    def detokenize(
+        self, global_tokens: mx.array, semantic_tokens: mx.array
+    ) -> np.array:
+        """detokenize the tokens to waveform
+
+        Args:
+            global_tokens: global tokens. shape: (batch_size, global_dim)
+            semantic_tokens: semantic tokens. shape: (batch_size, latent_dim)
+
+        Returns:
+            wav_rec: waveform. shape: (batch_size, seq_len) for batch or (seq_len,) for single
+        """
+        global_tokens = mx.expand_dims(global_tokens, 1)
+
+        # convert to mlx array
+        wav_rec = self.model.detokenize(semantic_tokens, global_tokens)
+        return wav_rec.squeeze()

nexaai/binds/metal/py-lib/mlx_audio/tts/models/spark/bicodec.py

@@ -0,0 +1,269 @@
+from pathlib import Path
+from typing import Any, Dict, Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+import numpy as np
+from omegaconf import DictConfig
+from safetensors.torch import load_file
+
+from mlx_audio.tts.models.spark.modules.encoder_decoder.feat_decoder import Decoder
+from mlx_audio.tts.models.spark.modules.encoder_decoder.feat_encoder import Encoder
+from mlx_audio.tts.models.spark.modules.encoder_decoder.wave_generator import (
+    WaveGenerator,
+)
+from mlx_audio.tts.models.spark.modules.residual import FactorizedVectorQuantize
+from mlx_audio.tts.models.spark.modules.speaker.speaker_encoder import SpeakerEncoder
+from mlx_audio.tts.models.spark.utils.file import load_config
+from mlx_audio.tts.utils import get_model_path
+from mlx_audio.utils import hanning, mel_filters, stft
+
+
+def mel_spectrogram(
+    audio: mx.array,
+    sample_rate: int = 16_000,
+    n_mels: int = 128,
+    n_fft: int = 1024,
+    f_min: int = 10,
+    f_max: Optional[int] = None,
+    hop_length: int = 320,
+    win_length: int = 640,
+    padding: int = 0,
+):
+    if not isinstance(audio, mx.array):
+        audio = mx.array(audio)
+    if padding > 0:
+        audio = mx.pad(audio, (0, padding))
+    window = hanning(win_length + 1)[:-1]
+    freqs = stft(
+        audio, window=window, win_length=win_length, hop_length=hop_length, n_fft=n_fft
+    )
+    magnitudes = freqs.abs()
+    filters = mel_filters(
+        sample_rate=sample_rate,
+        n_fft=n_fft,
+        n_mels=n_mels,
+        f_min=f_min,
+        f_max=f_max,
+        norm="slaney",
+        mel_scale="slaney",
+    )
+    mel_spec = magnitudes @ filters.T
+    return mx.expand_dims(mel_spec, axis=0)
+
+
+class BiCodec(nn.Module):
+    """
+    BiCodec model for speech synthesis, incorporating a speaker encoder, feature encoder/decoder,
+    quantizer, and wave generator.
+    """
+
+    def __init__(
+        self,
+        mel_params: Dict[str, Any],
+        encoder: nn.Module,
+        decoder: nn.Module,
+        quantizer: nn.Module,
+        speaker_encoder: nn.Module,
+        prenet: nn.Module,
+        postnet: nn.Module,
+        **kwargs,
+    ) -> None:
+        """
+        Initializes the BiCodec model with the required components.
+
+        Args:
+            mel_params (dict): Parameters for the mel-spectrogram transformer.
+            encoder (nn.Module): Encoder module.
+            decoder (nn.Module): Decoder module.
+            quantizer (nn.Module): Quantizer module.
+            speaker_encoder (nn.Module): Speaker encoder module.
+            prenet (nn.Module): Prenet network.
+            postnet (nn.Module): Postnet network.
+        """
+        super().__init__()
+        self.encoder = encoder
+        self.decoder = decoder
+        self.quantizer = quantizer
+        self.speaker_encoder = speaker_encoder
+        self.prenet = prenet
+        self.postnet = postnet
+        self.mel_params = mel_params
+
+    @classmethod
+    def load_from_checkpoint(cls, model_dir: Path, **kwargs) -> "BiCodec":
+        """
+        Loads the model from a checkpoint.
+
+        Args:
+            model_dir (Path): Path to the model directory containing checkpoint and config.
+
+        Returns:
+            BiCodec: The initialized BiCodec model.
+        """
+        ckpt_path = f"{model_dir}/model.safetensors"
+        config = load_config(f"{model_dir}/config.yaml")["audio_tokenizer"]
+        mel_params = config["mel_params"]
+
+        encoder = Encoder(**config["encoder"])
+        quantizer = FactorizedVectorQuantize(**config["quantizer"])
+        prenet = Decoder(**config["prenet"])
+        postnet = Decoder(**config["postnet"])
+        decoder = WaveGenerator(**config["decoder"])
+        speaker_encoder = SpeakerEncoder(**config["speaker_encoder"])
+
+        model = cls(
+            mel_params=mel_params,
+            encoder=encoder,
+            decoder=decoder,
+            quantizer=quantizer,
+            speaker_encoder=speaker_encoder,
+            prenet=prenet,
+            postnet=postnet,
+        )
+
+        weights = load_file(ckpt_path)
+
+        # Convert PyTorch weights to MLX arrays and sanitize
+        weights = {
+            k: mx.array(v) for k, v in weights.items() if "num_batches_tracked" not in k
+        }
+
+        for module in [encoder, decoder, quantizer, speaker_encoder]:
+            if hasattr(module, "sanitize"):
+                weights = module.sanitize(weights)
+
+        model.load_weights(list(weights.items()), strict=True)
+
+        return model
+
+    def __call__(self, batch: Dict[str, Any]) -> Dict[str, Any]:
+        """
+        Performs a forward pass through the model.
+
+        Args:
+            batch (dict): A dictionary containing features, reference waveform, and target waveform.
+
+        Returns:
+            dict: A dictionary containing the reconstruction, features, and other metrics.
+        """
+        feat = batch["feat"]
+        ref_wav = batch["ref_wav"]
+        mel = self.get_mel_spectrogram(ref_wav)
+        z = self.encoder(feat.transpose(0, 2, 1))
+        vq_outputs = self.quantizer(z)
+
+        x_vector, d_vector = self.speaker_encoder(mel)
+
+        conditions = d_vector
+        with_speaker_loss = False
+
+        # Ensure conditions is an integer type for embedding lookup
+        # The error shows that the embedding layer expects integral indices
+        if isinstance(conditions, mx.array) and conditions.dtype == mx.float32:
+            # Convert to integer type if needed for the embedding layer
+            # or ensure it's properly formatted for the prenet
+            conditions = conditions.astype(mx.int32)
+
+        x = self.prenet(vq_outputs["z_q"], conditions)
+        pred_feat = self.postnet(x)
+        x = x + conditions[..., None]
+        wav_recon = self.decoder(x)
+
+        return {
+            "vq_loss": vq_outputs["vq_loss"],
+            "perplexity": vq_outputs["perplexity"],
+            "cluster_size": vq_outputs["active_num"],
+            "recons": wav_recon,
+            "pred_feat": pred_feat,
+            "x_vector": x_vector,
+            "d_vector": d_vector,
+            "audios": batch["wav"][:, None],
+            "with_speaker_loss": with_speaker_loss,
+        }
+
+    def tokenize(self, batch: Dict[str, Any]):
+        """
+        Tokenizes the input audio into semantic and global tokens.
+
+        Args:
+            batch (dict): The input audio features and reference waveform.
+
+        Returns:
+            tuple: Semantic tokens and global tokens.
+        """
+        feat = batch["feat"]
+        ref_wav = mx.array(batch["ref_wav"])
+        mel = self.get_mel_spectrogram(ref_wav)
+        z = self.encoder(feat.transpose(0, 2, 1))
+        semantic_tokens = self.quantizer.tokenize(z)
+        global_tokens = self.speaker_encoder.tokenize(mel)
+
+        return semantic_tokens, global_tokens
+
+    def detokenize(self, semantic_tokens, global_tokens):
+        """
+        Detokenizes the semantic and global tokens into a waveform.
+
+        Args:
+            semantic_tokens (tensor): Semantic tokens.
+            global_tokens (tensor): Global tokens.
+
+        Returns:
+            tensor: Reconstructed waveform.
+        """
+
+        z_q = self.quantizer.detokenize(semantic_tokens.transpose(0, 1)).transpose(
+            0, 2, 1
+        )
+        d_vector = self.speaker_encoder.detokenize(global_tokens)
+        x = self.prenet(z_q, d_vector)
+        x = x + d_vector[..., None]
+        wav_recon = self.decoder(x)
+
+        return wav_recon  # Return MLX array directly
+
+    def get_mel_spectrogram(self, wav):
+        mels = []
+        for i in range(wav.shape[0]):
+            audio_sample = mx.squeeze(wav[i])
+            mel = mel_spectrogram(
+                audio=audio_sample,
+                sample_rate=self.mel_params["sample_rate"],
+                n_mels=self.mel_params["num_mels"],
+                n_fft=self.mel_params["n_fft"],
+                hop_length=self.mel_params["hop_length"],
+                win_length=self.mel_params["win_length"],
+                f_min=self.mel_params["mel_fmin"],
+                f_max=self.mel_params["mel_fmax"],
+            )
+            mels.append(mel)
+        return mx.concatenate(mels, axis=0)
+
+
+if __name__ == "__main__":
+
+    model_path = get_model_path("SparkAudio/Spark-TTS-0.5B")
+
+    model = BiCodec.load_from_checkpoint(model_path / "BiCodec")
+    model.eval()
+
+    # Generate random inputs for testing
+    duration = 0.96
+    x = mx.random.normal((20, 1, int(duration * 16000)), dtype=mx.float32)
+    feat = mx.random.normal((20, int(duration * 50), 1024), dtype=mx.float32)
+    inputs = {"feat": feat, "wav": x, "ref_wav": x}
+
+    # Forward pass
+    outputs = model(inputs)
+    semantic_tokens, global_tokens = model.tokenize(inputs)
+
+    wav_recon = model.detokenize(semantic_tokens, global_tokens)
+
+    print(outputs["recons"].shape)
+    print(wav_recon.shape)
+
+    if np.allclose(outputs["recons"], wav_recon):
+        print("Test successful")
+    else:
+        print("Test failed")

nexaai/binds/metal/py-lib/mlx_audio/tts/models/spark/modules/blocks/sampler.py

@@ -0,0 +1,111 @@
+import math
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from mlx_audio.codec.models.descript.nn.layers import WNConvTranspose1d
+
+
+class SamplingBlock(nn.Module):
+    """Sampling block for upsampling or downsampling"""
+
+    def __init__(
+        self,
+        dim: int,
+        groups: int = 1,
+        upsample_scale: int = 1,
+        downsample_scale: int = 1,
+    ) -> None:
+        """
+        Args:
+            dim: input dimension
+            groups: number of groups
+            upsample_scale: upsampling scale
+            downsample_scale: downsampling scale
+        """
+        super(SamplingBlock, self).__init__()
+
+        self.upsample_scale = upsample_scale
+        self.downsample_scale = downsample_scale
+
+        if self.upsample_scale > 1:
+            self.de_conv_upsampler = nn.Sequential(
+                nn.LeakyReLU(0.2),
+                WNConvTranspose1d(
+                    dim,
+                    dim,
+                    kernel_size=upsample_scale * 2,
+                    stride=upsample_scale,
+                    padding=upsample_scale // 2 + upsample_scale % 2,
+                    groups=groups,
+                ),
+            )
+
+        if self.downsample_scale > 1:
+            self.conv_downsampler = nn.Sequential(
+                nn.LeakyReLU(0.2),
+                nn.Conv1d(
+                    dim,
+                    dim,
+                    kernel_size=2 * downsample_scale,
+                    stride=downsample_scale,
+                    padding=downsample_scale // 2 + downsample_scale % 2,
+                    groups=groups,
+                ),
+            )
+
+    @staticmethod
+    def repeat_upsampler(x, upsample_scale):
+        # MLX doesn't have repeat_interleave, so we need to implement it manually
+        batch_size, seq_len, channels = x.shape
+        # Create a new tensor with the expanded shape
+        output = mx.zeros((batch_size, seq_len * upsample_scale, channels))
+        # Fill the output tensor by repeating each element
+        for i in range(seq_len):
+            for j in range(upsample_scale):
+                output[:, i * upsample_scale + j, :] = x[:, i, :]
+        return output
+
+    @staticmethod
+    def skip_downsampler(x, downsample_scale):
+        return nn.AvgPool1d(kernel_size=downsample_scale, stride=downsample_scale)(x)
+
+    def __call__(self, x):
+        x = x.transpose(0, 2, 1)
+        if self.upsample_scale > 1:
+            repeat_res = self.repeat_upsampler(x, self.upsample_scale)
+            deconv_res = self.de_conv_upsampler(x)
+            upmerge_res = repeat_res + deconv_res
+        else:
+            upmerge_res = x
+            repeat_res = x
+
+        if self.downsample_scale > 1:
+            conv_res = self.conv_downsampler(upmerge_res)
+            skip2_res = self.skip_downsampler(upmerge_res, self.downsample_scale)
+            skip1_res = self.skip_downsampler(repeat_res, self.downsample_scale)
+        else:
+            conv_res = upmerge_res
+            skip2_res = upmerge_res
+            skip1_res = repeat_res
+
+        final_res = conv_res + skip1_res + skip2_res
+
+        return final_res.transpose(0, 2, 1)
+
+
+# test
+if __name__ == "__main__":
+    test_input = mx.random.randint(
+        0, 100, (8, 1024, 50)
+    )  # Batch size = 8, 1024 channels, length = 50
+    model = SamplingBlock(1024, 1024, upsample_scale=2)
+    model_down = SamplingBlock(1024, 1024, downsample_scale=2)
+    output = model(test_input)
+    output_down = model_down(test_input)
+    print("shape after upsample * 2", output.shape)  # torch.Size([8, 1024, 100])
+    print("shape after downsample * 2", output_down.shape)  # torch.Size([8, 1024, 25])
+    if output.shape == (8, 1024, 100) and output_down.shape == (8, 1024, 25):
+        print("test successful")
+    else:
+        print("test failed")