File created successfully at: /home/ubuntu/lewm-tts/train.py

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

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

            # 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()
            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/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} | grad={grad_norm.item():.2f} | "
                    f"{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
        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"{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.01)
    # 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)


Now the inference pipeline.


Write

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

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

from model import LeWMTTS

class MelDecoder(nn.Module):
    """
    Decodes JEPA latent embeddings back to mel spectrograms.
    Simple transposed-conv decoder that reverses the AudioEncoder's downsampling.
    """

    def __init__(self, d_model=256, n_mels=80, upsample_factor=4):
        super().__init__()
        self.decoder = nn.Sequential(
            nn.Linear(d_model, d_model),
            nn.GELU(),
        )
        self.upsample = nn.Sequential(
            # Upsample by 2
            nn.ConvTranspose1d(d_model, d_model, kernel_size=4, stride=2, padding=1),
            nn.GELU(),
            # Upsample by 2 again (total 4x)
            nn.ConvTranspose1d(d_model, d_model, kernel_size=4, stride=2, padding=1),
            nn.GELU(),
            # Project to mel dims
            nn.Conv1d(d_model, n_mels, kernel_size=7, padding=3),
        )

    def forward(self, z):
        """
        Args:
            z: [B, T_down, d_model]
        Returns:
            mel: [B, n_mels, T_up] — reconstructed mel spectrogram
        """
        x = self.decoder(z)
        x = x.transpose(1, 2)  # [B, d_model, T_down]
        mel = self.upsample(x)  # [B, n_mels, T_up]
        return mel

class GriffinLimVocoder:
    """Simple Griffin-Lim vocoder as fallback when HiFi-GAN isn't available."""

    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):
        """
        Args:
            mel: [1, n_mels, T] log-mel spectrogram
        Returns:
            waveform: [1, samples]
        """
        # Un-log the mel
        mel = torch.exp(mel)
        # Inverse mel scale
        spec = self.inverse_mel(mel)
        # Griffin-Lim
        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()

        # Build mel decoder (needs to be trained alongside or separately)
        self.mel_decoder = MelDecoder(
            d_model=config["d_model"],
            n_mels=config["n_mels"],
        ).to(self.device)

        # Try loading mel decoder weights
        decoder_path = Path(checkpoint_path).parent / "mel_decoder.pt"
        if decoder_path.exists():
            self.mel_decoder.load_state_dict(
                torch.load(decoder_path, map_location=self.device, weights_only=True)
            )
        self.mel_decoder.eval()

        # Vocoder
        self.vocoder = GriffinLimVocoder(
            n_fft=1024, hop_length=256, n_mels=80, sample_rate=24000,
        )

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

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

    @torch.no_grad()
    def synthesize(self, text, max_steps=500, temperature=0.7):
        """
        Generate speech from text.

        Args:
            text: Input text string (Hindi)
            max_steps: Maximum number of AR steps
            temperature: Sampling temperature (0 = deterministic)
        Returns:
            waveform: numpy array
        """
        # Encode text
        text_tokens = self.text_to_tokens(text)
        text_emb = self.model.text_encoder(text_tokens)

        # Initialize with a learned start embedding
        d = self.config["d_model"]
        # Start with zeros (will be projected)
        current_emb = torch.zeros(1, 1, d, device=self.device)

        # Autoregressive generation
        all_embeddings = [current_emb]

        for step in range(max_steps):
            # Concatenate all embeddings so far
            context = torch.cat(all_embeddings, dim=1)  # [1, t, d]

            # Predict next
            next_emb = self.model.predict_next(context, text_emb)  # [1, 1, d]

            # Add noise for diversity
            if temperature > 0:
                noise = torch.randn_like(next_emb) * temperature * 0.1
                next_emb = next_emb + noise

            all_embeddings.append(next_emb)

            # Simple stopping: if embedding norm is very small
            if next_emb.norm() < 0.01 and step > 10:
                break

        # Combine all embeddings
        audio_embs = torch.cat(all_embeddings[1:], dim=1)  # Skip the zero start token

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

        # Vocoder: mel → waveform
        waveform = self.vocoder(mel.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}")

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--checkpoint", required=True, help="Path to model checkpoint")
    parser.add_argument("--text", required=True, help="Text to synthesize")
    parser.add_argument("--output", default="output.wav", help="Output wav path")
    parser.add_argument("--max_steps", type=int, default=500)
    parser.add_argument("--temperature", type=float, default=0.7)
    args = parser.parse_args()

    tts = LeWMTTSInference(args.checkpoint)
    waveform, sr = tts.synthesize(
        args.text, max_steps=args.max_steps, temperature=args.temperature
    )
    tts.save_audio(waveform, sr, args.output)

