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

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

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


Edit

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

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


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

"""
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
        epoch_recon_loss = 0.0
        t0 = time.time()

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

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

+            recon_loss = losses["recon_loss"]


            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_loss += total_loss.item()
            epoch_pred_loss += pred_loss.item()
            epoch_kl_loss += kl_loss.item()
            epoch_recon_loss += recon_loss.item()
            global_step += 1

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

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


            epoch_loss += total_loss.item()
            epoch_pred_loss += pred_loss.item()
            epoch_kl_loss += kl_loss.item()
            epoch_recon_loss += recon_loss.item()
            global_step += 1

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

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

"""
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
        epoch_recon_loss = 0.0
        t0 = time.time()

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

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

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

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

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

            optimizer.step()
            scheduler.step()

            # Logging
            epoch_loss += total_loss.item()
            epoch_pred_loss += pred_loss.item()
            epoch_kl_loss += kl_loss.item()
            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)

+            epoch_recon_loss += recon_loss.item()

+                writer.add_scalar("train/recon_loss", recon_loss.item(), global_step)

-                    f"kl={kl_loss.item():.4f} | grad={grad_norm.item():.2f} | "

+                    f"kl={kl_loss.item():.4f} recon={recon_loss.item():.4f} | "

+                    f"grad={grad_norm.item():.2f} | {elapsed:.1f}s"


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

        avg_loss = epoch_loss / n_batches
        avg_pred = epoch_pred_loss / n_batches
        avg_kl = epoch_kl_loss / n_batches
        avg_recon = epoch_recon_loss / n_batches
        epoch_time = time.time() - t0

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


        avg_loss = epoch_loss / n_batches
        avg_pred = epoch_pred_loss / n_batches
        avg_kl = epoch_kl_loss / n_batches
        avg_recon = epoch_recon_loss / n_batches
        epoch_time = time.time() - t0

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

"""
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
        epoch_recon_loss = 0.0
        t0 = time.time()

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

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

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

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

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

            optimizer.step()
            scheduler.step()

            # Logging
            epoch_loss += total_loss.item()
            epoch_pred_loss += pred_loss.item()
            epoch_kl_loss += kl_loss.item()
            epoch_recon_loss += recon_loss.item()
            global_step += 1

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

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

        # Epoch summary
        n_batches = len(loader)
        avg_loss = epoch_loss / n_batches
        avg_pred = epoch_pred_loss / n_batches
        avg_kl = epoch_kl_loss / n_batches
        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)

+        avg_recon = epoch_recon_loss / n_batches

-            f"avg_loss={avg_loss:.4f} pred={avg_pred:.4f} kl={avg_kl:.4f} | "

+            f"avg_loss={avg_loss:.4f} pred={avg_pred:.4f} kl={avg_kl:.4f} "

+            f"recon={avg_recon:.4f} | {epoch_time:.1f}s"


Also add `recon_weight` to the config in train.py.


        "kl_weight": args.kl_weight,
    }

        "kl_weight": args.kl_weight,
        "recon_weight": args.recon_weight,
    }


        "kl_weight": args.kl_weight,
    }

        "kl_weight": args.kl_weight,
        "recon_weight": args.recon_weight,
    }

"""
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
        epoch_recon_loss = 0.0
        t0 = time.time()

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

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

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

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

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

            optimizer.step()
            scheduler.step()

            # Logging
            epoch_loss += total_loss.item()
            epoch_pred_loss += pred_loss.item()
            epoch_kl_loss += kl_loss.item()
            epoch_recon_loss += recon_loss.item()
            global_step += 1

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

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

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

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

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

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

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

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    # Data
    parser.add_argument("--data_dir", default="/home/ubuntu/lewm-tts/processed_data")
    parser.add_argument("--output_dir", default="/home/ubuntu/lewm-tts/output")
    # Model
    parser.add_argument("--d_model", type=int, default=256)
    parser.add_argument("--nhead", type=int, default=4)
    parser.add_argument("--text_layers", type=int, default=4)
    parser.add_argument("--audio_layers", type=int, default=4)
    parser.add_argument("--predictor_layers", type=int, default=6)
    parser.add_argument("--dropout", type=float, default=0.1)
    parser.add_argument("--kl_weight", type=float, default=0.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)

