/home/ubuntu/lewm-tts/processed_data/mels/000105.pt
/home/ubuntu/lewm-tts/processed_data/mels/000086.pt
/home/ubuntu/lewm-tts/processed_data/mels/000027.pt
/home/ubuntu/lewm-tts/processed_data/mels/000021.pt
/home/ubuntu/lewm-tts/processed_data/mels/000060.pt
/home/ubuntu/lewm-tts/processed_data/mels/000069.pt
/home/ubuntu/lewm-tts/processed_data/mels/000031.pt
/home/ubuntu/lewm-tts/processed_data/mels/000100.pt
/home/ubuntu/lewm-tts/processed_data/mels/000011.pt
/home/ubuntu/lewm-tts/processed_data/mels/000079.pt
---
/home/ubuntu/lewm-tts/processed_data/mels/000105.pt
/home/ubuntu/lewm-tts/processed_data/mels/000086.pt
/home/ubuntu/lewm-tts/processed_data/mels/000027.pt
/home/ubuntu/lewm-tts/processed_data/mels/000021.pt
/home/ubuntu/lewm-tts/processed_data/mels/000060.pt
/home/ubuntu/lewm-tts/processed_data/mels/000069.pt
/home/ubuntu/lewm-tts/processed_data/mels/000031.pt
/home/ubuntu/lewm-tts/processed_data/mels/000100.pt
/home/ubuntu/lewm-tts/processed_data/mels/000011.pt
/home/ubuntu/lewm-tts/processed_data/mels/000079.pt
---
/home/ubuntu/transcripts/data/final_export_reference/manifest.json
/home/ubuntu/lewm-tts/processed_data_codec/manifest.json
/home/ubuntu/lewm-tts/v2/processed_data/manifest.json
/home/ubuntu/lewm-tts/processed_data_100mel/manifest.json
/home/ubuntu/lewm-tts/processed_data/manifest.json
/home/ubuntu/training/test_audio/te_shard/manifest.json


Bash

python3 -c "
import json
for path in [
    '/home/ubuntu/lewm-tts/processed_data_100mel/manifest.json',
    '/home/ubuntu/lewm-tts/v2/processed_data/manifest.json',
    '/home/ubuntu/lewm-tts/processed_data_codec/manifest.json',
]:
    with open(path) as f:
        d = json.load(f)
    print(f'{path}:')
    print(f'  Samples: {len(d)}')
    print(f'  Keys: {list(d[0].keys())}')
    print(f'  First mel_path: {d[0].get(\"mel_path\", \"N/A\")}')
    print()
" 2>&1

Check other manifests


/home/ubuntu/lewm-tts/processed_data_100mel/manifest.json:
  Samples: 7009
  Keys: ['id', 'audio_path', 'mel_path', 'text', 'duration', 'mel_frames']
  First mel_path: /home/ubuntu/lewm-tts/processed_data_100mel/mels/000000.pt

/home/ubuntu/lewm-tts/v2/processed_data/manifest.json:
  Samples: 7009
  Keys: ['id', 'latent_path', 'latent_frames', 'text', 'duration']
  First mel_path: N/A

/home/ubuntu/lewm-tts/processed_data_codec/manifest.json:
  Samples: 7009
  Keys: ['emb_path', 'text', 'audio_path', 'emb_frames', 'duration']
  First mel_path: N/A


**Found it.** `processed_data_100mel` has 7009 samples — that's your 20h dataset with mel files.


ls /home/ubuntu/lewm-tts/processed_data_100mel/mels/ | wc -l && python3 -c "
import torch, json
mel = torch.load('/home/ubuntu/lewm-tts/processed_data_100mel/mels/000000.pt', weights_only=True)
print(f'Mel shape: {mel.shape}, range: [{mel.min():.2f}, {mel.max():.2f}]')

with open('/home/ubuntu/lewm-tts/processed_data_100mel/manifest.json') as f:
    d = json.load(f)
total_dur = sum(e['duration'] for e in d)
print(f'Total: {len(d)} samples, {total_dur/3600:.2f}h')
print(f'Mel frames range: {min(e[\"mel_frames\"] for e in d)} - {max(e[\"mel_frames\"] for e in d)}')
" 2>&1

Verify 100mel dataset


Error: Exit code 1
ls: cannot access '/home/ubuntu/lewm-tts/processed_data_100mel/mels/': No such file or directory
0
Traceback (most recent call last):
  File "<string>", line 3, in <module>
  File "/home/ubuntu/.local/lib/python3.10/site-packages/torch/serialization.py", line 1484, in load
    with _open_file_like(f, "rb") as opened_file:
  File "/home/ubuntu/.local/lib/python3.10/site-packages/torch/serialization.py", line 759, in _open_file_like
    return _open_file(name_or_buffer, mode)
  File "/home/ubuntu/.local/lib/python3.10/site-packages/torch/serialization.py", line 740, in __init__
    super().__init__(open(name, mode))
FileNotFoundError: [Errno 2] No such file or directory: '/home/ubuntu/lewm-tts/processed_data_100mel/mels/000000.pt'


Check what's in 100mel dir


audio
config.json
manifest.json


ls /home/ubuntu/lewm-tts/processed_data_100mel/audio/ | wc -l && ls /home/ubuntu/lewm-tts/processed_data_100mel/audio/ | head -5

Check audio files


7009
000000.wav
000001.wav
000002.wav
000003.wav
000004.wav


Audio files exist (7009) but mels were never generated. Manifest references mel paths that don't exist. We need to extract mels from the audio. This is fast on A100:


python3 -c "
import torch, torchaudio, json, os
from pathlib import Path

# Mel config matching Vocos mel-24khz
mel_transform = torchaudio.transforms.MelSpectrogram(
    sample_rate=24000, n_fft=1024, hop_length=256, n_mels=100, power=1.0, f_min=0,
)

data_dir = Path('/home/ubuntu/lewm-tts/processed_data_100mel')
mel_dir = data_dir / 'mels'
mel_dir.mkdir(exist_ok=True)

with open(data_dir / 'manifest.json') as f:
    manifest = json.load(f)

done = 0
errors = 0
for entry in manifest:
    audio_path = entry['audio_path']
    mel_path = entry['mel_path']
    try:
        wav, sr = torchaudio.load(audio_path)
        if sr != 24000:
            wav = torchaudio.functional.resample(wav, sr, 24000)
        mel = mel_transform(wav)
        log_mel = torch.log(torch.clamp(mel, min=1e-7))
        torch.save(log_mel.squeeze(0), mel_path)  # [100, T]
        done += 1
    except Exception as e:
        errors += 1
        if errors <= 3:
            print(f'Error on {audio_path}: {e}')
    if done % 1000 == 0:
        print(f'  Extracted {done}/{len(manifest)} mels...')

print(f'Done: {done} mels extracted, {errors} errors')
" 2>&1

Extract mel spectrograms from audio


<task-notification>
<task-id>bcqly5525</task-id>
<tool-use-id>toolu_01BMWBW933TR2qKf9sgmTcHG</tool-use-id>
<output-file>/tmp/claude-1001/-home-ubuntu/e792b4a6-7519-4b8f-9922-071bec1295ca/tasks/bcqly5525.output</output-file>
<status>completed</status>
<summary>Background command "Full training run (100 epochs)" completed (exit code 0)</summary>
</task-notification>


what is happening


  Extracted 1000/7009 mels...
  Extracted 2000/7009 mels...
  Extracted 3000/7009 mels...
  Extracted 4000/7009 mels...
  Extracted 5000/7009 mels...
  Extracted 6000/7009 mels...
  Extracted 7000/7009 mels...
Done: 7009 mels extracted, 0 errors


7009 mels extracted, zero errors. Now let me kill the small training run and start the real one on the full 20h dataset.


pkill -f "train.py.*output_v4" 2>/dev/null; echo "killed"

Stop small training run


Exit code 144


python3 -c "
import torch, json
mel = torch.load('/home/ubuntu/lewm-tts/processed_data_100mel/mels/000000.pt', weights_only=True)
print(f'Mel: {mel.shape}, [{mel.min():.2f}, {mel.max():.2f}]')

with open('/home/ubuntu/lewm-tts/processed_data_100mel/manifest.json') as f:
    d = json.load(f)
total = sum(e['duration'] for e in d)
print(f'{len(d)} samples, {total/3600:.2f}h')
" 2>&1

Verify full dataset ready


Mel: torch.Size([100, 893]), [-8.69, 4.46]
7009 samples, 17.88h


17.88 hours, 7009 samples, mels ready. Starting the real training:


python3 train.py \
  --data_dir /home/ubuntu/lewm-tts/processed_data_100mel \
  --output_dir ./output_v4 \
  --d_model 256 \
  --nhead 4 \
  --text_layers 4 \
  --audio_layers 4 \
  --predictor_layers 6 \
  --downsample_factor 4 \
  --kl_weight 0.05 \
  --recon_weight 1.0 \
  --spectral_weight 0.5 \
  --pred_weight 10.0 \
  --cosine_weight 1.0 \
  --ema_decay 0.998 \
  --free_bits 2.0 \
  --input_noise_max 0.2 \
  --n_speakers 1 \
  --epochs 200 \
  --batch_size 16 \
  --lr 1e-4 \
  --log_every 10 \
  --print_every 50 \
  --save_every 25 \
  2>&1 | head -30

Start real training on 18h dataset


Device: cuda
LeWM TTS model: 24.07M parameters
Dataset: 6959 samples loaded
Training: 6959 samples, 434 batches/epoch
  [1/200] step 50 | loss=17.9194 pred=1.2876 kl=0.0004 recon=2.3804 spec=5.3255 spk=0.0000 | grad=30.77 | 5.3s
  [1/200] step 100 | loss=15.1931 pred=1.0352 kl=0.0004 recon=2.2976 spec=5.0873 spk=0.0000 | grad=21.81 | 9.7s
  [1/200] step 150 | loss=13.3163 pred=0.8435 kl=0.0006 recon=2.3041 spec=5.1554 spk=0.0000 | grad=17.46 | 14.0s
  [1/200] step 200 | loss=11.6255 pred=0.6790 kl=0.0011 recon=2.3504 spec=4.9699 spk=0.0000 | grad=14.35 | 18.4s
  [1/200] step 250 | loss=10.1685 pred=0.5633 kl=0.0047 recon=2.2794 spec=4.5117 spk=0.0000 | grad=10.66 | 22.6s
  [1/200] step 300 | loss=9.1086 pred=0.4591 kl=0.0203 recon=2.4211 spec=4.1929 spk=0.0000 | grad=8.17 | 26.8s
  [1/200] step 350 | loss=7.6631 pred=0.3747 kl=0.0637 recon=2.4235 spec=2.9858 spk=0.0000 | grad=6.27 | 30.9s
  [1/200] step 400 | loss=6.2237 pred=0.3007 kl=0.0664 recon=1.9667 spec=2.4997 spk=0.0000 | grad=5.08 | 35.0s
Epoch 1/200 | avg_loss=11.6912 pred=0.7245 kl=0.0196 recon=2.2753 spec=4.3421 spk=0.0000 | 37.9s
  New best model (loss=11.6912)
  [2/200] step 450 | loss=4.7332 pred=0.2574 kl=0.0595 recon=1.3293 spec=1.6598 spk=0.0000 | grad=4.10 | 1.5s
  [2/200] step 500 | loss=3.9343 pred=0.2105 kl=0.0464 recon=1.1591 spec=1.3410 spk=0.0000 | grad=2.85 | 5.6s
  [2/200] step 550 | loss=3.3680 pred=0.1760 kl=0.0272 recon=1.0509 spec=1.1141 spk=0.0000 | grad=2.80 | 9.8s
  [2/200] step 600 | loss=2.9655 pred=0.1518 kl=0.0217 recon=0.9710 spec=0.9537 spk=0.0000 | grad=3.60 | 14.0s
  [2/200] step 650 | loss=2.7948 pred=0.1322 kl=0.0173 recon=0.9897 spec=0.9660 spk=0.0000 | grad=3.45 | 18.1s
  [2/200] step 700 | loss=2.5423 pred=0.1168 kl=0.0163 recon=0.9171 spec=0.9150 spk=0.0000 | grad=2.81 | 22.3s
  [2/200] step 750 | loss=2.3605 pred=0.0986 kl=0.0127 recon=0.8810 spec=0.9870 spk=0.0000 | grad=5.97 | 26.4s
  [2/200] step 800 | loss=2.2424 pred=0.0845 kl=0.0162 recon=0.9148 spec=0.9650 spk=0.0000 | grad=2.48 | 30.5s
  [2/200] step 850 | loss=2.0832 pred=0.0754 kl=0.0175 recon=0.8789 spec=0.9013 spk=0.0000 | grad=6.01 | 34.6s
Epoch 2/200 | avg_loss=2.9760 pred=0.1427 kl=0.0256 recon=1.0074 spec=1.0834 spk=0.0000 | 36.1s
  New best model (loss=2.9760)
  [3/200] step 900 | loss=1.9889 pred=0.0661 kl=0.0168 recon=0.8717 spec=0.9122 spk=0.0000 | grad=4.14 | 2.8s
  [3/200] step 950 | loss=1.8172 pred=0.0562 kl=0.0160 recon=0.8442 spec=0.8210 spk=0.0000 | grad=5.61 | 6.9s
  [3/200] step 1000 | loss=1.7564 pred=0.0477 kl=0.0210 recon=0.8399 spec=0.8778 spk=0.0000 | grad=4.54 | 11.0s
  [3/200] step 1050 | loss=1.6775 pred=0.0412 kl=0.0185 recon=0.8233 spec=0.8853 spk=0.0000 | grad=3.52 | 15.2s
  [3/200] step 1100 | loss=1.6286 pred=0.0355 kl=0.0193 recon=0.8145 spec=0.9175 spk=0.0000 | grad=8.08 | 19.2s


Training is converging fast on the real data. Let me run it fully in the background:


pkill -f "train.py.*output_v4" 2>/dev/null; sleep 1; echo "stopped"

Stop to restart in background


nohup python3 train.py \
  --data_dir /home/ubuntu/lewm-tts/processed_data_100mel \
  --output_dir ./output_v4 \
  --d_model 256 \
  --nhead 4 \
  --text_layers 4 \
  --audio_layers 4 \
  --predictor_layers 6 \
  --downsample_factor 4 \
  --kl_weight 0.05 \
  --recon_weight 1.0 \
  --spectral_weight 0.5 \
  --pred_weight 10.0 \
  --cosine_weight 1.0 \
  --ema_decay 0.998 \
  --free_bits 2.0 \
  --input_noise_max 0.2 \
  --n_speakers 1 \
  --epochs 200 \
  --batch_size 16 \
  --lr 1e-4 \
  --log_every 10 \
  --print_every 100 \
  --save_every 25 \
  > /home/ubuntu/lewm-tts/train_v4.log 2>&1 &
echo "PID: $!"

Start training in background


PID: 2554789


sleep 10 && tail -5 /home/ubuntu/lewm-tts/train_v4.log

Check training started


Device: cuda
LeWM TTS model: 24.07M parameters
Dataset: 6959 samples loaded
Training: 6959 samples, 434 batches/epoch


Training is running. Here's where things stand:

## Summary

**Found your 20h dataset** at `/home/ubuntu/lewm-tts/processed_data_100mel/` — 7009 samples, 17.88 hours. Audio existed but mels were missing. Extracted all 7009 mel spectrograms (zero errors).

**Training is running** on A100:
- 6959 samples (after length filtering), 434 batches/epoch
- ~38s/epoch → **200 epochs ≈ 2.1 hours**
- Loss was already at 1.63 by step 1100 (epoch 3) — converging well
- Checkpoints saving every 25 epochs to `output_v4/checkpoints/`

**PID:** 2554789
**Log:** `/home/ubuntu/lewm-tts/train_v4.log`

You can monitor with:
```bash
tail -f ~/lewm-tts/train_v4.log
```


The small training run (116 samples) finished. The real training on 7009 samples is still running in the background. Let me check progress:


Check training progress


Check if training process is alive