if __name__ == "__main__":
    main()


File created successfully at: /home/ubuntu/lewm-tts/inference.py

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

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

from model import LeWMTTS

class MelDecoder(nn.Module):
    """
    Decodes JEPA latent embeddings back to mel spectrograms.
    Simple transposed-conv decoder that reverses the AudioEncoder's downsampling.
    """

    def __init__(self, d_model=256, n_mels=80, upsample_factor=4):
        super().__init__()
        self.decoder = nn.Sequential(
            nn.Linear(d_model, d_model),
            nn.GELU(),
        )
        self.upsample = nn.Sequential(
            # Upsample by 2
            nn.ConvTranspose1d(d_model, d_model, kernel_size=4, stride=2, padding=1),
            nn.GELU(),
            # Upsample by 2 again (total 4x)
            nn.ConvTranspose1d(d_model, d_model, kernel_size=4, stride=2, padding=1),
            nn.GELU(),
            # Project to mel dims
            nn.Conv1d(d_model, n_mels, kernel_size=7, padding=3),
        )

    def forward(self, z):
        """
        Args:
            z: [B, T_down, d_model]
        Returns:
            mel: [B, n_mels, T_up] — reconstructed mel spectrogram
        """
        x = self.decoder(z)
        x = x.transpose(1, 2)  # [B, d_model, T_down]
        mel = self.upsample(x)  # [B, n_mels, T_up]
        return mel

class GriffinLimVocoder:
    """Simple Griffin-Lim vocoder as fallback when HiFi-GAN isn't available."""

    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):
        """
        Args:
            mel: [1, n_mels, T] log-mel spectrogram
        Returns:
            waveform: [1, samples]
        """
        # Un-log the mel
        mel = torch.exp(mel)
        # Inverse mel scale
        spec = self.inverse_mel(mel)
        # Griffin-Lim
        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()

        # Build mel decoder (needs to be trained alongside or separately)
        self.mel_decoder = MelDecoder(
            d_model=config["d_model"],
            n_mels=config["n_mels"],
        ).to(self.device)

        # Try loading mel decoder weights
        decoder_path = Path(checkpoint_path).parent / "mel_decoder.pt"
        if decoder_path.exists():
            self.mel_decoder.load_state_dict(
                torch.load(decoder_path, map_location=self.device, weights_only=True)
            )
        self.mel_decoder.eval()

        # Vocoder
        self.vocoder = GriffinLimVocoder(
            n_fft=1024, hop_length=256, n_mels=80, sample_rate=24000,
        )

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

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

    @torch.no_grad()
    def synthesize(self, text, max_steps=500, temperature=0.7):
        """
        Generate speech from text.

        Args:
            text: Input text string (Hindi)
            max_steps: Maximum number of AR steps
            temperature: Sampling temperature (0 = deterministic)
        Returns:
            waveform: numpy array
        """
        # Encode text
        text_tokens = self.text_to_tokens(text)
        text_emb = self.model.text_encoder(text_tokens)

        # Initialize with a learned start embedding
        d = self.config["d_model"]
        # Start with zeros (will be projected)
        current_emb = torch.zeros(1, 1, d, device=self.device)

        # Autoregressive generation
        all_embeddings = [current_emb]

        for step in range(max_steps):
            # Concatenate all embeddings so far
            context = torch.cat(all_embeddings, dim=1)  # [1, t, d]

            # Predict next
            next_emb = self.model.predict_next(context, text_emb)  # [1, 1, d]

            # Add noise for diversity
            if temperature > 0:
                noise = torch.randn_like(next_emb) * temperature * 0.1
                next_emb = next_emb + noise

            all_embeddings.append(next_emb)

            # Simple stopping: if embedding norm is very small
            if next_emb.norm() < 0.01 and step > 10:
                break

        # Combine all embeddings
        audio_embs = torch.cat(all_embeddings[1:], dim=1)  # Skip the zero start token

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

        # Vocoder: mel → waveform
        waveform = self.vocoder(mel.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}")

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--checkpoint", required=True, help="Path to model checkpoint")
    parser.add_argument("--text", required=True, help="Text to synthesize")
    parser.add_argument("--output", default="output.wav", help="Output wav path")
    parser.add_argument("--max_steps", type=int, default=500)
    parser.add_argument("--temperature", type=float, default=0.7)
    args = parser.parse_args()

    tts = LeWMTTSInference(args.checkpoint)
    waveform, sr = tts.synthesize(
        args.text, max_steps=args.max_steps, temperature=args.temperature
    )
    tts.save_audio(waveform, sr, args.output)

