diff --git a/mlx_audio/tts/models/ace_step/ace_step.py b/mlx_audio/tts/models/ace_step/ace_step.py index bfa644f86..8377734fb 100644 --- a/mlx_audio/tts/models/ace_step/ace_step.py +++ b/mlx_audio/tts/models/ace_step/ace_step.py @@ -27,6 +27,7 @@ class Model(nn.Module): + custom_loading = True def __init__(self, config: ModelConfig): super().__init__() @@ -55,6 +56,12 @@ def __init__(self, config: ModelConfig): self._lm = None self._lm_config = None + # LoRA state + self._lora_loaded = False + self._lora_path = None + self._lora_scale = 1.0 + self._base_decoder_weights = None + @property def sample_rate(self) -> int: """Return audio sample rate.""" @@ -154,11 +161,24 @@ def post_load_hook(cls, model: "Model", model_path: Path) -> "Model": sampling_rate=vae_config.get("sampling_rate", 48000), ) - weights_path = model_path / "vae" / "diffusion_pytorch_model.safetensors" - if weights_path.exists(): - weights = mx.load(str(weights_path)) + mlx_weights_path = model_path / "vae" / "mlx_weights.safetensors" + pt_weights_path = model_path / "vae" / "diffusion_pytorch_model.safetensors" + + if mlx_weights_path.exists(): + from mlx.utils import tree_unflatten + + weights = mx.load(str(mlx_weights_path)) + model.vae.update(tree_unflatten(list(weights.items()))) + elif pt_weights_path.exists(): + weights = mx.load(str(pt_weights_path)) + sanitized = AutoencoderOobleck.sanitize(weights) model.vae.load_weights(weights, strict=False) + try: + mx.save_safetensors(str(mlx_weights_path), sanitized) + except Exception: + pass + model._compiled_vae_decode = mx.compile(model.vae.decode) model._sample_rate = vae_config.get("sampling_rate", 48000) # Load silence latent - try turbo subdirectory first, then root @@ -508,6 +528,169 @@ def _generate_lm_hints( return lm_hints + def load_lora( + self, lora_path: str, scale: float = 1.0, verbose: bool = True + ) -> str: + """Load a LoRA adapter into the decoder via weight fusion. + + Fuses LoRA weights directly into decoder parameters: + W_new = W_base + scale * (alpha/r) * B @ A + + Args: + lora_path: Path to the LoRA adapter directory containing + adapter_config.json and adapter_model.safetensors + scale: LoRA influence scale (0.0 to 1.0+, default 1.0) + verbose: Whether to print progress + + Returns: + Status message + """ + import os + import re + + from mlx.utils import tree_flatten, tree_unflatten + + lora_path = str(lora_path).strip() + if not os.path.exists(lora_path): + return f"LoRA path not found: {lora_path}" + + config_file = os.path.join(lora_path, "adapter_config.json") + if not os.path.exists(config_file): + return f"adapter_config.json not found in {lora_path}" + + # Load adapter config + with open(config_file) as f: + adapter_config = json.load(f) + + lora_alpha = adapter_config.get("lora_alpha", 1) + lora_r = adapter_config.get("r", 1) + use_rslora = adapter_config.get("use_rslora", False) + target_modules = adapter_config.get("target_modules", []) + + if use_rslora: + base_scaling = lora_alpha / math.sqrt(lora_r) + else: + base_scaling = lora_alpha / lora_r + + effective_scaling = base_scaling * scale + + if verbose: + print(f"Loading LoRA from {lora_path}") + print(f" rank={lora_r}, alpha={lora_alpha}, targets={target_modules}") + print(f" base_scaling={base_scaling:.4f}, lora_scale={scale}, effective={effective_scaling:.4f}") + + # Find and load adapter weights + weights_file = os.path.join(lora_path, "adapter_model.safetensors") + if not os.path.exists(weights_file): + weights_file = os.path.join(lora_path, "adapter_model.bin") + if not os.path.exists(weights_file): + return f"No adapter weights found in {lora_path}" + + lora_weights = mx.load(weights_file) + + # Back up base decoder weights before first LoRA application + if self._base_decoder_weights is None: + if verbose: + print("Backing up base decoder weights...") + self._base_decoder_weights = { + k: mx.array(v) for k, v in dict(tree_flatten(self.decoder.parameters())).items() + } + else: + # Restore base weights before applying new LoRA + if verbose: + print("Restoring base decoder weights before applying new LoRA...") + self.decoder.update(tree_unflatten(list(self._base_decoder_weights.items()))) + + # Get current decoder parameters + decoder_params = dict(tree_flatten(self.decoder.parameters())) + + # Fuse LoRA: W_new = W_base + effective_scaling * (B @ A) + applied = 0 + skipped = 0 + weight_updates = {} + + for key in sorted(lora_weights.keys()): + if not key.endswith(".lora_A.weight"): + continue + + b_key = key.replace(".lora_A.weight", ".lora_B.weight") + if b_key not in lora_weights: + skipped += 1 + continue + + # Map PEFT key to decoder parameter key + # PEFT: base_model.model.layers.0.self_attn.q_proj.lora_A.weight + # Decoder: layers.0.self_attn.q_proj.weight + base = key.replace(".lora_A.weight", "") + base = re.sub(r"^base_model\.model\.", "", base) + param_key = f"{base}.weight" + + if param_key not in decoder_params: + if verbose: + print(f" SKIP {param_key} (not found in decoder)") + skipped += 1 + continue + + A = lora_weights[key] # [r, in_features] + B = lora_weights[b_key] # [out_features, r] + W = decoder_params[param_key] + + # Verify shape compatibility + delta = (B @ A).astype(W.dtype) + if delta.shape != W.shape: + if verbose: + print(f" SKIP {param_key}: shape mismatch delta={delta.shape} vs W={W.shape}") + skipped += 1 + continue + + weight_updates[param_key] = W + effective_scaling * delta + applied += 1 + + if not weight_updates: + return f"No matching LoRA weights found (skipped={skipped})" + + # Apply fused weights directly to decoder (bypasses Model.sanitize) + self.decoder.update(tree_unflatten(list(weight_updates.items()))) + mx.eval(self.decoder.parameters()) + + self._lora_loaded = True + self._lora_path = lora_path + self._lora_scale = scale + + if verbose: + print(f"LoRA applied: {applied} layers fused, {skipped} skipped") + + return f"LoRA loaded: {applied} layers, scale={scale}" + + def unload_lora(self, verbose: bool = True) -> str: + """Unload LoRA adapter and restore base decoder weights. + + Returns: + Status message + """ + from mlx.utils import tree_unflatten + + if not self._lora_loaded: + return "No LoRA loaded" + + if self._base_decoder_weights is None: + return "No base weights backup found" + + if verbose: + print("Restoring base decoder weights...") + + self.decoder.update(tree_unflatten(list(self._base_decoder_weights.items()))) + mx.eval(self.decoder.parameters()) + + self._lora_loaded = False + self._lora_path = None + self._lora_scale = 1.0 + + if verbose: + print("Base decoder restored") + + return "LoRA unloaded, base decoder restored" + @staticmethod def sanitize(weights: Dict[str, mx.array]) -> Dict[str, mx.array]: """Convert PyTorch weights to MLX format. @@ -782,35 +965,14 @@ def generate_audio( fix_nfe: int = 8, infer_method: str = "ode", shift: float = 3.0, - guidance_scale: float = 15.0, - guidance_interval: float = 0.5, - omega_scale: float = 10.0, - cfg_type: str = "apg", lm_hints_25hz: Optional[mx.array] = None, + **kwargs, ) -> Dict: - """Generate audio latents. + """Generate audio latents via single-pass diffusion (turbo model). - Args: - text_hidden_states: Text embeddings - text_attention_mask: Text attention mask - lyric_hidden_states: Lyric embeddings - lyric_attention_mask: Lyric attention mask - refer_audio_acoustic_hidden_states_packed: Reference audio features - refer_audio_order_mask: Reference audio order mask - src_latents: Source latents - chunk_masks: Chunk masks - is_covers: Cover song flags - silence_latent: Silence latent for padding - attention_mask: Attention mask - seed: Random seed - fix_nfe: Number of function evaluations (steps) - infer_method: Inference method ('ode' or 'sde') - shift: Timestep schedule shift - guidance_scale: Classifier-free guidance scale (default 15.0) - guidance_interval: Fraction of steps where guidance is applied (0.5 = middle 50%) - omega_scale: Granularity scale for variance control (default 10.0) - cfg_type: CFG type ('cfg' for standard, 'apg' for Adaptive Projected Gradient) - lm_hints_25hz: Optional pre-computed LM hints at 25Hz rate + The turbo model was distilled without CFG, so no guidance is applied. + Any guidance_scale/cfg_type kwargs are accepted but ignored (matching + the original PyTorch turbo model behavior). Returns: Dictionary with 'target_latents' and 'time_costs' @@ -886,21 +1048,6 @@ def generate_audio( ) ) - # Prepare null (unconditional) embeddings for CFG - do_cfg = guidance_scale != 1.0 and guidance_scale != 0.0 - if do_cfg: - # Encode zeros through the encoder (same as PyTorch) - null_cond, _ = self.encoder( - text_hidden_states=mx.zeros_like(text_hidden_states), - text_attention_mask=text_attention_mask, - lyric_hidden_states=mx.zeros_like(lyric_hidden_states), - lyric_attention_mask=lyric_attention_mask, - refer_audio_acoustic_hidden_states_packed=mx.zeros_like( - refer_audio_acoustic_hidden_states_packed - ), - refer_audio_order_mask=refer_audio_order_mask, - ) - end_time = time.time() time_costs["encoder_time_cost"] = end_time - start_time start_time = end_time @@ -910,33 +1057,18 @@ def generate_audio( batch_size = context_latents.shape[0] dtype = context_latents.dtype - # Timestep schedule for turbo model num_steps = len(t_schedule_list) - - # Calculate guidance interval bounds - # guidance_interval=0.5 means guidance applied from 25% to 75% of steps - guidance_start = int(num_steps * (1 - guidance_interval) / 2) - guidance_end = int(num_steps * (1 + guidance_interval) / 2) - - # APG momentum buffer (running average for momentum-based guidance) - apg_momentum = 0.0 - apg_momentum_coef = -0.75 # Momentum coefficient from PyTorch reference - xt = noise - # Cross-attention caches: K,V computed on first step, reused for all subsequent steps - # Each cache auto-populates when first accessed, then returns cached values + # Cross-attention cache: K,V computed on first step, reused thereafter num_layers = len(self.decoder.layers) - cond_cache = make_cache(num_layers) - uncond_cache = make_cache(num_layers) if do_cfg else None + cache = make_cache(num_layers) for step_idx in range(num_steps): current_sigma = t_schedule_list[step_idx] t_curr = mx.full((batch_size,), current_sigma, dtype=dtype) - # Predict velocity with conditions - # Cache auto-populates on first call, reuses on subsequent calls - vt_cond = self.decoder( + vt = self.decoder( hidden_states=xt, timestep=t_curr, timestep_r=t_curr, @@ -944,90 +1076,17 @@ def generate_audio( encoder_hidden_states=encoder_hidden_states, encoder_attention_mask=encoder_attention_mask, context_latents=context_latents, - cache=cond_cache, + cache=cache, ) - # Check if we should apply guidance at this step - apply_guidance = do_cfg and guidance_start <= step_idx < guidance_end - - if apply_guidance: - # Predict velocity without conditions (null) - vt_uncond = self.decoder( - hidden_states=xt, - timestep=t_curr, - timestep_r=t_curr, - attention_mask=attention_mask, - encoder_hidden_states=null_cond, - encoder_attention_mask=encoder_attention_mask, - context_latents=context_latents, - cache=uncond_cache, - ) - - if cfg_type == "apg": - # Adaptive Projected Gradient (APG) guidance - # Projects diff onto orthogonal component of vt_cond - # PyTorch base model uses dims=[1] (time dimension) for [B, T, C] - diff = vt_cond - vt_uncond - - # Apply momentum - apg_momentum = apg_momentum_coef * apg_momentum + diff - diff = apg_momentum - - # Norm thresholding over time dimension (axis=1) - norm_threshold = 2.5 - diff_norm = mx.linalg.norm(diff, axis=1, keepdims=True) - scale_factor = mx.minimum( - mx.ones_like(diff_norm), norm_threshold / (diff_norm + 1e-8) - ) - diff = diff * scale_factor - - # Project over time dimension (axis=1) like PyTorch - # Use float32 for numerical stability - diff_f32 = diff.astype(mx.float32) - vt_cond_f32 = vt_cond.astype(mx.float32) - - # Normalize vt_cond over time dimension (axis=1) - vt_cond_normalized = vt_cond_f32 / ( - mx.linalg.norm(vt_cond_f32, axis=1, keepdims=True) + 1e-8 - ) - - # Parallel component: projection of diff onto vt_cond direction - diff_parallel = ( - mx.sum(diff_f32 * vt_cond_normalized, axis=1, keepdims=True) - * vt_cond_normalized - ) - - # Orthogonal component - diff_orthogonal = diff_f32 - diff_parallel - - # APG uses only the orthogonal component (eta=0 by default) - normalized_update = diff_orthogonal.astype(vt_cond.dtype) - - # Apply APG formula - vt = vt_cond + (guidance_scale - 1) * normalized_update - else: - # Standard CFG formula: v = v_uncond + scale * (v_cond - v_uncond) - vt = vt_uncond + guidance_scale * (vt_cond - vt_uncond) - else: - vt = vt_cond - - # Update x_t using Euler integration - # Turbo model uses simple Euler: xt = xt - vt * dt - next_t = ( - t_schedule_list[step_idx + 1] - if step_idx + 1 < len(t_schedule_list) - else 0.0 - ) - - # On final step, directly compute x0 from noise if step_idx == num_steps - 1: xt = self.get_x0_from_noise(xt, vt, t_curr) elif infer_method == "sde": - # Stochastic: predict clean, then renoise pred_clean = self.get_x0_from_noise(xt, vt, t_curr) + next_t = t_schedule_list[step_idx + 1] xt = self.renoise(pred_clean, next_t) - else: # ode - # Turbo model Euler: dx/dt = -v, so x_{t+1} = x_t - v_t * dt + else: + next_t = t_schedule_list[step_idx + 1] dt = current_sigma - next_t xt = xt - vt * dt @@ -1060,6 +1119,10 @@ def generate( cfg_type: str = "apg", vocal_language: str = "unknown", verbose: bool = True, + # Music metadata + bpm: Optional[int] = None, + keyscale: Optional[str] = None, + timesignature: Optional[str] = None, # Task parameters task_type: str = "text2music", source_audio: Optional[mx.array] = None, @@ -1184,7 +1247,9 @@ def generate( # Prepare text embeddings with task-specific instruction # Format: SFT_GEN_PROMPT with instruction, caption, and metas - text_hidden, text_mask = self._prepare_text_embeddings(text, duration=duration) + text_hidden, text_mask = self._prepare_text_embeddings( + text, duration=duration, bpm=bpm, keyscale=keyscale, timesignature=timesignature + ) lyric_hidden, lyric_mask = self._prepare_lyric_embeddings( lyrics, language=vocal_language ) @@ -1287,8 +1352,7 @@ def generate( print("Decoding audio...") decode_start = time.time() - latents_f32 = target_latents.astype(mx.float32) - audio = self.vae.decode(latents_f32) + audio = self._compiled_vae_decode(target_latents) mx.eval(audio) decode_time = time.time() - decode_start @@ -1297,7 +1361,7 @@ def generate( print(f"Audio shape: {audio.shape}") # Format output: [batch, channels, samples] -> [samples, channels] - audio = audio[0] + audio = audio[0].astype(mx.float32) audio = mx.clip(audio, -1.0, 1.0) audio = mx.transpose(audio) # [channels, samples] -> [samples, channels] num_samples = audio.shape[0] diff --git a/mlx_audio/tts/models/ace_step/lm.py b/mlx_audio/tts/models/ace_step/lm.py index f06dc4ae9..7266f5850 100644 --- a/mlx_audio/tts/models/ace_step/lm.py +++ b/mlx_audio/tts/models/ace_step/lm.py @@ -20,26 +20,25 @@ class LMConfig: """Configuration for the 5Hz Language Model. Available models: - - "0.6B": Smallest, fastest (default, recommended for most use cases) + - "0.6B": Full precision (default) + - "0.6B-8bit": 8-bit quantized, ~704 MB + - "0.6B-4bit": 4-bit quantized, ~373 MB - "4B": Largest, highest quality but slower - - Note: The 1.7B bundled model requires manual weight conversion and is not - currently supported. Use 0.6B or 4B instead. """ - model_size: str = "0.6B" # "0.6B" or "4B" + model_size: str = "0.6B" max_new_tokens: int = 3000 temperature: float = 0.8 top_k: int = 200 top_p: float = 0.95 repetition_penalty: float = 1.05 - # Base model path (set when loading from ACE-Step1.5) base_model_path: Optional[str] = None - # Model ID mapping for standalone models (HuggingFace repos) _STANDALONE_MODEL_IDS = { "0.6B": "ACE-Step/acestep-5Hz-lm-0.6B", + "0.6B-8bit": "WaveCut/acestep-5Hz-lm-0.6B-mlx_8bit", + "0.6B-4bit": "WaveCut/acestep-5Hz-lm-0.6B-mlx_4bit", "4B": "ACE-Step/acestep-5Hz-lm-4B", } diff --git a/mlx_audio/tts/models/ace_step/modules.py b/mlx_audio/tts/models/ace_step/modules.py index 231bb32dd..3f4b3a4bb 100644 --- a/mlx_audio/tts/models/ace_step/modules.py +++ b/mlx_audio/tts/models/ace_step/modules.py @@ -295,32 +295,12 @@ def __call__( kv_len = keys.shape[2] - # Repeat KV heads for grouped query attention - if self.num_kv_heads < self.num_heads: - n_rep = self.num_heads // self.num_kv_heads - keys = mx.repeat(keys, n_rep, axis=1) - values = mx.repeat(values, n_rep, axis=1) - - # ACE-Step turbo uses standard scaled dot-product attention (softmax) - # for both self-attention and cross-attention - # queries, keys, values: [B, H, S, D] scale = 1.0 / math.sqrt(self.head_dim) - # Compute attention scores - # scores: [B, H, seq_len, kv_len] - scores = (queries @ keys.transpose(0, 1, 3, 2)) * scale - - # Apply attention mask if provided - if attention_mask is not None: - scores = scores + attention_mask - - # Softmax - weights = mx.softmax(scores.astype(mx.float32), axis=-1).astype(queries.dtype) - - # Apply to values - output = weights @ values # [B, H, S, D] + output = mx.fast.scaled_dot_product_attention( + queries, keys, values, scale=scale, mask=attention_mask + ) - # Reshape output output = output.transpose(0, 2, 1, 3) # [B, S, H, D] output = output.reshape(batch_size, seq_len, -1) diff --git a/mlx_audio/tts/models/ace_step/vae.py b/mlx_audio/tts/models/ace_step/vae.py index 3a75149b8..247195b1f 100644 --- a/mlx_audio/tts/models/ace_step/vae.py +++ b/mlx_audio/tts/models/ace_step/vae.py @@ -1,89 +1,11 @@ -# Copyright (c) 2025, Prince Canuma and contributors (https://github.com/Blaizzy/mlx-audio) - import math -from typing import Dict, List +from typing import Dict, List, Optional import mlx.core as mx import mlx.nn as nn -class Snake1d(nn.Module): - - def __init__(self, channels: int, logscale: bool = True): - super().__init__() - self.alpha = mx.zeros((1, channels, 1)) - self.beta = mx.zeros((1, channels, 1)) - self.logscale = logscale - - def __call__(self, x: mx.array) -> mx.array: - # x: [batch, channels, time] - # Apply exp() to alpha/beta when logscale=True (matches diffusers) - alpha = mx.exp(self.alpha) if self.logscale else self.alpha - beta = mx.exp(self.beta) if self.logscale else self.beta - return x + (1.0 / (beta + 1e-9)) * mx.power(mx.sin(alpha * x), 2) - - -class WeightNormConv1d(nn.Module): - - def __init__( - self, - in_channels: int, - out_channels: int, - kernel_size: int, - stride: int = 1, - padding: int = 0, - dilation: int = 1, - bias: bool = True, - ): - super().__init__() - self.in_channels = in_channels - self.out_channels = out_channels - self.kernel_size = kernel_size - self.stride = stride - self.padding = padding - self.dilation = dilation - - # Weight normalization: W = g * V / ||V|| - self.weight_g = mx.ones((out_channels, 1, 1)) - self.weight_v = ( - mx.random.normal((out_channels, in_channels, kernel_size)) * 0.02 - ) - - if bias: - self.bias = mx.zeros((out_channels,)) - else: - self.bias = None - - def __call__(self, x: mx.array) -> mx.array: - # x: [batch, channels, time] - # Compute normalized weight - norm = mx.sqrt(mx.sum(self.weight_v**2, axis=(1, 2), keepdims=True) + 1e-12) - weight = self.weight_g * self.weight_v / norm - - # Transpose for MLX conv: [batch, time, channels] - x = x.transpose(0, 2, 1) - - # MLX conv1d expects weight as [out, kernel, in] - weight_mlx = weight.transpose(0, 2, 1) - - # Apply convolution - y = mx.conv1d( - x, - weight_mlx, - stride=self.stride, - padding=self.padding, - dilation=self.dilation, - ) - - # Add bias - if self.bias is not None: - y = y + self.bias - - # Transpose back: [batch, channels, time] - return y.transpose(0, 2, 1) - - -class WeightNormConvTranspose1d(nn.Module): +class FastConvTranspose1d(nn.Module): def __init__( self, @@ -100,60 +22,37 @@ def __init__( self.kernel_size = kernel_size self.stride = stride self.padding = padding - - self.weight_g = mx.ones((in_channels, 1, 1)) - self.weight_v = ( - mx.random.normal((in_channels, out_channels, kernel_size)) * 0.02 - ) - + # Weight stored as [in_ch, out_ch, K] for efficient matmul + self.weight = mx.zeros((in_channels, out_channels, kernel_size)) if bias: self.bias = mx.zeros((out_channels,)) - else: - self.bias = None def __call__(self, x: mx.array) -> mx.array: - # x: [batch, channels, time] - batch_size, in_ch, in_len = x.shape + # x: [B, L, in_ch] (NLC) + batch_size, in_len, _ = x.shape - # Compute normalized weight: [in, out, kernel] - norm = mx.sqrt(mx.sum(self.weight_v**2, axis=(1, 2), keepdims=True) + 1e-12) - weight = self.weight_g * self.weight_v / norm + # [in_ch, out_ch, K] -> [in_ch, out_ch * K] + w = self.weight.reshape(self.in_channels, self.out_channels * self.kernel_size) - # Transposed conv1d output length: (in_len - 1) * stride + kernel_size - # After removing 2*padding, final length is: (in_len - 1) * stride + kernel_size - 2*padding + # [B, L, in_ch] @ [in_ch, out_ch * K] -> [B, L, out_ch * K] + y = x @ w - # x: [batch, in_ch, in_len] -> [batch, in_len, in_ch] - x_t = x.transpose(0, 2, 1) - - # weight: [in_ch, out_ch, kernel] -> [in_ch, out_ch * kernel] - weight_reshaped = weight.reshape( - self.in_channels, self.out_channels * self.kernel_size - ) - - # Compute: [batch, in_len, out_ch * kernel] - y = x_t @ weight_reshaped - - # Reshape to [batch, in_len, out_ch, kernel] + # [B, L, out_ch, K] y = y.reshape(batch_size, in_len, self.out_channels, self.kernel_size) - # Transpose: [batch, out_ch, in_len, kernel] + # [B, out_ch, L, K] y = y.transpose(0, 2, 1, 3) - # Full output length before padding removal full_out_len = (in_len - 1) * self.stride + self.kernel_size - overlap = self.kernel_size - self.stride - # Split kernel contributions - first_part = y[:, :, :, : self.stride] # [batch, out_ch, in_len, stride] - second_part = y[:, :, :, self.stride :] # [batch, out_ch, in_len, overlap] - - # Flatten: [batch, out_ch, in_len * stride] + first_part = y[:, :, :, : self.stride] first_flat = first_part.reshape( batch_size, self.out_channels, in_len * self.stride ) if overlap > 0: + second_part = y[:, :, :, self.stride :] second_flat = second_part.reshape( batch_size, self.out_channels, in_len * overlap ) @@ -165,8 +64,6 @@ def __call__(self, x: mx.array) -> mx.array: ], axis=2, ) - - # Pad second_flat at the start (no overlap before first input) second_padded = mx.concatenate( [ mx.zeros( @@ -177,250 +74,186 @@ def __call__(self, x: mx.array) -> mx.array: axis=2, ) - # Trim to match full_out_len first_padded = first_padded[:, :, :full_out_len] second_padded = second_padded[:, :, :full_out_len] - output = first_padded + second_padded else: output = first_flat - # Trim padding if self.padding > 0: end_idx = full_out_len - self.padding output = output[:, :, self.padding : end_idx] - # Add bias - if self.bias is not None: + if "bias" in self: output = output + self.bias[None, :, None] - return output + # [B, out_ch, L_out] -> [B, L_out, out_ch] (NLC) + return output.transpose(0, 2, 1) + + +class Snake1d(nn.Module): + + def __init__(self, channels: int, logscale: bool = True): + super().__init__() + self.alpha = mx.zeros(channels) + self.beta = mx.zeros(channels) + self.logscale = logscale + + def __call__(self, x: mx.array) -> mx.array: + alpha = mx.exp(self.alpha) if self.logscale else self.alpha + beta = mx.exp(self.beta) if self.logscale else self.beta + return x + mx.reciprocal(beta + 1e-9) * mx.power(mx.sin(alpha * x), 2) class ResidualUnit(nn.Module): - def __init__(self, channels: int, kernel_size: int = 7, dilation: int = 1): + def __init__(self, dimension: int = 16, dilation: int = 1): super().__init__() - self.snake1 = Snake1d(channels) - self.conv1 = WeightNormConv1d( - channels, - channels, - kernel_size, - padding=(kernel_size - 1) * dilation // 2, - dilation=dilation, + pad = ((7 - 1) * dilation) // 2 + self.snake1 = Snake1d(dimension) + self.conv1 = nn.Conv1d( + dimension, dimension, kernel_size=7, dilation=dilation, padding=pad ) - self.snake2 = Snake1d(channels) - self.conv2 = WeightNormConv1d(channels, channels, 1) + self.snake2 = Snake1d(dimension) + self.conv2 = nn.Conv1d(dimension, dimension, kernel_size=1) - def __call__(self, x: mx.array) -> mx.array: - y = self.snake1(x) - y = self.conv1(y) - y = self.snake2(y) - y = self.conv2(y) - # Handle potential padding mismatch - padding = (x.shape[-1] - y.shape[-1]) // 2 + def __call__(self, hidden_state: mx.array) -> mx.array: + output = self.conv1(self.snake1(hidden_state)) + output = self.conv2(self.snake2(output)) + padding = (hidden_state.shape[1] - output.shape[1]) // 2 if padding > 0: - x = x[..., padding:-padding] - return x + y + hidden_state = hidden_state[:, padding:-padding, :] + return hidden_state + output class EncoderBlock(nn.Module): - """Encoder block for Oobleck VAE.""" - def __init__(self, in_channels: int, out_channels: int, stride: int): + def __init__(self, input_dim: int, output_dim: int, stride: int = 1): super().__init__() - self.res_unit1 = ResidualUnit(in_channels, dilation=1) - self.res_unit2 = ResidualUnit(in_channels, dilation=3) - self.res_unit3 = ResidualUnit(in_channels, dilation=9) - self.snake1 = Snake1d(in_channels) - self.conv1 = WeightNormConv1d( - in_channels, - out_channels, - kernel_size=stride * 2, + self.res_unit1 = ResidualUnit(input_dim, dilation=1) + self.res_unit2 = ResidualUnit(input_dim, dilation=3) + self.res_unit3 = ResidualUnit(input_dim, dilation=9) + self.snake1 = Snake1d(input_dim) + self.conv1 = nn.Conv1d( + input_dim, + output_dim, + kernel_size=2 * stride, stride=stride, padding=math.ceil(stride / 2), ) - def __call__(self, x: mx.array) -> mx.array: - x = self.res_unit1(x) - x = self.res_unit2(x) - x = self.res_unit3(x) - x = self.snake1(x) - x = self.conv1(x) - return x + def __call__(self, hidden_state: mx.array) -> mx.array: + hidden_state = self.res_unit1(hidden_state) + hidden_state = self.res_unit2(hidden_state) + hidden_state = self.snake1(self.res_unit3(hidden_state)) + hidden_state = self.conv1(hidden_state) + return hidden_state -class OobleckEncoder(nn.Module): - """Oobleck VAE Encoder for encoding audio to latents. - - The encoder outputs 2*latent_dim channels representing mean and scale - of a diagonal Gaussian distribution. - """ +class DecoderBlock(nn.Module): - def __init__( - self, - audio_channels: int = 2, - channels: int = 128, - latent_channels: int = 128, # Output channels (mean + scale) - channel_multiples: List[int] = None, - downsampling_ratios: List[int] = None, - ): + def __init__(self, input_dim: int, output_dim: int, stride: int = 1): super().__init__() - - if channel_multiples is None: - channel_multiples = [1, 2, 4, 8, 16] - if downsampling_ratios is None: - downsampling_ratios = [2, 4, 4, 6, 10] - - self.audio_channels = audio_channels - self.channels = channels - self.latent_channels = latent_channels - - # Input convolution - self.conv1 = WeightNormConv1d( - audio_channels, channels, kernel_size=7, padding=3 - ) - - # Encoder blocks - channel_multiples_with_1 = [1] + channel_multiples - self.block = [] - for i, stride in enumerate(downsampling_ratios): - in_ch = channels * channel_multiples_with_1[i] - out_ch = channels * channel_multiples_with_1[i + 1] - self.block.append(EncoderBlock(in_ch, out_ch, stride)) - - # Output layers - output encoder_hidden_size channels (mean + scale) - d_model = channels * channel_multiples[-1] - self.snake1 = Snake1d(d_model) - self.conv2 = WeightNormConv1d( - d_model, latent_channels, kernel_size=3, padding=1 + self.snake1 = Snake1d(input_dim) + self.conv_t1 = FastConvTranspose1d( + input_dim, + output_dim, + kernel_size=2 * stride, + stride=stride, + padding=stride // 2, ) + self.res_unit1 = ResidualUnit(output_dim, dilation=1) + self.res_unit2 = ResidualUnit(output_dim, dilation=3) + self.res_unit3 = ResidualUnit(output_dim, dilation=9) - def __call__(self, x: mx.array) -> mx.array: - """Encode audio to latent parameters (mean and scale). - - Args: - x: Audio tensor of shape [batch, channels, samples] or [batch, samples, channels] - - Returns: - Latent parameters of shape [batch, latent_channels, time] - where latent_channels = 2 * latent_dim (mean + scale) - """ - # Ensure channels-first format - if x.shape[-1] == self.audio_channels: - x = x.transpose(0, 2, 1) # [batch, samples, ch] -> [batch, ch, samples] + def __call__(self, hidden_state: mx.array) -> mx.array: + hidden_state = self.snake1(hidden_state) + hidden_state = self.conv_t1(hidden_state) + hidden_state = self.res_unit1(hidden_state) + hidden_state = self.res_unit2(hidden_state) + hidden_state = self.res_unit3(hidden_state) + return hidden_state - x = self.conv1(x) - for block in self.block: - x = block(x) - - x = self.snake1(x) - x = self.conv2(x) - - return x - - -class DecoderBlock(nn.Module): +class OobleckEncoder(nn.Module): def __init__( self, - in_channels: int, - out_channels: int, - stride: int, + encoder_hidden_size: int, + audio_channels: int, + downsampling_ratios: List[int], + channel_multiples: List[int], ): super().__init__() - # Snake activation before transposed conv - self.snake1 = Snake1d(in_channels) - - # Transposed conv for upsampling - self.conv_t1 = WeightNormConvTranspose1d( - in_channels, - out_channels, - kernel_size=stride * 2, - stride=stride, - padding=stride // 2, + cm = [1] + list(channel_multiples) + + self.conv1 = nn.Conv1d( + audio_channels, encoder_hidden_size, kernel_size=7, padding=3 ) - # Residual units with different dilations - self.res_unit1 = ResidualUnit(out_channels, dilation=1) - self.res_unit2 = ResidualUnit(out_channels, dilation=3) - self.res_unit3 = ResidualUnit(out_channels, dilation=9) + self.block = [] + for i, stride in enumerate(downsampling_ratios): + self.block.append( + EncoderBlock( + input_dim=encoder_hidden_size * cm[i], + output_dim=encoder_hidden_size * cm[i + 1], + stride=stride, + ) + ) - def __call__(self, x: mx.array) -> mx.array: - x = self.snake1(x) - x = self.conv_t1(x) - x = self.res_unit1(x) - x = self.res_unit2(x) - x = self.res_unit3(x) - return x + d_model = encoder_hidden_size * cm[-1] + self.snake1 = Snake1d(d_model) + self.conv2 = nn.Conv1d(d_model, encoder_hidden_size, kernel_size=3, padding=1) + + def __call__(self, hidden_state: mx.array) -> mx.array: + hidden_state = self.conv1(hidden_state) + for module in self.block: + hidden_state = module(hidden_state) + hidden_state = self.snake1(hidden_state) + hidden_state = self.conv2(hidden_state) + return hidden_state class OobleckDecoder(nn.Module): def __init__( self, - input_channels: int = 64, - channels: int = 128, - audio_channels: int = 2, - channel_multiples: List[int] = None, - downsampling_ratios: List[int] = None, + channels: int, + input_channels: int, + audio_channels: int, + upsampling_ratios: List[int], + channel_multiples: List[int], ): super().__init__() + strides = upsampling_ratios + cm = [1] + list(channel_multiples) - if channel_multiples is None: - channel_multiples = [1, 2, 4, 8, 16] - if downsampling_ratios is None: - downsampling_ratios = [2, 4, 4, 6, 10] - - self.input_channels = input_channels - self.channels = channels - self.audio_channels = audio_channels - - # Reverse for decoder (upsampling) - channel_multiples = list(reversed(channel_multiples)) - downsampling_ratios = list(reversed(downsampling_ratios)) - - # Input convolution (named conv1 to match PyTorch weights) - in_ch = input_channels - out_ch = channels * channel_multiples[0] - self.conv1 = WeightNormConv1d(in_ch, out_ch, kernel_size=7, padding=3) + self.conv1 = nn.Conv1d( + input_channels, channels * cm[-1], kernel_size=7, padding=3 + ) - # Decoder blocks self.block = [] - for i, (mult, stride) in enumerate(zip(channel_multiples, downsampling_ratios)): - in_ch = channels * mult - out_ch = ( - channels * channel_multiples[i + 1] - if i + 1 < len(channel_multiples) - else channels + for i, stride in enumerate(strides): + self.block.append( + DecoderBlock( + input_dim=channels * cm[len(strides) - i], + output_dim=channels * cm[len(strides) - i - 1], + stride=stride, + ) ) - self.block.append(DecoderBlock(in_ch, out_ch, stride)) - # Output layers (named snake1/conv2 to match PyTorch weights) self.snake1 = Snake1d(channels) - self.conv2 = WeightNormConv1d( + self.conv2 = nn.Conv1d( channels, audio_channels, kernel_size=7, padding=3, bias=False ) - def __call__(self, x: mx.array) -> mx.array: - - # Ensure channels-first format - if x.shape[-1] == self.input_channels: - x = x.transpose(0, 2, 1) # [batch, time, ch] -> [batch, ch, time] - - x = self.conv1(x) - - for block in self.block: - x = block(x) - - x = self.snake1(x) - x = self.conv2(x) - - # Note: PyTorch's Oobleck decoder does NOT apply tanh - # The audio values should already be in a reasonable range - - return x + def __call__(self, hidden_state: mx.array) -> mx.array: + hidden_state = self.conv1(hidden_state) + for layer in self.block: + hidden_state = layer(hidden_state) + hidden_state = self.snake1(hidden_state) + hidden_state = self.conv2(hidden_state) + return hidden_state class AutoencoderOobleck(nn.Module): @@ -428,19 +261,18 @@ class AutoencoderOobleck(nn.Module): def __init__( self, audio_channels: int = 2, - channel_multiples: List[int] = None, + channel_multiples: Optional[List[int]] = None, decoder_channels: int = 128, decoder_input_channels: int = 64, - downsampling_ratios: List[int] = None, + downsampling_ratios: Optional[List[int]] = None, encoder_hidden_size: int = 128, sampling_rate: int = 48000, ): super().__init__() - - if channel_multiples is None: - channel_multiples = [1, 2, 4, 8, 16] if downsampling_ratios is None: downsampling_ratios = [2, 4, 4, 6, 10] + if channel_multiples is None: + channel_multiples = [1, 2, 4, 8, 16] self.sampling_rate = sampling_rate self.hop_length = math.prod(downsampling_ratios) @@ -448,112 +280,102 @@ def __init__( self.latent_channels = decoder_input_channels self.encoder_hidden_size = encoder_hidden_size - # Encoder for audio-to-latent conversion - # Output encoder_hidden_size channels (mean + scale = 2 * latent_dim) self.encoder = OobleckEncoder( + encoder_hidden_size=encoder_hidden_size, audio_channels=audio_channels, - channels=decoder_channels, # Match decoder's internal channel width - latent_channels=encoder_hidden_size, # 128 = 64 mean + 64 scale - channel_multiples=channel_multiples, downsampling_ratios=downsampling_ratios, + channel_multiples=channel_multiples, ) - - # Decoder for latent-to-audio conversion self.decoder = OobleckDecoder( - input_channels=decoder_input_channels, channels=decoder_channels, + input_channels=decoder_input_channels, audio_channels=audio_channels, + upsampling_ratios=downsampling_ratios[::-1], channel_multiples=channel_multiples, - downsampling_ratios=downsampling_ratios, ) def encode(self, audio: mx.array, sample: bool = True) -> mx.array: - """Encode audio to latent representation. - - The encoder outputs a diagonal Gaussian distribution (mean + log_scale). - By default, this returns the mean (deterministic sampling). + # audio: [B, channels, samples] -> NLC [B, samples, channels] + if audio.shape[1] == self.audio_channels and audio.shape[-1] != self.audio_channels: + audio = audio.transpose(0, 2, 1) - Args: - audio: Audio tensor of shape [batch, channels, samples] - or [batch, samples, channels] - sample: If True, return mean (deterministic). If False, return - both mean and scale as a tuple. - - Returns: - Latent tensor of shape [batch, time, latent_dim] (if sample=True) - or tuple of (mean, scale) tensors (if sample=False) - """ - # Encode to latent parameters: [batch, encoder_hidden_size, time] - latent_params = self.encoder(audio) - - # Split into mean and log_scale (each latent_channels/2 = 64 dimensions) - # Shape: [batch, 128, time] -> [batch, 64, time] + [batch, 64, time] - mean, log_scale = mx.split(latent_params, 2, axis=1) + h = self.encoder(audio) + mean, _scale = mx.split(h, 2, axis=-1) if sample: - # Return mean (deterministic sampling, equivalent to mode) - # Transpose to [batch, time, latent_dim] for consistency - return mean.transpose(0, 2, 1) - else: - # Return distribution parameters - scale = mx.exp(log_scale) - return mean.transpose(0, 2, 1), scale.transpose(0, 2, 1) + return mean + scale = mx.exp(_scale) + return mean, scale def decode(self, latents: mx.array) -> mx.array: - """Decode latents to audio. + # latents: [B, T, latent_dim] NLC from ace_step.py + audio_nlc = self.decoder(latents) + # -> [B, channels, samples] for ace_step.py + return audio_nlc.transpose(0, 2, 1) - Args: - latents: Latent tensor of shape [batch, time, latent_dim] - or [batch, latent_dim, time] - - Returns: - Audio tensor of shape [batch, channels, samples] - """ - return self.decoder(latents) + @staticmethod + def _fuse_weight_norm(g: mx.array, v: mx.array) -> mx.array: + v_flat = v.reshape(v.shape[0], -1) + norm = mx.sqrt(mx.sum(v_flat * v_flat, axis=1, keepdims=True)).reshape(g.shape) + return g * v / (norm + 1e-9) @staticmethod def sanitize(weights: Dict[str, mx.array]) -> Dict[str, mx.array]: - """Convert PyTorch VAE weights to MLX format.""" sanitized = {} + processed = set() + all_keys = sorted(weights.keys()) + + for key in all_keys: + if key in processed: + continue + + if key.endswith(".weight_g"): + base = key[: -len(".weight_g")] + v_key = base + ".weight_v" + if v_key not in weights: + processed.add(key) + continue + + w = AutoencoderOobleck._fuse_weight_norm(weights[key], weights[v_key]) + + if "conv_t1" in base: + # FastConvTranspose1d stores [in, out, K] — keep as-is + pass + else: + # Conv1d: PT [out, in, K] -> MLX [out, K, in] + w = w.transpose(0, 2, 1) - for key, value in weights.items(): - new_key = key + sanitized[base + ".weight"] = w + processed.add(key) + processed.add(v_key) + continue - # Handle weight normalization - weight_v needs transposition - # PyTorch Conv1d: [out, in, kernel] - # MLX expects: [out, in, kernel] for weight_v (we transpose in forward) - if "weight_v" in key and len(value.shape) == 3: - # Keep as [out, in, kernel] - we'll handle in forward pass - pass + if key.endswith(".weight_v"): + continue - # Snake alpha/beta should be [1, channels, 1] - if "snake" in key and ("alpha" in key or "beta" in key): - if len(value.shape) == 1: - value = value[None, :, None] + if key.endswith(".alpha") or key.endswith(".beta"): + sanitized[key] = weights[key].squeeze() + processed.add(key) + continue - sanitized[new_key] = value + sanitized[key] = weights[key] + processed.add(key) return sanitized def load_weights(self, weights: Dict[str, mx.array], strict: bool = False): - """Load weights into the model.""" from mlx.utils import tree_unflatten - # Sanitize all weights - sanitized_weights = self.sanitize(weights) + sanitized = self.sanitize(weights) - # Load decoder weights decoder_weights = { - k: v for k, v in sanitized_weights.items() if k.startswith("decoder.") + k: v for k, v in sanitized.items() if k.startswith("decoder.") } if decoder_weights: - nested = tree_unflatten(list(decoder_weights.items())) - self.update(nested) + self.update(tree_unflatten(list(decoder_weights.items()))) - # Load encoder weights if available encoder_weights = { - k: v for k, v in sanitized_weights.items() if k.startswith("encoder.") + k: v for k, v in sanitized.items() if k.startswith("encoder.") } if encoder_weights: - nested = tree_unflatten(list(encoder_weights.items())) - self.update(nested) + self.update(tree_unflatten(list(encoder_weights.items()))) diff --git a/mlx_audio/tts/utils.py b/mlx_audio/tts/utils.py index 2933b2cbd..f8b486bf1 100644 --- a/mlx_audio/tts/utils.py +++ b/mlx_audio/tts/utils.py @@ -17,6 +17,7 @@ ) MODEL_REMAPPING = { + "acestep": "ace_step", "qwen3_tts": "qwen3_tts", "outetts": "outetts", "spark": "spark", diff --git a/mlx_audio/utils.py b/mlx_audio/utils.py index 57e2663a2..97e2e2fce 100644 --- a/mlx_audio/utils.py +++ b/mlx_audio/utils.py @@ -284,15 +284,15 @@ def get_model_class( if item.is_dir() and not item.name.startswith("__"): available_models.append(item.name) - if model_name is not None and model_type_mapped != model_type: + if model_type_mapped is not None: + model_type = model_type_mapped + elif model_name is not None: for part in model_name: if part in available_models: model_type = part if part in model_remapping: model_type = model_remapping[part] break - elif model_type_mapped is not None: - model_type = model_type_mapped try: module_path = f"mlx_audio.{category}.models.{model_type}" @@ -371,6 +371,10 @@ def base_load_model( model_remapping=model_remapping, ) + # Models with custom_loading delegate entirely to from_pretrained + if getattr(model_class.Model, "custom_loading", False): + return model_class.Model.from_pretrained(str(model_path)) + # Get model config from model class if it exists, otherwise use the config model_config = ( model_class.ModelConfig.from_dict(config)