Perfect! Now I have all the information needed. Let me create a comprehensive report:

## LeWM TTS Codebase - Comprehensive Exploration Report

I've completed a thorough exploration of the LeWM TTS codebase at `/home/ubuntu/lewm-tts`. Here's a complete breakdown of all the key components:

---

### 1. MAIN MODEL ARCHITECTURE FILES

#### **Primary Model: `/home/ubuntu/lewm-tts/model.py`** (708 lines)
The core JEPA-based TTS implementation with 4 main components:

**TextEncoder:**
- Byte-level character transformer for Hindi/UTF-8 text
- Config: 256d embedding, 4 attention heads, 4 transformer layers
- Character embedding + sinusoidal positional encoding
- Linear projection output

**AudioEncoder:**
- 1D CNN downsampler (stride-2 or stride-4) + Transformer contextualization
- Input: mel spectrogram [B, n_mels, T_mel]
- 1D convolutions with GELU activation for temporal compression
- Sinusoidal positional encoding + 4 transformer encoder layers
- **Gaussian latent space** with reparameterization:
  - `proj_mu`: projects to mean (μ)
  - `proj_logvar`: projects to log-variance (σ)
  - During training: z = μ + ε·σ (reparameterization trick)
  - During inference: z = μ (deterministic)

**JEPAPredictor:**
- Causal transformer decoder with **cross-attention to text**
- 6 transformer layers with:
  - Self-attention over audio embeddings (causal masking)
  - Cross-attention to text embeddings
  - Feed-forward networks
- Output projection to predict next audio embedding
- **KV-cache support** for efficient autoregressive inference (lines 209-323):
  - `_cached_mha()`: multi-head attention with cached K,V
  - `_cached_cross_attn()`: cross-attention with pre-computed text K,V
  - `init_cache()`: initialize empty KV cache
  - `inference_step()`: process one token in O(1) time using cached attention

**MelDecoder:**
- Reconstructs mel spectrograms from latent embeddings
- Architecture:
  - Linear projection to 512d (2× d_model)
  - Transposed convolution upsample (2×2 = 4×)
  - ResConvBlock residual blocks for sharpness
  - Final conv refinement → n_mels output
- Output: [B, n_mels, T_up]

**MultiResolutionSpectralLoss:**
- STFT loss at 3 resolutions: (64/16/64), (128/32/128), (256/64/256)
- For each resolution: spectral convergence + log magnitude loss
- Forces sharp, well-defined mel spectral structure

**Full LeWMTTS Model:**
- **EMA target encoder** (core JEPA component):
  - `ema_audio_encoder`: exponential moving average copy of online encoder
  - `ema_decay`: 0.998 (updated each training step)
  - Provides slow-moving, stable prediction targets
  - Prevents representation collapse
- **Learnable start embedding**: [1, 1, d_model] to seed AR generation
- **Speaker conditioning** (optional multi-speaker):
  - `speaker_embed`: embedding table (if n_speakers > 1)
  - Added to text embeddings via `_add_speaker()`
- **4 Loss functions** (lines 605-654):
  1. **Prediction Loss (MSE)**: MSE(predicted[t], ema_target[t+1])
  2. **KL Loss**: Free-bits KL divergence with annealing
     - `kl_free = max(0, kl_per_dim - free_bits)` (prevents collapse)
  3. **Reconstruction Loss (L1)**: mel_decoder output vs ground truth mel
  4. **Spectral Loss (STFT)**: multi-resolution spectral convergence

---

#### **V5 Model: `/home/ubuntu/lewm-tts/model_v5.py`** (463 lines)
Evolved codec-based variant using EnCodec instead of mel:

**Key Differences:**
- Encodes audio with **EnCodec (128-dim embeddings @ 75Hz)**
- Simple linear projections instead of convolutional encoder/decoder
- **Duration predictor** for text-to-audio alignment (FastSpeech-style)
- Frame-level text conditioning (text expanded then added to audio input)

**DurationPredictor:**
- Conv1d layers with LayerNorm and dropout
- Predicts log-duration per text character
- `compute_uniform_durations()`: spreads audio frames evenly across text

**JEPAPredictor (v5):**
- Uses **TransformerEncoder** (self-attention only)
- Text embeddings expanded to frame-level and **summed with audio input**
- Cannot ignore text—it's baked into the input
- Backward-compatible cross-attention interface

**LeWMTTSv5:**
- `proj_in`: codec_dim (128) → d_model (256)
- `proj_out`: d_model (256) → codec_dim (128)
- EMA on `proj_in` (provides targets)
- **3 Loss functions**:
  1. Prediction loss (MSE)
  2. Roundtrip loss (proj_in → proj_out → match original codec embedding)
  3. Duration loss (MSE on log-durations)

---

#### **V2 Model: `/home/ubuntu/lewm-tts/v2/model.py`** (242 lines)
Earlier DAC-based variant:

- DAC encoder (1024-dim latents @ 75Hz, frozen)
- Simple projection: 1024 ↔ 256
- Gaussian regularizer projections (mu/logvar)
- Cross-attention predictor
- No mel decoder—DAC handles reconstruction
- Losses: prediction (MSE) + KL + roundtrip reconstruction

---

### 2. SPEAKER CONDITIONING IMPLEMENTATION

**Location:** `/home/ubuntu/lewm-tts/model.py`, lines 510-540

**Method: Additive Embedding**
```python
self.speaker_embed = nn.Embedding(n_speakers, d_model)  # if n_speakers > 1

def _add_speaker(self, text_emb, speaker_id):
    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
```

- Speaker ID is converted to d_model-dimensional embedding
- Added to text embeddings (broadcast across time dimension)
- Called in forward pass before passing to predictor
- Currently designed but **not yet trained** (mentioned in README as "planned")

---

### 3. TRAINING LOOP & LOSS FUNCTIONS

**Main Training: `/home/ubuntu/lewm-tts/train.py`** (246 lines)

**Loss Weights (configurable):**
```python
pred_weight: 10.0        # Next-embedding prediction
kl_weight: 0.05-0.1      # Gaussian regularizer
recon_weight: 1.0        # Mel reconstruction
spectral_weight: 0.5     # Multi-resolution spectral
```

**Total Loss:**
```python
loss = (pred_weight × prediction_loss
        + kl_weight × kl_loss
        + recon_weight × recon_loss
        + spectral_weight × spectral_loss)
```

**Annealing Strategies:**
- **KL Annealing**: ramps from 0 to 1 over first 20% of training
- **Input Noise Annealing**: ramps from 0 to 0.2 over first 30% of training
  - Bridges train/inference gap by simulating imperfect AR predictions
  - Only applied to positions 1+ (not start_emb)

**Free-Bits KL:**
```python
kl_free = torch.clamp(kl_per_dim - free_bits, min=0.0)  # free_bits=1.0
```
- Allows KL up to threshold without penalty
- Prevents both collapse (too low KL) and explosion (too high KL)

**Optimizer & Scheduler:**
- AdamW: lr=1e-4, betas=(0.9, 0.98), weight_decay=0.01
- Warmup: min(4000 steps, 10% of total)
- Cosine annealing decay over remaining steps
- Gradient clipping: 1.0

**EMA Update:**
```python
model.update_ema()  # After each optimizer step
```

**V5 Training: `/home/ubuntu/lewm-tts/train_v5.py`** (similar structure)
- Codec embedding input instead of mel
- 3 losses: prediction, roundtrip, duration
- Input noise annealing only

---