if __name__ == "__main__":
    main()


Now let me test the full pipeline — verify everything works before running on real data.


Updated task #3 status


Updated task #1 status


Updated task #2 status


Updated task #5 status


Bash

cd /home/ubuntu/lewm-tts && python3 -c "
from model import build_model
import torch

# Test model build
model, config = build_model()

# Test forward pass with dummy data
B = 2
mel = torch.randn(B, 80, 400)       # 2 samples, 80 mels, 400 frames
text = torch.randint(0, 256, (B, 50)) # 50 byte tokens
mel_mask = torch.zeros(B, 400, dtype=torch.bool)
text_mask = torch.zeros(B, 50, dtype=torch.bool)

print('Forward pass...')
losses = model(mel, text, mel_mask, text_mask)
print(f'  total_loss: {losses[\"total_loss\"].item():.4f}')
print(f'  prediction_loss: {losses[\"prediction_loss\"].item():.4f}')
print(f'  kl_loss: {losses[\"kl_loss\"].item():.4f}')

# Test backward
losses['total_loss'].backward()
print('Backward pass OK')

# Test inference path
model.eval()
with torch.no_grad():
    text_emb = model.text_encoder(text[:1])
    audio_emb = torch.randn(1, 1, 256)
    next_emb = model.predict_next(audio_emb, text_emb)
    print(f'Inference predict_next shape: {next_emb.shape}')

print('\\nAll model tests passed!')
"

Test model forward/backward pass


LeWM TTS model: 13.77M parameters
Forward pass...
  total_loss: 1.4973
  prediction_loss: 1.4950
  kl_loss: 0.2324
Backward pass OK
Inference predict_next shape: torch.Size([1, 1, 256])

All model tests passed!


13.77M parameters, forward/backward pass works. Now let me run the preprocessing on the real dataset.


python3 preprocess.py --input /home/ubuntu/modi_dataset --output /home/ubuntu/lewm-tts/processed_data 2>&1

Preprocess Modi dataset into mel spectrograms


Loading dataset from /home/ubuntu/modi_dataset...
  Processed 100/2330 samples, 305 segments, 0.77h total
  Processed 200/2330 samples, 608 segments, 1.54h total
  Processed 300/2330 samples, 916 segments, 2.31h total
  Processed 400/2330 samples, 1216 segments, 3.08h total
  Processed 500/2330 samples, 1518 segments, 3.84h total
  Processed 600/2330 samples, 1821 segments, 4.61h total
  Processed 700/2330 samples, 2124 segments, 5.38h total
  Processed 800/2330 samples, 2417 segments, 6.13h total
  Processed 900/2330 samples, 2711 segments, 6.88h total
  Processed 1000/2330 samples, 3011 segments, 7.66h total
  Processed 1100/2330 samples, 3309 segments, 8.42h total
  Processed 1200/2330 samples, 3605 segments, 9.19h total
  Processed 1300/2330 samples, 3904 segments, 9.96h total
  Processed 1400/2330 samples, 4200 segments, 10.71h total
  Processed 1500/2330 samples, 4497 segments, 11.47h total
  Processed 1600/2330 samples, 4797 segments, 12.24h total
  Processed 1700/2330 samples, 5098 segments, 13.01h total
  Processed 1800/2330 samples, 5402 segments, 13.78h total
  Processed 1900/2330 samples, 5714 segments, 14.56h total
  Processed 2000/2330 samples, 6016 segments, 15.34h total
  Processed 2100/2330 samples, 6317 segments, 16.11h total
  Processed 2200/2330 samples, 6621 segments, 16.89h total
  Processed 2300/2330 samples, 6922 segments, 17.66h total

Done! 7009 segments, 17.88 hours
Manifest: /home/ubuntu/lewm-tts/processed_data/manifest.json
Config: /home/ubuntu/lewm-tts/processed_data/config.json


Preprocessing complete — 7,009 segments from 2,330 clips. Now let me verify the data quality and test the training loop.


python3 -c "
import json
import numpy as np

