The user confirms:
1. synth_pure is noise
2. new1 and new2 sound the same as original - text conditioning still not working

At epoch 30, cosine is still 0.9999. The 20% dropout isn't enough to force text conditioning. The problem is fundamental:

- 80% of the time, the model has audio context and can ignore text
- 20% of the time, audio is zeroed but the model still has to predict from text alone → this is a much harder task and the model might just learn to output mean/noise for those samples

The issue is that 20% audio dropout creates two very different modes:
1. With audio (80%): easy, just predict from previous frame

2. Without audio (20%): extremely hard, generate from text alone with no audio context

The model minimizes loss by excelling at mode 1 while ignoring mode 2 entirely—it never actually learns to condition on text. I'm realizing the JEPA next-frame prediction approach fundamentally doesn't work for TTS. The proven methods use either non-autoregressive alignment like FastSpeech and VITS, or discrete token prediction like VALL-E where text alignment emerges naturally through the discrete prediction task.

Flow and diffusion models like F5-TTS handle this differently—they explicitly condition on text through cross-attention while denoising. But JEPA's continuous frame prediction is too easy to solve from audio history alone, so the model never needs to use the text signal. I've tried multiple iterations now: the mel decoder was the bottleneck in v3, fixing the JEPA mechanics in v4 didn't help, switching to EnCodec in v5 failed, simplifying to linear projections got good roundtrip loss but the predictor still ignored text, and even classifier-free guidance couldn't force text conditioning. The core issue is that next-frame prediction doesn't inherently require text to work well.

I could try massive audio dropout—50-80%—to force the model to rely on text half the time, but then it's not really JEPA anymore. And even then, at inference when I provide audio context, the model might still prefer ignoring text. Classifier-free guidance could help here too, generating with both audio and text, then with just text, and blending them the way diffusion models do. But that's complicated for autoregressive models. Maybe I should just pivot to discrete token prediction, which is proven to work, or add explicit alignment mechanisms like a duration predictor.

Actually, there's a simpler idea: instead of dropping entire audio chunks, I could mask random spans of audio frames. This way the model gets some audio context for voice continuity but has gaps it needs to fill using text—similar to how BERT learns bidirectional context. The causal predictor would zero out random frames and predict the next one using whatever audio remains plus the text.

But I'm realizing the core problem might be more fundamental. The user wants text-to-speech generation, but JEPA is designed to predict the next audio frame from previous frames plus text as auxiliary context. Those are actually different tasks—TTS is about matching speech to text content, while JEPA is about frame-level sequence prediction where text is just extra information.

For TTS to actually work, text needs to be the primary driver, not secondary. The model has to learn what phoneme to produce, how long it lasts, and what prosody to apply—none of which naturally emerge from next-frame prediction alone. I think I need to be direct with the user that pure JEPA won't work for this. But maybe there's a middle ground: use a duration predictor to align text to audio length, then concatenate the expanded text embeddings with audio embeddings before doing next-frame prediction. This way text is frame-aligned and can't be ignored by the model.

Actually, that's essentially what VITS and FastSpeech do. I should just be honest with the user—the JEPA approach with text as optional cross-attention won't cut it for TTS. I'll suggest either adding explicit duration prediction and alignment, or switching to a proven TTS architecture altogether.


I need to be straight with you. The 20% audio dropout isn't enough — the model still solves 80% of the loss by just copying audio, ignoring text.

**The core problem**: JEPA next-frame prediction is the wrong task for TTS. At 75Hz, predicting the next frame from previous frames is trivially solvable WITHOUT text. The model will always take the shortcut of using audio context instead of text, no matter how much dropout we add.

This is why no production TTS uses "predict next continuous frame with text as optional context." They all force text to be the primary driver:

- **VALL-E/XTTS**: Discrete tokens where text→audio mapping is explicit
- **FastSpeech/VITS**: Duration predictor aligns text to audio frames
- **F5-TTS**: Flow matching with forced text conditioning

We have two real options:

**Option A**: Add a **duration predictor** — expand text to match audio length, concatenate with input. Text becomes part of every frame, can't be ignored. Keeps the JEPA-ish architecture.

**Option B**: Switch to **discrete token prediction** with SNAC (12Hz, ~120 tokens for 10s). Predict token IDs (classification), not continuous vectors. Proven approach.

Which direction do you want to go? Option A is closer to what we have. Option B is more of a rewrite but more proven.


option A


Bash

kill $(ps aux | grep train_v5_pred | grep python | grep -v grep | awk '{print $2}') 2>/dev/null; sleep 3; nvidia-smi | grep MiB

Kill current training


| N/A   42C    P0             79W /  400W |       1MiB /  81920MiB |      0%      Default |


Let me design and implement the duration predictor approach. The idea:

1. **Duration predictor**: text tokens → predicted duration per character → expand text to match audio length
2. **Frame-level text**: each audio frame gets a text embedding concatenated/added — text can't be ignored
3. **Predictor**: predicts next frame from (audio + aligned text)


cd /home/ubuntu/lewm-tts && python3 -c "
# Check text-audio length relationships in dataset
import json, numpy as np
manifest = json.load(open('processed_data_codec/manifest.json'))

ratios = []
for e in manifest[:500]:
    text_len = len(e['text'].encode('utf-8'))
    audio_frames = e['emb_frames']
    ratios.append(audio_frames / text_len)

print(f'Audio frames per text byte:')
print(f'  Mean: {np.mean(ratios):.1f}')
print(f'  Std:  {np.std(ratios):.1f}')
print(f'  Min:  {np.min(ratios):.1f}')
print(f'  Max:  {np.max(ratios):.1f}')
print(f'\\nExample: text={len(manifest[0][\"text\"].encode(\"utf-8\"))} bytes, frames={manifest[0][\"emb_frames\"]}')
print(f'  Ratio: {manifest[0][\"emb_frames\"]/len(manifest[0][\"text\"].encode(\"utf-8\")):.1f} frames/byte')
" 2>&1

Check text-to-audio length ratios


Audio frames per text byte:
  Mean: 3.5
  Std:  0.6
  Min:  2.5
  Max:  8.7

Example: text=175 bytes, frames=714
  Ratio: 4.1 frames/byte
Shell cwd was reset to /home/ubuntu/pipecat


