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

Architecture:
  - TextEncoder: Character-level transformer for Hindi/Devanagari text
  - AudioEncoder: 1D CNN + Transformer that encodes mel spectrograms to embeddings
  - JEPAPredictor: Transformer that predicts next audio embedding given text + audio context
  - Two losses: next-embedding prediction (MSE) + Gaussian regularizer (KL to N(0,I))
"""

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

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

class TextEncoder(nn.Module):
    """Character-level transformer encoder for Hindi text."""

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

        # Character embedding (covers all Unicode Hindi chars via byte-level encoding)
        self.char_embed = nn.Embedding(vocab_size, d_model)
        self.pos_embed = SinusoidalPositionalEncoding(d_model, max_len=2048)
        self.dropout = nn.Dropout(dropout)

        encoder_layer = nn.TransformerEncoderLayer(
            d_model=d_model, nhead=nhead,
            dim_feedforward=d_model * 4,
            dropout=dropout, batch_first=True,
            activation="gelu",
        )
        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
        self.proj = nn.Linear(d_model, d_model)

    def forward(self, text_tokens, text_mask=None):
        """
        Args:
            text_tokens: [B, T_text] long tensor of byte-level token IDs
            text_mask: [B, T_text] bool, True = padding
        Returns:
            text_emb: [B, T_text, d_model]
        """
        x = self.char_embed(text_tokens) * math.sqrt(self.d_model)
        x = self.pos_embed(x)
        x = self.dropout(x)
        x = self.transformer(x, src_key_padding_mask=text_mask)
        x = self.proj(x)
        return x

# ─── Audio Encoder ──────────────────────────────────────────────────────────

class AudioEncoder(nn.Module):
    """
    Encodes mel spectrograms into a sequence of audio embeddings.
    Uses 1D CNN for downsampling + Transformer for contextualization.
    """

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

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

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

        encoder_layer = nn.TransformerEncoderLayer(
            d_model=d_model, nhead=nhead,
            dim_feedforward=d_model * 4,
            dropout=dropout, batch_first=True,
            activation="gelu",
        )
        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)

        # Project to embedding with mu and logvar for Gaussian regularizer
        self.proj_mu = nn.Linear(d_model, d_model)
        self.proj_logvar = nn.Linear(d_model, d_model)

    def forward(self, mel, mel_mask=None):
        """
        Args:
            mel: [B, n_mels, T_mel] — log mel spectrogram
            mel_mask: [B, T_mel] bool, True = padding (before downsampling)
        Returns:
            z: [B, T_down, d_model] — sampled latent embeddings
            mu: [B, T_down, d_model]
            logvar: [B, T_down, d_model]
        """
        # CNN encoding: [B, n_mels, T] -> [B, d_model, T//4]
        x = self.conv_pre(mel)
        x = x.transpose(1, 2)  # [B, T_down, d_model]

        # Downsample mask if provided
        if mel_mask is not None:
            # Approximate downsampled mask
            T_down = x.shape[1]
            mel_mask = mel_mask[:, ::self.downsample_factor][:, :T_down]

        x = self.pos_embed(x)
        x = self.dropout(x)
        x = self.transformer(x, src_key_padding_mask=mel_mask)

        mu = self.proj_mu(x)
        logvar = self.proj_logvar(x)

        # Reparameterization trick
        if self.training:
            std = torch.exp(0.5 * logvar)
            eps = torch.randn_like(std)
            z = mu + eps * std
        else:
            z = mu

        return z, mu, logvar

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

class JEPAPredictor(nn.Module):
    """
    Predicts next audio embedding given:
    - Text context (from TextEncoder)
    - Previous audio embeddings (from AudioEncoder)

    Uses cross-attention to condition on text, and causal self-attention
    over the audio embedding sequence.
    """

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

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

        # Transformer decoder with causal self-attention + cross-attention to text
        decoder_layer = nn.TransformerDecoderLayer(
            d_model=d_model, nhead=nhead,
            dim_feedforward=d_model * 4,
            dropout=dropout, batch_first=True,
            activation="gelu",
        )
        self.transformer = nn.TransformerDecoder(decoder_layer, num_layers=num_layers)

        # Output projection to predict next embedding
        self.output_proj = nn.Sequential(
            nn.Linear(d_model, d_model),
            nn.GELU(),
            nn.Linear(d_model, d_model),
        )

    def forward(self, audio_emb, text_emb, audio_mask=None, text_mask=None):
        """
        Args:
            audio_emb: [B, T_audio, d_model] — audio embeddings (from encoder)
            text_emb: [B, T_text, d_model] — text embeddings
            audio_mask: [B, T_audio] bool, True = padding
            text_mask: [B, T_text] bool, True = padding
        Returns:
            predicted: [B, T_audio, d_model] — predicted next embeddings
        """
        x = self.audio_input_proj(audio_emb)
        x = self.pos_embed(x)
        x = self.dropout(x)

        # Causal mask: each position can only attend to previous positions
        T = x.shape[1]
        causal_mask = torch.triu(
            torch.ones(T, T, device=x.device, dtype=torch.bool), diagonal=1
        )

        predicted = self.transformer(
            x, text_emb,
            tgt_mask=causal_mask,
            tgt_key_padding_mask=audio_mask,
            memory_key_padding_mask=text_mask,
        )

        predicted = self.output_proj(predicted)
        return predicted

    def _cached_mha(self, mha, q_input, kv_input, cache_k, cache_v):
        """Run multi-head attention with KV cache.
        q_input: [B, 1, d], kv_input: [B, 1, d] (new token for K,V)
        cache_k, cache_v: [B, T_prev, d] projected keys/values
        Returns: output [B, 1, d], updated cache_k, updated cache_v
        """
        d = mha.embed_dim
        nhead = mha.num_heads
        head_dim = d // nhead
        B = q_input.shape[0]

        # Project Q from q_input, K/V from kv_input
        # in_proj_weight is [3d, d] = [Wq; Wk; Wv]
        Wq, Wk, Wv = mha.in_proj_weight.chunk(3, dim=0)
        bq, bk, bv = mha.in_proj_bias.chunk(3, dim=0)

        q = F.linear(q_input, Wq, bq)        # [B, 1, d]
        k_new = F.linear(kv_input, Wk, bk)    # [B, 1, d]
        v_new = F.linear(kv_input, Wv, bv)    # [B, 1, d]

        # Append to cache
        if cache_k is not None and cache_k.shape[1] > 0:
            k = torch.cat([cache_k, k_new], dim=1)
            v = torch.cat([cache_v, v_new], dim=1)
        else:
            k = k_new
            v = v_new

        # Multi-head attention
        q = q.view(B, 1, nhead, head_dim).transpose(1, 2)      # [B, nhead, 1, head_dim]
        k_mh = k.view(B, -1, nhead, head_dim).transpose(1, 2)  # [B, nhead, T, head_dim]
        v_mh = v.view(B, -1, nhead, head_dim).transpose(1, 2)

        attn = torch.matmul(q, k_mh.transpose(-2, -1)) / (head_dim ** 0.5)
        attn = F.softmax(attn, dim=-1)
        out = torch.matmul(attn, v_mh)  # [B, nhead, 1, head_dim]

        out = out.transpose(1, 2).contiguous().view(B, 1, d)
        out = F.linear(out, mha.out_proj.weight, mha.out_proj.bias)

        return out, k, v

    def _cached_cross_attn(self, mha, q_input, memory, cache_k, cache_v, memory_mask=None):
        """Cross-attention with cached memory K,V (computed once).
        memory_mask: [B, T_text] bool, True = padding (applied as -inf)."""
        d = mha.embed_dim
        nhead = mha.num_heads
        head_dim = d // nhead
        B = q_input.shape[0]

        Wq, Wk, Wv = mha.in_proj_weight.chunk(3, dim=0)
        bq, bk, bv = mha.in_proj_bias.chunk(3, dim=0)

        q = F.linear(q_input, Wq, bq)  # [B, 1, d]

        if cache_k is None:
            k = F.linear(memory, Wk, bk)  # [B, T_text, d]
            v = F.linear(memory, Wv, bv)
        else:
            k, v = cache_k, cache_v

        q = q.view(B, 1, nhead, head_dim).transpose(1, 2)
        k_mh = k.view(B, -1, nhead, head_dim).transpose(1, 2)
        v_mh = v.view(B, -1, nhead, head_dim).transpose(1, 2)

        attn = torch.matmul(q, k_mh.transpose(-2, -1)) / (head_dim ** 0.5)
        if memory_mask is not None:
            # memory_mask: [B, T_text] → [B, 1, 1, T_text]
            attn = attn.masked_fill(memory_mask.unsqueeze(1).unsqueeze(2), float('-inf'))
        attn = F.softmax(attn, dim=-1)
        out = torch.matmul(attn, v_mh)

        out = out.transpose(1, 2).contiguous().view(B, 1, d)
        out = F.linear(out, mha.out_proj.weight, mha.out_proj.bias)

        return out, k, v

    def init_cache(self, num_layers):
        """Initialize empty KV cache."""
        return {
            'self_k': [None] * num_layers,
            'self_v': [None] * num_layers,
            'cross_k': [None] * num_layers,
            'cross_v': [None] * num_layers,
        }

    def inference_step(self, new_emb, text_emb, step_idx, cache, text_mask=None):
        """Process one token with KV cache.
        new_emb: [B, 1, d_model] — raw audio embedding (not projected)
        text_mask: [B, T_text] bool, True = padding
        Returns: predicted [B, 1, d_model], updated cache
        """
        x = self.audio_input_proj(new_emb)
        # Add positional encoding for this specific position
        x = x + self.pos_embed.pe[:, step_idx:step_idx + 1]

        for i, layer in enumerate(self.transformer.layers):
            # Self-attention with cache (post-norm)
            sa_out, cache['self_k'][i], cache['self_v'][i] = self._cached_mha(
                layer.self_attn, x, x, cache['self_k'][i], cache['self_v'][i]
            )
            x = layer.norm1(x + layer.dropout1(sa_out))

            # Cross-attention with cached text KV
            ca_out, cache['cross_k'][i], cache['cross_v'][i] = self._cached_cross_attn(
                layer.multihead_attn, x, text_emb, cache['cross_k'][i], cache['cross_v'][i],
                memory_mask=text_mask,
            )
            x = layer.norm2(x + layer.dropout2(ca_out))

            # FFN
            ff_out = layer.linear2(layer.dropout(layer.activation(layer.linear1(x))))
            x = layer.norm3(x + layer.dropout3(ff_out))

        return self.output_proj(x), cache

# ─── Positional Encoding ────────────────────────────────────────────────────

class SinusoidalPositionalEncoding(nn.Module):
    def __init__(self, d_model, max_len=8192):
        super().__init__()
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        self.register_buffer("pe", pe.unsqueeze(0))  # [1, max_len, d_model]

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

# ─── Mel Decoder ────────────────────────────────────────────────────────────

class MelDecoder(nn.Module):
    """Decodes latent embeddings back to mel spectrograms (reverses AudioEncoder)."""

    def __init__(self, d_model=256, n_mels=100):
        super().__init__()
        self.proj = nn.Sequential(
            nn.Linear(d_model, d_model),
            nn.GELU(),
        )
        self.upsample = nn.Sequential(
            nn.ConvTranspose1d(d_model, d_model, kernel_size=4, stride=2, padding=1),
            nn.GELU(),
            nn.ConvTranspose1d(d_model, d_model, kernel_size=4, stride=2, padding=1),
            nn.GELU(),
            nn.Conv1d(d_model, n_mels, kernel_size=7, padding=3),
        )

    def forward(self, z):
        """z: [B, T_down, d_model] → mel: [B, n_mels, T_up]"""
        x = self.proj(z)
        x = x.transpose(1, 2)
        return self.upsample(x)

# ─── Full LeWM TTS Model ────────────────────────────────────────────────────

class LeWMTTS(nn.Module):
    """
    Complete LeWM TTS model.

    Training: text + mel → audio embeddings → predict next embedding
    Two losses:
        1. Next-embedding prediction: MSE(predicted[t], target[t+1])
        2. Gaussian regularizer: KL(q(z|x) || N(0, I))
    """

    def __init__(self, config):
        super().__init__()
        d_model = config.get("d_model", 256)
        nhead = config.get("nhead", 4)
        n_mels = config.get("n_mels", 100)
        text_vocab = config.get("text_vocab_size", 256)
        text_layers = config.get("text_encoder_layers", 4)
        audio_layers = config.get("audio_encoder_layers", 4)
        predictor_layers = config.get("predictor_layers", 6)
        dropout = config.get("dropout", 0.1)
        self.kl_weight = config.get("kl_weight", 0.01)
        n_speakers = config.get("n_speakers", 1)

        self.text_encoder = TextEncoder(
            vocab_size=text_vocab, d_model=d_model, nhead=nhead,
            num_layers=text_layers, dropout=dropout,
        )
        self.audio_encoder = AudioEncoder(
            n_mels=n_mels, d_model=d_model, nhead=nhead,
            num_layers=audio_layers, dropout=dropout,
        )
        self.predictor = JEPAPredictor(
            d_model=d_model, nhead=nhead,
            num_layers=predictor_layers, dropout=dropout,
        )

        # Mel decoder: reconstructs mel from latent embeddings
        # Needed so the latent space is grounded to actual audio
        self.mel_decoder = MelDecoder(d_model=d_model, n_mels=n_mels)
        self.recon_weight = config.get("recon_weight", 1.0)

        # Speaker embedding (for multi-speaker)
        self.n_speakers = n_speakers
        if n_speakers > 1:
            self.speaker_embed = nn.Embedding(n_speakers, d_model)
        else:
            self.speaker_embed = None

        self.config = config

    def _add_speaker(self, text_emb, speaker_id):
        """Add speaker embedding to text embeddings."""
        if self.speaker_embed is not None and speaker_id is not None:
            spk = self.speaker_embed(speaker_id).unsqueeze(1)  # [B, 1, d]
            text_emb = text_emb + spk
        return text_emb

    def forward(self, mel, text_tokens, mel_mask=None, text_mask=None, speaker_id=None):
        """
        Args:
            mel: [B, n_mels, T_mel]
            text_tokens: [B, T_text]
            mel_mask: [B, T_mel] bool
            text_mask: [B, T_text] bool
        Returns:
            loss_dict with 'prediction_loss', 'kl_loss', 'total_loss'
        """
        # Encode text + speaker conditioning
        text_emb = self.text_encoder(text_tokens, text_mask)
        text_emb = self._add_speaker(text_emb, speaker_id)

        # Encode audio → latent embeddings
        z, mu, logvar = self.audio_encoder(mel, mel_mask)

        # Predict next embedding: input z[:-1], target z[1:]
        input_emb = z[:, :-1]     # [B, T-1, d]
        target_emb = z[:, 1:]     # [B, T-1, d]

        # Adjust mask for downsampled sequence
        if mel_mask is not None:
            T_down = z.shape[1]
            ds_mask = mel_mask[:, ::self.audio_encoder.downsample_factor][:, :T_down]
            pred_mask = ds_mask[:, :-1]
        else:
            pred_mask = None

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

        # ─── Loss 1: Next-embedding prediction (MSE) ───
        # Mask out padding positions
        if pred_mask is not None:
            valid = ~pred_mask  # True = valid
            valid = valid.unsqueeze(-1)  # [B, T-1, 1]
            prediction_loss = F.mse_loss(
                predicted * valid, target_emb * valid, reduction="sum"
            ) / (valid.sum() * predicted.shape[-1] + 1e-8)
        else:
            prediction_loss = F.mse_loss(predicted, target_emb)

        # ─── Loss 2: Gaussian regularizer — KL(N(mu, sigma²) || N(0, I)) ───
        # KL = -0.5 * sum(1 + logvar - mu² - exp(logvar))
        if mel_mask is not None:
            T_down = mu.shape[1]
            ds_mask = mel_mask[:, ::self.audio_encoder.downsample_factor][:, :T_down]
            valid = (~ds_mask).unsqueeze(-1).float()
            kl = -0.5 * (1 + logvar - mu.pow(2) - logvar.exp())
            kl_loss = (kl * valid).sum() / (valid.sum() * mu.shape[-1] + 1e-8)
        else:
            kl_loss = -0.5 * torch.mean(1 + logvar - mu.pow(2) - logvar.exp())

        # ─── Loss 3: Mel reconstruction ───
        # Decode latents back to mel and compare with input
        mel_recon = self.mel_decoder(z)  # [B, n_mels, T_up]
        # Trim/pad to match input mel length
        T_mel = mel.shape[2]
        T_recon = mel_recon.shape[2]
        if T_recon > T_mel:
            mel_recon = mel_recon[:, :, :T_mel]
        elif T_recon < T_mel:
            mel = mel[:, :, :T_recon]
            if mel_mask is not None:
                mel_mask = mel_mask[:, :T_recon]

        if mel_mask is not None:
            valid_mel = (~mel_mask).unsqueeze(1).float()  # [B, 1, T]
            recon_loss = F.l1_loss(
                mel_recon * valid_mel, mel * valid_mel, reduction="sum"
            ) / (valid_mel.sum() * mel.shape[1] + 1e-8)
        else:
            recon_loss = F.l1_loss(mel_recon, mel)

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

        return {
            "total_loss": total_loss,
            "prediction_loss": prediction_loss,
            "kl_loss": kl_loss,
            "recon_loss": recon_loss,
        }

    def encode_audio(self, mel):
        """Encode mel to latent embeddings (for inference)."""
        z, mu, _ = self.audio_encoder(mel)
        return mu  # Use mean at inference

    def predict_next(self, audio_emb, text_emb, text_mask=None):
        """Predict next audio embedding autoregressively."""
        predicted = self.predictor(audio_emb, text_emb, text_mask=text_mask)
        return predicted[:, -1:]  # Only the last prediction

    def init_ar_cache(self):
        """Initialize KV cache for autoregressive inference."""
        num_layers = len(self.predictor.transformer.layers)
        return self.predictor.init_cache(num_layers)

    def predict_next_cached(self, new_emb, text_emb, step_idx, cache, text_mask=None):
        """Predict next embedding using KV cache. O(1) per step instead of O(n)."""
        return self.predictor.inference_step(new_emb, text_emb, step_idx, cache, text_mask=text_mask)

def count_parameters(model):
    return sum(p.numel() for p in model.parameters() if p.requires_grad)

def build_model(config=None):
    if config is None:
        config = {
            "d_model": 256,
            "nhead": 4,
            "n_mels": 80,
            "text_vocab_size": 256,
            "text_encoder_layers": 4,
            "audio_encoder_layers": 4,
            "predictor_layers": 6,
            "dropout": 0.1,
            "kl_weight": 0.01,
        }
    model = LeWMTTS(config)
    print(f"LeWM TTS model: {count_parameters(model)/1e6:.2f}M parameters")
    return model, config

config = {
"d_model": 256,
"nhead": 4,
-            "n_mels": 80,
+            "n_mels": 100,
"text_vocab_size": 256,
"text_encoder_layers": 4,
"audio_encoder_layers": 4,


Updated task #5 status


Updated task #3 status


Update train.py:


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

"""
Training loop for LeWM TTS.
"""

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

from model import build_model
from dataset import build_dataloader

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

    # ─── Build model ───
    config = {
        "d_model": args.d_model,
        "nhead": args.nhead,
        "n_mels": 80,
        "text_vocab_size": 256,
        "text_encoder_layers": args.text_layers,
        "audio_encoder_layers": args.audio_layers,
        "predictor_layers": args.predictor_layers,
        "dropout": args.dropout,
        "kl_weight": args.kl_weight,
        "recon_weight": args.recon_weight,
    }
    model, config = build_model(config)
    model = model.to(device)

    # ─── Build dataloader ───
    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,
    )
    print(f"Training: {len(dataset)} samples, {len(loader)} batches/epoch")

    # ─── Optimizer & Scheduler ───
    optimizer = torch.optim.AdamW(
        model.parameters(), lr=args.lr,
        betas=(0.9, 0.98), weight_decay=0.01,
    )

    # Warmup + cosine decay
    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)

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

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

    # ─── Resume from checkpoint ───
    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}")

    # ─── Training loop ───
    model.train()
    best_loss = float("inf")

    for epoch in range(start_epoch, args.epochs):
        epoch_loss = 0.0
        epoch_pred_loss = 0.0
        epoch_kl_loss = 0.0
        epoch_recon_loss = 0.0
        t0 = time.time()

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

            # KL annealing: ramp up over first 20% of training
            kl_anneal = min(1.0, global_step / (total_steps * 0.2))
            model.kl_weight = args.kl_weight * kl_anneal

            # Forward
            losses = model(mel, text_tokens, mel_mask, text_mask)

            total_loss = losses["total_loss"]
            pred_loss = losses["prediction_loss"]
            kl_loss = losses["kl_loss"]
            recon_loss = losses["recon_loss"]

            # Backward
            optimizer.zero_grad()
            total_loss.backward()

            # Gradient clipping
            grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)

            optimizer.step()
            scheduler.step()

            # Logging
            epoch_loss += total_loss.item()
            epoch_pred_loss += pred_loss.item()
            epoch_kl_loss += kl_loss.item()
            epoch_recon_loss += recon_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/kl_loss", kl_loss.item(), global_step)
                writer.add_scalar("train/recon_loss", recon_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"kl={kl_loss.item():.4f} recon={recon_loss.item():.4f} | "
                    f"grad={grad_norm.item():.2f} | {elapsed:.1f}s"
                )

        # Epoch summary
        n_batches = len(loader)
        avg_loss = epoch_loss / n_batches
        avg_pred = epoch_pred_loss / n_batches
        avg_kl = epoch_kl_loss / n_batches
        avg_recon = epoch_recon_loss / n_batches
        epoch_time = time.time() - t0

        print(
            f"Epoch {epoch+1}/{args.epochs} | "
            f"avg_loss={avg_loss:.4f} pred={avg_pred:.4f} kl={avg_kl:.4f} "
            f"recon={avg_recon:.4f} | {epoch_time:.1f}s"
        )

        # Save checkpoint
        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}")

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

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

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    # Data
    parser.add_argument("--data_dir", default="/home/ubuntu/lewm-tts/processed_data")
    parser.add_argument("--output_dir", default="/home/ubuntu/lewm-tts/output")
    # Model
    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("--audio_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("--kl_weight", type=float, default=0.1)
    parser.add_argument("--recon_weight", type=float, default=1.0)
    # Training
    parser.add_argument("--epochs", type=int, default=200)
    parser.add_argument("--batch_size", type=int, default=16)
    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)
    # Logging
    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)


Updated task #2 status


Now rewrite inference.py — remove Griffin-Lim entirely, use Vocos directly:


Write

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

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

from model import LeWMTTS

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

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

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

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

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

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

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

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

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

        d = self.config["d_model"]
        start_emb = torch.randn(1, 1, d, device=self.device) * self._start_emb_scale
        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

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

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

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

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

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

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

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

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

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

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

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

    tts = LeWMTTSInference(args.checkpoint)

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

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

if __name__ == "__main__":
    main()


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

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

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

from model import LeWMTTS

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

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

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

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

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

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

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

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

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

        d = self.config["d_model"]
        start_emb = torch.randn(1, 1, d, device=self.device) * self._start_emb_scale
        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

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

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

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

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

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

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

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

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

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

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

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

    tts = LeWMTTSInference(args.checkpoint)

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

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

if __name__ == "__main__":
    main()

-Text → JEPA predictor → mel reconstruction → vocoder → waveform

+Text → JEPA predictor → mel reconstruction → Vocos vocoder → waveform

-    """Griffin-Lim vocoder as fallback."""

-    def __init__(self, n_fft=1024, hop_length=256, n_mels=80, sample_rate=24000):

-        self.inverse_mel = torchaudio.transforms.InverseMelScale(

-        self.griffin_lim = torchaudio.transforms.GriffinLim(

-        waveform = self.griffin_lim(spec)

-        # Vocoder: Griffin-Lim → Vocos refinement

-        self.gl_vocoder = GriffinLimVocoder(

-            n_fft=1024, hop_length=256, n_mels=80, sample_rate=24000,

+        # Vocos vocoder (directly decodes 100-mel)

-        self._vocos = Vocos.from_pretrained("charactr/vocos-mel-24khz")

+        self.vocos = Vocos.from_pretrained("charactr/vocos-mel-24khz")

-        """Decode mel (B, 80, T) → waveform.

