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

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

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

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

# ─── FiLM Conditioning ────────────────────────────────────────────────────────

class FiLMConditioner(nn.Module):
    """Feature-wise Linear Modulation: x → gamma * x + beta.
    Learns per-speaker scale (gamma) and shift (beta) from a conditioning vector.
    Much more expressive than simple addition — can reshape the entire distribution."""

    def __init__(self, cond_dim, channels):
        super().__init__()
        self.proj = nn.Linear(cond_dim, channels * 2)
        # Initialize near-identity: gamma≈1, beta≈0
        nn.init.zeros_(self.proj.weight)
        nn.init.zeros_(self.proj.bias)
        self.proj.bias.data[:channels] = 1.0  # gamma starts at 1

    def forward(self, x, cond):
        """
        Args:
            x: [B, T, d] or [B, d, T] (set channel_last accordingly)
            cond: [B, d_cond] — speaker conditioning vector
        Returns: modulated x, same shape
        """
        params = self.proj(cond)  # [B, channels*2]
        gamma, beta = params.chunk(2, dim=-1)  # each [B, channels]

        if x.dim() == 3:
            if x.shape[-1] == gamma.shape[-1]:
                # [B, T, d] format
                gamma = gamma.unsqueeze(1)  # [B, 1, d]
                beta = beta.unsqueeze(1)
            else:
                # [B, d, T] format (conv layers)
                gamma = gamma.unsqueeze(2)  # [B, d, 1]
                beta = beta.unsqueeze(2)
        return gamma * x + beta