~3.5 frames per text byte on average. Now let me build the duration predictor + aligned text approach:


     1→"""
     2→LeWM TTS v5 — JEPA with EnCodec backend.
     3→
     4→Architecture:
     5→  - EnCodec encoder/decoder (frozen): audio ↔ 128d embeddings @ 75Hz
     6→  - Simple linear projections: 128d ↔ d_model (no lossy compression)
     7→  - TextEncoder: byte-level transformer for Hindi text
     8→  - JEPAPredictor: causal transformer, predicts next embedding
     9→  - EMA target on projected embeddings
    10→"""
    11→
    12→import math
    13→import copy
    14→import torch
    15→import torch.nn as nn
    16→import torch.nn.functional as F
    17→
    18→
    19→# ─── JEPA Predictor ─────────────────────────────────────────────────────────
    20→
    21→class JEPAPredictor(nn.Module):
    22→    def __init__(self, d_model=256, nhead=4, num_layers=6, dropout=0.1):
    23→        super().__init__()
    24→        self.d_model = d_model
    25→        self.audio_input_proj = nn.Linear(d_model, d_model)
    26→        self.pos_embed = SinusoidalPositionalEncoding(d_model, max_len=4096)
    27→        self.dropout = nn.Dropout(dropout)
    28→
    29→        decoder_layer = nn.TransformerDecoderLayer(
    30→            d_model=d_model, nhead=nhead, dim_feedforward=d_model * 4,

"""
LeWM TTS v5 — JEPA with EnCodec backend.

Architecture:
  - EnCodec encoder/decoder (frozen): audio ↔ 128d embeddings @ 75Hz
  - Simple linear projections: 128d ↔ d_model (no lossy compression)
  - TextEncoder: byte-level transformer for Hindi text
  - JEPAPredictor: causal transformer, predicts next embedding
  - EMA target on projected embeddings
"""

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

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

class JEPAPredictor(nn.Module):
    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)

        decoder_layer = nn.TransformerDecoderLayer(
            d_model=d_model, nhead=nhead, dim_feedforward=d_model * 4,


   215→        # EMA on proj_in (provides stable prediction targets)
   216→        self.ema_proj_in = copy.deepcopy(self.proj_in)
   217→        for p in self.ema_proj_in.parameters():
   218→            p.requires_grad = False
   219→        self.ema_decay = config.get("ema_decay", 0.998)
   220→
   221→        # Learnable start embedding
   222→        self.start_emb = nn.Parameter(torch.randn(1, 1, d_model) * 0.02)
   223→
   224→        # Loss weights
   225→        self.pred_weight = config.get("pred_weight", 10.0)
   226→        self.roundtrip_weight = config.get("roundtrip_weight", 1.0)
   227→
   228→        # Input noise for train/inference gap
   229→        self.input_noise = config.get("input_noise", 0.0)
   230→
   231→        self.config = config
   232→
   233→    @torch.no_grad()
   234→    def update_ema(self):
   235→        for p_online, p_ema in zip(self.proj_in.parameters(),
   236→                                   self.ema_proj_in.parameters()):
   237→            p_ema.data.mul_(self.ema_decay).add_(p_online.data, alpha=1 - self.ema_decay)
   238→
   239→    def forward(self, codec_emb, text_tokens, codec_mask=None, text_mask=None):
   240→        """
   241→        Args:
   242→            codec_emb: [B, 128, T] — continuous EnCodec embeddings
   243→            text_tokens: [B, T_text]
   244→            codec_mask: [B, T] bool, True = padding
   245→            text_mask: [B, T_text] bool
   246→        """
   247→        B = codec_emb.shape[0]
   248→        T = codec_emb.shape[2]
   249→
   250→        # codec_emb is [B, 128, T] → transpose to [B, T, 128]
   251→        codec_seq = codec_emb.transpose(1, 2)
   252→
   253→        # Text encoding
   254→        text_emb = self.text_encoder(text_tokens, text_mask)
   255→
   256→        # Project codec embeddings to predictor space
   257→        z = self.proj_in(codec_seq)  # [B, T, d_model]
   258→
   259→        # EMA targets (deterministic)
   260→        with torch.no_grad():
   261→            self.ema_proj_in.eval()
   262→            target_emb = self.ema_proj_in(codec_seq)  # [B, T, d_model]
   263→            self.ema_proj_in.train()
   264→
   265→        # Prepend start embedding, shift right
   266→        start = self.start_emb.expand(B, -1, -1)
   267→        input_emb = torch.cat([start, z[:, :-1]], dim=1)  # [B, T, d_model]
   268→
   269→        # Input noise
   270→        if self.training and self.input_noise > 0:
   271→            noise = torch.zeros_like(input_emb)
   272→            noise[:, 1:] = torch.randn(B, T - 1, input_emb.shape[-1],
   273→                                        device=input_emb.device) * self.input_noise
   274→            input_emb = input_emb + noise
   275→
   276→        # Masks
   277→        if codec_mask is not None:
   278→            start_mask = torch.zeros(B, 1, dtype=torch.bool, device=codec_emb.device)
   279→            pred_mask = torch.cat([start_mask, codec_mask[:, :-1]], dim=1)
   280→            loss_mask = codec_mask
   281→        else:
   282→            pred_mask = None
   283→            loss_mask = None
   284→
   285→        # JEPA prediction
   286→        predicted = self.predictor(input_emb, text_emb, pred_mask, text_mask)
   287→
   288→        # ─── Loss 1: Prediction (MSE vs EMA targets) ───
   289→        if loss_mask is not None:
   290→            valid = (~loss_mask).unsqueeze(-1)
   291→            prediction_loss = F.mse_loss(
   292→                predicted * valid, target_emb * valid, reduction="sum"
   293→            ) / (valid.sum() * predicted.shape[-1] + 1e-8)
   294→        else:
   295→            prediction_loss = F.mse_loss(predicted, target_emb)
   296→
   297→        # ─── Loss 2: Roundtrip (proj_in → proj_out → match original codec_emb) ───
   298→        codec_recon = self.proj_out(z)  # [B, T, 128]
   299→
   300→        if codec_mask is not None:
   301→            valid_rt = (~codec_mask).unsqueeze(-1).float()
   302→            roundtrip_loss = F.l1_loss(
   303→                codec_recon * valid_rt, codec_seq * valid_rt, reduction="sum"
   304→            ) / (valid_rt.sum() * self.codec_dim + 1e-8)
   305→        else:
   306→            roundtrip_loss = F.l1_loss(codec_recon, codec_seq)
   307→
   308→        total_loss = self.pred_weight * prediction_loss + self.roundtrip_weight * roundtrip_loss
   309→
   310→        return {
   311→            "total_loss": total_loss,
   312→            "prediction_loss": prediction_loss,
   313→            "roundtrip_loss": roundtrip_loss,
   314→        }

        # EMA on proj_in (provides stable prediction targets)
        self.ema_proj_in = copy.deepcopy(self.proj_in)
        for p in self.ema_proj_in.parameters():
            p.requires_grad = False
        self.ema_decay = config.get("ema_decay", 0.998)

        # Learnable start embedding
        self.start_emb = nn.Parameter(torch.randn(1, 1, d_model) * 0.02)

        # Loss weights
        self.pred_weight = config.get("pred_weight", 10.0)
        self.roundtrip_weight = config.get("roundtrip_weight", 1.0)

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

        self.config = config

    @torch.no_grad()
    def update_ema(self):
        for p_online, p_ema in zip(self.proj_in.parameters(),
                                   self.ema_proj_in.parameters()):
            p_ema.data.mul_(self.ema_decay).add_(p_online.data, alpha=1 - self.ema_decay)

    def forward(self, codec_emb, text_tokens, codec_mask=None, text_mask=None):
        """
        Args:
            codec_emb: [B, 128, T] — continuous EnCodec embeddings
            text_tokens: [B, T_text]
            codec_mask: [B, T] bool, True = padding
            text_mask: [B, T_text] bool
        """
        B = codec_emb.shape[0]
        T = codec_emb.shape[2]

        # codec_emb is [B, 128, T] → transpose to [B, T, 128]
        codec_seq = codec_emb.transpose(1, 2)

        # Text encoding
        text_emb = self.text_encoder(text_tokens, text_mask)

        # Project codec embeddings to predictor space
        z = self.proj_in(codec_seq)  # [B, T, d_model]

        # EMA targets (deterministic)
        with torch.no_grad():
            self.ema_proj_in.eval()
            target_emb = self.ema_proj_in(codec_seq)  # [B, T, d_model]
            self.ema_proj_in.train()

        # Prepend start embedding, shift right
        start = self.start_emb.expand(B, -1, -1)
        input_emb = torch.cat([start, z[:, :-1]], dim=1)  # [B, T, d_model]

        # Input noise
        if self.training and self.input_noise > 0:
            noise = torch.zeros_like(input_emb)
            noise[:, 1:] = torch.randn(B, T - 1, input_emb.shape[-1],
                                        device=input_emb.device) * self.input_noise
            input_emb = input_emb + noise

        # Masks
        if codec_mask is not None:
            start_mask = torch.zeros(B, 1, dtype=torch.bool, device=codec_emb.device)
            pred_mask = torch.cat([start_mask, codec_mask[:, :-1]], dim=1)
            loss_mask = codec_mask
        else:
            pred_mask = None
            loss_mask = None

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

        # ─── Loss 1: Prediction (MSE vs EMA targets) ───
        if loss_mask is not None:
            valid = (~loss_mask).unsqueeze(-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: Roundtrip (proj_in → proj_out → match original codec_emb) ───
        codec_recon = self.proj_out(z)  # [B, T, 128]

        if codec_mask is not None:
            valid_rt = (~codec_mask).unsqueeze(-1).float()
            roundtrip_loss = F.l1_loss(
                codec_recon * valid_rt, codec_seq * valid_rt, reduction="sum"
            ) / (valid_rt.sum() * self.codec_dim + 1e-8)
        else:
            roundtrip_loss = F.l1_loss(codec_recon, codec_seq)

        total_loss = self.pred_weight * prediction_loss + self.roundtrip_weight * roundtrip_loss

        return {
            "total_loss": total_loss,
            "prediction_loss": prediction_loss,
            "roundtrip_loss": roundtrip_loss,
        }


Conversation compacted

Agent Bash Edit Glob Grep Read TaskCreate TaskUpdate Write


This session is being continued from a previous conversation that ran out of context. The summary below covers the earlier portion of the conversation.

Summary:
1. Primary Request and Intent:
   The user is building a JEPA-based Hindi Text-to-Speech system at `/home/ubuntu/lewm-tts`. After v3/v3b/v3c/v4 attempts failed due to weak mel decoders, the user suggested using a neural codec (like "nanocodec") instead. The system was pivoted to use EnCodec as the audio backend. Through iterative testing, the roundtrip quality (encode→project→project_back→decode) became good, but the JEPA predictor completely ignores text conditioning (cosine similarity = 1.0 between different texts). The user chose **Option A: add a duration predictor** to force text-audio alignment, making text an essential part of each frame's prediction rather than optional cross-attention context.

2. Key Technical Concepts:
   - **JEPA (Joint Embedding Predictive Architecture)**: Predicts next audio embedding from previous embeddings + text. The core finding: at 75Hz, next-frame prediction doesn't need text — audio context alone is sufficient, causing text to be completely ignored.
   - **EnCodec 24kHz**: Meta's neural audio codec. Encoder produces 128-dim continuous embeddings at 75Hz. Decoder reconstructs audio from embeddings. Near-transparent quality (L1=0.01). Used frozen (pretrained).
   - **Linear projections (proj_in/proj_out)**: `Linear(128→256→256)` and `Linear(256→256→128)` to map between codec space (128d) and predictor space (256d). Roundtrip L1=0.27, audio sounds clean.
   - **Text conditioning failure**: Cosine similarity = 1.0000 between predictions with different texts. The predictor learned to predict next frame purely from audio context, ignoring cross-attention to text entirely.
   - **Classifier-Free Guidance (CFG)**: 20% audio dropout during training to force text usage. Still insufficient at epoch 30 (cosine=0.9999).
   - **Voice prompting**: Using real audio as prefix (3s) then AR-generating continuation. Sounds good but only because it continues audio pattern — doesn't speak new text.
   - **Duration predictor**: The chosen fix — expand text to match audio frame length, concatenate/add to each frame so text can't be ignored. Similar to FastSpeech/VITS approach.
   - **75Hz problem**: At 75Hz, 10 seconds needs ~750 AR steps. Too many for error accumulation. Pure AR synthesis collapses after ~50-150 steps.
   - **EMA divergence on small networks**: EMA on 2-layer linear proj_in caused prediction loss to increase (0.12→0.38) because the EMA and online projections diverge slightly, confusing the predictor.

3. Files and Code Sections:
   - **`/home/ubuntu/lewm-tts/model_v5.py`** — Core v5 model. Most important file.
     - Contains LeWMTTSv5 with: TextEncoder, linear proj_in/proj_out, JEPAPredictor (6L decoder with KV-cache), EMA, start_emb
     - The predictor uses TransformerDecoder with cross-attention to text (which is being ignored)
     - Was modified from complex down/upsampler CNN to simple linear projections
     - Key architecture:
     ```python
     class LeWMTTSv5(nn.Module):
         def __init__(self, config):
             # ...
             self.proj_in = nn.Sequential(
                 nn.Linear(codec_dim, d_model), nn.GELU(), nn.Linear(d_model, d_model),
             )
             self.proj_out = nn.Sequential(
                 nn.Linear(d_model, d_model), nn.GELU(), nn.Linear(d_model, codec_dim),
             )
             self.predictor = JEPAPredictor(d_model=256, nhead=4, num_layers=6)
             self.text_encoder = TextEncoder(vocab_size=256, d_model=256, num_layers=4)
             self.start_emb = nn.Parameter(torch.randn(1, 1, d_model) * 0.02)
     ```
     - Projections are FROZEN (trained in output_v5/), predictor is trained separately (output_v5_pred/)
     - `codec_dim` variable bug was fixed (`codec_dim` → `self.codec_dim` in forward)

   - **`/home/ubuntu/lewm-tts/train_v5_pred.py`** — Predictor-only training with frozen projections
     - Loads pretrained checkpoint, freezes proj_in/proj_out/ema_proj_in
     - Has `--reinit` flag to reinitialize predictor from scratch
     - Has `--audio_drop_prob` for CFG (classifier-free guidance) audio dropout
     - Key CFG code:
     ```python
     if args.audio_drop_prob > 0:
         drop_mask = torch.rand(B, device=device) < args.audio_drop_prob
         if drop_mask.any():
             input_emb = input_emb.clone()
             input_emb[drop_mask] = 0.0  # zero out audio for dropped samples
     ```
     - Uses frozen proj_in targets directly (no EMA) since proj_in is frozen

   - **`/home/ubuntu/lewm-tts/train_v5.py`** — Full v5 training loop (used for training projections)

   - **`/home/ubuntu/lewm-tts/dataset_v5.py`** — CodecDataset loading EnCodec embeddings + text, with collate_fn for padding

   - **`/home/ubuntu/lewm-tts/inference_v5.py`** — Inference with AR generation, voice prompting, reconstruction. Uses `model.proj_out(pred_embs).transpose(1, 2)` then `codec.decoder()`.

   - **`/home/ubuntu/lewm-tts/preprocess_codec.py`** — Extracts EnCodec continuous embeddings for all 7009 audio files. Output in `processed_data_codec/`.

   - **`/home/ubuntu/lewm-tts/model.py`** — Original v4 model (mel-based). Not used in v5.

   - **`/home/ubuntu/.claude/projects/-home-ubuntu-pipecat/memory/lewm-tts.md`** — Memory file with project notes

   - **Checkpoints**:
     - `output_v5/checkpoints/epoch_0120.pt` — Best projection training (roundtrip L1=0.27, "fucking good")
     - `output_v5_pred/checkpoints/best.pt` — Latest predictor training (CFG, epoch 30, pred_loss=0.62)

   - **Data**: `processed_data_codec/` — 7009 samples, EnCodec embeddings (128d, 75Hz), avg 689 frames/sample

4. Errors and fixes:
   - **`codec_dim` not defined in forward()**: Referenced local variable instead of instance variable. Fixed: `codec_dim` → `self.codec_dim`
   - **Down/upsampler CNN caused noisy roundtrip**: Complex stride-3 CNN down/upsampler couldn't faithfully compress/decompress EnCodec embeddings (range [-22, +24]). Fixed: replaced with simple linear projections (proj_in/proj_out).
   - **EMA on small proj_in caused pred_loss divergence (0.12→0.38)**: EMA of 2-layer linear network creates slightly different linear map, confusing predictor. Fixed: removed EMA, froze proj_in, use online proj_in directly as targets.
   - **Pretrained predictor stuck at 0.40**: Predictor trained against EMA targets was at local minimum for new frozen targets. Fixed: reinitialize predictor from scratch with `--reinit` flag.
   - **VQ quantization made audio worse, not better**: `roundtrip_vq.wav` sounded worse than `roundtrip.wav`. User: "rountrip_vq and synth fucked up". Abandoned VQ approach.
   - **Text conditioning completely ignored (cosine=1.0)**: JEPA next-frame prediction at 75Hz doesn't need text. User: "it is not saying this text." Attempted fix: CFG audio dropout 20%. Still not working (cosine=0.9999 at epoch 30). User: "new1 and new2 also cmng up the same text as original." Root cause: the prediction task itself doesn't require text. User chose Option A: duration predictor.
   - **QuantizedResult unpacking error**: `TypeError: cannot unpack non-iterable QuantizedResult object`. Fixed: accessed attributes directly (`result.quantized`, `result.codes`).

5. Problem Solving:
   - **v4 mel decoder bottleneck**: Diagnosed that even with perfect JEPA prediction (cosine_sim=0.97), mel decoder L1=0.57 was too high. JEPA predictor was working (6.2x better than copy baseline) but decoder was the bottleneck.
   - **Codec approach iteration**: Went through: complex down/upsampler → noisy → simplified linear projections → roundtrip good → predictor issues → frozen projections → text not used → CFG → still not working → duration predictor (pending).
   - **Text conditioning**: Verified with controlled experiments that different texts produce identical predictions (cosine=1.0). Confirmed it's a training problem, not inference. The architecture allows text but doesn't require it.
   - **Key insight**: At 75Hz, consecutive frames are only 13ms apart and highly correlated. The predictor can minimize loss by extrapolating from audio alone, never learning to use text. This is fundamental to the JEPA next-frame prediction approach.

6. All user messages:
   - "check again" (asking for training status)
   - "check samples now"
   - "recon_0.wav │ Reconstruction encode→decode (9.5s) │ ... tf_0.wav siundung same.. noice + voice and synth is pure noise"
   - "jepa is predicting good ? check that ... if it good .. we can use anither codec like nanocodec ryt"
   - "hw much total szie"
   - "run samples on latest epoch"
   - "pure noise .. and y we not used codec like nanocode or speechtoknizzer codec?"
   - "roundtrip.wav this fucking good .. tf.wav also cmng voice correctly but sounding robotic .. rountrip_vq and synth fucked up"
   - "lets wait for sometime at epoch 10 do we get good audio ?"
   - "what we using ?"
   - "total param size" (sent during wait)
   - "status"
   - "create samples on this epoch"
   - "are yiu sure it will work this time?"
   - "can u try samples on latest epoch"
   - "synth fucked up ... roundtrip ok but noice cmng .. tf is robotic voice fully"
   - "y we not seeing loss?" (pred loss barely moving 0.405→0.397)
   - "these are fucking good" (about voice-prompted samples)
   - "on what checkpoint u ran prompted?"
   - "can u test other text not this one"
   - "some neew text which is not is dataset"
   - "it is not saying this text check once code once"
   - "only noise . it is trying to speak i think .. when person speaking at that time noise cmng instead of voice"
   - "status"
   - "samples"
   - "synth_pure.wav │ Pure text-only synthesis (1.99s - noise . and new1 and new2 also cmng up the same text as original"
   - "but synth is not good y"
   - "cant we use tf at infernce ?"
   - "are u sure we dng wrong in traing ... or do we dng in infernce"
   - "not cleaner noise cmng along with voice .. and y dont u check on latest checkpoint"
   - "check npw" (sent during wait)
   - "option A" (choosing duration predictor over discrete token prediction)

7. Pending Tasks:
   - **Implement Option A: Duration predictor + text-audio alignment** for the JEPA predictor
     - Design duration predictor module that predicts duration per text character
     - Expand text embeddings to match audio frame length
     - Concatenate/add aligned text to each audio frame input
     - Modify training to include duration prediction loss
     - Retrain predictor with forced text alignment
   - Update memory file with v5 findings

8. Current Work:
   The user chose **Option A** — adding a duration predictor to force text-audio alignment. I was about to implement this when the summary was requested. The current state:
   - **Frozen projections** (proj_in/proj_out) from `output_v5/checkpoints/epoch_0120.pt` work well (roundtrip L1=0.27, "fucking good" audio)
   - **Predictor needs text alignment**: The JEPAPredictor's cross-attention to text is completely ignored. Text must be expanded to frame level and concatenated/added to audio input so it can't be bypassed.
   - **Data stats**: ~3.5 audio frames per text byte on average, with range 2.5-8.7
   - I had just killed the CFG training run and confirmed GPU is free (1MiB used)
   - I was reading model_v5.py to understand the current architecture before implementing the duration predictor modification
   - The last code read was the forward() method of LeWMTTSv5 (lines 239-314 of model_v5.py)

9. Optional Next Step:
   Implement the duration predictor approach in model_v5.py. The user explicitly chose "option A" when presented with:
   > "**Option A**: Add a **duration predictor** — expand text to match audio length, concatenate with input. Text becomes part of every frame, can't be ignored. Keeps the JEPA-ish architecture."

   This involves:
   1. Creating a DurationPredictor module (small conv/transformer network)
   2. During training: compute actual durations from text-audio alignment (or use uniform/learned alignment)
   3. Expand text embeddings to audio frame length using predicted durations
   4. Concatenate or add expanded text embeddings to the audio input in the predictor (instead of cross-attention)
   5. Modify the predictor to use frame-level text (self-attention on combined audio+text, not cross-attention)
   6. Add duration prediction loss (L1 or MSE on predicted vs actual durations)

If you need specific details from before compaction (like exact code snippets, error messages, or content you generated), read the full transcript at: /home/ubuntu/.claude/projects/-home-ubuntu-pipecat/a6ad5d6a-37c4-4b0c-8baa-65a4a8ef03b9.jsonl
Continue the conversation from where it left off without asking the user any further questions. Resume directly — do not acknowledge the summary, do not recap what was happening, do not preface with "I'll continue" or similar. Pick up the last task as if the break never happened.


Let me implement the duration predictor approach. I need to:

1. Create a DurationPredictor module
2. Modify the model to expand text to frame level
3. Update training to use aligned text
4. Update inference

Let me start by implementing this in model_v5.py.


ToolSearch

Question: select:Read,Write,Edit,Bash,Grep


Question: select:Read,Write,Edit,Bash,Grep


Let me implement the duration predictor approach. I need to modify model_v5.py to add a DurationPredictor, modify the predictor to use frame-level text instead of cross-attention, and create a new training script.

Let me first check the dataset to understand what alignment info we have.


     1→"""
     2→Dataset for codec-based JEPA TTS v5.
     3→Loads precomputed EnCodec embeddings + text.
     4→"""
     5→
     6→import json
     7→import torch
     8→from torch.utils.data import Dataset, DataLoader
     9→
    10→
    11→class CodecDataset(Dataset):
    12→    def __init__(self, manifest_path, max_codec_frames=900):
    13→        with open(manifest_path, "r") as f:
    14→            self.manifest = json.load(f)
    15→        self.max_codec_frames = max_codec_frames
    16→        # Filter out too-long samples
    17→        self.manifest = [e for e in self.manifest if e["emb_frames"] <= max_codec_frames]
    18→        print(f"CodecDataset: {len(self.manifest)} samples (max {max_codec_frames} frames)")
    19→
    20→    def __len__(self):
    21→        return len(self.manifest)
    22→
    23→    def __getitem__(self, idx):
    24→        entry = self.manifest[idx]
    25→        emb = torch.load(entry["emb_path"], weights_only=True)  # [128, T]
    26→        text_tokens = torch.tensor(list(entry["text"].encode("utf-8")), dtype=torch.long)
    27→        return {"codec_emb": emb, "text_tokens": text_tokens}
    28→
    29→
    30→def collate_fn(batch):
    31→    codec_embs = [b["codec_emb"] for b in batch]
    32→    text_tokens = [b["text_tokens"] for b in batch]
    33→
    34→    # Pad codec embeddings
    35→    codec_dim = codec_embs[0].shape[0]
    36→    max_codec_len = max(e.shape[1] for e in codec_embs)
    37→    padded_codec = torch.zeros(len(batch), codec_dim, max_codec_len)
    38→    codec_mask = torch.ones(len(batch), max_codec_len, dtype=torch.bool)
    39→    for i, e in enumerate(codec_embs):
    40→        T = e.shape[1]
    41→        padded_codec[i, :, :T] = e
    42→        codec_mask[i, :T] = False
    43→
    44→    # Pad text
    45→    max_text_len = max(t.shape[0] for t in text_tokens)
    46→    padded_text = torch.zeros(len(batch), max_text_len, dtype=torch.long)
    47→    text_mask = torch.ones(len(batch), max_text_len, dtype=torch.bool)
    48→    for i, t in enumerate(text_tokens):
    49→        padded_text[i, :t.shape[0]] = t
    50→        text_mask[i, :t.shape[0]] = False
    51→
    52→    return {
    53→        "codec_emb": padded_codec,
    54→        "text_tokens": padded_text,
    55→        "codec_mask": codec_mask,
    56→        "text_mask": text_mask,
    57→    }
    58→
    59→
    60→def build_dataloader(manifest_path, batch_size=64, num_workers=4, max_codec_frames=900):
    61→    dataset = CodecDataset(manifest_path, max_codec_frames=max_codec_frames)
    62→    loader = DataLoader(
    63→        dataset, batch_size=batch_size, shuffle=True,
    64→        num_workers=num_workers, collate_fn=collate_fn,
    65→        pin_memory=True, drop_last=True,
    66→    )
    67→    return loader, dataset
    68→

"""
Dataset for codec-based JEPA TTS v5.
Loads precomputed EnCodec embeddings + text.
"""

import json
import torch
from torch.utils.data import Dataset, DataLoader

class CodecDataset(Dataset):
    def __init__(self, manifest_path, max_codec_frames=900):
        with open(manifest_path, "r") as f:
            self.manifest = json.load(f)
        self.max_codec_frames = max_codec_frames
        # Filter out too-long samples
        self.manifest = [e for e in self.manifest if e["emb_frames"] <= max_codec_frames]
        print(f"CodecDataset: {len(self.manifest)} samples (max {max_codec_frames} frames)")

    def __len__(self):
        return len(self.manifest)

    def __getitem__(self, idx):
        entry = self.manifest[idx]
        emb = torch.load(entry["emb_path"], weights_only=True)  # [128, T]
        text_tokens = torch.tensor(list(entry["text"].encode("utf-8")), dtype=torch.long)
        return {"codec_emb": emb, "text_tokens": text_tokens}

def collate_fn(batch):
    codec_embs = [b["codec_emb"] for b in batch]
    text_tokens = [b["text_tokens"] for b in batch]

    # Pad codec embeddings
    codec_dim = codec_embs[0].shape[0]
    max_codec_len = max(e.shape[1] for e in codec_embs)
    padded_codec = torch.zeros(len(batch), codec_dim, max_codec_len)
    codec_mask = torch.ones(len(batch), max_codec_len, dtype=torch.bool)
    for i, e in enumerate(codec_embs):
        T = e.shape[1]
        padded_codec[i, :, :T] = e
        codec_mask[i, :T] = False

    # Pad text
    max_text_len = max(t.shape[0] for t in text_tokens)
    padded_text = torch.zeros(len(batch), max_text_len, dtype=torch.long)
    text_mask = torch.ones(len(batch), max_text_len, dtype=torch.bool)
    for i, t in enumerate(text_tokens):
        padded_text[i, :t.shape[0]] = t
        text_mask[i, :t.shape[0]] = False

    return {
        "codec_emb": padded_codec,
        "text_tokens": padded_text,
        "codec_mask": codec_mask,
        "text_mask": text_mask,
    }

def build_dataloader(manifest_path, batch_size=64, num_workers=4, max_codec_frames=900):
    dataset = CodecDataset(manifest_path, max_codec_frames=max_codec_frames)
    loader = DataLoader(
        dataset, batch_size=batch_size, shuffle=True,
        num_workers=num_workers, collate_fn=collate_fn,
        pin_memory=True, drop_last=True,
    )
    return loader, dataset


Now let me implement the duration predictor approach. The plan:

1. **DurationPredictor** module: takes text encoder output, predicts log-duration per character
2. **Length regulator**: expands text embeddings to audio frame length using durations
3. **Modified predictor**: concatenate expanded text with audio input (add, not concat, to keep d_model the same)
4. During training: compute actual durations from text/audio lengths (uniform distribution initially)
5. At inference: use predicted durations


Edit

"""
LeWM TTS v5 — JEPA with EnCodec backend.

Architecture:
  - EnCodec encoder/decoder (frozen): audio ↔ 128d embeddings @ 75Hz
  - Simple linear projections: 128d ↔ d_model (no lossy compression)
  - TextEncoder: byte-level transformer for Hindi text
  - JEPAPredictor: causal transformer, predicts next embedding
  - EMA target on projected embeddings
"""

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

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

class JEPAPredictor(nn.Module):
    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)

        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)
        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):
        x = self.audio_input_proj(audio_emb)
        x = self.pos_embed(x)
        x = self.dropout(x)
        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,
        )
        return self.output_proj(predicted)