ubuntu   2554789  138  1.4 18667652 1764224 ?    Rl   09:04   0:45 python3 train.py --data_dir /home/ubuntu/lewm-tts/processed_data_100mel --output_dir ./output_v4 --d_model 256 --nhead 4 --text_layers 4 --audio_layers 4 --predictor_layers 6 --downsample_factor 4 --kl_weight 0.05 --recon_weight 1.0 --spectral_weight 0.5 --pred_weight 10.0 --cosine_weight 1.0 --ema_decay 0.998 --free_bits 2.0 --input_noise_max 0.2 --n_speakers 1 --epochs 200 --batch_size 16 --lr 1e-4 --log_every 10 --print_every 100 --save_every 25
ubuntu   2554878  4.6  0.4 11718764 557804 ?     Sl   09:04   0:01 python3 train.py --data_dir /home/ubuntu/lewm-tts/processed_data_100mel --output_dir ./output_v4 --d_model 256 --nhead 4 --text_layers 4 --audio_layers 4 --predictor_layers 6 --downsample_factor 4 --kl_weight 0.05 --recon_weight 1.0 --spectral_weight 0.5 --pred_weight 10.0 --cosine_weight 1.0 --ema_decay 0.998 --free_bits 2.0 --input_noise_max 0.2 --n_speakers 1 --epochs 200 --batch_size 16 --lr 1e-4 --log_every 10 --print_every 100 --save_every 25
ubuntu   2554879  4.6  0.4 11718776 558524 ?     Sl   09:04   0:01 python3 train.py --data_dir /home/ubuntu/lewm-tts/processed_data_100mel --output_dir ./output_v4 --d_model 256 --nhead 4 --text_layers 4 --audio_layers 4 --predictor_layers 6 --downsample_factor 4 --kl_weight 0.05 --recon_weight 1.0 --spectral_weight 0.5 --pred_weight 10.0 --cosine_weight 1.0 --ema_decay 0.998 --free_bits 2.0 --input_noise_max 0.2 --n_speakers 1 --epochs 200 --batch_size 16 --lr 1e-4 --log_every 10 --print_every 100 --save_every 25
ubuntu   2554880  4.8  0.4 11718788 558812 ?     Sl   09:04   0:01 python3 train.py --data_dir /home/ubuntu/lewm-tts/processed_data_100mel --output_dir ./output_v4 --d_model 256 --nhead 4 --text_layers 4 --audio_layers 4 --predictor_layers 6 --downsample_factor 4 --kl_weight 0.05 --recon_weight 1.0 --spectral_weight 0.5 --pred_weight 10.0 --cosine_weight 1.0 --ema_decay 0.998 --free_bits 2.0 --input_noise_max 0.2 --n_speakers 1 --epochs 200 --batch_size 16 --lr 1e-4 --log_every 10 --print_every 100 --save_every 25
ubuntu   2554885  5.2  0.4 11718800 558472 ?     Sl   09:04   0:01 python3 train.py --data_dir /home/ubuntu/lewm-tts/processed_data_100mel --output_dir ./output_v4 --d_model 256 --nhead 4 --text_layers 4 --audio_layers 4 --predictor_layers 6 --downsample_factor 4 --kl_weight 0.05 --recon_weight 1.0 --spectral_weight 0.5 --pred_weight 10.0 --cosine_weight 1.0 --ema_decay 0.998 --free_bits 2.0 --input_noise_max 0.2 --n_speakers 1 --epochs 200 --batch_size 16 --lr 1e-4 --log_every 10 --print_every 100 --save_every 25


sleep 30 && tail -10 /home/ubuntu/lewm-tts/train_v4.log

Wait and check log output


Device: cuda
LeWM TTS model: 24.07M parameters
Dataset: 6959 samples loaded
Training: 6959 samples, 434 batches/epoch
  [1/200] step 100 | loss=14.8622 pred=0.9961 kl=0.0014 recon=2.3254 spec=5.1520 spk=0.0000 | grad=23.70 | 9.6s
  [1/200] step 200 | loss=10.9639 pred=0.6280 kl=0.0067 recon=2.3630 spec=4.6421 spk=0.0000 | grad=13.36 | 18.1s
  [1/200] step 300 | loss=8.6789 pred=0.4336 kl=0.0521 recon=2.2738 spec=4.1383 spk=0.0000 | grad=7.33 | 26.4s
  [1/200] step 400 | loss=5.7563 pred=0.2804 kl=2.6590 recon=2.0398 spec=1.8182 spk=0.0000 | grad=6.40 | 34.7s
Epoch 1/200 | avg_loss=11.3280 pred=0.6874 kl=0.3916 recon=2.3150 spec=4.2768 spk=0.0000 | 37.6s
  New best model (loss=11.3280)


Training is going well. ~38s/epoch, so **200 epochs ≈ 2.1 hours**. Loss already dropped from 14.8 to 5.7 within epoch 1. It's running — should finish around 11:10 AM.


what is status


Check training progress


  New best model (loss=0.8636)
  [21/200] step 8700 | loss=0.8527 pred=0.0005 kl=0.0257 recon=0.5064 spec=0.6812 spk=0.0000 | grad=1.59 | 1.9s
  [21/200] step 8800 | loss=0.8461 pred=0.0005 kl=0.0318 recon=0.5202 spec=0.6394 spk=0.0000 | grad=3.53 | 10.0s
  [21/200] step 8900 | loss=0.8694 pred=0.0005 kl=0.0266 recon=0.5109 spec=0.7054 spk=0.0000 | grad=1.86 | 18.1s
  [21/200] step 9000 | loss=0.8441 pred=0.0005 kl=0.0262 recon=0.5100 spec=0.6566 spk=0.0000 | grad=1.41 | 26.3s
  [21/200] step 9100 | loss=0.8304 pred=0.0005 kl=0.0284 recon=0.5014 spec=0.6473 spk=0.0000 | grad=2.74 | 34.5s
Epoch 21/200 | avg_loss=0.8514 pred=0.0005 kl=0.0266 recon=0.5078 spec=0.6755 spk=0.0000 | 35.7s
  New best model (loss=0.8514)
  [22/200] step 9200 | loss=0.8417 pred=0.0004 kl=0.0273 recon=0.5087 spec=0.6558 spk=0.0000 | grad=1.61 | 7.3s
  [22/200] step 9300 | loss=0.8632 pred=0.0005 kl=0.0266 recon=0.5128 spec=0.6897 spk=0.0000 | grad=2.40 | 15.6s
  [22/200] step 9400 | loss=0.8765 pred=0.0004 kl=0.0264 recon=0.5219 spec=0.6988 spk=0.0000 | grad=1.06 | 23.9s
  [22/200] step 9500 | loss=0.8346 pred=0.0004 kl=0.0261 recon=0.4927 spec=0.6748 spk=0.0000 | grad=1.94 | 32.1s
Epoch 22/200 | avg_loss=0.8417 pred=0.0004 kl=0.0259 recon=0.5017 spec=0.6698 spk=0.0000 | 36.0s
  New best model (loss=0.8417)
  [23/200] step 9600 | loss=0.8435 pred=0.0004 kl=0.0258 recon=0.4965 spec=0.6847 spk=0.0000 | grad=1.19 | 4.5s
  [23/200] step 9700 | loss=0.8172 pred=0.0004 kl=0.0256 recon=0.4961 spec=0.6332 spk=0.0000 | grad=2.78 | 12.7s
  [23/200] step 9800 | loss=0.8147 pred=0.0004 kl=0.0239 recon=0.4925 spec=0.6361 spk=0.0000 | grad=1.61 | 21.0s
  [23/200] step 9900 | loss=0.8168 pred=0.0004 kl=0.0264 recon=0.4990 spec=0.6264 spk=0.0000 | grad=1.20 | 29.2s
Epoch 23/200 | avg_loss=0.8347 pred=0.0004 kl=0.0245 recon=0.4970 spec=0.6667 spk=0.0000 | 35.8s
  New best model (loss=0.8347)


**Epoch 23/200** — about 11% done.

| Metric | Value |
|---|---|
| Loss | **0.83** (down from 11.3 at epoch 1) |
| Prediction loss | 0.0004 (converged) |
| Recon loss | 0.497 |
| Spectral loss | 0.667 |
| Speed | ~36s/epoch |
| ETA | ~106 min (~1h 46m remaining) |

Converging well. Prediction loss is essentially solved at 0.0004 — the JEPA encoder is learning good representations. The remaining loss is almost entirely mel reconstruction + spectral, which will keep improving gradually over the remaining 177 epochs.


can we genaret a sample and test ..or we need to wait


Let's test now — epoch 23 checkpoint exists. Won't be perfect but will tell us if the pipeline works end-to-end.


Check available checkpoints


best.pt
epoch_0010.pt
epoch_0020.pt
epoch_0025.pt
epoch_0050.pt
epoch_0075.pt
epoch_0100.pt


python3 -c "
import torch
ckpt = torch.load('/home/ubuntu/lewm-tts/output_v4/checkpoints/best.pt', map_location='cpu', weights_only=False)
print(f'Best checkpoint: epoch {ckpt[\"epoch\"]+1}, loss={ckpt[\"loss\"]:.4f}')
" 2>&1

Check best checkpoint info


Best checkpoint: epoch 26, loss=0.8203


python3 -c "
import torch, json

# Quick inference test — reconstruct a real mel first (easiest test)
from model import LeWMTTS

ckpt = torch.load('output_v4/checkpoints/best.pt', map_location='cuda', weights_only=False)
config = ckpt['config']
model = LeWMTTS(config).cuda()
model.load_state_dict(ckpt['model'])
model.eval()

# Load a real mel
mel = torch.load('processed_data_100mel/mels/000050.pt', weights_only=True).unsqueeze(0).cuda()
print(f'Input mel: {mel.shape}')

