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

    def forward(self, pred, target):
        """pred, target: [B, n_mels, T]"""
        loss = torch.tensor(0.0, device=pred.device)
        count = 0
        for n_fft, hop, win in self.resolutions:
            # Skip if T is too short for this resolution
            if pred.shape[2] >= n_fft:
                loss = loss + self._stft_loss(pred, target, n_fft, hop, win)
                count += 1
        if count > 0:
            loss = loss / count
        return loss

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

Architecture:
  - TextEncoder: Character-level transformer for Hindi/Devanagari text
  - AudioEncoder: 1D CNN + Transformer that encodes mel spectrograms to embeddings
  - JEPAPredictor: Transformer that predicts next audio embedding given text + audio context
  - Two losses: next-embedding prediction (MSE) + Gaussian regularizer (KL to N(0,I))
"""

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

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

class TextEncoder(nn.Module):
    """Character-level transformer encoder for Hindi text."""

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

        # Character embedding (covers all Unicode Hindi chars via byte-level encoding)
        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):
        """
        Args:
            text_tokens: [B, T_text] long tensor of byte-level token IDs
            text_mask: [B, T_text] bool, True = padding
        Returns:
            text_emb: [B, T_text, d_model]
        """
        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)
        x = self.proj(x)
        return x

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

class AudioEncoder(nn.Module):
    """
    Encodes mel spectrograms into a sequence of audio embeddings.
    Uses 1D CNN for downsampling + Transformer for contextualization.
    """

    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

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

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

        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)

        # Project to embedding with mu and logvar for Gaussian regularizer
        self.proj_mu = nn.Linear(d_model, d_model)
        self.proj_logvar = nn.Linear(d_model, d_model)

    def forward(self, mel, mel_mask=None):
        """
        Args:
            mel: [B, n_mels, T_mel] — log mel spectrogram
            mel_mask: [B, T_mel] bool, True = padding (before downsampling)
        Returns:
            z: [B, T_down, d_model] — sampled latent embeddings
            mu: [B, T_down, d_model]
            logvar: [B, T_down, d_model]
        """
        # CNN encoding: [B, n_mels, T] -> [B, d_model, T//4]
        x = self.conv_pre(mel)
        x = x.transpose(1, 2)  # [B, T_down, d_model]

        # Downsample mask if provided
        if mel_mask is not None:
            # Approximate downsampled mask
            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)

        mu = self.proj_mu(x)
        logvar = self.proj_logvar(x)

        # Reparameterization trick
        if self.training:
            std = torch.exp(0.5 * logvar)
            eps = torch.randn_like(std)
            z = mu + eps * std
        else:
            z = mu

        return z, mu, logvar

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

class JEPAPredictor(nn.Module):
    """
    Predicts next audio embedding given:
    - Text context (from TextEncoder)
    - Previous audio embeddings (from AudioEncoder)

    Uses cross-attention to condition on text, and causal self-attention
    over the audio embedding sequence.
    """

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

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

        # Transformer decoder with causal self-attention + cross-attention to text
        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 projection to predict next embedding
        self.output_proj = nn.Sequential(
            nn.Linear(d_model, d_model),
            nn.GELU(),
            nn.Linear(d_model, d_model),
        )

    def forward(self, audio_emb, text_emb, audio_mask=None, text_mask=None):
        """
        Args:
            audio_emb: [B, T_audio, d_model] — audio embeddings (from encoder)
            text_emb: [B, T_text, d_model] — text embeddings
            audio_mask: [B, T_audio] bool, True = padding
            text_mask: [B, T_text] bool, True = padding
        Returns:
            predicted: [B, T_audio, d_model] — predicted next embeddings
        """
        x = self.audio_input_proj(audio_emb)
        x = self.pos_embed(x)
        x = self.dropout(x)

        # Causal mask: each position can only attend to previous positions
        T = x.shape[1]
        causal_mask = torch.triu(
            torch.ones(T, T, device=x.device, dtype=torch.bool), diagonal=1
        )

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

        predicted = self.output_proj(predicted)
        return predicted

    def _cached_mha(self, mha, q_input, kv_input, cache_k, cache_v):
        """Run multi-head attention with KV cache.
        q_input: [B, 1, d], kv_input: [B, 1, d] (new token for K,V)
        cache_k, cache_v: [B, T_prev, d] projected keys/values
        Returns: output [B, 1, d], updated cache_k, updated cache_v
        """
        d = mha.embed_dim
        nhead = mha.num_heads
        head_dim = d // nhead
        B = q_input.shape[0]

        # Project Q from q_input, K/V from kv_input
        # in_proj_weight is [3d, d] = [Wq; Wk; Wv]
        Wq, Wk, Wv = mha.in_proj_weight.chunk(3, dim=0)
        bq, bk, bv = mha.in_proj_bias.chunk(3, dim=0)

        q = F.linear(q_input, Wq, bq)        # [B, 1, d]
        k_new = F.linear(kv_input, Wk, bk)    # [B, 1, d]
        v_new = F.linear(kv_input, Wv, bv)    # [B, 1, d]

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

        # Multi-head attention
        q = q.view(B, 1, nhead, head_dim).transpose(1, 2)      # [B, nhead, 1, head_dim]
        k_mh = k.view(B, -1, nhead, head_dim).transpose(1, 2)  # [B, nhead, T, head_dim]
        v_mh = v.view(B, -1, nhead, head_dim).transpose(1, 2)

        attn = torch.matmul(q, k_mh.transpose(-2, -1)) / (head_dim ** 0.5)
        attn = F.softmax(attn, dim=-1)
        out = torch.matmul(attn, v_mh)  # [B, nhead, 1, head_dim]

        out = out.transpose(1, 2).contiguous().view(B, 1, d)
        out = F.linear(out, mha.out_proj.weight, mha.out_proj.bias)

        return out, k, v

    def _cached_cross_attn(self, mha, q_input, memory, cache_k, cache_v, memory_mask=None):
        """Cross-attention with cached memory K,V (computed once).
        memory_mask: [B, T_text] bool, True = padding (applied as -inf)."""
        d = mha.embed_dim
        nhead = mha.num_heads
        head_dim = d // nhead
        B = q_input.shape[0]

        Wq, Wk, Wv = mha.in_proj_weight.chunk(3, dim=0)
        bq, bk, bv = mha.in_proj_bias.chunk(3, dim=0)

        q = F.linear(q_input, Wq, bq)  # [B, 1, d]

        if cache_k is None:
            k = F.linear(memory, Wk, bk)  # [B, T_text, d]
            v = F.linear(memory, Wv, bv)
        else:
            k, v = cache_k, cache_v

        q = q.view(B, 1, nhead, head_dim).transpose(1, 2)
        k_mh = k.view(B, -1, nhead, head_dim).transpose(1, 2)
        v_mh = v.view(B, -1, nhead, head_dim).transpose(1, 2)

        attn = torch.matmul(q, k_mh.transpose(-2, -1)) / (head_dim ** 0.5)
        if memory_mask is not None:
            # memory_mask: [B, T_text] → [B, 1, 1, T_text]
            attn = attn.masked_fill(memory_mask.unsqueeze(1).unsqueeze(2), float('-inf'))
        attn = F.softmax(attn, dim=-1)
        out = torch.matmul(attn, v_mh)

        out = out.transpose(1, 2).contiguous().view(B, 1, d)
        out = F.linear(out, mha.out_proj.weight, mha.out_proj.bias)

        return out, k, v

    def init_cache(self, num_layers):
        """Initialize empty KV cache."""
        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):
        """Process one token with KV cache.
        new_emb: [B, 1, d_model] — raw audio embedding (not projected)
        text_mask: [B, T_text] bool, True = padding
        Returns: predicted [B, 1, d_model], updated cache
        """
        x = self.audio_input_proj(new_emb)
        # Add positional encoding for this specific position
        x = x + self.pos_embed.pe[:, step_idx:step_idx + 1]

        for i, layer in enumerate(self.transformer.layers):
            # Self-attention with cache (post-norm)
            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 KV
            ca_out, cache['cross_k'][i], cache['cross_v'][i] = self._cached_cross_attn(
                layer.multihead_attn, x, text_emb, cache['cross_k'][i], cache['cross_v'][i],
                memory_mask=text_mask,
            )
            x = layer.norm2(x + layer.dropout2(ca_out))

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

        return self.output_proj(x), cache

# ─── Positional Encoding ────────────────────────────────────────────────────

class SinusoidalPositionalEncoding(nn.Module):
    def __init__(self, d_model, max_len=8192):
        super().__init__()
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        self.register_buffer("pe", pe.unsqueeze(0))  # [1, max_len, d_model]

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

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

class ResConvBlock(nn.Module):
    """Residual conv block for mel decoder."""

    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):
    """Decodes latent embeddings back to mel spectrograms.
    Improved architecture with residual blocks for sharper mel output."""

    def __init__(self, d_model=256, n_mels=100):
        super().__init__()
        hidden = d_model * 2  # 512 — more capacity

        self.proj = nn.Sequential(
            nn.Linear(d_model, hidden),
            nn.GELU(),
        )
        self.upsample = nn.Sequential(
            # 4x upsample (2x + 2x)
            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),
        )
        # Refinement + projection to mel bins
        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):
        """z: [B, T_down, d_model] → mel: [B, n_mels, T_up]"""
        x = self.proj(z)
        x = x.transpose(1, 2)  # [B, hidden, T_down]
        x = self.upsample(x)   # [B, hidden, T_up]
        return self.out(x)      # [B, n_mels, T_up]

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

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

    def __init__(self, resolutions=((512, 128, 512), (1024, 256, 1024), (2048, 512, 2048))):
        super().__init__()
        # Each resolution is (n_fft, hop_length, win_length)
        self.resolutions = resolutions

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

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

        pred_stft = torch.stft(pred_flat, n_fft, hop_length, win_length,
                               window=window, return_complex=True)
        target_stft = torch.stft(target_flat, n_fft, hop_length, win_length,
                                 window=window, return_complex=True)

        pred_mag = pred_stft.abs()
        target_mag = target_stft.abs()

        # Spectral convergence: Frobenius norm of difference / Frobenius norm of target
        sc_loss = torch.norm(target_mag - pred_mag, p="fro") / (torch.norm(target_mag, p="fro") + 1e-8)

        # Log magnitude loss
        log_mag_loss = F.l1_loss(
            torch.log(pred_mag + 1e-7),
            torch.log(target_mag + 1e-7),
        )

        return sc_loss + log_mag_loss

    def forward(self, pred, target):
        """pred, target: [B, n_mels, T]"""
        loss = 0.0
        for n_fft, hop, win in self.resolutions:
            # Skip if T is too short for this resolution
            if pred.shape[2] >= n_fft:
                loss = loss + self._stft_loss(pred, target, n_fft, hop, win)
        return loss / len(self.resolutions)

# ─── Full LeWM TTS Model ────────────────────────────────────────────────────

class LeWMTTS(nn.Module):
    """
    Complete LeWM TTS model.

    Training: text + mel → audio embeddings → predict next embedding
    Two losses:
        1. Next-embedding prediction: MSE(predicted[t], target[t+1])
        2. Gaussian regularizer: KL(q(z|x) || N(0, I))
    """

    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)
        self.kl_weight = config.get("kl_weight", 0.01)
        n_speakers = config.get("n_speakers", 1)

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

        # Mel decoder: reconstructs mel from latent embeddings
        # Needed so the latent space is grounded to actual audio
        self.mel_decoder = MelDecoder(d_model=d_model, n_mels=n_mels)
        self.recon_weight = config.get("recon_weight", 1.0)

        # Speaker embedding (for multi-speaker)
        self.n_speakers = n_speakers
        if n_speakers > 1:
            self.speaker_embed = nn.Embedding(n_speakers, d_model)
        else:
            self.speaker_embed = None

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

        self.config = config

    def _add_speaker(self, text_emb, speaker_id):
        """Add speaker embedding to text embeddings."""
        if self.speaker_embed is not None and speaker_id is not None:
            spk = self.speaker_embed(speaker_id).unsqueeze(1)  # [B, 1, d]
            text_emb = text_emb + spk
        return text_emb

    def forward(self, mel, text_tokens, mel_mask=None, text_mask=None, speaker_id=None):
        """
        Args:
            mel: [B, n_mels, T_mel]
            text_tokens: [B, T_text]
            mel_mask: [B, T_mel] bool
            text_mask: [B, T_text] bool
        Returns:
            loss_dict with 'prediction_loss', 'kl_loss', 'total_loss'
        """
        # Encode text + speaker conditioning
        text_emb = self.text_encoder(text_tokens, text_mask)
        text_emb = self._add_speaker(text_emb, speaker_id)

        # Encode audio → latent embeddings
        z, mu, logvar = self.audio_encoder(mel, mel_mask)

        # Predict next embedding: input z[:-1], target z[1:]
        input_emb = z[:, :-1]     # [B, T-1, d]
        target_emb = z[:, 1:]     # [B, T-1, d]

        # Adjust mask for downsampled sequence
        if mel_mask is not None:
            T_down = z.shape[1]
            ds_mask = mel_mask[:, ::self.audio_encoder.downsample_factor][:, :T_down]
            pred_mask = ds_mask[:, :-1]
        else:
            pred_mask = None

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

        # ─── Loss 1: Next-embedding prediction (MSE) ───
        # Mask out padding positions
        if pred_mask is not None:
            valid = ~pred_mask  # True = valid
            valid = valid.unsqueeze(-1)  # [B, T-1, 1]
            prediction_loss = F.mse_loss(
                predicted * valid, target_emb * valid, reduction="sum"
            ) / (valid.sum() * predicted.shape[-1] + 1e-8)
        else:
            prediction_loss = F.mse_loss(predicted, target_emb)

        # ─── Loss 2: Gaussian regularizer — KL(N(mu, sigma²) || N(0, I)) ───
        # KL = -0.5 * sum(1 + logvar - mu² - exp(logvar))
        if mel_mask is not None:
            T_down = mu.shape[1]
            ds_mask = mel_mask[:, ::self.audio_encoder.downsample_factor][:, :T_down]
            valid = (~ds_mask).unsqueeze(-1).float()
            kl = -0.5 * (1 + logvar - mu.pow(2) - logvar.exp())
            kl_loss = (kl * valid).sum() / (valid.sum() * mu.shape[-1] + 1e-8)
        else:
            kl_loss = -0.5 * torch.mean(1 + logvar - mu.pow(2) - logvar.exp())

        # ─── Loss 3: Mel reconstruction ───
        # Decode latents back to mel and compare with input
        mel_recon = self.mel_decoder(z)  # [B, n_mels, T_up]
        # Trim/pad to match input mel length
        T_mel = mel.shape[2]
        T_recon = mel_recon.shape[2]
        if T_recon > T_mel:
            mel_recon = mel_recon[:, :, :T_mel]
        elif T_recon < T_mel:
            mel = mel[:, :, :T_recon]
            if mel_mask is not None:
                mel_mask = mel_mask[:, :T_recon]

        if mel_mask is not None:
            valid_mel = (~mel_mask).unsqueeze(1).float()  # [B, 1, T]
            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: Multi-resolution spectral loss ───
        # Apply on valid (non-padded) regions only
        if mel_mask is not None:
            spectral_loss = self.spectral_loss(mel_recon * valid_mel, mel * valid_mel)
        else:
            spectral_loss = self.spectral_loss(mel_recon, mel)

        total_loss = (prediction_loss
                      + self.kl_weight * kl_loss
                      + self.recon_weight * recon_loss
                      + self.spectral_weight * spectral_loss)

        return {
            "total_loss": total_loss,
            "prediction_loss": prediction_loss,
            "kl_loss": kl_loss,
            "recon_loss": recon_loss,
            "spectral_loss": spectral_loss,
        }

    def encode_audio(self, mel):
        """Encode mel to latent embeddings (for inference)."""
        z, mu, _ = self.audio_encoder(mel)
        return mu  # Use mean at inference

    def predict_next(self, audio_emb, text_emb, text_mask=None):
        """Predict next audio embedding autoregressively."""
        predicted = self.predictor(audio_emb, text_emb, text_mask=text_mask)
        return predicted[:, -1:]  # Only the last prediction

    def init_ar_cache(self):
        """Initialize KV cache for autoregressive inference."""
        num_layers = len(self.predictor.transformer.layers)
        return self.predictor.init_cache(num_layers)

    def predict_next_cached(self, new_emb, text_emb, step_idx, cache, text_mask=None):
        """Predict next embedding using KV cache. O(1) per step instead of O(n)."""
        return self.predictor.inference_step(new_emb, text_emb, step_idx, cache, text_mask=text_mask)

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

def build_model(config=None):
    if config is None:
        config = {
            "d_model": 256,
            "nhead": 4,
            "n_mels": 100,
            "text_vocab_size": 256,
            "text_encoder_layers": 4,
            "audio_encoder_layers": 4,
            "predictor_layers": 6,
            "dropout": 0.1,
            "kl_weight": 0.01,
        }
    model = LeWMTTS(config)
    print(f"LeWM TTS model: {count_parameters(model)/1e6:.2f}M parameters")
    return model, config

+        loss = torch.tensor(0.0, device=pred.device)

-        return loss / len(self.resolutions)


Bash

python3 -c "
from model import build_model
model, config = build_model()

import torch
B, T_mel, T_text = 2, 400, 50
mel = torch.randn(B, 100, T_mel)
text = torch.randint(0, 256, (B, T_text))
mel_mask = torch.zeros(B, T_mel, dtype=torch.bool)
text_mask = torch.zeros(B, T_text, dtype=torch.bool)

losses = model(mel, text, mel_mask, text_mask)
for k, v in losses.items():
    print(f'  {k}: {v.item():.4f}')
print('Forward pass OK!')
"

Test model forward pass with fixed spectral loss


LeWM TTS model: 24.00M parameters
  total_loss: 2.3051
  prediction_loss: 1.5044
  kl_loss: 0.2520
  recon_loss: 0.7982
  spectral_loss: 0.0000
Forward pass OK!


Spectral loss is 0 because T=400, the decoder downsamples 4x to 100 then upsamples 4x back to 400, but the actual mel recon might be exactly 400. The issue is the resolutions — n_fft=512 > T for the per-mel-bin signals. Let me use smaller resolutions that work on mel-spectrogram temporal dimensions:


python3 -c "
from model import build_model
import torch
model, config = build_model()
mel = torch.randn(2, 100, 400)
text = torch.randint(0, 256, (2, 50))
z, mu, logvar = model.audio_encoder(mel)
print(f'z shape: {z.shape}')
mel_recon = model.mel_decoder(z)
print(f'mel_recon shape: {mel_recon.shape}')
"

Check mel shapes through encoder/decoder


LeWM TTS model: 24.00M parameters
z shape: torch.Size([2, 100, 256])
mel_recon shape: torch.Size([2, 100, 400])


OK so mel_recon is [B, 100, 400]. T=400. The spectral loss flattens to [B*100, 400] and does STFT. With n_fft=512 > T=400, it skips. I need to use resolutions appropriate for time series of length ~400. Let me use smaller FFT sizes:


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

Architecture:
  - TextEncoder: Character-level transformer for Hindi/Devanagari text
  - AudioEncoder: 1D CNN + Transformer that encodes mel spectrograms to embeddings
  - JEPAPredictor: Transformer that predicts next audio embedding given text + audio context
  - Two losses: next-embedding prediction (MSE) + Gaussian regularizer (KL to N(0,I))
"""

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

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