"""
LeWM TTS v5 — JEPA with EnCodec backend.

Architecture:
  - EnCodec encoder/decoder (frozen): audio ↔ 128d embeddings @ 75Hz
  - Simple linear projections: 128d ↔ d_model (no lossy compression)
  - TextEncoder: byte-level transformer for Hindi text
  - DurationPredictor: predicts per-character duration
  - JEPAPredictor: causal transformer with frame-level text conditioning
"""

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

# ─── Duration Predictor ─────────────────────────────────────────────────────

class DurationPredictor(nn.Module):
    """Predicts log-duration per text token. Similar to FastSpeech."""
    def __init__(self, d_model=256, kernel_size=3, num_layers=2, dropout=0.1):
        super().__init__()
        layers = []
        for _ in range(num_layers):
            layers.extend([
                nn.Conv1d(d_model, d_model, kernel_size, padding=kernel_size // 2),
                nn.ReLU(),
                nn.LayerNorm(d_model),
                nn.Dropout(dropout),
            ])
        self.convs = nn.ModuleList(layers)
        self.proj = nn.Linear(d_model, 1)

    def forward(self, text_emb, text_mask=None):
        """
        text_emb: [B, T_text, d_model]
        Returns: log_durations [B, T_text]
        """
        x = text_emb.transpose(1, 2)  # [B, d_model, T]
        for i in range(0, len(self.convs), 4):
            x = self.convs[i](x)      # Conv1d
            x = self.convs[i+1](x)    # ReLU
            x = x.transpose(1, 2)     # [B, T, d_model] for LayerNorm
            x = self.convs[i+2](x)    # LayerNorm
            x = self.convs[i+3](x)    # Dropout
            x = x.transpose(1, 2)     # back to [B, d_model, T]
        x = x.transpose(1, 2)  # [B, T, d_model]
        log_dur = self.proj(x).squeeze(-1)  # [B, T_text]
        if text_mask is not None:
            log_dur = log_dur.masked_fill(text_mask, 0.0)
        return log_dur

def length_regulate(text_emb, durations, text_mask=None):
    """
    Expand text embeddings by repeating each token according to its duration.
    text_emb: [B, T_text, d_model]
    durations: [B, T_text] integer durations
    Returns: expanded [B, T_audio, d_model], where T_audio = sum(durations) per sample
    """
    B, T_text, D = text_emb.shape
    max_audio_len = durations.sum(dim=1).max().item()
    expanded = torch.zeros(B, max_audio_len, D, device=text_emb.device)
    for b in range(B):
        pos = 0
        for t in range(T_text):
            if text_mask is not None and text_mask[b, t]:
                continue
            dur = durations[b, t].item()
            if dur > 0 and pos < max_audio_len:
                end = min(pos + dur, max_audio_len)
                expanded[b, pos:end] = text_emb[b, t]
                pos = end
    return expanded

def compute_uniform_durations(text_lengths, audio_lengths):
    """
    Compute uniform durations: spread audio frames evenly across text characters.
    text_lengths: [B] number of valid text tokens
    audio_lengths: [B] number of valid audio frames
    Returns: durations [B, max_text_len] integer
    """
    B = text_lengths.shape[0]
    max_text_len = text_lengths.max().item()
    durations = torch.zeros(B, max_text_len, dtype=torch.long, device=text_lengths.device)
    for b in range(B):
        tl = text_lengths[b].item()
        al = audio_lengths[b].item()
        if tl == 0:
            continue
        base = al // tl
        remainder = al % tl
        durations[b, :tl] = base
        durations[b, :remainder] += 1
    return durations

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

class JEPAPredictor(nn.Module):
    """Causal predictor with frame-level text added to audio input."""
    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.text_input_proj = nn.Linear(d_model, d_model)
        self.pos_embed = SinusoidalPositionalEncoding(d_model, max_len=4096)
        self.dropout = nn.Dropout(dropout)

        # Use encoder (self-attention only) since text is now part of input
        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.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_expanded, audio_mask=None):
        """
        audio_emb: [B, T, d_model] — shifted audio embeddings
        text_emb_expanded: [B, T, d_model] — frame-level text (already expanded)
        """
        x = self.audio_input_proj(audio_emb) + self.text_input_proj(text_emb_expanded)
        x = self.pos_embed(x)
        x = self.dropout(x)
        T = x.shape[1]
        causal_mask = torch.triu(torch.ones(T, T, device=x.device, dtype=torch.bool), diagonal=1)
        predicted = self.transformer(x, mask=causal_mask, src_key_padding_mask=audio_mask)
        return self.output_proj(predicted)

    # Keep old cross-attention interface for backward compat loading
    def forward_cross_attn(self, audio_emb, text_emb, audio_mask=None, text_mask=None):
        """Legacy cross-attention forward (not used in duration-predictor mode)."""
        x = self.audio_input_proj(audio_emb)
        x = self.pos_embed(x)
        x = self.dropout(x)
        return self.output_proj(x)


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

