# Copyright (c) 2025, Prince Canuma and contributors (https://github.com/Blaizzy/mlx-audio) import time from dataclasses import dataclass from pathlib import Path from typing import Dict, Generator, Optional, Union import mlx.core as mx import mlx.nn as nn import numpy as np from mlx_audio.utils import load_audio, resample_audio from ..base import GenerationResult from .config import ModelConfig from .s3gen import S3Token2Wav from .s3tokenizer import S3TokenizerV2, log_mel_spectrogram from .t3 import T3 from .t3.cond_enc import T3Cond from .tokenizer import MTLTokenizer from .voice_encoder import VoiceEncoder # Constants S3_SR = 16000 # Sample rate for speech tokenizer S3GEN_SR = 24000 # Sample rate for vocoder SPEECH_VOCAB_SIZE = 6561 # Size of speech token vocabulary # Special tokens SOT = "[START]" EOT = "[STOP]" SPACE = "[SPACE]" def punc_norm(text: str) -> str: """ Quick cleanup func for punctuation from LLMs or containing chars not seen often in the dataset. """ if len(text) == 0: return "You need to add some text for me to talk." # Capitalise first letter if text[0].islower(): text = text[0].upper() + text[1:] # Remove multiple space chars text = " ".join(text.split()) # Replace uncommon/llm punc punc_to_replace = [ ("...", ", "), ("…", ", "), (":", ","), (" - ", ", "), (";", ", "), ("—", "-"), ("–", "-"), (" ,", ","), ( """, "\""), (""", '"', ), ("'", "'"), ("'", "'"), ] for old_char_sequence, new_char in punc_to_replace: text = text.replace(old_char_sequence, new_char) # Add full stop if no ending punc text = text.rstrip(" ") sentence_enders = {".", "!", "?", "-", ","} if not any(text.endswith(p) for p in sentence_enders): text += "." return text def drop_invalid_tokens(x: mx.array) -> mx.array: """ Drop SOS and EOS tokens, extracting only the speech content between them. This matches the original PyTorch implementation which slices between the start-of-speech (SOS=6561) and end-of-speech (EOS=6562) markers. """ SOS = SPEECH_VOCAB_SIZE # 6561 EOS = SPEECH_VOCAB_SIZE + 1 # 6562 assert len(x.shape) == 1 or ( len(x.shape) == 2 and x.shape[0] == 1 ), "only batch size of one allowed for now" x_flat = x.flatten() # Find SOS position using argmax (returns first True index) sos_mask = x_flat == SOS s = 0 if mx.any(sos_mask): s = int(mx.argmax(sos_mask)) + 1 # Start after SOS # Find EOS position using argmax eos_mask = x_flat == EOS e = x_flat.shape[0] if mx.any(eos_mask): e = int(mx.argmax(eos_mask)) # End before EOS return x_flat[s:e] @dataclass class Conditionals: """ Conditionals for T3 and S3Gen. T3 conditionals: - speaker_emb - cond_prompt_speech_tokens - emotion_adv S3Gen conditionals: - prompt_token - prompt_token_len - prompt_feat - prompt_feat_len - embedding """ t3: T3Cond gen: dict class Model(nn.Module): """ Chatterbox Text-to-Speech model. This integrates: - T3: LLaMA-based text-to-speech-token generator - S3Gen: Flow matching decoder with HiFi-GAN vocoder - VoiceEncoder: Speaker embedding extractor - S3Tokenizer: Speech tokenizer for reference audio - Text Tokenizer: Text to token ID conversion Usage: model = Model.from_pretrained("path/to/weights") audio = model.generate("Hello world!", audio_prompt_path="reference.wav") """ ENC_COND_LEN = 6 * S3_SR # 6 seconds at 16kHz for encoder DEC_COND_LEN = 10 * S3GEN_SR # 10 seconds at 24kHz for decoder def __init__( self, config_or_t3: Union[ModelConfig, T3] = None, s3gen: Optional[S3Token2Wav] = None, ve: Optional[VoiceEncoder] = None, conds: Optional[Conditionals] = None, ): super().__init__() self.sr = S3GEN_SR # sample rate of synthesized audio # Check if first argument is a config if config_or_t3 is None or isinstance(config_or_t3, ModelConfig): # Initialize from config config = config_or_t3 or ModelConfig() self.config = config self.t3 = T3(config.t3_config) self.s3gen = S3Token2Wav() self.ve = VoiceEncoder() self._conds = None else: # Initialize with individual components self.config = None self.t3 = config_or_t3 self.s3gen = s3gen self.ve = ve self._conds = conds # S3 tokenizer for speech token extraction (initialized lazily or during load_weights) self._s3_tokenizer = S3TokenizerV2("speech_tokenizer_v2_25hz") # Text tokenizer (initialized during load_weights if model_path is available) self.tokenizer = None self.mtl_tokenizer = None def sanitize(self, weights: Dict[str, mx.array]) -> Dict[str, mx.array]: """ Sanitize PyTorch weights for MLX. This routes weights to the appropriate component's sanitize method based on the weight key prefix. S3Tokenizer weights are loaded separately from mlx-community/S3TokenizerV2. Handles two cases: 1. Pre-prefixed weights (from MLX converted models): - ve.* -> VoiceEncoder weights - t3.* -> T3 model weights - s3gen.* -> S3Gen (S3Token2Wav) weights 2. Original PyTorch weights (without prefix): - Infers component from key names Each component's sanitize method handles: - Conv1d/Conv2d weight transposition - LSTM weight renaming (for VoiceEncoder) - Perceiver attention weight renaming (for T3) - ConvTranspose weight transposition (for S3Gen) """ new_weights = {} # Separate weights by component (S3Tokenizer loaded separately) ve_weights = {} t3_weights = {} s3gen_weights = {} other_weights = {} for key, value in weights.items(): if key.startswith("ve."): # Remove prefix for component sanitize ve_weights[key[3:]] = value elif key.startswith("t3."): t3_weights[key[3:]] = value elif key.startswith("s3gen."): s3gen_weights[key[6:]] = value else: # If no prefix, infer from key names # VoiceEncoder keys: lstm.*, similarity_*, proj.* if ( key.startswith("lstm.") or key.startswith("similarity") or key == "proj.weight" or key == "proj.bias" ): ve_weights[key] = value # T3 keys: tfmr.*, text_emb.*, speech_emb.*, text_head.*, speech_head.*, # perceiver.*, cond_emb.*, prompt_pos_emb.*, cond_enc.*, text_pos_emb.*, speech_pos_emb.* elif ( key.startswith("tfmr.") or key.startswith("text_emb.") or key.startswith("speech_emb.") or key.startswith("text_head.") or key.startswith("speech_head.") or key.startswith("perceiver.") or key.startswith("cond_emb.") or key.startswith("prompt_pos_emb.") or key.startswith("cond_enc.") or key.startswith("text_pos_emb.") or key.startswith("speech_pos_emb.") ): t3_weights[key] = value # S3Gen keys: flow.*, mel2wav.*, speaker_encoder.*, f0_predictor.* elif ( key.startswith("flow.") or key.startswith("mel2wav.") or key.startswith("speaker_encoder.") or key.startswith("f0_predictor.") ): s3gen_weights[key] = value else: other_weights[key] = value # Sanitize each component's weights if ve_weights: ve_sanitized = self.ve.sanitize(ve_weights) for k, v in ve_sanitized.items(): new_weights[f"ve.{k}"] = v if t3_weights: t3_sanitized = self.t3.sanitize(t3_weights) for k, v in t3_sanitized.items(): new_weights[f"t3.{k}"] = v if s3gen_weights: s3gen_sanitized = self.s3gen.sanitize(s3gen_weights) for k, v in s3gen_sanitized.items(): new_weights[f"s3gen.{k}"] = v # Add other weights as-is new_weights.update(other_weights) return new_weights def load_weights( self, weights, strict: bool = True, ): """ Load weights into the model. Uses strict=False by default for components because Chatterbox has several non-checkpoint parameters (rand_noise, pos_enc.pe, stft_window, trim_fade) that are generated during initialization. Args: weights: List of (key, value) tuples or dict strict: If False, ignore missing/extra keys. Default True for compatibility with utils.load_model(). """ if isinstance(weights, dict): weights = list(weights.items()) # Split weights by component prefix (S3Tokenizer loaded separately) ve_weights = [] t3_weights = [] s3gen_weights = [] other_weights = [] for k, v in weights: if k.startswith("ve."): ve_weights.append((k[3:], v)) elif k.startswith("t3."): t3_weights.append((k[3:], v)) elif k.startswith("s3gen."): s3gen_weights.append((k[6:], v)) elif k.startswith("gen."): # Skip gen.* keys - these are conditionals, not model weights continue else: other_weights.append((k, v)) # Load each component with strict=False to handle non-checkpoint params if ve_weights: self.ve.load_weights(ve_weights, strict=False) if t3_weights: self.t3.load_weights(t3_weights, strict=False) if s3gen_weights: self.s3gen.load_weights(s3gen_weights, strict=False) # Handle any remaining weights at the top level if other_weights and strict: raise ValueError( f"Unrecognized weight keys: {[k for k, v in other_weights]}" ) @classmethod def from_pretrained( cls, ckpt_dir: Union[str, Path], s3_tokenizer_repo: str = "mlx-community/S3TokenizerV2", ) -> "Model": """ Load a pretrained Chatterbox model from a checkpoint directory. Expects the standard mlx-audio format: a single model.safetensors file with component prefixes (ve.*, t3.*, s3gen.*). S3Tokenizer weights are loaded separately from a shared repository. Automatically handles quantized weights if config.json contains quantization info. Use scripts/convert.py to convert from original PyTorch weights. Args: ckpt_dir: Path to checkpoint directory containing model.safetensors and tokenizer.json. s3_tokenizer_repo: Hugging Face repo for S3Tokenizer weights. Default: "mlx-community/S3TokenizerV2" Returns: Initialized ChatterboxTTS model """ import json from huggingface_hub import snapshot_download ckpt_dir = Path(ckpt_dir) # Download from Hub if path doesn't exist locally if not ckpt_dir.exists(): print(f"Downloading {ckpt_dir} from Hugging Face...") ckpt_dir = Path( snapshot_download( repo_id=str(ckpt_dir), allow_patterns=[ "model.safetensors", "tokenizer.json", "config.json", ], ) ) # Initialize models with default config model = cls() print("Loading MLX weights...") combined_path = ckpt_dir / "model.safetensors" if not combined_path.exists(): raise FileNotFoundError( f"model.safetensors not found in {ckpt_dir}. " "Use scripts/convert.py to convert from PyTorch weights." ) # Load config to check for quantization config_path = ckpt_dir / "config.json" config = {} if config_path.exists(): with open(config_path) as f: config = json.load(f) all_weights = mx.load(str(combined_path)) # Split weights by prefix (S3Tokenizer loaded separately) ve_weights = {k[3:]: v for k, v in all_weights.items() if k.startswith("ve.")} t3_weights = {k[3:]: v for k, v in all_weights.items() if k.startswith("t3.")} s3gen_weights = { k[6:]: v for k, v in all_weights.items() if k.startswith("s3gen.") } # Handle quantization if present in config if quantization := config.get("quantization"): print(f"Detected {quantization['bits']}-bit quantization...") # Quantize the model before loading weights # This converts Linear layers to QuantizedLinear so they can accept packed weights def should_quantize(path, module): """Check if this module should be quantized based on weights.""" if not hasattr(module, "to_quantized"): return False # Check if we have quantized weights (scales) for this path # Need to check in the appropriate weight dict based on prefix if path.startswith("ve."): weight_path = path[3:] # Remove "ve." prefix return f"{weight_path}.scales" in ve_weights elif path.startswith("t3."): weight_path = path[3:] # Remove "t3." prefix return f"{weight_path}.scales" in t3_weights elif path.startswith("s3gen."): weight_path = path[6:] # Remove "s3gen." prefix return f"{weight_path}.scales" in s3gen_weights return False nn.quantize( model, group_size=quantization["group_size"], bits=quantization["bits"], class_predicate=should_quantize, ) def load_component_weights(component, weights, strict: bool = True): """Load weights dict into component.""" if hasattr(component, "sanitize"): weights = component.sanitize(weights) component.load_weights(list(weights.items()), strict=strict) if ve_weights: load_component_weights(model.ve, ve_weights, strict=True) if t3_weights: load_component_weights(model.t3, t3_weights, strict=False) if s3gen_weights: load_component_weights(model.s3gen, s3gen_weights, strict=False) # Load S3Tokenizer from separate repo print(f"Loading S3Tokenizer from {s3_tokenizer_repo}...") s3tok_dir = Path( snapshot_download( repo_id=s3_tokenizer_repo, allow_patterns=["model.safetensors", "config.json"], ) ) s3tok_path = s3tok_dir / "model.safetensors" if not s3tok_path.exists(): raise FileNotFoundError( f"model.safetensors not found in {s3_tokenizer_repo}. " "S3Tokenizer weights are required for Chatterbox." ) s3tok_weights = mx.load(str(s3tok_path)) model._s3_tokenizer = S3TokenizerV2("speech_tokenizer_v2_25hz") load_component_weights(model._s3_tokenizer, s3tok_weights, strict=False) # Initialize text tokenizer (check config for multilingual setting) tokenizer_path = ckpt_dir / "tokenizer.json" config = None if tokenizer_path.exists(): try: import json from .tokenizer import EnTokenizer, MTLTokenizer # Check if multilingual model from config.json config_path = ckpt_dir / "config.json" if config_path.exists(): with open(config_path) as f: config = json.load(f) if config and config.get("multilingual", False): model.mtl_tokenizer = MTLTokenizer(tokenizer_path) print("Loaded multilingual tokenizer (MTLTokenizer)") model.tokenizer = EnTokenizer(tokenizer_path) print("Loaded English tokenizer (EnTokenizer)") except ImportError: print("Warning: tokenizers library not available") model.tokenizer = None else: print("Warning: tokenizer.json not found") model.tokenizer = None # Set to eval mode for inference (important for BatchNorm) model.eval() print("Model loaded successfully!") return model @staticmethod def post_load_hook(model: "Model", model_path: Path) -> "Model": """ Post-load hook called by load_model to initialize tokenizer and conditionals. Args: model: The loaded model instance model_path: Path to the model directory Returns: The model with tokenizer and conditionals initialized """ # Load text tokenizer (check config for multilingual setting) tokenizer_path = model_path / "tokenizer.json" config = None if tokenizer_path.exists(): try: import json from .tokenizer import EnTokenizer, MTLTokenizer # Check if multilingual model from config..json config_path = model_path / "config.json" if config_path.exists(): with open(config_path) as f: config = json.load(f) if config and config.get("multilingual", False): model.mtl_tokenizer = MTLTokenizer(tokenizer_path) print("Loaded multilingual tokenizer (MTLTokenizer)") model.tokenizer = EnTokenizer(tokenizer_path) print("Loaded English tokenizer (EnTokenizer)") except ImportError: print("Warning: tokenizers library not available") model.tokenizer = None else: print(f"Warning: tokenizer.json not found at {tokenizer_path}") model.tokenizer = None # Load S3Tokenizer from separate repo from huggingface_hub import snapshot_download s3_tokenizer_repo = "mlx-community/S3TokenizerV2" print(f"Loading S3Tokenizer from {s3_tokenizer_repo}...") s3tok_dir = Path( snapshot_download( repo_id=s3_tokenizer_repo, allow_patterns=["model.safetensors", "config.json"], ) ) s3tok_path = s3tok_dir / "model.safetensors" if s3tok_path.exists(): s3tok_weights = mx.load(str(s3tok_path)) model._s3_tokenizer = S3TokenizerV2("speech_tokenizer_v2_25hz") if hasattr(model._s3_tokenizer, "sanitize"): s3tok_weights = model._s3_tokenizer.sanitize(s3tok_weights) model._s3_tokenizer.load_weights(list(s3tok_weights.items()), strict=False) print("Loaded S3Tokenizer weights") else: print(f"Warning: S3Tokenizer weights not found at {s3tok_path}") # Load pre-computed conditionals from conds.safetensors conds_path = model_path / "conds.safetensors" if conds_path.exists(): conds_data = mx.load(str(conds_path)) print("Loaded pre-computed conditionals from conds.safetensors") # Extract T3 conditionals speaker_emb = conds_data.get("t3.speaker_emb") if speaker_emb is None: speaker_emb = mx.zeros((1, 256)) cond_tokens = conds_data.get("t3.cond_prompt_speech_tokens") emotion_adv = conds_data.get("t3.emotion_adv") if emotion_adv is None: emotion_adv = mx.ones((1, 1, 1)) * 0.5 t3_cond = T3Cond( speaker_emb=speaker_emb, cond_prompt_speech_tokens=cond_tokens, emotion_adv=emotion_adv, ) # Extract gen conditionals gen_dict = {} for k, v in conds_data.items(): if k.startswith("gen."): gen_dict[k.replace("gen.", "")] = v # Compute prompt_feat_len if missing if "prompt_feat_len" not in gen_dict and "prompt_feat" in gen_dict: prompt_feat = gen_dict["prompt_feat"] gen_dict["prompt_feat_len"] = mx.array([prompt_feat.shape[1]]) model._conds = Conditionals(t3_cond, gen_dict) else: print("Warning: conds.safetensors not found - ref_audio will be required") model._conds = None return model def prepare_conditionals( self, ref_wav: Union[str, mx.array, np.ndarray], ref_sr: int, exaggeration: float = 0.5, ) -> Conditionals: """ Prepare conditioning from a reference audio clip. Args: ref_wav: Reference waveform (samples,) or (1, samples) ref_sr: Reference sample rate exaggeration: Emotion exaggeration factor (0-1) Returns: Conditionals object with T3 and S3Gen conditioning Note: Following the original PyTorch implementation: - S3Gen uses up to 10s of audio (DEC_COND_LEN at 24kHz) - T3 encoder uses up to 6s of audio (ENC_COND_LEN at 16kHz) - S3Gen tokens are computed from 24kHz audio resampled to 16kHz - T3 tokens are computed from 16kHz audio (resampled from original) """ # Ensure 1D waveform if isinstance(ref_wav, str): ref_wav = load_audio(ref_wav, sample_rate=S3GEN_SR) ref_sr = S3GEN_SR if ref_wav.ndim == 2: ref_wav = ref_wav.squeeze(0) # Resample to 24kHz for S3Gen (mel extraction and base for speaker embedding) if ref_sr != S3GEN_SR: ref_wav_24k = resample_audio(ref_wav, ref_sr, S3GEN_SR) else: ref_wav_24k = ref_wav # Truncate to decoder conditioning length (10s) ref_wav_24k = ref_wav_24k[: self.DEC_COND_LEN] # Resample 24kHz to 16kHz for S3Gen tokenization and speaker embedding # This matches the original which resamples from 24kHz ref_wav_16k_from_24k = resample_audio(ref_wav_24k, S3GEN_SR, S3_SR) # Resample original to 16kHz for T3 encoder conditioning and VE embedding if ref_sr != S3_SR: ref_wav_16k_full = resample_audio(ref_wav, ref_sr, S3_SR) else: ref_wav_16k_full = ref_wav # Truncate to encoder conditioning length (6s) for T3 tokens only ref_wav_16k = ref_wav_16k_full[: self.ENC_COND_LEN] # Get S3Gen reference embeddings s3gen_ref_dict = {} t3_cond_prompt_tokens = None if self._s3_tokenizer is not None: # --- S3Gen tokens (from 10s audio, resampled 24k->16k) --- s3gen_mel = log_mel_spectrogram(ref_wav_16k_from_24k) s3gen_mel = mx.expand_dims(s3gen_mel, 0) # Add batch dim s3gen_mel_len = mx.array([s3gen_mel.shape[2]]) s3gen_tokens, s3gen_token_lens = self._s3_tokenizer( s3gen_mel, s3gen_mel_len ) # Get S3Gen embeddings (tokens may be truncated by embed_ref to match mel) s3gen_ref_dict = self.s3gen.embed_ref( ref_wav=mx.expand_dims(ref_wav_24k, 0), ref_sr=S3GEN_SR, ref_speech_tokens=s3gen_tokens, ref_speech_token_lens=s3gen_token_lens, ) # --- T3 conditioning tokens (from 6s audio, limited to plen) --- t3_mel = log_mel_spectrogram(ref_wav_16k) t3_mel = mx.expand_dims(t3_mel, 0) t3_mel_len = mx.array([t3_mel.shape[2]]) t3_tokens, t3_token_lens = self._s3_tokenizer(t3_mel, t3_mel_len) # Limit T3 tokens to prompt length plen = self.t3.hp.speech_cond_prompt_len if hasattr(self.t3, "hp") else 150 t3_cond_prompt_tokens = t3_tokens[:, :plen] # Voice encoder speaker embedding (from full 16kHz audio, not truncated) ve_embed = self.ve.embeds_from_wavs([ref_wav_16k_full], sample_rate=S3_SR) ve_embed = mx.mean(ve_embed, axis=0, keepdims=True) t3_cond = T3Cond( speaker_emb=ve_embed, cond_prompt_speech_tokens=t3_cond_prompt_tokens, emotion_adv=mx.ones((1, 1, 1)) * exaggeration, ) return Conditionals(t3_cond, s3gen_ref_dict) @property def sample_rate(self) -> int: """Output sample rate.""" return S3GEN_SR def generate( self, text: str, audio_prompt: Optional[mx.array] = None, audio_prompt_sr: Optional[int] = None, conds: Optional[Conditionals] = None, exaggeration: float = 0.1, cfg_weight: float = 0.5, temperature: float = 0.8, repetition_penalty: float = 1.2, min_p: float = 0.05, top_p: float = 1.0, max_new_tokens: int = 1000, # Standard mlx_audio.tts.generate parameters (mapped to Chatterbox params) ref_audio: Optional[Union[str, mx.array, np.ndarray]] = None, voice: Optional[str] = None, speed: float = 1.0, lang_code: str = "en", max_tokens: int = None, verbose: bool = True, stream: bool = False, streaming_interval: float = 2.0, **kwargs, ) -> Generator[GenerationResult, None, None]: """ Generate speech from text. Args: text: Input text to synthesize audio_prompt: Reference audio for voice cloning (optional) audio_prompt_sr: Sample rate of audio prompt conds: Pre-computed conditionals (optional) exaggeration: Emotion exaggeration factor (0-1) cfg_weight: Classifier-free guidance weight temperature: Sampling temperature repetition_penalty: Penalty for repeated tokens min_p: Minimum probability threshold top_p: Top-p (nucleus) sampling threshold max_new_tokens: Maximum number of tokens to generate ref_audio: Alias for audio_prompt (for mlx_audio.tts.generate compatibility) voice: Ignored (Chatterbox uses reference audio for voice cloning) speed: Ignored (Chatterbox doesn't support speed adjustment) lang_code: Ignored (Chatterbox is English-only) max_tokens: Alias for max_new_tokens verbose: Whether to print verbose output stream: Ignored (Chatterbox doesn't support streaming) streaming_interval: Ignored Yields: GenerationResult with generated audio waveform """ start_time = time.time() # Handle parameter aliases for mlx_audio.tts.generate compatibility if ref_audio is not None and audio_prompt is None: audio_prompt = ref_audio audio_prompt_sr = ( self.sample_rate ) # Assume ref_audio is already at correct sample rate if max_tokens is not None and max_new_tokens == 1000: max_new_tokens = max_tokens # Prepare conditionals if needed if conds is None: if audio_prompt is not None and audio_prompt_sr is not None: conds = self.prepare_conditionals( audio_prompt, audio_prompt_sr, exaggeration ) elif self._conds is not None: conds = self._conds else: raise ValueError( "No conditionals available. Either provide audio_prompt/audio_prompt_sr " "for voice cloning, or ensure conds.safetensors is in the model directory." ) # Update exaggeration if needed if exaggeration != float(conds.t3.emotion_adv[0, 0, 0]): conds.t3.emotion_adv = mx.ones((1, 1, 1)) * exaggeration # Normalize and tokenize text text = punc_norm(text) try: if lang_code == "en": text_tokens = self.tokenizer.text_to_tokens(text) elif isinstance(self.mtl_tokenizer, MTLTokenizer): text_tokens = self.mtl_tokenizer.text_to_tokens( text, language_id=lang_code ) else: if self.tokenizer is None and self.mtl_tokenizer is None: raise ValueError( "Text tokenizer or multilingual tokenizer not initialized.\n" "Load model with from_pretrained() or set model.tokenizer manually.\n" ) else: raise ValueError( "Invalid language code. Supported languages: " "ar (Arabic), da (Danish), de (German), el (Greek), en (English), " "es (Spanish), fi (Finnish), fr (French), he (Hebrew), hi (Hindi), " "it (Italian), ja (Japanese), ko (Korean), ms (Malay), nl (Dutch), " "no (Norwegian), pl (Polish), pt (Portuguese), ru (Russian), " "sv (Swedish), sw (Swahili), tr (Turkish), zh (Chinese)" ) except Exception as e: print(f"Error tokenizing text: {e}") raise e token_count = text_tokens.shape[1] if cfg_weight > 0.0: text_tokens = mx.concatenate([text_tokens, text_tokens], axis=0) # Add start/end tokens sot = self.t3.hp.start_text_token eot = self.t3.hp.stop_text_token # Pad with start token sot_tokens = mx.full((text_tokens.shape[0], 1), sot, dtype=mx.int32) eot_tokens = mx.full((text_tokens.shape[0], 1), eot, dtype=mx.int32) text_tokens = mx.concatenate([sot_tokens, text_tokens, eot_tokens], axis=1) # Clear any accumulated cache from previous generations mx.clear_cache() # Generate speech tokens with T3 speech_tokens = self.t3.inference( t3_cond=conds.t3, text_tokens=text_tokens, max_new_tokens=max_new_tokens, temperature=temperature, cfg_weight=cfg_weight, repetition_penalty=repetition_penalty, min_p=min_p, top_p=top_p, ) # Clear cache after T3 inference to free memory before S3Gen mx.clear_cache() # Extract conditional batch (first in CFG pair) speech_tokens = speech_tokens[0:1] # Drop invalid tokens speech_tokens = drop_invalid_tokens(speech_tokens) # Filter out tokens >= SPEECH_VOCAB_SIZE mask = speech_tokens < SPEECH_VOCAB_SIZE # Use argsort to get indices of valid tokens (preserves order via stable sort) valid_count = int(mx.sum(mask.astype(mx.int32))) sorted_indices = mx.argsort(-mask.astype(mx.int32)) valid_indices = sorted_indices[:valid_count] speech_tokens = mx.take(speech_tokens, valid_indices) # Reshape for S3Gen speech_tokens = mx.expand_dims(speech_tokens, 0) # Generate waveform with S3Gen wav = self.s3gen( speech_tokens=speech_tokens, ref_dict=conds.gen, finalize=True, ) # Flatten to 1D if needed if wav.ndim == 2: wav = wav.squeeze(0) # Calculate timing and metrics processing_time = time.time() - start_time samples = wav.shape[0] audio_duration_seconds = samples / self.sample_rate # Format duration as HH:MM:SS.mmm duration_hours = int(audio_duration_seconds // 3600) duration_mins = int((audio_duration_seconds % 3600) // 60) duration_secs = int(audio_duration_seconds % 60) duration_ms = int((audio_duration_seconds % 1) * 1000) duration_str = f"{duration_hours:02d}:{duration_mins:02d}:{duration_secs:02d}.{duration_ms:03d}" # Calculate real-time factor rtf = ( processing_time / audio_duration_seconds if audio_duration_seconds > 0 else 0 ) yield GenerationResult( audio=wav, samples=samples, sample_rate=self.sample_rate, segment_idx=0, token_count=token_count, audio_duration=duration_str, real_time_factor=round(rtf, 2), prompt={ "tokens": token_count, "tokens-per-sec": ( round(token_count / processing_time, 2) if processing_time > 0 else 0 ), }, audio_samples={ "samples": samples, "samples-per-sec": ( round(samples / processing_time, 2) if processing_time > 0 else 0 ), }, processing_time_seconds=processing_time, peak_memory_usage=mx.get_peak_memory() / 1e9, ) # Clear cache after generation mx.clear_cache()