48	    @torch.no_grad()
49	    def synthesize(self, text, max_steps=300, temperature=0.0, speaker_id=None):
50	        """Generate speech from text using AR prediction with KV-cache."""
51	        text_tokens = self.text_to_tokens(text)
52	        text_emb = self.model.text_encoder(text_tokens)
53
54	        # Add speaker conditioning to text embeddings
55	        if self.model.speaker_embed is not None and speaker_id is not None:
56	            spk_id = torch.tensor([speaker_id], device=self.device)
57	            text_emb = self.model._add_speaker(text_emb, spk_id)
58	            # Cache speaker embedding for adding to predicted latents
59	            self._spk_emb = self.model.speaker_embed(spk_id).unsqueeze(1)  # [1, 1, d]
60	        else:
61	            self._spk_emb = None
62
63	        # Use learned start embedding instead of random noise
64	        start_emb = self.model.start_emb.to(self.device)
65	        all_embeddings = [start_emb]
66
67	        cache = self.model.init_ar_cache()
68	        next_emb, cache = self.model.predict_next_cached(start_emb, text_emb, 0, cache)
69	        if temperature > 0:
70	            next_emb = next_emb + torch.randn_like(next_emb) * temperature
71	        # Add speaker embedding to predicted latent (matches training where z = z + spk)
72	        emb_for_decode = next_emb + self._spk_emb if self._spk_emb is not None else next_emb
73	        all_embeddings.append(emb_for_decode)
74
75	        prev_norms = [next_emb.norm().item()]
76
77	        for step in range(1, max_steps):
78	            # Feed speaker-conditioned embedding as input (matches training input_emb)
79	            next_input = emb_for_decode
80	            next_emb, cache = self.model.predict_next_cached(next_input, text_emb, step, cache)
81	            if temperature > 0:
82	                next_emb = next_emb + torch.randn_like(next_emb) * temperature
83	            emb_for_decode = next_emb + self._spk_emb if self._spk_emb is not None else next_emb
84	            all_embeddings.append(emb_for_decode)
85
86	            norm = next_emb.norm().item()
87	            prev_norms.append(norm)
88	            if len(prev_norms) > 20:
89	                recent = prev_norms[-20:]
90	                if np.std(recent) < 0.01 * np.mean(recent) and step > 30:
91	                    break
92

    @torch.no_grad()
    def synthesize(self, text, max_steps=300, temperature=0.0, speaker_id=None):
        """Generate speech from text using AR prediction with KV-cache."""
        text_tokens = self.text_to_tokens(text)
        text_emb = self.model.text_encoder(text_tokens)

        # Add speaker conditioning to text embeddings
        if self.model.speaker_embed is not None and speaker_id is not None:
            spk_id = torch.tensor([speaker_id], device=self.device)
            text_emb = self.model._add_speaker(text_emb, spk_id)
            # Cache speaker embedding for adding to predicted latents
            self._spk_emb = self.model.speaker_embed(spk_id).unsqueeze(1)  # [1, 1, d]
        else:
            self._spk_emb = None

        # Use learned start embedding instead of random noise
        start_emb = self.model.start_emb.to(self.device)
        all_embeddings = [start_emb]

        cache = self.model.init_ar_cache()
        next_emb, cache = self.model.predict_next_cached(start_emb, text_emb, 0, cache)
        if temperature > 0:
            next_emb = next_emb + torch.randn_like(next_emb) * temperature
        # Add speaker embedding to predicted latent (matches training where z = z + spk)
        emb_for_decode = next_emb + self._spk_emb if self._spk_emb is not None else next_emb
        all_embeddings.append(emb_for_decode)

        prev_norms = [next_emb.norm().item()]

        for step in range(1, max_steps):
            # Feed speaker-conditioned embedding as input (matches training input_emb)
            next_input = emb_for_decode
            next_emb, cache = self.model.predict_next_cached(next_input, text_emb, step, cache)
            if temperature > 0:
                next_emb = next_emb + torch.randn_like(next_emb) * temperature
            emb_for_decode = next_emb + self._spk_emb if self._spk_emb is not None else next_emb
            all_embeddings.append(emb_for_decode)

            norm = next_emb.norm().item()
            prev_norms.append(norm)
            if len(prev_norms) > 20:
                recent = prev_norms[-20:]
                if np.std(recent) < 0.01 * np.mean(recent) and step > 30:
                    break


