class LeWMTTSv5(nn.Module):
    """
    JEPA TTS with EnCodec backend + duration predictor.
    Text is expanded to frame level and added to audio input — can't be ignored.

    Training: codec_emb → proj_in → predict(audio + aligned_text) → loss vs targets
              Duration predictor trained with uniform alignment as supervision.
    Inference: text → duration predict → expand text → AR predict → proj_out → decode
    """

    def __init__(self, config):
        super().__init__()
        d_model = config.get("d_model", 256)
        nhead = config.get("nhead", 4)
        codec_dim = config.get("codec_dim", 128)
        text_vocab = config.get("text_vocab_size", 256)
        text_layers = config.get("text_encoder_layers", 4)
        predictor_layers = config.get("predictor_layers", 6)
        dropout = config.get("dropout", 0.1)

        self.codec_dim = codec_dim

        self.text_encoder = TextEncoder(
            vocab_size=text_vocab, d_model=d_model, nhead=nhead,
            num_layers=text_layers, dropout=dropout,
        )

        # Duration predictor
        self.duration_predictor = DurationPredictor(
            d_model=d_model, kernel_size=3, num_layers=2, dropout=dropout,
        )

        # Simple linear projections — no lossy compression
        self.proj_in = nn.Sequential(
            nn.Linear(codec_dim, d_model),
            nn.GELU(),
            nn.Linear(d_model, d_model),
        )
        self.proj_out = nn.Sequential(
            nn.Linear(d_model, d_model),
            nn.GELU(),
            nn.Linear(d_model, codec_dim),
        )

        self.predictor = JEPAPredictor(
            d_model=d_model, nhead=nhead,
            num_layers=predictor_layers, dropout=dropout,
        )

        # EMA on proj_in (provides stable prediction targets)
        self.ema_proj_in = copy.deepcopy(self.proj_in)
        for p in self.ema_proj_in.parameters():
            p.requires_grad = False
        self.ema_decay = config.get("ema_decay", 0.998)

        # Learnable start embedding
        self.start_emb = nn.Parameter(torch.randn(1, 1, d_model) * 0.02)

        # Loss weights
        self.pred_weight = config.get("pred_weight", 10.0)
        self.roundtrip_weight = config.get("roundtrip_weight", 1.0)
        self.dur_weight = config.get("dur_weight", 1.0)

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

        self.config = config

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

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

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

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

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

    def forward(self, text_emb, text_mask=None):
        """
        text_emb: [B, T_text, d_model]
        Returns: log_durations [B, T_text]
        """
        x = text_emb.transpose(1, 2)  # [B, d_model, T]
        for i in range(0, len(self.convs), 4):
            x = self.convs[i](x)      # Conv1d
            x = self.convs[i+1](x)    # ReLU
            x = x.transpose(1, 2)     # [B, T, d_model] for LayerNorm
            x = self.convs[i+2](x)    # LayerNorm
            x = self.convs[i+3](x)    # Dropout
            x = x.transpose(1, 2)     # back to [B, d_model, T]
        x = x.transpose(1, 2)  # [B, T, d_model]
        log_dur = self.proj(x).squeeze(-1)  # [B, T_text]
        if text_mask is not None:
            log_dur = log_dur.masked_fill(text_mask, 0.0)
        return log_dur

def length_regulate(text_emb, durations, text_mask=None):
    """
    Expand text embeddings by repeating each token according to its duration.
    text_emb: [B, T_text, d_model]
    durations: [B, T_text] integer durations
    Returns: expanded [B, T_audio, d_model], where T_audio = sum(durations) per sample
    """
    B, T_text, D = text_emb.shape
    max_audio_len = durations.sum(dim=1).max().item()
    expanded = torch.zeros(B, max_audio_len, D, device=text_emb.device)
    for b in range(B):
        pos = 0
        for t in range(T_text):
            if text_mask is not None and text_mask[b, t]:
                continue
            dur = durations[b, t].item()
            if dur > 0 and pos < max_audio_len:
                end = min(pos + dur, max_audio_len)
                expanded[b, pos:end] = text_emb[b, t]
                pos = end
    return expanded

def compute_uniform_durations(text_lengths, audio_lengths):
    """
    Compute uniform durations: spread audio frames evenly across text characters.
    text_lengths: [B] number of valid text tokens
    audio_lengths: [B] number of valid audio frames
    Returns: durations [B, max_text_len] integer
    """
    B = text_lengths.shape[0]
    max_text_len = text_lengths.max().item()
    durations = torch.zeros(B, max_text_len, dtype=torch.long, device=text_lengths.device)
    for b in range(B):
        tl = text_lengths[b].item()
        al = audio_lengths[b].item()
        if tl == 0:
            continue
        base = al // tl
        remainder = al % tl
        durations[b, :tl] = base
        durations[b, :remainder] += 1
    return durations

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

class JEPAPredictor(nn.Module):
    """Causal predictor with frame-level text added to audio input."""
    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.text_input_proj = nn.Linear(d_model, d_model)
        self.pos_embed = SinusoidalPositionalEncoding(d_model, max_len=4096)
        self.dropout = nn.Dropout(dropout)

        # Use encoder (self-attention only) since text is now part of input
        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.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_expanded, audio_mask=None):
        """
        audio_emb: [B, T, d_model] — shifted audio embeddings
        text_emb_expanded: [B, T, d_model] — frame-level text (already expanded)
        """
        x = self.audio_input_proj(audio_emb) + self.text_input_proj(text_emb_expanded)
        x = self.pos_embed(x)
        x = self.dropout(x)
        T = x.shape[1]
        causal_mask = torch.triu(torch.ones(T, T, device=x.device, dtype=torch.bool), diagonal=1)
        predicted = self.transformer(x, mask=causal_mask, src_key_padding_mask=audio_mask)
        return self.output_proj(predicted)

    # Keep old cross-attention interface for backward compat loading
    def forward_cross_attn(self, audio_emb, text_emb, audio_mask=None, text_mask=None):
        """Legacy cross-attention forward (not used in duration-predictor mode)."""
        x = self.audio_input_proj(audio_emb)
        x = self.pos_embed(x)
        x = self.dropout(x)
        return self.output_proj(x)

    def _cached_mha(self, mha, q_input, kv_input, cache_k, cache_v):
        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)
        k_new = F.linear(kv_input, Wk, bk)
        v_new = F.linear(kv_input, Wv, bv)
        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, v = k_new, v_new
        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)
        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):
        return {
            'self_k': [None] * num_layers,
            'self_v': [None] * num_layers,
        }

    def inference_step(self, new_emb, text_emb_frame, step_idx, cache):
        """
        new_emb: [B, 1, d_model] — audio embedding for this step
        text_emb_frame: [B, 1, d_model] — expanded text embedding for this step
        """
        x = self.audio_input_proj(new_emb) + self.text_input_proj(text_emb_frame)
        x = x + self.pos_embed.pe[:, step_idx:step_idx + 1]
        for i, layer in enumerate(self.transformer.layers):
            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))
            ff_out = layer.linear2(layer.dropout(layer.activation(layer.linear1(x))))
            x = layer.norm2(x + layer.dropout2(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))

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

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

class TextEncoder(nn.Module):
    def __init__(self, vocab_size=256, d_model=256, nhead=4, num_layers=4, dropout=0.1):
        super().__init__()
        self.d_model = d_model
        self.char_embed = nn.Embedding(vocab_size, d_model)
        self.pos_embed = SinusoidalPositionalEncoding(d_model, max_len=2048)
        self.dropout = nn.Dropout(dropout)
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=d_model, nhead=nhead, dim_feedforward=d_model * 4,
            dropout=dropout, batch_first=True, activation="gelu",
        )
        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
        self.proj = nn.Linear(d_model, d_model)

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

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

