import time from pathlib import Path from typing import Optional import mlx.core as mx import mlx.nn as nn from ..base import GenerationResult from .audio_vae import AudioVAE from .config import LMConfig, ModelArgs from .dit import UnifiedCFM, VoxCPMLocDiT from .encoder import VoxCPMLocEnc from .minicpm import MiniCPMModel class ScalarQuantizationLayer(nn.Module): def __init__(self, in_dim: int, out_dim: int, latent_dim: int = 64, scale: int = 9): super().__init__() self.scale = scale self.in_proj = nn.Linear(in_dim, latent_dim) self.out_proj = nn.Linear(latent_dim, out_dim) def __call__(self, x): x = self.in_proj(x) x = mx.tanh(x) # Rounding x = mx.round(x * self.scale) / self.scale return self.out_proj(x) class Model(nn.Module): def __init__(self, args: ModelArgs): super().__init__() self.args = args self.patch_size = args.patch_size self.feat_dim = args.feat_dim # LM Backbone self.base_lm = MiniCPMModel(args.lm_config) # Residual LM (vocab_size=0) res_config = LMConfig(**vars(args.lm_config)) res_config.num_hidden_layers = args.residual_lm_num_layers res_config.vocab_size = 0 self.residual_lm = MiniCPMModel(res_config) # Encoder enc_config = LMConfig(**vars(args.lm_config)) enc_config.hidden_size = args.encoder_config.hidden_dim enc_config.intermediate_size = args.encoder_config.ffn_dim enc_config.num_attention_heads = args.encoder_config.num_heads enc_config.num_hidden_layers = args.encoder_config.num_layers enc_config.vocab_size = 0 self.feat_encoder = VoxCPMLocEnc(enc_config, input_dim=args.feat_dim) # DiT / CFM dit_config = LMConfig(**vars(args.lm_config)) # base on LM config but override dit_config.hidden_size = args.dit_config.hidden_dim dit_config.intermediate_size = args.dit_config.ffn_dim dit_config.num_attention_heads = args.dit_config.num_heads dit_config.num_hidden_layers = args.dit_config.num_layers dit_config.vocab_size = 0 estimator = VoxCPMLocDiT(dit_config, in_channels=args.feat_dim) self.feat_decoder = UnifiedCFM( in_channels=args.feat_dim, cfm_params=args.dit_config.cfm_config, estimator=estimator, ) # Projections self.fsq_layer = ScalarQuantizationLayer( args.lm_config.hidden_size, args.lm_config.hidden_size, args.scalar_quantization_latent_dim, args.scalar_quantization_scale, ) self.enc_to_lm_proj = nn.Linear( args.encoder_config.hidden_dim, args.lm_config.hidden_size ) self.lm_to_dit_proj = nn.Linear( args.lm_config.hidden_size, args.dit_config.hidden_dim ) self.res_to_dit_proj = nn.Linear( args.lm_config.hidden_size, args.dit_config.hidden_dim ) # Stop Predictor self.stop_proj = nn.Linear( args.lm_config.hidden_size, args.lm_config.hidden_size ) self.stop_head = nn.Linear(args.lm_config.hidden_size, 2, bias=False) # Audio VAE self.audio_vae = AudioVAE(args.audio_vae_config) # Placeholder for tokenizer self.tokenizer = None @property def sample_rate(self): return self.args.audio_vae_config.sample_rate def sanitize(self, weights: dict): from mlx.utils import tree_flatten # Track whether VAE weights were already sanitized to prevent double processing vae_already_sanitized = False # 0. Check if audio_vae weights are present. If not, try to load from pth has_vae = any(k.startswith("audio_vae.") for k in weights.keys()) if not has_vae and self.args.model_path: p = Path(self.args.model_path) / "audiovae.pth" if p.exists(): try: import torch state = torch.load(p, map_location="cpu") if "state_dict" in state: state = state["state_dict"] # Convert to numpy/mlx and prefix vae_weights_pth = {} for k, v in state.items(): if k.startswith("module."): k = k[7:] # v is tensor arr = mx.array(v.numpy()) vae_weights_pth[k] = arr # Sanitize these VAE weights sanitized_vae = self.audio_vae.sanitize(vae_weights_pth) # Add to main weights for k, v in sanitized_vae.items(): weights[f"audio_vae.{k}"] = v # Mark that VAE weights have been sanitized vae_already_sanitized = True except ImportError: print(f"Warning: torch not installed, skipping loading {p}") except Exception as e: print(f"Error loading {p}: {e}") # Delegate AudioVAE sanitize (if keys exist now and not already sanitized) # Extract audio_vae weights vae_weights = {k: v for k, v in weights.items() if k.startswith("audio_vae.")} # Strip prefix vae_weights_stripped = { k[len("audio_vae.") :]: v for k, v in vae_weights.items() } if vae_weights_stripped and not vae_already_sanitized: # Sanitize VAE sanitized_vae = self.audio_vae.sanitize(vae_weights_stripped) # Put back for k in list(vae_weights.keys()): del weights[k] for k, v in sanitized_vae.items(): weights[f"audio_vae.{k}"] = v new_weights = {} curr_shapes = {k: v.shape for k, v in tree_flatten(self.parameters())} for k, v in weights.items(): if k not in curr_shapes: # might be skipped params or structure mismatch # keep it for now new_weights[k] = v continue target_shape = curr_shapes[k] if v.shape == target_shape: new_weights[k] = v else: # Try transpose if len(v.shape) == 2 and v.transpose().shape == target_shape: new_weights[k] = v.transpose() else: # Shape mismatch that transpose can't fix # Check for 1D weights (bias, RMSNorm) if ( len(v.shape) == 1 and len(target_shape) == 1 and v.shape != target_shape ): # e.g. embedding size mismatch print(f"Shape mismatch for {k}: {v.shape} vs {target_shape}") new_weights[k] = v # 3. Add computed buffers (RoPE) if missing # Since we are strict loading, we need to provide all params. # RoPE params are computed in __init__, so we can grab them from self. model_params = dict(tree_flatten(self.parameters())) for k, v in model_params.items(): if k not in new_weights and ("rope" in k): # It's a missing RoPE parameter that we have in memory. # Add it. new_weights[k] = v return new_weights @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 if model.tokenizer is None: model.tokenizer = AutoTokenizer.from_pretrained(str(model_path)) return model def _encode_prompt_audio(self, audio: mx.array) -> mx.array: """ Encode prompt audio into latent features. Args: audio: Audio waveform as mx.array, shape (T,) - already loaded and resampled Returns: audio_feat: Encoded audio features of shape (audio_length, patch_size, latent_dim) """ # Ensure proper length for patch alignment patch_len = self.patch_size * self.audio_vae.hop_length if audio.shape[0] % patch_len != 0: # Left padding to keep valid audio data at the end padding_size = patch_len - audio.shape[0] % patch_len audio = mx.pad(audio, [(padding_size, 0)]) audio_input = audio[None, None, :] # (1, 1, T) audio_feat = self.audio_vae.encode( audio_input, self.audio_vae.sample_rate ) # (1, T', D) in MLX format # audio_feat is (1, T', D) - batch, time, latent_dim audio_feat = audio_feat.squeeze(0) # (T', D) # Reshape into patches: (T', D) -> (audio_length, patch_size, D) T_prime = audio_feat.shape[0] audio_length = T_prime // self.patch_size audio_feat = audio_feat[ : audio_length * self.patch_size, : ] # Trim to exact multiple audio_feat = audio_feat.reshape( audio_length, self.patch_size, -1 ) # (audio_length, patch_size, D) return audio_feat def generate( self, text: str, max_tokens: int = 4096, ref_text: Optional[str] = None, ref_audio: Optional[str] = None, inference_timesteps: int = 10, cfg_value: float = 2.0, **kwargs, ): # Tokenize if self.tokenizer is None: raise ValueError("Tokenizer not loaded") start_time = time.perf_counter() # scale_emb scale_emb = ( self.args.lm_config.scale_emb if not self.args.lm_config.use_mup else 1.0 ) audio_start_token = 101 text_mask = None audio_mask = None feat_embed = None if ref_audio is not None and ref_text is not None: # Voice cloning mode: ref_text + target_text combined_text = ref_text + text input_ids = self.tokenizer.encode(combined_text) input_ids = mx.array(input_ids) input_ids = mx.concatenate([input_ids, mx.array([audio_start_token])]) text_length = len(input_ids) audio_feat = self._encode_prompt_audio( ref_audio ) # (audio_length, patch_size, D) audio_length = audio_feat.shape[0] text_pad_token = mx.zeros(audio_length, dtype=mx.int32) text_token = mx.concatenate([input_ids, text_pad_token]) audio_pad_feat = mx.zeros( (text_length, self.patch_size, self.feat_dim), dtype=mx.float32, ) audio_feat = mx.concatenate( [audio_pad_feat, audio_feat], axis=0 ) # (text_length + audio_length, patch_size, D) text_mask = mx.concatenate( [ mx.ones(text_length, dtype=mx.float32), mx.zeros(audio_length, dtype=mx.float32), ] ) audio_mask = mx.concatenate( [ mx.zeros(text_length, dtype=mx.float32), mx.ones(audio_length, dtype=mx.float32), ] ) text_token = text_token[None, :] # (1, T) audio_feat = audio_feat[None, :, :, :] # (1, T, P, D) text_mask = text_mask[None, :] # (1, T) audio_mask = audio_mask[None, :] # (1, T) feat_embed = self.feat_encoder(audio_feat) # (1, T, H) feat_embed = self.enc_to_lm_proj(feat_embed) # (1, T, H) text_embed = self.base_lm.embed_tokens(text_token) * scale_emb # (1, T, H) combined_embed = ( text_mask[:, :, None] * text_embed + audio_mask[:, :, None] * feat_embed ) # (1, T, H) prefix_feat_cond = audio_feat[:, -1, :, :] # (1, P, D) token_count = len(input_ids) else: # No voice cloning input_ids = self.tokenizer.encode(text) input_ids = mx.array(input_ids) token_count = len(input_ids) start_token = mx.array([audio_start_token]) input_ids = mx.concatenate([input_ids, start_token]) combined_embed = ( self.base_lm.embed_tokens(input_ids[None, :]) * scale_emb ) # (1, L, D) prefix_feat_cond = mx.zeros((1, self.patch_size, self.feat_dim)) enc_outputs, lm_cache = self.base_lm(combined_embed) if text_mask is not None and audio_mask is not None: enc_outputs = ( self.fsq_layer(enc_outputs) * audio_mask[:, :, None] + enc_outputs * text_mask[:, :, None] ) lm_hidden = enc_outputs[:, -1, :] if text_mask is None: lm_hidden = self.fsq_layer(lm_hidden) if text_mask is not None and audio_mask is not None: residual_input = enc_outputs + audio_mask[:, :, None] * feat_embed else: residual_input = enc_outputs residual_outputs, res_cache = self.residual_lm(residual_input) residual_hidden = residual_outputs[:, -1, :] # Generation Loop pred_feat_seq = [] for i in range(max_tokens): # DiT dit_h1 = self.lm_to_dit_proj(lm_hidden) dit_h2 = self.res_to_dit_proj(residual_hidden) dit_h = dit_h1 + dit_h2 # (1, H) cond_in = prefix_feat_cond.transpose(0, 2, 1) # (B, D, P) pred_feat = self.feat_decoder.sample( mu=dit_h, n_timesteps=inference_timesteps, patch_size=self.patch_size, cond=cond_in, cfg_value=cfg_value, ) pred_feat = pred_feat.transpose(0, 2, 1) pred_feat_seq.append(pred_feat) curr_embed = self.feat_encoder(pred_feat[:, None, :, :]) # (B, 1, H) curr_embed = self.enc_to_lm_proj(curr_embed) stop_logits = self.stop_head(nn.silu(self.stop_proj(lm_hidden))) stop_flag = mx.argmax(stop_logits, axis=-1).item() if i > 5 and stop_flag == 1: break curr_embed_step = curr_embed # (B, 1, H) new_lm_out, lm_cache = self.base_lm( inputs_embeds=curr_embed_step, cache=lm_cache ) lm_hidden_next = new_lm_out[:, -1, :] lm_hidden_next = self.fsq_layer(lm_hidden_next) res_in = lm_hidden_next[:, None, :] + curr_embed_step new_res_out, res_cache = self.residual_lm( inputs_embeds=res_in, cache=res_cache ) residual_hidden = new_res_out[:, -1, :] lm_hidden = lm_hidden_next prefix_feat_cond = pred_feat all_feats = mx.concatenate(pred_feat_seq, axis=1) # (B, Total_P, D) audio = self.audio_vae.decode(all_feats) audio = audio.flatten() end_time = time.perf_counter() elapsed_time = end_time - start_time samples = audio.shape[0] sample_rate = self.args.audio_vae_config.sample_rate # Use config value (44100) audio_duration_seconds = samples / sample_rate rtf = audio_duration_seconds / elapsed_time if elapsed_time > 0 else 0 duration_mins = int(audio_duration_seconds // 60) duration_secs = int(audio_duration_seconds % 60) duration_ms = int((audio_duration_seconds % 1) * 1000) duration_str = f"{int(audio_duration_seconds // 3600):02d}:{duration_mins:02d}:{duration_secs:02d}.{duration_ms:03d}" # Yield single result yield GenerationResult( audio=audio, samples=samples, sample_rate=sample_rate, segment_idx=0, token_count=token_count, audio_duration=duration_str, real_time_factor=rtf, prompt={ "tokens": token_count, "tokens-per-sec": ( round(token_count / elapsed_time, 2) if elapsed_time > 0 else 0 ), }, audio_samples={ "samples": samples, "samples-per-sec": ( round(samples / elapsed_time, 2) if elapsed_time > 0 else 0 ), }, processing_time_seconds=elapsed_time, peak_memory_usage=mx.get_peak_memory() / 1e9, )