class TextEncoder(nn.Module):
    """Character-level transformer encoder for Hindi text."""

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

        # Character embedding (covers all Unicode Hindi chars via byte-level encoding)
        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):
        """
        Args:
            text_tokens: [B, T_text] long tensor of byte-level token IDs
            text_mask: [B, T_text] bool, True = padding
        Returns:
            text_emb: [B, T_text, d_model]
        """
        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)
        x = self.proj(x)
        return x

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

class AudioEncoder(nn.Module):
    """
    Encodes mel spectrograms into a sequence of audio embeddings.
    Uses 1D CNN for downsampling + Transformer for contextualization.
    """

    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

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

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

        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)

        # Project to embedding with mu and logvar for Gaussian regularizer
        self.proj_mu = nn.Linear(d_model, d_model)
        self.proj_logvar = nn.Linear(d_model, d_model)

    def forward(self, mel, mel_mask=None):
        """
        Args:
            mel: [B, n_mels, T_mel] — log mel spectrogram
            mel_mask: [B, T_mel] bool, True = padding (before downsampling)
        Returns:
            z: [B, T_down, d_model] — sampled latent embeddings
            mu: [B, T_down, d_model]
            logvar: [B, T_down, d_model]
        """
        # CNN encoding: [B, n_mels, T] -> [B, d_model, T//4]
        x = self.conv_pre(mel)
        x = x.transpose(1, 2)  # [B, T_down, d_model]

        # Downsample mask if provided
        if mel_mask is not None:
            # Approximate downsampled mask
            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)

        mu = self.proj_mu(x)
        logvar = self.proj_logvar(x)

        # Reparameterization trick
        if self.training:
            std = torch.exp(0.5 * logvar)
            eps = torch.randn_like(std)
            z = mu + eps * std
        else:
            z = mu

        return z, mu, logvar

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