class LeWMTTSv5(nn.Module):
    """
    JEPA TTS with EnCodec backend.
    No down/upsampler — works directly at EnCodec's native 75Hz.
    Simple linear projections: 128d (codec) ↔ d_model (predictor).

    Training: codec_emb → proj_in → JEPA predict → loss vs EMA targets
              codec_emb → proj_in → proj_out → loss vs codec_emb (roundtrip)
    Inference: text → predictor (AR) → proj_out → EnCodec decoder → waveform
    """

    def __init__(self, config):
        super().__init__()
        d_model = config.get("d_model", 256)
        nhead = config.get("nhead", 4)
        codec_dim = config.get("codec_dim", 128)
        text_vocab = config.get("text_vocab_size", 256)
        text_layers = config.get("text_encoder_layers", 4)
        predictor_layers = config.get("predictor_layers", 6)
        dropout = config.get("dropout", 0.1)

        self.codec_dim = codec_dim

        self.text_encoder = TextEncoder(
            vocab_size=text_vocab, d_model=d_model, nhead=nhead,
            num_layers=text_layers, dropout=dropout,
        )

        # Simple linear projections — no lossy compression
        self.proj_in = nn.Sequential(
            nn.Linear(codec_dim, d_model),
            nn.GELU(),
            nn.Linear(d_model, d_model),
        )
        self.proj_out = nn.Sequential(
            nn.Linear(d_model, d_model),
            nn.GELU(),
            nn.Linear(d_model, codec_dim),
        )

        self.predictor = JEPAPredictor(
            d_model=d_model, nhead=nhead,
            num_layers=predictor_layers, dropout=dropout,
        )

        # EMA on proj_in (provides stable prediction targets)
        self.ema_proj_in = copy.deepcopy(self.proj_in)
        for p in self.ema_proj_in.parameters():
            p.requires_grad = False
        self.ema_decay = config.get("ema_decay", 0.998)

        # Learnable start embedding
        self.start_emb = nn.Parameter(torch.randn(1, 1, d_model) * 0.02)

        # Loss weights
        self.pred_weight = config.get("pred_weight", 10.0)
        self.roundtrip_weight = config.get("roundtrip_weight", 1.0)

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

        self.config = config

    @torch.no_grad()
    def update_ema(self):
        for p_online, p_ema in zip(self.proj_in.parameters(),
                                   self.ema_proj_in.parameters()):
            p_ema.data.mul_(self.ema_decay).add_(p_online.data, alpha=1 - self.ema_decay)

    def forward(self, codec_emb, text_tokens, codec_mask=None, text_mask=None):
        """
        Args:
            codec_emb: [B, 128, T] — continuous EnCodec embeddings
            text_tokens: [B, T_text]
            codec_mask: [B, T] bool, True = padding
            text_mask: [B, T_text] bool
        """
        B = codec_emb.shape[0]
        T = codec_emb.shape[2]

        # codec_emb is [B, 128, T] → transpose to [B, T, 128]
        codec_seq = codec_emb.transpose(1, 2)

        # Text encoding
        text_emb = self.text_encoder(text_tokens, text_mask)

        # Project codec embeddings to predictor space
        z = self.proj_in(codec_seq)  # [B, T, d_model]

        # EMA targets (deterministic)
        with torch.no_grad():
            self.ema_proj_in.eval()
            target_emb = self.ema_proj_in(codec_seq)  # [B, T, d_model]
            self.ema_proj_in.train()

        # Prepend start embedding, shift right
        start = self.start_emb.expand(B, -1, -1)
        input_emb = torch.cat([start, z[:, :-1]], dim=1)  # [B, T, d_model]

        # Input noise
        if self.training and self.input_noise > 0:
            noise = torch.zeros_like(input_emb)
            noise[:, 1:] = torch.randn(B, T - 1, input_emb.shape[-1],
                                        device=input_emb.device) * self.input_noise
            input_emb = input_emb + noise

        # Masks
        if codec_mask is not None:
            start_mask = torch.zeros(B, 1, dtype=torch.bool, device=codec_emb.device)
            pred_mask = torch.cat([start_mask, codec_mask[:, :-1]], dim=1)
            loss_mask = codec_mask
        else:
            pred_mask = None
            loss_mask = None

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

        # ─── Loss 1: Prediction (MSE vs EMA targets) ───
        if loss_mask is not None:
            valid = (~loss_mask).unsqueeze(-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: Roundtrip (proj_in → proj_out → match original codec_emb) ───
        codec_recon = self.proj_out(z)  # [B, T, 128]

        if codec_mask is not None:
            valid_rt = (~codec_mask).unsqueeze(-1).float()
            roundtrip_loss = F.l1_loss(
                codec_recon * valid_rt, codec_seq * valid_rt, reduction="sum"
            ) / (valid_rt.sum() * self.codec_dim + 1e-8)
        else:
            roundtrip_loss = F.l1_loss(codec_recon, codec_seq)

        total_loss = self.pred_weight * prediction_loss + self.roundtrip_weight * roundtrip_loss

        return {
            "total_loss": total_loss,
            "prediction_loss": prediction_loss,
            "roundtrip_loss": roundtrip_loss,
        }

    def init_ar_cache(self):
        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):
        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, "codec_dim": 128,
            "text_vocab_size": 256, "text_encoder_layers": 4,
            "predictor_layers": 6, "dropout": 0.1,
            "pred_weight": 10.0, "roundtrip_weight": 1.0,
            "ema_decay": 0.998, "input_noise": 0.0,
        }
    model = LeWMTTSv5(config)
    print(f"LeWM TTS v5: {count_parameters(model)/1e6:.2f}M trainable parameters")
    return model, config

-    No down/upsampler — works directly at EnCodec's native 75Hz.

-    Simple linear projections: 128d (codec) ↔ d_model (predictor).

+    JEPA TTS with EnCodec backend + duration predictor.

+    Text is expanded to frame level and added to audio input — can't be ignored.

-    Training: codec_emb → proj_in → JEPA predict → loss vs EMA targets

-              codec_emb → proj_in → proj_out → loss vs codec_emb (roundtrip)

-    Inference: text → predictor (AR) → proj_out → EnCodec decoder → waveform

+    Training: codec_emb → proj_in → predict(audio + aligned_text) → loss vs targets

+              Duration predictor trained with uniform alignment as supervision.

+    Inference: text → duration predict → expand text → AR predict → proj_out → decode

