nexaai 1.0.29__cp310-cp310-macosx_14_0_universal2.whl

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

@@ -0,0 +1,221 @@
+# 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
+    file reading and writing operations. It provides utilities to read from files,
+    write data to files, and perform file manipulation tasks.
+"""
+
+
+import csv
+import json
+import os
+from pathlib import Path
+from typing import Any, Dict, List, Set, Union
+
+from omegaconf import DictConfig, OmegaConf
+from tqdm import tqdm
+
+
+def resolve_symbolic_link(symbolic_link_path: Path) -> Path:
+    """
+    Resolves the absolute path of a symbolic link.
+
+    Args:
+        symbolic_link_path (Path): The path to the symbolic link.
+
+    Returns:
+        Path: The absolute path that the symbolic link points to.
+    """
+
+    link_directory = os.path.dirname(symbolic_link_path)
+    target_path_relative = os.readlink(symbolic_link_path)
+    return os.path.join(link_directory, target_path_relative)
+
+
+def write_jsonl(metadata: List[dict], file_path: Path) -> None:
+    """Writes a list of dictionaries to a JSONL file.
+
+    Args:
+    metadata : List[dict]
+        A list of dictionaries, each representing a piece of meta.
+    file_path : Path
+        The file path to save the JSONL file
+
+    This function writes each dictionary in the list to a new line in the specified file.
+    """
+    with open(file_path, "w", encoding="utf-8") as f:
+        for meta in tqdm(metadata, desc="writing jsonl"):
+            # Convert dictionary to JSON string and write it to the file with a newline
+            json_str = json.dumps(meta, ensure_ascii=False) + "\n"
+            f.write(json_str)
+    print(f"jsonl saved to {file_path}")
+
+
+def read_jsonl(file_path: Path) -> List[dict]:
+    """
+    Reads a JSONL file and returns a list of dictionaries.
+
+    Args:
+    file_path : Path
+        The path to the JSONL file to be read.
+
+    Returns:
+    List[dict]
+        A list of dictionaries parsed from each line of the JSONL file.
+    """
+    metadata = []
+    # Open the file for reading
+    with open(file_path, "r", encoding="utf-8") as f:
+        # Split the file into lines
+        lines = f.read().splitlines()
+    # Process each line
+    for line in lines:
+        # Convert JSON string back to dictionary and append to list
+        meta = json.loads(line)
+        metadata.append(meta)
+    # Return the list of metadata
+    return metadata
+
+
+def read_json_as_jsonl(file_path: Path) -> List[dict]:
+    metadata = []
+    with open(file_path, "r", encoding="utf-8") as infile:
+        data = json.load(infile)
+    for k in sorted(data.keys()):
+        meta = {"index": k}
+        meta.update(data[k])
+        metadata.append(meta)
+    return metadata
+
+
+def decode_unicode_strings(meta: Dict[str, Any]) -> Dict[str, Any]:
+    processed_meta = {}
+    for k, v in meta.items():
+        if isinstance(v, str):
+            processed_meta[k] = v.encode("utf-8").decode("unicode_escape")
+        else:
+            processed_meta[k] = v
+    return processed_meta
+
+
+def load_config(config_path: Path) -> DictConfig:
+    """Loads a configuration file and optionally merges it with a base configuration.
+
+    Args:
+    config_path (Path): Path to the configuration file.
+    """
+    # Load the initial configuration from the given path
+    config = OmegaConf.load(config_path)
+
+    # Check if there is a base configuration specified and merge if necessary
+    if config.get("base_config", None) is not None:
+        base_config = OmegaConf.load(config["base_config"])
+        config = OmegaConf.merge(base_config, config)
+
+    return config
+
+
+def jsonl_to_csv(jsonl_file_path: str, csv_file_path: str) -> None:
+    """
+    Converts a JSONL file to a CSV file.
+
+    This function reads a JSONL file, determines all unique keys present in the file,
+    and writes the data to a CSV file with columns for all these keys.
+    """
+
+    all_keys = set()
+    data_rows = []
+
+    # Read the JSONL file once to extract keys and collect data
+    with open(jsonl_file_path, "r") as file:
+        for line in file:
+            data = json.loads(line.strip())
+            data_rows.append(data)
+            all_keys.update(data.keys())
+
+    # Convert the set of keys to a sorted list for consistent column order
+    sorted_keys = sorted(all_keys)
+
+    # Write the data to a CSV file
+    with open(csv_file_path, "w", newline="") as csvfile:
+        writer = csv.DictWriter(csvfile, fieldnames=sorted_keys)
+
+        # Write the header row
+        writer.writeheader()
+
+        # Write each row of data
+        for data in data_rows:
+            writer.writerow(data)
+
+    print(f"CSV file has been created at {csv_file_path}")
+
+
+def save_metadata(data, filename, headers=None):
+    """
+    Save metadata to a file.
+
+    Args:
+        data (list of dict): Metadata to be saved.
+        filename (str): Name of the file to save the metadata.
+        headers (list of str): The order of column names to be saved; defaults to the keys from the first dictionary in data if not provided.
+    """
+    # Set headers to keys from the first dictionary in data if not explicitly provided
+    if headers is None:
+        headers = list(data[0].keys())
+
+    with open(filename, "w", encoding="utf-8") as file:
+        # Write the headers to the file
+        file.write("|".join(headers) + "\n")
+        for entry in data:
+            # Retrieve values in the order of headers, replacing any '|' characters with a space to prevent formatting errors
+            formatted_values = [
+                str(entry.get(key, "")).replace("|", " ") for key in headers
+            ]
+            # Write the formatted values to the file
+            file.write("|".join(formatted_values) + "\n")
+
+
+def read_metadata(filename, headers=None):
+    """
+    Read metadata from a file.
+
+    Args:
+        filename (str): The file from which to read the metadata.
+
+    Returns:
+        list of dict: The metadata read from the file.
+        list of str: The headers used in the file.
+    """
+    with open(filename, "r", encoding="utf-8") as file:
+        lines = file.readlines()
+
+    data = []
+    # Set headers from the first line of the file if not provided
+    if headers is None:
+        headers = lines[0].strip().split("|")
+        lines = lines[1:]
+
+    for line in lines:
+        line = line.strip()
+        # Skip empty lines
+        if not line:
+            continue
+        # Split the line by '|' and pair with headers to form a dictionary
+        entry_data = dict(zip(headers, line.split("|")))
+        data.append(entry_data)
+
+    return data, headers

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

@@ -0,0 +1,181 @@
+TASK_TOKEN_MAP = {
+    "vc": "<|task_vc|>",
+    "tts": "<|task_tts|>",
+    "asr": "<|task_asr|>",
+    "s2s": "<|task_s2s|>",
+    "t2s": "<|task_t2s|>",
+    "understand": "<|task_understand|>",
+    "caption": "<|task_cap|>",
+    "controllable_tts": "<|task_controllable_tts|>",
+    "prompt_tts": "<|task_prompt_tts|>",
+    "speech_edit": "<|task_edit|>",
+}
+
+LEVELS_MAP = {
+    "very_low": 0,
+    "low": 1,
+    "moderate": 2,
+    "high": 3,
+    "very_high": 4,
+}
+
+LEVELS_MAP_UI = {1: "very_low", 2: "low", 3: "moderate", 4: "high", 5: "very_high"}
+
+GENDER_MAP = {
+    "female": 0,
+    "male": 1,
+}
+
+AGE_MAP = {"Child": 0, "Teenager": 1, "Youth-Adult": 2, "Middle-aged": 3, "Elderly": 4}
+
+EMO_MAP = {
+    "UNKNOWN": 0,
+    "NEUTRAL": 1,
+    "ANGRY": 2,
+    "HAPPY": 3,
+    "SAD": 4,
+    "FEARFUL": 5,
+    "DISGUSTED": 6,
+    "SURPRISED": 7,
+    "SARCASTIC": 8,
+    "EXCITED": 9,
+    "SLEEPY": 10,
+    "CONFUSED": 11,
+    "EMPHASIS": 12,
+    "LAUGHING": 13,
+    "SINGING": 14,
+    "WORRIED": 15,
+    "WHISPER": 16,
+    "ANXIOUS": 17,
+    "NO-AGREEMENT": 18,
+    "APOLOGETIC": 19,
+    "CONCERNED": 20,
+    "ENUNCIATED": 21,
+    "ASSERTIVE": 22,
+    "ENCOURAGING": 23,
+    "CONTEMPT": 24,
+}
+
+
+class TokenParser:
+    """Turn label to special token"""
+
+    def __init__(self):
+        pass
+
+    """Parse the attributes of a person."""
+
+    def __init__(self):
+        pass
+
+    @staticmethod
+    def age(age: str) -> str:
+        """Turn age token."""
+        age_id = AGE_MAP[age]
+        return f"<|age_{age_id}|>"
+
+    @staticmethod
+    def gender(gender: str) -> str:
+        """Turn gender token."""
+        gender_id = GENDER_MAP[gender]
+        return f"<|gender_{gender_id}|>"
+
+    @staticmethod
+    def mel_value(mel: int):
+        """Turn special token of mel scale pitch."""
+        mel = max(0, int(mel))
+        mel = min(1000, int(mel))
+        return f"<|pitch_value_{mel}|>"
+
+    @staticmethod
+    def mel_level(level: str):
+        """Turn special token of mel level."""
+        level_tag = LEVELS_MAP[level]
+        return f"<|pitch_label_{level_tag}|>"
+
+    @staticmethod
+    def pitch_var_value(pitch_std: int):
+        """Turn special token of pitch_std value."""
+        assert isinstance(pitch_std, int)
+        pitch_std = max(0, int(pitch_std))
+        pitch_std = min(10, int(pitch_std))
+        return f"<|pitch_var_value_{pitch_std}|>"
+
+    @staticmethod
+    def pitch_var_level(level: str):
+        """Turn special token of pitch std level."""
+        level_tag = LEVELS_MAP[level]
+        return f"<|pitch_var_label_{level_tag}|>"
+
+    @staticmethod
+    def loudness_value(loudness: int):
+        """Turn special toak of loudness value [0, 30]"""
+        assert loudness >= 0
+        loudness = max(0, int(loudness))
+        loudness = min(30, int(loudness))
+        return f"<|loudness_value_{loudness}|>"
+
+    @staticmethod
+    def loudness_level(level: str):
+        """Turn special token of loudness level."""
+        level_tag = LEVELS_MAP[level]
+        return f"<|loudness_label_{level_tag}|>"
+
+    @staticmethod
+    def speed_value(speed: int):
+        """Turn special token of speed value."""
+        speed = max(0, int(speed))
+        speed = min(10, int(speed))
+        return f"<|speed_value_{speed}|>"
+
+    @staticmethod
+    def speed_level(level: str):
+        """Turn special token of speed level."""
+        level_tag = LEVELS_MAP[level]
+        return f"<|speed_label_{level_tag}|>"
+
+    @staticmethod
+    def task(task: str) -> str:
+        """Turn special token of task."""
+        assert task in TASK_TOKEN_MAP.keys()
+
+        return TASK_TOKEN_MAP[task]
+
+    @staticmethod
+    def emotion(emotion: str):
+        emo_id = EMO_MAP[emotion]
+
+        return f"<|emotion_{emo_id}|>"
+
+
+# test
+if __name__ == "__main__":
+    from transformers import AutoTokenizer
+
+    tokenizer = AutoTokenizer.from_pretrained(
+        "/aifs4su/xinshengwang/code/StyleCraft/tokenizer/stylecraft-bicodec-pitch-loudness-speed-emotion-tokenizer"
+    )
+
+    tasks = ["tts", "tts", "understand", "controllable_tts", "prompt_tts"]
+    ages = ["Child", "Teenager", "Youth-Adult", "Middle-aged", "Elderly"]
+    genders = ["female", "female", "female", "male", "male"]
+    mels = [100, 200, 300, 400, 500]
+    mel_levels = ["very_low", "low", "moderate", "high", "very_high"]
+    loudnesses = [1, 10, 23, 19, 30]
+    loudness_levels = ["very_low", "low", "moderate", "high", "very_high"]
+    emotions = ["UNKNOWN", "NEUTRAL", "ANGRY", "HAPPY", "SAD"]
+
+    for i in range(5):
+        task = TokenParser.task(tasks[i])
+        age = TokenParser.age(ages[i])
+        gender = TokenParser.gender(genders[i])
+        mel = TokenParser.mel_value(mels[i])
+        mel_level = TokenParser.mel_level(mel_levels[i])
+        loudness = TokenParser.loudness_value(loudnesses[i])
+        loudness_level = TokenParser.loudness_level(loudness_levels[i])
+        emotion = TokenParser.emotion(emotions[i])
+        inputs = [task, age, gender, mel, mel_level, loudness, loudness_level, emotion]
+        inputs = "".join(inputs)
+        ids = tokenizer.encode(inputs, add_special_tokens=False)
+        print(ids)
+        print("decode", tokenizer.decode(ids))

nexaai/binds/metal/py-lib/mlx_audio/tts/tests/test_base.py

@@ -0,0 +1,66 @@
+import unittest
+
+import mlx.core as mx
+import numpy as np
+
+from mlx_audio.tts.models.base import BaseModelArgs, check_array_shape
+
+
+class TestBaseModel(unittest.TestCase):
+    def test_base_model_args_from_dict(self):
+        """Test BaseModelArgs.from_dict method."""
+
+        # Define a test subclass
+        class TestArgs(BaseModelArgs):
+            def __init__(self, param1, param2, param3=None):
+                self.param1 = param1
+                self.param2 = param2
+                self.param3 = param3
+
+        # Test with all parameters
+        params = {"param1": 1, "param2": "test", "param3": True}
+        args = TestArgs.from_dict(params)
+        self.assertEqual(args.param1, 1)
+        self.assertEqual(args.param2, "test")
+        self.assertEqual(args.param3, True)
+
+        # Test with extra parameters (should be ignored)
+        params = {"param1": 1, "param2": "test", "param3": True, "extra": "ignored"}
+        args = TestArgs.from_dict(params)
+        self.assertEqual(args.param1, 1)
+        self.assertEqual(args.param2, "test")
+        self.assertEqual(args.param3, True)
+        self.assertFalse(hasattr(args, "extra"))
+
+        # Test with missing optional parameter
+        params = {"param1": 1, "param2": "test"}
+        args = TestArgs.from_dict(params)
+        self.assertEqual(args.param1, 1)
+        self.assertEqual(args.param2, "test")
+        self.assertIsNone(args.param3)
+
+    def test_check_array_shape(self):
+        """Test check_array_shape function."""
+        # Valid shape: out_channels >= kH == kW
+        valid_array = mx.array(np.zeros((64, 3, 3)))
+        self.assertTrue(check_array_shape(valid_array))
+
+        # Invalid shape: kH != kW
+        invalid_array1 = mx.array(np.zeros((64, 3, 4)))
+        self.assertFalse(check_array_shape(invalid_array1))
+
+        # Invalid shape: out_channels < kH
+        invalid_array2 = mx.array(np.zeros((2, 3, 3)))
+        self.assertFalse(check_array_shape(invalid_array2))
+
+        # Invalid shape: wrong number of dimensions
+        invalid_array3 = mx.array(np.zeros((64, 3)))
+        self.assertFalse(check_array_shape(invalid_array3))
+
+        # Invalid shape: wrong number of dimensions
+        invalid_array4 = mx.array(np.zeros((64, 3, 3, 3)))
+        self.assertFalse(check_array_shape(invalid_array4))
+
+
+if __name__ == "__main__":
+    unittest.main()

nexaai/binds/metal/py-lib/mlx_audio/tts/tests/test_convert.py

@@ -0,0 +1,173 @@
+import sys  # Import sys to patch argv
+import unittest
+from unittest.mock import MagicMock, patch
+
+from mlx_audio.tts.convert import configure_parser, main
+
+
+class TestConvert(unittest.TestCase):
+    def setUp(self):
+        self.parser = configure_parser()
+
+        # Mock the actual convert function
+        self.convert_mock = MagicMock()
+        self.patcher = patch("mlx_audio.tts.convert.convert", new=self.convert_mock)
+        self.patcher.start()
+
+    def tearDown(self):
+        self.patcher.stop()
+
+    def test_basic_conversion(self):
+        test_args = [
+            "--hf-path",
+            "dummy_hf",
+            "--mlx-path",
+            "dummy_mlx",
+            "--dtype",
+            "float16",
+        ]
+        # Patch sys.argv for this test run
+        with patch.object(sys, "argv", ["convert.py"] + test_args):
+            main()
+
+        self.convert_mock.assert_called_once_with(
+            hf_path="dummy_hf",
+            mlx_path="dummy_mlx",
+            quantize=False,
+            q_group_size=64,
+            q_bits=4,
+            quant_predicate=None,
+            dtype="float16",
+            upload_repo=None,
+            dequantize=False,
+        )
+
+    def test_quantized_conversion(self):
+        test_args = [
+            "--hf-path",
+            "dummy_hf",
+            "--quantize",
+            "--q-group-size",
+            "128",
+            "--q-bits",
+            "8",
+        ]
+        # Patch sys.argv for this test run
+        with patch.object(sys, "argv", ["convert.py"] + test_args):
+            main()
+
+        self.convert_mock.assert_called_once_with(
+            hf_path="dummy_hf",
+            mlx_path="mlx_model",  # Default mlx_path
+            quantize=True,
+            q_group_size=128,
+            q_bits=8,
+            quant_predicate=None,
+            dtype="float16",  # Should be ignored when quantize=True
+            upload_repo=None,
+            dequantize=False,
+        )
+
+    def test_quantized_conversion_invalid_group_size_raises_error(self):
+        """Tests if main raises ValueError for invalid group size."""
+        test_args = [
+            "--hf-path",
+            "dummy_hf",
+            "--quantize",
+            "--q-group-size",
+            "100",  # Invalid: not 64 or 128
+            "--q-bits",
+            "4",
+        ]
+
+        # Configure the mock to raise ValueError when called with q_group_size=100
+        def side_effect(*args, **kwargs):
+            if kwargs.get("q_group_size") == 100:
+                raise ValueError(
+                    "[quantize] The requested group size 100 is not supported."
+                )
+            return MagicMock()  # Default return for other calls if needed
+
+        self.convert_mock.side_effect = side_effect
+
+        # Patch sys.argv and assert ValueError is raised
+        with patch.object(sys, "argv", ["convert.py"] + test_args):
+            with self.assertRaisesRegex(
+                ValueError, "requested group size 100 is not supported"
+            ):
+                main()
+
+        # Verify the mock was called (even though it raised an error)
+        self.convert_mock.assert_called_once_with(
+            hf_path="dummy_hf",
+            mlx_path="mlx_model",
+            quantize=True,
+            q_group_size=100,
+            q_bits=4,
+            quant_predicate=None,
+            dtype="float16",
+            upload_repo=None,
+            dequantize=False,
+        )
+
+    def test_quantization_recipes(self):
+        for recipe in ["mixed_2_6", "mixed_3_6", "mixed_4_6"]:
+            with self.subTest(recipe=recipe):
+                self.convert_mock.reset_mock()  # Reset mock for each subtest
+                test_args = ["--hf-path", "dummy_hf", "--quant-predicate", recipe]
+                # Patch sys.argv for this test run
+                with patch.object(sys, "argv", ["convert.py"] + test_args):
+                    main()
+
+                self.convert_mock.assert_called_once_with(  # Changed to assert_called_once_with
+                    hf_path="dummy_hf",
+                    mlx_path="mlx_model",  # Default mlx_path
+                    quantize=False,  # Default quantize
+                    q_group_size=64,  # Default q_group_size
+                    q_bits=4,  # Default q_bits
+                    quant_predicate=recipe,
+                    dtype="float16",  # Default dtype
+                    upload_repo=None,  # Default upload_repo
+                    dequantize=False,  # Default dequantize
+                )
+                # No need to reset mock here, it's handled at the start of the loop
+
+    def test_dequantize_flag(self):
+        test_args = ["--hf-path", "dummy_hf", "--dequantize"]
+        # Patch sys.argv for this test run
+        with patch.object(sys, "argv", ["convert.py"] + test_args):
+            main()
+
+        self.convert_mock.assert_called_once_with(
+            hf_path="dummy_hf",
+            mlx_path="mlx_model",  # Default mlx_path
+            quantize=False,
+            q_group_size=64,
+            q_bits=4,
+            quant_predicate=None,
+            dtype="float16",
+            upload_repo=None,
+            dequantize=True,
+        )
+
+    def test_upload_repo_argument(self):
+        test_args = ["--hf-path", "dummy_hf", "--upload-repo", "my/repo"]
+        # Patch sys.argv for this test run
+        with patch.object(sys, "argv", ["convert.py"] + test_args):
+            main()
+
+        self.convert_mock.assert_called_once_with(
+            hf_path="dummy_hf",
+            mlx_path="mlx_model",  # Default mlx_path
+            quantize=False,
+            q_group_size=64,
+            q_bits=4,
+            quant_predicate=None,
+            dtype="float16",
+            upload_repo="my/repo",
+            dequantize=False,
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()