# ─── 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() * 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)

"""
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
  - JEPAPredictor: causal transformer, predicts next embedding
  - EMA target on projected embeddings
"""

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

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

class JEPAPredictor(nn.Module):
    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)

        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)
        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):
        x = self.audio_input_proj(audio_emb)
        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, text_emb, tgt_mask=causal_mask,
            tgt_key_padding_mask=audio_mask, memory_key_padding_mask=text_mask,
        )
        return self.output_proj(predicted)

    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 _cached_cross_attn(self, mha, q_input, memory, cache_k, cache_v, memory_mask=None):
        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)
        if cache_k is None:
            k = F.linear(memory, Wk, bk)
            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:
            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):
        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):
        x = self.audio_input_proj(new_emb)
        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))
            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))
            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))

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

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

class LeWMTTSv5(nn.Module):
    """
    JEPA TTS with EnCodec backend.

    Training flow:
        codec_emb → downsample → JEPA predict → loss vs EMA targets
        codec_emb → downsample → upsample → loss vs codec_emb (roundtrip)

    Inference flow:
        text → predictor (AR) → upsample → 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)
        ds_stride = config.get("ds_stride", 3)

        self.text_encoder = TextEncoder(
            vocab_size=text_vocab, d_model=d_model, nhead=nhead,
            num_layers=text_layers, dropout=dropout,
        )
        self.downsampler = CodecDownsampler(
            codec_dim=codec_dim, d_model=d_model, stride=ds_stride,
        )
        self.upsampler = CodecUpsampler(
            d_model=d_model, codec_dim=codec_dim, stride=ds_stride,
        )
        self.predictor = JEPAPredictor(
            d_model=d_model, nhead=nhead,
            num_layers=predictor_layers, dropout=dropout,
        )

        # EMA target on downsampler output
        self.ema_downsampler = copy.deepcopy(self.downsampler)
        for p in self.ema_downsampler.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.downsampler.parameters(),
                                   self.ema_downsampler.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_75Hz] — continuous EnCodec embeddings
            text_tokens: [B, T_text]
            codec_mask: [B, T_75Hz] bool, True = padding
            text_mask: [B, T_text] bool
        """
        B = codec_emb.shape[0]

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

        # Online downsampler: 75Hz×128d → 25Hz×256d
        z, ds_mask = self.downsampler(codec_emb, codec_mask)  # [B, T_down, d_model]
        T_down = z.shape[1]

        # EMA target downsampler (deterministic)
        with torch.no_grad():
            self.ema_downsampler.eval()
            target_emb, _ = self.ema_downsampler(codec_emb, codec_mask)
            self.ema_downsampler.train()

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

        # Input noise on 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

        # Masks
        if ds_mask is not None:
            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

        # 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 (downsample → upsample → match original codec_emb) ───
        T_codec = codec_emb.shape[2]
        codec_recon = self.upsampler(z, target_len=T_codec)  # [B, 128, T_75Hz]

        if codec_mask is not None:
            valid_codec = (~codec_mask).unsqueeze(1).float()
            roundtrip_loss = F.l1_loss(
                codec_recon * valid_codec, codec_emb * valid_codec, reduction="sum"
            ) / (valid_codec.sum() * codec_emb.shape[1] + 1e-8)
        else:
            roundtrip_loss = F.l1_loss(codec_recon, codec_emb)

        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, "ds_stride": 3,
            "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

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

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

+        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.transformer(x, src_key_padding_mask=text_mask)

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

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

-        codec_emb → downsample → JEPA predict → loss vs EMA targets

-        codec_emb → downsample → upsample → loss vs codec_emb (roundtrip)

-        text → predictor (AR) → upsample → EnCodec decoder → waveform

+    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

-        ds_stride = config.get("ds_stride", 3)