### 4. JEPA ENCODER/PREDICTOR/TARGET ENCODER INTERACTION

**Core JEPA Loop (lines 542-612 in model.py):**

1. **Online Encoder** processes mel spectrogram:
   ```python
   z, mu, logvar = self.audio_encoder(mel, mel_mask)  # [B, T_down, d]
   ```

2. **EMA Target Encoder** (no gradients, always eval):
   ```python
   with torch.no_grad():
       self.ema_audio_encoder.eval()
       _, mu_ema, _ = self.ema_audio_encoder(mel, mel_mask)
       self.ema_audio_encoder.train()
       target_emb = mu_ema  # [B, T_down, d]
   ```

3. **Prepare Input Sequence** (shifted right):
   ```python
   start = self.start_emb.expand(B, -1, -1)        # [B, 1, d]
   input_emb = torch.cat([start, z[:, :-1]], dim=1)  # [B, T_down, d]
   # Input:  [start_emb, z[0], z[1], ..., z[T-2]]
   # Target: [z_ema[0],  z_ema[1], ..., z_ema[T-1]]
   ```

4. **Optionally Add Input Noise** (training only):
   ```python
   if self.training and self.input_noise > 0:
       noise[:, 1:] = randn(...) * self.input_noise
       input_emb = input_emb + noise
   ```

5. **Predictor** predicts next embeddings:
   ```python
   predicted = self.predictor(input_emb, text_emb, pred_mask, text_mask)
   # predicted[t] should match target_emb[t]
   ```

6. **EMA Update**:
   ```python
   for p_online, p_ema in zip(audio_encoder.parameters(), ema_encoder.parameters()):
       p_ema = p_ema * 0.998 + p_online * 0.002
   ```

**Why This Design:**
- **Prevents collapse**: EMA encoder provides slowly-moving targets
- **Prediction learning**: Predictor learns to extrapolate from text context
- **Stability**: No gradient feedback to targets; prevents training instability

---

### 5. CONFIG FILES & MODEL DIMENSIONS

**Main Config (output_v3/config.json):**
```json
{
  "d_model": 256,
  "nhead": 4,
  "n_mels": 100,
  "text_vocab_size": 256,
  "text_encoder_layers": 4,
  "audio_encoder_layers": 4,
  "predictor_layers": 6,
  "downsample_factor": 4,
  "dropout": 0.1,
  "kl_weight": 0.1,
  "recon_weight": 1.0,
  "spectral_weight": 0.5,
  "pred_weight": 10.0,
  "ema_decay": 0.998,
  "free_bits": 1.0
}
```

**V5 Config (output_v5/config.json):**
```json
{
  "d_model": 256,
  "nhead": 4,
  "codec_dim": 128,
  "text_vocab_size": 256,
  "text_encoder_layers": 4,
  "predictor_layers": 6,
  "dropout": 0.1,
  "pred_weight": 10.0,
  "roundtrip_weight": 1.0,
  "ema_decay": 0.998,
  "input_noise": 0.0
}
```

**Key Dimensions:**
- d_model: 256 (embedding dimension)
- nhead: 4 (attention heads, head_dim=64)
- text_vocab_size: 256 (byte-level UTF-8 encoding, 0-255)
- n_mels: 100 (mel bins, matches Vocos)
- downsample_factor: 4 (CNN stride in audio encoder)
- Total parameters: ~24M (v3), similar for v5

---

### 6. INFERENCE & SYNTHESIS PIPELINE

**Basic Inference: `/home/ubuntu/lewm-tts/inference.py`** (151 lines)

**Class: LeWMTTSInference**

**Methods:**

1. **`synthesize(text, max_steps=300, temperature=0.0)`** (lines 48-82)
   - Text → byte-level tokens
   - Encode text: TextEncoder → [B, T_text, d]
   - Use learned start_emb to seed
   - **Autoregressive loop with KV-cache**:
     ```python
     all_embeddings = [start_emb]
     cache = model.init_ar_cache()

     for step in range(max_steps):
         next_emb, cache = model.predict_next_cached(
             next_emb, text_emb, step, cache
         )
         all_embeddings.append(next_emb)

         # Stop early if embeddings stabilize
         if embedding_norm_variance < threshold:
             break
     ```
   - Concatenate embeddings: [B, T_total, d]
   - MelDecoder → mel [B, 100, T]
   - Vocos vocoder → waveform

2. **`reconstruct(mel_path)`** (lines 85-93)
   - Encode real mel
   - MelDecoder reconstruction
   - Vocos decode

3. **`reconstruct_with_prediction(mel_path)`** (lines 96-110)
   - Encode real mel → z
   - Predictor with dummy text (teacher-forced)
   - Decode predicted embeddings

---

**V5 Inference: `/home/ubuntu/lewm-tts/inference_v5.py`** (263 lines)

**Methods:**

1. **`synthesize(text, duration_scale=1.0, temperature=0.0)`** (lines 40-87)
   - Text → TextEncoder → [B, T_text, d]
   - **DurationPredictor** → log_durations → durations
   - `length_regulate()`: expand text embeddings by durations
   - Autoregressive loop (similar to v1)
   - proj_out → codec embeddings
   - EnCodec decoder → waveform

2. **`synthesize_prompted(text, prompt_audio, ...)`** (lines 90-160)
   - Encode prompt audio with EnCodec
   - Process prompt frames through predictor cache
   - Continue AR generation from prompt
   - Allows voice cloning by providing reference audio

3. **`reconstruct(audio_path)`** (lines 163-173)
   - Tests roundtrip: encode → proj_in → proj_out → decode

4. **`reconstruct_with_prediction(audio_path, text)`** (lines 187-220)
   - Encode audio → z
   - Expand text uniformly to audio length
   - Predictor with teacher forcing
   - Decode

---

**Production Inference Engine: `/home/ubuntu/lewm-tts/inference_engine.py`** (partial view)

**Features:**
- Batched inference (concurrent requests)
- FP16 support for GPU efficiency
- KV-cache for O(1) per-step inference
- Audio prompt encoding for voice cloning
- Vocos vocoder integration
- Benchmarking utilities

**Performance (A100 80GB):**
- Single request: 9× real-time
- Batch size 64: 359× real-time throughput

---

### 7. DATASET & PREPROCESSING

**Dataset: `/home/ubuntu/lewm-tts/dataset.py`** (98 lines)

**TTSDataset:**
- Loads from manifest.json (list of entries with mel_path, text)
- Max length filtering: 1200 mel frames, 512 text tokens
- Byte-level text tokenization (UTF-8 encoding)
- Returns: {mel [n_mels, T], text_tokens [T_text]}

**Collate Function:**
- Pads mel specs to max length along time axis
- Pads text tokens to max length
- Creates boolean masks (True = padding)
- Returns: mel [B, n_mels, T_max], text_tokens [B, T_text_max], mel_mask, text_mask

**V5 Dataset: `/home/ubuntu/lewm-tts/dataset_v5.py`** (68 lines)

**CodecDataset:**
- Loads precomputed EnCodec embeddings [128, T]
- Byte-level text tokens
- Max 900 codec frames
- Collate: pad codec embeddings [B, 128, T_max], text tokens [B, T_text_max]

---

**Preprocessing: `/home/ubuntu/lewm-tts/preprocess.py`** (100+ lines)

**Mel Config:**
```python
sample_rate: 24000
n_fft: 1024
hop_length: 256
n_mels: 100
f_min: 0
f_max: None  # full bandwidth
power: 1.0
clip_val: 1e-7  # log clamp (matches Vocos)
```

