mlx-audio 0.0.1__tar.gz

mlx_audio-0.0.1/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 Prince Canuma
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

mlx_audio-0.0.1/MANIFEST.in

@@ -0,0 +1,2 @@
+include ./requirements.txt
+recursive-include mlx_audio/ *.py

mlx_audio-0.0.1/PKG-INFO

@@ -0,0 +1,131 @@
+Metadata-Version: 2.2
+Name: mlx-audio
+Version: 0.0.1
+Summary: MLX-Audio is a package for inference of text-to-speech (TTS) and speech-to-speech (STS) models locally on your Mac using MLX
+Home-page: https://github.com/Blaizzy/mlx-audio
+Author: Prince Canuma
+Author-email: prince.gdt@gmail.com
+License: MIT
+Classifier: Programming Language :: Python :: 3
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Requires-Python: >=3.8
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: misaki[en]>=0.8.2
+Requires-Dist: loguru>=0.7.3
+Requires-Dist: num2words>=0.5.14
+Requires-Dist: spacy>=3.8.4
+Requires-Dist: phonemizer>=3.3.0
+Requires-Dist: espeakng-loader>=0.2.4
+Requires-Dist: mlx>=0.22.0
+Requires-Dist: mlx-vlm>=0.1.14
+Requires-Dist: mlx-lm>=0.21.5
+Requires-Dist: numpy>=1.26.4
+Requires-Dist: torch>=2.5.1
+Requires-Dist: transformers>=4.49.0
+Requires-Dist: sentencepiece>=0.2.0
+Requires-Dist: huggingface_hub>=0.27.0
+Dynamic: author
+Dynamic: author-email
+Dynamic: classifier
+Dynamic: description
+Dynamic: description-content-type
+Dynamic: home-page
+Dynamic: license
+Dynamic: requires-dist
+Dynamic: requires-python
+Dynamic: summary
+
+# MLX-Audio
+
+A text-to-speech (TTS) and Speech-to-Speech (STS) library built on Apple's MLX framework, providing efficient speech synthesis on Apple Silicon.
+
+## Features
+
+- Fast inference on Apple Silicon (M series chips)
+- Multiple language support
+- Voice customization options
+- Quantization support for optimized performance
+
+## Installation
+
+```bash
+pip install mlx-audio
+```
+
+## Models
+
+### Kokoro
+
+Kokoro is a multilingual TTS model that supports various languages and voice styles.
+
+#### Example Usage
+
+```python
+from tts.models.kokoro import KokoroModel, KokoroPipeline
+from IPython.display import Audio
+import soundfile as sf
+
+# Initialize the model
+model = KokoroModel(repo_id='prince-canuma/Kokoro-82M')
+
+# Create a pipeline with American English
+pipeline = KokoroPipeline(lang_code='a', model=model)
+
+# Generate audio
+text = "The MLX King lives. Let him cook!"
+for _, _, audio in pipeline(text, voice='af_heart', speed=1, split_pattern=r'\n+'):
+    # Display audio in notebook (if applicable)
+    display(Audio(data=audio, rate=24000, autoplay=0))
+
+    # Save audio to file
+    sf.write('audio.wav', audio[0], 24000)
+```
+
+#### Language Options
+
+- 🇺🇸 `'a'` - American English
+- 🇬🇧 `'b'` - British English
+- 🇯🇵 `'j'` - Japanese (requires `pip install misaki[ja]`)
+- 🇨🇳 `'z'` - Mandarin Chinese (requires `pip install misaki[zh]`)
+
+## Advanced Features
+
+### Quantization
+
+You can quantize models for improved performance:
+
+```python
+from tts.models.kokoro import KokoroModel
+from tts.utils import quantize_model
+import json
+import mlx.core as mx
+
+model = KokoroModel(repo_id='prince-canuma/Kokoro-82M')
+config = model.config
+
+# Quantize to 8-bit
+weights, config = quantize_model(model, config, 64, 8)
+
+# Save quantized model
+with open('./8bit/config.json', 'w') as f:
+    json.dump(config, f)
+
+mx.save_safetensors("./8bit/kokoro-v1_0.safetensors", weights, metadata={"format": "mlx"})
+```
+
+## Requirements
+
+- MLX
+- Python 3.8+
+- Apple Silicon Mac (for optimal performance)
+
+## License
+
+[MIT License](LICENSE)
+
+## Acknowledgements
+
+- Thanks to the Apple MLX team for providing a great framework for building TTS and STS models.
+- This project uses the Kokoro model architecture for text-to-speech synthesis.