"""
LeWM TTS v5 — JEPA with EnCodec backend.

Architecture:
  - EnCodec encoder/decoder (frozen): audio ↔ 128d embeddings @ 75Hz
  - Simple linear projections: 128d ↔ d_model (no lossy compression)
  - TextEncoder: byte-level transformer for Hindi text
  - DurationPredictor: predicts per-character duration
  - JEPAPredictor: causal transformer with frame-level text conditioning
"""

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

# ─── Duration Predictor ─────────────────────────────────────────────────────

class DurationPredictor(nn.Module):
    """Predicts log-duration per text token. Similar to FastSpeech."""
    def __init__(self, d_model=256, kernel_size=3, num_layers=2, dropout=0.1):
        super().__init__()
        layers = []
        for _ in range(num_layers):
            layers.extend([
                nn.Conv1d(d_model, d_model, kernel_size, padding=kernel_size // 2),
                nn.ReLU(),
                nn.LayerNorm(d_model),
                nn.Dropout(dropout),
            ])
        self.convs = nn.ModuleList(layers)
        self.proj = nn.Linear(d_model, 1)

    def forward(self, text_emb, text_mask=None):
        """
        text_emb: [B, T_text, d_model]
        Returns: log_durations [B, T_text]
        """
        x = text_emb.transpose(1, 2)  # [B, d_model, T]
        for i in range(0, len(self.convs), 4):
            x = self.convs[i](x)      # Conv1d
            x = self.convs[i+1](x)    # ReLU
            x = x.transpose(1, 2)     # [B, T, d_model] for LayerNorm
            x = self.convs[i+2](x)    # LayerNorm
            x = self.convs[i+3](x)    # Dropout
            x = x.transpose(1, 2)     # back to [B, d_model, T]
        x = x.transpose(1, 2)  # [B, T, d_model]
        log_dur = self.proj(x).squeeze(-1)  # [B, T_text]
        if text_mask is not None:
            log_dur = log_dur.masked_fill(text_mask, 0.0)
        return log_dur

def length_regulate(text_emb, durations, text_mask=None):
    """
    Expand text embeddings by repeating each token according to its duration.
    text_emb: [B, T_text, d_model]
    durations: [B, T_text] integer durations
    Returns: expanded [B, T_audio, d_model], where T_audio = sum(durations) per sample
    """
    B, T_text, D = text_emb.shape
    max_audio_len = durations.sum(dim=1).max().item()
    expanded = torch.zeros(B, max_audio_len, D, device=text_emb.device)
    for b in range(B):
        pos = 0
        for t in range(T_text):
            if text_mask is not None and text_mask[b, t]:
                continue
            dur = durations[b, t].item()
            if dur > 0 and pos < max_audio_len:
                end = min(pos + dur, max_audio_len)
                expanded[b, pos:end] = text_emb[b, t]
                pos = end
    return expanded

def compute_uniform_durations(text_lengths, audio_lengths):
    """
    Compute uniform durations: spread audio frames evenly across text characters.
    text_lengths: [B] number of valid text tokens
    audio_lengths: [B] number of valid audio frames
    Returns: durations [B, max_text_len] integer
    """
    B = text_lengths.shape[0]
    max_text_len = text_lengths.max().item()
    durations = torch.zeros(B, max_text_len, dtype=torch.long, device=text_lengths.device)
    for b in range(B):
        tl = text_lengths[b].item()
        al = audio_lengths[b].item()
        if tl == 0:
            continue
        base = al // tl
        remainder = al % tl
        durations[b, :tl] = base
        durations[b, :remainder] += 1
    return durations

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

class JEPAPredictor(nn.Module):
    """Causal predictor with frame-level text added to audio input."""
    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.text_input_proj = nn.Linear(d_model, d_model)
        self.pos_embed = SinusoidalPositionalEncoding(d_model, max_len=4096)
        self.dropout = nn.Dropout(dropout)

        # Use encoder (self-attention only) since text is now part of input
        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.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_expanded, audio_mask=None):
        """
        audio_emb: [B, T, d_model] — shifted audio embeddings
        text_emb_expanded: [B, T, d_model] — frame-level text (already expanded)
        """
        x = self.audio_input_proj(audio_emb) + self.text_input_proj(text_emb_expanded)
        x = self.pos_embed(x)
        x = self.dropout(x)
        T = x.shape[1]
        causal_mask = torch.triu(torch.ones(T, T, device=x.device, dtype=torch.bool), diagonal=1)
        predicted = self.transformer(x, mask=causal_mask, src_key_padding_mask=audio_mask)
        return self.output_proj(predicted)

    # Keep old cross-attention interface for backward compat loading
    def forward_cross_attn(self, audio_emb, text_emb, audio_mask=None, text_mask=None):
        """Legacy cross-attention forward (not used in duration-predictor mode)."""
        x = self.audio_input_proj(audio_emb)
        x = self.pos_embed(x)
        x = self.dropout(x)
        return self.output_proj(x)