-        Griffin-Lim (80-mel → waveform) → Vocos refinement (audio → 100-mel → decode)."""

-        # Step 1: GL to get waveform from 80-mel

-        waveform = self.gl_vocoder(mel)  # [B, T]

-        # Step 2: Vocos copy-synthesis refinement

+        """Decode mel [B, 100, T] → waveform via Vocos."""

-            waveform = self._vocos(waveform)  # [B, T] → [B, T]

+            waveform = self.vocos.decode(mel)  # [B, T_audio]

-        Generate speech from text using AR prediction with KV-cache.

+        """Generate speech from text using AR prediction with KV-cache."""

-        # Start with a Gaussian sample matching training distribution

-        # Process start token through cache (step 0)

-            # Stopping: track norm stability — stop if converged

-        waveform = self.vocoder_decode(mel.cpu())

+        waveform = self.vocos_decode(mel.cpu())

-        Teacher-forced reconstruction: encode real mel → decode back.

-        Tests mel decoder + vocoder quality without AR drift.

+        """Teacher-forced reconstruction: encode real mel → decode back."""

-        waveform = self.vocoder_decode(mel_recon.cpu())

+        waveform = self.vocos_decode(mel_recon.cpu())

-        Hybrid: encode real mel, run predictor on real embeddings,

-        decode predicted embeddings. Tests prediction quality without AR drift.

+        """Hybrid: encode real mel → predictor → decode predicted embeddings."""

-        # We need text too — use dummy text (empty)

-        # This tests if the predictor can reproduce the audio embeddings

-        # Run predictor on real embeddings (teacher-forced)

-        # Use real first embedding + predicted rest

+            waveform = waveform / mx * 0.95

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

import torch
import torch.nn as nn
import torchaudio
import json
import argparse
import numpy as np
from pathlib import Path

from model import LeWMTTS

class GriffinLimVocoder:
    """Griffin-Lim vocoder as fallback."""

    def __init__(self, n_fft=1024, hop_length=256, n_mels=80, sample_rate=24000):
        self.n_fft = n_fft
        self.hop_length = hop_length
        self.sample_rate = sample_rate

        self.inverse_mel = torchaudio.transforms.InverseMelScale(
            n_stft=n_fft // 2 + 1,
            n_mels=n_mels,
            sample_rate=sample_rate,
        )
        self.griffin_lim = torchaudio.transforms.GriffinLim(
            n_fft=n_fft,
            hop_length=hop_length,
            n_iter=64,
            power=1.0,
        )

    def __call__(self, mel):
        mel = torch.exp(mel)
        spec = self.inverse_mel(mel)
        waveform = self.griffin_lim(spec)
        return waveform

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

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

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

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

        # Vocoder: Griffin-Lim → Vocos refinement
        self.gl_vocoder = GriffinLimVocoder(
            n_fft=1024, hop_length=256, n_mels=80, sample_rate=24000,
        )
        from vocos import Vocos
        self._vocos = Vocos.from_pretrained("charactr/vocos-mel-24khz")

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

    def vocoder_decode(self, mel):
        """Decode mel (B, 80, T) → waveform.
        Griffin-Lim (80-mel → waveform) → Vocos refinement (audio → 100-mel → decode)."""
        # Step 1: GL to get waveform from 80-mel
        waveform = self.gl_vocoder(mel)  # [B, T]
        # Step 2: Vocos copy-synthesis refinement
        with torch.no_grad():
            waveform = self._vocos(waveform)  # [B, T] → [B, T]
        return waveform

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

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

        d = self.config["d_model"]

        # Start with a Gaussian sample matching training distribution
        start_emb = torch.randn(1, 1, d, device=self.device) * self._start_emb_scale
        all_embeddings = [start_emb]

        # Initialize KV cache
        cache = self.model.init_ar_cache()

        # Process start token through cache (step 0)
        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)

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

        audio_embs = torch.cat(all_embeddings, dim=1)

        # Decode to mel
        mel = self.model.mel_decoder(audio_embs)

        # Vocoder
        waveform = self.vocoder_decode(mel.cpu())
        return waveform.squeeze().numpy(), 24000

    @torch.no_grad()
    def reconstruct(self, mel_path):
        """
        Teacher-forced reconstruction: encode real mel → decode back.
        Tests mel decoder + vocoder quality without AR drift.
        """
        mel = torch.load(mel_path, weights_only=True).unsqueeze(0).to(self.device)
        # [1, n_mels, T]

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

        # Decode from mean embeddings
        mel_recon = self.model.mel_decoder(mu)

        # Vocoder
        waveform = self.vocoder_decode(mel_recon.cpu())
        return waveform.squeeze().numpy(), 24000

    @torch.no_grad()
    def reconstruct_with_prediction(self, mel_path):
        """
        Hybrid: encode real mel, run predictor on real embeddings,
        decode predicted embeddings. Tests prediction quality without AR drift.
        """
        mel = torch.load(mel_path, weights_only=True).unsqueeze(0).to(self.device)

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

        # We need text too — use dummy text (empty)
        # This tests if the predictor can reproduce the audio embeddings
        text_tokens = torch.zeros(1, 1, dtype=torch.long, device=self.device)
        text_emb = self.model.text_encoder(text_tokens)

        # Run predictor on real embeddings (teacher-forced)
        predicted = self.model.predictor(mu[:, :-1], text_emb)

        # Use real first embedding + predicted rest
        combined = torch.cat([mu[:, :1], predicted], dim=1)

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

    def save_audio(self, waveform, sr, output_path):
        wav_tensor = torch.from_numpy(waveform).unsqueeze(0)
        torchaudio.save(output_path, wav_tensor, sr)
        print(f"Saved: {output_path} ({len(waveform)/sr:.2f}s)")

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

    tts = LeWMTTSInference(args.checkpoint)

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

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

if __name__ == "__main__":
    main()


Updated task #1 status


Now update inference_engine.py — remove Griffin-Lim, direct Vocos decode, 100 mels for prompt extraction:


Edit

class Vocoder:
    """Griffin-Lim → Vocos refinement pipeline for 80-mel → waveform."""

    def __init__(self, device="cpu"):
        self.inv_mel = torchaudio.transforms.InverseMelScale(
            n_stft=513, n_mels=80, sample_rate=24000,
        )
        self.griffin_lim = torchaudio.transforms.GriffinLim(
            n_fft=1024, hop_length=256, n_iter=64, power=1.0,
        )
        from vocos import Vocos
        self.vocos = Vocos.from_pretrained("charactr/vocos-mel-24khz")
        self.vocos.eval()

    def __call__(self, mel):
        """mel: [B, 80, T] → waveforms list of [T_audio] numpy arrays."""
        # Griffin-Lim per sample (handles variable lengths)
        waveforms = []
        for i in range(mel.shape[0]):
            m = mel[i:i+1]  # [1, 80, T]
            m_exp = torch.exp(m)
            spec = self.inv_mel(m_exp)
            wav = self.griffin_lim(spec)  # [1, T_audio]
            # Vocos refinement
            with torch.no_grad():
                wav = self.vocos(wav)  # [1, T_audio]
            waveforms.append(wav.squeeze(0).numpy())
        return waveforms

class Vocoder:
    """Vocos vocoder for 100-mel → waveform (direct decode, no Griffin-Lim)."""

    def __init__(self):
        from vocos import Vocos
        self.vocos = Vocos.from_pretrained("charactr/vocos-mel-24khz")
        self.vocos.eval()

    def __call__(self, mel):
        """mel: [B, 100, T] → waveforms list of [T_audio] numpy arrays."""
        with torch.no_grad():
            # Vocos decode handles batched input
            waveforms = []
            for i in range(mel.shape[0]):
                wav = self.vocos.decode(mel[i:i+1])  # [1, T_audio]
                waveforms.append(wav.squeeze(0).numpy())
        return waveforms


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

class Vocoder:
    """Vocos vocoder for 100-mel → waveform (direct decode, no Griffin-Lim)."""

    def __init__(self):
        from vocos import Vocos
        self.vocos = Vocos.from_pretrained("charactr/vocos-mel-24khz")
        self.vocos.eval()

    def __call__(self, mel):
        """mel: [B, 100, T] → waveforms list of [T_audio] numpy arrays."""
        with torch.no_grad():
            # Vocos decode handles batched input
            waveforms = []
            for i in range(mel.shape[0]):
                wav = self.vocos.decode(mel[i:i+1])  # [1, T_audio]
                waveforms.append(wav.squeeze(0).numpy())
        return waveforms

"""
Production inference engine for LeWM TTS.
Features:
  - Batched inference (many concurrent requests in parallel)
  - KV-cache (O(1) per step instead of O(n))
  - FP16
  - Audio prompt support (for multi-speaker / voice cloning)
  - Continuous batching (add/remove requests dynamically)
  - Griffin-Lim + Vocos vocoder pipeline