class JEPAPredictor(nn.Module):
    """
    Predicts next audio embedding given:
    - Text context (from TextEncoder)
    - Previous audio embeddings (from AudioEncoder)

    Uses cross-attention to condition on text, and causal self-attention
    over the audio embedding sequence.
    """

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

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

        # Transformer decoder with causal self-attention + cross-attention to text
        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 projection to predict next embedding
        self.output_proj = nn.Sequential(
            nn.Linear(d_model, d_model),
            nn.GELU(),
            nn.Linear(d_model, d_model),
        )

    def forward(self, audio_emb, text_emb, audio_mask=None, text_mask=None):
        """
        Args:
            audio_emb: [B, T_audio, d_model] — audio embeddings (from encoder)
            text_emb: [B, T_text, d_model] — text embeddings
            audio_mask: [B, T_audio] bool, True = padding
            text_mask: [B, T_text] bool, True = padding
        Returns:
            predicted: [B, T_audio, d_model] — predicted next embeddings
        """
        x = self.audio_input_proj(audio_emb)
        x = self.pos_embed(x)
        x = self.dropout(x)

        # Causal mask: each position can only attend to previous positions
        T = x.shape[1]
        causal_mask = torch.triu(
            torch.ones(T, T, device=x.device, dtype=torch.bool), diagonal=1
        )

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

        predicted = self.output_proj(predicted)
        return predicted

    def _cached_mha(self, mha, q_input, kv_input, cache_k, cache_v):
        """Run multi-head attention with KV cache.
        q_input: [B, 1, d], kv_input: [B, 1, d] (new token for K,V)
        cache_k, cache_v: [B, T_prev, d] projected keys/values
        Returns: output [B, 1, d], updated cache_k, updated cache_v
        """
        d = mha.embed_dim
        nhead = mha.num_heads
        head_dim = d // nhead
        B = q_input.shape[0]

        # Project Q from q_input, K/V from kv_input
        # in_proj_weight is [3d, d] = [Wq; Wk; Wv]
        Wq, Wk, Wv = mha.in_proj_weight.chunk(3, dim=0)
        bq, bk, bv = mha.in_proj_bias.chunk(3, dim=0)

        q = F.linear(q_input, Wq, bq)        # [B, 1, d]
        k_new = F.linear(kv_input, Wk, bk)    # [B, 1, d]
        v_new = F.linear(kv_input, Wv, bv)    # [B, 1, d]

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

        # Multi-head attention
        q = q.view(B, 1, nhead, head_dim).transpose(1, 2)      # [B, nhead, 1, head_dim]
        k_mh = k.view(B, -1, nhead, head_dim).transpose(1, 2)  # [B, nhead, T, head_dim]
        v_mh = v.view(B, -1, nhead, head_dim).transpose(1, 2)

        attn = torch.matmul(q, k_mh.transpose(-2, -1)) / (head_dim ** 0.5)
        attn = F.softmax(attn, dim=-1)
        out = torch.matmul(attn, v_mh)  # [B, nhead, 1, head_dim]

        out = out.transpose(1, 2).contiguous().view(B, 1, d)
        out = F.linear(out, mha.out_proj.weight, mha.out_proj.bias)

        return out, k, v

    def _cached_cross_attn(self, mha, q_input, memory, cache_k, cache_v, memory_mask=None):
        """Cross-attention with cached memory K,V (computed once).
        memory_mask: [B, T_text] bool, True = padding (applied as -inf)."""
        d = mha.embed_dim
        nhead = mha.num_heads
        head_dim = d // nhead
        B = q_input.shape[0]

        Wq, Wk, Wv = mha.in_proj_weight.chunk(3, dim=0)
        bq, bk, bv = mha.in_proj_bias.chunk(3, dim=0)

        q = F.linear(q_input, Wq, bq)  # [B, 1, d]

        if cache_k is None:
            k = F.linear(memory, Wk, bk)  # [B, T_text, d]
            v = F.linear(memory, Wv, bv)
        else:
            k, v = cache_k, cache_v

        q = q.view(B, 1, nhead, head_dim).transpose(1, 2)
        k_mh = k.view(B, -1, nhead, head_dim).transpose(1, 2)
        v_mh = v.view(B, -1, nhead, head_dim).transpose(1, 2)

        attn = torch.matmul(q, k_mh.transpose(-2, -1)) / (head_dim ** 0.5)
        if memory_mask is not None:
            # memory_mask: [B, T_text] → [B, 1, 1, T_text]
            attn = attn.masked_fill(memory_mask.unsqueeze(1).unsqueeze(2), float('-inf'))
        attn = F.softmax(attn, dim=-1)
        out = torch.matmul(attn, v_mh)

        out = out.transpose(1, 2).contiguous().view(B, 1, d)
        out = F.linear(out, mha.out_proj.weight, mha.out_proj.bias)

        return out, k, v

    def init_cache(self, num_layers):
        """Initialize empty KV cache."""
        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):
        """Process one token with KV cache.
        new_emb: [B, 1, d_model] — raw audio embedding (not projected)
        text_mask: [B, T_text] bool, True = padding
        Returns: predicted [B, 1, d_model], updated cache
        """
        x = self.audio_input_proj(new_emb)
        # Add positional encoding for this specific position
        x = x + self.pos_embed.pe[:, step_idx:step_idx + 1]

        for i, layer in enumerate(self.transformer.layers):
            # Self-attention with cache (post-norm)
            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 KV
            ca_out, cache['cross_k'][i], cache['cross_v'][i] = self._cached_cross_attn(
                layer.multihead_attn, x, text_emb, cache['cross_k'][i], cache['cross_v'][i],
                memory_mask=text_mask,
            )
            x = layer.norm2(x + layer.dropout2(ca_out))

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

        return self.output_proj(x), cache

# ─── Positional Encoding ────────────────────────────────────────────────────

class SinusoidalPositionalEncoding(nn.Module):
    def __init__(self, d_model, max_len=8192):
        super().__init__()
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        self.register_buffer("pe", pe.unsqueeze(0))  # [1, max_len, d_model]

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

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

class ResConvBlock(nn.Module):
    """Residual conv block for mel decoder."""

    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):
    """Decodes latent embeddings back to mel spectrograms.
    Improved architecture with residual blocks for sharper mel output."""

    def __init__(self, d_model=256, n_mels=100):
        super().__init__()
        hidden = d_model * 2  # 512 — more capacity

        self.proj = nn.Sequential(
            nn.Linear(d_model, hidden),
            nn.GELU(),
        )
        self.upsample = nn.Sequential(
            # 4x upsample (2x + 2x)
            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),
        )
        # Refinement + projection to mel bins
        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):
        """z: [B, T_down, d_model] → mel: [B, n_mels, T_up]"""
        x = self.proj(z)
        x = x.transpose(1, 2)  # [B, hidden, T_down]
        x = self.upsample(x)   # [B, hidden, T_up]
        return self.out(x)      # [B, n_mels, T_up]

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

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

    def __init__(self, resolutions=((512, 128, 512), (1024, 256, 1024), (2048, 512, 2048))):
        super().__init__()
        # Each resolution is (n_fft, hop_length, win_length)
        self.resolutions = resolutions

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

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

        pred_stft = torch.stft(pred_flat, n_fft, hop_length, win_length,
                               window=window, return_complex=True)
        target_stft = torch.stft(target_flat, n_fft, hop_length, win_length,
                                 window=window, return_complex=True)

        pred_mag = pred_stft.abs()
        target_mag = target_stft.abs()

        # Spectral convergence: Frobenius norm of difference / Frobenius norm of target
        sc_loss = torch.norm(target_mag - pred_mag, p="fro") / (torch.norm(target_mag, p="fro") + 1e-8)

        # Log magnitude loss
        log_mag_loss = F.l1_loss(
            torch.log(pred_mag + 1e-7),
            torch.log(target_mag + 1e-7),
        )

        return sc_loss + log_mag_loss

    def forward(self, pred, target):
        """pred, target: [B, n_mels, T]"""
        loss = torch.tensor(0.0, device=pred.device)
        count = 0
        for n_fft, hop, win in self.resolutions:
            # Skip if T is too short for this resolution
            if pred.shape[2] >= n_fft:
                loss = loss + self._stft_loss(pred, target, n_fft, hop, win)
                count += 1
        if count > 0:
            loss = loss / count
        return loss

# ─── Full LeWM TTS Model ────────────────────────────────────────────────────