# ─── 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.
    Returns intermediate layer outputs for multi-scale EMA targets.
    """

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

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

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

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

        # Layer norm after each layer (for multi-scale targets)
        self.layer_norms = nn.ModuleList([
            nn.LayerNorm(d_model) for _ in range(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, return_intermediates=False):
        """
        Args:
            mel: [B, n_mels, T_mel] — log mel spectrogram
            mel_mask: [B, T_mel] bool, True = padding (before downsampling)
            return_intermediates: if True, return list of per-layer outputs
        Returns:
            z: [B, T_down, d_model] — sampled latent embeddings
            mu: [B, T_down, d_model]
            logvar: [B, T_down, d_model]
            intermediates: list of [B, T_down, d_model] (only if return_intermediates)
        """
        # 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)

        intermediates = []
        for i, (layer, norm) in enumerate(zip(self.layers, self.layer_norms)):
            x = layer(x, src_key_padding_mask=mel_mask)
            if return_intermediates:
                intermediates.append(norm(x))

        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

        if return_intermediates:
            return z, mu, logvar, intermediates
        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 FiLMResConvBlock(nn.Module):
    """Residual conv block with FiLM speaker conditioning.
    When no speaker conditioning is provided, behaves like a normal ResConvBlock."""

    def __init__(self, channels, kernel_size=3, cond_dim=None):
        super().__init__()
        self.conv1 = nn.Conv1d(channels, channels, kernel_size, padding=kernel_size // 2)
        self.conv2 = nn.Conv1d(channels, channels, kernel_size, padding=kernel_size // 2)
        self.has_film = cond_dim is not None
        if self.has_film:
            self.film = FiLMConditioner(cond_dim, channels)

    def forward(self, x, spk_cond=None):
        """x: [B, C, T], spk_cond: [B, d_cond] or None"""
        h = F.gelu(self.conv1(x))
        h = self.conv2(h)
        if self.has_film and spk_cond is not None:
            h = self.film(h, spk_cond)  # FiLM on residual branch
        return x + h

class MelDecoder(nn.Module):
    """Decodes latent embeddings back to mel spectrograms.
    Speaker-conditioned via FiLM in residual blocks."""

    def __init__(self, d_model=256, n_mels=100, upsample_factor=4, n_speakers=1):
        super().__init__()
        hidden = d_model * 2  # 512
        cond_dim = d_model if n_speakers > 1 else None

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

        # FiLM on input projection (before conv)
        self.has_film = n_speakers > 1
        if self.has_film:
            self.input_film = FiLMConditioner(d_model, hidden)

        if upsample_factor == 2:
            self.up1 = nn.ConvTranspose1d(hidden, hidden, kernel_size=4, stride=2, padding=1)
            self.res1 = FiLMResConvBlock(hidden, kernel_size=3, cond_dim=cond_dim)
            self.up2 = None
            self.res2 = None
        else:
            self.up1 = nn.ConvTranspose1d(hidden, hidden, kernel_size=4, stride=2, padding=1)
            self.res1 = FiLMResConvBlock(hidden, kernel_size=3, cond_dim=cond_dim)
            self.up2 = nn.ConvTranspose1d(hidden, hidden, kernel_size=4, stride=2, padding=1)
            self.res2 = FiLMResConvBlock(hidden, kernel_size=3, cond_dim=cond_dim)

        # Refinement
        self.refine_conv = nn.Conv1d(hidden, hidden, kernel_size=7, padding=3)
        self.refine_res = FiLMResConvBlock(hidden, kernel_size=5, cond_dim=cond_dim)
        self.out_conv = nn.Conv1d(hidden, n_mels, kernel_size=7, padding=3)

    def forward(self, z, spk_cond=None):
        """z: [B, T_down, d_model], spk_cond: [B, d_model] or None → mel: [B, n_mels, T_up]"""
        x = self.proj(z)  # [B, T_down, hidden]

        if self.has_film and spk_cond is not None:
            x = self.input_film(x, spk_cond)  # FiLM on [B, T, hidden]

        x = x.transpose(1, 2)  # [B, hidden, T_down]

        x = F.gelu(self.up1(x))
        x = self.res1(x, spk_cond)
        if self.up2 is not None:
            x = F.gelu(self.up2(x))
            x = self.res2(x, spk_cond)

        x = F.gelu(self.refine_conv(x))
        x = self.refine_res(x, spk_cond)
        return self.out_conv(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 with speaker-preserving JEPA.

    Speaker fidelity is enforced through 4 mechanisms:
      1. FiLM conditioning (scale+shift) on latents and targets
      2. FiLM-conditioned MelDecoder (speaker-aware reconstruction)
      3. Multi-scale EMA targets (retain acoustic detail across layers)
      4. Cosine + MSE prediction loss (preserves direction = speaker identity)
      5. Speaker consistency classifier on raw encoder output
    """

    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: speaker-conditioned reconstruction
        self.mel_decoder = MelDecoder(
            d_model=d_model, n_mels=n_mels,
            upsample_factor=downsample_factor, n_speakers=n_speakers,
        )
        self.recon_weight = config.get("recon_weight", 1.0)

        # EMA target encoder — core JEPA component
        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)

        # Multi-scale target: which layers to average for EMA targets
        # data2vec-style: average top-K layers instead of just the last
        self.ema_target_layers = config.get("ema_target_layers", None)  # None = all layers
        self.ema_target_proj = nn.Linear(d_model, d_model)  # project averaged intermediates

        # Speaker embedding + FiLM conditioning (for multi-speaker)
        self.n_speakers = n_speakers
        if n_speakers > 1:
            self.speaker_embed = nn.Embedding(n_speakers, d_model)
            # FiLM for latent space (replaces additive injection)
            self.latent_film = FiLMConditioner(d_model, d_model)
            self.target_film = FiLMConditioner(d_model, d_model)
            # Speaker classifier — forces encoder to retain speaker info
            self.speaker_classifier = nn.Sequential(
                nn.Linear(d_model, d_model),
                nn.GELU(),
                nn.Linear(d_model, n_speakers),
            )
        else:
            self.speaker_embed = None
            self.latent_film = None
            self.target_film = None
            self.speaker_classifier = 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)

        # Cosine similarity weight in prediction loss
        self.cosine_weight = config.get("cosine_weight", 1.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_fn = MultiResolutionSpectralLoss()
        self.spectral_weight = config.get("spectral_weight", 0.5)

        # Speaker consistency loss weight
        self.speaker_weight = config.get("speaker_weight", 1.0)

        self.config = config

    def _get_speaker_cond(self, speaker_id):
        """Get speaker conditioning vector. Returns [B, d_model] or None."""
        if self.speaker_embed is not None and speaker_id is not None:
            return self.speaker_embed(speaker_id)  # [B, d_model]
        return None

    def _apply_speaker_film(self, x, spk_cond, film_module):
        """Apply FiLM conditioning to tensor. No-op if no speaker."""
        if film_module is not None and spk_cond is not None:
            return film_module(x, spk_cond)
        return x

    @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 _compute_multiscale_targets(self, mel, mel_mask):
        """Compute multi-scale EMA targets by averaging intermediate layer outputs.
        data2vec-style: captures both low-level acoustic detail (early layers)
        and high-level semantic structure (late layers)."""
        with torch.no_grad():
            self.ema_audio_encoder.eval()
            _, mu_ema, _, intermediates = self.ema_audio_encoder(
                mel, mel_mask, return_intermediates=True
            )
            self.ema_audio_encoder.train()

            # Select which layers to average
            if self.ema_target_layers is not None:
                selected = [intermediates[i] for i in self.ema_target_layers]
            else:
                selected = intermediates  # all layers

            # Average across selected layers
            target_emb = torch.stack(selected, dim=0).mean(dim=0)  # [B, T, d]

        # Project through learned projection (allows model to adapt target space)
        target_emb = self.ema_target_proj(target_emb)
        return target_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
            speaker_id: [B] long tensor of speaker IDs (optional)
        Returns:
            loss_dict with all individual losses
        """
        # Get speaker conditioning
        spk_cond = self._get_speaker_cond(speaker_id)  # [B, d] or None

        # Encode text + FiLM speaker conditioning on text
        text_emb = self.text_encoder(text_tokens, text_mask)
        if spk_cond is not None:
            # FiLM on text embeddings (replaces simple addition)
            text_emb = self._apply_speaker_film(text_emb, spk_cond, self.latent_film)

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

        # Save raw mu before speaker conditioning — for speaker classifier
        mu_raw = mu

        # Multi-scale EMA targets (data2vec-style: average layer outputs)
        target_emb = self._compute_multiscale_targets(mel, mel_mask)

        # Apply FiLM speaker conditioning to targets and encoder output
        # This makes prediction loss require speaker knowledge
        target_emb = self._apply_speaker_film(target_emb, spk_cond, self.target_film)
        z = self._apply_speaker_film(z, spk_cond, self.latent_film)

        # Prepend learnable start embedding to input sequence
        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
        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]
            start_mask = torch.zeros(B, 1, dtype=torch.bool, device=z.device)
            pred_mask = torch.cat([start_mask, ds_mask[:, :-1]], dim=1)
            loss_mask = ds_mask
        else:
            pred_mask = None
            loss_mask = None

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

        # ─── Loss 1: Prediction loss (MSE + cosine) against multi-scale EMA targets ───
        if loss_mask is not None:
            valid = (~loss_mask).unsqueeze(-1).float()  # [B, T_down, 1]
            n_valid = valid.sum() * predicted.shape[-1] + 1e-8

            # MSE component
            mse_loss = (((predicted - target_emb) ** 2) * valid).sum() / n_valid

            # Cosine component — preserves direction (speaker identity)
            pred_flat = (predicted * valid).reshape(-1, predicted.shape[-1])
            tgt_flat = (target_emb * valid).reshape(-1, target_emb.shape[-1])
            # Only compute on non-padding positions
            valid_flat = valid.squeeze(-1).reshape(-1) > 0
            if valid_flat.any():
                cosine_loss = 1.0 - F.cosine_similarity(
                    pred_flat[valid_flat], tgt_flat[valid_flat], dim=-1
                ).mean()
            else:
                cosine_loss = torch.tensor(0.0, device=mel.device)
        else:
            mse_loss = F.mse_loss(predicted, target_emb)
            cosine_loss = 1.0 - F.cosine_similarity(
                predicted.reshape(-1, predicted.shape[-1]),
                target_emb.reshape(-1, target_emb.shape[-1]),
                dim=-1,
            ).mean()

        prediction_loss = mse_loss + self.cosine_weight * cosine_loss

        # ─── Loss 2: Free-bits KL ───
        kl_per_dim = -0.5 * (1 + logvar - mu.pow(2) - logvar.exp())
        kl_free = torch.clamp(kl_per_dim - self.free_bits, min=0.0)
        if mel_mask is not None:
            T_down_kl = mu.shape[1]
            ds_mask_kl = mel_mask[:, ::self.audio_encoder.downsample_factor][:, :T_down_kl]
            valid_kl = (~ds_mask_kl).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 (speaker-conditioned decoder) ───
        mel_recon = self.mel_decoder(z, spk_cond)  # [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()
            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_fn(mel_recon * valid_mel, mel * valid_mel)
        else:
            spectral_loss = self.spectral_loss_fn(mel_recon, mel)

        # ─── Loss 5: Speaker consistency on raw encoder output ───
        speaker_loss = torch.tensor(0.0, device=mel.device)
        if self.speaker_classifier is not None and speaker_id is not None:
            if mel_mask is not None:
                ds_mask_spk = mel_mask[:, ::self.audio_encoder.downsample_factor][:, :T_down]
                valid_spk = (~ds_mask_spk).unsqueeze(-1).float()
                mu_pooled = (mu_raw * valid_spk).sum(dim=1) / (valid_spk.sum(dim=1) + 1e-8)
            else:
                mu_pooled = mu_raw.mean(dim=1)
            spk_logits = self.speaker_classifier(mu_pooled)
            speaker_loss = F.cross_entropy(spk_logits, speaker_id)

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

        return {
            "total_loss": total_loss,
            "prediction_loss": prediction_loss,
            "mse_loss": mse_loss,
            "cosine_loss": cosine_loss,
            "kl_loss": kl_loss,
            "recon_loss": recon_loss,
            "spectral_loss": spectral_loss,
            "speaker_loss": speaker_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,
            "cosine_weight": 1.0,
            "ema_decay": 0.998,
            "free_bits": 2.0,
            "input_noise": 0.0,
        }
    model = LeWMTTS(config)
    print(f"LeWM TTS model: {count_parameters(model)/1e6:.2f}M parameters")
    return model, config

- TextEncoder: Character-level transformer for Hindi/Devanagari text

+    - Returns intermediate layer outputs for multi-scale targets

-  - Two losses: next-embedding prediction (MSE) + Gaussian regularizer (KL to N(0,I))

+  - SpeakerConditioner: FiLM-based speaker conditioning (scale + shift)

+  - MelDecoder: FiLM-conditioned decoder with speaker-aware residual blocks

+  - Losses: prediction (MSE + cosine), KL, reconstruction, spectral, speaker consistency

+# ─── FiLM Conditioning ────────────────────────────────────────────────────────

+    """Feature-wise Linear Modulation: x → gamma * x + beta.

+    Learns per-speaker scale (gamma) and shift (beta) from a conditioning vector.

+    Much more expressive than simple addition — can reshape the entire distribution."""

+    def __init__(self, cond_dim, channels):

+        self.proj = nn.Linear(cond_dim, channels * 2)

+        # Initialize near-identity: gamma≈1, beta≈0

+        nn.init.zeros_(self.proj.weight)

+        self.proj.bias.data[:channels] = 1.0  # gamma starts at 1

+            x: [B, T, d] or [B, d, T] (set channel_last accordingly)

+            cond: [B, d_cond] — speaker conditioning vector

+        Returns: modulated x, same shape

+        params = self.proj(cond)  # [B, channels*2]

+        gamma, beta = params.chunk(2, dim=-1)  # each [B, channels]

+            if x.shape[-1] == gamma.shape[-1]:

+                gamma = gamma.unsqueeze(1)  # [B, 1, d]

+                beta = beta.unsqueeze(1)

+                # [B, d, T] format (conv layers)

+                gamma = gamma.unsqueeze(2)  # [B, d, 1]

+                beta = beta.unsqueeze(2)

Encodes mel spectrograms into a sequence of audio embeddings.

Uses 1D CNN for downsampling + Transformer for contextualization.

+    Returns intermediate layer outputs for multi-scale EMA targets.

def __init__(self, n_mels=100, d_model=256, nhead=4, num_layers=4,

# CNN downsampling: reduces temporal resolution by downsample_factor

self.pos_embed = SinusoidalPositionalEncoding(d_model, max_len=4096)

-        encoder_layer = nn.TransformerEncoderLayer(

-            d_model=d_model, nhead=nhead,

-            dim_feedforward=d_model * 4,

-            dropout=dropout, batch_first=True,

-        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)

+        # Build layers individually so we can capture intermediate outputs

+            nn.TransformerEncoderLayer(

+                d_model=d_model, nhead=nhead,

+                dim_feedforward=d_model * 4,

+                dropout=dropout, batch_first=True,

+        # Layer norm after each layer (for multi-scale targets)

+        self.layer_norms = nn.ModuleList([

+            nn.LayerNorm(d_model) for _ in range(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):

+    def forward(self, mel, mel_mask=None, return_intermediates=False):

mel: [B, n_mels, T_mel] — log mel spectrogram

mel_mask: [B, T_mel] bool, True = padding (before downsampling)

+            return_intermediates: if True, return list of per-layer outputs

z: [B, T_down, d_model] — sampled latent embeddings

+            intermediates: list of [B, T_down, d_model] (only if return_intermediates)

# CNN encoding: [B, n_mels, T] -> [B, d_model, T//4]

-        x = self.transformer(x, src_key_padding_mask=mel_mask)

+        for i, (layer, norm) in enumerate(zip(self.layers, self.layer_norms)):

+            x = layer(x, src_key_padding_mask=mel_mask)

+                intermediates.append(norm(x))

+            return z, mu, logvar, intermediates

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

+    """Residual conv block with FiLM speaker conditioning.

+    When no speaker conditioning is provided, behaves like a normal ResConvBlock."""

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

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

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

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

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

+        self.has_film = cond_dim is not None

+            self.film = FiLMConditioner(cond_dim, channels)

+    def forward(self, x, spk_cond=None):

+        """x: [B, C, T], spk_cond: [B, d_cond] or None"""

+        if self.has_film and spk_cond is not None:

+            h = self.film(h, spk_cond)  # FiLM on residual branch

"""Decodes latent embeddings back to mel spectrograms.

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

+    Speaker-conditioned via FiLM in residual blocks."""

-    def __init__(self, d_model=256, n_mels=100, upsample_factor=4):

+    def __init__(self, d_model=256, n_mels=100, upsample_factor=4, n_speakers=1):

-        hidden = d_model * 2  # 512 — more capacity

+        cond_dim = d_model if n_speakers > 1 else None

+        # FiLM on input projection (before conv)

+            self.input_film = FiLMConditioner(d_model, hidden)

-            self.upsample = nn.Sequential(

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

-                ResConvBlock(hidden, kernel_size=3),

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

+            self.res1 = FiLMResConvBlock(hidden, kernel_size=3, cond_dim=cond_dim)

-                # 4x upsample (2x + 2x)

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

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

+            self.res2 = FiLMResConvBlock(hidden, kernel_size=3, cond_dim=cond_dim)

-        """z: [B, T_down, d_model] → mel: [B, n_mels, T_up]"""

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

+        self.refine_res = FiLMResConvBlock(hidden, kernel_size=5, cond_dim=cond_dim)

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

+    def forward(self, z, spk_cond=None):

+        """z: [B, T_down, d_model], spk_cond: [B, d_model] or None → mel: [B, n_mels, T_up]"""

+        x = self.proj(z)  # [B, T_down, hidden]

+        if self.has_film and spk_cond is not None:

+            x = self.input_film(x, spk_cond)  # FiLM on [B, T, hidden]

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]

+        x = F.gelu(self.refine_conv(x))

+        x = self.refine_res(x, spk_cond)

+        return self.out_conv(x)  # [B, n_mels, T_up]

+    Complete LeWM TTS model with speaker-preserving JEPA.

-    Training: text + mel → audio embeddings → predict next embedding

-        1. Next-embedding prediction: MSE(predicted[t], target[t+1])

-        2. Gaussian regularizer: KL(q(z|x) || N(0, I))

+    Speaker fidelity is enforced through 4 mechanisms:

+      1. FiLM conditioning (scale+shift) on latents and targets

+      2. FiLM-conditioned MelDecoder (speaker-aware reconstruction)

+      3. Multi-scale EMA targets (retain acoustic detail across layers)

+      4. Cosine + MSE prediction loss (preserves direction = speaker identity)

+      5. Speaker consistency classifier on raw encoder output

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

+        # Mel decoder: speaker-conditioned reconstruction

+            d_model=d_model, n_mels=n_mels,

+            upsample_factor=downsample_factor, n_speakers=n_speakers,

self.recon_weight = config.get("recon_weight", 1.0)

# EMA target encoder — core JEPA component

-        # Provides slowly-moving prediction targets to prevent representation collapse

for p in self.ema_audio_encoder.parameters():

self.ema_decay = config.get("ema_decay", 0.998)

-        # Speaker embedding (for multi-speaker)

+        # Multi-scale target: which layers to average for EMA targets

+        # data2vec-style: average top-K layers instead of just the last

+        self.ema_target_layers = config.get("ema_target_layers", None)  # None = all layers

+        self.ema_target_proj = nn.Linear(d_model, d_model)  # project averaged intermediates

+        # Speaker embedding + FiLM conditioning (for multi-speaker)

self.speaker_embed = nn.Embedding(n_speakers, d_model)

-            # Speaker classifier on audio encoder output — forces encoder to retain speaker info

-            self.speaker_classifier = nn.Linear(d_model, n_speakers)

+            # FiLM for latent space (replaces additive injection)

+            self.latent_film = FiLMConditioner(d_model, d_model)

+            self.target_film = FiLMConditioner(d_model, d_model)

+            # Speaker classifier — forces encoder to retain speaker info

+            self.speaker_classifier = nn.Sequential(

+                nn.Linear(d_model, d_model),

+                nn.Linear(d_model, n_speakers),

# Learnable start-of-sequence embedding for AR inference

self.pred_weight = config.get("pred_weight", 10.0)

+        # Cosine similarity weight in prediction loss

+        self.cosine_weight = config.get("cosine_weight", 1.0)

# Input noise injection to bridge train/inference gap

self.input_noise = config.get("input_noise", 0.0)

self.free_bits = config.get("free_bits", 1.0)

# Multi-resolution spectral loss for sharper mel output

-        self.spectral_loss = MultiResolutionSpectralLoss()

+        self.spectral_loss_fn = MultiResolutionSpectralLoss()

self.spectral_weight = config.get("spectral_weight", 0.5)

-    def _add_speaker(self, text_emb, speaker_id):

-        """Add speaker embedding to text embeddings."""

+    def _get_speaker_cond(self, speaker_id):

+        """Get speaker conditioning vector. Returns [B, d_model] or None."""

if self.speaker_embed is not None and speaker_id is not None:

-            spk = self.speaker_embed(speaker_id).unsqueeze(1)  # [B, 1, d]

+            return self.speaker_embed(speaker_id)  # [B, d_model]

+    def _apply_speaker_film(self, x, spk_cond, film_module):

+        """Apply FiLM conditioning to tensor. No-op if no speaker."""

+        if film_module is not None and spk_cond is not None:

+            return film_module(x, spk_cond)

"""Update EMA target encoder weights."""

+    def _compute_multiscale_targets(self, mel, mel_mask):

+        """Compute multi-scale EMA targets by averaging intermediate layer outputs.

+        data2vec-style: captures both low-level acoustic detail (early layers)

+        and high-level semantic structure (late layers)."""

+            self.ema_audio_encoder.eval()

+            _, mu_ema, _, intermediates = self.ema_audio_encoder(

+                mel, mel_mask, return_intermediates=True

+            self.ema_audio_encoder.train()

+            # Select which layers to average

+            if self.ema_target_layers is not None:

+                selected = [intermediates[i] for i in self.ema_target_layers]

+                selected = intermediates  # all layers

+            # Average across selected layers

+            target_emb = torch.stack(selected, dim=0).mean(dim=0)  # [B, T, d]

+        # Project through learned projection (allows model to adapt target space)

+        target_emb = self.ema_target_proj(target_emb)

+            speaker_id: [B] long tensor of speaker IDs (optional)

-            loss_dict with 'prediction_loss', 'kl_loss', 'total_loss'

+            loss_dict with all individual losses

-        # Encode text + speaker conditioning

+        spk_cond = self._get_speaker_cond(speaker_id)  # [B, d] or None

+        # Encode text + FiLM speaker conditioning on text

text_emb = self.text_encoder(text_tokens, text_mask)

-        text_emb = self._add_speaker(text_emb, speaker_id)

+            # FiLM on text embeddings (replaces simple addition)

+            text_emb = self._apply_speaker_film(text_emb, spk_cond, self.latent_film)

-        # Online encoder → latent embeddings (used for predictor input + reconstruction)

+        # Online encoder → latent embeddings

z, mu, logvar = self.audio_encoder(mel, mel_mask)

-        # EMA target encoder → stable prediction targets (no gradient)

-        # Force eval mode so dropout is disabled — EMA targets must be deterministic

-            self.ema_audio_encoder.eval()

-            _, mu_ema, _ = self.ema_audio_encoder(mel, mel_mask)

-            self.ema_audio_encoder.train()

-            target_emb = mu_ema  # [B, T_down, d] — all EMA targets

-        # Save raw encoder mu before speaker injection — used for speaker classifier

-        # so the encoder itself must learn speaker info, not just parrot the added embedding

+        # Save raw mu before speaker conditioning — for speaker classifier

-        # Inject speaker into targets and encoder output so prediction loss

-        # requires speaker discrimination — two speakers saying the same word

-        # produce different targets, forcing the predictor to use speaker info

-        if self.speaker_embed is not None and speaker_id is not None:

-            spk = self.speaker_embed(speaker_id).unsqueeze(1)  # [B, 1, d]

-            target_emb = target_emb + spk

+        # Multi-scale EMA targets (data2vec-style: average layer outputs)

+        target_emb = self._compute_multiscale_targets(mel, mel_mask)

+        # Apply FiLM speaker conditioning to targets and encoder output

+        # This makes prediction loss require speaker knowledge

+        target_emb = self._apply_speaker_film(target_emb, spk_cond, self.target_film)

+        z = self._apply_speaker_film(z, spk_cond, self.latent_film)

# 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[:, 1:] = torch.randn(B, T_down - 1, input_emb.shape[-1],

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

+            pred_mask = torch.cat([start_mask, ds_mask[:, :-1]], dim=1)

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

-        # ─── Loss 1: Next-embedding prediction (MSE) against EMA targets ───

+        # ─── Loss 1: Prediction loss (MSE + cosine) against multi-scale EMA targets ───

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

+            valid = (~loss_mask).unsqueeze(-1).float()  # [B, T_down, 1]

+            n_valid = valid.sum() * predicted.shape[-1] + 1e-8

+            mse_loss = (((predicted - target_emb) ** 2) * valid).sum() / n_valid

+            # Cosine component — preserves direction (speaker identity)

+            pred_flat = (predicted * valid).reshape(-1, predicted.shape[-1])

+            tgt_flat = (target_emb * valid).reshape(-1, target_emb.shape[-1])

+            # Only compute on non-padding positions

+            valid_flat = valid.squeeze(-1).reshape(-1) > 0

+                cosine_loss = 1.0 - F.cosine_similarity(

+                    pred_flat[valid_flat], tgt_flat[valid_flat], dim=-1

+                cosine_loss = torch.tensor(0.0, device=mel.device)

-            prediction_loss = F.mse_loss(predicted, target_emb)

+            mse_loss = F.mse_loss(predicted, target_emb)

+            cosine_loss = 1.0 - F.cosine_similarity(

+                predicted.reshape(-1, predicted.shape[-1]),

+                target_emb.reshape(-1, target_emb.shape[-1]),

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

+        prediction_loss = mse_loss + self.cosine_weight * cosine_loss

+        kl_per_dim = -0.5 * (1 + logvar - mu.pow(2) - logvar.exp())

kl_free = torch.clamp(kl_per_dim - self.free_bits, min=0.0)

-            ds_mask = mel_mask[:, ::self.audio_encoder.downsample_factor][:, :T_down]

-            valid_kl = (~ds_mask).unsqueeze(-1).float()

+            ds_mask_kl = mel_mask[:, ::self.audio_encoder.downsample_factor][:, :T_down_kl]

+            valid_kl = (~ds_mask_kl).unsqueeze(-1).float()

-        # ─── Loss 3: Mel reconstruction ───

-        mel_recon = self.mel_decoder(z)  # [B, n_mels, T_up]

+        # ─── Loss 3: Mel reconstruction (speaker-conditioned decoder) ───

+        mel_recon = self.mel_decoder(z, spk_cond)  # [B, n_mels, T_up]

-            valid_mel = (~mel_mask).unsqueeze(1).float()  # [B, 1, T]

+            valid_mel = (~mel_mask).unsqueeze(1).float()

mel_recon * valid_mel, mel * valid_mel, reduction="sum"

) / (valid_mel.sum() * mel.shape[1] + 1e-8)

# ─── Loss 4: Multi-resolution spectral loss ───

-            spectral_loss = self.spectral_loss(mel_recon * valid_mel, mel * valid_mel)

+            spectral_loss = self.spectral_loss_fn(mel_recon * valid_mel, mel * valid_mel)

-            spectral_loss = self.spectral_loss(mel_recon, mel)

+            spectral_loss = self.spectral_loss_fn(mel_recon, mel)

-        # ─── Loss 5: Speaker consistency — force audio encoder to retain speaker identity ───

-        # Uses the raw mu (before speaker embed injection) so the encoder itself

-        # must learn to preserve speaker info, not just parrot the added embedding

+        # ─── Loss 5: Speaker consistency on raw encoder output ───

speaker_loss = torch.tensor(0.0, device=mel.device)

if self.speaker_classifier is not None and speaker_id is not None:

-            # Pool encoder output across time (exclude padding)

-                valid_spk = (~ds_mask_spk).unsqueeze(-1).float()  # [B, T_down, 1]

+                valid_spk = (~ds_mask_spk).unsqueeze(-1).float()

-                mu_pooled = mu_raw.mean(dim=1)  # [B, d]

-            spk_logits = self.speaker_classifier(mu_pooled)  # [B, n_speakers]

+                mu_pooled = mu_raw.mean(dim=1)

+            spk_logits = self.speaker_classifier(mu_pooled)

speaker_loss = F.cross_entropy(spk_logits, speaker_id)

total_loss = (self.pred_weight * prediction_loss

+            "cosine_loss": cosine_loss,

"dropout": 0.1,
"kl_weight": 0.05,
"pred_weight": 10.0,
+            "cosine_weight": 1.0,
"ema_decay": 0.998,
-            "free_bits": 1.0,
+            "free_bits": 2.0,
"input_noise": 0.0,
}
model = LeWMTTS(config)

"""
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)
            # Speaker classifier on audio encoder output — forces encoder to retain speaker info
            self.speaker_classifier = nn.Linear(d_model, n_speakers)
        else:
            self.speaker_embed = None
            self.speaker_classifier = 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)

        # Speaker consistency loss weight
        self.speaker_weight = config.get("speaker_weight", 1.0)

        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)
        # Force eval mode so dropout is disabled — EMA targets must be deterministic
        with torch.no_grad():
            self.ema_audio_encoder.eval()
            _, mu_ema, _ = self.ema_audio_encoder(mel, mel_mask)
            self.ema_audio_encoder.train()
            target_emb = mu_ema  # [B, T_down, d] — all EMA targets

        # Save raw encoder mu before speaker injection — used for speaker classifier
        # so the encoder itself must learn speaker info, not just parrot the added embedding
        mu_raw = mu

        # Inject speaker into targets and encoder output so prediction loss
        # requires speaker discrimination — two speakers saying the same word
        # produce different targets, forcing the predictor to use speaker info
        if self.speaker_embed is not None and speaker_id is not None:
            spk = self.speaker_embed(speaker_id).unsqueeze(1)  # [B, 1, d]
            target_emb = target_emb + spk
            z = z + spk

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

        # ─── Loss 5: Speaker consistency — force audio encoder to retain speaker identity ───
        # Uses the raw mu (before speaker embed injection) so the encoder itself
        # must learn to preserve speaker info, not just parrot the added embedding
        speaker_loss = torch.tensor(0.0, device=mel.device)
        if self.speaker_classifier is not None and speaker_id is not None:
            # Pool encoder output across time (exclude padding)
            if mel_mask is not None:
                ds_mask_spk = mel_mask[:, ::self.audio_encoder.downsample_factor][:, :T_down]
                valid_spk = (~ds_mask_spk).unsqueeze(-1).float()  # [B, T_down, 1]
                mu_pooled = (mu_raw * valid_spk).sum(dim=1) / (valid_spk.sum(dim=1) + 1e-8)
            else:
                mu_pooled = mu_raw.mean(dim=1)  # [B, d]
            spk_logits = self.speaker_classifier(mu_pooled)  # [B, n_speakers]
            speaker_loss = F.cross_entropy(spk_logits, speaker_id)

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

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