"""

import torch
import torch.nn.functional as F
import torchaudio
import numpy as np
import time
from dataclasses import dataclass, field
from typing import Optional
from pathlib import Path

from model import LeWMTTS

# ─── Vocoder ────────────────────────────────────────────────────────────────

class Vocoder:
    """Griffin-Lim → Vocos refinement pipeline for 80-mel → waveform."""

    def __init__(self, device="cpu"):
        self.inv_mel = torchaudio.transforms.InverseMelScale(
            n_stft=513, n_mels=80, sample_rate=24000,
        )
        self.griffin_lim = torchaudio.transforms.GriffinLim(
            n_fft=1024, hop_length=256, n_iter=64, power=1.0,
        )
        from vocos import Vocos
        self.vocos = Vocos.from_pretrained("charactr/vocos-mel-24khz")
        self.vocos.eval()

    def __call__(self, mel):
        """mel: [B, 80, T] → waveforms list of [T_audio] numpy arrays."""
        # Griffin-Lim per sample (handles variable lengths)
        waveforms = []
        for i in range(mel.shape[0]):
            m = mel[i:i+1]  # [1, 80, T]
            m_exp = torch.exp(m)
            spec = self.inv_mel(m_exp)
            wav = self.griffin_lim(spec)  # [1, T_audio]
            # Vocos refinement
            with torch.no_grad():
                wav = self.vocos(wav)  # [1, T_audio]
            waveforms.append(wav.squeeze(0).numpy())
        return waveforms

# ─── Request / Result ───────────────────────────────────────────────────────

@dataclass
class TTSRequest:
    text: str
    max_steps: int = 300
    temperature: float = 0.0
    prompt_mel: Optional[torch.Tensor] = None  # [1, 80, T] for voice cloning

@dataclass
class TTSResult:
    waveform: np.ndarray
    sample_rate: int = 24000
    steps_taken: int = 0
    generation_time: float = 0.0

    def save(self, path):
        wav = torch.from_numpy(self.waveform).unsqueeze(0)
        torchaudio.save(path, wav, self.sample_rate)

    @property
    def duration(self):
        return len(self.waveform) / self.sample_rate

# ─── Inference Engine ───────────────────────────────────────────────────────

class InferenceEngine:
    """
    High-throughput batched TTS inference engine.

    Usage:
        engine = InferenceEngine("checkpoint.pt")

        # Single request
        result = engine.synthesize("नमस्ते भारत")

        # Batched (high throughput)
        results = engine.synthesize_batch(["text1", "text2", ...])

        # With voice prompt
        result = engine.synthesize("text", prompt_audio="ref.wav")

        # Benchmark
        engine.benchmark()
    """

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

        # Load model
        ckpt = torch.load(checkpoint_path, map_location=self.device, weights_only=False)
        self.config = ckpt["config"]
        self.model = LeWMTTS(self.config).to(self.device).to(dtype)
        self.model.load_state_dict(ckpt["model"])
        self.model.eval()

        self.d_model = self.config["d_model"]

        # Vocoder (runs on CPU)
        self.vocoder = Vocoder(device="cpu")

        # Precompute start embedding scale
        self._start_scale = 5.0

        print(f"Engine ready: {checkpoint_path}")
        print(f"  Device: {self.device}, Dtype: {dtype}")
        print(f"  Model: {sum(p.numel() for p in self.model.parameters())/1e6:.1f}M params")

    # ─── Text encoding ──────────────────────────────────────────────────

    def _encode_texts(self, texts):
        """Encode list of texts → padded embeddings + mask.
        Returns: text_emb [B, T_max, d], text_mask [B, T_max] (True=padding)
        """
        # Byte-level tokenization
        token_lists = [list(t.encode("utf-8")) for t in texts]
        max_len = max(len(t) for t in token_lists)

        # Pad to same length
        padded = torch.zeros(len(texts), max_len, dtype=torch.long, device=self.device)
        mask = torch.ones(len(texts), max_len, dtype=torch.bool, device=self.device)

        for i, tokens in enumerate(token_lists):
            padded[i, :len(tokens)] = torch.tensor(tokens, dtype=torch.long)
            mask[i, :len(tokens)] = False

        with torch.no_grad():
            text_emb = self.model.text_encoder(padded, text_mask=mask)

        return text_emb, mask

    # ─── Audio prompt encoding ──────────────────────────────────────────

    def _encode_prompt(self, prompt_mel):
        """Encode prompt mel → embeddings for AR seeding.
        prompt_mel: [1, 80, T]
        Returns: prompt_embs [1, T_down, d]
        """
        prompt_mel = prompt_mel.to(self.device).to(self.dtype)
        with torch.no_grad():
            embs = self.model.encode_audio(prompt_mel)  # [1, T_down, d]
        return embs

    def load_prompt_audio(self, audio_path, max_seconds=3.0):
        """Load audio file → mel for prompting.
        Returns: mel [1, 80, T]
        """
        wav, sr = torchaudio.load(audio_path)
        if sr != 24000:
            wav = torchaudio.functional.resample(wav, sr, 24000)
        # Trim to max_seconds
        max_samples = int(max_seconds * 24000)
        wav = wav[:1, :max_samples]  # mono, trimmed

        mel_transform = torchaudio.transforms.MelSpectrogram(
            sample_rate=24000, n_fft=1024, hop_length=256,
            n_mels=80, power=2.0,
        )
        mel = mel_transform(wav)
        mel = torch.log(mel.clamp(min=1e-5))
        return mel  # [1, 80, T]

    # ─── Core AR generation (batched) ───────────────────────────────────

    def _generate_batch(self, text_emb, text_mask, batch_size,
                        max_steps=300, temperature=0.0,
                        prompt_embs=None):
        """
        Core batched AR generation with KV-cache.

        Args:
            text_emb: [B, T_text, d]
            text_mask: [B, T_text]
            batch_size: int
            max_steps: int
            temperature: float
            prompt_embs: list of [1, T_prompt, d] or None per request

        Returns:
            all_embs: [B, T_total, d]
            steps_taken: list of int per request
        """
        d = self.d_model
        cache = self.model.init_ar_cache()

        # Track per-request state
        active = torch.ones(batch_size, dtype=torch.bool, device=self.device)
        steps_taken = [0] * batch_size
        all_emb_lists = [[] for _ in range(batch_size)]
        prev_norms = [[] for _ in range(batch_size)]

        # Handle prompts: seed cache with prompt embeddings
        if prompt_embs is not None:
            max_prompt_len = max(
                (e.shape[1] for e in prompt_embs if e is not None), default=0
            )
            if max_prompt_len > 0:
                # Process prompt embeddings through cache step by step
                # Pad prompts to same length
                prompt_padded = torch.zeros(
                    batch_size, max_prompt_len, d,
                    device=self.device, dtype=self.dtype,
                )
                prompt_lengths = []
                for i, pe in enumerate(prompt_embs):
                    if pe is not None:
                        L = pe.shape[1]
                        prompt_padded[i, :L] = pe[0]
                        prompt_lengths.append(L)
                    else:
                        prompt_lengths.append(0)

                # Feed prompt through cache one step at a time
                for t in range(max_prompt_len):
                    step_emb = prompt_padded[:, t:t+1, :]  # [B, 1, d]
                    _, cache = self.model.predict_next_cached(
                        step_emb, text_emb, t, cache, text_mask=text_mask,
                    )
                    # Store embeddings
                    for i in range(batch_size):
                        if t < prompt_lengths[i]:
                            all_emb_lists[i].append(step_emb[i:i+1])

                step_offset = max_prompt_len
                # Use last prompt prediction as first AR input
                last_pred, cache = self.model.predict_next_cached(
                    prompt_padded[:, -1:, :], text_emb, step_offset, cache,
                    text_mask=text_mask,
                )
                current_emb = last_pred
                step_offset += 1
            else:
                step_offset = 0
                current_emb = None
        else:
            step_offset = 0
            current_emb = None

        # If no prompt, start with random embedding
        if current_emb is None:
            start_emb = torch.randn(
                batch_size, 1, d, device=self.device, dtype=self.dtype,
            ) * self._start_scale
            current_emb, cache = self.model.predict_next_cached(
                start_emb, text_emb, 0, cache, text_mask=text_mask,
            )
            for i in range(batch_size):
                all_emb_lists[i].append(start_emb[i:i+1])
            step_offset = 1

        if temperature > 0:
            current_emb = current_emb + torch.randn_like(current_emb) * temperature

        for i in range(batch_size):
            all_emb_lists[i].append(current_emb[i:i+1])

        # AR loop
        for step in range(1, max_steps):
            if not active.any():
                break

            next_emb, cache = self.model.predict_next_cached(
                current_emb, text_emb, step + step_offset, cache,
                text_mask=text_mask,
            )

            if temperature > 0:
                next_emb = next_emb + torch.randn_like(next_emb) * temperature

            # Store and check stopping per request
            for i in range(batch_size):
                if not active[i]:
                    continue

                all_emb_lists[i].append(next_emb[i:i+1])
                steps_taken[i] = step

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

            current_emb = next_emb

        # Concatenate embeddings per request
        results = []
        for i in range(batch_size):
            embs = torch.cat(all_emb_lists[i], dim=1)  # [1, T, d]
            results.append(embs)

        return results, steps_taken

    # ─── Public API ─────────────────────────────────────────────────────

    def synthesize(self, text, max_steps=300, temperature=0.0,
                   prompt_audio=None, prompt_seconds=3.0):
        """Synthesize single text.

        Args:
            text: Hindi text string
            max_steps: max AR steps
            temperature: sampling temperature (0 = greedy)
            prompt_audio: path to reference audio for voice cloning
            prompt_seconds: max seconds of prompt to use

        Returns: TTSResult
        """
        return self.synthesize_batch(
            [text], max_steps=max_steps, temperature=temperature,
            prompt_audios=[prompt_audio] if prompt_audio else None,
            prompt_seconds=prompt_seconds,
        )[0]

    def synthesize_batch(self, texts, max_steps=300, temperature=0.0,
                         prompt_audios=None, prompt_seconds=3.0):
        """Synthesize a batch of texts in parallel.

        Args:
            texts: list of text strings
            max_steps: max AR steps
            temperature: sampling temperature
            prompt_audios: list of audio paths (or None) per text
            prompt_seconds: max seconds of prompt

        Returns: list of TTSResult
        """
        B = len(texts)
        t0 = time.time()

        # Encode texts
        text_emb, text_mask = self._encode_texts(texts)

        # Encode prompts if given
        prompt_embs = None
        if prompt_audios is not None:
            prompt_embs = []
            for pa in prompt_audios:
                if pa is not None:
                    mel = self.load_prompt_audio(pa, max_seconds=prompt_seconds)
                    embs = self._encode_prompt(mel)
                    prompt_embs.append(embs)
                else:
                    prompt_embs.append(None)

        # Generate
        with torch.no_grad():
            emb_results, steps_taken = self._generate_batch(
                text_emb, text_mask, B,
                max_steps=max_steps, temperature=temperature,
                prompt_embs=prompt_embs,
            )

        # Decode to mel → waveform
        results = []
        for i in range(B):
            with torch.no_grad():
                mel = self.model.mel_decoder(emb_results[i].to(self.device))
                # mel: [1, 80, T]
            wav_list = self.vocoder(mel.float().cpu())
            waveform = wav_list[0]

            # Normalize
            mx = np.abs(waveform).max()
            if mx > 0:
                waveform = waveform / mx * 0.95

            results.append(TTSResult(
                waveform=waveform,
                steps_taken=steps_taken[i],
                generation_time=time.time() - t0,
            ))

        return results

    # ─── Benchmark ──────────────────────────────────────────────────────

    def benchmark(self, batch_sizes=None, steps=300):
        """Benchmark throughput at various batch sizes."""
        if batch_sizes is None:
            batch_sizes = [1, 4, 16, 32, 64, 128, 256]

        text = "भारत एक महान देश है और हम सब भारतवासी हैं।"
        audio_per_step = 4 * 256 / 24000  # seconds of audio per AR step

        print(f"\n{'='*70}")
        print(f"  LeWM TTS Inference Benchmark — {self.device}, {self.dtype}")
        print(f"  Model: {self.config['d_model']}d, {self.config['predictor_layers']}L predictor")
        print(f"  AR steps: {steps}")
        print(f"{'='*70}")
        print(f"  {'Batch':>6} {'Time(s)':>8} {'Audio(s)':>9} {'RTF':>8} "
              f"{'Speed':>8} {'Throughput':>12} {'Per-req':>10}")
        print(f"  {'-'*6} {'-'*8} {'-'*9} {'-'*8} {'-'*8} {'-'*12} {'-'*10}")

        for B in batch_sizes:
            try:
                texts = [text] * B
                token_lists = [list(t.encode("utf-8")) for t in texts]
                max_len = max(len(t) for t in token_lists)
                padded = torch.zeros(B, max_len, dtype=torch.long, device=self.device)
                mask = torch.ones(B, max_len, dtype=torch.bool, device=self.device)
                for i, tokens in enumerate(token_lists):
                    padded[i, :len(tokens)] = torch.tensor(tokens, dtype=torch.long)
                    mask[i, :len(tokens)] = False

                with torch.no_grad():
                    text_emb = self.model.text_encoder(padded, text_mask=mask)

                # Warmup
                if B == batch_sizes[0]:
                    cache = self.model.init_ar_cache()
                    start = torch.randn(B, 1, self.d_model,
                                        device=self.device, dtype=self.dtype) * self._start_scale
                    nxt = start
                    for s in range(5):
                        nxt, cache = self.model.predict_next_cached(
                            nxt, text_emb, s, cache, text_mask=mask)
                    torch.cuda.synchronize()

                # Benchmark
                cache = self.model.init_ar_cache()
                start = torch.randn(B, 1, self.d_model,
                                    device=self.device, dtype=self.dtype) * self._start_scale
                nxt = start

                torch.cuda.synchronize()
                t0 = time.time()
                for s in range(steps):
                    nxt, cache = self.model.predict_next_cached(
                        nxt, text_emb, s, cache, text_mask=mask)
                torch.cuda.synchronize()
                elapsed = time.time() - t0

                total_audio = B * steps * audio_per_step
                rtf = elapsed / (steps * audio_per_step)
                speed = (steps * audio_per_step) / elapsed
                throughput = total_audio / elapsed
                per_req = elapsed / B

                print(f"  {B:>6} {elapsed:>8.3f} {total_audio:>9.1f} {rtf:>8.4f} "
                      f"{speed:>7.0f}x {throughput:>10.0f}x {per_req:>9.3f}s")

            except torch.cuda.OutOfMemoryError:
                print(f"  {B:>6} {'OOM':>8}")
                torch.cuda.empty_cache()
                break

        print(f"{'='*70}")
        print(f"  Throughput = total audio-seconds generated per wall-second")
        print(f"  Speed = single-request real-time factor")
        print(f"  Per-req = wall-time per request in batch")
        print()

# ─── CLI ────────────────────────────────────────────────────────────────────

if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(description="LeWM TTS Inference Engine")
    parser.add_argument("--checkpoint", required=True)
    parser.add_argument("--text", default=None, help="Text to synthesize")
    parser.add_argument("--texts", nargs="+", help="Multiple texts for batch synthesis")
    parser.add_argument("--prompt", default=None, help="Reference audio for voice cloning")
    parser.add_argument("--output", default="output.wav")
    parser.add_argument("--output_dir", default=None, help="Output dir for batch mode")
    parser.add_argument("--max_steps", type=int, default=300)
    parser.add_argument("--temperature", type=float, default=0.0)
    parser.add_argument("--benchmark", action="store_true", help="Run throughput benchmark")
    parser.add_argument("--device", default="cuda")
    parser.add_argument("--fp32", action="store_true", help="Use FP32 instead of FP16")
    args = parser.parse_args()

    dtype = torch.float32 if args.fp32 else torch.float16
    engine = InferenceEngine(args.checkpoint, device=args.device, dtype=dtype)

    if args.benchmark:
        engine.benchmark()

    elif args.texts:
        results = engine.synthesize_batch(
            args.texts, max_steps=args.max_steps, temperature=args.temperature,
            prompt_audios=[args.prompt] * len(args.texts) if args.prompt else None,
        )
        out_dir = Path(args.output_dir or "output")
        out_dir.mkdir(exist_ok=True)
        for i, r in enumerate(results):
            path = out_dir / f"batch_{i:03d}.wav"
            r.save(str(path))
            print(f"  [{i}] {path} — {r.duration:.2f}s, {r.steps_taken} steps")

    elif args.text:
        result = engine.synthesize(
            args.text, max_steps=args.max_steps, temperature=args.temperature,
            prompt_audio=args.prompt,
        )
        result.save(args.output)
        print(f"Saved: {args.output} — {result.duration:.2f}s, "
              f"{result.steps_taken} steps, {result.generation_time:.2f}s")
    else:
        parser.print_help()

-    """Griffin-Lim → Vocos refinement pipeline for 80-mel → waveform."""

+    """Vocos vocoder for 100-mel → waveform (direct decode, no Griffin-Lim)."""

-        self.inv_mel = torchaudio.transforms.InverseMelScale(

-            n_stft=513, n_mels=80, sample_rate=24000,

-        self.griffin_lim = torchaudio.transforms.GriffinLim(

-        """mel: [B, 80, T] → waveforms list of [T_audio] numpy arrays."""

-        # Griffin-Lim per sample (handles variable lengths)

-            wav = self.griffin_lim(spec)  # [1, T_audio]

-                wav = self.vocos(wav)  # [1, T_audio]

-            waveforms.append(wav.squeeze(0).numpy())

+        """mel: [B, 100, T] → waveforms list of [T_audio] numpy arrays."""

+            # Vocos decode handles batched input

+            for i in range(mel.shape[0]):

+                wav = self.vocos.decode(mel[i:i+1])  # [1, T_audio]

+                waveforms.append(wav.squeeze(0).numpy())


Update the prompt mel extraction to use 100 mels matching Vocos:


        mel_transform = torchaudio.transforms.MelSpectrogram(
            sample_rate=24000, n_fft=1024, hop_length=256,
            n_mels=80, power=2.0,
        )

        mel_transform = torchaudio.transforms.MelSpectrogram(
            sample_rate=24000, n_fft=1024, hop_length=256,
            n_mels=100, power=1.0,
        )


        mel_transform = torchaudio.transforms.MelSpectrogram(
            sample_rate=24000, n_fft=1024, hop_length=256,
            n_mels=100, power=1.0,
        )

"""
Production inference engine for LeWM TTS.
Features:
  - Batched inference (many concurrent requests in parallel)
  - KV-cache (O(1) per step instead of O(n))
  - FP16
  - Audio prompt support (for multi-speaker / voice cloning)
  - Continuous batching (add/remove requests dynamically)
  - Griffin-Lim + Vocos vocoder pipeline