class LeWMTTS(nn.Module):
    """
    Complete LeWM TTS model.

    Training: text + mel → audio embeddings → predict next embedding
    Two losses:
        1. Next-embedding prediction: MSE(predicted[t], target[t+1])
        2. Gaussian regularizer: KL(q(z|x) || N(0, I))
    """

    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)
        self.kl_weight = config.get("kl_weight", 0.01)
        n_speakers = config.get("n_speakers", 1)

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

        # Mel decoder: reconstructs mel from latent embeddings
        # Needed so the latent space is grounded to actual audio
        self.mel_decoder = MelDecoder(d_model=d_model, n_mels=n_mels)
        self.recon_weight = config.get("recon_weight", 1.0)

        # Speaker embedding (for multi-speaker)
        self.n_speakers = n_speakers
        if n_speakers > 1:
            self.speaker_embed = nn.Embedding(n_speakers, d_model)
        else:
            self.speaker_embed = None

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

        self.config = config

    def _add_speaker(self, text_emb, speaker_id):
        """Add speaker embedding to text embeddings."""
        if self.speaker_embed is not None and speaker_id is not None:
            spk = self.speaker_embed(speaker_id).unsqueeze(1)  # [B, 1, d]
            text_emb = text_emb + spk
        return text_emb

    def forward(self, mel, text_tokens, mel_mask=None, text_mask=None, speaker_id=None):
        """
        Args:
            mel: [B, n_mels, T_mel]
            text_tokens: [B, T_text]
            mel_mask: [B, T_mel] bool
            text_mask: [B, T_text] bool
        Returns:
            loss_dict with 'prediction_loss', 'kl_loss', 'total_loss'
        """
        # Encode text + speaker conditioning
        text_emb = self.text_encoder(text_tokens, text_mask)
        text_emb = self._add_speaker(text_emb, speaker_id)

        # Encode audio → latent embeddings
        z, mu, logvar = self.audio_encoder(mel, mel_mask)

        # Predict next embedding: input z[:-1], target z[1:]
        input_emb = z[:, :-1]     # [B, T-1, d]
        target_emb = z[:, 1:]     # [B, T-1, d]

        # Adjust mask for downsampled sequence
        if mel_mask is not None:
            T_down = z.shape[1]
            ds_mask = mel_mask[:, ::self.audio_encoder.downsample_factor][:, :T_down]
            pred_mask = ds_mask[:, :-1]
        else:
            pred_mask = None

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

        # ─── Loss 1: Next-embedding prediction (MSE) ───
        # Mask out padding positions
        if pred_mask is not None:
            valid = ~pred_mask  # True = valid
            valid = valid.unsqueeze(-1)  # [B, T-1, 1]
            prediction_loss = F.mse_loss(
                predicted * valid, target_emb * valid, reduction="sum"
            ) / (valid.sum() * predicted.shape[-1] + 1e-8)
        else:
            prediction_loss = F.mse_loss(predicted, target_emb)

        # ─── Loss 2: Gaussian regularizer — KL(N(mu, sigma²) || N(0, I)) ───
        # KL = -0.5 * sum(1 + logvar - mu² - exp(logvar))
        if mel_mask is not None:
            T_down = mu.shape[1]
            ds_mask = mel_mask[:, ::self.audio_encoder.downsample_factor][:, :T_down]
            valid = (~ds_mask).unsqueeze(-1).float()
            kl = -0.5 * (1 + logvar - mu.pow(2) - logvar.exp())
            kl_loss = (kl * valid).sum() / (valid.sum() * mu.shape[-1] + 1e-8)
        else:
            kl_loss = -0.5 * torch.mean(1 + logvar - mu.pow(2) - logvar.exp())

        # ─── Loss 3: Mel reconstruction ───
        # Decode latents back to mel and compare with input
        mel_recon = self.mel_decoder(z)  # [B, n_mels, T_up]
        # Trim/pad to match input mel length
        T_mel = mel.shape[2]
        T_recon = mel_recon.shape[2]
        if T_recon > T_mel:
            mel_recon = mel_recon[:, :, :T_mel]
        elif T_recon < T_mel:
            mel = mel[:, :, :T_recon]
            if mel_mask is not None:
                mel_mask = mel_mask[:, :T_recon]

        if mel_mask is not None:
            valid_mel = (~mel_mask).unsqueeze(1).float()  # [B, 1, T]
            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: Multi-resolution spectral loss ───
        # Apply on valid (non-padded) regions only
        if mel_mask is not None:
            spectral_loss = self.spectral_loss(mel_recon * valid_mel, mel * valid_mel)
        else:
            spectral_loss = self.spectral_loss(mel_recon, mel)

        total_loss = (prediction_loss
                      + self.kl_weight * kl_loss
                      + self.recon_weight * recon_loss
                      + self.spectral_weight * spectral_loss)

        return {
            "total_loss": total_loss,
            "prediction_loss": prediction_loss,
            "kl_loss": kl_loss,
            "recon_loss": recon_loss,
            "spectral_loss": spectral_loss,
        }

    def encode_audio(self, mel):
        """Encode mel to latent embeddings (for inference)."""
        z, mu, _ = self.audio_encoder(mel)
        return mu  # Use mean at inference

    def predict_next(self, audio_emb, text_emb, text_mask=None):
        """Predict next audio embedding autoregressively."""
        predicted = self.predictor(audio_emb, text_emb, text_mask=text_mask)
        return predicted[:, -1:]  # Only the last prediction

    def init_ar_cache(self):
        """Initialize KV cache for autoregressive inference."""
        num_layers = len(self.predictor.transformer.layers)
        return self.predictor.init_cache(num_layers)

    def predict_next_cached(self, new_emb, text_emb, step_idx, cache, text_mask=None):
        """Predict next embedding using KV cache. O(1) per step instead of O(n)."""
        return self.predictor.inference_step(new_emb, text_emb, step_idx, cache, text_mask=text_mask)

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

def build_model(config=None):
    if config is None:
        config = {
            "d_model": 256,
            "nhead": 4,
            "n_mels": 100,
            "text_vocab_size": 256,
            "text_encoder_layers": 4,
            "audio_encoder_layers": 4,
            "predictor_layers": 6,
            "dropout": 0.1,
            "kl_weight": 0.01,
        }
    model = LeWMTTS(config)
    print(f"LeWM TTS model: {count_parameters(model)/1e6:.2f}M parameters")
    return model, config


python3 -c "
from model import build_model
import torch
model, config = build_model()
mel = torch.randn(2, 100, 400)
text = torch.randint(0, 256, (2, 50))
mel_mask = torch.zeros(2, 400, dtype=torch.bool)
text_mask = torch.zeros(2, 50, dtype=torch.bool)
losses = model(mel, text, mel_mask, text_mask)
for k, v in losses.items():
    print(f'  {k}: {v.item():.4f}')
print('Forward pass OK!')
"

Test with smaller STFT resolutions


LeWM TTS model: 24.00M parameters
  total_loss: 5.6785
  prediction_loss: 1.5380
  kl_loss: 0.2525
  recon_loss: 0.7981
  spectral_loss: 6.6796
Forward pass OK!


All losses working. Spectral loss is active at 6.68 (weighted by 0.5 → contributes ~3.34). Let me verify training works end-to-end with a quick dry run:


Check if 100mel dataset exists


python3 train.py --data_dir /home/ubuntu/lewm-tts/processed_data_100mel --output_dir /home/ubuntu/lewm-tts/output_v3 --epochs 1 --batch_size 16 --print_every 5 --save_every 1 --d_model 256 --nhead 4 2>&1 | head -30

Quick 1-epoch test to verify training works with new losses


Device: cuda
LeWM TTS model: 24.00M parameters
Dataset: 6959 samples loaded
Training: 6959 samples, 434 batches/epoch
  [1/1] step 5 | loss=6.4499 pred=1.4911 kl=0.2487 recon=2.2780 spec=5.3592 | grad=2.94 | 1.4s
  [1/1] step 10 | loss=6.2573 pred=1.3466 kl=0.2221 recon=2.3383 spec=5.1404 | grad=4.16 | 1.8s
  [1/1] step 15 | loss=6.4845 pred=1.1830 kl=0.1898 recon=2.5779 spec=5.4413 | grad=2.80 | 2.3s
  [1/1] step 20 | loss=6.0556 pred=1.0229 kl=0.1660 recon=2.4029 spec=5.2523 | grad=1.75 | 2.7s
  [1/1] step 25 | loss=5.7989 pred=0.8691 kl=0.1524 recon=2.3760 spec=5.0993 | grad=1.37 | 3.1s
  [1/1] step 30 | loss=5.7088 pred=0.7448 kl=0.1543 recon=2.3491 spec=5.2194 | grad=1.75 | 3.5s
  [1/1] step 35 | loss=5.3453 pred=0.6485 kl=0.1704 recon=2.3637 spec=4.6528 | grad=1.32 | 3.9s
  [1/1] step 40 | loss=5.2644 pred=0.5785 kl=0.2062 recon=2.4291 spec=4.4951 | grad=1.55 | 4.3s
  [1/1] step 45 | loss=4.5579 pred=0.5260 kl=0.2699 recon=2.1934 spec=3.6495 | grad=1.45 | 4.6s
  [1/1] step 50 | loss=4.6456 pred=0.4669 kl=0.3165 recon=2.8444 spec=2.6330 | grad=11.23 | 5.0s
  [1/1] step 55 | loss=4.0946 pred=0.4783 kl=0.2868 recon=2.1935 spec=2.8100 | grad=1.75 | 5.4s
  [1/1] step 60 | loss=3.8696 pred=0.4071 kl=0.3072 recon=2.1160 spec=2.6513 | grad=4.70 | 5.8s
  [1/1] step 65 | loss=3.6254 pred=0.3757 kl=0.3246 recon=1.9950 spec=2.4614 | grad=1.95 | 6.2s
  [1/1] step 70 | loss=3.2655 pred=0.3471 kl=0.3327 recon=1.8046 spec=2.1748 | grad=0.98 | 6.5s
  [1/1] step 75 | loss=2.8864 pred=0.3455 kl=0.3442 recon=1.6376 spec=1.7480 | grad=3.05 | 6.9s
  [1/1] step 80 | loss=2.9153 pred=0.3879 kl=0.3523 recon=1.5726 spec=1.8457 | grad=7.44 | 7.3s
  [1/1] step 85 | loss=2.5405 pred=0.3190 kl=0.3895 recon=1.3966 spec=1.5743 | grad=2.34 | 7.7s
  [1/1] step 90 | loss=2.4685 pred=0.2897 kl=0.4117 recon=1.4136 spec=1.4481 | grad=2.24 | 8.1s
  [1/1] step 95 | loss=2.3366 pred=0.2846 kl=0.4268 recon=1.3787 spec=1.2610 | grad=3.30 | 8.4s
  [1/1] step 100 | loss=2.2200 pred=0.2681 kl=0.4484 recon=1.3030 spec=1.2081 | grad=3.11 | 8.8s
  [1/1] step 105 | loss=2.2483 pred=0.2519 kl=0.4701 recon=1.3501 spec=1.1986 | grad=3.06 | 9.2s
  [1/1] step 110 | loss=2.2084 pred=0.2476 kl=0.4876 recon=1.2941 spec=1.2358 | grad=4.60 | 9.6s
  [1/1] step 115 | loss=2.1776 pred=0.2393 kl=0.5098 recon=1.3280 spec=1.1187 | grad=1.56 | 10.0s
  [1/1] step 120 | loss=2.0550 pred=0.2175 kl=0.5330 recon=1.2304 spec=1.1076 | grad=1.92 | 10.4s
  [1/1] step 125 | loss=2.0460 pred=0.2184 kl=0.5501 recon=1.2295 spec=1.0862 | grad=5.77 | 10.7s
  [1/1] step 130 | loss=2.0484 pred=0.2089 kl=0.5675 recon=1.2351 spec=1.0951 | grad=7.77 | 11.1s


Training works perfectly. All losses are decreasing well — spectral loss drops from 5.36 → 1.09 in just 130 steps. Now let me start full training:


Clean up test run output


(Bash completed with no output)


nohup python3 train.py \
  --data_dir /home/ubuntu/lewm-tts/processed_data_100mel \
  --output_dir /home/ubuntu/lewm-tts/output_v3 \
  --epochs 200 \
  --batch_size 256 \
  --d_model 256 \
  --nhead 4 \
  --text_layers 4 \
  --audio_layers 4 \
  --predictor_layers 6 \
  --lr 1e-4 \
  --kl_weight 0.1 \
  --recon_weight 1.0 \
  --spectral_weight 0.5 \
  --grad_clip 1.0 \
  --num_workers 4 \
  --log_every 10 \
  --print_every 50 \
  --save_every 10 \
  > /home/ubuntu/lewm-tts/train_v3.log 2>&1 &