**Segmentation Config:**
- Min duration: 3 seconds
- Max duration: 12 seconds
- Target chunk: 8 seconds
- Silence threshold: -40 dB
- Min silence: 0.3 seconds

**Process:**
1. Load audio from HuggingFace dataset
2. Detect silence points (RMS-based)
3. Segment into 3-12s chunks at silence boundaries
4. Extract mel spectrograms
5. Save preprocessed dataset + manifest.json

**Dataset Stats:**
- Source: 20 hours single-speaker Hindi audio
- Processed: 7,009 segments, 17.88 hours
- Mel extraction: matches Vocos `mel-24khz` exactly

---

### 8. DIRECTORY STRUCTURE

```
/home/ubuntu/lewm-tts/
├── model.py                    # Main v1 model (mel-based)
├── model_v5.py                 # V5 model (codec-based)
├── train.py                    # Training loop v1
├── train_v5.py                 # Training loop v5
├── inference.py                # Single-request inference v1
├── inference_v5.py             # Single-request inference v5
├── inference_engine.py         # Production batched inference
├── dataset.py                  # PyTorch dataset (mel-based)
├── dataset_v5.py               # PyTorch dataset (codec-based)
├── preprocess.py               # Data preprocessing
├── README.md                    # Documentation

├── v2/                         # Earlier DAC-based variant
│   ├── model.py
│   ├── train.py
│   ├── inference.py
│   ├── dataset.py
│   ├── preprocess.py
│   └── output/
│       └── checkpoints/

├── processed_data/             # Mel-based dataset
│   ├── manifest.json
│   └── config.json

├── processed_data_codec/       # Codec-based dataset
│   └── manifest.json

├── output_v3/                  # Latest mel-based checkpoint
│   ├── checkpoints/
│   ├── logs/
│   └── config.json

├── output_v5/                  # Latest codec-based checkpoint
│   ├── checkpoints/
│   ├── logs/
│   └── config.json

├── samples_v5/                 # Generated audio samples
├── samples_v5_pred/
├── samples_v5_prompted/        # Prompted synthesis outputs
└── [other output dirs...]
```

---

### 9. ARCHITECTURE SUMMARY DIAGRAM

```
TEXT INPUT (UTF-8 bytes)
         ↓
    TextEncoder [4 layers, 256d, 4 heads]
         ↓
   text_emb [B, T_text, 256d]
         ↓ (+ optional speaker embedding)

                            MEL INPUT [B, 100, T_mel]
                                      ↓
                         AudioEncoder (1D CNN + Transformer)
                         1. Conv1d stride-4 downsample
                         2. Transformer [4 layers, 256d, 4 heads]
                         3. Gaussian projections (mu/logvar)
                                      ↓
                         z [B, T_down, 256d] (sampled)
                         mu [B, T_down, 256d] (mean)
                         logvar [B, T_down, 256d] (log-variance)

                         EMA AudioEncoder (same, no gradient, 0.998 decay)
                                      ↓
                         target_emb [B, T_down, 256d]

    [start_emb] + z[:-1] ──→ JEPAPredictor ←── text_emb
                              6 layers:
                              - Causal self-attn
                              - Cross-attn to text
                              - FFN
                                      ↓
                         predicted [B, T_down, 256d]
                                      ↓
                         MelDecoder (linear + transpose conv + ResBlocks)
                                      ↓
                         mel_recon [B, 100, T_mel]
                                      ↓
                         Vocos Vocoder (24khz)
                                      ↓
                         waveform [24000 hz]
```

---

### 10. KEY TRAINING INSIGHTS

1. **Representation Learning**: The EMA target encoder provides stable targets that prevent the online encoder from collapsing to a meaningless representation.

2. **Text Conditioning**: Cross-attention from predictor to text embeddings allows the model to learn text-aligned predictions, grounding the latent space in linguistic structure.

3. **Multi-Scale Losses**: Combination of MSE (prediction), KL (regularization), L1 (reconstruction), and STFT (spectral detail) ensures well-formed latent representations and sharp mel output.

4. **Free-Bits Regularization**: Prevents both KL collapse (model ignores text) and excessive KL divergence (latent space becomes too stochastic).

5. **Input Noise & Annealing**: Simulates the AR inference distribution during training, bridging the train/test gap where errors accumulate.

6. **Causal Masking**: The predictor uses causal attention, ensuring it only attends to past audio and current text context—enforcing autoregressiveness.

7. **Vocab Design**: Byte-level encoding (256 tokens) handles any Unicode text (Hindi/English/etc.) without needing language-specific tokenizers.

---

### 11. INFERENCE OPTIMIZATION

- **KV-Cache**: Reduces per-step complexity from O(T) to O(1) by caching previous attention keys/values
- **FP16**: Reduces memory footprint and increases throughput
- **Batch Processing**: Concurrent inference on multiple requests with shared compute
- **Early Stopping**: Monitors embedding norm variance to stop AR generation when stable

All files are located in `/home/ubuntu/lewm-tts/` with subdirectories for different training variants (v2, v3, v5) and output checkpoints.

Prompt:
Thoroughly explore the LeWM TTS codebase in /home/ubuntu. Find:

1. The main model architecture files — encoder, decoder, predictor, any JEPA-related components
2. How speaker conditioning is implemented (speaker embeddings, FiLM, cross-attention, etc.)
3. The training loop and loss functions
4. How the JEPA encoder/predictor/target encoder interact
5. Any config files showing model dimensions, layers, etc.
6. The inference/synthesis pipeline

Look for directories like lewm*, tts*, model*, src*, etc. Check for Python files with model definitions, training scripts, config files (yaml, json, toml). Be very thorough — check all subdirectories.


Alright, I've read through your full codebase. Here's exactly where speaker signal gets lost:

## The Problem — 3 Layers Deep

**1. Speaker embedding is added to text, not audio**

```python
# model.py — _add_speaker()
spk = self.speaker_embed(speaker_id).unsqueeze(1)  # [B, 1, d]
text_emb = text_emb + spk  # added to text only
```

The predictor sees speaker info through cross-attention to text — but that's indirect. The **audio encoder** and **EMA target encoder** never see speaker identity. So the target embeddings (`mu_ema`) that drive the main loss are **speaker-agnostic**. The model is being trained to predict targets that don't encode who's speaking.

**2. EMA target = mu_ema, which has no speaker signal**

```python
with torch.no_grad():
    _, mu_ema, _ = self.ema_audio_encoder(mel, mel_mask)
    target_emb = mu_ema
```

The audio encoder is a CNN + transformer that processes raw mel. It *could* learn speaker info, but nothing in the loss incentivizes it. The prediction loss is:

```
MSE(predicted[t], ema_target[t+1])
```

If `ema_target` has already averaged out speaker detail (which EMA + L2 will do), the predictor learns to match that speaker-less target. Adding speaker to the text side doesn't help — the predictor will learn to *ignore* it because the targets don't need it.

**3. Gaussian bottleneck further erases speaker**

```python
z = mu + eps * exp(0.5 * logvar)  # reparameterization
```

The KL loss pushes toward N(0,1). Speaker identity is a **global, time-invariant** signal — exactly the kind of thing KL regularization removes first, because it's "redundant" information that inflates KL across every frame.

## The Fix

Three changes, in order of impact:

### Fix 1: Inject speaker into the audio encoder output (biggest impact)

The target embeddings need speaker info so the predictor is forced to use it:

```python
# In forward(), after EMA encoding:
target_emb = mu_ema  # [B, T, d]
if self.speaker_embed is not None and speaker_id is not None:
    spk = self.speaker_embed(speaker_id).unsqueeze(1)  # [B, 1, d]
    target_emb = target_emb + spk  # now targets are speaker-specific
```

And same for the online encoder output `z` before it goes into the predictor input. Now the prediction loss *requires* speaker discrimination — two speakers saying the same word produce different targets.

### Fix 2: Reduce KL pressure on speaker dimensions

Your free_bits=1.0 is reasonable, but consider **splitting the latent**:

```python
# Reserve first 32 dims for speaker (no KL), rest for content
z_speaker = mu[:, :, :32]   # no sampling, no KL
z_content = reparameterize(mu[:, :, 32:], logvar[:, :, 32:])
z = torch.cat([z_speaker, z_content], dim=-1)
```

Or simpler — just raise free_bits to 2.0-4.0. Your current setup penalizes the model for encoding anything the Gaussian prior doesn't expect, and speaker identity is exactly that.

### Fix 3: Add a speaker consistency loss

After the audio encoder, verify it preserves speaker:

```python
# Pool across time
z_pooled = z.mean(dim=1)  # [B, d]
spk_pred = self.speaker_classifier(z_pooled)  # [B, n_speakers]
spk_loss = F.cross_entropy(spk_pred, speaker_id)
```

This is a 1-layer classifier, negligible params. Weight it at ~1.0. It directly tells the encoder "keep speaker information in the latent."

## Priority

**Fix 1 alone** will probably solve 80% of your problem. The targets are the anchor of the whole system — if they're speaker-agnostic, nothing downstream can recover it. Fix 2 and 3 are insurance.

Want me to implement these changes in `model.py`?


yeah implement ..


