import math from typing import Dict, List, Optional import mlx.core as mx import mlx.nn as nn from .config import VoiceEncConfig from .melspec import melspectrogram def get_num_wins( n_frames: int, step: int, min_coverage: float, hp: VoiceEncConfig, ): """Calculate number of windows and target length for partial utterance splitting.""" assert n_frames > 0 win_size = hp.ve_partial_frames n_wins, remainder = divmod(max(n_frames - win_size + step, 0), step) if n_wins == 0 or (remainder + (win_size - step)) / win_size >= min_coverage: n_wins += 1 target_n = win_size + step * (n_wins - 1) return n_wins, target_n def get_frame_step( overlap: float, rate: Optional[float], hp: VoiceEncConfig, ): """Compute how many frames separate two partial utterances.""" assert 0 <= overlap < 1 if rate is None: frame_step = int(round(hp.ve_partial_frames * (1 - overlap))) else: frame_step = int(round((hp.sample_rate / rate) / hp.ve_partial_frames)) assert 0 < frame_step <= hp.ve_partial_frames return frame_step def sanitize_lstm_weights( key: str, value: mx.array, num_layers: int = 3 ) -> Dict[str, mx.array]: """ Convert PyTorch LSTM weight keys to MLX LSTM weight keys. PyTorch LSTM weight format: lstm.weight_ih_l{layer} -> input-hidden weights for layer lstm.weight_hh_l{layer} -> hidden-hidden weights for layer lstm.bias_ih_l{layer} -> input-hidden bias for layer lstm.bias_hh_l{layer} -> hidden-hidden bias for layer MLX LSTM weight format (for list of LSTM layers): lstm.{layer}.Wx -> input weights lstm.{layer}.Wh -> hidden weights lstm.{layer}.bias -> combined bias (ih + hh) """ result = {} # Extract layer number from key like "lstm.weight_ih_l0" import re match = re.search(r"lstm\.(weight_ih|weight_hh|bias_ih|bias_hh)_l(\d+)", key) if not match: result[key] = value return result weight_type = match.group(1) layer_idx = int(match.group(2)) if weight_type == "weight_ih": result[f"lstm.layers.{layer_idx}.Wx"] = value elif weight_type == "weight_hh": result[f"lstm.layers.{layer_idx}.Wh"] = value elif weight_type == "bias_ih": # MLX combines biases, so we need special handling result[f"lstm.layers.{layer_idx}.bias_ih"] = value elif weight_type == "bias_hh": result[f"lstm.layers.{layer_idx}.bias_hh"] = value return result class StackedLSTM(nn.Module): """Multi-layer LSTM that matches PyTorch's nn.LSTM(num_layers=N).""" def __init__(self, input_size: int, hidden_size: int, num_layers: int = 1): super().__init__() self.input_size = input_size self.hidden_size = hidden_size self.num_layers = num_layers # Create stacked LSTM layers self.layers = [ nn.LSTM(input_size if i == 0 else hidden_size, hidden_size) for i in range(num_layers) ] def __call__(self, x: mx.array, hidden=None): """ Args: x: Input tensor (seq_len, batch, input_size) or (batch, seq_len, input_size) hidden: Tuple of (h_0, c_0), each of shape (num_layers, batch, hidden_size) Returns: output: Output tensor (h_n, c_n): Final hidden and cell states """ if hidden is None: h_list = [None] * self.num_layers c_list = [None] * self.num_layers else: h_0, c_0 = hidden h_list = [h_0[i] for i in range(self.num_layers)] c_list = [c_0[i] for i in range(self.num_layers)] output = x new_h = [] new_c = [] for i, layer in enumerate(self.layers): all_h, all_c = layer(output, hidden=h_list[i], cell=c_list[i]) output = all_h # Extract final timestep: all_h has shape (batch, seq_len, hidden_size) # Use [:, -1, :] to get (batch, hidden_size) for the last timestep new_h.append(all_h[:, -1, :] if all_h.ndim == 3 else all_h) new_c.append(all_c[:, -1, :] if all_c.ndim == 3 else all_c) h_n = mx.stack(new_h, axis=0) c_n = mx.stack(new_c, axis=0) return output, (h_n, c_n) class VoiceEncoder(nn.Module): """LSTM-based voice encoder for speaker embeddings.""" def __init__(self, hp: VoiceEncConfig = None): super().__init__() self.hp = hp or VoiceEncConfig() # Network definition: 3-layer stacked LSTM self.lstm = StackedLSTM(self.hp.num_mels, self.hp.ve_hidden_size, num_layers=3) self.proj = nn.Linear(self.hp.ve_hidden_size, self.hp.speaker_embed_size) # Cosine similarity scaling (fixed initial parameter values) self.similarity_weight = mx.array([10.0]) self.similarity_bias = mx.array([-5.0]) def sanitize(self, weights: Dict[str, mx.array]) -> Dict[str, mx.array]: """ Sanitize PyTorch weights for MLX. Handles: - LSTM weight renaming (weight_ih_l0 -> layers.0.Wx, etc.) - Bias combining (PyTorch has separate ih/hh biases) This method is idempotent - key renaming only affects PyTorch-format keys, MLX-format keys pass through unchanged on subsequent runs. """ new_weights = {} bias_ih = {} bias_hh = {} for key, value in weights.items(): # Handle LSTM weights if "lstm." in key and any( x in key for x in ["weight_ih", "weight_hh", "bias_ih", "bias_hh"] ): import re match = re.search( r"lstm\.(weight_ih|weight_hh|bias_ih|bias_hh)_l(\d+)", key ) if match: weight_type = match.group(1) layer_idx = int(match.group(2)) if weight_type == "weight_ih": new_weights[f"lstm.layers.{layer_idx}.Wx"] = value elif weight_type == "weight_hh": new_weights[f"lstm.layers.{layer_idx}.Wh"] = value elif weight_type == "bias_ih": bias_ih[layer_idx] = value elif weight_type == "bias_hh": bias_hh[layer_idx] = value else: new_weights[key] = value # Combine ih and hh biases (MLX LSTM uses combined bias) for layer_idx in bias_ih: if layer_idx in bias_hh: combined_bias = bias_ih[layer_idx] + bias_hh[layer_idx] new_weights[f"lstm.layers.{layer_idx}.bias"] = combined_bias return new_weights def __call__(self, mels: mx.array) -> mx.array: """ Computes the embeddings of a batch of partial utterances. Args: mels: Batch of unscaled mel spectrograms (B, T, M) where T is hp.ve_partial_frames Returns: Embeddings as (B, E) where E is hp.speaker_embed_size. Embeddings are L2-normed. """ if self.hp.normalized_mels: min_val = mx.min(mels) max_val = mx.max(mels) if float(min_val) < 0 or float(max_val) > 1: raise Exception(f"Mels outside [0, 1]. Min={min_val}, Max={max_val}") # Pass full sequence through LSTM layers (vectorized, no per-timestep loop) _, (h_n, _) = self.lstm(mels) # Get final hidden state from last layer final_hidden = h_n[-1] # (B, H) # Project raw_embeds = self.proj(final_hidden) # Apply ReLU if configured if self.hp.ve_final_relu: raw_embeds = nn.relu(raw_embeds) # L2 normalize embeds = raw_embeds / mx.linalg.norm(raw_embeds, axis=1, keepdims=True) return embeds def inference( self, mels: mx.array, mel_lens: List[int], overlap: float = 0.5, rate: Optional[float] = None, min_coverage: float = 0.8, batch_size: Optional[int] = None, ) -> mx.array: """ Computes embeddings of a batch of full utterances. Args: mels: (B, T, M) unscaled mels mel_lens: List of mel lengths for each batch item overlap: Overlap between partial windows rate: Rate for frame step calculation min_coverage: Minimum coverage for partial windows batch_size: Batch size for processing partials Returns: (B, E) embeddings """ # Compute where to split the utterances into partials frame_step = get_frame_step(overlap, rate, self.hp) n_partials_list = [] target_lens = [] for l in mel_lens: n_p, t_l = get_num_wins(l, frame_step, min_coverage, self.hp) n_partials_list.append(n_p) target_lens.append(t_l) # Possibly pad the mels to reach the target lengths len_diff = max(target_lens) - mels.shape[1] if len_diff > 0: pad = mx.zeros((mels.shape[0], len_diff, self.hp.num_mels)) mels = mx.concatenate([mels, pad], axis=1) # Group all partials together (vectorized extraction) # For each mel, extract all partials at once using mx.take partial_list = [] for mel, n_partial in zip(mels, n_partials_list): if n_partial > 0: # Create indices for all partials at once: (n_partial, ve_partial_frames) partial_starts = mx.arange(n_partial) * frame_step # (n_partial,) frame_offsets = mx.arange( self.hp.ve_partial_frames ) # (ve_partial_frames,) # Broadcast to get all indices: (n_partial, ve_partial_frames) indices = partial_starts[:, None] + frame_offsets[None, :] # Extract all partials at once for each mel dimension # mel is (T, M), we need (n_partial, ve_partial_frames, M) mel_partials = mx.take(mel, indices.flatten(), axis=0).reshape( n_partial, self.hp.ve_partial_frames, mel.shape[1] ) partial_list.append(mel_partials) partials = mx.concatenate(partial_list, axis=0) # (total_partials, T, M) # Forward the partials (in batches if needed) if batch_size is None or batch_size >= len(partials): partial_embeds = self(partials) else: embed_chunks = [] for i in range(0, len(partials), batch_size): chunk = partials[i : i + batch_size] embed_chunks.append(self(chunk)) partial_embeds = mx.concatenate(embed_chunks, axis=0) # Reduce the partial embeds into full embeds # Compute slice indices without numpy slices = [0] for n in n_partials_list: slices.append(slices[-1] + n) raw_embeds = [] for start, end in zip(slices[:-1], slices[1:]): raw_embeds.append(mx.mean(partial_embeds[start:end], axis=0)) raw_embeds = mx.stack(raw_embeds) # L2-normalize the final embeds embeds = raw_embeds / mx.linalg.norm(raw_embeds, axis=1, keepdims=True) return embeds @staticmethod def utt_to_spk_embed(utt_embeds: mx.array) -> mx.array: """ Takes L2-normalized utterance embeddings, computes mean and L2-normalizes to get a speaker embedding. """ assert len(utt_embeds.shape) == 2 utt_embeds = mx.mean(utt_embeds, axis=0) return utt_embeds / mx.linalg.norm(utt_embeds) @staticmethod def voice_similarity(embeds_x: mx.array, embeds_y: mx.array) -> float: """Cosine similarity for L2-normalized embeddings.""" if len(embeds_x.shape) != 1: embeds_x = VoiceEncoder.utt_to_spk_embed(embeds_x) if len(embeds_y.shape) != 1: embeds_y = VoiceEncoder.utt_to_spk_embed(embeds_y) return float(embeds_x @ embeds_y) def embeds_from_mels( self, mels: List[mx.array], mel_lens: Optional[List[int]] = None, as_spk: bool = False, batch_size: int = 32, **kwargs, ) -> mx.array: """ Convenience function for deriving utterance or speaker embeddings from mel spectrograms. Args: mels: List of (Ti, M) arrays or stacked (B, T, M) array mel_lens: Individual mel lengths if mels is stacked as_spk: Whether to return speaker embedding (single) or utterance embeddings (per-utt) batch_size: Batch size for processing **kwargs: Additional args for inference() Returns: (B, E) embeddings if as_spk is False, else (E,) speaker embedding """ # Handle list of mels if isinstance(mels, list): mel_lens = [mel.shape[0] for mel in mels] max_len = max(mel_lens) # Pad and stack padded = [] for mel in mels: if mel.shape[0] < max_len: pad = mx.zeros((max_len - mel.shape[0], mel.shape[1])) mel = mx.concatenate([mel, pad], axis=0) padded.append(mel) mels = mx.stack(padded) # Embed them utt_embeds = self.inference(mels, mel_lens, batch_size=batch_size, **kwargs) return self.utt_to_spk_embed(utt_embeds) if as_spk else utt_embeds def embeds_from_wavs( self, wavs: List[mx.array], sample_rate: int, as_spk: bool = False, batch_size: int = 32, trim_top_db: Optional[float] = 20, **kwargs, ) -> mx.array: """ Wrapper around embeds_from_mels that first converts waveforms to mel spectrograms. Args: wavs: List of waveform arrays sample_rate: Sample rate of the waveforms as_spk: Whether to return speaker embedding batch_size: Batch size for processing trim_top_db: Trim silence below this threshold (None to disable) **kwargs: Additional args for inference() Returns: Embeddings """ import scipy.signal # Resample if needed if sample_rate != self.hp.sample_rate: import numpy as np resampled_wavs = [] for wav in wavs: # Convert to numpy for resampling wav_np = np.array(wav) # Calculate resampling ratio gcd = math.gcd(sample_rate, self.hp.sample_rate) up = self.hp.sample_rate // gcd down = sample_rate // gcd wav_resampled = scipy.signal.resample_poly(wav_np, up, down) resampled_wavs.append(mx.array(wav_resampled)) wavs = resampled_wavs # Trim silence if requested if trim_top_db is not None: import numpy as np trimmed_wavs = [] for wav in wavs: wav_np = np.array(wav) # Simple energy-based trimming # Calculate frame energy frame_length = 2048 hop_length = 512 # Compute RMS energy per frame n_frames = 1 + (len(wav_np) - frame_length) // hop_length if n_frames > 0: rms = np.array( [ math.sqrt( np.mean( wav_np[ i * hop_length : i * hop_length + frame_length ] ** 2 ) ) for i in range(n_frames) ] ) # Convert to dB rms_db = 20 * np.log10(np.maximum(rms, 1e-10)) threshold = np.max(rms_db) - trim_top_db # Find non-silent frames non_silent = np.where(rms_db >= threshold)[0] if len(non_silent) > 0: start_frame = non_silent[0] end_frame = non_silent[-1] + 1 start_sample = start_frame * hop_length end_sample = min( end_frame * hop_length + frame_length, len(wav_np) ) wav_np = wav_np[start_sample:end_sample] trimmed_wavs.append(mx.array(wav_np)) wavs = trimmed_wavs # Set default rate if not provided if "rate" not in kwargs: kwargs["rate"] = 1.3 # Resemble's default value # Convert waveforms to mel spectrograms mels = [] for wav in wavs: mel = melspectrogram(wav, self.hp) # Transpose from (M, T) to (T, M) mel = mx.transpose(mel) mels.append(mel) return self.embeds_from_mels( mels, as_spk=as_spk, batch_size=batch_size, **kwargs )