mlx_audio-0.0.1/README.md

@@ -0,0 +1,92 @@
+# MLX-Audio
+
+A text-to-speech (TTS) and Speech-to-Speech (STS) library built on Apple's MLX framework, providing efficient speech synthesis on Apple Silicon.
+
+## Features
+
+- Fast inference on Apple Silicon (M series chips)
+- Multiple language support
+- Voice customization options
+- Quantization support for optimized performance
+
+## Installation
+
+```bash
+pip install mlx-audio
+```
+
+## Models
+
+### Kokoro
+
+Kokoro is a multilingual TTS model that supports various languages and voice styles.
+
+#### Example Usage
+
+```python
+from tts.models.kokoro import KokoroModel, KokoroPipeline
+from IPython.display import Audio
+import soundfile as sf
+
+# Initialize the model
+model = KokoroModel(repo_id='prince-canuma/Kokoro-82M')
+
+# Create a pipeline with American English
+pipeline = KokoroPipeline(lang_code='a', model=model)
+
+# Generate audio
+text = "The MLX King lives. Let him cook!"
+for _, _, audio in pipeline(text, voice='af_heart', speed=1, split_pattern=r'\n+'):
+    # Display audio in notebook (if applicable)
+    display(Audio(data=audio, rate=24000, autoplay=0))
+
+    # Save audio to file
+    sf.write('audio.wav', audio[0], 24000)
+```
+
+#### Language Options
+
+- 🇺🇸 `'a'` - American English
+- 🇬🇧 `'b'` - British English
+- 🇯🇵 `'j'` - Japanese (requires `pip install misaki[ja]`)
+- 🇨🇳 `'z'` - Mandarin Chinese (requires `pip install misaki[zh]`)
+
+## Advanced Features
+
+### Quantization
+
+You can quantize models for improved performance:
+
+```python
+from tts.models.kokoro import KokoroModel
+from tts.utils import quantize_model
+import json
+import mlx.core as mx
+
+model = KokoroModel(repo_id='prince-canuma/Kokoro-82M')
+config = model.config
+
+# Quantize to 8-bit
+weights, config = quantize_model(model, config, 64, 8)
+
+# Save quantized model
+with open('./8bit/config.json', 'w') as f:
+    json.dump(config, f)
+
+mx.save_safetensors("./8bit/kokoro-v1_0.safetensors", weights, metadata={"format": "mlx"})
+```
+
+## Requirements
+
+- MLX
+- Python 3.8+
+- Apple Silicon Mac (for optimal performance)
+
+## License
+
+[MIT License](LICENSE)
+
+## Acknowledgements
+
+- Thanks to the Apple MLX team for providing a great framework for building TTS and STS models.
+- This project uses the Kokoro model architecture for text-to-speech synthesis.

mlx_audio-0.0.1/mlx_audio/tts/__init__.py

@@ -0,0 +1 @@
+

mlx_audio-0.0.1/mlx_audio/tts/generate.py