+        "recon_weight": args.recon_weight,


    parser.add_argument("--kl_weight", type=float, default=0.01)
    parser.add_argument("--recon_weight", type=float, default=1.0)


    parser.add_argument("--kl_weight", type=float, default=0.01)
    parser.add_argument("--recon_weight", type=float, default=1.0)

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

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

from model import build_model
from dataset import build_dataloader

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

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

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

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

    # Warmup + cosine decay
    total_steps = args.epochs * len(loader)
    warmup_steps = min(4000, total_steps // 10)

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

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

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

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

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

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

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

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

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

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

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

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

            optimizer.step()
            scheduler.step()

            # Logging
            epoch_loss += total_loss.item()
            epoch_pred_loss += pred_loss.item()
            epoch_kl_loss += kl_loss.item()
            epoch_recon_loss += recon_loss.item()
            global_step += 1

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

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

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

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

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

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

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

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    # Data
    parser.add_argument("--data_dir", default="/home/ubuntu/lewm-tts/processed_data")
    parser.add_argument("--output_dir", default="/home/ubuntu/lewm-tts/output")
    # Model
    parser.add_argument("--d_model", type=int, default=256)
    parser.add_argument("--nhead", type=int, default=4)
    parser.add_argument("--text_layers", type=int, default=4)
    parser.add_argument("--audio_layers", type=int, default=4)
    parser.add_argument("--predictor_layers", type=int, default=6)
    parser.add_argument("--dropout", type=float, default=0.1)
    parser.add_argument("--kl_weight", type=float, default=0.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)

+    parser.add_argument("--recon_weight", type=float, default=1.0)


Now update inference.py to use the model's built-in MelDecoder instead of a separate one.


    60→        self.inverse_mel = torchaudio.transforms.InverseMelScale(
    61→            n_stft=n_fft // 2 + 1,
    62→            n_mels=n_mels,
    63→            sample_rate=sample_rate,
    64→        )
    65→        self.griffin_lim = torchaudio.transforms.GriffinLim(
    66→            n_fft=n_fft,
    67→            hop_length=hop_length,
    68→            n_iter=64,
    69→            power=1.0,
    70→        )
    71→
    72→    def __call__(self, mel):
    73→        """
    74→        Args:
    75→            mel: [1, n_mels, T] log-mel spectrogram
    76→        Returns:
    77→            waveform: [1, samples]
    78→        """
    79→        # Un-log the mel
    80→        mel = torch.exp(mel)
    81→        # Inverse mel scale
    82→        spec = self.inverse_mel(mel)
    83→        # Griffin-Lim
    84→        waveform = self.griffin_lim(spec)
    85→        return waveform
    86→
    87→
    88→class LeWMTTSInference:
    89→    """Full inference pipeline."""
    90→
    91→    def __init__(self, checkpoint_path, device="cuda"):
    92→        self.device = torch.device(device if torch.cuda.is_available() else "cpu")
    93→
    94→        # Load checkpoint
    95→        ckpt = torch.load(checkpoint_path, map_location=self.device, weights_only=False)
    96→        config = ckpt["config"]
    97→
    98→        # Build model
    99→        self.model = LeWMTTS(config).to(self.device)

        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)


    99→        self.model = LeWMTTS(config).to(self.device)
   100→        self.model.load_state_dict(ckpt["model"])
   101→        self.model.eval()
   102→
   103→        # Build mel decoder (needs to be trained alongside or separately)
   104→        self.mel_decoder = MelDecoder(
   105→            d_model=config["d_model"],
   106→            n_mels=config["n_mels"],
   107→        ).to(self.device)
   108→
   109→        # Try loading mel decoder weights
   110→        decoder_path = Path(checkpoint_path).parent / "mel_decoder.pt"
   111→        if decoder_path.exists():
   112→            self.mel_decoder.load_state_dict(
   113→                torch.load(decoder_path, map_location=self.device, weights_only=True)
   114→            )
   115→        self.mel_decoder.eval()
   116→
   117→        # Vocoder
   118→        self.vocoder = GriffinLimVocoder(
   119→            n_fft=1024, hop_length=256, n_mels=80, sample_rate=24000,
   120→        )
   121→
   122→        self.config = config
   123→        print(f"Model loaded from {checkpoint_path}")
   124→
   125→    def text_to_tokens(self, text):
   126→        tokens = list(text.encode("utf-8"))
   127→        return torch.tensor(tokens, dtype=torch.long).unsqueeze(0).to(self.device)
   128→
   129→    @torch.no_grad()
   130→    def synthesize(self, text, max_steps=500, temperature=0.7):
   131→        """
   132→        Generate speech from text.
   133→
   134→        Args:
   135→            text: Input text string (Hindi)
   136→            max_steps: Maximum number of AR steps
   137→            temperature: Sampling temperature (0 = deterministic)
   138→        Returns:
   139→            waveform: numpy array
   140→        """
   141→        # Encode text
   142→        text_tokens = self.text_to_tokens(text)
   143→        text_emb = self.model.text_encoder(text_tokens)
   144→
   145→        # Initialize with a learned start embedding
   146→        d = self.config["d_model"]
   147→        # Start with zeros (will be projected)
   148→        current_emb = torch.zeros(1, 1, d, device=self.device)
   149→
   150→        # Autoregressive generation
   151→        all_embeddings = [current_emb]
   152→
   153→        for step in range(max_steps):
   154→            # Concatenate all embeddings so far
   155→            context = torch.cat(all_embeddings, dim=1)  # [1, t, d]
   156→
   157→            # Predict next
   158→            next_emb = self.model.predict_next(context, text_emb)  # [1, 1, d]
   159→
   160→            # Add noise for diversity
   161→            if temperature > 0:
   162→                noise = torch.randn_like(next_emb) * temperature * 0.1
   163→                next_emb = next_emb + noise
   164→
   165→            all_embeddings.append(next_emb)
   166→
   167→            # Simple stopping: if embedding norm is very small
   168→            if next_emb.norm() < 0.01 and step > 10:
   169→                break
   170→
   171→        # Combine all embeddings
   172→        audio_embs = torch.cat(all_embeddings[1:], dim=1)  # Skip the zero start token
   173→
   174→        # Decode to mel
   175→        mel = self.mel_decoder(audio_embs)
   176→
   177→        # Vocoder: mel → waveform
   178→        waveform = self.vocoder(mel.cpu())
   179→
   180→        return waveform.squeeze().numpy(), 24000
   181→
   182→    def save_audio(self, waveform, sr, output_path):
   183→        wav_tensor = torch.from_numpy(waveform).unsqueeze(0)
   184→        torchaudio.save(output_path, wav_tensor, sr)
   185→        print(f"Saved: {output_path}")
   186→
   187→
   188→def main():
   189→    parser = argparse.ArgumentParser()
   190→    parser.add_argument("--checkpoint", required=True, help="Path to model checkpoint")
   191→    parser.add_argument("--text", required=True, help="Text to synthesize")
   192→    parser.add_argument("--output", default="output.wav", help="Output wav path")
   193→    parser.add_argument("--max_steps", type=int, default=500)
   194→    parser.add_argument("--temperature", type=float, default=0.7)
   195→    args = parser.parse_args()
   196→
   197→    tts = LeWMTTSInference(args.checkpoint)
   198→    waveform, sr = tts.synthesize(
   199→        args.text, max_steps=args.max_steps, temperature=args.temperature
   200→    )
   201→    tts.save_audio(waveform, sr, args.output)
   202→
   203→
   204→if __name__ == "__main__":
   205→    main()
   206→

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


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