with open('processed_data/config.json') as f:
    config = json.load(f)
print('Config:', json.dumps(config, indent=2))

with open('processed_data/manifest.json', encoding='utf-8') as f:
    manifest = json.load(f)

durations = [e['duration'] for e in manifest]
mel_frames = [e['mel_frames'] for e in manifest]
text_lens = [len(e['text'].encode('utf-8')) for e in manifest]

print(f'\nSegment stats:')
print(f'  Duration: min={min(durations):.1f}s, max={max(durations):.1f}s, mean={np.mean(durations):.1f}s, median={np.median(durations):.1f}s')
print(f'  Mel frames: min={min(mel_frames)}, max={max(mel_frames)}, mean={np.mean(mel_frames):.0f}')
print(f'  Text bytes: min={min(text_lens)}, max={max(text_lens)}, mean={np.mean(text_lens):.0f}')

# Check how many fit in max_mel_frames=1200
fit = sum(1 for m in mel_frames if m <= 1200)
print(f'  Segments with <= 1200 mel frames: {fit}/{len(mel_frames)} ({100*fit/len(mel_frames):.1f}%)')
"

Verify preprocessed data stats


Config: {
  "mel": {
    "sample_rate": 24000,
    "n_fft": 1024,
    "hop_length": 256,
    "n_mels": 80,
    "f_min": 0,
    "f_max": 12000,
    "power": 1.0
  },
  "segmentation": {
    "min_duration": 3.0,
    "max_duration": 12.0,
    "target_duration": 8.0,
    "silence_threshold_db": -40,
    "min_silence_len": 0.3
  },
  "total_segments": 7009,
  "total_duration_hours": 17.88,
  "original_samples": 2330
}

Segment stats:
  Duration: min=3.0s, max=14.4s, mean=9.2s, median=9.2s
  Mel frames: min=282, max=1350, mean=861
  Text bytes: min=24, max=573, mean=218
  Segments with <= 1200 mel frames: 6959/7009 (99.3%)


99.3% of segments fit within the max frame limit. Now let me test the full training loop for a few steps.


python3 -c "
from dataset import build_dataloader
from model import build_model
import torch

# Test dataloader
loader, dataset = build_dataloader('processed_data/manifest.json', batch_size=8, num_workers=0)
batch = next(iter(loader))

print('Batch shapes:')
for k, v in batch.items():
    print(f'  {k}: {v.shape} ({v.dtype})')

# Test training step on GPU
device = torch.device('cuda')
model, config = build_model()
model = model.to(device)
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)

model.train()
for step in range(5):
    mel = batch['mel'].to(device)
    text = batch['text_tokens'].to(device)
    mel_mask = batch['mel_mask'].to(device)
    text_mask = batch['text_mask'].to(device)

    losses = model(mel, text, mel_mask, text_mask)
    losses['total_loss'].backward()
    optimizer.step()
    optimizer.zero_grad()

    print(f'Step {step}: loss={losses[\"total_loss\"].item():.4f} pred={losses[\"prediction_loss\"].item():.4f} kl={losses[\"kl_loss\"].item():.4f}')

print('\\nTraining loop test passed!')
" 2>&1

Test dataloader + training loop on GPU


Dataset: 6959 samples loaded
Batch shapes:
  mel: torch.Size([8, 80, 992]) (torch.float32)
  text_tokens: torch.Size([8, 221]) (torch.int64)
  mel_mask: torch.Size([8, 992]) (torch.bool)
  text_mask: torch.Size([8, 221]) (torch.bool)
LeWM TTS model: 13.77M parameters
Step 0: loss=1.5321 pred=1.5294 kl=0.2694
Step 1: loss=1.3030 pred=1.3008 kl=0.2228
Step 2: loss=1.1513 pred=1.1493 kl=0.1918
Step 3: loss=1.0426 pred=1.0408 kl=0.1729
Step 4: loss=0.9656 pred=0.9639 kl=0.1622

Training loop test passed!


Loss is dropping nicely even in 5 steps. Let me also test the inference pipeline.


python3 -c "
from inference import MelDecoder, GriffinLimVocoder
import torch

# Test mel decoder
decoder = MelDecoder(d_model=256, n_mels=80)
z = torch.randn(1, 50, 256)  # 50 downsampled frames
mel = decoder(z)
print(f'MelDecoder: input {z.shape} → output {mel.shape}')  # Expect [1, 80, 200]