"""
LeWM TTS v5 — JEPA with EnCodec backend.

Architecture:
  - EnCodec encoder/decoder (frozen): audio ↔ 128d embeddings @ 75Hz
  - Simple linear projections: 128d ↔ d_model (no lossy compression)
  - TextEncoder: byte-level transformer for Hindi text
  - JEPAPredictor: causal transformer, predicts next embedding
  - EMA target on projected embeddings
"""

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

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

class JEPAPredictor(nn.Module):
    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)

        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)
        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):
        x = self.audio_input_proj(audio_emb)
        x = self.pos_embed(x)
        x = self.dropout(x)
        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,
        )
        return self.output_proj(predicted)

    def _cached_mha(self, mha, q_input, kv_input, cache_k, cache_v):
        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)
        k_new = F.linear(kv_input, Wk, bk)
        v_new = F.linear(kv_input, Wv, bv)
        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, v = k_new, v_new
        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)
        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 _cached_cross_attn(self, mha, q_input, memory, cache_k, cache_v, memory_mask=None):
        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)
        if cache_k is None:
            k = F.linear(memory, Wk, bk)
            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:
            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):
        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):
        x = self.audio_input_proj(new_emb)
        x = x + self.pos_embed.pe[:, step_idx:step_idx + 1]
        for i, layer in enumerate(self.transformer.layers):
            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))
            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))
            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))

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

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

class TextEncoder(nn.Module):
    def __init__(self, vocab_size=256, d_model=256, nhead=4, num_layers=4, dropout=0.1):
        super().__init__()
        self.d_model = d_model
        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):
        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)
        return self.proj(x)

# ─── Full v5 Model ──────────────────────────────────────────────────────────

