nexaai 1.0.29__cp310-cp310-macosx_14_0_universal2.whl

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

@@ -0,0 +1,382 @@
+import re
+import time
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Dict, Optional, Tuple, Union
+
+import mlx.core as mx
+import mlx.nn as nn
+from mlx_lm.generate import stream_generate
+from mlx_lm.models.qwen2 import Model as Qwen2Model
+from mlx_lm.sample_utils import make_logits_processors, make_sampler
+from mlx_lm.tokenizer_utils import load_tokenizer
+from tqdm import tqdm
+
+from mlx_audio.tts.models.base import BaseModelArgs, GenerationResult
+from mlx_audio.tts.utils import get_model_path
+
+from .audio_tokenizer import BiCodecTokenizer
+from .utils.token_parser import GENDER_MAP, LEVELS_MAP, TASK_TOKEN_MAP
+
+PITCH_MAP = SPEED_MAP = {
+    0.0: "very_low",
+    0.5: "low",
+    1.0: "moderate",
+    1.5: "high",
+    2.0: "very_high",
+}
+
+
+@dataclass
+class ModelConfig(BaseModelArgs):
+    model_path: Path = None
+    sample_rate: int = 16000
+    bos_token_id: int = 151643
+    eos_token_id: int = 151645
+    hidden_act: str = "silu"
+    hidden_size: int = 896
+    initializer_range: float = 0.02
+    intermediate_size: int = 4864
+    max_position_embeddings: int = 32768
+    max_window_layers: int = 21
+    model_type: str = "qwen2"
+    num_attention_heads: int = 14
+    num_hidden_layers: int = 24
+    num_key_value_heads: int = 2
+    rms_norm_eps: float = 1e-06
+    rope_theta: float = 1000000.0
+    sliding_window: int = 32768
+    tie_word_embeddings: bool = True
+    torch_dtype: str = "bfloat16"
+    transformers_version: str = "4.43.1"
+    use_sliding_window: bool = False
+    vocab_size: int = 166000
+    rope_traditional: bool = False
+    rope_scaling: Optional[Dict[str, Union[float, str]]] = None
+
+
+class Model(nn.Module):
+    """
+    Spark-TTS for text-to-speech generation.
+    """
+
+    def __init__(self, config: ModelConfig):
+        """
+        Initializes the SparkTTS model with the provided configurations and device.
+
+        Args:
+            config (ModelConfig): The configuration for the model.
+        """
+        self.config = config
+
+        model_dir = config.model_path
+
+        self.model = Qwen2Model(config)
+        self.tokenizer = load_tokenizer(model_dir, eos_token_ids=config.eos_token_id)
+
+        self._audio_tokenizer = BiCodecTokenizer(model_dir)
+
+    def load_weights(self, weights, strict=True):
+        self.model.load_weights(weights, strict=strict)
+
+    def parameters(self):
+        return self.model.parameters()
+
+    def model_type(self):
+        return "spark"
+
+    def sanitize(self, weights):
+        return self.model.sanitize(weights)
+
+    @property
+    def sample_rate(self):
+        return self.config.sample_rate
+
+    @property
+    def layers(self):
+        return self.model.layers
+
+    def model_quant_predicate(self, p, m, config):
+        """
+        Model modules to skip during quantization
+        """
+        return not p.startswith("_audio_tokenizer")
+
+    def process_prompt(
+        self,
+        text: str,
+        ref_audio: Path,
+        ref_text: str,
+    ) -> Tuple[str, mx.array]:
+        """
+        Process input for voice cloning.
+
+        Args:
+            text (str): The text input to be converted to speech.
+            ref_audio (Path): Path to the audio file used as a reference.
+            ref_text (str, optional): Transcript of the reference audio.
+
+        Return:
+            Tuple[str, mx.array]: Input prompt; global tokens
+        """
+
+        global_token_ids, semantic_token_ids = self._audio_tokenizer.tokenize(ref_audio)
+        global_tokens = "".join(
+            [f"<|bicodec_global_{i}|>" for i in global_token_ids.squeeze()]
+        )
+
+        # Prepare the input tokens for the model
+        if ref_text is not None:
+            semantic_tokens = "".join(
+                [f"<|bicodec_semantic_{i}|>" for i in semantic_token_ids.squeeze()]
+            )
+            inputs = [
+                TASK_TOKEN_MAP["tts"],
+                "<|start_content|>",
+                ref_text,
+                text,
+                "<|end_content|>",
+                "<|start_global_token|>",
+                global_tokens,
+                "<|end_global_token|>",
+                "<|start_semantic_token|>",
+                semantic_tokens,
+            ]
+        else:
+            inputs = [
+                TASK_TOKEN_MAP["tts"],
+                "<|start_content|>",
+                text,
+                "<|end_content|>",
+                "<|start_global_token|>",
+                global_tokens,
+                "<|end_global_token|>",
+            ]
+
+        inputs = "".join(inputs)
+
+        return inputs, global_token_ids
+
+    def process_prompt_control(
+        self,
+        gender: str,
+        pitch: str,
+        speed: str,
+        text: str,
+    ):
+        """
+        Process input for voice creation.
+
+        Args:
+            gender (str): female | male.
+            pitch (str): very_low | low | moderate | high | very_high
+            speed (str): very_low | low | moderate | high | very_high
+            text (str): The text input to be converted to speech.
+
+        Return:
+            str: Input prompt
+        """
+        assert gender in GENDER_MAP.keys()
+        assert pitch in LEVELS_MAP.keys()
+        assert speed in LEVELS_MAP.keys()
+
+        gender_id = GENDER_MAP[gender]
+        pitch_level_id = LEVELS_MAP[pitch]
+        speed_level_id = LEVELS_MAP[speed]
+
+        pitch_label_tokens = f"<|pitch_label_{pitch_level_id}|>"
+        speed_label_tokens = f"<|speed_label_{speed_level_id}|>"
+        gender_tokens = f"<|gender_{gender_id}|>"
+
+        attribte_tokens = "".join(
+            [gender_tokens, pitch_label_tokens, speed_label_tokens]
+        )
+
+        control_tts_inputs = [
+            TASK_TOKEN_MAP["controllable_tts"],
+            "<|start_content|>",
+            text,
+            "<|end_content|>",
+            "<|start_style_label|>",
+            attribte_tokens,
+            "<|end_style_label|>",
+        ]
+
+        return "".join(control_tts_inputs)
+
+    def generate(
+        self,
+        text: str,
+        ref_audio: Path = None,
+        ref_text: str = None,
+        gender: str = "male",
+        pitch: float = 1.0,
+        speed: float = 1.0,
+        temperature: float = 0.8,
+        top_k: float = 50,
+        top_p: float = 0.95,
+        max_tokens: int = 3000,
+        verbose: bool = False,
+        split_pattern: str = "\n",
+        **kwargs,
+    ) -> GenerationResult:
+        """
+        Performs inference to generate speech from text, incorporating prompt audio and/or text.
+
+        Args:
+            text (str): The text input to be converted to speech.
+            ref_audio (Path): Path to the audio file used as a reference.
+            ref_text (str, optional): Transcript of the reference audio.
+            gender (str): female | male.
+            pitch (str): very_low | low | moderate | high | very_high
+            speed (str): very_low | low | moderate | high | very_high
+            temperature (float, optional): Sampling temperature for controlling randomness. Default is 0.8.
+            top_k (float, optional): Top-k sampling parameter. Default is 50.
+            top_p (float, optional): Top-p (nucleus) sampling parameter. Default is 0.95.
+
+        Returns:
+            GenerationResult: Generated waveform as a tensor.
+        """
+
+        speed_factor = SPEED_MAP[speed]
+        pitch_factor = PITCH_MAP[pitch]
+
+        if ref_audio is not None:  # voice cloning
+            gender = None
+
+        text_splits = text.split(split_pattern)
+
+        for text_split in text_splits:
+            if gender is not None:
+                prompt = self.process_prompt_control(
+                    gender, pitch_factor, speed_factor, text_split
+                )
+
+            else:
+                prompt, global_token_ids = self.process_prompt(
+                    text_split, ref_audio, ref_text
+                )
+
+            inputs = self.tokenizer._tokenizer([prompt], return_tensors="pt")
+
+            input_ids = mx.array(inputs.input_ids)
+
+            sampler = make_sampler(temperature, top_p=top_p, top_k=top_k)
+            logits_processors = make_logits_processors(
+                kwargs.get("logit_bias", None),
+                kwargs.get("repetition_penalty", 1.3),
+                kwargs.get("repetition_context_size", 20),
+            )
+
+            time_start = time.time()
+
+            generated_ids = []
+
+            # Generate speech using the model
+            for i, response in enumerate(
+                tqdm(
+                    stream_generate(
+                        self.model,
+                        tokenizer=self.tokenizer,
+                        prompt=input_ids.squeeze(0),
+                        max_tokens=max_tokens,
+                        sampler=sampler,
+                        logits_processors=logits_processors,
+                    ),
+                    total=max_tokens,
+                    disable=not verbose,
+                )
+            ):
+                next_token = mx.array([response.token])
+                input_ids = mx.concatenate([input_ids, next_token[None, :]], axis=1)
+                if i % 50 == 0:
+                    mx.clear_cache()
+
+                if next_token == 128258:
+                    break
+
+            time_end = time.time()
+            # Trim the output tokens to remove the input tokens
+            generated_ids = mx.array(
+                [
+                    output[len(input) :]
+                    for input, output in zip(inputs.input_ids, input_ids)
+                ]
+            ).tolist()
+
+            # Decode the generated tokens into text
+            predicts = self.tokenizer._tokenizer.batch_decode(
+                generated_ids, skip_special_tokens=True
+            )[0]
+
+            # Extract semantic token IDs from the generated text
+            pred_semantic_ids = mx.array(
+                [
+                    int(token)
+                    for token in re.findall(r"bicodec_semantic_(\d+)", predicts)
+                ]
+            )[None, ...]
+
+            if gender is not None:
+                global_token_ids = mx.array(
+                    [
+                        int(token)
+                        for token in re.findall(r"bicodec_global_(\d+)", predicts)
+                    ]
+                )[None, ...]
+
+            # Convert semantic tokens back to waveform
+            audio = self._audio_tokenizer.detokenize(
+                global_token_ids.astype(mx.int32),
+                pred_semantic_ids.astype(mx.int32),
+            )
+
+            # Clear cache
+            mx.clear_cache()
+
+            audio_samples = len(audio)
+            audio_duration_seconds = audio_samples / self.config.sample_rate
+
+            # Format duration as HH:MM:SS.mmm
+            duration_mins = int(audio_duration_seconds // 60)
+            duration_secs = int(audio_duration_seconds % 60)
+            duration_ms = int((audio_duration_seconds % 1) * 1000)
+            duration_hours = int(audio_duration_seconds // 3600)
+            duration_str = f"{duration_hours:02d}:{duration_mins:02d}:{duration_secs:02d}.{duration_ms:03d}"
+
+            yield GenerationResult(
+                audio=audio,
+                sample_rate=self.config.sample_rate,
+                samples=audio_samples,
+                segment_idx=0,  # Default segment index
+                token_count=len(pred_semantic_ids.squeeze()),
+                audio_samples={
+                    "samples": audio_samples,
+                    "samples-per-sec": (
+                        round(audio_samples / audio_duration_seconds, 2)
+                        if audio_duration_seconds > 0
+                        else 0
+                    ),
+                },
+                audio_duration=duration_str,
+                real_time_factor=(
+                    audio_duration_seconds / (time_end - time_start)
+                    if (time_end - time_start) > 0
+                    else 0
+                ),
+                prompt={
+                    "tokens": len(pred_semantic_ids.squeeze()),
+                    "tokens-per-sec": (
+                        round(
+                            len(pred_semantic_ids.squeeze()) / audio_duration_seconds, 2
+                        )
+                        if audio_duration_seconds > 0
+                        else 0
+                    ),
+                },
+                processing_time_seconds=time_end - time_start,
+                peak_memory_usage=mx.get_peak_memory() / 1e9,
+            )
+
+            # Clear cache after each segment to avoid memory leaks
+            mx.clear_cache()

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

@@ -0,0 +1,220 @@
+# 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.
+"""
+Description:
+    This script contains a collection of functions designed to handle various
+    audio processing.
+"""
+
+import random
+from pathlib import Path
+from typing import Tuple
+
+import mlx.core as mx
+import numpy as np
+import soundfile
+import soxr
+from numpy.lib.stride_tricks import sliding_window_view
+
+
+def audio_volume_normalize(audio: np.ndarray, coeff: float = 0.2) -> np.ndarray:
+    """
+    Normalize the volume of an audio signal.
+
+    Parameters:
+        audio (numpy array): Input audio signal array.
+        coeff (float): Target coefficient for normalization, default is 0.2.
+
+    Returns:
+        numpy array: The volume-normalized audio signal.
+    """
+    # Sort the absolute values of the audio signal
+    temp = np.sort(np.abs(audio))
+
+    # If the maximum value is less than 0.1, scale the array to have a maximum of 0.1
+    if temp[-1] < 0.1:
+        scaling_factor = max(
+            temp[-1], 1e-3
+        )  # Prevent division by zero with a small constant
+        audio = audio / scaling_factor * 0.1
+
+    # Filter out values less than 0.01 from temp
+    temp = temp[temp > 0.01]
+    L = temp.shape[0]  # Length of the filtered array
+
+    # If there are fewer than or equal to 10 significant values, return the audio without further processing
+    if L <= 10:
+        return audio
+
+    # Compute the average of the top 10% to 1% of values in temp
+    volume = np.mean(temp[int(0.9 * L) : int(0.99 * L)])
+
+    # Normalize the audio to the target coefficient level, clamping the scale factor between 0.1 and 10
+    audio = audio * np.clip(coeff / volume, a_min=0.1, a_max=10)
+
+    # Ensure the maximum absolute value in the audio does not exceed 1
+    max_value = np.max(np.abs(audio))
+    if max_value > 1:
+        audio = audio / max_value
+
+    return audio
+
+
+def load_audio(
+    adfile: Path,
+    sampling_rate: int = None,
+    length: int = None,
+    volume_normalize: bool = False,
+    segment_duration: int = None,
+) -> np.ndarray:
+    r"""Load audio file with target sampling rate and lsength
+
+    Args:
+        adfile (Path): path to audio file.
+        sampling_rate (int, optional): target sampling rate. Defaults to None.
+        length (int, optional): target audio length. Defaults to None.
+        volume_normalize (bool, optional): whether perform volume normalization. Defaults to False.
+        segment_duration (int): random select a segment with duration of {segment_duration}s.
+                                Defualt to None which means the whole audio will be used.
+
+    Returns:
+        audio (np.ndarray): audio
+    """
+
+    audio, sr = soundfile.read(adfile)
+    if len(audio.shape) > 1:
+        audio = audio[:, 0]
+
+    if sampling_rate is not None and sr != sampling_rate:
+        audio = soxr.resample(audio, sr, sampling_rate, quality="VHQ")
+        sr = sampling_rate
+
+    if segment_duration is not None:
+        seg_length = int(sr * segment_duration)
+        audio = random_select_audio_segment(audio, seg_length)
+
+    # Audio volume normalize
+    if volume_normalize:
+        audio = audio_volume_normalize(audio)
+    # check the audio length
+    if length is not None:
+        assert abs(audio.shape[0] - length) < 1000
+        if audio.shape[0] > length:
+            audio = audio[:length]
+        else:
+            audio = np.pad(audio, (0, int(length - audio.shape[0])))
+    return audio
+
+
+def random_select_audio_segment(audio: np.ndarray, length: int) -> np.ndarray:
+    """get an audio segment given the length
+
+    Args:
+        audio (np.ndarray):
+        length (int): audio length = sampling_rate * duration
+    """
+    if audio.shape[0] < length:
+        audio = np.pad(audio, (0, int(length - audio.shape[0])))
+    start_index = random.randint(0, audio.shape[0] - length)
+    end_index = int(start_index + length)
+
+    return audio[start_index:end_index]
+
+
+def detect_speech_boundaries(
+    wav: np.ndarray,
+    sample_rate: int,
+    window_duration: float = 0.1,
+    energy_threshold: float = 0.01,
+    margin_factor: int = 2,
+) -> Tuple[int, int]:
+    """Detect the start and end points of speech in an audio signal using RMS energy.
+
+    Args:
+        wav: Input audio signal array with values in [-1, 1]
+        sample_rate: Audio sample rate in Hz
+        window_duration: Duration of detection window in seconds
+        energy_threshold: RMS energy threshold for speech detection
+        margin_factor: Factor to determine extra margin around detected boundaries
+
+    Returns:
+        tuple: (start_index, end_index) of speech segment
+
+    Raises:
+        ValueError: If the audio contains only silence
+    """
+    window_size = int(window_duration * sample_rate)
+    margin = margin_factor * window_size
+    step_size = window_size // 10
+
+    # Create sliding windows using stride tricks to avoid loops
+    windows = sliding_window_view(wav, window_size)[::step_size]
+
+    # Calculate RMS energy for each window
+    energy = np.sqrt(np.mean(windows**2, axis=1))
+    speech_mask = energy >= energy_threshold
+
+    if not np.any(speech_mask):
+        raise ValueError("No speech detected in audio (only silence)")
+
+    start = max(0, np.argmax(speech_mask) * step_size - margin)
+    end = min(
+        len(wav),
+        (len(speech_mask) - 1 - np.argmax(speech_mask[::-1])) * step_size + margin,
+    )
+
+    return start, end
+
+
+def remove_silence_on_both_ends(
+    wav: np.ndarray,
+    sample_rate: int,
+    window_duration: float = 0.1,
+    volume_threshold: float = 0.01,
+) -> np.ndarray:
+    """Remove silence from both ends of an audio signal.
+
+    Args:
+        wav: Input audio signal array
+        sample_rate: Audio sample rate in Hz
+        window_duration: Duration of detection window in seconds
+        volume_threshold: Amplitude threshold for silence detection
+
+    Returns:
+        np.ndarray: Audio signal with silence removed from both ends
+
+    Raises:
+        ValueError: If the audio contains only silence
+    """
+    start, end = detect_speech_boundaries(
+        wav, sample_rate, window_duration, volume_threshold
+    )
+    return wav[start:end]
+
+
+def hertz_to_mel(pitch: float) -> float:
+    """
+    Converts a frequency from the Hertz scale to the Mel scale.
+
+    Parameters:
+    - pitch: float or ndarray
+        Frequency in Hertz.
+
+    Returns:
+    - mel: float or ndarray
+        Frequency in Mel scale.
+    """
+    mel = 2595 * np.log10(1 + pitch / 700)
+    return mel