+        self.duration_predictor = DurationPredictor(

+            d_model=d_model, kernel_size=3, num_layers=2, dropout=dropout,

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


Now update the `forward` method to use duration predictor:


Edit

    @torch.no_grad()
    def update_ema(self):
        for p_online, p_ema in zip(self.proj_in.parameters(),
                                   self.ema_proj_in.parameters()):
            p_ema.data.mul_(self.ema_decay).add_(p_online.data, alpha=1 - self.ema_decay)

    def forward(self, codec_emb, text_tokens, codec_mask=None, text_mask=None):
        """
        Args:
            codec_emb: [B, 128, T] — continuous EnCodec embeddings
            text_tokens: [B, T_text]
            codec_mask: [B, T] bool, True = padding
            text_mask: [B, T_text] bool
        """
        B = codec_emb.shape[0]
        T = codec_emb.shape[2]

        # codec_emb is [B, 128, T] → transpose to [B, T, 128]
        codec_seq = codec_emb.transpose(1, 2)

        # Text encoding
        text_emb = self.text_encoder(text_tokens, text_mask)

        # Project codec embeddings to predictor space
        z = self.proj_in(codec_seq)  # [B, T, d_model]

        # EMA targets (deterministic)
        with torch.no_grad():
            self.ema_proj_in.eval()
            target_emb = self.ema_proj_in(codec_seq)  # [B, T, d_model]
            self.ema_proj_in.train()

        # Prepend start embedding, shift right
        start = self.start_emb.expand(B, -1, -1)
        input_emb = torch.cat([start, z[:, :-1]], dim=1)  # [B, T, d_model]

        # Input noise
        if self.training and self.input_noise > 0:
            noise = torch.zeros_like(input_emb)
            noise[:, 1:] = torch.randn(B, T - 1, input_emb.shape[-1],
                                        device=input_emb.device) * self.input_noise
            input_emb = input_emb + noise

        # Masks
        if codec_mask is not None:
            start_mask = torch.zeros(B, 1, dtype=torch.bool, device=codec_emb.device)
            pred_mask = torch.cat([start_mask, codec_mask[:, :-1]], dim=1)
            loss_mask = codec_mask
        else:
            pred_mask = None
            loss_mask = None

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

        # ─── Loss 1: Prediction (MSE vs EMA targets) ───
        if loss_mask is not None:
            valid = (~loss_mask).unsqueeze(-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: Roundtrip (proj_in → proj_out → match original codec_emb) ───
        codec_recon = self.proj_out(z)  # [B, T, 128]

        if codec_mask is not None:
            valid_rt = (~codec_mask).unsqueeze(-1).float()
            roundtrip_loss = F.l1_loss(
                codec_recon * valid_rt, codec_seq * valid_rt, reduction="sum"
            ) / (valid_rt.sum() * self.codec_dim + 1e-8)
        else:
            roundtrip_loss = F.l1_loss(codec_recon, codec_seq)

        total_loss = self.pred_weight * prediction_loss + self.roundtrip_weight * roundtrip_loss

        return {
            "total_loss": total_loss,
            "prediction_loss": prediction_loss,
            "roundtrip_loss": roundtrip_loss,
        }

    @torch.no_grad()
    def update_ema(self):
        for p_online, p_ema in zip(self.proj_in.parameters(),
                                   self.ema_proj_in.parameters()):
            p_ema.data.mul_(self.ema_decay).add_(p_online.data, alpha=1 - self.ema_decay)

    def forward(self, codec_emb, text_tokens, codec_mask=None, text_mask=None):
        """
        Args:
            codec_emb: [B, 128, T] — continuous EnCodec embeddings
            text_tokens: [B, T_text]
            codec_mask: [B, T] bool, True = padding
            text_mask: [B, T_text] bool
        """
        B = codec_emb.shape[0]
        T = codec_emb.shape[2]

        # codec_emb is [B, 128, T] → transpose to [B, T, 128]
        codec_seq = codec_emb.transpose(1, 2)

        # Text encoding
        text_emb = self.text_encoder(text_tokens, text_mask)

        # Duration prediction + ground truth durations (uniform alignment)
        log_dur_pred = self.duration_predictor(text_emb, text_mask)

        # Compute actual text/audio lengths from masks
        if text_mask is not None:
            text_lengths = (~text_mask).sum(dim=1)  # [B]
        else:
            text_lengths = torch.full((B,), text_tokens.shape[1], device=text_tokens.device)
        if codec_mask is not None:
            audio_lengths = (~codec_mask).sum(dim=1)  # [B]
        else:
            audio_lengths = torch.full((B,), T, device=codec_emb.device)

        # Ground truth durations (uniform distribution)
        gt_durations = compute_uniform_durations(text_lengths, audio_lengths)
        log_dur_gt = torch.log(gt_durations.float().clamp(min=1))

        # Duration loss (MSE on log-durations)
        if text_mask is not None:
            valid_text = (~text_mask).float()
            dur_loss = F.mse_loss(
                log_dur_pred * valid_text, log_dur_gt * valid_text, reduction="sum"
            ) / (valid_text.sum() + 1e-8)
        else:
            dur_loss = F.mse_loss(log_dur_pred, log_dur_gt)

        # Expand text to frame level using GT durations (training uses GT)
        text_emb_expanded = length_regulate(text_emb, gt_durations, text_mask)
        # Trim/pad to match audio length T
        if text_emb_expanded.shape[1] > T:
            text_emb_expanded = text_emb_expanded[:, :T]
        elif text_emb_expanded.shape[1] < T:
            pad = torch.zeros(B, T - text_emb_expanded.shape[1], text_emb.shape[-1],
                            device=text_emb.device)
            text_emb_expanded = torch.cat([text_emb_expanded, pad], dim=1)

        # Project codec embeddings to predictor space
        z = self.proj_in(codec_seq)  # [B, T, d_model]

        # EMA targets (deterministic)
        with torch.no_grad():
            self.ema_proj_in.eval()
            target_emb = self.ema_proj_in(codec_seq)  # [B, T, d_model]
            self.ema_proj_in.train()

        # Prepend start embedding, shift right
        start = self.start_emb.expand(B, -1, -1)
        input_emb = torch.cat([start, z[:, :-1]], dim=1)  # [B, T, d_model]

        # Input noise
        if self.training and self.input_noise > 0:
            noise = torch.zeros_like(input_emb)
            noise[:, 1:] = torch.randn(B, T - 1, input_emb.shape[-1],
                                        device=input_emb.device) * self.input_noise
            input_emb = input_emb + noise

        # Masks
        if codec_mask is not None:
            start_mask = torch.zeros(B, 1, dtype=torch.bool, device=codec_emb.device)
            pred_mask = torch.cat([start_mask, codec_mask[:, :-1]], dim=1)
            loss_mask = codec_mask
        else:
            pred_mask = None
            loss_mask = None

        # JEPA prediction with frame-level text
        predicted = self.predictor(input_emb, text_emb_expanded, pred_mask)

        # ─── Loss 1: Prediction (MSE vs EMA targets) ───
        if loss_mask is not None:
            valid = (~loss_mask).unsqueeze(-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: Roundtrip (proj_in → proj_out → match original codec_emb) ───
        codec_recon = self.proj_out(z)  # [B, T, 128]

        if codec_mask is not None:
            valid_rt = (~codec_mask).unsqueeze(-1).float()
            roundtrip_loss = F.l1_loss(
                codec_recon * valid_rt, codec_seq * valid_rt, reduction="sum"
            ) / (valid_rt.sum() * self.codec_dim + 1e-8)
        else:
            roundtrip_loss = F.l1_loss(codec_recon, codec_seq)

        total_loss = (self.pred_weight * prediction_loss
                     + self.roundtrip_weight * roundtrip_loss
                     + self.dur_weight * dur_loss)

        return {
            "total_loss": total_loss,
            "prediction_loss": prediction_loss,
            "roundtrip_loss": roundtrip_loss,
            "dur_loss": dur_loss,
        }


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

    @torch.no_grad()
    def update_ema(self):
        for p_online, p_ema in zip(self.proj_in.parameters(),
                                   self.ema_proj_in.parameters()):
            p_ema.data.mul_(self.ema_decay).add_(p_online.data, alpha=1 - self.ema_decay)

    def forward(self, codec_emb, text_tokens, codec_mask=None, text_mask=None):
        """
        Args:
            codec_emb: [B, 128, T] — continuous EnCodec embeddings
            text_tokens: [B, T_text]
            codec_mask: [B, T] bool, True = padding
            text_mask: [B, T_text] bool
        """
        B = codec_emb.shape[0]
        T = codec_emb.shape[2]

        # codec_emb is [B, 128, T] → transpose to [B, T, 128]
        codec_seq = codec_emb.transpose(1, 2)

        # Text encoding
        text_emb = self.text_encoder(text_tokens, text_mask)

        # Duration prediction + ground truth durations (uniform alignment)
        log_dur_pred = self.duration_predictor(text_emb, text_mask)

        # Compute actual text/audio lengths from masks
        if text_mask is not None:
            text_lengths = (~text_mask).sum(dim=1)  # [B]
        else:
            text_lengths = torch.full((B,), text_tokens.shape[1], device=text_tokens.device)
        if codec_mask is not None:
            audio_lengths = (~codec_mask).sum(dim=1)  # [B]
        else:
            audio_lengths = torch.full((B,), T, device=codec_emb.device)

        # Ground truth durations (uniform distribution)
        gt_durations = compute_uniform_durations(text_lengths, audio_lengths)
        log_dur_gt = torch.log(gt_durations.float().clamp(min=1))

        # Duration loss (MSE on log-durations)
        if text_mask is not None:
            valid_text = (~text_mask).float()
            dur_loss = F.mse_loss(
                log_dur_pred * valid_text, log_dur_gt * valid_text, reduction="sum"
            ) / (valid_text.sum() + 1e-8)
        else:
            dur_loss = F.mse_loss(log_dur_pred, log_dur_gt)

        # Expand text to frame level using GT durations (training uses GT)
        text_emb_expanded = length_regulate(text_emb, gt_durations, text_mask)
        # Trim/pad to match audio length T
        if text_emb_expanded.shape[1] > T:
            text_emb_expanded = text_emb_expanded[:, :T]
        elif text_emb_expanded.shape[1] < T:
            pad = torch.zeros(B, T - text_emb_expanded.shape[1], text_emb.shape[-1],
                            device=text_emb.device)
            text_emb_expanded = torch.cat([text_emb_expanded, pad], dim=1)

        # Project codec embeddings to predictor space
        z = self.proj_in(codec_seq)  # [B, T, d_model]

        # EMA targets (deterministic)
        with torch.no_grad():
            self.ema_proj_in.eval()
            target_emb = self.ema_proj_in(codec_seq)  # [B, T, d_model]
            self.ema_proj_in.train()

        # Prepend start embedding, shift right
        start = self.start_emb.expand(B, -1, -1)
        input_emb = torch.cat([start, z[:, :-1]], dim=1)  # [B, T, d_model]

        # Input noise
        if self.training and self.input_noise > 0:
            noise = torch.zeros_like(input_emb)
            noise[:, 1:] = torch.randn(B, T - 1, input_emb.shape[-1],
                                        device=input_emb.device) * self.input_noise
            input_emb = input_emb + noise

        # Masks
        if codec_mask is not None:
            start_mask = torch.zeros(B, 1, dtype=torch.bool, device=codec_emb.device)
            pred_mask = torch.cat([start_mask, codec_mask[:, :-1]], dim=1)
            loss_mask = codec_mask
        else:
            pred_mask = None
            loss_mask = None

        # JEPA prediction with frame-level text
        predicted = self.predictor(input_emb, text_emb_expanded, pred_mask)

        # ─── Loss 1: Prediction (MSE vs EMA targets) ───
        if loss_mask is not None:
            valid = (~loss_mask).unsqueeze(-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: Roundtrip (proj_in → proj_out → match original codec_emb) ───
        codec_recon = self.proj_out(z)  # [B, T, 128]

        if codec_mask is not None:
            valid_rt = (~codec_mask).unsqueeze(-1).float()
            roundtrip_loss = F.l1_loss(
                codec_recon * valid_rt, codec_seq * valid_rt, reduction="sum"
            ) / (valid_rt.sum() * self.codec_dim + 1e-8)
        else:
            roundtrip_loss = F.l1_loss(codec_recon, codec_seq)

        total_loss = (self.pred_weight * prediction_loss
                     + self.roundtrip_weight * roundtrip_loss
                     + self.dur_weight * dur_loss)

        return {
            "total_loss": total_loss,
            "prediction_loss": prediction_loss,
            "roundtrip_loss": roundtrip_loss,
            "dur_loss": dur_loss,
        }

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

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

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

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

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

    def forward(self, text_emb, text_mask=None):
        """
        text_emb: [B, T_text, d_model]
        Returns: log_durations [B, T_text]
        """
        x = text_emb.transpose(1, 2)  # [B, d_model, T]
        for i in range(0, len(self.convs), 4):
            x = self.convs[i](x)      # Conv1d
            x = self.convs[i+1](x)    # ReLU
            x = x.transpose(1, 2)     # [B, T, d_model] for LayerNorm
            x = self.convs[i+2](x)    # LayerNorm
            x = self.convs[i+3](x)    # Dropout
            x = x.transpose(1, 2)     # back to [B, d_model, T]
        x = x.transpose(1, 2)  # [B, T, d_model]
        log_dur = self.proj(x).squeeze(-1)  # [B, T_text]
        if text_mask is not None:
            log_dur = log_dur.masked_fill(text_mask, 0.0)
        return log_dur

def length_regulate(text_emb, durations, text_mask=None):
    """
    Expand text embeddings by repeating each token according to its duration.
    text_emb: [B, T_text, d_model]
    durations: [B, T_text] integer durations
    Returns: expanded [B, T_audio, d_model], where T_audio = sum(durations) per sample
    """
    B, T_text, D = text_emb.shape
    max_audio_len = durations.sum(dim=1).max().item()
    expanded = torch.zeros(B, max_audio_len, D, device=text_emb.device)
    for b in range(B):
        pos = 0
        for t in range(T_text):
            if text_mask is not None and text_mask[b, t]:
                continue
            dur = durations[b, t].item()
            if dur > 0 and pos < max_audio_len:
                end = min(pos + dur, max_audio_len)
                expanded[b, pos:end] = text_emb[b, t]
                pos = end
    return expanded

def compute_uniform_durations(text_lengths, audio_lengths):
    """
    Compute uniform durations: spread audio frames evenly across text characters.
    text_lengths: [B] number of valid text tokens
    audio_lengths: [B] number of valid audio frames
    Returns: durations [B, max_text_len] integer
    """
    B = text_lengths.shape[0]
    max_text_len = text_lengths.max().item()
    durations = torch.zeros(B, max_text_len, dtype=torch.long, device=text_lengths.device)
    for b in range(B):
        tl = text_lengths[b].item()
        al = audio_lengths[b].item()
        if tl == 0:
            continue
        base = al // tl
        remainder = al % tl
        durations[b, :tl] = base
        durations[b, :remainder] += 1
    return durations

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

class JEPAPredictor(nn.Module):
    """Causal predictor with frame-level text added to audio input."""
    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.text_input_proj = nn.Linear(d_model, d_model)
        self.pos_embed = SinusoidalPositionalEncoding(d_model, max_len=4096)
        self.dropout = nn.Dropout(dropout)

        # Use encoder (self-attention only) since text is now part of input
        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.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_expanded, audio_mask=None):
        """
        audio_emb: [B, T, d_model] — shifted audio embeddings
        text_emb_expanded: [B, T, d_model] — frame-level text (already expanded)
        """
        x = self.audio_input_proj(audio_emb) + self.text_input_proj(text_emb_expanded)
        x = self.pos_embed(x)
        x = self.dropout(x)
        T = x.shape[1]
        causal_mask = torch.triu(torch.ones(T, T, device=x.device, dtype=torch.bool), diagonal=1)
        predicted = self.transformer(x, mask=causal_mask, src_key_padding_mask=audio_mask)
        return self.output_proj(predicted)

    # Keep old cross-attention interface for backward compat loading
    def forward_cross_attn(self, audio_emb, text_emb, audio_mask=None, text_mask=None):
        """Legacy cross-attention forward (not used in duration-predictor mode)."""
        x = self.audio_input_proj(audio_emb)
        x = self.pos_embed(x)
        x = self.dropout(x)
        return self.output_proj(x)

    def _cached_mha(self, mha, q_input, kv_input, cache_k, cache_v):
        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)
        k_new = F.linear(kv_input, Wk, bk)
        v_new = F.linear(kv_input, Wv, bv)
        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, v = k_new, v_new
        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)
        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):
        return {
            'self_k': [None] * num_layers,
            'self_v': [None] * num_layers,
        }

    def inference_step(self, new_emb, text_emb_frame, step_idx, cache):
        """
        new_emb: [B, 1, d_model] — audio embedding for this step
        text_emb_frame: [B, 1, d_model] — expanded text embedding for this step
        """
        x = self.audio_input_proj(new_emb) + self.text_input_proj(text_emb_frame)
        x = x + self.pos_embed.pe[:, step_idx:step_idx + 1]
        for i, layer in enumerate(self.transformer.layers):
            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))
            ff_out = layer.linear2(layer.dropout(layer.activation(layer.linear1(x))))
            x = layer.norm2(x + layer.dropout2(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))

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

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