Now update the inference to use FiLM instead of additive:


1	"""
2	Inference pipeline for LeWM TTS.
3	Text → JEPA predictor → mel reconstruction → Vocos vocoder → waveform
4	"""
5
6	import torch
7	import torchaudio
8	import numpy as np
9	import argparse
10	from pathlib import Path
11
12	from model import LeWMTTS
13
14
15	class LeWMTTSInference:
16	    """Full inference pipeline."""
17
18	    def __init__(self, checkpoint_path, device="cuda"):
19	        self.device = torch.device(device if torch.cuda.is_available() else "cpu")
20
21	        # Load checkpoint
22	        ckpt = torch.load(checkpoint_path, map_location=self.device, weights_only=False)
23	        config = ckpt["config"]
24
25	        # Build model
26	        self.model = LeWMTTS(config).to(self.device)
27	        self.model.load_state_dict(ckpt["model"])
28	        self.model.eval()
29
30	        # Vocos vocoder (directly decodes 100-mel)
31	        from vocos import Vocos
32	        self.vocos = Vocos.from_pretrained("charactr/vocos-mel-24khz")
33	        self.vocos.eval()
34
35	        self.config = config
36	        print(f"Model loaded from {checkpoint_path}")
37
38	    def vocos_decode(self, mel):
39	        """Decode mel [B, 100, T] → waveform via Vocos."""
40	        with torch.no_grad():
41	            waveform = self.vocos.decode(mel)  # [B, T_audio]
42	        return waveform
43
44	    def text_to_tokens(self, text):
45	        tokens = list(text.encode("utf-8"))
46	        return torch.tensor(tokens, dtype=torch.long).unsqueeze(0).to(self.device)
47
48	    @torch.no_grad()
49	    def synthesize(self, text, max_steps=300, temperature=0.0, speaker_id=None):
50	        """Generate speech from text using AR prediction with KV-cache."""
51	        text_tokens = self.text_to_tokens(text)
52	        text_emb = self.model.text_encoder(text_tokens)
53
54	        # Add speaker conditioning to text embeddings
55	        if self.model.speaker_embed is not None and speaker_id is not None:
56	            spk_id = torch.tensor([speaker_id], device=self.device)
57	            text_emb = self.model._add_speaker(text_emb, spk_id)
58	            # Cache speaker embedding for adding to predicted latents
59	            self._spk_emb = self.model.speaker_embed(spk_id).unsqueeze(1)  # [1, 1, d]
60	        else:
61	            self._spk_emb = None
62
63	        # Use learned start embedding instead of random noise
64	        start_emb = self.model.start_emb.to(self.device)
65	        all_embeddings = [start_emb]
66
67	        cache = self.model.init_ar_cache()
68	        next_emb, cache = self.model.predict_next_cached(start_emb, text_emb, 0, cache)
69	        if temperature > 0:
70	            next_emb = next_emb + torch.randn_like(next_emb) * temperature
71	        # Add speaker embedding to predicted latent (matches training where z = z + spk)
72	        emb_for_decode = next_emb + self._spk_emb if self._spk_emb is not None else next_emb
73	        all_embeddings.append(emb_for_decode)
74
75	        prev_norms = [next_emb.norm().item()]
76
77	        for step in range(1, max_steps):
78	            # Feed speaker-conditioned embedding as input (matches training input_emb)
79	            next_input = emb_for_decode
80	            next_emb, cache = self.model.predict_next_cached(next_input, text_emb, step, cache)
81	            if temperature > 0:
82	                next_emb = next_emb + torch.randn_like(next_emb) * temperature
83	            emb_for_decode = next_emb + self._spk_emb if self._spk_emb is not None else next_emb
84	            all_embeddings.append(emb_for_decode)
85
86	            norm = next_emb.norm().item()
87	            prev_norms.append(norm)
88	            if len(prev_norms) > 20:
89	                recent = prev_norms[-20:]
90	                if np.std(recent) < 0.01 * np.mean(recent) and step > 30:
91	                    break
92
93	        audio_embs = torch.cat(all_embeddings, dim=1)
94	        mel = self.model.mel_decoder(audio_embs)
95	        waveform = self.vocos_decode(mel.cpu())
96	        return waveform.squeeze().numpy(), 24000
97
98	    @torch.no_grad()
99	    def reconstruct(self, mel_path):
100	        """Teacher-forced reconstruction: encode real mel → decode back."""
101	        mel = torch.load(mel_path, weights_only=True).unsqueeze(0).to(self.device)
102
103	        z, mu, logvar = self.model.audio_encoder(mel)
104	        mel_recon = self.model.mel_decoder(mu)
105
106	        waveform = self.vocos_decode(mel_recon.cpu())
107	        return waveform.squeeze().numpy(), 24000
108
109	    @torch.no_grad()
110	    def reconstruct_with_prediction(self, mel_path):
111	        """Hybrid: encode real mel → predictor → decode predicted embeddings."""
112	        mel = torch.load(mel_path, weights_only=True).unsqueeze(0).to(self.device)
113
114	        z, mu, logvar = self.model.audio_encoder(mel)
115
116	        text_tokens = torch.zeros(1, 1, dtype=torch.long, device=self.device)
117	        text_emb = self.model.text_encoder(text_tokens)
118
119	        predicted = self.model.predictor(mu[:, :-1], text_emb)
120	        combined = torch.cat([mu[:, :1], predicted], dim=1)
121
122	        mel_recon = self.model.mel_decoder(combined)
123	        waveform = self.vocos_decode(mel_recon.cpu())
124	        return waveform.squeeze().numpy(), 24000
125
126	    def save_audio(self, waveform, sr, output_path):
127	        # Normalize
128	        mx = np.abs(waveform).max()
129	        if mx > 0:
130	            waveform = waveform / mx * 0.95
131	        wav_tensor = torch.from_numpy(waveform).unsqueeze(0)
132	        torchaudio.save(output_path, wav_tensor, sr)
133	        print(f"Saved: {output_path} ({len(waveform)/sr:.2f}s)")
134
135
136	def main():
137	    parser = argparse.ArgumentParser()
138	    parser.add_argument("--checkpoint", required=True)
139	    parser.add_argument("--text", default=None, help="Text to synthesize")
140	    parser.add_argument("--mel", default=None, help="Mel path for reconstruction test")
141	    parser.add_argument("--mode", choices=["synthesize", "reconstruct", "predict"],
142	                        default="synthesize")
143	    parser.add_argument("--output", default="output.wav")
144	    parser.add_argument("--max_steps", type=int, default=300)
145	    parser.add_argument("--temperature", type=float, default=0.0)
146	    parser.add_argument("--speaker_id", type=int, default=None)
147	    args = parser.parse_args()
148
149	    tts = LeWMTTSInference(args.checkpoint)
150
151	    if args.mode == "synthesize":
152	        assert args.text, "Need --text for synthesis"
153	        waveform, sr = tts.synthesize(args.text, args.max_steps, args.temperature,
154	                                       speaker_id=args.speaker_id)
155	    elif args.mode == "reconstruct":
156	        assert args.mel, "Need --mel for reconstruction"
157	        waveform, sr = tts.reconstruct(args.mel)
158	    elif args.mode == "predict":
159	        assert args.mel, "Need --mel for prediction test"
160	        waveform, sr = tts.reconstruct_with_prediction(args.mel)
161
162	    tts.save_audio(waveform, sr, args.output)
163
164
165	if __name__ == "__main__":
166	    main()
167

