total 1716
drwxrwxr-x  2 ubuntu ubuntu   4096 Mar 24 10:25 .
drwxrwxr-x 11 ubuntu ubuntu   4096 Mar 24 10:24 ..
-rw-rw-r--  1 ubuntu ubuntu 456270 Mar 24 10:25 predict_best.wav
-rw-rw-r--  1 ubuntu ubuntu 456270 Mar 24 10:24 recon_best.wav
-rw-rw-r--  1 ubuntu ubuntu 456270 Mar 24 10:25 recon_ep200.wav
-rw-rw-r--  1 ubuntu ubuntu 307790 Mar 24 10:25 synth_best_t05.wav
-rw-rw-r--  1 ubuntu ubuntu  59982 Mar 24 10:25 synth_best.wav
---
-rw-rw-r-- 1 ubuntu ubuntu 613966 Mar 24 07:01 /home/ubuntu/lewm-tts/sample_ar_prompted_v3b.wav
-rw-rw-r-- 1 ubuntu ubuntu 613966 Mar 24 06:50 /home/ubuntu/lewm-tts/sample_ar_prompted.wav
-rw-rw-r-- 1 ubuntu ubuntu 409166 Mar 24 08:12 /home/ubuntu/lewm-tts/sample_ar_v3c.wav
-rw-rw-r-- 1 ubuntu ubuntu 456782 Mar 24 06:46 /home/ubuntu/lewm-tts/sample_gt_0.wav
-rw-rw-r-- 1 ubuntu ubuntu 359502 Mar 24 06:46 /home/ubuntu/lewm-tts/sample_gt_100.wav
-rw-rw-r-- 1 ubuntu ubuntu 384078 Mar 24 06:46 /home/ubuntu/lewm-tts/sample_gt_500.wav
-rw-rw-r-- 1 ubuntu ubuntu 456782 Mar 24 06:34 /home/ubuntu/lewm-tts/sample_gt_v3b.wav
-rw-rw-r-- 1 ubuntu ubuntu 456270 Mar 24 06:17 /home/ubuntu/lewm-tts/sample_recon_ep43_0.wav
-rw-rw-r-- 1 ubuntu ubuntu 357966 Mar 24 06:17 /home/ubuntu/lewm-tts/sample_recon_ep43_100.wav
-rw-rw-r-- 1 ubuntu ubuntu 382542 Mar 24 06:17 /home/ubuntu/lewm-tts/sample_recon_ep43_500.wav
-rw-rw-r-- 1 ubuntu ubuntu 456270 Mar 24 06:07 /home/ubuntu/lewm-tts/sample_reconstruct_v3.wav
-rw-rw-r-- 1 ubuntu ubuntu 456270 Mar 24 07:01 /home/ubuntu/lewm-tts/sample_recon_v3b_0.wav
-rw-rw-r-- 1 ubuntu ubuntu 357966 Mar 24 07:01 /home/ubuntu/lewm-tts/sample_recon_v3b_100.wav
-rw-rw-r-- 1 ubuntu ubuntu 382542 Mar 24 07:01 /home/ubuntu/lewm-tts/sample_recon_v3b_500.wav
-rw-rw-r-- 1 ubuntu ubuntu 456270 Mar 24 06:34 /home/ubuntu/lewm-tts/sample_recon_v3b.wav
-rw-rw-r-- 1 ubuntu ubuntu 456270 Mar 24 08:09 /home/ubuntu/lewm-tts/sample_recon_v3c_0.wav
-rw-rw-r-- 1 ubuntu ubuntu 358990 Mar 24 08:09 /home/ubuntu/lewm-tts/sample_recon_v3c_100.wav
-rw-rw-r-- 1 ubuntu ubuntu 383566 Mar 24 08:09 /home/ubuntu/lewm-tts/sample_recon_v3c_500.wav
-rw-rw-r-- 1 ubuntu ubuntu 411214 Mar 24 06:17 /home/ubuntu/lewm-tts/sample_synth_ep43.wav
-rw-rw-r-- 1 ubuntu ubuntu 411214 Mar 24 06:46 /home/ubuntu/lewm-tts/sample_synth_v3b.wav
-rw-rw-r-- 1 ubuntu ubuntu 161358 Mar 24 06:07 /home/ubuntu/lewm-tts/sample_synth_v3.wav
-rw-rw-r-- 1 ubuntu ubuntu 456270 Mar 24 06:50 /home/ubuntu/lewm-tts/sample_tf_0.wav
-rw-rw-r-- 1 ubuntu ubuntu 357966 Mar 24 06:50 /home/ubuntu/lewm-tts/sample_tf_100.wav
-rw-rw-r-- 1 ubuntu ubuntu 382542 Mar 24 06:50 /home/ubuntu/lewm-tts/sample_tf_500.wav
-rw-rw-r-- 1 ubuntu ubuntu 456270 Mar 24 07:01 /home/ubuntu/lewm-tts/sample_tf_v3b_0.wav
-rw-rw-r-- 1 ubuntu ubuntu 357966 Mar 24 07:01 /home/ubuntu/lewm-tts/sample_tf_v3b_100.wav
-rw-rw-r-- 1 ubuntu ubuntu 456270 Mar 24 08:09 /home/ubuntu/lewm-tts/sample_tf_v3c_0.wav


Total samples: 7009
Duration: min=3.0s, max=14.4s, mean=9.2s, total=17.88h
Mel frames: min=282, max=1350, mean=861
Text tokens: min=24, max=573, mean=218
  id=0: 9.52s, 893 frames, text=मेरे प्यारे देशवासियों, नमस्कार। मन की बात में आपका स्वागत है, अभिन...
  id=1: 8.8s, 826 frames, text=ंदन है। मन की बात देश और देशवासियों की उपलब्धियों को सामने लान...
  id=2: 9.14s, 857 frames, text=े का एक मजबूत platform है। देश ने ऐसी ही उपलब्धि अभी दिल्ली में ह...


Now I have a complete picture. Here's my thorough analysis:

---

## Thorough Analysis: Why the JEPA TTS Model Isn't Getting Good Results

I've reviewed all three versions (v3, v3b, v3c), training logs, model architecture, dataset, and inference code. Here are the **fundamental problems** and what needs to be fixed.

---

### Problem 1: KL Divergence Explodes — Latent Space is Uncontrollable

| Version | KL Weight | Final KL | Final Pred Loss |
|---------|-----------|----------|-----------------|
| v3      | 0.1       | ~1.18    | 0.057           |
| v3b     | 0.01      | ~2.44    | 0.019           |
| v3c     | 0.01      | ~2.45    | 0.017           |

In v3b/v3c, `kl_weight=0.01` is **10x too low**. The KL divergence climbed to 2.45, meaning the encoder learns a latent space that is very far from N(0,I). At inference time, when you start AR generation from `torch.randn() * 5.0`, you're starting from a distribution that has **nothing to do** with what the predictor was trained on. The v3 version with `kl_weight=0.1` kept KL at ~1.18, which was better.

### Problem 2: Prediction Loss is Trivially Low — The Predictor Learns Nothing

The most damning signal: **prediction loss = 0.017** in v3c (0.86% of total loss). This means:

- Consecutive audio embeddings `z[t]` and `z[t+1]` are nearly identical (mel frames with stride-2 overlap heavily)
- The predictor learns an approximate **identity mapping** — just copy the input
- The predictor has learned nothing meaningful about speech dynamics
- At inference time, AR generation produces repetitive/degenerate output because "predict next" ≈ "copy current"

This is the **classic JEPA failure mode**: without proper constraints, the representation becomes trivially predictable.

### Problem 3: Loss Balance is Completely Wrong

In v3c at epoch 200:
- `pred = 0.017` (0.86% of total)
- `kl = 0.01 × 2.45 = 0.025` (1.2% of total)
- `recon = 2.0 × 0.50 = 1.00` (50.7% of total)
- `spectral = 1.5 × 0.62 = 0.93` (47.2% of total)

**97.9% of the loss is reconstruction/spectral**. The model is optimizing almost entirely for encoder→decoder quality. The JEPA predictor — the part that matters for synthesis — gets less than 1% of the training signal.

### Problem 4: Training vs Inference Mismatch (Exposure Bias)

During training: predictor sees **real encoder outputs** `z[t]` as input.
During inference: predictor sees **its own previous predictions** as input.

With no scheduled sampling or noise injection during training, errors compound immediately during AR inference. Even tiny prediction errors snowball.

### Problem 5: Random Start Embedding in Inference

```python
start_emb = torch.randn(1, 1, d) * self._start_emb_scale  # scale=5.0
```

The model never trains with random start embeddings. The predictor has never seen anything like `randn * 5.0` as input. This guarantees the first AR step produces garbage, and it cascades from there.

### Problem 6: v3c's Downsample Factor = 2 Made Things Worse

v3c reduced downsample from 4→2. This means:
- **More temporal redundancy** between consecutive embeddings (they overlap more)
- Prediction is even more trivially "copy input"
- Model dropped from 24M → 21.1M params (less capacity)
- Yet we train the same 200 epochs

### Problem 7: No Proper EMA Target Encoder (Missing Core JEPA Component)

The original JEPA/I-JEPA paper uses an **EMA target encoder** — the predictions target a slowly-moving average of the encoder, which prevents representation collapse. Here, both input and target come from the same encoder, which means the encoder can cheat by making all embeddings identical (which it nearly does — pred loss 0.017).

---

## What Needs to Be Fixed

Here are the concrete fixes, in priority order:

1. **Add EMA target encoder** — The encoder that produces prediction targets should be an exponential moving average of the main encoder (momentum ~0.996). This is the core of JEPA and it's missing.

2. **Increase prediction loss weight significantly** — The predictor should be the primary loss (weight 5-10x), not 1% of total.