-        self.downsampler = CodecDownsampler(

-            codec_dim=codec_dim, d_model=d_model, stride=ds_stride,

+        # Simple linear projections — no lossy compression

+            nn.Linear(codec_dim, d_model),

+            nn.Linear(d_model, d_model),

-        self.upsampler = CodecUpsampler(

-            d_model=d_model, codec_dim=codec_dim, stride=ds_stride,

+            nn.Linear(d_model, codec_dim),

-        # EMA target on downsampler output

-        self.ema_downsampler = copy.deepcopy(self.downsampler)

-        for p in self.ema_downsampler.parameters():

+        # 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():

-        for p_online, p_ema in zip(self.downsampler.parameters(),

-                                   self.ema_downsampler.parameters()):

+        for p_online, p_ema in zip(self.proj_in.parameters(),

+                                   self.ema_proj_in.parameters()):

-            codec_emb: [B, 128, T_75Hz] — continuous EnCodec embeddings

+            codec_emb: [B, 128, T] — continuous EnCodec embeddings

-            codec_mask: [B, T_75Hz] bool, True = padding

+            codec_mask: [B, T] bool, True = padding

+        # codec_emb is [B, 128, T] → transpose to [B, T, 128]

+        codec_seq = codec_emb.transpose(1, 2)

-        # Online downsampler: 75Hz×128d → 25Hz×256d

-        z, ds_mask = self.downsampler(codec_emb, codec_mask)  # [B, T_down, d_model]

+        # Project codec embeddings to predictor space

+        z = self.proj_in(codec_seq)  # [B, T, d_model]

-        # EMA target downsampler (deterministic)

-            self.ema_downsampler.eval()

-            target_emb, _ = self.ema_downsampler(codec_emb, codec_mask)

-            self.ema_downsampler.train()

+            target_emb = self.ema_proj_in(codec_seq)  # [B, T, d_model]

+        # Prepend start embedding, shift right

start = self.start_emb.expand(B, -1, -1)

-        input_emb = torch.cat([start, z[:, :-1]], dim=1)

+        input_emb = torch.cat([start, z[:, :-1]], dim=1)  # [B, T, d_model]

-        # Input noise on real embeddings (not start_emb)

if self.training and self.input_noise > 0:

-            noise[:, 1:] = torch.randn(B, T_down - 1, input_emb.shape[-1],

+            noise[:, 1:] = torch.randn(B, T - 1, input_emb.shape[-1],

-            start_mask = torch.zeros(B, 1, dtype=torch.bool, device=z.device)

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

+            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 2: Roundtrip (downsample → upsample → match original codec_emb) ───

-        codec_recon = self.upsampler(z, target_len=T_codec)  # [B, 128, T_75Hz]

+        # ─── Loss 2: Roundtrip (proj_in → proj_out → match original codec_emb) ───

+        codec_recon = self.proj_out(z)  # [B, T, 128]

-            valid_codec = (~codec_mask).unsqueeze(1).float()

+            valid_rt = (~codec_mask).unsqueeze(-1).float()

-                codec_recon * valid_codec, codec_emb * valid_codec, reduction="sum"

-            ) / (valid_codec.sum() * codec_emb.shape[1] + 1e-8)

+                codec_recon * valid_rt, codec_seq * valid_rt, reduction="sum"

+            ) / (valid_rt.sum() * codec_dim + 1e-8)

-            roundtrip_loss = F.l1_loss(codec_recon, codec_emb)

+            roundtrip_loss = F.l1_loss(codec_recon, codec_seq)


Now update the inference script for the simplified model:


    42→        """AR synthesis: text → predicted embeddings → upsample → EnCodec decode."""
    43→        text_tokens = self.text_to_tokens(text)
    44→        text_emb = self.model.text_encoder(text_tokens)
    45→
    46→        start_emb = self.model.start_emb.to(self.device)
    47→        all_embeddings = [start_emb]
    48→
    49→        cache = self.model.init_ar_cache()
    50→        next_emb, cache = self.model.predict_next_cached(start_emb, text_emb, 0, cache)
    51→        if temperature > 0:
    52→            next_emb = next_emb + torch.randn_like(next_emb) * temperature
    53→        all_embeddings.append(next_emb)
    54→
    55→        prev_norms = [next_emb.norm().item()]
    56→
    57→        for step in range(1, max_steps):
    58→            next_emb, cache = self.model.predict_next_cached(next_emb, text_emb, step, cache)
    59→            if temperature > 0:
    60→                next_emb = next_emb + torch.randn_like(next_emb) * temperature
    61→            all_embeddings.append(next_emb)
    62→
    63→            norm = next_emb.norm().item()
    64→            prev_norms.append(norm)
    65→            if len(prev_norms) > 20:
    66→                recent = prev_norms[-20:]
    67→                if np.std(recent) < 0.01 * np.mean(recent) and step > 30:
    68→                    break
    69→
    70→        # Concat predicted embeddings and upsample to codec space
    71→        pred_embs = torch.cat(all_embeddings, dim=1)  # [1, T_25Hz, d_model]
    72→        codec_emb = self.model.upsampler(pred_embs)    # [1, 128, T_75Hz]
    73→
    74→        # EnCodec decode
    75→        waveform = self.codec.decoder(codec_emb)  # [1, 1, T_audio]
    76→        waveform = waveform.squeeze().cpu().numpy()
    77→
    78→        return waveform, 24000
    79→
    80→    @torch.no_grad()
    81→    def reconstruct(self, audio_path):
    82→        """Encode audio → downsample → upsample → decode. Tests roundtrip quality."""
    83→        from encodec.utils import convert_audio
    84→        wav, sr = torchaudio.load(audio_path)
    85→        wav = convert_audio(wav, sr, 24000, 1).to(self.device)
    86→
    87→        # EnCodec encode
    88→        codec_emb = self.codec.encoder(wav.unsqueeze(0))  # [1, 128, T]
    89→
    90→        # Downsample → upsample roundtrip
    91→        z, _ = self.model.downsampler(codec_emb)

        """AR synthesis: text → predicted embeddings → upsample → EnCodec decode."""
        text_tokens = self.text_to_tokens(text)
        text_emb = self.model.text_encoder(text_tokens)

        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
        all_embeddings.append(next_emb)

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

        for step in range(1, max_steps):
            next_emb, cache = self.model.predict_next_cached(next_emb, text_emb, step, cache)
            if temperature > 0:
                next_emb = next_emb + torch.randn_like(next_emb) * temperature
            all_embeddings.append(next_emb)

            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

        # Concat predicted embeddings and upsample to codec space
        pred_embs = torch.cat(all_embeddings, dim=1)  # [1, T_25Hz, d_model]
        codec_emb = self.model.upsampler(pred_embs)    # [1, 128, T_75Hz]

        # EnCodec decode
        waveform = self.codec.decoder(codec_emb)  # [1, 1, T_audio]
        waveform = waveform.squeeze().cpu().numpy()

        return waveform, 24000

    @torch.no_grad()
    def reconstruct(self, audio_path):
        """Encode audio → downsample → upsample → decode. Tests roundtrip quality."""
        from encodec.utils import convert_audio
        wav, sr = torchaudio.load(audio_path)
        wav = convert_audio(wav, sr, 24000, 1).to(self.device)

        # EnCodec encode
        codec_emb = self.codec.encoder(wav.unsqueeze(0))  # [1, 128, T]

        # Downsample → upsample roundtrip
        z, _ = self.model.downsampler(codec_emb)


Edit

        # Concat predicted embeddings and upsample to codec space
        pred_embs = torch.cat(all_embeddings, dim=1)  # [1, T_25Hz, d_model]
        codec_emb = self.model.upsampler(pred_embs)    # [1, 128, T_75Hz]

        # EnCodec decode
        waveform = self.codec.decoder(codec_emb)  # [1, 1, T_audio]

        # Concat predicted embeddings and project to codec space
        pred_embs = torch.cat(all_embeddings, dim=1)  # [1, T, d_model]
        codec_emb = self.model.proj_out(pred_embs)     # [1, T, 128]
        codec_emb = codec_emb.transpose(1, 2)          # [1, 128, T]

        # EnCodec decode
        waveform = self.codec.decoder(codec_emb)  # [1, 1, T_audio]


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

        # Concat predicted embeddings and project to codec space
        pred_embs = torch.cat(all_embeddings, dim=1)  # [1, T, d_model]
        codec_emb = self.model.proj_out(pred_embs)     # [1, T, 128]
        codec_emb = codec_emb.transpose(1, 2)          # [1, 128, T]

        # EnCodec decode
        waveform = self.codec.decoder(codec_emb)  # [1, 1, T_audio]

"""
Inference for LeWM TTS v5 (codec-based JEPA).
Text → JEPA predictor → upsample → EnCodec decoder → waveform
"""

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

from model_v5 import LeWMTTSv5

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

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

        self.model = LeWMTTSv5(config).to(self.device)
        self.model.load_state_dict(ckpt["model"])
        self.model.eval()

        # EnCodec decoder (frozen)
        from encodec import EncodecModel
        self.codec = EncodecModel.encodec_model_24khz()
        self.codec.set_target_bandwidth(6.0)
        self.codec.eval().to(self.device)

        self.config = config
        self.ds_stride = config.get("ds_stride", 3)
        print(f"v5 model loaded from {checkpoint_path}")

    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):
        """AR synthesis: text → predicted embeddings → upsample → EnCodec decode."""
        text_tokens = self.text_to_tokens(text)
        text_emb = self.model.text_encoder(text_tokens)

        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
        all_embeddings.append(next_emb)

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

        for step in range(1, max_steps):
            next_emb, cache = self.model.predict_next_cached(next_emb, text_emb, step, cache)
            if temperature > 0:
                next_emb = next_emb + torch.randn_like(next_emb) * temperature
            all_embeddings.append(next_emb)

            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

        # Concat predicted embeddings and upsample to codec space
        pred_embs = torch.cat(all_embeddings, dim=1)  # [1, T_25Hz, d_model]
        codec_emb = self.model.upsampler(pred_embs)    # [1, 128, T_75Hz]

        # EnCodec decode
        waveform = self.codec.decoder(codec_emb)  # [1, 1, T_audio]
        waveform = waveform.squeeze().cpu().numpy()

        return waveform, 24000

    @torch.no_grad()
    def reconstruct(self, audio_path):
        """Encode audio → downsample → upsample → decode. Tests roundtrip quality."""
        from encodec.utils import convert_audio
        wav, sr = torchaudio.load(audio_path)
        wav = convert_audio(wav, sr, 24000, 1).to(self.device)

        # EnCodec encode
        codec_emb = self.codec.encoder(wav.unsqueeze(0))  # [1, 128, T]

        # Downsample → upsample roundtrip
        z, _ = self.model.downsampler(codec_emb)
        codec_recon = self.model.upsampler(z, target_len=codec_emb.shape[2])

        # EnCodec decode
        waveform = self.codec.decoder(codec_recon).squeeze().cpu().numpy()
        return waveform, 24000

    @torch.no_grad()
    def reconstruct_codec_only(self, audio_path):
        """Pure EnCodec roundtrip (no JEPA). Baseline quality."""
        from encodec.utils import convert_audio
        wav, sr = torchaudio.load(audio_path)
        wav = convert_audio(wav, sr, 24000, 1).to(self.device)

        codec_emb = self.codec.encoder(wav.unsqueeze(0))
        waveform = self.codec.decoder(codec_emb).squeeze().cpu().numpy()
        return waveform, 24000

    @torch.no_grad()
    def reconstruct_with_prediction(self, audio_path, text):
        """Encode audio → downsample → teacher-forced predict → upsample → decode."""
        from encodec.utils import convert_audio
        wav, sr = torchaudio.load(audio_path)
        wav = convert_audio(wav, sr, 24000, 1).to(self.device)

        codec_emb = self.codec.encoder(wav.unsqueeze(0))
        z, _ = self.model.downsampler(codec_emb)

        text_tokens = self.text_to_tokens(text)
        text_emb = self.model.text_encoder(text_tokens)

        B = z.shape[0]
        start = self.model.start_emb.expand(B, -1, -1)
        input_emb = torch.cat([start, z[:, :-1]], dim=1)
        predicted = self.model.predictor(input_emb, text_emb)

        codec_recon = self.model.upsampler(predicted, target_len=codec_emb.shape[2])
        waveform = self.codec.decoder(codec_recon).squeeze().cpu().numpy()
        return waveform, 24000

    def save_audio(self, waveform, sr, output_path):
        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)
    parser.add_argument("--audio", default=None, help="Audio path for reconstruction")
    parser.add_argument("--mode", choices=["synthesize", "reconstruct", "codec_only", "predict"],
                        default="synthesize")
    parser.add_argument("--output", default="output_v5.wav")
    parser.add_argument("--max_steps", type=int, default=300)
    parser.add_argument("--temperature", type=float, default=0.0)
    args = parser.parse_args()

    tts = LeWMTTSv5Inference(args.checkpoint)

    if args.mode == "synthesize":
        assert args.text, "Need --text"
        waveform, sr = tts.synthesize(args.text, args.max_steps, args.temperature)
    elif args.mode == "reconstruct":
        assert args.audio, "Need --audio"
        waveform, sr = tts.reconstruct(args.audio)
    elif args.mode == "codec_only":
        assert args.audio, "Need --audio"
        waveform, sr = tts.reconstruct_codec_only(args.audio)
    elif args.mode == "predict":
        assert args.audio and args.text, "Need --audio and --text"
        waveform, sr = tts.reconstruct_with_prediction(args.audio, args.text)

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