class LeWMTTSv5(nn.Module):
    """
    JEPA TTS with EnCodec backend.
    No down/upsampler — works directly at EnCodec's native 75Hz.
    Simple linear projections: 128d (codec) ↔ d_model (predictor).

    Training: codec_emb → proj_in → JEPA predict → loss vs EMA targets
              codec_emb → proj_in → proj_out → loss vs codec_emb (roundtrip)
    Inference: text → predictor (AR) → proj_out → EnCodec decoder → waveform
    """

    def __init__(self, config):
        super().__init__()
        d_model = config.get("d_model", 256)
        nhead = config.get("nhead", 4)
        codec_dim = config.get("codec_dim", 128)
        text_vocab = config.get("text_vocab_size", 256)
        text_layers = config.get("text_encoder_layers", 4)
        predictor_layers = config.get("predictor_layers", 6)
        dropout = config.get("dropout", 0.1)

        self.codec_dim = codec_dim

        self.text_encoder = TextEncoder(
            vocab_size=text_vocab, d_model=d_model, nhead=nhead,
            num_layers=text_layers, dropout=dropout,
        )

        # Simple linear projections — no lossy compression
        self.proj_in = nn.Sequential(
            nn.Linear(codec_dim, d_model),
            nn.GELU(),
            nn.Linear(d_model, d_model),
        )
        self.proj_out = nn.Sequential(
            nn.Linear(d_model, d_model),
            nn.GELU(),
            nn.Linear(d_model, codec_dim),
        )

        self.predictor = JEPAPredictor(
            d_model=d_model, nhead=nhead,
            num_layers=predictor_layers, dropout=dropout,
        )

        # EMA on proj_in (provides stable prediction targets)
        self.ema_proj_in = copy.deepcopy(self.proj_in)
        for p in self.ema_proj_in.parameters():
            p.requires_grad = False
        self.ema_decay = config.get("ema_decay", 0.998)

        # Learnable start embedding
        self.start_emb = nn.Parameter(torch.randn(1, 1, d_model) * 0.02)

        # Loss weights
        self.pred_weight = config.get("pred_weight", 10.0)
        self.roundtrip_weight = config.get("roundtrip_weight", 1.0)

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

        self.config = config

    @torch.no_grad()
    def update_ema(self):
        for p_online, p_ema in zip(self.proj_in.parameters(),
                                   self.ema_proj_in.parameters()):
            p_ema.data.mul_(self.ema_decay).add_(p_online.data, alpha=1 - self.ema_decay)

    def forward(self, codec_emb, text_tokens, codec_mask=None, text_mask=None):
        """
        Args:
            codec_emb: [B, 128, T] — continuous EnCodec embeddings
            text_tokens: [B, T_text]
            codec_mask: [B, T] bool, True = padding
            text_mask: [B, T_text] bool
        """
        B = codec_emb.shape[0]
        T = codec_emb.shape[2]

        # codec_emb is [B, 128, T] → transpose to [B, T, 128]
        codec_seq = codec_emb.transpose(1, 2)

        # Text encoding
        text_emb = self.text_encoder(text_tokens, text_mask)

        # Project codec embeddings to predictor space
        z = self.proj_in(codec_seq)  # [B, T, d_model]

        # EMA targets (deterministic)
        with torch.no_grad():
            self.ema_proj_in.eval()
            target_emb = self.ema_proj_in(codec_seq)  # [B, T, d_model]
            self.ema_proj_in.train()

        # Prepend start embedding, shift right
        start = self.start_emb.expand(B, -1, -1)
        input_emb = torch.cat([start, z[:, :-1]], dim=1)  # [B, T, d_model]

        # Input noise
        if self.training and self.input_noise > 0:
            noise = torch.zeros_like(input_emb)
            noise[:, 1:] = torch.randn(B, T - 1, input_emb.shape[-1],
                                        device=input_emb.device) * self.input_noise
            input_emb = input_emb + noise

        # Masks
        if codec_mask is not None:
            start_mask = torch.zeros(B, 1, dtype=torch.bool, device=codec_emb.device)
            pred_mask = torch.cat([start_mask, codec_mask[:, :-1]], dim=1)
            loss_mask = codec_mask
        else:
            pred_mask = None
            loss_mask = None

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

        # ─── Loss 1: Prediction (MSE vs EMA targets) ───
        if loss_mask is not None:
            valid = (~loss_mask).unsqueeze(-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: Roundtrip (proj_in → proj_out → match original codec_emb) ───
        codec_recon = self.proj_out(z)  # [B, T, 128]

        if codec_mask is not None:
            valid_rt = (~codec_mask).unsqueeze(-1).float()
            roundtrip_loss = F.l1_loss(
                codec_recon * valid_rt, codec_seq * valid_rt, reduction="sum"
            ) / (valid_rt.sum() * self.codec_dim + 1e-8)
        else:
            roundtrip_loss = F.l1_loss(codec_recon, codec_seq)

        total_loss = self.pred_weight * prediction_loss + self.roundtrip_weight * roundtrip_loss

        return {
            "total_loss": total_loss,
            "prediction_loss": prediction_loss,
            "roundtrip_loss": roundtrip_loss,
        }

    def init_ar_cache(self):
        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):
        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, "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,
        }
    model = LeWMTTSv5(config)
    print(f"LeWM TTS v5: {count_parameters(model)/1e6:.2f}M trainable parameters")
    return model, config

-  - JEPAPredictor: causal transformer, predicts next embedding

+  - DurationPredictor: predicts per-character duration

+  - JEPAPredictor: causal transformer with frame-level text conditioning

+# ─── Duration Predictor ─────────────────────────────────────────────────────

+    """Predicts log-duration per text token. Similar to FastSpeech."""

+    def __init__(self, d_model=256, kernel_size=3, num_layers=2, dropout=0.1):

+                nn.Conv1d(d_model, d_model, kernel_size, padding=kernel_size // 2),

+        self.convs = nn.ModuleList(layers)

+        self.proj = nn.Linear(d_model, 1)

+    def forward(self, text_emb, text_mask=None):

+        Returns: log_durations [B, T_text]

+        x = text_emb.transpose(1, 2)  # [B, d_model, T]

+        for i in range(0, len(self.convs), 4):

+            x = self.convs[i](x)      # Conv1d

+            x = self.convs[i+1](x)    # ReLU

+            x = x.transpose(1, 2)     # [B, T, d_model] for LayerNorm

+            x = self.convs[i+2](x)    # LayerNorm

+            x = self.convs[i+3](x)    # Dropout

+            x = x.transpose(1, 2)     # back to [B, d_model, T]

+        log_dur = self.proj(x).squeeze(-1)  # [B, T_text]

+            log_dur = log_dur.masked_fill(text_mask, 0.0)

+def length_regulate(text_emb, durations, text_mask=None):

+    Expand text embeddings by repeating each token according to its duration.

+    durations: [B, T_text] integer durations

+    Returns: expanded [B, T_audio, d_model], where T_audio = sum(durations) per sample

+    max_audio_len = durations.sum(dim=1).max().item()

+    expanded = torch.zeros(B, max_audio_len, D, device=text_emb.device)

+            if text_mask is not None and text_mask[b, t]:

+            dur = durations[b, t].item()

+            if dur > 0 and pos < max_audio_len:

+                end = min(pos + dur, max_audio_len)

+                expanded[b, pos:end] = text_emb[b, t]

+    Compute uniform durations: spread audio frames evenly across text characters.

+    text_lengths: [B] number of valid text tokens

+    audio_lengths: [B] number of valid audio frames

+    Returns: durations [B, max_text_len] integer

+    max_text_len = text_lengths.max().item()

+    durations = torch.zeros(B, max_text_len, dtype=torch.long, device=text_lengths.device)

+    """Causal predictor with frame-level text added to audio input."""

+        self.text_input_proj = nn.Linear(d_model, d_model)

-        decoder_layer = nn.TransformerDecoderLayer(

+        # Use encoder (self-attention only) since text is now part of input

+        encoder_layer = nn.TransformerEncoderLayer(

-        self.transformer = nn.TransformerDecoder(decoder_layer, num_layers=num_layers)

+        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)

-    def forward(self, audio_emb, text_emb, audio_mask=None, text_mask=None):

-        x = self.audio_input_proj(audio_emb)

+    def forward(self, audio_emb, text_emb_expanded, audio_mask=None):

+        audio_emb: [B, T, d_model] — shifted audio embeddings

+        text_emb_expanded: [B, T, d_model] — frame-level text (already expanded)

+        x = self.audio_input_proj(audio_emb) + self.text_input_proj(text_emb_expanded)

-            x, text_emb, tgt_mask=causal_mask,

-            tgt_key_padding_mask=audio_mask, memory_key_padding_mask=text_mask,

+        predicted = self.transformer(x, mask=causal_mask, src_key_padding_mask=audio_mask)

+    # Keep old cross-attention interface for backward compat loading

+    def forward_cross_attn(self, audio_emb, text_emb, audio_mask=None, text_mask=None):

+        """Legacy cross-attention forward (not used in duration-predictor mode)."""

+        x = self.audio_input_proj(audio_emb)


Now update the KV-cache inference methods and the rest of the predictor:


    def _cached_mha(self, mha, q_input, kv_input, cache_k, cache_v):
        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)
        k_new = F.linear(kv_input, Wk, bk)
        v_new = F.linear(kv_input, Wv, bv)
        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, v = k_new, v_new
        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)
        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 _cached_cross_attn(self, mha, q_input, memory, cache_k, cache_v, memory_mask=None):
        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)
        if cache_k is None:
            k = F.linear(memory, Wk, bk)
            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:
            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):
        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):
        x = self.audio_input_proj(new_emb)
        x = x + self.pos_embed.pe[:, step_idx:step_idx + 1]
        for i, layer in enumerate(self.transformer.layers):
            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))
            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))
            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

    def _cached_mha(self, mha, q_input, kv_input, cache_k, cache_v):
        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)
        k_new = F.linear(kv_input, Wk, bk)
        v_new = F.linear(kv_input, Wv, bv)
        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, v = k_new, v_new
        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)
        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):
        return {
            'self_k': [None] * num_layers,
            'self_v': [None] * num_layers,
        }

    def inference_step(self, new_emb, text_emb_frame, step_idx, cache):
        """
        new_emb: [B, 1, d_model] — audio embedding for this step
        text_emb_frame: [B, 1, d_model] — expanded text embedding for this step
        """
        x = self.audio_input_proj(new_emb) + self.text_input_proj(text_emb_frame)
        x = x + self.pos_embed.pe[:, step_idx:step_idx + 1]
        for i, layer in enumerate(self.transformer.layers):
            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))
            ff_out = layer.linear2(layer.dropout(layer.activation(layer.linear1(x))))
            x = layer.norm2(x + layer.dropout2(ff_out))
        return self.output_proj(x), cache


    def _cached_mha(self, mha, q_input, kv_input, cache_k, cache_v):
        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)
        k_new = F.linear(kv_input, Wk, bk)
        v_new = F.linear(kv_input, Wv, bv)
        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, v = k_new, v_new
        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)
        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):
        return {
            'self_k': [None] * num_layers,
            'self_v': [None] * num_layers,
        }

    def inference_step(self, new_emb, text_emb_frame, step_idx, cache):
        """
        new_emb: [B, 1, d_model] — audio embedding for this step
        text_emb_frame: [B, 1, d_model] — expanded text embedding for this step
        """
        x = self.audio_input_proj(new_emb) + self.text_input_proj(text_emb_frame)
        x = x + self.pos_embed.pe[:, step_idx:step_idx + 1]
        for i, layer in enumerate(self.transformer.layers):
            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))
            ff_out = layer.linear2(layer.dropout(layer.activation(layer.linear1(x))))
            x = layer.norm2(x + layer.dropout2(ff_out))
        return self.output_proj(x), cache