"""

import torch
import torch.nn.functional as F
import torchaudio
import numpy as np
import time
from dataclasses import dataclass, field
from typing import Optional
from pathlib import Path

from model import LeWMTTS

# ─── Vocoder ────────────────────────────────────────────────────────────────

class Vocoder:
    """Vocos vocoder for 100-mel → waveform (direct decode, no Griffin-Lim)."""

    def __init__(self):
        from vocos import Vocos
        self.vocos = Vocos.from_pretrained("charactr/vocos-mel-24khz")
        self.vocos.eval()

    def __call__(self, mel):
        """mel: [B, 100, T] → waveforms list of [T_audio] numpy arrays."""
        with torch.no_grad():
            # Vocos decode handles batched input
            waveforms = []
            for i in range(mel.shape[0]):
                wav = self.vocos.decode(mel[i:i+1])  # [1, T_audio]
                waveforms.append(wav.squeeze(0).numpy())
        return waveforms

# ─── Request / Result ───────────────────────────────────────────────────────

@dataclass
class TTSRequest:
    text: str
    max_steps: int = 300
    temperature: float = 0.0
    prompt_mel: Optional[torch.Tensor] = None  # [1, 80, T] for voice cloning

@dataclass
class TTSResult:
    waveform: np.ndarray
    sample_rate: int = 24000
    steps_taken: int = 0
    generation_time: float = 0.0

    def save(self, path):
        wav = torch.from_numpy(self.waveform).unsqueeze(0)
        torchaudio.save(path, wav, self.sample_rate)

    @property
    def duration(self):
        return len(self.waveform) / self.sample_rate

# ─── Inference Engine ───────────────────────────────────────────────────────

class InferenceEngine:
    """
    High-throughput batched TTS inference engine.

    Usage:
        engine = InferenceEngine("checkpoint.pt")

        # Single request
        result = engine.synthesize("नमस्ते भारत")

        # Batched (high throughput)
        results = engine.synthesize_batch(["text1", "text2", ...])

        # With voice prompt
        result = engine.synthesize("text", prompt_audio="ref.wav")

        # Benchmark
        engine.benchmark()
    """

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

        # Load model
        ckpt = torch.load(checkpoint_path, map_location=self.device, weights_only=False)
        self.config = ckpt["config"]
        self.model = LeWMTTS(self.config).to(self.device).to(dtype)
        self.model.load_state_dict(ckpt["model"])
        self.model.eval()

        self.d_model = self.config["d_model"]

        # Vocoder (runs on CPU)
        self.vocoder = Vocoder(device="cpu")

        # Precompute start embedding scale
        self._start_scale = 5.0

        print(f"Engine ready: {checkpoint_path}")
        print(f"  Device: {self.device}, Dtype: {dtype}")
        print(f"  Model: {sum(p.numel() for p in self.model.parameters())/1e6:.1f}M params")

    # ─── Text encoding ──────────────────────────────────────────────────

    def _encode_texts(self, texts):
        """Encode list of texts → padded embeddings + mask.
        Returns: text_emb [B, T_max, d], text_mask [B, T_max] (True=padding)
        """
        # Byte-level tokenization
        token_lists = [list(t.encode("utf-8")) for t in texts]
        max_len = max(len(t) for t in token_lists)

        # Pad to same length
        padded = torch.zeros(len(texts), max_len, dtype=torch.long, device=self.device)
        mask = torch.ones(len(texts), max_len, dtype=torch.bool, device=self.device)

        for i, tokens in enumerate(token_lists):
            padded[i, :len(tokens)] = torch.tensor(tokens, dtype=torch.long)
            mask[i, :len(tokens)] = False

        with torch.no_grad():
            text_emb = self.model.text_encoder(padded, text_mask=mask)

        return text_emb, mask

    # ─── Audio prompt encoding ──────────────────────────────────────────

    def _encode_prompt(self, prompt_mel):
        """Encode prompt mel → embeddings for AR seeding.
        prompt_mel: [1, 80, T]
        Returns: prompt_embs [1, T_down, d]
        """
        prompt_mel = prompt_mel.to(self.device).to(self.dtype)
        with torch.no_grad():
            embs = self.model.encode_audio(prompt_mel)  # [1, T_down, d]
        return embs

    def load_prompt_audio(self, audio_path, max_seconds=3.0):
        """Load audio file → mel for prompting.
        Returns: mel [1, 80, T]
        """
        wav, sr = torchaudio.load(audio_path)
        if sr != 24000:
            wav = torchaudio.functional.resample(wav, sr, 24000)
        # Trim to max_seconds
        max_samples = int(max_seconds * 24000)
        wav = wav[:1, :max_samples]  # mono, trimmed

        mel_transform = torchaudio.transforms.MelSpectrogram(
            sample_rate=24000, n_fft=1024, hop_length=256,
            n_mels=80, power=2.0,
        )
        mel = mel_transform(wav)
        mel = torch.log(mel.clamp(min=1e-5))
        return mel  # [1, 80, T]

    # ─── Core AR generation (batched) ───────────────────────────────────

    def _generate_batch(self, text_emb, text_mask, batch_size,
                        max_steps=300, temperature=0.0,
                        prompt_embs=None):
        """
        Core batched AR generation with KV-cache.

        Args:
            text_emb: [B, T_text, d]
            text_mask: [B, T_text]
            batch_size: int
            max_steps: int
            temperature: float
            prompt_embs: list of [1, T_prompt, d] or None per request

        Returns:
            all_embs: [B, T_total, d]
            steps_taken: list of int per request
        """
        d = self.d_model
        cache = self.model.init_ar_cache()

        # Track per-request state
        active = torch.ones(batch_size, dtype=torch.bool, device=self.device)
        steps_taken = [0] * batch_size
        all_emb_lists = [[] for _ in range(batch_size)]
        prev_norms = [[] for _ in range(batch_size)]

        # Handle prompts: seed cache with prompt embeddings
        if prompt_embs is not None:
            max_prompt_len = max(
                (e.shape[1] for e in prompt_embs if e is not None), default=0
            )
            if max_prompt_len > 0:
                # Process prompt embeddings through cache step by step
                # Pad prompts to same length
                prompt_padded = torch.zeros(
                    batch_size, max_prompt_len, d,
                    device=self.device, dtype=self.dtype,
                )
                prompt_lengths = []
                for i, pe in enumerate(prompt_embs):
                    if pe is not None:
                        L = pe.shape[1]
                        prompt_padded[i, :L] = pe[0]
                        prompt_lengths.append(L)
                    else:
                        prompt_lengths.append(0)

                # Feed prompt through cache one step at a time
                for t in range(max_prompt_len):
                    step_emb = prompt_padded[:, t:t+1, :]  # [B, 1, d]
                    _, cache = self.model.predict_next_cached(
                        step_emb, text_emb, t, cache, text_mask=text_mask,
                    )
                    # Store embeddings
                    for i in range(batch_size):
                        if t < prompt_lengths[i]:
                            all_emb_lists[i].append(step_emb[i:i+1])

                step_offset = max_prompt_len
                # Use last prompt prediction as first AR input
                last_pred, cache = self.model.predict_next_cached(
                    prompt_padded[:, -1:, :], text_emb, step_offset, cache,
                    text_mask=text_mask,
                )
                current_emb = last_pred
                step_offset += 1
            else:
                step_offset = 0
                current_emb = None
        else:
            step_offset = 0
            current_emb = None

        # If no prompt, start with random embedding
        if current_emb is None:
            start_emb = torch.randn(
                batch_size, 1, d, device=self.device, dtype=self.dtype,
            ) * self._start_scale
            current_emb, cache = self.model.predict_next_cached(
                start_emb, text_emb, 0, cache, text_mask=text_mask,
            )
            for i in range(batch_size):
                all_emb_lists[i].append(start_emb[i:i+1])
            step_offset = 1

        if temperature > 0:
            current_emb = current_emb + torch.randn_like(current_emb) * temperature

        for i in range(batch_size):
            all_emb_lists[i].append(current_emb[i:i+1])

        # AR loop
        for step in range(1, max_steps):
            if not active.any():
                break

            next_emb, cache = self.model.predict_next_cached(
                current_emb, text_emb, step + step_offset, cache,
                text_mask=text_mask,
            )

            if temperature > 0:
                next_emb = next_emb + torch.randn_like(next_emb) * temperature

            # Store and check stopping per request
            for i in range(batch_size):
                if not active[i]:
                    continue

                all_emb_lists[i].append(next_emb[i:i+1])
                steps_taken[i] = step

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

            current_emb = next_emb

        # Concatenate embeddings per request
        results = []
        for i in range(batch_size):
            embs = torch.cat(all_emb_lists[i], dim=1)  # [1, T, d]
            results.append(embs)

        return results, steps_taken

    # ─── Public API ─────────────────────────────────────────────────────

    def synthesize(self, text, max_steps=300, temperature=0.0,
                   prompt_audio=None, prompt_seconds=3.0):
        """Synthesize single text.

        Args:
            text: Hindi text string
            max_steps: max AR steps
            temperature: sampling temperature (0 = greedy)
            prompt_audio: path to reference audio for voice cloning
            prompt_seconds: max seconds of prompt to use

        Returns: TTSResult
        """
        return self.synthesize_batch(
            [text], max_steps=max_steps, temperature=temperature,
            prompt_audios=[prompt_audio] if prompt_audio else None,
            prompt_seconds=prompt_seconds,
        )[0]

    def synthesize_batch(self, texts, max_steps=300, temperature=0.0,
                         prompt_audios=None, prompt_seconds=3.0):
        """Synthesize a batch of texts in parallel.

        Args:
            texts: list of text strings
            max_steps: max AR steps
            temperature: sampling temperature
            prompt_audios: list of audio paths (or None) per text
            prompt_seconds: max seconds of prompt

        Returns: list of TTSResult
        """
        B = len(texts)
        t0 = time.time()

        # Encode texts
        text_emb, text_mask = self._encode_texts(texts)

        # Encode prompts if given
        prompt_embs = None
        if prompt_audios is not None:
            prompt_embs = []
            for pa in prompt_audios:
                if pa is not None:
                    mel = self.load_prompt_audio(pa, max_seconds=prompt_seconds)
                    embs = self._encode_prompt(mel)
                    prompt_embs.append(embs)
                else:
                    prompt_embs.append(None)

        # Generate
        with torch.no_grad():
            emb_results, steps_taken = self._generate_batch(
                text_emb, text_mask, B,
                max_steps=max_steps, temperature=temperature,
                prompt_embs=prompt_embs,
            )

        # Decode to mel → waveform
        results = []
        for i in range(B):
            with torch.no_grad():
                mel = self.model.mel_decoder(emb_results[i].to(self.device))
                # mel: [1, 80, T]
            wav_list = self.vocoder(mel.float().cpu())
            waveform = wav_list[0]

            # Normalize
            mx = np.abs(waveform).max()
            if mx > 0:
                waveform = waveform / mx * 0.95

            results.append(TTSResult(
                waveform=waveform,
                steps_taken=steps_taken[i],
                generation_time=time.time() - t0,
            ))

        return results

    # ─── Benchmark ──────────────────────────────────────────────────────

    def benchmark(self, batch_sizes=None, steps=300):
        """Benchmark throughput at various batch sizes."""
        if batch_sizes is None:
            batch_sizes = [1, 4, 16, 32, 64, 128, 256]

        text = "भारत एक महान देश है और हम सब भारतवासी हैं।"
        audio_per_step = 4 * 256 / 24000  # seconds of audio per AR step

        print(f"\n{'='*70}")
        print(f"  LeWM TTS Inference Benchmark — {self.device}, {self.dtype}")
        print(f"  Model: {self.config['d_model']}d, {self.config['predictor_layers']}L predictor")
        print(f"  AR steps: {steps}")
        print(f"{'='*70}")
        print(f"  {'Batch':>6} {'Time(s)':>8} {'Audio(s)':>9} {'RTF':>8} "
              f"{'Speed':>8} {'Throughput':>12} {'Per-req':>10}")
        print(f"  {'-'*6} {'-'*8} {'-'*9} {'-'*8} {'-'*8} {'-'*12} {'-'*10}")

        for B in batch_sizes:
            try:
                texts = [text] * B
                token_lists = [list(t.encode("utf-8")) for t in texts]
                max_len = max(len(t) for t in token_lists)
                padded = torch.zeros(B, max_len, dtype=torch.long, device=self.device)
                mask = torch.ones(B, max_len, dtype=torch.bool, device=self.device)
                for i, tokens in enumerate(token_lists):
                    padded[i, :len(tokens)] = torch.tensor(tokens, dtype=torch.long)
                    mask[i, :len(tokens)] = False

                with torch.no_grad():
                    text_emb = self.model.text_encoder(padded, text_mask=mask)

                # Warmup
                if B == batch_sizes[0]:
                    cache = self.model.init_ar_cache()
                    start = torch.randn(B, 1, self.d_model,
                                        device=self.device, dtype=self.dtype) * self._start_scale
                    nxt = start
                    for s in range(5):
                        nxt, cache = self.model.predict_next_cached(
                            nxt, text_emb, s, cache, text_mask=mask)
                    torch.cuda.synchronize()

                # Benchmark
                cache = self.model.init_ar_cache()
                start = torch.randn(B, 1, self.d_model,
                                    device=self.device, dtype=self.dtype) * self._start_scale
                nxt = start

                torch.cuda.synchronize()
                t0 = time.time()
                for s in range(steps):
                    nxt, cache = self.model.predict_next_cached(
                        nxt, text_emb, s, cache, text_mask=mask)
                torch.cuda.synchronize()
                elapsed = time.time() - t0

                total_audio = B * steps * audio_per_step
                rtf = elapsed / (steps * audio_per_step)
                speed = (steps * audio_per_step) / elapsed
                throughput = total_audio / elapsed
                per_req = elapsed / B

                print(f"  {B:>6} {elapsed:>8.3f} {total_audio:>9.1f} {rtf:>8.4f} "
                      f"{speed:>7.0f}x {throughput:>10.0f}x {per_req:>9.3f}s")

            except torch.cuda.OutOfMemoryError:
                print(f"  {B:>6} {'OOM':>8}")
                torch.cuda.empty_cache()
                break

        print(f"{'='*70}")
        print(f"  Throughput = total audio-seconds generated per wall-second")
        print(f"  Speed = single-request real-time factor")
        print(f"  Per-req = wall-time per request in batch")
        print()

# ─── CLI ────────────────────────────────────────────────────────────────────

if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(description="LeWM TTS Inference Engine")
    parser.add_argument("--checkpoint", required=True)
    parser.add_argument("--text", default=None, help="Text to synthesize")
    parser.add_argument("--texts", nargs="+", help="Multiple texts for batch synthesis")
    parser.add_argument("--prompt", default=None, help="Reference audio for voice cloning")
    parser.add_argument("--output", default="output.wav")
    parser.add_argument("--output_dir", default=None, help="Output dir for batch mode")
    parser.add_argument("--max_steps", type=int, default=300)
    parser.add_argument("--temperature", type=float, default=0.0)
    parser.add_argument("--benchmark", action="store_true", help="Run throughput benchmark")
    parser.add_argument("--device", default="cuda")
    parser.add_argument("--fp32", action="store_true", help="Use FP32 instead of FP16")
    args = parser.parse_args()

    dtype = torch.float32 if args.fp32 else torch.float16
    engine = InferenceEngine(args.checkpoint, device=args.device, dtype=dtype)

    if args.benchmark:
        engine.benchmark()

    elif args.texts:
        results = engine.synthesize_batch(
            args.texts, max_steps=args.max_steps, temperature=args.temperature,
            prompt_audios=[args.prompt] * len(args.texts) if args.prompt else None,
        )
        out_dir = Path(args.output_dir or "output")
        out_dir.mkdir(exist_ok=True)
        for i, r in enumerate(results):
            path = out_dir / f"batch_{i:03d}.wav"
            r.save(str(path))
            print(f"  [{i}] {path} — {r.duration:.2f}s, {r.steps_taken} steps")

    elif args.text:
        result = engine.synthesize(
            args.text, max_steps=args.max_steps, temperature=args.temperature,
            prompt_audio=args.prompt,
        )
        result.save(args.output)
        print(f"Saved: {args.output} — {result.duration:.2f}s, "
              f"{result.steps_taken} steps, {result.generation_time:.2f}s")
    else:
        parser.print_help()


Also update the log clamp to match Vocos:


   170→    # ─── Core AR generation (batched) ───────────────────────────────────
   171→
   172→    def _generate_batch(self, text_emb, text_mask, batch_size,
   173→                        max_steps=300, temperature=0.0,
   174→                        prompt_embs=None):
   175→        """
   176→        Core batched AR generation with KV-cache.
   177→
   178→        Args:
   179→            text_emb: [B, T_text, d]

    # ─── Core AR generation (batched) ───────────────────────────────────

    def _generate_batch(self, text_emb, text_mask, batch_size,
                        max_steps=300, temperature=0.0,
                        prompt_embs=None):
        """
        Core batched AR generation with KV-cache.

        Args:
            text_emb: [B, T_text, d]


Bash

grep -n "1e-5\|clip_val\|clamp" /home/ubuntu/lewm-tts/inference_engine.py

Check clamp values in engine


167:        mel = torch.log(mel.clamp(min=1e-5))