if __name__ == "__main__":
    main()

-        # Concat predicted embeddings and upsample to codec space

-        pred_embs = torch.cat(all_embeddings, dim=1)  # [1, T_25Hz, d_model]

-        codec_emb = self.model.upsampler(pred_embs)    # [1, 128, T_75Hz]

+        # Concat predicted embeddings and project to codec space

+        pred_embs = torch.cat(all_embeddings, dim=1)  # [1, T, d_model]

+        codec_emb = self.model.proj_out(pred_embs)     # [1, T, 128]

+        codec_emb = codec_emb.transpose(1, 2)          # [1, 128, T]


    @torch.no_grad()
    def reconstruct(self, audio_path):
        """Encode audio → downsample → upsample → decode. Tests roundtrip quality."""
        from encodec.utils import convert_audio
        wav, sr = torchaudio.load(audio_path)
        wav = convert_audio(wav, sr, 24000, 1).to(self.device)

        # EnCodec encode
        codec_emb = self.codec.encoder(wav.unsqueeze(0))  # [1, 128, T]

        # Downsample → upsample roundtrip
        z, _ = self.model.downsampler(codec_emb)
        codec_recon = self.model.upsampler(z, target_len=codec_emb.shape[2])

    @torch.no_grad()
    def reconstruct(self, audio_path):
        """Encode audio → proj_in → proj_out → decode. Tests roundtrip quality."""
        from encodec.utils import convert_audio
        wav, sr = torchaudio.load(audio_path)
        wav = convert_audio(wav, sr, 24000, 1).to(self.device)

        # EnCodec encode
        codec_emb = self.codec.encoder(wav.unsqueeze(0))  # [1, 128, T]

        # proj_in → proj_out roundtrip
        z = self.model.proj_in(codec_emb.transpose(1, 2))  # [1, T, d_model]
        codec_recon = self.model.proj_out(z).transpose(1, 2)  # [1, 128, T]


    @torch.no_grad()
    def reconstruct(self, audio_path):
        """Encode audio → proj_in → proj_out → decode. Tests roundtrip quality."""
        from encodec.utils import convert_audio
        wav, sr = torchaudio.load(audio_path)
        wav = convert_audio(wav, sr, 24000, 1).to(self.device)

        # EnCodec encode
        codec_emb = self.codec.encoder(wav.unsqueeze(0))  # [1, 128, T]

        # proj_in → proj_out roundtrip
        z = self.model.proj_in(codec_emb.transpose(1, 2))  # [1, T, d_model]
        codec_recon = self.model.proj_out(z).transpose(1, 2)  # [1, 128, T]