# Test vocoder
vocoder = GriffinLimVocoder()
waveform = vocoder(mel)
print(f'Vocoder: mel {mel.shape} → waveform {waveform.shape}')
print(f'Duration: {waveform.shape[-1]/24000:.2f}s')

print('\\nInference pipeline test passed!')
" 2>&1

Test inference pipeline


MelDecoder: input torch.Size([1, 50, 256]) → output torch.Size([1, 80, 200])
Vocoder: mel torch.Size([1, 80, 200]) → waveform torch.Size([1, 50944])
Duration: 2.12s

Inference pipeline test passed!


Everything works end-to-end. Now let me review the code once more for issues and add the mel decoder to training (it needs to be jointly trained).


     1→"""
     2→LeWM TTS Model — JEPA-based Text-to-Speech
     3→
     4→Architecture:
     5→  - TextEncoder: Character-level transformer for Hindi/Devanagari text
     6→  - AudioEncoder: 1D CNN + Transformer that encodes mel spectrograms to embeddings
     7→  - JEPAPredictor: Transformer that predicts next audio embedding given text + audio context
     8→  - Two losses: next-embedding prediction (MSE) + Gaussian regularizer (KL to N(0,I))
     9→"""
    10→
    11→import math
    12→import torch
    13→import torch.nn as nn
    14→import torch.nn.functional as F
    15→
    16→
    17→# ─── Text Encoder ───────────────────────────────────────────────────────────
    18→
    19→class TextEncoder(nn.Module):
    20→    """Character-level transformer encoder for Hindi text."""
    21→
    22→    def __init__(self, vocab_size=256, d_model=256, nhead=4, num_layers=4, dropout=0.1):
    23→        super().__init__()
    24→        self.d_model = d_model
    25→
    26→        # Character embedding (covers all Unicode Hindi chars via byte-level encoding)
    27→        self.char_embed = nn.Embedding(vocab_size, d_model)
    28→        self.pos_embed = SinusoidalPositionalEncoding(d_model, max_len=2048)
    29→        self.dropout = nn.Dropout(dropout)
    30→
    31→        encoder_layer = nn.TransformerEncoderLayer(
    32→            d_model=d_model, nhead=nhead,
    33→            dim_feedforward=d_model * 4,
    34→            dropout=dropout, batch_first=True,
    35→            activation="gelu",
    36→        )
    37→        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
    38→        self.proj = nn.Linear(d_model, d_model)
    39→
    40→    def forward(self, text_tokens, text_mask=None):
    41→        """
    42→        Args:
    43→            text_tokens: [B, T_text] long tensor of byte-level token IDs
    44→            text_mask: [B, T_text] bool, True = padding
    45→        Returns:
    46→            text_emb: [B, T_text, d_model]
    47→        """
    48→        x = self.char_embed(text_tokens) * math.sqrt(self.d_model)
    49→        x = self.pos_embed(x)
    50→        x = self.dropout(x)

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