with torch.no_grad():
    # Test 1: Encode → Decode (reconstruction)
    z, mu, logvar = model.audio_encoder(mel)
    mel_recon = model.mel_decoder(mu, None)
    print(f'Recon mel: {mel_recon.shape}')

    # Test 2: AR synthesis
    text = 'नमस्कार, मैं भारत का प्रधानमंत्री हूं।'
    tokens = list(text.encode('utf-8'))
    text_tokens = torch.tensor(tokens, dtype=torch.long).unsqueeze(0).cuda()
    text_emb = model.text_encoder(text_tokens)

    start_emb = model.start_emb
    cache = model.init_ar_cache()
    all_embs = [start_emb]

    next_emb, cache = model.predict_next_cached(start_emb, text_emb, 0, cache)
    all_embs.append(next_emb)

    for step in range(1, 150):
        next_emb, cache = model.predict_next_cached(next_emb, text_emb, step, cache)
        all_embs.append(next_emb)

    audio_embs = torch.cat(all_embs, dim=1)
    mel_synth = model.mel_decoder(audio_embs, None)
    print(f'Synth mel: {mel_synth.shape}')
    print(f'Synth mel range: [{mel_synth.min():.2f}, {mel_synth.max():.2f}]')

# Decode both with Vocos
from vocos import Vocos
vocos = Vocos.from_pretrained('charactr/vocos-mel-24khz')
vocos.eval()

import torchaudio
with torch.no_grad():
    # Reconstruction
    wav_recon = vocos.decode(mel_recon.cpu())
    torchaudio.save('test_recon.wav', wav_recon.unsqueeze(0) if wav_recon.dim()==1 else wav_recon, 24000)
    print(f'Reconstruction: {wav_recon.shape}, {wav_recon.shape[-1]/24000:.2f}s')

    # Synthesis
    wav_synth = vocos.decode(mel_synth.cpu())
    torchaudio.save('test_synth.wav', wav_synth.unsqueeze(0) if wav_synth.dim()==1 else wav_synth, 24000)
    print(f'Synthesis: {wav_synth.shape}, {wav_synth.shape[-1]/24000:.2f}s')

print('Done! Check test_recon.wav and test_synth.wav')
" 2>&1

Generate test samples from best checkpoint


Input mel: torch.Size([1, 100, 848])
Recon mel: torch.Size([1, 100, 848])
Synth mel: torch.Size([1, 100, 604])
Synth mel range: [-4.16, 1.33]
Reconstruction: torch.Size([1, 216832]), 9.03s
Synthesis: torch.Size([1, 154368]), 6.43s
Done! Check test_recon.wav and test_synth.wav


Both worked. Let me check the audio quality quickly:


python3 -c "
import torch, torchaudio

# Check reconstruction quality
wav_r, sr = torchaudio.load('test_recon.wav')
wav_s, sr = torchaudio.load('test_synth.wav')

print(f'Recon: {wav_r.shape[1]/sr:.2f}s, peak={wav_r.abs().max():.4f}, rms={wav_r.pow(2).mean().sqrt():.4f}')
print(f'Synth: {wav_s.shape[1]/sr:.2f}s, peak={wav_s.abs().max():.4f}, rms={wav_s.pow(2).mean().sqrt():.4f}')

# Check if synth is not silence/noise
if wav_s.abs().max() < 0.01:
    print('WARNING: Synthesis output is near-silent')
elif wav_s.pow(2).mean().sqrt() < 0.005:
    print('WARNING: Synthesis output is very quiet')
else:
    print('Synthesis has signal — not silent/collapsed')

# Compare original
wav_orig, sr_orig = torchaudio.load('processed_data_100mel/audio/000050.wav')
print(f'Original: {wav_orig.shape[1]/sr_orig:.2f}s, peak={wav_orig.abs().max():.4f}, rms={wav_orig.pow(2).mean().sqrt():.4f}')
" 2>&1

Check audio signal quality


Recon: 9.03s, peak=0.9381, rms=0.0730
Synth: 6.43s, peak=0.0562, rms=0.0164
Synthesis has signal — not silent/collapsed
Original: 9.04s, peak=0.8817, rms=0.0838


**Pipeline works end-to-end.** Results at epoch 26:

| | Duration | Peak | RMS |
|---|---|---|---|
| **Original** | 9.04s | 0.88 | 0.084 |
| **Reconstruction** | 9.03s | 0.94 | 0.073 |
| **Synthesis** | 6.43s | 0.06 | 0.016 |

- **Reconstruction** is solid — close to original signal levels. The encode→decode path is working well.
- **Synthesis** has signal (not collapsed/silent) but is quiet — peak 0.06 vs 0.88 original. Expected at epoch 26 — the AR predictor is still learning. The mel values are in a narrower range [-4.16, 1.33] vs what real mels look like. This will improve as training continues.

Files saved:
- `test_recon.wav` — reconstruction (should sound close to original)
- `test_synth.wav` — AR synthesis (early, will be rough/quiet)

You can listen to them now. The reconstruction should sound recognizable. Synthesis will improve a lot by epoch 100+. Training is still running — about 1.5 hours left.


how reconctrution is dng ,, it is audiable and close to origina but some nise .. but synth is total noice


That's actually a good sign at epoch 26.

**Reconstruction working = the core architecture is sound.** The audio encoder is learning good representations, the mel decoder can turn them back into recognizable speech, and Vocos can decode it. The noise will clean up as recon_loss (currently 0.49) keeps dropping.

**Synthesis being noise = expected.** The prediction loss is 0.0004 which looks converged, but that's deceptive. Here's why:

The predictor is trained with **teacher forcing** — at each step it gets the *real* encoder output as input and predicts the next one. But at inference, it feeds its *own* predictions back in. Any small error compounds across 150 steps. This is the classic **exposure bias** problem.

The `input_noise` annealing is supposed to bridge this gap, but at epoch 26 it's only ramped up to `0.2 * min(1.0, 26*434 / (200*434*0.3))` ≈ `0.2 * 0.43` = **0.086**. It hasn't reached full strength yet. By epoch ~60 it'll be at the full 0.2, and the predictor will have trained with enough noise to handle its own imperfect inputs.

**What to expect:**
- **Epoch 50-75:** Synthesis starts being intelligible but robotic
- **Epoch 100-150:** Speaker character emerges, sounds recognizable
- **Epoch 150-200:** Quality stabilizes, noise drops

Let it cook. Want me to auto-generate samples at epoch 50, 100, and 150 so you can hear the progression?


status


Check training progress


  [66/200] step 28400 | loss=0.7147 pred=0.0000 kl=0.0190 recon=0.4001 spec=0.6270 spk=0.0000 | grad=1.32 | 15.4s
  [66/200] step 28500 | loss=0.7185 pred=0.0000 kl=0.0215 recon=0.4079 spec=0.6187 spk=0.0000 | grad=0.75 | 23.6s
  [66/200] step 28600 | loss=0.7154 pred=0.0000 kl=0.0192 recon=0.4089 spec=0.6107 spk=0.0000 | grad=1.86 | 31.6s
Epoch 66/200 | avg_loss=0.7068 pred=0.0000 kl=0.0194 recon=0.4029 spec=0.6054 spk=0.0000 | 35.1s
  New best model (loss=0.7068)


**Epoch 66/200** — 33% done.

| Metric | Epoch 23 | Epoch 66 | Trend |
|---|---|---|---|
| Total loss | 0.83 | **0.71** | still dropping |
| Prediction | 0.0004 | 0.0000 | converged |
| Recon | 0.50 | **0.40** | improving |
| Spectral | 0.67 | **0.61** | improving |
| Speed | 36s/ep | 35s/ep | steady |