@@ -0,0 +1,99 @@
+import argparse
+import json
+import os
+import sys
+
+import mlx.core as mx
+import soundfile as sf
+
+from .utils import load_model
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--model",
+        type=str,
+        default="prince-canuma/Kokoro-82M",
+        help="Path or repo id of the model",
+    )
+    parser.add_argument(
+        "--text", type=str, default="The sky above the port", help="Text to generate"
+    )
+    parser.add_argument("--voice", type=str, default="af_heart", help="Voice name")
+    parser.add_argument("--speed", type=float, default=1.0, help="Speed of the audio")
+    parser.add_argument("--lang_code", type=str, default="a", help="Language code")
+    parser.add_argument(
+        "--file_prefix", type=str, default="audio", help="Output file name prefix"
+    )
+    parser.add_argument("--verbose", action="store_false", help="Print verbose output")
+    parser.add_argument(
+        "--join_audio", action="store_true", help="Join all audio files into one"
+    )
+    return parser.parse_args()
+
+
+def main():
+    args = parse_args()
+    try:
+        model = load_model(model_path=args.model)
+        print(
+            f"\n\033[94mModel:\033[0m {args.model}\n"
+            f"\033[94mText:\033[0m {args.text}\n"
+            f"\033[94mVoice:\033[0m {args.voice}\n"
+            f"\033[94mSpeed:\033[0m {args.speed}x\n"
+            f"\033[94mLanguage:\033[0m {args.lang_code}"
+        )
+        print("==========")
+        results = model.generate(
+            text=args.text,
+            voice=args.voice,
+            speed=args.speed,
+            lang_code=args.lang_code,
+            verbose=True,
+        )
+        print(
+            f"\033[92mAudio generated successfully, saving to\033[0m {args.file_prefix}!"
+        )
+
+        audio_list = []
+        for i, result in enumerate(results):
+            if args.join_audio:
+                audio_list.append(result.audio)
+            else:
+                sf.write(f"{args.file_prefix}_{i:03d}.wav", result.audio, 24000)
+
+            if args.verbose:
+                print("==========")
+                print(f"Duration:              {result.audio_duration}")
+                print(
+                    f"Samples/sec:           {result.audio_samples['samples-per-sec']:.1f}"
+                )
+                print(
+                    f"Prompt:                {result.token_count} tokens, {result.prompt['tokens-per-sec']:.1f} tokens-per-sec"
+                )
+                print(
+                    f"Audio:                 {result.audio_samples['samples']} samples, {result.audio_samples['samples-per-sec']:.1f} samples-per-sec"
+                )
+                print(f"Real-time factor:      {result.real_time_factor:.2f}x")
+                print(f"Processing time:       {result.processing_time_seconds:.2f}s")
+                print(f"Peak memory usage:     {result.peak_memory_usage:.2f}GB")
+
+        if args.join_audio:
+            print(f"Joining {len(audio_list)} audio files")
+            audio = mx.concatenate(audio_list, axis=0)
+            sf.write(f"{args.file_prefix}.wav", audio, 24000)
+    except ImportError as e:
+        print(f"Import error: {e}")
+        print(
+            "This might be due to incorrect Python path. Check your project structure."
+        )
+    except Exception as e:
+        print(f"Error loading model: {e}")
+        import traceback
+
+        traceback.print_exc()
+
+
+if __name__ == "__main__":
+    main()

mlx_audio-0.0.1/mlx_audio/tts/models/base.py

@@ -0,0 +1,48 @@
+import inspect
+from dataclasses import dataclass
+
+import mlx.core as mx
+
+
+@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 check_array_shape(arr):
+    shape = arr.shape
+
+    # Check if the shape has 4 dimensions
+    if len(shape) != 3:
+        return False
+
+    out_channels, kH, KW = shape
+
+    # Check if out_channels is the largest, and kH and KW are the same
+    if (out_channels >= kH) and (out_channels >= KW) and (kH == KW):
+        return True
+    else:
+        return False
+
+
+@dataclass
+class GenerationResult:
+    audio: mx.array
+    samples: int
+    segment_idx: int
+    token_count: int
+    audio_samples: int
+    audio_duration: str
+    real_time_factor: float
+    prompt: dict
+    audio_samples: dict
+    processing_time_seconds: float
+    peak_memory_usage: float

mlx_audio-0.0.1/mlx_audio/tts/models/interpolate.py