I need to add the MelDecoder to the model so it trains jointly — this is critical for good inference. Let me update the model and training to include a reconstruction loss for the decoder.


   140→
   141→    Uses cross-attention to condition on text, and causal self-attention
   142→    over the audio embedding sequence.
   143→    """
   144→
   145→    def __init__(self, d_model=256, nhead=4, num_layers=6, dropout=0.1):
   146→        super().__init__()
   147→        self.d_model = d_model
   148→
   149→        self.audio_input_proj = nn.Linear(d_model, d_model)
   150→        self.pos_embed = SinusoidalPositionalEncoding(d_model, max_len=4096)
   151→        self.dropout = nn.Dropout(dropout)
   152→
   153→        # Transformer decoder with causal self-attention + cross-attention to text
   154→        decoder_layer = nn.TransformerDecoderLayer(
   155→            d_model=d_model, nhead=nhead,
   156→            dim_feedforward=d_model * 4,
   157→            dropout=dropout, batch_first=True,
   158→            activation="gelu",
   159→        )
   160→        self.transformer = nn.TransformerDecoder(decoder_layer, num_layers=num_layers)
   161→
   162→        # Output projection to predict next embedding
   163→        self.output_proj = nn.Sequential(
   164→            nn.Linear(d_model, d_model),
   165→            nn.GELU(),
   166→            nn.Linear(d_model, d_model),
   167→        )
   168→
   169→    def forward(self, audio_emb, text_emb, audio_mask=None, text_mask=None):
   170→        """
   171→        Args:
   172→            audio_emb: [B, T_audio, d_model] — audio embeddings (from encoder)
   173→            text_emb: [B, T_text, d_model] — text embeddings
   174→            audio_mask: [B, T_audio] bool, True = padding
   175→            text_mask: [B, T_text] bool, True = padding
   176→        Returns:
   177→            predicted: [B, T_audio, d_model] — predicted next embeddings
   178→        """
   179→        x = self.audio_input_proj(audio_emb)
   180→        x = self.pos_embed(x)
   181→        x = self.dropout(x)
   182→
   183→        # Causal mask: each position can only attend to previous positions
   184→        T = x.shape[1]
   185→        causal_mask = torch.triu(
   186→            torch.ones(T, T, device=x.device, dtype=torch.bool), diagonal=1
   187→        )
   188→
   189→        predicted = self.transformer(
   190→            x, text_emb,
   191→            tgt_mask=causal_mask,
   192→            tgt_key_padding_mask=audio_mask,
   193→            memory_key_padding_mask=text_mask,
   194→        )
   195→
   196→        predicted = self.output_proj(predicted)
   197→        return predicted
   198→
   199→
   200→# ─── Positional Encoding ────────────────────────────────────────────────────
   201→
   202→class SinusoidalPositionalEncoding(nn.Module):
   203→    def __init__(self, d_model, max_len=8192):
   204→        super().__init__()
   205→        pe = torch.zeros(max_len, d_model)
   206→        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
   207→        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
   208→        pe[:, 0::2] = torch.sin(position * div_term)
   209→        pe[:, 1::2] = torch.cos(position * div_term)
   210→        self.register_buffer("pe", pe.unsqueeze(0))  # [1, max_len, d_model]
   211→
   212→    def forward(self, x):
   213→        return x + self.pe[:, :x.shape[1]]
   214→
   215→
   216→# ─── Full LeWM TTS Model ────────────────────────────────────────────────────
   217→
   218→class LeWMTTS(nn.Module):
   219→    """
   220→    Complete LeWM TTS model.
   221→
   222→    Training: text + mel → audio embeddings → predict next embedding
   223→    Two losses:
   224→        1. Next-embedding prediction: MSE(predicted[t], target[t+1])
   225→        2. Gaussian regularizer: KL(q(z|x) || N(0, I))
   226→    """
   227→
   228→    def __init__(self, config):
   229→        super().__init__()
   230→        d_model = config.get("d_model", 256)
   231→        nhead = config.get("nhead", 4)
   232→        n_mels = config.get("n_mels", 80)
   233→        text_vocab = config.get("text_vocab_size", 256)
   234→        text_layers = config.get("text_encoder_layers", 4)
   235→        audio_layers = config.get("audio_encoder_layers", 4)
   236→        predictor_layers = config.get("predictor_layers", 6)
   237→        dropout = config.get("dropout", 0.1)
   238→        self.kl_weight = config.get("kl_weight", 0.01)
   239→
   240→        self.text_encoder = TextEncoder(
   241→            vocab_size=text_vocab, d_model=d_model, nhead=nhead,
   242→            num_layers=text_layers, dropout=dropout,
   243→        )
   244→        self.audio_encoder = AudioEncoder(
   245→            n_mels=n_mels, d_model=d_model, nhead=nhead,
   246→            num_layers=audio_layers, dropout=dropout,
   247→        )
   248→        self.predictor = JEPAPredictor(
   249→            d_model=d_model, nhead=nhead,
   250→            num_layers=predictor_layers, dropout=dropout,
   251→        )
   252→
   253→        self.config = config
   254→
   255→    def forward(self, mel, text_tokens, mel_mask=None, text_mask=None):
   256→        """
   257→        Args:
   258→            mel: [B, n_mels, T_mel]
   259→            text_tokens: [B, T_text]
   260→            mel_mask: [B, T_mel] bool
   261→            text_mask: [B, T_text] bool
   262→        Returns:
   263→            loss_dict with 'prediction_loss', 'kl_loss', 'total_loss'
   264→        """
   265→        # Encode text
   266→        text_emb = self.text_encoder(text_tokens, text_mask)
   267→
   268→        # Encode audio → latent embeddings
   269→        z, mu, logvar = self.audio_encoder(mel, mel_mask)
   270→
   271→        # Predict next embedding: input z[:-1], target z[1:]
   272→        input_emb = z[:, :-1]     # [B, T-1, d]
   273→        target_emb = z[:, 1:]     # [B, T-1, d]
   274→
   275→        # Adjust mask for downsampled sequence
   276→        if mel_mask is not None:
   277→            T_down = z.shape[1]
   278→            ds_mask = mel_mask[:, ::self.audio_encoder.downsample_factor][:, :T_down]
   279→            pred_mask = ds_mask[:, :-1]
   280→        else:
   281→            pred_mask = None
   282→
   283→        predicted = self.predictor(input_emb, text_emb, pred_mask, text_mask)
   284→
   285→        # ─── Loss 1: Next-embedding prediction (MSE) ───
   286→        # Mask out padding positions
   287→        if pred_mask is not None:
   288→            valid = ~pred_mask  # True = valid
   289→            valid = valid.unsqueeze(-1)  # [B, T-1, 1]
   290→            prediction_loss = F.mse_loss(
   291→                predicted * valid, target_emb * valid, reduction="sum"
   292→            ) / (valid.sum() * predicted.shape[-1] + 1e-8)
   293→        else:
   294→            prediction_loss = F.mse_loss(predicted, target_emb)
   295→
   296→        # ─── Loss 2: Gaussian regularizer — KL(N(mu, sigma²) || N(0, I)) ───
   297→        # KL = -0.5 * sum(1 + logvar - mu² - exp(logvar))
   298→        if mel_mask is not None:
   299→            T_down = mu.shape[1]
   300→            ds_mask = mel_mask[:, ::self.audio_encoder.downsample_factor][:, :T_down]
   301→            valid = (~ds_mask).unsqueeze(-1).float()
   302→            kl = -0.5 * (1 + logvar - mu.pow(2) - logvar.exp())
   303→            kl_loss = (kl * valid).sum() / (valid.sum() * mu.shape[-1] + 1e-8)
   304→        else:
   305→            kl_loss = -0.5 * torch.mean(1 + logvar - mu.pow(2) - logvar.exp())
   306→
   307→        total_loss = prediction_loss + self.kl_weight * kl_loss
   308→
   309→        return {
   310→            "total_loss": total_loss,
   311→            "prediction_loss": prediction_loss,
   312→            "kl_loss": kl_loss,
   313→        }
   314→
   315→    def encode_audio(self, mel):
   316→        """Encode mel to latent embeddings (for inference)."""
   317→        z, mu, _ = self.audio_encoder(mel)
   318→        return mu  # Use mean at inference
   319→
   320→    def predict_next(self, audio_emb, text_emb, text_mask=None):
   321→        """Predict next audio embedding autoregressively."""
   322→        predicted = self.predictor(audio_emb, text_emb, text_mask=text_mask)
   323→        return predicted[:, -1:]  # Only the last prediction
   324→
   325→
   326→def count_parameters(model):
   327→    return sum(p.numel() for p in model.parameters() if p.requires_grad)
   328→
   329→
   330→def build_model(config=None):
   331→    if config is None:
   332→        config = {
   333→            "d_model": 256,
   334→            "nhead": 4,
   335→            "n_mels": 80,
   336→            "text_vocab_size": 256,
   337→            "text_encoder_layers": 4,
   338→            "audio_encoder_layers": 4,
   339→            "predictor_layers": 6,
   340→            "dropout": 0.1,
   341→            "kl_weight": 0.01,
   342→        }
   343→    model = LeWMTTS(config)
   344→    print(f"LeWM TTS model: {count_parameters(model)/1e6:.2f}M parameters")
   345→    return model, config
   346→

    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

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

# ─── 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", 80)
        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)

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

        self.config = config

    def forward(self, mel, text_tokens, mel_mask=None, text_mask=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
        text_emb = self.text_encoder(text_tokens, text_mask)

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

        total_loss = prediction_loss + self.kl_weight * kl_loss

        return {
            "total_loss": total_loss,
            "prediction_loss": prediction_loss,
            "kl_loss": kl_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 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


I need to add a MelDecoder to the main model and include a reconstruction loss — this is essential. Without it, the latent space has no grounding to actual mel spectrograms, and inference won't produce usable audio.


Edit

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

        self.config = config

        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)

        self.config = config


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

        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)

        self.config = config

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

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

import math
import 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=80, 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

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

# ─── 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", 80)
        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)

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

        self.config = config

    def forward(self, mel, text_tokens, mel_mask=None, text_mask=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
        text_emb = self.text_encoder(text_tokens, text_mask)

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

        total_loss = prediction_loss + self.kl_weight * kl_loss

        return {
            "total_loss": total_loss,
            "prediction_loss": prediction_loss,
            "kl_loss": kl_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 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

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


Now add the MelDecoder class and reconstruction loss.


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

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

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

    def __init__(self, d_model=256, n_mels=80):
        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 ────────────────────────────────────────────────────


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

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

    def __init__(self, d_model=256, n_mels=80):
        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 ────────────────────────────────────────────────────

"""
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=80, 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

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