class TextEncoder(nn.Module):
    def __init__(self, vocab_size=256, d_model=256, nhead=4, num_layers=4, dropout=0.1):
        super().__init__()
        self.d_model = d_model
        self.char_embed = nn.Embedding(vocab_size, d_model)
        self.pos_embed = SinusoidalPositionalEncoding(d_model, max_len=2048)
        self.dropout = nn.Dropout(dropout)
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=d_model, nhead=nhead, dim_feedforward=d_model * 4,
            dropout=dropout, batch_first=True, activation="gelu",
        )
        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
        self.proj = nn.Linear(d_model, d_model)

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

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

class LeWMTTSv5(nn.Module):
    """
    JEPA TTS with EnCodec backend + duration predictor.
    Text is expanded to frame level and added to audio input — can't be ignored.

    Training: codec_emb → proj_in → predict(audio + aligned_text) → loss vs targets
              Duration predictor trained with uniform alignment as supervision.
    Inference: text → duration predict → expand text → AR predict → proj_out → decode
    """

    def __init__(self, config):
        super().__init__()
        d_model = config.get("d_model", 256)
        nhead = config.get("nhead", 4)
        codec_dim = config.get("codec_dim", 128)
        text_vocab = config.get("text_vocab_size", 256)
        text_layers = config.get("text_encoder_layers", 4)
        predictor_layers = config.get("predictor_layers", 6)
        dropout = config.get("dropout", 0.1)

        self.codec_dim = codec_dim

        self.text_encoder = TextEncoder(
            vocab_size=text_vocab, d_model=d_model, nhead=nhead,
            num_layers=text_layers, dropout=dropout,
        )

        # Duration predictor
        self.duration_predictor = DurationPredictor(
            d_model=d_model, kernel_size=3, num_layers=2, dropout=dropout,
        )

        # Simple linear projections — no lossy compression
        self.proj_in = nn.Sequential(
            nn.Linear(codec_dim, d_model),
            nn.GELU(),
            nn.Linear(d_model, d_model),
        )
        self.proj_out = nn.Sequential(
            nn.Linear(d_model, d_model),
            nn.GELU(),
            nn.Linear(d_model, codec_dim),
        )

        self.predictor = JEPAPredictor(
            d_model=d_model, nhead=nhead,
            num_layers=predictor_layers, dropout=dropout,
        )

        # EMA on proj_in (provides stable prediction targets)
        self.ema_proj_in = copy.deepcopy(self.proj_in)
        for p in self.ema_proj_in.parameters():
            p.requires_grad = False
        self.ema_decay = config.get("ema_decay", 0.998)

        # Learnable start embedding
        self.start_emb = nn.Parameter(torch.randn(1, 1, d_model) * 0.02)

        # Loss weights
        self.pred_weight = config.get("pred_weight", 10.0)
        self.roundtrip_weight = config.get("roundtrip_weight", 1.0)
        self.dur_weight = config.get("dur_weight", 1.0)

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

        self.config = config

    @torch.no_grad()
    def update_ema(self):
        for p_online, p_ema in zip(self.proj_in.parameters(),
                                   self.ema_proj_in.parameters()):
            p_ema.data.mul_(self.ema_decay).add_(p_online.data, alpha=1 - self.ema_decay)

    def forward(self, codec_emb, text_tokens, codec_mask=None, text_mask=None):
        """
        Args:
            codec_emb: [B, 128, T] — continuous EnCodec embeddings
            text_tokens: [B, T_text]
            codec_mask: [B, T] bool, True = padding
            text_mask: [B, T_text] bool
        """
        B = codec_emb.shape[0]
        T = codec_emb.shape[2]

        # codec_emb is [B, 128, T] → transpose to [B, T, 128]
        codec_seq = codec_emb.transpose(1, 2)

        # Text encoding
        text_emb = self.text_encoder(text_tokens, text_mask)

        # Project codec embeddings to predictor space
        z = self.proj_in(codec_seq)  # [B, T, d_model]

        # EMA targets (deterministic)
        with torch.no_grad():
            self.ema_proj_in.eval()
            target_emb = self.ema_proj_in(codec_seq)  # [B, T, d_model]
            self.ema_proj_in.train()

        # Prepend start embedding, shift right
        start = self.start_emb.expand(B, -1, -1)
        input_emb = torch.cat([start, z[:, :-1]], dim=1)  # [B, T, d_model]

        # Input noise
        if self.training and self.input_noise > 0:
            noise = torch.zeros_like(input_emb)
            noise[:, 1:] = torch.randn(B, T - 1, input_emb.shape[-1],
                                        device=input_emb.device) * self.input_noise
            input_emb = input_emb + noise

        # Masks
        if codec_mask is not None:
            start_mask = torch.zeros(B, 1, dtype=torch.bool, device=codec_emb.device)
            pred_mask = torch.cat([start_mask, codec_mask[:, :-1]], dim=1)
            loss_mask = codec_mask
        else:
            pred_mask = None
            loss_mask = None

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

        # ─── Loss 1: Prediction (MSE vs EMA targets) ───
        if loss_mask is not None:
            valid = (~loss_mask).unsqueeze(-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: Roundtrip (proj_in → proj_out → match original codec_emb) ───
        codec_recon = self.proj_out(z)  # [B, T, 128]

        if codec_mask is not None:
            valid_rt = (~codec_mask).unsqueeze(-1).float()
            roundtrip_loss = F.l1_loss(
                codec_recon * valid_rt, codec_seq * valid_rt, reduction="sum"
            ) / (valid_rt.sum() * self.codec_dim + 1e-8)
        else:
            roundtrip_loss = F.l1_loss(codec_recon, codec_seq)

        total_loss = self.pred_weight * prediction_loss + self.roundtrip_weight * roundtrip_loss

        return {
            "total_loss": total_loss,
            "prediction_loss": prediction_loss,
            "roundtrip_loss": roundtrip_loss,
        }

    def init_ar_cache(self):
        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):
        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, "codec_dim": 128,
            "text_vocab_size": 256, "text_encoder_layers": 4,
            "predictor_layers": 6, "dropout": 0.1,
            "pred_weight": 10.0, "roundtrip_weight": 1.0,
            "ema_decay": 0.998, "input_noise": 0.0,
        }
    model = LeWMTTSv5(config)
    print(f"LeWM TTS v5: {count_parameters(model)/1e6:.2f}M trainable parameters")
    return model, config