Good — inference adds speaker to embeddings before mel decode, matching training. The flow is consistent.

Now let me run the full test again with the `mu_raw` fix:


Bash

python3 -c "
import torch
from model import build_model

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': 2.0, 'input_noise': 0.1,
    'n_speakers': 15, 'speaker_weight': 1.0,
    'recon_weight': 1.0, 'spectral_weight': 0.5,
}
model, _ = build_model(config)
model.train()

B = 4
mel = torch.randn(B, 100, 400)
text = torch.randint(0, 256, (B, 30))
mel_mask = torch.zeros(B, 400, dtype=torch.bool)
text_mask = torch.zeros(B, 30, dtype=torch.bool)
spk = torch.tensor([0, 3, 7, 14])

# Test forward + backward
losses = model(mel, text, mel_mask, text_mask, speaker_id=spk)
losses['total_loss'].backward()

# Verify speaker classifier gradient flows to audio encoder (not speaker embed)
# The classifier should be classifying from raw mu, not speaker-injected z
ae_grad = model.audio_encoder.proj_mu.weight.grad.norm().item()
cls_grad = model.speaker_classifier.weight.grad.norm().item()
spk_grad = model.speaker_embed.weight.grad.norm().item()

print(f'AudioEncoder proj_mu grad: {ae_grad:.4f} (should be > 0)')
print(f'Speaker classifier grad:   {cls_grad:.4f} (should be > 0)')
print(f'Speaker embed grad:        {spk_grad:.4f} (should be > 0)')

for k, v in losses.items():
    print(f'{k}: {v.item():.4f}')

# Verify no speaker_id still works
model.zero_grad()
losses2 = model(mel, text, mel_mask, text_mask)
losses2['total_loss'].backward()
print(f'No-speaker backward OK, speaker_loss={losses2[\"speaker_loss\"].item():.4f}')

print('ALL OK')
" 2>&1

Test mu_raw fix