"""
Inference pipeline for LeWM TTS.
Text → JEPA predictor → mel reconstruction → Vocos vocoder → waveform
"""

import torch
import torchaudio
import numpy as np
import argparse
from pathlib import Path

from model import LeWMTTS

class LeWMTTSInference:
    """Full inference pipeline."""

    def __init__(self, checkpoint_path, device="cuda"):
        self.device = torch.device(device if torch.cuda.is_available() else "cpu")

        # Load checkpoint
        ckpt = torch.load(checkpoint_path, map_location=self.device, weights_only=False)
        config = ckpt["config"]

        # Build model
        self.model = LeWMTTS(config).to(self.device)
        self.model.load_state_dict(ckpt["model"])
        self.model.eval()

        # Vocos vocoder (directly decodes 100-mel)
        from vocos import Vocos
        self.vocos = Vocos.from_pretrained("charactr/vocos-mel-24khz")
        self.vocos.eval()

        self.config = config
        print(f"Model loaded from {checkpoint_path}")

    def vocos_decode(self, mel):
        """Decode mel [B, 100, T] → waveform via Vocos."""
        with torch.no_grad():
            waveform = self.vocos.decode(mel)  # [B, T_audio]
        return waveform

    def text_to_tokens(self, text):
        tokens = list(text.encode("utf-8"))
        return torch.tensor(tokens, dtype=torch.long).unsqueeze(0).to(self.device)

    @torch.no_grad()
    def synthesize(self, text, max_steps=300, temperature=0.0, speaker_id=None):
        """Generate speech from text using AR prediction with KV-cache."""
        text_tokens = self.text_to_tokens(text)
        text_emb = self.model.text_encoder(text_tokens)

        # Add speaker conditioning to text embeddings
        if self.model.speaker_embed is not None and speaker_id is not None:
            spk_id = torch.tensor([speaker_id], device=self.device)
            text_emb = self.model._add_speaker(text_emb, spk_id)
            # Cache speaker embedding for adding to predicted latents
            self._spk_emb = self.model.speaker_embed(spk_id).unsqueeze(1)  # [1, 1, d]
        else:
            self._spk_emb = None

        # Use learned start embedding instead of random noise
        start_emb = self.model.start_emb.to(self.device)
        all_embeddings = [start_emb]

        cache = self.model.init_ar_cache()
        next_emb, cache = self.model.predict_next_cached(start_emb, text_emb, 0, cache)
        if temperature > 0:
            next_emb = next_emb + torch.randn_like(next_emb) * temperature
        # Add speaker embedding to predicted latent (matches training where z = z + spk)
        emb_for_decode = next_emb + self._spk_emb if self._spk_emb is not None else next_emb
        all_embeddings.append(emb_for_decode)

        prev_norms = [next_emb.norm().item()]

        for step in range(1, max_steps):
            # Feed speaker-conditioned embedding as input (matches training input_emb)
            next_input = emb_for_decode
            next_emb, cache = self.model.predict_next_cached(next_input, text_emb, step, cache)
            if temperature > 0:
                next_emb = next_emb + torch.randn_like(next_emb) * temperature
            emb_for_decode = next_emb + self._spk_emb if self._spk_emb is not None else next_emb
            all_embeddings.append(emb_for_decode)

            norm = next_emb.norm().item()
            prev_norms.append(norm)
            if len(prev_norms) > 20:
                recent = prev_norms[-20:]
                if np.std(recent) < 0.01 * np.mean(recent) and step > 30:
                    break

        audio_embs = torch.cat(all_embeddings, dim=1)
        mel = self.model.mel_decoder(audio_embs)
        waveform = self.vocos_decode(mel.cpu())
        return waveform.squeeze().numpy(), 24000

    @torch.no_grad()
    def reconstruct(self, mel_path):
        """Teacher-forced reconstruction: encode real mel → decode back."""
        mel = torch.load(mel_path, weights_only=True).unsqueeze(0).to(self.device)

        z, mu, logvar = self.model.audio_encoder(mel)
        mel_recon = self.model.mel_decoder(mu)

        waveform = self.vocos_decode(mel_recon.cpu())
        return waveform.squeeze().numpy(), 24000

    @torch.no_grad()
    def reconstruct_with_prediction(self, mel_path):
        """Hybrid: encode real mel → predictor → decode predicted embeddings."""
        mel = torch.load(mel_path, weights_only=True).unsqueeze(0).to(self.device)

        z, mu, logvar = self.model.audio_encoder(mel)

        text_tokens = torch.zeros(1, 1, dtype=torch.long, device=self.device)
        text_emb = self.model.text_encoder(text_tokens)

        predicted = self.model.predictor(mu[:, :-1], text_emb)
        combined = torch.cat([mu[:, :1], predicted], dim=1)

        mel_recon = self.model.mel_decoder(combined)
        waveform = self.vocos_decode(mel_recon.cpu())
        return waveform.squeeze().numpy(), 24000

    def save_audio(self, waveform, sr, output_path):
        # Normalize
        mx = np.abs(waveform).max()
        if mx > 0:
            waveform = waveform / mx * 0.95
        wav_tensor = torch.from_numpy(waveform).unsqueeze(0)
        torchaudio.save(output_path, wav_tensor, sr)
        print(f"Saved: {output_path} ({len(waveform)/sr:.2f}s)")

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--checkpoint", required=True)
    parser.add_argument("--text", default=None, help="Text to synthesize")
    parser.add_argument("--mel", default=None, help="Mel path for reconstruction test")
    parser.add_argument("--mode", choices=["synthesize", "reconstruct", "predict"],
                        default="synthesize")
    parser.add_argument("--output", default="output.wav")
    parser.add_argument("--max_steps", type=int, default=300)
    parser.add_argument("--temperature", type=float, default=0.0)
    parser.add_argument("--speaker_id", type=int, default=None)
    args = parser.parse_args()

    tts = LeWMTTSInference(args.checkpoint)

    if args.mode == "synthesize":
        assert args.text, "Need --text for synthesis"
        waveform, sr = tts.synthesize(args.text, args.max_steps, args.temperature,
                                       speaker_id=args.speaker_id)
    elif args.mode == "reconstruct":
        assert args.mel, "Need --mel for reconstruction"
        waveform, sr = tts.reconstruct(args.mel)
    elif args.mode == "predict":
        assert args.mel, "Need --mel for prediction test"
        waveform, sr = tts.reconstruct_with_prediction(args.mel)

    tts.save_audio(waveform, sr, args.output)

if __name__ == "__main__":
    main()