+        # Duration prediction + ground truth durations (uniform alignment)

+        log_dur_pred = self.duration_predictor(text_emb, text_mask)

+        # Compute actual text/audio lengths from masks

+            text_lengths = (~text_mask).sum(dim=1)  # [B]

+            text_lengths = torch.full((B,), text_tokens.shape[1], device=text_tokens.device)

+            audio_lengths = (~codec_mask).sum(dim=1)  # [B]

+            audio_lengths = torch.full((B,), T, device=codec_emb.device)

+        # Ground truth durations (uniform distribution)

+        gt_durations = compute_uniform_durations(text_lengths, audio_lengths)

+        log_dur_gt = torch.log(gt_durations.float().clamp(min=1))

+        # Duration loss (MSE on log-durations)

+            valid_text = (~text_mask).float()

+                log_dur_pred * valid_text, log_dur_gt * valid_text, reduction="sum"

+            ) / (valid_text.sum() + 1e-8)

+            dur_loss = F.mse_loss(log_dur_pred, log_dur_gt)

+        # Expand text to frame level using GT durations (training uses GT)

+        text_emb_expanded = length_regulate(text_emb, gt_durations, text_mask)

+        # Trim/pad to match audio length T

+        if text_emb_expanded.shape[1] > T:

+            text_emb_expanded = text_emb_expanded[:, :T]

+        elif text_emb_expanded.shape[1] < T:

+            pad = torch.zeros(B, T - text_emb_expanded.shape[1], text_emb.shape[-1],

+                            device=text_emb.device)

+            text_emb_expanded = torch.cat([text_emb_expanded, pad], dim=1)

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

+        # JEPA prediction with frame-level text

+        predicted = self.predictor(input_emb, text_emb_expanded, pred_mask)

-        total_loss = self.pred_weight * prediction_loss + self.roundtrip_weight * roundtrip_loss

+        total_loss = (self.pred_weight * prediction_loss

+                     + self.roundtrip_weight * roundtrip_loss

+                     + self.dur_weight * dur_loss)


Now update the inference methods:


    def init_ar_cache(self):
        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):
        return self.predictor.inference_step(new_emb, text_emb, step_idx, cache, text_mask=text_mask)

    def predict_durations(self, text_emb, text_mask=None):
        """Predict durations from text embeddings (for inference)."""
        log_dur = self.duration_predictor(text_emb, text_mask)
        durations = torch.exp(log_dur).round().long().clamp(min=1)
        if text_mask is not None:
            durations = durations.masked_fill(text_mask, 0)
        return durations

    def init_ar_cache(self):
        num_layers = len(self.predictor.transformer.layers)
        return self.predictor.init_cache(num_layers)

    def predict_next_cached(self, new_emb, text_emb_frame, step_idx, cache):
        """
        new_emb: [B, 1, d_model]
        text_emb_frame: [B, 1, d_model] — expanded text for this frame
        """
        return self.predictor.inference_step(new_emb, text_emb_frame, step_idx, cache)


    def predict_durations(self, text_emb, text_mask=None):
        """Predict durations from text embeddings (for inference)."""
        log_dur = self.duration_predictor(text_emb, text_mask)
        durations = torch.exp(log_dur).round().long().clamp(min=1)
        if text_mask is not None:
            durations = durations.masked_fill(text_mask, 0)
        return durations

    def init_ar_cache(self):
        num_layers = len(self.predictor.transformer.layers)
        return self.predictor.init_cache(num_layers)

    def predict_next_cached(self, new_emb, text_emb_frame, step_idx, cache):
        """
        new_emb: [B, 1, d_model]
        text_emb_frame: [B, 1, d_model] — expanded text for this frame
        """
        return self.predictor.inference_step(new_emb, text_emb_frame, step_idx, cache)

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

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

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

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

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

    def forward(self, text_emb, text_mask=None):
        """
        text_emb: [B, T_text, d_model]
        Returns: log_durations [B, T_text]
        """
        x = text_emb.transpose(1, 2)  # [B, d_model, T]
        for i in range(0, len(self.convs), 4):
            x = self.convs[i](x)      # Conv1d
            x = self.convs[i+1](x)    # ReLU
            x = x.transpose(1, 2)     # [B, T, d_model] for LayerNorm
            x = self.convs[i+2](x)    # LayerNorm
            x = self.convs[i+3](x)    # Dropout
            x = x.transpose(1, 2)     # back to [B, d_model, T]
        x = x.transpose(1, 2)  # [B, T, d_model]
        log_dur = self.proj(x).squeeze(-1)  # [B, T_text]
        if text_mask is not None:
            log_dur = log_dur.masked_fill(text_mask, 0.0)
        return log_dur