# ─── 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", 80)
        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)

        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)

        self.config = config

    def forward(self, mel, text_tokens, mel_mask=None, text_mask=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
        text_emb = self.text_encoder(text_tokens, text_mask)

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

        total_loss = prediction_loss + self.kl_weight * kl_loss

        return {
            "total_loss": total_loss,
            "prediction_loss": prediction_loss,
            "kl_loss": kl_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 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

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

+    """Decodes latent embeddings back to mel spectrograms (reverses AudioEncoder)."""

+    def __init__(self, d_model=256, n_mels=80):

+            nn.Linear(d_model, d_model),

+            nn.ConvTranspose1d(d_model, d_model, kernel_size=4, stride=2, padding=1),

+            nn.Conv1d(d_model, n_mels, kernel_size=7, padding=3),

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


Now add reconstruction loss to the forward method.


        total_loss = prediction_loss + self.kl_weight * kl_loss

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

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


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

"""
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=80, 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

# ─── 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=80):
        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", 80)
        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)

        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)

        self.config = config

    def forward(self, mel, text_tokens, mel_mask=None, text_mask=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
        text_emb = self.text_encoder(text_tokens, text_mask)

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

        total_loss = prediction_loss + self.kl_weight * kl_loss

        return {
            "total_loss": total_loss,
            "prediction_loss": prediction_loss,
            "kl_loss": kl_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 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

-        total_loss = prediction_loss + self.kl_weight * kl_loss

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

+            mel_recon = mel_recon[:, :, :T_mel]

+                mel_mask = mel_mask[:, :T_recon]

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

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

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

+            recon_loss = F.l1_loss(mel_recon, mel)

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


Now update train.py to log the reconstruction loss too.


    70→
    71→    # ─── Resume from checkpoint ───
    72→    start_epoch = 0
    73→    global_step = 0
    74→    if args.resume:
    75→        ckpt = torch.load(args.resume, map_location=device, weights_only=False)
    76→        model.load_state_dict(ckpt["model"])
    77→        optimizer.load_state_dict(ckpt["optimizer"])
    78→        scheduler.load_state_dict(ckpt["scheduler"])
    79→        start_epoch = ckpt["epoch"] + 1
    80→        global_step = ckpt["global_step"]
    81→        print(f"Resumed from epoch {start_epoch}, step {global_step}")
    82→
    83→    # ─── Training loop ───
    84→    model.train()
    85→    best_loss = float("inf")
    86→
    87→    for epoch in range(start_epoch, args.epochs):
    88→        epoch_loss = 0.0
    89→        epoch_pred_loss = 0.0
    90→        epoch_kl_loss = 0.0
    91→        t0 = time.time()
    92→
    93→        for batch_idx, batch in enumerate(loader):
    94→            mel = batch["mel"].to(device)
    95→            text_tokens = batch["text_tokens"].to(device)
    96→            mel_mask = batch["mel_mask"].to(device)
    97→            text_mask = batch["text_mask"].to(device)
    98→
    99→            # Forward
   100→            losses = model(mel, text_tokens, mel_mask, text_mask)
   101→
   102→            total_loss = losses["total_loss"]
   103→            pred_loss = losses["prediction_loss"]
   104→            kl_loss = losses["kl_loss"]
   105→
   106→            # Backward
   107→            optimizer.zero_grad()
   108→            total_loss.backward()
   109→
   110→            # Gradient clipping
   111→            grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
   112→
   113→            optimizer.step()
   114→            scheduler.step()
   115→
   116→            # Logging
   117→            epoch_loss += total_loss.item()
   118→            epoch_pred_loss += pred_loss.item()
   119→            epoch_kl_loss += kl_loss.item()

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

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

            # 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_loss = 0.0
        epoch_pred_loss = 0.0
        epoch_kl_loss = 0.0

        epoch_loss = 0.0
        epoch_pred_loss = 0.0
        epoch_kl_loss = 0.0
        epoch_recon_loss = 0.0