batch_size=64: OK, GPU=7.2GB, pred=1.0555, rt=3.7690
batch_size=128: OK, GPU=15.9GB, pred=1.0479, rt=3.7867
batch_size=256: OK, GPU=32.3GB, pred=1.0355, rt=3.7771
batch_size=384: OK, GPU=48.3GB, pred=1.0791, rt=3.7733
Shell cwd was reset to /home/ubuntu/pipecat


batch_size=384 uses 48GB — still has room. Let me try 512:


Bash

cd /home/ubuntu/lewm-tts && python3 -c "
import torch
from model_v5 import build_model
from dataset_v5 import build_dataloader

device = torch.device('cuda')
torch.cuda.reset_peak_memory_stats()
model, config = build_model()
model = model.to(device)
model.train()
loader, _ = build_dataloader('processed_data_codec/manifest.json', batch_size=512)
batch = next(iter(loader))
losses = model(batch['codec_emb'].to(device), batch['text_tokens'].to(device),
               batch['codec_mask'].to(device), batch['text_mask'].to(device))
losses['total_loss'].backward()
print(f'batch_size=512: GPU={torch.cuda.max_memory_allocated()/1e9:.1f}GB')
" 2>&1 | grep -E "batch_size|OOM"

Test batch_size=512


batch_size=512: GPU=66.8GB
Shell cwd was reset to /home/ubuntu/pipecat


512 uses 67GB — fits well on A100 80GB. Let's launch v5 training:


cd /home/ubuntu/lewm-tts && nohup python3 -u train_v5.py \
    --data_dir processed_data_codec \
    --output_dir output_v5 \
    --epochs 200 \
    --batch_size 512 \
    --lr 1e-4 \
    --pred_weight 10.0 \
    --roundtrip_weight 1.0 \
    --ema_decay 0.998 \
    --input_noise_max 0.1 \
    --ds_stride 3 \
    --num_workers 4 \
    --print_every 20 \
    --log_every 5 \
    --save_every 10 \
    > train_v5.log 2>&1 &
echo "PID: $!"

Launch v5 training with batch_size=512


PID: 3690515


sleep 15 && tail -10 /home/ubuntu/lewm-tts/train_v5.log

Check v5 training started


