import time from dataclasses import dataclass from numbers import Number from typing import Dict, Optional, Union import mlx.core as mx import mlx.nn as nn from ..base import BaseModelArgs, GenerationResult, check_array_shape from .istftnet import Decoder from .modules import AlbertModelArgs, CustomAlbert, ProsodyPredictor, TextEncoder from .pipeline import KokoroPipeline def sanitize_lstm_weights(key: str, state_dict: mx.array) -> dict: """Convert PyTorch LSTM weight keys to MLX LSTM weight keys.""" base_key = key.rsplit(".", 1)[0] # Mapping of PyTorch LSTM weight suffixes to MLX weight names weight_map = { "weight_ih_l0_reverse": "Wx_backward", "weight_hh_l0_reverse": "Wh_backward", "bias_ih_l0_reverse": "bias_ih_backward", "bias_hh_l0_reverse": "bias_hh_backward", "weight_ih_l0": "Wx_forward", "weight_hh_l0": "Wh_forward", "bias_ih_l0": "bias_ih_forward", "bias_hh_l0": "bias_hh_forward", } for suffix, new_suffix in weight_map.items(): if key.endswith(suffix): return {f"{base_key}.{new_suffix}": state_dict} return {key: state_dict} @dataclass class ModelConfig(BaseModelArgs): istftnet: dict dim_in: int dropout: float hidden_dim: int max_conv_dim: int max_dur: int multispeaker: bool n_layer: int n_mels: int n_token: int style_dim: int text_encoder_kernel_size: int plbert: dict vocab: Dict[str, int] sample_rate: int = 24000 class Model(nn.Module): """ KokoroModel is a mlx.nn.Module with 2 main responsibilities: 1. Init weights, downloading config.json + model.pth from HF if needed 2. forward(phonemes: str, ref_s: FloatTensor) -> (audio: FloatTensor) You likely only need one KokoroModel instance, and it can be reused across multiple KokoroPipelines to avoid redundant memory allocation. Unlike KokoroPipeline, KokoroModel is language-blind. KokoroModel stores self.vocab and thus knows how to map phonemes -> input_ids, so there is no need to repeatedly download config.json outside of KokoroModel. """ REPO_ID = "prince-canuma/Kokoro-82M" def __init__(self, config: ModelConfig, repo_id: str = None): super().__init__() self.repo_id = repo_id self.config = config self.vocab = config.vocab self.bert = CustomAlbert( AlbertModelArgs(vocab_size=config.n_token, **config.plbert) ) self.bert_encoder = nn.Linear(self.bert.config.hidden_size, config.hidden_dim) self.context_length = self.bert.config.max_position_embeddings self.predictor = ProsodyPredictor( style_dim=config.style_dim, d_hid=config.hidden_dim, nlayers=config.n_layer, max_dur=config.max_dur, dropout=config.dropout, ) self.text_encoder = TextEncoder( channels=config.hidden_dim, kernel_size=config.text_encoder_kernel_size, depth=config.n_layer, n_symbols=config.n_token, ) self.decoder = Decoder( dim_in=config.hidden_dim, style_dim=config.style_dim, dim_out=config.n_mels, **config.istftnet, ) self._pipelines: Dict[str, KokoroPipeline] = {} # Cache for pipelines @dataclass class Output: audio: mx.array pred_dur: Optional[mx.array] = None def __call__( self, phonemes: str, ref_s: mx.array, speed: Number = 1, return_output: bool = False, # MARK: BACKWARD COMPAT decoder=None, ) -> Union["KokoroModel.Output", mx.array]: input_ids = list( filter(lambda i: i is not None, map(lambda p: self.vocab.get(p), phonemes)) ) assert len(input_ids) + 2 <= self.context_length, ( len(input_ids) + 2, self.context_length, ) input_ids = mx.array([[0, *input_ids, 0]]) input_lengths = mx.array([input_ids.shape[-1]]) text_mask = mx.arange(int(input_lengths.max()))[None, ...] text_mask = mx.repeat(text_mask, input_lengths.shape[0], axis=0).astype( input_lengths.dtype ) text_mask = text_mask + 1 > input_lengths[:, None] bert_dur, _ = self.bert(input_ids, attention_mask=(~text_mask).astype(mx.int32)) d_en = self.bert_encoder(bert_dur).transpose(0, 2, 1) ref_s = ref_s s = ref_s[:, 128:] d = self.predictor.text_encoder(d_en, s, input_lengths, text_mask) x, _ = self.predictor.lstm(d) duration = self.predictor.duration_proj(x) max_frames_per_phoneme = 100 duration = mx.sigmoid(duration).sum(axis=-1) / speed duration = mx.nan_to_num( duration, nan=1.0, posinf=max_frames_per_phoneme, neginf=1.0 ) pred_dur = mx.clip( mx.round(duration), a_min=1, a_max=max_frames_per_phoneme ).astype(mx.int32)[0] indices_list = [] for i, n in enumerate(pred_dur): count = min(max(int(n), 0), max_frames_per_phoneme) if count > 0: indices_list.append(mx.repeat(mx.array(i), count)) if not indices_list: silence = mx.zeros((1, 1)) return ( self.Output(audio=silence, pred_dur=pred_dur) if return_output else silence ) indices = mx.concatenate(indices_list) pred_aln_trg = mx.zeros((input_ids.shape[1], indices.shape[0])) pred_aln_trg[indices, mx.arange(indices.shape[0])] = 1 pred_aln_trg = pred_aln_trg[None, :] en = d.transpose(0, 2, 1) @ pred_aln_trg F0_pred, N_pred = self.predictor.F0Ntrain(en, s) t_en = self.text_encoder( input_ids, input_lengths, text_mask ) # Working fine in MLX asr = t_en @ pred_aln_trg decoder = mx.compile(decoder) if decoder is not None else self.decoder audio = decoder(asr, F0_pred, N_pred, ref_s[:, :128])[0] # Working fine in MLX # Evaluate the computation graph for audio and pred_dur before returning mx.eval(audio, pred_dur) return self.Output(audio=audio, pred_dur=pred_dur) if return_output else audio def sanitize(self, weights): sanitized_weights = {} has_packed_quantized_weights = any( key.endswith(".scales") or key.endswith(".biases") for key in weights ) for key, state_dict in weights.items(): if key.startswith("bert"): if "position_ids" in key: # Remove unused position_ids continue else: # print(k, v.shape) sanitized_weights[key] = state_dict if key.startswith("bert_encoder"): sanitized_weights[key] = state_dict if key.startswith("text_encoder"): if key.endswith((".gamma", ".beta")): base_key = key.rsplit(".", 1)[0] if key.endswith(".gamma"): new_key = f"{base_key}.weight" else: new_key = f"{base_key}.bias" sanitized_weights[new_key] = state_dict elif "weight_v" in key: if has_packed_quantized_weights: sanitized_weights[key] = state_dict elif check_array_shape(state_dict): sanitized_weights[key] = state_dict else: sanitized_weights[key] = state_dict.transpose(0, 2, 1) # Replace weight_ih_l0_reverse and weight_hh_l0_reverse with Wx and Wh elif key.endswith( ( ".weight_ih_l0_reverse", ".weight_hh_l0_reverse", ".bias_ih_l0_reverse", ".bias_hh_l0_reverse", ".weight_ih_l0", ".weight_hh_l0", ".bias_ih_l0", ".bias_hh_l0", ) ): sanitized_weights.update(sanitize_lstm_weights(key, state_dict)) else: sanitized_weights[key] = state_dict if key.startswith("predictor"): if "F0_proj.weight" in key: sanitized_weights[key] = ( state_dict if has_packed_quantized_weights else state_dict.transpose(0, 2, 1) ) elif "N_proj.weight" in key: sanitized_weights[key] = ( state_dict if has_packed_quantized_weights else state_dict.transpose(0, 2, 1) ) elif "weight_v" in key: if has_packed_quantized_weights: sanitized_weights[key] = state_dict elif check_array_shape(state_dict): sanitized_weights[key] = state_dict else: sanitized_weights[key] = state_dict.transpose(0, 2, 1) # Replace weight_ih_l0_reverse and weight_hh_l0_reverse with Wx and Wh elif key.endswith( ( ".weight_ih_l0_reverse", ".weight_hh_l0_reverse", ".bias_ih_l0_reverse", ".bias_hh_l0_reverse", ".weight_ih_l0", ".weight_hh_l0", ".bias_ih_l0", ".bias_hh_l0", ) ): sanitized_weights.update(sanitize_lstm_weights(key, state_dict)) else: sanitized_weights[key] = state_dict if key.startswith("decoder"): sanitized_weights[key] = self.decoder.sanitize( key, state_dict, has_packed_quantized_weights=has_packed_quantized_weights, ) return sanitized_weights @property def sample_rate(self): return self.config.sample_rate def _get_pipeline(self, lang_code: str) -> KokoroPipeline: """Retrieves or creates a cached KokoroPipeline for the given language code.""" if lang_code not in self._pipelines: print(f"Creating new KokoroPipeline for language: {lang_code}") self._pipelines[lang_code] = KokoroPipeline( model=self, repo_id=self.REPO_ID if self.repo_id is None else self.repo_id, lang_code=lang_code, ) return self._pipelines[lang_code] def generate( self, text: str, voice: str = None, speed: float = 1.0, lang_code: str = "a", split_pattern: str = r"\n+", **kwargs, ): pipeline = self._get_pipeline(lang_code) pipeline.voices = {} # Reset voices if voice is None: voice = "af_heart" # Track overall generation time start_time = time.time() for segment_idx, (graphenes, phonemes, audio) in enumerate( pipeline(text, voice=voice, speed=speed, split_pattern=split_pattern) ): # Track per-segment generation time now = time.time() segment_time = now - start_time start_time = now samples = audio.shape[0] if audio is not None else 0 assert samples > 0, "No audio generated" # Calculate token count token_count = len(phonemes) if phonemes is not None else 0 # Calculate audio duration in seconds sample_rate = self.config.sample_rate audio_duration_seconds = samples / sample_rate * audio.shape[1] # Calculate real-time factor (RTF) rtf = ( segment_time / audio_duration_seconds if audio_duration_seconds > 0 else 0 ) # 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[0], samples=samples, sample_rate=sample_rate, segment_idx=segment_idx, token_count=token_count, audio_duration=duration_str, real_time_factor=round(rtf, 2), prompt={ "tokens": token_count, "tokens-per-sec": ( round(token_count / segment_time, 2) if segment_time > 0 else 0 ), }, audio_samples={ "samples": samples, "samples-per-sec": ( round(samples / segment_time, 2) if segment_time > 0 else 0 ), }, processing_time_seconds=segment_time, peak_memory_usage=mx.get_peak_memory() / 1e9, ) # Clear cache after each segment to avoid memory leaks mx.clear_cache()