# MelDecoder is now part of the model (trained jointly)


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

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

-        # Build mel decoder (needs to be trained alongside or separately)

+        # MelDecoder is now part of the model (trained jointly)

-        # Try loading mel decoder weights

-        decoder_path = Path(checkpoint_path).parent / "mel_decoder.pt"

-            self.mel_decoder.load_state_dict(

-                torch.load(decoder_path, map_location=self.device, weights_only=True)


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

        # Decode to mel using the model's built-in decoder
        mel = self.model.mel_decoder(audio_embs)


        # Decode to mel using the model's built-in decoder
        mel = self.model.mel_decoder(audio_embs)

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

        # MelDecoder is now part of the model (trained jointly)

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

-        mel = self.mel_decoder(audio_embs)

+        # Decode to mel using the model's built-in decoder

+        mel = self.model.mel_decoder(audio_embs)


Also remove the now-unused MelDecoder import/class from inference.py since it's in model.py.


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


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

        # MelDecoder is now part of the model (trained jointly)

        # 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 using the model's built-in decoder
        mel = self.model.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()

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

-            nn.Linear(d_model, d_model),

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

-            # Upsample by 2 again (total 4x)

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

-            mel: [B, n_mels, T_up] — reconstructed mel spectrogram

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

-        mel = self.upsample(x)  # [B, n_mels, T_up]


Now let me re-test everything with the updated model.


Bash

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

# Build model
model, config = build_model()
device = torch.device('cuda')
model = model.to(device)

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

model.train()
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)

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()
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    optimizer.step()
    optimizer.zero_grad()

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