"""
Inference for LeWM TTS v5 (codec-based JEPA).
Text → JEPA predictor → upsample → EnCodec decoder → waveform
"""

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

from model_v5 import LeWMTTSv5

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

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

        self.model = LeWMTTSv5(config).to(self.device)
        self.model.load_state_dict(ckpt["model"])
        self.model.eval()

        # EnCodec decoder (frozen)
        from encodec import EncodecModel
        self.codec = EncodecModel.encodec_model_24khz()
        self.codec.set_target_bandwidth(6.0)
        self.codec.eval().to(self.device)

        self.config = config
        self.ds_stride = config.get("ds_stride", 3)
        print(f"v5 model loaded from {checkpoint_path}")

    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):
        """AR synthesis: text → predicted embeddings → upsample → EnCodec decode."""
        text_tokens = self.text_to_tokens(text)
        text_emb = self.model.text_encoder(text_tokens)

        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
        all_embeddings.append(next_emb)

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

        for step in range(1, max_steps):
            next_emb, cache = self.model.predict_next_cached(next_emb, text_emb, step, cache)
            if temperature > 0:
                next_emb = next_emb + torch.randn_like(next_emb) * temperature
            all_embeddings.append(next_emb)

            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

        # Concat predicted embeddings and project to codec space
        pred_embs = torch.cat(all_embeddings, dim=1)  # [1, T, d_model]
        codec_emb = self.model.proj_out(pred_embs)     # [1, T, 128]
        codec_emb = codec_emb.transpose(1, 2)          # [1, 128, T]

        # EnCodec decode
        waveform = self.codec.decoder(codec_emb)  # [1, 1, T_audio]
        waveform = waveform.squeeze().cpu().numpy()

        return waveform, 24000

    @torch.no_grad()
    def reconstruct(self, audio_path):
        """Encode audio → downsample → upsample → decode. Tests roundtrip quality."""
        from encodec.utils import convert_audio
        wav, sr = torchaudio.load(audio_path)
        wav = convert_audio(wav, sr, 24000, 1).to(self.device)

        # EnCodec encode
        codec_emb = self.codec.encoder(wav.unsqueeze(0))  # [1, 128, T]

        # Downsample → upsample roundtrip
        z, _ = self.model.downsampler(codec_emb)
        codec_recon = self.model.upsampler(z, target_len=codec_emb.shape[2])

        # EnCodec decode
        waveform = self.codec.decoder(codec_recon).squeeze().cpu().numpy()
        return waveform, 24000

    @torch.no_grad()
    def reconstruct_codec_only(self, audio_path):
        """Pure EnCodec roundtrip (no JEPA). Baseline quality."""
        from encodec.utils import convert_audio
        wav, sr = torchaudio.load(audio_path)
        wav = convert_audio(wav, sr, 24000, 1).to(self.device)

        codec_emb = self.codec.encoder(wav.unsqueeze(0))
        waveform = self.codec.decoder(codec_emb).squeeze().cpu().numpy()
        return waveform, 24000

    @torch.no_grad()
    def reconstruct_with_prediction(self, audio_path, text):
        """Encode audio → downsample → teacher-forced predict → upsample → decode."""
        from encodec.utils import convert_audio
        wav, sr = torchaudio.load(audio_path)
        wav = convert_audio(wav, sr, 24000, 1).to(self.device)

        codec_emb = self.codec.encoder(wav.unsqueeze(0))
        z, _ = self.model.downsampler(codec_emb)

        text_tokens = self.text_to_tokens(text)
        text_emb = self.model.text_encoder(text_tokens)

        B = z.shape[0]
        start = self.model.start_emb.expand(B, -1, -1)
        input_emb = torch.cat([start, z[:, :-1]], dim=1)
        predicted = self.model.predictor(input_emb, text_emb)

        codec_recon = self.model.upsampler(predicted, target_len=codec_emb.shape[2])
        waveform = self.codec.decoder(codec_recon).squeeze().cpu().numpy()
        return waveform, 24000

    def save_audio(self, waveform, sr, output_path):
        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)
    parser.add_argument("--audio", default=None, help="Audio path for reconstruction")
    parser.add_argument("--mode", choices=["synthesize", "reconstruct", "codec_only", "predict"],
                        default="synthesize")
    parser.add_argument("--output", default="output_v5.wav")
    parser.add_argument("--max_steps", type=int, default=300)
    parser.add_argument("--temperature", type=float, default=0.0)
    args = parser.parse_args()

    tts = LeWMTTSv5Inference(args.checkpoint)

    if args.mode == "synthesize":
        assert args.text, "Need --text"
        waveform, sr = tts.synthesize(args.text, args.max_steps, args.temperature)
    elif args.mode == "reconstruct":
        assert args.audio, "Need --audio"
        waveform, sr = tts.reconstruct(args.audio)
    elif args.mode == "codec_only":
        assert args.audio, "Need --audio"
        waveform, sr = tts.reconstruct_codec_only(args.audio)
    elif args.mode == "predict":
        assert args.audio and args.text, "Need --audio and --text"
        waveform, sr = tts.reconstruct_with_prediction(args.audio, args.text)

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