**ETA: ~47 min remaining.** Input noise should be near full strength now (epoch 60 = 0.3 * 200 = 60, so it's at max 0.2). Good time to test synthesis again — want me to generate a sample from the current best checkpoint?


generate samples


python3 -c "
import torch, torchaudio, numpy as np
from model import LeWMTTS
from vocos import Vocos

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

vocos = Vocos.from_pretrained('charactr/vocos-mel-24khz').eval()

def synthesize(model, text, max_steps=200):
    tokens = list(text.encode('utf-8'))
    text_tokens = torch.tensor(tokens, dtype=torch.long).unsqueeze(0).cuda()
    text_emb = model.text_encoder(text_tokens)

    cache = model.init_ar_cache()
    start = model.start_emb
    all_embs = [start]

    emb = start
    for step in range(max_steps):
        emb, cache = model.predict_next_cached(emb, text_emb, step, cache)
        all_embs.append(emb)

    audio_embs = torch.cat(all_embs, dim=1)
    mel = model.mel_decoder(audio_embs, None)
    wav = vocos.decode(mel.cpu())
    return wav.squeeze().numpy()

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

with torch.no_grad():
    # Reconstruction test
    mel_orig = torch.load('processed_data_100mel/mels/000050.pt', weights_only=True).unsqueeze(0).cuda()
    z, mu, _ = model.audio_encoder(mel_orig)
    mel_recon = model.mel_decoder(mu, None)
    wav_recon = vocos.decode(mel_recon.cpu()).squeeze().numpy()
    save(wav_recon, 'sample_ep66_recon.wav')

    # Synthesis - Hindi texts
    print('Generating synthesis samples...')

    texts = [
        'नमस्कार, मैं भारत का प्रधानमंत्री हूं।',
        'मेरे प्यारे देशवासियों, आप सभी को मेरा नमस्कार।',
        'हम सब मिलकर भारत को आगे बढ़ाएंगे।',
    ]

    for i, text in enumerate(texts):
        wav = synthesize(model, text, max_steps=200)
        save(wav, f'sample_ep66_synth_{i+1}.wav')
        print(f'    text: {text}')

print('Done!')
" 2>&1

Generate samples from epoch 66 checkpoint


Loaded epoch 68, loss=0.7060
  sample_ep66_recon.wav: 9.03s
Generating synthesis samples...
  sample_ep66_synth_1.wav: 8.57s
    text: नमस्कार, मैं भारत का प्रधानमंत्री हूं।
  sample_ep66_synth_2.wav: 8.57s
    text: मेरे प्यारे देशवासियों, आप सभी को मेरा नमस्कार।
  sample_ep66_synth_3.wav: 8.57s
    text: हम सब मिलकर भारत को आगे बढ़ाएंगे।
Done!


python3 -c "
import torchaudio
for f in ['sample_ep66_recon.wav', 'sample_ep66_synth_1.wav', 'sample_ep66_synth_2.wav', 'sample_ep66_synth_3.wav']:
    wav, sr = torchaudio.load(f)
    print(f'{f}: {wav.shape[1]/sr:.2f}s  peak={wav.abs().max():.3f}  rms={wav.pow(2).mean().sqrt():.4f}')
" 2>&1

Check sample signal levels


sample_ep66_recon.wav: 9.03s  peak=0.950  rms=0.0821
sample_ep66_synth_1.wav: 8.57s  peak=0.950  rms=0.2731
sample_ep66_synth_2.wav: 8.57s  peak=0.950  rms=0.2731
sample_ep66_synth_3.wav: 8.57s  peak=0.950  rms=0.2731


Samples generated from **epoch 68** (best checkpoint, loss=0.706):

| File | Duration | RMS | Notes |
|---|---|---|---|
| `sample_ep66_recon.wav` | 9.03s | 0.082 | Reconstruction — should sound close to original |
| `sample_ep66_synth_1.wav` | 8.57s | 0.273 | "नमस्कार, मैं भारत का प्रधानमंत्री हूं।" |
| `sample_ep66_synth_2.wav` | 8.57s | 0.273 | "मेरे प्यारे देशवासियों, आप सभी को मेरा नमस्कार।" |
| `sample_ep66_synth_3.wav` | 8.57s | 0.273 | "हम सब मिलकर भारत को आगे बढ़ाएंगे।" |

Signal levels look much better than epoch 26 — synthesis RMS is 0.273 now vs 0.016 before. The synth samples all have identical RMS/duration though, which suggests the AR might be producing similar output regardless of text. That could mean:

1. The predictor is ignoring text (cross-attention hasn't learned alignment yet)
2. Or it's producing a learned "average" pattern

Have a listen — tell me what you hear and I'll diagnose from there.


recon is good .. synth pure noise .. let wait .. we will get audio ryt after 200 epochs


Reconstruction being good confirms the encoder + decoder path is solid. The synth noise means the AR predictor hasn't learned to generate coherent sequences on its own yet — it still relies on teacher-forced real embeddings.

Honestly though, I'm not confident just more epochs will fix synthesis. The prediction loss is already 0.0000 — it can't go lower. The problem isn't that the predictor can't predict the next embedding given real inputs. It's that it **can't recover from its own errors** during free-running generation.

Signs that worry me:
- All 3 synth samples have **identical RMS and duration** — the AR is collapsing to a fixed pattern
- Pred loss at 0.0000 but synth is noise — classic exposure bias

Let it finish the 200 epochs. But if synth is still noise at epoch 200, the fix will be:
1. **Increase input_noise_max** from 0.2 to 0.5+ (force the predictor to handle worse inputs)
2. **Scheduled sampling** — sometimes feed predicted embeddings during training instead of always real ones
3. **Shorter generation** — try max_steps=50 instead of 200 to see if short utterances work

We'll know in ~40 minutes. I'll generate samples at epoch 200 automatically.


Check training progress


  [151/200] step 65500 | loss=0.6268 pred=0.0000 kl=0.0255 recon=0.3545 spec=0.5420 spk=0.0000 | grad=0.92 | 32.8s
Epoch 151/200 | avg_loss=0.6375 pred=0.0000 kl=0.0252 recon=0.3530 spec=0.5665 spk=0.0000 | 35.6s
  New best model (loss=0.6375)


**Epoch 151/200** — 75% done. ~17 min left.

| Metric | Ep 23 | Ep 66 | Ep 151 |
|---|---|---|---|
| Total | 0.83 | 0.71 | **0.64** |
| Recon | 0.50 | 0.40 | **0.35** |
| Spectral | 0.67 | 0.61 | **0.57** |

Still improving. Will generate samples as soon as it finishes.


run samples


python3 -c "
import torch, torchaudio, numpy as np
from model import LeWMTTS
from vocos import Vocos

ckpt = torch.load('output_v4/checkpoints/best.pt', map_location='cuda', weights_only=False)
model = LeWMTTS(ckpt['config']).cuda()
model.load_state_dict(ckpt['model'])
model.eval()
print(f'Loaded epoch {ckpt[\"epoch\"]+1}, loss={ckpt[\"loss\"]:.4f}')

vocos = Vocos.from_pretrained('charactr/vocos-mel-24khz').eval()

def synthesize(model, text, max_steps=200, temperature=0.0):
    tokens = list(text.encode('utf-8'))
    text_tokens = torch.tensor(tokens, dtype=torch.long).unsqueeze(0).cuda()
    text_emb = model.text_encoder(text_tokens)
    cache = model.init_ar_cache()
    emb = model.start_emb
    all_embs = [emb]
    for step in range(max_steps):
        emb, cache = model.predict_next_cached(emb, text_emb, step, cache)
        if temperature > 0:
            emb = emb + torch.randn_like(emb) * temperature
        all_embs.append(emb)
    mel = model.mel_decoder(torch.cat(all_embs, dim=1), None)
    return vocos.decode(mel.cpu()).squeeze().numpy()

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

with torch.no_grad():
    # Recon
    mel_orig = torch.load('processed_data_100mel/mels/000050.pt', weights_only=True).unsqueeze(0).cuda()
    z, mu, _ = model.audio_encoder(mel_orig)
    mel_recon = model.mel_decoder(mu, None)
    save(vocos.decode(mel_recon.cpu()).squeeze().numpy(), 'ep151_recon.wav')

    # Synth at different temperatures and lengths
    text = 'नमस्कार, मैं भारत का प्रधानमंत्री हूं।'

    save(synthesize(model, text, max_steps=50), 'ep151_synth_short.wav')
    save(synthesize(model, text, max_steps=200), 'ep151_synth_200.wav')
    save(synthesize(model, text, max_steps=200, temperature=0.3), 'ep151_synth_temp03.wav')
    save(synthesize(model, text, max_steps=200, temperature=0.6), 'ep151_synth_temp06.wav')

print('Done')
" 2>&1

Generate samples at epoch 151


[W330 10:35:47.760654721 NNPACK.cpp:56] Could not initialize NNPACK! Reason: Unsupported hardware.
Loaded epoch 151, loss=0.6375
  ep151_recon.wav: 9.03s rms=0.0834
  ep151_synth_short.wav: 2.17s rms=0.2236
  ep151_synth_200.wav: 8.57s rms=0.2239
  ep151_synth_temp03.wav: 8.57s rms=0.1368
  ep151_synth_temp06.wav: 8.57s rms=0.1150
Done


Generated 5 samples from **epoch 151**:

| File | What | Duration |
|---|---|---|
| `ep151_recon.wav` | Reconstruction | 9.03s |
| `ep151_synth_short.wav` | AR, 50 steps only | 2.17s |
| `ep151_synth_200.wav` | AR, 200 steps, temp=0 | 8.57s |
| `ep151_synth_temp03.wav` | AR, 200 steps, temp=0.3 | 8.57s |
| `ep151_synth_temp06.wav` | AR, 200 steps, temp=0.6 | 8.57s |

The short 50-step and the 200-step still have identical-ish RMS patterns (0.224 vs 0.224). Temperature adds some variation. Have a listen — particularly compare the short one vs the long one. If the short one sounds any better, it means error accumulation is the problem.


still noise full noise all synth ... recon is sounding like original ..


Recon working + synth noise = **exposure bias**. The predictor perfectly predicts given real encoder outputs, but has never seen its own predictions as input. Let me diagnose exactly how bad the mismatch is:


python3 -c "
import torch
from model import LeWMTTS

ckpt = torch.load('output_v4/checkpoints/best.pt', map_location='cuda', weights_only=False)
model = LeWMTTS(ckpt['config']).cuda()
model.load_state_dict(ckpt['model'])
model.eval()

# Load a real sample
mel = torch.load('processed_data_100mel/mels/000050.pt', weights_only=True).unsqueeze(0).cuda()
tokens = list('नमस्कार मेरे प्यारे देशवासियों'.encode('utf-8'))
text_tokens = torch.tensor(tokens, dtype=torch.long).unsqueeze(0).cuda()

with torch.no_grad():
    text_emb = model.text_encoder(text_tokens)
    z, mu, _ = model.audio_encoder(mel)

    # What the predictor sees during TRAINING (real encoder outputs)
    start = model.start_emb.expand(1, -1, -1)
    train_input = torch.cat([start, z[:, :-1]], dim=1)

    # What the predictor PRODUCES
    predicted = model.predictor(train_input, text_emb)

    # Compare distributions
    print('=== Distribution mismatch ===')
    print(f'Encoder z  - mean: {z.mean():.4f}, std: {z.std():.4f}, range: [{z.min():.3f}, {z.max():.3f}]')
    print(f'Predicted  - mean: {predicted.mean():.4f}, std: {predicted.std():.4f}, range: [{predicted.min():.3f}, {predicted.max():.3f}]')

    # EMA targets (what predicted is trained to match)
    target = model._compute_multiscale_targets(mel, None)
    print(f'EMA target - mean: {target.mean():.4f}, std: {target.std():.4f}, range: [{target.min():.3f}, {target.max():.3f}]')

    # The gap: predicted is in EMA-target space, but gets fed back as input (encoder-z space)
    gap = (predicted - z[:, :predicted.shape[1]]).pow(2).mean().sqrt()
    print(f'')
    print(f'RMS gap (predicted vs encoder z): {gap:.4f}')
    print(f'RMS of encoder z: {z.pow(2).mean().sqrt():.4f}')
    print(f'Gap as % of signal: {gap/z.pow(2).mean().sqrt()*100:.1f}%')

    # Cosine similarity between predicted and z
    cos = torch.nn.functional.cosine_similarity(
        predicted.reshape(-1, 256), z[:, :predicted.shape[1]].reshape(-1, 256), dim=-1
    ).mean()
    print(f'Cosine sim (predicted vs z): {cos:.4f}')

    # Now simulate AR for a few steps and see how fast it diverges
    print('')
    print('=== AR divergence test ===')
    cache = model.init_ar_cache()
    emb = model.start_emb

    for step in range(10):
        emb, cache = model.predict_next_cached(emb, text_emb, step, cache)
        if step < z.shape[1]:
            real = z[:, step:step+1]
            step_gap = (emb - real).pow(2).mean().sqrt().item()
            step_cos = torch.nn.functional.cosine_similarity(emb.squeeze(), real.squeeze(), dim=-1).item()
            print(f'  Step {step:2d}: gap={step_gap:.4f}, cos={step_cos:.4f}, emb_norm={emb.norm():.4f}, real_norm={real.norm():.4f}')
" 2>&1

Diagnose distribution mismatch


=== Distribution mismatch ===
Encoder z  - mean: 0.0093, std: 0.5354, range: [-1.774, 1.864]
Predicted  - mean: 0.0107, std: 0.1293, range: [-0.260, 0.278]
EMA target - mean: 0.0106, std: 0.1290, range: [-0.261, 0.279]

RMS gap (predicted vs encoder z): 0.5502
RMS of encoder z: 0.5355
Gap as % of signal: 102.7%
Cosine sim (predicted vs z): 0.0086

=== AR divergence test ===
  Step  0: gap=0.6451, cos=0.0117, emb_norm=2.0619, real_norm=10.1370
  Step  1: gap=0.6114, cos=-0.1191, emb_norm=2.0612, real_norm=9.3206
  Step  2: gap=0.6252, cos=-0.1521, emb_norm=2.0665, real_norm=9.4776
  Step  3: gap=0.6265, cos=-0.0996, emb_norm=2.0728, real_norm=9.6035
  Step  4: gap=0.5160, cos=0.0062, emb_norm=2.0761, real_norm=8.0045
  Step  5: gap=0.5019, cos=0.0563, emb_norm=2.0787, real_norm=7.8748
  Step  6: gap=0.5149, cos=0.1054, emb_norm=2.0804, real_norm=8.1932
  Step  7: gap=0.4591, cos=0.0618, emb_norm=2.0809, real_norm=7.1745
  Step  8: gap=0.4701, cos=0.1226, emb_norm=2.0815, real_norm=7.4867
  Step  9: gap=0.5264, cos=0.0842, emb_norm=2.0836, real_norm=8.3379


**Found the problem.** It's crystal clear:

```
Encoder z   — std: 0.535, norm: ~9.0
Predicted   — std: 0.129, norm: ~2.0
```

The predictor outputs are **4x smaller** than encoder outputs. They live in completely different scales. Cosine similarity is ~0.01 (essentially random). The predictor learned to match EMA targets (std=0.129) perfectly, but encoder z (std=0.535) is what gets fed as input during training. At inference, the predictor feeds its own outputs (std=0.129) back in — a scale the predictor has never seen as input.

**Root cause:** Multi-scale EMA targets + `ema_target_proj` compressed the target space to a much smaller range than the encoder output space. The predictor learns to predict in this compressed space, but AR inference needs it to produce embeddings in the encoder's scale.

**The fix:** Make predicted outputs match encoder z scale, not EMA target scale. Two options:

1. **Normalize targets to match encoder z scale** (simplest)
2. **Feed predicted embeddings during training** (scheduled sampling)

Let me do both — they're complementary:


470	        x = self.res1(x, spk_cond)
471	        if self.up2 is not None:
472	            x = F.gelu(self.up2(x))
473	            x = self.res2(x, spk_cond)
474
475	        x = F.gelu(self.refine_conv(x))
476	        x = self.refine_res(x, spk_cond)
477	        return self.out_conv(x)  # [B, n_mels, T_up]
478
479
480	# ─── Multi-Resolution Spectral Loss ─────────────────────────────────────────
481
482	class MultiResolutionSpectralLoss(nn.Module):
483	    """Computes spectral convergence + log magnitude loss at multiple STFT resolutions.
484	    This forces the mel decoder to produce sharp, well-defined spectral structure
485	    rather than blurry averages that Vocos can't decode properly."""
486
487	    def __init__(self, resolutions=((64, 16, 64), (128, 32, 128), (256, 64, 256))):
488	        super().__init__()
489	        # Each resolution is (n_fft, hop_length, win_length)
490	        self.resolutions = resolutions
491
492	    def _stft_loss(self, pred, target, n_fft, hop_length, win_length):
493	        """Compute spectral convergence + log magnitude loss for one resolution."""
494	        window = torch.hann_window(win_length, device=pred.device)
495
496	        # pred/target: [B, n_mels, T] — treat as 1D signal per mel bin
497	        # Flatten mel bins into batch: [B*n_mels, T]
498	        B, M, T = pred.shape
499	        pred_flat = pred.reshape(B * M, T)

        x = self.res1(x, spk_cond)
        if self.up2 is not None:
            x = F.gelu(self.up2(x))
            x = self.res2(x, spk_cond)

        x = F.gelu(self.refine_conv(x))
        x = self.refine_res(x, spk_cond)
        return self.out_conv(x)  # [B, n_mels, T_up]

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

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

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

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

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


557	        predictor_layers = config.get("predictor_layers", 6)
558	        dropout = config.get("dropout", 0.1)
559	        self.kl_weight = config.get("kl_weight", 0.01)
560	        n_speakers = config.get("n_speakers", 1)
561
562	        self.text_encoder = TextEncoder(
563	            vocab_size=text_vocab, d_model=d_model, nhead=nhead,
564	            num_layers=text_layers, dropout=dropout,
565	        )
566	        downsample_factor = config.get("downsample_factor", 4)
567	        self.audio_encoder = AudioEncoder(
568	            n_mels=n_mels, d_model=d_model, nhead=nhead,
569	            num_layers=audio_layers, downsample_factor=downsample_factor,
570	            dropout=dropout,
571	        )
572	        self.predictor = JEPAPredictor(
573	            d_model=d_model, nhead=nhead,
574	            num_layers=predictor_layers, dropout=dropout,
575	        )
576
577	        # Mel decoder: speaker-conditioned reconstruction
578	        self.mel_decoder = MelDecoder(
579	            d_model=d_model, n_mels=n_mels,
580	            upsample_factor=downsample_factor, n_speakers=n_speakers,
581	        )
582	        self.recon_weight = config.get("recon_weight", 1.0)
583
584	        # EMA target encoder — core JEPA component
585	        self.ema_audio_encoder = copy.deepcopy(self.audio_encoder)
586	        for p in self.ema_audio_encoder.parameters():
587	            p.requires_grad = False
588	        self.ema_decay = config.get("ema_decay", 0.998)
589
590	        # Multi-scale target: which layers to average for EMA targets
591	        # data2vec-style: average top-K layers instead of just the last
592	        self.ema_target_layers = config.get("ema_target_layers", None)  # None = all layers
593	        self.ema_target_proj = nn.Linear(d_model, d_model)  # project averaged intermediates
594
595	        # Speaker embedding + FiLM conditioning (for multi-speaker)
596	        self.n_speakers = n_speakers
597	        if n_speakers > 1:
598	            self.speaker_embed = nn.Embedding(n_speakers, d_model)
599	            # FiLM for latent space (replaces additive injection)
600	            self.latent_film = FiLMConditioner(d_model, d_model)
601	            self.target_film = FiLMConditioner(d_model, d_model)
602	            # Speaker classifier — forces encoder to retain speaker info
603	            self.speaker_classifier = nn.Sequential(
604	                nn.Linear(d_model, d_model),
605	                nn.GELU(),
606	                nn.Linear(d_model, n_speakers),
607	            )
608	        else:
609	            self.speaker_embed = None
610	            self.latent_film = None
611	            self.target_film = None
612	            self.speaker_classifier = None
613
614	        # Learnable start-of-sequence embedding for AR inference
615	        self.start_emb = nn.Parameter(torch.randn(1, 1, d_model) * 0.02)
616
617	        # Prediction weight — primary JEPA loss
618	        self.pred_weight = config.get("pred_weight", 10.0)
619
620	        # Cosine similarity weight in prediction loss
621	        self.cosine_weight = config.get("cosine_weight", 1.0)
622
623	        # Input noise injection to bridge train/inference gap
624	        self.input_noise = config.get("input_noise", 0.0)
625
626	        # Free-bits KL: allow KL up to this threshold without penalty
627	        self.free_bits = config.get("free_bits", 1.0)
628
629	        # Multi-resolution spectral loss for sharper mel output
630	        self.spectral_loss_fn = MultiResolutionSpectralLoss()
631	        self.spectral_weight = config.get("spectral_weight", 0.5)
632
633	        # Speaker consistency loss weight
634	        self.speaker_weight = config.get("speaker_weight", 1.0)
635
636	        self.config = config
637
638	    def _get_speaker_cond(self, speaker_id):
639	        """Get speaker conditioning vector. Returns [B, d_model] or None."""
640	        if self.speaker_embed is not None and speaker_id is not None:
641	            return self.speaker_embed(speaker_id)  # [B, d_model]
642	        return None
643
644	    def _apply_speaker_film(self, x, spk_cond, film_module):
645	        """Apply FiLM conditioning to tensor. No-op if no speaker."""
646	        if film_module is not None and spk_cond is not None:

        predictor_layers = config.get("predictor_layers", 6)
        dropout = config.get("dropout", 0.1)
        self.kl_weight = config.get("kl_weight", 0.01)
        n_speakers = config.get("n_speakers", 1)

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

        # Mel decoder: speaker-conditioned reconstruction
        self.mel_decoder = MelDecoder(
            d_model=d_model, n_mels=n_mels,
            upsample_factor=downsample_factor, n_speakers=n_speakers,
        )
        self.recon_weight = config.get("recon_weight", 1.0)

        # EMA target encoder — core JEPA component
        self.ema_audio_encoder = copy.deepcopy(self.audio_encoder)
        for p in self.ema_audio_encoder.parameters():
            p.requires_grad = False
        self.ema_decay = config.get("ema_decay", 0.998)

        # Multi-scale target: which layers to average for EMA targets
        # data2vec-style: average top-K layers instead of just the last
        self.ema_target_layers = config.get("ema_target_layers", None)  # None = all layers
        self.ema_target_proj = nn.Linear(d_model, d_model)  # project averaged intermediates

        # Speaker embedding + FiLM conditioning (for multi-speaker)
        self.n_speakers = n_speakers
        if n_speakers > 1:
            self.speaker_embed = nn.Embedding(n_speakers, d_model)
            # FiLM for latent space (replaces additive injection)
            self.latent_film = FiLMConditioner(d_model, d_model)
            self.target_film = FiLMConditioner(d_model, d_model)
            # Speaker classifier — forces encoder to retain speaker info
            self.speaker_classifier = nn.Sequential(
                nn.Linear(d_model, d_model),
                nn.GELU(),
                nn.Linear(d_model, n_speakers),
            )
        else:
            self.speaker_embed = None
            self.latent_film = None
            self.target_film = None
            self.speaker_classifier = None

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

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

        # Cosine similarity weight in prediction loss
        self.cosine_weight = config.get("cosine_weight", 1.0)

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

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

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

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

        self.config = config

    def _get_speaker_cond(self, speaker_id):
        """Get speaker conditioning vector. Returns [B, d_model] or None."""
        if self.speaker_embed is not None and speaker_id is not None:
            return self.speaker_embed(speaker_id)  # [B, d_model]
        return None

    def _apply_speaker_film(self, x, spk_cond, film_module):
        """Apply FiLM conditioning to tensor. No-op if no speaker."""
        if film_module is not None and spk_cond is not None:


648	        return x
649
650	    @torch.no_grad()
651	    def update_ema(self):
652	        """Update EMA target encoder weights."""
653	        for p_online, p_ema in zip(self.audio_encoder.parameters(),
654	                                   self.ema_audio_encoder.parameters()):
655	            p_ema.data.mul_(self.ema_decay).add_(p_online.data, alpha=1 - self.ema_decay)
656
657	    def _compute_multiscale_targets(self, mel, mel_mask):
658	        """Compute multi-scale EMA targets by averaging intermediate layer outputs.
659	        data2vec-style: captures both low-level acoustic detail (early layers)
660	        and high-level semantic structure (late layers)."""
661	        with torch.no_grad():
662	            self.ema_audio_encoder.eval()
663	            _, mu_ema, _, intermediates = self.ema_audio_encoder(
664	                mel, mel_mask, return_intermediates=True
665	            )
666	            self.ema_audio_encoder.train()
667
668	            # Select which layers to average
669	            if self.ema_target_layers is not None:
670	                selected = [intermediates[i] for i in self.ema_target_layers]
671	            else:
672	                selected = intermediates  # all layers
673
674	            # Average across selected layers
675	            target_emb = torch.stack(selected, dim=0).mean(dim=0)  # [B, T, d]
676
677	        # Project through learned projection (allows model to adapt target space)
678	        target_emb = self.ema_target_proj(target_emb)
679	        return target_emb
680
681	    def forward(self, mel, text_tokens, mel_mask=None, text_mask=None, speaker_id=None):
682	        """
683	        Args:
684	            mel: [B, n_mels, T_mel]
685	            text_tokens: [B, T_text]
686	            mel_mask: [B, T_mel] bool
687	            text_mask: [B, T_text] bool
688	            speaker_id: [B] long tensor of speaker IDs (optional)
689	        Returns:
690	            loss_dict with all individual losses
691	        """
692	        # Get speaker conditioning
693	        spk_cond = self._get_speaker_cond(speaker_id)  # [B, d] or None
694
695	        # Encode text + FiLM speaker conditioning on text
696	        text_emb = self.text_encoder(text_tokens, text_mask)
697	        if spk_cond is not None:
698	            # FiLM on text embeddings (replaces simple addition)
699	            text_emb = self._apply_speaker_film(text_emb, spk_cond, self.latent_film)
700
701	        # Online encoder → latent embeddings
702	        z, mu, logvar = self.audio_encoder(mel, mel_mask)
703	        B = z.shape[0]
704	        T_down = z.shape[1]
705
706	        # Save raw mu before speaker conditioning — for speaker classifier
707	        mu_raw = mu
708
709	        # Multi-scale EMA targets (data2vec-style: average layer outputs)
710	        target_emb = self._compute_multiscale_targets(mel, mel_mask)
711
712	        # Apply FiLM speaker conditioning to targets and encoder output
713	        # This makes prediction loss require speaker knowledge
714	        target_emb = self._apply_speaker_film(target_emb, spk_cond, self.target_film)
715	        z = self._apply_speaker_film(z, spk_cond, self.latent_film)
716
717	        # Prepend learnable start embedding to input sequence
718	        start = self.start_emb.expand(B, -1, -1)  # [B, 1, d]
719	        input_emb = torch.cat([start, z[:, :-1]], dim=1)  # [B, T_down, d]
720
721	        # Input noise: only on positions 1+ (real embeddings), not start_emb
722	        if self.training and self.input_noise > 0:
723	            noise = torch.zeros_like(input_emb)
724	            noise[:, 1:] = torch.randn(B, T_down - 1, input_emb.shape[-1],
725	                                        device=input_emb.device) * self.input_noise
726	            input_emb = input_emb + noise
727
728	        # Adjust masks for downsampled sequence
729	        if mel_mask is not None:
730	            ds_mask = mel_mask[:, ::self.audio_encoder.downsample_factor][:, :T_down]
731	            start_mask = torch.zeros(B, 1, dtype=torch.bool, device=z.device)
732	            pred_mask = torch.cat([start_mask, ds_mask[:, :-1]], dim=1)
733	            loss_mask = ds_mask
734	        else:
735	            pred_mask = None
736	            loss_mask = None
737
738	        predicted = self.predictor(input_emb, text_emb, pred_mask, text_mask)
739
740	        # ─── Loss 1: Prediction loss (MSE + cosine) against multi-scale EMA targets ───
741	        if loss_mask is not None:
742	            valid = (~loss_mask).unsqueeze(-1).float()  # [B, T_down, 1]
743	            n_valid = valid.sum() * predicted.shape[-1] + 1e-8
744
745	            # MSE component
746	            mse_loss = (((predicted - target_emb) ** 2) * valid).sum() / n_valid
747
748	            # Cosine component — preserves direction (speaker identity)
749	            pred_flat = (predicted * valid).reshape(-1, predicted.shape[-1])
750	            tgt_flat = (target_emb * valid).reshape(-1, target_emb.shape[-1])
751	            # Only compute on non-padding positions
752	            valid_flat = valid.squeeze(-1).reshape(-1) > 0
753	            if valid_flat.any():
754	                cosine_loss = 1.0 - F.cosine_similarity(
755	                    pred_flat[valid_flat], tgt_flat[valid_flat], dim=-1
756	                ).mean()
757	            else:
758	                cosine_loss = torch.tensor(0.0, device=mel.device)
759	        else:
760	            mse_loss = F.mse_loss(predicted, target_emb)
761	            cosine_loss = 1.0 - F.cosine_similarity(
762	                predicted.reshape(-1, predicted.shape[-1]),
763	                target_emb.reshape(-1, target_emb.shape[-1]),
764	                dim=-1,
765	            ).mean()
766
767	        prediction_loss = mse_loss + self.cosine_weight * cosine_loss

        return x

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

    def _compute_multiscale_targets(self, mel, mel_mask):
        """Compute multi-scale EMA targets by averaging intermediate layer outputs.
        data2vec-style: captures both low-level acoustic detail (early layers)
        and high-level semantic structure (late layers)."""
        with torch.no_grad():
            self.ema_audio_encoder.eval()
            _, mu_ema, _, intermediates = self.ema_audio_encoder(
                mel, mel_mask, return_intermediates=True
            )
            self.ema_audio_encoder.train()

            # Select which layers to average
            if self.ema_target_layers is not None:
                selected = [intermediates[i] for i in self.ema_target_layers]
            else:
                selected = intermediates  # all layers

            # Average across selected layers
            target_emb = torch.stack(selected, dim=0).mean(dim=0)  # [B, T, d]

        # Project through learned projection (allows model to adapt target space)
        target_emb = self.ema_target_proj(target_emb)
        return target_emb

    def forward(self, mel, text_tokens, mel_mask=None, text_mask=None, speaker_id=None):
        """
        Args:
            mel: [B, n_mels, T_mel]
            text_tokens: [B, T_text]
            mel_mask: [B, T_mel] bool
            text_mask: [B, T_text] bool
            speaker_id: [B] long tensor of speaker IDs (optional)
        Returns:
            loss_dict with all individual losses
        """
        # Get speaker conditioning
        spk_cond = self._get_speaker_cond(speaker_id)  # [B, d] or None

        # Encode text + FiLM speaker conditioning on text
        text_emb = self.text_encoder(text_tokens, text_mask)
        if spk_cond is not None:
            # FiLM on text embeddings (replaces simple addition)
            text_emb = self._apply_speaker_film(text_emb, spk_cond, self.latent_film)

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

        # Save raw mu before speaker conditioning — for speaker classifier
        mu_raw = mu

        # Multi-scale EMA targets (data2vec-style: average layer outputs)
        target_emb = self._compute_multiscale_targets(mel, mel_mask)

        # Apply FiLM speaker conditioning to targets and encoder output
        # This makes prediction loss require speaker knowledge
        target_emb = self._apply_speaker_film(target_emb, spk_cond, self.target_film)
        z = self._apply_speaker_film(z, spk_cond, self.latent_film)

        # Prepend learnable start embedding to input sequence
        start = self.start_emb.expand(B, -1, -1)  # [B, 1, d]
        input_emb = torch.cat([start, z[:, :-1]], dim=1)  # [B, T_down, d]

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

        # Adjust masks for downsampled sequence
        if mel_mask is not None:
            ds_mask = mel_mask[:, ::self.audio_encoder.downsample_factor][:, :T_down]
            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

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

        # ─── Loss 1: Prediction loss (MSE + cosine) against multi-scale EMA targets ───
        if loss_mask is not None:
            valid = (~loss_mask).unsqueeze(-1).float()  # [B, T_down, 1]
            n_valid = valid.sum() * predicted.shape[-1] + 1e-8

            # MSE component
            mse_loss = (((predicted - target_emb) ** 2) * valid).sum() / n_valid

            # Cosine component — preserves direction (speaker identity)
            pred_flat = (predicted * valid).reshape(-1, predicted.shape[-1])
            tgt_flat = (target_emb * valid).reshape(-1, target_emb.shape[-1])
            # Only compute on non-padding positions
            valid_flat = valid.squeeze(-1).reshape(-1) > 0
            if valid_flat.any():
                cosine_loss = 1.0 - F.cosine_similarity(
                    pred_flat[valid_flat], tgt_flat[valid_flat], dim=-1
                ).mean()
            else:
                cosine_loss = torch.tensor(0.0, device=mel.device)
        else:
            mse_loss = F.mse_loss(predicted, target_emb)
            cosine_loss = 1.0 - F.cosine_similarity(
                predicted.reshape(-1, predicted.shape[-1]),
                target_emb.reshape(-1, target_emb.shape[-1]),
                dim=-1,
            ).mean()

        prediction_loss = mse_loss + self.cosine_weight * cosine_loss


Now I see the full picture. The fix is simple — **predict in encoder z space, not EMA target space**. The predictor should be trained to predict the next encoder output (what it'll actually receive as input during AR), not a multi-scale EMA target in a different scale.


Edit

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

        # Save raw mu before speaker conditioning — for speaker classifier
        mu_raw = mu

        # Multi-scale EMA targets (data2vec-style: average layer outputs)
        target_emb = self._compute_multiscale_targets(mel, mel_mask)

        # Apply FiLM speaker conditioning to targets and encoder output
        # This makes prediction loss require speaker knowledge
        target_emb = self._apply_speaker_film(target_emb, spk_cond, self.target_film)
        z = self._apply_speaker_film(z, spk_cond, self.latent_film)

        # Prepend learnable start embedding to input sequence
        start = self.start_emb.expand(B, -1, -1)  # [B, 1, d]
        input_emb = torch.cat([start, z[:, :-1]], dim=1)  # [B, T_down, d]

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

        # Save raw mu before speaker conditioning — for speaker classifier
        mu_raw = mu

        # Apply FiLM speaker conditioning to encoder output
        z = self._apply_speaker_film(z, spk_cond, self.latent_film)

        # Prediction target = encoder z (FiLM'd), shifted by 1
        # CRITICAL: target must be in the SAME space as predictor input (encoder z),
        # not a different space (EMA targets). Otherwise AR inference diverges because
        # predicted outputs (in target space) get fed back as input (in encoder space).
        # We use EMA encoder to get stable targets, but keep them in encoder scale.
        with torch.no_grad():
            self.ema_audio_encoder.eval()
            _, mu_ema, _ = self.ema_audio_encoder(mel, mel_mask)
            self.ema_audio_encoder.train()
            target_emb = mu_ema
        # Apply same FiLM as encoder output so target space matches input space
        target_emb = self._apply_speaker_film(target_emb, spk_cond, self.target_film)

        # Prepend learnable start embedding to input sequence
        start = self.start_emb.expand(B, -1, -1)  # [B, 1, d]
        input_emb = torch.cat([start, z[:, :-1]], dim=1)  # [B, T_down, d]