@@ -0,0 +1,108 @@
+from typing import List, Optional, Tuple, Union
+
+import mlx.core as mx
+
+
+def interpolate(
+    input: mx.array,
+    size: Optional[Union[int, Tuple[int, ...], List[int]]] = None,
+    scale_factor: Optional[Union[float, List[float], Tuple[float, ...]]] = None,
+    mode: str = "nearest",
+    align_corners: Optional[bool] = None,
+) -> mx.array:
+    """Interpolate array with correct shape handling.
+
+    Args:
+        input (mx.array): Input tensor [N, C, ...] where ... represents spatial dimensions
+        size (int or tuple): Output size
+        scale_factor (float or tuple): Multiplier for spatial size
+        mode (str): 'nearest' or 'linear'
+        align_corners (bool): If True, align corners of input and output tensors
+    """
+    ndim = input.ndim
+    if ndim < 3:
+        raise ValueError(f"Expected at least 3D input (N, C, D1), got {ndim}D")
+
+    spatial_dims = ndim - 2
+
+    # Handle size and scale_factor
+    if size is not None and scale_factor is not None:
+        raise ValueError("Only one of size or scale_factor should be defined")
+    elif size is None and scale_factor is None:
+        raise ValueError("One of size or scale_factor must be defined")
+
+    # Convert single values to tuples
+    if size is not None and not isinstance(size, (list, tuple)):
+        size = [size] * spatial_dims
+    if scale_factor is not None and not isinstance(scale_factor, (list, tuple)):
+        scale_factor = [scale_factor] * spatial_dims
+
+    # Calculate output size from scale factor if needed
+    if size is None:
+        size = []
+        for i in range(spatial_dims):
+            # Use ceiling instead of floor to match PyTorch behavior
+            curr_size = max(1, int(mx.ceil(input.shape[i + 2] * scale_factor[i])))
+            size.append(curr_size)
+
+    # Handle 1D case (N, C, W)
+    if spatial_dims == 1:
+        return interpolate1d(input, size[0], mode, align_corners)
+    else:
+        raise ValueError(
+            f"Only 1D interpolation currently supported, got {spatial_dims}D"
+        )
+
+
+def interpolate1d(
+    input: mx.array,
+    size: int,
+    mode: str = "linear",
+    align_corners: Optional[bool] = None,
+) -> mx.array:
+    """1D interpolation implementation."""
+    batch_size, channels, in_width = input.shape
+
+    # Handle edge cases
+    if size < 1:
+        size = 1
+    if in_width < 1:
+        in_width = 1
+
+    if mode == "nearest":
+        if size == 1:
+            indices = mx.array([0])
+        else:
+            scale = in_width / size
+            indices = mx.floor(mx.arange(size) * scale).astype(mx.int32)
+            indices = mx.clip(indices, 0, in_width - 1)
+        return input[:, :, indices]
+
+    # Linear interpolation
+    if align_corners and size > 1:
+        x = mx.arange(size) * ((in_width - 1) / (size - 1))
+    else:
+        if size == 1:
+            x = mx.array([0.0])
+        else:
+            x = mx.arange(size) * (in_width / size)
+            if not align_corners:
+                x = x + 0.5 * (in_width / size) - 0.5
+
+    # Handle the case where input width is 1
+    if in_width == 1:
+        output = mx.broadcast_to(input, (batch_size, channels, size))
+        return output
+
+    x_low = mx.floor(x).astype(mx.int32)
+    x_high = mx.minimum(x_low + 1, in_width - 1)
+    x_frac = x - x_low
+
+    # Pre-compute indices to avoid repeated computation
+    y_low = input[:, :, x_low]
+    y_high = input[:, :, x_high]
+
+    # Vectorized interpolation
+    output = y_low * (1 - x_frac)[None, None, :] + y_high * x_frac[None, None, :]
+
+    return output

mlx_audio-0.0.1/mlx_audio/tts/models/kokoro/__init__.py

@@ -0,0 +1,4 @@
+from .kokoro import Model
+from .pipeline import KokoroPipeline
+
+__all__ = ["KokoroPipeline", "Model"]