def length_regulate(text_emb, durations, text_mask=None):
    """
    Expand text embeddings by repeating each token according to its duration.
    text_emb: [B, T_text, d_model]
    durations: [B, T_text] integer durations
    Returns: expanded [B, T_audio, d_model], where T_audio = sum(durations) per sample
    """
    B, T_text, D = text_emb.shape
    max_audio_len = durations.sum(dim=1).max().item()
    expanded = torch.zeros(B, max_audio_len, D, device=text_emb.device)
    for b in range(B):
        pos = 0
        for t in range(T_text):
            if text_mask is not None and text_mask[b, t]:
                continue
            dur = durations[b, t].item()
            if dur > 0 and pos < max_audio_len:
                end = min(pos + dur, max_audio_len)
                expanded[b, pos:end] = text_emb[b, t]
                pos = end
    return expanded

def compute_uniform_durations(text_lengths, audio_lengths):
    """
    Compute uniform durations: spread audio frames evenly across text characters.
    text_lengths: [B] number of valid text tokens
    audio_lengths: [B] number of valid audio frames
    Returns: durations [B, max_text_len] integer
    """
    B = text_lengths.shape[0]
    max_text_len = text_lengths.max().item()
    durations = torch.zeros(B, max_text_len, dtype=torch.long, device=text_lengths.device)
    for b in range(B):
        tl = text_lengths[b].item()
        al = audio_lengths[b].item()
        if tl == 0:
            continue
        base = al // tl
        remainder = al % tl
        durations[b, :tl] = base
        durations[b, :remainder] += 1
    return durations

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

class JEPAPredictor(nn.Module):
    """Causal predictor with frame-level text added to audio input."""
    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.text_input_proj = nn.Linear(d_model, d_model)
        self.pos_embed = SinusoidalPositionalEncoding(d_model, max_len=4096)
        self.dropout = nn.Dropout(dropout)

        # Use encoder (self-attention only) since text is now part of input
        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.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_expanded, audio_mask=None):
        """
        audio_emb: [B, T, d_model] — shifted audio embeddings
        text_emb_expanded: [B, T, d_model] — frame-level text (already expanded)
        """
        x = self.audio_input_proj(audio_emb) + self.text_input_proj(text_emb_expanded)
        x = self.pos_embed(x)
        x = self.dropout(x)
        T = x.shape[1]
        causal_mask = torch.triu(torch.ones(T, T, device=x.device, dtype=torch.bool), diagonal=1)
        predicted = self.transformer(x, mask=causal_mask, src_key_padding_mask=audio_mask)
        return self.output_proj(predicted)

    # Keep old cross-attention interface for backward compat loading
    def forward_cross_attn(self, audio_emb, text_emb, audio_mask=None, text_mask=None):
        """Legacy cross-attention forward (not used in duration-predictor mode)."""
        x = self.audio_input_proj(audio_emb)
        x = self.pos_embed(x)
        x = self.dropout(x)
        return self.output_proj(x)

    def _cached_mha(self, mha, q_input, kv_input, cache_k, cache_v):
        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)
        k_new = F.linear(kv_input, Wk, bk)
        v_new = F.linear(kv_input, Wv, bv)
        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, v = k_new, v_new
        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)
        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):
        return {
            'self_k': [None] * num_layers,
            'self_v': [None] * num_layers,
        }

    def inference_step(self, new_emb, text_emb_frame, step_idx, cache):
        """
        new_emb: [B, 1, d_model] — audio embedding for this step
        text_emb_frame: [B, 1, d_model] — expanded text embedding for this step
        """
        x = self.audio_input_proj(new_emb) + self.text_input_proj(text_emb_frame)
        x = x + self.pos_embed.pe[:, step_idx:step_idx + 1]
        for i, layer in enumerate(self.transformer.layers):
            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))
            ff_out = layer.linear2(layer.dropout(layer.activation(layer.linear1(x))))
            x = layer.norm2(x + layer.dropout2(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))

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

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

class TextEncoder(nn.Module):
    def __init__(self, vocab_size=256, d_model=256, nhead=4, num_layers=4, dropout=0.1):
        super().__init__()
        self.d_model = d_model
        self.char_embed = nn.Embedding(vocab_size, d_model)
        self.pos_embed = SinusoidalPositionalEncoding(d_model, max_len=2048)
        self.dropout = nn.Dropout(dropout)
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=d_model, nhead=nhead, dim_feedforward=d_model * 4,
            dropout=dropout, batch_first=True, activation="gelu",
        )
        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
        self.proj = nn.Linear(d_model, d_model)

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

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