echo "PID: $!"

Start full training with upgraded mel decoder + spectral loss


PID: 3518002


sleep 5 && tail -20 /home/ubuntu/lewm-tts/train_v3.log

Check training started successfully


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


Training is running. 27 batches/epoch with batch_size=256. Let me check after a bit more progress:


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

Check training progress


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

Check training after more time


Device: cuda
LeWM TTS model: 24.00M parameters
Dataset: 6959 samples loaded
Training: 6959 samples, 27 batches/epoch
Epoch 1/200 | avg_loss=6.4157 pred=1.5698 kl=0.2692 recon=2.3557 spec=4.9796 | 27.5s
  New best model (loss=6.4157)
  [2/200] step 50 | loss=5.9234 pred=1.1925 kl=0.2006 recon=2.3602 spec=4.7394 | grad=1.48 | 22.9s
Epoch 2/200 | avg_loss=6.1036 pred=1.3063 kl=0.2213 recon=2.3550 spec=4.8829 | 26.7s
  New best model (loss=6.1036)


Training is progressing well. Let me update the memory file and let it run:


     1→# LeWM TTS Project
     2→
     3→## Overview
     4→Adapting LeWM (LeWorldModel) JEPA architecture for Text-to-Speech synthesis.
     5→- Paper: https://le-wm.github.io/
     6→- Project folder: `/home/ubuntu/lewm-tts`
     7→- Goal: 10-20 speaker multi-speaker TTS on this 14M model
     8→
     9→## Architecture (Implemented)
    10→- **TextEncoder**: Byte-level char transformer (4 layers, 256d, 4 heads)
    11→- **AudioEncoder**: 1D CNN (stride 4 downsample) + Transformer (4 layers) → mu/logvar → reparameterization
    12→- **JEPAPredictor**: Transformer decoder (6 layers) with causal self-attn + cross-attn to text + **KV-cache**
    13→- **MelDecoder**: Transposed-conv upsampler (4x) → reconstructs mel from latents
    14→- **Vocoder**: Griffin-Lim → Vocos refinement (NOT raw Vocos — model outputs 80 mels, Vocos expects 100)
    15→- **3 losses**: prediction (MSE), KL (Gaussian regularizer), reconstruction (L1 mel)
    16→- **Model size**: 14.50M parameters
    17→- **KL annealing**: ramps up over first 20% of training
    18→
    19→## Files
    20→- `preprocess.py` — segments audio, extracts mels, saves manifest
    21→- `model.py` — TextEncoder, AudioEncoder, JEPAPredictor (with KV-cache), MelDecoder, LeWMTTS
    22→- `dataset.py` — TTSDataset, collate_fn, build_dataloader
    23→- `train.py` — training loop with logging, checkpointing, KL annealing
    24→- `inference.py` — AR generation with GL→Vocos vocoder (single request)
    25→- `inference_engine.py` — Production engine: batched inference, KV-cache, FP16, voice prompts
    26→
    27→## Dataset
    28→- **Raw**: `/home/ubuntu/modi_dataset` (2,330 clips, 18.25h, Hindi, 24kHz)
    29→- **Processed**: `/home/ubuntu/lewm-tts/processed_data` (7,009 segments, 17.88h)
    30→- Mel config: n_fft=1024, hop=256, n_mels=80, sr=24000
    31→
    32→## Vocoder Notes (IMPORTANT)
    33→- Model outputs 80-mel, Vocos `charactr/vocos-mel-24khz` expects 100-mel
    34→- Cannot pad 80→100 with zeros (produces noise)
    35→- Cannot interpolate 80→100 (produces quiet/bad audio)
    36→- Correct pipeline: Griffin-Lim (80-mel → waveform) → Vocos refinement (audio → 100-mel → decode)
    37→- Original good files: `output/vocos_predict_tf.wav` etc. (normalized to max=1.0)
    38→
    39→## Inference Findings
    40→- **Predict (teacher-forced)**: works well, good quality (a bit robotic at early epochs)
    41→- **AR synthesis**: broken with random start token — model never saw random noise during training
    42→- **Audio prompt fixes AR**: seeding with real encoder embeddings gives proper generation
    43→- **Best early checkpoints**: v1 epoch 70-100 (later epochs collapse for AR due to norm stopping)
    44→- v2 (DAC-based) in `/home/ubuntu/lewm-tts/v2/`
    45→
    46→## Inference Performance (A100 80GB)
    47→- KV-cache: 9x real-time single request
    48→- Batch=32: 185x throughput
    49→- Batch=64: 359x throughput
    50→- Batch=128: OOM
    51→
    52→## Next Steps for Multi-Speaker
    53→- Audio prompt = voice cloning (reference clip → encoder → seed AR generation)
    54→- Scheduled sampling during training (robustify AR)
    55→- Learned start token (alternative to audio prompt)
    56→- Train on 10-20 speaker combined dataset
    57→

# LeWM TTS Project

## Overview
Adapting LeWM (LeWorldModel) JEPA architecture for Text-to-Speech synthesis.
- Paper: https://le-wm.github.io/
- Project folder: `/home/ubuntu/lewm-tts`
- Goal: 10-20 speaker multi-speaker TTS on this 14M model

## Architecture (Implemented)
- **TextEncoder**: Byte-level char transformer (4 layers, 256d, 4 heads)
- **AudioEncoder**: 1D CNN (stride 4 downsample) + Transformer (4 layers) → mu/logvar → reparameterization
- **JEPAPredictor**: Transformer decoder (6 layers) with causal self-attn + cross-attn to text + **KV-cache**
- **MelDecoder**: Transposed-conv upsampler (4x) → reconstructs mel from latents
- **Vocoder**: Griffin-Lim → Vocos refinement (NOT raw Vocos — model outputs 80 mels, Vocos expects 100)
- **3 losses**: prediction (MSE), KL (Gaussian regularizer), reconstruction (L1 mel)
- **Model size**: 14.50M parameters
- **KL annealing**: ramps up over first 20% of training

## Files
- `preprocess.py` — segments audio, extracts mels, saves manifest
- `model.py` — TextEncoder, AudioEncoder, JEPAPredictor (with KV-cache), MelDecoder, LeWMTTS
- `dataset.py` — TTSDataset, collate_fn, build_dataloader
- `train.py` — training loop with logging, checkpointing, KL annealing
- `inference.py` — AR generation with GL→Vocos vocoder (single request)
- `inference_engine.py` — Production engine: batched inference, KV-cache, FP16, voice prompts

## Dataset
- **Raw**: `/home/ubuntu/modi_dataset` (2,330 clips, 18.25h, Hindi, 24kHz)
- **Processed**: `/home/ubuntu/lewm-tts/processed_data` (7,009 segments, 17.88h)
- Mel config: n_fft=1024, hop=256, n_mels=80, sr=24000

## Vocoder Notes (IMPORTANT)
- Model outputs 80-mel, Vocos `charactr/vocos-mel-24khz` expects 100-mel
- Cannot pad 80→100 with zeros (produces noise)
- Cannot interpolate 80→100 (produces quiet/bad audio)
- Correct pipeline: Griffin-Lim (80-mel → waveform) → Vocos refinement (audio → 100-mel → decode)
- Original good files: `output/vocos_predict_tf.wav` etc. (normalized to max=1.0)

## Inference Findings
- **Predict (teacher-forced)**: works well, good quality (a bit robotic at early epochs)
- **AR synthesis**: broken with random start token — model never saw random noise during training
- **Audio prompt fixes AR**: seeding with real encoder embeddings gives proper generation
- **Best early checkpoints**: v1 epoch 70-100 (later epochs collapse for AR due to norm stopping)
- v2 (DAC-based) in `/home/ubuntu/lewm-tts/v2/`

## Inference Performance (A100 80GB)
- KV-cache: 9x real-time single request
- Batch=32: 185x throughput
- Batch=64: 359x throughput
- Batch=128: OOM

## Next Steps for Multi-Speaker
- Audio prompt = voice cloning (reference clip → encoder → seed AR generation)
- Scheduled sampling during training (robustify AR)
- Learned start token (alternative to audio prompt)
- Train on 10-20 speaker combined dataset


Write

# LeWM TTS Project

## Overview
Adapting LeWM (LeWorldModel) JEPA architecture for Text-to-Speech synthesis.
- Paper: https://le-wm.github.io/
- Project folder: `/home/ubuntu/lewm-tts`
- Goal: 10-20 speaker multi-speaker TTS on this 14M→24M model

## Architecture (v3 — Current)
- **TextEncoder**: Byte-level char transformer (4 layers, 256d, 4 heads)
- **AudioEncoder**: 1D CNN (stride 4 downsample) + Transformer (4 layers) → mu/logvar → reparameterization
- **JEPAPredictor**: Transformer decoder (6 layers) with causal self-attn + cross-attn to text + **KV-cache**
- **MelDecoder**: Upgraded — Linear→512 + 2x ConvTranspose1d + ResConvBlock residual blocks + refinement conv
- **Vocoder**: Direct Vocos decode (model outputs 100 mels matching `charactr/vocos-mel-24khz`)
- **4 losses**: prediction (MSE), KL (Gaussian reg), reconstruction (L1 mel), **multi-resolution spectral loss**
- **Speaker embedding**: support for n_speakers > 1 (not yet trained)
- **Model size**: 24.00M parameters
- **KL annealing**: ramps up over first 20% of training

## Multi-Resolution Spectral Loss
- STFT at 3 resolutions: (64,16,64), (128,32,128), (256,64,256)
- Applied per mel-bin (flattens [B, n_mels, T] → [B*n_mels, T])
- Spectral convergence + log magnitude L1
- Weight: 0.5 (configurable via --spectral_weight)

## Files
- `preprocess.py` — segments audio, extracts mels, saves manifest
- `model.py` — TextEncoder, AudioEncoder, JEPAPredictor (with KV-cache), MelDecoder (upgraded), MultiResolutionSpectralLoss, LeWMTTS
- `dataset.py` — TTSDataset, collate_fn, build_dataloader
- `train.py` — training loop with logging, checkpointing, KL annealing, spectral loss
- `inference.py` — AR generation with Vocos vocoder (single request)
- `inference_engine.py` — Production engine: batched inference, KV-cache, FP16, voice prompts