# Test inference path
model.eval()
with torch.no_grad():
    text_tokens = torch.tensor([list('test'.encode('utf-8'))], dtype=torch.long).to(device)
    text_emb = model.text_encoder(text_tokens)
    audio_emb = torch.randn(1, 10, 256).to(device)
    next_emb = model.predict_next(audio_emb, text_emb)
    mel_out = model.mel_decoder(torch.cat([audio_emb, next_emb], dim=1))
    print(f'Inference: audio_emb {audio_emb.shape} → mel {mel_out.shape}')

print(f'\\nTotal params: {count_parameters(model)/1e6:.2f}M')
print('All tests passed!')
" 2>&1

Test updated model with reconstruction loss


LeWM TTS model: 14.50M parameters
Dataset: 6959 samples loaded
Step 0: total=3.7948 pred=1.5424 kl=0.2635 recon=2.2498
Step 1: total=3.5743 pred=1.3237 kl=0.2201 recon=2.2484
Step 2: total=3.4397 pred=1.1905 kl=0.1968 recon=2.2472
Step 3: total=3.3433 pred=1.0954 kl=0.1818 recon=2.2460
Step 4: total=3.2668 pred=1.0204 kl=0.1709 recon=2.2447
Inference: audio_emb torch.Size([1, 10, 256]) → mel torch.Size([1, 80, 44])

Total params: 14.50M
All tests passed!


14.5M parameters, all 3 losses working, inference pipeline functional. Let me now do a quick training run (5 epochs) to verify training stability.


python3 train.py --epochs 3 --batch_size 16 --print_every 20 --log_every 5 --save_every 1 --num_workers 4 2>&1

Quick 3-epoch training run to verify stability


Device: cuda
LeWM TTS model: 14.50M parameters
Dataset: 6959 samples loaded
Training: 6959 samples, 434 batches/epoch
  [1/3] step 20 | loss=3.5727 pred=1.3530 kl=0.2258 recon=2.2175 | grad=1.76 | 2.0s
  [1/3] step 40 | loss=3.1927 pred=0.8972 kl=0.1610 recon=2.2940 | grad=1.08 | 3.2s
  [1/3] step 60 | loss=2.8307 pred=0.5960 kl=0.1723 recon=2.2329 | grad=0.66 | 4.5s
  [1/3] step 80 | loss=2.3449 pred=0.4625 kl=0.2891 recon=1.8795 | grad=0.62 | 5.8s
  [1/3] step 100 | loss=1.9692 pred=0.3647 kl=0.3499 recon=1.6010 | grad=0.91 | 7.0s
  [1/3] step 120 | loss=1.5899 pred=0.3424 kl=0.4099 recon=1.2434 | grad=1.51 | 8.3s
  [1/3] step 140 | loss=1.4116 pred=0.2451 kl=0.4939 recon=1.1615 | grad=0.90 | 9.6s
  [1/3] step 160 | loss=1.3433 pred=0.2005 kl=0.5768 recon=1.1371 | grad=1.05 | 10.8s
  [1/3] step 180 | loss=1.2628 pred=0.1690 kl=0.6653 recon=1.0872 | grad=4.33 | 12.0s
  [1/3] step 200 | loss=1.1285 pred=0.1492 kl=0.7376 recon=0.9720 | grad=1.61 | 13.2s
  [1/3] step 220 | loss=1.0944 pred=0.1334 kl=0.8331 recon=0.9527 | grad=3.83 | 14.4s
  [1/3] step 240 | loss=1.0491 pred=0.1160 kl=0.9214 recon=0.9238 | grad=5.67 | 15.7s
  [1/3] step 260 | loss=0.9481 pred=0.1048 kl=0.9948 recon=0.8334 | grad=1.37 | 16.9s
  [1/3] step 280 | loss=0.9424 pred=0.0951 kl=1.0549 recon=0.8368 | grad=0.78 | 18.1s
  [1/3] step 300 | loss=0.9423 pred=0.0874 kl=1.1145 recon=0.8437 | grad=4.97 | 19.3s
  [1/3] step 320 | loss=0.9548 pred=0.0789 kl=1.1497 recon=0.8644 | grad=6.02 | 20.5s
  [1/3] step 340 | loss=0.9855 pred=0.0727 kl=1.2011 recon=0.9007 | grad=2.54 | 21.7s
  [1/3] step 360 | loss=0.8973 pred=0.0593 kl=1.2942 recon=0.8251 | grad=4.16 | 22.9s
  [1/3] step 380 | loss=0.8991 pred=0.0529 kl=1.3852 recon=0.8324 | grad=6.50 | 24.1s
  [1/3] step 400 | loss=0.8540 pred=0.0479 kl=1.4658 recon=0.7915 | grad=1.16 | 25.3s
  [1/3] step 420 | loss=0.8726 pred=0.0416 kl=1.5614 recon=0.8154 | grad=2.69 | 26.5s