Device: cuda
LeWM TTS v5: 14.80M trainable parameters
CodecDataset: 6879 samples (max 900 frames)
Training: 6879 samples, 13 batches/epoch
/home/ubuntu/.local/lib/python3.10/site-packages/torch/nn/modules/transformer.py:515: UserWarning: The PyTorch API of nested tensors is in prototype stage and will change in the near future. We recommend specifying layout=torch.jagged when constructing a nested tensor, as this layout receives active development, has better operator coverage, and works with torch.compile. (Triggered internally at /pytorch/aten/src/ATen/NestedTensorImpl.cpp:178.)
  output = torch._nested_tensor_from_mask(


sleep 30 && tail -15 /home/ubuntu/lewm-tts/train_v5.log

Check training progress


Device: cuda
LeWM TTS v5: 14.80M trainable parameters
CodecDataset: 6879 samples (max 900 frames)
Training: 6879 samples, 13 batches/epoch
/home/ubuntu/.local/lib/python3.10/site-packages/torch/nn/modules/transformer.py:515: UserWarning: The PyTorch API of nested tensors is in prototype stage and will change in the near future. We recommend specifying layout=torch.jagged when constructing a nested tensor, as this layout receives active development, has better operator coverage, and works with torch.compile. (Triggered internally at /pytorch/aten/src/ATen/NestedTensorImpl.cpp:178.)
  output = torch._nested_tensor_from_mask(
Epoch 1/200 | avg_loss=13.8438 pred=1.0075 rt=3.7691 | 21.1s
  New best (loss=13.8438)
  [2/200] step 20 | loss=13.0210 pred=0.9269 rt=3.7518 | grad=10.34 | 11.1s
Epoch 2/200 | avg_loss=13.0226 pred=0.9271 rt=3.7518 | 19.8s
  New best (loss=13.0226)


V5 training is running. ~20s/epoch. Let me wait for a few more epochs to see the trend:


sleep 120 && tail -20 /home/ubuntu/lewm-tts/train_v5.log

Check v5 training after a few minutes


  New best (loss=13.8438)
  [2/200] step 20 | loss=13.0210 pred=0.9269 rt=3.7518 | grad=10.34 | 11.1s
Epoch 2/200 | avg_loss=13.0226 pred=0.9271 rt=3.7518 | 19.8s
  New best (loss=13.0226)
Epoch 3/200 | avg_loss=12.0493 pred=0.8338 rt=3.7111 | 19.9s
  New best (loss=12.0493)
  [4/200] step 40 | loss=11.5790 pred=0.7907 rt=3.6718 | grad=8.18 | 2.2s
Epoch 4/200 | avg_loss=11.1822 pred=0.7578 rt=3.6042 | 19.8s
  New best (loss=11.1822)
  [5/200] step 60 | loss=10.0740 pred=0.6856 rt=3.2181 | grad=8.77 | 12.3s
Epoch 5/200 | avg_loss=10.1431 pred=0.6907 rt=3.2360 | 19.7s
  New best (loss=10.1431)
Epoch 6/200 | avg_loss=8.5957 pred=0.6251 rt=2.3442 | 20.0s
  New best (loss=8.5957)
  [7/200] step 80 | loss=7.7253 pred=0.5840 rt=1.8848 | grad=10.12 | 3.6s
Epoch 7/200 | avg_loss=7.2125 pred=0.5581 rt=1.6318 | 19.8s
  New best (loss=7.2125)
  [8/200] step 100 | loss=5.9818 pred=0.4731 rt=1.2512 | grad=8.81 | 14.0s
Epoch 8/200 | avg_loss=6.1338 pred=0.4849 rt=1.2845 | 19.9s
  New best (loss=6.1338)


Losses are dropping fast:

| Epoch | Pred Loss | Roundtrip Loss | Total |
|-------|-----------|----------------|-------|
| 1 | 1.008 | 3.77 | 13.84 |
| 4 | 0.758 | 3.60 | 11.18 |
| 8 | 0.485 | 1.28 | 6.13 |

Both prediction and roundtrip losses are converging well. The roundtrip (downsample→upsample) is learning quickly — from 3.77 to 1.28 in 8 epochs.


Updated task #6 status


14.8M trainable params, 20s/epoch. Training should take ~67 minutes for 200 epochs. I'll check back at a milestone to generate samples. The key advantage here: even if prediction isn't perfect, EnCodec's decoder will produce clean audio (no mel decoder noise).


Edit

## Training Command (v4)
```bash
python3 train.py \
  --data_dir processed_data_100mel \
  --output_dir output_v4 \
  --epochs 200 --batch_size 256 \
  --pred_weight 10.0 --kl_weight 0.05 \
  --recon_weight 1.0 --spectral_weight 0.5 \
  --ema_decay 0.998 --free_bits 1.0 --input_noise_max 0.2
```

## Architecture
TextEncoder(4L,256d) + AudioEncoder(CNN+4L,256d,stride4) + JEPAPredictor(6L decoder) + MelDecoder(512h,ResConv) → Vocos 24kHz

## Dataset
7009 segments, 17.88h single-speaker Hindi, 100-mel, processed_data_100mel/

## v4 Results
- MelDecoder was the bottleneck (L1=0.57 even with perfect embeddings)
- JEPA predictor worked well: cosine_sim=0.97, 6.2x better than copy baseline
- Recon had noise overlay, synthesis was pure noise

## v5 Architecture (codec-based)
- EnCodec 24kHz (frozen): audio ↔ 128d embeddings @ 75Hz, near-transparent quality
- CodecDownsampler(stride3): 75Hz×128d → 25Hz×256d (trainable, 2L transformer)
- CodecUpsampler(stride3): 25Hz×256d → 75Hz×128d (trainable)
- TextEncoder(4L,256d) + JEPAPredictor(6L decoder)
- EMA on downsampler, learnable start_emb
- Losses: pred_weight=10.0 (MSE vs EMA targets) + roundtrip_weight=1.0 (L1 down→up vs codec_emb)
- 14.8M trainable params, ~20s/epoch with batch_size=512 on A100
- Files: model_v5.py, train_v5.py, dataset_v5.py, inference_v5.py, preprocess_codec.py
- Data: processed_data_codec/ (7009 samples, EnCodec embeddings)

## Dataset
7009 segments, 17.88h single-speaker Hindi, processed_data_100mel/ and processed_data_codec/


The file /home/ubuntu/.claude/projects/-home-ubuntu-pipecat/memory/lewm-tts.md has been updated successfully.

## v4 Results
- MelDecoder was the bottleneck (L1=0.57 even with perfect embeddings)
- JEPA predictor worked well: cosine_sim=0.97, 6.2x better than copy baseline
- Recon had noise overlay, synthesis was pure noise

## v5 Architecture (codec-based)
- EnCodec 24kHz (frozen): audio ↔ 128d embeddings @ 75Hz, near-transparent quality
- CodecDownsampler(stride3): 75Hz×128d → 25Hz×256d (trainable, 2L transformer)
- CodecUpsampler(stride3): 25Hz×256d → 75Hz×128d (trainable)
- TextEncoder(4L,256d) + JEPAPredictor(6L decoder)
- EMA on downsampler, learnable start_emb
- Losses: pred_weight=10.0 (MSE vs EMA targets) + roundtrip_weight=1.0 (L1 down→up vs codec_emb)
- 14.8M trainable params, ~20s/epoch with batch_size=512 on A100
- Files: model_v5.py, train_v5.py, dataset_v5.py, inference_v5.py, preprocess_codec.py
- Data: processed_data_codec/ (7009 samples, EnCodec embeddings)

## Dataset
7009 segments, 17.88h single-speaker Hindi, processed_data_100mel/ and processed_data_codec/

# LeWM TTS - JEPA-based Hindi TTS

## Key Issues Found (v3/v3b/v3c)
1. Missing EMA target encoder (core JEPA component)
2. Prediction loss was <1% of total loss — predictor learned identity mapping
3. KL divergence exploded to 2.45 (kl_weight=0.01 too low)
4. v3c's downsample_factor=2 made consecutive embeddings too similar
5. Random start embedding at inference (never trained)
6. No input noise to bridge train/inference gap

## v4 Fixes Applied
- **EMA target encoder**: deepcopy of AudioEncoder, updated with momentum 0.998
- **Rebalanced losses**: pred_weight=10.0 (now ~36% of total), kl_weight=0.05
- **Free-bits KL**: threshold=1.0, prevents both collapse and explosion
- **Learnable start_emb**: prepended to training sequence, receives gradients
- **Input noise**: annealed 0→0.2 over first 30% of training
- **Downsample factor**: back to 4 (more temporal separation)

## Training Command (v4)
```bash
python3 train.py \
  --data_dir processed_data_100mel \
  --output_dir output_v4 \
  --epochs 200 --batch_size 256 \
  --pred_weight 10.0 --kl_weight 0.05 \
  --recon_weight 1.0 --spectral_weight 0.5 \
  --ema_decay 0.998 --free_bits 1.0 --input_noise_max 0.2
```

## Architecture
TextEncoder(4L,256d) + AudioEncoder(CNN+4L,256d,stride4) + JEPAPredictor(6L decoder) + MelDecoder(512h,ResConv) → Vocos 24kHz

## Dataset
7009 segments, 17.88h single-speaker Hindi, 100-mel, processed_data_100mel/

-  --pred_weight 10.0 --kl_weight 0.05 \

-  --recon_weight 1.0 --spectral_weight 0.5 \

-  --ema_decay 0.998 --free_bits 1.0 --input_noise_max 0.2

+- MelDecoder was the bottleneck (L1=0.57 even with perfect embeddings)

+- JEPA predictor worked well: cosine_sim=0.97, 6.2x better than copy baseline

+- Recon had noise overlay, synthesis was pure noise

-TextEncoder(4L,256d) + AudioEncoder(CNN+4L,256d,stride4) + JEPAPredictor(6L decoder) + MelDecoder(512h,ResConv) → Vocos 24kHz

+- EnCodec 24kHz (frozen): audio ↔ 128d embeddings @ 75Hz, near-transparent quality

+- CodecDownsampler(stride3): 75Hz×128d → 25Hz×256d (trainable, 2L transformer)

+- CodecUpsampler(stride3): 25Hz×256d → 75Hz×128d (trainable)

+- TextEncoder(4L,256d) + JEPAPredictor(6L decoder)

+- EMA on downsampler, learnable start_emb

+- Losses: pred_weight=10.0 (MSE vs EMA targets) + roundtrip_weight=1.0 (L1 down→up vs codec_emb)

+- 14.8M trainable params, ~20s/epoch with batch_size=512 on A100

+- Files: model_v5.py, train_v5.py, dataset_v5.py, inference_v5.py, preprocess_codec.py

+- Data: processed_data_codec/ (7009 samples, EnCodec embeddings)

-7009 segments, 17.88h single-speaker Hindi, 100-mel, processed_data_100mel/

+7009 segments, 17.88h single-speaker Hindi, processed_data_100mel/ and processed_data_codec/


Training is running. Want me to check progress and generate samples once it hits epoch 50 or so?


status


Bash

Check v5 training progress


Epoch 31/200 | avg_loss=1.4079 pred=0.0610 rt=0.7982 | 19.9s
  New best (loss=1.4079)
Epoch 32/200 | avg_loss=1.3721 pred=0.0579 rt=0.7930 | 19.8s
  New best (loss=1.3721)
  [33/200] step 420 | loss=1.3376 pred=0.0553 rt=0.7851 | grad=1.72 | 6.7s
Epoch 33/200 | avg_loss=1.3387 pred=0.0549 rt=0.7894 | 20.0s
  New best (loss=1.3387)
  [34/200] step 440 | loss=1.2999 pred=0.0516 rt=0.7843 | grad=3.03 | 16.8s
Epoch 34/200 | avg_loss=1.3111 pred=0.0525 rt=0.7862 | 19.7s
  New best (loss=1.3111)


Epoch 34/200. Losses dropping steadily:

| Epoch | Pred Loss | Roundtrip Loss | Total |
|-------|-----------|----------------|-------|
| 1 | 1.008 | 3.77 | 13.84 |
| 8 | 0.485 | 1.28 | 6.13 |
| 34 | 0.053 | 0.79 | 1.31 |

Prediction loss already at **0.053** — lower than v4 ever achieved (0.073 at epoch 120). Roundtrip loss still converging. About 55 minutes left.


what we using ?


The user is asking what codec/model we're using for v5.


**EnCodec 24kHz** (from Meta/Facebook) as the audio backbone.

The setup:
- **EnCodec encoder** (frozen) — converts audio → 128-dim embeddings at 75Hz. Near-transparent quality (L1=0.01 vs our old mel decoder's 0.53)
- **Our JEPA predictor** — predicts next embedding from text + previous embeddings (same architecture as before, just working in codec space now)
- **Stride-3 down/upsampler** (trainable) — compresses 75Hz → 25Hz for efficient AR prediction, then back up
- **EnCodec decoder** (frozen) — converts predicted embeddings back to audio

The key insight: our JEPA predictor was already working well (cosine_sim=0.97). The problem was our weak mel decoder + Vocos pipeline adding noise. EnCodec's decoder is pre-trained on massive data — it produces clean audio as long as the embeddings are reasonable.

14.8M trainable params. EnCodec itself is frozen, we only train the text encoder, predictor, and down/upsampler.


hw much total szie


cd /home/ubuntu/lewm-tts && python3 -c "
from model_v5 import LeWMTTSv5
from encodec import EncodecModel

# Trainable params
config = {'d_model': 256, 'nhead': 4, 'codec_dim': 128, 'text_vocab_size': 256,
          'text_encoder_layers': 4, 'predictor_layers': 6, 'dropout': 0.1, 'ds_stride': 3}
model = LeWMTTSv5(config)
trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
total_ours = sum(p.numel() for p in model.parameters())

# EnCodec params
codec = EncodecModel.encodec_model_24khz()
codec_params = sum(p.numel() for p in codec.parameters())

print(f'Our trainable:  {trainable/1e6:.2f}M')
print(f'Our total (inc EMA): {total_ours/1e6:.2f}M')
print(f'EnCodec (frozen): {codec_params/1e6:.2f}M')
print(f'Grand total:  {(total_ours + codec_params)/1e6:.2f}M')
" 2>&1 | grep -v Warning

Count total model parameters


  WeightNorm.apply(module, name, dim)
Our trainable:  14.80M
Our total (inc EMA): 18.12M
EnCodec (frozen): 14.85M
Grand total:  32.97M
Shell cwd was reset to /home/ubuntu/pipecat


- **Our model**: 14.8M trainable + 3.3M EMA (frozen copy) = 18.1M
- **EnCodec**: 14.9M (frozen, pretrained)
- **Grand total: ~33M params**

At inference you need both (~33M). For comparison, v4 was ~24M but with a garbage mel decoder. This is still very lightweight for a TTS system.


run samples on latest epoch


tail -3 /home/ubuntu/lewm-tts/train_v5.log && ls -t /home/ubuntu/lewm-tts/output_v5/checkpoints/ | head -3

Check latest epoch and checkpoint


  New best (loss=1.2359)
Epoch 38/200 | avg_loss=1.2157 pred=0.0447 rt=0.7684 | 20.2s
  New best (loss=1.2157)
best.pt
epoch_0030.pt
epoch_0020.pt


Only epoch 30 checkpoint saved so far (epoch 38 running, saves every 10). Let me use `best.pt` which is the latest:


cd /home/ubuntu/lewm-tts && python3 -c "
import torch, torchaudio, numpy as np, json
from model_v5 import LeWMTTSv5
from encodec import EncodecModel
from encodec.utils import convert_audio

device = torch.device('cuda:0')

# Load model
ckpt = torch.load('output_v5/checkpoints/best.pt', map_location='cpu', weights_only=False)
config = ckpt['config']
model = LeWMTTSv5(config).to(device)
model.load_state_dict(ckpt['model'])
model.eval()
print(f'Loaded epoch {ckpt[\"epoch\"]+1}, loss={ckpt[\"loss\"]:.4f}')

# Load EnCodec
codec = EncodecModel.encodec_model_24khz()
codec.set_target_bandwidth(6.0)
codec.eval().to(device)

manifest = json.load(open('processed_data_codec/manifest.json'))
import os; os.makedirs('samples_v5', exist_ok=True)

def save_wav(waveform, path, sr=24000):
    mx = np.abs(waveform).max()
    if mx > 0: waveform = waveform / mx * 0.95
    torchaudio.save(path, torch.from_numpy(waveform).unsqueeze(0), sr)
    print(f'Saved: {path} ({len(waveform)/sr:.2f}s)')

def text_to_tokens(text):
    return torch.tensor(list(text.encode('utf-8')), dtype=torch.long).unsqueeze(0).to(device)

entry = manifest[0]
text = entry['text']
audio_path = entry['audio_path']
print(f'Text: {text[:80]}...')

# 1. Ground truth (original audio)
wav, sr = torchaudio.load(audio_path)
wav_24k = convert_audio(wav, sr, 24000, 1)
save_wav(wav_24k.squeeze().numpy(), 'samples_v5/gt_0.wav')

# 2. Pure EnCodec roundtrip (baseline — best possible quality)
with torch.no_grad():
    codec_emb = codec.encoder(wav_24k.unsqueeze(0).to(device))
    wav_codec = codec.decoder(codec_emb).squeeze().cpu().numpy()
save_wav(wav_codec, 'samples_v5/codec_only_0.wav')

# 3. Roundtrip: codec → downsample → upsample → codec decode
with torch.no_grad():
    codec_emb = codec.encoder(wav_24k.unsqueeze(0).to(device))
    z, _ = model.downsampler(codec_emb)
    codec_recon = model.upsampler(z, target_len=codec_emb.shape[2])
    wav_rt = codec.decoder(codec_recon).squeeze().cpu().numpy()
save_wav(wav_rt, 'samples_v5/roundtrip_0.wav')

# 4. Teacher-forced prediction: codec → downsample → predict → upsample → decode
with torch.no_grad():
    codec_emb = codec.encoder(wav_24k.unsqueeze(0).to(device))
    z, _ = model.downsampler(codec_emb)
    text_emb = model.text_encoder(text_to_tokens(text))
    start = model.start_emb.expand(1, -1, -1)
    input_emb = torch.cat([start, z[:, :-1]], dim=1)
    predicted = model.predictor(input_emb, text_emb)
    codec_pred = model.upsampler(predicted, target_len=codec_emb.shape[2])
    wav_tf = codec.decoder(codec_pred).squeeze().cpu().numpy()
save_wav(wav_tf, 'samples_v5/tf_0.wav')

# 5. AR synthesis temp=0.0
with torch.no_grad():
    text_emb = model.text_encoder(text_to_tokens(text))
    start_emb = model.start_emb.to(device)
    all_embs = [start_emb]
    cache = model.init_ar_cache()
    next_emb, cache = model.predict_next_cached(start_emb, text_emb, 0, cache)
    all_embs.append(next_emb)
    for step in range(1, 300):
        next_emb, cache = model.predict_next_cached(next_emb, text_emb, step, cache)
        all_embs.append(next_emb)
        if step > 30:
            norms = [e.norm().item() for e in all_embs[-20:]]
            if np.std(norms) < 0.01 * np.mean(norms): break
    pred_embs = torch.cat(all_embs, dim=1)
    codec_synth = model.upsampler(pred_embs)
    wav_synth = codec.decoder(codec_synth).squeeze().cpu().numpy()
save_wav(wav_synth, 'samples_v5/synth_t00.wav')
print(f'Synthesis: {step+1} steps')

# 6. AR synthesis temp=0.3
with torch.no_grad():
    all_embs = [start_emb]
    cache = model.init_ar_cache()
    next_emb, cache = model.predict_next_cached(start_emb, text_emb, 0, cache)
    next_emb = next_emb + torch.randn_like(next_emb) * 0.3
    all_embs.append(next_emb)
    for step in range(1, 300):
        next_emb, cache = model.predict_next_cached(next_emb, text_emb, step, cache)
        next_emb = next_emb + torch.randn_like(next_emb) * 0.3
        all_embs.append(next_emb)
        if step > 30:
            norms = [e.norm().item() for e in all_embs[-20:]]
            if np.std(norms) < 0.01 * np.mean(norms): break
    pred_embs = torch.cat(all_embs, dim=1)
    codec_synth = model.upsampler(pred_embs)
    wav_synth = codec.decoder(codec_synth).squeeze().cpu().numpy()
save_wav(wav_synth, 'samples_v5/synth_t03.wav')
print(f'Synthesis t=0.3: {step+1} steps')

print('\\nAll samples in samples_v5/')
" 2>&1

Generate v5 audio samples from latest checkpoint


/home/ubuntu/.local/lib/python3.10/site-packages/torch/nn/utils/weight_norm.py:144: FutureWarning: `torch.nn.utils.weight_norm` is deprecated in favor of `torch.nn.utils.parametrizations.weight_norm`.
  WeightNorm.apply(module, name, dim)
Loaded epoch 39, loss=1.1969
Text: मेरे प्यारे देशवासियों, नमस्कार। मन की बात में आपका स्वागत है, अभिन...
Saved: samples_v5/gt_0.wav (9.52s)
Saved: samples_v5/codec_only_0.wav (9.52s)
Saved: samples_v5/roundtrip_0.wav (9.52s)
Saved: samples_v5/tf_0.wav (9.52s)
Saved: samples_v5/synth_t00.wav (1.33s)
Synthesis: 32 steps
Saved: samples_v5/synth_t03.wav (12.05s)
Synthesis t=0.3: 300 steps

All samples in samples_v5/
Shell cwd was reset to /home/ubuntu/pipecat


Samples in `samples_v5/`:

| File | Description |
|------|-------------|
| `gt_0.wav` | Ground truth (9.52s) |
| `codec_only_0.wav` | Pure EnCodec roundtrip — best possible quality (9.52s) |
| `roundtrip_0.wav` | Codec → downsample → upsample → decode (9.52s) |
| `tf_0.wav` | Teacher-forced prediction (9.52s) |
| `synth_t00.wav` | AR synthesis, temp=0.0 (1.3s — stopped early) |
| `synth_t03.wav` | AR synthesis, temp=0.3 (12s) |

The AR synthesis at temp=0.0 collapsed to 1.3s (32 steps) — the predictor converged too fast to a fixed point. This is expected at epoch 39, it'll improve. Temp=0.3 ran the full 300 steps.

Key thing to listen for: compare `codec_only_0.wav` vs `roundtrip_0.wav` — this tells you how good the down/upsampler is. And `tf_0.wav` shows the prediction quality.


pure noise .. and y we not used codec like nanocode or speechtoknizzer codec?


The user is saying the samples are pure noise, and asking why we didn't use a speech-specific codec like SpeechTokenizer or something like "nanocodec".

Let me think about this:

1. The roundtrip loss at 0.77 is still high - meaning the down/upsampler isn't reconstructing the codec embeddings well yet. This could explain noise.

2. More importantly, the user is suggesting speech-specific codecs. EnCodec is a general audio codec. Speech-specific codecs like SpeechTokenizer, or semantic codecs might work better for TTS.

Let me check what speech codecs are available. The user mentioned "nanocodec" and "speechtokenizer". Let me check what's installed.

Actually, I'm realizing the core issue might be that EnCodec's continuous embeddings are too high-dimensional and complex for prediction, whereas a speech-specific codec would have embeddings designed to capture speech structure more efficiently. The down/upsampler is adding another layer of potential error on top of that. Let me see what codecs are actually available to work with. The upsampler's stride-3 transposed convolution might be creating checkerboard artifacts, and the roundtrip loss of 0.77 could actually be quite high relative to the embedding magnitudes—especially since the roundtrip audio came out as pure noise, suggesting the down/upsampler isn't trained well enough or has a fundamental bug. I should either fix the upsampler architecture or consider skipping it entirely and working with a speech codec that naturally produces lower-rate embeddings.


You're right, let me check what's actually happening and what speech codecs are available.


cd /home/ubuntu/lewm-tts && python3 -c "
import torch
# Check codec embedding ranges to understand if roundtrip is failing
codec_emb = torch.load('processed_data_codec/embeddings/000000.pt')
print(f'Codec emb shape: {codec_emb.shape}')
print(f'Range: [{codec_emb.min():.2f}, {codec_emb.max():.2f}]')
print(f'Mean: {codec_emb.mean():.4f}, Std: {codec_emb.std():.4f}')
print(f'Norm per frame: {codec_emb.norm(dim=0).mean():.2f}')
print(f'Roundtrip L1=0.77 is {0.77/codec_emb.abs().mean():.1f}x the mean abs value — terrible')
" 2>&1

Check codec embedding statistics


Codec emb shape: torch.Size([128, 714])
Range: [-21.58, 23.89]
Mean: -0.5224, Std: 4.8446
Norm per frame: 53.99
Roundtrip L1=0.77 is 0.2x the mean abs value — terrible
Shell cwd was reset to /home/ubuntu/pipecat


pip list 2>/dev/null | grep -iE "speech|snac|dac|codec|xcodec|token"

Check available speech codecs


asttokens                                3.0.1
audiocodecs                              0.0.1
curated-tokenizers                       0.0.9
dacvae                                   1.0.0
descript-audio-codec                     1.0.0
encodec                                  0.1.1
fish-speech                              0.1.0
google-cloud-speech                      2.37.0
google-cloud-texttospeech                2.34.0
linacodec                                0.0.11
orpheus-speech                           0.1.0
python_speech_features                   0.6
s3tokenizer                              0.3.0
snac                                     1.2.1
solospeech                               0.1.0
speechbrain                              1.0.3
SudachiDict-core                         20260116
SudachiPy                                0.6.10
tiktoken                                 0.12.0
tokenizers                               0.22.2
torchcodec                               0.8.0
xcodec2                                  0.1.5


cd /home/ubuntu/lewm-tts && python3 -c "
# Check SNAC - multi-scale neural audio codec
import torch
from snac import SNAC
model = SNAC.from_pretrained('hubertsiuzdak/snac_24khz').eval().cuda()
import torchaudio, json
from encodec.utils import convert_audio

manifest = json.load(open('processed_data_codec/manifest.json'))
wav, sr = torchaudio.load(manifest[0]['audio_path'])
wav = convert_audio(wav, sr, 24000, 1).cuda()

with torch.no_grad():
    codes = model.encode(wav.unsqueeze(0))
    print('SNAC codes:')
    for i, c in enumerate(codes):
        print(f'  Level {i}: {c.shape} ({c.shape[-1]} tokens)')

    # Reconstruct
    wav_recon = model.decode(codes)
    print(f'Recon shape: {wav_recon.shape}')
    T = min(wav.shape[-1], wav_recon.shape[-1])
    l1 = (wav[:,:T] - wav_recon[0,:,:T]).abs().mean().item()
    print(f'SNAC L1: {l1:.6f}')
" 2>&1 | grep -v Warning

Test SNAC codec


SNAC codes:
  Level 0: torch.Size([1, 112]) (112 tokens)
  Level 1: torch.Size([1, 224]) (224 tokens)
  Level 2: torch.Size([1, 448]) (448 tokens)
Recon shape: torch.Size([1, 1, 229376])
SNAC L1: 0.028372
Shell cwd was reset to /home/ubuntu/pipecat


cd /home/ubuntu/lewm-tts && python3 -c "
# Check XCodec2 - speech-specific codec
import torch
from xcodec2.modeling_xcodec2 import XCodec2Model
model = XCodec2Model.from_pretrained('HKUSTAudio/xcodec2')
model.eval().cuda()
import torchaudio, json
from encodec.utils import convert_audio

manifest = json.load(open('processed_data_codec/manifest.json'))
wav, sr = torchaudio.load(manifest[0]['audio_path'])
wav = convert_audio(wav, sr, 16000, 1).cuda()  # xcodec2 uses 16kHz

with torch.no_grad():
    codes = model.encode_code(wav.unsqueeze(0))
    print(f'XCodec2 codes: {codes.shape}')  # [B, n_q, T]

    # Check if we can get continuous embeddings
    # Get encoder output before VQ
    vq_emb = model.encode(wav.unsqueeze(0))
    print(f'XCodec2 continuous emb: {vq_emb.shape}')

    wav_recon = model.decode_code(codes)
    print(f'Recon: {wav_recon.shape}')
    T = min(wav.shape[-1], wav_recon.shape[-1])
    l1 = (wav[:,:T] - wav_recon[0,:,:T]).abs().mean().item()
    print(f'XCodec2 L1: {l1:.6f}')
" 2>&1 | grep -v -E "Warning|weight_norm"

Test XCodec2 speech codec


You are using a model of type xcodec2 to instantiate a model of type xcodec. This is not supported for all configurations of models and can yield errors.
Some weights of the model checkpoint at HKUSTAudio/xcodec2 were not used when initializing XCodec2Model: ['CodecEnc.conv_blocks.1.block.0.block.0.act.beta', 'CodecEnc.conv_blocks.1.block.0.block.2.act.beta', 'CodecEnc.conv_blocks.1.block.1.block.0.act.beta', 'CodecEnc.conv_blocks.1.block.1.block.2.act.beta', 'CodecEnc.conv_blocks.1.block.2.block.0.act.beta', 'CodecEnc.conv_blocks.1.block.2.block.2.act.beta', 'CodecEnc.conv_blocks.1.block.3.act.beta', 'CodecEnc.conv_blocks.2.block.0.block.0.act.beta', 'CodecEnc.conv_blocks.2.block.0.block.2.act.beta', 'CodecEnc.conv_blocks.2.block.1.block.0.act.beta', 'CodecEnc.conv_blocks.2.block.1.block.2.act.beta', 'CodecEnc.conv_blocks.2.block.2.block.0.act.beta', 'CodecEnc.conv_blocks.2.block.2.block.2.act.beta', 'CodecEnc.conv_blocks.2.block.3.act.beta', 'CodecEnc.conv_blocks.3.block.0.block.0.act.beta', 'CodecEnc.conv_blocks.3.block.0.block.2.act.beta', 'CodecEnc.conv_blocks.3.block.1.block.0.act.beta', 'CodecEnc.conv_blocks.3.block.1.block.2.act.beta', 'CodecEnc.conv_blocks.3.block.2.block.0.act.beta', 'CodecEnc.conv_blocks.3.block.2.block.2.act.beta', 'CodecEnc.conv_blocks.3.block.3.act.beta', 'CodecEnc.conv_blocks.4.block.0.block.0.act.beta', 'CodecEnc.conv_blocks.4.block.0.block.2.act.beta', 'CodecEnc.conv_blocks.4.block.1.block.0.act.beta', 'CodecEnc.conv_blocks.4.block.1.block.2.act.beta', 'CodecEnc.conv_blocks.4.block.2.block.0.act.beta', 'CodecEnc.conv_blocks.4.block.2.block.2.act.beta', 'CodecEnc.conv_blocks.4.block.3.act.beta', 'CodecEnc.conv_blocks.5.block.0.block.0.act.beta', 'CodecEnc.conv_blocks.5.block.0.block.2.act.beta', 'CodecEnc.conv_blocks.5.block.1.block.0.act.beta', 'CodecEnc.conv_blocks.5.block.1.block.2.act.beta', 'CodecEnc.conv_blocks.5.block.2.block.0.act.beta', 'CodecEnc.conv_blocks.5.block.2.block.2.act.beta', 'CodecEnc.conv_blocks.5.block.3.act.beta', 'CodecEnc.conv_final_block.0.act.beta']
- This IS expected if you are initializing XCodec2Model from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).
- This IS NOT expected if you are initializing XCodec2Model from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).
Some weights of XCodec2Model were not initialized from the model checkpoint at HKUSTAudio/xcodec2 and are newly initialized: ['CodecEnc.conv_blocks.1.block.0.block.0.act.bias', 'CodecEnc.conv_blocks.1.block.0.block.2.act.bias', 'CodecEnc.conv_blocks.1.block.1.block.0.act.bias', 'CodecEnc.conv_blocks.1.block.1.block.2.act.bias', 'CodecEnc.conv_blocks.1.block.2.block.0.act.bias', 'CodecEnc.conv_blocks.1.block.2.block.2.act.bias', 'CodecEnc.conv_blocks.1.block.3.act.bias', 'CodecEnc.conv_blocks.2.block.0.block.0.act.bias', 'CodecEnc.conv_blocks.2.block.0.block.2.act.bias', 'CodecEnc.conv_blocks.2.block.1.block.0.act.bias', 'CodecEnc.conv_blocks.2.block.1.block.2.act.bias', 'CodecEnc.conv_blocks.2.block.2.block.0.act.bias', 'CodecEnc.conv_blocks.2.block.2.block.2.act.bias', 'CodecEnc.conv_blocks.2.block.3.act.bias', 'CodecEnc.conv_blocks.3.block.0.block.0.act.bias', 'CodecEnc.conv_blocks.3.block.0.block.2.act.bias', 'CodecEnc.conv_blocks.3.block.1.block.0.act.bias', 'CodecEnc.conv_blocks.3.block.1.block.2.act.bias', 'CodecEnc.conv_blocks.3.block.2.block.0.act.bias', 'CodecEnc.conv_blocks.3.block.2.block.2.act.bias', 'CodecEnc.conv_blocks.3.block.3.act.bias', 'CodecEnc.conv_blocks.4.block.0.block.0.act.bias', 'CodecEnc.conv_blocks.4.block.0.block.2.act.bias', 'CodecEnc.conv_blocks.4.block.1.block.0.act.bias', 'CodecEnc.conv_blocks.4.block.1.block.2.act.bias', 'CodecEnc.conv_blocks.4.block.2.block.0.act.bias', 'CodecEnc.conv_blocks.4.block.2.block.2.act.bias', 'CodecEnc.conv_blocks.4.block.3.act.bias', 'CodecEnc.conv_blocks.5.block.0.block.0.act.bias', 'CodecEnc.conv_blocks.5.block.0.block.2.act.bias', 'CodecEnc.conv_blocks.5.block.1.block.0.act.bias', 'CodecEnc.conv_blocks.5.block.1.block.2.act.bias', 'CodecEnc.conv_blocks.5.block.2.block.0.act.bias', 'CodecEnc.conv_blocks.5.block.2.block.2.act.bias', 'CodecEnc.conv_blocks.5.block.3.act.bias', 'CodecEnc.conv_final_block.0.act.bias']
You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.
Traceback (most recent call last):
  File "<string>", line 6, in <module>
  File "/home/ubuntu/.local/lib/python3.10/site-packages/transformers/modeling_utils.py", line 4276, in cuda
    return super().cuda(*args, **kwargs)
  File "/home/ubuntu/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1084, in cuda
    return self._apply(lambda t: t.cuda(device))
  File "/home/ubuntu/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 930, in _apply
    module._apply(fn)
  File "/home/ubuntu/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 930, in _apply
    module._apply(fn)
  File "/home/ubuntu/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 930, in _apply
    module._apply(fn)
  [Previous line repeated 3 more times]
  File "/home/ubuntu/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 957, in _apply
    param_applied = fn(param)
  File "/home/ubuntu/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1084, in <lambda>
    return self._apply(lambda t: t.cuda(device))
torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 16.00 MiB. GPU 0 has a total capacity of 79.25 GiB of which 3.19 MiB is free. Process 3690516 has 76.75 GiB memory in use. Including non-PyTorch memory, this process has 2.48 GiB memory in use. Of the allocated memory 2.07 GiB is allocated by PyTorch, and 5.82 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation.  See documentation for Memory Management  (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables)
Shell cwd was reset to /home/ubuntu/pipecat


GPU is occupied by v5 training. Let me check SNAC more carefully - it seems promising. Let me think about the approach:

SNAC produces discrete tokens at 3 levels:
- Level 0: 112 tokens for 9.52s → 11.8 Hz (very low rate!)
- Level 1: 224 tokens → 23.5 Hz
- Level 2: 448 tokens → 47.1 Hz

This is multi-scale. The level 0 has just ~12 tokens/sec which is very manageable for AR prediction.

But SNAC uses discrete tokens, not continuous embeddings.

For JEPA I need continuous embeddings to do MSE prediction. Maybe I should reconsider - next-token prediction on codec tokens is what VALL-E does, but the user specifically wants JEPA. I could either use SNAC's codebook embeddings and predict those continuously at the low 12Hz rate, or work directly with EnCodec and project the embeddings to a lower dimension.

The real issue with the previous approach is that the down/upsampler is struggling with codec embeddings - they have high dynamic range and the transposed convolution creates artifacts. The cleanest solution is to skip the down/upsampler entirely and predict directly in codec space at 75Hz, which gives me 4 seconds for 300 steps. Alternatively, I could use SNAC at a lower level for better compression.

Now I'm checking whether SNAC provides continuous embeddings before quantization, since I need continuous representations rather than discrete tokens.


XCodec2 OOM because v5 training is using the GPU. Let me think about this differently.

The core problem with our v5: the down/upsampler is a lossy bottleneck we're training from scratch. SNAC looks much better — **12 tokens/sec at level 0** vs EnCodec's 75Hz. But both use discrete tokens.

Let me check SNAC's continuous embeddings (before quantization) since that's what JEPA needs:


cd /home/ubuntu/lewm-tts && python3 -c "
import torch
from snac import SNAC
model = SNAC.from_pretrained('hubertsiuzdak/snac_24khz').eval().cpu()

# Inspect SNAC architecture to find continuous embeddings
print('SNAC encoder:')
for name, mod in model.named_modules():
    if hasattr(mod, 'weight') and 'encoder' in name.lower():
        if hasattr(mod, 'out_channels'):
            print(f'  {name}: {mod}' if len(str(mod)) < 100 else f'  {name}: out={mod.out_channels}')

# Check the encoder output (before quantization)
import torchaudio, json
from encodec.utils import convert_audio
manifest = json.load(open('processed_data_codec/manifest.json'))
wav, sr = torchaudio.load(manifest[0]['audio_path'])
wav = convert_audio(wav, sr, 24000, 1)

with torch.no_grad():
    # SNAC internal: encoder → quantizer
    # Let's trace the encoder
    z = model.encoder(wav.unsqueeze(0))
    print(f'\\nSNAC encoder output: {z.shape}')  # continuous embedding before VQ

    # Check quantizer structure
    print(f'\\nQuantizer type: {type(model.quantizer)}')
    for name, param in model.quantizer.named_parameters():
        if 'codebook' in name or 'embed' in name:
            print(f'  {name}: {param.shape}')
" 2>&1 | grep -v -E "Warning|weight_norm"

Inspect SNAC continuous embeddings


SNAC encoder:
  encoder.block.0: out=48
  encoder.block.1.block.0.block.1: out=48
  encoder.block.1.block.0.block.3: out=48
  encoder.block.1.block.1.block.1: out=48
  encoder.block.1.block.1.block.3: out=48
  encoder.block.1.block.2.block.1: out=48
  encoder.block.1.block.2.block.3: out=48
  encoder.block.1.block.4: out=96
  encoder.block.2.block.0.block.1: out=96
  encoder.block.2.block.0.block.3: out=96
  encoder.block.2.block.1.block.1: out=96
  encoder.block.2.block.1.block.3: out=96
  encoder.block.2.block.2.block.1: out=96
  encoder.block.2.block.2.block.3: out=96
  encoder.block.2.block.4: out=192
  encoder.block.3.block.0.block.1: out=192
  encoder.block.3.block.0.block.3: out=192
  encoder.block.3.block.1.block.1: out=192
  encoder.block.3.block.1.block.3: out=192
  encoder.block.3.block.2.block.1: out=192
  encoder.block.3.block.2.block.3: out=192
  encoder.block.3.block.4: out=384
  encoder.block.4.block.0.block.1: out=384
  encoder.block.4.block.0.block.3: out=384
  encoder.block.4.block.1.block.1: out=384
  encoder.block.4.block.1.block.3: out=384
  encoder.block.4.block.2.block.1: out=384
  encoder.block.4.block.2.block.3: out=384
  encoder.block.4.block.4: out=768
  encoder.block.5: out=768

SNAC encoder output: torch.Size([1, 768, 446])

Quantizer type: <class 'snac.vq.ResidualVectorQuantize'>
  quantizers.0.codebook.weight: torch.Size([4096, 8])
  quantizers.1.codebook.weight: torch.Size([4096, 8])
  quantizers.2.codebook.weight: torch.Size([4096, 8])
Shell cwd was reset to /home/ubuntu/pipecat


Very interesting! SNAC:
- Encoder output: **768-dim at 446 frames** for 9.52s = **46.8 Hz**
- Codebook: 4096 entries × 8-dim (very small codebook embeddings)
- 3 quantizers at different temporal scales (112, 224, 448 tokens)

The 768-dim continuous embedding at 46.8Hz is interesting. But 768 dims is large and 46.8Hz is still quite fast.

The SNAC quantizer has multi-scale structure:
- Level 0: 112 tokens (12Hz) — coarse/semantic

- Level 1: 224 tokens (23.5Hz) — medium detail
- Level 2: 448 tokens (47Hz) — fine detail

For JEPA TTS, I could predict the coarse level first with minimal autoregressive steps, then use finer levels for refinement. But those 8-dim codebook embeddings feel too constrained for meaningful continuous prediction. I'm reconsidering the whole approach—maybe the user's frustration points to something more fundamental that needs rethinking. Let me look at what actually works in the literature instead.

Most successful TTS systems—VALL-E, SoundStorm, VoiceCraft, XTTS, Tortoise, F5-TTS, E2 TTS—all predict discrete tokens or work with mel spectrograms, not continuous embeddings. The JEPA approach with continuous prediction is far less established for speech synthesis. But since the user specifically wants JEPA, I need to figure out what could actually make it viable. The real question is whether JEPA's strength in vision—where prediction spaces are smooth and high-dimensional—translates to audio at all.

The issue with audio codecs is that their continuous embeddings before quantization aren't necessarily smooth; they're optimized for quantization, not for smooth variation. A better path might be using a speech encoder like HuBERT or WavLM that produces genuinely smooth, structured semantic embeddings designed for representation learning, then predicting those with JEPA and decoding with a vocoder. That's essentially what SpeechTokenizer does—separating semantic and acoustic information. But this is getting complicated. Let me focus on the simplest approach that could actually work.

Instead of predicting continuous embeddings, I could use SNAC and predict discrete tokens at level 0 (12Hz, so only ~120 tokens for 10 seconds), then predict higher levels non-autoregressively from there. This is what successful TTS models do, and it's efficient. But the user wants JEPA, which is fundamentally about continuous embeddings, not discrete tokens. I could try a hybrid approach where I use the SNAC codebook embeddings themselves as continuous targets—each level 0 token maps to an 8-dimensional codebook vector that I predict continuously.

Actually, I'm realizing the real issue: the SNAC encoder outputs 768-dimensional features at 47Hz before the multi-scale quantization happens. If I predict those continuous features directly and then feed them through SNAC's quantizer and decoder, I'd be introducing quantization errors since the decoder expects already-quantized inputs. The upsampler approach was flawed because it's too hard to train. Let me reconsider what the user actually asked about—why not just use nanocodec or speechtokenizer instead of building this from scratch?

—and then propose a simpler fix: drop the down/upsampler entirely and work directly in codec space with just linear projections at 75Hz, which trades more AR steps for eliminating the lossy compression bottleneck.


OK, here's the issue. Two problems:

**Problem 1**: The down/upsampler (stride-3 CNN) is too hard to train — it's a lossy bottleneck causing noise. EnCodec embeddings have range [-22, +24] and our upsampler can't reconstruct them accurately at epoch 39.

**Problem 2**: You're right about codec choice. EnCodec is a general audio codec, not speech-specific. SNAC gives us:
- **Level 0: 12 Hz** (112 tokens for 10s!) — perfect for AR
- **Level 1: 23 Hz** — medium detail
- **Level 2: 47 Hz** — fine detail

But the real question is: JEPA predicts **continuous** embeddings, while codecs use **discrete** tokens. The most successful TTS models (VALL-E, etc.) predict discrete tokens, not continuous.

Let me take the simplest approach that fixes the noise: **remove the down/upsampler entirely**, predict directly in EnCodec space with just linear projections (128→256→128). No compression, no information loss.


kill 3690516 && sleep 5 && nvidia-smi | grep MiB

Kill v5 training to free GPU


| N/A   40C    P0             61W /  400W |   78598MiB /  81920MiB |      3%      Default |


ps aux | grep train_v5 | grep -v grep | awk '{print $2}' | xargs -r kill -9; sleep 5; nvidia-smi | grep MiB

Force kill training processes


| N/A   35C    P0             52W /  400W |       1MiB /  81920MiB |      0%      Default |


Good. Now let me build a much simpler v5b — no down/upsampler, just linear projections:


   280→        # EMA target on downsampler output
   281→        self.ema_downsampler = copy.deepcopy(self.downsampler)
   282→        for p in self.ema_downsampler.parameters():
   283→            p.requires_grad = False
   284→        self.ema_decay = config.get("ema_decay", 0.998)
   285→
   286→        # Learnable start embedding
   287→        self.start_emb = nn.Parameter(torch.randn(1, 1, d_model) * 0.02)
   288→
   289→        # Loss weights
   290→        self.pred_weight = config.get("pred_weight", 10.0)
   291→        self.roundtrip_weight = config.get("roundtrip_weight", 1.0)
   292→
   293→        # Input noise for train/inference gap
   294→        self.input_noise = config.get("input_noise", 0.0)
   295→
   296→        self.config = config
   297→
   298→    @torch.no_grad()
   299→    def update_ema(self):
   300→        for p_online, p_ema in zip(self.downsampler.parameters(),
   301→                                   self.ema_downsampler.parameters()):
   302→            p_ema.data.mul_(self.ema_decay).add_(p_online.data, alpha=1 - self.ema_decay)
   303→
   304→    def forward(self, codec_emb, text_tokens, codec_mask=None, text_mask=None):
   305→        """
   306→        Args:
   307→            codec_emb: [B, 128, T_75Hz] — continuous EnCodec embeddings
   308→            text_tokens: [B, T_text]
   309→            codec_mask: [B, T_75Hz] bool, True = padding
   310→            text_mask: [B, T_text] bool
   311→        """
   312→        B = codec_emb.shape[0]
   313→
   314→        # Text encoding
   315→        text_emb = self.text_encoder(text_tokens, text_mask)
   316→
   317→        # Online downsampler: 75Hz×128d → 25Hz×256d
   318→        z, ds_mask = self.downsampler(codec_emb, codec_mask)  # [B, T_down, d_model]
   319→        T_down = z.shape[1]
   320→
   321→        # EMA target downsampler (deterministic)
   322→        with torch.no_grad():
   323→            self.ema_downsampler.eval()
   324→            target_emb, _ = self.ema_downsampler(codec_emb, codec_mask)
   325→            self.ema_downsampler.train()
   326→
   327→        # Prepend start embedding
   328→        start = self.start_emb.expand(B, -1, -1)
   329→        input_emb = torch.cat([start, z[:, :-1]], dim=1)
   330→
   331→        # Input noise on real embeddings (not start_emb)
   332→        if self.training and self.input_noise > 0:
   333→            noise = torch.zeros_like(input_emb)
   334→            noise[:, 1:] = torch.randn(B, T_down - 1, input_emb.shape[-1],
   335→                                        device=input_emb.device) * self.input_noise
   336→            input_emb = input_emb + noise
   337→
   338→        # Masks
   339→        if ds_mask is not None:
   340→            start_mask = torch.zeros(B, 1, dtype=torch.bool, device=z.device)
   341→            pred_mask = torch.cat([start_mask, ds_mask[:, :-1]], dim=1)
   342→            loss_mask = ds_mask
   343→        else:
   344→            pred_mask = None
   345→            loss_mask = None
   346→
   347→        # JEPA prediction
   348→        predicted = self.predictor(input_emb, text_emb, pred_mask, text_mask)
   349→
   350→        # ─── Loss 1: Prediction (MSE vs EMA targets) ───
   351→        if loss_mask is not None:
   352→            valid = (~loss_mask).unsqueeze(-1)
   353→            prediction_loss = F.mse_loss(
   354→                predicted * valid, target_emb * valid, reduction="sum"
   355→            ) / (valid.sum() * predicted.shape[-1] + 1e-8)
   356→        else:
   357→            prediction_loss = F.mse_loss(predicted, target_emb)
   358→
   359→        # ─── Loss 2: Roundtrip (downsample → upsample → match original codec_emb) ───
   360→        T_codec = codec_emb.shape[2]
   361→        codec_recon = self.upsampler(z, target_len=T_codec)  # [B, 128, T_75Hz]
   362→
   363→        if codec_mask is not None:
   364→            valid_codec = (~codec_mask).unsqueeze(1).float()
   365→            roundtrip_loss = F.l1_loss(
   366→                codec_recon * valid_codec, codec_emb * valid_codec, reduction="sum"
   367→            ) / (valid_codec.sum() * codec_emb.shape[1] + 1e-8)
   368→        else:
   369→            roundtrip_loss = F.l1_loss(codec_recon, codec_emb)
   370→
   371→        total_loss = self.pred_weight * prediction_loss + self.roundtrip_weight * roundtrip_loss
   372→
   373→        return {
   374→            "total_loss": total_loss,
   375→            "prediction_loss": prediction_loss,
   376→            "roundtrip_loss": roundtrip_loss,
   377→        }
   378→
   379→    def init_ar_cache(self):
   380→        num_layers = len(self.predictor.transformer.layers)
   381→        return self.predictor.init_cache(num_layers)
   382→
   383→    def predict_next_cached(self, new_emb, text_emb, step_idx, cache, text_mask=None):
   384→        return self.predictor.inference_step(new_emb, text_emb, step_idx, cache, text_mask=text_mask)
   385→
   386→
   387→def count_parameters(model):
   388→    return sum(p.numel() for p in model.parameters() if p.requires_grad)
   389→
   390→
   391→def build_model(config=None):
   392→    if config is None:
   393→        config = {
   394→            "d_model": 256, "nhead": 4, "codec_dim": 128,
   395→            "text_vocab_size": 256, "text_encoder_layers": 4,
   396→            "predictor_layers": 6, "dropout": 0.1, "ds_stride": 3,
   397→            "pred_weight": 10.0, "roundtrip_weight": 1.0,
   398→            "ema_decay": 0.998, "input_noise": 0.0,
   399→        }
   400→    model = LeWMTTSv5(config)
   401→    print(f"LeWM TTS v5: {count_parameters(model)/1e6:.2f}M trainable parameters")
   402→    return model, config
   403→

        # EMA target on downsampler output
        self.ema_downsampler = copy.deepcopy(self.downsampler)
        for p in self.ema_downsampler.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.downsampler.parameters(),
                                   self.ema_downsampler.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_75Hz] — continuous EnCodec embeddings
            text_tokens: [B, T_text]
            codec_mask: [B, T_75Hz] bool, True = padding
            text_mask: [B, T_text] bool
        """
        B = codec_emb.shape[0]

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

        # Online downsampler: 75Hz×128d → 25Hz×256d
        z, ds_mask = self.downsampler(codec_emb, codec_mask)  # [B, T_down, d_model]
        T_down = z.shape[1]

        # EMA target downsampler (deterministic)
        with torch.no_grad():
            self.ema_downsampler.eval()
            target_emb, _ = self.ema_downsampler(codec_emb, codec_mask)
            self.ema_downsampler.train()

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

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

        # Masks
        if ds_mask is not None:
            start_mask = torch.zeros(B, 1, dtype=torch.bool, device=z.device)
            pred_mask = torch.cat([start_mask, ds_mask[:, :-1]], dim=1)
            loss_mask = ds_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 (downsample → upsample → match original codec_emb) ───
        T_codec = codec_emb.shape[2]
        codec_recon = self.upsampler(z, target_len=T_codec)  # [B, 128, T_75Hz]

        if codec_mask is not None:
            valid_codec = (~codec_mask).unsqueeze(1).float()
            roundtrip_loss = F.l1_loss(
                codec_recon * valid_codec, codec_emb * valid_codec, reduction="sum"
            ) / (valid_codec.sum() * codec_emb.shape[1] + 1e-8)
        else:
            roundtrip_loss = F.l1_loss(codec_recon, codec_emb)

        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, "ds_stride": 3,
            "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


Now let me rewrite model_v5.py with the simplified approach — no down/upsampler, just linear projections, working at EnCodec's native 75Hz:


Edit

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

Architecture:
  - EnCodec encoder/decoder (frozen): audio ↔ 128d embeddings @ 75Hz
  - CodecDownsampler: 75Hz → 25Hz, 128d → 256d (trainable)
  - CodecUpsampler: 25Hz → 75Hz, 256d → 128d (trainable)
  - TextEncoder: byte-level transformer for Hindi text
  - JEPAPredictor: causal transformer, predicts next embedding
  - EMA target on downsampled embeddings
"""

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

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

