import re import time from dataclasses import dataclass from pathlib import Path from typing import Dict, Optional, Tuple, Union import mlx.core as mx import mlx.nn as nn from mlx_lm.generate import stream_generate from mlx_lm.models.qwen2 import Model as Qwen2Model from mlx_lm.sample_utils import make_logits_processors, make_sampler from tqdm import tqdm from mlx_audio.tts.models.base import BaseModelArgs, GenerationResult from .audio_tokenizer import BiCodecTokenizer from .utils.token_parser import GENDER_MAP, LEVELS_MAP, TASK_TOKEN_MAP PITCH_MAP = SPEED_MAP = { 0.0: "very_low", 0.5: "low", 1.0: "moderate", 1.5: "high", 2.0: "very_high", } @dataclass class ModelConfig(BaseModelArgs): sample_rate: int = 16000 bos_token_id: int = 151643 eos_token_id: int = 151645 hidden_act: str = "silu" hidden_size: int = 896 initializer_range: float = 0.02 intermediate_size: int = 4864 max_position_embeddings: int = 32768 max_window_layers: int = 21 model_type: str = "qwen2" num_attention_heads: int = 14 num_hidden_layers: int = 24 num_key_value_heads: int = 2 rms_norm_eps: float = 1e-06 rope_theta: float = 1000000.0 sliding_window: int = 32768 tie_word_embeddings: bool = True torch_dtype: str = "bfloat16" transformers_version: str = "4.43.1" use_sliding_window: bool = False vocab_size: int = 166000 rope_traditional: bool = False rope_scaling: Optional[Dict[str, Union[float, str]]] = None class Model(nn.Module): """ Spark-TTS for text-to-speech generation. """ def __init__(self, config: ModelConfig): """ Initializes the SparkTTS model with the provided configurations and device. Args: config (ModelConfig): The configuration for the model. """ self.config = config self.model = Qwen2Model(config) self.tokenizer = None self._audio_tokenizer = None @classmethod def post_load_hook(cls, model: "Model", model_path: Path): """ Hook called after model weights are loaded. Used to initialize the tokenizer which is required for text input. """ from transformers import AutoTokenizer print( f"Loading tokenizer from {model_path} with eos_token_ids={model.config.eos_token_id}" ) model.tokenizer = AutoTokenizer.from_pretrained( str(model_path), eos_token_ids=model.config.eos_token_id ) model._audio_tokenizer = BiCodecTokenizer(model_path) return model def load_weights(self, weights, strict=True): self.model.load_weights(weights, strict=strict) def parameters(self): return self.model.parameters() def model_type(self): return "spark" def sanitize(self, weights): return self.model.sanitize(weights) @property def sample_rate(self): return self.config.sample_rate @property def layers(self): return self.model.layers def model_quant_predicate(self, p, m): """ Model modules to skip during quantization """ return not p.startswith("_audio_tokenizer") def process_prompt( self, text: str, ref_audio: Path, ref_text: str, ) -> Tuple[str, mx.array]: """ Process input for voice cloning. Args: text (str): The text input to be converted to speech. ref_audio (Path): Path to the audio file used as a reference. ref_text (str, optional): Transcript of the reference audio. Return: Tuple[str, mx.array]: Input prompt; global tokens """ global_token_ids, semantic_token_ids = self._audio_tokenizer.tokenize(ref_audio) global_tokens = "".join( [f"<|bicodec_global_{i}|>" for i in global_token_ids.squeeze()] ) # Prepare the input tokens for the model if ref_text is not None: semantic_tokens = "".join( [f"<|bicodec_semantic_{i}|>" for i in semantic_token_ids.squeeze()] ) inputs = [ TASK_TOKEN_MAP["tts"], "<|start_content|>", ref_text, text, "<|end_content|>", "<|start_global_token|>", global_tokens, "<|end_global_token|>", "<|start_semantic_token|>", semantic_tokens, ] else: inputs = [ TASK_TOKEN_MAP["tts"], "<|start_content|>", text, "<|end_content|>", "<|start_global_token|>", global_tokens, "<|end_global_token|>", ] inputs = "".join(inputs) return inputs, global_token_ids def process_prompt_control( self, gender: str, pitch: str, speed: str, text: str, ): """ Process input for voice creation. Args: gender (str): female | male. pitch (str): very_low | low | moderate | high | very_high speed (str): very_low | low | moderate | high | very_high text (str): The text input to be converted to speech. Return: str: Input prompt """ assert gender in GENDER_MAP.keys() assert pitch in LEVELS_MAP.keys() assert speed in LEVELS_MAP.keys() gender_id = GENDER_MAP[gender] pitch_level_id = LEVELS_MAP[pitch] speed_level_id = LEVELS_MAP[speed] pitch_label_tokens = f"<|pitch_label_{pitch_level_id}|>" speed_label_tokens = f"<|speed_label_{speed_level_id}|>" gender_tokens = f"<|gender_{gender_id}|>" attribte_tokens = "".join( [gender_tokens, pitch_label_tokens, speed_label_tokens] ) control_tts_inputs = [ TASK_TOKEN_MAP["controllable_tts"], "<|start_content|>", text, "<|end_content|>", "<|start_style_label|>", attribte_tokens, "<|end_style_label|>", ] return "".join(control_tts_inputs) def generate( self, text: str, ref_audio: Path = None, ref_text: str = None, gender: str = "male", pitch: float = 1.0, speed: float = 1.0, temperature: float = 0.8, top_k: float = 50, top_p: float = 0.95, max_tokens: int = 3000, verbose: bool = False, split_pattern: str = "\n", **kwargs, ) -> GenerationResult: """ Performs inference to generate speech from text, incorporating prompt audio and/or text. Args: text (str): The text input to be converted to speech. ref_audio (Path): Path to the audio file used as a reference. ref_text (str, optional): Transcript of the reference audio. gender (str): female | male. pitch (str): very_low | low | moderate | high | very_high speed (str): very_low | low | moderate | high | very_high temperature (float, optional): Sampling temperature for controlling randomness. Default is 0.8. top_k (float, optional): Top-k sampling parameter. Default is 50. top_p (float, optional): Top-p (nucleus) sampling parameter. Default is 0.95. Returns: GenerationResult: Generated waveform as a tensor. """ speed_factor = SPEED_MAP[speed] pitch_factor = PITCH_MAP[pitch] if ref_audio is not None: # voice cloning gender = None text_splits = text.split(split_pattern) for text_split in text_splits: if gender is not None: prompt = self.process_prompt_control( gender, pitch_factor, speed_factor, text_split ) else: prompt, global_token_ids = self.process_prompt( text_split, ref_audio, ref_text ) inputs = self.tokenizer.encode( prompt, add_special_tokens=False, return_tensors="mlx" ) input_ids = mx.array(inputs) sampler = make_sampler(temperature, top_p=top_p, top_k=top_k) logits_processors = make_logits_processors( kwargs.get("logit_bias", None), kwargs.get("repetition_penalty", 1.3), kwargs.get("repetition_context_size", 20), ) time_start = time.time() generated_ids = [] # Generate speech using the model for i, response in enumerate( tqdm( stream_generate( self.model, tokenizer=self.tokenizer, prompt=input_ids.squeeze(0), max_tokens=max_tokens, sampler=sampler, logits_processors=logits_processors, ), total=max_tokens, disable=not verbose, ) ): next_token = mx.array([response.token]) input_ids = mx.concatenate([input_ids, next_token[None, :]], axis=1) if i % 50 == 0: mx.clear_cache() if next_token == 128258: break time_end = time.time() # Trim the output tokens to remove the input tokens generated_ids = mx.array( [output[len(input) :] for input, output in zip(inputs, input_ids)] ).tolist() # Decode the generated tokens into text predicts = self.tokenizer.batch_decode( generated_ids, skip_special_tokens=False )[0] # Extract semantic token IDs from the generated text pred_semantic_ids = mx.array( [ int(token) for token in re.findall(r"bicodec_semantic_(\d+)", predicts) ] )[None, ...] if gender is not None: global_token_ids = mx.array( [ int(token) for token in re.findall(r"bicodec_global_(\d+)", predicts) ] )[None, ...] # Convert semantic tokens back to waveform audio = self._audio_tokenizer.detokenize( global_token_ids.astype(mx.int32), pred_semantic_ids.astype(mx.int32), ) # Clear cache mx.clear_cache() audio_samples = len(audio) audio_duration_seconds = audio_samples / self.config.sample_rate # 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, sample_rate=self.config.sample_rate, samples=audio_samples, segment_idx=0, # Default segment index token_count=len(pred_semantic_ids.squeeze()), audio_samples={ "samples": audio_samples, "samples-per-sec": ( round(audio_samples / audio_duration_seconds, 2) if audio_duration_seconds > 0 else 0 ), }, audio_duration=duration_str, real_time_factor=( audio_duration_seconds / (time_end - time_start) if (time_end - time_start) > 0 else 0 ), prompt={ "tokens": len(pred_semantic_ids.squeeze()), "tokens-per-sec": ( round( len(pred_semantic_ids.squeeze()) / audio_duration_seconds, 2 ) if audio_duration_seconds > 0 else 0 ), }, processing_time_seconds=time_end - time_start, peak_memory_usage=mx.get_peak_memory() / 1e9, ) # Clear cache after each segment to avoid memory leaks mx.clear_cache()