if __name__ == "__main__":
    main()

-        """Encode audio → downsample → upsample → decode. Tests roundtrip quality."""

+        """Encode audio → proj_in → proj_out → decode. Tests roundtrip quality."""

-        # Downsample → upsample roundtrip

-        z, _ = self.model.downsampler(codec_emb)

-        codec_recon = self.model.upsampler(z, target_len=codec_emb.shape[2])

+        z = self.model.proj_in(codec_emb.transpose(1, 2))  # [1, T, d_model]

+        codec_recon = self.model.proj_out(z).transpose(1, 2)  # [1, 128, T]


    @torch.no_grad()
    def reconstruct_with_prediction(self, audio_path, text):
        """Encode audio → downsample → teacher-forced predict → upsample → decode."""
        from encodec.utils import convert_audio
        wav, sr = torchaudio.load(audio_path)
        wav = convert_audio(wav, sr, 24000, 1).to(self.device)

        codec_emb = self.codec.encoder(wav.unsqueeze(0))
        z, _ = self.model.downsampler(codec_emb)

        text_tokens = self.text_to_tokens(text)
        text_emb = self.model.text_encoder(text_tokens)

        B = z.shape[0]
        start = self.model.start_emb.expand(B, -1, -1)
        input_emb = torch.cat([start, z[:, :-1]], dim=1)
        predicted = self.model.predictor(input_emb, text_emb)

        codec_recon = self.model.upsampler(predicted, target_len=codec_emb.shape[2])

    @torch.no_grad()
    def reconstruct_with_prediction(self, audio_path, text):
        """Encode audio → proj_in → teacher-forced predict → proj_out → decode."""
        from encodec.utils import convert_audio
        wav, sr = torchaudio.load(audio_path)
        wav = convert_audio(wav, sr, 24000, 1).to(self.device)

        codec_emb = self.codec.encoder(wav.unsqueeze(0))  # [1, 128, T]
        z = self.model.proj_in(codec_emb.transpose(1, 2))  # [1, T, d_model]

        text_tokens = self.text_to_tokens(text)
        text_emb = self.model.text_encoder(text_tokens)

        B = z.shape[0]
        start = self.model.start_emb.expand(B, -1, -1)
        input_emb = torch.cat([start, z[:, :-1]], dim=1)
        predicted = self.model.predictor(input_emb, text_emb)

        codec_recon = self.model.proj_out(predicted).transpose(1, 2)  # [1, 128, T]


    @torch.no_grad()
    def reconstruct_with_prediction(self, audio_path, text):
        """Encode audio → proj_in → teacher-forced predict → proj_out → decode."""
        from encodec.utils import convert_audio
        wav, sr = torchaudio.load(audio_path)
        wav = convert_audio(wav, sr, 24000, 1).to(self.device)

        codec_emb = self.codec.encoder(wav.unsqueeze(0))  # [1, 128, T]
        z = self.model.proj_in(codec_emb.transpose(1, 2))  # [1, T, d_model]

        text_tokens = self.text_to_tokens(text)
        text_emb = self.model.text_encoder(text_tokens)

        B = z.shape[0]
        start = self.model.start_emb.expand(B, -1, -1)
        input_emb = torch.cat([start, z[:, :-1]], dim=1)
        predicted = self.model.predictor(input_emb, text_emb)

        codec_recon = self.model.proj_out(predicted).transpose(1, 2)  # [1, 128, T]

"""
Inference for LeWM TTS v5 (codec-based JEPA).
Text → JEPA predictor → upsample → EnCodec decoder → waveform
"""

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