"""
LeWM TTS v5 — JEPA with EnCodec backend.

Architecture:
  - EnCodec encoder/decoder (frozen): audio ↔ 128d embeddings @ 75Hz
  - Simple linear projections: 128d ↔ d_model (no lossy compression)
  - TextEncoder: byte-level transformer for Hindi text
  - DurationPredictor: predicts per-character duration
  - JEPAPredictor: causal transformer with frame-level text conditioning
"""

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

# ─── Duration Predictor ─────────────────────────────────────────────────────

class DurationPredictor(nn.Module):
    """Predicts log-duration per text token. Similar to FastSpeech."""
    def __init__(self, d_model=256, kernel_size=3, num_layers=2, dropout=0.1):
        super().__init__()
        layers = []
        for _ in range(num_layers):
            layers.extend([
                nn.Conv1d(d_model, d_model, kernel_size, padding=kernel_size // 2),
                nn.ReLU(),
                nn.LayerNorm(d_model),
                nn.Dropout(dropout),
            ])
        self.convs = nn.ModuleList(layers)
        self.proj = nn.Linear(d_model, 1)

    def forward(self, text_emb, text_mask=None):
        """
        text_emb: [B, T_text, d_model]
        Returns: log_durations [B, T_text]
        """
        x = text_emb.transpose(1, 2)  # [B, d_model, T]
        for i in range(0, len(self.convs), 4):
            x = self.convs[i](x)      # Conv1d
            x = self.convs[i+1](x)    # ReLU
            x = x.transpose(1, 2)     # [B, T, d_model] for LayerNorm
            x = self.convs[i+2](x)    # LayerNorm
            x = self.convs[i+3](x)    # Dropout
            x = x.transpose(1, 2)     # back to [B, d_model, T]
        x = x.transpose(1, 2)  # [B, T, d_model]
        log_dur = self.proj(x).squeeze(-1)  # [B, T_text]
        if text_mask is not None:
            log_dur = log_dur.masked_fill(text_mask, 0.0)
        return log_dur

def length_regulate(text_emb, durations, text_mask=None):
    """
    Expand text embeddings by repeating each token according to its duration.
    text_emb: [B, T_text, d_model]
    durations: [B, T_text] integer durations
    Returns: expanded [B, T_audio, d_model], where T_audio = sum(durations) per sample
    """
    B, T_text, D = text_emb.shape
    max_audio_len = durations.sum(dim=1).max().item()
    expanded = torch.zeros(B, max_audio_len, D, device=text_emb.device)
    for b in range(B):
        pos = 0
        for t in range(T_text):
            if text_mask is not None and text_mask[b, t]:
                continue
            dur = durations[b, t].item()
            if dur > 0 and pos < max_audio_len:
                end = min(pos + dur, max_audio_len)
                expanded[b, pos:end] = text_emb[b, t]
                pos = end
    return expanded

def compute_uniform_durations(text_lengths, audio_lengths):
    """
    Compute uniform durations: spread audio frames evenly across text characters.
    text_lengths: [B] number of valid text tokens
    audio_lengths: [B] number of valid audio frames
    Returns: durations [B, max_text_len] integer
    """
    B = text_lengths.shape[0]
    max_text_len = text_lengths.max().item()
    durations = torch.zeros(B, max_text_len, dtype=torch.long, device=text_lengths.device)
    for b in range(B):
        tl = text_lengths[b].item()
        al = audio_lengths[b].item()
        if tl == 0:
            continue
        base = al // tl
        remainder = al % tl
        durations[b, :tl] = base
        durations[b, :remainder] += 1
    return durations

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

class JEPAPredictor(nn.Module):
    """Causal predictor with frame-level text added to audio input."""
    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.text_input_proj = nn.Linear(d_model, d_model)
        self.pos_embed = SinusoidalPositionalEncoding(d_model, max_len=4096)
        self.dropout = nn.Dropout(dropout)

        # Use encoder (self-attention only) since text is now part of input
        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.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_expanded, audio_mask=None):
        """
        audio_emb: [B, T, d_model] — shifted audio embeddings
        text_emb_expanded: [B, T, d_model] — frame-level text (already expanded)
        """
        x = self.audio_input_proj(audio_emb) + self.text_input_proj(text_emb_expanded)
        x = self.pos_embed(x)
        x = self.dropout(x)
        T = x.shape[1]
        causal_mask = torch.triu(torch.ones(T, T, device=x.device, dtype=torch.bool), diagonal=1)
        predicted = self.transformer(x, mask=causal_mask, src_key_padding_mask=audio_mask)
        return self.output_proj(predicted)

    # Keep old cross-attention interface for backward compat loading
    def forward_cross_attn(self, audio_emb, text_emb, audio_mask=None, text_mask=None):
        """Legacy cross-attention forward (not used in duration-predictor mode)."""
        x = self.audio_input_proj(audio_emb)
        x = self.pos_embed(x)
        x = self.dropout(x)
        return self.output_proj(x)

    def _cached_mha(self, mha, q_input, kv_input, cache_k, cache_v):
        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)
        k_new = F.linear(kv_input, Wk, bk)
        v_new = F.linear(kv_input, Wv, bv)
        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, v = k_new, v_new
        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)
        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 _cached_cross_attn(self, mha, q_input, memory, cache_k, cache_v, memory_mask=None):
        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)
        if cache_k is None:
            k = F.linear(memory, Wk, bk)
            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:
            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):
        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):
        x = self.audio_input_proj(new_emb)
        x = x + self.pos_embed.pe[:, step_idx:step_idx + 1]
        for i, layer in enumerate(self.transformer.layers):
            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))
            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))
            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))

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

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

class TextEncoder(nn.Module):
    def __init__(self, vocab_size=256, d_model=256, nhead=4, num_layers=4, dropout=0.1):
        super().__init__()
        self.d_model = d_model
        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):
        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)
        return self.proj(x)

# ─── Full v5 Model ──────────────────────────────────────────────────────────