3. **Increase KL weight or use target-KL** — Use `kl_weight=0.1` minimum, or better, use "free bits" (target KL of ~2.0) to prevent both collapse and explosion.

4. **Increase downsample factor back to 4 or even 8** — More temporal compression means consecutive embeddings are more different, giving the predictor a meaningful prediction task.

5. **Add scheduled sampling / teacher forcing noise** — During training, randomly replace some input embeddings with predicted ones (curriculum: start at 0%, increase to 50%). This bridges the training/inference gap.

6. **Learn a proper start embedding** — Replace `randn * 5.0` with a learnable `nn.Parameter` for the start-of-sequence token.

7. **Predict multiple steps ahead** — Instead of just next-step prediction, predict 2-4 steps ahead. This forces the predictor to learn actual speech dynamics, not just copy.

---

Want me to implement these fixes? I'm confident that with an EMA target encoder, proper loss balancing, and scheduled sampling, this JEPA architecture can produce genuinely good results. The core architecture is sound — it's the training setup that's undermining it.


yeah do .. after dng again i wnat u to check in a researcher pov


The user wants me to implement the fixes and approach this from a researcher's perspective. Let me plan the implementation carefully, then implement the changes.

Let me think about what changes are needed:

1. **EMA target encoder** - Add an EMA copy of the AudioEncoder that produces targets for the predictor
2. **Increase prediction loss weight** - Make it the primary loss
3. **Better KL handling** - Use free-bits or higher kl_weight
4. **Increase downsample factor** - Back to 4 or higher
5. **Scheduled sampling** - Mix predicted embeddings during training
6. **Learnable start embedding** - Instead of random noise

7. **Multi-step prediction** - Predict further ahead

Now I'm reconsidering the core JEPA architecture more carefully. The key insight from I-JEPA is that the target encoder—an EMA copy of the context encoder—provides stable targets that prevent collapse. For audio JEPA, I need a context encoder that processes the mel spectrogram and a target encoder that slowly tracks it, with a predictor learning to forecast target representations from context.

To implement this properly, I'll need to add an EMA target encoder to the model, introduce a learnable start embedding, update the forward pass to use EMA targets, and add an EMA update mechanism. On the training side, I should rebalance the loss weights, incorporate scheduled sampling, call the EMA update after each step, and restore the downsample factor to 4. For inference, I'll switch to using the learnable start embedding instead of random initialization.

Implementing these changes now... about the implementation more carefully:

**EMA Target Encoder:**
- Copy of AudioEncoder with same architecture
- Updated with EMA after each optimizer step
- `ema_decay` typically 0.996-0.999
- Targets come from this encoder (no gradient)

**Loss rebalancing:**
- pred_weight: 10.0 (primary loss)
- kl: use free-bits with target ~1.0-2.0, weight 0.05
- recon_weight: 1.0
- spectral_weight: 0.5

**Scheduled sampling:**
- During training, with probability `p_scheduled`, replace input embedding at position t with the predictor's prediction from t-1
- Curriculum: p starts at 0, linearly increases to 0.5 over training

**Start embedding:**
- Learnable parameter initialized with small random values

I'm reconsidering the multi-step prediction approach—it adds unnecessary complexity. The core improvements that matter most are the EMA target encoder, rebalancing the losses with prediction as the primary objective, scheduled sampling with curriculum learning, the learnable start embedding, and a downsample factor of 4. For the architecture itself, the predictor takes all embeddings except the last as input and predicts the next embeddings using causal self-attention plus cross-attention to the text conditioning.

With the EMA setup, the online encoder produces embeddings for the prediction task while the EMA encoder provides the target embeddings with stopped gradients. The reconstruction loss still uses the online encoder's output through the mel decoder. Let me trace through the current forward pass to make sure I understand the flow correctly before implementing these changes. (L1) + spectral loss

With EMA:
1. Encode text → text_emb
2. Encode audio with online encoder → z_online, mu, logvar
3. Encode audio with EMA encoder → z_ema (no grad, use mu directly)
4. Predict: input z_online[:-1], target z_ema[1:] → prediction loss
5. KL loss on mu, logvar
6. Reconstruct mel from z_online → recon loss + spectral loss

Actually, I need to think about this differently. In the original I-JEPA:
- The context encoder processes visible context
- The predictor predicts target representation from context
- The target encoder (EMA) provides the targets
- Loss = MSE(predictor_output, EMA_target)

The key property is that the EMA encoder provides slowly-changing targets, which prevents the online encoder from collapsing to trivial solutions. But in our case, the real issue isn't representation collapse — the prediction loss is already low at 0.017. The problem is that consecutive embeddings are nearly identical because mel frames overlap significantly. With a downsample factor of 2, consecutive embeddings are only about 21ms apart, so the predictor can easily learn a trivial mapping. The EMA target encoder would help by introducing a slight difference between the online and target representations, forcing the predictor to learn something more meaningful.

The better solutions are to increase the temporal distance between embeddings, weight the prediction loss more heavily, or use scheduled sampling to bridge the training-inference gap. I could also consider cosine similarity loss instead of MSE to handle the scale issues that come with high-dimensional vectors, but I'll keep it simple for now.

Let me start implementing these changes. I'll modify the model and training code to add the EMA encoder and adjust the loss weighting.

Modifying model.py and train.py... For inference, I'll use the learned start_emb parameter instead of random initialization. But I need to think through scheduled sampling more carefully—implementing it with a sequence predictor that uses causal attention is complex since we can't easily swap individual tokens mid-sequence. A simpler alternative would be to inject noise into the input embeddings during training, which achieves a similar regularization effect without the architectural complications.

I could add Gaussian noise to the input embeddings before the predictor, annealing the noise scale from zero during training to simulate the imperfect predictions at inference time. Or even simpler—use input dropout to randomly zero out embeddings and force the predictor to rely more on text and broader audio context rather than just copying the previous frame.

I'll combine these approaches: input noise annealing, an EMA target encoder, higher prediction weight, and the learnable start embedding. Let me start implementing...

Implementing the training loop modifications...

I should also consider free-bits KL to prevent KL collapse—clamping the per-dimension KL loss so it only incurs a penalty above a threshold like 1.0-2.0 bits. Now I'll update model.py, train.py, and inference.py with the minimal necessary changes.

For model.py, I need to add the EMA audio encoder as a deepcopy without gradients, introduce learnable start embeddings and configuration parameters for EMA decay, prediction weight, and input noise, implement an update_ema method, and modify the forward pass to use EMA targets for the prediction loss, apply input noise to the encoder, weight the prediction loss appropriately, and apply free-bits KL. For train.py, I'll add the new hyperparameters to the config and adjust the default weights for KL, reconstruction, and spectral losses.

Now I'm planning the specific edits to make—I'll add the import statement at the top, then insert the new parameters and EMA encoder initialization in the __init__ method, create the update_ema method, and finally modify the forward pass accordingly. For inference.py, I just need to swap out the random initialization for the learned start embeddings. after _add_speaker
4. Modify forward() significantly

The forward() method is complex. Let me think about what the new forward looks like:

I'm setting up the forward pass to handle the online and EMA encoders, then feeding shifted embeddings into the predictor with noise injection during training. The prediction loss becomes the primary objective, with additional regularization terms like KL divergence to follow.

Now I'm computing the KL loss with free bits to prevent posterior collapse, then reconstructing the mel-spectrogram from the online embeddings and calculating reconstruction and spectral losses. Finally, I'm combining all weighted losses into the total loss.

I need to plan out the edits carefully across multiple files. The forward method is quite long, so I'll break this into separate edits for model.py, then train.py, then inference.py, starting with adding the necessary imports and modifying the initialization to include the new EMA components.

Now I'm considering the loss function for the prediction task—MSE is standard for JEPA-style approaches, though I should think about whether normalizing the target embeddings would help, similar to what BYOL and DINO do.

Actually, let me not overcomplicate this. The core improvements I need to focus on are the EMA target encoder, rebalancing the prediction loss weight to 10, adding free-bits KL, injecting input noise, making the start embedding learnable, and using a downsample factor of 4. These changes will have the most impact.

Now I'm ready to start implementing. I'll work through model.py first, being careful with the line numbers as I add the import for copy, modify the initialization, add the EMA update method, replace the forward pass, and update the build_model configuration.

Actually, I realize the forward method changes are extensive enough that I should replace whole sections rather than trying to patch individual lines. Let me start with the imports and initialization changes, then handle the forward pass as a complete replacement.

Adding the copy import and setting up the EMA encoder with its decay parameter...

Setting up additional components... I'm mapping out the exact line numbers in the model file to understand the structure—the audio encoder is around line 486, mel decoder at 499, and the forward method spans lines 522-614. My plan is to make five targeted edits: first adding the copy import, then modifying the initialization to include EMA components and the new parameters, adding an update_ema method, replacing the forward logic, and finally updating build_model defaults. I'm mapping out the exact insertion points in the `__init__` method—after the reconstruction weight, I need to add the EMA audio encoder, learnable start embedding, and several hyperparameters like prediction weight and free bits. Then after storing the config, I'll add the `update_ema` method. I should also import `copy` at the top of the file to support the EMA functionality. I'm realizing that doing sequential edits will shift line numbers after each change, so I need to be strategic about the order. Let me consolidate this into fewer, larger edits to avoid recalculating offsets—I'll add the import at the top, then make one substantial edit covering the __init__ modifications through the forward method replacement, and handle the build_model update separately. I'm working through how to apply multiple edits to a Python file, but the line number approach is getting complicated since each edit shifts subsequent line numbers. Let me use the Edit tool with old_string/new_string matching instead, which avoids this dependency—I'll start by checking the tool's exact schema first, then proceed through the edits sequentially: adding the import, expanding the initialization section with EMA logic, inserting the update function, and modifying forward and build_model.