Epoch 1/3 | avg_loss=1.5304 pred=0.2986 kl=0.8091 recon=1.2236 | 27.3s
  Saved checkpoint: /home/ubuntu/lewm-tts/output/checkpoints/epoch_0001.pt
  New best model (loss=1.5304)
  [2/3] step 440 | loss=0.8742 pred=0.0383 kl=1.6456 recon=0.8194 | grad=2.66 | 0.5s
  [2/3] step 460 | loss=0.8750 pred=0.0366 kl=1.7271 recon=0.8211 | grad=2.47 | 1.7s
  [2/3] step 480 | loss=0.8796 pred=0.0347 kl=1.8037 recon=0.8269 | grad=1.99 | 2.9s
  [2/3] step 500 | loss=0.8248 pred=0.0316 kl=1.8731 recon=0.7745 | grad=1.80 | 4.1s
  [2/3] step 520 | loss=0.8125 pred=0.0300 kl=1.9457 recon=0.7630 | grad=1.81 | 5.3s
  [2/3] step 540 | loss=0.7959 pred=0.0286 kl=2.0073 recon=0.7473 | grad=0.61 | 6.5s
  [2/3] step 560 | loss=0.8056 pred=0.0274 kl=2.0713 recon=0.7575 | grad=1.10 | 7.7s
  [2/3] step 580 | loss=0.7954 pred=0.0247 kl=2.1359 recon=0.7493 | grad=2.53 | 8.9s
  [2/3] step 600 | loss=0.7879 pred=0.0271 kl=2.1625 recon=0.7392 | grad=1.61 | 10.1s
  [2/3] step 620 | loss=0.7843 pred=0.0233 kl=2.2027 recon=0.7390 | grad=1.86 | 11.3s
  [2/3] step 640 | loss=0.8120 pred=0.0254 kl=2.2330 recon=0.7643 | grad=1.46 | 12.5s
  [2/3] step 660 | loss=0.7844 pred=0.0228 kl=2.2593 recon=0.7390 | grad=1.66 | 13.7s
  [2/3] step 680 | loss=0.7784 pred=0.0215 kl=2.2941 recon=0.7339 | grad=1.26 | 14.9s
  [2/3] step 700 | loss=0.7517 pred=0.0212 kl=2.3116 recon=0.7074 | grad=1.75 | 16.1s
  [2/3] step 720 | loss=0.7537 pred=0.0203 kl=2.3380 recon=0.7100 | grad=1.30 | 17.3s
  [2/3] step 740 | loss=0.7883 pred=0.0199 kl=2.3572 recon=0.7447 | grad=0.97 | 18.5s
  [2/3] step 760 | loss=0.7193 pred=0.0193 kl=2.3618 recon=0.6764 | grad=2.40 | 19.7s
  [2/3] step 780 | loss=0.7678 pred=0.0198 kl=2.3695 recon=0.7243 | grad=1.50 | 20.9s
  [2/3] step 800 | loss=0.7276 pred=0.0183 kl=2.3827 recon=0.6855 | grad=1.27 | 22.1s
  [2/3] step 820 | loss=0.7316 pred=0.0177 kl=2.3926 recon=0.6900 | grad=0.70 | 23.2s
  [2/3] step 840 | loss=0.7163 pred=0.0176 kl=2.3981 recon=0.6747 | grad=0.59 | 24.4s
  [2/3] step 860 | loss=0.7321 pred=0.0183 kl=2.3970 recon=0.6899 | grad=2.48 | 25.6s