class LeWMTTSv5(nn.Module):
    """
    JEPA TTS with EnCodec backend + duration predictor.
    Text is expanded to frame level and added to audio input — can't be ignored.

    Training: codec_emb → proj_in → predict(audio + aligned_text) → loss vs targets
              Duration predictor trained with uniform alignment as supervision.
    Inference: text → duration predict → expand text → AR predict → proj_out → decode
    """

    def __init__(self, config):
        super().__init__()
        d_model = config.get("d_model", 256)
        nhead = config.get("nhead", 4)
        codec_dim = config.get("codec_dim", 128)
        text_vocab = config.get("text_vocab_size", 256)
        text_layers = config.get("text_encoder_layers", 4)
        predictor_layers = config.get("predictor_layers", 6)
        dropout = config.get("dropout", 0.1)

        self.codec_dim = codec_dim

        self.text_encoder = TextEncoder(
            vocab_size=text_vocab, d_model=d_model, nhead=nhead,
            num_layers=text_layers, dropout=dropout,
        )

        # Duration predictor
        self.duration_predictor = DurationPredictor(
            d_model=d_model, kernel_size=3, num_layers=2, dropout=dropout,
        )

        # Simple linear projections — no lossy compression
        self.proj_in = nn.Sequential(
            nn.Linear(codec_dim, d_model),
            nn.GELU(),
            nn.Linear(d_model, d_model),
        )
        self.proj_out = nn.Sequential(
            nn.Linear(d_model, d_model),
            nn.GELU(),
            nn.Linear(d_model, codec_dim),
        )

        self.predictor = JEPAPredictor(
            d_model=d_model, nhead=nhead,
            num_layers=predictor_layers, dropout=dropout,
        )

        # EMA on proj_in (provides stable prediction targets)
        self.ema_proj_in = copy.deepcopy(self.proj_in)
        for p in self.ema_proj_in.parameters():
            p.requires_grad = False
        self.ema_decay = config.get("ema_decay", 0.998)

        # Learnable start embedding
        self.start_emb = nn.Parameter(torch.randn(1, 1, d_model) * 0.02)

        # Loss weights
        self.pred_weight = config.get("pred_weight", 10.0)
        self.roundtrip_weight = config.get("roundtrip_weight", 1.0)
        self.dur_weight = config.get("dur_weight", 1.0)

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

        self.config = config

    @torch.no_grad()
    def update_ema(self):
        for p_online, p_ema in zip(self.proj_in.parameters(),
                                   self.ema_proj_in.parameters()):
            p_ema.data.mul_(self.ema_decay).add_(p_online.data, alpha=1 - self.ema_decay)

    def forward(self, codec_emb, text_tokens, codec_mask=None, text_mask=None):
        """
        Args:
            codec_emb: [B, 128, T] — continuous EnCodec embeddings
            text_tokens: [B, T_text]
            codec_mask: [B, T] bool, True = padding
            text_mask: [B, T_text] bool
        """
        B = codec_emb.shape[0]
        T = codec_emb.shape[2]

        # codec_emb is [B, 128, T] → transpose to [B, T, 128]
        codec_seq = codec_emb.transpose(1, 2)

        # Text encoding
        text_emb = self.text_encoder(text_tokens, text_mask)

        # Duration prediction + ground truth durations (uniform alignment)
        log_dur_pred = self.duration_predictor(text_emb, text_mask)

        # Compute actual text/audio lengths from masks
        if text_mask is not None:
            text_lengths = (~text_mask).sum(dim=1)  # [B]
        else:
            text_lengths = torch.full((B,), text_tokens.shape[1], device=text_tokens.device)
        if codec_mask is not None:
            audio_lengths = (~codec_mask).sum(dim=1)  # [B]
        else:
            audio_lengths = torch.full((B,), T, device=codec_emb.device)

        # Ground truth durations (uniform distribution)
        gt_durations = compute_uniform_durations(text_lengths, audio_lengths)
        log_dur_gt = torch.log(gt_durations.float().clamp(min=1))

        # Duration loss (MSE on log-durations)
        if text_mask is not None:
            valid_text = (~text_mask).float()
            dur_loss = F.mse_loss(
                log_dur_pred * valid_text, log_dur_gt * valid_text, reduction="sum"
            ) / (valid_text.sum() + 1e-8)
        else:
            dur_loss = F.mse_loss(log_dur_pred, log_dur_gt)

        # Expand text to frame level using GT durations (training uses GT)
        text_emb_expanded = length_regulate(text_emb, gt_durations, text_mask)
        # Trim/pad to match audio length T
        if text_emb_expanded.shape[1] > T:
            text_emb_expanded = text_emb_expanded[:, :T]
        elif text_emb_expanded.shape[1] < T:
            pad = torch.zeros(B, T - text_emb_expanded.shape[1], text_emb.shape[-1],
                            device=text_emb.device)
            text_emb_expanded = torch.cat([text_emb_expanded, pad], dim=1)

        # Project codec embeddings to predictor space
        z = self.proj_in(codec_seq)  # [B, T, d_model]

        # EMA targets (deterministic)
        with torch.no_grad():
            self.ema_proj_in.eval()
            target_emb = self.ema_proj_in(codec_seq)  # [B, T, d_model]
            self.ema_proj_in.train()

        # Prepend start embedding, shift right
        start = self.start_emb.expand(B, -1, -1)
        input_emb = torch.cat([start, z[:, :-1]], dim=1)  # [B, T, d_model]

        # Input noise
        if self.training and self.input_noise > 0:
            noise = torch.zeros_like(input_emb)
            noise[:, 1:] = torch.randn(B, T - 1, input_emb.shape[-1],
                                        device=input_emb.device) * self.input_noise
            input_emb = input_emb + noise

        # Masks
        if codec_mask is not None:
            start_mask = torch.zeros(B, 1, dtype=torch.bool, device=codec_emb.device)
            pred_mask = torch.cat([start_mask, codec_mask[:, :-1]], dim=1)
            loss_mask = codec_mask
        else:
            pred_mask = None
            loss_mask = None

        # JEPA prediction with frame-level text
        predicted = self.predictor(input_emb, text_emb_expanded, pred_mask)

        # ─── Loss 1: Prediction (MSE vs EMA targets) ───
        if loss_mask is not None:
            valid = (~loss_mask).unsqueeze(-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: Roundtrip (proj_in → proj_out → match original codec_emb) ───
        codec_recon = self.proj_out(z)  # [B, T, 128]

        if codec_mask is not None:
            valid_rt = (~codec_mask).unsqueeze(-1).float()
            roundtrip_loss = F.l1_loss(
                codec_recon * valid_rt, codec_seq * valid_rt, reduction="sum"
            ) / (valid_rt.sum() * self.codec_dim + 1e-8)
        else:
            roundtrip_loss = F.l1_loss(codec_recon, codec_seq)

        total_loss = (self.pred_weight * prediction_loss
                     + self.roundtrip_weight * roundtrip_loss
                     + self.dur_weight * dur_loss)

        return {
            "total_loss": total_loss,
            "prediction_loss": prediction_loss,
            "roundtrip_loss": roundtrip_loss,
            "dur_loss": dur_loss,
        }

    def init_ar_cache(self):
        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):
        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, "codec_dim": 128,
            "text_vocab_size": 256, "text_encoder_layers": 4,
            "predictor_layers": 6, "dropout": 0.1,
            "pred_weight": 10.0, "roundtrip_weight": 1.0,
            "ema_decay": 0.998, "input_noise": 0.0,
        }
    model = LeWMTTSv5(config)
    print(f"LeWM TTS v5: {count_parameters(model)/1e6:.2f}M trainable parameters")
    return model, config

+    def predict_durations(self, text_emb, text_mask=None):

+        """Predict durations from text embeddings (for inference)."""

+        log_dur = self.duration_predictor(text_emb, text_mask)

+        durations = torch.exp(log_dur).round().long().clamp(min=1)

+            durations = durations.masked_fill(text_mask, 0)

-    def predict_next_cached(self, new_emb, text_emb, step_idx, cache, text_mask=None):

-        return self.predictor.inference_step(new_emb, text_emb, step_idx, cache, text_mask=text_mask)

+    def predict_next_cached(self, new_emb, text_emb_frame, step_idx, cache):

+        text_emb_frame: [B, 1, d_model] — expanded text for this frame

+        return self.predictor.inference_step(new_emb, text_emb_frame, step_idx, cache)


Now update the training script `train_v5_pred.py` to use the new architecture:


Write

"""
Train predictor + text encoder + duration predictor with frozen projections.
Loads proj_in/proj_out from a v5 checkpoint (already good at roundtrip).
Text is expanded to frame level via duration predictor — can't be ignored.
"""

import os
import json
import time
import argparse
import torch
import torch.nn.functional as F
from torch.utils.tensorboard import SummaryWriter
from pathlib import Path

from model_v5 import (
    LeWMTTSv5, build_model, JEPAPredictor, TextEncoder,
    DurationPredictor, compute_uniform_durations, length_regulate,
)
from dataset_v5 import build_dataloader