/home/ubuntu/.local/lib/python3.10/site-packages/torch/nn/modules/transformer.py:515: UserWarning: The PyTorch API of nested tensors is in prototype stage and will change in the near future. We recommend specifying layout=torch.jagged when constructing a nested tensor, as this layout receives active development, has better operator coverage, and works with torch.compile. (Triggered internally at /pytorch/aten/src/ATen/NestedTensorImpl.cpp:178.)
  output = torch._nested_tensor_from_mask(
LeWM TTS model: 24.01M parameters
AudioEncoder proj_mu grad: 4.1307 (should be > 0)
Speaker classifier grad:   3.8774 (should be > 0)
Speaker embed grad:        0.7417 (should be > 0)
total_loss: 21.2905
prediction_loss: 1.4223
kl_loss: 0.0012
recon_loss: 0.7987
spectral_loss: 6.5493
speaker_loss: 2.9938
No-speaker backward OK, speaker_loss=0.0000
ALL OK


All clean. Here's the summary:

## Current Model Size

| | Params | Notes |
|---|---|---|
| **Total (training)** | **24.01M** trainable + 3.99M EMA = 28M | Well under 80M budget |
| **Inference** | **17.48M** | No predictor, no EMA, no classifier |
| **On disk (fp16)** | ~63MB | State dict including EMA |
| **Inference on disk (fp16)** | ~35MB | Very edge-friendly |

You have **~52M of headroom** left within your 80M target. Room to scale up d_model to 384 or add layers if needed.

## Improvement Made

**Fixed the speaker classifier to use `mu_raw` (pre-injection) instead of `mu` (post-injection).**

The previous version classified speaker from `mu` which already had `spk_embed` added — so the classifier could cheat by just reading back the added embedding instead of forcing the encoder to learn speaker-discriminative features. Now it classifies from the raw encoder output, so the audio encoder itself is forced to retain speaker identity.

## Architecture Summary (with all fixes)

Speaker signal is now preserved through **3 independent mechanisms**:

1. **Target injection** — `target_emb = ema_output + spk_embed`, so prediction loss requires speaker knowledge
2. **Input injection** — `z = z + spk_embed`, so the predictor input carries speaker identity (matching inference)
3. **Classifier on raw mu** — cross-entropy on time-pooled encoder output, gradient flows directly into the audio encoder's weights


See how, from 0 to 10 points, how confident you are that this will work. I can give 10 speaker audios, and I will get the same similar audio in the inference.


Honestly? **4/10.**

Here's why:

## What will work (the +4)

- The speaker **classifier loss** will make the encoder retain speaker identity in latents — that's well-proven (SpeakerNet, GE2GE, etc.)
- The **target injection** will force the predictor to use speaker info — mathematically the loss can't go to zero without it
- The architecture is sound for multi-speaker — cross-attention + FiLM-style additive conditioning works

## What will likely break (the -6)

**1. Additive speaker embedding is too weak (biggest risk)**

You're doing `z = z + spk_embed` — a single 256-dim vector added globally. Speaker identity isn't a single offset. It's a nonlinear transformation of the entire latent space — pitch range, formant structure, speaking rate, breathiness. A learned 256-dim addition can shift the mean but can't reshape the distribution.

Models that actually do multi-speaker well (VITS, YourTTS, StyleTTS) use **FiLM conditioning** (scale + shift, not just shift) or **cross-attention to a speaker encoder**, not additive.

**2. 20 hours / 10 speakers = 2 hours per speaker**

That's thin. You need the model to generalize speaker characteristics from limited data. With additive conditioning and only 2h, you'll likely get "vaguely in the right pitch range" but not "sounds like that person."

**3. No reference encoder**

Right now speaker identity comes from a lookup table (`nn.Embedding(15, 256)`). At inference you pass `speaker_id=3` and hope. There's no mechanism to capture what speaker 3 actually *sounds* like from audio. A reference encoder (encode a few seconds of target speaker → speaker embedding) would be much more robust and is what production multi-speaker systems use.

**4. Mel decoder sees speaker only through the latent**

The mel decoder takes `z` (which has spk added) and must reconstruct speaker-specific spectrograms. But it's a small ConvTranspose stack — it may not have enough capacity to disentangle the speaker signal from the content signal in the latent.

**5. JEPA's L2 loss is inherently smoothing**

Even with speaker injection, the L2 prediction loss averages across variation. Speaker-specific details (micro-prosody, glottal characteristics) are high-frequency — exactly what L2 penalizes least.

## What would get you to 7-8/10

**Replace additive with FiLM** — scale and shift, not just shift:

```python
# Instead of: z = z + spk
gamma = self.spk_scale(speaker_emb)  # [B, 1, d]
beta = self.spk_shift(speaker_emb)   # [B, 1, d]
z = z * (1 + gamma) + beta
```

**Add a reference encoder** — even a small 3-layer CNN that takes a mel clip from the target speaker and produces the speaker embedding, instead of a lookup table. This generalizes better and captures what the speaker actually sounds like.

**Inject speaker into the mel decoder too** — not just the latent. FiLM condition the ResConvBlocks in MelDecoder on speaker.

**Use cosine similarity loss alongside L2** for the prediction loss — it preserves direction (speaker identity) while L2 only cares about distance.

Want me to implement the FiLM conditioning + mel decoder speaker injection? That's probably the highest-impact change to get you closer to 7-8.


Now I will give every speaker around 10 to 20 hours.


That changes things. 10-20 hours per speaker is **a lot** — that's what single-speaker TTS models train on successfully.

Revised: **6/10.**

The +2 is because data scarcity is no longer a problem. 10-20h is enough for the model to memorize speaker characteristics even with a weak conditioning mechanism.

But the fundamental architecture issues remain:

- **Additive embedding** still can't reshape the latent distribution, it can only shift it. With enough data the model will learn to *work around* this limitation, but it's fighting the architecture
- **L2 smoothing** still washes out fine speaker detail
- **No FiLM in mel decoder** — the decoder still has to recover speaker from a shifted latent

The thing is — with 10-20h per speaker, even a mediocre architecture will produce *something*. The question is whether "something" meets your bar for "same similar audio."

If your bar is "clearly the right speaker, recognizable voice" — 6/10, probably gets there.

If your bar is "could fool someone into thinking it's the real person" — 3/10, additive conditioning won't get you there.

Want me to do the FiLM upgrade? With your data budget it would push you to a solid **8/10**. It's a small code change — replace addition with scale+shift in 3 places (latent, targets, mel decoder).


Can we make it like 10 by 10, like fix all the things? I want it to be, because we're doing this because we can do a compact model with this architecture. We can do a compact model. If there is a compact model, it won't matter if we don't produce in art regressive mode, because anyway the hardware will be produced in milliseconds; it feels like very fast. In the production, let's say there are 10 companies only; only one in 10 companies, nine companies, only need fixed speakers, so this model will work for so many of them.


You're right — compact model + fixed speakers + 10-20h per speaker is a strong product thesis. Let me do the full overhaul.


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	            # Speaker classifier on audio encoder output — forces encoder to retain speaker info
515	            self.speaker_classifier = nn.Linear(d_model, n_speakers)
516	        else:
517	            self.speaker_embed = None
518	            self.speaker_classifier = None
519
520	        # Learnable start-of-sequence embedding for AR inference
521	        self.start_emb = nn.Parameter(torch.randn(1, 1, d_model) * 0.02)
522
523	        # Prediction weight — primary JEPA loss
524	        self.pred_weight = config.get("pred_weight", 10.0)
525
526	        # Input noise injection to bridge train/inference gap
527	        self.input_noise = config.get("input_noise", 0.0)
528
529	        # Free-bits KL: allow KL up to this threshold without penalty
530	        self.free_bits = config.get("free_bits", 1.0)
531
532	        # Multi-resolution spectral loss for sharper mel output
533	        self.spectral_loss = MultiResolutionSpectralLoss()
534	        self.spectral_weight = config.get("spectral_weight", 0.5)
535
536	        # Speaker consistency loss weight
537	        self.speaker_weight = config.get("speaker_weight", 1.0)
538
539	        self.config = config
540
541	    def _add_speaker(self, text_emb, speaker_id):
542	        """Add speaker embedding to text embeddings."""
543	        if self.speaker_embed is not None and speaker_id is not None:
544	            spk = self.speaker_embed(speaker_id).unsqueeze(1)  # [B, 1, d]
545	            text_emb = text_emb + spk
546	        return text_emb
547
548	    @torch.no_grad()
549	    def update_ema(self):
550	        """Update EMA target encoder weights."""
551	        for p_online, p_ema in zip(self.audio_encoder.parameters(),
552	                                   self.ema_audio_encoder.parameters()):
553	            p_ema.data.mul_(self.ema_decay).add_(p_online.data, alpha=1 - self.ema_decay)
554
555	    def forward(self, mel, text_tokens, mel_mask=None, text_mask=None, speaker_id=None):
556	        """
557	        Args:
558	            mel: [B, n_mels, T_mel]
559	            text_tokens: [B, T_text]
560	            mel_mask: [B, T_mel] bool
561	            text_mask: [B, T_text] bool
562	        Returns:
563	            loss_dict with 'prediction_loss', 'kl_loss', 'total_loss'
564	        """
565	        # Encode text + speaker conditioning
566	        text_emb = self.text_encoder(text_tokens, text_mask)
567	        text_emb = self._add_speaker(text_emb, speaker_id)
568
569	        # Online encoder → latent embeddings (used for predictor input + reconstruction)
570	        z, mu, logvar = self.audio_encoder(mel, mel_mask)
571	        B = z.shape[0]
572	        T_down = z.shape[1]
573
574	        # EMA target encoder → stable prediction targets (no gradient)
575	        # Force eval mode so dropout is disabled — EMA targets must be deterministic
576	        with torch.no_grad():
577	            self.ema_audio_encoder.eval()
578	            _, mu_ema, _ = self.ema_audio_encoder(mel, mel_mask)
579	            self.ema_audio_encoder.train()
580	            target_emb = mu_ema  # [B, T_down, d] — all EMA targets
581
582	        # Save raw encoder mu before speaker injection — used for speaker classifier
583	        # so the encoder itself must learn speaker info, not just parrot the added embedding
584	        mu_raw = mu
585
586	        # Inject speaker into targets and encoder output so prediction loss
587	        # requires speaker discrimination — two speakers saying the same word
588	        # produce different targets, forcing the predictor to use speaker info
589	        if self.speaker_embed is not None and speaker_id is not None:
590	            spk = self.speaker_embed(speaker_id).unsqueeze(1)  # [B, 1, d]
591	            target_emb = target_emb + spk
592	            z = z + spk
593
594	        # Prepend learnable start embedding to input sequence
595	        # Input:  [start_emb, z[0], z[1], ..., z[T-2]]  → T_down positions
596	        # Target: [z_ema[0],  z_ema[1], ..., z_ema[T-1]] → T_down positions
597	        # This trains the model to predict the first audio embedding from start_emb + text
598	        start = self.start_emb.expand(B, -1, -1)  # [B, 1, d]
599	        input_emb = torch.cat([start, z[:, :-1]], dim=1)  # [B, T_down, d]
600
601	        # Input noise: only on positions 1+ (real embeddings), not start_emb
602	        # Simulates the imperfect predictions seen during AR inference
603	        if self.training and self.input_noise > 0:
604	            noise = torch.zeros_like(input_emb)
605	            noise[:, 1:] = torch.randn(B, T_down - 1, input_emb.shape[-1],
606	                                        device=input_emb.device) * self.input_noise
607	            input_emb = input_emb + noise
608
609	        # Adjust masks for downsampled sequence
610	        if mel_mask is not None:
611	            ds_mask = mel_mask[:, ::self.audio_encoder.downsample_factor][:, :T_down]
612	            # Predictor input mask: start_emb (pos 0) is always unmasked
613	            start_mask = torch.zeros(B, 1, dtype=torch.bool, device=z.device)
614	            pred_mask = torch.cat([start_mask, ds_mask[:, :-1]], dim=1)  # [B, T_down]
615	            # Loss mask: follows original ds_mask for target positions
616	            loss_mask = ds_mask
617	        else:
618	            pred_mask = None
619	            loss_mask = None
620
621	        predicted = self.predictor(input_emb, text_emb, pred_mask, text_mask)
622
623	        # ─── Loss 1: Next-embedding prediction (MSE) against EMA targets ───
624	        if loss_mask is not None:
625	            valid = (~loss_mask).unsqueeze(-1)  # [B, T_down, 1]
626	            prediction_loss = F.mse_loss(
627	                predicted * valid, target_emb * valid, reduction="sum"
628	            ) / (valid.sum() * predicted.shape[-1] + 1e-8)
629	        else:
630	            prediction_loss = F.mse_loss(predicted, target_emb)
631
632	        # ─── Loss 2: Free-bits KL — only penalize KL above threshold ───
633	        # Prevents both KL collapse (too low) and KL explosion (too high)
634	        kl_per_dim = -0.5 * (1 + logvar - mu.pow(2) - logvar.exp())  # [B, T, d]
635	        kl_free = torch.clamp(kl_per_dim - self.free_bits, min=0.0)
636	        if mel_mask is not None:
637	            T_down = mu.shape[1]
638	            ds_mask = mel_mask[:, ::self.audio_encoder.downsample_factor][:, :T_down]
639	            valid_kl = (~ds_mask).unsqueeze(-1).float()
640	            kl_loss = (kl_free * valid_kl).sum() / (valid_kl.sum() * mu.shape[-1] + 1e-8)
641	        else:
642	            kl_loss = kl_free.mean()
643
644	        # ─── Loss 3: Mel reconstruction ───
645	        mel_recon = self.mel_decoder(z)  # [B, n_mels, T_up]
646	        T_mel = mel.shape[2]
647	        T_recon = mel_recon.shape[2]
648	        if T_recon > T_mel:
649	            mel_recon = mel_recon[:, :, :T_mel]
650	        elif T_recon < T_mel:
651	            mel = mel[:, :, :T_recon]
652	            if mel_mask is not None:
653	                mel_mask = mel_mask[:, :T_recon]
654
655	        if mel_mask is not None:
656	            valid_mel = (~mel_mask).unsqueeze(1).float()  # [B, 1, T]
657	            recon_loss = F.l1_loss(
658	                mel_recon * valid_mel, mel * valid_mel, reduction="sum"
659	            ) / (valid_mel.sum() * mel.shape[1] + 1e-8)
660	        else:
661	            recon_loss = F.l1_loss(mel_recon, mel)
662
663	        # ─── Loss 4: Multi-resolution spectral loss ───
664	        if mel_mask is not None:
665	            spectral_loss = self.spectral_loss(mel_recon * valid_mel, mel * valid_mel)
666	        else:
667	            spectral_loss = self.spectral_loss(mel_recon, mel)
668
669	        # ─── Loss 5: Speaker consistency — force audio encoder to retain speaker identity ───
670	        # Uses the raw mu (before speaker embed injection) so the encoder itself
671	        # must learn to preserve speaker info, not just parrot the added embedding
672	        speaker_loss = torch.tensor(0.0, device=mel.device)
673	        if self.speaker_classifier is not None and speaker_id is not None:
674	            # Pool encoder output across time (exclude padding)
675	            if mel_mask is not None:
676	                ds_mask_spk = mel_mask[:, ::self.audio_encoder.downsample_factor][:, :T_down]
677	                valid_spk = (~ds_mask_spk).unsqueeze(-1).float()  # [B, T_down, 1]
678	                mu_pooled = (mu_raw * valid_spk).sum(dim=1) / (valid_spk.sum(dim=1) + 1e-8)
679	            else:
680	                mu_pooled = mu_raw.mean(dim=1)  # [B, d]
681	            spk_logits = self.speaker_classifier(mu_pooled)  # [B, n_speakers]
682	            speaker_loss = F.cross_entropy(spk_logits, speaker_id)
683
684	        total_loss = (self.pred_weight * prediction_loss
685	                      + self.kl_weight * kl_loss
686	                      + self.recon_weight * recon_loss
687	                      + self.spectral_weight * spectral_loss
688	                      + self.speaker_weight * speaker_loss)
689
690	        return {
691	            "total_loss": total_loss,
692	            "prediction_loss": prediction_loss,
693	            "kl_loss": kl_loss,
694	            "recon_loss": recon_loss,
695	            "spectral_loss": spectral_loss,
696	            "speaker_loss": speaker_loss,
697	        }
698
699	    def encode_audio(self, mel):
700	        """Encode mel to latent embeddings (for inference)."""
701	        z, mu, _ = self.audio_encoder(mel)
702	        return mu  # Use mean at inference
703
704	    def predict_next(self, audio_emb, text_emb, text_mask=None):
705	        """Predict next audio embedding autoregressively."""
706	        predicted = self.predictor(audio_emb, text_emb, text_mask=text_mask)
707	        return predicted[:, -1:]  # Only the last prediction
708
709	    def init_ar_cache(self):
710	        """Initialize KV cache for autoregressive inference."""
711	        num_layers = len(self.predictor.transformer.layers)
712	        return self.predictor.init_cache(num_layers)
713
714	    def predict_next_cached(self, new_emb, text_emb, step_idx, cache, text_mask=None):
715	        """Predict next embedding using KV cache. O(1) per step instead of O(n)."""
716	        return self.predictor.inference_step(new_emb, text_emb, step_idx, cache, text_mask=text_mask)
717
718
719	def count_parameters(model):
720	    return sum(p.numel() for p in model.parameters() if p.requires_grad)
721
722
723	def build_model(config=None):
724	    if config is None:
725	        config = {
726	            "d_model": 256,
727	            "nhead": 4,
728	            "n_mels": 100,
729	            "text_vocab_size": 256,
730	            "text_encoder_layers": 4,
731	            "audio_encoder_layers": 4,
732	            "predictor_layers": 6,
733	            "dropout": 0.1,
734	            "kl_weight": 0.05,
735	            "pred_weight": 10.0,
736	            "ema_decay": 0.998,
737	            "free_bits": 1.0,
738	            "input_noise": 0.0,
739	        }
740	    model = LeWMTTS(config)
741	    print(f"LeWM TTS model: {count_parameters(model)/1e6:.2f}M parameters")
742	    return model, config
743

"""
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)
            # Speaker classifier on audio encoder output — forces encoder to retain speaker info
            self.speaker_classifier = nn.Linear(d_model, n_speakers)
        else:
            self.speaker_embed = None
            self.speaker_classifier = 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)

        # Speaker consistency loss weight
        self.speaker_weight = config.get("speaker_weight", 1.0)

        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

        # Save raw encoder mu before speaker injection — used for speaker classifier
        # so the encoder itself must learn speaker info, not just parrot the added embedding
        mu_raw = mu

        # Inject speaker into targets and encoder output so prediction loss
        # requires speaker discrimination — two speakers saying the same word
        # produce different targets, forcing the predictor to use speaker info
        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
            z = z + spk

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

        # ─── Loss 5: Speaker consistency — force audio encoder to retain speaker identity ───
        # Uses the raw mu (before speaker embed injection) so the encoder itself
        # must learn to preserve speaker info, not just parrot the added embedding
        speaker_loss = torch.tensor(0.0, device=mel.device)
        if self.speaker_classifier is not None and speaker_id is not None:
            # Pool encoder output across time (exclude padding)
            if mel_mask is not None:
                ds_mask_spk = mel_mask[:, ::self.audio_encoder.downsample_factor][:, :T_down]
                valid_spk = (~ds_mask_spk).unsqueeze(-1).float()  # [B, T_down, 1]
                mu_pooled = (mu_raw * valid_spk).sum(dim=1) / (valid_spk.sum(dim=1) + 1e-8)
            else:
                mu_pooled = mu_raw.mean(dim=1)  # [B, d]
            spk_logits = self.speaker_classifier(mu_pooled)  # [B, n_speakers]
            speaker_loss = F.cross_entropy(spk_logits, speaker_id)

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

        return {
            "total_loss": total_loss,
            "prediction_loss": prediction_loss,
            "kl_loss": kl_loss,
            "recon_loss": recon_loss,
            "spectral_loss": spectral_loss,
            "speaker_loss": speaker_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 I have the full picture. Let me rewrite the model with all fixes. This is a significant overhaul touching multiple components.