     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 copy
    13→import torch
    14→import torch.nn as nn
    15→import torch.nn.functional as F
    16→import torchaudio
    17→
    18→
    19→# ─── Text Encoder ───────────────────────────────────────────────────────────
    20→
    21→class TextEncoder(nn.Module):
    22→    """Character-level transformer encoder for Hindi text."""
    23→
    24→    def __init__(self, vocab_size=256, d_model=256, nhead=4, num_layers=4, dropout=0.1):
    25→        super().__init__()
    26→        self.d_model = d_model
    27→
    28→        # Character embedding (covers all Unicode Hindi chars via byte-level encoding)
    29→        self.char_embed = nn.Embedding(vocab_size, d_model)
    30→        self.pos_embed = SinusoidalPositionalEncoding(d_model, max_len=2048)
    31→        self.dropout = nn.Dropout(dropout)
    32→
    33→        encoder_layer = nn.TransformerEncoderLayer(
    34→            d_model=d_model, nhead=nhead,
    35→            dim_feedforward=d_model * 4,
    36→            dropout=dropout, batch_first=True,
    37→            activation="gelu",
    38→        )
    39→        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
    40→        self.proj = nn.Linear(d_model, d_model)
    41→
    42→    def forward(self, text_tokens, text_mask=None):
    43→        """
    44→        Args:
    45→            text_tokens: [B, T_text] long tensor of byte-level token IDs
    46→            text_mask: [B, T_text] bool, True = padding
    47→        Returns:
    48→            text_emb: [B, T_text, d_model]
    49→        """
    50→        x = self.char_embed(text_tokens) * math.sqrt(self.d_model)
    51→        x = self.pos_embed(x)
    52→        x = self.dropout(x)
    53→        x = self.transformer(x, src_key_padding_mask=text_mask)
    54→        x = self.proj(x)
    55→        return x
    56→
    57→
    58→# ─── Audio Encoder ──────────────────────────────────────────────────────────
    59→
    60→class AudioEncoder(nn.Module):
    61→    """
    62→    Encodes mel spectrograms into a sequence of audio embeddings.
    63→    Uses 1D CNN for downsampling + Transformer for contextualization.
    64→    """
    65→
    66→    def __init__(self, n_mels=100, d_model=256, nhead=4, num_layers=4,
    67→                 downsample_factor=4, dropout=0.1):
    68→        super().__init__()
    69→        self.d_model = d_model
    70→        self.downsample_factor = downsample_factor
    71→
    72→        # CNN downsampling: reduces temporal resolution by downsample_factor
    73→        if downsample_factor == 2:
    74→            self.conv_pre = nn.Sequential(
    75→                nn.Conv1d(n_mels, d_model, kernel_size=7, padding=3),
    76→                nn.GELU(),
    77→                nn.Conv1d(d_model, d_model, kernel_size=4, stride=2, padding=1),
    78→                nn.GELU(),
    79→            )  # Total stride = 2
    80→        else:
    81→            self.conv_pre = nn.Sequential(
    82→                nn.Conv1d(n_mels, d_model, kernel_size=7, padding=3),
    83→                nn.GELU(),
    84→                nn.Conv1d(d_model, d_model, kernel_size=4, stride=2, padding=1),
    85→                nn.GELU(),
    86→                nn.Conv1d(d_model, d_model, kernel_size=4, stride=2, padding=1),
    87→                nn.GELU(),
    88→            )  # Total stride = 4
    89→
    90→        self.pos_embed = SinusoidalPositionalEncoding(d_model, max_len=4096)
    91→        self.dropout = nn.Dropout(dropout)
    92→
    93→        encoder_layer = nn.TransformerEncoderLayer(
    94→            d_model=d_model, nhead=nhead,
    95→            dim_feedforward=d_model * 4,
    96→            dropout=dropout, batch_first=True,
    97→            activation="gelu",
    98→        )
    99→        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
   100→
   101→        # Project to embedding with mu and logvar for Gaussian regularizer
   102→        self.proj_mu = nn.Linear(d_model, d_model)
   103→        self.proj_logvar = nn.Linear(d_model, d_model)
   104→
   105→    def forward(self, mel, mel_mask=None):
   106→        """
   107→        Args:
   108→            mel: [B, n_mels, T_mel] — log mel spectrogram
   109→            mel_mask: [B, T_mel] bool, True = padding (before downsampling)
   110→        Returns:
   111→            z: [B, T_down, d_model] — sampled latent embeddings
   112→            mu: [B, T_down, d_model]
   113→            logvar: [B, T_down, d_model]
   114→        """
   115→        # CNN encoding: [B, n_mels, T] -> [B, d_model, T//4]
   116→        x = self.conv_pre(mel)
   117→        x = x.transpose(1, 2)  # [B, T_down, d_model]
   118→
   119→        # Downsample mask if provided
   120→        if mel_mask is not None:
   121→            # Approximate downsampled mask
   122→            T_down = x.shape[1]
   123→            mel_mask = mel_mask[:, ::self.downsample_factor][:, :T_down]
   124→
   125→        x = self.pos_embed(x)
   126→        x = self.dropout(x)
   127→        x = self.transformer(x, src_key_padding_mask=mel_mask)
   128→
   129→        mu = self.proj_mu(x)
   130→        logvar = self.proj_logvar(x)
   131→
   132→        # Reparameterization trick
   133→        if self.training:
   134→            std = torch.exp(0.5 * logvar)
   135→            eps = torch.randn_like(std)
   136→            z = mu + eps * std
   137→        else:
   138→            z = mu
   139→
   140→        return z, mu, logvar
   141→
   142→
   143→# ─── JEPA Predictor ─────────────────────────────────────────────────────────
   144→
   145→class JEPAPredictor(nn.Module):
   146→    """
   147→    Predicts next audio embedding given:
   148→    - Text context (from TextEncoder)
   149→    - Previous audio embeddings (from AudioEncoder)
   150→
   151→    Uses cross-attention to condition on text, and causal self-attention
   152→    over the audio embedding sequence.
   153→    """
   154→
   155→    def __init__(self, d_model=256, nhead=4, num_layers=6, dropout=0.1):
   156→        super().__init__()
   157→        self.d_model = d_model
   158→
   159→        self.audio_input_proj = nn.Linear(d_model, d_model)
   160→        self.pos_embed = SinusoidalPositionalEncoding(d_model, max_len=4096)
   161→        self.dropout = nn.Dropout(dropout)
   162→
   163→        # Transformer decoder with causal self-attention + cross-attention to text
   164→        decoder_layer = nn.TransformerDecoderLayer(
   165→            d_model=d_model, nhead=nhead,
   166→            dim_feedforward=d_model * 4,
   167→            dropout=dropout, batch_first=True,
   168→            activation="gelu",
   169→        )
   170→        self.transformer = nn.TransformerDecoder(decoder_layer, num_layers=num_layers)
   171→
   172→        # Output projection to predict next embedding
   173→        self.output_proj = nn.Sequential(
   174→            nn.Linear(d_model, d_model),
   175→            nn.GELU(),
   176→            nn.Linear(d_model, d_model),
   177→        )
   178→
   179→    def forward(self, audio_emb, text_emb, audio_mask=None, text_mask=None):
   180→        """
   181→        Args:
   182→            audio_emb: [B, T_audio, d_model] — audio embeddings (from encoder)
   183→            text_emb: [B, T_text, d_model] — text embeddings
   184→            audio_mask: [B, T_audio] bool, True = padding
   185→            text_mask: [B, T_text] bool, True = padding
   186→        Returns:
   187→            predicted: [B, T_audio, d_model] — predicted next embeddings
   188→        """
   189→        x = self.audio_input_proj(audio_emb)
   190→        x = self.pos_embed(x)
   191→        x = self.dropout(x)
   192→
   193→        # Causal mask: each position can only attend to previous positions
   194→        T = x.shape[1]
   195→        causal_mask = torch.triu(
   196→            torch.ones(T, T, device=x.device, dtype=torch.bool), diagonal=1
   197→        )
   198→
   199→        predicted = self.transformer(
   200→            x, text_emb,
   201→            tgt_mask=causal_mask,
   202→            tgt_key_padding_mask=audio_mask,
   203→            memory_key_padding_mask=text_mask,
   204→        )
   205→
   206→        predicted = self.output_proj(predicted)
   207→        return predicted
   208→
   209→    def _cached_mha(self, mha, q_input, kv_input, cache_k, cache_v):
   210→        """Run multi-head attention with KV cache.
   211→        q_input: [B, 1, d], kv_input: [B, 1, d] (new token for K,V)
   212→        cache_k, cache_v: [B, T_prev, d] projected keys/values
   213→        Returns: output [B, 1, d], updated cache_k, updated cache_v
   214→        """
   215→        d = mha.embed_dim
   216→        nhead = mha.num_heads
   217→        head_dim = d // nhead
   218→        B = q_input.shape[0]
   219→
   220→        # Project Q from q_input, K/V from kv_input
   221→        # in_proj_weight is [3d, d] = [Wq; Wk; Wv]
   222→        Wq, Wk, Wv = mha.in_proj_weight.chunk(3, dim=0)
   223→        bq, bk, bv = mha.in_proj_bias.chunk(3, dim=0)
   224→
   225→        q = F.linear(q_input, Wq, bq)        # [B, 1, d]
   226→        k_new = F.linear(kv_input, Wk, bk)    # [B, 1, d]
   227→        v_new = F.linear(kv_input, Wv, bv)    # [B, 1, d]
   228→
   229→        # Append to cache
   230→        if cache_k is not None and cache_k.shape[1] > 0:
   231→            k = torch.cat([cache_k, k_new], dim=1)
   232→            v = torch.cat([cache_v, v_new], dim=1)
   233→        else:
   234→            k = k_new
   235→            v = v_new
   236→
   237→        # Multi-head attention
   238→        q = q.view(B, 1, nhead, head_dim).transpose(1, 2)      # [B, nhead, 1, head_dim]
   239→        k_mh = k.view(B, -1, nhead, head_dim).transpose(1, 2)  # [B, nhead, T, head_dim]
   240→        v_mh = v.view(B, -1, nhead, head_dim).transpose(1, 2)
   241→
   242→        attn = torch.matmul(q, k_mh.transpose(-2, -1)) / (head_dim ** 0.5)
   243→        attn = F.softmax(attn, dim=-1)
   244→        out = torch.matmul(attn, v_mh)  # [B, nhead, 1, head_dim]
   245→
   246→        out = out.transpose(1, 2).contiguous().view(B, 1, d)
   247→        out = F.linear(out, mha.out_proj.weight, mha.out_proj.bias)
   248→
   249→        return out, k, v
   250→
   251→    def _cached_cross_attn(self, mha, q_input, memory, cache_k, cache_v, memory_mask=None):
   252→        """Cross-attention with cached memory K,V (computed once).
   253→        memory_mask: [B, T_text] bool, True = padding (applied as -inf)."""
   254→        d = mha.embed_dim
   255→        nhead = mha.num_heads
   256→        head_dim = d // nhead
   257→        B = q_input.shape[0]
   258→
   259→        Wq, Wk, Wv = mha.in_proj_weight.chunk(3, dim=0)
   260→        bq, bk, bv = mha.in_proj_bias.chunk(3, dim=0)
   261→
   262→        q = F.linear(q_input, Wq, bq)  # [B, 1, d]
   263→
   264→        if cache_k is None:
   265→            k = F.linear(memory, Wk, bk)  # [B, T_text, d]
   266→            v = F.linear(memory, Wv, bv)
   267→        else:
   268→            k, v = cache_k, cache_v
   269→
   270→        q = q.view(B, 1, nhead, head_dim).transpose(1, 2)
   271→        k_mh = k.view(B, -1, nhead, head_dim).transpose(1, 2)
   272→        v_mh = v.view(B, -1, nhead, head_dim).transpose(1, 2)
   273→
   274→        attn = torch.matmul(q, k_mh.transpose(-2, -1)) / (head_dim ** 0.5)
   275→        if memory_mask is not None:
   276→            # memory_mask: [B, T_text] → [B, 1, 1, T_text]
   277→            attn = attn.masked_fill(memory_mask.unsqueeze(1).unsqueeze(2), float('-inf'))
   278→        attn = F.softmax(attn, dim=-1)
   279→        out = torch.matmul(attn, v_mh)
   280→
   281→        out = out.transpose(1, 2).contiguous().view(B, 1, d)
   282→        out = F.linear(out, mha.out_proj.weight, mha.out_proj.bias)
   283→
   284→        return out, k, v
   285→
   286→    def init_cache(self, num_layers):
   287→        """Initialize empty KV cache."""
   288→        return {
   289→            'self_k': [None] * num_layers,
   290→            'self_v': [None] * num_layers,
   291→            'cross_k': [None] * num_layers,
   292→            'cross_v': [None] * num_layers,
   293→        }
   294→
   295→    def inference_step(self, new_emb, text_emb, step_idx, cache, text_mask=None):
   296→        """Process one token with KV cache.
   297→        new_emb: [B, 1, d_model] — raw audio embedding (not projected)
   298→        text_mask: [B, T_text] bool, True = padding
   299→        Returns: predicted [B, 1, d_model], updated cache
   300→        """
   301→        x = self.audio_input_proj(new_emb)
   302→        # Add positional encoding for this specific position
   303→        x = x + self.pos_embed.pe[:, step_idx:step_idx + 1]
   304→
   305→        for i, layer in enumerate(self.transformer.layers):
   306→            # Self-attention with cache (post-norm)
   307→            sa_out, cache['self_k'][i], cache['self_v'][i] = self._cached_mha(
   308→                layer.self_attn, x, x, cache['self_k'][i], cache['self_v'][i]
   309→            )
   310→            x = layer.norm1(x + layer.dropout1(sa_out))
   311→
   312→            # Cross-attention with cached text KV
   313→            ca_out, cache['cross_k'][i], cache['cross_v'][i] = self._cached_cross_attn(
   314→                layer.multihead_attn, x, text_emb, cache['cross_k'][i], cache['cross_v'][i],
   315→                memory_mask=text_mask,
   316→            )
   317→            x = layer.norm2(x + layer.dropout2(ca_out))
   318→
   319→            # FFN
   320→            ff_out = layer.linear2(layer.dropout(layer.activation(layer.linear1(x))))
   321→            x = layer.norm3(x + layer.dropout3(ff_out))
   322→
   323→        return self.output_proj(x), cache
   324→
   325→
   326→# ─── Positional Encoding ────────────────────────────────────────────────────
   327→
   328→class SinusoidalPositionalEncoding(nn.Module):
   329→    def __init__(self, d_model, max_len=8192):
   330→        super().__init__()
   331→        pe = torch.zeros(max_len, d_model)
   332→        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
   333→        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
   334→        pe[:, 0::2] = torch.sin(position * div_term)
   335→        pe[:, 1::2] = torch.cos(position * div_term)
   336→        self.register_buffer("pe", pe.unsqueeze(0))  # [1, max_len, d_model]
   337→
   338→    def forward(self, x):
   339→        return x + self.pe[:, :x.shape[1]]
   340→
   341→
   342→# ─── Mel Decoder ────────────────────────────────────────────────────────────
   343→
   344→class ResConvBlock(nn.Module):
   345→    """Residual conv block for mel decoder."""
   346→
   347→    def __init__(self, channels, kernel_size=3):
   348→        super().__init__()
   349→        self.net = nn.Sequential(
   350→            nn.Conv1d(channels, channels, kernel_size, padding=kernel_size // 2),
   351→            nn.GELU(),
   352→            nn.Conv1d(channels, channels, kernel_size, padding=kernel_size // 2),
   353→        )
   354→
   355→    def forward(self, x):
   356→        return x + self.net(x)
   357→
   358→
   359→class MelDecoder(nn.Module):
   360→    """Decodes latent embeddings back to mel spectrograms.
   361→    Improved architecture with residual blocks for sharper mel output."""
   362→
   363→    def __init__(self, d_model=256, n_mels=100, upsample_factor=4):
   364→        super().__init__()
   365→        hidden = d_model * 2  # 512 — more capacity
   366→
   367→        self.proj = nn.Sequential(
   368→            nn.Linear(d_model, hidden),
   369→            nn.GELU(),
   370→        )
   371→        if upsample_factor == 2:
   372→            self.upsample = nn.Sequential(
   373→                nn.ConvTranspose1d(hidden, hidden, kernel_size=4, stride=2, padding=1),
   374→                nn.GELU(),
   375→                ResConvBlock(hidden, kernel_size=3),
   376→            )
   377→        else:
   378→            self.upsample = nn.Sequential(
   379→                # 4x upsample (2x + 2x)
   380→                nn.ConvTranspose1d(hidden, hidden, kernel_size=4, stride=2, padding=1),
   381→                nn.GELU(),
   382→                ResConvBlock(hidden, kernel_size=3),
   383→                nn.ConvTranspose1d(hidden, hidden, kernel_size=4, stride=2, padding=1),
   384→                nn.GELU(),
   385→                ResConvBlock(hidden, kernel_size=3),
   386→            )
   387→        # Refinement + projection to mel bins
   388→        self.out = nn.Sequential(
   389→            nn.Conv1d(hidden, hidden, kernel_size=7, padding=3),
   390→            nn.GELU(),
   391→            ResConvBlock(hidden, kernel_size=5),
   392→            nn.Conv1d(hidden, n_mels, kernel_size=7, padding=3),
   393→        )
   394→
   395→    def forward(self, z):
   396→        """z: [B, T_down, d_model] → mel: [B, n_mels, T_up]"""
   397→        x = self.proj(z)
   398→        x = x.transpose(1, 2)  # [B, hidden, T_down]
   399→        x = self.upsample(x)   # [B, hidden, T_up]
   400→        return self.out(x)      # [B, n_mels, T_up]
   401→
   402→
   403→# ─── Multi-Resolution Spectral Loss ─────────────────────────────────────────
   404→
   405→class MultiResolutionSpectralLoss(nn.Module):
   406→    """Computes spectral convergence + log magnitude loss at multiple STFT resolutions.
   407→    This forces the mel decoder to produce sharp, well-defined spectral structure
   408→    rather than blurry averages that Vocos can't decode properly."""
   409→
   410→    def __init__(self, resolutions=((64, 16, 64), (128, 32, 128), (256, 64, 256))):
   411→        super().__init__()
   412→        # Each resolution is (n_fft, hop_length, win_length)
   413→        self.resolutions = resolutions
   414→
   415→    def _stft_loss(self, pred, target, n_fft, hop_length, win_length):
   416→        """Compute spectral convergence + log magnitude loss for one resolution."""
   417→        window = torch.hann_window(win_length, device=pred.device)
   418→
   419→        # pred/target: [B, n_mels, T] — treat as 1D signal per mel bin
   420→        # Flatten mel bins into batch: [B*n_mels, T]
   421→        B, M, T = pred.shape
   422→        pred_flat = pred.reshape(B * M, T)
   423→        target_flat = target.reshape(B * M, T)
   424→
   425→        pred_stft = torch.stft(pred_flat, n_fft, hop_length, win_length,
   426→                               window=window, return_complex=True)
   427→        target_stft = torch.stft(target_flat, n_fft, hop_length, win_length,
   428→                                 window=window, return_complex=True)
   429→
   430→        pred_mag = pred_stft.abs()
   431→        target_mag = target_stft.abs()
   432→
   433→        # Spectral convergence: Frobenius norm of difference / Frobenius norm of target
   434→        sc_loss = torch.norm(target_mag - pred_mag, p="fro") / (torch.norm(target_mag, p="fro") + 1e-8)
   435→
   436→        # Log magnitude loss
   437→        log_mag_loss = F.l1_loss(
   438→            torch.log(pred_mag + 1e-7),
   439→            torch.log(target_mag + 1e-7),
   440→        )
   441→
   442→        return sc_loss + log_mag_loss
   443→
   444→    def forward(self, pred, target):
   445→        """pred, target: [B, n_mels, T]"""
   446→        loss = torch.tensor(0.0, device=pred.device)
   447→        count = 0
   448→        for n_fft, hop, win in self.resolutions:
   449→            # Skip if T is too short for this resolution
   450→            if pred.shape[2] >= n_fft:
   451→                loss = loss + self._stft_loss(pred, target, n_fft, hop, win)
   452→                count += 1
   453→        if count > 0:
   454→            loss = loss / count
   455→        return loss
   456→
   457→
   458→# ─── Full LeWM TTS Model ────────────────────────────────────────────────────
   459→
   460→class LeWMTTS(nn.Module):
   461→    """
   462→    Complete LeWM TTS model.
   463→
   464→    Training: text + mel → audio embeddings → predict next embedding
   465→    Two losses:
   466→        1. Next-embedding prediction: MSE(predicted[t], target[t+1])
   467→        2. Gaussian regularizer: KL(q(z|x) || N(0, I))
   468→    """
   469→
   470→    def __init__(self, config):
   471→        super().__init__()
   472→        d_model = config.get("d_model", 256)
   473→        nhead = config.get("nhead", 4)
   474→        n_mels = config.get("n_mels", 100)
   475→        text_vocab = config.get("text_vocab_size", 256)
   476→        text_layers = config.get("text_encoder_layers", 4)
   477→        audio_layers = config.get("audio_encoder_layers", 4)
   478→        predictor_layers = config.get("predictor_layers", 6)
   479→        dropout = config.get("dropout", 0.1)
   480→        self.kl_weight = config.get("kl_weight", 0.01)
   481→        n_speakers = config.get("n_speakers", 1)
   482→
   483→        self.text_encoder = TextEncoder(
   484→            vocab_size=text_vocab, d_model=d_model, nhead=nhead,
   485→            num_layers=text_layers, dropout=dropout,
   486→        )
   487→        downsample_factor = config.get("downsample_factor", 4)
   488→        self.audio_encoder = AudioEncoder(
   489→            n_mels=n_mels, d_model=d_model, nhead=nhead,
   490→            num_layers=audio_layers, downsample_factor=downsample_factor,
   491→            dropout=dropout,
   492→        )
   493→        self.predictor = JEPAPredictor(
   494→            d_model=d_model, nhead=nhead,
   495→            num_layers=predictor_layers, dropout=dropout,
   496→        )
   497→
   498→        # Mel decoder: reconstructs mel from latent embeddings
   499→        # Needed so the latent space is grounded to actual audio
   500→        self.mel_decoder = MelDecoder(d_model=d_model, n_mels=n_mels, upsample_factor=downsample_factor)
   501→        self.recon_weight = config.get("recon_weight", 1.0)
   502→
   503→        # EMA target encoder — core JEPA component
   504→        # Provides slowly-moving prediction targets to prevent representation collapse
   505→        self.ema_audio_encoder = copy.deepcopy(self.audio_encoder)
   506→        for p in self.ema_audio_encoder.parameters():
   507→            p.requires_grad = False
   508→        self.ema_decay = config.get("ema_decay", 0.998)
   509→
   510→        # Speaker embedding (for multi-speaker)
   511→        self.n_speakers = n_speakers
   512→        if n_speakers > 1:
   513→            self.speaker_embed = nn.Embedding(n_speakers, d_model)
   514→        else:
   515→            self.speaker_embed = None
   516→
   517→        # Learnable start-of-sequence embedding for AR inference
   518→        self.start_emb = nn.Parameter(torch.randn(1, 1, d_model) * 0.02)
   519→
   520→        # Prediction weight — primary JEPA loss
   521→        self.pred_weight = config.get("pred_weight", 10.0)
   522→
   523→        # Input noise injection to bridge train/inference gap
   524→        self.input_noise = config.get("input_noise", 0.0)
   525→
   526→        # Free-bits KL: allow KL up to this threshold without penalty
   527→        self.free_bits = config.get("free_bits", 1.0)
   528→
   529→        # Multi-resolution spectral loss for sharper mel output
   530→        self.spectral_loss = MultiResolutionSpectralLoss()
   531→        self.spectral_weight = config.get("spectral_weight", 0.5)
   532→
   533→        self.config = config
   534→
   535→    def _add_speaker(self, text_emb, speaker_id):
   536→        """Add speaker embedding to text embeddings."""
   537→        if self.speaker_embed is not None and speaker_id is not None:
   538→            spk = self.speaker_embed(speaker_id).unsqueeze(1)  # [B, 1, d]
   539→            text_emb = text_emb + spk
   540→        return text_emb
   541→
   542→    @torch.no_grad()
   543→    def update_ema(self):
   544→        """Update EMA target encoder weights."""
   545→        for p_online, p_ema in zip(self.audio_encoder.parameters(),
   546→                                   self.ema_audio_encoder.parameters()):
   547→            p_ema.data.mul_(self.ema_decay).add_(p_online.data, alpha=1 - self.ema_decay)
   548→
   549→    def forward(self, mel, text_tokens, mel_mask=None, text_mask=None, speaker_id=None):
   550→        """
   551→        Args:
   552→            mel: [B, n_mels, T_mel]
   553→            text_tokens: [B, T_text]
   554→            mel_mask: [B, T_mel] bool
   555→            text_mask: [B, T_text] bool
   556→        Returns:
   557→            loss_dict with 'prediction_loss', 'kl_loss', 'total_loss'
   558→        """
   559→        # Encode text + speaker conditioning
   560→        text_emb = self.text_encoder(text_tokens, text_mask)
   561→        text_emb = self._add_speaker(text_emb, speaker_id)
   562→
   563→        # Online encoder → latent embeddings (used for predictor input + reconstruction)
   564→        z, mu, logvar = self.audio_encoder(mel, mel_mask)
   565→        B = z.shape[0]
   566→        T_down = z.shape[1]
   567→
   568→        # EMA target encoder → stable prediction targets (no gradient)
   569→        with torch.no_grad():
   570→            _, mu_ema, _ = self.ema_audio_encoder(mel, mel_mask)
   571→            target_emb = mu_ema  # [B, T_down, d] — all EMA targets
   572→
   573→        # Prepend learnable start embedding to input sequence
   574→        # Input:  [start_emb, z[0], z[1], ..., z[T-2]]  → T_down positions
   575→        # Target: [z_ema[0],  z_ema[1], ..., z_ema[T-1]] → T_down positions
   576→        # This trains the model to predict the first audio embedding from start_emb + text
   577→        start = self.start_emb.expand(B, -1, -1)  # [B, 1, d]
   578→        input_emb = torch.cat([start, z[:, :-1]], dim=1)  # [B, T_down, d]
   579→
   580→        # Input noise: only on positions 1+ (real embeddings), not start_emb
   581→        # Simulates the imperfect predictions seen during AR inference
   582→        if self.training and self.input_noise > 0:
   583→            noise = torch.zeros_like(input_emb)
   584→            noise[:, 1:] = torch.randn(B, T_down - 1, input_emb.shape[-1],
   585→                                        device=input_emb.device) * self.input_noise
   586→            input_emb = input_emb + noise
   587→
   588→        # Adjust masks for downsampled sequence
   589→        if mel_mask is not None:
   590→            ds_mask = mel_mask[:, ::self.audio_encoder.downsample_factor][:, :T_down]
   591→            # Predictor input mask: start_emb (pos 0) is always unmasked
   592→            start_mask = torch.zeros(B, 1, dtype=torch.bool, device=z.device)
   593→            pred_mask = torch.cat([start_mask, ds_mask[:, :-1]], dim=1)  # [B, T_down]
   594→            # Loss mask: follows original ds_mask for target positions
   595→            loss_mask = ds_mask
   596→        else:
   597→            pred_mask = None
   598→            loss_mask = None
   599→
   600→        predicted = self.predictor(input_emb, text_emb, pred_mask, text_mask)
   601→
   602→        # ─── Loss 1: Next-embedding prediction (MSE) against EMA targets ───
   603→        if loss_mask is not None:
   604→            valid = (~loss_mask).unsqueeze(-1)  # [B, T_down, 1]
   605→            prediction_loss = F.mse_loss(
   606→                predicted * valid, target_emb * valid, reduction="sum"
   607→            ) / (valid.sum() * predicted.shape[-1] + 1e-8)
   608→        else:
   609→            prediction_loss = F.mse_loss(predicted, target_emb)
   610→
   611→        # ─── Loss 2: Free-bits KL — only penalize KL above threshold ───
   612→        # Prevents both KL collapse (too low) and KL explosion (too high)
   613→        kl_per_dim = -0.5 * (1 + logvar - mu.pow(2) - logvar.exp())  # [B, T, d]
   614→        kl_free = torch.clamp(kl_per_dim - self.free_bits, min=0.0)
   615→        if mel_mask is not None:
   616→            T_down = mu.shape[1]
   617→            ds_mask = mel_mask[:, ::self.audio_encoder.downsample_factor][:, :T_down]
   618→            valid_kl = (~ds_mask).unsqueeze(-1).float()
   619→            kl_loss = (kl_free * valid_kl).sum() / (valid_kl.sum() * mu.shape[-1] + 1e-8)
   620→        else:
   621→            kl_loss = kl_free.mean()
   622→
   623→        # ─── Loss 3: Mel reconstruction ───
   624→        mel_recon = self.mel_decoder(z)  # [B, n_mels, T_up]
   625→        T_mel = mel.shape[2]
   626→        T_recon = mel_recon.shape[2]
   627→        if T_recon > T_mel:
   628→            mel_recon = mel_recon[:, :, :T_mel]
   629→        elif T_recon < T_mel:
   630→            mel = mel[:, :, :T_recon]
   631→            if mel_mask is not None:
   632→                mel_mask = mel_mask[:, :T_recon]
   633→
   634→        if mel_mask is not None:
   635→            valid_mel = (~mel_mask).unsqueeze(1).float()  # [B, 1, T]
   636→            recon_loss = F.l1_loss(
   637→                mel_recon * valid_mel, mel * valid_mel, reduction="sum"
   638→            ) / (valid_mel.sum() * mel.shape[1] + 1e-8)
   639→        else:
   640→            recon_loss = F.l1_loss(mel_recon, mel)
   641→
   642→        # ─── Loss 4: Multi-resolution spectral loss ───
   643→        if mel_mask is not None:
   644→            spectral_loss = self.spectral_loss(mel_recon * valid_mel, mel * valid_mel)
   645→        else:
   646→            spectral_loss = self.spectral_loss(mel_recon, mel)
   647→
   648→        total_loss = (self.pred_weight * prediction_loss
   649→                      + self.kl_weight * kl_loss
   650→                      + self.recon_weight * recon_loss
   651→                      + self.spectral_weight * spectral_loss)
   652→
   653→        return {
   654→            "total_loss": total_loss,
   655→            "prediction_loss": prediction_loss,
   656→            "kl_loss": kl_loss,
   657→            "recon_loss": recon_loss,
   658→            "spectral_loss": spectral_loss,
   659→        }
   660→
   661→    def encode_audio(self, mel):
   662→        """Encode mel to latent embeddings (for inference)."""
   663→        z, mu, _ = self.audio_encoder(mel)
   664→        return mu  # Use mean at inference
   665→
   666→    def predict_next(self, audio_emb, text_emb, text_mask=None):
   667→        """Predict next audio embedding autoregressively."""
   668→        predicted = self.predictor(audio_emb, text_emb, text_mask=text_mask)
   669→        return predicted[:, -1:]  # Only the last prediction
   670→
   671→    def init_ar_cache(self):
   672→        """Initialize KV cache for autoregressive inference."""
   673→        num_layers = len(self.predictor.transformer.layers)
   674→        return self.predictor.init_cache(num_layers)
   675→
   676→    def predict_next_cached(self, new_emb, text_emb, step_idx, cache, text_mask=None):
   677→        """Predict next embedding using KV cache. O(1) per step instead of O(n)."""
   678→        return self.predictor.inference_step(new_emb, text_emb, step_idx, cache, text_mask=text_mask)
   679→
   680→
   681→def count_parameters(model):
   682→    return sum(p.numel() for p in model.parameters() if p.requires_grad)
   683→
   684→
   685→def build_model(config=None):
   686→    if config is None:
   687→        config = {
   688→            "d_model": 256,
   689→            "nhead": 4,
   690→            "n_mels": 100,
   691→            "text_vocab_size": 256,
   692→            "text_encoder_layers": 4,
   693→            "audio_encoder_layers": 4,
   694→            "predictor_layers": 6,
   695→            "dropout": 0.1,
   696→            "kl_weight": 0.05,
   697→            "pred_weight": 10.0,
   698→            "ema_decay": 0.998,
   699→            "free_bits": 1.0,
   700→            "input_noise": 0.0,
   701→        }
   702→    model = LeWMTTS(config)
   703→    print(f"LeWM TTS model: {count_parameters(model)/1e6:.2f}M parameters")
   704→    return model, config
   705→

"""
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 copy
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
        if downsample_factor == 2:
            self.conv_pre = nn.Sequential(
                nn.Conv1d(n_mels, d_model, kernel_size=7, padding=3),
                nn.GELU(),
                nn.Conv1d(d_model, d_model, kernel_size=4, stride=2, padding=1),
                nn.GELU(),
            )  # Total stride = 2
        else:
            self.conv_pre = nn.Sequential(
                nn.Conv1d(n_mels, d_model, kernel_size=7, padding=3),
                nn.GELU(),
                nn.Conv1d(d_model, d_model, kernel_size=4, stride=2, padding=1),
                nn.GELU(),
                nn.Conv1d(d_model, d_model, kernel_size=4, stride=2, padding=1),
                nn.GELU(),
            )  # Total stride = 4

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

        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, upsample_factor=4):
        super().__init__()
        hidden = d_model * 2  # 512 — more capacity

        self.proj = nn.Sequential(
            nn.Linear(d_model, hidden),
            nn.GELU(),
        )
        if upsample_factor == 2:
            self.upsample = nn.Sequential(
                nn.ConvTranspose1d(hidden, hidden, kernel_size=4, stride=2, padding=1),
                nn.GELU(),
                ResConvBlock(hidden, kernel_size=3),
            )
        else:
            self.upsample = nn.Sequential(
                # 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=((64, 16, 64), (128, 32, 128), (256, 64, 256))):
        super().__init__()
        # Each resolution is (n_fft, hop_length, win_length)
        self.resolutions = resolutions

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

        # pred/target: [B, n_mels, T] — treat as 1D signal per mel bin
        # Flatten mel bins into batch: [B*n_mels, T]
        B, M, T = pred.shape
        pred_flat = pred.reshape(B * M, T)
        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,
        )
        downsample_factor = config.get("downsample_factor", 4)
        self.audio_encoder = AudioEncoder(
            n_mels=n_mels, d_model=d_model, nhead=nhead,
            num_layers=audio_layers, downsample_factor=downsample_factor,
            dropout=dropout,
        )
        self.predictor = JEPAPredictor(
            d_model=d_model, nhead=nhead,
            num_layers=predictor_layers, dropout=dropout,
        )

        # Mel decoder: 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, upsample_factor=downsample_factor)
        self.recon_weight = config.get("recon_weight", 1.0)

        # EMA target encoder — core JEPA component
        # Provides slowly-moving prediction targets to prevent representation collapse
        self.ema_audio_encoder = copy.deepcopy(self.audio_encoder)
        for p in self.ema_audio_encoder.parameters():
            p.requires_grad = False
        self.ema_decay = config.get("ema_decay", 0.998)

        # 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

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

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

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

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

        # Multi-resolution spectral loss for sharper mel output
        self.spectral_loss = 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

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

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

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

        # EMA target encoder → stable prediction targets (no gradient)
        with torch.no_grad():
            _, mu_ema, _ = self.ema_audio_encoder(mel, mel_mask)
            target_emb = mu_ema  # [B, T_down, d] — all EMA targets

        # Prepend learnable start embedding to input sequence
        # Input:  [start_emb, z[0], z[1], ..., z[T-2]]  → T_down positions
        # Target: [z_ema[0],  z_ema[1], ..., z_ema[T-1]] → T_down positions
        # This trains the model to predict the first audio embedding from start_emb + text
        start = self.start_emb.expand(B, -1, -1)  # [B, 1, d]
        input_emb = torch.cat([start, z[:, :-1]], dim=1)  # [B, T_down, d]

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

        # Adjust masks for downsampled sequence
        if mel_mask is not None:
            ds_mask = mel_mask[:, ::self.audio_encoder.downsample_factor][:, :T_down]
            # Predictor input mask: start_emb (pos 0) is always unmasked
            start_mask = torch.zeros(B, 1, dtype=torch.bool, device=z.device)
            pred_mask = torch.cat([start_mask, ds_mask[:, :-1]], dim=1)  # [B, T_down]
            # Loss mask: follows original ds_mask for target positions
            loss_mask = ds_mask
        else:
            pred_mask = None
            loss_mask = None

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

        # ─── Loss 1: Next-embedding prediction (MSE) against EMA targets ───
        if loss_mask is not None:
            valid = (~loss_mask).unsqueeze(-1)  # [B, T_down, 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: Free-bits KL — only penalize KL above threshold ───
        # Prevents both KL collapse (too low) and KL explosion (too high)
        kl_per_dim = -0.5 * (1 + logvar - mu.pow(2) - logvar.exp())  # [B, T, d]
        kl_free = torch.clamp(kl_per_dim - self.free_bits, min=0.0)
        if mel_mask is not None:
            T_down = mu.shape[1]
            ds_mask = mel_mask[:, ::self.audio_encoder.downsample_factor][:, :T_down]
            valid_kl = (~ds_mask).unsqueeze(-1).float()
            kl_loss = (kl_free * valid_kl).sum() / (valid_kl.sum() * mu.shape[-1] + 1e-8)
        else:
            kl_loss = kl_free.mean()

        # ─── Loss 3: Mel reconstruction ───
        mel_recon = self.mel_decoder(z)  # [B, n_mels, T_up]
        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 ───
        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 = (self.pred_weight * 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.05,
            "pred_weight": 10.0,
            "ema_decay": 0.998,
            "free_bits": 1.0,
            "input_noise": 0.0,
        }
    model = LeWMTTS(config)
    print(f"LeWM TTS model: {count_parameters(model)/1e6:.2f}M parameters")
    return model, config


Let me now carefully audit every aspect of this from a researcher's perspective. Let me look for subtle bugs and potential issues:

1. **Forward pass flow**:
   - text_emb comes from text_encoder, used in predictor cross-attention. ✓
   - z (reparameterized) is used for predictor input AND mel reconstruction. ✓
   - mu_ema (EMA mean, no sampling) is used for prediction targets. ✓
   - start_emb is prepended to input, targets are all T_down positions from EMA. ✓

2. **Mask handling**:
   - ds_mask: mel_mask downsampled by stride factor. ✓
   - pred_mask: [False, ds_mask[0], ..., ds_mask[T-2]] — start_emb is always unmasked, then follows the mask of z[0:T-2]. ✓
   - loss_mask: ds_mask — masks target positions that are padding. ✓
   - Wait, I'm catching something here: the pred_mask controls which positions the predictor's self-attention can see, but there's a potential mismatch at position i=0 where start_emb is always unmasked regardless of whether the first actual token should be masked.

Actually, looking more carefully at the indexing: at position i in pred_mask, we have start_emb at i=0 (always unmasked), then z[k-1] at position k where it's valid if ds_mask[k-1] is False. For loss_mask, we compute loss on target z_ema[k] when ds_mask[k] is False. This alignment checks out.

Now I'm noticing a potential issue with mel mutation — the code slices mel to T_recon, which could modify the original tensor if it's not handled carefully. Looking at the sequence of operations, the mel_mask slicing at line 632 happens after the ds_mask computations, so those use the original unmodified mask—which is correct. The sliced mel_mask then properly aligns with the shortened mel tensor for the reconstruction loss, so that's also fine. Now I need to verify whether this mel_mask modification affects the KL computation that occurs earlier. BEFORE the mel truncation at 630-632. So it uses the original mel_mask. ✓

Actually, I realize there could be a different subtle issue. If we enter the branch `T_recon < T_mel` at line 629, we slice mel to be shorter. But mel was the input parameter. If we're in a training loop and the dataloader reuses tensors... actually no, in PyTorch, slicing creates a new tensor (a view). The original tensor is not modified. ✓

4. **Reconstruction uses z (reparameterized), not mu**:
   - Line 624: `mel_recon = self.mel_decoder(z)` — uses z which has sampling noise during training. This is correct for a VAE — the decoder should learn to reconstruct from noisy samples. ✓

5. **EMA encoder uses training/eval mode**:
   - The ema_audio_encoder runs within `torch.no_grad()`, but the model's training flag is still True, so dropout layers are active. This means the EMA targets get random variation from dropout, which isn't ideal for a target encoder that should produce deterministic outputs like in I-JEPA. Let me check the AudioEncoder forward to see exactly where dropout is applied. The transformer encoder layers also have dropout internally.

And in training mode, dropout is active. So the EMA encoder outputs would be noisy even though we use mu_ema (not z_ema with reparameterization). The mu projection itself doesn't have dropout, but the features feeding into it do.

This IS a potential issue. We should put the EMA encoder in eval mode. Let me fix this.

During inference, both encoders are in eval mode so dropout is disabled and predictions are deterministic. Actually, the EMA encoder is only used during training for computing targets anyway — at inference we just use the online encoder and predictor. ✓

The input noise injection looks correct too — it creates a zero tensor, fills positions 1+ with scaled random noise, and adds it to the input embeddings, leaving the start token unperturbed. ✓

Now I'm checking the free-bits KL calculation... decay)` — this implements the standard exponential moving average where the EMA parameters blend the online parameters with exponential decay, applied after the optimizer step so targets lag appropriately. For the predictor's positional encoding with start_emb, the start embedding occupies position 0 during training just as it does during inference, maintaining consistency across the predicted sequence positions. Now checking the causal mask...

The causal mask correctly restricts each position to attend only to itself and previous positions, so position 0 (start_emb) sees only itself, position 1 sees start_emb and itself, matching the inference behavior where each step has access to all prior context.

I found the main issue: the EMA encoder runs with dropout enabled during training. I need to wrap the EMA forward call with eval mode and restore it afterward, though actually since it's already inside a `torch.no_grad()` block, I should just set it to eval mode there.

Setting eval/train on a submodule only affects that submodule, so that's fine. I should also check the inference_engine.py file to see if it needs updates. LayerNorm doesn't change behavior between train and eval modes, and the AudioEncoder uses Conv1d without batch norm, so dropout is the only thing that matters here.


Good, I'm reading through line by line. I found a subtle bug. Let me also check the inference engine.