1	"""
2	LeWM TTS Model — JEPA-based Text-to-Speech
3
4	Architecture:
5	  - TextEncoder: Character-level transformer for Hindi/Devanagari text
6	  - AudioEncoder: 1D CNN + Transformer that encodes mel spectrograms to embeddings
7	  - JEPAPredictor: Transformer that predicts next audio embedding given text + audio context
8	  - Two losses: next-embedding prediction (MSE) + Gaussian regularizer (KL to N(0,I))
9	"""
10
11	import math
12	import copy
13	import torch
14	import torch.nn as nn
15	import torch.nn.functional as F
16	import torchaudio
17
18
19	# ─── Text Encoder ───────────────────────────────────────────────────────────
20
21	class TextEncoder(nn.Module):
22	    """Character-level transformer encoder for Hindi text."""
23
24	    def __init__(self, vocab_size=256, d_model=256, nhead=4, num_layers=4, dropout=0.1):
25	        super().__init__()
26	        self.d_model = d_model
27
28	        # Character embedding (covers all Unicode Hindi chars via byte-level encoding)
29	        self.char_embed = nn.Embedding(vocab_size, d_model)
30	        self.pos_embed = SinusoidalPositionalEncoding(d_model, max_len=2048)
31	        self.dropout = nn.Dropout(dropout)
32
33	        encoder_layer = nn.TransformerEncoderLayer(
34	            d_model=d_model, nhead=nhead,
35	            dim_feedforward=d_model * 4,
36	            dropout=dropout, batch_first=True,
37	            activation="gelu",
38	        )
39	        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
40	        self.proj = nn.Linear(d_model, d_model)
41
42	    def forward(self, text_tokens, text_mask=None):
43	        """
44	        Args:
45	            text_tokens: [B, T_text] long tensor of byte-level token IDs
46	            text_mask: [B, T_text] bool, True = padding
47	        Returns:
48	            text_emb: [B, T_text, d_model]
49	        """
50	        x = self.char_embed(text_tokens) * math.sqrt(self.d_model)
51	        x = self.pos_embed(x)
52	        x = self.dropout(x)
53	        x = self.transformer(x, src_key_padding_mask=text_mask)
54	        x = self.proj(x)
55	        return x
56
57
58	# ─── Audio Encoder ──────────────────────────────────────────────────────────
59
60	class AudioEncoder(nn.Module):
61	    """
62	    Encodes mel spectrograms into a sequence of audio embeddings.
63	    Uses 1D CNN for downsampling + Transformer for contextualization.
64	    """
65
66	    def __init__(self, n_mels=100, d_model=256, nhead=4, num_layers=4,
67	                 downsample_factor=4, dropout=0.1):
68	        super().__init__()
69	        self.d_model = d_model
70	        self.downsample_factor = downsample_factor
71
72	        # CNN downsampling: reduces temporal resolution by downsample_factor
73	        if downsample_factor == 2:
74	            self.conv_pre = nn.Sequential(
75	                nn.Conv1d(n_mels, d_model, kernel_size=7, padding=3),
76	                nn.GELU(),
77	                nn.Conv1d(d_model, d_model, kernel_size=4, stride=2, padding=1),
78	                nn.GELU(),
79	            )  # Total stride = 2
80	        else:
81	            self.conv_pre = nn.Sequential(
82	                nn.Conv1d(n_mels, d_model, kernel_size=7, padding=3),
83	                nn.GELU(),
84	                nn.Conv1d(d_model, d_model, kernel_size=4, stride=2, padding=1),
85	                nn.GELU(),
86	                nn.Conv1d(d_model, d_model, kernel_size=4, stride=2, padding=1),
87	                nn.GELU(),
88	            )  # Total stride = 4
89
90	        self.pos_embed = SinusoidalPositionalEncoding(d_model, max_len=4096)
91	        self.dropout = nn.Dropout(dropout)
92
93	        encoder_layer = nn.TransformerEncoderLayer(
94	            d_model=d_model, nhead=nhead,
95	            dim_feedforward=d_model * 4,
96	            dropout=dropout, batch_first=True,
97	            activation="gelu",
98	        )
99	        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
100
101	        # Project to embedding with mu and logvar for Gaussian regularizer
102	        self.proj_mu = nn.Linear(d_model, d_model)
103	        self.proj_logvar = nn.Linear(d_model, d_model)
104
105	    def forward(self, mel, mel_mask=None):
106	        """
107	        Args:
108	            mel: [B, n_mels, T_mel] — log mel spectrogram
109	            mel_mask: [B, T_mel] bool, True = padding (before downsampling)
110	        Returns:
111	            z: [B, T_down, d_model] — sampled latent embeddings
112	            mu: [B, T_down, d_model]
113	            logvar: [B, T_down, d_model]
114	        """
115	        # CNN encoding: [B, n_mels, T] -> [B, d_model, T//4]
116	        x = self.conv_pre(mel)
117	        x = x.transpose(1, 2)  # [B, T_down, d_model]
118
119	        # Downsample mask if provided
120	        if mel_mask is not None:
121	            # Approximate downsampled mask
122	            T_down = x.shape[1]
123	            mel_mask = mel_mask[:, ::self.downsample_factor][:, :T_down]
124
125	        x = self.pos_embed(x)
126	        x = self.dropout(x)
127	        x = self.transformer(x, src_key_padding_mask=mel_mask)
128
129	        mu = self.proj_mu(x)
130	        logvar = self.proj_logvar(x)
131
132	        # Reparameterization trick
133	        if self.training:
134	            std = torch.exp(0.5 * logvar)
135	            eps = torch.randn_like(std)
136	            z = mu + eps * std
137	        else:
138	            z = mu
139
140	        return z, mu, logvar
141
142
143	# ─── JEPA Predictor ─────────────────────────────────────────────────────────
144
145	class JEPAPredictor(nn.Module):
146	    """
147	    Predicts next audio embedding given:
148	    - Text context (from TextEncoder)
149	    - Previous audio embeddings (from AudioEncoder)
150
151	    Uses cross-attention to condition on text, and causal self-attention
152	    over the audio embedding sequence.
153	    """
154
155	    def __init__(self, d_model=256, nhead=4, num_layers=6, dropout=0.1):
156	        super().__init__()
157	        self.d_model = d_model
158
159	        self.audio_input_proj = nn.Linear(d_model, d_model)
160	        self.pos_embed = SinusoidalPositionalEncoding(d_model, max_len=4096)
161	        self.dropout = nn.Dropout(dropout)
162
163	        # Transformer decoder with causal self-attention + cross-attention to text
164	        decoder_layer = nn.TransformerDecoderLayer(
165	            d_model=d_model, nhead=nhead,
166	            dim_feedforward=d_model * 4,
167	            dropout=dropout, batch_first=True,
168	            activation="gelu",
169	        )
170	        self.transformer = nn.TransformerDecoder(decoder_layer, num_layers=num_layers)
171
172	        # Output projection to predict next embedding
173	        self.output_proj = nn.Sequential(
174	            nn.Linear(d_model, d_model),
175	            nn.GELU(),
176	            nn.Linear(d_model, d_model),
177	        )
178
179	    def forward(self, audio_emb, text_emb, audio_mask=None, text_mask=None):
180	        """
181	        Args:
182	            audio_emb: [B, T_audio, d_model] — audio embeddings (from encoder)
183	            text_emb: [B, T_text, d_model] — text embeddings
184	            audio_mask: [B, T_audio] bool, True = padding
185	            text_mask: [B, T_text] bool, True = padding
186	        Returns:
187	            predicted: [B, T_audio, d_model] — predicted next embeddings
188	        """
189	        x = self.audio_input_proj(audio_emb)
190	        x = self.pos_embed(x)
191	        x = self.dropout(x)
192
193	        # Causal mask: each position can only attend to previous positions
194	        T = x.shape[1]
195	        causal_mask = torch.triu(
196	            torch.ones(T, T, device=x.device, dtype=torch.bool), diagonal=1
197	        )
198
199	        predicted = self.transformer(
200	            x, text_emb,
201	            tgt_mask=causal_mask,
202	            tgt_key_padding_mask=audio_mask,
203	            memory_key_padding_mask=text_mask,
204	        )
205
206	        predicted = self.output_proj(predicted)
207	        return predicted
208
209	    def _cached_mha(self, mha, q_input, kv_input, cache_k, cache_v):
210	        """Run multi-head attention with KV cache.
211	        q_input: [B, 1, d], kv_input: [B, 1, d] (new token for K,V)
212	        cache_k, cache_v: [B, T_prev, d] projected keys/values
213	        Returns: output [B, 1, d], updated cache_k, updated cache_v
214	        """
215	        d = mha.embed_dim
216	        nhead = mha.num_heads
217	        head_dim = d // nhead
218	        B = q_input.shape[0]
219
220	        # Project Q from q_input, K/V from kv_input
221	        # in_proj_weight is [3d, d] = [Wq; Wk; Wv]
222	        Wq, Wk, Wv = mha.in_proj_weight.chunk(3, dim=0)
223	        bq, bk, bv = mha.in_proj_bias.chunk(3, dim=0)
224
225	        q = F.linear(q_input, Wq, bq)        # [B, 1, d]
226	        k_new = F.linear(kv_input, Wk, bk)    # [B, 1, d]
227	        v_new = F.linear(kv_input, Wv, bv)    # [B, 1, d]
228
229	        # Append to cache
230	        if cache_k is not None and cache_k.shape[1] > 0:
231	            k = torch.cat([cache_k, k_new], dim=1)
232	            v = torch.cat([cache_v, v_new], dim=1)
233	        else:
234	            k = k_new
235	            v = v_new
236
237	        # Multi-head attention
238	        q = q.view(B, 1, nhead, head_dim).transpose(1, 2)      # [B, nhead, 1, head_dim]
239	        k_mh = k.view(B, -1, nhead, head_dim).transpose(1, 2)  # [B, nhead, T, head_dim]
240	        v_mh = v.view(B, -1, nhead, head_dim).transpose(1, 2)
241
242	        attn = torch.matmul(q, k_mh.transpose(-2, -1)) / (head_dim ** 0.5)
243	        attn = F.softmax(attn, dim=-1)
244	        out = torch.matmul(attn, v_mh)  # [B, nhead, 1, head_dim]
245
246	        out = out.transpose(1, 2).contiguous().view(B, 1, d)
247	        out = F.linear(out, mha.out_proj.weight, mha.out_proj.bias)
248
249	        return out, k, v
250
251	    def _cached_cross_attn(self, mha, q_input, memory, cache_k, cache_v, memory_mask=None):
252	        """Cross-attention with cached memory K,V (computed once).
253	        memory_mask: [B, T_text] bool, True = padding (applied as -inf)."""
254	        d = mha.embed_dim
255	        nhead = mha.num_heads
256	        head_dim = d // nhead
257	        B = q_input.shape[0]
258
259	        Wq, Wk, Wv = mha.in_proj_weight.chunk(3, dim=0)
260	        bq, bk, bv = mha.in_proj_bias.chunk(3, dim=0)
261
262	        q = F.linear(q_input, Wq, bq)  # [B, 1, d]
263
264	        if cache_k is None:
265	            k = F.linear(memory, Wk, bk)  # [B, T_text, d]
266	            v = F.linear(memory, Wv, bv)
267	        else:
268	            k, v = cache_k, cache_v
269
270	        q = q.view(B, 1, nhead, head_dim).transpose(1, 2)
271	        k_mh = k.view(B, -1, nhead, head_dim).transpose(1, 2)
272	        v_mh = v.view(B, -1, nhead, head_dim).transpose(1, 2)
273
274	        attn = torch.matmul(q, k_mh.transpose(-2, -1)) / (head_dim ** 0.5)
275	        if memory_mask is not None:
276	            # memory_mask: [B, T_text] → [B, 1, 1, T_text]
277	            attn = attn.masked_fill(memory_mask.unsqueeze(1).unsqueeze(2), float('-inf'))
278	        attn = F.softmax(attn, dim=-1)
279	        out = torch.matmul(attn, v_mh)
280
281	        out = out.transpose(1, 2).contiguous().view(B, 1, d)
282	        out = F.linear(out, mha.out_proj.weight, mha.out_proj.bias)
283
284	        return out, k, v
285
286	    def init_cache(self, num_layers):
287	        """Initialize empty KV cache."""
288	        return {
289	            'self_k': [None] * num_layers,
290	            'self_v': [None] * num_layers,
291	            'cross_k': [None] * num_layers,
292	            'cross_v': [None] * num_layers,
293	        }
294
295	    def inference_step(self, new_emb, text_emb, step_idx, cache, text_mask=None):
296	        """Process one token with KV cache.
297	        new_emb: [B, 1, d_model] — raw audio embedding (not projected)
298	        text_mask: [B, T_text] bool, True = padding
299	        Returns: predicted [B, 1, d_model], updated cache
300	        """
301	        x = self.audio_input_proj(new_emb)
302	        # Add positional encoding for this specific position
303	        x = x + self.pos_embed.pe[:, step_idx:step_idx + 1]
304
305	        for i, layer in enumerate(self.transformer.layers):
306	            # Self-attention with cache (post-norm)
307	            sa_out, cache['self_k'][i], cache['self_v'][i] = self._cached_mha(
308	                layer.self_attn, x, x, cache['self_k'][i], cache['self_v'][i]
309	            )
310	            x = layer.norm1(x + layer.dropout1(sa_out))
311
312	            # Cross-attention with cached text KV
313	            ca_out, cache['cross_k'][i], cache['cross_v'][i] = self._cached_cross_attn(
314	                layer.multihead_attn, x, text_emb, cache['cross_k'][i], cache['cross_v'][i],
315	                memory_mask=text_mask,
316	            )
317	            x = layer.norm2(x + layer.dropout2(ca_out))
318
319	            # FFN
320	            ff_out = layer.linear2(layer.dropout(layer.activation(layer.linear1(x))))
321	            x = layer.norm3(x + layer.dropout3(ff_out))
322
323	        return self.output_proj(x), cache
324
325
326	# ─── Positional Encoding ────────────────────────────────────────────────────
327
328	class SinusoidalPositionalEncoding(nn.Module):
329	    def __init__(self, d_model, max_len=8192):
330	        super().__init__()
331	        pe = torch.zeros(max_len, d_model)
332	        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
333	        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
334	        pe[:, 0::2] = torch.sin(position * div_term)
335	        pe[:, 1::2] = torch.cos(position * div_term)
336	        self.register_buffer("pe", pe.unsqueeze(0))  # [1, max_len, d_model]
337
338	    def forward(self, x):
339	        return x + self.pe[:, :x.shape[1]]
340
341
342	# ─── Mel Decoder ────────────────────────────────────────────────────────────
343
344	class ResConvBlock(nn.Module):
345	    """Residual conv block for mel decoder."""
346
347	    def __init__(self, channels, kernel_size=3):
348	        super().__init__()
349	        self.net = nn.Sequential(
350	            nn.Conv1d(channels, channels, kernel_size, padding=kernel_size // 2),
351	            nn.GELU(),
352	            nn.Conv1d(channels, channels, kernel_size, padding=kernel_size // 2),
353	        )
354
355	    def forward(self, x):
356	        return x + self.net(x)
357
358
359	class MelDecoder(nn.Module):
360	    """Decodes latent embeddings back to mel spectrograms.
361	    Improved architecture with residual blocks for sharper mel output."""
362
363	    def __init__(self, d_model=256, n_mels=100, upsample_factor=4):
364	        super().__init__()
365	        hidden = d_model * 2  # 512 — more capacity
366
367	        self.proj = nn.Sequential(
368	            nn.Linear(d_model, hidden),
369	            nn.GELU(),
370	        )
371	        if upsample_factor == 2:
372	            self.upsample = nn.Sequential(
373	                nn.ConvTranspose1d(hidden, hidden, kernel_size=4, stride=2, padding=1),
374	                nn.GELU(),
375	                ResConvBlock(hidden, kernel_size=3),
376	            )
377	        else:
378	            self.upsample = nn.Sequential(
379	                # 4x upsample (2x + 2x)
380	                nn.ConvTranspose1d(hidden, hidden, kernel_size=4, stride=2, padding=1),
381	                nn.GELU(),
382	                ResConvBlock(hidden, kernel_size=3),
383	                nn.ConvTranspose1d(hidden, hidden, kernel_size=4, stride=2, padding=1),
384	                nn.GELU(),
385	                ResConvBlock(hidden, kernel_size=3),
386	            )
387	        # Refinement + projection to mel bins
388	        self.out = nn.Sequential(
389	            nn.Conv1d(hidden, hidden, kernel_size=7, padding=3),
390	            nn.GELU(),
391	            ResConvBlock(hidden, kernel_size=5),
392	            nn.Conv1d(hidden, n_mels, kernel_size=7, padding=3),
393	        )
394
395	    def forward(self, z):
396	        """z: [B, T_down, d_model] → mel: [B, n_mels, T_up]"""
397	        x = self.proj(z)
398	        x = x.transpose(1, 2)  # [B, hidden, T_down]
399	        x = self.upsample(x)   # [B, hidden, T_up]
400	        return self.out(x)      # [B, n_mels, T_up]
401
402
403	# ─── Multi-Resolution Spectral Loss ─────────────────────────────────────────
404
405	class MultiResolutionSpectralLoss(nn.Module):
406	    """Computes spectral convergence + log magnitude loss at multiple STFT resolutions.
407	    This forces the mel decoder to produce sharp, well-defined spectral structure
408	    rather than blurry averages that Vocos can't decode properly."""
409
410	    def __init__(self, resolutions=((64, 16, 64), (128, 32, 128), (256, 64, 256))):
411	        super().__init__()
412	        # Each resolution is (n_fft, hop_length, win_length)
413	        self.resolutions = resolutions
414
415	    def _stft_loss(self, pred, target, n_fft, hop_length, win_length):
416	        """Compute spectral convergence + log magnitude loss for one resolution."""
417	        window = torch.hann_window(win_length, device=pred.device)
418
419	        # pred/target: [B, n_mels, T] — treat as 1D signal per mel bin
420	        # Flatten mel bins into batch: [B*n_mels, T]
421	        B, M, T = pred.shape
422	        pred_flat = pred.reshape(B * M, T)
423	        target_flat = target.reshape(B * M, T)
424
425	        pred_stft = torch.stft(pred_flat, n_fft, hop_length, win_length,
426	                               window=window, return_complex=True)
427	        target_stft = torch.stft(target_flat, n_fft, hop_length, win_length,
428	                                 window=window, return_complex=True)
429
430	        pred_mag = pred_stft.abs()
431	        target_mag = target_stft.abs()
432
433	        # Spectral convergence: Frobenius norm of difference / Frobenius norm of target
434	        sc_loss = torch.norm(target_mag - pred_mag, p="fro") / (torch.norm(target_mag, p="fro") + 1e-8)
435
436	        # Log magnitude loss
437	        log_mag_loss = F.l1_loss(
438	            torch.log(pred_mag + 1e-7),
439	            torch.log(target_mag + 1e-7),
440	        )
441
442	        return sc_loss + log_mag_loss
443
444	    def forward(self, pred, target):
445	        """pred, target: [B, n_mels, T]"""
446	        loss = torch.tensor(0.0, device=pred.device)
447	        count = 0
448	        for n_fft, hop, win in self.resolutions:
449	            # Skip if T is too short for this resolution
450	            if pred.shape[2] >= n_fft:
451	                loss = loss + self._stft_loss(pred, target, n_fft, hop, win)
452	                count += 1
453	        if count > 0:
454	            loss = loss / count
455	        return loss
456
457
458	# ─── Full LeWM TTS Model ────────────────────────────────────────────────────
459
460	class LeWMTTS(nn.Module):
461	    """
462	    Complete LeWM TTS model.
463
464	    Training: text + mel → audio embeddings → predict next embedding
465	    Two losses:
466	        1. Next-embedding prediction: MSE(predicted[t], target[t+1])
467	        2. Gaussian regularizer: KL(q(z|x) || N(0, I))
468	    """
469
470	    def __init__(self, config):
471	        super().__init__()
472	        d_model = config.get("d_model", 256)
473	        nhead = config.get("nhead", 4)
474	        n_mels = config.get("n_mels", 100)
475	        text_vocab = config.get("text_vocab_size", 256)
476	        text_layers = config.get("text_encoder_layers", 4)
477	        audio_layers = config.get("audio_encoder_layers", 4)
478	        predictor_layers = config.get("predictor_layers", 6)
479	        dropout = config.get("dropout", 0.1)
480	        self.kl_weight = config.get("kl_weight", 0.01)
481	        n_speakers = config.get("n_speakers", 1)
482
483	        self.text_encoder = TextEncoder(
484	            vocab_size=text_vocab, d_model=d_model, nhead=nhead,
485	            num_layers=text_layers, dropout=dropout,
486	        )
487	        downsample_factor = config.get("downsample_factor", 4)
488	        self.audio_encoder = AudioEncoder(
489	            n_mels=n_mels, d_model=d_model, nhead=nhead,
490	            num_layers=audio_layers, downsample_factor=downsample_factor,
491	            dropout=dropout,
492	        )
493	        self.predictor = JEPAPredictor(
494	            d_model=d_model, nhead=nhead,
495	            num_layers=predictor_layers, dropout=dropout,
496	        )
497
498	        # Mel decoder: reconstructs mel from latent embeddings
499	        # Needed so the latent space is grounded to actual audio
500	        self.mel_decoder = MelDecoder(d_model=d_model, n_mels=n_mels, upsample_factor=downsample_factor)
501	        self.recon_weight = config.get("recon_weight", 1.0)
502
503	        # EMA target encoder — core JEPA component
504	        # Provides slowly-moving prediction targets to prevent representation collapse
505	        self.ema_audio_encoder = copy.deepcopy(self.audio_encoder)
506	        for p in self.ema_audio_encoder.parameters():
507	            p.requires_grad = False
508	        self.ema_decay = config.get("ema_decay", 0.998)
509
510	        # Speaker embedding (for multi-speaker)
511	        self.n_speakers = n_speakers
512	        if n_speakers > 1:
513	            self.speaker_embed = nn.Embedding(n_speakers, d_model)
514	        else:
515	            self.speaker_embed = None
516
517	        # Learnable start-of-sequence embedding for AR inference
518	        self.start_emb = nn.Parameter(torch.randn(1, 1, d_model) * 0.02)
519
520	        # Prediction weight — primary JEPA loss
521	        self.pred_weight = config.get("pred_weight", 10.0)
522
523	        # Input noise injection to bridge train/inference gap
524	        self.input_noise = config.get("input_noise", 0.0)
525
526	        # Free-bits KL: allow KL up to this threshold without penalty
527	        self.free_bits = config.get("free_bits", 1.0)
528
529	        # Multi-resolution spectral loss for sharper mel output
530	        self.spectral_loss = MultiResolutionSpectralLoss()
531	        self.spectral_weight = config.get("spectral_weight", 0.5)
532
533	        self.config = config
534
535	    def _add_speaker(self, text_emb, speaker_id):
536	        """Add speaker embedding to text embeddings."""
537	        if self.speaker_embed is not None and speaker_id is not None:
538	            spk = self.speaker_embed(speaker_id).unsqueeze(1)  # [B, 1, d]
539	            text_emb = text_emb + spk
540	        return text_emb
541
542	    @torch.no_grad()
543	    def update_ema(self):
544	        """Update EMA target encoder weights."""
545	        for p_online, p_ema in zip(self.audio_encoder.parameters(),
546	                                   self.ema_audio_encoder.parameters()):
547	            p_ema.data.mul_(self.ema_decay).add_(p_online.data, alpha=1 - self.ema_decay)
548
549	    def forward(self, mel, text_tokens, mel_mask=None, text_mask=None, speaker_id=None):
550	        """
551	        Args:
552	            mel: [B, n_mels, T_mel]
553	            text_tokens: [B, T_text]
554	            mel_mask: [B, T_mel] bool
555	            text_mask: [B, T_text] bool
556	        Returns:
557	            loss_dict with 'prediction_loss', 'kl_loss', 'total_loss'
558	        """
559	        # Encode text + speaker conditioning
560	        text_emb = self.text_encoder(text_tokens, text_mask)
561	        text_emb = self._add_speaker(text_emb, speaker_id)
562
563	        # Online encoder → latent embeddings (used for predictor input + reconstruction)
564	        z, mu, logvar = self.audio_encoder(mel, mel_mask)
565	        B = z.shape[0]
566	        T_down = z.shape[1]
567
568	        # EMA target encoder → stable prediction targets (no gradient)
569	        # Force eval mode so dropout is disabled — EMA targets must be deterministic
570	        with torch.no_grad():
571	            self.ema_audio_encoder.eval()
572	            _, mu_ema, _ = self.ema_audio_encoder(mel, mel_mask)
573	            self.ema_audio_encoder.train()
574	            target_emb = mu_ema  # [B, T_down, d] — all EMA targets
575
576	        # Prepend learnable start embedding to input sequence
577	        # Input:  [start_emb, z[0], z[1], ..., z[T-2]]  → T_down positions
578	        # Target: [z_ema[0],  z_ema[1], ..., z_ema[T-1]] → T_down positions
579	        # This trains the model to predict the first audio embedding from start_emb + text
580	        start = self.start_emb.expand(B, -1, -1)  # [B, 1, d]
581	        input_emb = torch.cat([start, z[:, :-1]], dim=1)  # [B, T_down, d]
582
583	        # Input noise: only on positions 1+ (real embeddings), not start_emb
584	        # Simulates the imperfect predictions seen during AR inference
585	        if self.training and self.input_noise > 0:
586	            noise = torch.zeros_like(input_emb)
587	            noise[:, 1:] = torch.randn(B, T_down - 1, input_emb.shape[-1],
588	                                        device=input_emb.device) * self.input_noise
589	            input_emb = input_emb + noise
590
591	        # Adjust masks for downsampled sequence
592	        if mel_mask is not None:
593	            ds_mask = mel_mask[:, ::self.audio_encoder.downsample_factor][:, :T_down]
594	            # Predictor input mask: start_emb (pos 0) is always unmasked
595	            start_mask = torch.zeros(B, 1, dtype=torch.bool, device=z.device)
596	            pred_mask = torch.cat([start_mask, ds_mask[:, :-1]], dim=1)  # [B, T_down]
597	            # Loss mask: follows original ds_mask for target positions
598	            loss_mask = ds_mask
599	        else:
600	            pred_mask = None
601	            loss_mask = None
602
603	        predicted = self.predictor(input_emb, text_emb, pred_mask, text_mask)
604
605	        # ─── Loss 1: Next-embedding prediction (MSE) against EMA targets ───
606	        if loss_mask is not None:
607	            valid = (~loss_mask).unsqueeze(-1)  # [B, T_down, 1]
608	            prediction_loss = F.mse_loss(
609	                predicted * valid, target_emb * valid, reduction="sum"
610	            ) / (valid.sum() * predicted.shape[-1] + 1e-8)
611	        else:
612	            prediction_loss = F.mse_loss(predicted, target_emb)
613
614	        # ─── Loss 2: Free-bits KL — only penalize KL above threshold ───
615	        # Prevents both KL collapse (too low) and KL explosion (too high)
616	        kl_per_dim = -0.5 * (1 + logvar - mu.pow(2) - logvar.exp())  # [B, T, d]
617	        kl_free = torch.clamp(kl_per_dim - self.free_bits, min=0.0)
618	        if mel_mask is not None:
619	            T_down = mu.shape[1]
620	            ds_mask = mel_mask[:, ::self.audio_encoder.downsample_factor][:, :T_down]
621	            valid_kl = (~ds_mask).unsqueeze(-1).float()
622	            kl_loss = (kl_free * valid_kl).sum() / (valid_kl.sum() * mu.shape[-1] + 1e-8)
623	        else:
624	            kl_loss = kl_free.mean()
625
626	        # ─── Loss 3: Mel reconstruction ───
627	        mel_recon = self.mel_decoder(z)  # [B, n_mels, T_up]
628	        T_mel = mel.shape[2]
629	        T_recon = mel_recon.shape[2]
630	        if T_recon > T_mel:
631	            mel_recon = mel_recon[:, :, :T_mel]
632	        elif T_recon < T_mel:
633	            mel = mel[:, :, :T_recon]
634	            if mel_mask is not None:
635	                mel_mask = mel_mask[:, :T_recon]
636
637	        if mel_mask is not None:
638	            valid_mel = (~mel_mask).unsqueeze(1).float()  # [B, 1, T]
639	            recon_loss = F.l1_loss(
640	                mel_recon * valid_mel, mel * valid_mel, reduction="sum"
641	            ) / (valid_mel.sum() * mel.shape[1] + 1e-8)
642	        else:
643	            recon_loss = F.l1_loss(mel_recon, mel)
644
645	        # ─── Loss 4: Multi-resolution spectral loss ───
646	        if mel_mask is not None:
647	            spectral_loss = self.spectral_loss(mel_recon * valid_mel, mel * valid_mel)
648	        else:
649	            spectral_loss = self.spectral_loss(mel_recon, mel)
650
651	        total_loss = (self.pred_weight * prediction_loss
652	                      + self.kl_weight * kl_loss
653	                      + self.recon_weight * recon_loss
654	                      + self.spectral_weight * spectral_loss)
655
656	        return {
657	            "total_loss": total_loss,
658	            "prediction_loss": prediction_loss,
659	            "kl_loss": kl_loss,
660	            "recon_loss": recon_loss,
661	            "spectral_loss": spectral_loss,
662	        }
663
664	    def encode_audio(self, mel):
665	        """Encode mel to latent embeddings (for inference)."""
666	        z, mu, _ = self.audio_encoder(mel)
667	        return mu  # Use mean at inference
668
669	    def predict_next(self, audio_emb, text_emb, text_mask=None):
670	        """Predict next audio embedding autoregressively."""
671	        predicted = self.predictor(audio_emb, text_emb, text_mask=text_mask)
672	        return predicted[:, -1:]  # Only the last prediction
673
674	    def init_ar_cache(self):
675	        """Initialize KV cache for autoregressive inference."""
676	        num_layers = len(self.predictor.transformer.layers)
677	        return self.predictor.init_cache(num_layers)
678
679	    def predict_next_cached(self, new_emb, text_emb, step_idx, cache, text_mask=None):
680	        """Predict next embedding using KV cache. O(1) per step instead of O(n)."""
681	        return self.predictor.inference_step(new_emb, text_emb, step_idx, cache, text_mask=text_mask)
682
683
684	def count_parameters(model):
685	    return sum(p.numel() for p in model.parameters() if p.requires_grad)
686
687
688	def build_model(config=None):
689	    if config is None:
690	        config = {
691	            "d_model": 256,
692	            "nhead": 4,
693	            "n_mels": 100,
694	            "text_vocab_size": 256,
695	            "text_encoder_layers": 4,
696	            "audio_encoder_layers": 4,
697	            "predictor_layers": 6,
698	            "dropout": 0.1,
699	            "kl_weight": 0.05,
700	            "pred_weight": 10.0,
701	            "ema_decay": 0.998,
702	            "free_bits": 1.0,
703	            "input_noise": 0.0,
704	        }
705	    model = LeWMTTS(config)
706	    print(f"LeWM TTS model: {count_parameters(model)/1e6:.2f}M parameters")
707	    return model, config
708

"""
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)

        # 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

    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, 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)

        # Online encoder → latent embeddings (used for predictor input + reconstruction)
        z, mu, logvar = self.audio_encoder(mel, mel_mask)
        B = z.shape[0]
        T_down = z.shape[1]

        # EMA target encoder → stable prediction targets (no gradient)
        # Force eval mode so dropout is disabled — EMA targets must be deterministic
        with torch.no_grad():
            self.ema_audio_encoder.eval()
            _, mu_ema, _ = self.ema_audio_encoder(mel, mel_mask)
            self.ema_audio_encoder.train()
            target_emb = mu_ema  # [B, T_down, d] — all EMA targets

        # Prepend learnable start embedding to input sequence
        # Input:  [start_emb, z[0], z[1], ..., z[T-2]]  → T_down positions
        # Target: [z_ema[0],  z_ema[1], ..., z_ema[T-1]] → T_down positions
        # This trains the model to predict the first audio embedding from start_emb + text
        start = self.start_emb.expand(B, -1, -1)  # [B, 1, d]
        input_emb = torch.cat([start, z[:, :-1]], dim=1)  # [B, T_down, d]

        # Input noise: only on positions 1+ (real embeddings), not start_emb
        # Simulates the imperfect predictions seen during AR inference
        if self.training and self.input_noise > 0:
            noise = torch.zeros_like(input_emb)
            noise[:, 1:] = torch.randn(B, T_down - 1, input_emb.shape[-1],
                                        device=input_emb.device) * self.input_noise
            input_emb = input_emb + noise

        # Adjust masks for downsampled sequence
        if mel_mask is not None:
            ds_mask = mel_mask[:, ::self.audio_encoder.downsample_factor][:, :T_down]
            # Predictor input mask: start_emb (pos 0) is always unmasked
            start_mask = torch.zeros(B, 1, dtype=torch.bool, device=z.device)
            pred_mask = torch.cat([start_mask, ds_mask[:, :-1]], dim=1)  # [B, T_down]
            # Loss mask: follows original ds_mask for target positions
            loss_mask = ds_mask
        else:
            pred_mask = None
            loss_mask = None

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

        # ─── Loss 1: Next-embedding prediction (MSE) against EMA targets ───
        if loss_mask is not None:
            valid = (~loss_mask).unsqueeze(-1)  # [B, T_down, 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: Free-bits KL — only penalize KL above threshold ───
        # Prevents both KL collapse (too low) and KL explosion (too high)
        kl_per_dim = -0.5 * (1 + logvar - mu.pow(2) - logvar.exp())  # [B, T, d]
        kl_free = torch.clamp(kl_per_dim - self.free_bits, min=0.0)
        if mel_mask is not None:
            T_down = mu.shape[1]
            ds_mask = mel_mask[:, ::self.audio_encoder.downsample_factor][:, :T_down]
            valid_kl = (~ds_mask).unsqueeze(-1).float()
            kl_loss = (kl_free * valid_kl).sum() / (valid_kl.sum() * mu.shape[-1] + 1e-8)
        else:
            kl_loss = kl_free.mean()

        # ─── Loss 3: Mel reconstruction ───
        mel_recon = self.mel_decoder(z)  # [B, n_mels, T_up]
        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 ───
        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 = (self.pred_weight * 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.05,
            "pred_weight": 10.0,
            "ema_decay": 0.998,
            "free_bits": 1.0,
            "input_noise": 0.0,
        }
    model = LeWMTTS(config)
    print(f"LeWM TTS model: {count_parameters(model)/1e6:.2f}M parameters")
    return model, config


Now let me also check the training script to see how speaker_id flows through.