Let me get the tools I need for implementing the fixes.


ToolSearch

Question: select:Edit,Write,TaskCreate,TaskUpdate


Question: select:Edit,Write,TaskCreate,TaskUpdate


Let me create tasks to track the implementation, then start implementing.


Let me create a task plan and start implementing.


TaskCreate

Add EMA target encoder to model.py

Add EMA (exponential moving average) copy of AudioEncoder as target encoder for stable prediction targets. Add update_ema() method. This is the core missing JEPA component.


Rebalance losses and add free-bits KL

Add pred_weight (10.0) to make prediction the primary loss. Replace standard KL with free-bits KL (threshold ~1.0). Add input noise injection for train/inference gap bridging. Add learnable start embedding.


Update train.py with new training loop

Update default hyperparams (downsample_factor=4, pred_weight=10, kl_weight=0.05, recon_weight=1.0, spectral_weight=0.5). Add EMA update after each optimizer step. Add input noise annealing schedule.


Update inference.py to use learnable start embedding

Replace random start_emb with model's learned start_emb parameter. Update EMA encoder handling for inference.


Research review of all changes

Review all changes from a researcher POV: verify loss math, check for bugs, ensure EMA correctness, validate training dynamics will improve, check for potential issues.


Task #1 created successfully: Add EMA target encoder to model.py

Add EMA target encoder to model.py


Task #2 created successfully: Rebalance losses and add free-bits KL

Rebalance losses and add free-bits KL


Task #3 created successfully: Update train.py with new training loop

Update train.py with new training loop


Task #4 created successfully: Research review of all changes

Research review of all changes


Task #5 created successfully: Update inference.py to use learnable start embedding


Updated task #1 status


Updated task #2 status


I'll implement tasks 1 and 2 together since they're both in model.py. Let me make the edits.


Edit

import math
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchaudio

import math
import copy
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchaudio


The file /home/ubuntu/lewm-tts/model.py has been updated successfully.

import math
import copy
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchaudio

"""
LeWM TTS Model — JEPA-based Text-to-Speech

Architecture:
  - TextEncoder: Character-level transformer for Hindi/Devanagari text
  - AudioEncoder: 1D CNN + Transformer that encodes mel spectrograms to embeddings
  - JEPAPredictor: Transformer that predicts next audio embedding given text + audio context
  - Two losses: next-embedding prediction (MSE) + Gaussian regularizer (KL to N(0,I))
"""

import math
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchaudio

# ─── Text Encoder ───────────────────────────────────────────────────────────

class TextEncoder(nn.Module):
    """Character-level transformer encoder for Hindi text."""

    def __init__(self, vocab_size=256, d_model=256, nhead=4, num_layers=4, dropout=0.1):
        super().__init__()
        self.d_model = d_model

        # Character embedding (covers all Unicode Hindi chars via byte-level encoding)
        self.char_embed = nn.Embedding(vocab_size, d_model)
        self.pos_embed = SinusoidalPositionalEncoding(d_model, max_len=2048)
        self.dropout = nn.Dropout(dropout)

        encoder_layer = nn.TransformerEncoderLayer(
            d_model=d_model, nhead=nhead,
            dim_feedforward=d_model * 4,
            dropout=dropout, batch_first=True,
            activation="gelu",
        )
        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
        self.proj = nn.Linear(d_model, d_model)

    def forward(self, text_tokens, text_mask=None):
        """
        Args:
            text_tokens: [B, T_text] long tensor of byte-level token IDs
            text_mask: [B, T_text] bool, True = padding
        Returns:
            text_emb: [B, T_text, d_model]
        """
        x = self.char_embed(text_tokens) * math.sqrt(self.d_model)
        x = self.pos_embed(x)
        x = self.dropout(x)
        x = self.transformer(x, src_key_padding_mask=text_mask)
        x = self.proj(x)
        return x

# ─── Audio Encoder ──────────────────────────────────────────────────────────

class AudioEncoder(nn.Module):
    """
    Encodes mel spectrograms into a sequence of audio embeddings.
    Uses 1D CNN for downsampling + Transformer for contextualization.
    """

    def __init__(self, n_mels=100, d_model=256, nhead=4, num_layers=4,
                 downsample_factor=4, dropout=0.1):
        super().__init__()
        self.d_model = d_model
        self.downsample_factor = downsample_factor

        # CNN downsampling: reduces temporal resolution by downsample_factor
        if downsample_factor == 2:
            self.conv_pre = nn.Sequential(
                nn.Conv1d(n_mels, d_model, kernel_size=7, padding=3),
                nn.GELU(),
                nn.Conv1d(d_model, d_model, kernel_size=4, stride=2, padding=1),
                nn.GELU(),
            )  # Total stride = 2
        else:
            self.conv_pre = nn.Sequential(
                nn.Conv1d(n_mels, d_model, kernel_size=7, padding=3),
                nn.GELU(),
                nn.Conv1d(d_model, d_model, kernel_size=4, stride=2, padding=1),
                nn.GELU(),
                nn.Conv1d(d_model, d_model, kernel_size=4, stride=2, padding=1),
                nn.GELU(),
            )  # Total stride = 4

        self.pos_embed = SinusoidalPositionalEncoding(d_model, max_len=4096)
        self.dropout = nn.Dropout(dropout)

        encoder_layer = nn.TransformerEncoderLayer(
            d_model=d_model, nhead=nhead,
            dim_feedforward=d_model * 4,
            dropout=dropout, batch_first=True,
            activation="gelu",
        )
        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)

        # Project to embedding with mu and logvar for Gaussian regularizer
        self.proj_mu = nn.Linear(d_model, d_model)
        self.proj_logvar = nn.Linear(d_model, d_model)

    def forward(self, mel, mel_mask=None):
        """
        Args:
            mel: [B, n_mels, T_mel] — log mel spectrogram
            mel_mask: [B, T_mel] bool, True = padding (before downsampling)
        Returns:
            z: [B, T_down, d_model] — sampled latent embeddings
            mu: [B, T_down, d_model]
            logvar: [B, T_down, d_model]
        """
        # CNN encoding: [B, n_mels, T] -> [B, d_model, T//4]
        x = self.conv_pre(mel)
        x = x.transpose(1, 2)  # [B, T_down, d_model]

        # Downsample mask if provided
        if mel_mask is not None:
            # Approximate downsampled mask
            T_down = x.shape[1]
            mel_mask = mel_mask[:, ::self.downsample_factor][:, :T_down]

        x = self.pos_embed(x)
        x = self.dropout(x)
        x = self.transformer(x, src_key_padding_mask=mel_mask)

        mu = self.proj_mu(x)
        logvar = self.proj_logvar(x)

        # Reparameterization trick
        if self.training:
            std = torch.exp(0.5 * logvar)
            eps = torch.randn_like(std)
            z = mu + eps * std
        else:
            z = mu

        return z, mu, logvar

# ─── JEPA Predictor ─────────────────────────────────────────────────────────