class LeWMTTSv5(nn.Module):
    """
    JEPA TTS with EnCodec backend.
    No down/upsampler — works directly at EnCodec's native 75Hz.
    Simple linear projections: 128d (codec) ↔ d_model (predictor).

    Training: codec_emb → proj_in → JEPA predict → loss vs EMA targets
              codec_emb → proj_in → proj_out → loss vs codec_emb (roundtrip)
    Inference: text → predictor (AR) → proj_out → EnCodec decoder → waveform
    """

    def __init__(self, config):
        super().__init__()
        d_model = config.get("d_model", 256)
        nhead = config.get("nhead", 4)
        codec_dim = config.get("codec_dim", 128)
        text_vocab = config.get("text_vocab_size", 256)
        text_layers = config.get("text_encoder_layers", 4)
        predictor_layers = config.get("predictor_layers", 6)
        dropout = config.get("dropout", 0.1)

        self.codec_dim = codec_dim

        self.text_encoder = TextEncoder(
            vocab_size=text_vocab, d_model=d_model, nhead=nhead,
            num_layers=text_layers, dropout=dropout,
        )

        # Simple linear projections — no lossy compression
        self.proj_in = nn.Sequential(
            nn.Linear(codec_dim, d_model),
            nn.GELU(),
            nn.Linear(d_model, d_model),
        )
        self.proj_out = nn.Sequential(
            nn.Linear(d_model, d_model),
            nn.GELU(),
            nn.Linear(d_model, codec_dim),
        )

        self.predictor = JEPAPredictor(
            d_model=d_model, nhead=nhead,
            num_layers=predictor_layers, dropout=dropout,
        )

        # EMA on proj_in (provides stable prediction targets)
        self.ema_proj_in = copy.deepcopy(self.proj_in)
        for p in self.ema_proj_in.parameters():
            p.requires_grad = False
        self.ema_decay = config.get("ema_decay", 0.998)

        # Learnable start embedding
        self.start_emb = nn.Parameter(torch.randn(1, 1, d_model) * 0.02)

        # Loss weights
        self.pred_weight = config.get("pred_weight", 10.0)
        self.roundtrip_weight = config.get("roundtrip_weight", 1.0)

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

        self.config = config

    @torch.no_grad()
    def update_ema(self):
        for p_online, p_ema in zip(self.proj_in.parameters(),
                                   self.ema_proj_in.parameters()):
            p_ema.data.mul_(self.ema_decay).add_(p_online.data, alpha=1 - self.ema_decay)

    def forward(self, codec_emb, text_tokens, codec_mask=None, text_mask=None):
        """
        Args:
            codec_emb: [B, 128, T] — continuous EnCodec embeddings
            text_tokens: [B, T_text]
            codec_mask: [B, T] bool, True = padding
            text_mask: [B, T_text] bool
        """
        B = codec_emb.shape[0]
        T = codec_emb.shape[2]

        # codec_emb is [B, 128, T] → transpose to [B, T, 128]
        codec_seq = codec_emb.transpose(1, 2)

        # Text encoding
        text_emb = self.text_encoder(text_tokens, text_mask)

        # Project codec embeddings to predictor space
        z = self.proj_in(codec_seq)  # [B, T, d_model]

        # EMA targets (deterministic)
        with torch.no_grad():
            self.ema_proj_in.eval()
            target_emb = self.ema_proj_in(codec_seq)  # [B, T, d_model]
            self.ema_proj_in.train()

        # Prepend start embedding, shift right
        start = self.start_emb.expand(B, -1, -1)
        input_emb = torch.cat([start, z[:, :-1]], dim=1)  # [B, T, d_model]

        # Input noise
        if self.training and self.input_noise > 0:
            noise = torch.zeros_like(input_emb)
            noise[:, 1:] = torch.randn(B, T - 1, input_emb.shape[-1],
                                        device=input_emb.device) * self.input_noise
            input_emb = input_emb + noise

        # Masks
        if codec_mask is not None:
            start_mask = torch.zeros(B, 1, dtype=torch.bool, device=codec_emb.device)
            pred_mask = torch.cat([start_mask, codec_mask[:, :-1]], dim=1)
            loss_mask = codec_mask
        else:
            pred_mask = None
            loss_mask = None

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

        # ─── Loss 1: Prediction (MSE vs EMA targets) ───
        if loss_mask is not None:
            valid = (~loss_mask).unsqueeze(-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: Roundtrip (proj_in → proj_out → match original codec_emb) ───
        codec_recon = self.proj_out(z)  # [B, T, 128]

        if codec_mask is not None:
            valid_rt = (~codec_mask).unsqueeze(-1).float()
            roundtrip_loss = F.l1_loss(
                codec_recon * valid_rt, codec_seq * valid_rt, reduction="sum"
            ) / (valid_rt.sum() * self.codec_dim + 1e-8)
        else:
            roundtrip_loss = F.l1_loss(codec_recon, codec_seq)

        total_loss = self.pred_weight * prediction_loss + self.roundtrip_weight * roundtrip_loss

        return {
            "total_loss": total_loss,
            "prediction_loss": prediction_loss,
            "roundtrip_loss": roundtrip_loss,
        }

    def init_ar_cache(self):
        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):
        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, "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,
        }
    model = LeWMTTSv5(config)
    print(f"LeWM TTS v5: {count_parameters(model)/1e6:.2f}M trainable parameters")
    return model, config

-    def _cached_cross_attn(self, mha, q_input, memory, cache_k, cache_v, memory_mask=None):

-        Wq, Wk, Wv = mha.in_proj_weight.chunk(3, dim=0)

-        bq, bk, bv = mha.in_proj_bias.chunk(3, dim=0)

-            k = F.linear(memory, Wk, bk)

-            v = F.linear(memory, Wv, bv)

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

-            attn = attn.masked_fill(memory_mask.unsqueeze(1).unsqueeze(2), float('-inf'))

-        out = out.transpose(1, 2).contiguous().view(B, 1, d)

-        out = F.linear(out, mha.out_proj.weight, mha.out_proj.bias)

-            'cross_k': [None] * num_layers,

-            'cross_v': [None] * num_layers,

-    def inference_step(self, new_emb, text_emb, step_idx, cache, text_mask=None):

-        x = self.audio_input_proj(new_emb)

+    def inference_step(self, new_emb, text_emb_frame, step_idx, cache):

+        new_emb: [B, 1, d_model] — audio embedding for this step

+        text_emb_frame: [B, 1, d_model] — expanded text embedding for this step

+        x = self.audio_input_proj(new_emb) + self.text_input_proj(text_emb_frame)

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

-            x = layer.norm2(x + layer.dropout2(ca_out))

-            x = layer.norm3(x + layer.dropout3(ff_out))

+            x = layer.norm2(x + layer.dropout2(ff_out))


Now update the LeWMTTSv5 class to include the duration predictor:


class LeWMTTSv5(nn.Module):
    """
    JEPA TTS with EnCodec backend.
    No down/upsampler — works directly at EnCodec's native 75Hz.
    Simple linear projections: 128d (codec) ↔ d_model (predictor).

    Training: codec_emb → proj_in → JEPA predict → loss vs EMA targets
              codec_emb → proj_in → proj_out → loss vs codec_emb (roundtrip)
    Inference: text → predictor (AR) → proj_out → EnCodec decoder → waveform
    """

    def __init__(self, config):
        super().__init__()
        d_model = config.get("d_model", 256)
        nhead = config.get("nhead", 4)
        codec_dim = config.get("codec_dim", 128)
        text_vocab = config.get("text_vocab_size", 256)
        text_layers = config.get("text_encoder_layers", 4)
        predictor_layers = config.get("predictor_layers", 6)
        dropout = config.get("dropout", 0.1)

        self.codec_dim = codec_dim

        self.text_encoder = TextEncoder(
            vocab_size=text_vocab, d_model=d_model, nhead=nhead,
            num_layers=text_layers, dropout=dropout,
        )

        # Simple linear projections — no lossy compression
        self.proj_in = nn.Sequential(
            nn.Linear(codec_dim, d_model),
            nn.GELU(),
            nn.Linear(d_model, d_model),
        )
        self.proj_out = nn.Sequential(
            nn.Linear(d_model, d_model),
            nn.GELU(),
            nn.Linear(d_model, codec_dim),
        )

        self.predictor = JEPAPredictor(
            d_model=d_model, nhead=nhead,
            num_layers=predictor_layers, dropout=dropout,
        )

        # EMA on proj_in (provides stable prediction targets)
        self.ema_proj_in = copy.deepcopy(self.proj_in)
        for p in self.ema_proj_in.parameters():
            p.requires_grad = False
        self.ema_decay = config.get("ema_decay", 0.998)

        # Learnable start embedding
        self.start_emb = nn.Parameter(torch.randn(1, 1, d_model) * 0.02)

        # Loss weights
        self.pred_weight = config.get("pred_weight", 10.0)
        self.roundtrip_weight = config.get("roundtrip_weight", 1.0)

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

        self.config = config

class LeWMTTSv5(nn.Module):
    """
    JEPA TTS with EnCodec backend + duration predictor.
    Text is expanded to frame level and added to audio input — can't be ignored.

    Training: codec_emb → proj_in → predict(audio + aligned_text) → loss vs targets
              Duration predictor trained with uniform alignment as supervision.
    Inference: text → duration predict → expand text → AR predict → proj_out → decode
    """

    def __init__(self, config):
        super().__init__()
        d_model = config.get("d_model", 256)
        nhead = config.get("nhead", 4)
        codec_dim = config.get("codec_dim", 128)
        text_vocab = config.get("text_vocab_size", 256)
        text_layers = config.get("text_encoder_layers", 4)
        predictor_layers = config.get("predictor_layers", 6)
        dropout = config.get("dropout", 0.1)

        self.codec_dim = codec_dim

        self.text_encoder = TextEncoder(
            vocab_size=text_vocab, d_model=d_model, nhead=nhead,
            num_layers=text_layers, dropout=dropout,
        )

        # Duration predictor
        self.duration_predictor = DurationPredictor(
            d_model=d_model, kernel_size=3, num_layers=2, dropout=dropout,
        )

        # Simple linear projections — no lossy compression
        self.proj_in = nn.Sequential(
            nn.Linear(codec_dim, d_model),
            nn.GELU(),
            nn.Linear(d_model, d_model),
        )
        self.proj_out = nn.Sequential(
            nn.Linear(d_model, d_model),
            nn.GELU(),
            nn.Linear(d_model, codec_dim),
        )

        self.predictor = JEPAPredictor(
            d_model=d_model, nhead=nhead,
            num_layers=predictor_layers, dropout=dropout,
        )

        # EMA on proj_in (provides stable prediction targets)
        self.ema_proj_in = copy.deepcopy(self.proj_in)
        for p in self.ema_proj_in.parameters():
            p.requires_grad = False
        self.ema_decay = config.get("ema_decay", 0.998)

        # Learnable start embedding
        self.start_emb = nn.Parameter(torch.randn(1, 1, d_model) * 0.02)

        # Loss weights
        self.pred_weight = config.get("pred_weight", 10.0)
        self.roundtrip_weight = config.get("roundtrip_weight", 1.0)
        self.dur_weight = config.get("dur_weight", 1.0)

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

        self.config = config