{"repo": "blaizzy/mlx-audio", "files": {"mlx_audio/tts/models/longcat_audiodit/longcat_audiodit.py": {"223afd602c6d583455bdb9f181b91fc090f4aff9": {"t": 1785674297.5328586, "n": 24183, "h": "57f7b826864b1c08", "c": "\"\"\"LongCat-AudioDiT: Conditional Flow Matching TTS with DiT backbone.\n\nPorts meituan-longcat/LongCat-AudioDiT to MLX. Components:\n  1. UMT5 text encoder (frozen)\n  2. DiT transformer (CrossDiT with AdaLN, RoPE, cross-attention)\n  3. WAV-VAE audio codec (frozen, 24kHz, latent_dim=64)\n\"\"\"\n\nimport math\nimport re\nimport time as time_module\nfrom typing import Generator, Optional\n\nimport mlx.core as mx\nimport mlx.nn as nn\nimport numpy as np\n\nfrom mlx_audio.tts.models.base import GenerationResult\n\nfrom .config import ModelConfig\nfrom .dit import AudioDiTTransformer\nfrom .text_encoder import UMT5Encoder\nfrom .vae import AudioDiTVae\n\n# ---------------------------------------------------------------------------\n# Duration estimation heuristic\n# ---------------------------------------------------------------------------\n\nEN_DUR_PER_CHAR = 0.082\nZH_DUR_PER_CHAR = 0.21\n\n\ndef _normalize_text(text: str) -> str:\n    text = text.lower()\n    text = re.sub(r'[\"\"\"\\u2018\\u2019]', \" \", text)\n    text = re.sub(r\"\\s+\", \" \", text)\n    return text.strip()\n\n\ndef _approx_duration(text: str, max_duration: float = 30.0) -> float:\n    text = re.sub(r\"\\s+\", \"\", text)\n    num_zh = num_en = num_other = 0\n    for c in text:\n        if \"\\u4e00\" <= c <= \"\\u9fff\":\n            num_zh += 1\n        elif c.isalpha():\n            num_en += 1\n        else:\n            num_other += 1\n    if num_zh > num_en:\n        num_zh += num_other\n    else:\n        num_en += num_other\n    return min(max_duration, num_zh * ZH_DUR_PER_CHAR + num_en * EN_DUR_PER_CHAR)\n\n\n# ---------------------------------------------------------------------------\n# APG helpers (Adaptive Projected Guidance)\n# ---------------------------------------------------------------------------\n\n\nclass _MomentumBuffer:\n    def __init__(self, momentum: float = -0.75):\n        self.momentum = momentum\n        self.running_average = 0\n\n    def update(self, update_value: mx.array):\n        new_average = self.momentum * self.running_average\n        self.running_average = update_value + new_average\n\n\ndef _project(v0: mx.array, v1: mx.array):\n    v0 = v0.astype(mx.float32)\n    v1 = v1.astype(mx.float32)\n    v1_norm = v1 / (mx.sqrt(mx.sum(v1 * v1, axis=(-1, -2), keepdims=True)) + 1e-8)\n    v0_parallel = mx.sum(v0 * v1_norm, axis=(-1, -2), keepdims=True) * v1_norm\n    v0_orthogonal = v0 - v0_parallel\n    return v0_parallel, v0_orthogonal\n\n\ndef _apg_forward(\n    pred_cond,\n    pred_uncond,\n    guidance_scale,\n    momentum_buffer=None,\n    eta=0.0,\n    norm_threshold=2.5,\n):\n    diff = pred_cond - pred_uncond\n    if momentum_buffer is not None:\n        momentum_buffer.update(diff)\n        diff = momentum_buffer.running_average\n    if norm_threshold > 0:\n        diff_norm = mx.sqrt(mx.sum(diff * diff, axis=(-1, -2), keepdims=True))\n        scale_factor = mx.minimum(mx.ones_like(diff_norm), norm_threshold / diff_norm)\n        diff = diff * scale_factor\n    diff_parallel, diff_orthogonal = _project(diff, pred_cond)\n    normalized_update = diff_orthogonal + eta * diff_parallel\n    return pred_cond + guidance_scale * normalized_update\n\n\n# ---------------------------------------------------------------------------\n# Euler ODE solver\n# ---------------------------------------------------------------------------\n\n\ndef _odeint_euler(fn, y0, t_steps):\n    y = y0\n    for i in range(len(t_steps) - 1):\n        dt = t_steps[i + 1] - t_steps[i]\n        y = y + fn(t_steps[i], y) * dt\n    return y\n\n\n# ---------------------------------------------------------------------------\n# Main model\n# ---------------------------------------------------------------------------\n\n\nclass Model(nn.Module):\n    def __init__(self, config: ModelConfig):\n        super().__init__()\n        self.config = config\n\n        # Text encoder (UMT5)\n        self.text_encoder = UMT5Encoder(config.text_encoder_config)\n\n        # DiT transformer\n        self.transformer = AudioDiTTransformer(config)\n\n        # WAV-VAE\n        self.vae = AudioDiTVae(config.vae_config)\n\n    def model_quant_predicate(self, p, m):\n        \"\"\"Skip quantization for VAE (conv-based codec).\"\"\"\n        return not p.startswith(\"vae\")\n\n    @property\n    def sample_rate(self) -> int:\n        return self.config.sampling_rate\n\n    def encode_text(self, input_ids: mx.array, attention_mask: mx.array):\n        \"\"\"Encode text with UMT5 and apply text_add_embed / text_norm_feat.\"\"\"\n        last_hidden, initial_embed = self.text_encoder(input_ids, attention_mask)\n        d_model = self.config.text_encoder_config.d_model\n\n        if self.config.text_norm_feat:\n            last_hidden = _layer_norm(last_hidden, d_model)\n\n        if self.config.text_add_embed:\n            if self.config.text_norm_feat:\n                initial_embed = _layer_norm(initial_embed, d_model)\n            last_hidden = last_hidden + initial_embed\n\n        return last_hidden.astype(mx.float32)\n\n    def encode_prompt_audio(self, prompt_audio: mx.array):\n        \"\"\"Encode prompt audio to latent space. Returns (latent, num_frames).\"\"\"\n        full_hop = self.config.latent_hop\n        off = 3\n        wav = prompt_audio\n        if wav.ndim == 2:\n            wav = wav[..., None]  # (B, L) -> (B, L, 1)\n        if wav.ndim == 3 and wav.shape[-1] != 1:\n            wav = wav.transpose(0, 2, 1)  # (B, 1, L) -> (B, L, 1)\n\n        L = wav.shape[1]\n        if L % full_hop != 0:\n            pad_amount = full_hop - L % full_hop\n            wav = mx.pad(wav, [(0, 0), (0, pad_amount), (0, 0)])\n        wav = mx.pad(wav, [(0, 0), (0, full_hop * off), (0, 0)])\n\n        latent = self.vae.encode(wav)  # (B, T, latent_dim)\n        if off != 0:\n            latent = latent[:, :-off, :]\n        prompt_dur = latent.shape[1]\n        return latent, prompt_dur\n\n    @staticmethod\n    def _format_duration(seconds: float) -> str:\n        return f\"{int(seconds // 3600):02d}:{int((seconds % 3600) // 60):02d}:{seconds % 60:06.3f}\"\n\n    def _stream_decode(\n        self,\n        pred_latent: mx.array,\n        sr: int,\n        start_time: float,\n        chunk_seconds: float = 2.0,\n        overlap_seconds: float = 0.5,\n    ) -> Generator[GenerationResult, None, None]:\n        \"\"\"Decode latents in overlapping chunks with cosine crossfade.\n\n        Runs the VAE decoder on small overlapping windows of the latent\n        sequence and yields audio as soon as each chunk is ready, giving\n        much lower time-to-first-audio than full-sequence decoding.\n        \"\"\"\n        ratio = self.config.vae_config.downsampling_ratio\n        chunk_frames = max(1, int(chunk_seconds * sr / ratio))\n        overlap_frames = max(0, int(overlap_seconds * sr / ratio))\n        hop_frames = max(1, chunk_frames - overlap_frames)\n        overlap_samples = overlap_frames * ratio\n        # Context frames fed to the VAE decoder on each side so its\n        # convolutions have proper context, then trimmed from the output.\n        context_frames = overlap_frames\n\n        total_frames = pred_latent.shape[1]\n        prev_tail = None\n        chunk_idx = 0\n        cumulative_samples = 0\n\n        start = 0\n        while start < total_frames:\n            end = min(start + chunk_frames, total_frames)\n            is_last = end >= total_frames\n\n            # Extend chunk with context on both sides for decoder conv context\n            left_ctx = min(context_frames, start)\n            right_ctx = min(context_frames, total_frames - end)\n            ctx_start = start - left_ctx\n            ctx_end = end + right_ctx\n\n            # Decode extended chunk, then trim context audio\n            audio_full = self.vae.decode(pred_latent[:, ctx_start:ctx_end, :])\n            audio_full = audio_full.squeeze(-1).squeeze(0)  # (L,)\n            mx.eval(audio_full)\n            left_trim = left_ctx * ratio\n            right_trim = right_ctx * ratio\n            if right_trim > 0:\n                audio_chunk = audio_full[left_trim:-right_trim]\n            else:\n                audio_chunk = audio_full[left_trim:]\n\n            if prev_tail is not None and overlap_samples > 0:\n                # Cosine crossfade with previous chunk's tail\n                ol = min(overlap_samples, prev_tail.shape[0], audio_chunk.shape[0])\n                t = mx.linspace(0, 1, ol)\n                fade_in = 0.5 * (1 - mx.cos(mx.array(math.pi) * t))\n                fade_out = 1.0 - fade_in\n                blended = prev_tail[:ol] * fade_out + audio_chunk[:ol] * fade_in\n\n                if is_last:\n                    output = mx.concatenate([blended, audio_chunk[ol:]])\n                else:\n                    output = mx.concatenate([blended, audio_chunk[ol:-overlap_samples]])\n                    prev_tail = audio_chunk[-overlap_samples:]\n                    mx.eval(prev_tail)\n            else:\n                if is_last or overlap_frames == 0:\n                    output = audio_chunk\n                else:\n                    output = audio_chunk[:-overlap_samples]\n                    prev_tail = audio_chunk[-overlap_samples:]\n                    mx.eval(prev_tail)\n\n            mx.eval(output)\n            num_samples = output.shape[0]\n            cumulative_samples += num_samples\n            processing_time = time_module.time() - start_time\n            audio_dur = cumulative_samples / sr\n\n            yield GenerationResult(\n                audio=output,\n                samples=num_samples,\n                sample_rate=sr,\n                segment_idx=chunk_idx,\n                token_count=0,\n                audio_duration=self._format_duration(audio_dur),\n                real_time_factor=processing_time / max(audio_dur, 1e-6),\n                prompt={\"tokens\": 0, \"tokens-per-sec\": 0},\n                audio_samples={\n                    \"samples\": num_samples,\n                    \"samples-per-sec\": num_samples / max(processing_time, 1e-6),\n                },\n                processing_time_seconds=processing_time,\n                peak_memory_usage=mx.get_peak_memory() / 1e9,\n                is_streaming_chunk=True,\n                is_final_chunk=is_last,\n            )\n\n            chunk_idx += 1\n            start += hop_frames\n\n    def generate(\n        self,\n        text: str,\n        voice: Optional[str] = None,\n        speed: float = 1.0,\n        lang_code: str = \"en\",\n        ref_audio: Optional[mx.array] = None,\n        ref_text: Optional[str] = None,\n        steps: int = 16,\n        cfg_strength: float = 4.0,\n        guidance_method: str = \"cfg\",\n        seed: int = 1024,\n        stream: bool = False,\n        streaming_interval: float = 2.0,\n        chunk_seconds: float = 2.0,\n        overlap_seconds: float = 0.5,\n        **kwargs,\n    ) -> Generator[GenerationResult, None, None]:\n        \"\"\"Generate audio from text.\n\n        Args:\n            text: Text to synthesize.\n            voice: Unused (kept for API compat).\n            speed: Unused (kept for API compat).\n            ref_audio: Optional reference audio array for voice cloning (samples at 24kHz).\n            ref_text: Optional text of the reference audio.\n            steps: Number of ODE Euler steps (default 16).\n            cfg_strength: Classifier-free guidance strength.\n            guidance_method: \"cfg\" or \"apg\".\n            seed: Random seed.\n            stream: Stream audio chunks during VAE decoding.\n            streaming_interval: Alias for chunk_seconds (CLI compat).\n            chunk_seconds: Chunk length in seconds for streaming decode.\n            overlap_seconds: Overlap between chunks for cosine crossfade.\n        \"\"\"\n        import transformers\n\n        start_time = time_module.time()\n        mx.random.seed(seed)\n\n        sr = self.config.sampling_rate\n        full_hop = self.config.latent_hop\n        max_duration = self.config.max_wav_duration\n        repa_layer = self.config.repa_dit_layer\n\n        # Tokenize\n        tokenizer = transformers.AutoTokenizer.from_pretrained(\n            self.config.text_encoder_model\n        )\n\n        text = _normalize_text(text)\n        no_prompt = ref_audio is None\n\n        if not no_prompt and ref_text is not None:\n            prompt_text = _normalize_text(ref_text)\n            full_text = f\"{prompt_text} {text}\"\n        else:\n            full_text = text\n\n        inputs = tokenizer([full_text], padding=\"longest\", return_tensors=\"np\")\n        input_ids = mx.array(inputs[\"input_ids\"])\n        attention_mask = mx.array(inputs[\"attention_mask\"]).astype(mx.float32)\n\n        # Text encoding\n        text_condition = self.encode_text(input_ids, attention_mask)\n        text_condition_len = mx.sum(attention_mask, axis=1)\n\n        batch = text_condition.shape[0]\n\n        # Prompt audio encoding\n        if not no_prompt:\n            prompt_audio_mx = (\n                mx.array(np.array(ref_audio))\n                if not isinstance(ref_audio, mx.array)\n                else ref_audio\n            )\n            prompt_latent, prompt_dur = self.encode_prompt_audio(\n                prompt_audio_mx[None] if prompt_audio_mx.ndim < 3 else prompt_audio_mx\n            )\n        else:\n            prompt_latent = mx.zeros((batch, 0, self.config.latent_dim))\n            prompt_dur = 0\n\n        # Duration estimation\n        prompt_time = prompt_dur * full_hop / sr\n        dur_sec = _approx_duration(text, max_duration=max_duration - prompt_time)\n        if not no_prompt and ref_text is not None:\n            approx_pd = _approx_duration(ref_text, max_duration=max_duration)\n            ratio = float(np.clip(prompt_time / max(approx_pd, 1e-6), 1.0, 1.5))\n            dur_sec = dur_sec * ratio\n\n        duration = int(dur_sec * sr // full_hop)\n        total_duration = min(duration + prompt_dur, int(max_duration * sr // full_hop))\n\n        # Masks\n        mask = mx.arange(total_duration)[None, :] < mx.array([total_duration])\n        text_mask = (\n            mx.arange(text_condition.shape[1])[None, :] < text_condition_len[:, None]\n        )\n\n        # Conditioning\n        neg_text = mx.zeros_like(text_condition)\n\n        if not no_prompt:\n            gen_len = total_duration - prompt_dur\n            latent_cond = mx.pad(prompt_latent, [(0, 0), (0, gen_len), (0, 0)])\n            empty_latent_cond = mx.zeros_like(latent_cond)\n        else:\n            latent_cond = mx.zeros((batch, total_duration, self.config.latent_dim))\n            empty_latent_cond = latent_cond\n\n        # APG buffer\n        apg_buffer = (\n            _MomentumBuffer(momentum=-0.3) if guidance_method == \"apg\" else None\n        )\n\n        # Initial noise\n        y = mx.random.normal((batch, total_duration, self.config.latent_dim))\n        prompt_noise = y[:, :prompt_dur] if prompt_dur > 0 else None\n\n        # Time steps\n        t_steps = np.linspace(0, 1, steps).tolist()\n\n        # Euler ODE integration\n        for i in range(len(t_steps) - 1):\n            t_val = t_steps[i]\n            dt = t_steps[i + 1] - t_val\n            t = mx.array([t_val])\n\n            # Set prompt region\n            if prompt_dur > 0:\n                y_prompt = (\n                    prompt_noise * (1 - t_val) + latent_cond[:, :prompt_dur] * t_val\n                )\n                y = mx.concatenate([y_prompt, y[:, prompt_dur:]], axis=1)\n\n            output = self.transformer(\n                x=y,\n                text=text_condition,\n                text_len=text_condition_len,\n                time=t,\n                mask=mask,\n                cond_mask=text_mask,\n                return_ith_layer=repa_layer,\n                latent_cond=latent_cond,\n            )\n            pred = output[\"last_hidden_state\"]\n\n            if cfg_strength >= 1e-5:\n                # Null forward for guidance\n                y_null = mx.array(y)\n                if prompt_dur > 0:\n                    y_null = mx.concatenate(\n                        [mx.zeros_like(y_null[:, :prompt_dur]), y_null[:, prompt_dur:]],\n                        axis=1,\n                    )\n\n                null_output = self.transformer(\n                    x=y_null,\n                    text=neg_text,\n                    text_len=text_condition_len,\n                    time=t,\n                    mask=mask,\n                    cond_mask=text_mask,\n                    return_ith_layer=repa_layer,\n                    latent_cond=empty_latent_cond,\n                )\n                null_pred = null_output[\"last_hidden_state\"]\n\n                if guidance_method == \"cfg\":\n                    pred = pred + (pred - null_pred) * cfg_strength\n                else:\n                    # APG\n                    x_s = y[:, prompt_dur:]\n                    pred_s = pred[:, prompt_dur:]\n                    null_s = null_pred[:, prompt_dur:]\n                    pred_sample = x_s + (1 - t_val) * pred_s\n                    null_sample = x_s + (1 - t_val) * null_s\n                    out = _apg_forward(\n                        pred_sample,\n                        null_sample,\n                        cfg_strength,\n                        apg_buffer,\n                        eta=0.5,\n                        norm_threshold=0.0,\n                    )\n                    out = (out - x_s) / (1 - t_val)\n                    pred = mx.pad(out, [(0, 0), (prompt_dur, 0), (0, 0)])\n\n            y = y + pred * dt\n            mx.eval(y)\n\n        # Decode\n        pred_latent = y[:, prompt_dur:] if not no_prompt else y\n\n        if stream:\n            # streaming_interval from CLI takes priority over chunk_seconds\n            cs = streaming_interval if streaming_interval != 2.0 else chunk_seconds\n            yield from self._stream_decode(\n                pred_latent, sr, start_time, cs, overlap_seconds\n            )\n        else:\n            waveform = self.vae.decode(pred_latent)  # (B, L, 1)\n            waveform = waveform.squeeze(-1)  # (B, L)\n\n            mx.eval(waveform)\n            processing_time = time_module.time() - start_time\n\n            audio = waveform[0]\n            num_samples = audio.shape[0]\n            audio_duration = num_samples / sr\n\n            yield GenerationResult(\n                audio=audio,\n                samples=num_samples,\n                sample_rate=sr,\n                segment_idx=0,\n                token_count=0,\n                audio_duration=self._format_duration(audio_duration),\n                real_time_factor=processing_time / max(audio_duration, 1e-6),\n                prompt={\"tokens\": 0, \"tokens-per-sec\": 0},\n                audio_samples={\n                    \"samples\": num_samples,\n                    \"samples-per-sec\": num_samples / max(processing_time, 1e-6),\n                },\n                processing_time_seconds=processing_time,\n                peak_memory_usage=mx.get_peak_memory() / 1e9,\n            )\n\n    def sanitize(self, weights: dict) -> dict:\n        \"\"\"Convert PyTorch AudioDiT weights to MLX format.\"\"\"\n        sanitized = {}\n\n        # Collect weight_v/weight_g pairs for reconstruction\n        weight_v_keys = [k for k in weights if k.endswith(\".weight_v\")]\n        processed = set()\n\n        for wv_key in weight_v_keys:\n            base = wv_key[: -len(\".weight_v\")]\n            wg_key = base + \".weight_g\"\n            if wg_key not in weights:\n                continue\n\n            weight_v = weights[wv_key]\n            weight_g = weights[wg_key]\n\n            # Determine if this is ConvTranspose1d\n            # ConvTranspose1d is at decoder.layers.{block}.layers.1\n            is_conv_transpose = bool(\n                re.search(r\"vae\\.decoder\\.layers\\.\\d+\\.layers\\.1\\.weight\", wv_key)\n            )\n\n            if weight_v.ndim == 3:\n                # Reconstruct weight from weight_norm: w = g * v / ||v||\n                if is_conv_transpose:\n                    # PyTorch ConvTranspose1d: (in_ch, out_ch, ksize)\n                    # norm over dims 1,2 (out_ch and ksize)\n                    norm_axes = (1, 2)\n                    norm = mx.sqrt(\n                        mx.sum(weight_v * weight_v, axis=norm_axes, keepdims=True)\n                        + 1e-12\n                    )\n                    weight = weight_g * weight_v / norm\n                    # Transpose: (in_ch, out_ch, ksize) -> (out_ch, ksize, in_ch)\n                    weight = weight.transpose(1, 2, 0)\n                else:\n                    # PyTorch Conv1d: (out_ch, in_ch, ksize)\n                    # norm over dims 1,2 (in_ch and ksize)\n                    norm_axes = (1, 2)\n                    norm = mx.sqrt(\n                        mx.sum(weight_v * weight_v, axis=norm_axes, keepdims=True)\n                        + 1e-12\n                    )\n                    weight = weight_g * weight_v / norm\n                    # Transpose: (out_ch, in_ch, ksize) -> (out_ch, ksize, in_ch)\n                    weight = weight.transpose(0, 2, 1)\n\n                sanitized[base + \".weight\"] = weight\n            else:\n                sanitized[base + \".weight\"] = weight_v\n\n            processed.add(wv_key)\n            processed.add(wg_key)\n\n        # Process remaining weights\n        for key, value in weights.items():\n            if key in processed:\n                continue\n\n            new_key = key\n\n            # --- Text encoder remapping ---\n            if new_key.startswith(\"text_encoder.encoder.embed_tokens.\"):\n                new_key = new_key.replace(\n                    \"text_encoder.encoder.embed_tokens.\", \"text_encoder.shared.\"\n                )\n\n            # T5 block layer remapping:\n            # HF: block.{i}.layer.0.SelfAttention -> block.{i}.SelfAttention\n            # HF: block.{i}.layer.0.layer_norm -> block.{i}.layer_norm_sa\n            # HF: block.{i}.layer.1.DenseReluDense -> block.{i}.DenseReluDense\n            # HF: block.{i}.layer.1.layer_norm -> block.{i}.layer_norm_ff\n            new_key = re.sub(\n                r\"text_encoder\\.encoder\\.block\\.(\\d+)\\.layer\\.0\\.SelfAttention\\.\",\n                r\"text_encoder.block.\\1.SelfAttention.\",\n                new_key,\n            )\n            new_key = re.sub(\n                r\"text_encoder\\.encoder\\.block\\.(\\d+)\\.layer\\.0\\.layer_norm\\.\",\n                r\"text_encoder.block.\\1.layer_norm_sa.\",\n                new_key,\n            )\n            new_key = re.sub(\n                r\"text_encoder\\.encoder\\.block\\.(\\d+)\\.layer\\.1\\.DenseReluDense\\.\",\n                r\"text_encoder.block.\\1.DenseReluDense.\",\n                new_key,\n            )\n            new_key = re.sub(\n                r\"text_encoder\\.encoder\\.block\\.(\\d+)\\.layer\\.1\\.layer_norm\\.\",\n                r\"text_encoder.block.\\1.layer_norm_ff.\",\n                new_key,\n            )\n            new_key = new_key.replace(\n                \"text_encoder.encoder.final_layer_norm.\",\n                \"text_encoder.final_layer_norm.\",\n            )\n\n            # --- Transformer Sequential index remapping ---\n            # Embedder: proj.0 -> proj.0, proj.2 -> proj.1\n            new_key = re.sub(r\"\\.proj\\.2\\.\", \".proj.1.\", new_key)\n            # TimestepEmbedding: time_mlp.0 -> time_mlp.0, time_mlp.2 -> time_mlp.1\n            new_key = re.sub(r\"\\.time_mlp\\.2\\.\", \".time_mlp.1.\", new_key)\n            # AdaLNMLP: mlp.1 -> mlp.0\n            new_key = re.sub(r\"\\.mlp\\.1\\.\", \".mlp.0.\", new_key)\n            # Attention to_out: to_out.0 -> to_out\n            new_key = re.sub(r\"\\.to_out\\.0\\.\", \".to_out.\", new_key)\n            # FeedForward: ff.3 -> ff.1\n            new_key = re.sub(r\"\\.ff\\.3\\.\", \".ff.1.\", new_key)\n\n            # --- ConvNeXtV2 depthwise conv weight ---\n            # dwconv.weight: (dim, 1, ksize) -> (dim, ksize, 1) for channels-last\n            if \"dwconv.weight\" in new_key and value.ndim == 3:\n                value = value.transpose(0, 2, 1)\n                new_key = new_key.replace(\".dwconv.weight\", \".dwconv_weight\")\n            elif \"dwconv.bias\" in new_key:\n                new_key = new_key.replace(\".dwconv.bias\", \".dwconv_bias\")\n\n            # --- Handle standalone Conv1d weights (non weight-normed) ---\n            # These might be in the ConvNeXtV2 dwconv or other places\n\n            sanitized[new_key] = value\n\n        return sanitized\n\n\ndef _layer_norm(x: mx.array, dim: int, eps: float = 1e-6) -> mx.array:\n    \"\"\"Layer normalization.\"\"\"\n    mean = mx.mean(x, axis=-1, keepdims=True)\n    var = mx.var(x, axis=-1, keepdims=True)\n    return (x - mean) * mx.rsqrt(var + eps)\n"}, "2c9461f5d8315fa8e7013ab2729495b2bb83d384": {"t": 1785846870.95548, "n": 24183, "h": "57f7b826864b1c08", "c": "\"\"\"LongCat-AudioDiT: Conditional Flow Matching TTS with DiT backbone.\n\nPorts meituan-longcat/LongCat-AudioDiT to MLX. Components:\n  1. UMT5 text encoder (frozen)\n  2. DiT transformer (CrossDiT with AdaLN, RoPE, cross-attention)\n  3. WAV-VAE audio codec (frozen, 24kHz, latent_dim=64)\n\"\"\"\n\nimport math\nimport re\nimport time as time_module\nfrom typing import Generator, Optional\n\nimport mlx.core as mx\nimport mlx.nn as nn\nimport numpy as np\n\nfrom mlx_audio.tts.models.base import GenerationResult\n\nfrom .config import ModelConfig\nfrom .dit import AudioDiTTransformer\nfrom .text_encoder import UMT5Encoder\nfrom .vae import AudioDiTVae\n\n# ---------------------------------------------------------------------------\n# Duration estimation heuristic\n# ---------------------------------------------------------------------------\n\nEN_DUR_PER_CHAR = 0.082\nZH_DUR_PER_CHAR = 0.21\n\n\ndef _normalize_text(text: str) -> str:\n    text = text.lower()\n    text = re.sub(r'[\"\"\"\\u2018\\u2019]', \" \", text)\n    text = re.sub(r\"\\s+\", \" \", text)\n    return text.strip()\n\n\ndef _approx_duration(text: str, max_duration: float = 30.0) -> float:\n    text = re.sub(r\"\\s+\", \"\", text)\n    num_zh = num_en = num_other = 0\n    for c in text:\n        if \"\\u4e00\" <= c <= \"\\u9fff\":\n            num_zh += 1\n        elif c.isalpha():\n            num_en += 1\n        else:\n            num_other += 1\n    if num_zh > num_en:\n        num_zh += num_other\n    else:\n        num_en += num_other\n    return min(max_duration, num_zh * ZH_DUR_PER_CHAR + num_en * EN_DUR_PER_CHAR)\n\n\n# ---------------------------------------------------------------------------\n# APG helpers (Adaptive Projected Guidance)\n# ---------------------------------------------------------------------------\n\n\nclass _MomentumBuffer:\n    def __init__(self, momentum: float = -0.75):\n        self.momentum = momentum\n        self.running_average = 0\n\n    def update(self, update_value: mx.array):\n        new_average = self.momentum * self.running_average\n        self.running_average = update_value + new_average\n\n\ndef _project(v0: mx.array, v1: mx.array):\n    v0 = v0.astype(mx.float32)\n    v1 = v1.astype(mx.float32)\n    v1_norm = v1 / (mx.sqrt(mx.sum(v1 * v1, axis=(-1, -2), keepdims=True)) + 1e-8)\n    v0_parallel = mx.sum(v0 * v1_norm, axis=(-1, -2), keepdims=True) * v1_norm\n    v0_orthogonal = v0 - v0_parallel\n    return v0_parallel, v0_orthogonal\n\n\ndef _apg_forward(\n    pred_cond,\n    pred_uncond,\n    guidance_scale,\n    momentum_buffer=None,\n    eta=0.0,\n    norm_threshold=2.5,\n):\n    diff = pred_cond - pred_uncond\n    if momentum_buffer is not None:\n        momentum_buffer.update(diff)\n        diff = momentum_buffer.running_average\n    if norm_threshold > 0:\n        diff_norm = mx.sqrt(mx.sum(diff * diff, axis=(-1, -2), keepdims=True))\n        scale_factor = mx.minimum(mx.ones_like(diff_norm), norm_threshold / diff_norm)\n        diff = diff * scale_factor\n    diff_parallel, diff_orthogonal = _project(diff, pred_cond)\n    normalized_update = diff_orthogonal + eta * diff_parallel\n    return pred_cond + guidance_scale * normalized_update\n\n\n# ---------------------------------------------------------------------------\n# Euler ODE solver\n# ---------------------------------------------------------------------------\n\n\ndef _odeint_euler(fn, y0, t_steps):\n    y = y0\n    for i in range(len(t_steps) - 1):\n        dt = t_steps[i + 1] - t_steps[i]\n        y = y + fn(t_steps[i], y) * dt\n    return y\n\n\n# ---------------------------------------------------------------------------\n# Main model\n# ---------------------------------------------------------------------------\n\n\nclass Model(nn.Module):\n    def __init__(self, config: ModelConfig):\n        super().__init__()\n        self.config = config\n\n        # Text encoder (UMT5)\n        self.text_encoder = UMT5Encoder(config.text_encoder_config)\n\n        # DiT transformer\n        self.transformer = AudioDiTTransformer(config)\n\n        # WAV-VAE\n        self.vae = AudioDiTVae(config.vae_config)\n\n    def model_quant_predicate(self, p, m):\n        \"\"\"Skip quantization for VAE (conv-based codec).\"\"\"\n        return not p.startswith(\"vae\")\n\n    @property\n    def sample_rate(self) -> int:\n        return self.config.sampling_rate\n\n    def encode_text(self, input_ids: mx.array, attention_mask: mx.array):\n        \"\"\"Encode text with UMT5 and apply text_add_embed / text_norm_feat.\"\"\"\n        last_hidden, initial_embed = self.text_encoder(input_ids, attention_mask)\n        d_model = self.config.text_encoder_config.d_model\n\n        if self.config.text_norm_feat:\n            last_hidden = _layer_norm(last_hidden, d_model)\n\n        if self.config.text_add_embed:\n            if self.config.text_norm_feat:\n                initial_embed = _layer_norm(initial_embed, d_model)\n            last_hidden = last_hidden + initial_embed\n\n        return last_hidden.astype(mx.float32)\n\n    def encode_prompt_audio(self, prompt_audio: mx.array):\n        \"\"\"Encode prompt audio to latent space. Returns (latent, num_frames).\"\"\"\n        full_hop = self.config.latent_hop\n        off = 3\n        wav = prompt_audio\n        if wav.ndim == 2:\n            wav = wav[..., None]  # (B, L) -> (B, L, 1)\n        if wav.ndim == 3 and wav.shape[-1] != 1:\n            wav = wav.transpose(0, 2, 1)  # (B, 1, L) -> (B, L, 1)\n\n        L = wav.shape[1]\n        if L % full_hop != 0:\n            pad_amount = full_hop - L % full_hop\n            wav = mx.pad(wav, [(0, 0), (0, pad_amount), (0, 0)])\n        wav = mx.pad(wav, [(0, 0), (0, full_hop * off), (0, 0)])\n\n        latent = self.vae.encode(wav)  # (B, T, latent_dim)\n        if off != 0:\n            latent = latent[:, :-off, :]\n        prompt_dur = latent.shape[1]\n        return latent, prompt_dur\n\n    @staticmethod\n    def _format_duration(seconds: float) -> str:\n        return f\"{int(seconds // 3600):02d}:{int((seconds % 3600) // 60):02d}:{seconds % 60:06.3f}\"\n\n    def _stream_decode(\n        self,\n        pred_latent: mx.array,\n        sr: int,\n        start_time: float,\n        chunk_seconds: float = 2.0,\n        overlap_seconds: float = 0.5,\n    ) -> Generator[GenerationResult, None, None]:\n        \"\"\"Decode latents in overlapping chunks with cosine crossfade.\n\n        Runs the VAE decoder on small overlapping windows of the latent\n        sequence and yields audio as soon as each chunk is ready, giving\n        much lower time-to-first-audio than full-sequence decoding.\n        \"\"\"\n        ratio = self.config.vae_config.downsampling_ratio\n        chunk_frames = max(1, int(chunk_seconds * sr / ratio))\n        overlap_frames = max(0, int(overlap_seconds * sr / ratio))\n        hop_frames = max(1, chunk_frames - overlap_frames)\n        overlap_samples = overlap_frames * ratio\n        # Context frames fed to the VAE decoder on each side so its\n        # convolutions have proper context, then trimmed from the output.\n        context_frames = overlap_frames\n\n        total_frames = pred_latent.shape[1]\n        prev_tail = None\n        chunk_idx = 0\n        cumulative_samples = 0\n\n        start = 0\n        while start < total_frames:\n            end = min(start + chunk_frames, total_frames)\n            is_last = end >= total_frames\n\n            # Extend chunk with context on both sides for decoder conv context\n            left_ctx = min(context_frames, start)\n            right_ctx = min(context_frames, total_frames - end)\n            ctx_start = start - left_ctx\n            ctx_end = end + right_ctx\n\n            # Decode extended chunk, then trim context audio\n            audio_full = self.vae.decode(pred_latent[:, ctx_start:ctx_end, :])\n            audio_full = audio_full.squeeze(-1).squeeze(0)  # (L,)\n            mx.eval(audio_full)\n            left_trim = left_ctx * ratio\n            right_trim = right_ctx * ratio\n            if right_trim > 0:\n                audio_chunk = audio_full[left_trim:-right_trim]\n            else:\n                audio_chunk = audio_full[left_trim:]\n\n            if prev_tail is not None and overlap_samples > 0:\n                # Cosine crossfade with previous chunk's tail\n                ol = min(overlap_samples, prev_tail.shape[0], audio_chunk.shape[0])\n                t = mx.linspace(0, 1, ol)\n                fade_in = 0.5 * (1 - mx.cos(mx.array(math.pi) * t))\n                fade_out = 1.0 - fade_in\n                blended = prev_tail[:ol] * fade_out + audio_chunk[:ol] * fade_in\n\n                if is_last:\n                    output = mx.concatenate([blended, audio_chunk[ol:]])\n                else:\n                    output = mx.concatenate([blended, audio_chunk[ol:-overlap_samples]])\n                    prev_tail = audio_chunk[-overlap_samples:]\n                    mx.eval(prev_tail)\n            else:\n                if is_last or overlap_frames == 0:\n                    output = audio_chunk\n                else:\n                    output = audio_chunk[:-overlap_samples]\n                    prev_tail = audio_chunk[-overlap_samples:]\n                    mx.eval(prev_tail)\n\n            mx.eval(output)\n            num_samples = output.shape[0]\n            cumulative_samples += num_samples\n            processing_time = time_module.time() - start_time\n            audio_dur = cumulative_samples / sr\n\n            yield GenerationResult(\n                audio=output,\n                samples=num_samples,\n                sample_rate=sr,\n                segment_idx=chunk_idx,\n                token_count=0,\n                audio_duration=self._format_duration(audio_dur),\n                real_time_factor=processing_time / max(audio_dur, 1e-6),\n                prompt={\"tokens\": 0, \"tokens-per-sec\": 0},\n                audio_samples={\n                    \"samples\": num_samples,\n                    \"samples-per-sec\": num_samples / max(processing_time, 1e-6),\n                },\n                processing_time_seconds=processing_time,\n                peak_memory_usage=mx.get_peak_memory() / 1e9,\n                is_streaming_chunk=True,\n                is_final_chunk=is_last,\n            )\n\n            chunk_idx += 1\n            start += hop_frames\n\n    def generate(\n        self,\n        text: str,\n        voice: Optional[str] = None,\n        speed: float = 1.0,\n        lang_code: str = \"en\",\n        ref_audio: Optional[mx.array] = None,\n        ref_text: Optional[str] = None,\n        steps: int = 16,\n        cfg_strength: float = 4.0,\n        guidance_method: str = \"cfg\",\n        seed: int = 1024,\n        stream: bool = False,\n        streaming_interval: float = 2.0,\n        chunk_seconds: float = 2.0,\n        overlap_seconds: float = 0.5,\n        **kwargs,\n    ) -> Generator[GenerationResult, None, None]:\n        \"\"\"Generate audio from text.\n\n        Args:\n            text: Text to synthesize.\n            voice: Unused (kept for API compat).\n            speed: Unused (kept for API compat).\n            ref_audio: Optional reference audio array for voice cloning (samples at 24kHz).\n            ref_text: Optional text of the reference audio.\n            steps: Number of ODE Euler steps (default 16).\n            cfg_strength: Classifier-free guidance strength.\n            guidance_method: \"cfg\" or \"apg\".\n            seed: Random seed.\n            stream: Stream audio chunks during VAE decoding.\n            streaming_interval: Alias for chunk_seconds (CLI compat).\n            chunk_seconds: Chunk length in seconds for streaming decode.\n            overlap_seconds: Overlap between chunks for cosine crossfade.\n        \"\"\"\n        import transformers\n\n        start_time = time_module.time()\n        mx.random.seed(seed)\n\n        sr = self.config.sampling_rate\n        full_hop = self.config.latent_hop\n        max_duration = self.config.max_wav_duration\n        repa_layer = self.config.repa_dit_layer\n\n        # Tokenize\n        tokenizer = transformers.AutoTokenizer.from_pretrained(\n            self.config.text_encoder_model\n        )\n\n        text = _normalize_text(text)\n        no_prompt = ref_audio is None\n\n        if not no_prompt and ref_text is not None:\n            prompt_text = _normalize_text(ref_text)\n            full_text = f\"{prompt_text} {text}\"\n        else:\n            full_text = text\n\n        inputs = tokenizer([full_text], padding=\"longest\", return_tensors=\"np\")\n        input_ids = mx.array(inputs[\"input_ids\"])\n        attention_mask = mx.array(inputs[\"attention_mask\"]).astype(mx.float32)\n\n        # Text encoding\n        text_condition = self.encode_text(input_ids, attention_mask)\n        text_condition_len = mx.sum(attention_mask, axis=1)\n\n        batch = text_condition.shape[0]\n\n        # Prompt audio encoding\n        if not no_prompt:\n            prompt_audio_mx = (\n                mx.array(np.array(ref_audio))\n                if not isinstance(ref_audio, mx.array)\n                else ref_audio\n            )\n            prompt_latent, prompt_dur = self.encode_prompt_audio(\n                prompt_audio_mx[None] if prompt_audio_mx.ndim < 3 else prompt_audio_mx\n            )\n        else:\n            prompt_latent = mx.zeros((batch, 0, self.config.latent_dim))\n            prompt_dur = 0\n\n        # Duration estimation\n        prompt_time = prompt_dur * full_hop / sr\n        dur_sec = _approx_duration(text, max_duration=max_duration - prompt_time)\n        if not no_prompt and ref_text is not None:\n            approx_pd = _approx_duration(ref_text, max_duration=max_duration)\n            ratio = float(np.clip(prompt_time / max(approx_pd, 1e-6), 1.0, 1.5))\n            dur_sec = dur_sec * ratio\n\n        duration = int(dur_sec * sr // full_hop)\n        total_duration = min(duration + prompt_dur, int(max_duration * sr // full_hop))\n\n        # Masks\n        mask = mx.arange(total_duration)[None, :] < mx.array([total_duration])\n        text_mask = (\n            mx.arange(text_condition.shape[1])[None, :] < text_condition_len[:, None]\n        )\n\n        # Conditioning\n        neg_text = mx.zeros_like(text_condition)\n\n        if not no_prompt:\n            gen_len = total_duration - prompt_dur\n            latent_cond = mx.pad(prompt_latent, [(0, 0), (0, gen_len), (0, 0)])\n            empty_latent_cond = mx.zeros_like(latent_cond)\n        else:\n            latent_cond = mx.zeros((batch, total_duration, self.config.latent_dim))\n            empty_latent_cond = latent_cond\n\n        # APG buffer\n        apg_buffer = (\n            _MomentumBuffer(momentum=-0.3) if guidance_method == \"apg\" else None\n        )\n\n        # Initial noise\n        y = mx.random.normal((batch, total_duration, self.config.latent_dim))\n        prompt_noise = y[:, :prompt_dur] if prompt_dur > 0 else None\n\n        # Time steps\n        t_steps = np.linspace(0, 1, steps).tolist()\n\n        # Euler ODE integration\n        for i in range(len(t_steps) - 1):\n            t_val = t_steps[i]\n            dt = t_steps[i + 1] - t_val\n            t = mx.array([t_val])\n\n            # Set prompt region\n            if prompt_dur > 0:\n                y_prompt = (\n                    prompt_noise * (1 - t_val) + latent_cond[:, :prompt_dur] * t_val\n                )\n                y = mx.concatenate([y_prompt, y[:, prompt_dur:]], axis=1)\n\n            output = self.transformer(\n                x=y,\n                text=text_condition,\n                text_len=text_condition_len,\n                time=t,\n                mask=mask,\n                cond_mask=text_mask,\n                return_ith_layer=repa_layer,\n                latent_cond=latent_cond,\n            )\n            pred = output[\"last_hidden_state\"]\n\n            if cfg_strength >= 1e-5:\n                # Null forward for guidance\n                y_null = mx.array(y)\n                if prompt_dur > 0:\n                    y_null = mx.concatenate(\n                        [mx.zeros_like(y_null[:, :prompt_dur]), y_null[:, prompt_dur:]],\n                        axis=1,\n                    )\n\n                null_output = self.transformer(\n                    x=y_null,\n                    text=neg_text,\n                    text_len=text_condition_len,\n                    time=t,\n                    mask=mask,\n                    cond_mask=text_mask,\n                    return_ith_layer=repa_layer,\n                    latent_cond=empty_latent_cond,\n                )\n                null_pred = null_output[\"last_hidden_state\"]\n\n                if guidance_method == \"cfg\":\n                    pred = pred + (pred - null_pred) * cfg_strength\n                else:\n                    # APG\n                    x_s = y[:, prompt_dur:]\n                    pred_s = pred[:, prompt_dur:]\n                    null_s = null_pred[:, prompt_dur:]\n                    pred_sample = x_s + (1 - t_val) * pred_s\n                    null_sample = x_s + (1 - t_val) * null_s\n                    out = _apg_forward(\n                        pred_sample,\n                        null_sample,\n                        cfg_strength,\n                        apg_buffer,\n                        eta=0.5,\n                        norm_threshold=0.0,\n                    )\n                    out = (out - x_s) / (1 - t_val)\n                    pred = mx.pad(out, [(0, 0), (prompt_dur, 0), (0, 0)])\n\n            y = y + pred * dt\n            mx.eval(y)\n\n        # Decode\n        pred_latent = y[:, prompt_dur:] if not no_prompt else y\n\n        if stream:\n            # streaming_interval from CLI takes priority over chunk_seconds\n            cs = streaming_interval if streaming_interval != 2.0 else chunk_seconds\n            yield from self._stream_decode(\n                pred_latent, sr, start_time, cs, overlap_seconds\n            )\n        else:\n            waveform = self.vae.decode(pred_latent)  # (B, L, 1)\n            waveform = waveform.squeeze(-1)  # (B, L)\n\n            mx.eval(waveform)\n            processing_time = time_module.time() - start_time\n\n            audio = waveform[0]\n            num_samples = audio.shape[0]\n            audio_duration = num_samples / sr\n\n            yield GenerationResult(\n                audio=audio,\n                samples=num_samples,\n                sample_rate=sr,\n                segment_idx=0,\n                token_count=0,\n                audio_duration=self._format_duration(audio_duration),\n                real_time_factor=processing_time / max(audio_duration, 1e-6),\n                prompt={\"tokens\": 0, \"tokens-per-sec\": 0},\n                audio_samples={\n                    \"samples\": num_samples,\n                    \"samples-per-sec\": num_samples / max(processing_time, 1e-6),\n                },\n                processing_time_seconds=processing_time,\n                peak_memory_usage=mx.get_peak_memory() / 1e9,\n            )\n\n    def sanitize(self, weights: dict) -> dict:\n        \"\"\"Convert PyTorch AudioDiT weights to MLX format.\"\"\"\n        sanitized = {}\n\n        # Collect weight_v/weight_g pairs for reconstruction\n        weight_v_keys = [k for k in weights if k.endswith(\".weight_v\")]\n        processed = set()\n\n        for wv_key in weight_v_keys:\n            base = wv_key[: -len(\".weight_v\")]\n            wg_key = base + \".weight_g\"\n            if wg_key not in weights:\n                continue\n\n            weight_v = weights[wv_key]\n            weight_g = weights[wg_key]\n\n            # Determine if this is ConvTranspose1d\n            # ConvTranspose1d is at decoder.layers.{block}.layers.1\n            is_conv_transpose = bool(\n                re.search(r\"vae\\.decoder\\.layers\\.\\d+\\.layers\\.1\\.weight\", wv_key)\n            )\n\n            if weight_v.ndim == 3:\n                # Reconstruct weight from weight_norm: w = g * v / ||v||\n                if is_conv_transpose:\n                    # PyTorch ConvTranspose1d: (in_ch, out_ch, ksize)\n                    # norm over dims 1,2 (out_ch and ksize)\n                    norm_axes = (1, 2)\n                    norm = mx.sqrt(\n                        mx.sum(weight_v * weight_v, axis=norm_axes, keepdims=True)\n                        + 1e-12\n                    )\n                    weight = weight_g * weight_v / norm\n                    # Transpose: (in_ch, out_ch, ksize) -> (out_ch, ksize, in_ch)\n                    weight = weight.transpose(1, 2, 0)\n                else:\n                    # PyTorch Conv1d: (out_ch, in_ch, ksize)\n                    # norm over dims 1,2 (in_ch and ksize)\n                    norm_axes = (1, 2)\n                    norm = mx.sqrt(\n                        mx.sum(weight_v * weight_v, axis=norm_axes, keepdims=True)\n                        + 1e-12\n                    )\n                    weight = weight_g * weight_v / norm\n                    # Transpose: (out_ch, in_ch, ksize) -> (out_ch, ksize, in_ch)\n                    weight = weight.transpose(0, 2, 1)\n\n                sanitized[base + \".weight\"] = weight\n            else:\n                sanitized[base + \".weight\"] = weight_v\n\n            processed.add(wv_key)\n            processed.add(wg_key)\n\n        # Process remaining weights\n        for key, value in weights.items():\n            if key in processed:\n                continue\n\n            new_key = key\n\n            # --- Text encoder remapping ---\n            if new_key.startswith(\"text_encoder.encoder.embed_tokens.\"):\n                new_key = new_key.replace(\n                    \"text_encoder.encoder.embed_tokens.\", \"text_encoder.shared.\"\n                )\n\n            # T5 block layer remapping:\n            # HF: block.{i}.layer.0.SelfAttention -> block.{i}.SelfAttention\n            # HF: block.{i}.layer.0.layer_norm -> block.{i}.layer_norm_sa\n            # HF: block.{i}.layer.1.DenseReluDense -> block.{i}.DenseReluDense\n            # HF: block.{i}.layer.1.layer_norm -> block.{i}.layer_norm_ff\n            new_key = re.sub(\n                r\"text_encoder\\.encoder\\.block\\.(\\d+)\\.layer\\.0\\.SelfAttention\\.\",\n                r\"text_encoder.block.\\1.SelfAttention.\",\n                new_key,\n            )\n            new_key = re.sub(\n                r\"text_encoder\\.encoder\\.block\\.(\\d+)\\.layer\\.0\\.layer_norm\\.\",\n                r\"text_encoder.block.\\1.layer_norm_sa.\",\n                new_key,\n            )\n            new_key = re.sub(\n                r\"text_encoder\\.encoder\\.block\\.(\\d+)\\.layer\\.1\\.DenseReluDense\\.\",\n                r\"text_encoder.block.\\1.DenseReluDense.\",\n                new_key,\n            )\n            new_key = re.sub(\n                r\"text_encoder\\.encoder\\.block\\.(\\d+)\\.layer\\.1\\.layer_norm\\.\",\n                r\"text_encoder.block.\\1.layer_norm_ff.\",\n                new_key,\n            )\n            new_key = new_key.replace(\n                \"text_encoder.encoder.final_layer_norm.\",\n                \"text_encoder.final_layer_norm.\",\n            )\n\n            # --- Transformer Sequential index remapping ---\n            # Embedder: proj.0 -> proj.0, proj.2 -> proj.1\n            new_key = re.sub(r\"\\.proj\\.2\\.\", \".proj.1.\", new_key)\n            # TimestepEmbedding: time_mlp.0 -> time_mlp.0, time_mlp.2 -> time_mlp.1\n            new_key = re.sub(r\"\\.time_mlp\\.2\\.\", \".time_mlp.1.\", new_key)\n            # AdaLNMLP: mlp.1 -> mlp.0\n            new_key = re.sub(r\"\\.mlp\\.1\\.\", \".mlp.0.\", new_key)\n            # Attention to_out: to_out.0 -> to_out\n            new_key = re.sub(r\"\\.to_out\\.0\\.\", \".to_out.\", new_key)\n            # FeedForward: ff.3 -> ff.1\n            new_key = re.sub(r\"\\.ff\\.3\\.\", \".ff.1.\", new_key)\n\n            # --- ConvNeXtV2 depthwise conv weight ---\n            # dwconv.weight: (dim, 1, ksize) -> (dim, ksize, 1) for channels-last\n            if \"dwconv.weight\" in new_key and value.ndim == 3:\n                value = value.transpose(0, 2, 1)\n                new_key = new_key.replace(\".dwconv.weight\", \".dwconv_weight\")\n            elif \"dwconv.bias\" in new_key:\n                new_key = new_key.replace(\".dwconv.bias\", \".dwconv_bias\")\n\n            # --- Handle standalone Conv1d weights (non weight-normed) ---\n            # These might be in the ConvNeXtV2 dwconv or other places\n\n            sanitized[new_key] = value\n\n        return sanitized\n\n\ndef _layer_norm(x: mx.array, dim: int, eps: float = 1e-6) -> mx.array:\n    \"\"\"Layer normalization.\"\"\"\n    mean = mx.mean(x, axis=-1, keepdims=True)\n    var = mx.var(x, axis=-1, keepdims=True)\n    return (x - mean) * mx.rsqrt(var + eps)\n"}}}}