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

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

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

# ─── 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 MelDecoder(nn.Module):
    """Decodes latent embeddings back to mel spectrograms (reverses AudioEncoder)."""

    def __init__(self, d_model=256, n_mels=100):
        super().__init__()
        self.proj = nn.Sequential(
            nn.Linear(d_model, d_model),
            nn.GELU(),
        )
        self.upsample = nn.Sequential(
            nn.ConvTranspose1d(d_model, d_model, kernel_size=4, stride=2, padding=1),
            nn.GELU(),
            nn.ConvTranspose1d(d_model, d_model, kernel_size=4, stride=2, padding=1),
            nn.GELU(),
            nn.Conv1d(d_model, 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)
        return self.upsample(x)

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

        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)

        total_loss = prediction_loss + self.kl_weight * kl_loss + self.recon_weight * recon_loss

        return {
            "total_loss": total_loss,
            "prediction_loss": prediction_loss,
            "kl_loss": kl_loss,
            "recon_loss": recon_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

+    """Residual conv block for mel decoder."""

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

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

-    """Decodes latent embeddings back to mel spectrograms (reverses AudioEncoder)."""

+    """Decodes latent embeddings back to mel spectrograms.

+    Improved architecture with residual blocks for sharper mel output."""

+        hidden = d_model * 2  # 512 — more capacity

-            nn.Linear(d_model, d_model),

+            nn.Linear(d_model, hidden),

-            nn.ConvTranspose1d(d_model, d_model, kernel_size=4, stride=2, padding=1),

+            nn.ConvTranspose1d(hidden, hidden, kernel_size=4, stride=2, padding=1),

+            ResConvBlock(hidden, kernel_size=3),

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

+        # Refinement + projection to mel bins

+            nn.Conv1d(hidden, hidden, kernel_size=7, padding=3),

+            ResConvBlock(hidden, kernel_size=5),

+            nn.Conv1d(hidden, n_mels, kernel_size=7, padding=3),

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


Conversation compacted

Agent Bash Edit Glob Read TaskCreate TaskOutput TaskUpdate Write


This session is being continued from a previous conversation that ran out of context. The summary below covers the earlier portion of the conversation.

Summary:
1. Primary Request and Intent:
   - Run inference on 2 versions (v1 mel-based, v2 DAC-based) of LeWM TTS with latest checkpoints using Hindi text
   - Fix vocoder to use Vocos instead of Griffin-Lim (as was done previously)
   - Implement KV-cache for fast autoregressive inference
   - Build production batched inference engine with voice prompting support
   - Change model from 80 mels to 100 mels to match Vocos mel-24khz directly (eliminate Griffin-Lim middleman causing "rustic" sound)
   - Add speaker embedding support for future 10-20 speaker multi-speaker TTS
   - Reprocess dataset with 100 mels and retrain
   - Upgrade mel decoder architecture and add spectral loss because model produces blurry/structured mel artifacts that Vocos handles poorly (unlike random noise)
   - Final goal: 10-20 speaker multi-speaker TTS on 14M parameter model with high throughput (200x+ real-time)

2. Key Technical Concepts:
   - **LeWM TTS Architecture**: JEPA-based TTS with TextEncoder, AudioEncoder, JEPAPredictor, MelDecoder
   - **Vocos vocoder**: `charactr/vocos-mel-24khz` expects 100 mels, n_fft=1024, hop=256, sr=24000, f_max=None, power=1, log clamp=1e-7
   - **KV-cache**: Custom implementation for JEPAPredictor - caches projected K,V for self-attention and cross-attention per layer, O(1) per step instead of O(n)
   - **Batched inference**: Multiple concurrent requests processed in parallel on GPU - batch=64 achieves 359x real-time throughput on A100
   - **Inference modes**: synthesize (pure AR), reconstruct (encode→decode), predict_tf (teacher-forced prediction)
   - **AR synthesis failure**: Random start token causes noise; audio prompting (voice cloning) fixes this by seeding with real encoder embeddings
   - **Exposure bias**: Predictor trained on real encoder embeddings but sees its own predictions during AR generation
   - **Mel decoder quality**: Simple transposed conv decoder produces structured/blurry distortion; Vocos is sensitive to this unlike Griffin-Lim
   - **Speaker embedding**: Learned embedding table added to text embeddings for multi-speaker conditioning
   - **Noise tolerance test**: GT mel + random noise at various levels decoded through Vocos to determine quality threshold

3. Files and Code Sections:

   - **`/home/ubuntu/lewm-tts/model.py`** — Core model file, heavily modified
     - Changed all `n_mels=80` defaults to `n_mels=100`
     - Added KV-cache methods to JEPAPredictor: `_cached_mha()`, `_cached_cross_attn()`, `init_cache()`, `inference_step()`
     - Added speaker embedding to LeWMTTS: `speaker_embed`, `_add_speaker()`, `n_speakers` config
     - Added `predict_next_cached()` and `init_ar_cache()` to LeWMTTS
     - **Currently being upgraded**: MelDecoder replaced with improved version using ResConvBlock
     - New MelDecoder (partially inserted):
     ```python
     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):
         def __init__(self, d_model=256, n_mels=100):
             super().__init__()
             hidden = d_model * 2  # 512
             self.proj = nn.Sequential(nn.Linear(d_model, hidden), nn.GELU())
             self.upsample = nn.Sequential(
                 nn.ConvTranspose1d(hidden, hidden, kernel_size=4, stride=2, padding=1),
                 nn.GELU(),
                 ResConvBlock(hidden, kernel_size=3),
                 nn.ConvTranspose1d(hidden, hidden, kernel_size=4, stride=2, padding=1),
                 nn.GELU(),
                 ResConvBlock(hidden, kernel_size=3),
             )
             self.out = nn.Sequential(
                 nn.Conv1d(hidden, hidden, kernel_size=7, padding=3),
                 nn.GELU(),
                 ResConvBlock(hidden, kernel_size=5),
                 nn.Conv1d(hidden, n_mels, kernel_size=7, padding=3),
             )
         def forward(self, z):
             x = self.proj(z)
             x = x.transpose(1, 2)
             x = self.upsample(x)
             return self.out(x)
     ```
     - **Multi-resolution spectral loss NOT yet added** — this was the next step

   - **`/home/ubuntu/lewm-tts/preprocess.py`** — Data preprocessing
     - Changed MEL_CONFIG: `n_mels=100`, `f_max=None`, added `clip_val=1e-7`
     - Updated `extract_mel()` to handle f_max=None and use clip_val from config
     - Dataset reprocessed to `/home/ubuntu/lewm-tts/processed_data_100mel/` (7009 segments, 17.88h)

   - **`/home/ubuntu/lewm-tts/inference.py`** — Single-request inference
     - Complete rewrite: removed Griffin-Lim entirely, uses `vocos.decode()` directly
     - Uses KV-cache for AR synthesis
     - Added audio normalization to `save_audio()`

   - **`/home/ubuntu/lewm-tts/inference_engine.py`** — Production batched inference engine
     - Created new file with: batched inference, KV-cache, FP16, voice prompts, CLI
     - Vocoder class uses Vocos direct decode (no Griffin-Lim)
     - `InferenceEngine` class with `synthesize()`, `synthesize_batch()`, `benchmark()`
     - Prompt mel extraction uses n_mels=100, power=1.0, clamp=1e-7
     - Benchmark results: batch=64 → 359x throughput, 36ms per request on A100

   - **`/home/ubuntu/lewm-tts/train.py`** — Training loop
     - Changed `n_mels: 80` → `n_mels: 100` in config

   - **`/home/ubuntu/lewm-tts/dataset.py`** — No changes needed (reads n_mels from data dynamically)

   - **`/home/ubuntu/.claude/projects/-home-ubuntu/memory/lewm-tts.md`** — Project memory updated with vocoder notes, inference findings, performance benchmarks

4. Errors and fixes:
   - **English text used for Hindi model**: User corrected to use Hindi text since dataset is Modi speeches
   - **Vocos expects 100 mels, model outputs 80**: Tried zero-padding (noise), interpolation (quiet/bad), GL→Vocos pipeline (rustic sound). Final fix: retrain model with 100 mels
   - **`vocos_decode` method broke init**: Inserted `vocos_decode` method with `return` statement that orphaned `_start_emb_scale` and `print` lines from `__init__`. Fixed by moving them back above the method
   - **Vocos forward shape mismatch**: Tried `[1, 1, T]` but Vocos expects `[B, T]`. Fixed by removing extra dimension
   - **Griffin-Lim vocoder returns different shape**: GL returns `[1, T]`, squeeze made it 1D. Fixed shape handling
   - **Embedding concatenation shape bug in inference_engine.py**: `torch.cat(dim=0)` on `[1,1,d]` tensors gave `[T,1,d]` instead of `[1,T,d]`. Fixed to `torch.cat(dim=1)`
   - **Model recon produces structured distortion**: Even at same L1 as random noise, model output sounds worse because distortion is structured (blurry mels). Solution: upgrade mel decoder + add spectral loss (in progress)
   - **`replace_all` for n_mels missed config.get default**: Had to manually fix `config.get("n_mels", 80)` → 100

5. Problem Solving:
   - **AR synthesis noise**: Diagnosed as random start token issue. Model never sees random noise during training. Audio prompting (voice cloning) fixes this by providing real encoder embeddings
   - **Later checkpoint collapse**: Epochs 350+ produce very short AR output due to norm-based stopping criterion firing early (embeddings too uniform). Early checkpoints (70-100) better for AR
   - **Rustic/rough sound**: Caused by Griffin-Lim artifacts in GL→Vocos pipeline. Fixed by switching to 100 mels for direct Vocos decode
   - **Slow single-request inference**: 6.5x real-time without KV-cache. Added KV-cache → 9x. Added batching → 359x at batch=64
   - **Blurry mel reconstruction**: Current mel decoder too simple (2 transposed convs). Upgrading with ResConvBlock and more capacity + planning multi-resolution spectral loss
   - **Noise tolerance test**: Proved GT mel + random noise at L1=0.3 sounds good/natural, but model's structured error at L1=0.6 sounds worse. Need better decoder architecture

6. All user messages:
   - "hi"
   - "run infernce on 2 version latets checkpoints '/home/ubuntu/lewm-tts'"
   - "fuck brother use hindi .. and aso u fucjed up soemwhere ..check meory once on this"
   - "Predict (teacher-forced) │ output/v1_e350_vocos_predict.wav │ 9.50s │ ... Synthesize (Hindi) wt the differnce between both"
   - "Synthesize - this is noise ..at 25th/75th epoch we goyt good result but not here... u dng infernce something worng and for v2 also"
   - "but everuuthing pure noise .. u dng infernce some wrong"
   - "'/home/ubuntu/lewm-tts/output/vocos_predict_tf.wav' i am talking aboit this ,,, this came good ...i tested it yesterday at 25 or 70 eoch not sure ... it is abit robotic i thought because of early checkpoint"
   - "yeah its good .. what is preoblem ..are we using weights to predict?"
   - "but it can't generate autoregressively yet ? why"
   - "use a real audio embedding - is it liek cloning?"
   - "what other options ?"
   - "and see my final goal is to train atleast 10 -20 speakers on this 14m model"
   - "but it will slow at inferne ryt ?"
   - "in normal how much rtf we wil get?"
   - "0.15 is very less ... i am expeting somethinglike 200x atlease .. and are u run n gpu or cpu"
   - "yeah do" (re: KV-cache)
   - "no i will do so kany concurrnet request so 200x is less .. and i am stressing on infernce beause we got to know this will work"
   - "setup all" (re: batched inference engine)
   - "check with latest checkpoint and lets see the voice how it coems"
   - "brotehr y u using 70 checpoint we usine before 350 ryt ... just to check voice .. use tf_precit .. anyway we need to retain toatl model again with spekaer emebeds"
   - "it is ok but some rustic sound cmng .. voice good but rustic y so"
   - "go with option 1... and please make sure this time correctly" (re: 100 mels)
   - "status"
   - "thats good stop old v1 first"
   - "what about v2"
   - "what is batch size running?"
   - "remove previous checkpoint of v1... dont remove v2 ... and do a quick test to adjust bacth size"
   - "do 512 .."
   - "stop" (interrupted batch 512 training)
   - "what is happening"
   - "do 256 batch size"
   - "status"
   - "checkmgpu utilization ... and test latest chepcoint"
   - "pure noise it is"
   - "gt is good.. remianig 2 voice cmng started a liitle bit"
   - "test now" (at epoch 69)
   - "still noice" (at epoch 102)
   - "are u sure we able to produce good voice"
   - "vocos_noise_0.1.wav this sounds good but cmng robotic voice ..vocos_noise_0.3.wav this is good with out any robotiic sound .... just to mention 200 also good not best sounds natural.. but the thing we genarted before doesnt have any quality iof these genarted"
   - "yeah dobit make sure u do correct" (re: upgrade mel decoder + spectral loss)

7. Pending Tasks:
   - **Complete mel decoder upgrade in model.py** — ResConvBlock written and inserted, but need to verify it compiles
   - **Add multi-resolution spectral loss** — Not yet implemented. Needs to be added to the `forward()` method in LeWMTTS
   - **Restart training** with upgraded mel decoder + spectral loss on `processed_data_100mel`
   - **Test quality** after training converges
   - **v2 checkpoints** are preserved at `/home/ubuntu/lewm-tts/v2/output/checkpoints/`

8. Current Work:
   Training was stopped (`pkill -f "train\.py"`). The mel decoder in `model.py` was being upgraded. The new `MelDecoder` class with `ResConvBlock` has been written and inserted into model.py (replacing the old simple version). The old MelDecoder was:
   ```python
   # Simple: Linear → 2x ConvTranspose → Conv → n_mels
   ```
   The new one uses hidden=d_model*2 (512), residual conv blocks, and more refinement layers.

   **What's NOT done yet**: Multi-resolution spectral loss needs to be added to the `forward()` method. This was explicitly mentioned as needed ("Better loss — multi-resolution spectral loss instead of just L1") and user confirmed ("yeah dobit make sure u do correct").

   The old v1 checkpoints were removed. V2 checkpoints preserved. Dataset at `processed_data_100mel/` is ready (7009 segments, 17.88h, 100-mel).

9. Optional Next Step:
   Add multi-resolution spectral loss to `model.py`'s `LeWMTTS.forward()` method, then verify the model builds correctly with the new MelDecoder + spectral loss, and restart training.

   User's last explicit request: "yeah dobit make sure u do correct" in response to assistant saying "The problem is the mel decoder — it's too simple... Two fixes: 1. Better mel decoder — add residual layers, more capacity. 2. Better loss — multi-resolution spectral loss instead of just L1. Both are needed."

   The MelDecoder upgrade was inserted into model.py. The spectral loss still needs to be implemented and added to the forward pass. Then training needs to restart with `processed_data_100mel`.

If you need specific details from before compaction (like exact code snippets, error messages, or content you generated), read the full transcript at: /home/ubuntu/.claude/projects/-home-ubuntu/1b244b10-1345-40d2-b03a-2de470236690.jsonl
Continue the conversation from where it left off without asking the user any further questions. Resume directly — do not acknowledge the summary, do not recap what was happening, do not preface with "I'll continue" or similar. Pick up the last task as if the break never happened.


Let me check the current state of model.py to see where the mel decoder upgrade left off and what still needs to be done (multi-resolution spectral loss).


ToolSearch

Question: select:Read,Edit,Write,Bash,Grep


Question: select:Read,Edit,Write,Bash,Grep


Let me check the current state of model.py to see where the mel decoder upgrade is and what needs to be added (multi-resolution spectral loss).


     1→"""
     2→LeWM TTS Model — JEPA-based Text-to-Speech
     3→
     4→Architecture:
     5→  - TextEncoder: Character-level transformer for Hindi/Devanagari text
     6→  - AudioEncoder: 1D CNN + Transformer that encodes mel spectrograms to embeddings
     7→  - JEPAPredictor: Transformer that predicts next audio embedding given text + audio context
     8→  - Two losses: next-embedding prediction (MSE) + Gaussian regularizer (KL to N(0,I))
     9→"""
    10→
    11→import math
    12→import torch
    13→import torch.nn as nn
    14→import torch.nn.functional as F
    15→
    16→
    17→# ─── Text Encoder ───────────────────────────────────────────────────────────
    18→
    19→class TextEncoder(nn.Module):
    20→    """Character-level transformer encoder for Hindi text."""
    21→
    22→    def __init__(self, vocab_size=256, d_model=256, nhead=4, num_layers=4, dropout=0.1):
    23→        super().__init__()
    24→        self.d_model = d_model
    25→
    26→        # Character embedding (covers all Unicode Hindi chars via byte-level encoding)
    27→        self.char_embed = nn.Embedding(vocab_size, d_model)
    28→        self.pos_embed = SinusoidalPositionalEncoding(d_model, max_len=2048)
    29→        self.dropout = nn.Dropout(dropout)
    30→
    31→        encoder_layer = nn.TransformerEncoderLayer(
    32→            d_model=d_model, nhead=nhead,
    33→            dim_feedforward=d_model * 4,
    34→            dropout=dropout, batch_first=True,
    35→            activation="gelu",
    36→        )
    37→        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
    38→        self.proj = nn.Linear(d_model, d_model)
    39→
    40→    def forward(self, text_tokens, text_mask=None):
    41→        """
    42→        Args:
    43→            text_tokens: [B, T_text] long tensor of byte-level token IDs
    44→            text_mask: [B, T_text] bool, True = padding
    45→        Returns:
    46→            text_emb: [B, T_text, d_model]
    47→        """
    48→        x = self.char_embed(text_tokens) * math.sqrt(self.d_model)
    49→        x = self.pos_embed(x)
    50→        x = self.dropout(x)
    51→        x = self.transformer(x, src_key_padding_mask=text_mask)
    52→        x = self.proj(x)
    53→        return x
    54→
    55→
    56→# ─── Audio Encoder ──────────────────────────────────────────────────────────
    57→
    58→class AudioEncoder(nn.Module):
    59→    """
    60→    Encodes mel spectrograms into a sequence of audio embeddings.
    61→    Uses 1D CNN for downsampling + Transformer for contextualization.
    62→    """
    63→
    64→    def __init__(self, n_mels=100, d_model=256, nhead=4, num_layers=4,
    65→                 downsample_factor=4, dropout=0.1):
    66→        super().__init__()
    67→        self.d_model = d_model
    68→        self.downsample_factor = downsample_factor
    69→
    70→        # CNN downsampling: reduces temporal resolution by downsample_factor
    71→        self.conv_pre = nn.Sequential(
    72→            nn.Conv1d(n_mels, d_model, kernel_size=7, padding=3),
    73→            nn.GELU(),
    74→            nn.Conv1d(d_model, d_model, kernel_size=4, stride=2, padding=1),
    75→            nn.GELU(),
    76→            nn.Conv1d(d_model, d_model, kernel_size=4, stride=2, padding=1),
    77→            nn.GELU(),
    78→        )  # Total stride = 4
    79→
    80→        self.pos_embed = SinusoidalPositionalEncoding(d_model, max_len=4096)
    81→        self.dropout = nn.Dropout(dropout)
    82→
    83→        encoder_layer = nn.TransformerEncoderLayer(
    84→            d_model=d_model, nhead=nhead,
    85→            dim_feedforward=d_model * 4,
    86→            dropout=dropout, batch_first=True,
    87→            activation="gelu",
    88→        )
    89→        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
    90→
    91→        # Project to embedding with mu and logvar for Gaussian regularizer
    92→        self.proj_mu = nn.Linear(d_model, d_model)
    93→        self.proj_logvar = nn.Linear(d_model, d_model)
    94→
    95→    def forward(self, mel, mel_mask=None):
    96→        """
    97→        Args:
    98→            mel: [B, n_mels, T_mel] — log mel spectrogram
    99→            mel_mask: [B, T_mel] bool, True = padding (before downsampling)
   100→        Returns:
   101→            z: [B, T_down, d_model] — sampled latent embeddings
   102→            mu: [B, T_down, d_model]
   103→            logvar: [B, T_down, d_model]
   104→        """
   105→        # CNN encoding: [B, n_mels, T] -> [B, d_model, T//4]
   106→        x = self.conv_pre(mel)
   107→        x = x.transpose(1, 2)  # [B, T_down, d_model]
   108→
   109→        # Downsample mask if provided
   110→        if mel_mask is not None:
   111→            # Approximate downsampled mask
   112→            T_down = x.shape[1]
   113→            mel_mask = mel_mask[:, ::self.downsample_factor][:, :T_down]
   114→
   115→        x = self.pos_embed(x)
   116→        x = self.dropout(x)
   117→        x = self.transformer(x, src_key_padding_mask=mel_mask)
   118→
   119→        mu = self.proj_mu(x)
   120→        logvar = self.proj_logvar(x)
   121→
   122→        # Reparameterization trick
   123→        if self.training:
   124→            std = torch.exp(0.5 * logvar)
   125→            eps = torch.randn_like(std)
   126→            z = mu + eps * std
   127→        else:
   128→            z = mu
   129→
   130→        return z, mu, logvar
   131→
   132→
   133→# ─── JEPA Predictor ─────────────────────────────────────────────────────────
   134→
   135→class JEPAPredictor(nn.Module):
   136→    """
   137→    Predicts next audio embedding given:
   138→    - Text context (from TextEncoder)
   139→    - Previous audio embeddings (from AudioEncoder)
   140→
   141→    Uses cross-attention to condition on text, and causal self-attention
   142→    over the audio embedding sequence.
   143→    """
   144→
   145→    def __init__(self, d_model=256, nhead=4, num_layers=6, dropout=0.1):
   146→        super().__init__()
   147→        self.d_model = d_model
   148→
   149→        self.audio_input_proj = nn.Linear(d_model, d_model)
   150→        self.pos_embed = SinusoidalPositionalEncoding(d_model, max_len=4096)
   151→        self.dropout = nn.Dropout(dropout)
   152→
   153→        # Transformer decoder with causal self-attention + cross-attention to text
   154→        decoder_layer = nn.TransformerDecoderLayer(
   155→            d_model=d_model, nhead=nhead,
   156→            dim_feedforward=d_model * 4,
   157→            dropout=dropout, batch_first=True,
   158→            activation="gelu",
   159→        )
   160→        self.transformer = nn.TransformerDecoder(decoder_layer, num_layers=num_layers)
   161→
   162→        # Output projection to predict next embedding
   163→        self.output_proj = nn.Sequential(
   164→            nn.Linear(d_model, d_model),
   165→            nn.GELU(),
   166→            nn.Linear(d_model, d_model),
   167→        )
   168→
   169→    def forward(self, audio_emb, text_emb, audio_mask=None, text_mask=None):
   170→        """
   171→        Args:
   172→            audio_emb: [B, T_audio, d_model] — audio embeddings (from encoder)
   173→            text_emb: [B, T_text, d_model] — text embeddings
   174→            audio_mask: [B, T_audio] bool, True = padding
   175→            text_mask: [B, T_text] bool, True = padding
   176→        Returns:
   177→            predicted: [B, T_audio, d_model] — predicted next embeddings
   178→        """
   179→        x = self.audio_input_proj(audio_emb)
   180→        x = self.pos_embed(x)
   181→        x = self.dropout(x)
   182→
   183→        # Causal mask: each position can only attend to previous positions
   184→        T = x.shape[1]
   185→        causal_mask = torch.triu(
   186→            torch.ones(T, T, device=x.device, dtype=torch.bool), diagonal=1
   187→        )
   188→
   189→        predicted = self.transformer(
   190→            x, text_emb,
   191→            tgt_mask=causal_mask,
   192→            tgt_key_padding_mask=audio_mask,
   193→            memory_key_padding_mask=text_mask,
   194→        )
   195→
   196→        predicted = self.output_proj(predicted)
   197→        return predicted
   198→
   199→    def _cached_mha(self, mha, q_input, kv_input, cache_k, cache_v):
   200→        """Run multi-head attention with KV cache.
   201→        q_input: [B, 1, d], kv_input: [B, 1, d] (new token for K,V)
   202→        cache_k, cache_v: [B, T_prev, d] projected keys/values
   203→        Returns: output [B, 1, d], updated cache_k, updated cache_v
   204→        """
   205→        d = mha.embed_dim
   206→        nhead = mha.num_heads
   207→        head_dim = d // nhead
   208→        B = q_input.shape[0]
   209→
   210→        # Project Q from q_input, K/V from kv_input
   211→        # in_proj_weight is [3d, d] = [Wq; Wk; Wv]
   212→        Wq, Wk, Wv = mha.in_proj_weight.chunk(3, dim=0)
   213→        bq, bk, bv = mha.in_proj_bias.chunk(3, dim=0)
   214→
   215→        q = F.linear(q_input, Wq, bq)        # [B, 1, d]
   216→        k_new = F.linear(kv_input, Wk, bk)    # [B, 1, d]
   217→        v_new = F.linear(kv_input, Wv, bv)    # [B, 1, d]
   218→
   219→        # Append to cache
   220→        if cache_k is not None and cache_k.shape[1] > 0:
   221→            k = torch.cat([cache_k, k_new], dim=1)
   222→            v = torch.cat([cache_v, v_new], dim=1)
   223→        else:
   224→            k = k_new
   225→            v = v_new
   226→
   227→        # Multi-head attention
   228→        q = q.view(B, 1, nhead, head_dim).transpose(1, 2)      # [B, nhead, 1, head_dim]
   229→        k_mh = k.view(B, -1, nhead, head_dim).transpose(1, 2)  # [B, nhead, T, head_dim]
   230→        v_mh = v.view(B, -1, nhead, head_dim).transpose(1, 2)
   231→
   232→        attn = torch.matmul(q, k_mh.transpose(-2, -1)) / (head_dim ** 0.5)
   233→        attn = F.softmax(attn, dim=-1)
   234→        out = torch.matmul(attn, v_mh)  # [B, nhead, 1, head_dim]
   235→
   236→        out = out.transpose(1, 2).contiguous().view(B, 1, d)
   237→        out = F.linear(out, mha.out_proj.weight, mha.out_proj.bias)
   238→
   239→        return out, k, v
   240→
   241→    def _cached_cross_attn(self, mha, q_input, memory, cache_k, cache_v, memory_mask=None):
   242→        """Cross-attention with cached memory K,V (computed once).
   243→        memory_mask: [B, T_text] bool, True = padding (applied as -inf)."""
   244→        d = mha.embed_dim
   245→        nhead = mha.num_heads
   246→        head_dim = d // nhead
   247→        B = q_input.shape[0]
   248→
   249→        Wq, Wk, Wv = mha.in_proj_weight.chunk(3, dim=0)
   250→        bq, bk, bv = mha.in_proj_bias.chunk(3, dim=0)
   251→
   252→        q = F.linear(q_input, Wq, bq)  # [B, 1, d]
   253→
   254→        if cache_k is None:
   255→            k = F.linear(memory, Wk, bk)  # [B, T_text, d]
   256→            v = F.linear(memory, Wv, bv)
   257→        else:
   258→            k, v = cache_k, cache_v
   259→
   260→        q = q.view(B, 1, nhead, head_dim).transpose(1, 2)
   261→        k_mh = k.view(B, -1, nhead, head_dim).transpose(1, 2)
   262→        v_mh = v.view(B, -1, nhead, head_dim).transpose(1, 2)
   263→
   264→        attn = torch.matmul(q, k_mh.transpose(-2, -1)) / (head_dim ** 0.5)
   265→        if memory_mask is not None:
   266→            # memory_mask: [B, T_text] → [B, 1, 1, T_text]
   267→            attn = attn.masked_fill(memory_mask.unsqueeze(1).unsqueeze(2), float('-inf'))
   268→        attn = F.softmax(attn, dim=-1)
   269→        out = torch.matmul(attn, v_mh)
   270→
   271→        out = out.transpose(1, 2).contiguous().view(B, 1, d)
   272→        out = F.linear(out, mha.out_proj.weight, mha.out_proj.bias)
   273→
   274→        return out, k, v
   275→
   276→    def init_cache(self, num_layers):
   277→        """Initialize empty KV cache."""
   278→        return {
   279→            'self_k': [None] * num_layers,
   280→            'self_v': [None] * num_layers,
   281→            'cross_k': [None] * num_layers,
   282→            'cross_v': [None] * num_layers,
   283→        }
   284→
   285→    def inference_step(self, new_emb, text_emb, step_idx, cache, text_mask=None):
   286→        """Process one token with KV cache.
   287→        new_emb: [B, 1, d_model] — raw audio embedding (not projected)
   288→        text_mask: [B, T_text] bool, True = padding
   289→        Returns: predicted [B, 1, d_model], updated cache
   290→        """
   291→        x = self.audio_input_proj(new_emb)
   292→        # Add positional encoding for this specific position
   293→        x = x + self.pos_embed.pe[:, step_idx:step_idx + 1]
   294→
   295→        for i, layer in enumerate(self.transformer.layers):
   296→            # Self-attention with cache (post-norm)
   297→            sa_out, cache['self_k'][i], cache['self_v'][i] = self._cached_mha(
   298→                layer.self_attn, x, x, cache['self_k'][i], cache['self_v'][i]
   299→            )
   300→            x = layer.norm1(x + layer.dropout1(sa_out))
   301→
   302→            # Cross-attention with cached text KV
   303→            ca_out, cache['cross_k'][i], cache['cross_v'][i] = self._cached_cross_attn(
   304→                layer.multihead_attn, x, text_emb, cache['cross_k'][i], cache['cross_v'][i],
   305→                memory_mask=text_mask,
   306→            )
   307→            x = layer.norm2(x + layer.dropout2(ca_out))
   308→
   309→            # FFN
   310→            ff_out = layer.linear2(layer.dropout(layer.activation(layer.linear1(x))))
   311→            x = layer.norm3(x + layer.dropout3(ff_out))
   312→
   313→        return self.output_proj(x), cache
   314→
   315→
   316→# ─── Positional Encoding ────────────────────────────────────────────────────
   317→
   318→class SinusoidalPositionalEncoding(nn.Module):
   319→    def __init__(self, d_model, max_len=8192):
   320→        super().__init__()
   321→        pe = torch.zeros(max_len, d_model)
   322→        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
   323→        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
   324→        pe[:, 0::2] = torch.sin(position * div_term)
   325→        pe[:, 1::2] = torch.cos(position * div_term)
   326→        self.register_buffer("pe", pe.unsqueeze(0))  # [1, max_len, d_model]
   327→
   328→    def forward(self, x):
   329→        return x + self.pe[:, :x.shape[1]]
   330→
   331→
   332→# ─── Mel Decoder ────────────────────────────────────────────────────────────
   333→
   334→class ResConvBlock(nn.Module):
   335→    """Residual conv block for mel decoder."""
   336→
   337→    def __init__(self, channels, kernel_size=3):
   338→        super().__init__()
   339→        self.net = nn.Sequential(
   340→            nn.Conv1d(channels, channels, kernel_size, padding=kernel_size // 2),
   341→            nn.GELU(),
   342→            nn.Conv1d(channels, channels, kernel_size, padding=kernel_size // 2),
   343→        )
   344→
   345→    def forward(self, x):
   346→        return x + self.net(x)
   347→
   348→
   349→class MelDecoder(nn.Module):
   350→    """Decodes latent embeddings back to mel spectrograms.
   351→    Improved architecture with residual blocks for sharper mel output."""
   352→
   353→    def __init__(self, d_model=256, n_mels=100):
   354→        super().__init__()
   355→        hidden = d_model * 2  # 512 — more capacity
   356→
   357→        self.proj = nn.Sequential(
   358→            nn.Linear(d_model, hidden),
   359→            nn.GELU(),
   360→        )
   361→        self.upsample = nn.Sequential(
   362→            # 4x upsample (2x + 2x)
   363→            nn.ConvTranspose1d(hidden, hidden, kernel_size=4, stride=2, padding=1),
   364→            nn.GELU(),
   365→            ResConvBlock(hidden, kernel_size=3),
   366→            nn.ConvTranspose1d(hidden, hidden, kernel_size=4, stride=2, padding=1),
   367→            nn.GELU(),
   368→            ResConvBlock(hidden, kernel_size=3),
   369→        )
   370→        # Refinement + projection to mel bins
   371→        self.out = nn.Sequential(
   372→            nn.Conv1d(hidden, hidden, kernel_size=7, padding=3),
   373→            nn.GELU(),
   374→            ResConvBlock(hidden, kernel_size=5),
   375→            nn.Conv1d(hidden, n_mels, kernel_size=7, padding=3),
   376→        )
   377→
   378→    def forward(self, z):
   379→        """z: [B, T_down, d_model] → mel: [B, n_mels, T_up]"""
   380→        x = self.proj(z)
   381→        x = x.transpose(1, 2)  # [B, hidden, T_down]
   382→        x = self.upsample(x)   # [B, hidden, T_up]
   383→        return self.out(x)      # [B, n_mels, T_up]
   384→
   385→
   386→# ─── Full LeWM TTS Model ────────────────────────────────────────────────────
   387→
   388→class LeWMTTS(nn.Module):
   389→    """
   390→    Complete LeWM TTS model.
   391→
   392→    Training: text + mel → audio embeddings → predict next embedding
   393→    Two losses:
   394→        1. Next-embedding prediction: MSE(predicted[t], target[t+1])
   395→        2. Gaussian regularizer: KL(q(z|x) || N(0, I))
   396→    """
   397→
   398→    def __init__(self, config):
   399→        super().__init__()
   400→        d_model = config.get("d_model", 256)
   401→        nhead = config.get("nhead", 4)
   402→        n_mels = config.get("n_mels", 100)
   403→        text_vocab = config.get("text_vocab_size", 256)
   404→        text_layers = config.get("text_encoder_layers", 4)
   405→        audio_layers = config.get("audio_encoder_layers", 4)
   406→        predictor_layers = config.get("predictor_layers", 6)
   407→        dropout = config.get("dropout", 0.1)
   408→        self.kl_weight = config.get("kl_weight", 0.01)
   409→        n_speakers = config.get("n_speakers", 1)
   410→
   411→        self.text_encoder = TextEncoder(
   412→            vocab_size=text_vocab, d_model=d_model, nhead=nhead,
   413→            num_layers=text_layers, dropout=dropout,
   414→        )
   415→        self.audio_encoder = AudioEncoder(
   416→            n_mels=n_mels, d_model=d_model, nhead=nhead,
   417→            num_layers=audio_layers, dropout=dropout,
   418→        )
   419→        self.predictor = JEPAPredictor(
   420→            d_model=d_model, nhead=nhead,
   421→            num_layers=predictor_layers, dropout=dropout,
   422→        )
   423→
   424→        # Mel decoder: reconstructs mel from latent embeddings
   425→        # Needed so the latent space is grounded to actual audio
   426→        self.mel_decoder = MelDecoder(d_model=d_model, n_mels=n_mels)
   427→        self.recon_weight = config.get("recon_weight", 1.0)
   428→
   429→        # Speaker embedding (for multi-speaker)
   430→        self.n_speakers = n_speakers
   431→        if n_speakers > 1:
   432→            self.speaker_embed = nn.Embedding(n_speakers, d_model)
   433→        else:
   434→            self.speaker_embed = None
   435→
   436→        self.config = config
   437→
   438→    def _add_speaker(self, text_emb, speaker_id):
   439→        """Add speaker embedding to text embeddings."""
   440→        if self.speaker_embed is not None and speaker_id is not None:
   441→            spk = self.speaker_embed(speaker_id).unsqueeze(1)  # [B, 1, d]
   442→            text_emb = text_emb + spk
   443→        return text_emb
   444→
   445→    def forward(self, mel, text_tokens, mel_mask=None, text_mask=None, speaker_id=None):
   446→        """
   447→        Args:
   448→            mel: [B, n_mels, T_mel]
   449→            text_tokens: [B, T_text]
   450→            mel_mask: [B, T_mel] bool
   451→            text_mask: [B, T_text] bool
   452→        Returns:
   453→            loss_dict with 'prediction_loss', 'kl_loss', 'total_loss'
   454→        """
   455→        # Encode text + speaker conditioning
   456→        text_emb = self.text_encoder(text_tokens, text_mask)
   457→        text_emb = self._add_speaker(text_emb, speaker_id)
   458→
   459→        # Encode audio → latent embeddings
   460→        z, mu, logvar = self.audio_encoder(mel, mel_mask)
   461→
   462→        # Predict next embedding: input z[:-1], target z[1:]
   463→        input_emb = z[:, :-1]     # [B, T-1, d]
   464→        target_emb = z[:, 1:]     # [B, T-1, d]
   465→
   466→        # Adjust mask for downsampled sequence
   467→        if mel_mask is not None:
   468→            T_down = z.shape[1]
   469→            ds_mask = mel_mask[:, ::self.audio_encoder.downsample_factor][:, :T_down]
   470→            pred_mask = ds_mask[:, :-1]
   471→        else:
   472→            pred_mask = None
   473→
   474→        predicted = self.predictor(input_emb, text_emb, pred_mask, text_mask)
   475→
   476→        # ─── Loss 1: Next-embedding prediction (MSE) ───
   477→        # Mask out padding positions
   478→        if pred_mask is not None:
   479→            valid = ~pred_mask  # True = valid
   480→            valid = valid.unsqueeze(-1)  # [B, T-1, 1]
   481→            prediction_loss = F.mse_loss(
   482→                predicted * valid, target_emb * valid, reduction="sum"
   483→            ) / (valid.sum() * predicted.shape[-1] + 1e-8)
   484→        else:
   485→            prediction_loss = F.mse_loss(predicted, target_emb)
   486→
   487→        # ─── Loss 2: Gaussian regularizer — KL(N(mu, sigma²) || N(0, I)) ───
   488→        # KL = -0.5 * sum(1 + logvar - mu² - exp(logvar))
   489→        if mel_mask is not None:
   490→            T_down = mu.shape[1]
   491→            ds_mask = mel_mask[:, ::self.audio_encoder.downsample_factor][:, :T_down]
   492→            valid = (~ds_mask).unsqueeze(-1).float()
   493→            kl = -0.5 * (1 + logvar - mu.pow(2) - logvar.exp())
   494→            kl_loss = (kl * valid).sum() / (valid.sum() * mu.shape[-1] + 1e-8)
   495→        else:
   496→            kl_loss = -0.5 * torch.mean(1 + logvar - mu.pow(2) - logvar.exp())
   497→
   498→        # ─── Loss 3: Mel reconstruction ───
   499→        # Decode latents back to mel and compare with input
   500→        mel_recon = self.mel_decoder(z)  # [B, n_mels, T_up]
   501→        # Trim/pad to match input mel length
   502→        T_mel = mel.shape[2]
   503→        T_recon = mel_recon.shape[2]
   504→        if T_recon > T_mel:
   505→            mel_recon = mel_recon[:, :, :T_mel]
   506→        elif T_recon < T_mel:
   507→            mel = mel[:, :, :T_recon]
   508→            if mel_mask is not None:
   509→                mel_mask = mel_mask[:, :T_recon]
   510→
   511→        if mel_mask is not None:
   512→            valid_mel = (~mel_mask).unsqueeze(1).float()  # [B, 1, T]
   513→            recon_loss = F.l1_loss(
   514→                mel_recon * valid_mel, mel * valid_mel, reduction="sum"
   515→            ) / (valid_mel.sum() * mel.shape[1] + 1e-8)
   516→        else:
   517→            recon_loss = F.l1_loss(mel_recon, mel)
   518→
   519→        total_loss = prediction_loss + self.kl_weight * kl_loss + self.recon_weight * recon_loss
   520→
   521→        return {
   522→            "total_loss": total_loss,
   523→            "prediction_loss": prediction_loss,
   524→            "kl_loss": kl_loss,
   525→            "recon_loss": recon_loss,
   526→        }
   527→
   528→    def encode_audio(self, mel):
   529→        """Encode mel to latent embeddings (for inference)."""
   530→        z, mu, _ = self.audio_encoder(mel)
   531→        return mu  # Use mean at inference
   532→
   533→    def predict_next(self, audio_emb, text_emb, text_mask=None):
   534→        """Predict next audio embedding autoregressively."""
   535→        predicted = self.predictor(audio_emb, text_emb, text_mask=text_mask)
   536→        return predicted[:, -1:]  # Only the last prediction
   537→
   538→    def init_ar_cache(self):
   539→        """Initialize KV cache for autoregressive inference."""
   540→        num_layers = len(self.predictor.transformer.layers)
   541→        return self.predictor.init_cache(num_layers)
   542→
   543→    def predict_next_cached(self, new_emb, text_emb, step_idx, cache, text_mask=None):
   544→        """Predict next embedding using KV cache. O(1) per step instead of O(n)."""
   545→        return self.predictor.inference_step(new_emb, text_emb, step_idx, cache, text_mask=text_mask)
   546→
   547→
   548→def count_parameters(model):
   549→    return sum(p.numel() for p in model.parameters() if p.requires_grad)
   550→
   551→
   552→def build_model(config=None):
   553→    if config is None:
   554→        config = {
   555→            "d_model": 256,
   556→            "nhead": 4,
   557→            "n_mels": 100,
   558→            "text_vocab_size": 256,
   559→            "text_encoder_layers": 4,
   560→            "audio_encoder_layers": 4,
   561→            "predictor_layers": 6,
   562→            "dropout": 0.1,
   563→            "kl_weight": 0.01,
   564→        }
   565→    model = LeWMTTS(config)
   566→    print(f"LeWM TTS model: {count_parameters(model)/1e6:.2f}M parameters")
   567→    return model, config
   568→

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

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

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

        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)

        total_loss = prediction_loss + self.kl_weight * kl_loss + self.recon_weight * recon_loss

        return {
            "total_loss": total_loss,
            "prediction_loss": prediction_loss,
            "kl_loss": kl_loss,
            "recon_loss": recon_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


The upgraded MelDecoder with ResConvBlock is already in place. Now I need to add multi-resolution spectral loss to the `forward()` method.


Edit

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

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