class JEPAPredictor(nn.Module):
    """
    Predicts next audio embedding given:
    - Text context (from TextEncoder)
    - Previous audio embeddings (from AudioEncoder)

    Uses cross-attention to condition on text, and causal self-attention
    over the audio embedding sequence.
    """

    def __init__(self, d_model=256, nhead=4, num_layers=6, dropout=0.1):
        super().__init__()
        self.d_model = d_model

        self.audio_input_proj = nn.Linear(d_model, d_model)
        self.pos_embed = SinusoidalPositionalEncoding(d_model, max_len=4096)
        self.dropout = nn.Dropout(dropout)

        # Transformer decoder with causal self-attention + cross-attention to text
        decoder_layer = nn.TransformerDecoderLayer(
            d_model=d_model, nhead=nhead,
            dim_feedforward=d_model * 4,
            dropout=dropout, batch_first=True,
            activation="gelu",
        )
        self.transformer = nn.TransformerDecoder(decoder_layer, num_layers=num_layers)

        # Output projection to predict next embedding
        self.output_proj = nn.Sequential(
            nn.Linear(d_model, d_model),
            nn.GELU(),
            nn.Linear(d_model, d_model),
        )

    def forward(self, audio_emb, text_emb, audio_mask=None, text_mask=None):
        """
        Args:
            audio_emb: [B, T_audio, d_model] — audio embeddings (from encoder)
            text_emb: [B, T_text, d_model] — text embeddings
            audio_mask: [B, T_audio] bool, True = padding
            text_mask: [B, T_text] bool, True = padding
        Returns:
            predicted: [B, T_audio, d_model] — predicted next embeddings
        """
        x = self.audio_input_proj(audio_emb)
        x = self.pos_embed(x)
        x = self.dropout(x)

        # Causal mask: each position can only attend to previous positions
        T = x.shape[1]
        causal_mask = torch.triu(
            torch.ones(T, T, device=x.device, dtype=torch.bool), diagonal=1
        )

        predicted = self.transformer(
            x, text_emb,
            tgt_mask=causal_mask,
            tgt_key_padding_mask=audio_mask,
            memory_key_padding_mask=text_mask,
        )

        predicted = self.output_proj(predicted)
        return predicted

    def _cached_mha(self, mha, q_input, kv_input, cache_k, cache_v):
        """Run multi-head attention with KV cache.
        q_input: [B, 1, d], kv_input: [B, 1, d] (new token for K,V)
        cache_k, cache_v: [B, T_prev, d] projected keys/values
        Returns: output [B, 1, d], updated cache_k, updated cache_v
        """
        d = mha.embed_dim
        nhead = mha.num_heads
        head_dim = d // nhead
        B = q_input.shape[0]

        # Project Q from q_input, K/V from kv_input
        # in_proj_weight is [3d, d] = [Wq; Wk; Wv]
        Wq, Wk, Wv = mha.in_proj_weight.chunk(3, dim=0)
        bq, bk, bv = mha.in_proj_bias.chunk(3, dim=0)

        q = F.linear(q_input, Wq, bq)        # [B, 1, d]
        k_new = F.linear(kv_input, Wk, bk)    # [B, 1, d]
        v_new = F.linear(kv_input, Wv, bv)    # [B, 1, d]

        # Append to cache
        if cache_k is not None and cache_k.shape[1] > 0:
            k = torch.cat([cache_k, k_new], dim=1)
            v = torch.cat([cache_v, v_new], dim=1)
        else:
            k = k_new
            v = v_new

        # Multi-head attention
        q = q.view(B, 1, nhead, head_dim).transpose(1, 2)      # [B, nhead, 1, head_dim]
        k_mh = k.view(B, -1, nhead, head_dim).transpose(1, 2)  # [B, nhead, T, head_dim]
        v_mh = v.view(B, -1, nhead, head_dim).transpose(1, 2)

        attn = torch.matmul(q, k_mh.transpose(-2, -1)) / (head_dim ** 0.5)
        attn = F.softmax(attn, dim=-1)
        out = torch.matmul(attn, v_mh)  # [B, nhead, 1, head_dim]

        out = out.transpose(1, 2).contiguous().view(B, 1, d)
        out = F.linear(out, mha.out_proj.weight, mha.out_proj.bias)

        return out, k, v

    def _cached_cross_attn(self, mha, q_input, memory, cache_k, cache_v, memory_mask=None):
        """Cross-attention with cached memory K,V (computed once).
        memory_mask: [B, T_text] bool, True = padding (applied as -inf)."""
        d = mha.embed_dim
        nhead = mha.num_heads
        head_dim = d // nhead
        B = q_input.shape[0]

        Wq, Wk, Wv = mha.in_proj_weight.chunk(3, dim=0)
        bq, bk, bv = mha.in_proj_bias.chunk(3, dim=0)

        q = F.linear(q_input, Wq, bq)  # [B, 1, d]

        if cache_k is None:
            k = F.linear(memory, Wk, bk)  # [B, T_text, d]
            v = F.linear(memory, Wv, bv)
        else:
            k, v = cache_k, cache_v

        q = q.view(B, 1, nhead, head_dim).transpose(1, 2)
        k_mh = k.view(B, -1, nhead, head_dim).transpose(1, 2)
        v_mh = v.view(B, -1, nhead, head_dim).transpose(1, 2)

        attn = torch.matmul(q, k_mh.transpose(-2, -1)) / (head_dim ** 0.5)
        if memory_mask is not None:
            # memory_mask: [B, T_text] → [B, 1, 1, T_text]
            attn = attn.masked_fill(memory_mask.unsqueeze(1).unsqueeze(2), float('-inf'))
        attn = F.softmax(attn, dim=-1)
        out = torch.matmul(attn, v_mh)

        out = out.transpose(1, 2).contiguous().view(B, 1, d)
        out = F.linear(out, mha.out_proj.weight, mha.out_proj.bias)

        return out, k, v

    def init_cache(self, num_layers):
        """Initialize empty KV cache."""
        return {
            'self_k': [None] * num_layers,
            'self_v': [None] * num_layers,
            'cross_k': [None] * num_layers,
            'cross_v': [None] * num_layers,
        }

    def inference_step(self, new_emb, text_emb, step_idx, cache, text_mask=None):
        """Process one token with KV cache.
        new_emb: [B, 1, d_model] — raw audio embedding (not projected)
        text_mask: [B, T_text] bool, True = padding
        Returns: predicted [B, 1, d_model], updated cache
        """
        x = self.audio_input_proj(new_emb)
        # Add positional encoding for this specific position
        x = x + self.pos_embed.pe[:, step_idx:step_idx + 1]

        for i, layer in enumerate(self.transformer.layers):
            # Self-attention with cache (post-norm)
            sa_out, cache['self_k'][i], cache['self_v'][i] = self._cached_mha(
                layer.self_attn, x, x, cache['self_k'][i], cache['self_v'][i]
            )
            x = layer.norm1(x + layer.dropout1(sa_out))

            # Cross-attention with cached text KV
            ca_out, cache['cross_k'][i], cache['cross_v'][i] = self._cached_cross_attn(
                layer.multihead_attn, x, text_emb, cache['cross_k'][i], cache['cross_v'][i],
                memory_mask=text_mask,
            )
            x = layer.norm2(x + layer.dropout2(ca_out))

            # FFN
            ff_out = layer.linear2(layer.dropout(layer.activation(layer.linear1(x))))
            x = layer.norm3(x + layer.dropout3(ff_out))

        return self.output_proj(x), cache

# ─── Positional Encoding ────────────────────────────────────────────────────

class SinusoidalPositionalEncoding(nn.Module):
    def __init__(self, d_model, max_len=8192):
        super().__init__()
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        self.register_buffer("pe", pe.unsqueeze(0))  # [1, max_len, d_model]

    def forward(self, x):
        return x + self.pe[:, :x.shape[1]]

# ─── Mel Decoder ────────────────────────────────────────────────────────────