Epoch 2/3 | avg_loss=0.7816 pred=0.0245 kl=2.1712 recon=0.7354 | 26.2s
  Saved checkpoint: /home/ubuntu/lewm-tts/output/checkpoints/epoch_0002.pt
  New best model (loss=0.7816)
  [3/3] step 880 | loss=0.7069 pred=0.0175 kl=2.4026 recon=0.6654 | grad=0.46 | 0.9s
  [3/3] step 900 | loss=0.7088 pred=0.0170 kl=2.4101 recon=0.6677 | grad=1.00 | 2.1s
  [3/3] step 920 | loss=0.7412 pred=0.0178 kl=2.4106 recon=0.6992 | grad=0.52 | 3.3s
  [3/3] step 940 | loss=0.7046 pred=0.0172 kl=2.4180 recon=0.6632 | grad=0.71 | 4.5s
  [3/3] step 960 | loss=0.6987 pred=0.0169 kl=2.4203 recon=0.6576 | grad=0.46 | 5.7s
  [3/3] step 980 | loss=0.7127 pred=0.0174 kl=2.4203 recon=0.6711 | grad=0.96 | 6.8s
  [3/3] step 1000 | loss=0.7740 pred=0.0165 kl=2.4292 recon=0.7332 | grad=0.48 | 8.0s
  [3/3] step 1020 | loss=0.7529 pred=0.0166 kl=2.4276 recon=0.7120 | grad=0.45 | 9.2s
  [3/3] step 1040 | loss=0.6999 pred=0.0168 kl=2.4239 recon=0.6588 | grad=0.75 | 10.4s
  [3/3] step 1060 | loss=0.7209 pred=0.0173 kl=2.4195 recon=0.6795 | grad=0.40 | 11.5s
  [3/3] step 1080 | loss=0.6981 pred=0.0159 kl=2.4304 recon=0.6579 | grad=0.51 | 12.7s
  [3/3] step 1100 | loss=0.6966 pred=0.0163 kl=2.4288 recon=0.6560 | grad=0.43 | 13.9s
  [3/3] step 1120 | loss=0.7354 pred=0.0165 kl=2.4266 recon=0.6946 | grad=0.46 | 15.0s
  [3/3] step 1140 | loss=0.7077 pred=0.0160 kl=2.4304 recon=0.6674 | grad=0.54 | 16.2s
  [3/3] step 1160 | loss=0.7134 pred=0.0165 kl=2.4265 recon=0.6727 | grad=0.20 | 17.3s
  [3/3] step 1180 | loss=0.7031 pred=0.0163 kl=2.4296 recon=0.6625 | grad=0.35 | 18.5s
  [3/3] step 1200 | loss=0.7292 pred=0.0161 kl=2.4294 recon=0.6888 | grad=0.26 | 19.7s
  [3/3] step 1220 | loss=0.7154 pred=0.0165 kl=2.4302 recon=0.6746 | grad=0.34 | 20.8s
  [3/3] step 1240 | loss=0.7048 pred=0.0164 kl=2.4292 recon=0.6641 | grad=0.23 | 22.0s
  [3/3] step 1260 | loss=0.7140 pred=0.0165 kl=2.4278 recon=0.6732 | grad=0.43 | 23.2s
  [3/3] step 1280 | loss=0.7225 pred=0.0164 kl=2.4298 recon=0.6818 | grad=0.43 | 24.3s
  [3/3] step 1300 | loss=0.6939 pred=0.0164 kl=2.4278 recon=0.6532 | grad=0.27 | 25.5s