# ─── Codec Down/Upsampler ──────────────────────────────────────────────────

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

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

class CodecDownsampler(nn.Module):
    """Downsample codec embeddings: 75Hz×128d → 25Hz×d_model."""
    def __init__(self, codec_dim=128, d_model=256, stride=3):
        super().__init__()
        self.stride = stride
        self.net = nn.Sequential(
            nn.Conv1d(codec_dim, d_model, kernel_size=7, padding=3),
            nn.GELU(),
            ResConvBlock(d_model, kernel_size=5),
            nn.Conv1d(d_model, d_model, kernel_size=2 * stride + 1, stride=stride, padding=stride),
            nn.GELU(),
            ResConvBlock(d_model, kernel_size=3),
        )
        # Transformer layers for contextual understanding
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=d_model, nhead=4, dim_feedforward=d_model * 4,
            dropout=0.1, batch_first=True, activation="gelu",
        )
        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=2)
        self.pos_embed = SinusoidalPositionalEncoding(d_model, max_len=4096)

    def forward(self, codec_emb, mask=None):
        """codec_emb: [B, 128, T_75Hz] → [B, T_25Hz, d_model]"""
        x = self.net(codec_emb)  # [B, d_model, T_25Hz]
        x = x.transpose(1, 2)   # [B, T_25Hz, d_model]
        if mask is not None:
            mask = mask[:, ::self.stride][:, :x.shape[1]]
        x = self.pos_embed(x)
        x = self.transformer(x, src_key_padding_mask=mask)
        return x, mask