from model_v5 import LeWMTTSv5

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

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

        self.model = LeWMTTSv5(config).to(self.device)
        self.model.load_state_dict(ckpt["model"])
        self.model.eval()

        # EnCodec decoder (frozen)
        from encodec import EncodecModel
        self.codec = EncodecModel.encodec_model_24khz()
        self.codec.set_target_bandwidth(6.0)
        self.codec.eval().to(self.device)

        self.config = config
        self.ds_stride = config.get("ds_stride", 3)
        print(f"v5 model loaded from {checkpoint_path}")

    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):
        """AR synthesis: text → predicted embeddings → upsample → EnCodec decode."""
        text_tokens = self.text_to_tokens(text)
        text_emb = self.model.text_encoder(text_tokens)

        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
        all_embeddings.append(next_emb)

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

        for step in range(1, max_steps):
            next_emb, cache = self.model.predict_next_cached(next_emb, text_emb, step, cache)
            if temperature > 0:
                next_emb = next_emb + torch.randn_like(next_emb) * temperature
            all_embeddings.append(next_emb)

            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

        # Concat predicted embeddings and project to codec space
        pred_embs = torch.cat(all_embeddings, dim=1)  # [1, T, d_model]
        codec_emb = self.model.proj_out(pred_embs)     # [1, T, 128]
        codec_emb = codec_emb.transpose(1, 2)          # [1, 128, T]

        # EnCodec decode
        waveform = self.codec.decoder(codec_emb)  # [1, 1, T_audio]
        waveform = waveform.squeeze().cpu().numpy()

        return waveform, 24000

    @torch.no_grad()
    def reconstruct(self, audio_path):
        """Encode audio → proj_in → proj_out → decode. Tests roundtrip quality."""
        from encodec.utils import convert_audio
        wav, sr = torchaudio.load(audio_path)
        wav = convert_audio(wav, sr, 24000, 1).to(self.device)

        # EnCodec encode
        codec_emb = self.codec.encoder(wav.unsqueeze(0))  # [1, 128, T]

        # proj_in → proj_out roundtrip
        z = self.model.proj_in(codec_emb.transpose(1, 2))  # [1, T, d_model]
        codec_recon = self.model.proj_out(z).transpose(1, 2)  # [1, 128, T]

        # EnCodec decode
        waveform = self.codec.decoder(codec_recon).squeeze().cpu().numpy()
        return waveform, 24000

    @torch.no_grad()
    def reconstruct_codec_only(self, audio_path):
        """Pure EnCodec roundtrip (no JEPA). Baseline quality."""
        from encodec.utils import convert_audio
        wav, sr = torchaudio.load(audio_path)
        wav = convert_audio(wav, sr, 24000, 1).to(self.device)

        codec_emb = self.codec.encoder(wav.unsqueeze(0))
        waveform = self.codec.decoder(codec_emb).squeeze().cpu().numpy()
        return waveform, 24000

    @torch.no_grad()
    def reconstruct_with_prediction(self, audio_path, text):
        """Encode audio → downsample → teacher-forced predict → upsample → decode."""
        from encodec.utils import convert_audio
        wav, sr = torchaudio.load(audio_path)
        wav = convert_audio(wav, sr, 24000, 1).to(self.device)

        codec_emb = self.codec.encoder(wav.unsqueeze(0))
        z, _ = self.model.downsampler(codec_emb)

        text_tokens = self.text_to_tokens(text)
        text_emb = self.model.text_encoder(text_tokens)

        B = z.shape[0]
        start = self.model.start_emb.expand(B, -1, -1)
        input_emb = torch.cat([start, z[:, :-1]], dim=1)
        predicted = self.model.predictor(input_emb, text_emb)

        codec_recon = self.model.upsampler(predicted, target_len=codec_emb.shape[2])
        waveform = self.codec.decoder(codec_recon).squeeze().cpu().numpy()
        return waveform, 24000

    def save_audio(self, waveform, sr, output_path):
        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)
    parser.add_argument("--audio", default=None, help="Audio path for reconstruction")
    parser.add_argument("--mode", choices=["synthesize", "reconstruct", "codec_only", "predict"],
                        default="synthesize")
    parser.add_argument("--output", default="output_v5.wav")
    parser.add_argument("--max_steps", type=int, default=300)
    parser.add_argument("--temperature", type=float, default=0.0)
    args = parser.parse_args()

    tts = LeWMTTSv5Inference(args.checkpoint)

    if args.mode == "synthesize":
        assert args.text, "Need --text"
        waveform, sr = tts.synthesize(args.text, args.max_steps, args.temperature)
    elif args.mode == "reconstruct":
        assert args.audio, "Need --audio"
        waveform, sr = tts.reconstruct(args.audio)
    elif args.mode == "codec_only":
        assert args.audio, "Need --audio"
        waveform, sr = tts.reconstruct_codec_only(args.audio)
    elif args.mode == "predict":
        assert args.audio and args.text, "Need --audio and --text"
        waveform, sr = tts.reconstruct_with_prediction(args.audio, args.text)

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

if __name__ == "__main__":
    main()

-        """Encode audio → downsample → teacher-forced predict → upsample → decode."""

+        """Encode audio → proj_in → teacher-forced predict → proj_out → decode."""

-        codec_emb = self.codec.encoder(wav.unsqueeze(0))

+        codec_emb = self.codec.encoder(wav.unsqueeze(0))  # [1, 128, T]

-        codec_recon = self.model.upsampler(predicted, target_len=codec_emb.shape[2])

+        codec_recon = self.model.proj_out(predicted).transpose(1, 2)  # [1, 128, T]


Also remove `ds_stride` from train_v5.py since we don't need it anymore:


    25→        "text_vocab_size": 256,
    26→        "text_encoder_layers": args.text_layers,
    27→        "predictor_layers": args.predictor_layers,
    28→        "dropout": args.dropout,
    29→        "ds_stride": args.ds_stride,
    30→        "pred_weight": args.pred_weight,
    31→        "roundtrip_weight": args.roundtrip_weight,
    32→        "ema_decay": args.ema_decay,
    33→        "input_noise": 0.0,
    34→    }
    35→    model, config = build_model(config)
    36→    model = model.to(device)
    37→
    38→    manifest_path = os.path.join(args.data_dir, "manifest.json")
    39→    loader, dataset = build_dataloader(
    40→        manifest_path, batch_size=args.batch_size,
    41→        num_workers=args.num_workers, max_codec_frames=args.max_codec_frames,
    42→    )
    43→    print(f"Training: {len(dataset)} samples, {len(loader)} batches/epoch")
    44→

        "text_vocab_size": 256,
        "text_encoder_layers": args.text_layers,
        "predictor_layers": args.predictor_layers,
        "dropout": args.dropout,
        "ds_stride": args.ds_stride,
        "pred_weight": args.pred_weight,
        "roundtrip_weight": args.roundtrip_weight,
        "ema_decay": args.ema_decay,
        "input_noise": 0.0,
    }
    model, config = build_model(config)
    model = model.to(device)

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


        "ds_stride": args.ds_stride,
        "pred_weight": args.pred_weight,


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

"""
Training loop for LeWM TTS v5 (codec-based JEPA).
"""

import os
import json
import time
import argparse
import torch
from torch.utils.tensorboard import SummaryWriter
from pathlib import Path

