"""Utility functions for mlx_audio. This module provides a unified interface for loading TTS, STT, and STS models, with lazy imports to avoid loading unnecessary dependencies. """ import dataclasses import glob import importlib import importlib.util import json import keyword import logging import re from pathlib import Path from typing import ( List, Optional, Tuple, Type, TypeVar, Union, get_origin, get_type_hints, ) import mlx.core as mx import mlx.nn as nn from huggingface_hub import snapshot_download from mlx_audio.dsp import ( STR_TO_WINDOW_FN, bartlett, blackman, hamming, hanning, istft, mel_filters, stft, ) T = TypeVar("T") def from_dict(data_class: Type[T], data: dict) -> T: if not dataclasses.is_dataclass(data_class): raise TypeError(f"{data_class} is not a dataclass") field_types = get_type_hints(data_class) kwargs = {} for field in dataclasses.fields(data_class): field_name = field.name if field_name not in data: continue value = data[field_name] field_type = field_types[field_name] # Handle Optional types origin = get_origin(field_type) if origin is Union: # For Optional[X], get the non-None type args = [a for a in field_type.__args__ if a is not type(None)] if args: field_type = args[0] # Recursively convert nested dataclasses if dataclasses.is_dataclass(field_type) and isinstance(value, dict): value = from_dict(field_type, value) kwargs[field_name] = value return data_class(**kwargs) # ============================================================================= # Shared Model Loading Utilities # ============================================================================= # Default file patterns to download from HuggingFace DEFAULT_ALLOW_PATTERNS = [ "*.json", "*.safetensors", "*.py", "*.model", "*.tiktoken", "*.txt", "*.jinja", "*.jsonl", "*.yaml", "*.npz", "*.pth", ] def _is_local_path(path: str) -> bool: """Check if the path looks like a local filesystem path.""" return ( path.startswith(".") or path.startswith("/") or path.startswith("~") or (len(path) > 1 and path[1] == ":") # Windows drive letter ) def get_model_path( path_or_hf_repo: str, revision: Optional[str] = None, force_download: bool = False, allow_patterns: Optional[List[str]] = None, ) -> Path: """ Ensures the model is available locally. If the path does not exist locally, it is downloaded from the Hugging Face Hub. Args: path_or_hf_repo: The local path or Hugging Face repository ID of the model. revision: A revision id which can be a branch name, a tag, or a commit hash. force_download: Force re-download even if cached. allow_patterns: File patterns to download. Defaults to common model files. Returns: Path: The path to the model. Raises: FileNotFoundError: If a local path is provided but doesn't exist. """ model_path = Path(path_or_hf_repo).expanduser() if model_path.exists(): return model_path # If it looks like a local path but doesn't exist, raise an error if _is_local_path(path_or_hf_repo): raise FileNotFoundError(f"Local path not found: {path_or_hf_repo}") # Use default patterns if none provided if allow_patterns is None: allow_patterns = DEFAULT_ALLOW_PATTERNS model_path = Path( snapshot_download( path_or_hf_repo, revision=revision, allow_patterns=allow_patterns, force_download=force_download, ) ) return model_path def load_config(model_path: Union[str, Path], **kwargs) -> dict: """Load model configuration from a path or Hugging Face repo. Args: model_path: Local path or Hugging Face repo ID to load config from **kwargs: Additional keyword arguments (revision, force_download) Returns: dict: Model configuration Raises: FileNotFoundError: If config.json is not found at the path """ if isinstance(model_path, str): model_path = get_model_path(model_path, **kwargs) config_file = model_path / "config.json" if config_file.exists(): with open(config_file, encoding="utf-8") as f: return json.load(f) raise FileNotFoundError(f"Config not found at {model_path}") def load_weights(model_path: Path) -> dict: """Load model weights from safetensors or npz files. Args: model_path: Path to the model directory Returns: dict: Dictionary of weight name -> array Raises: FileNotFoundError: If no weight files found """ # Try safetensors first, then npz weight_files = glob.glob(str(model_path / "*.safetensors")) if not weight_files: weight_files = glob.glob(str(model_path / "*.npz")) if not weight_files: raise FileNotFoundError( f"No weight files (safetensors or npz) found in {model_path}" ) weights = {} for wf in weight_files: weights.update(mx.load(wf)) return weights def apply_quantization( model: nn.Module, config: dict, weights: dict, model_quant_predicate: Optional[callable] = None, ) -> None: """Apply quantization to a model if specified in config. Args: model: The model to quantize config: Model config dict (should contain 'quantization' key if quantized) weights: Loaded weights dict (used to check which layers have scales) model_quant_predicate: Optional model-specific predicate for which layers to quantize """ quantization = config.get("quantization", None) if quantization is None: quantization = config.get("quantization_config", None) if quantization is None: return group_size = quantization.get("group_size", 64) def get_class_predicate(p, m): # Skip layers without quantization capability if not hasattr(m, "to_quantized"): return False # Skip layers not divisible by configured group size if hasattr(m, "weight") and m.weight.shape[-1] % group_size != 0: return False # Use model-specific predicate if available if model_quant_predicate is not None: pred_result = model_quant_predicate(p, m) if isinstance(pred_result, dict): return pred_result if not pred_result: return False # Handle custom per layer quantizations if p in quantization: return quantization[p] # Handle legacy models which may not have everything quantized return f"{p}.scales" in weights nn.quantize( model, group_size=group_size, bits=quantization["bits"], mode=quantization.get("mode", "affine"), class_predicate=get_class_predicate, ) def get_model_class( model_type: str, model_name: List[str], category: str, model_remapping: dict, ) -> Tuple: """ Retrieve the model architecture module based on the model type and name. Args: model_type: The type of model to load (e.g., "whisper", "voxtral"). model_name: List of model name components for remapping hints. category: Either "tts" or "stt". model_remapping: Dictionary mapping model names to architecture names. Returns: Tuple[module, str]: The imported architecture module and resolved model_type. Raises: ValueError: If the model type is not supported. """ # Stage 1: Check if the model type is in the remapping model_type_mapped = model_remapping.get(model_type, None) # Stage 2: Check for partial matches in segments of the model name # Only do this if the initial mapping didn't find a match models_dir = Path(__file__).parent / category / "models" available_models = [] if models_dir.exists() and models_dir.is_dir(): for item in models_dir.iterdir(): if item.is_dir() and not item.name.startswith("__"): available_models.append(item.name) if model_name is not None and model_type_mapped != model_type: for part in model_name: if part in available_models: model_type = part if part in model_remapping: model_type = model_remapping[part] break elif model_type_mapped is not None: model_type = model_type_mapped try: module_path = f"mlx_audio.{category}.models.{model_type}" arch = importlib.import_module(module_path) except ImportError as e: if e.name != module_path: print("\n", flush=True) raise ImportError( f"\nMissing dependency while loading {model_type}: {e}\n" f"Please install it using: pip install {e.name}" ) from e msg = f"Model type {model_type} not supported for {category}." logging.error(msg) raise ValueError(msg) return arch, model_type def base_load_model( model_path: Union[str, Path], category: str, model_remapping: dict, lazy: bool = False, strict: bool = False, **kwargs, ) -> nn.Module: """ Base implementation for loading models (shared between TTS, STT, and STS). Args: model_path: The path or HuggingFace repo to load the model from. category: Either "tts" or "stt". model_remapping: Dictionary mapping model names to architecture names. lazy: If False, evaluate model parameters immediately. strict: If True, raise an error if any weights are missing. **kwargs: Additional keyword arguments (revision, force_download). Returns: nn.Module: The loaded and initialized model. """ model_name = kwargs.pop("model_name_parts", None) model_type = kwargs.pop("model_type", None) allow_patterns = kwargs.pop("allow_patterns", None) if isinstance(model_path, str): if model_name is None: model_name = get_model_name_parts(model_path) revision = kwargs.get("revision", None) force_download = kwargs.get("force_download", False) model_path = get_model_path( model_path, revision=revision, force_download=force_download, allow_patterns=allow_patterns, ) elif isinstance(model_path, Path): if model_name is None: model_name = get_model_name_parts(model_path) else: raise ValueError(f"Invalid model path type: {type(model_path)}") config = load_config(model_path) config["model_path"] = str(model_path) # Determine model_type from config or model_name if model_type is None: model_type = config.get("model_type", None) if model_type is None: model_type = config.get("architecture", None) if model_type is None: model_type = model_name[0].lower() if model_name is not None else None # Override model_type for TADA models (config says "llama" but it's TADA) if model_type == "llama" and "acoustic_dim" in config: model_type = "tada" model_class, model_type = get_model_class( model_type=model_type, model_name=model_name, category=category, model_remapping=model_remapping, ) # Get model config from model class if it exists, otherwise use the config model_config = ( model_class.ModelConfig.from_dict(config) if hasattr(model_class, "ModelConfig") else config ) model = model_class.Model(model_config) # Load weights weights = load_weights(model_path) # Sanitize weights if the model has a sanitize method if hasattr(model, "sanitize"): weights = model.sanitize(weights) # Apply quantization if specified model_quant_predicate = getattr(model, "model_quant_predicate", None) apply_quantization(model, config, weights, model_quant_predicate) model.load_weights(list(weights.items()), strict=strict) if not lazy: mx.eval(model.parameters()) model.eval() # Call post-load hook if the model defines one if hasattr(model_class.Model, "post_load_hook"): model = model_class.Model.post_load_hook(model, model_path) return model # Lazy-loaded modules _stt_utils = None _tts_utils = None _sts_utils = None _vad_utils = None _lid_utils = None def _get_stt_utils(): """Lazy load STT utils.""" global _stt_utils if _stt_utils is None: from mlx_audio.stt import utils as stt_utils _stt_utils = stt_utils return _stt_utils def _get_tts_utils(): """Lazy load TTS utils.""" global _tts_utils if _tts_utils is None: from mlx_audio.tts import utils as tts_utils _tts_utils = tts_utils return _tts_utils def _get_sts_utils(): """Lazy load STS utils.""" global _sts_utils if _sts_utils is None: from mlx_audio.sts import utils as sts_utils _sts_utils = sts_utils return _sts_utils def _get_vad_utils(): """Lazy load VAD utils.""" global _vad_utils if _vad_utils is None: from mlx_audio.vad import utils as vad_utils _vad_utils = vad_utils return _vad_utils def _get_lid_utils(): """Lazy load LID utils.""" global _lid_utils if _lid_utils is None: from mlx_audio.lid import utils as lid_utils _lid_utils = lid_utils return _lid_utils def audio_volume_normalize(audio, coeff: float = 0.2): """Normalize the volume of an audio signal. Args: audio: Input audio signal array (numpy array). coeff: Target coefficient for normalization, default is 0.2. Returns: numpy array: The volume-normalized audio signal. """ import numpy as np # 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) audio = audio / scaling_factor * 0.1 # Filter out values less than 0.01 from temp temp = temp[temp > 0.01] L = temp.shape[0] # If there are fewer than or equal to 10 significant values, return as-is 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 audio = audio * np.clip(coeff / volume, a_min=0.1, a_max=10) # Ensure the maximum absolute value does not exceed 1 max_value = np.max(np.abs(audio)) if max_value > 1: audio = audio / max_value return audio def random_select_audio_segment(audio, length: int): """Get a random audio segment of given length. Args: audio: Input audio array (numpy array). length: Desired segment length (sample_rate * duration). Returns: numpy array: Audio segment of specified length. """ import random import numpy as np 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 resample_audio( audio: Union[mx.array, "np.ndarray"], orig_sample_rate: int, sample_rate: int, axis: int = -1, ): """Resample audio with polyphase filtering. Uses a Kaiser-windowed sinc anti-aliasing filter equivalent to ``librosa``'s ``kaiser_best`` (resampy parameters) instead of ``resample_poly``'s default Kaiser(5.0). The default filter has a wide transition band that leaves significant energy near the new Nyquist; the sharper filter band-limits cleanly so resampled audio matches the librosa-equivalent featurizers that ASR reference pipelines (e.g. NeMo) assume. See issue #24. Args: audio: Audio array as numpy or MLX. orig_sample_rate: Original sample rate. sample_rate: Target sample rate. axis: Axis containing the time dimension. Returns: Audio resampled to ``sample_rate``. The return type matches the input type. """ import math import numpy as np from scipy import signal if orig_sample_rate == sample_rate: return audio audio_np = np.asarray(audio) gcd = math.gcd(int(orig_sample_rate), int(sample_rate)) up = sample_rate // gcd down = orig_sample_rate // gcd # kaiser_best-equivalent anti-aliasing FIR (resampy defaults): a long, # high-attenuation Kaiser sinc designed at the upsampled rate. Cutoff is at # ``rolloff / max(up, down)`` of the upsampled Nyquist. max_rate = max(up, down) num_zeros, rolloff, beta = 64, 0.9475937167399596, 14.769656459379492 fir = signal.firwin( 2 * num_zeros * max_rate + 1, rolloff / max_rate, window=("kaiser", beta), ) resampled = signal.resample_poly( audio_np, up, down, axis=axis, window=fir, padtype="edge", ).astype(np.float32, copy=False) if isinstance(audio, mx.array): return mx.array(resampled) return resampled def trim_silence( audio: Union[mx.array, "np.ndarray"], top_db: float = 20, frame_length: int = 2048, hop_length: int = 512, ): """Trim leading/trailing low-energy regions using a simple RMS gate.""" import numpy as np audio_np = np.asarray(audio) n_frames = 1 + (len(audio_np) - frame_length) // hop_length if n_frames <= 0: return audio rms = np.array( [ np.sqrt( np.mean(audio_np[i * hop_length : i * hop_length + frame_length] ** 2) ) for i in range(n_frames) ] ) rms_db = 20 * np.log10(np.maximum(rms, 1e-10)) threshold = np.max(rms_db) - top_db non_silent = np.where(rms_db >= threshold)[0] if len(non_silent) == 0: return audio start_frame = int(non_silent[0]) end_frame = int(non_silent[-1]) + 1 start_sample = start_frame * hop_length end_sample = min(end_frame * hop_length + frame_length, len(audio_np)) trimmed = audio_np[start_sample:end_sample] if isinstance(audio, mx.array): return mx.array(trimmed.astype(np.float32, copy=False)) return trimmed.astype(np.float32, copy=False) def load_audio( audio: Union[str, mx.array], sample_rate: int = 24000, length: Optional[int] = None, volume_normalize: bool = False, segment_duration: Optional[int] = None, ) -> mx.array: """Load audio from file path or return mx.array as-is. Args: audio: Audio input - can be: - str: Path to audio file (will be loaded and resampled to sample_rate) - mx.array: MLX array (returned as-is) sample_rate: Target sample rate (default 24000) length: Target length in samples (pad or truncate if specified) volume_normalize: Whether to normalize audio volume segment_duration: If specified, randomly select a segment of this duration (seconds) Returns: mx.array: Audio waveform at target sample rate Raises: FileNotFoundError: If audio file path does not exist TypeError: If audio is not str or mx.array """ if isinstance(audio, mx.array): return audio if not isinstance(audio, str): raise TypeError(f"audio must be str or mx.array, got {type(audio)}") import os import numpy as np from mlx_audio.audio_io import read as audio_read if not os.path.exists(audio): raise FileNotFoundError(f"Audio file not found: {audio}") samples, orig_sample_rate = audio_read(audio) shape = samples.shape # Collapse multi channel as mono if len(shape) > 1: samples = samples.sum(axis=1) samples = samples / shape[1] # Resample if needed if sample_rate != orig_sample_rate: samples = resample_audio(samples, orig_sample_rate, sample_rate) # Random segment selection if segment_duration is not None: seg_length = int(sample_rate * segment_duration) samples = random_select_audio_segment(samples, seg_length) # Volume normalization if volume_normalize: samples = audio_volume_normalize(samples) # Length adjustment if length is not None: if samples.shape[0] > length: samples = samples[:length] else: samples = np.pad(samples, (0, int(length - samples.shape[0]))) return mx.array(samples, dtype=mx.float32) __all__ = [ # DSP functions (re-exported from dsp.py) "hanning", "hamming", "blackman", "bartlett", "STR_TO_WINDOW_FN", "stft", "istft", "mel_filters", # Audio utilities "load_audio", "resample_audio", "trim_silence", "audio_volume_normalize", "random_select_audio_segment", # Model utilities "from_dict", "is_valid_module_name", "get_model_category", "get_model_name_parts", "load_model", # Shared loading utilities "get_model_path", "load_config", "load_weights", "apply_quantization", "get_model_class", "base_load_model", ] def is_valid_module_name(name: str) -> bool: """Check if a string is a valid Python module name.""" return isinstance(name, str) and name.isidentifier() and not keyword.iskeyword(name) def _has_model_module(module_path: str) -> bool: try: return importlib.util.find_spec(module_path) is not None except ModuleNotFoundError as exc: missing_name = exc.name or "" if module_path == missing_name or module_path.startswith(f"{missing_name}."): return False raise def get_model_category(model_type: str, model_name: List[str]) -> Optional[str]: """Determine whether a model belongs to the TTS, STT, STS, LID, or VAD category.""" stt_utils = _get_stt_utils() tts_utils = _get_tts_utils() sts_utils = _get_sts_utils() vad_utils = _get_vad_utils() lid_utils = _get_lid_utils() candidates = [model_type] + (model_name or []) categories = [ ("tts", tts_utils.MODEL_REMAPPING), ("stt", stt_utils.MODEL_REMAPPING), ("sts", sts_utils.MODEL_REMAPPING), ("lid", lid_utils.MODEL_REMAPPING), ("vad", vad_utils.MODEL_REMAPPING), ] # If the repo name already includes an explicit category token like "lid", # prefer that category as long as we can resolve any candidate hint there. for category, remap in categories: if category not in candidates: continue for hint in candidates: arch = remap.get(hint, hint) if not is_valid_module_name(arch): continue module_path = f"mlx_audio.{category}.models.{arch}" if _has_model_module(module_path): return category # First pass: check for explicit remapping matches (higher priority) for category, remap in categories: for hint in candidates: if hint in remap: arch = remap[hint] if not is_valid_module_name(arch): continue module_path = f"mlx_audio.{category}.models.{arch}" if _has_model_module(module_path): return category # Second pass: check for direct module matches (fallback) for category, remap in categories: for hint in candidates: if hint not in remap and is_valid_module_name(hint): module_path = f"mlx_audio.{category}.models.{hint}" if _has_model_module(module_path): return category return None def get_model_name_parts(model_path: Union[str, Path]) -> List[str]: model_name = None if isinstance(model_path, str): model_name = model_path.lower().split("/")[-1] elif isinstance(model_path, Path): try: index = model_path.parts.index("hub") model_name = model_path.parts[index + 1].lower().split("--")[-1] except ValueError: model_name = model_path.name.lower() else: raise ValueError(f"Invalid model path type: {type(model_path)}") parts = [] seen = set() dash_parts = [part for part in model_name.split("-") if part] for part in dash_parts: if not part or part in seen: continue parts.append(part) seen.add(part) if "_" in part: for subpart in part.split("_"): if subpart and subpart not in seen: parts.append(subpart) seen.add(subpart) normalized = re.sub(r"[^a-z0-9]+", "", part) if normalized and normalized not in seen: parts.append(normalized) seen.add(normalized) for start in range(len(dash_parts)): for end in range(start + 2, len(dash_parts) + 1): segment = dash_parts[start:end] for combined in ("_".join(segment), "".join(segment)): if combined and combined not in seen: parts.append(combined) seen.add(combined) return parts def load_model(model_name: str): """Load a TTS, STT, STS, LID, or VAD model based on its configuration and name. Args: model_name (str): Name or path of the model to load Returns: The loaded model instance Raises: ValueError: If the model type cannot be determined or is not supported """ model_name_parts = get_model_name_parts(model_name) load_error = None try: config = load_config(model_name) except FileNotFoundError as exc: config = {} load_error = exc # Try to determine model type from config first, then from name model_type = config.get("model_type", None) if model_type is None: model_type = config.get("architecture", None) if model_type is None: model_type = _get_sts_utils().infer_model_type_from_config(config) model_category = get_model_category(model_type, model_name_parts) if not model_category: if load_error is not None: raise load_error raise ValueError(f"Could not determine model type for {model_name}") model_loaders = { "tts": _get_tts_utils().load_model, "stt": _get_stt_utils().load_model, "sts": _get_sts_utils().load_model, "lid": _get_lid_utils().load_model, "vad": _get_vad_utils().load_model, } if model_category not in model_loaders: raise ValueError(f"Model type '{model_category}' not supported") return model_loaders[model_category](model_name)