## Dataset
- **Raw**: `/home/ubuntu/modi_dataset` (2,330 clips, 18.25h, Hindi, 24kHz)
- **100-mel processed**: `/home/ubuntu/lewm-tts/processed_data_100mel` (7,009 segments, 17.88h)
- Mel config: n_fft=1024, hop=256, n_mels=100, sr=24000, f_max=None, power=1.0, clip_val=1e-7

## Training History
- **v1**: 80 mels, simple decoder, GL→Vocos pipeline. Rustic sound.
- **v2**: DAC-based, checkpoints preserved at `/home/ubuntu/lewm-tts/v2/output/checkpoints/`
- **v3** (current): 100 mels, upgraded MelDecoder + spectral loss, direct Vocos. Training at output_v3/.

## Inference Findings
- **Predict (teacher-forced)**: works well, good quality
- **AR synthesis**: broken with random start token — needs audio prompt (voice cloning)
- **Best early checkpoints**: epoch 70-100 (later epochs collapse for AR)
- Noise tolerance test: GT mel + noise at L1=0.3 sounds natural; model's structured error sounds worse

## Inference Performance (A100 80GB)
- KV-cache: 9x real-time single request
- Batch=64: 359x throughput

## Next Steps
- Wait for v3 training to converge, test quality
- If mel quality good with Vocos, proceed to multi-speaker training
- Scheduled sampling during training (robustify AR)


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

# LeWM TTS Project

## Overview
Adapting LeWM (LeWorldModel) JEPA architecture for Text-to-Speech synthesis.
- Paper: https://le-wm.github.io/
- Project folder: `/home/ubuntu/lewm-tts`
- Goal: 10-20 speaker multi-speaker TTS on this 14M→24M model

## Architecture (v3 — Current)
- **TextEncoder**: Byte-level char transformer (4 layers, 256d, 4 heads)
- **AudioEncoder**: 1D CNN (stride 4 downsample) + Transformer (4 layers) → mu/logvar → reparameterization
- **JEPAPredictor**: Transformer decoder (6 layers) with causal self-attn + cross-attn to text + **KV-cache**
- **MelDecoder**: Upgraded — Linear→512 + 2x ConvTranspose1d + ResConvBlock residual blocks + refinement conv
- **Vocoder**: Direct Vocos decode (model outputs 100 mels matching `charactr/vocos-mel-24khz`)
- **4 losses**: prediction (MSE), KL (Gaussian reg), reconstruction (L1 mel), **multi-resolution spectral loss**
- **Speaker embedding**: support for n_speakers > 1 (not yet trained)
- **Model size**: 24.00M parameters
- **KL annealing**: ramps up over first 20% of training

## Multi-Resolution Spectral Loss
- STFT at 3 resolutions: (64,16,64), (128,32,128), (256,64,256)
- Applied per mel-bin (flattens [B, n_mels, T] → [B*n_mels, T])
- Spectral convergence + log magnitude L1
- Weight: 0.5 (configurable via --spectral_weight)

## Files
- `preprocess.py` — segments audio, extracts mels, saves manifest
- `model.py` — TextEncoder, AudioEncoder, JEPAPredictor (with KV-cache), MelDecoder (upgraded), MultiResolutionSpectralLoss, LeWMTTS
- `dataset.py` — TTSDataset, collate_fn, build_dataloader
- `train.py` — training loop with logging, checkpointing, KL annealing, spectral loss
- `inference.py` — AR generation with Vocos vocoder (single request)
- `inference_engine.py` — Production engine: batched inference, KV-cache, FP16, voice prompts

## Dataset
- **Raw**: `/home/ubuntu/modi_dataset` (2,330 clips, 18.25h, Hindi, 24kHz)
- **100-mel processed**: `/home/ubuntu/lewm-tts/processed_data_100mel` (7,009 segments, 17.88h)
- Mel config: n_fft=1024, hop=256, n_mels=100, sr=24000, f_max=None, power=1.0, clip_val=1e-7

## Training History
- **v1**: 80 mels, simple decoder, GL→Vocos pipeline. Rustic sound.
- **v2**: DAC-based, checkpoints preserved at `/home/ubuntu/lewm-tts/v2/output/checkpoints/`
- **v3** (current): 100 mels, upgraded MelDecoder + spectral loss, direct Vocos. Training at output_v3/.

## Inference Findings
- **Predict (teacher-forced)**: works well, good quality
- **AR synthesis**: broken with random start token — needs audio prompt (voice cloning)
- **Best early checkpoints**: epoch 70-100 (later epochs collapse for AR)
- Noise tolerance test: GT mel + noise at L1=0.3 sounds natural; model's structured error sounds worse

## Inference Performance (A100 80GB)
- KV-cache: 9x real-time single request
- Batch=64: 359x throughput

## Next Steps
- Wait for v3 training to converge, test quality
- If mel quality good with Vocos, proceed to multi-speaker training
- Scheduled sampling during training (robustify AR)

-- Goal: 10-20 speaker multi-speaker TTS on this 14M model

+- Goal: 10-20 speaker multi-speaker TTS on this 14M→24M model

-- **MelDecoder**: Transposed-conv upsampler (4x) → reconstructs mel from latents

-- **Vocoder**: Griffin-Lim → Vocos refinement (NOT raw Vocos — model outputs 80 mels, Vocos expects 100)

-- **3 losses**: prediction (MSE), KL (Gaussian regularizer), reconstruction (L1 mel)

+- **MelDecoder**: Upgraded — Linear→512 + 2x ConvTranspose1d + ResConvBlock residual blocks + refinement conv

+- **Vocoder**: Direct Vocos decode (model outputs 100 mels matching `charactr/vocos-mel-24khz`)

+- **4 losses**: prediction (MSE), KL (Gaussian reg), reconstruction (L1 mel), **multi-resolution spectral loss**

+- **Speaker embedding**: support for n_speakers > 1 (not yet trained)

+- Applied per mel-bin (flattens [B, n_mels, T] → [B*n_mels, T])

+- Spectral convergence + log magnitude L1

+- Weight: 0.5 (configurable via --spectral_weight)

-- `model.py` — TextEncoder, AudioEncoder, JEPAPredictor (with KV-cache), MelDecoder, LeWMTTS

+- `model.py` — TextEncoder, AudioEncoder, JEPAPredictor (with KV-cache), MelDecoder (upgraded), MultiResolutionSpectralLoss, LeWMTTS

-- `train.py` — training loop with logging, checkpointing, KL annealing

-- `inference.py` — AR generation with GL→Vocos vocoder (single request)

+- `train.py` — training loop with logging, checkpointing, KL annealing, spectral loss

+- `inference.py` — AR generation with Vocos vocoder (single request)

-- **Processed**: `/home/ubuntu/lewm-tts/processed_data` (7,009 segments, 17.88h)

+- **100-mel processed**: `/home/ubuntu/lewm-tts/processed_data_100mel` (7,009 segments, 17.88h)

+- Mel config: n_fft=1024, hop=256, n_mels=100, sr=24000, f_max=None, power=1.0, clip_val=1e-7

-- Model outputs 80-mel, Vocos `charactr/vocos-mel-24khz` expects 100-mel

-- Cannot pad 80→100 with zeros (produces noise)

-- Cannot interpolate 80→100 (produces quiet/bad audio)

-- Correct pipeline: Griffin-Lim (80-mel → waveform) → Vocos refinement (audio → 100-mel → decode)

-- Original good files: `output/vocos_predict_tf.wav` etc. (normalized to max=1.0)

+- **v1**: 80 mels, simple decoder, GL→Vocos pipeline. Rustic sound.

+- **v2**: DAC-based, checkpoints preserved at `/home/ubuntu/lewm-tts/v2/output/checkpoints/`

+- **v3** (current): 100 mels, upgraded MelDecoder + spectral loss, direct Vocos. Training at output_v3/.

-- **Predict (teacher-forced)**: works well, good quality (a bit robotic at early epochs)

-- **AR synthesis**: broken with random start token — model never saw random noise during training

-- **Audio prompt fixes AR**: seeding with real encoder embeddings gives proper generation

-- **Best early checkpoints**: v1 epoch 70-100 (later epochs collapse for AR due to norm stopping)

-- v2 (DAC-based) in `/home/ubuntu/lewm-tts/v2/`

+- **Predict (teacher-forced)**: works well, good quality

+- **AR synthesis**: broken with random start token — needs audio prompt (voice cloning)

+- **Best early checkpoints**: epoch 70-100 (later epochs collapse for AR)

+- Noise tolerance test: GT mel + noise at L1=0.3 sounds natural; model's structured error sounds worse

-- Audio prompt = voice cloning (reference clip → encoder → seed AR generation)

+- Wait for v3 training to converge, test quality

+- If mel quality good with Vocos, proceed to multi-speaker training

-- Learned start token (alternative to audio prompt)

-- Train on 10-20 speaker combined dataset

# LeWM TTS Project

## Overview
Adapting LeWM (LeWorldModel) JEPA architecture for Text-to-Speech synthesis.
- Paper: https://le-wm.github.io/
- Project folder: `/home/ubuntu/lewm-tts`
- Goal: 10-20 speaker multi-speaker TTS on this 14M model

## Architecture (Implemented)
- **TextEncoder**: Byte-level char transformer (4 layers, 256d, 4 heads)
- **AudioEncoder**: 1D CNN (stride 4 downsample) + Transformer (4 layers) → mu/logvar → reparameterization
- **JEPAPredictor**: Transformer decoder (6 layers) with causal self-attn + cross-attn to text + **KV-cache**
- **MelDecoder**: Transposed-conv upsampler (4x) → reconstructs mel from latents
- **Vocoder**: Griffin-Lim → Vocos refinement (NOT raw Vocos — model outputs 80 mels, Vocos expects 100)
- **3 losses**: prediction (MSE), KL (Gaussian regularizer), reconstruction (L1 mel)
- **Model size**: 14.50M parameters
- **KL annealing**: ramps up over first 20% of training

## Files
- `preprocess.py` — segments audio, extracts mels, saves manifest
- `model.py` — TextEncoder, AudioEncoder, JEPAPredictor (with KV-cache), MelDecoder, LeWMTTS
- `dataset.py` — TTSDataset, collate_fn, build_dataloader
- `train.py` — training loop with logging, checkpointing, KL annealing
- `inference.py` — AR generation with GL→Vocos vocoder (single request)
- `inference_engine.py` — Production engine: batched inference, KV-cache, FP16, voice prompts

## Dataset
- **Raw**: `/home/ubuntu/modi_dataset` (2,330 clips, 18.25h, Hindi, 24kHz)
- **Processed**: `/home/ubuntu/lewm-tts/processed_data` (7,009 segments, 17.88h)
- Mel config: n_fft=1024, hop=256, n_mels=80, sr=24000