def train(args):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Device: {device}")

    # Load pretrained checkpoint for projections
    ckpt = torch.load(args.pretrained, map_location=device, weights_only=False)
    config = ckpt["config"]
    config["pred_weight"] = args.pred_weight
    config["dur_weight"] = args.dur_weight
    config["input_noise"] = 0.0

    model = LeWMTTSv5(config).to(device)

    # Load what we can from pretrained (projections, etc.)
    pretrained_state = ckpt["model"]
    model_state = model.state_dict()
    loaded_keys = []
    skipped_keys = []
    for k, v in pretrained_state.items():
        if k in model_state and model_state[k].shape == v.shape:
            model_state[k] = v
            loaded_keys.append(k)
        else:
            skipped_keys.append(k)
    model.load_state_dict(model_state)
    print(f"Loaded {len(loaded_keys)} params from {args.pretrained} (skipped {len(skipped_keys)})")
    if skipped_keys:
        print(f"  Skipped: {skipped_keys[:10]}{'...' if len(skipped_keys) > 10 else ''}")

    # Freeze proj_in and proj_out — they're already good
    for p in model.proj_in.parameters():
        p.requires_grad = False
    for p in model.proj_out.parameters():
        p.requires_grad = False
    for p in model.ema_proj_in.parameters():
        p.requires_grad = False

    # Reinitialize predictor + text encoder + duration predictor from scratch
    if args.reinit:
        print("Reinitializing predictor + text encoder + duration predictor from scratch")
        d_model = config.get("d_model", 256)
        nhead = config.get("nhead", 4)
        predictor_layers = config.get("predictor_layers", 6)
        dropout = config.get("dropout", 0.1)
        model.predictor = JEPAPredictor(
            d_model=d_model, nhead=nhead,
            num_layers=predictor_layers, dropout=dropout,
        ).to(device)
        model.text_encoder = TextEncoder(
            vocab_size=config.get("text_vocab_size", 256),
            d_model=d_model, nhead=nhead,
            num_layers=config.get("text_encoder_layers", 4),
            dropout=dropout,
        ).to(device)
        model.duration_predictor = DurationPredictor(
            d_model=d_model, kernel_size=3, num_layers=2, dropout=dropout,
        ).to(device)
        model.start_emb = torch.nn.Parameter(torch.randn(1, 1, d_model, device=device) * 0.02)

    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
    print(f"Trainable params: {trainable/1e6:.2f}M (predictor + text encoder + dur predictor + start_emb)")

    manifest_path = os.path.join(args.data_dir, "manifest.json")
    loader, dataset = build_dataloader(
        manifest_path, batch_size=args.batch_size,
        num_workers=args.num_workers, max_codec_frames=args.max_codec_frames,
    )
    print(f"Training: {len(dataset)} samples, {len(loader)} batches/epoch")

    # Only optimize trainable params
    optimizer = torch.optim.AdamW(
        [p for p in model.parameters() if p.requires_grad],
        lr=args.lr, betas=(0.9, 0.98), weight_decay=0.01,
    )

    total_steps = args.epochs * len(loader)
    warmup_steps = min(2000, total_steps // 10)

    def lr_lambda(step):
        if step < warmup_steps:
            return step / max(1, warmup_steps)
        progress = (step - warmup_steps) / max(1, total_steps - warmup_steps)
        return 0.5 * (1 + torch.cos(torch.tensor(progress * 3.14159)).item())

    scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)

    ckpt_dir = Path(args.output_dir) / "checkpoints"
    ckpt_dir.mkdir(parents=True, exist_ok=True)
    writer = SummaryWriter(Path(args.output_dir) / "logs")

    with open(Path(args.output_dir) / "config.json", "w") as f:
        json.dump(config, f, indent=2)

    model.train()
    model.proj_in.eval()
    model.proj_out.eval()
    best_loss = float("inf")
    global_step = 0

    for epoch in range(args.epochs):
        epoch_loss = 0.0
        epoch_pred = 0.0
        epoch_dur = 0.0
        t0 = time.time()

        for batch_idx, batch in enumerate(loader):
            codec_emb = batch["codec_emb"].to(device)
            text_tokens = batch["text_tokens"].to(device)
            codec_mask = batch["codec_mask"].to(device)
            text_mask = batch["text_mask"].to(device)

            B = codec_emb.shape[0]
            T = codec_emb.shape[2]
            codec_seq = codec_emb.transpose(1, 2)  # [B, T, 128]

            # Text encoding (trainable)
            text_emb = model.text_encoder(text_tokens, text_mask)

            # Duration prediction
            log_dur_pred = model.duration_predictor(text_emb, text_mask)

            # Ground truth durations (uniform)
            text_lengths = (~text_mask).sum(dim=1)
            audio_lengths = (~codec_mask).sum(dim=1)
            gt_durations = compute_uniform_durations(text_lengths, audio_lengths)
            log_dur_gt = torch.log(gt_durations.float().clamp(min=1))

            # Duration loss
            valid_text = (~text_mask).float()
            dur_loss = F.mse_loss(
                log_dur_pred * valid_text, log_dur_gt * valid_text, reduction="sum"
            ) / (valid_text.sum() + 1e-8)

            # Expand text to frame level using GT durations
            text_emb_expanded = length_regulate(text_emb, gt_durations, text_mask)
            if text_emb_expanded.shape[1] > T:
                text_emb_expanded = text_emb_expanded[:, :T]
            elif text_emb_expanded.shape[1] < T:
                pad = torch.zeros(B, T - text_emb_expanded.shape[1], text_emb.shape[-1],
                                device=device)
                text_emb_expanded = torch.cat([text_emb_expanded, pad], dim=1)

            # Project with FROZEN proj_in — targets are stable
            with torch.no_grad():
                z = model.proj_in(codec_seq)  # [B, T, d_model]
                target_emb = z.clone()

            # Prepend start, shift right
            start = model.start_emb.expand(B, -1, -1)
            input_emb = torch.cat([start, z[:, :-1]], dim=1)

            # Input noise
            noise_progress = min(1.0, global_step / (total_steps * 0.3))
            noise_level = args.input_noise_max * noise_progress
            if noise_level > 0:
                noise = torch.zeros_like(input_emb)
                noise[:, 1:] = torch.randn(B, T - 1, input_emb.shape[-1],
                                            device=device) * noise_level
                input_emb = input_emb + noise

            # Masks
            start_mask = torch.zeros(B, 1, dtype=torch.bool, device=device)
            pred_mask = torch.cat([start_mask, codec_mask[:, :-1]], dim=1)

            # Predict with frame-level text
            predicted = model.predictor(input_emb, text_emb_expanded, pred_mask)

            # Prediction loss (MSE vs frozen proj_in targets)
            valid = (~codec_mask).unsqueeze(-1)
            pred_loss = F.mse_loss(
                predicted * valid, target_emb * valid, reduction="sum"
            ) / (valid.sum() * predicted.shape[-1] + 1e-8)

            total_loss = args.pred_weight * pred_loss + args.dur_weight * dur_loss

            optimizer.zero_grad()
            total_loss.backward()
            grad_norm = torch.nn.utils.clip_grad_norm_(
                [p for p in model.parameters() if p.requires_grad], args.grad_clip
            )
            optimizer.step()
            scheduler.step()

            epoch_loss += pred_loss.item()
            epoch_pred += pred_loss.item()
            epoch_dur += dur_loss.item()
            global_step += 1

            if global_step % args.log_every == 0:
                writer.add_scalar("train/pred_loss", pred_loss.item(), global_step)
                writer.add_scalar("train/dur_loss", dur_loss.item(), global_step)
                writer.add_scalar("train/total_loss", total_loss.item(), global_step)
                writer.add_scalar("train/grad_norm", grad_norm.item(), global_step)
                writer.add_scalar("train/lr", optimizer.param_groups[0]["lr"], global_step)

            if global_step % args.print_every == 0:
                elapsed = time.time() - t0
                print(
                    f"  [{epoch+1}/{args.epochs}] step {global_step} | "
                    f"pred={pred_loss.item():.4f} dur={dur_loss.item():.4f} | "
                    f"grad={grad_norm.item():.2f} | {elapsed:.1f}s"
                )

        avg_loss = epoch_loss / len(loader)
        avg_dur = epoch_dur / len(loader)
        epoch_time = time.time() - t0
        print(f"Epoch {epoch+1}/{args.epochs} | pred={avg_loss:.4f} dur={avg_dur:.4f} | {epoch_time:.1f}s")

        if (epoch + 1) % args.save_every == 0 or (epoch + 1) == args.epochs:
            ckpt_path = ckpt_dir / f"epoch_{epoch+1:04d}.pt"
            torch.save({
                "epoch": epoch,
                "global_step": global_step,
                "model": model.state_dict(),
                "optimizer": optimizer.state_dict(),
                "scheduler": scheduler.state_dict(),
                "config": config,
                "loss": avg_loss,
            }, ckpt_path)
            print(f"  Saved: {ckpt_path}")

        if avg_loss < best_loss:
            best_loss = avg_loss
            torch.save({
                "epoch": epoch,
                "global_step": global_step,
                "model": model.state_dict(),
                "config": config,
                "loss": avg_loss,
            }, ckpt_dir / "best.pt")
            print(f"  New best (pred={avg_loss:.4f})")

    writer.close()
    print(f"\nDone. Best pred loss: {best_loss:.4f}")

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--pretrained", default="output_v5/checkpoints/epoch_0120.pt")
    parser.add_argument("--data_dir", default="processed_data_codec")
    parser.add_argument("--output_dir", default="output_v5_pred")
    parser.add_argument("--pred_weight", type=float, default=1.0)
    parser.add_argument("--dur_weight", type=float, default=1.0)
    parser.add_argument("--input_noise_max", type=float, default=0.05)
    parser.add_argument("--max_codec_frames", type=int, default=900)
    parser.add_argument("--epochs", type=int, default=200)
    parser.add_argument("--batch_size", type=int, default=256)
    parser.add_argument("--lr", type=float, default=5e-5)
    parser.add_argument("--grad_clip", type=float, default=1.0)
    parser.add_argument("--num_workers", type=int, default=4)
    parser.add_argument("--log_every", type=int, default=10)
    parser.add_argument("--print_every", type=int, default=20)
    parser.add_argument("--save_every", type=int, default=10)
    parser.add_argument("--reinit", action="store_true", help="Reinitialize predictor from scratch")
    args = parser.parse_args()
    train(args)