class ResConvBlock(nn.Module):
    """Residual conv block for mel decoder."""

    def __init__(self, channels, kernel_size=3):
        super().__init__()
        self.net = nn.Sequential(
            nn.Conv1d(channels, channels, kernel_size, padding=kernel_size // 2),
            nn.GELU(),
            nn.Conv1d(channels, channels, kernel_size, padding=kernel_size // 2),
        )

    def forward(self, x):
        return x + self.net(x)

class MelDecoder(nn.Module):
    """Decodes latent embeddings back to mel spectrograms.
    Improved architecture with residual blocks for sharper mel output."""

    def __init__(self, d_model=256, n_mels=100, upsample_factor=4):
        super().__init__()
        hidden = d_model * 2  # 512 — more capacity

        self.proj = nn.Sequential(
            nn.Linear(d_model, hidden),
            nn.GELU(),
        )
        if upsample_factor == 2:
            self.upsample = nn.Sequential(
                nn.ConvTranspose1d(hidden, hidden, kernel_size=4, stride=2, padding=1),
                nn.GELU(),
                ResConvBlock(hidden, kernel_size=3),
            )
        else:
            self.upsample = nn.Sequential(
                # 4x upsample (2x + 2x)
                nn.ConvTranspose1d(hidden, hidden, kernel_size=4, stride=2, padding=1),
                nn.GELU(),
                ResConvBlock(hidden, kernel_size=3),
                nn.ConvTranspose1d(hidden, hidden, kernel_size=4, stride=2, padding=1),
                nn.GELU(),
                ResConvBlock(hidden, kernel_size=3),
            )
        # Refinement + projection to mel bins
        self.out = nn.Sequential(
            nn.Conv1d(hidden, hidden, kernel_size=7, padding=3),
            nn.GELU(),
            ResConvBlock(hidden, kernel_size=5),
            nn.Conv1d(hidden, n_mels, kernel_size=7, padding=3),
        )

    def forward(self, z):
        """z: [B, T_down, d_model] → mel: [B, n_mels, T_up]"""
        x = self.proj(z)
        x = x.transpose(1, 2)  # [B, hidden, T_down]
        x = self.upsample(x)   # [B, hidden, T_up]
        return self.out(x)      # [B, n_mels, T_up]

# ─── Multi-Resolution Spectral Loss ─────────────────────────────────────────

class MultiResolutionSpectralLoss(nn.Module):
    """Computes spectral convergence + log magnitude loss at multiple STFT resolutions.
    This forces the mel decoder to produce sharp, well-defined spectral structure
    rather than blurry averages that Vocos can't decode properly."""

    def __init__(self, resolutions=((64, 16, 64), (128, 32, 128), (256, 64, 256))):
        super().__init__()
        # Each resolution is (n_fft, hop_length, win_length)
        self.resolutions = resolutions

    def _stft_loss(self, pred, target, n_fft, hop_length, win_length):
        """Compute spectral convergence + log magnitude loss for one resolution."""
        window = torch.hann_window(win_length, device=pred.device)

        # pred/target: [B, n_mels, T] — treat as 1D signal per mel bin
        # Flatten mel bins into batch: [B*n_mels, T]
        B, M, T = pred.shape
        pred_flat = pred.reshape(B * M, T)
        target_flat = target.reshape(B * M, T)

        pred_stft = torch.stft(pred_flat, n_fft, hop_length, win_length,
                               window=window, return_complex=True)
        target_stft = torch.stft(target_flat, n_fft, hop_length, win_length,
                                 window=window, return_complex=True)

        pred_mag = pred_stft.abs()
        target_mag = target_stft.abs()

        # Spectral convergence: Frobenius norm of difference / Frobenius norm of target
        sc_loss = torch.norm(target_mag - pred_mag, p="fro") / (torch.norm(target_mag, p="fro") + 1e-8)

        # Log magnitude loss
        log_mag_loss = F.l1_loss(
            torch.log(pred_mag + 1e-7),
            torch.log(target_mag + 1e-7),
        )

        return sc_loss + log_mag_loss

    def forward(self, pred, target):
        """pred, target: [B, n_mels, T]"""
        loss = torch.tensor(0.0, device=pred.device)
        count = 0
        for n_fft, hop, win in self.resolutions:
            # Skip if T is too short for this resolution
            if pred.shape[2] >= n_fft:
                loss = loss + self._stft_loss(pred, target, n_fft, hop, win)
                count += 1
        if count > 0:
            loss = loss / count
        return loss

# ─── Full LeWM TTS Model ────────────────────────────────────────────────────

class LeWMTTS(nn.Module):
    """
    Complete LeWM TTS model.

    Training: text + mel → audio embeddings → predict next embedding
    Two losses:
        1. Next-embedding prediction: MSE(predicted[t], target[t+1])
        2. Gaussian regularizer: KL(q(z|x) || N(0, I))
    """

    def __init__(self, config):
        super().__init__()
        d_model = config.get("d_model", 256)
        nhead = config.get("nhead", 4)
        n_mels = config.get("n_mels", 100)
        text_vocab = config.get("text_vocab_size", 256)
        text_layers = config.get("text_encoder_layers", 4)
        audio_layers = config.get("audio_encoder_layers", 4)
        predictor_layers = config.get("predictor_layers", 6)
        dropout = config.get("dropout", 0.1)
        self.kl_weight = config.get("kl_weight", 0.01)
        n_speakers = config.get("n_speakers", 1)

        self.text_encoder = TextEncoder(
            vocab_size=text_vocab, d_model=d_model, nhead=nhead,
            num_layers=text_layers, dropout=dropout,
        )
        downsample_factor = config.get("downsample_factor", 4)
        self.audio_encoder = AudioEncoder(
            n_mels=n_mels, d_model=d_model, nhead=nhead,
            num_layers=audio_layers, downsample_factor=downsample_factor,
            dropout=dropout,
        )
        self.predictor = JEPAPredictor(
            d_model=d_model, nhead=nhead,
            num_layers=predictor_layers, dropout=dropout,
        )

        # Mel decoder: reconstructs mel from latent embeddings
        # Needed so the latent space is grounded to actual audio
        self.mel_decoder = MelDecoder(d_model=d_model, n_mels=n_mels, upsample_factor=downsample_factor)
        self.recon_weight = config.get("recon_weight", 1.0)

        # Speaker embedding (for multi-speaker)
        self.n_speakers = n_speakers
        if n_speakers > 1:
            self.speaker_embed = nn.Embedding(n_speakers, d_model)
        else:
            self.speaker_embed = None

        # Multi-resolution spectral loss for sharper mel output
        self.spectral_loss = MultiResolutionSpectralLoss()
        self.spectral_weight = config.get("spectral_weight", 0.5)

        self.config = config

    def _add_speaker(self, text_emb, speaker_id):
        """Add speaker embedding to text embeddings."""
        if self.speaker_embed is not None and speaker_id is not None:
            spk = self.speaker_embed(speaker_id).unsqueeze(1)  # [B, 1, d]
            text_emb = text_emb + spk
        return text_emb

    def forward(self, mel, text_tokens, mel_mask=None, text_mask=None, speaker_id=None):
        """
        Args:
            mel: [B, n_mels, T_mel]
            text_tokens: [B, T_text]
            mel_mask: [B, T_mel] bool
            text_mask: [B, T_text] bool
        Returns:
            loss_dict with 'prediction_loss', 'kl_loss', 'total_loss'
        """
        # Encode text + speaker conditioning
        text_emb = self.text_encoder(text_tokens, text_mask)
        text_emb = self._add_speaker(text_emb, speaker_id)

        # Encode audio → latent embeddings
        z, mu, logvar = self.audio_encoder(mel, mel_mask)

        # Predict next embedding: input z[:-1], target z[1:]
        input_emb = z[:, :-1]     # [B, T-1, d]
        target_emb = z[:, 1:]     # [B, T-1, d]

        # Adjust mask for downsampled sequence
        if mel_mask is not None:
            T_down = z.shape[1]
            ds_mask = mel_mask[:, ::self.audio_encoder.downsample_factor][:, :T_down]
            pred_mask = ds_mask[:, :-1]
        else:
            pred_mask = None

        predicted = self.predictor(input_emb, text_emb, pred_mask, text_mask)

        # ─── Loss 1: Next-embedding prediction (MSE) ───
        # Mask out padding positions
        if pred_mask is not None:
            valid = ~pred_mask  # True = valid
            valid = valid.unsqueeze(-1)  # [B, T-1, 1]
            prediction_loss = F.mse_loss(
                predicted * valid, target_emb * valid, reduction="sum"
            ) / (valid.sum() * predicted.shape[-1] + 1e-8)
        else:
            prediction_loss = F.mse_loss(predicted, target_emb)

        # ─── Loss 2: Gaussian regularizer — KL(N(mu, sigma²) || N(0, I)) ───
        # KL = -0.5 * sum(1 + logvar - mu² - exp(logvar))
        if mel_mask is not None:
            T_down = mu.shape[1]
            ds_mask = mel_mask[:, ::self.audio_encoder.downsample_factor][:, :T_down]
            valid = (~ds_mask).unsqueeze(-1).float()
            kl = -0.5 * (1 + logvar - mu.pow(2) - logvar.exp())
            kl_loss = (kl * valid).sum() / (valid.sum() * mu.shape[-1] + 1e-8)
        else:
            kl_loss = -0.5 * torch.mean(1 + logvar - mu.pow(2) - logvar.exp())

        # ─── Loss 3: Mel reconstruction ───
        # Decode latents back to mel and compare with input
        mel_recon = self.mel_decoder(z)  # [B, n_mels, T_up]
        # Trim/pad to match input mel length
        T_mel = mel.shape[2]
        T_recon = mel_recon.shape[2]
        if T_recon > T_mel:
            mel_recon = mel_recon[:, :, :T_mel]
        elif T_recon < T_mel:
            mel = mel[:, :, :T_recon]
            if mel_mask is not None:
                mel_mask = mel_mask[:, :T_recon]

        if mel_mask is not None:
            valid_mel = (~mel_mask).unsqueeze(1).float()  # [B, 1, T]
            recon_loss = F.l1_loss(
                mel_recon * valid_mel, mel * valid_mel, reduction="sum"
            ) / (valid_mel.sum() * mel.shape[1] + 1e-8)
        else:
            recon_loss = F.l1_loss(mel_recon, mel)

        # ─── Loss 4: Multi-resolution spectral loss ───
        # Apply on valid (non-padded) regions only
        if mel_mask is not None:
            spectral_loss = self.spectral_loss(mel_recon * valid_mel, mel * valid_mel)
        else:
            spectral_loss = self.spectral_loss(mel_recon, mel)

        total_loss = (prediction_loss
                      + self.kl_weight * kl_loss
                      + self.recon_weight * recon_loss
                      + self.spectral_weight * spectral_loss)

        return {
            "total_loss": total_loss,
            "prediction_loss": prediction_loss,
            "kl_loss": kl_loss,
            "recon_loss": recon_loss,
            "spectral_loss": spectral_loss,
        }

    def encode_audio(self, mel):
        """Encode mel to latent embeddings (for inference)."""
        z, mu, _ = self.audio_encoder(mel)
        return mu  # Use mean at inference

    def predict_next(self, audio_emb, text_emb, text_mask=None):
        """Predict next audio embedding autoregressively."""
        predicted = self.predictor(audio_emb, text_emb, text_mask=text_mask)
        return predicted[:, -1:]  # Only the last prediction

    def init_ar_cache(self):
        """Initialize KV cache for autoregressive inference."""
        num_layers = len(self.predictor.transformer.layers)
        return self.predictor.init_cache(num_layers)

    def predict_next_cached(self, new_emb, text_emb, step_idx, cache, text_mask=None):
        """Predict next embedding using KV cache. O(1) per step instead of O(n)."""
        return self.predictor.inference_step(new_emb, text_emb, step_idx, cache, text_mask=text_mask)

def count_parameters(model):
    return sum(p.numel() for p in model.parameters() if p.requires_grad)

def build_model(config=None):
    if config is None:
        config = {
            "d_model": 256,
            "nhead": 4,
            "n_mels": 100,
            "text_vocab_size": 256,
            "text_encoder_layers": 4,
            "audio_encoder_layers": 4,
            "predictor_layers": 6,
            "dropout": 0.1,
            "kl_weight": 0.01,
        }
    model = LeWMTTS(config)
    print(f"LeWM TTS model: {count_parameters(model)/1e6:.2f}M parameters")
    return model, config

"""
import math
+import copy
import torch
import torch.nn as nn
import torch.nn.functional as F


Now update the LeWMTTS `__init__` to add EMA encoder, learnable start embedding, and new config params:


        # Mel decoder: reconstructs mel from latent embeddings
        # Needed so the latent space is grounded to actual audio
        self.mel_decoder = MelDecoder(d_model=d_model, n_mels=n_mels, upsample_factor=downsample_factor)
        self.recon_weight = config.get("recon_weight", 1.0)

        # Speaker embedding (for multi-speaker)
        self.n_speakers = n_speakers
        if n_speakers > 1:
            self.speaker_embed = nn.Embedding(n_speakers, d_model)
        else:
            self.speaker_embed = None

        # Multi-resolution spectral loss for sharper mel output
        self.spectral_loss = MultiResolutionSpectralLoss()
        self.spectral_weight = config.get("spectral_weight", 0.5)

        self.config = config

        # Mel decoder: reconstructs mel from latent embeddings
        # Needed so the latent space is grounded to actual audio
        self.mel_decoder = MelDecoder(d_model=d_model, n_mels=n_mels, upsample_factor=downsample_factor)
        self.recon_weight = config.get("recon_weight", 1.0)

        # EMA target encoder — core JEPA component
        # Provides slowly-moving prediction targets to prevent representation collapse
        self.ema_audio_encoder = copy.deepcopy(self.audio_encoder)
        for p in self.ema_audio_encoder.parameters():
            p.requires_grad = False
        self.ema_decay = config.get("ema_decay", 0.998)

        # Speaker embedding (for multi-speaker)
        self.n_speakers = n_speakers
        if n_speakers > 1:
            self.speaker_embed = nn.Embedding(n_speakers, d_model)
        else:
            self.speaker_embed = None

        # Learnable start-of-sequence embedding for AR inference
        self.start_emb = nn.Parameter(torch.randn(1, 1, d_model) * 0.02)

        # Prediction weight — primary JEPA loss
        self.pred_weight = config.get("pred_weight", 10.0)

        # Input noise injection to bridge train/inference gap
        self.input_noise = config.get("input_noise", 0.0)

        # Free-bits KL: allow KL up to this threshold without penalty
        self.free_bits = config.get("free_bits", 1.0)

        # Multi-resolution spectral loss for sharper mel output
        self.spectral_loss = MultiResolutionSpectralLoss()
        self.spectral_weight = config.get("spectral_weight", 0.5)

        self.config = config


        # Mel decoder: reconstructs mel from latent embeddings
        # Needed so the latent space is grounded to actual audio
        self.mel_decoder = MelDecoder(d_model=d_model, n_mels=n_mels, upsample_factor=downsample_factor)
        self.recon_weight = config.get("recon_weight", 1.0)

        # EMA target encoder — core JEPA component
        # Provides slowly-moving prediction targets to prevent representation collapse
        self.ema_audio_encoder = copy.deepcopy(self.audio_encoder)
        for p in self.ema_audio_encoder.parameters():
            p.requires_grad = False
        self.ema_decay = config.get("ema_decay", 0.998)

        # Speaker embedding (for multi-speaker)
        self.n_speakers = n_speakers
        if n_speakers > 1:
            self.speaker_embed = nn.Embedding(n_speakers, d_model)
        else:
            self.speaker_embed = None

        # Learnable start-of-sequence embedding for AR inference
        self.start_emb = nn.Parameter(torch.randn(1, 1, d_model) * 0.02)

        # Prediction weight — primary JEPA loss
        self.pred_weight = config.get("pred_weight", 10.0)

        # Input noise injection to bridge train/inference gap
        self.input_noise = config.get("input_noise", 0.0)

        # Free-bits KL: allow KL up to this threshold without penalty
        self.free_bits = config.get("free_bits", 1.0)

        # Multi-resolution spectral loss for sharper mel output
        self.spectral_loss = MultiResolutionSpectralLoss()
        self.spectral_weight = config.get("spectral_weight", 0.5)

        self.config = config

"""
LeWM TTS Model — JEPA-based Text-to-Speech

Architecture:
  - TextEncoder: Character-level transformer for Hindi/Devanagari text
  - AudioEncoder: 1D CNN + Transformer that encodes mel spectrograms to embeddings
  - JEPAPredictor: Transformer that predicts next audio embedding given text + audio context
  - Two losses: next-embedding prediction (MSE) + Gaussian regularizer (KL to N(0,I))
"""

import math
import copy
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchaudio

# ─── Text Encoder ───────────────────────────────────────────────────────────

class TextEncoder(nn.Module):
    """Character-level transformer encoder for Hindi text."""

    def __init__(self, vocab_size=256, d_model=256, nhead=4, num_layers=4, dropout=0.1):
        super().__init__()
        self.d_model = d_model

        # Character embedding (covers all Unicode Hindi chars via byte-level encoding)
        self.char_embed = nn.Embedding(vocab_size, d_model)
        self.pos_embed = SinusoidalPositionalEncoding(d_model, max_len=2048)
        self.dropout = nn.Dropout(dropout)

        encoder_layer = nn.TransformerEncoderLayer(
            d_model=d_model, nhead=nhead,
            dim_feedforward=d_model * 4,
            dropout=dropout, batch_first=True,
            activation="gelu",
        )
        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
        self.proj = nn.Linear(d_model, d_model)

    def forward(self, text_tokens, text_mask=None):
        """
        Args:
            text_tokens: [B, T_text] long tensor of byte-level token IDs
            text_mask: [B, T_text] bool, True = padding
        Returns:
            text_emb: [B, T_text, d_model]
        """
        x = self.char_embed(text_tokens) * math.sqrt(self.d_model)
        x = self.pos_embed(x)
        x = self.dropout(x)
        x = self.transformer(x, src_key_padding_mask=text_mask)
        x = self.proj(x)
        return x

# ─── Audio Encoder ──────────────────────────────────────────────────────────

class AudioEncoder(nn.Module):
    """
    Encodes mel spectrograms into a sequence of audio embeddings.
    Uses 1D CNN for downsampling + Transformer for contextualization.
    """

    def __init__(self, n_mels=100, d_model=256, nhead=4, num_layers=4,
                 downsample_factor=4, dropout=0.1):
        super().__init__()
        self.d_model = d_model
        self.downsample_factor = downsample_factor

        # CNN downsampling: reduces temporal resolution by downsample_factor
        if downsample_factor == 2:
            self.conv_pre = nn.Sequential(
                nn.Conv1d(n_mels, d_model, kernel_size=7, padding=3),
                nn.GELU(),
                nn.Conv1d(d_model, d_model, kernel_size=4, stride=2, padding=1),
                nn.GELU(),
            )  # Total stride = 2
        else:
            self.conv_pre = nn.Sequential(
                nn.Conv1d(n_mels, d_model, kernel_size=7, padding=3),
                nn.GELU(),
                nn.Conv1d(d_model, d_model, kernel_size=4, stride=2, padding=1),
                nn.GELU(),
                nn.Conv1d(d_model, d_model, kernel_size=4, stride=2, padding=1),
                nn.GELU(),
            )  # Total stride = 4

        self.pos_embed = SinusoidalPositionalEncoding(d_model, max_len=4096)
        self.dropout = nn.Dropout(dropout)

        encoder_layer = nn.TransformerEncoderLayer(
            d_model=d_model, nhead=nhead,
            dim_feedforward=d_model * 4,
            dropout=dropout, batch_first=True,
            activation="gelu",
        )
        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)

        # Project to embedding with mu and logvar for Gaussian regularizer
        self.proj_mu = nn.Linear(d_model, d_model)
        self.proj_logvar = nn.Linear(d_model, d_model)

    def forward(self, mel, mel_mask=None):
        """
        Args:
            mel: [B, n_mels, T_mel] — log mel spectrogram
            mel_mask: [B, T_mel] bool, True = padding (before downsampling)
        Returns:
            z: [B, T_down, d_model] — sampled latent embeddings
            mu: [B, T_down, d_model]
            logvar: [B, T_down, d_model]
        """
        # CNN encoding: [B, n_mels, T] -> [B, d_model, T//4]
        x = self.conv_pre(mel)
        x = x.transpose(1, 2)  # [B, T_down, d_model]

        # Downsample mask if provided
        if mel_mask is not None:
            # Approximate downsampled mask
            T_down = x.shape[1]
            mel_mask = mel_mask[:, ::self.downsample_factor][:, :T_down]

        x = self.pos_embed(x)
        x = self.dropout(x)
        x = self.transformer(x, src_key_padding_mask=mel_mask)

        mu = self.proj_mu(x)
        logvar = self.proj_logvar(x)

        # Reparameterization trick
        if self.training:
            std = torch.exp(0.5 * logvar)
            eps = torch.randn_like(std)
            z = mu + eps * std
        else:
            z = mu

        return z, mu, logvar

# ─── JEPA Predictor ─────────────────────────────────────────────────────────

class JEPAPredictor(nn.Module):
    """
    Predicts next audio embedding given:
    - Text context (from TextEncoder)
    - Previous audio embeddings (from AudioEncoder)

    Uses cross-attention to condition on text, and causal self-attention
    over the audio embedding sequence.
    """

    def __init__(self, d_model=256, nhead=4, num_layers=6, dropout=0.1):
        super().__init__()
        self.d_model = d_model

        self.audio_input_proj = nn.Linear(d_model, d_model)
        self.pos_embed = SinusoidalPositionalEncoding(d_model, max_len=4096)
        self.dropout = nn.Dropout(dropout)

        # Transformer decoder with causal self-attention + cross-attention to text
        decoder_layer = nn.TransformerDecoderLayer(
            d_model=d_model, nhead=nhead,
            dim_feedforward=d_model * 4,
            dropout=dropout, batch_first=True,
            activation="gelu",
        )
        self.transformer = nn.TransformerDecoder(decoder_layer, num_layers=num_layers)

        # Output projection to predict next embedding
        self.output_proj = nn.Sequential(
            nn.Linear(d_model, d_model),
            nn.GELU(),
            nn.Linear(d_model, d_model),
        )

    def forward(self, audio_emb, text_emb, audio_mask=None, text_mask=None):
        """
        Args:
            audio_emb: [B, T_audio, d_model] — audio embeddings (from encoder)
            text_emb: [B, T_text, d_model] — text embeddings
            audio_mask: [B, T_audio] bool, True = padding
            text_mask: [B, T_text] bool, True = padding
        Returns:
            predicted: [B, T_audio, d_model] — predicted next embeddings
        """
        x = self.audio_input_proj(audio_emb)
        x = self.pos_embed(x)
        x = self.dropout(x)

        # Causal mask: each position can only attend to previous positions
        T = x.shape[1]
        causal_mask = torch.triu(
            torch.ones(T, T, device=x.device, dtype=torch.bool), diagonal=1
        )

        predicted = self.transformer(
            x, text_emb,
            tgt_mask=causal_mask,
            tgt_key_padding_mask=audio_mask,
            memory_key_padding_mask=text_mask,
        )

        predicted = self.output_proj(predicted)
        return predicted

    def _cached_mha(self, mha, q_input, kv_input, cache_k, cache_v):
        """Run multi-head attention with KV cache.
        q_input: [B, 1, d], kv_input: [B, 1, d] (new token for K,V)
        cache_k, cache_v: [B, T_prev, d] projected keys/values
        Returns: output [B, 1, d], updated cache_k, updated cache_v
        """
        d = mha.embed_dim
        nhead = mha.num_heads
        head_dim = d // nhead
        B = q_input.shape[0]

        # Project Q from q_input, K/V from kv_input
        # in_proj_weight is [3d, d] = [Wq; Wk; Wv]
        Wq, Wk, Wv = mha.in_proj_weight.chunk(3, dim=0)
        bq, bk, bv = mha.in_proj_bias.chunk(3, dim=0)

        q = F.linear(q_input, Wq, bq)        # [B, 1, d]
        k_new = F.linear(kv_input, Wk, bk)    # [B, 1, d]
        v_new = F.linear(kv_input, Wv, bv)    # [B, 1, d]

        # Append to cache
        if cache_k is not None and cache_k.shape[1] > 0:
            k = torch.cat([cache_k, k_new], dim=1)
            v = torch.cat([cache_v, v_new], dim=1)
        else:
            k = k_new
            v = v_new

        # Multi-head attention
        q = q.view(B, 1, nhead, head_dim).transpose(1, 2)      # [B, nhead, 1, head_dim]
        k_mh = k.view(B, -1, nhead, head_dim).transpose(1, 2)  # [B, nhead, T, head_dim]
        v_mh = v.view(B, -1, nhead, head_dim).transpose(1, 2)

        attn = torch.matmul(q, k_mh.transpose(-2, -1)) / (head_dim ** 0.5)
        attn = F.softmax(attn, dim=-1)
        out = torch.matmul(attn, v_mh)  # [B, nhead, 1, head_dim]

        out = out.transpose(1, 2).contiguous().view(B, 1, d)
        out = F.linear(out, mha.out_proj.weight, mha.out_proj.bias)

        return out, k, v

    def _cached_cross_attn(self, mha, q_input, memory, cache_k, cache_v, memory_mask=None):
        """Cross-attention with cached memory K,V (computed once).
        memory_mask: [B, T_text] bool, True = padding (applied as -inf)."""
        d = mha.embed_dim
        nhead = mha.num_heads
        head_dim = d // nhead
        B = q_input.shape[0]

        Wq, Wk, Wv = mha.in_proj_weight.chunk(3, dim=0)
        bq, bk, bv = mha.in_proj_bias.chunk(3, dim=0)

        q = F.linear(q_input, Wq, bq)  # [B, 1, d]

        if cache_k is None:
            k = F.linear(memory, Wk, bk)  # [B, T_text, d]
            v = F.linear(memory, Wv, bv)
        else:
            k, v = cache_k, cache_v

        q = q.view(B, 1, nhead, head_dim).transpose(1, 2)
        k_mh = k.view(B, -1, nhead, head_dim).transpose(1, 2)
        v_mh = v.view(B, -1, nhead, head_dim).transpose(1, 2)

        attn = torch.matmul(q, k_mh.transpose(-2, -1)) / (head_dim ** 0.5)
        if memory_mask is not None:
            # memory_mask: [B, T_text] → [B, 1, 1, T_text]
            attn = attn.masked_fill(memory_mask.unsqueeze(1).unsqueeze(2), float('-inf'))
        attn = F.softmax(attn, dim=-1)
        out = torch.matmul(attn, v_mh)

        out = out.transpose(1, 2).contiguous().view(B, 1, d)
        out = F.linear(out, mha.out_proj.weight, mha.out_proj.bias)

        return out, k, v

    def init_cache(self, num_layers):
        """Initialize empty KV cache."""
        return {
            'self_k': [None] * num_layers,
            'self_v': [None] * num_layers,
            'cross_k': [None] * num_layers,
            'cross_v': [None] * num_layers,
        }

    def inference_step(self, new_emb, text_emb, step_idx, cache, text_mask=None):
        """Process one token with KV cache.
        new_emb: [B, 1, d_model] — raw audio embedding (not projected)
        text_mask: [B, T_text] bool, True = padding
        Returns: predicted [B, 1, d_model], updated cache
        """
        x = self.audio_input_proj(new_emb)
        # Add positional encoding for this specific position
        x = x + self.pos_embed.pe[:, step_idx:step_idx + 1]

        for i, layer in enumerate(self.transformer.layers):
            # Self-attention with cache (post-norm)
            sa_out, cache['self_k'][i], cache['self_v'][i] = self._cached_mha(
                layer.self_attn, x, x, cache['self_k'][i], cache['self_v'][i]
            )
            x = layer.norm1(x + layer.dropout1(sa_out))

            # Cross-attention with cached text KV
            ca_out, cache['cross_k'][i], cache['cross_v'][i] = self._cached_cross_attn(
                layer.multihead_attn, x, text_emb, cache['cross_k'][i], cache['cross_v'][i],
                memory_mask=text_mask,
            )
            x = layer.norm2(x + layer.dropout2(ca_out))

            # FFN
            ff_out = layer.linear2(layer.dropout(layer.activation(layer.linear1(x))))
            x = layer.norm3(x + layer.dropout3(ff_out))

        return self.output_proj(x), cache

# ─── Positional Encoding ────────────────────────────────────────────────────

class SinusoidalPositionalEncoding(nn.Module):
    def __init__(self, d_model, max_len=8192):
        super().__init__()
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        self.register_buffer("pe", pe.unsqueeze(0))  # [1, max_len, d_model]

    def forward(self, x):
        return x + self.pe[:, :x.shape[1]]

# ─── Mel Decoder ────────────────────────────────────────────────────────────

class ResConvBlock(nn.Module):
    """Residual conv block for mel decoder."""

    def __init__(self, channels, kernel_size=3):
        super().__init__()
        self.net = nn.Sequential(
            nn.Conv1d(channels, channels, kernel_size, padding=kernel_size // 2),
            nn.GELU(),
            nn.Conv1d(channels, channels, kernel_size, padding=kernel_size // 2),
        )

    def forward(self, x):
        return x + self.net(x)

class MelDecoder(nn.Module):
    """Decodes latent embeddings back to mel spectrograms.
    Improved architecture with residual blocks for sharper mel output."""

    def __init__(self, d_model=256, n_mels=100, upsample_factor=4):
        super().__init__()
        hidden = d_model * 2  # 512 — more capacity

        self.proj = nn.Sequential(
            nn.Linear(d_model, hidden),
            nn.GELU(),
        )
        if upsample_factor == 2:
            self.upsample = nn.Sequential(
                nn.ConvTranspose1d(hidden, hidden, kernel_size=4, stride=2, padding=1),
                nn.GELU(),
                ResConvBlock(hidden, kernel_size=3),
            )
        else:
            self.upsample = nn.Sequential(
                # 4x upsample (2x + 2x)
                nn.ConvTranspose1d(hidden, hidden, kernel_size=4, stride=2, padding=1),
                nn.GELU(),
                ResConvBlock(hidden, kernel_size=3),
                nn.ConvTranspose1d(hidden, hidden, kernel_size=4, stride=2, padding=1),
                nn.GELU(),
                ResConvBlock(hidden, kernel_size=3),
            )
        # Refinement + projection to mel bins
        self.out = nn.Sequential(
            nn.Conv1d(hidden, hidden, kernel_size=7, padding=3),
            nn.GELU(),
            ResConvBlock(hidden, kernel_size=5),
            nn.Conv1d(hidden, n_mels, kernel_size=7, padding=3),
        )

    def forward(self, z):
        """z: [B, T_down, d_model] → mel: [B, n_mels, T_up]"""
        x = self.proj(z)
        x = x.transpose(1, 2)  # [B, hidden, T_down]
        x = self.upsample(x)   # [B, hidden, T_up]
        return self.out(x)      # [B, n_mels, T_up]

# ─── Multi-Resolution Spectral Loss ─────────────────────────────────────────

class MultiResolutionSpectralLoss(nn.Module):
    """Computes spectral convergence + log magnitude loss at multiple STFT resolutions.
    This forces the mel decoder to produce sharp, well-defined spectral structure
    rather than blurry averages that Vocos can't decode properly."""

    def __init__(self, resolutions=((64, 16, 64), (128, 32, 128), (256, 64, 256))):
        super().__init__()
        # Each resolution is (n_fft, hop_length, win_length)
        self.resolutions = resolutions

    def _stft_loss(self, pred, target, n_fft, hop_length, win_length):
        """Compute spectral convergence + log magnitude loss for one resolution."""
        window = torch.hann_window(win_length, device=pred.device)

        # pred/target: [B, n_mels, T] — treat as 1D signal per mel bin
        # Flatten mel bins into batch: [B*n_mels, T]
        B, M, T = pred.shape
        pred_flat = pred.reshape(B * M, T)
        target_flat = target.reshape(B * M, T)

        pred_stft = torch.stft(pred_flat, n_fft, hop_length, win_length,
                               window=window, return_complex=True)
        target_stft = torch.stft(target_flat, n_fft, hop_length, win_length,
                                 window=window, return_complex=True)

        pred_mag = pred_stft.abs()
        target_mag = target_stft.abs()

        # Spectral convergence: Frobenius norm of difference / Frobenius norm of target
        sc_loss = torch.norm(target_mag - pred_mag, p="fro") / (torch.norm(target_mag, p="fro") + 1e-8)

        # Log magnitude loss
        log_mag_loss = F.l1_loss(
            torch.log(pred_mag + 1e-7),
            torch.log(target_mag + 1e-7),
        )

        return sc_loss + log_mag_loss

    def forward(self, pred, target):
        """pred, target: [B, n_mels, T]"""
        loss = torch.tensor(0.0, device=pred.device)
        count = 0
        for n_fft, hop, win in self.resolutions:
            # Skip if T is too short for this resolution
            if pred.shape[2] >= n_fft:
                loss = loss + self._stft_loss(pred, target, n_fft, hop, win)
                count += 1
        if count > 0:
            loss = loss / count
        return loss

# ─── Full LeWM TTS Model ────────────────────────────────────────────────────

class LeWMTTS(nn.Module):
    """
    Complete LeWM TTS model.

    Training: text + mel → audio embeddings → predict next embedding
    Two losses:
        1. Next-embedding prediction: MSE(predicted[t], target[t+1])
        2. Gaussian regularizer: KL(q(z|x) || N(0, I))
    """

    def __init__(self, config):
        super().__init__()
        d_model = config.get("d_model", 256)
        nhead = config.get("nhead", 4)
        n_mels = config.get("n_mels", 100)
        text_vocab = config.get("text_vocab_size", 256)
        text_layers = config.get("text_encoder_layers", 4)
        audio_layers = config.get("audio_encoder_layers", 4)
        predictor_layers = config.get("predictor_layers", 6)
        dropout = config.get("dropout", 0.1)
        self.kl_weight = config.get("kl_weight", 0.01)
        n_speakers = config.get("n_speakers", 1)

        self.text_encoder = TextEncoder(
            vocab_size=text_vocab, d_model=d_model, nhead=nhead,
            num_layers=text_layers, dropout=dropout,
        )
        downsample_factor = config.get("downsample_factor", 4)
        self.audio_encoder = AudioEncoder(
            n_mels=n_mels, d_model=d_model, nhead=nhead,
            num_layers=audio_layers, downsample_factor=downsample_factor,
            dropout=dropout,
        )
        self.predictor = JEPAPredictor(
            d_model=d_model, nhead=nhead,
            num_layers=predictor_layers, dropout=dropout,
        )

        # Mel decoder: reconstructs mel from latent embeddings
        # Needed so the latent space is grounded to actual audio
        self.mel_decoder = MelDecoder(d_model=d_model, n_mels=n_mels, upsample_factor=downsample_factor)
        self.recon_weight = config.get("recon_weight", 1.0)

        # Speaker embedding (for multi-speaker)
        self.n_speakers = n_speakers
        if n_speakers > 1:
            self.speaker_embed = nn.Embedding(n_speakers, d_model)
        else:
            self.speaker_embed = None

        # Multi-resolution spectral loss for sharper mel output
        self.spectral_loss = MultiResolutionSpectralLoss()
        self.spectral_weight = config.get("spectral_weight", 0.5)

        self.config = config

    def _add_speaker(self, text_emb, speaker_id):
        """Add speaker embedding to text embeddings."""
        if self.speaker_embed is not None and speaker_id is not None:
            spk = self.speaker_embed(speaker_id).unsqueeze(1)  # [B, 1, d]
            text_emb = text_emb + spk
        return text_emb

    def forward(self, mel, text_tokens, mel_mask=None, text_mask=None, speaker_id=None):
        """
        Args:
            mel: [B, n_mels, T_mel]
            text_tokens: [B, T_text]
            mel_mask: [B, T_mel] bool
            text_mask: [B, T_text] bool
        Returns:
            loss_dict with 'prediction_loss', 'kl_loss', 'total_loss'
        """
        # Encode text + speaker conditioning
        text_emb = self.text_encoder(text_tokens, text_mask)
        text_emb = self._add_speaker(text_emb, speaker_id)

        # Encode audio → latent embeddings
        z, mu, logvar = self.audio_encoder(mel, mel_mask)

        # Predict next embedding: input z[:-1], target z[1:]
        input_emb = z[:, :-1]     # [B, T-1, d]
        target_emb = z[:, 1:]     # [B, T-1, d]

        # Adjust mask for downsampled sequence
        if mel_mask is not None:
            T_down = z.shape[1]
            ds_mask = mel_mask[:, ::self.audio_encoder.downsample_factor][:, :T_down]
            pred_mask = ds_mask[:, :-1]
        else:
            pred_mask = None

        predicted = self.predictor(input_emb, text_emb, pred_mask, text_mask)

        # ─── Loss 1: Next-embedding prediction (MSE) ───
        # Mask out padding positions
        if pred_mask is not None:
            valid = ~pred_mask  # True = valid
            valid = valid.unsqueeze(-1)  # [B, T-1, 1]
            prediction_loss = F.mse_loss(
                predicted * valid, target_emb * valid, reduction="sum"
            ) / (valid.sum() * predicted.shape[-1] + 1e-8)
        else:
            prediction_loss = F.mse_loss(predicted, target_emb)

        # ─── Loss 2: Gaussian regularizer — KL(N(mu, sigma²) || N(0, I)) ───
        # KL = -0.5 * sum(1 + logvar - mu² - exp(logvar))
        if mel_mask is not None:
            T_down = mu.shape[1]
            ds_mask = mel_mask[:, ::self.audio_encoder.downsample_factor][:, :T_down]
            valid = (~ds_mask).unsqueeze(-1).float()
            kl = -0.5 * (1 + logvar - mu.pow(2) - logvar.exp())
            kl_loss = (kl * valid).sum() / (valid.sum() * mu.shape[-1] + 1e-8)
        else:
            kl_loss = -0.5 * torch.mean(1 + logvar - mu.pow(2) - logvar.exp())

        # ─── Loss 3: Mel reconstruction ───
        # Decode latents back to mel and compare with input
        mel_recon = self.mel_decoder(z)  # [B, n_mels, T_up]
        # Trim/pad to match input mel length
        T_mel = mel.shape[2]
        T_recon = mel_recon.shape[2]
        if T_recon > T_mel:
            mel_recon = mel_recon[:, :, :T_mel]
        elif T_recon < T_mel:
            mel = mel[:, :, :T_recon]
            if mel_mask is not None:
                mel_mask = mel_mask[:, :T_recon]

        if mel_mask is not None:
            valid_mel = (~mel_mask).unsqueeze(1).float()  # [B, 1, T]
            recon_loss = F.l1_loss(
                mel_recon * valid_mel, mel * valid_mel, reduction="sum"
            ) / (valid_mel.sum() * mel.shape[1] + 1e-8)
        else:
            recon_loss = F.l1_loss(mel_recon, mel)

        # ─── Loss 4: Multi-resolution spectral loss ───
        # Apply on valid (non-padded) regions only
        if mel_mask is not None:
            spectral_loss = self.spectral_loss(mel_recon * valid_mel, mel * valid_mel)
        else:
            spectral_loss = self.spectral_loss(mel_recon, mel)

        total_loss = (prediction_loss
                      + self.kl_weight * kl_loss
                      + self.recon_weight * recon_loss
                      + self.spectral_weight * spectral_loss)

        return {
            "total_loss": total_loss,
            "prediction_loss": prediction_loss,
            "kl_loss": kl_loss,
            "recon_loss": recon_loss,
            "spectral_loss": spectral_loss,
        }

    def encode_audio(self, mel):
        """Encode mel to latent embeddings (for inference)."""
        z, mu, _ = self.audio_encoder(mel)
        return mu  # Use mean at inference

    def predict_next(self, audio_emb, text_emb, text_mask=None):
        """Predict next audio embedding autoregressively."""
        predicted = self.predictor(audio_emb, text_emb, text_mask=text_mask)
        return predicted[:, -1:]  # Only the last prediction

    def init_ar_cache(self):
        """Initialize KV cache for autoregressive inference."""
        num_layers = len(self.predictor.transformer.layers)
        return self.predictor.init_cache(num_layers)

    def predict_next_cached(self, new_emb, text_emb, step_idx, cache, text_mask=None):
        """Predict next embedding using KV cache. O(1) per step instead of O(n)."""
        return self.predictor.inference_step(new_emb, text_emb, step_idx, cache, text_mask=text_mask)

def count_parameters(model):
    return sum(p.numel() for p in model.parameters() if p.requires_grad)

def build_model(config=None):
    if config is None:
        config = {
            "d_model": 256,
            "nhead": 4,
            "n_mels": 100,
            "text_vocab_size": 256,
            "text_encoder_layers": 4,
            "audio_encoder_layers": 4,
            "predictor_layers": 6,
            "dropout": 0.1,
            "kl_weight": 0.01,
        }
    model = LeWMTTS(config)
    print(f"LeWM TTS model: {count_parameters(model)/1e6:.2f}M parameters")
    return model, config

+        # EMA target encoder — core JEPA component

+        # Provides slowly-moving prediction targets to prevent representation collapse

+        self.ema_audio_encoder = copy.deepcopy(self.audio_encoder)

+        for p in self.ema_audio_encoder.parameters():

+        self.ema_decay = config.get("ema_decay", 0.998)

+        # Learnable start-of-sequence embedding for AR inference

+        self.start_emb = nn.Parameter(torch.randn(1, 1, d_model) * 0.02)

+        # Prediction weight — primary JEPA loss

+        self.pred_weight = config.get("pred_weight", 10.0)

+        # Input noise injection to bridge train/inference gap

+        self.input_noise = config.get("input_noise", 0.0)

+        # Free-bits KL: allow KL up to this threshold without penalty

+        self.free_bits = config.get("free_bits", 1.0)


Now add the `update_ema` method after `_add_speaker`:


    def _add_speaker(self, text_emb, speaker_id):
        """Add speaker embedding to text embeddings."""
        if self.speaker_embed is not None and speaker_id is not None:
            spk = self.speaker_embed(speaker_id).unsqueeze(1)  # [B, 1, d]
            text_emb = text_emb + spk
        return text_emb

    def forward(self

    def _add_speaker(self, text_emb, speaker_id):
        """Add speaker embedding to text embeddings."""
        if self.speaker_embed is not None and speaker_id is not None:
            spk = self.speaker_embed(speaker_id).unsqueeze(1)  # [B, 1, d]
            text_emb = text_emb + spk
        return text_emb

    @torch.no_grad()
    def update_ema(self):
        """Update EMA target encoder weights."""
        for p_online, p_ema in zip(self.audio_encoder.parameters(),
                                   self.ema_audio_encoder.parameters()):
            p_ema.data.mul_(self.ema_decay).add_(p_online.data, alpha=1 - self.ema_decay)

    def forward(self