Epoch 3/3 | avg_loss=0.7169 pred=0.0167 kl=2.4235 recon=0.6759 | 25.6s
  Saved checkpoint: /home/ubuntu/lewm-tts/output/checkpoints/epoch_0003.pt
  New best model (loss=0.7169)

Training complete. Best loss: 0.7169


Training is very stable. Key observations:

- **Prediction loss**: 1.54 → 0.017 (dropped 99% — model is learning to predict next embeddings well)
- **Reconstruction loss**: 2.25 → 0.68 (mel reconstruction improving steadily)
- **KL loss**: Growing as expected (latents becoming more Gaussian)
- **No NaNs, no spikes** — stable training
- **~27s per epoch** on A100 — full 200 epochs would take ~1.5 hours

There's one issue I want to fix: the KL loss is growing unbounded while prediction loss is tiny. I should adjust the KL weight to keep it in check. Let me also check the `kl_weight` balance.


python3 -c "
# KL is at 2.43, weighted at 0.01 = 0.024 contribution to loss
# pred at 0.017, recon at 0.68
# The KL is growing because the regularizer is too weak —
# the encoder is not being pushed toward N(0,I)
# Let's check: with kl_weight=0.01, kl contributes 0.024 to total
# recon contributes 0.68, pred 0.017
# Total = 0.717, KL is only 3.3% — too small
# Should increase kl_weight to ~0.1 for better balance
print('Current loss breakdown:')
print(f'  pred: 0.017 (2.4%)')
print(f'  kl*w: 2.43*0.01 = 0.024 (3.4%)')
print(f'  recon: 0.676 (94.2%)')
print(f'  total: 0.717')
print()
print('Recommendation: increase kl_weight to 0.1')
print('This gives:')
kl_contrib = 2.43 * 0.1
total = 0.017 + kl_contrib + 0.676
print(f'  pred: 0.017 ({0.017/total*100:.1f}%)')
print(f'  kl*w: {kl_contrib:.3f} ({kl_contrib/total*100:.1f}%)')
print(f'  recon: 0.676 ({0.676/total*100:.1f}%)')
print(f'  total: {total:.3f}')
"

Analyze loss balance


Current loss breakdown:
  pred: 0.017 (2.4%)
  kl*w: 2.43*0.01 = 0.024 (3.4%)
  recon: 0.676 (94.2%)
  total: 0.717

Recommendation: increase kl_weight to 0.1
This gives:
  pred: 0.017 (1.8%)
  kl*w: 0.243 (26.0%)
  recon: 0.676 (72.2%)
  total: 0.936


The KL is too weak at 0.01. Let me update the default and also add KL annealing (ramp up KL weight gradually) which is the standard practice for VAE-style models.


    83→
    84→    # ─── Training loop ───
    85→    model.train()
    86→    best_loss = float("inf")
    87→
    88→    for epoch in range(start_epoch, args.epochs):
    89→        epoch_loss = 0.0
    90→        epoch_pred_loss = 0.0
    91→        epoch_kl_loss = 0.0
    92→        epoch_recon_loss = 0.0
    93→        t0 = time.time()
    94→
    95→        for batch_idx, batch in enumerate(loader):
    96→            mel = batch["mel"].to(device)
    97→            text_tokens = batch["text_tokens"].to(device)
    98→            mel_mask = batch["mel_mask"].to(device)
    99→            text_mask = batch["text_mask"].to(device)
   100→
   101→            # Forward
   102→            losses = model(mel, text_tokens, mel_mask, text_mask)
   103→
   104→            total_loss = losses["total_loss"]
   105→            pred_loss = losses["prediction_loss"]
   106→            kl_loss = losses["kl_loss"]
   107→            recon_loss = losses["recon_loss"]
   108→
   109→            # Backward
   110→            optimizer.zero_grad()
   111→            total_loss.backward()
   112→
   113→            # Gradient clipping
   114→            grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
   115→
   116→            optimizer.step()
   117→            scheduler.step()

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

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

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

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

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

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

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

            optimizer.step()
            scheduler.step()


I'll add KL annealing — the KL weight gradually ramps up during training. This lets the model first learn good representations, then gradually enforce the Gaussian prior.


Edit

        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)

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

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

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