## Vocoder Notes (IMPORTANT)
- Model outputs 80-mel, Vocos `charactr/vocos-mel-24khz` expects 100-mel
- Cannot pad 80→100 with zeros (produces noise)
- Cannot interpolate 80→100 (produces quiet/bad audio)
- Correct pipeline: Griffin-Lim (80-mel → waveform) → Vocos refinement (audio → 100-mel → decode)
- Original good files: `output/vocos_predict_tf.wav` etc. (normalized to max=1.0)

## Inference Findings
- **Predict (teacher-forced)**: works well, good quality (a bit robotic at early epochs)
- **AR synthesis**: broken with random start token — model never saw random noise during training
- **Audio prompt fixes AR**: seeding with real encoder embeddings gives proper generation
- **Best early checkpoints**: v1 epoch 70-100 (later epochs collapse for AR due to norm stopping)
- v2 (DAC-based) in `/home/ubuntu/lewm-tts/v2/`

## Inference Performance (A100 80GB)
- KV-cache: 9x real-time single request
- Batch=32: 185x throughput
- Batch=64: 359x throughput
- Batch=128: OOM

## Next Steps for Multi-Speaker
- Audio prompt = voice cloning (reference clip → encoder → seed AR generation)
- Scheduled sampling during training (robustify AR)
- Learned start token (alternative to audio prompt)
- Train on 10-20 speaker combined dataset


Training is running (PID 3518002). Here's the summary:

**Done:**
- Upgraded MelDecoder with ResConvBlock residual layers (d_model×2 hidden = 512 channels)
- Added multi-resolution spectral loss (3 STFT resolutions: 64/128/256) — forces sharp mel structure instead of blurry averages
- Model now 24M params (up from 14.5M due to larger decoder)
- Updated train.py to log spectral loss
- Verified forward pass works — all 4 losses decreasing

**Training v3 running:**
- `output_v3/`, batch_size=256, 27 batches/epoch, 200 epochs
- Epoch 2: loss=6.10, pred=1.31, recon=2.35, spectral=4.88 (all decreasing nicely)
- Spectral loss already dropped from ~5.0 → trending down, which means the decoder is learning to produce sharper spectra


<task-notification>
<task-id>bvnjhsy1w</task-id>
<tool-use-id>toolu_vrtx_015W8xFokRRtfd8dxXSDyDQx</tool-use-id>
<output-file>/tmp/claude-1001/-home-ubuntu/1b244b10-1345-40d2-b03a-2de470236690/tasks/bvnjhsy1w.output</output-file>
<status>failed</status>
<summary>Background command "Start training with batch 256" failed with exit code 144</summary>
</task-notification>
Read the output file to retrieve the result: /tmp/claude-1001/-home-ubuntu/1b244b10-1345-40d2-b03a-2de470236690/tasks/bvnjhsy1w.output


Bash

Check training log status


Check if training process is still running


Device: cuda
LeWM TTS model: 24.00M parameters
Dataset: 6959 samples loaded
Training: 6959 samples, 27 batches/epoch
Epoch 1/200 | avg_loss=6.4157 pred=1.5698 kl=0.2692 recon=2.3557 spec=4.9796 | 27.5s
  New best model (loss=6.4157)
  [2/200] step 50 | loss=5.9234 pred=1.1925 kl=0.2006 recon=2.3602 spec=4.7394 | grad=1.48 | 22.9s
Epoch 2/200 | avg_loss=6.1036 pred=1.3063 kl=0.2213 recon=2.3550 spec=4.8829 | 26.7s
  New best model (loss=6.1036)
Epoch 3/200 | avg_loss=5.6631 pred=0.9758 kl=0.1740 recon=2.3543 spec=4.6637 | 26.8s
  New best model (loss=5.6631)


ubuntu   3518002  105  2.2 82121824 2780012 ?    Sl   05:57   1:39 python3 train.py --data_dir /home/ubuntu/lewm-tts/processed_data_100mel --output_dir /home/ubuntu/lewm-tts/output_v3 --epochs 200 --batch_size 256 --d_model 256 --nhead 4 --text_layers 4 --audio_layers 4 --predictor_layers 6 --lr 1e-4 --kl_weight 0.1 --recon_weight 1.0 --spectral_weight 0.5 --grad_clip 1.0 --num_workers 4 --log_every 10 --print_every 50 --save_every 10
ubuntu   3518848 12.0  1.0 82175608 1255316 ?    Sl   05:59   0:01 python3 train.py --data_dir /home/ubuntu/lewm-tts/processed_data_100mel --output_dir /home/ubuntu/lewm-tts/output_v3 --epochs 200 --batch_size 256 --d_model 256 --nhead 4 --text_layers 4 --audio_layers 4 --predictor_layers 6 --lr 1e-4 --kl_weight 0.1 --recon_weight 1.0 --spectral_weight 0.5 --grad_clip 1.0 --num_workers 4 --log_every 10 --print_every 50 --save_every 10
ubuntu   3518849 12.2  1.0 82173952 1253780 ?    Sl   05:59   0:01 python3 train.py --data_dir /home/ubuntu/lewm-tts/processed_data_100mel --output_dir /home/ubuntu/lewm-tts/output_v3 --epochs 200 --batch_size 256 --d_model 256 --nhead 4 --text_layers 4 --audio_layers 4 --predictor_layers 6 --lr 1e-4 --kl_weight 0.1 --recon_weight 1.0 --spectral_weight 0.5 --grad_clip 1.0 --num_workers 4 --log_every 10 --print_every 50 --save_every 10
ubuntu   3518850 12.2  1.0 82175964 1255648 ?    Sl   05:59   0:01 python3 train.py --data_dir /home/ubuntu/lewm-tts/processed_data_100mel --output_dir /home/ubuntu/lewm-tts/output_v3 --epochs 200 --batch_size 256 --d_model 256 --nhead 4 --text_layers 4 --audio_layers 4 --predictor_layers 6 --lr 1e-4 --kl_weight 0.1 --recon_weight 1.0 --spectral_weight 0.5 --grad_clip 1.0 --num_workers 4 --log_every 10 --print_every 50 --save_every 10
ubuntu   3518851 10.5  1.0 82175392 1255056 ?    Sl   05:59   0:01 python3 train.py --data_dir /home/ubuntu/lewm-tts/processed_data_100mel --output_dir /home/ubuntu/lewm-tts/output_v3 --epochs 200 --batch_size 256 --d_model 256 --nhead 4 --text_layers 4 --audio_layers 4 --predictor_layers 6 --lr 1e-4 --kl_weight 0.1 --recon_weight 1.0 --spectral_weight 0.5 --grad_clip 1.0 --num_workers 4 --log_every 10 --print_every 50 --save_every 10


Training is running fine (PID 3518002). The failed task notification is from a previous session's background task — ignore it.

**Current status** — Epoch 3/200, losses all decreasing:
- total: 6.42 → 5.66
- prediction: 1.57 → 0.98
- recon: 2.36 → 2.35
- spectral: 4.98 → 4.66

~27s per epoch, so full 200 epochs will take ~1.5 hours.


can u add .md file and update what we dng ..als also update ur md


The user wants me to create a markdown file documenting what we're doing in the project, and also update my memory file.


ToolSearch

Question: select:Write


Question: select:Write


Bash

Get latest training numbers for the doc


Epoch 3/200 | avg_loss=5.6631 pred=0.9758 kl=0.1740 recon=2.3543 spec=4.6637 | 26.8s
  New best model (loss=5.6631)
  [4/200] step 100 | loss=4.8707 pred=0.7100 kl=0.2010 recon=2.3181 spec=3.6816 | grad=1.11 | 18.9s
Epoch 4/200 | avg_loss=5.0403 pred=0.7430 kl=0.1915 recon=2.3460 spec=3.8993 | 26.7s
  New best model (loss=5.0403)


Write

# LeWM TTS

