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

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

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

def synth(model, text, steps=200, temp=0.0):
    tokens = list(text.encode('utf-8'))
    tt = torch.tensor(tokens, dtype=torch.long).unsqueeze(0).cuda()
    te = model.text_encoder(tt)
    cache = model.init_ar_cache()
    emb = model.start_emb
    embs = [emb]
    for s in range(steps):
        emb, cache = model.predict_next_cached(emb, te, s, cache)
        if temp > 0: emb = emb + torch.randn_like(emb) * temp
        embs.append(emb)
    mel = model.mel_decoder(torch.cat(embs, dim=1), None)
    return vocos.decode(mel.cpu()).squeeze().numpy()

def save(wav, path):
    mx = np.abs(wav).max()
    if mx > 0: wav = wav / mx * 0.95
    torchaudio.save(path, torch.from_numpy(wav).unsqueeze(0), 24000)
    print(f'  {path}: {len(wav)/24000:.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()
    _, mu, _ = model.audio_encoder(mel_orig)
    save(vocos.decode(model.mel_decoder(mu, None).cpu()).squeeze().numpy(), 'final_recon.wav')

    # Synth temp=0
    save(synth(model, 'नमस्कार, मैं भारत का प्रधानमंत्री हूं।', temp=0.0), 'final_synth1_t0.wav')
    save(synth(model, 'मेरे प्यारे देशवासियों, आप सभी को नमस्कार।', temp=0.0), 'final_synth2_t0.wav')
    save(synth(model, 'हम सब मिलकर भारत को आगे बढ़ाएंगे।', temp=0.0), 'final_synth3_t0.wav')

    # Synth temp=0.3 and 0.8
    save(synth(model, 'नमस्कार, मैं भारत का प्रधानमंत्री हूं।', temp=0.3), 'final_synth1_t03.wav')
    save(synth(model, 'नमस्कार, मैं भारत का प्रधानमंत्री हूं।', temp=0.8), 'final_synth1_t08.wav')

print('Done')
" 2>&1

Generate final samples from best checkpoint


Epoch 4, loss=2.2614
  final_recon.wav: 9.03s rms=0.0773
  final_synth1_t0.wav: 8.57s rms=0.0985
  final_synth2_t0.wav: 8.57s rms=0.1025
  final_synth3_t0.wav: 8.57s rms=0.0999
  final_synth1_t03.wav: 8.57s rms=0.1008
  final_synth1_t08.wav: 8.57s rms=0.1088
Done


Wait — "Epoch 4" on best.pt? The best loss was saved early. Let me use the final epoch 200 checkpoint instead:


Bash

ls /home/ubuntu/lewm-tts/output_v4_fix2/checkpoints/ && python3 -c "
import torch
for f in ['best.pt', 'epoch_0200.pt', 'epoch_0175.pt', 'epoch_0150.pt']:
    try:
        ckpt = torch.load(f'output_v4_fix2/checkpoints/{f}', map_location='cpu', weights_only=False)
        print(f'{f}: epoch {ckpt[\"epoch\"]+1}, loss={ckpt[\"loss\"]:.4f}')
    except: pass
"

Check all checkpoints


best.pt
epoch_0025.pt
epoch_0050.pt
epoch_0075.pt
epoch_0100.pt
epoch_0125.pt
epoch_0150.pt
epoch_0175.pt
epoch_0200.pt
best.pt: epoch 4, loss=2.2614
epoch_0200.pt: epoch 200, loss=2.6099
epoch_0175.pt: epoch 175, loss=2.5489
epoch_0150.pt: epoch 150, loss=2.4997


Interesting — best loss was epoch 4 (2.26) before scheduled sampling kicked in and made losses higher. The later epochs have higher total loss because the task is harder (predicting from own outputs), but that's what makes synthesis work. Let me use epoch 200:


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

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

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

def synth(model, text, steps=200, temp=0.0):
    tokens = list(text.encode('utf-8'))
    tt = torch.tensor(tokens, dtype=torch.long).unsqueeze(0).cuda()
    te = model.text_encoder(tt)
    cache = model.init_ar_cache()
    emb = model.start_emb
    embs = [emb]
    for s in range(steps):
        emb, cache = model.predict_next_cached(emb, te, s, cache)
        if temp > 0: emb = emb + torch.randn_like(emb) * temp
        embs.append(emb)
    mel = model.mel_decoder(torch.cat(embs, dim=1), None)
    return vocos.decode(mel.cpu()).squeeze().numpy()

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

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

    save(synth(model, 'नमस्कार, मैं भारत का प्रधानमंत्री हूं।', temp=0.0), 'final200_synth1_t0.wav')
    save(synth(model, 'मेरे प्यारे देशवासियों, आप सभी को नमस्कार।', temp=0.0), 'final200_synth2_t0.wav')
    save(synth(model, 'हम सब मिलकर भारत को आगे बढ़ाएंगे।', temp=0.0), 'final200_synth3_t0.wav')
    save(synth(model, 'नमस्कार, मैं भारत का प्रधानमंत्री हूं।', temp=0.3), 'final200_synth1_t03.wav')
    save(synth(model, 'नमस्कार, मैं भारत का प्रधानमंत्री हूं।', temp=0.8), 'final200_synth1_t08.wav')
print('Done')
" 2>&1

Generate samples from epoch 200


Epoch 200, loss=2.6099
  final200_recon.wav: 9.03s rms=0.0817
  final200_synth1_t0.wav: 8.57s rms=0.1562
  final200_synth2_t0.wav: 8.57s rms=0.1526
  final200_synth3_t0.wav: 8.57s rms=0.1509
  final200_synth1_t03.wav: 8.57s rms=0.0504
  final200_synth1_t08.wav: 8.57s rms=0.0586
Done


**Epoch 200 final samples ready.** RMS at temp=0 is 0.15 — strongest signal we've seen in synthesis. And temp=0 is clearly louder than temp=0.8 — confirming temperature is adding noise that hurts.

| File | Temp | RMS |
|---|---|---|
| `final200_recon.wav` | - | 0.082 |
| `final200_synth1_t0.wav` | 0.0 | **0.156** |
| `final200_synth2_t0.wav` | 0.0 | **0.153** |
| `final200_synth3_t0.wav` | 0.0 | **0.151** |
| `final200_synth1_t03.wav` | 0.3 | 0.050 |
| `final200_synth1_t08.wav` | 0.8 | 0.059 |

Listen to the temp=0 ones — those should be the cleanest. Tell me what you hear.


fucking noise ... check for any auido jera alrady reied by somebody


Agent

Research audio JEPA TTS implementations

Prompt:
Search the web thoroughly for any existing implementations or papers that use JEPA (Joint-Embedding Predictive Architecture) for TTS (text-to-speech) or audio generation. I need to find:

1. Any GitHub repos that implement JEPA for TTS or audio synthesis
2. Any papers that use JEPA specifically for generating audio (not just representation learning)
3. A-JEPA, Audio-JEPA, or any variant used for generation
4. How other JEPA audio models handle the autoregressive generation problem (exposure bias, embedding space mismatch)
5. Any TTS models that use JEPA-style latent prediction instead of traditional approaches

Search for:
- "JEPA TTS"
- "JEPA audio generation"
- "JEPA speech synthesis"
- "A-JEPA generation"
- "joint embedding predictive architecture speech"
- GitHub repos with JEPA + audio/TTS

Also look for how SoundStorm, VALL-E, or other modern TTS handle the similar problem of predicting in latent space and avoiding exposure bias during autoregressive generation.

Be thorough - check arxiv, GitHub, papers with code, blog posts.


Now I have comprehensive results. Here is my thorough synthesis:

---

## Research Findings: JEPA for TTS / Audio Generation

### 1. GitHub Repos with JEPA for Audio

**No JEPA-for-TTS repo exists.** Nobody has published a JEPA-based TTS system. What exists:

- **[SonyCSLParis/audio-representations](https://github.com/SonyCSLParis/audio-representations)** -- JEPA for general audio *representation learning* (ICASSP 2024 SASB workshop). Closest to what you're doing architecturally, but representation-only.
- **[labhamlet/wavjepa](https://github.com/labhamlet/wavjepa)** -- WavJEPA official code. Time-domain audio foundation model, representation learning only.
- **[SonyCSLParis/Stem-JEPA](https://github.com/SonyCSLParis/Stem-JEPA)** -- JEPA for musical stem compatibility (ISMIR 2024). Not generative.
- **[facebookresearch/eb_jepa](https://github.com/facebookresearch/eb_jepa)** -- Meta's educational JEPA library with image, video, and action-conditioned planning examples. Demonstrates JEPA for *planning* (generation of action sequences in latent space), which is conceptually closest to what you'd need for autoregressive latent generation.
- **[facebookresearch/ijepa](https://github.com/facebookresearch/ijepa)** and **[facebookresearch/jepa](https://github.com/facebookresearch/jepa)** (V-JEPA) -- Official Meta codebases for image/video JEPA.

**Your LeWM-TTS project appears to be the first attempt at JEPA-based TTS.** This is genuinely novel.

### 2. Papers Using JEPA for Audio (Representation, Not Generation)

All existing JEPA audio papers are for *representation learning*, not generation:

- **[A-JEPA](https://arxiv.org/abs/2311.15830)** (Nov 2023) -- Masked spectrogram prediction in latent space with curriculum masking (time-frequency aware). SOTA on audio/speech classification. No generation capability.
- **[Audio-JEPA](https://arxiv.org/abs/2507.02915)** (Tuncay et al., July 2025) -- ViT backbone predicting latent representations of masked spectrogram patches. Comparable to wav2vec 2.0 with 1/5th the data. Representation only.
- **[WavJEPA](https://arxiv.org/abs/2509.23238)** (Sep 2025) -- Waveform-domain JEPA. Outperforms time-domain audio foundation models. Has [code](https://github.com/labhamlet/wavjepa). Representation only.
- **[JEPA as a Neural Tokenizer](https://arxiv.org/abs/2512.07168)** (Dec 2025) -- **Most relevant to your work.** Two-stage: (1) JEPA+DAAM learns semantic features at 2.5 Hz via masked latent prediction, (2) FSQ tokenization + HiFi-GAN decoder for waveform reconstruction. Produces 47.5 tokens/sec, vocabulary of 16,384. Designed to be "LM-friendly" -- explicitly intended as a tokenizer for downstream LM-based TTS. But the paper itself does not build the TTS system on top.
- **[JEPA-v0](https://startpinch.com/research/en/jepa-encoder-translation/)** (Pinch Research, 2025-2026) -- JEPA encoder for speech-to-speech translation. Preserves speaker identity unlike ASR-MT-TTS cascades. Encoder only; no generative decoder published.

### 3. D-JEPA and Discrete JEPA -- Attempts at Making JEPA Generative

These are the closest to what you're trying to do (making JEPA generate, not just represent):

- **[D-JEPA](https://d-jepa.github.io/)** -- Reinterprets JEPA as masked image modeling + next-token prediction, then uses diffusion loss to model per-token probability distribution for image generation. **This is the key bridge paper** -- it shows how to go from JEPA representations to generation via autoregressive + diffusion.
- **[Discrete JEPA](https://arxiv.org/pdf/2506.14373)** (June 2025) -- Learns discrete semantic tokens via latent predictive coding. Key finding: **operating in discrete index space prevents error accumulation** over long horizons, unlike continuous prediction which suffers compounding errors. Trains autoregressive world models using standard ViT. Directly relevant to your exposure bias concern.
- **[JEPA-Reasoner](https://arxiv.org/abs/2512.19171)** (Dec 2025) -- **Most architecturally relevant to your problem.** Autoregressively reasons in continuous latent space, with a separate "Talker" decoder for token generation. Key insight: **decoupling latent-space prediction from output generation** contains errors -- token-level failures cannot propagate into the latent reasoning chain. Uses normalized latent space. This is essentially the same architecture pattern as your JEPA encoder + separate decoder.

### 4. How the Exposure Bias / Embedding Space Mismatch Problem Is Handled

This is the core challenge for your architecture. Here's how different systems address it:

**Direct approaches to exposure bias in latent-prediction TTS:**

- **[Bridging the Gap Between Training and Inference in LM-Based TTS](https://arxiv.org/abs/2509.17021)** (Sep 2025) -- Hybrid training combining teacher forcing with "free running" (using self-generated tokens during training). Uses an EOS prediction mechanism to adaptively control how much the model sees its own outputs. Directly applicable to your problem.

- **[Self-Forcing](https://arxiv.org/abs/2506.08009)** (June 2025, video diffusion but principle applies) -- During training, conditions on self-generated outputs rather than ground truth, with holistic sequence-level loss. Uses stochastic gradient truncation for efficiency. The **Self-Forcing++** extension addresses semantic drift over extended generation.

- **[KALL-E](https://arxiv.org/abs/2412.16846)** (Dec 2024) -- Predicts the **next distribution** (not next token) of continuous speech frames using Flow-VAE + KL divergence loss. Avoids quantization entirely. The distribution prediction sidesteps the discrete mismatch problem. Outperforms VALL-E, NaturalSpeech 2, CosyVoice.

- **[CLEAR](https://arxiv.org/abs/2508.19098)** (Aug 2025) -- Continuous Latent Autoregressive TTS. Maps waveforms to compact continuous latents via enhanced VAE, then uses MLP-based rectified flow head per hidden state. Single-stage training. 96ms first-frame delay for streaming. The continuous latent representation avoids the discrete token mismatch.

**How SoundStorm and VALL-E handle it:**

- **[SoundStorm](https://arxiv.org/abs/2305.09636)** -- Sidesteps exposure bias entirely by using **non-autoregressive parallel decoding** with confidence-based iterative refinement (MaskGIT-style). Takes semantic tokens from AudioLM, generates codec tokens in parallel. 2 orders of magnitude faster than autoregressive. The key insight: avoid autoregression altogether for the acoustic token generation stage.

- **VALL-E** -- Hybrid: first RVQ level is autoregressive, subsequent levels are non-autoregressive (greedy level-by-level). This limits exposure bias to only the first level.

### 5. Key Takeaways for Your LeWM-TTS Architecture

Based on this research, here is what matters for your project:

**You are in genuinely uncharted territory.** No one has published a JEPA-based TTS system. The closest prior work is:
- JEPA as Neural Tokenizer (provides the encoder/tokenizer half)
- JEPA-Reasoner (provides the autoregressive-in-latent-space + separate decoder pattern)
- D-JEPA (shows how to bridge JEPA to generation)

**For the exposure bias / embedding mismatch problem, the most promising strategies from the literature are:**

1. **Hybrid teacher-forcing / free-running training** (from the "Bridging the Gap" paper) -- gradually introduce self-generated latents during training
2. **Self-Forcing** -- train directly on self-generated rollouts with sequence-level loss
3. **Predict distributions, not points** (KALL-E approach) -- predict a distribution over the next latent rather than a single embedding, using KL divergence loss
4. **Discrete JEPA's finding** -- discretizing the latent space (via FSQ or VQ) may inherently limit error accumulation vs continuous prediction
5. **Non-autoregressive refinement** (SoundStorm approach) -- use AR only for coarse semantic tokens, then parallel decode the fine acoustic tokens

Sources:
- [A-JEPA: Joint-Embedding Predictive Architecture Can Listen](https://arxiv.org/abs/2311.15830)
- [Audio-JEPA: Joint-Embedding Predictive Architecture for Audio Representation Learning](https://arxiv.org/abs/2507.02915)
- [WavJEPA: Semantic learning unlocks robust audio foundation models for raw waveforms](https://arxiv.org/abs/2509.23238)
- [JEPA as a Neural Tokenizer (arxiv 2512.07168)](https://arxiv.org/abs/2512.07168)
- [JEPA-v0: Pinch Research](https://startpinch.com/research/en/jepa-encoder-translation/)
- [D-JEPA: Denoising with JEPA](https://d-jepa.github.io/)
- [Discrete JEPA (arxiv 2506.14373)](https://arxiv.org/pdf/2506.14373)
- [JEPA-Reasoner: Decoupling Latent Reasoning from Token Generation](https://arxiv.org/abs/2512.19171)
- [Bridging the Gap Between Training and Inference in LM-Based TTS](https://arxiv.org/abs/2509.17021)
- [Self-Forcing: Bridging the Train-Test Gap in Autoregressive Video Diffusion](https://arxiv.org/abs/2506.08009)
- [KALL-E: Autoregressive Speech Synthesis with Next-Distribution Prediction](https://arxiv.org/abs/2412.16846)
- [CLEAR: Continuous Latent Autoregressive Modeling for Speech Synthesis](https://arxiv.org/abs/2508.19098)
- [SoundStorm: Efficient Parallel Audio Generation](https://arxiv.org/abs/2305.09636)
- [SonyCSLParis/audio-representations (GitHub)](https://github.com/SonyCSLParis/audio-representations)
- [labhamlet/wavjepa (GitHub)](https://github.com/labhamlet/wavjepa)
- [facebookresearch/eb_jepa (GitHub)](https://github.com/facebookresearch/eb_jepa)
- [JEPA-Reasoner OpenReview](https://openreview.net/forum?id=7ruA2rXG42)
- [Discrete JEPA OpenReview](https://openreview.net/pdf?id=91kfdrEafN)
- [JEPA as Neural Tokenizer UniReps Blog](https://unireps.org/blog/2025/jepa-daam_fsqvae/)
- [Exploring JEPA for Real-Time Speech Translation (Hacker News)](https://news.ycombinator.com/item?id=47332883)

agentId: a1b5cfd74c8dcb644 (use SendMessage with to: 'a1b5cfd74c8dcb644' to continue this agent)
<usage>total_tokens: 35791
tool_uses: 21
duration_ms: 141950</usage>

Prompt:
Search the web thoroughly for any existing implementations or papers that use JEPA (Joint-Embedding Predictive Architecture) for TTS (text-to-speech) or audio generation. I need to find:

1. Any GitHub repos that implement JEPA for TTS or audio synthesis
2. Any papers that use JEPA specifically for generating audio (not just representation learning)
3. A-JEPA, Audio-JEPA, or any variant used for generation
4. How other JEPA audio models handle the autoregressive generation problem (exposure bias, embedding space mismatch)
5. Any TTS models that use JEPA-style latent prediction instead of traditional approaches

Search for:
- "JEPA TTS"
- "JEPA audio generation"
- "JEPA speech synthesis"
- "A-JEPA generation"
- "joint embedding predictive architecture speech"
- GitHub repos with JEPA + audio/TTS

Also look for how SoundStorm, VALL-E, or other modern TTS handle the similar problem of predicting in latent space and avoiding exposure bias during autoregressive generation.

Be thorough - check arxiv, GitHub, papers with code, blog posts.


OK here's the reality check:

## Nobody has done JEPA for TTS. You're the first attempt.

All existing JEPA audio work (A-JEPA, Audio-JEPA, WavJEPA, Stem-JEPA) is **representation learning only** — encode audio, classify it, done. None of them generate audio.

The reason our AR synthesis produces noise is the **fundamental problem** of continuous autoregressive prediction — errors compound every step. This isn't a bug we can fix with more training. It's a known unsolved challenge.

## What actually works (from research)

Three approaches that handle this:

**1. Discretize the latent space (most promising for you)**
- Paper: "JEPA as Neural Tokenizer" + "Discrete JEPA"
- Key finding: **discrete tokens prevent error accumulation** — a wrong token is still a valid token, not a drifting float
- You'd VQ/FSQ quantize your JEPA encoder output into discrete codes, then AR predict tokens instead of continuous embeddings
- This is basically what VALL-E, SoundStorm do — they work

**2. Predict distributions, not points (KALL-E)**
- Instead of predicting one embedding, predict a distribution (mean + variance)
- Use flow matching or diffusion to sample from it
- Outperforms VALL-E

**3. Non-autoregressive (SoundStorm)**
- Skip AR entirely — parallel decode with iterative refinement
- 100x faster, no exposure bias at all
- But needs semantic tokens as input (separate model)

## My recommendation

**Option 1 — add VQ to your JEPA encoder.** Your architecture stays mostly the same. Add a VQ bottleneck after the audio encoder, predict discrete token IDs instead of continuous vectors. The mel decoder learns from discrete codes. This is the smallest change that would actually make synthesis work.

Want me to implement the VQ approach? It's a direct modification of your existing model — add a codebook, change the prediction head to classification over codebook entries instead of regression.


yeah try


1	"""
2	LeWM TTS Model — JEPA-based Text-to-Speech
3
4	Architecture:
5	  - TextEncoder: Character-level transformer for Hindi/Devanagari text
6	  - AudioEncoder: 1D CNN + Transformer that encodes mel spectrograms to embeddings
7	    - Returns intermediate layer outputs for multi-scale targets
8	  - JEPAPredictor: Transformer that predicts next audio embedding given text + audio context
9	  - SpeakerConditioner: FiLM-based speaker conditioning (scale + shift)
10	  - MelDecoder: FiLM-conditioned decoder with speaker-aware residual blocks
11	  - Losses: prediction (MSE + cosine), KL, reconstruction, spectral, speaker consistency
12	"""
13
14	import math
15	import copy

"""
LeWM TTS Model — JEPA-based Text-to-Speech

Architecture:
  - TextEncoder: Character-level transformer for Hindi/Devanagari text
  - AudioEncoder: 1D CNN + Transformer that encodes mel spectrograms to embeddings
    - Returns intermediate layer outputs for multi-scale targets
  - JEPAPredictor: Transformer that predicts next audio embedding given text + audio context
  - SpeakerConditioner: FiLM-based speaker conditioning (scale + shift)
  - MelDecoder: FiLM-conditioned decoder with speaker-aware residual blocks
  - Losses: prediction (MSE + cosine), KL, reconstruction, spectral, speaker consistency
"""

import math
import copy


Let me read the key sections I need to understand for the VQ integration:


95	        x = self.proj(x)
96	        return x
97
98
99	# ─── Audio Encoder ──────────────────────────────────────────────────────────
100
101	class AudioEncoder(nn.Module):
102	    """
103	    Encodes mel spectrograms into a sequence of audio embeddings.
104	    Uses 1D CNN for downsampling + Transformer for contextualization.
105	    Returns intermediate layer outputs for multi-scale EMA targets.
106	    """
107
108	    def __init__(self, n_mels=100, d_model=256, nhead=4, num_layers=4,
109	                 downsample_factor=4, dropout=0.1):
110	        super().__init__()
111	        self.d_model = d_model
112	        self.downsample_factor = downsample_factor
113	        self.num_layers = num_layers
114
115	        # CNN downsampling: reduces temporal resolution by downsample_factor
116	        if downsample_factor == 2:
117	            self.conv_pre = nn.Sequential(
118	                nn.Conv1d(n_mels, d_model, kernel_size=7, padding=3),
119	                nn.GELU(),
120	                nn.Conv1d(d_model, d_model, kernel_size=4, stride=2, padding=1),
121	                nn.GELU(),
122	            )  # Total stride = 2
123	        else:
124	            self.conv_pre = nn.Sequential(
125	                nn.Conv1d(n_mels, d_model, kernel_size=7, padding=3),
126	                nn.GELU(),
127	                nn.Conv1d(d_model, d_model, kernel_size=4, stride=2, padding=1),
128	                nn.GELU(),
129	                nn.Conv1d(d_model, d_model, kernel_size=4, stride=2, padding=1),
130	                nn.GELU(),
131	            )  # Total stride = 4
132
133	        self.pos_embed = SinusoidalPositionalEncoding(d_model, max_len=4096)
134	        self.dropout = nn.Dropout(dropout)
135
136	        # Build layers individually so we can capture intermediate outputs
137	        self.layers = nn.ModuleList([
138	            nn.TransformerEncoderLayer(
139	                d_model=d_model, nhead=nhead,
140	                dim_feedforward=d_model * 4,
141	                dropout=dropout, batch_first=True,
142	                activation="gelu",
143	            )
144	            for _ in range(num_layers)

        x = self.proj(x)
        return x

# ─── Audio Encoder ──────────────────────────────────────────────────────────

class AudioEncoder(nn.Module):
    """
    Encodes mel spectrograms into a sequence of audio embeddings.
    Uses 1D CNN for downsampling + Transformer for contextualization.
    Returns intermediate layer outputs for multi-scale EMA targets.
    """

    def __init__(self, n_mels=100, d_model=256, nhead=4, num_layers=4,
                 downsample_factor=4, dropout=0.1):
        super().__init__()
        self.d_model = d_model
        self.downsample_factor = downsample_factor
        self.num_layers = num_layers

        # CNN downsampling: reduces temporal resolution by downsample_factor
        if downsample_factor == 2:
            self.conv_pre = nn.Sequential(
                nn.Conv1d(n_mels, d_model, kernel_size=7, padding=3),
                nn.GELU(),
                nn.Conv1d(d_model, d_model, kernel_size=4, stride=2, padding=1),
                nn.GELU(),
            )  # Total stride = 2
        else:
            self.conv_pre = nn.Sequential(
                nn.Conv1d(n_mels, d_model, kernel_size=7, padding=3),
                nn.GELU(),
                nn.Conv1d(d_model, d_model, kernel_size=4, stride=2, padding=1),
                nn.GELU(),
                nn.Conv1d(d_model, d_model, kernel_size=4, stride=2, padding=1),
                nn.GELU(),
            )  # Total stride = 4

        self.pos_embed = SinusoidalPositionalEncoding(d_model, max_len=4096)
        self.dropout = nn.Dropout(dropout)

        # Build layers individually so we can capture intermediate outputs
        self.layers = nn.ModuleList([
            nn.TransformerEncoderLayer(
                d_model=d_model, nhead=nhead,
                dim_feedforward=d_model * 4,
                dropout=dropout, batch_first=True,
                activation="gelu",
            )
            for _ in range(num_layers)


I'll create a clean new model file — `model_vq.py` — instead of further mutating the existing one. This keeps the VQ approach separate and testable.


Write

"""
LeWM TTS v6 — JEPA + VQ: Discrete token prediction for stable AR synthesis.

Key insight: continuous AR prediction accumulates errors every step → noise.
Discrete tokens can't drift — a predicted token is always a valid codebook entry.

Architecture:
  - TextEncoder: byte-level transformer (same as before)
  - AudioEncoder: CNN + Transformer → continuous embeddings (same as before)
  - VectorQuantizer: continuous → discrete codebook indices + quantized embeddings
  - JEPAPredictor: predicts next codebook INDEX (classification, not regression)
  - MelDecoder: codebook embeddings → mel spectrogram
  - Vocos: mel → waveform (external, frozen)
"""

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

# ─── Vector Quantizer ─────────────────────────────────────────────────────────

class VectorQuantizer(nn.Module):
    """
    Discretizes continuous embeddings into codebook indices.
    Uses EMA codebook updates (no codebook loss needed in optimizer).
    Straight-through estimator for gradients.
    """

    def __init__(self, d_model=256, n_codes=1024, commitment_weight=0.25, ema_decay=0.99):
        super().__init__()
        self.d_model = d_model
        self.n_codes = n_codes
        self.commitment_weight = commitment_weight
        self.ema_decay = ema_decay

        # Codebook: [n_codes, d_model]
        self.register_buffer("codebook", torch.randn(n_codes, d_model))
        self.register_buffer("ema_count", torch.ones(n_codes))
        self.register_buffer("ema_sum", self.codebook.clone())
        self._initialized = False

    def _init_codebook(self, flat_z):
        """Initialize codebook from first batch using k-means++ style."""
        if self._initialized:
            return
        n = min(flat_z.shape[0], self.n_codes)
        indices = torch.randperm(flat_z.shape[0], device=flat_z.device)[:n]
        self.codebook[:n] = flat_z[indices].detach()
        self.ema_sum.copy_(self.codebook)
        self._initialized = True

    def forward(self, z):
        """
        Args:
            z: [B, T, d_model] continuous embeddings
        Returns:
            z_q: [B, T, d_model] quantized embeddings (straight-through)
            indices: [B, T] codebook indices
            commit_loss: scalar commitment loss
        """
        B, T, D = z.shape
        flat_z = z.reshape(-1, D)  # [B*T, D]

        # Initialize codebook from first batch
        if self.training and not self._initialized:
            self._init_codebook(flat_z)

        # Find nearest codebook entry: ||z - e||^2 = ||z||^2 - 2<z,e> + ||e||^2
        dists = (
            flat_z.pow(2).sum(dim=1, keepdim=True)
            - 2 * flat_z @ self.codebook.t()
            + self.codebook.pow(2).sum(dim=1, keepdim=True).t()
        )  # [B*T, n_codes]

        indices = dists.argmin(dim=1)  # [B*T]
        z_q = self.codebook[indices]   # [B*T, D]

        # EMA codebook update (only during training)
        if self.training:
            with torch.no_grad():
                one_hot = F.one_hot(indices, self.n_codes).float()  # [B*T, n_codes]
                counts = one_hot.sum(dim=0)  # [n_codes]
                sums = one_hot.t() @ flat_z  # [n_codes, D]

                self.ema_count.mul_(self.ema_decay).add_(counts, alpha=1 - self.ema_decay)
                self.ema_sum.mul_(self.ema_decay).add_(sums, alpha=1 - self.ema_decay)

                # Laplace smoothing to avoid dead codes
                n = self.ema_count.sum()
                count_smooth = (self.ema_count + 1e-5) / (n + self.n_codes * 1e-5) * n
                self.codebook.copy_(self.ema_sum / count_smooth.unsqueeze(1))

        # Commitment loss: push encoder output toward codebook
        commit_loss = F.mse_loss(flat_z, z_q.detach())

        # Straight-through estimator: z_q gets gradients of z
        z_q = z + (z_q - z).detach()

        z_q = z_q.reshape(B, T, D)
        indices = indices.reshape(B, T)

        return z_q, indices, self.commitment_weight * commit_loss

# ─── Shared components (same as model.py) ─────────────────────────────────────

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

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)

class AudioEncoder(nn.Module):
    def __init__(self, n_mels=100, d_model=256, nhead=4, num_layers=4,
                 downsample_factor=4, dropout=0.1):
        super().__init__()
        self.d_model = d_model
        self.downsample_factor = downsample_factor

        if downsample_factor == 2:
            self.conv_pre = nn.Sequential(
                nn.Conv1d(n_mels, d_model, kernel_size=7, padding=3), nn.GELU(),
                nn.Conv1d(d_model, d_model, kernel_size=4, stride=2, padding=1), nn.GELU(),
            )
        else:
            self.conv_pre = nn.Sequential(
                nn.Conv1d(n_mels, d_model, kernel_size=7, padding=3), nn.GELU(),
                nn.Conv1d(d_model, d_model, kernel_size=4, stride=2, padding=1), nn.GELU(),
                nn.Conv1d(d_model, d_model, kernel_size=4, stride=2, padding=1), nn.GELU(),
            )

        self.pos_embed = SinusoidalPositionalEncoding(d_model, max_len=4096)
        self.dropout = nn.Dropout(dropout)
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=d_model, nhead=nhead, dim_feedforward=d_model * 4,
            dropout=dropout, batch_first=True, activation="gelu",
        )
        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
        self.proj = nn.Linear(d_model, d_model)

    def forward(self, mel, mel_mask=None):
        x = self.conv_pre(mel)
        x = x.transpose(1, 2)
        if mel_mask is not None:
            T_down = x.shape[1]
            mel_mask = mel_mask[:, ::self.downsample_factor][:, :T_down]
        x = self.pos_embed(x)
        x = self.dropout(x)
        x = self.transformer(x, src_key_padding_mask=mel_mask)
        return self.proj(x)  # [B, T_down, d_model]

# ─── Predictor (now predicts codebook logits) ─────────────────────────────────

class TokenPredictor(nn.Module):
    """
    Causal transformer that predicts next CODEBOOK INDEX.
    Input: quantized embeddings (codebook lookups)
    Output: logits over codebook entries [B, T, n_codes]
    """

    def __init__(self, d_model=256, nhead=4, num_layers=6, n_codes=1024, dropout=0.1):
        super().__init__()
        self.d_model = d_model
        self.n_codes = n_codes

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

        # Output head: predict codebook index (classification, not regression!)
        self.output_head = nn.Sequential(
            nn.Linear(d_model, d_model),
            nn.GELU(),
            nn.Linear(d_model, n_codes),
        )

    def forward(self, audio_emb, text_emb, audio_mask=None, text_mask=None):
        x = self.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)

        x = self.transformer(
            x, text_emb,
            tgt_mask=causal_mask,
            tgt_key_padding_mask=audio_mask,
            memory_key_padding_mask=text_mask,
        )

        return self.output_head(x)  # [B, T, n_codes]

    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):
        """One-step cached inference. Returns logits [B, 1, n_codes] and updated cache."""
        x = self.input_proj(new_emb)
        x = x + self.pos_embed.pe[:, step_idx:step_idx + 1]

        for i, layer in enumerate(self.transformer.layers):
            # Self-attention with cache
            sa_out, cache['self_k'][i], cache['self_v'][i] = self._cached_mha(
                layer.self_attn, x, x, cache['self_k'][i], cache['self_v'][i]
            )
            x = layer.norm1(x + layer.dropout1(sa_out))

            # Cross-attention with cached text
            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], text_mask
            )
            x = layer.norm2(x + layer.dropout2(ca_out))

            # FFN
            ff_out = layer.linear2(layer.dropout(layer.activation(layer.linear1(x))))
            x = layer.norm3(x + layer.dropout3(ff_out))

        return self.output_head(x), cache  # [B, 1, n_codes]

    def _cached_mha(self, mha, q, kv, cache_k, cache_v):
        d, nhead = mha.embed_dim, mha.num_heads
        head_dim = d // nhead
        B = q.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_proj = F.linear(q, Wq, bq)
        k_new = F.linear(kv, Wk, bk)
        v_new = F.linear(kv, Wv, bv)

        k = torch.cat([cache_k, k_new], dim=1) if cache_k is not None and cache_k.shape[1] > 0 else k_new
        v = torch.cat([cache_v, v_new], dim=1) if cache_v is not None and cache_v.shape[1] > 0 else v_new

        q_proj = q_proj.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_proj, k_mh.transpose(-2, -1)) / (head_dim ** 0.5)
        out = torch.matmul(F.softmax(attn, dim=-1), v_mh)
        out = out.transpose(1, 2).contiguous().view(B, 1, d)
        return F.linear(out, mha.out_proj.weight, mha.out_proj.bias), k, v

    def _cached_cross_attn(self, mha, q, memory, cache_k, cache_v, memory_mask=None):
        d, nhead = mha.embed_dim, mha.num_heads
        head_dim = d // nhead
        B = q.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_proj = F.linear(q, 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_proj = q_proj.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_proj, 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'))
        out = torch.matmul(F.softmax(attn, dim=-1), v_mh)
        out = out.transpose(1, 2).contiguous().view(B, 1, d)
        return F.linear(out, mha.out_proj.weight, mha.out_proj.bias), k, v

# ─── Mel Decoder ──────────────────────────────────────────────────────────────

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 MelDecoder(nn.Module):
    def __init__(self, d_model=256, n_mels=100, upsample_factor=4):
        super().__init__()
        hidden = d_model * 2

        self.proj = nn.Sequential(nn.Linear(d_model, hidden), nn.GELU())

        if upsample_factor == 2:
            self.upsample = nn.Sequential(
                nn.ConvTranspose1d(hidden, hidden, kernel_size=4, stride=2, padding=1), nn.GELU(),
                ResConvBlock(hidden, kernel_size=3),
            )
        else:
            self.upsample = nn.Sequential(
                nn.ConvTranspose1d(hidden, hidden, kernel_size=4, stride=2, padding=1), nn.GELU(),
                ResConvBlock(hidden, kernel_size=3),
                nn.ConvTranspose1d(hidden, hidden, kernel_size=4, stride=2, padding=1), nn.GELU(),
                ResConvBlock(hidden, kernel_size=3),
            )

        self.out = nn.Sequential(
            nn.Conv1d(hidden, hidden, kernel_size=7, padding=3), nn.GELU(),
            ResConvBlock(hidden, kernel_size=5),
            nn.Conv1d(hidden, n_mels, kernel_size=7, padding=3),
        )

    def forward(self, z):
        x = self.proj(z).transpose(1, 2)
        x = self.upsample(x)
        return self.out(x)

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

class MultiResolutionSpectralLoss(nn.Module):
    def __init__(self, resolutions=((64, 16, 64), (128, 32, 128), (256, 64, 256))):
        super().__init__()
        self.resolutions = resolutions

    def _stft_loss(self, pred, target, n_fft, hop_length, win_length):
        window = torch.hann_window(win_length, device=pred.device)
        B, M, T = pred.shape
        pred_stft = torch.stft(pred.reshape(B*M, T), n_fft, hop_length, win_length, window=window, return_complex=True)
        target_stft = torch.stft(target.reshape(B*M, T), n_fft, hop_length, win_length, window=window, return_complex=True)
        sc = torch.norm(target_stft.abs() - pred_stft.abs(), p="fro") / (torch.norm(target_stft.abs(), p="fro") + 1e-8)
        log_mag = F.l1_loss(torch.log(pred_stft.abs() + 1e-7), torch.log(target_stft.abs() + 1e-7))
        return sc + log_mag

    def forward(self, pred, target):
        loss = torch.tensor(0.0, device=pred.device)
        count = 0
        for n_fft, hop, win in self.resolutions:
            if pred.shape[2] >= n_fft:
                loss = loss + self._stft_loss(pred, target, n_fft, hop, win)
                count += 1
        return loss / max(count, 1)

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

class LeWMTTSvq(nn.Module):
    """
    JEPA + VQ TTS model.

    Training flow:
      mel → AudioEncoder → continuous z → VQ → discrete indices + quantized z_q
      [start, z_q[:-1]] + text → TokenPredictor → logits → cross_entropy(logits, indices)
      z_q → MelDecoder → mel_recon → L1 + spectral loss

    Inference flow:
      text → TextEncoder → text_emb
      start_emb → Predictor → logits → argmax → codebook lookup → next input → repeat
      all codebook embeddings → MelDecoder → mel → Vocos → waveform

    No continuous drift possible — every predicted embedding is a valid codebook entry.
    """

    def __init__(self, config):
        super().__init__()
        d_model = config.get("d_model", 256)
        nhead = config.get("nhead", 4)
        n_mels = config.get("n_mels", 100)
        text_vocab = config.get("text_vocab_size", 256)
        text_layers = config.get("text_encoder_layers", 4)
        audio_layers = config.get("audio_encoder_layers", 4)
        predictor_layers = config.get("predictor_layers", 6)
        dropout = config.get("dropout", 0.1)
        n_codes = config.get("n_codes", 1024)
        downsample_factor = config.get("downsample_factor", 4)

        self.text_encoder = TextEncoder(
            vocab_size=text_vocab, d_model=d_model, nhead=nhead,
            num_layers=text_layers, dropout=dropout,
        )
        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.vq = VectorQuantizer(
            d_model=d_model, n_codes=n_codes,
            commitment_weight=config.get("commitment_weight", 0.25),
        )
        self.predictor = TokenPredictor(
            d_model=d_model, nhead=nhead,
            num_layers=predictor_layers, n_codes=n_codes, dropout=dropout,
        )
        self.mel_decoder = MelDecoder(
            d_model=d_model, n_mels=n_mels, upsample_factor=downsample_factor,
        )

        # 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", 1.0)
        self.recon_weight = config.get("recon_weight", 1.0)
        self.spectral_weight = config.get("spectral_weight", 0.5)

        # Spectral loss
        self.spectral_loss_fn = MultiResolutionSpectralLoss()

        # Label smoothing for token prediction (prevents overconfidence)
        self.label_smoothing = config.get("label_smoothing", 0.1)

        self.config = config
        self.downsample_factor = downsample_factor

    def forward(self, mel, text_tokens, mel_mask=None, text_mask=None, **kwargs):
        B = mel.shape[0]

        # Encode text
        text_emb = self.text_encoder(text_tokens, text_mask)

        # Encode audio → continuous → quantize to discrete
        z_continuous = self.audio_encoder(mel, mel_mask)  # [B, T_down, d]
        z_q, indices, commit_loss = self.vq(z_continuous)  # [B, T_down, d], [B, T_down]

        T_down = z_q.shape[1]

        # Predictor input: [start, z_q[0], z_q[1], ..., z_q[T-2]]
        # Predictor target: [idx[0], idx[1], ..., idx[T-1]]
        start = self.start_emb.expand(B, -1, -1)
        input_emb = torch.cat([start, z_q[:, :-1]], dim=1)  # [B, T_down, d]

        # Masks
        if mel_mask is not None:
            ds_mask = mel_mask[:, ::self.downsample_factor][:, :T_down]
            start_mask = torch.zeros(B, 1, dtype=torch.bool, device=mel.device)
            pred_mask = torch.cat([start_mask, ds_mask[:, :-1]], dim=1)
            loss_mask = ds_mask
        else:
            pred_mask = None
            loss_mask = None

        # Predict next token logits
        logits = self.predictor(input_emb, text_emb, pred_mask, text_mask)  # [B, T_down, n_codes]

        # ─── Loss 1: Token prediction (cross-entropy with label smoothing) ───
        if loss_mask is not None:
            # Flatten valid positions only
            valid = ~loss_mask  # [B, T_down]
            logits_flat = logits[valid]  # [N_valid, n_codes]
            targets_flat = indices[valid]  # [N_valid]
            if logits_flat.shape[0] > 0:
                token_loss = F.cross_entropy(logits_flat, targets_flat,
                                              label_smoothing=self.label_smoothing)
            else:
                token_loss = torch.tensor(0.0, device=mel.device)
        else:
            token_loss = F.cross_entropy(
                logits.reshape(-1, logits.shape[-1]), indices.reshape(-1),
                label_smoothing=self.label_smoothing,
            )

        # ─── Loss 2: VQ commitment loss ───
        # (already computed in VQ forward)

        # ─── Loss 3: Mel reconstruction from quantized embeddings ───
        mel_recon = self.mel_decoder(z_q)
        T_mel = mel.shape[2]
        T_recon = mel_recon.shape[2]
        T_min = min(T_mel, T_recon)
        mel = mel[:, :, :T_min]
        mel_recon = mel_recon[:, :, :T_min]
        if mel_mask is not None:
            mel_mask = mel_mask[:, :T_min]

        if mel_mask is not None:
            valid_mel = (~mel_mask).unsqueeze(1).float()
            recon_loss = F.l1_loss(
                mel_recon * valid_mel, mel * valid_mel, reduction="sum"
            ) / (valid_mel.sum() * mel.shape[1] + 1e-8)
        else:
            recon_loss = F.l1_loss(mel_recon, mel)

        # ─── Loss 4: Spectral loss ───
        if mel_mask is not None:
            spectral_loss = self.spectral_loss_fn(mel_recon * valid_mel, mel * valid_mel)
        else:
            spectral_loss = self.spectral_loss_fn(mel_recon, mel)

        # ─── Token accuracy (for monitoring) ───
        with torch.no_grad():
            if loss_mask is not None and logits_flat.shape[0] > 0:
                token_acc = (logits_flat.argmax(dim=-1) == targets_flat).float().mean()
            else:
                token_acc = (logits.argmax(dim=-1) == indices).float().mean()

        total_loss = (self.pred_weight * token_loss
                      + commit_loss
                      + self.recon_weight * recon_loss
                      + self.spectral_weight * spectral_loss)

        return {
            "total_loss": total_loss,
            "token_loss": token_loss,
            "commit_loss": commit_loss,
            "recon_loss": recon_loss,
            "spectral_loss": spectral_loss,
            "token_accuracy": token_acc,
        }

    # ─── Inference ────────────────────────────────────────────────────────────

    def encode_to_tokens(self, mel):
        """Encode mel → discrete token indices."""
        z = self.audio_encoder(mel)
        _, indices, _ = self.vq(z)
        return indices

    @torch.no_grad()
    def synthesize_tokens(self, text_tokens, max_steps=300, temperature=1.0,
                          top_k=50, text_mask=None):
        """AR token generation: text → token indices."""
        text_emb = self.text_encoder(text_tokens, text_mask)

        cache = self.predictor.init_cache(len(self.predictor.transformer.layers))
        emb = self.start_emb
        all_indices = []

        for step in range(max_steps):
            logits, cache = self.predictor.inference_step(
                emb, text_emb, step, cache, text_mask
            )  # [B, 1, n_codes]

            logits = logits.squeeze(1)  # [B, n_codes]

            # Temperature + top-k sampling
            logits = logits / max(temperature, 1e-8)
            if top_k > 0:
                topk_vals, _ = logits.topk(top_k, dim=-1)
                logits[logits < topk_vals[:, -1:]] = float('-inf')

            probs = F.softmax(logits, dim=-1)
            idx = torch.multinomial(probs, 1)  # [B, 1]
            all_indices.append(idx)

            # Lookup codebook for next input
            emb = self.vq.codebook[idx]  # [B, 1, d]

        return torch.cat(all_indices, dim=1)  # [B, max_steps]

    @torch.no_grad()
    def tokens_to_mel(self, indices):
        """Convert token indices → mel spectrogram."""
        z_q = self.vq.codebook[indices]  # [B, T, d]
        return self.mel_decoder(z_q)  # [B, n_mels, T_up]

    @torch.no_grad()
    def synthesize_mel(self, text_tokens, max_steps=300, temperature=1.0,
                       top_k=50, text_mask=None):
        """Full pipeline: text → tokens → mel."""
        indices = self.synthesize_tokens(
            text_tokens, max_steps, temperature, top_k, text_mask
        )
        return self.tokens_to_mel(indices), indices

def count_parameters(model):
    return sum(p.numel() for p in model.parameters() if p.requires_grad)

def build_model_vq(config=None):
    if config is None:
        config = {
            "d_model": 256,
            "nhead": 4,
            "n_mels": 100,
            "text_vocab_size": 256,
            "text_encoder_layers": 4,
            "audio_encoder_layers": 4,
            "predictor_layers": 6,
            "n_codes": 1024,
            "dropout": 0.1,
            "pred_weight": 1.0,
            "recon_weight": 1.0,
            "spectral_weight": 0.5,
            "commitment_weight": 0.25,
            "label_smoothing": 0.1,
            "downsample_factor": 4,
        }
    model = LeWMTTSvq(config)
    print(f"LeWM TTS VQ model: {count_parameters(model)/1e6:.2f}M parameters")
    return model, config


File created successfully at: /home/ubuntu/lewm-tts/model_vq.py

"""
LeWM TTS v6 — JEPA + VQ: Discrete token prediction for stable AR synthesis.

Key insight: continuous AR prediction accumulates errors every step → noise.
Discrete tokens can't drift — a predicted token is always a valid codebook entry.

Architecture:
  - TextEncoder: byte-level transformer (same as before)
  - AudioEncoder: CNN + Transformer → continuous embeddings (same as before)
  - VectorQuantizer: continuous → discrete codebook indices + quantized embeddings
  - JEPAPredictor: predicts next codebook INDEX (classification, not regression)
  - MelDecoder: codebook embeddings → mel spectrogram
  - Vocos: mel → waveform (external, frozen)
"""

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

# ─── Vector Quantizer ─────────────────────────────────────────────────────────

class VectorQuantizer(nn.Module):
    """
    Discretizes continuous embeddings into codebook indices.
    Uses EMA codebook updates (no codebook loss needed in optimizer).
    Straight-through estimator for gradients.
    """

    def __init__(self, d_model=256, n_codes=1024, commitment_weight=0.25, ema_decay=0.99):
        super().__init__()
        self.d_model = d_model
        self.n_codes = n_codes
        self.commitment_weight = commitment_weight
        self.ema_decay = ema_decay

        # Codebook: [n_codes, d_model]
        self.register_buffer("codebook", torch.randn(n_codes, d_model))
        self.register_buffer("ema_count", torch.ones(n_codes))
        self.register_buffer("ema_sum", self.codebook.clone())
        self._initialized = False

    def _init_codebook(self, flat_z):
        """Initialize codebook from first batch using k-means++ style."""
        if self._initialized:
            return
        n = min(flat_z.shape[0], self.n_codes)
        indices = torch.randperm(flat_z.shape[0], device=flat_z.device)[:n]
        self.codebook[:n] = flat_z[indices].detach()
        self.ema_sum.copy_(self.codebook)
        self._initialized = True

    def forward(self, z):
        """
        Args:
            z: [B, T, d_model] continuous embeddings
        Returns:
            z_q: [B, T, d_model] quantized embeddings (straight-through)
            indices: [B, T] codebook indices
            commit_loss: scalar commitment loss
        """
        B, T, D = z.shape
        flat_z = z.reshape(-1, D)  # [B*T, D]

        # Initialize codebook from first batch
        if self.training and not self._initialized:
            self._init_codebook(flat_z)

        # Find nearest codebook entry: ||z - e||^2 = ||z||^2 - 2<z,e> + ||e||^2
        dists = (
            flat_z.pow(2).sum(dim=1, keepdim=True)
            - 2 * flat_z @ self.codebook.t()
            + self.codebook.pow(2).sum(dim=1, keepdim=True).t()
        )  # [B*T, n_codes]

        indices = dists.argmin(dim=1)  # [B*T]
        z_q = self.codebook[indices]   # [B*T, D]

        # EMA codebook update (only during training)
        if self.training:
            with torch.no_grad():
                one_hot = F.one_hot(indices, self.n_codes).float()  # [B*T, n_codes]
                counts = one_hot.sum(dim=0)  # [n_codes]
                sums = one_hot.t() @ flat_z  # [n_codes, D]

                self.ema_count.mul_(self.ema_decay).add_(counts, alpha=1 - self.ema_decay)
                self.ema_sum.mul_(self.ema_decay).add_(sums, alpha=1 - self.ema_decay)

                # Laplace smoothing to avoid dead codes
                n = self.ema_count.sum()
                count_smooth = (self.ema_count + 1e-5) / (n + self.n_codes * 1e-5) * n
                self.codebook.copy_(self.ema_sum / count_smooth.unsqueeze(1))

        # Commitment loss: push encoder output toward codebook
        commit_loss = F.mse_loss(flat_z, z_q.detach())

        # Straight-through estimator: z_q gets gradients of z
        z_q = z + (z_q - z).detach()

        z_q = z_q.reshape(B, T, D)
        indices = indices.reshape(B, T)

        return z_q, indices, self.commitment_weight * commit_loss

# ─── Shared components (same as model.py) ─────────────────────────────────────

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

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)

class AudioEncoder(nn.Module):
    def __init__(self, n_mels=100, d_model=256, nhead=4, num_layers=4,
                 downsample_factor=4, dropout=0.1):
        super().__init__()
        self.d_model = d_model
        self.downsample_factor = downsample_factor

        if downsample_factor == 2:
            self.conv_pre = nn.Sequential(
                nn.Conv1d(n_mels, d_model, kernel_size=7, padding=3), nn.GELU(),
                nn.Conv1d(d_model, d_model, kernel_size=4, stride=2, padding=1), nn.GELU(),
            )
        else:
            self.conv_pre = nn.Sequential(
                nn.Conv1d(n_mels, d_model, kernel_size=7, padding=3), nn.GELU(),
                nn.Conv1d(d_model, d_model, kernel_size=4, stride=2, padding=1), nn.GELU(),
                nn.Conv1d(d_model, d_model, kernel_size=4, stride=2, padding=1), nn.GELU(),
            )

        self.pos_embed = SinusoidalPositionalEncoding(d_model, max_len=4096)
        self.dropout = nn.Dropout(dropout)
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=d_model, nhead=nhead, dim_feedforward=d_model * 4,
            dropout=dropout, batch_first=True, activation="gelu",
        )
        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
        self.proj = nn.Linear(d_model, d_model)

    def forward(self, mel, mel_mask=None):
        x = self.conv_pre(mel)
        x = x.transpose(1, 2)
        if mel_mask is not None:
            T_down = x.shape[1]
            mel_mask = mel_mask[:, ::self.downsample_factor][:, :T_down]
        x = self.pos_embed(x)
        x = self.dropout(x)
        x = self.transformer(x, src_key_padding_mask=mel_mask)
        return self.proj(x)  # [B, T_down, d_model]

# ─── Predictor (now predicts codebook logits) ─────────────────────────────────

class TokenPredictor(nn.Module):
    """
    Causal transformer that predicts next CODEBOOK INDEX.
    Input: quantized embeddings (codebook lookups)
    Output: logits over codebook entries [B, T, n_codes]
    """

    def __init__(self, d_model=256, nhead=4, num_layers=6, n_codes=1024, dropout=0.1):
        super().__init__()
        self.d_model = d_model
        self.n_codes = n_codes

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

        # Output head: predict codebook index (classification, not regression!)
        self.output_head = nn.Sequential(
            nn.Linear(d_model, d_model),
            nn.GELU(),
            nn.Linear(d_model, n_codes),
        )

    def forward(self, audio_emb, text_emb, audio_mask=None, text_mask=None):
        x = self.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)

        x = self.transformer(
            x, text_emb,
            tgt_mask=causal_mask,
            tgt_key_padding_mask=audio_mask,
            memory_key_padding_mask=text_mask,
        )

        return self.output_head(x)  # [B, T, n_codes]

    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):
        """One-step cached inference. Returns logits [B, 1, n_codes] and updated cache."""
        x = self.input_proj(new_emb)
        x = x + self.pos_embed.pe[:, step_idx:step_idx + 1]

        for i, layer in enumerate(self.transformer.layers):
            # Self-attention with cache
            sa_out, cache['self_k'][i], cache['self_v'][i] = self._cached_mha(
                layer.self_attn, x, x, cache['self_k'][i], cache['self_v'][i]
            )
            x = layer.norm1(x + layer.dropout1(sa_out))

            # Cross-attention with cached text
            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], text_mask
            )
            x = layer.norm2(x + layer.dropout2(ca_out))

            # FFN
            ff_out = layer.linear2(layer.dropout(layer.activation(layer.linear1(x))))
            x = layer.norm3(x + layer.dropout3(ff_out))

        return self.output_head(x), cache  # [B, 1, n_codes]

    def _cached_mha(self, mha, q, kv, cache_k, cache_v):
        d, nhead = mha.embed_dim, mha.num_heads
        head_dim = d // nhead
        B = q.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_proj = F.linear(q, Wq, bq)
        k_new = F.linear(kv, Wk, bk)
        v_new = F.linear(kv, Wv, bv)

        k = torch.cat([cache_k, k_new], dim=1) if cache_k is not None and cache_k.shape[1] > 0 else k_new
        v = torch.cat([cache_v, v_new], dim=1) if cache_v is not None and cache_v.shape[1] > 0 else v_new

        q_proj = q_proj.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_proj, k_mh.transpose(-2, -1)) / (head_dim ** 0.5)
        out = torch.matmul(F.softmax(attn, dim=-1), v_mh)
        out = out.transpose(1, 2).contiguous().view(B, 1, d)
        return F.linear(out, mha.out_proj.weight, mha.out_proj.bias), k, v

    def _cached_cross_attn(self, mha, q, memory, cache_k, cache_v, memory_mask=None):
        d, nhead = mha.embed_dim, mha.num_heads
        head_dim = d // nhead
        B = q.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_proj = F.linear(q, 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_proj = q_proj.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_proj, 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'))
        out = torch.matmul(F.softmax(attn, dim=-1), v_mh)
        out = out.transpose(1, 2).contiguous().view(B, 1, d)
        return F.linear(out, mha.out_proj.weight, mha.out_proj.bias), k, v

# ─── Mel Decoder ──────────────────────────────────────────────────────────────

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 MelDecoder(nn.Module):
    def __init__(self, d_model=256, n_mels=100, upsample_factor=4):
        super().__init__()
        hidden = d_model * 2

        self.proj = nn.Sequential(nn.Linear(d_model, hidden), nn.GELU())

        if upsample_factor == 2:
            self.upsample = nn.Sequential(
                nn.ConvTranspose1d(hidden, hidden, kernel_size=4, stride=2, padding=1), nn.GELU(),
                ResConvBlock(hidden, kernel_size=3),
            )
        else:
            self.upsample = nn.Sequential(
                nn.ConvTranspose1d(hidden, hidden, kernel_size=4, stride=2, padding=1), nn.GELU(),
                ResConvBlock(hidden, kernel_size=3),
                nn.ConvTranspose1d(hidden, hidden, kernel_size=4, stride=2, padding=1), nn.GELU(),
                ResConvBlock(hidden, kernel_size=3),
            )

        self.out = nn.Sequential(
            nn.Conv1d(hidden, hidden, kernel_size=7, padding=3), nn.GELU(),
            ResConvBlock(hidden, kernel_size=5),
            nn.Conv1d(hidden, n_mels, kernel_size=7, padding=3),
        )

    def forward(self, z):
        x = self.proj(z).transpose(1, 2)
        x = self.upsample(x)
        return self.out(x)

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

class MultiResolutionSpectralLoss(nn.Module):
    def __init__(self, resolutions=((64, 16, 64), (128, 32, 128), (256, 64, 256))):
        super().__init__()
        self.resolutions = resolutions

    def _stft_loss(self, pred, target, n_fft, hop_length, win_length):
        window = torch.hann_window(win_length, device=pred.device)
        B, M, T = pred.shape
        pred_stft = torch.stft(pred.reshape(B*M, T), n_fft, hop_length, win_length, window=window, return_complex=True)
        target_stft = torch.stft(target.reshape(B*M, T), n_fft, hop_length, win_length, window=window, return_complex=True)
        sc = torch.norm(target_stft.abs() - pred_stft.abs(), p="fro") / (torch.norm(target_stft.abs(), p="fro") + 1e-8)
        log_mag = F.l1_loss(torch.log(pred_stft.abs() + 1e-7), torch.log(target_stft.abs() + 1e-7))
        return sc + log_mag

    def forward(self, pred, target):
        loss = torch.tensor(0.0, device=pred.device)
        count = 0
        for n_fft, hop, win in self.resolutions:
            if pred.shape[2] >= n_fft:
                loss = loss + self._stft_loss(pred, target, n_fft, hop, win)
                count += 1
        return loss / max(count, 1)

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

class LeWMTTSvq(nn.Module):
    """
    JEPA + VQ TTS model.

    Training flow:
      mel → AudioEncoder → continuous z → VQ → discrete indices + quantized z_q
      [start, z_q[:-1]] + text → TokenPredictor → logits → cross_entropy(logits, indices)
      z_q → MelDecoder → mel_recon → L1 + spectral loss

    Inference flow:
      text → TextEncoder → text_emb
      start_emb → Predictor → logits → argmax → codebook lookup → next input → repeat
      all codebook embeddings → MelDecoder → mel → Vocos → waveform

    No continuous drift possible — every predicted embedding is a valid codebook entry.
    """

    def __init__(self, config):
        super().__init__()
        d_model = config.get("d_model", 256)
        nhead = config.get("nhead", 4)
        n_mels = config.get("n_mels", 100)
        text_vocab = config.get("text_vocab_size", 256)
        text_layers = config.get("text_encoder_layers", 4)
        audio_layers = config.get("audio_encoder_layers", 4)
        predictor_layers = config.get("predictor_layers", 6)
        dropout = config.get("dropout", 0.1)
        n_codes = config.get("n_codes", 1024)
        downsample_factor = config.get("downsample_factor", 4)

        self.text_encoder = TextEncoder(
            vocab_size=text_vocab, d_model=d_model, nhead=nhead,
            num_layers=text_layers, dropout=dropout,
        )
        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.vq = VectorQuantizer(
            d_model=d_model, n_codes=n_codes,
            commitment_weight=config.get("commitment_weight", 0.25),
        )
        self.predictor = TokenPredictor(
            d_model=d_model, nhead=nhead,
            num_layers=predictor_layers, n_codes=n_codes, dropout=dropout,
        )
        self.mel_decoder = MelDecoder(
            d_model=d_model, n_mels=n_mels, upsample_factor=downsample_factor,
        )

        # 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", 1.0)
        self.recon_weight = config.get("recon_weight", 1.0)
        self.spectral_weight = config.get("spectral_weight", 0.5)

        # Spectral loss
        self.spectral_loss_fn = MultiResolutionSpectralLoss()

        # Label smoothing for token prediction (prevents overconfidence)
        self.label_smoothing = config.get("label_smoothing", 0.1)

        self.config = config
        self.downsample_factor = downsample_factor

    def forward(self, mel, text_tokens, mel_mask=None, text_mask=None, **kwargs):
        B = mel.shape[0]

        # Encode text
        text_emb = self.text_encoder(text_tokens, text_mask)

        # Encode audio → continuous → quantize to discrete
        z_continuous = self.audio_encoder(mel, mel_mask)  # [B, T_down, d]
        z_q, indices, commit_loss = self.vq(z_continuous)  # [B, T_down, d], [B, T_down]

        T_down = z_q.shape[1]

        # Predictor input: [start, z_q[0], z_q[1], ..., z_q[T-2]]
        # Predictor target: [idx[0], idx[1], ..., idx[T-1]]
        start = self.start_emb.expand(B, -1, -1)
        input_emb = torch.cat([start, z_q[:, :-1]], dim=1)  # [B, T_down, d]

        # Masks
        if mel_mask is not None:
            ds_mask = mel_mask[:, ::self.downsample_factor][:, :T_down]
            start_mask = torch.zeros(B, 1, dtype=torch.bool, device=mel.device)
            pred_mask = torch.cat([start_mask, ds_mask[:, :-1]], dim=1)
            loss_mask = ds_mask
        else:
            pred_mask = None
            loss_mask = None

        # Predict next token logits
        logits = self.predictor(input_emb, text_emb, pred_mask, text_mask)  # [B, T_down, n_codes]

        # ─── Loss 1: Token prediction (cross-entropy with label smoothing) ───
        if loss_mask is not None:
            # Flatten valid positions only
            valid = ~loss_mask  # [B, T_down]
            logits_flat = logits[valid]  # [N_valid, n_codes]
            targets_flat = indices[valid]  # [N_valid]
            if logits_flat.shape[0] > 0:
                token_loss = F.cross_entropy(logits_flat, targets_flat,
                                              label_smoothing=self.label_smoothing)
            else:
                token_loss = torch.tensor(0.0, device=mel.device)
        else:
            token_loss = F.cross_entropy(
                logits.reshape(-1, logits.shape[-1]), indices.reshape(-1),
                label_smoothing=self.label_smoothing,
            )

        # ─── Loss 2: VQ commitment loss ───
        # (already computed in VQ forward)

        # ─── Loss 3: Mel reconstruction from quantized embeddings ───
        mel_recon = self.mel_decoder(z_q)
        T_mel = mel.shape[2]
        T_recon = mel_recon.shape[2]
        T_min = min(T_mel, T_recon)
        mel = mel[:, :, :T_min]
        mel_recon = mel_recon[:, :, :T_min]
        if mel_mask is not None:
            mel_mask = mel_mask[:, :T_min]

        if mel_mask is not None:
            valid_mel = (~mel_mask).unsqueeze(1).float()
            recon_loss = F.l1_loss(
                mel_recon * valid_mel, mel * valid_mel, reduction="sum"
            ) / (valid_mel.sum() * mel.shape[1] + 1e-8)
        else:
            recon_loss = F.l1_loss(mel_recon, mel)

        # ─── Loss 4: Spectral loss ───
        if mel_mask is not None:
            spectral_loss = self.spectral_loss_fn(mel_recon * valid_mel, mel * valid_mel)
        else:
            spectral_loss = self.spectral_loss_fn(mel_recon, mel)

        # ─── Token accuracy (for monitoring) ───
        with torch.no_grad():
            if loss_mask is not None and logits_flat.shape[0] > 0:
                token_acc = (logits_flat.argmax(dim=-1) == targets_flat).float().mean()
            else:
                token_acc = (logits.argmax(dim=-1) == indices).float().mean()

        total_loss = (self.pred_weight * token_loss
                      + commit_loss
                      + self.recon_weight * recon_loss
                      + self.spectral_weight * spectral_loss)

        return {
            "total_loss": total_loss,
            "token_loss": token_loss,
            "commit_loss": commit_loss,
            "recon_loss": recon_loss,
            "spectral_loss": spectral_loss,
            "token_accuracy": token_acc,
        }

    # ─── Inference ────────────────────────────────────────────────────────────

    def encode_to_tokens(self, mel):
        """Encode mel → discrete token indices."""
        z = self.audio_encoder(mel)
        _, indices, _ = self.vq(z)
        return indices

    @torch.no_grad()
    def synthesize_tokens(self, text_tokens, max_steps=300, temperature=1.0,
                          top_k=50, text_mask=None):
        """AR token generation: text → token indices."""
        text_emb = self.text_encoder(text_tokens, text_mask)

        cache = self.predictor.init_cache(len(self.predictor.transformer.layers))
        emb = self.start_emb
        all_indices = []

        for step in range(max_steps):
            logits, cache = self.predictor.inference_step(
                emb, text_emb, step, cache, text_mask
            )  # [B, 1, n_codes]

            logits = logits.squeeze(1)  # [B, n_codes]

            # Temperature + top-k sampling
            logits = logits / max(temperature, 1e-8)
            if top_k > 0:
                topk_vals, _ = logits.topk(top_k, dim=-1)
                logits[logits < topk_vals[:, -1:]] = float('-inf')

            probs = F.softmax(logits, dim=-1)
            idx = torch.multinomial(probs, 1)  # [B, 1]
            all_indices.append(idx)

            # Lookup codebook for next input
            emb = self.vq.codebook[idx]  # [B, 1, d]

        return torch.cat(all_indices, dim=1)  # [B, max_steps]

    @torch.no_grad()
    def tokens_to_mel(self, indices):
        """Convert token indices → mel spectrogram."""
        z_q = self.vq.codebook[indices]  # [B, T, d]
        return self.mel_decoder(z_q)  # [B, n_mels, T_up]

    @torch.no_grad()
    def synthesize_mel(self, text_tokens, max_steps=300, temperature=1.0,
                       top_k=50, text_mask=None):
        """Full pipeline: text → tokens → mel."""
        indices = self.synthesize_tokens(
            text_tokens, max_steps, temperature, top_k, text_mask
        )
        return self.tokens_to_mel(indices), indices

def count_parameters(model):
    return sum(p.numel() for p in model.parameters() if p.requires_grad)

def build_model_vq(config=None):
    if config is None:
        config = {
            "d_model": 256,
            "nhead": 4,
            "n_mels": 100,
            "text_vocab_size": 256,
            "text_encoder_layers": 4,
            "audio_encoder_layers": 4,
            "predictor_layers": 6,
            "n_codes": 1024,
            "dropout": 0.1,
            "pred_weight": 1.0,
            "recon_weight": 1.0,
            "spectral_weight": 0.5,
            "commitment_weight": 0.25,
            "label_smoothing": 0.1,
            "downsample_factor": 4,
        }
    model = LeWMTTSvq(config)
    print(f"LeWM TTS VQ model: {count_parameters(model)/1e6:.2f}M parameters")
    return model, config


Now the training script:


"""Training loop for LeWM TTS VQ model."""

import os
import json
import time
import argparse
import torch
from torch.utils.tensorboard import SummaryWriter
from pathlib import Path

from model_vq import build_model_vq
from dataset import build_dataloader

def train(args):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Device: {device}")

    config = {
        "d_model": args.d_model,
        "nhead": args.nhead,
        "n_mels": 100,
        "text_vocab_size": 256,
        "text_encoder_layers": args.text_layers,
        "audio_encoder_layers": args.audio_layers,
        "predictor_layers": args.predictor_layers,
        "n_codes": args.n_codes,
        "dropout": args.dropout,
        "pred_weight": args.pred_weight,
        "recon_weight": args.recon_weight,
        "spectral_weight": args.spectral_weight,
        "commitment_weight": args.commitment_weight,
        "label_smoothing": args.label_smoothing,
        "downsample_factor": args.downsample_factor,
    }
    model, config = build_model_vq(config)
    model = model.to(device)

    manifest_path = os.path.join(args.data_dir, "manifest.json")
    loader, dataset = build_dataloader(
        manifest_path, batch_size=args.batch_size, num_workers=args.num_workers,
    )
    print(f"Training: {len(dataset)} samples, {len(loader)} batches/epoch")

    optimizer = torch.optim.AdamW(
        model.parameters(), lr=args.lr, betas=(0.9, 0.98), weight_decay=0.01,
    )

    total_steps = args.epochs * len(loader)
    warmup_steps = min(4000, total_steps // 10)

    def lr_lambda(step):
        if step < warmup_steps:
            return step / max(1, warmup_steps)
        progress = (step - warmup_steps) / max(1, total_steps - warmup_steps)
        return 0.5 * (1 + torch.cos(torch.tensor(progress * 3.14159)).item())

    scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)

    ckpt_dir = Path(args.output_dir) / "checkpoints"
    ckpt_dir.mkdir(parents=True, exist_ok=True)
    writer = SummaryWriter(Path(args.output_dir) / "logs")

    with open(Path(args.output_dir) / "config.json", "w") as f:
        json.dump(config, f, indent=2)

    start_epoch = 0
    global_step = 0
    if args.resume:
        ckpt = torch.load(args.resume, map_location=device, weights_only=False)
        model.load_state_dict(ckpt["model"])
        optimizer.load_state_dict(ckpt["optimizer"])
        scheduler.load_state_dict(ckpt["scheduler"])
        start_epoch = ckpt["epoch"] + 1
        global_step = ckpt["global_step"]
        print(f"Resumed from epoch {start_epoch}, step {global_step}")

    model.train()
    best_loss = float("inf")

    for epoch in range(start_epoch, args.epochs):
        ep_loss = ep_tok = ep_commit = ep_recon = ep_spec = ep_acc = 0.0
        t0 = time.time()

        for batch_idx, batch in enumerate(loader):
            mel = batch["mel"].to(device)
            text_tokens = batch["text_tokens"].to(device)
            mel_mask = batch["mel_mask"].to(device)
            text_mask = batch["text_mask"].to(device)

            losses = model(mel, text_tokens, mel_mask, text_mask)

            total_loss = losses["total_loss"]

            optimizer.zero_grad()
            total_loss.backward()
            grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
            optimizer.step()
            scheduler.step()

            ep_loss += total_loss.item()
            ep_tok += losses["token_loss"].item()
            ep_commit += losses["commit_loss"].item()
            ep_recon += losses["recon_loss"].item()
            ep_spec += losses["spectral_loss"].item()
            ep_acc += losses["token_accuracy"].item()
            global_step += 1

            if global_step % args.log_every == 0:
                lr = optimizer.param_groups[0]["lr"]
                for k, v in losses.items():
                    writer.add_scalar(f"train/{k}", v.item(), global_step)
                writer.add_scalar("train/grad_norm", grad_norm.item(), global_step)
                writer.add_scalar("train/lr", lr, global_step)

            if global_step % args.print_every == 0:
                print(
                    f"  [{epoch+1}/{args.epochs}] step {global_step} | "
                    f"loss={total_loss.item():.4f} tok={losses['token_loss'].item():.4f} "
                    f"commit={losses['commit_loss'].item():.4f} "
                    f"recon={losses['recon_loss'].item():.4f} "
                    f"spec={losses['spectral_loss'].item():.4f} "
                    f"acc={losses['token_accuracy'].item():.3f} | "
                    f"grad={grad_norm.item():.2f} | {time.time()-t0:.1f}s"
                )

        n = len(loader)
        epoch_time = time.time() - t0
        print(
            f"Epoch {epoch+1}/{args.epochs} | "
            f"loss={ep_loss/n:.4f} tok={ep_tok/n:.4f} commit={ep_commit/n:.4f} "
            f"recon={ep_recon/n:.4f} spec={ep_spec/n:.4f} "
            f"acc={ep_acc/n:.3f} | {epoch_time:.1f}s"
        )

        if (epoch + 1) % args.save_every == 0 or (epoch + 1) == args.epochs:
            torch.save({
                "epoch": epoch, "global_step": global_step,
                "model": model.state_dict(),
                "optimizer": optimizer.state_dict(),
                "scheduler": scheduler.state_dict(),
                "config": config, "loss": ep_loss / n,
            }, ckpt_dir / f"epoch_{epoch+1:04d}.pt")

        if ep_loss / n < best_loss:
            best_loss = ep_loss / n
            torch.save({
                "epoch": epoch, "global_step": global_step,
                "model": model.state_dict(), "config": config, "loss": best_loss,
            }, ckpt_dir / "best.pt")
            print(f"  New best (loss={best_loss:.4f})")

    writer.close()
    print(f"\nDone. Best loss: {best_loss:.4f}")

if __name__ == "__main__":
    p = argparse.ArgumentParser()
    p.add_argument("--data_dir", default="/home/ubuntu/lewm-tts/processed_data_100mel")
    p.add_argument("--output_dir", default="/home/ubuntu/lewm-tts/output_vq")
    p.add_argument("--d_model", type=int, default=256)
    p.add_argument("--nhead", type=int, default=4)
    p.add_argument("--text_layers", type=int, default=4)
    p.add_argument("--audio_layers", type=int, default=4)
    p.add_argument("--predictor_layers", type=int, default=6)
    p.add_argument("--n_codes", type=int, default=1024)
    p.add_argument("--dropout", type=float, default=0.1)
    p.add_argument("--downsample_factor", type=int, default=4)
    p.add_argument("--pred_weight", type=float, default=1.0)
    p.add_argument("--recon_weight", type=float, default=1.0)
    p.add_argument("--spectral_weight", type=float, default=0.5)
    p.add_argument("--commitment_weight", type=float, default=0.25)
    p.add_argument("--label_smoothing", type=float, default=0.1)
    p.add_argument("--epochs", type=int, default=200)
    p.add_argument("--batch_size", type=int, default=16)
    p.add_argument("--lr", type=float, default=1e-4)
    p.add_argument("--grad_clip", type=float, default=1.0)
    p.add_argument("--num_workers", type=int, default=4)
    p.add_argument("--log_every", type=int, default=10)
    p.add_argument("--print_every", type=int, default=100)
    p.add_argument("--save_every", type=int, default=25)
    p.add_argument("--resume", default=None)
    train(p.parse_args())


File created successfully at: /home/ubuntu/lewm-tts/train_vq.py

"""Training loop for LeWM TTS VQ model."""

import os
import json
import time
import argparse
import torch
from torch.utils.tensorboard import SummaryWriter
from pathlib import Path

from model_vq import build_model_vq
from dataset import build_dataloader

def train(args):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Device: {device}")

    config = {
        "d_model": args.d_model,
        "nhead": args.nhead,
        "n_mels": 100,
        "text_vocab_size": 256,
        "text_encoder_layers": args.text_layers,
        "audio_encoder_layers": args.audio_layers,
        "predictor_layers": args.predictor_layers,
        "n_codes": args.n_codes,
        "dropout": args.dropout,
        "pred_weight": args.pred_weight,
        "recon_weight": args.recon_weight,
        "spectral_weight": args.spectral_weight,
        "commitment_weight": args.commitment_weight,
        "label_smoothing": args.label_smoothing,
        "downsample_factor": args.downsample_factor,
    }
    model, config = build_model_vq(config)
    model = model.to(device)

    manifest_path = os.path.join(args.data_dir, "manifest.json")
    loader, dataset = build_dataloader(
        manifest_path, batch_size=args.batch_size, num_workers=args.num_workers,
    )
    print(f"Training: {len(dataset)} samples, {len(loader)} batches/epoch")

    optimizer = torch.optim.AdamW(
        model.parameters(), lr=args.lr, betas=(0.9, 0.98), weight_decay=0.01,
    )

    total_steps = args.epochs * len(loader)
    warmup_steps = min(4000, total_steps // 10)

    def lr_lambda(step):
        if step < warmup_steps:
            return step / max(1, warmup_steps)
        progress = (step - warmup_steps) / max(1, total_steps - warmup_steps)
        return 0.5 * (1 + torch.cos(torch.tensor(progress * 3.14159)).item())

    scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)

    ckpt_dir = Path(args.output_dir) / "checkpoints"
    ckpt_dir.mkdir(parents=True, exist_ok=True)
    writer = SummaryWriter(Path(args.output_dir) / "logs")

    with open(Path(args.output_dir) / "config.json", "w") as f:
        json.dump(config, f, indent=2)

    start_epoch = 0
    global_step = 0
    if args.resume:
        ckpt = torch.load(args.resume, map_location=device, weights_only=False)
        model.load_state_dict(ckpt["model"])
        optimizer.load_state_dict(ckpt["optimizer"])
        scheduler.load_state_dict(ckpt["scheduler"])
        start_epoch = ckpt["epoch"] + 1
        global_step = ckpt["global_step"]
        print(f"Resumed from epoch {start_epoch}, step {global_step}")

    model.train()
    best_loss = float("inf")

    for epoch in range(start_epoch, args.epochs):
        ep_loss = ep_tok = ep_commit = ep_recon = ep_spec = ep_acc = 0.0
        t0 = time.time()

        for batch_idx, batch in enumerate(loader):
            mel = batch["mel"].to(device)
            text_tokens = batch["text_tokens"].to(device)
            mel_mask = batch["mel_mask"].to(device)
            text_mask = batch["text_mask"].to(device)

            losses = model(mel, text_tokens, mel_mask, text_mask)

            total_loss = losses["total_loss"]

            optimizer.zero_grad()
            total_loss.backward()
            grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
            optimizer.step()
            scheduler.step()

            ep_loss += total_loss.item()
            ep_tok += losses["token_loss"].item()
            ep_commit += losses["commit_loss"].item()
            ep_recon += losses["recon_loss"].item()
            ep_spec += losses["spectral_loss"].item()
            ep_acc += losses["token_accuracy"].item()
            global_step += 1

            if global_step % args.log_every == 0:
                lr = optimizer.param_groups[0]["lr"]
                for k, v in losses.items():
                    writer.add_scalar(f"train/{k}", v.item(), global_step)
                writer.add_scalar("train/grad_norm", grad_norm.item(), global_step)
                writer.add_scalar("train/lr", lr, global_step)

            if global_step % args.print_every == 0:
                print(
                    f"  [{epoch+1}/{args.epochs}] step {global_step} | "
                    f"loss={total_loss.item():.4f} tok={losses['token_loss'].item():.4f} "
                    f"commit={losses['commit_loss'].item():.4f} "
                    f"recon={losses['recon_loss'].item():.4f} "
                    f"spec={losses['spectral_loss'].item():.4f} "
                    f"acc={losses['token_accuracy'].item():.3f} | "
                    f"grad={grad_norm.item():.2f} | {time.time()-t0:.1f}s"
                )

        n = len(loader)
        epoch_time = time.time() - t0
        print(
            f"Epoch {epoch+1}/{args.epochs} | "
            f"loss={ep_loss/n:.4f} tok={ep_tok/n:.4f} commit={ep_commit/n:.4f} "
            f"recon={ep_recon/n:.4f} spec={ep_spec/n:.4f} "
            f"acc={ep_acc/n:.3f} | {epoch_time:.1f}s"
        )

        if (epoch + 1) % args.save_every == 0 or (epoch + 1) == args.epochs:
            torch.save({
                "epoch": epoch, "global_step": global_step,
                "model": model.state_dict(),
                "optimizer": optimizer.state_dict(),
                "scheduler": scheduler.state_dict(),
                "config": config, "loss": ep_loss / n,
            }, ckpt_dir / f"epoch_{epoch+1:04d}.pt")

        if ep_loss / n < best_loss:
            best_loss = ep_loss / n
            torch.save({
                "epoch": epoch, "global_step": global_step,
                "model": model.state_dict(), "config": config, "loss": best_loss,
            }, ckpt_dir / "best.pt")
            print(f"  New best (loss={best_loss:.4f})")

    writer.close()
    print(f"\nDone. Best loss: {best_loss:.4f}")

if __name__ == "__main__":
    p = argparse.ArgumentParser()
    p.add_argument("--data_dir", default="/home/ubuntu/lewm-tts/processed_data_100mel")
    p.add_argument("--output_dir", default="/home/ubuntu/lewm-tts/output_vq")
    p.add_argument("--d_model", type=int, default=256)
    p.add_argument("--nhead", type=int, default=4)
    p.add_argument("--text_layers", type=int, default=4)
    p.add_argument("--audio_layers", type=int, default=4)
    p.add_argument("--predictor_layers", type=int, default=6)
    p.add_argument("--n_codes", type=int, default=1024)
    p.add_argument("--dropout", type=float, default=0.1)
    p.add_argument("--downsample_factor", type=int, default=4)
    p.add_argument("--pred_weight", type=float, default=1.0)
    p.add_argument("--recon_weight", type=float, default=1.0)
    p.add_argument("--spectral_weight", type=float, default=0.5)
    p.add_argument("--commitment_weight", type=float, default=0.25)
    p.add_argument("--label_smoothing", type=float, default=0.1)
    p.add_argument("--epochs", type=int, default=200)
    p.add_argument("--batch_size", type=int, default=16)
    p.add_argument("--lr", type=float, default=1e-4)
    p.add_argument("--grad_clip", type=float, default=1.0)
    p.add_argument("--num_workers", type=int, default=4)
    p.add_argument("--log_every", type=int, default=10)
    p.add_argument("--print_every", type=int, default=100)
    p.add_argument("--save_every", type=int, default=25)
    p.add_argument("--resume", default=None)
    train(p.parse_args())


Test it compiles and forward/backward works:


Bash

python3 -c "
import torch
from model_vq import build_model_vq

model, config = build_model_vq()
model.train()

B = 4
mel = torch.randn(B, 100, 400)
text = torch.randint(0, 256, (B, 30))
mel_mask = torch.zeros(B, 400, dtype=torch.bool)
text_mask = torch.zeros(B, 30, dtype=torch.bool)

losses = model(mel, text, mel_mask, text_mask)
losses['total_loss'].backward()
print('Forward+backward OK')
for k, v in losses.items():
    print(f'  {k}: {v.item():.4f}')

# Test inference
model.eval()
with torch.no_grad():
    text_tok = torch.randint(0, 256, (1, 20))
    mel_out, indices = model.synthesize_mel(text_tok, max_steps=50, temperature=1.0, top_k=50)
    print(f'Synth mel: {mel_out.shape}')
    print(f'Synth indices: {indices.shape}, unique tokens: {indices.unique().shape[0]}')

# Size
total = sum(p.numel() for p in model.parameters())
print(f'Total params: {total/1e6:.2f}M')
" 2>&1

Test VQ model


Error: Exit code 1
Traceback (most recent call last):
  File "<string>", line 14, in <module>
  File "/home/ubuntu/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1775, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/ubuntu/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1786, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/ubuntu/lewm-tts/model_vq.py", line 467, in forward
    z_q, indices, commit_loss = self.vq(z_continuous)  # [B, T_down, d], [B, T_down]
  File "/home/ubuntu/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1775, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/ubuntu/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1786, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/ubuntu/lewm-tts/model_vq.py", line 100, in forward
    z_q = z + (z_q - z).detach()
RuntimeError: The size of tensor a (400) must match the size of tensor b (100) at non-singleton dimension 1
LeWM TTS VQ model: 24.13M parameters


The issue is that `z_q` was flattened and reshaped but `z` wasn't. The straight-through needs to happen on the flat tensors before reshaping:


Edit

        # Straight-through estimator: z_q gets gradients of z
        z_q = z + (z_q - z).detach()

        z_q = z_q.reshape(B, T, D)
        indices = indices.reshape(B, T)

        # Straight-through estimator: z_q gets gradients of z (both are flat here)
        z_q = flat_z + (z_q - flat_z).detach()

        z_q = z_q.reshape(B, T, D)
        indices = indices.reshape(B, T)