# MLX diffusion generation loop for AceStep DiT decoder. # # Replicates the timestep scheduling and ODE/SDE stepping from # ``AceStepConditionGenerationModel.generate_audio`` using pure MLX arrays. # # Enhanced sampling modes (see issue #957): # - ``euler``: First-order Euler ODE/SDE step (default, original behaviour). # - ``heun``: Second-order Heun predictor-corrector — evaluates the model # twice per step and averages the predictions for higher accuracy, which # matters especially with 8-step turbo inference. # # Optional stabilisation techniques (work with *any* sampler mode): # - ``velocity_norm_threshold``: Clamp the L2 norm of velocity predictions # relative to the input norm. Prevents outlier predictions that cause # audio artefacts. # - ``velocity_ema_factor``: Exponential moving average blending between # the current and previous velocity prediction, smoothing the denoising # trajectory. import logging import math import time from typing import Dict, List, Optional, Tuple, Union import numpy as np from tqdm import tqdm logger = logging.getLogger(__name__) VALID_SAMPLER_MODES = {"euler", "heun"} # Pre-defined timestep schedules (from modeling_acestep_v15_turbo.py) VALID_SHIFTS = [1.0, 2.0, 3.0] VALID_TIMESTEPS = [ 1.0, 0.9545454545454546, 0.9333333333333333, 0.9, 0.875, 0.8571428571428571, 0.8333333333333334, 0.7692307692307693, 0.75, 0.6666666666666666, 0.6428571428571429, 0.625, 0.5454545454545454, 0.5, 0.4, 0.375, 0.3, 0.25, 0.2222222222222222, 0.125, ] SHIFT_TIMESTEPS = { 1.0: [1.0, 0.875, 0.75, 0.625, 0.5, 0.375, 0.25, 0.125], 2.0: [1.0, 0.9333333333333333, 0.8571428571428571, 0.7692307692307693, 0.6666666666666666, 0.5454545454545454, 0.4, 0.2222222222222222], 3.0: [1.0, 0.9545454545454546, 0.9, 0.8333333333333334, 0.75, 0.6428571428571429, 0.5, 0.3], } def get_timestep_schedule( shift: float = 3.0, timesteps: Optional[list] = None, infer_steps: Optional[int] = None, ) -> List[float]: """Compute the timestep schedule for diffusion sampling. When ``infer_steps`` is provided and ``timesteps`` is None, a continuous linspace schedule is generated (matching the PyTorch base-model behaviour). The legacy lookup-table path (8-step ``SHIFT_TIMESTEPS``) is used only when neither ``timesteps`` nor ``infer_steps`` is supplied. Args: shift: Diffusion timestep shift (applied via ``shift*t / (1+(shift-1)*t)``). timesteps: Optional custom list of timesteps. infer_steps: Number of diffusion steps. When given, overrides the fixed 8-step lookup table. Returns: List of timestep values (descending, without trailing 0). """ t_schedule_list = None if timesteps is not None: ts_list = list(timesteps) while ts_list and ts_list[-1] == 0: ts_list.pop() if len(ts_list) < 1: logger.warning("timesteps empty after removing zeros; using default shift=%s", shift) else: if len(ts_list) > 20: logger.warning("timesteps length=%d > 20; truncating", len(ts_list)) ts_list = ts_list[:20] mapped = [min(VALID_TIMESTEPS, key=lambda x, t=t: abs(x - t)) for t in ts_list] t_schedule_list = mapped if t_schedule_list is None and infer_steps is not None and infer_steps > 0: raw = [1.0 - i / infer_steps for i in range(infer_steps)] if shift != 1.0: raw = [shift * t / (1.0 + (shift - 1.0) * t) for t in raw] t_schedule_list = raw if t_schedule_list is None: original_shift = shift shift = min(VALID_SHIFTS, key=lambda x: abs(x - shift)) if original_shift != shift: logger.warning("shift=%.2f rounded to nearest valid shift=%.1f", original_shift, shift) t_schedule_list = SHIFT_TIMESTEPS[shift] return t_schedule_list def _mlx_apg_forward( pred_cond, pred_uncond, guidance_scale: float, momentum_state: Optional[Dict] = None, norm_threshold: float = 2.5, ): """APG (Adaptive Projected Guidance) in pure MLX — mirrors the PyTorch ``apg_forward``. Projection is performed along axis 1 (the time/sequence dimension) to match the PyTorch implementation which calls ``apg_forward(..., dims=[1])``. """ import mlx.core as mx proj_axis = 1 diff = pred_cond - pred_uncond if momentum_state is not None: diff = diff + momentum_state.get("running", 0) momentum_state["running"] = diff if norm_threshold > 0: diff_norm = mx.sqrt((diff * diff).sum(axis=proj_axis, keepdims=True)) scale_factor = mx.minimum(mx.ones_like(diff_norm), norm_threshold / (diff_norm + 1e-8)) diff = diff * scale_factor v1 = pred_cond / (mx.sqrt((pred_cond * pred_cond).sum(axis=proj_axis, keepdims=True)) + 1e-8) parallel = (diff * v1).sum(axis=proj_axis, keepdims=True) * v1 orthogonal = diff - parallel return pred_cond + (guidance_scale - 1) * orthogonal def _mlx_repaint_step_injection(xt, clean_src, mask, t_next, noise): """Replace non-repaint regions of *xt* with noised source latents (MLX).""" import mlx.core as mx zt = t_next * noise + (1.0 - t_next) * clean_src m = mx.expand_dims(mask, axis=-1) return mx.where(m, xt, zt) def _mlx_repaint_boundary_blend(x_gen, clean_src, mask_np, cf_frames): """Blend generated latents with source at repaint boundaries (MLX).""" import mlx.core as mx soft = mask_np.astype(np.float32).copy() if cf_frames <= 0: m = mx.expand_dims(mx.array(soft), axis=-1) return m * x_gen + (1.0 - m) * clean_src B, T = mask_np.shape for b in range(B): row = mask_np[b] if row.all() or not row.any(): continue idx = np.nonzero(row)[0] if len(idx) == 0: continue left, right = int(idx[0]), int(idx[-1]) + 1 fs = max(left - cf_frames, 0) if left - fs > 0: soft[b, fs:left] = np.linspace(0, 1, left - fs + 2)[1:-1] fe = min(right + cf_frames, T) if fe - right > 0: soft[b, right:fe] = np.linspace(1, 0, fe - right + 2)[1:-1] m = mx.expand_dims(mx.array(soft), axis=-1) return m * x_gen + (1.0 - m) * clean_src def mlx_generate_diffusion( mlx_decoder, encoder_hidden_states_np: np.ndarray, context_latents_np: np.ndarray, src_latents_shape: Tuple[int, ...], seed: Optional[Union[int, List[int]]] = None, infer_method: str = "ode", shift: float = 3.0, timesteps: Optional[list] = None, infer_steps: Optional[int] = None, guidance_scale: float = 1.0, null_condition_emb_np: Optional[np.ndarray] = None, cfg_interval_start: float = 0.0, cfg_interval_end: float = 1.0, audio_cover_strength: float = 1.0, encoder_hidden_states_non_cover_np: Optional[np.ndarray] = None, context_latents_non_cover_np: Optional[np.ndarray] = None, retake_seed: Optional[Union[int, List[int]]] = None, retake_variance: float = 0.0, compile_model: bool = False, disable_tqdm: bool = False, sampler_mode: str = "euler", velocity_norm_threshold: float = 0.0, velocity_ema_factor: float = 0.0, dcw_enabled: bool = True, dcw_mode: str = "double", dcw_scaler: float = 0.05, dcw_high_scaler: float = 0.02, dcw_wavelet: str = "haar", repaint_mask_np: Optional[np.ndarray] = None, clean_src_latents_np: Optional[np.ndarray] = None, repaint_crossfade_frames: int = 10, repaint_injection_ratio: float = 0.5, ) -> Dict[str, object]: """Run the complete MLX diffusion loop with optional CFG guidance. This is the core generation function. It accepts numpy arrays (converted from PyTorch tensors by the handler) and returns numpy arrays that the handler converts back to PyTorch. Args: mlx_decoder: ``MLXDiTDecoder`` instance with loaded weights. encoder_hidden_states_np: [B, enc_L, D] from prepare_condition (numpy). context_latents_np: [B, T, C] from prepare_condition (numpy). src_latents_shape: shape tuple [B, T, 64] for noise generation. seed: random seed (int, list[int], or None). infer_method: "ode" or "sde". shift: timestep shift factor. timesteps: optional custom timestep list. infer_steps: number of diffusion steps. guidance_scale: CFG guidance strength (>1.0 enables CFG). null_condition_emb_np: [1, 1, D] null condition embedding for CFG. cfg_interval_start: timestep ratio below which CFG is disabled. cfg_interval_end: timestep ratio above which CFG is disabled. audio_cover_strength: cover strength (0-1). encoder_hidden_states_non_cover_np: optional [B, enc_L, D] for non-cover. context_latents_non_cover_np: optional [B, T, C] for non-cover. compile_model: If True, compile the decoder step with ``mx.compile``. disable_tqdm: If True, suppress the diffusion progress bar. sampler_mode: Sampler algorithm — ``"euler"`` (first-order, default) or ``"heun"`` (second-order predictor-corrector for cleaner output). velocity_norm_threshold: Clamp velocity prediction L2 norm relative to input norm at each step. 0 disables (default). Values around 1.5–3.0 reduce outlier artefacts. velocity_ema_factor: Blend current velocity prediction with the previous step's prediction via EMA (``vt = (1-f)*vt + f*prev``). 0 disables (default). Values around 0.05–0.2 smooth the trajectory. Returns: Dict with ``"target_latents"`` (numpy) and ``"time_costs"`` dict. """ import mlx.core as mx from .dit_model import MLXCrossAttentionCache if sampler_mode not in VALID_SAMPLER_MODES: raise ValueError(f"Unsupported sampler_mode '{sampler_mode}'. Expected one of {VALID_SAMPLER_MODES}.") use_heun = sampler_mode == "heun" use_norm_clamp = velocity_norm_threshold > 0 use_ema = velocity_ema_factor > 0 if use_heun: if infer_method == "sde": logger.warning( "[MLX-DiT] Heun sampler is not supported with SDE inference method. " "Falling back to Euler for SDE steps. Use infer_method='ode' for Heun." ) else: logger.info("[MLX-DiT] Using Heun (second-order) sampler for higher-quality output.") if use_norm_clamp: logger.info("[MLX-DiT] Velocity norm clamping enabled (threshold=%.2f).", velocity_norm_threshold) if use_ema: logger.info("[MLX-DiT] Velocity EMA smoothing enabled (factor=%.3f).", velocity_ema_factor) time_costs = {} total_start = time.time() enc_hs = mx.array(encoder_hidden_states_np) ctx = mx.array(context_latents_np) enc_hs_nc = mx.array(encoder_hidden_states_non_cover_np) if encoder_hidden_states_non_cover_np is not None else None ctx_nc = mx.array(context_latents_non_cover_np) if context_latents_non_cover_np is not None else None # ---- Repaint setup ---- do_repaint = repaint_mask_np is not None and clean_src_latents_np is not None repaint_mask_mx = mx.array(repaint_mask_np) if do_repaint else None clean_src_mx = mx.array(clean_src_latents_np) if do_repaint else None bsz = src_latents_shape[0] T = src_latents_shape[1] C = src_latents_shape[2] # ---- CFG setup ---- do_cfg = guidance_scale > 1.0 and null_condition_emb_np is not None null_cond = mx.array(null_condition_emb_np) if do_cfg else None if do_cfg: null_expanded = mx.broadcast_to(null_cond, enc_hs.shape) enc_hs = mx.concatenate([enc_hs, null_expanded], axis=0) ctx = mx.concatenate([ctx, ctx], axis=0) if enc_hs_nc is not None: null_expanded_nc = mx.broadcast_to(null_cond, enc_hs_nc.shape) enc_hs_nc = mx.concatenate([enc_hs_nc, null_expanded_nc], axis=0) if ctx_nc is not None: ctx_nc = mx.concatenate([ctx_nc, ctx_nc], axis=0) momentum_state: Optional[Dict] = {} if do_cfg else None # ---- Noise preparation ---- def _draw_noise(_seed): if _seed is None: return mx.random.normal((bsz, T, C)) if isinstance(_seed, list): parts = [] for s in _seed: if s is None or s < 0: parts.append(mx.random.normal((1, T, C))) else: key = mx.random.key(int(s)) parts.append(mx.random.normal((1, T, C), key=key)) return mx.concatenate(parts, axis=0) key = mx.random.key(int(_seed)) return mx.random.normal((bsz, T, C), key=key) noise = _draw_noise(seed) # Retake mixing: variance-preserving blend with an independent noise draw. # v=0 -> noise unchanged; v=1 -> equivalent to using retake_seed as the main seed. if retake_variance > 0.0: retake_noise = _draw_noise(retake_seed) v_rad = retake_variance * (math.pi / 2.0) noise = math.cos(v_rad) * noise + math.sin(v_rad) * retake_noise # ---- Timestep schedule ---- t_schedule_list = get_timestep_schedule(shift, timesteps, infer_steps=infer_steps) num_steps = len(t_schedule_list) cover_steps = int(num_steps * audio_cover_strength) # ---- Prepare decoder step (compiled or plain with KV cache) ---- _compiled_step = None if compile_model: def _raw_step(xt, t, tr, enc, ctx): vt, _ = mlx_decoder( hidden_states=xt, timestep=t, timestep_r=tr, encoder_hidden_states=enc, context_latents=ctx, cache=None, use_cache=False, ) return vt try: _compiled_step = mx.compile(_raw_step) logger.info("[MLX-DiT] Diffusion step compiled with mx.compile().") except Exception as exc: logger.warning( "[MLX-DiT] mx.compile() failed (%s); using uncompiled path.", exc ) # Note: Heun solver requires two model evaluations per step with # different inputs, so we disable KV caching when using it. if use_heun: cache = None else: cache = MLXCrossAttentionCache() if _compiled_step is None else None xt = noise prev_vt = None # for EMA smoothing def _model_eval(x_input, t_val, enc, ctx_in, step_cache): """Single model evaluation helper.""" t_arr = mx.full((x_input.shape[0],), t_val) if _compiled_step is not None: return _compiled_step(x_input, t_arr, t_arr, enc, ctx_in), step_cache vt_out, step_cache = mlx_decoder( hidden_states=x_input, timestep=t_arr, timestep_r=t_arr, encoder_hidden_states=enc, context_latents=ctx_in, cache=step_cache, use_cache=(not do_cfg and not use_heun), ) return vt_out, step_cache def _apply_cfg(vt_raw, current_t_val): """Apply CFG guidance if enabled.""" if not do_cfg: return vt_raw pred_cond = vt_raw[:bsz] pred_uncond = vt_raw[bsz:] if cfg_interval_start <= current_t_val <= cfg_interval_end: return _mlx_apg_forward(pred_cond, pred_uncond, guidance_scale, momentum_state) return pred_cond def _apply_stabilisation(vt_guided, xt_current, prev_velocity): """Apply optional norm clamping and EMA smoothing.""" # Velocity norm clamping — prevents outlier predictions if use_norm_clamp: vt_norm = mx.sqrt((vt_guided * vt_guided).sum(axis=(1, 2), keepdims=True)) xt_norm = mx.sqrt((xt_current * xt_current).sum(axis=(1, 2), keepdims=True)) + 1e-10 scale = mx.minimum( mx.ones_like(vt_norm), (velocity_norm_threshold * xt_norm) / (vt_norm + 1e-10), ) vt_guided = vt_guided * scale # Velocity EMA smoothing — stabilises denoising trajectory if use_ema and prev_velocity is not None: vt_guided = (1.0 - velocity_ema_factor) * vt_guided + velocity_ema_factor * prev_velocity return vt_guided diff_start = time.time() _switched_to_non_cover = False # DCW — opt-in per-band wavelet-domain correction (CVPR 2026). On MLX, # `haar` runs natively; other wavelets bridge through pytorch_wavelets # for output parity with the CUDA/CPU PyTorch path. See # `acestep.models.mlx.dcw_correction_mlx`. from acestep.models.mlx.dcw_correction_mlx import apply_mlx_dcw dcw_active = dcw_enabled and ( dcw_scaler != 0.0 or (dcw_mode == "double" and dcw_high_scaler != 0.0) ) if dcw_active: _backend = "MLX-native Haar" if dcw_wavelet == "haar" else f"torch bridge ({dcw_wavelet})" logger.info( "[MLX-DiT] DCW enabled (mode=%s, scaler=%.3f, high_scaler=%.3f, wavelet=%s, backend=%s).", dcw_mode, dcw_scaler, dcw_high_scaler, dcw_wavelet, _backend, ) for step_idx in tqdm(range(num_steps), desc="MLX DiT diffusion", disable=disable_tqdm): current_t = t_schedule_list[step_idx] # Switch to non-cover conditions when appropriate if step_idx >= cover_steps and not _switched_to_non_cover: _switched_to_non_cover = True if enc_hs_nc is not None: enc_hs = enc_hs_nc ctx = ctx_nc if cache is not None: cache = MLXCrossAttentionCache() # Build input: double batch for CFG x_in = mx.concatenate([xt, xt], axis=0) if do_cfg else xt # ---- First model evaluation (predictor) ---- vt, cache = _model_eval(x_in, current_t, enc_hs, ctx, cache) mx.eval(vt) vt = _apply_cfg(vt, current_t) vt = _apply_stabilisation(vt, xt, prev_vt) # Cache pre-step latent so DCW can reconstruct the predicted clean # sample ``denoised = x_before - v * t`` after the sampler update. # Also stash the raw velocity (pre-Heun-averaging) so the x0 # reconstruction uses the single-evaluation ``v(t_curr)``, matching # the reference FLUX scheduler's ``x0 = sample - sigma * v``. xt_before_step = xt vt_for_denoise = vt # Final step: compute x0 if step_idx == num_steps - 1: t_unsq = mx.full((bsz, 1, 1), current_t) xt = xt - vt * t_unsq mx.eval(xt) else: next_t = t_schedule_list[step_idx + 1] if use_heun and infer_method == "ode": # ---- Heun (second-order) ODE step ---- # Predictor: Euler step to get xt_predicted at next_t dt = current_t - next_t dt_arr = mx.full((bsz, 1, 1), dt) xt_predicted = xt - vt * dt_arr mx.eval(xt_predicted) # Corrector: evaluate model at the predicted point x_in2 = mx.concatenate([xt_predicted, xt_predicted], axis=0) if do_cfg else xt_predicted vt2, cache = _model_eval(x_in2, next_t, enc_hs, ctx, cache) mx.eval(vt2) vt2 = _apply_cfg(vt2, next_t) vt2 = _apply_stabilisation(vt2, xt_predicted, vt) # Average the two velocity predictions (trapezoidal rule) vt_avg = 0.5 * (vt + vt2) xt = xt - vt_avg * dt_arr vt = vt_avg # store averaged velocity for EMA elif infer_method == "sde": t_unsq = mx.full((bsz, 1, 1), current_t) pred_clean = xt - vt * t_unsq new_noise = mx.random.normal(xt.shape) xt = next_t * new_noise + (1.0 - next_t) * pred_clean else: # ---- Standard Euler ODE step ---- dt = current_t - next_t dt_arr = mx.full((bsz, 1, 1), dt) xt = xt - vt * dt_arr mx.eval(xt) # DCW correction — push x_next's frequency bands away from the # predicted clean sample. Scaler decays with t_curr so this is # identity at t=0 and strongest at t≈1. if dcw_active: t_unsq_d = mx.full((bsz, 1, 1), current_t) denoised = xt_before_step - vt_for_denoise * t_unsq_d xt = apply_mlx_dcw( xt, denoised, t_curr=current_t, enabled=True, mode=dcw_mode, scaler=dcw_scaler, high_scaler=dcw_high_scaler, wavelet=dcw_wavelet, ) mx.eval(xt) prev_vt = vt # store for EMA # ---- Repaint step injection ---- if do_repaint: injection_cutoff = round(repaint_injection_ratio * num_steps) if step_idx < injection_cutoff: t_after = t_schedule_list[step_idx + 1] if step_idx < num_steps - 1 else 0.0 xt = _mlx_repaint_step_injection(xt, clean_src_mx, repaint_mask_mx, t_after, noise) mx.eval(xt) # ---- Repaint boundary blend (post-loop) ---- if do_repaint and repaint_crossfade_frames > 0: xt = _mlx_repaint_boundary_blend(xt, clean_src_mx, repaint_mask_np, repaint_crossfade_frames) mx.eval(xt) diff_end = time.time() total_end = time.time() time_costs["diffusion_time_cost"] = diff_end - diff_start time_costs["diffusion_per_step_time_cost"] = time_costs["diffusion_time_cost"] / max(num_steps, 1) time_costs["total_time_cost"] = total_end - total_start time_costs["sampler_mode"] = sampler_mode result_np = np.array(xt) return { "target_latents": result_np, "time_costs": time_costs, }