JEPA-based Text-to-Speech system adapted from [LeWorldModel](https://le-wm.github.io/).

## Architecture

A 24M parameter model with 4 core components:

```
Text (Hindi/UTF-8 bytes) ──> TextEncoder ──> text embeddings ──┐
                                                                ├──> JEPAPredictor ──> predicted embeddings
Audio (mel spectrogram) ──> AudioEncoder ──> audio embeddings ──┘            │
                                                                             v
                                                                        MelDecoder ──> mel spectrogram ──> Vocos ──> waveform
```

| Component | Description |
|-----------|-------------|
| **TextEncoder** | Byte-level character transformer (4 layers, 256d, 4 heads). Encodes Hindi text via UTF-8 bytes — no tokenizer needed. |
| **AudioEncoder** | 1D CNN (stride-4 downsample) + Transformer (4 layers). Outputs mu/logvar for Gaussian latent space with reparameterization. |
| **JEPAPredictor** | Transformer decoder (6 layers) with causal self-attention + cross-attention to text. Predicts next audio embedding autoregressively. Has KV-cache for fast inference. |
| **MelDecoder** | Upgraded architecture — Linear projection to 512d, 2x transposed convolution (4x upsample), residual conv blocks (ResConvBlock), and refinement layers. Outputs 100-mel spectrogram. |
| **Vocoder** | [Vocos](https://github.com/gemelo-ai/vocos) `charactr/vocos-mel-24khz` — directly decodes 100-mel to 24kHz waveform. No Griffin-Lim. |

### Speaker Embedding (planned)
Built-in support for multi-speaker via learned embedding table added to text embeddings. Not yet trained — will be used for 10-20 speaker multi-speaker TTS.

## Training Losses

4 losses, all active during training:

1. **Prediction Loss (MSE)** — next-embedding prediction: `MSE(predicted[t], target[t+1])`
2. **KL Loss** — Gaussian regularizer: `KL(q(z|x) || N(0, I))` with annealing over first 20% of training
3. **Reconstruction Loss (L1)** — mel decoder output vs ground truth mel spectrogram
4. **Multi-Resolution Spectral Loss** — STFT at 3 resolutions (64/128/256 FFT sizes), combines spectral convergence + log magnitude L1. Forces sharp spectral detail instead of blurry averages.

## Dataset

- **Source**: 20 hours of studio-quality single-speaker Hindi audio (Modi speeches)
- **Preprocessing**: silence-based segmentation into 3-12s chunks, mel extraction
- **Processed**: 7,009 segments, 17.88 hours
- **Mel config**: 100 mels, n_fft=1024, hop=256, sr=24000, f_max=None (full bandwidth), log-clamped at 1e-7 — matches Vocos `mel-24khz` exactly

## Files

| File | Purpose |
|------|---------|
| `model.py` | All model components: TextEncoder, AudioEncoder, JEPAPredictor, MelDecoder, MultiResolutionSpectralLoss, LeWMTTS |
| `train.py` | Training loop with KL annealing, gradient clipping, TensorBoard logging, checkpointing |
| `inference.py` | Single-request inference — AR synthesis, teacher-forced reconstruction |
| `inference_engine.py` | Production batched inference engine — KV-cache, FP16, voice prompts, benchmarking |
| `preprocess.py` | Data preprocessing — audio segmentation, mel extraction, manifest creation |
| `dataset.py` | PyTorch dataset + dataloader with dynamic batching and padding |

## Inference

### Modes
- **Synthesize**: Pure autoregressive generation from text. Requires audio prompt (voice cloning) for best results.
- **Reconstruct**: Encode real mel → decode back (tests encoder/decoder quality)
- **Predict (teacher-forced)**: Encode mel → predictor → decode (tests full pipeline without AR drift)

### Performance (A100 80GB)
- Single request with KV-cache: 9x real-time
- Batched (batch=64): 359x real-time throughput

### Voice Prompting
The inference engine supports audio prompts — a reference clip is encoded and used to seed autoregressive generation, effectively providing voice cloning without explicit speaker conditioning.

## Training History

### v1 — 80 mels + Griffin-Lim + Vocos
- Simple MelDecoder (2x transposed conv), 80-mel output
- Vocoder: Griffin-Lim → Vocos refinement pipeline
- Result: teacher-forced quality decent but "rustic" sound from GL artifacts

### v2 — DAC-based
- Experimental DAC codec variant
- Checkpoints preserved at `v2/output/checkpoints/`

### v3 — 100 mels + upgraded decoder + spectral loss (current)
- Direct 100-mel output matching Vocos natively (no GL middleman)
- Upgraded MelDecoder with ResConvBlock residual layers, 512 hidden channels
- Multi-resolution spectral loss for sharp mel output
- 24M parameters, training on `processed_data_100mel/`
- Output: `output_v3/`

## Usage

### Preprocessing
```bash
python3 preprocess.py --input /path/to/dataset --output /path/to/processed_data
```

### Training
```bash
python3 train.py \
  --data_dir processed_data_100mel \
  --output_dir output_v3 \
  --epochs 200 \
  --batch_size 256 \
  --spectral_weight 0.5
```

### Inference
```bash
# Synthesize from text
python3 inference.py --checkpoint output_v3/checkpoints/best.pt --text "your text" --mode synthesize

# Teacher-forced reconstruction
python3 inference.py --checkpoint output_v3/checkpoints/best.pt --mel path/to/mel.pt --mode predict
```

### Batched Inference
```bash
python3 inference_engine.py --checkpoint output_v3/checkpoints/best.pt --benchmark
```


<tool_use_error>File has not been read yet. Read it first before writing to it.</tool_use_error>

Error: File has not been read yet. Read it first before writing to it.


Bash

ls /home/ubuntu/lewm-tts/README.md 2>/dev/null; echo "---"; ls /home/ubuntu/lewm-tts/*.md 2>/dev/null

Check if README exists


/home/ubuntu/lewm-tts/README.md
---
/home/ubuntu/lewm-tts/README.md


     1→# LeWM TTS
     2→
     3→Adapting the LeWM (LeWorldModel) JEPA architecture for single-voice Text-to-Speech synthesis.
     4→
     5→## Architecture
     6→
     7→LeWM's core idea: predict future latent representations using only two losses:
     8→1. **Next-embedding prediction** — given text + previous audio latents, predict the next audio latent
     9→2. **Gaussian regularizer** — prevents latent space collapse
    10→
    11→### Components
    12→- **Text Encoder** — character-level transformer
    13→- **Audio Encoder** — mel spectrogram patch encoder
    14→- **JEPA Predictor** — conditioned on text, predicts next audio embedding
    15→- **Vocoder** — pretrained HiFi-GAN (waveform reconstruction)
    16→
    17→## Dataset
    18→- 20 hours, single speaker, studio quality
    19→- Preprocessing: silence trim, segment 3-15s, forced alignment, 22050/24000 Hz
    20→
    21→## Reference
    22→- LeWM paper: https://le-wm.github.io/
    23→

# LeWM TTS

Adapting the LeWM (LeWorldModel) JEPA architecture for single-voice Text-to-Speech synthesis.

## Architecture

LeWM's core idea: predict future latent representations using only two losses:
1. **Next-embedding prediction** — given text + previous audio latents, predict the next audio latent
2. **Gaussian regularizer** — prevents latent space collapse

### Components
- **Text Encoder** — character-level transformer
- **Audio Encoder** — mel spectrogram patch encoder
- **JEPA Predictor** — conditioned on text, predicts next audio embedding
- **Vocoder** — pretrained HiFi-GAN (waveform reconstruction)

## Dataset
- 20 hours, single speaker, studio quality
- Preprocessing: silence trim, segment 3-15s, forced alignment, 22050/24000 Hz

## Reference
- LeWM paper: https://le-wm.github.io/


Write

# LeWM TTS

JEPA-based Text-to-Speech system adapted from [LeWorldModel](https://le-wm.github.io/).

## Architecture

A 24M parameter model with 4 core components:

```
Text (Hindi/UTF-8 bytes) --> TextEncoder --> text embeddings ---+
                                                                |--> JEPAPredictor --> predicted embeddings
Audio (mel spectrogram)  --> AudioEncoder --> audio embeddings --+           |
                                                                            v
                                                                       MelDecoder --> mel spectrogram --> Vocos --> waveform
```

| Component | Description |
|-----------|-------------|
| **TextEncoder** | Byte-level character transformer (4 layers, 256d, 4 heads). Encodes Hindi text via UTF-8 bytes — no tokenizer needed. |
| **AudioEncoder** | 1D CNN (stride-4 downsample) + Transformer (4 layers). Outputs mu/logvar for Gaussian latent space with reparameterization. |
| **JEPAPredictor** | Transformer decoder (6 layers) with causal self-attention + cross-attention to text. Predicts next audio embedding autoregressively. Has KV-cache for fast inference. |
| **MelDecoder** | Linear projection to 512d, 2x transposed convolution (4x upsample), residual conv blocks (ResConvBlock), and refinement layers. Outputs 100-mel spectrogram. |
| **Vocoder** | [Vocos](https://github.com/gemelo-ai/vocos) `charactr/vocos-mel-24khz` — directly decodes 100-mel to 24kHz waveform. No Griffin-Lim. |

### Speaker Embedding (planned)
Built-in support for multi-speaker via learned embedding table added to text embeddings. Not yet trained — target is 10-20 speaker multi-speaker TTS.

## Training Losses

4 losses, all active during training:

1. **Prediction Loss (MSE)** — next-embedding prediction: `MSE(predicted[t], target[t+1])`
2. **KL Loss** — Gaussian regularizer: `KL(q(z|x) || N(0, I))` with annealing over first 20% of training
3. **Reconstruction Loss (L1)** — mel decoder output vs ground truth mel spectrogram
4. **Multi-Resolution Spectral Loss** — STFT at 3 resolutions (64/128/256 FFT sizes), combines spectral convergence + log magnitude L1. Forces sharp spectral detail instead of blurry averages.

## Dataset

- **Source**: 20 hours of studio-quality single-speaker Hindi audio
- **Preprocessing**: silence-based segmentation into 3-12s chunks, mel extraction
- **Processed**: 7,009 segments, 17.88 hours
- **Mel config**: 100 mels, n_fft=1024, hop=256, sr=24000, f_max=None (full bandwidth), log-clamped at 1e-7 — matches Vocos `mel-24khz` exactly

## Files

| File | Purpose |
|------|---------|
| `model.py` | All model components: TextEncoder, AudioEncoder, JEPAPredictor, MelDecoder, MultiResolutionSpectralLoss, LeWMTTS |
| `train.py` | Training loop with KL annealing, gradient clipping, TensorBoard logging, checkpointing |
| `inference.py` | Single-request inference — AR synthesis, teacher-forced reconstruction |
| `inference_engine.py` | Production batched inference engine — KV-cache, FP16, voice prompts, benchmarking |
| `preprocess.py` | Data preprocessing — audio segmentation, mel extraction, manifest creation |
| `dataset.py` | PyTorch dataset + dataloader with dynamic batching and padding |

## Inference

### Modes
- **Synthesize**: Pure autoregressive generation from text. Requires audio prompt (voice cloning) for best results.
- **Reconstruct**: Encode real mel -> decode back (tests encoder/decoder quality)
- **Predict (teacher-forced)**: Encode mel -> predictor -> decode (tests full pipeline without AR drift)

### Performance (A100 80GB)
- Single request with KV-cache: 9x real-time
- Batched (batch=64): 359x real-time throughput

### Voice Prompting
The inference engine supports audio prompts — a reference clip is encoded and used to seed autoregressive generation, effectively providing voice cloning without explicit speaker conditioning.

## Training History

### v1 — 80 mels + Griffin-Lim + Vocos
- Simple MelDecoder (2x transposed conv), 80-mel output
- Vocoder: Griffin-Lim -> Vocos refinement pipeline
- Result: teacher-forced quality decent but "rustic" sound from GL artifacts

### v2 — DAC-based
- Experimental DAC codec variant
- Checkpoints preserved at `v2/output/checkpoints/`

### v3 — 100 mels + upgraded decoder + spectral loss (current)
- Direct 100-mel output matching Vocos natively (no GL middleman)
- Upgraded MelDecoder with ResConvBlock residual layers, 512 hidden channels
- Multi-resolution spectral loss for sharp mel output
- 24M parameters, training on `processed_data_100mel/`
- Output: `output_v3/`

## Usage

### Preprocessing
```bash
python3 preprocess.py --input /path/to/dataset --output /path/to/processed_data
```

### Training
```bash
python3 train.py \
  --data_dir processed_data_100mel \
  --output_dir output_v3 \
  --epochs 200 \
  --batch_size 256 \
  --spectral_weight 0.5
```

### Inference
```bash
# Synthesize from text
python3 inference.py --checkpoint output_v3/checkpoints/best.pt --text "your text" --mode synthesize

# Teacher-forced reconstruction
python3 inference.py --checkpoint output_v3/checkpoints/best.pt --mel path/to/mel.pt --mode predict
```

### Batched Inference
```bash
python3 inference_engine.py --checkpoint output_v3/checkpoints/best.pt --benchmark
```

## Reference
- LeWM paper: https://le-wm.github.io/