class CodecUpsampler(nn.Module):
    """Upsample predicted embeddings: 25Hz×d_model → 75Hz×128d."""
    def __init__(self, d_model=256, codec_dim=128, stride=3):
        super().__init__()
        self.net = nn.Sequential(
            nn.ConvTranspose1d(d_model, d_model, kernel_size=2 * stride, stride=stride, padding=stride // 2),
            nn.GELU(),
            ResConvBlock(d_model, kernel_size=5),
            ResConvBlock(d_model, kernel_size=3),
            nn.Conv1d(d_model, codec_dim, kernel_size=7, padding=3),
        )

    def forward(self, z, target_len=None):
        """z: [B, T_25Hz, d_model] → [B, 128, T_75Hz]"""
        x = z.transpose(1, 2)  # [B, d_model, T_25Hz]
        x = self.net(x)        # [B, 128, T_75Hz_approx]
        if target_len is not None:
            if x.shape[2] > target_len:
                x = x[:, :, :target_len]
            elif x.shape[2] < target_len:
                x = F.pad(x, (0, target_len - x.shape[2]))
        return 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


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

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

Architecture:
  - EnCodec encoder/decoder (frozen): audio ↔ 128d embeddings @ 75Hz
  - CodecDownsampler: 75Hz → 25Hz, 128d → 256d (trainable)
  - CodecUpsampler: 25Hz → 75Hz, 256d → 128d (trainable)
  - TextEncoder: byte-level transformer for Hindi text
  - JEPAPredictor: causal transformer, predicts next embedding
  - EMA target on downsampled embeddings
"""

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

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

# ─── Codec Down/Upsampler ──────────────────────────────────────────────────

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

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

class CodecDownsampler(nn.Module):
    """Downsample codec embeddings: 75Hz×128d → 25Hz×d_model."""
    def __init__(self, codec_dim=128, d_model=256, stride=3):
        super().__init__()
        self.stride = stride
        self.net = nn.Sequential(
            nn.Conv1d(codec_dim, d_model, kernel_size=7, padding=3),
            nn.GELU(),
            ResConvBlock(d_model, kernel_size=5),
            nn.Conv1d(d_model, d_model, kernel_size=2 * stride + 1, stride=stride, padding=stride),
            nn.GELU(),
            ResConvBlock(d_model, kernel_size=3),
        )
        # Transformer layers for contextual understanding
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=d_model, nhead=4, dim_feedforward=d_model * 4,
            dropout=0.1, batch_first=True, activation="gelu",
        )
        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=2)
        self.pos_embed = SinusoidalPositionalEncoding(d_model, max_len=4096)

    def forward(self, codec_emb, mask=None):
        """codec_emb: [B, 128, T_75Hz] → [B, T_25Hz, d_model]"""
        x = self.net(codec_emb)  # [B, d_model, T_25Hz]
        x = x.transpose(1, 2)   # [B, T_25Hz, d_model]
        if mask is not None:
            mask = mask[:, ::self.stride][:, :x.shape[1]]
        x = self.pos_embed(x)
        x = self.transformer(x, src_key_padding_mask=mask)
        return x, mask

class CodecUpsampler(nn.Module):
    """Upsample predicted embeddings: 25Hz×d_model → 75Hz×128d."""
    def __init__(self, d_model=256, codec_dim=128, stride=3):
        super().__init__()
        self.net = nn.Sequential(
            nn.ConvTranspose1d(d_model, d_model, kernel_size=2 * stride, stride=stride, padding=stride // 2),
            nn.GELU(),
            ResConvBlock(d_model, kernel_size=5),
            ResConvBlock(d_model, kernel_size=3),
            nn.Conv1d(d_model, codec_dim, kernel_size=7, padding=3),
        )

    def forward(self, z, target_len=None):
        """z: [B, T_25Hz, d_model] → [B, 128, T_75Hz]"""
        x = z.transpose(1, 2)  # [B, d_model, T_25Hz]
        x = self.net(x)        # [B, 128, T_75Hz_approx]
        if target_len is not None:
            if x.shape[2] > target_len:
                x = x[:, :, :target_len]
            elif x.shape[2] < target_len:
                x = F.pad(x, (0, target_len - x.shape[2]))
        return x

# ─── 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]]

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

class LeWMTTSv5(nn.Module):
    """
    JEPA TTS with EnCodec backend.

    Training flow:
        codec_emb → downsample → JEPA predict → loss vs EMA targets
        codec_emb → downsample → upsample → loss vs codec_emb (roundtrip)

    Inference flow:
        text → predictor (AR) → upsample → 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)
        ds_stride = config.get("ds_stride", 3)

        self.text_encoder = TextEncoder(
            vocab_size=text_vocab, d_model=d_model, nhead=nhead,
            num_layers=text_layers, dropout=dropout,
        )
        self.downsampler = CodecDownsampler(
            codec_dim=codec_dim, d_model=d_model, stride=ds_stride,
        )
        self.upsampler = CodecUpsampler(
            d_model=d_model, codec_dim=codec_dim, stride=ds_stride,
        )
        self.predictor = JEPAPredictor(
            d_model=d_model, nhead=nhead,
            num_layers=predictor_layers, dropout=dropout,
        )

        # EMA target on downsampler output
        self.ema_downsampler = copy.deepcopy(self.downsampler)
        for p in self.ema_downsampler.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.downsampler.parameters(),
                                   self.ema_downsampler.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_75Hz] — continuous EnCodec embeddings
            text_tokens: [B, T_text]
            codec_mask: [B, T_75Hz] bool, True = padding
            text_mask: [B, T_text] bool
        """
        B = codec_emb.shape[0]

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

        # Online downsampler: 75Hz×128d → 25Hz×256d
        z, ds_mask = self.downsampler(codec_emb, codec_mask)  # [B, T_down, d_model]
        T_down = z.shape[1]

        # EMA target downsampler (deterministic)
        with torch.no_grad():
            self.ema_downsampler.eval()
            target_emb, _ = self.ema_downsampler(codec_emb, codec_mask)
            self.ema_downsampler.train()

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

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

        # Masks
        if ds_mask is not None:
            start_mask = torch.zeros(B, 1, dtype=torch.bool, device=z.device)
            pred_mask = torch.cat([start_mask, ds_mask[:, :-1]], dim=1)
            loss_mask = ds_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 (downsample → upsample → match original codec_emb) ───
        T_codec = codec_emb.shape[2]
        codec_recon = self.upsampler(z, target_len=T_codec)  # [B, 128, T_75Hz]

        if codec_mask is not None:
            valid_codec = (~codec_mask).unsqueeze(1).float()
            roundtrip_loss = F.l1_loss(
                codec_recon * valid_codec, codec_emb * valid_codec, reduction="sum"
            ) / (valid_codec.sum() * codec_emb.shape[1] + 1e-8)
        else:
            roundtrip_loss = F.l1_loss(codec_recon, codec_emb)

        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, "ds_stride": 3,
            "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

-  - CodecDownsampler: 75Hz → 25Hz, 128d → 256d (trainable)

-  - CodecUpsampler: 25Hz → 75Hz, 256d → 128d (trainable)

+  - Simple linear projections: 128d ↔ d_model (no lossy compression)

-  - EMA target on downsampled embeddings

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

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

-        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.transformer(x, src_key_padding_mask=text_mask)

-# ─── Codec Down/Upsampler ──────────────────────────────────────────────────

-    def __init__(self, channels, kernel_size=3):

-            nn.Conv1d(channels, channels, kernel_size, padding=kernel_size // 2),

-    """Downsample codec embeddings: 75Hz×128d → 25Hz×d_model."""

-    def __init__(self, codec_dim=128, d_model=256, stride=3):

-            nn.Conv1d(codec_dim, d_model, kernel_size=7, padding=3),

-            ResConvBlock(d_model, kernel_size=5),

-            nn.Conv1d(d_model, d_model, kernel_size=2 * stride + 1, stride=stride, padding=stride),

-            ResConvBlock(d_model, kernel_size=3),

-        # Transformer layers for contextual understanding

-            d_model=d_model, nhead=4, dim_feedforward=d_model * 4,

-            dropout=0.1, batch_first=True, activation="gelu",

-        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=2)

-        self.pos_embed = SinusoidalPositionalEncoding(d_model, max_len=4096)

-    def forward(self, codec_emb, mask=None):

-        x = self.net(codec_emb)  # [B, d_model, T_25Hz]

-        x = x.transpose(1, 2)   # [B, T_25Hz, d_model]

-            mask = mask[:, ::self.stride][:, :x.shape[1]]

-        x = self.transformer(x, src_key_padding_mask=mask)

-    """Upsample predicted embeddings: 25Hz×d_model → 75Hz×128d."""

-    def __init__(self, d_model=256, codec_dim=128, stride=3):

-            nn.ConvTranspose1d(d_model, d_model, kernel_size=2 * stride, stride=stride, padding=stride // 2),

-            ResConvBlock(d_model, kernel_size=5),

-            nn.Conv1d(d_model, codec_dim, kernel_size=7, padding=3),

-    def forward(self, z, target_len=None):

-        x = z.transpose(1, 2)  # [B, d_model, T_25Hz]

-        x = self.net(x)        # [B, 128, T_75Hz_approx]

-            if x.shape[2] > target_len:

-                x = x[:, :, :target_len]

-            elif x.shape[2] < target_len:

-                x = F.pad(x, (0, target_len - x.shape[2]))


Now replace the model class:


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

class LeWMTTSv5(nn.Module):
    """
    JEPA TTS with EnCodec backend.

    Training flow:
        codec_emb → downsample → JEPA predict → loss vs EMA targets
        codec_emb → downsample → upsample → loss vs codec_emb (roundtrip)

    Inference flow:
        text → predictor (AR) → upsample → 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)
        ds_stride = config.get("ds_stride", 3)

        self.text_encoder = TextEncoder(
            vocab_size=text_vocab, d_model=d_model, nhead=nhead,
            num_layers=text_layers, dropout=dropout,
        )
        self.downsampler = CodecDownsampler(
            codec_dim=codec_dim, d_model=d_model, stride=ds_stride,
        )
        self.upsampler = CodecUpsampler(
            d_model=d_model, codec_dim=codec_dim, stride=ds_stride,
        )
        self.predictor = JEPAPredictor(
            d_model=d_model, nhead=nhead,
            num_layers=predictor_layers, dropout=dropout,
        )

        # EMA target on downsampler output
        self.ema_downsampler = copy.deepcopy(self.downsampler)
        for p in self.ema_downsampler.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.downsampler.parameters(),
                                   self.ema_downsampler.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_75Hz] — continuous EnCodec embeddings
            text_tokens: [B, T_text]
            codec_mask: [B, T_75Hz] bool, True = padding
            text_mask: [B, T_text] bool
        """
        B = codec_emb.shape[0]

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

        # Online downsampler: 75Hz×128d → 25Hz×256d
        z, ds_mask = self.downsampler(codec_emb, codec_mask)  # [B, T_down, d_model]
        T_down = z.shape[1]

        # EMA target downsampler (deterministic)
        with torch.no_grad():
            self.ema_downsampler.eval()
            target_emb, _ = self.ema_downsampler(codec_emb, codec_mask)
            self.ema_downsampler.train()

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

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

        # Masks
        if ds_mask is not None:
            start_mask = torch.zeros(B, 1, dtype=torch.bool, device=z.device)
            pred_mask = torch.cat([start_mask, ds_mask[:, :-1]], dim=1)
            loss_mask = ds_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 (downsample → upsample → match original codec_emb) ───
        T_codec = codec_emb.shape[2]
        codec_recon = self.upsampler(z, target_len=T_codec)  # [B, 128, T_75Hz]

        if codec_mask is not None:
            valid_codec = (~codec_mask).unsqueeze(1).float()
            roundtrip_loss = F.l1_loss(
                codec_recon * valid_codec, codec_emb * valid_codec, reduction="sum"
            ) / (valid_codec.sum() * codec_emb.shape[1] + 1e-8)
        else:
            roundtrip_loss = F.l1_loss(codec_recon, codec_emb)

        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)

# ─── 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() * 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)