from model_v5 import build_model
from dataset_v5 import build_dataloader

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

    config = {
        "d_model": args.d_model,
        "nhead": args.nhead,
        "codec_dim": 128,
        "text_vocab_size": 256,
        "text_encoder_layers": args.text_layers,
        "predictor_layers": args.predictor_layers,
        "dropout": args.dropout,
        "ds_stride": args.ds_stride,
        "pred_weight": args.pred_weight,
        "roundtrip_weight": args.roundtrip_weight,
        "ema_decay": args.ema_decay,
        "input_noise": 0.0,
    }
    model, config = build_model(config)
    model = model.to(device)

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

    optimizer = torch.optim.AdamW(
        model.parameters(), lr=args.lr,
        betas=(0.9, 0.98), weight_decay=0.01,
    )

    total_steps = args.epochs * len(loader)
    warmup_steps = min(4000, 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)

    start_epoch = 0
    global_step = 0
    if args.resume:
        ckpt = torch.load(args.resume, map_location=device, weights_only=False)
        model.load_state_dict(ckpt["model"])
        optimizer.load_state_dict(ckpt["optimizer"])
        scheduler.load_state_dict(ckpt["scheduler"])
        start_epoch = ckpt["epoch"] + 1
        global_step = ckpt["global_step"]
        print(f"Resumed from epoch {start_epoch}, step {global_step}")

    model.train()
    best_loss = float("inf")

    for epoch in range(start_epoch, args.epochs):
        epoch_loss = 0.0
        epoch_pred = 0.0
        epoch_rt = 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)

            # Input noise annealing
            noise_progress = min(1.0, global_step / (total_steps * 0.3))
            model.input_noise = args.input_noise_max * noise_progress

            losses = model(codec_emb, text_tokens, codec_mask, text_mask)

            total_loss = losses["total_loss"]
            pred_loss = losses["prediction_loss"]
            rt_loss = losses["roundtrip_loss"]

            optimizer.zero_grad()
            total_loss.backward()
            grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
            optimizer.step()
            scheduler.step()
            model.update_ema()

            epoch_loss += total_loss.item()
            epoch_pred += pred_loss.item()
            epoch_rt += rt_loss.item()
            global_step += 1

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

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

        n_batches = len(loader)
        avg_loss = epoch_loss / n_batches
        avg_pred = epoch_pred / n_batches
        avg_rt = epoch_rt / n_batches
        epoch_time = time.time() - t0

        print(
            f"Epoch {epoch+1}/{args.epochs} | "
            f"avg_loss={avg_loss:.4f} pred={avg_pred:.4f} rt={avg_rt:.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 checkpoint: {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 (loss={avg_loss:.4f})")

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

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--data_dir", default="processed_data_codec")
    parser.add_argument("--output_dir", default="output_v5")
    parser.add_argument("--d_model", type=int, default=256)
    parser.add_argument("--nhead", type=int, default=4)
    parser.add_argument("--text_layers", type=int, default=4)
    parser.add_argument("--predictor_layers", type=int, default=6)
    parser.add_argument("--dropout", type=float, default=0.1)
    parser.add_argument("--ds_stride", type=int, default=3)
    parser.add_argument("--pred_weight", type=float, default=10.0)
    parser.add_argument("--roundtrip_weight", type=float, default=1.0)
    parser.add_argument("--ema_decay", type=float, default=0.998)
    parser.add_argument("--input_noise_max", type=float, default=0.1)
    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=64)
    parser.add_argument("--lr", type=float, default=1e-4)
    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=50)
    parser.add_argument("--save_every", type=int, default=10)
    parser.add_argument("--resume", default=None)
    args = parser.parse_args()
    train(args)


Also update the build_model defaults:


        config = {
            "d_model": 256, "nhead": 4, "codec_dim": 128,
            "text_vocab_size": 256, "text_encoder_layers": 4,
            "predictor_layers": 6, "dropout": 0.1, "ds_stride": 3,
            "pred_weight": 10.0, "roundtrip_weight": 1.0,
            "ema_decay": 0.998, "input_noise": 0.0,
        }

        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,
        }


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

        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,
        }

"""
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
  - JEPAPredictor: causal transformer, predicts next embedding
  - EMA target on projected embeddings
"""

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

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

class JEPAPredictor(nn.Module):
    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)

        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)
        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):
        x = self.audio_input_proj(audio_emb)
        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, text_emb, tgt_mask=causal_mask,
            tgt_key_padding_mask=audio_mask, memory_key_padding_mask=text_mask,
        )
        return self.output_proj(predicted)

    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 _cached_cross_attn(self, mha, q_input, memory, cache_k, cache_v, memory_mask=None):
        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)
        if cache_k is None:
            k = F.linear(memory, Wk, bk)
            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:
            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):
        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):
        x = self.audio_input_proj(new_emb)
        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))
            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))
            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))

    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() * 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, "ds_stride": 3,
            "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

-            "predictor_layers": 6, "dropout": 0.1, "ds_stride": 3,

+            "predictor_layers": 6, "dropout": 0.1,


Now test and find optimal batch size for the simplified model:


Bash

cd /home/ubuntu/lewm-tts && python3 -c "
import torch
from model_v5 import build_model
from dataset_v5 import build_dataloader

device = torch.device('cuda')

for bs in [32, 64, 128]:
    torch.cuda.reset_peak_memory_stats()
    try:
        model, config = build_model()
        model = model.to(device)
        model.train()
        loader, _ = build_dataloader('processed_data_codec/manifest.json', batch_size=bs, max_codec_frames=900)
        batch = next(iter(loader))
        losses = model(batch['codec_emb'].to(device), batch['text_tokens'].to(device),
                       batch['codec_mask'].to(device), batch['text_mask'].to(device))
        losses['total_loss'].backward()
        mem = torch.cuda.max_memory_allocated()/1e9
        print(f'batch_size={bs}: GPU={mem:.1f}GB, pred={losses[\"prediction_loss\"].item():.4f}, rt={losses[\"roundtrip_loss\"].item():.4f}')
        del model, loader, batch, losses
        torch.cuda.empty_cache()
    except torch.cuda.OutOfMemoryError:
        print(f'batch_size={bs}: OOM')
        torch.cuda.empty_cache()
        break
" 2>&1 | grep -E "batch_size|OOM|parameters"

Find batch size for simplified v5 (75Hz, no downsampler)


LeWM TTS v5: 10.01M trainable parameters
Shell cwd was reset to /home/ubuntu/pipecat


It seems like only bs=32 printed, then nothing. The 75Hz sequences are much longer (avg 689 frames vs ~230 before), so self-attention in the predictor is O(T²) which uses a lot more memory. Let me try smaller batch sizes or limit max_codec_frames.


The 75Hz sequences are ~690 frames long — the causal attention is O(T²) so memory explodes. Let me limit sequence length or use smaller batches:


cd /home/ubuntu/lewm-tts && python3 -c "
import torch
from model_v5 import build_model
from dataset_v5 import build_dataloader

device = torch.device('cuda')

# Try with shorter max frames
for max_frames, bs in [(300, 128), (300, 256), (450, 64), (450, 128), (600, 32), (600, 64)]:
    torch.cuda.reset_peak_memory_stats()
    try:
        model, config = build_model()
        model = model.to(device).train()
        loader, ds = build_dataloader('processed_data_codec/manifest.json', batch_size=bs, max_codec_frames=max_frames)
        batch = next(iter(loader))
        losses = model(batch['codec_emb'].to(device), batch['text_tokens'].to(device),
                       batch['codec_mask'].to(device), batch['text_mask'].to(device))
        losses['total_loss'].backward()
        mem = torch.cuda.max_memory_allocated()/1e9
        print(f'max_frames={max_frames} bs={bs}: GPU={mem:.1f}GB, samples={len(ds)}, batches={len(loader)}')
        del model, loader, ds, batch, losses
        torch.cuda.empty_cache()
    except torch.cuda.OutOfMemoryError:
        print(f'max_frames={max_frames} bs={bs}: OOM')
        del model
        torch.cuda.empty_cache()
" 2>&1 | grep -E "max_frames|OOM"

Find optimal max_frames and batch_size combo


(Bash completed with no output)


cd /home/ubuntu/lewm-tts && python3 -c "
import torch
from model_v5 import build_model
from dataset_v5 import build_dataloader
device = torch.device('cuda')

# Start small
for max_frames, bs in [(300, 32), (300, 64), (300, 128)]:
    torch.cuda.reset_peak_memory_stats()
    model, _ = build_model()
    model = model.to(device).train()
    loader, ds = build_dataloader('processed_data_codec/manifest.json', batch_size=bs, max_codec_frames=max_frames)
    batch = next(iter(loader))
    try:
        losses = model(batch['codec_emb'].to(device), batch['text_tokens'].to(device),
                       batch['codec_mask'].to(device), batch['text_mask'].to(device))
        losses['total_loss'].backward()
        mem = torch.cuda.max_memory_allocated()/1e9
        print(f'max_frames={max_frames} bs={bs}: GPU={mem:.1f}GB, samples={len(ds)}')
    except Exception as e:
        print(f'max_frames={max_frames} bs={bs}: ERROR {e}')
    del model, loader, ds, batch
    torch.cuda.empty_cache()
" 2>&1 | grep -E "max_frames|ERROR"

Test smaller configs


max_frames=300 bs=32: ERROR name 'codec_dim' is not defined
max_frames=300 bs=64: ERROR name 'codec_dim' is not defined
max_frames=300 bs=128: ERROR name 'codec_dim' is not defined
Shell cwd was reset to /home/ubuntu/pipecat


Bug in the model — `codec_dim` variable reference in forward. Let me fix it:


Edit

) / (valid_rt.sum() * self.codec_dim + 1e-8)


"""
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
  - JEPAPredictor: causal transformer, predicts next embedding
  - EMA target on projected embeddings
"""

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

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

class JEPAPredictor(nn.Module):
    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)

        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)
        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):
        x = self.audio_input_proj(audio_emb)
        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, text_emb, tgt_mask=causal_mask,
            tgt_key_padding_mask=audio_mask, memory_key_padding_mask=text_mask,
        )
        return self.output_proj(predicted)

    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 _cached_cross_attn(self, mha, q_input, memory, cache_k, cache_v, memory_mask=None):
        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)
        if cache_k is None:
            k = F.linear(memory, Wk, bk)
            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:
            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):
        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):
        x = self.audio_input_proj(new_emb)
        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))
            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))
            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))

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

-            ) / (valid_rt.sum() * codec_dim + 1e-8)

+            ) / (valid_rt.sum() * self.codec_dim + 1e-8)


Bash

cd /home/ubuntu/lewm-tts && python3 -c "
import torch
from model_v5 import build_model
from dataset_v5 import build_dataloader
device = torch.device('cuda')

for max_frames, bs in [(300, 64), (300, 128), (300, 256), (450, 64), (450, 128), (600, 64)]:
    torch.cuda.reset_peak_memory_stats()
    model, _ = build_model()
    model = model.to(device).train()
    loader, ds = build_dataloader('processed_data_codec/manifest.json', batch_size=bs, max_codec_frames=max_frames)
    batch = next(iter(loader))
    try:
        losses = model(batch['codec_emb'].to(device), batch['text_tokens'].to(device),
                       batch['codec_mask'].to(device), batch['text_mask'].to(device))
        losses['total_loss'].backward()
        mem = torch.cuda.max_memory_allocated()/1e9
        print(f'max_frames={max_frames} bs={bs}: GPU={mem:.1f}GB, samples={len(ds)}, batches={len(loader)}')
    except torch.cuda.OutOfMemoryError:
        print(f'max_frames={max_frames} bs={bs}: OOM')
    del model, loader, ds, batch
    torch.cuda.empty_cache()
" 2>&1 | grep -E "max_frames|OOM"

Find optimal config


max_frames=300 bs=64: GPU=4.6GB, samples=141, batches=2
max_frames=300 bs=128: GPU=9.5GB, samples=141, batches=1
Shell cwd was reset to /home/ubuntu/pipecat